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

cache_dit/cache_factory/cache_adapters.py

@@ -3,14 +3,16 @@ import torch
 import unittest
 import functools
 
-from typing import Dict
 from contextlib import ExitStack
+from typing import Dict, List, Tuple, Any
+
 from diffusers import DiffusionPipeline
+
 from cache_dit.cache_factory import CacheType
-from cache_dit.cache_factory import CachedContext
-from cache_dit.cache_factory import ForwardPattern
 from cache_dit.cache_factory import BlockAdapter
+from cache_dit.cache_factory import ParamsModifier
 from cache_dit.cache_factory import BlockAdapterRegistry
+from cache_dit.cache_factory import CachedContextManager
 from cache_dit.cache_factory import CachedBlocks
 
 from cache_dit.logger import init_logger
@@ -29,7 +31,6 @@ class CachedAdapter:
         cls,
         pipe: DiffusionPipeline = None,
         block_adapter: BlockAdapter = None,
-        # forward_pattern: ForwardPattern = ForwardPattern.Pattern_0,
         **cache_context_kwargs,
     ) -> DiffusionPipeline:
         assert (
@@ -74,54 +75,67 @@ class CachedAdapter:
             )
 
         if BlockAdapter.check_block_adapter(block_adapter):
+
+            # 0. Must normalize block_adapter before apply cache
             block_adapter = BlockAdapter.normalize(block_adapter)
-            # 0. Apply cache on pipeline: wrap cache context
+            if BlockAdapter.is_cached(block_adapter):
+                return block_adapter.pipe
+
+            # 1. Apply cache on pipeline: wrap cache context, must
+            # call create_context before mock_blocks.
             cls.create_context(
                 block_adapter,
                 **cache_context_kwargs,
             )
-            # 1. Apply cache on transformer: mock cached transformer blocks
+
+            # 2. Apply cache on transformer: mock cached blocks
             cls.mock_blocks(
                 block_adapter,
             )
-            cls.patch_params(
-                block_adapter,
-                **cache_context_kwargs,
-            )
+
         return block_adapter.pipe
 
     @classmethod
     def patch_params(
         cls,
         block_adapter: BlockAdapter,
-        **cache_context_kwargs,
+        contexts_kwargs: List[Dict],
     ):
-        block_adapter.transformer._forward_pattern = (
-            block_adapter.forward_pattern
-        )
-        block_adapter.transformer._has_separate_cfg = (
-            block_adapter.has_separate_cfg
-        )
-        block_adapter.transformer._cache_context_kwargs = cache_context_kwargs
-        block_adapter.pipe.__class__._cache_context_kwargs = (
-            cache_context_kwargs
-        )
-        for blocks, forward_pattern in zip(
-            block_adapter.blocks, block_adapter.forward_pattern
-        ):
-            blocks._forward_pattern = forward_pattern
-            blocks._cache_context_kwargs = cache_context_kwargs
+        block_adapter.pipe._cache_context_kwargs = contexts_kwargs[0]
+
+        params_shift = 0
+        for i in range(len(block_adapter.transformer)):
+
+            block_adapter.transformer[i]._forward_pattern = (
+                block_adapter.forward_pattern
+            )
+            block_adapter.transformer[i]._has_separate_cfg = (
+                block_adapter.has_separate_cfg
+            )
+            block_adapter.transformer[i]._cache_context_kwargs = (
+                contexts_kwargs[params_shift]
+            )
+
+            blocks = block_adapter.blocks[i]
+            for j in range(len(blocks)):
+                blocks[j]._forward_pattern = block_adapter.forward_pattern[i][j]
+                blocks[j]._cache_context_kwargs = contexts_kwargs[
+                    params_shift + j
+                ]
+
+            params_shift += len(blocks)
 
     @classmethod
     def check_context_kwargs(cls, pipe, **cache_context_kwargs):
         # Check cache_context_kwargs
-        if not cache_context_kwargs["do_separate_cfg"]:
+        if not cache_context_kwargs["enable_spearate_cfg"]:
             # Check cfg for some specific case if users don't set it as True
