Compare commits
18
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
81de3f3308 | ||
|
|
f9c9de8d16 | ||
|
|
ff2e33a329 | ||
|
|
0b5b0979d8 | ||
|
|
f04ff76c54 | ||
|
|
b220df8397 | ||
|
|
520e058b95 | ||
|
|
0288c3f56e | ||
|
|
477083a103 | ||
|
|
e1b2f589c1 | ||
|
|
44c8d1c449 | ||
|
|
2298359f67 | ||
|
|
d1c8beeb18 | ||
|
|
e2b83db683 | ||
|
|
49c6fc1540 | ||
|
|
a65e89cbd4 | ||
|
|
7992c2213e | ||
|
|
8e26dce1f4 |
No files matched your search
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
@@ -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
@@ -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:
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -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,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.
|
||||
@@ -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
|
||||
|
||||
@@ -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).
|
||||
@@ -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` | ✅︎ | ✅︎ | | | — |
|
||||
|
||||
@@ -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.
|
||||
@@ -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 \
|
||||
|
||||
@@ -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)
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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,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 |
|
||||
|
||||
@@ -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.
|
||||
@@ -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) | ✅ | ❌ |
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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.
|
||||
@@ -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)
|
||||
@@ -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 |
|
||||
|
||||
@@ -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.
|
||||
@@ -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"
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 = []
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -1461,6 +1461,7 @@ class OnlineOmniClient:
|
||||
"seed",
|
||||
"instructions",
|
||||
"speed",
|
||||
"sample_rate",
|
||||
"stream_format",
|
||||
"x_vector_only_mode",
|
||||
):
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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 = (
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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,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:
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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,
|
||||
|
||||
@@ -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()}
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
Reference in new issue
Block a user