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

cache_dit/cache_factory/cache_blocks/pattern_base.py

@@ -2,7 +2,10 @@ import inspect
 import torch
 import torch.distributed as dist
 
-from cache_dit.cache_factory import CachedContext
+from cache_dit.cache_factory.cache_contexts.cache_context import CachedContext
+from cache_dit.cache_factory.cache_contexts.cache_manager import (
+    CachedContextManager,
+)
 from cache_dit.cache_factory import ForwardPattern
 from cache_dit.logger import init_logger
 
@@ -18,29 +21,34 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
 
     def __init__(
         self,
+        # 0. Transformer blocks configuration
         transformer_blocks: torch.nn.ModuleList,
-        # 'transformer_blocks', 'blocks', 'single_transformer_blocks',
-        # 'layers', 'single_stream_blocks', 'double_stream_blocks'
-        blocks_name: str,
-        # Usually, blocks_name, etc.
-        cache_context: str,
-        *,
         transformer: torch.nn.Module = None,
         forward_pattern: ForwardPattern = ForwardPattern.Pattern_0,
         check_num_outputs: bool = True,
+        # 1. Cache context configuration
+        cache_prefix: str = None,  # maybe un-need.
+        cache_context: CachedContext | str = None,
+        cache_manager: CachedContextManager = None,
+        **kwargs,
     ):
         super().__init__()
 
+        # 0. Transformer blocks configuration
         self.transformer = transformer
         self.transformer_blocks = transformer_blocks
-        self.blocks_name = blocks_name
-        self.cache_context = cache_context
         self.forward_pattern = forward_pattern
         self.check_num_outputs = check_num_outputs
+        # 1. Cache context configuration
+        self.cache_prefix = cache_prefix
+        self.cache_context = cache_context
+        self.cache_manager = cache_manager
+
         self._check_forward_pattern()
         logger.info(
             f"Match Cached Blocks: {self.__class__.__name__}, for "
-            f"{self.blocks_name}, context: {self.cache_context}"
+            f"{self.cache_prefix}, cache_context: {self.cache_context}, "
+            f"cache_manager: {self.cache_manager.name}."
         )
 
     def _check_forward_pattern(self):
@@ -79,9 +87,7 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
         *args,
         **kwargs,
     ):
-        CachedContext.set_cache_context(
-            self.cache_context,
-        )
+        self.cache_manager.set_context(self.cache_context)
 
         original_hidden_states = hidden_states
         # Call first `n` blocks to process the hidden states for
@@ -96,39 +102,39 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
         Fn_hidden_states_residual = hidden_states - original_hidden_states
         del original_hidden_states
 
-        CachedContext.mark_step_begin()
+        self.cache_manager.mark_step_begin()
         # Residual L1 diff or Hidden States L1 diff
