[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>

.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

.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/

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:

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)

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

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)

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)

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
 ]
 

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).
+"""

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", {})

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

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

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",
 ]

vllm_omni/diffusion/cache/apply.py

@@ -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

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})"

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

vllm_omni/diffusion/cache/registry.py

@@ -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)

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'")

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",

vllm_omni/diffusion/cache/teacache/adapter.py

@@ -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)

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")

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)
 
 

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]

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