cache-dit 1.0.2__py3-none-any.whl → 1.0.4__py3-none-any.whl

cache_dit/cache_factory/cache_contexts/prune_context.py

@@ -0,0 +1,155 @@
+import torch
+import logging
+import dataclasses
+from typing import List
+from cache_dit.cache_factory.cache_types import CacheType
+from cache_dit.cache_factory.cache_contexts.prune_config import (
+    DBPruneConfig,
+)
+from cache_dit.cache_factory.cache_contexts.cache_context import (
+    CachedContext,
+)
+from cache_dit.logger import init_logger
+
+
+logger = init_logger(__name__)
+
+
+@dataclasses.dataclass
+class PrunedContext(CachedContext):
+    # Overwrite the cache_config type for PrunedContext
+    cache_config: DBPruneConfig = dataclasses.field(
+        default_factory=DBPruneConfig,
+    )
+    # Specially for Dynamic Block Prune
+    pruned_blocks: List[int] = dataclasses.field(default_factory=list)
+    actual_blocks: List[int] = dataclasses.field(default_factory=list)
+    cfg_pruned_blocks: List[int] = dataclasses.field(default_factory=list)
+    cfg_actual_blocks: List[int] = dataclasses.field(default_factory=list)
+
+    def __post_init__(self):
+        super().__post_init__()
+        if not isinstance(self.cache_config, DBPruneConfig):
+            raise ValueError(
+                "PrunedContext only supports DBPruneConfig as cache_config."
+            )
+
+        if self.cache_config.cache_type == CacheType.DBPrune:
+            if (
+                self.calibrator_config is not None
+                and self.cache_config.force_reduce_calibrator_vram
+            ):
+                # May reduce VRAM usage for Dynamic Block Prune
+                self.extra_cache_config.downsample_factor = max(
+                    4, self.extra_cache_config.downsample_factor
+                )
+
+    def get_residual_diff_threshold(self):
+        # Overwite this func for Dynamic Block Prune
+        residual_diff_threshold = self.cache_config.residual_diff_threshold
+        if isinstance(residual_diff_threshold, torch.Tensor):
+            residual_diff_threshold = residual_diff_threshold.item()
+        if self.cache_config.enable_dynamic_prune_threshold:
+            # Compute the dynamic prune threshold based on the mean of the
+            # residual diffs of the previous computed or pruned blocks.
+            step = str(self.get_current_step())
+            if int(step) >= 0 and str(step) in self.residual_diffs:
+                assert isinstance(self.residual_diffs[step], list)
+                # Use all the recorded diffs for this step
+                # NOTE: Should we only use the last 5 diffs?
+                diffs = self.residual_diffs[step][:5]
+                diffs = [d for d in diffs if d > 0.0]
+                if diffs:
+                    mean_diff = sum(diffs) / len(diffs)
+                    relaxed_diff = (
+                        mean_diff
+                        * self.cache_config.dynamic_prune_threshold_relax_ratio
+                    )
+                    if self.cache_config.max_dynamic_prune_threshold is None:
+                        max_dynamic_prune_threshold = (
+                            2 * residual_diff_threshold
+                        )
+                    else:
+                        max_dynamic_prune_threshold = (
+                            self.cache_config.max_dynamic_prune_threshold
+                        )
+                    if relaxed_diff < max_dynamic_prune_threshold:
+                        # If the mean diff is less than twice the threshold,
+                        # we can use it as the dynamic prune threshold.
+                        residual_diff_threshold = (
+                            relaxed_diff
+                            if relaxed_diff > residual_diff_threshold
+                            else residual_diff_threshold
+                        )
+                    if logger.isEnabledFor(logging.DEBUG):
+                        logger.debug(
+                            f"Dynamic prune threshold for step {step}: "
+                            f"{residual_diff_threshold:.6f}"
+                        )
+        return residual_diff_threshold
+
+    def mark_step_begin(self):
+        # Overwite this func for Dynamic Block Prune
+        super().mark_step_begin()
+        # Reset pruned_blocks and actual_blocks at the beginning
+        # of each transformer step.
+        if self.get_current_transformer_step() == 0:
+            self.pruned_blocks.clear()
+            self.actual_blocks.clear()
+
+    def add_residual_diff(self, diff: float | torch.Tensor):
+        # Overwite this func for Dynamic Block Prune
+        if isinstance(diff, torch.Tensor):
+            diff = diff.item()
+        # step: executed_steps - 1, not transformer_steps - 1
+        step = str(self.get_current_step())
+        # For Dynamic Block Prune, we will record all the diffs for this step
+        # Only add the diff if it is not already recorded for this step
+        if not self.is_separate_cfg_step():
+            if step not in self.residual_diffs:
+                self.residual_diffs[step] = []
+            self.residual_diffs[step].append(diff)
+        else:
+            if step not in self.cfg_residual_diffs:
+                self.cfg_residual_diffs[step] = []
+            self.cfg_residual_diffs[step].append(diff)
+
+    def add_pruned_step(self):
+        curr_cached_step = self.get_current_step()
+        # Avoid adding the same step multiple times
+        if not self.is_separate_cfg_step():
+            if curr_cached_step not in self.cached_steps:
+                self.add_cached_step()
+        else:
+            if curr_cached_step not in self.cfg_cached_steps:
+                self.add_cached_step()
+
+    def add_pruned_block(self, num_blocks):
+        if not self.is_separate_cfg_step():
+            self.pruned_blocks.append(num_blocks)
+        else:
+            self.cfg_pruned_blocks.append(num_blocks)
+
+    def add_actual_block(self, num_blocks):
+        if not self.is_separate_cfg_step():
+            self.actual_blocks.append(num_blocks)
+        else:
+            self.cfg_actual_blocks.append(num_blocks)
+
+    def get_pruned_blocks(self):
+        return self.pruned_blocks.copy()
+
+    def get_cfg_pruned_blocks(self):
+        return self.cfg_pruned_blocks.copy()
+
+    def get_actual_blocks(self):
+        return self.actual_blocks.copy()
+
+    def get_cfg_actual_blocks(self):
+        return self.cfg_actual_blocks.copy()
+
+    def get_pruned_steps(self):
+        return self.get_cached_steps()
+
+    def get_cfg_pruned_steps(self):
+        return self.get_cfg_cached_steps()

