Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Hongsheng Liu
|
01e7d09119 | ||
|
hsliuustc0106
|
69766960f2 | ||
|
hsliuustc0106
|
80fd96c1da |
No files matched your search
@@ -47,17 +47,15 @@ steps:
|
||||
- timeout 15m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: "Diffusion · Distributed & Attention Test"
|
||||
- label: "Diffusion · Distributed Test"
|
||||
source_file_dependencies:
|
||||
- tests/diffusion/distributed/
|
||||
- tests/diffusion/attention/test_attention_sp.py
|
||||
- vllm_omni/diffusion/distributed/
|
||||
- vllm_omni/diffusion/attention/
|
||||
- vllm_omni/diffusion/attention/parallel/
|
||||
- vllm_omni/diffusion/data.py
|
||||
- vllm_omni/diffusion/models/
|
||||
commands:
|
||||
- timeout 20m pytest -sv tests/diffusion/distributed -m 'core_model and cuda' --run-level "core_model"
|
||||
- timeout 20m pytest -sv tests/diffusion/attention/test_attention_sp.py -m 'core_model and cuda' --run-level "core_model"
|
||||
mirror_hardwares: l4_4
|
||||
|
||||
- label: "Diffusion · Model Test"
|
||||
|
||||
@@ -9,7 +9,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
|
||||
</h3>
|
||||
|
||||
<p align="center">
|
||||
| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://deepwiki.com/vllm-project/vllm-omni"><b>DeepWiki</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1aPj0OGl_-ZVoib-Qne5dGDAlrRFB-PdHl6E-EE99g8E/edit?usp=sharing"><b>Slides</b></a> |
|
||||
| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://deepwiki.com/vllm-project/vllm-omni"><b>DeepWiki</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/111-L8zF7A1j_YI_cR8JsblofdScdRr2f/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
|
||||
</p>
|
||||
|
||||
|
||||
|
||||
+1
-1
@@ -63,5 +63,5 @@ vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, in
|
||||
|
||||
For more information, checkout the following:
|
||||
|
||||
- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1aPj0OGl_-ZVoib-Qne5dGDAlrRFB-PdHl6E-EE99g8E/edit?usp=sharing)
|
||||
- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1)
|
||||
- [vllm-omni announcement blogpost](https://blog.vllm.ai/2025/11/30/vllm-omni.html)
|
||||
@@ -1,6 +1,6 @@
|
||||
# Architecture Overview
|
||||
|
||||
This document outlines the architecture of vLLM-Omni.
|
||||
This document outlines the architectural design for vLLM-Omni.
|
||||
|
||||
<p align="center">
|
||||
<picture>
|
||||
@@ -9,456 +9,194 @@ This document outlines the architecture of vLLM-Omni.
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
## Design goals
|
||||
# Goals
|
||||
|
||||
The primary goal of vLLM-Omni is to provide a fast and easy-to-use inference
|
||||
and serving engine for omni-modality models. vLLM-Omni extends vLLM's
|
||||
text-oriented autoregressive (AR) runtime with stage-based execution for
|
||||
non-textual outputs and non-autoregressive model components.
|
||||
The primary goal of the vLLM-Omni project is to build the fastest and easiest-to-use open-source Omni-Modality model inference & serving engine. vLLM-Omni extends the original vLLM, which was created to support large language models for text-based autoregressive (AR) generation tasks. vLLM-Omni is designed to support:
|
||||
|
||||
The architecture is designed to:
|
||||
* **Non-textual Output:** Enables the integration, efficient processing and output of various data types, including but not limited to, images, audio, video, and action trajectories, alongside text.
|
||||
* **Non-Autoregressive Structure:** Support model structure beyond autoregressive, especially Diffusion Transformer (DiT), which is widely used in visual and audio generation.
|
||||
* **Integration with vLLM Core:** Maintain compatibility and leverage existing vLLM key modules and optimizations where applicable.
|
||||
* **Extensibility:** Design a modular and flexible architecture that can easily accommodate new modalities, model architectures, and output formats.
|
||||
|
||||
* support text, image, audio, video, and action inputs and outputs;
|
||||
* compose autoregressive, generation, and diffusion stages in one pipeline;
|
||||
* reuse vLLM scheduling, cache, distributed-execution, and serving
|
||||
primitives where they fit; and
|
||||
* keep model topology, deployment placement, runtime lifecycle, transport,
|
||||
and public API concerns in separate layers.
|
||||
|
||||
## Model execution
|
||||
# Representative omni-modality models
|
||||
|
||||
These four examples show how vLLM-Omni maps model components to stage-local
|
||||
execution policies. Each summary covers the model-owned path, batching,
|
||||
attention, parallelism, quantization, and primary serving targets. The profiles
|
||||
are representative; a model may support multiple `PipelineConfig` and
|
||||
`DeployConfig` combinations.
|
||||
According to analysis for current popular open-source models, most of them have the combination of AR+DiT. Specifically, they can be further categorized into 3 types below:
|
||||
|
||||
| Model | Model-owned path | vLLM-Omni stage view | Primary serving target |
|
||||
| --- | --- | --- | --- |
|
||||
| Qwen3-Omni | Thinker → Talker → Code2Wav | Three stages with asynchronous chunk hand-off | TTFT/TTFP, TPOT, E2EL, audio RTF, and concurrent throughput |
|
||||
| HunyuanImage-3.0 | Multimodal AR understanding/reasoning plus image-generation DiT | AR-only, DiT-only, or split AR → DiT deployment | DiT E2EL/denoising latency and image throughput; TTFT/TPOT for AR tasks |
|
||||
| MiniMax-H3 | Text encoder → task-specific FL2VA or Ref2VA DiT → video/audio VAE decoder | Three logical component/phase boundaries within one registered diffusion stage | Video/audio E2EL and media throughput; RTF when audio streaming is measured |
|
||||
| Cosmos3 | Unified MoT reasoner tower plus diffusion generator tower | One diffusion pipeline materializes both towers inside the stage | Image/video/action E2EL and throughput; RTF for synchronized audio |
|
||||
|
||||
### Serving metric vocabulary
|
||||
|
||||
The target metric follows the output modality, not just the model name.
|
||||
|
||||
| Metric | Meaning in this document | Most useful for |
|
||||
| --- | --- | --- |
|
||||
| **TTFT** | Time to the first text token | AR reasoning or text output |
|
||||
| **TPOT** | Time per output token after the first token | AR decode stages |
|
||||
| **TTFP** | Time to the first streamed media packet, such as audio | Interactive speech and multimodal chat |
|
||||
| **E2EL** | End-to-end latency from request admission to the final output | Image, video, audio, and action requests |
|
||||
| **RTF** | Wall-clock processing time divided by generated audio duration; lower is faster than real time | Audio generation |
|
||||
| **Throughput** | Requests, tokens, or media seconds produced per wall-clock second | Offline and concurrent serving |
|
||||
|
||||
### Qwen3-Omni: streaming AR pipeline
|
||||
|
||||
Qwen3-Omni illustrates a multi-stage streaming AR pipeline: Thinker → Talker →
|
||||
Code2Wav. The [technical report](https://arxiv.org/abs/2509.17765) describes
|
||||
the model components, and the figure shows the stage-level streaming schedule
|
||||
used by vLLM-Omni.
|
||||
**DiT as a main structure, with AR as text encoder (e.g.: Qwen-Image)**
|
||||
A powerful image generation foundation model capable of complex text rendering and precise image editing.
|
||||
|
||||
<p align="center">
|
||||
<a href="https://github.com/vllm-project/vllm-omni/blob/main/docs/source/architecture/qwen3-omni-async-chunk.png">
|
||||
<img alt="Qwen3-Omni Thinker, Talker, and Code2Wav streaming stages" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-async-chunk.png" width="80%">
|
||||
</a>
|
||||
<picture>
|
||||
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/dit-main-architecture.png">
|
||||
<img alt="Qwen-Image" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/dit-main-architecture.png" width=30%>
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
#### Model architecture by stage
|
||||
|
||||
| Stage | Model component | Stage output |
|
||||
| --- | --- | --- |
|
||||
| **Thinker** | AuT audio encoder, SigLIP2 vision encoder, and the 30B-total/3B-active Thinker MoE Transformer | Text tokens and conditioning for the Talker |
|
||||
| **Talker** | 3B-total/0.3B-active Talker MoE Transformer with the MTP code-prediction path | Text tokens and audio codec tokens |
|
||||
| **Code2Wav** | Approximately 200M-parameter ConvNet audio decoder | Waveform/audio chunks |
|
||||
|
||||
#### Representative stage execution
|
||||
|
||||
| Stage | Batching | Attention and execution | Parallelism | Quantization |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| **Thinker** | Continuous batching with stage-local `max_num_seqs` and token budgets | vLLM KV-cached AR attention; CUDA Graph execution when the stage is not eager | Independent stage placement; use the AR runtime's configured tensor parallelism | Supported ModelOpt FP8/NVFP4 or AutoRound checkpoint paths target the Thinker; encoders remain BF16 in the documented paths |
|
||||
| **Talker** | Continuous batching over autoregressive decode requests | KV-cached token decode with async-chunk output to Code2Wav | Independent stage placement; typically a separate GPU or replica from Thinker | BF16 baseline; no generic Talker quantization path is assumed |
|
||||
| **Code2Wav** | Static/chunk batching for codec-to-waveform decode | Non-AR chunked audio decode; eager execution is commonly retained | Separate stage or colocated placement, depending on the deployment | BF16 baseline; keep codec/vocoder precision model-specific |
|
||||
|
||||
For this pipeline, the main serving target is interactive responsiveness:
|
||||
TTFT for text, TTFP for the first audio packet, TPOT for autoregressive decode,
|
||||
E2EL for the complete answer, and RTF/throughput for sustained audio serving.
|
||||
Async chunking and streaming primarily reduce TTFP and improve stage overlap;
|
||||
batching and CUDA Graphs primarily improve E2EL and throughput. See the
|
||||
[Qwen3-Omni recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen3-Omni.md)
|
||||
and [async chunk design](feature/async_chunk.md) for the deployment details.
|
||||
|
||||
### HunyuanImage-3.0: shared multimodal model, split deployment choices
|
||||
|
||||
HunyuanImage-3.0 illustrates a shared multimodal model with flexible AR, DiT,
|
||||
and split AR → DiT deployment. The official [technical
|
||||
report](https://arxiv.org/abs/2509.23951) provides the model view; the tables
|
||||
below map it to vLLM-Omni stages.
|
||||
**AR as a main structure, with DiT as multi-modal generator (e.g. BAGEL)**
|
||||
A unified multimodal comprehension and generation model, with cot text output and visual generation.
|
||||
|
||||
<p align="center">
|
||||
<a href="https://github.com/Tencent-Hunyuan/HunyuanImage-3.0/blob/main/assets/framework.png">
|
||||
<img alt="HunyuanImage-3.0 framework for understanding, language modeling, and image generation" src="https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/framework.png" width="95%">
|
||||
</a>
|
||||
<picture>
|
||||
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/ar-main-architecture.png">
|
||||
<img alt="Bagel" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/ar-main-architecture.png" width=30%>
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
#### Model architecture by stage
|
||||
|
||||
| Stage or component | Model component | Stage output |
|
||||
| --- | --- | --- |
|
||||
| **AR understanding/reasoning** | Understanding and generation encoders feeding the Hunyuan-A13B decoder-only Transformer and text detokenizer | Text responses or image-generation conditioning |
|
||||
| **Image generation** | Generation encoder, diffusion prediction path/Gen. Decoder, and the image VAE decode path | Image latents and final images |
|
||||
|
||||
The vLLM-Omni deployment may expose the AR path only, the DiT path only, or
|
||||
both as an AR → DiT pipeline with a connector between them. That is a runtime
|
||||
decomposition of the model's capabilities; it should not be read as a claim
|
||||
that the published framework figure contains two unrelated backbones.
|
||||
|
||||
#### Representative stage execution
|
||||
|
||||
| Stage | Batching | Attention and execution | Parallelism | Quantization |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| **AR** | Standard vLLM continuous batching when serving understanding/text tasks | KV-cached causal attention | Tensor/expert parallelism follows the selected AR deployment | Checkpoint-specific; do not infer DiT FP8 settings for the AR checkpoint |
|
||||
| **DiT + VAE** | Request batching or step-wise batching; Hunyuan step batching with more than one request requires `TORCH_SDPA` | `TORCH_SDPA` is used because the model mixes causal and full attention | Validated profiles include TP4, TP2 + sequence parallelism, TP2 + CFG parallelism, and expert parallelism for MoE weights | Online FP8 and ModelOpt mixed FP8/NVFP4 are documented for the DiT; VAE precision stays separate |
|
||||
|
||||
Image generation is therefore measured mainly with per-image E2EL/denoising
|
||||
latency, peak memory, and images-per-second throughput. TTFT and TPOT remain
|
||||
meaningful for the optional AR understanding path, but they are not the right
|
||||
headline metrics for the DiT path. See the
|
||||
[HunyuanImage-3.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanImage-3.0-Instruct.md)
|
||||
for the validated batching, attention, and quantization combinations.
|
||||
|
||||
### MiniMax-H3: conditioning, denoising, and decode
|
||||
|
||||
MiniMax-H3 illustrates logical component and phase boundaries between text
|
||||
encoding, task-specific denoising, and media decoding:
|
||||
`Text Encoder → DiT → VAE Decoder`. The current registered topology co-locates
|
||||
these components in one diffusion stage. These boundaries are useful for
|
||||
profiling and offload; independent placement would require a future
|
||||
multi-stage `PipelineConfig`. The stage mapping follows the implementation and
|
||||
serving recipe; the [release overview](https://minimaxi.com/blog/minimax-h3)
|
||||
notes that a detailed technical report is forthcoming.
|
||||
|
||||
#### Model architecture by component
|
||||
|
||||
| Logical component | Model component | Component output |
|
||||
| --- | --- | --- |
|
||||
| **Text Encoder** | Tokenizer/processor and Qwen3-VL text encoder; vision/audio references are prepared alongside the text condition | Packed multimodal conditioning for denoising |
|
||||
| **DiT** | Task-specific `FL2VA` DiT for text/first-frame generation or `Ref2VA` DiT for mixed references | Video and audio latent tokens |
|
||||
| **VAE Decoder** | Video VAE and audio VAE decode the generated latents and assemble the synchronized output | Decoded video frames and stereo waveform, with FPS and audio sample rate metadata |
|
||||
|
||||
Both DiT partitions remain behind one task-selection point; `task` chooses
|
||||
which partition runs. In the current implementation, the three logical
|
||||
components execute inside one diffusion stage. This is different from
|
||||
Qwen3-Omni's three independently scheduled stages, even though both models
|
||||
produce synchronized audio and video-like outputs.
|
||||
|
||||
#### Representative component execution
|
||||
|
||||
| Logical component | Batching | Attention and execution | Parallelism | Quantization |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| **Text Encoder** | The current co-located pipeline follows the diffusion request scheduler; independent prompt/reference batching would require a separate multi-stage topology | Qwen3-VL attention and multimodal preprocessing | `--text-encoder-tp-size N` shards the encoder across the first `N` DiT ranks; independent placement is not part of the current topology | BF16/FP32 baseline; the H3 DiT FP8 path does not quantize the text encoder |
|
||||
| **DiT** | The current H3 implementation executes one generation request per diffusion batch; use concurrency for service-level throughput | cuDNN attention for the validated two-GPU consumer profile; TRTLLM attention or FlashAttention-4 on supported Blackwell profiles | TP2 plus text-encoder TP and Ulysses/VAE parallel groups on multi-GPU profiles; DLO/CPU offload trade memory for transfer time | Online FP8 applies to eligible DiT linears and is incompatible with layerwise offload |
|
||||
| **VAE Decoder** | Decode follows each generation request; tile/patch work can be distributed even when denoising is not batched | VAE decode kernels rather than DiT attention | VAE patch parallelism and native tiled decode within the diffusion stage | BF16/FP32 baseline; the documented H3 FP8 path leaves both VAEs unchanged |
|
||||
|
||||
The primary target is video/audio E2EL and media throughput, with the logical
|
||||
component boundaries making prompt encoding, denoising, and VAE decode costs
|
||||
separately measurable. RTF is useful when the audio stream is evaluated
|
||||
independently. TTFT and TPOT do not describe the main H3 generation path
|
||||
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.
|
||||
|
||||
### Cosmos3: unified MoT reasoner and generator
|
||||
|
||||
Cosmos3 illustrates a unified Mixture-of-Transformers (MoT) reasoner and
|
||||
diffusion generator. The official [technical
|
||||
report](https://research.nvidia.com/labs/cosmos-lab/cosmos3/technical-report.pdf)
|
||||
and NVIDIA's [technical explanation](https://developer.nvidia.com/blog/develop-physical-ai-reasoning-world-and-action-models-with-nvidia-cosmos-3/)
|
||||
provide the two-tower model view used below.
|
||||
**AR+DiT (e.g. Qwen-Omni)**
|
||||
A natively end-to-end omni-modal LLM for multimodal inputs (text/image/audio/video...) and outputs (text/audio...).
|
||||
|
||||
<p align="center">
|
||||
<a href="https://developer.nvidia.com/blog/develop-physical-ai-reasoning-world-and-action-models-with-nvidia-cosmos-3/">
|
||||
<img alt="Cosmos3 autoregressive reasoner and diffusion generator towers" src="https://developer-blogs.nvidia.com/wp-content/uploads/2026/05/image-11-e1780000686151.webp" width="95%">
|
||||
</a>
|
||||
<picture>
|
||||
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/ar-dit-main-architecture.png">
|
||||
<img alt="Qwen-Omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/ar-dit-main-architecture.png" width=30%>
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
#### Model architecture by stage
|
||||
# vLLM-Omni main architecture
|
||||
|
||||
| Stage or component | Model component | Stage output |
|
||||
| --- | --- | --- |
|
||||
| **Reasoner tower** | Autoregressive VLM path with causal attention over language and optional vision/action context; runs once per generation request | Discrete reasoning/context tokens |
|
||||
| **Generator tower** | Diffusion path with full attention over the generator tokens and the reasoner K/V context; runs for each denoising step | Image, video, audio, or action tokens |
|
||||
| **Shared pipeline** | `Cosmos3OmniDiffusersPipeline` plus modality encoders/decoders and VAE paths | Final image/video/audio or action response |
|
||||
<p align="center">
|
||||
<picture>
|
||||
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-main-architecture.png">
|
||||
<img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-main-architecture.png" width=55%>
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
The reasoner and generator are model-owned components, not two independent
|
||||
vLLM-Omni stages by default. The current Cosmos3 pipeline materializes both
|
||||
inside one diffusion stage; model-level CPU offload can swap the reasoner and
|
||||
generator components between their phases.
|
||||
## Key Components
|
||||
|
||||
#### Representative stage execution
|
||||
| Component | Description |
|
||||
| ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| **OmniRouter** | provide an intelligent router for Omni-modality requests dispatch |
|
||||
| **EntryPoints** | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni), while `AsyncOmniEngine` and `Orchestrator` coordinate multi-stage AR/DiT execution |
|
||||
| **AR** | adapted for omni-modality models while inheriting efficient features from vLLM, such as cache management |
|
||||
| **Diffusion** | natively implemented and optimized using acceleration components |
|
||||
| **OmniConnector** | supports fully disaggregation based on E/P/D/G (Encoding/Processing/Decoding/Generation) disaggregation across stages |
|
||||
|
||||
| Stage | Batching | Attention and execution | Parallelism | Quantization |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| **Reasoner + generator pipeline** | Use request batching or step execution only when the selected pipeline advertises that capability; otherwise use `max_num_seqs=1` | Causal attention for reasoner tokens and full attention for diffusion tokens; the validated recipe uses a platform-selected aiter FlashAttention path with `TORCH_SDPA` fallback | Ulysses/context parallelism or tensor parallelism for the transformer; VAE tiling and layerwise/model-level offload address activation and weight memory | Online FP8 is supported for the diffusion pipeline; AutoRound/checkpoint-specific paths and quality validation remain separate from the BF16 baseline |
|
||||
|
||||
For Cosmos3, generation E2EL and output throughput are the primary targets.
|
||||
Action-policy workloads add a control-loop latency target, while synchronized
|
||||
audio workloads add RTF. TTFT/TPOT are appropriate for a reasoner-only service,
|
||||
but not for the full diffusion generation path. See the
|
||||
[Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-Nano.md)
|
||||
and [diffusion execution modes](../user_guide/diffusion/execution_modes.md)
|
||||
for the current runtime choices.
|
||||
|
||||
## Why the stage-based architecture follows from these examples
|
||||
|
||||
The four representative pipelines show that omni-modality serving is not one
|
||||
execution problem. Autoregressive decoding, diffusion denoising, multimodal
|
||||
encoding, and media decoding have different scheduling, attention, parallelism,
|
||||
memory, and quantization requirements. The architecture therefore needs a
|
||||
unit that can be optimized and resourced independently while remaining part of
|
||||
one request pipeline.
|
||||
|
||||
| Observation | Design implication | vLLM-Omni response |
|
||||
| --- | --- | --- |
|
||||
| AR, diffusion, encoder, and decoder components have different execution loops | One scheduler and one execution policy cannot fit every component | Specialized AR, generation, and diffusion stage runtimes |
|
||||
| A model can be deployed as AR-only, DiT-only, or AR → DiT | Model structure must be separate from deployment topology | `PipelineConfig` describes logical stages and relationships; `DeployConfig` describes placement and resources |
|
||||
| Text encoders, DiTs, and VAE decoders have different batch and memory behavior | Batching, attention, parallelism, and quantization must be stage-local | Each stage carries its own runtime and engine configuration |
|
||||
| MiniMax-H3 can co-locate or logically split its text encoder and VAE decoder | A stage boundary should be configurable rather than forcing a process boundary | Logical stages can be co-located or assigned independent runtime placement |
|
||||
| Qwen3-Omni transfers hidden states and codec chunks across stages | Large intermediate data needs typed transport and synchronization | `OmniConnector` transports payloads and KV-cache data without owning model routing |
|
||||
| Requests span multiple stages and may be cancelled or produce ordered streaming output | Cross-stage request state needs a separate control plane | `Orchestrator` owns correlation, cancellation, routing state, and output ordering |
|
||||
| Qwen3-Omni optimizes first-packet latency while image/video models optimize E2EL and throughput | Performance objectives must be model- and stage-aware | TTFT/TPOT/TTFP, E2EL, RTF, and throughput are measured according to the output modality |
|
||||
|
||||
This leads to a precise definition of a stage: it is a logical execution unit
|
||||
with its own model component, execution policy, resource ownership, and
|
||||
performance objectives. A stage is not necessarily a separate process. Cosmos3
|
||||
keeps its reasoner and generator inside one diffusion pipeline, while
|
||||
MiniMax-H3 exposes useful logical boundaries between text encoding, denoising,
|
||||
and VAE decoding. Both are valid mappings of the same abstraction.
|
||||
|
||||
The resulting separation can be summarized as:
|
||||
|
||||
```text
|
||||
Model structure What is computed and how components are related
|
||||
↓
|
||||
PipelineConfig Which logical stages exist and how they are connected
|
||||
↓
|
||||
Stage runtime policy How each stage batches, attends, parallelizes, and quantizes
|
||||
↓
|
||||
DeployConfig Where stages run, with which devices, replicas, and connectors
|
||||
```
|
||||
|
||||
`AsyncOmniEngine` and `Orchestrator` coordinate the pipeline, `OmniConnector`
|
||||
transports stage outputs, and `StageRuntime` materializes the selected
|
||||
placement. This separation lets one logical model support multiple serving
|
||||
topologies without coupling model code to process layout or deployment
|
||||
resources.
|
||||
|
||||
## System architecture
|
||||
|
||||
The figure shows the main components of the stage-based runtime: entrypoints,
|
||||
the shared `AsyncOmniEngine`, request orchestration, stage lifecycle
|
||||
management, AR and diffusion execution modules, model/layer operations,
|
||||
connector transport, and multimodal outputs. `AsyncOmniEngine` is the
|
||||
composition root between the entrypoints and the orchestrator.
|
||||
|
||||
```mermaid
|
||||
flowchart TB
|
||||
subgraph omni["vLLM-Omni"]
|
||||
subgraph entry["Entry points"]
|
||||
api["API server"]
|
||||
client["Omni / AsyncOmni"]
|
||||
end
|
||||
engine["AsyncOmniEngine<br/>composition root and background loop"]
|
||||
orchestrator["Orchestrator<br/>request state and cross-stage routing"]
|
||||
runtime["StageRuntime<br/>placement, replicas, readiness, and lifecycle"]
|
||||
ar["AR module<br/>vLLM engine, scheduler, cache, executor, worker, and model runner"]
|
||||
diffusion["Diffusion module<br/>scheduler, worker, model runner, and pipeline"]
|
||||
ops["Model / layer / ops<br/>attention, parallelism, and quantization"]
|
||||
connector["OmniConnector<br/>payload and KV transport plus synchronization"]
|
||||
outputs["Multimodal outputs<br/>streaming chunks and final artifacts"]
|
||||
|
||||
api --> engine
|
||||
client --> engine
|
||||
engine --> orchestrator
|
||||
engine --> runtime
|
||||
orchestrator --> runtime
|
||||
runtime --> ar
|
||||
runtime --> diffusion
|
||||
ar -. uses .-> ops
|
||||
diffusion -. uses .-> ops
|
||||
ar -. transports through .-> connector
|
||||
diffusion -. transports through .-> connector
|
||||
ar --> outputs
|
||||
diffusion --> outputs
|
||||
end
|
||||
```
|
||||
|
||||
### Key components
|
||||
|
||||
| Component | Responsibility |
|
||||
| --- | --- |
|
||||
| **Entrypoints** | Translate offline, CLI, OpenAI-compatible, and duplex requests into engine operations and render outputs back to public protocols. |
|
||||
| **Configuration resolution** | Combines pipeline topology, deployment settings, model metadata, and user overrides into a validated control-plane configuration. |
|
||||
| **AsyncOmniEngine** | Owns engine composition, the background event loop, stage initialization, request submission, and output collection. |
|
||||
| **Orchestrator** | Owns cross-stage request state, stage-to-stage routing, request correlation, cancellation, and output ordering. It does not own model selection or deployment placement. |
|
||||
| **StageRuntime** | Expands logical stages into local or distributed replicas, starts stage clients and processes, and manages readiness, affinity, failure, and shutdown. |
|
||||
| **AR runtime** | Extends vLLM's scheduler, KV-cache, worker, and model-runner path for omni-modality inputs and inter-stage outputs. |
|
||||
| **Diffusion runtime** | Schedules and executes denoising workloads through diffusion executors, workers, pipelines, acceleration backends, and output materialization. |
|
||||
| **OmniConnector** | Transports stage payloads and KV-cache data and provides synchronization. Connectors transport data; they do not choose the next logical stage. |
|
||||
| **Multimodal outputs** | `MultimodalPayload` separates tensor content from metadata, while `OmniRequestOutput` carries pipeline and diffusion results through the common output path. |
|
||||
|
||||
## Configuration and runtime resolution
|
||||
|
||||
The control plane has five conceptual layers. Authoring inputs are resolved
|
||||
once into a complete, transport-safe configuration before runtime processes are
|
||||
started. This keeps stage topology and model capabilities distinct from
|
||||
deployment placement and from process-local engine objects.
|
||||
|
||||
The figure shows only the primary hand-off object for each layer; the detailed
|
||||
inputs, fields, and ownership rules are described below.
|
||||
|
||||
```mermaid
|
||||
flowchart TB
|
||||
layer1["Layer 1 · Authoring inputs<br/>PipelineConfig + DeployConfig"]
|
||||
layer2["Layer 2 · Resolve once<br/>StageConfigFactory.create_from_model()<br/>VllmOmniConfig.from_pipeline_config()"]
|
||||
layer3["Layer 3 · Transport-safe control plane<br/>VllmOmniConfig"]
|
||||
layer4["Layer 4 · Runtime launch planning<br/>StageRuntime"]
|
||||
layer5["Layer 5 · Engine materialization<br/>VllmConfig / OmniDiffusionConfig"]
|
||||
|
||||
layer1 --> layer2 --> layer3 --> layer4 --> layer5
|
||||
```
|
||||
|
||||
The single resolution boundary is implemented by
|
||||
`StageConfigFactory.create_from_model()` and
|
||||
`VllmOmniConfig.from_pipeline_config()` in
|
||||
[`vllm_omni/config`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/config).
|
||||
The legacy `stage_args` YAML path remains only for models that have not yet
|
||||
migrated to `PipelineConfig` and `DeployConfig`.
|
||||
|
||||
In the typed path, each stage configuration derives from
|
||||
`BaseVllmOmniStageConfig` and is specialized as
|
||||
`VllmOmniARStageConfig`, `VllmOmniGenerationStageConfig`, or
|
||||
`VllmOmniDiffusionStageConfig`. The request and engine-spec fields belong to the
|
||||
control-plane boundary; the current implementation stores their equivalent
|
||||
projections in the structured stage configuration and materializes
|
||||
backend-specific engine objects during stage initialization.
|
||||
|
||||
The important ownership rules are:
|
||||
|
||||
1. `PipelineConfig` is the source of truth for stage topology, execution type,
|
||||
model capabilities, and stage relationships.
|
||||
2. `DeployConfig` describes placement, replicas, devices, connectors, and
|
||||
deploy-time defaults; it does not redefine the model graph.
|
||||
3. CLI and Python overrides are applied at the resolution boundary, with
|
||||
per-stage overrides taking precedence over global values where supported.
|
||||
4. `StageRuntime` owns launch planning and replica lifecycle. `ReplicaInitPlan`
|
||||
is runtime-private state, not a user configuration object.
|
||||
5. `VllmConfig` and the enriched `OmniDiffusionConfig` are materialized in the
|
||||
process that owns the corresponding engine.
|
||||
Disaggregated stages are managed through stage configuration. In Qwen3-Omni, Thinker/Talker/Code2wav are declared as separate configured stages, and runtime routing is handled by `Orchestrator` over `StageEngineCoreClient` / `StageDiffusionClient`.
|
||||
|
||||
## Main features
|
||||
|
||||
The feature surface is grouped to match the
|
||||
[feature design documents](index.md#feature-design-documents). This page
|
||||
summarizes each feature's architectural role; the linked design document is
|
||||
the source for configuration, compatibility, and implementation details.
|
||||
vLLM-Omni aims to be fast, flexible, and easy to use with the following features:
|
||||
|
||||
### Runtime and stage execution
|
||||
### Performance and Acceleration
|
||||
|
||||
* **Disaggregated inference:** Logical stages can run in separate processes,
|
||||
devices, or nodes while the orchestrator preserves their declared
|
||||
relationships. `OmniConnector` implementations transfer stage data and
|
||||
control-plane metadata. See [Disaggregated Inference](feature/disaggregated_inference.md).
|
||||
* **Asynchronous stage and output execution:** [Async Chunk](feature/async_chunk.md)
|
||||
forwards partial stage outputs as they become available. [Async Diffusion
|
||||
Output](feature/async_diffusion_output.md) overlaps device-to-host output
|
||||
packing with the next diffusion request, while [Async Omni Output
|
||||
Materialization](feature/omni_async_output_materialization.md) moves
|
||||
CPU-side payload construction off the AR decode critical path.
|
||||
* **Automatic prefix caching:** [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md)
|
||||
reuses KV-cache-aligned stage outputs and multimodal tensors for requests
|
||||
with common prefixes.
|
||||
The framework achieves high performance through several optimization techniques:
|
||||
|
||||
### Communication
|
||||
* **Efficient AR Support:** Leverages efficient KV cache management inherited from vLLM.
|
||||
* **Pipelined Execution:** Uses pipelined stage execution overlapping to ensure high throughput.
|
||||
* **Full Disaggregation:** Relies on the OmniConnector and dynamic resource allocation across stages.
|
||||
* **Diffusion Acceleration:** Includes integrated support for diffusion acceleration. This is managed by the acceleration layer, which handles:
|
||||
* **Cache:** Includes DBCache, TeaCache and third-party integration(e.g., [cache-dit](https://github.com/vipshop/cache-dit)).
|
||||
* **Parallelism:** Supports TP, CP, USP, and CFG.
|
||||
* **Attention:** Provides an interface for third-party integration (e.g., FA3, SAGE, MindIE-SD).
|
||||
* **Quantization:** Supports various quantization implementations including FP8 and AWQ.
|
||||
* **FusedOps:** Allows for custom and third-party integration.
|
||||
|
||||
* **OmniConnector transport:** The connector contract carries tensors, KV-cache
|
||||
data, and transport metadata across stage boundaries. The available
|
||||
implementations cover shared memory and multi-node Mooncake, Mori, and
|
||||
Yuanrong transports; see [Disaggregated Inference](feature/disaggregated_inference.md)
|
||||
for the connector choices and configuration model.
|
||||
### Classifier-Free Guidance (CFG) Companion Flow
|
||||
|
||||
### Diffusion acceleration
|
||||
vLLM-Omni natively models Classifier-Free Guidance (CFG) across disaggregated multi-stage setups via a "companion request" paradigm, eliminating redundant textual/multimodal context computation boundaries:
|
||||
1. **Prompt Expansion:** In the initial autoregressive (AR) stage, a customized `prompt_expand_func` hook intercepts incoming generation prompts and pairs them directly with negative companion prompts (e.g., a default negative prompt) on the fly, tagging the secondary prompt with a specific internal role (`cfg_text`).
|
||||
2. **Synchronized KV Cache Transfer:** The AR stage evaluates both the primary and companion sequence batches concurrently. The `OmniConnector` captures these specific structural dependencies and reliably passes the positive and negative outcome KV caches seamlessly across stage boundaries via shared memory or network protocols.
|
||||
3. **KV Cache Collection & Injection:** Upon reaching the downstream Diffusion (DiT) Engine, an assigned `cfg_kv_collect_func` automatically intercepts the mapped companion caches (`cfg_text_past_key_values`). These auxiliary dependencies are natively gathered and seamlessly bound to the primary generation sequence variables, enabling the DiT Engine to cleanly implement cross-attention CFG guidance over accurate conditioning and unconditioning structures in parallel.
|
||||
|
||||
* **Request and step batching:** [Diffusion Continuous Batching](feature/diffusion_continuous_batching.md)
|
||||
defines request-batch and step-batch execution, scheduler admission, and the
|
||||
common streaming output path.
|
||||
* **Composable parallelism:** Diffusion stages can combine [CFG-Parallel](feature/cfg_parallel.md),
|
||||
[Expert Parallel](feature/expert_parallel.md), [HSDP](feature/hsdp.md),
|
||||
[Pipeline Parallel](feature/pipeline_parallel.md), [Sequence Parallel](feature/sequence_parallel.md),
|
||||
[Tensor Parallel](feature/tensor_parallel.md), and [VAE Patch Parallelism](feature/vae_parallel.md)
|
||||
according to the pipeline and hardware topology.
|
||||
* **Attention and cache acceleration:** [Skip-Softmax](feature/skip_softmax.md),
|
||||
[Cache-DiT](feature/cache_dit.md), and [TeaCache](feature/teacache.md)
|
||||
provide backend and denoising-step optimizations without changing the
|
||||
stage contract.
|
||||
* **Quantization and memory efficiency:** [Quantization](feature/quantization.md)
|
||||
resolves per-pipeline or per-component quantization configurations, while
|
||||
[Distributed Layerwise Offload](feature/distributed_layerwise_offload.md)
|
||||
streams diffusion blocks from host memory within the existing parallel
|
||||
topology.
|
||||
### Flexibility and Usability
|
||||
|
||||
## Interfaces
|
||||
vLLM-Omni is designed to be flexible and straightforward for users:
|
||||
|
||||
The public interfaces map onto the same engine and stage boundaries:
|
||||
* **Heterogeneous Pipeline Abstraction:** Manages complex model workflows effectively.
|
||||
* **Hugging Face Integration:** Offers seamless integration with popular Hugging Face models.
|
||||
* **Distributed Inference:** Supports tensor, pipeline, data, and expert parallelism.
|
||||
* **Streaming Outputs:** Supports streaming outputs.
|
||||
* **Unified API:** Provides a consistent and unified API interface compatible with vLLM.
|
||||
* **OpenAI-compatible API Server:** Includes a FastAPI-based server for online serving that is compatible with the OpenAI API.
|
||||
|
||||
# Interface design
|
||||
|
||||
If you use vLLM, then you know how to use vLLM-Omni from Day 0:
|
||||
|
||||
<p align="center">
|
||||
<picture>
|
||||
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-user-interface.png">
|
||||
<img alt="vLLM-Omni interface design" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-user-interface.png" width=55%>
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
Taking **Qwen3-Omni** as an example:
|
||||
|
||||
## Offline Inference
|
||||
The **Omni** class provides a Python interface for offline batched inference. Users initialize the Omni class with a Hugging Face model name and use the generate method, passing inputs that include both text prompts and multi-modal data:
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
offline["Offline Python<br/>Omni.generate()"] --> engine["AsyncOmniEngine"]
|
||||
online["OpenAI-compatible serving<br/>vllm serve ... --omni"] --> engine
|
||||
duplex["Experimental duplex WebSocket<br/>/v1/duplex or realtime duplex"] --> engine
|
||||
engine --> stages["Configured AR and diffusion stages"]
|
||||
stages --> result["Streaming or final multimodal output"]
|
||||
```
|
||||
|
||||
### Offline inference
|
||||
|
||||
The **Omni** class provides a Python interface for offline batched inference:
|
||||
|
||||
```python
|
||||
# Create an omni runtime with HF model name.
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
omni = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
|
||||
|
||||
om_inputs = {
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {
|
||||
"video": video_frames,
|
||||
"audio": audio_signal,
|
||||
},
|
||||
}
|
||||
# Example prompts.
|
||||
om_inputs = {"prompt": prompt,
|
||||
"multi_modal_data": {
|
||||
"video": video_frames,
|
||||
"audio": audio_signal,
|
||||
}}
|
||||
|
||||
# Generate texts and audio from the multi-modality inputs.
|
||||
outputs = omni.generate(om_inputs, sampling_params_list)
|
||||
```
|
||||
|
||||
### Online serving
|
||||
## Online Serving
|
||||
Similar to vLLM, vLLM-Omni also provides a FastAPI-based server for online serving. Users can launch the server using the vllm serve command with the `--omni` flag:
|
||||
|
||||
The OpenAI-compatible server uses the same stage configuration and engine
|
||||
boundaries:
|
||||
|
||||
```bash
|
||||
```
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||
```
|
||||
|
||||
For example, a Qwen3-Omni chat request can contain text, image, audio, or
|
||||
video content and a `sampling_params_list` for its configured stages. See the
|
||||
[Qwen3-Omni serving example](../user_guide/examples/online_serving/qwen3_omni.md)
|
||||
and the [examples](https://github.com/vllm-project/vllm-omni/tree/main/examples)
|
||||
for complete requests.
|
||||
Users can send requests to the server using curl:
|
||||
|
||||
Some pipelines expose additional OpenAI-compatible endpoints, such as joint
|
||||
video/audio generation. Endpoint support remains model-specific; consult the
|
||||
relevant model guide before assuming that every OpenAI route applies to every
|
||||
pipeline.
|
||||
```
|
||||
# prepare user content
|
||||
user_content='[
|
||||
{
|
||||
"type": "video_url",
|
||||
"video_url": {
|
||||
"url": "'"$SAMPLE_VIDEO_URL"'"
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Why is this video funny?"
|
||||
}
|
||||
]'
|
||||
sampling_params_list='[
|
||||
'"$thinker_sampling_params"',
|
||||
'"$talker_sampling_params"',
|
||||
'"$code2wav_sampling_params"'
|
||||
]'
|
||||
mm_processor_kwargs="{}"
|
||||
|
||||
# send the request
|
||||
curl -sS -X POST http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d @- <<EOF
|
||||
{
|
||||
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
|
||||
"sampling_params_list": $sampling_params_list,
|
||||
"mm_processor_kwargs": $mm_processor_kwargs,
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech."
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": $user_content
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
For more usages, please refer to [examples](https://github.com/vllm-project/vllm-omni/tree/main/examples).
|
||||
@@ -17,7 +17,7 @@ The full set of backends and their platform defaults is in the **Backend Options
|
||||
| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). Dense BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Packed paths can provide `cu_seqlens` directly. Supports optional **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. |
|
||||
| `FLASH_ATTN` | Wraps FlashAttention 4 on Blackwell when `flash-attn-4` is installed, then falls back to FlashAttention 3/2. Default on Hopper / Ada / Ampere when a compatible FlashAttention package is installed. |
|
||||
| `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). |
|
||||
| `FLASHINFER_ATTN` | Uses FlashInfer's batch-prefill wrapper and supports mixed Q/K and V input dtypes through backend-specific configuration. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer` >= 0.6.16rc1 for mixed-dtype configurations. |
|
||||
| `FLASHINFER_ATTN` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. |
|
||||
| `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
|
||||
| `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. |
|
||||
| `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. |
|
||||
@@ -84,7 +84,7 @@ A backend that needs configuration exposes it as a typed field on the spec:
|
||||
When constructing `OmniDiffusionConfig` directly:
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, OmniDiffusionConfig
|
||||
|
||||
config = OmniDiffusionConfig(
|
||||
diffusion_attention_config=AttentionConfig(
|
||||
@@ -99,27 +99,6 @@ config = OmniDiffusionConfig(
|
||||
|
||||
A plain dict is also accepted and normalized to `AttentionConfig`.
|
||||
|
||||
#### FlashInfer QK16/V8 quantized attention
|
||||
|
||||
`FLASHINFER_ATTN` accepts quantization options through `AttentionSpec.quant`. For QK16/V8:
|
||||
|
||||
```python
|
||||
config = OmniDiffusionConfig(
|
||||
diffusion_attention_config=AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="FLASHINFER_ATTN",
|
||||
quant=AttnQuantSpec(
|
||||
dtype_qk="bfloat16",
|
||||
dtype_vo="fp8_e4m3",
|
||||
),
|
||||
),
|
||||
),
|
||||
...,
|
||||
)
|
||||
```
|
||||
|
||||
`dtype_qk` configures the Q and K input dtypes. `dtype_vo` configures the V input dtype.
|
||||
|
||||
## Platform Defaults
|
||||
|
||||
### Blackwell (sm_100 / sm_103 / sm_120 / sm_121)
|
||||
|
||||
@@ -6,7 +6,6 @@
|
||||
- [Overview](#overview)
|
||||
- [Quick Start](#quick-start)
|
||||
- [Example Script](#example-script)
|
||||
- [Request-Scoped Quality](#request-scoped-quality-minimax-h3)
|
||||
- [Acceleration Methods](#acceleration-methods)
|
||||
- [Configuration Parameters](#configuration-parameters)
|
||||
- [Best Practices](#best-practices)
|
||||
@@ -141,43 +140,6 @@ vllm serve Qwen/Qwen-Image --omni --port 8091 \
|
||||
--cache-config '{"Fn_compute_blocks": 1, "residual_diff_threshold": 0.12}'
|
||||
```
|
||||
|
||||
## Request-Scoped Quality (MiniMax H3)
|
||||
|
||||
MiniMax H3 can switch Cache-DiT at the request boundary without enabling a
|
||||
cache backend at server startup. Add a quality field to an online request:
|
||||
|
||||
```bash
|
||||
# Native reference path for this request.
|
||||
-F 'quality=lossless'
|
||||
|
||||
# H3's conservative Cache-DiT profile for this request.
|
||||
-F 'quality=high'
|
||||
```
|
||||
|
||||
For offline inference, set the same field on the sampling parameters:
|
||||
|
||||
```python
|
||||
OmniDiffusionSamplingParams(
|
||||
quality="high",
|
||||
num_inference_steps=50,
|
||||
extra_args={"task": "t2va", "duration": 5.0, "audio_flow_shift": 3.0},
|
||||
)
|
||||
```
|
||||
|
||||
| Server startup | Request quality | Behavior |
|
||||
|---|---|---|
|
||||
| `--cache-backend cache_dit` | omitted | Use or restore the startup Cache-DiT profile |
|
||||
| `--cache-backend cache_dit` | `lossless` | Remove Cache-DiT hooks and run the reference path |
|
||||
| `--cache-backend cache_dit` | `high` | Use or install H3's conservative Cache-DiT profile |
|
||||
| no Cache-DiT | omitted or `lossless` | Run the reference path |
|
||||
| no Cache-DiT | `high` | Install H3's conservative Cache-DiT profile |
|
||||
|
||||
The startup option is only needed when omitted-quality requests should use the
|
||||
server-configured Cache-DiT profile by default.
|
||||
|
||||
See the [MiniMax H3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md#request-scoped-quality)
|
||||
for a complete request and measured trade-off.
|
||||
|
||||
---
|
||||
|
||||
## Acceleration Methods
|
||||
|
||||
@@ -1,240 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Reproduce MiniMax-H3 request-mode and DLO/DP2 lifecycle behavior.
|
||||
|
||||
Examples:
|
||||
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
|
||||
VLLM_OMNI_DLO_DP_WAVE_TIMEOUT=600 \
|
||||
CUDA_VISIBLE_DEVICES=0,1 \
|
||||
python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
|
||||
--model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2
|
||||
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
|
||||
CUDA_VISIBLE_DEVICES=0,1 \
|
||||
python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
|
||||
--model /path/to/MiniMax-H3/FL2VA --mode request
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import copy
|
||||
import json
|
||||
import multiprocessing
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
|
||||
DEFAULT_PROMPTS = (
|
||||
"At night, three cats march into a bedroom playing tiny brass instruments, "
|
||||
"then abruptly file out, with synchronized room ambience.",
|
||||
"A paper boat crosses a rain-filled street while distant traffic and water sounds remain synchronized.",
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--model", required=True, help="Path to MiniMax-H3/FL2VA")
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
choices=("dlo-dp2", "request"),
|
||||
required=True,
|
||||
help="DLO AllGather DP2 or ordinary non-DLO TP2 request mode",
|
||||
)
|
||||
parser.add_argument("--steps", type=int, default=2)
|
||||
parser.add_argument("--duration", type=float, default=5.0)
|
||||
parser.add_argument("--width", type=int, default=1344)
|
||||
parser.add_argument("--height", type=int, default=768)
|
||||
parser.add_argument("--batch-wait-ms", type=float, default=500.0)
|
||||
parser.add_argument("--init-timeout", type=float, default=1800.0)
|
||||
parser.add_argument("--output", type=Path)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
|
||||
common: dict[str, Any] = {
|
||||
"model": args.model,
|
||||
"trust_remote_code": True,
|
||||
"num_gpus": 2,
|
||||
"usp": 1,
|
||||
"ring": 1,
|
||||
"vae_parallel_mode": "tile",
|
||||
"vae_use_tiling": True,
|
||||
"diffusion_attention_backend": "CUDNN_ATTN",
|
||||
"request_batch_max_wait_ms": args.batch_wait_ms,
|
||||
"enforce_eager": True,
|
||||
"stage_init_timeout": args.init_timeout,
|
||||
"init_timeout": args.init_timeout,
|
||||
}
|
||||
if args.mode == "dlo-dp2":
|
||||
common.update(
|
||||
tensor_parallel_size=1,
|
||||
data_parallel_size=2,
|
||||
text_encoder_tp_size=1,
|
||||
vae_patch_parallel_size=1,
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=True,
|
||||
dlo_resident_layers=0,
|
||||
)
|
||||
else:
|
||||
common.update(
|
||||
tensor_parallel_size=2,
|
||||
data_parallel_size=1,
|
||||
text_encoder_tp_size=2,
|
||||
vae_patch_parallel_size=2,
|
||||
enable_distributed_layerwise_offload=False,
|
||||
)
|
||||
return common
|
||||
|
||||
|
||||
def sampling_params(
|
||||
engine: AsyncOmni,
|
||||
args: argparse.Namespace,
|
||||
seed: int,
|
||||
) -> list[Any]:
|
||||
params = copy.deepcopy(engine.default_sampling_params_list)
|
||||
diffusion = params[0]
|
||||
diffusion.width = args.width
|
||||
diffusion.height = args.height
|
||||
diffusion.fps = 24
|
||||
diffusion.num_inference_steps = args.steps
|
||||
diffusion.seed = seed
|
||||
diffusion.extra_args = {
|
||||
"task": "t2va",
|
||||
"duration": args.duration,
|
||||
"aspect_ratio": "16:9",
|
||||
"flow_shift": 12.0,
|
||||
"audio_flow_shift": 3.0,
|
||||
}
|
||||
return params
|
||||
|
||||
|
||||
async def generate_one(
|
||||
engine: AsyncOmni,
|
||||
args: argparse.Namespace,
|
||||
*,
|
||||
request_id: str,
|
||||
prompt: str,
|
||||
seed: int,
|
||||
) -> Any:
|
||||
final_output = None
|
||||
async for output in engine.generate(
|
||||
prompt=prompt,
|
||||
request_id=request_id,
|
||||
sampling_params_list=sampling_params(engine, args, seed),
|
||||
):
|
||||
if output.finished:
|
||||
final_output = output
|
||||
if final_output is None:
|
||||
raise RuntimeError(f"{request_id} finished without an output")
|
||||
return final_output
|
||||
|
||||
|
||||
def output_summary(output: Any, args: argparse.Namespace) -> dict[str, Any]:
|
||||
if not output.images:
|
||||
raise RuntimeError(f"{output.request_id} returned no video")
|
||||
frames = np.asarray(output.images[0])
|
||||
audio = np.asarray(output.multimodal_output.get("audio"))
|
||||
if frames.ndim != 4 or tuple(frames.shape[1:]) != (
|
||||
args.height,
|
||||
args.width,
|
||||
3,
|
||||
):
|
||||
raise RuntimeError(f"{output.request_id} returned invalid video shape {tuple(frames.shape)}")
|
||||
if audio.ndim != 3 or tuple(audio.shape[:2]) != (1, 2):
|
||||
raise RuntimeError(f"{output.request_id} returned invalid audio shape {tuple(audio.shape)}")
|
||||
if args.duration == 5.0 and args.height == 768 and args.width == 1344:
|
||||
if tuple(frames.shape) != (124, 768, 1344, 3) or tuple(audio.shape) != (1, 2, 165600):
|
||||
raise RuntimeError(
|
||||
f"{output.request_id} default shape mismatch: video={tuple(frames.shape)}, audio={tuple(audio.shape)}"
|
||||
)
|
||||
return {
|
||||
"request_id": output.request_id,
|
||||
"frames_shape": list(frames.shape),
|
||||
"audio_shape": list(audio.shape),
|
||||
"peak_memory_mb": output.peak_memory_mb,
|
||||
"stage_durations": output.stage_durations,
|
||||
}
|
||||
|
||||
|
||||
async def run(args: argparse.Namespace) -> dict[str, Any]:
|
||||
engine = AsyncOmni(**engine_kwargs(args))
|
||||
summary: dict[str, Any] = {
|
||||
"mode": args.mode,
|
||||
"engine_kwargs": engine_kwargs(args),
|
||||
}
|
||||
try:
|
||||
if args.mode == "dlo-dp2":
|
||||
started = time.perf_counter()
|
||||
asymmetric = await asyncio.gather(
|
||||
generate_one(
|
||||
engine,
|
||||
args,
|
||||
request_id="invalid-empty",
|
||||
prompt="",
|
||||
seed=1001,
|
||||
),
|
||||
generate_one(
|
||||
engine,
|
||||
args,
|
||||
request_id="invalid-peer",
|
||||
prompt=DEFAULT_PROMPTS[0],
|
||||
seed=1002,
|
||||
),
|
||||
return_exceptions=True,
|
||||
)
|
||||
summary["asymmetric_wave_s"] = time.perf_counter() - started
|
||||
summary["asymmetric_errors"] = [
|
||||
f"{type(result).__name__}: {result}" for result in asymmetric if isinstance(result, BaseException)
|
||||
]
|
||||
if len(summary["asymmetric_errors"]) != 2:
|
||||
raise RuntimeError("Both requests in the asymmetric DP wave must fail before dispatch")
|
||||
|
||||
started = time.perf_counter()
|
||||
request_count = 2 if args.mode == "dlo-dp2" else 1
|
||||
outputs = await asyncio.gather(
|
||||
*(
|
||||
generate_one(
|
||||
engine,
|
||||
args,
|
||||
request_id=f"recovery-{index}",
|
||||
prompt=DEFAULT_PROMPTS[index],
|
||||
seed=2000 + index,
|
||||
)
|
||||
for index in range(request_count)
|
||||
)
|
||||
)
|
||||
summary["valid_wave_s"] = time.perf_counter() - started
|
||||
summary["outputs"] = [output_summary(output, args) for output in outputs]
|
||||
finally:
|
||||
started = time.perf_counter()
|
||||
engine.close()
|
||||
summary["shutdown_s"] = time.perf_counter() - started
|
||||
|
||||
children = multiprocessing.active_children()
|
||||
summary["active_children"] = [{"name": child.name, "pid": child.pid} for child in children]
|
||||
if children:
|
||||
raise RuntimeError(f"Worker processes remain after shutdown: {children}")
|
||||
return summary
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
summary = asyncio.run(run(args))
|
||||
rendered = json.dumps(summary, indent=2, sort_keys=True)
|
||||
print(f"E2E_RESULT {rendered}", flush=True)
|
||||
if args.output is not None:
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.output.write_text(rendered + "\n", encoding="utf-8")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -574,43 +574,10 @@ default. `short_edge` controls the 768-pixel canvas and must be `768`.
|
||||
outputs; the synchronous raw-MP4 endpoint returns the first output when more
|
||||
than one is requested.
|
||||
|
||||
## Request-scoped quality
|
||||
|
||||
Add one of these fields to any HTTP request above. No Cache-DiT startup option
|
||||
is required; H3 installs its conservative profile when a `high` request
|
||||
arrives and removes it for a `lossless` request.
|
||||
|
||||
```bash
|
||||
-F 'quality=lossless' # Native reference path for this request.
|
||||
-F 'quality=high' # H3's conservative Cache-DiT profile.
|
||||
```
|
||||
|
||||
Omitting `quality` preserves the startup default: it uses the reference path
|
||||
normally, or the server-configured profile when the server was started with
|
||||
`--cache-backend cache_dit`.
|
||||
|
||||
The following result was measured on 4× NVIDIA H200 with SP4, text-encoder
|
||||
TP4, 1344×768, 124 frames, 24 FPS, and 50 inference steps. One full
|
||||
`lossless` warmup was excluded, followed by three fixed prompt/seed pairs in
|
||||
balanced switch order.
|
||||
|
||||
| `quality` | Median inference latency | Speedup | SSIM vs `lossless` | PSNR vs `lossless` | Expected trade-off |
|
||||
|---|---:|---:|---:|---:|---|
|
||||
| `lossless` | 85.49 s | 1.00× | 1.0000 | exact | Native reference path |
|
||||
| `high` | 63.36 s | 1.35× | 0.9709 | 34.98 dB | Faster with measured same-seed deviation |
|
||||
|
||||
> [!NOTE]
|
||||
> `high` selects a fixed Cache-DiT profile, but its cache hit rate and
|
||||
> resulting latency/quality trade-off may vary by hardware, topology, and
|
||||
> workload. The values above apply to this deployment and are not universal
|
||||
> guarantees. `lossless` remains the exact reference path.
|
||||
|
||||
## Key parameters
|
||||
|
||||
| Parameter | Recommended value | Notes |
|
||||
|-----------|-------------------|-------|
|
||||
| `quality` | omitted or `lossless` | Request-level quality intent; `high` dynamically installs H3's conservative Cache-DiT profile |
|
||||
| `extra_params.force_refresh_step_hint` | omitted | Optional positive 1-based denoising-step hint for an active Cache-DiT request; pair with `extra_params.force_refresh_step_policy`=`once` or `repeat` |
|
||||
| `task` | `t2va`, `fl2va`, or `ref2va` | Passed in `extra_params`; selects the task-specific DiT |
|
||||
| `duration` | Workload-specific | Decimal seconds in `extra_params`; converted to H3-compatible frame count |
|
||||
| `fps` | `24` | H3 output FPS is fixed |
|
||||
|
||||
@@ -31,7 +31,6 @@ import tempfile
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, QueryRange
|
||||
from vllm_omni.diffusion.attention.layer import Attention
|
||||
from vllm_omni.diffusion.attention.parallel.allgather_kv import (
|
||||
@@ -410,7 +409,6 @@ def test_allgather_kv_keeps_gathered_kv_compressed_for_gqa():
|
||||
torch.testing.assert_close(v_full, expected_value)
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=4)
|
||||
@pytest.mark.parametrize(
|
||||
"test_model_cls",
|
||||
[
|
||||
@@ -418,12 +416,10 @@ def test_allgather_kv_keeps_gathered_kv_compressed_for_gqa():
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
("ulysses_degree", "ring_degree", "allgather_degree", "num_kv_heads"),
|
||||
("ulysses_degree", "ring_degree", "allgather_degree"),
|
||||
[
|
||||
pytest.param(2, 2, 1, None, id="ulysses-ring"),
|
||||
pytest.param(1, 1, 2, None, id="allgather-kv"),
|
||||
pytest.param(1, 2, 1, 2, id="ring-gqa"),
|
||||
pytest.param(1, 2, 1, 1, id="ring-mqa"),
|
||||
pytest.param(2, 2, 1, id="ulysses-ring"),
|
||||
pytest.param(1, 1, 2, id="allgather-kv"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("batch_size", [2])
|
||||
@@ -449,7 +445,6 @@ def test_sequence_parallel(
|
||||
seq_len: int,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
num_kv_heads: int | None,
|
||||
):
|
||||
"""Compare Ulysses/Ring and AllGather-KV SP against a single-rank run."""
|
||||
sequence_parallel_size = allgather_degree if allgather_degree > 1 else ulysses_degree * ring_degree
|
||||
@@ -481,7 +476,6 @@ def test_sequence_parallel(
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_size,
|
||||
num_kv_heads,
|
||||
dtype,
|
||||
causal,
|
||||
use_sync,
|
||||
@@ -513,7 +507,6 @@ def test_sequence_parallel(
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_size,
|
||||
num_kv_heads,
|
||||
dtype,
|
||||
causal,
|
||||
use_sync,
|
||||
@@ -610,7 +603,6 @@ def ulysses_attention_on_test_model(
|
||||
seq_len: int,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
num_kv_heads: int | None,
|
||||
dtype: torch.dtype,
|
||||
causal: bool,
|
||||
use_sync: bool,
|
||||
@@ -666,13 +658,10 @@ def ulysses_attention_on_test_model(
|
||||
cfg_parallel_size=1,
|
||||
)
|
||||
|
||||
od_config = OmniDiffusionConfig.from_kwargs(
|
||||
od_config = OmniDiffusionConfig(
|
||||
model="test_model",
|
||||
dtype=dtype,
|
||||
parallel_config=parallel_config,
|
||||
# This regression targets pytorch_attn_forward(). Do not let an
|
||||
# installed FA/FA3 backend silently bypass the SDPA Ring path.
|
||||
diffusion_attention_backend="TORCH_SDPA",
|
||||
)
|
||||
|
||||
# Initialize model parallel
|
||||
@@ -698,7 +687,7 @@ def ulysses_attention_on_test_model(
|
||||
"head_size": head_size,
|
||||
"hidden_size": hidden_size,
|
||||
"causal": causal,
|
||||
"num_kv_heads": num_kv_heads,
|
||||
"num_kv_heads": None,
|
||||
"scatter_idx": 2,
|
||||
"gather_idx": 1,
|
||||
"use_sync": use_sync,
|
||||
|
||||
@@ -6,14 +6,12 @@ from __future__ import annotations
|
||||
import os
|
||||
import socket
|
||||
import tempfile
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.attention.layer import Attention
|
||||
from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention
|
||||
from vllm_omni.diffusion.config import set_current_diffusion_config
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
@@ -41,33 +39,6 @@ def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None:
|
||||
os.environ["MASTER_PORT"] = str(master_port)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_advanced_uaa_hybrid_rejects_padded_mqa_heads(monkeypatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.attention.parallel.ulysses.get_ulysses_mode",
|
||||
lambda default: "advanced_uaa",
|
||||
)
|
||||
sp_group = SimpleNamespace(
|
||||
ulysses_group=None,
|
||||
ulysses_world_size=2,
|
||||
ulysses_rank=0,
|
||||
ring_world_size=2,
|
||||
)
|
||||
strategy = UlyssesParallelAttention(
|
||||
sp_group=sp_group,
|
||||
scatter_idx=2,
|
||||
gather_idx=1,
|
||||
use_sync=False,
|
||||
)
|
||||
query = torch.zeros(1, 2, 8, 4)
|
||||
key = torch.zeros(1, 2, 1, 4)
|
||||
value = torch.zeros_like(key)
|
||||
|
||||
with pytest.raises(ValueError, match="does not support K/V head padding"):
|
||||
strategy.pre_attention(query, key, value, None)
|
||||
|
||||
|
||||
def _run_attention_case(
|
||||
local_rank: int,
|
||||
world_size: int,
|
||||
@@ -270,9 +241,7 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None:
|
||||
batch_size = 2
|
||||
head_size = 8
|
||||
seq_len = 10
|
||||
# This is the positive Hybrid case: K/V heads must not require padding,
|
||||
# which is unsupported by advanced_uaa when Ring is also enabled.
|
||||
num_heads = 4
|
||||
num_heads = 3 # head_cnt not divisible by ulysses_degree=2 -> triggers head padding
|
||||
|
||||
# Ensure ring ranks see equal post-Ulysses seq_len:
|
||||
# rank0/1 -> 3+2=5, rank2/3 -> 3+2=5
|
||||
|
||||
+3
-80
@@ -11,7 +11,6 @@ This module tests the cache backend implementations:
|
||||
- DiffusionCacheConfig: configuration dataclass
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import cache_dit
|
||||
@@ -22,7 +21,6 @@ from vllm_omni.diffusion.cache.cachedit import (
|
||||
CacheDiTAdapterConfig,
|
||||
CacheDiTBackend,
|
||||
CacheDiTConfig,
|
||||
cache_summary,
|
||||
)
|
||||
from vllm_omni.diffusion.cache.magcache import MagCacheBackend
|
||||
from vllm_omni.diffusion.cache.selector import get_cache_backend
|
||||
@@ -93,8 +91,8 @@ class TestCacheDiTBackend:
|
||||
|
||||
@patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
|
||||
@patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
|
||||
def test_lifecycle_applies_to_every_declared_dit(self, mock_cache_dit, mock_block_adapter):
|
||||
"""Combined pipelines manage every exact Cache-DiT target."""
|
||||
def test_enable_and_refresh_every_declared_dit(self, mock_cache_dit, mock_block_adapter):
|
||||
"""Combined pipelines enable and refresh Cache-DiT for every DiT."""
|
||||
pipeline = Mock()
|
||||
pipeline.__class__.__name__ = "MiniMaxH3Pipeline"
|
||||
pipeline._dit_modules = ["transformer", "transformers_ref"]
|
||||
@@ -104,67 +102,18 @@ class TestCacheDiTBackend:
|
||||
transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
|
||||
block_forward_patterns={"blocks": ForwardPattern.Pattern_3}
|
||||
)
|
||||
installed_adapters = [Mock(), Mock()]
|
||||
mock_block_adapter.side_effect = installed_adapters
|
||||
|
||||
backend = CacheDiTBackend({"Fn_compute_blocks": 2})
|
||||
backend.enable(pipeline)
|
||||
backend.refresh(pipeline, num_inference_steps=20)
|
||||
backend.disable(pipeline)
|
||||
|
||||
assert len(backend._refresh_funcs) == 2
|
||||
assert mock_cache_dit.enable_cache.call_count == 2
|
||||
assert mock_cache_dit.refresh_context.call_count == 2
|
||||
assert {call.args[0] for call in mock_cache_dit.refresh_context.call_args_list} == {
|
||||
pipeline.transformer,
|
||||
pipeline.transformers_ref,
|
||||
}
|
||||
assert [call.args[0] for call in mock_cache_dit.disable_cache.call_args_list] == installed_adapters
|
||||
assert backend._refresh_funcs == []
|
||||
assert backend._cache_targets == []
|
||||
assert not backend.is_enabled()
|
||||
|
||||
@patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
|
||||
def test_summary_skips_uncached_nested_dit(self, mock_cache_dit):
|
||||
"""Only DiT modules with an active Cache-DiT context are summarized."""
|
||||
language_model = Mock()
|
||||
language_model._is_cached = False
|
||||
transformer = Mock()
|
||||
transformer._is_cached = True
|
||||
transformer.language_model = language_model
|
||||
pipeline = SimpleNamespace(
|
||||
_dit_modules=["transformer.language_model", "transformer"],
|
||||
transformer=transformer,
|
||||
)
|
||||
|
||||
cache_summary(pipeline)
|
||||
|
||||
mock_cache_dit.summary.assert_called_once_with(transformer, details=True)
|
||||
|
||||
@patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
|
||||
@patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
|
||||
def test_enable_and_refresh_nested_declared_dit(self, mock_cache_dit, mock_block_adapter):
|
||||
"""Dotted component paths resolve to their nested DiT module."""
|
||||
transformer = Mock()
|
||||
transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
|
||||
block_forward_patterns={"blocks": ForwardPattern.Pattern_3}
|
||||
)
|
||||
pipeline = SimpleNamespace(
|
||||
_dit_modules=["language_model.model"],
|
||||
language_model=SimpleNamespace(model=transformer),
|
||||
)
|
||||
|
||||
backend = CacheDiTBackend({"Fn_compute_blocks": 2})
|
||||
backend.enable(pipeline)
|
||||
backend.refresh(pipeline, num_inference_steps=20)
|
||||
|
||||
cache_summary(pipeline)
|
||||
mock_cache_dit.enable_cache.assert_called_once()
|
||||
mock_cache_dit.refresh_context.assert_called_once_with(
|
||||
transformer,
|
||||
num_inference_steps=20,
|
||||
verbose=True,
|
||||
)
|
||||
mock_cache_dit.summary.assert_called_once_with(transformer, details=True)
|
||||
|
||||
@patch("vllm_omni.diffusion.cache.cachedit.backend.logger")
|
||||
@patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
|
||||
@@ -297,32 +246,6 @@ class TestCacheDiTBackend:
|
||||
|
||||
assert mock_cache_dit.refresh_context.call_count == 2
|
||||
|
||||
@patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
|
||||
@patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
|
||||
def test_refresh_rearms_force_refresh_hint(self, mock_cache_dit, mock_block_adapter):
|
||||
"""A once hint must be restored for every repeated request."""
|
||||
pipeline = Mock()
|
||||
pipeline.__class__.__name__ = "DiTPipeline"
|
||||
pipeline.transformer = Mock()
|
||||
pipeline.transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
|
||||
block_forward_patterns={"layers": ForwardPattern.Pattern_0}
|
||||
)
|
||||
|
||||
config = DiffusionCacheConfig(
|
||||
force_refresh_step_hint=7,
|
||||
force_refresh_step_policy="once",
|
||||
)
|
||||
backend = CacheDiTBackend(config)
|
||||
backend.enable(pipeline)
|
||||
backend.refresh(pipeline, num_inference_steps=20)
|
||||
backend.refresh(pipeline, num_inference_steps=20)
|
||||
|
||||
assert mock_cache_dit.refresh_context.call_count == 2
|
||||
for call in mock_cache_dit.refresh_context.call_args_list:
|
||||
refresh_config = call.kwargs["cache_config"]
|
||||
assert refresh_config.force_refresh_step_hint == 7
|
||||
assert refresh_config.force_refresh_step_policy == "once"
|
||||
|
||||
@patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
|
||||
@patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
|
||||
def test_enable_hunyuan_pipeline_uses_model_transformer(self, mock_cache_dit, mock_block_adapter):
|
||||
|
||||
-2
@@ -81,8 +81,6 @@ def test_cachedit_public_api_is_explicit():
|
||||
"CacheDiTAdapterConfig",
|
||||
"CacheDiTBackend",
|
||||
"CacheDiTConfig",
|
||||
"CacheDiTRequestSpec",
|
||||
"RequestScopedCacheDiTRuntime",
|
||||
"SensenovaCachedAdapter",
|
||||
"cache_summary",
|
||||
"enable_cache_for_dit",
|
||||
|
||||
@@ -1,101 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion.cache.cachedit import runtime as runtime_module
|
||||
from vllm_omni.diffusion.cache.cachedit.runtime import (
|
||||
CacheDiTRequestSpec,
|
||||
RequestScopedCacheDiTRuntime,
|
||||
)
|
||||
from vllm_omni.diffusion.data import DiffusionCacheConfig
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
def _spec(key: str, steps: int = 50) -> CacheDiTRequestSpec:
|
||||
return CacheDiTRequestSpec(
|
||||
installation_key=key,
|
||||
cache_config=DiffusionCacheConfig(),
|
||||
num_inference_steps=steps,
|
||||
)
|
||||
|
||||
|
||||
def _recording_backend(events):
|
||||
class FakeBackend:
|
||||
def __init__(self, config):
|
||||
events.append(("create", config))
|
||||
self.enabled = False
|
||||
|
||||
def enable(self, pipeline):
|
||||
events.append(("enable", pipeline))
|
||||
self.enabled = True
|
||||
|
||||
def refresh(self, pipeline, num_inference_steps):
|
||||
events.append(("refresh", pipeline, num_inference_steps))
|
||||
|
||||
def disable(self, pipeline):
|
||||
events.append(("disable", pipeline))
|
||||
self.enabled = False
|
||||
|
||||
def is_enabled(self):
|
||||
return self.enabled
|
||||
|
||||
return FakeBackend
|
||||
|
||||
|
||||
def test_startup_cache_follows_omit_lossless_high_omit_state_machine(monkeypatch):
|
||||
events = []
|
||||
FakeBackend = _recording_backend(events)
|
||||
monkeypatch.setattr(runtime_module, "CacheDiTBackend", FakeBackend)
|
||||
pipeline = object()
|
||||
config = DiffusionCacheConfig()
|
||||
runtime = RequestScopedCacheDiTRuntime(pipeline)
|
||||
startup_backend = FakeBackend(config)
|
||||
startup_backend.enabled = True
|
||||
|
||||
runtime.adopt(startup_backend, installation_key="generic")
|
||||
runtime.prepare(_spec("generic", 50)) # omitted: keep startup cache
|
||||
runtime.prepare(None) # lossless: uninstall
|
||||
runtime.prepare(None) # repeated lossless: no-op
|
||||
runtime.prepare(_spec("high", 40)) # high: install H3 profile
|
||||
runtime.prepare(_spec("high", 30)) # repeated high: refresh only
|
||||
runtime.prepare(_spec("generic", 20)) # omitted: restore startup profile
|
||||
|
||||
assert [event[0] for event in events] == [
|
||||
"create",
|
||||
"refresh",
|
||||
"disable",
|
||||
"create",
|
||||
"enable",
|
||||
"refresh",
|
||||
"refresh",
|
||||
"disable",
|
||||
"create",
|
||||
"enable",
|
||||
"refresh",
|
||||
]
|
||||
assert events[1][-1] == 50
|
||||
assert events[5][-1] == 40
|
||||
assert events[6][-1] == 30
|
||||
assert events[10][-1] == 20
|
||||
assert runtime.is_enabled
|
||||
|
||||
|
||||
def test_high_installs_without_startup_cache_and_lossless_uninstalls(monkeypatch):
|
||||
events = []
|
||||
monkeypatch.setattr(runtime_module, "CacheDiTBackend", _recording_backend(events))
|
||||
pipeline = object()
|
||||
runtime = RequestScopedCacheDiTRuntime(pipeline)
|
||||
|
||||
runtime.prepare(_spec("high", 40))
|
||||
runtime.prepare(None)
|
||||
|
||||
assert [event[0] for event in events] == [
|
||||
"create",
|
||||
"enable",
|
||||
"refresh",
|
||||
"disable",
|
||||
]
|
||||
assert events[2][-1] == 40
|
||||
assert not runtime.is_enabled
|
||||
@@ -13,70 +13,6 @@ import torch
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
pipeline.partition = "fl2va"
|
||||
pipeline.supported_tasks = frozenset({"t2va"})
|
||||
pipeline.default_video_shift = 12.0
|
||||
pipeline.default_audio_shift = 3.0
|
||||
pipeline.device = torch.device("cpu")
|
||||
pipeline.od_config = SimpleNamespace()
|
||||
cache_spec = object()
|
||||
quality_plan = SimpleNamespace(cache_dit=cache_spec)
|
||||
pipeline._quality_policy = Mock()
|
||||
pipeline._quality_policy.resolve.return_value = quality_plan
|
||||
events = []
|
||||
pipeline._cache_dit_runtime = SimpleNamespace(prepare=lambda spec: events.append(("prepare", spec)))
|
||||
pipeline.encode_prompt = Mock(
|
||||
return_value=(
|
||||
torch.ones(1, 2),
|
||||
torch.ones(1, dtype=torch.long),
|
||||
)
|
||||
)
|
||||
|
||||
def diffuse(**kwargs):
|
||||
events.append(("diffuse", kwargs))
|
||||
return torch.zeros(1), torch.zeros(1)
|
||||
|
||||
pipeline.diffuse = diffuse
|
||||
pipeline.decode = Mock(return_value=(torch.zeros(1), torch.zeros(1)))
|
||||
sampling = OmniDiffusionSamplingParams(
|
||||
quality="high",
|
||||
width=1344,
|
||||
height=768,
|
||||
fps=24,
|
||||
num_frames=124,
|
||||
num_inference_steps=50,
|
||||
extra_args={"task": "t2va", "aspect_ratio": "16:9"},
|
||||
)
|
||||
batch = DiffusionRequestBatch(
|
||||
[
|
||||
OmniDiffusionRequest(
|
||||
prompt="quality boundary",
|
||||
sampling_params=sampling,
|
||||
request_id="quality-boundary",
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
output = pipeline.forward(batch)
|
||||
|
||||
assert events[0] == ("prepare", cache_spec)
|
||||
assert events[1][0] == "diffuse"
|
||||
pipeline._quality_policy.resolve.assert_called_once_with(
|
||||
quality="high",
|
||||
num_inference_steps=50,
|
||||
extra_args={"task": "t2va", "aspect_ratio": "16:9"},
|
||||
)
|
||||
assert output.output == pipeline.decode.return_value
|
||||
|
||||
|
||||
def test_pipeline_import_registry_and_component_discovery():
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
from vllm_omni.diffusion.registry import (
|
||||
|
||||
@@ -1,212 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionCacheConfig
|
||||
from vllm_omni.errors import OmniClientError
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
def _quality_od_config(*, cache_backend: str = "cache_dit", **overrides):
|
||||
values = {
|
||||
"cache_backend": cache_backend,
|
||||
"cache_config": DiffusionCacheConfig(),
|
||||
"diffusion_attention_config": SimpleNamespace(default=None, per_role={}),
|
||||
"enable_cpu_offload": False,
|
||||
"enable_layerwise_offload": False,
|
||||
"enable_distributed_layerwise_offload": False,
|
||||
"enforce_eager": True,
|
||||
"num_gpus": 4,
|
||||
"quantization_config": None,
|
||||
"parallel_config": SimpleNamespace(
|
||||
allgather_degree=1,
|
||||
cfg_parallel_size=1,
|
||||
data_parallel_size=1,
|
||||
pipeline_parallel_size=1,
|
||||
ring_degree=1,
|
||||
sequence_parallel_size=4,
|
||||
tensor_parallel_size=1,
|
||||
ulysses_degree=4,
|
||||
use_hsdp=False,
|
||||
),
|
||||
}
|
||||
values.update(overrides)
|
||||
return SimpleNamespace(**values)
|
||||
|
||||
|
||||
def _resolve_quality(policy, **overrides):
|
||||
values = {
|
||||
"quality": "high",
|
||||
"num_inference_steps": 50,
|
||||
"extra_args": None,
|
||||
}
|
||||
values.update(overrides)
|
||||
return policy.resolve(**values)
|
||||
|
||||
|
||||
def test_high_quality_policy_emits_cache_profile_without_deployment_gates():
|
||||
from vllm_omni.diffusion.models.minimax_h3.quality_policy import (
|
||||
MiniMaxH3QualityPolicy,
|
||||
)
|
||||
|
||||
policy = MiniMaxH3QualityPolicy(
|
||||
_quality_od_config(
|
||||
cache_backend="none",
|
||||
enable_layerwise_offload=True,
|
||||
enforce_eager=False,
|
||||
num_gpus=1,
|
||||
)
|
||||
)
|
||||
|
||||
plan = _resolve_quality(policy, num_inference_steps=17)
|
||||
spec = plan.cache_dit
|
||||
|
||||
assert spec is not None
|
||||
assert spec.installation_key == "minimax_h3.high"
|
||||
assert spec.num_inference_steps == 17
|
||||
assert spec.cache_config.Fn_compute_blocks == 1
|
||||
assert spec.cache_config.Bn_compute_blocks == 0
|
||||
assert spec.cache_config.max_warmup_steps == 4
|
||||
assert spec.cache_config.residual_diff_threshold == pytest.approx(0.04)
|
||||
assert spec.cache_config.max_continuous_cached_steps == 1
|
||||
assert spec.cache_config.enable_taylorseer is False
|
||||
assert spec.cache_config.scm_steps_mask_policy is None
|
||||
|
||||
|
||||
def test_h3_force_refresh_is_model_owned_and_request_scoped():
|
||||
from vllm_omni.diffusion.models.minimax_h3.quality_policy import MiniMaxH3QualityPolicy
|
||||
|
||||
policy = MiniMaxH3QualityPolicy(_quality_od_config())
|
||||
plan = _resolve_quality(
|
||||
policy,
|
||||
quality="high",
|
||||
num_inference_steps=17,
|
||||
extra_args={"force_refresh_step_hint": 7},
|
||||
)
|
||||
|
||||
assert plan.cache_dit is not None
|
||||
assert plan.cache_dit.installation_key == "minimax_h3.high:force_refresh=7:once"
|
||||
assert plan.cache_dit.cache_config.force_refresh_step_hint == 7
|
||||
assert plan.cache_dit.cache_config.force_refresh_step_policy == "once"
|
||||
|
||||
|
||||
def test_h3_force_refresh_policy_supports_repeat_and_does_not_mutate_generic_config():
|
||||
from vllm_omni.diffusion.models.minimax_h3.quality_policy import MiniMaxH3QualityPolicy
|
||||
|
||||
config = _quality_od_config()
|
||||
plan = _resolve_quality(
|
||||
MiniMaxH3QualityPolicy(config),
|
||||
quality=None,
|
||||
num_inference_steps=20,
|
||||
extra_args={"force_refresh_step_hint": 5, "force_refresh_step_policy": "repeat"},
|
||||
)
|
||||
|
||||
assert plan.cache_dit is not None
|
||||
assert plan.cache_dit.installation_key == "minimax_h3.generic:force_refresh=5:repeat"
|
||||
assert plan.cache_dit.cache_config is not config.cache_config
|
||||
assert plan.cache_dit.cache_config.force_refresh_step_hint == 5
|
||||
assert plan.cache_dit.cache_config.force_refresh_step_policy == "repeat"
|
||||
assert config.cache_config.force_refresh_step_hint is None
|
||||
assert config.cache_config.force_refresh_step_policy == "once"
|
||||
|
||||
restored = _resolve_quality(MiniMaxH3QualityPolicy(config), quality=None)
|
||||
assert restored.cache_dit is not None
|
||||
assert restored.cache_dit.installation_key == "minimax_h3.generic"
|
||||
assert restored.cache_dit.cache_config is config.cache_config
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("extra_args", "message"),
|
||||
[
|
||||
({"force_refresh_step_hint": True}, "positive integer"),
|
||||
({"force_refresh_step_hint": 0}, "between 1"),
|
||||
({"force_refresh_step_hint": 51}, "between 1"),
|
||||
({"force_refresh_step_hint": 5, "force_refresh_step_policy": "always"}, "one of"),
|
||||
({"force_refresh_step_policy": "repeat"}, "requires force_refresh_step_hint"),
|
||||
],
|
||||
)
|
||||
def test_h3_force_refresh_rejects_invalid_values(extra_args, message):
|
||||
from vllm_omni.diffusion.models.minimax_h3.quality_policy import MiniMaxH3QualityPolicy
|
||||
|
||||
with pytest.raises(OmniClientError, match=message):
|
||||
_resolve_quality(
|
||||
MiniMaxH3QualityPolicy(_quality_od_config()),
|
||||
extra_args=extra_args,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("cache_backend", "quality"),
|
||||
[
|
||||
("cache_dit", "lossless"),
|
||||
("none", None),
|
||||
],
|
||||
)
|
||||
def test_h3_force_refresh_requires_active_cache_target(cache_backend, quality):
|
||||
from vllm_omni.diffusion.models.minimax_h3.quality_policy import MiniMaxH3QualityPolicy
|
||||
|
||||
policy = MiniMaxH3QualityPolicy(_quality_od_config(cache_backend=cache_backend))
|
||||
|
||||
with pytest.raises(OmniClientError, match="require an active Cache-DiT"):
|
||||
_resolve_quality(
|
||||
policy,
|
||||
quality=quality,
|
||||
extra_args={"force_refresh_step_hint": 7},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("cache_backend", "quality", "expected_key"),
|
||||
[
|
||||
("cache_dit", None, "minimax_h3.generic"),
|
||||
("cache_dit", "lossless", None),
|
||||
("cache_dit", "high", "minimax_h3.high"),
|
||||
("none", None, None),
|
||||
("none", "lossless", None),
|
||||
("none", "high", "minimax_h3.high"),
|
||||
],
|
||||
)
|
||||
def test_model_policy_owns_request_cache_target(cache_backend, quality, expected_key):
|
||||
from vllm_omni.diffusion.models.minimax_h3.quality_policy import (
|
||||
MiniMaxH3QualityPolicy,
|
||||
)
|
||||
|
||||
config = _quality_od_config(cache_backend=cache_backend)
|
||||
policy = MiniMaxH3QualityPolicy(config)
|
||||
|
||||
plan = _resolve_quality(
|
||||
policy,
|
||||
quality=quality,
|
||||
)
|
||||
|
||||
if expected_key is None:
|
||||
assert plan.cache_dit is None
|
||||
else:
|
||||
assert plan.cache_dit is not None
|
||||
assert plan.cache_dit.installation_key == expected_key
|
||||
if quality is None:
|
||||
assert plan.cache_dit.cache_config is config.cache_config
|
||||
|
||||
|
||||
def test_h3_adopts_runner_installed_cache_dit_backend():
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
backend = object()
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
pipeline._cache_dit_runtime = SimpleNamespace(
|
||||
adopt=lambda adopted, *, installation_key: setattr(
|
||||
pipeline,
|
||||
"adopted",
|
||||
(adopted, installation_key),
|
||||
),
|
||||
is_enabled=True,
|
||||
)
|
||||
|
||||
pipeline.adopt_cache_dit_backend(backend)
|
||||
|
||||
assert pipeline.adopted == (backend, "minimax_h3.generic")
|
||||
assert pipeline.is_cache_dit_enabled()
|
||||
@@ -178,9 +178,10 @@ class TestAsyncOutputLoopLogic:
|
||||
t = threading.Thread(target=proc._async_output_loop, daemon=True)
|
||||
t.start()
|
||||
|
||||
# Give the loop time to process, then send the shutdown sentinel.
|
||||
# Give the loop time to process, then unblock and stop.
|
||||
time.sleep(0.5)
|
||||
proc._async_output_queue.put(None)
|
||||
proc._running = False
|
||||
proc._async_output_queue.put((DiffusionOutput(), "sentinel", None))
|
||||
t.join(timeout=2.0)
|
||||
|
||||
mock_pack.assert_called()
|
||||
@@ -190,43 +191,3 @@ class TestAsyncOutputLoopLogic:
|
||||
for c in proc.result_mq.enqueue.call_args_list
|
||||
)
|
||||
assert output_ready_enqueued, "Expected OUTPUT_READY with async_output_id='abc123' in enqueue calls"
|
||||
|
||||
|
||||
class TestWorkerProcShutdown:
|
||||
def test_shutdown_stops_async_thread_and_releases_queues(self):
|
||||
import threading
|
||||
|
||||
proc = _make_worker_proc(step_execution=False)
|
||||
proc.worker = MagicMock()
|
||||
proc.mq = MagicMock()
|
||||
broadcast_mq = proc.mq
|
||||
result_mq = proc.result_mq
|
||||
|
||||
thread = threading.Thread(target=proc._async_output_queue.get)
|
||||
proc._async_output_thread = thread
|
||||
thread.start()
|
||||
|
||||
proc.shutdown()
|
||||
|
||||
assert not thread.is_alive()
|
||||
proc.worker.shutdown.assert_called_once_with()
|
||||
broadcast_mq.shutdown.assert_called_once_with()
|
||||
result_mq.shutdown.assert_called_once_with()
|
||||
assert proc.mq is None
|
||||
assert proc.result_mq is None
|
||||
|
||||
def test_shutdown_releases_queues_when_worker_shutdown_fails(self):
|
||||
proc = _make_worker_proc(step_execution=True)
|
||||
proc.worker = MagicMock()
|
||||
proc.worker.shutdown.side_effect = RuntimeError("shutdown failed")
|
||||
proc.mq = MagicMock()
|
||||
broadcast_mq = proc.mq
|
||||
result_mq = proc.result_mq
|
||||
|
||||
with pytest.raises(RuntimeError, match="shutdown failed"):
|
||||
proc.shutdown()
|
||||
|
||||
broadcast_mq.shutdown.assert_called_once_with()
|
||||
result_mq.shutdown.assert_called_once_with()
|
||||
assert proc.mq is None
|
||||
assert proc.result_mq is None
|
||||
@@ -15,10 +15,7 @@ from vllm_omni.diffusion.data import (
|
||||
DiffusionParallelConfig,
|
||||
OmniDiffusionConfig,
|
||||
)
|
||||
from vllm_omni.diffusion.model_metadata import (
|
||||
HUNYUAN_IMAGE3_MAX_INPUT_IMAGES,
|
||||
QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES,
|
||||
)
|
||||
from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
@@ -139,30 +136,6 @@ def test_task_type_roundtrip():
|
||||
assert od.task_type == "model-defined-task"
|
||||
|
||||
|
||||
def test_architecture_name_resolves_via_pipeline_class_fallback():
|
||||
# Direct hit: the architecture name itself is the metadata table key.
|
||||
qwen_od_config = OmniDiffusionConfig(
|
||||
model="Qwen/Qwen-Image-Edit-2511",
|
||||
model_class_name="QwenImageEditPlusPipeline",
|
||||
)
|
||||
qwen_od_config.update_multimodal_support()
|
||||
|
||||
assert qwen_od_config.supports_multimodal_inputs is True
|
||||
assert qwen_od_config.max_multimodal_image_inputs == QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
|
||||
|
||||
# Fallback: the HF architecture name differs from the internal pipeline class
|
||||
# name, so ``get_diffusion_model_metadata`` must consult ``_DIFFUSION_MODELS``
|
||||
# and resolve the metadata via the pipeline class name.
|
||||
hunyuan_od_config = OmniDiffusionConfig(
|
||||
model="tencent/HunyuanImage-3.0-Instruct",
|
||||
model_class_name="HunyuanImage3ForCausalMM",
|
||||
)
|
||||
hunyuan_od_config.update_multimodal_support()
|
||||
|
||||
assert hunyuan_od_config.supports_multimodal_inputs is True
|
||||
assert hunyuan_od_config.max_multimodal_image_inputs == HUNYUAN_IMAGE3_MAX_INPUT_IMAGES
|
||||
|
||||
|
||||
def test_additional_config_roundtrip():
|
||||
additional_config = {"torchair_graph_config": {"enabled": True}}
|
||||
od = _roundtrip_diffusion_config(model="x", additional_config=additional_config)
|
||||
|
||||
@@ -331,55 +331,6 @@ class TestRequestBatchCapability:
|
||||
DiffusionEngine(od_config)
|
||||
fake_executor_cls.assert_not_called()
|
||||
|
||||
def test_engine_allows_independent_dlo_dp_requests_for_single_request_pipeline(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
mocker: MockerFixture,
|
||||
) -> None:
|
||||
od_config = SimpleNamespace(
|
||||
model_class_name="SinglePipeline",
|
||||
custom_pipeline_args=None,
|
||||
streaming_output=False,
|
||||
max_num_seqs=2,
|
||||
request_batch_max_wait_ms=0,
|
||||
parallel_config=SimpleNamespace(data_parallel_size=2),
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=True,
|
||||
)
|
||||
fake_executor = SimpleNamespace(
|
||||
execute_request=mocker.Mock(return_value="per-request"),
|
||||
execute_batch=mocker.Mock(return_value="batch"),
|
||||
execute_step=mocker.Mock(return_value="step"),
|
||||
)
|
||||
fake_executor_cls = mocker.Mock(return_value=fake_executor)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
|
||||
lambda *args, **kwargs: None,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
|
||||
lambda *args, **kwargs: None,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
|
||||
lambda *args, **kwargs: fake_executor_cls,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
diffusion_engine_module.DiffusionModelRegistry,
|
||||
"_try_load_model_cls",
|
||||
lambda model_class_name: _SingleRequestPipeline,
|
||||
)
|
||||
|
||||
engine = DiffusionEngine(od_config)
|
||||
|
||||
assert engine.execution_mode == DiffusionExecutionMode.REQUEST_BATCH
|
||||
assert engine.supports_request_batch is False
|
||||
assert engine.dp_concurrent is True
|
||||
assert engine.scheduler.max_num_running_reqs == 2
|
||||
assert od_config.request_batch_max_wait_ms == 0
|
||||
fake_executor_cls.assert_called_once_with(od_config)
|
||||
|
||||
@pytest.mark.parametrize("request_ids", [("req-a",), ("req-a", "req-b")])
|
||||
def test_engine_enables_batch_dispatch_for_request_batch_pipeline(
|
||||
self,
|
||||
|
||||
@@ -2,19 +2,11 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import contextlib
|
||||
import multiprocessing as mp
|
||||
import queue
|
||||
import threading
|
||||
import time
|
||||
from multiprocessing import shared_memory
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput
|
||||
from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
|
||||
from vllm_omni.diffusion.ipc import (
|
||||
_SHM_TENSOR_THRESHOLD,
|
||||
DIFFUSION_RPC_RESULT_ENVELOPE,
|
||||
@@ -38,129 +30,6 @@ def _cleanup_shm_handle(value: object) -> None:
|
||||
_unpack_if_shm_handle(value)
|
||||
|
||||
|
||||
def _result_queue_worker(connection, rank: int) -> None:
|
||||
"""Send one nested NumPy result through a worker-owned MessageQueue."""
|
||||
result_mq = MessageQueue(n_reader=1, n_local_reader=1, local_reader_ranks=[0])
|
||||
try:
|
||||
connection.send(result_mq.export_handle())
|
||||
assert connection.recv() == "reader-ready"
|
||||
|
||||
frames = np.full(300_000, rank, dtype=np.float32)
|
||||
audio = np.arange(300_000, dtype=np.float32) + rank
|
||||
output = DiffusionOutput(
|
||||
output={
|
||||
"rank": rank,
|
||||
"nested": {"frames": frames},
|
||||
"audio": [audio],
|
||||
}
|
||||
)
|
||||
pack_diffusion_output_shm(output)
|
||||
connection.send(
|
||||
[
|
||||
output.output["nested"]["frames"]["name"],
|
||||
output.output["audio"][0]["name"],
|
||||
]
|
||||
)
|
||||
result_mq.enqueue(output)
|
||||
assert connection.recv() == "shutdown"
|
||||
finally:
|
||||
result_mq.shutdown()
|
||||
connection.close()
|
||||
|
||||
|
||||
def test_per_worker_result_queues_release_nested_numpy_shm_and_processes() -> None:
|
||||
"""Exercise the production per-worker queue/pump/SHM lifecycle."""
|
||||
ctx = mp.get_context("spawn")
|
||||
parent_connections = []
|
||||
processes = []
|
||||
result_mqs = []
|
||||
shm_names = []
|
||||
executor = None
|
||||
|
||||
try:
|
||||
for rank in range(2):
|
||||
parent_connection, child_connection = ctx.Pipe()
|
||||
process = ctx.Process(
|
||||
target=_result_queue_worker,
|
||||
args=(child_connection, rank),
|
||||
name=f"IPCResultWorker-{rank}",
|
||||
)
|
||||
process.start()
|
||||
child_connection.close()
|
||||
parent_connections.append(parent_connection)
|
||||
processes.append(process)
|
||||
|
||||
for connection in parent_connections:
|
||||
handle = connection.recv()
|
||||
result_mqs.append(MessageQueue.create_from_handle(handle, 0))
|
||||
connection.send("reader-ready")
|
||||
shm_names.extend(connection.recv())
|
||||
|
||||
executor = object.__new__(MultiprocDiffusionExecutor)
|
||||
executor._result_mqs = result_mqs
|
||||
executor._result_mq = result_mqs[0]
|
||||
executor._pump_running = False
|
||||
executor._pump_stop = threading.Event()
|
||||
executor._sync_result_buffer = queue.Queue()
|
||||
executor._is_failed = False
|
||||
executor._futures_lock = threading.RLock()
|
||||
executor._rpc_futures = {}
|
||||
executor._output_futures = {}
|
||||
executor._completed_outputs = {}
|
||||
executor._batch_split_map = {}
|
||||
executor._start_result_pump()
|
||||
|
||||
deadline = time.monotonic() + 10.0
|
||||
while executor._sync_result_buffer.qsize() < 2 and time.monotonic() < deadline:
|
||||
time.sleep(0.01)
|
||||
assert executor._sync_result_buffer.qsize() == 2
|
||||
|
||||
outputs = []
|
||||
for _ in range(2):
|
||||
output = executor._sync_result_buffer.get_nowait()
|
||||
unpack_diffusion_output_shm(output)
|
||||
outputs.append(output)
|
||||
|
||||
assert sorted(output.output["rank"] for output in outputs) == [0, 1]
|
||||
for output in outputs:
|
||||
rank = output.output["rank"]
|
||||
assert output.output["nested"]["frames"][0] == rank
|
||||
assert output.output["audio"][0][0] == rank
|
||||
|
||||
executor._pump_stop.set()
|
||||
for thread in executor._result_pump_threads:
|
||||
thread.join(timeout=3.0)
|
||||
assert all(not thread.is_alive() for thread in executor._result_pump_threads)
|
||||
finally:
|
||||
if executor is not None:
|
||||
executor._pump_stop.set()
|
||||
for thread in getattr(executor, "_result_pump_threads", []):
|
||||
thread.join(timeout=3.0)
|
||||
|
||||
for connection in parent_connections:
|
||||
with contextlib.suppress(BrokenPipeError, EOFError, OSError):
|
||||
connection.send("shutdown")
|
||||
for process in processes:
|
||||
process.join(timeout=10.0)
|
||||
if process.is_alive():
|
||||
process.terminate()
|
||||
process.join(timeout=5.0)
|
||||
for result_mq in result_mqs:
|
||||
with contextlib.suppress(Exception):
|
||||
result_mq.shutdown()
|
||||
for connection in parent_connections:
|
||||
connection.close()
|
||||
|
||||
assert all(process.exitcode == 0 for process in processes)
|
||||
for name in shm_names:
|
||||
try:
|
||||
leaked = shared_memory.SharedMemory(name=name)
|
||||
except FileNotFoundError:
|
||||
continue
|
||||
leaked.close()
|
||||
pytest.fail(f"shared memory segment {name} still exists after unpack")
|
||||
|
||||
|
||||
def test_diffusion_output_dict_tensors_round_trip_through_shm() -> None:
|
||||
image = torch.arange(300_000, dtype=torch.float32)
|
||||
video = torch.arange(300_000, dtype=torch.float32) * 2
|
||||
@@ -214,29 +83,6 @@ def test_diffusion_output_list_tensors_round_trip_through_shm() -> None:
|
||||
torch.testing.assert_close(output.output[1], frames[1])
|
||||
|
||||
|
||||
def test_diffusion_output_numpy_array_round_trips_through_shm() -> None:
|
||||
frames = np.arange(300_000, dtype=np.float32)
|
||||
output = DiffusionOutput(output=frames)
|
||||
|
||||
pack_diffusion_output_shm(output)
|
||||
|
||||
assert output.output["__ndarray_shm__"] is True
|
||||
|
||||
unpack_diffusion_output_shm(output)
|
||||
|
||||
assert isinstance(output.output, np.ndarray)
|
||||
np.testing.assert_array_equal(output.output, frames)
|
||||
|
||||
|
||||
def test_numpy_object_array_stays_inline() -> None:
|
||||
values = np.empty(200_000, dtype=object)
|
||||
values[:] = "safe-to-pickle"
|
||||
|
||||
packed = _pack_value_if_large(values)
|
||||
|
||||
assert packed is values
|
||||
|
||||
|
||||
def test_rpc_result_envelope_diffusion_output_round_trips_through_shm() -> None:
|
||||
tensor = torch.arange(300_000, dtype=torch.float32)
|
||||
envelope = {
|
||||
@@ -262,25 +108,6 @@ def test_rpc_result_envelope_diffusion_output_round_trips_through_shm() -> None:
|
||||
assert unpacked["rank_statuses"] == [{"rank": 0, "ok": True}]
|
||||
|
||||
|
||||
def test_rpc_result_envelope_dp_tagged_output_round_trips_through_shm() -> None:
|
||||
frames = np.arange(300_000, dtype=np.float32)
|
||||
envelope = {
|
||||
"type": DIFFUSION_RPC_RESULT_ENVELOPE,
|
||||
"result": {"dp_rank": 1, "output": DiffusionOutput(output=frames)},
|
||||
"rank_statuses": [{"rank": 1, "ok": True}],
|
||||
}
|
||||
|
||||
packed = pack_diffusion_output_shm(envelope)
|
||||
tagged = packed["result"]
|
||||
assert tagged["dp_rank"] == 1
|
||||
assert tagged["output"].output["__ndarray_shm__"] is True
|
||||
|
||||
unpacked = unpack_diffusion_output_shm(packed)
|
||||
tagged = unpacked["result"]
|
||||
assert tagged["dp_rank"] == 1
|
||||
np.testing.assert_array_equal(tagged["output"].output, frames)
|
||||
|
||||
|
||||
def test_batch_runner_output_round_trips_nested_results_through_shm() -> None:
|
||||
first = torch.arange(_large_numel(torch.float32), dtype=torch.float32)
|
||||
second = torch.arange(_large_numel(torch.float32), dtype=torch.float32) + 1
|
||||
|
||||
@@ -162,21 +162,6 @@ def _fake_platform_for_peak_memory():
|
||||
)
|
||||
|
||||
|
||||
def test_request_scoped_cache_dit_lifecycle_is_pipeline_opt_in():
|
||||
events = []
|
||||
plain_pipeline = object()
|
||||
request_scoped_pipeline = SimpleNamespace(
|
||||
adopt_cache_dit_backend=lambda backend: events.append(backend),
|
||||
is_cache_dit_enabled=lambda: True,
|
||||
)
|
||||
backend = object()
|
||||
|
||||
assert not model_runner_module.adopt_request_scoped_cache_dit(plain_pipeline, backend)
|
||||
assert model_runner_module.adopt_request_scoped_cache_dit(request_scoped_pipeline, backend)
|
||||
assert model_runner_module.is_request_scoped_cache_dit_enabled(request_scoped_pipeline)
|
||||
assert events == [backend]
|
||||
|
||||
|
||||
def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summary: bool = True):
|
||||
runner = object.__new__(DiffusionModelRunner)
|
||||
runner.vllm_config = object()
|
||||
@@ -434,33 +419,6 @@ def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeyp
|
||||
assert cache_backend.refresh_calls == [(runner.pipeline, 4, True)]
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.parametrize("cache_dit_enabled", [False, True])
|
||||
def test_execute_model_cache_summary_follows_pipeline_owned_cache_dit_state(monkeypatch, cache_dit_enabled):
|
||||
runner = _make_runner(cache_backend=None, cache_backend_name="cache_dit")
|
||||
runner.pipeline.adopt_cache_dit_backend = lambda backend: None
|
||||
runner.pipeline.is_cache_dit_enabled = lambda: cache_dit_enabled
|
||||
req = _make_request()
|
||||
|
||||
cache_summary_calls = []
|
||||
|
||||
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
||||
monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
|
||||
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
|
||||
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
|
||||
monkeypatch.setattr(
|
||||
model_runner_module,
|
||||
"cache_summary",
|
||||
lambda pipeline, details: cache_summary_calls.append((pipeline, details)),
|
||||
)
|
||||
|
||||
output = DiffusionModelRunner.execute_model(runner, req)
|
||||
|
||||
assert output.output == "ok"
|
||||
assert cache_summary_calls == ([(runner.pipeline, True)] if cache_dit_enabled else [])
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_execute_model_passes_single_request_batch_to_non_admission_pipeline(monkeypatch):
|
||||
|
||||
@@ -380,63 +380,6 @@ class TestRequestScheduler:
|
||||
assert first.num_running_reqs == 1
|
||||
assert first.num_waiting_reqs == 1
|
||||
|
||||
def test_batches_different_quality_levels_separately(self) -> None:
|
||||
scheduler = RequestScheduler()
|
||||
scheduler.initialize(SimpleNamespace(max_num_seqs=2))
|
||||
|
||||
high = scheduler.add_request(
|
||||
_make_step_request(
|
||||
"high",
|
||||
sampling_params=OmniDiffusionSamplingParams(
|
||||
num_inference_steps=2,
|
||||
quality="high",
|
||||
),
|
||||
)
|
||||
)
|
||||
scheduler.add_request(
|
||||
_make_step_request(
|
||||
"lossless",
|
||||
sampling_params=OmniDiffusionSamplingParams(
|
||||
num_inference_steps=2,
|
||||
quality="lossless",
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
first = scheduler.schedule()
|
||||
|
||||
assert _new_ids(first) == [high]
|
||||
assert first.num_running_reqs == 1
|
||||
assert first.num_waiting_reqs == 1
|
||||
|
||||
def test_batches_omitted_and_explicit_lossless_separately(self) -> None:
|
||||
scheduler = RequestScheduler()
|
||||
scheduler.initialize(SimpleNamespace(max_num_seqs=2))
|
||||
|
||||
omitted = scheduler.add_request(
|
||||
_make_step_request(
|
||||
"omitted",
|
||||
sampling_params=OmniDiffusionSamplingParams(
|
||||
num_inference_steps=2,
|
||||
),
|
||||
)
|
||||
)
|
||||
scheduler.add_request(
|
||||
_make_step_request(
|
||||
"explicit",
|
||||
sampling_params=OmniDiffusionSamplingParams(
|
||||
num_inference_steps=2,
|
||||
quality="lossless",
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
first = scheduler.schedule()
|
||||
|
||||
assert _new_ids(first) == [omitted]
|
||||
assert first.num_running_reqs == 1
|
||||
assert first.num_waiting_reqs == 1
|
||||
|
||||
def test_batches_different_request_local_seed_together(self) -> None:
|
||||
scheduler = RequestScheduler()
|
||||
scheduler.initialize(SimpleNamespace(max_num_seqs=2))
|
||||
@@ -987,35 +930,6 @@ class TestStepScheduler:
|
||||
assert second.num_running_reqs == 1
|
||||
assert second.num_waiting_reqs == 1
|
||||
|
||||
def test_step_batch_rejects_different_quality_levels(self) -> None:
|
||||
scheduler = StepScheduler()
|
||||
scheduler.initialize(SimpleNamespace(max_num_seqs=2))
|
||||
|
||||
high = scheduler.add_request(
|
||||
_make_step_request(
|
||||
"high",
|
||||
sampling_params=OmniDiffusionSamplingParams(
|
||||
num_inference_steps=4,
|
||||
quality="high",
|
||||
),
|
||||
)
|
||||
)
|
||||
scheduler.add_request(
|
||||
_make_step_request(
|
||||
"lossless",
|
||||
sampling_params=OmniDiffusionSamplingParams(
|
||||
num_inference_steps=4,
|
||||
quality="lossless",
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
first = scheduler.schedule()
|
||||
|
||||
assert _new_ids(first) == [high]
|
||||
assert first.num_running_reqs == 1
|
||||
assert first.num_waiting_reqs == 1
|
||||
|
||||
def test_step_batch_co_schedules_requests_sharing_lora(self) -> None:
|
||||
"""Multiple requests with the same LoRA (id + scale) co-batch."""
|
||||
from vllm_omni.lora.request import LoRARequest
|
||||
|
||||
@@ -75,32 +75,6 @@ async def test_inline_dispatch_request_success(client, mock_engine):
|
||||
assert output.request_id == "req-1"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inline_requests_clone_shared_sampling_params(client, mock_engine):
|
||||
requests = []
|
||||
|
||||
async def _step_streaming(request):
|
||||
requests.append(request)
|
||||
yield [OmniRequestOutput.from_diffusion(request_id=request.request_id, images=[MagicMock()])]
|
||||
|
||||
mock_engine.step_streaming = _step_streaming
|
||||
|
||||
shared = OmniDiffusionSamplingParams()
|
||||
await client.add_request_async("req-1", "First prompt", shared)
|
||||
await client.add_request_async("req-2", "Second prompt", shared)
|
||||
|
||||
for _ in range(10):
|
||||
if len(requests) == 2:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
assert len(requests) == 2
|
||||
assert requests[0].sampling_params is not requests[1].sampling_params
|
||||
assert requests[0].sampling_params.guidance_scale_provided is False
|
||||
assert requests[1].sampling_params.guidance_scale_provided is False
|
||||
assert shared.guidance_scale is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inline_dispatch_request_streaming_success(client, mock_engine):
|
||||
"""Inline StageDiffusionClient correctly receives streaming chunk output from Diffusion Engine"""
|
||||
@@ -163,20 +137,6 @@ def test_inline_shutdown(client, mock_engine):
|
||||
mock_engine.close.assert_called_once()
|
||||
|
||||
|
||||
def test_inline_shutdown_runs_after_orchestrator_premark(client, mock_engine):
|
||||
# The orchestrator pre-marks clients to suppress false worker-death errors
|
||||
# before StagePool invokes the actual teardown.
|
||||
client._shutting_down = True
|
||||
|
||||
client.shutdown()
|
||||
|
||||
mock_engine.close.assert_called_once()
|
||||
assert client._shutdown_complete
|
||||
|
||||
client.shutdown()
|
||||
mock_engine.close.assert_called_once()
|
||||
|
||||
|
||||
def test_inline_registers_executor_failure_callback(client, mock_engine):
|
||||
mock_engine.executor.register_failure_callback.assert_called_once()
|
||||
|
||||
|
||||
@@ -137,18 +137,6 @@ class TestPackDiffusionOutputShmWithD2hStream:
|
||||
d2h_stream=d2h,
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("wrapped_in_rpc_envelope", [False, True])
|
||||
def test_dp_tagged_output_path_with_stream(self, mocker, wrapped_in_rpc_envelope):
|
||||
mock_pack_fields = mocker.patch(
|
||||
"vllm_omni.diffusion.ipc._pack_diffusion_fields",
|
||||
)
|
||||
d2h = mocker.MagicMock()
|
||||
result = DiffusionOutput()
|
||||
tagged = {"dp_rank": 1, "output": result}
|
||||
value = {"type": "diffusion_rpc_result", "result": tagged} if wrapped_in_rpc_envelope else tagged
|
||||
pack_diffusion_output_shm(value, d2h_stream=d2h)
|
||||
mock_pack_fields.assert_called_once_with(result, d2h_stream=d2h)
|
||||
|
||||
def test_runner_output_path_with_stream(self, mocker):
|
||||
mock_pack_fields = mocker.patch(
|
||||
"vllm_omni.diffusion.ipc._pack_diffusion_fields",
|
||||
|
||||
@@ -276,8 +276,6 @@ class TestRequestModeDispatch:
|
||||
def worker():
|
||||
request = req_q.get(timeout=2)
|
||||
assert "rpc_id" not in request
|
||||
assert request["output_rank"] is None
|
||||
assert request["collect_rank_status"] is False
|
||||
wave_id = request["wave_id"]
|
||||
for dp_rank in range(2):
|
||||
res_q.put(
|
||||
@@ -299,181 +297,6 @@ class TestRequestModeDispatch:
|
||||
|
||||
assert [result.error for result in results] == ["0", "1"]
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("tp_size", "sp_size", "pp_size", "cfg_size", "expected_primary_ranks"),
|
||||
[
|
||||
(2, 1, 1, 1, [0, 2]),
|
||||
(2, 2, 1, 1, [0, 4]),
|
||||
(1, 2, 2, 1, [0, 4]),
|
||||
(1, 1, 2, 2, [0, 4]),
|
||||
],
|
||||
)
|
||||
def test_step_execution_reads_each_dp_primary_result_queue(
|
||||
self,
|
||||
tp_size,
|
||||
sp_size,
|
||||
pp_size,
|
||||
cfg_size,
|
||||
expected_primary_ranks,
|
||||
):
|
||||
dp_size = 2
|
||||
num_gpus = dp_size * tp_size * sp_size * pp_size * cfg_size
|
||||
executor, _, _ = _make_executor(num_gpus=num_gpus)
|
||||
executor.od_config = SimpleNamespace(
|
||||
step_execution=True,
|
||||
parallel_config=SimpleNamespace(
|
||||
data_parallel_size=dp_size,
|
||||
tensor_parallel_size=tp_size,
|
||||
sequence_parallel_size=sp_size,
|
||||
pipeline_parallel_size=pp_size,
|
||||
cfg_parallel_size=cfg_size,
|
||||
),
|
||||
)
|
||||
request: dict = {}
|
||||
executor._broadcast_mq = SimpleNamespace(enqueue=lambda value: request.update(value))
|
||||
|
||||
def result_queue(global_rank):
|
||||
result_mq = MagicMock()
|
||||
result_mq.dequeue.side_effect = lambda timeout=None: {
|
||||
"dp_rank": global_rank,
|
||||
"output": _tagged_output(str(global_rank)),
|
||||
"wave_id": request["wave_id"],
|
||||
}
|
||||
return result_mq
|
||||
|
||||
# Every queue returns a valid, rank-tagged result. If the executor
|
||||
# selects a non-primary queue, the returned rank exposes the mistake.
|
||||
executor._result_mqs = [result_queue(rank) for rank in range(num_gpus)]
|
||||
executor._result_mq = executor._result_mqs[0]
|
||||
|
||||
results = executor.collective_rpc(
|
||||
"execute_model",
|
||||
unique_reply_rank=None,
|
||||
exec_all_ranks=True,
|
||||
)
|
||||
|
||||
assert [result.error for result in results] == [str(rank) for rank in expected_primary_ranks]
|
||||
for rank, result_mq in enumerate(executor._result_mqs):
|
||||
assert result_mq.dequeue.call_count == (1 if rank in expected_primary_ranks else 0)
|
||||
|
||||
@pytest.mark.parametrize("empty_prompt", ["", {"prompt": ""}])
|
||||
def test_dlo_dp_rejects_empty_prompt_before_worker_dispatch(self, empty_prompt):
|
||||
executor, _, _ = _make_executor(num_gpus=2)
|
||||
executor.od_config = SimpleNamespace(
|
||||
step_execution=False,
|
||||
parallel_config=SimpleNamespace(data_parallel_size=2),
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=True,
|
||||
)
|
||||
executor.collective_rpc = Mock()
|
||||
scheduler_output = _make_sched_output("invalid", "valid")
|
||||
scheduler_output.scheduled_new_reqs[0].req.prompt = empty_prompt
|
||||
|
||||
with pytest.raises(ValueError, match="non-empty prompt"):
|
||||
executor.execute_request(scheduler_output)
|
||||
|
||||
executor.collective_rpc.assert_not_called()
|
||||
|
||||
def test_dlo_dp_valid_wave_still_runs_after_rejected_wave(self):
|
||||
executor, _, _ = _make_executor(num_gpus=2)
|
||||
executor.od_config = SimpleNamespace(
|
||||
step_execution=False,
|
||||
parallel_config=SimpleNamespace(data_parallel_size=2),
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=True,
|
||||
)
|
||||
executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
|
||||
invalid_wave = _make_sched_output("invalid", "valid")
|
||||
invalid_wave.scheduled_new_reqs[0].req.prompt = ""
|
||||
|
||||
with pytest.raises(ValueError, match="non-empty prompt"):
|
||||
executor.execute_request(invalid_wave)
|
||||
|
||||
result = executor.execute_request(_make_sched_output("A", "B"))
|
||||
|
||||
assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
|
||||
executor.collective_rpc.assert_called_once()
|
||||
|
||||
def test_dlo_dp_allows_shared_default_denoise_steps(self):
|
||||
executor, _, _ = _make_executor(num_gpus=2)
|
||||
executor.od_config = SimpleNamespace(
|
||||
step_execution=False,
|
||||
parallel_config=SimpleNamespace(data_parallel_size=2),
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=True,
|
||||
)
|
||||
executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
|
||||
scheduler_output = _make_sched_output("A", "B")
|
||||
for new_req in scheduler_output.scheduled_new_reqs:
|
||||
new_req.req.sampling_params.num_inference_steps = None
|
||||
|
||||
result = executor.execute_request(scheduler_output)
|
||||
|
||||
assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
|
||||
executor.collective_rpc.assert_called_once()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("field", "value"),
|
||||
[
|
||||
("num_inference_steps", 2),
|
||||
("guidance_scale", 7.5),
|
||||
("width", 1024),
|
||||
],
|
||||
)
|
||||
def test_dlo_dp_rejects_incompatible_collective_wave(self, field, value):
|
||||
executor, _, _ = _make_executor(num_gpus=2)
|
||||
executor.od_config = SimpleNamespace(
|
||||
step_execution=False,
|
||||
parallel_config=SimpleNamespace(data_parallel_size=2),
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=True,
|
||||
)
|
||||
executor.collective_rpc = Mock()
|
||||
scheduler_output = _make_sched_output("A", "B")
|
||||
setattr(scheduler_output.scheduled_new_reqs[1].req.sampling_params, field, value)
|
||||
|
||||
with pytest.raises(ValueError, match="compatible shape, CFG"):
|
||||
executor.execute_request(scheduler_output)
|
||||
|
||||
executor.collective_rpc.assert_not_called()
|
||||
|
||||
def test_dlo_dp_partial_reply_times_out_and_fails_closed(self, monkeypatch):
|
||||
from vllm_omni.diffusion.executor import multiproc_executor as executor_module
|
||||
|
||||
executor, req_q, res_q = _make_executor(num_gpus=2)
|
||||
executor.od_config = SimpleNamespace(
|
||||
step_execution=False,
|
||||
parallel_config=SimpleNamespace(data_parallel_size=2),
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=True,
|
||||
)
|
||||
executor._sync_result_buffer = res_q
|
||||
executor._fail_closed_on_dp_wave_timeout = Mock()
|
||||
monkeypatch.setattr(executor_module, "_DLO_DP_WAVE_TIMEOUT_S", 0.05)
|
||||
|
||||
def reply_from_only_one_dp_rank():
|
||||
request = req_q.get(timeout=2.0)
|
||||
res_q.put(
|
||||
{
|
||||
"dp_rank": 0,
|
||||
"output": _tagged_output("rank-0"),
|
||||
"wave_id": request["wave_id"],
|
||||
}
|
||||
)
|
||||
|
||||
worker = threading.Thread(target=reply_from_only_one_dp_rank, daemon=True)
|
||||
worker.start()
|
||||
started = time.monotonic()
|
||||
|
||||
result = executor.execute_request(_make_sched_output("A", "B"))
|
||||
|
||||
elapsed = time.monotonic() - started
|
||||
worker.join(timeout=2.0)
|
||||
assert elapsed < 1.0
|
||||
assert all("timed out" in output.result.error for output in result.runner_outputs)
|
||||
executor._fail_closed_on_dp_wave_timeout.assert_called_once()
|
||||
assert isinstance(executor._fail_closed_on_dp_wave_timeout.call_args.args[0], TimeoutError)
|
||||
|
||||
|
||||
# ───────────────── concurrent collective RPC ─────────────────
|
||||
|
||||
@@ -1125,42 +948,6 @@ class TestStageDiffusionClientErrorPropagation:
|
||||
assert rpc_id not in client._pending_rpcs
|
||||
|
||||
|
||||
class TestExecutorShutdownCleaner:
|
||||
def test_worker_joins_share_one_global_deadline(self, monkeypatch):
|
||||
from vllm_omni.diffusion.executor import multiproc_executor as executor_module
|
||||
|
||||
class FakeProcess:
|
||||
def __init__(self, name):
|
||||
self.name = name
|
||||
self.alive = True
|
||||
self.terminated = False
|
||||
self.join_timeouts = []
|
||||
|
||||
def is_alive(self):
|
||||
return self.alive
|
||||
|
||||
def join(self, timeout):
|
||||
self.join_timeouts.append(timeout)
|
||||
if self.terminated:
|
||||
self.alive = False
|
||||
|
||||
def terminate(self):
|
||||
self.terminated = True
|
||||
|
||||
monotonic = Mock(side_effect=[100.0, 100.0, 110.0, 120.0, 120.0, 124.0])
|
||||
monkeypatch.setattr(executor_module, "time", SimpleNamespace(monotonic=monotonic))
|
||||
first = FakeProcess("worker-0")
|
||||
second = FakeProcess("worker-1")
|
||||
cleaner = executor_module._ExecutorShutdownCleaner(processes=[first, second])
|
||||
|
||||
cleaner()
|
||||
|
||||
assert first.join_timeouts == [15.0, 5.0]
|
||||
assert second.join_timeouts == [5.0, 1.0]
|
||||
assert first.terminated and second.terminated
|
||||
assert not first.is_alive() and not second.is_alive()
|
||||
|
||||
|
||||
# ───────── monitor thread & death sentinel integration tests ─────────
|
||||
|
||||
|
||||
|
||||
@@ -35,7 +35,6 @@ def _make_executor(step_execution=False):
|
||||
executor._pump_stop = threading.Event()
|
||||
executor._sync_result_buffer = queue.Queue()
|
||||
executor._result_mq = MagicMock()
|
||||
executor._result_mqs = []
|
||||
executor._broadcast_mq = MagicMock()
|
||||
executor._closed = False
|
||||
executor._is_failed = False
|
||||
@@ -158,41 +157,6 @@ class TestResultPumpDispatch:
|
||||
retrieved = executor._sync_result_buffer.get_nowait()
|
||||
assert isinstance(retrieved, DiffusionOutput)
|
||||
|
||||
def test_start_result_pump_reads_every_worker_queue(self):
|
||||
executor = _make_executor()
|
||||
messages = [DiffusionOutput(output="rank0"), DiffusionOutput(output="rank1")]
|
||||
result_mqs = [MagicMock(), MagicMock()]
|
||||
|
||||
def make_dequeue(message):
|
||||
emitted = False
|
||||
|
||||
def dequeue(timeout=None):
|
||||
nonlocal emitted
|
||||
if not emitted:
|
||||
emitted = True
|
||||
return message
|
||||
executor._pump_stop.wait(0.01)
|
||||
raise TimeoutError
|
||||
|
||||
return dequeue
|
||||
|
||||
for result_mq, message in zip(result_mqs, messages, strict=True):
|
||||
result_mq.dequeue = make_dequeue(message)
|
||||
|
||||
executor._result_mqs = result_mqs
|
||||
executor._result_mq = result_mqs[0]
|
||||
executor._start_result_pump()
|
||||
deadline = time.monotonic() + 2.0
|
||||
while executor._sync_result_buffer.qsize() < 2 and time.monotonic() < deadline:
|
||||
time.sleep(0.01)
|
||||
executor._pump_stop.set()
|
||||
for thread in executor._result_pump_threads:
|
||||
thread.join(timeout=2.0)
|
||||
|
||||
assert len(executor._result_pump_threads) == 2
|
||||
received = [executor._sync_result_buffer.get_nowait().output for _ in range(2)]
|
||||
assert sorted(received) == ["rank0", "rank1"]
|
||||
|
||||
def test_compute_done_routes_to_rpc_future(self):
|
||||
executor = _make_executor()
|
||||
rpc_id = "42"
|
||||
@@ -269,20 +233,6 @@ class TestResultPumpDispatch:
|
||||
class TestShutdownCleansUpFutures:
|
||||
"""Test that shutdown cancels pending futures."""
|
||||
|
||||
def test_shutdown_joins_result_pump_threads(self):
|
||||
executor = _make_executor()
|
||||
pump = threading.Thread(
|
||||
target=executor._pump_stop.wait,
|
||||
name="test-result-pump",
|
||||
)
|
||||
executor._result_pump_threads = [pump]
|
||||
pump.start()
|
||||
|
||||
executor.shutdown()
|
||||
|
||||
assert not pump.is_alive()
|
||||
assert executor._result_pump_threads == []
|
||||
|
||||
def test_shutdown_sets_exception_on_pending_futures(self):
|
||||
executor = _make_executor()
|
||||
|
||||
|
||||
@@ -29,7 +29,7 @@ HEIGHT = 512
|
||||
NUM_INFERENCE_STEPS = 20
|
||||
TRUE_CFG_SCALE = 4.0
|
||||
SEED = 42
|
||||
SSIM_THRESHOLD = 0.97
|
||||
SSIM_THRESHOLD = 0.94
|
||||
PSNR_THRESHOLD = 30.0
|
||||
|
||||
MODEL_2512_ID = "Qwen/Qwen-Image-2512"
|
||||
@@ -69,15 +69,7 @@ def _local_files_only(model: str) -> bool:
|
||||
|
||||
|
||||
def _run_vllm_omni_qwen_image(*, model: str, output_path: Path) -> Image.Image:
|
||||
server_args = [
|
||||
"--num-gpus",
|
||||
"1",
|
||||
"--stage-init-timeout",
|
||||
"300",
|
||||
"--init-timeout",
|
||||
"900",
|
||||
"--fa-deterministic",
|
||||
]
|
||||
server_args = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"]
|
||||
with OmniServer(model, server_args, use_omni=True) as omni_server:
|
||||
response = requests.post(
|
||||
f"http://{omni_server.host}:{omni_server.port}/v1/images/generations",
|
||||
|
||||
@@ -136,36 +136,3 @@ def test_flux_2_dev(
|
||||
}
|
||||
|
||||
openai_client.send_diffusion_request(request_config)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"omni_server",
|
||||
_get_flux_2_dev_feature_cases(MODEL),
|
||||
indirect=True,
|
||||
)
|
||||
def test_flux_2_dev_batched_chat_completions(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
):
|
||||
"""Validate batched chat completions for diffusion models."""
|
||||
messages = [
|
||||
[{"role": "user", "content": PROMPT}],
|
||||
[{"role": "user", "content": "A sunset over the ocean."}],
|
||||
]
|
||||
responses = openai_client.send_batched_chat_completions_http_request(
|
||||
{
|
||||
"json": {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"num_inference_steps": 2,
|
||||
"height": 512,
|
||||
"width": 512,
|
||||
"seed": 42,
|
||||
},
|
||||
},
|
||||
)
|
||||
assert responses and len(responses) == 1
|
||||
resp = responses[0]
|
||||
assert resp.success
|
||||
choices = resp.json_body["choices"]
|
||||
assert len(choices) == len(messages)
|
||||
@@ -250,101 +250,3 @@ def test_completions_rejected_for_thinker_talker(omni_server, openai_client) ->
|
||||
err_code=400,
|
||||
)
|
||||
assert not responses[0].success
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_batched_completions_text(omni_server, openai_client) -> None:
|
||||
"""Ensure that we can make a batch chat completions request (text only)."""
|
||||
responses = openai_client.send_batched_chat_completions_http_request(
|
||||
{
|
||||
"json": {
|
||||
"model": omni_server.model,
|
||||
"messages": [
|
||||
[{"role": "user", "content": "What color is the sky? Answer in one word."}],
|
||||
[{"role": "user", "content": "What is 2+2? Answer in one word."}],
|
||||
],
|
||||
"modalities": ["text"],
|
||||
"max_tokens": 20,
|
||||
},
|
||||
},
|
||||
)
|
||||
assert responses
|
||||
resp = responses[0]
|
||||
assert resp.success
|
||||
body = resp.json_body
|
||||
choices = body["choices"]
|
||||
assert len(choices) == 2
|
||||
for choice in choices:
|
||||
assert choice["message"]["content"]
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_batched_completions_audio_out(omni_server, openai_client) -> None:
|
||||
"""Ensure that we can make a batch chat completions request (audio + text)."""
|
||||
messages = [
|
||||
[{"role": "user", "content": "Say hello."}],
|
||||
[{"role": "user", "content": "Say goodbye."}],
|
||||
]
|
||||
num_messages = len(messages)
|
||||
responses = openai_client.send_batched_chat_completions_http_request(
|
||||
{
|
||||
"json": {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"modalities": ["text", "audio"],
|
||||
"max_tokens": 200,
|
||||
},
|
||||
},
|
||||
)
|
||||
assert responses and len(responses) == 1
|
||||
resp = responses[0]
|
||||
assert resp.success
|
||||
choices = resp.json_body["choices"]
|
||||
|
||||
assert len(choices) == num_messages
|
||||
|
||||
# Every message in the batch should contain both the audio and text content
|
||||
for choice in choices:
|
||||
assert choice["message"]["audio"]["data"]
|
||||
assert choice["message"]["content"]
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
@pytest.mark.parametrize(
|
||||
"sad_opts",
|
||||
[
|
||||
{"stream": "not-a-bool"}, # Fails validation
|
||||
{"n": 2}, # vLLM Batched endpoints doesn't support n > 1
|
||||
],
|
||||
)
|
||||
def test_batched_completions_with_bad_values(omni_server, openai_client, sad_opts) -> None:
|
||||
"""Ensure that a bad values are correctly handled as 400s."""
|
||||
messages = [
|
||||
[{"role": "user", "content": "Say hello."}],
|
||||
[{"role": "user", "content": "Say goodbye."}],
|
||||
]
|
||||
responses = openai_client.send_batched_chat_completions_http_request(
|
||||
{
|
||||
"json": {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
**sad_opts,
|
||||
},
|
||||
},
|
||||
)
|
||||
assert responses and len(responses) == 1
|
||||
resp = responses[0]
|
||||
assert not responses[0].success
|
||||
assert resp.status_code == 400
|
||||
@@ -1,150 +0,0 @@
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
from openai.types.chat import ChatCompletionUserMessageParam
|
||||
from openai.types.chat.chat_completion_audio import ChatCompletionAudio
|
||||
from starlette.requests import Request
|
||||
from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import (
|
||||
BatchChatCompletionRequest,
|
||||
ChatCompletionResponse,
|
||||
ChatCompletionResponseChoice,
|
||||
ChatMessage,
|
||||
)
|
||||
from vllm.entrypoints.openai.engine.protocol import ErrorResponse, UsageInfo
|
||||
|
||||
from vllm_omni.entrypoints.openai.batch_serving import OmniOpenAIServingChatBatch
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
MSG_BATCH: list[list[ChatCompletionMessageParam]] = [
|
||||
[ChatCompletionUserMessageParam(role="user", content="What color is the sky?")],
|
||||
[ChatCompletionUserMessageParam(role="user", content="What is 2+2?")],
|
||||
]
|
||||
collapse = OmniOpenAIServingChatBatch._maybe_collapse_choices
|
||||
|
||||
|
||||
# Helpers for testing creating a __new__ serving instance and ensuring IDs are correct
|
||||
def _make_raw_request(headers: list[tuple[bytes, bytes]]):
|
||||
scope = {
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"path": "/v1/chat/completions/batch",
|
||||
"headers": headers,
|
||||
}
|
||||
return Request(scope)
|
||||
|
||||
|
||||
def _make_completion(text):
|
||||
return ChatCompletionResponse(
|
||||
model="test-model",
|
||||
choices=[
|
||||
ChatCompletionResponseChoice(
|
||||
index=0,
|
||||
message=ChatMessage(role="assistant", content=text),
|
||||
)
|
||||
],
|
||||
usage=UsageInfo(prompt_tokens=0, completion_tokens=0, total_tokens=0),
|
||||
)
|
||||
|
||||
|
||||
def _make_handler(responses):
|
||||
handler = OmniOpenAIServingChatBatch.__new__(OmniOpenAIServingChatBatch)
|
||||
handler.create_chat_completion = AsyncMock(side_effect=responses)
|
||||
return handler
|
||||
|
||||
|
||||
# Helpers for creating choices of different modality types
|
||||
def _text_choice(text="hello"):
|
||||
return ChatCompletionResponseChoice(
|
||||
index=0,
|
||||
message=ChatMessage(role="assistant", content=text),
|
||||
)
|
||||
|
||||
|
||||
def _audio_choice():
|
||||
audio = ChatCompletionAudio(id="a1", data="base64audio", expires_at=0, transcript="")
|
||||
return ChatCompletionResponseChoice(
|
||||
index=0,
|
||||
message=ChatMessage(role="assistant", content=None, audio=audio),
|
||||
)
|
||||
|
||||
|
||||
def _image_choice():
|
||||
content = [{"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}}]
|
||||
msg = ChatMessage.model_construct(role="assistant")
|
||||
object.__setattr__(msg, "content", content)
|
||||
return ChatCompletionResponseChoice.model_construct(
|
||||
index=0,
|
||||
message=msg,
|
||||
)
|
||||
|
||||
|
||||
def _get_subrequest_ids(handler):
|
||||
return [call.args[0].request_id for call in handler.create_chat_completion.call_args_list]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("has_header", [True, False])
|
||||
def test_subrequest_ids_are_unique(has_header: bool):
|
||||
"""Ensure that a submitted batch request creates unique subrequest IDs"""
|
||||
header = None if not has_header else b"user-123"
|
||||
headers = [] if not header else [(b"x-request-id", header)]
|
||||
|
||||
handler = _make_handler([_make_completion("hello"), _make_completion("world")])
|
||||
request = BatchChatCompletionRequest(
|
||||
model="test-model",
|
||||
messages=MSG_BATCH,
|
||||
seed=42,
|
||||
)
|
||||
result = asyncio.run(handler.create_batch_chat_completion(request, _make_raw_request(headers)))
|
||||
assert not isinstance(result, ErrorResponse)
|
||||
assert result.id.startswith("chatcmpl-batch")
|
||||
|
||||
# Ensure request IDs are unique and that if we had a header, it's in all reqs
|
||||
req_ids = _get_subrequest_ids(handler)
|
||||
assert len(set(req_ids)) == len(req_ids) == 2
|
||||
|
||||
if header is not None:
|
||||
assert all(header.decode("utf-8") in req_id for req_id in req_ids)
|
||||
|
||||
|
||||
### Tests for ensuring that the (currently batch specific) choice collapsing is correct
|
||||
def test_single_text_passthrough():
|
||||
result = collapse([_text_choice()])
|
||||
assert result.message.content == "hello"
|
||||
|
||||
|
||||
def test_single_image_passthrough():
|
||||
result = collapse([_image_choice()])
|
||||
assert isinstance(result.message.content, list)
|
||||
assert result.message.content[0]["type"] == "image_url"
|
||||
|
||||
|
||||
def test_text_plus_audio():
|
||||
result = collapse([_text_choice(), _audio_choice()])
|
||||
assert result.message.content == "hello"
|
||||
assert result.message.audio.data == "base64audio"
|
||||
|
||||
|
||||
def test_audio_plus_text_order_independent():
|
||||
result = collapse([_audio_choice(), _text_choice()])
|
||||
assert result.message.content == "hello"
|
||||
assert result.message.audio.data == "base64audio"
|
||||
|
||||
|
||||
def test_image_plus_audio():
|
||||
result = collapse([_image_choice(), _audio_choice()])
|
||||
assert isinstance(result.message.content, list)
|
||||
assert result.message.content[0]["type"] == "image_url"
|
||||
assert result.message.audio.data == "base64audio"
|
||||
|
||||
|
||||
def test_two_content_choices_raises():
|
||||
with pytest.raises(ValueError, match="Multiple content choices cannot be set"):
|
||||
collapse([_text_choice(), _text_choice()])
|
||||
|
||||
|
||||
def test_three_choices_raises():
|
||||
with pytest.raises(ValueError, match="got 3"):
|
||||
collapse([_text_choice(), _audio_choice(), _text_choice()])
|
||||
@@ -95,7 +95,6 @@ def test_serving_boundary_normalizes_declared_root_and_nested_extras(serving_cha
|
||||
messages=[],
|
||||
modalities=["image"],
|
||||
num_inference_steps=7,
|
||||
quality="high",
|
||||
size="768x512",
|
||||
negative_prompt="avoid blur",
|
||||
lora={"name": "adapter"},
|
||||
@@ -111,7 +110,6 @@ def test_serving_boundary_normalizes_declared_root_and_nested_extras(serving_cha
|
||||
}
|
||||
assert diffusion_request_args["cfg_text_scale"] == 7.0
|
||||
assert diffusion_request_args["num_inference_steps"] == 7
|
||||
assert diffusion_request_args["quality"] == "high"
|
||||
assert diffusion_request_args["size"] == "768x512"
|
||||
assert diffusion_request_args["negative_prompt"] == "avoid blur"
|
||||
assert diffusion_request_args["lora"] == {"name": "adapter"}
|
||||
@@ -135,19 +133,6 @@ def test_unknown_root_extra_does_not_claim_canonical_extra(serving_chat):
|
||||
assert "pipeline_option" not in diffusion_request_args
|
||||
|
||||
|
||||
def test_serving_boundary_rejects_invalid_quality(serving_chat):
|
||||
serving_chat._diffusion_extra_body_params = frozenset()
|
||||
request = ChatCompletionRequest(
|
||||
model="test",
|
||||
messages=[],
|
||||
modalities=["image"],
|
||||
quality="medium",
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="quality must be one of"):
|
||||
serving_chat._normalize_diffusion_request_args(request)
|
||||
|
||||
|
||||
def test_unregistered_cfg_scale_aliases_common_true_cfg_scale(serving_chat):
|
||||
serving_chat._diffusion_extra_body_params = frozenset()
|
||||
request = ChatCompletionRequest(
|
||||
@@ -241,7 +226,6 @@ def test_pure_consumer_preserves_defaults_and_separate_cfg_owners(serving_chat,
|
||||
negative_prompt="avoid blur",
|
||||
cfg_scale=7.0,
|
||||
true_cfg_scale=2.0,
|
||||
quality="high",
|
||||
extra_body={"extra_args": {"solver": "ddim"}},
|
||||
)
|
||||
|
||||
@@ -255,7 +239,6 @@ def test_pure_consumer_preserves_defaults_and_separate_cfg_owners(serving_chat,
|
||||
}
|
||||
assert sampling_params.true_cfg_scale == 2.0
|
||||
assert sampling_params.num_inference_steps == 7
|
||||
assert sampling_params.quality == "high"
|
||||
assert (sampling_params.height, sampling_params.width) == (512, 768)
|
||||
assert captured["prompt"]["negative_prompt"] == "avoid blur"
|
||||
assert response.error.message == "No output generated from AsyncOmni"
|
||||
@@ -327,7 +310,6 @@ def test_mixed_consumer_keeps_root_common_args_with_nested_extras(serving_chat,
|
||||
assert sampling_params_list[0].extra_args == {"target_h": 512, "target_w": 768}
|
||||
diffusion_params = sampling_params_list[1]
|
||||
assert diffusion_params.num_inference_steps == 7
|
||||
assert diffusion_params.quality is None
|
||||
assert (diffusion_params.height, diffusion_params.width) == (512, 768)
|
||||
assert diffusion_params.extra_args == {"solver": "euler"}
|
||||
assert captured["prompt"]["negative_prompt"] == "avoid blur"
|
||||
|
||||
@@ -17,7 +17,6 @@ from pytest_mock import MockerFixture
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
from tests.helpers.media import generate_synthetic_image
|
||||
from vllm_omni.entrypoints.openai.protocol.videos import VideoGenerationRequest
|
||||
from vllm_omni.entrypoints.openai.serving_video_output_stream import OmniStreamingVideoOutputHandler
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
@@ -104,55 +103,6 @@ def _build_test_app(
|
||||
class TestStreamingVideoOutputWebSocket:
|
||||
"""WebSocket protocol tests for streaming generated video output."""
|
||||
|
||||
def test_sampling_params_quality_reaches_h3_pipeline_boundary(
|
||||
self,
|
||||
mocker: MockerFixture,
|
||||
) -> None:
|
||||
async def mock_generate(*_args, **_kwargs):
|
||||
if False:
|
||||
yield None
|
||||
|
||||
_app, handler, _engine_client = _build_test_app(
|
||||
mocker=mocker,
|
||||
streaming_chunks=[],
|
||||
mock_generate=mock_generate,
|
||||
)
|
||||
request_data = handler._coerce_request_data(
|
||||
{
|
||||
"type": "session.start",
|
||||
"prompt": "H3 streaming quality",
|
||||
"sampling_params": {"quality": "high"},
|
||||
}
|
||||
)
|
||||
request = VideoGenerationRequest(**request_data)
|
||||
|
||||
_prompt, sampling_params, _video_params = asyncio.run(handler._build_prompt_and_sampling_params(request))
|
||||
|
||||
assert sampling_params.quality == "high"
|
||||
|
||||
def test_omitted_quality_preserves_engine_default(self, mocker: MockerFixture) -> None:
|
||||
async def mock_generate(*_args, **_kwargs):
|
||||
if False:
|
||||
yield None
|
||||
|
||||
_app, handler, engine_client = _build_test_app(
|
||||
mocker=mocker,
|
||||
streaming_chunks=[],
|
||||
mock_generate=mock_generate,
|
||||
)
|
||||
engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams(quality="high")]
|
||||
request_data = handler._coerce_request_data(
|
||||
{
|
||||
"type": "session.start",
|
||||
"prompt": "H3 streaming default quality",
|
||||
}
|
||||
)
|
||||
request = VideoGenerationRequest(**request_data)
|
||||
|
||||
_prompt, sampling_params, _video_params = asyncio.run(handler._build_prompt_and_sampling_params(request))
|
||||
|
||||
assert sampling_params.quality == "high"
|
||||
|
||||
def test_streaming_session_emits_video_start_binary_chunks_and_done(self, mocker: MockerFixture):
|
||||
"""A full session delivers video.start, binary chunks, then session.done."""
|
||||
|
||||
|
||||
@@ -868,7 +868,6 @@ def test_default_sampling_params_apply_to_video_requests(test_client, mocker: Mo
|
||||
OmniDiffusionSamplingParams(
|
||||
num_inference_steps=4,
|
||||
guidance_scale=7.5,
|
||||
quality="high",
|
||||
generator_device="cpu",
|
||||
enable_frame_interpolation=True,
|
||||
frame_interpolation_exp=2,
|
||||
@@ -891,7 +890,6 @@ def test_default_sampling_params_apply_to_video_requests(test_client, mocker: Mo
|
||||
captured = engine.captured_sampling_params_list[0]
|
||||
assert captured.num_inference_steps == 4
|
||||
assert captured.guidance_scale == 7.5
|
||||
assert captured.quality == "high"
|
||||
assert captured.generator_device == "cpu"
|
||||
assert captured.enable_frame_interpolation is True
|
||||
assert captured.frame_interpolation_exp == 2
|
||||
@@ -1447,7 +1445,6 @@ def test_invalid_uploaded_input_reference_returns_400(test_client):
|
||||
def test_video_request_validation():
|
||||
req = VideoGenerationRequest(prompt="test")
|
||||
assert req.prompt == "test"
|
||||
assert req.quality is None
|
||||
assert req.generate_sound is False
|
||||
assert req.sound_duration is None
|
||||
assert VideoGenerationRequest(prompt="test", generate_sound=True, sound_duration=1.5).generate_sound is True
|
||||
@@ -1467,8 +1464,6 @@ def test_video_request_validation():
|
||||
VideoGenerationRequest(prompt="test", frame_interpolation_scale=0)
|
||||
with pytest.raises(ValueError):
|
||||
VideoGenerationRequest(prompt="test", sound_duration=0)
|
||||
with pytest.raises(ValueError):
|
||||
VideoGenerationRequest(prompt="test", quality="medium")
|
||||
|
||||
|
||||
def test_list_videos_supports_order_after_and_limit(test_client, mocker: MockerFixture):
|
||||
@@ -1980,7 +1975,6 @@ def test_sync_sampling_params_pass_through(test_client, mocker: MockerFixture):
|
||||
"num_inference_steps": "30",
|
||||
"guidance_scale": "6.5",
|
||||
"seed": "42",
|
||||
"quality": "high",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
@@ -1989,7 +1983,6 @@ def test_sync_sampling_params_pass_through(test_client, mocker: MockerFixture):
|
||||
assert captured.num_inference_steps == 30
|
||||
assert captured.guidance_scale == 6.5
|
||||
assert captured.seed == 42
|
||||
assert captured.quality == "high"
|
||||
|
||||
|
||||
def test_sync_frame_interpolation_params_pass_to_sampling_params(test_client, mocker: MockerFixture):
|
||||
@@ -2025,7 +2018,6 @@ def test_sync_default_sampling_params_apply_to_video_requests(test_client, mocke
|
||||
OmniDiffusionSamplingParams(
|
||||
num_inference_steps=4,
|
||||
guidance_scale=7.5,
|
||||
quality="high",
|
||||
enable_frame_interpolation=True,
|
||||
frame_interpolation_exp=2,
|
||||
frame_interpolation_scale=0.5,
|
||||
@@ -2046,7 +2038,6 @@ def test_sync_default_sampling_params_apply_to_video_requests(test_client, mocke
|
||||
captured = engine.captured_sampling_params_list[0]
|
||||
assert captured.num_inference_steps == 4
|
||||
assert captured.guidance_scale == 7.5
|
||||
assert captured.quality == "high"
|
||||
assert captured.enable_frame_interpolation is True
|
||||
assert captured.frame_interpolation_exp == 2
|
||||
assert captured.frame_interpolation_scale == 0.5
|
||||
|
||||
@@ -1252,28 +1252,6 @@ class OpenAIClientHandler:
|
||||
)
|
||||
return [resp]
|
||||
|
||||
def send_batched_chat_completions_http_request(
|
||||
self,
|
||||
request_config: dict[str, Any],
|
||||
*,
|
||||
err_code: int | tuple[int, ...] | list[int] | None = None,
|
||||
err_message: str | tuple[str, ...] | list[str] | None = None,
|
||||
) -> list[HttpResponse]:
|
||||
"""Post to batched chat completions."""
|
||||
cfg = _merge_http_expectation_kwargs(
|
||||
request_config,
|
||||
err_code=err_code,
|
||||
err_message=err_message,
|
||||
)
|
||||
r = self._post_json_endpoint("/v1/chat/completions/batch", cfg, default_timeout=120.0)
|
||||
resp = self._http_response_from_requests(r)
|
||||
assert_http_error(
|
||||
resp,
|
||||
err_code=cfg.get("err_code"),
|
||||
err_message=cfg.get("err_message"),
|
||||
)
|
||||
return [resp]
|
||||
|
||||
def send_completions_http_request(
|
||||
self,
|
||||
request_config: dict[str, Any],
|
||||
|
||||
@@ -25,23 +25,6 @@ def test_from_params_returns_existing_omni_params() -> None:
|
||||
assert converted is params
|
||||
|
||||
|
||||
@pytest.mark.parametrize("quality", ["lossless", "high"])
|
||||
def test_quality_accepts_supported_request_levels(quality: str) -> None:
|
||||
params = OmniDiffusionSamplingParams(quality=quality)
|
||||
|
||||
assert params.quality == quality
|
||||
|
||||
|
||||
def test_quality_defaults_to_model_owned_policy() -> None:
|
||||
params = OmniDiffusionSamplingParams()
|
||||
assert params.quality is None
|
||||
|
||||
|
||||
def test_quality_rejects_unsupported_request_level() -> None:
|
||||
with pytest.raises(ValueError, match="quality must be one of"):
|
||||
OmniDiffusionSamplingParams(quality="medium")
|
||||
|
||||
|
||||
def test_from_params_converts_sampling_params_seed_and_known_extra_args() -> None:
|
||||
params = _sampling_params(
|
||||
seed=1234,
|
||||
@@ -50,7 +33,6 @@ def test_from_params_converts_sampling_params_seed_and_known_extra_args() -> Non
|
||||
"height": 1024,
|
||||
"width": 768,
|
||||
"guidance_scale": 7.5,
|
||||
"quality": "high",
|
||||
},
|
||||
)
|
||||
|
||||
@@ -61,7 +43,6 @@ def test_from_params_converts_sampling_params_seed_and_known_extra_args() -> Non
|
||||
assert converted.height == 1024
|
||||
assert converted.width == 768
|
||||
assert converted.guidance_scale == 7.5
|
||||
assert converted.quality == "high"
|
||||
assert converted.extra_args == {}
|
||||
|
||||
|
||||
|
||||
@@ -21,6 +21,7 @@ class RouteTarget(NamedTuple):
|
||||
class OmniServingCapability(Enum):
|
||||
"""Serving capabilities that pipelines can shut down."""
|
||||
|
||||
CHAT_COMPLETIONS_BATCH = RouteTarget("/v1/chat/completions/batch", frozenset({"POST"}))
|
||||
COMPLETIONS = RouteTarget("/v1/completions", frozenset({"POST"}))
|
||||
|
||||
@property
|
||||
@@ -38,6 +39,16 @@ class EndpointRestriction:
|
||||
reason: str
|
||||
|
||||
|
||||
# Routes that are not supported for any model, but are supported in vLLM.
|
||||
# This is only temporary to avoid 500s for batched chat completions.
|
||||
UNSUPPORTED_ROUTES: tuple[EndpointRestriction, ...] = (
|
||||
EndpointRestriction(
|
||||
OmniServingCapability.CHAT_COMPLETIONS_BATCH,
|
||||
"Batched chat completions are not yet supported by vLLM Omni.",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def build_rejection_handler(reason: str):
|
||||
"""Build a rejection handler for a given endpoint for the provided reason."""
|
||||
|
||||
@@ -60,7 +71,11 @@ def shutdown_unsupported_routes(
|
||||
"""
|
||||
from vllm_omni.entrypoints.openai.api_server import _remove_route_from_app
|
||||
|
||||
for end_restrict in endpoint_restrictions:
|
||||
# Generally these should not overlap since there is no point. If they do,
|
||||
# we use the reason message in UNSUPPORTED_ROUTES, for consistent error messages.
|
||||
restricted_endpoints = (*endpoint_restrictions, *UNSUPPORTED_ROUTES)
|
||||
|
||||
for end_restrict in restricted_endpoints:
|
||||
capability = end_restrict.capability
|
||||
# Remove the route from the app
|
||||
_remove_route_from_app(app, capability.path, capability.methods)
|
||||
|
||||
@@ -568,7 +568,6 @@ class _DiffusionConfigProjection:
|
||||
pin_cpu_memory: bool = True
|
||||
diffusion_compile_granularity: Literal["regional", "full"] = "regional"
|
||||
diffusion_compile_dynamic: bool = Field(default=True, strict=True)
|
||||
fa_deterministic: bool = False
|
||||
vae_use_slicing: bool = False
|
||||
vae_use_tiling: bool = False
|
||||
mask_strategy_file_path: str | None = None
|
||||
|
||||
@@ -394,7 +394,6 @@ class StageDeployConfig:
|
||||
# Diffusion execution, cache, and VAE behavior.
|
||||
diffusion_compile_granularity: str | None = None
|
||||
diffusion_compile_dynamic: bool | None = None
|
||||
fa_deterministic: bool | None = None
|
||||
cache_backend: str | None = None
|
||||
cache_config: dict[str, Any] | None = None
|
||||
enable_cache_dit_summary: bool | None = None
|
||||
|
||||
@@ -15,7 +15,6 @@ from vllm_omni.diffusion.attention.backends.sdpa import _maybe_reshape_attn_mask
|
||||
from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
|
||||
piecewise_attn,
|
||||
)
|
||||
from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
@@ -72,20 +71,9 @@ class FlashAttentionImpl(AttentionImpl):
|
||||
self.causal = causal
|
||||
self.softmax_scale = softmax_scale
|
||||
self.qkv_layout = qkv_layout
|
||||
cfg = get_current_diffusion_config_or_none()
|
||||
self.fa_deterministic = bool(getattr(cfg, "fa_deterministic", False)) if cfg is not None else False
|
||||
if backend_kwargs:
|
||||
logger.warning("FlashAttentionImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys()))
|
||||
|
||||
def _warn_fa_deterministic_non_dense(self, path: str) -> None:
|
||||
if not self.fa_deterministic:
|
||||
return
|
||||
logger.warning_once(
|
||||
"fa_deterministic=True is ignored on the %s FlashAttention path; "
|
||||
"only the dense flash_attn_func path passes deterministic=True.",
|
||||
path,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _unwrap_flash_output(out: torch.Tensor | tuple[torch.Tensor, ...]) -> torch.Tensor:
|
||||
# FA3 may return (out, lse), FA2 returns out
|
||||
@@ -262,7 +250,6 @@ class FlashAttentionImpl(AttentionImpl):
|
||||
|
||||
# Try piecewise attention
|
||||
if full_attn_spans is not None:
|
||||
self._warn_fa_deterministic_non_dense("piecewise")
|
||||
logger.debug("Using piecewise Flash Attention for mixed causal/full mask")
|
||||
if flash_attn_func is not None:
|
||||
attn_func = partial(
|
||||
@@ -293,7 +280,6 @@ class FlashAttentionImpl(AttentionImpl):
|
||||
if len(present_packed_keys) != len(packed_keys):
|
||||
missing = sorted(set(packed_keys) - set(present_packed_keys))
|
||||
raise ValueError(f"Incomplete packed FlashAttention metadata; missing {missing}")
|
||||
self._warn_fa_deterministic_non_dense("packed-varlen")
|
||||
return self._forward_varlen_packed(
|
||||
query,
|
||||
key,
|
||||
@@ -305,7 +291,6 @@ class FlashAttentionImpl(AttentionImpl):
|
||||
)
|
||||
|
||||
if attention_mask is not None and torch.any(~attention_mask):
|
||||
self._warn_fa_deterministic_non_dense("masked-varlen")
|
||||
return self._forward_varlen_masked(
|
||||
query,
|
||||
key,
|
||||
@@ -314,16 +299,15 @@ class FlashAttentionImpl(AttentionImpl):
|
||||
)
|
||||
|
||||
if flash_attn_func is not None:
|
||||
fa_kwargs = {
|
||||
"causal": self.causal,
|
||||
"softmax_scale": self.softmax_scale,
|
||||
}
|
||||
if self.fa_deterministic:
|
||||
fa_kwargs["deterministic"] = True
|
||||
out = flash_attn_func(query, key, value, **fa_kwargs)
|
||||
out = flash_attn_func(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
causal=self.causal,
|
||||
softmax_scale=self.softmax_scale,
|
||||
)
|
||||
return self._unwrap_flash_output(out)
|
||||
|
||||
self._warn_fa_deterministic_non_dense("dense-varlen-fallback")
|
||||
return self._forward_varlen_dense(
|
||||
query,
|
||||
key,
|
||||
|
||||
@@ -1,10 +1,7 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
from packaging.version import InvalidVersion, Version
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.abstract import (
|
||||
@@ -16,8 +13,7 @@ from vllm_omni.diffusion.attention.backends.abstract import (
|
||||
logger = init_logger(__name__)
|
||||
|
||||
try:
|
||||
import flashinfer
|
||||
from flashinfer.prefill import BatchPrefillWithRaggedKVCacheWrapper
|
||||
from flashinfer.prefill import single_prefill_with_kv_cache
|
||||
|
||||
HAS_FLASHINFER = True
|
||||
except Exception as e:
|
||||
@@ -33,8 +29,10 @@ class FlashInferAttentionBackend(AttentionBackend):
|
||||
|
||||
@classmethod
|
||||
def supports_attention_mask(cls) -> bool:
|
||||
# Non-CuTe batch-prefill backends accept boolean custom masks. CuTe-DSL
|
||||
# falls back to SDPA for nontrivial masks.
|
||||
# FlashInfer single_prefill_with_kv_cache accepts a ``custom_mask`` kwarg
|
||||
# (2D boolean, ``True`` = keep) for non-causal attention. See reviewer
|
||||
# comment on #3079 and the flashinfer docs:
|
||||
# https://docs.flashinfer.ai/generated/flashinfer.prefill.single_prefill_with_kv_cache.html
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
@@ -54,27 +52,6 @@ class FlashInferAttentionBackend(AttentionBackend):
|
||||
|
||||
|
||||
class FlashInferAttentionImpl(AttentionImpl):
|
||||
_QK_DTYPES = {torch.float16, torch.bfloat16}
|
||||
_VO_DTYPES = {torch.float16, torch.bfloat16, torch.float8_e4m3fn}
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _WrapperPlanKey:
|
||||
batch_size: int
|
||||
qo_len: int
|
||||
kv_len: int
|
||||
num_q_heads: int
|
||||
num_kv_heads: int
|
||||
head_dim_qk: int
|
||||
head_dim_k: int
|
||||
head_dim_vo: int
|
||||
q_dtype: torch.dtype
|
||||
k_dtype: torch.dtype
|
||||
v_dtype: torch.dtype
|
||||
o_dtype: torch.dtype
|
||||
causal: bool
|
||||
softmax_scale: float
|
||||
has_custom_mask: bool
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_heads: int,
|
||||
@@ -83,92 +60,10 @@ class FlashInferAttentionImpl(AttentionImpl):
|
||||
causal: bool = False,
|
||||
num_kv_heads: int | None = None,
|
||||
prefix: str = "",
|
||||
backend_kwargs: dict | None = None,
|
||||
**extra_impl_args,
|
||||
) -> None:
|
||||
self.causal = causal
|
||||
self.softmax_scale = softmax_scale
|
||||
self.device = torch.device("cuda", torch.accelerator.current_device_index())
|
||||
backend_kwargs = backend_kwargs or {}
|
||||
quant = backend_kwargs.get("quant") or {}
|
||||
self.dtype_qk = self._check_dtype(quant.get("dtype_qk"), "dtype_qk", self._QK_DTYPES)
|
||||
self.dtype_vo = self._check_dtype(quant.get("dtype_vo"), "dtype_vo", self._VO_DTYPES)
|
||||
requested_backend = quant.get("flashinfer_backend", "auto")
|
||||
|
||||
if not HAS_FLASHINFER:
|
||||
raise ImportError("FLASHINFER_ATTN backend requires flashinfer")
|
||||
|
||||
self._check_flashinfer_version()
|
||||
|
||||
self.flashinfer_backend = self._select_backend(requested_backend, self.device)
|
||||
workspace_size = 0 if self.flashinfer_backend == "cute-dsl" else 128 * 1024 * 1024
|
||||
self._workspace = torch.empty(
|
||||
workspace_size,
|
||||
device=self.device,
|
||||
dtype=torch.uint8,
|
||||
)
|
||||
self._wrapper = BatchPrefillWithRaggedKVCacheWrapper(
|
||||
self._workspace,
|
||||
kv_layout="NHD",
|
||||
backend=self.flashinfer_backend,
|
||||
)
|
||||
self._qo_indptr: torch.Tensor | None = None
|
||||
self._kv_indptr: torch.Tensor | None = None
|
||||
self._plan_key: FlashInferAttentionImpl._WrapperPlanKey | None = None
|
||||
|
||||
if self.dtype_qk is not None or self.dtype_vo is not None:
|
||||
logger.info_once(
|
||||
"FLASHINFER_ATTN dtype override: Q/K=%s, V=%s.",
|
||||
self.dtype_qk,
|
||||
self.dtype_vo,
|
||||
)
|
||||
|
||||
logger.info_once(
|
||||
"FLASHINFER_ATTN initialized backend=%s on %s.",
|
||||
self.flashinfer_backend,
|
||||
self.device,
|
||||
)
|
||||
|
||||
def _check_flashinfer_version(self) -> None:
|
||||
if self.dtype_qk == self.dtype_vo:
|
||||
return
|
||||
try:
|
||||
flashinfer_version = Version(flashinfer.__version__)
|
||||
except (AttributeError, InvalidVersion):
|
||||
return
|
||||
if flashinfer_version <= Version("0.6.15"):
|
||||
raise RuntimeError(
|
||||
f"FlashInfer {flashinfer_version} is too old for reliable mixed "
|
||||
f"QK/V dtype attention (Q/K={self.dtype_qk}, V={self.dtype_vo}); "
|
||||
"install flashinfer >= 0.6.16rc1."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _select_backend(requested_backend: str, device: torch.device) -> str:
|
||||
if requested_backend != "auto":
|
||||
return requested_backend
|
||||
major, _minor = torch.cuda.get_device_capability(device)
|
||||
if major >= 10:
|
||||
return "cute-dsl"
|
||||
if major >= 9:
|
||||
return "fa3"
|
||||
return "fa2"
|
||||
|
||||
@classmethod
|
||||
def _check_dtype(
|
||||
cls,
|
||||
dtype: torch.dtype | str | None,
|
||||
option_name: str,
|
||||
allowed: set[torch.dtype],
|
||||
) -> torch.dtype | None:
|
||||
if dtype is None:
|
||||
return None
|
||||
if isinstance(dtype, str):
|
||||
dtype = torch.float8_e4m3fn if dtype == "fp8_e4m3" else getattr(torch, dtype)
|
||||
if dtype not in allowed:
|
||||
choices = ", ".join(sorted(str(item) for item in allowed))
|
||||
raise ValueError(f"Unsupported {option_name}={dtype}; expected one of: {choices}")
|
||||
return dtype
|
||||
|
||||
@staticmethod
|
||||
def _pack_mask_for_flashinfer(
|
||||
@@ -233,126 +128,6 @@ class FlashInferAttentionImpl(AttentionImpl):
|
||||
)
|
||||
return impl.forward_cuda(query, key, value, attn_metadata)
|
||||
|
||||
@torch.compiler.disable
|
||||
def _plan_wrapper(
|
||||
self,
|
||||
key: _WrapperPlanKey,
|
||||
flat_mask: torch.Tensor | None,
|
||||
) -> None:
|
||||
self._wrapper.plan(
|
||||
self._qo_indptr,
|
||||
self._kv_indptr,
|
||||
key.num_q_heads,
|
||||
key.num_kv_heads,
|
||||
key.head_dim_qk,
|
||||
head_dim_vo=key.head_dim_vo,
|
||||
custom_mask=flat_mask,
|
||||
causal=key.causal,
|
||||
sm_scale=key.softmax_scale,
|
||||
q_data_type=key.q_dtype,
|
||||
# FlashInfer versions newer than 0.6.15 can dispatch K and V with
|
||||
# independent runtime dtypes even though plan() names this K/V.
|
||||
kv_data_type=key.k_dtype,
|
||||
o_data_type=key.o_dtype,
|
||||
)
|
||||
|
||||
@torch.compiler.disable
|
||||
def _make_indptr(self, batch_size: int, sequence_length: int) -> torch.Tensor:
|
||||
return (
|
||||
torch.arange(
|
||||
batch_size + 1,
|
||||
device=self.device,
|
||||
dtype=torch.int32,
|
||||
)
|
||||
* sequence_length
|
||||
)
|
||||
|
||||
def _ensure_plan(
|
||||
self,
|
||||
key: _WrapperPlanKey,
|
||||
flat_mask: torch.Tensor | None,
|
||||
) -> None:
|
||||
key_changed = key != self._plan_key
|
||||
if not key_changed and flat_mask is None:
|
||||
return
|
||||
|
||||
if key_changed:
|
||||
self._qo_indptr = self._make_indptr(key.batch_size, key.qo_len)
|
||||
self._kv_indptr = self._make_indptr(key.batch_size, key.kv_len)
|
||||
|
||||
# A custom mask's contents may change without its shape changing, so
|
||||
# copy/re-plan it for every masked invocation.
|
||||
self._plan_wrapper(key, flat_mask)
|
||||
self._plan_key = key
|
||||
|
||||
def _run_batch_prefill(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
custom_mask: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
if query.device != self.device or key.device != self.device or value.device != self.device:
|
||||
raise ValueError(
|
||||
"FLASHINFER_ATTN inputs must remain on the layer initialization "
|
||||
f"device {self.device}; got Q={query.device}, K={key.device}, "
|
||||
f"V={value.device}"
|
||||
)
|
||||
|
||||
batch_size, qo_len, num_q_heads, head_dim_qk = query.shape
|
||||
kv_len = key.shape[1]
|
||||
num_kv_heads = key.shape[2]
|
||||
head_dim_k = key.shape[3]
|
||||
head_dim_vo = value.shape[3]
|
||||
|
||||
q = query.reshape(batch_size * qo_len, num_q_heads, head_dim_qk)
|
||||
k = key.reshape(batch_size * kv_len, num_kv_heads, head_dim_k)
|
||||
v = value.reshape(batch_size * kv_len, num_kv_heads, head_dim_vo)
|
||||
if self.dtype_qk is not None:
|
||||
q = q.to(self.dtype_qk)
|
||||
k = k.to(self.dtype_qk)
|
||||
if self.dtype_vo is not None:
|
||||
v = v.to(self.dtype_vo)
|
||||
|
||||
flat_mask = None
|
||||
if custom_mask is not None:
|
||||
if custom_mask.dim() == 2:
|
||||
custom_mask = custom_mask.unsqueeze(0).expand(batch_size, -1, -1)
|
||||
flat_mask = custom_mask.contiguous().view(-1)
|
||||
|
||||
self._ensure_plan(
|
||||
FlashInferAttentionImpl._WrapperPlanKey(
|
||||
batch_size=batch_size,
|
||||
qo_len=qo_len,
|
||||
kv_len=kv_len,
|
||||
num_q_heads=num_q_heads,
|
||||
num_kv_heads=num_kv_heads,
|
||||
head_dim_qk=head_dim_qk,
|
||||
head_dim_k=head_dim_k,
|
||||
head_dim_vo=head_dim_vo,
|
||||
q_dtype=q.dtype,
|
||||
k_dtype=k.dtype,
|
||||
v_dtype=v.dtype,
|
||||
o_dtype=query.dtype,
|
||||
causal=self.causal,
|
||||
softmax_scale=self.softmax_scale,
|
||||
has_custom_mask=flat_mask is not None,
|
||||
),
|
||||
flat_mask,
|
||||
)
|
||||
out = self._run_wrapper(q, k, v)
|
||||
out = out.reshape(batch_size, qo_len, num_q_heads, head_dim_vo)
|
||||
return out.to(query.dtype) if out.dtype != query.dtype else out
|
||||
|
||||
@torch.compiler.disable
|
||||
def _run_wrapper(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
return self._wrapper.run(query, key, value)
|
||||
|
||||
def forward_cuda(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
@@ -360,7 +135,18 @@ class FlashInferAttentionImpl(AttentionImpl):
|
||||
value: torch.Tensor,
|
||||
attn_metadata: AttentionMetadata | None = None,
|
||||
) -> torch.Tensor:
|
||||
if not HAS_FLASHINFER:
|
||||
raise ImportError(
|
||||
"FLASHINFER_ATTN backend requires flashinfer. "
|
||||
"Install it or set DIFFUSION_ATTENTION_BACKEND to another backend."
|
||||
)
|
||||
|
||||
# Try the custom_mask path; if the mask can't be packed into the
|
||||
# (qo_len, kv_len) layout FlashInfer expects, fall back to SDPA
|
||||
# instead of risking an illegal-memory-access crash in the kernel.
|
||||
# Input layout is (B, S, H, D); FlashInfer dense prefill takes (S, H, D).
|
||||
batch_size = query.shape[0]
|
||||
|
||||
custom_mask = None
|
||||
if attn_metadata is not None and attn_metadata.attn_mask is not None:
|
||||
try:
|
||||
@@ -370,14 +156,28 @@ class FlashInferAttentionImpl(AttentionImpl):
|
||||
qo_len=query.shape[1],
|
||||
kv_len=key.shape[1],
|
||||
)
|
||||
except ValueError as error:
|
||||
logger.debug("Falling back to SDPA for mask path: %s", error)
|
||||
except ValueError as e:
|
||||
logger.debug("Falling back to SDPA for mask path: %s", e)
|
||||
return self._sdpa_fallback(query, key, value, attn_metadata)
|
||||
# FlashInfer cannot combine causal masking with a custom_mask; rather
|
||||
# than silently dropping the explicit mask (diverging from SDPA), let
|
||||
# SDPA handle the causal+mask case correctly.
|
||||
if custom_mask is not None and self.causal:
|
||||
logger.debug("causal=True with explicit attn_mask; deferring to SDPA")
|
||||
return self._sdpa_fallback(query, key, value, attn_metadata)
|
||||
|
||||
if custom_mask is not None and self.flashinfer_backend == "cute-dsl":
|
||||
logger.debug("CuTe DSL does not support custom masks; deferring to SDPA")
|
||||
return self._sdpa_fallback(query, key, value, attn_metadata)
|
||||
outputs = []
|
||||
for b in range(batch_size):
|
||||
kwargs: dict = {
|
||||
"sm_scale": self.softmax_scale,
|
||||
"causal": self.causal,
|
||||
"return_lse": False,
|
||||
}
|
||||
if custom_mask is not None:
|
||||
kwargs["custom_mask"] = custom_mask if custom_mask.dim() == 2 else custom_mask[b]
|
||||
out = single_prefill_with_kv_cache(query[b], key[b], value[b], **kwargs)
|
||||
outputs.append(out)
|
||||
|
||||
return self._run_batch_prefill(query, key, value, custom_mask)
|
||||
if batch_size == 1:
|
||||
return outputs[0].unsqueeze(0)
|
||||
return torch.stack(outputs, dim=0)
|
||||
@@ -83,19 +83,6 @@ def pytorch_attn_forward(
|
||||
k = k.transpose(1, 2)
|
||||
v = v.transpose(1, 2)
|
||||
|
||||
# GQA/MQA support. Ring attention only ever *communicates* K/V with H_kv
|
||||
# heads -- that is precisely the benefit of GQA -- so expanding here does
|
||||
# not change communication volume. The expansion is needed because the aten
|
||||
# SDPA ops below are required for the LSE output that ring accumulation
|
||||
# consumes, and they do not accept `enable_gqa`.
|
||||
if k.shape[1] != q.shape[1]:
|
||||
num_heads, num_kv_heads = q.shape[1], k.shape[1]
|
||||
if num_heads % num_kv_heads != 0:
|
||||
raise ValueError(f"num_heads ({num_heads}) must be divisible by num_kv_heads ({num_kv_heads}) for GQA/MQA.")
|
||||
n_rep = num_heads // num_kv_heads
|
||||
k = k.repeat_interleave(n_rep, dim=1)
|
||||
v = v.repeat_interleave(n_rep, dim=1)
|
||||
|
||||
if op_type == "flash":
|
||||
out, lse = _scaled_dot_product_flash_attention(
|
||||
q,
|
||||
|
||||
@@ -204,25 +204,6 @@ class UlyssesParallelAttention:
|
||||
attn_metadata: AttentionMetadata | None,
|
||||
):
|
||||
mode = get_ulysses_mode(default="strict")
|
||||
ulysses_world_size = self._sp_group.ulysses_world_size
|
||||
|
||||
# advanced_uaa pads non-divisible head counts before the Ulysses
|
||||
# all-to-all. Padding K/V is not valid in hybrid Ulysses+Ring: the Ring
|
||||
# GQA/MQA path would repeat the padded zero heads as real K/V heads.
|
||||
# Reject this layout until K/V replication is performed before Ulysses.
|
||||
if mode == "advanced_uaa" and self._sp_group.ring_world_size > 1:
|
||||
for name, tensor in (("key", key), ("value", value)):
|
||||
kv_head_cnt = int(tensor.shape[2])
|
||||
if kv_head_cnt % ulysses_world_size != 0:
|
||||
raise ValueError(
|
||||
"ulysses_mode='advanced_uaa' with hybrid Ulysses+Ring "
|
||||
"does not support K/V head padding. "
|
||||
f"{name}_head_cnt={kv_head_cnt}, "
|
||||
f"ulysses_degree={ulysses_world_size}. "
|
||||
"Use ring_degree=1, choose a compatible ulysses_degree, "
|
||||
"or replicate K/V heads before Ulysses."
|
||||
)
|
||||
|
||||
joint_tensor_query = joint_tensor_key = joint_tensor_value = None
|
||||
joint_strategy = "front"
|
||||
joint_len = 0
|
||||
@@ -245,6 +226,7 @@ class UlyssesParallelAttention:
|
||||
|
||||
# Slice joint_query for this Ulysses rank
|
||||
# joint_query is (B, S, H, D). We split H (dim 2).
|
||||
ulysses_world_size = self._sp_group.ulysses_world_size
|
||||
ulysses_rank = self._sp_group.ulysses_rank
|
||||
joint_head_cnt = int(joint_tensor_query.shape[-2])
|
||||
joint_orig_head_cnt = joint_head_cnt
|
||||
|
||||
@@ -20,10 +20,6 @@ from vllm_omni.diffusion.cache.cachedit.model_specific import (
|
||||
from vllm_omni.diffusion.cache.cachedit.model_specific import (
|
||||
register_custom_dit_enablers as _register_custom_dit_enablers,
|
||||
)
|
||||
from vllm_omni.diffusion.cache.cachedit.runtime import (
|
||||
CacheDiTRequestSpec,
|
||||
RequestScopedCacheDiTRuntime,
|
||||
)
|
||||
|
||||
_register_custom_dit_enablers()
|
||||
|
||||
@@ -33,8 +29,6 @@ __all__ = [
|
||||
"CacheDiTAdapterConfig",
|
||||
"CacheDiTBackend",
|
||||
"CacheDiTConfig",
|
||||
"CacheDiTRequestSpec",
|
||||
"RequestScopedCacheDiTRuntime",
|
||||
"SensenovaCachedAdapter",
|
||||
"cache_summary",
|
||||
"enable_cache_for_dit",
|
||||
|
||||
+12
-68
@@ -6,7 +6,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from operator import attrgetter
|
||||
from typing import TYPE_CHECKING, Any, TypeAlias, cast
|
||||
|
||||
import cache_dit
|
||||
@@ -39,27 +38,14 @@ def _dit_module_names(pipeline: SupportsComponentDiscovery) -> tuple[str, ...]:
|
||||
names = getattr(pipeline, "_dit_modules", None)
|
||||
if not isinstance(names, (list, tuple)):
|
||||
names = ("transformer",)
|
||||
|
||||
resolved_names = []
|
||||
for name in names:
|
||||
if not isinstance(name, str):
|
||||
continue
|
||||
try:
|
||||
module = attrgetter(name)(pipeline)
|
||||
except AttributeError:
|
||||
continue
|
||||
if module is not None:
|
||||
resolved_names.append(name)
|
||||
return tuple(resolved_names)
|
||||
return tuple(name for name in names if isinstance(name, str) and getattr(pipeline, name, None) is not None)
|
||||
|
||||
|
||||
def cache_summary(pipeline: SupportsComponentDiscovery, details: bool = True) -> None:
|
||||
"""Log Cache-DiT statistics for every transformer on the pipeline."""
|
||||
|
||||
transformers = [attrgetter(name)(pipeline) for name in _dit_module_names(pipeline)]
|
||||
transformers = [getattr(pipeline, name) for name in _dit_module_names(pipeline)]
|
||||
for transformer in transformers:
|
||||
if not BlockAdapter.is_cached(transformer):
|
||||
continue
|
||||
cache_dit.summary(transformer, details=details)
|
||||
|
||||
if not transformers:
|
||||
@@ -78,7 +64,7 @@ def _make_pipeline_transformer_getter(
|
||||
def get_pipeline_transformer(
|
||||
pipeline: SupportsComponentDiscovery,
|
||||
) -> nn.Module:
|
||||
return cast(nn.Module, attrgetter(name)(pipeline))
|
||||
return cast(nn.Module, getattr(pipeline, name))
|
||||
|
||||
return get_pipeline_transformer
|
||||
|
||||
@@ -99,37 +85,19 @@ def _build_cache_context_refresh(
|
||||
# Cache-DiT has no predefined SCM mask for these small step counts.
|
||||
scm_supported_steps = num_inference_steps >= 8 or num_inference_steps in (4, 6)
|
||||
if projected_config.scm_steps_mask_policy is None or not scm_supported_steps:
|
||||
if projected_config.force_refresh_step_hint is not None:
|
||||
# Cache-DiT's ``once`` policy clears the hint after it fires.
|
||||
# Reapply it at every request boundary so a repeated request
|
||||
# with the same configured policy receives the same behavior.
|
||||
cache_dit.refresh_context(
|
||||
transformer,
|
||||
cache_config=DBCacheConfig().reset(
|
||||
num_inference_steps=num_inference_steps,
|
||||
force_refresh_step_hint=projected_config.force_refresh_step_hint,
|
||||
force_refresh_step_policy=projected_config.force_refresh_step_policy,
|
||||
),
|
||||
verbose=verbose,
|
||||
)
|
||||
return
|
||||
cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
|
||||
return
|
||||
|
||||
refresh_config = DBCacheConfig().reset(
|
||||
num_inference_steps=num_inference_steps,
|
||||
steps_computation_mask=cache_dit.steps_mask(
|
||||
mask_policy=projected_config.scm_steps_mask_policy,
|
||||
total_steps=num_inference_steps,
|
||||
),
|
||||
steps_computation_policy=projected_config.scm_steps_policy,
|
||||
)
|
||||
if projected_config.force_refresh_step_hint is not None:
|
||||
refresh_config.force_refresh_step_hint = projected_config.force_refresh_step_hint
|
||||
refresh_config.force_refresh_step_policy = projected_config.force_refresh_step_policy
|
||||
cache_dit.refresh_context(
|
||||
transformer,
|
||||
cache_config=refresh_config,
|
||||
cache_config=DBCacheConfig().reset(
|
||||
num_inference_steps=num_inference_steps,
|
||||
steps_computation_mask=cache_dit.steps_mask(
|
||||
mask_policy=projected_config.scm_steps_mask_policy,
|
||||
total_steps=num_inference_steps,
|
||||
),
|
||||
steps_computation_policy=projected_config.scm_steps_policy,
|
||||
),
|
||||
verbose=verbose,
|
||||
)
|
||||
|
||||
@@ -245,23 +213,19 @@ class CacheDiTBackend(CacheBackend):
|
||||
|
||||
super().__init__(config)
|
||||
self._refresh_funcs: list[RefreshCacheContextFunc] = []
|
||||
self._cache_targets: list[Any] = []
|
||||
|
||||
def enable(self, pipeline: SupportsComponentDiscovery) -> None:
|
||||
pipeline_name = type(pipeline).__name__
|
||||
custom_enabler = CUSTOM_DIT_ENABLERS.get(pipeline_name)
|
||||
self._refresh_funcs = []
|
||||
self._cache_targets = []
|
||||
if custom_enabler is not None:
|
||||
logger.info("Using custom cache-dit enabler for model: %s", pipeline_name)
|
||||
self._refresh_funcs = [custom_enabler(pipeline, self.config)]
|
||||
self._cache_targets = [_default_get_pipeline_transformer(pipeline)]
|
||||
else:
|
||||
self._refresh_funcs = []
|
||||
for name in _dit_module_names(pipeline):
|
||||
get_transformer = _make_pipeline_transformer_getter(name)
|
||||
block_adapter = _maybe_build_block_adapter(pipeline, get_transformer)
|
||||
adapter_cls = _maybe_get_cached_adapter_cls(pipeline, get_transformer)
|
||||
cache_target = get_transformer(pipeline) if block_adapter is None else block_adapter
|
||||
self._refresh_funcs.append(
|
||||
enable_cache_for_dit(
|
||||
pipeline,
|
||||
@@ -271,32 +235,12 @@ class CacheDiTBackend(CacheBackend):
|
||||
get_transformer,
|
||||
)
|
||||
)
|
||||
self._cache_targets.append(cache_target)
|
||||
if not self._refresh_funcs:
|
||||
raise ValueError(f"Pipeline {pipeline_name} has no declared DiT modules for Cache-DiT")
|
||||
|
||||
self.enabled = True
|
||||
logger.info("Cache-dit enabled successfully on %s", pipeline_name)
|
||||
|
||||
def disable(self, pipeline: SupportsComponentDiscovery) -> None:
|
||||
"""Remove Cache-DiT hooks so later requests use native forwards."""
|
||||
|
||||
if not self.enabled:
|
||||
return
|
||||
|
||||
logger.info(
|
||||
"Disabling cache-dit on %d DiT target(s) for %s",
|
||||
len(self._cache_targets),
|
||||
type(pipeline).__name__,
|
||||
)
|
||||
try:
|
||||
for target in self._cache_targets:
|
||||
cache_dit.disable_cache(target)
|
||||
finally:
|
||||
self._refresh_funcs = []
|
||||
self._cache_targets = []
|
||||
self.enabled = False
|
||||
|
||||
def refresh(self, pipeline: SupportsComponentDiscovery, num_inference_steps: int, verbose: bool = True) -> None:
|
||||
if not self.enabled or not self._refresh_funcs:
|
||||
logger.warning("Cache-dit is not enabled. Cannot refresh cache context.")
|
||||
|
||||
-94
@@ -1,94 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Declarative request-scoped Cache-DiT lifecycle."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from vllm_omni.diffusion.cache.cachedit.backend import CacheDiTBackend
|
||||
from vllm_omni.diffusion.data import DiffusionCacheConfig
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CacheDiTRequestSpec:
|
||||
"""Desired Cache-DiT state for one request.
|
||||
|
||||
``installation_key`` identifies configurations that require reinstalling
|
||||
hooks when switching between them. Step-count changes only require a
|
||||
context refresh.
|
||||
"""
|
||||
|
||||
installation_key: str
|
||||
cache_config: DiffusionCacheConfig
|
||||
num_inference_steps: int
|
||||
|
||||
|
||||
class RequestScopedCacheDiTRuntime:
|
||||
"""Reconcile declarative request state with installed Cache-DiT hooks."""
|
||||
|
||||
def __init__(self, pipeline: Any) -> None:
|
||||
self._pipeline = pipeline
|
||||
self._backend: CacheDiTBackend | None = None
|
||||
self._installation_key: str | None = None
|
||||
|
||||
@property
|
||||
def is_enabled(self) -> bool:
|
||||
return self._backend is not None and self._backend.is_enabled()
|
||||
|
||||
def adopt(self, backend: CacheDiTBackend, *, installation_key: str) -> None:
|
||||
"""Assume ownership of an already-enabled startup backend."""
|
||||
|
||||
if self._backend is not None:
|
||||
if self._backend is backend and self._installation_key == installation_key:
|
||||
return
|
||||
raise RuntimeError(
|
||||
f"Cache-DiT is already installed with key {self._installation_key!r}; cannot adopt {installation_key!r}"
|
||||
)
|
||||
if not backend.is_enabled():
|
||||
raise ValueError("Cannot adopt a Cache-DiT backend before it is enabled")
|
||||
|
||||
self._backend = backend
|
||||
self._installation_key = installation_key
|
||||
|
||||
def prepare(self, spec: CacheDiTRequestSpec | None) -> None:
|
||||
"""Apply the minimal transition needed before the request denoise.
|
||||
|
||||
``None`` disables installed hooks. An unchanged installation key only
|
||||
refreshes request context; a changed key disables the previous hooks
|
||||
before installing the requested profile.
|
||||
"""
|
||||
|
||||
desired_key = None if spec is None else spec.installation_key
|
||||
if self._installation_key != desired_key:
|
||||
self.disable()
|
||||
|
||||
if spec is None:
|
||||
return
|
||||
|
||||
if self._backend is None:
|
||||
backend = CacheDiTBackend(spec.cache_config)
|
||||
backend.enable(self._pipeline)
|
||||
self.adopt(backend, installation_key=spec.installation_key)
|
||||
|
||||
assert self._backend is not None
|
||||
self._backend.refresh(
|
||||
self._pipeline,
|
||||
num_inference_steps=spec.num_inference_steps,
|
||||
)
|
||||
|
||||
def disable(self) -> None:
|
||||
try:
|
||||
if self._backend is not None:
|
||||
self._backend.disable(self._pipeline)
|
||||
finally:
|
||||
self._backend = None
|
||||
self._installation_key = None
|
||||
|
||||
|
||||
__all__ = [
|
||||
"CacheDiTRequestSpec",
|
||||
"RequestScopedCacheDiTRuntime",
|
||||
]
|
||||
@@ -613,7 +613,6 @@ class OmniDiffusionConfig:
|
||||
|
||||
# Attention
|
||||
diffusion_attention_config: "AttentionConfig" = field(default_factory=lambda: AttentionConfig())
|
||||
fa_deterministic: bool = False
|
||||
|
||||
# Running mode
|
||||
# mode: ExecutionMode = ExecutionMode.INFERENCE
|
||||
|
||||
@@ -116,16 +116,6 @@ def _max_num_seqs(od_config: OmniDiffusionConfig) -> int:
|
||||
return 1
|
||||
|
||||
|
||||
def _uses_dlo_dp_concurrency(od_config: OmniDiffusionConfig) -> bool:
|
||||
parallel_config = getattr(od_config, "parallel_config", None)
|
||||
dp_size = getattr(parallel_config, "data_parallel_size", 1)
|
||||
return (
|
||||
dp_size > 1
|
||||
and getattr(od_config, "enable_distributed_layerwise_offload", False)
|
||||
and getattr(od_config, "dlo_use_allgather", True)
|
||||
)
|
||||
|
||||
|
||||
def _move_tensor_tree_to_cpu(value: object) -> object:
|
||||
if isinstance(value, torch.Tensor):
|
||||
return value.cpu() if value.device.type != "cpu" else value
|
||||
@@ -212,7 +202,7 @@ class DiffusionEngine:
|
||||
return DiffusionExecutionMode.STEP_BATCH
|
||||
|
||||
self.supports_request_batch = supports_request_batch(od_config)
|
||||
if not self.supports_request_batch and _max_num_seqs(od_config) > 1 and not _uses_dlo_dp_concurrency(od_config):
|
||||
if not self.supports_request_batch and _max_num_seqs(od_config) > 1:
|
||||
raise ValueError(
|
||||
f"{getattr(od_config, 'model_class_name', None)!r} does not support request-level batching. "
|
||||
"Use max_num_seqs=1 for serial request execution, or choose a pipeline with "
|
||||
@@ -243,10 +233,17 @@ class DiffusionEngine:
|
||||
# for distributed layerwise offload (which shards weights and
|
||||
# needs all ranks active simultaneously). Ordinary DP with a
|
||||
# non-batch pipeline should not schedule multiple requests.
|
||||
dp_size = getattr(getattr(self.od_config, "parallel_config", None), "data_parallel_size", 1)
|
||||
if _uses_dlo_dp_concurrency(self.od_config):
|
||||
dp_size = 1
|
||||
if getattr(self.od_config, "parallel_config", None) is not None:
|
||||
dp_size = getattr(self.od_config.parallel_config, "data_parallel_size", 1)
|
||||
dist_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False)
|
||||
if dp_size > 1 and dist_offload and getattr(self.od_config, "dlo_use_allgather", True):
|
||||
self.scheduler.max_num_running_reqs = dp_size
|
||||
self.dp_concurrent = True
|
||||
# Set batch admission wait so concurrent requests accumulate
|
||||
# before scheduling. The scheduler reads this from od_config.
|
||||
if getattr(self.od_config, "request_batch_max_wait_ms", 0) == 0:
|
||||
self.od_config.request_batch_max_wait_ms = 500.0
|
||||
logger.info(
|
||||
"dp_concurrent: max_num_running_reqs=%d, batch_wait=%sms",
|
||||
dp_size,
|
||||
|
||||
@@ -4,7 +4,6 @@ import concurrent.futures
|
||||
import json
|
||||
import multiprocessing as mp
|
||||
import multiprocessing.connection
|
||||
import os
|
||||
import queue
|
||||
import threading
|
||||
import time
|
||||
@@ -22,7 +21,6 @@ from vllm.v1.engine.exceptions import EngineDeadError
|
||||
from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, AsyncDiffusionOutput, AsyncOutputKind, DiffusionOutput
|
||||
from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
|
||||
from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, unpack_diffusion_output_shm
|
||||
from vllm_omni.diffusion.sched.request_scheduler import build_request_batch_sampling_params_key
|
||||
from vllm_omni.diffusion.worker import WorkerProc
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -32,21 +30,6 @@ if TYPE_CHECKING:
|
||||
logger = init_logger(__name__)
|
||||
|
||||
_DEQUEUE_TIMEOUT_S = 5.0
|
||||
_DLO_DP_WAVE_TIMEOUT_S = float(os.environ.get("VLLM_OMNI_DLO_DP_WAVE_TIMEOUT", 600.0))
|
||||
_WORKER_SHUTDOWN_GRACE_S = 15.0
|
||||
_WORKER_TERMINATE_GRACE_S = 5.0
|
||||
_RESULT_PUMP_JOIN_TIMEOUT_S = 2.0
|
||||
|
||||
|
||||
def _is_empty_dp_prompt(prompt: object) -> bool:
|
||||
"""Return whether a DP request has no usable text prompt."""
|
||||
if prompt is None:
|
||||
return True
|
||||
if isinstance(prompt, str):
|
||||
return not prompt
|
||||
if isinstance(prompt, dict) and "prompt" in prompt:
|
||||
return not prompt["prompt"]
|
||||
return False
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -69,20 +52,14 @@ class _ExecutorShutdownCleaner:
|
||||
logger.warning("Failed to send shutdown signal: %s", exc)
|
||||
|
||||
if self.processes:
|
||||
join_deadline = time.monotonic() + _WORKER_SHUTDOWN_GRACE_S
|
||||
for proc in self.processes:
|
||||
if not proc.is_alive():
|
||||
continue
|
||||
proc.join(max(0.0, join_deadline - time.monotonic()))
|
||||
|
||||
alive = [proc for proc in self.processes if proc.is_alive()]
|
||||
for proc in alive:
|
||||
logger.warning("Terminating diffusion worker %s after timeout", proc.name)
|
||||
proc.terminate()
|
||||
|
||||
terminate_deadline = time.monotonic() + _WORKER_TERMINATE_GRACE_S
|
||||
for proc in alive:
|
||||
proc.join(max(0.0, terminate_deadline - time.monotonic()))
|
||||
proc.join(5)
|
||||
if proc.is_alive():
|
||||
logger.warning("Terminating diffusion worker %s after timeout", proc.name)
|
||||
proc.terminate()
|
||||
proc.join(5)
|
||||
|
||||
|
||||
class MultiprocDiffusionExecutor(DiffusionExecutor):
|
||||
@@ -98,7 +75,6 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
|
||||
self._is_failed = False
|
||||
self._failure_callbacks: list[Callable[[], None]] = []
|
||||
self._result_mq: MessageQueue | None = None
|
||||
self._result_mqs: list[MessageQueue] = []
|
||||
self._rpc_wave_id: int = 0
|
||||
|
||||
num_workers = cast(int, self.od_config.num_gpus)
|
||||
@@ -108,7 +84,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
|
||||
broadcast_handle = self._broadcast_mq.export_handle()
|
||||
|
||||
# Launch workers
|
||||
processes, result_handles = self._launch_workers(broadcast_handle, self.wake_events)
|
||||
processes, result_handle = self._launch_workers(broadcast_handle, self.wake_events)
|
||||
self._processes = processes
|
||||
|
||||
shutdown_cleaner = _ExecutorShutdownCleaner(
|
||||
@@ -120,8 +96,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
|
||||
self._finalizer = weakref.finalize(self, shutdown_cleaner)
|
||||
|
||||
try:
|
||||
self._result_mqs = [self._init_result_queue(handle) for handle in result_handles]
|
||||
self._result_mq = self._result_mqs[0]
|
||||
self._result_mq = self._init_result_queue(result_handle)
|
||||
except Exception:
|
||||
self.shutdown()
|
||||
raise
|
||||
@@ -137,8 +112,8 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
|
||||
self._futures_lock = threading.RLock()
|
||||
self._pump_running = False
|
||||
self._pump_stop = threading.Event()
|
||||
# When pumps are active they are the sole readers of the worker result
|
||||
# queues; non-async messages are placed here for collective_rpc().
|
||||
# When pump is active it is the sole reader of result_mq; non-async
|
||||
# messages are placed here for collective_rpc() to consume.
|
||||
self._sync_result_buffer: queue.Queue = queue.Queue()
|
||||
if not self.od_config.step_execution:
|
||||
self._start_result_pump()
|
||||
@@ -165,17 +140,12 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
|
||||
if self._broadcast_mq is None:
|
||||
raise RuntimeError("Broadcast queue is closed")
|
||||
|
||||
def _dequeue_one_with_failure_polling(
|
||||
self,
|
||||
deadline: float | None,
|
||||
method: str,
|
||||
result_mq: MessageQueue | None = None,
|
||||
) -> Any:
|
||||
def _dequeue_one_with_failure_polling(self, deadline: float | None, method: str) -> Any:
|
||||
"""Block until one result message, polling ``_is_failed`` between chunk timeouts.
|
||||
|
||||
When async output is enabled, the pumps are the sole readers of the
|
||||
worker result queues; non-async messages are placed in
|
||||
_sync_result_buffer, so this method reads from there instead.
|
||||
When async output is enabled, pump is the sole reader of result_mq;
|
||||
non-async messages are placed in _sync_result_buffer, so this method
|
||||
reads from there instead of result_mq directly.
|
||||
"""
|
||||
while True:
|
||||
if deadline is None:
|
||||
@@ -193,11 +163,8 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
|
||||
raise EngineDeadError()
|
||||
continue
|
||||
else:
|
||||
queue_to_read = result_mq if result_mq is not None else self._result_mq
|
||||
if queue_to_read is None:
|
||||
raise RuntimeError("Result queue is closed")
|
||||
try:
|
||||
return queue_to_read.dequeue(timeout=chunk_timeout)
|
||||
return self._result_mq.dequeue(timeout=chunk_timeout) # pyright: ignore[reportOptionalMemberAccess] MQ is not None before shutdown
|
||||
except (TimeoutError, zmq.error.Again):
|
||||
if self._is_failed:
|
||||
raise EngineDeadError()
|
||||
@@ -256,14 +223,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
|
||||
|
||||
_MAX_STALE_DISCARDS = 16
|
||||
|
||||
def _validate_wave_id(
|
||||
self,
|
||||
response: Any,
|
||||
expected_wave_id: int,
|
||||
deadline: float | None,
|
||||
method: str,
|
||||
result_mq: MessageQueue | None = None,
|
||||
) -> Any:
|
||||
def _validate_wave_id(self, response: Any, expected_wave_id: int, deadline: float | None, method: str) -> Any:
|
||||
"""Discard stale RPC responses from a previous wave."""
|
||||
discards = 0
|
||||
while discards < self._MAX_STALE_DISCARDS:
|
||||
@@ -279,7 +239,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
|
||||
method,
|
||||
)
|
||||
discards += 1
|
||||
response = self._dequeue_one_with_failure_polling(deadline, method, result_mq)
|
||||
response = self._dequeue_one_with_failure_polling(deadline, method)
|
||||
raise TimeoutError(
|
||||
f"Discarded {self._MAX_STALE_DISCARDS} stale RPC responses "
|
||||
f"without finding wave_id={expected_wave_id} for method={method}."
|
||||
@@ -289,7 +249,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
|
||||
self,
|
||||
broadcast_handle: Handle,
|
||||
wake_events: list[Event],
|
||||
) -> tuple[list[mp.Process], list[Handle]]:
|
||||
) -> tuple[list[mp.Process], Handle | None]:
|
||||
od_config = self.od_config
|
||||
logger.info("Starting server...")
|
||||
|
||||
@@ -327,7 +287,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
|
||||
processes.append(process)
|
||||
|
||||
# Wait for all workers to be ready
|
||||
result_handles: list[Handle] = []
|
||||
result_handle = None
|
||||
for writer in scheduler_pipe_writers:
|
||||
writer.close()
|
||||
|
||||
@@ -343,16 +303,14 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
|
||||
if data["status"] != "ready":
|
||||
raise RuntimeError("Initialization failed. Please see the error messages above.")
|
||||
|
||||
result_handle = data.get("result_handle")
|
||||
if result_handle is None:
|
||||
raise RuntimeError(f"Rank {i} did not provide a result queue handle")
|
||||
result_handles.append(result_handle)
|
||||
if i == 0:
|
||||
result_handle = data.get("result_handle")
|
||||
|
||||
reader.close()
|
||||
|
||||
logger.debug("All workers are ready")
|
||||
|
||||
return processes, result_handles
|
||||
return processes, result_handle
|
||||
|
||||
@property
|
||||
def is_dead(self) -> bool:
|
||||
@@ -411,27 +369,6 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
|
||||
t = threading.Thread(target=_monitor, daemon=True, name="diffusion-worker-monitor")
|
||||
t.start()
|
||||
|
||||
def _fail_closed_on_dp_wave_timeout(self, exc: TimeoutError) -> None:
|
||||
"""Shut down every rank after a partial DP wave times out.
|
||||
|
||||
Once one rank is stuck in a DLO AllGather, that process group cannot be
|
||||
reused safely. Shutting down the executor converts an otherwise
|
||||
permanent request hang into a bounded engine failure and lets the
|
||||
caller restart.
|
||||
"""
|
||||
logger.error(
|
||||
"DLO DP collective wave timed out after %.1fs; shutting down the worker group: %s",
|
||||
_DLO_DP_WAVE_TIMEOUT_S,
|
||||
exc,
|
||||
)
|
||||
self._is_failed = True
|
||||
self.shutdown()
|
||||
for callback in self._failure_callbacks:
|
||||
try:
|
||||
callback()
|
||||
except Exception:
|
||||
logger.exception("failure_callback raised")
|
||||
|
||||
def register_failure_callback(
|
||||
self,
|
||||
callback: Callable[[], None],
|
||||
@@ -461,17 +398,21 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
|
||||
and getattr(self.od_config, "enable_distributed_layerwise_offload", False)
|
||||
and getattr(self.od_config, "dlo_use_allgather", True)
|
||||
):
|
||||
# Reuse the request scheduler's complete compatibility key. DLO
|
||||
# AllGather requires every DP rank to execute the same collective
|
||||
# schedule, including shape, CFG, denoise steps, output count,
|
||||
# and LoRA settings. Two default (None) step counts are valid and
|
||||
# resolve identically inside the same pipeline.
|
||||
compatibility_keys = [build_request_batch_sampling_params_key(nr.req) for nr in new_reqs]
|
||||
if any(key != compatibility_keys[0] for key in compatibility_keys[1:]):
|
||||
# Validate: all concurrent requests must share the same
|
||||
# num_inference_steps and identical extra_args, because
|
||||
# AllGather is a collective that requires every rank to
|
||||
# participate at each step.
|
||||
step_counts = {
|
||||
nr.req.sampling_params.num_inference_steps
|
||||
for nr in new_reqs
|
||||
if nr.req.sampling_params.num_inference_steps is not None
|
||||
}
|
||||
has_none = any(nr.req.sampling_params.num_inference_steps is None for nr in new_reqs)
|
||||
if (len(step_counts) > 1) or has_none:
|
||||
raise ValueError(
|
||||
"DLO DP multi-concurrency requires compatible shape, CFG, "
|
||||
"denoise schedule, output count, and LoRA settings for all "
|
||||
"requests in one collective wave."
|
||||
"DP multi-concurrency requires all concurrent requests to have "
|
||||
"the same explicit num_inference_steps (None is not allowed), got "
|
||||
f"{[nr.req.sampling_params.num_inference_steps for nr in new_reqs]}."
|
||||
)
|
||||
extra_args_signatures: set = set()
|
||||
for nr in new_reqs:
|
||||
@@ -483,18 +424,11 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
|
||||
"share identical extra_args. Different extra_args can change "
|
||||
"the forward schedule and cause AllGather deadlock."
|
||||
)
|
||||
empty_prompt_ids = [nr.request_id for nr in new_reqs if _is_empty_dp_prompt(nr.req.prompt)]
|
||||
if empty_prompt_ids:
|
||||
raise ValueError(
|
||||
"DP multi-concurrency requires a non-empty prompt for every request; "
|
||||
f"empty prompt request IDs: {empty_prompt_ids}."
|
||||
)
|
||||
|
||||
reqs_list = [nr.req for nr in new_reqs]
|
||||
try:
|
||||
results = self.collective_rpc(
|
||||
"execute_model",
|
||||
timeout=_DLO_DP_WAVE_TIMEOUT_S,
|
||||
args=(reqs_list, self.od_config, scheduler_output.kv_prefetch_job),
|
||||
unique_reply_rank=None,
|
||||
exec_all_ranks=True,
|
||||
@@ -514,8 +448,6 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
|
||||
else:
|
||||
raise RuntimeError(f"Unexpected response type [{i}]: {type(res)!r}")
|
||||
except Exception as exc:
|
||||
if isinstance(exc, TimeoutError):
|
||||
self._fail_closed_on_dp_wave_timeout(exc)
|
||||
for new_req in new_reqs:
|
||||
runner_outputs.append(
|
||||
RunnerOutput(
|
||||
@@ -659,18 +591,14 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
|
||||
deadline = None if timeout is None else time.monotonic() + timeout
|
||||
kwargs = kwargs or {}
|
||||
|
||||
multi_rank_reply = unique_reply_rank is None and exec_all_ranks
|
||||
execute_all_ranks = unique_reply_rank is None or exec_all_ranks
|
||||
# Status aggregation is for control-plane RPCs, where rank 0 sends
|
||||
# one envelope representing every rank. DP request concurrency needs
|
||||
# one independent reply per DP primary instead.
|
||||
collect_rank_status = unique_reply_rank is None and not multi_rank_reply
|
||||
collect_rank_status = unique_reply_rank is None
|
||||
rpc_request = {
|
||||
"type": "rpc",
|
||||
"method": method,
|
||||
"args": args,
|
||||
"kwargs": kwargs,
|
||||
"output_rank": None if multi_rank_reply else (unique_reply_rank if unique_reply_rank is not None else 0),
|
||||
"output_rank": unique_reply_rank if unique_reply_rank is not None else 0,
|
||||
"exec_all_ranks": execute_all_ranks,
|
||||
"collect_rank_status": collect_rank_status,
|
||||
}
|
||||
@@ -680,6 +608,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
|
||||
# so the D2H copy runs on a side stream in the worker background
|
||||
# thread while the default stream is free for the next forward.
|
||||
# pump routes the result back to a Future via rpc_id.
|
||||
multi_rank_reply = unique_reply_rank is None and exec_all_ranks
|
||||
if (
|
||||
not self.od_config.step_execution
|
||||
and method in ("execute_model", "execute_model_batch")
|
||||
@@ -722,28 +651,11 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
|
||||
responses: list = []
|
||||
if unique_reply_rank is None and exec_all_ranks and num_responses > 1:
|
||||
# DP multi-concurrency: collect num_responses replies, sort by dp_rank.
|
||||
result_mqs: list[MessageQueue | None] = [None] * num_responses
|
||||
if self.od_config.step_execution:
|
||||
parallel_config = self.od_config.parallel_config
|
||||
replica_size = (
|
||||
getattr(parallel_config, "tensor_parallel_size", 1)
|
||||
* getattr(parallel_config, "sequence_parallel_size", 1)
|
||||
* getattr(parallel_config, "pipeline_parallel_size", 1)
|
||||
* getattr(parallel_config, "cfg_parallel_size", 1)
|
||||
)
|
||||
primary_ranks = [dp_rank * replica_size for dp_rank in range(num_responses)]
|
||||
all_result_mqs = getattr(self, "_result_mqs", [])
|
||||
if not all_result_mqs or primary_ranks[-1] >= len(all_result_mqs):
|
||||
raise RuntimeError(
|
||||
"Missing result queues for DP primary ranks "
|
||||
f"{primary_ranks}; available queues: {len(all_result_mqs)}"
|
||||
)
|
||||
result_mqs = [all_result_mqs[rank] for rank in primary_ranks]
|
||||
tagged: list[tuple[int, Any]] = []
|
||||
collected_errors: list[str] = []
|
||||
for result_mq in result_mqs:
|
||||
response = self._dequeue_one_with_failure_polling(deadline, method, result_mq)
|
||||
response = self._validate_wave_id(response, wave_id, deadline, method, result_mq)
|
||||
for _ in range(num_responses):
|
||||
response = self._dequeue_one_with_failure_polling(deadline, method)
|
||||
response = self._validate_wave_id(response, wave_id, deadline, method)
|
||||
try:
|
||||
unpack_diffusion_output_shm(response)
|
||||
except Exception as e:
|
||||
@@ -786,35 +698,20 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
|
||||
def _start_result_pump(self) -> None:
|
||||
self._pump_running = True
|
||||
self._pump_stop.clear()
|
||||
result_mqs = self._result_mqs or ([self._result_mq] if self._result_mq is not None else [])
|
||||
self._result_pump_threads = [
|
||||
threading.Thread(
|
||||
target=self._result_pump,
|
||||
args=(result_mq,),
|
||||
daemon=True,
|
||||
name=f"DiffusionResultPump-{rank}",
|
||||
)
|
||||
for rank, result_mq in enumerate(result_mqs)
|
||||
]
|
||||
for thread in self._result_pump_threads:
|
||||
thread.start()
|
||||
self._result_pump_thread = self._result_pump_threads[0] if self._result_pump_threads else None
|
||||
logger.info("Async result pump started for %d worker queue(s)", len(self._result_pump_threads))
|
||||
self._result_pump_thread = threading.Thread(target=self._result_pump, daemon=True, name="DiffusionResultPump")
|
||||
self._result_pump_thread.start()
|
||||
logger.info("Async result pump started")
|
||||
|
||||
def _result_pump(self, result_mq: MessageQueue | None = None) -> None:
|
||||
"""Sole reader of one worker result queue when async output is enabled.
|
||||
def _result_pump(self) -> None:
|
||||
"""Sole reader of result_mq when async output is enabled.
|
||||
|
||||
Dispatches AsyncDiffusionOutput messages to the appropriate future:
|
||||
* RPC_RESULT / COMPUTE_DONE → _rpc_futures[rpc_id]
|
||||
* OUTPUT_READY → _output_futures[async_output_id]
|
||||
"""
|
||||
result_mq = result_mq or self._result_mq
|
||||
if result_mq is None:
|
||||
return
|
||||
|
||||
while not self._pump_stop.is_set():
|
||||
try:
|
||||
msg = result_mq.dequeue(timeout=1.0)
|
||||
msg = self._result_mq.dequeue(timeout=1.0)
|
||||
except TimeoutError:
|
||||
if self._is_failed:
|
||||
break
|
||||
@@ -926,18 +823,8 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
|
||||
try:
|
||||
self._finalizer()
|
||||
finally:
|
||||
pump_threads = getattr(self, "_result_pump_threads", [])
|
||||
for thread in pump_threads:
|
||||
if thread is threading.current_thread():
|
||||
continue
|
||||
thread.join(timeout=_RESULT_PUMP_JOIN_TIMEOUT_S)
|
||||
if thread.is_alive():
|
||||
logger.warning("Result pump thread %s did not stop before shutdown", thread.name)
|
||||
self._pump_running = False
|
||||
self._broadcast_mq = None
|
||||
self._result_mq = None
|
||||
self._result_mqs = []
|
||||
self._result_pump_threads = []
|
||||
with self._futures_lock:
|
||||
for fut in self._rpc_futures.values():
|
||||
if not fut.done():
|
||||
|
||||
@@ -8,7 +8,6 @@ IPC overhead. Used when there is only a single diffusion stage.
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import TYPE_CHECKING, Any
|
||||
@@ -67,8 +66,6 @@ class InlineStageDiffusionClient(StageClientBase):
|
||||
self._tasks: dict[str, asyncio.Task] = {}
|
||||
self._engine_dead = False
|
||||
self._shutting_down = False
|
||||
self._shutdown_complete = False
|
||||
self._shutdown_lock = threading.Lock()
|
||||
|
||||
self._engine.executor.register_failure_callback(self._mark_engine_dead)
|
||||
|
||||
@@ -104,10 +101,6 @@ class InlineStageDiffusionClient(StageClientBase):
|
||||
sampling_params: OmniDiffusionSamplingParams,
|
||||
kv_sender_info: dict[int, dict[str, Any]] | None = None,
|
||||
) -> None:
|
||||
# Each request mutates its sampling state while it is normalized and
|
||||
# executed. Callers commonly reuse one params object for concurrent
|
||||
# requests, so take the copy synchronously before either task starts.
|
||||
sampling_params = sampling_params.clone()
|
||||
logger.debug(
|
||||
"[InlineStageDiffusionClient] stage-%s [rep-%s] add request: %s",
|
||||
self.stage_id,
|
||||
@@ -305,25 +298,19 @@ class InlineStageDiffusionClient(StageClientBase):
|
||||
raise
|
||||
|
||||
def shutdown(self) -> None:
|
||||
with self._shutdown_lock:
|
||||
if self._shutdown_complete:
|
||||
return
|
||||
self._shutting_down = True
|
||||
self._shutting_down = True
|
||||
|
||||
# Cancel all pending tasks
|
||||
for task in self._tasks.values():
|
||||
task.cancel()
|
||||
# Cancel all pending tasks
|
||||
for task in self._tasks.values():
|
||||
task.cancel()
|
||||
|
||||
try:
|
||||
# Stop the engine first so any control RPC running in the thread
|
||||
# pool can observe shutdown instead of keeping stage teardown
|
||||
# blocked while the executor waits for that RPC.
|
||||
self._engine.close()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
# Cancel queued futures and wait for the running one to complete deterministically
|
||||
self._executor.shutdown(wait=True, cancel_futures=True)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
try:
|
||||
self._executor.shutdown(wait=True, cancel_futures=True)
|
||||
except Exception:
|
||||
pass
|
||||
self._shutdown_complete = True
|
||||
try:
|
||||
self._engine.close()
|
||||
except Exception:
|
||||
pass
|
||||
+42
-74
@@ -13,7 +13,6 @@ from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput
|
||||
@@ -22,45 +21,6 @@ _SHM_TENSOR_THRESHOLD = 1_000_000 # 1 MB
|
||||
DIFFUSION_RPC_RESULT_ENVELOPE = "diffusion_rpc_result"
|
||||
|
||||
|
||||
def _array_to_shm(array: np.ndarray) -> dict[str, Any]:
|
||||
"""Copy a contiguous NumPy-compatible array into shared memory."""
|
||||
from multiprocessing import shared_memory
|
||||
|
||||
if array.dtype.hasobject:
|
||||
raise TypeError("NumPy object arrays cannot be transferred through raw shared memory")
|
||||
|
||||
array = np.ascontiguousarray(array)
|
||||
nbytes = array.nbytes
|
||||
shm = shared_memory.SharedMemory(create=True, size=nbytes)
|
||||
shm_array = np.ndarray(array.shape, dtype=array.dtype, buffer=shm.buf[:nbytes])
|
||||
np.copyto(shm_array, array)
|
||||
handle = {
|
||||
"name": shm.name,
|
||||
"shape": list(array.shape),
|
||||
"numpy_dtype": str(array.dtype),
|
||||
"nbytes": nbytes,
|
||||
}
|
||||
shm.close()
|
||||
return handle
|
||||
|
||||
|
||||
def _array_from_shm(handle: dict[str, Any]) -> np.ndarray:
|
||||
"""Copy an array from shared memory, then close and unlink its segment."""
|
||||
from multiprocessing import shared_memory
|
||||
|
||||
shm = shared_memory.SharedMemory(name=handle["name"])
|
||||
try:
|
||||
array = np.ndarray(
|
||||
handle["shape"],
|
||||
dtype=np.dtype(handle["numpy_dtype"]),
|
||||
buffer=shm.buf[: handle["nbytes"]],
|
||||
).copy()
|
||||
finally:
|
||||
shm.close()
|
||||
shm.unlink()
|
||||
return array
|
||||
|
||||
|
||||
def _tensor_to_shm(
|
||||
tensor: torch.Tensor,
|
||||
d2h_stream: torch.Stream | None = None,
|
||||
@@ -75,6 +35,10 @@ def _tensor_to_shm(
|
||||
path. The caller must synchronize *d2h_stream* after all tensors are
|
||||
packed.
|
||||
"""
|
||||
from multiprocessing import shared_memory
|
||||
|
||||
import numpy as np
|
||||
|
||||
original_dtype = tensor.dtype
|
||||
if d2h_stream is not None:
|
||||
# Non-blocking D2H: copy on side stream to pinned CPU memory.
|
||||
@@ -94,26 +58,44 @@ def _tensor_to_shm(
|
||||
tensor = tensor.detach().cpu().contiguous()
|
||||
if original_dtype == torch.bfloat16:
|
||||
tensor = tensor.to(torch.float32)
|
||||
handle = _array_to_shm(tensor.numpy())
|
||||
handle.update(
|
||||
{
|
||||
"__tensor_shm__": True,
|
||||
"torch_dtype": str(original_dtype),
|
||||
}
|
||||
)
|
||||
arr = tensor.numpy()
|
||||
nbytes = arr.nbytes
|
||||
shm = shared_memory.SharedMemory(create=True, size=nbytes)
|
||||
shm_arr = np.ndarray(arr.shape, dtype=arr.dtype, buffer=shm.buf[:nbytes])
|
||||
np.copyto(shm_arr, arr)
|
||||
handle = {
|
||||
"__tensor_shm__": True,
|
||||
"name": shm.name,
|
||||
"shape": list(tensor.shape),
|
||||
"torch_dtype": str(original_dtype),
|
||||
"numpy_dtype": str(arr.dtype),
|
||||
"nbytes": nbytes,
|
||||
}
|
||||
shm.close()
|
||||
return handle
|
||||
|
||||
|
||||
def _tensor_from_shm(handle: dict[str, Any]) -> torch.Tensor:
|
||||
"""Reconstruct a tensor from a shared-memory handle and free the segment."""
|
||||
tensor = torch.from_numpy(_array_from_shm(handle))
|
||||
# Restore the original dtype if it differs from the numpy-compatible
|
||||
# dtype used for the SHM transfer (e.g. bfloat16 → float32 → bfloat16).
|
||||
torch_dtype_str = handle.get("torch_dtype", "")
|
||||
if torch_dtype_str:
|
||||
original_dtype = getattr(torch, torch_dtype_str.replace("torch.", ""), None)
|
||||
if original_dtype is not None and tensor.dtype != original_dtype:
|
||||
tensor = tensor.to(original_dtype)
|
||||
from multiprocessing import shared_memory
|
||||
|
||||
import numpy as np
|
||||
|
||||
shm = shared_memory.SharedMemory(name=handle["name"])
|
||||
try:
|
||||
np_dtype = np.dtype(handle["numpy_dtype"])
|
||||
arr = np.ndarray(handle["shape"], dtype=np_dtype, buffer=shm.buf[: handle["nbytes"]])
|
||||
tensor = torch.from_numpy(arr.copy())
|
||||
# Restore the original dtype if it differs from the numpy-compatible
|
||||
# dtype used for the SHM transfer (e.g. bfloat16 → float32 → bfloat16).
|
||||
torch_dtype_str = handle.get("torch_dtype", "")
|
||||
if torch_dtype_str:
|
||||
original_dtype = getattr(torch, torch_dtype_str.replace("torch.", ""), None)
|
||||
if original_dtype is not None and tensor.dtype != original_dtype:
|
||||
tensor = tensor.to(original_dtype)
|
||||
finally:
|
||||
shm.close()
|
||||
shm.unlink()
|
||||
return tensor
|
||||
|
||||
|
||||
@@ -127,18 +109,6 @@ def _pack_tensor_if_large(
|
||||
return val
|
||||
|
||||
|
||||
def _ndarray_to_shm(array: np.ndarray) -> dict[str, Any]:
|
||||
"""Copy a contiguous NumPy array into POSIX shared memory."""
|
||||
handle = _array_to_shm(array)
|
||||
handle["__ndarray_shm__"] = True
|
||||
return handle
|
||||
|
||||
|
||||
def _ndarray_from_shm(handle: dict[str, Any]) -> np.ndarray:
|
||||
"""Reconstruct a NumPy array from shared memory and free the segment."""
|
||||
return _array_from_shm(handle)
|
||||
|
||||
|
||||
def _pack_value_if_large(
|
||||
val: object,
|
||||
d2h_stream: torch.Stream | None = None,
|
||||
@@ -153,8 +123,6 @@ def _pack_value_if_large(
|
||||
"""
|
||||
if isinstance(val, torch.Tensor):
|
||||
return _pack_tensor_if_large(val, d2h_stream=d2h_stream)
|
||||
if isinstance(val, np.ndarray):
|
||||
return _ndarray_to_shm(val) if not val.dtype.hasobject and val.nbytes > _SHM_TENSOR_THRESHOLD else val
|
||||
if isinstance(val, dict):
|
||||
return {key: _pack_value_if_large(value, d2h_stream=d2h_stream) for key, value in val.items()}
|
||||
if isinstance(val, list):
|
||||
@@ -168,8 +136,6 @@ def _unpack_if_shm_handle(val: object) -> object:
|
||||
"""Reconstruct tensors from SHM handles, mirroring ``_pack_value_if_large``."""
|
||||
if isinstance(val, dict) and val.get("__tensor_shm__"):
|
||||
return _tensor_from_shm(val)
|
||||
if isinstance(val, dict) and val.get("__ndarray_shm__"):
|
||||
return _ndarray_from_shm(val)
|
||||
if isinstance(val, dict):
|
||||
return {key: _unpack_if_shm_handle(value) for key, value in val.items()}
|
||||
if isinstance(val, list):
|
||||
@@ -220,12 +186,13 @@ def pack_diffusion_output_shm(
|
||||
if isinstance(output, dict) and "dp_rank" in output and "output" in output:
|
||||
inner = output["output"]
|
||||
if isinstance(inner, DiffusionOutput):
|
||||
output["output"] = _pack_diffusion_fields(inner, d2h_stream=d2h_stream)
|
||||
output["output"] = _pack_diffusion_fields(inner)
|
||||
return output
|
||||
|
||||
if _is_rpc_result_envelope(output):
|
||||
result = output.get("result")
|
||||
output["result"] = pack_diffusion_output_shm(result, d2h_stream=d2h_stream)
|
||||
if isinstance(result, DiffusionOutput):
|
||||
output["result"] = _pack_diffusion_fields(result, d2h_stream=d2h_stream)
|
||||
return output
|
||||
|
||||
result = getattr(output, "result", None)
|
||||
@@ -261,7 +228,8 @@ def unpack_diffusion_output_shm(output: object) -> object:
|
||||
|
||||
if _is_rpc_result_envelope(output):
|
||||
result = output.get("result")
|
||||
output["result"] = unpack_diffusion_output_shm(result)
|
||||
if isinstance(result, DiffusionOutput):
|
||||
output["result"] = _unpack_diffusion_fields(result)
|
||||
return output
|
||||
|
||||
result = getattr(output, "result", None)
|
||||
|
||||
@@ -64,21 +64,4 @@ def get_diffusion_model_metadata(model_class_name: str | None) -> DiffusionModel
|
||||
# pipelines do not accidentally inherit limits meant for other models.
|
||||
if model_class_name is None:
|
||||
return DiffusionModelMetadata()
|
||||
metadata = _DIFFUSION_MODEL_METADATA.get(model_class_name)
|
||||
if metadata is not None:
|
||||
return metadata
|
||||
# Some checkpoints report the HF architecture name diff from internal pipeline class name
|
||||
# (e.g. HunyuanImage3ForCausalMM, WanVACEPipeline, OmniVoice ...).
|
||||
from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
|
||||
|
||||
entry = _DIFFUSION_MODELS.get(model_class_name)
|
||||
if entry is not None:
|
||||
# Unpack instead of indexing so a future change to the registry tuple
|
||||
# shape fails loudly instead of silently reading the wrong element.
|
||||
_, _, pipeline_cls_name = entry
|
||||
# Note: the registry ``cls_name`` and the metadata keys are two separate
|
||||
# key spaces. Aliases whose pipeline class has no metadata entry (e.g.
|
||||
# Wan22VACEPipeline, OmniVoicePipeline) still fall back to the defaults;
|
||||
# that is not a regression, it just means no capability override.
|
||||
return _DIFFUSION_MODEL_METADATA.get(pipeline_cls_name, DiffusionModelMetadata())
|
||||
return DiffusionModelMetadata()
|
||||
return _DIFFUSION_MODEL_METADATA.get(model_class_name, DiffusionModelMetadata())
|
||||
@@ -17,7 +17,6 @@ if TYPE_CHECKING:
|
||||
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.cache.cachedit import CacheDiTBackend
|
||||
from vllm_omni.diffusion.data import DiffusionOutput
|
||||
from vllm_omni.diffusion.worker.input_batch import InputBatch
|
||||
from vllm_omni.diffusion.worker.utils import StepRequestState
|
||||
@@ -132,31 +131,3 @@ def supports_prompt_update(pipeline: object) -> bool:
|
||||
"""Return whether ``pipeline`` implements :class:`SupportsPromptUpdate`."""
|
||||
|
||||
return isinstance(pipeline, SupportsPromptUpdate)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class SupportsRequestScopedCacheDiT(Protocol):
|
||||
"""Optional protocol for pipelines that own Cache-DiT hook transitions."""
|
||||
|
||||
def adopt_cache_dit_backend(self, backend: CacheDiTBackend) -> None:
|
||||
"""Assume ownership of an enabled Cache-DiT backend."""
|
||||
...
|
||||
|
||||
def is_cache_dit_enabled(self) -> bool:
|
||||
"""Return whether this pipeline currently has Cache-DiT installed."""
|
||||
...
|
||||
|
||||
|
||||
def adopt_request_scoped_cache_dit(pipeline: object, backend: CacheDiTBackend) -> bool:
|
||||
"""Transfer an enabled Cache-DiT backend to an opted-in pipeline."""
|
||||
|
||||
if not isinstance(pipeline, SupportsRequestScopedCacheDiT):
|
||||
return False
|
||||
pipeline.adopt_cache_dit_backend(backend)
|
||||
return True
|
||||
|
||||
|
||||
def is_request_scoped_cache_dit_enabled(pipeline: object) -> bool:
|
||||
"""Read Cache-DiT state from a pipeline that owns its lifecycle."""
|
||||
|
||||
return isinstance(pipeline, SupportsRequestScopedCacheDiT) and pipeline.is_cache_dit_enabled()
|
||||
@@ -22,10 +22,6 @@ from transformers import Qwen2TokenizerFast, Qwen3VLProcessor
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.diffusion import envs
|
||||
from vllm_omni.diffusion.cache.cachedit import (
|
||||
CacheDiTBackend,
|
||||
RequestScopedCacheDiTRuntime,
|
||||
)
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
get_dit_group,
|
||||
@@ -77,7 +73,6 @@ from .presentation import (
|
||||
minimax_h3_ref2va_video_presentation,
|
||||
minimax_h3_text_only_ids,
|
||||
)
|
||||
from .quality_policy import MINIMAX_H3_GENERIC_CACHE_KEY, MiniMaxH3QualityPolicy
|
||||
from .reference_video import (
|
||||
load_audio_file,
|
||||
load_video_audio,
|
||||
@@ -543,19 +538,6 @@ class MiniMaxH3Pipeline(
|
||||
]
|
||||
dummy_run_num_frames: ClassVar[int] = 0
|
||||
|
||||
def adopt_cache_dit_backend(self, backend: CacheDiTBackend) -> None:
|
||||
"""Adopt runner-installed generic Cache-DiT for request transitions."""
|
||||
|
||||
self._cache_dit_runtime.adopt(
|
||||
backend,
|
||||
installation_key=MINIMAX_H3_GENERIC_CACHE_KEY,
|
||||
)
|
||||
|
||||
def is_cache_dit_enabled(self) -> bool:
|
||||
"""Return the request-scoped Cache-DiT installation state."""
|
||||
|
||||
return self._cache_dit_runtime.is_enabled
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
@@ -694,9 +676,6 @@ class MiniMaxH3Pipeline(
|
||||
# Registry-side VAE patch-parallel discovery uses ``pipeline.vae``.
|
||||
self.vae = self.video_vae
|
||||
|
||||
self._quality_policy = MiniMaxH3QualityPolicy(od_config)
|
||||
self._cache_dit_runtime = RequestScopedCacheDiTRuntime(self)
|
||||
|
||||
self.setup_diffusion_pipeline_profiler(
|
||||
enable_diffusion_pipeline_profiler=(od_config.enable_diffusion_pipeline_profiler)
|
||||
)
|
||||
@@ -1572,8 +1551,6 @@ class MiniMaxH3Pipeline(
|
||||
raise OmniClientError("MiniMax H3 requires a non-empty prompt")
|
||||
|
||||
sampling = request.sampling_params
|
||||
quality = sampling.quality
|
||||
logger.debug("MiniMax H3 request quality=%s", quality)
|
||||
extra = sampling.extra_args or {}
|
||||
task = self._resolve_task(extra.get("task"), multi_modal_data)
|
||||
|
||||
@@ -1755,12 +1732,6 @@ class MiniMaxH3Pipeline(
|
||||
num_steps = int(sampling.num_inference_steps or 50)
|
||||
video_shift = float(extra.get("flow_shift", self.default_video_shift))
|
||||
audio_shift = float(extra.get("audio_flow_shift", self.default_audio_shift))
|
||||
quality_plan = self._quality_policy.resolve(
|
||||
quality=quality,
|
||||
num_inference_steps=num_steps,
|
||||
extra_args=extra,
|
||||
)
|
||||
self._cache_dit_runtime.prepare(quality_plan.cache_dit)
|
||||
num_outputs = _resolve_minimax_h3_num_outputs(sampling.num_outputs_per_prompt)
|
||||
videos = []
|
||||
audios = []
|
||||
|
||||
@@ -1,178 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Model-owned request quality policy for MiniMax H3."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass, replace
|
||||
from numbers import Integral
|
||||
from typing import Any
|
||||
|
||||
from vllm_omni.diffusion.cache.cachedit import CacheDiTRequestSpec
|
||||
from vllm_omni.diffusion.data import DiffusionCacheConfig
|
||||
from vllm_omni.errors import OmniClientError
|
||||
|
||||
MINIMAX_H3_GENERIC_CACHE_KEY = "minimax_h3.generic"
|
||||
MINIMAX_H3_HIGH_CACHE_KEY = "minimax_h3.high"
|
||||
MINIMAX_H3_FORCE_REFRESH_STEP_HINT_ARG = "force_refresh_step_hint"
|
||||
MINIMAX_H3_FORCE_REFRESH_STEP_POLICY_ARG = "force_refresh_step_policy"
|
||||
MINIMAX_H3_FORCE_REFRESH_POLICIES = ("once", "repeat")
|
||||
|
||||
|
||||
def _high_quality_cache_config() -> DiffusionCacheConfig:
|
||||
return DiffusionCacheConfig(
|
||||
Fn_compute_blocks=1,
|
||||
Bn_compute_blocks=0,
|
||||
max_warmup_steps=4,
|
||||
residual_diff_threshold=0.04,
|
||||
max_continuous_cached_steps=1,
|
||||
enable_taylorseer=False,
|
||||
scm_steps_mask_policy=None,
|
||||
)
|
||||
|
||||
|
||||
def _resolve_force_refresh(
|
||||
extra_args: Mapping[str, Any] | None,
|
||||
*,
|
||||
num_inference_steps: int,
|
||||
) -> tuple[int, str] | None:
|
||||
"""Resolve H3's request-scoped Cache-DiT refresh hint.
|
||||
|
||||
These are model-specific ``extra_args`` rather than global sampling fields,
|
||||
keeping other diffusion models unchanged.
|
||||
"""
|
||||
|
||||
if not extra_args:
|
||||
return None
|
||||
|
||||
raw_hint = extra_args.get(MINIMAX_H3_FORCE_REFRESH_STEP_HINT_ARG)
|
||||
raw_policy = extra_args.get(MINIMAX_H3_FORCE_REFRESH_STEP_POLICY_ARG)
|
||||
|
||||
if raw_hint is None:
|
||||
if raw_policy is not None:
|
||||
raise OmniClientError("MiniMax H3 force_refresh_step_policy requires force_refresh_step_hint")
|
||||
return None
|
||||
if isinstance(raw_hint, bool) or not isinstance(raw_hint, Integral):
|
||||
raise OmniClientError("MiniMax H3 force_refresh_step_hint must be a positive integer")
|
||||
|
||||
hint = int(raw_hint)
|
||||
if not 1 <= hint <= num_inference_steps:
|
||||
raise OmniClientError(
|
||||
"MiniMax H3 force_refresh_step_hint must be between 1 and "
|
||||
f"num_inference_steps ({num_inference_steps}), got {hint}"
|
||||
)
|
||||
|
||||
policy = "once" if raw_policy is None else raw_policy
|
||||
if not isinstance(policy, str) or policy not in MINIMAX_H3_FORCE_REFRESH_POLICIES:
|
||||
raise OmniClientError(
|
||||
"MiniMax H3 force_refresh_step_policy must be one of "
|
||||
f"{list(MINIMAX_H3_FORCE_REFRESH_POLICIES)}, got {policy!r}"
|
||||
)
|
||||
return hint, policy
|
||||
|
||||
|
||||
def _has_force_refresh_args(extra_args: Mapping[str, Any] | None) -> bool:
|
||||
if not extra_args:
|
||||
return False
|
||||
return any(
|
||||
extra_args.get(name) is not None
|
||||
for name in (
|
||||
MINIMAX_H3_FORCE_REFRESH_STEP_HINT_ARG,
|
||||
MINIMAX_H3_FORCE_REFRESH_STEP_POLICY_ARG,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _with_force_refresh(
|
||||
cache_config: DiffusionCacheConfig,
|
||||
force_refresh: tuple[int, str] | None,
|
||||
) -> DiffusionCacheConfig:
|
||||
if force_refresh is None:
|
||||
return cache_config
|
||||
hint, policy = force_refresh
|
||||
return replace(
|
||||
cache_config,
|
||||
force_refresh_step_hint=hint,
|
||||
force_refresh_step_policy=policy,
|
||||
)
|
||||
|
||||
|
||||
def _cache_installation_key(base_key: str, force_refresh: tuple[int, str] | None) -> str:
|
||||
"""Make refresh-policy transitions explicit to the request runtime.
|
||||
|
||||
Cache-DiT cannot clear an existing ``force_refresh_step_hint`` through its
|
||||
incremental context update API because ``None`` is treated as "keep the
|
||||
old value". Including the request hint in the installation key therefore
|
||||
makes a hint change perform a safe hook reinstall, while repeated requests
|
||||
with the same hint still only refresh the context.
|
||||
"""
|
||||
|
||||
if force_refresh is None:
|
||||
return base_key
|
||||
hint, policy = force_refresh
|
||||
return f"{base_key}:force_refresh={hint}:{policy}"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MiniMaxH3QualityPlan:
|
||||
"""Resolved execution choices for one MiniMax H3 request."""
|
||||
|
||||
cache_dit: CacheDiTRequestSpec | None
|
||||
|
||||
|
||||
class MiniMaxH3QualityPolicy:
|
||||
"""Resolve H3 request quality into a declarative Cache-DiT target.
|
||||
|
||||
When the server starts with Cache-DiT, omitted quality selects the
|
||||
server-configured profile. Otherwise, omitted quality selects no cache.
|
||||
In either case, ``lossless`` selects no cache and ``high`` selects H3's
|
||||
high-quality profile. The policy therefore owns whether a request installs
|
||||
Cache-DiT; the startup backend only controls the omitted-quality default.
|
||||
H3-specific Cache-DiT refresh hints are read from request ``extra_args``;
|
||||
they do not change the global diffusion request contract.
|
||||
The pipeline owns applying the resulting target at the request boundary.
|
||||
"""
|
||||
|
||||
def __init__(self, od_config: Any) -> None:
|
||||
self._od_config = od_config
|
||||
self._configured_backend = str(getattr(od_config, "cache_backend", "none") or "none").lower()
|
||||
|
||||
def resolve(
|
||||
self,
|
||||
*,
|
||||
quality: str | None,
|
||||
num_inference_steps: int,
|
||||
extra_args: Mapping[str, Any] | None = None,
|
||||
) -> MiniMaxH3QualityPlan:
|
||||
if quality == "high":
|
||||
base_key = MINIMAX_H3_HIGH_CACHE_KEY
|
||||
base_config = _high_quality_cache_config()
|
||||
elif self._configured_backend == "cache_dit" and quality is None:
|
||||
base_key = MINIMAX_H3_GENERIC_CACHE_KEY
|
||||
base_config = self._od_config.cache_config
|
||||
else:
|
||||
if _has_force_refresh_args(extra_args):
|
||||
raise OmniClientError("MiniMax H3 force-refresh arguments require an active Cache-DiT request target")
|
||||
return MiniMaxH3QualityPlan(cache_dit=None)
|
||||
|
||||
force_refresh = _resolve_force_refresh(
|
||||
extra_args,
|
||||
num_inference_steps=num_inference_steps,
|
||||
)
|
||||
return MiniMaxH3QualityPlan(
|
||||
cache_dit=CacheDiTRequestSpec(
|
||||
installation_key=_cache_installation_key(base_key, force_refresh),
|
||||
cache_config=_with_force_refresh(base_config, force_refresh),
|
||||
num_inference_steps=num_inference_steps,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MINIMAX_H3_GENERIC_CACHE_KEY",
|
||||
"MINIMAX_H3_HIGH_CACHE_KEY",
|
||||
"MiniMaxH3QualityPlan",
|
||||
"MiniMaxH3QualityPolicy",
|
||||
]
|
||||
@@ -55,10 +55,6 @@ class StepBatchSamplingParamsKey:
|
||||
true_cfg_scale: float | None = None
|
||||
cfg_normalize: bool = False
|
||||
|
||||
# Request-scoped execution policy. A batch must use one quality mode
|
||||
# because model acceleration hooks are shared by the whole worker batch.
|
||||
quality: str | None = None
|
||||
|
||||
# Output count. Requests with different num_outputs_per_prompt produce
|
||||
# differently shaped outputs and cannot share a batch.
|
||||
num_outputs_per_prompt: int = 1
|
||||
@@ -112,11 +108,6 @@ class RequestBatchSamplingParamsKey:
|
||||
layers: int = 4
|
||||
use_en_prompt: bool = False
|
||||
|
||||
# Request-scoped execution policy. Keep accelerated and reference-path
|
||||
# requests in separate worker batches. ``None`` delegates the default to
|
||||
# the model and must remain distinct from explicit ``lossless``.
|
||||
quality: str | None = None
|
||||
|
||||
# LoRA identity.
|
||||
lora_int_id: int | None = None
|
||||
lora_scale: float = 1.0
|
||||
|
||||
@@ -24,21 +24,17 @@ _REQUEST_BATCH_SAMPLING_PARAMS_KEY_FIELD_NAMES = frozenset(
|
||||
) - {"lora_int_id"}
|
||||
|
||||
|
||||
def build_request_batch_sampling_params_key(request: OmniDiffusionRequest) -> RequestBatchSamplingParamsKey:
|
||||
"""Build the compatibility key shared by scheduling and DP dispatch."""
|
||||
sampling = request.sampling_params
|
||||
# LoRA identity is optional on sampling params (and on test stubs).
|
||||
lora_request = getattr(sampling, "lora_request", None)
|
||||
key_kwargs = {name: getattr(sampling, name) for name in _REQUEST_BATCH_SAMPLING_PARAMS_KEY_FIELD_NAMES}
|
||||
key_kwargs["lora_int_id"] = lora_request.lora_int_id if lora_request is not None else None
|
||||
return RequestBatchSamplingParamsKey(**key_kwargs)
|
||||
|
||||
|
||||
class RequestScheduler(BaseScheduler):
|
||||
"""Diffusion scheduler with vLLM-style waiting/running queues."""
|
||||
|
||||
def _build_sampling_params_key(self, request: OmniDiffusionRequest) -> RequestBatchSamplingParamsKey:
|
||||
return build_request_batch_sampling_params_key(request)
|
||||
"""Build a request-batch compatibility key from sampling parameters."""
|
||||
sampling = request.sampling_params
|
||||
# LoRA identity is optional on sampling params (and on test stubs).
|
||||
lora_request = getattr(sampling, "lora_request", None)
|
||||
key_kwargs = {name: getattr(sampling, name) for name in _REQUEST_BATCH_SAMPLING_PARAMS_KEY_FIELD_NAMES}
|
||||
key_kwargs["lora_int_id"] = lora_request.lora_int_id if lora_request is not None else None
|
||||
return RequestBatchSamplingParamsKey(**key_kwargs)
|
||||
|
||||
def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
|
||||
scheduled_request_ids = sched_output.scheduled_request_ids
|
||||
|
||||
@@ -32,13 +32,7 @@ from vllm_omni.diffusion.compile import regionally_compile
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.forward_context import set_forward_context
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
from vllm_omni.diffusion.models.interface import (
|
||||
SupportsPromptUpdate,
|
||||
adopt_request_scoped_cache_dit,
|
||||
is_request_scoped_cache_dit_enabled,
|
||||
supports_prompt_update,
|
||||
supports_step_execution,
|
||||
)
|
||||
from vllm_omni.diffusion.models.interface import SupportsPromptUpdate, supports_prompt_update, supports_step_execution
|
||||
from vllm_omni.diffusion.offloader import get_offload_backend
|
||||
from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
@@ -310,19 +304,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
|
||||
self.cache_backend = None
|
||||
self.od_config.cache_backend = None
|
||||
else:
|
||||
# Install configured cache capability once at startup. A model
|
||||
# may explicitly adopt the enabled Cache-DiT backend and then
|
||||
# own all later request-boundary enable/disable transitions.
|
||||
self.cache_backend.enable(self.pipeline)
|
||||
if str(self.od_config.cache_backend).lower() == "cache_dit" and adopt_request_scoped_cache_dit(
|
||||
self.pipeline,
|
||||
self.cache_backend,
|
||||
):
|
||||
logger.info(
|
||||
"Pipeline %s owns request-scoped Cache-DiT transitions.",
|
||||
type(self.pipeline).__name__,
|
||||
)
|
||||
self.cache_backend = None
|
||||
|
||||
# Install prompt-embedding cache (transparent wrapper around
|
||||
# ``pipeline.encode_prompt``). Enabled via config or env var; a no-op
|
||||
@@ -541,11 +523,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
|
||||
if is_primary and prompt_embed_cache is not None:
|
||||
logger.debug("prompt-embed cache: %s", prompt_embed_cache.stats())
|
||||
|
||||
runner_cache_dit_enabled = self.cache_backend is not None and self.cache_backend.is_enabled()
|
||||
if (
|
||||
od_config.cache_backend == "cache_dit"
|
||||
self.cache_backend is not None
|
||||
and self.cache_backend.is_enabled()
|
||||
and od_config.cache_backend == "cache_dit"
|
||||
and od_config.enable_cache_dit_summary
|
||||
and (runner_cache_dit_enabled or is_request_scoped_cache_dit_enabled(self.pipeline))
|
||||
):
|
||||
cache_summary(self.pipeline, details=True)
|
||||
|
||||
|
||||
@@ -63,8 +63,6 @@ from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
_ASYNC_OUTPUT_THREAD_JOIN_TIMEOUT_S = 10.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class _DiffusionVllmModelConfig:
|
||||
@@ -431,8 +429,8 @@ class DiffusionWorker:
|
||||
ranks stay synchronised at each AllGather call while computing
|
||||
different activations.
|
||||
|
||||
Each rank returns its OWN DiffusionOutput (no gather). The executor
|
||||
collects N responses via the per-worker result queues.
|
||||
Each rank returns its OWN DiffusionOutput (no gather). The executor
|
||||
collects N responses via result_mq (all ranks share one queue).
|
||||
"""
|
||||
assert self.model_runner is not None, "Model runner not initialized"
|
||||
|
||||
@@ -897,11 +895,8 @@ class WorkerProc:
|
||||
"""
|
||||
device = torch.device(torch.accelerator.current_accelerator().type, self.gpu_id)
|
||||
d2h_stream = torch.Stream(device=device)
|
||||
while True:
|
||||
item = self._async_output_queue.get()
|
||||
if item is None:
|
||||
break
|
||||
output, async_output_id, gpu_event = item
|
||||
while self._running:
|
||||
output, async_output_id, gpu_event = self._async_output_queue.get()
|
||||
try:
|
||||
# Cross-stream ordering: wait for default stream to finish
|
||||
# writing the output tensors before the side stream reads.
|
||||
@@ -930,40 +925,6 @@ class WorkerProc:
|
||||
)
|
||||
)
|
||||
|
||||
def shutdown(self) -> None:
|
||||
"""Stop background work and release worker-owned IPC resources."""
|
||||
self._running = False
|
||||
|
||||
if self._async_output_queue is not None:
|
||||
self._async_output_queue.put(None)
|
||||
if self._async_output_thread is not None:
|
||||
self._async_output_thread.join(timeout=_ASYNC_OUTPUT_THREAD_JOIN_TIMEOUT_S)
|
||||
if self._async_output_thread.is_alive():
|
||||
logger.warning(
|
||||
"Worker %d: Async output thread did not stop before shutdown",
|
||||
self.gpu_id,
|
||||
)
|
||||
|
||||
try:
|
||||
self.worker.shutdown()
|
||||
finally:
|
||||
if self.mq is not None:
|
||||
self.mq.shutdown()
|
||||
self.mq = None
|
||||
|
||||
# The worker creates this queue's shared-memory ring buffer.
|
||||
# Dropping the final creator reference invokes
|
||||
# ShmRingBuffer.__del__, which closes and unlinks it before
|
||||
# multiprocessing.resource_tracker exits.
|
||||
if self.result_mq is not None and (
|
||||
self._async_output_thread is None or not self._async_output_thread.is_alive()
|
||||
):
|
||||
result_mq = self.result_mq
|
||||
self.result_mq = None
|
||||
result_mq.shutdown()
|
||||
del result_mq
|
||||
gc.collect()
|
||||
|
||||
def _gather_rpc_rank_statuses(self, status: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
if not torch.distributed.is_initialized():
|
||||
return [status]
|
||||
@@ -1235,7 +1196,7 @@ class WorkerProc:
|
||||
finally:
|
||||
if worker_proc is not None:
|
||||
try:
|
||||
worker_proc.shutdown()
|
||||
worker_proc.worker.shutdown()
|
||||
except Exception as exc:
|
||||
logger.warning("Worker %d: Shutdown encountered an error: %s", rank, exc)
|
||||
worker_proc.context.term()
|
||||
|
||||
@@ -188,7 +188,6 @@ class OmniEngineArgs(EngineArgs):
|
||||
omni: bool = False
|
||||
# Diffusion request-mode batch admission (forwarded to OmniDiffusionConfig).
|
||||
request_batch_max_wait_ms: float = 0.0
|
||||
fa_deterministic: bool = False
|
||||
|
||||
@classmethod
|
||||
def _add_omni_specific_args(cls, parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
|
||||
|
||||
@@ -1043,7 +1043,6 @@ class AsyncOmniEngine:
|
||||
"diffusion_compile_dynamic": (
|
||||
True if kwargs.get("diffusion_compile_dynamic") is None else kwargs["diffusion_compile_dynamic"]
|
||||
),
|
||||
"fa_deterministic": bool(kwargs.get("fa_deterministic", False)),
|
||||
"boundary_ratio": kwargs.get("boundary_ratio", None),
|
||||
"flow_shift": kwargs.get("flow_shift", None),
|
||||
"diffusion_load_format": kwargs.get("diffusion_load_format", "default"),
|
||||
|
||||
@@ -434,17 +434,6 @@ class OmniServeCommand(CLISubcommand):
|
||||
"Disable for fixed-shape workloads with --no-diffusion-compile-dynamic."
|
||||
),
|
||||
)
|
||||
omni_config_group.add_argument(
|
||||
"--fa-deterministic",
|
||||
dest="fa_deterministic",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help=(
|
||||
"Request FlashAttention deterministic=True on the local FLASH_ATTN dense path. "
|
||||
"Slower than the library default deterministic=False; intended for accuracy CI. "
|
||||
"Serving default remains non-deterministic."
|
||||
),
|
||||
)
|
||||
omni_config_group.add_argument(
|
||||
"--diffusers-load-kwargs",
|
||||
dest="diffusers_load_kwargs",
|
||||
|
||||
@@ -39,7 +39,6 @@ from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, Tool
|
||||
from vllm.entrypoints.openai.api_server import build_app as build_openai_app
|
||||
from vllm.entrypoints.openai.api_server import setup_server as setup_openai_server
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import (
|
||||
BatchChatCompletionRequest,
|
||||
ChatCompletionRequest,
|
||||
ChatCompletionResponse,
|
||||
)
|
||||
@@ -96,7 +95,6 @@ from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
|
||||
from vllm_omni.config.endpoint_policy import shutdown_unsupported_routes
|
||||
from vllm_omni.diffusion.models.interface import ReferenceVideoDecodeSpec
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
from vllm_omni.entrypoints.openai.batch_serving import OmniOpenAIServingChatBatch
|
||||
from vllm_omni.entrypoints.openai.duplex_capability import should_enable_duplex_endpoint
|
||||
from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
|
||||
from vllm_omni.entrypoints.openai.image_api_utils import (
|
||||
@@ -515,7 +513,6 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
|
||||
|
||||
# OMNI: Remove upstream routes that we override with omni-specific handlers
|
||||
_remove_route_from_app(app, "/v1/chat/completions", {"POST"})
|
||||
_remove_route_from_app(app, "/v1/chat/completions/batch", {"POST"})
|
||||
_remove_route_from_app(app, "/v1/models", {"GET"}) # Remove upstream /v1/models to use omni's handler
|
||||
app.include_router(router)
|
||||
|
||||
@@ -780,10 +777,6 @@ async def omni_init_app_state(
|
||||
diffusion_engine=engine_client, # type: ignore
|
||||
model_name=model_name,
|
||||
)
|
||||
state.openai_serving_chat_batch = OmniOpenAIServingChatBatch.for_diffusion(
|
||||
diffusion_engine=engine_client, # type: ignore
|
||||
model_name=model_name,
|
||||
)
|
||||
|
||||
# audio related
|
||||
state.openai_serving_speech = None
|
||||
@@ -957,33 +950,30 @@ async def omni_init_app_state(
|
||||
if "generate" in supported_tasks
|
||||
else None
|
||||
)
|
||||
|
||||
_chat_kwargs = dict(
|
||||
engine_client=engine_client,
|
||||
models=state.openai_serving_models,
|
||||
response_role=args.response_role,
|
||||
online_renderer=state.online_renderer,
|
||||
request_logger=request_logger,
|
||||
chat_template=resolved_chat_template,
|
||||
chat_template_content_format=args.chat_template_content_format,
|
||||
default_chat_template_kwargs=args.default_chat_template_kwargs,
|
||||
trust_request_chat_template=args.trust_request_chat_template,
|
||||
return_tokens_as_token_ids=args.return_tokens_as_token_ids,
|
||||
enable_auto_tools=args.enable_auto_tool_choice,
|
||||
exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,
|
||||
tool_parser=args.tool_call_parser,
|
||||
reasoning_parser=args.structured_outputs_config.reasoning_parser,
|
||||
enable_prompt_tokens_details=args.enable_prompt_tokens_details,
|
||||
enable_force_include_usage=args.enable_force_include_usage,
|
||||
enable_log_outputs=args.enable_log_outputs,
|
||||
enable_log_deltas=args.enable_log_deltas,
|
||||
state.openai_serving_chat = (
|
||||
OmniOpenAIServingChat(
|
||||
engine_client,
|
||||
state.openai_serving_models,
|
||||
args.response_role,
|
||||
online_renderer=state.online_renderer,
|
||||
request_logger=request_logger,
|
||||
chat_template=resolved_chat_template,
|
||||
chat_template_content_format=args.chat_template_content_format,
|
||||
default_chat_template_kwargs=args.default_chat_template_kwargs,
|
||||
trust_request_chat_template=args.trust_request_chat_template,
|
||||
return_tokens_as_token_ids=args.return_tokens_as_token_ids,
|
||||
enable_auto_tools=args.enable_auto_tool_choice,
|
||||
exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,
|
||||
tool_parser=args.tool_call_parser,
|
||||
reasoning_parser=args.structured_outputs_config.reasoning_parser,
|
||||
enable_prompt_tokens_details=args.enable_prompt_tokens_details,
|
||||
enable_force_include_usage=args.enable_force_include_usage,
|
||||
enable_log_outputs=args.enable_log_outputs,
|
||||
enable_log_deltas=args.enable_log_deltas,
|
||||
)
|
||||
if "generate" in supported_tasks
|
||||
else None
|
||||
)
|
||||
|
||||
state.openai_serving_chat = OmniOpenAIServingChat(**_chat_kwargs) if "generate" in supported_tasks else None
|
||||
state.openai_serving_chat_batch = (
|
||||
OmniOpenAIServingChatBatch(**_chat_kwargs) if "generate" in supported_tasks else None
|
||||
)
|
||||
|
||||
# Warm up chat template processing to avoid first-request latency
|
||||
if state.openai_serving_chat is not None:
|
||||
state.openai_serving_chat.warmup()
|
||||
@@ -1172,10 +1162,6 @@ def Omnichat(request: Request) -> OmniOpenAIServingChat | None:
|
||||
return request.app.state.openai_serving_chat
|
||||
|
||||
|
||||
def OmniBatchChat(request: Request) -> OmniOpenAIServingChatBatch | None:
|
||||
return request.app.state.openai_serving_chat_batch
|
||||
|
||||
|
||||
def Omnispeech(request: Request) -> OmniOpenAIServingSpeech | None:
|
||||
return request.app.state.openai_serving_speech
|
||||
|
||||
@@ -1259,45 +1245,6 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re
|
||||
return StreamingResponse(content=generator, media_type="text/event-stream")
|
||||
|
||||
|
||||
@router.post(
|
||||
"/v1/chat/completions/batch",
|
||||
dependencies=[Depends(validate_json_request)],
|
||||
responses={
|
||||
HTTPStatus.OK.value: {},
|
||||
HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
|
||||
HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
|
||||
HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
|
||||
HTTPStatus.NOT_IMPLEMENTED.value: {"model": ErrorResponse},
|
||||
},
|
||||
)
|
||||
@with_cancellation
|
||||
@load_aware_call
|
||||
async def create_batch_chat_completion(request: BatchChatCompletionRequest, raw_request: Request):
|
||||
handler = OmniBatchChat(raw_request)
|
||||
if handler is None:
|
||||
base_server = getattr(raw_request.app.state, "serving_tokenization", None)
|
||||
if base_server is None:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.NOT_FOUND.value,
|
||||
detail="The model does not support Chat Completions API",
|
||||
)
|
||||
return base_server.create_error_response(message="The model does not support Chat Completions API")
|
||||
try:
|
||||
result = await handler.create_batch_chat_completion(request, raw_request)
|
||||
except (EngineGenerateError, EngineDeadError) as exc:
|
||||
return _create_engine_error_json_response(raw_request, exc)
|
||||
except Exception as e:
|
||||
logger.exception("Batched chat completion failed: %s", e)
|
||||
raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)) from e
|
||||
|
||||
if isinstance(result, ErrorResponse):
|
||||
return JSONResponse(
|
||||
content=result.model_dump(),
|
||||
status_code=result.error.code if result.error else 400,
|
||||
)
|
||||
return JSONResponse(content=result.model_dump(mode="json"))
|
||||
|
||||
|
||||
_remove_route_from_router(router, "/v1/audio/speech", {"POST"})
|
||||
|
||||
|
||||
@@ -2864,7 +2811,6 @@ def video_response_from_request(model_name: str, req: VideoGenerationRequest) ->
|
||||
status=VideoGenerationStatus.QUEUED,
|
||||
size=req.size,
|
||||
prompt=req.prompt,
|
||||
quality=req.quality or "default",
|
||||
)
|
||||
resp.seconds = str(req.seconds or resp.seconds)
|
||||
return resp
|
||||
@@ -3204,7 +3150,6 @@ async def _parse_video_form(
|
||||
short_edge: int | None = Form(default=None, ge=1),
|
||||
num_outputs_per_prompt: int = Form(default=1, ge=1, le=10),
|
||||
start_time_seconds: float | None = Form(default=None, ge=0.0),
|
||||
quality: str | None = Form(default=None),
|
||||
num_inference_steps: int | None = Form(default=None),
|
||||
guidance_scale: float | None = Form(default=None),
|
||||
guidance_scale_2: float | None = Form(default=None),
|
||||
@@ -3273,7 +3218,6 @@ async def _parse_video_form(
|
||||
"short_edge": short_edge,
|
||||
"num_outputs_per_prompt": num_outputs_per_prompt,
|
||||
"start_time_seconds": start_time_seconds,
|
||||
"quality": quality,
|
||||
"num_inference_steps": num_inference_steps,
|
||||
"guidance_scale": guidance_scale,
|
||||
"guidance_scale_2": guidance_scale_2,
|
||||
|
||||
@@ -1,151 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import asyncio
|
||||
import time
|
||||
|
||||
from fastapi import Request
|
||||
from pydantic import ValidationError
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import (
|
||||
BatchChatCompletionRequest,
|
||||
ChatCompletionRequest,
|
||||
ChatCompletionResponse,
|
||||
ChatCompletionResponseChoice,
|
||||
)
|
||||
from vllm.entrypoints.openai.engine.protocol import (
|
||||
ErrorResponse,
|
||||
UsageInfo,
|
||||
)
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class OmniOpenAIServingChatBatch(OmniOpenAIServingChat):
|
||||
@staticmethod
|
||||
def _maybe_collapse_choices(choices: list[ChatCompletionResponseChoice]):
|
||||
"""Collapse content + audio choices into one choice.
|
||||
|
||||
NOTE: This is a hack for models that produce separate text and
|
||||
audio choices to ensure we have 1:1 mapping between number of
|
||||
inputs in the batch and responses.
|
||||
"""
|
||||
content_choice, audio_choice = None, None
|
||||
if len(choices) == 1:
|
||||
return choices[0]
|
||||
if len(choices) != 2:
|
||||
raise ValueError(f"Expected 1 or 2 choices to collapse, got {len(choices)}. ")
|
||||
|
||||
for choice in choices:
|
||||
if choice.message.audio:
|
||||
if audio_choice is not None:
|
||||
raise ValueError("Multiple audio choices cannot be set in one completion message")
|
||||
audio_choice = choice
|
||||
else:
|
||||
if content_choice is not None:
|
||||
raise ValueError("Multiple content choices cannot be set in one completion message")
|
||||
content_choice = choice
|
||||
|
||||
if content_choice is None or audio_choice is None:
|
||||
raise ValueError(
|
||||
"Expected one content and one audio choice, but got: "
|
||||
f"content={content_choice is not None}, audio={audio_choice is not None}"
|
||||
)
|
||||
content_choice.message.audio = audio_choice.message.audio
|
||||
return content_choice
|
||||
|
||||
@staticmethod
|
||||
def _get_subrequest_ids(base_id, request: BatchChatCompletionRequest) -> list[str]:
|
||||
"""Get the request ID for each entry in the batch request to be resubmitted to chat completions."""
|
||||
return [f"{base_id}-idx-{idx}" for idx in range(len(request.messages))]
|
||||
|
||||
async def create_batch_chat_completion(
|
||||
self,
|
||||
request: BatchChatCompletionRequest,
|
||||
raw_request: Request,
|
||||
) -> ChatCompletionResponse | ErrorResponse:
|
||||
"""Given a request, submit each request to chat completions & collect the results."""
|
||||
model = ""
|
||||
enabled_streaming = False
|
||||
chat_requests: list[ChatCompletionRequest] = []
|
||||
choices: list[ChatCompletionResponseChoice] = []
|
||||
total_prompt_tokens = 0
|
||||
total_completion_tokens = 0
|
||||
|
||||
# Get the base request from the raw_request before maybe mutating the header
|
||||
base_id = OmniOpenAIServingChatBatch._base_request_id(
|
||||
raw_request,
|
||||
request.request_id,
|
||||
)
|
||||
batch_req_id = f"chatcmpl-batch-{base_id}"
|
||||
sub_req_ids = self._get_subrequest_ids(base_id, request)
|
||||
|
||||
# Remove the raw request header if it exists now that we have subreq IDs,
|
||||
# since this takes priority over the object request_ids
|
||||
if raw_request.headers.get("X-Request-Id"):
|
||||
copy_headers = raw_request.headers.mutablecopy()
|
||||
del copy_headers["X-Request-Id"]
|
||||
raw_request._headers = copy_headers
|
||||
|
||||
for idx, msg in enumerate(request.messages):
|
||||
# Update the current request ID to avoid subrequest collisions
|
||||
request.request_id = sub_req_ids[idx]
|
||||
try:
|
||||
chat_cmp_request = request.to_chat_completion_request(msg)
|
||||
except ValidationError as e:
|
||||
return self._create_error_response(
|
||||
f"Message {idx} could not be converted to a chat completion request: {e}",
|
||||
)
|
||||
# Streaming isn't supported for batch chat completions,
|
||||
# so always set to False & warn if it was requested.
|
||||
enabled_streaming |= chat_cmp_request.stream
|
||||
chat_cmp_request.stream = False
|
||||
chat_requests.append(chat_cmp_request)
|
||||
|
||||
if enabled_streaming:
|
||||
logger.warning("Streaming is not supported for batched chat completions; ignoring stream=True.")
|
||||
|
||||
# Submit each chat completion request as a task, then gather results.
|
||||
# TODO (Alex): optimize this
|
||||
tasks = [asyncio.create_task(self.create_chat_completion(c, raw_request)) for c in chat_requests]
|
||||
try:
|
||||
results = await asyncio.gather(*tasks)
|
||||
finally:
|
||||
# Ensure we cancel remaining tasks if needed, e.g., early exit due to bad behavior
|
||||
for t in tasks:
|
||||
if not t.done():
|
||||
t.cancel()
|
||||
|
||||
for i, resp in enumerate(results):
|
||||
if isinstance(resp, ErrorResponse):
|
||||
return resp
|
||||
completion: ChatCompletionResponse = resp
|
||||
model = completion.model
|
||||
# FIXME (Alex): We should probably handle this in chat completions,
|
||||
# not here, but we need to ensure streaming is properly handled
|
||||
try:
|
||||
collapsed_choice = self._maybe_collapse_choices(completion.choices)
|
||||
except ValueError as e:
|
||||
return self._create_error_response(
|
||||
f"Failed to collapse choices with error: {e}",
|
||||
)
|
||||
collapsed_choice.index = i
|
||||
choices.append(collapsed_choice)
|
||||
if completion.usage:
|
||||
total_prompt_tokens += completion.usage.prompt_tokens
|
||||
total_completion_tokens += completion.usage.completion_tokens
|
||||
|
||||
usage = UsageInfo(
|
||||
prompt_tokens=total_prompt_tokens,
|
||||
completion_tokens=total_completion_tokens,
|
||||
total_tokens=total_prompt_tokens + total_completion_tokens,
|
||||
)
|
||||
|
||||
return ChatCompletionResponse(
|
||||
id=batch_req_id,
|
||||
created=int(time.time()),
|
||||
model=model,
|
||||
choices=choices,
|
||||
usage=usage,
|
||||
)
|
||||
@@ -7,10 +7,7 @@ from collections.abc import Collection, Mapping
|
||||
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.inputs.data import (
|
||||
DIFFUSION_QUALITY_LEVELS,
|
||||
OmniDiffusionSamplingParams,
|
||||
)
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
@@ -93,9 +90,6 @@ def normalize_diffusion_request_args(
|
||||
for alias, canonical_name in aliases.items():
|
||||
if alias in request_args:
|
||||
request_args[canonical_name] = request_args.pop(alias)
|
||||
quality = request_args.get("quality")
|
||||
if quality is not None and quality not in DIFFUSION_QUALITY_LEVELS:
|
||||
raise ValueError(f"quality must be one of {list(DIFFUSION_QUALITY_LEVELS)}, got {quality!r}")
|
||||
return normalized_extra_args, request_args
|
||||
|
||||
|
||||
|
||||
@@ -15,10 +15,9 @@ from enum import Enum
|
||||
from functools import lru_cache
|
||||
from typing import Annotated, Any, Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, StringConstraints, field_validator
|
||||
from pydantic import BaseModel, ConfigDict, Field, StringConstraints
|
||||
|
||||
from vllm_omni.entrypoints.openai.image_api_utils import parse_size
|
||||
from vllm_omni.inputs.data import DIFFUSION_QUALITY_LEVELS
|
||||
|
||||
|
||||
class VideoGenerationStatus(str, Enum):
|
||||
@@ -165,24 +164,6 @@ class VideoGenerationRequest(BaseModel):
|
||||
)
|
||||
|
||||
# vllm-omni extensions for diffusion control
|
||||
quality: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Request-level generation quality intent. Supported values are "
|
||||
"'lossless' and 'high'; exact behavior is model-specific. "
|
||||
"When omitted, the model chooses its default policy."
|
||||
),
|
||||
)
|
||||
|
||||
@field_validator("quality")
|
||||
@classmethod
|
||||
def validate_quality(cls, value: str | None) -> str | None:
|
||||
if value is None:
|
||||
return None
|
||||
if value not in DIFFUSION_QUALITY_LEVELS:
|
||||
raise ValueError(f"quality must be one of {list(DIFFUSION_QUALITY_LEVELS)}, got {value!r}")
|
||||
return value
|
||||
|
||||
negative_prompt: str | None = Field(default=None, description="Text describing what to avoid in the video")
|
||||
num_inference_steps: int | None = Field(
|
||||
default=None,
|
||||
|
||||
@@ -170,7 +170,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
"width",
|
||||
"num_outputs_per_prompt",
|
||||
"seed",
|
||||
"quality",
|
||||
"num_inference_steps",
|
||||
"guidance_scale",
|
||||
"true_cfg_scale",
|
||||
@@ -760,11 +759,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
if hasattr(sp, "num_inference_steps") and num_inference_steps is not None:
|
||||
sp.num_inference_steps = num_inference_steps
|
||||
if isinstance(sp, OmniDiffusionSamplingParams):
|
||||
quality = extra_body.get("quality")
|
||||
self._set_if_supported(
|
||||
sp,
|
||||
quality=quality,
|
||||
)
|
||||
apply_normalized_diffusion_request_extra_args(sp, normalized_extra_args)
|
||||
else:
|
||||
apply_declared_extra_args(
|
||||
@@ -2943,7 +2937,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
generator_device = gen_params.generator_device
|
||||
num_outputs_per_prompt = gen_params.num_outputs_per_prompt
|
||||
num_inference_steps = extra_body.get("num_inference_steps")
|
||||
quality = extra_body.get("quality")
|
||||
guidance_scale = extra_body.get("guidance_scale")
|
||||
true_cfg_scale = extra_body.get("true_cfg_scale") or extra_body.get("cfg_scale")
|
||||
negative_prompt = extra_body.get("negative_prompt")
|
||||
@@ -3094,7 +3087,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
seed=seed,
|
||||
generator_device=generator_device,
|
||||
num_outputs_per_prompt=num_outputs_per_prompt,
|
||||
quality=quality,
|
||||
num_inference_steps=num_inference_steps,
|
||||
guidance_scale=guidance_scale,
|
||||
true_cfg_scale=true_cfg_scale,
|
||||
@@ -3159,7 +3151,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
self._set_if_supported(
|
||||
gen_params,
|
||||
generator_device=generator_device,
|
||||
quality=extra_body.get("quality"),
|
||||
num_inference_steps=extra_body.get("num_inference_steps"),
|
||||
guidance_scale=extra_body.get("guidance_scale"),
|
||||
true_cfg_scale=extra_body.get("true_cfg_scale") or extra_body.get("cfg_scale"),
|
||||
@@ -3498,7 +3489,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
# method apply its own model-specific default when the user does
|
||||
# not provide a value.
|
||||
num_inference_steps = extra_body.get("num_inference_steps")
|
||||
quality = extra_body.get("quality")
|
||||
guidance_scale = extra_body.get("guidance_scale")
|
||||
true_cfg_scale = extra_body.get("true_cfg_scale")
|
||||
seed = extra_body.get("seed")
|
||||
@@ -3553,8 +3543,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
# Only override defaults when the user explicitly provides values
|
||||
if num_inference_steps is not None:
|
||||
gen_params.num_inference_steps = num_inference_steps
|
||||
if quality is not None:
|
||||
gen_params.quality = quality
|
||||
if guidance_scale is not None:
|
||||
gen_params.guidance_scale = guidance_scale
|
||||
if true_cfg_scale is not None:
|
||||
|
||||
@@ -200,8 +200,6 @@ class OmniOpenAIServingVideo:
|
||||
|
||||
if "num_inference_steps" in provided_fields and request.num_inference_steps is not None:
|
||||
gen_params.num_inference_steps = request.num_inference_steps
|
||||
if "quality" in provided_fields:
|
||||
gen_params.quality = request.quality
|
||||
if "guidance_scale" in provided_fields and request.guidance_scale is not None:
|
||||
gen_params.guidance_scale = request.guidance_scale
|
||||
if "guidance_scale_2" in provided_fields and request.guidance_scale_2 is not None:
|
||||
|
||||
@@ -535,8 +535,6 @@ class OmniStreamingVideoOutputHandler:
|
||||
provided_fields = request.model_fields_set
|
||||
if "num_inference_steps" in provided_fields and request.num_inference_steps is not None:
|
||||
gen_params.num_inference_steps = request.num_inference_steps
|
||||
if "quality" in provided_fields:
|
||||
gen_params.quality = request.quality
|
||||
if "guidance_scale" in provided_fields and request.guidance_scale is not None:
|
||||
gen_params.guidance_scale = request.guidance_scale
|
||||
if "guidance_scale_2" in provided_fields and request.guidance_scale_2 is not None:
|
||||
|
||||
@@ -11,8 +11,6 @@ from vllm.sampling_params import SamplingParams
|
||||
|
||||
from vllm_omni.lora.request import LoRARequest
|
||||
|
||||
DIFFUSION_QUALITY_LEVELS: tuple[str, ...] = ("lossless", "high")
|
||||
|
||||
|
||||
class OmniTextPrompt(TextPrompt):
|
||||
"""Text prompt with optional embeddings and stage payloads.
|
||||
@@ -210,10 +208,6 @@ class OmniDiffusionSamplingParams:
|
||||
do_classifier_free_guidance: bool = False
|
||||
output_type: str | None = None
|
||||
|
||||
# Request-scoped quality intent. ``None`` delegates the default behavior
|
||||
# to the model; explicit levels select a model-owned quality policy.
|
||||
quality: str | None = None
|
||||
|
||||
# Batch info
|
||||
num_outputs_per_prompt: int = 1
|
||||
seed: int | None = None
|
||||
@@ -343,10 +337,6 @@ class OmniDiffusionSamplingParams:
|
||||
# results
|
||||
output: torch.Tensor | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.quality is not None and self.quality not in DIFFUSION_QUALITY_LEVELS:
|
||||
raise ValueError(f"quality must be one of {list(DIFFUSION_QUALITY_LEVELS)}, got {self.quality!r}")
|
||||
|
||||
@property
|
||||
def batch_size(self):
|
||||
# This class is changed to only represent a single prompt request
|
||||
|
||||
Reference in new issue
Block a user