cache-dit 0.2.27__py3-none-any.whl → 0.2.29__py3-none-any.whl

cache_dit/cache_factory/cache_contexts/cache_context.py

@@ -1,11 +1,9 @@
 import logging
-import contextlib
 import dataclasses
 from collections import defaultdict
 from typing import Any, DefaultDict, Dict, List, Optional, Union, Tuple
 
 import torch
-import torch.distributed as dist
 
 from cache_dit.cache_factory.cache_contexts.taylorseer import TaylorSeer
 from cache_dit.logger import init_logger
@@ -14,7 +12,7 @@ logger = init_logger(__name__)
 
 
 @dataclasses.dataclass
-class _CachedContext:  # Internal CachedContext Impl class
+class CachedContext:  # Internal CachedContext Impl class
     name: str = "default"
     # Dual Block Cache
     # Fn=1, Bn=0, means FB Cache, otherwise, Dual Block Cache
@@ -67,15 +65,16 @@ class _CachedContext:  # Internal CachedContext Impl class
     enable_encoder_taylorseer: bool = False
     # NOTE: use residual cache for taylorseer may incur precision loss
     taylorseer_cache_type: str = "hidden_states"  # residual or hidden_states
+    taylorseer_order: int = 2  # The order for TaylorSeer
     taylorseer_kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict)
     taylorseer: Optional[TaylorSeer] = None
     encoder_tarlorseer: Optional[TaylorSeer] = None
 
-    # Support do_separate_cfg, such as Wan 2.1,
+    # Support enable_spearate_cfg, such as Wan 2.1,
     # Qwen-Image. For model that fused CFG and non-CFG into single
-    # forward step, should set do_separate_cfg as False.
+    # forward step, should set enable_spearate_cfg as False.
     # For example: CogVideoX, HunyuanVideo, Mochi.
-    do_separate_cfg: bool = False
+    enable_spearate_cfg: bool = 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
@@ -103,10 +102,10 @@ class _CachedContext:  # Internal CachedContext Impl class
         if logger.isEnabledFor(logging.DEBUG):
             logger.info(f"Created _CacheContext: {self.name}")
         # Some checks for settings
-        if self.do_separate_cfg:
+        if self.enable_spearate_cfg:
             assert self.enable_alter_cache is False, (
                 "enable_alter_cache must set as False if "
-                "do_separate_cfg is enabled."
+                "enable_spearate_cfg is enabled."
             )
             if self.cfg_diff_compute_separate:
                 assert self.cfg_compute_first is False, (
@@ -121,20 +120,17 @@ class _CachedContext:  # Internal CachedContext Impl class
                 self.max_warmup_steps if self.max_warmup_steps > 0 else 1
             )
 
-        # Only set n_derivatives as 2 or 3, which is enough for most cases.
-        if "n_derivatives" not in self.taylorseer_kwargs:
-            self.taylorseer_kwargs["n_derivatives"] = max(
-                2, min(3, self.taylorseer_kwargs["max_warmup_steps"])
-            )
+        # Overwrite the 'n_derivatives' by 'taylorseer_order', default: 2.
+        self.taylorseer_kwargs["n_derivatives"] = self.taylorseer_order
 
         if self.enable_taylorseer:
             self.taylorseer = TaylorSeer(**self.taylorseer_kwargs)
-            if self.do_separate_cfg:
+            if self.enable_spearate_cfg:
                 self.cfg_taylorseer = TaylorSeer(**self.taylorseer_kwargs)
 
         if self.enable_encoder_taylorseer:
             self.encoder_tarlorseer = TaylorSeer(**self.taylorseer_kwargs)
-            if self.do_separate_cfg:
+            if self.enable_spearate_cfg:
                 self.cfg_encoder_taylorseer = TaylorSeer(
                     **self.taylorseer_kwargs
                 )
@@ -181,7 +177,7 @@ class _CachedContext:  # Internal CachedContext Impl class
         # incr    step: prev 0 -> 1; prev 1 -> 2
         # current step: incr step - 1
         self.transformer_executed_steps += 1
-        if not self.do_separate_cfg:
+        if not self.enable_spearate_cfg:
             self.executed_steps += 1
         else:
             # 0,1 -> 0 + 1, 2,3 -> 1 + 1, ...
