cache-dit 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

cache_dit/cache_factory/dynamic_block_prune/prune_context.py

@@ -55,6 +55,7 @@ class DBPPruneContext:
         default_factory=lambda: defaultdict(list),
     )
 
+    @torch.compiler.disable
     def get_residual_diff_threshold(self):
         residual_diff_threshold = self.residual_diff_threshold
         if self.l1_hidden_states_diff_threshold is not None:
@@ -98,19 +99,24 @@ class DBPPruneContext:
                         )
         return residual_diff_threshold
 
+    @torch.compiler.disable
     def get_buffer(self, name):
         return self.buffers.get(name)
 
+    @torch.compiler.disable
     def set_buffer(self, name, buffer):
         self.buffers[name] = buffer
 
+    @torch.compiler.disable
     def remove_buffer(self, name):
         if name in self.buffers:
             del self.buffers[name]
 
+    @torch.compiler.disable
     def clear_buffers(self):
         self.buffers.clear()
 
+    @torch.compiler.disable
     def mark_step_begin(self):
         self.executed_steps += 1
         if self.get_current_step() == 0:
@@ -118,12 +124,15 @@ class DBPPruneContext:
             self.actual_blocks.clear()
             self.residual_diffs.clear()
 
+    @torch.compiler.disable
     def add_pruned_block(self, num_blocks):
         self.pruned_blocks.append(num_blocks)
 
+    @torch.compiler.disable
     def add_actual_block(self, num_blocks):
         self.actual_blocks.append(num_blocks)
 
+    @torch.compiler.disable
     def add_residual_diff(self, diff):
         if isinstance(diff, torch.Tensor):
             diff = diff.item()
@@ -141,9 +150,11 @@ class DBPPruneContext:
                 f"residual diff: {diff:.6f}"
             )
 
+    @torch.compiler.disable
     def get_current_step(self):
         return self.executed_steps - 1
 
+    @torch.compiler.disable
     def is_in_warmup(self):
         return self.get_current_step() < self.warmup_steps
 
@@ -348,11 +359,19 @@ def collect_prune_kwargs(default_attrs: dict, **kwargs):
         )
         for attr in prune_attrs
     }
+
     # Manually set sequence fields, such as non_prune_blocks_ids
-    prune_kwargs["non_prune_blocks_ids"] = kwargs.pop(
-        "non_prune_blocks_ids",
-        [],
-    )
+    def _safe_set_sequence_field(
+        field_name: str,
+        default_value: Any = None,
+    ):
+        if field_name not in prune_kwargs:
+            prune_kwargs[field_name] = kwargs.pop(
+                field_name,
+                default_value,
+            )
+
+    _safe_set_sequence_field("non_prune_blocks_ids", [])
 
     assert default_attrs is not None, "default_attrs must be set before"
     for attr in prune_attrs:
@@ -627,10 +646,6 @@ class DBPrunedTransformerBlocks(torch.nn.Module):
         ]
         return sorted(non_prune_blocks_ids)
 
-    # @torch.compile(dynamic=True)
-    # mark this function as compile with dynamic=True will
-    # cause precision degradate, so, we choose to disable it
-    # now, until we find a better solution or fixed the bug.
     @torch.compiler.disable
     def _compute_single_hidden_states_residual(
         self,
@@ -667,10 +682,6 @@ class DBPrunedTransformerBlocks(torch.nn.Module):
             single_encoder_hidden_states_residual,
         )
 
