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

cache_dit/cache_factory/cache_contexts/cache_context.py

@@ -5,6 +5,11 @@ from typing import Any, DefaultDict, Dict, List, Optional, Union, Tuple
 
 import torch
 
+from cache_dit.cache_factory.cache_contexts.cache_config import (
+    BasicCacheConfig,
+    ExtraCacheConfig,
+    DBCacheConfig,
+)
 from cache_dit.cache_factory.cache_contexts.calibrators import (
     Calibrator,
     CalibratorBase,
@@ -15,101 +20,16 @@ from cache_dit.logger import init_logger
 logger = init_logger(__name__)
 
 
-@dataclasses.dataclass
-class BasicCacheConfig:
-    # Dual Block Cache with Flexible FnBn configuration.
-
-    # Fn_compute_blocks: (`int`, *required*, defaults to 8):
-    #     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. Please check https://github.com/vipshop/cache-dit/blob/main/docs/DBCache.md
-    #     for more details of DBCache.
-    Fn_compute_blocks: int = 8
-    # Bn_compute_blocks: (`int`, *required*, defaults to 0):
-    #     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.
-    Bn_compute_blocks: int = 0
-    # residual_diff_threshold (`float`, *required*, defaults to 0.08):
-    #     the value of residual diff threshold, a higher value leads to faster performance at the
-    #     cost of lower precision.
-    residual_diff_threshold: Union[torch.Tensor, float] = 0.08
-    # 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.
-    max_warmup_steps: int = 8  # DON'T Cache in warmup steps
-    # 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.
-    warmup_interval: int = 1  # skip interval in warmup steps
-    # 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.
-    max_cached_steps: int = -1  # for both CFG and non-CFG
-    # max_continuous_cached_steps (`int`, *required*, defaults to -1):
-    #     DBCache disables the caching strategy when the previous continous cached steps exceed this value to
-    #     prevent precision degradation.
-    max_continuous_cached_steps: int = -1  # the max continuous cached steps
-    # enable_separate_cfg (`bool`, *required*,  defaults to None):
-    #     Whether to do separate cfg or not, such as Wan 2.1, Qwen-Image. For model that fused CFG
-    #     and non-CFG into single forward step, should set enable_separate_cfg as False, for example:
-    #     CogVideoX, HunyuanVideo, Mochi, etc.
-    enable_separate_cfg: Optional[bool] = None
-    # cfg_compute_first (`bool`, *required*,  defaults to False):
-    #     Compute cfg forward first or not, default False, namely, 0, 2, 4, ..., -> non-CFG step;
-    #     1, 3, 5, ... -> CFG step.
-    cfg_compute_first: bool = False
-    # cfg_diff_compute_separate (`bool`, *required*,  defaults to True):
-    #     Compute separate diff values for CFG and non-CFG step, default True. If False, we will
-    #     use the computed diff from current non-CFG transformer step for current CFG step.
-    cfg_diff_compute_separate: bool = True
-
-    def update(self, **kwargs) -> "BasicCacheConfig":
-        for key, value in kwargs.items():
-            if hasattr(self, key):
-                setattr(self, key, value)
-        return self
-
-    def strify(self) -> str:
-        return (
-            f"DBCACHE_F{self.Fn_compute_blocks}"
-            f"B{self.Bn_compute_blocks}_"
-            f"W{self.max_warmup_steps}"
-            f"I{self.warmup_interval}"
-            f"M{max(0, self.max_cached_steps)}"
-            f"MC{max(0, self.max_continuous_cached_steps)}_"
-            f"R{self.residual_diff_threshold}"
-        )
-
-
-@dataclasses.dataclass
-class ExtraCacheConfig:
-    # Some other not very important settings for Dual Block Cache.
-    # NOTE: These flags maybe deprecated in the future and users
-    # should never use these extra configurations in their cases.
-
-    # l1_hidden_states_diff_threshold (`float`, *optional*, defaults to None):
-    #     The hidden states diff threshold for DBCache if use hidden_states as
-    #     cache (not residual).
-    l1_hidden_states_diff_threshold: float = None
-    # important_condition_threshold (`float`, *optional*, defaults to 0.0):
-    #     Only select the most important tokens while calculating the l1 diff.
-    important_condition_threshold: float = 0.0
-    # downsample_factor (`int`, *optional*, defaults to 1):
-    #     Downsample factor for Fn buffer, in order the save GPU memory.
-    downsample_factor: int = 1
-    # num_inference_steps (`int`, *optional*, defaults to -1):
-    #     num_inference_steps for DiffusionPipeline, for future use.
-    num_inference_steps: int = -1
-
-
 @dataclasses.dataclass
 class CachedContext:
     name: str = "default"
     # Buffer for storing the residuals and other tensors
     buffers: Dict[str, Any] = dataclasses.field(default_factory=dict)
     # Basic Dual Block Cache Config
-    cache_config: BasicCacheConfig = dataclasses.field(
+    cache_config: Union[
+        BasicCacheConfig,
+        DBCacheConfig,
+    ] = dataclasses.field(
         default_factory=BasicCacheConfig,
     )
     # NOTE: Users should never use these extra configurations.
@@ -131,14 +51,14 @@ class CachedContext:
     # be double of executed_steps.
     transformer_executed_steps: int = 0
 
-    # CFG & non-CFG cached steps
+    # CFG & non-CFG cached/pruned steps
     cached_steps: List[int] = dataclasses.field(default_factory=list)
-    residual_diffs: DefaultDict[str, float] = dataclasses.field(
+    residual_diffs: DefaultDict[str, float | list] = dataclasses.field(
         default_factory=lambda: defaultdict(float),
     )
     continuous_cached_steps: int = 0
     cfg_cached_steps: List[int] = dataclasses.field(default_factory=list)
-    cfg_residual_diffs: DefaultDict[str, float] = dataclasses.field(
+    cfg_residual_diffs: DefaultDict[str, float | list] = dataclasses.field(
         default_factory=lambda: defaultdict(float),
     )
     cfg_continuous_cached_steps: int = 0
@@ -286,7 +206,9 @@ class CachedContext:
     def get_cfg_calibrators(self) -> Tuple[CalibratorBase, CalibratorBase]:
         return self.cfg_calibrator, self.cfg_encoder_calibrator
 
-    def add_residual_diff(self, diff):
+    def add_residual_diff(self, diff: float | torch.Tensor):
+        if isinstance(diff, torch.Tensor):
+            diff = diff.item()
         # step: executed_steps - 1, not transformer_steps - 1
         step = str(self.get_current_step())
         # Only add the diff if it is not already recorded for this step

cache_dit/cache_factory/cache_contexts/cache_manager.py

@@ -14,7 +14,7 @@ from cache_dit.logger import init_logger
 logger = init_logger(__name__)
 
 
-class CacheNotExistError(Exception):
+class ContextNotExistError(Exception):
     pass
 
 
@@ -36,14 +36,14 @@ class CachedContextManager:
             self._current_context = cached_context
         else:
             if cached_context not in self._cached_context_manager:
-                raise CacheNotExistError("Context not exist!")
+                raise ContextNotExistError("Context not exist!")
             self._current_context = self._cached_context_manager[cached_context]
         return self._current_context
 
     def get_context(self, name: str = None) -> CachedContext:
         if name is not None:
             if name not in self._cached_context_manager:
-                raise CacheNotExistError("Context not exist!")
+                raise ContextNotExistError("Context not exist!")
             return self._cached_context_manager[name]
         return self._current_context
 
@@ -482,7 +482,7 @@ class CachedContextManager:
 
             if calibrator is not None:
                 # Use calibrator to update the buffer
-                calibrator.update(buffer)
+                calibrator.update(buffer, name=prefix)
             else:
                 if logger.isEnabledFor(logging.DEBUG):
                     logger.debug(
@@ -513,7 +513,7 @@ class CachedContextManager:
                 calibrator, _ = self.get_calibrator()
 
             if calibrator is not None:
-                return calibrator.approximate()
+                return calibrator.approximate(name=prefix)
             else:
                 if logger.isEnabledFor(logging.DEBUG):
                     logger.debug(
@@ -551,7 +551,7 @@ class CachedContextManager:
 
             if encoder_calibrator is not None:
                 # Use CalibratorBase to update the buffer
-                encoder_calibrator.update(buffer)
+                encoder_calibrator.update(buffer, name=prefix)
             else:
                 if logger.isEnabledFor(logging.DEBUG):
                     logger.debug(
@@ -582,7 +582,7 @@ class CachedContextManager:
 
             if encoder_calibrator is not None:
                 # Use calibrator to approximate the value
-                return encoder_calibrator.approximate()
+                return encoder_calibrator.approximate(name=prefix)
             else:
                 if logger.isEnabledFor(logging.DEBUG):
                     logger.debug(

cache_dit/cache_factory/cache_contexts/calibrators/taylorseer.py

@@ -10,13 +10,12 @@ from cache_dit.logger import init_logger
 logger = init_logger(__name__)
 
 
-class TaylorSeerCalibrator(CalibratorBase):
+class TaylorSeerState:
     def __init__(
         self,
         n_derivatives=1,
         max_warmup_steps=1,
         skip_interval_steps=1,
-        **kwargs,
     ):
         self.n_derivatives = n_derivatives
         self.order = n_derivatives + 1
@@ -28,9 +27,8 @@ class TaylorSeerCalibrator(CalibratorBase):
             "dY_prev": [None] * self.order,
             "dY_current": [None] * self.order,
         }
-        self.reset_cache()
 
-    def reset_cache(self):  # NEED
+    def reset(self):
         self.state: Dict[str, List[torch.Tensor]] = {
             "dY_prev": [None] * self.order,
             "dY_current": [None] * self.order,
@@ -38,6 +36,9 @@ class TaylorSeerCalibrator(CalibratorBase):
         self.current_step = -1
         self.last_non_approximated_step = -1
 
+    def mark_step_begin(self):  # NEED
+        self.current_step += 1
+
     def should_compute(self, step=None):
         step = self.current_step if step is None else step
         if (
@@ -56,7 +57,7 @@ class TaylorSeerCalibrator(CalibratorBase):
         window = self.current_step - self.last_non_approximated_step
         if self.state["dY_prev"][0] is not None:
             if dY_current[0].shape != self.state["dY_prev"][0].shape:
-                self.reset_cache()
+                self.reset()
 
         for i in range(self.n_derivatives):
             if self.state["dY_prev"][i] is not None and self.current_step > 1:
@@ -77,9 +78,6 @@ class TaylorSeerCalibrator(CalibratorBase):
                 break
         return output
 
-    def mark_step_begin(self):  # NEED
-        self.current_step += 1
-
     def update(self, Y: torch.Tensor):  # NEED
         # Directly call this method will ingnore the warmup
         # policy and force full computation.
@@ -106,5 +104,77 @@ class TaylorSeerCalibrator(CalibratorBase):
         else:
             return self.approximate()
 
+
+class TaylorSeerCalibrator(CalibratorBase):
+    def __init__(
+        self,
+        n_derivatives=1,
+        max_warmup_steps=1,
+        skip_interval_steps=1,
+        **kwargs,
+    ):
+        self.n_derivatives = n_derivatives
+        self.max_warmup_steps = max_warmup_steps
+        self.skip_interval_steps = skip_interval_steps
+        self.states: Dict[str, TaylorSeerState] = {}
+        self.reset_cache()
+
+    def reset_cache(self):  # NEED
+        if self.states:
+            for state in self.states.values():
+                state.reset()
+
+    def maybe_init_state(
+        self,
+        name: str = "default",
+    ):
+        if name not in self.states:
+            self.states[name] = TaylorSeerState(
+                n_derivatives=self.n_derivatives,
+                max_warmup_steps=self.max_warmup_steps,
+                skip_interval_steps=self.skip_interval_steps,
+            )
+
+    def mark_step_begin(self, *args, **kwargs):
+        if self.states:
+            for state in self.states.values():
+                state.mark_step_begin()
+
+    def derivative(
+        self,
+        Y: torch.Tensor,
+        name: str = "default",
+    ) -> List[torch.Tensor]:
+        self.maybe_init_state(name)
+        state = self.states[name]
+        state.derivative(Y)
+        return state.state["dY_current"]
+
+    def approximate(
+        self,
+        name: str = "default",
+    ) -> torch.Tensor:  # NEED
+        assert name in self.states, f"State '{name}' not found."
+        state = self.states[name]
+        return state.approximate()
+
+    def update(
+        self,
+        Y: torch.Tensor,
+        name: str = "default",
+    ):  # NEED
+        self.maybe_init_state(name)
+        state = self.states[name]
+        state.update(Y)
+
+    def step(
+        self,
+        Y: torch.Tensor,
+        name: str = "default",
+    ):
+        self.maybe_init_state(name)
+        state = self.states[name]
+        return state.step(Y)
+
     def __repr__(self):
         return f"TaylorSeerCalibrator_O({self.n_derivatives})"

cache_dit/cache_factory/cache_contexts/context_manager.py

@@ -0,0 +1,29 @@
+from cache_dit.cache_factory.cache_types import CacheType
+from cache_dit.cache_factory.cache_contexts.cache_manager import (
+    CachedContextManager,
+)
+from cache_dit.cache_factory.cache_contexts.prune_manager import (
+    PrunedContextManager,
+)
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class ContextManager:
+    _supported_managers = (
+        CachedContextManager,
+        PrunedContextManager,
+    )
+
+    def __new__(
+        cls,
+        cache_type: CacheType,
+        name: str = "default",
+    ) -> CachedContextManager | PrunedContextManager:
+        if cache_type == CacheType.DBCache:
+            return CachedContextManager(name)
+        elif cache_type == CacheType.DBPrune:
+            return PrunedContextManager(name)
+        else:
+            raise ValueError(f"Unsupported cache_type: {cache_type}.")

cache_dit/cache_factory/cache_contexts/prune_config.py

@@ -0,0 +1,69 @@
+import dataclasses
+from typing import List
+from cache_dit.cache_factory.cache_types import CacheType
+from cache_dit.cache_factory.cache_contexts.cache_config import (
+    BasicCacheConfig,
+)
+
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclasses.dataclass
+class DBPruneConfig(BasicCacheConfig):
+    # Dyanamic Block Prune specific configurations
+    cache_type: CacheType = CacheType.DBPrune  # DBPrune
+
+    # enable_dynamic_prune_threshold (`bool`, *required*, defaults to False):
+    #     Whether to enable the dynamic prune threshold or not. If True, we will
+    #     compute the dynamic prune threshold based on the mean of the residual
+    #     diffs of the previous computed or pruned blocks.
+    #     But, also limit mean_diff to be at least 2x the residual_diff_threshold
+    #     to avoid too aggressive pruning.
+    enable_dynamic_prune_threshold: bool = False
+    # max_dynamic_prune_threshold (`float`, *optional*, defaults to None):
+    #     The max dynamic prune threshold, if not None, the dynamic prune threshold
+    #     will not exceed this value. If None, we will limit it to be at least 2x
+    #     the residual_diff_threshold to avoid too aggressive pruning.
+    max_dynamic_prune_threshold: float = None
+    # dynamic_prune_threshold_relax_ratio (`float`, *optional*, defaults to 1.25):
+    #     The relax ratio for dynamic prune threshold, the dynamic prune threshold
+    #     will be set as:
+    #         dynamic_prune_threshold = mean_diff * dynamic_prune_threshold_relax_ratio
+    #     to avoid too aggressive pruning.
+    #     The default value is 1.25, which means the dynamic prune threshold will
+    #     be 1.25 times the mean of the residual diffs of the previous computed
+    #     or pruned blocks.
+    #     Users can tune this value to achieve a better trade-off between speedup
+    #     and precision. A higher value leads to more aggressive pruning
+    #     and faster speedup, but may also lead to lower precision.
+    dynamic_prune_threshold_relax_ratio: float = 1.25
+    # non_prune_block_ids (`List[int]`, *optional*, defaults to []):
+    #     The list of block ids that will not be pruned, even if their residual
+    #     diffs are below the prune threshold. This can be useful for the first
+    #     few blocks, which are usually more important for the model performance.
+    non_prune_block_ids: List[int] = dataclasses.field(default_factory=list)
+    # force_reduce_calibrator_vram (`bool`, *optional*, defaults to True):
+    #     Whether to force reduce the VRAM usage of the calibrator for Dynamic Block
+    #     Prune. If True, we will set the downsample_factor of the extra_cache_config
+    #     to at least 2 to reduce the VRAM usage of the calibrator.
+    force_reduce_calibrator_vram: bool = False
+
+    def update(self, **kwargs) -> "DBPruneConfig":
+        for key, value in kwargs.items():
+            if hasattr(self, key):
+                setattr(self, key, value)
+        return self
+
+    def strify(self) -> str:
+        return (
+            f"{self.cache_type}_"
+            f"F{self.Fn_compute_blocks}"
+            f"B{self.Bn_compute_blocks}_"
+            f"W{self.max_warmup_steps}"
+            f"I{self.warmup_interval}"
+            f"M{max(0, self.max_cached_steps)}"
+            f"MC{max(0, self.max_continuous_cached_steps)}_"
+            f"R{self.residual_diff_threshold}"
+        )

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()