Compare commits

..
Author SHA1 Message Date
Kevin H. Luu 61e2c3e458 push
Signed-off-by: Kevin H. Luu <khluu000@gmail.com>
2025-12-08 16:44:19 -08:00
Kevin H. Luu 560ecc0af2 push
Signed-off-by: Kevin H. Luu <khluu000@gmail.com>
2025-12-08 16:41:04 -08:00
Kevin H. Luu c9647d4688 push
Signed-off-by: Kevin H. Luu <khluu000@gmail.com>
2025-12-08 16:39:58 -08:00
Kevin H. Luu a8b5f0b9ae push
Signed-off-by: Kevin H. Luu <khluu000@gmail.com>
2025-12-08 16:35:25 -08:00
Kevin H. Luu 84c0a26bc9 push
Signed-off-by: Kevin H. Luu <khluu000@gmail.com>
2025-12-08 16:04:10 -08:00
Kevin H. Luu 6940d7ce91 push
Signed-off-by: Kevin H. Luu <khluu000@gmail.com>
2025-12-08 15:34:04 -08:00
Kevin H. Luu c0ca921bd7 push
Signed-off-by: Kevin H. Luu <khluu000@gmail.com>
2025-12-08 15:32:47 -08:00
Kevin H. Luu 07282be582 push
Signed-off-by: Kevin H. Luu <khluu000@gmail.com>
2025-12-08 15:29:32 -08:00
Kevin H. Luu 6b7982afb0 push
Signed-off-by: Kevin H. Luu <khluu000@gmail.com>
2025-12-08 15:20:47 -08:00
103 changed files with 949 additions and 6450 deletions

No files matched your search

