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

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

@@ -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/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/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.2.dist-info → cache_dit-0.2.3.dist-info}/METADATA

@@ -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.dist-info → cache_dit-0.2.3.dist-info}/RECORD

@@ -1,20 +1,20 @@
 cache_dit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cache_dit/_version.py,sha256=OjGGK5TcHVG44Y62aAqeJH4CskkZoY9ydbHOtCDew50,511
-cache_dit/logger.py,sha256=dKfNe_RRk9HJwfgHGeRR1f0LbskJpKdGmISCbL9roQs,3443
+cache_dit/_version.py,sha256=wD8hnA5gV5UmPkQnpT3xR6V2csgj9K5NEADogbLK79M,511
+cache_dit/logger.py,sha256=0zsu42hN-3-rgGC_C29ms1IvVpV4_b4_SwJCKSenxBE,4304
 cache_dit/primitives.py,sha256=A2iG9YLot3gOsZSPp-_gyjqjLgJvWQRx8aitD4JQ23Y,3877
 cache_dit/cache_factory/__init__.py,sha256=5RNuhWakvvqrOV4vkqrEBA7d-V1LwcNSsjtW14mkqK8,5255
-cache_dit/cache_factory/taylorseer.py,sha256=G3Bu7tPqKP1zt4pzN4gJ1T4VJobyaktC3zNSR1MACqc,4005
+cache_dit/cache_factory/taylorseer.py,sha256=LKSNo2ode69EVo9xrxjxAMEjz0yDGiGADeDYnEqddA8,3987
 cache_dit/cache_factory/utils.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cache_dit/cache_factory/dual_block_cache/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cache_dit/cache_factory/dual_block_cache/cache_context.py,sha256=ip3bY79lEGo2xDcGZuRxDu065q1eYabXLJUV8BsfhvM,59698
+cache_dit/cache_factory/dual_block_cache/cache_context.py,sha256=EV0REfwLuXTqnkZ9vD2nFFjRceRLrXhl1b1SO5N4os8,59272
 cache_dit/cache_factory/dual_block_cache/diffusers_adapters/__init__.py,sha256=krNAICf-aS3JLmSG8vOB9tpLa04uYRcABsC8PMbVUKY,1870
 cache_dit/cache_factory/dual_block_cache/diffusers_adapters/cogvideox.py,sha256=fibkeU-FHa30BNT-uPV2Eqcd5IRli07EKb25tMDp23c,2270
 cache_dit/cache_factory/dual_block_cache/diffusers_adapters/flux.py,sha256=fddSpTHXU24COMGAY-Z21EmHHAEArZBv_-XLRFD6ADU,2625
 cache_dit/cache_factory/dual_block_cache/diffusers_adapters/hunyuan_video.py,sha256=wcZdBhjUB8WSfz40A268BtSe3nr_hRsIi2BNlg1FHRU,9965
 cache_dit/cache_factory/dual_block_cache/diffusers_adapters/mochi.py,sha256=Cmy0KHRDgwXqtmqfkrr7kw0CP6CmkSnuz29gDHcD6sQ,2262
-cache_dit/cache_factory/dual_block_cache/diffusers_adapters/wan.py,sha256=oOb-NNesuyBAMwpRbFDG93TCwffLO7-TXZ4bvEtvGJc,2604
+cache_dit/cache_factory/dual_block_cache/diffusers_adapters/wan.py,sha256=kNoCZshNUxgwS9Di84Mz8Js1QAcs_U665x6wcSKYE2A,2594
 cache_dit/cache_factory/dynamic_block_prune/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cache_dit/cache_factory/dynamic_block_prune/prune_context.py,sha256=YRDwZ_16yjThpgVgDv6YaIB4QCE9nEkE-MOru0jOd50,35026
+cache_dit/cache_factory/dynamic_block_prune/prune_context.py,sha256=so1wGdb8W0ATwrjv7E5IEZLPcobybaY1HJa6hBYlOOQ,34698
 cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/__init__.py,sha256=hVBTXj9MMGFGVezT3j8MntFRBiphSaUL4YhSOd8JtuY,1870
 cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/cogvideox.py,sha256=KP8NxtHAKzzBOoX0lhvlMgY_5dmP4Z3T5TOfwl4SSyg,2273
 cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/flux.py,sha256=kCB7lL4OIq8TZn-baMIF8D_PVPTFW60omCMVQCb8ebs,2628
@@ -22,15 +22,19 @@ cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/hunyuan_video.py,
 cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/mochi.py,sha256=zXgoRDDjus3a2WSjtNh4ERtQp20ceb6nzohHMDlo2zY,2265
 cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/wan.py,sha256=PA7nuLgfAelnaI8usQx0Kxi8XATzMapyR1WndEdFoZA,2604
 cache_dit/cache_factory/first_block_cache/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cache_dit/cache_factory/first_block_cache/cache_context.py,sha256=oeOmVDho8aJa86p8LGACAWltu6Fe4chOW2OW8aPtB5c,23643
+cache_dit/cache_factory/first_block_cache/cache_context.py,sha256=NeAfDJlJVVUAL4btax5_iOLTuue1x4qeXwk0pM-QH28,23219
 cache_dit/cache_factory/first_block_cache/diffusers_adapters/__init__.py,sha256=-FFgA2MoudEo7uDacg4aWgm1KwfLZFsEDTVxatgbq9M,2146
 cache_dit/cache_factory/first_block_cache/diffusers_adapters/cogvideox.py,sha256=qO5CWyurtwW30mvOe6cxeQPTSXLDlPJcezm72zEjDq8,2375
 cache_dit/cache_factory/first_block_cache/diffusers_adapters/flux.py,sha256=Dcd4OzABCtyQCZNX2KNnUTdVoO1E1ApM7P8gcVYzcK0,2733
 cache_dit/cache_factory/first_block_cache/diffusers_adapters/hunyuan_video.py,sha256=OL7W4ukYlZz0IDmBR1zVV6XT3Mgciglj9Hqzv1wUAkQ,10092
 cache_dit/cache_factory/first_block_cache/diffusers_adapters/mochi.py,sha256=lQTClo52OwPbNEE4jiBZQhfC7hbtYqnYIABp_vbm_dk,2363
 cache_dit/cache_factory/first_block_cache/diffusers_adapters/wan.py,sha256=dBNzHBECAuTTA1a7kLdvZL20YzaKTAS3iciVLzKKEWA,2638
-cache_dit-0.2.2.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
-cache_dit-0.2.2.dist-info/METADATA,sha256=-4Sl6vCYl5bPBQVtlHGhHvEz0Nn-o3R5mYArGARsjLU,24521
-cache_dit-0.2.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-cache_dit-0.2.2.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
-cache_dit-0.2.2.dist-info/RECORD,,
+cache_dit/compile/__init__.py,sha256=DfMdPleFFGADXLsr7zXui8BTz_y9futY6rNmNdh9y7k,63
+cache_dit/compile/utils.py,sha256=KU60xc474Anbj7Y_FLRFmNxEjVYLLXkhbtCLXO7o_Tc,3699
+cache_dit/custom_ops/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+cache_dit/custom_ops/triton_taylorseer.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+cache_dit-0.2.3.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
+cache_dit-0.2.3.dist-info/METADATA,sha256=XH7Wn7GdTRth6g0yAY5heArGqXmrrC7CflOJ8ZXH-1k,24867
+cache_dit-0.2.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+cache_dit-0.2.3.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
+cache_dit-0.2.3.dist-info/RECORD,,