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

cache_dit/cache_factory/cache_blocks/pattern_base.py

@@ -2,17 +2,17 @@ import inspect
 import torch
 import torch.distributed as dist
 
-from cache_dit.cache_factory import cache_context
-from cache_dit.cache_factory import ForwardPattern
-from cache_dit.cache_factory.cache_blocks.utils import (
-    patch_cached_stats,
+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
 
 logger = init_logger(__name__)
 
 
-class DBCachedBlocks_Pattern_Base(torch.nn.Module):
+class CachedBlocks_Pattern_Base(torch.nn.Module):
     _supported_patterns = [
         ForwardPattern.Pattern_0,
         ForwardPattern.Pattern_1,
@@ -21,18 +21,35 @@ class DBCachedBlocks_Pattern_Base(torch.nn.Module):
 
     def __init__(
         self,
+        # 0. Transformer blocks configuration
         transformer_blocks: torch.nn.ModuleList,
-        *,
         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.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__}")
+        logger.info(
+            f"Match Cached Blocks: {self.__class__.__name__}, for "
+            f"{self.cache_prefix}, cache_context: {self.cache_context}, "
+            f"cache_manager: {self.cache_manager.name}."
+        )
 
     def _check_forward_pattern(self):
         assert (
@@ -45,16 +62,18 @@ class DBCachedBlocks_Pattern_Base(torch.nn.Module):
                 forward_parameters = set(
                     inspect.signature(block.forward).parameters.keys()
                 )
-                num_outputs = str(
-                    inspect.signature(block.forward).return_annotation
-                ).count("torch.Tensor")
-
-                if num_outputs > 0:
-                    assert len(self.forward_pattern.Out) == num_outputs, (
-                        f"The number of block's outputs is {num_outputs} don't not "
-                        f"match the number of the pattern: {self.forward_pattern}, "
-                        f"Out: {len(self.forward_pattern.Out)}."
-                    )
+
+                if self.check_num_outputs:
+                    num_outputs = str(
+                        inspect.signature(block.forward).return_annotation
+                    ).count("torch.Tensor")
+
+                    if num_outputs > 0:
+                        assert len(self.forward_pattern.Out) == num_outputs, (
+                            f"The number of block's outputs is {num_outputs} don't not "
+                            f"match the number of the pattern: {self.forward_pattern}, "
+                            f"Out: {len(self.forward_pattern.Out)}."
+                        )
 
                 for required_param in self.forward_pattern.In:
                     assert (
@@ -68,6 +87,8 @@ class DBCachedBlocks_Pattern_Base(torch.nn.Module):
         *args,
         **kwargs,
     ):
+        self.cache_manager.set_context(self.cache_context)
+
         original_hidden_states = hidden_states
         # Call first `n` blocks to process the hidden states for
         # more stable diff calculation.
@@ -81,39 +102,39 @@ class DBCachedBlocks_Pattern_Base(torch.nn.Module):
         Fn_hidden_states_residual = hidden_states - original_hidden_states
         del original_hidden_states
 
-        cache_context.mark_step_begin()
+        self.cache_manager.mark_step_begin()
         # Residual L1 diff or Hidden States L1 diff
-        can_use_cache = cache_context.get_can_use_cache(
+        can_use_cache = self.cache_manager.can_cache(
             (
                 Fn_hidden_states_residual
-                if not cache_context.is_l1_diff_enabled()
+                if not self.cache_manager.is_l1_diff_enabled()
                 else hidden_states
             ),
             parallelized=self._is_parallelized(),
             prefix=(
-                "Fn_residual"
-                if not cache_context.is_l1_diff_enabled()
-                else "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:
-            cache_context.add_cached_step()
+            self.cache_manager.add_cached_step()
             del Fn_hidden_states_residual
             hidden_states, encoder_hidden_states = (
-                cache_context.apply_hidden_states_residual(
+                self.cache_manager.apply_cache(
                     hidden_states,
                     encoder_hidden_states,
                     prefix=(
-                        "Bn_residual"
-                        if cache_context.is_cache_residual()
-                        else "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=(
-                        "Bn_residual"
-                        if cache_context.is_encoder_cache_residual()
-                        else "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"
                     ),
                 )
             )
@@ -127,12 +148,16 @@ class DBCachedBlocks_Pattern_Base(torch.nn.Module):
                 **kwargs,
             )
         else:
-            cache_context.set_Fn_buffer(
-                Fn_hidden_states_residual, prefix="Fn_residual"
+            self.cache_manager.set_Fn_buffer(
+                Fn_hidden_states_residual,
+                prefix=f"{self.cache_prefix}_Fn_residual",
             )
-            if cache_context.is_l1_diff_enabled():
+            if self.cache_manager.is_l1_diff_enabled():
                 # for hidden states L1 diff
-                cache_context.set_Fn_buffer(hidden_states, "Fn_hidden_states")
+                self.cache_manager.set_Fn_buffer(
+                    hidden_states,
+                    f"{self.cache_prefix}_Fn_hidden_states",
+                )
             del Fn_hidden_states_residual
             torch._dynamo.graph_break()
             (
@@ -147,27 +172,27 @@ class DBCachedBlocks_Pattern_Base(torch.nn.Module):
                 **kwargs,
             )
             torch._dynamo.graph_break()
-            if cache_context.is_cache_residual():
-                cache_context.set_Bn_buffer(
+            if self.cache_manager.is_cache_residual():
+                self.cache_manager.set_Bn_buffer(
                     hidden_states_residual,
-                    prefix="Bn_residual",
+                    prefix=f"{self.cache_prefix}_Bn_residual",
                 )
             else:
                 # TaylorSeer
-                cache_context.set_Bn_buffer(
+                self.cache_manager.set_Bn_buffer(
                     hidden_states,
-                    prefix="Bn_hidden_states",
+                    prefix=f"{self.cache_prefix}_Bn_hidden_states",
                 )
-            if cache_context.is_encoder_cache_residual():
-                cache_context.set_Bn_encoder_buffer(
+            if self.cache_manager.is_encoder_cache_residual():
+                self.cache_manager.set_Bn_encoder_buffer(
                     encoder_hidden_states_residual,
-                    prefix="Bn_residual",
+                    prefix=f"{self.cache_prefix}_Bn_residual",
                 )
             else:
                 # TaylorSeer
-                cache_context.set_Bn_encoder_buffer(
+                self.cache_manager.set_Bn_encoder_buffer(
                     encoder_hidden_states,
-                    prefix="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
@@ -179,7 +204,7 @@ class DBCachedBlocks_Pattern_Base(torch.nn.Module):
                 **kwargs,
             )
 
-        patch_cached_stats(self.transformer)
+        # patch cached stats for blocks or remove it.
         torch._dynamo.graph_break()
 
         return (
@@ -213,10 +238,11 @@ class DBCachedBlocks_Pattern_Base(torch.nn.Module):
         # If so, we can skip some Bn blocks and directly
         # use the cached values.
         return (
-            cache_context.get_current_step() in cache_context.get_cached_steps()
+            self.cache_manager.get_current_step()
+            in self.cache_manager.get_cached_steps()
         ) or (
-            cache_context.get_current_step()
-            in cache_context.get_cfg_cached_steps()
+            self.cache_manager.get_current_step()
+            in self.cache_manager.get_cfg_cached_steps()
         )
 
     @torch.compiler.disable
@@ -225,20 +251,20 @@ class DBCachedBlocks_Pattern_Base(torch.nn.Module):
         # more stable diff calculation.
         # Fn: [0,...,n-1]
         selected_Fn_blocks = self.transformer_blocks[
-            : cache_context.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 cache_context.Bn_compute_blocks() == 0:  # WARN: x[:-0] = []
+        if self.cache_manager.Bn_compute_blocks() == 0:  # WARN: x[:-0] = []
             selected_Mn_blocks = self.transformer_blocks[
-                cache_context.Fn_compute_blocks() :
+                self.cache_manager.Fn_compute_blocks() :
             ]
         else:
             selected_Mn_blocks = self.transformer_blocks[
-                cache_context.Fn_compute_blocks() : -cache_context.Bn_compute_blocks()
+                self.cache_manager.Fn_compute_blocks() : -self.cache_manager.Bn_compute_blocks()
             ]
         return selected_Mn_blocks
 
@@ -246,7 +272,7 @@ class DBCachedBlocks_Pattern_Base(torch.nn.Module):
     def _Bn_blocks(self):
         # Bn: transformer_blocks [N-n+1,...,N-1]
         selected_Bn_blocks = self.transformer_blocks[
-            -cache_context.Bn_compute_blocks() :
+            -self.cache_manager.Bn_compute_blocks() :
         ]
         return selected_Bn_blocks
 
@@ -257,10 +283,10 @@ class DBCachedBlocks_Pattern_Base(torch.nn.Module):
         *args,
         **kwargs,
     ):
-        assert cache_context.Fn_compute_blocks() <= len(
+        assert self.cache_manager.Fn_compute_blocks() <= len(
             self.transformer_blocks
         ), (
-            f"Fn_compute_blocks {cache_context.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():
@@ -357,7 +383,7 @@ class DBCachedBlocks_Pattern_Base(torch.nn.Module):
                     )
             # Cache residuals for the non-compute Bn blocks for
             # subsequent cache steps.
-            if block_id not in cache_context.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
                 )
@@ -366,22 +392,22 @@ class DBCachedBlocks_Pattern_Base(torch.nn.Module):
                 )
 
                 # Save original_hidden_states for diff calculation.
-                cache_context.set_Bn_buffer(
+                self.cache_manager.set_Bn_buffer(
                     Bn_i_original_hidden_states,
-                    prefix=f"Bn_{block_id}_original",
+                    prefix=f"{self.cache_prefix}_Bn_{block_id}_original",
                 )
-                cache_context.set_Bn_encoder_buffer(
+                self.cache_manager.set_Bn_encoder_buffer(
                     Bn_i_original_encoder_hidden_states,
-                    prefix=f"Bn_{block_id}_original",
+                    prefix=f"{self.cache_prefix}_Bn_{block_id}_original",
                 )
 
-                cache_context.set_Bn_buffer(
+                self.cache_manager.set_Bn_buffer(
                     Bn_i_hidden_states_residual,
-                    prefix=f"Bn_{block_id}_residual",
+                    prefix=f"{self.cache_prefix}_Bn_{block_id}_residual",
                 )
-                cache_context.set_Bn_encoder_buffer(
+                self.cache_manager.set_Bn_encoder_buffer(
                     Bn_i_encoder_hidden_states_residual,
-                    prefix=f"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
@@ -392,7 +418,7 @@ class DBCachedBlocks_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 cache_context.Bn_compute_blocks_ids():
+            if block_id in self.cache_manager.Bn_compute_blocks_ids():
                 hidden_states = block(
                     hidden_states,
                     encoder_hidden_states,
@@ -410,25 +436,25 @@ class DBCachedBlocks_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 cache_context.get_can_use_cache(
+                if self.cache_manager.can_cache(
                     hidden_states,  # curr step
                     parallelized=self._is_parallelized(),
-                    threshold=cache_context.non_compute_blocks_diff_threshold(),
-                    prefix=f"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 = (
-                        cache_context.apply_hidden_states_residual(
+                        self.cache_manager.apply_cache(
                             hidden_states,
                             encoder_hidden_states,
                             prefix=(
-                                f"Bn_{block_id}_residual"
-                                if cache_context.is_cache_residual()
-                                else f"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"Bn_{block_id}_residual"
-                                if cache_context.is_encoder_cache_residual()
-                                else f"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"
                             ),
                         )
                     )
@@ -455,16 +481,16 @@ class DBCachedBlocks_Pattern_Base(torch.nn.Module):
         *args,
         **kwargs,
     ):
-        if cache_context.Bn_compute_blocks() == 0:
+        if self.cache_manager.Bn_compute_blocks() == 0:
             return hidden_states, encoder_hidden_states
 
-        assert cache_context.Bn_compute_blocks() <= len(
+        assert self.cache_manager.Bn_compute_blocks() <= len(
             self.transformer_blocks
         ), (
-            f"Bn_compute_blocks {cache_context.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(cache_context.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

@@ -1,19 +1,25 @@
 import torch
 
-from cache_dit.cache_factory import cache_context
+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(
-    transformer,
+    module: torch.nn.Module | Any,
+    cache_context: CachedContext | str = None,
+    cache_manager: CachedContextManager = None,
 ):
-    # Patch the cached stats to the transformer, the cached stats
+    # Patch the cached stats to the module, the cached stats
     # will be reset for each calling of pipe.__call__(**kwargs).
-    if transformer is None:
+    if module is None or cache_manager is None:
         return
 
-    # TODO: Patch more cached stats to the transformer
-    transformer._cached_steps = cache_context.get_cached_steps()
-    transformer._residual_diffs = cache_context.get_residual_diffs()
-    transformer._cfg_cached_steps = cache_context.get_cfg_cached_steps()
-    transformer._cfg_residual_diffs = cache_context.get_cfg_residual_diffs()
+    if cache_context is not None:
+        cache_manager.set_context(cache_context)
+
+    # TODO: Patch more cached stats to the module
+    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

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