cache_dit/cache_factory/cache_contexts/prune_manager.py

@@ -0,0 +1,154 @@
+import torch
+import functools
+from typing import Dict, List, Tuple, Union
+
+from cache_dit.cache_factory.cache_contexts.cache_manager import (
+    CachedContextManager,
+)
+from cache_dit.cache_factory.cache_contexts.prune_context import (
+    PrunedContext,
+)
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class PrunedContextManager(CachedContextManager):
+    # Reuse CachedContextManager for Dynamic Block Prune
+
+    def __init__(self, name: str = None):
+        super().__init__(name)
+        # Overwrite for Dynamic Block Prune
+        self._current_context: PrunedContext = None
+        self._cached_context_manager: Dict[str, PrunedContext] = {}
+
+    # Overwrite for Dynamic Block Prune
+    def new_context(self, *args, **kwargs) -> PrunedContext:
+        _context = PrunedContext(*args, **kwargs)
+        self._cached_context_manager[_context.name] = _context
+        return _context
+
+    def set_context(
+        self,
+        cached_context: PrunedContext | str,
+    ) -> PrunedContext:
+        return super().set_context(cached_context)
+
+    def get_context(self, name: str = None) -> PrunedContext:
+        return super().get_context(name)
+
+    def reset_context(
+        self,
+        cached_context: PrunedContext | str,
+        *args,
+        **kwargs,
+    ) -> PrunedContext:
+        return super().reset_context(cached_context, *args, **kwargs)
+
+    # Specially for Dynamic Block Prune
+    @torch.compiler.disable
+    def add_pruned_step(self):
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        cached_context.add_pruned_step()
+
+    @torch.compiler.disable
+    def add_pruned_block(self, num_blocks):
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        cached_context.add_pruned_block(num_blocks)
+
+    @torch.compiler.disable
+    def add_actual_block(self, num_blocks):
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        cached_context.add_actual_block(num_blocks)
+
+    @torch.compiler.disable
+    def get_pruned_steps(self) -> List[int]:
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        return cached_context.get_pruned_steps()
+
+    @torch.compiler.disable
+    def get_cfg_pruned_steps(self) -> List[int]:
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        return cached_context.get_cfg_pruned_steps()
+
+    @torch.compiler.disable
+    def get_pruned_blocks(self) -> List[int]:
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        return cached_context.get_pruned_blocks()
+
+    @torch.compiler.disable
+    def get_actual_blocks(self) -> List[int]:
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        return cached_context.get_actual_blocks()
+
+    @torch.compiler.disable
+    def get_cfg_pruned_blocks(self) -> List[int]:
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        return cached_context.get_cfg_pruned_blocks()
+
+    @torch.compiler.disable
+    def get_cfg_actual_blocks(self) -> List[int]:
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        return cached_context.get_cfg_actual_blocks()
+
+    @torch.compiler.disable
+    @functools.lru_cache(maxsize=8)
+    def get_non_prune_blocks_ids(self, num_blocks: int) -> List[int]:
+        assert num_blocks is not None, "num_blocks must be provided"
+        assert num_blocks > 0, "num_blocks must be greater than 0"
+        # Get the non-prune block ids for current context
+        # Never prune the first `Fn` and last `Bn` blocks.
+        Fn_compute_blocks_ids = list(
+            range(
+                self.Fn_compute_blocks()
+                if self.Fn_compute_blocks() < num_blocks
+                else num_blocks
+            )
+        )
+
+        Bn_compute_blocks_ids = list(
+            range(
+                num_blocks
+                - (
+                    self.Bn_compute_blocks()
+                    if self.Bn_compute_blocks() < num_blocks
+                    else num_blocks
+                ),
+                num_blocks,
+            )
+        )
+        context = self.get_context()
+        assert context is not None, "cached_context must be set before"
+
+        non_prune_blocks_ids = list(
+            set(
+                Fn_compute_blocks_ids
+                + Bn_compute_blocks_ids
+                + context.cache_config.non_prune_block_ids
+            )
+        )
+        non_prune_blocks_ids = [
+            d for d in non_prune_blocks_ids if d < num_blocks
+        ]
+        return sorted(non_prune_blocks_ids)
+
+    @torch.compiler.disable
+    def can_prune(self, *args, **kwargs) -> bool:
+        # Directly reuse can_cache for Dynamic Block Prune
+        return self.can_cache(*args, **kwargs)
+
+    @torch.compiler.disable
+    def apply_prune(
+        self, *args, **kwargs
+    ) -> Tuple[torch.Tensor, Union[torch.Tensor, None]]:
+        # Directly reuse apply_cache for Dynamic Block Prune
+        return self.apply_cache(*args, **kwargs)