File renamed without changes.
+10
View File
@@ -0,0 +1,10 @@
name: vllm_omni_ci
github_repo_name: vllm-project/vllm-omni
job_dirs:
- ".buildkite/jobs"
run_all_patterns: []
run_all_exclude_patterns: []
registries: public.ecr.aws/q9t5s3a7
repositories:
main: "vllm-ci-postmerge-repo"
premerge: "vllm-ci-test-repo"
+10
View File
@@ -0,0 +1,10 @@
group: Build
steps:
- label: ":docker: build vllm-omni-ci image"
key: image-build
depends_on: []
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "docker build --file docker/Dockerfile.ci -t vllm-omni-ci ."
- "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
- "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
+30
View File
@@ -0,0 +1,30 @@
group: Tests
depends_on:
- image-build
steps:
- label: "Simple Unit Test"
commands:
- ".buildkite/scripts/simple_test.sh"
no_gpu: true
no_plugin: true
- label: "Diffusion Model Test"
timeout_in_minutes: 15
commands:
- pytest -s -v tests/single_stage/test_diffusion_model.py
- label: "Omni Model Test"
timeout_in_minutes: 15
num_gpus: 4
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/multi_stages/
- label: "Omni Model Test with H100"
timeout_in_minutes: 20
gpu: h100
num_gpus: 2
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/multi_stages_h100/
-87
View File
@@ -1,87 +0,0 @@
steps:
- label: ":docker: Build image"
key: image-build
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "docker build --file docker/Dockerfile.ci -t vllm-omni-ci ."
- "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
- "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
agents:
queue: "cpu_queue_premerge_us_east_1"
- label: "Simple Unit Test"
commands:
- ".buildkite/scripts/simple_test.sh"
agents:
queue: "cpu_queue_premerge"
- label: "Diffusion Model Test"
timeout_in_minutes: 15
depends_on: image-build
commands:
- 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:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Omni Model Test"
timeout_in_minutes: 15
depends_on: image-build
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
agents:
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
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Omni Model Test with H100"
timeout_in_minutes: 20
depends_on: image-build
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- 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:
- kubernetes:
podSpec:
containers:
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
-3
View File
@@ -179,7 +179,6 @@ checkpoints/
# Cache directories
cache/
!vllm_omni/diffusion/cache/
.cache/
diffusion_cache/
kv_cache/
@@ -230,5 +229,3 @@ configs/development.yaml
# Docker
.dockerignore
Dockerfile.dev
discussion
tmp_test
+6 -10
View File
@@ -7,21 +7,17 @@ nav:
- Examples:
- examples/README.md
- Offline Inference:
- 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
- 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
- Online Serving:
- Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
- Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
- 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
- General:
- usage/*
- Configuration:
- configuration/README.md
- configuration/*
- Diffusion Acceleration:
- TeaCache: user_guide/teacache.md
- Models:
- models/supported_models.md
- Developer Guide:
@@ -35,10 +31,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
-1
View File
@@ -96,7 +96,6 @@ Model execution components.
Configuration classes.
- [vllm_omni.config.model.OmniModelConfig][]
- [vllm_omni.diffusion.cache.teacache.config.TeaCacheConfig][]
## Workers
-4
View File
@@ -9,7 +9,3 @@ 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
+75 -157
View File
@@ -1,195 +1,113 @@
# Architecture Overview
This document outlines the architectural design for vLLM-Omni.
# Introduction
<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>
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.
# Goals
The primary goal of the vLLM-Omni project is to build the fastest and easiest-to-use open-source Omni-Modality model inference & serving engine. vLLM-Omni extends the original vLLM, which was created to support large language models for text-based autoregressive (AR) generation tasks. vLLM-Omni is designed to support:
The primary goals for vLLM-Omni are:
* **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.
* **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.
* **Extensibility:** Design a modular and flexible architecture that can easily accommodate new modalities, model architectures, and output formats.
# Expected Gains
# Representative omni-modality models
* 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
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 powerful image generation foundation model capable of complex text rendering and precise image editing.
* **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]
<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.
<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>
# Architectural Overview
**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...).
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]
<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>
## General ideas
# 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>
* 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.
## Key Components
| 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 |
* **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
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.
# Examples data flow of specific models
## Main features
1. Qwen-omni
![][image7]![][image8]
2. Qwen-Image
![][image9] ![][image10]
vLLM-Omni aims to be fast, flexible, and easy to use with the following features:
# Overview of function call dependency
### Performance and Acceleration
![][image11]
The framework achieves high performance through several optimization techniques:
* **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.
[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
### Flexibility and Usability
[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
vLLM-Omni is designed to be flexible and straightforward for users:
[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
* **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.
[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
# Interface design
[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
If you use vLLM, then you know how to use vLLM-Omni from Day 0:
[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
<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>
[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
Taking **Qwen3-Omni** as an example:
[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
## 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:
[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
```
# Create an omni_lm with HF model name.
from vllm_omni.entrypoints.omni import Omni
[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
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/).
[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
+1 -200
View File
@@ -1,200 +1 @@
# 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.
## OmniConnector for fully disaggregated serving (update soon)
@@ -1,61 +0,0 @@
# 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]"
```
+1 -1
View File
@@ -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
+509
View File
@@ -0,0 +1,509 @@
# 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,5 +4,4 @@ vLLM-Omni supports the following hardware platforms:
- [GPU](gpu.md)
- [NVIDIA CUDA](gpu.md)
- [AMD ROCm](gpu.md)
- [NPU](npu.md)
-23
View File
@@ -14,10 +14,6 @@ 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
@@ -30,11 +26,6 @@ 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
@@ -43,22 +34,8 @@ 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,9 +74,6 @@ 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]
@@ -1,115 +0,0 @@
# --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]
+1 -2
View File
@@ -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 / AMD GPU
## List of Supported Models for Nvidia GPU
<style>
th {
@@ -22,7 +22,6 @@ 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.

Before

Width:  |  Height:  |  Size: 107 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 97 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 126 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 367 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 480 KiB

@@ -1,30 +0,0 @@
# 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"
``````
@@ -1,9 +1,9 @@
# Qwen-Image Offline Inference
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_image>.
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen_image>.
This folder provides two simple entrypoints for experimenting with `Qwen/Qwen-Image` `Tongyi-MAI/Z-Image-Turbo` using vLLM-Omni:
This folder provides two simple entrypoints for experimenting with `Qwen/Qwen-Image` 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,7 +13,6 @@ 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 \
@@ -34,6 +33,8 @@ 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:
@@ -48,9 +49,9 @@ Then open `http://localhost:7862/` on your local browser to interact with the we
??? abstract "gradio_demo.py"
``````py
--8<-- "examples/offline_inference/text_to_image/gradio_demo.py"
--8<-- "examples/offline_inference/qwen_image/gradio_demo.py"
``````
??? abstract "text_to_image.py"
``````py
--8<-- "examples/offline_inference/text_to_image/text_to_image.py"
--8<-- "examples/offline_inference/qwen_image/text_to_image.py"
``````
@@ -1,38 +0,0 @@
# 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"
``````
-99
View File
@@ -1,99 +0,0 @@
# 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
@@ -1,30 +0,0 @@
# 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,15 +496,7 @@ 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()
@@ -1,38 +0,0 @@
# 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"
``````
@@ -1,105 +0,0 @@
# 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,58 +91,9 @@ 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)]
@@ -549,8 +500,11 @@ def main():
print(f"Using custom stage configs: {args.stage_configs_path}")
sampling_params = build_sampling_params(SEED, model_name)
cli_args = build_async_omni_cli_args(args)
omni = AsyncOmni(model=args.model, cli_args=cli_args)
omni = AsyncOmni(
model=args.model,
stage_configs_path=args.stage_configs_path,
init_timeout=ASYNC_INIT_TIMEOUT,
)
print("✓ AsyncOmni initialized successfully")
prompt_args_template = create_prompt_args(args)
@@ -94,58 +94,9 @@ 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)]
@@ -555,8 +506,11 @@ def main():
print(f"Using custom stage configs: {args.stage_configs_path}")
sampling_params = build_sampling_params(SEED, model_name)
cli_args = build_async_omni_cli_args(args)
omni = AsyncOmni(model=args.model, cli_args=cli_args)
omni = AsyncOmni(
model=args.model,
stage_configs_path=args.stage_configs_path,
init_timeout=ASYNC_INIT_TIMEOUT,
)
print("✓ AsyncOmni initialized successfully")
prompt_args_template = create_prompt_args(args)
@@ -1,236 +0,0 @@
# 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
@@ -1,99 +0,0 @@
# 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
@@ -1,64 +0,0 @@
# 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}")
@@ -1,48 +0,0 @@
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
@@ -1,98 +0,0 @@
# 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
-206
View File
@@ -1,206 +0,0 @@
# 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
-58
View File
@@ -1,58 +0,0 @@
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
File renamed without changes.
File renamed without changes.
@@ -1,6 +1,6 @@
# stage config for running qwen2.5-omni with architecture of OmniLLM.
# The following config has been verified on 2x 24GB GPU (L4/RTX3090/RTX4090).
# The following config has been verified on 1x 24GB GPU (L4/RTX3090).
# This config is optimized for CI e2e tests.
stage_args:
- stage_id: 0
@@ -18,19 +18,15 @@ from .utils import create_new_process_for_each_test
models = ["Qwen/Qwen2.5-Omni-3B"]
# CI stage config optimized for 24GB GPU (L4/RTX3090)
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]
CI_STAGE_CONFIG_PATH = str(Path(__file__).parent / "stage_configs" / "qwen2_5_omni_ci.yaml")
@pytest.mark.core_model
@pytest.mark.parametrize("test_config", test_params)
@pytest.mark.parametrize("model", models)
@create_new_process_for_each_test()
def test_mixed_modalities_to_audio(omni_runner: type[OmniRunner], test_config: tuple[str, str]) -> None:
def test_mixed_modalities_to_audio(omni_runner: type[OmniRunner], model: str) -> None:
"""Test processing audio, image, and video together, generating audio output."""
model, stage_config_path = test_config
with omni_runner(model, seed=42, stage_configs_path=stage_config_path) as runner:
with omni_runner(model, seed=42, stage_configs_path=CI_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
File renamed without changes.
View File
Whitespace-only changes.
+14
View File
@@ -0,0 +1,14 @@
# 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"]
@@ -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 → 16-layer RVQ codec codes)
# 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.
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
E2E offline tests for Omni model with video input and audio output.
E2E tests for Qwen3-Omni model with audio/video/image input and audio output.
"""
import os
@@ -10,24 +10,20 @@ from pathlib import Path
import pytest
from vllm.assets.video import VideoAsset
from .conftest import OmniRunner
from ..multi_stages.conftest import OmniRunner
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
# CI stage config for 2xH100-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]
CI_STAGE_CONFIG_PATH = str(Path(__file__).parent / "stage_configs" / "qwen3_omni_ci.yaml")
@pytest.mark.parametrize("test_config", test_params)
def test_video_to_audio(omni_runner: type[OmniRunner], test_config) -> None:
@pytest.mark.parametrize("model", models)
def test_video_to_audio(omni_runner: type[OmniRunner], model: str) -> None:
"""Test processing video, generating audio output."""
model, stage_config_path = test_config
with omni_runner(model, seed=42, stage_configs_path=stage_config_path) as runner:
with omni_runner(model, seed=42, stage_configs_path=CI_STAGE_CONFIG_PATH) as runner:
# Prepare inputs
question = "Describe the video briefly."
video = VideoAsset(name="baby_reading", num_frames=4).np_ndarrays
-1
View File
@@ -104,7 +104,6 @@ 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
View File
@@ -1,41 +0,0 @@
# 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",
]
-82
View File
@@ -1,82 +0,0 @@
# 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
-67
View File
@@ -1,67 +0,0 @@
# 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
View File
@@ -1,100 +0,0 @@
# 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
View File
@@ -1,45 +0,0 @@
# 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
View File
@@ -1,96 +0,0 @@
# 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
View File
@@ -1,67 +0,0 @@
# 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
View File
@@ -1,334 +0,0 @@
# 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
View File
@@ -1,246 +0,0 @@
# 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
View File
@@ -1,40 +0,0 @@
# 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
+1 -17
View File
@@ -2,7 +2,6 @@
# 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
@@ -61,13 +60,9 @@ class OmniDiffusionConfig:
# Workload type
# workload_type: WorkloadType = WorkloadType.T2V
# Cache strategy (legacy)
# Cache strategy
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
@@ -167,8 +162,6 @@ 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"
@@ -213,17 +206,8 @@ 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)
-101
View File
@@ -1,101 +0,0 @@
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,9 +270,6 @@ 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
@@ -549,46 +546,31 @@ class QwenImagePipeline(
if self.interrupt:
continue
self._current_timestep = t
# Broadcast timestep to match batch size
# broadcast to batch dimension and place on same device/dtype as latents
timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype)
# 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
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]
if do_true_cfg:
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]
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]
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(od_config=od_config)
self.transformer = QwenImageTransformer2DModel()
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
@@ -1,4 +0,0 @@
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"]
@@ -1,426 +0,0 @@
# 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.")
@@ -1,704 +0,0 @@
# 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
-6
View File
@@ -24,11 +24,6 @@ _DIFFUSION_MODELS = {
"pipeline_z_image",
"ZImagePipeline",
),
"WanPipeline": (
"wan2_2",
"pipeline_wan2_2",
"Wan22Pipeline",
),
}
@@ -66,7 +61,6 @@ _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 = {
+4
View File
@@ -124,6 +124,10 @@ 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
-17
View File
@@ -84,29 +84,12 @@ 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)
-25
View File
@@ -1,25 +0,0 @@
# 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",
]
@@ -1,37 +0,0 @@
# 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",
]
@@ -1,168 +0,0 @@
# 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
@@ -1,2 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
@@ -1,72 +0,0 @@
# 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)
@@ -1,170 +0,0 @@
# 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)
@@ -1,118 +0,0 @@
# 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}
@@ -1,87 +0,0 @@
# 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)
@@ -1,2 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
@@ -1,40 +0,0 @@
# 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
@@ -1,322 +0,0 @@
# 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
@@ -1,14 +0,0 @@
# 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)
@@ -1,79 +0,0 @@
# 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)
@@ -1,10 +0,0 @@
# 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",
]
-183
View File
@@ -1,183 +0,0 @@
# 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
+62 -95
View File
@@ -33,16 +33,6 @@ 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
@@ -53,12 +43,10 @@ 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.utils import (
load_stage_configs_from_model,
load_stage_configs_from_yaml,
resolve_model_config_path,
)
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.outputs import OmniRequestOutput
logger = init_logger(__name__)
@@ -105,29 +93,16 @@ 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
@@ -163,56 +138,34 @@ 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))
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._ctx = mp.get_context("spawn")
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_stages(model)
self._start_stage_processes(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_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"
)
def _start_stage_processes(self, model: str) -> None:
for stage_id, stage in enumerate(self.stage_list):
# Use unbounded queues to avoid deadlock when seeding many requests
in_q = self._queue_cls()
out_q = self._queue_cls()
in_q: mp.Queue = self._ctx.Queue(maxsize=0)
out_q: mp.Queue = self._ctx.Queue(maxsize=0)
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 if self.worker_backend != "ray" else None,
ctx=self._ctx,
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)
@@ -238,9 +191,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:
@@ -313,14 +266,12 @@ 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 in range(last_stage_id, -1, -1):
if getattr(self.stage_list[_sid], "final_output", False):
final_stage_id_for_e2e = _sid
break
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)
if final_stage_id_for_e2e < 0:
final_stage_id_for_e2e = last_stage_id
final_stage_id_for_e2e = len(self.stage_list) - 1
except Exception as e:
logger.debug(
"[Orchestrator] Failed to determine final stage for E2E; \
@@ -328,7 +279,7 @@ class AsyncOmni(EngineClient):
e,
exc_info=True,
)
final_stage_id_for_e2e = last_stage_id
final_stage_id_for_e2e = len(self.stage_list) - 1
# Metrics/aggregation helper
metrics = OrchestratorMetrics(
num_stages,
@@ -411,8 +362,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_key = str(req_id)
if stage_id == final_stage_id_for_e2e and rid_key not in metrics.e2e_done:
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:
metrics.on_finalize_request(
stage_id,
req_id,
@@ -441,36 +392,52 @@ 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]
# 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,
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",
)
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."
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,
}
)
logger.error(error_msg)
raise RuntimeError(error_msg)
logger.debug(
"[Orchestrator] Forwarded request %s to stage-%s",
req_id,
-13
View File
@@ -95,19 +95,6 @@ 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
+27 -18
View File
@@ -301,7 +301,7 @@ def compute_and_log_stage_request_stats(
def record_stage_metrics(
per_request: dict[str, dict[str, Any]],
per_request: dict[int, 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_key = str(req_id)
pr = per_request.setdefault(rid_key, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
rid_int = int(req_id)
pr = per_request.setdefault(rid_int, {"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, str], dict[str, float]],
transfer_edge_req: dict[tuple[int, int, int], dict[str, float]],
transfer_agg: dict[tuple[int, int], dict[str, float]],
per_request: dict[str, dict[str, Any]],
per_request: dict[int, dict[str, Any]],
stage_id: int,
req_id: Any,
rx_bytes: float,
@@ -355,8 +355,11 @@ def aggregate_rx_and_maybe_total(
agg["rx_count"] += 1.0
# Try combine with sender-side timing if present
rid_key = str(req_id)
s = transfer_edge_req.get((stage_id - 1, stage_id, rid_key))
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))
if s is None:
return None
tx_ms = float(s.get("tx_ms", 0.0))
@@ -366,7 +369,7 @@ def aggregate_rx_and_maybe_total(
agg["total_count"] += 1.0
# accumulate per-request transfer totals
try:
pr = per_request.setdefault(rid_key, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
pr = per_request.setdefault(rid_int, {"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:
@@ -379,7 +382,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, str], dict[str, float]],
transfer_edge_req: dict[tuple[int, int, int], dict[str, float]],
from_stage: int,
to_stage: int,
req_id: Any,
@@ -405,8 +408,11 @@ def record_sender_transfer_agg(
agg["sum_ms"] += float(tx_ms)
agg["count"] += 1.0
# Store sender-side timing for per-request combination
rid_key = str(req_id)
transfer_edge_req[(from_stage, to_stage, rid_key)] = {
try:
rid_int = int(req_id)
except Exception:
rid_int = req_id
transfer_edge_req[(from_stage, to_stage, rid_int)] = {
"tx_ms": float(tx_ms),
"size_bytes": float(size_bytes),
}
@@ -506,12 +512,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, str], dict[str, float]] = {}
self.transfer_edge_req: dict[tuple[int, int, int], 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[str] = set()
self.per_request: dict[str, dict[str, Any]] = {}
self.e2e_done: set[int] = set()
self.per_request: dict[int, 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)
@@ -608,7 +614,10 @@ class OrchestratorMetrics:
)
def on_finalize_request(self, stage_id: int, req_id: Any, engine_outputs: list[Any], req_start_ts: float) -> None:
rid_key = str(req_id)
try:
rid_int = int(req_id)
except Exception:
rid_int = req_id
_t0 = float(req_start_ts)
_t1 = time.time()
# Update last output time for this stage
@@ -620,11 +629,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_key)
pr = self.per_request.setdefault(rid_key, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
self.e2e_done.add(rid_int)
pr = self.per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
per_req_record = {
"type": "overall_request",
"request_id": rid_key,
"request_id": rid_int,
"e2e_time_ms": e2e_ms,
"num_tokens_out": int(num_tokens),
"transfers_total_time_ms": float(pr.get("transfers_ms", 0.0)),
+65 -91
View File
@@ -1,7 +1,6 @@
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
@@ -20,18 +19,7 @@ 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
@@ -42,12 +30,10 @@ 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.utils import (
load_stage_configs_from_model,
load_stage_configs_from_yaml,
resolve_model_config_path,
)
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.outputs import OmniRequestOutput
logger = init_logger(__name__)
@@ -98,25 +84,14 @@ 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
@@ -149,51 +124,33 @@ class OmniLLM:
self.stage_list = [st for _, st in results]
logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list))
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._ctx = mp.get_context("spawn")
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_stages(model)
self._start_stage_processes(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_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"
)
def _start_stage_processes(self, model: str) -> None:
for stage_id, stage in enumerate(self.stage_list):
in_q = self._queue_cls()
out_q = self._queue_cls()
# 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)
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 if self.worker_backend != "ray" else None,
ctx=self._ctx,
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)
@@ -218,8 +175,6 @@ 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()
@@ -282,31 +237,28 @@ class OmniLLM:
# Orchestrator keeps stage objects for input derivation
num_stages = len(self.stage_list)
# 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)}
# Map from request_id to original prompt
request_id_to_prompt: dict[int, PromptType] = {i: p for i, p in enumerate(request_prompts)}
# Track per-request start time for end-to-end timing
_req_start_ts: dict[str, float] = {}
_req_start_ts: dict[int, 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 in range(last_stage_id, -1, -1):
if getattr(self.stage_list[_sid], "final_output", False):
final_stage_id_for_e2e = _sid
break
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)
if final_stage_id_for_e2e < 0:
final_stage_id_for_e2e = last_stage_id
final_stage_id_for_e2e = len(self.stage_list) - 1
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 = last_stage_id
final_stage_id_for_e2e = len(self.stage_list) - 1
# Metrics/aggregation helper
metrics = OrchestratorMetrics(
num_stages,
@@ -405,8 +357,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_key = str(req_id)
if stage_id == final_stage_id_for_e2e and rid_key not in metrics.e2e_done:
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:
metrics.on_finalize_request(
stage_id,
req_id,
@@ -435,28 +387,50 @@ class OmniLLM:
)
continue
sp_next: SamplingParams = sampling_params_list[next_stage_id] # type: ignore[index]
# 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,
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",
)
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."
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,
}
)
logger.debug(
"[Orchestrator] Forwarded request %s to stage-%s",
+28 -128
View File
@@ -15,7 +15,6 @@ import importlib
import logging
import multiprocessing as mp
import os
import sys
from typing import Any, Optional, Union
from vllm.inputs import TextPrompt
@@ -28,11 +27,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
@@ -164,9 +163,6 @@ 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.
@@ -180,23 +176,13 @@ 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
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
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)
@@ -206,74 +192,43 @@ class OmniStage:
"engine_args": engine_args,
"runtime": runtime_cfg,
"shm_threshold_bytes": self._shm_threshold_bytes,
"connectors_config": connectors_config or {},
}
if worker_backend == "ray":
if is_async:
self._ray_actor = start_ray_actor(
_stage_worker_async_entry,
ray_placement_group,
self.stage_id,
if is_async:
self._proc = ctx.Process(
target=_stage_worker_async_entry,
args=(
self,
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,
)
model,
stage_payload,
batch_timeout,
),
)
else:
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()
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 hasattr(self, "_ray_actor") and self._ray_actor:
kill_ray_actor(self._ray_actor)
self._ray_actor = None
elif self._proc is not None:
if self._proc is not None:
try:
self._proc.join(timeout=5)
except Exception as e:
@@ -283,7 +238,6 @@ 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:
@@ -380,8 +334,6 @@ 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,
@@ -395,7 +347,6 @@ 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:
@@ -489,18 +440,6 @@ 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})
@@ -557,25 +496,9 @@ def _stage_worker(
_in_flight_ms_by_rid[rid] = 0.0
except Exception:
_in_flight_ms_by_rid[rid] = 0.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))
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))
batch_request_ids.append(rid)
if isinstance(ein, list):
batch_engine_inputs.extend(ein)
@@ -748,8 +671,6 @@ 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,
@@ -764,7 +685,6 @@ 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
@@ -809,17 +729,6 @@ 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",
@@ -879,16 +788,7 @@ 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 = 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."
)
ein, _rx_metrics = maybe_load_from_ipc_with_metrics(task, 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))
+13 -39
View File
@@ -3,9 +3,11 @@ 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__)
@@ -25,11 +27,8 @@ def set_stage_devices(
Args:
stage_id: Stage identifier for logging
devices: Device specification:
- 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.
- 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.
- 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.
@@ -64,7 +63,7 @@ def set_stage_devices(
get_device_name_fn = torch.npu.get_device_name
device_type_label = "NPU"
elif device_type == "cuda":
import torch
import torch # noqa: WPS433
is_available_fn = torch.cuda.is_available
set_device_fn = torch.cuda.set_device
@@ -82,35 +81,16 @@ 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(mapped_devices[0])
selected_physical = int(toks[0])
logger.debug(
"[Stage-%s] Set %s to %s; logical 0 -> physical %s",
stage_id,
env_var,
mapped_devices_str,
devices,
selected_physical,
)
except Exception as e:
@@ -145,7 +125,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
import torch # noqa: WPS433
if is_available_fn():
try:
@@ -175,10 +155,8 @@ def set_stage_devices(
def serialize_obj(obj: Any) -> bytes:
"""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)
"""Serialize a Python object to bytes using cloudpickle."""
return cloudpickle.dumps(obj)
def shm_write_bytes(payload: bytes) -> dict[str, Any]:
@@ -260,9 +238,7 @@ 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:
from vllm_omni.distributed.omni_connectors.utils.serialization import OmniSerializer
return OmniSerializer.deserialize(shm_read_bytes(container[shm_key]))
return pickle.loads(shm_read_bytes(container[shm_key]))
return container[obj_key]
@@ -277,13 +253,11 @@ 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 = OmniSerializer.deserialize(payload)
obj = pickle.loads(payload)
try:
rx_bytes = int(meta.get("size", len(payload))) # type: ignore[call-arg]
except Exception:
+23 -44
View File
@@ -2,7 +2,7 @@ import os
from collections import Counter
from dataclasses import asdict, is_dataclass
from pathlib import Path
from typing import Any, Optional
from typing import Any
from omegaconf import OmegaConf
from vllm.transformers_utils.config import get_config
@@ -62,42 +62,7 @@ def _convert_dataclasses_to_dict(obj: Any) -> Any:
return obj
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:
def load_stage_configs_from_model(model: str, base_engine_args: dict = {}) -> list:
"""Load stage configurations from model's default config file.
Loads stage configurations based on the model type and device type.
@@ -113,14 +78,30 @@ def load_stage_configs_from_model(model: str, base_engine_args: Optional[dict] =
Raises:
FileNotFoundError: If no stage config file exists for the model type
"""
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)
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)
return stage_configs
def load_stage_configs_from_yaml(config_path: str, base_engine_args: Optional[dict] = None) -> list:
def load_stage_configs_from_yaml(config_path: str, base_engine_args: dict = {}) -> list:
"""Load stage configurations from a YAML file.
Args:
@@ -129,8 +110,6 @@ def load_stage_configs_from_yaml(config_path: str, base_engine_args: Optional[di
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 2x H100-80G GPU.
# The following config has been verified on 1x H100-80G GPU.
stage_args:
- stage_id: 0
runtime:
@@ -29,7 +29,6 @@ stage_args:
seed: 42
detokenize: True
repetition_penalty: 1.1
- stage_id: 1
runtime:
process: true
@@ -57,11 +56,10 @@ stage_args:
detokenize: True
repetition_penalty: 1.05
stop_token_ids: [8294]
- stage_id: 2
runtime:
process: true
devices: "2" # Example: use a different GPU than the previous stage; use "0" if single GPU
devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU
max_batch_size: 1
engine_args:
model_stage: code2wav
@@ -92,7 +90,6 @@ 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
@@ -1,130 +0,0 @@
# 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
Loaded 100 of 103 files, more files were not shown because too many files have changed in this diff. Show more