-        can_use_cache = CachedContext.get_can_use_cache(
+        can_use_cache = self.cache_manager.can_cache(
             (
                 Fn_hidden_states_residual
-                if not CachedContext.is_l1_diff_enabled()
+                if not self.cache_manager.is_l1_diff_enabled()
                 else hidden_states
             ),
             parallelized=self._is_parallelized(),
             prefix=(
-                f"{self.blocks_name}_Fn_residual"
-                if not CachedContext.is_l1_diff_enabled()
-                else f"{self.blocks_name}_Fn_hidden_states"
+                f"{self.cache_prefix}_Fn_residual"
+                if not self.cache_manager.is_l1_diff_enabled()
+                else f"{self.cache_prefix}_Fn_hidden_states"
             ),
         )
 
         torch._dynamo.graph_break()
         if can_use_cache:
-            CachedContext.add_cached_step()
+            self.cache_manager.add_cached_step()
             del Fn_hidden_states_residual
             hidden_states, encoder_hidden_states = (
-                CachedContext.apply_hidden_states_residual(
+                self.cache_manager.apply_cache(
                     hidden_states,
                     encoder_hidden_states,
                     prefix=(
-                        f"{self.blocks_name}_Bn_residual"
-                        if CachedContext.is_cache_residual()
-                        else f"{self.blocks_name}_Bn_hidden_states"
+                        f"{self.cache_prefix}_Bn_residual"
+                        if self.cache_manager.is_cache_residual()
+                        else f"{self.cache_prefix}_Bn_hidden_states"
                     ),
                     encoder_prefix=(
-                        f"{self.blocks_name}_Bn_residual"
-                        if CachedContext.is_encoder_cache_residual()
-                        else f"{self.blocks_name}_Bn_hidden_states"
+                        f"{self.cache_prefix}_Bn_residual"
+                        if self.cache_manager.is_encoder_cache_residual()
+                        else f"{self.cache_prefix}_Bn_hidden_states"
                     ),
                 )
             )
@@ -142,15 +148,15 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
                 **kwargs,
             )
         else:
-            CachedContext.set_Fn_buffer(
+            self.cache_manager.set_Fn_buffer(
                 Fn_hidden_states_residual,
-                prefix=f"{self.blocks_name}_Fn_residual",
+                prefix=f"{self.cache_prefix}_Fn_residual",
             )
-            if CachedContext.is_l1_diff_enabled():
+            if self.cache_manager.is_l1_diff_enabled():
                 # for hidden states L1 diff
-                CachedContext.set_Fn_buffer(
+                self.cache_manager.set_Fn_buffer(
                     hidden_states,
-                    f"{self.blocks_name}_Fn_hidden_states",
+                    f"{self.cache_prefix}_Fn_hidden_states",
                 )
             del Fn_hidden_states_residual
             torch._dynamo.graph_break()
@@ -166,27 +172,27 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
                 **kwargs,
             )
             torch._dynamo.graph_break()
-            if CachedContext.is_cache_residual():
-                CachedContext.set_Bn_buffer(
+            if self.cache_manager.is_cache_residual():
+                self.cache_manager.set_Bn_buffer(
                     hidden_states_residual,
-                    prefix=f"{self.blocks_name}_Bn_residual",
+                    prefix=f"{self.cache_prefix}_Bn_residual",
                 )
             else:
                 # TaylorSeer
-                CachedContext.set_Bn_buffer(
+                self.cache_manager.set_Bn_buffer(
                     hidden_states,
-                    prefix=f"{self.blocks_name}_Bn_hidden_states",
+                    prefix=f"{self.cache_prefix}_Bn_hidden_states",
                 )
-            if CachedContext.is_encoder_cache_residual():
-                CachedContext.set_Bn_encoder_buffer(
+            if self.cache_manager.is_encoder_cache_residual():
+                self.cache_manager.set_Bn_encoder_buffer(
                     encoder_hidden_states_residual,
-                    prefix=f"{self.blocks_name}_Bn_residual",
+                    prefix=f"{self.cache_prefix}_Bn_residual",
                 )
             else:
                 # TaylorSeer
-                CachedContext.set_Bn_encoder_buffer(
+                self.cache_manager.set_Bn_encoder_buffer(
                     encoder_hidden_states,
-                    prefix=f"{self.blocks_name}_Bn_hidden_states",
+                    prefix=f"{self.cache_prefix}_Bn_hidden_states",
                 )
             torch._dynamo.graph_break()
             # Call last `n` blocks to further process the hidden states
@@ -232,10 +238,11 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
         # If so, we can skip some Bn blocks and directly
         # use the cached values.
         return (
-            CachedContext.get_current_step() in CachedContext.get_cached_steps()
+            self.cache_manager.get_current_step()
+            in self.cache_manager.get_cached_steps()
         ) or (
-            CachedContext.get_current_step()
-            in CachedContext.get_cfg_cached_steps()
+            self.cache_manager.get_current_step()
+            in self.cache_manager.get_cfg_cached_steps()
         )
 
     @torch.compiler.disable
@@ -244,20 +251,20 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
         # more stable diff calculation.
         # Fn: [0,...,n-1]
         selected_Fn_blocks = self.transformer_blocks[
-            : CachedContext.Fn_compute_blocks()
+            : self.cache_manager.Fn_compute_blocks()
         ]
         return selected_Fn_blocks
 
     @torch.compiler.disable
     def _Mn_blocks(self):  # middle blocks
         # M(N-2n): only transformer_blocks [n,...,N-n], middle
-        if CachedContext.Bn_compute_blocks() == 0:  # WARN: x[:-0] = []
+        if self.cache_manager.Bn_compute_blocks() == 0:  # WARN: x[:-0] = []
             selected_Mn_blocks = self.transformer_blocks[
-                CachedContext.Fn_compute_blocks() :
+                self.cache_manager.Fn_compute_blocks() :
             ]
         else:
             selected_Mn_blocks = self.transformer_blocks[
-                CachedContext.Fn_compute_blocks() : -CachedContext.Bn_compute_blocks()
+                self.cache_manager.Fn_compute_blocks() : -self.cache_manager.Bn_compute_blocks()
             ]
         return selected_Mn_blocks
 
@@ -265,7 +272,7 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
     def _Bn_blocks(self):
         # Bn: transformer_blocks [N-n+1,...,N-1]
         selected_Bn_blocks = self.transformer_blocks[
-            -CachedContext.Bn_compute_blocks() :
+            -self.cache_manager.Bn_compute_blocks() :
         ]
         return selected_Bn_blocks
 
@@ -276,10 +283,10 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
         *args,
         **kwargs,
     ):
-        assert CachedContext.Fn_compute_blocks() <= len(
+        assert self.cache_manager.Fn_compute_blocks() <= len(
             self.transformer_blocks
         ), (
-            f"Fn_compute_blocks {CachedContext.Fn_compute_blocks()} must be less than "
+            f"Fn_compute_blocks {self.cache_manager.Fn_compute_blocks()} must be less than "
             f"the number of transformer blocks {len(self.transformer_blocks)}"
         )
         for block in self._Fn_blocks():
@@ -376,7 +383,7 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
                     )
             # Cache residuals for the non-compute Bn blocks for
             # subsequent cache steps.
-            if block_id not in CachedContext.Bn_compute_blocks_ids():
+            if block_id not in self.cache_manager.Bn_compute_blocks_ids():
                 Bn_i_hidden_states_residual = (
                     hidden_states - Bn_i_original_hidden_states
                 )
@@ -385,22 +392,22 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
                 )
 
                 # Save original_hidden_states for diff calculation.
-                CachedContext.set_Bn_buffer(
+                self.cache_manager.set_Bn_buffer(
                     Bn_i_original_hidden_states,
-                    prefix=f"{self.blocks_name}_Bn_{block_id}_original",
+                    prefix=f"{self.cache_prefix}_Bn_{block_id}_original",
                 )
-                CachedContext.set_Bn_encoder_buffer(
+                self.cache_manager.set_Bn_encoder_buffer(
                     Bn_i_original_encoder_hidden_states,
-                    prefix=f"{self.blocks_name}_Bn_{block_id}_original",
+                    prefix=f"{self.cache_prefix}_Bn_{block_id}_original",
                 )
 
-                CachedContext.set_Bn_buffer(
+                self.cache_manager.set_Bn_buffer(
                     Bn_i_hidden_states_residual,
-                    prefix=f"{self.blocks_name}_Bn_{block_id}_residual",
+                    prefix=f"{self.cache_prefix}_Bn_{block_id}_residual",
                 )
-                CachedContext.set_Bn_encoder_buffer(
+                self.cache_manager.set_Bn_encoder_buffer(
                     Bn_i_encoder_hidden_states_residual,
-                    prefix=f"{self.blocks_name}_Bn_{block_id}_residual",
+                    prefix=f"{self.cache_prefix}_Bn_{block_id}_residual",
                 )
                 del Bn_i_hidden_states_residual
                 del Bn_i_encoder_hidden_states_residual
@@ -411,7 +418,7 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
         else:
             # Cache steps: Reuse the cached residuals.
             # Check if the block is in the Bn_compute_blocks_ids.
-            if block_id in CachedContext.Bn_compute_blocks_ids():
+            if block_id in self.cache_manager.Bn_compute_blocks_ids():
                 hidden_states = block(
                     hidden_states,
                     encoder_hidden_states,
@@ -429,25 +436,25 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
                 # Skip the block if it is not in the Bn_compute_blocks_ids.
                 # Use the cached residuals instead.
                 # Check if can use the cached residuals.
-                if CachedContext.get_can_use_cache(
+                if self.cache_manager.can_cache(
                     hidden_states,  # curr step
                     parallelized=self._is_parallelized(),
-                    threshold=CachedContext.non_compute_blocks_diff_threshold(),
-                    prefix=f"{self.blocks_name}_Bn_{block_id}_original",  # prev step
+                    threshold=self.cache_manager.non_compute_blocks_diff_threshold(),
+                    prefix=f"{self.cache_prefix}_Bn_{block_id}_original",  # prev step
                 ):
                     hidden_states, encoder_hidden_states = (
-                        CachedContext.apply_hidden_states_residual(
+                        self.cache_manager.apply_cache(
                             hidden_states,
                             encoder_hidden_states,
                             prefix=(
-                                f"{self.blocks_name}_Bn_{block_id}_residual"
-                                if CachedContext.is_cache_residual()
-                                else f"{self.blocks_name}_Bn_{block_id}_original"
+                                f"{self.cache_prefix}_Bn_{block_id}_residual"
+                                if self.cache_manager.is_cache_residual()
+                                else f"{self.cache_prefix}_Bn_{block_id}_original"
                             ),
                             encoder_prefix=(
-                                f"{self.blocks_name}_Bn_{block_id}_residual"
-                                if CachedContext.is_encoder_cache_residual()
-                                else f"{self.blocks_name}_Bn_{block_id}_original"
+                                f"{self.cache_prefix}_Bn_{block_id}_residual"
+                                if self.cache_manager.is_encoder_cache_residual()
+                                else f"{self.cache_prefix}_Bn_{block_id}_original"
                             ),
                         )
                     )
@@ -474,16 +481,16 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
         *args,
         **kwargs,
     ):
-        if CachedContext.Bn_compute_blocks() == 0:
+        if self.cache_manager.Bn_compute_blocks() == 0:
             return hidden_states, encoder_hidden_states
 
-        assert CachedContext.Bn_compute_blocks() <= len(
+        assert self.cache_manager.Bn_compute_blocks() <= len(
             self.transformer_blocks
         ), (
-            f"Bn_compute_blocks {CachedContext.Bn_compute_blocks()} must be less than "
+            f"Bn_compute_blocks {self.cache_manager.Bn_compute_blocks()} must be less than "
             f"the number of transformer blocks {len(self.transformer_blocks)}"
         )
-        if len(CachedContext.Bn_compute_blocks_ids()) > 0:
+        if len(self.cache_manager.Bn_compute_blocks_ids()) > 0:
             for i, block in enumerate(self._Bn_blocks()):
                 hidden_states, encoder_hidden_states = (
                     self._compute_or_cache_block(

cache_dit/cache_factory/cache_blocks/utils.py

@@ -2,22 +2,24 @@ import torch
 
 from typing import Any
 from cache_dit.cache_factory import CachedContext
+from cache_dit.cache_factory import CachedContextManager
 
 
-@torch.compiler.disable
 def patch_cached_stats(
-    module: torch.nn.Module | Any, cache_context: str = None
+    module: torch.nn.Module | Any,
+    cache_context: CachedContext | str = None,
+    cache_manager: CachedContextManager = None,
 ):
     # Patch the cached stats to the module, the cached stats
     # will be reset for each calling of pipe.__call__(**kwargs).
-    if module is None:
+    if module is None or cache_manager is None:
         return
 
     if cache_context is not None:
-        CachedContext.set_cache_context(cache_context)
+        cache_manager.set_context(cache_context)
 
     # TODO: Patch more cached stats to the module
-    module._cached_steps = CachedContext.get_cached_steps()
-    module._residual_diffs = CachedContext.get_residual_diffs()
-    module._cfg_cached_steps = CachedContext.get_cfg_cached_steps()
-    module._cfg_residual_diffs = CachedContext.get_cfg_residual_diffs()
+    module._cached_steps = cache_manager.get_cached_steps()
+    module._residual_diffs = cache_manager.get_residual_diffs()
+    module._cfg_cached_steps = cache_manager.get_cfg_cached_steps()
+    module._cfg_residual_diffs = cache_manager.get_cfg_residual_diffs()

cache_dit/cache_factory/cache_contexts/__init__.py

@@ -1,2 +1,5 @@
 # namespace alias: for _CachedContext and many others' cache context funcs.
-import cache_dit.cache_factory.cache_contexts.cache_context as CachedContext
+from cache_dit.cache_factory.cache_contexts.cache_context import CachedContext
+from cache_dit.cache_factory.cache_contexts.cache_manager import (
+    CachedContextManager,
+)