Compare commits
17
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
baa56e48d4 | ||
|
|
49f8f5cffb | ||
|
|
50dcc5086b | ||
|
|
65ca13102d | ||
|
|
7aaec2d7f4 | ||
|
|
0548674d30 | ||
|
|
cc716e139a | ||
|
|
4128d63c51 | ||
|
|
3e4f7eb9ba | ||
|
|
3f699a7b4d | ||
|
|
f995211da0 | ||
|
|
cafe52e8ec | ||
|
|
71f9f9a9a7 | ||
|
|
9a7f7b7de5 | ||
|
|
75a84c8127 | ||
|
|
c45cdea512 | ||
|
|
811b5f919d |
No files matched your search
@@ -19,7 +19,7 @@ steps:
|
||||
timeout_in_minutes: 15
|
||||
depends_on: image-build
|
||||
commands:
|
||||
- pytest -s -v tests/single_stage/test_diffusion_model.py
|
||||
- pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
|
||||
plugins:
|
||||
@@ -38,9 +38,9 @@ steps:
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/multi_stages/
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
@@ -56,7 +56,7 @@ steps:
|
||||
depends_on: image-build
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/multi_stages_h100/
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
|
||||
@@ -179,6 +179,7 @@ checkpoints/
|
||||
|
||||
# Cache directories
|
||||
cache/
|
||||
!vllm_omni/diffusion/cache/
|
||||
.cache/
|
||||
diffusion_cache/
|
||||
kv_cache/
|
||||
@@ -229,3 +230,5 @@ configs/development.yaml
|
||||
# Docker
|
||||
.dockerignore
|
||||
Dockerfile.dev
|
||||
discussion
|
||||
tmp_test
|
||||
+10
-6
@@ -7,17 +7,21 @@ nav:
|
||||
- Examples:
|
||||
- examples/README.md
|
||||
- Offline Inference:
|
||||
- Offline Example of vLLM-Omni for Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
|
||||
- Offline Example of vLLM-Omni for Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
|
||||
- Qwen-Image Offline Inference: user_guide/examples/offline_inference/qwen_image.md
|
||||
- Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
|
||||
- Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
|
||||
- Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
|
||||
- Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
|
||||
- Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
|
||||
- Online Serving:
|
||||
- Online serving Example of vLLM-Omni for Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
|
||||
- Online serving Example of vLLM-Omni for Qwen3-omni: user_guide/examples/online_serving/qwen3_omni.md
|
||||
- Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
|
||||
- Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
|
||||
- General:
|
||||
- usage/*
|
||||
- Configuration:
|
||||
- configuration/README.md
|
||||
- configuration/*
|
||||
- Diffusion Acceleration:
|
||||
- TeaCache: user_guide/teacache.md
|
||||
- Models:
|
||||
- models/supported_models.md
|
||||
- Developer Guide:
|
||||
@@ -31,10 +35,10 @@ nav:
|
||||
- Design Documents:
|
||||
- design/index.md
|
||||
- design/architecture_overview.md
|
||||
- design/vllm_omni_design.md
|
||||
- Feature Design:
|
||||
- design/feature/multi_request_streaming_design.md
|
||||
- design/feature/omni_connector_design.md
|
||||
- design/feature/ray_based_execution.md
|
||||
- Module Design:
|
||||
- design/module/ar_module_design.md
|
||||
- design/module/dit_module_design.md
|
||||
|
||||
@@ -96,6 +96,7 @@ Model execution components.
|
||||
Configuration classes.
|
||||
|
||||
- [vllm_omni.config.model.OmniModelConfig][]
|
||||
- [vllm_omni.diffusion.cache.teacache.config.TeaCacheConfig][]
|
||||
|
||||
## Workers
|
||||
|
||||
|
||||
@@ -9,3 +9,7 @@ Currently, the main options are maintained by stage configs for each model.
|
||||
For specific example, please refer to [Qwen2.5-omni stage config](stage_configs/qwen2_5_omni.yaml)
|
||||
|
||||
For introduction, please check [Introduction for stage config](./stage_configs.md)
|
||||
|
||||
## Optimization Features
|
||||
|
||||
- **[TeaCache Configuration](../user_guide/teacache.md)** - Enable TeaCache adaptive caching for DiT models to achieve 1.5x-2.0x speedup with minimal quality loss
|
||||
@@ -1,113 +1,195 @@
|
||||
# Architecture Overview
|
||||
|
||||
# Introduction
|
||||
This document outlines the architectural design for vLLM-Omni.
|
||||
|
||||
This document outlines the architectural design for vLLM-Omni, a new repository developed under the vLLM project. vLLM-Omni supports multi-modality models with non-autoregressive structures and non-textual outputs, extending vLLM beyond traditional text-based, autoregressive generation.
|
||||
<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/omni-modality-model-architecture.png">
|
||||
<img alt="Omni-Modality Model Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/omni-modality-model-architecture.png" width=55%>
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
# Goals
|
||||
|
||||
The primary goals for vLLM-Omni are:
|
||||
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:
|
||||
|
||||
* **Support for Non-textual Output:** Enables the integration, efficient processing and output of various data types, including but not limited to, images, audio, and video, 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.
|
||||
* **Non-textual Output:** Enables the integration, efficient processing and output of various data types, including but not limited to, images, audio, and video, 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.
|
||||
|
||||
# Expected Gains
|
||||
|
||||
* It breaks through the limitation of vLLM that only supports autoregressive text generation models, and provides architectural support for popular model structures, especially multimodal generation.
|
||||
* By supporting new output formats and new model architectures, developers can easily connect to mature ecosystems such as Higress Gateway and Ray, improving the ease of deployment and efficiency of corresponding models.
|
||||
* Integrating efficient cache management, multiple parallel strategies, disaggregation from vLLM, to improve inference efficiency of corresponding models.
|
||||
|
||||
# Non-Autoregressive Structure
|
||||
|
||||
* **Diffusion Transformer:** initialized with a random latent noise, the transformer model iteratively predicts latent output for multiple steps, under control of sampling timestamps. Finally, it is decoded by the decoder of Variational AutoEncoder (VAE). It is not supported by vLLM now.
|
||||
* **Autoregressive:** dominant generation paradigm for text, now supported by vLLM. It generates tokens conditioned on previous tokens one by one. The efficient KV Cache management provided by vLLM can accelerate inference effectively.
|
||||
![][image1] ![][image2]
|
||||
|
||||
|
||||
| | Autoregressive | DiT |
|
||||
| :---- | ----- | ----- |
|
||||
| Scenarios | Text generation | Multimodal generation |
|
||||
| Generation | Token-by-token, left to right | Continuous vector, iteratively refine |
|
||||
| Key advantages | Smooth, good local consistency | Global consistency, diversity and controlability |
|
||||
| Efficiency | More efficient with KV cache | Need generation repeatedly, slow |
|
||||
| Error accumulation | Easy to accumulate | Can fix possibly |
|
||||
| Key challenges | Limited global planning, long context generation | Large global computational cost |
|
||||
| Sequence length | Varied length | Fixed length |
|
||||
| Attention mask | Lower triangular | Full bi-directional matrix |
|
||||
| Parallel strategy | TP, DP, PP, .etc | DP, CP, .etc |
|
||||
| vLLM support | Supported | Unsupported |
|
||||
|
||||
# Targeted models to support
|
||||
# Representative omni-modality models
|
||||
|
||||
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:
|
||||
|
||||
* **DiT as a main structure, with AR as text encoder (e.g.: Qwen-Image)**
|
||||
A popular model for image generation and editing. Main visual generation models are similar to the structure, such as Flux.
|
||||
![][image3]
|
||||
* **AR as a main structure, with DiT as multi-modal generator (e.g. BAGEL)**
|
||||
A unified multimodal comprehension and generation model. Visual generation can utilize CoT text generation output. Main unified multimodal models are similar to the structure.
|
||||
![][image4]
|
||||
* **Extension, multiple AR+DiT (e.g. Qwen-omni)**
|
||||
A new model for multimodal input and output. It designs the Thinker-Talker-Codec structure, which is a double-AR+DiT format.
|
||||
![][image5]
|
||||
**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">
|
||||
<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="Qwen-Image" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/ar-main-architecture.png" width=30%>
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
**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.
|
||||
|
||||
# Architectural Overview
|
||||
<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/dit-main-architecture.png">
|
||||
<img alt="Bagel" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/dit-main-architecture.png" width=30%>
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
vLLM-Omni will import core modules from existing vLLM framework and introduce several new components and modifications to achieve its goals. The core architecture will be based on a modular design, allowing for independent development and integration of different modality handlers and output generators.
|
||||
![][image6]
|
||||
**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...).
|
||||
|
||||
## General ideas
|
||||
<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/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>
|
||||
|
||||
* It divides models in vLLM-Omni into multiple stages, which are processed by different Engine Cores. Each stage can choose AR Engine Core for AR model part, and Diffusion Engine Core for diffusion model part.
|
||||
* Designed a specialized diffusion engine core for DiT or dLLM.
|
||||
* Plug-in with external DiT acceleration modules such as xDiT, Cache-DiT.
|
||||
* Output linked with Higress API, Gradio for online demo, ComfyUI for less-code process.
|
||||
# vLLM-Omni main architecture
|
||||
|
||||
<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>
|
||||
|
||||
## Key Components
|
||||
|
||||
* **Processor:** Input processing, tokenization or encoding
|
||||
* **Scheduler:** for order and queue of incoming requests
|
||||
* **Executor:** Prepare model input and run
|
||||
* **Multimodal output processor:** support multimodal outputs
|
||||
* **Hidden State Output processor:** support intermediate hidden vectors as request output
|
||||
* **Engine Core:** accelerate AR modules, imported from vLLM main repo
|
||||
* **Diffusion Engine Core:** Wrap and accelerate DiT modules.
|
||||
* **Diffusion model:** Wrap up diffusion models with vLLM format
|
||||
| Component | Description |
|
||||
|-----------|-------------|
|
||||
| **OmniRouter** | provide an intelligent router for Omni-modality requests dispatch |
|
||||
| **EntryPoints** | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni) and provide the OmniStage abstraction for different AR/DiT stages |
|
||||
| **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 |
|
||||
|
||||
# Examples data flow of specific models
|
||||
Disaggregated stages are managed through configuration, such as in the Qwen3-Omni example, where stages like Thinker, Talker, and Code2wav are defined as separate OmniStage instances with specific resources and input/output type.
|
||||
|
||||
1. Qwen-omni
|
||||
![][image7]![][image8]
|
||||
2. Qwen-Image
|
||||
![][image9] ![][image10]
|
||||
## Main features
|
||||
|
||||
# Overview of function call dependency
|
||||
vLLM-Omni aims to be fast, flexible, and easy to use with the following features:
|
||||
|
||||
![][image11]
|
||||
### Performance and Acceleration
|
||||
|
||||
The framework achieves high performance through several optimization techniques:
|
||||
|
||||
[image1]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAALYAAAFwCAYAAAAVNoEZAABxw0lEQVR4XuydBVRcWb6ve90709OStnTScel4x93dIUaEuBJ3dzciRPDgEtzdCRpCBCIEEkIgQHCCBHe+t6uSvNvNnZ55r2/fng5zvrX+69TZp4qkan3nz2/vKup88sm/OWqXL3HxwnnOnj7JubOnuaB6jqtXLqJ28TxXRWmrX8Ha0gIzE0OuX1bFQFsTcxMDLMyMMDXS44bWVRxsLLG5aYKuxlVsLMxwc3bAWF8bNzdnHGwt8fb2wNXZHlcnW3HMDjsbC4Ju3cLO2hwvT3ecnexwsLcmIMAHHy83+X18PVxxtLMUt+1xcbTBw8URRwdr+f1tLM1wd3PB08NZlAs24ue4uTrKf0bD5yfxb8qWDWvZtnE9W0Vt37yBvTu3cWDPTlE7OHviCBfPnsTEwECIfg7Vk0e5fkkVoxtamBneQF1NlSsXTmMrZJaN6elcF7LexFaIZ6SngbOjvfy2l5DUSwhoa2mK9U1jfL09ibwTKaS3IPx2BA4ONnh7uuHr4yUX1d/Pi0B/X1xd7PHz8cTT3UUc85SfIMHBQUJ4a3EfH/mJ4OsjThoXB/njPNydJbEl3rFisTKrli5i9fLFrF21jC0bVdi2aR07tqzn+KF9nD91HF1NDU4eO8CZIwe5fO4UBjoacpEvnDnG5fMnsDQ1Rk/rOje0rwmxLeQd3chAC0d7W6zMjfBwc8JHiGlmpIu5oS6hISHcibgtOrE99+4+wE50dR8hraeHm1zOW4G+BAb4yjtyYICf/FiAED0wMICw0DDxODvR8QNEBeLv7yO6txPeotN7e7lLYku8Y8akccycOhGlGVNRVprJsoXzWLlEmbUrl7Jv5xbOHD8sBD/ApnWrOLZvN2ePH0FDRJVronMf3rdDiH1aiK4p9s+JE0AWS6xE59bA2EAHO6ub3DTRw1UI7Ovjg/oVEWV01Llz567ouH6i+wbz8FEMtiKaBPr7yzuvn6+nED9I3rH9/bwJDgqUSy3r1OHh4YSEBOPu7io6/W1CQ0PFSeAvxHeTd3B/X19JbIl3TB4zjOnjRzJz8jiUFCazYPZ0li1QYrXo4nt3bObE4QPs3LIRlRWLObh7h9jfj5rqGdGtj7N/11Yh9hk0r14W+yfQ0bgiZLZC8/pljPR1ROwwwdxYF2d7GyGfN1cunkJfiB0h4oe3lweRkfeIjn4iF9vf1wcnkbX9fL0ICb4ll1kutrgtkzdEyB4ZGUlQ0C35Y+/cuSPv3oHimOxkkHX2W6KjN3x+Ev+mjBzQmzGD+zJxxCAmjxrC9AmjUBZyL1dWEkJvEF17G6tEB1++aB47Nq0Xsm+Vy71PbHdv2yTP4LL9Q3t3onVdTYhsyCXV02LyqCuyuS762texlk0oXV05e/IQetrqBInuK5v8RUc/JiLiDvYiisgmf7Y2N0UssZJv1a9fFQJ7yiUPkXXrsBAePHiAn+jMsghy//590dlDRRTxk8cW2VhoWLAktsQ7Bv7UleH9egi5+8kFl3Xw+TOnsmTBLDGpXMsOMaFcunCuiChKbF27ml1CbtnEcvt6FbnY504eE517uxB+CxpXLwlxNbhwXjbhvCGXWFOMmZsZibztwMmj+7khsniA6Kwurs4iSoi87OKCuuj4urqa3DQ3xdHBVt6pXZwd0NJQJ1AmbGiw6NC3xYkQLbK6lxA6WC55yP8VWxZrbhEWHiKJLfGOPl06MKhnVwb06Cy2XZg0aqiIJJNYNHcG61cvl2dtZXF78byZrF22iHUrl7B53WqRwxcI6dfLO/WW9Wvk3f3i2VOonj7BeRFT9HW0ULtwRoh/FIMb2hgZ6otMLuvql9klos3COYri3xrOiEH96dujK907d+Snrp3o37c3Gzdu4MKFCxjqG7Bn9y7uRt6Vd+gHUQ9xFsKHiMmnbEy29ReRRbaaEhYWJo8qDZ+fxL8pfbt2pF+3TqJ+ZHCvLiKODGX2tAly8VRWLJFPJBcoKcpFX6E8j9WL300slyrPkXd0WdfeIE4A2UqKbHlQtpJy7tQxtNWvckZILYspOprX0dFSZ8+2DezcqMKsCaOZOGwQvTp1oHPbVnRo2YL2rZrz/TdN+O6rJrRp8YMQvA+q5y9w9NARbgUFEf3wEffuR2FrZyUXOiIiQnTyEPmkUSa2LNLcu/dAElviHb06taF357bIBB/cqzNjh/YXOXs0cxWnsWjebPlKyVwRTebNmILyLAWWzJ3NMuW5LJijIO/o61YvY/WShWxUWcmB3dvZLgQ/emAvZ04cZc/2zULmTaiKuLJv22ZmjxvJqH496dK6pagW9Gnfhp7tWtH5h+9p//23tPruG1qI+v6br/n2qy/47K9/odn337FGRQUHJ2fCbt/GyFhPxI4Q+eRRJrivr7d8AvngQZTo6tGS2BLv6P6jELtLO7nY/bv/yPCBPZkiBJylMIV5QuS5M6czR3EKStMnMU8m+5yZLFGazdwZU1mzfAmrly2UTyzXr1omjyMb1qxgn8jch/fvYdsGFXZsXMMusV0i7j+6V1eGdm3PGBF5pvXtztxBPZjetyujxL8/VPw/erRuzo9C8pZNheBNv+b7r7/iy88/o2PHjixcspTbtyO4IbK4bMnvzp13KySyzB0cIptYRktiS/wXnzdpzpdNmtGtQyt+6tiK/j06MXbEYCaPG8WMKROYOWUc0yaOY6ro4orittKMacxWmMysaeNZtmAei+fPEZ19FquWLmSdkHv5ovnyzL1dxJQ1Qvr1ohaOG8rxJTO4ukYJ/a1LODprDEsGdmdkux8Y2KopHb/+knZffk6LJp/xvRD5e3G7+Zdf0PKbr2gtJG8jZP/uqy9Zu24tqmdOEn77jrxjy96wkb1BExERyd2794m4c1cSW+Idn33ZlL99IaJA6x/o1KaZ6N7tGTmkP+NGDpXLPG38KLnkk8eNYPqkMXKpZ0ydIKQfK2LJHHkHV1aaIYReIM/jMtFlgss691KRzeeKn6G7YwVup3fgr3YAw83KLO/VkbX9OrNhYBfmd23JrC6tmNT+B7p89zVdRafu1uwbOjf9ivbffUUHsS/r4h3F/69Fs2YsWbhAZOn78owdGBgol1vWyWWTyYg7EZLYEu9o/cVXrO71E6aL53Nx1mT6iLw9sE9Phg/qx/jhQ+Q1doTYjhryXu6xTJNtx4scrjCVOSKiKImoIsvispJ19MULlEQGn8GsYf05vWI2+aGOZHqb4LJvNY+v7iP2+hEeXdpL2NF1uG9V5ti4nuwb1pUlPVoz88cWTGzblHGihrVuypA23zOsfXOGCPE7t/qBZt9+g7qmFmFh4XKxZXlb9kbPbdG1ZZ87afj8JP5NGd2mFfuHD8BCeQ43F8/hzMxxzB3al5H9fmJo/z4MH9CXEYMHMEpMKscNFxFl9AjGjxjERNHRFSbIJB/N9MljmDV9MjOnTUJh8nhmT5/CxNHDMNu/luI77pRGB5Dhb0t1zB2qnoRREuxEnpsx989sI3T/apzXKXBz6WRMl01Bf9F4VKcP5uj4Pmwd1p21g7ui3Lsj83p1oEerZrQV3fv7pt9y4uRxbgUHERIWRlh4BMFC9FvBoZLYEu/YObgfp0YO4ZLowtemjhViKXB1wVT2KIxBccRAxorOPVgIPnRQX8YMHSjv4GPeSz5p7EgmjhkhjywyoadPkmXxMfJMPmHoAOJstKhLjKLiaTiVCQ+pSY6jJj6K0jA3CjxNSdI/S4L2MR5c2sW9CzsIOqKC156lGC2bivr8MRye0I+dI3uyelBXlvXvxMA2zWnb7Ds+//QvjBkziiAxaQwJDRWTx3D5kqC/9Ja6xAcU27Rkcgvx67/59yzo2JrDo/thvVYZj51rCT20leCDm7DfvALN5XNRHj2UcQP7MKRPV9HJe4suLuvkQvZRw+SCTxg9XH57pDgZji6bSXVcJHWZidS+SaMuN43alGdysasfBlER6UNpsCNv/azJctQh1eoacdrHuX9hF67bFmK+SpFz04dwZHw/to3sxfqh3ZgssniXZl/LlwT/+p//ia+vn5g8+uPt7SX/YJSrq4sktsQ7un39Fd2/+ZoB337NpJZNWdezI/pKY3FSmU+o2mVCjc3xd/LCw9kHbSNrDl/QQFFpAaNHDGNov3dZfIzo3rLJ5tiR7/L4zJGDCBQ5ui4tntr8TOpF1WWnUJsUQ+2z+1RHBVF1L4DSUFeKAuzJdjEkzUaT+BtniLq8F+ftypitVOCsEPvoxH7sHt2XzUN7MLtHGwa0aEpLIfYXf/sMS0sLPDw8cHR0xN7ORl4Nn5/Evymd2rXgxzY/0EF0wj5dO3NE2xTj5CLMC+rQL6zDtKQe+/I6rEvrMS6owfBNFQY5leimlnLI0pOl2/aLDN6X0UPede9TG1dxS/sMiO5cHO5DcYgHed4OZNkZURbqSfWDYKrvBlN+25/iIDcK/Rx542JEpp0uTzWPEXl+Jw6b52O8fDrnFYZxcupg9o7pK7p2Dxb3aYdil5a0+eYrvvnicw4e3I/lTQsM9G+gr6eL3g1dSWyJd3T44Rt+bN2C+bsOc9w7ApMiuFkMN/JrsSyuwbm8FuuiWkzF/s3Caizf1mCZX4VpbgV6yW+5/iCRpVv3MkFMJBdMn0BehDf50SGQLaLH6xfUy+JHVCiVYZ6UeltRGexO9T2xHxFI0S2RtX0cyHMxI9P2Bo9Flw8/vQ3bjfMwWDqVc0LsE1MGsXd0H7aP6M6SPu2Z3bUVbb9pwjdffsnKlcvQ1rjO1SuXuHZVVpclsf8gNomSvdgfKuaXh//1qD3O4fLrKjRya9HJrUE3uxqr4mrcy2vwLKvDMKsKSyG0s9i3Lxad+20txqJj3xRlJeQ2yyhBNy6Dq0GPWLN8Kbn3g6jNEdEjL5uHphZcO3QRNyNHEnyCqY2NpfbFC2rjnlBxL0x0clkMuUmSmQZPdS9wR3U/t4Line truncated
|
||||
* **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.
|
||||
|
||||
[image2]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAARUAAAEOCAMAAACQIXOIAAADAFBMVEUAAAAAAAD/AAAAAAAAAADfACAAAAAAAADvADAAAAAAAAAAAADqACsAAADgACkAAAAAAAAAAAAAAAAAAAAAAAAtAAbmAC0AAAAAAADqADAAAAAAAAAAAAAAAAAAAADnADAAAAAAAAAAAADqAC4AAAAAAAAAAAAAAADpADAAAAAAAADoAC4AAADnADAAAAAAAAAAAACxACQAAAAAAADoADAQAALkAC8AAAAAAAAAAADoAC8AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAADqAC4AAAAAAAB/ABoAAADpAC8AAAAAAAAAAADqAC8AAAAAAAAAAAAAAADpAC8AAAAAAADpAC8AAAAAAAAAAAAAAAAAAAAAAAAAAADpAC8AAAA0AAvhAC0AAABXABLqAC8AAAAAAADpAC8AAAAAAADqAC+rACLpAC/IACgAAADGACjqADDpAC8AAADVACsAAAAdHRoAAAAdHRojHBowGhw2GR08GB1JAA9XFSBdFCFmABVpEiKQDSWWDCenACKq4+q8Biu9ACbCBSvTACvcAi7jAS/pAC/4Y4L9gpv+h6D/kKcAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAn++oGAAAAdHRSTlMAAQEECAgKCxATFRYYGRkaHSAhJCgoKCsvMDY3Ojw+QEJFR0hJTE1QUFRYWF1gYmRoaWttcHJyc3R3eHyAgYSGiYyNj5CSk5SXmZ2en6Klpq+vsrO3uLu/wMXMzc/S1tbX19fg4eft7u7v7/H09Pb3+f3+/iGfBl4AABcuSURBVHhe7Z1dcBNXlsePaeE4MW23jLyOhASOsPiwwYAXg5kwZAOJcajNkDDxUMkkVcmmaqt2HlJTNTOv+7Cvm62amtpK1W7NJFWbbDIDMyRVmYUAkxAIWcAQMAx2+Su2wY6FQTYysh0jS2bvR7fUfdV2d0stS2Lv70G6ffuclvqv27dvdx+dC8DhcDgcDofD4TwcFLEVOWNlcR9bZRHJlSjGB1T11nGwFblC2Op/m62zyMrmRDH4kKhiA/c6AMp8oVFUDLPrrPEwqTI4CFDvCx9j661TMKqIq93SdNdAFB1rK4pvxn114pEoSOt8UVKnh7DBI00PXSZF7NIiXhmgLsfm8UhQKKo81YhfA6N/moai5mVHdvgAWg97D+C64GHdndzwpIheA9Hr6LXoWRG7eIN/XPHCEoBlf7zPWmspFFWWD6MfunpvVfMnaMHxk+73wN3y3OOdF6D6aXdNJ2uNqG+Gtg4o2tIyR9ZiF3eLu+XxLuyyWs9DxRK2Ik85/fvuKAwehdJivBQ8PjZ2HbWco2Nj31yGNawxoqIJjp4ZGwudhC3EA7tcP4FcjiEXWGMwICkUVcbI61TESRr3FXLMxNrwawhKE1ZJvGVBOvy5V0lHMdjl7lzs4gO8QLWdn0I5ggBcFStcS8UZXIyRE29kKXm7qzZKUAM3aG8T8ksjILtMTRYTlzmVoS6Foor/Wdx34v0zy7hSqFDXmqNQVDkQ+7pnDCpal7IrskKBqCLApXNs3YKElSbiTDYa8xRIb1tE982LDyNThGAz7YSdMzeYVSYokLYSg331E0vW90bMHkHXhl/6p9hAdMn6ecZ4C1MgbQUugrdu/fmv2Or5GT805AjUrY+lI0oe3V8xQJDkMYt5kMuc/nmbw+FwOBxOIVEw45V/YyvU/IKtyJA8GPHbub8LbutXhvdVFAplxL+4cFX04KrowVXRg6uiB1dFD66KHlwVPbgqenBV9OCq6MFV0YOrogdXRQ+uih5cFT24KnpwVfTgqujBVdGDq6IHV0UProoeXBU9uCp6FKAqcdOPANMm/54zrxueZKu07IHP2aoENj1RzT+W/WwZW8WwZw9bYzf5dwRNdm9nqxgurDXSLVPyTxW002wNg7FumaLuV36lKueSx6bZGoYlUtjWLuJt8qchFXnYVmDG6EvNzZSwVfZi9AVywZzhPs+UZPd7Z3fraWK4z9luLEafnxOM9zn6aFa/OBsX96/Mck74m5d/ZzCSgxfvzT+Us8IvdNXVrcw1tw+9aTQi+bgsm0O5vFQFgjcNRyTnN7A1NpKfqqB9NvrTWPAmW2MjeapK8OYmtorlvJFuGZCnqsD5LUY7HTTULX3yVZVgyHCnDXVLn3xVxURjMdYtbfJWFePGYqxb2uStKsY7baxb2uSvKsY7bahb2uSvKvCxbytbpSX459ezJEseq2LcFoybU5rksyrGO22oW5rksypop9kaBmPd0iOvVQmG2BqWLDWWvFbF+Fon+H1WGkt+q2J8rXMpK40lv1UxvtbpykpjyXNVjPc5K40lz1Ux3ueu7/VSC2ZInqvS9dVBtorlgy3r2KqMyYEq9b/8qUHCQxVdQPd5AadLBhcGacA++UiT3Q0fKnn/6kme5tjoZBtOSpwxl7Z2sVUMXfarYk9bcfhhG1OzYu1r2po0URrLAlxiKzLGnrZSK4E3GUkQ/HwWwLdd3HlWbZMuJhrLOiMLq9jSVoQauFRSnVweHxsba/9deNV8PYElTDQW2w8hW1Qpr4z1KNl1E8TaaC7ajDHRmxrqZhFbvrhXRNdxlS42z6pDYsJzxNXrH1zpVpaEjQEHDH4j5/6rDrge3OgiSe+EljJ4cENZIZY43/j+iuIEwgHFabO7v7u60XF6xMRBZhE7VHFsgxvRsORnVHHR5LxJSI5ar7IHFa041eaK7b/GC9JeH3r1Crgnoiu82/+C84Bjp+KqsPcUdaIrV2w/0YnW+COPoPOdNILOQjb3LHaoskvqOwtHWpvGVTm6xcpnxXc0TUXY7w++h35jHxmXoaW+/yDvL3yCF2bel8/j2IyuaNl0OEqd1m394JfKSuwl7N+3BmcVb+qTp3v44FWwVRYbVKkIQA/AxB1xjayK+y3y1qk9fjz+2Cnc8Ef6/XSJZL2O9+16bNrjhy+UwQ02Iysu+t01ndQJtQXipHjF+/zklDdzQnay+xiyobf1ijOD+Juik7OqNvbZUfwm1BCKcfrdO+QmUpxk+Q7ISyTttRPk1N/qFUGcuFd2urSVrldWku4HYDihu4kTlRUybyvotDzhQe9LImIDHaAEUcvf2dgi4aXyPbiTiB0aARdMaZKqyi0KYtK0dpVcfjAuuRSn5PBV8SI9uepWXZetjSVzVcoroeoFWlzVpuxc/LTQsBWrMkFDkkI4raoatPQdLT3QrEo1I9BdxisTXlqubrGzw81clc2i/E3LRPXJ+YxfKo4qxwsirj0hhaEP95fKQim2xcRVZaSl4nT1aXmlykvD7BU7G0vGqqBLoPP0wHG8rj45x44f3MaM+OlFATriAO+w6hIhBE5HUgm5XFRBcmVTs9mAvFLlpcXWxpJxb+uWYv20hN43q/rbocvqJYBeoBcF6IgD3F+qLhGGIyUNKeVaCXXiihMUk5GzxkvL7BUbQ8IyVUVoVM4Zyf2WaS8h9xQURvqhGfXK0gGSJx0vkXh9AdWN90LTTlSW1tIyrvfvgPZpxQmkZY8oXsTph7hSy1Wnm61Km0yPIPTLy1NH4HOppAxZMONQU6tajH9ZKb7S4fQEr+A+Ai89/1w3OD0foqUzlb6mrbfKxfPocuC05P/5QMwrQudZ2Sk06vR89qIge/0cO8W+TW5YZvZK08dsXbqw/27IbrytWDquPjsvd96NJJaXl95SyqIzel8ektAl5LX09T8HqZ3aScNWn3VZSLxtyr8bMm0r1ohENItj6qzpqgWtGV1KNAWNk4arRg9gTZNpv7J4XHUZ9RuzySvrDCkcVVBjYatYrhrpZpbCUcVMYzHUzSQFpIqJxmKom0kKSBUTjcVYN3MUkirGjcVYN3MUkirGO22smzkKShXjnTbUzRwFpQpcGn2RrdIy+9FP7JClsFSBaysNdnryulFzMkOBqWK80xd8BrqZocBUgQsr2RqGyQ4j3UxQaKpMkkdnC3Ghqgy9rlvF1luh0FSBC+Tfqwvs9GRvI1r/BpYmbQpOlcntaJ9fXWinLwQ2vvoGW2mNxb2/Ygdrb2/ayNZp8MZfZausUniqzC68z+u2LqyZKdg7lPlN8WPKrzj1vWaFgpCYeXkeA11S7lAWVr9imP4rHpnncZE1CkwVw51+MB0xUs4EhaUK3mm2iuX+vRhbZZkCUwXttOE+xyYsdCn65F8WPSOW7sFP2T74K1uvpu4Zz8IGBhRcW4HZz/6LDUtMoeMjtsYa+DlloXFnYLnr2m22VsNUycTCBg8jS//ecKhmmBpqIbT9yr9olhj+OVk0a7eQYdLMlBGTd/OB5osnZ0W0Jz1n4fUrlOyeJQpVlezCVdGDq6IHV0UProoeXBU9uCp6cFX04KrowVXRg6uiB1dFD66KHlwVPbgqenBV9OCq6MFV0YOrogdXRQ+uih5cFT24KnpwVTgcjt2YiNQQfJWTcbby4UYbb1u/610ajlfffJT+mb+++TdzzcsOGcbR2APJhwXwG/3/ti8eWlWGS6qpGC6oIO9CDUQXMVLine truncated
|
||||
### Flexibility and Usability
|
||||
|
||||
[image3]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAPwAAAEqCAYAAAAxuz3hAACAAElEQVR4Xuy9B3QUV57/246AMTkpC6Gcc84555xzQgkFQBIIRAaTsQkGg43JGWMwwQRhG2MbbGzjgMceezyzO7sz7+y+8z9n98y+s+/N9/3ubZUolVpCIAkEut9zvqdbVdW3qqvv5/5+v1vVLZVKSOgR9e3p43hBpcLoF1UYO2YMtKcbQl/bgmyNmbqOmKXvAjMjd1gYe8LazBv2VoFwsA6Cs00YXG0j4WEfA0+HOHg7JcLHOQn+rqkIcEtHkEcWQrxyEOqdg3DfAkT4FSLSvwhR5OigEsQGlyIupALxoZVICJtNrkJSRA2SI2uRElWHtJh67oy4RmTGNyE7cR53TtJ85CTOR25SKz22dDk7obnb86z4FmTGNSMjdj7SY+ZRm01IJSdHNCApvB4JoXWIC6tDbGgtYkJqEBVUhcjA2YgIqES4fxnC/EoR6luCEJ9iBHsXcQd55Xe6AAEeefB3z0GAeza930z4OafBl+zlnAovpxR4OibTeWFOIifCwy4BbnbxcLGNpXPHHANn6yg4WUXC0TIcjhbhsLMIga15MDkQNmYBsLUIoPWBoI/peeXnJiT0SLp7+gQHfsxLz2PSuHHQ1zGBkZ4dge4IE0NXDruliResTH2oA/rDwSoYjtYhcLONgrtdNLwcGewJ8HVJhp9rCgLdM7pgD/PJ47BH+jLYCxEVUIzoQAZ7GYEuwV6NxPBqDjsDPTV6TjfQsxLmdoGem9yMvOQW5KUwL6Dn980GAOlRPRi0EvitHHwGfSpBnxo9FymRaugTw+Zw6JnjybEEfWxwNQ1IlTQ4VSDcr5wDL0GvBr+QO8irEIGezLkEfi78CHx/l3T4kRn0PhL0nfZg8Nsnwd0+Ea4c+jhuJ+sYOFpFwZmgd7KMIOhDYU/Q21kEqW0ZQANDsABeaPD0bSfwr7z8PKZMmAQDHXOC3QnGBiyye8B8lidFdy+C3pt3QDvLIIo6YRx2Ft0Z7CyyM9j9OyN7sGd2t8gexeyvhj0miEX2cm457OrITrBH1yM9tkEz7AS6BHR+ysIuS9DL/2bbdAM+Vh3lGfDMLMpL0MfJgI8OmI1IGfDB3sUEOIvwDPpCsjraM/AZ7Mws0vu5UpR3zaKBrxN6p1RuL0cp2ifBw6En8E7WauCZHcxDYG8Z2gU8y6ZcBPBCgykG/EsvqPDqmJegNVUbM/VsYDrTjWBnoJNn+WBOZSOWtS2jNNMfh/cewNIFy6mDRlDUiqVOHE8dOh6+BHyAezqHXR3d8zvT+EJ6ZLCTA0opnS8lsMrJ6jSeOTmyjiCsQ3Pda1i9ZDt/nhU/9z7wiQz2ZhoE5uGTa7dx8J0zBDWDu60T7s4on0iwdz7PSVxAbbDUXp3S56e24PrlGzRwtCIxlFL6kDk06NQS6LWU2tfQ8VUjMqCKBiYC3p8BX4ni3Ga8vm4XNq3diezkeg5+CIHPonuQpxp4P48c+Lplw8clCz6u6TQApqrtSGaRniJ7kGcGDYapcGMR3lad1jM7WsbAwTKaR3cHy0isWrQKrvYRPLW3swgm+INo+1ABvNDgiQH/MtXv48e+BN3pehTd7dQ1+ywvbjMjL1x67zi+/PQyjzofX76A9pbF1NkT4e4QQ/VuFpKjizpr98zO6J7LozszAz01ppLS+RLkpVL9HFmNyoIFKMtt5ZE9L6WJIno9CtPnY9XiHWitX4uKvEWYXbgUORTVm2avRmV+OwFL0Zpq819//hUH9pxAWc5iAn4Bf6zMX8oje2X+Esyrfa1zAGhFddGyTvib0TZ3A3758Seq61soujdSPd+A0qxFdOz1vI7PjKesIoGl+jUEfgXBXYqOi9fw8ZWr+OLGDRpkjiAquAQRgaWICS2her8YkUH0vj1pcPPLo9cUwI/ef7hfDg1w2QR4GgfemyJ70+wWZCWWEeTxiAjIoOyIojvV8X5u8Qj1SyXwKboT9Ls2vU7Ah/Mob2/JgA8WwAsNru6+dwJjRlE6P340jHRnwdjQkafwViY+vHZf0tqOvPQCHNv3LjxdonD1/PuoKqunx7Norm/D9s3bMa9mIZa0rKYUNgUhBDyr3SP8iuixiOrmKpw59gHam9fh4N5jOLr/JBY1v4bjh89g+6a9+PLmF6gtXYKvb93htfvONw7gP/7+77RuH4F2C3t3HcVP399DTsICAnQx/vzrr3h313H88d7PuHn9C3R8eINeexenjpyn7X7C6qU78ObrR/DNnbu4crED1698ijtffoszxy/g+6/vUjSfQ4NOI76j52to23vf/YAFTevx6fVPsXThFpyi4wom2LMS5+L8+xeQGl+H/XuO4EZHBz66cg2b1mzDv//lD9izfQ+dgwvY99Z+Op49OH34KN7cvBOH9+ynQacNW9e+TmAnID6sBPt3vY20+CJsXLEBZXk12L/zLRRlVmFl2wpU5NdgYWMbwR2JhXMXwdlWAl4NvatNiABeaPDEIvyEV16A1qTRMDGwhKmRG6xMvQl2H2xdtwnnTx7HwuaF+ODkMaQl5GHjqg1U26Zg26atePvN3aitaMKurW+hJK8RLtbx2PXG20iPm40IqtmjKIUPdMvFLz99j2+/+grleQvw5edfoWF2O25+9BnBvZ/A/QP27jxMsH5P9fQc3Lj+Oa5e/Ag7Nu3HX//8Ozas2IUlzVv5JFzT7HW4ce0mDS6v4zOC/U+//Er7O4A///E3rF+5m2DcjZ9/+IkGkW/wze2vsap9F5Uja3Dl/DV8/833+OC9y0iMmIu0uPkUuT9B/ewVdBxf4K3th7Bo3ibUlC3H+uW76P2VoyCjFRcJ+KKsJnz31R1sXL0Tp4+cxm8/36VBaB9uf3qDBrwVlPl8hN1bd1PmUImFTStRlt+E8oJ5mFe9GB52yVRC1OG9w4eRGleG948dR0JkPna/8SYW1i9Cflo5qovrUVFQy2v5mJBMOFoR8BZhZAG80BCIAT9x7AvQnjwGZjNtYW7swWfkrU39CPIT1Klb4GIXhk2rN/NavrFqPuLCcwmQZWisXYgPz52nqH0I8+cspYiWhHWrXuepO0vhowLK4O+Si7//2+84/94FxIfOxu5tB/Htl9/g2qWPKOpvxq8//YrbN29ThP0ciSH1uPXpHRzZ+z7amjbhG4rcn16/hdfXvstr9tqS1fj9l9/x092f+Wt/+u5nitQ/4ucff8bZkx/S9t/jzq1v8da247h390dcufAp5lSsxK9/+JWWf43bn32LBErnU6Ln4S5lAN/c+Y4ymK3YuukgastWYnHzG1RSrKPjrkJMQCUNUnfw1c1blD28h8ykRhzaexz/8tsfsXTBOvzxx+8oU6jEiYMncLPjGhY0rMSyhWuQnjAbzXOWUQSvh4dDKjLiqnD60FEaQOZgdftayiAOU7SnMieqiAaC/Ti5fz+/0uHuEI/lC5ZT2h/ZDXgXWwG80CCKpfRTxr0EvWmvwtLYlUd3W/MAfi3Y3IiBHwAHijo25sE8+rjYxsCNOqe7QwI8HJPoMQmeTqm8fmWpfFxoGcK8C3h0T4upxe7tB7Fy8RYkhFUiIYRACpyNaIIpNqiKauk6xIfU8MieFM5m5ucjLbqJoqX6OjubcGP1N6vVi9IXkxdx8JnZ37lJbeSFXY/ZlPZnJ7QgN3ERvwafGcdm9Bchi+r2lOj5VK/PRWpsM79Ex1J75oQwquFD5lAdTw6uo+OrQWxgLT1W8dQ+MrAc65bt5DV+sHcZAjyKEOxTggA39fV4f8pg/F1z6P1nwd89m8/Y+7qyCbxMPpnn55zOr8fz2XoaEF1t2fV49hhP5zUGDhaRcLaOJUfDjkB3tIrgwLOJOwa8uCwnNKhiwE+b8DL0p42DpakHrM18OfBslpjdBGJnHkIdkzqiNcFuF8UjkYdjPLyck+HtkkKdms1AZyLII5sAyEGoF6vf2TX3Uh7RC7PmUx1bxaN9QiiDuxqJYbUEeC1SIhso0jVyyNlMegbByC6lsfSdAcwm4gpSF6MwjWDvcjvB3s4fC7lpfWo7ClIWIy+JnEwDQWIb1fyLqJ02Dj9rKyN+AdJiFxDsDP5mJEc1UYRuohS/iUf9OAZ8SB3iOqGPJrNIH+ZTSduyQaGKX58P8y1DODnMp5TeazGC+Yx9AQI98hHomYcA987LdDQAcPNLdQQ/m8CjgdHTIQXu7Jq8XSKfxHNml+dsorkZ7ErgnQTwQoOp786cwPSJo2A4YyJFcfXNNeoZYmniKJR3QlZjutrF8sju6UQRyzmFolgaRbcMPjsf5JGDEC92o00RvwzHgI8JKueODZpNwFQTVDU0CNTwyJ4SOYfMZsubCPi5ncDP74zM7G45Nruujt55yYtQQC4kqIsIbmYGeX4Km6knJ9Pz5CVqyMnZ8cyd0JOz4hcS7K3cXcBHzkdKxDwkRzTy7IIBz03QxwQR7IHVHPho6VKdXyV3ONX33N5q4ENYxPcg4Bn0nuwOPAn4bG4fgp3Z25ldskuDN4v0BLwHAe9G0d7FJq7zzrtYfo4d2fV4Swa9+nq8o3WQAF5o8MSAnzFpNAE/mWD34zfX2Fuqo7qDpdqsEzpTBHKzj+NpPAPe2yWVX3dmwLObTlhHD/ZUAx/pX4IIcnRQGcFTQVGzilJmAp4iewKP7vX8bjcJdgn49Bh2K6wG4JPuA1/Aonmn1bCzyN7OrQZ9scxq6DM7gWewq4FvuQ98uBr4eErtlcBLsHcHvoI7zJuVLqUI8SwiqyN8EIEf4M5uub0P/X3gM+DllN4NeHeW3tvEy4BnUT5KDbwli/JB/DZmlQBeaLD0/ZmTBPwYGGlPJ9j9qbNRJ7MKVQNPNbsEO7tRxN0hsVt093FRR3d/Sl0DGfBe+Qgl4EMpwkewSbvASl6zx1N0Z+awR8wh2Ot5Kn8fePY4D2lRc5HBoW+hqMzulOse5ZmlqM6e53aCnpvIovt90LPiWNLine truncated
|
||||
vLLM-Omni is designed to be flexible and straightforward for users:
|
||||
|
||||
[image4]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAO0AAAFdCAYAAADv8OpOAABCCElEQVR4Xu2dB3gUxRvGlwRCh9BTCb2KVAWkKiDSiwKiCPxV7IACKlbABohBESk2pAkC0kR6SWjSW+i9Kr0TapL3P7N3c+x911I2l1z4fs/zPrmd/XZmd2fendkaTWMYhmEYhmEYhmEYhmEYhmEYhmEYhmEYhmEYhmEYJoUA8OtVqnqTj8rUfappcMmmOfyzNhXJKVJw1hxNWwSXbdo9vFK9IdUb59UYhjGHryrUfapFSIlnS2XPs6NOgdA9Imm3WSqcNfvuanmLrK0WGNxLYxgmZcyr1SL020qNhvlrmdCocDHMr9MR0fWfx+qGL7jVGqEV9Z/TtbJBF/1vVAPny8nYn6s1R818IagaWKhnEa1ITroeDMMkko/K1u1QPTDo4FslquNEszeBZz5AwtP9cbplT6FeLvVfi5641eZd3GjVD9fb9EVs63643fY9nG3Z2yFW6lLrdzC71tMQRUZny6zVoevBMEwi6BRSrmWpHHn3BmfLiXhh1GvCfBVzF5TGSpSWDWoDzOuL7Hn8sPLL9ujfqZpDDFWmTPrf2IjQnG9Pi6yVX2MYJvGUyJ63XW7/zIdr5w8Vpv1ADI076Mbq1KIEpg5vhCmRT1g03FGThz2Oazu7Aadfw8JfmuH6ru44uvJZTPu2kUOsUVPF/O8/fgz+fpkOly0V2puuE8Mw7mkrdLBx4QjsaPwSauYNwdLfWgBnegGHXsPtnS/jg1eroGeXiuhlot7u/hAqlgpEZj8NT7Ro+3vjlq3/R1eMYRgn+BtM+8/jXVExRyFsmN4WOPUmTq3qghfalHYY3pqtXLnzxGXLnn1f7wHfBmoMw7jHaNp1j3fTTbt+mjTtW3i1U3ndVN3efAf9Ph+GPoOGmqp+n32Nd78Yhi6vvomQsHCEFS/1xoABA9i4DOMOl6Y92wuBuTOjRr0ncPQOsP9aQop18AZ00fTj94C1B/9DoSJBl1s/+z++j8sw7vBk2gZPtdGNtufyvRTp0E1g75U47Lsaj/3X7fM7GAusP3YOgQUKXmrS8mk2LcO4w5Np6zdtnSLTHr4FLNq6Dy06dLadw/YdNATbTl+3M+26o2flvIuBgfl72q8hwzB2pMS0+8TQ9thd6L0onae08fgFNGvX0XJu3L40nn6yGDL5Z8ankaP0XpeaVohNyzDuSI5pd1+6hw3HzqPPwK9QskRxNG7eCgfEkPfA9QRi2jj0GfClbtj+r1bWbyHh8OsILxyArNlz68vsEUNmfXh8/LyMOyX0hv0aMgxjR3JMK3vXds93swx3C4Xpf4f9PFE3HzXt+18M0+e/9Exp/TYSTr6JIvkzI3uuQJtpZU+9at/xe8FhRd8SseXtVpBhGHuSYtoj4vx0+op1qFKjJrRSlaENnQNtTTy0OSd1YwYFh+gGlMNmW698ORa9P/5MzPfHIw8VQMVSeVAkJAwT/l6mX5jS870NLIs5dGPBgQNZ6foxDENIrGnlFd+f/pyPEqWsD1uMXgVtRSy0hRegLb8BrVUPPX3k73/aneMeEr3vplOXMOi70fDLkhu5Agvjt3lL7a4gy4tVy2IOxu5JSChAVo9hGEpiTLv3yj30/yrSYtZmYlg87QC0RRcthlWSxpXpIqZ4iZK6WdWFJil10Urek5U97K4Lt7FX3v6R92+tt3zEsleFmtutIMMw9ngy7RMt2uON9z5BtoAAaJ37iqHwKWHQ6/aG1XXRYtxGz+rGHTd3MQ6Tq8qyd5Xq/kZvVKxUCQOHj8bW/67i0C3b1ePrQi3s15BhGDvcmTZ/YFb4+flBKxgCrf/P0JZccWJWouViyPz7Lt245R96WDeu7EnlkPmN9z5G3pw5oFVpCG3A79AKhOhxE/5ernraK4ULF25tv4YMw1DaCR1q5My0eQMsQ+LBs6AtvuxoUKe6aDnXfayFvuy0Zf+Ic+K7mLlyoyWvnHmhTT8EbZnorUcshVa5HoIKFcJb73+MHNmyXQwLC3vTfvUYhrEjOCB72xz+/ocq5SmEnU1eRiWjacXwWCtUTBjsmjDtJScGdaMJ26HVbIpHHquLdz79AqVKl4UWUR7ax+OhLb1qOSdeIYbTU/dCy5YLosdHYGDgjR49ekyh68gwjAEMGOD3aGBwvWyZMs8pn7sA8vtntzftI80sJqOmTIyib0Gr2tDSwxYIhrY2wbn5ZZzozbOKofjSpUuRkJBQiq4nwzCECrkK1dUyaZeFEjbMaGeOaaV+3gDt4ToW4/60ztJr0xg5VJ59HFlFzOzZs6Vp36brxzCMEyZ81ahA+WKBH3dvVwZXdr6MgnmypNy0UtG3oY1ZDS0gK7TIBY7zlwojTz+InH7arSlTpqwUpq1K141hGBe0bhxeI1s2/7stGxRFzqx+0Go1T7lppeTFqS9nQPvziOM8q2kLZvW/2LZ1C/5mFMMklbLFC1fKlyvLy+LnXa1mM8vwVV44Si3J+77CtFrhkMs5m/D7tAyTXLJoWXNt1IKLWe69Lr5iue2TGlp5B9rwhdCCwo/6PdbieboiDMMkkor1Hn9MCys1TcscIM5Hs0PLki115JcZWbLm+Dc8vEAIXQeGYZJIoT7fNG3+ZBM8Xq8uGtatkypq0qBefL0XekTRshmGSSL37t1rGR8fvxxeQpT1XUJCQkm6HgyT4RENv49V7em8xCI8lEWYKIr4KlUR5V0UGkHXhWEyPAYTXBbG/VT8HCj/CtUSv/2E/KXEdFZnEvNyiV62lcrHm8TFxV0U61Dauh4B1r/+dBsZJkNBjeAKYeqrwiSXqcSsKzTWm4j1ui3XQ/y9JHTNmnZEGDiMbivDZAioCdI1wpPnb9DE+1w7fwZnrt6Vpo0Rpo2g28owGQLa8G2cmQGtwPP4S76Jng7YNPYl1KzUBDfpDAMbv2yCoGLNdNNOB4rSbWWYDAFt+IrYqCEo+MxP2HMpgc5KA+LQrHxpNB8UTWfYcWHtp6gZVAL32LRMRoY2fAvnMLNfQ7w5Zy9ui6mXH6uCRzoNRO0yRRFcpAiCitbFxpVj0aNhBRQuVBghD3cRtrIw5IV6KF00DEFBwQgNL4n6Lw615bpqTC80figcwSFhKN+gC16p/yhavT8HuHMKS0e+hvDQMISFh4rlyiLmYrx1qVvY/ctz6D7hMOThI+HKbvz+YRuEFyuGYkVD0azvROy8ZIm9c+Ef9K1REPPuxscMOMymZTIoNkcZubseg9sWx09bzuuTFXNp0HKVwLglG7Bu7rdoV8wPDz/eCj0+m4yZn7REsMjmrHTUjc3Im7cMnhk4GYdiVmDM241RwL8A/pOZJFxFxby5UbH9R4hZ/Qc+eroycmYLQqtv9yDqm2cQmisPJq3ciyNHtmHqhy3Q+puVOCu8GHdlHfrXCULUVcuqbR7fA6VzhWLnieM4tnspcmQqgI5jtuBigozdh8/r58PPV+JjfrzJpmUyKMqnRs4tGISGBTTs0y/6XBZGLI93l57T513a+Qc6BflhtfVU95/IjiiWORyx+lQCDh/Yi107tiFq8SyM7N8RZTLnwRVcx9/9G6L3zL24Y1kMuL0XBct1xZwT/0LzD0erj//Eju3bsV1oR8xi+Id2wsy9N7C036MoUKiyWgo7ZryLGnkz4/XIGdhw8LRIice9e3FIEKa9d3k7PqmdD7/cYNMyGRibGwysG/kCimsB0PvZu6uRv+LrWHHpnpi4g4OL3kc5/8LQ760IfnulBnLkeRLSsJd2z0eVMsUQHhaGsuXKoGRwLmgBFUUXuBlftQzB9Jj7l37jrq5GyZaRiPl3BzIF5EZIibIoWbKkrtJlyyLsiQ8QffwWvm4ahoIlnrUth4QziBrdE/kDNGQNfgT7r6iBuRhIn5mPl0rnxyIeHjMZmftuuM+g5kWgBXXUf1/88w00HroV1+Vp47WD+L1rBAJK3X+WokNpP4R3+Q1357+JMvlz3O9JD89Gt5IactV8TRh/A4a0K4rx2y39Me4dwYTuFdB16jncPbcdOav3wJwLasEE7Jo5BH/tkbeAgedLi3Pmpt/ov+PP7kSn93/FdlXIv7OghXbF3H2WsfOeMW1RJF8JxMXH71jA92mZjIq1+dvRpWwmZHt8EOTQc9WnT+DtJaeh97Mn1+DdatkR3KCvLbZGzmxoMiwatxb2Qrm8WXFKBgq++V9lZNH8UO7pwZBD7DWRHVD9peFYsikGP/ZriQjRUw7blSCvHiGrfwhaD1qAy6f3Y9XkfqhTJDPmH7EYUZo2tPlw/XfChS3Qcj+C92bs16dvn5qFLBV7Yckxy42gb5sXRb6yHeUtn50H2LRMRkVv7YSxg4di+vpj4tcprJwyFidiLU68emIVfo0chT9WHbFGXkXk8D+w+aTlBPfSrvlo07QJmrXvivUxe7Doh4/x1bgl1lhgxe8j8eWgz7A26js0L14IJ+9a0mNPrMPU7z5B15d7YsD307Dh5P27sUObRyC46lu26cOrfsfgPi+iYaMWeO71j3BJjY7jT6CgMOzEffJ6t/5U1E66rQyTIbC5wSSuXbqEaxbf2Di1diI+6NrdNn10THuE5iqCe4m4Bbykb20EFWtBk3H+0nXbbSbJqXn98NyoTbaHL+RjjXRbGSZDYGj3qcaFpV+gjhgOf/LzdPw1bwYeC82OQo/1oWFOOT2vJyqHVqTJ9pxZgj71y2IXecSRbivDZAjsm3kqEheLzevWYNWq1Th+4Rad65qEWzi6YbHLxxcTYk9g3cIFWHvQcuHKCN1WhskQ0IaeXnE5kpY3aF1At5VhmBQQpWmBUjSdYZh0iDRrtKZtk2LjMkw6x2pYGMXGZZh0iqGHtTOtTPtb0/IZYxmGSWOc9bBUbFyGSSe46WGpuMdlmLQmMT2sE23TF2YYxvsI07YVGigVrWlHox0NqnRUxUnZMmAYJu2IFh6OdjSrUtT9SIZh0gXRbFqG8S2i2bQM41Line truncated
|
||||
* **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.
|
||||
|
||||
[image5]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAR8AAAEwCAYAAACOrXWuAAA5IElEQVR4Xu2dB5xU1fXHH21FsSCisrTV2LBFMAhG/EfFilEES2LDQiIREZKoUdSIGI1ExW6CGAEBFYMoTXrZoQlIbwssbQHpZWlL3z3/d97M2b1z3ntT9s3dmTdzfp/P77Mzt5z35s6737n3vrKGIRKlkHZt2wt+t6lK/HOJRKIUF+/IfrTARyTyoXhHdvLu7ftgz84DtvRUscBHJPKheEd2c+GO/ba0VLHARyTyoXhHdnJg0jRoddvtsGbleji496gtP9kW+IhEPhTvyE4e8r+hYBaFvMUr4ciBYlt+si3wEYl8KN6RyTjNKtwenGp9/+1wyMrKguVL8i34qHmpYIGPSORD8Y6M3rGlEIYOGQGjR461Xg/7bmQQPktXwYE9R2DQV4PNvHGwd9dBW91kWOAjEvlQvCOjt2zcAS2uuRZa33GX9fqHYaMt+KxYthp2bN0Dl15yqZV39GCJra6b9+4qgn279cBK4CMS+VC8I6M3b9gOJ5xQHc4551zr9ajhYyz4rFy2BrZt2mmt/2BeyVGw1eXGKdpOE1jjx0yC774dbsuPx3jKH83TBT4ikQ9FHRjXcMCEyeEDxRZwatQ4Gc4773xX+GBe8RGwAFRyJBxCCBuETvFhjHcctm/eBfUbNLDq4TQO141wW04gcTOeaZs8caq1D/y0v8BHJPKhDu8/bk2Htm7aBT+YkJk1Y27s8DEBgms/Y0dNCAMJwmHrzzthwtjJsH7NJgs4F154kVUPy82aMcfaFk7pOGRU7y88BIfM/cMp25gfJkDbu+6GwYO+M9MPQ9HeI1Y+xhP4iEQ+1EMPtoOB/QfBqhXrLDjc1PJmCxyxwMfq9ZUqWe/xCmha10GgLVu8ErKqZUH/fl9agLjookZWOQTHTTfebL1esmC5VQ/r0GjGOpNmGtM//7QvPPLwozB39kLo93l/qFa1GrzdoycUHy+BF194GT75sBeszd8g8BGJ/CgERFezI6/NX28BAS8mLA98cCSCp+Q/793PGung4vSpp5wKXw0cZINPq9t+G7puKB9mTp9j1ulrgigPdmwutEYz5Kef6gInVj/RHF2Nh/59B8JJJ54EPd96D0pKSqDZVc3hkXaPwdKFKwQ+IpFPBX97rmspfG5v9du44YPGNZxfXdm0FCqrlq+D0049zYTPNzb43N7qDut1ft5aePnFV6z6A8wREq4VNWncBDr88UlYtmgldPxTJ6hSuSqMGjHWBp8mja8EHLUJfEQi/yph8LmqaTPrdTzw+ftL3aw6A7/42oTPXrjil1fAH9s/AUsXrRD4iERprsTA5xjAxY0usfKWLFweM3z+9uwLVv3PPwtO1x5/tD289eY71tktgY9IlMaqV68+vP6PN83OXmAB4Zabb40bPvger3Zev3YTrF25AYr2HLHuA8P1mgH9v7LBB7eBr5eakNqycbsFPtwmTrtwoZkWobs+/xLkNDwHJo7LDYdPcQncdmsr6Nzpz7B8ySqBj0jkR21av80Eyi7rlPjdbe6GHv982zoFHit87v/9g/DgAw+VXgCIxsXidas3Wms30wIzYee2PaXwQbC81aMn3NP2XutMFZ7ZcrveB/cLt4+v+5gjIwQd7h+eTUNYYT4CS+AjEvlQCAoEwvbNu2HKpOmweH5ezNf5lBwHmDF1Nvw4bXYYNBAmeGHhnJnzrfrbt+yG88+/wKqH+TjimTp5hgUPDhzVwRHQIQtQP/04D9547U0ImPu4zxxl4Wn9PTuL5Dofkciv4tBAxwyf0FXKbiMXTEdI7DDB1vHJp+GO394Zth1ePpKxfOGOA471BD4ikQ/FOzIagZOoe7vIa/LXQ37eGlt6IizwEYl8KN6R0TgdwsXcl158xXo94vsfgne1L10Ne8wp0BN/+JOVd2j/EVvdZFjgIxL5ULwjk+k2B3xtf5jYAdvNncm0wEck8qF4Ryar6zJ2+LifoUqGBT4ikQ/FO7KTvx88DKpUqWJduSzPcBaJRAkR78hOxpHOgT2HU2q0o1rgIxKJRCKRSCQSiUQikUhUsco1jMfQPF0kEom0CaETMAxATzKMx3m+SCQSaVHAMEoCIfigcwVAIpFIp0IjnjDwkAVAIpFIi9SplpsFQCKRKKGKNOLhFgCJRCItynWBEBiGXI0sEon0SeAjEomSIoGPSCRKigQ+IpEoKRL4iESipEjgIxKJkiKBj0gkSooEPiKRKCkS+IhEoqRI4CMSiUQiUSqrQYO6IBanm81DuzI/1kUpJv6licXpYEPgk/riX5pYnA42BD6pL/ULg4AhFvvWAh+fSeAjThcLfHwmgY84XSzw8ZkEPuJ0scDHZxL4iNPFAh+fSeAjThcLfHwmgY84XSzw8ZkEPuJ0scDHZxL4iNPFAh+fSeAjThcLfHwmgY84XSzw8ZkEPuJ0scDHZ8oE+Gz+zoBVXxmw7ht7nhf3ed6AKy+0pyfb119hwKz/2NPT3QIfnyld4XNskoEHoKMHdbOXL491wWftIANqnWrAtqHB97gd/hnQWdUMODA2vO6h8cG8NtcG3/d1qUvG/X/29wZ0e9SAklz7vvjJAh+fKV3hc33jYOe6/BcGFIU6KALp1qtSHz4Ig0qGCYPQe9rOgs+DI7gRbxrQ+e4yeBwYE15fHeHhZ8Y66OUDgnXoPY0GEVgIMr+PlgQ+PlM6wodGCjw90dYFHwTP16+UvaftcMhYZSsFP2ssoxYaFRHUVOfUCQKIp/vJAh+fKd3gM+3jYAeLdXSDUxwsr9qpE459xz5d+aiLHT48HsZSQYK+oYl9m+ooB6dMKkwiwYemaDRtoukmvudlI8EH2w1B5rQNv1jg4zOlG3xo1LNtmD3PybVOCU7Nhr4RnIbQdO1rBV7UMdVyWAatwodAgGnqNEeFGcGpR4eyMjiFIiBY6y+PhQMiEnwQNq1blAGrvPDBPBz9+HnqJfDxmdINPndeE+xgtM4TyQgVp1EOjZ4IQAgeXo46vQofp3IIDkzHxWOCGB8Jkd0AEAk+aAQW5ZcXPhQH980tP9Ut8PGZMhk+z/zOPm1CUyd99bHge1yDcSrH13ywDi9HwEH4EIjcIMLPcvHtONVTIZgI+Pj5rJfAx2dKN/jQtGt2L3ueaqvTXlN2SprnYQzMU1/zcip81FP711xm96xeQUhgvlvnLg98qA6NpgQ+Ah/fKN3gg50PRz98+uNk7NQ1qtvTqZPStAtf83KO0y6Hcnx7NArieejyTLv42a7ywoc+j0y7RBWmdIMPmtZsln5hz1NNi788naBEkKAOzuvS4jKlOZVzqoMLyjwPrQJATXeDz8ddgttT96G88HEDn58s8PGZ0hE+aJp+kXHqQ6/V0/A49cI0HCk93qqsjNpBCRqY3qihAb+6MNjh+al2tVyHO4Om2FSGAIBp111hwOmnhG8PwclBwz8LGWNwkJQXPnglNMLTr1MutMDHZ0pX+KjG09nje0ZehMb7vxb2safzMlM+jByHPPqt4Kl0nk7GPLdYCAGnNSadRmhGmjL6wQIfnykT4OM3V/TVxjRa4tM9v1ng4zMJfFLTdLEhT9dhhJ3bIrifLPDxmQQ+qetEPwLEzZGmh36ywMdnEviI08UCH59J4CNOFwt8fCaBjzhdLPDxmQQ+4nSxwMdnEviI08UCH59J4CNOFwt8fCb1CxOL08WGwCf1xb80sTgdbAh8Ul/8SxOL08GGwEckEolEIpFI5D/xMwterDO2ZcOohHFLAkaJLc+LIRhX5CyzvYttbebFIFMnkSnbgeHBOmNb9jl8SnL17rct36NL91vgI9Ih24HhwTpjWxb4hFvgI/KzbAeGB+uMbVngE26Bj8jPsh0YHqwztmWBT7gFPiI/y3ZgeLDO2JYFPuEW+Ij8LNuB4cE6Y1sW+IRb4CPys2wHhgfrjG1Z4BNugY/Iz7IdGB6sM7ZlgU+4BT4iP8t2YHiwztiWBT7hFviI/CzbgeHBOmNbFviEW+Aj8rNsB4YH64xtWeATboGPyM+yHRgerDO2ZYFPuAU+Ij/LdmB4sM7YlgU+4Rb4iPws24HhwTpjWxb4hFvgI/KzbAeGB+uMbVngE26Bj8jPsh0YHqwztmWBT7gFPiI/y3ZgeLDO2JYFPuEW+IhEolgk8AlZ4CMSiUQikUgUr/j/1RGLxWJdNtT/28UzxWKxWJcFPmKxOCl2hU9pokgkEiVIAh+RSJQUCXxEIlFSJPARiURJkcDHX2pvep7pk3hGEoXHRyrtj6rhRnD/SFVD7yvktgtRZAl84ld1I3gAR2oXyufuoxYqhzIBPueaPmD6LMPeftw1QnXcVJHw6Wn6NZ4ocpfAJz7hwTvC9Fum15u+Pzy7VNhmCAlVWN5rR013+FxruMcrMMJBEosqEj6oEkNf7LSTwCc+4a/ybtPNTL9r2AFDcoIPggPTz2Tp8Sjd4YNt6na8FRjueW6qaPgUmG7OE0XOEvjEJxq9oGj65SQn+FDHihU+txrB8mSM1zn0V+3sCES13BHDPiLLDeWpVoVA5fntwkoExcvhvuBfDp9doXR1n0gEAIIxmvYXXw8NveYqMOz7jSMNvt/qVCwafPA9llH3D9uT7/8DSj7FVPefhCM3bJNo00GRYWOMwCeKeOfAzsg7OgrLEXyyjTKQLC4tEVk4tcPyaqemKYkKHyqnSk2jaSKuoVxaWqJMlM+BRp2UPoNbOT5NonJqG/E0is0hg+kIARxVOqnAsH9WDj0EvAr9aPDB1075fP/VOhST7z+KyqqwErlI4BO7aIShwoZ3NBId1KpxsZl3FjdhefXXGOUEAKdy6vSO4OAESBTlO3V42m+UWznqbLQ/buVwnxCAKKqDC8qqcCQZL3y41O2gIsGH2ukRJR/3H/eBT51whEX7SzH5/pOwrMAnBgl8Yhd2fBy5tFL8pBE8ELuXFbOEafQL/GbofffS3KCcpkvU8fLine truncated
|
||||
# Interface design
|
||||
|
||||
[image6]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAdAAAAHUCAYAAACKzHkWAACAAElEQVR4XuydB1wUW5q3Z77dmdndibuzOzM36DXca85iwAiIZDEAKqJiwpxzzllERUUlGlCM16yIIkGSmHPOCkoGQQWhn6+qoBXLcFFE6dvn7+//6+pTp9Ir3U+/p8459ZvfCL0hhIRKWNnZ2bx8+VJZzs3N5cWLF8p7uTwvL08pf/78ORqNRlmfnp5OZmam8l5ISOhtyZ8X+fMjS/4Mye+1r4XryGXy5ygjI0Ox9vP2qVLzQ++lDpCQkJCQUOmWDEctLGVAZmVlKa8yVOUfqNrXzy01P/Re6gAJCRVX8fHxDBgw4JX79++Pi4uLstyvXz/69OmjlMmW3/ft2/eN+sLCuugRI0aQmpqq/jiUiLTZpjYTlQEqt97Ikl+1y59ban7ovdQB0kXJf0zyH5Hc7Kdrfvr06aummF+LLl++zG9/+1v5w8Zv/+23b/n/vaNMWFhXLf89y3/rf/jDH7h//77641BikjNO+btP+/2nvU1SklLzQ++lDpAuKjQ0lI4dO9K5c2ecnJx0yl27dlXO+9ekK1euKAB13biAiznnOJ15Qlj4V+urXGLEvGH8+c9/5sGDB+qPQ4lJBqicgebk5CjNtV+iz4CaH3ovdYB0UcHBwdjY2CjNJ+oMr7T71q1bdOjQQX1JOi0tQGd5Ticq4RhhD48KC/9qHZsazcAp/b84QGXJ2eeXyDy1UvND76UOkC5KC9CSavcvSd25cwc7Ozt1sU5LAFRYn/w1AfqlpeaH3ksdIF2UFqBf8pfY59Lt27cFQIWFddgCoHosdYB0UQKgpUsCoML6ZAFQPZY6QLooAdDSpXcC9NFRjidHEpMkXBxHJ0YQ+iD4rS9x4a9nAVA9ljpAuigB0NKldwH0yK1DjJg9nDELRgt/oscuHCPFdIYCUfWXuPDXswCoHksdIF2UAGjp0rsAeujGQarWrEblimWoWa2C8Ee6Xq2f+Nc//k8BaYyUyau/xIW/ngVA9VjqAOmiPidAk0KXMNRtL68mwdLksX1qJ2ytLTA3N8fcypmDd9LIS7nF9c1jWRKRUXjzj5Y+AbRilSrMneLA7TPzuXl6nnARffvCQk4cnUK92j8xZr4AaGmzAKgeSx0gXdTnA2gq8xpWpYbLWM4nFky6nJfNINNvWHjgLklJSSRHe1Cr6zye3TvBqWnVaLs+4c1dfKT0CaAVqlTGdZajFGYvSPYULqqz/Lh9bj61qlcUAC2FFgDVY6kDpIv6bAC9tgCjgQcInt+Pqfsv5pdpshlsVRH/ywV1EnZT3nAgcfdPcGJ2Xez8E19t/inSN4AumumIRoJCbuIa4SKaDB8lExUALZ0WANVjqQOki/pcAA0d1pZF16SFuI00GhVQUKphYvPfU7ZKHerXr0/tCoZM33mWvIcniJlcUwD0HRIA/bwWAC3dFgDVY6kDpIv6PAB9zhKL+rTuP46xg2z4t6qjySxYM87oTxhaOdGjRw96uszlRhq8fBBLxIQaAqDvkADo57UAaOm2AKgeSx0gXdTnAOiz85uoZupIWFg44cei8OpanaGhGuKfahhsXYG1Z9+cqDn73nEiJgqAvkslDVBSvCR7okl6XaYpKMt7R33ZecleSn3tq1KW9P76Rbenck/3ffspfLxPtQBo6bYAqB5LHSBdVPEBmse+6daYj96NFpPJO50oY7qMVZEPGWD8f7gfS39ji+w7UYSNKoel35M3yj9WAqAfYQl4pMxniFFFatp0ZWP4YglOEjhfrGShWQVqm9lwN2WD9GvIR4KafAypfoavtLyGwJkmTN68jCPjqjLxgC+axJVsXuTIweilkLUWnkr1nkuvadL+Mvykfa5VgJybIME51Sd/XZZUJ2l1wfnIx10PcZMwr96UMKkeyd7528p1n8n78GL/gDqM3bYcTdqaTwapAGjptgCoHksdIF1U8QH6krP7fybibsrrouzb+PsFsCM2gaj967kYn/N6naTcjMfERQdw8Fr+U+E/Vb9GgN64caMEATqXYa1McRnUmJ6+EyQw+cGpfpj3bUU/q3Y8SliE18whHApfhCZhOfu8B3D+6gqOLLRiyvTuOLUsi2FvZ+Zvm8/pXSM4d8WNkOWd8VjsTL/e1uw+PJsL3m1x7NSW4FMSXJ96kxY9mvE9jRg+xZm4RAnQqTIoXVk6xJypY81pVNOUiBQJsremcCOgA10djRkyrBtxEmAPD63PxB3LJTALgP5aLQCqx1IHSBdVfIB+PX1tgEZFRSnPU/1cDg8PZ926dfIHjakrJ5cAQOcw2NgKD3czWjo5kpjlzfHZ7Vm03YWhZm15Ej8R2zo2zPGfDgluTO7Zit0xywieZ8Y0vzlsGlSVIf5LeXB9Ca4DWuMfPouxBj8wZMUM7qyz4O/1m2M9YxSn17THamBvsm4toK9pTdbtm8XOsY3oMGUEEVJmu2tQQ2ymjOXBvu40LGvAqcduuHczwqKNA9ceeRLQryoDNi4ickwDAdBfuQVA9VjqAOmiBEA/XXXq1GH48OGMGjXqs3n06NFYmFkwfM4wohM/P0AHtTDj5+hheHRrRMAJN6YOsef4yeEMMLWVADoJu4a2LNg0AxLdmNHPnL3HJYDObc30ADcOjq3BhL1evHjkjttgczZJAJ1pXZNNZzagOdeXzrY92H9zBZlH+uIwSso4Tw6jdydbHjxZT3rsCOyc7Ri3fhoTLVsRdN8D0ucysLmZlIFKWentedzbNwjf1QOZ37MOvX3nETVWAPTXbgFQPZY6QLooAdBPVyuTVpTEk+zX+a1jwOR+JQTQ1kTfXEBqYDsshndj9BAH7t4YT/9XAG2H65aZEOfKGMfm7It1/yBAF7Srzf5rW8g97YKdTTd+vrycp4ddsB/pTLwE5j5agB4fTofuEkA3TGOCRSsO3csH6KDm5sQ8dGV+t6aYmVmxedtoVgyo+8UBGh4X8laZcMlbAFSPpQ6QLuqjAZp8Etf5c5k7dy4+O6O58mrevqIr4WwQ2zcGka5JZe+2fdzPLOKxVfraADVtbcqDU0GsmJsfD9lu7h4cPH2fW+H++PwcTHrBpEyy0q8E4rk2kPTn+dDNOn+QhfPyt1vitQ/tneI1q9YoXyolBdATcW5w2YUKteszakI/Eh9NUACaEDcB0xotGL5aAmjCfBwb1eHAqeWvALpvdHVmBq+XYLaCxYPyATq/bW32Xt38ToCmnJ9Cd6umXH2yibjtDhh1cWRp5BJ8OzfHLXIVxE+mfa1mRF+fSddWjWk0aJJ09Ts5NKgevdbNJ/oLAVRuKs9/Yot4WsuXtgCoHksdIF3URwH0xgbqGdRk0ro97NmziwltTajWcgZx6nq/oFOzzKnwp9Yc1zzj8rnLpGYXosxH6EMAffbsmbros6uVqSlPbpzh8B45HnvYe+AQQ1r+D108b7LH5U/8ppw1gTde/8LY6vT/+M333bmWlMuj42uoXMmGTT/nb+vZrxF1Bi4jVarnWWIAnY1LQ2MiUj3JfuGK44/fM2WzG5rHY+jZxIK7GX5sG1gF40aVsbNqhpmUKe4/sYygmcaM3+DKmTm1+d7WgZFe01k5xJT1oTOZZVGdXVekDPRUb2xad2HbJQmgh6TlIV3Jlo57YEZrjBpXp0XzOvhHrObp07VknexLk5o1adOzGVY/mnA8dTmbhjbHosyPNG5hyqLRTbGfMYHDo+oxZmsJ9cJ9dJSI+DBOpMbgHehJ7ca1mb92LtEJ4mktX9ICoHosdYB0UUUHaDaetpXovOX1/LWZ57exc0hbdibB4cGGlKlUixZ9l3D16laa/LMyVSpXpUyF+niEPCYvL5WVg834oVJVrJsZ8be/dORa3i16tOnJsYQXnPWfSJ3vK/FjxXK4n0gmowhMVQM0JSVF6cXq6OjI//7v/5ZI82phmZiY8DRTO2WElJxv6kkV65XIfYsPjCxPjTZjcPeP5lyuVJAdSpfydjQd3p/Lj16wpHsdZuwr9IXxIo5B5ob4nkzEa3UJALTAeYW2l/eVpx3XWVAul2lhpayX18ljPuVxoIXGfmqXC5dpVMvafSjHS5L3+3qbwsfQHl8pS8hfVs6r4Ljqa/gYqwEqf2HLEN17cRdjFo6ikXFDpdPWv77/lxJr0ZT7ZS0AqsdSB0gXVXSAnse2ohVHlHbGNDZPGMGAYWMZM3sZUfEaDvc3oM7c/ElvT65zJzI7f6u4vRPpNHobT66tpm7t4TyVyoLGNaPh/1pyJe+2BNBehN26Ri+Diux7ApqTrlRq5MzAQ2naA79XMkC7dOmiLK9atYoGDRrw7//+78oX4o4dO1S1P79MpQz0aUY+QNNPbKBpo04cLxjyumd4JZwmzMdqiBsjj+WRsc0O28WH6e/izPn7DxloWpvttwrtjBy2jnJgxIYzeHt6lhhA9c2FATrRbQKhD47Qb6ILZcp/r/ydaL1400IuZJ/lVMZxybHCX8iX8y4wfM5Q/vKXv5Ce/uZ48V+b3qSHkJ4B9CoOlRsRoMx9kMOjK5eICZiLm93/MjRSQ8iAegzcW5CNnVtJjaqVqVypEuW++z86uYVyP6A/9R39kI9ybdcM2lfqLgFUzkD7cVjKWG2qWXH2mQTQ1Bt0tbfEbPWd14d+j+7evav0hJXn2ZX/O7SWP4z29vZ0796dbt26vdfy+o+1s7Oz4l69elG9enWyNXmk7B7EQGcpEz//OhvdPeRHxu85RR/HmfRfuJtFFsasO3mZrt2duPAgjfF2NdlyvaBynpwpZ+PV14IRARcEQD+jtQCtXeNHuvRz5H/+8T9v/K3I/v1//B4zu9ZYOJhJNhf+grZ2tKJK7cr87ne/w8HBQZnyU/sZ+1RrP6elLaMthA4hWeoA6aKKDlCIHGtIpYFb0d5dTAycxhSzPzMkIh+gQw/kSKUJ9P/ub7jdz69zcmVXnKbsI+GCH63Npynbxiy0o9m37QoA2oLine truncated
|
||||
If you use vLLM, then you know how to use vLLM-Omni from Day 0:
|
||||
|
||||
[image7]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAQwAAAEbCAYAAAAvX7qwAAA00ElEQVR4Xu2dCXgUVbbHC5AgoIgITkIIwYERFB0RFxj1jeuMy3MBlTcgIouKKIIL7jMjqOMyioyCDqCCAgqIIIiyCzQgAsq+LwJhFRDClrAn59Wp7pPcPlXVqXR1JVWd8/u+/5fue2+dqq6u8+97by3RNEHwmH27D0IQpVOOfxZBEDyGJ2JQJIYhCKUAT0Su7D2H4MDeHNj/22FTXWlKDEMQSgGeiKrQLDZv2ApzZ82DxT8tgwO/5ZjalJbEMAShFOCJqOrA3lwITZ8DXbs8Dm+98Q4cOXjC1Ka0JIYhCKUAT0RVB/cdgTGjxsFlTS+Hdm3bw/GcPFOb0pIYhiCUAjwRScdyTkGu3qP4+qtv4IrLr4T27TrCidx8yDlwDI4ePgl7dx0wLVOSEsMQhFKAJyIawq/bfoOrr7oGxo7+Fr4bN9EwjA4PdILfdJN49ZXX4c7b74ITR/JNSWwlnAc5lH3EVO5WYhiCUArwRMw9eBx2bt0DVaueAYM+HgITvplkGEbH9p1g94690PnhLlC/fgPIPwGmJLYS9kT27Mw2lbuVGIYglAI8ERNpGAf35cKUid/DZZddYbw/lH3U1KYo4fAH4/ByMQxBKAWOHT5lDBl27dgH3+nmsH/PYUeGkacbxsRvp8DkCdOMYQclMl6vsXTRSuPsypaNO+CrkV+DvprI9Ry5xjpwyMMNgITtjurbhGdk8PqPlnfdDaNGjIHD+48ZZnZ4/1GjjRiGIJQCbe9rB8OGjIANazcbib12xQbYtX1vkYZhZGy5csYymNjYC8C/aEBdH30cMjPrwZBPPzcmTbENJvz2LbuM1yuWrDHaolHRBWH4F7Vx3RZ44P728GyPF+DUsTyoeFpFePvN3pB3Kh9efP7v8MH7/WHT+q1iGIJQGqRUTIEX9ETctH6LkczrVv5SbMPAX340hmmTZ8Bvv+6HJ7s/Deef3xC+GDYiyjB2bNltvF69fD3M++Fn+GTgYN08Vhtmgz0H1Mqla6Hy6ZXhkj9eApAHUKVyFej97z6Qn58PV17RDB5o18FoI4YhCKUDPPvMC3EbBgr04Qleq3Ff67aGGTz1RA9o2LCRbhgjLQ1j/epN8PcX/2ksO1TvhYwbMx4ubXIpdH6oi7EdFcqfBhdecKHJMC5t0hSwRySGIQilR0IMA9u0bXO/Y8P4x0svG8sO+2w4jB093uhRPNTpYWO4IYYhCP4FntCHEG4M40TuKXj6qWeMxMaJzqIMY62+jqGffQF33dkCpk2aYdyvMnXSdFjw4yJYtWyd0SazbiZAfrRhNNSHOffc3cqYAxHDEIRSID29Drz26huwcV1WOJnjmPTES8i3bNoB2zb/CrkHjkO3rt3h97+vD0OHfGFpGCuXroFft+0xTArXhZOdOAmKWrNivW4W9eCG626A/JP5hYaRlw+33HyrHvsJvc0GMQxBKA0wiXfv2Gf0DO5ucbeR9Hja04lhtP7bfXBfm7bGaU4STlwOGTwMnu3xPMwJzTPuRdFXY5gBToje0/JeY9iBJqGejiXhhV54WvdX3UjwqlM0pTdff9s4+4LmgtuKbcQwBKEUUK9/OHLohCEn12HghVt4dgTFkx7PeuToPYocvbcxauQY0PSkxyTHthjfyijsdCwnzzj9yu9dEcMQhFKAJyhd6Vmp0unwyUefmQzjoQc7Q7165zm60hOFvYL1qzcar4tjFEVJDEMQSgGeiGQY5cqVh4H9B0UZxq7tv0EH/W/t2ukAJ50ZBg5Fjuc6u1GtOBLDEIRSgCciCec09up/6fb2Du06GvMIO7fuhq2bdia0txCPxDAEoRTgiUgybkvfdzTqeRj4AB2a3OTtS1piGIJQCvBEVIWnS7+WJ24JgkDwRFSF8w/zf/gJev7zFRjw34+N05y8TWlJDEMQBEEQBEEQBEEQhEQxU9M68DJBEAQTIU3L1wUoViUIglBISDELUlQDQRAEBIchIWYWKL28Y1RDQRDKNiGLngVXQWNBEMouIQdmQYosIgiCYJiHySR0l5ArMwVBMBMSwxAEwSkhMQxBEJwSEsMQBMEpITEMQRCcEhLDEATBKSExDEEQnBISwxAEwSkhMQxBEARBsCEjozaIRH6XfqiW58euUArwL0Yk8qM0MQx/wL8YkciP0sQw/IH6pUBIE4l8o8y6Yhi+QwxD5FeJYfgQMQyRXyWG4UPEMER+lRiGDxHDEPlVYhg+RAxD5FeJYfgQMQyRXyWG4UPEMER+lRiGDxHDEPlVYhg+RAxD5FeJYfgQMQyRXyWG4UOSzTDWDDX+gVCUFvQ3tyuuBj2nQdPzzeVuNKefBuXKaZAzKRxf3earLtKgzY3mZVA1qoW3BZcbzJbjmtpbg7uu0SB/pjmO3yWG4UOSyTC63W1OGNSIl81tiysvDAMTv+rp4dfcMEgpFXVjmBy9XGaqBi0iJlCUYaCpoCm93N68fr9LDMOHJIth3HFVOEGGvmSuS4QSbRjYu6BeAr6n+PSe1K97YeLzGFxHp4bb5rNyNBU0jaD1MsQwfEgyGAYlCi+30/Q+4S4/6q3O5nrUzjHh+msvCb+3MoyT0zW4+YrYcbDXQ3E2jywsx97F7rGF7+0MA9Xjb+Gexvz/ht8PeFqDsf8yt7MzDIoRtF6GGIYPSQbDwGSjrn1RqnGmBhf/XoMNX4R1XROz2eCvMZZhUmIbfI1SDWPTiHDZkkHhNndeFU5qXv9m53A9zq307FCYzOW06MSOZRhoTBgLhyH4/s6rw8nPewyxDAPr1O0LgsQwfEgyGMbT/2f+9bcSJZ46J8DL6L0ajyYnqQzbYNJSAqvLURJjvZ2JYVuevE4Mg9Yfj2EUbJ9FnV8lhuFDksEwcP7CiWHYGQtOIuKvP74uSFw20agOSXr8X8RkWHKjqdAwg85+8HWhsPfBzcRrw0AFbR5DDMOHJINhYLLxX2wuo1dwVXSvgETlmGjqa74ONWH1XVcwD6Jqfn9zb4OruIZBw5vh/yxcvxiGUCokg2FQghZ1vYXdXAf2MIZHTr1iLwTbqBOSNAQhw8A4OAehtuFSextcVvMJsQyDTouqw53iGoYMSYSEkAyGgdI/ipHoKz8z15Hol5qXq8mP8xXYhgyElqOLpagNJrHahgvbv9zBXI6i5FXL7AzD6rRqPIZhZVJ+lxiGD0kWw0DpH6dAODyg1+qFWzjkwDK8irLjreHXnW6LjoPJjuWN6mpw2fnhRO7bPXr+g8wHk7DzHeF2akJS8mIZnlI9+8ywgRScJWFzHGgY6vaTcHluAPEYBl6LgVd88nI/SwzDhySTYZDwGgq8JDqXTVyqmvW+BksHmct5m1gxUHjKdOK/7dthvVUcNAi6WpMvk2hRD8luiORXiWH4kGQ0jKAIexlWk7CJlnpKOEgSw/AhYhilJ7z6kl/AlWjhMIVP4gZFYhg+RAxD5FeJYfgQMQyRXyWG4UPEMER+lRiGDxHDEPlVYhg+RAxD5FeJYfgQMQyRXyWG4UNUwxCJ/CpNDMMf8C9GJPKjNDEMf8C/GJHIj9LEMARBEIRSJT+k5fPJs3jFYwuC4ACeSG7kdVwxDG/g+yZugQwtkh7Tl+5CXscNkmHw9cUt0Mp5ErdnYXKb6uKVGEbyY/rSXcjruGIYWuLiimEI8WD60l3I67hiGFri4ophCPFg+tJdyOu4Yhha4uKKYQjxYPrSXcjruGIYWuLiimEI8WD60l3I67hiGFri4ophCPFg+tJdyOu4Yhha4uKKYQjxYPrSXcjruGIYWuLiimEI8WD60l3I67hiGFri4ophCPFg+tJdyOu4Yhha4uKKYQjxYPrSXcjruGIYWuLiimEI8WD60l3I67hiGFri4ophCPFg+tJdyOu4Yhha4uKKYQjxYPrSXcjruGIYWuLiimEI8WD60l3I67hiGFri4ophCPFg+tJdyOu4Yhha4uKKYQiCIAiCIAiCIJjgjxQXiUQiVZr6bwZ4pUgkEqmyNYyCQkEQyjS6H+SJYQiC4AgxDEEQHCOGIQiCY8QwiuZdXcd5oQLuHxK2GxpdXWw6aeFYVXhFKZDI7bhG1yIten9xYX1VWsCC07RwO/4+6qpKl+TrGssLhTBiGLE5XdcWLfog5fCDHjUoqkXxSFbD2KfrXM28r1T5wTB6a2HTSGTMpEEMIzbjtfABiabRmtURPKnQZHhZcUhGw6DeBSdLizaAoigJw0DwO5BehgViGLGhXz084O2GJVZJhWVXsjI70nTNjaiFZm8YXbRwm+ct6pAbtMI42EYFE2twpC6kWS+P20HrwO1AYm0HSq3DdQyMvKZtQc7TCnsXnCzNbBgYZ7IWvR1qXVGG8Y6uD7TCnooaj+8X/Cy0jkm6mil10suwQAzDHjSJHF21Iu+tkgfh5ZggTvcZrgPb4kHbQNcbkfdqTIyXrYWNC9tQr0c1JKxHQ8N6VF+ljraH1oHJjO/VHhNtx/JIG6vtQNTtaBSpp+2g5MU2IS2cjAjOAdntjywtug7NEj/HJVp4HbQNRFGGgcMJ1djpszfRwvG+YfX4frgW3uZWum5X6rK0aAMRNDGMWOCBrnZL8aCyGpaoSYW/0tgOE88JmPy8q84TlQyCwCTBMto2SprGBS0KobZ8HVSmrkN9j5CJUBlPVsRqO3hXHtvY9c6ytOiYf9ait4HMRu0t8H1BhvF65HUbpR4NQd0ebjBYb7XNCNapsQRNDCMW+Lmpd4Hgr5/a4yCwnap20dXGQXorE4Htea+FD0nwNU946v0g2J7XE5RwfB00z9Ir8t6qDaKWYyy+HnU7eDITOByhNpwszXoZBHswX2nhehrO8HXQezK3XkodgmV8EhWHGhSPDMVq6IG9lVd4YVlHDMMa6spaiSeNmlT0Xh0u2BkGP/gJ1TCoDYrmDUihSL36K8/hvRNC7Q3YbQeibgfFstoOxC5OcQyDfvGxR/Line truncated
|
||||
<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>
|
||||
|
||||
[image8]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAVwAAAGJCAYAAADG5YuwAACAAElEQVR4XuydBXgV19aG23tv7//f+u3f21svFHd3d5fgkgRCcHeXIIFAkKDB3d1dgwa3lhYo7m4tUOT797fCnJ6chJaE5JCTs77nWU8ye/bImVnzztpr75l5CyqVSqVyit5yLFCpVCpV7EiBq1KpVE6SAlelUqmcJAWuSqVSOUkKXJVKpXKSFLgqlUrlJClwVSqVyklS4KpUKpWTpMBVqVQqJ0mBq1KpVE6SAlelUqmcJAWuSqVSOUkKXJVKpXKSFLgqlUrlJClwVSqVyklS4KpUKpWTpMBVqVQqJ0mBq1KpVE6SAlelUqmcJAWuSqVSOUkKXJVKpXKSFLgqlUrlJClwVSqVyklS4KpUKpWTpMB1Ez1//hyzZ89G+vTp8d133yFlypSoWrUqLl686Fg1zmjJkiVInDix7G/r1q3lN/yZqlWrJnXtrXTp0o7VIojHJUOGDFI/VapUmD9/Pq5du+ZYLdaUMWNGpE2bFvfv33ecpYpnUuC6iaZOnYoECRLgrbfestk//vEP5MqVCzdu3HCs/sb1448/IlGiRLZ9/fDDDzFy5EjHauH07bffhvt9tKRJkzpWC6d58+ZFOC4JEyZEvnz58PDhQ8fqsaJ//vOf+Nvf/oY7d+44zlLFMylw3UATJkywwaRGjRro0qULatasKREdy/7973/j3r17jou9MS1YsABJkiSRffP29kabNm3k5kAonTlzxrG66ObNmwIuRsT8fZb9GaQZQVvHpVKlSlK/Vq1atrKvvvoKjx8/dlwsxjVmzBgEBwc7ZVuqNysFbjzXgwcP8MknnwhA5s6dG27ekydP8P3338u8Dh06YOXKlfJ/y5Ytbc3b/v37S1n+/PltyzECZBmb+RSb5O+//76U5cmTBz///LOt+U94spwwIzj5/7vvvovy5cvj9u3btnXaq0yZMlIvderUtjJfX18p27Ztm13NP7R161aZ37x5c+zdu1fsz5rohNvXX38ty/CGZK/Vq1fLtjmvR48eUpY7d26ZXrp0qS2S/u9//4vz589j//79tmMcGBhoW8/169flRsEoe/HixfK7WYfbPXjwoK3exx9/LMfv7t27tjJV/JQCN56LMOBFzmZzZLKASKDwgv/ggw+QLl06LFq0SOYToBZc2MS+cuWKTNMOHz4s0Zk1/fe//13+/u///q9EihTXwzJGp/Z1aFYdRzHFce7cOdmWJQKf6yDMI1NQUFC47dAYuR86dMixqujEiRNS54svvnCcJdq4caPMZ36VKlCggG29b7/9tu13MNVhRd/WfOumwN9BkFrz7Otwu7///rvU+89//iP7qsCN/1LgxnMxIuMFXrhwYcdZoj179sj8ZMmSyXSOHDkEDJ6ennj06JHAwILF8ePHbemJzz77TKJERsisY0W0zIlyPoHEnKQF3G+++caWoyTQWcZtvYp69uwp9QsWLPjSjrPixYtLHUbdY8eORf369WWa0Xhk2rx5s8zPnj274yzRsWPHbMs/ffrUBlwrqqesqHbYsGEy3aRJE5m2ImYLuCzj9ij+Ft68mP5gBEwpcN1HCtx4ru3bt8sFz86xyLRz506Zb6UMrIiVAN63b5/8X7duXfnLPCNHAjCiq1y5sqQk3nvvPZkXmbHZbAGXALfk4eEhQH7ZPtmrVatWsjwjTSsifFVZnWGMlh1l/TaO2ohMR44csR0HQt4CLlMVltKkSYN33nkHFy5ckOk+ffpInUmTJsm0BVymDCwtXLhQ0hX/8z//Y+usVOC6jxS48VzMMRIKbAZbOVNGYhMnTpSOJquTaOjQoTKPEGBdRrT16tWTeadOnRJwMO3AiI+Q2rBhgwCQwGWUy0432m+//SbljH6fPXtmA26jRo1s+8T87asAl5DmsnXq1JF1vUycR4hzP7h9S1YkffLkSbvaYeLvZE6Vv5XHyF6MxFu0aCHLcr2UBdyjR4/a6hGcPLYW0CMDLo+PfdqCKRwFrvtKgRvPxeisUKFCAgJe6JcuXZJOLk5b5ti0z5Ytm20eYUBVr17dVtaxY0dbj3qKFCkEztOmTZNpNueLFi2KEiVKSEoiOsC9deuWdLJxOXbavYoyZ84s9fv27SvTjNwZiX/66acvhXXFihVlGUbCVq6XoyA4Rpkg5jxGwpQCVxUTUuC6gWbOnGmDJWHHiNG+84pDqUJDQ231vby8bPOsHKefn5+tzH60A1MMLCN4CVVrBANBx5RDdIDL/bXG4FaoUAGNGzeW5WkHDhyQOuzU4vSmTZtkmjcBC54sz5s3r0wTqi/TihUrbL+JNyUuxwclrDKORrDSGApcVUxIgesmsi50CyY09rCzM4v/sxPHGkbFIUzWOFhrmNOOHTtsy9k32wlVC240RrvdunXD1atXZT7hzHL7EQnswGMZn65yFDuoPv/883D7aW+zZs2SetZwtQEDBtiWK1myZLi6HGHxV+OLV61aJSkS++V48yhVqlS4Bx+sCNp+1IM1PMwaG9y1a1eZZh6cYqcYp5m6sMT9t5azOs2slI8++BD/pcB1IzG9wIuaTXb+tfKtgwcPRtOmTW312AT/9ddfpZ4V4XFZ5oBfFoWxnPUdx75yeZZzfZY4NpjrigyG9vsYmVn7w3QFp/nXXtZ+RAVejsclshQEfxfnE+yWuC3+Dqs+Ac06VrrFOmb2+8L9t5azRlxwvv20Kv5KgatSqVROkgJXpVKpnCQFrkqlUjlJClyVSqVykhS4KpVK5SQpcFUqlcpJUuCqVCqVk6TAValUKidJgatSqVROkgJXpVKpnCQFrkqlUjlJClyVSqVykhS4KpVK5SQpcFUqlcpJUuCqVCqVk6TAValUKidJgatSqVROkgLXiJ+DuXbtWpw061M1KpXK9aXANfrXv/6Fd/7xT/zj7/+IO/a3f+Cdv78jHyFUqVTxQwpco0/+/QmWHVmMvfdC44wd+HUPNp9bj/f+V4EbX8QPbNauXdtl7JNPPnH8CarXlAIXYcCds3MmtlzYGGds6+XNchN4/1/vO+6uykXFT7+7kpImTepYpHpNKXChwFU5RwpclQIXClyVc6TAVSlwocBVOUcKXJUCF+GBu+vmdoTe3IHQWzuw68a2CCB0lilw458UuCoFLv4AbsilTchdKBcy58iEDFkzoGTlkthxfauUb720OQyE5q/1v9jFsLKQi5sinW8/zTr2da35XL/Mt9+OAjfeSYGrUuAiPHDzl8iH/Q9244enh9BnbC8MnDEA7fu3hVeTmth2dQtqNfeGb2sfiYQJxfl75qK8ZzlMWjcB265sQZ1WPmjQoR52394p8PRtUwc1GlbH9qshmL19BsrXLIvBMwOx49pWAatn4xroMqSTLNt5cCczXRMbzqzFtstbIgXumDFjsHTp0nBlKtfQXwH38OHDSJEiBVq1aoXOnTsjffr0ePDggWO119a9e/fQtm1bx+IIUuDGvBS4CA9cRrf9Jvqj1xg/lKlWGnNDZ5toN73As7V/SwTNGSwgbtW7JTadW48iZQtj373daBvQBu0D2mLapskYNCMQ7cz0kNmDMGrJSCw9vAhjVwQjb9Hc+OHZYfQZ10tAXaVOZcwImQafFrUQbOplz5sN60+tCRfhvvcCuIsXL5Zxke+++67D3qtcRX8F3KNHj6Jdu3a26caNG+PHH3/EtGnTkDlzZlSsWFHKL168iGTJksHT0xN169bFpEmTsGnTJuzYsUNA/fjxY9SsWRNZs2bFyJEjcfPmTWTKlAnp0qXDDz/8gJ49e+LLL7/E9evXbduKTArcmJcCF+GBm6tgTizYO1ds5Y/LsOn8BhQuUwg7r2+DdxNPgeea4yslch27YjQqeHlItLrv/m54mig4VfrUSJc5nUS6y39YigIl8iN3odySRhi1eASSp0oGv5HdMXrZKOQpkhvrTq7G8AVD0SWoE0pXLhkupbDcAPfL/34poH3rrbfEeNE2atRILsYmTZqgefPmaNmyJdq0McBv3x4dO3ZEly5d0L17d/To0QO9evWCv78/AgICEBgYiEGDBiEoKAjDhw/HiBEjEBwcLFHz+PHj5cKdOnUqpk+fjpkzZ2L27NmYN28eFi5cKMBftmwZVq5ciTVr1mDt2rXYsGGDXOghISHYvn07du7cid27d2Pfvn3Yv38/Dh06JBAhNH7++WecPHkSp0+fxrlz53D+/HlcunQJV65cwY0bNwQK/Buf7fvvv3d0vXA6duwYPvroIyRMmBBfffWVnFsqQ4YMEpXyYQQe4759++LgwYO4fPkyGjZsKOeL54DHvnfv3vK3WbNmuH//PgoUKCDnjOf49u3bePLkiRzrTp06OWw9ohS4MS8FLsIDN2/RPNhsIGvlWDeeDYtit18LMVFtc4xYOBQBE/uiVa8WWHNiFfIWyYODv+1Dsx5N0K5/W8zdNdtExbPQZ0xPDJ8/FGtNnWmbJqHjwPZo2bM5TuMkvJt5Yf7uOahsItzZO2bCt3UdjFgwDCUrFRfQ2ke4H77/IUJDQ5E/f34BboIECeQdC7zYaIx2Lly4IBA7e/asAO2XX34RuJ04cUJAxwuZ0CP8jhw5IiA8cOCAGOG4d+9euUi5HV7QhOe2bduwdetWbNmyBZs3b8bGjRsFsOvWrRPg0latWiUX8/LlyyXNsWTJEixatEgAvWDBAsyfPx9z584VcBPgM2bMkGhtypQpmDx5sgB+4sSJmDBhgoB+3Lhx8dq+++47R9cLJ54f3jQpRqKEKcXolseIx5jnkhEqzysBypsvgctztn79egEuz2OdOnXkJsrzRvF8MXrm+eWNztrOn0mBG/NS4MIOuCYKbd2nZTjgbjq3QXK4jDy3XwlBtXpVJe/KEQzMu07eMAnFKxTDhDXjJO/q2agGKvlUlBzuZrNs1bpVJP/LacKadRnpMirm/MqmbruAtrKujoHtbR1okXWaEYDM8RGYKtLine truncated
|
||||
Taking **Qwen3-Omni** as an example:
|
||||
|
||||
[image9]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAOsAAAEWCAYAAACHcWCmAAB5cklEQVR4XuydB3gVVfrG4/5XRXonjYQQ0nvvvffee++kQwKkQOgdpIogiPQmIBBAJXRUFERQV1111+7aV1d33fX9f2duJlwmNw0IIXDe53lz751y5szM+c33nZlzb5SUuLgUSH3cSAz4sxIe//MjUBmrifEq+mQjTFCzwCRNW+hq2cFgkiNM9F1hZugBC2Nv2JgEwM40CA7moXCyjICzVSTcbGLgYZcAT/tEeDsmw9c5FX4u6Qhwy0QgOcg9CyFeOeRchPkUINy3kFyESP8SRAVMRnRgKWKDyxEXUoH40EokhlcjKWKK4OSIqUiJnCY4OaJWcFJ4zS3vE0JrEB8yVXB0UBViAqsQHVCJSL9yRPiWIdS3FCE+kxHsXYJAzyIEeBTC370Afm658HHJFuzllNnqDHg6ppHT4W6fCje7ZLjbJcHdNgGuVrFwIjtaxcDRIpqOQRQcLMjmEbA3DYetaRisTULIwbAyJhsFwtIwABYGfjDX94GpvjdM9DxgrOtOr+4wM/AAnYY/Sc8LF1c7qY8dgSce/RMGPvZ/GK86CVrqppg4noFqA72J9gKoxnoMVHeYG3pRA/RrBTWEQA2Hi3UUXG2iCdR4eDkktYHq75KBAGYGqkc2gj1zCNJ8weG+xYjwKxZAZZDGBJUJoDJIE8KqBFCTIwnQqBrBqdG15OkyR8nMwBVfZRBPR2JYLUE7BXHBU6jMagK2GlH+FYLDfUoFh5FDCNhgjyK6iBTA3zWfQL0JrLdzlmAZsBnwcGBOIWhT4ErQulnHwZXswqC1jIGDZbRgewasWSTszCJgQ8AyW5uEwtIoGBaGgbAiYC0N/GEhAOvZZgsjTw4rV/ekNno4Bg34Pwx54lFoqhlCW8MaOhPsCFQH6Gs7wUjXRYgAFkbe8HKJkIFqFkxRJUSIqK4UUd1s44SI6uOUQqCmUESVQSoDNYtAzaOIWohQ7wKCpVAAVYimAaUEZhUyE6ZRRKVoGiYXSaOYa5ERV4fJ2XOQHlNHnoG0aNF15OnC+9QoGaxJ4bVURg1iKbrmpTYgO7GOQK0iV9IFokwANZRFV69iqlchRft8BLgWEKQ58CZQ/d1zCVCKrI4M2NYIy0ywutkxJ9P+JsDZmlkGrDMB60wR1p4irL15pGBbs3CCNRTWxgzWIHIAweoPawLW1sRXiK43YfXisHJ1T2pjhmMwgTp00OOYoG5MoLLUl4HqQJHUA2UFlQjwIiAdwrBzyzZK+UIJ1hDYUwRxsAiDC0VVd7s4IarKYE2jaJVBziRnERDZCKKoysyADfMpEtLfKAKVwVpTuhDzZ66VgRtW3Zb+soiaQPBWFC7A9s2HCMgZAqyyaNoaXSNkn5mTCNLEsKkCrNEBVTj94nls23QQET7lBGs5XSgmU0SdTOlwCYFaDH+3QoqshXRByRciaVZKDW1nF5KiyimKpsHbKUuIqp5kIaraJ8PFNokgTRRAZVHV0TyGUmKClSKrp0M8XbRiCFSKrCZhAqgsHbYwCIa5QZAQVaOD0jGvfh6lwF4EqhelwF4cVq7uazz1WYcPeQyjhw2iqGoupL762o706ohdW5/F06tWYW7DbJTkV+Hci8dw4vAhikxJKMgox9kXTyI3pZIibDSlvwzUdAFUb+rvHdx7BFs27MSWp3Zidt1yHNx9BG++fhUb12zH4b0nkJM0HYf3HUdN2SLMrV+FG9fewUH6/MqZ13Hj6g2CeLEQMZfO3YRrr7+FF4+cQUHaLJw/fRmXL13Fivlb8NqFN7D72WMoypiNK6+8hS3r96Ewo0lY/rOPP0bDlBUUxauQHj8DV169imc37kU8XQzOtVygZa5i6Zz18HXJw47N+3D9yhWKsJm4ePo0VixZj2kV8/DysZNYsXA9ntu4Azkp1bQ/W/HchuewZM4KHNq1H80HDlCkjCAww3Fs3z4saVqCyIB0PL9zB3Y9s4VS4mBse3oj9j37HF0oEhAVlIrFTQsoU/EmUJkZrN4cVq7uSVNlFMYOHwDVUUMwaYKV0Ec1mOSMEL9YnD5+COtXPonnqMFtWruBGvUWrF22DhtWr0dL8zFkJpehurgRjmZRlOpm3RJRd23dj5NHTqD5hRO4cPoCDuw6is0E7xuvvY6rr72BGZVL8PabbyGM+o8H9x7Hm5evEZgbCaLrOPPSRQJzJooz5+Kl5nOUBs/HX999D9uf2SdAeuPKX+gisgtvvnYNjTWrcebkBYLsNSxbsBmvE8hZKTNx+uQ55KU1ISakChfOvIpkSqk/ePd9PLlkM9at2Ipzpy7g2Q27KeXNweXzr2L7lj0UdYvwwr4XUFU8my4kB+iicJb2eTvt/3O42PIiTp9oxqJZy3D80CHUVc/FC3v2CrDamERj61MbEeydiv3btlMfNoyWW4CmaU2IDqYLmEcc5tbNo0gdicykYoLUp9XeHFau7muC6mgojxgAtdHDoKNlC0MdJwHW1YuXY+HseZhRM4Mi4S4c2L4Dy+Ytxb5tu/HM+o0ESDNB+zSyUythbRSGDaueodQxldLLLAS65xAQW/DBX25gYdNavP7KG3h+z2FcPHMJL+xvFsBjEfcvN24g3KeMouUrePn4Gaxb/hz2bDtEEW+D0B+tKlxMcL+FmbVP4tLpVykq76LtbMPf//oRlhDYyxdsxN7tR/D6xTexZumzlLIvwI0336Z6P4WjB19EZmIDYkOn4pVzr6O8cA5t/1U8vXYH6qcsp4vDm1gyewNF0zxcOHUexw83o2nGciyb/xQBXoKP3n9LiKSXL5xHTfkcvH7hDDau3oiE0AKcOt6M3LQqvEhZhr1pFMEaiYM7dyImNBeHCeDwgDRs37QVM8rrkRabh+KscuSnT6Y+bwQBnQAzfT8yh5Wrh9JWH0ugDoSm8lgY6DhSf8pVeKwQ6hsPw0keBK8nwgNTEOafQpEhGjHB2fBwikNSdCE1zjyE++XAzyWVQKmjVwYq9VE9cnHsUDNFzjdpWib1PStRUTCPUuY6RPoXUzRegHL6XJI1W7hzW1W0kFLZJmTE1mNew3qKSKvpfR2yExtRkb8QCxqfRmZcPXKTZmH29KcoKq9Cbeky6v/RxSJhJvJSZmN+w9PIim9Aac4CSq2XE2ArkBJTT/3EqRRVZ6CBIjB7pJMeT9OCqwigWYgNrkQI9V+jg8pQmDkTb1+9RhlCPnyc81CUU4/IwBIUZs2Al1M64kInU9l1lMoW0r5PpmkpVG4J9VkTqBsQT/tYQOuxm2tJlOJTH9knnfr14UiPK6F+eA4s9APh5xaHMN8U6sP688jK1XPpaSpj/Jgh0FbXbLvzy25+mOixGyA+MDf0haVxIGyp/2VnHkYpXgScrKPhbBMLN1v2XDUJno7J8HFMRUBrVGV3fLOSa6lhFlGaW0TRs5hcgki/ycLNnpjASsQSOAzU+JAaJIWzu7nsEcwMpEezu74NyIyVc1wjvcqcQZ8zYhppuQakRjbQOnVIiahDcjh7rad+bh3BMR2xIdPpwlJLINaQp1BfsgoR/lV0cakkkMqEG06hXqUI9ihBIAEb6F6EIM8i6lsWURqfS33ZXAKPXp1zaN+y4CXcbEqHh32acMPJXbg7nARX20S42JCtE+BiFSN7nGPO7g5Hws6U+rMmYbBij3CM2V3hQFgY+guwyu4IewmPw5Q4rFzdkcEEZWiMpRRYQ1t4nmqq7yE0JNkNEB9qXH6wooZmaxbSCmsUwRoDV5s4YZCAl0My9ftYRGEDILIFWFlk9XFmAyHyCAgGa4ng6IAyMoO1Sg5WdgeXDWyQPX5Ji64XnE7OJCiZs2JnEqgzZZCS0wjUtCgCNaKBIK1HUtgMMgHLoI2YgfiwaRJYp9JFopqiOhskUUEXkzKCspSAJVg9b8Lq5yJ77hroWgA/1zz4UYT1c6J9ccyEt30mvOwJVgasQ6oArPAoxzapDVbn1mevThbRcGSDJcwiKUUOlwFrHCw8xrEwDGiFVTZAgt1kUuKwcnVHhloq0Bw3EgZaxjA1YCNqvIVoam7A7Cc0LjYax848HPYWkXC0YlE1jmCNF545skbr5cBGK2VSfzVbcJBnLkI8CwRQQ3wpqvpORqRvqTA4gY0qig2qFsxgjQtmI4/YCCQG7HQBNhYt0yLrkRHdGkUlkKZGNgpOCmOgNrSaRVf2uY5SVhmoMlhrERU4FdH+UxDl1wqrL0tTWXS9CWuQOxskUSg8d2X2c8kX7Euw+jrlwNuBgG2NrB52FF3t2MimlDZY2bNXJ0qHHS3j2mC1p8hqR7DaGjNYQwSLwJoZyJ63cli5ui2jiarQUhkLI21rgtOTQJWlvuYUUVnaJouqYQKosqgaSw2TQKWo6kYpoAeD1TFNeOzhx2ClPmughwzUMO9i4W5vpH8ZgVompL+yqMpgZa9TEBtYjXgBWDaggUXXGW3ApkaxKNsgmL1PITNIUyJmEpgySBND69ucQPDGE6yxITMQFzSdLEbWmtY0uJLSYDaaiSIrpcChXpPbYA0gWNmzVz8hqlKEJVCZb8KaRWlwJu0vG4KYRqDeCqsIqqNFnDCayc4sWoistqZsRBOLrqFtwAqw6t8cHKHEYeXqjlhknaiqAmMdWwK0NaoSqGZG/kI/i0VVNiKHgcrsSLD6O0YhwDmK0uAkYXSPj1MafF2zCNYsgjXnFlhZ+hvlz6IqS4Erb0ZWgjQyoAbTs9IJ4GkCsAxW+egqAis6uS31pahKsCYyWMkM0gSCNT6URdX6DmEVQG1NgxmkIZ4M1snCIAmWBjNYRWBvwppHaXB2K6xZt8DqaptCoLKBEjejqoN5rDBgxI5AFSKrMPzwJqyy4Yc3YeWRlavbMqbIqjNeEya6DrAwZg/p/ZFs6IEDZq7wNKL+FcHKRuLYsyF1FDGivSOwOF4PJX7mcKSIwm66+DqmI8E/Fb4uORSh8gjWPASz0UreDNLine truncated
|
||||
## 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:
|
||||
|
||||
[image10]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAXYAAAFiCAYAAADiGfPoAAB7FElEQVR4XuydB1gV19aGc2tuv3/KTUwxsSb2rthFELtSFMUKduzYC1aw94rSRFBBUUBFEVEEsffeYktssbfEWKJ+//4WzvF4wIIVcH/Ps55zZs2ePYXDO2vW7PIBtLS0tLQylT6wdGhpaWlpZWxpsGtpaWllMmmwa2lpaWUyabBraWlpZTJpsGtpaWllMmmwa2lpaWUyabBraWlpZTJpsGtpaWllMmmwa2lpaWUyabBraWlpZTJpsGtpaWllMmmwa2lpaWUyabBraWlpZTJpsGtpaWllMmmwa2lpaWUyabBraWlpZTJpsGtpaWllMmmwa2lpaWUyabBraWlpZTJpsGtpaWllMmmwa2lpaWUyabBraWlpZTJpsGtpaWllMmmwa2lpaWUyabBrvbJOnDiBDh06IEuWLPjkk09QpkwZxMTE4N69e5ZF37muXbsmx2dtbY2HDx9artbSyhTSYNd6JY0YMQIffPBBqlaqVCmBfnrS1atXkTdvXhQpUkSDXSvTSoNd66UVGhpqgnixYsWwefNmHDt2DD169DD5K1WqhN9++02i999///2J7R88eJDCd//+/RQ+LtNY3lzG9gQ0zSj3LKUGdvPtUquDx0RLTUZ54zhSk3mdxvlZnotR5mn70dJKizTYtV5Kd+7cQenSpQXe/LQE1ahRo/DHP/5R1q9evRr//ve/5fuyZctkPSFYoEAB8a1bt058BNv//vc/8d29excnT56Eo6Oj6Sbx1Vdf4fDhw6Z9uLm5ib9hw4ZyYzHKtW3bVlIuqckS7KdOncIf/vAH2W7q1KmmOpo2bYpffvkF48aNk+U//elPst4c3tHR0abyxnXYvn27af2NGzfQrl0703pnZ2dUrFhRvk+ePNlULjIy0lTmP//5j1w7nr+W1stKg13rpXTw4EEBEaG4atUqy9USpWfLlk3K9OzZE35+fvLd1tZW1u/cudMEMwKaioqKkuXixYtj3759JuDyCYCAtbKykuXExEQp37FjR1Mdffv2Ffvb3/4my23atDEdi7kswX7mzBnTDejDDz/E2LFj8fXXX5vq5bFw33/+85/leBYvXiz1+Pv7y/r4+HhZDgkJkeWPP/5Y9nHz5k3897//FR/3xzrc3d1N9c6cOVO243FyuXXr1lKGuX8uFy1aVK6hltbLSINd66WUkJAgAGIku2nTJsvVknr57rvvpAyBxuibAP373/8uwDJy8/R98803EvG3atVKfN7e3ujSpYt8b968uUTfBOWWLVvER9BTBti5H0NNmjQRn4uLi8lnrmeBPSAgQMrwBsHl//u//8OFCxfElzt3bvEx/UTdvn0bt27dkjp4bEuWLJH1fHnMbfjymDeCf/7zn3LuhniDYDne6BiVcx9cvnLlitRz9OhR09PNgQMHTNtpaaVFGuxaL6VDhw4JfAivuLg4y9UCPiNinzZtmgDQSEMwHWNnZ4d//etfqFmzpvi2bdsmqZm//vWvcqNwcHAQf2pWvnx52YcBdnt7e9N+27dvL75GjRqZfOZ6Fth37NghZYz0C1/+GqkX4yYVFhZmqisiIkKOkzcr49gMsDPVwmWmiMzTN4ULFxY/I/5Lly7JU4Ll+RnGVI+W1stIg13rpcQI88svvxQAeXh4iI8Aq1evnsCSKQpC+h//+Ad2794t6728vEwgzpo1K/LkyWOK3F1dXSV6zZcvn+S2jdx6586dpT7mws+ePYuff/5Z9k0ZYGeO3ZCR7ngZsDP9QzEdw+WSJUs+FeyEP29qjMC5r9GjR8t6A+w+Pj6yzHM0fyHKfdNPsF++fFnAzqce4xz5yXOk8T2GltbLSINd66VE4BnRMSPvBQsWSPqFIKOPLzr56enpadpm48aNphw4jfBes2aN3AAMX4sWLaRsnz59ZJkpkIsXL4qPTScJxL1798qyAfYGDRqY9vEqYDfqfR7YCeoKFSrIsgHtDRs2yDLBzuPluRrvCJYuXSplmHoxfEzF8Hp9+umnsswUDMWbF1NCQUFBcoPT0noZabBrvbTOnz9vitoJR0bbRqsWw9gZyGgpQrB9//33pnVM0TBqNeBGW758uZQ9fvy4KU3BfRQsWFCgyWWjvncFdvr4Etg4Fp4X00Nc5svT06dPC/D54pU+5tl5/N9++63pPI2Xp0OGDJFlXjeW+fzzz2WZzUR1xK71stJg13plnTt3Dl27dpXWJART9erVpZWIAUOmUwzNmjVLyhByRkqFL0Ppq1y5sqkcRThOmjQJ2bNnl/VsTmjexpwvOeln80ZDvXv3Fh8Bn5quX78uECaYCWge+xdffCHbsKUPxTQKl6tVq2YCe9myZcXHljsUj437582GrXUYsbM5I8vwqcIQmzLmzJlT9sGnGsKb14TXwRCPgWW4LdfzKUZL61Wkwa71RsWUQnrrffo2xDbsQ4cOlTbp5m3bjRYvvBFoab0pabBrab0B8cnAaC3DF6wcR+ezzz6TZUbmT+tApaX1OqTBrqX1hsS27IUKFTLl1Wls/mh0sNLSelPSYNfSeoNip6MffvhB2v3TdEsXrbchDXYtLS2tTCYNdi0tLa1MJg12LS0trUwmDXYtLS2tTCYNdi0tLa1MJg12LS0trUwmDXYtLS2tTCYNdi0tLa1MJg12LS0trUwmDXYtLS2tTCYNdi0tLa1MJg12LS0trUwmDXYtLS2tTCYNdi0tLa1MJg12LS0trUymNw72n3/+GQcOHMDhw4fTrZ06dcrysLW0tLQyrN442Ldt24YCBQogd+7c6c9y5ZaJkhMSEiwPW0tLSyvD6o2DfdeuXciZKydGBAzDpAUTMDFsfLqxyeGT4NzaGevWrbM8bK10pNu3b8t0cosXL053xtmRtLTSm94K2HPkzAGfxdOw4XwS1p5Zk25s06X1cGjuoMGeznXmzBm0b99eAP/rr7+mG7ty5QpGjx5tebhaWu9cGuwa7OleBHubNm0s3e9cvNFosGulR2mwa7Cne6VXsP/2228a7FrpUhrsGuzpXhrsWlppkwa7Bnu6lwa7llba9NbBvu5cIjZcSML682tlOelsQgrgvi3TYM8Y0mDX0kqb3irYN15ch4FTPNGqR0ssO7AU/cb1wZDpg5FwKh7L1fKKQ8uw+sRKLNkThVXHYwX6SecSsOLwMsQeXi43BQKZ62NUWa5fdSxWLUdi5dEYKRt/Mg5L90YhUdXJ5dXHWV+k1JtwarXaLhrR+xYj8XR8qmDnC7HAwEDEx8ebnYXWu9TLgD0kJAQ5c+bEZ599hhw5cqBjx46WRZ4QW7i4uLhYup8pDXat9Kp3AvZeo3ti373dGDV7BDoP6giPoV1QuqIVBk31RHXHarBvXBe2tSoL0BdsCkXdRnVQsWoFLNq2AF29uqKWc02UrmAlkHZs6oD6bvVQt3FtKd+oXSNUc6wq++HNoUptW9m+qr0d5q4NwbfZvoVr52aI/zEOmy9vELDzGO/evQtfX198//33+NOf/oStW7danorWO9LLgN1QlSpVTN8J+759++LYsWMIDw/HzZs3xT9//ny5mfMGcP/+fVP550mDXSu96q2DfXiAt3RYKlKyCKxrWCN6/2J0GdwJg6cNkvVNOjTG1qsb0bZ3Gwzz90KP4d0wOXwiVv4QgzmJs1FZbcPIfXrkVFXnVNRr7iTRNyG/4nA0qjtUk2h9xZFl0gmp04AO2HFzK5p3aga/5TNRzaGqKdfPiL1eCye0bNkSefLkwQcffCBGGBw8eBA7d+7E7t27sWfPHuzfv198HILg6NGjOH78OE6ePImffvoJp0+fxtmzZ3H+/HlcuHABly9flgjw+vXruHHjBm7duiUQeJrxKeFpdufOnWcab0hPs3v37qVqv//++1ONYHuaPXjw4Kn28OHDZ9qr6Ny5cy8N9kqVKsm1iIqKQpcuXZCUlIQKFSrI37B3795YsmQJpk+fLh2N6E/LsWqwa6VXvXWwS8Q+qid2/7YD265vxrqfE9F9mAcmhU9A/Ik4tOjmho2X1sFzUn94zRiCVt1bwCdqGrZe24SI7QtRupIV7JvYo55rPSzbvwQzlk5HdadqaNaxCVafjMP8DfMkOh8ROAwj1E2kx4ju2Hp9k9pnD8xc6gOXNg0lv2+AvX6LeqhXrx4++ugjE9hz5cqFfv36oUePHujWrRs8PDzQuXNndOrUSTrKtG3bFq1bt0arVq3QokULuLq6olmzZmjSpAkaNWqEhg0bwtnZWep1dHRE3bp1YW9vL+bg4CBGP83JyUmMZevXry/WoEEDMdZDY4qAxrpp3I9hTZs2FeP+mzdvLsbjcXNzE+Px0XisNB43jaCk8Vxo7dq1g7u7uxjPsUOHDmJMYfC8abwGNAKya9euYrw2NF6n7t27yzWj9ezZU6xXr15ihGifPn0kYjasf//+Yp6enmIDBgwQGzRoEAYOHCifgwcPlv3w+ryMCHbezLhvXq/x48fLcBKUv78/6tSpI995U65Ro4b5ps+VBrtWetU7AXtnFaEb+XLmwbt5d5Xomi9VrWtUknw44cwIu8+Y3ug/sR96j+mFoT6DUc2xmmw7ft5YTI+YqsDdDbtubYd7v7YCbkb6i7YuUJ+tVYQfLPWs+XEVKtpVQHB8EBq0cn4C7EzFbNq0SaJuAurvf/87smbNikuXLlmeyjuXEf0yQmYEbUTZBBejUkbwRqRP6PBJgcZekr/88oukHmh8iuDTxLVr18SuXr0qTxh80uB5X7x4UYxPHwQeB3KjMXLmkwlTIzReMw6gRuOTy48//ihPMSdOnJAnGhrTHoyOGRHTjhw5Ik89hw4dEuNTEAeJo/GpaN++fdi7d688JdH4xBQbGys3q5eRAXYvLy9MnDgRGzZskCcxasyYMXKz5fXi+Wqwa2UWvV2wX1iHKQsnYeSs4abWMPwcEzwSs+ICJHof5jcUFezKq4i9n9wIWLine truncated
|
||||
```
|
||||
# Create an omni_lm with HF model name.
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
[image11]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAABM8AAARlCAYAAACukOBGAACAAElEQVR4Xuy9B7hV13WuHSd5ch3Zia17rcQl13+K88dyXOLYlmPLRb13CVUkUEMIEB3Rkeiig+iiI3rvvffeO4gmeu+dM+56xz4TbRZ7IxDt7M33Ps98zll97rXGnGvOb4055l+YEEIIIYQQQgghhBAiJX8RXyGEEEIIIYQQQgghhEgg8UwIIYQQQgghhBBCiDRIPBNCCCGEEEIIIYQQIg0Sz4QQQgghhBBCCCGESIPEMyGEEEIIIYQQQggh0iDxTAghhBBCCCGEEEKINEg8E0IIIYQQQgghhBAiDRLPhBBCCCGEEEIIIYRIg8QzIYQQQgghhBBCCCHSIPFMCCGEEEIIIYQQQog0SDwTQgghhBBCCCGEECINEs+EuAhnz5jlKCnF09k0Kb49/B9fH98/1bazSfaXYhvr025PPm98e3xd8nJ826VcK36eVNeN/e/nih+Xajl+vaT1550jvu+lbIv/TfV/fDn5nPHz5yYhhBBCCCFE9iHxTIh05JgdP5Bjx3YrKcXSntTpeG5KXk61/lzafZFtsXPE07n1KfJz7pwp8nvB9ZL2uWBb8j7ptuWm+DWT973geqmOS943zfW+7Den+j++HD9Hym25/8eXk895wfmj8506mhOvRYQQQgghhBBZgMQzIdKBeLY/x45sz7GjO5SUvkhH0qT49vB/fH18/1TbLnbclx2bvC2+Pb4ueTm+Lb5Pqm2p9onvG///Upfj10ten25b/P/4cvzcqbalukb8nBecf7vEMyGEEEIIIbIViWdCXIQT+y3ROd6ppKSklD5RT0g8E0IIIYQQIjuReCZEOnIQz3ISHigpOstKSkpKIUk8E0IIIYQQInuReCZEOsKwTYlnSkpKX5IkngkhhBBCCJG9SDwTIh3yPFNSUrrEJPFMCCGEEEKI7EXimRDpkOeZkpLSJSaJZ0IIIYQQQmQvEs+ESIfEMyUlpUtMEs+EEEIIIYTIXiSeCZEOiWdKSkqXmCSeCSGEEEIIkb1IPBPiIpzYbxLPbmA6vjt6CEej57Dni+dgh83OHkgsu2CxL9rnpPlflk9wzIlon4Nf7HNyb2IfO56bonMe22UXXC+euD7nuZR9SX79KK/k8cuOYd/TkX3lHJKNZUOSeCaEEEIIIUT2IvFMiHRowoAbmhC8tq3bby0afWKfLd3mQhPrenToZ2MGT3ax7PRBs3lTllmJImVsxbz1dvpQ4pjypSrb8H7jfDv7LZu91t59s5gVzk1Vy9ew3RtQuOycoIZQhuCFsBYEtu1r91vfbkNs/+fsYC7K2ZGEKOYiGeuOJY4NQt6GKK9d2va0PZuOfnGu4wkB0MXAIOBFl58zabGNHDDBBTpSOD/58GXyd+ILYZB74PvkCm5nDnyxHL9/Stc3STwTQgghhBAie5F4JkQ6NGzzhiaEorWLt9hD9z9mcycvdVGJde++Ucwa1m7m4hMi1IAew+373/uBNW/Y1h/bkN6jfblRneYuPuHZNbDHCLvvrgds5viFtmDaUlsya5Xt2XjE5kxaZD079bcB3YfZwa2nXOBCdOvZaUC03zL7fNVuGzFgvK1bssUmjphhndv0sCmjZrvnG/v36NjPBvYc4fk9sSfHBa/5U5ZZwfxv2cblO23RjOXWu8tA69NlkO3//IQd3n7Gl7u172N7N+2zSmWq2aMPPm57Nh2wxTNXW6fW3WzCsOkuni2ZtSbK93Dr3qFvlI89djr6vRuWbY/y0N1mTljowtyaRZs9T3OnLPU8xe+h0vVLEs+EEEIIIYTIXiSeCXERNGzzxiWEMjzO8j39ki2YtsLFMtaVKVbeWjRs66IYAtKwvmPtyUeftVJFy9mhHcetecM2lu+pF619y09dcEPQGtpnjL30XH4XoXatP2CHo/3WR+f+nzvudE+0P995t/XpOtjGD5saXe9Fe+O1QvbgfY/YzAkLrGzxCtamWQe741f/YyXeLWP3/Pk+Wz5njbVq3N5KFStnBV9509o07eBeYORp0fSV5zzh7r/nQXvztXeiY+73/HRq3d3ee6ekVS73gQtolcpWs8ceftKFOfLfpF5Le+bxfDZ93Hx75YUC9nz025945Gmr8n4N2711X/S3ur1fspIVi84xacxUK/1eeav9QT17KV9+mzNxkZ2WgHbDksQzIYQQQgghsheJZ0JcBIlnNy5dqng2uNcoK1mkjL1fqrIN6D7CalerbzUq17HWTTu4eMZ+YwZNstv/4z8t/wsFrMDLr1vLxm1s7aLP7bmnXrAd6w5Yp1bdrH7NJlauRAUfErp1zT4r9EYRmzZmrpUrWTHa/xN7r3ApP1fZaJ9P2/exJx992lo1aW8fVKhlTz/+nB3YctLzE8SzZXPWuaj12dKt1rfrYKsV5avtx52sepS3kQNG277NB21g9+HWqM7HdmLP2SjvQ21Qr5H2wrMvW99Ph1iRQiVsanT9GePmW8WyVaPtw/28eNMd2XPUOrfuYX+68y4b2meYPfLA4+5pR/7i91Hp+iSJZ0IIIYQQQmQvEs+ESIeGbd7QhFC2bulWe+HZV1zock4kxLMOLbuee0yDeo6wmlU/cq8uvLVqVatnrZu0TwzjzPU8Q2B7u0DhRKyxXH1j3eIt7nW2c/2h6HyfutcXItzEETNt+7oDVrRQiYR4VqKifdK8k4tfDKdkqGW39n3tuadfsJpVPrKPG7T0oZOHtp2xnMNfiGcMu+TvphU7rVfHAVbnw4Z2cNux6Nge7s3W9ZNePuyzWb1WNnviYnv5+VetWf1W9uC9D/tQ0MrlPrSF01fYhOHTo99X14eSFn+39Lk4bd2jPNz9p/tcCKxfs7FNGjnTTu298D4qXZ8k8UwIIYQQQojsReKZEOmQeHZDE+LZ+mXb7Q+/v8uFq95detnMCXOtyFvF7eV8r1qvzr1sxMARLnzh/TV2yBS77Tu32bC+Y6xFo7ZWv1ZT9wTDGwvx7Df//VsX2Lq1722Deg6xuZOXWIGX37Bta/db8wZtrHHdFja490h7q0Bhq/NBA3vmiXw2bew8K/p2cR8KWrFMNc8XotqwfmPdY6xhnWY+bJLjT+xNzARKzLO3Cha2hdNX+vDPz5ZstS5te9iHlWpbuxado3O19mM/ad7Z+n06xAq88oa1b9HFr9e/21B75IHHXIwrWbSsTygwatBEK1uivK2cv87PhyhYqlhZ69Wlr4tzvToPsBLvlrbJo2b5PYvfR6XrkySeCSGEEEIIkb1IPBMiHRLPbmg6vjvH9m85YZ+262N1PmxgtarWs8G9R3nA/no1Gvsyccfw2po1YaHt3njExadd6w/ZwmkrfBbOk/uYtTPHA+03rdfSaler555pzeq3ttULNrlXF9fAW2zBtOW2d9Mxa1i7oRV45U2PZcZxU0fPskUzVvrwSfJFfLItq/b49Zi4oEGtprZpxS47tS8xacC2Nfv8vMzUOX7YNNuz8agtn7POxbrN0X51qzf0vByIrrt55S6Pp8YQT2boRFQbPWiSX2/K6Nm2dfVej82GMMaEALMmLvJhqUP6jLacI4nZOvk9vToNiOz0rN+z+H1Uuj5J4pkQQgghhBDZi8QzIdIh8eyGJ4ZJMlTTTuemI5aYZTMsn0gMy8yJ0ond5sMZmTGT4P2ITTw70sm9SceQTpqd2pc4F9dgf87F7JYli5axd94o4gIY52afMBkAeWKZ83EdO5VILIdr+frDudujv8d3J2K1cS7PB8ecTKz35ZOJ7f47o//5jZ6fQ4lznGKfQ4nze5y33PtwlGXyx/KxxO+I3z+l65ckngkhhBBCCJG9SDwTIh0Sz266hACFCIewpiGQSpeTJJ4JIYQQQgiRvUg8EyIdEs9uyoSAJi8upctNEs+EEEIIIYTIXiSeCZEOiWdKSkqXmCSeCSGEEEIIkb1IPBMiHRLPlJSULjFJPBNCCCGEECJ7kXgmRDoknikpKV1ikngmhBBCCCFE9iLxTIh0SDxTUlK6xCTxTAghhBBCiOxF4pkQ6ZB4pqSkdIlJ4pkQQgghhBDZi8QzIdIh8UxJKU8kZj89nsdnQJV4JoQQQgghRPYi8UyIdCCe7Ys6xdsTHWMlJaXrn45G6cRus5N77YJteSpF9cSpIxLPhBBCCCGEyEYkngmRhpyoH4wnCR1iJaXz0lG7cDms+7JtydtTLce3pTrvxY5LtRzfFj82+RrxY+P7xrelOzbVvun+j58nadvZE2Zd2/ewMsUr2IlDZ+x0/Nj4NeLnSZe/VMelOybVtvj/UTpzQuKZEEIIIYQQ2YjEMyGEEHmacePHWdu2bS0n52x8kxBCCCGEEEJccySeCSGEEEIIIYQQQgiRBolnQgghhBBCCCGEEEKkQeKZEEIIIYQQQgghhBBpkHgmhBAiTzNixAj76KOP7OxZxTwTQgghhBBCXH8kngkhhMjTFC9e3G699VY7ffp0fJMQQgghhBBCXHMkngkhhMjTVKxY0f7t3/7NTp06Fd8khBBCCCGEENcciWdCCCHyNJUqVbJ//dd/lXgmhBBCCCGEuCFIPBNCCJGnqVy5sv3Lv/yLxDMhhBBCCCHEDUHimRBCiDwNwzZ/+MMfSjwTQgghhBBC3BAkngkhhMjT4Hkm8UwIIYQQQghxo5B4JoQQIk+za9cuW7NmjeXk5MQ3CSGEEEIIIcQ1R+KZEEIIIYQQQgghhBBpkHgmhBBCCCGEEEIIIUQaJJ4JIYQQQgghhBBCCJEGiWdCCCHyNK1bt7Z77rlHEwYIIYQQQgghbggSz4QQQuRpqlSpYn//939vJ0+ejG8SQgghhBBCiGuOxDMhhBB5mpo1a9p3v/tdiWdCCCGEEEKIG4LEMyGEEHmaWrVq2T/+4z9KPBNCCCGEEELcECSeCSGEyNMgnv3DP/yDxDMhhBBCCCHEDUHimRBCiDxNtWrV7JZbbpF4JoQQQgghhLghSLine truncated
|
||||
omni_lm = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
|
||||
|
||||
# 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_lm.generate(om_inputs, sampling_params_list)
|
||||
```
|
||||
|
||||
## 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:
|
||||
|
||||
```
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||
```
|
||||
|
||||
Users can send requests to the server using curl:
|
||||
|
||||
```
|
||||
# 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](../user_guide/examples/).
|
||||
@@ -1 +1,200 @@
|
||||
## OmniConnector for fully disaggregated serving (update soon)
|
||||
# OmniConnector: a unified connector interface of disaggregated inference for omni-modality models
|
||||
|
||||
This guide explains how to configure and use distributed connectors (vllm_omni/distributed/connectors) in vllm-omni for multi-stage pipelines.
|
||||
|
||||
## 1. Overview
|
||||
|
||||
Connectors enable data transfer between pipeline stages (e.g., Thinker -> Talker).
|
||||
Currently supported connectors operate in **D2H2D (Device-to-Host-to-Device)** mode:
|
||||
1. **SharedMemoryConnector**: Uses system shared memory.
|
||||
2. **MooncakeConnector**: Uses [Mooncake](https://github.com/kvcache-ai/Mooncake).
|
||||
|
||||
* **SharedMemoryConnector (Default)**: Zero-copy (on host), lowest latency. Best for **single-node** deployments. Auto-configured if no connectors are specified.
|
||||
* **MooncakeConnector**: TCP/RDMA based. Best for **multi-node** distributed deployments. Requires a Mooncake Master service.
|
||||
|
||||
## 2. API Design
|
||||
|
||||
The connector system is built around the `OmniConnectorBase` abstraction, which decouples data transport from stage logic.
|
||||
|
||||
### Core Interface
|
||||
|
||||
```python
|
||||
class OmniConnectorBase(ABC):
|
||||
@abstractmethod
|
||||
def put(self, from_stage: str, to_stage: str, request_id: str, data: Any) -> tuple[bool, int, Optional[dict]]:
|
||||
"""
|
||||
Store data.
|
||||
Returns: (success, serialized_size, metadata)
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get(self, from_stage: str, to_stage: str, request_id: str, metadata: Optional[dict] = None) -> Optional[tuple[Any, int]]:
|
||||
"""
|
||||
Retrieve data.
|
||||
Args: metadata - Transport-specific handles returned by put() (e.g., SHM name).
|
||||
Returns: (object, serialized_size)
|
||||
"""
|
||||
pass
|
||||
```
|
||||
|
||||
### Key Concept: Metadata Passing
|
||||
Unlike a pure key-value store, some connectors (like `SharedMemoryConnector`) generate transient resources (e.g., a shared memory block name) during `put()`. This `metadata` **must be passed via the control plane** (e.g., HTTP headers, queue messages) from the producer stage to the consumer stage so `get()` can locate the data.
|
||||
|
||||
## 3. Backends & Use Cases
|
||||
|
||||
### 3.1 SharedMemoryConnector
|
||||
**Best for:** Single-node, high-performance IPC.
|
||||
|
||||
* **Mechanism:**
|
||||
* **Small Payloads (< Threshold)**: Data is serialized and passed directly "inline" within the `metadata` dictionary. This avoids the overhead of creating SHM blocks for tiny messages.
|
||||
* **Large Payloads (>= Threshold)**: Data is written to a named System Shared Memory block. The block name is returned in `metadata`.
|
||||
* **Configuration:**
|
||||
* `shm_threshold_bytes`: Size in bytes to switch from inline to SHM (default: 64KB).
|
||||
|
||||
### 3.2 MooncakeConnector
|
||||
**Best for:** Multi-node distributed inference.
|
||||
|
||||
* **Mechanism:** Uses Mooncake's distributed KVCache store.
|
||||
* **Data Plane**: TCP or RDMA for high-bandwidth transfer.
|
||||
* **Control Plane**: Uses a centralized Mooncake Master and Metadata Server.
|
||||
* **Keying**: Deterministic keys based on `request_id/from_stage_to_stage`.
|
||||
* **Requirements**: Requires a running Mooncake Master service.
|
||||
|
||||
## 4. Relationship with vLLM
|
||||
|
||||
vLLM provides specialized distributed mechanisms for specific artifacts:
|
||||
* **KV Transfer** (`vllm.distributed.kv_transfer`): Optimized for transferring KV caches between prefill and decode instances (using NCCL, Mooncake, etc.).
|
||||
* **EC Transfer** (`vllm.distributed.ec_transfer`): Optimized for sharing encoder embeddings.
|
||||
* **Device Communicators** (`vllm.distributed.device_communicators`): Low-level primitives (NCCL, SHM) for Tensor/Pipeline Parallelism.
|
||||
|
||||
`vllm-omni` complements this by introducing a **Generalized Connector Abstraction** (`OmniConnector`) for multimodal pipelines. While vLLM's connectors are artifact-specific, `vllm-omni`:
|
||||
|
||||
1. **Unifies Transport**: Provides a single API (`put`/`get`) to transport *any* stage artifact (Input Embeddings, Hidden States, Audio/Image Tensors, KV Cache, Final Output) between arbitrary pipeline stages (e.g., AudioEncoder -> LLM -> AudioGenerator).
|
||||
2. **Extends Connectivity**: Enables flexible construction of complex DAGs (Directed Acyclic Graphs) where stages can run in the same process, same node, or across nodes, using the most appropriate backend (SHM, Mooncake, etc.) for each edge.
|
||||
3. **Wraps & Adapts**: Can internally utilize vLLM's specialized `kv_transfer` for KV paths while using generic transports (SHM/Mooncake) for other data types, presenting a consistent interface to the application layer.
|
||||
|
||||
## 5. Installation (Mooncake)
|
||||
|
||||
If using `MooncakeConnector`, install the library first:
|
||||
|
||||
```bash
|
||||
# For CUDA-enabled systems (Recommended)
|
||||
pip install mooncake-transfer-engine
|
||||
|
||||
# For non-CUDA systems
|
||||
pip install mooncake-transfer-engine-non-cuda
|
||||
```
|
||||
|
||||
## 6. Using MooncakeConnector
|
||||
|
||||
### 6.1 Start Mooncake Master
|
||||
|
||||
Start the master service on your primary node:
|
||||
|
||||
```bash
|
||||
# if you use mooncake SSD storage
|
||||
mkdir -p ./mc_storage
|
||||
|
||||
mooncake_master \
|
||||
--rpc_port=50051 \
|
||||
--enable_http_metadata_server=true \
|
||||
--http_metadata_server_host=0.0.0.0 \
|
||||
--http_metadata_server_port=8080 \
|
||||
--metrics_port=9003 \
|
||||
--root_fs_dir=./mc_storage/ \
|
||||
--cluster_id=mc-local-1 &
|
||||
```
|
||||
|
||||
### 6.2 Configuration (YAML)
|
||||
|
||||
Edit your stage config (e.g., `qwen2_5_omni.yaml`).
|
||||
|
||||
**Step 1: Define Connector in Global Runtime**
|
||||
|
||||
```yaml
|
||||
runtime:
|
||||
connectors:
|
||||
connector_of_mooncake:
|
||||
name: MooncakeConnector
|
||||
extra:
|
||||
host: "127.0.0.1" # Local Worker IP
|
||||
metadata_server: "http://<MASTER_IP>:8080/metadata"
|
||||
master: "<MASTER_IP>:50051"
|
||||
segment: 512000000 # 512MB segment
|
||||
localbuf: 64000000 # 64MB buffer
|
||||
proto: "tcp" # "tcp" or "rdma"
|
||||
```
|
||||
|
||||
**Mooncake Configuration Parameters:**
|
||||
|
||||
* **host**: The hostname or IP address of the local machine (worker). Mooncake uses this to register itself in the metadata server so other nodes can find it.
|
||||
* **metadata_server**: The URL of the metadata server. This is used for service discovery and connection establishment (e.g., exchanging QP information for RDMA).
|
||||
* **master**: The address of the Mooncake Master Server (e.g., `<MASTER_IP>:50051`). This is used for global state management and control plane operations.
|
||||
* **segment**: The size of the global memory segment in bytes (default: ~512MB). This defines the shared memory region accessible by Mooncake for data transfer.
|
||||
* **localbuf**: The size of the local buffer in bytes (default: ~64MB). Used for local data buffering during transfer operations.
|
||||
* **proto**: The transport protocol to use. Options:
|
||||
* `tcp`: Standard TCP/IP (easier setup, universal compatibility).
|
||||
* `rdma`: Remote Direct Memory Access (higher performance, requires RDMA-capable hardware).
|
||||
|
||||
For more details, refer to the [Mooncake Repository](https://github.com/kvcache-ai/Mooncake).
|
||||
|
||||
**Step 2: Reference in Stages**
|
||||
|
||||
Explicitly link stages using `input_connectors` and `output_connectors`:
|
||||
|
||||
```yaml
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
# ...
|
||||
output_connectors:
|
||||
to_stage_1: connector_of_mooncake
|
||||
|
||||
- stage_id: 1
|
||||
# ...
|
||||
input_connectors:
|
||||
from_stage_0: connector_of_mooncake
|
||||
```
|
||||
|
||||
## 7. Using SharedMemoryConnector (Auto-Mode)
|
||||
|
||||
**Best for single-node.**
|
||||
|
||||
The system will automatically create `SharedMemoryConnector`s for any pipeline edge that does not have an explicit connector defined. This is inferred from:
|
||||
1. `runtime.edges` list in the config.
|
||||
2. `engine_input_source` dependencies defined in `stage_args`.
|
||||
|
||||
### Threshold Configuration
|
||||
By default, payloads larger than **64KB** (default threshold) are transferred via shared memory, while smaller ones use the control queue (inline).
|
||||
|
||||
To adjust this threshold (e.g., to 1GB), add the following to your `runtime.connectors`:
|
||||
|
||||
```yaml
|
||||
runtime:
|
||||
connectors:
|
||||
connector_of_shared_memory:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
shm_threshold_bytes: 1024 # 1KB threshold
|
||||
```
|
||||
|
||||
## 8. Summary
|
||||
|
||||
| Use Case | Recommended Connector | Configuration |
|
||||
| :--- | :--- | :--- |
|
||||
| **Single Node** | `SharedMemoryConnector` | **None** (Automatic) or Custom Threshold |
|
||||
| **Multi Node** | `MooncakeConnector` | Explicit YAML + Mooncake Master |
|
||||
|
||||
## 9. Operational Notes (important)
|
||||
|
||||
- **Fail-fast config validation**: the loader raises if any expected edge is missing a connector. Define `input_connectors`/`output_connectors` or rely on auto-SHM filling; otherwise startup aborts.
|
||||
- **Missing payloads halt the stage**: workers expect connector payloads; if metadata or connector config is missing, the stage raises and stops. Verify connector wiring and metadata propagation before production.
|
||||
|
||||
## 10. Future Roadmap: Device-to-Device (D2D) Transport
|
||||
|
||||
The current implementations (`SharedMemoryConnector`, `MooncakeConnector`) utilize a **D2H2D (Device-to-Host-to-Device)** data path. Tensors are moved to CPU memory (Host) for transport, which incurs PCIe overhead.
|
||||
|
||||
As outlined in the design RFC, future versions will introduce **D2D (Device-to-Device)** connectors:
|
||||
|
||||
* **Goal**: Direct GPU-to-GPU transfer (via NCCL, UCX, or IPC) to minimize latency for large tensor payloads.
|
||||
* **Mechanism**: The `OmniConnector` API allows `put()` to initiate a transfer and return a lightweight handle (metadata) via the control plane, while the heavy payload flows directly between devices.
|
||||
@@ -0,0 +1,61 @@
|
||||
# VLLM-Omni Distributed Utilities
|
||||
|
||||
This directory (vllm_omni/distributed/ray_utils) contains utilities for distributed execution in vllm-omni, supporting both **Ray** and **Multiprocessing** backends.
|
||||
|
||||
## 1. Ray Utils
|
||||
|
||||
The `ray_utils` module provides helper functions for managing Ray clusters and actors, which is essential for:
|
||||
* **Multi-node deployment**: Running pipeline stages across different physical machines.
|
||||
* **Resource management**: Efficient GPU/CPU allocation.
|
||||
|
||||
### 1.1 Basic Usage
|
||||
|
||||
To use the Ray backend, specify `worker_backend="ray"` when initializing the engine.
|
||||
|
||||
**Command Line Example:**
|
||||
```bash
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--worker-backend ray \
|
||||
--ray-address auto
|
||||
```
|
||||
|
||||
### 1.2 Cluster Setup
|
||||
|
||||
**Step 1: Start Head Node**
|
||||
Run this on your primary machine:
|
||||
```bash
|
||||
ray start --head --port=6399
|
||||
```
|
||||
|
||||
**Step 2: Connect Worker Nodes**
|
||||
Run this on each worker machine:
|
||||
```bash
|
||||
ray start --address=<HEAD_NODE_IP>:6399
|
||||
```
|
||||
|
||||
> **Tip**: For a complete cluster setup script, refer to the vLLM example:
|
||||
> [run_cluster.sh](https://github.com/vllm-project/vllm/blob/main/examples/online_serving/run_cluster.sh)
|
||||
|
||||
### 1.3 Distributed Connector Support
|
||||
|
||||
When running on Ray, the system automatically adapts its communication strategy:
|
||||
|
||||
* **Cross-Node**: Recommended to use `MooncakeConnector` (requires separate configuration).
|
||||
* **Same-Node**: Can still use `SharedMemoryConnector` for efficiency, or Ray's native object store (plasma).
|
||||
* **SHM threshold default differs**: when `worker_backend="ray"`, the SharedMemoryConnector default threshold is set to `sys.maxsize`, which forces payloads to go inline (no SHM). Override `shm_threshold_bytes` in the connector config if you want SHM for Ray runs.
|
||||
|
||||
### 1.4 Internal Helpers
|
||||
|
||||
* **`initialize_ray_cluster`**: Connects to an existing Ray cluster or starts a local one.
|
||||
|
||||
## 2. Troubleshooting
|
||||
|
||||
* **Connection Issues**: Ensure the Ray head node is accessible and ports (default 6399 in this example) are open.
|
||||
* **Version Mismatch**: Ensure all nodes run the same version of Ray and Python.
|
||||
|
||||
### Installation
|
||||
```bash
|
||||
pip install "ray[default]"
|
||||
```
|
||||
@@ -5,12 +5,12 @@ This section contains design documents and architecture specifications for vLLM-
|
||||
## Architecture Documents
|
||||
|
||||
- [Architecture Overview](architecture_overview.md)
|
||||
- [vLLM-Omni Software Design Document](vllm_omni_design.md)
|
||||
|
||||
## Feature Design Documents
|
||||
|
||||
- [Multi-Request Streaming](feature/multi_request_streaming_design.md)
|
||||
- [Omni Connector Design](feature/omni_connector_design.md)
|
||||
- [Ray-based Execution](feature/ray_based_execution.md)
|
||||
|
||||
## Module Design Documents
|
||||
|
||||
|
||||
@@ -1,509 +0,0 @@
|
||||
# Architecture Design Principles
|
||||
|
||||
## Overview
|
||||
|
||||
vLLM-Omni is a multi-modality extension for vLLM that supports non-autoregressive structures and non-textual outputs. This document outlines the key software abstractions, APIs, and dependencies for the system, designed to maximize reuse of vLLM's proven architecture.
|
||||
|
||||
## Architecture Principles
|
||||
|
||||
1. **vLLM V1 Compatibility**: Built on vLLM's Engine V1 architecture with AsyncLLM and EngineCore patterns
|
||||
2. **Stage-based Processing**: Models are divided into multiple stages, each processed by different Engine Cores
|
||||
3. **Multiple Engine Core Support**: Each stage can use either AR Engine Core (reusing vLLM) or Diffusion Engine Core (new DiT support), or other new Cores
|
||||
4. **Worker Process Pattern**: Follows vLLM's multiprocess worker architecture for scalability
|
||||
5. **Extensibility**: Easy integration of new modalities, model architectures, and output formats
|
||||
|
||||
## Key Data Flow
|
||||
|
||||
API Server --> OmniLLM/AsyncOmni (New, including multi engines) --> LLMEngine/AsyncLLM --> Engine Core
|
||||
--> Scheduler (New one for DiT) --> Executor (New one for diffusers) --> Worker (New one for DiT)
|
||||
--> ModelRunner (New one for AR hiddenstate, New one for DiT) --> RequestState --> OutputProcessor (New one for final multimodal output)
|
||||
|
||||
## Core Components (Ordered by Data Flow)
|
||||
|
||||
### 1. Installation
|
||||
```bash
|
||||
pip install vllm
|
||||
pip install vllm-omni
|
||||
```
|
||||
|
||||
### 2. Online inference launch (Entry Point)
|
||||
|
||||
Keep consistency with the vllm main branch
|
||||
```bash
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8000
|
||||
```
|
||||
|
||||
#### Design Architecture
|
||||
```mermaid
|
||||
graph TD
|
||||
A[vllm serve model --omni] --> B[vLLM-Omni CLI Wrapper]
|
||||
B --> C{Detect --omni flag}
|
||||
C -->|Yes| D[Parse OmniConfig]
|
||||
C -->|No| E[Forward to vLLM CLI]
|
||||
D --> F[Initialize AsyncOmni]
|
||||
F --> G[Start omni Server]
|
||||
G --> H[Multi-stage Processing]
|
||||
E --> I[Standard vLLM Pipeline]
|
||||
|
||||
subgraph "vLLM-Omni Components"
|
||||
F
|
||||
G
|
||||
H
|
||||
end
|
||||
|
||||
subgraph "vLLM Components"
|
||||
I
|
||||
end
|
||||
```
|
||||
|
||||
#### omni Serve Command Implementation
|
||||
File: vllm_omni/entrypoints/cli/main.py
|
||||
```python
|
||||
import sys
|
||||
import argparse
|
||||
from typing import List, Optional
|
||||
from vllm_omni.entrypoints.omni import OmniServeCommand
|
||||
|
||||
def main():
|
||||
"""Main CLI entry point that intercepts vLLM commands"""
|
||||
# Check if --omni flag is present
|
||||
if "--omni" in sys.argv:
|
||||
# Remove --omni flag and process with vLLM-Omni
|
||||
omni_args = [arg for arg in sys.argv[1:] if arg != "--omni"]
|
||||
omni_serve = OmniServeCommand()
|
||||
omni_serve.run(omni_args)
|
||||
else:
|
||||
# Forward to original vLLM CLI
|
||||
from vllm.entrypoints.cli.main import main as vllm_main
|
||||
vllm_main()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
```
|
||||
|
||||
#### vLLM Plugin Integration
|
||||
File: vllm_omni/pyproject.toml
|
||||
```python
|
||||
[project.scripts]
|
||||
# Override vLLM's CLI entry point
|
||||
vllm = "vllm_omni.cli:main"
|
||||
|
||||
# Add entry point for vLLM integration
|
||||
[project.entry-points."vllm.plugins"]
|
||||
omni = "vllm_omni.plugin:OmniPlugin"
|
||||
```
|
||||
|
||||
### 3. Offline inference launch (Entry Point)
|
||||
Design an upper level class to incorporate multi Engines, each engine has a engine call.
|
||||
```python
|
||||
from vllm.entrypoints.llm import LLM
|
||||
from vllm.v1.engine.llm_engine import LLMEngine
|
||||
|
||||
class OmniLLM(LLM):
|
||||
"""Extended LLM supporting multiple engines and stage-based processing"""
|
||||
|
||||
def __init__(self, stage_configs: List[StageConfig]):
|
||||
super().__init__()
|
||||
self.stage_configs = stage_configs
|
||||
self.engine_list = [] # List of AsyncLLM instances for each stage
|
||||
self.output_processor = MultimodalOutputProcessor()
|
||||
self._initialize_stage_engines()
|
||||
|
||||
def _initialize_stage_engines(self) -> None:
|
||||
"""Initialize LLMEngine instances for each stage"""
|
||||
for stage_config in self.stage_configs:
|
||||
stage_llm = LLMEngine.from_vllm_config(
|
||||
vllm_config=stage_config.vllm_config,
|
||||
executor_class=self.executor_class,
|
||||
log_stats=self.log_stats
|
||||
)
|
||||
self.engine_list.append(stage_llm)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
stage_args_list: List[stage_args],
|
||||
use_tqdm: Union[bool, Callable[..., tqdm]] = True,
|
||||
lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
|
||||
priority: Optional[list[int]] = None,
|
||||
) -> list[RequestOutput]:
|
||||
"""Main generation interface - orchestrates multi-stage processing"""
|
||||
|
||||
# Process through each stage sequentially
|
||||
for i, stage_config in enumerate(self.stage_configs):
|
||||
stage_engine = self.engine_list[i]
|
||||
|
||||
# Prepare input for this stage
|
||||
stage_args = stage_args_list[i]
|
||||
prompt_str, engine_request, tokenization_kwargs = self._process_stage_inputs(stage_config, **stage_args)
|
||||
|
||||
# Add inputs to Engine
|
||||
stage_engine.add_request(request_id, prompt_str, tokenization_kwargs)
|
||||
# Run Engine
|
||||
stage_output = stage_engine.step()
|
||||
|
||||
# Update stage input and output of each stage for later usage
|
||||
self._update_stage_io(stage_output, stage_config)
|
||||
|
||||
|
||||
# Process final output
|
||||
final_output = self.output_processor.process_output(
|
||||
self.stage_configs[-1].stage_output
|
||||
)
|
||||
return final_output
|
||||
|
||||
def _process_stage_inputs(stage_config, **stage_args) -> Any:
|
||||
"""Prepare input for specific stage"""
|
||||
if stage_config.engine_type == "AR":
|
||||
return self._process_ar_inputs(**stage_args)
|
||||
elif stage_config.engine_type == "DiT":
|
||||
return self._process_dit_inputs(**stage_args)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
def _process_dit_inputs(**stage_args)-> Any:
|
||||
image_latent = self.vae.encode(**stage_args)
|
||||
|
||||
return image_latent
|
||||
```
|
||||
|
||||
### 4. Stage config for setting of different model stages
|
||||
```python
|
||||
@dataclass
|
||||
class StageConfig:
|
||||
"""Configuration for a processing stage"""
|
||||
stage_id: int
|
||||
engine_type: str # "AR" or "DiT"
|
||||
model_path: str
|
||||
input_modalities: List[str]
|
||||
output_modalities: List[str]
|
||||
vllm_config: Optional[VllmConfig] = None # For engine config of corresponding stage
|
||||
executor_class: type[Executor] # For execute class config of corresponding stage
|
||||
dit_config: Optional[DiTConfig] = None # For diffusion stages
|
||||
cache_config: Optional[DiTCacheConfig] = None
|
||||
```
|
||||
For AR stage, the setting is:
|
||||
```python
|
||||
from vllm.v1.core.sched.scheduler import Scheduler
|
||||
from vllm.v1.executor.multiproc_executor import MultiprocExecutor
|
||||
ar_stage_config = StageConfig()
|
||||
ar_stage_config.vllm_config.scheduler_config.scheduler_cls = Scheduler # original vllm scheduler for AR
|
||||
ar_stage_config.executor_class = MultiprocExecutor # original vllm executor for AR
|
||||
```
|
||||
For DiT stage, the setting is:
|
||||
```python
|
||||
from vllm.v1.executor.multiproc_executor import MultiprocExecutor
|
||||
dit_stage_config = StageConfig()
|
||||
dit_stage_config.vllm_config.scheduler_config.scheduler_cls = DiffusionScheduler # New scheduler for DiT
|
||||
|
||||
# For diffusion models without using diffusers
|
||||
dit_stage_config.executor_class = MultiprocExecutor
|
||||
|
||||
# For diffusion models using diffusers
|
||||
dit_stage_config.executor_class = DiffusersPipelineExecutor
|
||||
```
|
||||
|
||||
### 5. Online inference main class
|
||||
|
||||
Similar to OmniLLM in offline inference, add some asynchronous processing, referring to AsyncLLM
|
||||
```python
|
||||
from vllm.v1.engine.async_llm import AsyncLLM
|
||||
|
||||
|
||||
class AsyncOmni(AsyncLLM):
|
||||
"""Extended AsyncLLM supporting multiple engines and stage-based processing"""
|
||||
|
||||
def __init__(self, stage_configs: List[StageConfig]):
|
||||
super().__init__()
|
||||
```
|
||||
|
||||
### 6. Engine Core
|
||||
No need to change. The specific executor, scheduler and output type can be transferred into it with new configs.
|
||||
```python
|
||||
# class: EngineCore
|
||||
# EngineCore.step (simplified)
|
||||
class EngineCore:
|
||||
def step(self):
|
||||
scheduler_output = scheduler.schedule()
|
||||
model_output = executor.execute_model(scheduler_output)
|
||||
engine_outputs = scheduler.update_from_output(
|
||||
scheduler_output, model_output
|
||||
)
|
||||
return engine_outputs
|
||||
```
|
||||
|
||||
### 7. Scheduler
|
||||
Create a new Diffusion Scheduler for DiT, which is inherited from original vllm AR scheduler. At first, no complicated strategy.
|
||||
Just process the request first in first out. Then create a class of DiT Cache Manager for optimization in the future.
|
||||
```python
|
||||
from vllm.v1.core.sched.scheduler import Scheduler
|
||||
class DiffusionScheduler(Scheduler):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vllm_config: VllmConfig,
|
||||
kv_cache_config: KVCacheConfig,
|
||||
dit_cache_config: DiTCacheConfig,
|
||||
structured_output_manager: StructuredOutputManager,
|
||||
mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
|
||||
include_finished_set: bool = False,
|
||||
log_stats: bool = False,
|
||||
) -> None:
|
||||
|
||||
self.dit_cache_manager = DiTCacheManager(dit_cache_config)
|
||||
```
|
||||
For DiT Cache Manager (Can refer to xDiT):
|
||||
```python
|
||||
class DiTCacheManager:
|
||||
"""Manages DiT-specific caching"""
|
||||
|
||||
def __init__(self, config: DiTCacheConfig):
|
||||
self.cache_tensors: Dict[str, torch.Tensor] = {}
|
||||
self.cache_groups: List[DiTCacheTensor] = config.dit_cache_tensors
|
||||
|
||||
def allocate_cache(self, request_id: str, size: int) -> torch.Tensor
|
||||
def get_cache(self, request_id: str) -> Optional[torch.Tensor]
|
||||
def release_cache(self, request_id: str) -> None
|
||||
def clear_expired_cache(self) -> None
|
||||
```
|
||||
|
||||
### 8. Executor for DiT without diffusers
|
||||
|
||||
No need to change. Just import original AR executor.
|
||||
|
||||
### 9. Diffusers Pipeline Executor (no worker)
|
||||
|
||||
A single-process executor that directly runs the Diffusers pipeline without spawning workers or using RPC. Interfaces remain identical to `Executor` so the EngineCore loop is unchanged.
|
||||
|
||||
#### Function map(Pipeline Executor)
|
||||
|
||||
##### 1) Inherited and overridden
|
||||
```python
|
||||
from concurrent.futures import Future
|
||||
from typing import Optional, Union, Callable, Any
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from vllm.v1.executor.abstract import Executor
|
||||
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
|
||||
from vllm.v1.outputs import ModelRunnerOutput
|
||||
from vllm.tasks import SupportedTask
|
||||
from diffusers import DiffusionPipeline
|
||||
|
||||
class DiffusersPipelineExecutor(Executor):
|
||||
supports_pp: bool = False # Single-process, no TP/PP/DP
|
||||
|
||||
def _init_executor(self) -> None:
|
||||
# Called by ExecutorBase.__init__
|
||||
self._failure_callback: Optional[Callable[[], None]] = None
|
||||
self._device = self._resolve_device()
|
||||
self._dtype = self._resolve_dtype()
|
||||
self._pipeline = self._build_pipeline(device=self._device, dtype=self._dtype)
|
||||
self._profiler = None
|
||||
self._is_failed = False
|
||||
self.is_sleeping = False
|
||||
self.sleeping_tags: set[str] = set()
|
||||
|
||||
# major functions to build/run diffusers pipeline
|
||||
def _build_pipeline(self, device:torch.device, dtype:torch.dtype)->DiffusionPipeline:
|
||||
model_name = "Qwen/Qwen-Image"
|
||||
|
||||
self.pipe = DiffusionPipeline.from_pretrained(model_name, torch_dtype=dtype)
|
||||
self.pipe = pipe.to(device)
|
||||
|
||||
def _run_pipeline(self, scheduler_output) -> ModelRunnerOutput: ...
|
||||
positive_magic = {
|
||||
"en": ", Ultra HD, 4K, cinematic composition.", # for english prompt
|
||||
"zh": ", 超清,4K,电影级构图." # for chinese prompt
|
||||
}
|
||||
|
||||
# Generate image
|
||||
prompt_embeds = self._get_and_process_prompt_embeds(scheduler_output, positive_magic)
|
||||
negative_prompt_embeds = self.pipe.embed_prompt(" ")
|
||||
|
||||
width, height = aspect_ratios["16:9"]
|
||||
|
||||
image = pipe(
|
||||
prompt_embeds=prompt_embeds,
|
||||
negative_prompt_embeds=negative_prompt_embeds,
|
||||
width=width,
|
||||
height=height,
|
||||
num_inference_steps=50,
|
||||
true_cfg_scale=4.0,
|
||||
generator=torch.Generator(device="cuda").manual_seed(42)
|
||||
).images[0]
|
||||
|
||||
output = self.wrap_image_as_ModelRunnerOutput(image)
|
||||
return output
|
||||
|
||||
# ---- Internal helpers (implementation-specific, not public API) ----
|
||||
def _resolve_device(self): ...
|
||||
def _resolve_dtype(self): ...
|
||||
def _get_model(self) -> nn.Module: ...
|
||||
def _get_and_process_prompt_embeds(self, scheduler_output, positive_magic):
|
||||
...
|
||||
#append the positive_magic to prompt and embed them to prompt embed tensors
|
||||
|
||||
# Functions related to workers should either raise NotImplementedError
|
||||
# or return empty defaults in this no-worker executor.
|
||||
|
||||
def collective_rpc(self, method, timeout=None, args=(), kwargs=None) -> list[Any]:
|
||||
# No workers in pipeline executor
|
||||
raise NotImplementedError("No workers in DiffusersPipelineExecutor")
|
||||
|
||||
def initialize_from_config(self, kv_cache_configs: list[KVCacheConfig]) -> None:
|
||||
return # no-op (pipeline already built in _init_executor)
|
||||
|
||||
def register_failure_callback(self, callback):
|
||||
self._failure_callback = callback
|
||||
|
||||
def determine_available_memory(self) -> list[int]: # bytes
|
||||
# Single device; return [available_bytes]. If CPU-only, return [0].
|
||||
return [self._determine_available_bytes(self._device)]
|
||||
|
||||
```
|
||||
|
||||
### 10. Worker
|
||||
|
||||
#### Inheritance strategy
|
||||
Prefer reusing the mature GPU Worker end-to-end. Worker class is selected by configuration (vllm_config). Do not add a new executor-specific worker binding. If customization is needed, override only `init_device` to construct the `DiffusionModelRunner`; all other behaviors (device init details, profiling, sleep/wake, PP/TP comms, execute path) remain from `vllm/v1/worker/gpu_worker.py::Worker`.
|
||||
|
||||
|
||||
Inherited and overridden
|
||||
```python
|
||||
# Optional: only if you need to plug a custom DiffusionModelRunner.
|
||||
from vllm.v1.worker.gpu_worker import Worker as GPUWorker
|
||||
from vllm.v1.worker.diffusion_model_runner import DiffusionModelRunner
|
||||
|
||||
class Worker(GPUWorker):
|
||||
def init_device(self) -> None:
|
||||
#those related to device check and init
|
||||
...
|
||||
|
||||
self.model_runner = DiffusionModelRunner(self.vllm_config, ...)
|
||||
```
|
||||
|
||||
### 11. Model Runner
|
||||
|
||||
### Function map (Model Runner)
|
||||
#### 1) Inherited and overridden
|
||||
Those parts relied to the KV Cache will be omitted if we do not register the model to the vllm config. The engine core will view it as do not require KV Cache, and handle it properly
|
||||
|
||||
Reuse `vllm/v1/outputs.py::ModelRunnerOutput`:
|
||||
- DiffusionModelRunner: Use the `pooler_output=[Tensor,...]` to return multi modal tensors
|
||||
- ARModelRunner: Use the `pooler_output=[Tensor,...]` to return hidden states.
|
||||
```python
|
||||
from typing import Optional, Union
|
||||
import torch
|
||||
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
|
||||
from vllm.v1.outputs import ModelRunnerOutput
|
||||
|
||||
|
||||
class DiffusionModelRunner(GPUModelRunner):
|
||||
@torch.inference_mode()
|
||||
def execute_model(
|
||||
self,
|
||||
scheduler_output: "SchedulerOutput",
|
||||
intermediate_tensors: Optional[IntermediateTensors] = None,
|
||||
) -> Union[ModelRunnerOutput, IntermediateTensors]:
|
||||
...
|
||||
return ModelRunnerOutput(
|
||||
req_ids=[...],
|
||||
req_id_to_index={...},
|
||||
sampled_token_ids=[],
|
||||
spec_token_ids=None,
|
||||
logprobs=None,
|
||||
prompt_logprobs_dict={},
|
||||
pooler_output=[Tensor,...], # Return Hidden states
|
||||
kv_connector_output=None,
|
||||
num_nans_in_logits=None,
|
||||
)# return multi modal tensors via pooler_output=[Tensor,...]
|
||||
|
||||
|
||||
class ARModelRunner(GPUModelRunner):
|
||||
@torch.inference_mode()
|
||||
def execute_model(
|
||||
self,
|
||||
scheduler_output: "SchedulerOutput",
|
||||
intermediate_tensors: Optional[IntermediateTensors] = None,
|
||||
) -> Union[ModelRunnerOutput, IntermediateTensors]:
|
||||
...
|
||||
return ModelRunnerOutput(
|
||||
req_ids=[...],
|
||||
req_id_to_index={...},
|
||||
sampled_token_ids=[],
|
||||
spec_token_ids=None,
|
||||
logprobs=None,
|
||||
prompt_logprobs_dict={},
|
||||
pooler_output=[Tensor,...], # Return Hidden states
|
||||
kv_connector_output=None,
|
||||
num_nans_in_logits=None,
|
||||
)# return hidden states via pooler_output=[Tensor,...]
|
||||
```
|
||||
|
||||
### 12. RequestState Role and Usage
|
||||
```bash
|
||||
# RequestState serves as the per-request state tracker in OutputProcessor:
|
||||
# - Maintains request-specific state (tokens, logprobs, detokenizer, etc.)
|
||||
# - Converts EngineCoreOutput → RequestOutput/PoolingRequestOutput
|
||||
# - Manages request lifecycle from registration to completion
|
||||
AsyncLLM.add_request()
|
||||
↓
|
||||
OutputProcessor.add_request()
|
||||
↓
|
||||
RequestState.from_new_request() → create request state
|
||||
|
||||
AsyncLLM.__init__() / AsyncLLM.generate() / AsyncLLM.encode() → create a background loop that continuously pulls output from EngineCore
|
||||
↓
|
||||
OutputProcessor.process_outputs() → update state and process outputs
|
||||
↓
|
||||
RequestState.make_request_output() → convert to the final output, formatted as RequestOutput or PoolingRequestOutput
|
||||
↓
|
||||
RequestOutputCollector.put() → push to the queue (AsyncLLM)
|
||||
```
|
||||
|
||||
Need to add implementation for an existing method
|
||||
```python
|
||||
class RequestState:
|
||||
def __init__(self, request_id: str, parent_req: Optional[ParentRequest],
|
||||
request_index: int, lora_name: Optional[str],
|
||||
output_kind: RequestOutputKind, prompt: Optional[str],
|
||||
prompt_token_ids: list[int], logprobs_processor: Optional[LogprobsProcessor],
|
||||
detokenizer: Optional[IncrementalDetokenizer],
|
||||
max_tokens_param: Optional[int], arrival_time: float,
|
||||
queue: Optional[RequestOutputCollector], log_stats: bool):
|
||||
|
||||
def _new_pooling_output(self, pooling_output: torch.Tensor) -> PoolingOutput:
|
||||
"""Create PoolingOutput for multimodal/pooling requests"""
|
||||
return PoolingOutput(data=pooling_output)
|
||||
```
|
||||
|
||||
|
||||
### 13. OutputProcessor
|
||||
|
||||
For hidden state output, original OutputProcessor already support.
|
||||
|
||||
We only need to add new one for final multimodal output.
|
||||
```python
|
||||
class MultimodalOutputProcessor(OutputProcessor):
|
||||
"""Handles multimodal output processing"""
|
||||
|
||||
def __init__(self):
|
||||
self.output_handlers: Dict[str, OutputProcessor] = {}
|
||||
|
||||
def process_outputs(self, engine_core_outputs: list[EngineCoreOutput], ...):
|
||||
for engine_core_output in engine_core_outputs:
|
||||
# Option 1: Use output_type field (if EngineCoreOutput is extended)
|
||||
if engine_core_output.output_type == "image":
|
||||
self._process_image_output(engine_core_output)
|
||||
elif engine_core_output.output_type == "text+image":
|
||||
self._process_text_image_output(engine_core_output)
|
||||
elif engine_core_output.output_type == "latents":
|
||||
self._process_latents_output(engine_core_output)
|
||||
elif engine_core_output.output_type == "text":
|
||||
self._process_text_output(engine_core_output)
|
||||
else:
|
||||
# Fallback: use existing pooling_output logic
|
||||
if engine_core_output.pooling_output is not None:
|
||||
self._process_pooling_output(engine_core_output)
|
||||
else:
|
||||
self._process_text_output(engine_core_output)
|
||||
```
|
||||
@@ -4,4 +4,5 @@ vLLM-Omni supports the following hardware platforms:
|
||||
|
||||
- [GPU](gpu.md)
|
||||
- [NVIDIA CUDA](gpu.md)
|
||||
- [AMD ROCm](gpu.md)
|
||||
- [NPU](npu.md)
|
||||
@@ -14,6 +14,10 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu/cuda.inc.md:requirements"
|
||||
|
||||
=== "AMD ROCm"
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu/rocm.inc.md:requirements"
|
||||
|
||||
## Set up using Python
|
||||
|
||||
### Create a new Python environment
|
||||
@@ -26,6 +30,11 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu/cuda.inc.md:pre-built-wheels"
|
||||
|
||||
|
||||
=== "AMD ROCm"
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu/rocm.inc.md:pre-built-wheels"
|
||||
|
||||
[](){ #build-from-source }
|
||||
|
||||
### Build wheel from source
|
||||
@@ -34,8 +43,22 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu/cuda.inc.md:build-wheel-from-source"
|
||||
|
||||
=== "AMD ROCm"
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu/rocm.inc.md:build-wheel-from-source"
|
||||
|
||||
## Set up using Docker
|
||||
|
||||
### Build wheel from source
|
||||
|
||||
=== "NVIDIA CUDA"
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu/cuda.inc.md:build-wheel-from-source-in-docker"
|
||||
|
||||
=== "AMD ROCm"
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu/rocm.inc.md:build-wheel-from-source-in-docker"
|
||||
|
||||
### Pre-built images
|
||||
|
||||
=== "NVIDIA CUDA"
|
||||
|
||||
@@ -74,6 +74,9 @@ uv pip install --no-build-isolation --editable .
|
||||
|
||||
# --8<-- [end:build-wheel-from-source]
|
||||
|
||||
# --8<-- [start:build-wheel-from-source-in-docker]
|
||||
|
||||
# --8<-- [end:build-wheel-from-source-in-docker]
|
||||
|
||||
# --8<-- [start:pre-built-images]
|
||||
|
||||
|
||||
@@ -0,0 +1,115 @@
|
||||
# --8<-- [start:requirements]
|
||||
|
||||
- GPU: Validated on gfx942 (It should be supported on the AMD GPUs that are supported by vLLM.)
|
||||
|
||||
# --8<-- [end:requirements]
|
||||
# --8<-- [start:set-up-using-python]
|
||||
|
||||
vLLM-Omni current recommends the steps in under setup through Docker Images.
|
||||
|
||||
# --8<-- [start:pre-built-wheels]
|
||||
|
||||
# --8<-- [end:pre-built-wheels]
|
||||
|
||||
# --8<-- [start:build-wheel-from-source]
|
||||
|
||||
# --8<-- [end:build-wheel-from-source]
|
||||
|
||||
# --8<-- [start:build-wheel-from-source-in-docker]
|
||||
|
||||
#### Get the correct docker environment
|
||||
|
||||
The v0.11.0 releases is on Oct 3, 2025, so we will pull a AMD nightly docker image close to the release period. `rocm/vllm-dev:nightly_main_20251005` is picked as it has a correct `transformers` version that supports qwen3_omni.
|
||||
|
||||
`docker pull rocm/vllm-dev:nightly_main_20251005`
|
||||
|
||||
#### Launch the docker environment
|
||||
|
||||
In this docker environment, we have all the dependencies for vLLM installed, like `flash-attn` and `aiter`
|
||||
|
||||
```bash
|
||||
#!/bin/bash
|
||||
docker run -it \
|
||||
--privileged \
|
||||
--network=host \
|
||||
--group-add=video \
|
||||
--ipc=host \
|
||||
--cap-add=SYS_PTRACE \
|
||||
--security-opt seccomp=unconfined \
|
||||
--device /dev/kfd \
|
||||
--device /dev/dri \
|
||||
--name vllmomni \
|
||||
rocm/vllm-dev:nightly_main_20251005 \
|
||||
bash
|
||||
```
|
||||
|
||||
#### Update docker dependencies
|
||||
|
||||
```bash
|
||||
sudo apt update
|
||||
sudo apt install ffmpeg -y
|
||||
```
|
||||
|
||||
#### Installation of vLLM
|
||||
Install the v0.11.0 release of vLLM as vLLM-Omni is currently depend on v0.11.0.
|
||||
|
||||
```bash
|
||||
# uninstall existing vllm which is based on the main branch, not on the releases branch
|
||||
python3 -m pip uninstall -y vllm
|
||||
git clone https://github.com/vllm-project/vllm.git vllm0110
|
||||
cd vllm0110/
|
||||
|
||||
# checkout the version tag v0.11.0
|
||||
git checkout -b v0.11.0 v0.11.0
|
||||
|
||||
# you should see 0.11.0
|
||||
python -c "import setuptools_scm; print(setuptools_scm.get_version())"
|
||||
|
||||
# PYTORCH_ROCM_ARCH=<your-gpu-arch> python3 setup.py develop
|
||||
PYTORCH_ROCM_ARCH=gfx942 python3 setup.py develop
|
||||
```
|
||||
|
||||
!!! note
|
||||
vLLM release wheels based on the branch with prefix `releases/`, not from the tag as vLLM may cherry pick bugfixes after cutting a branch.
|
||||
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
|
||||
There are two approaches to installing `vllm-omni`
|
||||
|
||||
|
||||
1. Install from PyPI
|
||||
|
||||
```bash
|
||||
pip install vllm-omni
|
||||
|
||||
# you might want to also download the source code of the same version to run the examples
|
||||
```
|
||||
|
||||
2. Install from source
|
||||
|
||||
```bash
|
||||
# clone and checkout to the release version.
|
||||
git clone https://github.com/vllm-project/vllm-omni.git
|
||||
cd vllm-omni
|
||||
git checkout v0.11.0rc1
|
||||
pip install -e .
|
||||
```
|
||||
|
||||
#### Setup Environment Variables
|
||||
|
||||
After installing `vllm` and `vllm-omni` please setup the environment variable as well before running any of the examples and command.
|
||||
|
||||
```bash
|
||||
export MIOPEN_FIND_MODE=FAST
|
||||
export VLLM_ROCM_USE_AITER=1
|
||||
export VLLM_ROCM_USE_AITER_MHA=1
|
||||
export VLLM_ROCM_USE_AITER_LINEAR=0
|
||||
export VLLM_ROCM_USE_AITER_RMSNORM=0
|
||||
```
|
||||
|
||||
# --8<-- [end:build-wheel-from-source-in-docker]
|
||||
|
||||
# --8<-- [start:pre-built-images]
|
||||
|
||||
# --8<-- [end:pre-built-images]
|
||||
@@ -6,7 +6,7 @@ vLLM-Omni supports unified multimodal comprehension and generation models across
|
||||
|
||||
If vLLM-Omni natively supports a model, its implementation can be found in <gh-file:vllm_omni/model_executor/models> and <gh-file:vllm_omni/diffusion/models>.
|
||||
|
||||
## List of Supported Models for Nvidia GPU
|
||||
## List of Supported Models for Nvidia GPU / AMD GPU
|
||||
|
||||
<style>
|
||||
th {
|
||||
@@ -22,6 +22,7 @@ th {
|
||||
| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
|
||||
| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
|
||||
|`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
|
||||
| `WanPipeline` | Wan2.2 | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` |
|
||||
|
||||
|
||||
## List of Supported Models for NPU
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 107 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 97 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 126 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 367 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 480 KiB |
@@ -0,0 +1,30 @@
|
||||
# Qwen-Image-Edit Offline Inference
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_image>. Added in PR [#196](https://github.com/vllm-project/vllm-omni/pull/196).
|
||||
|
||||
This example edits an input image with `Qwen/Qwen-Image-Edit` using the `image_edit.py` CLI.
|
||||
|
||||
## Local CLI Usage
|
||||
|
||||
```bash
|
||||
python image_edit.py \
|
||||
--image qwen_bear.png \
|
||||
--prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \
|
||||
--output output_image_edit.png \
|
||||
--num_inference_steps 50 \
|
||||
--cfg_scale 4.0
|
||||
```
|
||||
|
||||
Key arguments:
|
||||
- `--image`: path to the source image (PNG/JPG, converted to RGB).
|
||||
- `--prompt` / `--negative_prompt`: text description (string).
|
||||
- `--cfg_scale`: true CFG scale for Qwen-Image-Edit (quality vs. fidelity).
|
||||
- `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
|
||||
- `--output`: path to save the generated PNG.
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "image_edit.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/image_to_image/image_edit.py"
|
||||
``````
|
||||
+5
-6
@@ -1,9 +1,9 @@
|
||||
# Qwen-Image Offline Inference
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen_image>.
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_image>.
|
||||
|
||||
|
||||
This folder provides two simple entrypoints for experimenting with `Qwen/Qwen-Image` using vLLM-Omni:
|
||||
This folder provides two simple entrypoints for experimenting with `Qwen/Qwen-Image` `Tongyi-MAI/Z-Image-Turbo` using vLLM-Omni:
|
||||
|
||||
- `text_to_image.py`: command-line script for single image generation.
|
||||
- `web_demo.py`: lightweight Gradio UI for interactive prompt/seed/CFG exploration.
|
||||
@@ -13,6 +13,7 @@ This folder provides two simple entrypoints for experimenting with `Qwen/Qwen-Im
|
||||
|
||||
```bash
|
||||
python text_to_image.py \
|
||||
--model Tongyi-MAI/Z-Image-Turbo \
|
||||
--prompt "a cup of coffee on the table" \
|
||||
--seed 42 \
|
||||
--cfg_scale 4.0 \
|
||||
@@ -33,8 +34,6 @@ Key arguments:
|
||||
- `--height/--width`: output resolution (defaults 1024x1024).
|
||||
- `--output`: path to save the generated PNG.
|
||||
|
||||
> ℹ️ Qwen-Image currently publishes best-effort presets at `1328x1328`, `1664x928`, `928x1664`, `1472x1140`, `1140x1472`, `1584x1056`, and `1056x1584`. Adjust `--height/--width` accordingly for the most reliable outcomes.
|
||||
|
||||
## Web UI Demo
|
||||
|
||||
Launch the gradio demo:
|
||||
@@ -49,9 +48,9 @@ Then open `http://localhost:7862/` on your local browser to interact with the we
|
||||
|
||||
??? abstract "gradio_demo.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/qwen_image/gradio_demo.py"
|
||||
--8<-- "examples/offline_inference/text_to_image/gradio_demo.py"
|
||||
``````
|
||||
??? abstract "text_to_image.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/qwen_image/text_to_image.py"
|
||||
--8<-- "examples/offline_inference/text_to_image/text_to_image.py"
|
||||
``````
|
||||
@@ -0,0 +1,38 @@
|
||||
# Wan2.2 Text-to-Video Offline Inference
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_video>. Added in PR [#202](https://github.com/vllm-project/vllm-omni/pull/202).
|
||||
|
||||
The `Wan-AI/Wan2.2-T2V-A14B-Diffusers` pipeline generates short videos from text prompts.
|
||||
|
||||
## Local CLI Usage
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
|
||||
--negative_prompt "<optional quality filter>" \
|
||||
--height 480 \
|
||||
--width 640 \
|
||||
--num_frames 32 \
|
||||
--guidance_scale 4.0 \
|
||||
--guidance_scale_high 3.0 \
|
||||
--num_inference_steps 40 \
|
||||
--fps 16 \
|
||||
--output t2v_out.mp4
|
||||
```
|
||||
|
||||
Key arguments:
|
||||
- `--prompt`: text description (string).
|
||||
- `--height/--width`: output resolution (defaults 720x1280). Dimensions should align with Wan VAE downsampling (multiples of 8).
|
||||
- `--num_frames`: Number of frames (Wan default is 81).
|
||||
- `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high)..
|
||||
- `--negative_prompt`: optional list of artifacts to suppress (the PR demo used a long Chinese string).
|
||||
- `--boundary_ratio`: Boundary split ratio for low/high DiT.
|
||||
- `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video).
|
||||
- `--output`: path to save the generated video.
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "text_to_video.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/wan22/text_to_video.py"
|
||||
``````
|
||||
@@ -0,0 +1,99 @@
|
||||
# TeaCache Configuration Guide
|
||||
|
||||
TeaCache speeds up diffusion model inference by caching transformer computations when consecutive timesteps are similar. This typically provides **1.5x-2.0x speedup** with minimal quality loss.
|
||||
|
||||
## Quick Start
|
||||
|
||||
Enable TeaCache by setting `cache_adapter` to `"tea_cache"`:
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
# Simple configuration - model_class_name and model_type are auto-detected
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
cache_adapter="tea_cache",
|
||||
cache_config={
|
||||
"rel_l1_thresh": 0.2 # Optional, defaults to 0.2
|
||||
}
|
||||
)
|
||||
outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
|
||||
```
|
||||
|
||||
### Using Environment Variable
|
||||
|
||||
You can also enable TeaCache via environment variable:
|
||||
|
||||
```bash
|
||||
export DIFFUSION_CACHE_ADAPTER=tea_cache
|
||||
```
|
||||
|
||||
Then initialize without explicitly setting `cache_adapter`:
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
cache_config={"rel_l1_thresh": 0.2} # Optional
|
||||
)
|
||||
```
|
||||
|
||||
## Configuration Parameters
|
||||
|
||||
### `rel_l1_thresh` (float, default: `0.2`)
|
||||
|
||||
Controls the balance between speed and quality. Lower values prioritize quality, higher values prioritize speed.
|
||||
|
||||
**Recommended values:**
|
||||
|
||||
- `0.2` - **~1.5x speedup** with minimal quality loss (recommended)
|
||||
- `0.4` - **~1.8x speedup** with slight quality loss
|
||||
- `0.6` - **~2.0x speedup** with noticeable quality loss
|
||||
- `0.8` - **~2.25x speedup** with significant quality loss
|
||||
|
||||
## Examples
|
||||
|
||||
### Python API
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
cache_adapter="tea_cache",
|
||||
cache_config={"rel_l1_thresh": 0.2}
|
||||
)
|
||||
outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
|
||||
```
|
||||
|
||||
## Performance Tuning
|
||||
|
||||
Start with the default `rel_l1_thresh=0.2` and adjust based on your needs:
|
||||
|
||||
- **Maximum quality**: Use `0.1-0.2`
|
||||
- **Balanced**: Use `0.2-0.4` (recommended)
|
||||
- **Maximum speed**: Use `0.6-0.8` (may reduce quality)
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Quality Degradation
|
||||
|
||||
If you notice quality issues, lower the threshold:
|
||||
|
||||
```python
|
||||
cache_config={"rel_l1_thresh": 0.1} # More conservative caching
|
||||
```
|
||||
|
||||
## Supported Models
|
||||
|
||||
Currently supported models:
|
||||
|
||||
- **QwenImagePipeline**
|
||||
- Model identifier: `"Qwen/Qwen-Image"`
|
||||
- Example: `model="Qwen/Qwen-Image"`
|
||||
|
||||
Future support (extractors in development):
|
||||
|
||||
- **FluxPipeline** - Coming soon
|
||||
- **CogVideoXPipeline** - Coming soon
|
||||
File renamed without changes.
@@ -0,0 +1,30 @@
|
||||
# Qwen-Image-Edit Offline Inference
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_image>. Added in PR [#196](https://github.com/vllm-project/vllm-omni/pull/196).
|
||||
|
||||
This example edits an input image with `Qwen/Qwen-Image-Edit` using the `image_edit.py` CLI.
|
||||
|
||||
## Local CLI Usage
|
||||
|
||||
```bash
|
||||
python image_edit.py \
|
||||
--image qwen_bear.png \
|
||||
--prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \
|
||||
--output output_image_edit.png \
|
||||
--num_inference_steps 50 \
|
||||
--cfg_scale 4.0
|
||||
```
|
||||
|
||||
Key arguments:
|
||||
- `--image`: path to the source image (PNG/JPG, converted to RGB).
|
||||
- `--prompt` / `--negative_prompt`: text description (string).
|
||||
- `--cfg_scale`: true CFG scale for Qwen-Image-Edit (quality vs. fidelity).
|
||||
- `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
|
||||
- `--output`: path to save the generated PNG.
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "image_edit.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/image_to_image/image_edit.py"
|
||||
``````
|
||||
@@ -496,7 +496,15 @@ def parse_args():
|
||||
default=16000,
|
||||
help="Sampling rate for audio loading (default: 16000).",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--worker-backend", type=str, default="multi_process", choices=["multi_process", "ray"], help="backend"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ray-address",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Address of the Ray cluster.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
File renamed without changes.
File renamed without changes.
File renamed without changes.
@@ -0,0 +1,38 @@
|
||||
# Wan2.2 Text-to-Video Offline Inference
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_video>. Added in PR [#202](https://github.com/vllm-project/vllm-omni/pull/202).
|
||||
|
||||
The `Wan-AI/Wan2.2-T2V-A14B-Diffusers` pipeline generates short videos from text prompts.
|
||||
|
||||
## Local CLI Usage
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
|
||||
--negative_prompt "<optional quality filter>" \
|
||||
--height 480 \
|
||||
--width 640 \
|
||||
--num_frames 32 \
|
||||
--guidance_scale 4.0 \
|
||||
--guidance_scale_high 3.0 \
|
||||
--num_inference_steps 40 \
|
||||
--fps 16 \
|
||||
--output t2v_out.mp4
|
||||
```
|
||||
|
||||
Key arguments:
|
||||
- `--prompt`: text description (string).
|
||||
- `--height/--width`: output resolution (defaults 720x1280). Dimensions should align with Wan VAE downsampling (multiples of 8).
|
||||
- `--num_frames`: Number of frames (Wan default is 81).
|
||||
- `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high)..
|
||||
- `--negative_prompt`: optional list of artifacts to suppress (the PR demo used a long Chinese string).
|
||||
- `--boundary_ratio`: Boundary split ratio for low/high DiT.
|
||||
- `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video).
|
||||
- `--output`: path to save the generated video.
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "text_to_video.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/wan22/text_to_video.py"
|
||||
``````
|
||||
@@ -0,0 +1,105 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.utils.platform_utils import detect_device_type, is_npu
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Generate a video with Wan2.2 T2V.")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
help="Diffusers Wan2.2 model ID or local path.",
|
||||
)
|
||||
parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
|
||||
parser.add_argument("--negative_prompt", default="", help="Negative prompt.")
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
|
||||
parser.add_argument("--guidance_scale", type=float, default=4.0, help="CFG scale (applied to low/high).")
|
||||
parser.add_argument("--guidance_scale_high", type=float, default=None, help="Optional separate CFG for high-noise.")
|
||||
parser.add_argument("--height", type=int, default=720, help="Video height.")
|
||||
parser.add_argument("--width", type=int, default=1280, help="Video width.")
|
||||
parser.add_argument("--num_frames", type=int, default=81, help="Number of frames (Wan default is 81).")
|
||||
parser.add_argument("--num_inference_steps", type=int, default=40, help="Sampling steps.")
|
||||
parser.add_argument("--boundary_ratio", type=float, default=0.875, help="Boundary split ratio for low/high DiT.")
|
||||
parser.add_argument(
|
||||
"--flow_shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
|
||||
)
|
||||
parser.add_argument("--output", type=str, default="wan22_output.mp4", help="Path to save the video (mp4).")
|
||||
parser.add_argument("--fps", type=int, default=24, help="Frames per second for the output video.")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
device = detect_device_type()
|
||||
generator = torch.Generator(device=device).manual_seed(args.seed)
|
||||
|
||||
# Enable VAE memory optimizations on NPU
|
||||
vae_use_slicing = is_npu()
|
||||
vae_use_tiling = is_npu()
|
||||
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
vae_use_slicing=vae_use_slicing,
|
||||
vae_use_tiling=vae_use_tiling,
|
||||
boundary_ratio=args.boundary_ratio,
|
||||
flow_shift=args.flow_shift,
|
||||
)
|
||||
|
||||
frames = omni.generate(
|
||||
args.prompt,
|
||||
negative_prompt=args.negative_prompt,
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
generator=generator,
|
||||
guidance_scale=args.guidance_scale,
|
||||
guidance_scale_2=args.guidance_scale_high,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
num_frames=args.num_frames,
|
||||
)
|
||||
|
||||
output_path = Path(args.output)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
try:
|
||||
from diffusers.utils import export_to_video
|
||||
except ImportError:
|
||||
raise ImportError("diffusers is required for export_to_video.")
|
||||
|
||||
# frames may be np.ndarray (preferred) or torch.Tensor
|
||||
# export_to_video expects a list of frames with values in [0, 1]
|
||||
if isinstance(frames, torch.Tensor):
|
||||
video_tensor = frames.detach().cpu()
|
||||
if video_tensor.dim() == 5:
|
||||
# [B, C, F, H, W] or [B, F, H, W, C]
|
||||
if video_tensor.shape[1] in (3, 4):
|
||||
video_tensor = video_tensor[0].permute(1, 2, 3, 0)
|
||||
else:
|
||||
video_tensor = video_tensor[0]
|
||||
elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
|
||||
video_tensor = video_tensor.permute(1, 2, 3, 0)
|
||||
# If float, assume [-1,1] and normalize to [0,1]
|
||||
if video_tensor.is_floating_point():
|
||||
video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
|
||||
video_array = video_tensor.float().numpy()
|
||||
else:
|
||||
video_array = frames
|
||||
if hasattr(video_array, "shape") and video_array.ndim == 5:
|
||||
video_array = video_array[0]
|
||||
|
||||
# Convert 4D array (frames, H, W, C) to list of frames for export_to_video
|
||||
if isinstance(video_array, np.ndarray) and video_array.ndim == 4:
|
||||
video_array = list(video_array)
|
||||
|
||||
export_to_video(video_array, str(output_path), fps=args.fps)
|
||||
print(f"Saved generated video to {output_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -91,9 +91,58 @@ def parse_args():
|
||||
default=None,
|
||||
help="Path to custom stage configs YAML file (optional).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
action="store_true",
|
||||
help="Enable statistics logging for AsyncOmni.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-file",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path prefix for AsyncOmni log files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--init-sleep-seconds",
|
||||
type=int,
|
||||
default=30,
|
||||
help="Seconds to sleep between starting stage processes.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--shm-threshold-bytes",
|
||||
type=int,
|
||||
default=65536,
|
||||
help="Threshold in bytes for using shared memory IPC.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch-timeout",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Batching timeout (seconds) inside each stage.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--init-timeout",
|
||||
type=int,
|
||||
default=ASYNC_INIT_TIMEOUT,
|
||||
help="Timeout (seconds) for initializing all stages.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def build_async_omni_cli_args(base_args: argparse.Namespace) -> argparse.Namespace:
|
||||
"""Construct the minimal CLI args Namespace expected by AsyncOmni."""
|
||||
return argparse.Namespace(
|
||||
model=base_args.model,
|
||||
stage_configs_path=getattr(base_args, "stage_configs_path", None),
|
||||
log_stats=bool(getattr(base_args, "log_stats", False)),
|
||||
log_file=getattr(base_args, "log_file", None),
|
||||
init_sleep_seconds=int(getattr(base_args, "init_sleep_seconds", 30)),
|
||||
shm_threshold_bytes=int(getattr(base_args, "shm_threshold_bytes", 65536)),
|
||||
batch_timeout=int(getattr(base_args, "batch_timeout", 10)),
|
||||
init_timeout=int(getattr(base_args, "init_timeout", ASYNC_INIT_TIMEOUT)),
|
||||
)
|
||||
|
||||
|
||||
def build_sampling_params(seed: int, model_key: str) -> list[SamplingParams]:
|
||||
"""Build SamplingParams objects by reusing the dict definitions."""
|
||||
return [SamplingParams(**params_dict) for params_dict in build_sampling_params_dict(seed, model_key)]
|
||||
@@ -500,11 +549,8 @@ def main():
|
||||
print(f"Using custom stage configs: {args.stage_configs_path}")
|
||||
|
||||
sampling_params = build_sampling_params(SEED, model_name)
|
||||
omni = AsyncOmni(
|
||||
model=args.model,
|
||||
stage_configs_path=args.stage_configs_path,
|
||||
init_timeout=ASYNC_INIT_TIMEOUT,
|
||||
)
|
||||
cli_args = build_async_omni_cli_args(args)
|
||||
omni = AsyncOmni(model=args.model, cli_args=cli_args)
|
||||
print("✓ AsyncOmni initialized successfully")
|
||||
prompt_args_template = create_prompt_args(args)
|
||||
|
||||
|
||||
@@ -94,9 +94,58 @@ def parse_args():
|
||||
default=None,
|
||||
help="Path to custom stage configs YAML file (optional).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
action="store_true",
|
||||
help="Enable statistics logging for AsyncOmni.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-file",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path prefix for AsyncOmni log files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--init-sleep-seconds",
|
||||
type=int,
|
||||
default=30,
|
||||
help="Seconds to sleep between starting stage processes.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--shm-threshold-bytes",
|
||||
type=int,
|
||||
default=65536,
|
||||
help="Threshold in bytes for using shared memory IPC.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch-timeout",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Batching timeout (seconds) inside each stage.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--init-timeout",
|
||||
type=int,
|
||||
default=ASYNC_INIT_TIMEOUT,
|
||||
help="Timeout (seconds) for initializing all stages.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def build_async_omni_cli_args(base_args: argparse.Namespace) -> argparse.Namespace:
|
||||
"""Construct the minimal CLI args Namespace expected by AsyncOmni."""
|
||||
return argparse.Namespace(
|
||||
model=base_args.model,
|
||||
stage_configs_path=getattr(base_args, "stage_configs_path", None),
|
||||
log_stats=bool(getattr(base_args, "log_stats", False)),
|
||||
log_file=getattr(base_args, "log_file", None),
|
||||
init_sleep_seconds=int(getattr(base_args, "init_sleep_seconds", 30)),
|
||||
shm_threshold_bytes=int(getattr(base_args, "shm_threshold_bytes", 65536)),
|
||||
batch_timeout=int(getattr(base_args, "batch_timeout", 10)),
|
||||
init_timeout=int(getattr(base_args, "init_timeout", ASYNC_INIT_TIMEOUT)),
|
||||
)
|
||||
|
||||
|
||||
def build_sampling_params(seed: int, model_key: str) -> list[SamplingParams]:
|
||||
"""Build SamplingParams objects by reusing the dict definitions."""
|
||||
return [SamplingParams(**params_dict) for params_dict in build_sampling_params_dict(seed, model_key)]
|
||||
@@ -506,11 +555,8 @@ def main():
|
||||
print(f"Using custom stage configs: {args.stage_configs_path}")
|
||||
|
||||
sampling_params = build_sampling_params(SEED, model_name)
|
||||
omni = AsyncOmni(
|
||||
model=args.model,
|
||||
stage_configs_path=args.stage_configs_path,
|
||||
init_timeout=ASYNC_INIT_TIMEOUT,
|
||||
)
|
||||
cli_args = build_async_omni_cli_args(args)
|
||||
omni = AsyncOmni(model=args.model, cli_args=cli_args)
|
||||
print("✓ AsyncOmni initialized successfully")
|
||||
prompt_args_template = create_prompt_args(args)
|
||||
|
||||
|
||||
@@ -0,0 +1,236 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.distributed.omni_connectors.adapter import try_recv_via_connector, try_send_via_connector
|
||||
from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector
|
||||
from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec, OmniTransferConfig
|
||||
from vllm_omni.distributed.omni_connectors.utils.initialization import get_connectors_config_for_stage
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_objects():
|
||||
return {"connector": MagicMock(), "metrics": MagicMock(), "queue_fn": MagicMock()}
|
||||
|
||||
|
||||
def test_send_success(mock_objects):
|
||||
"""Test try_send_via_connector success path."""
|
||||
# Setup
|
||||
mock_connector = mock_objects["connector"]
|
||||
mock_metrics = mock_objects["metrics"]
|
||||
mock_queue_fn = mock_objects["queue_fn"]
|
||||
|
||||
stage_id = 0
|
||||
next_stage_id = 1
|
||||
req_id = "req_123"
|
||||
inputs = {"input_ids": [1, 2, 3]}
|
||||
sampling_params = {"temperature": 0.7}
|
||||
prompt = "test prompt"
|
||||
|
||||
# Mock connector.put return
|
||||
# Returns: (success, size, metadata)
|
||||
mock_metadata = {"handle": "xyz"}
|
||||
mock_connector.put.return_value = (True, 100, mock_metadata)
|
||||
|
||||
# Execute
|
||||
result = try_send_via_connector(
|
||||
connector=mock_connector,
|
||||
stage_id=stage_id,
|
||||
next_stage_id=next_stage_id,
|
||||
req_id=req_id,
|
||||
next_inputs=inputs,
|
||||
sampling_params=sampling_params,
|
||||
original_prompt=prompt,
|
||||
next_stage_queue_submit_fn=mock_queue_fn,
|
||||
metrics=mock_metrics,
|
||||
)
|
||||
|
||||
# Verify
|
||||
assert result is True
|
||||
|
||||
# 1. Verify connector.put called correctly
|
||||
mock_connector.put.assert_called_once()
|
||||
args, _ = mock_connector.put.call_args
|
||||
assert args[0] == "0" # from_stage
|
||||
assert args[1] == "1" # to_stage
|
||||
assert args[2] == req_id
|
||||
# Verify payload structure in put
|
||||
payload = args[3]
|
||||
assert payload["engine_inputs"] == inputs
|
||||
assert payload["sampling_params"] == sampling_params
|
||||
|
||||
# 2. Verify queue notification submitted
|
||||
mock_queue_fn.assert_called_once()
|
||||
notify_payload = mock_queue_fn.call_args[0][0]
|
||||
assert notify_payload["request_id"] == req_id
|
||||
assert notify_payload["from_connector"] is True
|
||||
assert notify_payload["connector_metadata"] == mock_metadata
|
||||
|
||||
# 3. Verify metrics recorded
|
||||
mock_metrics.on_forward.assert_called_once()
|
||||
|
||||
|
||||
def test_send_fail(mock_objects):
|
||||
"""Test try_send_via_connector when connector fails."""
|
||||
mock_connector = mock_objects["connector"]
|
||||
mock_metrics = mock_objects["metrics"]
|
||||
mock_queue_fn = mock_objects["queue_fn"]
|
||||
|
||||
mock_connector.put.return_value = (False, 0, None)
|
||||
|
||||
result = try_send_via_connector(
|
||||
connector=mock_connector,
|
||||
stage_id=0,
|
||||
next_stage_id=1,
|
||||
req_id="req_fail",
|
||||
next_inputs={},
|
||||
sampling_params={},
|
||||
original_prompt="",
|
||||
next_stage_queue_submit_fn=mock_queue_fn,
|
||||
metrics=mock_metrics,
|
||||
)
|
||||
|
||||
assert result is False
|
||||
mock_queue_fn.assert_not_called()
|
||||
|
||||
|
||||
def test_recv_success(mock_objects):
|
||||
"""Test try_recv_via_connector success path."""
|
||||
mock_connector = mock_objects["connector"]
|
||||
|
||||
# Setup task received from queue
|
||||
task = {
|
||||
"request_id": "req_recv",
|
||||
"from_connector": True,
|
||||
"from_stage": "0",
|
||||
"connector_metadata": {"handle": "xyz"},
|
||||
}
|
||||
|
||||
# Setup connectors dict
|
||||
connectors = {("0", "1"): mock_connector}
|
||||
|
||||
# Mock connector.get return
|
||||
expected_data = {"engine_inputs": {"ids": [1]}}
|
||||
# get returns: (data_obj, size)
|
||||
mock_connector.get.return_value = (expected_data, 50)
|
||||
# serialize_obj needed for metrics calculation if size not returned directly
|
||||
mock_connector.serialize_obj.return_value = b"bytes"
|
||||
|
||||
# Execute
|
||||
# We are stage 1 receiving from stage 0
|
||||
inputs, rx_metrics = try_recv_via_connector(task, connectors, stage_id=1)
|
||||
|
||||
# Verify
|
||||
assert inputs == expected_data["engine_inputs"]
|
||||
assert rx_metrics is not None
|
||||
mock_connector.get.assert_called_once_with("0", "1", "req_recv", metadata={"handle": "xyz"})
|
||||
|
||||
|
||||
def test_recv_no_connector():
|
||||
"""Test recv fails when no connector exists for edge."""
|
||||
task = {"request_id": "req_missing", "from_connector": True, "from_stage": "0"}
|
||||
connectors = {} # Empty connectors
|
||||
|
||||
inputs, _ = try_recv_via_connector(task, connectors, stage_id=1)
|
||||
assert inputs is None
|
||||
|
||||
|
||||
def test_shm_connector_flow():
|
||||
"""
|
||||
Verify the full flow: Send -> Adapter -> Connector -> Adapter -> Recv.
|
||||
Using real SharedMemoryConnector (inline mode for simplicity).
|
||||
"""
|
||||
# 1. Setup Connector
|
||||
config = {"shm_threshold_bytes": 1024} # Large threshold to use inline
|
||||
connector = SharedMemoryConnector(config)
|
||||
connectors_map = {("0", "1"): connector}
|
||||
|
||||
# 2. Setup Data
|
||||
stage_id = 0
|
||||
next_stage_id = 1
|
||||
req_id = "flow_req"
|
||||
inputs = {"tokens": [10, 20, 30]}
|
||||
sampling_params = {"n": 1}
|
||||
|
||||
# Queue capture mechanism
|
||||
queue_capture = []
|
||||
|
||||
def mock_submit(payload):
|
||||
queue_capture.append(payload)
|
||||
|
||||
mock_metrics = MagicMock()
|
||||
|
||||
# 3. Send
|
||||
success = try_send_via_connector(
|
||||
connector=connector,
|
||||
stage_id=stage_id,
|
||||
next_stage_id=next_stage_id,
|
||||
req_id=req_id,
|
||||
next_inputs=inputs,
|
||||
sampling_params=sampling_params,
|
||||
original_prompt="prompt",
|
||||
next_stage_queue_submit_fn=mock_submit,
|
||||
metrics=mock_metrics,
|
||||
)
|
||||
assert success is True
|
||||
assert len(queue_capture) == 1
|
||||
|
||||
# 4. Recv
|
||||
# The 'task' is what would be popped from the queue
|
||||
received_task = queue_capture[0]
|
||||
|
||||
# Verify queue payload contains what we expect
|
||||
assert received_task["from_connector"] is True
|
||||
assert received_task["from_stage"] == "0"
|
||||
|
||||
# Decode
|
||||
decoded_inputs, _ = try_recv_via_connector(received_task, connectors_map, stage_id=1)
|
||||
|
||||
# 5. Verify Data Integrity
|
||||
assert decoded_inputs == inputs
|
||||
|
||||
|
||||
def test_get_connectors_for_stage():
|
||||
"""Test filtering logic for stage config."""
|
||||
# Config has edges: 0->1, 1->2
|
||||
config = OmniTransferConfig(connectors={("0", "1"): ConnectorSpec(name="C1"), ("1", "2"): ConnectorSpec(name="C2")})
|
||||
|
||||
# Get config for Stage 1
|
||||
# Stage 1 receives from 0 (input) and sends to 2 (output)
|
||||
# get_connectors_config_for_stage ONLY returns INPUT connectors for the worker to initialize
|
||||
|
||||
stage_config = get_connectors_config_for_stage(config, stage_id=1)
|
||||
|
||||
# Should contain "from_stage_0"
|
||||
assert "from_stage_0" in stage_config
|
||||
assert stage_config["from_stage_0"]["spec"]["name"] == "C1"
|
||||
|
||||
# Should NOT contain "from_stage_1" or related to output
|
||||
assert "from_stage_1" not in stage_config
|
||||
|
||||
# Verify Stage 2
|
||||
stage_2_config = get_connectors_config_for_stage(config, stage_id=2)
|
||||
assert "from_stage_1" in stage_2_config
|
||||
assert stage_2_config["from_stage_1"]["spec"]["name"] == "C2"
|
||||
|
||||
|
||||
def test_recv_with_missing_metadata():
|
||||
"""Test recv when queue payload is malformed (missing metadata)."""
|
||||
# Connector expects metadata but task doesn't have it
|
||||
task = {
|
||||
"request_id": "req_bad",
|
||||
"from_connector": True,
|
||||
"from_stage": "0",
|
||||
# Missing "connector_metadata"
|
||||
}
|
||||
mock_conn = MagicMock()
|
||||
# If get is called with None metadata, connector usually handles it or adapter handles exception
|
||||
mock_conn.get.side_effect = Exception("Get failed")
|
||||
|
||||
connectors = {("0", "1"): mock_conn}
|
||||
|
||||
inputs, _ = try_recv_via_connector(task, connectors, stage_id=1)
|
||||
assert inputs is None
|
||||
@@ -0,0 +1,99 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector
|
||||
from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory
|
||||
from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
|
||||
from vllm_omni.distributed.omni_connectors.utils.serialization import OmniSerializer
|
||||
|
||||
|
||||
def test_pickle_serialization():
|
||||
"""Test basic pickle serialization."""
|
||||
data = {"key": "value", "list": [1, 2, 3]}
|
||||
serialized = OmniSerializer.serialize(data, method="cloudpickle")
|
||||
assert isinstance(serialized, bytes)
|
||||
|
||||
deserialized = OmniSerializer.deserialize(serialized, method="cloudpickle")
|
||||
assert data == deserialized
|
||||
|
||||
|
||||
def test_create_shm_connector():
|
||||
"""Test creating SharedMemoryConnector via Factory."""
|
||||
spec = ConnectorSpec(name="SharedMemoryConnector", extra={"shm_threshold_bytes": 1024})
|
||||
connector = OmniConnectorFactory.create_connector(spec)
|
||||
assert isinstance(connector, SharedMemoryConnector)
|
||||
assert connector.threshold == 1024
|
||||
|
||||
|
||||
def test_create_unknown_connector():
|
||||
"""Test error when creating unknown connector."""
|
||||
spec = ConnectorSpec(name="UnknownConnector")
|
||||
with pytest.raises(ValueError):
|
||||
OmniConnectorFactory.create_connector(spec)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def shm_connector():
|
||||
config = {"shm_threshold_bytes": 100} # Small threshold for testing
|
||||
return SharedMemoryConnector(config)
|
||||
|
||||
|
||||
def test_put_get_inline(shm_connector):
|
||||
"""Test inline transfer for small data."""
|
||||
data = {"small": "data"}
|
||||
# Ensure data is smaller than threshold (100 bytes)
|
||||
|
||||
success, size, metadata = shm_connector.put("stage_0", "stage_1", "req_1", data)
|
||||
assert success is True
|
||||
assert "inline_bytes" in metadata
|
||||
assert "shm" not in metadata
|
||||
|
||||
# Retrieve
|
||||
retrieved_data, ret_size = shm_connector.get("stage_0", "stage_1", "req_1", metadata)
|
||||
assert data == retrieved_data
|
||||
assert size == ret_size
|
||||
|
||||
|
||||
def test_put_get_shm(shm_connector, monkeypatch):
|
||||
"""Test SHM transfer logic for large data (Mocked)."""
|
||||
# Create data larger than 100 bytes
|
||||
data = {"large": "x" * 200}
|
||||
|
||||
# Mock SHM return values
|
||||
mock_handle = {"name": "test_shm", "size": 200}
|
||||
mock_write = MagicMock(return_value=mock_handle)
|
||||
monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_connector.shm_write_bytes", mock_write)
|
||||
|
||||
# When reading, return the serialized bytes of the data
|
||||
serialized_data = shm_connector.serialize_obj(data)
|
||||
mock_read = MagicMock(return_value=serialized_data)
|
||||
monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_connector.shm_read_bytes", mock_read)
|
||||
|
||||
# Put
|
||||
success, size, metadata = shm_connector.put("stage_0", "stage_1", "req_2", data)
|
||||
|
||||
assert success is True
|
||||
# Should use SHM because data > threshold
|
||||
assert "shm" in metadata
|
||||
assert metadata["shm"] == mock_handle
|
||||
assert "inline_bytes" not in metadata
|
||||
|
||||
mock_write.assert_called_once()
|
||||
|
||||
# Get
|
||||
retrieved_data, ret_size = shm_connector.get("stage_0", "stage_1", "req_2", metadata)
|
||||
|
||||
assert data == retrieved_data
|
||||
mock_read.assert_called_once_with(mock_handle)
|
||||
|
||||
|
||||
def test_get_invalid_metadata(shm_connector):
|
||||
"""Test get with invalid metadata."""
|
||||
result = shm_connector.get("stage_0", "stage_1", "req_3", {})
|
||||
assert result is None
|
||||
|
||||
result = shm_connector.get("stage_0", "stage_1", "req_3", {"unknown": "format"})
|
||||
assert result is None
|
||||
@@ -0,0 +1,64 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
# Use the new import path for initialization utilities
|
||||
from vllm_omni.distributed.omni_connectors.utils.initialization import load_omni_transfer_config
|
||||
|
||||
|
||||
def get_config_files():
|
||||
"""Helper to find config files."""
|
||||
# Go up two levels from 'tests/distributed/omni_connectors' (approx) to 'vllm-omni' root
|
||||
# Adjust based on file location: vllm-omni/tests/distributed/omni_connectors/test_omni_connector_configs.py
|
||||
# This file is 4 levels deep from root if we count from tests?
|
||||
# vllm-omni/tests/distributed/omni_connectors -> parent -> distributed -> parent -> tests -> parent -> vllm-omni
|
||||
# Let's use resolve to be safe.
|
||||
|
||||
# Path(__file__) = .../vllm-omni/tests/distributed/omni_connectors/test_omni_connector_configs.py
|
||||
# .parent = omni_connectors
|
||||
# .parent = distributed
|
||||
# .parent = tests
|
||||
# .parent = vllm-omni
|
||||
|
||||
base_dir = Path(__file__).resolve().parent.parent.parent.parent
|
||||
config_dir = base_dir / "vllm_omni" / "model_executor" / "stage_configs"
|
||||
|
||||
if not config_dir.exists():
|
||||
return []
|
||||
|
||||
return list(config_dir.glob("qwen*.yaml"))
|
||||
|
||||
|
||||
# Collect files at module level for parametrization
|
||||
config_files = get_config_files()
|
||||
|
||||
|
||||
@pytest.mark.skipif(len(config_files) == 0, reason="No config files found or directory missing")
|
||||
@pytest.mark.parametrize("yaml_file", config_files, ids=lambda p: p.name)
|
||||
def test_load_qwen_yaml_configs(yaml_file):
|
||||
"""
|
||||
Scan and test loading of all qwen*.yaml config files.
|
||||
This ensures that existing stage configs are compatible with the OmniConnector system.
|
||||
"""
|
||||
print(f"Testing config load: {yaml_file.name}")
|
||||
try:
|
||||
# Attempt to load the config
|
||||
# default_shm_threshold doesn't matter much for loading correctness, using default
|
||||
config = load_omni_transfer_config(yaml_file)
|
||||
|
||||
assert config is not None, "Config should not be None"
|
||||
|
||||
# Basic validation
|
||||
# Note: Some configs might not have 'runtime' or 'connectors' section if they rely on auto-shm
|
||||
# but the load function should succeed regardless.
|
||||
|
||||
# If the config defines stages, we expect connectors to be populated (either explicit or auto SHM)
|
||||
# We can't strictly assert len(config.connectors) > 0 because a single stage pipeline might have 0 edges.
|
||||
|
||||
print(f" -> Successfully loaded. Connectors: {len(config.connectors)}")
|
||||
|
||||
except Exception as e:
|
||||
pytest.fail(f"Failed to load config {yaml_file.name}: {e}")
|
||||
File renamed without changes.
File renamed without changes.
+1
-1
@@ -1,6 +1,6 @@
|
||||
# stage config for running qwen2.5-omni with architecture of OmniLLM.
|
||||
|
||||
# The following config has been verified on 1x 24GB GPU (L4/RTX3090).
|
||||
# The following config has been verified on 2x 24GB GPU (L4/RTX3090/RTX4090).
|
||||
# This config is optimized for CI e2e tests.
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
|
||||
# Stage 0: Thinker (multimodal understanding + text generation)
|
||||
# Stage 1: Talker (text embeddings → 8-layer RVQ codec codes)
|
||||
# Stage 1: Talker (text embeddings → 16-layer RVQ codec codes)
|
||||
# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
|
||||
|
||||
# The following config has been verified on 2x H100-80G GPUs.
|
||||
+8
-4
@@ -18,15 +18,19 @@ from .utils import create_new_process_for_each_test
|
||||
models = ["Qwen/Qwen2.5-Omni-3B"]
|
||||
|
||||
# CI stage config optimized for 24GB GPU (L4/RTX3090)
|
||||
CI_STAGE_CONFIG_PATH = str(Path(__file__).parent / "stage_configs" / "qwen2_5_omni_ci.yaml")
|
||||
stage_configs = [str(Path(__file__).parent / "stage_configs" / "qwen2_5_omni_ci.yaml")]
|
||||
|
||||
# Create parameter combinations for model and stage config
|
||||
test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.parametrize("model", models)
|
||||
@pytest.mark.parametrize("test_config", test_params)
|
||||
@create_new_process_for_each_test()
|
||||
def test_mixed_modalities_to_audio(omni_runner: type[OmniRunner], model: str) -> None:
|
||||
def test_mixed_modalities_to_audio(omni_runner: type[OmniRunner], test_config: tuple[str, str]) -> None:
|
||||
"""Test processing audio, image, and video together, generating audio output."""
|
||||
with omni_runner(model, seed=42, stage_configs_path=CI_STAGE_CONFIG_PATH) as runner:
|
||||
model, stage_config_path = test_config
|
||||
with omni_runner(model, seed=42, stage_configs_path=stage_config_path) as runner:
|
||||
# Prepare multimodal inputs
|
||||
question = "What is recited in the audio? What is in this image? Describe the video briefly."
|
||||
audio = AudioAsset("mary_had_lamb").audio_and_sample_rate
|
||||
+10
-6
@@ -1,7 +1,7 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
E2E tests for Qwen3-Omni model with audio/video/image input and audio output.
|
||||
E2E offline tests for Omni model with video input and audio output.
|
||||
"""
|
||||
|
||||
import os
|
||||
@@ -10,20 +10,24 @@ from pathlib import Path
|
||||
import pytest
|
||||
from vllm.assets.video import VideoAsset
|
||||
|
||||
from ..multi_stages.conftest import OmniRunner
|
||||
from .conftest import OmniRunner
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
|
||||
|
||||
# CI stage config for 2xH100-80G GPUs
|
||||
CI_STAGE_CONFIG_PATH = str(Path(__file__).parent / "stage_configs" / "qwen3_omni_ci.yaml")
|
||||
stage_configs = [str(Path(__file__).parent / "stage_configs" / "qwen3_omni_ci.yaml")]
|
||||
|
||||
# Create parameter combinations for model and stage config
|
||||
test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", models)
|
||||
def test_video_to_audio(omni_runner: type[OmniRunner], model: str) -> None:
|
||||
@pytest.mark.parametrize("test_config", test_params)
|
||||
def test_video_to_audio(omni_runner: type[OmniRunner], test_config) -> None:
|
||||
"""Test processing video, generating audio output."""
|
||||
with omni_runner(model, seed=42, stage_configs_path=CI_STAGE_CONFIG_PATH) as runner:
|
||||
model, stage_config_path = test_config
|
||||
with omni_runner(model, seed=42, stage_configs_path=stage_config_path) as runner:
|
||||
# Prepare inputs
|
||||
question = "Describe the video briefly."
|
||||
video = VideoAsset(name="baby_reading", num_frames=4).np_ndarrays
|
||||
File renamed without changes.
@@ -0,0 +1,48 @@
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
# ruff: noqa: E402
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from vllm_omni import Omni
|
||||
|
||||
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
|
||||
|
||||
models = ["Wan-AI/Wan2.2-T2V-A14B-Diffusers"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", models)
|
||||
def test_video_diffusion_model(model_name: str):
|
||||
m = Omni(
|
||||
model=model_name,
|
||||
boundary_ratio=0.875,
|
||||
flow_shift=5.0,
|
||||
)
|
||||
# Use minimal settings for testing
|
||||
# num_frames must satisfy: num_frames % vae_scale_factor_temporal == 1
|
||||
# For Wan2.2, vae_scale_factor_temporal=4, so valid values are 5, 9, 13, 17, ...
|
||||
height = 480
|
||||
width = 640
|
||||
num_frames = 5
|
||||
frames = m.generate(
|
||||
"A cat sitting on a table",
|
||||
height=height,
|
||||
width=width,
|
||||
num_frames=num_frames,
|
||||
num_inference_steps=2,
|
||||
guidance_scale=1.0,
|
||||
generator=torch.Generator("cuda").manual_seed(42),
|
||||
)
|
||||
|
||||
assert frames is not None
|
||||
assert hasattr(frames, "shape")
|
||||
# frames shape: (batch, num_frames, height, width, channels)
|
||||
assert frames.shape[1] == num_frames
|
||||
assert frames.shape[2] == height
|
||||
assert frames.shape[3] == width
|
||||
File renamed without changes.
File renamed without changes.
@@ -0,0 +1,98 @@
|
||||
# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
|
||||
# Stage 0: Thinker (multimodal understanding + text generation)
|
||||
# Stage 1: Talker (text embeddings → 16-layer RVQ codec codes)
|
||||
# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
|
||||
|
||||
# The following config has been verified on 2x H100-80G GPUs.
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
runtime:
|
||||
devices: "0,1"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: thinker
|
||||
model_arch: Qwen3OmniMoeForConditionalGeneration
|
||||
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.6
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: latent # Output hidden states for talker
|
||||
distributed_executor_backend: "mp"
|
||||
enable_prefix_caching: false
|
||||
hf_config_name: thinker_config
|
||||
tensor_parallel_size: 2
|
||||
load_format: dummy
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
is_comprehension: true
|
||||
default_sampling_params:
|
||||
temperature: 0.4
|
||||
top_p: 0.9
|
||||
top_k: 1
|
||||
max_tokens: 100
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.05
|
||||
|
||||
- stage_id: 1
|
||||
runtime:
|
||||
devices: "1"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: talker
|
||||
model_arch: Qwen3OmniMoeForConditionalGeneration
|
||||
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.3
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: latent # Output codec codes for code2wav
|
||||
# tensor_parallel_size: 2
|
||||
enable_prefix_caching: false
|
||||
distributed_executor_backend: "mp"
|
||||
hf_config_name: talker_config
|
||||
load_format: dummy
|
||||
engine_input_source: [0]
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
|
||||
# final_output: true
|
||||
# final_output_type: text
|
||||
default_sampling_params:
|
||||
temperature: 0.9
|
||||
top_k: 50
|
||||
max_tokens: 100
|
||||
seed: 42
|
||||
detokenize: False
|
||||
repetition_penalty: 1.05
|
||||
stop_token_ids: [2150]
|
||||
|
||||
- stage_id: 2
|
||||
runtime:
|
||||
devices: "0"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: code2wav
|
||||
model_arch: Qwen3OmniMoeForConditionalGeneration
|
||||
worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: audio # Final output: audio waveform
|
||||
gpu_memory_utilization: 0.1
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 1000000
|
||||
hf_config_name: thinker_config
|
||||
load_format: dummy
|
||||
engine_input_source: [1]
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 200
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.1
|
||||
@@ -0,0 +1,206 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
E2E Online tests for Qwen3-Omni model with video input and audio output.
|
||||
"""
|
||||
|
||||
import os
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
from vllm.assets.video import VideoAsset
|
||||
from vllm.utils import get_open_port
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
|
||||
|
||||
# CI stage config for 2*H100-80G GPUs
|
||||
stage_configs = [str(Path(__file__).parent / "stage_configs" / "qwen3_omni_ci.yaml")]
|
||||
|
||||
# Create parameter combinations for model and stage config
|
||||
test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
|
||||
|
||||
|
||||
class OmniServer:
|
||||
"""Omniserver for vLLM-Omni tests."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: str,
|
||||
serve_args: list[str],
|
||||
*,
|
||||
env_dict: dict[str, str] | None = None,
|
||||
) -> None:
|
||||
self.model = model
|
||||
self.serve_args = serve_args
|
||||
self.env_dict = env_dict
|
||||
self.proc: subprocess.Popen | None = None
|
||||
self.host = "127.0.0.1"
|
||||
self.port = get_open_port()
|
||||
|
||||
def _start_server(self) -> None:
|
||||
"""Start the vLLM-Omni server subprocess."""
|
||||
env = os.environ.copy()
|
||||
env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
if self.env_dict is not None:
|
||||
env.update(self.env_dict)
|
||||
|
||||
cmd = [
|
||||
sys.executable,
|
||||
"-m",
|
||||
"vllm_omni.entrypoints.cli.main",
|
||||
"serve",
|
||||
self.model,
|
||||
"--omni",
|
||||
"--host",
|
||||
self.host,
|
||||
"--port",
|
||||
str(self.port),
|
||||
] + self.serve_args
|
||||
|
||||
print(f"Launching OmniServer with: {' '.join(cmd)}")
|
||||
self.proc = subprocess.Popen(
|
||||
cmd,
|
||||
env=env,
|
||||
cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), # Set working directory to vllm-omni root
|
||||
)
|
||||
|
||||
# Wait for server to be ready
|
||||
max_wait = 600 # 10 minutes
|
||||
start_time = time.time()
|
||||
while time.time() - start_time < max_wait:
|
||||
try:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
|
||||
sock.settimeout(1)
|
||||
result = sock.connect_ex((self.host, self.port))
|
||||
if result == 0:
|
||||
print(f"Server ready on {self.host}:{self.port}")
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
time.sleep(2)
|
||||
|
||||
raise RuntimeError(f"Server failed to start within {max_wait} seconds")
|
||||
|
||||
def __enter__(self):
|
||||
self._start_server()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
if self.proc:
|
||||
self.proc.terminate()
|
||||
try:
|
||||
self.proc.wait(timeout=30)
|
||||
except subprocess.TimeoutExpired:
|
||||
self.proc.kill()
|
||||
self.proc.wait()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def omni_server(request):
|
||||
"""Start vLLM-Omni server as a subprocess with actual model weights.
|
||||
Uses module scope so the server starts only once for all tests.
|
||||
Multi-stage initialization can take 10-20+ minutes.
|
||||
"""
|
||||
model, stage_config_path = request.param
|
||||
with OmniServer(model, ["--stage-configs-path", stage_config_path]) as server:
|
||||
yield server
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(omni_server):
|
||||
"""OpenAI client for the running vLLM-Omni server."""
|
||||
return openai.OpenAI(
|
||||
base_url=f"http://{omni_server.host}:{omni_server.port}/v1",
|
||||
api_key="EMPTY",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def base64_encoded_video() -> str:
|
||||
"""Base64 encoded video for testing."""
|
||||
import base64
|
||||
|
||||
video = VideoAsset(name="baby_reading", num_frames=4)
|
||||
with open(video.video_path, "rb") as f:
|
||||
content = f.read()
|
||||
return base64.b64encode(content).decode("utf-8")
|
||||
|
||||
|
||||
def get_system_prompt():
|
||||
return {
|
||||
"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."
|
||||
),
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def dummy_messages_from_video_data(
|
||||
video_data_url: str,
|
||||
content_text: str = "Describe the video briefly.",
|
||||
):
|
||||
"""Create messages with video data URL for OpenAI API."""
|
||||
return [
|
||||
get_system_prompt(),
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "video_url", "video_url": {"url": video_data_url}},
|
||||
{"type": "text", "text": content_text},
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_video_to_audio(
|
||||
client: openai.OpenAI,
|
||||
omni_server,
|
||||
base64_encoded_video: str,
|
||||
) -> None:
|
||||
"""Test processing video, generating audio output via OpenAI API."""
|
||||
# Create data URL for the base64 encoded video
|
||||
video_data_url = f"data:video/mp4;base64,{base64_encoded_video}"
|
||||
|
||||
messages = dummy_messages_from_video_data(video_data_url)
|
||||
|
||||
# Test single completion
|
||||
chat_completion = client.chat.completions.create(
|
||||
model=omni_server.model,
|
||||
messages=messages,
|
||||
)
|
||||
|
||||
assert len(chat_completion.choices) == 2 # 1 for text output, 1 for audio output
|
||||
|
||||
# Verify text output
|
||||
text_choice = chat_completion.choices[0]
|
||||
assert text_choice.finish_reason == "length"
|
||||
|
||||
# Verify we got a response
|
||||
text_message = text_choice.message
|
||||
assert text_message.content is not None and len(text_message.content) >= 10
|
||||
assert text_message.role == "assistant"
|
||||
|
||||
# Verify audio output
|
||||
audio_choice = chat_completion.choices[1]
|
||||
assert audio_choice.finish_reason == "stop"
|
||||
audio_message = audio_choice.message
|
||||
|
||||
# Check if audio was generated
|
||||
if hasattr(audio_message, "audio") and audio_message.audio:
|
||||
assert audio_message.audio.data is not None
|
||||
assert len(audio_message.audio.data) > 0
|
||||
@@ -0,0 +1,58 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.entrypoints.stage_utils import set_stage_devices
|
||||
|
||||
|
||||
def _make_dummy_torch(call_log):
|
||||
class _Props:
|
||||
def __init__(self, total):
|
||||
self.total_memory = total
|
||||
|
||||
class _Cuda:
|
||||
@staticmethod
|
||||
def is_available():
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def set_device(idx):
|
||||
call_log.append(idx)
|
||||
|
||||
@staticmethod
|
||||
def device_count():
|
||||
return 2
|
||||
|
||||
@staticmethod
|
||||
def get_device_properties(idx):
|
||||
return _Props(total=16000)
|
||||
|
||||
@staticmethod
|
||||
def mem_get_info(idx):
|
||||
return (8000, 16000)
|
||||
|
||||
@staticmethod
|
||||
def get_device_name(idx):
|
||||
return f"gpu-{idx}"
|
||||
|
||||
class _Torch:
|
||||
cuda = _Cuda
|
||||
|
||||
return _Torch
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("clean_gpu_memory_between_tests")
|
||||
def test_set_stage_devices_respects_logical_ids(monkeypatch):
|
||||
# Preserve an existing logical mapping and ensure devices "0,1" map through it.
|
||||
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "6,7")
|
||||
call_log: list[int] = []
|
||||
dummy_torch = _make_dummy_torch(call_log)
|
||||
monkeypatch.setitem(sys.modules, "torch", dummy_torch)
|
||||
monkeypatch.setattr("vllm_omni.utils.detect_device_type", lambda: "cuda")
|
||||
monkeypatch.setattr("vllm_omni.utils.get_device_control_env_var", lambda: "CUDA_VISIBLE_DEVICES")
|
||||
|
||||
set_stage_devices(stage_id=0, devices="0,1")
|
||||
|
||||
assert os.environ["CUDA_VISIBLE_DEVICES"] == "6,7"
|
||||
assert call_log and call_log[0] == 0 # current device set after remap
|
||||
@@ -1,14 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Pytest configuration for multi_stages_h100 tests.
|
||||
|
||||
This conftest imports fixtures from the parent multi_stages conftest.
|
||||
"""
|
||||
|
||||
# Import the omni_runner fixture from the parent directory's conftest
|
||||
# This makes the fixture available to tests in this directory
|
||||
from ..multi_stages.conftest import omni_runner # noqa: F401
|
||||
|
||||
# Explicitly expose the fixture for pytest discovery
|
||||
__all__ = ["omni_runner"]
|
||||
@@ -104,6 +104,7 @@ class _FakeStage:
|
||||
shm_threshold_bytes: int = 65536,
|
||||
ctx=None,
|
||||
batch_timeout: int = 10,
|
||||
**kwargs,
|
||||
):
|
||||
"""Mock init_stage_worker: don't start real process, just send stage_ready message."""
|
||||
# Create a mock process object
|
||||
|
||||
+41
@@ -0,0 +1,41 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Cache module for diffusion model inference acceleration.
|
||||
|
||||
This module provides a unified cache adapter system for different caching strategies:
|
||||
- TeaCache: Timestep Embedding Aware Cache for adaptive transformer caching
|
||||
- (Future: DeepCache, DistriFusion, etc.)
|
||||
|
||||
The cache adapter system uses a registry pattern where cache types can be
|
||||
dynamically registered and instantiated. Cache is configured via OmniDiffusionConfig
|
||||
or DIFFUSION_CACHE_ADAPTER environment variable.
|
||||
"""
|
||||
|
||||
from vllm_omni.diffusion.cache.apply import setup_cache
|
||||
from vllm_omni.diffusion.cache.base import CacheAdapter
|
||||
from vllm_omni.diffusion.cache.registry import (
|
||||
CacheType,
|
||||
get_cache_adapter,
|
||||
register_cache_adapter,
|
||||
)
|
||||
from vllm_omni.diffusion.cache.teacache import (
|
||||
CacheContext,
|
||||
TeaCacheConfig,
|
||||
apply_teacache_hook,
|
||||
)
|
||||
|
||||
# Import teacache adapter to trigger registration
|
||||
from vllm_omni.diffusion.cache.teacache.adapter import TeaCacheAdapter # noqa: F401
|
||||
|
||||
__all__ = [
|
||||
"CacheAdapter",
|
||||
"CacheType",
|
||||
"TeaCacheConfig",
|
||||
"CacheContext",
|
||||
"setup_cache",
|
||||
"get_cache_adapter",
|
||||
"register_cache_adapter",
|
||||
"apply_teacache_hook",
|
||||
]
|
||||
Vendored
+82
@@ -0,0 +1,82 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Unified cache setup function for diffusion models.
|
||||
|
||||
This module provides the main entry point for setting up cache adapters
|
||||
on transformer models. It handles cache type selection, adapter instantiation,
|
||||
and application to the model.
|
||||
"""
|
||||
|
||||
from typing import Any, Optional
|
||||
|
||||
import torch
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.diffusion.cache.base import CacheAdapter
|
||||
from vllm_omni.diffusion.cache.registry import get_cache_adapter
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def setup_cache(
|
||||
transformer: torch.nn.Module,
|
||||
cache_type: str = "none",
|
||||
cache_config: Optional[dict[str, Any]] = None,
|
||||
) -> Optional[CacheAdapter]:
|
||||
"""
|
||||
Setup cache adapter for transformer (one-time initialization).
|
||||
|
||||
This is the main entry point for applying caching to diffusion transformers.
|
||||
It's called once during pipeline initialization and applies the appropriate
|
||||
cache adapter based on configuration.
|
||||
|
||||
Args:
|
||||
transformer: Transformer module to apply cache to
|
||||
cache_type: Type of cache to apply. Options:
|
||||
- "none": No caching (default)
|
||||
- "tea_cache": TeaCache adaptive caching
|
||||
- Additional types can be registered via register_cache_adapter()
|
||||
cache_config: Cache-specific configuration dictionary. Contents depend
|
||||
on cache type:
|
||||
- For "tea_cache": {"rel_l1_thresh": 0.2, "model_type": "QwenImagePipeline"}
|
||||
- For other types: see their respective adapter documentation
|
||||
|
||||
Returns:
|
||||
CacheAdapter instance for state management, or None if cache_type="none"
|
||||
|
||||
Example:
|
||||
>>> # Setup TeaCache during pipeline initialization
|
||||
>>> adapter = setup_cache(
|
||||
... transformer,
|
||||
... cache_type="tea_cache",
|
||||
... cache_config={"rel_l1_thresh": 0.2}
|
||||
... )
|
||||
>>>
|
||||
>>> # Reset cache state before each generation
|
||||
>>> if adapter is not None:
|
||||
... adapter.reset(transformer)
|
||||
"""
|
||||
# No caching requested
|
||||
if cache_type == "none":
|
||||
logger.debug("No cache adapter requested (cache_type='none')")
|
||||
return None
|
||||
|
||||
# Get adapter instance from registry
|
||||
try:
|
||||
adapter = get_cache_adapter(cache_type, cache_config or {})
|
||||
except ValueError as e:
|
||||
logger.error(f"Failed to get cache adapter: {e}")
|
||||
raise
|
||||
|
||||
# Apply cache to transformer
|
||||
logger.info(f"Applying {cache_type} cache to transformer")
|
||||
try:
|
||||
adapter.apply(transformer)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to apply cache adapter: {e}")
|
||||
raise
|
||||
|
||||
logger.info(f"Cache adapter '{cache_type}' successfully applied")
|
||||
return adapter
|
||||
Vendored
+67
@@ -0,0 +1,67 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Base cache adapter interface for diffusion models.
|
||||
|
||||
This module defines the abstract base class that all cache adapters must implement.
|
||||
Cache adapters provide a unified interface for applying different caching strategies
|
||||
(TeaCache, DeepCache, etc.) to transformer models using hooks.
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class CacheAdapter(ABC):
|
||||
"""
|
||||
Abstract base class for cache adapters.
|
||||
|
||||
Cache adapters apply caching strategies to transformer models to accelerate
|
||||
inference. Each adapter type (TeaCache, DeepCache, etc.) implements the
|
||||
apply() and reset() methods to manage cache lifecycle.
|
||||
|
||||
Attributes:
|
||||
config: Dictionary containing cache-specific configuration parameters
|
||||
"""
|
||||
|
||||
def __init__(self, config: dict[str, Any]):
|
||||
"""
|
||||
Initialize cache adapter with configuration.
|
||||
|
||||
Args:
|
||||
config: Cache-specific configuration dictionary
|
||||
"""
|
||||
self.config = config
|
||||
|
||||
@abstractmethod
|
||||
def apply(self, transformer: torch.nn.Module) -> None:
|
||||
"""
|
||||
Apply cache to transformer using hooks.
|
||||
|
||||
This method should register the appropriate hooks on the transformer
|
||||
to enable caching during inference. Called once during pipeline
|
||||
initialization.
|
||||
|
||||
Args:
|
||||
transformer: Transformer module to apply cache to
|
||||
"""
|
||||
raise NotImplementedError("Subclasses must implement apply()")
|
||||
|
||||
@abstractmethod
|
||||
def reset(self, transformer: torch.nn.Module) -> None:
|
||||
"""
|
||||
Reset cache state for new generation.
|
||||
|
||||
This method should clear any cached values and reset counters/accumulators.
|
||||
Called at the start of each generation to ensure clean state.
|
||||
|
||||
Args:
|
||||
transformer: Transformer module to reset cache on
|
||||
"""
|
||||
raise NotImplementedError("Subclasses must implement reset()")
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"{self.__class__.__name__}(config={self.config})"
|
||||
+100
@@ -0,0 +1,100 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Cache adapter registry for diffusion models.
|
||||
|
||||
This module provides a registry pattern for cache adapters, allowing dynamic
|
||||
registration and instantiation of different cache types (TeaCache, DeepCache, etc.).
|
||||
"""
|
||||
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.diffusion.cache.base import CacheAdapter
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class CacheType(Enum):
|
||||
"""Supported cache adapter types."""
|
||||
|
||||
NONE = "none"
|
||||
TEA_CACHE = "tea_cache"
|
||||
# Future cache types can be added here:
|
||||
# DEEP_CACHE = "deep_cache"
|
||||
# DISTRI_FUSION = "distri_fusion"
|
||||
|
||||
|
||||
# Global registry mapping cache types to adapter classes
|
||||
CACHE_ADAPTER_REGISTRY: dict[CacheType, type[CacheAdapter]] = {}
|
||||
|
||||
|
||||
def register_cache_adapter(cache_type: CacheType, adapter_class: type[CacheAdapter]) -> None:
|
||||
"""
|
||||
Register a cache adapter class for a given cache type.
|
||||
|
||||
This allows extending the system with new cache types without modifying
|
||||
the core cache infrastructure.
|
||||
|
||||
Args:
|
||||
cache_type: CacheType enum value
|
||||
adapter_class: CacheAdapter subclass to register
|
||||
|
||||
Example:
|
||||
>>> register_cache_adapter(CacheType.TEA_CACHE, TeaCacheAdapter)
|
||||
"""
|
||||
if not issubclass(adapter_class, CacheAdapter):
|
||||
raise TypeError(f"{adapter_class} must be a subclass of CacheAdapter")
|
||||
|
||||
CACHE_ADAPTER_REGISTRY[cache_type] = adapter_class
|
||||
logger.debug(f"Registered cache adapter: {cache_type.value} -> {adapter_class.__name__}")
|
||||
|
||||
|
||||
def get_cache_adapter(cache_type: str, config: dict[str, Any]) -> CacheAdapter:
|
||||
"""
|
||||
Factory function to get cache adapter instance.
|
||||
|
||||
Converts string cache type to enum, looks up in registry, and instantiates
|
||||
the appropriate adapter class with the provided configuration.
|
||||
|
||||
Args:
|
||||
cache_type: String name of cache type ("tea_cache", "deep_cache", etc.)
|
||||
config: Configuration dictionary to pass to adapter constructor
|
||||
|
||||
Returns:
|
||||
Instantiated CacheAdapter subclass
|
||||
|
||||
Raises:
|
||||
ValueError: If cache_type is unknown or not registered
|
||||
|
||||
Example:
|
||||
>>> adapter = get_cache_adapter("tea_cache", {"rel_l1_thresh": 0.2})
|
||||
>>> adapter.apply(transformer)
|
||||
"""
|
||||
# Normalize cache type string
|
||||
cache_type_str = cache_type.lower().strip()
|
||||
|
||||
# Convert string to enum
|
||||
try:
|
||||
cache_enum = CacheType(cache_type_str)
|
||||
except ValueError:
|
||||
available = [ct.value for ct in CacheType if ct != CacheType.NONE]
|
||||
raise ValueError(f"Unknown cache type: '{cache_type}'. Available types: {available}")
|
||||
|
||||
# Check if it's the special "none" case
|
||||
if cache_enum == CacheType.NONE:
|
||||
raise ValueError("Cannot instantiate adapter for cache_type='none'. Use setup_cache() which handles this case.")
|
||||
|
||||
# Lookup in registry
|
||||
if cache_enum not in CACHE_ADAPTER_REGISTRY:
|
||||
raise ValueError(
|
||||
f"Cache type '{cache_type}' is not registered. Registered types: {list(CACHE_ADAPTER_REGISTRY.keys())}"
|
||||
)
|
||||
|
||||
adapter_class = CACHE_ADAPTER_REGISTRY[cache_enum]
|
||||
|
||||
# Instantiate and return
|
||||
return adapter_class(config)
|
||||
+45
@@ -0,0 +1,45 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
TeaCache: Timestep Embedding Aware Cache for diffusion model acceleration.
|
||||
|
||||
TeaCache speeds up diffusion inference by reusing transformer block computations
|
||||
when consecutive timestep embeddings are similar.
|
||||
|
||||
This implementation uses a hooks-based approach that requires zero changes to
|
||||
model code. Model developers only need to add an extractor function to support
|
||||
new models.
|
||||
|
||||
Usage:
|
||||
from vllm_omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
cache_adapter="tea_cache",
|
||||
cache_config={"rel_l1_thresh": 0.2}
|
||||
)
|
||||
images = omni.generate("a cat")
|
||||
|
||||
# Alternative: Using environment variable
|
||||
# export DIFFUSION_CACHE_ADAPTER=tea_cache
|
||||
"""
|
||||
|
||||
from vllm_omni.diffusion.cache.teacache.adapter import TeaCacheAdapter
|
||||
from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
|
||||
from vllm_omni.diffusion.cache.teacache.extractors import (
|
||||
CacheContext,
|
||||
register_extractor,
|
||||
)
|
||||
from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook, apply_teacache_hook
|
||||
from vllm_omni.diffusion.cache.teacache.state import TeaCacheState
|
||||
|
||||
__all__ = [
|
||||
"TeaCacheAdapter",
|
||||
"TeaCacheConfig",
|
||||
"TeaCacheState",
|
||||
"TeaCacheHook",
|
||||
"apply_teacache_hook",
|
||||
"register_extractor",
|
||||
"CacheContext",
|
||||
]
|
||||
+96
@@ -0,0 +1,96 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
TeaCache adapter implementation.
|
||||
|
||||
This module provides the TeaCache adapter that implements the CacheAdapter
|
||||
interface using the hooks-based TeaCache system.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.diffusion.cache.base import CacheAdapter
|
||||
from vllm_omni.diffusion.cache.registry import CacheType, register_cache_adapter
|
||||
from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
|
||||
from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook, apply_teacache_hook
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class TeaCacheAdapter(CacheAdapter):
|
||||
"""
|
||||
TeaCache implementation using hooks.
|
||||
|
||||
TeaCache (Timestep Embedding Aware Cache) is an adaptive caching technique
|
||||
that speeds up diffusion inference by reusing transformer block computations
|
||||
when consecutive timestep embeddings are similar.
|
||||
|
||||
The adapter applies TeaCache hooks to the transformer which intercept the
|
||||
forward pass and implement the caching logic transparently.
|
||||
|
||||
Example:
|
||||
>>> adapter = TeaCacheAdapter({"rel_l1_thresh": 0.2})
|
||||
>>> adapter.apply(transformer)
|
||||
>>> # Generate with cache enabled
|
||||
>>> adapter.reset(transformer) # Reset before each generation
|
||||
"""
|
||||
|
||||
def apply(self, transformer: torch.nn.Module) -> None:
|
||||
"""
|
||||
Apply TeaCache to transformer using hooks.
|
||||
|
||||
This creates a TeaCacheConfig from the adapter's config dict and
|
||||
applies the TeaCache hook to the transformer.
|
||||
|
||||
Args:
|
||||
transformer: Transformer module to apply TeaCache to
|
||||
"""
|
||||
# Validate that model_type is provided
|
||||
if "model_type" not in self.config:
|
||||
raise ValueError(
|
||||
"model_type must be provided in cache_config. "
|
||||
"Should be set to OmniDiffusionConfig.model_class_name (e.g., 'QwenImagePipeline')"
|
||||
)
|
||||
|
||||
# Create TeaCacheConfig from dict
|
||||
try:
|
||||
teacache_config = TeaCacheConfig(**self.config)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create TeaCacheConfig: {e}")
|
||||
raise ValueError(
|
||||
f"Invalid TeaCache configuration: {e}. "
|
||||
f"Expected keys: rel_l1_thresh, model_type (pipeline class name), coefficients (optional)"
|
||||
)
|
||||
|
||||
# Apply hook to transformer
|
||||
apply_teacache_hook(transformer, teacache_config)
|
||||
logger.info(
|
||||
f"TeaCache applied with rel_l1_thresh={teacache_config.rel_l1_thresh}, "
|
||||
f"model_type={teacache_config.model_type}"
|
||||
)
|
||||
|
||||
def reset(self, transformer: torch.nn.Module) -> None:
|
||||
"""
|
||||
Reset TeaCache state for new generation.
|
||||
|
||||
Clears all cached residuals and resets counters/accumulators.
|
||||
Should be called before each generation to ensure clean state.
|
||||
|
||||
Args:
|
||||
transformer: Transformer module to reset cache on
|
||||
"""
|
||||
if hasattr(transformer, "_hook_registry"):
|
||||
hook = transformer._hook_registry.get_hook(TeaCacheHook._HOOK_NAME)
|
||||
if hook is not None:
|
||||
transformer._hook_registry.reset_hook(TeaCacheHook._HOOK_NAME)
|
||||
logger.debug("TeaCache state reset")
|
||||
else:
|
||||
logger.warning("TeaCache hook not found, nothing to reset")
|
||||
else:
|
||||
logger.warning("Transformer has no hook registry, TeaCache may not be applied")
|
||||
|
||||
|
||||
# Register TeaCache adapter in the global registry
|
||||
register_cache_adapter(CacheType.TEA_CACHE, TeaCacheAdapter)
|
||||
+67
@@ -0,0 +1,67 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
# Model-specific polynomial coefficients for rescaling L1 distances
|
||||
# These coefficients account for model-specific characteristics in how embeddings change
|
||||
# Source: TeaCache paper and ComfyUI-TeaCache empirical tuning
|
||||
_MODEL_COEFFICIENTS = {
|
||||
# FLUX model coefficients from TeaCache paper
|
||||
"FluxPipeline": [
|
||||
4.98651651e02,
|
||||
-2.83781631e02,
|
||||
5.58554382e01,
|
||||
-3.82021401e00,
|
||||
2.64230861e-01,
|
||||
],
|
||||
# Qwen-Image model coefficients from ComfyUI-TeaCache
|
||||
# Tuned specifically for Qwen's dual-stream transformer architecture
|
||||
"QwenImagePipeline": [
|
||||
-4.50000000e02,
|
||||
2.80000000e02,
|
||||
-4.50000000e01,
|
||||
3.20000000e00,
|
||||
-2.00000000e-02,
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class TeaCacheConfig:
|
||||
"""
|
||||
Configuration for TeaCache applied to transformer models.
|
||||
|
||||
TeaCache (Timestep Embedding Aware Cache) is an adaptive caching technique that speeds up
|
||||
diffusion model inference by reusing transformer block computations when consecutive
|
||||
timestep embeddings are similar.
|
||||
|
||||
Args:
|
||||
rel_l1_thresh: Threshold for accumulated relative L1 distance. When below threshold,
|
||||
cached residual is reused. Values in [0.1, 0.3] work best:
|
||||
- 0.2: ~1.5x speedup with minimal quality loss
|
||||
- 0.4: ~1.8x speedup with slight quality loss
|
||||
- 0.6: ~2.0x speedup with noticeable quality loss
|
||||
coefficients: Polynomial coefficients for rescaling L1 distance. If None, uses
|
||||
model-specific defaults based on model_type.
|
||||
model_type: Pipeline class name (e.g., "QwenImagePipeline", "FluxPipeline").
|
||||
Must match OmniDiffusionConfig.model_class_name for proper extractor lookup.
|
||||
Defaults to "QwenImagePipeline".
|
||||
"""
|
||||
|
||||
rel_l1_thresh: float = 0.2
|
||||
coefficients: Optional[list[float]] = None
|
||||
model_type: str = "QwenImagePipeline"
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
"""Validate and set default coefficients."""
|
||||
if self.rel_l1_thresh <= 0:
|
||||
raise ValueError(f"rel_l1_thresh must be positive, got {self.rel_l1_thresh}")
|
||||
|
||||
if self.coefficients is None:
|
||||
# Use model-specific coefficients, fallback to FluxPipeline if not found
|
||||
self.coefficients = _MODEL_COEFFICIENTS.get(self.model_type, _MODEL_COEFFICIENTS["FluxPipeline"])
|
||||
|
||||
if len(self.coefficients) != 5:
|
||||
raise ValueError(f"coefficients must contain exactly 5 elements, got {len(self.coefficients)}")
|
||||
+334
@@ -0,0 +1,334 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Model-specific extractors for TeaCache.
|
||||
|
||||
This module provides a registry of extractor functions that know how to extract
|
||||
modulated inputs from different transformer architectures. Adding support for
|
||||
a new model requires only adding a new extractor function to the registry.
|
||||
|
||||
With Option B enhancement, extractors now return a CacheContext object containing
|
||||
all model-specific information needed for generic caching, including preprocessing,
|
||||
transformer execution, and postprocessing logic.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Callable, Union
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
@dataclass
|
||||
class CacheContext:
|
||||
"""
|
||||
Context object containing all model-specific information for caching.
|
||||
|
||||
This allows the TeaCacheHook to remain completely generic - all model-specific
|
||||
logic is encapsulated in the extractor that returns this context.
|
||||
|
||||
Attributes:
|
||||
modulated_input: Tensor used for cache decision (similarity comparison).
|
||||
Must be a torch.Tensor extracted from the first transformer block,
|
||||
typically after applying normalization and modulation.
|
||||
|
||||
hidden_states: Current hidden states (will be modified by caching).
|
||||
Must be a torch.Tensor representing the main image/latent states
|
||||
after preprocessing but before transformer blocks.
|
||||
|
||||
encoder_hidden_states: Optional encoder states (for dual-stream models).
|
||||
Set to None for single-stream models (e.g., Flux).
|
||||
For dual-stream models (e.g., Qwen), contains text encoder outputs.
|
||||
|
||||
temb: Timestep embedding tensor.
|
||||
Must be a torch.Tensor containing the timestep conditioning.
|
||||
|
||||
run_transformer_blocks: Callable that executes model-specific transformer blocks.
|
||||
Signature: () -> tuple[torch.Tensor, ...]
|
||||
|
||||
Returns:
|
||||
tuple containing:
|
||||
- [0]: processed hidden_states (required)
|
||||
- [1]: processed encoder_hidden_states (optional, only for dual-stream)
|
||||
|
||||
Example for single-stream:
|
||||
def run_blocks():
|
||||
h = hidden_states
|
||||
for block in module.transformer_blocks:
|
||||
h = block(h, temb=temb)
|
||||
return (h,)
|
||||
|
||||
Example for dual-stream:
|
||||
def run_blocks():
|
||||
h, e = hidden_states, encoder_hidden_states
|
||||
for block in module.transformer_blocks:
|
||||
e, h = block(h, e, temb=temb)
|
||||
return (h, e)
|
||||
|
||||
postprocess: Callable that does model-specific output postprocessing.
|
||||
Signature: (torch.Tensor) -> Union[torch.Tensor, Transformer2DModelOutput, tuple]
|
||||
|
||||
Takes the processed hidden_states and applies final transformations
|
||||
(normalization, projection) to produce the model output.
|
||||
|
||||
Example:
|
||||
def postprocess(h):
|
||||
h = module.norm_out(h, temb)
|
||||
output = module.proj_out(h)
|
||||
return Transformer2DModelOutput(sample=output)
|
||||
|
||||
extra_states: Optional dict for additional model-specific state.
|
||||
Use this for models that need to pass additional context beyond
|
||||
the standard fields.
|
||||
"""
|
||||
|
||||
modulated_input: torch.Tensor
|
||||
hidden_states: torch.Tensor
|
||||
encoder_hidden_states: torch.Tensor | None
|
||||
temb: torch.Tensor
|
||||
run_transformer_blocks: Callable[[], tuple[torch.Tensor, ...]]
|
||||
postprocess: Callable[[torch.Tensor], Any]
|
||||
extra_states: dict[str, Any] | None = None
|
||||
|
||||
def validate(self) -> None:
|
||||
"""
|
||||
Validate that the CacheContext contains valid data.
|
||||
|
||||
Raises:
|
||||
TypeError: If fields have wrong types
|
||||
ValueError: If tensors have invalid properties
|
||||
RuntimeError: If callables fail basic invocation tests
|
||||
|
||||
This method should be called after creating a CacheContext to catch
|
||||
common developer errors early with clear error messages.
|
||||
"""
|
||||
# Validate tensor fields
|
||||
if not isinstance(self.modulated_input, torch.Tensor):
|
||||
raise TypeError(f"modulated_input must be torch.Tensor, got {type(self.modulated_input)}")
|
||||
|
||||
if not isinstance(self.hidden_states, torch.Tensor):
|
||||
raise TypeError(f"hidden_states must be torch.Tensor, got {type(self.hidden_states)}")
|
||||
|
||||
if self.encoder_hidden_states is not None and not isinstance(self.encoder_hidden_states, torch.Tensor):
|
||||
raise TypeError(
|
||||
f"encoder_hidden_states must be torch.Tensor or None, got {type(self.encoder_hidden_states)}"
|
||||
)
|
||||
|
||||
if not isinstance(self.temb, torch.Tensor):
|
||||
raise TypeError(f"temb must be torch.Tensor, got {type(self.temb)}")
|
||||
|
||||
# Validate callables
|
||||
if not callable(self.run_transformer_blocks):
|
||||
raise TypeError(f"run_transformer_blocks must be callable, got {type(self.run_transformer_blocks)}")
|
||||
|
||||
if not callable(self.postprocess):
|
||||
raise TypeError(f"postprocess must be callable, got {type(self.postprocess)}")
|
||||
|
||||
# Validate tensor shapes are compatible
|
||||
if self.modulated_input.shape[0] != self.hidden_states.shape[0]:
|
||||
raise ValueError(
|
||||
f"Batch size mismatch: modulated_input has batch size "
|
||||
f"{self.modulated_input.shape[0]}, but hidden_states has "
|
||||
f"{self.hidden_states.shape[0]}"
|
||||
)
|
||||
|
||||
# Validate devices match
|
||||
if self.modulated_input.device != self.hidden_states.device:
|
||||
raise ValueError(
|
||||
f"Device mismatch: modulated_input on {self.modulated_input.device}, "
|
||||
f"hidden_states on {self.hidden_states.device}"
|
||||
)
|
||||
|
||||
|
||||
def extract_qwen_context(
|
||||
module: nn.Module,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
encoder_hidden_states_mask: torch.Tensor,
|
||||
timestep: Union[torch.Tensor, float, int],
|
||||
img_shapes: torch.Tensor,
|
||||
txt_seq_lens: torch.Tensor,
|
||||
guidance: torch.Tensor | None = None,
|
||||
attention_kwargs: dict[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> CacheContext:
|
||||
"""
|
||||
Extract cache context for QwenImageTransformer2DModel.
|
||||
|
||||
This is the ONLY Qwen-specific code needed for TeaCache support.
|
||||
It encapsulates preprocessing, modulated input extraction, transformer execution,
|
||||
and postprocessing logic.
|
||||
|
||||
Args:
|
||||
module: QwenImageTransformer2DModel instance
|
||||
hidden_states: Input hidden states tensor
|
||||
encoder_hidden_states: Text encoder outputs
|
||||
encoder_hidden_states_mask: Mask for text encoder
|
||||
timestep: Current diffusion timestep
|
||||
img_shapes: Image shapes for position embedding
|
||||
txt_seq_lens: Text sequence lengths
|
||||
guidance: Optional guidance scale for CFG
|
||||
attention_kwargs: Additional attention arguments
|
||||
**kwargs: Additional keyword arguments ignored by this extractor
|
||||
|
||||
Returns:
|
||||
CacheContext with all information needed for generic caching
|
||||
"""
|
||||
from diffusers.models.modeling_outputs import Transformer2DModelOutput
|
||||
|
||||
if not hasattr(module, "transformer_blocks") or len(module.transformer_blocks) == 0:
|
||||
raise ValueError("Module must have transformer_blocks")
|
||||
|
||||
# ============================================================================
|
||||
# PREPROCESSING (Qwen-specific)
|
||||
# ============================================================================
|
||||
hidden_states = module.img_in(hidden_states)
|
||||
timestep = timestep.to(device=hidden_states.device, dtype=hidden_states.dtype)
|
||||
encoder_hidden_states = module.txt_norm(encoder_hidden_states)
|
||||
encoder_hidden_states = module.txt_in(encoder_hidden_states)
|
||||
|
||||
if guidance is not None:
|
||||
guidance = guidance.to(hidden_states.dtype) * 1000
|
||||
|
||||
temb = (
|
||||
module.time_text_embed(timestep, hidden_states)
|
||||
if guidance is None
|
||||
else module.time_text_embed(timestep, guidance, hidden_states)
|
||||
)
|
||||
|
||||
image_rotary_emb = module.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device)
|
||||
|
||||
# ============================================================================
|
||||
# EXTRACT MODULATED INPUT (for cache decision)
|
||||
# ============================================================================
|
||||
block = module.transformer_blocks[0]
|
||||
img_mod_params = block.img_mod(temb)
|
||||
img_mod1, _ = img_mod_params.chunk(2, dim=-1)
|
||||
img_normed = block.img_norm1(hidden_states)
|
||||
img_modulated, _ = block._modulate(img_normed, img_mod1)
|
||||
|
||||
# ============================================================================
|
||||
# DEFINE TRANSFORMER EXECUTION (Qwen-specific)
|
||||
# ============================================================================
|
||||
def run_transformer_blocks():
|
||||
"""Execute all Qwen transformer blocks."""
|
||||
h = hidden_states
|
||||
e = encoder_hidden_states
|
||||
for block in module.transformer_blocks:
|
||||
e, h = block(
|
||||
hidden_states=h,
|
||||
encoder_hidden_states=e,
|
||||
encoder_hidden_states_mask=encoder_hidden_states_mask,
|
||||
temb=temb,
|
||||
image_rotary_emb=image_rotary_emb,
|
||||
joint_attention_kwargs=attention_kwargs,
|
||||
)
|
||||
return (h, e)
|
||||
|
||||
# ============================================================================
|
||||
# DEFINE POSTPROCESSING (Qwen-specific)
|
||||
# ============================================================================
|
||||
return_dict = kwargs.get("return_dict", True)
|
||||
|
||||
def postprocess(h):
|
||||
"""Apply Qwen-specific output postprocessing."""
|
||||
h = module.norm_out(h, temb)
|
||||
output = module.proj_out(h)
|
||||
if not return_dict:
|
||||
return (output,)
|
||||
return Transformer2DModelOutput(sample=output)
|
||||
|
||||
# ============================================================================
|
||||
# RETURN CONTEXT
|
||||
# ============================================================================
|
||||
return CacheContext(
|
||||
modulated_input=img_modulated,
|
||||
hidden_states=hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
temb=temb,
|
||||
run_transformer_blocks=run_transformer_blocks,
|
||||
postprocess=postprocess,
|
||||
)
|
||||
|
||||
|
||||
# Registry for model-specific extractors
|
||||
# Key: Pipeline class name (from OmniDiffusionConfig.model_class_name)
|
||||
# Value: extractor function with signature (module, *args, **kwargs) -> CacheContext
|
||||
#
|
||||
# Note: Use the exact pipeline class name as specified in OmniDiffusionConfig.model_class_name
|
||||
# to ensure consistent lookup without substring matching.
|
||||
EXTRACTOR_REGISTRY: dict[str, Callable] = {
|
||||
"QwenImagePipeline": extract_qwen_context,
|
||||
# Future models:
|
||||
# "FluxPipeline": extract_flux_context,
|
||||
# "CogVideoXPipeline": extract_cogvideox_context,
|
||||
}
|
||||
|
||||
|
||||
def register_extractor(model_identifier: str, extractor_fn: Callable) -> None:
|
||||
"""
|
||||
Register a new extractor function for a model type.
|
||||
|
||||
This allows extending TeaCache support to new models without modifying
|
||||
the core TeaCache code.
|
||||
|
||||
Args:
|
||||
model_identifier: Model type identifier (class name or type string)
|
||||
extractor_fn: Function with signature (module, *args, **kwargs) -> CacheContext
|
||||
|
||||
Example:
|
||||
>>> def extract_flux_context(module, hidden_states, timestep, guidance=None, **kwargs):
|
||||
... # Preprocessing
|
||||
... temb = module.time_text_embed(timestep, guidance)
|
||||
... # Extract modulated input
|
||||
... modulated = module.transformer_blocks[0].norm1(hidden_states, emb=temb)
|
||||
... # Define execution
|
||||
... def run_blocks():
|
||||
... h = hidden_states
|
||||
... for block in module.transformer_blocks:
|
||||
... h = block(h, temb=temb)
|
||||
... return (h,)
|
||||
... # Define postprocessing
|
||||
... def postprocess(h):
|
||||
... return module.proj_out(module.norm_out(h, temb))
|
||||
... # Return context
|
||||
... return CacheContext(modulated, hidden_states, None, temb, run_blocks, postprocess)
|
||||
>>> register_extractor("FluxTransformer2DModel", extract_flux_context)
|
||||
"""
|
||||
EXTRACTOR_REGISTRY[model_identifier] = extractor_fn
|
||||
|
||||
|
||||
def get_extractor(model_type: str) -> Callable:
|
||||
"""
|
||||
Get extractor function for given model type.
|
||||
|
||||
This function looks up the extractor based on the exact model_type string,
|
||||
which should match the pipeline class name from OmniDiffusionConfig.model_class_name.
|
||||
|
||||
Args:
|
||||
model_type: Model type string (e.g., "QwenImagePipeline", "FluxPipeline")
|
||||
Must exactly match a key in EXTRACTOR_REGISTRY.
|
||||
|
||||
Returns:
|
||||
Extractor function with signature (module, *args, **kwargs) -> CacheContext
|
||||
|
||||
Raises:
|
||||
ValueError: If model type not found in registry
|
||||
|
||||
Example:
|
||||
>>> # Get extractor for Qwen model
|
||||
>>> extractor = get_extractor("QwenImagePipeline")
|
||||
>>> ctx = extractor(transformer, hidden_states, encoder_hidden_states, timestep, ...)
|
||||
"""
|
||||
# Direct lookup - no substring matching
|
||||
if model_type in EXTRACTOR_REGISTRY:
|
||||
return EXTRACTOR_REGISTRY[model_type]
|
||||
|
||||
# No match found
|
||||
available_types = list(EXTRACTOR_REGISTRY.keys())
|
||||
raise ValueError(
|
||||
f"Unknown model type: '{model_type}'. "
|
||||
f"Available types: {available_types}\n"
|
||||
f"To add support for a new model, use register_extractor() or add to EXTRACTOR_REGISTRY."
|
||||
)
|
||||
+246
@@ -0,0 +1,246 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Hook-based TeaCache implementation for vLLM-Omni.
|
||||
|
||||
This module implements a diffusers-style hook system that completely intercepts
|
||||
the transformer forward pass, eliminating the need for any TeaCache-specific
|
||||
code in model definitions. Model developers only need to add an extractor function
|
||||
to support new models.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
|
||||
from vllm_omni.diffusion.cache.teacache.extractors import get_extractor
|
||||
from vllm_omni.diffusion.cache.teacache.state import TeaCacheState
|
||||
from vllm_omni.diffusion.hooks import HookRegistry, ModelHook, StateManager
|
||||
|
||||
|
||||
class TeaCacheHook(ModelHook):
|
||||
"""
|
||||
ModelHook implementing TeaCache for transformer models.
|
||||
|
||||
This hook completely intercepts the transformer's forward pass and implements
|
||||
adaptive caching based on timestep embedding similarity. It's model-agnostic
|
||||
and supports multiple model types through extractor functions.
|
||||
|
||||
Key features:
|
||||
- Zero changes to model code
|
||||
- CFG-aware with separate states for positive/negative branches
|
||||
- Model-specific polynomial rescaling
|
||||
- Auto-detection of model types
|
||||
|
||||
Attributes:
|
||||
config: TeaCache configuration with thresholds and callbacks
|
||||
rescale_func: Polynomial function for rescaling L1 distances
|
||||
state_manager: Manages TeaCacheState across forward passes
|
||||
extractor_fn: Model-specific function to extract modulated input
|
||||
"""
|
||||
|
||||
_HOOK_NAME = "teacache"
|
||||
|
||||
def __init__(self, config: TeaCacheConfig):
|
||||
"""
|
||||
Initialize TeaCacheHook.
|
||||
|
||||
Args:
|
||||
config: TeaCache configuration object.
|
||||
"""
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.rescale_func = np.poly1d(config.coefficients)
|
||||
self.state_manager = StateManager(TeaCacheState)
|
||||
self.extractor_fn = None
|
||||
|
||||
def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module:
|
||||
"""
|
||||
Initialize hook with extractor from config model type.
|
||||
|
||||
Args:
|
||||
module: The module to initialize the hook for.
|
||||
|
||||
Returns:
|
||||
The initialized module.
|
||||
"""
|
||||
# Get extractor function based on model_type from config
|
||||
# model_type should be the pipeline class name (e.g., "QwenImagePipeline")
|
||||
self.extractor_fn = get_extractor(self.config.model_type)
|
||||
|
||||
# Set default context
|
||||
self.state_manager.set_context("teacache")
|
||||
|
||||
return module
|
||||
|
||||
def new_forward(self, module: torch.nn.Module, *args: Any, **kwargs: Any) -> Any:
|
||||
"""
|
||||
Generic forward handler that works for ANY model.
|
||||
|
||||
This method is completely model-agnostic. All model-specific logic
|
||||
is encapsulated in the extractor function that returns a CacheContext.
|
||||
|
||||
The extractor does:
|
||||
- Model-specific preprocessing
|
||||
- Extraction of modulated input for cache decision
|
||||
- Providing transformer execution callable
|
||||
- Providing postprocessing callable
|
||||
|
||||
This hook does:
|
||||
- CFG-aware state management
|
||||
- Cache decision logic (generic)
|
||||
- Residual caching and reuse
|
||||
|
||||
Args:
|
||||
module: Transformer module (any architecture)
|
||||
*args: Positional arguments for model forward
|
||||
**kwargs: Keyword arguments for model forward
|
||||
|
||||
Returns:
|
||||
Model output (format depends on model)
|
||||
"""
|
||||
# Get model-specific context from extractor
|
||||
# The extractor encapsulates ALL model-specific logic
|
||||
ctx = self.extractor_fn(module, *args, **kwargs)
|
||||
|
||||
# ============================================================================
|
||||
# GENERIC CACHING LOGIC (works for all models)
|
||||
# ============================================================================
|
||||
# Set context based on CFG branch for separate state tracking
|
||||
cache_branch = kwargs.get("cache_branch", "default")
|
||||
context_name = f"teacache_{cache_branch}"
|
||||
self.state_manager.set_context(context_name)
|
||||
state = self.state_manager.get_state()
|
||||
|
||||
# Decide whether to compute or cache based on modulated input similarity
|
||||
should_compute = self._should_compute_full_transformer(state, ctx.modulated_input)
|
||||
|
||||
if not should_compute and state.previous_residual is not None:
|
||||
# ============================================================================
|
||||
# FAST PATH: Reuse cached residuals
|
||||
# ============================================================================
|
||||
ctx.hidden_states = ctx.hidden_states + state.previous_residual
|
||||
if state.previous_residual_encoder is not None and ctx.encoder_hidden_states is not None:
|
||||
ctx.encoder_hidden_states = ctx.encoder_hidden_states + state.previous_residual_encoder
|
||||
output = ctx.hidden_states
|
||||
else:
|
||||
# ============================================================================
|
||||
# SLOW PATH: Full transformer computation
|
||||
# ============================================================================
|
||||
ori_hidden_states = ctx.hidden_states.clone()
|
||||
ori_encoder_hidden_states = (
|
||||
ctx.encoder_hidden_states.clone() if ctx.encoder_hidden_states is not None else None
|
||||
)
|
||||
|
||||
# Run transformer blocks using model-specific callable
|
||||
outputs = ctx.run_transformer_blocks()
|
||||
|
||||
# Update context with outputs
|
||||
ctx.hidden_states = outputs[0]
|
||||
if len(outputs) > 1 and ctx.encoder_hidden_states is not None:
|
||||
ctx.encoder_hidden_states = outputs[1]
|
||||
|
||||
# Cache residuals for next timestep
|
||||
state.previous_residual = (ctx.hidden_states - ori_hidden_states).detach()
|
||||
if ori_encoder_hidden_states is not None:
|
||||
state.previous_residual_encoder = (ctx.encoder_hidden_states - ori_encoder_hidden_states).detach()
|
||||
|
||||
output = ctx.hidden_states
|
||||
|
||||
# Update state
|
||||
state.previous_modulated_input = ctx.modulated_input.detach()
|
||||
state.cnt += 1
|
||||
|
||||
# ============================================================================
|
||||
# POSTPROCESSING (model-specific, via callable)
|
||||
# ============================================================================
|
||||
return ctx.postprocess(output)
|
||||
|
||||
def _should_compute_full_transformer(self, state: TeaCacheState, modulated_inp: torch.Tensor) -> bool:
|
||||
"""
|
||||
Determine whether to compute full transformer or reuse cached residual.
|
||||
|
||||
This implements the core TeaCache algorithm:
|
||||
1. Always compute first timestep
|
||||
2. For intermediate steps:
|
||||
- Compute relative L1 distance between current and previous modulated inputs
|
||||
- Apply polynomial rescaling with model-specific coefficients
|
||||
- Accumulate rescaled distances
|
||||
- Compare to threshold: below = cache, above = compute
|
||||
|
||||
Args:
|
||||
state: Current TeaCacheState containing counters and cached values
|
||||
modulated_inp: Modulated input extracted from first transformer block
|
||||
|
||||
Returns:
|
||||
True to compute full transformer, False to reuse cached residual
|
||||
"""
|
||||
# First timestep: always compute
|
||||
if state.cnt == 0:
|
||||
state.accumulated_rel_l1_distance = 0.0
|
||||
return True
|
||||
|
||||
# Need previous input for comparison
|
||||
if state.previous_modulated_input is None:
|
||||
return True
|
||||
|
||||
# Compute relative L1 distance between consecutive modulated inputs
|
||||
rel_distance = (
|
||||
(
|
||||
(modulated_inp - state.previous_modulated_input).abs().mean()
|
||||
/ (state.previous_modulated_input.abs().mean() + 1e-8)
|
||||
)
|
||||
.cpu()
|
||||
.item()
|
||||
)
|
||||
|
||||
# Apply model-specific polynomial rescaling
|
||||
rescaled_distance = float(self.rescale_func(rel_distance))
|
||||
state.accumulated_rel_l1_distance += abs(rescaled_distance)
|
||||
|
||||
# Decision: below threshold = cache, above = compute
|
||||
if state.accumulated_rel_l1_distance < self.config.rel_l1_thresh:
|
||||
return False # Use cache
|
||||
else:
|
||||
state.accumulated_rel_l1_distance = 0.0 # Reset accumulator
|
||||
return True # Compute
|
||||
|
||||
def reset_state(self, module: torch.nn.Module) -> torch.nn.Module:
|
||||
"""
|
||||
Reset all cached states for a new inference run.
|
||||
|
||||
Args:
|
||||
module: The module to reset state for.
|
||||
|
||||
Returns:
|
||||
The module with reset state.
|
||||
"""
|
||||
self.state_manager.reset()
|
||||
return module
|
||||
|
||||
|
||||
def apply_teacache_hook(module: torch.nn.Module, config: TeaCacheConfig) -> None:
|
||||
"""
|
||||
Apply TeaCache optimization to a transformer module.
|
||||
|
||||
This function registers a TeaCacheHook that completely intercepts the
|
||||
module's forward pass, implementing adaptive caching without any changes
|
||||
to the model code.
|
||||
|
||||
Args:
|
||||
module: Transformer model to optimize (e.g., QwenImageTransformer2DModel)
|
||||
config: TeaCacheConfig specifying caching parameters
|
||||
|
||||
Example:
|
||||
>>> config = TeaCacheConfig(rel_l1_thresh=0.2, model_type="Qwen")
|
||||
>>> apply_teacache_hook(transformer, config)
|
||||
>>> # Model now uses TeaCache automatically, no code changes needed!
|
||||
"""
|
||||
registry = HookRegistry.get_or_create(module)
|
||||
hook = TeaCacheHook(config)
|
||||
registry.register_hook(TeaCacheHook._HOOK_NAME, hook)
|
||||
+40
@@ -0,0 +1,40 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
TeaCache state management.
|
||||
|
||||
This module manages the state for TeaCache hooks across diffusion timesteps.
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class TeaCacheState:
|
||||
"""
|
||||
State management for TeaCache hook.
|
||||
|
||||
Tracks caching state across diffusion timesteps, managing counters,
|
||||
accumulated distances, and cached residuals for the TeaCache algorithm.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize empty TeaCache state."""
|
||||
# Timestep tracking
|
||||
self.cnt = 0
|
||||
|
||||
# Caching state
|
||||
self.accumulated_rel_l1_distance = 0.0
|
||||
self.previous_modulated_input: Optional[torch.Tensor] = None
|
||||
self.previous_residual: Optional[torch.Tensor] = None
|
||||
self.previous_residual_encoder: Optional[torch.Tensor] = None
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset all state variables for a new inference run."""
|
||||
self.cnt = 0
|
||||
self.accumulated_rel_l1_distance = 0.0
|
||||
self.previous_modulated_input = None
|
||||
self.previous_residual = None
|
||||
self.previous_residual_encoder = None
|
||||
@@ -2,6 +2,7 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import enum
|
||||
import os
|
||||
import random
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable
|
||||
@@ -60,9 +61,13 @@ class OmniDiffusionConfig:
|
||||
# Workload type
|
||||
# workload_type: WorkloadType = WorkloadType.T2V
|
||||
|
||||
# Cache strategy
|
||||
# Cache strategy (legacy)
|
||||
cache_strategy: str = "none"
|
||||
|
||||
# Cache adapter configuration (NEW)
|
||||
cache_adapter: str = "none" # "tea_cache", "deep_cache", etc.
|
||||
cache_config: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
# Distributed executor backend
|
||||
distributed_executor_backend: str = "mp"
|
||||
nccl_port: int | None = None
|
||||
@@ -162,6 +167,8 @@ class OmniDiffusionConfig:
|
||||
|
||||
# MoE parameters used by Wan2.2
|
||||
boundary_ratio: float | None = None
|
||||
# Scheduler flow_shift for Wan2.2 (12.0 for 480p, 5.0 for 720p)
|
||||
flow_shift: float | None = None
|
||||
|
||||
# Logging
|
||||
log_level: str = "info"
|
||||
@@ -206,8 +213,17 @@ class OmniDiffusionConfig:
|
||||
initial_master_port = (self.master_port or 30005) + random.randint(0, 100)
|
||||
self.master_port = self.settle_port(initial_master_port, 37)
|
||||
|
||||
# Automatically inject model_class_name into cache_config if not present
|
||||
if self.cache_adapter != "none" and self.model_class_name:
|
||||
if "model_type" not in self.cache_config:
|
||||
self.cache_config["model_type"] = self.model_class_name
|
||||
logger.debug(f"Auto-injected model_type='{self.model_class_name}' into cache_config")
|
||||
|
||||
@classmethod
|
||||
def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig":
|
||||
# Check environment variable as fallback for cache_adapter
|
||||
if "cache_adapter" not in kwargs:
|
||||
kwargs["cache_adapter"] = os.environ.get("DIFFUSION_CACHE_ADAPTER", "none").lower()
|
||||
return cls(**kwargs)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,101 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Callable
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class BaseState:
|
||||
"""Base class for hook state containers."""
|
||||
|
||||
def reset(self) -> None: # pragma: no cover - default is no-op
|
||||
pass
|
||||
|
||||
|
||||
class StateManager:
|
||||
"""Manage per-context hook state instances."""
|
||||
|
||||
def __init__(self, state_cls: Callable[[], BaseState]):
|
||||
self._state_cls = state_cls
|
||||
self._states: dict[str, BaseState] = {}
|
||||
self._context: str = "default"
|
||||
|
||||
def set_context(self, name: str) -> None:
|
||||
self._context = name or "default"
|
||||
|
||||
def get_state(self) -> BaseState:
|
||||
if self._context not in self._states:
|
||||
self._states[self._context] = self._state_cls()
|
||||
return self._states[self._context]
|
||||
|
||||
def reset(self) -> None:
|
||||
self._states.clear()
|
||||
|
||||
|
||||
class ModelHook:
|
||||
"""Base class for model hooks that can override a module's forward."""
|
||||
|
||||
def initialize_hook(self, module: nn.Module) -> nn.Module:
|
||||
return module
|
||||
|
||||
def new_forward(self, module: nn.Module, *args: Any, **kwargs: Any):
|
||||
raise NotImplementedError
|
||||
|
||||
def reset_state(self, module: nn.Module) -> nn.Module:
|
||||
return module
|
||||
|
||||
|
||||
@dataclass
|
||||
class _WrappedForward:
|
||||
module: nn.Module
|
||||
|
||||
def __call__(self, *args: Any, **kwargs: Any):
|
||||
registry: HookRegistry | None = getattr(self.module, "_hook_registry", None)
|
||||
if registry is None or not registry._hooks:
|
||||
return self.module._original_forward(*args, **kwargs)
|
||||
return registry.dispatch(*args, **kwargs)
|
||||
|
||||
|
||||
class HookRegistry:
|
||||
"""Registry of hooks attached to a module."""
|
||||
|
||||
def __init__(self, module: nn.Module):
|
||||
self.module = module
|
||||
self._hooks: dict[str, ModelHook] = {}
|
||||
|
||||
@classmethod
|
||||
def get_or_create(cls, module: nn.Module) -> HookRegistry:
|
||||
registry: HookRegistry | None = getattr(module, "_hook_registry", None)
|
||||
if registry is None:
|
||||
registry = cls(module)
|
||||
setattr(module, "_hook_registry", registry)
|
||||
|
||||
# Wrap module.forward once so hooks can intercept calls.
|
||||
if not hasattr(module, "_original_forward"):
|
||||
module._original_forward = module.forward # type: ignore[attr-defined]
|
||||
module.forward = _WrappedForward(module) # type: ignore[assignment]
|
||||
|
||||
return registry
|
||||
|
||||
def register_hook(self, name: str, hook: ModelHook) -> None:
|
||||
hook.initialize_hook(self.module)
|
||||
self._hooks[name] = hook
|
||||
|
||||
def get_hook(self, name: str) -> ModelHook | None:
|
||||
return self._hooks.get(name)
|
||||
|
||||
def dispatch(self, *args: Any, **kwargs: Any):
|
||||
# For now we support a single active hook and call it directly.
|
||||
# This can be extended to a chain if needed.
|
||||
if not self._hooks:
|
||||
return self.module._original_forward(*args, **kwargs) # type: ignore[attr-defined]
|
||||
# Deterministic order: sort by name.
|
||||
name = sorted(self._hooks.keys())[0]
|
||||
hook = self._hooks[name]
|
||||
return hook.new_forward(self.module, *args, **kwargs)
|
||||
|
||||
def reset_hook(self, name: str) -> None:
|
||||
hook = self._hooks.get(name)
|
||||
if hook is not None:
|
||||
hook.reset_state(self.module)
|
||||
@@ -270,6 +270,9 @@ class QwenImagePipeline(
|
||||
self.transformer = QwenImageTransformer2DModel(od_config=od_config)
|
||||
self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
|
||||
|
||||
# Initialize cache adapter to None (will be set by worker or setup_cache if needed)
|
||||
self._cache_adapter = None
|
||||
|
||||
self.stage = None
|
||||
|
||||
self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
|
||||
@@ -546,31 +549,46 @@ class QwenImagePipeline(
|
||||
if self.interrupt:
|
||||
continue
|
||||
self._current_timestep = t
|
||||
# broadcast to batch dimension and place on same device/dtype as latents
|
||||
|
||||
# Broadcast timestep to match batch size
|
||||
timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype)
|
||||
noise_pred = self.transformer(
|
||||
hidden_states=latents,
|
||||
timestep=timestep / 1000,
|
||||
guidance=guidance,
|
||||
encoder_hidden_states_mask=prompt_embeds_mask,
|
||||
encoder_hidden_states=prompt_embeds,
|
||||
img_shapes=img_shapes,
|
||||
txt_seq_lens=txt_seq_lens,
|
||||
attention_kwargs=self.attention_kwargs,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
|
||||
# Forward pass for positive prompt (or unconditional if no CFG)
|
||||
# cache_branch is passed to hook for CFG-aware state management
|
||||
transformer_kwargs = {
|
||||
"hidden_states": latents,
|
||||
"timestep": timestep / 1000,
|
||||
"guidance": guidance,
|
||||
"encoder_hidden_states_mask": prompt_embeds_mask,
|
||||
"encoder_hidden_states": prompt_embeds,
|
||||
"img_shapes": img_shapes,
|
||||
"txt_seq_lens": txt_seq_lens,
|
||||
"attention_kwargs": self.attention_kwargs,
|
||||
"return_dict": False,
|
||||
}
|
||||
if self._cache_adapter is not None:
|
||||
transformer_kwargs["cache_branch"] = "positive"
|
||||
|
||||
noise_pred = self.transformer(**transformer_kwargs)[0]
|
||||
|
||||
# Forward pass for negative prompt (CFG)
|
||||
# cache_branch is passed to hook for CFG-aware state management
|
||||
if do_true_cfg:
|
||||
neg_noise_pred = self.transformer(
|
||||
hidden_states=latents,
|
||||
timestep=timestep / 1000,
|
||||
guidance=guidance,
|
||||
encoder_hidden_states_mask=negative_prompt_embeds_mask,
|
||||
encoder_hidden_states=negative_prompt_embeds,
|
||||
img_shapes=img_shapes,
|
||||
txt_seq_lens=negative_txt_seq_lens,
|
||||
attention_kwargs=self.attention_kwargs,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
neg_transformer_kwargs = {
|
||||
"hidden_states": latents,
|
||||
"timestep": timestep / 1000,
|
||||
"guidance": guidance,
|
||||
"encoder_hidden_states_mask": negative_prompt_embeds_mask,
|
||||
"encoder_hidden_states": negative_prompt_embeds,
|
||||
"img_shapes": img_shapes,
|
||||
"txt_seq_lens": negative_txt_seq_lens,
|
||||
"attention_kwargs": self.attention_kwargs,
|
||||
"return_dict": False,
|
||||
}
|
||||
if self._cache_adapter is not None:
|
||||
neg_transformer_kwargs["cache_branch"] = "negative"
|
||||
|
||||
neg_noise_pred = self.transformer(**neg_transformer_kwargs)[0]
|
||||
comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
|
||||
cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
|
||||
noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
|
||||
|
||||
@@ -224,7 +224,7 @@ class QwenImageEditPipeline(
|
||||
self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
|
||||
self.device
|
||||
)
|
||||
self.transformer = QwenImageTransformer2DModel()
|
||||
self.transformer = QwenImageTransformer2DModel(od_config=od_config)
|
||||
self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
|
||||
self.processor = Qwen2VLProcessor.from_pretrained(
|
||||
model, subfolder="processor", local_files_only=local_files_only
|
||||
|
||||
@@ -0,0 +1,4 @@
|
||||
from .pipeline_wan2_2 import Wan22Pipeline, get_wan22_post_process_func
|
||||
from .wan2_2_transformer import WanTransformer3DModel
|
||||
|
||||
__all__ = ["Wan22Pipeline", "get_wan22_post_process_func", "WanTransformer3DModel"]
|
||||
@@ -0,0 +1,426 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from collections.abc import Iterable
|
||||
|
||||
import torch
|
||||
from diffusers import AutoencoderKLWan, FlowMatchEulerDiscreteScheduler
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
from torch import nn
|
||||
from transformers import AutoTokenizer, UMT5EncoderModel
|
||||
from vllm.model_executor.models.utils import AutoWeightsLoader
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.utils import get_local_device
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
|
||||
|
||||
def get_wan22_post_process_func(
|
||||
od_config: OmniDiffusionConfig,
|
||||
):
|
||||
from diffusers.video_processor import VideoProcessor
|
||||
|
||||
video_processor = VideoProcessor(vae_scale_factor=8)
|
||||
|
||||
def post_process_func(
|
||||
video: torch.Tensor,
|
||||
output_type: str = "np",
|
||||
):
|
||||
if output_type == "latent":
|
||||
return video
|
||||
return video_processor.postprocess_video(video, output_type=output_type)
|
||||
|
||||
return post_process_func
|
||||
|
||||
|
||||
class Wan22Pipeline(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
od_config: OmniDiffusionConfig,
|
||||
prefix: str = "",
|
||||
):
|
||||
super().__init__()
|
||||
self.od_config = od_config
|
||||
|
||||
self.device = get_local_device()
|
||||
dtype = getattr(od_config, "dtype", torch.bfloat16)
|
||||
|
||||
model = od_config.model
|
||||
local_files_only = os.path.exists(model)
|
||||
|
||||
# Set up weights sources for transformer(s)
|
||||
self.weights_sources = [
|
||||
DiffusersPipelineLoader.ComponentSource(
|
||||
model_or_path=od_config.model,
|
||||
subfolder="transformer",
|
||||
revision=None,
|
||||
prefix="transformer.",
|
||||
fall_back_to_pt=True,
|
||||
),
|
||||
DiffusersPipelineLoader.ComponentSource(
|
||||
model_or_path=od_config.model,
|
||||
subfolder="transformer_2",
|
||||
revision=None,
|
||||
prefix="transformer_2.",
|
||||
fall_back_to_pt=True,
|
||||
),
|
||||
]
|
||||
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
|
||||
self.text_encoder = UMT5EncoderModel.from_pretrained(
|
||||
model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only
|
||||
).to(self.device)
|
||||
self.vae = AutoencoderKLWan.from_pretrained(
|
||||
model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
|
||||
).to(self.device)
|
||||
|
||||
# Initialize transformers without from_pretrained (weights loaded via load_weights)
|
||||
self.transformer = WanTransformer3DModel()
|
||||
self.transformer_2 = WanTransformer3DModel()
|
||||
|
||||
self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
|
||||
model, subfolder="scheduler", local_files_only=local_files_only
|
||||
)
|
||||
# Apply flow_shift if specified (12.0 for 480p, 5.0 for 720p recommended for Wan2.2)
|
||||
if od_config.flow_shift is not None:
|
||||
self.scheduler.config.flow_shift = od_config.flow_shift
|
||||
|
||||
self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal if getattr(self, "vae", None) else 4
|
||||
self.vae_scale_factor_spatial = self.vae.config.scale_factor_spatial if getattr(self, "vae", None) else 8
|
||||
self.boundary_ratio = od_config.boundary_ratio
|
||||
|
||||
self._guidance_scale = None
|
||||
self._guidance_scale_2 = None
|
||||
self._num_timesteps = None
|
||||
self._current_timestep = None
|
||||
|
||||
@property
|
||||
def guidance_scale(self):
|
||||
return self._guidance_scale
|
||||
|
||||
@property
|
||||
def do_classifier_free_guidance(self):
|
||||
return self._guidance_scale is not None and self._guidance_scale > 1.0
|
||||
|
||||
@property
|
||||
def num_timesteps(self):
|
||||
return self._num_timesteps
|
||||
|
||||
@property
|
||||
def current_timestep(self):
|
||||
return self._current_timestep
|
||||
|
||||
@torch.no_grad()
|
||||
def forward(
|
||||
self,
|
||||
req: OmniDiffusionRequest,
|
||||
prompt: str | None = None,
|
||||
negative_prompt: str | None = None,
|
||||
height: int | None = None,
|
||||
width: int | None = None,
|
||||
num_inference_steps: int | None = None,
|
||||
guidance_scale: float | tuple[float, float] = 4.0,
|
||||
frame_num: int | None = None,
|
||||
output_type: str | None = "np",
|
||||
generator: torch.Generator | None = None,
|
||||
prompt_embeds: torch.Tensor | None = None,
|
||||
negative_prompt_embeds: torch.Tensor | None = None,
|
||||
attention_kwargs: dict | None = None,
|
||||
**kwargs,
|
||||
) -> DiffusionOutput:
|
||||
prompt = req.prompt if req.prompt is not None else prompt
|
||||
negative_prompt = req.negative_prompt if req.negative_prompt is not None else negative_prompt
|
||||
if prompt is None and prompt_embeds is None:
|
||||
raise ValueError("Prompt or prompt_embeds is required for Wan2.2 generation.")
|
||||
|
||||
height = req.height or height or 720
|
||||
width = req.width or width or 1280
|
||||
num_frames = req.num_frames if req.num_frames else frame_num or 81
|
||||
num_steps = req.num_inference_steps or num_inference_steps or 40
|
||||
|
||||
guidance_low = guidance_scale if isinstance(guidance_scale, (int, float)) else guidance_scale[0]
|
||||
guidance_high = (
|
||||
req.guidance_scale_2
|
||||
if req.guidance_scale_2 is not None
|
||||
else (
|
||||
guidance_scale[1]
|
||||
if isinstance(guidance_scale, (list, tuple)) and len(guidance_scale) > 1
|
||||
else guidance_low
|
||||
)
|
||||
)
|
||||
|
||||
# record guidance for properties
|
||||
self._guidance_scale = guidance_low
|
||||
self._guidance_scale_2 = guidance_high
|
||||
|
||||
# validate shapes
|
||||
self.check_inputs(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
height=height,
|
||||
width=width,
|
||||
prompt_embeds=prompt_embeds,
|
||||
negative_prompt_embeds=negative_prompt_embeds,
|
||||
guidance_scale_2=guidance_high if self.boundary_ratio is not None else None,
|
||||
)
|
||||
|
||||
if num_frames % self.vae_scale_factor_temporal != 1:
|
||||
num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1
|
||||
num_frames = max(num_frames, 1)
|
||||
|
||||
device = self.device
|
||||
dtype = self.transformer.dtype
|
||||
|
||||
# Seed / generator
|
||||
if generator is None:
|
||||
generator = req.generator
|
||||
if generator is None and req.seed is not None:
|
||||
generator = torch.Generator(device=device).manual_seed(req.seed)
|
||||
|
||||
# Encode prompts
|
||||
if prompt_embeds is None:
|
||||
prompt_embeds, negative_prompt_embeds = self.encode_prompt(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0,
|
||||
num_videos_per_prompt=req.num_outputs_per_prompt or 1,
|
||||
max_sequence_length=req.max_sequence_length or 512,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
else:
|
||||
prompt_embeds = prompt_embeds.to(device=device, dtype=dtype)
|
||||
if negative_prompt_embeds is not None:
|
||||
negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype)
|
||||
elif guidance_low > 1.0 or guidance_high > 1.0:
|
||||
raise ValueError(
|
||||
"negative_prompt_embeds must be provided when prompt_embeds are given and guidance > 1."
|
||||
)
|
||||
|
||||
# Timesteps
|
||||
self.scheduler.set_timesteps(num_steps, device=device)
|
||||
timesteps = self.scheduler.timesteps
|
||||
self._num_timesteps = len(timesteps)
|
||||
boundary_timestep = None
|
||||
if self.boundary_ratio is not None:
|
||||
boundary_timestep = self.boundary_ratio * self.scheduler.config.num_train_timesteps
|
||||
|
||||
# Latents
|
||||
num_channels_latents = self.transformer.config.in_channels
|
||||
latents = self.prepare_latents(
|
||||
batch_size=prompt_embeds.shape[0],
|
||||
num_channels_latents=num_channels_latents,
|
||||
height=height,
|
||||
width=width,
|
||||
num_frames=num_frames,
|
||||
dtype=torch.float32,
|
||||
device=device,
|
||||
generator=generator,
|
||||
latents=req.latents,
|
||||
)
|
||||
|
||||
if attention_kwargs is None:
|
||||
attention_kwargs = {}
|
||||
|
||||
# Denoising
|
||||
for t in timesteps:
|
||||
self._current_timestep = t
|
||||
current_model = self.transformer
|
||||
current_guidance_scale = guidance_low
|
||||
if boundary_timestep is not None and t < boundary_timestep:
|
||||
current_model = self.transformer_2
|
||||
current_guidance_scale = guidance_high
|
||||
|
||||
latent_model_input = latents.to(dtype)
|
||||
timestep = t.expand(latents.shape[0])
|
||||
|
||||
noise_pred = current_model(
|
||||
hidden_states=latent_model_input,
|
||||
timestep=timestep,
|
||||
encoder_hidden_states=prompt_embeds,
|
||||
attention_kwargs=attention_kwargs,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
|
||||
if current_guidance_scale > 1.0 and negative_prompt_embeds is not None:
|
||||
noise_uncond = current_model(
|
||||
hidden_states=latent_model_input,
|
||||
timestep=timestep,
|
||||
encoder_hidden_states=negative_prompt_embeds,
|
||||
attention_kwargs=attention_kwargs,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
noise_pred = noise_uncond + current_guidance_scale * (noise_pred - noise_uncond)
|
||||
|
||||
latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
|
||||
|
||||
self._current_timestep = None
|
||||
|
||||
# Decode
|
||||
if output_type == "latent":
|
||||
output = latents
|
||||
else:
|
||||
latents = latents.to(self.vae.dtype)
|
||||
latents_mean = (
|
||||
torch.tensor(self.vae.config.latents_mean)
|
||||
.view(1, self.vae.config.z_dim, 1, 1, 1)
|
||||
.to(latents.device, latents.dtype)
|
||||
)
|
||||
latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
|
||||
latents.device, latents.dtype
|
||||
)
|
||||
latents = latents / latents_std + latents_mean
|
||||
output = self.vae.decode(latents, return_dict=False)[0]
|
||||
|
||||
return DiffusionOutput(output=output)
|
||||
|
||||
def encode_prompt(
|
||||
self,
|
||||
prompt: str | list[str],
|
||||
negative_prompt: str | list[str] | None = None,
|
||||
do_classifier_free_guidance: bool = True,
|
||||
num_videos_per_prompt: int = 1,
|
||||
max_sequence_length: int = 512,
|
||||
device: torch.device | None = None,
|
||||
dtype: torch.dtype | None = None,
|
||||
):
|
||||
device = device or self.device
|
||||
dtype = dtype or self.text_encoder.dtype
|
||||
|
||||
prompt = [prompt] if isinstance(prompt, str) else prompt
|
||||
prompt_clean = [self._prompt_clean(p) for p in prompt]
|
||||
batch_size = len(prompt_clean)
|
||||
|
||||
text_inputs = self.tokenizer(
|
||||
prompt_clean,
|
||||
padding="max_length",
|
||||
max_length=max_sequence_length,
|
||||
truncation=True,
|
||||
add_special_tokens=True,
|
||||
return_attention_mask=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
ids, mask = text_inputs.input_ids, text_inputs.attention_mask
|
||||
seq_lens = mask.gt(0).sum(dim=1).long()
|
||||
|
||||
prompt_embeds = self.text_encoder(ids.to(device), mask.to(device)).last_hidden_state
|
||||
prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
|
||||
prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)]
|
||||
prompt_embeds = torch.stack(
|
||||
[torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0
|
||||
)
|
||||
|
||||
_, seq_len, _ = prompt_embeds.shape
|
||||
prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
|
||||
prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
|
||||
|
||||
negative_prompt_embeds = None
|
||||
if do_classifier_free_guidance:
|
||||
negative_prompt = negative_prompt or ""
|
||||
negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
|
||||
neg_text_inputs = self.tokenizer(
|
||||
[self._prompt_clean(p) for p in negative_prompt],
|
||||
padding="max_length",
|
||||
max_length=max_sequence_length,
|
||||
truncation=True,
|
||||
add_special_tokens=True,
|
||||
return_attention_mask=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
ids_neg, mask_neg = neg_text_inputs.input_ids, neg_text_inputs.attention_mask
|
||||
seq_lens_neg = mask_neg.gt(0).sum(dim=1).long()
|
||||
negative_prompt_embeds = self.text_encoder(ids_neg.to(device), mask_neg.to(device)).last_hidden_state
|
||||
negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device)
|
||||
negative_prompt_embeds = [u[:v] for u, v in zip(negative_prompt_embeds, seq_lens_neg)]
|
||||
negative_prompt_embeds = torch.stack(
|
||||
[
|
||||
torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))])
|
||||
for u in negative_prompt_embeds
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_videos_per_prompt, 1)
|
||||
negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
|
||||
|
||||
return prompt_embeds, negative_prompt_embeds
|
||||
|
||||
@staticmethod
|
||||
def _prompt_clean(text: str) -> str:
|
||||
return " ".join(text.strip().split())
|
||||
|
||||
def prepare_latents(
|
||||
self,
|
||||
batch_size: int,
|
||||
num_channels_latents: int,
|
||||
height: int,
|
||||
width: int,
|
||||
num_frames: int,
|
||||
dtype: torch.dtype | None,
|
||||
device: torch.device | None,
|
||||
generator: torch.Generator | list[torch.Generator] | None,
|
||||
latents: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
if latents is not None:
|
||||
return latents.to(device=device, dtype=dtype)
|
||||
|
||||
num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
|
||||
shape = (
|
||||
batch_size,
|
||||
num_channels_latents,
|
||||
num_latent_frames,
|
||||
int(height) // self.vae_scale_factor_spatial,
|
||||
int(width) // self.vae_scale_factor_spatial,
|
||||
)
|
||||
if isinstance(generator, list) and len(generator) != batch_size:
|
||||
raise ValueError(f"Generator list length {len(generator)} does not match batch size {batch_size}.")
|
||||
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
|
||||
return latents
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
"""Load weights using AutoWeightsLoader for vLLM integration."""
|
||||
loader = AutoWeightsLoader(self)
|
||||
return loader.load_weights(weights)
|
||||
|
||||
def check_inputs(
|
||||
self,
|
||||
prompt,
|
||||
negative_prompt,
|
||||
height,
|
||||
width,
|
||||
prompt_embeds=None,
|
||||
negative_prompt_embeds=None,
|
||||
guidance_scale_2=None,
|
||||
):
|
||||
if height % 16 != 0 or width % 16 != 0:
|
||||
raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.")
|
||||
|
||||
if prompt is not None and prompt_embeds is not None:
|
||||
raise ValueError(
|
||||
f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
|
||||
" only forward one of the two."
|
||||
)
|
||||
elif negative_prompt is not None and negative_prompt_embeds is not None:
|
||||
raise ValueError(
|
||||
f"Cannot forward both `negative_prompt`: {negative_prompt} and "
|
||||
f"`negative_prompt_embeds`: {negative_prompt_embeds}. "
|
||||
"Please make sure to only forward one of the two."
|
||||
)
|
||||
elif prompt is None and prompt_embeds is None:
|
||||
raise ValueError(
|
||||
"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
|
||||
)
|
||||
elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
|
||||
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
|
||||
elif negative_prompt is not None and (
|
||||
not isinstance(negative_prompt, str) and not isinstance(negative_prompt, list)
|
||||
):
|
||||
raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}")
|
||||
|
||||
if self.boundary_ratio is None and guidance_scale_2 is not None:
|
||||
raise ValueError("`guidance_scale_2` is only supported when `boundary_ratio` is set.")
|
||||
@@ -0,0 +1,704 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import math
|
||||
from collections.abc import Iterable
|
||||
from typing import Any, Optional, Union
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from diffusers.models.attention import FeedForward
|
||||
from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps
|
||||
from diffusers.models.modeling_outputs import Transformer2DModelOutput
|
||||
from diffusers.models.normalization import FP32LayerNorm
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||||
from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinear
|
||||
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||||
|
||||
from vllm_omni.diffusion.attention.layer import Attention
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def apply_rotary_emb_wan(
|
||||
hidden_states: torch.Tensor,
|
||||
freqs_cos: torch.Tensor,
|
||||
freqs_sin: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Apply rotary embeddings to input tensors using the given frequency tensors.
|
||||
|
||||
Args:
|
||||
hidden_states: Input tensor of shape [B, S, H, D]
|
||||
freqs_cos: Cosine frequencies
|
||||
freqs_sin: Sine frequencies
|
||||
|
||||
Returns:
|
||||
Tensor with rotary embeddings applied
|
||||
"""
|
||||
x1, x2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1)
|
||||
cos = freqs_cos[..., 0::2]
|
||||
sin = freqs_sin[..., 1::2]
|
||||
out = torch.empty_like(hidden_states)
|
||||
out[..., 0::2] = x1 * cos - x2 * sin
|
||||
out[..., 1::2] = x1 * sin + x2 * cos
|
||||
return out.type_as(hidden_states)
|
||||
|
||||
|
||||
class WanRotaryPosEmbed(nn.Module):
|
||||
"""
|
||||
Rotary position embeddings for 3D video data (temporal + spatial dimensions).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
attention_head_dim: int,
|
||||
patch_size: tuple[int, int, int],
|
||||
max_seq_len: int,
|
||||
theta: float = 10000.0,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.attention_head_dim = attention_head_dim
|
||||
self.patch_size = patch_size
|
||||
self.max_seq_len = max_seq_len
|
||||
|
||||
# Split dimensions for temporal, height, width
|
||||
h_dim = w_dim = 2 * (attention_head_dim // 6)
|
||||
t_dim = attention_head_dim - h_dim - w_dim
|
||||
freqs_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64
|
||||
|
||||
freqs_cos = []
|
||||
freqs_sin = []
|
||||
|
||||
for dim in [t_dim, h_dim, w_dim]:
|
||||
freq_cos, freq_sin = self._get_1d_rotary_pos_embed(dim, max_seq_len, theta, freqs_dtype)
|
||||
freqs_cos.append(freq_cos)
|
||||
freqs_sin.append(freq_sin)
|
||||
|
||||
self.register_buffer("freqs_cos", torch.cat(freqs_cos, dim=1), persistent=False)
|
||||
self.register_buffer("freqs_sin", torch.cat(freqs_sin, dim=1), persistent=False)
|
||||
|
||||
@staticmethod
|
||||
def _get_1d_rotary_pos_embed(
|
||||
dim: int,
|
||||
max_seq_len: int,
|
||||
theta: float,
|
||||
freqs_dtype: torch.dtype,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Generate 1D rotary position embeddings."""
|
||||
freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=freqs_dtype) / dim))
|
||||
t = torch.arange(max_seq_len, dtype=freqs_dtype)
|
||||
freqs = torch.outer(t, freqs)
|
||||
# Repeat interleave for real representation
|
||||
freqs_cos = freqs.cos().repeat_interleave(2, dim=-1)
|
||||
freqs_sin = freqs.sin().repeat_interleave(2, dim=-1)
|
||||
return freqs_cos.float(), freqs_sin.float()
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
batch_size, num_channels, num_frames, height, width = hidden_states.shape
|
||||
p_t, p_h, p_w = self.patch_size
|
||||
ppf, pph, ppw = num_frames // p_t, height // p_h, width // p_w
|
||||
|
||||
split_sizes = [
|
||||
self.attention_head_dim - 2 * (self.attention_head_dim // 3),
|
||||
self.attention_head_dim // 3,
|
||||
self.attention_head_dim // 3,
|
||||
]
|
||||
|
||||
freqs_cos = self.freqs_cos.split(split_sizes, dim=1)
|
||||
freqs_sin = self.freqs_sin.split(split_sizes, dim=1)
|
||||
|
||||
freqs_cos_f = freqs_cos[0][:ppf].view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1)
|
||||
freqs_cos_h = freqs_cos[1][:pph].view(1, pph, 1, -1).expand(ppf, pph, ppw, -1)
|
||||
freqs_cos_w = freqs_cos[2][:ppw].view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1)
|
||||
|
||||
freqs_sin_f = freqs_sin[0][:ppf].view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1)
|
||||
freqs_sin_h = freqs_sin[1][:pph].view(1, pph, 1, -1).expand(ppf, pph, ppw, -1)
|
||||
freqs_sin_w = freqs_sin[2][:ppw].view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1)
|
||||
|
||||
freqs_cos = torch.cat([freqs_cos_f, freqs_cos_h, freqs_cos_w], dim=-1).reshape(1, ppf * pph * ppw, 1, -1)
|
||||
freqs_sin = torch.cat([freqs_sin_f, freqs_sin_h, freqs_sin_w], dim=-1).reshape(1, ppf * pph * ppw, 1, -1)
|
||||
|
||||
return freqs_cos, freqs_sin
|
||||
|
||||
|
||||
class WanImageEmbedding(nn.Module):
|
||||
"""Image embedding module for I2V tasks."""
|
||||
|
||||
def __init__(self, in_features: int, out_features: int, pos_embed_seq_len: Optional[int] = None):
|
||||
super().__init__()
|
||||
|
||||
self.norm1 = FP32LayerNorm(in_features)
|
||||
self.ff = FeedForward(in_features, out_features, mult=1, activation_fn="gelu")
|
||||
self.norm2 = FP32LayerNorm(out_features)
|
||||
if pos_embed_seq_len is not None:
|
||||
self.pos_embed = nn.Parameter(torch.zeros(1, pos_embed_seq_len, in_features))
|
||||
else:
|
||||
self.pos_embed = None
|
||||
|
||||
def forward(self, encoder_hidden_states_image: torch.Tensor) -> torch.Tensor:
|
||||
if self.pos_embed is not None:
|
||||
batch_size, seq_len, embed_dim = encoder_hidden_states_image.shape
|
||||
encoder_hidden_states_image = encoder_hidden_states_image.view(-1, 2 * seq_len, embed_dim)
|
||||
encoder_hidden_states_image = encoder_hidden_states_image + self.pos_embed
|
||||
|
||||
hidden_states = self.norm1(encoder_hidden_states_image)
|
||||
hidden_states = self.ff(hidden_states)
|
||||
hidden_states = self.norm2(hidden_states)
|
||||
return hidden_states
|
||||
|
||||
|
||||
class WanTimeTextImageEmbedding(nn.Module):
|
||||
"""Combined time, text, and image condition embeddings."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
time_freq_dim: int,
|
||||
time_proj_dim: int,
|
||||
text_embed_dim: int,
|
||||
image_embed_dim: Optional[int] = None,
|
||||
pos_embed_seq_len: Optional[int] = None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.timesteps_proj = Timesteps(num_channels=time_freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0)
|
||||
self.time_embedder = TimestepEmbedding(in_channels=time_freq_dim, time_embed_dim=dim)
|
||||
self.act_fn = nn.SiLU()
|
||||
self.time_proj = nn.Linear(dim, time_proj_dim)
|
||||
self.text_embedder = PixArtAlphaTextProjection(text_embed_dim, dim, act_fn="gelu_tanh")
|
||||
|
||||
self.image_embedder = None
|
||||
if image_embed_dim is not None:
|
||||
self.image_embedder = WanImageEmbedding(image_embed_dim, dim, pos_embed_seq_len=pos_embed_seq_len)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
timestep: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
encoder_hidden_states_image: Optional[torch.Tensor] = None,
|
||||
timestep_seq_len: Optional[int] = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
|
||||
timestep = self.timesteps_proj(timestep)
|
||||
if timestep_seq_len is not None:
|
||||
timestep = timestep.unflatten(0, (-1, timestep_seq_len))
|
||||
|
||||
time_embedder_dtype = next(iter(self.time_embedder.parameters())).dtype
|
||||
if timestep.dtype != time_embedder_dtype and time_embedder_dtype != torch.int8:
|
||||
timestep = timestep.to(time_embedder_dtype)
|
||||
temb = self.time_embedder(timestep).type_as(encoder_hidden_states)
|
||||
timestep_proj = self.time_proj(self.act_fn(temb))
|
||||
|
||||
encoder_hidden_states = self.text_embedder(encoder_hidden_states)
|
||||
if encoder_hidden_states_image is not None:
|
||||
encoder_hidden_states_image = self.image_embedder(encoder_hidden_states_image)
|
||||
|
||||
return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image
|
||||
|
||||
|
||||
class WanSelfAttention(nn.Module):
|
||||
"""
|
||||
Optimized self-attention module using vLLM layers.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
eps: float = 1e-5,
|
||||
dropout: float = 0.0,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.dim = dim
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = head_dim
|
||||
self.inner_dim = num_heads * head_dim
|
||||
|
||||
# Fused QKV projection using vLLM's optimized layer
|
||||
self.to_qkv = QKVParallelLinear(
|
||||
hidden_size=dim,
|
||||
head_size=head_dim,
|
||||
total_num_heads=num_heads,
|
||||
bias=True,
|
||||
disable_tp=True,
|
||||
)
|
||||
|
||||
# QK normalization using vLLM's RMSNorm
|
||||
self.norm_q = RMSNorm(self.inner_dim, eps=eps)
|
||||
self.norm_k = RMSNorm(self.inner_dim, eps=eps)
|
||||
|
||||
# Output projection
|
||||
self.to_out = nn.ModuleList(
|
||||
[
|
||||
ReplicatedLinear(self.inner_dim, dim, bias=True),
|
||||
nn.Dropout(dropout),
|
||||
]
|
||||
)
|
||||
|
||||
# Unified attention layer
|
||||
self.attn = Attention(
|
||||
num_heads=num_heads,
|
||||
head_size=head_dim,
|
||||
softmax_scale=1.0 / (head_dim**0.5),
|
||||
causal=False,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
rotary_emb: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
|
||||
) -> torch.Tensor:
|
||||
# Fused QKV projection
|
||||
qkv, _ = self.to_qkv(hidden_states)
|
||||
query, key, value = qkv.chunk(3, dim=-1)
|
||||
|
||||
# Apply QK normalization
|
||||
query = self.norm_q(query)
|
||||
key = self.norm_k(key)
|
||||
|
||||
# Reshape for multi-head attention
|
||||
query = query.unflatten(2, (self.num_heads, -1))
|
||||
key = key.unflatten(2, (self.num_heads, -1))
|
||||
value = value.unflatten(2, (self.num_heads, -1))
|
||||
|
||||
# Apply rotary embeddings
|
||||
if rotary_emb is not None:
|
||||
freqs_cos, freqs_sin = rotary_emb
|
||||
query = apply_rotary_emb_wan(query, freqs_cos, freqs_sin)
|
||||
key = apply_rotary_emb_wan(key, freqs_cos, freqs_sin)
|
||||
|
||||
# Compute attention using unified attention layer
|
||||
hidden_states = self.attn(query, key, value)
|
||||
hidden_states = hidden_states.flatten(2, 3)
|
||||
hidden_states = hidden_states.type_as(query)
|
||||
|
||||
# Output projection
|
||||
hidden_states, _ = self.to_out[0](hidden_states)
|
||||
hidden_states = self.to_out[1](hidden_states)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class WanCrossAttention(nn.Module):
|
||||
"""
|
||||
Optimized cross-attention module using vLLM layers.
|
||||
Handles both text cross-attention and optional image cross-attention (I2V).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
eps: float = 1e-5,
|
||||
dropout: float = 0.0,
|
||||
added_kv_proj_dim: Optional[int] = None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.dim = dim
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = head_dim
|
||||
self.inner_dim = num_heads * head_dim
|
||||
self.kv_inner_dim = head_dim * num_heads # For cross-attention, K/V come from encoder
|
||||
|
||||
# Query projection
|
||||
self.to_q = ReplicatedLinear(dim, self.inner_dim, bias=True)
|
||||
|
||||
# Separate K and V projections for cross-attention
|
||||
self.to_k = ReplicatedLinear(dim, self.kv_inner_dim, bias=True)
|
||||
self.to_v = ReplicatedLinear(dim, self.kv_inner_dim, bias=True)
|
||||
|
||||
# QK normalization
|
||||
self.norm_q = RMSNorm(self.inner_dim, eps=eps)
|
||||
self.norm_k = RMSNorm(self.kv_inner_dim, eps=eps)
|
||||
|
||||
# Optional added KV projections for I2V (image embeddings)
|
||||
self.added_kv_proj_dim = added_kv_proj_dim
|
||||
if added_kv_proj_dim is not None:
|
||||
self.add_k_proj = ReplicatedLinear(added_kv_proj_dim, self.inner_dim, bias=True)
|
||||
self.add_v_proj = ReplicatedLinear(added_kv_proj_dim, self.inner_dim, bias=True)
|
||||
self.norm_added_k = RMSNorm(self.inner_dim, eps=eps)
|
||||
else:
|
||||
self.add_k_proj = None
|
||||
self.add_v_proj = None
|
||||
self.norm_added_k = None
|
||||
|
||||
# Output projection
|
||||
self.to_out = nn.ModuleList(
|
||||
[
|
||||
ReplicatedLinear(self.inner_dim, dim, bias=True),
|
||||
nn.Dropout(dropout),
|
||||
]
|
||||
)
|
||||
|
||||
# Unified attention layer
|
||||
self.attn = Attention(
|
||||
num_heads=num_heads,
|
||||
head_size=head_dim,
|
||||
softmax_scale=1.0 / (head_dim**0.5),
|
||||
causal=False,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
# Handle I2V case where encoder_hidden_states contains both image and text
|
||||
encoder_hidden_states_img = None
|
||||
if self.add_k_proj is not None:
|
||||
# 512 is the context length of the text encoder, hardcoded for now
|
||||
image_context_length = encoder_hidden_states.shape[1] - 512
|
||||
encoder_hidden_states_img = encoder_hidden_states[:, :image_context_length]
|
||||
encoder_hidden_states = encoder_hidden_states[:, image_context_length:]
|
||||
|
||||
# Query projection
|
||||
query, _ = self.to_q(hidden_states)
|
||||
query = self.norm_q(query)
|
||||
|
||||
# KV projection from encoder
|
||||
key, _ = self.to_k(encoder_hidden_states)
|
||||
value, _ = self.to_v(encoder_hidden_states)
|
||||
key = self.norm_k(key)
|
||||
|
||||
# Reshape for multi-head attention
|
||||
query = query.unflatten(2, (self.num_heads, -1))
|
||||
key = key.unflatten(2, (self.num_heads, -1))
|
||||
value = value.unflatten(2, (self.num_heads, -1))
|
||||
|
||||
# I2V: Additional attention with image embeddings
|
||||
hidden_states_img = None
|
||||
if encoder_hidden_states_img is not None:
|
||||
key_img, _ = self.add_k_proj(encoder_hidden_states_img)
|
||||
value_img, _ = self.add_v_proj(encoder_hidden_states_img)
|
||||
key_img = self.norm_added_k(key_img)
|
||||
|
||||
key_img = key_img.unflatten(2, (self.num_heads, -1))
|
||||
value_img = value_img.unflatten(2, (self.num_heads, -1))
|
||||
|
||||
hidden_states_img = self.attn(query, key_img, value_img)
|
||||
hidden_states_img = hidden_states_img.flatten(2, 3)
|
||||
hidden_states_img = hidden_states_img.type_as(query)
|
||||
|
||||
# Main cross-attention using unified attention layer
|
||||
hidden_states = self.attn(query, key, value)
|
||||
hidden_states = hidden_states.flatten(2, 3)
|
||||
hidden_states = hidden_states.type_as(query)
|
||||
|
||||
# Add image attention output if present
|
||||
if hidden_states_img is not None:
|
||||
hidden_states = hidden_states + hidden_states_img
|
||||
|
||||
# Output projection
|
||||
hidden_states, _ = self.to_out[0](hidden_states)
|
||||
hidden_states = self.to_out[1](hidden_states)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class WanTransformerBlock(nn.Module):
|
||||
"""
|
||||
Transformer block for Wan model with self-attention, cross-attention, and FFN.
|
||||
Uses scale-shift modulation from timestep embeddings.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
ffn_dim: int,
|
||||
num_heads: int,
|
||||
eps: float = 1e-6,
|
||||
added_kv_proj_dim: Optional[int] = None,
|
||||
cross_attn_norm: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
head_dim = dim // num_heads
|
||||
|
||||
# 1. Self-attention
|
||||
self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
|
||||
self.attn1 = WanSelfAttention(
|
||||
dim=dim,
|
||||
num_heads=num_heads,
|
||||
head_dim=head_dim,
|
||||
eps=eps,
|
||||
)
|
||||
|
||||
# 2. Cross-attention
|
||||
self.attn2 = WanCrossAttention(
|
||||
dim=dim,
|
||||
num_heads=num_heads,
|
||||
head_dim=head_dim,
|
||||
eps=eps,
|
||||
added_kv_proj_dim=added_kv_proj_dim,
|
||||
)
|
||||
self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
|
||||
|
||||
# 3. Feed-forward
|
||||
self.ffn = FeedForward(dim, inner_dim=ffn_dim, activation_fn="gelu-approximate")
|
||||
self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False)
|
||||
|
||||
# Scale-shift table for modulation
|
||||
self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
temb: torch.Tensor,
|
||||
rotary_emb: tuple[torch.Tensor, torch.Tensor],
|
||||
) -> torch.Tensor:
|
||||
if temb.ndim == 4:
|
||||
# temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v)
|
||||
shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
|
||||
self.scale_shift_table.unsqueeze(0) + temb.float()
|
||||
).chunk(6, dim=2)
|
||||
shift_msa = shift_msa.squeeze(2)
|
||||
scale_msa = scale_msa.squeeze(2)
|
||||
gate_msa = gate_msa.squeeze(2)
|
||||
c_shift_msa = c_shift_msa.squeeze(2)
|
||||
c_scale_msa = c_scale_msa.squeeze(2)
|
||||
c_gate_msa = c_gate_msa.squeeze(2)
|
||||
else:
|
||||
# temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B)
|
||||
shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
|
||||
self.scale_shift_table + temb.float()
|
||||
).chunk(6, dim=1)
|
||||
|
||||
# 1. Self-attention
|
||||
norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).type_as(hidden_states)
|
||||
attn_output = self.attn1(norm_hidden_states, rotary_emb)
|
||||
hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as(hidden_states)
|
||||
|
||||
# 2. Cross-attention
|
||||
norm_hidden_states = self.norm2(hidden_states.float()).type_as(hidden_states)
|
||||
attn_output = self.attn2(norm_hidden_states, encoder_hidden_states)
|
||||
hidden_states = hidden_states + attn_output
|
||||
|
||||
# 3. Feed-forward
|
||||
norm_hidden_states = (self.norm3(hidden_states.float()) * (1 + c_scale_msa) + c_shift_msa).type_as(
|
||||
hidden_states
|
||||
)
|
||||
ff_output = self.ffn(norm_hidden_states)
|
||||
hidden_states = (hidden_states.float() + ff_output.float() * c_gate_msa).type_as(hidden_states)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class WanTransformer3DModel(nn.Module):
|
||||
"""
|
||||
Optimized Wan Transformer model for video generation using vLLM layers.
|
||||
|
||||
This is an optimized version of the diffusers WanTransformer3DModel that uses
|
||||
vLLM's efficient QKVParallelLinear and RMSNorm implementations.
|
||||
|
||||
Args:
|
||||
patch_size: 3D patch dimensions for video embedding (t_patch, h_patch, w_patch)
|
||||
num_attention_heads: Number of attention heads
|
||||
attention_head_dim: Dimension of each attention head
|
||||
in_channels: Number of input channels
|
||||
out_channels: Number of output channels
|
||||
text_dim: Input dimension for text embeddings
|
||||
freq_dim: Dimension for sinusoidal time embeddings
|
||||
ffn_dim: Intermediate dimension in feed-forward network
|
||||
num_layers: Number of transformer blocks
|
||||
cross_attn_norm: Enable cross-attention normalization
|
||||
eps: Epsilon value for normalization layers
|
||||
image_dim: Optional image embedding dimension for I2V
|
||||
added_kv_proj_dim: Optional added KV projection dimension for I2V
|
||||
rope_max_seq_len: Maximum sequence length for rotary embeddings
|
||||
pos_embed_seq_len: Optional position embedding sequence length
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
patch_size: tuple[int, int, int] = (1, 2, 2),
|
||||
num_attention_heads: int = 40,
|
||||
attention_head_dim: int = 128,
|
||||
in_channels: int = 16,
|
||||
out_channels: int = 16,
|
||||
text_dim: int = 4096,
|
||||
freq_dim: int = 256,
|
||||
ffn_dim: int = 13824,
|
||||
num_layers: int = 40,
|
||||
cross_attn_norm: bool = True,
|
||||
eps: float = 1e-6,
|
||||
image_dim: Optional[int] = None,
|
||||
added_kv_proj_dim: Optional[int] = None,
|
||||
rope_max_seq_len: int = 1024,
|
||||
pos_embed_seq_len: Optional[int] = None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# Store config for compatibility
|
||||
self.config = type(
|
||||
"Config",
|
||||
(),
|
||||
{
|
||||
"patch_size": patch_size,
|
||||
"num_attention_heads": num_attention_heads,
|
||||
"attention_head_dim": attention_head_dim,
|
||||
"in_channels": in_channels,
|
||||
"out_channels": out_channels,
|
||||
"text_dim": text_dim,
|
||||
"freq_dim": freq_dim,
|
||||
"ffn_dim": ffn_dim,
|
||||
"num_layers": num_layers,
|
||||
"cross_attn_norm": cross_attn_norm,
|
||||
"eps": eps,
|
||||
"image_dim": image_dim,
|
||||
"added_kv_proj_dim": added_kv_proj_dim,
|
||||
"rope_max_seq_len": rope_max_seq_len,
|
||||
"pos_embed_seq_len": pos_embed_seq_len,
|
||||
},
|
||||
)()
|
||||
|
||||
inner_dim = num_attention_heads * attention_head_dim
|
||||
out_channels = out_channels or in_channels
|
||||
|
||||
# 1. Patch & position embedding
|
||||
self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len)
|
||||
self.patch_embedding = nn.Conv3d(in_channels, inner_dim, kernel_size=patch_size, stride=patch_size)
|
||||
|
||||
# 2. Condition embeddings
|
||||
self.condition_embedder = WanTimeTextImageEmbedding(
|
||||
dim=inner_dim,
|
||||
time_freq_dim=freq_dim,
|
||||
time_proj_dim=inner_dim * 6,
|
||||
text_embed_dim=text_dim,
|
||||
image_embed_dim=image_dim,
|
||||
pos_embed_seq_len=pos_embed_seq_len,
|
||||
)
|
||||
|
||||
# 3. Transformer blocks
|
||||
self.blocks = nn.ModuleList(
|
||||
[
|
||||
WanTransformerBlock(inner_dim, ffn_dim, num_attention_heads, eps, added_kv_proj_dim, cross_attn_norm)
|
||||
for _ in range(num_layers)
|
||||
]
|
||||
)
|
||||
|
||||
# 4. Output norm & projection
|
||||
self.norm_out = FP32LayerNorm(inner_dim, eps, elementwise_affine=False)
|
||||
self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size))
|
||||
self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5)
|
||||
|
||||
@property
|
||||
def dtype(self) -> torch.dtype:
|
||||
"""Return the dtype of the model parameters."""
|
||||
return next(self.parameters()).dtype
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
timestep: torch.LongTensor,
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
encoder_hidden_states_image: Optional[torch.Tensor] = None,
|
||||
return_dict: bool = True,
|
||||
attention_kwargs: Optional[dict[str, Any]] = None,
|
||||
) -> Union[torch.Tensor, Transformer2DModelOutput]:
|
||||
batch_size, num_channels, num_frames, height, width = hidden_states.shape
|
||||
p_t, p_h, p_w = self.config.patch_size
|
||||
post_patch_num_frames = num_frames // p_t
|
||||
post_patch_height = height // p_h
|
||||
post_patch_width = width // p_w
|
||||
|
||||
rotary_emb = self.rope(hidden_states)
|
||||
|
||||
hidden_states = self.patch_embedding(hidden_states)
|
||||
hidden_states = hidden_states.flatten(2).transpose(1, 2)
|
||||
|
||||
# Handle timestep shape
|
||||
if timestep.ndim == 2:
|
||||
ts_seq_len = timestep.shape[1]
|
||||
timestep = timestep.flatten()
|
||||
else:
|
||||
ts_seq_len = None
|
||||
|
||||
temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
|
||||
timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len
|
||||
)
|
||||
if ts_seq_len is not None:
|
||||
timestep_proj = timestep_proj.unflatten(2, (6, -1))
|
||||
else:
|
||||
timestep_proj = timestep_proj.unflatten(1, (6, -1))
|
||||
|
||||
if encoder_hidden_states_image is not None:
|
||||
encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1)
|
||||
|
||||
# Transformer blocks
|
||||
for block in self.blocks:
|
||||
hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb)
|
||||
|
||||
# Output norm, projection & unpatchify
|
||||
if temb.ndim == 3:
|
||||
shift, scale = (self.scale_shift_table.unsqueeze(0).to(temb.device) + temb.unsqueeze(2)).chunk(2, dim=2)
|
||||
shift = shift.squeeze(2)
|
||||
scale = scale.squeeze(2)
|
||||
else:
|
||||
shift, scale = (self.scale_shift_table.to(temb.device) + temb.unsqueeze(1)).chunk(2, dim=1)
|
||||
|
||||
shift = shift.to(hidden_states.device)
|
||||
scale = scale.to(hidden_states.device)
|
||||
|
||||
hidden_states = (self.norm_out(hidden_states.float()) * (1 + scale) + shift).type_as(hidden_states)
|
||||
hidden_states = self.proj_out(hidden_states)
|
||||
|
||||
hidden_states = hidden_states.reshape(
|
||||
batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1
|
||||
)
|
||||
hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
|
||||
output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
|
||||
|
||||
if not return_dict:
|
||||
return (output,)
|
||||
|
||||
return Transformer2DModelOutput(sample=output)
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
"""
|
||||
Load weights from a pretrained model, handling the mapping from
|
||||
separate Q/K/V projections to fused QKV projections for self-attention.
|
||||
|
||||
Diffusers weight names:
|
||||
- blocks.N.attn1.to_q/to_k/to_v -> fused to blocks.N.attn1.to_qkv (self-attention)
|
||||
- blocks.N.attn2.to_q/to_k/to_v -> kept separate (cross-attention)
|
||||
- blocks.N.attn1.norm_q/norm_k -> QK normalization for self-attention
|
||||
|
||||
Returns:
|
||||
Set of parameter names that were successfully loaded.
|
||||
"""
|
||||
# Stacked params mapping for self-attention QKV fusion
|
||||
# Format: (param_name, shard_name, shard_id)
|
||||
# Note: Only fuse attn1 (self-attention), NOT attn2 (cross-attention)
|
||||
stacked_params_mapping = [
|
||||
# self-attention QKV fusion (attn1 only)
|
||||
(".attn1.to_qkv", ".attn1.to_q", "q"),
|
||||
(".attn1.to_qkv", ".attn1.to_k", "k"),
|
||||
(".attn1.to_qkv", ".attn1.to_v", "v"),
|
||||
]
|
||||
|
||||
params_dict = dict(self.named_parameters())
|
||||
loaded_params: set[str] = set()
|
||||
|
||||
for name, loaded_weight in weights:
|
||||
for param_name, weight_name, shard_id in stacked_params_mapping:
|
||||
if weight_name not in name:
|
||||
continue
|
||||
name = name.replace(weight_name, param_name)
|
||||
param = params_dict[name]
|
||||
weight_loader = param.weight_loader
|
||||
weight_loader(param, loaded_weight, shard_id)
|
||||
break
|
||||
else:
|
||||
param = params_dict[name]
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
weight_loader(param, loaded_weight)
|
||||
loaded_params.add(name)
|
||||
|
||||
return loaded_params
|
||||
@@ -24,6 +24,11 @@ _DIFFUSION_MODELS = {
|
||||
"pipeline_z_image",
|
||||
"ZImagePipeline",
|
||||
),
|
||||
"WanPipeline": (
|
||||
"wan2_2",
|
||||
"pipeline_wan2_2",
|
||||
"Wan22Pipeline",
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
@@ -61,6 +66,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
|
||||
"QwenImagePipeline": "get_qwen_image_post_process_func",
|
||||
"QwenImageEditPipeline": "get_qwen_image_edit_post_process_func",
|
||||
"ZImagePipeline": "get_post_process_func",
|
||||
"WanPipeline": "get_wan22_post_process_func",
|
||||
}
|
||||
|
||||
_DIFFUSION_PRE_PROCESS_FUNCS = {
|
||||
|
||||
@@ -124,10 +124,6 @@ class OmniDiffusionRequest:
|
||||
save_output: bool = True
|
||||
return_frames: bool = False
|
||||
|
||||
# TeaCache parameters
|
||||
enable_teacache: bool = False
|
||||
# teacache_params: TeaCacheParams | WanTeaCacheParams | None = None
|
||||
|
||||
# STA parameters
|
||||
STA_param: list | None = None
|
||||
is_cfg_negative: bool = False
|
||||
|
||||
@@ -84,12 +84,29 @@ class GPUWorker:
|
||||
)
|
||||
logger.info(f"Worker {self.rank}: Model loaded successfully.")
|
||||
|
||||
# Apply cache adapter (model_type is auto-injected in OmniDiffusionConfig.__post_init__)
|
||||
from vllm_omni.diffusion.cache.apply import setup_cache
|
||||
|
||||
self.od_config.cache_config["model_type"] = self.od_config.model_class_name
|
||||
|
||||
self.pipeline._cache_adapter = setup_cache(
|
||||
self.pipeline.transformer,
|
||||
cache_type=self.od_config.cache_adapter,
|
||||
cache_config=self.od_config.cache_config,
|
||||
)
|
||||
|
||||
def maybe_reset_cache(self) -> None:
|
||||
"""Reset cache state before each generation if applicable."""
|
||||
if self.pipeline._cache_adapter is not None:
|
||||
self.pipeline._cache_adapter.reset(self.pipeline.transformer)
|
||||
|
||||
@torch.inference_mode()
|
||||
def execute_model(self, reqs: list[OmniDiffusionRequest], od_config: OmniDiffusionConfig) -> DiffusionOutput:
|
||||
"""
|
||||
Execute a forward pass.
|
||||
"""
|
||||
assert self.pipeline is not None
|
||||
self.maybe_reset_cache()
|
||||
# TODO: dealing with first req for now
|
||||
req = reqs[0]
|
||||
output = self.pipeline.forward(req)
|
||||
|
||||
@@ -0,0 +1,25 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from .omni_connectors import (
|
||||
ConnectorSpec,
|
||||
MooncakeConnector,
|
||||
OmniConnectorBase,
|
||||
OmniConnectorFactory,
|
||||
OmniTransferConfig,
|
||||
SharedMemoryConnector,
|
||||
load_omni_transfer_config,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# Config
|
||||
"ConnectorSpec",
|
||||
"OmniTransferConfig",
|
||||
# Connectors
|
||||
"OmniConnectorBase",
|
||||
"OmniConnectorFactory",
|
||||
"MooncakeConnector",
|
||||
"SharedMemoryConnector",
|
||||
# Utilities
|
||||
"load_omni_transfer_config",
|
||||
]
|
||||
@@ -0,0 +1,37 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from .connectors.base import OmniConnectorBase
|
||||
from .connectors.mooncake_connector import MooncakeConnector
|
||||
from .connectors.shm_connector import SharedMemoryConnector
|
||||
from .factory import OmniConnectorFactory
|
||||
from .utils.config import ConnectorSpec, OmniTransferConfig
|
||||
from .utils.initialization import (
|
||||
build_stage_connectors,
|
||||
get_connectors_config_for_stage,
|
||||
get_stage_connector_config,
|
||||
initialize_connectors_from_config,
|
||||
initialize_orchestrator_connectors,
|
||||
load_omni_transfer_config,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# Config
|
||||
"ConnectorSpec",
|
||||
"OmniTransferConfig",
|
||||
# Base classes and implementations
|
||||
"OmniConnectorBase",
|
||||
# Factory
|
||||
"OmniConnectorFactory",
|
||||
# Specific implementations
|
||||
"MooncakeConnector",
|
||||
"SharedMemoryConnector",
|
||||
# Utilities
|
||||
"load_omni_transfer_config",
|
||||
"initialize_connectors_from_config",
|
||||
"get_connectors_config_for_stage",
|
||||
# Manager helpers
|
||||
"initialize_orchestrator_connectors",
|
||||
"get_stage_connector_config",
|
||||
"build_stage_connectors",
|
||||
]
|
||||
@@ -0,0 +1,168 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# temporary for compatibility with vllm_omni.entrypoints.omni_stage.py
|
||||
# and vllm_omni.entrypoints.omni_llm.py
|
||||
|
||||
import time
|
||||
from typing import Any, Callable, Optional
|
||||
|
||||
from .utils.logging import get_connector_logger
|
||||
|
||||
logger = get_connector_logger(__name__)
|
||||
|
||||
|
||||
def try_send_via_connector(
|
||||
connector: Any,
|
||||
stage_id: int,
|
||||
next_stage_id: int,
|
||||
req_id: str,
|
||||
next_inputs: Any,
|
||||
sampling_params: Any,
|
||||
original_prompt: Any,
|
||||
next_stage_queue_submit_fn: Callable[[dict[str, Any]], None],
|
||||
metrics: Any,
|
||||
) -> bool:
|
||||
"""
|
||||
Attempts to send data via OmniConnector.
|
||||
Returns True if successful, False otherwise.
|
||||
Encapsulates the logic of preparing payload, sending via connector,
|
||||
sending notification, and recording metrics.
|
||||
"""
|
||||
try:
|
||||
t0 = time.time()
|
||||
|
||||
# Prepare data for connector
|
||||
payload_data = {
|
||||
"engine_inputs": next_inputs,
|
||||
"sampling_params": sampling_params,
|
||||
"metadata": {
|
||||
"original_prompt": original_prompt,
|
||||
"stage_transition": f"{stage_id}->{next_stage_id}",
|
||||
"timestamp": time.time(),
|
||||
},
|
||||
}
|
||||
|
||||
# Send data via connector
|
||||
success, serialized_size, metadata = connector.put(str(stage_id), str(next_stage_id), str(req_id), payload_data)
|
||||
|
||||
if success:
|
||||
# Send lightweight notification via queue
|
||||
notify_payload = {
|
||||
"request_id": req_id,
|
||||
"sampling_params": sampling_params,
|
||||
"from_connector": True,
|
||||
"from_stage": str(stage_id),
|
||||
"to_stage": str(next_stage_id),
|
||||
"sent_ts": time.time(),
|
||||
}
|
||||
# Merge connector metadata (e.g. shm handle or inline data) into queue payload
|
||||
if metadata:
|
||||
notify_payload["connector_metadata"] = metadata
|
||||
|
||||
next_stage_queue_submit_fn(notify_payload)
|
||||
|
||||
t1 = time.time()
|
||||
tx_ms = (t1 - t0) * 1000.0
|
||||
|
||||
metrics.on_forward(
|
||||
stage_id,
|
||||
next_stage_id,
|
||||
req_id,
|
||||
serialized_size, # Use size from connector
|
||||
float(tx_ms),
|
||||
True, # Mark as using connector
|
||||
)
|
||||
return True
|
||||
else:
|
||||
# If put returned False, we let the caller handle fallback
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"[Orchestrator] OmniConnector failed for req %s: %s; falling back to queue",
|
||||
req_id,
|
||||
e,
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def try_recv_via_connector(
|
||||
task: dict[str, Any],
|
||||
connectors: dict[Any, Any],
|
||||
stage_id: int,
|
||||
) -> tuple[Any, Optional[dict[str, Any]]]:
|
||||
"""
|
||||
Attempts to resolve input data from either connector or IPC.
|
||||
Returns (engine_inputs, rx_metrics) or (None, None) if failed/skipped.
|
||||
"""
|
||||
rid = task["request_id"]
|
||||
|
||||
if task.get("from_connector"):
|
||||
from_stage = task.get("from_stage")
|
||||
to_stage = str(stage_id)
|
||||
|
||||
if not from_stage:
|
||||
logger.error(
|
||||
"[Stage-%s] 'from_connector' is true but 'from_stage' is missing for request %s", stage_id, rid
|
||||
)
|
||||
return None, None
|
||||
|
||||
# Get connector for this edge
|
||||
connector_key = (from_stage, to_stage)
|
||||
connector = connectors.get(connector_key)
|
||||
|
||||
if connector:
|
||||
try:
|
||||
# Get data from connector with timeout
|
||||
_t_start = time.time()
|
||||
connector_metadata = task.get("connector_metadata")
|
||||
payload = connector.get(from_stage, to_stage, str(rid), metadata=connector_metadata)
|
||||
_t_end = time.time()
|
||||
|
||||
if payload:
|
||||
if isinstance(payload, tuple):
|
||||
payload_data, serialized_size = payload
|
||||
else:
|
||||
payload_data = payload
|
||||
serialized_size = len(connector.serialize_obj(payload_data))
|
||||
else:
|
||||
payload_data = None
|
||||
serialized_size = 0
|
||||
|
||||
if payload_data and isinstance(payload_data, dict):
|
||||
ein = payload_data.get("engine_inputs")
|
||||
decode_ms = (_t_end - _t_start) * 1000.0
|
||||
|
||||
rx_metrics = {"rx_decode_time_ms": decode_ms, "rx_transfer_bytes": serialized_size}
|
||||
return ein, rx_metrics
|
||||
else:
|
||||
logger.error(
|
||||
"[Stage-%s] Failed to get data from connector for request %s or payload is empty", stage_id, rid
|
||||
)
|
||||
return None, None
|
||||
except Exception as e:
|
||||
logger.error("[Stage-%s] Error retrieving data from connector for request %s: %s", stage_id, rid, e)
|
||||
return None, None
|
||||
else:
|
||||
logger.error(
|
||||
"[Stage-%s] No connector found for edge %s -> %s for request %s", stage_id, from_stage, to_stage, rid
|
||||
)
|
||||
return None, None
|
||||
else:
|
||||
# Data comes from queue as usual (e.g. seed request for Stage-0)
|
||||
# Since fallback logic is deprecated, we assume this is a direct inputs payload.
|
||||
# We still need to decode it if it used SHM (via legacy stage_utils logic, or new shm_connector format)
|
||||
# For Stage-0 specifically, 'engine_inputs' is often directly in the task dict.
|
||||
|
||||
# Try to use the new stage_utils which uses OmniSerializer
|
||||
from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc_with_metrics
|
||||
|
||||
try:
|
||||
ein, metrics = maybe_load_from_ipc_with_metrics(task, "engine_inputs", "engine_inputs_shm")
|
||||
# If metrics are empty or zero, we might want to populate dummy metrics
|
||||
return ein, metrics
|
||||
except Exception:
|
||||
# If engine_inputs is missing, it might be a different kind of payload,
|
||||
# but for Stage-0 seed it should be there.
|
||||
# We'll return None to let caller handle error if strictly required.
|
||||
return None, None
|
||||
@@ -0,0 +1,2 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -0,0 +1,72 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Optional
|
||||
|
||||
from ..utils.logging import get_connector_logger
|
||||
|
||||
logger = get_connector_logger(__name__)
|
||||
|
||||
|
||||
class OmniConnectorBase(ABC):
|
||||
"""Base class for all OmniConnectors."""
|
||||
|
||||
@abstractmethod
|
||||
def put(
|
||||
self, from_stage: str, to_stage: str, request_id: str, data: Any
|
||||
) -> tuple[bool, int, Optional[dict[str, Any]]]:
|
||||
"""Store Python object, internal serialization handled by connector.
|
||||
|
||||
Args:
|
||||
from_stage: Source stage identifier
|
||||
to_stage: Destination stage identifier
|
||||
request_id: Unique request identifier
|
||||
data: Python object to store
|
||||
|
||||
Returns:
|
||||
tuple: (success: bool, serialized_size: int, metadata: Optional[dict])
|
||||
Metadata may contain transport-specific handles or inline data.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get(
|
||||
self, from_stage: str, to_stage: str, request_id: str, metadata: Optional[dict[str, Any]] = None
|
||||
) -> Optional[tuple[Any, int]]:
|
||||
"""Retrieve Python object and payload size (bytes).
|
||||
|
||||
Args:
|
||||
from_stage: Source stage identifier
|
||||
to_stage: Destination stage identifier
|
||||
request_id: Unique request identifier
|
||||
metadata: Optional transport-specific metadata from the put operation
|
||||
|
||||
Returns:
|
||||
Tuple of (Python object, serialized byte size) if found, None otherwise
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def cleanup(self, request_id: str) -> None:
|
||||
"""Clean up resources for a request."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def health(self) -> dict[str, Any]:
|
||||
"""Return health status and metrics."""
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def serialize_obj(obj: Any) -> bytes:
|
||||
"""Serialize a Python object to bytes using centralized serializer."""
|
||||
from ..utils.serialization import OmniSerializer
|
||||
|
||||
return OmniSerializer.serialize(obj)
|
||||
|
||||
@staticmethod
|
||||
def deserialize_obj(data: bytes) -> Any:
|
||||
"""Deserialize bytes to Python object using centralized serializer."""
|
||||
from ..utils.serialization import OmniSerializer
|
||||
|
||||
return OmniSerializer.deserialize(data)
|
||||
@@ -0,0 +1,170 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import time
|
||||
from typing import Any, Optional
|
||||
|
||||
from ..utils.logging import get_connector_logger
|
||||
from .base import OmniConnectorBase
|
||||
|
||||
logger = get_connector_logger(__name__)
|
||||
|
||||
try:
|
||||
from mooncake.store import MooncakeDistributedStore, ReplicateConfig
|
||||
except ImportError:
|
||||
try:
|
||||
from mooncake import MooncakeDistributedStore, ReplicateConfig
|
||||
except ImportError:
|
||||
logger.warning("Mooncake not available, MooncakeOmniConnector will not work")
|
||||
MooncakeDistributedStore = None
|
||||
ReplicateConfig = None
|
||||
|
||||
|
||||
class MooncakeConnector(OmniConnectorBase):
|
||||
"""Mooncake-based distributed connector for OmniConnector."""
|
||||
|
||||
def __init__(self, config: dict[str, Any]):
|
||||
if MooncakeDistributedStore is None:
|
||||
raise ImportError("Mooncake not available")
|
||||
|
||||
self.config = config
|
||||
self.host = config.get("host", "127.0.0.1")
|
||||
self.metadata = config.get("metadata_server", "http://127.0.0.1:8080/metadata")
|
||||
self.master = config.get("master", "127.0.0.1:50051")
|
||||
self.segment = config.get("segment", 512 * 1024 * 1024) # 512MB
|
||||
self.localbuf = config.get("localbuf", 64 * 1024 * 1024) # 64MB
|
||||
self.proto = config.get("proto", "tcp")
|
||||
self.rdma = config.get("rdma", "")
|
||||
|
||||
self.store: Optional[MooncakeDistributedStore] = None
|
||||
self.pin: Optional[ReplicateConfig] = None
|
||||
|
||||
self._metrics = {
|
||||
"puts": 0,
|
||||
"gets": 0,
|
||||
"bytes_transferred": 0,
|
||||
"errors": 0,
|
||||
"timeouts": 0,
|
||||
}
|
||||
|
||||
self._init_store()
|
||||
|
||||
def _make_key(self, rid: str, from_stage: str, to_stage: str) -> str:
|
||||
"""Generate store key for request between stages."""
|
||||
return f"{rid}/{from_stage}_to_{to_stage}"
|
||||
|
||||
def _init_store(self):
|
||||
"""Initialize Mooncake store."""
|
||||
try:
|
||||
self.store = MooncakeDistributedStore()
|
||||
rc = self.store.setup(
|
||||
self.host, self.metadata, self.segment, self.localbuf, self.proto, self.rdma, self.master
|
||||
)
|
||||
if rc != 0:
|
||||
raise RuntimeError(f"Mooncake setup failed: {rc}")
|
||||
|
||||
self.pin = ReplicateConfig()
|
||||
self.pin.with_soft_pin = True
|
||||
logger.info("MooncakeConnector initialized successfully")
|
||||
except Exception as e:
|
||||
logger.error("Failed to initialize Mooncake store: %s", e)
|
||||
raise
|
||||
|
||||
# Use base class serialization methods for consistency
|
||||
|
||||
def put(
|
||||
self, from_stage: str, to_stage: str, request_id: str, data: Any
|
||||
) -> tuple[bool, int, Optional[dict[str, Any]]]:
|
||||
if not self.store:
|
||||
logger.error("Store not initialized")
|
||||
return False, 0, None
|
||||
|
||||
try:
|
||||
serialized_data = self.serialize_obj(data)
|
||||
key = self._make_key(request_id, from_stage, to_stage)
|
||||
self.store.put(key, serialized_data, self.pin)
|
||||
|
||||
self._metrics["puts"] += 1
|
||||
self._metrics["bytes_transferred"] += len(serialized_data)
|
||||
|
||||
logger.debug(
|
||||
"MooncakeConnector: stored %s (%s -> %s) %d bytes",
|
||||
key,
|
||||
from_stage,
|
||||
to_stage,
|
||||
len(serialized_data),
|
||||
)
|
||||
return True, len(serialized_data), None
|
||||
|
||||
except Exception as e:
|
||||
self._metrics["errors"] += 1
|
||||
logger.error("MooncakeConnector put failed: %s", e)
|
||||
return False, 0, None
|
||||
|
||||
def get(
|
||||
self, from_stage: str, to_stage: str, request_id: str, metadata: Optional[dict[str, Any]] = None
|
||||
) -> Optional[tuple[Any, int]]:
|
||||
if not self.store:
|
||||
logger.error("Store not initialized")
|
||||
return None
|
||||
|
||||
retries = 20
|
||||
sleep_s = 0.05
|
||||
key = self._make_key(request_id, from_stage, to_stage)
|
||||
|
||||
for attempt in range(retries):
|
||||
try:
|
||||
raw_data = self.store.get(key)
|
||||
|
||||
if raw_data:
|
||||
data = self.deserialize_obj(raw_data)
|
||||
self._metrics["gets"] += 1
|
||||
payload_size = len(raw_data)
|
||||
logger.debug(
|
||||
"MooncakeConnector: retrieved %s (%s -> %s) %d bytes",
|
||||
key,
|
||||
from_stage,
|
||||
to_stage,
|
||||
payload_size,
|
||||
)
|
||||
return data, payload_size
|
||||
|
||||
except Exception as e:
|
||||
logger.debug("MooncakeConnector get attempt %s failed: %s", attempt, e)
|
||||
|
||||
if attempt < retries - 1:
|
||||
time.sleep(sleep_s)
|
||||
|
||||
self._metrics["timeouts"] += 1
|
||||
logger.warning("MooncakeConnector: timeout waiting for %s", key)
|
||||
return None
|
||||
|
||||
def cleanup(self, request_id: str) -> None:
|
||||
if not self.store:
|
||||
return
|
||||
|
||||
# Note: Mooncake doesn't have explicit delete, data will be garbage collected
|
||||
# We could implement a cleanup mechanism by storing deletion markers
|
||||
logger.debug("MooncakeConnector: cleanup requested for %s (no-op)", request_id)
|
||||
|
||||
def health(self) -> dict[str, Any]:
|
||||
if not self.store:
|
||||
return {"status": "unhealthy", "error": "Store not initialized"}
|
||||
|
||||
return {
|
||||
"status": "healthy",
|
||||
"host": self.host,
|
||||
"metadata_server": self.metadata,
|
||||
"master": self.master,
|
||||
**self._metrics,
|
||||
}
|
||||
|
||||
def close(self):
|
||||
"""Clean shutdown."""
|
||||
if self.store:
|
||||
try:
|
||||
self.store.close()
|
||||
self.store = None
|
||||
logger.info("MooncakeConnector closed")
|
||||
except Exception as e:
|
||||
logger.error("Error closing Mooncake store: %s", e)
|
||||
@@ -0,0 +1,118 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from typing import Any, Optional
|
||||
|
||||
from vllm_omni.entrypoints.stage_utils import shm_read_bytes, shm_write_bytes
|
||||
|
||||
from ..utils.logging import get_connector_logger
|
||||
from .base import OmniConnectorBase
|
||||
|
||||
logger = get_connector_logger(__name__)
|
||||
|
||||
|
||||
class SharedMemoryConnector(OmniConnectorBase):
|
||||
"""
|
||||
Connector that uses SharedMemory for large objects and inline data for small objects.
|
||||
Acts as a unified replacement for the legacy IPC fallback logic.
|
||||
"""
|
||||
|
||||
def __init__(self, config: dict[str, Any]):
|
||||
self.config = config
|
||||
# Default threshold matches legacy behavior (64KB)
|
||||
self.threshold = int(config.get("shm_threshold_bytes", 65536))
|
||||
self._metrics = {
|
||||
"puts": 0,
|
||||
"gets": 0,
|
||||
"bytes_transferred": 0,
|
||||
"shm_writes": 0,
|
||||
"inline_writes": 0,
|
||||
}
|
||||
|
||||
def put(
|
||||
self, from_stage: str, to_stage: str, request_id: str, data: Any
|
||||
) -> tuple[bool, int, Optional[dict[str, Any]]]:
|
||||
try:
|
||||
# Always serialize first to check size (and for SHM writing)
|
||||
# Note: For extremely large objects in "inline" mode (e.g. Ray),
|
||||
# we might double-serialize if we're not careful, but here we assume
|
||||
# if it's huge we use SHM, or if Ray, threshold is maxsize.
|
||||
payload = self.serialize_obj(data)
|
||||
size = len(payload)
|
||||
|
||||
if size > self.threshold:
|
||||
# Use Shared Memory
|
||||
meta = shm_write_bytes(payload)
|
||||
# meta contains {'name': ..., 'size': ...}
|
||||
metadata = {"shm": meta, "size": size}
|
||||
self._metrics["shm_writes"] += 1
|
||||
else:
|
||||
# Inline - pass bytes directly to avoid double serialization of the object
|
||||
# We already serialized it to check size, so we pass the bytes.
|
||||
# The Queue will pickle these bytes (fast), avoiding re-serializing the complex object.
|
||||
metadata = {"inline_bytes": payload, "size": size}
|
||||
self._metrics["inline_writes"] += 1
|
||||
|
||||
self._metrics["puts"] += 1
|
||||
self._metrics["bytes_transferred"] += size
|
||||
|
||||
return True, size, metadata
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"SharedMemoryConnector put failed for req {request_id}: {e}")
|
||||
return False, 0, None
|
||||
|
||||
def get(
|
||||
self, from_stage: str, to_stage: str, request_id: str, metadata: Optional[dict[str, Any]] = None
|
||||
) -> Optional[tuple[Any, int]]:
|
||||
if not metadata:
|
||||
logger.error(f"SharedMemoryConnector get called without metadata for req {request_id}")
|
||||
return None
|
||||
|
||||
try:
|
||||
obj = None
|
||||
size = 0
|
||||
|
||||
if "shm" in metadata:
|
||||
meta = metadata["shm"]
|
||||
# shm_read_bytes handles reading and unlinking
|
||||
data_bytes = shm_read_bytes(meta)
|
||||
obj = self.deserialize_obj(data_bytes)
|
||||
size = metadata.get("size", len(data_bytes))
|
||||
elif "inline_bytes" in metadata:
|
||||
# Deserialize bytes back to object
|
||||
payload = metadata["inline_bytes"]
|
||||
obj = self.deserialize_obj(payload)
|
||||
size = metadata.get("size", len(payload))
|
||||
elif "inline" in metadata:
|
||||
obj = metadata["inline"]
|
||||
size = metadata.get("size", 0)
|
||||
if size == 0:
|
||||
# Fallback if size wasn't recorded
|
||||
try:
|
||||
size = len(self.serialize_obj(obj))
|
||||
except Exception:
|
||||
pass
|
||||
else:
|
||||
logger.error(
|
||||
f"Unknown metadata format in SharedMemoryConnector for req {request_id}: {list(metadata.keys())}"
|
||||
)
|
||||
return None
|
||||
|
||||
self._metrics["gets"] += 1
|
||||
return obj, size
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"SharedMemoryConnector get failed for req {request_id}: {e}")
|
||||
return None
|
||||
|
||||
def cleanup(self, request_id: str) -> None:
|
||||
# SHM segments are automatically unlinked during 'get' (shm_read_bytes).
|
||||
# If 'get' is never called (e.g. error flow), the SHM segment might leak.
|
||||
# A robust implementation might track created segments and unlink them here
|
||||
# if they haven't been consumed.
|
||||
# For now, we rely on the consumer to read and unlink.
|
||||
pass
|
||||
|
||||
def health(self) -> dict[str, Any]:
|
||||
return {"status": "healthy", "threshold": self.threshold, **self._metrics}
|
||||
@@ -0,0 +1,87 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from .utils.logging import get_connector_logger
|
||||
|
||||
try:
|
||||
from .connectors.base import OmniConnectorBase
|
||||
from .utils.config import ConnectorSpec
|
||||
except ImportError:
|
||||
# Fallback for direct execution
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
|
||||
from omni_connectors.connectors.base import OmniConnectorBase
|
||||
from omni_connectors.utils.config import ConnectorSpec
|
||||
|
||||
logger = get_connector_logger(__name__)
|
||||
|
||||
|
||||
class OmniConnectorFactory:
|
||||
"""Factory for creating OmniConnectors."""
|
||||
|
||||
_registry: dict[str, Callable[[dict[str, Any]], OmniConnectorBase]] = {}
|
||||
|
||||
@classmethod
|
||||
def register_connector(cls, name: str, constructor: Callable[[dict[str, Any]], OmniConnectorBase]) -> None:
|
||||
"""Register a connector constructor."""
|
||||
if name in cls._registry:
|
||||
raise ValueError(f"Connector '{name}' is already registered.")
|
||||
cls._registry[name] = constructor
|
||||
logger.debug(f"Registered connector: {name}")
|
||||
|
||||
@classmethod
|
||||
def create_connector(cls, spec: ConnectorSpec) -> OmniConnectorBase:
|
||||
"""Create a connector from specification."""
|
||||
if spec.name not in cls._registry:
|
||||
raise ValueError(f"Unknown connector: {spec.name}. Available: {list(cls._registry.keys())}")
|
||||
|
||||
constructor = cls._registry[spec.name]
|
||||
try:
|
||||
connector = constructor(spec.extra)
|
||||
logger.info(f"Created connector: {spec.name}")
|
||||
return connector
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create connector {spec.name}: {e}")
|
||||
raise
|
||||
|
||||
@classmethod
|
||||
def list_registered_connectors(cls) -> list[str]:
|
||||
"""List all registered connector names."""
|
||||
return list(cls._registry.keys())
|
||||
|
||||
|
||||
# Register built-in connectors with lazy imports
|
||||
def _create_mooncake_connector(config: dict[str, Any]) -> OmniConnectorBase:
|
||||
try:
|
||||
from .connectors.mooncake_connector import MooncakeConnector
|
||||
except ImportError:
|
||||
# Fallback import
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
|
||||
from omni_connectors.connectors.mooncake_connector import MooncakeConnector
|
||||
return MooncakeConnector(config)
|
||||
|
||||
|
||||
def _create_shm_connector(config: dict[str, Any]) -> OmniConnectorBase:
|
||||
try:
|
||||
from .connectors.shm_connector import SharedMemoryConnector
|
||||
except ImportError:
|
||||
# Fallback import
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
|
||||
from omni_connectors.connectors.shm_connector import SharedMemoryConnector
|
||||
return SharedMemoryConnector(config)
|
||||
|
||||
|
||||
# Register connectors
|
||||
OmniConnectorFactory.register_connector("MooncakeConnector", _create_mooncake_connector)
|
||||
OmniConnectorFactory.register_connector("SharedMemoryConnector", _create_shm_connector)
|
||||
@@ -0,0 +1,2 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -0,0 +1,40 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Optional
|
||||
|
||||
from .logging import get_connector_logger
|
||||
|
||||
logger = get_connector_logger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ConnectorSpec:
|
||||
"""Specification for a connector instance."""
|
||||
|
||||
name: str # e.g., "MooncakeConnector", "SharedMemoryConnector"
|
||||
extra: dict[str, Any] = field(default_factory=dict) # backend-specific config
|
||||
|
||||
|
||||
@dataclass
|
||||
class OmniTransferConfig:
|
||||
"""
|
||||
Top-level configuration for OmniConnector system.
|
||||
Members:
|
||||
connectors: A dictionary of connectors, keyed by (from_stage, to_stage).
|
||||
default_connector: The default connector to use if no connector is specified for an edge.
|
||||
"""
|
||||
|
||||
# Direct mapping: (from_stage, to_stage) -> connector
|
||||
connectors: dict[tuple[str, str], ConnectorSpec] = field(default_factory=dict)
|
||||
default_connector: Optional[ConnectorSpec] = None
|
||||
|
||||
def get_connector_for_edge(self, from_stage: str, to_stage: str) -> Optional[ConnectorSpec]:
|
||||
"""Get connector spec for a specific edge."""
|
||||
edge_key = (from_stage, to_stage)
|
||||
return self.connectors.get(edge_key, self.default_connector)
|
||||
|
||||
def has_connector_for_edge(self, from_stage: str, to_stage: str) -> bool:
|
||||
"""Check if there's a connector configured for the edge."""
|
||||
return self.get_connector_for_edge(from_stage, to_stage) is not None
|
||||
@@ -0,0 +1,322 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Utilities for OmniConnector configuration and validation."""
|
||||
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Optional, Union
|
||||
|
||||
from ..factory import OmniConnectorFactory
|
||||
from .config import ConnectorSpec, OmniTransferConfig
|
||||
from .logging import get_connector_logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..connectors.base import OmniConnectorBase
|
||||
else:
|
||||
OmniConnectorBase = Any
|
||||
|
||||
logger = get_connector_logger(__name__)
|
||||
|
||||
|
||||
def initialize_connectors_from_config(
|
||||
config_path: Optional[Union[str, Path]] = None, default_shm_threshold: int = 65536
|
||||
) -> tuple[Optional[OmniTransferConfig], dict[tuple[str, str], OmniConnectorBase]]:
|
||||
"""
|
||||
Initialize connectors from configuration file.
|
||||
|
||||
Returns:
|
||||
tuple: (OmniTransferConfig, dict of {(from, to): connector_instance})
|
||||
"""
|
||||
transfer_config = load_omni_transfer_config(config_path, default_shm_threshold=default_shm_threshold)
|
||||
|
||||
if not transfer_config:
|
||||
logger.info("No OmniTransferConfig provided")
|
||||
return None, {}
|
||||
|
||||
# create connectors from config
|
||||
connectors = create_connectors_from_config(transfer_config.connectors)
|
||||
return transfer_config, connectors
|
||||
|
||||
|
||||
def create_connectors_from_config(
|
||||
connectors_config: dict[tuple[str, str], ConnectorSpec],
|
||||
) -> dict[tuple[str, str], OmniConnectorBase]:
|
||||
"""
|
||||
Create connectors from config.
|
||||
|
||||
Args:
|
||||
connectors_config: A dictionary of connector configurations.
|
||||
|
||||
Returns:
|
||||
A dictionary of connectors.
|
||||
"""
|
||||
connectors = {}
|
||||
for edge_key, connector_spec in connectors_config.items():
|
||||
try:
|
||||
connector = OmniConnectorFactory.create_connector(connector_spec)
|
||||
connectors[edge_key] = connector
|
||||
logger.info(f"Created connector for {edge_key[0]} -> {edge_key[1]}: {type(connector).__name__}")
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to initialize connector for edge {edge_key}: {e}") from e
|
||||
|
||||
return connectors
|
||||
|
||||
|
||||
def get_connectors_config_for_stage(
|
||||
transfer_config: Optional[OmniTransferConfig], stage_id: Union[str, int]
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Extract connector configurations relevant for a specific stage worker.
|
||||
|
||||
Returns a dict compatible with worker initialization:
|
||||
{
|
||||
"from_stage_X": {
|
||||
"spec": {
|
||||
"name": "ConnectorName",
|
||||
"extra": {...}
|
||||
}
|
||||
},
|
||||
...
|
||||
}
|
||||
"""
|
||||
if not transfer_config:
|
||||
return {}
|
||||
|
||||
stage_connectors_config = {}
|
||||
target_stage = str(stage_id)
|
||||
|
||||
# Iterate through all configured edges
|
||||
for (from_stage, to_stage), spec in transfer_config.connectors.items():
|
||||
# We only care about incoming edges for the worker process
|
||||
# (Worker needs to create connectors to receive data)
|
||||
if to_stage == target_stage:
|
||||
stage_connectors_config[f"from_stage_{from_stage}"] = {"spec": {"name": spec.name, "extra": spec.extra}}
|
||||
|
||||
return stage_connectors_config
|
||||
|
||||
|
||||
def load_omni_transfer_config(
|
||||
config_path: Optional[Union[str, Path]] = None,
|
||||
config_dict: Optional[dict[str, Any]] = None,
|
||||
default_shm_threshold: int = 65536,
|
||||
) -> Optional[OmniTransferConfig]:
|
||||
"""Load OmniTransferConfig from file or dict."""
|
||||
if config_path is None and config_dict is None:
|
||||
# Even if no config provided, we might want to return a default config with SHM connectors
|
||||
# But without stage info we can't do much.
|
||||
return None
|
||||
|
||||
if config_path is not None:
|
||||
config_path = Path(config_path)
|
||||
if not config_path.exists():
|
||||
raise FileNotFoundError(f"Config file not found: {config_path}")
|
||||
|
||||
with open(config_path, encoding="utf-8") as f:
|
||||
if config_path.suffix.lower() == ".json":
|
||||
config_dict = json.load(f)
|
||||
elif config_path.suffix.lower() in [".yaml", ".yml"]:
|
||||
try:
|
||||
import yaml
|
||||
|
||||
config_dict = yaml.safe_load(f)
|
||||
except ImportError:
|
||||
raise ImportError("PyYAML required for YAML config files")
|
||||
else:
|
||||
raise ValueError(f"Unsupported config file format: {config_path.suffix}")
|
||||
|
||||
if config_dict is None:
|
||||
return None
|
||||
|
||||
# Parse connectors
|
||||
connectors = {}
|
||||
runtime_config = config_dict.get("runtime", {})
|
||||
|
||||
# Parse global connectors (from runtime.connectors)
|
||||
global_connectors = runtime_config.get("connectors", {})
|
||||
|
||||
# Parse stage-level connectors
|
||||
stage_args = config_dict.get("stage_args", [])
|
||||
expected_edges: set[tuple[str, str]] = set()
|
||||
for stage_config in stage_args:
|
||||
stage_id = str(stage_config["stage_id"])
|
||||
|
||||
# Input connectors
|
||||
for input_key, conn_ref in stage_config.get("input_connectors", {}).items():
|
||||
if isinstance(conn_ref, str):
|
||||
# Reference to global connector
|
||||
if conn_ref in global_connectors:
|
||||
conn_config = global_connectors[conn_ref]
|
||||
connector = ConnectorSpec(name=conn_config["name"], extra=conn_config.get("extra", {}))
|
||||
else:
|
||||
raise ValueError(f"Undefined connector reference: {conn_ref}")
|
||||
else:
|
||||
# Inline connector definition
|
||||
connector = ConnectorSpec(name=conn_ref["name"], extra=conn_ref.get("extra", {}))
|
||||
|
||||
# Parse from_stage from key (e.g., "from_stage_0" -> "0")
|
||||
from_stage = input_key.replace("from_stage_", "")
|
||||
edge_key = (from_stage, stage_id)
|
||||
connectors[edge_key] = connector
|
||||
expected_edges.add(edge_key)
|
||||
|
||||
# Output connectors
|
||||
for output_key, conn_ref in stage_config.get("output_connectors", {}).items():
|
||||
if isinstance(conn_ref, str):
|
||||
# Reference to global connector
|
||||
if conn_ref in global_connectors:
|
||||
conn_config = global_connectors[conn_ref]
|
||||
connector = ConnectorSpec(name=conn_config["name"], extra=conn_config.get("extra", {}))
|
||||
else:
|
||||
raise ValueError(f"Undefined connector reference: {conn_ref}")
|
||||
else:
|
||||
# Inline connector definition
|
||||
connector = ConnectorSpec(name=conn_ref["name"], extra=conn_ref.get("extra", {}))
|
||||
|
||||
# Parse to_stage from key (e.g., "to_stage_1" -> "1")
|
||||
to_stage = output_key.replace("to_stage_", "")
|
||||
edge_key = (stage_id, to_stage)
|
||||
connectors[edge_key] = connector
|
||||
expected_edges.add(edge_key)
|
||||
|
||||
# Auto-configure SharedMemoryConnector for missing edges based on runtime edges / engine_input_source
|
||||
if stage_args:
|
||||
try:
|
||||
# Prefer explicit runtime edges if provided
|
||||
runtime_edges = runtime_config.get("edges", [])
|
||||
if isinstance(runtime_edges, list) and runtime_edges:
|
||||
for edge in runtime_edges:
|
||||
from_stage = edge.get("from")
|
||||
to_stage = edge.get("to")
|
||||
if from_stage is None or to_stage is None:
|
||||
continue
|
||||
edge_key = (str(from_stage), str(to_stage))
|
||||
expected_edges.add(edge_key)
|
||||
if edge_key not in connectors:
|
||||
logger.info(f"Auto-configuring SharedMemoryConnector for edge {edge_key}")
|
||||
connectors[edge_key] = ConnectorSpec(
|
||||
name="SharedMemoryConnector",
|
||||
extra={"shm_threshold_bytes": default_shm_threshold},
|
||||
)
|
||||
|
||||
# Fallback: infer edges from engine_input_source for each stage
|
||||
for stage_config in stage_args:
|
||||
to_stage = str(stage_config["stage_id"])
|
||||
# Check explicit input sources
|
||||
sources = stage_config.get("engine_input_source", [])
|
||||
|
||||
for from_stage in sources:
|
||||
from_stage_str = str(from_stage)
|
||||
edge_key = (from_stage_str, to_stage)
|
||||
expected_edges.add(edge_key)
|
||||
|
||||
if edge_key not in connectors:
|
||||
logger.info(f"Auto-configuring SharedMemoryConnector for edge {edge_key}")
|
||||
connectors[edge_key] = ConnectorSpec(
|
||||
name="SharedMemoryConnector", extra={"shm_threshold_bytes": default_shm_threshold}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to auto-configure SHM connectors: {e}")
|
||||
|
||||
# Fail fast if any expected edge is still missing a connector
|
||||
missing_edges = [edge for edge in expected_edges if edge not in connectors]
|
||||
if missing_edges:
|
||||
missing_str = ", ".join([f"{f}->{t}" for f, t in missing_edges])
|
||||
raise ValueError(
|
||||
"Connector configuration missing for edges: "
|
||||
f"{missing_str}. Define connectors or allow auto SHM creation for these edges."
|
||||
)
|
||||
|
||||
config = OmniTransferConfig(connectors=connectors)
|
||||
|
||||
logger.info(f"Loaded OmniTransferConfig with {len(connectors)} connector configurations")
|
||||
return config
|
||||
|
||||
|
||||
# High-level management functions
|
||||
|
||||
|
||||
def initialize_orchestrator_connectors(
|
||||
config_path: Optional[str], worker_backend: Optional[str] = "multi_process", shm_threshold_bytes: int = 65536
|
||||
) -> tuple[Optional[OmniTransferConfig], dict[tuple[str, str], OmniConnectorBase]]:
|
||||
"""Initialize connectors shared at orchestrator level.
|
||||
Args:
|
||||
config_path: The path to the configuration file.
|
||||
worker_backend: The backend to use for the worker.
|
||||
Returns:
|
||||
A tuple containing the OmniTransferConfig and a dictionary of connectors.
|
||||
"""
|
||||
if worker_backend == "ray":
|
||||
default_shm_threshold = sys.maxsize
|
||||
else:
|
||||
default_shm_threshold = max(0, shm_threshold_bytes)
|
||||
transfer_config, connectors = initialize_connectors_from_config(
|
||||
config_path, default_shm_threshold=default_shm_threshold
|
||||
)
|
||||
return transfer_config, connectors
|
||||
|
||||
|
||||
def get_stage_connector_config(
|
||||
transfer_config: Optional[OmniTransferConfig],
|
||||
stage_id: int,
|
||||
) -> dict[str, Any]:
|
||||
"""Return the serialized connector config payload for a specific stage."""
|
||||
if transfer_config is None:
|
||||
return {}
|
||||
|
||||
try:
|
||||
return get_connectors_config_for_stage(transfer_config, stage_id)
|
||||
except Exception as exc: # pragma: no cover - defensive logging
|
||||
logger.warning(
|
||||
"Failed to build connector config for stage %s: %s. Using IPC fallback.",
|
||||
stage_id,
|
||||
exc,
|
||||
)
|
||||
return {}
|
||||
|
||||
|
||||
def build_stage_connectors(
|
||||
stage_id: int,
|
||||
connectors_config: dict[str, Any],
|
||||
) -> Optional[dict[tuple[str, str], Any]]:
|
||||
"""Instantiate OmniConnectors for a stage based on config."""
|
||||
if not connectors_config:
|
||||
return {}
|
||||
|
||||
logger.info(
|
||||
"[Stage-%s] Initializing OmniConnectors with config keys: %s",
|
||||
stage_id,
|
||||
list(connectors_config.keys()),
|
||||
)
|
||||
|
||||
from .config import ConnectorSpec
|
||||
|
||||
connectors: dict[tuple[str, str], Any] = {}
|
||||
# 将字典格式的配置转换为ConnectorSpec对象
|
||||
stage_connector_specs = {}
|
||||
for input_key, config in connectors_config.items():
|
||||
if not input_key.startswith("from_stage_"):
|
||||
continue
|
||||
|
||||
from_stage = input_key.replace("from_stage_", "")
|
||||
spec_dict = config.get("spec", {})
|
||||
if not spec_dict:
|
||||
continue
|
||||
|
||||
connector_spec = ConnectorSpec(
|
||||
name=spec_dict.get("name", "SharedMemoryConnector"),
|
||||
extra=spec_dict.get("extra", {}),
|
||||
)
|
||||
stage_connector_specs[(str(from_stage), str(stage_id))] = connector_spec
|
||||
|
||||
try:
|
||||
# 使用统一的连接器创建逻辑
|
||||
connectors = create_connectors_from_config(stage_connector_specs)
|
||||
except Exception as exc: # pragma: no cover - defensive logging
|
||||
# Fail fast so the stage does not start with missing connectors.
|
||||
logger.exception("[Stage-%s] Failed to initialize connectors: %s", stage_id, exc)
|
||||
raise
|
||||
|
||||
return connectors
|
||||
@@ -0,0 +1,14 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import logging
|
||||
|
||||
try:
|
||||
from vllm.logger import init_logger as _vllm_init_logger
|
||||
except Exception: # pragma: no cover - optional dependency
|
||||
_vllm_init_logger = None
|
||||
|
||||
|
||||
def get_connector_logger(name: str) -> logging.Logger:
|
||||
"""Return a logger preferring vLLM's init_logger when available."""
|
||||
return _vllm_init_logger(name) if _vllm_init_logger else logging.getLogger(name)
|
||||
@@ -0,0 +1,79 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import io
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
import torch
|
||||
|
||||
_has_torch = True
|
||||
except ImportError:
|
||||
_has_torch = False
|
||||
|
||||
try:
|
||||
import cloudpickle
|
||||
|
||||
_has_cloudpickle = True
|
||||
except ImportError:
|
||||
_has_cloudpickle = False
|
||||
|
||||
from .logging import get_connector_logger
|
||||
|
||||
logger = get_connector_logger(__name__)
|
||||
|
||||
|
||||
class OmniSerializer:
|
||||
"""
|
||||
Centralized serialization handler for OmniConnectors.
|
||||
Supports multiple backends (torch, cloudpickle) to ensure consistency across connectors.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def serialize(obj: Any, method: str = "cloudpickle") -> bytes:
|
||||
"""
|
||||
Serialize an object to bytes.
|
||||
|
||||
Args:
|
||||
obj: The object to serialize.
|
||||
method: Serialization method ("cloudpickle" or "torch").
|
||||
Defaults to "cloudpickle" (consistent with legacy stage_utils).
|
||||
"""
|
||||
if method == "torch" and _has_torch:
|
||||
buf = io.BytesIO()
|
||||
torch.save(obj, buf)
|
||||
return buf.getvalue()
|
||||
elif method == "cloudpickle" or (method == "torch" and not _has_torch):
|
||||
if not _has_cloudpickle:
|
||||
# Fallback to standard pickle if cloudpickle is missing
|
||||
import pickle
|
||||
|
||||
return pickle.dumps(obj)
|
||||
return cloudpickle.dumps(obj)
|
||||
else:
|
||||
# Default fallback
|
||||
import pickle
|
||||
|
||||
return pickle.dumps(obj)
|
||||
|
||||
@staticmethod
|
||||
def deserialize(data: bytes, method: str = "cloudpickle") -> Any:
|
||||
"""
|
||||
Deserialize bytes to an object.
|
||||
|
||||
Args:
|
||||
data: The bytes to deserialize.
|
||||
method: Serialization method used ("cloudpickle" or "torch").
|
||||
"""
|
||||
if method == "torch" and _has_torch:
|
||||
return torch.load(io.BytesIO(data), map_location="cpu", weights_only=False)
|
||||
elif method == "cloudpickle" or (method == "torch" and not _has_torch):
|
||||
if not _has_cloudpickle:
|
||||
import pickle
|
||||
|
||||
return pickle.loads(data)
|
||||
return cloudpickle.loads(data)
|
||||
else:
|
||||
import pickle
|
||||
|
||||
return pickle.loads(data)
|
||||
@@ -0,0 +1,10 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from .utils import calculate_total_bytes, is_ray_initialized, maybe_disable_pin_memory_for_ray
|
||||
|
||||
__all__ = [
|
||||
"calculate_total_bytes",
|
||||
"is_ray_initialized",
|
||||
"maybe_disable_pin_memory_for_ray",
|
||||
]
|
||||
@@ -0,0 +1,183 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import logging
|
||||
import os
|
||||
from contextlib import contextmanager
|
||||
from typing import Any, Optional
|
||||
|
||||
import torch
|
||||
|
||||
try:
|
||||
import ray
|
||||
from ray.util.queue import Queue as RayQueue
|
||||
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
|
||||
|
||||
RAY_AVAILABLE = True
|
||||
from ray.util.placement_group import PlacementGroup
|
||||
except ImportError:
|
||||
ray = None
|
||||
RayQueue = None
|
||||
PlacementGroupSchedulingStrategy = None
|
||||
RAY_AVAILABLE = False
|
||||
PlacementGroup = Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def is_ray_initialized():
|
||||
"""Check if Ray is initialized without hard dependency on Ray."""
|
||||
# 1. Try standard API
|
||||
if RAY_AVAILABLE:
|
||||
if ray.is_initialized():
|
||||
return True
|
||||
# 2. Fallback: Check environment variables typical for Ray Workers
|
||||
# RAY_RAYLET_PID is always set in Ray workers
|
||||
if "RAY_RAYLET_PID" in os.environ:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def calculate_total_bytes(size_args, dtype):
|
||||
"""
|
||||
Calculate total bytes for a tensor allocation, handling nested tuples in size args.
|
||||
"""
|
||||
num_elements = 1
|
||||
for s in size_args:
|
||||
if isinstance(s, (tuple, list)):
|
||||
for inner in s:
|
||||
num_elements *= inner
|
||||
else:
|
||||
num_elements *= s
|
||||
|
||||
element_size = torch.tensor([], dtype=dtype).element_size()
|
||||
return num_elements * element_size
|
||||
|
||||
|
||||
@contextmanager
|
||||
def maybe_disable_pin_memory_for_ray(obj, size_bytes, threshold=32 * 1024 * 1024):
|
||||
"""
|
||||
Context manager to temporarily disable pin_memory if running in Ray and
|
||||
the allocation size exceeds the threshold.
|
||||
|
||||
This is a workaround for Ray workers often having low ulimit -l (locked memory),
|
||||
causing OS call failed errors when allocating large pinned buffers.
|
||||
"""
|
||||
should_disable = False
|
||||
old_pin = False
|
||||
|
||||
# Check 1: Are we in a Ray-like environment?
|
||||
in_ray = is_ray_initialized()
|
||||
|
||||
# Check 2: Is the size large enough to worry?
|
||||
is_large = size_bytes > threshold
|
||||
|
||||
# Check 3: Is pinning currently enabled?
|
||||
is_pinned = getattr(obj, "pin_memory", False)
|
||||
|
||||
if in_ray and is_large and is_pinned:
|
||||
should_disable = True
|
||||
old_pin = obj.pin_memory
|
||||
obj.pin_memory = False
|
||||
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
if should_disable:
|
||||
obj.pin_memory = old_pin
|
||||
|
||||
|
||||
# --- Ray specific utilities ---
|
||||
|
||||
|
||||
def get_ray_queue_class():
|
||||
if not RAY_AVAILABLE:
|
||||
raise ImportError("ray is required for worker_backend='ray'")
|
||||
return lambda: RayQueue(maxsize=0)
|
||||
|
||||
|
||||
def initialize_ray_cluster(address: Optional[str] = None):
|
||||
if not RAY_AVAILABLE:
|
||||
logger.warning("Ray is not available, skipping initialization.")
|
||||
return
|
||||
|
||||
if not ray.is_initialized():
|
||||
# Pass current PYTHONPATH to workers to ensure they can find vllm_omni
|
||||
runtime_env = {"env_vars": {"PYTHONPATH": os.environ.get("PYTHONPATH", "")}}
|
||||
ray.init(address=address, ignore_reinit_error=True, runtime_env=runtime_env)
|
||||
|
||||
|
||||
def create_placement_group(
|
||||
number_of_stages: int, address: Optional[str] = None, strategy: str = "PACK"
|
||||
) -> PlacementGroup:
|
||||
"""Create a placement group for the given number of stages.
|
||||
Args:
|
||||
number_of_stages: The number of stages to create the placement group for.
|
||||
strategy: The strategy to use for the placement group.
|
||||
Returns:
|
||||
The placement group.
|
||||
"""
|
||||
if not RAY_AVAILABLE:
|
||||
raise ImportError("ray is required for creating placement group")
|
||||
|
||||
# Initialize Ray if not already initialized (using default args if needed)
|
||||
if not ray.is_initialized():
|
||||
logger.warning("[Orchestrator] Ray is not initialized. Initializing with default settings.")
|
||||
initialize_ray_cluster(address)
|
||||
|
||||
bundles = [{"GPU": 1.0, "CPU": 1.0} for _ in range(number_of_stages)]
|
||||
pg = ray.util.placement_group(bundles, strategy=strategy)
|
||||
ray.get(pg.ready())
|
||||
logger.info("[Orchestrator] Ray Placement Group created")
|
||||
return pg
|
||||
|
||||
|
||||
def remove_placement_group(pg):
|
||||
if pg and RAY_AVAILABLE:
|
||||
try:
|
||||
ray.util.remove_placement_group(pg)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to remove placement group: {e}")
|
||||
|
||||
|
||||
def try_close_ray(pg=None):
|
||||
"""Try to clean up Ray resources including placement group and shutdown."""
|
||||
if pg:
|
||||
remove_placement_group(pg)
|
||||
# Note: We typically don't shutdown ray.init() here as it might be used by other components
|
||||
# or the user might want it to persist. If full shutdown is needed, ray.shutdown() can be called.
|
||||
|
||||
|
||||
def kill_ray_actor(actor):
|
||||
if actor and RAY_AVAILABLE:
|
||||
try:
|
||||
ray.kill(actor)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to kill actor: {e}")
|
||||
|
||||
|
||||
def start_ray_actor(
|
||||
worker_entry_fn,
|
||||
placement_group,
|
||||
placement_group_bundle_index: int,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
if not RAY_AVAILABLE:
|
||||
raise ImportError("ray is required for starting ray actor")
|
||||
|
||||
@ray.remote(num_gpus=1)
|
||||
class OmniStageRayWorker:
|
||||
def run(self, func, *args, **kwargs):
|
||||
return func(*args, **kwargs)
|
||||
|
||||
worker_actor = OmniStageRayWorker.options(
|
||||
scheduling_strategy=PlacementGroupSchedulingStrategy(
|
||||
placement_group=placement_group, placement_group_bundle_index=placement_group_bundle_index
|
||||
),
|
||||
runtime_env={"env_vars": {"PYTHONPATH": os.environ.get("PYTHONPATH", "")}, "CUDA_LAUNCH_BLOCKING": "1"},
|
||||
).remote()
|
||||
|
||||
worker_actor.run.remote(worker_entry_fn, *args, **kwargs)
|
||||
|
||||
return worker_actor
|
||||
@@ -33,6 +33,16 @@ from vllm.v1.metrics.loggers import StatLoggerFactory, StatLoggerManager
|
||||
|
||||
# Internal imports (our code)
|
||||
from vllm_omni.config import OmniModelConfig
|
||||
from vllm_omni.distributed.omni_connectors import (
|
||||
get_stage_connector_config,
|
||||
initialize_orchestrator_connectors,
|
||||
)
|
||||
from vllm_omni.distributed.omni_connectors.adapter import try_send_via_connector
|
||||
from vllm_omni.distributed.ray_utils.utils import (
|
||||
create_placement_group,
|
||||
get_ray_queue_class,
|
||||
try_close_ray,
|
||||
)
|
||||
from vllm_omni.engine.arg_utils import AsyncOmniEngineArgs
|
||||
from vllm_omni.engine.output_processor import MultimodalOutputProcessor
|
||||
from vllm_omni.engine.processor import OmniProcessor
|
||||
@@ -43,10 +53,12 @@ from vllm_omni.entrypoints.log_utils import (
|
||||
remove_old_logs,
|
||||
)
|
||||
from vllm_omni.entrypoints.omni_stage import OmniStage
|
||||
from vllm_omni.entrypoints.stage_utils import encode_for_ipc as _encode
|
||||
from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
|
||||
from vllm_omni.entrypoints.stage_utils import serialize_obj as _set
|
||||
from vllm_omni.entrypoints.utils import load_stage_configs_from_model, load_stage_configs_from_yaml
|
||||
from vllm_omni.entrypoints.utils import (
|
||||
load_stage_configs_from_model,
|
||||
load_stage_configs_from_yaml,
|
||||
resolve_model_config_path,
|
||||
)
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
logger = init_logger(__name__)
|
||||
@@ -93,16 +105,29 @@ class AsyncOmni(EngineClient):
|
||||
cli_args: Namespace,
|
||||
**kwargs: Any,
|
||||
):
|
||||
self.worker_backend = cli_args.worker_backend
|
||||
self.ray_address = cli_args.ray_address
|
||||
self._ray_pg = None
|
||||
|
||||
self.batch_timeout = cli_args.batch_timeout
|
||||
self._enable_stats: bool = bool(cli_args.log_stats)
|
||||
|
||||
base_engine_args = AsyncOmniEngineArgs.from_cli_args(cli_args).__dict__.copy()
|
||||
|
||||
if cli_args.stage_configs_path is None:
|
||||
self.config_path = resolve_model_config_path(model)
|
||||
self.stage_configs = load_stage_configs_from_model(model, base_engine_args)
|
||||
else:
|
||||
self.config_path = cli_args.stage_configs_path
|
||||
self.stage_configs = load_stage_configs_from_yaml(cli_args.stage_configs_path, base_engine_args)
|
||||
|
||||
shm_threshold_bytes = cli_args.shm_threshold_bytes
|
||||
|
||||
# Initialize connectors
|
||||
self.omni_transfer_config, self.connectors = initialize_orchestrator_connectors(
|
||||
self.config_path, worker_backend=self.worker_backend, shm_threshold_bytes=shm_threshold_bytes
|
||||
)
|
||||
|
||||
self.stage_list: list[OmniStage] = []
|
||||
self.default_sampling_params_list: list[SamplingParams] = []
|
||||
# Optional file handler for orchestrator
|
||||
@@ -138,34 +163,56 @@ class AsyncOmni(EngineClient):
|
||||
self.default_sampling_params_list = [st.default_sampling_params for st in self.stage_list]
|
||||
logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list))
|
||||
|
||||
self._ctx = mp.get_context("spawn")
|
||||
if self.worker_backend == "ray":
|
||||
self._queue_cls = get_ray_queue_class()
|
||||
else:
|
||||
self._ctx = mp.get_context("spawn")
|
||||
self._queue_cls = lambda: self._ctx.Queue(maxsize=0)
|
||||
|
||||
self._stage_in_queues: list[mp.Queue] = []
|
||||
self._stage_out_queues: list[mp.Queue] = []
|
||||
self._init_sleep_seconds = max(0, int(init_sleep_seconds))
|
||||
self._shm_threshold_bytes = max(0, int(shm_threshold_bytes))
|
||||
self._start_stage_processes(model)
|
||||
self._start_stages(model)
|
||||
# Wait for all stages to report readiness before seeding
|
||||
self._stages_ready: set[int] = set()
|
||||
self._wait_for_stages_ready(timeout=init_timeout)
|
||||
|
||||
def _start_stage_processes(self, model: str) -> None:
|
||||
def _start_stages(self, model: str) -> None:
|
||||
if self.worker_backend == "ray":
|
||||
# Initialize Ray Cluster
|
||||
self._ray_pg = create_placement_group(
|
||||
number_of_stages=len(self.stage_list), address=self.ray_address, strategy="PACK"
|
||||
)
|
||||
|
||||
for stage_id, stage in enumerate(self.stage_list):
|
||||
# Use unbounded queues to avoid deadlock when seeding many requests
|
||||
in_q: mp.Queue = self._ctx.Queue(maxsize=0)
|
||||
out_q: mp.Queue = self._ctx.Queue(maxsize=0)
|
||||
in_q = self._queue_cls()
|
||||
out_q = self._queue_cls()
|
||||
self._stage_in_queues.append(in_q)
|
||||
self._stage_out_queues.append(out_q)
|
||||
|
||||
# Attach queues and start Stage-owned worker process
|
||||
stage.attach_queues(in_q, out_q)
|
||||
|
||||
# Build connectors config for this stage
|
||||
stage_connectors_config = get_stage_connector_config(
|
||||
self.omni_transfer_config,
|
||||
stage_id,
|
||||
)
|
||||
|
||||
stage.init_stage_worker(
|
||||
model,
|
||||
is_async=True,
|
||||
log_file=self._log_file,
|
||||
shm_threshold_bytes=self._shm_threshold_bytes,
|
||||
ctx=self._ctx,
|
||||
ctx=self._ctx if self.worker_backend != "ray" else None,
|
||||
batch_timeout=self.batch_timeout,
|
||||
connectors_config=stage_connectors_config,
|
||||
worker_backend=self.worker_backend,
|
||||
ray_placement_group=self._ray_pg,
|
||||
)
|
||||
|
||||
logger.debug("[Orchestrator] Stage-%s process started", stage_id)
|
||||
time.sleep(self._init_sleep_seconds)
|
||||
|
||||
@@ -191,9 +238,9 @@ class AsyncOmni(EngineClient):
|
||||
except Exception as e:
|
||||
logger.warning("[Orchestrator] Failed to stop stage worker: %s", e)
|
||||
|
||||
try_close_ray(self._ray_pg)
|
||||
|
||||
def __del__(self) -> None: # best-effort
|
||||
print("[AsyncOmni] __del__ close()", flush=True)
|
||||
raise Exception("test")
|
||||
try:
|
||||
self.close()
|
||||
except Exception as e:
|
||||
@@ -266,12 +313,14 @@ class AsyncOmni(EngineClient):
|
||||
# Determine the final stage for E2E stats (highest stage_id with
|
||||
# final_output=True; fallback to last stage)
|
||||
final_stage_id_for_e2e = -1
|
||||
last_stage_id = len(self.stage_list) - 1
|
||||
try:
|
||||
for _sid, _st in enumerate(self.stage_list):
|
||||
if getattr(_st, "final_output", False):
|
||||
final_stage_id_for_e2e = max(final_stage_id_for_e2e, _sid)
|
||||
for _sid in range(last_stage_id, -1, -1):
|
||||
if getattr(self.stage_list[_sid], "final_output", False):
|
||||
final_stage_id_for_e2e = _sid
|
||||
break
|
||||
if final_stage_id_for_e2e < 0:
|
||||
final_stage_id_for_e2e = len(self.stage_list) - 1
|
||||
final_stage_id_for_e2e = last_stage_id
|
||||
except Exception as e:
|
||||
logger.debug(
|
||||
"[Orchestrator] Failed to determine final stage for E2E; \
|
||||
@@ -279,7 +328,7 @@ class AsyncOmni(EngineClient):
|
||||
e,
|
||||
exc_info=True,
|
||||
)
|
||||
final_stage_id_for_e2e = len(self.stage_list) - 1
|
||||
final_stage_id_for_e2e = last_stage_id
|
||||
# Metrics/aggregation helper
|
||||
metrics = OrchestratorMetrics(
|
||||
num_stages,
|
||||
@@ -362,8 +411,8 @@ class AsyncOmni(EngineClient):
|
||||
# End-to-end timing and time-per-token for final output
|
||||
# (only once per request at the designated final stage)
|
||||
try:
|
||||
rid_int = int(req_id) if isinstance(req_id, (int, str)) and str(req_id).isdigit() else req_id
|
||||
if stage_id == final_stage_id_for_e2e and rid_int not in metrics.e2e_done:
|
||||
rid_key = str(req_id)
|
||||
if stage_id == final_stage_id_for_e2e and rid_key not in metrics.e2e_done:
|
||||
metrics.on_finalize_request(
|
||||
stage_id,
|
||||
req_id,
|
||||
@@ -392,52 +441,36 @@ class AsyncOmni(EngineClient):
|
||||
next_stage: OmniStage = self.stage_list[next_stage_id]
|
||||
next_inputs = next_stage.process_engine_inputs(self.stage_list, prompt)
|
||||
sp_next: SamplingParams = sampling_params_list[next_stage_id]
|
||||
try:
|
||||
# Measure transfer size and time (encode + enqueue)
|
||||
size_bytes = 0
|
||||
try:
|
||||
size_bytes = len(_set(next_inputs))
|
||||
except Exception:
|
||||
size_bytes = 0
|
||||
t0 = time.time()
|
||||
ipc_payload = _encode(
|
||||
next_inputs,
|
||||
getattr(self, "_shm_threshold_bytes", 65536),
|
||||
obj_key="engine_inputs",
|
||||
shm_key="engine_inputs_shm",
|
||||
|
||||
# Check if we have a connector for this edge
|
||||
connector_key = (str(stage_id), str(next_stage_id))
|
||||
connector = self.connectors.get(connector_key)
|
||||
|
||||
sent_via_connector = False
|
||||
if connector:
|
||||
sent_via_connector = try_send_via_connector(
|
||||
connector=connector,
|
||||
stage_id=stage_id,
|
||||
next_stage_id=next_stage_id,
|
||||
req_id=req_id,
|
||||
next_inputs=next_inputs,
|
||||
sampling_params=sp_next,
|
||||
original_prompt=prompt,
|
||||
next_stage_queue_submit_fn=self.stage_list[next_stage_id].submit,
|
||||
metrics=metrics,
|
||||
)
|
||||
ipc_payload.update(
|
||||
{
|
||||
"request_id": req_id,
|
||||
"sampling_params": sp_next,
|
||||
"sent_ts": time.time(),
|
||||
}
|
||||
)
|
||||
self.stage_list[next_stage_id].submit(ipc_payload)
|
||||
t1 = time.time()
|
||||
tx_ms = (t1 - t0) * 1000.0
|
||||
metrics.on_forward(
|
||||
stage_id,
|
||||
next_stage_id,
|
||||
req_id,
|
||||
int(size_bytes),
|
||||
float(tx_ms),
|
||||
bool("engine_inputs_shm" in ipc_payload),
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"[Orchestrator] IPC encode failed for req %s: %s; \
|
||||
falling back to inline payload",
|
||||
req_id,
|
||||
e,
|
||||
)
|
||||
self.stage_list[next_stage_id].submit(
|
||||
{
|
||||
"request_id": req_id,
|
||||
"engine_inputs": next_inputs,
|
||||
"sampling_params": sp_next,
|
||||
}
|
||||
|
||||
if not sent_via_connector:
|
||||
# Fallback logic removed as we now enforce connector usage.
|
||||
# If no connector is found or send fails, we log an error and raise,
|
||||
# because continuing would cause the request to be silently dropped
|
||||
# and the orchestrator to hang waiting for completion.
|
||||
error_msg = (
|
||||
f"[Orchestrator] Failed to send request {req_id} to stage-{next_stage_id} via connector. "
|
||||
"Configure a connector for this edge or inspect connector logs for details."
|
||||
)
|
||||
logger.error(error_msg)
|
||||
raise RuntimeError(error_msg)
|
||||
logger.debug(
|
||||
"[Orchestrator] Forwarded request %s to stage-%s",
|
||||
req_id,
|
||||
|
||||
@@ -95,6 +95,19 @@ class OmniServeCommand(CLISubcommand):
|
||||
default=10,
|
||||
help="The timeout for the batch.",
|
||||
)
|
||||
serve_parser.add_argument(
|
||||
"--worker-backend",
|
||||
type=str,
|
||||
default="multi_process",
|
||||
choices=["multi_process", "ray"],
|
||||
help="The backend to use for stage workers.",
|
||||
)
|
||||
serve_parser.add_argument(
|
||||
"--ray-address",
|
||||
type=str,
|
||||
default=None,
|
||||
help="The address of the Ray cluster to connect to.",
|
||||
)
|
||||
return serve_parser
|
||||
|
||||
|
||||
|
||||
@@ -301,7 +301,7 @@ def compute_and_log_stage_request_stats(
|
||||
|
||||
|
||||
def record_stage_metrics(
|
||||
per_request: dict[int, dict[str, Any]],
|
||||
per_request: dict[str, dict[str, Any]],
|
||||
stage_req_counts: list[int],
|
||||
stage_total_time_ms: list[float],
|
||||
stage_total_tokens: list[int],
|
||||
@@ -312,8 +312,8 @@ def record_stage_metrics(
|
||||
try:
|
||||
stage_req_counts[stage_id] += 1
|
||||
stage_total_tokens[stage_id] += int(metrics.get("num_tokens_out", 0))
|
||||
rid_int = int(req_id)
|
||||
pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
|
||||
rid_key = str(req_id)
|
||||
pr = per_request.setdefault(rid_key, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
|
||||
pr_stages = pr["stages"] # type: ignore[index]
|
||||
pr_stages[stage_id] = {
|
||||
"stage_gen_time_ms": float(metrics.get("stage_gen_time_ms", 0.0)),
|
||||
@@ -324,9 +324,9 @@ def record_stage_metrics(
|
||||
|
||||
|
||||
def aggregate_rx_and_maybe_total(
|
||||
transfer_edge_req: dict[tuple[int, int, int], dict[str, float]],
|
||||
transfer_edge_req: dict[tuple[int, int, str], dict[str, float]],
|
||||
transfer_agg: dict[tuple[int, int], dict[str, float]],
|
||||
per_request: dict[int, dict[str, Any]],
|
||||
per_request: dict[str, dict[str, Any]],
|
||||
stage_id: int,
|
||||
req_id: Any,
|
||||
rx_bytes: float,
|
||||
@@ -355,11 +355,8 @@ def aggregate_rx_and_maybe_total(
|
||||
agg["rx_count"] += 1.0
|
||||
|
||||
# Try combine with sender-side timing if present
|
||||
try:
|
||||
rid_int = int(req_id)
|
||||
except Exception:
|
||||
rid_int = req_id # best effort
|
||||
s = transfer_edge_req.get((stage_id - 1, stage_id, rid_int))
|
||||
rid_key = str(req_id)
|
||||
s = transfer_edge_req.get((stage_id - 1, stage_id, rid_key))
|
||||
if s is None:
|
||||
return None
|
||||
tx_ms = float(s.get("tx_ms", 0.0))
|
||||
@@ -369,7 +366,7 @@ def aggregate_rx_and_maybe_total(
|
||||
agg["total_count"] += 1.0
|
||||
# accumulate per-request transfer totals
|
||||
try:
|
||||
pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
|
||||
pr = per_request.setdefault(rid_key, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
|
||||
pr["transfers_ms"] = float(pr.get("transfers_ms", 0.0)) + total_ms # type: ignore[index]
|
||||
pr["transfers_bytes"] = int(pr.get("transfers_bytes", 0)) + int(rx_bytes) # type: ignore[index]
|
||||
except Exception:
|
||||
@@ -382,7 +379,7 @@ def aggregate_rx_and_maybe_total(
|
||||
|
||||
def record_sender_transfer_agg(
|
||||
transfer_agg: dict[tuple[int, int], dict[str, float]],
|
||||
transfer_edge_req: dict[tuple[int, int, int], dict[str, float]],
|
||||
transfer_edge_req: dict[tuple[int, int, str], dict[str, float]],
|
||||
from_stage: int,
|
||||
to_stage: int,
|
||||
req_id: Any,
|
||||
@@ -408,11 +405,8 @@ def record_sender_transfer_agg(
|
||||
agg["sum_ms"] += float(tx_ms)
|
||||
agg["count"] += 1.0
|
||||
# Store sender-side timing for per-request combination
|
||||
try:
|
||||
rid_int = int(req_id)
|
||||
except Exception:
|
||||
rid_int = req_id
|
||||
transfer_edge_req[(from_stage, to_stage, rid_int)] = {
|
||||
rid_key = str(req_id)
|
||||
transfer_edge_req[(from_stage, to_stage, rid_key)] = {
|
||||
"tx_ms": float(tx_ms),
|
||||
"size_bytes": float(size_bytes),
|
||||
}
|
||||
@@ -512,12 +506,12 @@ class OrchestratorMetrics:
|
||||
self.stage_total_tokens: list[int] = [0 for _ in range(self.num_stages)]
|
||||
self.stage_req_counts: list[int] = [0 for _ in range(self.num_stages)]
|
||||
self.transfer_agg: dict[tuple[int, int], dict[str, float]] = {}
|
||||
self.transfer_edge_req: dict[tuple[int, int, int], dict[str, float]] = {}
|
||||
self.transfer_edge_req: dict[tuple[int, int, str], dict[str, float]] = {}
|
||||
self.e2e_total_ms: float = 0.0
|
||||
self.e2e_total_tokens: int = 0
|
||||
self.e2e_count: int = 0
|
||||
self.e2e_done: set[int] = set()
|
||||
self.per_request: dict[int, dict[str, Any]] = {}
|
||||
self.e2e_done: set[str] = set()
|
||||
self.per_request: dict[str, dict[str, Any]] = {}
|
||||
self.sum_per_request_transfer_ms: float = 0.0
|
||||
self.wall_start_ts: float = float(wall_start_ts)
|
||||
self.last_finish_ts: float = float(wall_start_ts)
|
||||
@@ -614,10 +608,7 @@ class OrchestratorMetrics:
|
||||
)
|
||||
|
||||
def on_finalize_request(self, stage_id: int, req_id: Any, engine_outputs: list[Any], req_start_ts: float) -> None:
|
||||
try:
|
||||
rid_int = int(req_id)
|
||||
except Exception:
|
||||
rid_int = req_id
|
||||
rid_key = str(req_id)
|
||||
_t0 = float(req_start_ts)
|
||||
_t1 = time.time()
|
||||
# Update last output time for this stage
|
||||
@@ -629,11 +620,11 @@ class OrchestratorMetrics:
|
||||
self.e2e_total_ms += e2e_ms
|
||||
self.e2e_total_tokens += int(num_tokens)
|
||||
self.e2e_count += 1
|
||||
self.e2e_done.add(rid_int)
|
||||
pr = self.per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
|
||||
self.e2e_done.add(rid_key)
|
||||
pr = self.per_request.setdefault(rid_key, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
|
||||
per_req_record = {
|
||||
"type": "overall_request",
|
||||
"request_id": rid_int,
|
||||
"request_id": rid_key,
|
||||
"e2e_time_ms": e2e_ms,
|
||||
"num_tokens_out": int(num_tokens),
|
||||
"transfers_total_time_ms": float(pr.get("transfers_ms", 0.0)),
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Sequence
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from typing import Any, Optional, Union
|
||||
@@ -19,7 +20,18 @@ from vllm.usage.usage_lib import UsageContext
|
||||
from vllm.utils import Counter
|
||||
from vllm.v1.engine.llm_engine import LLMEngine
|
||||
|
||||
from vllm_omni.distributed.omni_connectors import (
|
||||
get_stage_connector_config,
|
||||
initialize_orchestrator_connectors,
|
||||
)
|
||||
|
||||
# Internal imports (our code)
|
||||
from vllm_omni.distributed.omni_connectors.adapter import try_send_via_connector
|
||||
from vllm_omni.distributed.ray_utils.utils import (
|
||||
create_placement_group,
|
||||
get_ray_queue_class,
|
||||
try_close_ray,
|
||||
)
|
||||
from vllm_omni.engine.arg_utils import OmniEngineArgs
|
||||
from vllm_omni.engine.output_processor import MultimodalOutputProcessor
|
||||
from vllm_omni.engine.processor import OmniProcessor
|
||||
@@ -30,10 +42,12 @@ from vllm_omni.entrypoints.log_utils import (
|
||||
remove_old_logs,
|
||||
)
|
||||
from vllm_omni.entrypoints.omni_stage import OmniStage
|
||||
from vllm_omni.entrypoints.stage_utils import encode_for_ipc as _encode
|
||||
from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
|
||||
from vllm_omni.entrypoints.stage_utils import serialize_obj as _set
|
||||
from vllm_omni.entrypoints.utils import load_stage_configs_from_model, load_stage_configs_from_yaml
|
||||
from vllm_omni.entrypoints.utils import (
|
||||
load_stage_configs_from_model,
|
||||
load_stage_configs_from_yaml,
|
||||
resolve_model_config_path,
|
||||
)
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
logger = init_logger(__name__)
|
||||
@@ -84,14 +98,25 @@ class OmniLLM:
|
||||
init_timeout: int = 300,
|
||||
**kwargs: Any,
|
||||
):
|
||||
self.worker_backend = kwargs.get("worker_backend", "multi_process")
|
||||
self.ray_address = kwargs.get("ray_address", None)
|
||||
self._ray_pg = None
|
||||
self.batch_timeout = batch_timeout
|
||||
self._enable_stats: bool = bool(log_stats)
|
||||
|
||||
# Do NOT call super().__init__ to avoid creating OmniStageLLM instances in parent.
|
||||
if stage_configs_path is None:
|
||||
self.config_path = resolve_model_config_path(model)
|
||||
self.stage_configs = load_stage_configs_from_model(model)
|
||||
else:
|
||||
self.config_path = stage_configs_path
|
||||
self.stage_configs = load_stage_configs_from_yaml(stage_configs_path)
|
||||
|
||||
# Initialize connectors
|
||||
self.omni_transfer_config, self.connectors = initialize_orchestrator_connectors(
|
||||
self.config_path, worker_backend=self.worker_backend, shm_threshold_bytes=shm_threshold_bytes
|
||||
)
|
||||
|
||||
# Optional file handler for orchestrator
|
||||
self._log_file = log_file
|
||||
|
||||
@@ -124,33 +149,51 @@ class OmniLLM:
|
||||
self.stage_list = [st for _, st in results]
|
||||
logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list))
|
||||
|
||||
self._ctx = mp.get_context("spawn")
|
||||
if self.worker_backend == "ray":
|
||||
self._queue_cls = get_ray_queue_class()
|
||||
else:
|
||||
self._ctx = mp.get_context("spawn")
|
||||
self._queue_cls = lambda: self._ctx.Queue(maxsize=0)
|
||||
|
||||
self._stage_in_queues: list[mp.Queue] = []
|
||||
self._stage_out_queues: list[mp.Queue] = []
|
||||
self._init_sleep_seconds = max(0, int(init_sleep_seconds))
|
||||
self._shm_threshold_bytes = max(0, int(shm_threshold_bytes))
|
||||
self._start_stage_processes(model)
|
||||
self._start_stages(model)
|
||||
# Wait for all stages to report readiness before seeding
|
||||
self._stages_ready: set[int] = set()
|
||||
self._wait_for_stages_ready(timeout=init_timeout)
|
||||
|
||||
def _start_stage_processes(self, model: str) -> None:
|
||||
def _start_stages(self, model: str) -> None:
|
||||
if self.worker_backend == "ray":
|
||||
# Initialize Ray Cluster
|
||||
self._ray_pg = create_placement_group(
|
||||
number_of_stages=len(self.stage_list), address=self.ray_address, strategy="PACK"
|
||||
)
|
||||
|
||||
for stage_id, stage in enumerate(self.stage_list):
|
||||
# Use unbounded queues to avoid deadlock when seeding many requests
|
||||
in_q: mp.Queue = self._ctx.Queue(maxsize=0)
|
||||
out_q: mp.Queue = self._ctx.Queue(maxsize=0)
|
||||
in_q = self._queue_cls()
|
||||
out_q = self._queue_cls()
|
||||
self._stage_in_queues.append(in_q)
|
||||
self._stage_out_queues.append(out_q)
|
||||
|
||||
# Attach queues and start Stage-owned worker process
|
||||
stage.attach_queues(in_q, out_q)
|
||||
|
||||
stage_connectors_config = get_stage_connector_config(
|
||||
self.omni_transfer_config,
|
||||
stage_id,
|
||||
)
|
||||
|
||||
stage.init_stage_worker(
|
||||
model,
|
||||
log_file=self._log_file,
|
||||
shm_threshold_bytes=self._shm_threshold_bytes,
|
||||
ctx=self._ctx,
|
||||
ctx=self._ctx if self.worker_backend != "ray" else None,
|
||||
batch_timeout=self.batch_timeout,
|
||||
connectors_config=stage_connectors_config,
|
||||
worker_backend=self.worker_backend,
|
||||
ray_placement_group=self._ray_pg,
|
||||
)
|
||||
|
||||
logger.debug("[Orchestrator] Stage-%s process started", stage_id)
|
||||
time.sleep(self._init_sleep_seconds)
|
||||
|
||||
@@ -175,6 +218,8 @@ class OmniLLM:
|
||||
except Exception as e:
|
||||
logger.warning("[Orchestrator] Failed to stop stage worker: %s", e)
|
||||
|
||||
try_close_ray(self._ray_pg)
|
||||
|
||||
def __del__(self) -> None: # best-effort
|
||||
try:
|
||||
self.close()
|
||||
@@ -237,28 +282,31 @@ class OmniLLM:
|
||||
# Orchestrator keeps stage objects for input derivation
|
||||
num_stages = len(self.stage_list)
|
||||
|
||||
# Map from request_id to original prompt
|
||||
request_id_to_prompt: dict[int, PromptType] = {i: p for i, p in enumerate(request_prompts)}
|
||||
# Generate globally unique request IDs and map them to original prompts
|
||||
request_ids: list[str] = [f"{i}_{uuid.uuid4()}" for i in range(len(request_prompts))]
|
||||
request_id_to_prompt: dict[str, PromptType] = {rid: p for rid, p in zip(request_ids, request_prompts)}
|
||||
|
||||
# Track per-request start time for end-to-end timing
|
||||
_req_start_ts: dict[int, float] = {}
|
||||
_req_start_ts: dict[str, float] = {}
|
||||
_wall_start_ts: float = time.time()
|
||||
|
||||
# Determine the final stage for E2E stats (highest stage_id with final_output=True; fallback to last stage)
|
||||
final_stage_id_for_e2e = -1
|
||||
last_stage_id = len(self.stage_list) - 1
|
||||
try:
|
||||
for _sid, _st in enumerate(self.stage_list):
|
||||
if getattr(_st, "final_output", False):
|
||||
final_stage_id_for_e2e = max(final_stage_id_for_e2e, _sid)
|
||||
for _sid in range(last_stage_id, -1, -1):
|
||||
if getattr(self.stage_list[_sid], "final_output", False):
|
||||
final_stage_id_for_e2e = _sid
|
||||
break
|
||||
if final_stage_id_for_e2e < 0:
|
||||
final_stage_id_for_e2e = len(self.stage_list) - 1
|
||||
final_stage_id_for_e2e = last_stage_id
|
||||
except Exception as e:
|
||||
logger.debug(
|
||||
"[Orchestrator] Failed to determine final stage for E2E; falling back to last: %s",
|
||||
e,
|
||||
exc_info=True,
|
||||
)
|
||||
final_stage_id_for_e2e = len(self.stage_list) - 1
|
||||
final_stage_id_for_e2e = last_stage_id
|
||||
# Metrics/aggregation helper
|
||||
metrics = OrchestratorMetrics(
|
||||
num_stages,
|
||||
@@ -357,8 +405,8 @@ class OmniLLM:
|
||||
# End-to-end timing and time-per-token for final output
|
||||
# (only once per request at the designated final stage)
|
||||
try:
|
||||
rid_int = int(req_id) if isinstance(req_id, (int, str)) and str(req_id).isdigit() else req_id
|
||||
if stage_id == final_stage_id_for_e2e and rid_int not in metrics.e2e_done:
|
||||
rid_key = str(req_id)
|
||||
if stage_id == final_stage_id_for_e2e and rid_key not in metrics.e2e_done:
|
||||
metrics.on_finalize_request(
|
||||
stage_id,
|
||||
req_id,
|
||||
@@ -387,50 +435,28 @@ class OmniLLM:
|
||||
)
|
||||
continue
|
||||
sp_next: SamplingParams = sampling_params_list[next_stage_id] # type: ignore[index]
|
||||
try:
|
||||
# Measure transfer size and time (encode + enqueue)
|
||||
size_bytes = 0
|
||||
try:
|
||||
size_bytes = len(_set(next_inputs))
|
||||
except Exception:
|
||||
size_bytes = 0
|
||||
t0 = time.time()
|
||||
ipc_payload = _encode(
|
||||
next_inputs,
|
||||
getattr(self, "_shm_threshold_bytes", 65536),
|
||||
obj_key="engine_inputs",
|
||||
shm_key="engine_inputs_shm",
|
||||
|
||||
# Check if we have a connector for this edge
|
||||
connector_key = (str(stage_id), str(next_stage_id))
|
||||
connector = self.connectors.get(connector_key)
|
||||
sent_via_connector = False
|
||||
if connector:
|
||||
sent_via_connector = try_send_via_connector(
|
||||
connector=connector,
|
||||
stage_id=stage_id,
|
||||
next_stage_id=next_stage_id,
|
||||
req_id=req_id,
|
||||
next_inputs=next_inputs,
|
||||
sampling_params=sp_next,
|
||||
original_prompt=request_id_to_prompt[req_id],
|
||||
next_stage_queue_submit_fn=self.stage_list[next_stage_id].submit,
|
||||
metrics=metrics,
|
||||
)
|
||||
ipc_payload.update(
|
||||
{
|
||||
"request_id": req_id,
|
||||
"sampling_params": sp_next,
|
||||
"sent_ts": time.time(),
|
||||
}
|
||||
)
|
||||
self.stage_list[next_stage_id].submit(ipc_payload)
|
||||
t1 = time.time()
|
||||
tx_ms = (t1 - t0) * 1000.0
|
||||
metrics.on_forward(
|
||||
stage_id,
|
||||
next_stage_id,
|
||||
req_id,
|
||||
int(size_bytes),
|
||||
float(tx_ms),
|
||||
bool("engine_inputs_shm" in ipc_payload),
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"[Orchestrator] IPC encode failed for req %s: %s; falling back to inline payload",
|
||||
req_id,
|
||||
e,
|
||||
)
|
||||
self.stage_list[next_stage_id].submit(
|
||||
{
|
||||
"request_id": req_id,
|
||||
"engine_inputs": next_inputs,
|
||||
"sampling_params": sp_next,
|
||||
}
|
||||
|
||||
if not sent_via_connector:
|
||||
raise RuntimeError(
|
||||
f"[Orchestrator] Failed to send request {req_id} to stage-{next_stage_id} via connector. "
|
||||
"Configure a connector for this edge or inspect connector logs for details."
|
||||
)
|
||||
logger.debug(
|
||||
"[Orchestrator] Forwarded request %s to stage-%s",
|
||||
|
||||
@@ -15,6 +15,7 @@ import importlib
|
||||
import logging
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import sys
|
||||
from typing import Any, Optional, Union
|
||||
|
||||
from vllm.inputs import TextPrompt
|
||||
@@ -27,11 +28,11 @@ from vllm.v1.engine import EngineCoreOutput
|
||||
from vllm.v1.engine.async_llm import AsyncLLM
|
||||
from vllm.v1.engine.llm_engine import LLMEngine
|
||||
|
||||
from vllm_omni.distributed.ray_utils.utils import kill_ray_actor, start_ray_actor
|
||||
from vllm_omni.engine.arg_utils import AsyncOmniEngineArgs
|
||||
from vllm_omni.entrypoints.stage_utils import (
|
||||
_to_dict,
|
||||
maybe_dump_to_shm,
|
||||
maybe_load_from_ipc_with_metrics,
|
||||
set_stage_devices,
|
||||
)
|
||||
from vllm_omni.inputs.data import OmniTokensPrompt
|
||||
@@ -163,6 +164,9 @@ class OmniStage:
|
||||
shm_threshold_bytes: int = 65536,
|
||||
ctx: Optional[mp.context.BaseContext] = None,
|
||||
batch_timeout: int = 10,
|
||||
connectors_config: Optional[dict] = None,
|
||||
worker_backend: str = "multi_process",
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initialize and start the stage worker process.
|
||||
|
||||
@@ -176,13 +180,23 @@ class OmniStage:
|
||||
shm_threshold_bytes: Threshold for using shared memory for IPC
|
||||
ctx: Optional multiprocessing context (default: spawn)
|
||||
batch_timeout: Timeout in seconds for batching requests
|
||||
connectors_config: Configuration for stage connectors
|
||||
worker_backend: Backend type ("multi_process" or "ray")
|
||||
**kwargs: Additional arguments (e.g. ray_placement_group)
|
||||
|
||||
Raises:
|
||||
AssertionError: If queues are not attached before calling this method
|
||||
"""
|
||||
assert self._in_q is not None and self._out_q is not None, "Queues must be attached before start_process"
|
||||
self._log_file = log_file
|
||||
self._shm_threshold_bytes = shm_threshold_bytes
|
||||
|
||||
if worker_backend == "ray":
|
||||
ray_placement_group = kwargs.get("ray_placement_group", None)
|
||||
assert ray_placement_group is not None, "Ray placement group must be provided"
|
||||
self._shm_threshold_bytes = sys.maxsize
|
||||
else:
|
||||
self._shm_threshold_bytes = shm_threshold_bytes
|
||||
|
||||
ctx = ctx or mp.get_context("spawn")
|
||||
# Prepare lightweight dict config for worker
|
||||
engine_args = _to_dict(self.engine_args)
|
||||
@@ -192,43 +206,74 @@ class OmniStage:
|
||||
"engine_args": engine_args,
|
||||
"runtime": runtime_cfg,
|
||||
"shm_threshold_bytes": self._shm_threshold_bytes,
|
||||
"connectors_config": connectors_config or {},
|
||||
}
|
||||
if is_async:
|
||||
self._proc = ctx.Process(
|
||||
target=_stage_worker_async_entry,
|
||||
args=(
|
||||
|
||||
if worker_backend == "ray":
|
||||
if is_async:
|
||||
self._ray_actor = start_ray_actor(
|
||||
_stage_worker_async_entry,
|
||||
ray_placement_group,
|
||||
self.stage_id,
|
||||
self,
|
||||
model,
|
||||
stage_payload,
|
||||
batch_timeout,
|
||||
),
|
||||
)
|
||||
model=model,
|
||||
stage_payload=stage_payload,
|
||||
batch_timeout=batch_timeout,
|
||||
)
|
||||
else:
|
||||
self._ray_actor = start_ray_actor(
|
||||
_stage_worker,
|
||||
ray_placement_group,
|
||||
self.stage_id,
|
||||
model=model,
|
||||
stage_payload=stage_payload,
|
||||
in_q=self._in_q,
|
||||
out_q=self._out_q,
|
||||
log_file=self._log_file,
|
||||
batch_timeout=batch_timeout,
|
||||
)
|
||||
else:
|
||||
self._proc = ctx.Process(
|
||||
target=_stage_worker,
|
||||
args=(
|
||||
model,
|
||||
stage_payload,
|
||||
self._in_q,
|
||||
self._out_q,
|
||||
self._log_file,
|
||||
batch_timeout,
|
||||
),
|
||||
)
|
||||
self._proc.start()
|
||||
if is_async:
|
||||
self._proc = ctx.Process(
|
||||
target=_stage_worker_async_entry,
|
||||
args=(
|
||||
self,
|
||||
model,
|
||||
stage_payload,
|
||||
batch_timeout,
|
||||
),
|
||||
)
|
||||
else:
|
||||
self._proc = ctx.Process(
|
||||
target=_stage_worker,
|
||||
args=(
|
||||
model,
|
||||
stage_payload,
|
||||
self._in_q,
|
||||
self._out_q,
|
||||
self._log_file,
|
||||
batch_timeout,
|
||||
),
|
||||
)
|
||||
self._proc.start()
|
||||
|
||||
def stop_stage_worker(self) -> None:
|
||||
"""Stop the stage worker process gracefully.
|
||||
|
||||
Sends shutdown signal to the worker and waits for it to terminate.
|
||||
If graceful shutdown fails, forcefully terminates the process.
|
||||
Handles both multiprocessing Process and Ray Actor.
|
||||
"""
|
||||
if self._in_q is not None:
|
||||
try:
|
||||
self._in_q.put_nowait(None)
|
||||
except Exception as e:
|
||||
self._logger.warning("[Stage-%s] Failed to send shutdown to in_q: %s", self.stage_id, e)
|
||||
if self._proc is not None:
|
||||
|
||||
if hasattr(self, "_ray_actor") and self._ray_actor:
|
||||
kill_ray_actor(self._ray_actor)
|
||||
self._ray_actor = None
|
||||
elif self._proc is not None:
|
||||
try:
|
||||
self._proc.join(timeout=5)
|
||||
except Exception as e:
|
||||
@@ -238,6 +283,7 @@ class OmniStage:
|
||||
self._proc.terminate()
|
||||
except Exception as e:
|
||||
self._logger.warning("[Stage-%s] terminate() failed: %s", self.stage_id, e)
|
||||
|
||||
# Cleanup temporary stage log if we created one (only when no log_file provided)
|
||||
try:
|
||||
if not self._log_file:
|
||||
@@ -334,6 +380,8 @@ def _stage_worker(
|
||||
import os as _os
|
||||
import time as _time
|
||||
|
||||
from vllm_omni.distributed.omni_connectors import build_stage_connectors
|
||||
from vllm_omni.distributed.omni_connectors.adapter import try_recv_via_connector
|
||||
from vllm_omni.entrypoints.log_utils import (
|
||||
compute_and_log_stage_request_stats,
|
||||
count_tokens_from_outputs,
|
||||
@@ -347,6 +395,7 @@ def _stage_worker(
|
||||
engine_args = stage_payload.get("engine_args", {})
|
||||
runtime_cfg = stage_payload.get("runtime", {})
|
||||
shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536))
|
||||
connectors_config = stage_payload.get("connectors_config", {})
|
||||
|
||||
# Per-stage logger: clear inherited handlers to avoid broken parent streams
|
||||
try:
|
||||
@@ -440,6 +489,18 @@ def _stage_worker(
|
||||
)
|
||||
stage_engine = OmniStageLLM(model=model, **engine_args)
|
||||
_logging.getLogger(__name__).debug("[Stage-%s] Engine initialized", stage_id)
|
||||
|
||||
# Initialize OmniConnectors if configured
|
||||
connectors = {}
|
||||
if connectors_config:
|
||||
built_connectors = build_stage_connectors(
|
||||
stage_id=stage_id,
|
||||
connectors_config=connectors_config,
|
||||
)
|
||||
if built_connectors is None:
|
||||
return
|
||||
connectors = built_connectors
|
||||
|
||||
# Signal readiness to orchestrator
|
||||
try:
|
||||
out_q.put({"type": "stage_ready", "stage_id": stage_id})
|
||||
@@ -496,9 +557,25 @@ def _stage_worker(
|
||||
_in_flight_ms_by_rid[rid] = 0.0
|
||||
except Exception:
|
||||
_in_flight_ms_by_rid[rid] = 0.0
|
||||
ein, _rx_metrics = maybe_load_from_ipc_with_metrics(t, obj_key="engine_inputs", shm_key="engine_inputs_shm")
|
||||
_rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0))
|
||||
_rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
|
||||
|
||||
# Resolve input data strictly via connectors if payload
|
||||
# is larger than shm_threshold_bytes or using other connectors
|
||||
ein, _rx_metrics = try_recv_via_connector(
|
||||
task=t,
|
||||
connectors=connectors,
|
||||
stage_id=stage_id,
|
||||
)
|
||||
|
||||
if ein is None or _rx_metrics is None:
|
||||
raise RuntimeError(
|
||||
f"[Stage-{stage_id}] Missing connector payload for request {rid}. "
|
||||
"Ensure connectors are configured for all incoming edges."
|
||||
)
|
||||
|
||||
if _rx_metrics:
|
||||
_rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0))
|
||||
_rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
|
||||
|
||||
batch_request_ids.append(rid)
|
||||
if isinstance(ein, list):
|
||||
batch_engine_inputs.extend(ein)
|
||||
@@ -671,6 +748,8 @@ async def _stage_worker_async(
|
||||
import logging as _logging
|
||||
import time as _time
|
||||
|
||||
from vllm_omni.distributed.omni_connectors import build_stage_connectors
|
||||
from vllm_omni.distributed.omni_connectors.adapter import try_recv_via_connector
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmniStageLLM
|
||||
from vllm_omni.entrypoints.log_utils import (
|
||||
compute_and_log_stage_request_stats,
|
||||
@@ -685,6 +764,7 @@ async def _stage_worker_async(
|
||||
engine_args = stage_payload.get("engine_args", {})
|
||||
runtime_cfg = stage_payload.get("runtime", {})
|
||||
shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536))
|
||||
connectors_config = stage_payload.get("connectors_config", {})
|
||||
|
||||
log_file = omni_stage._log_file
|
||||
in_q = omni_stage._in_q
|
||||
@@ -729,6 +809,17 @@ async def _stage_worker_async(
|
||||
except Exception as e:
|
||||
_logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e)
|
||||
|
||||
# Initialize OmniConnectors if configured to match sync worker behavior
|
||||
connectors: dict[Any, Any] = {}
|
||||
if connectors_config:
|
||||
built_connectors = build_stage_connectors(
|
||||
stage_id=stage_id,
|
||||
connectors_config=connectors_config,
|
||||
)
|
||||
if built_connectors is None:
|
||||
return
|
||||
connectors = built_connectors
|
||||
|
||||
# Init LLM
|
||||
_logging.getLogger(__name__).debug(
|
||||
"[Stage-%s] Initializing engine with args keys=%s",
|
||||
@@ -788,7 +879,16 @@ async def _stage_worker_async(
|
||||
_in_flight_ms_by_rid[rid] = 0.0
|
||||
except Exception:
|
||||
_in_flight_ms_by_rid[rid] = 0.0
|
||||
ein, _rx_metrics = maybe_load_from_ipc_with_metrics(task, obj_key="engine_inputs", shm_key="engine_inputs_shm")
|
||||
ein, _rx_metrics = try_recv_via_connector(
|
||||
task=task,
|
||||
connectors=connectors,
|
||||
stage_id=stage_id,
|
||||
)
|
||||
if ein is None or _rx_metrics is None:
|
||||
raise RuntimeError(
|
||||
f"[Stage-{stage_id}] Missing connector payload for request {rid}. "
|
||||
"Ensure connectors are configured for all incoming edges."
|
||||
)
|
||||
_rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0))
|
||||
_rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
|
||||
|
||||
|
||||
@@ -3,11 +3,9 @@ from __future__ import annotations
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import pickle
|
||||
from multiprocessing import shared_memory as _shm
|
||||
from typing import Any
|
||||
|
||||
import cloudpickle
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -27,8 +25,11 @@ def set_stage_devices(
|
||||
Args:
|
||||
stage_id: Stage identifier for logging
|
||||
devices: Device specification:
|
||||
- Comma-separated string (e.g. "2,5,7"): set device visibility env var
|
||||
exactly to this list; logical index 0 is used as current device.
|
||||
- Comma-separated string (e.g. "2,5,7"): interpreted as logical
|
||||
indices against the current device visibility env var (e.g.
|
||||
CUDA_VISIBLE_DEVICES/ASCEND_RT_VISIBLE_DEVICES) when present;
|
||||
falls back to physical IDs if no mapping exists. Logical index 0
|
||||
is used as current device.
|
||||
- Integer or digit-string: treat as logical index (0-based) into the
|
||||
current device visibility mapping; map to physical device, then set
|
||||
env var to this single device.
|
||||
@@ -63,7 +64,7 @@ def set_stage_devices(
|
||||
get_device_name_fn = torch.npu.get_device_name
|
||||
device_type_label = "NPU"
|
||||
elif device_type == "cuda":
|
||||
import torch # noqa: WPS433
|
||||
import torch
|
||||
|
||||
is_available_fn = torch.cuda.is_available
|
||||
set_device_fn = torch.cuda.set_device
|
||||
@@ -81,16 +82,35 @@ def set_stage_devices(
|
||||
logical_idx: int | None = None
|
||||
|
||||
if isinstance(devices, str) and "," in devices:
|
||||
os.environ[env_var] = devices
|
||||
toks = [t.strip() for t in devices.split(",") if t.strip() != ""]
|
||||
vis = os.environ.get(env_var)
|
||||
mapped_devices: list[str] = []
|
||||
mapping: list[int] = []
|
||||
if vis:
|
||||
try:
|
||||
mapping = [int(x) for x in vis.split(",") if x.strip() != ""]
|
||||
except Exception as e:
|
||||
logger.debug("[Stage-%s] Failed to parse existing %s: %s", stage_id, env_var, e)
|
||||
for tok in toks:
|
||||
try:
|
||||
idx = int(tok)
|
||||
except Exception:
|
||||
mapped_devices.append(tok)
|
||||
continue
|
||||
if mapping and 0 <= idx < len(mapping):
|
||||
mapped_devices.append(str(mapping[idx]))
|
||||
else:
|
||||
mapped_devices.append(str(idx))
|
||||
mapped_devices_str = ",".join(mapped_devices)
|
||||
os.environ[env_var] = mapped_devices_str
|
||||
if toks:
|
||||
try:
|
||||
selected_physical = int(toks[0])
|
||||
selected_physical = int(mapped_devices[0])
|
||||
logger.debug(
|
||||
"[Stage-%s] Set %s to %s; logical 0 -> physical %s",
|
||||
stage_id,
|
||||
env_var,
|
||||
devices,
|
||||
mapped_devices_str,
|
||||
selected_physical,
|
||||
)
|
||||
except Exception as e:
|
||||
@@ -125,7 +145,7 @@ def set_stage_devices(
|
||||
logger.debug("[Stage-%s] Set %s to single device %s (fallback)", stage_id, env_var, selected_physical)
|
||||
|
||||
try:
|
||||
import torch # noqa: WPS433
|
||||
import torch
|
||||
|
||||
if is_available_fn():
|
||||
try:
|
||||
@@ -155,8 +175,10 @@ def set_stage_devices(
|
||||
|
||||
|
||||
def serialize_obj(obj: Any) -> bytes:
|
||||
"""Serialize a Python object to bytes using cloudpickle."""
|
||||
return cloudpickle.dumps(obj)
|
||||
"""Serialize a Python object to bytes using centralized serializer (defaults to cloudpickle)."""
|
||||
from vllm_omni.distributed.omni_connectors.utils.serialization import OmniSerializer
|
||||
|
||||
return OmniSerializer.serialize(obj)
|
||||
|
||||
|
||||
def shm_write_bytes(payload: bytes) -> dict[str, Any]:
|
||||
@@ -238,7 +260,9 @@ def maybe_load_from_ipc(container: dict[str, Any], obj_key: str, shm_key: str) -
|
||||
decode-time and size metrics.
|
||||
"""
|
||||
if shm_key in container:
|
||||
return pickle.loads(shm_read_bytes(container[shm_key]))
|
||||
from vllm_omni.distributed.omni_connectors.utils.serialization import OmniSerializer
|
||||
|
||||
return OmniSerializer.deserialize(shm_read_bytes(container[shm_key]))
|
||||
return container[obj_key]
|
||||
|
||||
|
||||
@@ -253,11 +277,13 @@ def maybe_load_from_ipc_with_metrics(
|
||||
"""
|
||||
import time as _time # local import to avoid overhead at module import
|
||||
|
||||
from vllm_omni.distributed.omni_connectors.utils.serialization import OmniSerializer
|
||||
|
||||
t0 = _time.time()
|
||||
if shm_key in container:
|
||||
meta = container[shm_key] # type: ignore[index]
|
||||
payload = shm_read_bytes(meta)
|
||||
obj = pickle.loads(payload)
|
||||
obj = OmniSerializer.deserialize(payload)
|
||||
try:
|
||||
rx_bytes = int(meta.get("size", len(payload))) # type: ignore[call-arg]
|
||||
except Exception:
|
||||
|
||||
@@ -2,7 +2,7 @@ import os
|
||||
from collections import Counter
|
||||
from dataclasses import asdict, is_dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from typing import Any, Optional
|
||||
|
||||
from omegaconf import OmegaConf
|
||||
from vllm.transformers_utils.config import get_config
|
||||
@@ -62,7 +62,42 @@ def _convert_dataclasses_to_dict(obj: Any) -> Any:
|
||||
return obj
|
||||
|
||||
|
||||
def load_stage_configs_from_model(model: str, base_engine_args: dict = {}) -> list:
|
||||
def resolve_model_config_path(model: str) -> str:
|
||||
"""Resolve the stage config file path from the model name.
|
||||
|
||||
Resolves stage configuration path based on the model type and device type.
|
||||
First tries to find a device-specific YAML file from stage_configs/{device_type}/
|
||||
directory. If not found, falls back to the default config file.
|
||||
|
||||
Args:
|
||||
model: Model name or path (used to determine model_type)
|
||||
|
||||
Returns:
|
||||
String path to the stage configuration file
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If no stage config file exists for the model type
|
||||
"""
|
||||
hf_config = get_config(model, trust_remote_code=True)
|
||||
model_type = hf_config.model_type
|
||||
device_type = detect_device_type()
|
||||
|
||||
# Try device-specific config first
|
||||
if device_type != "cuda":
|
||||
device_config_file = f"vllm_omni/model_executor/stage_configs/{device_type}/{model_type}.yaml"
|
||||
device_config_path = PROJECT_ROOT / device_config_file
|
||||
if os.path.exists(device_config_path):
|
||||
return str(device_config_path)
|
||||
|
||||
# Fall back to default config
|
||||
stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml"
|
||||
stage_config_path = PROJECT_ROOT / stage_config_file
|
||||
if not os.path.exists(stage_config_path):
|
||||
raise FileNotFoundError(f"Stage config file {stage_config_path} not found")
|
||||
return str(stage_config_path)
|
||||
|
||||
|
||||
def load_stage_configs_from_model(model: str, base_engine_args: Optional[dict] = None) -> list:
|
||||
"""Load stage configurations from model's default config file.
|
||||
|
||||
Loads stage configurations based on the model type and device type.
|
||||
@@ -78,30 +113,14 @@ def load_stage_configs_from_model(model: str, base_engine_args: dict = {}) -> li
|
||||
Raises:
|
||||
FileNotFoundError: If no stage config file exists for the model type
|
||||
"""
|
||||
hf_config = get_config(model, trust_remote_code=True)
|
||||
model_type = hf_config.model_type
|
||||
device_type = detect_device_type()
|
||||
|
||||
# Try device-specific config first
|
||||
if device_type != "cuda":
|
||||
device_config_file = f"vllm_omni/model_executor/stage_configs/{device_type}/{model_type}.yaml"
|
||||
device_config_path = PROJECT_ROOT / device_config_file
|
||||
if os.path.exists(device_config_path):
|
||||
stage_configs = load_stage_configs_from_yaml(
|
||||
config_path=str(device_config_path), base_engine_args=base_engine_args
|
||||
)
|
||||
return stage_configs
|
||||
|
||||
# Fall back to default config
|
||||
stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml"
|
||||
stage_config_path = PROJECT_ROOT / stage_config_file
|
||||
if not os.path.exists(stage_config_path):
|
||||
raise FileNotFoundError(f"Stage config file {stage_config_path} not found")
|
||||
stage_configs = load_stage_configs_from_yaml(config_path=str(stage_config_path), base_engine_args=base_engine_args)
|
||||
if base_engine_args is None:
|
||||
base_engine_args = {}
|
||||
stage_config_path = resolve_model_config_path(model)
|
||||
stage_configs = load_stage_configs_from_yaml(config_path=stage_config_path, base_engine_args=base_engine_args)
|
||||
return stage_configs
|
||||
|
||||
|
||||
def load_stage_configs_from_yaml(config_path: str, base_engine_args: dict = {}) -> list:
|
||||
def load_stage_configs_from_yaml(config_path: str, base_engine_args: Optional[dict] = None) -> list:
|
||||
"""Load stage configurations from a YAML file.
|
||||
|
||||
Args:
|
||||
@@ -110,6 +129,8 @@ def load_stage_configs_from_yaml(config_path: str, base_engine_args: dict = {})
|
||||
Returns:
|
||||
List of stage configuration dictionaries from the file's stage_args
|
||||
"""
|
||||
if base_engine_args is None:
|
||||
base_engine_args = {}
|
||||
config_data = OmegaConf.load(config_path)
|
||||
stage_args = config_data.stage_args
|
||||
# Convert any nested dataclass objects to dicts before creating OmegaConf
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
# stage config for running qwen2.5-omni with architecture of OmniLLM.
|
||||
|
||||
# The following config has been verified on 1x H100-80G GPU.
|
||||
# The following config has been verified on 2x H100-80G GPU.
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
runtime:
|
||||
@@ -29,6 +29,7 @@ stage_args:
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.1
|
||||
|
||||
- stage_id: 1
|
||||
runtime:
|
||||
process: true
|
||||
@@ -56,10 +57,11 @@ stage_args:
|
||||
detokenize: True
|
||||
repetition_penalty: 1.05
|
||||
stop_token_ids: [8294]
|
||||
|
||||
- stage_id: 2
|
||||
runtime:
|
||||
process: true
|
||||
devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU
|
||||
devices: "2" # Example: use a different GPU than the previous stage; use "0" if single GPU
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: code2wav
|
||||
@@ -90,6 +92,7 @@ runtime:
|
||||
defaults:
|
||||
window_size: -1 # Simplified: trigger downstream only after full upstream completion
|
||||
max_inflight: 1 # Simplified: process serially within each stage
|
||||
|
||||
edges:
|
||||
- from: 0 # thinker → talker: trigger only after receiving full input (-1)
|
||||
to: 1
|
||||
|
||||
@@ -0,0 +1,130 @@
|
||||
# stage config for running qwen2.5-omni with architecture of OmniLLM.
|
||||
|
||||
# The following config has been verified on 1x H100-80G GPU.
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
runtime:
|
||||
process: true # Run this stage in a separate process
|
||||
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: thinker
|
||||
model_arch: Qwen2_5OmniForConditionalGeneration
|
||||
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.8
|
||||
enforce_eager: true # Now we only support eager mode
|
||||
trust_remote_code: true
|
||||
engine_output_type: latent
|
||||
enable_prefix_caching: false
|
||||
is_comprehension: true
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 2048
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.1
|
||||
# Distributed connector configuration (optional)
|
||||
output_connectors:
|
||||
to_stage_1: mooncake_connector
|
||||
- stage_id: 1
|
||||
runtime:
|
||||
process: true
|
||||
devices: "1"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: talker
|
||||
model_arch: Qwen2_5OmniForConditionalGeneration
|
||||
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.8
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: latent
|
||||
engine_input_source: [0]
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
|
||||
default_sampling_params:
|
||||
temperature: 0.9
|
||||
top_p: 0.8
|
||||
top_k: 40
|
||||
max_tokens: 2048
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.05
|
||||
stop_token_ids: [8294]
|
||||
# Distributed connector configuration (optional)
|
||||
input_connectors:
|
||||
from_stage_0: mooncake_connector
|
||||
output_connectors:
|
||||
to_stage_2: mooncake_connector
|
||||
- stage_id: 2
|
||||
runtime:
|
||||
process: true
|
||||
devices: "2" # Example: use a different GPU than the previous stage; use "0" if single GPU
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: code2wav
|
||||
model_arch: Qwen2_5OmniForConditionalGeneration
|
||||
worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
gpu_memory_utilization: 0.3
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
max_num_batched_tokens: 32768
|
||||
engine_output_type: audio
|
||||
engine_input_source: [1]
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 2048
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.1
|
||||
# Distributed connector configuration (optional)
|
||||
input_connectors:
|
||||
from_stage_1: mooncake_connector
|
||||
|
||||
# Top-level runtime config (concise): default windows and stage edges
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1 # Simplified: trigger downstream only after full upstream completion
|
||||
max_inflight: 1 # Simplified: process serially within each stage
|
||||
|
||||
# Distributed connectors configuration (optional)
|
||||
# More connectors will be supported in the future.
|
||||
connectors:
|
||||
# Mooncake connector for cross-node/intra-node communication
|
||||
mooncake_connector:
|
||||
name: MooncakeConnector
|
||||
extra:
|
||||
host: "127.0.0.1"
|
||||
metadata_server: "http://10.90.67.86:8080/metadata"
|
||||
master: "10.90.67.86:50051"
|
||||
segment: 512000000 # 512MB
|
||||
localbuf: 64000000 # 64MB
|
||||
proto: "tcp"
|
||||
|
||||
# SharedMemory connector for intra-node communication
|
||||
#Alternative SHM connector with different threshold
|
||||
shared_memory_connector:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
shm_threshold_bytes: 65536 # 64KB threshold
|
||||
|
||||
edges:
|
||||
- from: 0 # thinker → talker: trigger only after receiving full input (-1)
|
||||
to: 1
|
||||
window_size: -1
|
||||
- from: 1 # talker → code2wav: trigger only after receiving full input (-1)
|
||||
to: 2
|
||||
window_size: -1
|
||||
@@ -20,6 +20,7 @@ stage_args:
|
||||
engine_output_type: latent # Output hidden states for talker
|
||||
distributed_executor_backend: "mp"
|
||||
enable_prefix_caching: false
|
||||
max_num_batched_tokens: 32768
|
||||
hf_config_name: thinker_config
|
||||
tensor_parallel_size: 2
|
||||
final_output: true
|
||||
@@ -36,21 +37,22 @@ stage_args:
|
||||
|
||||
- stage_id: 1
|
||||
runtime:
|
||||
devices: "1"
|
||||
max_batch_size: 1
|
||||
devices: "1"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: talker
|
||||
model_arch: Qwen3OmniMoeForConditionalGeneration
|
||||
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.5
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: latent # Output codec codes for code2wav
|
||||
# tensor_parallel_size: 2
|
||||
enable_prefix_caching: false
|
||||
distributed_executor_backend: "mp"
|
||||
hf_config_name: talker_config
|
||||
model_stage: talker
|
||||
model_arch: Qwen3OmniMoeForConditionalGeneration
|
||||
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.3
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: latent # Output codec codes for code2wav
|
||||
# tensor_parallel_size: 2
|
||||
enable_prefix_caching: false
|
||||
max_num_batched_tokens: 32768
|
||||
distributed_executor_backend: "mp"
|
||||
hf_config_name: talker_config
|
||||
engine_input_source: [0]
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
|
||||
# final_output: true
|
||||
|
||||
@@ -0,0 +1,140 @@
|
||||
# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
|
||||
# Stage 0: Thinker (multimodal understanding + text generation)
|
||||
# Stage 1: Talker (text embeddings -> 8-layer RVQ codec codes)
|
||||
# Stage 2: Code2Wav (8-layer RVQ codes -> audio waveform)
|
||||
|
||||
# The following config has been verified on 2x H100-80G GPUs.
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
runtime:
|
||||
devices: "0,1"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: thinker
|
||||
model_arch: Qwen3OmniMoeForConditionalGeneration
|
||||
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
|
||||
scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
|
||||
gpu_memory_utilization: 0.6
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: latent # Output hidden states for talker
|
||||
distributed_executor_backend: "mp"
|
||||
enable_prefix_caching: false
|
||||
hf_config_name: thinker_config
|
||||
tensor_parallel_size: 2
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
is_comprehension: true
|
||||
default_sampling_params:
|
||||
temperature: 0.4
|
||||
top_p: 0.9
|
||||
top_k: 1
|
||||
max_tokens: 2048
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.05
|
||||
# Distributed connector configuration
|
||||
output_connectors:
|
||||
to_stage_1: connector_of_mooncake
|
||||
|
||||
- stage_id: 1
|
||||
runtime:
|
||||
devices: "1"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: talker
|
||||
model_arch: Qwen3OmniMoeForConditionalGeneration
|
||||
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
|
||||
scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
|
||||
gpu_memory_utilization: 0.3
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: latent # Output codec codes for code2wav
|
||||
# tensor_parallel_size: 2
|
||||
enable_prefix_caching: false
|
||||
distributed_executor_backend: "mp"
|
||||
hf_config_name: talker_config
|
||||
engine_input_source: [0]
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
|
||||
# final_output: true
|
||||
# final_output_type: text
|
||||
default_sampling_params:
|
||||
temperature: 0.9
|
||||
top_k: 50
|
||||
max_tokens: 4096
|
||||
seed: 42
|
||||
detokenize: False
|
||||
repetition_penalty: 1.05
|
||||
stop_token_ids: [2150]
|
||||
# Distributed connector configuration
|
||||
input_connectors:
|
||||
from_stage_0: connector_of_mooncake
|
||||
output_connectors:
|
||||
to_stage_2: connector_of_mooncake
|
||||
|
||||
- stage_id: 2
|
||||
runtime:
|
||||
devices: "0"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: code2wav
|
||||
model_arch: Qwen3OmniMoeForConditionalGeneration
|
||||
worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
|
||||
scheduler_cls: vllm_omni.core.sched.generation_scheduler.GenerationScheduler
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: audio # Final output: audio waveform
|
||||
gpu_memory_utilization: 0.1
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 1000000
|
||||
hf_config_name: thinker_config
|
||||
engine_input_source: [1]
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 65536
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.1
|
||||
# Distributed connector configuration
|
||||
input_connectors:
|
||||
from_stage_1: connector_of_mooncake
|
||||
|
||||
# Top-level runtime config: default windows and stage edges
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1
|
||||
max_inflight: 1
|
||||
|
||||
# Distributed connectors configuration
|
||||
connectors:
|
||||
# Mooncake connector for cross-node/intra-node communication
|
||||
connector_of_mooncake:
|
||||
name: MooncakeConnector
|
||||
extra:
|
||||
host: "127.0.0.1"
|
||||
metadata_server: "http://10.90.67.86:8080/metadata"
|
||||
master: "10.90.67.86:50051"
|
||||
segment: 512000000 # 512MB
|
||||
localbuf: 64000000 # 64MB
|
||||
proto: "tcp"
|
||||
|
||||
# SharedMemory connector for intra-node communication
|
||||
connector_of_shared_memory:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
shm_threshold_bytes: 65536 # 64KB threshold
|
||||
|
||||
edges:
|
||||
- from: 0
|
||||
to: 1
|
||||
window_size: -1
|
||||
- from: 1
|
||||
to: 2
|
||||
window_size: -1
|
||||
@@ -47,6 +47,19 @@ class GPUARModelRunner(OmniGPUModelRunner):
|
||||
self.hidden_size = self.model_config.hf_text_config.hidden_size
|
||||
self.inputs_embeds = self._make_buffer(self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False)
|
||||
|
||||
def _make_buffer(self, *size, dtype, numpy=True):
|
||||
# Prevent ray from pinning the buffer due to large size
|
||||
from vllm_omni.distributed.ray_utils.utils import (
|
||||
calculate_total_bytes,
|
||||
maybe_disable_pin_memory_for_ray,
|
||||
)
|
||||
|
||||
total_bytes = calculate_total_bytes(size, dtype)
|
||||
|
||||
# Use the context manager to temporarily disable pinning if needed
|
||||
with maybe_disable_pin_memory_for_ray(self, total_bytes):
|
||||
return super()._make_buffer(*size, dtype=dtype, numpy=numpy)
|
||||
|
||||
@torch.inference_mode()
|
||||
def execute_model(
|
||||
self,
|
||||
|
||||
Reference in new issue
Block a user