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

cache_dit/cache_factory/cache_blocks/__init__.py

@@ -1,20 +1,18 @@
 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,
 )
 
 
-class DBCachedBlocks:
+class CachedBlocks:
     def __new__(cls, *args, **kwargs):
         forward_pattern = kwargs.get("forward_pattern", 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)
+        if forward_pattern in CachedBlocks_Pattern_0_1_2._supported_patterns:
+            return CachedBlocks_Pattern_0_1_2(*args, **kwargs)
+        elif forward_pattern in CachedBlocks_Pattern_3_4_5._supported_patterns:
+            return CachedBlocks_Pattern_3_4_5(*args, **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,16 @@
 import torch
 
-from cache_dit.cache_factory import cache_context
+from cache_dit.cache_factory import CachedContext
 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 +23,11 @@ class DBCachedBlocks_Pattern_3_4_5(DBCachedBlocks_Pattern_Base):
         *args,
         **kwargs,
     ):
+        # Use it's own cache context.
+        CachedContext.set_cache_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 +41,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()
+        CachedContext.mark_step_begin()
         # Residual L1 diff or Hidden States L1 diff
-        can_use_cache = cache_context.get_can_use_cache(
+        can_use_cache = CachedContext.get_can_use_cache(
             (
                 Fn_hidden_states_residual
-                if not cache_context.is_l1_diff_enabled()
+                if not CachedContext.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.blocks_name}_Fn_residual"
+                if not CachedContext.is_l1_diff_enabled()
+                else f"{self.blocks_name}_Fn_hidden_states"
             ),
         )
 
         torch._dynamo.graph_break()
         if can_use_cache:
-            cache_context.add_cached_step()
+            CachedContext.add_cached_step()
             del Fn_hidden_states_residual
             hidden_states, encoder_hidden_states = (
-                cache_context.apply_hidden_states_residual(
+                CachedContext.apply_hidden_states_residual(
                     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.blocks_name}_Bn_residual"
+                        if CachedContext.is_cache_residual()
+                        else f"{self.blocks_name}_Bn_hidden_states"
                     ),
                     encoder_prefix=(
-                        "Bn_residual"
-                        if cache_context.is_encoder_cache_residual()
-                        else "Bn_hidden_states"
+                        f"{self.blocks_name}_Bn_residual"
+                        if CachedContext.is_encoder_cache_residual()
+                        else f"{self.blocks_name}_Bn_hidden_states"
                     ),
                 )
             )
@@ -86,12 +88,16 @@ class DBCachedBlocks_Pattern_3_4_5(DBCachedBlocks_Pattern_Base):
                 **kwargs,
             )
         else:
