cache-dit 1.0.3__py3-none-any.whl → 1.0.5__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.
@@ -136,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),
@@ -159,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.
@@ -202,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

@@ -7,10 +7,6 @@ def load_cache_options_from_yaml(yaml_file_path):
             kwargs: dict = yaml.safe_load(f)
 
         required_keys = [
-            "max_warmup_steps",
-            "max_cached_steps",
-            "Fn_compute_blocks",
-            "Bn_compute_blocks",
             "residual_diff_threshold",
         ]
         for key in required_keys:
@@ -38,10 +34,25 @@ def load_cache_options_from_yaml(yaml_file_path):
                 )
             )
 
-        from cache_dit.cache_factory.cache_contexts import BasicCacheConfig
+        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)
+            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 cache_context_kwargs
 

cache_dit/quantize/quantize_ao.py

@@ -83,11 +83,18 @@ def quantize_ao(
     def _quantization_fn():
         try:
             if quant_type == "fp8_w8a8_dq":
-                from torchao.quantization import (
-                    float8_dynamic_activation_float8_weight,
-                    PerTensor,
-                    PerRow,
-                )
+                try:
+                    from torchao.quantization import (
+                        float8_dynamic_activation_float8_weight,
+                        PerTensor,
+                        PerRow,
+                    )
+                except ImportError:
+                    from torchao.quantization import (
+                        Float8DynamicActivationFloat8WeightConfig as float8_dynamic_activation_float8_weight,
+                        PerTensor,
+                        PerRow,
+                    )
 
                 if per_row:  # Ensure bfloat16
                     module.to(torch.bfloat16)
@@ -109,7 +116,12 @@ def quantize_ao(
                 )
 
             elif quant_type == "fp8_w8a16_wo":
-                from torchao.quantization import float8_weight_only
+                try:
+                    from torchao.quantization import float8_weight_only
+                except ImportError:
+                    from torchao.quantization import (
+                        Float8WeightOnlyConfig as float8_weight_only,
+                    )
 
                 quantization_fn = float8_weight_only(
                     weight_dtype=kwargs.get(
@@ -119,14 +131,25 @@ def quantize_ao(
                 )
 
             elif quant_type == "int8_w8a8_dq":
-                from torchao.quantization import (
-                    int8_dynamic_activation_int8_weight,
-                )
+                try:
+                    from torchao.quantization import (
+                        int8_dynamic_activation_int8_weight,
+                    )
+                except ImportError:
+                    from torchao.quantization import (
+                        Int8DynamicActivationInt8WeightConfig as int8_dynamic_activation_int8_weight,
+                    )
 
                 quantization_fn = int8_dynamic_activation_int8_weight()
 
             elif quant_type == "int8_w8a16_wo":
-                from torchao.quantization import int8_weight_only
+
+                try:
+                    from torchao.quantization import int8_weight_only
+                except ImportError:
+                    from torchao.quantization import (
+                        Int8WeightOnlyConfig as int8_weight_only,
+                    )
 
                 quantization_fn = int8_weight_only(
                     # group_size is None -> per_channel, else per group
@@ -134,23 +157,41 @@ def quantize_ao(
                 )
 
             elif quant_type == "int4_w4a8_dq":
-                from torchao.quantization import (
-                    int8_dynamic_activation_int4_weight,
-                )
+
+                try:
+                    from torchao.quantization import (
+                        int8_dynamic_activation_int4_weight,
+                    )
+                except ImportError:
+                    from torchao.quantization import (
+                        Int8DynamicActivationInt4WeightConfig as int8_dynamic_activation_int4_weight,
+                    )
 
                 quantization_fn = int8_dynamic_activation_int4_weight(
                     group_size=kwargs.get("group_size", 32),
                 )
 
             elif quant_type == "int4_w4a4_dq":
-                from torchao.quantization import (
-                    int4_dynamic_activation_int4_weight,
-                )
+
+                try:
+                    from torchao.quantization import (
+                        int4_dynamic_activation_int4_weight,
+                    )
+                except ImportError:
+                    from torchao.quantization import (
+                        Int4DynamicActivationInt4WeightConfig as int4_dynamic_activation_int4_weight,
+                    )
 
                 quantization_fn = int4_dynamic_activation_int4_weight()
 
             elif quant_type == "int4_w4a16_wo":
-                from torchao.quantization import int4_weight_only
+
+                try:
+                    from torchao.quantization import int4_weight_only
+                except ImportError:
+                    from torchao.quantization import (
+                        Int4WeightOnlyConfig as int4_weight_only,
+                    )
 
                 quantization_fn = int4_weight_only(
                     group_size=kwargs.get("group_size", 32),