-            cache_context_kwargs["do_separate_cfg"] = (
+            cache_context_kwargs["enable_spearate_cfg"] = (
                 BlockAdapterRegistry.has_separate_cfg(pipe)
             )
             logger.info(
-                f"Use default 'do_separate_cfg': {cache_context_kwargs['do_separate_cfg']}, "
+                f"Use default 'enable_spearate_cfg': "
+                f"{cache_context_kwargs['enable_spearate_cfg']}, "
                 f"Pipeline: {pipe.__class__.__name__}."
             )
 
@@ -138,7 +152,10 @@ class CachedAdapter:
         block_adapter: BlockAdapter,
         **cache_context_kwargs,
     ) -> DiffusionPipeline:
-        if getattr(block_adapter.pipe, "_is_cached", False):
+
+        BlockAdapter.assert_normalized(block_adapter)
+
+        if BlockAdapter.is_cached(block_adapter.pipe):
             return block_adapter.pipe
 
         # Check cache_context_kwargs
@@ -147,22 +164,35 @@ class CachedAdapter:
             **cache_context_kwargs,
         )
         # Apply cache on pipeline: wrap cache context
-        cache_kwargs, _ = CachedContext.collect_cache_kwargs(
-            default_attrs={},
-            **cache_context_kwargs,
+        pipe_cls_name = block_adapter.pipe.__class__.__name__
+
+        # Each Pipeline should have it's own context manager instance.
+        # Different transformers (Wan2.2, etc) should shared the same
+        # cache manager but with different cache context (according
+        # to their unique instance id).
+        cache_manager = CachedContextManager(
+            name=f"{pipe_cls_name}_{hash(id(block_adapter.pipe))}",
+        )
+        block_adapter.pipe._cache_manager = cache_manager  # instance level
+
+        flatten_contexts, contexts_kwargs = cls.modify_context_params(
+            block_adapter, cache_manager, **cache_context_kwargs
         )
+
         original_call = block_adapter.pipe.__class__.__call__
 
         @functools.wraps(original_call)
         def new_call(self, *args, **kwargs):
             with ExitStack() as stack:
-                # cache context will reset for each pipe inference
-                for blocks_name in block_adapter.blocks_name:
+                # cache context will be reset for each pipe inference
+                for context_name, context_kwargs in zip(
+                    flatten_contexts, contexts_kwargs
+                ):
                     stack.enter_context(
-                        CachedContext.cache_context(
-                            CachedContext.reset_cache_context(
-                                blocks_name,
-                                **cache_kwargs,
+                        cache_manager.enter_context(
+                            cache_manager.reset_context(
+                                context_name,
+                                **context_kwargs,
                             ),
                         )
                     )
@@ -171,109 +201,194 @@ class CachedAdapter:
                 return outputs
 
         block_adapter.pipe.__class__.__call__ = new_call
+        block_adapter.pipe.__class__._original_call = original_call
         block_adapter.pipe.__class__._is_cached = True
+
+        cls.patch_params(block_adapter, contexts_kwargs)
+
         return block_adapter.pipe
 
     @classmethod
-    def patch_stats(cls, block_adapter: BlockAdapter):
+    def modify_context_params(
+        cls,
+        block_adapter: BlockAdapter,
+        cache_manager: CachedContextManager,
+        **cache_context_kwargs,
+    ) -> Tuple[List[str], List[Dict[str, Any]]]:
+
+        flatten_contexts = BlockAdapter.flatten(
+            block_adapter.unique_blocks_name
+        )
+        contexts_kwargs = [
+            cache_context_kwargs.copy()
+            for _ in range(
+                len(flatten_contexts),
+            )
+        ]
+
+        for i in range(len(contexts_kwargs)):
+            contexts_kwargs[i]["name"] = flatten_contexts[i]
+
+        if block_adapter.params_modifiers is None:
+            return flatten_contexts, contexts_kwargs
+
+        flatten_modifiers: List[ParamsModifier] = BlockAdapter.flatten(
+            block_adapter.params_modifiers,
+        )
+
+        for i in range(
+            min(len(contexts_kwargs), len(flatten_modifiers)),
+        ):
+            contexts_kwargs[i].update(
+                flatten_modifiers[i]._context_kwargs,
+            )
+            contexts_kwargs[i], _ = cache_manager.collect_cache_kwargs(
+                default_attrs={}, **contexts_kwargs[i]
+            )
+
+        return flatten_contexts, contexts_kwargs
+
+    @classmethod
+    def patch_stats(
+        cls,
+        block_adapter: BlockAdapter,
+    ):
         from cache_dit.cache_factory.cache_blocks.utils import (
             patch_cached_stats,
         )
 
-        patch_cached_stats(block_adapter.transformer)
-        for blocks, blocks_name in zip(
-            block_adapter.blocks, block_adapter.blocks_name
-        ):
-            patch_cached_stats(blocks, blocks_name)
+        cache_manager = block_adapter.pipe._cache_manager
+
+        for i in range(len(block_adapter.transformer)):
+            patch_cached_stats(
+                block_adapter.transformer[i],
+                cache_context=block_adapter.unique_blocks_name[i][-1],
+                cache_manager=cache_manager,
+            )
+            for blocks, unique_name in zip(
+                block_adapter.blocks[i],
+                block_adapter.unique_blocks_name[i],
+            ):
+                patch_cached_stats(
+                    blocks,
+                    cache_context=unique_name,
+                    cache_manager=cache_manager,
+                )
 
     @classmethod
     def mock_blocks(
         cls,
         block_adapter: BlockAdapter,
-    ) -> torch.nn.Module:
+    ) -> List[torch.nn.Module]:
+
+        BlockAdapter.assert_normalized(block_adapter)
 
-        if getattr(block_adapter.transformer, "_is_cached", False):
+        if BlockAdapter.is_cached(block_adapter.transformer):
             return block_adapter.transformer
 
-        # Check block forward pattern matching
-        block_adapter = BlockAdapter.normalize(block_adapter)
-        for forward_pattern, blocks in zip(
-            block_adapter.forward_pattern, block_adapter.blocks
+        # Apply cache on transformer: mock cached transformer blocks
+        for (
+            cached_blocks,
+            transformer,
+            blocks_name,
+            unique_blocks_name,
+            dummy_blocks_names,
+        ) in zip(
+            cls.collect_cached_blocks(block_adapter),
+            block_adapter.transformer,
+            block_adapter.blocks_name,
+            block_adapter.unique_blocks_name,
+            block_adapter.dummy_blocks_names,
         ):
-            assert BlockAdapter.match_blocks_pattern(
-                blocks,
-                forward_pattern=forward_pattern,
-                check_num_outputs=block_adapter.check_num_outputs,
-            ), (
-                "No block forward pattern matched, "
-                f"supported lists: {ForwardPattern.supported_patterns()}"
+            cls.mock_transformer(
+                cached_blocks,
+                transformer,
+                blocks_name,
+                unique_blocks_name,
+                dummy_blocks_names,
             )
 
-        # Apply cache on transformer: mock cached transformer blocks
-        # TODO: Use blocks_name to spearate cached context for different
-        # blocks list. For example, single_transformer_blocks and
-        # transformer_blocks should have different cached context and
-        # forward pattern.
-        cached_blocks = cls.collect_cached_blocks(
-            block_adapter=block_adapter,
-        )
+        return block_adapter.transformer
+
+    @classmethod
+    def mock_transformer(
+        cls,
+        cached_blocks: Dict[str, torch.nn.ModuleList],
+        transformer: torch.nn.Module,
+        blocks_name: List[str],
+        unique_blocks_name: List[str],
+        dummy_blocks_names: List[str],
+    ) -> torch.nn.Module:
         dummy_blocks = torch.nn.ModuleList()
 
-        original_forward = block_adapter.transformer.forward
+        original_forward = transformer.forward
 
-        assert isinstance(block_adapter.dummy_blocks_names, list)
+        assert isinstance(dummy_blocks_names, list)
 
         @functools.wraps(original_forward)
         def new_forward(self, *args, **kwargs):
             with ExitStack() as stack:
-                for blocks_name in block_adapter.blocks_name:
+                for name, context_name in zip(
+                    blocks_name,
+                    unique_blocks_name,
+                ):
                     stack.enter_context(
                         unittest.mock.patch.object(
-                            self,
-                            blocks_name,
-                            cached_blocks[blocks_name],
+                            self, name, cached_blocks[context_name]
                         )
                     )
-                for dummy_name in block_adapter.dummy_blocks_names:
+                for dummy_name in dummy_blocks_names:
                     stack.enter_context(
                         unittest.mock.patch.object(
-                            self,
-                            dummy_name,
-                            dummy_blocks,
+                            self, dummy_name, dummy_blocks
                         )
                     )
                 return original_forward(*args, **kwargs)
 
-        block_adapter.transformer.forward = new_forward.__get__(
-            block_adapter.transformer
-        )
-        block_adapter.transformer._is_cached = True
+        transformer.forward = new_forward.__get__(transformer)
+        transformer._original_forward = original_forward
+        transformer._is_cached = True
 
-        return block_adapter.transformer
+        return transformer
 
     @classmethod
     def collect_cached_blocks(
         cls,
         block_adapter: BlockAdapter,
-    ) -> Dict[str, torch.nn.ModuleList]:
-        block_adapter = BlockAdapter.normalize(block_adapter)
+    ) -> List[Dict[str, torch.nn.ModuleList]]:
+
+        BlockAdapter.assert_normalized(block_adapter)
 
-        cached_blocks_bind_context = {}
+        total_cached_blocks: List[Dict[str, torch.nn.ModuleList]] = []
+        assert hasattr(block_adapter.pipe, "_cache_manager")
+        assert isinstance(
+            block_adapter.pipe._cache_manager, CachedContextManager
+        )
+
+        for i in range(len(block_adapter.transformer)):
 
-        for i in range(len(block_adapter.blocks)):
-            cached_blocks_bind_context[block_adapter.blocks_name[i]] = (
-                torch.nn.ModuleList(
+            cached_blocks_bind_context = {}
+            for j in range(len(block_adapter.blocks[i])):
+                cached_blocks_bind_context[
+                    block_adapter.unique_blocks_name[i][j]
+                ] = torch.nn.ModuleList(
                     [
                         CachedBlocks(
-                            block_adapter.blocks[i],
-                            block_adapter.blocks_name[i],
-                            block_adapter.blocks_name[i],  # context name
-                            transformer=block_adapter.transformer,
-                            forward_pattern=block_adapter.forward_pattern[i],
+                            # 0. Transformer blocks configuration
+                            block_adapter.blocks[i][j],
+                            transformer=block_adapter.transformer[i],
+                            forward_pattern=block_adapter.forward_pattern[i][j],
                             check_num_outputs=block_adapter.check_num_outputs,
+                            # 1. Cache context configuration
+                            cache_prefix=block_adapter.blocks_name[i][j],
+                            cache_context=block_adapter.unique_blocks_name[i][
+                                j
+                            ],
+                            cache_manager=block_adapter.pipe._cache_manager,
                         )
                     ]
                 )
-            )
 
-        return cached_blocks_bind_context
+            total_cached_blocks.append(cached_blocks_bind_context)
+
+        return total_cached_blocks

cache_dit/cache_factory/cache_blocks/__init__.py

@@ -1,3 +1,11 @@
+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 (
     CachedBlocks_Pattern_0_1_2,
 )
@@ -5,14 +13,57 @@ from cache_dit.cache_factory.cache_blocks.pattern_3_4_5 import (
     CachedBlocks_Pattern_3_4_5,
 )
 
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
 
 class CachedBlocks:
-    def __new__(cls, *args, **kwargs):
-        forward_pattern = kwargs.get("forward_pattern", None)
+    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."
+        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(*args, **kwargs)
+            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(*args, **kwargs)
+            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_3_4_5.py

@@ -1,6 +1,5 @@
 import torch
 
-from cache_dit.cache_factory import CachedContext
 from cache_dit.cache_factory import ForwardPattern
 from cache_dit.cache_factory.cache_blocks.pattern_base import (
     CachedBlocks_Pattern_Base,
@@ -24,7 +23,7 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
         **kwargs,
     ):
         # Use it's own cache context.
-        CachedContext.set_cache_context(
+        self.cache_manager.set_context(
             self.cache_context,
         )
 
@@ -41,40 +40,40 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
         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,
                     # None Pattern 3, else 4, 5
                     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"
                     ),
                 )
             )
@@ -88,15 +87,15 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
                 **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()
@@ -114,29 +113,29 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
                 **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(
                     # None Pattern 3, else 4, 5
                     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(
                     # None Pattern 3, else 4, 5
                     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
@@ -167,10 +166,10 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
         *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)}"
         )
         encoder_hidden_states = None  # Pattern 3
@@ -242,16 +241,16 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
         *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:
             raise ValueError(
                 f"Bn_compute_blocks_ids is not support for "
                 f"patterns: {self._supported_patterns}."