-            cache_context.set_Fn_buffer(
-                Fn_hidden_states_residual, prefix="Fn_residual"
+            CachedContext.set_Fn_buffer(
+                Fn_hidden_states_residual,
+                prefix=f"{self.blocks_name}_Fn_residual",
             )
-            if cache_context.is_l1_diff_enabled():
+            if CachedContext.is_l1_diff_enabled():
                 # for hidden states L1 diff
-                cache_context.set_Fn_buffer(hidden_states, "Fn_hidden_states")
+                CachedContext.set_Fn_buffer(
+                    hidden_states,
+                    f"{self.blocks_name}_Fn_hidden_states",
+                )
             del Fn_hidden_states_residual
             torch._dynamo.graph_break()
             (
@@ -108,29 +114,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 CachedContext.is_cache_residual():
+                CachedContext.set_Bn_buffer(
                     hidden_states_residual,
-                    prefix="Bn_residual",
+                    prefix=f"{self.blocks_name}_Bn_residual",
                 )
             else:
                 # TaylorSeer
-                cache_context.set_Bn_buffer(
+                CachedContext.set_Bn_buffer(
                     hidden_states,
-                    prefix="Bn_hidden_states",
+                    prefix=f"{self.blocks_name}_Bn_hidden_states",
                 )
-            if cache_context.is_encoder_cache_residual():
-                cache_context.set_Bn_encoder_buffer(
+            if CachedContext.is_encoder_cache_residual():
+                CachedContext.set_Bn_encoder_buffer(
                     # None Pattern 3, else 4, 5
                     encoder_hidden_states_residual,
-                    prefix="Bn_residual",
+                    prefix=f"{self.blocks_name}_Bn_residual",
                 )
             else:
                 # TaylorSeer
-                cache_context.set_Bn_encoder_buffer(
+                CachedContext.set_Bn_encoder_buffer(
                     # None Pattern 3, else 4, 5
                     encoder_hidden_states,
-                    prefix="Bn_hidden_states",
+                    prefix=f"{self.blocks_name}_Bn_hidden_states",
                 )
             torch._dynamo.graph_break()
             # Call last `n` blocks to further process the hidden states
@@ -143,7 +149,6 @@ class DBCachedBlocks_Pattern_3_4_5(DBCachedBlocks_Pattern_Base):
                 **kwargs,
             )
 
-        patch_cached_stats(self.transformer)
         torch._dynamo.graph_break()
 
         return (
@@ -162,10 +167,10 @@ class DBCachedBlocks_Pattern_3_4_5(DBCachedBlocks_Pattern_Base):
         *args,
         **kwargs,
     ):
-        assert cache_context.Fn_compute_blocks() <= len(
+        assert CachedContext.Fn_compute_blocks() <= len(
             self.transformer_blocks
         ), (
-            f"Fn_compute_blocks {cache_context.Fn_compute_blocks()} must be less than "
+            f"Fn_compute_blocks {CachedContext.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 +242,16 @@ class DBCachedBlocks_Pattern_3_4_5(DBCachedBlocks_Pattern_Base):
         *args,
         **kwargs,
     ):
-        if cache_context.Bn_compute_blocks() == 0:
+        if CachedContext.Bn_compute_blocks() == 0:
             return hidden_states, encoder_hidden_states
 
-        assert cache_context.Bn_compute_blocks() <= len(
+        assert CachedContext.Bn_compute_blocks() <= len(
             self.transformer_blocks
         ), (
-            f"Bn_compute_blocks {cache_context.Bn_compute_blocks()} must be less than "
+            f"Bn_compute_blocks {CachedContext.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(CachedContext.Bn_compute_blocks_ids()) > 0:
             raise ValueError(
                 f"Bn_compute_blocks_ids is not support for "
                 f"patterns: {self._supported_patterns}."

cache_dit/cache_factory/cache_blocks/pattern_base.py

@@ -2,17 +2,14 @@ import inspect
 import torch
 import torch.distributed as dist
 
-from cache_dit.cache_factory import cache_context
+from cache_dit.cache_factory import CachedContext
 from cache_dit.cache_factory import ForwardPattern
-from cache_dit.cache_factory.cache_blocks.utils import (
-    patch_cached_stats,
-)
 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,
@@ -22,17 +19,29 @@ class DBCachedBlocks_Pattern_Base(torch.nn.Module):
     def __init__(
         self,
         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,
     ):
         super().__init__()
 
         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
         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.blocks_name}, context: {self.cache_context}"
+        )
 
     def _check_forward_pattern(self):
         assert (
@@ -45,16 +54,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 +79,10 @@ class DBCachedBlocks_Pattern_Base(torch.nn.Module):
         *args,
         **kwargs,
     ):
+        CachedContext.set_cache_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 +96,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()
+        CachedContext.mark_step_begin()
         # Residual L1 diff or Hidden States L1 diff
-        can_use_cache = cache_context.get_can_use_cache(
+        can_use_cache = CachedContext.get_can_use_cache(
             (
                 Fn_hidden_states_residual
-                if not cache_context.is_l1_diff_enabled()
+                if not CachedContext.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.blocks_name}_Fn_residual"
+                if not CachedContext.is_l1_diff_enabled()
+                else f"{self.blocks_name}_Fn_hidden_states"
             ),
         )
 
         torch._dynamo.graph_break()
         if can_use_cache:
-            cache_context.add_cached_step()
+            CachedContext.add_cached_step()
             del Fn_hidden_states_residual
             hidden_states, encoder_hidden_states = (
-                cache_context.apply_hidden_states_residual(
+                CachedContext.apply_hidden_states_residual(
                     hidden_states,
                     encoder_hidden_states,
                     prefix=(
-                        "Bn_residual"
-                        if cache_context.is_cache_residual()
-                        else "Bn_hidden_states"
+                        f"{self.blocks_name}_Bn_residual"
+                        if CachedContext.is_cache_residual()
+                        else f"{self.blocks_name}_Bn_hidden_states"
                     ),
                     encoder_prefix=(
-                        "Bn_residual"
-                        if cache_context.is_encoder_cache_residual()
-                        else "Bn_hidden_states"
+                        f"{self.blocks_name}_Bn_residual"
+                        if CachedContext.is_encoder_cache_residual()
+                        else f"{self.blocks_name}_Bn_hidden_states"
                     ),
                 )
             )
@@ -127,12 +142,16 @@ class DBCachedBlocks_Pattern_Base(torch.nn.Module):
                 **kwargs,
             )
         else:
-            cache_context.set_Fn_buffer(
-                Fn_hidden_states_residual, prefix="Fn_residual"
+            CachedContext.set_Fn_buffer(
+                Fn_hidden_states_residual,
+                prefix=f"{self.blocks_name}_Fn_residual",
             )
-            if cache_context.is_l1_diff_enabled():
+            if CachedContext.is_l1_diff_enabled():
                 # for hidden states L1 diff
-                cache_context.set_Fn_buffer(hidden_states, "Fn_hidden_states")
+                CachedContext.set_Fn_buffer(
+                    hidden_states,
+                    f"{self.blocks_name}_Fn_hidden_states",
+                )
             del Fn_hidden_states_residual
             torch._dynamo.graph_break()
             (
@@ -147,27 +166,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 CachedContext.is_cache_residual():
+                CachedContext.set_Bn_buffer(
                     hidden_states_residual,
-                    prefix="Bn_residual",
+                    prefix=f"{self.blocks_name}_Bn_residual",
                 )
             else:
                 # TaylorSeer
-                cache_context.set_Bn_buffer(
+                CachedContext.set_Bn_buffer(
                     hidden_states,
-                    prefix="Bn_hidden_states",
+                    prefix=f"{self.blocks_name}_Bn_hidden_states",
                 )
-            if cache_context.is_encoder_cache_residual():
-                cache_context.set_Bn_encoder_buffer(
+            if CachedContext.is_encoder_cache_residual():
+                CachedContext.set_Bn_encoder_buffer(
                     encoder_hidden_states_residual,
-                    prefix="Bn_residual",
+                    prefix=f"{self.blocks_name}_Bn_residual",
                 )
             else:
                 # TaylorSeer
-                cache_context.set_Bn_encoder_buffer(
+                CachedContext.set_Bn_encoder_buffer(
                     encoder_hidden_states,
-                    prefix="Bn_hidden_states",
+                    prefix=f"{self.blocks_name}_Bn_hidden_states",
                 )
             torch._dynamo.graph_break()
             # Call last `n` blocks to further process the hidden states
@@ -179,7 +198,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 +232,10 @@ 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()
+            CachedContext.get_current_step() in CachedContext.get_cached_steps()
         ) or (
-            cache_context.get_current_step()
-            in cache_context.get_cfg_cached_steps()
+            CachedContext.get_current_step()
+            in CachedContext.get_cfg_cached_steps()
         )
 
     @torch.compiler.disable
@@ -225,20 +244,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()
+            : CachedContext.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 CachedContext.Bn_compute_blocks() == 0:  # WARN: x[:-0] = []
             selected_Mn_blocks = self.transformer_blocks[
-                cache_context.Fn_compute_blocks() :
+                CachedContext.Fn_compute_blocks() :
             ]
         else:
             selected_Mn_blocks = self.transformer_blocks[
-                cache_context.Fn_compute_blocks() : -cache_context.Bn_compute_blocks()
+                CachedContext.Fn_compute_blocks() : -CachedContext.Bn_compute_blocks()
             ]
         return selected_Mn_blocks
 
@@ -246,7 +265,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() :
+            -CachedContext.Bn_compute_blocks() :
         ]
         return selected_Bn_blocks
 
@@ -257,10 +276,10 @@ class DBCachedBlocks_Pattern_Base(torch.nn.Module):
         *args,
         **kwargs,
     ):
-        assert cache_context.Fn_compute_blocks() <= len(
+        assert CachedContext.Fn_compute_blocks() <= len(
             self.transformer_blocks
         ), (
-            f"Fn_compute_blocks {cache_context.Fn_compute_blocks()} must be less than "
+            f"Fn_compute_blocks {CachedContext.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 +376,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 CachedContext.Bn_compute_blocks_ids():
                 Bn_i_hidden_states_residual = (
                     hidden_states - Bn_i_original_hidden_states
                 )
@@ -366,22 +385,22 @@ class DBCachedBlocks_Pattern_Base(torch.nn.Module):
                 )
 
                 # Save original_hidden_states for diff calculation.
-                cache_context.set_Bn_buffer(
+                CachedContext.set_Bn_buffer(
                     Bn_i_original_hidden_states,
-                    prefix=f"Bn_{block_id}_original",
+                    prefix=f"{self.blocks_name}_Bn_{block_id}_original",
                 )
-                cache_context.set_Bn_encoder_buffer(
+                CachedContext.set_Bn_encoder_buffer(
                     Bn_i_original_encoder_hidden_states,
-                    prefix=f"Bn_{block_id}_original",
+                    prefix=f"{self.blocks_name}_Bn_{block_id}_original",
                 )
 
-                cache_context.set_Bn_buffer(
+                CachedContext.set_Bn_buffer(
                     Bn_i_hidden_states_residual,
-                    prefix=f"Bn_{block_id}_residual",
+                    prefix=f"{self.blocks_name}_Bn_{block_id}_residual",
                 )
-                cache_context.set_Bn_encoder_buffer(
+                CachedContext.set_Bn_encoder_buffer(
                     Bn_i_encoder_hidden_states_residual,
-                    prefix=f"Bn_{block_id}_residual",
+                    prefix=f"{self.blocks_name}_Bn_{block_id}_residual",
                 )
                 del Bn_i_hidden_states_residual
                 del Bn_i_encoder_hidden_states_residual
@@ -392,7 +411,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 CachedContext.Bn_compute_blocks_ids():
                 hidden_states = block(
                     hidden_states,
                     encoder_hidden_states,
@@ -410,25 +429,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 CachedContext.get_can_use_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=CachedContext.non_compute_blocks_diff_threshold(),
+                    prefix=f"{self.blocks_name}_Bn_{block_id}_original",  # prev step
                 ):
                     hidden_states, encoder_hidden_states = (
-                        cache_context.apply_hidden_states_residual(
+                        CachedContext.apply_hidden_states_residual(
                             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.blocks_name}_Bn_{block_id}_residual"
+                                if CachedContext.is_cache_residual()
+                                else f"{self.blocks_name}_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.blocks_name}_Bn_{block_id}_residual"
+                                if CachedContext.is_encoder_cache_residual()
+                                else f"{self.blocks_name}_Bn_{block_id}_original"
                             ),
                         )
                     )
@@ -455,16 +474,16 @@ class DBCachedBlocks_Pattern_Base(torch.nn.Module):
         *args,
         **kwargs,
     ):
-        if cache_context.Bn_compute_blocks() == 0:
+        if CachedContext.Bn_compute_blocks() == 0:
             return hidden_states, encoder_hidden_states
 
-        assert cache_context.Bn_compute_blocks() <= len(
+        assert CachedContext.Bn_compute_blocks() <= len(
             self.transformer_blocks
         ), (
-            f"Bn_compute_blocks {cache_context.Bn_compute_blocks()} must be less than "
+            f"Bn_compute_blocks {CachedContext.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(CachedContext.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,23 @@
 import torch
 
-from cache_dit.cache_factory import cache_context
+from typing import Any
+from cache_dit.cache_factory import CachedContext
 
 
 @torch.compiler.disable
 def patch_cached_stats(
-    transformer,
+    module: torch.nn.Module | Any, cache_context: str = 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:
         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:
+        CachedContext.set_cache_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()

cache_dit/cache_factory/cache_contexts/__init__.py

@@ -0,0 +1,2 @@
+# namespace alias: for _CachedContext and many others' cache context funcs.
+import cache_dit.cache_factory.cache_contexts.cache_context as CachedContext