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

cache_dit/cache_factory/cache_blocks/__init__.py

@@ -1,20 +1,69 @@
+import torch
+
+from cache_dit.cache_factory import ForwardPattern
+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.cache_blocks.pattern_0_1_2 import (
-    DBCachedBlocks_Pattern_0_1_2,
+    CachedBlocks_Pattern_0_1_2,
 )
 from cache_dit.cache_factory.cache_blocks.pattern_3_4_5 import (
-    DBCachedBlocks_Pattern_3_4_5,
+    CachedBlocks_Pattern_3_4_5,
 )
 
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
 
-class DBCachedBlocks:
-    def __new__(cls, *args, **kwargs):
-        forward_pattern = kwargs.get("forward_pattern", None)
+class CachedBlocks:
+    def __new__(
+        cls,
+        # 0. Transformer blocks configuration
+        transformer_blocks: torch.nn.ModuleList,
+        transformer: torch.nn.Module = None,
+        forward_pattern: ForwardPattern = None,
+        check_num_outputs: bool = True,
+        # 1. Cache context configuration
+        # 'transformer_blocks', 'blocks', 'single_transformer_blocks',
+        # 'layers', 'single_stream_blocks', 'double_stream_blocks'
+        cache_prefix: str = None,  # cache_prefix maybe un-need.
+        # Usually, blocks_name, etc.
+        cache_context: CachedContext | str = None,
+        cache_manager: CachedContextManager = None,
+        **kwargs,
+    ):
+        assert transformer is not None, "transformer can't be None."
         assert forward_pattern is not None, "forward_pattern can't be None."
-        if forward_pattern in DBCachedBlocks_Pattern_0_1_2._supported_patterns:
-            return DBCachedBlocks_Pattern_0_1_2(*args, **kwargs)
-        elif (
-            forward_pattern in DBCachedBlocks_Pattern_3_4_5._supported_patterns
-        ):
-            return DBCachedBlocks_Pattern_3_4_5(*args, **kwargs)
+        assert cache_context is not None, "cache_context can't be None."
+        assert cache_manager is not None, "cache_manager can't be None."
+        if forward_pattern in CachedBlocks_Pattern_0_1_2._supported_patterns:
+            return CachedBlocks_Pattern_0_1_2(
+                # 0. Transformer blocks configuration
+                transformer_blocks,
+                transformer=transformer,
+                forward_pattern=forward_pattern,
+                check_num_outputs=check_num_outputs,
+                # 1. Cache context configuration
+                cache_prefix=cache_prefix,
+                cache_context=cache_context,
+                cache_manager=cache_manager,
+                **kwargs,
+            )
+        elif forward_pattern in CachedBlocks_Pattern_3_4_5._supported_patterns:
+            return CachedBlocks_Pattern_3_4_5(
+                # 0. Transformer blocks configuration
+                transformer_blocks,
+                transformer=transformer,
+                forward_pattern=forward_pattern,
+                check_num_outputs=check_num_outputs,
+                # 1. Cache context configuration
+                cache_prefix=cache_prefix,
+                cache_context=cache_context,
+                cache_manager=cache_manager,
+                **kwargs,
+            )
         else:
             raise ValueError(f"Pattern {forward_pattern} is not supported now!")

cache_dit/cache_factory/cache_blocks/pattern_0_1_2.py

@@ -1,13 +1,13 @@
 from cache_dit.cache_factory import ForwardPattern
 from cache_dit.cache_factory.cache_blocks.pattern_base import (
-    DBCachedBlocks_Pattern_Base,
+    CachedBlocks_Pattern_Base,
 )
 from cache_dit.logger import init_logger
 
 logger = init_logger(__name__)
 
 
-class DBCachedBlocks_Pattern_0_1_2(DBCachedBlocks_Pattern_Base):
+class CachedBlocks_Pattern_0_1_2(CachedBlocks_Pattern_Base):
     _supported_patterns = [
         ForwardPattern.Pattern_0,
         ForwardPattern.Pattern_1,

cache_dit/cache_factory/cache_blocks/pattern_3_4_5.py

@@ -1,19 +1,15 @@
 import torch
 
-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_blocks.pattern_base import (
-    DBCachedBlocks_Pattern_Base,
+    CachedBlocks_Pattern_Base,
 )
 from cache_dit.logger import init_logger
 
 logger = init_logger(__name__)
 
 
-class DBCachedBlocks_Pattern_3_4_5(DBCachedBlocks_Pattern_Base):
+class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
     _supported_patterns = [
         ForwardPattern.Pattern_3,
         ForwardPattern.Pattern_4,
@@ -26,6 +22,11 @@ class DBCachedBlocks_Pattern_3_4_5(DBCachedBlocks_Pattern_Base):
         *args,
         **kwargs,
     ):
+        # Use it's own 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
         # more stable diff calculation.
@@ -39,40 +40,40 @@ class DBCachedBlocks_Pattern_3_4_5(DBCachedBlocks_Pattern_Base):
         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,
                     # None Pattern 3, else 4, 5
                     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"
                     ),
                 )
             )
@@ -86,12 +87,16 @@ class DBCachedBlocks_Pattern_3_4_5(DBCachedBlocks_Pattern_Base):
                 **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()
             (
@@ -108,29 +113,29 @@ class DBCachedBlocks_Pattern_3_4_5(DBCachedBlocks_Pattern_Base):
                 **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(
                     # None Pattern 3, else 4, 5
                     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(
                     # None Pattern 3, else 4, 5
                     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
@@ -143,7 +148,6 @@ class DBCachedBlocks_Pattern_3_4_5(DBCachedBlocks_Pattern_Base):
                 **kwargs,
             )
 
-        patch_cached_stats(self.transformer)
         torch._dynamo.graph_break()
 
         return (
@@ -162,10 +166,10 @@ class DBCachedBlocks_Pattern_3_4_5(DBCachedBlocks_Pattern_Base):
         *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)}"
         )
         encoder_hidden_states = None  # Pattern 3
@@ -237,16 +241,16 @@ class DBCachedBlocks_Pattern_3_4_5(DBCachedBlocks_Pattern_Base):
         *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:
             raise ValueError(
                 f"Bn_compute_blocks_ids is not support for "
                 f"patterns: {self._supported_patterns}."