diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index a663f4c93..952954e81 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -32,6 +32,23 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" + - label: "Diffusion Cache Backend Test" + timeout_in_minutes: 15 + depends_on: image-build + commands: + - pytest -s -v tests/e2e/offline_inference/test_cache_dit.py tests/e2e/offline_inference/test_teacache.py + agents: + queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + environment: + - "HF_HOME=/fsx/hf_cache" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" + - label: "Omni Model Test" timeout_in_minutes: 15 depends_on: image-build diff --git a/.buildkite/scripts/simple_test.sh b/.buildkite/scripts/simple_test.sh index 58379c964..157ad9ca0 100755 --- a/.buildkite/scripts/simple_test.sh +++ b/.buildkite/scripts/simple_test.sh @@ -50,3 +50,4 @@ VENV_PYTHON="${VENV_DIR}/bin/python" "${UV_BIN}" pip install --python "${VENV_PYTHON}" vllm==0.11.0 "${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/ diff --git a/docs/.nav.yml b/docs/.nav.yml index d54517df9..fb34f1742 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -21,7 +21,9 @@ nav: - configuration/README.md - configuration/* - Diffusion Acceleration: + - Overview: user_guide/diffusion_acceleration.md - TeaCache: user_guide/teacache.md + - Cache-DiT: user_guide/cache_dit_acceleration.md - Models: - models/supported_models.md - Developer Guide: diff --git a/docs/user_guide/cache_dit_acceleration.md b/docs/user_guide/cache_dit_acceleration.md new file mode 100644 index 000000000..506dff731 --- /dev/null +++ b/docs/user_guide/cache_dit_acceleration.md @@ -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) diff --git a/docs/user_guide/diffusion_acceleration.md b/docs/user_guide/diffusion_acceleration.md new file mode 100644 index 000000000..9dc5aecac --- /dev/null +++ b/docs/user_guide/diffusion_acceleration.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 diff --git a/docs/user_guide/teacache.md b/docs/user_guide/teacache.md index 43afef3a9..5df0a4e1a 100644 --- a/docs/user_guide/teacache.md +++ b/docs/user_guide/teacache.md @@ -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) diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index 966679215..4626e647f 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import argparse +import time from pathlib import Path import torch @@ -12,7 +13,11 @@ from vllm_omni.utils.platform_utils import detect_device_type, is_npu def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Generate an image with Qwen-Image.") - parser.add_argument("--model", default="Qwen/Qwen-Image", help="Diffusion model name or local path.") + parser.add_argument( + "--model", + default="Qwen/Qwen-Image", + help="Diffusion model name or local path. Supported models: Qwen/Qwen-Image, Tongyi-MAI/Z-Image-Turbo", + ) parser.add_argument("--prompt", default="a cup of coffee on the table", help="Text prompt for image generation.") parser.add_argument("--seed", type=int, default=42, help="Random seed for deterministic results.") parser.add_argument( @@ -41,6 +46,17 @@ def parse_args() -> argparse.Namespace: default=50, help="Number of denoising steps for the diffusion sampler.", ) + parser.add_argument( + "--cache_backend", + type=str, + default=None, + choices=["cache_dit", "tea_cache"], + help=( + "Cache backend to use for acceleration. " + "Options: 'cache_dit' (DBCache + SCM + TaylorSeer), 'tea_cache' (Timestep Embedding Aware Cache). " + "Default: None (no cache acceleration)." + ), + ) return parser.parse_args() @@ -53,11 +69,53 @@ def main(): vae_use_slicing = is_npu() vae_use_tiling = is_npu() + # Configure cache based on backend type + cache_config = None + if args.cache_backend == "cache_dit": + # cache-dit configuration: Hybrid DBCache + SCM + TaylorSeer + # All parameters marked with [cache-dit only] in DiffusionCacheConfig + 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) + "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": "fast", # SCM mask policy: "slow", "medium", "fast", "ultra" + "scm_steps_policy": "dynamic", # SCM steps policy: "dynamic" or "static" + } + elif args.cache_backend == "tea_cache": + # TeaCache configuration + # All parameters marked with [tea_cache only] in DiffusionCacheConfig + cache_config = { + # TeaCache parameters [tea_cache only] + "rel_l1_thresh": 0.2, # Threshold for accumulated relative L1 distance + # Note: coefficients will use model-specific defaults based on model_type + # (e.g., QwenImagePipeline or FluxPipeline) + } + omni = Omni( model=args.model, vae_use_slicing=vae_use_slicing, vae_use_tiling=vae_use_tiling, + cache_backend=args.cache_backend, + cache_config=cache_config, ) + + # Time profiling for generation + print(f"\n{'=' * 60}") + print("Generation Configuration:") + print(f" Model: {args.model}") + print(f" Inference steps: {args.num_inference_steps}") + print(f" Cache backend: {args.cache_backend if args.cache_backend else 'None (no acceleration)'}") + print(f" Image size: {args.width}x{args.height}") + print(f"{'=' * 60}\n") + + generation_start = time.perf_counter() images = omni.generate( args.prompt, height=args.height, @@ -67,6 +125,11 @@ def main(): 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 + + # Print profiling results + print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)") output_path = Path(args.output) output_path.parent.mkdir(parents=True, exist_ok=True) diff --git a/pyproject.toml b/pyproject.toml index de3200df3..ad014c400 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,7 +37,8 @@ dependencies = [ "diffusers==0.35.2", "accelerate==1.12.0", "gradio==5.50", - "soundfile>=0.13.1" + "soundfile>=0.13.1", + "cache-dit==1.1.8", # "vllm==0.11.0", # TODO: fix the entrypoints overwrite problem ] diff --git a/tests/diffusion/cache/__init__.py b/tests/diffusion/cache/__init__.py new file mode 100644 index 000000000..7d7fc7fcd --- /dev/null +++ b/tests/diffusion/cache/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Unit tests for cache backends (cache-dit and teacache). +""" diff --git a/tests/diffusion/cache/test_cache_backends.py b/tests/diffusion/cache/test_cache_backends.py new file mode 100644 index 000000000..5713c6ba7 --- /dev/null +++ b/tests/diffusion/cache/test_cache_backends.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Unit tests for cache backends (cache-dit and teacache). + +This module tests the cache backend implementations: +- CacheDiTBackend: cache-dit acceleration backend +- TeaCacheBackend: TeaCache hook-based backend +- Cache selector function: get_cache_backend +- DiffusionCacheConfig: configuration dataclass +""" + +from unittest.mock import Mock, patch + +import pytest + +from vllm_omni.diffusion.cache.cache_dit_backend import ( + CacheDiTBackend, +) +from vllm_omni.diffusion.cache.selector import get_cache_backend +from vllm_omni.diffusion.cache.teacache.backend import TeaCacheBackend +from vllm_omni.diffusion.data import DiffusionCacheConfig + + +class TestCacheDiTBackend: + """Test CacheDiTBackend implementation.""" + + def test_init_with_dict(self): + """Test initialization with dictionary config.""" + config_dict = {"Fn_compute_blocks": 4, "max_warmup_steps": 8} + backend = CacheDiTBackend(config_dict) + assert backend.config.Fn_compute_blocks == 4 + assert backend.config.max_warmup_steps == 8 + assert backend.enabled is False + + def test_init_with_config_object(self): + """Test initialization with DiffusionCacheConfig object.""" + config = DiffusionCacheConfig(Fn_compute_blocks=4) + backend = CacheDiTBackend(config) + assert backend.config.Fn_compute_blocks == 4 + assert backend.enabled is False + + @patch("vllm_omni.diffusion.cache.cache_dit_backend.CACHE_DIT_AVAILABLE", True) + @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit") + def test_enable_single_transformer(self, mock_cache_dit): + """Test enabling cache-dit on single-transformer pipeline.""" + # Mock pipeline + mock_pipeline = Mock() + mock_pipeline.__class__.__name__ = "DiTPipeline" + mock_transformer = Mock() + mock_pipeline.transformer = mock_transformer + + # Mock cache_dit functions + mock_cache_dit.enable_cache = Mock() + mock_cache_dit.refresh_context = Mock() + + backend = CacheDiTBackend({"Fn_compute_blocks": 2}) + backend.enable(mock_pipeline) + + # Verify cache-dit was enabled + assert backend.enabled is True + assert backend._refresh_func is not None + mock_cache_dit.enable_cache.assert_called_once() + + @patch("vllm_omni.diffusion.cache.cache_dit_backend.CACHE_DIT_AVAILABLE", True) + @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit") + def test_refresh(self, mock_cache_dit): + """Test refreshing cache context with SCM mask policy updates when num_inference_steps changes.""" + # Mock pipeline + mock_pipeline = Mock() + mock_pipeline.__class__.__name__ = "DiTPipeline" + mock_transformer = Mock() + mock_pipeline.transformer = mock_transformer + + # Mock cache_dit functions + mock_cache_dit.enable_cache = Mock() + mock_cache_dit.refresh_context = Mock() + mock_steps_mask_50 = [1, 0, 1, 0, 1] * 10 # Mock mask for 50 steps + mock_steps_mask_100 = [1, 0, 1, 0, 1] * 20 # Mock mask for 100 steps + mock_cache_dit.steps_mask = Mock(side_effect=[mock_steps_mask_50, mock_steps_mask_100]) + + # Enable cache-dit with SCM enabled (using mask policy) + config = DiffusionCacheConfig( + scm_steps_mask_policy="fast", + scm_steps_policy="dynamic", + ) + backend = CacheDiTBackend(config) + backend.enable(mock_pipeline) + + # First refresh with 50 steps + backend.refresh(mock_pipeline, num_inference_steps=50) + assert backend._last_num_inference_steps == 50 + + # Verify steps_mask was called with mask policy (not direct steps mask) + mock_cache_dit.steps_mask.assert_called_with(mask_policy="fast", total_steps=50) + assert mock_cache_dit.steps_mask.call_count == 1 + + # Verify refresh_context was called with cache_config (SCM path) + mock_cache_dit.refresh_context.assert_called_once() + call_args = mock_cache_dit.refresh_context.call_args + assert call_args[0][0] == mock_transformer + # Check that cache_config was passed (not num_inference_steps directly when SCM is enabled) + assert "cache_config" in call_args[1] + cache_config_arg = call_args[1]["cache_config"] + assert cache_config_arg is not None + + # Change num_inference_steps and refresh again + mock_cache_dit.refresh_context.reset_mock() + backend.refresh(mock_pipeline, num_inference_steps=100) + + # Verify steps_mask was called again with new num_inference_steps (using mask policy) + assert mock_cache_dit.steps_mask.call_count == 2 + # Check the last call was with 100 steps and mask policy + assert mock_cache_dit.steps_mask.call_args_list[-1].kwargs["total_steps"] == 100 + assert mock_cache_dit.steps_mask.call_args_list[-1].kwargs["mask_policy"] == "fast" + + # Verify refresh_context was called again with updated mask + mock_cache_dit.refresh_context.assert_called_once() + call_args = mock_cache_dit.refresh_context.call_args + assert call_args[0][0] == mock_transformer + assert "cache_config" in call_args[1] + assert backend._last_num_inference_steps == 100 + + +class TestTeaCacheBackend: + """Test TeaCacheBackend implementation.""" + + def test_init(self): + """Test initialization.""" + config = DiffusionCacheConfig(rel_l1_thresh=0.3) + backend = TeaCacheBackend(config) + assert backend.config.rel_l1_thresh == 0.3 + assert backend.enabled is False + + @patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook") + def test_enable(self, mock_apply_hook): + """Test enabling TeaCache on pipeline.""" + # Mock pipeline + mock_pipeline = Mock() + mock_pipeline.__class__.__name__ = "QwenImagePipeline" + mock_transformer = Mock() + mock_pipeline.transformer = mock_transformer + + config = DiffusionCacheConfig(rel_l1_thresh=0.3) + backend = TeaCacheBackend(config) + backend.enable(mock_pipeline) + + # Verify hook was applied + assert backend.enabled is True + assert hasattr(mock_pipeline, "_cache_backend") + assert mock_pipeline._cache_backend is backend + mock_apply_hook.assert_called_once() + + @patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook") + def test_enable_with_coefficients(self, mock_apply_hook): + """Test enabling TeaCache with custom coefficients.""" + mock_pipeline = Mock() + mock_pipeline.__class__.__name__ = "QwenImagePipeline" + mock_transformer = Mock() + mock_pipeline.transformer = mock_transformer + + config = DiffusionCacheConfig(rel_l1_thresh=0.3, coefficients=[1.0, 0.5, 0.2, 0.1, 0.05]) + backend = TeaCacheBackend(config) + backend.enable(mock_pipeline) + + assert backend.enabled is True + mock_apply_hook.assert_called_once() + + @patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook") + def test_refresh(self, mock_apply_hook): + """Test refreshing TeaCache state.""" + mock_pipeline = Mock() + mock_pipeline.__class__.__name__ = "QwenImagePipeline" + mock_transformer = Mock() + mock_pipeline.transformer = mock_transformer + + # Mock hook registry + mock_hook = Mock() + mock_registry = Mock() + mock_registry.get_hook = Mock(return_value=mock_hook) + mock_registry.reset_hook = Mock() + mock_transformer._hook_registry = mock_registry + + config = DiffusionCacheConfig() + backend = TeaCacheBackend(config) + backend.enable(mock_pipeline) + + # Test refresh + backend.refresh(mock_pipeline, num_inference_steps=50) + mock_registry.reset_hook.assert_called_once() + + +class TestCacheSelector: + """Test cache backend selector function.""" + + def test_get_cache_backend_none(self): + """Test getting None backend.""" + backend = get_cache_backend(None, None) + assert backend is None + + backend = get_cache_backend("none", None) + assert backend is None + + def test_get_cache_backend_cache_dit(self): + """Test getting cache-dit backend.""" + config_dict = {"Fn_compute_blocks": 4} + backend = get_cache_backend("cache_dit", config_dict) + assert isinstance(backend, CacheDiTBackend) + assert backend.config.Fn_compute_blocks == 4 + + def test_get_cache_backend_tea_cache(self): + """Test getting teacache backend.""" + config_dict = {"rel_l1_thresh": 0.3} + backend = get_cache_backend("tea_cache", config_dict) + assert isinstance(backend, TeaCacheBackend) + assert backend.config.rel_l1_thresh == 0.3 + + def test_get_cache_backend_invalid(self): + """Test getting invalid backend raises error.""" + with pytest.raises(ValueError, match="Unsupported cache backend"): + get_cache_backend("invalid_backend", {}) diff --git a/tests/e2e/offline_inference/test_cache_dit.py b/tests/e2e/offline_inference/test_cache_dit.py new file mode 100644 index 000000000..0113b20e2 --- /dev/null +++ b/tests/e2e/offline_inference/test_cache_dit.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +System test for cache-dit backend. + +This test verifies that cache-dit acceleration works correctly with diffusion models. +It uses minimal settings to keep test time short for CI. +""" + +import os +import sys +from pathlib import Path + +import pytest +import torch + +# ruff: noqa: E402 +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from vllm_omni import Omni + +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" + +# Use random weights model for testing +models = ["riverclouds/qwen_image_random"] + + +@pytest.mark.parametrize("model_name", models) +def test_cache_dit(model_name: str): + """Test cache-dit backend with diffusion model.""" + # Configure cache-dit with minimal settings for fast testing + cache_config = { + "Fn_compute_blocks": 1, + "Bn_compute_blocks": 0, + "max_warmup_steps": 2, # Minimal warmup for fast test + "residual_diff_threshold": 0.24, + "max_continuous_cached_steps": 3, + } + + m = Omni( + model=model_name, + cache_backend="cache_dit", + cache_config=cache_config, + ) + + # Use minimal settings for fast testing + height = 256 + width = 256 + num_inference_steps = 4 # Minimal steps for fast test + + images = 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 + ) + + # Verify generation succeeded + assert images is not None + assert len(images) == 1 + # Check image size + assert images[0].width == width + assert images[0].height == height diff --git a/tests/e2e/offline_inference/test_teacache.py b/tests/e2e/offline_inference/test_teacache.py new file mode 100644 index 000000000..9e925f5e4 --- /dev/null +++ b/tests/e2e/offline_inference/test_teacache.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +System test for TeaCache backend. + +This test verifies that TeaCache acceleration works correctly with diffusion models. +It uses minimal settings to keep test time short for CI. +""" + +import os +import sys +from pathlib import Path + +import pytest +import torch + +# ruff: noqa: E402 +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from vllm_omni import Omni + +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" + +# Use random weights model for testing +models = ["riverclouds/qwen_image_random"] + + +@pytest.mark.parametrize("model_name", models) +def test_teacache(model_name: str): + """Test TeaCache backend with diffusion model.""" + # Configure TeaCache with default settings for fast testing + cache_config = { + "rel_l1_thresh": 0.2, # Default threshold + } + + m = Omni( + model=model_name, + cache_backend="tea_cache", + cache_config=cache_config, + ) + + # Use minimal settings for fast testing + height = 256 + width = 256 + num_inference_steps = 4 # Minimal steps for fast test + + images = 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 + ) + + # Verify generation succeeded + assert images is not None + assert len(images) == 1 + # Check image size + assert images[0].width == width + assert images[0].height == height diff --git a/vllm_omni/diffusion/cache/__init__.py b/vllm_omni/diffusion/cache/__init__.py index 5a3e969b1..a5968f612 100644 --- a/vllm_omni/diffusion/cache/__init__.py +++ b/vllm_omni/diffusion/cache/__init__.py @@ -1,41 +1,27 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project - """ Cache module for diffusion model inference acceleration. -This module provides a unified cache adapter system for different caching strategies: +This module provides a unified cache backend system for different caching strategies: - TeaCache: Timestep Embedding Aware Cache for adaptive transformer caching -- (Future: DeepCache, DistriFusion, etc.) +- cache-dit: DBCache, SCM, and TaylorSeer caching strategies -The cache adapter system uses a registry pattern where cache types can be -dynamically registered and instantiated. Cache is configured via OmniDiffusionConfig -or DIFFUSION_CACHE_ADAPTER environment variable. +Cache backends are instantiated directly via their constructors and configured via OmniDiffusionConfig. """ -from vllm_omni.diffusion.cache.apply import setup_cache -from vllm_omni.diffusion.cache.base import CacheAdapter -from vllm_omni.diffusion.cache.registry import ( - CacheType, - get_cache_adapter, - register_cache_adapter, -) +from vllm_omni.diffusion.cache.base import CacheBackend from vllm_omni.diffusion.cache.teacache import ( CacheContext, TeaCacheConfig, apply_teacache_hook, ) - -# Import teacache adapter to trigger registration -from vllm_omni.diffusion.cache.teacache.adapter import TeaCacheAdapter # noqa: F401 +from vllm_omni.diffusion.cache.teacache.backend import TeaCacheBackend __all__ = [ - "CacheAdapter", - "CacheType", + "CacheBackend", "TeaCacheConfig", "CacheContext", - "setup_cache", - "get_cache_adapter", - "register_cache_adapter", + "TeaCacheBackend", "apply_teacache_hook", ] diff --git a/vllm_omni/diffusion/cache/apply.py b/vllm_omni/diffusion/cache/apply.py deleted file mode 100644 index ec91180c1..000000000 --- a/vllm_omni/diffusion/cache/apply.py +++ /dev/null @@ -1,82 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -""" -Unified cache setup function for diffusion models. - -This module provides the main entry point for setting up cache adapters -on transformer models. It handles cache type selection, adapter instantiation, -and application to the model. -""" - -from typing import Any, Optional - -import torch -from vllm.logger import init_logger - -from vllm_omni.diffusion.cache.base import CacheAdapter -from vllm_omni.diffusion.cache.registry import get_cache_adapter - -logger = init_logger(__name__) - - -def setup_cache( - transformer: torch.nn.Module, - cache_type: str = "none", - cache_config: Optional[dict[str, Any]] = None, -) -> Optional[CacheAdapter]: - """ - Setup cache adapter for transformer (one-time initialization). - - This is the main entry point for applying caching to diffusion transformers. - It's called once during pipeline initialization and applies the appropriate - cache adapter based on configuration. - - Args: - transformer: Transformer module to apply cache to - cache_type: Type of cache to apply. Options: - - "none": No caching (default) - - "tea_cache": TeaCache adaptive caching - - Additional types can be registered via register_cache_adapter() - cache_config: Cache-specific configuration dictionary. Contents depend - on cache type: - - For "tea_cache": {"rel_l1_thresh": 0.2, "model_type": "QwenImagePipeline"} - - For other types: see their respective adapter documentation - - Returns: - CacheAdapter instance for state management, or None if cache_type="none" - - Example: - >>> # Setup TeaCache during pipeline initialization - >>> adapter = setup_cache( - ... transformer, - ... cache_type="tea_cache", - ... cache_config={"rel_l1_thresh": 0.2} - ... ) - >>> - >>> # Reset cache state before each generation - >>> if adapter is not None: - ... adapter.reset(transformer) - """ - # No caching requested - if cache_type == "none": - logger.debug("No cache adapter requested (cache_type='none')") - return None - - # Get adapter instance from registry - try: - adapter = get_cache_adapter(cache_type, cache_config or {}) - except ValueError as e: - logger.error(f"Failed to get cache adapter: {e}") - raise - - # Apply cache to transformer - logger.info(f"Applying {cache_type} cache to transformer") - try: - adapter.apply(transformer) - except Exception as e: - logger.error(f"Failed to apply cache adapter: {e}") - raise - - logger.info(f"Cache adapter '{cache_type}' successfully applied") - return adapter diff --git a/vllm_omni/diffusion/cache/base.py b/vllm_omni/diffusion/cache/base.py index de58fe1fa..2d7c64442 100644 --- a/vllm_omni/diffusion/cache/base.py +++ b/vllm_omni/diffusion/cache/base.py @@ -2,66 +2,99 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ -Base cache adapter interface for diffusion models. +Base cache backend interface for diffusion models. -This module defines the abstract base class that all cache adapters must implement. -Cache adapters provide a unified interface for applying different caching strategies -(TeaCache, DeepCache, etc.) to transformer models using hooks. +This module defines the abstract base class that all cache backends must implement. +Cache backends provide a unified interface for applying different caching strategies +to transformer models. + +Main cache backend implementations: +1. CacheDiTBackend: Implements cache-dit acceleration (DBCache, SCM, TaylorSeer) using + the cache-dit library. Inherits from CacheBackend. Used via cache_backend="cache_dit". +2. TeaCacheBackend: Hook-based backend for TeaCache acceleration. Inherits from + CacheBackend. Used via cache_backend="tea_cache". + +All backends implement the same interface: +- enable(pipeline): Enable cache on the pipeline +- refresh(pipeline, num_inference_steps, verbose): Refresh cache state +- is_enabled(): Check if cache is enabled """ from abc import ABC, abstractmethod from typing import Any -import torch +from vllm_omni.diffusion.data import DiffusionCacheConfig -class CacheAdapter(ABC): +class CacheBackend(ABC): """ - Abstract base class for cache adapters. + Abstract base class for cache backends. - Cache adapters apply caching strategies to transformer models to accelerate - inference. Each adapter type (TeaCache, DeepCache, etc.) implements the - apply() and reset() methods to manage cache lifecycle. + All cache backend implementations (CacheDiTBackend, TeaCacheBackend, etc.) inherit + from this base class and implement the enable() and refresh() methods to manage + cache lifecycle. + + Cache backends apply caching strategies to transformer models to accelerate + inference. Different backends use different underlying mechanisms (e.g., cache-dit + library for CacheDiTBackend, hooks for TeaCacheBackend), but all share the same + unified interface. Attributes: - config: Dictionary containing cache-specific configuration parameters + config: DiffusionCacheConfig instance containing cache-specific configuration parameters + enabled: Boolean flag indicating whether cache is enabled (set to True after enable() is called) """ - def __init__(self, config: dict[str, Any]): + def __init__(self, config: DiffusionCacheConfig): """ - Initialize cache adapter with configuration. + Initialize cache backend with configuration. Args: - config: Cache-specific configuration dictionary + config: DiffusionCacheConfig instance with cache-specific parameters """ self.config = config + self.enabled = False @abstractmethod - def apply(self, transformer: torch.nn.Module) -> None: + def enable(self, pipeline: Any) -> None: """ - Apply cache to transformer using hooks. + Enable cache on the pipeline. - This method should register the appropriate hooks on the transformer - to enable caching during inference. Called once during pipeline - initialization. + This method applies the caching strategy to the transformer(s) in the pipeline. + The specific implementation depends on the backend (e.g., hooks for TeaCacheBackend, + cache-dit library for CacheDiTBackend). Called once during pipeline initialization. Args: - transformer: Transformer module to apply cache to + pipeline: Diffusion pipeline instance. The backend can extract: + - transformer: via pipeline.transformer + - model_type: via pipeline.__class__.__name__ """ - raise NotImplementedError("Subclasses must implement apply()") + raise NotImplementedError("Subclasses must implement enable()") @abstractmethod - def reset(self, transformer: torch.nn.Module) -> None: + def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: """ - Reset cache state for new generation. + Refresh cache state for new generation. This method should clear any cached values and reset counters/accumulators. Called at the start of each generation to ensure clean state. Args: - transformer: Transformer module to reset cache on + pipeline: Diffusion pipeline instance. The backend can extract: + - transformer: via pipeline.transformer + num_inference_steps: Number of inference steps for the current generation. + May be used for cache context updates. + verbose: Whether to log refresh operations (default: True) """ - raise NotImplementedError("Subclasses must implement reset()") + raise NotImplementedError("Subclasses must implement refresh()") + + def is_enabled(self) -> bool: + """ + Check if cache is enabled on this backend. + + Returns: + True if cache is enabled, False otherwise. + """ + return self.enabled def __repr__(self) -> str: return f"{self.__class__.__name__}(config={self.config})" diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py new file mode 100644 index 000000000..63362159e --- /dev/null +++ b/vllm_omni/diffusion/cache/cache_dit_backend.py @@ -0,0 +1,394 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +cache-dit integration backend for vllm-omni. + +This module provides a CacheDiTBackend class to enable cache-dit acceleration on diffusion +pipelines in vllm-omni, supporting both single and dual-transformer architectures. +""" + +from typing import Any, Callable, Optional + +from vllm.logger import init_logger + +from vllm_omni.diffusion.cache.base import CacheBackend +from vllm_omni.diffusion.data import DiffusionCacheConfig, OmniDiffusionConfig + +logger = init_logger(__name__) + +try: + import cache_dit + from cache_dit import BlockAdapter, DBCacheConfig, ForwardPattern, ParamsModifier, TaylorSeerCalibratorConfig + + CACHE_DIT_AVAILABLE = True +except ImportError: + CACHE_DIT_AVAILABLE = False + logger.warning("cache-dit is not installed. Cache-dit acceleration will not be available.") + + +# 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) +# Will be populated after function definitions +CUSTOM_DIT_ENABLERS: dict[str, Callable] = {} + + +def _build_db_cache_config(cache_config: Any) -> DBCacheConfig: + """Build DBCacheConfig with optional SCM (Step Computation Masking) support. + + Args: + cache_config: DiffusionCacheConfig instance. + + Returns: + DBCacheConfig instance with SCM support if configured. + """ + + return DBCacheConfig( + # we will refresh the context when gets num_inference_steps in the first inference request + num_inference_steps=None, + 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, + ) + + +def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int], None]: + """Enable cache-dit for Wan2.2 dual-transformer architecture. + + Wan2.2 uses two transformers (transformer and transformer_2) that need + to be enabled together using BlockAdapter. + + Args: + pipeline: The Wan2.2 pipeline instance. + cache_config: DiffusionCacheConfig instance with cache configuration. + + Returns: + 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], + has_separate_cfg=True, + ), + ) + + # from https://github.com/vipshop/cache-dit/pull/542 + def _split_inference_steps(num_inference_steps: int) -> tuple[int, int]: + """Split inference steps into high-noise and low-noise steps for Wan2.2. + + This is an internal helper function specific to Wan2.2's dual-transformer + architecture that uses boundary_ratio to determine the split point. + + Args: + num_inference_steps: Total number of inference steps. + + 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 + else: + boundary_timestep = None + + # Set timesteps to calculate the split + device = next(pipeline.transformer.parameters()).device + pipeline.scheduler.set_timesteps(num_inference_steps, device=device) + + timesteps = pipeline.scheduler.timesteps + num_high_noise_steps = 0 # high-noise steps for transformer + for t in timesteps: + if boundary_timestep is None or t >= boundary_timestep: + num_high_noise_steps += 1 + # low-noise steps for transformer_2 + num_low_noise_steps = num_inference_steps - num_high_noise_steps + return num_high_noise_steps, num_low_noise_steps + + def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: + """Refresh cache context for both transformers with new num_inference_steps. + + Args: + 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) + else: + cache_dit.refresh_context( + pipeline.transformer, + cache_config=DBCacheConfig().reset( + num_inference_steps=num_high_noise_steps, + steps_computation_mask=cache_dit.steps_mask( + mask_policy=cache_config.scm_steps_mask_policy, total_steps=num_high_noise_steps + ), + steps_computation_policy=cache_config.scm_steps_policy, + ), + 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( + num_inference_steps=num_low_noise_steps, + steps_computation_mask=cache_dit.steps_mask( + mask_policy=cache_config.scm_steps_mask_policy, total_steps=num_low_noise_steps + ), + steps_computation_policy=cache_config.scm_steps_policy, + ), + verbose=verbose, + ) + + return refresh_cache_context + + +def enable_cache_for_flux(pipeline: Any, cache_config: Any) -> Callable[[int], None]: + """Enable cache-dit for Flux dual-transformer architecture. + + Flux uses two transformers (transformer and transformer_2) that need + to be enabled together using BlockAdapter. + + Args: + pipeline: The Flux pipeline instance. + cache_config: DiffusionCacheConfig instance with cache configuration. + + Returns: + A refresh function that can be called to update cache context with new num_inference_steps. + """ + raise NotImplementedError("cache-dit is not implemented for Flux pipeline.") + + +def enable_cache_for_dit(pipeline: Any, cache_config: Any) -> Callable[[int], None]: + """Enable cache-dit for regular single-transformer DiT models. + + Args: + pipeline: The diffusion pipeline instance. + cache_config: DiffusionCacheConfig instance with cache configuration. + + Returns: + A refresh function that can be called to update cache context with new num_inference_steps. + """ + # Build DBCacheConfig with optional SCM support + db_cache_config = _build_db_cache_config(cache_config) + + # Build calibrator config if TaylorSeer is enabled + calibrator_config = None + if cache_config.enable_taylorseer: + taylorseer_order = cache_config.taylorseer_order + calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) + logger.info(f"TaylorSeer enabled with order={taylorseer_order}") + + logger.info( + f"Enabling cache-dit on transformer: " + f"Fn={db_cache_config.Fn_compute_blocks}, " + f"Bn={db_cache_config.Bn_compute_blocks}, " + f"W={db_cache_config.max_warmup_steps}, " + ) + + # Enable cache-dit on the transformer + cache_dit.enable_cache( + pipeline.transformer, + cache_config=db_cache_config, + calibrator_config=calibrator_config, + ) + + def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: + """Refresh cache context for the transformer with new num_inference_steps. + + Args: + pipeline: The diffusion pipeline instance. + num_inference_steps: New number of inference steps. + """ + if cache_config.scm_steps_mask_policy is None: + cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose) + else: + cache_dit.refresh_context( + pipeline.transformer, + cache_config=DBCacheConfig().reset( + num_inference_steps=num_inference_steps, + steps_computation_mask=cache_dit.steps_mask( + mask_policy=cache_config.scm_steps_mask_policy, + total_steps=num_inference_steps, + ), + steps_computation_policy=cache_config.scm_steps_policy, + ), + verbose=verbose, + ) + + return refresh_cache_context + + +# Register custom cache-dit enablers after function definitions +CUSTOM_DIT_ENABLERS.update( + { + "WanPipeline": enable_cache_for_wan22, + "FluxPipeline": enable_cache_for_flux, + } +) + + +class CacheDiTBackend(CacheBackend): + """Backend class for cache-dit acceleration on diffusion pipelines. + + This class implements cache-dit acceleration (DBCache, SCM, TaylorSeer) using + the cache-dit library. It inherits from CacheBackend and provides a unified + interface for managing cache-dit acceleration on diffusion models. + + Attributes: + config: Cache configuration (DiffusionCacheConfig instance), inherited from CacheBackend. + enabled: Whether cache-dit is enabled on this pipeline, inherited from CacheBackend. + _refresh_func: Internal refresh function for updating cache context. + _last_num_inference_steps: Last num_inference_steps used for refresh optimization. + """ + + def __init__(self, cache_config: Any = None): + """Initialize the cache-dit backend. + + Args: + cache_config: Cache configuration (DiffusionCacheConfig instance, dict, or None). + If None or empty, uses default DiffusionCacheConfig(). + """ + # Use default config if cache_config is not provided or is empty + if cache_config is None: + config = DiffusionCacheConfig() + elif isinstance(cache_config, dict): + # Convert dict to DiffusionCacheConfig, using defaults for missing keys + config = DiffusionCacheConfig.from_dict(cache_config) + else: + config = cache_config + + # Initialize base class with normalized config + super().__init__(config) + + # Cache-dit specific attributes + self._refresh_func: Optional[Callable[[Any, int, bool], None]] = None + self._last_num_inference_steps: Optional[int] = None + + def enable(self, pipeline: Any) -> None: + """Enable cache-dit on the pipeline if configured. + + This method applies cache-dit acceleration to the appropriate transformer(s) + in the pipeline. It handles both single-transformer and dual-transformer + architectures (e.g., Wan2.2). + + Args: + pipeline: The diffusion pipeline instance. + """ + if not CACHE_DIT_AVAILABLE: + logger.warning("cache-dit is not available, skipping cache-dit setup.") + return + + # Extract pipeline name from pipeline + pipeline_name = pipeline.__class__.__name__ + # Check if this model has a custom cache-dit enabler + if pipeline_name in CUSTOM_DIT_ENABLERS: + logger.info(f"Using custom cache-dit enabler for model: {pipeline_name}") + self._refresh_func = CUSTOM_DIT_ENABLERS[pipeline_name](pipeline, self.config) + else: + # For regular single-transformer models + self._refresh_func = enable_cache_for_dit(pipeline, self.config) + + self.enabled = True + logger.info(f"Cache-dit enabled successfully on {pipeline_name}") + + def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: + """Refresh cache context with new num_inference_steps. + + This method updates the cache context when num_inference_steps changes + during inference. For dual-transformer models (e.g., Wan2.2), it automatically + splits the steps based on boundary_ratio. + + Args: + pipeline: The diffusion pipeline instance. + num_inference_steps: New number of inference steps. + verbose: Whether to log refresh operations. + """ + if not self.enabled or self._refresh_func is None: + logger.warning("Cache-dit is not enabled. Cannot refresh cache context.") + return + + # Only refresh if num_inference_steps has changed + if self._last_num_inference_steps is None or num_inference_steps != self._last_num_inference_steps: + if verbose: + logger.info(f"Refreshing cache context for transformer with num_inference_steps: {num_inference_steps}") + self._refresh_func(pipeline, num_inference_steps, verbose) + self._last_num_inference_steps = num_inference_steps + + def is_enabled(self) -> bool: + """Check if cache-dit is enabled on this pipeline. + + Returns: + True if cache-dit is enabled, False otherwise. + """ + return self.enabled + + +def may_enable_cache_dit(pipeline: Any, od_config: OmniDiffusionConfig) -> Optional["CacheDiTBackend"]: + """Enable cache-dit on the pipeline if configured (convenience function). + + This is a convenience function that creates and enables a CacheDiTBackend. + For new code, consider using CacheDiTBackend directly. + + Args: + pipeline: The diffusion pipeline instance. + od_config: OmniDiffusionConfig with cache configuration. + + Returns: + A CacheDiTBackend instance if cache-dit is enabled, None otherwise. + """ + if od_config.cache_backend != "cache-dit" or not od_config.cache_config: + return None + + backend = CacheDiTBackend(od_config.cache_config) + backend.enable(pipeline) + return backend if backend.is_enabled() else None diff --git a/vllm_omni/diffusion/cache/registry.py b/vllm_omni/diffusion/cache/registry.py deleted file mode 100644 index 168e07dc5..000000000 --- a/vllm_omni/diffusion/cache/registry.py +++ /dev/null @@ -1,100 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -""" -Cache adapter registry for diffusion models. - -This module provides a registry pattern for cache adapters, allowing dynamic -registration and instantiation of different cache types (TeaCache, DeepCache, etc.). -""" - -from enum import Enum -from typing import Any - -from vllm.logger import init_logger - -from vllm_omni.diffusion.cache.base import CacheAdapter - -logger = init_logger(__name__) - - -class CacheType(Enum): - """Supported cache adapter types.""" - - NONE = "none" - TEA_CACHE = "tea_cache" - # Future cache types can be added here: - # DEEP_CACHE = "deep_cache" - # DISTRI_FUSION = "distri_fusion" - - -# Global registry mapping cache types to adapter classes -CACHE_ADAPTER_REGISTRY: dict[CacheType, type[CacheAdapter]] = {} - - -def register_cache_adapter(cache_type: CacheType, adapter_class: type[CacheAdapter]) -> None: - """ - Register a cache adapter class for a given cache type. - - This allows extending the system with new cache types without modifying - the core cache infrastructure. - - Args: - cache_type: CacheType enum value - adapter_class: CacheAdapter subclass to register - - Example: - >>> register_cache_adapter(CacheType.TEA_CACHE, TeaCacheAdapter) - """ - if not issubclass(adapter_class, CacheAdapter): - raise TypeError(f"{adapter_class} must be a subclass of CacheAdapter") - - CACHE_ADAPTER_REGISTRY[cache_type] = adapter_class - logger.debug(f"Registered cache adapter: {cache_type.value} -> {adapter_class.__name__}") - - -def get_cache_adapter(cache_type: str, config: dict[str, Any]) -> CacheAdapter: - """ - Factory function to get cache adapter instance. - - Converts string cache type to enum, looks up in registry, and instantiates - the appropriate adapter class with the provided configuration. - - Args: - cache_type: String name of cache type ("tea_cache", "deep_cache", etc.) - config: Configuration dictionary to pass to adapter constructor - - Returns: - Instantiated CacheAdapter subclass - - Raises: - ValueError: If cache_type is unknown or not registered - - Example: - >>> adapter = get_cache_adapter("tea_cache", {"rel_l1_thresh": 0.2}) - >>> adapter.apply(transformer) - """ - # Normalize cache type string - cache_type_str = cache_type.lower().strip() - - # Convert string to enum - try: - cache_enum = CacheType(cache_type_str) - except ValueError: - available = [ct.value for ct in CacheType if ct != CacheType.NONE] - raise ValueError(f"Unknown cache type: '{cache_type}'. Available types: {available}") - - # Check if it's the special "none" case - if cache_enum == CacheType.NONE: - raise ValueError("Cannot instantiate adapter for cache_type='none'. Use setup_cache() which handles this case.") - - # Lookup in registry - if cache_enum not in CACHE_ADAPTER_REGISTRY: - raise ValueError( - f"Cache type '{cache_type}' is not registered. Registered types: {list(CACHE_ADAPTER_REGISTRY.keys())}" - ) - - adapter_class = CACHE_ADAPTER_REGISTRY[cache_enum] - - # Instantiate and return - return adapter_class(config) diff --git a/vllm_omni/diffusion/cache/selector.py b/vllm_omni/diffusion/cache/selector.py new file mode 100644 index 000000000..7c09bf664 --- /dev/null +++ b/vllm_omni/diffusion/cache/selector.py @@ -0,0 +1,38 @@ +from typing import Any + +from vllm_omni.diffusion.cache.base import CacheBackend +from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTBackend +from vllm_omni.diffusion.cache.teacache.backend import TeaCacheBackend +from vllm_omni.diffusion.data import DiffusionCacheConfig + + +def get_cache_backend(cache_backend: str | None, cache_config: Any) -> CacheBackend | None: + """Get cache backend instance based on cache_backend string. + + This is a selector function that routes to the appropriate backend implementation. + - cache_dit: Uses CacheDiTBackend with enable()/refresh() interface + - tea_cache: Uses TeaCacheBackend with enable()/refresh() interface + + Args: + cache_backend: Cache backend name ("cache_dit", "tea_cache", or None). + cache_config: Cache configuration (dict or DiffusionCacheConfig instance). + + Returns: + Cache backend instance (CacheDiTBackend or TeaCacheBackend) if cache_backend is set, + None otherwise. + + Raises: + ValueError: If cache_backend is unsupported. + """ + if cache_backend is None or cache_backend == "none": + return None + + if isinstance(cache_config, dict): + cache_config = DiffusionCacheConfig.from_dict(cache_config) + + if cache_backend == "cache_dit": + return CacheDiTBackend(cache_config) + elif cache_backend == "tea_cache": + return TeaCacheBackend(cache_config) + else: + raise ValueError(f"Unsupported cache backend: {cache_backend}. Supported: 'cache_dit', 'tea_cache'") diff --git a/vllm_omni/diffusion/cache/teacache/__init__.py b/vllm_omni/diffusion/cache/teacache/__init__.py index ba57ebdb5..4f4ba5c6c 100644 --- a/vllm_omni/diffusion/cache/teacache/__init__.py +++ b/vllm_omni/diffusion/cache/teacache/__init__.py @@ -16,16 +16,16 @@ Usage: omni = Omni( model="Qwen/Qwen-Image", - cache_adapter="tea_cache", + cache_backend="tea_cache", cache_config={"rel_l1_thresh": 0.2} ) images = omni.generate("a cat") # Alternative: Using environment variable - # export DIFFUSION_CACHE_ADAPTER=tea_cache + # export DIFFUSION_CACHE_BACKEND=tea_cache """ -from vllm_omni.diffusion.cache.teacache.adapter import TeaCacheAdapter +from vllm_omni.diffusion.cache.teacache.backend import TeaCacheBackend from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig from vllm_omni.diffusion.cache.teacache.extractors import ( CacheContext, @@ -35,7 +35,7 @@ from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook, apply_teacache from vllm_omni.diffusion.cache.teacache.state import TeaCacheState __all__ = [ - "TeaCacheAdapter", + "TeaCacheBackend", "TeaCacheConfig", "TeaCacheState", "TeaCacheHook", diff --git a/vllm_omni/diffusion/cache/teacache/adapter.py b/vllm_omni/diffusion/cache/teacache/adapter.py deleted file mode 100644 index f94774842..000000000 --- a/vllm_omni/diffusion/cache/teacache/adapter.py +++ /dev/null @@ -1,96 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -""" -TeaCache adapter implementation. - -This module provides the TeaCache adapter that implements the CacheAdapter -interface using the hooks-based TeaCache system. -""" - -import torch -from vllm.logger import init_logger - -from vllm_omni.diffusion.cache.base import CacheAdapter -from vllm_omni.diffusion.cache.registry import CacheType, register_cache_adapter -from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig -from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook, apply_teacache_hook - -logger = init_logger(__name__) - - -class TeaCacheAdapter(CacheAdapter): - """ - TeaCache implementation using hooks. - - TeaCache (Timestep Embedding Aware Cache) is an adaptive caching technique - that speeds up diffusion inference by reusing transformer block computations - when consecutive timestep embeddings are similar. - - The adapter applies TeaCache hooks to the transformer which intercept the - forward pass and implement the caching logic transparently. - - Example: - >>> adapter = TeaCacheAdapter({"rel_l1_thresh": 0.2}) - >>> adapter.apply(transformer) - >>> # Generate with cache enabled - >>> adapter.reset(transformer) # Reset before each generation - """ - - def apply(self, transformer: torch.nn.Module) -> None: - """ - Apply TeaCache to transformer using hooks. - - This creates a TeaCacheConfig from the adapter's config dict and - applies the TeaCache hook to the transformer. - - Args: - transformer: Transformer module to apply TeaCache to - """ - # Validate that model_type is provided - if "model_type" not in self.config: - raise ValueError( - "model_type must be provided in cache_config. " - "Should be set to OmniDiffusionConfig.model_class_name (e.g., 'QwenImagePipeline')" - ) - - # Create TeaCacheConfig from dict - try: - teacache_config = TeaCacheConfig(**self.config) - except Exception as e: - logger.error(f"Failed to create TeaCacheConfig: {e}") - raise ValueError( - f"Invalid TeaCache configuration: {e}. " - f"Expected keys: rel_l1_thresh, model_type (pipeline class name), coefficients (optional)" - ) - - # Apply hook to transformer - apply_teacache_hook(transformer, teacache_config) - logger.info( - f"TeaCache applied with rel_l1_thresh={teacache_config.rel_l1_thresh}, " - f"model_type={teacache_config.model_type}" - ) - - def reset(self, transformer: torch.nn.Module) -> None: - """ - Reset TeaCache state for new generation. - - Clears all cached residuals and resets counters/accumulators. - Should be called before each generation to ensure clean state. - - Args: - transformer: Transformer module to reset cache on - """ - if hasattr(transformer, "_hook_registry"): - hook = transformer._hook_registry.get_hook(TeaCacheHook._HOOK_NAME) - if hook is not None: - transformer._hook_registry.reset_hook(TeaCacheHook._HOOK_NAME) - logger.debug("TeaCache state reset") - else: - logger.warning("TeaCache hook not found, nothing to reset") - else: - logger.warning("Transformer has no hook registry, TeaCache may not be applied") - - -# Register TeaCache adapter in the global registry -register_cache_adapter(CacheType.TEA_CACHE, TeaCacheAdapter) diff --git a/vllm_omni/diffusion/cache/teacache/backend.py b/vllm_omni/diffusion/cache/teacache/backend.py new file mode 100644 index 000000000..71f9ad653 --- /dev/null +++ b/vllm_omni/diffusion/cache/teacache/backend.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +TeaCache backend implementation. + +This module provides the TeaCache backend that implements the CacheBackend +interface using the hooks-based TeaCache system. +""" + +from typing import Any + +from vllm.logger import init_logger + +from vllm_omni.diffusion.cache.base import CacheBackend +from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig +from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook, apply_teacache_hook + +logger = init_logger(__name__) + + +class TeaCacheBackend(CacheBackend): + """ + TeaCache implementation using hooks. + + TeaCache (Timestep Embedding Aware Cache) is an adaptive caching technique + that speeds up diffusion inference by reusing transformer block computations + when consecutive timestep embeddings are similar. + + The backend applies TeaCache hooks to the transformer which intercept the + forward pass and implement the caching logic transparently. + + Example: + >>> from vllm_omni.diffusion.data import DiffusionCacheConfig + >>> backend = TeaCacheBackend(DiffusionCacheConfig(rel_l1_thresh=0.2)) + >>> backend.enable(pipeline) + >>> # Generate with cache enabled + >>> backend.refresh(pipeline, num_inference_steps=50) # Refresh before each generation + >>> # Access config attributes: backend.config.rel_l1_thresh + """ + + def enable(self, pipeline: Any) -> None: + """ + Enable TeaCache on transformer using hooks. + + This creates a TeaCacheConfig from the backend's DiffusionCacheConfig + and applies the TeaCache hook to the transformer. + + Args: + pipeline: Diffusion pipeline instance. Extracts transformer and model_type: + - transformer: pipeline.transformer + - model_type: pipeline.__class__.__name__ + """ + # Extract transformer and model_type from pipeline + transformer = pipeline.transformer + model_type = pipeline.__class__.__name__ + + # Create TeaCacheConfig from DiffusionCacheConfig with model_type + # Access parameters via attribute access: config.rel_l1_thresh + # rel_l1_thresh already has a default value of 0.2 in DiffusionCacheConfig + try: + teacache_config = TeaCacheConfig( + model_type=model_type, + rel_l1_thresh=self.config.rel_l1_thresh, + coefficients=self.config.coefficients, + ) + except Exception as e: + logger.error(f"Failed to create TeaCacheConfig: {e}") + raise ValueError( + f"Invalid TeaCache configuration: {e}. " + f"Expected keys: rel_l1_thresh, coefficients (optional). " + f"model_type is automatically extracted from pipeline.__class__.__name__." + ) + + # Apply hook to transformer + apply_teacache_hook(transformer, teacache_config) + + # Store backend reference on pipeline for pipeline's internal use + # (e.g., checking if cache is enabled for cache_branch parameter) + pipeline._cache_backend = self + + # Mark as enabled + self.enabled = True + + logger.info( + f"TeaCache applied with rel_l1_thresh={teacache_config.rel_l1_thresh}, " + f"model_type={teacache_config.model_type}" + ) + + def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: + """ + Refresh TeaCache state for new generation. + + Clears all cached residuals and resets counters/accumulators. + Should be called before each generation to ensure clean state. + + Args: + pipeline: Diffusion pipeline instance. Extracts transformer via pipeline.transformer. + num_inference_steps: Number of inference steps for the current generation. + Currently not used by TeaCache but accepted for interface consistency. + verbose: Whether to log refresh operations (default: True) + """ + # Extract transformer from pipeline + transformer = pipeline.transformer + + if hasattr(transformer, "_hook_registry"): + hook = transformer._hook_registry.get_hook(TeaCacheHook._HOOK_NAME) + if hook is not None: + transformer._hook_registry.reset_hook(TeaCacheHook._HOOK_NAME) + if verbose: + logger.debug(f"TeaCache state refreshed (num_inference_steps={num_inference_steps})") + else: + if verbose: + logger.warning("TeaCache hook not found, nothing to refresh") + else: + if verbose: + logger.warning("Transformer has no hook registry, TeaCache may not be applied") diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 41c6d2ae9..ca315935a 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -4,7 +4,7 @@ import enum import os import random -from dataclasses import dataclass, field +from dataclasses import dataclass, field, fields from typing import Any, Callable import torch @@ -41,6 +41,122 @@ class TransformerConfig: raise AttributeError(item) from exc +@dataclass +class DiffusionCacheConfig: + """ + Configuration for cache adapters (TeaCache, cache-dit, etc.). + + This dataclass provides a unified interface for cache configuration parameters. + It can be initialized from a dictionary and accessed via attributes. + + Common parameters: + - TeaCache: rel_l1_thresh, coefficients (optional) + - cache-dit: Fn_compute_blocks, Bn_compute_blocks, max_warmup_steps, + residual_diff_threshold, enable_taylorseer, taylorseer_order, + scm_steps_mask_policy, scm_steps_policy + + Example: + >>> # From dict (user-facing API) - partial config uses defaults for missing keys + >>> config = DiffusionCacheConfig.from_dict({"rel_l1_thresh": 0.3}) + >>> # Access via attribute + >>> print(config.rel_l1_thresh) # 0.3 (from dict) + >>> print(config.Fn_compute_blocks) # 8 (default) + >>> # Empty dict uses all defaults + >>> default_config = DiffusionCacheConfig.from_dict({}) + >>> print(default_config.rel_l1_thresh) # 0.2 (default) + """ + + # TeaCache parameters [tea_cache only] + # Default: 0.2 provides ~1.5x speedup with minimal quality loss (optimal balance) + rel_l1_thresh: float = 0.2 + coefficients: list[float] | None = None # Uses model-specific defaults if None + + # cache-dit parameters [cache-dit only] + # Default: 1 forward compute block (optimized for single-transformer models) + # Use 1 as default instead of cache-dit's 8, optimized for single-transformer models + # This provides better performance while maintaining quality for most use cases + Fn_compute_blocks: int = 1 + # Default: 0 backward compute blocks (no fusion by default) + Bn_compute_blocks: int = 0 + # Default: 4 warmup steps (optimized for few-step distilled models like Z-Image with 8 steps) + # Use 4 as default warmup steps instead of 8 in cache-dit, making DBCache work + # for few-step distilled models (e.g., Z-Image with 8 steps) + max_warmup_steps: int = 4 + # Default: -1 (unlimited cached steps) - DBCache disables caching when previous cached steps exceed this value + # to prevent precision degradation. Set to -1 for unlimited caching (cache-dit default). + max_cached_steps: int = -1 + # Default: 0.24 residual difference threshold (higher for more aggressive caching) + # Use a relatively higher residual diff threshold (0.24) as default to allow more + # aggressive caching. This is safe because we have max_continuous_cached_steps limit. + # Without this limit, a lower threshold like 0.12 would be needed. + residual_diff_threshold: float = 0.24 + # Default: Limit consecutive cached steps to 3 to prevent precision degradation + # This allows us to use a higher residual_diff_threshold for more aggressive caching + max_continuous_cached_steps: int = 3 + # Default: Disable TaylorSeer (not suitable for few-step distilled models) + # TaylorSeer is not suitable for few-step distilled models, so we disable it by default. + # References: + # - From Reusing to Forecasting: Accelerating Diffusion Models with TaylorSeers + # - Forecast then Calibrate: Feature Caching as ODE for Efficient Diffusion Transformers + enable_taylorseer: bool = False + # Default: 1st order TaylorSeer polynomial + taylorseer_order: int = 1 + # Default: "fast" SCM mask policy for good speed/quality balance + scm_steps_mask_policy: str = "fast" + # Default: "dynamic" steps policy for adaptive caching + scm_steps_policy: str = "dynamic" + # Used by cache-dit for scm mask generation. If this value changes during inference, + # we will re-generate the scm mask and refresh the cache context. + num_inference_steps: int | None = None + + # Additional parameters that may be passed but not explicitly defined + _extra_params: dict[str, Any] = field(default_factory=dict, repr=False) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "DiffusionCacheConfig": + """ + Create DiffusionCacheConfig from a dictionary. + + Args: + data: Dictionary containing cache configuration parameters + + Returns: + DiffusionCacheConfig instance with parameters set from dict + """ + if not isinstance(data, dict): + raise TypeError(f"Expected cache config dict, got {type(data)!r}") + + # Get all dataclass field names automatically + field_names = {f.name for f in fields(cls)} + + # Extract parameters that match dataclass fields (excluding private fields) + known_params = {k: v for k, v in data.items() if k in field_names and not k.startswith("_")} + + # Store extra parameters + extra_params = {k: v for k, v in data.items() if k not in field_names} + + # Create instance with known params (missing ones will use defaults) + # Then update _extra_params after creation since it's a private field + instance = cls(**known_params, _extra_params=extra_params) + return instance + + def __getattr__(self, item: str) -> Any: + """ + Allow access to extra parameters via attribute access. + + This enables accessing parameters that weren't explicitly defined + in the dataclass fields but were passed in the dict. + """ + if item == "_extra_params" or item.startswith("_"): + return object.__getattribute__(self, item) + + extra = object.__getattribute__(self, "_extra_params") + if item in extra: + return extra[item] + + raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{item}'") + + @dataclass class OmniDiffusionConfig: # Model and path configuration (for convenience) @@ -64,9 +180,9 @@ class OmniDiffusionConfig: # Cache strategy (legacy) cache_strategy: str = "none" - # Cache adapter configuration (NEW) - cache_adapter: str = "none" # "tea_cache", "deep_cache", etc. - cache_config: dict[str, Any] = field(default_factory=dict) + # Cache backend configuration (NEW) + cache_backend: str = "none" # "tea_cache", "deep_cache", etc. + cache_config: DiffusionCacheConfig | dict[str, Any] = field(default_factory=dict) # Distributed executor backend distributed_executor_backend: str = "mp" @@ -213,17 +329,20 @@ class OmniDiffusionConfig: initial_master_port = (self.master_port or 30005) + random.randint(0, 100) self.master_port = self.settle_port(initial_master_port, 37) - # Automatically inject model_class_name into cache_config if not present - if self.cache_adapter != "none" and self.model_class_name: - if "model_type" not in self.cache_config: - self.cache_config["model_type"] = self.model_class_name - logger.debug(f"Auto-injected model_type='{self.model_class_name}' into cache_config") + # Convert cache_config dict to DiffusionCacheConfig if needed + if isinstance(self.cache_config, dict): + self.cache_config = DiffusionCacheConfig.from_dict(self.cache_config) + elif not isinstance(self.cache_config, DiffusionCacheConfig): + # If it's neither dict nor DiffusionCacheConfig, convert to empty config + self.cache_config = DiffusionCacheConfig() @classmethod def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig": - # Check environment variable as fallback for cache_adapter - if "cache_adapter" not in kwargs: - kwargs["cache_adapter"] = os.environ.get("DIFFUSION_CACHE_ADAPTER", "none").lower() + # Check environment variable as fallback for cache_backend + # Support both old DIFFUSION_CACHE_ADAPTER and new DIFFUSION_CACHE_BACKEND for backwards compatibility + if "cache_backend" not in kwargs: + cache_backend = os.environ.get("DIFFUSION_CACHE_BACKEND") or os.environ.get("DIFFUSION_CACHE_ADAPTER") + kwargs["cache_backend"] = cache_backend.lower() if cache_backend else "none" return cls(**kwargs) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index 2a5678c7f..84967c6cf 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py @@ -268,10 +268,11 @@ class QwenImagePipeline( self.device ) self.transformer = QwenImageTransformer2DModel(od_config=od_config) + self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) - # Initialize cache adapter to None (will be set by worker or setup_cache if needed) - self._cache_adapter = None + # Initialize cache backend to None (will be set by worker if needed) + self._cache_backend = None self.stage = None @@ -566,7 +567,7 @@ class QwenImagePipeline( "attention_kwargs": self.attention_kwargs, "return_dict": False, } - if self._cache_adapter is not None: + if self._cache_backend is not None: transformer_kwargs["cache_branch"] = "positive" noise_pred = self.transformer(**transformer_kwargs)[0] @@ -585,7 +586,7 @@ class QwenImagePipeline( "attention_kwargs": self.attention_kwargs, "return_dict": False, } - if self._cache_adapter is not None: + if self._cache_backend is not None: neg_transformer_kwargs["cache_branch"] = "negative" neg_noise_pred = self.transformer(**neg_transformer_kwargs)[0] diff --git a/vllm_omni/diffusion/worker/gpu_worker.py b/vllm_omni/diffusion/worker/gpu_worker.py index 147629d1f..868611841 100644 --- a/vllm_omni/diffusion/worker/gpu_worker.py +++ b/vllm_omni/diffusion/worker/gpu_worker.py @@ -15,6 +15,7 @@ from vllm.distributed.parallel_state import ( from vllm.logger import init_logger from vllm.utils import DeviceMemoryProfiler, GiB_bytes +from vllm_omni.diffusion.cache.selector import get_cache_backend from vllm_omni.diffusion.data import ( SHUTDOWN_MESSAGE, DiffusionOutput, @@ -84,21 +85,11 @@ class GPUWorker: ) logger.info(f"Worker {self.rank}: Model loaded successfully.") - # Apply cache adapter (model_type is auto-injected in OmniDiffusionConfig.__post_init__) - from vllm_omni.diffusion.cache.apply import setup_cache + # Setup cache backend based on type (both backends use enable()/reset() interface) + self.cache_backend = get_cache_backend(self.od_config.cache_backend, self.od_config.cache_config) - self.od_config.cache_config["model_type"] = self.od_config.model_class_name - - self.pipeline._cache_adapter = setup_cache( - self.pipeline.transformer, - cache_type=self.od_config.cache_adapter, - cache_config=self.od_config.cache_config, - ) - - def maybe_reset_cache(self) -> None: - """Reset cache state before each generation if applicable.""" - if self.pipeline._cache_adapter is not None: - self.pipeline._cache_adapter.reset(self.pipeline.transformer) + if self.cache_backend is not None: + self.cache_backend.enable(self.pipeline) @torch.inference_mode() def execute_model(self, reqs: list[OmniDiffusionRequest], od_config: OmniDiffusionConfig) -> DiffusionOutput: @@ -106,9 +97,13 @@ class GPUWorker: Execute a forward pass. """ assert self.pipeline is not None - self.maybe_reset_cache() # TODO: dealing with first req for now req = reqs[0] + + # Refresh cache context if needed + if self.cache_backend is not None and self.cache_backend.is_enabled(): + self.cache_backend.refresh(self.pipeline, req.num_inference_steps) + output = self.pipeline.forward(req) return output