-    # @torch.compile(dynamic=True)
-    # mark this function as compile with dynamic=True will
-    # cause precision degradate, so, we choose to disable it
-    # now, until we find a better solution or fixed the bug.
     @torch.compiler.disable
     def _split_single_hidden_states(
         self,
@@ -969,17 +980,6 @@ def patch_pruned_stats(
     if transformer is None:
         return
 
-    pruned_transformer_blocks = getattr(transformer, "transformer_blocks", None)
-    if pruned_transformer_blocks is None:
-        return
-
-    if isinstance(pruned_transformer_blocks, torch.nn.ModuleList):
-        pruned_transformer_blocks = pruned_transformer_blocks[0]
-    if not isinstance(
-        pruned_transformer_blocks, DBPrunedTransformerBlocks
-    ) or not isinstance(transformer, torch.nn.Module):
-        return
-
     # TODO: Patch more pruned stats to the transformer
     transformer._pruned_blocks = get_pruned_blocks()
     transformer._pruned_steps = get_pruned_steps()

cache_dit/cache_factory/first_block_cache/cache_context.py

@@ -370,6 +370,9 @@ def apply_prev_hidden_states_residual(
         hidden_states = hidden_states_residual + hidden_states
 
         hidden_states = hidden_states.contiguous()
+        # NOTE: We should also support taylorseer for 
+        # encoder_hidden_states approximation. Please 
+        # use DBCache instead.
     else:
         hidden_states_residual = get_hidden_states_residual()
         assert (
@@ -711,17 +714,6 @@ def patch_cached_stats(
     if transformer is None:
         return
 
-    cached_transformer_blocks = getattr(transformer, "transformer_blocks", None)
-    if cached_transformer_blocks is None:
-        return
-
-    if isinstance(cached_transformer_blocks, torch.nn.ModuleList):
-        cached_transformer_blocks = cached_transformer_blocks[0]
-    if not isinstance(
-        cached_transformer_blocks, CachedTransformerBlocks
-    ) or not isinstance(transformer, torch.nn.Module):
-        return
-
     # TODO: Patch more cached stats to the transformer
     transformer._cached_steps = get_cached_steps()
     transformer._residual_diffs = get_residual_diffs()

cache_dit/cache_factory/taylorseer.py

@@ -1,5 +1,7 @@
 # Adapted from: https://github.com/chengzeyi/ParaAttention/blob/main/src/para_attn/first_block_cache/taylorseer.py
+# Reference: https://github.com/Shenyi-Z/TaylorSeer/TaylorSeer-FLUX/src/flux/taylor_utils/__init__.py
 import math
+import torch
 
 
 class TaylorSeer:
@@ -17,6 +19,7 @@ class TaylorSeer:
         self.compute_step_map = compute_step_map
         self.reset_cache()
 
+    @torch.compiler.disable
     def reset_cache(self):
         self.state = {
             "dY_prev": [None] * self.ORDER,
@@ -25,6 +28,7 @@ class TaylorSeer:
         self.current_step = -1
         self.last_non_approximated_step = -1
 
+    @torch.compiler.disable
     def should_compute_full(self, step=None):
         step = self.current_step if step is None else step
         if self.compute_step_map is not None:
@@ -36,7 +40,13 @@ class TaylorSeer:
             return True
         return False
 
+    @torch.compiler.disable
     def approximate_derivative(self, Y):
+        # n-th order Taylor expansion:
+        # Y(t) = Y(0) + dY(0)/dt * t + d^2Y(0)/dt^2 * t^2 / 2!
+        #        + ... + d^nY(0)/dt^n * t^n / n!
+        # TODO: Custom Triton/CUDA kernel for better performance,
+        # especially for large n_derivatives.
         dY_current = [None] * self.ORDER
         dY_current[0] = Y
         window = self.current_step - self.last_non_approximated_step
@@ -49,7 +59,10 @@ class TaylorSeer:
                 break
         return dY_current
 
+    @torch.compiler.disable
     def approximate_value(self):
+        # TODO: Custom Triton/CUDA kernel for better performance,
+        # especially for large n_derivatives.
         elapsed = self.current_step - self.last_non_approximated_step
         output = 0
         for i, derivative in enumerate(self.state["dY_current"]):
@@ -59,14 +72,30 @@ class TaylorSeer:
                 break
         return output
 
+    @torch.compiler.disable
     def mark_step_begin(self):
         self.current_step += 1
 
+    @torch.compiler.disable
     def update(self, Y):
+        # Directly call this method will ingnore the warmup
+        # policy and force full computation.
+        # Assume warmup steps is 3, and n_derivatives is 3.
+        # step 0: dY_prev    = [None, None,   None,    None   ]
+        #         dY_current = [Y0,   None,   None,    None   ]
+        # step 1: dY_prev    = [Y0,   None,   None,    None   ]
+        #         dY_current = [Y1,   dY1,    None,    None   ]
+        # step 2: dY_prev    = [Y1,   dY1,    None,    None   ]
+        #         dY_current = [Y2,   dY2/Y1, dY2/dY1, None   ]
+        # step 3: dY_prev    = [Y2,   dY2/Y1, dY2/dY1, None   ],
+        #         dY_current = [Y3,   dY3/Y2, dY3/dY2, dY3/dY1]
+        # step 4: dY_prev    = [Y3,   dY3/Y2, dY3/dY2, dY3/dY1]
+        #         dY_current = [Y4,   dY4/Y3, dY4/dY3, dY4/dY2]
         self.state["dY_prev"] = self.state["dY_current"]
         self.state["dY_current"] = self.approximate_derivative(Y)
         self.last_non_approximated_step = self.current_step
 
+    @torch.compiler.disable
     def step(self, Y):
         self.mark_step_begin()
         if self.should_compute_full():

cache_dit/compile/__init__.py

@@ -0,0 +1 @@
+from cache_dit.compile.utils import set_custom_compile_configs

cache_dit/compile/utils.py

@@ -0,0 +1,94 @@
+import os
+
+import torch
+from cache_dit.logger import init_logger, logging_rank_0
+
+logger = init_logger(__name__)
+
+
+def epilogue_prologue_fusion_enabled(**kwargs) -> bool:
+    mode = kwargs.get("epilogue_prologue_fusion", False)
+    CACHE_DIT_EPILOGUE_PROLOGUE_FUSION = bool(
+        int(os.environ.get("CACHE_DIT_EPILOGUE_PROLOGUE_FUSION", "0"))
+    )
+
+    if CACHE_DIT_EPILOGUE_PROLOGUE_FUSION:
+        logging_rank_0(
+            logger,
+            "CACHE_DIT_EPILOGUE_PROLOGUE_FUSION is set to 1. \n"
+            "Force enable epilogue and prologue fusion.",
+        )
+
+    return CACHE_DIT_EPILOGUE_PROLOGUE_FUSION or mode
+
+
+def set_custom_compile_configs(
+    cuda_graphs: bool = False,
+    force_disable_compile_caches: bool = False,
+    use_fast_math: bool = False,
+    **kwargs,  # other kwargs
+):
+    # Alway increase recompile_limit for dynamic shape compilation
+    torch._dynamo.config.recompile_limit = 96  # default is 8
+    torch._dynamo.config.accumulated_recompile_limit = 2048  # default is 256
+    # Handle compiler caches
+    # https://github.com/vllm-project/vllm/blob/23baa2180b0ebba5ae94073ba9b8e93f88b75486/vllm/compilation/compiler_interface.py#L270
+    torch._inductor.config.fx_graph_cache = True
+    torch._inductor.config.fx_graph_remote_cache = False
+    # https://github.com/pytorch/pytorch/issues/153791
+    torch._inductor.config.autotune_local_cache = False
+
+    FORCE_DISABLE_CUSTOM_COMPILE_CONFIG = (
+        os.environ.get("CACHE_DIT_FORCE_DISABLE_CUSTOM_COMPILE_CONFIG", "0")
+        == "1"
+    )
+    if FORCE_DISABLE_CUSTOM_COMPILE_CONFIG:
+        logging_rank_0(
+            logger,
+            "CACHE_DIT_FORCE_DISABLE_CUSTOM_COMPILE_CONFIG is set to 1. \n"
+            "Force disable custom compile config.",
+        )
+        return
+
+    # Enable compute comm overlap
+    torch._inductor.config.reorder_for_compute_comm_overlap = True
+    # L20 64 GB/s, PCIe; A100/A800 NVLink 300 GB/s.
+    torch._inductor.config.intra_node_bw = (
+        64 if "L20" in torch.cuda.get_device_name() else 300
+    )
+
+    # Below are default settings for torch.compile, you can change
+    # them to your needs and test the performance
+    torch._inductor.config.max_fusion_size = 64
+    torch._inductor.config.max_pointwise_cat_inputs = 8
+    torch._inductor.config.triton.cudagraphs = cuda_graphs
+    torch._inductor.config.triton.use_block_ptr = False
+    torch._inductor.config.triton.codegen_upcast_to_fp32 = True
+
+    # Copy from https://pytorch.org/blog/accelerating-generative-ai-3/
+    torch._inductor.config.conv_1x1_as_mm = True
+    torch._inductor.config.coordinate_descent_tuning = True
+    torch._inductor.config.coordinate_descent_check_all_directions = True
+    torch._inductor.config.epilogue_fusion = False
+
+    # Enable epilogue and prologue fusion
+    if epilogue_prologue_fusion_enabled(**kwargs):
+        torch._inductor.config.epilogue_fusion = True
+        torch._inductor.config.prologue_fusion = True
+        torch._inductor.config.epilogue_fusion_first = True
+
+    # Dead code elimination
+    torch._inductor.config.dce = True  # default is False
+
+    # May need to force disable all cache
+    if force_disable_compile_caches:
+        torch._inductor.config.force_disable_caches = True
+        torch._inductor.config.fx_graph_cache = False
+        torch._inductor.config.fx_graph_remote_cache = False
+        torch._inductor.config.autotune_local_cache = False  # default is True
+
+    # Use fast math
+    if hasattr(torch._inductor.config, "use_fast_math"):
+        torch._inductor.config.use_fast_math = use_fast_math
+    if hasattr(torch._inductor.config, "cuda.use_fast_math"):
+        torch._inductor.config.cuda.use_fast_math = use_fast_math

cache_dit/logger.py

@@ -1,6 +1,7 @@
 import logging
 import os
 import sys
+import torch.distributed as dist
 
 _FORMAT = "%(levelname)s %(asctime)s [%(filename)s:%(lineno)d] %(message)s"
 _DATE_FORMAT = "%m-%d %H:%M:%S"
@@ -95,3 +96,30 @@ def init_logger(name: str):
             logger.addHandler(_inference_log_file_handler[pid])
     logger.propagate = False
     return logger
+
+
+def logging_rank_0(
+    logger: logging.Logger, message: str, level: int = logging.INFO
+):
+    if not isinstance(logger, logging.Logger):
+        raise TypeError("logger must be an instance of logging.Logger")
+    if not isinstance(message, str):
+        raise TypeError("message must be a string")
+    if not isinstance(level, int):
+        raise TypeError("level must be an integer representing a logging level")
+
+    def _logging_msg():
+        if level == logging.DEBUG:
+            logger.debug(message)
+        elif level == logging.WARNING:
+            logger.warning(message)
+        elif level == logging.ERROR:
+            logger.error(message)
+        else:
+            logger.info(message)
+
+    if dist.is_initialized():
+        if dist.get_rank() == 0:
+            _logging_msg()
+    else:
+        _logging_msg()

{cache_dit-0.2.1.dist-info → cache_dit-0.2.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cache_dit
-Version: 0.2.1
+Version: 0.2.3
 Summary: 🤗 CacheDiT: A Training-free and Easy-to-use Cache Acceleration Toolbox for Diffusion Transformers
 Author: DefTruth, vipshop.com, etc.
 Maintainer: DefTruth, vipshop.com, etc
@@ -37,31 +37,31 @@ Dynamic: requires-python
   <p align="center">
     <h2>🤗 CacheDiT: A Training-free and Easy-to-use Cache Acceleration <br>Toolbox for Diffusion Transformers</h2>
   </p>
-  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/cache-dit.png >
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/cache-dit-v1.png >
   <div align='center'>
       <img src=https://img.shields.io/badge/Language-Python-brightgreen.svg >
       <img src=https://img.shields.io/badge/PRs-welcome-9cf.svg >
       <img src=https://img.shields.io/badge/PyPI-pass-brightgreen.svg >
       <img src=https://static.pepy.tech/badge/cache-dit >
       <img src=https://img.shields.io/badge/Python-3.10|3.11|3.12-9cf.svg >
-      <img src=https://img.shields.io/badge/Release-v0.2.1-brightgreen.svg >
+      <img src=https://img.shields.io/badge/Release-v0.2.2-brightgreen.svg >
  </div>
   <p align="center">
-    DeepCache is for UNet not DiT. Most DiT cache speedups are complex and not training-free. CacheDiT <br>offers a set of training-free cache accelerators for DiT: 🔥DBCache, DBPrune, FBCache, etc🔥
+    DeepCache is for UNet not DiT. Most DiT cache speedups are complex and not training-free. CacheDiT offers <br>a set of training-free cache accelerators for DiT: <b>🔥<a href="#dbcache">DBCache</a>, <a href="#dbprune">DBPrune</a>, <a href="#taylorseer">TaylorSeer</a>, <a href="#fbcache">FBCache</a></b>, etc🔥
   </p>
 </div>
 
-## 👋 Highlight
-
-<div id="reference"></div>
-
-The **CacheDiT** codebase is adapted from [FBCache](https://github.com/chengzeyi/ParaAttention/tree/main/src/para_attn/first_block_cache). Special thanks to their excellent work! The **FBCache** support for Mochi, FLUX.1, CogVideoX, Wan2.1, and HunyuanVideo is directly adapted from the original [FBCache](https://github.com/chengzeyi/ParaAttention/tree/main/src/para_attn/first_block_cache).
+<div align="center">
+  <p align="center">
+    <b>♥️ Please consider to leave a ⭐️ Star to support us ~ ♥️</b>
+  </p>
+</div> 
 
 ## 🤗 Introduction 
 
 <div align="center">
   <p align="center">
-    <h3>🔥 DBCache: Dual Block Caching for Diffusion Transformers</h3>
+    <h3>🔥DBCache: Dual Block Caching for Diffusion Transformers</h3>
   </p>
 </div> 
 
@@ -77,9 +77,9 @@ The **CacheDiT** codebase is adapted from [FBCache](https://github.com/chengzeyi
 |:---:|:---:|:---:|:---:|:---:|:---:|
 |24.85s|15.59s|8.58s|15.41s|15.11s|17.74s|
 |<img src=https://github.com/vipshop/cache-dit/raw/main/assets/NONE_R0.08_S0.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F1B0S1_R0.08_S11.png width=105px> | <img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F1B0S1_R0.2_S19.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F8B8S1_R0.15_S15.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F12B12S4_R0.2_S16.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F16B16S4_R0.2_S13.png width=105px>|
-|**Baseline(L20x1)**|**F1B0 (0.08)**|**F8B8 (0.12)**|**F8B12 (0.20)**|**F8B16 (0.20)**|**F8B20 (0.20)**|
+|**Baseline(L20x1)**|**F1B0 (0.08)**|**F8B8 (0.12)**|**F8B12 (0.12)**|**F8B16 (0.20)**|**F8B20 (0.20)**|
 |27.85s|6.04s|5.88s|5.77s|6.01s|6.20s|
-|<img src=https://github.com/user-attachments/assets/70ea57f4-d8f2-415b-8a96-d8315974a5e6 width=105px>|<img src=https://github.com/user-attachments/assets/fc0e1a67-19cc-44aa-bf50-04696e7978a0 width=105px> |<img src=https://github.com/user-attachments/assets/d1434896-628c-436b-95ad-43c085a8629e width=105px>|<img src=https://github.com/user-attachments/assets/aaa42cd2-57de-4c4e-8bfb-913018a8251d width=105px>|<img src=https://github.com/user-attachments/assets/dc0ba2a4-ef7c-436d-8a39-67055deab92f width=105px>|<img src=https://github.com/user-attachments/assets/aede466f-61ed-4256-8df0-fecf8020c5ca width=105px>|
+|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/TEXTURE_NONE_R0.08.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/TEXTURE_DBCACHE_F1B0_R0.08.png width=105px> |<img src=https://github.com/vipshop/cache-dit/raw/main/assets/TEXTURE_DBCACHE_F8B8_R0.12.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/TEXTURE_DBCACHE_F8B12_R0.12.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/TEXTURE_DBCACHE_F8B16_R0.2.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/TEXTURE_DBCACHE_F8B20_R0.2.png width=105px>|
 
 <div align="center">
   <p align="center">
@@ -91,7 +91,7 @@ These case studies demonstrate that even with relatively high thresholds (such a
 
 <div align="center">
   <p align="center">
-    <h3>🔥 DBPrune: Dynamic Block Prune with Residual Caching</h3>
+    <h3>🔥DBPrune: Dynamic Block Prune with Residual Caching</h3>
   </p>
 </div> 
 
@@ -110,11 +110,11 @@ These case studies demonstrate that even with relatively high thresholds (such a
 
 <div align="center">
   <p align="center">
-    <h3>🔥 Context Parallelism and Torch Compile</h3>
+    <h3>🔥Context Parallelism and Torch Compile</h3>
   </p>
 </div> 
 
-Moreover, **CacheDiT** are **plug-and-play** solutions that works hand-in-hand with [ParaAttention](https://github.com/chengzeyi/ParaAttention). Users can easily tap into its **Context Parallelism** features for distributed inference. By the way, CacheDiT is designed to work compatibly with **torch.compile.** You can easily use CacheDiT with torch.compile to further achieve a better performance.
+Moreover, **CacheDiT** are **plug-and-play** solutions that works hand-in-hand with [ParaAttention](https://github.com/chengzeyi/ParaAttention). Users can easily tap into its **Context Parallelism** features for distributed inference. CacheDiT is designed to work compatibly with **torch.compile.** You can easily use CacheDiT with torch.compile to further achieve a better performance.
 
 <div align="center">
   <p align="center">
@@ -128,12 +128,6 @@ Moreover, **CacheDiT** are **plug-and-play** solutions that works hand-in-hand w
 |+L20x4:7.75s|6.62s|6.03s|5.81s|5.24s|3.93s|
 |<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C1_NONE_R0.08_S0_T20.43s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C1_DBPRUNE_F1B0_R0.03_P24.0_T16.25s.png width=105px> | <img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C1_DBPRUNE_F1B0_R0.04_P34.6_T14.12s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C1_DBPRUNE_F1B0_R0.045_P38.2_T13.41s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C1_DBPRUNE_F1B0_R0.055_P45.1_T12.00s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C1_DBPRUNE_F1B0_R0.2_P59.5_T8.86s.png width=105px>|
 
-<div align="center">
-  <p align="center">
-    <b>♥️ Please consider to leave a ⭐️ Star to support us ~ ♥️</b>
-  </p>
-</div> 
-
 ## ©️Citations
 
 ```BibTeX
@@ -146,6 +140,12 @@ Moreover, **CacheDiT** are **plug-and-play** solutions that works hand-in-hand w
 }
 ```
 
+## 👋Reference
+
+<div id="reference"></div>
+
+The **CacheDiT** codebase is adapted from [FBCache](https://github.com/chengzeyi/ParaAttention/tree/main/src/para_attn/first_block_cache). Special thanks to their excellent work! 
+
 ## 📖Contents 
 
 <div id="contents"></div>  
@@ -153,6 +153,7 @@ Moreover, **CacheDiT** are **plug-and-play** solutions that works hand-in-hand w
 - [⚙️Installation](#️installation)
 - [🔥Supported Models](#supported)
 - [⚡️Dual Block Cache](#dbcache)
+- [🔥Hybrid TaylorSeer](#taylorseer)
 - [🎉First Block Cache](#fbcache)
 - [⚡️Dynamic Block Prune](#dbprune)
 - [🎉Context Parallelism](#context-parallelism)  
@@ -187,28 +188,19 @@ pip3 install git+https://github.com/vipshop/cache-dit.git
 - [🚀HunyuanVideo](https://github.com/vipshop/cache-dit/raw/main/examples)
 
 
-<!--
-<p align="center">
-  <h4> 🔥Supported Models🔥</h4>
-  <a href=https://github.com/vipshop/cache-dit/raw/main/examples> <b>🚀FLUX.1</b>: ✔️DBCache, ✔️DBPrune, ✔️FBCache🔥</a> <br>
-  <a href=https://github.com/vipshop/cache-dit/raw/main/examples> <b>🚀Mochi</b>: ✔️DBCache, ✔️DBPrune, ✔️FBCache🔥</a> <br>
-  <a href=https://github.com/vipshop/cache-dit/raw/main/examples> <b>🚀CogVideoX</b>: ✔️DBCache, ✔️DBPrune, ✔️FBCache🔥</a> <br>
-  <a href=https://github.com/vipshop/cache-dit/raw/main/examples> <b>🚀CogVideoX1.5</b>: ✔️DBCache, ✔️DBPrune, ✔️FBCache🔥</a> <br>
-  <a href=https://github.com/vipshop/cache-dit/raw/main/examples> <b>🚀Wan2.1</b>: ✔️DBCache, ✔️DBPrune, ✔️FBCache🔥</a> <br>
-  <a href=https://github.com/vipshop/cache-dit/raw/main/examples> <b>🚀HunyuanVideo</b>: ✔️DBCache, ✔️DBPrune, ✔️FBCache🔥</a> <br>
-</p>
--->
-
 ## ⚡️DBCache: Dual Block Cache  
 
 <div id="dbcache"></div>
 
-![](https://github.com/user-attachments/assets/c2a382b9-0ccd-46f4-aacc-87857b4a4de8)
+![](https://github.com/vipshop/cache-dit/raw/main/assets/dbcache-v1.png)
 
 **DBCache** provides configurable parameters for custom optimization, enabling a balanced trade-off between performance and precision:
 
 - **Fn**: Specifies that DBCache uses the **first n** Transformer blocks to fit the information at time step t, enabling the calculation of a more stable L1 diff and delivering more accurate information to subsequent blocks.
 - **Bn**: Further fuses approximate information in the **last n** Transformer blocks to enhance prediction accuracy. These blocks act as an auto-scaler for approximate hidden states that use residual cache.
+
+![](https://github.com/vipshop/cache-dit/raw/main/assets/dbcache-fnbn-v1.png)
+
 - **warmup_steps**: (default: 0) DBCache does not apply the caching strategy when the number of running steps is less than or equal to this value, ensuring the model sufficiently learns basic features during warmup.
 - **max_cached_steps**:  (default: -1) DBCache disables the caching strategy when the previous cached steps exceed this value to prevent precision degradation.
 - **residual_diff_threshold**: The value of residual diff threshold, a higher value leads to faster performance at the cost of lower precision.
@@ -264,11 +256,54 @@ cache_options = {
 |24.85s|15.59s|8.58s|15.41s|15.11s|17.74s|
 |<img src=https://github.com/vipshop/cache-dit/raw/main/assets/NONE_R0.08_S0.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F1B0S1_R0.08_S11.png width=105px> | <img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F1B0S1_R0.2_S19.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F8B8S1_R0.15_S15.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F12B12S4_R0.2_S16.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F16B16S4_R0.2_S13.png width=105px>|
 
+## 🔥Hybrid TaylorSeer
+
+<div id="taylorseer"></div>
+
+We have supported the [TaylorSeers: From Reusing to Forecasting: Accelerating Diffusion Models with TaylorSeers](https://arxiv.org/pdf/2503.06923) algorithm to further improve the precision of DBCache in cases where the cached steps are large, namely, **Hybrid TaylorSeer + DBCache**. At timesteps with significant intervals, the feature similarity in diffusion models decreases substantially, significantly harming the generation quality. 
+
+$$
+\mathcal{F}\_{\text {pred }, m}\left(x_{t-k}^l\right)=\mathcal{F}\left(x_t^l\right)+\sum_{i=1}^m \frac{\Delta^i \mathcal{F}\left(x_t^l\right)}{i!\cdot N^i}(-k)^i
+$$
+
+**TaylorSeer** employs a differential method to approximate the higher-order derivatives of features and predict features in future timesteps with Taylor series expansion. The TaylorSeer implemented in CacheDiT supports both hidden states and residual cache types. That is $\mathcal{F}\_{\text {pred }, m}\left(x_{t-k}^l\right)$ can be a residual cache or a hidden-state cache.
+
+```python
+cache_options = {
+    # TaylorSeer options
+    "enable_taylorseer": True,
+    "enable_encoder_taylorseer": True,
+    # Taylorseer cache type cache be hidden_states or residual.
+    "taylorseer_cache_type": "residual",
+    # Higher values of n_derivatives will lead to longer 
+    # computation time but may improve precision significantly.
+    "taylorseer_kwargs": {
+        "n_derivatives": 2, # default is 2.
+    },
+    "warmup_steps": 3, # n_derivatives + 1
+    "residual_diff_threshold": 0.12,
+}
+``` 
+
+> [!Important]
+> Please note that if you have used TaylorSeer as the calibrator for approximate hidden states, the **Bn** param of DBCache can be set to **0**. In essence, DBCache's Bn is also act as a calibrator, so you can choose either Bn > 0 or TaylorSeer. We recommend using the configuration scheme of **TaylorSeer** + **DBCache FnB0**.
+
+<div align="center">
+  <p align="center">
+    <b>DBCache F1B0 + TaylorSeer</b>, L20x1, Steps: 28, <br>"A cat holding a sign that says hello world with complex background"
+  </p>
+</div>
+
+|Baseline(L20x1)|F1B0 (0.12)|+TaylorSeer|F1B0 (0.15)|+TaylorSeer|+compile| 
+|:---:|:---:|:---:|:---:|:---:|:---:|
+|24.85s|12.85s|12.86s|10.27s|10.28s|8.48s|
+|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/NONE_R0.08_S0.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C0_DBCACHE_F1B0S1W0T0ET0_R0.12_S14_T12.85s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C0_DBCACHE_F1B0S1W0T1ET1_R0.12_S14_T12.86s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C0_DBCACHE_F1B0S1W0T0ET0_R0.15_S17_T10.27s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C0_DBCACHE_F1B0S1W0T1ET1_R0.15_S17_T10.28s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C1_DBCACHE_F1B0S1W0T1ET1_R0.15_S17_T8.48s.png width=105px>|
+
 ## 🎉FBCache: First Block Cache  
 
 <div id="fbcache"></div>
 
-![](https://github.com/user-attachments/assets/0fb66656-b711-457a-92a7-a830f134272d)
+![](https://github.com/vipshop/cache-dit/raw/main/assets/fbcache-v1.png)
 
 **DBCache** is a more general cache algorithm than **FBCache**. When Fn=1 and Bn=0, DBCache behaves identically to FBCache. Therefore, you can either use the original FBCache implementation directly or configure **DBCache** with **F1B0** settings to achieve the same functionality.
 
@@ -302,7 +337,7 @@ apply_cache_on_pipe(pipe, **cache_options)
 
 <div id="dbprune"></div>  
 
-![](https://github.com/user-attachments/assets/932b6360-9533-4352-b176-4c4d84bd4695)
+![](https://github.com/vipshop/cache-dit/raw/main/assets/dbprune-v1.png)
 
 We have further implemented a new **Dynamic Block Prune** algorithm based on **Residual Caching** for Diffusion Transformers, which is referred to as **DBPrune**. DBPrune caches each block's hidden states and residuals, then dynamically prunes blocks during inference by computing the L1 distance between previous hidden states. When a block is pruned, its output is approximated using the cached residuals. DBPrune is currently in the experimental phase, and we kindly invite you to stay tuned for upcoming updates.
 
@@ -389,7 +424,7 @@ from para_attn.context_parallel import init_context_parallel_mesh
 from para_attn.context_parallel.diffusers_adapters import parallelize_pipe
 from cache_dit.cache_factory import apply_cache_on_pipe, CacheType
 
- # Init distributed process group
+# Init distributed process group
 dist.init_process_group()
 torch.cuda.set_device(dist.get_rank())
 
@@ -436,14 +471,16 @@ torch._dynamo.config.recompile_limit = 96  # default is 8
 torch._dynamo.config.accumulated_recompile_limit = 2048  # default is 256
 ```
 
+Please check [bench.py](./bench/bench.py) for more details.
+
 ## 👋Contribute 
 <div id="contribute"></div>
 
-How to contribute? Star ⭐️ this repo to support us or check [CONTRIBUTE.md](./CONTRIBUTE.md).
+How to contribute? Star ⭐️ this repo to support us or check [CONTRIBUTE.md](https://github.com/vipshop/cache-dit/raw/main/CONTRIBUTE.md).
 
 ## ©️License   
 
 <div id="license"></div>
 
 
-We have followed the original License from [ParaAttention](https://github.com/chengzeyi/ParaAttention), please check [LICENSE](./LICENSE) for more details.
+We have followed the original License from [ParaAttention](https://github.com/chengzeyi/ParaAttention), please check [LICENSE](https://github.com/vipshop/cache-dit/raw/main/LICENSE) for more details.