[Diffusion] Add cache-dit and unify diffusion cache backend interface (#250)
Signed-off-by: samithuang <285365963@qq.com> Signed-off-by: SamitHuang <285365963@qq.com> Signed-off-by: Samit <285365963@qq.com> Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
This commit is contained in:
SamitandJiangyun Zhu authored and GitHub committed
1 parent
63f1e347ad
commit
2b4a9fe729
24 files changed
+1644
-367
No files matched your search
@@ -0,0 +1,327 @@
|
||||
# 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:
|
||||
- `Fn_compute_blocks=8` (aligned with cache-dit's recommended default)
|
||||
- `Bn_compute_blocks=0`
|
||||
- `max_warmup_steps=8`
|
||||
- `residual_diff_threshold=0.08`
|
||||
- No TaylorSeer or SCM enabled
|
||||
|
||||
### 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,
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
### Example Script
|
||||
|
||||
See `examples/offline_inference/qwen_image/text_to_image.py` for a complete working example with cache-dit acceleration.
|
||||
|
||||
```bash
|
||||
# Enable cache-dit with default parameters
|
||||
cd examples/offline_inference/qwen_image
|
||||
python text_to_image.py \
|
||||
--prompt "a cup of coffee on the table" \
|
||||
--enable_cache_dit \
|
||||
--num_inference_steps 50
|
||||
```
|
||||
|
||||
The `--enable_cache_dit` flag enables cache-dit acceleration with these customized parameters:
|
||||
|
||||
```python
|
||||
omni = Omni(
|
||||
...
|
||||
cache_backend="cache_dit" if args.enable_cache_dit else None,
|
||||
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",
|
||||
"scm_steps_policy": "dynamic",
|
||||
},
|
||||
)
|
||||
|
||||
```
|
||||
|
||||
## 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: 8): Number of initial steps where caching is disabled to ensure the model learns sufficient features before caching begins.
|
||||
- `residual_diff_threshold` (float, default: 0.08): Threshold for residual difference. Higher values lead to faster performance but may reduce precision.
|
||||
- `max_cached_steps` (int, default: -1): Maximum number of cached steps. Set to -1 for unlimited caching.
|
||||
- `max_continuous_cached_steps` (int, default: -1): Maximum number of consecutive cached steps. Set to -1 for unlimited consecutive caching.
|
||||
|
||||
**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**:
|
||||
- Start with `Fn_compute_blocks=8` and adjust based on your model size
|
||||
- Increase `residual_diff_threshold` (e.g., 0.12-0.15) for faster inference with slight quality trade-off
|
||||
- Reduce `max_warmup_steps` (e.g., 4-6) for faster startup, but ensure sufficient warmup
|
||||
|
||||
### 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:
|
||||
- `"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**:
|
||||
- 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
|
||||
|
||||
## Complete Configuration Examples
|
||||
|
||||
### Example 1: DBCache Only (Recommended Starting Point)
|
||||
|
||||
```python
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
cache_backend="cache_dit",
|
||||
cache_config={
|
||||
"Fn_compute_blocks": 8,
|
||||
"Bn_compute_blocks": 0,
|
||||
"max_warmup_steps": 8,
|
||||
"residual_diff_threshold": 0.08,
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
### Example 2: DBCache + TaylorSeer (Balanced Speed/Quality)
|
||||
|
||||
```python
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
cache_backend="cache_dit",
|
||||
cache_config={
|
||||
# DBCache
|
||||
"Fn_compute_blocks": 8,
|
||||
"Bn_compute_blocks": 0,
|
||||
"max_warmup_steps": 8,
|
||||
"residual_diff_threshold": 0.12,
|
||||
# TaylorSeer
|
||||
"enable_taylorseer": True,
|
||||
"taylorseer_order": 1,
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
### Example 3: DBCache + SCM
|
||||
|
||||
```python
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
cache_backend="cache_dit",
|
||||
cache_config={
|
||||
# DBCache
|
||||
"Fn_compute_blocks": 8,
|
||||
"Bn_compute_blocks": 0,
|
||||
"max_warmup_steps": 4,
|
||||
"residual_diff_threshold": 0.12,
|
||||
# SCM
|
||||
"scm_steps_mask_policy": "fast",
|
||||
"scm_steps_policy": "dynamic",
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
### Example 4: All Methods Combined
|
||||
|
||||
```python
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
cache_backend="cache_dit",
|
||||
cache_config={
|
||||
# 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",
|
||||
"scm_steps_policy": "dynamic",
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
## Dynamic num_inference_steps
|
||||
|
||||
Cache-dit automatically handles changes in `num_inference_steps` during inference. When you call `generate()` with a different `num_inference_steps` than the previous request, cache-dit will:
|
||||
|
||||
1. Automatically refresh the cache context
|
||||
2. Regenerate SCM masks if SCM is enabled
|
||||
3. Update internal configurations to match the new step count
|
||||
|
||||
**Example**:
|
||||
|
||||
```python
|
||||
# First request with 50 steps
|
||||
images1 = omni.generate("prompt 1", num_inference_steps=50)
|
||||
|
||||
# Second request with 28 steps - cache context automatically refreshes
|
||||
images2 = omni.generate("prompt 2", num_inference_steps=28)
|
||||
```
|
||||
|
||||
No manual intervention is required - the refresh happens automatically and efficiently.
|
||||
|
||||
## Configuration Reference
|
||||
|
||||
### DiffusionCacheConfig Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `Fn_compute_blocks` | int | 8 | First n blocks for difference computation (aligned with cache-dit's default) |
|
||||
| `Bn_compute_blocks` | int | 0 | Last n blocks for fusion |
|
||||
| `max_warmup_steps` | int | 8 | Steps before caching starts |
|
||||
| `max_cached_steps` | int | -1 | Max cached steps (-1 = unlimited) |
|
||||
| `max_continuous_cached_steps` | int | -1 | Max consecutive cached steps (-1 = unlimited) |
|
||||
| `residual_diff_threshold` | float | 0.08 | Residual difference threshold |
|
||||
| `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 |
|
||||
| `taylorseer_order` | int | 1 | Taylor expansion order |
|
||||
| `scm_steps_mask_policy` | str \| None | None | SCM mask policy ("slow", "medium", "fast", "ultra") |
|
||||
| `scm_steps_policy` | str | "dynamic" | SCM computation policy ("dynamic" or "static") |
|
||||
|
||||
## Tips and Best Practices
|
||||
|
||||
### 1. Start Simple
|
||||
The easiest way to get started is to enable cache-dit without any configuration:
|
||||
|
||||
```python
|
||||
omni = Omni(model="Qwen/Qwen-Image", cache_backend="cache_dit")
|
||||
```
|
||||
|
||||
This uses cache-dit's recommended defaults (`Fn_compute_blocks=8`). Once you understand the baseline performance, you can customize the configuration to add other acceleration methods or fine-tune parameters.
|
||||
|
||||
### 2. Quality vs Speed Trade-offs
|
||||
- **Higher quality**: Lower `residual_diff_threshold` (0.08), more `Fn_compute_blocks` (8-12), SCM `"slow"` or `"medium"`
|
||||
- **Faster inference**: Higher `residual_diff_threshold` (0.12-0.15), fewer `max_warmup_steps` (4-6), SCM `"fast"` or `"ultra"`
|
||||
|
||||
### 3. Model-Specific Tuning
|
||||
Different models may benefit from different configurations:
|
||||
- **Smaller models**: May work well with `Fn_compute_blocks=4-6`
|
||||
- **Larger models**: Often benefit from `Fn_compute_blocks=8-12`
|
||||
- **High-resolution images**: May need lower `residual_diff_threshold` for quality
|
||||
|
||||
### 4. Combining Methods
|
||||
- **DBCache + TaylorSeer**: Good balance, works well for most cases
|
||||
- **DBCache + SCM**: Can try when quality can be slightly compromised
|
||||
- **All three**: Maximum acceleration if the parameters are carefully tuned
|
||||
|
||||
## 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,99 @@
|
||||
# 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.
|
||||
|
||||
## 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
|
||||
|
||||
Both methods can provide significant speedups (typically **1.5x-2.0x**) while maintaining high output quality.
|
||||
|
||||
## Quick Comparison
|
||||
|
||||
| Method | Configuration | Description | Best For |
|
||||
|--------|--------------|-------------|----------|
|
||||
| **TeaCache** | `cache_backend="tea_cache"` | Simple, adaptive caching with minimal configuration | Quick setup, balanced speed/quality |
|
||||
| **Cache-DiT** | `cache_backend="cache_dit"` | Advanced caching with multiple techniques (DBCache, TaylorSeer, SCM) | Maximum acceleration, fine-grained control |
|
||||
|
||||
## Supported Models
|
||||
|
||||
The following table shows which models are currently supported by each cache backend:
|
||||
|
||||
| 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` | ❌ | ✅* |
|
||||
|
||||
\* Qwen-Image-Edit support is not fully verified
|
||||
|
||||
## Performance Benchmarks
|
||||
|
||||
The following benchmarks were measured on **Qwen/Qwen-Image** model generating 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:
|
||||
|
||||
- Specific model and use case
|
||||
- Hardware configuration
|
||||
- Careful parameter tuning
|
||||
- Different inference settings (e.g., number of steps, image resolution)
|
||||
|
||||
For optimal performance in your specific scenario, we recommend experimenting with different parameter configurations as described in the detailed guides below.
|
||||
|
||||
| Configuration | Generation Time | Speedup | Notes |
|
||||
|---------------|----------------|---------|-------|
|
||||
| **Baseline (diffusers)** | 20.0s | 1.0x | No acceleration |
|
||||
| **TeaCache** (`rel_l1_thresh=0.2`) | 10.47s | **1.91x** | Recommended default setting |
|
||||
| **Cache-DiT (DBCache + TaylorSeer)** | 10.8s | **1.85x** | Fn=1, Bn=0, W=8, TaylorSeer order=1 |
|
||||
| **Cache-DiT (DBCache + TaylorSeer + SCM)** | 14.0s | **1.43x** | Fn=8, Bn=0, W=4, TaylorSeer order=1, SCM fast |
|
||||
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Using TeaCache
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
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 Cache-DiT
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
cache_backend="cache_dit",
|
||||
cache_config={
|
||||
"Fn_compute_blocks": 1,
|
||||
"Bn_compute_blocks": 0,
|
||||
"max_warmup_steps": 8,
|
||||
"enable_taylorseer": True,
|
||||
"taylorseer_order": 1,
|
||||
}
|
||||
)
|
||||
|
||||
outputs = omni.generate(prompt="A cat sitting on a windowsill", 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
|
||||
@@ -4,15 +4,15 @@ TeaCache speeds up diffusion model inference by caching transformer computations
|
||||
|
||||
## Quick Start
|
||||
|
||||
Enable TeaCache by setting `cache_adapter` to `"tea_cache"`:
|
||||
Enable TeaCache by setting `cache_backend` to `"tea_cache"`:
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
# Simple configuration - model_class_name and model_type are auto-detected
|
||||
# Simple configuration - model_type is automatically extracted from pipeline.__class__.__name__
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
cache_adapter="tea_cache",
|
||||
cache_backend="tea_cache",
|
||||
cache_config={
|
||||
"rel_l1_thresh": 0.2 # Optional, defaults to 0.2
|
||||
}
|
||||
@@ -25,10 +25,10 @@ outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_st
|
||||
You can also enable TeaCache via environment variable:
|
||||
|
||||
```bash
|
||||
export DIFFUSION_CACHE_ADAPTER=tea_cache
|
||||
export DIFFUSION_CACHE_BACKEND=tea_cache
|
||||
```
|
||||
|
||||
Then initialize without explicitly setting `cache_adapter`:
|
||||
Then initialize without explicitly setting `cache_backend`:
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
@@ -61,7 +61,7 @@ from vllm_omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
cache_adapter="tea_cache",
|
||||
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)
|
||||
|
||||
Reference in new issue
Block a user