anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Merge branch 'main' into dev/vllm-align

Browse Source 
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
This commit is contained in:
tzhouam committed 2026-01-30 02:53:03 +00:00
2 parents 62dbf07fb6 0e07eb6ff9
commit 4b895459e4
196 files changed
+9406 -4132
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.buildkite/pipeline.yml
+18 -4
View File
@@ -18,11 +18,25 @@ steps:
# queue: "ascend"
- label: "Simple Unit Test"
depends_on: ~
depends_on: image-build
commands:
- ".buildkite/scripts/simple_test.sh"
- pytest -v -s tests/entrypoints/
- pytest -v -s tests/diffusion/cache/
- pytest -v -s tests/diffusion/lora/
- pytest -v -s tests/model_executor/models/qwen2_5_omni/test_audio_length.py
- pytest -v -s tests/worker/
- pytest -v -s tests/distributed/omni_connectors/test_kv_flow.py
agents:
queue: "cpu_queue_premerge"
queue: "gpu_1_queue"
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: "Diffusion Model Test"
timeout_in_minutes: 20
@@ -149,7 +163,7 @@ steps:
timeout_in_minutes: 20
depends_on: image-build
commands:
- pytest -s -v tests/diffusion/test_gpu_diffusion_worker.py
- pytest -s -v tests/diffusion/test_diffusion_worker.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
.buildkite/scripts/hardware_ci/run-amd-test.sh
+6
View File
@@ -116,6 +116,9 @@ if [[ $commands == *"--shard-id="* ]]; then
--shm-size=16gb \
--group-add "$render_gid" \
--rm \
-e MIOPEN_DEBUG_CONV_DIRECT=0 \
-e MIOPEN_DEBUG_CONV_GEMM=0 \
-e VLLM_ROCM_USE_AITER=1 \
-e HIP_VISIBLE_DEVICES="${GPU}" \
-e HF_TOKEN \
-e AWS_ACCESS_KEY_ID \
@@ -148,6 +151,9 @@ else
--shm-size=16gb \
--group-add "$render_gid" \
--rm \
-e MIOPEN_DEBUG_CONV_DIRECT=0 \
-e MIOPEN_DEBUG_CONV_GEMM=0 \
-e VLLM_ROCM_USE_AITER=1 \
-e HF_TOKEN \
-e AWS_ACCESS_KEY_ID \
-e AWS_SECRET_ACCESS_KEY \
.buildkite/scripts/simple_test.sh
-58
View File
@@ -1,58 +0,0 @@
#!/usr/bin/env bash
set -euo pipefail
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../.." && pwd)"
UV_BIN="uv"
PYTHON_BOOTSTRAP="${PYTHON:-python3}"
VENV_DIR="${VENV_DIR:-${REPO_ROOT}/.venv-simple-test}"
TARGET_PY_VERSION="3.12"
if ! command -v "${PYTHON_BOOTSTRAP}" >/dev/null 2>&1; then
PYTHON_BOOTSTRAP="python"
fi
cd "${REPO_ROOT}"
# Ensure pip-installed scripts (including uv) are on PATH
PYTHON_SCRIPTS="$(${PYTHON_BOOTSTRAP} -c 'import sysconfig; print(sysconfig.get_path("scripts"))')"
USER_BASE="$(${PYTHON_BOOTSTRAP} -m site --user-base 2>/dev/null || true)"
PATH="${PYTHON_SCRIPTS}:${PATH}"
if [[ -n "${USER_BASE}" ]]; then
PATH="${USER_BASE}/bin:${PATH}"
fi
if ! command -v "${UV_BIN}" >/dev/null 2>&1; then
"${PYTHON_BOOTSTRAP}" -m pip install --upgrade pip
"${PYTHON_BOOTSTRAP}" -m pip install uv
hash -r
fi
TARGET_PY_PATH="$(${UV_BIN} python find "${TARGET_PY_VERSION}" 2>/dev/null | head -n1 || true)"
if [[ -z "${TARGET_PY_PATH}" ]]; then
"${UV_BIN}" python install "${TARGET_PY_VERSION}"
TARGET_PY_PATH="$(${UV_BIN} python find "${TARGET_PY_VERSION}" 2>/dev/null | head -n1 || true)"
fi
if [[ -z "${TARGET_PY_PATH}" ]]; then
echo "Failed to locate python ${TARGET_PY_VERSION}" >&2
exit 1
fi
if [[ ! -d "${VENV_DIR}" ]]; then
"${UV_BIN}" venv --python "${TARGET_PY_PATH}" "${VENV_DIR}"
fi
VENV_PYTHON="${VENV_DIR}/bin/python"
[[ -x "${VENV_PYTHON}" ]] || { echo "Python not found in ${VENV_DIR}"; exit 1; }
"${UV_BIN}" pip install --python "${VENV_PYTHON}" --upgrade --force-reinstall vllm --extra-index-url https://wheels.vllm.ai/706f123b23b75097da9a3f51ce538bd784fdc75b
"${UV_BIN}" pip install --python "${VENV_PYTHON}" -e ".[dev]"
"${VENV_PYTHON}" -m pytest -v -s tests/entrypoints/
"${VENV_PYTHON}" -m pytest -v -s tests/diffusion/cache/
"${VENV_PYTHON}" -m pytest -v -s tests/diffusion/lora/
"${VENV_PYTHON}" -m pytest -v -s tests/model_executor/models/qwen2_5_omni/test_audio_length.py
"${VENV_PYTHON}" -m pytest -v -s tests/worker/
"${VENV_PYTHON}" -m pytest -v -s tests/distributed/omni_connectors/test_kv_flow.py
.buildkite/test-amd.yaml
+44 -34
View File
@@ -8,14 +8,32 @@ steps:
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- export MIOPEN_DEBUG_CONV_DIRECT=0
- export MIOPEN_DEBUG_CONV_GEMM=0
- 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
- pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
- label: "Diffusion Images API LoRA E2E"
timeout_in_minutes: 20
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
- label: "Diffusion Model CPU offloading Test"
timeout_in_minutes: 20
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
- label: "Diffusion Cache Backend Test"
timeout_in_minutes: 15
agent_pool: mi325_1
@@ -26,25 +44,30 @@ steps:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export MIOPEN_DEBUG_CONV_DIRECT=0
- export MIOPEN_DEBUG_CONV_GEMM=0
- 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
- pytest -s -v tests/e2e/offline_inference/test_cache_dit.py tests/e2e/offline_inference/test_teacache.py
- label: "Diffusion Parallelism Test"
timeout_in_minutes: 15
- label: "Diffusion Sequence Parallelism Test"
timeout_in_minutes: 20
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export MIOPEN_DEBUG_CONV_DIRECT=0
- export MIOPEN_DEBUG_CONV_GEMM=0
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
- label: "Diffusion Tensor Parallelism Test"
timeout_in_minutes: 20
agent_pool: mi325_2
depends_on: amd-build
commands:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_zimage_tensor_parallel.py
- label: "Diffusion GPU Worker Test"
timeout_in_minutes: 20
agent_pool: mi325_2
@@ -52,9 +75,7 @@ steps:
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export MIOPEN_DEBUG_CONV_DIRECT=0
- export MIOPEN_DEBUG_CONV_GEMM=0
- pytest -s -v tests/diffusion/test_gpu_diffusion_worker.py
- pytest -s -v tests/diffusion/test_diffusion_worker.py
- label: "Omni Model Test Qwen2-5-Omni"
timeout_in_minutes: 15
@@ -66,12 +87,6 @@ steps:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export MIOPEN_DEBUG_CONV_DIRECT=0
- export MIOPEN_DEBUG_CONV_GEMM=0
- 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
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
- label: "Omni Model Test Qwen3-Omni"
@@ -83,9 +98,10 @@ steps:
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export MIOPEN_DEBUG_CONV_DIRECT=0
- export MIOPEN_DEBUG_CONV_GEMM=0
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
- pytest -s -v tests/e2e/online_serving/test_async_omni.py
- label: "Diffusion Image Edit Test"
timeout_in_minutes: 15
@@ -97,10 +113,4 @@ steps:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export MIOPEN_DEBUG_CONV_DIRECT=0
- export MIOPEN_DEBUG_CONV_GEMM=0
- 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
- pytest -s -v tests/e2e/online_serving/test_i2i_multi_image_input.py
docs/configuration/stage_configs.md
+3 -3
View File
@@ -40,7 +40,7 @@ stage_args:
engine_args: # Engine arguments for a certain engine
model_stage: thinker
model_arch: Qwen2_5OmniForConditionalGeneration # The model implementation registered in model_executor/models/registry.py
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker # The specific worker used
worker_type: ar # The specific worker used
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler # The specific scehduler used
gpu_memory_utilization: 0.8 # The gpu memory allocation for the stage within a single chip
enforce_eager: true # Now we only support eager mode
@@ -66,7 +66,7 @@ stage_args:
engine_args:
model_stage: talker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.8
enforce_eager: true
@@ -92,7 +92,7 @@ stage_args:
engine_args:
model_stage: code2wav
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
gpu_memory_utilization: 0.15
enforce_eager: true
docs/configuration/stage_configs/qwen2_5_omni.yaml
+3 -3
View File
@@ -8,7 +8,7 @@ stage_args:
engine_args:
model_stage: thinker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.8
enforce_eager: true # Now we only support eager mode
@@ -34,7 +34,7 @@ stage_args:
engine_args:
model_stage: talker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.8
enforce_eager: true
@@ -60,7 +60,7 @@ stage_args:
engine_args:
model_stage: code2wav
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
gpu_memory_utilization: 0.15
enforce_eager: true
docs/contributing/profiling.md
+9 -26
View File
@@ -2,7 +2,7 @@
> **Warning:** Profiling incurs significant overhead. Use only for development and debugging, never in production.
vLLM-Omni uses the PyTorch Profiler to analyze performance across both **Multi-Stage LLMs** and **Diffusion Models**.
vLLM-Omni uses the PyTorch Profiler to analyze performance across both **multi-stage omni-modality models** and **diffusion models**.
### 1. Set the Output Directory
Before running any script, set this environment variable. The system detects this and automatically saves traces here.
@@ -11,7 +11,7 @@ Before running any script, set this environment variable. The system detects thi
export VLLM_TORCH_PROFILER_DIR=./profiles
```
### 2. Start Profiling
### 2. Profiling Omni-Modality Models
It is best to limit profiling to one iteration to keep trace files manageable.
@@ -77,12 +77,12 @@ omni_llm.close()
**Examples**:
1. **Qwen-omni 2.5**: [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py)
1. **Qwen2.5-Omni**: [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py)
2. **Qwen-omni 3.0**: [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py)
2. **Qwen3-Omni**: [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py)
**For Diffusion Models as a single stage**
### 3. Profiling diffusion models
Diffusion profiling is End-to-End, capturing encoding, denoising loops, and decoding.
@@ -131,27 +131,10 @@ python image_to_video.py \
2. **Wan-AI/Wan2.2-I2V-A14B-Diffusers**: [https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)
> **Note:**
As of now, asynchronous (online) profiling is not fully supported in vLLM-Omni. While start_profile() and stop_profile() methods exist, they are only reliable in offline inference scripts (e.g., the provided end2end.py examples). Do not use them in server-mode or streaming scenarios—traces may be incomplete or fail to flush.
**Online Inference(Async)**
For online serving using AsyncOmni, the methods are asynchronous. This allows you to toggle profiling dynamically without restarting the server.
```python
from vllm_omni import AsyncOmni
# Inside an async function:
async_omni = AsyncOmni.from_engine_args(engine_args)
await async_omni.start_profile()
async for output in async_omni.generate(prompt, sampling_params, request_id):
# Process outputs...
pass
await async_omni.stop_profile()
```
### 3. Analyzing Omni Traces
### 4. Analyzing Omni Traces
Output files are saved to your configured ```VLLM_TORCH_PROFILER_DIR```.
@@ -163,4 +146,4 @@ Output files are saved to your configured ```VLLM_TORCH_PROFILER_DIR```.
- [Perfetto](https://ui.perfetto.dev/)(recommended)
- ```chrome://tracing```(Chrome only)
**Note**: vLLM-Omni reuses the PyTorch Profiler infrastructure from vLLM. See the official vLLM profiler documentation: [vLLM Profiling Guide](https://docs.vllm.ai/en/latest/dev/profiling.html)
**Note**: vLLM-Omni reuses the PyTorch Profiler infrastructure from vLLM. See the official vLLM profiler documentation: [vLLM Profiling Guide](https://docs.vllm.ai/en/stable/contributing/profiling/)
docs/getting_started/quickstart.md
+40
View File
@@ -46,6 +46,46 @@ if __name__ == "__main__":
images[0].save("coffee.png")
```
You can pass a list of prompts and wait for them to process altogether, shown below.
!!! info
However, it is not currently recommended to do so
because not all models support batch inference,
and batch requesting mostly does not provide significant performance improvement (despite the impression that it does).
This feature is primarily for the sake of interface compatibility with vLLM and to allow for future improvements.
```python
from vllm_omni.entrypoints.omni import Omni
if __name__ == "__main__":
omni = Omni(
model="Tongyi-MAI/Z-Image-Turbo",
# stage_configs_path="./stage-config.yaml", # See below
)
prompts = [
"a cup of coffee on a table",
"a toy dinosaur on a sandy beach",
"a fox waking up in bed and yawning",
]
omni_outputs = omni.generate(prompts)
for i_prompt, prompt_output in enumerate(omni_outputs):
this_request_output = prompt_output.request_output[0]
this_images = this_request_output.images
for i_image, image in enumerate(this_images):
image.save(f"p{i_prompt}-img{i_image}.jpg")
print("saved to", f"p{i_prompt}-img{i_image}.jpg")
# saved to p0-img0.jpg
# saved to p1-img0.jpg
# saved to p2-img0.jpg
```
!!! info
For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
batched inputs, you can [modify this configuration](../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
For more usages, please refer to [offline inference](../user_guide/examples/offline_inference/qwen2_5_omni.md)
## Online Serving with OpenAI-Completions API
docs/models/supported_models.md
+1
View File
@@ -33,6 +33,7 @@ th {
|`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` |
|`StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` |
|`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
|`FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` |
|`StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` |
|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` |
|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` |
docs/user_guide/diffusion/cache_dit_acceleration.md
+2 -1
View File
@@ -18,6 +18,7 @@ Enable cache-dit acceleration by simply setting `cache_backend="cache_dit"`. Cac
```python
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
# Simplest way: just enable cache-dit with default parameters
omni = Omni(
@@ -27,7 +28,7 @@ omni = Omni(
images = omni.generate(
"a beautiful landscape",
num_inference_steps=50,
OmniDiffusionSamplingParams(num_inference_steps=50),
)
```
docs/user_guide/diffusion/parallelism_acceleration.md
+31 -12
View File
@@ -30,6 +30,7 @@ The following table shows which models are currently supported by parallelism me
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) |
| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ❌ |
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ |
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ❌ | ✅ |
!!! note "TP Limitations for Diffusion Models"
@@ -67,10 +68,12 @@ omni = Omni(
)
outputs = omni.generate(
prompt="a cat reading a book",
num_inference_steps=9,
width=512,
height=512,
"a cat reading a book",
OmniDiffusionSamplingParams(
num_inference_steps=9,
width=512,
height=512,
),
)
```
@@ -83,6 +86,7 @@ outputs = omni.generate(
An example of offline inference script using [Ulysses-SP](https://arxiv.org/pdf/2309.14509) is shown below:
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
ulysses_degree = 2
@@ -91,7 +95,10 @@ omni = Omni(
parallel_config=DiffusionParallelConfig(ulysses_degree=2)
)
outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
)
```
See `examples/offline_inference/text_to_image/text_to_image.py` for a complete working example.
@@ -133,6 +140,7 @@ Ring-Attention ([arxiv paper](https://arxiv.org/abs/2310.01889)) splits the inpu
An example of offline inference script using Ring-Attention is shown below:
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
ring_degree = 2
@@ -141,7 +149,10 @@ omni = Omni(
parallel_config=DiffusionParallelConfig(ring_degree=2)
)
outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
)
```
See `examples/offline_inference/text_to_image/text_to_image.py` for a complete working example.
@@ -183,6 +194,7 @@ You can combine both Ulysses-SP and Ring-Attention for larger scale parallelism.
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
# Hybrid: 2 Ulysses × 2 Ring = 4 GPUs total
@@ -191,7 +203,10 @@ omni = Omni(
parallel_config=DiffusionParallelConfig(ulysses_degree=2, ring_degree=2)
)
outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
)
```
##### Online Serving
@@ -374,11 +389,15 @@ omni = Omni(
)
outputs = omni.generate(
prompt="turn this cat to a dog",
negative_prompt="low quality, blurry",
true_cfg_scale=4.0,
pil_image=input_image,
num_inference_steps=50,
{
"prompt": "turn this cat to a dog",
"negative_prompt": "low quality, blurry",
},
OmniDiffusionSamplingParams(
true_cfg_scale=4.0,
pil_image=input_image,
num_inference_steps=50,
),
)
```
docs/user_guide/diffusion/teacache.md
+14 -2
View File
@@ -8,6 +8,7 @@ Enable TeaCache by setting `cache_backend` to `"tea_cache"`:
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
# Simple configuration - model_type is automatically extracted from pipeline.__class__.__name__
omni = Omni(
@@ -17,7 +18,12 @@ omni = Omni(
"rel_l1_thresh": 0.2 # Optional, defaults to 0.2
}
)
outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(
num_inference_steps=50,
),
)
```
### Using Environment Variable
@@ -68,13 +74,19 @@ Controls the balance between speed and quality. Lower values prioritize quality,
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
omni = Omni(
model="Qwen/Qwen-Image",
cache_backend="tea_cache",
cache_config={"rel_l1_thresh": 0.2}
)
outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(
num_inference_steps=50,
),
)
```
## Performance Tuning
docs/user_guide/diffusion_acceleration.md
+40 -8
View File
@@ -96,6 +96,7 @@ To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** m
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
omni = Omni(
model="Qwen/Qwen-Image",
@@ -103,13 +104,19 @@ omni = Omni(
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)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(
num_inference_steps=50,
),
)
```
### Using Cache-DiT
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
omni = Omni(
model="Qwen/Qwen-Image",
@@ -123,7 +130,12 @@ omni = Omni(
}
)
outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(
num_inference_steps=50,
),
)
```
### Using Ulysses-SP
@@ -131,6 +143,7 @@ outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_st
Run text-to-image:
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
ulysses_degree = 2
@@ -139,13 +152,17 @@ omni = Omni(
parallel_config=DiffusionParallelConfig(ulysses_degree=ulysses_degree)
)
outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
)
```
Run image-to-image:
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
ulysses_degree = 2
@@ -154,8 +171,13 @@ omni = Omni(
parallel_config=DiffusionParallelConfig(ulysses_degree=ulysses_degree)
)
outputs = omni.generate(prompt="turn this cat to a dog",
pil_image=input_image, num_inference_steps=50)
outputs = omni.generate(
{
"prompt": "turn this cat to a dog",
"multi_modal_data": {"image": input_image}
},
OmniDiffusionSamplingParams(num_inference_steps=50),
)
```
### Using Ring-Attention
@@ -163,6 +185,7 @@ outputs = omni.generate(prompt="turn this cat to a dog",
Run text-to-image:
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
ring_degree = 2
@@ -171,7 +194,10 @@ omni = Omni(
parallel_config=DiffusionParallelConfig(ring_degree=2)
)
outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
)
```
### Using CFG-Parallel
@@ -182,6 +208,7 @@ CFG-Parallel splits the CFG positive/negative branches across GPUs. Use it when
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
cfg_parallel_size = 2
@@ -190,8 +217,13 @@ omni = Omni(
parallel_config=DiffusionParallelConfig(cfg_parallel_size=cfg_parallel_size)
)
outputs = omni.generate(prompt="turn this cat to a dog",
pil_image=input_image, num_inference_steps=50, true_cfg_scale=4.0)
outputs = omni.generate(
{
"prompt": "turn this cat to a dog",
"multi_modal_data": {"image": input_image}
},
OmniDiffusionSamplingParams(num_inference_steps=50, true_cfg_scale=4.0),
)
```
## Documentation
docs/user_guide/examples/offline_inference/bagel.md
+179
View File
@@ -0,0 +1,179 @@
# BAGEL-7B-MoT
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/bagel>.
## Set up
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
## Run examples
**Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, please modify the stage configuration to distribute the model across devices.
Get into the bagel folder
```bash
cd examples/offline_inference/bagel
```
### Modality Control
BAGEL-7B-MoT supports multiple modality modes. You can control the mode using the `--modality` argument:
#### Text to Image (text2img)
- **Pipeline**: Text → Thinker → DiT → VAE Decode → Image
- **Stages Used**: Stage 0 (Thinker) + Stage 1 (DiT)
- **KV Transfer**: Thinker sends KV cache to DiT for conditioned generation
Generate images from text prompts:
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat"
```
#### Image to Image (img2img)
- **Pipeline**: Image → VAE Encode → DiT → VAE Decode → New Image
- **Stages Used**: Stage 1 (DiT) only
- **Special**: Bypasses the Thinker stage, direct image-to-image transformation
Transform images based on text prompts:
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality img2img \
--image-path /path/to/image.jpg \
--prompts "Let the woman wear a blue dress"
```
#### Image to Text (img2text)
- **Pipeline**: Image → ViT + VAE Encode → Thinker → Text Output
- **Stages Used**: Stage 0 (Thinker) only
- **Special**: Uses both VAE latent encoding AND ViT semantic encoding for comprehensive image understanding
Generate text descriptions from images:
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality img2text \
--image-path /path/to/image.jpg \
--prompts "Describe this image in detail"
```
#### Text to Text (text2text)
- **Pipeline**: Text → Thinker → Text Output
- **Stages Used**: Stage 0 (Thinker) only
- **Special**: No visual components involved, operates as pure language model
Pure text generation:
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2text \
--prompts "What is the capital of France?"
# You can load prompts from a text file (one prompt per line):
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2text \
--txt-prompts /path/to/prompts.txt
```
### Inference Steps
Control the number of inference steps for image generation:
```bash
# You can adjust steps to 100 to improve image quality
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--steps 50 \
--prompts "A cute cat"
```
### Key arguments
BAGEL-7B-MoT supports **multiple modality modes** for different use cases.
The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)
#### 📌 Command Line Arguments (end2end.py)
| Argument | Type | Default | Description |
| :--------------------- | :----- | :---------------------------- | :----------------------------------------------------------- |
| `--model` | string | `ByteDance-Seed/BAGEL-7B-MoT` | Model path or name |
| `--modality` | choice | `text2img` | Modality mode: `text2img`, `img2img`, `img2text`, `text2text` |
| `--prompts` | list | `None` | Input text prompts directly |
| `--txt-prompts` | string | `None` | Path to txt file with one prompt per line |
| `--image-path` | string | `None` | Input image path (for `img2img`/`img2text`) |
| `--steps` | int | `50` | Number of inference steps |
| `--stage-configs-path` | string | `None` | Custom stage config file path |
| `--worker-backend` | choice | `process` | Worker backend: `process` or `ray` |
| `--ray-address` | string | `None` | Ray cluster address |
| `--enable-stats` | flag | `False` | Enable statistics logging |
| `--init-sleep-seconds` | int | `20` | Initialization sleep time |
| `--batch-timeout` | int | `5` | Batch timeout |
| `--init-timeout` | int | `300` | Initialization timeout |
------
#### ⚙️ Stage Configuration Parameters (bagel.yaml)
**Stage 0 - Thinker (LLM Stage)**
| Parameter | Value | Description |
| :------------------------------- | :------------------------------ | :----------------------- |
| `stage_type` | `llm` | Stage type |
| `devices` | `"0"` | GPU device ID |
| `max_batch_size` | `1` | Maximum batch size |
| `model_stage` | `thinker` | Model stage identifier |
| `model_arch` | `BagelForConditionalGeneration` | Model architecture |
| `gpu_memory_utilization` | `0.4` | GPU memory utilization |
| `tensor_parallel_size` | `1` | Tensor parallel size |
| `max_num_batched_tokens` | `32768` | Maximum batched tokens |
| `omni_kv_config.need_send_cache` | `true` | Whether to send KV cache |
------
**Stage 1 - DiT (Diffusion Stage)**
| Parameter | Value | Description |
| :------------------------------- | :---------- | :-------------------------- |
| `stage_type` | `diffusion` | Stage type |
| `devices` | `"0"` | GPU device ID |
| `max_batch_size` | `1` | Maximum batch size |
| `model_stage` | `dit` | Model stage identifier |
| `gpu_memory_utilization` | `0.4` | GPU memory utilization |
| `omni_kv_config.need_recv_cache` | `true` | Whether to receive KV cache |
| `engine_input_source` | `[0]` | Input source from Stage 0 |
------
#### 🔗 Runtime Configuration
| Parameter | Value | Description |
| :-------------------- | :------ | :------------------------------- |
| `window_size` | `-1` | Window size (-1 means unlimited) |
| `max_inflight` | `1` | Maximum inflight requests |
| `shm_threshold_bytes` | `65536` | Shared memory threshold (64KB) |
## FAQ
- If you encounter an error about the backend of librosa, try to install ffmpeg with the command below.
```bash
sudo apt update
sudo apt install ffmpeg
```
- If you don’t know how much VRAM is needed for the model or encounter the OOM error, you can try to decrease the max_model_len.
| Stage | VRAM |
| :------------------ | :--------------------------- |
| Stage-0 (Thinker) | **15.04 GiB** **+ KV Cache** |
| Stage-1 (DiT) | **26.50 GiB** |
| Total | **~42 GiB + KV Cache** |
docs/user_guide/examples/offline_inference/text_to_image.md
+36 -1
View File
@@ -23,7 +23,7 @@ if __name__ == "__main__":
images[0].save("coffee.png")
```
Or put more than one prompt in a request, processing them sequentially.
Or put more than one prompt in a request.
```python
from vllm_omni.entrypoints.omni import Omni
@@ -40,6 +40,41 @@ if __name__ == "__main__":
image = output.request_output[0].images[0].save(f"{i}.jpg")
```
!!! info
However, it is not currently recommended to do so
because not all models support batch inference,
and batch requesting mostly does not provide significant performance improvement (despite the impression that it does).
This feature is primarily for the sake of interface compatibility with vLLM and to allow for future improvements.
!!! info
For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
batched inputs, you can [modify this configuration](../../../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
Apart from string prompt, vLLM-Omni also supports dictionary prompts in the same style as vLLM.
This is useful for models that support negative prompts.
```python
from vllm_omni.entrypoints.omni import Omni
if __name__ == "__main__":
omni = Omni(model="Qwen/Qwen-Image")
outputs = omni.generate([
{
"prompt": "a cup of coffee on a table"，
"negative_prompt": "low resolution"
},
{
"prompt": "a toy dinosaur on a sandy beach"，
"negative_prompt": "cinematic, realistic"
}
])
for i, output in enumerate(outputs):
image = output.request_output[0].images[0].save(f"{i}.jpg")
```
## Local CLI Usage
```bash
docs/user_guide/examples/online_serving/bagel.md
+232
View File
@@ -0,0 +1,232 @@
# BAGEL-7B-MoT
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/bagel>.
## 🛠️ Installation
Please refer to [README.md](../../../README.md)
## Run examples (BAGEL-7B-MoT)
**Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, please modify the stage configuration to distribute the model across devices.
### Launch the Server
```bash
# Use default configuration
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091
```
Or use the convenience script:
```bash
cd /workspace/vllm-omni/examples/online_serving/bagel
bash run_server.sh
```
If you have a custom stage configs file, launch the server with the command below:
```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
```
### Send Multi-modal Request
Get into the bagel folder:
```bash
cd examples/online_serving/bagel
```
Send request via Python
```bash
python openai_chat_client.py --prompt "A cute cat" --modality text2img
```
The Python client supports the following command-line arguments:
- `--prompt` (or `-p`): Text prompt for generation (default: `A cute cat`)
- `--output` (or `-o`): Output file path for image results (default: `bagel_output.png`)
- `--server` (or `-s`): Server URL (default: `http://localhost:8091`)
- `--image-url` (or `-i`): Input image URL or local file path (for img2img/img2text modes)
- `--modality` (or `-m`): Task modality (default: `text2img`). Options: `text2img`, `img2img`, `img2text`, `text2text`
- `--height`: Image height in pixels (default: 512)
- `--width`: Image width in pixels (default: 512)
- `--steps`: Number of inference steps (default: 25)
- `--seed`: Random seed (default: 42)
- `--negative`: Negative prompt for image generation
Example with custom parameters:
```bash
python openai_chat_client.py \
--prompt "A futuristic city" \
--modality text2img \
--height 768 \
--width 768 \
--steps 50 \
--seed 42 \
--negative "blurry, low quality"
```
## Modality Control
BAGEL-7B-MoT supports **multiple modality modes** for different use cases.
The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)
| Modality | Input | Output | Description |
| ----------- | ------------ | ------ | -------------------------------------- |
| `text2img` | Text | Image | Generate images from text prompts |
| `img2img` | Image + Text | Image | Transform images using text guidance |
| `img2text` | Image + Text | Text | Generate text descriptions from images |
| `text2text` | Text | Text | Pure text generation |
### Text to Image (text2img)
Generate images from text prompts:
**Using Python client**
```bash
python openai_chat_client.py \
--prompt "A beautiful sunset over mountains" \
--modality text2img \
--output sunset.png \
--steps 50
```
**Using curl**
```bash
curl http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>A beautiful sunset over mountains<|im_end|>"}]}],
"modalities": ["image"],
"height": 512,
"width": 512,
"num_inference_steps": 50,
"seed": 42
}'
```
### Image to Image (img2img)
Transform images based on text prompts:
**Using Python client**
```bash
python openai_chat_client.py \
--prompt "Make the cat stand up" \
--modality img2img \
--image-url /path/to/input.jpg \
--output transformed.png
```
**Using curl**
```bash
IMAGE_BASE64=$(base64 -w 0 cat.jpg)
cat <<EOF > payload.json
{
"messages": [{
"role": "user",
"content": [
{"type": "text", "text": "<|im_start|>Make the cat stand up<|im_end|>"},
{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}}
]
}],
"modalities": ["image"],
"height": 512,
"width": 512,
"num_inference_steps": 50,
"seed": 42
}
EOF
curl http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d @payload.json
```
### Image to Text (img2text)
Generate text descriptions from images:
**Using Python client**
```bash
python openai_chat_client.py \
--prompt "Describe this image in detail" \
--modality img2text \
--image-url /path/to/image.jpg
```
**Using curl**
```bash
IMAGE_BASE64=$(base64 -w 0 cat.jpg)
cat <<EOF > payload.json
{
"messages": [{
"role": "user",
"content": [
{"type": "text", "text": "<|im_start|>user\n<|image_pad|>\nDescribe this image in detail<|im_end|>\n<|im_start|>assistant\n"},
{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}}
]
}],
"modalities": ["text"]
}
EOF
curl http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d @payload.json
```
### Text to Text (text2text)
Pure text generation:
**Using Python client**
```bash
python openai_chat_client.py \
--prompt "What is the capital of France?" \
--modality text2text
```
**Using curl**
```bash
curl http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n"}]}]
"modalities": ["text"]
}'
```
## FAQ
- If you encounter an error about the backend of librosa, try to install ffmpeg with the command below.
```bash
sudo apt update
sudo apt install ffmpeg
```
- If you don’t know how much VRAM is needed for the model or encounter the OOM error, you can try to decrease the max_model_len.
| Stage | VRAM |
| :------------------ | :--------------------------- |
| Stage-0 (Thinker) | **15.04 GiB** **+ KV Cache** |
| Stage-1 (DiT) | **26.50 GiB** |
| Total | **~42 GiB + KV Cache** |
examples/offline_inference/bagel/README.md
+177
View File
@@ -0,0 +1,177 @@
# BAGEL-7B-MoT
## Set up
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
## Run examples
**Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, please modify the stage configuration to distribute the model across devices.
Get into the bagel folder
```bash
cd examples/offline_inference/bagel
```
### Modality Control
BAGEL-7B-MoT supports multiple modality modes. You can control the mode using the `--modality` argument:
#### Text to Image (text2img)
- **Pipeline**: Text → Thinker → DiT → VAE Decode → Image
- **Stages Used**: Stage 0 (Thinker) + Stage 1 (DiT)
- **KV Transfer**: Thinker sends KV cache to DiT for conditioned generation
Generate images from text prompts:
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat"
```
#### Image to Image (img2img)
- **Pipeline**: Image → VAE Encode → DiT → VAE Decode → New Image
- **Stages Used**: Stage 1 (DiT) only
- **Special**: Bypasses the Thinker stage, direct image-to-image transformation
Transform images based on text prompts:
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality img2img \
--image-path /path/to/image.jpg \
--prompts "Let the woman wear a blue dress"
```
#### Image to Text (img2text)
- **Pipeline**: Image → ViT + VAE Encode → Thinker → Text Output
- **Stages Used**: Stage 0 (Thinker) only
- **Special**: Uses both VAE latent encoding AND ViT semantic encoding for comprehensive image understanding
Generate text descriptions from images:
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality img2text \
--image-path /path/to/image.jpg \
--prompts "Describe this image in detail"
```
#### Text to Text (text2text)
- **Pipeline**: Text → Thinker → Text Output
- **Stages Used**: Stage 0 (Thinker) only
- **Special**: No visual components involved, operates as pure language model
Pure text generation:
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2text \
--prompts "What is the capital of France?"
# You can load prompts from a text file (one prompt per line):
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2text \
--txt-prompts /path/to/prompts.txt
```
### Inference Steps
Control the number of inference steps for image generation:
```bash
# You can adjust steps to 100 to improve image quality
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--steps 50 \
--prompts "A cute cat"
```
### Key arguments
BAGEL-7B-MoT supports **multiple modality modes** for different use cases.
The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)
#### 📌 Command Line Arguments (end2end.py)
| Argument | Type | Default | Description |
| :--------------------- | :----- | :---------------------------- | :----------------------------------------------------------- |
| `--model` | string | `ByteDance-Seed/BAGEL-7B-MoT` | Model path or name |
| `--modality` | choice | `text2img` | Modality mode: `text2img`, `img2img`, `img2text`, `text2text` |
| `--prompts` | list | `None` | Input text prompts directly |
| `--txt-prompts` | string | `None` | Path to txt file with one prompt per line |
| `--image-path` | string | `None` | Input image path (for `img2img`/`img2text`) |
| `--steps` | int | `50` | Number of inference steps |
| `--stage-configs-path` | string | `None` | Custom stage config file path |
| `--worker-backend` | choice | `process` | Worker backend: `process` or `ray` |
| `--ray-address` | string | `None` | Ray cluster address |
| `--enable-stats` | flag | `False` | Enable statistics logging |
| `--init-sleep-seconds` | int | `20` | Initialization sleep time |
| `--batch-timeout` | int | `5` | Batch timeout |
| `--init-timeout` | int | `300` | Initialization timeout |
------
#### ⚙️ Stage Configuration Parameters (bagel.yaml)
**Stage 0 - Thinker (LLM Stage)**
| Parameter | Value | Description |
| :------------------------------- | :------------------------------ | :----------------------- |
| `stage_type` | `llm` | Stage type |
| `devices` | `"0"` | GPU device ID |
| `max_batch_size` | `1` | Maximum batch size |
| `model_stage` | `thinker` | Model stage identifier |
| `model_arch` | `BagelForConditionalGeneration` | Model architecture |
| `gpu_memory_utilization` | `0.4` | GPU memory utilization |
| `tensor_parallel_size` | `1` | Tensor parallel size |
| `max_num_batched_tokens` | `32768` | Maximum batched tokens |
| `omni_kv_config.need_send_cache` | `true` | Whether to send KV cache |
------
**Stage 1 - DiT (Diffusion Stage)**
| Parameter | Value | Description |
| :------------------------------- | :---------- | :-------------------------- |
| `stage_type` | `diffusion` | Stage type |
| `devices` | `"0"` | GPU device ID |
| `max_batch_size` | `1` | Maximum batch size |
| `model_stage` | `dit` | Model stage identifier |
| `gpu_memory_utilization` | `0.4` | GPU memory utilization |
| `omni_kv_config.need_recv_cache` | `true` | Whether to receive KV cache |
| `engine_input_source` | `[0]` | Input source from Stage 0 |
------
#### 🔗 Runtime Configuration
| Parameter | Value | Description |
| :-------------------- | :------ | :------------------------------- |
| `window_size` | `-1` | Window size (-1 means unlimited) |
| `max_inflight` | `1` | Maximum inflight requests |
| `shm_threshold_bytes` | `65536` | Shared memory threshold (64KB) |
## FAQ
- If you encounter an error about the backend of librosa, try to install ffmpeg with the command below.
```bash
sudo apt update
sudo apt install ffmpeg
```
- If you don’t know how much VRAM is needed for the model or encounter the OOM error, you can try to decrease the max_model_len.
| Stage | VRAM |
| :------------------ | :--------------------------- |
| Stage-0 (Thinker) | **15.04 GiB** **+ KV Cache** |
| Stage-1 (DiT) | **26.50 GiB** |
| Total | **~42 GiB + KV Cache** |
examples/offline_inference/bagel/end2end.py
+29 -19
View File
@@ -1,5 +1,8 @@
import argparse
import os
from typing import cast
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
def parse_args():
@@ -50,21 +53,24 @@ def parse_args():
def main():
args = parse_args()
model_name = args.model
prompts: list[OmniPromptType] = []
try:
# Preferred: load from txt file (one prompt per line)
if getattr(args, "txt_prompts", None) and args.prompt_type == "text":
with open(args.txt_prompts, encoding="utf-8") as f:
lines = [ln.strip() for ln in f.readlines()]
args.prompts = [ln for ln in lines if ln != ""]
print(f"[Info] Loaded {len(args.prompts)} prompts from {args.txt_prompts}")
prompts = [ln for ln in lines if ln != ""]
print(f"[Info] Loaded {len(prompts)} prompts from {args.txt_prompts}")
else:
prompts = args.prompts
except Exception as e:
print(f"[Error] Failed to load prompts: {e}")
raise
if args.prompts is None:
if not prompts:
# Default prompt for text2img test if none provided
args.prompts = ["<|im_start|>A cute cat<|im_end|>"]
print(f"[Info] No prompts provided, using default: {args.prompts}")
prompts = ["<|im_start|>A cute cat<|im_end|>"]
print(f"[Info] No prompts provided, using default: {prompts}")
omni_outputs = []
from PIL import Image
@@ -77,21 +83,27 @@ def main():
print("[Info] Running in img2img mode (Stage 1 only)")
client = OmniDiffusion(model=model_name)
generate_kwargs = {
"prompt": args.prompts,
"seed": 52,
"need_kv_receive": False,
"num_inference_steps": args.steps,
}
if args.image_path:
if os.path.exists(args.image_path):
loaded_image = Image.open(args.image_path).convert("RGB")
generate_kwargs["pil_image"] = loaded_image
prompts = [
{
"prompt": cast(str, p),
"multi_modal_data": {"image": loaded_image},
}
for p in prompts
]
else:
print(f"[Warning] Image path {args.image_path} does not exist.")
result = client.generate(**generate_kwargs)
result = client.generate(
prompts,
OmniDiffusionSamplingParams(
seed=52,
need_kv_receive=False,
num_inference_steps=args.steps,
),
)
# Ensure result is a list for iteration
if not isinstance(result, list):
@@ -100,8 +112,6 @@ def main():
omni_outputs = result
else:
import copy
from vllm_omni.entrypoints.omni import Omni
omni_kwargs = {}
@@ -144,11 +154,11 @@ def main():
prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]}
formatted_prompts.append(prompt_dict)
params_list = copy.deepcopy(omni.default_sampling_params_list)
params_list = omni.default_sampling_params_list
if args.modality == "text2img":
params_list[0]["max_tokens"] = 1
params_list[0].max_tokens = 1 # type: ignore # The first stage is a SamplingParam (vllm)
if len(params_list) > 1:
params_list[1]["num_inference_steps"] = args.steps
params_list[1].num_inference_steps = args.steps # type: ignore # The second stage is an OmniDiffusionSamplingParam
omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
examples/offline_inference/bagel/run_single_prompt.sh
-6
View File
@@ -1,6 +0,0 @@
prompt="<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n"
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--prompt_type text \
--init-sleep-seconds 0 \
--prompts ${prompt}
examples/offline_inference/image_to_image/image_edit.py
+33 -24
View File
@@ -77,10 +77,11 @@ from pathlib import Path
import torch
from PIL import Image
from vllm_omni.diffusion.data import DiffusionParallelConfig, logger
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.utils.platform_utils import detect_device_type, is_npu
from vllm_omni.platforms import current_omni_platform
def parse_args() -> argparse.Namespace:
@@ -279,6 +280,16 @@ def parse_args() -> argparse.Namespace:
action="store_true",
help="Disable torch.compile and force eager execution.",
)
parser.add_argument(
"--vae_use_slicing",
action="store_true",
help="Enable VAE slicing for memory optimization.",
)
parser.add_argument(
"--vae_use_tiling",
action="store_true",
help="Enable VAE tiling for memory optimization.",
)
parser.add_argument(
"--enable-cpu-offload",
action="store_true",
@@ -305,12 +316,8 @@ def main():
else:
input_image = input_images
device = detect_device_type()
generator = torch.Generator(device=device).manual_seed(args.seed)
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
# Enable VAE memory optimizations on NPU
vae_use_slicing = is_npu()
vae_use_tiling = is_npu()
parallel_config = DiffusionParallelConfig(
ulysses_degree=args.ulysses_degree,
ring_degree=args.ring_degree,
@@ -343,8 +350,8 @@ def main():
# Initialize Omni with appropriate pipeline
omni = Omni(
model=args.model,
vae_use_slicing=vae_use_slicing,
vae_use_tiling=vae_use_tiling,
vae_use_slicing=args.vae_use_slicing,
vae_use_tiling=args.vae_use_tiling,
cache_backend=args.cache_backend,
cache_config=cache_config,
parallel_config=parallel_config,
@@ -374,25 +381,28 @@ def main():
print(f"{'=' * 60}\n")
generation_start = time.perf_counter()
# Generate edited image
generate_kwargs = {
"prompt": args.prompt,
"pil_image": input_image,
"negative_prompt": args.negative_prompt,
"generator": generator,
"true_cfg_scale": args.cfg_scale,
"guidance_scale": args.guidance_scale,
"num_inference_steps": args.num_inference_steps,
"num_outputs_per_prompt": args.num_outputs_per_prompt,
"layers": args.layers,
"resolution": args.resolution,
}
if profiler_enabled:
print("[Profiler] Starting profiling...")
omni.start_profile()
outputs = omni.generate(**generate_kwargs)
# Generate edited image
outputs = omni.generate(
{
"prompt": args.prompt,
"negative_prompt": args.negative_prompt,
"multi_modal_data": {"image": input_image},
},
OmniDiffusionSamplingParams(
generator=generator,
true_cfg_scale=args.cfg_scale,
guidance_scale=args.guidance_scale,
num_inference_steps=args.num_inference_steps,
num_outputs_per_prompt=args.num_outputs_per_prompt,
layers=args.layers,
resolution=args.resolution,
),
)
generation_end = time.perf_counter()
generation_time = generation_end - generation_start
@@ -418,7 +428,6 @@ def main():
if not outputs:
raise ValueError("No output generated from omni.generate()")
logger.info("Outputs: %s", outputs)
# Extract images from OmniRequestOutput
# omni.generate() returns list[OmniRequestOutput], extract images from request_output[0].images
examples/offline_inference/image_to_image/image_to_image.md
+4
View File
@@ -47,4 +47,8 @@ Key arguments:
- `--guidance_scale`: guidance scale for guidance-distilled models (default: 1.0, disabled). Unlike classifier-free guidance (--cfg_scale), guidance-distilled models take the guidance scale directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models.
- `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
- `--output`: path to save the generated PNG.
- `--vae_use_slicing`: enable VAE slicing for memory optimization.
- `--vae_use_tiling`: enable VAE tiling for memory optimization.
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
examples/offline_inference/image_to_video/README.md
+4
View File
@@ -54,4 +54,8 @@ Key arguments:
- `--num_inference_steps`: Number of denoising steps (default 50).
- `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video).
- `--output`: Path to save the generated video.
- `--vae_use_slicing`: Enable VAE slicing for memory optimization.
- `--vae_use_tiling`: Enable VAE tiling for memory optimization.
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
examples/offline_inference/image_to_video/image_to_video.py
+29 -19
View File
@@ -27,8 +27,9 @@ import PIL.Image
import torch
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.utils.platform_utils import detect_device_type, is_npu
from vllm_omni.platforms import current_omni_platform
def parse_args() -> argparse.Namespace:
@@ -58,6 +59,16 @@ def parse_args() -> argparse.Namespace:
)
parser.add_argument("--output", type=str, default="i2v_output.mp4", help="Path to save the video (mp4).")
parser.add_argument("--fps", type=int, default=16, help="Frames per second for the output video.")
parser.add_argument(
"--vae_use_slicing",
action="store_true",
help="Enable VAE slicing for memory optimization.",
)
parser.add_argument(
"--vae_use_tiling",
action="store_true",
help="Enable VAE tiling for memory optimization.",
)
parser.add_argument(
"--enable-cpu-offload",
action="store_true",
@@ -79,8 +90,7 @@ def calculate_dimensions(image: PIL.Image.Image, max_area: int = 480 * 832) -> t
def main():
args = parse_args()
device = detect_device_type()
generator = torch.Generator(device=device).manual_seed(args.seed)
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
# Load input image
image = PIL.Image.open(args.image).convert("RGB")
@@ -97,17 +107,13 @@ def main():
# Resize image to target dimensions
image = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
# Enable VAE memory optimizations on NPU
vae_use_slicing = is_npu()
vae_use_tiling = is_npu()
# Check if profiling is requested via environment variable
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
omni = Omni(
model=args.model,
vae_use_slicing=vae_use_slicing,
vae_use_tiling=vae_use_tiling,
vae_use_slicing=args.vae_use_slicing,
vae_use_tiling=args.vae_use_tiling,
boundary_ratio=args.boundary_ratio,
flow_shift=args.flow_shift,
enable_cpu_offload=args.enable_cpu_offload,
@@ -119,16 +125,20 @@ def main():
# omni.generate() returns Generator[OmniRequestOutput, None, None]
frames = omni.generate(
args.prompt,
negative_prompt=args.negative_prompt,
pil_image=image,
height=height,
width=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,
{
"prompt": args.prompt,
"negative_prompt": args.negative_prompt,
"multi_modal_data": {"image": image},
},
OmniDiffusionSamplingParams(
height=height,
width=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,
),
)
# Extract video frames from OmniRequestOutput
examples/offline_inference/lora_inference/lora_inference.py
+9 -9
View File
@@ -5,6 +5,7 @@ import argparse
from pathlib import Path
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.lora.request import LoRARequest
from vllm_omni.lora.utils import stable_lora_int_id
@@ -95,18 +96,17 @@ def main():
)
print(f"Activating pre-loaded LoRA: id={lora_request_id}, scale={args.lora_scale}")
gen_kwargs = {
"prompt": args.prompt,
"height": args.height,
"width": args.width,
"num_inference_steps": args.num_inference_steps,
}
sampling_params = OmniDiffusionSamplingParams(
height=args.height,
width=args.width,
num_inference_steps=args.num_inference_steps,
)
if lora_request:
gen_kwargs["lora_request"] = lora_request
gen_kwargs["lora_scale"] = args.lora_scale
sampling_params.lora_request = lora_request
sampling_params.lora_scale = args.lora_scale
outputs = omni.generate(**gen_kwargs)
outputs = omni.generate(args.prompt, sampling_params)
if not outputs or len(outputs) == 0:
raise ValueError("No output generated from omni.generate()")
examples/offline_inference/text_to_audio/text_to_audio.py
+16 -12
View File
@@ -21,7 +21,8 @@ import numpy as np
import torch
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.utils.platform_utils import detect_device_type
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.platforms import current_omni_platform
def parse_args() -> argparse.Namespace:
@@ -117,8 +118,7 @@ def save_audio(audio_data: np.ndarray, output_path: str, sample_rate: int = 4410
def main():
args = parse_args()
device = detect_device_type()
generator = torch.Generator(device=device).manual_seed(args.seed)
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
print(f"\n{'=' * 60}")
print("Stable Audio Open - Text-to-Audio Generation")
@@ -143,16 +143,20 @@ def main():
# Generate audio
outputs = omni.generate(
args.prompt,
negative_prompt=args.negative_prompt,
generator=generator,
guidance_scale=args.guidance_scale,
num_inference_steps=args.num_inference_steps,
num_outputs_per_prompt=args.num_waveforms,
extra={
"audio_start_in_s": args.audio_start,
"audio_end_in_s": audio_end_in_s,
{
"prompt": args.prompt,
"negative_prompt": args.negative_prompt,
},
OmniDiffusionSamplingParams(
generator=generator,
guidance_scale=args.guidance_scale,
num_inference_steps=args.num_inference_steps,
num_outputs_per_prompt=args.num_waveforms,
extra_args={
"audio_start_in_s": args.audio_start,
"audio_end_in_s": audio_end_in_s,
},
),
)
generation_end = time.perf_counter()
examples/offline_inference/text_to_image/README.md
+40 -1
View File
@@ -20,7 +20,7 @@ if __name__ == "__main__":
images[0].save("coffee.png")
```
Or put more than one prompt in a request, processing them sequentially.
Or put more than one prompt in a request.
```python
from vllm_omni.entrypoints.omni import Omni
@@ -37,6 +37,41 @@ if __name__ == "__main__":
image = output.request_output[0].images[0].save(f"{i}.jpg")
```
!!! info
However, it is not currently recommended to do so
because not all models support batch inference,
and batch requesting mostly does not provide significant performance improvement (despite the impression that it does).
This feature is primarily for the sake of interface compatibility with vLLM and to allow for future improvements.
!!! info
For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
batched inputs, you can [modify this configuration](../../../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
Apart from string prompt, vLLM-Omni also supports dictionary prompts in the same style as vLLM.
This is useful for models that support negative prompts.
```python
from vllm_omni.entrypoints.omni import Omni
if __name__ == "__main__":
omni = Omni(model="Qwen/Qwen-Image")
outputs = omni.generate([
{
"prompt": "a cup of coffee on a table"，
"negative_prompt": "low resolution"
},
{
"prompt": "a toy dinosaur on a sandy beach"，
"negative_prompt": "cinematic, realistic"
}
])
for i, output in enumerate(outputs):
image = output.request_output[0].images[0].save(f"{i}.jpg")
```
## Local CLI Usage
```bash
@@ -61,8 +96,12 @@ Key arguments:
- `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
- `--height/--width`: output resolution (defaults 1024x1024).
- `--output`: path to save the generated PNG.
- `--vae_use_slicing`: enable VAE slicing for memory optimization.
- `--vae_use_tiling`: enable VAE tiling for memory optimization.
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
> ℹ️ 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
examples/offline_inference/text_to_image/gradio_demo.py
+13 -11
View File
@@ -5,8 +5,9 @@ import gradio as gr
import torch
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.utils.platform_utils import detect_device_type, is_npu
from vllm_omni.platforms import current_omni_platform
ASPECT_RATIOS: dict[str, tuple[int, int]] = {
"1:1": (1328, 1328),
@@ -61,8 +62,8 @@ def parse_args() -> argparse.Namespace:
@lru_cache(maxsize=1)
def get_omni(model_name: str) -> Omni:
# Enable VAE memory optimizations on NPU
vae_use_slicing = is_npu()
vae_use_tiling = is_npu()
vae_use_slicing = current_omni_platform.is_npu()
vae_use_tiling = current_omni_platform.is_npu()
return Omni(
model=model_name,
vae_use_slicing=vae_use_slicing,
@@ -71,7 +72,6 @@ def get_omni(model_name: str) -> Omni:
def build_demo(args: argparse.Namespace) -> gr.Blocks:
device = detect_device_type()
omni = get_omni(args.model)
def run_inference(
@@ -98,15 +98,17 @@ def build_demo(args: argparse.Namespace) -> gr.Blocks:
raise gr.Error("Inference steps must be a positive integer.")
if num_images not in {1, 2, 3, 4}:
raise gr.Error("Number of images must be 1, 2, 3, or 4.")
generator = torch.Generator(device=device).manual_seed(seed)
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
outputs = omni.generate(
prompt.strip(),
height=height,
width=width,
generator=generator,
true_cfg_scale=float(cfg_scale_value),
num_inference_steps=num_steps,
num_outputs_per_prompt=num_images,
OmniDiffusionSamplingParams(
height=height,
width=width,
generator=generator,
true_cfg_scale=float(cfg_scale_value),
num_inference_steps=num_steps,
num_outputs_per_prompt=num_images,
),
)
images_outputs = []
for output in outputs:
examples/offline_inference/text_to_image/text_to_image.py
+34 -18
View File
@@ -10,8 +10,9 @@ import torch
from vllm_omni.diffusion.data import DiffusionParallelConfig, logger
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.utils.platform_utils import detect_device_type, is_npu
from vllm_omni.platforms import current_omni_platform
def parse_args() -> argparse.Namespace:
@@ -72,6 +73,11 @@ def parse_args() -> argparse.Namespace:
"Default: None (no cache acceleration)."
),
)
parser.add_argument(
"--enable-cache-dit-summary",
action="store_true",
help="Enable cache-dit summary logging after diffusion forward passes.",
)
parser.add_argument(
"--ulysses_degree",
type=int,
@@ -107,17 +113,22 @@ def parse_args() -> argparse.Namespace:
default=1,
help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
)
parser.add_argument(
"--vae_use_slicing",
action="store_true",
help="Enable VAE slicing for memory optimization.",
)
parser.add_argument(
"--vae_use_tiling",
action="store_true",
help="Enable VAE tiling for memory optimization.",
)
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()
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
# Configure cache based on backend type
cache_config = None
@@ -161,10 +172,11 @@ def main():
omni = Omni(
model=args.model,
vae_use_slicing=vae_use_slicing,
vae_use_tiling=vae_use_tiling,
vae_use_slicing=args.vae_use_slicing,
vae_use_tiling=args.vae_use_tiling,
cache_backend=args.cache_backend,
cache_config=cache_config,
enable_cache_dit_summary=args.enable_cache_dit_summary,
parallel_config=parallel_config,
enforce_eager=args.enforce_eager,
enable_cpu_offload=args.enable_cpu_offload,
@@ -189,15 +201,19 @@ def main():
generation_start = time.perf_counter()
outputs = omni.generate(
args.prompt,
negative_prompt=args.negative_prompt,
height=args.height,
width=args.width,
generator=generator,
true_cfg_scale=args.cfg_scale,
guidance_scale=args.guidance_scale,
num_inference_steps=args.num_inference_steps,
num_outputs_per_prompt=args.num_images_per_prompt,
{
"prompt": args.prompt,
"negative_prompt": args.negative_prompt,
},
OmniDiffusionSamplingParams(
height=args.height,
width=args.width,
generator=generator,
true_cfg_scale=args.cfg_scale,
guidance_scale=args.guidance_scale,
num_inference_steps=args.num_inference_steps,
num_outputs_per_prompt=args.num_images_per_prompt,
),
)
generation_end = time.perf_counter()
generation_time = generation_end - generation_start
examples/offline_inference/text_to_video/text_to_video.md
+5 -1
View File
@@ -26,7 +26,11 @@ Key arguments:
- `--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.
- `--boundary_ratio`: Boundary split ratio for low/high DiT. Default `0.875` uses both transformers for best quality. Set to `1.0` to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited). Set to `0.0` loads only the high-noise transformer (not recommended, lower quality).
- `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video).
- `--output`: path to save the generated video.
- `--vae_use_slicing`: enable VAE slicing for memory optimization.
- `--vae_use_tiling`: enable VAE tiling for memory optimization.
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
examples/offline_inference/text_to_video/text_to_video.py
+76 -24
View File
@@ -11,8 +11,9 @@ import torch
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.utils.platform_utils import detect_device_type, is_npu
from vllm_omni.platforms import current_omni_platform
def parse_args() -> argparse.Namespace:
@@ -31,12 +32,48 @@ def parse_args() -> argparse.Namespace:
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(
"--boundary_ratio",
type=float,
default=0.875,
help="Boundary split ratio for low/high DiT. Default 0.875 uses both transformers for best quality. Set to 1.0 to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited).",
)
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(
"--cache_backend",
type=str,
default=None,
choices=["cache_dit"],
help=(
"Cache backend to use for acceleration. "
"Options: 'cache_dit' (DBCache + SCM + TaylorSeer). "
"Default: None (no cache acceleration)."
),
)
parser.add_argument(
"--enable-cache-dit-summary",
action="store_true",
help="Enable cache-dit summary logging after diffusion forward passes.",
)
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.")
parser.add_argument(
"--vae_use_slicing",
action="store_true",
help="Enable VAE slicing for memory optimization.",
)
parser.add_argument(
"--vae_use_tiling",
action="store_true",
help="Enable VAE tiling for memory optimization.",
)
parser.add_argument(
"--enforce_eager",
action="store_true",
help="Disable torch.compile and force eager execution.",
)
parser.add_argument(
"--enable-cpu-offload",
action="store_true",
@@ -54,23 +91,31 @@ def parse_args() -> argparse.Namespace:
default=1,
help="Number of GPUs used for ring sequence parallelism.",
)
parser.add_argument(
"--enforce_eager",
action="store_true",
help="Disable torch.compile and force eager execution.",
)
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()
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
# Wan2.2 cache-dit tuning (from cache-dit examples and cache_alignment).
cache_config = None
if args.cache_backend == "cache_dit":
cache_config = {
# DBCache parameters [cache-dit only]
"Fn_compute_blocks": 1, # Optimized for single-transformer models
"Bn_compute_blocks": 0, # Number of backward compute blocks
"max_warmup_steps": 4, # Maximum warmup steps (works for few-step models)
"max_cached_steps": 20,
"residual_diff_threshold": 0.24, # Higher threshold for more aggressive caching
"max_continuous_cached_steps": 3, # Limit to prevent precision degradation
# TaylorSeer parameters [cache-dit only]
"enable_taylorseer": False, # Disabled by default (not suitable for few-step models)
"taylorseer_order": 1, # TaylorSeer polynomial order
# SCM (Step Computation Masking) parameters [cache-dit only]
"scm_steps_mask_policy": None, # SCM mask policy: None (disabled), "slow", "medium", "fast", "ultra"
"scm_steps_policy": "dynamic", # SCM steps policy: "dynamic" or "static"
}
# Configure parallel settings (only SP is supported for Wan)
# Note: cfg_parallel and tensor_parallel are not implemented for Wan models
parallel_config = DiffusionParallelConfig(
@@ -83,10 +128,13 @@ def main():
omni = Omni(
model=args.model,
vae_use_slicing=vae_use_slicing,
vae_use_tiling=vae_use_tiling,
vae_use_slicing=args.vae_use_slicing,
vae_use_tiling=args.vae_use_tiling,
boundary_ratio=args.boundary_ratio,
flow_shift=args.flow_shift,
cache_backend=args.cache_backend,
cache_config=cache_config,
enable_cache_dit_summary=args.enable_cache_dit_summary,
enable_cpu_offload=args.enable_cpu_offload,
parallel_config=parallel_config,
enforce_eager=args.enforce_eager,
@@ -108,15 +156,19 @@ def main():
generation_start = time.perf_counter()
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,
{
"prompt": args.prompt,
"negative_prompt": args.negative_prompt,
},
OmniDiffusionSamplingParams(
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,
),
)
generation_end = time.perf_counter()
generation_time = generation_end - generation_start
examples/online_serving/bagel/README.md
+230
View File
@@ -0,0 +1,230 @@
# BAGEL-7B-MoT
## 🛠️ Installation
Please refer to [README.md](../../../README.md)
## Run examples (BAGEL-7B-MoT)
**Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, please modify the stage configuration to distribute the model across devices.
### Launch the Server
```bash
# Use default configuration
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091
```
Or use the convenience script:
```bash
cd /workspace/vllm-omni/examples/online_serving/bagel
bash run_server.sh
```
If you have a custom stage configs file, launch the server with the command below:
```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
```
### Send Multi-modal Request
Get into the bagel folder:
```bash
cd examples/online_serving/bagel
```
Send request via Python
```bash
python openai_chat_client.py --prompt "A cute cat" --modality text2img
```
The Python client supports the following command-line arguments:
- `--prompt` (or `-p`): Text prompt for generation (default: `A cute cat`)
- `--output` (or `-o`): Output file path for image results (default: `bagel_output.png`)
- `--server` (or `-s`): Server URL (default: `http://localhost:8091`)
- `--image-url` (or `-i`): Input image URL or local file path (for img2img/img2text modes)
- `--modality` (or `-m`): Task modality (default: `text2img`). Options: `text2img`, `img2img`, `img2text`, `text2text`
- `--height`: Image height in pixels (default: 512)
- `--width`: Image width in pixels (default: 512)
- `--steps`: Number of inference steps (default: 25)
- `--seed`: Random seed (default: 42)
- `--negative`: Negative prompt for image generation
Example with custom parameters:
```bash
python openai_chat_client.py \
--prompt "A futuristic city" \
--modality text2img \
--height 768 \
--width 768 \
--steps 50 \
--seed 42 \
--negative "blurry, low quality"
```
## Modality Control
BAGEL-7B-MoT supports **multiple modality modes** for different use cases.
The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)
| Modality | Input | Output | Description |
| ----------- | ------------ | ------ | -------------------------------------- |
| `text2img` | Text | Image | Generate images from text prompts |
| `img2img` | Image + Text | Image | Transform images using text guidance |
| `img2text` | Image + Text | Text | Generate text descriptions from images |
| `text2text` | Text | Text | Pure text generation |
### Text to Image (text2img)
Generate images from text prompts:
**Using Python client**
```bash
python openai_chat_client.py \
--prompt "A beautiful sunset over mountains" \
--modality text2img \
--output sunset.png \
--steps 50
```
**Using curl**
```bash
curl http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>A beautiful sunset over mountains<|im_end|>"}]}],
"modalities": ["image"],
"height": 512,
"width": 512,
"num_inference_steps": 50,
"seed": 42
}'
```
### Image to Image (img2img)
Transform images based on text prompts:
**Using Python client**
```bash
python openai_chat_client.py \
--prompt "Make the cat stand up" \
--modality img2img \
--image-url /path/to/input.jpg \
--output transformed.png
```
**Using curl**
```bash
IMAGE_BASE64=$(base64 -w 0 cat.jpg)
cat <<EOF > payload.json
{
"messages": [{
"role": "user",
"content": [
{"type": "text", "text": "<|im_start|>Make the cat stand up<|im_end|>"},
{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}}
]
}],
"modalities": ["image"],
"height": 512,
"width": 512,
"num_inference_steps": 50,
"seed": 42
}
EOF
curl http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d @payload.json
```
### Image to Text (img2text)
Generate text descriptions from images:
**Using Python client**
```bash
python openai_chat_client.py \
--prompt "Describe this image in detail" \
--modality img2text \
--image-url /path/to/image.jpg
```
**Using curl**
```bash
IMAGE_BASE64=$(base64 -w 0 cat.jpg)
cat <<EOF > payload.json
{
"messages": [{
"role": "user",
"content": [
{"type": "text", "text": "<|im_start|>user\n<|image_pad|>\nDescribe this image in detail<|im_end|>\n<|im_start|>assistant\n"},
{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}}
]
}],
"modalities": ["text"]
}
EOF
curl http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d @payload.json
```
### Text to Text (text2text)
Pure text generation:
**Using Python client**
```bash
python openai_chat_client.py \
--prompt "What is the capital of France?" \
--modality text2text
```
**Using curl**
```bash
curl http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n"}]}]
"modalities": ["text"]
}'
```
## FAQ
- If you encounter an error about the backend of librosa, try to install ffmpeg with the command below.
```bash
sudo apt update
sudo apt install ffmpeg
```
- If you don’t know how much VRAM is needed for the model or encounter the OOM error, you can try to decrease the max_model_len.
| Stage | VRAM |
| :------------------ | :--------------------------- |
| Stage-0 (Thinker) | **15.04 GiB** **+ KV Cache** |
| Stage-1 (DiT) | **26.50 GiB** |
| Total | **~42 GiB + KV Cache** |
examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml
+1 -1
View File
@@ -13,7 +13,7 @@ stage_args:
engine_args:
model_stage: thinker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.9
enforce_eager: true
examples/online_serving/qwen3_tts/README.md
+2
View File
@@ -41,12 +41,14 @@ python openai_speech_client.py \
# VoiceDesign: Describe the voice style
python openai_speech_client.py \
--model Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign \
--task-type VoiceDesign \
--text "哥哥，你回来啦" \
--instructions "体现撒娇稚嫩的萝莉女声，音调偏高"
# Base: Voice cloning
python openai_speech_client.py \
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
--task-type Base \
--text "Hello, this is a cloned voice" \
--ref-audio /path/to/reference.wav \
examples/online_serving/qwen3_tts/openai_speech_client.py
+1
View File
@@ -90,6 +90,7 @@ def run_tts_generation(args) -> None:
if args.x_vector_only:
payload["x_vector_only_mode"] = True
print(f"Model: {args.model}")
print(f"Task type: {args.task_type or 'CustomVoice'}")
print(f"Text: {args.text}")
print(f"Voice: {args.voice}")
pyproject.toml
+3 -1
View File
@@ -39,9 +39,11 @@ dependencies = [
"cache-dit==1.2.0",
"tqdm>=4.66.0",
"torchsde>=0.2.6", # Required for Stable Audio scheduler
"fa3-fwd", # flash attention 3, maintained by @ZJY0516
"fa3-fwd==0.0.1", # flash attention 3, maintained by @ZJY0516
"openai-whisper>=20250625",
"imageio[ffmpeg]>=2.37.2",
"onnxruntime>=1.19.0",
"sox>=1.5.0",
# "vllm==0.14.0", # TODO: fix the entrypoints overwrite problem
]
tests/conftest.py
+42 -27
View File
@@ -10,6 +10,7 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
if "VLLM_TARGET_DEVICE" not in os.environ:
os.environ["VLLM_TARGET_DEVICE"] = "cpu"
import gc
import socket
import subprocess
import sys
@@ -59,13 +60,12 @@ def clean_gpu_memory_between_tests():
_run_post_test_cleanup()
def _run_pre_test_cleanup():
if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1":
def _run_pre_test_cleanup(enable_force=False):
if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1" and not enable_force:
print("GPU cleanup disabled")
return
print("Pre-test GPU status:")
_print_simple_gpu_status()
num_gpus = torch.cuda.device_count()
if num_gpus > 0:
@@ -74,44 +74,25 @@ def _run_pre_test_cleanup():
wait_for_gpu_memory_to_clear(
devices=list(range(num_gpus)),
threshold_ratio=0.1,
threshold_ratio=0.05,
)
except Exception as e:
print(f"Pre-test cleanup note: {e}")
def _run_post_test_cleanup():
if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1":
def _run_post_test_cleanup(enable_force=False):
if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1" and not enable_force:
print("GPU cleanup disabled")
return
import gc
if torch.cuda.is_available():
gc.collect()
torch.cuda.empty_cache()
print("Post-test GPU status:")
_print_simple_gpu_status()
_print_gpu_processes()
def _print_simple_gpu_status():
"""Print simple GPU memory status"""
if not torch.cuda.is_available():
print(" CUDA not available")
return
num_devices = torch.cuda.device_count()
for device_id in range(num_devices):
try:
torch.cuda.set_device(device_id)
allocated = torch.cuda.memory_allocated(device_id) / (1024**2)
reserved = torch.cuda.memory_reserved(device_id) / (1024**2)
print(f" GPU {device_id}: Allocated: {allocated:.1f}MB, Reserved: {reserved:.1f}MB")
except Exception:
print(f" GPU {device_id}: Error reading status")
def _print_gpu_processes():
"""Print GPU information including nvidia-smi and system processes"""
@@ -623,6 +604,36 @@ def convert_audio_to_text(audio_data):
return ""
def merge_base64_and_convert_to_text(base64_list):
"""
Merge a list of base64 encoded audio data and convert to text.
"""
import whisper
from pydub import AudioSegment
merged_audio = None
for base64_data in base64_list:
audio_data = base64.b64decode(base64_data.split(",", 1)[-1])
seg = AudioSegment.from_file(io.BytesIO(audio_data))
if merged_audio is None:
merged_audio = seg
else:
merged_audio += seg
output_path = f"./test_{int(time.time())}"
merged_audio.export(output_path, format="wav")
model = whisper.load_model("base")
text = model.transcribe(
output_path,
temperature=0.0,
word_timestamps=True,
condition_on_previous_text=False,
)["text"]
if text:
return text
else:
return ""
def modify_stage_config(
yaml_path: str,
updates: dict[str, Any],
@@ -871,6 +882,9 @@ class OmniServer:
*,
env_dict: dict[str, str] | None = None,
) -> None:
_run_pre_test_cleanup(enable_force=True)
_run_post_test_cleanup(enable_force=True)
cleanup_dist_env_and_memory()
self.model = model
self.serve_args = serve_args
self.env_dict = env_dict
@@ -986,5 +1000,6 @@ class OmniServer:
def __exit__(self, exc_type, exc_val, exc_tb):
if self.proc:
self._kill_process_tree(self.proc.pid)
_run_post_test_cleanup()
_run_pre_test_cleanup(enable_force=True)
_run_post_test_cleanup(enable_force=True)
cleanup_dist_env_and_memory()
tests/diffusion/attention/test_attention_sp.py
+6 -15
View File
@@ -6,7 +6,6 @@ import tempfile
import pytest
import torch
from vllm.platforms import current_platform
from vllm_omni.diffusion.attention.layer import Attention
from vllm_omni.diffusion.data import (
@@ -19,15 +18,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
initialize_model_parallel,
)
from vllm_omni.diffusion.forward_context import set_forward_context
from vllm_omni.utils.platform_utils import detect_device_type
device_type = detect_device_type()
if device_type == "cuda":
torch_device = torch.cuda
elif device_type == "npu":
torch_device = torch.npu
else:
raise ValueError(f"Unsupported device type: {device_type} for this test script! Expected GPU or NPU.")
from vllm_omni.platforms import current_omni_platform
def update_environment_variables(envs_dict: dict[str, str]):
@@ -171,7 +162,7 @@ def test_sequence_parallel(
sequence_parallel_size = ulysses_degree * ring_degree
# Skip if not enough GPUs available
available_gpus = torch_device.device_count()
available_gpus = current_omni_platform.get_device_count()
if available_gpus < sequence_parallel_size:
pytest.skip(f"Test requires {sequence_parallel_size} GPUs but only {available_gpus} available")
@@ -338,13 +329,13 @@ def ulysses_attention_on_test_model(
"""Run Ulysses attention test on a test model and save results for comparison."""
# Use fixed seed for reproducibility across baseline and SP runs
RANDOM_SEED = 42
current_platform.seed_everything(RANDOM_SEED)
current_omni_platform.seed_everything(RANDOM_SEED)
mode_str = "Baseline (no SP)" if is_baseline else f"SP (ulysses={ulysses_degree}, ring={ring_degree})"
print(f"\n[{mode_str}] Rank {local_rank}/{world_size} - Random seed set to {RANDOM_SEED}")
device = torch.device(f"{device_type}:{local_rank}")
torch_device.set_device(device)
device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
current_omni_platform.set_device(device)
torch.set_default_device(device)
torch.set_default_dtype(dtype)
@@ -423,7 +414,7 @@ def ulysses_attention_on_test_model(
# Generate and save full input data with fixed seed
# Reinitialize RNG to ensure reproducibility
torch.manual_seed(42)
torch_device.manual_seed_all(42)
current_omni_platform.seed_everything(42)
full_hidden_states = torch.randn(
(batch_size, seq_len, hidden_size),
dtype=dtype,
tests/diffusion/distributed/test_comm.py
+10 -18
View File
@@ -14,15 +14,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
init_distributed_environment,
initialize_model_parallel,
)
from vllm_omni.utils.platform_utils import detect_device_type
device_type = detect_device_type()
if device_type == "cuda":
torch_device = torch.cuda
elif device_type == "npu":
torch_device = torch.npu
else:
raise ValueError(f"Unsupported device type: {device_type} for this test script! Expected GPU or NPU.")
from vllm_omni.platforms import current_omni_platform
def update_environment_variables(envs_dict: dict[str, str]):
@@ -49,7 +41,7 @@ def test_4d_identity(
):
"""Test that two consecutive all-to-all operations return the original input."""
# Skip if not enough GPUs available
available_gpus = torch_device.device_count()
available_gpus = current_omni_platform.get_device_count()
if available_gpus < world_size:
pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
@@ -85,8 +77,8 @@ def _test_4d_identity_worker(
):
"""Worker function for test_4d_identity."""
# Set device
device = torch.device(f"{device_type}:{local_rank}")
torch_device.set_device(device)
device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
current_omni_platform.set_device(device)
# Set environment variables for distributed training
update_environment_variables(
@@ -183,7 +175,7 @@ def test_5d_identity(
):
"""Test that two consecutive all-to-all operations return the original input."""
# Skip if not enough GPUs available
available_gpus = torch_device.device_count()
available_gpus = current_omni_platform.get_device_count()
if available_gpus < world_size:
pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
@@ -219,8 +211,8 @@ def _test_5d_identity_worker(
):
"""Worker function for test_5d_identity."""
# Set device
device = torch.device(f"{device_type}:{local_rank}")
torch_device.set_device(device)
device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
current_omni_platform.set_device(device)
# Set environment variables for distributed training
update_environment_variables(
@@ -316,7 +308,7 @@ def test_ring_p2p(
):
"""Test Ring P2P communication (send_recv)."""
# Skip if not enough GPUs available
available_gpus = torch_device.device_count()
available_gpus = current_omni_platform.get_device_count()
if available_gpus < world_size:
pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
@@ -339,8 +331,8 @@ def _test_ring_p2p_worker(
import sys
# Set device
device = torch.device(f"{device_type}:{local_rank}")
torch_device.set_device(device)
device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
current_omni_platform.set_device(device)
# Set env vars
# Use a different port to avoid conflict with other tests if run in parallel
tests/diffusion/test_gpu_diffusion_worker.py → tests/diffusion/test_diffusion_worker.py
+33 -29
View File
@@ -2,9 +2,9 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Unit tests for GPUDiffusionWorker class.
Unit tests for DiffusionWorker class.
This module tests the GPUDiffusionWorker implementation:
This module tests the DiffusionWorker implementation:
- load_weights: Loading model weights
- sleep: Putting worker into sleep mode (levels 1 and 2)
- wake_up: Waking worker from sleep mode
@@ -15,7 +15,7 @@ from unittest.mock import Mock, patch
import pytest
import torch
from vllm_omni.diffusion.worker.gpu_diffusion_worker import GPUDiffusionWorker
from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
@pytest.fixture
@@ -33,18 +33,19 @@ def mock_od_config():
@pytest.fixture
def mock_gpu_worker(mock_od_config):
"""Create a GPUDiffusionWorker with mocked initialization."""
with patch.object(GPUDiffusionWorker, "init_device"):
worker = GPUDiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config)
"""Create a DiffusionWorker with mocked initialization."""
with patch.object(DiffusionWorker, "init_device"):
worker = DiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config)
# Mock the model_runner with pipeline
worker.model_runner = Mock()
worker.model_runner.pipeline = Mock()
worker.device = torch.device("cuda", 0)
worker._sleep_saved_buffers = {}
return worker
class TestGPUDiffusionWorkerLoadWeights:
"""Test GPUDiffusionWorker.load_weights method."""
class TestDiffusionWorkerLoadWeights:
"""Test DiffusionWorker.load_weights method."""
def test_load_weights_calls_pipeline(self, mock_gpu_worker):
"""Test that load_weights delegates to model_runner.load_weights."""
@@ -75,20 +76,21 @@ class TestGPUDiffusionWorkerLoadWeights:
assert result == set()
class TestGPUDiffusionWorkerSleep:
"""Test GPUDiffusionWorker.sleep method."""
class TestDiffusionWorkerSleep:
"""Test DiffusionWorker.sleep method."""
@patch("vllm_omni.diffusion.worker.gpu_diffusion_worker.torch.cuda.mem_get_info")
@patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
@patch("vllm.device_allocator.cumem.CuMemAllocator")
def test_sleep_level_1(self, mock_allocator_class, mock_mem_info, mock_gpu_worker):
def test_sleep_level_1(self, mock_allocator_class, mock_platform, mock_gpu_worker):
"""Test sleep mode level 1 (offload weights only)."""
# Setup memory info mocks
# Before sleep: 1GB free, 8GB total
# After sleep: 3GB free, 8GB total (freed 2GB)
mock_mem_info.side_effect = [
(1 * 1024**3, 8 * 1024**3), # Before sleep
(3 * 1024**3, 8 * 1024**3), # After sleep
# Before sleep: 1GB free
# After sleep: 3GB free (freed 2GB)
mock_platform.get_free_memory.side_effect = [
1 * 1024**3, # Before sleep
3 * 1024**3, # After sleep
]
mock_platform.get_device_total_memory.return_value = 8 * 1024**3
# Setup allocator mock
mock_allocator = Mock()
@@ -104,15 +106,16 @@ class TestGPUDiffusionWorkerSleep:
# Verify buffers were NOT saved (level 1 doesn't save buffers)
assert len(mock_gpu_worker._sleep_saved_buffers) == 0
@patch("vllm_omni.diffusion.worker.gpu_diffusion_worker.torch.cuda.mem_get_info")
@patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
@patch("vllm.device_allocator.cumem.CuMemAllocator")
def test_sleep_level_2(self, mock_allocator_class, mock_mem_info, mock_gpu_worker):
def test_sleep_level_2(self, mock_allocator_class, mock_platform, mock_gpu_worker):
"""Test sleep mode level 2 (offload all, save buffers)."""
# Setup memory info mocks
mock_mem_info.side_effect = [
(1 * 1024**3, 8 * 1024**3), # Before sleep
(5 * 1024**3, 8 * 1024**3), # After sleep (freed 4GB)
mock_platform.get_free_memory.side_effect = [
1 * 1024**3, # Before sleep
5 * 1024**3, # After sleep (freed 4GB)
]
mock_platform.get_device_total_memory.return_value = 8 * 1024**3
# Setup allocator mock
mock_allocator = Mock()
@@ -141,15 +144,16 @@ class TestGPUDiffusionWorkerSleep:
assert "buffer1" in mock_gpu_worker._sleep_saved_buffers
assert "buffer2" in mock_gpu_worker._sleep_saved_buffers
@patch("vllm_omni.diffusion.worker.gpu_diffusion_worker.torch.cuda.mem_get_info")
@patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
@patch("vllm.device_allocator.cumem.CuMemAllocator")
def test_sleep_memory_freed_validation(self, mock_allocator_class, mock_mem_info, mock_gpu_worker):
def test_sleep_memory_freed_validation(self, mock_allocator_class, mock_platform, mock_gpu_worker):
"""Test that sleep validates memory was actually freed."""
# Simulate memory increase (should trigger assertion error)
mock_mem_info.side_effect = [
(3 * 1024**3, 8 * 1024**3), # Before sleep: 3GB free
(1 * 1024**3, 8 * 1024**3), # After sleep: 1GB free (negative freed!)
mock_platform.get_free_memory.side_effect = [
3 * 1024**3, # Before sleep: 3GB free
1 * 1024**3, # After sleep: 1GB free (negative freed!)
]
mock_platform.get_device_total_memory.return_value = 8 * 1024**3
mock_allocator = Mock()
mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
@@ -160,8 +164,8 @@ class TestGPUDiffusionWorkerSleep:
mock_gpu_worker.sleep(level=1)
class TestGPUDiffusionWorkerWakeUp:
"""Test GPUDiffusionWorker.wake_up method."""
class TestDiffusionWorkerWakeUp:
"""Test DiffusionWorker.wake_up method."""
@patch("vllm.device_allocator.cumem.CuMemAllocator")
def test_wake_up_without_buffers(self, mock_allocator_class, mock_gpu_worker):
tests/distributed/omni_connectors/test_kv_flow.py
+241 -179
View File
@@ -1,189 +1,251 @@
import unittest
from unittest.mock import MagicMock, patch
import pytest
import torch
from vllm_omni.diffusion.models.bagel.pipeline_bagel import BagelPipeline
from tests.utils import hardware_test
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
OmniKVCacheConfig,
OmniKVTransferManager,
)
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
class MockGPUARModelRunner(GPUARModelRunner):
"""Subclass to bypass heavy initialization."""
def __init__(self, kv_caches, input_batch):
self.kv_caches = kv_caches
self.input_batch = input_batch
self.device = "cpu"
self.cache_config = MagicMock()
self.cache_config.block_size = 16
self.cache_config.cache_dtype = "auto"
self.logger = MagicMock()
class MockBagelPipeline(BagelPipeline):
"""Subclass to bypass heavy initialization."""
class MockConnector:
def __init__(self):
self.device = "cpu"
self.od_config = MagicMock()
self.bagel = MagicMock()
self.tokenizer = MagicMock()
self.language_model = MagicMock()
self.new_token_ids = {}
self.store = {}
def put(self, from_stage, to_stage, put_key, data):
# The manager now passes full key as put_key
key = f"{from_stage}->{to_stage}:{put_key}"
self.store[key] = data
return True, len(str(data)), None # (success, size, metadata)
def get(self, from_stage, to_stage, get_key, metadata=None):
# The manager now passes full key as get_key
key = f"{from_stage}->{to_stage}:{get_key}"
if key in self.store:
return self.store[key], len(str(self.store[key]))
return None
class TestKVFlow(unittest.TestCase):
def setUp(self):
# Common constants
self.num_layers = 2
self.num_heads = 4
self.head_dim = 16
self.block_size = 8
self.x = 8 # PagedAttention factor
self.seq_len = 20
self.req_id = "req_test_1"
def test_sender_extraction_logic(self):
"""Test extraction logic in GPUARModelRunner."""
if GPUARModelRunner is object:
self.skipTest("vLLM not installed")
# 1. Setup KV Cache (List of tuples (K, V))
# Shape: [num_blocks, block_size, num_heads, head_dim] matching 4D expectation
num_blocks = 10
kv_caches = []
for _ in range(self.num_layers):
k_cache = torch.randn(num_blocks, self.block_size, self.num_heads, self.head_dim)
v_cache = torch.randn(num_blocks, self.block_size, self.num_heads, self.head_dim)
# Stack K and V to create [2, num_blocks, block_size, n_heads, head_dim]
layer_cache = torch.stack([k_cache, v_cache], dim=0)
kv_caches.append(layer_cache)
# 2. Setup Input Batch Mock
block_ids = [1, 3, 5]
mock_input_batch = MagicMock()
mock_input_batch.req_id_to_index = {self.req_id: 0}
mock_input_batch.block_table.get_row.return_value = block_ids
# 3. Instantiate Runner
runner = MockGPUARModelRunner(kv_caches, mock_input_batch)
# 4. Run Extraction
# calling _extract_kv_cache directly: (req_id, block_ids, seq_len)
result = runner._extract_kv_cache(self.req_id, block_ids, self.seq_len)
# 5. Verify Result
self.assertIsNotNone(result)
self.assertEqual(result.request_id, self.req_id)
# Check keys "key_cache" and "value_cache" exist (length 2)
self.assertEqual(len(result.layer_blocks["key_cache"]), 2)
self.assertEqual(len(result.layer_blocks["value_cache"]), 2)
# Check Tensor Shape: [seq_len, num_heads, head_dim]
# result.layer_blocks["key_cache"] is a list of tensors
expected_shape = (self.seq_len, self.num_heads, self.head_dim)
self.assertEqual(result.layer_blocks["key_cache"][0].shape, expected_shape)
return result # Return for use in next test
def test_receiver_injection_logic(self):
"""Test injection logic in BagelPipeline."""
if BagelPipeline is object:
self.skipTest("vLLM not installed")
# 1. Get Data from Sender Test (simulate transfer)
# Re-create data manually to be independent
key_cache = []
value_cache = []
expected_shape = (self.seq_len, self.num_heads, self.head_dim)
for i in range(self.num_layers):
key_cache.append(torch.randn(expected_shape))
value_cache.append(torch.randn(expected_shape))
layer_blocks = {"key_cache": key_cache, "value_cache": value_cache}
transfer_data = MagicMock() # Mock KVCacheTransferData
transfer_data.layer_blocks = layer_blocks
transfer_data.metadata = {"kv_lens": [self.seq_len], "ropes": [0]}
# 2. Setup Request with Injected Data
req = OmniDiffusionRequest(prompt="test")
from types import SimpleNamespace
req.past_key_values = SimpleNamespace(**layer_blocks)
req.kv_metadata = transfer_data.metadata
# 3. Setup Pipeline
pipeline = MockBagelPipeline()
# Mock Bagel's NaiveCache
class RealNaiveCache: # Minimal impl
def __init__(self, n):
self.key_cache = {i: None for i in range(n)}
self.value_cache = {i: None for i in range(n)}
pipeline.bagel.config.llm_config.num_hidden_layers = self.num_layers
captured_context = {}
def mock_prepare_prompts(curr_kvlens, curr_rope, **kwargs):
# Capture the state passed to prepare_prompts
captured_context["kv_lens"] = curr_kvlens
captured_context["ropes"] = curr_rope
return {}, [0], [0]
pipeline.bagel.prepare_prompts = MagicMock(side_effect=mock_prepare_prompts)
with patch("vllm_omni.diffusion.models.bagel.pipeline_bagel.NaiveCache") as MockNaiveCacheCls:
# Setup the instance returned by the constructor
mock_cache_instance = RealNaiveCache(self.num_layers)
MockNaiveCacheCls.return_value = mock_cache_instance
# Verification of Injection Logic (Simulation)
current_cache = RealNaiveCache(self.num_layers)
# --- Logic from Source Code ---
injected_kv = req.past_key_values
if isinstance(current_cache, RealNaiveCache) and hasattr(injected_kv, "key_cache"):
# Assuming injected_kv is SimpleNamespace or object with list attrs
for layer_idx in range(len(injected_kv.key_cache)):
if injected_kv.key_cache[layer_idx] is not None:
k_tensor = injected_kv.key_cache[layer_idx]
v_tensor = injected_kv.value_cache[layer_idx]
if k_tensor.device != pipeline.device:
k_tensor = k_tensor.to(pipeline.device)
if v_tensor.device != pipeline.device:
v_tensor = v_tensor.to(pipeline.device)
current_cache.key_cache[layer_idx] = k_tensor
current_cache.value_cache[layer_idx] = v_tensor
self.assertTrue(torch.allclose(current_cache.key_cache[0], layer_blocks["key_cache"][0]))
self.assertTrue(torch.allclose(current_cache.value_cache[1], layer_blocks["value_cache"][1]))
def test_integration(self):
"""Simulate the flow from Sender -> Connector (Dict) -> Receiver."""
if GPUARModelRunner is object or BagelPipeline is object:
self.skipTest("vLLM not installed")
# 1. Sender (Extraction)
runner_test_result = self.test_sender_extraction_logic() # Get KVCacheTransferData
# 2. Connector (Serialize/Deserialize Simulation)
# KVCacheTransferData has to_dict
data_dict = runner_test_result.to_dict()
# 3. Receiver (Request Setup)
req = OmniDiffusionRequest(prompt="integration_test")
req.past_key_values = data_dict["layer_blocks"]
req.kv_metadata = data_dict["metadata"]
# 4. Receiver (Injection Simulation)
# Use the logic verification again
pass
@pytest.fixture
def mock_connector():
return MockConnector()
if __name__ == "__main__":
unittest.main()
@pytest.fixture
def kv_config():
return OmniKVCacheConfig(
connector_config={"type": "mock"},
from_stage="stage1",
to_stage="stage2",
stage_id="stage2", # Acting as receiver for some tests
need_recv_cache=True,
need_send_cache=True,
recv_timeout=1.0, # Short timeout for tests
)
@pytest.fixture
def common_constants():
return {
"num_layers": 2,
"num_heads": 4,
"head_dim": 16,
"block_size": 8,
"seq_len": 20,
"req_id": "req_test_1",
}
@pytest.mark.cache
@hardware_test(
res={"cuda": "L4"},
num_cards=2,
)
def test_manager_extraction(kv_config, mock_connector, common_constants):
"""Test extraction and sending logic in OmniKVTransferManager."""
num_layers = common_constants["num_layers"]
block_size = common_constants["block_size"]
num_heads = common_constants["num_heads"]
head_dim = common_constants["head_dim"]
seq_len = common_constants["seq_len"]
req_id = common_constants["req_id"]
num_blocks = 10
kv_caches = []
for _ in range(num_layers):
k_cache = torch.randn(num_blocks, block_size, num_heads, head_dim)
v_cache = torch.randn(num_blocks, block_size, num_heads, head_dim)
# Stack K and V to create [2, num_blocks, block_size, n_heads, head_dim]
layer_cache = torch.stack([k_cache, v_cache], dim=0)
kv_caches.append(layer_cache)
block_ids = [1, 3, 5]
finished_reqs = {req_id: {"block_ids": block_ids, "seq_len": seq_len}}
manager = OmniKVTransferManager(kv_config)
# Mock the connector factory or injection
manager._connector = mock_connector
processed = manager.handle_finished_requests_kv_transfer(finished_reqs, kv_caches, block_size, "float32")
assert req_id in processed
# Check if data was put into connector
# Manager builds full key: omni_{from}_to_{to}_kv_cache_{req_id}
full_request_id = f"omni_stage1_to_stage2_kv_cache_{req_id}"
expected_key = f"stage1->stage2:{full_request_id}"
assert expected_key in mock_connector.store
data = mock_connector.store[expected_key]
assert data["request_id"] == req_id
assert "layer_blocks" in data
assert len(data["layer_blocks"]["key_cache"]) == num_layers
# Verify shape of extracted tensor: [seq_len, heads, dim]
# Note: Manager detaches and moves to CPU
expected_shape = (seq_len, num_heads, head_dim)
assert data["layer_blocks"]["key_cache"][0].shape == expected_shape
@pytest.mark.cache
@hardware_test(
res={"cuda": "L4"},
num_cards=2,
)
def test_manager_reception(kv_config, mock_connector, common_constants):
"""Test reception and injection logic in OmniKVTransferManager."""
num_layers = common_constants["num_layers"]
block_size = common_constants["block_size"]
num_heads = common_constants["num_heads"]
head_dim = common_constants["head_dim"]
seq_len = common_constants["seq_len"]
req_id = common_constants["req_id"]
expected_shape = (seq_len, num_heads, head_dim)
key_cache = [torch.randn(expected_shape) for _ in range(num_layers)]
value_cache = [torch.randn(expected_shape) for _ in range(num_layers)]
layer_blocks = {"key_cache": key_cache, "value_cache": value_cache}
metadata = {
"block_size": block_size,
"num_layers": num_layers,
"dtype": "float32",
"seq_len": seq_len,
}
data_to_receive = {
"request_id": req_id,
"layer_blocks": layer_blocks,
"metadata": metadata,
"block_ids": [],
}
# In setUp, from_stage="stage1", stage_id="stage2". recv_stages=("stage1", "stage2")
manager = OmniKVTransferManager(kv_config)
manager._connector = mock_connector
# Pre-populate connector with data
# Manager builds full key: omni_{from}_to_{to}_kv_cache_{req_id}
full_request_id = f"omni_stage1_to_stage2_kv_cache_{req_id}"
store_key = f"stage1->stage2:{full_request_id}"
mock_connector.store[store_key] = data_to_receive
req = OmniDiffusionRequest(
prompts=["test_recv"],
sampling_params=OmniDiffusionSamplingParams(),
request_ids=[req_id],
)
# req.need_kv_receive = True # Implicitly handled by receive_kv_cache check? No, manager doesn't check it, runner does.
# But receive_kv_cache in manager checks request_id. Which we need to fix in manager next.
success = manager.receive_kv_cache(req, target_device=torch.device("cpu"))
assert success
assert hasattr(req, "past_key_values")
assert hasattr(req, "kv_metadata")
assert len(req.past_key_values.key_cache) == num_layers
assert torch.allclose(req.past_key_values.key_cache[0], key_cache[0])
assert req.kv_metadata["seq_len"] == seq_len
@pytest.mark.cache
@hardware_test(
res={"cuda": "L4"},
num_cards=2,
)
def test_integration_flow(common_constants):
"""Simulate extraction -> connector -> reception."""
num_layers = common_constants["num_layers"]
block_size = common_constants["block_size"]
num_heads = common_constants["num_heads"]
head_dim = common_constants["head_dim"]
req_id = common_constants["req_id"]
sender_config = OmniKVCacheConfig(
connector_config={"type": "mock"}, from_stage="sender", to_stage="receiver", need_send_cache=True
)
sender_manager = OmniKVTransferManager(sender_config)
connector = MockConnector()
sender_manager._connector = connector # Shared connector
# Create Data
num_blocks = 5
kv_caches = []
for _ in range(num_layers):
layer = torch.randn(2, num_blocks, block_size, num_heads, head_dim)
kv_caches.append(layer)
finished_reqs = {req_id: {"block_ids": [0, 1], "seq_len": 10}}
# Send
sender_manager.handle_finished_requests_kv_transfer(finished_reqs, kv_caches, block_size, "float32")
receiver_config = OmniKVCacheConfig(
connector_config={"type": "mock"},
from_stage="sender",
stage_id="receiver",
need_recv_cache=True,
recv_timeout=1.0,
)
receiver_manager = OmniKVTransferManager(receiver_config)
receiver_manager._connector = connector # Share the same mock connector instance
req = OmniDiffusionRequest(
prompts=["test_integ"],
sampling_params=OmniDiffusionSamplingParams(),
request_ids=[req_id],
)
# Receive
success = receiver_manager.receive_kv_cache(req)
# Verify
assert success
assert req.past_key_values is not None
assert req.kv_metadata["seq_len"] == 10
@pytest.mark.cache
@hardware_test(
res={"cuda": "L4"},
num_cards=2,
)
def test_manager_extraction_no_connector(kv_config, common_constants):
"""Test extraction when connector is unavailable (should still return IDs)."""
block_size = common_constants["block_size"]
req_id = common_constants["req_id"]
manager = OmniKVTransferManager(kv_config)
# Force connector to be None
manager._connector = None
manager.config.connector_config = None
finished_reqs = {req_id: {"block_ids": [1, 2], "seq_len": 10}}
processed = manager.handle_finished_requests_kv_transfer(
finished_reqs, kv_caches=[], block_size=block_size, cache_dtype="float32"
)
assert req_id in processed
tests/e2e/offline_inference/conftest.py
+16 -12
View File
@@ -7,11 +7,12 @@ Pytest configuration and fixtures for vllm-omni tests.
from typing import Any
import pytest
from vllm import TextPrompt
from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
from vllm.sampling_params import SamplingParams
from tests.conftest import _run_post_test_cleanup
from tests.conftest import _run_post_test_cleanup, _run_pre_test_cleanup
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniSamplingParams
from vllm_omni.outputs import OmniRequestOutput
PromptAudioInput = list[tuple[Any, int]] | tuple[Any, int] | None
@@ -50,6 +51,9 @@ class OmniRunner:
stage_configs_path: Optional path to YAML stage config file
**kwargs: Additional arguments passed to Omni
"""
cleanup_dist_env_and_memory()
_run_pre_test_cleanup(enable_force=True)
_run_post_test_cleanup(enable_force=True)
self.model_name = model_name
self.seed = seed
@@ -64,7 +68,7 @@ class OmniRunner:
**kwargs,
)
def get_default_sampling_params_list(self) -> list[SamplingParams]:
def get_default_sampling_params_list(self) -> list[OmniSamplingParams]:
"""
Get a list of default sampling parameters for all stages.
@@ -82,7 +86,7 @@ class OmniRunner:
videos: PromptVideoInput = None,
mm_processor_kwargs: dict[str, Any] | None = None,
modalities: list[str] | None = None,
) -> list[dict[str, Any]]:
) -> list[TextPrompt]:
"""
Construct Omni input format from prompts and multimodal data.
@@ -175,7 +179,7 @@ class OmniRunner:
f"<|im_start|>assistant\n"
)
input_dict: dict[str, Any] = {"prompt": full_prompt}
input_dict: TextPrompt = {"prompt": full_prompt}
if multi_modal_data:
input_dict["multi_modal_data"] = multi_modal_data
if modalities:
@@ -189,8 +193,8 @@ class OmniRunner:
def generate(
self,
prompts: list[dict[str, Any]],
sampling_params_list: list[SamplingParams] | None = None,
prompts: list[TextPrompt],
sampling_params_list: list[OmniSamplingParams] | None = None,
) -> list[OmniRequestOutput]:
"""
Generate outputs for the given prompts.
@@ -212,7 +216,7 @@ class OmniRunner:
def generate_multimodal(
self,
prompts: list[str] | str,
sampling_params_list: list[SamplingParams] | None = None,
sampling_params_list: list[OmniSamplingParams] | None = None,
system_prompt: str | None = None,
audios: PromptAudioInput = None,
images: PromptImageInput = None,
@@ -249,7 +253,7 @@ class OmniRunner:
def generate_audio(
self,
prompts: list[str] | str,
sampling_params_list: list[SamplingParams] | None = None,
sampling_params_list: list[OmniSamplingParams] | None = None,
system_prompt: str | None = None,
audios: PromptAudioInput = None,
mm_processor_kwargs: dict[str, Any] | None = None,
@@ -276,7 +280,7 @@ class OmniRunner:
def generate_video(
self,
prompts: list[str] | str,
sampling_params_list: list[SamplingParams] | None = None,
sampling_params_list: list[OmniSamplingParams] | None = None,
system_prompt: str | None = None,
videos: PromptVideoInput = None,
mm_processor_kwargs: dict[str, Any] | None = None,
@@ -303,7 +307,7 @@ class OmniRunner:
def generate_image(
self,
prompts: list[str] | str,
sampling_params_list: list[SamplingParams] | None = None,
sampling_params_list: list[OmniSamplingParams] | None = None,
system_prompt: str | None = None,
images: PromptImageInput = None,
mm_processor_kwargs: dict[str, Any] | None = None,
@@ -336,7 +340,7 @@ class OmniRunner:
self.close()
del self.omni
cleanup_dist_env_and_memory()
_run_post_test_cleanup()
_run_post_test_cleanup(enable_force=True)
def close(self):
"""Close and cleanup the Omni instance."""
tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml
+3 -3
View File
@@ -10,7 +10,7 @@ stage_args:
engine_args:
model_stage: thinker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
max_model_len: 896
max_num_batched_tokens: 896
@@ -40,7 +40,7 @@ stage_args:
engine_args:
model_stage: talker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
max_model_len: 896
max_num_batched_tokens: 896
@@ -70,7 +70,7 @@ stage_args:
engine_args:
model_stage: code2wav
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
gpu_memory_utilization: 0.15
enforce_eager: true
tests/e2e/offline_inference/stage_configs/qwen2_5_omni_ci.yaml
+3 -3
View File
@@ -11,7 +11,7 @@ stage_args:
engine_args:
model_stage: thinker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
max_model_len: 896
max_num_batched_tokens: 896
@@ -41,7 +41,7 @@ stage_args:
engine_args:
model_stage: talker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
max_model_len: 896
max_num_batched_tokens: 896
@@ -71,7 +71,7 @@ stage_args:
engine_args:
model_stage: code2wav
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
gpu_memory_utilization: 0.15
enforce_eager: true
tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml
+8 -8
View File
@@ -7,21 +7,21 @@
stage_args:
- stage_id: 0
runtime:
devices: "0,1"
devices: "0"
max_batch_size: 1
engine_args:
model_stage: thinker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.6
gpu_memory_utilization: 0.9
enforce_eager: false
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
tensor_parallel_size: 1
load_format: dummy
final_output: true
final_output_type: text
@@ -42,9 +42,9 @@ stage_args:
engine_args:
model_stage: talker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.3
gpu_memory_utilization: 0.6
enforce_eager: true
trust_remote_code: true
engine_output_type: latent # Output codec codes for code2wav
@@ -68,12 +68,12 @@ stage_args:
- stage_id: 2
runtime:
devices: "0"
devices: "1"
max_batch_size: 1
engine_args:
model_stage: code2wav
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
enforce_eager: true
trust_remote_code: true
tests/e2e/offline_inference/stage_configs/rocm/qwen2_5_omni_ci.yaml
+3 -3
View File
@@ -11,7 +11,7 @@ stage_args:
engine_args:
model_stage: thinker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
max_model_len: 896
max_num_batched_tokens: 896
@@ -41,7 +41,7 @@ stage_args:
engine_args:
model_stage: talker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
max_model_len: 896
max_num_batched_tokens: 896
@@ -71,7 +71,7 @@ stage_args:
engine_args:
model_stage: code2wav
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
gpu_memory_utilization: 0.15
enforce_eager: true
tests/e2e/offline_inference/stage_configs/rocm/qwen3_omni_ci.yaml
+8 -8
View File
@@ -7,21 +7,21 @@
stage_args:
- stage_id: 0
runtime:
devices: "0,1"
devices: "0"
max_batch_size: 1
engine_args:
model_stage: thinker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.6
gpu_memory_utilization: 0.9
enforce_eager: false
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
tensor_parallel_size: 1
load_format: dummy
final_output: true
final_output_type: text
@@ -42,9 +42,9 @@ stage_args:
engine_args:
model_stage: talker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.3
gpu_memory_utilization: 0.6
enforce_eager: true
trust_remote_code: true
engine_output_type: latent # Output codec codes for code2wav
@@ -68,12 +68,12 @@ stage_args:
- stage_id: 2
runtime:
devices: "0"
devices: "1"
max_batch_size: 1
engine_args:
model_stage: code2wav
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
enforce_eager: true
trust_remote_code: true
tests/e2e/offline_inference/test_cache_dit.py
+10 -6
View File
@@ -15,6 +15,8 @@ from pathlib import Path
import pytest
import torch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
# ruff: noqa: E402
REPO_ROOT = Path(__file__).resolve().parents[2]
if str(REPO_ROOT) not in sys.path:
@@ -55,12 +57,14 @@ def test_cache_dit(model_name: str):
outputs = m.generate(
"a photo of a cat sitting on a laptop keyboard",
height=height,
width=width,
num_inference_steps=num_inference_steps,
guidance_scale=0.0,
generator=torch.Generator("cuda").manual_seed(42),
num_outputs_per_prompt=1, # Single output for speed
OmniDiffusionSamplingParams(
height=height,
width=width,
num_inference_steps=num_inference_steps,
guidance_scale=0.0,
generator=torch.Generator("cuda").manual_seed(42),
num_outputs_per_prompt=1, # Single output for speed
),
)
# Extract images from request_output[0]['images']
first_output = outputs[0]
tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+11 -8
View File
@@ -6,7 +6,8 @@ import torch
from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
from tests.utils import GPUMemoryMonitor
from vllm_omni.utils.platform_utils import is_npu, is_rocm
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.platforms import current_omni_platform
# ruff: noqa: E402
REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -19,7 +20,7 @@ models = ["riverclouds/qwen_image_random"]
def inference(model_name: str, offload: bool = True):
torch.cuda.empty_cache()
current_omni_platform.empty_cache()
device_index = torch.cuda.current_device()
monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
monitor.start()
@@ -30,11 +31,13 @@ def inference(model_name: str, offload: bool = True):
m.generate(
"a photo of a cat sitting on a laptop keyboard",
height=height,
width=width,
num_inference_steps=9,
guidance_scale=0.0,
generator=torch.Generator("cuda").manual_seed(42),
OmniDiffusionSamplingParams(
height=height,
width=width,
num_inference_steps=9,
guidance_scale=0.0,
generator=torch.Generator("cuda").manual_seed(42),
),
)
peak = monitor.peak_used_mb
monitor.stop()
@@ -42,7 +45,7 @@ def inference(model_name: str, offload: bool = True):
return peak
@pytest.mark.skipif(is_npu() or is_rocm(), reason="Hardware not supported")
@pytest.mark.skipif(current_omni_platform.is_npu() or current_omni_platform.is_rocm(), reason="Hardware not supported")
@pytest.mark.parametrize("model_name", models)
def test_cpu_offload_diffusion_model(model_name: str):
try:
tests/e2e/offline_inference/test_diffusion_lora.py
+19 -19
View File
@@ -7,8 +7,8 @@ import pytest
import torch
from safetensors.torch import save_file
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.utils.platform_utils import is_npu
# ruff: noqa: E402
REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -24,10 +24,6 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
# model to reduce runtime and avoid extra downloads.
models = ["Tongyi-MAI/Z-Image-Turbo"]
# NPU still can't run Tongyi-MAI/Z-Image-Turbo properly.
if is_npu():
pytest.skip("Tongyi-MAI/Z-Image-Turbo is not supported on NPU yet.", allow_module_level=True)
@pytest.mark.parametrize("model_name", models)
def test_diffusion_model(model_name: str, tmp_path: Path):
@@ -88,12 +84,14 @@ def test_diffusion_model(model_name: str, tmp_path: Path):
outputs = m.generate(
prompt,
height=height,
width=width,
num_inference_steps=2,
guidance_scale=0.0,
generator=torch.Generator("cuda").manual_seed(42),
num_outputs_per_prompt=1,
OmniDiffusionSamplingParams(
height=height,
width=width,
num_inference_steps=2,
guidance_scale=0.0,
generator=torch.Generator("cuda").manual_seed(42),
num_outputs_per_prompt=1,
),
)
images = _extract_images(outputs)
@@ -116,14 +114,16 @@ def test_diffusion_model(model_name: str, tmp_path: Path):
)
outputs_lora = m.generate(
prompt,
height=height,
width=width,
num_inference_steps=2,
guidance_scale=0.0,
generator=torch.Generator("cuda").manual_seed(42),
num_outputs_per_prompt=1,
lora_request=lora_request,
lora_scale=2.0,
OmniDiffusionSamplingParams(
height=height,
width=width,
num_inference_steps=2,
guidance_scale=0.0,
generator=torch.Generator("cuda").manual_seed(42),
num_outputs_per_prompt=1,
lora_request=lora_request,
lora_scale=2.0,
),
)
images_lora = _extract_images(outputs_lora)
assert len(images_lora) == 1
tests/e2e/offline_inference/test_ovis_image.py
+25 -14
View File
@@ -25,6 +25,7 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.models.ovis_image.pipeline_ovis_image import OvisImagePipeline
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@pytest.fixture
@@ -162,11 +163,13 @@ def test_basic_generation(ovis_pipeline):
"""Test the forward pass logic."""
# Setup request
req = OmniDiffusionRequest(
prompt="A photo of a cat",
height=256,
width=256,
num_inference_steps=2,
guidance_scale=1.0,
prompts=["A photo of a cat"],
sampling_params=OmniDiffusionSamplingParams(
height=256,
width=256,
num_inference_steps=2,
guidance_scale=1.0,
),
)
output = ovis_pipeline(req)
@@ -184,12 +187,18 @@ def test_basic_generation(ovis_pipeline):
def test_guidance_scale(ovis_pipeline):
"""Test that classifier-free guidance path is taken when scale > 1.0."""
req = OmniDiffusionRequest(
prompt="A photo of a cat",
negative_prompt="bad quality",
height=256,
width=256,
num_inference_steps=1,
guidance_scale=2.0, # Trigger CFG
prompts=[
{
"prompt": "A photo of a cat",
"negative_prompt": "bad quality",
}
],
sampling_params=OmniDiffusionSamplingParams(
height=256,
width=256,
num_inference_steps=1,
guidance_scale=2.0, # Trigger CFG
),
)
ovis_pipeline(req)
@@ -200,9 +209,11 @@ def test_resolution_check(ovis_pipeline):
"""Test resolution divisible validation logic if present."""
# Pass odd resolution
req = OmniDiffusionRequest(
prompt="test",
height=250, # Not divisible by 16 (8*2)
width=250,
prompts=["test"],
sampling_params=OmniDiffusionSamplingParams(
height=250, # Not divisible by 16 (8*2)
width=250,
),
)
# Should warn but proceed (as per code I read earlier) or resize?
tests/e2e/offline_inference/test_qwen2_5_omni.py
+3 -3
View File
@@ -13,7 +13,7 @@ from vllm.assets.video import VideoAsset
from vllm.envs import VLLM_USE_MODELSCOPE
from vllm.multimodal.image import convert_image_mode
from vllm_omni.utils import is_npu, is_rocm
from vllm_omni.platforms import current_omni_platform
from .conftest import OmniRunner
from .utils import create_new_process_for_each_test
@@ -21,9 +21,9 @@ from .utils import create_new_process_for_each_test
models = ["Qwen/Qwen2.5-Omni-3B"]
# CI stage config optimized for 24GB GPU (L4/RTX3090) or NPU
if is_npu():
if current_omni_platform.is_npu():
stage_config = str(Path(__file__).parent / "stage_configs" / "npu" / "qwen2_5_omni_ci.yaml")
elif is_rocm():
elif current_omni_platform.is_rocm():
# ROCm stage config optimized for MI325 GPU
stage_config = str(Path(__file__).parent / "stage_configs" / "rocm" / "qwen2_5_omni_ci.yaml")
else:
tests/e2e/offline_inference/test_qwen3_omni.py
+2 -2
View File
@@ -14,14 +14,14 @@ from pathlib import Path
import pytest
from vllm.assets.video import VideoAsset
from vllm_omni.utils import is_rocm
from vllm_omni.platforms import current_omni_platform
from .conftest import OmniRunner
models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
# CI stage config for 2xH100-80G GPUs or AMD GPU MI325
if is_rocm():
if current_omni_platform.is_rocm():
# ROCm stage config optimized for MI325 GPU
stage_configs = [str(Path(__file__).parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")]
else:
tests/e2e/offline_inference/test_sequence_parallel.py
+32 -25
View File
@@ -21,6 +21,8 @@ import torch
import torch.distributed as dist
from PIL import Image
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
# ruff: noqa: E402
REPO_ROOT = Path(__file__).resolve().parents[3]
if str(REPO_ROOT) not in sys.path:
@@ -28,8 +30,7 @@ if str(REPO_ROOT) not in sys.path:
from vllm_omni import Omni
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.diffusion.distributed.parallel_state import device_count
from vllm_omni.diffusion.envs import get_device_name
from vllm_omni.platforms import current_omni_platform
models = ["riverclouds/qwen_image_random"]
@@ -61,9 +62,9 @@ def _cleanup_distributed():
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.synchronize()
if current_omni_platform.is_available():
current_omni_platform.empty_cache()
current_omni_platform.synchronize()
time.sleep(5)
@@ -81,12 +82,14 @@ def _run_baseline(model_name: str, dtype: torch.dtype, attn_backend: str, height
try:
outputs = baseline.generate(
PROMPT,
height=height,
width=width,
num_inference_steps=4,
guidance_scale=0.0,
generator=torch.Generator(get_device_name()).manual_seed(seed),
num_outputs_per_prompt=1,
OmniDiffusionSamplingParams(
height=height,
width=width,
num_inference_steps=4,
guidance_scale=0.0,
generator=torch.Generator(current_omni_platform.device_type).manual_seed(seed),
num_outputs_per_prompt=1,
),
)
return outputs[0].request_output[0].images
finally:
@@ -116,12 +119,14 @@ def _run_sp(
try:
outputs = sp.generate(
PROMPT,
height=height,
width=width,
num_inference_steps=4,
guidance_scale=0.0,
generator=torch.Generator(get_device_name()).manual_seed(seed),
num_outputs_per_prompt=1,
OmniDiffusionSamplingParams(
height=height,
width=width,
num_inference_steps=4,
guidance_scale=0.0,
generator=torch.Generator(current_omni_platform.device_type).manual_seed(seed),
num_outputs_per_prompt=1,
),
)
return outputs[0].request_output[0].images
finally:
@@ -151,8 +156,8 @@ def test_baseline_only(model_name: str, dtype: torch.dtype, attn_backend: str):
@pytest.mark.parametrize("attn_backend", ["sdpa"])
def test_sp_ulysses2_only(model_name: str, dtype: torch.dtype, attn_backend: str):
"""Test SP inference only (ulysses=2)."""
if device_count() < 2:
pytest.skip(f"Test requires 2 GPUs but only {device_count()} available")
if current_omni_platform.get_device_count() < 2:
pytest.skip(f"Test requires 2 GPUs but only {current_omni_platform.get_device_count()} available")
height = 256
width = 256
@@ -171,8 +176,8 @@ def test_sp_ulysses2_only(model_name: str, dtype: torch.dtype, attn_backend: str
@pytest.mark.parametrize("attn_backend", ["sdpa"])
def test_sp_ring2_only(model_name: str, dtype: torch.dtype, attn_backend: str):
"""Test SP inference only (ring=2)."""
if device_count() < 2:
pytest.skip(f"Test requires 2 GPUs but only {device_count()} available")
if current_omni_platform.get_device_count() < 2:
pytest.skip(f"Test requires 2 GPUs but only {current_omni_platform.get_device_count()} available")
height = 256
width = 256
@@ -206,8 +211,8 @@ def test_sequence_parallel(
)
sp_size = ulysses_degree * ring_degree
if device_count() < sp_size:
pytest.skip(f"Test requires {sp_size} GPUs but only {device_count()} available")
if current_omni_platform.get_device_count() < sp_size:
pytest.skip(f"Test requires {sp_size} GPUs but only {current_omni_platform.get_device_count()} available")
height = 256
width = 256
@@ -237,8 +242,10 @@ def test_sequence_parallel_ulysses4(model_name: str, dtype: torch.dtype, attn_ba
ulysses_degree = 4
ring_degree = 1
if device_count() < ulysses_degree * ring_degree:
pytest.skip(f"Test requires {ulysses_degree * ring_degree} GPUs but only {device_count()} available")
if current_omni_platform.get_device_count() < ulysses_degree * ring_degree:
pytest.skip(
f"Test requires {ulysses_degree * ring_degree} GPUs but only {current_omni_platform.get_device_count()} available"
)
height = 272
width = 272
tests/e2e/offline_inference/test_stable_audio_model.py
+14 -9
View File
@@ -5,6 +5,7 @@ import numpy as np
import pytest
import torch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
# ruff: noqa: E402
@@ -29,16 +30,20 @@ def test_stable_audio_model(model_name: str):
sample_rate = 44100 # Stable Audio uses 44100 Hz
outputs = m.generate(
"The sound of a dog barking",
negative_prompt="Low quality.",
num_inference_steps=4, # Minimal steps for speed
guidance_scale=7.0,
generator=torch.Generator("cuda").manual_seed(42),
num_outputs_per_prompt=1,
extra={
"audio_start_in_s": audio_start_in_s,
"audio_end_in_s": audio_end_in_s,
prompts={
"prompt": "The sound of a dog barking",
"negative_prompt": "Low quality.",
},
sampling_params_list=OmniDiffusionSamplingParams(
num_inference_steps=4, # Minimal steps for speed
guidance_scale=7.0,
generator=torch.Generator("cuda").manual_seed(42),
num_outputs_per_prompt=1,
extra_args={
"audio_start_in_s": audio_start_in_s,
"audio_end_in_s": audio_end_in_s,
},
),
)
# Extract audio from OmniRequestOutput
tests/e2e/offline_inference/test_t2i_model.py
+13 -11
View File
@@ -5,8 +5,9 @@ from pathlib import Path
import pytest
import torch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.utils.platform_utils import is_npu, is_rocm
from vllm_omni.platforms import current_omni_platform
# ruff: noqa: E402
REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -20,12 +21,11 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
models = ["Tongyi-MAI/Z-Image-Turbo", "riverclouds/qwen_image_random"]
# NPU still can't run Tongyi-MAI/Z-Image-Turbo properly
# Modelscope can't find riverclouds/qwen_image_random
# TODO: When NPU support is ready, remove this branch.
if is_npu():
models = ["Qwen/Qwen-Image"]
elif is_rocm():
if current_omni_platform.is_npu():
models = ["Tongyi-MAI/Z-Image-Turbo", "Qwen/Qwen-Image"]
elif current_omni_platform.is_rocm():
# TODO: When ROCm support is ready, remove this branch.
# vLLM V0.11.0 has issues running riverclouds/qwen_image_random
# on ROCm
@@ -42,12 +42,14 @@ def test_diffusion_model(model_name: str):
width = 256
outputs = m.generate(
"a photo of a cat sitting on a laptop keyboard",
height=height,
width=width,
num_inference_steps=2,
guidance_scale=0.0,
generator=torch.Generator("cuda").manual_seed(42),
num_outputs_per_prompt=2,
OmniDiffusionSamplingParams(
height=height,
width=width,
num_inference_steps=2,
guidance_scale=0.0,
generator=torch.Generator("cuda").manual_seed(42),
num_outputs_per_prompt=2,
),
)
# Extract images from request_output[0]['images']
first_output = outputs[0]
tests/e2e/offline_inference/test_t2v_model.py
+11 -7
View File
@@ -5,6 +5,8 @@ from pathlib import Path
import pytest
import torch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
# ruff: noqa: E402
REPO_ROOT = Path(__file__).resolve().parents[2]
if str(REPO_ROOT) not in sys.path:
@@ -33,13 +35,15 @@ def test_video_diffusion_model(model_name: str):
width = 640
num_frames = 5
outputs = 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),
prompts="A cat sitting on a table",
sampling_params_list=OmniDiffusionSamplingParams(
height=height,
width=width,
num_frames=num_frames,
num_inference_steps=2,
guidance_scale=1.0,
generator=torch.Generator("cuda").manual_seed(42),
),
)
first_output = outputs[0]
assert first_output.final_output_type == "image"
tests/e2e/offline_inference/test_teacache.py
+10 -6
View File
@@ -15,6 +15,8 @@ from pathlib import Path
import pytest
import torch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
# ruff: noqa: E402
REPO_ROOT = Path(__file__).resolve().parents[2]
if str(REPO_ROOT) not in sys.path:
@@ -51,12 +53,14 @@ def test_teacache(model_name: str):
outputs = m.generate(
"a photo of a cat sitting on a laptop keyboard",
height=height,
width=width,
num_inference_steps=num_inference_steps,
guidance_scale=0.0,
generator=torch.Generator("cuda").manual_seed(42),
num_outputs_per_prompt=1, # Single output for speed
OmniDiffusionSamplingParams(
height=height,
width=width,
num_inference_steps=num_inference_steps,
guidance_scale=0.0,
generator=torch.Generator("cuda").manual_seed(42),
num_outputs_per_prompt=1, # Single output for speed
),
)
# Extract images from request_output[0]['images']
first_output = outputs[0]
tests/e2e/offline_inference/test_zimage_tensor_parallel.py
+12 -8
View File
@@ -12,6 +12,8 @@ import torch
from PIL import Image
from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
# ruff: noqa: E402
REPO_ROOT = Path(__file__).resolve().parents[2]
if str(REPO_ROOT) not in sys.path:
@@ -21,7 +23,7 @@ from tests.utils import GPUMemoryMonitor
from vllm_omni import Omni
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.utils.platform_utils import is_npu, is_rocm
from vllm_omni.platforms import current_omni_platform
# os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
@@ -90,12 +92,14 @@ def _run_zimage_generate(
num_requests = 4 # 1 warmup + 3 timed
gen = m.generate(
[PROMPT] * num_requests,
height=height,
width=width,
num_inference_steps=num_inference_steps,
guidance_scale=0.0,
seed=seed,
num_outputs_per_prompt=1,
OmniDiffusionSamplingParams(
height=height,
width=width,
num_inference_steps=num_inference_steps,
guidance_scale=0.0,
seed=seed,
num_outputs_per_prompt=1,
),
py_generator=True,
)
@@ -125,7 +129,7 @@ def _run_zimage_generate(
@pytest.mark.integration
def test_zimage_tensor_parallel_tp2(tmp_path: Path):
if is_npu() or is_rocm():
if current_omni_platform.is_npu() or current_omni_platform.is_rocm():
pytest.skip("Z-Image TP e2e test is only supported on CUDA for now.")
if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
pytest.skip("Z-Image TP=2 requires >= 2 CUDA devices.")
tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml
+12 -8
View File
@@ -6,22 +6,24 @@
# The following config has been verified on 2x H100-80G GPUs.
stage_args:
- stage_id: 0
stage_type: llm # Use llm stage type to launch OmniLLM
runtime:
devices: "0,1"
devices: "0"
max_batch_size: 5
engine_args:
model_stage: thinker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.6
gpu_memory_utilization: 0.9
enforce_eager: false
trust_remote_code: true
engine_output_type: latent # Output hidden states for talker
distributed_executor_backend: "mp"
enable_prefix_caching: false
max_num_batched_tokens: 32768
hf_config_name: thinker_config
tensor_parallel_size: 2
tensor_parallel_size: 1
load_format: dummy
final_output: true
final_output_type: text
@@ -36,15 +38,16 @@ stage_args:
repetition_penalty: 1.05
- stage_id: 1
stage_type: llm # Use llm stage type to launch OmniLLM
runtime:
devices: "1"
max_batch_size: 5
engine_args:
model_stage: talker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.3
gpu_memory_utilization: 0.6
enforce_eager: false
trust_remote_code: true
engine_output_type: latent # Output codec codes for code2wav
@@ -67,13 +70,14 @@ stage_args:
stop_token_ids: [2150]
- stage_id: 2
stage_type: llm # Use llm stage type to launch OmniLLM
runtime:
devices: "0"
devices: "1"
max_batch_size: 1
engine_args:
model_stage: code2wav
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
enforce_eager: true
trust_remote_code: true
tests/e2e/online_serving/stage_configs/qwen3_omni_thinker_ci.yaml
+1 -1
View File
@@ -7,7 +7,7 @@ stage_args:
engine_args:
model_stage: thinker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.6
enforce_eager: true
tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml
+10 -14
View File
@@ -2,27 +2,25 @@
# 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"
devices: "0"
max_batch_size: 5
engine_args:
model_stage: thinker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.6
gpu_memory_utilization: 0.9
enforce_eager: false
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
tensor_parallel_size: 1
final_output: true
final_output_type: text
is_comprehension: true
@@ -42,9 +40,9 @@ stage_args:
engine_args:
model_stage: talker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.3
gpu_memory_utilization: 0.6
enforce_eager: true
trust_remote_code: true
engine_output_type: latent # Output codec codes for code2wav
@@ -52,7 +50,6 @@ stage_args:
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
@@ -60,7 +57,7 @@ stage_args:
default_sampling_params:
temperature: 0.9
top_k: 50
max_tokens: 100
max_tokens: 1000
seed: 42
detokenize: False
repetition_penalty: 1.05
@@ -68,12 +65,12 @@ stage_args:
- stage_id: 2
runtime:
devices: "0"
devices: "1"
max_batch_size: 1
engine_args:
model_stage: code2wav
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
enforce_eager: true
trust_remote_code: true
@@ -83,7 +80,6 @@ stage_args:
distributed_executor_backend: "mp"
max_num_batched_tokens: 1000000
hf_config_name: thinker_config
load_format: dummy
async_scheduling: false
engine_input_source: [1]
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
@@ -93,7 +89,7 @@ stage_args:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 200
max_tokens: 2000
seed: 42
detokenize: True
repetition_penalty: 1.1
tests/e2e/online_serving/test_async_omni.py
+10 -9
View File
@@ -5,6 +5,7 @@ from contextlib import ExitStack
from pathlib import Path
import pytest
from vllm import SamplingParams
from vllm.inputs import PromptType
from vllm_omni.entrypoints.async_omni import AsyncOmni, ClientRequestState
@@ -25,15 +26,15 @@ async def generate(
) -> tuple[int, str]:
# Ensure generate doesn't complete too fast for cancellation test.
await asyncio.sleep(0.2)
thinker_sampling_params = {
"temperature": 0.4, # Deterministic
"top_p": 0.9,
"top_k": 1,
"max_tokens": max_tokens,
"repetition_penalty": 1.05,
"stop_token_ids": [151645], # Qwen EOS token <|im_end|>
"seed": SEED,
}
thinker_sampling_params = SamplingParams(
temperature=0.4, # Deterministic
top_p=0.9,
top_k=1,
max_tokens=max_tokens,
repetition_penalty=1.05,
stop_token_ids=[151645], # Qwen EOS token <|im_end|>
seed=SEED,
)
sampling_params_list = [
thinker_sampling_params,
tests/e2e/online_serving/test_images_generations_lora.py
+32 -98
View File
@@ -12,10 +12,6 @@ This validates:
import base64
import json
import os
import signal
import subprocess
import sys
import time
from io import BytesIO
from pathlib import Path
@@ -25,9 +21,8 @@ import requests
import torch
from PIL import Image
from safetensors.torch import save_file
from vllm.utils.network_utils import get_open_port
from vllm_omni.utils.platform_utils import is_npu
from tests.conftest import OmniServer
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
@@ -39,91 +34,8 @@ SIZE = "256x256"
SEED = 42
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:
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__))), # vllm-omni root
start_new_session=True,
)
# Wait for server to be ready.
max_wait = 1200
url = f"http://{self.host}:{self.port}/v1/models"
start_time = time.time()
while time.time() - start_time < max_wait:
try:
resp = requests.get(url, headers={"Authorization": "Bearer EMPTY"}, timeout=10)
if resp.status_code == 200:
print(f"Server ready on {self.host}:{self.port}")
return
except Exception:
pass
time.sleep(2)
raise RuntimeError(f"Server failed to become ready within {max_wait} seconds")
def __enter__(self):
self._start_server()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if self.proc is None:
return
try:
os.killpg(self.proc.pid, signal.SIGTERM)
except ProcessLookupError:
pass
try:
self.proc.wait(timeout=30)
except subprocess.TimeoutExpired:
try:
os.killpg(self.proc.pid, signal.SIGKILL)
except ProcessLookupError:
pass
self.proc.wait()
@pytest.fixture(scope="module")
def omni_server():
if is_npu():
pytest.skip("Tongyi-MAI/Z-Image-Turbo is not supported on NPU yet.")
with OmniServer(MODEL, ["--num-gpus", "1"]) as server:
yield server
@@ -187,17 +99,30 @@ def _image_blue_tail_slice(img: Image.Image) -> np.ndarray:
return tail
def _assert_slice_close(actual: np.ndarray, expected: np.ndarray, *, label: str) -> None:
def _slice_diff_stats(actual: np.ndarray, expected: np.ndarray) -> tuple[float, float]:
diff = np.abs(actual - expected)
return float(diff.max()), float(diff.mean())
def _assert_slice_close(
actual: np.ndarray,
expected: np.ndarray,
*,
label: str,
base_max: float,
base_mean: float,
) -> None:
assert actual.shape == (3, 3)
assert expected.shape == (3, 3)
diff = np.abs(actual - expected)
max_diff = float(diff.max())
mean_diff = float(diff.mean())
max_diff, mean_diff = _slice_diff_stats(actual, expected)
# NOTE: Different attention backends / torch.compile can introduce small
# floating-point drift that shows up as a few LSBs in uint8 pixels. Keep
# the reset check tolerant but bounded to avoid flaky CI.
assert max_diff <= 5.0 and mean_diff <= 3.0, (
f"{label} slice mismatch (max={max_diff:.1f}, mean={mean_diff:.1f}): {actual.tolist()}"
max_thresh = max(10.0, base_max + 4.0)
mean_thresh = max(6.0, base_mean + 4.0)
assert max_diff <= max_thresh and mean_diff <= mean_thresh, (
f"{label} slice mismatch (max={max_diff:.1f} > {max_thresh:.1f} or "
f"mean={mean_diff:.1f} > {mean_thresh:.1f}): {actual.tolist()}"
)
@@ -223,6 +148,9 @@ def test_images_generations_per_request_lora_switching(omni_server: OmniServer,
# Base generation.
base_img = _post_images(omni_server, _basic_payload())
base_slice = _image_blue_tail_slice(base_img)
base_ref_img = _post_images(omni_server, _basic_payload())
base_ref_slice = _image_blue_tail_slice(base_ref_img)
base_ref_max, base_ref_mean = _slice_diff_stats(base_ref_slice, base_slice)
# Adapter A: apply delta to V slice only.
lora_a_dir = tmp_path / "zimage_lora_a"
@@ -249,11 +177,17 @@ def test_images_generations_per_request_lora_switching(omni_server: OmniServer,
# Ensure switching back to no-LoRA restores the base output.
base_img_2 = _post_images(omni_server, _basic_payload())
base_slice_2 = _image_blue_tail_slice(base_img_2)
_assert_slice_close(base_slice_2, base_slice, label="base_after_reset")
base_reset = float(np.abs(base_slice_2 - base_slice).mean())
_, base_reset_mean = _slice_diff_stats(base_slice_2, base_slice)
_assert_slice_close(
base_slice_2,
base_slice,
label="base_after_reset",
base_max=base_ref_max,
base_mean=base_ref_mean,
)
# Ensure LoRA effects are clearly above the baseline drift.
min_delta = base_reset + 0.5
min_delta = max(base_reset_mean + 1.0, 1.5)
assert a_vs_base > min_delta, f"lora_a_vs_base drift too small: {a_vs_base} <= {min_delta}"
assert b_vs_base > min_delta, f"lora_b_vs_base drift too small: {b_vs_base} <= {min_delta}"
assert b_vs_a > min_delta, f"lora_b_vs_lora_a drift too small: {b_vs_a} <= {min_delta}"
tests/e2e/online_serving/test_qwen3_omni.py
+7 -17
View File
@@ -25,9 +25,10 @@ from tests.conftest import (
generate_synthetic_audio,
generate_synthetic_image,
generate_synthetic_video,
merge_base64_and_convert_to_text,
modify_stage_config,
)
from vllm_omni.utils import is_rocm
from vllm_omni.platforms import current_omni_platform
models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
@@ -57,7 +58,7 @@ def get_chunk_config():
CHUNK_CONFIG_PATH = get_chunk_config()
# CI stage config for 2xH100-80G GPUs or AMD GPU MI325
if is_rocm():
if current_omni_platform.is_rocm():
# ROCm stage config optimized for MI325 GPU
stage_configs = [str(Path(__file__).parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")]
else:
@@ -133,7 +134,7 @@ def dummy_messages_from_video_data(
def get_prompt(prompt_type="text_only"):
prompts = {
"text_only": "What is the capital of China?",
"text_only": "What is the capital of China? Answer in 20 words.",
"mix": "What is recited in the audio? What is in this image? Describe the video briefly.",
}
return prompts.get(prompt_type, prompts["text_only"])
@@ -144,7 +145,6 @@ def get_max_batch_size(size_type="few"):
return batch_sizes.get(size_type, 5)
@pytest.mark.skipif(is_rocm(), reason="Test skipped on AMD environment due to known output issues")
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> None:
"""
@@ -155,11 +155,6 @@ def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> N
Input Setting: stream=True
Datasets: single request
"""
# TODO：This skip will be removed when the chunk scenario supports multimodal input.
param = request.node.callspec.params.get("omni_server")
if param[1] == CHUNK_CONFIG_PATH:
pytest.skip("The current chunk scenario does not support multimodal.")
# Test single completion
e2e_list = list()
@@ -179,7 +174,7 @@ def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> N
chat_completion = client.chat.completions.create(model=omni_server.model, messages=messages, stream=True)
text_content = ""
audio_data = None
audio_data = []
for chunk in chat_completion:
for choice in chunk.choices:
if hasattr(choice, "delta"):
@@ -190,11 +185,7 @@ def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> N
modality = getattr(chunk, "modality", None)
if modality == "audio" and content:
# Audio chunk - content
if audio_data is None:
audio_data = content
else:
audio_data += content
audio_data.append(content)
elif modality == "text" and content:
# Text chunk - accumulate text content
text_content += content if content else ""
@@ -216,7 +207,7 @@ def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> N
), "The output does not contain any of the keywords."
# Verify text output same as audio output
audio_content = convert_audio_to_text(audio_data)
audio_content = merge_base64_and_convert_to_text(audio_data)
print(f"text content is: {text_content}")
print(f"audio content is: {audio_content}")
similarity = cosine_similarity_text(audio_content.lower(), text_content.lower())
@@ -224,7 +215,6 @@ def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> N
assert similarity > 0.9, "The audio content is not same as the text"
@pytest.mark.skipif(is_rocm(), reason="Test skipped on AMD environment due to known output issues")
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_text_to_text_audio_001(client: openai.OpenAI, omni_server) -> None:
"""
tests/e2e/stage_configs/qwen3_omni_ci.yaml
+10 -8
View File
@@ -7,22 +7,23 @@
stage_args:
- stage_id: 0
runtime:
devices: "0,1"
devices: "0"
max_batch_size: 5
engine_args:
model_stage: thinker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.6
gpu_memory_utilization: 0.9
enforce_eager: false
trust_remote_code: true
engine_output_type: latent # Output hidden states for talker
distributed_executor_backend: "mp"
max_num_batched_tokens: 32768
max_model_len: 32768
enable_prefix_caching: false
hf_config_name: thinker_config
tensor_parallel_size: 2
tensor_parallel_size: 1
final_output: true
final_output_type: text
is_comprehension: true
@@ -43,14 +44,15 @@ stage_args:
engine_args:
model_stage: talker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.3
gpu_memory_utilization: 0.6
enforce_eager: false
trust_remote_code: true
engine_output_type: latent # Output codec codes for code2wav
enable_prefix_caching: false
max_num_batched_tokens: 32768
max_model_len: 32768
distributed_executor_backend: "mp"
hf_config_name: talker_config
engine_input_source: [0]
@@ -66,12 +68,12 @@ stage_args:
- stage_id: 2
runtime:
devices: "0"
devices: "1"
max_batch_size: 1
engine_args:
model_stage: code2wav
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
enforce_eager: true
trust_remote_code: true
tests/entrypoints/openai_api/test_image_server.py
+70 -35
View File
@@ -14,11 +14,13 @@ from unittest.mock import AsyncMock, Mock
import pytest
from fastapi.testclient import TestClient
from PIL import Image
from vllm import SamplingParams
from vllm_omni.entrypoints.openai.image_api_utils import (
encode_image_base64,
parse_size,
)
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
# Unit Tests
@@ -111,7 +113,7 @@ class FakeAsyncOmni:
def __init__(self):
self.stage_list = ["llm", "diffusion"]
self.default_sampling_params_list = [{"temperature": 0.1}, {"top_p": 0.9}]
self.default_sampling_params_list = [SamplingParams(temperature=0.1), OmniDiffusionSamplingParams()]
self.captured_sampling_params_list = None
async def generate(self, prompt, request_id, sampling_params_list):
@@ -124,10 +126,13 @@ class FakeAsyncOmni:
def mock_async_diffusion():
"""Mock AsyncOmniDiffusion instance that returns fake images"""
mock = Mock()
mock.is_running = True # For health endpoint
mock.check_health = AsyncMock() # For LLM mode health check
async def generate(**kwargs):
# Return n PIL images wrapped in result object
n = kwargs.get("num_outputs_per_prompt", 1)
print("!!!!!!!!!!!!!!!!!!!!! kwargs", kwargs)
n = kwargs["sampling_params_list"][0].num_outputs_per_prompt
images = [Image.new("RGB", (64, 64), color="blue") for _ in range(n)]
return MockGenerationResult(images)
@@ -147,8 +152,9 @@ def test_client(mock_async_diffusion):
# Set up app state with diffusion engine
app.state.engine_client = mock_async_diffusion
app.state.diffusion_engine = mock_async_diffusion # Also set for health endpoint
app.state.stage_configs = [{"stage_type": "diffusion"}]
app.state.served_model_names = "Qwen/Qwen-Image"
app.state.diffusion_model_name = "Qwen/Qwen-Image" # For models endpoint
return TestClient(app)
@@ -169,10 +175,58 @@ def async_omni_test_client():
return TestClient(app)
@pytest.mark.skip(reason="Async API server uses different health check mechanism")
def test_health_endpoint(test_client):
"""Test health check endpoint - skipped for async server"""
pass
"""Test health check endpoint for diffusion mode"""
response = test_client.get("/health")
assert response.status_code == 200
data = response.json()
assert data["status"] == "healthy"
def test_health_endpoint_no_engine():
"""Test health check endpoint when no engine is initialized"""
from fastapi import FastAPI
from vllm_omni.entrypoints.openai.api_server import router
app = FastAPI()
app.include_router(router)
# Don't set any engine
client = TestClient(app)
response = client.get("/health")
assert response.status_code == 503
data = response.json()
assert data["status"] == "unhealthy"
def test_models_endpoint(test_client):
"""Test /v1/models endpoint for diffusion mode"""
response = test_client.get("/v1/models")
assert response.status_code == 200
data = response.json()
assert data["object"] == "list"
assert len(data["data"]) == 1
assert data["data"][0]["id"] == "Qwen/Qwen-Image"
assert data["data"][0]["object"] == "model"
def test_models_endpoint_no_engine():
"""Test /v1/models endpoint when no engine is initialized"""
from fastapi import FastAPI
from vllm_omni.entrypoints.openai.api_server import router
app = FastAPI()
app.include_router(router)
# Don't set any engine
client = TestClient(app)
response = client.get("/v1/models")
assert response.status_code == 200
data = response.json()
assert data["object"] == "list"
assert len(data["data"]) == 0
def test_generate_single_image(test_client):
@@ -217,11 +271,11 @@ def test_generate_images_async_omni_sampling_params(async_omni_test_client):
captured = engine.captured_sampling_params_list
assert captured is not None
assert len(captured) == 2
assert captured[0] == {"temperature": 0.1}
assert captured[1]["num_outputs_per_prompt"] == 2
assert captured[1]["height"] == 256
assert captured[1]["width"] == 256
assert captured[1]["seed"] == 7
assert captured[0].temperature == 0.1
assert captured[1].num_outputs_per_prompt == 2
assert captured[1].height == 256
assert captured[1].width == 256
assert captured[1].seed == 7
def test_generate_multiple_images(test_client):
@@ -454,30 +508,11 @@ def test_parameters_passed_through(test_client, mock_async_diffusion):
# Ensure generate() was called exactly once
mock_async_diffusion.generate.assert_awaited_once()
call_kwargs = mock_async_diffusion.generate.call_args[1]
assert call_kwargs["num_inference_steps"] == 100
assert call_kwargs["guidance_scale"] == 7.5
assert call_kwargs["true_cfg_scale"] == 3.0
assert call_kwargs["seed"] == 42
def test_optional_parameters_omitted(test_client, mock_async_diffusion):
"""Verify optional parameters not passed when omitted"""
response = test_client.post(
"/v1/images/generations",
json={
"prompt": "test",
"size": "512x512",
},
)
assert response.status_code == 200
# Ensure generate() was called exactly once
mock_async_diffusion.generate.assert_awaited_once()
call_kwargs = mock_async_diffusion.generate.call_args[1]
assert "num_inference_steps" not in call_kwargs
assert "guidance_scale" not in call_kwargs
assert "true_cfg_scale" not in call_kwargs
call_kwargs = mock_async_diffusion.generate.call_args[1]["sampling_params_list"][0]
assert call_kwargs.num_inference_steps == 100
assert call_kwargs.guidance_scale == 7.5
assert call_kwargs.true_cfg_scale == 3.0
assert call_kwargs.seed == 42
def test_model_field_omitted_works(test_client):
tests/entrypoints/test_omni_diffusion.py
+1103
View File
File diff suppressed because it is too large. Load diff
tests/entrypoints/test_omni_llm.py
+15 -7
View File
@@ -5,6 +5,7 @@ from typing import Any
from unittest.mock import MagicMock
import pytest
from vllm import SamplingParams
from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK
@@ -82,8 +83,9 @@ class _FakeStage:
self.stage_id = getattr(config, "stage_id", 0)
self.engine_args = config.engine_args
self.model_stage = getattr(config.engine_args, "model_stage", None)
self.stage_type = "llm"
# set default sampling params
self.default_sampling_params = {"temperature": 1.0}
self.default_sampling_params = SamplingParams(temperature=1.0)
# Allow configuring final_output and final_output_type
self.final_output = config.final_output if hasattr(config, "final_output") else False
self.final_output_type = getattr(config, "final_output_type", None)
@@ -637,8 +639,10 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
}
)
# Use dicts instead of object() for serializable sampling params
sampling_params_list = [{"temperature": 0.7}, {"temperature": 0.8}]
sampling_params_list = [
SamplingParams(temperature=0.7),
SamplingParams(temperature=0.8),
]
prompts = ["hi"]
outputs = omni.generate(prompts=prompts, sampling_params_list=sampling_params_list)
@@ -722,8 +726,13 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
}
)
# Use dicts instead of object() for serializable sampling params
outputs = omni.generate(prompts=["p"], sampling_params_list=[{"temperature": 0.7}, {"temperature": 0.8}])
outputs = omni.generate(
prompts=["p"],
sampling_params_list=[
SamplingParams(temperature=0.7),
SamplingParams(temperature=0.8),
],
)
assert outputs == []
@@ -922,8 +931,7 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
)
# Generate should handle error gracefully (log but continue)
# Use dict instead of object() for serializable sampling params
sampling_params_list = [{"temperature": 0.7}]
sampling_params_list = [SamplingParams(temperature=0.7)]
outputs = omni.generate(prompts=["hi"], sampling_params_list=sampling_params_list)
# Should return final output (error was logged but didn't stop processing)
assert isinstance(outputs, list)
tests/entrypoints/test_stage_utils.py
+53 -2
View File
@@ -1,5 +1,6 @@
import os
import sys
from unittest.mock import MagicMock
import pytest
@@ -42,6 +43,14 @@ def _make_dummy_torch(call_log):
return _Torch
def _make_mock_platform(device_type: str = "cuda", env_var: str = "CUDA_VISIBLE_DEVICES"):
"""Create a mock platform for testing."""
mock_platform = MagicMock()
mock_platform.device_type = device_type
mock_platform.device_control_env_var = env_var
return mock_platform
@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.
@@ -49,9 +58,51 @@ def test_set_stage_devices_respects_logical_ids(monkeypatch):
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")
# Mock the platform at the source module where it's defined
mock_platform = _make_mock_platform(device_type="cuda", env_var="CUDA_VISIBLE_DEVICES")
monkeypatch.setattr(
"vllm_omni.platforms.current_omni_platform",
mock_platform,
)
set_stage_devices(stage_id=0, devices="0,1")
assert os.environ["CUDA_VISIBLE_DEVICES"] == "6,7"
@pytest.mark.usefixtures("clean_gpu_memory_between_tests")
def test_set_stage_devices_npu_platform(monkeypatch):
"""Test that set_stage_devices works correctly for NPU platform."""
monkeypatch.setenv("ASCEND_RT_VISIBLE_DEVICES", "4,5")
call_log: list[int] = []
# Create NPU mock torch
class _Npu:
@staticmethod
def is_available():
return True
@staticmethod
def set_device(idx):
call_log.append(idx)
@staticmethod
def device_count():
return 2
class _NpuTorch:
npu = _Npu
monkeypatch.setitem(sys.modules, "torch", _NpuTorch)
# Mock NPU platform at the source module where it's defined
mock_platform = _make_mock_platform(device_type="npu", env_var="ASCEND_RT_VISIBLE_DEVICES")
monkeypatch.setattr(
"vllm_omni.platforms.current_omni_platform",
mock_platform,
)
set_stage_devices(stage_id=0, devices="0,1")
assert os.environ["ASCEND_RT_VISIBLE_DEVICES"] == "4,5"
tests/test_outputs.py
+171
View File
@@ -0,0 +1,171 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for OmniRequestOutput class."""
from unittest.mock import Mock
from PIL import Image
from vllm_omni.outputs import OmniRequestOutput
class TestOmniRequestOutput:
"""Tests for OmniRequestOutput class."""
def test_from_diffusion(self):
"""Test creating output from diffusion model."""
images = [Image.new("RGB", (64, 64), color="red")]
output = OmniRequestOutput.from_diffusion(
request_id="test-123",
images=images,
prompt="a cat",
metrics={"steps": 50},
)
assert output.request_id == "test-123"
assert output.images == images
assert output.prompt == "a cat"
assert output.metrics == {"steps": 50}
assert output.is_diffusion_output
assert output.num_images == 1
def test_from_pipeline(self):
"""Test creating output from pipeline stage."""
mock_request_output = Mock()
mock_request_output.request_id = "pipeline-123"
mock_request_output.prompt_token_ids = [1, 2, 3]
mock_request_output.outputs = [Mock()]
mock_request_output.encoder_prompt_token_ids = None
mock_request_output.prompt_logprobs = None
mock_request_output.num_cached_tokens = 10
mock_request_output.kv_transfer_params = None
mock_request_output.multimodal_output = {"image": Mock()}
output = OmniRequestOutput.from_pipeline(
stage_id=0,
final_output_type="text",
request_output=mock_request_output,
)
assert output.request_id == "pipeline-123"
assert output.stage_id == 0
assert output.final_output_type == "text"
assert output.is_pipeline_output
def test_prompt_token_ids_property(self):
"""Test prompt_token_ids property for streaming compatibility."""
mock_request_output = Mock()
mock_request_output.prompt_token_ids = [1, 2, 3, 4, 5]
output = OmniRequestOutput.from_pipeline(
stage_id=0,
final_output_type="text",
request_output=mock_request_output,
)
assert output.prompt_token_ids == [1, 2, 3, 4, 5]
def test_prompt_token_ids_none_when_no_request_output(self):
"""Test prompt_token_ids returns None when no request_output."""
output = OmniRequestOutput.from_diffusion(
request_id="test-123",
images=[],
prompt="a cat",
)
assert output.prompt_token_ids is None
def test_outputs_property(self):
"""Test outputs property for chat completion compatibility."""
mock_output = Mock()
mock_request_output = Mock()
mock_request_output.outputs = [mock_output]
output = OmniRequestOutput.from_pipeline(
stage_id=0,
final_output_type="text",
request_output=mock_request_output,
)
assert output.outputs == [mock_output]
def test_outputs_empty_when_no_request_output(self):
"""Test outputs returns empty list when no request_output."""
output = OmniRequestOutput.from_diffusion(
request_id="test-123",
images=[],
prompt="a cat",
)
assert output.outputs == []
def test_encoder_prompt_token_ids_property(self):
"""Test encoder_prompt_token_ids property."""
mock_request_output = Mock()
mock_request_output.encoder_prompt_token_ids = [10, 20, 30]
output = OmniRequestOutput.from_pipeline(
stage_id=0,
final_output_type="text",
request_output=mock_request_output,
)
assert output.encoder_prompt_token_ids == [10, 20, 30]
def test_num_cached_tokens_property(self):
"""Test num_cached_tokens property."""
mock_request_output = Mock()
mock_request_output.num_cached_tokens = 42
output = OmniRequestOutput.from_pipeline(
stage_id=0,
final_output_type="text",
request_output=mock_request_output,
)
assert output.num_cached_tokens == 42
def test_multimodal_output_property(self):
"""Test multimodal_output property."""
mock_request_output = Mock()
mock_audio = Mock()
expected_output = {"audio": mock_audio}
mock_request_output.multimodal_output = expected_output
output = OmniRequestOutput.from_pipeline(
stage_id=0,
final_output_type="audio",
request_output=mock_request_output,
)
assert output.multimodal_output is expected_output
def test_to_dict_diffusion(self):
"""Test to_dict for diffusion output."""
output = OmniRequestOutput.from_diffusion(
request_id="test-123",
images=[Image.new("RGB", (64, 64), color="red")],
prompt="a cat",
metrics={"steps": 50},
)
result = output.to_dict()
assert result["request_id"] == "test-123"
assert result["finished"] is True
assert result["final_output_type"] == "image"
assert result["num_images"] == 1
assert result["prompt"] == "a cat"
def test_to_dict_pipeline(self):
"""Test to_dict for pipeline output."""
mock_request_output = Mock()
mock_request_output.request_id = "pipeline-123"
output = OmniRequestOutput.from_pipeline(
stage_id=0,
final_output_type="text",
request_output=mock_request_output,
)
result = output.to_dict()
assert result["request_id"] == "pipeline-123"
assert result["finished"] is True
assert result["final_output_type"] == "text"
assert result["stage_id"] == 0
tests/utils.py
+56 -20
View File
@@ -77,11 +77,36 @@ def wait_for_gpu_memory_to_clear(
threshold_ratio: float | None = None,
timeout_s: float = 120,
) -> None:
import gc
assert threshold_bytes is not None or threshold_ratio is not None
# Use nvml instead of pytorch to reduce measurement error from torch cuda
# context.
devices = get_physical_device_indices(devices)
start_time = time.time()
# Print waiting start information
device_list = ", ".join(str(d) for d in devices)
if threshold_bytes is not None:
threshold_str = f"{threshold_bytes / 2**30:.2f} GiB"
condition_str = f"Memory usage ≤ {threshold_str}"
else:
threshold_percent = threshold_ratio * 100
threshold_str = f"{threshold_percent:.1f}%"
condition_str = f"Memory usage ratio ≤ {threshold_str}"
print(f"[GPU Memory Monitor] Waiting for GPU {device_list} to free memory, Condition: {condition_str}")
# Define the is_free function based on threshold type
if threshold_bytes is not None:
def is_free(used, total):
return used <= threshold_bytes / 2**30
else:
def is_free(used, total):
return used / total <= threshold_ratio
while True:
output: dict[int, str] = {}
output_raw: dict[int, tuple[float, float]] = {}
@@ -97,33 +122,44 @@ def wait_for_gpu_memory_to_clear(
gb_used = mem_info.used / 2**30
gb_total = mem_info.total / 2**30
output_raw[device] = (gb_used, gb_total)
output[device] = f"{gb_used:.02f}/{gb_total:.02f}"
# Format to more readable form
usage_percent = (gb_used / gb_total) * 100 if gb_total > 0 else 0
output[device] = f"{gb_used:.1f}GiB/{gb_total:.1f}GiB ({usage_percent:.1f}%)"
print("gpu memory used/total (GiB): ", end="")
for k, v in output.items():
print(f"{k}={v}; ", end="")
print("")
if threshold_bytes is not None:
def is_free(used, total):
return used <= threshold_bytes / 2**30 # noqa E731
threshold = f"{threshold_bytes / 2**30} GiB"
else:
def is_free(used, total):
return used / total <= threshold_ratio # noqa E731
threshold = f"{threshold_ratio:.2f}"
# Optimized GPU memory status print
print("[GPU Memory Status] Current usage:")
for device_id, mem_info in output.items():
print(f" GPU {device_id}: {mem_info}")
# Calculate waiting duration
dur_s = time.time() - start_time
elapsed_minutes = dur_s / 60
# Check if all devices meet the condition
if all(is_free(used, total) for used, total in output_raw.values()):
print(f"Done waiting for free GPU memory on devices {devices=} ({threshold=}) {dur_s=:.02f}")
# Optimized completion message
print(f"[GPU Memory Freed] Devices {device_list} meet memory condition")
print(f" Condition: {condition_str}")
print(f" Wait time: {dur_s:.1f} seconds ({elapsed_minutes:.1f} minutes)")
print(" Final status:")
for device_id, mem_info in output.items():
print(f" GPU {device_id}: {mem_info}")
break
# Check timeout
if dur_s >= timeout_s:
raise ValueError(f"Memory of devices {devices=} not free after {dur_s=:.02f} ({threshold=})")
raise ValueError(
f"[GPU Memory Timeout] Devices {device_list} still don't meet memory condition after {dur_s:.1f} seconds\n"
f"Condition: {condition_str}\n"
f"Current status:\n" + "\n".join(f" GPU {device}: {output[device]}" for device in devices)
)
# Add waiting hint (optional)
if dur_s > 10 and int(dur_s) % 10 == 0: # Show hint every 10 seconds
print(f"Waiting... Already waited {dur_s:.1f} seconds ({elapsed_minutes:.1f} minutes)")
gc.collect()
torch.cuda.empty_cache()
time.sleep(5)
vllm_omni/diffusion/attention/backends/abstract.py
+54 -2
View File
@@ -7,12 +7,18 @@ from typing import Generic, TypeVar
import torch
from vllm_omni.platforms import current_omni_platform
class AttentionBackend(ABC):
"""Abstract class for diffusion attention backends."""
accept_output_buffer: bool = False
@classmethod
def supports_attention_mask(cls) -> bool:
return False
@staticmethod
@abstractmethod
def get_name() -> str:
@@ -77,12 +83,58 @@ class AttentionImpl(ABC, Generic[T]):
) -> None:
raise NotImplementedError
@abstractmethod
def forward(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attn_metadata: T,
attn_metadata: T | None = None,
) -> torch.Tensor:
"""Dispatch to platform-specific forward implementation."""
if current_omni_platform.is_rocm():
return self.forward_hip(query, key, value, attn_metadata)
elif current_omni_platform.is_cuda():
return self.forward_cuda(query, key, value, attn_metadata)
elif current_omni_platform.is_npu():
return self.forward_npu(query, key, value, attn_metadata)
elif current_omni_platform.is_xpu():
return self.forward_xpu(query, key, value, attn_metadata)
else:
raise NotImplementedError(f"No forward implementation for platform: {current_omni_platform}")
def forward_cuda(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attn_metadata: T | None = None,
) -> torch.Tensor:
raise NotImplementedError
def forward_npu(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attn_metadata: T | None = None,
) -> torch.Tensor:
raise NotImplementedError
def forward_xpu(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attn_metadata: T | None = None,
) -> torch.Tensor:
raise NotImplementedError
def forward_hip(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attn_metadata: T | None = None,
) -> torch.Tensor:
# By default, HIP ops are compatible with CUDA ops.
return self.forward_cuda(query, key, value, attn_metadata)
vllm_omni/diffusion/attention/backends/ascend_attn.py
-92
View File
@@ -1,92 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from importlib.util import find_spec
import torch
from vllm.logger import init_logger
from vllm_omni.diffusion.attention.backends.abstract import (
AttentionBackend,
AttentionImpl,
AttentionMetadata,
)
logger = init_logger(__name__)
class AscendAttentionBackend(AttentionBackend):
accept_output_buffer: bool = True
@staticmethod
def get_supported_head_sizes() -> list[int]:
return [x for x in range(1024)] # todo
@staticmethod
def get_name() -> str:
return "ASCEND"
@staticmethod
def get_impl_cls() -> type["AscendAttentionBackendImpl"]:
return AscendAttentionBackendImpl
class AscendAttentionBackendImpl(AttentionImpl):
def __init__(
self,
num_heads: int,
head_size: int,
softmax_scale: float,
causal: bool = False,
num_kv_heads: int | None = None,
prefix: str = "",
**extra_impl_args,
) -> None:
self.causal = causal
self.softmax_scale = softmax_scale
def forward_native(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attn_metadata: AttentionMetadata = None,
) -> torch.Tensor:
query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value))
attention_mask = attn_metadata.attn_mask if attn_metadata else None
output = torch.nn.functional.scaled_dot_product_attention(
query,
key,
value,
attn_mask=attention_mask,
dropout_p=0.0,
is_causal=self.causal,
scale=self.softmax_scale,
)
out = output.permute(0, 2, 1, 3)
return out
def forward(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attn_metadata: AttentionMetadata = None,
) -> torch.Tensor:
if find_spec("mindiesd"):
from mindiesd import attention_forward
attention_mask = attn_metadata.attn_mask if attn_metadata else None
output = attention_forward(
query,
key,
value,
attn_mask=attention_mask,
opt_mode="manual",
op_type="fused_attn_score",
layout="BNSD",
)
return output
else:
return self.forward_native(query, key, value, attn_metadata)
vllm_omni/diffusion/attention/backends/flash_attn.py
+43 -14
View File
@@ -9,7 +9,11 @@ from vllm_omni.diffusion.attention.backends.abstract import (
AttentionImpl,
AttentionMetadata,
)
# Import flash attention functions with fallback chain from utils/fa.py
# FA3 (fa3_fwd_interface) -> FA3 (flash_attn_interface) -> FA2 (flash_attn)
from vllm_omni.diffusion.attention.backends.utils.fa import (
HAS_FLASH_ATTN,
_pad_input,
_unpad_input,
_upad_input,
@@ -19,10 +23,21 @@ from vllm_omni.diffusion.attention.backends.utils.fa import (
logger = init_logger(__name__)
if not HAS_FLASH_ATTN:
raise ImportError(
"FlashAttentionBackend requires Flash Attention. "
"Please install one of: fa3-fwd, flash-attention, or flash-attn. "
"Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
)
class FlashAttentionBackend(AttentionBackend):
accept_output_buffer: bool = True
@classmethod
def supports_attention_mask(cls) -> bool:
return True
@staticmethod
def get_supported_head_sizes() -> list[int]:
return [64, 96, 128, 192, 256]
@@ -51,25 +66,14 @@ class FlashAttentionImpl(AttentionImpl):
self.causal = causal
self.softmax_scale = softmax_scale
def forward(
def forward_cuda(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attn_metadata: AttentionMetadata = None,
) -> torch.Tensor:
"""
Flash attention implementation.
Args:
query: (batch_size, seq_len, num_heads, head_dim)
key: (batch_size, seq_len, num_heads, head_dim)
value: (batch_size, seq_len, num_heads, head_dim)
attn_metadata: AttentionMetadata. Attention mask is supported as attn_metadata.attn_mask
Returns:
(batch_size, seq_len, num_heads, head_dim)
"""
"""CUDA/ROCm flash attention implementation."""
query_length = query.size(1)
attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
# Contains at least one padding token in the sequence
@@ -98,11 +102,36 @@ class FlashAttentionImpl(AttentionImpl):
out = _pad_input(out_unpad, indices_q, query.size(0), query_length)
else:
out: torch.Tensor = flash_attn_func(
out = flash_attn_func(
query,
key,
value,
causal=self.causal,
softmax_scale=self.softmax_scale,
)
# FA3 may return (out, lse) tuple, FA2 returns just out
if isinstance(out, tuple):
out = out[0]
return out
def forward_npu(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attn_metadata: AttentionMetadata = None,
) -> torch.Tensor:
"""NPU attention implementation using mindiesd."""
from mindiesd import attention_forward
attention_mask = attn_metadata.attn_mask if attn_metadata else None
output = attention_forward(
query,
key,
value,
attn_mask=attention_mask,
opt_mode="manual",
op_type="fused_attn_score",
layout="BNSD",
)
return output
vllm_omni/diffusion/attention/backends/registry.py
+152
View File
@@ -0,0 +1,152 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Diffusion attention backend registry.
This module provides an enum-based registry for diffusion attention backends,
similar to vLLM's AttentionBackendEnum. Each backend registers its class path,
and platforms can override or extend backends using register_backend().
"""
from collections.abc import Callable
from enum import Enum, EnumMeta
from typing import TYPE_CHECKING
from vllm.logger import init_logger
from vllm.utils.import_utils import resolve_obj_by_qualname
if TYPE_CHECKING:
from vllm_omni.diffusion.attention.backends.abstract import AttentionBackend
logger = init_logger(__name__)
class _DiffusionBackendEnumMeta(EnumMeta):
"""Metaclass for DiffusionAttentionBackendEnum to provide better error messages."""
def __getitem__(cls, name: str) -> "DiffusionAttentionBackendEnum":
"""Get backend by name with helpful error messages."""
try:
return super().__getitem__(name) # type: ignore[return-value]
except KeyError:
members = list(cls.__members__.keys())
valid_backends = ", ".join(members)
raise ValueError(
f"Unknown diffusion attention backend: '{name}'. Valid options are: {valid_backends}"
) from None
class DiffusionAttentionBackendEnum(Enum, metaclass=_DiffusionBackendEnumMeta):
"""Enumeration of all supported diffusion attention backends.
The enum value is the default class path, but this can be overridden
at runtime using register_backend().
To get the actual backend class (respecting overrides), use:
backend.get_class()
Example:
# Get backend class
backend = DiffusionAttentionBackendEnum.FLASH_ATTN
backend_cls = backend.get_class()
# Register custom backend
@register_diffusion_backend(DiffusionAttentionBackendEnum.CUSTOM)
class MyCustomBackend:
...
"""
# Common backends (available on most platforms)
FLASH_ATTN = "vllm_omni.diffusion.attention.backends.flash_attn.FlashAttentionBackend"
TORCH_SDPA = "vllm_omni.diffusion.attention.backends.sdpa.SDPABackend"
SAGE_ATTN = "vllm_omni.diffusion.attention.backends.sage_attn.SageAttentionBackend"
def get_path(self, include_classname: bool = True) -> str:
"""Get the class path for this backend (respects overrides).
Returns:
The fully qualified class path string
Raises:
ValueError: If backend has empty path and is not registered
"""
path = _DIFFUSION_ATTN_OVERRIDES.get(self, self.value)
if not path:
raise ValueError(
f"Backend {self.name} must be registered before use. "
f"Use register_diffusion_backend(DiffusionAttentionBackendEnum.{self.name}, "
f"'your.module.YourClass')"
)
if not include_classname:
path = path.rsplit(".", 1)[0]
return path
def get_class(self) -> "type[AttentionBackend]":
"""Get the backend class (respects overrides).
Returns:
The backend class
Raises:
ImportError: If the backend class cannot be imported
ValueError: If backend has empty path and is not registered
"""
return resolve_obj_by_qualname(self.get_path())
def is_overridden(self) -> bool:
"""Check if this backend has been overridden.
Returns:
True if the backend has a registered override
"""
return self in _DIFFUSION_ATTN_OVERRIDES
def clear_override(self) -> None:
"""Clear any override for this backend, reverting to the default."""
_DIFFUSION_ATTN_OVERRIDES.pop(self, None)
# Override registry
_DIFFUSION_ATTN_OVERRIDES: dict[DiffusionAttentionBackendEnum, str] = {}
def register_diffusion_backend(
backend: DiffusionAttentionBackendEnum,
class_path: str | None = None,
) -> Callable[[type], type]:
"""Register or override a diffusion backend implementation.
Args:
backend: The DiffusionAttentionBackendEnum member to register
class_path: Optional class path. If not provided and used as
decorator, will be auto-generated from the class.
Returns:
Decorator function if class_path is None, otherwise a no-op
Examples:
# Override an existing backend
@register_diffusion_backend(DiffusionAttentionBackendEnum.FLASH_ATTN)
class MyCustomFlashAttn:
...
# Override an existing backend (e.g., ASCEND_ATTN)
@register_diffusion_backend(DiffusionAttentionBackendEnum.ASCEND_ATTN)
class CustomAscendAttentionBackend:
...
# Direct registration
register_diffusion_backend(
DiffusionAttentionBackendEnum.CUSTOM,
"my.module.MyCustomBackend"
)
"""
def decorator(cls: type) -> type:
_DIFFUSION_ATTN_OVERRIDES[backend] = f"{cls.__module__}.{cls.__qualname__}"
return cls
if class_path is not None:
_DIFFUSION_ATTN_OVERRIDES[backend] = class_path
return lambda x: x
return decorator
vllm_omni/diffusion/attention/backends/ring/ring_globals.py
+39 -16
View File
@@ -3,30 +3,53 @@
# Copyright (c) 2024, Jiarui Fang.
# Adapted from https://github.com/feifeibear/long-context-attention
# test if flash_attn is available
# test if flash_attn (FA2) is available
try:
import flash_attn # noqa: F401
from flash_attn.flash_attn_interface import _flash_attn_backward, _flash_attn_forward # noqa: F401
from flash_attn.flash_attn_interface import _flash_attn_forward # noqa: F401
HAS_FLASH_ATTN = True
except ImportError:
except (ImportError, ModuleNotFoundError):
HAS_FLASH_ATTN = False
try:
from flash_attn_interface import _flash_attn_backward as flash_attn_func_hopper_backward # noqa: F401
from flash_attn_interface import _flash_attn_forward as flash_attn_forward_hopper # noqa: F401
from flash_attn_interface import flash_attn_func as flash3_attn_func # noqa: F401
# FA3 detection: try multiple sources (forward only, no backward needed for inference)
# Source 1: flash_attn_interface (from flash-attention source build)
# Source 2: fa3_fwd_interface (from fa3-fwd PyPI package, supports Ampere/Ada/Hopper)
# Note: FA3 high-level API may or may not return softmax_lse depending on version.
# For Ring Attention which requires LSE, we fall back to low-level API if needed.
HAS_FA3 = False
fa3_fwd_func = None # Low-level forward function (_flash_attn_forward)
fa3_attn_func = None # High-level attention function (flash_attn_func)
HAS_FLASH_ATTN_HOPPER = True
except ImportError:
HAS_FLASH_ATTN_HOPPER = False
# Try flash_attn_interface first (from flash-attention source build)
try:
from flash_attn_interface import _flash_attn_forward as fa3_fwd_func # noqa: F401
from flash_attn_interface import flash_attn_func as fa3_attn_func # noqa: F401
HAS_FA3 = True
except (ImportError, ModuleNotFoundError):
pass
# Fallback: try fa3_fwd_interface (PyPI package, supports Ampere/Ada/Hopper)
if not HAS_FA3:
try:
from fa3_fwd_interface import _flash_attn_forward as fa3_fwd_func # noqa: F401
from fa3_fwd_interface import flash_attn_func as fa3_attn_func # noqa: F401
HAS_FA3 = True
except (ImportError, ModuleNotFoundError):
pass
# Legacy aliases for backward compatibility
HAS_FLASH_ATTN_HOPPER = HAS_FA3
flash_attn_forward_hopper = fa3_fwd_func
flash3_attn_func = fa3_attn_func
try:
from flashinfer.prefill import single_prefill_with_kv_cache # noqa: F401
HAS_FLASHINFER = True
except ImportError:
except (ImportError, ModuleNotFoundError):
HAS_FLASHINFER = False
try:
@@ -34,26 +57,26 @@ try:
from aiter import flash_attn_func as flash_attn_func_aiter # noqa: F401
HAS_AITER = True
except ImportError:
except (ImportError, ModuleNotFoundError):
HAS_AITER = False
try:
import sageattention # noqa: F401
HAS_SAGE_ATTENTION = True
except ImportError:
except (ImportError, ModuleNotFoundError):
HAS_SAGE_ATTENTION = False
try:
import spas_sage_attn # noqa: F401
HAS_SPARSE_SAGE_ATTENTION = True
except ImportError:
except (ImportError, ModuleNotFoundError):
HAS_SPARSE_SAGE_ATTENTION = False
try:
import torch_npu # noqa: F401
HAS_NPU = True
except ImportError:
except (ImportError, ModuleNotFoundError):
HAS_NPU = False
vllm_omni/diffusion/attention/backends/ring/ring_kernels.py
+29 -69
View File
@@ -6,7 +6,13 @@ import math
import torch
from .ring_globals import HAS_AITER, HAS_FLASH_ATTN, HAS_FLASH_ATTN_HOPPER, HAS_FLASHINFER, HAS_NPU
from .ring_globals import (
HAS_AITER,
HAS_FA3,
HAS_FLASH_ATTN,
HAS_FLASHINFER,
fa3_fwd_func,
)
_scaled_dot_product_flash_attention = torch.ops.aten._scaled_dot_product_flash_attention
_scaled_dot_product_efficient_attention = torch.ops.aten._scaled_dot_product_efficient_attention
@@ -26,23 +32,11 @@ if HAS_FLASH_ATTN:
import flash_attn
from flash_attn.flash_attn_interface import _flash_attn_forward
if HAS_FLASH_ATTN_HOPPER:
from flash_attn_interface import _flash_attn_backward as flash_attn_func_hopper_backward
from flash_attn_interface import _flash_attn_forward as flash_attn_forward_hopper
from flash_attn_interface import flash_attn_func as flash3_attn_func
else:
flash_attn_forward_hopper = None
flash_attn_func_hopper_backward = None
flash3_attn_func = None
if HAS_FLASHINFER:
from flashinfer.prefill import single_prefill_with_kv_cache
_LOG2_E = math.log2(math.e)
if HAS_NPU:
import torch_npu
def pytorch_attn_forward(
q: torch.Tensor,
@@ -146,52 +140,35 @@ def flash_attn_forward(
return block_out, block_lse
def flash_attn3_func_forward(
q, k, v, dropout_p, softmax_scale, causal, window_size, softcap, alibi_slopes, return_softmax
):
assert HAS_FLASH_ATTN_HOPPER
# current signature of flash_attn_forward_hopper:
# (q, k, v, softmax_scale, causal, window_size, descale_q=None, descale_k=None, descale_v=None, gqa_parallel=False)
def fa3_forward(q, k, v, dropout_p, softmax_scale, causal, window_size, softcap, alibi_slopes, return_softmax):
"""FA3 forward pass for inference.
out, softmax_lse, *unused = flash_attn_forward_hopper(
q=q,
k=k,
v=v,
k_new=None,
v_new=None,
qv=None,
out=None,
cu_seqlens_q=None,
cu_seqlens_k=None,
cu_seqlens_k_new=None,
seqused_q=None,
seqused_k=None,
max_seqlen_q=None,
max_seqlen_k=None,
page_table=None,
kv_batch_idx=None,
leftpad_k=None,
rotary_cos=None,
rotary_sin=None,
seqlens_rotary=None,
q_descale=None,
k_descale=None,
v_descale=None,
FA3 supports Ampere, Ada, and Hopper GPUs. Dropout is ignored since FA3 is inference-only.
Uses low-level API (_flash_attn_forward) which always returns softmax_lse,
required for Ring Attention's correct accumulation.
"""
assert HAS_FA3, "FA3 is not available"
assert fa3_fwd_func is not None, "FA3 low-level API (fa3_fwd_func) not available"
# Low-level API always returns (out, softmax_lse, S_dmask, rng_state)
out, softmax_lse, *_ = fa3_fwd_func(
q,
k,
v,
softmax_scale=softmax_scale,
causal=False,
window_size=(-1, -1),
attention_chunk=0,
softcap=0.0,
rotary_interleaved=True,
scheduler_metadata=None,
num_splits=0,
pack_gqa=None,
sm_margin=0,
causal=causal,
window_size_left=window_size[0] if window_size else -1,
window_size_right=window_size[1] if window_size else -1,
softcap=softcap if softcap else 0.0,
)
return out, softmax_lse
# Legacy alias for backward compatibility
flash_attn3_func_forward = fa3_forward
def flash_attn_forward_aiter(
q,
k,
@@ -264,20 +241,3 @@ def flashinfer_attn_forward(
raise ValueError(f"Invalid input shape: {q.shape}")
lse = lse / _LOG2_E
return out, lse
def npu_attn_forward(q, k, v, softmax_scale=None, layout="BSND"):
assert HAS_NPU, "torch_npu is not available"
softmax_scale = q.shape[-1] ** (-0.5)
block_out, block_lse = torch_npu.npu_fused_infer_attention_score(
q,
k,
v,
num_heads=q.shape[-2],
input_layout=layout,
scale=softmax_scale,
softmax_lse_flag=True,
pre_tokens=65535,
next_tokens=65535,
)
return block_out, block_lse.squeeze(dim=-1)
vllm_omni/diffusion/attention/backends/ring/ring_selector.py
-10
View File
@@ -10,7 +10,6 @@ from functools import partial
import torch
from .ring_globals import (
HAS_NPU,
HAS_SAGE_ATTENTION,
HAS_SPARSE_SAGE_ATTENTION,
)
@@ -28,9 +27,6 @@ if HAS_SAGE_ATTENTION:
if HAS_SPARSE_SAGE_ATTENTION:
from spas_sage_attn.autotune import SparseAttentionMeansim
if HAS_NPU:
from torch_npu import npu_fused_infer_attention_score
class AttnType(Enum):
AITER = "aiter"
@@ -44,7 +40,6 @@ class AttnType(Enum):
SAGE_FP8 = "sage_fp8"
SAGE_FP8_SM90 = "sage_fp8_sm90"
SPARSE_SAGE = "sparse_sage"
NPU = "npu"
@classmethod
def from_string(cls, s: str):
@@ -157,11 +152,6 @@ def select_flash_attn_impl(
return fn
elif impl_type == AttnType.NPU:
if not HAS_NPU:
raise ImportError("torch_npu is not available!")
return npu_fused_infer_attention_score
elif attn_processor is not None:
return attn_processor
vllm_omni/diffusion/attention/backends/sage_attn.py
+1 -1
View File
@@ -54,7 +54,7 @@ class SageAttentionImpl(AttentionImpl):
self.causal = causal
self.softmax_scale = softmax_scale
def forward(
def forward_cuda(
self,
query: torch.Tensor,
key: torch.Tensor,
vllm_omni/diffusion/attention/backends/sdpa.py
+4
View File
@@ -16,6 +16,10 @@ logger = init_logger(__name__)
class SDPABackend(AttentionBackend):
accept_output_buffer: bool = True
@classmethod
def supports_attention_mask(cls) -> bool:
return True
@staticmethod
def get_supported_head_sizes() -> list[int]:
return [x for x in range(1024)] # todo
vllm_omni/diffusion/attention/backends/utils/fa.py
+46 -10
View File
@@ -15,19 +15,55 @@
import torch
import torch.nn.functional as F
from vllm_omni.utils.platform_utils import is_rocm
from vllm_omni.platforms import current_omni_platform
if is_rocm():
from vllm._aiter_ops import is_aiter_found_and_supported
# Flash Attention function detection with fallback chain
flash_attn_func = None
flash_attn_varlen_func = None
# Choose to enable this by default on ROCm
# Whenever possible as it is the fastest backend
if is_aiter_found_and_supported():
from aiter import flash_attn_func, flash_attn_varlen_func # noqa: F401
else:
raise ImportError("Aiter is not found and supported on currentROCm device.")
if current_omni_platform.is_rocm():
# ROCm: try Aiter first
try:
from vllm._aiter_ops import is_aiter_found_and_supported
if is_aiter_found_and_supported():
from aiter import flash_attn_func, flash_attn_varlen_func # noqa: F401
except (ImportError, ModuleNotFoundError):
pass
else:
from fa3_fwd_interface import flash_attn_func, flash_attn_varlen_func # noqa: F401
# CUDA: try FA3 -> FA2 fallback chain
# Try FA3 from fa3-fwd PyPI package
try:
from fa3_fwd_interface import flash_attn_func, flash_attn_varlen_func # noqa: F401
except (ImportError, ModuleNotFoundError):
pass
# Fallback: Try FA3 from flash-attention source build
if flash_attn_func is None:
try:
from flash_attn_interface import flash_attn_func, flash_attn_varlen_func # noqa: F401
except (ImportError, ModuleNotFoundError):
pass
# Fallback: Try FA2 from flash-attn package (try multiple import paths)
if flash_attn_func is None:
try:
from flash_attn import flash_attn_func, flash_attn_varlen_func # noqa: F401
except (ImportError, ModuleNotFoundError):
pass
if flash_attn_func is None:
try:
from flash_attn.flash_attn_interface import ( # noqa: F401
flash_attn_func,
flash_attn_varlen_func,
)
except (ImportError, ModuleNotFoundError):
pass
# If no FA backend available, SDPA backend will be selected at the platform level
# flash_attn_func and flash_attn_varlen_func will be None
HAS_FLASH_ATTN = flash_attn_func is not None
def _index_first_axis(tensor, indices):
vllm_omni/diffusion/attention/layer.py
+13 -4
View File
@@ -12,6 +12,7 @@ import torch.nn as nn
from vllm.logger import init_logger
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
from vllm_omni.diffusion.attention.parallel import build_parallel_attention_strategy
from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention
from vllm_omni.diffusion.attention.selector import get_attn_backend
@@ -45,6 +46,14 @@ class Attention(nn.Module):
causal=causal,
num_kv_heads=num_kv_heads,
)
# Instantiate fallback backend for float32 support
self.sdpa_fallback = SDPABackend.get_impl_cls()(
num_heads=num_heads,
head_size=head_size,
softmax_scale=softmax_scale,
causal=causal,
num_kv_heads=num_kv_heads,
)
self.backend_pref = None
self.softmax_scale = softmax_scale
@@ -104,12 +113,12 @@ class Attention(nn.Module):
return out
def _run_local_attention(self, query, key, value, attn_metadata):
if self.backend_pref == "flash_attn" and query.dtype == torch.float32:
logger.warning(
"Flash Attention does not support float32. Overriding user config "
if query.dtype == torch.float32:
logger.warning_once(
f"Only SDPA supports float32. Overriding user config {type(self.attention)} "
f"attention_backend='{self.backend_pref}' to 'sdpa' for dtype={query.dtype}."
)
self.backend_pref = "sdpa"
return self.sdpa_fallback.forward(query, key, value, attn_metadata)
# Fallback to standard attention
return self.attention.forward(query, key, value, attn_metadata)
vllm_omni/diffusion/attention/parallel/ring.py
+12 -7
View File
@@ -10,7 +10,7 @@ import torch
from vllm.logger import init_logger
# import torch.distributed as dist # Not used directly here, but good practice if needed
from vllm_omni.diffusion.attention.backends.ring.ring_globals import HAS_FLASH_ATTN
from vllm_omni.diffusion.attention.backends.ring.ring_globals import HAS_FA3, HAS_FLASH_ATTN
from vllm_omni.diffusion.attention.backends.ring.ring_selector import AttnType
from vllm_omni.diffusion.attention.parallel.base import (
ParallelAttentionContext,
@@ -116,11 +116,14 @@ class RingParallelAttention:
except Exception:
backend_pref = None
# Fallback for FP32 or if Flash Attention is not available
if query.dtype == torch.float32 or not HAS_FLASH_ATTN:
if not HAS_FLASH_ATTN and backend_pref != "sdpa":
# Determine attention type with fallback chain: FA3 -> FA2 -> SDPA
# FP32 is not supported by Flash Attention, force SDPA
if query.dtype == torch.float32:
backend_pref = "sdpa"
elif not HAS_FA3 and not HAS_FLASH_ATTN:
if backend_pref != "sdpa":
logger = init_logger(__name__)
logger.warning_once("Flash Attention is not available! Force enabling SDPA.")
logger.warning_once("Flash Attention (FA2/FA3) is not available! Force enabling SDPA.")
backend_pref = "sdpa"
# Extract joint tensors
@@ -150,6 +153,9 @@ class RingParallelAttention:
from vllm_omni.diffusion.attention.backends.ring_flash_attn import ring_flash_attn_func
# Prefer FA3 over FA2 for better performance (FA3 supports Ampere/Ada/Hopper)
attn_type = AttnType.FA3 if HAS_FA3 else AttnType.FA
return ring_flash_attn_func(
query,
key,
@@ -161,9 +167,8 @@ class RingParallelAttention:
softcap=0.0,
alibi_slopes=None,
deterministic=False,
# return_attn_probs=False, # Removed as it might not be supported in signature
group=self._sp_group.ring_group,
attn_type=AttnType.FA,
attn_type=attn_type,
joint_tensor_key=joint_key,
joint_tensor_value=joint_value,
joint_strategy=joint_strategy,
vllm_omni/diffusion/attention/selector.py
+47 -80
View File
@@ -1,52 +1,54 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Diffusion attention backend selector.
This module provides the interface for selecting diffusion attention backends.
The actual backend selection logic is delegated to the platform layer
(vllm_omni.platforms), similar to how vLLM handles attention backend selection.
Usage:
from vllm_omni.diffusion.attention.selector import get_attn_backend
# Get the appropriate backend for current platform
backend_cls = get_attn_backend(head_size=64)
# Or override via environment variable
# export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
"""
import importlib
import os
from functools import cache
import torch
from vllm.logger import init_logger
from vllm_omni.diffusion.attention.backends.abstract import (
AttentionBackend,
)
from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
from vllm_omni.utils.platform_utils import detect_device_type, is_rocm
logger = init_logger(__name__)
# environment variable value -> backend module and class
_BACKEND_CONFIG = {
"FLASH_ATTN": {
"module": "vllm_omni.diffusion.attention.backends.flash_attn",
"class": "FlashAttentionBackend",
},
"TORCH_SDPA": {
"module": "vllm_omni.diffusion.attention.backends.sdpa",
"class": "SDPABackend",
},
"SAGE_ATTN": {
"module": "vllm_omni.diffusion.attention.backends.sage_attn",
"class": "SageAttentionBackend",
},
"ASCEND": {"module": "vllm_omni.diffusion.attention.backends.ascend_attn", "class": "AscendAttentionBackend"},
}
_BACKENDS_SUPPORT_ATTENTION_MASK = ["SDPA", "ASCEND", "FLASH_ATTN"]
def _load_backend_cls(cls_path: str) -> type[AttentionBackend]:
"""Load a backend class from its fully qualified path.
Args:
cls_path: Fully qualified class path (e.g.,
"vllm_omni.diffusion.attention.backends.sdpa.SDPABackend")
def load_backend(backend_name: str) -> type[AttentionBackend]:
config = _BACKEND_CONFIG[backend_name]
Returns:
The loaded backend class
"""
module_path, class_name = cls_path.rsplit(".", 1)
try:
module = importlib.import_module(config["module"])
backend_class = getattr(module, config["class"])
module = importlib.import_module(module_path)
backend_class = getattr(module, class_name)
return backend_class
except ImportError as e:
raise ImportError(f"Failed to import module {config['module']}: {e}")
raise ImportError(f"Failed to import module {module_path}: {e}")
except AttributeError as e:
raise AttributeError(f"Class {config['class']} not found in module: {e}")
raise AttributeError(f"Class {class_name} not found in module: {e}")
@cache
@@ -54,65 +56,30 @@ def get_attn_backend(head_size: int) -> type[AttentionBackend]:
"""
Get attention backend for diffusion models.
The backend is selected based on the following priority:
1. DIFFUSION_ATTENTION_BACKEND environment variable (if set, e.g. export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN)
2. Default backend (SDPA)
The backend selection is delegated to the current platform
(vllm_omni.platforms.current_omni_platform), which selects the
appropriate backend based on:
1. User override via DIFFUSION_ATTENTION_BACKEND environment variable
2. Platform-specific defaults and capabilities
This is similar to how vLLM's get_attn_backend_cls works, where the
platform layer decides which backend to use based on hardware capabilities.
Args:
head_size: Head size (currently not used for selection, but kept for API compatibility)
head_size: Head size for attention computation (may affect backend selection)
Returns:
The selected attention backend class
"""
# Check environment variable
from vllm_omni.platforms import current_omni_platform
backend_name = os.environ.get("DIFFUSION_ATTENTION_BACKEND", None)
# Check environment variable for user override
selected_backend = os.environ.get("DIFFUSION_ATTENTION_BACKEND")
if detect_device_type() == "cuda" and not is_rocm():
compute_capability = torch.cuda.get_device_capability()
major, minor = compute_capability
if 80 <= major * 10 + minor < 100:
if backend_name is None:
backend_name = "FLASH_ATTN"
else:
if backend_name == "FLASH_ATTN":
logger.warning(
"""Flash Attention requires GPU with compute capability >= 8.0 or < 10.0. "
"Falling back to TORCH_SDPA backend."""
)
backend_name = "TORCH_SDPA"
elif detect_device_type() == "cuda" and is_rocm():
from vllm._aiter_ops import is_aiter_found_and_supported
# Delegate to platform for backend selection
backend_cls_path = current_omni_platform.get_diffusion_attn_backend_cls(
selected_backend=selected_backend,
head_size=head_size,
)
compute_capability = torch.cuda.get_device_capability()
major, minor = compute_capability
# Choose to enable this by default on ROCm
# Whenever possible as it is the fastest backend
# is_aiter_found_and_supported() checks if aiter library is found
# and is aiter supported on the current platform
# aiter currently only is supported on gfx942 and gfx950
# https://github.com/vllm-project/vllm/blob/main/vllm/_aiter_ops.py
if is_aiter_found_and_supported() and 90 < major * 10 + minor < 100:
if backend_name is None:
backend_name = "FLASH_ATTN"
else:
if backend_name == "FLASH_ATTN":
logger.warning(
"Flash Attention requires `aiter` library which is only supported "
"on gfx942 and gfx950. Falling back to TORCH_SDPA backend."
)
backend_name = "TORCH_SDPA"
if backend_name is not None:
backend_name_upper = backend_name.upper()
if backend_name_upper not in _BACKEND_CONFIG:
valid_backends = list(_BACKEND_CONFIG.keys())
raise ValueError(
f"Invalid attention backend for diffusion: '{backend_name}'. Valid backends are: {valid_backends}"
)
logger.info(f"Using attention backend '{backend_name_upper}' for diffusion")
return load_backend(backend_name_upper)
# Default to SDPA
return SDPABackend
return _load_backend_cls(backend_cls_path)
vllm_omni/diffusion/cache/cache_dit_backend.py
Vendored
+60 -51
View File
@@ -28,6 +28,13 @@ from vllm_omni.diffusion.data import DiffusionCacheConfig, OmniDiffusionConfig
logger = init_logger(__name__)
# Small helper to centralize cache-dit summaries.
def cache_summary(pipeline: Any, details: bool = True) -> None:
cache_dit.summary(pipeline.transformer, details=details)
if hasattr(pipeline, "transformer_2"):
cache_dit.summary(pipeline.transformer_2, details=details)
# Registry of custom cache-dit enablers for specific models
# Maps pipeline names to their cache-dit enablement functions
# Models in this registry require custom handling (e.g., dual-transformer architectures)
@@ -71,53 +78,46 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int],
A refresh function that can be called to update cache context with new num_inference_steps.
"""
# Build DBCacheConfig for primary transformer
primary_cache_config = _build_db_cache_config(cache_config)
# FIXME: secondary cache shares the same config with primary cache for now,
# but we should support different config for secondary transformer in the future
# Build DBCacheConfig for secondary transformer (can use same or different config)
secondary_cache_config = _build_db_cache_config(cache_config)
# Build calibrator configs if TaylorSeer is enabled
primary_calibrator = None
secondary_calibrator = None
if cache_config.enable_taylorseer:
taylorseer_order = cache_config.taylorseer_order
primary_calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
secondary_calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
# Build ParamsModifier for each transformer
primary_modifier = ParamsModifier(
cache_config=primary_cache_config,
calibrator_config=primary_calibrator,
)
secondary_modifier = ParamsModifier(
cache_config=secondary_cache_config,
calibrator_config=secondary_calibrator,
)
logger.info(
"Enabling cache-dit on Wan2.2 dual transformers with BlockAdapter: "
f"Fn={primary_cache_config.Fn_compute_blocks}, "
f"Bn={primary_cache_config.Bn_compute_blocks}, "
f"W={primary_cache_config.max_warmup_steps}, "
)
transformer = pipeline.transformer
transformer_2 = pipeline.transformer_2
transformer_blocks = transformer.blocks
transformer_2_blocks = transformer_2.blocks
# Enable cache-dit using BlockAdapter for both transformers simultaneously
cache_dit.enable_cache(
BlockAdapter(
transformer=[transformer, transformer_2],
blocks=[transformer_blocks, transformer_2_blocks],
forward_pattern=[ForwardPattern.Pattern_2, ForwardPattern.Pattern_2],
params_modifiers=[primary_modifier, secondary_modifier],
transformer=[
pipeline.transformer,
pipeline.transformer_2,
],
blocks=[
pipeline.transformer.blocks,
pipeline.transformer_2.blocks,
],
forward_pattern=[
ForwardPattern.Pattern_2,
ForwardPattern.Pattern_2,
],
params_modifiers=[
# high-noise transformer only have 30% steps
ParamsModifier(
cache_config=DBCacheConfig().reset(
max_warmup_steps=cache_config.max_warmup_steps,
max_cached_steps=cache_config.max_cached_steps,
),
),
ParamsModifier(
cache_config=DBCacheConfig().reset(
max_warmup_steps=2,
max_cached_steps=20,
),
),
],
has_separate_cfg=True,
),
cache_config=DBCacheConfig(
Fn_compute_blocks=cache_config.Fn_compute_blocks,
Bn_compute_blocks=cache_config.Bn_compute_blocks,
max_warmup_steps=cache_config.max_warmup_steps,
max_cached_steps=cache_config.max_cached_steps,
max_continuous_cached_steps=cache_config.max_continuous_cached_steps,
residual_diff_threshold=cache_config.residual_diff_threshold,
num_inference_steps=None,
),
)
# from https://github.com/vipshop/cache-dit/pull/542
@@ -133,8 +133,8 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int],
Returns:
A tuple of (num_high_noise_steps, num_low_noise_steps).
"""
if pipeline.config.boundary_ratio is not None:
boundary_timestep = pipeline.config.boundary_ratio * pipeline.scheduler.config.num_train_timesteps
if pipeline.boundary_ratio is not None:
boundary_timestep = pipeline.boundary_ratio * pipeline.scheduler.config.num_train_timesteps
else:
boundary_timestep = None
@@ -158,10 +158,21 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int],
pipeline: The Wan2.2 pipeline instance.
num_inference_steps: New number of inference steps.
"""
num_high_noise_steps, num_low_noise_steps = _split_inference_steps(num_inference_steps)
# Refresh context for high-noise transformer
if cache_config.scm_steps_mask_policy is None:
cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_high_noise_steps, verbose=verbose)
# cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_high_noise_steps, verbose=verbose)
cache_dit.refresh_context(
pipeline.transformer,
num_inference_steps=num_high_noise_steps,
verbose=verbose,
)
cache_dit.refresh_context(
pipeline.transformer_2,
num_inference_steps=num_low_noise_steps,
verbose=verbose,
)
else:
cache_dit.refresh_context(
pipeline.transformer,
@@ -175,10 +186,6 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int],
verbose=verbose,
)
# Refresh context for low-noise transformer
if cache_config.scm_steps_mask_policy is None:
cache_dit.refresh_context(pipeline.transformer_2, num_inference_steps=num_low_noise_steps, verbose=verbose)
else:
cache_dit.refresh_context(
pipeline.transformer_2,
cache_config=DBCacheConfig().reset(
@@ -790,7 +797,9 @@ def enable_cache_for_bagel(pipeline: Any, cache_config: Any) -> Callable[[int],
# Register custom cache-dit enablers after function definitions
CUSTOM_DIT_ENABLERS.update(
{
"WanPipeline": enable_cache_for_wan22,
"Wan22Pipeline": enable_cache_for_wan22,
"Wan22I2VPipeline": enable_cache_for_wan22,
"Wan22TI2VPipeline": enable_cache_for_wan22,
"FluxPipeline": enable_cache_for_flux,
"LongCatImagePipeline": enable_cache_for_longcat_image,
"LongCatImageEditPipeline": enable_cache_for_longcat_image,
vllm_omni/diffusion/data.py
+1
View File
@@ -258,6 +258,7 @@ class OmniDiffusionConfig:
# Cache backend configuration (NEW)
cache_backend: str = "none" # "tea_cache", "deep_cache", etc.
cache_config: DiffusionCacheConfig | dict[str, Any] = field(default_factory=dict)
enable_cache_dit_summary: bool = False
# Distributed executor backend
distributed_executor_backend: str = "mp"
vllm_omni/diffusion/diffusion_engine.py
+54 -55
View File
@@ -17,6 +17,7 @@ from vllm_omni.diffusion.registry import (
get_diffusion_pre_process_func,
)
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
from vllm_omni.outputs import OmniRequestOutput
logger = init_logger(__name__)
@@ -29,6 +30,11 @@ def supports_image_input(model_class_name: str) -> bool:
return bool(getattr(model_cls, "support_image_input", False))
def image_color_format(model_class_name: str) -> str:
model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
return getattr(model_cls, "color_format", "RGB")
def supports_audio_output(model_class_name: str) -> bool:
model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
if model_cls is None:
@@ -60,36 +66,31 @@ class DiffusionEngine:
self.close()
raise e
def step(self, requests: list[OmniDiffusionRequest]):
def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
# Apply pre-processing if available
if self.pre_process_func is not None:
preprocess_start_time = time.time()
requests = self.pre_process_func(requests)
request = self.pre_process_func(request)
preprocess_time = time.time() - preprocess_start_time
logger.info(f"Pre-processing completed in {preprocess_time:.4f} seconds")
output = self.add_req_and_wait_for_response(requests)
output = self.add_req_and_wait_for_response(request)
if output.error:
raise Exception(f"{output.error}")
logger.info("Generation completed successfully.")
if output.output is None:
logger.warning("Output is None, returning empty OmniRequestOutput")
# Return empty output for the first request
if len(requests) > 0:
request = requests[0]
request_id = request.request_id or ""
prompt = request.prompt
if isinstance(prompt, list):
prompt = prompt[0] if prompt else None
return OmniRequestOutput.from_diffusion(
request_id=request_id,
return [
OmniRequestOutput.from_diffusion(
request_id=request.request_ids[i] if i < len(request.request_ids) else "",
images=[],
prompt=prompt,
metrics={},
latents=None,
)
return None
for i, prompt in enumerate(request.prompts)
]
postprocess_start_time = time.time()
outputs = self.post_process_func(output.output) if self.post_process_func is not None else output.output
@@ -102,13 +103,10 @@ class DiffusionEngine:
outputs = [outputs] if outputs is not None else []
# Handle single request or multiple requests
if len(requests) == 1:
if len(request.prompts) == 1:
# Single request: return single OmniRequestOutput
request = requests[0]
request_id = request.request_id or ""
prompt = request.prompt
if isinstance(prompt, list):
prompt = prompt[0] if prompt else None
prompt = request.prompts[0]
request_id = request.request_ids[0] if request.request_ids else ""
metrics = {}
if output.trajectory_timesteps is not None:
@@ -116,37 +114,38 @@ class DiffusionEngine:
if supports_audio_output(self.od_config.model_class_name):
audio_payload = outputs[0] if len(outputs) == 1 else outputs
return OmniRequestOutput.from_diffusion(
request_id=request_id,
images=[],
prompt=prompt,
metrics=metrics,
latents=output.trajectory_latents,
multimodal_output={"audio": audio_payload},
final_output_type="audio",
)
return [
OmniRequestOutput.from_diffusion(
request_id=request_id,
images=[],
prompt=prompt,
metrics=metrics,
latents=output.trajectory_latents,
multimodal_output={"audio": audio_payload},
final_output_type="audio",
),
]
else:
return OmniRequestOutput.from_diffusion(
request_id=request_id,
images=outputs,
prompt=prompt,
metrics=metrics,
latents=output.trajectory_latents,
)
return [
OmniRequestOutput.from_diffusion(
request_id=request_id,
images=outputs,
prompt=prompt,
metrics=metrics,
latents=output.trajectory_latents,
),
]
else:
# Multiple requests: return list of OmniRequestOutput
# Split images based on num_outputs_per_prompt for each request
results = []
output_idx = 0
for request in requests:
request_id = request.request_id or ""
prompt = request.prompt
if isinstance(prompt, list):
prompt = prompt[0] if prompt else None
for i, prompt in enumerate(request.prompts):
request_id = request.request_ids[i] if i < len(request.request_ids) else ""
# Get images for this request
num_outputs = request.num_outputs_per_prompt
num_outputs = request.sampling_params.num_outputs_per_prompt
request_outputs = outputs[output_idx : output_idx + num_outputs] if output_idx < len(outputs) else []
output_idx += num_outputs
@@ -192,8 +191,8 @@ class DiffusionEngine:
"""
return DiffusionEngine(config)
def add_req_and_wait_for_response(self, requests: list[OmniDiffusionRequest]):
return self.executor.add_req(requests)
def add_req_and_wait_for_response(self, request: OmniDiffusionRequest):
return self.executor.add_req(request)
def start_profile(self, trace_filename: str | None = None) -> None:
"""
@@ -316,28 +315,28 @@ class DiffusionEngine:
def _dummy_run(self):
"""A dummy run to warm up the model."""
prompt = "dummy run"
# note that num_inference_steps=1 will cause timestep and temb None in the pipeline
num_inference_steps = 1
height = 1024
width = 1024
if supports_image_input(self.od_config.model_class_name):
# Provide a dummy image input if the model supports it
dummy_image = PIL.Image.new("RGB", (width, height), color=(0, 0, 0))
color_format = image_color_format(self.od_config.model_class_name)
dummy_image = PIL.Image.new(color_format, (width, height))
else:
dummy_image = None
prompt: OmniTextPrompt = {"prompt": "dummy run", "multi_modal_data": {"image": dummy_image}}
req = OmniDiffusionRequest(
prompt=prompt,
height=height,
width=width,
pil_image=dummy_image,
num_inference_steps=num_inference_steps,
num_outputs_per_prompt=1,
prompts=[prompt],
sampling_params=OmniDiffusionSamplingParams(
height=height,
width=width,
num_inference_steps=num_inference_steps,
num_outputs_per_prompt=1,
),
)
logger.info("dummy run to warm up the model")
requests = self.pre_process_func([req]) if self.pre_process_func is not None else [req]
self.add_req_and_wait_for_response(requests)
request = self.pre_process_func(req) if self.pre_process_func is not None else req
self.add_req_and_wait_for_response(request)
def collective_rpc(
self,
vllm_omni/diffusion/distributed/comm.py
+6 -4
View File
@@ -8,6 +8,8 @@ import torch
import torch.distributed as dist
from torch import Tensor
from vllm_omni.platforms import current_omni_platform
__all__ = ["all_to_all_4D", "all_to_all_5D", "SeqAllToAll4D", "SeqAllToAll5D", "RingComm"]
@@ -48,7 +50,7 @@ def all_to_all_4D(
if seq_world_size > 1:
dist.all_to_all_single(output, input_t, group=group)
if use_sync:
torch.cuda.synchronize()
current_omni_platform.synchronize()
else:
output = input_t
# if scattering the seq-dim, transpose the heads back to the original dimension
@@ -83,7 +85,7 @@ def all_to_all_4D(
if seq_world_size > 1:
dist.all_to_all_single(output, input_t, group=group)
if use_sync:
torch.cuda.synchronize()
current_omni_platform.synchronize()
else:
output = input_t
@@ -155,7 +157,7 @@ def all_to_all_5D(
if seq_world_size > 1:
dist.all_to_all_single(output, input_t, group=group)
if use_sync:
torch.cuda.synchronize()
current_omni_platform.synchronize()
else:
output = input_t
@@ -190,7 +192,7 @@ def all_to_all_5D(
if seq_world_size > 1:
dist.all_to_all_single(output, input_t, group=group)
if use_sync:
torch.cuda.synchronize()
current_omni_platform.synchronize()
else:
output = input_t
vllm_omni/diffusion/distributed/group_coordinator.py
+5 -9
View File
@@ -9,18 +9,14 @@ from typing import Any
import torch
import torch.distributed
from torch.cuda import synchronize
from torch.distributed import Backend, ProcessGroup
from vllm.logger import init_logger
from vllm_omni.diffusion import envs
from vllm_omni.platforms import current_omni_platform
logger = init_logger(__name__)
if envs._is_npu():
logger.info("torch.npu synchronize")
from torch.npu import synchronize
TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"])
@@ -124,7 +120,7 @@ class GroupCoordinator:
assert self.cpu_group is not None
assert self.device_group is not None
self.device = envs.get_device(local_rank)
self.device = current_omni_platform.get_torch_device(local_rank)
@property
def first_rank(self):
@@ -618,7 +614,7 @@ class PipelineGroupCoordinator(GroupCoordinator):
assert self.cpu_group is not None
assert self.device_group is not None
self.device = envs.get_device(local_rank)
self.device = current_omni_platform.get_torch_device(local_rank)
self.recv_buffer_set: bool = False
self.recv_tasks_queue: list[tuple[str, int]] = []
@@ -768,7 +764,7 @@ class PipelineGroupCoordinator(GroupCoordinator):
# To protect against race condition when using batch_isend_irecv().
# should take this out once the bug with batch_isend_irecv is resolved.
synchronize()
current_omni_platform.synchronize()
ops = []
recv_prev_shape_tensor = None
@@ -799,7 +795,7 @@ class PipelineGroupCoordinator(GroupCoordinator):
for req in reqs:
req.wait()
synchronize()
current_omni_platform.synchronize()
recv_prev_shape = [0, 0, 0]
if recv_prev_shape_tensor is not None:
vllm_omni/diffusion/distributed/parallel_state.py
+6 -12
View File
@@ -35,6 +35,7 @@ from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size
from vllm.logger import init_logger
from vllm_omni.diffusion import envs
from vllm_omni.platforms import current_omni_platform
from .group_coordinator import (
GroupCoordinator,
@@ -42,14 +43,6 @@ from .group_coordinator import (
SequenceParallelGroupCoordinator,
)
if envs._is_npu():
from torch.npu import device_count, set_device
elif envs._is_musa():
from torch_musa.core.device import device_count, set_device
else:
from torch.cuda import device_count, set_device
env_info = envs.PACKAGES_CHECKER.get_packages_info()
HAS_FLASH_ATTN = env_info["has_flash_attn"]
@@ -396,7 +389,7 @@ def init_distributed_environment(
backend: str | None = None,
):
if backend is None:
backend = envs.get_torch_distributed_backend()
backend = current_omni_platform.dist_backend
logger.debug(
"world_size=%d rank=%d local_rank=%d distributed_init_method=%s backend=%s",
world_size,
@@ -416,7 +409,8 @@ def init_distributed_environment(
world_size=world_size,
rank=rank,
)
set_device(torch.distributed.get_rank() % device_count())
device_id = torch.distributed.get_rank() % current_omni_platform.get_device_count()
current_omni_platform.set_device(current_omni_platform.get_torch_device(device_id))
# set the local rank
# local_rank is not available in torch ProcessGroup,
# see https://github.com/pytorch/pytorch/issues/122816
@@ -572,7 +566,7 @@ def initialize_model_parallel(
backend: str | None = None,
) -> None:
if backend is None:
backend = envs.get_torch_distributed_backend()
backend = current_omni_platform.dist_backend
"""
Initialize model parallel groups.
@@ -632,7 +626,7 @@ def initialize_model_parallel(
# FIXME: Since the async p2p communication operation of NPU is not same as cuda in torch,
# the pipefusion is not ready for npu yet
if envs._is_npu():
if current_omni_platform.is_npu():
assert pipeline_parallel_size == 1, "Current pipefusion is not ready for NPU"
dit_parallel_size = (
vllm_omni/diffusion/distributed/utils.py
+3 -4
View File
@@ -5,11 +5,10 @@ import os
import torch
from vllm_omni.utils.platform_utils import detect_device_type
from vllm_omni.platforms import current_omni_platform
def get_local_device() -> torch.device:
"""Return the torch device for the current rank based on detected device type."""
device_type = detect_device_type()
local_rank = os.environ.get("LOCAL_RANK", 0)
return torch.device(f"{device_type}:{local_rank}")
local_rank = int(os.environ.get("LOCAL_RANK", 0))
return current_omni_platform.get_torch_device(local_rank)
vllm_omni/diffusion/envs.py
+41 -142
View File
@@ -5,20 +5,15 @@ import os
from collections.abc import Callable
from typing import TYPE_CHECKING, Any
import torch
from packaging import version
from vllm.logger import init_logger
logger = init_logger(__name__)
from vllm_omni.platforms import current_omni_platform
if TYPE_CHECKING:
MASTER_ADDR: str = ""
MASTER_PORT: int | None = None
CUDA_HOME: str | None = None
LOCAL_RANK: int = 0
CUDA_VISIBLE_DEVICES: str | None = None
CUDA_VERSION: version.Version
TORCH_VERSION: version.Version
environment_variables: dict[str, Callable[[], Any]] = {
# ================== Runtime Env Vars ==================
@@ -32,127 +27,14 @@ environment_variables: dict[str, Callable[[], Any]] = {
# local rank of the process in the distributed setting, used to determine
# the GPU device id
"LOCAL_RANK": lambda: int(os.environ.get("LOCAL_RANK", "0")),
# used to control the visible devices in the distributed setting
"CUDA_VISIBLE_DEVICES": lambda: os.environ.get("CUDA_VISIBLE_DEVICES", None),
}
def _is_hip():
has_rocm = torch.version.hip is not None
return has_rocm
def _is_cuda():
has_cuda = torch.version.cuda is not None
return has_cuda
def _is_musa():
try:
if hasattr(torch, "musa") and torch.musa.is_available():
return True
except ModuleNotFoundError:
return False
def _is_mps():
return torch.backends.mps.is_available()
def _is_npu():
try:
if hasattr(torch, "npu") and torch.npu.is_available():
return True
except ModuleNotFoundError:
return False
def get_device(local_rank: int) -> torch.device:
if _is_cuda() or _is_hip():
return torch.device("cuda", local_rank)
elif _is_musa():
return torch.device("musa", local_rank)
elif _is_mps():
return torch.device("mps")
elif _is_npu():
return torch.device("npu", local_rank)
else:
return torch.device("cpu")
def get_device_name() -> str:
if _is_cuda() or _is_hip():
return "cuda"
elif _is_musa():
return "musa"
elif _is_mps():
return "mps"
elif _is_npu():
return "npu"
else:
return "cpu"
def get_device_version():
if _is_hip():
hip_version = torch.version.hip
hip_version = hip_version.split("-")[0]
return hip_version
elif _is_cuda():
return torch.version.cuda
elif _is_musa():
return torch.version.musa
elif _is_mps():
return None
elif _is_npu():
return None
else:
raise NotImplementedError("No Accelerators(AMD/NV/MTT GPU, AMD MI instinct accelerators) available")
def get_torch_distributed_backend() -> str:
if _is_cuda() or _is_hip():
return "nccl"
elif _is_musa():
return "mccl"
elif _is_mps():
return "gloo"
elif _is_npu():
return "hccl"
else:
raise NotImplementedError("No Accelerators(AMD/NV/MTT GPU, AMD MI instinct accelerators) available")
variables: dict[str, Callable[[], Any]] = {
# ================== Other Vars ==================
# used in version checking
"CUDA_VERSION": lambda: version.parse(get_device_version() or "0.0"),
"TORCH_VERSION": lambda: version.parse(version.parse(torch.__version__).base_version),
}
def _setup_musa(environment_variables, variables):
musa = getattr(torch, "musa", None)
if musa is None:
return
try:
if musa.is_available():
environment_variables["MUSA_HOME"] = lambda: os.environ.get("MUSA_HOME", None)
environment_variables["MUSA_VISIBLE_DEVICES"] = lambda: os.environ.get("MUSA_VISIBLE_DEVICES", None)
musa_ver = getattr(getattr(torch, "version", None), "musa", None)
if musa_ver:
variables["MUSA_VERSION"] = lambda: version.parse(musa_ver)
except Exception:
pass
try:
_setup_musa(environment_variables, variables)
except (AttributeError, ModuleNotFoundError):
pass
logger = init_logger(__name__)
class PackagesEnvChecker:
"""Singleton class for checking package availability."""
_instance = None
def __new__(cls):
@@ -163,38 +45,57 @@ class PackagesEnvChecker:
def initialize(self):
packages_info = {}
packages_info["has_flash_attn"] = self.check_flash_attn(packages_info)
packages_info["has_flash_attn"] = self._check_flash_attn(packages_info)
self.packages_info = packages_info
def check_flash_attn(self, packages_info):
if not torch.cuda.is_available():
def _check_flash_attn(self, packages_info) -> bool:
"""Check if flash attention is available and compatible."""
platform = current_omni_platform
# Flash attention requires CUDA-like platforms (CUDA or ROCm)
if not platform.is_cuda_alike():
return False
# Check if torch_npu is available
if _is_npu():
logger.info("`flash_attn` is not ready on torch_npu for now")
# Check if devices are available
if platform.get_device_count() == 0:
return False
if _is_musa():
logger.info("Flash Attention library is not supported on MUSA for the moment.")
return False
try:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
gpu_name = torch.cuda.get_device_name(device)
gpu_name = platform.get_device_name()
# Turing/Tesla/T4 GPUs don't support flash attention well
if "Turing" in gpu_name or "Tesla" in gpu_name or "T4" in gpu_name:
return False
else:
from flash_attn import __version__
if __version__ < "2.6.0":
raise ImportError("install flash_attn >= 2.6.0")
# Check for any FA backend: FA3 (fa3_fwd_interface, flash_attn_interface) or FA2 (flash_attn)
# Try FA3 from fa3-fwd PyPI package
try:
import fa3_fwd_interface # noqa: F401
return True
except ImportError:
except (ImportError, ModuleNotFoundError):
pass
# Try FA3 from flash-attention source build
try:
import flash_attn_interface # noqa: F401
return True
except (ImportError, ModuleNotFoundError):
pass
# Try FA2 from flash-attn package
from flash_attn import __version__
if __version__ < "2.6.0":
raise ImportError("install flash_attn >= 2.6.0")
return True
except (ImportError, ModuleNotFoundError):
if not packages_info.get("has_aiter", False):
logger.warning('Flash Attention library "flash_attn" not found, using pytorch attention implementation')
logger.warning("No Flash Attention backend found, using pytorch SDPA implementation")
return False
def get_packages_info(self):
def get_packages_info(self) -> dict:
"""Get the packages info dictionary."""
return self.packages_info
@@ -205,8 +106,6 @@ def __getattr__(name):
# lazy evaluation of environment variables
if name in environment_variables:
return environment_variables[name]()
if name in variables:
return variables[name]()
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
vllm_omni/diffusion/executor/abstract.py
+2 -2
View File
@@ -3,7 +3,7 @@ from typing import Any
from vllm.utils.import_utils import resolve_obj_by_qualname
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -59,7 +59,7 @@ class DiffusionExecutor(ABC):
pass
@abstractmethod
def add_req(self, requests: list[OmniDiffusionRequest]):
def add_req(self, requests: OmniDiffusionRequest) -> DiffusionOutput:
"""Add requests to the execution queue."""
pass
vllm_omni/diffusion/executor/multiproc_executor.py
+5 -8
View File
@@ -6,11 +6,11 @@ from typing import Any
from vllm.logger import init_logger
from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE
from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, DiffusionOutput
from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.scheduler import Scheduler
from vllm_omni.utils.platform_utils import get_diffusion_worker_class
from vllm_omni.diffusion.worker import WorkerProc
logger = init_logger(__name__)
@@ -77,9 +77,6 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
mp.set_start_method("spawn", force=True)
processes = []
# Get the appropriate worker class for current device
worker_proc = get_diffusion_worker_class()
# Launch all worker processes
scheduler_pipe_readers = []
scheduler_pipe_writers = []
@@ -88,7 +85,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
reader, writer = mp.Pipe(duplex=False)
scheduler_pipe_writers.append(writer)
process = mp.Process(
target=worker_proc.worker_main,
target=WorkerProc.worker_main,
args=(
i, # rank
od_config,
@@ -130,8 +127,8 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
return processes, result_handle
def add_req(self, requests: list[OmniDiffusionRequest]):
return self.scheduler.add_req(requests)
def add_req(self, request: OmniDiffusionRequest) -> DiffusionOutput:
return self.scheduler.add_req(request)
def collective_rpc(
self,
vllm_omni/diffusion/hooks/sequence_parallel.py
+2 -5
View File
@@ -335,10 +335,7 @@ class SequenceParallelSplitHook(ModelHook):
2. Creates an attention mask indicating valid vs padding positions
3. Stores the mask and padding info in ForwardContext
"""
from vllm_omni.diffusion.attention.selector import (
_BACKENDS_SUPPORT_ATTENTION_MASK,
get_attn_backend,
)
from vllm_omni.diffusion.attention.selector import get_attn_backend
from vllm_omni.diffusion.distributed.parallel_state import (
get_ring_parallel_world_size,
get_sequence_parallel_rank,
@@ -359,7 +356,7 @@ class SequenceParallelSplitHook(ModelHook):
# Check backend compatibility
attn_backend = get_attn_backend(-1)
if attn_backend.get_name() not in _BACKENDS_SUPPORT_ATTENTION_MASK:
if not attn_backend.supports_attention_mask:
raise ValueError(
f"Sequence length ({seq_len}) is not divisible by SP world size ({world_size}). "
f"Cannot use {attn_backend.get_name()} which does not support attention_mask. "
vllm_omni/diffusion/layers/adalayernorm.py
+8
View File
@@ -91,6 +91,14 @@ class AdaLayerNorm(CustomOp):
return output, gate_result
def forward_xpu(
self,
x: torch.Tensor,
mod_params: torch.Tensor,
index: torch.Tensor = None,
) -> torch.Tensor:
return self.forward_native(x, mod_params, index)
def forward_native(
self,
x: torch.Tensor,
vllm_omni/diffusion/layers/custom_op.py
+9 -5
View File
@@ -3,7 +3,7 @@ from typing import Any
import torch.nn as nn
from vllm_omni.utils.platform_utils import detect_device_type, is_npu, is_rocm
from vllm_omni.platforms import current_omni_platform
class CustomOp(nn.Module):
@@ -14,16 +14,17 @@ class CustomOp(nn.Module):
def __init__(self) -> None:
super().__init__()
self.is_cuda = detect_device_type() == "cuda"
self._forward_method = self.dispatch_forward()
def dispatch_forward(self) -> Callable:
if is_rocm():
if current_omni_platform.is_rocm():
return self.forward_hip
elif self.is_cuda:
elif current_omni_platform.is_cuda():
return self.forward_cuda
elif is_npu():
elif current_omni_platform.is_npu():
return self.forward_npu
elif current_omni_platform.is_xpu():
return self.forward_xpu
else:
return self.forward_native
@@ -44,6 +45,9 @@ class CustomOp(nn.Module):
def forward_npu(self, *args, **kwargs):
raise NotImplementedError
def forward_xpu(self, *args, **kwargs):
raise NotImplementedError
def forward_hip(self, *args, **kwargs):
# By default, we assume that HIP ops are compatible with CUDA ops.
return self.forward_cuda(*args, **kwargs)
vllm_omni/diffusion/layers/rope.py
+8
View File
@@ -137,6 +137,14 @@ class RotaryEmbedding(CustomOp):
else:
return self.forward_native(x, cos, sin)
def forward_xpu(
self,
x: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
) -> torch.Tensor:
return self.forward_native(x, cos, sin)
def forward_native(
self,
x: torch.Tensor,
vllm_omni/diffusion/models/bagel/pipeline_bagel.py
+23 -13
View File
@@ -272,16 +272,22 @@ class BagelPipeline(nn.Module):
@torch.inference_mode()
def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
prompt = req.prompt or ""
if isinstance(prompt, list):
# vllm-omni request supports list; Bagel pipeline currently supports first prompt.
prompt = prompt[0] if prompt else ""
if len(req.prompts) > 1:
logger.warning(
"""This model only supports a single prompt, not a batched request.""",
"""Taking only the first image for now.""",
)
# TODO: In online mode, sometimes it receives [{"prompts": None}, {...}], so cannot use .get("...", "")
# TODO: May be some data formatting operations on the API side. Hack for now.
first_prompt = req.prompts[0]
prompt = first_prompt if isinstance(req.prompts[0], str) else (req.prompts[0].get("prompt") or "")
max_hw = int(self.bagel.max_latent_size * self.bagel.latent_downsample)
if req.height is None and req.width is None:
if req.sampling_params.height is None and req.sampling_params.width is None:
height = width = max_hw
else:
height = int(req.height) if req.height is not None else max_hw
width = int(req.width) if req.width is not None else max_hw
height = int(req.sampling_params.height) if req.sampling_params.height is not None else max_hw
width = int(req.sampling_params.width) if req.sampling_params.width is not None else max_hw
if height > max_hw or width > max_hw:
raise ValueError(
f"Requested resolution {height}x{width} exceeds Bagel checkpoint limit "
@@ -292,7 +298,7 @@ class BagelPipeline(nn.Module):
# Map request params to Bagel gen params (defaults follow Bagel inferencer)
gen_params = BagelGenParams(
num_timesteps=int(req.num_inference_steps or 50),
num_timesteps=int(req.sampling_params.num_inference_steps or 50),
timestep_shift=3.0,
)
@@ -304,7 +310,7 @@ class BagelPipeline(nn.Module):
# Add text prompt (prefill) on gen context.
# [Omni] Check for injected KV Cache from remote transfer
injected_kv = getattr(req, "past_key_values", None)
injected_kv = req.sampling_params.past_key_values
if injected_kv is not None:
logger.info("Using injected KV Cache (direct)")
gen_context["past_key_values"] = injected_kv
@@ -316,9 +322,13 @@ class BagelPipeline(nn.Module):
gen_context["ropes"] = [seq_len]
else:
image_input = getattr(req, "pil_image", None)
image_input = (
None if isinstance(first_prompt, str) else (first_prompt.get("multi_modal_data") or {}).get("image")
)
if image_input and not isinstance(image_input, list):
image_input = [image_input]
if image_input:
image_input = [Image.open(image) if isinstance(image, str) else image for image in image_input]
if image_input:
# If we have an image, we prefill with it
@@ -414,10 +424,10 @@ class BagelPipeline(nn.Module):
gen_context["kv_lens"] = newlens
gen_context["ropes"] = new_rope
if req.seed is not None:
torch.manual_seed(req.seed)
if req.sampling_params.seed is not None:
torch.manual_seed(req.sampling_params.seed)
if self.device.type == "cuda":
torch.cuda.manual_seed(req.seed)
torch.cuda.manual_seed(req.sampling_params.seed)
# Prepare latent query and run flow
generation_input = self.bagel.prepare_vae_latent(
vllm_omni/diffusion/models/flux/__init__.py
+17
View File
@@ -0,0 +1,17 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""FLUX.1-dev diffusion model components."""
from vllm_omni.diffusion.models.flux.flux_transformer import (
FluxTransformer2DModel,
)
from vllm_omni.diffusion.models.flux.pipeline_flux import (
FluxPipeline,
get_flux_post_process_func,
)
__all__ = [
"FluxPipeline",
"FluxTransformer2DModel",
"get_flux_post_process_func",
]
Loaded 100 of 196 files, more files were not shown because too many files have changed in this diff. Show more