cache_dit/cache_factory/cache_interface.py

@@ -5,6 +5,8 @@ from cache_dit.cache_factory.block_adapters import BlockAdapter
 from cache_dit.cache_factory.block_adapters import BlockAdapterRegistry
 from cache_dit.cache_factory.cache_adapters import CachedAdapter
 from cache_dit.cache_factory.cache_contexts import BasicCacheConfig
+from cache_dit.cache_factory.cache_contexts import DBCacheConfig
+from cache_dit.cache_factory.cache_contexts import DBPruneConfig
 from cache_dit.cache_factory.cache_contexts import CalibratorConfig
 from cache_dit.cache_factory.params_modifier import ParamsModifier
 
@@ -19,8 +21,12 @@ def enable_cache(
         DiffusionPipeline,
         BlockAdapter,
     ],
-    # Basic DBCache config: BasicCacheConfig
-    cache_config: BasicCacheConfig = BasicCacheConfig(),
+    # BasicCacheConfig, DBCacheConfig, DBPruneConfig, etc.
+    cache_config: Union[
+        BasicCacheConfig,
+        DBCacheConfig,
+        DBPruneConfig,
+    ] = DBCacheConfig(),
     # Calibrator config: TaylorSeerCalibratorConfig, etc.
     calibrator_config: Optional[CalibratorConfig] = None,
     # Modify cache context params for specific blocks.
