[Doc] refactor diffusion doc (#753)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
This commit is contained in:
Jiangyun Zhu authored and GitHub committed
1 parent
f7611198ae
commit
d7cd00e26c
8 files changed
+20
-20
No files matched your search
@@ -0,0 +1,227 @@
|
||||
# Cache-DiT Acceleration Guide
|
||||
|
||||
This guide explains how to use cache-dit acceleration in vLLM-Omni to speed up diffusion model inference.
|
||||
|
||||
## Overview
|
||||
|
||||
Cache-dit is a library that accelerates diffusion transformer models through intelligent caching mechanisms. It supports multiple acceleration techniques that can be combined for optimal performance:
|
||||
|
||||
- **DBCache**: Dual Block Cache for reducing redundant computations
|
||||
- **TaylorSeer**: Taylor expansion-based forecasting for faster inference
|
||||
- **SCM**: Step Computation Masking for selective step computation
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Basic Usage
|
||||
|
||||
Enable cache-dit acceleration by simply setting `cache_backend="cache_dit"`. Cache-dit will use its recommended default parameters:
|
||||
|
||||
```python
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
# Simplest way: just enable cache-dit with default parameters
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
cache_backend="cache_dit",
|
||||
)
|
||||
|
||||
images = omni.generate(
|
||||
"a beautiful landscape",
|
||||
num_inference_steps=50,
|
||||
)
|
||||
```
|
||||
|
||||
**Default Parameters**: When `cache_config` is not provided, cache-dit uses optimized default values. See the [Configuration Reference](#configuration-reference) section for a complete list of all parameters and their default values.
|
||||
|
||||
### Custom Configuration
|
||||
|
||||
To customize cache-dit settings, provide a `cache_config` dictionary, for example:
|
||||
|
||||
```python
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
cache_backend="cache_dit",
|
||||
cache_config={
|
||||
"Fn_compute_blocks": 1,
|
||||
"Bn_compute_blocks": 0,
|
||||
"max_warmup_steps": 4,
|
||||
"residual_diff_threshold": 0.12,
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
## Online Serving (OpenAI-Compatible)
|
||||
|
||||
Enable Cache-DiT for online serving by passing `--cache-backend cache_dit` when starting the server:
|
||||
|
||||
```bash
|
||||
# Use Cache-DiT default (recommended) parameters
|
||||
vllm serve Qwen/Qwen-Image --omni --port 8091 --cache-backend cache_dit
|
||||
```
|
||||
|
||||
To customize Cache-DiT settings for online serving, pass a JSON string via `--cache-config`:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image --omni --port 8091 \
|
||||
--cache-backend cache_dit \
|
||||
--cache-config '{"Fn_compute_blocks": 1, "Bn_compute_blocks": 0, "max_warmup_steps": 4, "residual_diff_threshold": 0.12}'
|
||||
```
|
||||
|
||||
## Acceleration Methods
|
||||
|
||||
For comprehensive illustration, please view cache-dit [User_Guide](https://github.com/vipshop/cache-dit/blob/main/docs/User_Guide.md)
|
||||
|
||||
### 1. DBCache (Dual Block Cache)
|
||||
|
||||
DBCache intelligently caches intermediate transformer block outputs when the residual differences between consecutive steps are small, reducing redundant computations without sacrificing quality.
|
||||
|
||||
**Key Parameters**:
|
||||
|
||||
- `Fn_compute_blocks` (int, default: 1): Number of **first n** transformer blocks used to compute stable feature differences. Higher values provide more accurate caching decisions but increase computation.
|
||||
- `Bn_compute_blocks` (int, default: 0): Number of **last n** transformer blocks used for additional fusion. These blocks act as an auto-scaler for approximate hidden states.
|
||||
- `max_warmup_steps` (int, default: 4): Number of initial steps where caching is disabled to ensure the model learns sufficient features before caching begins. Optimized for few-step distilled models.
|
||||
- `residual_diff_threshold` (float, default: 0.24): Threshold for residual difference. Higher values lead to faster performance but may reduce precision. Default uses a relatively higher threshold for more aggressive caching.
|
||||
- `max_cached_steps` (int, default: -1): Maximum number of cached steps. Set to -1 for unlimited caching.
|
||||
- `max_continuous_cached_steps` (int, default: 3): Maximum number of consecutive cached steps. Limits consecutive caching to prevent precision degradation.
|
||||
|
||||
**Example Configuration**:
|
||||
|
||||
```python
|
||||
cache_config={
|
||||
"Fn_compute_blocks": 8, # Use first 8 blocks for difference computation
|
||||
"Bn_compute_blocks": 0, # No additional fusion blocks
|
||||
"max_warmup_steps": 8, # Cache after 8 warmup steps
|
||||
"residual_diff_threshold": 0.12, # Higher threshold for faster inference
|
||||
"max_cached_steps": -1, # No limit on cached steps
|
||||
}
|
||||
```
|
||||
|
||||
**Performance Tips**:
|
||||
|
||||
- Default `Fn_compute_blocks=1` works well for most cases. Increase to 8-12 for larger models or when more accuracy is needed
|
||||
- Increase `residual_diff_threshold` (e.g., 0.12-0.15) for faster inference with slight quality trade-off, or decrease from default 0.24 for higher quality
|
||||
- Default `max_warmup_steps=4` is optimized for few-step models. Increase to 6-8 for more steps if needed
|
||||
|
||||
### 2. TaylorSeer
|
||||
|
||||
TaylorSeer uses Taylor expansion to forecast future hidden states, allowing the model to skip some computation steps while maintaining quality.
|
||||
|
||||
**Key Parameters**:
|
||||
|
||||
- `enable_taylorseer` (bool, default: False): Enable TaylorSeer acceleration
|
||||
- `taylorseer_order` (int, default: 1): Order of Taylor expansion. Higher orders provide better accuracy but require more computation.
|
||||
|
||||
**Example Configuration**:
|
||||
|
||||
```python
|
||||
cache_config={
|
||||
"enable_taylorseer": True,
|
||||
"taylorseer_order": 1, # First-order Taylor expansion
|
||||
}
|
||||
```
|
||||
|
||||
**Performance Tips**:
|
||||
|
||||
- Use `taylorseer_order=1` for most cases (good balance of speed and quality)
|
||||
- Combine with DBCache for maximum acceleration
|
||||
- Higher orders (2-3) may improve quality but reduce speed gains
|
||||
|
||||
### 3. SCM (Step Computation Masking)
|
||||
|
||||
SCM allows you to specify which steps must be computed and which can use cached results, similar to LeMiCa/EasyCache style acceleration.
|
||||
|
||||
**Key Parameters**:
|
||||
|
||||
- `scm_steps_mask_policy` (str | None, default: None): Predefined mask policy. Options:
|
||||
- `None`: SCM disabled (default)
|
||||
- `"slow"`: More compute steps, higher quality (18 compute steps out of 28)
|
||||
- `"medium"`: Balanced (15 compute steps out of 28)
|
||||
- `"fast"`: More cache steps, faster inference (11 compute steps out of 28)
|
||||
- `"ultra"`: Maximum speed (8 compute steps out of 28)
|
||||
- `scm_steps_policy` (str, default: "dynamic"): Policy for cached steps:
|
||||
- `"dynamic"`: Use dynamic cache for masked steps (recommended)
|
||||
- `"static"`: Use static cache for masked steps
|
||||
|
||||
**Example Configuration**:
|
||||
|
||||
```python
|
||||
cache_config={
|
||||
"scm_steps_mask_policy": "medium", # Balanced speed/quality
|
||||
"scm_steps_policy": "dynamic", # Use dynamic cache
|
||||
}
|
||||
```
|
||||
|
||||
**Performance Tips**:
|
||||
|
||||
- SCM is disabled by default (`scm_steps_mask_policy=None`). Enable it by setting a policy value if you need additional acceleration
|
||||
- Start with `"medium"` policy and adjust based on quality requirements
|
||||
- Use `"fast"` or `"ultra"` for maximum speed when quality can be slightly compromised
|
||||
- `"dynamic"` policy generally provides better quality than `"static"`
|
||||
- SCM mask is automatically regenerated when `num_inference_steps` changes during inference
|
||||
|
||||
## Configuration Reference
|
||||
|
||||
### DiffusionCacheConfig Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `Fn_compute_blocks` | int | 1 | First n blocks for difference computation (optimized for single-transformer models) |
|
||||
| `Bn_compute_blocks` | int | 0 | Last n blocks for fusion |
|
||||
| `max_warmup_steps` | int | 4 | Steps before caching starts (optimized for few-step distilled models) |
|
||||
| `max_cached_steps` | int | -1 | Max cached steps (-1 = unlimited) |
|
||||
| `max_continuous_cached_steps` | int | 3 | Max consecutive cached steps (prevents precision degradation) |
|
||||
| `residual_diff_threshold` | float | 0.24 | Residual difference threshold (higher for more aggressive caching) |
|
||||
| `num_inference_steps` | int \| None | None | Initial inference steps for SCM mask generation (optional, auto-refreshed during inference) |
|
||||
| `enable_taylorseer` | bool | False | Enable TaylorSeer acceleration (not suitable for few-step distilled models) |
|
||||
| `taylorseer_order` | int | 1 | Taylor expansion order |
|
||||
| `scm_steps_mask_policy` | str \| None | None | SCM mask policy (None, "slow", "medium", "fast", "ultra") |
|
||||
| `scm_steps_policy` | str | "dynamic" | SCM computation policy ("dynamic" or "static") |
|
||||
|
||||
## Example: Accelerate Text-to-Image Generation with CacheDiT
|
||||
|
||||
See `examples/offline_inference/text_to_image/text_to_image.py` for a complete working example with cache-dit acceleration.
|
||||
|
||||
```bash
|
||||
# Enable cache-dit with hybrid acceleration
|
||||
cd examples/offline_inference/text_to_image
|
||||
python text_to_image.py \
|
||||
--model Qwen/Qwen-Image \
|
||||
--prompt "a cup of coffee on the table" \
|
||||
--cache_backend cache_dit \
|
||||
--num_inference_steps 50
|
||||
```
|
||||
|
||||
|
||||
The script uses cache-dit acceleration with a hybrid configuration combining DBCache, SCM, and TaylorSeer:
|
||||
|
||||
```python
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
cache_backend="cache_dit",
|
||||
cache_config={
|
||||
# Scheme: Hybrid DBCache + SCM + TaylorSeer
|
||||
# DBCache
|
||||
"Fn_compute_blocks": 8,
|
||||
"Bn_compute_blocks": 0,
|
||||
"max_warmup_steps": 4,
|
||||
"residual_diff_threshold": 0.12,
|
||||
# TaylorSeer
|
||||
"enable_taylorseer": True,
|
||||
"taylorseer_order": 1,
|
||||
# SCM
|
||||
"scm_steps_mask_policy": "fast", # Set to None to disable SCM
|
||||
"scm_steps_policy": "dynamic",
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
You can customize the configuration by modifying the `cache_config` dictionary to use only specific methods (e.g., DBCache only, DBCache + SCM, etc.) based on your quality and speed requirements.
|
||||
|
||||
To test another model, you can modify `--model` with the target model identifier like `Tongyi-MAI/Z-Image-Turbo` and update `cache_config` according the model architecture (e.g., number of transformer blocks).
|
||||
|
||||
|
||||
## Additional Resources
|
||||
|
||||
- [Cache-DiT User Guide](https://github.com/vipshop/cache-dit/blob/main/docs/User_Guide.md)
|
||||
- [Cache-DiT Benchmark](https://github.com/vipshop/cache-dit/tree/main/bench)
|
||||
- [DBCache Technical Details](https://github.com/vipshop/cache-dit/blob/main/docs/DBCache.md)
|
||||
@@ -0,0 +1,29 @@
|
||||
# CPU Offloading for Diffusion Model
|
||||
|
||||
## Overview
|
||||
CPU offload lets the diffusion worker move large model components between GPU and CPU memory on demand. It keeps the DiT transformer resident on GPU only while it is actively running, and swaps it out when encoders modules need the device. This reduces peak VRAM usage so bigger checkpoints run on smaller GPUs, or multiple requests can share the same GPU.
|
||||
|
||||
## Execution Model
|
||||
1. Text encoders run on GPU while the DiT transformer is offloaded to CPU.
|
||||
2. Before denoising, weights are prefetched back to GPU, honoring pinned-memory copies for speed.
|
||||
3. After the diffusion step, the transformer returns to CPU and the process repeats as needed.
|
||||
|
||||
Transfers use pinned host buffers, and the worker coordinates swaps via mutex-style hooks so components never compete for memory.
|
||||
|
||||
## Configuration
|
||||
You can enable CPU offload in two ways:
|
||||
|
||||
- **Python API**: set `enable_cpu_offload=True`.
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
m = Omni(model="Qwen/Qwen-Image",enable_cpu_offload=True)
|
||||
```
|
||||
|
||||
- **CLI**: pass `--enable-cpu-offload` to the diffusion service entrypoint.
|
||||
|
||||
## Known Limitations
|
||||
- Cold start latency increases for over one minute for some models(e.g., Qwen-Image)
|
||||
@@ -0,0 +1,404 @@
|
||||
# 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 (DeepSpeed Ulysses-SP) ([arxiv 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.
|
||||
|
||||
2. [Ring-Attention](#ring-attention) - splits the input along the sequence dimension and uses ring-based P2P communication to accumulate attention results, keeping the sequence dimension sharded
|
||||
|
||||
3. Classifier-Free-Guidance Parallel (CFG-Parallel): CFG-Parallel runs the positive/negative prompts of classifier-free guidance (CFG) on different devices, then merges on a single device to perform the scheduler step.
|
||||
|
||||
4. [Tensor Parallelism](#tensor-parallelism): Tensor parallelism shards model weights across devices. This can reduce per-GPU memory usage. Note that for diffusion models we currently shard the majority of layers within the DiT.
|
||||
|
||||
The following table shows which models are currently supported by parallelism method:
|
||||
|
||||
### ImageGen
|
||||
|
||||
| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel |
|
||||
|-------|------------------|------------|---------|--------------|--------------------------|
|
||||
| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ❌ |
|
||||
| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ❌ |
|
||||
| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ❌ |
|
||||
| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ❌ |
|
||||
| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ❌ |
|
||||
| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ❌ |
|
||||
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ❌ | ❌ | ❌ | ✅ (TP=2 only) |
|
||||
| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ❌ |
|
||||
|
||||
!!! note "Why Z-Image is TP=2 only"
|
||||
Z-Image Turbo is currently limited to `tensor_parallel_size` of **1 or 2** due to model shape divisibility constraints.
|
||||
For example, the model has `n_heads=30` and a final projection out dimension of `64`, so valid TP sizes must divide both 30 and 64; the only common divisors are **1 and 2**.
|
||||
|
||||
### VideoGen
|
||||
|
||||
| Model | Model Identifier | Ulysses-SP | Ring-SP | Tensor-Parallel |
|
||||
|-------|------------------|------------|---------|--------------------------|
|
||||
| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ❌ | ❌ |
|
||||
|
||||
### Tensor Parallelism
|
||||
|
||||
Tensor parallelism splits model parameters across GPUs. In vLLM-Omni, tensor parallelism is configured via `DiffusionParallelConfig.tensor_parallel_size`.
|
||||
|
||||
#### Offline Inference
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
|
||||
omni = Omni(
|
||||
model="Tongyi-MAI/Z-Image-Turbo",
|
||||
parallel_config=DiffusionParallelConfig(tensor_parallel_size=2),
|
||||
)
|
||||
|
||||
outputs = omni.generate(
|
||||
prompt="a cat reading a book",
|
||||
num_inference_steps=9,
|
||||
width=512,
|
||||
height=512,
|
||||
)
|
||||
```
|
||||
|
||||
### Sequence Parallelism
|
||||
|
||||
#### Ulysses-SP
|
||||
|
||||
##### Offline Inference
|
||||
|
||||
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.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.
|
||||
|
||||
##### Online Serving
|
||||
|
||||
You can enable Ulysses-SP in online serving for diffusion models via `--usp`:
|
||||
|
||||
```bash
|
||||
# Text-to-image (requires >= 2 GPUs)
|
||||
vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2
|
||||
```
|
||||
|
||||
##### 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 |
|
||||
|
||||
#### Ring-Attention
|
||||
|
||||
Ring-Attention ([arxiv paper](https://arxiv.org/abs/2310.01889)) splits the input along the sequence dimension and uses ring-based P2P communication to accumulate attention results. Unlike Ulysses-SP which uses all-to-all communication, Ring-Attention keeps the sequence dimension sharded throughout the computation and circulates Key/Value blocks through a ring topology.
|
||||
|
||||
##### Offline Inference
|
||||
|
||||
An example of offline inference script using Ring-Attention is shown below:
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
ring_degree = 2
|
||||
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
parallel_config=DiffusionParallelConfig(ring_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.
|
||||
|
||||
|
||||
##### Online Serving
|
||||
|
||||
You can enable Ring-Attention in online serving for diffusion models via `--ring`:
|
||||
|
||||
```bash
|
||||
# Text-to-image (requires >= 2 GPUs)
|
||||
vllm serve Qwen/Qwen-Image --omni --port 8091 --ring 2
|
||||
```
|
||||
|
||||
##### 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 (**1024x1024** as long sequence input) with 50 inference steps. The hardware devices are NVIDIA A100 GPUs. `flash_attn` is the attention backends.
|
||||
|
||||
| Configuration | Ring degree |Generation Time | Speedup |
|
||||
|---------------|----------------|---------|---------|
|
||||
| **Baseline (diffusers)** | - | 45.2s | 1.0x |
|
||||
| Ring-Attention | 2 | 29.9s | 1.51x |
|
||||
| Ring-Attention | 4 | 23.3s | 1.94x |
|
||||
|
||||
|
||||
#### Hybrid Ulysses + Ring
|
||||
|
||||
You can combine both Ulysses-SP and Ring-Attention for larger scale parallelism. The total sequence parallel size equals `ulysses_degree × ring_degree`.
|
||||
|
||||
##### Offline Inference
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
|
||||
# Hybrid: 2 Ulysses × 2 Ring = 4 GPUs total
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
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)
|
||||
```
|
||||
|
||||
##### Online Serving
|
||||
|
||||
```bash
|
||||
# Text-to-image (requires >= 4 GPUs)
|
||||
vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2 --ring 2
|
||||
```
|
||||
|
||||
##### 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 (**1024x1024** as long sequence input) with 50 inference steps. The hardware devices are NVIDIA A100 GPUs. `flash_attn` is the attention backends.
|
||||
|
||||
| Configuration | Ulysses degree | Ring degree | Generation Time | Speedup |
|
||||
|---------------|----------------|-------------|-----------------|---------|
|
||||
| **Baseline (diffusers)** | - | - | 45.2s | 1.0x |
|
||||
| Hybrid Ulysses + Ring | 2 | 2 | 24.3s | 1.87x |
|
||||
|
||||
|
||||
##### 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 instructions to parallelize it with [Ulysses-SP](https://arxiv.org/pdf/2309.14509).
|
||||
|
||||
|
||||
This section uses **Qwen-Image** (`QwenImageTransformer2DModel`) as the reference implementation. Qwen-Image is a **dual-stream** transformer (text + image) that performs **joint attention** across the concatenated sequences. Because of that, when enabling sequence parallel you typically:
|
||||
|
||||
- Chunk **image tokens** (`hidden_states`) across SP ranks along the **sequence dimension**.
|
||||
- Keep **text embeddings** (`encoder_hidden_states`) **replicated** on all SP ranks for correctness (unless you implement full joint-SP semantics and explicitly split text too).
|
||||
- Chunk **image RoPE freqs** to match the chunked image tokens.
|
||||
- `all_gather` the final output back along the sequence dimension.
|
||||
|
||||
First, add the sequence-parallel helpers and (for Qwen-Image) the forward-context flag. Then, in the transformer's `forward()`:
|
||||
|
||||
- Chunk `hidden_states` by SP world size.
|
||||
- Set `get_forward_context().split_text_embed_in_sp = False` because Qwen-Image uses joint attention and we would use a full text embedding in each rank.
|
||||
- After `pos_embed`, chunk `img_freqs` on `dim=0` (token axis) to match the chunked image tokens; only chunk `txt_freqs` if you explicitly split text embeddings in SP.
|
||||
|
||||
Taking `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` as an example:
|
||||
|
||||
```diff
|
||||
from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
get_sequence_parallel_rank,
|
||||
get_sequence_parallel_world_size,
|
||||
get_sp_group,
|
||||
)
|
||||
from vllm_omni.diffusion.forward_context import get_forward_context
|
||||
|
||||
class QwenImageTransformer2DModel(...):
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor = None,
|
||||
encoder_hidden_states_mask: torch.Tensor = None,
|
||||
timestep: torch.LongTensor = None,
|
||||
img_shapes: list[tuple[int, int, int]] | None = None,
|
||||
txt_seq_lens: list[int] | None = None,
|
||||
...
|
||||
):
|
||||
+ if self.parallel_config.sequence_parallel_size > 1:
|
||||
+ # Chunk image tokens along sequence dimension.
|
||||
+ hidden_states = torch.chunk(
|
||||
+ hidden_states,
|
||||
+ get_sequence_parallel_world_size(),
|
||||
+ dim=-2,
|
||||
+ )[get_sequence_parallel_rank()]
|
||||
+
|
||||
+ # Qwen-Image uses *dual-stream* (text + image) and runs *joint attention*.
|
||||
+ # Text embeddings should be replicated across SP ranks for correctness.
|
||||
+ get_forward_context().split_text_embed_in_sp = False
|
||||
|
||||
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):
|
||||
+ return torch.chunk(freqs, get_sequence_parallel_world_size(), dim=0)[
|
||||
+ get_sequence_parallel_rank()
|
||||
+ ]
|
||||
+
|
||||
+ if self.parallel_config.sequence_parallel_size > 1:
|
||||
+ img_freqs, txt_freqs = image_rotary_emb
|
||||
+ img_freqs = get_rotary_emb_chunk(img_freqs)
|
||||
+ if get_forward_context().split_text_embed_in_sp:
|
||||
+ txt_freqs = get_rotary_emb_chunk(txt_freqs)
|
||||
+ image_rotary_emb = (img_freqs, txt_freqs)
|
||||
```
|
||||
|
||||
Next, at the end of the `forward()` function, call `get_sp_group().all_gather` to gather the chunked outputs across devices and concatenate them along the sequence dimension (matching the earlier `dim=-2` chunking):
|
||||
|
||||
```diff
|
||||
class QwenImageTransformer2DModel(...):
|
||||
def forward(...):
|
||||
...
|
||||
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)
|
||||
```
|
||||
|
||||
|
||||
### CFG-Parallel
|
||||
|
||||
##### Offline Inference
|
||||
|
||||
CFG-Parallel is enabled through `DiffusionParallelConfig(cfg_parallel_size=...)`. The recommended configuration is `cfg_parallel_size=2` (one rank for the positive branch and one rank for the negative branch).
|
||||
|
||||
An example of offline inference using CFG-Parallel (image-to-image) is shown below:
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image-Edit",
|
||||
parallel_config=DiffusionParallelConfig(cfg_parallel_size=2),
|
||||
)
|
||||
|
||||
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,
|
||||
)
|
||||
```
|
||||
|
||||
Notes:
|
||||
|
||||
- CFG-Parallel is only effective when **true CFG** is enabled (i.e., `true_cfg_scale > 1` and a `negative_prompt` is provided).
|
||||
|
||||
#### How to parallelize a pipeline
|
||||
|
||||
This section describes how to add CFG-Parallel to a diffusion **pipeline**. We use the Qwen-Image pipeline (`vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py`) as the reference implementation.
|
||||
|
||||
In `QwenImagePipeline`, each diffusion step runs two denoiser forward passes sequentially:
|
||||
|
||||
- positive (prompt-conditioned)
|
||||
- negative (negative-prompt-conditioned)
|
||||
|
||||
CFG-Parallel assigns these two branches to different ranks in the **CFG group** and synchronizes the results.
|
||||
|
||||
Below is an example of CFG-Parallel implementation:
|
||||
|
||||
```python
|
||||
def diffuse(
|
||||
self,
|
||||
...
|
||||
):
|
||||
# Enable CFG-parallel: rank0 computes positive, rank1 computes negative.
|
||||
cfg_parallel_ready = do_true_cfg and get_classifier_free_guidance_world_size() > 1
|
||||
|
||||
self.transformer.do_true_cfg = do_true_cfg
|
||||
|
||||
if cfg_parallel_ready:
|
||||
cfg_group = get_cfg_group()
|
||||
cfg_rank = get_classifier_free_guidance_rank()
|
||||
|
||||
if cfg_rank == 0:
|
||||
local_pred = self.transformer(
|
||||
hidden_states=latents,
|
||||
timestep=timestep / 1000,
|
||||
guidance=guidance,
|
||||
encoder_hidden_states_mask=prompt_embeds_mask,
|
||||
encoder_hidden_states=prompt_embeds,
|
||||
img_shapes=img_shapes,
|
||||
txt_seq_lens=txt_seq_lens,
|
||||
attention_kwargs=self.attention_kwargs,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
else:
|
||||
local_pred = self.transformer(
|
||||
hidden_states=latents,
|
||||
timestep=timestep / 1000,
|
||||
guidance=guidance,
|
||||
encoder_hidden_states_mask=negative_prompt_embeds_mask,
|
||||
encoder_hidden_states=negative_prompt_embeds,
|
||||
img_shapes=img_shapes,
|
||||
txt_seq_lens=negative_txt_seq_lens,
|
||||
attention_kwargs=self.attention_kwargs,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
|
||||
gathered = cfg_group.all_gather(local_pred, separate_tensors=True)
|
||||
if cfg_rank == 0:
|
||||
noise_pred = gathered[0]
|
||||
neg_noise_pred = gathered[1]
|
||||
comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
|
||||
cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
|
||||
noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
|
||||
noise_pred = comb_pred * (cond_norm / noise_norm)
|
||||
latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
|
||||
cfg_group.broadcast(latents, src=0)
|
||||
else:
|
||||
# fallback: run positive then negative sequentially on one rank
|
||||
...
|
||||
```
|
||||
@@ -0,0 +1,132 @@
|
||||
# TeaCache Configuration Guide
|
||||
|
||||
TeaCache speeds up diffusion model inference by caching transformer computations when consecutive timesteps are similar. This typically provides **1.5x-2.0x speedup** with minimal quality loss.
|
||||
|
||||
## Quick Start
|
||||
|
||||
Enable TeaCache by setting `cache_backend` to `"tea_cache"`:
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
# Simple configuration - model_type is automatically extracted from pipeline.__class__.__name__
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
cache_backend="tea_cache",
|
||||
cache_config={
|
||||
"rel_l1_thresh": 0.2 # Optional, defaults to 0.2
|
||||
}
|
||||
)
|
||||
outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
|
||||
```
|
||||
|
||||
### Using Environment Variable
|
||||
|
||||
You can also enable TeaCache via environment variable:
|
||||
|
||||
```bash
|
||||
export DIFFUSION_CACHE_BACKEND=tea_cache
|
||||
```
|
||||
|
||||
Then initialize without explicitly setting `cache_backend`:
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
cache_config={"rel_l1_thresh": 0.2} # Optional
|
||||
)
|
||||
```
|
||||
|
||||
## Online Serving (OpenAI-Compatible)
|
||||
|
||||
Enable TeaCache for online serving by passing `--cache-backend tea_cache` when starting the server:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image --omni --port 8091 \
|
||||
--cache-backend tea_cache \
|
||||
--cache-config '{"rel_l1_thresh": 0.2}'
|
||||
```
|
||||
|
||||
## Configuration Parameters
|
||||
|
||||
### `rel_l1_thresh` (float, default: `0.2`)
|
||||
|
||||
Controls the balance between speed and quality. Lower values prioritize quality, higher values prioritize speed.
|
||||
|
||||
**Recommended values:**
|
||||
|
||||
- `0.2` - **~1.5x speedup** with minimal quality loss (recommended)
|
||||
- `0.4` - **~1.8x speedup** with slight quality loss
|
||||
- `0.6` - **~2.0x speedup** with noticeable quality loss
|
||||
- `0.8` - **~2.25x speedup** with significant quality loss
|
||||
|
||||
## Examples
|
||||
|
||||
### Python API
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
cache_backend="tea_cache",
|
||||
cache_config={"rel_l1_thresh": 0.2}
|
||||
)
|
||||
outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
|
||||
```
|
||||
|
||||
## Performance Tuning
|
||||
|
||||
Start with the default `rel_l1_thresh=0.2` and adjust based on your needs:
|
||||
|
||||
- **Maximum quality**: Use `0.1-0.2`
|
||||
- **Balanced**: Use `0.2-0.4` (recommended)
|
||||
- **Maximum speed**: Use `0.6-0.8` (may reduce quality)
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Quality Degradation
|
||||
|
||||
If you notice quality issues, lower the threshold:
|
||||
|
||||
```python
|
||||
cache_config={"rel_l1_thresh": 0.1} # More conservative caching
|
||||
```
|
||||
|
||||
## Supported Models
|
||||
|
||||
### ImageGen
|
||||
|
||||
<style>
|
||||
th {
|
||||
white-space: nowrap;
|
||||
min-width: 0 !important;
|
||||
}
|
||||
</style>
|
||||
|
||||
| Architecture | Models | Example HF Models |
|
||||
|--------------|--------|-------------------|
|
||||
| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
|
||||
| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
|
||||
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` |
|
||||
| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
|
||||
|
||||
### VideoGen
|
||||
|
||||
No VideoGen models are supported by TeaCache yet.
|
||||
|
||||
### Coming Soon
|
||||
|
||||
<style>
|
||||
th {
|
||||
white-space: nowrap;
|
||||
min-width: 0 !important;
|
||||
}
|
||||
</style>
|
||||
|
||||
| Architecture | Models | Example HF Models |
|
||||
|--------------|--------|-------------------|
|
||||
| `FluxPipeline` | Flux | - |
|
||||
| `CogVideoXPipeline` | CogVideoX | - |
|
||||
Reference in new issue
Block a user