@@ -223,7 +219,7 @@ class _CachedContext:  # Internal CachedContext Impl class
 
         # mark_step_begin of TaylorSeer must be called after the cache is reset.
         if self.enable_taylorseer or self.enable_encoder_taylorseer:
-            if self.do_separate_cfg:
+            if self.enable_spearate_cfg:
                 # Assume non-CFG steps: 0, 2, 4, 6, ...
                 if not self.is_separate_cfg_step():
                     taylorseer, encoder_taylorseer = self.get_taylorseers()
@@ -318,7 +314,7 @@ class _CachedContext:  # Internal CachedContext Impl class
 
     @torch.compiler.disable
     def is_separate_cfg_step(self):
-        if not self.do_separate_cfg:
+        if not self.enable_spearate_cfg:
             return False
         if self.cfg_compute_first:
             # CFG steps: 0, 2, 4, 6, ...
@@ -329,861 +325,3 @@ class _CachedContext:  # Internal CachedContext Impl class
     @torch.compiler.disable
     def is_in_warmup(self):
         return self.get_current_step() < self.max_warmup_steps
-
-
-# TODO: Support context manager for different cache_context
-_current_cache_context: _CachedContext = None
-
-_cache_context_manager: Dict[str, _CachedContext] = {}
-
-
-def create_cache_context(*args, **kwargs):
-    global _cache_context_manager
-    _context = _CachedContext(*args, **kwargs)
-    _cache_context_manager[_context.name] = _context
-    return _context
-
-
-def get_cache_context():
-    return _current_cache_context
-
-
-def set_cache_context(cache_context: _CachedContext | str):
-    global _current_cache_context, _cache_context_manager
-    if isinstance(cache_context, _CachedContext):
-        _current_cache_context = cache_context
-    else:
-        _current_cache_context = _cache_context_manager[cache_context]
-
-
-def reset_cache_context(cache_context: _CachedContext | str, *args, **kwargs):
-    global _cache_context_manager
-    if isinstance(cache_context, _CachedContext):
-        old_context_name = cache_context.name
-        if cache_context.name in _cache_context_manager:
-            del _cache_context_manager[cache_context.name]
-        # force use old_context name
-        kwargs["name"] = old_context_name
-        _context = _CachedContext(*args, **kwargs)
-        _cache_context_manager[_context.name] = _context
-    else:
-        old_context_name = cache_context
-        if cache_context in _cache_context_manager:
-            del _cache_context_manager[cache_context]
-        # force use old_context name
-        kwargs["name"] = old_context_name
-        _context = _CachedContext(*args, **kwargs)
-        _cache_context_manager[_context.name] = _context
-
-    return _context
-
-
-@contextlib.contextmanager
-def cache_context(cache_context: _CachedContext | str):
-    global _current_cache_context, _cache_context_manager
-    old_cache_context = _current_cache_context
-    if isinstance(cache_context, _CachedContext):
-        _current_cache_context = cache_context
-    else:
-        _current_cache_context = _cache_context_manager[cache_context]
-    try:
-        yield
-    finally:
-        _current_cache_context = old_cache_context
-
-
-@torch.compiler.disable
-def get_residual_diff_threshold():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.get_residual_diff_threshold()
-
-
-@torch.compiler.disable
-def get_buffer(name):
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.get_buffer(name)
-
-
-@torch.compiler.disable
-def set_buffer(name, buffer):
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    cache_context.set_buffer(name, buffer)
-
-
-@torch.compiler.disable
-def remove_buffer(name):
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    cache_context.remove_buffer(name)
-
-
-@torch.compiler.disable
-def mark_step_begin():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    cache_context.mark_step_begin()
-
-
-@torch.compiler.disable
-def get_current_step():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.get_current_step()
-
-
-@torch.compiler.disable
-def get_current_step_residual_diff():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    step = str(get_current_step())
-    residual_diffs = get_residual_diffs()
-    if step in residual_diffs:
-        return residual_diffs[step]
-    return None
-
-
-@torch.compiler.disable
-def get_current_step_cfg_residual_diff():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    step = str(get_current_step())
-    cfg_residual_diffs = get_cfg_residual_diffs()
-    if step in cfg_residual_diffs:
-        return cfg_residual_diffs[step]
-    return None
-
-
-@torch.compiler.disable
-def get_current_transformer_step():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.get_current_transformer_step()
-
-
-@torch.compiler.disable
-def get_cached_steps():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.get_cached_steps()
-
-
-@torch.compiler.disable
-def get_cfg_cached_steps():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.get_cfg_cached_steps()
-
-
-@torch.compiler.disable
-def get_max_cached_steps():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.max_cached_steps
-
-
-@torch.compiler.disable
-def get_max_continuous_cached_steps():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.max_continuous_cached_steps
-
-
-@torch.compiler.disable
-def get_continuous_cached_steps():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.continuous_cached_steps
-
-
-@torch.compiler.disable
-def get_cfg_continuous_cached_steps():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.cfg_continuous_cached_steps
-
-
-@torch.compiler.disable
-def add_cached_step():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    cache_context.add_cached_step()
-
-
-@torch.compiler.disable
-def add_residual_diff(diff):
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    cache_context.add_residual_diff(diff)
-
-
-@torch.compiler.disable
-def get_residual_diffs():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.get_residual_diffs()
-
-
-@torch.compiler.disable
-def get_cfg_residual_diffs():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.get_cfg_residual_diffs()
-
-
-@torch.compiler.disable
-def is_taylorseer_enabled():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.enable_taylorseer
-
-
-@torch.compiler.disable
-def is_encoder_taylorseer_enabled():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.enable_encoder_taylorseer
-
-
-def get_taylorseers() -> Tuple[TaylorSeer, TaylorSeer]:
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.get_taylorseers()
-
-
-def get_cfg_taylorseers() -> Tuple[TaylorSeer, TaylorSeer]:
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.get_cfg_taylorseers()
-
-
-@torch.compiler.disable
-def is_taylorseer_cache_residual():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.taylorseer_cache_type == "residual"
-
-
-@torch.compiler.disable
-def is_cache_residual():
-    if is_taylorseer_enabled():
-        # residual or hidden_states
-        return is_taylorseer_cache_residual()
-    return True
-
-
-@torch.compiler.disable
-def is_encoder_cache_residual():
-    if is_encoder_taylorseer_enabled():
-        # residual or hidden_states
-        return is_taylorseer_cache_residual()
-    return True
-
-
-@torch.compiler.disable
-def is_alter_cache_enabled():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.enable_alter_cache
-
-
-@torch.compiler.disable
-def is_alter_cache():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.is_alter_cache
-
-
-@torch.compiler.disable
-def is_in_warmup():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.is_in_warmup()
-
-
-@torch.compiler.disable
-def is_l1_diff_enabled():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return (
-        cache_context.l1_hidden_states_diff_threshold is not None
-        and cache_context.l1_hidden_states_diff_threshold > 0.0
-    )
-
-
-@torch.compiler.disable
-def get_important_condition_threshold():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.important_condition_threshold
-
-
-@torch.compiler.disable
-def non_compute_blocks_diff_threshold():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.non_compute_blocks_diff_threshold
-
-
-@torch.compiler.disable
-def Fn_compute_blocks():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    assert (
-        cache_context.Fn_compute_blocks >= 1
-    ), "Fn_compute_blocks must be >= 1"
-    if cache_context.max_Fn_compute_blocks > 0:
-        # NOTE: Fn_compute_blocks can be 1, which means FB Cache
-        # but it must be less than or equal to max_Fn_compute_blocks
-        assert (
-            cache_context.Fn_compute_blocks
-            <= cache_context.max_Fn_compute_blocks
-        ), (
-            f"Fn_compute_blocks must be <= {cache_context.max_Fn_compute_blocks}, "
-            f"but got {cache_context.Fn_compute_blocks}"
-        )
-    return cache_context.Fn_compute_blocks
-
-
-@torch.compiler.disable
-def Fn_compute_blocks_ids():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    assert (
-        len(cache_context.Fn_compute_blocks_ids)
-        <= cache_context.Fn_compute_blocks
-    ), (
-        "The num of Fn_compute_blocks_ids must be <= Fn_compute_blocks "
-        f"{cache_context.Fn_compute_blocks}, but got "
-        f"{len(cache_context.Fn_compute_blocks_ids)}"
-    )
-    return cache_context.Fn_compute_blocks_ids
-
-
-@torch.compiler.disable
-def Bn_compute_blocks():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    assert (
-        cache_context.Bn_compute_blocks >= 0
-    ), "Bn_compute_blocks must be >= 0"
-    if cache_context.max_Bn_compute_blocks > 0:
-        # NOTE: Bn_compute_blocks can be 0, which means FB Cache
-        # but it must be less than or equal to max_Bn_compute_blocks
-        assert (
-            cache_context.Bn_compute_blocks
-            <= cache_context.max_Bn_compute_blocks
-        ), (
-            f"Bn_compute_blocks must be <= {cache_context.max_Bn_compute_blocks}, "
-            f"but got {cache_context.Bn_compute_blocks}"
-        )
-    return cache_context.Bn_compute_blocks
-
-
-@torch.compiler.disable
-def Bn_compute_blocks_ids():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    assert (
-        len(cache_context.Bn_compute_blocks_ids)
-        <= cache_context.Bn_compute_blocks
-    ), (
-        "The num of Bn_compute_blocks_ids must be <= Bn_compute_blocks "
-        f"{cache_context.Bn_compute_blocks}, but got "
-        f"{len(cache_context.Bn_compute_blocks_ids)}"
-    )
-    return cache_context.Bn_compute_blocks_ids
-
-
-@torch.compiler.disable
-def do_separate_cfg():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.do_separate_cfg
-
-
-@torch.compiler.disable
-def is_separate_cfg_step():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.is_separate_cfg_step()
-
-
-@torch.compiler.disable
-def cfg_diff_compute_separate():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.cfg_diff_compute_separate
-
-
-def collect_cache_kwargs(default_attrs: dict, **kwargs):
-    # NOTE: This API will split kwargs into cache_kwargs and other_kwargs
-    # default_attrs: specific settings for different pipelines
-    cache_attrs = dataclasses.fields(_CachedContext)
-    cache_attrs = [
-        attr
-        for attr in cache_attrs
-        if hasattr(
-            _CachedContext,
-            attr.name,
-        )
-    ]
-    cache_kwargs = {
-        attr.name: kwargs.pop(
-            attr.name,
-            getattr(_CachedContext, attr.name),
-        )
-        for attr in cache_attrs
-    }
-
-    def _safe_set_sequence_field(
-        field_name: str,
-        default_value: Any = None,
-    ):
-        if field_name not in cache_kwargs:
-            cache_kwargs[field_name] = kwargs.pop(
-                field_name,
-                default_value,
-            )
-
-    # Manually set sequence fields, namely, Fn_compute_blocks_ids
-    # and Bn_compute_blocks_ids, which are lists or sets.
-    _safe_set_sequence_field("Fn_compute_blocks_ids", [])
-    _safe_set_sequence_field("Bn_compute_blocks_ids", [])
-    _safe_set_sequence_field("taylorseer_kwargs", {})
-
-    for attr in cache_attrs:
-        if attr.name in default_attrs:  # can be empty {}
-            cache_kwargs[attr.name] = default_attrs[attr.name]
-
-    if logger.isEnabledFor(logging.DEBUG):
-        logger.debug(f"Collected DBCache kwargs: {cache_kwargs}")
-
-    return cache_kwargs, kwargs
-
-
-@torch.compiler.disable
-def are_two_tensors_similar(
-    t1: torch.Tensor,  # prev residual R(t-1,n) = H(t-1,n) - H(t-1,0)
-    t2: torch.Tensor,  # curr residual R(t  ,n) = H(t  ,n) - H(t  ,0)
-    *,
-    threshold: float,
-    parallelized: bool = False,
-    prefix: str = "Fn",  # for debugging
-):
-    # Special case for threshold, 0.0 means the threshold is disabled, -1.0 means
-    # the threshold is always enabled, -2.0 means the shape is not matched.
-    if threshold <= 0.0:
-        add_residual_diff(-0.0)
-        return False
-
-    if threshold >= 1.0:
-        # If threshold is 1.0 or more, we consider them always similar.
-        add_residual_diff(-1.0)
-        return True
-
-    if t1.shape != t2.shape:
-        if logger.isEnabledFor(logging.DEBUG):
-            logger.debug(f"{prefix}, shape error: {t1.shape} != {t2.shape}")
-        add_residual_diff(-2.0)
-        return False
-
-    if all(
-        (
-            do_separate_cfg(),
-            is_separate_cfg_step(),
-            not cfg_diff_compute_separate(),
-            get_current_step_residual_diff() is not None,
-        )
-    ):
-        # Reuse computed diff value from non-CFG step
-        diff = get_current_step_residual_diff()
-    else:
-        # Find the most significant token through t1 and t2, and
-        # consider the diff of the significant token. The more significant,
-        # the more important.
-        condition_thresh = get_important_condition_threshold()
-        if condition_thresh > 0.0:
-            raw_diff = (t1 - t2).abs()  # [B, seq_len, d]
-            token_m_df = raw_diff.mean(dim=-1)  # [B, seq_len]
-            token_m_t1 = t1.abs().mean(dim=-1)  # [B, seq_len]
-            # D = (t1 - t2) / t1 = 1 - (t2 / t1), if D = 0, then t1 = t2.
-            token_diff = token_m_df / token_m_t1  # [B, seq_len]
-            condition = token_diff > condition_thresh  # [B, seq_len]
-            if condition.sum() > 0:
-                condition = condition.unsqueeze(-1)  # [B, seq_len, 1]
-                condition = condition.expand_as(raw_diff)  # [B, seq_len, d]
-                mean_diff = raw_diff[condition].mean()
-                mean_t1 = t1[condition].abs().mean()
-            else:
-                mean_diff = (t1 - t2).abs().mean()
-                mean_t1 = t1.abs().mean()
-        else:
-            # Use the mean of the absolute difference of the tensors
-            mean_diff = (t1 - t2).abs().mean()
-            mean_t1 = t1.abs().mean()
-
-        if parallelized:
-            dist.all_reduce(mean_diff, op=dist.ReduceOp.AVG)
-            dist.all_reduce(mean_t1, op=dist.ReduceOp.AVG)
-
-        # D = (t1 - t2) / t1 = 1 - (t2 / t1), if D = 0, then t1 = t2.
-        # Futher, if we assume that (H(t,  0) - H(t-1,0)) ~ 0, then,
-        # H(t-1,n) ~ H(t  ,n), which means the hidden states are similar.
-        diff = (mean_diff / mean_t1).item()
-
-    if logger.isEnabledFor(logging.DEBUG):
-        logger.debug(f"{prefix}, diff: {diff:.6f}, threshold: {threshold:.6f}")
-
-    add_residual_diff(diff)
-
-    return diff < threshold
-
-
-@torch.compiler.disable
-def _debugging_set_buffer(prefix):
-    if logger.isEnabledFor(logging.DEBUG):
-        logger.debug(
-            f"set {prefix}, "
-            f"transformer step: {get_current_transformer_step()}, "
-            f"executed step: {get_current_step()}"
-        )
-
-
-@torch.compiler.disable
-def _debugging_get_buffer(prefix):
-    if logger.isEnabledFor(logging.DEBUG):
-        logger.debug(
-            f"get {prefix}, "
-            f"transformer step: {get_current_transformer_step()}, "
-            f"executed step: {get_current_step()}"
-        )
-
-
-# Fn buffers
-@torch.compiler.disable
-def set_Fn_buffer(buffer: torch.Tensor, prefix: str = "Fn"):
-    # Set hidden_states or residual for Fn blocks.
-    # This buffer is only use for L1 diff calculation.
-    downsample_factor = get_downsample_factor()
-    if downsample_factor > 1:
-        buffer = buffer[..., ::downsample_factor]
-        buffer = buffer.contiguous()
-    if is_separate_cfg_step():
-        _debugging_set_buffer(f"{prefix}_buffer_cfg")
-        set_buffer(f"{prefix}_buffer_cfg", buffer)
-    else:
-        _debugging_set_buffer(f"{prefix}_buffer")
-        set_buffer(f"{prefix}_buffer", buffer)
-
-
-@torch.compiler.disable
-def get_Fn_buffer(prefix: str = "Fn"):
-    if is_separate_cfg_step():
-        _debugging_get_buffer(f"{prefix}_buffer_cfg")
-        return get_buffer(f"{prefix}_buffer_cfg")
-    _debugging_get_buffer(f"{prefix}_buffer")
-    return get_buffer(f"{prefix}_buffer")
-
-
-@torch.compiler.disable
-def set_Fn_encoder_buffer(buffer: torch.Tensor, prefix: str = "Fn"):
-    if is_separate_cfg_step():
-        _debugging_set_buffer(f"{prefix}_encoder_buffer_cfg")
-        set_buffer(f"{prefix}_encoder_buffer_cfg", buffer)
-    else:
-        _debugging_set_buffer(f"{prefix}_encoder_buffer")
-        set_buffer(f"{prefix}_encoder_buffer", buffer)
-
-
-@torch.compiler.disable
-def get_Fn_encoder_buffer(prefix: str = "Fn"):
-    if is_separate_cfg_step():
-        _debugging_get_buffer(f"{prefix}_encoder_buffer_cfg")
-        return get_buffer(f"{prefix}_encoder_buffer_cfg")
-    _debugging_get_buffer(f"{prefix}_encoder_buffer")
-    return get_buffer(f"{prefix}_encoder_buffer")
-
-
-# Bn buffers
-@torch.compiler.disable
-def set_Bn_buffer(buffer: torch.Tensor, prefix: str = "Bn"):
-    # Set hidden_states or residual for Bn blocks.
-    # This buffer is use for hidden states approximation.
-    if is_taylorseer_enabled():
-        # taylorseer, encoder_taylorseer
-        if is_separate_cfg_step():
-            taylorseer, _ = get_cfg_taylorseers()
-        else:
-            taylorseer, _ = get_taylorseers()
-
-        if taylorseer is not None:
-            # Use TaylorSeer to update the buffer
-            taylorseer.update(buffer)
-        else:
-            if logger.isEnabledFor(logging.DEBUG):
-                logger.debug(
-                    "TaylorSeer is enabled but not set in the cache context. "
-                    "Falling back to default buffer retrieval."
-                )
-            if is_separate_cfg_step():
-                _debugging_set_buffer(f"{prefix}_buffer_cfg")
-                set_buffer(f"{prefix}_buffer_cfg", buffer)
-            else:
-                _debugging_set_buffer(f"{prefix}_buffer")
-                set_buffer(f"{prefix}_buffer", buffer)
-    else:
-        if is_separate_cfg_step():
-            _debugging_set_buffer(f"{prefix}_buffer_cfg")
-            set_buffer(f"{prefix}_buffer_cfg", buffer)
-        else:
-            _debugging_set_buffer(f"{prefix}_buffer")
-            set_buffer(f"{prefix}_buffer", buffer)
-
-
-@torch.compiler.disable
-def get_Bn_buffer(prefix: str = "Bn"):
-    if is_taylorseer_enabled():
-        # taylorseer, encoder_taylorseer
-        if is_separate_cfg_step():
-            taylorseer, _ = get_cfg_taylorseers()
-        else:
-            taylorseer, _ = get_taylorseers()
-
-        if taylorseer is not None:
-            return taylorseer.approximate_value()
-        else:
-            if logger.isEnabledFor(logging.DEBUG):
-                logger.debug(
-                    "TaylorSeer is enabled but not set in the cache context. "
-                    "Falling back to default buffer retrieval."
-                )
-            # Fallback to default buffer retrieval
-            if is_separate_cfg_step():
-                _debugging_get_buffer(f"{prefix}_buffer_cfg")
-                return get_buffer(f"{prefix}_buffer_cfg")
-            _debugging_get_buffer(f"{prefix}_buffer")
-            return get_buffer(f"{prefix}_buffer")
-    else:
-        if is_separate_cfg_step():
-            _debugging_get_buffer(f"{prefix}_buffer_cfg")
-            return get_buffer(f"{prefix}_buffer_cfg")
-        _debugging_get_buffer(f"{prefix}_buffer")
-        return get_buffer(f"{prefix}_buffer")
-
-
-@torch.compiler.disable
-def set_Bn_encoder_buffer(buffer: torch.Tensor | None, prefix: str = "Bn"):
-    # DON'T set None Buffer
-    if buffer is None:
-        return
-
-    # This buffer is use for encoder hidden states approximation.
-    if is_encoder_taylorseer_enabled():
-        # taylorseer, encoder_taylorseer
-        if is_separate_cfg_step():
-            _, encoder_taylorseer = get_cfg_taylorseers()
-        else:
-            _, encoder_taylorseer = get_taylorseers()
-
-        if encoder_taylorseer is not None:
-            # Use TaylorSeer to update the buffer
-            encoder_taylorseer.update(buffer)
-        else:
-            if logger.isEnabledFor(logging.DEBUG):
-                logger.debug(
-                    "TaylorSeer is enabled but not set in the cache context. "
-                    "Falling back to default buffer retrieval."
-                )
-            if is_separate_cfg_step():
-                _debugging_set_buffer(f"{prefix}_encoder_buffer_cfg")
-                set_buffer(f"{prefix}_encoder_buffer_cfg", buffer)
-            else:
-                _debugging_set_buffer(f"{prefix}_encoder_buffer")
-                set_buffer(f"{prefix}_encoder_buffer", buffer)
-    else:
-        if is_separate_cfg_step():
-            _debugging_set_buffer(f"{prefix}_encoder_buffer_cfg")
-            set_buffer(f"{prefix}_encoder_buffer_cfg", buffer)
-        else:
-            _debugging_set_buffer(f"{prefix}_encoder_buffer")
-            set_buffer(f"{prefix}_encoder_buffer", buffer)
-
-
-@torch.compiler.disable
-def get_Bn_encoder_buffer(prefix: str = "Bn"):
-    if is_encoder_taylorseer_enabled():
-        if is_separate_cfg_step():
-            _, encoder_taylorseer = get_cfg_taylorseers()
-        else:
-            _, encoder_taylorseer = get_taylorseers()
-
-        if encoder_taylorseer is not None:
-            # Use TaylorSeer to approximate the value
-            return encoder_taylorseer.approximate_value()
-        else:
-            if logger.isEnabledFor(logging.DEBUG):
-                logger.debug(
-                    "TaylorSeer is enabled but not set in the cache context. "
-                    "Falling back to default buffer retrieval."
-                )
-            # Fallback to default buffer retrieval
-            if is_separate_cfg_step():
-                _debugging_get_buffer(f"{prefix}_encoder_buffer_cfg")
-                return get_buffer(f"{prefix}_encoder_buffer_cfg")
-            _debugging_get_buffer(f"{prefix}_encoder_buffer")
-            return get_buffer(f"{prefix}_encoder_buffer")
-    else:
-        if is_separate_cfg_step():
-            _debugging_get_buffer(f"{prefix}_encoder_buffer_cfg")
-            return get_buffer(f"{prefix}_encoder_buffer_cfg")
-        _debugging_get_buffer(f"{prefix}_encoder_buffer")
-        return get_buffer(f"{prefix}_encoder_buffer")
-
-
-@torch.compiler.disable
-def apply_hidden_states_residual(
-    hidden_states: torch.Tensor,
-    encoder_hidden_states: torch.Tensor = None,
-    prefix: str = "Bn",
-    encoder_prefix: str = "Bn_encoder",
-):
-    # Allow Bn and Fn prefix to be used for residual cache.
-    if "Bn" in prefix:
-        hidden_states_prev = get_Bn_buffer(prefix)
-    else:
-        hidden_states_prev = get_Fn_buffer(prefix)
-
-    assert hidden_states_prev is not None, f"{prefix}_buffer must be set before"
-
-    if is_cache_residual():
-        hidden_states = hidden_states_prev + hidden_states
-    else:
-        # If cache is not residual, we use the hidden states directly
-        hidden_states = hidden_states_prev
-
-    hidden_states = hidden_states.contiguous()
-
-    if encoder_hidden_states is not None:
-        if "Bn" in encoder_prefix:
-            encoder_hidden_states_prev = get_Bn_encoder_buffer(encoder_prefix)
-        else:
-            encoder_hidden_states_prev = get_Fn_encoder_buffer(encoder_prefix)
-
-        assert (
-            encoder_hidden_states_prev is not None
-        ), f"{prefix}_encoder_buffer must be set before"
-
-        if is_encoder_cache_residual():
-            encoder_hidden_states = (
-                encoder_hidden_states_prev + encoder_hidden_states
-            )
-        else:
-            # If encoder cache is not residual, we use the encoder hidden states directly
-            encoder_hidden_states = encoder_hidden_states_prev
-
-        encoder_hidden_states = encoder_hidden_states.contiguous()
-
-    return hidden_states, encoder_hidden_states
-
-
-@torch.compiler.disable
-def get_downsample_factor():
-    cache_context = get_cache_context()
-    assert cache_context is not None, "cache_context must be set before"
-    return cache_context.downsample_factor
-
-
-@torch.compiler.disable
-def get_can_use_cache(
-    states_tensor: torch.Tensor,  # hidden_states or residual
-    parallelized: bool = False,
-    threshold: Optional[float] = None,  # can manually set threshold
-    prefix: str = "Fn",
-):
-    if is_in_warmup():
-        return False
-
-    # max cached steps
-    max_cached_steps = get_max_cached_steps()
-    if not is_separate_cfg_step():
-        cached_steps = get_cached_steps()
-    else:
-        cached_steps = get_cfg_cached_steps()
-
-    if max_cached_steps >= 0 and (len(cached_steps) >= max_cached_steps):
-        if logger.isEnabledFor(logging.DEBUG):
-            logger.debug(
-                f"{prefix}, max_cached_steps reached: {max_cached_steps}, "
-                "can not use cache."
-            )
-        return False
-
-    # max continuous cached steps
-    max_continuous_cached_steps = get_max_continuous_cached_steps()
-    if not is_separate_cfg_step():
-        continuous_cached_steps = get_continuous_cached_steps()
-    else:
-        continuous_cached_steps = get_cfg_continuous_cached_steps()
-
-    if max_continuous_cached_steps >= 0 and (
-        continuous_cached_steps >= max_continuous_cached_steps
-    ):
-        if logger.isEnabledFor(logging.DEBUG):
-            logger.debug(
-                f"{prefix}, max_continuous_cached_steps "
-                f"reached: {max_continuous_cached_steps}, "
-                "can not use cache."
-            )
-        # reset continuous cached steps stats
-        cache_context = get_cache_context()
-        if not is_separate_cfg_step():
-            cache_context.continuous_cached_steps = 0
-        else:
-            cache_context.cfg_continuous_cached_steps = 0
-        return False
-
-    if threshold is None or threshold <= 0.0:
-        threshold = get_residual_diff_threshold()
-    if threshold <= 0.0:
-        return False
-
-    downsample_factor = get_downsample_factor()
-    if downsample_factor > 1 and "Bn" not in prefix:
-        states_tensor = states_tensor[..., ::downsample_factor]
-        states_tensor = states_tensor.contiguous()
-
-    # Allow Bn and Fn prefix to be used for diff calculation.
-    if "Bn" in prefix:
-        prev_states_tensor = get_Bn_buffer(prefix)
-    else:
-        prev_states_tensor = get_Fn_buffer(prefix)
-
-    if not is_alter_cache_enabled():
-        # Dynamic cache according to the residual diff
-        can_use_cache = (
-            prev_states_tensor is not None
-            and are_two_tensors_similar(
-                prev_states_tensor,
-                states_tensor,
-                threshold=threshold,
-                parallelized=parallelized,
-                prefix=prefix,
-            )
-        )
-    else:
-        # Only cache in the alter cache steps
-        can_use_cache = (
-            prev_states_tensor is not None
-            and are_two_tensors_similar(
-                prev_states_tensor,
-                states_tensor,
-                threshold=threshold,
-                parallelized=parallelized,
-                prefix=prefix,
-            )
-            and is_alter_cache()
-        )
-    return can_use_cache