@@ -86,6 +92,9 @@ def enable_cache(
                 max_warmup_steps (`int`, *required*, defaults to 8):
                     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.
+                warmup_interval (`int`, *required*, defaults to 1):
+                    Skip interval in warmup steps, e.g., when warmup_interval is 2, only 0, 2, 4, ... steps
+                    in warmup steps will be computed, others will use dynamic cache.
                 max_cached_steps (`int`, *required*, defaults to -1):
                     DBCache disables the caching strategy when the previous cached steps exceed this value to
                     prevent precision degradation.
@@ -133,14 +142,14 @@ def enable_cache(
     >>> cache_dit.disable_cache(pipe) # Disable cache and run original pipe.
     """
     # Collect cache context kwargs
-    cache_context_kwargs = {}
-    if (cache_type := cache_context_kwargs.pop("cache_type", None)) is not None:
+    context_kwargs = {}
+    if (cache_type := context_kwargs.get("cache_type", None)) is not None:
         if cache_type == CacheType.NONE:
             return pipe_or_adapter
 
     # WARNING: Deprecated cache config params. These parameters are now retained
     # for backward compatibility but will be removed in the future.
-    deprecated_cache_kwargs = {
+    deprecated_kwargs = {
         "Fn_compute_blocks": kwargs.get("Fn_compute_blocks", None),
         "Bn_compute_blocks": kwargs.get("Bn_compute_blocks", None),
         "max_warmup_steps": kwargs.get("max_warmup_steps", None),
@@ -156,23 +165,23 @@ def enable_cache(
         ),
     }
 
-    deprecated_cache_kwargs = {
-        k: v for k, v in deprecated_cache_kwargs.items() if v is not None
+    deprecated_kwargs = {
+        k: v for k, v in deprecated_kwargs.items() if v is not None
     }
 
-    if deprecated_cache_kwargs:
+    if deprecated_kwargs:
         logger.warning(
             "Manually settup DBCache context without BasicCacheConfig is "
             "deprecated and will be removed in the future, please use "
             "`cache_config` parameter instead!"
         )
         if cache_config is not None:
-            cache_config.update(**deprecated_cache_kwargs)
+            cache_config.update(**deprecated_kwargs)
         else:
-            cache_config = BasicCacheConfig(**deprecated_cache_kwargs)
+            cache_config = BasicCacheConfig(**deprecated_kwargs)
 
     if cache_config is not None:
-        cache_context_kwargs["cache_config"] = cache_config
+        context_kwargs["cache_config"] = cache_config
 
     # WARNING: Deprecated taylorseer params. These parameters are now retained
     # for backward compatibility but will be removed in the future.
@@ -199,15 +208,15 @@ def enable_cache(
         )
 
     if calibrator_config is not None:
-        cache_context_kwargs["calibrator_config"] = calibrator_config
+        context_kwargs["calibrator_config"] = calibrator_config
 
     if params_modifiers is not None:
-        cache_context_kwargs["params_modifiers"] = params_modifiers
+        context_kwargs["params_modifiers"] = params_modifiers
 
     if isinstance(pipe_or_adapter, (DiffusionPipeline, BlockAdapter)):
         return CachedAdapter.apply(
             pipe_or_adapter,
-            **cache_context_kwargs,
+            **context_kwargs,
         )
     else:
         raise ValueError(

cache_dit/cache_factory/cache_types.py

@@ -6,7 +6,8 @@ logger = init_logger(__name__)
 
 class CacheType(Enum):
     NONE = "NONE"
-    DBCache = "Dual_Block_Cache"
+    DBCache = "DBCache"  # "Dual_Block_Cache"
+    DBPrune = "DBPrune"  # "Dynamic_Block_Prune"
 
     @staticmethod
     def type(type_hint: "CacheType | str") -> "CacheType":
@@ -14,6 +15,9 @@ class CacheType(Enum):
             return type_hint
         return cache_type(type_hint)
 
+    def __str__(self) -> str:
+        return self.value
+
 
 def cache_type(type_hint: "CacheType | str") -> "CacheType":
     if type_hint is None:
@@ -21,7 +25,6 @@ def cache_type(type_hint: "CacheType | str") -> "CacheType":
 
     if isinstance(type_hint, CacheType):
         return type_hint
-
     elif type_hint.upper() in (
         "DUAL_BLOCK_CACHE",
         "DB_CACHE",
@@ -29,6 +32,20 @@ def cache_type(type_hint: "CacheType | str") -> "CacheType":
         "DB",
     ):
         return CacheType.DBCache
+    elif type_hint.upper() in (
+        "DYNAMIC_BLOCK_PRUNE",
+        "DB_PRUNE",
+        "DBPRUNE",
+        "DBP",
+    ):
+        return CacheType.DBPrune
+    elif type_hint.upper() in (
+        "NONE",
+        "NO_CACHE",
+        "NOCACHE",
+        "NC",
+    ):
+        return CacheType.NONE
     return CacheType.NONE
 
 

cache_dit/cache_factory/params_modifier.py

@@ -11,7 +11,7 @@ logger = init_logger(__name__)
 class ParamsModifier:
     def __init__(
         self,
-        # Basic DBCache config: BasicCacheConfig
+        # BasicCacheConfig, DBCacheConfig, DBPruneConfig, etc.
         cache_config: BasicCacheConfig = None,
         # Calibrator config: TaylorSeerCalibratorConfig, etc.
         calibrator_config: Optional[CalibratorConfig] = None,
@@ -22,7 +22,7 @@ class ParamsModifier:
 
         # WARNING: Deprecated cache config params. These parameters are now retained
         # for backward compatibility but will be removed in the future.
-        deprecated_cache_kwargs = {
+        deprecated_kwargs = {
             "Fn_compute_blocks": kwargs.get("Fn_compute_blocks", None),
             "Bn_compute_blocks": kwargs.get("Bn_compute_blocks", None),
             "max_warmup_steps": kwargs.get("max_warmup_steps", None),
@@ -40,20 +40,20 @@ class ParamsModifier:
             ),
         }
 
-        deprecated_cache_kwargs = {
-            k: v for k, v in deprecated_cache_kwargs.items() if v is not None
+        deprecated_kwargs = {
+            k: v for k, v in deprecated_kwargs.items() if v is not None
         }
 
-        if deprecated_cache_kwargs:
+        if deprecated_kwargs:
             logger.warning(
                 "Manually settup DBCache context without BasicCacheConfig is "
                 "deprecated and will be removed in the future, please use "
                 "`cache_config` parameter instead!"
             )
             if cache_config is not None:
-                cache_config.update(**deprecated_cache_kwargs)
+                cache_config.update(**deprecated_kwargs)
             else:
-                cache_config = BasicCacheConfig(**deprecated_cache_kwargs)
+                cache_config = BasicCacheConfig(**deprecated_kwargs)
 
         if cache_config is not None:
             self._context_kwargs["cache_config"] = cache_config

cache_dit/cache_factory/utils.py

@@ -1,49 +1,60 @@
 import yaml
-from cache_dit.cache_factory import CacheType
 
 
 def load_cache_options_from_yaml(yaml_file_path):
     try:
         with open(yaml_file_path, "r") as f:
-            config = yaml.safe_load(f)
+            kwargs: dict = yaml.safe_load(f)
 
         required_keys = [
-            "cache_type",
-            "max_warmup_steps",
-            "max_cached_steps",
-            "Fn_compute_blocks",
-            "Bn_compute_blocks",
             "residual_diff_threshold",
         ]
         for key in required_keys:
-            if key not in config:
+            if key not in kwargs:
                 raise ValueError(
                     f"Configuration file missing required item: {key}"
                 )
 
-        # Convert cache_type to CacheType enum
-        if isinstance(config["cache_type"], str):
-            try:
-                config["cache_type"] = CacheType[config["cache_type"]]
-            except KeyError:
-                valid_types = [ct.name for ct in CacheType]
-                raise ValueError(
-                    f"Invalid cache_type value: {config['cache_type']}, "
-                    f"valid values are: {valid_types}"
+        cache_context_kwargs = {}
+        if kwargs.get("enable_taylorseer", False):
+            from cache_dit.cache_factory.cache_contexts.calibrators import (
+                TaylorSeerCalibratorConfig,
+            )
+
+            cache_context_kwargs["calibrator_config"] = (
+                TaylorSeerCalibratorConfig(
+                    enable_calibrator=kwargs.pop("enable_taylorseer"),
+                    enable_encoder_calibrator=kwargs.pop(
+                        "enable_encoder_taylorseer", False
+                    ),
+                    calibrator_cache_type=kwargs.pop(
+                        "taylorseer_cache_type", "residual"
+                    ),
+                    taylorseer_order=kwargs.pop("taylorseer_order", 1),
                 )
-        elif not isinstance(config["cache_type"], CacheType):
-            raise ValueError(
-                f"cache_type must be a string or CacheType enum, "
-                f"got: {type(config['cache_type'])}"
             )
 
-        # Handle default value for taylorseer_kwargs
-        if "taylorseer_kwargs" not in config and config.get(
-            "enable_taylorseer", False
-        ):
-            config["taylorseer_kwargs"] = {"n_derivatives": 2}
+        if "cache_type" not in kwargs:
+            from cache_dit.cache_factory.cache_contexts import BasicCacheConfig
+
+            cache_context_kwargs["cache_config"] = BasicCacheConfig()
+            cache_context_kwargs["cache_config"].update(**kwargs)
+        else:
+            cache_type = kwargs.pop("cache_type")
+            if cache_type == "DBCache":
+                from cache_dit.cache_factory.cache_contexts import DBCacheConfig
+
+                cache_context_kwargs["cache_config"] = DBCacheConfig()
+                cache_context_kwargs["cache_config"].update(**kwargs)
+            elif cache_type == "DBPrune":
+                from cache_dit.cache_factory.cache_contexts import DBPruneConfig
+
+                cache_context_kwargs["cache_config"] = DBPruneConfig()
+                cache_context_kwargs["cache_config"].update(**kwargs)
+            else:
+                raise ValueError(f"Unsupported cache_type: {cache_type}.")
 
-        return config
+        return cache_context_kwargs
 
     except FileNotFoundError:
         raise FileNotFoundError(