cache-dit 0.2.2__tar.gz → 0.2.3__tar.gz

{cache_dit-0.2.2 → cache_dit-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cache_dit
-Version: 0.2.2
+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
@@ -283,7 +283,11 @@ cache_options = {
     "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"

{cache_dit-0.2.2 → cache_dit-0.2.3}/README.md

@@ -248,7 +248,11 @@ cache_options = {
     "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"

{cache_dit-0.2.2 → cache_dit-0.2.3}/bench/bench.py

@@ -21,9 +21,6 @@ def get_args() -> argparse.ArgumentParser:
     parser.add_argument("--alter", action="store_true", default=False)
     parser.add_argument("--taylorseer", action="store_true", default=False)
     parser.add_argument("--taylorseer-order", "--order", type=int, default=2)
-    parser.add_argument(
-        "--encoder-taylorseer", action="store_true", default=False
-    )
     parser.add_argument("--l1-diff", action="store_true", default=False)
     parser.add_argument("--rdt", type=float, default=0.08)
     parser.add_argument("--Fn-compute-blocks", "--Fn", type=int, default=1)
@@ -32,9 +29,15 @@ def get_args() -> argparse.ArgumentParser:
     parser.add_argument("--warmup-steps", type=int, default=0)
     parser.add_argument("--max-cached-steps", type=int, default=-1)
     parser.add_argument("--max-pruned-steps", type=int, default=-1)
+    parser.add_argument("--gen-device", type=str, default="cuda")
     parser.add_argument("--ulysses", type=int, default=None)
     parser.add_argument("--compile", action="store_true", default=False)
-    parser.add_argument("--gen-device", type=str, default="cuda")
+    parser.add_argument(
+        "--force-compile-all",
+        "--compile-all",
+        action="store_true",
+        default=False,
+    )
     return parser.parse_args()
 
 
@@ -52,12 +55,7 @@ def get_cache_options(cache_type: CacheType, args: argparse.Namespace):
     elif cache_type == CacheType.DBCache:
         cache_options = {
             "cache_type": CacheType.DBCache,
-            "warmup_steps": (
-                # TaylorSeer needs at least order + 1 warmup steps
-                max(args.warmup_steps, args.taylorseer_order + 1)
-                if (args.taylorseer or args.encoder_taylorseer)
-                else args.warmup_steps
-            ),
+            "warmup_steps": args.warmup_steps,
             "max_cached_steps": args.max_cached_steps,  # -1 means no limit
             # Fn=1, Bn=0, means FB Cache, otherwise, Dual Block Cache
             "Fn_compute_blocks": args.Fn_compute_blocks,  # Fn, F8, etc.
@@ -81,7 +79,7 @@ def get_cache_options(cache_type: CacheType, args: argparse.Namespace):
             "important_condition_threshold": 0.00,
             # TaylorSeer options
             "enable_taylorseer": args.taylorseer,
-            "enable_encoder_taylorseer": args.encoder_taylorseer,
+            "enable_encoder_taylorseer": args.taylorseer,
             # Taylorseer cache type cache be hidden_states or residual
             "taylorseer_cache_type": "residual",
             "taylorseer_kwargs": {
@@ -90,7 +88,7 @@ def get_cache_options(cache_type: CacheType, args: argparse.Namespace):
         }
     elif cache_type == CacheType.DBPrune:
         assert (
-            args.taylorseer is False and args.encoder_taylorseer is False
+            args.taylorseer is False
         ), "DBPrune does not support TaylorSeer yet."
         cache_options = {
             "cache_type": CacheType.DBPrune,
@@ -122,7 +120,6 @@ def get_cache_options(cache_type: CacheType, args: argparse.Namespace):
             f"{cache_type_str}_F{args.Fn_compute_blocks}"
             f"B{args.Bn_compute_blocks}S{args.Bn_steps}"
             f"W{args.warmup_steps}T{int(args.taylorseer)}"
-            f"ET{int(args.encoder_taylorseer)}"
             f"O{args.taylorseer_order}"
         )
     elif cache_type == CacheType.DBPrune:
@@ -132,7 +129,7 @@ def get_cache_options(cache_type: CacheType, args: argparse.Namespace):
         )
     elif cache_type == CacheType.FBCache:
         cache_type_str = (
-            f"{cache_type_str}_W{args.warmup_steps}" f"T{int(args.taylorseer)}"
+            f"{cache_type_str}_W{args.warmup_steps}T{int(args.taylorseer)}"
         )
     return cache_options, cache_type_str
 
@@ -201,8 +198,10 @@ def main():
                 "Only compile transformer blocks not the whole model "
                 "for FluxTransformer2DModel to keep higher precision."
             )
-            if args.taylorseer_order <= 2 or (
-                not args.taylorseer and not args.encoder_taylorseer
+            if (
+                args.taylorseer_order <= 2
+                or not args.taylorseer
+                or args.force_compile_all
             ):
                 # NOTE: Seems like compiling the whole transformer
                 # will cause precision issues while using TaylorSeer

{cache_dit-0.2.2 → cache_dit-0.2.3}/src/cache_dit/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.2'
-__version_tuple__ = version_tuple = (0, 2, 2)
+__version__ = version = '0.2.3'
+__version_tuple__ = version_tuple = (0, 2, 3)

{cache_dit-0.2.2 → cache_dit-0.2.3}/src/cache_dit/cache_factory/dual_block_cache/cache_context.py

@@ -1598,17 +1598,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, DBCachedTransformerBlocks
-    ) 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-0.2.2 → cache_dit-0.2.3}/src/cache_dit/cache_factory/dual_block_cache/diffusers_adapters/wan.py

@@ -92,6 +92,6 @@ def apply_db_cache_on_pipe(
         pipe.__class__._is_cached = True
 
     if not shallow_patch:
-        apply_db_cache_on_transformer(pipe.transformer, **kwargs)
+        apply_db_cache_on_transformer(pipe.transformer)
 
     return pipe

{cache_dit-0.2.2 → cache_dit-0.2.3}/src/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-0.2.2 → cache_dit-0.2.3}/src/cache_dit/cache_factory/first_block_cache/cache_context.py

@@ -714,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-0.2.2 → cache_dit-0.2.3}/src/cache_dit/cache_factory/taylorseer.py

