Compare commits

...
Author SHA1 Message Date
hsliuustc0106 81de3f3308 [Bugfix][MiniMax-H3] Apply pre-commit formatting
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-09-03 08:44:26 +08:00
hsliuustc0106 f9c9de8d16 [Bugfix][MiniMax-H3] Reconcile temporal decode plans across VAE ranks
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
(cherry picked from commit ac4ec9c0fcf3394e8476d3d2d913d7c86f5df8e3)
2026-09-03 08:20:00 +08:00
hsliuustc0106 ff2e33a329 [Docs][MiniMax-H3] Document temporal VAE chunks
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-09-02 07:08:40 +08:00
hsliuustc0106 0b5b0979d8 [Bugfix][MiniMax-H3] Reconcile chunk callback ownership
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-09-01 23:11:29 +08:00
hsliuustc0106 f04ff76c54 [Model][MiniMax-H3] Isolate temporal chunk coordination
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-09-01 23:11:29 +08:00
hsliuustc0106 b220df8397 [Model][MiniMax-H3] Support released VAE chunks directly
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-09-01 23:11:28 +08:00
hsliuustc0106 520e058b95 [Model][MiniMax-H3] Delegate temporal VAE chunks
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-09-01 23:11:27 +08:00
DanaerLeeandHongsheng Liu 0288c3f56e [Model] Support SenseNova-U1.5-8B-MoT and its distilled 8-step LoRA (#6516)
Signed-off-by: MrlixiangWE <mrdanaer@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-09-01 22:53:53 +08:00
andyluo7andHongsheng Liu 477083a103 [CI][ROCm] Fix deterministic AMD L3 blockers (#6884)
Signed-off-by: andyluo7 <andy.luo@amd.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-09-01 22:53:00 +08:00
Rahul SteigerandRahul Steiger e1b2f589c1 [Bugfix][Cosmos3] Avoid logging prompts at INFO level (#6913)
Signed-off-by: Rahul Steiger <rsteiger@nvl72d233-T16.cm.cluster>
Co-authored-by: Rahul Steiger <rsteiger@nvl72d233-T16.cm.cluster>
2026-09-01 22:52:15 +08:00
Hongsheng Liuandhsliu_ustc 44c8d1c449 docs: add concise API server endpoint guide (#6114)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com>
2026-09-01 20:34:09 +08:00
Andy Zhou 2298359f67 [Bugfix][Diffusion] Preserve stage attention backend (#6645)
Signed-off-by: AndyZhou952 <jzhoubc@connect.ust.hk>
2026-09-01 20:18:44 +08:00
Hongsheng Liu d1c8beeb18 [Doc] Update README and docs for v0.28.0 (#6858)
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-09-01 19:50:26 +08:00
WeiQing Chen e2b83db683 [Test] Add MiniMax-H3 DLO DP2 variants and Turbo LoRA L3 test (#6556)
Signed-off-by: david6666666 <530634352@qq.com>
2026-09-01 10:48:52 +00:00
AuFlowandHongsheng Liu 49c6fc1540 [Frontend] Add configurable output sample rate to Speech API (#6553)
Signed-off-by: 郑倩倩 <2458983186@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-09-01 17:25:09 +08:00
Alex Brooks a65e89cbd4 [Core] Validate PipelineConfig in __post_init__ (#6829)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
2026-09-01 11:17:57 +08:00
Yuxuan RhoandHongsheng Liu 7992c2213e [Config] Remove diffusion_batch_size and use max_num_seqs instead. (#6484)
Signed-off-by: LyxWxj <yuxuanlou621@foxmail.com>
Signed-off-by: Yuxuan Rho <104440472+LyxWxj@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-09-01 01:02:54 +00:00
XRay2016 8e26dce1f4 [Diffusion] Add Boogu Edit-Turbo support docs and tests (#6701)
Signed-off-by: xRay2016 <1150722393@qq.com>
2026-09-01 08:41:23 +08:00
106 changed files with 7080 additions and 507 deletions

No files matched your search

+4 -7
View File
@@ -238,13 +238,10 @@ steps:
grade: Blocking
timeout_in_minutes: 50
commands:
- |
timeout 40m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m 'advanced_model and cuda' --run-level 'advanced_model'
'
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- timeout 40m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model and cuda" --run-level "advanced_model"
# TODO: Bagel test on ROCm is very unstable. @tjtanaa
# Need to debug before reneable numerical changes across large PRs
+35
View File
@@ -237,3 +237,38 @@ steps:
pytest -s -v tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
"
mirror_hardwares: h100_1
- label: "Diffusion · MiniMax-H3 L3 DLO DP2 and FastH3"
source_file_dependencies:
- tests/e2e/online_serving/minimax_h3/
- tests/helpers/assertions.py
- tests/helpers/mark.py
- tests/helpers/media.py
- tests/helpers/runtime.py
- vllm_omni/diffusion/models/minimax_h3/
- vllm_omni/diffusion/lora/
- vllm_omni/diffusion/offloader/
- vllm_omni/diffusion/distributed/
- vllm_omni/entrypoints/openai/
# Keep four server lifecycles in one Buildkite job, but never in one
# pytest process. DLO uses multiprocessing shared-memory queues; the
# parent's resource_tracker can retain a stale /psm_* entry during
# teardown and strand the next Ref2VA collective.
# A per-file process plus a hard timeout also makes a regression fail
# instead of leaving the merge queue stuck indefinitely.
timeout_in_minutes: 150
commands:
- |
set -eu
export VLLM_IMAGE_FETCH_TIMEOUT=60
if [ -f vllm_omni/diffusion/models/minimax_h3/lora.py ]; then
export VLLM_TEST_MINIMAX_H3_TURBO_LORA=$(python3 -c 'from huggingface_hub import hf_hub_download; print(hf_hub_download(repo_id="lightx2v/Minimax-h3-Turbo", filename="minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"))')
fi
if [ -f vllm_omni/diffusion/models/minimax_h3/fasth3.py ]; then
export VLLM_TEST_MINIMAX_H3_FASTH3_LORA=$(python3 -c 'from huggingface_hub import hf_hub_download; print(hf_hub_download(repo_id="FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA", filename="dense-datafree/adapter_model.safetensors"))')
fi
timeout 35m pytest -s -v tests/e2e/online_serving/minimax_h3/test_minimax_h3_dlo_dp2_fl2va.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
timeout 35m pytest -s -v tests/e2e/online_serving/minimax_h3/test_minimax_h3_dlo_dp2_ref2va.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
timeout 35m pytest -s -v tests/e2e/online_serving/minimax_h3/test_minimax_h3_turbo_lora.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
timeout 35m pytest -s -v tests/e2e/online_serving/minimax_h3/test_minimax_h3_fasth3_lora.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
mirror_hardwares: h100_2
+5 -3
View File
@@ -16,6 +16,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
*Latest News* 🔥
- [2026/08] We released [0.28.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.28.0), featuring production-ready [MiniMax H3](recipes/MiniMaxAI/MiniMax-H3.md) serving on GPU and NPU, a unified AR/DiT paged KV cache runtime, and enhanced realtime full-duplex serving for the [MiniCPM-o series](recipes/OpenBMB/MiniCPM-o-4_5.md).
- [2026/08] [VeRL-Omni](https://github.com/verl-project/verl-omni) `v0.2.0` is released: faster diffusion RL powered by vLLM-Omni (request-level/step-wise batching with FA3), rebuilt Qwen3-Omni multimodal training (DPO & GSPO), plus LTX-2.3, Qwen-Image-Edit support and more. See the [release notes](https://github.com/verl-project/verl-omni/releases/tag/v0.2.0).
- [2026/08] We released [0.26.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.26.0) - aligned with the vLLM 0.26 release line, featuring [MiniMax H3](recipes/MiniMaxAI/MiniMax-H3.md) joint video/audio generation, an experimental full-duplex realtime runtime for [MiniCPM-o 4.5](recipes/OpenBMB/MiniCPM-o-4_5.md), distributed layerwise diffusion offload, and broader model, hardware, streaming, TTS, and quantization support.
- [2026/07] We released [0.24.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.24.0) - aligned with the vLLM 0.24 release line, expanding production-ready coverage across TTS, speech, diffusion, image/video generation, and robot-policy serving, with major Omni stage runtime refactoring, diffusion request-level batching, async output materialization, quantization/cache/memory improvements, and broad CUDA/ROCm/XPU/NPU support.
@@ -57,9 +58,9 @@ vLLM-Omni is flexible and easy to use with:
vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
- **Omni-modality models** (e.g. Qwen3-Omni, MiniCPM-o 4.5, Cosmos3, HunyuanImage, BAGEL)
- **TTS models** (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
- **Diffusion models** — image, video, and audio generation (e.g. MiniMax H3, Qwen-Image, Wan2.2, FLUX)
- **Robot-policy and action models** (e.g. GR00T-N1.7, DreamZero-DROID, InternVLA-A1, Cosmos3 action policy)
- **TTS models** (e.g. Qwen3-TTS, IndexTTS 2.5, dots.tts, CosyVoice3)
- **Diffusion models** — image, video, and audio generation (e.g. MiniMax H3, LTX-2.5, SANA-Video, Wan2.2)
- **Robot-policy and action models** (e.g. π0, GR00T-N1.7, DreamZero-DROID, InternVLA-A1)
## Getting Started
@@ -67,6 +68,7 @@ Visit our [documentation](https://vllm-omni.readthedocs.io/en/latest/) to learn
- [Installation](https://vllm-omni.readthedocs.io/en/latest/getting_started/installation/)
- [Quickstart](https://vllm-omni.readthedocs.io/en/latest/getting_started/quickstart/)
- [Nightly CUDA Docker images](https://hub.docker.com/r/vllm/vllm-omni/tags?name=nightly) built from the latest `main`
- [List of Supported Models](https://vllm-omni.readthedocs.io/en/latest/models/supported_models/)
- [Deployment Recipes](https://recipes.vllm.ai) for vLLM-Omni model serving
+19 -9
View File
@@ -5,14 +5,22 @@ nav:
- getting_started/quickstart.md
- getting_started/installation/*
- Serving:
- OpenAI-Compatible API:
- Chat Completions API: serving/chat_completions_api.md
- Image Generation: serving/image_generation_api.md
- Image Edit: serving/image_edit_api.md
- Text to Speech: serving/speech_api.md
- Videos: serving/videos_api.md
- Audio Generation: serving/audio_generate_api.md
- Streaming Video Input: serving/video_stream_api.md
- API Server:
- Overview: serving/README.md
- Core APIs:
- Chat Completions: serving/chat_completions_api.md
- Text to Speech: serving/speech_api.md
- Audio Generation: serving/audio_generate_api.md
- Image Generation: serving/image_generation_api.md
- Image Edit: serving/image_edit_api.md
- Videos: serving/videos_api.md
- Streaming and Realtime:
- Realtime Audio: serving/realtime_api.md
- Full Duplex: serving/full_duplex_api.md
- Streaming Video Input: serving/video_stream_api.md
- Streaming Video Output: serving/streaming_video_output_api.md
- OpenPI Robot Policy: serving/openpi_api.md
- Standalone Experimental Servers: serving/standalone_servers.md
- Examples:
- examples/README.md
- Offline Inference:
@@ -29,7 +37,7 @@ nav:
- Image-To-Video: user_guide/examples/online_serving/image_to_video.md
- Speech-To-Video: user_guide/examples/online_serving/speech_to_video.md
- Text-To-Image: user_guide/examples/online_serving/text_to_image.md
- Text-To-Speech: user_guide/examples/online_serving/text_to_speech.md
- Text-To-Speech (Online Serving): user_guide/examples/online_serving/text_to_speech.md
- Text-To-Video: user_guide/examples/online_serving/text_to_video.md
- General:
- usage/*
@@ -119,6 +127,8 @@ nav:
- Async Omni Output Materialization: design/feature/omni_async_output_materialization.md
- Automatic Prefix Caching: design/feature/prefix_caching.md
- Realtime AR-Diffusion: design/feature/realtime_ar_diffusion.md
- Model-Specific Capabilities:
- MiniMax H3 Temporal VAE Chunks: design/feature/minimax_h3_temporal_vae_chunks.md
- Communication:
- OmniConnector Implementations:
- Mooncake Store Connector: design/feature/omni_connectors/mooncake_store_connector.md
+3 -3
View File
@@ -57,9 +57,9 @@ vLLM-Omni is flexible and easy to use with:
vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
- Omni-modality models (e.g. Qwen3-Omni, MiniCPM-o 4.5, Cosmos3, HunyuanImage, BAGEL)
- TTS models (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
- Diffusion models — image, video, and audio generation (e.g. MiniMax H3, Qwen-Image, Wan2.2, FLUX)
- Robot-policy and action models (e.g. GR00T-N1.7, DreamZero-DROID, InternVLA-A1, Cosmos3 action policy)
- TTS models (e.g. Qwen3-TTS, IndexTTS 2.5, dots.tts, CosyVoice3)
- Diffusion models — image, video, and audio generation (e.g. MiniMax H3, LTX-2.5, SANA-Video, Wan2.2)
- Robot-policy and action models (e.g. π0, GR00T-N1.7, DreamZero-DROID, InternVLA-A1)
For more information, checkout the following:
+1 -1
View File
@@ -2,7 +2,7 @@
This section lists the most common options for running vLLM-Omni.
For options within a vLLM Engine, please refer to the [vLLM 0.26 configuration guide](https://docs.vllm.ai/en/v0.26.0/configuration/index.html).
For options within a vLLM Engine, please refer to the [vLLM 0.28 configuration guide](https://docs.vllm.ai/en/v0.28.0/configuration/index.html).
Each model defines fixed topology in a registered `PipelineConfig` and runtime
overrides in a deploy YAML.
+1 -1
View File
@@ -151,7 +151,7 @@ their keys only.
## Inherited vLLM variables
vLLM-Omni also reads variables through its aligned vLLM dependency. Refer to
the [vLLM 0.26 environment-variable reference](https://docs.vllm.ai/en/v0.26.0/configuration/env_vars.html)
the [vLLM 0.28 environment-variable reference](https://docs.vllm.ai/en/v0.28.0/configuration/env_vars.html)
for their definitions. This includes vLLM launch, cache, logging, plugin, ROCm,
XPU, ModelScope, and FlashInfer workspace settings.
+4
View File
@@ -170,6 +170,10 @@ because the user-visible output is diffusion media rather than a token stream.
The
[MiniMax-H3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md)
contains the hardware-specific profiles and warmup requirements.
The internal, opt-in
[temporal VAE chunk callback](feature/minimax_h3_temporal_vae_chunks.md)
documents the narrower released-source contract for publishing committed video
frames during decode. It does not change the complete-output serving topology.
### Cosmos3: unified MoT reasoner and generator
@@ -0,0 +1,260 @@
# MiniMax H3 temporal VAE chunk callback
!!! warning "Experimental internal capability"
This feature is an opt-in model integration seam. The current MiniMax H3
pipeline still uses complete-output VAE decode, and no CLI option, serving
path, transport, or MP4 sink enables temporal chunks by default.
## Status and scope
MiniMax H3 decodes video latents in temporal windows, but the released VAE
normally exposes only the complete reconstructed video. The temporal callback
publishes frames after the released decoder has completed overlap blending and
padding removal. This gives a future runtime consumer a stable point at which
it can schedule representation conversion, D2H transfer, and encoding while
the VAE continues decoding later windows.
The implementation remains model-local:
- `vae.py` owns VAE topology and the thin
`decode_latent_with_chunks()` facade;
- `chunked_decode.py` owns the normalized callback contract, metadata
validation, and fixed VAE-group rendezvous points; and
- `temporal_chunks.py` owns the source-gated compatibility planner and the
released temporal assembly control flow.
Crop policy, typed media, transport, backpressure, encoding, and public output
remain runtime responsibilities. The model code does not schedule requests or
select a presentation format.
## Execution flow
The default and opt-in paths are intentionally separate:
```text
Default:
latent -> MiniMaxH3VideoVAE.decode_latent() -> complete RGB video
Opt-in producer seam:
latent
-> MiniMaxH3VideoVAE.decode_latent_with_chunks()
-> group-local preflight and callback-owner validation
-> released temporal decode / overlap blending / padding removal
-> borrowed decoder-domain chunk
-> normalize + clone into owned RGB float32 BCTHW [0, 1]
-> synchronous callback
-> complete RGB video is still returned
```
Calling `decode_latent_with_chunks()` is the opt-in. Supplying no callback on a
VAE peer does not select the default decoder; every rank in a VAE patch group
must enter the chunked method, and only the output rank receives the callback.
## Producer contract
The normalized callback has this structural signature:
```python
def callback(
frames: torch.Tensor,
/,
*,
chunk_index: int,
total_chunks: int,
frame_start: int,
is_final: bool,
) -> None:
...
```
For a successful decode:
- `frames` is owned, contiguous RGB float32 in `BCTHW` layout and `[0, 1]`;
- callback storage does not alias the returned complete video;
- the consumer may retain or mutate a chunk;
- `chunk_index` is zero-based, contiguous, and strictly increasing;
- `total_chunks` is positive and stable for the decode;
- `frame_start` equals the end of the previous published range;
- chunks are nonempty and contain only temporally committed frames;
- exactly one chunk has `is_final=True`; and
- concatenating chunks along time exactly reconstructs the returned video.
The callback runs synchronously on the producer's current device stream. A
consumer that reads the tensor on another stream must establish ordering, for
example with an event, and retain the allocation for that stream's lifetime.
A slow callback also blocks VAE production. An overlap-oriented runtime
consumer should therefore acquire bounded capacity, enqueue side-stream work,
and return promptly instead of encoding frames inside the callback.
Chunks retain the decoded VAE canvas dimensions. A future runtime integration
must apply the same requested-height and requested-width crop as the complete
MiniMax H3 pipeline before declaring the media transport-ready.
The raw compatibility callback is deliberately narrower. It receives borrowed
decoder-domain tensors and must not retain or mutate them. The coordinator is
the ownership boundary that normalizes and clones raw chunks before invoking
the public callback.
## Temporal planning and terminal publication
The released configuration advances five latent tokens per temporal window and
includes a two-token overlap, so each `_adaptive_decode` input normally contains
seven tokens. After temporal scaling, pre-padding removal, and five-frame
overlap blending, the normal committed unit is 17 output frames. The final
remainder depends on the requested frame plan.
Representative released frame plans include:
| Output frames | Published chunks |
| ---: | --- |
| 124 | `7 x 17 + 5` |
| 209 | `12 x 17 + 5` |
| 243 | `14 x 17 + 5` |
| 362 | `21 x 17 + 5` |
The terminal callback is held until the complete frame plan passes its logical
frame, padding, written-frame, chunk-count, and finality checks. Earlier chunks
are semantically stable, but they do not imply that the overall request has
succeeded; a consumer must discard or abort partial output if a later decoder
failure terminates the request.
## Source and configuration gate
The released VAE has no public temporal-callback API. vLLM-Omni therefore owns
a compatibility implementation for one exact released code contract. The
reference revision identifies its provenance, but runtime acceptance is based
on source content, configuration, and structure rather than repository or
weight metadata:
- reference Hugging Face revision, not independently checked at runtime:
`42ed227ee7df40d41602854ae760620d6eb651fe`;
- sorted 15-file top-level `video_vae/*.py` manifest (excluding
`__init__.py`) SHA-256:
`30e64afbaa940696bf8bfe2c86003a4b1666c84b22fa3aa12403a3e6a03f705d`;
- clip length `17`;
- token chunk, overlap, and drop `5`, `2`, and `3`;
- temporal ratio `4`;
- frame pre-padding and overlap `3` and `5`; and
- no isolated first or last frame.
The adapter does not verify Hugging Face revision metadata, repository identity,
or model weights. The manifest pin detects Python compatibility drift; it is
not a security boundary for remote-code execution. The source manifest, listed
temporal attributes, structural contract, plan, and dtype are validated before
the first `_adaptive_decode` call. A mismatch disables only the opt-in chunk path.
`decode_latent()` remains available and performs no source fingerprint or chunk
rendezvous.
## Supported models
Support in this table means the internal producer contract above, not public
serving or end-to-end transport support.
| Model or component | Source contract | Status | Validation |
| --- | --- | --- | --- |
| MiniMax H3 FL2VA video VAE | Released manifest/listed temporal attributes above (reference `42ed227e`) | Supported, opt-in | Same-environment full-output parity and L20X/B300 feasibility measured on the implementation through PR commit `60e9bbad` |
| MiniMax H3 Ref2VA video VAE | Same released manifest/listed temporal attributes | Source-compatible; not separately benchmarked | Ref2VA and FL2VA use byte-identical video-VAE source and configuration in the reference snapshot |
| MiniMax H3 with a different Python manifest or listed temporal attribute | Different enforced contract | Unsupported for chunk mode | Fails closed before decoder collectives; complete decode remains available |
| Same manifest/temporal attributes with different weights or other component metadata | Not checked by this gate | Not independently qualified | The adapter does not inspect weights, repository identity, revision metadata, or arbitrary component configuration |
| Other video models | Any | Unsupported | No generic temporal VAE producer ABI is declared by this feature |
## Feature compatibility
| Dimension or combination | Status | Constraint |
| --- | --- | --- |
| Existing complete-output decode | Supported and default | No source check, callback object, or callback rendezvous |
| Explicit chunk decode | Supported internally | Caller invokes `decode_latent_with_chunks()` directly |
| VAE patch parallel size 1 | Supported by the model seam | Caller supplies the callback only on the request output owner; no WORLD rendezvous is added |
| VAE patch parallel size equal to the full DiT group | Supported when data parallel size is 1 | Every VAE-group rank calls; exactly group rank 0 owns the callback |
| Partial VAE patch subgroup | Unsupported | MiniMax H3 currently accepts size 1 or the full DiT group only |
| Data parallel size greater than 1 with VAE patch parallel size 1 | Model seam only | Per-replica callback ownership is possible, but no serving/runtime consumer is wired |
| Data parallel size greater than 1 with VAE patch parallel size greater than 1 | Unsupported | Fails before decoder collectives because no replica-local VAE patch group is defined |
| Too few spatial tiles for the VAE group | Correctness fallback | Decodes rank-locally and restores native tiling state; slower than patch parallel decode |
| Device memory and full-output assembly | Not reduced or bounded | The compatibility path still assembles and returns the full decode while allocating owned callback chunks; no peak-HBM reduction is claimed |
| DiT cache, attention, TP, SP, HSDP, LoRA, or quantization features | No direct callback dependency | This seam runs after denoising and does not independently qualify combinations beyond the MiniMax H3 model's own matrix |
| Layerwise or distributed component offload | Not separately validated | Component residency must place the video VAE before decode |
| Step execution or request batching | Not integrated | No pipeline caller binds a runtime chunk consumer |
| Requested-size crop and device uint8 preparation | Not implemented here | Runtime-owned follow-up |
| Pinned D2H, SHM/ZMQ, or cross-process transport | Not implemented here | Runtime-owned follow-up |
| Complete MP4 or fragmented MP4 sink | Not implemented here | Existing serving behavior still waits for complete output |
| Audio VAE chunks and incremental A/V composition | Not supported | Video chunks do not imply audio or request finality |
| Cancellation and multi-output lease ownership | Not implemented here | Requires a bounded runtime lifecycle |
| NVIDIA CUDA | Feasibility evidence | L20X and B300 measurements apply to the implementation through `60e9bbad`; the current coordination refactor was not reprofiled |
| AMD ROCm | Not validated for this feature | Base MiniMax H3 support is documented separately and does not imply temporal-chunk validation |
| Ascend, Intel, or other accelerators | No feature support claim | Consult the model support table independently; this design adds no platform qualification |
## Distributed and failure semantics
When VAE patch parallelism is active, all VAE-group ranks run the same
preflight before entering checkpoint collectives. Source/planner errors and the
resolved temporal concat dtype are reconciled at fixed group rendezvous points.
Callback ownership is also a group decision. Rank `r` contributes `r + 1` when
it owns a callback and zero otherwise. A group SUM equals one only when rank 0
is the sole owner; every invalid owner configuration therefore fails
identically on all ranks before decode.
Callback exceptions are retained on the output rank. Further publication stops,
but every rank completes the same remaining native decoder collective schedule.
The original exception is then raised on the owner and
`MiniMaxH3ChunkCallbackPeerError` on peers.
Exceptions raised inside the checkpoint's native distributed decode cannot in
general be reconciled by the adapter: a peer may already be blocked inside a
checkpoint-owned collective. Resolving such failures requires native collective
handling or supervisor-level timeout and worker-group teardown. The process
group must not be assumed reusable after such a failure.
## Validation evidence
The contract suite covers ordinary and padded frame plans, blend boundaries,
owned storage and mutation isolation, source-manifest/listed-temporal-attribute
drift, callback failure, and healthy subsequent requests. Real Gloo processes cover
rank-asymmetric preflight failures, dtype disagreement, two-rank callback
failure propagation, three-rank wrong-owner rejection, insufficient-tile
fallback, and group recovery for failures that occur at adapter-owned
rendezvous points.
Four-device L20X and independent four- and eight-device B300 experiments on the
implementation through PR commit `60e9bbad` showed exact serial-versus-chunked
frame/MP4 parity within each environment and lower VAE-to-complete-MP4 latency
when an experimental consumer was attached. Hashes are not claimed to match
across different hardware/software environments. The later coordination and
module-boundary refactors were revalidated with CPU and real Gloo tests but were
not reprofiled on GPUs. These are feasibility measurements, not a promise for
default serving, because this change does not contain the transport or encoder
consumer. See [RFC #6872](https://github.com/vllm-project/vllm-omni/issues/6872)
and [PR #6885](https://github.com/vllm-project/vllm-omni/pull/6885) for the
measurement provenance and limitations.
## Generalization boundary and follow-ups
This feature does not define a repository-wide `SupportsVideoChunkDecode`
interface. H3 temporal blending, Wan causal caches, LTX temporal/spatial halos,
and pipeline-level generation chunks have different commit and failure
boundaries. A generic producer contract should be extracted only after a second
model proves the same semantic event without `Any` arguments or model-name
dispatch.
The reusable future boundary is the emitted, committed video-media chunk, not a
universal latent-to-VAE method. It should extend the typed media work in
[RFC #6541](https://github.com/vllm-project/vllm-omni/issues/6541) and
[PR #6615](https://github.com/vllm-project/vllm-omni/pull/6615), while keeping
model-specific halo, cache, padding, blending, distributed, and source-gating
logic inside each model adapter.
Runtime follow-ups own:
- typed video-plus-audio media and request/output identity;
- requested-size crop and device preparation;
- bounded leases, backpressure, cancellation, and cleanup;
- D2H/IPC transport and observability; and
- complete-MP4 and fragmented-MP4 consumers.
## Related documentation
- [Diffusion model integration](../module/diffusion/diffusion_model_integration.md)
- [VAE patch parallelism](vae_parallel.md)
- [Async diffusion output](async_diffusion_output.md)
- [MiniMax H3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md)
+4
View File
@@ -25,6 +25,10 @@ implementation contract; it is not, by itself, a general support claim.
- [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md)
- [Realtime AR-Diffusion Sessions](feature/realtime_ar_diffusion.md)
### Model-specific capabilities
- [MiniMax H3 Temporal VAE Chunks](feature/minimax_h3_temporal_vae_chunks.md)
### Communication
#### OmniConnector implementations
@@ -57,3 +57,10 @@ cross-stage routing.
Test registry selection, checkpoint loading, minimal inference, input and output
contracts, and every declared optional capability.
## Model-specific capability documents
- [MiniMax H3 temporal VAE chunk callback](../../feature/minimax_h3_temporal_vae_chunks.md)
defines the source-gated producer seam and its boundary with future runtime
media transport. It is an internal opt-in capability, not a public serving
feature.
+3 -1
View File
@@ -3,4 +3,6 @@
vLLM-Omni's examples are split into two categories:
- If you are using vLLM-Omni from within Python code, see the *Offline Inference* section.
- If you are using vLLM-Omni from an HTTP application or client, see the *Online Serving* section.
- If you are using vLLM-Omni from an HTTP application or client, start with
the [API Server guide](../serving/README.md) to choose an endpoint, then use
the *Online Serving* section for model-specific examples.
+1 -1
View File
@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
### Pre-built wheels
Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, and 0.26.0; the 0.28.0 wheels are published together with the v0.28.0 release this branch produces. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, 0.26.0, and 0.28.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
=== "NVIDIA CUDA"
@@ -5,7 +5,7 @@
# --8<-- [end:requirements]
# --8<-- [start:set-up-using-python]
vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni 0.26.x release line uses vLLM 0.26.x.
vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni 0.28.x release line uses vLLM 0.28.x.
!!! note
PyTorch installed via `conda` will statically link `NCCL` library, which can cause issues when vLLM tries to use `NCCL`. See <gh-issue:8420> for more details.
@@ -95,6 +95,13 @@ uv pip install --no-build-isolation --editable .
vLLM-Omni offers an official docker image for deployment. These images are built on top of vLLM docker images and available on Docker Hub as [vllm/vllm-omni](https://hub.docker.com/r/vllm/vllm-omni/tags). The version of vLLM-Omni indicates which release of vLLM it is based on.
Nightly CUDA images built from the latest `main` are also available:
- [`vllm/vllm-omni:nightly`](https://hub.docker.com/r/vllm/vllm-omni/tags?name=nightly) is the rolling multi-architecture tag.
- [`vllm/vllm-omni:nightly-<commit>`](https://hub.docker.com/r/vllm/vllm-omni/tags?name=nightly-) provides commit-pinned multi-architecture tags. The newest 14 are retained.
Use a commit-pinned tag when you need a reproducible nightly environment.
Here's an example deployment command that has been verified on 2 x H100's:
```bash
docker run --runtime nvidia --gpus 2 \
@@ -9,8 +9,8 @@ The recommended way to use vLLM-Omni on NPU is through the vLLM-Ascend pre-built
```bash
# vLLM-Ascend image
# Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.26.0
# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.26.0-a3
# Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.28.0
# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.28.0-a3
docker run --rm \
--name vllm-omni-npu \
--shm-size=1g \
@@ -28,11 +28,11 @@ docker run --rm \
-v /etc/ascend_install.info:/etc/ascend_install.info \
-v /root/.cache:/root/.cache \
-p 8000:8000 \
-it quay.io/atlas-ci/vllm-ascend:v0.26.0-a3 bash
-it quay.io/atlas-ci/vllm-ascend:v0.28.0-a3 bash
# Inside the container, install vLLM-Omni from source
cd /vllm-workspace
git clone -b v0.26.0 https://github.com/vllm-project/vllm-omni.git
git clone -b v0.28.0 https://github.com/vllm-project/vllm-omni.git
cd vllm-omni
pip install -v -e . --no-build-isolation
# or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
@@ -50,21 +50,11 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886
# --8<-- [start:installation-main]
You can also build vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. (But sometimes it will break for a while. You can check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest commit of vLLM-Omni main branch on NPU.)
You can also install vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. Use the aligned vLLM-Ascend v0.28.0 image described above, and check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest vLLM-Omni commit on NPU.
```bash
# Pin vLLM and vLLM-Ascend to the v0.26 release line
git clone -b v0.26.0 https://github.com/vllm-project/vllm.git
cd vllm
VLLM_TARGET_DEVICE=empty pip install -v -e .
cd ..
git clone -b releases/v0.26.0rc https://github.com/vllm-project/vllm-ascend.git
cd vllm-ascend
pip install -v -e .
cd ..
# Install vLLM-Omni from the latest main branch
# Inside the vLLM-Ascend v0.28.0 container, install vLLM-Omni from main
cd /vllm-workspace
git clone https://github.com/vllm-project/vllm-omni.git
cd vllm-omni
pip install -v -e . --no-build-isolation
@@ -88,8 +78,8 @@ Supported images as following.
Here's an example deployment command that has been verified on 4 x NPUs:
```bash
# Atlas A2: quay.io/ascend/vllm-omni:v0.26.0
# Atlas A3: quay.io/ascend/vllm-omni:v0.26.0-a3
# Atlas A2: quay.io/ascend/vllm-omni:v0.28.0
# Atlas A3: quay.io/ascend/vllm-omni:v0.28.0-a3
docker run --rm \
--name vllm-omni-a3 \
--shm-size=4g \
@@ -114,7 +104,7 @@ docker run --rm \
-v /etc/ascend_install.info:/etc/ascend_install.info \
-v ~/.cache:/root/.cache \
-p 8091:8091 \
-it quay.io/ascend/vllm-omni:v0.26.0-a3 bash
-it quay.io/ascend/vllm-omni:v0.28.0-a3 bash
```
!!! tip
+10 -14
View File
@@ -91,7 +91,7 @@ if __name__ == "__main__":
For more usages, please refer to [offline inference](../user_guide/examples/offline_inference/qwen2_5_omni.md)
## Online Serving with OpenAI-Completions API
## Online Serving with the API Server
Text-to-image generation quickstart with vLLM-Omni:
@@ -100,20 +100,16 @@ vllm serve Tongyi-MAI/Z-Image-Turbo --omni --port 8091
```
```bash
curl -s http://localhost:8091/v1/chat/completions \
curl -s http://localhost:8091/v1/images/generations \
-H "Content-Type: application/json" \
-d '{
"messages": [
{"role": "user", "content": "a cup of coffee on the table"}
],
"extra_body": {
"height": 1024,
"width": 1024,
"num_inference_steps": 50,
"guidance_scale": 4.0,
"seed": 42
}
}' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > coffee.png
"prompt": "a cup of coffee on the table",
"size": "1024x1024",
"response_format": "b64_json",
"seed": 42
}' | jq -r '.data[0].b64_json' | base64 -d > coffee.png
```
For more details, please refer to [online serving](../user_guide/examples/online_serving/text_to_image.md).
See the [API Server guide](../serving/README.md) to choose an endpoint. For
model-specific details, refer to the [text-to-image online serving
example](../user_guide/examples/online_serving/text_to_image.md).
+3 -3
View File
@@ -21,7 +21,7 @@ th {
|--------------|--------|-------------------|------------|---------|-----|-----------|--------|
| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen3-Omni.md) |
| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ | | | | — |
| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ | | | | — |
| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ | — |
| `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | | — |
| `Gr00tN1d7Pipeline` | GR00T N1.7 | `nvidia/GR00T-N1.7-3B` | ✅︎ | | | | — |
@@ -62,7 +62,7 @@ th {
| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `LongCatVideoAvatarPipeline` | LongCat-Video-Avatar-1.5 A2V/AI2V (native single-speaker and multi-speaker AI2V/AVC) | `meituan-longcat/LongCat-Video-Avatar-1.5` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/meituan-longcat/LongCat-Video-Avatar-1.5.md) |
| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit` | ✅︎ | | | ✅︎ | — |
| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit`, `Boogu/Boogu-Image-0.1-Edit-Turbo` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Boogu/Boogu-Image.md) |
| `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ | — |
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — |
@@ -95,7 +95,7 @@ th {
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — |
| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | — |
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — |
| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | | — |
| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only), SenseNova-U1.5 | `SenseNova/SenseNova-U1-8B-MoT`, `sensenova/SenseNova-U1.5-8B-MoT` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/SenseNova/SenseNova-U1.5.md) |
| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — |
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — |
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | | — |
+187
View File
@@ -0,0 +1,187 @@
# API Server
vLLM-Omni exposes OpenAI-compatible HTTP APIs plus vLLM-Omni extensions for
image, audio, video, realtime, and robot-policy workloads. Use this page to
choose an endpoint; follow the linked reference page for model-specific fields,
wire protocols, and response details.
Each server instance hosts one model. An endpoint is usable only when the
loaded model supports that task.
This guide focuses on vLLM-Omni's primary task endpoints. Compatible models
may also expose standard routes inherited from vLLM, such as
`POST /v1/completions`; those routes remain model- and task-dependent.
## Start and Verify the Server
Start a model with the shared serving command:
```bash
vllm serve <model> --omni --port 8091
```
Some models require additional flags or a deployment configuration. Use the
model-specific guide when one is provided.
After the server starts, check its health and served model name:
```bash
export VLLM_OMNI_BASE_URL=http://localhost:8091
curl "$VLLM_OMNI_BASE_URL/health"
curl "$VLLM_OMNI_BASE_URL/v1/models" | jq .
```
Use `http://localhost:8091/v1` as the `base_url` for an OpenAI SDK client. If
the server was started with `--api-key`, include an `Authorization` header with
`Bearer <api-key>` in requests.
## Choose a Core HTTP Endpoint
Prefer a task-specific endpoint when one matches your workload. Use Chat
Completions for conversational or heterogeneous omni pipelines, rather than as
the default wrapper for every generation task.
| Task | Endpoint | Request | Response | Details |
|------|----------|---------|----------|---------|
| Conversation or multimodal understanding/generation | `POST /v1/chat/completions` | JSON | OpenAI-style JSON or SSE | [Chat Completions](chat_completions_api.md) |
| Text-to-speech | `POST /v1/audio/speech` | JSON | Audio bytes or SSE | [Speech](speech_api.md) |
| Sound, music, or ambient audio generation | `POST /v1/audio/generate` | JSON | Audio bytes | [Audio Generation](audio_generate_api.md) |
| Text-to-image | `POST /v1/images/generations` | JSON | Base64 JSON or an image file | [Image Generation](image_generation_api.md) |
| Image editing | `POST /v1/images/edits` | Multipart form | Base64 JSON or SSE | [Image Edit](image_edit_api.md) |
| Video generation (recommended) | `POST /v1/videos` | Multipart form | Asynchronous job JSON | [Videos](videos_api.md) |
| Video generation (blocking) | `POST /v1/videos/sync` | Multipart form | Video bytes | [Videos](videos_api.md#synchronous-response) |
`POST /v1/videos/sync` is intended for tests and simple scripts. Use the
asynchronous job API for production requests so clients can poll, download,
and delete generated videos independently.
## Minimal Requests
The examples below assume a compatible model is already running. TTS voices,
media inputs, generation controls, and output capabilities vary by model.
=== "Chat"
```bash
curl "$VLLM_OMNI_BASE_URL/v1/chat/completions" \
-H "Content-Type: application/json" \
-d '{
"messages": [{"role": "user", "content": "Describe vLLM-Omni briefly."}],
"modalities": ["text"]
}'
```
=== "Speech"
```bash
curl "$VLLM_OMNI_BASE_URL/v1/audio/speech" \
-H "Content-Type: application/json" \
-d '{
"input": "Hello from vLLM-Omni.",
"voice": "vivian"
}' \
--output speech.wav
```
=== "Audio"
```bash
curl "$VLLM_OMNI_BASE_URL/v1/audio/generate" \
-H "Content-Type: application/json" \
-d '{
"input": "Ocean waves on a quiet beach",
"audio_length": 5.0
}' \
--output audio.wav
```
=== "Image"
```bash
curl "$VLLM_OMNI_BASE_URL/v1/images/generations" \
-H "Content-Type: application/json" \
-d '{
"prompt": "A small robot reading beside a window",
"size": "1024x1024",
"response_format": "file"
}' \
--output image.png
```
=== "Image edit"
```bash
curl "$VLLM_OMNI_BASE_URL/v1/images/edits" \
-F "image=@input.png" \
-F "prompt=Turn the background into a snowy mountain" \
| jq -r '.data[0].b64_json' \
| base64 --decode > edited.png
```
=== "Video"
```bash
curl "$VLLM_OMNI_BASE_URL/v1/videos" \
-F "prompt=A cinematic tracking shot of a mountain lake at sunrise" \
-F "size=1280x720" \
-F "seconds=5"
```
The response contains a video job `id`. Poll and download it with:
```bash
curl "$VLLM_OMNI_BASE_URL/v1/videos/<video_id>"
curl -L "$VLLM_OMNI_BASE_URL/v1/videos/<video_id>/content" --output video.mp4
```
## Choose a Streaming or Realtime Endpoint
These WebSocket APIs have different event schemas and cannot be used
interchangeably.
| Workload | Endpoint | Interaction model | Details |
|----------|----------|-------------------|---------|
| Incremental text input for speech synthesis | `WS /v1/audio/speech/stream` | Send text events and receive audio | [Streaming Text to Speech](speech_api.md#streaming-text-input-websocket) |
| Live video understanding | `WS /v1/video/chat/stream` | Send video frames and receive text/audio | [Streaming Video Input](video_stream_api.md) |
| Turn-based realtime audio | `WS /v1/realtime` | Stream one audio input and receive transcript/audio events | [Realtime Audio](realtime_api.md) |
| Continuous speech-to-speech interaction | `WS /v1/realtime?duplex=1` or `WS /v1/duplex` | Listen and speak concurrently with session control | [Full Duplex](full_duplex_api.md) |
| Generated video chunks | `WS /v1/realtime/video` | Start a diffusion request and receive fragmented MP4 | [Streaming Video Output](streaming_video_output_api.md) |
| Robot policy inference | `WS /v1/realtime/robot/openpi` | Send MessagePack observations and receive action arrays | [OpenPI Robot Policy](openpi_api.md) |
All six routes are model- or configuration-dependent. In particular,
`/v1/realtime` is not full duplex unless the client sets `duplex=1` and the
deployment explicitly enables duplex sessions. Clients should also verify the
duplex capability payload because the query-parameter form falls back to the
ordinary realtime handler when duplex is unavailable.
## Related Endpoints
| Purpose | Endpoints | Reference |
|---------|-----------|-----------|
| Discovery and readiness | `GET /health`, `GET /v1/models` | This page |
| Batched conversations | `POST /v1/chat/completions/batch` | [Batch requests](chat_completions_api.md#batch-requests) |
| Batched speech | `POST /v1/audio/speech/batch` | [Batch speech generation](speech_api.md#batch-speech-generation) |
| TTS voice management | `GET/POST /v1/audio/voices`, `DELETE /v1/audio/voices/{name}` | [Voices](speech_api.md#voices-endpoint) |
| Video job lifecycle | `GET /v1/videos`, `GET/DELETE /v1/videos/{video_id}`, `GET /v1/videos/{video_id}/content` | [Video endpoints](videos_api.md#endpoints) |
| Release and restore stage memory | `POST /v1/omni/sleep`, `POST /v1/omni/wakeup` | [Sleep Mode](../features/sleep_mode.md) |
## Standalone Experimental Servers
PersonaPlex also provides a Moshi-compatible WebSocket server, while JoyVL
provides a stateful interaction orchestrator in front of another model server.
These are separate processes with their own routes; they are not additional
paths on every `vllm serve --omni` instance. See [Standalone Experimental
Servers](standalone_servers.md) before deploying either one.
## Compatibility Notes
- A registered path does not mean every model supports it. Unsupported tasks
return an OpenAI-style error or a WebSocket `error` event.
- Use JSON for Chat, Speech, Audio Generation, and Image Generation. Use
`multipart/form-data` when uploading images, videos, or other files.
- Generation can take minutes for large diffusion models. Configure client
timeouts accordingly; prefer asynchronous video jobs over long blocking
requests.
- Model-specific parameters are vLLM-Omni extensions. Check the endpoint
reference before assuming an OpenAI SDK exposes them directly.
+3 -3
View File
@@ -4,14 +4,14 @@ vLLM-Omni provides an API for text-to-audio generation using diffusion-based mod
Unlike the [Speech API](speech_api.md) which targets text-to-speech synthesis, the Audio Generate API is designed for general-purpose audio generation from text descriptions (sound effects, music, ambient soundscapes, etc.).
Each server instance runs a single model (specified at startup via `vllm-omni serve <model> --omni`).
Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
## Quick Start
### Start the Server
```bash
vllm-omni serve stabilityai/stable-audio-open-1.0 \
vllm serve stabilityai/stable-audio-open-1.0 \
--host 0.0.0.0 \
--port 8091 \
--gpu-memory-utilization 0.9 \
@@ -326,7 +326,7 @@ curl http://localhost:8091/health
Enable debug logging:
```bash
vllm-omni serve stabilityai/stable-audio-open-1.0 \
vllm serve stabilityai/stable-audio-open-1.0 \
--host 0.0.0.0 \
--port 8091 \
--gpu-memory-utilization 0.9 \
+86 -54
View File
@@ -1,78 +1,110 @@
# Chat Completions API
vLLM-Omni supports generating and editing images via the `/v1/chat/completions`
endpoint using diffusion models. This page explains how to pass generation
parameters (such as `num_inference_steps`, `height`, `width`) to diffusion
models through this endpoint.
Use `/v1/chat/completions` for conversational and multimodal pipelines that
follow the OpenAI Chat API schema. Depending on the loaded model, a request can
contain text, image, audio, or video input and can return text, audio, images,
or other model-specific output.
!!! tip
For dedicated endpoints that accept generation parameters as top-level
fields, see [Image Generation API](image_generation_api.md) and
[Image Edit API](image_edit_api.md).
Some diffusion pipelines also support image generation and editing through
Chat Completions. Prefer the dedicated [Image Generation
API](image_generation_api.md) or [Image Edit API](image_edit_api.md) for
task-specific clients; their request fields and response contracts are more
direct.
## Passing Generation Parameters
The `/v1/chat/completions` endpoint follows the OpenAI Chat API schema, which
does not natively include diffusion-specific fields like `num_inference_steps`
or `height`. How you pass these extra fields depends on your client.
### curl / Python `requests`
Wrap generation parameters inside an `"extra_body"` key in the JSON body:
## Basic Request
```bash
curl -s http://localhost:8091/v1/chat/completions \
curl http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"messages": [
{"role": "user", "content": "A beautiful landscape painting"}
{"role": "user", "content": "Describe vLLM-Omni briefly."}
],
"extra_body": {
"num_inference_steps": 50,
"seed": 42
}
"modalities": ["text"]
}'
```
### OpenAI Python SDK
Use `stream: true` for Server-Sent Events. Input media syntax, supported output
`modalities`, and response choices depend on the model; see the
[model-specific online serving examples](../user_guide/examples/online_serving/qwen3_omni.md).
Use the `extra_body` **keyword argument**. The SDK automatically merges these
fields into the top-level request body:
!!! tip
Each server hosts one model. Query `GET /v1/models` for its served model
name when your client requires the `model` field.
```python
response = client.chat.completions.create(
model="Qwen/Qwen-Image",
messages=[{"role": "user", "content": "A beautiful landscape painting"}],
extra_body={
"num_inference_steps": 50,
"seed": 42,
},
)
## vLLM Extension Parameters
Send standard and vLLM-specific Chat parameters as top-level fields in a
direct HTTP request. When using the OpenAI Python SDK, pass vLLM-specific
fields through the SDK's `extra_body` keyword; the SDK merges them into the
top-level JSON sent to the server.
=== "curl"
```bash
curl http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"messages": [{"role": "user", "content": "Write a short haiku."}],
"top_k": 40
}'
```
=== "OpenAI Python SDK"
```python
from openai import OpenAI
client = OpenAI(base_url="http://localhost:8091/v1", api_key="none")
response = client.chat.completions.create(
model="your-served-model",
messages=[{"role": "user", "content": "Write a short haiku."}],
extra_body={"top_k": 40},
)
```
??? note "Diffusion compatibility"
Diffusion models exposed through Chat Completions can accept supported
fields such as `num_inference_steps`, `seed`, `height`, and `width`. Send
them at the top level with direct HTTP, or in the SDK's `extra_body`
keyword argument.
A literal nested `"extra_body"` JSON object is accepted for compatibility,
but it is not recommended for new direct HTTP clients. Do not provide the
same parameter in multiple locations; the server returns a `400` error for
duplicate diffusion parameters. Prefer the dedicated [Image Generation
API](image_generation_api.md) or [Image Edit API](image_edit_api.md) when
either endpoint matches the task.
## Batch Requests
`POST /v1/chat/completions/batch` accepts the same shared generation fields,
but `messages` is a list of conversations. The response contains one choice
per conversation in input order.
```bash
curl http://localhost:8091/v1/chat/completions/batch \
-H "Content-Type: application/json" \
-d '{
"messages": [
[{"role": "user", "content": "Summarize vLLM in one sentence."}],
[{"role": "user", "content": "Summarize vLLM-Omni in one sentence."}]
],
"max_tokens": 64
}'
```
!!! note "SDK `extra_body` vs. JSON `extra_body`"
These two `extra_body` usages look similar but work differently under the
hood. The SDK flattens the dict into the top-level request JSON, while the
curl/requests approach sends it as a nested `"extra_body"` key. Both are
handled correctly by the server.
!!! note "About the `ignored fields` warning"
You may see a log message like:
```
WARNING: The following fields were present in the request but ignored: {'height', 'width', ...}
```
This is **harmless**. It is emitted by vLLM's request validation layer
because these fields are not part of the standard OpenAI
`ChatCompletionRequest` schema. The fields are still stored internally
and correctly forwarded to the diffusion pipeline.
Batch Chat Completions does not support streaming, tools, beam search, or
`n > 1`.
## Model-Specific Examples
For complete examples with full request/response details, see the model-specific
guides:
For complete examples with model-specific inputs and outputs, see:
- [Qwen3-Omni](../user_guide/examples/online_serving/qwen3_omni.md)
- [Qwen2.5-Omni](../user_guide/examples/online_serving/qwen2_5_omni.md)
- [Text-to-Image (Qwen-Image)](../user_guide/examples/online_serving/text_to_image.md)
- [Image-to-Image (Qwen-Image-Edit, Qwen-Image-Layered)](../user_guide/examples/online_serving/image_to_image.md)
- [GLM-Image](../user_guide/examples/online_serving/glm_image.md)
+116
View File
@@ -0,0 +1,116 @@
# Full-Duplex WebSocket API
vLLM-Omni provides an experimental full-duplex runtime for models that can
continue receiving speech while producing speech. It adds persistent session
state, model-specific turn policy, overlap handling, playback acknowledgement,
and optional session resume.
Full duplex is distinct from the turn-based [Realtime Audio API](realtime_api.md).
## Choose an Endpoint
| Endpoint | Protocol | Recommended use |
|----------|----------|-----------------|
| `WS /v1/realtime?duplex=1` | OpenAI Realtime-style event projection | Applications and browser clients |
| `WS /v1/duplex` | Native vLLM-Omni duplex events | Runtime integration and low-level testing |
Both endpoints use the same duplex engine and require the same model-side
adapter. Prefer the Realtime projection unless the native lifecycle events are
specifically required.
## Enable Full Duplex
The route becomes usable only when the deployment configuration explicitly
sets:
```yaml
session_mode: duplex
```
The selected model pipeline must also provide a duplex serving adapter.
Model-native deployments configure the corresponding engine runtime extension
and control plane as part of their registered pipeline.
!!! warning
`WS /v1/duplex` fails with `Duplex API is not available` when duplex is not
enabled. By contrast, `/v1/realtime?duplex=1` falls back to the ordinary
turn-based realtime handler when no duplex handler exists. Confirm that
`session.created.session.capabilities` is present before treating the
connection as full duplex.
The current unified-runtime integrations are:
- MiniCPM-o 4.5, using `vllm_omni/deploy/minicpmo_4_5_duplex.yaml`;
- PersonaPlex, whose default `vllm_omni/deploy/personaplex.yaml` enables duplex.
JoyVL is a separate HTTP interaction orchestrator and does not use these
WebSocket endpoints. See [Standalone Experimental Servers](standalone_servers.md).
## MiniCPM-o Quick Start
Start the duplex deployment:
```bash
vllm serve openbmb/MiniCPM-o-4_5 --omni \
--deploy-config vllm_omni/deploy/minicpmo_4_5_duplex.yaml \
--trust-remote-code \
--port 8091
```
Stream a mono, PCM16, 16 kHz WAV file with the provided client:
```bash
python examples/online_serving/minicpmo/realtime_duplex_demo.py \
--url 'ws://localhost:8091/v1/realtime?duplex=1' \
--model openbmb/MiniCPM-o-4_5 \
--input-wav input_16k_mono.wav \
--ref-audio reference_voice.wav \
--output-dir /tmp/minicpmo-duplex
```
## Realtime Event Lifecycle
A typical `/v1/realtime?duplex=1` session follows this lifecycle:
1. Send `session.update` with the model, modalities, audio formats, and session
options.
2. Wait for `session.created`; inspect `session.capabilities` instead of
assuming every model supports the same controls.
3. Send `input_audio_buffer.append` events while microphone audio arrives.
4. Send `input_audio_buffer.commit` at a user-turn boundary when required by
the model policy.
5. Consume `response.created`, transcript deltas, `response.audio.delta`, and
`response.done` or `response.listen` events.
6. Send `playback.ack` after audio has been played when the session advertises
playback acknowledgement support.
7. Send `session.close` and wait for `session.closed`.
Unlike the turn-based realtime endpoint, input may continue while a response
is active. The server can emit `overlap.decision` to describe whether input was
deferred, treated as a short acknowledgement, or used to interrupt output.
## Capabilities and Model Differences
The `session.created` payload includes capability fields such as
`supports_barge_in`, `supports_playback_ack`, `supports_multi_session`,
`supports_session_resume`, and `chunk_period_ms`. Treat this payload as the
runtime contract.
For example, PersonaPlex supports native overlapping speech but currently
advertises `supports_barge_in=false`; destructive output interruption and
model-state rewind have not been validated for that integration. Capacity and
session-resume behavior also depend on the selected deployment configuration.
## Native Protocol
`WS /v1/duplex` exposes lower-level lifecycle names including
`session.create`, `input_audio_buffer.append`, `turn.signal`, `playback.ack`,
and `session.close`. It returns native session, input, response, overlap, and
error events. This protocol is experimental and may evolve with the runtime;
applications should use the provided Realtime client where possible.
See the [MiniCPM-o example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/minicpmo),
[PersonaPlex example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/personaplex),
and [full-duplex runtime design](https://github.com/vllm-project/vllm-omni/blob/main/vllm_omni/experimental/fullduplex/DESIGN.md)
for model-specific validation and architecture details.
+109
View File
@@ -0,0 +1,109 @@
# OpenPI Robot Policy WebSocket API
Use `WS /v1/realtime/robot/openpi` for low-latency robot-policy inference with
OpenPI-compatible clients. A persistent WebSocket carries observations to the
loaded policy model and returns action arrays.
Despite the `/realtime` prefix, this endpoint does not use the OpenAI Realtime
event schema. It uses binary MessagePack frames with NumPy extensions.
## Availability
The endpoint is available only when the loaded diffusion policy configuration
contains `policy_server_config`. The server sends that model-specific
configuration to the client as the first binary frame.
For example, DreamZero declares:
```yaml
policy_server_config:
image_resolution: [180, 320]
n_external_cameras: 2
needs_wrist_camera: true
needs_stereo_camera: false
needs_session_id: true
action_space: joint_position
```
Without this configuration, the WebSocket returns `Robot policy not available`
and closes.
## DreamZero Quick Start
Start the bundled DreamZero deployment:
```bash
vllm serve GEAR-Dreams/DreamZero-DROID --omni --port 8091 \
--served-model-name dreamzero-droid \
--deploy-config vllm_omni/deploy/dreamzero_tp1_cfg2.yaml \
--enforce-eager --disable-log-stats
```
Install the optional client dependencies, download the sample camera inputs,
and run the client:
```bash
pip install openpi-client websockets opencv-python
hf download YangshenDeng/vllm-omni-dreamzero-assets \
--repo-type dataset \
--local-dir outputs/dreamzero/assets
python examples/online_serving/dreamzero/openpi_client.py \
--host 127.0.0.1 \
--port 8091 \
--video-dir outputs/dreamzero/assets
```
## Protocol
The connection is request-response after an initial server handshake:
```text
connect
<- msgpack(policy_server_config)
infer
-> msgpack({"endpoint": "infer", "session_id": "...", ...observation})
<- msgpack(ndarray | dict[str, ndarray])
reset
-> msgpack({"endpoint": "reset"})
<- msgpack({"status": "reset successful"})
```
If `endpoint` is omitted, the server treats the message as `infer`. Each
inference message produces one action response after the engine request
completes; action tokens or intermediate tensors are not streamed.
NumPy arrays use the marker format implemented by `openpi-client`. The server
also accepts the legacy vLLM NumPy marker representation on input. JSON text
frames are not observation messages.
## Sessions and Reset
- The API layer's current-session and first-call counters are scoped to each
WebSocket connection.
- `session_id` identifies model-side state across observations on that
connection. If omitted, it defaults to `default`.
- The first inference for a session is sent to the model with `reset=true`.
- Changing `session_id`, or sending the `reset` command, causes the next
inference to start with `reset=true`.
- The policy pipeline owns observation transforms and persistent model state,
normally keyed by `session_id`; the API layer forwards the raw observation
dictionary.
## Limits and Errors
- Maximum inbound payload size is 64 MiB.
- The server closes an idle connection after 30 seconds.
- Invalid binary input returns a MessagePack
`{"type":"error","message":"Invalid request payload"}` response.
- Inference failures return a generic `Internal inference error` without
exposing an internal traceback.
- Unsupported NumPy object, structured, and complex dtypes are rejected.
Observation keys, camera layout, state tensors, and action shapes are defined
by the loaded policy rather than by this transport. Read the handshake before
constructing observations. See the [DreamZero example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/dreamzero)
for a complete OpenPI client and DROID simulation loop.
+82
View File
@@ -0,0 +1,82 @@
# Realtime Audio WebSocket API
Use `WS /v1/realtime` to stream PCM audio into a compatible omni model and
receive incremental transcription and synthesized response audio. This is a
turn-based streaming API: it starts one generation over an audio stream and
finishes after the client closes that input stream.
For a model that listens while it is already speaking, use the experimental
[Full-Duplex API](full_duplex_api.md) instead.
## Quick Start
Start a compatible model. The currently documented example uses Qwen3-Omni:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
Run the provided client with a mono, 16-bit PCM, 16 kHz WAV file:
```bash
python examples/online_serving/qwen3_omni/openai_realtime_client.py \
--url ws://localhost:8091/v1/realtime \
--model Qwen/Qwen3-Omni-30B-A3B-Instruct \
--input-wav input_16k_mono.wav \
--output-wav response.wav \
--output-text response.txt
```
The client dependency is `websockets`.
## Protocol
Messages in both directions are JSON text frames. Audio carried inside an
event is base64-encoded raw PCM16.
| Direction | Event | Purpose |
|-----------|-------|---------|
| Server to client | `session.created` | Confirms that the WebSocket connection is ready |
| Client to server | `session.update` | Selects and validates the served model |
| Client to server | `input_audio_buffer.commit` with `final: false` | Starts generation over the incoming stream |
| Client to server | `input_audio_buffer.append` | Appends base64 PCM16 audio |
| Client to server | `input_audio_buffer.commit` with `final: true` | Marks the end of input |
| Server to client | `transcription.delta` | Carries incremental response text |
| Server to client | `transcription.done` | Carries final text and token usage |
| Server to client | `response.audio.delta` | Carries incremental PCM16 response audio |
| Server to client | `response.audio.done` | Marks the end of response audio |
| Server to client | `error` | Reports an invalid event, model, or audio payload |
A minimal client sends events in this order:
```json
{"type":"session.update","model":"Qwen/Qwen3-Omni-30B-A3B-Instruct"}
{"type":"input_audio_buffer.commit","final":false}
{"type":"input_audio_buffer.append","audio":"<base64-pcm16>"}
{"type":"input_audio_buffer.commit","final":true}
```
The initial non-final commit intentionally starts generation before all audio
has arrived. Continue sending `input_audio_buffer.append` events while the
engine is consuming the stream.
## Audio Handling
- Input is mono PCM16 at 16 kHz for the Qwen3-Omni example.
- `response.audio.delta.audio` contains base64-encoded PCM16 bytes.
- Read `sample_rate_hz` from each audio event instead of assuming an output
rate. Qwen3-Omni output is typically 24 kHz.
- Concatenate audio deltas in receive order to construct the output waveform.
- Set `input_audio_buffer.commit.final` to `true` after the last input chunk so
the server can terminate the streaming request.
## Availability and Limitations
The path is registered on the unified API server, but it is usable only when
the loaded pipeline implements realtime audio input and produces compatible
audio output. Unsupported deployments return an `error` event. The endpoint
does not provide duplex session resume, playback acknowledgement, overlap
policy, or barge-in controls.
See the [Qwen3-Omni online serving example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen3_omni)
for concurrency options, chunk pacing, and per-delta audio debugging.
+2 -1
View File
@@ -2,7 +2,7 @@
vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation. Supported TTS models include:
- **Qwen3-TTS** (`Qwen/Qwen3-TTS-12Hz-*`) -- Qwen3-based TTS with CustomVoice, VoiceDesign, and Base (voice cloning) task types. Output: 24 kHz.
- **Qwen3-TTS** (`Qwen/Qwen3-TTS-12Hz-*`) -- Qwen3-based TTS with CustomVoice, VoiceDesign, and Base (voice cloning) task types. Native output: 24 kHz; API output can be resampled to 8 kHz.
- **Fish Speech S2 Pro** (`fishaudio/s2-pro`) -- Dual-AR TTS with DAC codec. Supports text-to-speech and voice cloning via reference audio. Output: 44.1 kHz.
- **Voxtral TTS** (`mistralai/Voxtral-4B-TTS-2603`) -- AR + FlowMatching TTS with preset voices. Output: 24 kHz.
- **CosyVoice3** (`FunAudioLLM/Fun-CosyVoice3-0.5B-2512`) -- 2-stage talker + flow-matching code2wav. Voice cloning via `ref_audio` + `ref_text` (no presets). Output: 24 kHz.
@@ -116,6 +116,7 @@ Content-Type: application/json
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `sample_rate` | integer | model native | Target output sample rate. Qwen3-TTS supports 8000 or 24000 Hz; the model remains native 24 kHz internally and is resampled before encoding. |
| `task_type` | string | "CustomVoice" | TTS task type: CustomVoice, VoiceDesign, or Base |
| `language` | string | "Auto" | Language (see supported languages below) |
| `instructions` | string | "" | Voice style/emotion instructions |
+109
View File
@@ -0,0 +1,109 @@
# Standalone Experimental Servers
Most vLLM-Omni endpoints run on the unified server started with
`vllm serve <model> --omni`. The experimental full-duplex package also contains
two separate user-facing server processes. Their routes are not added to the
unified API server.
| Server | Transport | Role |
|--------|-----------|------|
| PersonaPlex compatibility server | WebSocket | Hosts the official PersonaPlex browser protocol and a raw-PCM alternative |
| JoyVL interaction server | HTTP | Adds state, memory, proactive decisions, and delegation in front of a separate OpenAI-compatible model backend |
## PersonaPlex Compatibility Server
Use this server when a client requires the official Moshi/PersonaPlex wire
protocol or the bundled browser UI. For the engine-owned unified duplex path,
use [`WS /v1/realtime?duplex=1`](full_duplex_api.md) instead.
Start it with:
```bash
HF_TOKEN=... CUDA_VISIBLE_DEVICES=0 python -m \
vllm_omni.experimental.fullduplex.personaplex.serving.server \
--port 8091
```
| Route | Purpose |
|-------|---------|
| `GET /` | Official PersonaPlex browser client, when its assets are available |
| `GET /health` | Readiness check |
| `WS /api/chat` | Moshi-compatible tagged binary protocol with Opus audio |
| `WS /v1/audio/duplex` | Simpler float32 raw-PCM protocol |
### `/api/chat` Protocol
The first byte of every binary frame is a message tag:
| Direction | Tag | Payload |
|-----------|-----|---------|
| Server to client | `0x00` | Ready handshake |
| Client to server | `0x01` | Opus microphone audio at 24 kHz |
| Server to client | `0x01` | Opus response audio |
| Server to client | `0x02` | UTF-8 inner-monologue text |
The optional `text_prompt` and `voice_prompt` query parameters select a persona
and voice. The default server admits one conversation; use `--batch-size N` to
configure multiple lockstep slots. Connections beyond capacity close with
WebSocket code 1013.
### `/v1/audio/duplex` Protocol
Send `{"type":"open","persona":"...","voice":"..."}` as a JSON text frame,
then exchange binary float32 PCM frames. The server returns generated float32
audio as binary frames and text as `{"type":"text","text":"..."}`. Send
`{"type":"close"}` and wait for `{"type":"done"}` to flush and finish.
See the [PersonaPlex example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/personaplex)
for the browser and headless clients.
## JoyVL Interaction Server
JoyVL is an orchestration layer, not a model-serving engine. It calls a
separate OpenAI-compatible backend and adds per-session frame history, memory,
persona policy, and optional delegation.
Start the model backend and orchestrator separately:
```bash
vllm serve jdopensource/JoyAI-VL-Interaction-Preview \
--served-model-name JoyAI-VL-Interaction-Preview \
--port 8092 \
--max-model-len 131072 \
--enable-prefix-caching \
--limit-mm-per-prompt '{"image":256,"video":1}'
python -m vllm_omni.experimental.fullduplex.joyvl.serving.server \
--port 8091 \
--main-backend-url http://127.0.0.1:8092/v1 \
--main-model JoyAI-VL-Interaction-Preview
```
| Method and route | Purpose |
|------------------|---------|
| `GET /health` | Readiness check |
| `GET /v1/models` | Reports the configured interaction model |
| `POST /v1/chat/completions` | Processes one frame or interaction tick |
| `POST /reset` | Resets a session |
| `POST /v1/streaming/reset` | Alias for `/reset` |
| `POST /v1/streaming/persona` | Changes the session persona |
Send one multimodal Chat Completions request per video frame, normally around
one frame per second, and identify the session with `x-session-id`. Responses
include an `interaction` block whose action is `silence`, `response`, or
`delegate`.
```bash
curl http://127.0.0.1:8091/v1/chat/completions \
-H 'content-type: application/json' \
-H 'x-session-id: demo' \
-d '{
"messages": [{"role": "user", "content": [
{"type": "text", "text": "Alert me if a fire appears"},
{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,..."}}
]}]
}'
```
See the [JoyVL recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/JD/JoyAI-VL-Interaction.md)
for memory, personas, delegation backends, and the browser UI.
@@ -0,0 +1,92 @@
# Streaming Video Output WebSocket API
Use `WS /v1/realtime/video` to receive fragmented MP4 chunks while a compatible
diffusion video model is generating. This endpoint streams generated video
output; it is unrelated to [`WS /v1/video/chat/stream`](video_stream_api.md),
which streams live video input into a conversational model.
## Quick Start
Start a model that supports step execution and enable diffusion streaming
output:
```bash
vllm serve BestWishYsh/Helios-Distilled --omni \
--diffusion-streaming-output \
--port 8091
```
Install the client dependency and run the provided client:
```bash
pip install av websockets
python examples/online_serving/streaming_video_generation/streaming_video_client.py \
--host 127.0.0.1 \
--port 8091 \
--model BestWishYsh/Helios-Distilled \
--prompt "A serene lakeside sunrise with mist over the water" \
--output output.mp4
```
## Protocol
Control events are JSON text frames. Each `video.chunk_metadata` event is
immediately followed by one binary WebSocket frame containing fragmented MP4
bytes.
| Direction | Event or frame | Purpose |
|-----------|----------------|---------|
| Client to server | `session.start` | Starts generation with a prompt and video parameters |
| Server to client | `video.start` | Confirms the request ID, format, and accepted configuration |
| Server to client | `video.chunk_metadata` | Describes the next binary media frame |
| Server to client | Binary frame | Carries one fragmented MP4 chunk |
| Client to server | `session.interaction` | Queues an optional prompt update during generation |
| Server to client | `session.interaction.queued` | Confirms that an interaction was queued |
| Client to server | `session.ping` | Refreshes the stall clock |
| Server to client | `session.pong` | Acknowledges a ping |
| Client to server | `session.stop` | Cancels the active generation |
| Server to client | `session.done` | Reports completion or cancellation and the chunk count |
| Server to client | `error` | Reports validation, generation, or timeout failures |
Start a minimal session with:
```json
{
"type": "session.start",
"model": "BestWishYsh/Helios-Distilled",
"prompt": "A serene lakeside sunrise",
"format": "m4s",
"width": 640,
"height": 384,
"fps": 16,
"num_frames": 99
}
```
Only `m4s` output is currently accepted. Concatenate the binary frames in
receive order, then remux the fragmented stream if the target player requires
a progressive MP4 file.
## Mid-Generation Interaction
Models that implement interaction updates can change the active prompt:
```json
{
"type": "session.interaction",
"interaction": {
"event_id": "sunlight-update",
"event": {"prompt": "Sunlight breaks through the mist"},
"transition_chunks": 2
}
}
```
This feature is pipeline-dependent. An enabled endpoint does not imply that
the loaded model supports midway prompt updates.
The default start timeout is 10 seconds and the server reports a stall after
about 60 seconds without generation progress or a `session.ping`. See the
[streaming video generation example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/streaming_video_generation)
for prompt-update scheduling and browser playback.
+1 -1
View File
@@ -138,7 +138,7 @@ The following tables show which models support each feature:
| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ |
| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **SenseNova-U1** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **SenseNova-U1 / U1.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **Stable-Diffusion-XL** | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
+1 -1
View File
@@ -49,7 +49,7 @@ scenario below for live barge-in validation on the target GPU.
Default:
```bash
vllm-omni serve openbmb/MiniCPM-o-4_5 \
vllm serve openbmb/MiniCPM-o-4_5 \
--omni \
--deploy-config vllm_omni/deploy/minicpmo_4_5.yaml \
--trust-remote-code \
@@ -37,7 +37,7 @@ During generation the client normally sends only `session.start` and then receiv
## Install Client Dependency
```bash
pip install websockets
pip install av websockets
# For the Gradio demo:
pip install 'vllm-omni[demo]' websockets
```
@@ -8,7 +8,7 @@ the received bytes to disk when the session finishes or is interrupted.
Image/reference input is intentionally not supported in this example yet.
Requirements:
pip install websockets
pip install av websockets
"""
import argparse
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""OpenAI-compatible client for Qwen3-TTS via /v1/audio/speech endpoint.
This script demonstrates how to use the OpenAI-compatible speech API
@@ -74,6 +77,8 @@ def run_tts_generation(args) -> None:
"voice": args.speaker,
"response_format": args.response_format,
}
if args.sample_rate is not None:
payload["sample_rate"] = args.sample_rate
# Add optional parameters
if args.instructions:
@@ -246,6 +251,13 @@ def parse_args():
choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
help="Audio output format (default: wav)",
)
parser.add_argument(
"--sample-rate",
type=int,
default=None,
choices=[8000, 24000],
help="Output sample rate in Hz (default: model native)",
)
parser.add_argument(
"--output",
"-o",
+86 -13
View File
@@ -1,11 +1,14 @@
# Boogu-Image
> Text-to-image and image-editing online serving (Boogu-Image-0.1-Base / -Edit)
> Text-to-image and image-editing online serving
(Boogu-Image-0.1-Base / -Edit / -Edit-Turbo)
## Summary
- Vendor: Boogu
- Model: `Boogu/Boogu-Image-0.1-Base` (text-to-image), `Boogu/Boogu-Image-0.1-Edit` (image editing)
- Model: `Boogu/Boogu-Image-0.1-Base` (text-to-image),
`Boogu/Boogu-Image-0.1-Edit` (image editing), and
`Boogu/Boogu-Image-0.1-Edit-Turbo` (four-step image editing)
- Task: Text-to-image generation and text-guided image editing (TI2I)
- Mode: Online serving with the OpenAI-compatible API
- Maintainer: Community
@@ -13,19 +16,23 @@
## When to use this recipe
Use this recipe when you want a known-good starting point for serving
`Boogu/Boogu-Image-0.1-Base` with vLLM-Omni's native pipeline (no
`Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit`, or
`Boogu/Boogu-Image-0.1-Edit-Turbo` with vLLM-Omni's native pipeline (no
`--diffusion-load-format diffusers`, and the upstream `boogu` package is not
required).
Boogu-Image-0.1 is an Apache-2.0 unified image generation and editing model
family. This recipe covers the Base text-to-image checkpoint, which pairs a
Qwen3-VL multimodal encoder with a Diffusion Transformer (DiT) and a flow-match
Euler scheduler with time-shift. It handles photorealistic generation and
Chinese/English text rendering.
family. The Base text-to-image checkpoint pairs a Qwen3-VL multimodal encoder
with a Diffusion Transformer (DiT) and a flow-match Euler scheduler with
time-shift. It handles photorealistic generation and Chinese/English text
rendering. The Edit and Edit-Turbo checkpoints use the same native
`BooguImagePipeline` for image editing.
## References
- Upstream model card: <https://huggingface.co/Boogu/Boogu-Image-0.1-Base>
- Edit model card: <https://huggingface.co/Boogu/Boogu-Image-0.1-Edit>
- Edit-Turbo 1K hotfix: <https://huggingface.co/Boogu/Boogu-Image-0.1-Edit-Turbo/tree/hotfix-1k-20260708>
- Project page: <https://boogu.org>
- GitHub: <https://github.com/boogu-project/Boogu-Image>
- Related example: [`examples/online_serving/text_to_image/`](../../examples/online_serving/text_to_image/README.md)
@@ -93,7 +100,7 @@ curl -s http://localhost:8091/v1/chat/completions \
- **Memory usage:** ~34.6 GiB on GPU. Use `--vae-use-slicing --vae-use-tiling`
to trim peak VRAM if needed.
- **Key flags:**
- `--omni` — enables vLLM-Omni diffusion serving.
- `--omni` — enables vLLM-Omni diffusion serving.
- **Guidance:** Boogu-Image uses `guidance_scale` (mapped to the upstream
`text_guidance_scale`); the default is `4.0`. Classifier-free guidance is
active whenever `guidance_scale > 1.0`.
@@ -111,13 +118,13 @@ the image-editing (TI2I) path activates automatically when a request carries a
reference image. The Base text-to-image path is unaffected (no reference image
is sent).
#### Command
### Command
```bash
vllm serve Boogu/Boogu-Image-0.1-Edit --omni --port 8091
```
#### Verification
### Verification
Edit an image with `/v1/images/edits` (the model-card example — change a photo
to a colored-pencil drawing). Diffusion parameters are plain multipart form
@@ -157,14 +164,14 @@ curl -s http://localhost:8091/v1/chat/completions \
}' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > edited.png
```
#### Notes
### Notes
- **Single reference image:** only one input image is supported for now (the
upstream "Only support 1 reference image for now" limit).
- **Guidance semantics:**
- `guidance_scale` = text guidance (upstream `text_guidance_scale`, default
- `guidance_scale` = text guidance (upstream `text_guidance_scale`, default
`4.0`); `> 1.0` enables text CFG. Editing typically uses `5.0`.
- `guidance_scale_2` = image guidance (upstream `image_guidance_scale`,
- `guidance_scale_2` = image guidance (upstream `image_guidance_scale`,
default `1.0` = off). Setting it `> 1.0` enables the double-guidance path
(3 model predictions per step), steering more strongly toward the reference
image.
@@ -178,3 +185,69 @@ curl -s http://localhost:8091/v1/chat/completions \
supports Boogu-Image-Edit directly
(`--model Boogu/Boogu-Image-0.1-Edit --guidance-scale 5.0`, optional
`--guidance-scale-2`).
## Fast image editing (Boogu-Image-0.1-Edit-Turbo)
Edit-Turbo is the distilled image-editing checkpoint. Its `model_index.json`
declares `BooguImagePipeline`, so it uses the same native TI2I path as Edit.
The upstream 1K hotfix is recommended for stable results and is pinned below
to avoid silently loading the older checkpoint from the repository's default
revision.
### Command
```bash
vllm serve Boogu/Boogu-Image-0.1-Edit-Turbo \
--omni \
--revision hotfix-1k-20260708 \
--port 8091
```
### Verification
Edit an image with the distilled four-step settings:
```bash
curl -s http://localhost:8091/v1/images/edits \
-F model="Boogu/Boogu-Image-0.1-Edit-Turbo" \
-F image="@input.png" \
-F prompt="Change the style to a colored pencil drawing." \
-F num_inference_steps=4 \
-F guidance_scale=1.0 \
-F seed=42 \
| jq -r '.data[0].b64_json' | base64 -d > edited-turbo.png
```
Or via chat completions:
```bash
curl -s http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"messages": [
{"role": "user", "content": [
{"type": "image_url", "image_url": {"url": "data:image/png;base64,<BASE64>"}},
{"type": "text", "text": "Change the style to a colored pencil drawing."}
]}
],
"extra_body": {
"num_inference_steps": 4,
"guidance_scale": 1.0,
"seed": 42
}
}' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > edited-turbo.png
```
### Notes
- **Pinned revision:** use `hotfix-1k-20260708`; upstream recommends the 1K
hotfix over the 1.5K variant for more stable results.
- **Recommended settings:** `num_inference_steps=4` and
`guidance_scale=1.0`. Edit-Turbo is guidance-distilled, meaning the guidance
behavior is already baked into the distilled checkpoint. A scale of `1.0`
avoids applying additional classifier-free guidance; unlike the regular Edit
checkpoint, it should not be increased to `5.0`.
- **Reference images:** the same single-reference and `align_res` behavior as
the regular Edit checkpoint applies.
- **Known limitations:** the same single-GPU limitations as Base and Edit apply;
CPU offload, Cache-DiT, and multi-GPU parallelism are not yet validated.
+1
View File
@@ -1142,3 +1142,4 @@ vllm serve "${MODEL_ROOT}/FL2VA" \
- [Supported models](../../docs/models/supported_models.md)
- [Video API](../../docs/serving/videos_api.md)
- [Diffusion parallelism](../../docs/user_guide/diffusion/parallelism/overview.md)
- [Temporal VAE chunk callback design](../../docs/design/feature/minimax_h3_temporal_vae_chunks.md)
+1
View File
@@ -78,6 +78,7 @@ recipes/
| [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat | 1x A100 80GB |
| [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB / 1x AMD MI300X |
| [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x or 2x H200 / 1x AMD MI300X |
| [`SenseNova/SenseNova-U1.5.md`](./SenseNova/SenseNova-U1.5.md) | Unified image generation and understanding, with an 8-step distilled LoRA | 1x A800 80GB |
| [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB / 1x AMD MI300X |
| [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB |
| [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 |
+159
View File
@@ -0,0 +1,159 @@
# SenseNova-U1.5
> Unified image generation and understanding, with an official 8-step distilled LoRA
## Summary
- Vendor: SenseNova
- Model: `sensenova/SenseNova-U1.5-8B-MoT`
- LoRA: `sensenova/SenseNova-U1.5-8B-MoT-LoRAs` (`SenseNova-U1.5-8B-MoT-LoRA-8step.safetensors`)
- Task: text2img, img2img, img2text (visual understanding), text2text (chat)
- Mode: Offline inference, Online serving (OpenAI-compatible API)
- Maintainer: Community
## When to use this recipe
U1.5 runs on the same `SenseNovaU1Pipeline` as U1 — the checkpoint keeps
`model_type: neo_chat`, so it resolves without `--model-class-name`. Relative to U1 the
config flips two fields, `use_pixel_head` to `true` (the flow-matching head becomes a
`ConvDecoder`) and `noise_scale_max_value` from 8.0 to 16.0; both are already read by
`SenseNovaU1Config`. The checkpoint is 13 shards / 50.2 GB on disk, 30.3 GB of that fp32, and
loads to roughly 34 GB in bf16.
Use this recipe for U1.5 specifically, including its distilled few-step LoRA. For U1 see
[SenseNova-U1](SenseNova-U1.md).
## Hardware Support
### GPU
#### 1x A800 80GB
- 1x NVIDIA A800 80GB PCIe with an Intel Xeon Gold 6336Y, driver 595.84, Linux
- Python 3.12.13, vLLM 0.28.0, torch 2.13.0+cu130, CUDA 13.0, diffusers 0.40.0,
transformers 5.14.1, flashinfer 0.6.16.post3, BF16, TP=1
- Peak GPU memory: 34.3 GB at 1024x1024, 36.4 GB at 1536x2720
##### Text-to-Image
```bash
python examples/offline_inference/text_to_image/text_to_image.py \
--model sensenova/SenseNova-U1.5-8B-MoT \
--prompt "Close portrait of an elderly woman by a farmhouse window, warm natural light." \
--width 1024 --height 1024 \
--seed 42 --num-inference-steps 50 --cfg-scale 4.0 \
--extra-body '{"think": false, "cfg_norm": "none", "timestep_shift": 3.0, "t_eps": 0.02}' \
--output sensenova_u15_t2i.png
```
Think mode (`"think": true`) is recommended for higher image quality.
##### Image-to-Image Editing
```bash
python examples/offline_inference/image_to_image/image_edit.py \
--model sensenova/SenseNova-U1.5-8B-MoT \
--prompt "Turn this into an oil painting" \
--image input.png --resolution 1024 \
--seed 42 --num-inference-steps 25 --cfg-scale 4.0 \
--extra-args '{"think": true, "img_cfg_scale": 1.0, "cfg_norm": "none", "timestep_shift": 3.0}' \
--output sensenova_u15_edit.png
```
##### 8-step distilled LoRA
The distilled LoRA is fused into the generation tower at load time:
```bash
python examples/offline_inference/text_to_image/text_to_image.py \
--model sensenova/SenseNova-U1.5-8B-MoT \
--lora-path SenseNova-U1.5-8B-MoT-LoRA-8step.safetensors --lora-backend distill \
--prompt "Close portrait of an elderly woman by a farmhouse window, warm natural light." \
--width 1024 --height 1024 \
--seed 42 --num-inference-steps 8 --cfg-scale 1.0 \
--extra-body '{"think": false, "cfg_norm": "none", "timestep_shift": 3.0, "t_eps": 0.02}' \
--output sensenova_u15_lora8.png
```
**Use `--cfg-scale 1.0` with this LoRA.** It is distilled with DMD and runs without
classifier-free guidance; the default `4.0` applies guidance twice and produces a
blown-out, posterised image.
##### Online serving
```bash
vllm serve sensenova/SenseNova-U1.5-8B-MoT --omni --port 8091
python examples/online_serving/sensenova_u1/openai_chat_client.py \
-s http://127.0.0.1:8091 -m img2text -i input.png -p "Describe this image."
```
`-s` takes the base URL; the client appends `/v1` itself.
#### Measured latency (1x A800 80GB, 25 steps, median of 3 after a warmup)
| Resolution | Step latency | Total |
| --- | --- | --- |
| 1024x1024 | 209.5 ms | 5.24 s |
| 1536x1536 | 458.5 ms | 11.46 s |
At 1536x2720 with 50 steps, end-to-end is 43.0 s.
#### Measured latency (1x NVIDIA H200 139GB, single run, no warmup)
Reported by @hsliuustc0106 against PR head `29c090d`, on one reserved CUDA device.
- Python 3.12.13, vLLM 0.28.0, torch 2.13.0+cu130, CUDA 13.0, diffusers 0.40.0,
transformers 5.14.1, flashinfer 0.6.16.post3, BF16, TP=1
- `sensenova/SenseNova-U1.5-8B-MoT` with `SenseNova-U1.5-8B-MoT-LoRA-8step.safetensors`,
seed 42, 1024x1024
| Case | Stage latency | Peak GPU memory |
| --- | --- | --- |
| 50 steps, cfg 4.0, think off | 9017 ms | 34384 MiB |
| 8 steps, cfg 1.0, no LoRA control | 624.00 ms, 77.85 ms/step | 34376 MiB |
| 8 steps, cfg 1.0, distilled LoRA | 623.85 ms, 77.98 ms/step | 34364 MiB |
| 50 steps, cfg 4.0, think on, `VLLM_OMNI_SENSENOVA_PAGED_DECODE=1` | 81622 ms | 34594 MiB |
Each row runs the Text-to-Image command above with `--num-inference-steps` and
`--cfg-scale` as listed and `think` set through `--extra-body`; the LoRA row adds
`--lora-path` and `--lora-backend distill`.
Single runs without a preceding warmup, so the think-on row carries the first-request
compile rather than a steady state. The distilled LoRA fused into 168 parameters, and its
image differs from the same-seed no-LoRA control in 1,048,574 of 1,048,576 pixels, MAE 58.62.
```bash
pytest -m "cpu and not cuda" tests/diffusion/models/sensenova_u1/ tests/diffusion/lora/test_loader.py -q
pytest -m "cpu and not cuda" tests/config/test_environment_variables.py tests/diffusion/test_diffusion_worker.py -q
pytest -m cuda tests/diffusion/models/sensenova_u1/test_sensenova_u1_attention_gqa.py \
tests/diffusion/models/sensenova_u1/test_sensenova_u1_paged_decode.py -q --tb=short
```
71, 13 and 9 passed. Serving reached `/health`; `img2text` returned a 1921-character
description and `text2text` the expected answer.
#### Verification
```bash
pytest -q tests/diffusion/models/sensenova_u1/
```
## Notes
- Both `use_pixel_head` and `noise_scale_max_value` come from `config.json`; no flag is needed.
- The checkpoint carries no `configuration_neo_chat.py`, so the loader logs a few
"does not appear to have a file named configuration_neo_chat.py" errors during startup and
then proceeds on the in-tree config. Generation is unaffected.
- The LoRA targets the generation tower only (`*_mot_gen`); understanding-tower weights are
untouched.
- Autoregressive decode (think, text-to-text and image-to-text) runs on a paged K/V cache under
a captured CUDA graph. It falls back to the ordinary cache when the device or the bundled
`flash_attn_varlen_func` cannot support it; set `VLLM_OMNI_SENSENOVA_PAGED_DECODE=0` to force
that fallback.
- The first request after startup costs about 0.7 s more than the steady state whether the paged
path is on or off. Measured on one A800 with the inductor, triton and vLLM compile caches all
cleared, median of three runs: 718 ms above steady with the path on, 679 ms with it off.
- Each request captures its own graphs, and a think request captures twice because the sequence
grows past the 512 bucket, so the capture cost is paid per request rather than once at
startup.
+42
View File
@@ -0,0 +1,42 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from pathlib import Path
from shlex import split
import pytest
import yaml
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
AMD_MERGE_PIPELINE = Path(".buildkite/amd/test-amd-merge.yml")
def _find_step(label: str) -> dict:
pipeline = yaml.safe_load(AMD_MERGE_PIPELINE.read_text(encoding="utf-8"))
def walk(steps: list[dict]) -> dict | None:
for step in steps:
if step.get("label") == label:
return step
if nested := walk(step.get("steps", [])):
return nested
return None
step = walk(pipeline.get("steps", []))
assert step is not None, f"missing AMD pipeline step: {label}"
return step
def test_qwen3_tts_base_preserves_advanced_model_arguments() -> None:
step = _find_step("Qwen3-TTS Base E2E Test")
commands = step["commands"]
assert all("bash -c" not in command for command in commands)
pytest_command = next(command for command in commands if "pytest" in command)
argv = split(pytest_command)
marker_index = argv.index("-m")
run_level_index = argv.index("--run-level")
assert argv[marker_index + 1] == "advanced_model and cuda"
assert argv[run_level_index + 1] == "advanced_model"
+142 -43
View File
@@ -41,6 +41,11 @@ from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
SINGLE_STAGE_PIPE_CFG = (StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),)
# Validation strings to match against in post init
NO_STAGES_MATCH_STR = "no stages"
NO_TERMINAL_STAGE_MATCH_STR = "No terminal stage"
@pytest.fixture(autouse=True)
def _stable_test_platform(monkeypatch):
@@ -618,6 +623,7 @@ class TestPipelineDiscovery:
p = PipelineConfig(
model_type="custom_collide",
hf_architectures=("SomeCollidingArch",),
stages=SINGLE_STAGE_PIPE_CFG,
)
assert p.hf_architectures == ("SomeCollidingArch",)
@@ -639,41 +645,42 @@ class TestStagePipelineConfig:
class TestPipelineConfigNew:
def test_simple_init(self):
# Ensure pipeline config validates in post init.
PipelineConfig(
model_type="t",
model_arch="A",
stages=SINGLE_STAGE_PIPE_CFG,
)
def test_frozen(self):
p = PipelineConfig(model_type="t", model_arch="A")
p = PipelineConfig(
model_type="t",
model_arch="A",
stages=SINGLE_STAGE_PIPE_CFG,
)
with pytest.raises(AttributeError):
p.model_type = "changed"
def test_validate_valid(self):
p = PipelineConfig(
model_type="t",
model_arch="A",
stages=(
StagePipelineConfig(stage_id=0, model_stage="a"),
StagePipelineConfig(stage_id=1, model_stage="b", input_sources=(0,), final_output=True),
),
)
assert p.validate() == []
def test_validate_no_stages(self):
p = PipelineConfig(model_type="t", model_arch="A")
assert any("no stages" in e.lower() for e in p.validate())
with pytest.raises(ValueError, match=NO_STAGES_MATCH_STR):
PipelineConfig(model_type="t", model_arch="A")
def test_validate_no_terminal_stage(self):
"""A pipeline with no ``final_output`` stage can never emit a result."""
p = PipelineConfig(
model_type="t",
model_arch="A",
stages=(
StagePipelineConfig(stage_id=0, model_stage="a"),
StagePipelineConfig(stage_id=1, model_stage="b", input_sources=(0,)),
),
)
assert any("no terminal stage" in e.lower() for e in p.validate())
with pytest.raises(ValueError, match=NO_TERMINAL_STAGE_MATCH_STR):
PipelineConfig(
model_type="t",
model_arch="A",
stages=(
StagePipelineConfig(stage_id=0, model_stage="a"),
StagePipelineConfig(stage_id=1, model_stage="b", input_sources=(0,)),
),
)
def test_validate_terminal_stage_need_not_be_last(self):
"""The check is cycle-agnostic: any stage may carry ``final_output``."""
p = PipelineConfig(
def test_pipeline_can_have_final_output_in_any_stage(self):
"""Ensure any stage may carry ``final_output``."""
PipelineConfig(
model_type="t",
model_arch="A",
stages=(
@@ -681,15 +688,20 @@ class TestPipelineConfigNew:
StagePipelineConfig(stage_id=1, model_stage="b", input_sources=(0,)),
),
)
assert p.validate() == []
class TestPipelineRegistration:
def test_resolve_pipeline_prefers_deploy_pipeline_key(self, clean_pipeline_registry, tmp_path):
deploy_key = "deploy_selected_pipeline"
model_type_key = "hf_model_type_pipeline"
deploy_pipe = PipelineConfig(model_type=deploy_key)
model_type_pipe = PipelineConfig(model_type=model_type_key)
deploy_pipe = PipelineConfig(
model_type=deploy_key,
stages=SINGLE_STAGE_PIPE_CFG,
)
model_type_pipe = PipelineConfig(
model_type=model_type_key,
stages=SINGLE_STAGE_PIPE_CFG,
)
register_pipeline(deploy_pipe)
register_pipeline(model_type_pipe)
@@ -715,7 +727,11 @@ class TestPipelineRegistration:
tmp_path,
):
model_type_key = "registered_model_type_pipeline"
register_pipeline(PipelineConfig(model_type=model_type_key))
pipeline_cfg = PipelineConfig(
model_type=model_type_key,
stages=SINGLE_STAGE_PIPE_CFG,
)
register_pipeline(pipeline_cfg)
deploy_path = tmp_path / "deploy.yaml"
deploy_path.write_text("pipeline: missing_pipeline\n", encoding="utf-8")
@@ -750,6 +766,7 @@ class TestPipelineRegistration:
pipe_cfg = PipelineConfig(
model_type=pipeline_key,
hf_architectures=("ArchitectureFallbackForTest",),
stages=SINGLE_STAGE_PIPE_CFG,
)
register_pipeline(pipe_cfg)
@@ -780,15 +797,18 @@ class TestPipelineRegistration:
model_type="rejected_by_predicate",
hf_architectures=(shared_arch,),
hf_config_predicate=rejecting_predicate,
stages=SINGLE_STAGE_PIPE_CFG,
)
raises_cfg = PipelineConfig(
model_type="raises_in_predicate",
hf_architectures=(shared_arch,),
hf_config_predicate=raising_predicate,
stages=SINGLE_STAGE_PIPE_CFG,
)
accept_cfg = PipelineConfig(
model_type="accepted_by_predicate",
hf_architectures=(shared_arch,),
stages=SINGLE_STAGE_PIPE_CFG,
)
register_pipeline(reject_cfg)
register_pipeline(raises_cfg)
@@ -822,6 +842,7 @@ class TestPipelineRegistration:
model_type="predicate_without_arch_match",
hf_architectures=("DifferentArchitectureForTest",),
hf_config_predicate=predicate,
stages=SINGLE_STAGE_PIPE_CFG,
)
register_pipeline(pipe_cfg)
@@ -845,6 +866,7 @@ class TestPipelineRegistration:
resolved_cfg = PipelineConfig(
model_type="callable_resolved_pipeline",
hf_architectures=("CallableArchitectureForTest",),
stages=SINGLE_STAGE_PIPE_CFG,
)
seen_hf_configs = []
@@ -918,7 +940,10 @@ class TestPipelineRegistration:
def test_pipeline_registration(self, clean_pipeline_registry):
"""Ensure that we can register and create a custom pipeline config."""
new_model_type = "new_model_type"
pipe_cfg = PipelineConfig(model_type=new_model_type)
pipe_cfg = PipelineConfig(
model_type=new_model_type,
stages=SINGLE_STAGE_PIPE_CFG,
)
# Register the new PipelineConfig
assert new_model_type not in OMNI_PIPELINES
@@ -961,7 +986,10 @@ class TestPipelineRegistration:
def custom_resolver(
hf_config: FakeConfig,
) -> PipelineConfig:
return PipelineConfig(model_type=resolved_type)
return PipelineConfig(
model_type=resolved_type,
stages=SINGLE_STAGE_PIPE_CFG,
)
# Register the new PipelineConfig
assert new_model_type not in OMNI_PIPELINES
@@ -1018,7 +1046,7 @@ class TestPipelineRegistration:
pipe_cfg = PipelineConfig(
model_type=pipeline_key,
model_arch="DeployOnlyArch",
stages=(StagePipelineConfig(stage_id=0, model_stage="diffusion"),),
stages=SINGLE_STAGE_PIPE_CFG,
)
register_pipeline(pipe_cfg)
@@ -1051,7 +1079,11 @@ class TestPipelineRegistration:
def test_structured_path_loads_explicit_deploy_config_once(self, clean_pipeline_registry, tmp_path):
pipeline_key = "single_load_pipeline"
register_pipeline(PipelineConfig(model_type=pipeline_key))
pipeline_cfg = PipelineConfig(
model_type=pipeline_key,
stages=SINGLE_STAGE_PIPE_CFG,
)
register_pipeline(pipeline_cfg)
deploy_path = tmp_path / "single_load.yaml"
deploy_path.write_text(f"pipeline: {pipeline_key}\n", encoding="utf-8")
@@ -1081,6 +1113,7 @@ class TestPipelineRegistration:
pipeline = PipelineConfig(
model_type="pipeline_with_default",
default_deploy_config_name=default_name,
stages=SINGLE_STAGE_PIPE_CFG,
)
with patch(
@@ -1100,8 +1133,16 @@ class TestPipelineRegistration:
OmniServingCapability.COMPLETIONS,
"pipeline_a blocks completions",
)
pipe_a = PipelineConfig(model_type="detect_type", endpoint_restrictions=(restriction,))
pipe_b = PipelineConfig(model_type="override_type", endpoint_restrictions=())
pipe_a = PipelineConfig(
model_type="detect_type",
endpoint_restrictions=(restriction,),
stages=SINGLE_STAGE_PIPE_CFG,
)
pipe_b = PipelineConfig(
model_type="override_type",
endpoint_restrictions=(),
stages=SINGLE_STAGE_PIPE_CFG,
)
register_pipeline(pipe_a)
register_pipeline(pipe_b)
@@ -1641,6 +1682,7 @@ stages:
model_stage="ar",
execution_type=StageExecutionType.LLM_AR,
requires_multimodal_data=True,
final_output=True,
),
),
)
@@ -1661,6 +1703,7 @@ stages:
model_stage="ar",
execution_type=StageExecutionType.LLM_AR,
scheduler_cls=scheduler_cls,
final_output=True,
),
),
)
@@ -1906,7 +1949,6 @@ class TestQwen3OmniPipeline:
assert isinstance(p, PipelineConfig)
assert p.model_arch == "Qwen3OmniMoeForConditionalGeneration"
assert len(p.stages) == 3
assert p.validate() == []
def test_thinker(self):
p = resolve_pipeline_config(
@@ -1960,7 +2002,6 @@ class TestQwen2_5OmniPipeline:
assert isinstance(p, PipelineConfig)
assert p.model_arch == "Qwen2_5OmniForConditionalGeneration"
assert len(p.stages) == 3
assert p.validate() == []
def test_thinker(self):
p = resolve_pipeline_config("qwen2_5_omni")
@@ -2009,7 +2050,6 @@ class TestQwen3TTSPipeline:
assert p is not None
assert p.model_arch == "Qwen3TTSTalkerForConditionalGeneration"
assert len(p.stages) == 2
assert p.validate() == []
def test_talker_stage(self):
p = resolve_pipeline_config("qwen3_tts")
@@ -2089,7 +2129,6 @@ class TestMingFlashOmniPipeline:
assert isinstance(p, PipelineConfig)
assert p.model_arch == "MingFlashOmniForConditionalGeneration"
assert len(p.stages) == 2
assert p.validate() == []
def test_thinker_stage(self):
p = resolve_pipeline_config("ming_flash_omni")
@@ -2146,7 +2185,6 @@ class TestMingFlashOmniPipeline:
assert isinstance(p, PipelineConfig)
assert p.model_arch == "MingFlashOmniTalkerForConditionalGeneration"
assert len(p.stages) == 1
assert p.validate() == []
def test_tts_stage(self):
p = resolve_pipeline_config("ming_flash_omni_tts")
@@ -2203,7 +2241,6 @@ class TestMingFlashOmniPipeline:
assert isinstance(p, PipelineConfig)
assert p.model_arch == "MingFlashOmniForConditionalGeneration"
assert len(p.stages) == 1
assert p.validate() == []
def test_thinker_only_stage(self):
p = resolve_pipeline_config("ming_flash_omni_thinker_only")
@@ -2242,7 +2279,6 @@ class TestMingFlashOmniPipeline:
assert p is not None
assert p.model_arch == "MingFlashOmniForConditionalGeneration"
assert len(p.stages) == 2
assert p.validate() == []
def test_image_thinker_stage(self):
s = resolve_pipeline_config("ming_flash_omni_image").get_stage(0)
@@ -2363,6 +2399,21 @@ class TestBaseConfigInheritance:
# CI overrides max_tokens
assert s0["max_tokens"] == 150
def test_qwen3_omni_colocate_async_bounds_only_rocm_kv_cache(self):
ci_path = Path(get_deploy_config_path("ci/qwen3_omni_moe_colocate_async.yaml"))
pipeline = resolve_pipeline_config("qwen3_omni_moe_thinker_only")
assert isinstance(pipeline, PipelineConfig)
cuda = _apply_platform_overrides(load_deploy_config(ci_path), platform="cuda")
cuda_stage = merge_pipeline_deploy(pipeline, cuda)[0]
assert cuda_stage.yaml_engine_args["gpu_memory_utilization"] == 0.9
assert "kv_cache_memory_bytes" not in cuda_stage.yaml_engine_args
rocm = _apply_platform_overrides(load_deploy_config(ci_path), platform="rocm")
rocm_stage = merge_pipeline_deploy(pipeline, rocm)[0]
assert "gpu_memory_utilization" not in rocm_stage.yaml_engine_args
assert rocm_stage.yaml_engine_args["kv_cache_memory_bytes"] == 2 * 1024**3
def test_pure_inheritance_overlay(self, tmp_path):
"""An overlay with only ``base_config`` inherits everything."""
base = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
@@ -2824,6 +2875,54 @@ stages:
assert omega_config.engine_args.max_model_len == 8192
assert omega_config.engine_args.enforce_eager is True
def test_cli_attention_shorthand_replaces_yaml_structured_default(self):
stage = StageConfig(
stage_id=0,
model_stage="dit",
stage_type=StageType.DIFFUSION,
yaml_engine_args={
"diffusion_attention_config": {
"default": {"backend": "FLASH_ATTN"},
"per_role": {"cross": {"backend": "TORCH_SDPA"}},
},
},
runtime_overrides={"diffusion_attention_backend": "SAGE_ATTN"},
)
engine_args = stage.to_omegaconf().engine_args
assert engine_args.diffusion_attention_backend == "SAGE_ATTN"
assert "default" not in engine_args.diffusion_attention_config
assert engine_args.diffusion_attention_config.per_role.cross.backend == "TORCH_SDPA"
def test_cli_attention_shorthand_keeps_yaml_per_role_only_config(self):
stage = StageConfig(
stage_id=0,
model_stage="dit",
stage_type=StageType.DIFFUSION,
yaml_engine_args={"diffusion_attention_config": {"per_role": {"cross": {"backend": "TORCH_SDPA"}}}},
runtime_overrides={"diffusion_attention_backend": "SAGE_ATTN"},
)
engine_args = stage.to_omegaconf().engine_args
assert engine_args.diffusion_attention_backend == "SAGE_ATTN"
assert engine_args.diffusion_attention_config == {"per_role": {"cross": {"backend": "TORCH_SDPA"}}}
def test_cli_attention_config_replaces_yaml_shorthand(self):
stage = StageConfig(
stage_id=0,
model_stage="dit",
stage_type=StageType.DIFFUSION,
yaml_engine_args={"diffusion_attention_backend": "TORCH_SDPA"},
runtime_overrides={"diffusion_attention_config": {"default": {"backend": "SAGE_ATTN"}}},
)
engine_args = stage.to_omegaconf().engine_args
assert "diffusion_attention_backend" not in engine_args
assert engine_args.diffusion_attention_config.default.backend == "SAGE_ATTN"
class TestSentinelDefaultPrecedence:
"""Caller-typed (non-None) values win over YAML; None values fall through
+2 -2
View File
@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
EnvironmentVariableCategory.PUBLIC_OMNI: 24,
EnvironmentVariableCategory.INHERITED_VLLM: 20,
EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
EnvironmentVariableCategory.MODEL_SPECIFIC: 56,
EnvironmentVariableCategory.MODEL_SPECIFIC: 57,
EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20,
EnvironmentVariableCategory.INTERNAL: 2,
}
@@ -178,7 +178,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC
)
assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
ModelEnvironmentVariableDisposition.PROMOTE: 33,
ModelEnvironmentVariableDisposition.PROMOTE: 34,
ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6,
ModelEnvironmentVariableDisposition.EXTERNAL: 0,
ModelEnvironmentVariableDisposition.INTERNALIZE: 11,
+6 -1
View File
@@ -45,6 +45,7 @@ from vllm_omni.config.stage_config import (
PipelineConfig,
StageDeployConfig,
StageExecutionType,
StagePipelineConfig,
load_deploy_config,
merge_pipeline_deploy,
)
@@ -1039,7 +1040,10 @@ def test_from_pipeline_config_uses_hf_config_for_callable_resolver():
def test_from_pipeline_config_accepts_pre_resolved_pipeline():
resolved_pipeline = PipelineConfig(model_type="callable_resolved_variant")
resolved_pipeline = PipelineConfig(
model_type="callable_resolved_variant",
stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),),
)
omni_config = VllmOmniConfig.from_pipeline_config(resolved_pipeline)
@@ -1072,6 +1076,7 @@ def test_from_pipeline_config_default_deploy_name_ignores_cwd(monkeypatch, tmp_p
pipeline = PipelineConfig(
model_type="pipeline_with_default",
default_deploy_config_name=default_name,
stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),),
)
loaded_paths = []
+11 -22
View File
@@ -1,16 +1,22 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Tests for out of tree registration to OMNI_PIPELINES."""
import pytest
from transformers import PretrainedConfig
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline
from vllm_omni.config.stage_config import PipelineConfig, pipeline_cfg_resolver
from vllm_omni.config.stage_config import PipelineConfig, StagePipelineConfig, pipeline_cfg_resolver
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def build_fake_pipeline_config(model_type: str) -> PipelineConfig:
return PipelineConfig(
model_type=model_type, stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),)
)
@pytest.fixture
def custom_resolver():
"""Build a reusable custom resolver for PipelineConfigs."""
@@ -22,7 +28,7 @@ def custom_resolver():
def custom_resolver(
hf_config: CustomConfigType,
) -> PipelineConfig:
return PipelineConfig(model_type="resolved_type")
return build_fake_pipeline_config("resolved_type")
return custom_resolver
@@ -30,7 +36,7 @@ def custom_resolver():
def test_register_pipeline_config(clean_pipeline_registry):
"""Ensure that we can register a custom pipeline config to OMNI_PIPELINES."""
new_model_type = "new_model_type"
pipe_cfg = PipelineConfig(model_type=new_model_type)
pipe_cfg = build_fake_pipeline_config(new_model_type)
assert new_model_type not in OMNI_PIPELINES
register_pipeline(pipe_cfg)
assert new_model_type in OMNI_PIPELINES
@@ -41,7 +47,7 @@ def test_register_pipeline_config_with_model_type(clean_pipeline_registry):
"""Ensure that we can register a custom pipeline config with an explicit model_type to OMNI_PIPELINES."""
new_model_type = "new_model_type"
unused_model_type = "foo"
pipe_cfg = PipelineConfig(model_type=unused_model_type)
pipe_cfg = build_fake_pipeline_config(unused_model_type)
assert new_model_type not in OMNI_PIPELINES
assert unused_model_type not in OMNI_PIPELINES
@@ -73,20 +79,3 @@ def test_minimax_h3_disaggregation_is_explicit_opt_in():
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
assert isinstance(pipeline, PipelineConfig)
assert pipeline.model_type == "minimax_h3_disaggregated"
def test_in_tree_pipelines_pass_topology_validation():
"""Every statically registered in-tree pipeline must have a valid topology.
``PipelineConfig.validate()`` only runs on ``register_pipeline``, which
in-tree pipelines never go through, so a shipped topology error would
otherwise stay invisible until serving hangs or misroutes.
Resolver entries are skipped: they need an ``hf_config`` to produce a
``PipelineConfig``.
"""
failures = {
model_type: errors
for model_type, pipeline in OMNI_PIPELINES.items()
if isinstance(pipeline, PipelineConfig) and (errors := pipeline.validate())
}
assert not failures, f"in-tree pipelines failed topology validation: {failures}"
@@ -7,7 +7,7 @@ End-to-end test for diffusion batching via AsyncOmni.
This test fires multiple concurrent ``AsyncOmni.generate()`` calls for a
diffusion model and validates that every caller receives its correct
individual result. When the underlying diffusion stage is configured with
``batch_size > 1`` (via stage config or ``StageDiffusionClient``), the
``max_num_seqs > 1``, the
requests will be batched internally.
Even without explicit batching config this test is useful for verifying
@@ -260,7 +260,7 @@ async def compare_single_vs_parallel(
) -> None:
"""Run the same prompts sequentially then in parallel and print a comparison."""
omni = AsyncOmni(model=model, max_num_seqs=batch_size, diffusion_batch_size=batch_size)
omni = AsyncOmni(model=model, max_num_seqs=batch_size)
try:
await warmup(omni, WARMUP_PROMPTS)
single_time = await run_single(omni, prompts)
@@ -288,7 +288,7 @@ async def main(model: str, num_prompts: int, mode: str, batch_size: int = 1) ->
await compare_single_vs_parallel(model, prompts, batch_size=batch_size)
return
omni = AsyncOmni(model=model, max_num_seqs=batch_size, diffusion_batch_size=batch_size)
omni = AsyncOmni(model=model, max_num_seqs=batch_size)
try:
await warmup(omni, WARMUP_PROMPTS)
@@ -389,7 +389,7 @@ def test_diffusion_batching_async_concurrent(model_name: str):
fired via asyncio.gather. Each request_id must appear in the results."""
async def _inner():
omni = AsyncOmni(model=model_name, diffusion_batch_size=1)
omni = AsyncOmni(model=model_name, max_num_seqs=1)
try:
prompts = TEST_PROMPTS[:4]
sp = _default_sampling_params()
@@ -430,7 +430,7 @@ def test_diffusion_batching_list_prompt_rejected(model_name: str):
"""
async def _inner():
omni = AsyncOmni(model=model_name, diffusion_batch_size=4)
omni = AsyncOmni(model=model_name, max_num_seqs=4)
try:
prompts = TEST_PROMPTS[:4]
sp = _default_sampling_params()
@@ -528,7 +528,7 @@ if __name__ == "__main__":
"--batch-size",
type=int,
default=1,
help="Scheduler in-flight width (--max-num-seqs) and client diffusion_batch_size",
help="Scheduler in-flight width (--max-num-seqs)",
)
parser.add_argument(
"--mode",
+55 -1
View File
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from collections import defaultdict
import pytest
@@ -106,7 +109,7 @@ def make_dummy_lora_state_dict(
"""
key_a = f"{base_key}.{lora_a_suffix}"
key_b = f"{base_key}.{lora_b_suffix}"
key_bias = f"{base_key}.{lora_bias_suffix}" if bias else None
key_bias = f"{base_key}.{lora_bias_suffix}"
if weight_a is None:
weight_a = torch.randn(RANK, HEAD_DIM)
@@ -436,3 +439,54 @@ def test_wan_load_lora_weights_targets_both_transformers(mocker: MockerFixture):
assert_params_equal(pipeline.transformer, original_high)
assert_params_equal(pipeline.transformer_2, original_low)
assert "adapter0" not in pipeline.lora_loaded
def test_prepare_lora_delta_disambiguates_prefix_overlapping_stacked_names():
"""A fused name can be a prefix of another fused name.
SenseNova-U1 carries both `.qkv_proj` (understanding) and `.qkv_proj_mot_gen`
(generation). Selecting stacked mappings by substring matched both rules for
the generation weight and concatenated six shards into a delta twice as tall
as the parameter.
"""
param_to_weight_names = {
".qkv_proj_mot_gen": [".q_proj_mot_gen", ".k_proj_mot_gen", ".v_proj_mot_gen"],
".qkv_proj": [".q_proj", ".k_proj", ".v_proj"],
}
state_dict = {}
for part in ("q", "k", "v"):
for tower in ("_mot_gen", ""):
base = f"layers.0.self_attn.{part}_proj{tower}"
state_dict[f"{base}.lora_A.weight"] = torch.ones(RANK, HEAD_DIM)
state_dict[f"{base}.lora_B.weight"] = torch.ones(HEAD_DIM, RANK)
for base_key in ("layers.0.self_attn.qkv_proj_mot_gen", "layers.0.self_attn.qkv_proj"):
delta, _ = _prepare_lora_delta(state_dict, base_key, param_to_weight_names)
assert delta.shape == (3 * HEAD_DIM, HEAD_DIM), base_key
@pytest.mark.parametrize(
"fused, shards",
[
# qwen_image_transformer.py stacked_params_mapping
(".to_qkv", (".to_q", ".to_k", ".to_v")),
(".add_kv_proj", (".add_q_proj", ".add_k_proj", ".add_v_proj")),
# wan2_2_transformer.py stacked_params_mapping
(".attn1.to_qkv", (".attn1.to_q", ".attn1.to_k", ".attn1.to_v")),
# wan2_2_s2v_transformer.py stacked_params_mapping
(".self_attn.to_qkv", (".self_attn.q", ".self_attn.k", ".self_attn.v")),
],
)
def test_prepare_lora_delta_unchanged_for_existing_fused_layouts(fused, shards):
"""None of the fused names shipped today overlap by prefix, so matching on the
tail must select exactly the shards substring matching selected."""
param_to_weight_names = {fused: list(shards)}
base_key = f"blocks.0{fused}"
state_dict = {}
for shard in shards:
state_dict[f"blocks.0{shard}.lora_A.weight"] = torch.ones(RANK, HEAD_DIM)
state_dict[f"blocks.0{shard}.lora_B.weight"] = torch.ones(HEAD_DIM, RANK)
delta, used_keys = _prepare_lora_delta(state_dict, base_key, param_to_weight_names)
assert delta.shape == (len(shards) * HEAD_DIM, HEAD_DIM)
assert len(used_keys) == 2 * len(shards)
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""L1 unit tests for the native Boogu-Image pipeline.
@@ -133,6 +133,51 @@ def test_constructor_strips_mllm_lm_head(boogu_pipeline, mock_dependencies):
assert boogu_pipeline.mllm is mock_dependencies["inner_encoder"]
def test_constructor_forwards_revision_to_all_component_loaders(mock_dependencies, mocker):
from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
BooguImagePipeline,
)
revision = "hotfix-1k-20260708"
prefetch = mocker.patch(f"{_MODULE}.prefetch_subfolders")
scheduler_loader = mocker.patch(
f"{_MODULE}.FlowMatchEulerDiscreteScheduler.from_pretrained",
return_value=mock_dependencies["scheduler"],
)
mllm_loader = mocker.patch(
f"{_MODULE}.Qwen3VLForConditionalGeneration.from_pretrained",
return_value=mock_dependencies["mllm_wrapper"],
)
processor_loader = mocker.patch(
f"{_MODULE}.Qwen3VLProcessor.from_pretrained",
return_value=mock_dependencies["processor"],
)
vae_loader = mocker.patch(
f"{_MODULE}.AutoencoderKL.from_pretrained",
return_value=mock_dependencies["vae"],
)
od_config = OmniDiffusionConfig(
model="dummy-boogu",
revision=revision,
tf_model_config=TransformerConfig(params={}),
dtype=torch.float32,
num_gpus=1,
)
pipeline = BooguImagePipeline(od_config=od_config)
(source,) = pipeline.weights_sources
assert source.revision == revision
prefetch.assert_called_once_with(
"dummy-boogu",
["scheduler", "vae", "mllm", "processor"],
local_files_only=True,
revision=revision,
)
for loader in (scheduler_loader, mllm_loader, processor_loader, vae_loader):
assert loader.call_args.kwargs["revision"] == revision
def test_constructor_weights_sources(boogu_pipeline):
(source,) = boogu_pipeline.weights_sources
assert source.model_or_path == "dummy-boogu"
File diff suppressed because it is too large. Load diff
@@ -0,0 +1,195 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""CPU tests for SenseNova-U1 attention head handling.
The model used to expand K/V from 8 heads to 32 with ``repeat_interleave``
before calling the attention backend, so the backend's own GQA dispatch never
saw a GQA shape, and it passed an all-zeros mask during single-token decode,
which only served to keep SDPA off its fused kernels. These pin both contracts:
K/V must reach the backend uncompressed, decode must carry no mask, and the
compressed path must compute the same thing the expanded one did.
"""
import pytest
import torch
from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import (
SenseNovaU1Attention,
SenseNovaU1Model,
)
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
N_HEADS, N_KV_HEADS, HEAD_DIM, SEQ = 32, 8, 8, 6
GROUPS = N_HEADS // N_KV_HEADS
class _Recorder:
"""Stands in for the Attention module and records what it was handed."""
def __init__(self):
self.calls: list[tuple[torch.Size, torch.Size, torch.Tensor | None]] = []
def __call__(self, query, key, value, attn_metadata):
mask = None if attn_metadata is None else attn_metadata.attn_mask
self.calls.append((query.shape, key.shape, mask))
return query
class _AttnHost:
"""Carries only the attributes ``_run_attn``/``_run_attn_bshd`` touch."""
num_kv_groups = GROUPS
_align_mask_dtype = staticmethod(SenseNovaU1Attention._align_mask_dtype)
_run_attn = SenseNovaU1Attention._run_attn
_run_attn_bshd = SenseNovaU1Attention._run_attn_bshd
def __init__(self):
self.attn = _Recorder()
def test_bhsd_path_hands_the_backend_compressed_kv():
host = _AttnHost()
q = torch.randn(1, N_HEADS, SEQ, HEAD_DIM)
k = torch.randn(1, N_KV_HEADS, SEQ, HEAD_DIM)
host._run_attn(q, k, k.clone(), None)
_, key_shape, _ = host.attn.calls[-1]
# [B, S, H, D] after the transpose, so heads sit on dim 2.
assert key_shape[2] == N_KV_HEADS, f"K reached the backend with {key_shape[2]} heads"
def test_bshd_path_hands_the_backend_compressed_kv():
host = _AttnHost()
q = torch.randn(1, SEQ, N_HEADS, HEAD_DIM)
k = torch.randn(1, SEQ, N_KV_HEADS, HEAD_DIM)
host._run_attn_bshd(q, k, k.clone(), None)
_, key_shape, _ = host.attn.calls[-1]
assert key_shape[2] == N_KV_HEADS, f"K reached the backend with {key_shape[2]} heads"
def test_compressed_kv_computes_what_expansion_computed():
"""The backend must return the same tensor either way, or this is a bugfix
dressed up as a perf change."""
torch.manual_seed(0)
q = torch.randn(1, SEQ, N_HEADS, HEAD_DIM)
k = torch.randn(1, SEQ, N_KV_HEADS, HEAD_DIM)
v = torch.randn(1, SEQ, N_KV_HEADS, HEAD_DIM)
impl = SDPAImpl(
num_heads=N_HEADS,
head_size=HEAD_DIM,
softmax_scale=HEAD_DIM**-0.5,
num_kv_heads=N_KV_HEADS,
)
compressed = impl._forward_impl(q, k, v)
expanded = impl._forward_impl(q, k.repeat_interleave(GROUPS, dim=2), v.repeat_interleave(GROUPS, dim=2))
torch.testing.assert_close(compressed, expanded, atol=1e-6, rtol=1e-6)
class _MaskProbe(torch.nn.Module):
def __init__(self):
super().__init__()
self.masks: list[torch.Tensor | None] = []
def forward(self, hidden_states, **kwargs):
self.masks.append(kwargs["attention_mask"]["full_attention"])
return hidden_states
def _fake_rope(hidden_states, position_ids):
"""The forward builds the three RoPE tables before the layer loop; the probe
ignores them, so any 2-tuple will do."""
return hidden_states, hidden_states
class _ModelHost:
"""Carries only the attributes ``SenseNovaU1Model.forward`` touches."""
forward = SenseNovaU1Model.forward
def __init__(self, probe):
self.layers = [probe]
self.norm = torch.nn.Identity()
self.norm_mot_gen = torch.nn.Identity()
self.rotary_emb = _fake_rope
self.rotary_emb_hw = _fake_rope
def _run_model(seq_len: int):
probe = _MaskProbe()
_ModelHost(probe).forward(
inputs_embeds=torch.zeros(1, seq_len, 4),
indexes=torch.zeros(3, seq_len, dtype=torch.long),
)
return probe.masks[-1]
def test_single_token_decode_carries_no_mask():
assert _run_model(1) is None, "decode still builds the all-zeros mask"
def test_prefill_still_gets_a_causal_mask():
mask = _run_model(3)
assert mask is not None and mask.shape == (1, 1, 3, 3)
assert mask[0, 0, 0, 0] == 0.0
assert torch.isinf(mask[0, 0, 0, 1]) and mask[0, 0, 0, 1] < 0
assert mask[0, 0, 2, 1] == 0.0
# ---------------------------------------------------------------------------
# The decode path changes which SDPA kernel runs, so it is not bit-identical to
# what it replaced. "Different" would be a regression if the new result were
# worse, so pin the direction against an independent float64 reference. This
# needs a fused GQA kernel, which only exists on CUDA.
# ---------------------------------------------------------------------------
def _float64_attention(q, k, v):
q64 = q.double()
k64 = k.double().repeat_interleave(GROUPS, dim=1)
v64 = v.double().repeat_interleave(GROUPS, dim=1)
scores = (q64 @ k64.transpose(-1, -2)) * (q.shape[-1] ** -0.5)
return torch.softmax(scores, dim=-1) @ v64
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for fused GQA")
@pytest.mark.cuda
@pytest.mark.L4
@pytest.mark.parametrize("kv_len", [271, 512, 2048])
def test_decode_without_the_no_op_mask_is_closer_to_float64(kv_len):
"""The old decode path expanded K/V and passed an all-zeros mask, which kept
SDPA off its fused kernels. Dropping the mask is only defensible if the
kernel it unlocks is at least as accurate."""
dev, dtype = "cuda", torch.bfloat16
heads, dim = 32, 128
new_errs: list[float] = []
old_errs: list[float] = []
trials = 8
for t in range(trials):
torch.manual_seed(2000 + t)
q = torch.randn(1, heads, 1, dim, device=dev, dtype=dtype)
k = torch.randn(1, heads // GROUPS, kv_len, dim, device=dev, dtype=dtype)
v = torch.randn(1, heads // GROUPS, kv_len, dim, device=dev, dtype=dtype)
ref = _float64_attention(q, k, v)
zeros_mask = torch.zeros(1, 1, 1, kv_len, device=dev, dtype=dtype)
old = torch.nn.functional.scaled_dot_product_attention(
q,
k.repeat_interleave(GROUPS, dim=1),
v.repeat_interleave(GROUPS, dim=1),
attn_mask=zeros_mask,
)
new = torch.nn.functional.scaled_dot_product_attention(q, k, v, enable_gqa=True)
new_errs.append((new.double() - ref).abs().mean().item())
old_errs.append((old.double() - ref).abs().mean().item())
# A single trial can tie either way depending on the GPU and the CUDA build,
# so gate on the aggregate: the maskless path must not be worse on average.
new_mean = sum(new_errs) / trials
old_mean = sum(old_errs) / trials
wins = sum(n <= o for n, o in zip(new_errs, old_errs))
assert new_mean <= old_mean, (
f"kv_len={kv_len}: maskless mean error {new_mean:.4e} is worse than masked {old_mean:.4e} "
f"({wins}/{trials} per-trial wins)"
)
@@ -0,0 +1,196 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""CPU tests for SenseNova-U1 distilled-LoRA fusion.
The checkpoints ship unfused ``q_proj_mot_gen`` / ``gate_proj`` keys while the
pipeline runs fused layers, so each delta lands in its own row slice.
"""
import pytest
import torch
import torch.nn as nn
from safetensors.torch import save_file
from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import SenseNovaU1Pipeline
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
HIDDEN, KV, INTER, RANK, ALPHA = 8, 4, 16, 2, 4.0
class _Layer(nn.Module):
def __init__(self):
super().__init__()
self.self_attn = nn.Module()
self.self_attn.qkv_proj_mot_gen = nn.Linear(HIDDEN, HIDDEN + 2 * KV, bias=False)
self.self_attn.o_proj_mot_gen = nn.Linear(HIDDEN, HIDDEN, bias=False)
self.mlp_mot_gen = nn.Module()
self.mlp_mot_gen.gate_up_proj = nn.Linear(HIDDEN, 2 * INTER, bias=False)
self.mlp_mot_gen.down_proj = nn.Linear(INTER, HIDDEN, bias=False)
class _Stub(SenseNovaU1Pipeline):
"""Only the parameter tree matters here, so skip the real __init__."""
def __init__(self):
nn.Module.__init__(self)
self.language_model = nn.Module()
self.language_model.model = nn.Module()
self.language_model.model.layers = nn.ModuleList([_Layer()])
def _triple(out: int, in_: int, fill: float) -> dict[str, torch.Tensor]:
return {
"lora_down.weight": torch.full((RANK, in_), fill, dtype=torch.float32),
"lora_up.weight": torch.full((out, RANK), 1.0, dtype=torch.float32),
"alpha": torch.tensor(ALPHA),
}
def _write(tmp_path, entries: dict[str, dict[str, torch.Tensor]]):
flat = {f"{mod}.{leaf}": t for mod, parts in entries.items() for leaf, t in parts.items()}
path = tmp_path / "lora.safetensors"
save_file(flat, str(path))
return str(path)
def _prefix(i: int = 0) -> str:
return f"language_model.model.layers.{i}"
def test_qkv_shards_land_in_their_own_row_slices(tmp_path):
pipe = _Stub()
qkv = pipe.language_model.model.layers[0].self_attn.qkv_proj_mot_gen.weight
with torch.no_grad():
qkv.zero_()
entries = {
f"{_prefix()}.self_attn.q_proj_mot_gen": _triple(HIDDEN, HIDDEN, 1.0),
f"{_prefix()}.self_attn.k_proj_mot_gen": _triple(KV, HIDDEN, 2.0),
f"{_prefix()}.self_attn.v_proj_mot_gen": _triple(KV, HIDDEN, 3.0),
}
pipe.load_lora_weights(_write(tmp_path, entries))
scale = ALPHA / RANK
assert torch.allclose(qkv[:HIDDEN], torch.full((HIDDEN, HIDDEN), scale * RANK * 1.0))
assert torch.allclose(qkv[HIDDEN : HIDDEN + KV], torch.full((KV, HIDDEN), scale * RANK * 2.0))
assert torch.allclose(qkv[HIDDEN + KV :], torch.full((KV, HIDDEN), scale * RANK * 3.0))
def test_gate_and_up_land_in_their_own_halves(tmp_path):
pipe = _Stub()
gu = pipe.language_model.model.layers[0].mlp_mot_gen.gate_up_proj.weight
with torch.no_grad():
gu.zero_()
entries = {
f"{_prefix()}.mlp_mot_gen.gate_proj": _triple(INTER, HIDDEN, 1.0),
f"{_prefix()}.mlp_mot_gen.up_proj": _triple(INTER, HIDDEN, 5.0),
}
pipe.load_lora_weights(_write(tmp_path, entries))
scale = ALPHA / RANK
assert torch.allclose(gu[:INTER], torch.full((INTER, HIDDEN), scale * RANK * 1.0))
assert torch.allclose(gu[INTER:], torch.full((INTER, HIDDEN), scale * RANK * 5.0))
def test_unfused_targets_are_added_whole(tmp_path):
pipe = _Stub()
o = pipe.language_model.model.layers[0].self_attn.o_proj_mot_gen.weight
before = o.detach().clone()
entries = {f"{_prefix()}.self_attn.o_proj_mot_gen": _triple(HIDDEN, HIDDEN, 1.0)}
pipe.load_lora_weights(_write(tmp_path, entries))
assert torch.allclose(o - before, torch.full((HIDDEN, HIDDEN), (ALPHA / RANK) * RANK))
def test_lora_that_matches_nothing_raises_instead_of_no_op(tmp_path):
"""Without a loader the worker only warned and generated without the
adapter, so a wrong path produced a plausible image and no error."""
pipe = _Stub()
entries = {f"{_prefix()}.self_attn.does_not_exist": _triple(HIDDEN, HIDDEN, 1.0)}
with pytest.raises(ValueError, match="matched no parameter"):
pipe.load_lora_weights(_write(tmp_path, entries))
def test_shape_mismatch_raises(tmp_path):
pipe = _Stub()
entries = {f"{_prefix()}.self_attn.o_proj_mot_gen": _triple(HIDDEN + 1, HIDDEN, 1.0)}
with pytest.raises(RuntimeError, match="must match the size"):
pipe.load_lora_weights(_write(tmp_path, entries))
def test_multiple_files_rejected(tmp_path):
pipe = _Stub()
path = _write(tmp_path, {f"{_prefix()}.self_attn.o_proj_mot_gen": _triple(HIDDEN, HIDDEN, 1.0)})
with pytest.raises(ValueError, match="exactly one"):
pipe.load_lora_weights([path, path])
class _ShardedLayer(nn.Module):
"""One rank of a two-way sharded fused QKV, with a vLLM-style weight loader.
The parameter holds half of each of q, k and v. The loader is handed the whole
layer's tensor and keeps only this rank's slice, which is what vLLM does.
"""
def __init__(self):
super().__init__()
self.weight = nn.Parameter(torch.zeros(HIDDEN // 2 + KV, HIDDEN))
self.seen: list[tuple[str, tuple[int, ...]]] = []
self.weight.weight_loader = self._weight_loader
def _weight_loader(self, param, loaded_weight, shard_id):
self.seen.append((shard_id, tuple(loaded_weight.shape)))
sizes = {"q": HIDDEN // 2, "k": KV // 2, "v": KV // 2}
offsets = {"q": 0, "k": HIDDEN // 2, "v": HIDDEN // 2 + KV // 2}
rows = sizes[shard_id]
param.data[offsets[shard_id] : offsets[shard_id] + rows] = loaded_weight[:rows]
class _ShardedStub(SenseNovaU1Pipeline):
def __init__(self):
nn.Module.__init__(self)
self.language_model = nn.Module()
self.language_model.model = nn.Module()
layer = nn.Module()
layer.self_attn = nn.Module()
layer.self_attn.qkv_proj_mot_gen = _ShardedLayer()
self.language_model.model.layers = nn.ModuleList([layer])
def test_fused_delta_is_sharded_through_the_weight_loader(tmp_path):
"""Under tensor parallelism the parameter is a rank-local shard while the
checkpoint delta spans the whole layer, so the delta has to go through the
layer's weight loader. Adding it directly raises a shape error."""
pipe = _ShardedStub()
layer = pipe.language_model.model.layers[0].self_attn.qkv_proj_mot_gen
entries = {
f"{_prefix()}.self_attn.q_proj_mot_gen": _triple(HIDDEN, HIDDEN, 1.0),
f"{_prefix()}.self_attn.k_proj_mot_gen": _triple(KV, HIDDEN, 2.0),
f"{_prefix()}.self_attn.v_proj_mot_gen": _triple(KV, HIDDEN, 3.0),
}
pipe.load_lora_weights(_write(tmp_path, entries))
# the loader saw the whole layer's tensors, not a pre-sharded slice
assert layer.seen == [("q", (HIDDEN, HIDDEN)), ("k", (KV, HIDDEN)), ("v", (KV, HIDDEN))]
scale = ALPHA / RANK * RANK
got = layer.weight
assert torch.allclose(got[: HIDDEN // 2], torch.full((HIDDEN // 2, HIDDEN), scale * 1.0))
assert torch.allclose(got[HIDDEN // 2 : HIDDEN // 2 + KV // 2], torch.full((KV // 2, HIDDEN), scale * 2.0))
assert torch.allclose(got[HIDDEN // 2 + KV // 2 :], torch.full((KV // 2, HIDDEN), scale * 3.0))
def test_fusion_does_not_retain_the_adapter_tensors(tmp_path):
"""Fusion is one-way and this pipeline has no unload path, so holding the
fp32 state dict would keep about 1.5 GiB alive for the shipped checkpoint."""
pipe = _Stub()
path = _write(tmp_path, {f"{_prefix()}.self_attn.o_proj_mot_gen": _triple(HIDDEN, HIDDEN, 1.0)})
pipe.load_lora_weights(path, adapter_name="distill")
kept = pipe.lora_loaded["distill"]
assert not isinstance(kept, dict)
assert not isinstance(kept, torch.Tensor)
# the sentinel still makes a second load a no-op
before = pipe.language_model.model.layers[0].self_attn.o_proj_mot_gen.weight.detach().clone()
pipe.load_lora_weights(path, adapter_name="distill")
torch.testing.assert_close(pipe.language_model.model.layers[0].self_attn.o_proj_mot_gen.weight, before)
@@ -0,0 +1,800 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Tests for the paged decode cache and its CUDA graph.
The cache exists so a decode step has static shapes and can be captured. That
puts the burden on two things a normal test would not look at: the prefix must
survive the hand-off from the ordinary cache, and every quantity that varies
between steps must live in a device tensor. A Python-level index is correct
eagerly and silently wrong once captured, because capture bakes it in -- which
is exactly what the last test here pins.
"""
from types import SimpleNamespace
import pytest
import torch
from vllm.lora.layers import BaseLayerWithLoRA
import vllm_omni.diffusion.models.sensenova_u1.paged_decode as paged_decode
from vllm_omni.diffusion.models.sensenova_u1.paged_decode import (
BLOCK_SIZE,
BUCKETS,
TAIL_STEP,
PagedDecodeCache,
_bucket_for,
)
from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import (
SenseNovaU1Attention,
)
from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
LAYERS, KV_HEADS, HEAD_DIM = 2, 4, 8
N_HEADS = KV_HEADS * 2
class _Layer:
def __init__(self, keys=None, values=None):
self.keys, self.values = keys, values
class _Cache:
"""Stands in for the transformers cache: layers with [B, H, S, D] tensors."""
def __init__(self, layers):
self.layers = layers
def _dyn_cache(prefix, device=None):
torch.manual_seed(0)
return _Cache(
[
_Layer(
torch.randn(1, KV_HEADS, prefix, HEAD_DIM, device=device),
torch.randn(1, KV_HEADS, prefix, HEAD_DIM, device=device),
)
for _ in range(LAYERS)
]
)
def test_bucket_rounds_up_and_is_block_aligned():
assert _bucket_for(1) == 512
assert _bucket_for(512) == 512
assert _bucket_for(513) == 1024
assert _bucket_for(100_000) % BLOCK_SIZE == 0
def test_past_the_last_bucket_the_schedule_grows_in_steps():
"""Every bucket change reallocates the cache and re-captures the graph, so
the tail of the schedule must not step by BLOCK_SIZE. A think edit with two
2048x2048 inputs runs to about 9,100 tokens, so this regime is reachable,
and a request decodes at most 1024 steps.
"""
last = BUCKETS[-1]
assert _bucket_for(last + 1) == last + TAIL_STEP
assert _bucket_for(100_000) % BLOCK_SIZE == 0
for start in (last + 1, last + TAIL_STEP - 1, last + TAIL_STEP + 7):
bucket, grows = _bucket_for(start), 0
for n in range(start, start + 1025):
if _bucket_for(n) != bucket:
bucket = _bucket_for(n)
grows += 1
assert grows <= 1, f"start {start}: {grows} reallocations over 1024 decode steps"
def test_prefix_survives_the_hand_off():
dyn = _dyn_cache(37)
paged = PagedDecodeCache.from_dynamic_cache(dyn, LAYERS, torch.device("cpu"), torch.float32)
assert paged.length == 37
for i in range(LAYERS):
stored = paged.k[i].view(-1, KV_HEADS, HEAD_DIM)[:37]
torch.testing.assert_close(stored, dyn.layers[i].keys[0].transpose(0, 1))
def test_write_back_restores_what_was_handed_over():
"""Decode is paged but the stage after it reads the ordinary cache."""
dyn = _dyn_cache(21)
original = [layer.keys.clone() for layer in dyn.layers]
paged = PagedDecodeCache.from_dynamic_cache(dyn, LAYERS, torch.device("cpu"), torch.float32)
paged.to_dynamic_cache(dyn)
for before, layer in zip(original, dyn.layers):
torch.testing.assert_close(layer.keys, before)
def test_growing_keeps_the_tokens_already_stored():
dyn = _dyn_cache(500)
paged = PagedDecodeCache.from_dynamic_cache(dyn, LAYERS, torch.device("cpu"), torch.float32)
assert paged.bucket == 512
before = [k.view(-1, KV_HEADS, HEAD_DIM)[:500].clone() for k in paged.k]
gen = paged.generation
assert paged.grow(513) is True
assert paged.bucket == 1024
assert paged.generation > gen, "a reallocation must invalidate captured graphs"
for old, new in zip(before, paged.k):
torch.testing.assert_close(new.view(-1, KV_HEADS, HEAD_DIM)[:500], old)
def test_a_failed_grow_leaves_the_cache_on_its_old_buffers(monkeypatch):
"""A half-applied resize would strand a capture on freed storage.
``generation`` is half of the key a captured graph is stored under, so
publishing new K without publishing V, the block table, the bucket and the
generation leaves the old key selecting a graph whose recorded addresses
have just been freed. A later request replays it. Injecting a failure into
the last allocation is the cheapest way to reach that window.
"""
dyn = _dyn_cache(500)
paged = PagedDecodeCache.from_dynamic_cache(dyn, LAYERS, torch.device("cpu"), torch.float32)
before_k, before_v = list(paged.k), list(paged.v)
before_table, before_bucket = paged.block_table, paged.bucket
before_generation, before_length = paged.generation, paged.length
real_zeros = torch.zeros
calls = []
def flaky(*args, **kwargs):
calls.append(1)
if len(calls) == 2 * LAYERS: # the last of the two allocations per layer
raise RuntimeError("injected allocation failure")
return real_zeros(*args, **kwargs)
monkeypatch.setattr(torch, "zeros", flaky)
with pytest.raises(RuntimeError, match="injected"):
paged.grow(513)
monkeypatch.undo()
assert [id(t) for t in paged.k] == [id(t) for t in before_k], "K was published on its own"
assert [id(t) for t in paged.v] == [id(t) for t in before_v], "V was published on its own"
assert paged.block_table is before_table
assert paged.bucket == before_bucket
assert paged.generation == before_generation, "the graph key moved without new buffers"
assert paged.length == before_length
def test_length_lives_in_device_tensors():
"""Both the attended length and the write slot must be tensors.
A captured graph reads them at replay; a Python int would be frozen at
capture time and every replay would attend the wrong span and overwrite the
same slot.
"""
dyn = _dyn_cache(10)
paged = PagedDecodeCache.from_dynamic_cache(dyn, LAYERS, torch.device("cpu"), torch.float32)
paged.set_length(15)
assert torch.is_tensor(paged.seqused) and int(paged.seqused) == 15
assert torch.is_tensor(paged.pos) and int(paged.pos) == 14
# ---------------------------------------------------------------------------
# The capture contract. These need CUDA and the bundled flash-attention kernel.
# ---------------------------------------------------------------------------
cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@cuda_only
@pytest.mark.cuda
@pytest.mark.L4
def test_paged_attention_matches_sdpa_over_the_same_buffer():
import torch.nn.functional as F
from vllm_omni.diffusion.models.sensenova_u1.paged_decode import PagedDecodeCache
heads, kv_heads, dim, n = 8, 2, 64, 100
dev, dt = torch.device("cuda"), torch.bfloat16
torch.manual_seed(0)
dyn = _Cache(
[
_Layer(
torch.randn(1, kv_heads, n, dim, device=dev, dtype=dt),
torch.randn(1, kv_heads, n, dim, device=dev, dtype=dt),
)
]
)
paged = PagedDecodeCache.from_dynamic_cache(dyn, 1, dev, dt)
paged.set_length(n)
q = torch.randn(1, heads, 1, dim, device=dev, dtype=dt)
k = torch.zeros(1, kv_heads, 1, dim, device=dev, dtype=dt)
scale = dim**-0.5
# write the last slot back as itself so attend() is a pure read
k[0, :, 0] = dyn.layers[0].keys[0, :, n - 1]
v = torch.zeros_like(k)
v[0, :, 0] = dyn.layers[0].values[0, :, n - 1]
got = paged.attend(0, q, k, v, scale)
flat_k = paged.k[0].view(-1, kv_heads, dim)[:n]
flat_v = paged.v[0].view(-1, kv_heads, dim)[:n]
want = F.scaled_dot_product_attention(
q,
flat_k.unsqueeze(0).transpose(1, 2).contiguous(),
flat_v.unsqueeze(0).transpose(1, 2).contiguous(),
enable_gqa=True,
scale=scale,
).transpose(1, 2)
torch.testing.assert_close(got.float(), want.float(), atol=2e-3, rtol=2e-3)
@cuda_only
@pytest.mark.cuda
@pytest.mark.L4
def test_a_captured_graph_follows_a_later_length():
"""One capture has to serve every length in the bucket.
If the attended span were a Python value the graph would keep attending the
capture-time length, and the run would still look fast.
"""
import torch.nn.functional as F
from vllm_omni.diffusion.models.sensenova_u1.paged_decode import PagedDecodeCache
heads, kv_heads, dim, n = 8, 2, 64, 200
dev, dt = torch.device("cuda"), torch.bfloat16
torch.manual_seed(0)
dyn = _Cache(
[
_Layer(
torch.randn(1, kv_heads, n, dim, device=dev, dtype=dt),
torch.randn(1, kv_heads, n, dim, device=dev, dtype=dt),
)
]
)
paged = PagedDecodeCache.from_dynamic_cache(dyn, 1, dev, dt)
q = torch.randn(1, heads, 1, dim, device=dev, dtype=dt)
k = torch.zeros(1, kv_heads, 1, dim, device=dev, dtype=dt)
v = torch.zeros_like(k)
scale = dim**-0.5
paged.set_length(n)
k[0, :, 0] = dyn.layers[0].keys[0, :, n - 1]
v[0, :, 0] = dyn.layers[0].values[0, :, n - 1]
side = torch.cuda.Stream()
side.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(side):
for _ in range(3):
paged.attend(0, q, k, v, scale)
torch.cuda.current_stream().wait_stream(side)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
out = paged.attend(0, q, k, v, scale)
shorter = n - 43
paged.set_length(shorter)
graph.replay()
torch.accelerator.synchronize()
flat_k = paged.k[0].view(-1, kv_heads, dim)[:shorter]
flat_v = paged.v[0].view(-1, kv_heads, dim)[:shorter]
want = F.scaled_dot_product_attention(
q,
flat_k.unsqueeze(0).transpose(1, 2).contiguous(),
flat_v.unsqueeze(0).transpose(1, 2).contiguous(),
enable_gqa=True,
scale=scale,
).transpose(1, 2)
torch.testing.assert_close(out.float(), want.float(), atol=2e-3, rtol=2e-3)
@cuda_only
@pytest.mark.cuda
@pytest.mark.L4
def test_a_captured_graph_writes_each_step_to_its_own_slot():
"""The write slot has to be a device tensor, not a Python index.
``test_a_captured_graph_follows_a_later_length`` pins the attended span.
The slot written is the other half: a Python index is frozen at capture, so
every replay would overwrite the capture-time slot and the cache would stop
growing -- correct eagerly, silently wrong once captured.
"""
from vllm_omni.diffusion.models.sensenova_u1.paged_decode import PagedDecodeCache
heads, kv_heads, dim, n = 8, 2, 64, 100
dev, dt = torch.device("cuda"), torch.bfloat16
torch.manual_seed(0)
dyn = _Cache(
[
_Layer(
torch.randn(1, kv_heads, n, dim, device=dev, dtype=dt),
torch.randn(1, kv_heads, n, dim, device=dev, dtype=dt),
)
]
)
paged = PagedDecodeCache.from_dynamic_cache(dyn, 1, dev, dt)
q = torch.randn(1, heads, 1, dim, device=dev, dtype=dt)
k = torch.zeros(1, kv_heads, 1, dim, device=dev, dtype=dt)
v = torch.zeros_like(k)
scale = dim**-0.5
# Capture on the step that writes slot n-1, restoring what is already there
# so the capture itself changes nothing.
paged.set_length(n)
k[0, :, 0] = dyn.layers[0].keys[0, :, n - 1]
v[0, :, 0] = dyn.layers[0].values[0, :, n - 1]
side = torch.cuda.Stream()
side.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(side):
for _ in range(3):
paged.attend(0, q, k, v, scale)
torch.cuda.current_stream().wait_stream(side)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
paged.attend(0, q, k, v, scale)
kept = paged.k[0].view(-1, kv_heads, dim)[n - 1].clone()
# The next decode step: one token further, different K/V.
paged.set_length(n + 1)
k[0, :, 0] = 1.5
v[0, :, 0] = -2.5
graph.replay()
current_omni_platform.synchronize()
flat_k = paged.k[0].view(-1, kv_heads, dim)
flat_v = paged.v[0].view(-1, kv_heads, dim)
torch.testing.assert_close(flat_k[n].float(), torch.full_like(flat_k[n], 1.5).float())
torch.testing.assert_close(flat_v[n].float(), torch.full_like(flat_v[n], -2.5).float())
torch.testing.assert_close(flat_k[n - 1].float(), kept.float())
@cuda_only
@pytest.mark.cuda
@pytest.mark.L4
def test_a_whole_decode_loop_matches_the_unpaged_path():
"""Per-step kernel equivalence is not loop equivalence.
With the switch off a single-token step attends over a cache grown by
``torch.cat`` and no mask; with it on the same step reads a bucket-sized
buffer through ``seqused_k``. Run the same sequence of steps through both,
crossing a bucket boundary on the way, and require the same argmax at every
step off a fixed projection -- which is the token id the pipeline would emit.
"""
import torch.nn.functional as F
from vllm_omni.diffusion.models.sensenova_u1.paged_decode import PagedDecodeCache
heads, kv_heads, dim = 8, 2, 64
prefill, steps = 500, 40 # crosses the 512 bucket at step 13
dev, dt = torch.device("cuda"), torch.bfloat16
scale = dim**-0.5
torch.manual_seed(7)
keys = torch.randn(1, kv_heads, prefill, dim, device=dev, dtype=dt)
values = torch.randn(1, kv_heads, prefill, dim, device=dev, dtype=dt)
paged = PagedDecodeCache.from_dynamic_cache(_Cache([_Layer(keys, values)]), 1, dev, dt)
ref_k, ref_v = keys.clone(), values.clone()
proj = torch.randn(dim * heads, 128, device=dev, dtype=torch.float32)
paged_ids, ref_ids, worst = [], [], 0.0
for step in range(steps):
q = torch.randn(1, heads, 1, dim, device=dev, dtype=dt)
k = torch.randn(1, kv_heads, 1, dim, device=dev, dtype=dt)
v = torch.randn(1, kv_heads, 1, dim, device=dev, dtype=dt)
if paged.length + 1 > paged.bucket:
paged.grow(paged.length + 1)
paged.set_length(paged.length + 1)
got = paged.attend(0, q, k, v, scale)
ref_k = torch.cat([ref_k, k], dim=2)
ref_v = torch.cat([ref_v, v], dim=2)
want = F.scaled_dot_product_attention(q, ref_k, ref_v, enable_gqa=True, scale=scale).transpose(1, 2)
worst = max(worst, (got.float() - want.float()).abs().max().item())
paged_ids.append(int((got.float().reshape(1, -1) @ proj).argmax()))
ref_ids.append(int((want.float().reshape(1, -1) @ proj).argmax()))
assert paged.bucket == 1024, "the loop never crossed a bucket boundary"
assert paged_ids == ref_ids, (
f"argmax diverges at step {next(i for i, (a, b) in enumerate(zip(paged_ids, ref_ids)) if a != b)}"
)
assert worst < 5e-2, f"largest per-step deviation {worst:.3e}"
class _StubLM(torch.nn.Module):
"""Just enough CUDA work inside the capture for a pool to be allocated."""
def __init__(self, device):
super().__init__()
self.w = torch.randn(16, 64, device=device)
def forward(self, input_ids, indexes, past_key_values=None, use_cache=False, paged_cache=None):
x = input_ids.float().expand(1, 16) @ self.w
return SimpleNamespace(logits=x.unsqueeze(0))
@cuda_only
@pytest.mark.cuda
@pytest.mark.L4
def test_the_decode_graph_captures_into_the_shared_platform_pool(monkeypatch):
"""``_decode_context`` builds a runner per request, so a private capture pool
hands every request its own arena and never gives it back. Measured on a
real model that was about 40 MiB of device memory per request, rising
linearly over twelve sequential think requests while the same twelve with
the switch off held flat.
"""
from vllm.platforms import current_platform
from vllm_omni.diffusion.models.sensenova_u1.paged_decode import DecodeGraphRunner
pools = []
real_graph = torch.cuda.graph
def spy(graph, pool=None, **kwargs):
pools.append(pool)
return real_graph(graph, pool=pool, **kwargs)
monkeypatch.setattr(torch.cuda, "graph", spy)
dev = torch.device("cuda")
cache = PagedDecodeCache.from_dynamic_cache(_dyn_cache(8, dev), LAYERS, dev, torch.float32)
runner = DecodeGraphRunner(_StubLM(dev), cache, dev)
runner.step(0, 0)
assert pools, "no graph was captured"
assert pools[0] is not None, "captured into a fresh private pool"
assert pools[0] == current_platform.get_global_graph_pool()
assert not hasattr(runner, "_pool"), "the runner still keeps a pool of its own"
# ---------------------------------------------------------------------------
# The cache is only worth anything if production actually reaches it, and the
# kwargs it needs are not in every wheel that exports the kernel. The tests
# above would all still pass with the `forward_und` branch deleted, so these
# drive the real dispatch instead of the cache object.
# ---------------------------------------------------------------------------
class _PagedProbe:
def __init__(self, out):
self.out = out
self.calls: list[tuple[int, torch.Size, float]] = []
def attend(self, layer_idx, query, key, value, scaling):
self.calls.append((layer_idx, query.shape, scaling))
return self.out
class _AttnHost:
"""Carries only what ``forward_und`` touches on the way to the branch."""
forward_und = SenseNovaU1Attention.forward_und
qkv_proj = q_norm = k_norm = q_norm_hw = k_norm_hw = None
layer_idx = 3
scaling = HEAD_DIM**-0.5
def __init__(self, seq, paged_out):
self._qkv = (
torch.randn(1, N_HEADS, seq, HEAD_DIM),
torch.randn(1, KV_HEADS, seq, HEAD_DIM),
torch.randn(1, KV_HEADS, seq, HEAD_DIM),
)
self.paged = _PagedProbe(paged_out)
self.sdpa_calls: list[torch.Tensor | None] = []
def _project_and_rope(self, *args, **kwargs):
return self._qkv
def _run_attn(self, query, key, value, mask):
self.sdpa_calls.append(mask)
return torch.zeros(1, N_HEADS, query.shape[2], HEAD_DIM)
def o_proj(self, hidden_states):
return hidden_states * 2, None
def test_forward_und_sends_a_single_token_decode_to_the_paged_cache():
"""Goes through the real dispatch, so deleting the branch turns this red."""
out = torch.randn(1, N_HEADS, 1, HEAD_DIM)
host = _AttnHost(seq=1, paged_out=out)
got = host.forward_und(torch.zeros(1, 1, 16), None, None, position_embeddings=None, paged_cache=host.paged)
assert len(host.paged.calls) == 1, "the paged hook was never reached"
assert host.sdpa_calls == [], "the same step also ran the SDPA path"
assert host.paged.calls[0][0] == host.layer_idx
torch.testing.assert_close(got, out.reshape(1, 1, -1).contiguous() * 2)
@pytest.mark.parametrize(
"seq,mask",
[(2, None), (1, torch.zeros(1, 1, 1, 1))],
ids=["multi_token", "masked"],
)
def test_only_an_unmasked_single_token_step_takes_the_paged_path(seq, mask):
"""Both other shapes are uncapturable, and the cache holds one row per step."""
host = _AttnHost(seq=seq, paged_out=torch.zeros(1, N_HEADS, seq, HEAD_DIM))
host.forward_und(torch.zeros(1, seq, 16), None, mask, position_embeddings=None, paged_cache=host.paged)
assert host.paged.calls == [], "a step that cannot be captured reached the paged cache"
assert len(host.sdpa_calls) == 1
def test_a_kernel_missing_any_kwarg_we_pass_is_not_supported(monkeypatch):
"""The call names every argument, so the probe has to cover every name.
An older wheel exports ``flash_attn_varlen_func`` without the paged kwargs,
and a future one may rename the standard ones. Either way the probe is what
keeps that install on SDPA instead of raising TypeError at the first decode
step, with the feature on by default.
"""
def no_paged_kwargs(q, k, v, max_seqlen_q, cu_seqlens_q, max_seqlen_k, cu_seqlens_k, causal=False):
raise AssertionError("must not be called") # pragma: no cover
def renamed_standard_args(q, k, v, max_q, cu_q, max_k, seqused_k=None, block_table=None, causal=False):
raise AssertionError("must not be called") # pragma: no cover
for kernel in (no_paged_kwargs, renamed_standard_args):
monkeypatch.setattr(paged_decode, "_flash_varlen", lambda kernel=kernel: kernel)
assert paged_decode.paged_decode_supported(torch.device("cuda"), HEAD_DIM) is False
def current_kernel(
q,
k,
v,
max_seqlen_q,
cu_seqlens_q,
max_seqlen_k,
seqused_k=None,
block_table=None,
softmax_scale=None,
causal=False,
):
raise AssertionError("must not be called") # pragma: no cover
monkeypatch.setattr(paged_decode, "_flash_varlen", lambda: current_kernel)
assert paged_decode.paged_decode_supported(torch.device("cuda"), HEAD_DIM) is True
def test_only_q_k_and_v_reach_the_kernel_positionally(monkeypatch):
"""A positional standard argument binds to whatever a future wheel puts in
that slot, and the probe above only checks that the names exist."""
seen = {}
def recorder(*args, **kwargs):
seen["args"], seen["kwargs"] = args, kwargs
return torch.zeros(1, N_HEADS, HEAD_DIM)
monkeypatch.setattr(paged_decode, "_flash_varlen", lambda: recorder)
cache = PagedDecodeCache(1, KV_HEADS, HEAD_DIM, 8, torch.device("cpu"), torch.float32)
cache.set_length(4)
cache.attend(
0,
torch.zeros(1, N_HEADS, 1, HEAD_DIM),
torch.zeros(1, KV_HEADS, 1, HEAD_DIM),
torch.zeros(1, KV_HEADS, 1, HEAD_DIM),
1.0,
)
assert len(seen["args"]) == 3, "only q, k and v may be positional"
assert set(seen["kwargs"]) == {
"max_seqlen_q",
"cu_seqlens_q",
"max_seqlen_k",
"seqused_k",
"block_table",
"softmax_scale",
"causal",
}
def test_the_kill_switch_returns_no_cache(monkeypatch):
"""The recipe documents `VLLM_OMNI_SENSENOVA_PAGED_DECODE=0` as the way back
to the ordinary cache, so the switch needs a test rather than only prose."""
from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import (
SenseNovaU1Pipeline,
)
monkeypatch.setenv("VLLM_OMNI_SENSENOVA_PAGED_DECODE", "0")
# The env check returns before anything else on the pipeline is touched.
host = object.__new__(SenseNovaU1Pipeline)
assert SenseNovaU1Pipeline._decode_context(host, object()) is None
def _pipeline_host(monkeypatch, device="cpu", language_model=None):
"""A bare pipeline carrying only what ``_decode_context`` touches."""
from vllm_omni.diffusion.models.sensenova_u1 import pipeline_sensenova_u1 as pipe_mod
monkeypatch.setenv("VLLM_OMNI_SENSENOVA_PAGED_DECODE", "1")
monkeypatch.setattr(pipe_mod, "paged_decode_supported", lambda dev, head_dim: True)
layer = SimpleNamespace(self_attn=SimpleNamespace(head_dim=HEAD_DIM))
host = object.__new__(pipe_mod.SenseNovaU1Pipeline)
# The pipeline is an nn.Module, and the language model assigned below is one
# too, so the bookkeeping has to exist before the assignment.
torch.nn.Module.__init__(host)
host.device = device
lm = language_model if language_model is not None else torch.nn.Module()
lm.model = SimpleNamespace(layers=[layer] * LAYERS)
host.language_model = lm
return pipe_mod, host
def test_the_decode_context_is_reused_across_requests(monkeypatch):
"""A capture binds the addresses it recorded, so a cache built per request
forces a capture per request. Two requests whose prefill fits the same
bucket have to come back with the same cache and the same runner, and with
the generation unchanged so the captures stay valid.
"""
pipe_mod, host = _pipeline_host(monkeypatch)
monkeypatch.setattr(pipe_mod, "DecodeGraphRunner", lambda lm, cache, dev: SimpleNamespace(cache=cache))
ctx = pipe_mod.SenseNovaU1Pipeline._decode_context
first = ctx(host, _dyn_cache(10))
assert first is not None
second_input = _dyn_cache(20)
second = ctx(host, second_input)
assert second[0] is first[0], "the cache was rebuilt for a request that fits it"
assert second[1] is first[1], "the runner was rebuilt, so its captures are gone"
assert second[0].generation == first[0].generation, "reuse must not invalidate a capture"
assert second[0].length == 20
# The reused buffers must carry the new request's prefill, not the old one.
flat = second[0].k[0].view(-1, KV_HEADS, HEAD_DIM)[:20]
torch.testing.assert_close(flat, second_input.layers[0].keys[0].transpose(0, 1))
# A prefill past the bucket cannot share the buffers.
bigger = ctx(host, _dyn_cache(BUCKETS[0] + 1))
assert bigger[0] is not first[0], "a prefill past the bucket must get new buffers"
def test_releasing_the_captures_forces_a_rebuild(monkeypatch):
"""Sleep level 2 discards the memory a capture recorded, so what the
pipeline reuses across requests has to be droppable."""
pipe_mod, host = _pipeline_host(monkeypatch)
monkeypatch.setattr(pipe_mod, "DecodeGraphRunner", lambda lm, cache, dev: SimpleNamespace(cache=cache))
ctx = pipe_mod.SenseNovaU1Pipeline._decode_context
first = ctx(host, _dyn_cache(10))
assert ctx(host, _dyn_cache(10))[0] is first[0], "the fixture is not reusing to begin with"
pipe_mod.SenseNovaU1Pipeline.release_captured_graphs(host)
after = ctx(host, _dyn_cache(10))
assert after[0] is not first[0], "the cache survived the release"
assert after[1] is not first[1], "the runner survived the release"
class _FakeLoRAWrapper(BaseLayerWithLoRA):
"""Stands in for what ``replace_submodule`` leaves in the module tree."""
def __init__(self):
super().__init__()
def test_no_capture_crosses_a_request_once_lora_wrappers_are_installed(monkeypatch):
"""Base -> LoRA -> base. The reuse test cannot see an adapter change.
``DiffusionLoRAManager`` replaces the decode path's linear layers with
``BaseLayerWithLoRA`` and binds, rescales or resets slot 0 per request, and
the wrappers stay in the tree afterwards. A graph captured before an adapter
was bound would replay without its matmuls; one captured with it bound would
survive its removal. Neither moves the shape, dtype or bucket that decide
reuse, so the wrappers themselves have to end it.
"""
pipe_mod, host = _pipeline_host(monkeypatch)
monkeypatch.setattr(pipe_mod, "DecodeGraphRunner", lambda lm, cache, dev: SimpleNamespace(cache=cache))
ctx = pipe_mod.SenseNovaU1Pipeline._decode_context
base = ctx(host, _dyn_cache(10))
assert ctx(host, _dyn_cache(10))[0] is base[0], "the fixture is not reusing to begin with"
host.language_model.add_module("q_proj", _FakeLoRAWrapper())
adapted = ctx(host, _dyn_cache(10))
assert adapted[0] is not base[0], "an adapted request reused the base capture"
assert adapted[1] is not base[1], "an adapted request reused the base runner"
back = ctx(host, _dyn_cache(10))
assert back[0] is not adapted[0], "a later request reused the adapted capture"
assert back[1] is not adapted[1], "a later request reused the adapted runner"
assert getattr(host, "_paged_decode", None) is None, "an adapted capture was left on the pipeline"
@cuda_only
@pytest.mark.cuda
@pytest.mark.L4
def test_a_second_request_replays_the_first_capture(monkeypatch):
"""Reuse is only worth anything if the capture survives it."""
dev = torch.device("cuda")
pipe_mod, host = _pipeline_host(monkeypatch, device="cuda", language_model=_StubLM(dev))
ctx = pipe_mod.SenseNovaU1Pipeline._decode_context
cache, runner = ctx(host, _dyn_cache(10, dev))
cache.set_length(cache.length + 1)
runner.step(0, 0)
assert runner.captures == 1
cache2, runner2 = ctx(host, _dyn_cache(12, dev))
assert cache2 is cache and runner2 is runner
cache2.set_length(cache2.length + 1)
runner2.step(0, 0)
assert runner.captures == 1, "the second request captured the graph again"
class _Tok:
"""The two tokenizer calls ``_generate_text`` makes."""
EOS = 7
def convert_tokens_to_ids(self, token):
return self.EOS
def decode(self, ids, skip_special_tokens=False):
return " ".join(str(i) for i in ids)
class _Out:
def __init__(self, logits, past_key_values):
self.logits = logits
self.past_key_values = past_key_values
class _TextHost:
"""Carries only what ``_generate_text`` touches."""
def __init__(self, context):
self._context = context
self.seen: list[object] = []
self.tokenizer = _Tok()
def _decode_context(self, past_key_values):
return self._context
def _ar_step(self, next_token, t_idx, past_key_values, decode=None):
self.seen.append(decode)
# Stop the loop: argmax of this row is the EOS id.
logits = torch.full((1, 1, 8), -1.0)
logits[0, 0, _Tok.EOS] = 1.0
return _Out(logits, past_key_values)
def test_text_decoding_uses_the_paged_context_too():
"""Image-to-text and text-to-text run their own decode loop. It used to call
``_ar_step`` without the context, so the paged path was documented but never
reached outside think."""
from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import (
SenseNovaU1Pipeline,
)
context = object()
host = _TextHost(context)
prefix = torch.full((1, 1, 8), -1.0)
prefix[0, 0, 5] = 1.0 # not EOS, so one step runs
SenseNovaU1Pipeline._generate_text(host, prefix, past_key_values=None, t_idx=0)
assert host.seen == [context], f"text decoding ran with decode={host.seen}"
class _WarmupReq:
"""The engine's dummy request, which carries the reserved id."""
def __init__(self, request_id):
self.request_id = request_id
self.prompts = [{"prompt": "", "modalities": ["image"]}]
def test_the_dummy_warmup_request_drives_one_decode_step(monkeypatch):
"""The dummy request is a text-to-image run with think off, so it exercises
the prefill shape and never the decode one. Without this hook whatever the
deploy config compiles for decode lands on the first real request."""
from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import (
SenseNovaU1Pipeline,
)
from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID
calls: list[str] = []
host = object.__new__(SenseNovaU1Pipeline)
monkeypatch.setattr(host, "_warm_ar_decode", lambda: calls.append("warm"), raising=False)
monkeypatch.setattr(
SenseNovaU1Pipeline,
"_parse_request",
lambda self, req: (_ for _ in ()).throw(_StopForwardError()),
raising=False,
)
with pytest.raises(_StopForwardError):
SenseNovaU1Pipeline.forward(host, _WarmupReq(DUMMY_DIFFUSION_REQUEST_ID))
assert calls == ["warm"], "the dummy warmup request did not drive a decode step"
calls.clear()
with pytest.raises(_StopForwardError):
SenseNovaU1Pipeline.forward(host, _WarmupReq("a-real-request"))
assert calls == [], "a real request paid for the warmup"
class _StopForwardError(Exception):
"""Stops ``forward`` right after the warmup hook, so the test needs no model."""
@@ -0,0 +1,107 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""CPU tests for SenseNova-U1 RMSNorm and the U1.5 config surface.
``Qwen3RMSNorm`` was a hand-rolled fp32 up/down-cast chain; it now calls the fused
``F.rms_norm``. These pin the mathematics against an independent reference so a
future kernel swap cannot silently change what the layer computes, and pin the two
config fields that distinguish U1.5 from U1.
"""
import pytest
import torch
from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import Qwen3RMSNorm
from vllm_omni.transformers_utils.configs.sensenova_u1 import SenseNovaU1Config
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
def _reference(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
"""RMSNorm written out in float64 so it is independent of the implementation."""
x64 = x.to(torch.float64)
var = x64.pow(2).mean(-1, keepdim=True)
return (weight.to(torch.float64) * (x64 * torch.rsqrt(var + eps))).to(x.dtype)
@pytest.mark.parametrize("hidden", [8, 128, 4096])
@pytest.mark.parametrize(
"dtype,atol",
[(torch.float32, 1e-5), (torch.bfloat16, 3e-2), (torch.float16, 4e-3)],
)
def test_rmsnorm_matches_float64_reference(hidden: int, dtype: torch.dtype, atol: float):
torch.manual_seed(0)
norm = Qwen3RMSNorm(hidden, eps=1e-6).to(dtype)
with torch.no_grad():
norm.weight.copy_(torch.rand(hidden, dtype=dtype) + 0.5)
x = torch.randn(7, hidden, dtype=dtype)
got = norm(x)
want = _reference(x, norm.weight, norm.variance_epsilon)
assert got.dtype == dtype
assert got.shape == x.shape
torch.testing.assert_close(got.to(torch.float32), want.to(torch.float32), atol=atol, rtol=atol)
def _cast_chain(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
"""The implementation this PR replaced, kept so its accuracy can be compared."""
var = x.to(torch.float32).pow(2).mean(-1, keepdim=True)
return weight * (x.to(torch.float32) * torch.rsqrt(var + eps)).to(weight.dtype)
@pytest.mark.parametrize("hidden", [512, 4096])
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
def test_rmsnorm_is_closer_to_float64_than_the_cast_chain(hidden: int, dtype: torch.dtype):
"""The cast chain rounds to ``dtype`` before the weight multiply and so rounds
twice, while ``F.rms_norm`` rounds once. Restoring the cast chain fails here."""
torch.manual_seed(0)
norm = Qwen3RMSNorm(hidden, eps=1e-6).to(dtype)
with torch.no_grad():
norm.weight.copy_(torch.rand(hidden, dtype=dtype) + 0.5)
x = torch.randn(64, hidden, dtype=dtype)
reference = _reference(x, norm.weight, norm.variance_epsilon).to(torch.float64)
scale = reference.abs().mean()
def relative_error(out: torch.Tensor) -> float:
return ((out.to(torch.float64) - reference).abs().mean() / scale).item()
fused = relative_error(norm(x))
legacy = relative_error(_cast_chain(x, norm.weight, norm.variance_epsilon))
assert fused < legacy, f"{dtype} hidden={hidden}: fused {fused:.3e} not better than {legacy:.3e}"
def test_rmsnorm_is_scale_equivariant():
"""RMSNorm divides by the RMS, so scaling the input must not change the output."""
torch.manual_seed(0)
norm = Qwen3RMSNorm(64, eps=1e-12)
x = torch.randn(4, 64)
torch.testing.assert_close(norm(x), norm(x * 7.5), atol=1e-4, rtol=1e-4)
def test_rmsnorm_preserves_leading_dims():
norm = Qwen3RMSNorm(16, eps=1e-6)
for shape in [(16,), (3, 16), (2, 3, 16)]:
assert norm(torch.randn(*shape)).shape == shape
def test_rmsnorm_weight_is_applied():
norm = Qwen3RMSNorm(4, eps=1e-6)
x = torch.randn(2, 4)
base = norm(x)
with torch.no_grad():
norm.weight.mul_(3.0)
torch.testing.assert_close(norm(x), base * 3.0, atol=1e-5, rtol=1e-5)
def test_u15_config_fields_default_to_u1_values():
"""U1.5 flips exactly these two; the defaults must stay U1's so old checkpoints
keep their behaviour when the keys are absent."""
cfg = SenseNovaU1Config()
assert cfg.use_pixel_head is False
assert cfg.noise_scale_max_value == 10.0
def test_u15_config_fields_are_read_from_kwargs():
cfg = SenseNovaU1Config(use_pixel_head=True, noise_scale_max_value=16.0)
assert cfg.use_pixel_head is True
assert cfg.noise_scale_max_value == 16.0
@@ -179,6 +179,30 @@ def test_request_scoped_cache_dit_lifecycle_is_pipeline_opt_in():
assert events == [backend]
def test_release_captured_graphs_clears_runners_and_delegates_to_the_pipeline():
"""Both halves in one place, so a pipeline that keeps captures is collected
without the caller knowing which pipelines have one."""
released = []
runner = object.__new__(DiffusionModelRunner)
runner.graph_runners = {"decode": object()}
runner.pipeline = SimpleNamespace(release_captured_graphs=lambda: released.append(True))
runner.release_captured_graphs()
assert runner.graph_runners == {}
assert released == [True]
def test_release_captured_graphs_tolerates_a_pipeline_without_captures():
"""Most pipelines keep none, and the runner may not carry `graph_runners`."""
runner = object.__new__(DiffusionModelRunner)
runner.pipeline = object()
runner.release_captured_graphs()
assert not hasattr(runner, "graph_runners")
def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summary: bool = True):
runner = object.__new__(DiffusionModelRunner)
runner.vllm_config = object()
@@ -441,7 +441,6 @@ class TestPipelineStreamingOutputToEntrypoint:
model="test_model",
od_config=od_config,
metadata=metadata,
batch_size=1,
)
@staticmethod
@@ -526,6 +525,7 @@ class TestSupportedPipelines:
runner = object.__new__(DiffusionModelRunner)
runner.vllm_config = _make_vllm_config()
runner.od_config = SimpleNamespace(
max_num_seqs=1,
enable_cpu_offload=False,
enable_layerwise_offload=False,
enforce_eager=True,
+16 -1
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Unit tests for DiffusionWorker class.
@@ -143,6 +143,21 @@ class TestDiffusionWorkerSleep:
assert "buffer1" in mock_gpu_worker._sleep_saved_buffers
assert "buffer2" in mock_gpu_worker._sleep_saved_buffers
def test_sleep_level_2_releases_captures_through_the_model_runner(self, mocker: MockerFixture, mock_gpu_worker):
"""Level 2 discards the memory captures were recorded against, and the
runner owns the release. The worker asks it once and does not reach past
it into the pipeline."""
mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
mock_platform.get_free_memory.side_effect = [10 * 1024**3, 12 * 1024**3]
mock_platform.get_device_total_memory.return_value = 80 * 1024**3
mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory", side_effect=[0, 0])
mock_gpu_worker.model_runner.pipeline.named_buffers = mocker.Mock(return_value=[])
mock_gpu_worker.sleep(level=2)
mock_gpu_worker.model_runner.release_captured_graphs.assert_called_once_with()
mock_gpu_worker.model_runner.pipeline.release_captured_graphs.assert_not_called()
def test_sleep_memory_freed_validation(self, mocker: MockerFixture, mock_gpu_worker):
"""Test that sleep validates memory was actually freed."""
mock_allocator_class = patch_cumem_allocator(mocker)
@@ -47,7 +47,7 @@ def client(mock_engine):
with patch.object(InlineStageDiffusionClient, "_enrich_config"):
od_config = MagicMock(spec=OmniDiffusionConfig)
od_config.streaming_output = False
c = InlineStageDiffusionClient(model="test_model", od_config=od_config, metadata=metadata, batch_size=1)
c = InlineStageDiffusionClient(model="test_model", od_config=od_config, metadata=metadata)
yield c
c.shutdown()
@@ -234,5 +234,4 @@ def test_inline_client_requires_replica_id(mock_engine):
model="test_model",
od_config=od_config,
metadata=metadata,
batch_size=1,
)
@@ -1171,7 +1171,6 @@ class TestStageDiffusionClientErrorPropagation:
metadata,
"tcp://req",
"tcp://resp",
batch_size=1,
)
@pytest.mark.asyncio
-1
View File
@@ -387,7 +387,6 @@ class TestPromptUpdateIntegration:
model="test_model",
od_config=od_config,
metadata=metadata,
batch_size=1,
)
@staticmethod
@@ -22,7 +22,7 @@ from vllm.inputs import TokensPrompt
from tests.helpers.clean import wait_for_gpu_memory_to_clear
from tests.helpers.mark import hardware_test
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni.entrypoints.async_omni import AsyncOmni
from vllm_omni.platforms import current_omni_platform
@@ -52,26 +52,7 @@ def _colocate_async_deploy() -> str:
CI yaml plus CuMem on every stage fails engine-core startup in the L3
merge job: Stage 1 resolves ``devices: "1"`` while only GPU 0 is visible.
"""
return modify_stage_config(
get_deploy_config_path("qwen3_omni_moe_thinking.yaml"),
updates={
"stages": {
0: {
"enable_sleep_mode": True,
"enforce_eager": True,
"trust_remote_code": True,
"tensor_parallel_size": 1,
"devices": "0",
"gpu_memory_utilization": 0.9,
"max_num_seqs": 1,
"max_model_len": 8192,
"max_num_batched_tokens": 8192,
"enable_prefix_caching": False,
"skip_mm_profiling": True,
},
},
},
)
return get_deploy_config_path("ci/qwen3_omni_moe_colocate_async.yaml")
@pytest.mark.advanced_model
@@ -0,0 +1,4 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""MiniMax online-serving tests."""
@@ -0,0 +1,369 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Shared request and server helpers for MiniMax-H3 L3 tests."""
from __future__ import annotations
import base64
import concurrent.futures
import io
import json
import os
from collections.abc import Callable
from pathlib import Path
from typing import Any
import av
import requests
from tests.helpers.assertions import assert_video_valid
from tests.helpers.media import generate_synthetic_image
from tests.helpers.runtime import OmniServerParams, OpenAIClientHandler
# DLO starts worker processes and uses shared-memory queues. Always select
# spawn for these cases so a parent fork cannot inherit stale resource-tracker
# state from a preceding server lifecycle.
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
MODEL = os.environ.get("VLLM_TEST_MINIMAX_H3_MODEL", "MiniMaxAI/MiniMax-H3")
WIDTH = 1344
HEIGHT = 768
FPS = 24
NUM_INFERENCE_STEPS = 4
# Keep a failed collective/preprocessing regression bounded well below the
# per-file Buildkite timeout. Healthy H100 runs complete in under two minutes
# after startup; ten minutes leaves ample room for a cold compile.
REQUEST_TIMEOUT_SECONDS = 600
# FastH3's four-step preview is a load-time fusion adapter. Its model-level
# deltas cannot use the dynamic PEFT LoRA manager, so this test deliberately
# skips on branches that predate PR #6714 instead of starting a base-H3 server.
_REPO_ROOT = Path(__file__).resolve().parents[4]
FASTH3_SUPPORTED = (_REPO_ROOT / "vllm_omni/diffusion/models/minimax_h3/fasth3.py").is_file()
# The Turbo artifact is handled by the MiniMax-specific loader introduced in
# PR #6550; the generic diffusion LoRA manager cannot parse that checkpoint.
TURBO_SUPPORTED = (_REPO_ROOT / "vllm_omni/diffusion/models/minimax_h3/lora.py").is_file()
FASTH3_WIDTH = 1024
FASTH3_HEIGHT = 576
FASTH3_DURATION = 4.4
DLO_SERVER_ARGS = [
"--trust-remote-code",
"--num-gpus",
"2",
"--tensor-parallel-size",
"1",
"--data-parallel-size",
"2",
"--request-batch-max-wait-ms",
"500",
"--usp",
"1",
"--ring",
"1",
"--text-encoder-tp-size",
"1",
"--vae-patch-parallel-size",
"1",
"--vae-parallel-mode",
"tile",
"--vae-use-tiling",
"--enable-distributed-layerwise-offload",
]
_FL2VA_IMAGE = base64.b64decode(generate_synthetic_image(WIDTH, HEIGHT, seed=42)["base64"])
# Keep the DLO/DP2 Ref2VA wave on the image-only path. The separate video
# reference path performs per-request ffmpeg/Qwen video preprocessing and has
# repeatedly stalled this merge wave before the first pipeline log; it remains
# covered by the Ref2VA contract/accuracy suites.
_REF2VA_IMAGE = base64.b64decode(generate_synthetic_image(512, 288, seed=43)["base64"])
def resolve_turbo_lora() -> str | None:
"""Resolve the legacy Turbo LoRA artifact without network at collection."""
configured = os.environ.get("VLLM_TEST_MINIMAX_H3_TURBO_LORA")
if configured and Path(configured).is_file():
return configured
try:
from huggingface_hub import hf_hub_download
return hf_hub_download(
repo_id="lightx2v/Minimax-h3-Turbo",
filename="minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors",
local_files_only=True,
)
except Exception:
return None
def resolve_fasth3_lora() -> str | None:
"""Resolve the dense/data-free FastH3 artifact without network at collection."""
configured = os.environ.get("VLLM_TEST_MINIMAX_H3_FASTH3_LORA")
if configured and Path(configured).is_file():
return configured
try:
from huggingface_hub import hf_hub_download
return hf_hub_download(
repo_id="FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA",
filename="dense-datafree/adapter_model.safetensors",
local_files_only=True,
)
except Exception:
# Buildkite downloads the artifact before pytest. Local collection
# must stay offline-friendly when that optional cache is absent.
return None
FASTH3_LORA = resolve_fasth3_lora()
TURBO_LORA = resolve_turbo_lora()
def _assert_audio_stream_present(video: bytes) -> None:
"""Assert that the generated MP4 contains decodable audio samples."""
with av.open(io.BytesIO(video)) as container:
audio_streams = [stream for stream in container.streams if stream.type == "audio"]
assert audio_streams, "MiniMax-H3 MP4 has no audio stream"
audio_frame = next(container.decode(audio=0), None)
assert audio_frame is not None and audio_frame.samples > 0, "MiniMax-H3 MP4 audio stream is empty"
def assert_h3_video(video: bytes, *, width: int, height: int) -> None:
"""Check the common H3 video/audio response contract."""
assert_video_valid(video, width=width, height=height, fps=FPS)
_assert_audio_stream_present(video)
def post_sync(
client: OpenAIClientHandler,
form_data: dict[str, str],
files: Any = None,
) -> bytes:
"""Submit one synchronous video request and return its MP4 body."""
response = requests.post(
f"{client.base_url.rstrip('/')}/v1/videos/sync",
data=form_data,
files=files,
headers={"Accept": "video/mp4"},
timeout=REQUEST_TIMEOUT_SECONDS,
)
response.raise_for_status()
assert response.headers.get("content-type", "").startswith("video/mp4")
assert response.content, "MiniMax-H3 returned an empty video body"
return response.content
def _h3_form(task: str, seed: int) -> dict[str, str]:
return {
"model": MODEL,
"prompt": (
"A cinematic live-action scene with a clear subject moving naturally; "
"the atmosphere includes synchronized environmental sound."
),
"width": str(WIDTH),
"height": str(HEIGHT),
"fps": str(FPS),
"num_inference_steps": str(NUM_INFERENCE_STEPS),
"flow_shift": "12",
"seed": str(seed),
"extra_params": json.dumps(
{
"task": task,
"duration": 4.0,
"aspect_ratio": "16:9",
"audio_flow_shift": 3.0,
},
separators=(",", ":"),
),
}
def run_fl2va(client: OpenAIClientHandler, seed: int) -> bytes:
return post_sync(
client,
_h3_form("fl2va", seed),
files=fl2va_files(),
)
def run_ref2va(client: OpenAIClientHandler, seed: int) -> bytes:
return post_sync(
client,
_h3_form("ref2va", seed),
files=[("input_references", ("reference.jpg", io.BytesIO(_REF2VA_IMAGE), "image/jpeg"))],
)
def fl2va_files() -> dict[str, tuple[str, io.BytesIO, str]]:
"""Return a fresh multipart first-frame payload for an FL2VA request."""
return {"input_reference": ("first_frame.jpg", io.BytesIO(_FL2VA_IMAGE), "image/jpeg")}
def run_dlo_wave(
client: OpenAIClientHandler,
request_fn: Callable[[OpenAIClientHandler, int], bytes],
) -> list[bytes]:
"""Run one complete two-request DLO/DP2 wave concurrently."""
with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
return list(executor.map(request_fn, (client, client), (2101, 2102)))
def dlo_params(task_type: str) -> OmniServerParams:
return OmniServerParams(
model=MODEL,
server_args=[*DLO_SERVER_ARGS, "--task-type", task_type],
stage_init_timeout=1800,
init_timeout=1800,
)
def turbo_params(lora_path: str | None) -> OmniServerParams:
"""Build the two-card TP2 + DLO-no-AllGather Turbo layout."""
return OmniServerParams(
model=MODEL,
server_args=[
"--trust-remote-code",
"--task-type",
"fl2va",
"--num-gpus",
"2",
"--tensor-parallel-size",
"2",
"--usp",
"1",
"--ring",
"1",
"--text-encoder-tp-size",
"2",
"--vae-patch-parallel-size",
"2",
"--vae-parallel-mode",
"tile",
"--vae-use-tiling",
"--enable-distributed-layerwise-offload",
"--dlo-no-use-allgather",
"--lora-backend",
"peft",
"--lora-path",
lora_path or "",
],
stage_init_timeout=1800,
init_timeout=1800,
)
def turbo_form(seed: int) -> dict[str, str]:
"""Build the official five-sigma-point Turbo FL2VA request."""
return {
"model": MODEL,
"prompt": (
"A man stands beside a yellow car at night; the car drives away "
"as he follows it with his eyes and begins singing sadly."
),
"width": str(WIDTH),
"height": str(HEIGHT),
"fps": str(FPS),
# The public Turbo artifact has four denoiser evaluations but its API
# contract requests five sigma points.
"num_inference_steps": "5",
"flow_shift": "6",
"seed": str(seed),
"extra_params": json.dumps(
{"task": "fl2va", "duration": FASTH3_DURATION, "aspect_ratio": "16:9", "audio_flow_shift": 3.0},
separators=(",", ":"),
),
"lora": json.dumps(
{"name": "h3-turbo-v1.0", "path": TURBO_LORA, "scale": 1.0},
separators=(",", ":"),
),
}
def fasth3_params(lora_path: str | None) -> OmniServerParams:
"""Build the requested two-card FastH3 parallel layout.
HSDP owns DiT sharding (the model tensor-parallel degree stays at its
default 1), USP2 shards the sequence, text-encoder TP2 shards Qwen3-VL,
and VAE patch parallel 2 is the project's VPP2 setting.
"""
return OmniServerParams(
model=MODEL,
server_args=[
"--trust-remote-code",
"--task-type",
"fl2va",
"--num-gpus",
"2",
"--usp",
"2",
"--ring",
"1",
"--use-hsdp",
"--hsdp-shard-size",
"2",
"--text-encoder-tp-size",
"2",
"--vae-patch-parallel-size",
"2",
"--vae-parallel-mode",
"tile",
"--vae-use-tiling",
"--lora-path",
lora_path or "",
],
stage_init_timeout=1800,
init_timeout=1800,
)
def fasth3_form(seed: int) -> dict[str, str]:
"""Build the four-step T2VA request accepted by the FastH3 preview."""
return {
"model": MODEL,
"prompt": (
"A cinematic live-action scene with a clear subject moving naturally; "
"the atmosphere includes synchronized environmental sound."
),
"width": str(FASTH3_WIDTH),
"height": str(FASTH3_HEIGHT),
"fps": str(FPS),
"num_inference_steps": "4",
"seed": str(seed),
# FastH3 is fused at startup, so a request-level `lora` field is both
# unnecessary and rejected by the load-time adapter contract.
"extra_params": json.dumps(
{
"task": "t2va",
"duration": FASTH3_DURATION,
"aspect_ratio": "16:9",
},
separators=(",", ":"),
),
}
__all__ = [
"FASTH3_LORA",
"FASTH3_SUPPORTED",
"FASTH3_DURATION",
"FASTH3_HEIGHT",
"FASTH3_WIDTH",
"FPS",
"MODEL",
"assert_h3_video",
"dlo_params",
"fasth3_form",
"fasth3_params",
"fl2va_files",
"post_sync",
"run_dlo_wave",
"run_fl2va",
"run_ref2va",
"TURBO_LORA",
"TURBO_SUPPORTED",
"turbo_form",
"turbo_params",
]
@@ -0,0 +1,27 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""MiniMax-H3 FL2VA DLO + DP2 L3 online-serving case."""
from __future__ import annotations
import pytest
from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import OmniServer, OpenAIClientHandler
from ._common import assert_h3_video, dlo_params, run_dlo_wave, run_fl2va
pytestmark = [pytest.mark.core_model, pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.slow]
H100_TWO_CARD_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize(
"omni_server",
[pytest.param(dlo_params("fl2va"), id="minimax_h3_dlo_dp2_fl2va", marks=H100_TWO_CARD_MARKS)],
indirect=True,
)
def test_minimax_h3_dlo_dp2_fl2va(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
"""Validate a complete concurrent DLO/DP2 wave with two FL2VA jobs."""
for video in run_dlo_wave(openai_client, run_fl2va):
assert_h3_video(video, width=1344, height=768)
@@ -0,0 +1,27 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""MiniMax-H3 Ref2VA DLO + DP2 L3 online-serving case."""
from __future__ import annotations
import pytest
from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import OmniServer, OpenAIClientHandler
from ._common import assert_h3_video, dlo_params, run_dlo_wave, run_ref2va
pytestmark = [pytest.mark.core_model, pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.slow]
H100_TWO_CARD_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize(
"omni_server",
[pytest.param(dlo_params("ref2va"), id="minimax_h3_dlo_dp2_ref2va", marks=H100_TWO_CARD_MARKS)],
indirect=True,
)
def test_minimax_h3_dlo_dp2_ref2va(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
"""Validate a complete concurrent DLO/DP2 wave with two image-only Ref2VA jobs."""
for video in run_dlo_wave(openai_client, run_ref2va):
assert_h3_video(video, width=1344, height=768)
@@ -0,0 +1,56 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""MiniMax-H3 FastH3 four-step LoRA L3 online-serving case."""
from __future__ import annotations
import pytest
from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import OmniServer, OpenAIClientHandler
from ._common import (
FASTH3_HEIGHT,
FASTH3_LORA,
FASTH3_SUPPORTED,
FASTH3_WIDTH,
assert_h3_video,
fasth3_form,
fasth3_params,
post_sync,
)
pytestmark = [pytest.mark.core_model, pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.slow]
H100_TWO_CARD_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.skipif(not FASTH3_SUPPORTED, reason="requires FastH3 support from MiniMax-H3 PR #6714")
@pytest.mark.skipif(
FASTH3_LORA is None,
reason="set VLLM_TEST_MINIMAX_H3_FASTH3_LORA or populate the local HF cache",
)
@pytest.mark.parametrize(
"omni_server",
[
pytest.param(
fasth3_params(FASTH3_LORA),
id="minimax_h3_fasth3_lora_hsdp2_usp2_vpp2",
marks=H100_TWO_CARD_MARKS,
)
],
indirect=True,
)
def test_minimax_h3_fasth3_lora_hsdp2_usp2_vpp2(
omni_server: OmniServer,
openai_client: OpenAIClientHandler,
) -> None:
"""Run the dense/data-free FastH3 adapter on two H100-class cards.
The preview adapter is T2VA-only and must be fused through ``--lora-path``;
it is not a request-switchable PEFT adapter. 1024x576 is intentionally
below the 1344x768 H3 reference shape to stay inside an 80 GiB H100 after
regional compile cache growth.
"""
video = post_sync(openai_client, fasth3_form(seed=3101))
assert_h3_video(video, width=FASTH3_WIDTH, height=FASTH3_HEIGHT)
@@ -0,0 +1,56 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""MiniMax-H3 Turbo LoRA + DLO L3 online-serving case."""
from __future__ import annotations
import pytest
from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import OmniServer, OpenAIClientHandler
from ._common import (
TURBO_LORA,
TURBO_SUPPORTED,
assert_h3_video,
fl2va_files,
post_sync,
turbo_form,
turbo_params,
)
pytestmark = [pytest.mark.core_model, pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.slow]
H100_TWO_CARD_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.skipif(
not TURBO_SUPPORTED,
reason="requires the MiniMax-H3 Turbo loader from PR #6550",
)
@pytest.mark.skipif(
TURBO_LORA is None,
reason="set VLLM_TEST_MINIMAX_H3_TURBO_LORA or populate the local HF cache",
)
@pytest.mark.parametrize(
"omni_server",
[
pytest.param(
turbo_params(TURBO_LORA),
id="minimax_h3_dlo_turbo_lora_tp2",
marks=H100_TWO_CARD_MARKS,
)
],
indirect=True,
)
def test_minimax_h3_dlo_turbo_lora_fl2va(
omni_server: OmniServer,
openai_client: OpenAIClientHandler,
) -> None:
"""Validate the official four-evaluation Turbo LoRA + DLO FL2VA request."""
video = post_sync(
openai_client,
turbo_form(seed=2201),
files=fl2va_files(),
)
assert_h3_video(video, width=1344, height=768)
@@ -1,15 +1,15 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Online serving tests for Boogu-Image-0.1-Edit (image-to-image via chat completions).
Online serving tests for Boogu-Image-0.1-Edit and Edit-Turbo.
The Edit checkpoint shares the ``BooguImagePipeline`` class with the Base
Both Edit checkpoints share the ``BooguImagePipeline`` class with the Base
text-to-image checkpoint; the editing (TI2I) path activates when the request
carries a reference image. Text guidance rides on ``guidance_scale`` (upstream
default 4.0), and image guidance rides on ``guidance_scale_2`` (upstream default
1.0 = off; a value > 1 enables the double-guidance path). Only a single
reference image is supported.
carries a reference image. The regular Edit checkpoint uses text guidance via
``guidance_scale`` and optional image guidance via ``guidance_scale_2``.
Edit-Turbo is guidance-distilled and uses four steps with guidance scale 1.0.
Only a single reference image is supported.
- ``test_single_image_to_image_001``: one reference image, text guidance only.
- ``test_double_guidance_001``: one reference image, text + image guidance.
@@ -40,6 +40,10 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
EDIT_PROMPT = "Change the style to a colored pencil drawing."
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
EDIT_MODEL = "Boogu/Boogu-Image-0.1-Edit"
EDIT_TURBO_MODEL = "Boogu/Boogu-Image-0.1-Edit-Turbo"
EDIT_TURBO_REVISION = "hotfix-1k-20260708"
def _get_diffusion_feature_cases(model: str):
"""Return one ``default`` ``OmniServerParams`` row for ``model`` (no extra ``server_args``)."""
@@ -54,13 +58,44 @@ def _get_diffusion_feature_cases(model: str):
]
def _get_single_edit_cases():
"""Return the regular and distilled Edit server/request configurations."""
return [
pytest.param(
OmniServerParams(model=EDIT_MODEL),
{
"num_inference_steps": 2,
"guidance_scale": 5.0,
},
id="edit",
marks=SINGLE_CARD_FEATURE_MARKS,
),
pytest.param(
OmniServerParams(
model=EDIT_TURBO_MODEL,
server_args=["--revision", EDIT_TURBO_REVISION],
),
{
"num_inference_steps": 4,
"guidance_scale": 1.0,
},
id="edit-turbo-hotfix-1k",
marks=SINGLE_CARD_FEATURE_MARKS,
),
]
@pytest.mark.parametrize(
"omni_server",
_get_diffusion_feature_cases("Boogu/Boogu-Image-0.1-Edit"),
indirect=True,
("omni_server", "sampling_params"),
_get_single_edit_cases(),
indirect=["omni_server"],
)
def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
"""Single-reference text-guided edit smoke for ``Boogu/Boogu-Image-0.1-Edit``."""
def test_single_image_to_image_001(
omni_server: OmniServer,
sampling_params: dict,
openai_client: OpenAIClientHandler,
):
"""Single-reference text-guided edit smoke for regular and Turbo Edit."""
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT)
@@ -69,9 +104,7 @@ def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenA
"model": omni_server.model,
"messages": messages,
"extra_body": {
"num_inference_steps": 2,
# Boogu text guidance (upstream default 4.0); > 1 enables CFG.
"guidance_scale": 5.0,
**sampling_params,
"seed": 42,
},
}
@@ -81,7 +114,7 @@ def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenA
@pytest.mark.parametrize(
"omni_server",
_get_diffusion_feature_cases("Boogu/Boogu-Image-0.1-Edit"),
_get_diffusion_feature_cases(EDIT_MODEL),
indirect=True,
)
def test_double_guidance_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
@@ -105,15 +138,20 @@ def test_double_guidance_001(omni_server: OmniServer, openai_client: OpenAIClien
openai_client.send_diffusion_request(request_config)
def _edits_request_config(model: str, *, guidance_scale_2: float | None = None) -> dict:
def _edits_request_config(
model: str,
*,
num_inference_steps: int = 2,
guidance_scale: float = 5.0,
guidance_scale_2: float | None = None,
) -> dict:
"""Multipart ``/v1/images/edits`` config: one synthetic reference + smoke-depth params."""
image_bytes = base64.b64decode(generate_synthetic_image(512, 512)["base64"])
data = {
"model": model,
"prompt": EDIT_PROMPT,
"num_inference_steps": 2,
# Boogu text guidance (upstream default 4.0); > 1 enables CFG.
"guidance_scale": 5.0,
"num_inference_steps": num_inference_steps,
"guidance_scale": guidance_scale,
"seed": 42,
}
if guidance_scale_2 is not None:
@@ -133,19 +171,25 @@ def _assert_edits_response_has_image(responses) -> None:
@pytest.mark.parametrize(
"omni_server",
_get_diffusion_feature_cases("Boogu/Boogu-Image-0.1-Edit"),
indirect=True,
("omni_server", "sampling_params"),
_get_single_edit_cases(),
indirect=["omni_server"],
)
def test_images_edits_endpoint_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
"""Single-reference text-guided edit via the OpenAI-compatible ``/v1/images/edits`` endpoint."""
responses = openai_client.send_images_edits_http_request(_edits_request_config(omni_server.model))
def test_images_edits_endpoint_001(
omni_server: OmniServer,
sampling_params: dict,
openai_client: OpenAIClientHandler,
):
"""Single-reference regular and Turbo Edit via ``/v1/images/edits``."""
responses = openai_client.send_images_edits_http_request(
_edits_request_config(omni_server.model, **sampling_params)
)
_assert_edits_response_has_image(responses)
@pytest.mark.parametrize(
"omni_server",
_get_diffusion_feature_cases("Boogu/Boogu-Image-0.1-Edit"),
_get_diffusion_feature_cases(EDIT_MODEL),
indirect=True,
)
def test_images_edits_endpoint_double_guidance_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
@@ -1,8 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
L4 expansion coverage for ``Boogu/Boogu-Image-0.1-Edit`` (image editing / TI2I).
L4 expansion coverage for Boogu-Image Edit and Edit-Turbo (image editing / TI2I).
The nightly smoke module (``test_boogu_image_edit.py``) already covers single-image
text-guided editing and the text+image double-guidance path. This expansion
@@ -16,11 +16,12 @@ module adds the two edit paths smoke does not exercise:
``height`` / ``width``: asserts the output resolution follows the reference
(upstream ``align_res``, on by default for a single reference).
The Edit checkpoint shares ``BooguImagePipeline`` with the Base checkpoint; the
TI2I path activates when the request carries a reference image. Only a single
reference image is supported. Cases stay smoke-depth (``num_inference_steps=2``).
CPU offload, Cache-DiT, and multi-GPU parallelism are not yet supported for this
model and are intentionally omitted.
Both Edit checkpoints share ``BooguImagePipeline`` with the Base checkpoint;
the TI2I path activates when the request carries a reference image. Only a
single reference image is supported. The regular Edit cases stay smoke-depth
(``num_inference_steps=2``), while Edit-Turbo uses its tuned four-step
configuration. CPU offload, Cache-DiT, and multi-GPU parallelism are not yet
supported for these models and are intentionally omitted.
From ``tests/``::
@@ -41,6 +42,8 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
MODEL = "Boogu/Boogu-Image-0.1-Edit"
EDIT_TURBO_MODEL = "Boogu/Boogu-Image-0.1-Edit-Turbo"
EDIT_TURBO_REVISION = "hotfix-1k-20260708"
EDIT_PROMPT = "Change the style to a colored pencil drawing."
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
@@ -57,6 +60,33 @@ def _get_diffusion_feature_cases(model: str):
]
def _get_align_res_cases():
"""Return regular and Turbo Edit configurations for align-res coverage."""
return [
pytest.param(
OmniServerParams(model=MODEL),
{
"num_inference_steps": 2,
"guidance_scale": 5.0,
},
id="edit",
marks=SINGLE_CARD_FEATURE_MARKS,
),
pytest.param(
OmniServerParams(
model=EDIT_TURBO_MODEL,
server_args=["--revision", EDIT_TURBO_REVISION],
),
{
"num_inference_steps": 4,
"guidance_scale": 1.0,
},
id="edit-turbo-hotfix-1k",
marks=SINGLE_CARD_FEATURE_MARKS,
),
]
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
def test_image_only_guidance(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
"""Image-only guidance: text CFG off (1.0), image CFG on (2.0) -> image-only CFG branch."""
@@ -76,8 +106,16 @@ def test_image_only_guidance(omni_server: OmniServer, openai_client: OpenAIClien
openai_client.send_diffusion_request(request_config)
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
def test_align_res_output_size(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
@pytest.mark.parametrize(
("omni_server", "sampling_params"),
_get_align_res_cases(),
indirect=["omni_server"],
)
def test_align_res_output_size(
omni_server: OmniServer,
sampling_params: dict,
openai_client: OpenAIClientHandler,
) -> None:
"""Output resolution follows the reference image (upstream ``align_res``).
No ``height`` / ``width`` is requested. The reference is 512x384 — both
@@ -92,8 +130,7 @@ def test_align_res_output_size(omni_server: OmniServer, openai_client: OpenAICli
"model": omni_server.model,
"messages": messages,
"extra_body": {
"num_inference_steps": 2,
"guidance_scale": 5.0,
**sampling_params,
"seed": 42,
},
}
@@ -138,6 +138,23 @@ def test_voice_003(omni_server, online_client) -> None:
online_client.send_audio_speech_request(request_config)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
def test_sample_rate_001(omni_server, online_client) -> None:
"""The speech API resamples Qwen3-TTS output from native 24 kHz to 8 kHz."""
request_config = {
"model": omni_server.model,
"input": "This response should be encoded as eight kilohertz audio.",
"stream": False,
"response_format": "wav",
"sample_rate": 8000,
"task_type": "CustomVoice",
"voice": "vivian",
}
online_client.send_audio_speech_request(request_config)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
def test_language_001(omni_server, online_client) -> None:
@@ -11,7 +11,7 @@ import types
import pytest
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, normalize_omni_diffusion_kwargs
from vllm_omni.engine import async_omni_engine as async_omni_engine_module
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.engine.stage_init_utils import (
@@ -172,7 +172,6 @@ def _make_stage_runtime() -> StageRuntime:
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
diffusion_batch_size=1,
async_chunk=False,
)
@@ -196,7 +195,6 @@ def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeyp
engine.model = "dummy-model"
engine.config_path = "dummy-config"
engine.single_stage_mode = False
engine.diffusion_batch_size = 1
engine.async_chunk = False
engine.tokenizer = None
engine._single_stage_id_filter = None
@@ -282,7 +280,6 @@ def test_build_logical_stage_init_plans_handles_stage_without_devices(monkeypatc
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
diffusion_batch_size=1,
async_chunk=False,
)
@@ -371,7 +368,6 @@ def test_initialize_local_diffusion_replica_restores_device_visibility_after_loc
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
diffusion_batch_size=1,
async_chunk=False,
)
@@ -416,7 +412,6 @@ def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
diffusion_batch_size=4,
async_chunk=False,
)
@@ -429,7 +424,6 @@ def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline
def _capture_launch_diffusion_stage_replica(**kwargs):
captured["stage_id"] = kwargs["metadata"].stage_id
captured["stage_init_timeout"] = kwargs["stage_init_timeout"]
captured["batch_size"] = kwargs["batch_size"]
captured["use_inline"] = kwargs["use_inline"]
captured["omni_master_server"] = kwargs["omni_master_server"]
return types.SimpleNamespace(), StageReplicaResources()
@@ -441,7 +435,6 @@ def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline
assert captured == {
"stage_id": 0,
"stage_init_timeout": 302,
"batch_size": 4,
"use_inline": expected_inline,
"omni_master_server": None,
}
@@ -511,21 +504,20 @@ def test_initialize_local_llm_replica_scopes_runtime_env(monkeypatch):
assert os.environ[runtime_env_var] == "parent-value"
def test_initialize_diffusion_stage_forwards_batch_size_without_writing_max_num_seqs(monkeypatch):
def test_initialize_diffusion_stage_preserves_configured_max_num_seqs(monkeypatch):
import vllm_omni.diffusion.stage_diffusion_client as client_mod
import vllm_omni.engine.stage_init_utils as init_mod
od_config = types.SimpleNamespace(max_num_seqs=1)
od_config = types.SimpleNamespace(max_num_seqs=4)
captured: dict[str, object] = {}
monkeypatch.setattr(init_mod, "build_diffusion_config", lambda *args: od_config)
def _capture_client(model, config, metadata, stage_init_timeout, batch_size, use_inline):
def _capture_client(model, config, metadata, stage_init_timeout, use_inline):
captured.update(
model=model,
config=config,
metadata=metadata,
stage_init_timeout=stage_init_timeout,
batch_size=batch_size,
use_inline=use_inline,
)
return object()
@@ -539,27 +531,25 @@ def test_initialize_diffusion_stage_forwards_batch_size_without_writing_max_num_
types.SimpleNamespace(),
metadata,
stage_init_timeout=12,
batch_size=4,
use_inline=True,
)
assert od_config.max_num_seqs == 1
assert od_config.max_num_seqs == 4
assert captured == {
"model": "dummy-model",
"config": od_config,
"metadata": metadata,
"stage_init_timeout": 12,
"batch_size": 4,
"use_inline": True,
}
def test_launch_diffusion_stage_replica_does_not_write_max_num_seqs_from_batch_size(monkeypatch):
def test_launch_diffusion_stage_replica_preserves_configured_max_num_seqs(monkeypatch):
import vllm_omni.diffusion.stage_diffusion_client as client_mod
import vllm_omni.diffusion.stage_diffusion_proc as proc_mod
import vllm_omni.engine.stage_engine_startup as startup_mod
od_config = types.SimpleNamespace(max_num_seqs=1, parallel_config=types.SimpleNamespace(world_size=1))
od_config = types.SimpleNamespace(max_num_seqs=4, parallel_config=types.SimpleNamespace(world_size=1))
monkeypatch.setattr(startup_mod, "build_diffusion_config", lambda *args: od_config)
monkeypatch.setattr(startup_mod, "acquire_device_locks", lambda *args: [])
monkeypatch.setattr(
@@ -596,13 +586,12 @@ def test_launch_diffusion_stage_replica_does_not_write_max_num_seqs_from_batch_s
stage_config=types.SimpleNamespace(),
metadata=types.SimpleNamespace(stage_id=0),
stage_init_timeout=12,
batch_size=4,
use_inline=False,
omni_master_server=omni_master_server,
)
assert result is sentinel_client
assert od_config.max_num_seqs == 1
assert od_config.max_num_seqs == 4
assert resources.manager is proc_manager
@@ -621,7 +610,6 @@ def test_initialize_diffusion_stage_does_not_write_max_num_seqs(monkeypatch):
types.SimpleNamespace(),
metadata,
stage_init_timeout=12,
batch_size=4,
use_inline=True,
)
@@ -672,7 +660,6 @@ def test_launch_diffusion_stage_replica_preserves_step_execution_max_num_seqs(mo
stage_config=types.SimpleNamespace(),
metadata=types.SimpleNamespace(stage_id=0),
stage_init_timeout=12,
batch_size=4,
use_inline=False,
omni_master_server=omni_master_server,
)
@@ -688,7 +675,6 @@ def test_stage_runtime_initializes_stage_pools(monkeypatch):
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
diffusion_batch_size=1,
async_chunk=False,
)
@@ -754,7 +740,6 @@ def test_stage_runtime_passes_log_stats_to_llm_replica_launch(monkeypatch):
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
diffusion_batch_size=1,
async_chunk=False,
log_stats=True,
)
@@ -798,7 +783,6 @@ def test_stage_runtime_passes_log_stats_to_output_processor(monkeypatch):
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
diffusion_batch_size=1,
async_chunk=False,
log_stats=True,
)
@@ -851,7 +835,6 @@ def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatc
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
diffusion_batch_size=1,
async_chunk=False,
)
@@ -892,7 +875,6 @@ def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monk
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=123,
diffusion_batch_size=1,
async_chunk=False,
)
@@ -930,7 +912,6 @@ def test_remote_replicas_use_distinct_init_group_keys():
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=123,
diffusion_batch_size=1,
async_chunk=False,
)
plan = _make_llm_plan(
@@ -956,7 +937,6 @@ def test_initialize_stages_cleans_up_successful_replicas_after_partial_multi_rep
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
diffusion_batch_size=1,
async_chunk=False,
)
@@ -993,7 +973,6 @@ def test_initialize_stages_cleans_up_late_successful_replicas_after_early_multi_
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
diffusion_batch_size=1,
async_chunk=False,
)
@@ -1032,7 +1011,6 @@ def test_initialize_local_llm_replica_passes_stage_init_timeout_to_complete_stag
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
diffusion_batch_size=1,
async_chunk=False,
)
stage_init_timeout = 302
@@ -1817,6 +1795,41 @@ def test_resolve_stage_configs_preserves_stage_diffusion_attention(monkeypatch):
assert stage_configs[0].engine_args.diffusion_attention_config is existing_attention
def test_resolve_stage_configs_does_not_inject_over_stage_diffusion_attention_backend(monkeypatch):
import vllm_omni.engine.async_omni_engine as engine_mod
engine = object.__new__(AsyncOmniEngine)
stage_cfg = types.SimpleNamespace(
stage_type="diffusion",
engine_args=types.SimpleNamespace(
diffusion_attention_backend="TORCH_SDPA",
diffusion_attention_config=None,
lora_path=None,
lora_scale=None,
enable_sleep_mode=None,
quantization_config=None,
),
)
monkeypatch.setattr(
engine_mod,
"load_and_resolve_stage_configs",
lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
)
_config_path, stage_configs = engine._resolve_stage_configs(
model="dummy-model",
kwargs={"diffusion_attention_backend": "FLASH_ATTN"},
trust_remote_code=False,
)
engine_args = stage_configs[0].engine_args
assert engine_args.diffusion_attention_backend == "TORCH_SDPA"
assert engine_args.diffusion_attention_config is None
normalized = normalize_omni_diffusion_kwargs(vars(engine_args))
assert normalized["diffusion_attention_config"].default.backend == "TORCH_SDPA"
def test_resolve_stage_configs_does_not_inject_diffusion_attention_into_llm_stage(monkeypatch):
import vllm_omni.engine.async_omni_engine as engine_mod
+2 -13
View File
@@ -575,7 +575,6 @@ class TestSingleStageInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
diffusion_batch_size=2,
async_chunk=False,
single_stage_id_filter=stage_id_filter,
omni_master_address="127.0.0.1",
@@ -807,7 +806,6 @@ class TestSingleStageInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
diffusion_batch_size=2,
async_chunk=False,
)
mocker.patch.object(runtime_mod, "prepare_engine_environment")
@@ -860,7 +858,6 @@ class TestSingleStageReplicaInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
diffusion_batch_size=2,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",
@@ -928,7 +925,6 @@ class TestSingleStageReplicaInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
diffusion_batch_size=2,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",
@@ -950,7 +946,6 @@ class TestSingleStageReplicaInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
diffusion_batch_size=2,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",
@@ -1002,7 +997,6 @@ class TestSingleStageReplicaInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
diffusion_batch_size=2,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",
@@ -1068,7 +1062,6 @@ class TestSingleStageReplicaInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
diffusion_batch_size=4,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",
@@ -1120,7 +1113,6 @@ class TestSingleStageReplicaInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
diffusion_batch_size=4,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",
@@ -1141,7 +1133,7 @@ class TestSingleStageReplicaInitialization:
runtime._init_visible_devices_baseline = "0"
mocker.patch.object(runtime_mod, "inject_kv_stage_info")
od_config = SimpleNamespace(max_num_seqs=None, parallel_config=SimpleNamespace(world_size=1))
od_config = SimpleNamespace(max_num_seqs=4, parallel_config=SimpleNamespace(world_size=1))
mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value=od_config)
mock_register = mocker.patch(
"vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
@@ -1178,7 +1170,7 @@ class TestSingleStageReplicaInitialization:
os.environ[device_env_var] = prev_device_env
assert result is sentinel_client
assert od_config.max_num_seqs is None
assert od_config.max_num_seqs == 4
mock_register.assert_called_once_with(
omni_master_address="127.0.0.1",
omni_master_port=25000,
@@ -1203,7 +1195,6 @@ class TestSingleStageReplicaInitialization:
request_address="tcp://127.0.0.1:26002",
response_address="tcp://127.0.0.1:26003",
proc_manager=mocker.ANY,
batch_size=4,
)
assert mock_from_addresses.call_args.kwargs["proc_manager"].proc is proc
@@ -1216,7 +1207,6 @@ class TestSingleStageReplicaInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
diffusion_batch_size=4,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",
@@ -1258,7 +1248,6 @@ class TestSingleStageReplicaInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
diffusion_batch_size=4,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",
+58 -1
View File
@@ -133,7 +133,7 @@ def _stable_engine_arg_environment(monkeypatch, tmp_path):
for worker_type in ("ar", "generation"):
worker_module = types.ModuleType(f"test.{worker_type}")
worker_module.Worker = _TestWorker
setattr(worker_module, "Worker", _TestWorker)
monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module)
@@ -434,6 +434,63 @@ def test_typed_diffusion_engine_args_use_structured_diffusion_config(tmp_path):
assert typed_args["diffusion_attention_config"].per_role["cross"].backend == "TORCH_SDPA"
def test_engine_args_consume_stage_diffusion_attention_shorthand(tmp_path):
"""A stage-level ``diffusion_attention_backend`` shorthand is folded into the structured config."""
pipeline, deploy, model = _engine_arg_inputs(tmp_path)
deploy.stages[2] = replace(
deploy.stages[2],
diffusion_attention_config=None,
diffusion_attention_backend="TORCH_SDPA",
)
legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model)
legacy_args = build_legacy_engine_args_dict(legacy_stages[2], model)
typed_args = build_engine_args_dict_from_omni_stage_config(
omni_config.stage_by_id(2),
model,
)
for engine_args in (legacy_args, typed_args):
assert engine_args.get("diffusion_attention_backend") is None
assert isinstance(engine_args["diffusion_attention_config"], AttentionConfig)
assert engine_args["diffusion_attention_config"].default.backend == "TORCH_SDPA"
od_config = OmniDiffusionConfig.from_kwargs(**engine_args)
assert od_config.diffusion_attention_config.default.backend == "TORCH_SDPA"
@pytest.mark.parametrize(
"yaml_attention_config",
[
{"default": {"backend": "FLASH_ATTN"}, "per_role": {"cross": {"backend": "SAGE_ATTN"}}},
{"per_role": {"cross": {"backend": "SAGE_ATTN"}}},
],
ids=["yaml-default", "yaml-per-role-only"],
)
def test_engine_args_apply_cli_attention_shorthand_over_yaml_config(tmp_path, yaml_attention_config):
pipeline, deploy, model = _engine_arg_inputs(tmp_path)
deploy.stages[2] = replace(deploy.stages[2], diffusion_attention_config=yaml_attention_config)
legacy_stages, omni_config = _legacy_and_typed_stages(
pipeline,
deploy,
model,
cli_overrides={"stage_2_diffusion_attention_backend": "TORCH_SDPA"},
)
legacy_args = build_legacy_engine_args_dict(legacy_stages[2], model)
typed_args = build_engine_args_dict_from_omni_stage_config(
omni_config.stage_by_id(2),
model,
)
for engine_args in (legacy_args, typed_args):
assert engine_args.get("diffusion_attention_backend") is None
attention_config = engine_args["diffusion_attention_config"]
assert attention_config.default.backend == "TORCH_SDPA"
assert attention_config.per_role["cross"].backend == "SAGE_ATTN"
od_config = OmniDiffusionConfig.from_kwargs(**engine_args)
assert od_config.diffusion_attention_config.default.backend == "TORCH_SDPA"
def test_typed_engine_args_preserve_explicit_backend_default_overrides(tmp_path):
pipeline, deploy, model = _engine_arg_inputs(tmp_path)
deploy.enable_prefix_caching = False
@@ -1,4 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Tests for audio output format handling in chat completions.
Covers:
@@ -11,10 +12,15 @@ Covers:
from __future__ import annotations
from io import BytesIO
import numpy as np
import pytest
import soundfile
import torch
import torchaudio
from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin, StreamingAudioResampler
from vllm_omni.entrypoints.openai.protocol.audio import (
DEFAULT_AUDIO_FORMAT,
SUPPORTED_AUDIO_FORMATS,
@@ -122,6 +128,81 @@ class TestCreateAudio:
decoded = base64.b64decode(response.audio_data)
assert decoded[:4] == b"RIFF"
def test_resamples_wav_to_requested_output_rate(self, mixin, audio_tensor):
response = mixin.create_audio(
CreateAudio(
audio_tensor=audio_tensor,
sample_rate=24000,
output_sample_rate=8000,
response_format="wav",
speed=1.0,
base64_encode=False,
)
)
with soundfile.SoundFile(BytesIO(response.audio_data)) as audio_file:
assert audio_file.samplerate == 8000
assert audio_file.frames == 8000
class TestStreamingAudioResampler:
@staticmethod
def _resample(waveform):
resampler = StreamingAudioResampler(24000, 8000)
return np.concatenate(
(
resampler.process(waveform),
resampler.process(np.empty(0), final=True),
)
)
def test_chunk_boundaries_do_not_change_output(self):
waveform = np.sin(np.linspace(0, 200 * np.pi, 24000, endpoint=False)).astype(np.float32)
expected = self._resample(waveform)
chunked = StreamingAudioResampler(24000, 8000)
pieces = [chunked.process(chunk) for chunk in np.split(waveform, [137, 2048, 9001, 17003])]
pieces.append(chunked.process(np.empty(0), final=True))
actual = np.concatenate(pieces)
assert actual.shape == (8000,)
np.testing.assert_allclose(actual, expected, atol=1e-6)
reference = torchaudio.functional.resample(torch.from_numpy(waveform), 24000, 8000).numpy()
np.testing.assert_allclose(actual[100:-100], reference[100:-100], atol=2e-3, rtol=1e-3)
def test_preserves_passband_signal(self):
samples = np.arange(24000, dtype=np.float32)
waveform = np.sin(2 * np.pi * 3000 * samples / 24000).astype(np.float32)
output = self._resample(waveform)
input_rms = np.sqrt(np.mean(waveform**2))
output_rms = np.sqrt(np.mean(output[100:-100] ** 2))
assert output_rms == pytest.approx(input_rms, rel=0.02)
@pytest.mark.parametrize("frequency", [6000, 10000])
def test_attenuates_aliasing_frequencies(self, frequency):
samples = np.arange(24000, dtype=np.float32)
waveform = np.sin(2 * np.pi * frequency * samples / 24000).astype(np.float32)
output = self._resample(waveform)
input_rms = np.sqrt(np.mean(waveform**2))
output_rms = np.sqrt(np.mean(output[100:-100] ** 2))
assert output_rms < input_rms * 0.01
def test_rejects_non_integer_downsampling_ratio(self):
with pytest.raises(ValueError, match="integer downsampling ratio"):
StreamingAudioResampler(24000, 16000)
def test_rejects_multichannel_audio(self):
resampler = StreamingAudioResampler(24000, 8000)
with pytest.raises(ValueError, match="only supports mono audio"):
resampler.process(np.zeros((2, 240), dtype=np.float32))
class TestResolveAudioFormat:
"""Test _resolve_audio_format via the serving chat class."""
@@ -2953,8 +2953,10 @@ class TestStreamingResponse:
usage_acc=None,
tts_params=None,
collect=None,
target_sample_rate=None,
):
captured["tts_params"] = tts_params
captured["target_sample_rate"] = target_sample_rate
yield b"\x00\x00"
monkeypatch.setattr(speech_server, "_prepare_speech_generation", prepare)
@@ -2967,6 +2969,7 @@ class TestStreamingResponse:
assert response.status_code == 200
assert captured["tts_params"] is finalized_tts_params
assert captured["target_sample_rate"] is None
def test_sse_streaming(self, streaming_app):
"""stream_format=sse without stream=True returns audio deltas as SSE."""
@@ -3292,6 +3295,28 @@ class TestMergeBatchItem:
assert merged.non_streaming_mode is False
def test_sample_rate_batch_default_used(self):
"""Batch-level sample_rate should be used when the item omits it."""
batch = BatchSpeechRequest(
items=[SpeechBatchItem(input="hi")],
sample_rate=8000,
)
merged = OmniOpenAIServingSpeech._merge_batch_item(batch, batch.items[0])
assert merged.sample_rate == 8000
def test_sample_rate_item_override_wins(self):
"""Per-item sample_rate should override the batch-level default."""
batch = BatchSpeechRequest(
items=[SpeechBatchItem(input="hi", sample_rate=24000)],
sample_rate=8000,
)
merged = OmniOpenAIServingSpeech._merge_batch_item(batch, batch.items[0])
assert merged.sample_rate == 24000
def test_streaming_speech_session_config_accepts_non_streaming_mode():
config = StreamingSpeechSessionConfig(non_streaming_mode=True)
@@ -3299,6 +3324,33 @@ def test_streaming_speech_session_config_accepts_non_streaming_mode():
assert config.non_streaming_mode is True
@pytest.mark.parametrize("sample_rate", [8000, 24000])
def test_speech_requests_accept_supported_sample_rates(sample_rate):
assert OpenAICreateSpeechRequest(input="hello", sample_rate=sample_rate).sample_rate == sample_rate
@pytest.mark.parametrize("sample_rate", [0, -8000])
def test_speech_request_rejects_non_positive_sample_rate(sample_rate):
with pytest.raises(ValidationError):
OpenAICreateSpeechRequest(input="hello", sample_rate=sample_rate)
def test_sample_rate_uses_adapter_capabilities():
server = object.__new__(OmniOpenAIServingSpeech)
request = OpenAICreateSpeechRequest(input="hello", sample_rate=8000)
server._tts_model_type = "qwen3_tts"
assert server._validate_speech_sample_rate(request) is None
request.sample_rate = 16000
assert server._validate_speech_sample_rate(request) == (
"sample_rate=16000 is not supported by the current TTS model; supported rates: 8000, 24000"
)
server._tts_model_type = "fish_speech"
assert server._validate_speech_sample_rate(request) == ("sample_rate is not supported by the current TTS model")
def test_streaming_speech_session_config_scopes_native_speed_to_http():
with pytest.raises(
ValidationError,
@@ -5027,6 +5079,78 @@ class TestTTSAsyncOffloading:
assert "partial_audio" not in payload["error"]
assert "action" not in payload["error"]
@pytest.mark.asyncio
async def test_generate_audio_chunks_resamples_pcm_to_8khz(self, qwen3_tts_server):
async def pcm_generator():
output = create_mock_audio_output_for_test()
output.multimodal_output.update(
audio=torch.linspace(-0.5, 0.5, 2400, dtype=torch.float32),
sr=24000,
)
yield output
chunks = [
chunk
async for chunk in qwen3_tts_server._generate_audio_chunks(
pcm_generator(),
"req-8khz",
target_sample_rate=8000,
)
]
assert len(chunks) >= 2
assert len(b"".join(chunks)) == 800 * 2
@pytest.mark.asyncio
async def test_generate_audio_chunks_wav_header_uses_target_sample_rate(self, qwen3_tts_server):
async def pcm_generator():
output = create_mock_audio_output_for_test()
output.multimodal_output.update(
audio=torch.zeros(2400, dtype=torch.float32),
sr=24000,
)
yield output
chunks = [
chunk
async for chunk in qwen3_tts_server._generate_audio_chunks(
pcm_generator(),
"req-8khz-wav",
response_format="wav",
target_sample_rate=8000,
)
]
assert struct.unpack("<I", chunks[0][24:28])[0] == 8000
@pytest.mark.asyncio
async def test_streaming_resampler_retains_first_chunk_sample_rate_metadata(self, qwen3_tts_server):
async def pcm_generator():
first = create_mock_audio_output_for_test()
first.multimodal_output.update(
audio=torch.zeros(10, dtype=torch.float32),
sr=24000,
)
yield first
second = create_mock_audio_output_for_test()
second.multimodal_output.update(audio=torch.zeros(230, dtype=torch.float32))
second.multimodal_output.pop("sr", None)
yield second
chunks = [
chunk
async for chunk in qwen3_tts_server._generate_audio_chunks(
pcm_generator(),
"req-retained-sr",
response_format="wav",
target_sample_rate=8000,
)
]
assert chunks[0][:4] == b"RIFF"
assert struct.unpack("<I", chunks[0][24:28])[0] == 8000
@pytest.mark.asyncio
async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_close(self, qwen3_tts_server):
async def pcm_generator():
@@ -46,6 +46,16 @@ class TestOmniBaseProfiler:
stage_ids=None,
)
def test_removed_diffusion_batch_size_fails_loudly(self, mock_engine, mocker: MockerFixture):
"""Legacy diffusion_batch_size callers receive an actionable error."""
mocker.patch("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", return_value=mock_engine)
mocker.patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", side_effect=lambda x: x)
mocker.patch("vllm_omni.entrypoints.omni_base.weakref.finalize")
from vllm_omni.entrypoints.omni_base import OmniBase
with pytest.raises(TypeError, match="diffusion_batch_size.*max_num_seqs"):
OmniBase(model="test-model", diffusion_batch_size=8)
def test_start_profile_with_prefix(self, omni_base_instance, mock_engine):
"""Test that start_profile passes profile_prefix to collective_rpc."""
omni_base_instance.start_profile(profile_prefix="test_trace")
+5
View File
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for vllm_omni.entrypoints.utils module."""
import os
@@ -420,6 +423,7 @@ class TestLoadAndResolveStageConfigs:
"""Ensure that kwargs survive default stage creation."""
engine_backend = "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"
kwargs = {"dtype": torch.float32, "engine_backend": engine_backend}
kwargs["revision"] = "pinned-revision"
mocker.patch("vllm_omni.entrypoints.utils.resolve_model_config_path", return_value=None)
mocker.patch("vllm_omni.entrypoints.utils.load_stage_configs_from_model", return_value=([], None))
@@ -433,6 +437,7 @@ class TestLoadAndResolveStageConfigs:
assert len(stage_configs) == 1
assert "dtype" in stage_configs[0]["engine_args"]
assert stage_configs[0]["engine_args"]["engine_backend"] == engine_backend
assert stage_configs[0]["engine_args"]["revision"] == "pinned-revision"
def test_deploy_config_preserves_cli_overrides_and_replicas(self, tmp_path, mocker: MockerFixture):
deploy_path = tmp_path / "qwen3_multi.yaml"
+7
View File
@@ -873,6 +873,13 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any],
elif req_fmt == "wav" and response.audio_format:
assert req_fmt in response.audio_format
expected_sample_rate = request_config.get("sample_rate")
if expected_sample_rate is not None and req_fmt == "wav" and response.audio_bytes:
info = sf.info(io.BytesIO(response.audio_bytes))
assert info.samplerate == int(expected_sample_rate), (
f"Expected sample_rate={expected_sample_rate}, got {info.samplerate}"
)
if run_level in {"advanced_model", "full_model"}:
if req_fmt == "pcm" and response.audio_bytes:
min_hnr_db = float(request_config.get("min_hnr_db", _MIN_PCM_SPEECH_HNR_DB))
+1
View File
@@ -1461,6 +1461,7 @@ class OnlineOmniClient:
"seed",
"instructions",
"speed",
"sample_rate",
"stream_format",
"x_vector_only_mode",
):
+34
View File
@@ -381,6 +381,40 @@ _CI_OVERLAYS: dict[str, dict[str, Any]] = {
},
},
},
"qwen3_omni_moe_colocate_async": {
"base_config": "qwen3_omni_moe_thinking.yaml",
"stages": [
{
"stage_id": 0,
"enable_sleep_mode": True,
"enforce_eager": True,
"trust_remote_code": True,
"tensor_parallel_size": 1,
"devices": "0",
"max_num_seqs": 1,
"max_model_len": 8192,
"max_num_batched_tokens": 8192,
"enable_prefix_caching": False,
"skip_mm_profiling": True,
},
],
"platforms": {
"rocm": {
"stages": [
{
"stage_id": 0,
# This control-plane test needs only 8K tokens. On MI300,
# percentage sizing attempted a ~125 GiB KV allocation
# after the 64.5 GiB model load and then failed while
# allocating the cache. Keep CUDA's existing sizing and
# give only ROCm a bounded cache with ample headroom.
"gpu_memory_utilization": None,
"kv_cache_memory_bytes": 2 * 1024**3,
},
],
},
},
},
"qwen3_omni_moe_multi_replicas_4gpu": {
"base_config": "qwen3_omni_moe.yaml",
"async_chunk": True,
@@ -26,6 +26,7 @@ from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
from vllm_omni.config.stage_config import (
PipelineConfig,
StageExecutionType,
StagePipelineConfig,
_apply_platform_overrides,
load_deploy_config,
merge_pipeline_deploy,
@@ -65,7 +66,9 @@ class TestRegistryDeclaration:
assert not hasattr(pipeline, "max_native_duplex_sessions")
def test_ordinary_pipeline_defaults_to_no_duplex_control(self) -> None:
pipeline = PipelineConfig(model_type="ordinary")
pipeline = PipelineConfig(
model_type="ordinary", stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),)
)
assert pipeline.duplex_control_enabled is False
assert pipeline.duplex_serving_adapter is None
assert not hasattr(pipeline, "max_native_duplex_sessions")
@@ -80,10 +83,6 @@ class TestPipelineTopology:
assert len(pipeline.stages) == 3
assert [s.stage_id for s in pipeline.stages] == [0, 1, 2]
def test_topology_validates(self, pipeline: PipelineConfig) -> None:
# ``validate`` returns a list of structural errors; empty == valid.
assert pipeline.validate() == []
def test_thinker_stage(self, pipeline: PipelineConfig) -> None:
thinker = pipeline.get_stage(0)
assert thinker is not None
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""L1 registration + config-shim tests for NemotronVoiceChat (CPU, no weights)."""
import subprocess
@@ -59,7 +59,6 @@ def test_pipeline_registered_with_three_stages() -> None:
assert code2wav.execution_type == StageExecutionType.LLM_GENERATION
assert code2wav.final_output and code2wav.final_output_type == "audio"
assert code2wav.input_sources == (1,)
assert pipeline.validate() == []
def test_architectures_registered() -> None:
+4 -1
View File
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Tests for TrackingArgumentParser and related utilities."""
import argparse
@@ -35,7 +38,7 @@ _TEST_PIPELINE = PipelineConfig(
stages=(
StagePipelineConfig(stage_id=0, model_stage="stage0"),
StagePipelineConfig(stage_id=1, model_stage="stage1"),
StagePipelineConfig(stage_id=2, model_stage="stage2"),
StagePipelineConfig(stage_id=2, model_stage="stage2", final_output=True),
),
)
@@ -167,6 +167,7 @@ _MODEL_PROMOTE = (
"VLLM_GEPARD_LOOKBACK_FRAMES",
"VLLM_OMNI_FISH_KVCACHE_ATTN",
"VLLM_OMNI_SANA_WM_STAGE1_TEXT_ENCODER",
"VLLM_OMNI_SENSENOVA_PAGED_DECODE",
)
_MODEL_REQUEST_SCOPE = (
+4
View File
@@ -47,6 +47,7 @@ from vllm_omni.config.stage_config import (
build_stage_runtime_overrides,
load_deploy_config,
normalize_pipeline_cli_overrides,
reconcile_diffusion_attention_overrides,
)
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
@@ -1207,6 +1208,9 @@ def _stage_engine_values(
if topology.omni_kv_config:
engine["omni_kv_config"] = _copy_value(topology.omni_kv_config)
if stage_cli_overrides:
if topology.execution_type == StageExecutionType.DIFFUSION:
# Mirror StageConfig.to_omegaconf so both projections resolve alike.
reconcile_diffusion_attention_overrides(engine, stage_cli_overrides)
engine.update(_copy_value(stage_cli_overrides))
_validate_stage_engine_override_ownership(
topology.stage_id,
+4 -9
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Pipeline registry and factory for vllm-omni.
``OMNI_PIPELINES`` maps each ``model_type`` to either a ``PipelineConfig``
@@ -196,18 +196,13 @@ def register_pipeline(pipeline: PipelineConfig | PipelineResolverFunc, model_typ
since resolvers can return multiple different PipelineConfigs depending on the
consumed config.
"""
errors: list[str] = []
if isinstance(pipeline, PipelineConfig):
errors = pipeline.validate()
model_type = model_type if model_type is not None else pipeline.model_type
else:
if model_type is None:
raise ValueError("Model type must be explicitly provided when registering a pipeline resolver")
elif model_type is None:
raise ValueError("Model type must be explicitly provided when registering a pipeline resolver")
if model_type in OMNI_PIPELINES:
errors.append(f"Model type {model_type} is already registered; the old mapping will be clobbered")
if errors:
logger.warning("Registration for pipeline of type %s produced the following issues: %s", model_type, errors)
logger.warning(f"Model type {model_type} is already registered; the old mapping will be clobbered")
OMNI_PIPELINES[model_type] = pipeline
+36 -2
View File
@@ -7,7 +7,7 @@ from __future__ import annotations
import functools
import re
import warnings
from collections.abc import Callable
from collections.abc import Callable, Mapping
from dataclasses import asdict, dataclass, field, fields
from enum import Enum
from pathlib import Path
@@ -161,6 +161,28 @@ def _apply_diffusion_parallel_runtime_overrides(
engine_args["parallel_config"] = parallel_config_dict
def reconcile_diffusion_attention_overrides(
engine_args: dict[str, Any],
runtime_overrides: Mapping[str, Any],
) -> None:
"""Apply CLI precedence across the two diffusion attention representations.
``diffusion_attention_backend`` and ``diffusion_attention_config.default``
express the same setting and are rejected downstream when both are set, so
a CLI value in one form replaces the YAML value in the other.
"""
if runtime_overrides.get("diffusion_attention_backend") is not None:
yaml_config = engine_args.get("diffusion_attention_config")
if isinstance(yaml_config, Mapping) and yaml_config.get("default") is not None:
remaining = {k: v for k, v in yaml_config.items() if k != "default"}
if remaining:
engine_args["diffusion_attention_config"] = remaining
else:
engine_args.pop("diffusion_attention_config")
if runtime_overrides.get("diffusion_attention_config") is not None:
engine_args.pop("diffusion_attention_backend", None)
class StageType(str, Enum):
"""Type of processing stage in the Omni pipeline."""
@@ -302,6 +324,17 @@ class PipelineConfig:
# Global CLI spelling -> (stage id, stage-local spelling).
stage_cli_aliases: dict[str, tuple[int, str]] = field(default_factory=dict)
def __post_init__(self) -> None:
errors = self.get_validation_errors()
if errors:
# If we have multiple errors, put them on separate indented lines for readability
multi_sep = "\n\t"
error_str = multi_sep.join(errors)
if len(errors) > 1:
error_str = f"{multi_sep}{error_str}"
raise ValueError(f"PipelineConfig initialization failed with the following error(s): {error_str}")
def get_stage(self, stage_id: int) -> StagePipelineConfig | None:
"""Look up a stage by its ID."""
for stage in self.stages:
@@ -309,7 +342,7 @@ class PipelineConfig:
return stage
return None
def validate(self) -> list[str]:
def get_validation_errors(self) -> list[str]:
"""Return list of topology errors (empty if valid)."""
errors: list[str] = []
if not self.stages:
@@ -1081,6 +1114,7 @@ class StageConfig:
if StageType(self.stage_type) == StageType.DIFFUSION:
_apply_diffusion_parallel_runtime_overrides(engine_args, runtime_overrides)
reconcile_diffusion_attention_overrides(engine_args, runtime_overrides)
# CLI overrides take precedence over YAML defaults
for key, value in runtime_overrides.items():
-1
View File
@@ -1,7 +1,6 @@
# Wan2.2 TI2V-5B deploy.
#
# Single GPU by default. Increase tensor_parallel_size and devices for A14B MoE variants.
# Request-level batching uses scheduler max_num_seqs, not diffusion_batch_size.
async_chunk: false
trust_remote_code: true
distributed_executor_backend: mp
@@ -48,7 +48,6 @@ class InlineStageDiffusionClient(StageClientBase):
model: str,
od_config: OmniDiffusionConfig,
metadata: StageMetadata,
batch_size: int = 1,
) -> None:
self.model = model
self.od_config = od_config
@@ -63,7 +62,6 @@ class InlineStageDiffusionClient(StageClientBase):
self.prompt_transform_func = None
self.prompt_expand_func = None
self.engine_input_source = metadata.engine_input_source
self.batch_size = batch_size
self._enrich_config()
self._engine = DiffusionEngine.make_engine(self.od_config)
@@ -79,10 +77,10 @@ class InlineStageDiffusionClient(StageClientBase):
self._engine.executor.register_failure_callback(self._mark_engine_dead)
logger.info(
"[InlineStageDiffusionClient] stage-%s [rep-%s] initialized inline (batch_size=%d)",
"[InlineStageDiffusionClient] stage-%s [rep-%s] initialized inline (max_num_seqs=%d)",
self.stage_id,
self.replica_id,
self.batch_size,
self.od_config.max_num_seqs,
)
def _enrich_config(self) -> None:
+14 -3
View File
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import os
from collections import OrderedDict, defaultdict
from collections.abc import Callable
@@ -41,7 +44,7 @@ def _prepare_lora_delta(
lora_b_suffix: str = "lora_B.weight",
lora_bias_suffix: str = "bias",
):
used_keys = set()
used_keys: set[str] = set()
# stacked_params_mapping param_to_weight_names
# [(".to_qkv", ".to_q.", "q")
# (".to_qkv", ".to_k.", "k") ========> {".to_qkv": [".to_q", ".to_k", ".to_v"]}
@@ -54,7 +57,10 @@ def _prepare_lora_delta(
if param_to_weight_names is None:
param_to_weight_names = defaultdict(list)
for param_name, weight_names in param_to_weight_names.items():
if param_name not in base_key:
# Fused names can overlap by prefix (`.qkv_proj` is a substring of
# `.qkv_proj_mot_gen`), and a substring match stacks both mappings into
# one delta of twice the height.
if not base_key.endswith(param_name):
continue
is_stacked_param = True
# handle lora_a_key and lora_b_key together
@@ -78,7 +84,7 @@ def _prepare_lora_delta(
stacked_deltas.append(delta)
used_keys.add(lora_a_key)
used_keys.add(lora_b_key)
continue
break
if is_stacked_param:
return torch.concat(stacked_deltas), used_keys
@@ -200,6 +206,10 @@ def _apply_diffusers_lora_alpha_scaling(state_dict: dict[str, torch.Tensor]) ->
class LoraLoaderMixin:
transformer_name = "transformer"
# Set by the pipeline this mixin is combined with; annotated, not assigned.
transformer: torch.nn.Module
_lora_loaded: dict[str | None, dict[str, torch.Tensor]]
# Lazy initialization to avoid MRO issues: __init__ may not be called
# when mixin with nn.Module
@property
@@ -373,6 +383,7 @@ class QwenImageLoraLoaderMixin(LoraLoaderMixin):
class WanLoraLoaderMixin(LoraLoaderMixin):
transformer_2_name = "transformer_2"
has_transformer_2: bool
def load_lora_weights(
self,
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Native vLLM-Omni pipeline for Boogu-Image-0.1.
@@ -220,7 +220,7 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery
DiffusersPipelineLoader.ComponentSource(
model_or_path=od_config.model,
subfolder="transformer",
revision=None,
revision=od_config.revision,
prefix="transformer.",
fall_back_to_pt=True,
)
@@ -233,10 +233,18 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery
# See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
# race; prefetch the whole component set before any from_pretrained.
boogu_subfolders = ["scheduler", "vae", "mllm", "processor"]
prefetch_subfolders(model, boogu_subfolders, local_files_only=local_files_only)
prefetch_subfolders(
model,
boogu_subfolders,
local_files_only=local_files_only,
revision=od_config.revision,
)
self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
model, subfolder="scheduler", local_files_only=local_files_only
model,
subfolder="scheduler",
local_files_only=local_files_only,
revision=od_config.revision,
)
mllm = from_pretrained_with_prefetch(
@@ -246,6 +254,7 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery
prefetch_list=boogu_subfolders,
local_files_only=local_files_only,
torch_dtype=od_config.dtype,
revision=od_config.revision,
)
# Upstream reuses the full VLM as an optional instruction rewriter and
# encodes with its inner model (no ``lm_head``); the rewriter is not
@@ -255,7 +264,10 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery
self.mllm = mllm.to(self._execution_device)
self.processor = Qwen3VLProcessor.from_pretrained(
model, subfolder="processor", local_files_only=local_files_only
model,
subfolder="processor",
local_files_only=local_files_only,
revision=od_config.revision,
)
self.vae = from_pretrained_with_prefetch(
@@ -264,6 +276,7 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery
subfolder="vae",
prefetch_list=boogu_subfolders,
local_files_only=local_files_only,
revision=od_config.revision,
).to(self._execution_device)
self.transformer = BooguImageTransformer2DModel(od_config=od_config)
@@ -2223,7 +2223,7 @@ class Cosmos3OmniDiffusersPipeline(
if prompt_suffix:
prompt = f"{prompt.rstrip()} {prompt_suffix.lstrip()}".strip()
if _is_rank_zero():
logger.info("Final prompt: '%s'", prompt)
logger.debug("Final prompt: '%s'", prompt)
if negative_metadata_mode == "none":
negative_dur_tmpl = None
@@ -0,0 +1,391 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Model-local coordination for MiniMax H3 temporal chunk decoding.
This boundary owns released-code compatibility and reconciles preflight and
callback failures at fixed VAE-group rendezvous points. Cropping,
representation conversion, transport, and encoding remain runtime
responsibilities.
"""
from __future__ import annotations
from contextlib import AbstractContextManager
from typing import Protocol
import torch
import torch.distributed as dist
import torch.nn as nn
from .temporal_chunks import (
MiniMaxH3TemporalChunkCompatibilityError,
MiniMaxH3TemporalDecodePlan,
_TemporalChunkCallbackSignature,
decode_minimax_h3_temporal_chunks_compat,
prepare_minimax_h3_temporal_chunks_compat,
)
class MiniMaxH3VideoChunkCallback(_TemporalChunkCallbackSignature, Protocol):
"""Consume one ordered, temporally committed H3 video chunk.
``frames`` is an owned, contiguous float32 RGB tensor in ``BCTHW`` layout
and the ``[0, 1]`` value range. The callback runs synchronously on the
producer's current device stream. Before reading it on another stream, the
consumer must establish ordering from the producer stream (for example via
an event) and retain the allocation for that stream's lifetime. The
consumer may retain or mutate the tensor; it does not alias the returned
complete decode.
"""
class MiniMaxH3ChunkedDecodeUnsupportedError(RuntimeError):
"""Raised when loaded MiniMax H3 code cannot honor the chunk contract."""
class MiniMaxH3ChunkCallbackPeerError(RuntimeError):
"""Raised on peer VAE ranks when the output-rank callback fails."""
class _MiniMaxH3ChunkDecodeHost(Protocol):
"""Narrow VAE capabilities required by the chunk coordinator."""
model: nn.Module
def _denormalize_latent(self, latent: torch.Tensor) -> torch.Tensor: ...
def _decode_tiling_context(
self,
latent: torch.Tensor,
) -> AbstractContextManager: ...
def _normalize_decoded_frames(
self,
decoded: torch.Tensor,
) -> torch.Tensor: ...
class _MiniMaxH3ChunkDecodeCoordinator:
"""Run source-gated chunk decode with group-local failure semantics.
The coordinator is a plain Python object. It caches only successful source
validation; request callbacks, tensors, metadata, and process groups remain
local to :meth:`decode`.
Errors raised inside the checkpoint decoder after native collectives begin
cannot be reconciled here: another rank may already be blocked inside that
collective. Resolving such a failure requires native collective handling or
supervisor-level timeout and worker-group teardown; the process group cannot
be assumed reusable.
"""
def __init__(self) -> None:
self._sources_validated = False
@staticmethod
def _build_execution_signature(
latent: torch.Tensor,
plan: MiniMaxH3TemporalDecodePlan,
) -> tuple[int, ...]:
"""Describe every shape and plan field used by the temporal schedule."""
dtype_codes: dict[torch.dtype | None, int] = {
None: 0,
torch.float16: 1,
torch.bfloat16: 2,
torch.float32: 3,
torch.float64: 4,
}
return (
# Keep the layout versioned so adding a field cannot silently make
# an older and newer worker appear compatible.
1,
*[int(size) for size in latent.shape],
dtype_codes.get(latent.dtype, -1),
*[int(size) for size in plan.latent.shape],
dtype_codes.get(plan.latent.dtype, -1),
int(plan.num_chunks),
int(plan.total_frames),
int(plan.pad_frames),
int(plan.output_frames),
len(plan.chunk_frames),
*[int(frame_count) for frame_count in plan.chunk_frames],
dtype_codes.get(plan.temporal_dtype, -1),
)
@staticmethod
def _is_output_rank(
chunk_callback: MiniMaxH3VideoChunkCallback | None,
*,
group: dist.ProcessGroup | None,
device: torch.device,
) -> bool:
if group is None:
return chunk_callback is not None
group_rank = dist.get_rank(group)
# A valid configuration has one contribution: rank 0 contributes 1.
# A non-zero owner contributes at least 2, and multiple owners sum to
# at least 3, so every rank makes the same decision from one reduction.
callback_owner = torch.tensor(
[group_rank + 1 if chunk_callback is not None else 0],
dtype=torch.int64,
device=device,
)
dist.all_reduce(callback_owner, op=dist.ReduceOp.SUM, group=group)
if int(callback_owner.item()) != 1:
raise ValueError("MiniMax-H3 chunk decode requires a callback on exactly VAE group rank 0")
return group_rank == 0
@staticmethod
def _synchronize_callback_error(
callback_error: BaseException | None,
*,
group: dist.ProcessGroup | None,
device: torch.device,
) -> None:
if group is None:
if callback_error is not None:
raise callback_error
return
is_output_rank = dist.get_rank(group) == 0
failed = torch.tensor(
[int(is_output_rank and callback_error is not None)],
dtype=torch.int32,
device=device,
)
dist.broadcast(
failed,
src=dist.get_global_rank(group, 0),
group=group,
)
if int(failed.item()) == 0:
return
if is_output_rank:
assert callback_error is not None
raise callback_error
raise MiniMaxH3ChunkCallbackPeerError("MiniMax-H3 video chunk callback failed on this VAE group's output rank")
@staticmethod
def _synchronize_preflight_error(
preflight_error: Exception | None,
temporal_dtype: torch.dtype | None,
execution_signature: tuple[int, ...] | None,
*,
group: dist.ProcessGroup | None,
device: torch.device,
) -> None:
if group is None:
if preflight_error is not None:
raise preflight_error
return
failed = torch.tensor(
[int(preflight_error is not None)],
dtype=torch.int32,
device=device,
)
dist.all_reduce(failed, op=dist.ReduceOp.MAX, group=group)
if int(failed.item()) == 0:
dtype_codes = {
None: 0,
torch.float16: 1,
torch.bfloat16: 2,
torch.float32: 3,
}
code = torch.tensor(
[dtype_codes.get(temporal_dtype, -1)],
dtype=torch.int32,
device=device,
)
lowest = code.clone()
highest = code.clone()
dist.all_reduce(lowest, op=dist.ReduceOp.MIN, group=group)
dist.all_reduce(highest, op=dist.ReduceOp.MAX, group=group)
if int(lowest.item()) == int(highest.item()) and int(lowest.item()) >= 0:
if execution_signature is None:
raise RuntimeError("MiniMax-H3 temporal decode preflight produced no execution signature")
# The plan length is shape-dependent, so first agree on a
# common tensor size and then compare the complete signatures.
# This keeps all ranks in tensor collectives and catches a
# schedule mismatch before entering the decoder's collectives.
signature_length = torch.tensor(
[len(execution_signature)],
dtype=torch.int64,
device=device,
)
max_signature_length = signature_length.clone()
dist.all_reduce(
max_signature_length,
op=dist.ReduceOp.MAX,
group=group,
)
padded_signature = torch.full(
(int(max_signature_length.item()),),
-2,
dtype=torch.int64,
device=device,
)
padded_signature[: len(execution_signature)] = torch.tensor(
execution_signature,
dtype=torch.int64,
device=device,
)
gathered_signatures = [torch.empty_like(padded_signature) for _ in range(dist.get_world_size(group))]
dist.all_gather(gathered_signatures, padded_signature, group=group)
if any(not torch.equal(padded_signature, peer) for peer in gathered_signatures):
raise MiniMaxH3ChunkedDecodeUnsupportedError(
"MiniMax-H3 VAE ranks built incompatible temporal decode plans"
)
return
raise MiniMaxH3ChunkedDecodeUnsupportedError(
"MiniMax-H3 VAE ranks selected incompatible temporal concat dtypes"
)
if preflight_error is not None:
raise preflight_error
raise MiniMaxH3ChunkedDecodeUnsupportedError("A peer VAE rank rejected the MiniMax-H3 temporal chunk preflight")
def decode(
self,
host: _MiniMaxH3ChunkDecodeHost,
latent: torch.Tensor,
chunk_callback: MiniMaxH3VideoChunkCallback | None,
*,
group: dist.ProcessGroup | None,
) -> torch.Tensor:
"""Decode one latent while publishing committed normalized chunks."""
plan: MiniMaxH3TemporalDecodePlan | None = None
preflight_error: Exception | None = None
if latent.ndim != 5:
preflight_error = MiniMaxH3ChunkedDecodeUnsupportedError(
f"MiniMax-H3 video latent must be rank 5, got {tuple(latent.shape)}"
)
else:
try:
denormalized = host._denormalize_latent(latent)
plan = prepare_minimax_h3_temporal_chunks_compat(
host.model,
denormalized,
validate_sources=not self._sources_validated,
)
self._sources_validated = True
del denormalized
except MiniMaxH3TemporalChunkCompatibilityError as exc:
preflight_error = MiniMaxH3ChunkedDecodeUnsupportedError(str(exc))
except Exception as exc: # noqa: BLE001
# Every VAE rank must reach the preflight reduction even when
# one node's private planner or environment fails locally.
exc.__traceback__ = None
preflight_error = exc
self._synchronize_preflight_error(
preflight_error,
plan.temporal_dtype if plan is not None else None,
self._build_execution_signature(latent, plan) if plan is not None else None,
group=group,
device=latent.device,
)
assert plan is not None
is_output_rank = self._is_output_rank(
chunk_callback,
group=group,
device=latent.device,
)
callback_error: BaseException | None = None
next_chunk_index = 0
expected_total_chunks: int | None = None
next_frame_start = 0
saw_final = False
def publish(
decoded_chunk: torch.Tensor,
*,
chunk_index: int,
total_chunks: int,
frame_start: int,
is_final: bool,
) -> None:
nonlocal callback_error, expected_total_chunks
nonlocal next_chunk_index, next_frame_start, saw_final
if not is_output_rank or callback_error is not None:
return
try:
if expected_total_chunks is None:
expected_total_chunks = total_chunks
if (
saw_final
or total_chunks <= 0
or total_chunks != expected_total_chunks
or chunk_index != next_chunk_index
or chunk_index >= total_chunks
or frame_start != next_frame_start
or is_final != (chunk_index + 1 == total_chunks)
):
raise RuntimeError(
"MiniMax-H3 temporal backend emitted discontinuous "
"chunk metadata: "
f"index={chunk_index}/{next_chunk_index}, "
f"total={total_chunks}/{expected_total_chunks}, "
f"frame_start={frame_start}/{next_frame_start}, "
f"final={is_final}, after_final={saw_final}"
)
frames = host._normalize_decoded_frames(decoded_chunk)
frames = frames.clone(memory_format=torch.contiguous_format)
frame_count = int(frames.shape[2])
if frame_count <= 0:
raise RuntimeError("MiniMax-H3 temporal backend emitted an empty chunk")
next_chunk_index += 1
next_frame_start += frame_count
saw_final = is_final
assert chunk_callback is not None
chunk_callback(
frames,
chunk_index=chunk_index,
total_chunks=total_chunks,
frame_start=frame_start,
is_final=is_final,
)
except BaseException as exc: # noqa: BLE001
# Keep a failed GPU chunk and its callback locals out of the
# traceback while peer ranks finish the remaining collectives.
exc.__traceback__ = None
callback_error = exc
# Do not catch and reduce exceptions from this region. A peer may
# already be blocked in one of the checkpoint's native tile
# collectives, in which case a new adapter-level collective would also
# deadlock. A native backend or supervisor must handle timeout and
# worker-group teardown.
with host._decode_tiling_context(latent):
decoded = decode_minimax_h3_temporal_chunks_compat(
host.model,
plan,
publish,
)
frames = host._normalize_decoded_frames(decoded)
if (
is_output_rank
and callback_error is None
and (not saw_final or next_chunk_index != expected_total_chunks or next_frame_start != int(frames.shape[2]))
):
callback_error = RuntimeError(
"MiniMax-H3 temporal backend did not reconstruct the full "
"decode: "
f"final={saw_final}, "
f"chunks={next_chunk_index}/{expected_total_chunks}, "
f"frames={next_frame_start}/{frames.shape[2]}"
)
self._synchronize_callback_error(
callback_error,
group=group,
device=frames.device,
)
return frames
__all__ = [
"MiniMaxH3ChunkCallbackPeerError",
"MiniMaxH3ChunkedDecodeUnsupportedError",
"MiniMaxH3VideoChunkCallback",
]
@@ -0,0 +1,414 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Source-gated temporal chunk compatibility for the released MiniMax H3 VAE."""
from __future__ import annotations
import hashlib
import importlib
import inspect
from dataclasses import dataclass
from pathlib import Path
from typing import Protocol
import torch
import torch.nn as nn
MINIMAX_H3_RELEASED_REVISION = "42ed227ee7df40d41602854ae760620d6eb651fe"
# Compatibility pins detect upstream code drift; they are not a trust or
# security boundary for remote code execution.
MINIMAX_H3_RELEASED_BUNDLE_SHA256 = "30e64afbaa940696bf8bfe2c86003a4b1666c84b22fa3aa12403a3e6a03f705d"
_RELEASED_TEMPORAL_CONFIG = {
"clip_length": 17,
"tokens_chunk_size": 5,
"token_overlap": 2,
"token_drop": 3,
"vae_ratio_t": 4,
"frame_pre_padding": 3,
"frame_overlap": 5,
"isolated_first_frame": False,
"isolated_last_frame": False,
}
class MiniMaxH3TemporalChunkCompatibilityError(RuntimeError):
"""Raised when remote H3 code cannot use the local compatibility path."""
class _TemporalChunkCallbackSignature(Protocol):
"""Shared callable shape; each producer defines tensor semantics."""
def __call__(
self,
frames: torch.Tensor,
/,
*,
chunk_index: int,
total_chunks: int,
frame_start: int,
is_final: bool,
) -> None: ...
@dataclass(slots=True)
class MiniMaxH3TemporalDecodePlan:
latent: torch.Tensor
num_chunks: int
total_frames: int
pad_frames: int
output_frames: int
chunk_frames: tuple[int, ...]
temporal_dtype: torch.dtype | None
class _RawChunkAssembly:
def __init__(
self,
*,
callback: _TemporalChunkCallbackSignature,
plan: MiniMaxH3TemporalDecodePlan,
) -> None:
self.callback = callback
self.plan = plan
self.frames: torch.Tensor | None = None
self.write_pos = 0
self.logical_frames = 0
self.dropped_frames = 0
self.chunk_index = 0
self.pending_final: tuple[torch.Tensor, int, int] | None = None
self.contract_error: RuntimeError | None = None
def write(self, part: torch.Tensor) -> None:
part_frames = int(part.shape[2])
if part_frames <= 0:
return
self.logical_frames += part_frames
remaining = self.plan.output_frames - self.write_pos
copy_frames = min(part_frames, max(0, remaining))
if copy_frames <= 0:
self.dropped_frames += part_frames
return
if self.chunk_index >= len(self.plan.chunk_frames):
self.contract_error = RuntimeError("MiniMax-H3 compatibility decode emitted an extra chunk")
else:
expected_frames = self.plan.chunk_frames[self.chunk_index]
if copy_frames != expected_frames and self.contract_error is None:
self.contract_error = RuntimeError(
"MiniMax-H3 compatibility chunk size mismatch: "
f"chunk={self.chunk_index}, frames={copy_frames}/{expected_frames}"
)
emitted = part[:, :, :copy_frames]
if self.frames is None:
output_shape = list(emitted.shape)
output_shape[2] = self.plan.output_frames
self.frames = torch.empty(
output_shape,
dtype=emitted.dtype,
device=emitted.device,
)
frame_start = self.write_pos
frame_stop = frame_start + copy_frames
self.frames[:, :, frame_start:frame_stop].copy_(emitted)
if frame_stop == self.plan.output_frames:
self.pending_final = (emitted, self.chunk_index, frame_start)
elif self.contract_error is None:
self.callback(
emitted,
chunk_index=self.chunk_index,
total_chunks=len(self.plan.chunk_frames),
frame_start=frame_start,
is_final=False,
)
self.chunk_index += 1
self.write_pos = frame_stop
self.dropped_frames += part_frames - copy_frames
def finish(self) -> torch.Tensor:
if self.contract_error is not None:
raise self.contract_error
if self.frames is None or self.pending_final is None:
raise RuntimeError("MiniMax-H3 compatibility decode produced no terminal chunk")
if (
self.logical_frames != self.plan.total_frames
or self.dropped_frames != self.plan.pad_frames
or self.write_pos != self.plan.output_frames
or self.chunk_index != len(self.plan.chunk_frames)
):
raise RuntimeError(
"MiniMax-H3 compatibility frame plan mismatch: "
f"logical={self.logical_frames}/{self.plan.total_frames}, "
f"dropped={self.dropped_frames}/{self.plan.pad_frames}, "
f"written={self.write_pos}/{self.plan.output_frames}, "
f"chunks={self.chunk_index}/{len(self.plan.chunk_frames)}"
)
final_frames, final_index, final_start = self.pending_final
self.callback(
final_frames,
chunk_index=final_index,
total_chunks=len(self.plan.chunk_frames),
frame_start=final_start,
is_final=True,
)
return self.frames
def _source_bundle_sha256(model: nn.Module) -> tuple[str | None, str | None]:
try:
source = inspect.getsourcefile(type(model))
except TypeError:
return None, None
if source is None:
return None, None
package_dir = Path(source).parent
files = sorted(path for path in package_dir.glob("*.py") if path.name != "__init__.py")
digest = hashlib.sha256()
try:
for path in files:
digest.update(path.name.encode())
digest.update(b"\0")
digest.update(path.read_bytes())
digest.update(b"\0")
except OSError:
return str(package_dir), None
return str(package_dir), digest.hexdigest()
def _validate_released_sources(model: nn.Module) -> None:
source_dir, actual = _source_bundle_sha256(model)
if actual != MINIMAX_H3_RELEASED_BUNDLE_SHA256:
raise MiniMaxH3TemporalChunkCompatibilityError(
"The vLLM-owned MiniMax-H3 chunk adapter supports only remote code "
f"from MiniMaxAI/MiniMax-H3@{MINIMAX_H3_RELEASED_REVISION}: "
f"bundle={actual!r} ({source_dir}). "
"The normal complete-output decode remains available."
)
def _validate_released_config(model: nn.Module) -> None:
actual = {name: getattr(model, name, None) for name in _RELEASED_TEMPORAL_CONFIG}
if actual != _RELEASED_TEMPORAL_CONFIG:
raise MiniMaxH3TemporalChunkCompatibilityError(
"The vLLM-owned MiniMax-H3 chunk adapter requires the released temporal "
f"configuration {_RELEASED_TEMPORAL_CONFIG}, found {actual}. "
"The normal complete-output decode remains available."
)
def _validate_compat_contract(model: nn.Module, latent: torch.Tensor) -> None:
if latent.ndim != 5 or int(latent.shape[2]) <= 0:
raise MiniMaxH3TemporalChunkCompatibilityError(
f"MiniMax-H3 video latent must have positive rank-5 time, got {tuple(latent.shape)}"
)
if not bool(getattr(model, "use_3d_conv", False)):
raise MiniMaxH3TemporalChunkCompatibilityError("MiniMax-H3 temporal chunk decode requires the 3D decoder")
if bool(getattr(model, "training", False)):
raise MiniMaxH3TemporalChunkCompatibilityError("MiniMax-H3 temporal chunk decode is inference-only")
required_methods = (
"_adaptive_decode",
"_decode_temporal_output_frame_plan",
"_decode_temporal_pad_frames",
"blend",
)
missing = [name for name in required_methods if not callable(getattr(model, name, None))]
if missing:
raise MiniMaxH3TemporalChunkCompatibilityError(
"Loaded MiniMax-H3 VAE lacks required released-code methods: " + ", ".join(missing)
)
def _truncate_chunk_frames(
logical_chunk_frames: list[int],
output_frames: int,
) -> tuple[int, ...]:
remaining = output_frames
result: list[int] = []
for logical_frames in logical_chunk_frames:
emitted_frames = min(logical_frames, max(0, remaining))
if emitted_frames > 0:
result.append(emitted_frames)
remaining -= emitted_frames
if remaining != 0:
raise RuntimeError(f"MiniMax-H3 compatibility chunk plan cannot cover output: remaining={remaining}")
return tuple(result)
def _build_decode_plan(
model: nn.Module,
latent: torch.Tensor,
temporal_dtype: torch.dtype | None,
) -> MiniMaxH3TemporalDecodePlan:
pseudo_total_tokens = int(latent.shape[2]) + int(model.token_drop) # type: ignore[attr-defined]
tokens_chunk_size = int(model.tokens_chunk_size) # type: ignore[attr-defined]
remainder = pseudo_total_tokens % tokens_chunk_size
pad_tokens = 0 if remainder == 0 else tokens_chunk_size - remainder
pseudo_total_tokens += pad_tokens
num_chunks = pseudo_total_tokens // tokens_chunk_size - int(model.token_drop > 0) # type: ignore[attr-defined]
if num_chunks <= 0:
raise MiniMaxH3TemporalChunkCompatibilityError(
f"MiniMax-H3 temporal chunk plan is empty for latent T={latent.shape[2]}"
)
temporal_latent = latent
if pad_tokens:
pad = temporal_latent[:, :, -1:].repeat(1, 1, pad_tokens, 1, 1)
temporal_latent = torch.cat((temporal_latent, pad), dim=2)
total_frames, pad_frames, output_frames = model._decode_temporal_output_frame_plan( # type: ignore[attr-defined]
temporal_latent,
None,
None,
num_chunks,
pad_tokens,
)
if output_frames <= 0:
raise MiniMaxH3TemporalChunkCompatibilityError(
f"MiniMax-H3 temporal chunk plan has {output_frames} output frames"
)
chunk_decoded_frames = tokens_chunk_size * int(model.vae_ratio_t) # type: ignore[attr-defined]
split_count = int(model.token_drop > 0) + 1 # type: ignore[attr-defined]
logical_chunk_frames: list[int] = []
final_overlap_frames = 0
for temporal_index in range(num_chunks):
token_start = temporal_index * tokens_chunk_size
token_stop = token_start + tokens_chunk_size + int(model.token_overlap) # type: ignore[attr-defined]
clip_tokens = max(
0,
min(token_stop, int(temporal_latent.shape[2])) - min(token_start, int(temporal_latent.shape[2])),
)
clip_frames = clip_tokens * int(model.vae_ratio_t) # type: ignore[attr-defined]
for split_index in range(split_count):
frame_start = split_index * chunk_decoded_frames
frame_stop = min(frame_start + chunk_decoded_frames, clip_frames)
part_frames = max(
0,
frame_stop - frame_start - int(model.frame_pre_padding), # type: ignore[attr-defined]
)
if split_index == 0:
if part_frames > 0:
logical_chunk_frames.append(part_frames)
else:
final_overlap_frames = part_frames
if final_overlap_frames > 0:
logical_chunk_frames.append(final_overlap_frames)
if sum(logical_chunk_frames) != int(total_frames):
raise RuntimeError(
"MiniMax-H3 compatibility planner disagrees with released remote code: "
f"frames={sum(logical_chunk_frames)}/{total_frames}"
)
chunk_frames = _truncate_chunk_frames(logical_chunk_frames, int(output_frames))
return MiniMaxH3TemporalDecodePlan(
latent=temporal_latent,
num_chunks=num_chunks,
total_frames=int(total_frames),
pad_frames=int(pad_frames),
output_frames=int(output_frames),
chunk_frames=chunk_frames,
temporal_dtype=temporal_dtype,
)
def resolve_minimax_h3_temporal_cat_dtype(model: nn.Module) -> torch.dtype | None:
module = importlib.import_module(type(model).__module__)
resolver = getattr(module, "_resolve_temporal_cat_dtype", None)
if resolver is None:
return None
if not callable(resolver):
raise MiniMaxH3TemporalChunkCompatibilityError(
"Loaded MiniMax-H3 VAE has a non-callable temporal dtype resolver"
)
try:
dtype = resolver()
except ValueError as exc:
raise MiniMaxH3TemporalChunkCompatibilityError(str(exc)) from None
if dtype is not None and not isinstance(dtype, torch.dtype):
raise MiniMaxH3TemporalChunkCompatibilityError(
f"Loaded MiniMax-H3 VAE returned invalid temporal dtype {dtype!r}"
)
return dtype
def decode_minimax_h3_temporal_chunks_compat(
model: nn.Module,
plan: MiniMaxH3TemporalDecodePlan,
callback: _TemporalChunkCallbackSignature,
) -> torch.Tensor:
"""Mirror the released Apache-2.0 H3 temporal assembler with callbacks.
Keep this implementation and the source/config fingerprints above in sync
with ``MiniMaxAI/MiniMax-H3@42ed227e``. Unknown remote code fails closed
before this function is called.
Callback tensors are borrowed decoder-domain views and must not be retained
or mutated. The high-level adapter owns normalization and storage isolation.
"""
chunk_decoded_frames = int(model.tokens_chunk_size) * int(model.vae_ratio_t) # type: ignore[attr-defined]
split_count = int(model.token_drop > 0) + 1 # type: ignore[attr-defined]
assembly = _RawChunkAssembly(callback=callback, plan=plan)
overlap: torch.Tensor | None = None
for temporal_index in range(plan.num_chunks):
token_start = temporal_index * int(model.tokens_chunk_size) # type: ignore[attr-defined]
token_stop = token_start + int(model.tokens_chunk_size) + int(model.token_overlap) # type: ignore[attr-defined]
chunk_latent = plan.latent[:, :, token_start:token_stop]
decoded = model._adaptive_decode(chunk_latent) # type: ignore[attr-defined]
if plan.temporal_dtype is not None and decoded.dtype != plan.temporal_dtype:
decoded = decoded.to(plan.temporal_dtype)
if decoded.device != plan.latent.device:
decoded = decoded.to(plan.latent.device)
for split_index in range(split_count):
frame_start = split_index * chunk_decoded_frames
frame_stop = min(frame_start + chunk_decoded_frames, int(decoded.shape[2]))
part = decoded[:, :, frame_start:frame_stop]
part = part[:, :, int(model.frame_pre_padding) :] # type: ignore[attr-defined]
if split_index == 0:
if overlap is not None:
part = model.blend( # type: ignore[attr-defined]
overlap,
part,
int(model.frame_overlap), # type: ignore[attr-defined]
dim=-3,
)
overlap = None
assembly.write(part)
else:
overlap = part.contiguous()
del part
if temporal_index == plan.num_chunks - 1 and overlap is not None:
assembly.write(overlap)
overlap = None
del decoded, chunk_latent
return assembly.finish()
def prepare_minimax_h3_temporal_chunks_compat(
model: nn.Module,
latent: torch.Tensor,
*,
validate_sources: bool = True,
) -> MiniMaxH3TemporalDecodePlan:
"""Build a validated plan without entering decoder collectives."""
if validate_sources:
_validate_released_sources(model)
_validate_released_config(model)
_validate_compat_contract(model, latent)
temporal_dtype = resolve_minimax_h3_temporal_cat_dtype(model)
return _build_decode_plan(model, latent, temporal_dtype)
__all__ = [
"MINIMAX_H3_RELEASED_REVISION",
"MINIMAX_H3_RELEASED_BUNDLE_SHA256",
"MiniMaxH3TemporalDecodePlan",
"MiniMaxH3TemporalChunkCompatibilityError",
"decode_minimax_h3_temporal_chunks_compat",
"prepare_minimax_h3_temporal_chunks_compat",
"resolve_minimax_h3_temporal_cat_dtype",
]
+105 -39
View File
@@ -21,12 +21,22 @@ from vllm.logger import init_logger
from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
DistributedVaeMixin,
)
from vllm_omni.diffusion.distributed.parallel_state import get_world_group
from vllm_omni.diffusion.distributed.parallel_state import (
get_data_parallel_world_size,
get_world_group,
model_parallel_is_initialized,
)
from vllm_omni.diffusion.offloader.module_residency import (
BoundedAllocatorCache,
PinnedModuleStager,
)
from .chunked_decode import (
MiniMaxH3ChunkCallbackPeerError,
MiniMaxH3ChunkedDecodeUnsupportedError,
MiniMaxH3VideoChunkCallback,
_MiniMaxH3ChunkDecodeCoordinator,
)
from .ops import install_h3_vae_optimizations
from .packed_tokens import minimax_h3_patchify_video_latent
@@ -153,6 +163,7 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
self.use_slicing = False
self.parallel_size = 1
self.device_module = torch.get_device_module()
self._chunk_decode_coordinator: _MiniMaxH3ChunkDecodeCoordinator | None = None
def load_to_device(self) -> None:
if self._stager is not None:
@@ -265,6 +276,59 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
def is_distributed_enabled(self) -> bool:
return self.parallel_size > 1 and dist.is_initialized()
def _denormalize_latent(self, latent: torch.Tensor) -> torch.Tensor:
channels = int(self.config_dict["latent_channels"])
mean = torch.tensor(
self.config_dict["latents_mean"],
device=latent.device,
dtype=latent.dtype,
).view(1, channels, 1, 1, 1)
std = torch.tensor(
self.config_dict["latents_std"],
device=latent.device,
dtype=latent.dtype,
).view(1, channels, 1, 1, 1)
return latent * std + mean
def _decode_tiling_context(self, latent: torch.Tensor) -> AbstractContextManager:
# The checkpoint hands rank r the tiles ``range(r, num_tiles, sp_size)``
# and then rejects an empty share inside the gather. A rank with no
# tiles raises and leaves the collective while the others block in it
# forever, so too few tiles hangs the whole stage rather than failing
# it. Tile count depends only on the latent shape, so every rank takes
# this branch together.
num_tiles = self._decoder_tile_count(latent)
if self.parallel_size > 1 and num_tiles < self.parallel_size:
logger.warning_once(
"MiniMax-H3 VAE decode splits into %d tile(s) but the tile group has "
"%d ranks; decoding rank-locally for this shape instead, which is "
"slower but avoids ranks without tiles hanging the collective.",
num_tiles,
self.parallel_size,
)
return self._rank_local_tiling()
return nullcontext()
def _normalize_decoded_frames(self, decoded: torch.Tensor) -> torch.Tensor:
frames = self.model.processor.revert_tensor(decoded)
if frames.ndim == 4:
frames = frames.unsqueeze(0).transpose(1, 2)
if frames.ndim != 5:
raise ValueError(f"unexpected decoded video shape {tuple(frames.shape)}")
return frames.float()
def _chunk_process_group(self) -> dist.ProcessGroup | None:
if not self.is_distributed_enabled():
return None
if model_parallel_is_initialized() and get_data_parallel_world_size() > 1:
raise MiniMaxH3ChunkedDecodeUnsupportedError(
"MiniMax-H3 temporal chunks do not yet support combining data parallelism with VAE patch parallelism"
)
group = self._native_parallel_state().get("sp_process_group")
if group is None or dist.get_world_size(group) != self.parallel_size:
raise RuntimeError("MiniMax-H3 VAE chunk decode has an invalid spatial-parallel group")
return group
@torch.inference_mode()
def encode_image(self, image: Image.Image) -> torch.Tensor:
previous_parallel = self.model.parallel_tiling
@@ -357,44 +421,40 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
@torch.inference_mode()
def decode_latent(self, latent: torch.Tensor) -> torch.Tensor:
channels = int(self.config_dict["latent_channels"])
mean = torch.tensor(
self.config_dict["latents_mean"],
device=latent.device,
dtype=latent.dtype,
).view(1, channels, 1, 1, 1)
std = torch.tensor(
self.config_dict["latents_std"],
device=latent.device,
dtype=latent.dtype,
).view(1, channels, 1, 1, 1)
# The checkpoint hands rank r the tiles ``range(r, num_tiles, sp_size)``
# and then rejects an empty share inside the gather. A rank with no
# tiles raises and leaves the collective while the others block in it
# forever, so too few tiles hangs the whole stage rather than failing
# it. Tile count depends only on the latent shape, so every rank takes
# this branch together.
num_tiles = self._decoder_tile_count(latent)
if self.parallel_size > 1 and num_tiles < self.parallel_size:
logger.warning_once(
"MiniMax-H3 VAE decode splits into %d tile(s) but the tile group has "
"%d ranks; decoding rank-locally for this shape instead, which is "
"slower but avoids ranks without tiles hanging the collective.",
num_tiles,
self.parallel_size,
)
tiling_context: AbstractContextManager = self._rank_local_tiling()
else:
tiling_context = nullcontext()
with self._decode_tiling_context(latent):
decoded = self.model.decode_base(self._denormalize_latent(latent))
return self._normalize_decoded_frames(decoded)
with tiling_context:
decoded = self.model.decode_base(latent * std + mean)
frames = self.model.processor.revert_tensor(decoded)
if frames.ndim == 4:
frames = frames.unsqueeze(0).transpose(1, 2)
if frames.ndim != 5:
raise ValueError(f"unexpected decoded video shape {tuple(frames.shape)}")
return frames.float()
@torch.inference_mode()
def decode_latent_with_chunks(
self,
latent: torch.Tensor,
chunk_callback: MiniMaxH3VideoChunkCallback | None,
) -> torch.Tensor:
"""Decode a video while publishing committed temporal chunks.
Every rank collaborating in the configured H3 VAE tile group must call
this method, with ``chunk_callback`` supplied only on group rank 0. A
callback failure is deferred until the compatibility decoder finishes its
remaining collectives, then propagated to every group rank. With
``parallel_size=1`` the runtime designates each request's output owner
by supplying a callback there and ``None`` on replica peers.
This is a model-local producer seam before the typed media boundary in
RFC #6541 / PR #6615. Representation conversion, transport,
presentation policy, and encoding remain runtime-owned. Chunks retain
the VAE canvas dimensions; a runtime consumer must apply the same
requested-size crop as the complete H3 pipeline output.
"""
if self._chunk_decode_coordinator is None:
self._chunk_decode_coordinator = _MiniMaxH3ChunkDecodeCoordinator()
return self._chunk_decode_coordinator.decode(
self,
latent,
chunk_callback,
group=self._chunk_process_group(),
)
class MiniMaxH3AudioVAE(nn.Module):
@@ -516,4 +576,10 @@ class MiniMaxH3AudioVAE(nn.Module):
return waveform.permute(1, 0, 2).contiguous().float()
__all__ = ["MiniMaxH3AudioVAE", "MiniMaxH3VideoVAE"]
__all__ = [
"MiniMaxH3AudioVAE",
"MiniMaxH3ChunkCallbackPeerError",
"MiniMaxH3ChunkedDecodeUnsupportedError",
"MiniMaxH3VideoChunkCallback",
"MiniMaxH3VideoVAE",
]
@@ -0,0 +1,368 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Paged KV cache for SenseNova-U1 autoregressive decode, and a CUDA graph over it.
Autoregressive decode dominates a think request -- 80% of the wall clock, at
about 20 ms per token for 42 layers -- and almost all of the GPU idle in that
window is Python/ATen dispatch rather than any CUDA call. A CUDA graph retires
those dispatches, but capture needs static shapes, and the default cache grows
K/V with ``torch.cat`` on every step.
Padding the cache to a bucket and masking the tail does give static shapes, but
the mask is what costs: measured on an A800 at real decode shapes, a masked
bucket attention runs 7.96-11.55 ms per step against 1.03 ms for an exact-length
unmasked one, which is more than the graph saves.
The way out is a paged cache. ``flash_attn_varlen_func`` takes the used length
as a *tensor* (``seqused_k``) alongside a ``block_table``, so the buffers stay
bucket-sized and capturable while the kernel reads only the valid prefix. One
captured graph then serves every length in the bucket.
Scope, and when to delete this. The cache is model-local on purpose:
``DiffusionKVCacheManager`` reserves once per scheduler request, and
``ARDiffusionModelRunner`` lives under ``vllm_omni/experimental``. SenseNova runs
its whole autoregressive loop inside one pipeline ``forward``, so the decode
steps never surface to the scheduler and it cannot admit, pool, evict, reuse a
prefix for, or continuously batch them. What is here is therefore one set of
buffers behind an identity block table, reused by whichever request fits them,
not general paged-KV support: it holds while the pipeline serves one sequence
per forward, and is released when sleep discards the memory it captured
against. Once the decode loop is
driven by a scheduler-visible runner, or the pipeline handles more than one
sequence per forward, delete this file and use the manager instead.
"""
from __future__ import annotations
import inspect
import torch
from vllm.logger import init_logger
from vllm.platforms import current_platform
logger = init_logger(__name__)
BLOCK_SIZE = 16
# Growth schedule. Each new bucket costs one capture, so keep the list short and
# let the last entry cover the tail rather than doubling forever.
BUCKETS = (512, 1024, 2048, 4096, 8192)
# Past the last bucket the schedule grows by this much at a time. A request
# decodes at most `max_think_tokens` (1024) or `max_tokens` (512) steps, so a
# step this size costs at most one reallocation per request. Doubling the last
# bucket instead would also cost at most one, but a 9,100-token request would
# then reserve 766 MiB it never uses, against 118 MiB here.
TAIL_STEP = 2048
# Every argument to the kernel is passed by keyword, so the probe covers all of
# them. An older wheel can export `flash_attn_varlen_func` without the paged
# ones, or under different names, and a version string would not say which;
# probe the signature and fall back to SDPA when any is missing.
_REQUIRED_KWARGS = (
"max_seqlen_q",
"cu_seqlens_q",
"max_seqlen_k",
"seqused_k",
"block_table",
"softmax_scale",
"causal",
)
def _flash_varlen():
"""Import lazily: the bundled kernel is CUDA-only and optional."""
from vllm.vllm_flash_attn import flash_attn_varlen_func
return flash_attn_varlen_func
def _accepts_paged_kwargs(fn) -> bool:
try:
params = inspect.signature(fn).parameters
except (TypeError, ValueError): # pragma: no cover - no introspectable signature
return False
return all(name in params for name in _REQUIRED_KWARGS)
def paged_decode_supported(device: torch.device, head_dim: int) -> bool:
if device.type != "cuda":
return False
if head_dim % 8 or head_dim > 256:
return False
try:
fn = _flash_varlen()
except Exception as exc: # pragma: no cover - depends on the wheel
logger.debug("Paged decode unavailable: %s", exc)
return False
if not _accepts_paged_kwargs(fn):
logger.debug("Paged decode unavailable: flash_attn_varlen_func lacks %s", _REQUIRED_KWARGS)
return False
return True
def dynamic_lora_wrappers_present(module) -> bool:
"""True once ``DiffusionLoRAManager`` has wrapped layers under ``module``.
The manager replaces linear layers with ``BaseLayerWithLoRA`` and then
binds, rescales or resets slot 0 per request; the wrappers stay in the tree
afterwards. Capture records the module tree and the branches that ran during
it, so a graph taken before an adapter was bound would replay without the
adapter's matmuls and one taken with it bound would survive its removal --
both silently, and neither visible in the shape, dtype and bucket that decide
reuse. Nothing captured may outlive a request once these exist.
The distilled LoRA is not affected: it is fused one-way into the weights and
leaves no wrapper behind.
"""
try:
from vllm.lora.layers import BaseLayerWithLoRA
except ImportError: # pragma: no cover - depends on the wheel
return False
return any(isinstance(m, BaseLayerWithLoRA) for m in module.modules())
def _bucket_for(length: int) -> int:
for b in BUCKETS:
if length <= b:
return b
# Past the schedule, grow in TAIL_STEP-sized steps. Rounding to a whole
# number of blocks here would reallocate and re-capture every BLOCK_SIZE
# tokens: a think edit with two 2048x2048 inputs runs to about 9,100 tokens
# and paid 41 reallocations and 42 captures in one request.
over = length - BUCKETS[-1]
return BUCKETS[-1] + ((over + TAIL_STEP - 1) // TAIL_STEP) * TAIL_STEP
class PagedDecodeCache:
"""Per-layer paged K/V for single-token decode.
Layout follows what ``flash_attn_varlen_func`` accepts with a block table:
``(num_blocks, BLOCK_SIZE, kv_heads, head_dim)``. ``seqused`` holds the
number of valid tokens and is the only thing that changes between decode
steps, which is what lets a captured graph be replayed unchanged.
"""
def __init__(self, num_layers, kv_heads, head_dim, length, device, dtype):
self.kv_heads = kv_heads
self.head_dim = head_dim
self.device = device
self.dtype = dtype
self.bucket = _bucket_for(length)
nblocks = self.bucket // BLOCK_SIZE
shape = (nblocks, BLOCK_SIZE, kv_heads, head_dim)
self.k = [torch.zeros(shape, device=device, dtype=dtype) for _ in range(num_layers)]
self.v = [torch.zeros(shape, device=device, dtype=dtype) for _ in range(num_layers)]
self.block_table = torch.arange(nblocks, device=device, dtype=torch.int32).unsqueeze(0)
self.cu_seqlens_q = torch.tensor([0, 1], device=device, dtype=torch.int32)
self.seqused = torch.zeros(1, device=device, dtype=torch.int32)
# The write slot has to live in a tensor, not a Python int. Under graph
# capture a Python index is baked in at capture time, so every replay
# would overwrite the same slot -- correct eagerly, silently wrong once
# captured.
self.pos = torch.zeros(1, device=device, dtype=torch.int64)
self._length = 0
# Bumped whenever the buffers are reallocated, so a graph runner can
# tell that its capture no longer points at live memory.
self.generation = 0
@property
def length(self) -> int:
return self._length
def set_length(self, n: int) -> None:
self._length = n
self.seqused.fill_(n)
self.pos.fill_(n - 1)
def reusable_for(self, num_layers, kv_heads, head_dim, dtype, length) -> bool:
"""Can this cache serve another request without reallocating?
Reusing the buffers is what lets captured graphs survive across
requests: a capture binds the addresses it recorded, so a fresh cache
forces a fresh capture.
"""
return (
len(self.k) == num_layers
and self.kv_heads == kv_heads
and self.head_dim == head_dim
and self.dtype == dtype
and _bucket_for(length) <= self.bucket
)
def load_prefix(self, cache) -> bool:
"""Copy a prefill cache into the paged buffers.
Only the first ``seqused`` rows are ever read, so the tail is left as
whatever the previous request wrote. Returns False for an empty cache.
"""
layer0 = cache.layers[0]
if layer0.keys is None:
return False
prefix = layer0.keys.shape[2]
for i in range(len(self.k)):
keys, values = cache.layers[i].keys, cache.layers[i].values
# [B, H, S, D] -> [S, H, D], then into the paged blocks
self._write_prefix(i, keys[0].transpose(0, 1), values[0].transpose(0, 1))
self.set_length(prefix)
return True
@classmethod
def from_dynamic_cache(cls, cache, num_layers, device, dtype, min_length=0):
"""Allocate buffers for a prefill cache and copy it in."""
layer0 = cache.layers[0]
if layer0.keys is None:
return None
_, kv_heads, prefix, head_dim = layer0.keys.shape
obj = cls(num_layers, kv_heads, head_dim, max(prefix + 1, min_length), device, dtype)
obj.load_prefix(cache)
return obj
def _write_prefix(self, layer_idx, keys_shd, values_shd):
n = keys_shd.shape[0]
flat_k = self.k[layer_idx].view(-1, self.kv_heads, self.head_dim)
flat_v = self.v[layer_idx].view(-1, self.kv_heads, self.head_dim)
flat_k[:n].copy_(keys_shd)
flat_v[:n].copy_(values_shd)
def grow(self, length: int) -> bool:
"""Move to the next bucket, keeping the tokens already stored.
Everything is built into locals first and published in one step. A
resize that failed part way -- an allocation for a later layer, say --
would otherwise leave new K beside old V under the ``generation`` a
captured graph is keyed on, so the next request would keep selecting
that graph and replay it over storage that had just been freed.
"""
new_bucket = _bucket_for(length)
if new_bucket == self.bucket:
return False
nblocks = new_bucket // BLOCK_SIZE
shape = (nblocks, BLOCK_SIZE, self.kv_heads, self.head_dim)
old_len = self._length
new_k, new_v = [], []
for old_k, old_v in zip(self.k, self.v):
k = torch.zeros(shape, device=self.device, dtype=self.dtype)
v = torch.zeros(shape, device=self.device, dtype=self.dtype)
k.view(-1, self.kv_heads, self.head_dim)[:old_len].copy_(
old_k.view(-1, self.kv_heads, self.head_dim)[:old_len]
)
v.view(-1, self.kv_heads, self.head_dim)[:old_len].copy_(
old_v.view(-1, self.kv_heads, self.head_dim)[:old_len]
)
new_k.append(k)
new_v.append(v)
block_table = torch.arange(nblocks, device=self.device, dtype=torch.int32).unsqueeze(0)
# Publish. New buffers mean any graph captured against the old ones is
# stale, so the generation moves in the same step that hands them over.
self.k, self.v, self.block_table, self.bucket = new_k, new_v, block_table, new_bucket
self.generation += 1
self.set_length(old_len)
logger.debug("Paged decode grew to bucket %d at length %d (generation %d)", new_bucket, length, self.generation)
return True
def attend(self, layer_idx, query_bhsd, key_bhsd, value_bhsd, softmax_scale):
"""Append this step's K/V and attend over the whole valid prefix.
Writes at ``length - 1`` because the caller has already advanced the
length for this step. Both the slot and the attended length come from
device tensors, so a captured graph reads whatever they hold at replay
rather than whatever they held at capture.
"""
flat_k = self.k[layer_idx].view(-1, self.kv_heads, self.head_dim)
flat_v = self.v[layer_idx].view(-1, self.kv_heads, self.head_dim)
flat_k.index_copy_(0, self.pos, key_bhsd[0, :, 0].unsqueeze(0))
flat_v.index_copy_(0, self.pos, value_bhsd[0, :, 0].unsqueeze(0))
heads, head_dim = query_bhsd.shape[1], query_bhsd.shape[3]
q = query_bhsd[0].transpose(0, 1).reshape(1, heads, head_dim)
out = _flash_varlen()(
q,
self.k[layer_idx],
self.v[layer_idx],
max_seqlen_q=1,
cu_seqlens_q=self.cu_seqlens_q,
max_seqlen_k=self.bucket,
seqused_k=self.seqused,
block_table=self.block_table,
softmax_scale=softmax_scale,
causal=False,
)
return out.reshape(1, 1, heads, head_dim)
def to_dynamic_cache(self, cache) -> None:
"""Write the decoded K/V back so the generation stage sees one cache.
Only decode is paged; the DiT stage that follows reads the ordinary
cache, so the two are reconciled once at the hand-off rather than kept
in sync on every step.
"""
n = self._length
for i, layer in enumerate(cache.layers):
keys = self.k[i].view(-1, self.kv_heads, self.head_dim)[:n]
values = self.v[i].view(-1, self.kv_heads, self.head_dim)[:n]
layer.keys = keys.transpose(0, 1).unsqueeze(0).contiguous()
layer.values = values.transpose(0, 1).unsqueeze(0).contiguous()
class DecodeGraphRunner:
"""One captured decode step per KV bucket, replayed for every token.
The whole point of the paged cache is that everything the step reads which
varies -- the token, its position, and how much of the cache is live -- sits
in device tensors. So a single capture serves every step in a bucket: fill
the tensors, replay, read the logits out of the static output.
A capture is invalidated when the cache reallocates (tracked by
``PagedDecodeCache.generation``), which happens once per bucket boundary.
"""
def __init__(self, language_model, cache, device):
self.lm = language_model
self.cache = cache
self.device = device
self.input_ids = torch.zeros(1, 1, dtype=torch.long, device=device)
self.indexes = torch.zeros(3, 1, dtype=torch.long, device=device)
self._graphs: dict[tuple[int, int], tuple[torch.cuda.CUDAGraph, torch.Tensor]] = {}
self.captures = 0
def _forward(self):
return self.lm(
input_ids=self.input_ids,
indexes=self.indexes,
past_key_values=None,
use_cache=False,
paged_cache=self.cache,
).logits
def _capture(self, key):
# Warm up on a side stream first so cuBLAS workspaces and any lazy
# allocation happen outside the capture. The warm-up recomputes this
# same step, writing the same K/V to the same slot, so it is idempotent.
side = torch.cuda.Stream()
side.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(side):
for _ in range(3):
self._forward()
torch.cuda.current_stream().wait_stream(side)
graph = torch.cuda.CUDAGraph()
# A private pool would hand every runner its own arena and never give it
# back, and a runner is rebuilt whenever the cache reallocates or a LoRA
# wrapper appears: measured, that grew device memory by about 40 MiB per
# request and never levelled off. The platform pool is shared, which is
# what every other captured path in the tree uses.
with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()):
logits = self._forward()
self.captures += 1
logger.debug("Captured decode graph for bucket=%d generation=%d", self.cache.bucket, self.cache.generation)
self._graphs[key] = (graph, logits)
return self._graphs[key]
def step(self, token, t_index):
"""Run one decode step. Returns the static logits tensor."""
self.input_ids[0, 0] = token
self.indexes[0, 0] = t_index
key = (self.cache.bucket, self.cache.generation)
entry = self._graphs.get(key) or self._capture(key)
entry[0].replay()
return entry[1]
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""SenseNova-U1 Pipeline for vLLM-Omni.
SenseNova-U1 is a unified Qwen3-based model that uses Mixture-of-Tokenizers
@@ -35,6 +35,13 @@ from vllm.logger import init_logger
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.lora.loader import (
LoraLoaderMixin,
_apply_diffusers_lora_alpha_scaling,
_load_lora_state_dict,
_prepare_lora_delta,
_remap_state_dict_keys,
)
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
@@ -44,7 +51,14 @@ from vllm_omni.transformers_utils.configs.sensenova_u1 import (
SenseNovaU1Config,
)
from .paged_decode import (
DecodeGraphRunner,
PagedDecodeCache,
dynamic_lora_wrappers_present,
paged_decode_supported,
)
from .sensenova_u1_transformer import (
SenseNovaU1CausalLMOutput,
SenseNovaU1ForCausalLM,
clear_flash_kv_cache,
create_block_causal_mask,
@@ -446,7 +460,13 @@ class SenseNovaU1DenoisingAdapter(nn.Module):
return self.language_model(*args, **kwargs)
class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProfilerMixin, CFGParallelMixin):
class SenseNovaU1Pipeline(
nn.Module,
SupportsComponentDiscovery,
DiffusionPipelineProfilerMixin,
CFGParallelMixin,
LoraLoaderMixin,
):
"""SenseNova-U1 text-to-image and image-to-image pipeline for vllm-omni.
Builds the full model graph internally:
@@ -709,8 +729,71 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
sigma = shift * sigma / (1 + (shift - 1) * sigma)
return 1 - sigma
def _ar_step(self, next_token, t_idx, past_key_values):
def _decode_context(self, past_key_values):
"""Paged cache plus a captured graph for the AR loop, or ``None``.
Decode is where a think request spends most of its time, and every step
issues the same work at the same shapes, so it is worth capturing. The
cache has to be paged for that: it keeps the buffers a fixed size while
the kernel reads only ``seqused_k`` of them, which is what lets one
capture serve every step until the next bucket.
Both live on the pipeline and are reused across requests unless a LoRA
adapter is being served dynamically. A capture binds the addresses it
recorded, so a cache built per request forces a capture per request:
measured, that left about 40 MiB of device memory behind on every
request and re-captured every time. Reuse rests on the same thing the
paged path already rests on -- one sequence in flight per pipeline
forward.
"""
if os.environ.get("VLLM_OMNI_SENSENOVA_PAGED_DECODE", "1") != "1":
return None
head_dim = self.language_model.model.layers[0].self_attn.head_dim
if not paged_decode_supported(torch.device(self.device), head_dim):
return None
if past_key_values is None or not past_key_values.layers:
return None
layer0 = past_key_values.layers[0]
if layer0.keys is None:
return None
_, kv_heads, prefix, kv_head_dim = layer0.keys.shape
device = torch.device(self.device)
num_layers = len(self.language_model.model.layers)
# Reuse is off once the LoRA manager holds the decode path: the test
# below sees shape, dtype and bucket, none of which move when an adapter
# is bound, rescaled or reset between requests.
reuse = not dynamic_lora_wrappers_present(self.language_model)
decode = getattr(self, "_paged_decode", None) if reuse else None
if decode is None or not decode[0].reusable_for(
num_layers, kv_heads, kv_head_dim, layer0.keys.dtype, prefix + 1
):
cache = PagedDecodeCache(num_layers, kv_heads, kv_head_dim, prefix + 1, device, layer0.keys.dtype)
decode = (cache, DecodeGraphRunner(self.language_model, cache, device))
self._paged_decode = decode if reuse else None
if not decode[0].load_prefix(past_key_values):
return None
return decode
def release_captured_graphs(self) -> None:
"""Drop the reused paged cache and the graphs captured against it.
Sleep level 2 discards the memory a capture recorded, so anything held
across requests has to go with it. The next request rebuilds both.
"""
self._paged_decode = None
def _ar_step(self, next_token, t_idx, past_key_values, decode=None):
"""One autoregressive token step. Constructs single-token `indexes` explicitly."""
if decode is not None:
cache, runner = decode
if cache.length + 1 > cache.bucket:
cache.grow(cache.length + 1)
cache.set_length(cache.length + 1)
logits = runner.step(int(next_token), t_idx + 1)
# The caller reassigns its cache handle from this field; the paged
# tokens are merged into it once, at the hand-off below.
return SenseNovaU1CausalLMOutput(logits=logits, past_key_values=past_key_values)
indexes = torch.tensor([[t_idx + 1], [0], [0]], dtype=torch.long, device=self.device)
outputs = self.language_model(
input_ids=next_token.unsqueeze(0),
@@ -725,24 +808,28 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
think_end_token_id = self.tokenizer.convert_tokens_to_ids("</think>")
think_token_ids = []
next_token = torch.argmax(prefix_outputs.logits[:, -1, :], dim=-1)
decode = self._decode_context(past_key_values)
for _ in range(max_think_tokens):
token_item = next_token.item()
if token_item == eos_token_id:
break
if token_item == think_end_token_id:
outputs = self._ar_step(next_token, t_idx, past_key_values)
outputs = self._ar_step(next_token, t_idx, past_key_values, decode)
past_key_values = outputs.past_key_values
t_idx += 1
think_token_ids.append(token_item)
break
think_token_ids.append(token_item)
outputs = self._ar_step(next_token, t_idx, past_key_values)
outputs = self._ar_step(next_token, t_idx, past_key_values, decode)
past_key_values = outputs.past_key_values
t_idx += 1
next_token = torch.argmax(outputs.logits[:, -1, :], dim=-1)
if decode is not None:
decode[0].to_dynamic_cache(past_key_values)
# Append "\n\n<img>" tokens to cache
append_ids = self.tokenizer(
"\n\n<img>",
@@ -797,6 +884,10 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
"""Autoregressive text decoding seeded from prefix logits."""
eos_token_id = self.tokenizer.convert_tokens_to_ids("<|im_end|>")
generated_ids: list[int] = []
# Same paged decode as the think loop. No write-back at the end: the
# caller returns the decoded string and never reads the cache again,
# unlike the think path, which hands it to the generation stage.
decode = self._decode_context(past_key_values)
logits = prefix_logits[:, -1, :]
for _ in range(max_tokens):
@@ -810,7 +901,7 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
if token_id == eos_token_id:
break
generated_ids.append(token_id)
outputs = self._ar_step(next_token, t_idx, past_key_values)
outputs = self._ar_step(next_token, t_idx, past_key_values, decode)
past_key_values = outputs.past_key_values
logits = outputs.logits[:, -1, :]
t_idx += 1
@@ -1174,8 +1265,37 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
layer.values = layer.values.expand(batch_size, *layer.values.shape[1:])
prepare_flash_kv_cache(kv, current_len=token_hw, batch_size=batch_size)
@staticmethod
def _is_warmup_request(req) -> bool:
is_dummy_run = getattr(req, "is_dummy_run", None)
if callable(is_dummy_run):
return bool(is_dummy_run())
return OmniDiffusionRequest.is_dummy_run_request_id(getattr(req, "request_id", None))
def _warm_ar_decode(self) -> None:
"""Run one single-token decode at startup so its compiled region is built.
The engine's dummy request is a text-to-image run with think off. It
exercises the prefill shape but never the decode one, so whatever the
deploy config compiles for decode lands on the first real request
instead of on readiness. Warmup is best effort and must not fail startup.
"""
try:
lm = self.language_model
device = torch.device(self.device)
ids = torch.zeros(1, 2, dtype=torch.long, device=device)
idx = torch.zeros(3, 2, dtype=torch.long, device=device)
prefill = lm(input_ids=ids, indexes=idx, use_cache=True)
nxt = torch.zeros(1, 1, dtype=torch.long, device=device)
nidx = torch.tensor([[2], [0], [0]], dtype=torch.long, device=device)
lm(input_ids=nxt, indexes=nidx, past_key_values=prefill.past_key_values, use_cache=True)
except Exception as exc: # pragma: no cover - warmup is best effort
logger.warning("Autoregressive decode warmup skipped: %s", exc)
@torch.inference_mode()
def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
if self._is_warmup_request(req):
self._warm_ar_decode()
p = self._parse_request(req)
input_images = self._extract_input_images(p.first_prompt)
@@ -1408,21 +1528,129 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
# Weight loading
# -----------------------------------------------------------------------
# Shard name -> fused parameter, used by weight loading and LoRA loading.
# More specific _mot_gen patterns FIRST to avoid substring ambiguity
# (e.g. `.q_proj` is a substring of `.q_proj_mot_gen`).
stacked_params_mapping: ClassVar[list[tuple[str, str, str | int]]] = [
(".qkv_proj_mot_gen", ".q_proj_mot_gen", "q"),
(".qkv_proj_mot_gen", ".k_proj_mot_gen", "k"),
(".qkv_proj_mot_gen", ".v_proj_mot_gen", "v"),
(".qkv_proj", ".q_proj", "q"),
(".qkv_proj", ".k_proj", "k"),
(".qkv_proj", ".v_proj", "v"),
# MLP gate/up fused into MergedColumnParallelLinear
(".gate_up_proj", ".gate_proj", 0),
(".gate_up_proj", ".up_proj", 1),
]
def load_lora_weights(
self,
pretrained_model_name_or_path: str | list[str],
adapter_name: str | None = None,
) -> None:
"""Fuse a distilled few-step LoRA into the weights.
The checkpoints ship kohya `lora_down`/`lora_up`/`alpha` names, renamed
here to the Diffusers names before `load_lora_into_module` routes each
delta into its slice of the fused projections.
"""
if adapter_name in self.lora_loaded:
return
paths = (
[pretrained_model_name_or_path]
if isinstance(pretrained_model_name_or_path, str)
else list(pretrained_model_name_or_path)
)
if len(paths) != 1:
raise ValueError(f"Expected exactly one distilled LoRA file, got {len(paths)}: {paths}")
state_dict = _load_lora_state_dict(paths[0])
state_dict = _remap_state_dict_keys(
state_dict,
[(".lora_down.weight", ".lora_A.weight"), (".lora_up.weight", ".lora_B.weight")],
)
# Fold alpha and the low-rank product in fp32; scaling B in bf16 would
# round once before the matmul and once after.
state_dict = {k: v.to(torch.float32) for k, v in state_dict.items()}
state_dict = _apply_diffusers_lora_alpha_scaling(state_dict)
applied = self._add_lora_deltas(state_dict)
if not applied:
raise ValueError(f"Distilled LoRA {paths[0]} matched no parameter of this pipeline.")
logger.info("Fused distilled LoRA %s into %d parameters", paths[0], applied)
# Fusion is one-way and there is no unload path, so keep only a sentinel;
# the fp32 state dict is ~1.5 GiB and would stay resident for the process.
self.lora_loaded[adapter_name] = paths[0]
def _fused_shards(self, base_key: str) -> list[tuple[str, str | int]] | None:
"""Unfused source names and shard ids for a fused parameter, else None."""
for param_name, weight_name, shard_id in self.stacked_params_mapping:
if base_key.endswith(param_name):
return [
(base_key[: -len(param_name)] + w, sid)
for p, w, sid in self.stacked_params_mapping
if p == param_name
]
return None
def _add_lora_deltas(self, state_dict: dict[str, torch.Tensor]) -> int:
"""Add every matching delta into its parameter, sharding through vLLM.
The delta read from the checkpoint spans the whole layer. A parameter of
a parallel layer only holds this rank's slice of it, so the delta goes
through the layer's own weight loader -- into a zeroed copy of the
parameter, which is then added -- instead of being added directly.
Each parameter is fused in place as it is reached, so a failure part way
through leaves the earlier ones fused and the pipeline has to be
reloaded. Deferring the writes would mean holding every delta at once,
and the deltas are full rank: 15.09 GiB for the 8B checkpoint, against
the 192 MiB one parameter costs here. Fusion runs at load time, before
the pipeline serves anything, so a partially fused model is never
reachable from a request.
"""
applied = 0
for name, param in self.named_parameters():
if not name.endswith(".weight"):
continue
base_key = name[: -len(".weight")]
shards = self._fused_shards(base_key)
deltas: list[tuple[torch.Tensor, str | int | None]] = []
for src, shard_id in shards or [(base_key, None)]:
delta, _ = _prepare_lora_delta(state_dict, src, None)
if delta is None:
deltas = []
break
deltas.append((delta, shard_id))
if not deltas:
continue
weight_loader = getattr(param, "weight_loader", None)
buffer = torch.zeros_like(param.data)
if weight_loader is None:
# Not a parallel layer: the delta is already rank-local.
buffer = torch.cat([d for d, _ in deltas]).to(device=param.device, dtype=param.dtype)
else:
saved = param.data
try:
param.data = buffer
for delta, shard_id in deltas:
delta = delta.to(device=param.device, dtype=param.dtype)
if shard_id is None:
weight_loader(param, delta)
else:
weight_loader(param, delta, shard_id)
finally:
param.data = saved
with torch.no_grad():
param.data.add_(buffer)
applied += 1
return applied
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
stacked_params_mapping = [
# More specific _mot_gen patterns FIRST to avoid substring
# ambiguity (e.g. `.q_proj` is a substring of `.q_proj_mot_gen`).
(".qkv_proj_mot_gen", ".q_proj_mot_gen", "q"),
(".qkv_proj_mot_gen", ".k_proj_mot_gen", "k"),
(".qkv_proj_mot_gen", ".v_proj_mot_gen", "v"),
(".qkv_proj", ".q_proj", "q"),
(".qkv_proj", ".k_proj", "k"),
(".qkv_proj", ".v_proj", "v"),
# MLP gate/up fused into MergedColumnParallelLinear
(".gate_up_proj", ".gate_proj", 0),
(".gate_up_proj", ".up_proj", 1),
]
self.stacked_params_mapping = stacked_params_mapping
stacked_params_mapping = self.stacked_params_mapping
params_dict = dict(self.named_parameters())
loaded_params: set[str] = set()
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Qwen3 LLM with Mixture-of-Tokenizers (MoT) for SenseNova-U1.
Ported from the sensenova_u1 package with vllm tensor-parallel support:
@@ -16,6 +16,7 @@ from dataclasses import dataclass
import torch
import torch.nn as nn
import torch.nn.functional as F
from cache_dit import ForwardPattern
from transformers.cache_utils import DynamicCache
from vllm.logger import init_logger
@@ -163,6 +164,18 @@ class Qwen3RotaryEmbedding(nn.Module):
return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
def _build_3d_rope(config) -> tuple[Qwen3RotaryEmbedding, Qwen3RotaryEmbedding]:
"""The t and h/w rotary modules: t uses half head_dim, h/w a quarter each."""
t_config = copy.deepcopy(config)
t_config.head_dim = config.head_dim // 2
hw_config = copy.deepcopy(config)
hw_config.head_dim = config.head_dim // 4
hw_config.rope_theta = config.rope_theta_hw
hw_config.max_position_embeddings = config.max_position_embeddings_hw
return Qwen3RotaryEmbedding(config=t_config), Qwen3RotaryEmbedding(config=hw_config)
# ---------------------------------------------------------------------------
# Norms
# ---------------------------------------------------------------------------
@@ -175,11 +188,9 @@ class Qwen3RMSNorm(nn.Module):
self.variance_epsilon = eps
def forward(self, hidden_states):
input_dtype = hidden_states.dtype
hidden_states = hidden_states.to(torch.float32)
variance = hidden_states.pow(2).mean(-1, keepdim=True)
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
return self.weight * hidden_states.to(input_dtype)
# F.rms_norm keeps the fp32 accumulation internal and rounds once; the
# cast chain it replaces rounded to bf16 before the weight multiply.
return F.rms_norm(hidden_states, self.weight.shape, self.weight, self.variance_epsilon)
# ---------------------------------------------------------------------------
@@ -303,23 +314,16 @@ class SenseNovaU1Attention(nn.Module):
self.q_norm_hw_mot_gen = Qwen3RMSNorm(self.head_dim // 2, eps=config.rms_norm_eps)
self.k_norm_hw_mot_gen = Qwen3RMSNorm(self.head_dim // 2, eps=config.rms_norm_eps)
# 3D RoPE: t uses half head_dim, h/w each use quarter head_dim
t_config = copy.deepcopy(config)
t_config.head_dim = config.head_dim // 2
self.rotary_emb = Qwen3RotaryEmbedding(config=t_config)
hw_config = copy.deepcopy(config)
hw_config.head_dim = config.head_dim // 4
hw_config.rope_theta = config.rope_theta_hw
hw_config.max_position_embeddings = config.max_position_embeddings_hw
self.rotary_emb_hw = Qwen3RotaryEmbedding(config=hw_config)
self.attn = Attention(
num_heads=self.num_heads,
head_size=self.head_dim,
causal=False,
softmax_scale=self.scaling,
num_kv_heads=self.num_heads,
# Keep K/V at their real head count. The SDPA backend already picks
# a fused GQA kernel when the runtime shape allows one and expands
# K/V itself when it does not; expanding here hides the GQA shape
# from that check, so it never fires.
num_kv_heads=self.num_kv_heads,
prefix=f"{prefix}.attn",
)
self.attn.attention = self.attn.sdpa_fallback
@@ -339,10 +343,6 @@ class SenseNovaU1Attention(nn.Module):
attention_mask: torch.Tensor | None,
) -> torch.Tensor:
"""Run unified attention with [B, H, S, D] inputs. Returns [B, S, H, D]."""
if self.num_kv_groups > 1:
n = self.num_kv_groups
key_bhsd = key_bhsd.repeat_interleave(n, dim=1)
value_bhsd = value_bhsd.repeat_interleave(n, dim=1)
q = query_bhsd.transpose(1, 2).contiguous()
k = key_bhsd.transpose(1, 2).contiguous()
v = value_bhsd.transpose(1, 2).contiguous()
@@ -358,16 +358,16 @@ class SenseNovaU1Attention(nn.Module):
attention_mask: torch.Tensor | None,
) -> torch.Tensor:
"""Run unified attention with [B, S, H, D] inputs. Returns [B, S, H, D]."""
if self.num_kv_groups > 1:
n = self.num_kv_groups
key_bshd = key_bshd.repeat_interleave(n, dim=2)
value_bshd = value_bshd.repeat_interleave(n, dim=2)
attention_mask = self._align_mask_dtype(attention_mask, query_bshd)
attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None
return self.attn(query_bshd, key_bshd, value_bshd, attn_metadata)
def _project_and_rope(self, hidden_states, indexes, qkv_proj, q_norm, k_norm, q_norm_hw, k_norm_hw):
"""Project Q/K/V via the given QKVParallelLinear and apply 3D RoPE."""
def _project_and_rope(self, hidden_states, position_embeddings, qkv_proj, q_norm, k_norm, q_norm_hw, k_norm_hw):
"""Project Q/K/V via the given QKVParallelLinear and apply 3D RoPE.
The three tables come from the model, which builds them once per forward
because every layer is handed the same `indexes`.
"""
input_shape = hidden_states.shape[:-1] # (B, S)
qkv, _ = qkv_proj(hidden_states)
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
@@ -376,9 +376,7 @@ class SenseNovaU1Attention(nn.Module):
k = k.view(*input_shape, self.num_kv_heads, self.head_dim)
v = v.view(*input_shape, self.num_kv_heads, self.head_dim)
cos_t, sin_t = self.rotary_emb(hidden_states, indexes[0].unsqueeze(0))
cos_h, sin_h = self.rotary_emb_hw(hidden_states, indexes[1].unsqueeze(0))
cos_w, sin_w = self.rotary_emb_hw(hidden_states, indexes[2].unsqueeze(0))
(cos_t, sin_t), (cos_h, sin_h), (cos_w, sin_w) = position_embeddings
value_states = v.transpose(1, 2) # [B, H, S, D]
try:
@@ -433,13 +431,23 @@ class SenseNovaU1Attention(nn.Module):
input_shape = hidden_states.shape[:-1]
query_states, key_states, value_states = self._project_and_rope(
hidden_states,
indexes,
kwargs["position_embeddings"],
self.qkv_proj,
self.q_norm,
self.k_norm,
self.q_norm_hw,
self.k_norm_hw,
)
# Single-token decode with a paged cache: append this step's K/V and let
# the kernel read the valid prefix via `seqused_k`. Shapes stay fixed,
# which is what makes the step capturable as a CUDA graph.
paged = kwargs.get("paged_cache")
if paged is not None and attention_mask is None and query_states.shape[2] == 1:
attn_output = paged.attend(self.layer_idx, query_states, key_states, value_states, self.scaling)
attn_output = attn_output.reshape(*input_shape, -1).contiguous()
attn_output, _ = self.o_proj(attn_output)
return attn_output
update_cache = kwargs.get("update_cache", True)
if past_key_values is not None:
if update_cache:
@@ -460,7 +468,7 @@ class SenseNovaU1Attention(nn.Module):
input_shape = hidden_states.shape[:-1]
query_states, key_states, value_states = self._project_and_rope(
hidden_states,
indexes,
kwargs["position_embeddings"],
self.qkv_proj_mot_gen,
self.q_norm_mot_gen,
self.k_norm_mot_gen,
@@ -659,6 +667,10 @@ class SenseNovaU1Model(nn.Module):
)
self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.norm_mot_gen = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
# Every layer receives the same `indexes`, so the tables are the same for
# all of them: build them once per forward here and thread them down,
# rather than giving every layer its own pair.
self.rotary_emb, self.rotary_emb_hw = _build_3d_rope(config)
def forward(
self,
@@ -690,15 +702,25 @@ class SenseNovaU1Model(nn.Module):
past_len = past_key_values.get_seq_length() if past_key_values else 0
seq_len = inputs_embeds.shape[1]
total_len = past_len + seq_len
mask = torch.zeros(1, 1, seq_len, total_len, device=inputs_embeds.device)
if seq_len > 1:
mask = torch.zeros(1, 1, seq_len, total_len, device=inputs_embeds.device)
causal = torch.tril(torch.ones(seq_len, seq_len, device=inputs_embeds.device))
mask[:, :, :, past_len:] = torch.where(causal == 1, 0.0, float("-inf"))
else:
# A single query token attends to every cached key, so the mask
# this branch used to build was all zeros. Passing it changed no
# result, only which SDPA kernel was reachable.
mask = None
causal_mask_mapping = {"full_attention": mask}
else:
causal_mask_mapping = attention_mask
hidden_states = inputs_embeds
position_embeddings = (
self.rotary_emb(hidden_states, indexes[0].unsqueeze(0)),
self.rotary_emb_hw(hidden_states, indexes[1].unsqueeze(0)),
self.rotary_emb_hw(hidden_states, indexes[2].unsqueeze(0)),
)
for layer in self.layers:
hidden_states = layer(
hidden_states,
@@ -706,6 +728,7 @@ class SenseNovaU1Model(nn.Module):
exist_und=exist_und,
exist_gen=exist_gen,
indexes=indexes,
position_embeddings=position_embeddings,
attention_mask=causal_mask_mapping,
past_key_values=past_key_values,
use_cache=use_cache,
+2 -10
View File
@@ -49,14 +49,13 @@ def create_diffusion_client(
od_config: OmniDiffusionConfig,
metadata: StageMetadata,
stage_init_timeout: int,
batch_size: int = 1,
use_inline: bool = False,
) -> Any:
"""Factory to create either an inline or out-of-process diffusion client."""
if use_inline:
from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
return InlineStageDiffusionClient(model, od_config, metadata, batch_size=batch_size)
return InlineStageDiffusionClient(model, od_config, metadata)
proc_manager = StageDiffusionProcManager(
model=model,
od_config=od_config,
@@ -67,7 +66,6 @@ def create_diffusion_client(
request_address=proc_manager.addresses.inputs[0],
response_address=proc_manager.addresses.outputs[0],
proc_manager=proc_manager,
batch_size=batch_size,
)
@@ -91,14 +89,12 @@ class StageDiffusionClient(StageClientBase):
response_address: str,
*,
proc_manager: StageDiffusionProcManager | None = None,
batch_size: int = 1,
) -> None:
self._initialize_client(
metadata,
request_address,
response_address,
proc_manager=proc_manager,
batch_size=batch_size,
)
@classmethod
@@ -109,7 +105,6 @@ class StageDiffusionClient(StageClientBase):
response_address: str,
*,
proc_manager: StageDiffusionProcManager | None = None,
batch_size: int = 1,
) -> StageDiffusionClient:
"""Create a client for an already-running diffusion subprocess."""
return cls(
@@ -117,7 +112,6 @@ class StageDiffusionClient(StageClientBase):
request_address,
response_address,
proc_manager=proc_manager,
batch_size=batch_size,
)
def _initialize_client(
@@ -127,7 +121,6 @@ class StageDiffusionClient(StageClientBase):
response_address: str,
*,
proc_manager: StageDiffusionProcManager | None = None,
batch_size: int,
) -> None:
self._set_stage_metadata(metadata)
self._proc_manager = proc_manager
@@ -144,11 +137,10 @@ class StageDiffusionClient(StageClientBase):
self._start_proc_monitor()
logger.info(
"[StageDiffusionClient] stage-%s [rep-%s] initialized (owns_process=%s, batch_size=%d)",
"[StageDiffusionClient] stage-%s [rep-%s] initialized (owns_process=%s)",
self.stage_id,
self.replica_id,
self._proc_manager is not None,
batch_size,
)
def _set_stage_metadata(self, metadata: StageMetadata) -> None:
@@ -519,6 +519,23 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
if self.prompt_embed_cache is not None:
self.prompt_embed_cache.clear()
def release_captured_graphs(self) -> None:
"""Drop every CUDA graph held for this model, wherever it is kept.
Sleep level 2 discards the memory a capture was recorded against, so a
graph that outlives it replays over freed storage. The runner owns
compilation and execution resources, so it owns the release: a pipeline
that captures graphs of its own implements ``release_captured_graphs``
and is collected here, instead of every caller having to know which
pipelines have one.
"""
runners = getattr(self, "graph_runners", None)
if runners is not None:
runners.clear()
release = getattr(self.pipeline, "release_captured_graphs", None)
if callable(release):
release()
def get_prompt_embed_cache_stats(self) -> dict | None:
"""Return hit/miss statistics for the prompt-embedding cache, if enabled.
@@ -786,9 +786,8 @@ class DiffusionWorker:
usage_before = allocator.get_current_usage()
if level == 2 and self.model_runner is not None:
if hasattr(self.model_runner, "graph_runners"):
self.model_runner.graph_runners.clear()
logger.info(f"[Worker {self.rank}] CUDA Graphs cleared.")
self.model_runner.release_captured_graphs()
logger.info(f"[Worker {self.rank}] CUDA Graphs cleared.")
model = self.model_runner.pipeline
self._sleep_saved_buffers = {name: buffer.cpu().clone() for name, buffer in model.named_buffers()}
+2 -5
View File
@@ -129,7 +129,6 @@ class AsyncOmniEngine:
model: str,
stage_init_timeout: int = 300,
init_timeout: int = 600,
diffusion_batch_size: int = 1,
single_stage_mode: bool = False,
transfer_emitter: Any = None,
prom_metrics: Any = None,
@@ -140,7 +139,6 @@ class AsyncOmniEngine:
) -> None:
self.model = model
self.tokenizer = tokenizer
self.diffusion_batch_size = diffusion_batch_size
# Cached by get_diffusion_od_config().
self._diffusion_od_config_view: Any = None
startup_timeout = int(init_timeout)
@@ -320,7 +318,6 @@ class AsyncOmniEngine:
config_path=self.config_path,
single_stage_mode=self.single_stage_mode,
stage_init_timeout=stage_init_timeout,
diffusion_batch_size=self.diffusion_batch_size,
async_chunk=self.async_chunk,
tokenizer=self.tokenizer,
single_stage_id_filter=self._single_stage_id_filter,
@@ -1352,8 +1349,8 @@ class AsyncOmniEngine:
or kwargs.get("fastvideo_vsa_topk") is not None
):
has_stage_attention = (
hasattr(cfg.engine_args, "diffusion_attention_config")
and cfg.engine_args.diffusion_attention_config is not None
getattr(cfg.engine_args, "diffusion_attention_config", None) is not None
or getattr(cfg.engine_args, "diffusion_attention_backend", None) is not None
)
if not has_stage_attention:
cfg.engine_args.diffusion_attention_config = parse_attention_config(
-3
View File
@@ -1546,7 +1546,6 @@ def launch_diffusion_stage_replica(
stage_config: Any,
metadata: Any,
stage_init_timeout: int,
batch_size: int,
use_inline: bool,
replica_id: int = 0,
omni_master_server: OmniMasterServer | None = None,
@@ -1565,7 +1564,6 @@ def launch_diffusion_stage_replica(
stage_config,
metadata,
stage_init_timeout=stage_init_timeout,
batch_size=batch_size,
use_inline=use_inline,
)
return client, StageReplicaResources()
@@ -1624,7 +1622,6 @@ def launch_diffusion_stage_replica(
request_address=proc_manager.addresses.inputs[0],
response_address=proc_manager.addresses.outputs[0],
proc_manager=proc_manager,
batch_size=batch_size,
)
return client, StageReplicaResources(
manager=proc_manager,
+13 -7
View File
@@ -1253,9 +1253,19 @@ def _finalize_engine_args_dict(
if is_diffusion:
from vllm_omni.diffusion.data import parse_attention_config
if engine_args_dict.get("diffusion_attention_config") is not None:
# Fold the attention shorthand into the structured config so only one
# representation reaches OmniDiffusionConfig.from_kwargs.
attention_backend = engine_args_dict.pop("diffusion_attention_backend", None)
fastvideo_vsa_topk = engine_args_dict.pop("fastvideo_vsa_topk", None)
if (
engine_args_dict.get("diffusion_attention_config") is not None
or attention_backend is not None
or fastvideo_vsa_topk is not None
):
engine_args_dict["diffusion_attention_config"] = parse_attention_config(
engine_args_dict["diffusion_attention_config"],
engine_args_dict.get("diffusion_attention_config"),
attention_backend=attention_backend,
fastvideo_vsa_topk=fastvideo_vsa_topk,
)
else:
resolve_worker_cls(engine_args_dict)
@@ -1813,7 +1823,6 @@ def initialize_diffusion_stage(
stage_cfg: Any,
metadata: StageMetadata,
stage_init_timeout: int,
batch_size: int = 1,
use_inline: bool = False,
) -> Any:
"""Build a diffusion stage client.
@@ -1823,15 +1832,12 @@ def initialize_diffusion_stage(
stage_cfg: Stage configuration.
metadata: Extracted stage metadata.
stage_init_timeout: Timeout in seconds for stage initialization handshake
batch_size: Client-side request batch width. Does not set scheduler
``max_num_seqs``; pass ``--max-num-seqs`` or stage YAML for that.
Forwarded to ``StageDiffusionClient``.
use_inline: If True, uses the inline diffusion client instead of subprocess.
"""
from vllm_omni.diffusion.stage_diffusion_client import create_diffusion_client
od_config = build_diffusion_config(model, stage_cfg, metadata)
return create_diffusion_client(model, od_config, metadata, stage_init_timeout, batch_size, use_inline)
return create_diffusion_client(model, od_config, metadata, stage_init_timeout, use_inline)
def _stage_declares_cfg_pairs(model_config: Any) -> bool:
+1 -14
View File
@@ -87,7 +87,6 @@ class StageRemoteFactoryContext:
base_metadata: Any
vllm_config: Any | None = None
executor_class: type | None = None
diffusion_batch_size: int = 1
def _build_load_balancer_factory(policy: str) -> Callable[[], LoadBalancer]:
@@ -124,7 +123,6 @@ class StageRuntime:
config_path: str,
*,
stage_init_timeout: int,
diffusion_batch_size: int,
async_chunk: bool,
tokenizer: str | None = None,
log_stats: bool = False,
@@ -133,7 +131,6 @@ class StageRuntime:
self._model = model
self._config_path = config_path
self._stage_init_timeout = stage_init_timeout
self._diffusion_batch_size = diffusion_batch_size
self._async_chunk = async_chunk
self._tokenizer = tokenizer
self._log_stats = log_stats
@@ -672,7 +669,6 @@ class StageRuntime:
stage_config=plan.stage_cfg,
metadata=plan.metadata,
stage_init_timeout=stage_init_timeout,
batch_size=self._diffusion_batch_size,
use_inline=plan.num_replicas == 1
and bool(self._num_stages == 1 or inline_diffusion or custom_pipeline_args),
replica_id=plan.replica_id,
@@ -681,10 +677,9 @@ class StageRuntime:
)
logger.info(
"[StageRuntime] Stage %s replica %s initialized (diffusion, batch_size=%d)",
"[StageRuntime] Stage %s replica %s initialized (diffusion)",
plan.metadata.stage_id,
plan.replica_id,
self._diffusion_batch_size,
)
return cast(StagePoolClient, client)
except Exception:
@@ -768,7 +763,6 @@ class DistStageRuntime(StageRuntime):
config_path: str,
*,
stage_init_timeout: int,
diffusion_batch_size: int,
async_chunk: bool,
single_stage_id_filter: int | None,
omni_master_address: str,
@@ -785,7 +779,6 @@ class DistStageRuntime(StageRuntime):
model=model,
config_path=config_path,
stage_init_timeout=stage_init_timeout,
diffusion_batch_size=diffusion_batch_size,
async_chunk=async_chunk,
tokenizer=tokenizer,
log_stats=log_stats,
@@ -914,7 +907,6 @@ class DistStageRuntime(StageRuntime):
base_metadata=metadata,
vllm_config=plan.stage_vllm_config,
executor_class=plan.executor_class,
diffusion_batch_size=self._diffusion_batch_size,
)
return self._create_remote_replica_client(ctx, plan.replica_id)
@@ -977,7 +969,6 @@ class DistStageRuntime(StageRuntime):
base_metadata=template.metadata,
vllm_config=template.stage_vllm_config,
executor_class=template.executor_class,
diffusion_batch_size=self._diffusion_batch_size,
)
return contexts
@@ -1045,7 +1036,6 @@ class DistStageRuntime(StageRuntime):
metadata,
request_address=resources.addresses.inputs[0],
response_address=resources.addresses.outputs[0],
batch_size=ctx.diffusion_batch_size,
)
logger.info(
"[DistStageRuntime] Remote diffusion replica attached stage=%d replica=%d",
@@ -1109,7 +1099,6 @@ def create_stage_runtime(
*,
single_stage_mode: bool,
stage_init_timeout: int,
diffusion_batch_size: int,
async_chunk: bool,
tokenizer: str | None = None,
# Distributed-only params:
@@ -1131,7 +1120,6 @@ def create_stage_runtime(
model=model,
config_path=config_path,
stage_init_timeout=stage_init_timeout,
diffusion_batch_size=diffusion_batch_size,
async_chunk=async_chunk,
tokenizer=tokenizer,
log_stats=log_stats,
@@ -1148,7 +1136,6 @@ def create_stage_runtime(
model=model,
config_path=config_path,
stage_init_timeout=stage_init_timeout,
diffusion_batch_size=diffusion_batch_size,
async_chunk=async_chunk,
tokenizer=tokenizer,
log_stats=log_stats,
+3 -2
View File
@@ -187,7 +187,9 @@ class OmniBase(PDDisaggregationMixin):
# override the deploy YAML's ``async_chunk: true`` default.
async_chunk = kwargs.get("async_chunk")
output_modalities = kwargs.pop("output_modalities", None)
diffusion_batch_size: int = kwargs.pop("diffusion_batch_size", 1)
if "diffusion_batch_size" in kwargs:
raise TypeError("`diffusion_batch_size` has been removed in Omni/AsyncOmni. Use `max_num_seqs`.")
if "log_requests" in kwargs:
raise TypeError("`log_requests` has been removed in Omni/AsyncOmni. Use `log_stats`.")
@@ -213,7 +215,6 @@ class OmniBase(PDDisaggregationMixin):
model=model,
init_timeout=init_timeout,
stage_init_timeout=stage_init_timeout,
diffusion_batch_size=diffusion_batch_size,
transfer_emitter=self.transfer_metrics,
prom_metrics=self.prom_metrics,
log_stats=log_stats,
Loaded 100 of 106 files, more files were not shown because too many files have changed in this diff. Show more