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

[Diffusion]: Diffusion Ulysses-Sequence-Parallelism support (#189)

Browse Source 
Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
This commit is contained in:
Didan Deng authored and GitHub committed 2025-12-17 15:41:31 +08:00
1 parent f785a5d401
commit 85f0c1177f
21 files changed
+3513 -61
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/user_guide/cache_dit_acceleration.md → docs/user_guide/acceleration/cache_dit_acceleration.md
View File
File renamed without changes.
docs/user_guide/acceleration/parallelism_acceleration.md
+128
View File
@@ -0,0 +1,128 @@
# Parallelism Acceleration Guide
This guide includes how to use parallelism methods in vLLM-Omni to speed up diffusion model inference as well as reduce the memory requirement on each device.
## Overview
The following parallelism methods are currently supported in vLLM-Omni:
1. DeepSpeed Ulysses Sequence Parallel (Ulysses-SP) ([paper](https://arxiv.org/pdf/2309.14509)): Ulysses-SP splits the input along the sequence dimension and uses all-to-all communication to allow each device to compute only a subset of attention heads.
The following table shows which models are currently supported by parallelism method:
| Model | Model Identifier | Ulysses-SP |
|-------|-----------------|-----------|
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ |
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ❌ |
| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ |
### Sequence Parallelism
#### Ulysses-SP
##### Quick Start
An example of using Ulysses-SP is shown below:
```python
from vllm_omni import Omni
from vllm_omni.diffusion.data import DiffusionParallelConfig
ulysses_degree = 2
omni = Omni(
model="Qwen/Qwen-Image",
parallel_config=DiffusionParallelConfig(ulysses_degree=2)
)
outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
```
See `examples/offline_inference/text_to_image/text_to_image.py` for a complete working example.
##### Benchmarks
!!! note "Benchmark Disclaimer"
These benchmarks are provided for **general reference only**. The configurations shown use default or common parameter settings and have not been exhaustively optimized for maximum performance. Actual performance may vary based on:
- Specific model and use case
- Hardware configuration
- Careful parameter tuning
- Different inference settings (e.g., number of steps, image resolution)
To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** model generating images (**2048x2048** as long sequence input) with 50 inference steps. The hardware devices are NVIDIA H800 GPUs. `sdpa` is the attention backends.
| Configuration | Ulysses degree |Generation Time | Speedup |
|---------------|----------------|---------|---------|
| **Baseline (diffusers)** | - | 112.5s | 1.0x |
| Ulysses-SP | 2 | 65.2s | 1.73x |
| Ulysses-SP | 4 | 39.6s | 2.84x |
| Ulysses-SP | 8 | 30.8s | 3.65x |
##### How to parallelize a new model
If a diffusion model has been deployed in vLLM-Omni and supports single-card inference, you can refer to the following instruction on how to parallelize this model with Ulysses-SP.
First, please edit the `TransformerModel`'s `forward` function in the `xxx_model_transformer.py` to make the inputs (image hidden states, positional embeddings, etc.) as chunks separated at the sequence dimension. Taking `qwen_image_transformer.py` as an example:
```diff
class QwenImageTransformer2DModel(nn.Module):
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor = None,
...
):
+ if self.parallel_config.sequence_parallel_size > 1:
+ hidden_states = torch.chunk(hidden_states, get_sequence_parallel_world_size(), dim=-2)[
+ get_sequence_parallel_rank()
+ ]
hidden_states = self.img_in(hidden_states)
...
image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device)
+ def get_rotary_emb_chunk(freqs):
+ freqs = torch.chunk(freqs, get_sequence_parallel_world_size(), dim=0)[get_sequence_parallel_rank()]
+ return freqs
+ if self.parallel_config.sequence_parallel_size > 1:
+ img_freqs, txt_freqs = image_rotary_emb
+ img_freqs = get_rotary_emb_chunk(img_freqs)
+ image_rotary_emb = (img_freqs, txt_freqs)
```
Next, at the end of the `forward` function, please call `get_sp_group().all_gather` to gather the chunked outputs across devices, and concatenate them at the sequence dimension.
```diff
class QwenImageTransformer2DModel(nn.Module):
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor = None,
...
):
# Use only the image part (hidden_states) from the dual-stream blocks
hidden_states = self.norm_out(hidden_states, temb)
output = self.proj_out(hidden_states)
+ if self.parallel_config.sequence_parallel_size > 1:
+ output = get_sp_group().all_gather(output, dim=-2)
return Transformer2DModelOutput(sample=output)
```
Finally, you can set the parallel configuration and pass it to `Omni` and start parallel inference with:
```diff
from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
ulysses_degree = 2
omni = Omni(
model="Qwen/Qwen-Image",
+ parallel_config=DiffusionParallelConfig(ulysses_degree=2)
)
outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
```
docs/user_guide/teacache.md → docs/user_guide/acceleration/teacache.md
View File
File renamed without changes.
docs/user_guide/diffusion_acceleration.md
+62 -15
View File
@@ -1,21 +1,27 @@
# Diffusion Acceleration Overview
vLLM-Omni supports various cache acceleration methods to speed up diffusion model inference with minimal quality degradation. These methods intelligently cache intermediate computations to avoid redundant work across diffusion timesteps.
vLLM-Omni supports various cache acceleration methods to speed up diffusion model inference with minimal quality degradation. These methods include **cache methods** that intelligently cache intermediate computations to avoid redundant work across diffusion timesteps, and **parallelism methods** that distribute the computation across multiple devices.
## Supported Acceleration Methods
vLLM-Omni currently supports two main cache acceleration backends:
1. **[TeaCache](teacache.md)** - Hook-based adaptive caching that caches transformer computations when consecutive timesteps are similar
2. **[Cache-DiT](cache_dit_acceleration.md)** - Library-based acceleration using multiple techniques:
- **DBCache** (Dual Block Cache): Caches intermediate transformer block outputs based on residual differences
- **TaylorSeer**: Uses Taylor expansion-based forecasting for faster inference
- **SCM** (Step Computation Masking): Selectively computes steps based on adaptive masking
1. **[TeaCache](acceleration/teacache.md)** - Hook-based adaptive caching that caches transformer computations when consecutive timesteps are similar
2. **[Cache-DiT](acceleration/cache_dit_acceleration.md)** - Library-based acceleration using multiple techniques:
- **DBCache** (Dual Block Cache): Caches intermediate transformer block outputs based on residual differences
- **TaylorSeer**: Uses Taylor expansion-based forecasting for faster inference
- **SCM** (Step Computation Masking): Selectively computes steps based on adaptive masking
Both methods can provide significant speedups (typically **1.5x-2.0x**) while maintaining high output quality.
vLLM-Omni also supports the sequence parallelism (SP) for the diffusion model, that includes:
1. [Ulysses-SP](acceleration/parallelism_acceleration.md#ulysses-sp) - splits the input along the sequence dimension and uses all-to-all communication to allow each device to compute only a subset of attention heads.
## Quick Comparison
### Cache Methods
| Method | Configuration | Description | Best For |
|--------|--------------|-------------|----------|
| **TeaCache** | `cache_backend="tea_cache"` | Simple, adaptive caching with minimal configuration | Quick setup, balanced speed/quality |
@@ -23,18 +29,18 @@ Both methods can provide significant speedups (typically **1.5x-2.0x**) while ma
## Supported Models
The following table shows which models are currently supported by each cache backend:
The following table shows which models are currently supported by each acceleration method:
| Model | Model Identifier | TeaCache | Cache-DiT |
|-------|-----------------|----------|-----------|
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ |
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ❌ | ✅ |
| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ❌ | ✅ |
| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP |
|-------|-----------------|----------|-----------|-----------|
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ |
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ❌ | ✅ |❌ |
| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ❌ | ✅ |✅ |
## Performance Benchmarks
The following benchmarks were measured on **Qwen/Qwen-Image** and **Qwen/Qwen-Image-Edit** models with 50 inference steps:
The following benchmarks were measured on **Qwen/Qwen-Image** and **Qwen/Qwen-Image-Edit** models generating 1024x1024 images with 50 inference steps:
!!! note "Benchmark Disclaimer"
These benchmarks are provided for **general reference only**. The configurations shown use default or common parameter settings and have not been exhaustively optimized for maximum performance. Actual performance may vary based on:
@@ -55,6 +61,14 @@ The following benchmarks were measured on **Qwen/Qwen-Image** and **Qwen/Qwen-Im
| **Qwen/Qwen-Image-Edit** | None | No acceleration | 51.5s | 1.0x | Baseline (diffusers) |
| **Qwen/Qwen-Image-Edit** | Cache-DiT | Default (Fn=1, Bn=0, W=4, TaylorSeer disabled, SCM disabled) | 21.6s | **2.38x** | - |
To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** model generating images (**2048x2048** as long sequence input) with 50 inference steps. The hardware devices are NVIDIA H800 GPUs. `sdpa` is the attention backends.
| Configuration | Ulysses degree |Generation Time | Speedup |
|---------------|----------------|---------|---------|
| **Baseline (diffusers)** | - | 112.5s | 1.0x |
| Ulysses-SP | 2 | 65.2s | 1.73x |
| Ulysses-SP | 4 | 39.6s | 2.84x |
| Ulysses-SP | 8 | 30.8s | 3.65x |
## Quick Start
@@ -92,9 +106,42 @@ omni = Omni(
outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
```
### Using Ulysses-SP
Run text-to-image:
```python
from vllm_omni import Omni
from vllm_omni.diffusion.data import DiffusionParallelConfig
ulysses_degree = 2
omni = Omni(
model="Qwen/Qwen-Image",
parallel_config=DiffusionParallelConfig(ulysses_degree=2)
)
outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
```
Run image-to-image:
```python
from vllm_omni import Omni
from vllm_omni.diffusion.data import DiffusionParallelConfig
ulysses_degree = 2
omni = Omni(
model="Qwen/Qwen-Image-Edit",
parallel_config=DiffusionParallelConfig(ulysses_degree=2)
)
outputs = omni.generate(prompt="turn this cat to a dog",
pil_image=input_image, num_inference_steps=50)
```
## Documentation
For detailed information on each acceleration method:
- **[TeaCache Guide](teacache.md)** - Complete TeaCache documentation, configuration options, and best practices
- **[Cache-DiT Acceleration Guide](cache_dit_acceleration.md)** - Comprehensive Cache-DiT guide covering DBCache, TaylorSeer, SCM, and configuration parameters
- **[TeaCache Guide](acceleration/teacache.md)** - Complete TeaCache documentation, configuration options, and best practices
- **[Cache-DiT Acceleration Guide](acceleration/cache_dit_acceleration.md)** - Comprehensive Cache-DiT guide covering DBCache, TaylorSeer, SCM, and configuration parameters
- **[Sequence Parallelism](acceleration/parallelism_acceleration.md#sequence-parallelism)** - Guidance on how to set sequence parallelism with configuration.