@@ -1,4 +1,5 @@
 # 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
 
@@ -44,8 +45,6 @@ class TaylorSeer:
         # 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!
-        # reference: https://github.com/Shenyi-Z/TaylorSeer
-        # TaylorSeer-FLUX/src/flux/taylor_utils/__init__.py
         # TODO: Custom Triton/CUDA kernel for better performance,
         # especially for large n_derivatives.
         dY_current = [None] * self.ORDER

cache_dit-0.2.3/src/cache_dit/compile/__init__.py

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

cache_dit-0.2.3/src/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-0.2.2 → cache_dit-0.2.3}/src/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.2 → cache_dit-0.2.3}/src/cache_dit.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cache_dit
-Version: 0.2.2
+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
@@ -283,7 +283,11 @@ cache_options = {
     "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"

{cache_dit-0.2.2 → cache_dit-0.2.3}/src/cache_dit.egg-info/SOURCES.txt

@@ -21,7 +21,6 @@ assets/DBCACHE_F8B16S1_R0.2_S18.png
 assets/DBCACHE_F8B8S1_R0.08_S9.png
 assets/DBCACHE_F8B8S1_R0.12_S12.png
 assets/DBCACHE_F8B8S1_R0.15_S15.png
-assets/DBCache.png
 assets/DBPRUNE_F1B0_R0.03_P24.0_T19.43s.png
 assets/DBPRUNE_F1B0_R0.04_P34.6_T16.82s.png
 assets/DBPRUNE_F1B0_R0.05_P38.3_T15.95s.png
@@ -68,7 +67,6 @@ assets/U4_C1_DBPRUNE_F1B0_R0.2_P59.5_T3.96s.png
 assets/U4_C1_NONE_R0.08_S0_T7.78s.png
 assets/U4_C1_NONE_R0.08_S0_T7.79s.png
 assets/cache-dit-v1.png
-assets/cache-dit.png
 assets/dbcache-fnbn-v1.png
 assets/dbcache-v1.png
 assets/dbprune-v1.png
@@ -123,6 +121,10 @@ src/cache_dit/cache_factory/first_block_cache/diffusers_adapters/flux.py
 src/cache_dit/cache_factory/first_block_cache/diffusers_adapters/hunyuan_video.py
 src/cache_dit/cache_factory/first_block_cache/diffusers_adapters/mochi.py
 src/cache_dit/cache_factory/first_block_cache/diffusers_adapters/wan.py
+src/cache_dit/compile/__init__.py
+src/cache_dit/compile/utils.py
+src/cache_dit/custom_ops/__init__.py
+src/cache_dit/custom_ops/triton_taylorseer.py
 tests/.gitignore
 tests/README.md
 tests/taylorseer_approximation_order_2.png