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

cache_dit/cache_factory/cache_adapters.py

@@ -4,7 +4,7 @@ import unittest
 import functools
 
 from contextlib import ExitStack
-from typing import Dict, List, Tuple, Any
+from typing import Dict, List, Tuple, Any, Union, Callable
 
 from diffusers import DiffusionPipeline
 
@@ -14,7 +14,10 @@ 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.cache_factory.cache_blocks.utils import (
+    patch_cached_stats,
+    remove_cached_stats,
+)
 from cache_dit.logger import init_logger
 
 logger = init_logger(__name__)
@@ -29,36 +32,45 @@ class CachedAdapter:
     @classmethod
     def apply(
         cls,
-        pipe: DiffusionPipeline = None,
-        block_adapter: BlockAdapter = None,
+        pipe_or_adapter: Union[
+            DiffusionPipeline,
+            BlockAdapter,
+        ],
         **cache_context_kwargs,
-    ) -> DiffusionPipeline:
+    ) -> Union[
+        DiffusionPipeline,
+        BlockAdapter,
+    ]:
         assert (
-            pipe is not None or block_adapter is not None
+            pipe_or_adapter is not None
         ), "pipe or block_adapter can not both None!"
 
-        if pipe is not None:
-            if BlockAdapterRegistry.is_supported(pipe):
+        if isinstance(pipe_or_adapter, DiffusionPipeline):
+            if BlockAdapterRegistry.is_supported(pipe_or_adapter):
                 logger.info(
-                    f"{pipe.__class__.__name__} is officially supported by cache-dit. "
-                    "Use it's pre-defined BlockAdapter directly!"
+                    f"{pipe_or_adapter.__class__.__name__} is officially "
+                    "supported by cache-dit. Use it's pre-defined BlockAdapter "
+                    "directly!"
+                )
+                block_adapter = BlockAdapterRegistry.get_adapter(
+                    pipe_or_adapter
                 )
-                block_adapter = BlockAdapterRegistry.get_adapter(pipe)
                 return cls.cachify(
                     block_adapter,
                     **cache_context_kwargs,
-                )
+                ).pipe
             else:
                 raise ValueError(
-                    f"{pipe.__class__.__name__} is not officially supported "
+                    f"{pipe_or_adapter.__class__.__name__} is not officially supported "
                     "by cache-dit, please set BlockAdapter instead!"
                 )
         else:
+            assert isinstance(pipe_or_adapter, BlockAdapter)
             logger.info(
-                "Adapting cache acceleration using custom BlockAdapter!"
+                "Adapting Cache Acceleration using custom BlockAdapter!"
             )
             return cls.cachify(
-                block_adapter,
+                pipe_or_adapter,
                 **cache_context_kwargs,
             )
 
@@ -67,7 +79,7 @@ class CachedAdapter:
         cls,
         block_adapter: BlockAdapter,
         **cache_context_kwargs,
-    ) -> DiffusionPipeline:
+    ) -> BlockAdapter:
 
         if block_adapter.auto:
             block_adapter = BlockAdapter.auto_block_adapter(
@@ -79,7 +91,7 @@ class CachedAdapter:
             # 0. Must normalize block_adapter before apply cache
             block_adapter = BlockAdapter.normalize(block_adapter)
             if BlockAdapter.is_cached(block_adapter):
-                return block_adapter.pipe
+                return block_adapter
 
             # 1. Apply cache on pipeline: wrap cache context, must
             # call create_context before mock_blocks.
@@ -93,53 +105,36 @@ class CachedAdapter:
                 block_adapter,
             )
 
-        return block_adapter.pipe
+        return block_adapter
 
     @classmethod
-    def patch_params(
+    def check_context_kwargs(
         cls,
         block_adapter: BlockAdapter,
-        contexts_kwargs: List[Dict],
+        **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["enable_spearate_cfg"]:
             # Check cfg for some specific case if users don't set it as True
-            cache_context_kwargs["enable_spearate_cfg"] = (
-                BlockAdapterRegistry.has_separate_cfg(pipe)
-            )
-            logger.info(
-                f"Use default 'enable_spearate_cfg': "
-                f"{cache_context_kwargs['enable_spearate_cfg']}, "
-                f"Pipeline: {pipe.__class__.__name__}."
-            )
+            if BlockAdapterRegistry.has_separate_cfg(block_adapter):
+                cache_context_kwargs["enable_spearate_cfg"] = True
+                logger.info(
+                    f"Use custom 'enable_spearate_cfg' from BlockAdapter: True. "
+                    f"Pipeline: {block_adapter.pipe.__class__.__name__}."
+                )
+            else:
+                cache_context_kwargs["enable_spearate_cfg"] = (
+                    BlockAdapterRegistry.has_separate_cfg(block_adapter.pipe)
+                )
+                logger.info(
+                    f"Use default 'enable_spearate_cfg' from block adapter "
+                    f"register: {cache_context_kwargs['enable_spearate_cfg']}, "
+                    f"Pipeline: {block_adapter.pipe.__class__.__name__}."
+                )
 
-        if cache_type := cache_context_kwargs.pop("cache_type", None):
+        if (
+            cache_type := cache_context_kwargs.pop("cache_type", None)
+        ) is not None:
             assert (
                 cache_type == CacheType.DBCache
             ), "Custom cache setting only support for DBCache now!"
@@ -160,8 +155,7 @@ class CachedAdapter:
 
         # Check cache_context_kwargs
         cache_context_kwargs = cls.check_context_kwargs(
-            block_adapter.pipe,
-            **cache_context_kwargs,
+            block_adapter, **cache_context_kwargs
         )
         # Apply cache on pipeline: wrap cache context
         pipe_cls_name = block_adapter.pipe.__class__.__name__
@@ -197,14 +191,14 @@ class CachedAdapter:
                         )
                     )
                 outputs = original_call(self, *args, **kwargs)
-                cls.patch_stats(block_adapter)
+                cls.apply_stats_hooks(block_adapter)
                 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)
+        cls.apply_params_hooks(block_adapter, contexts_kwargs)
 
         return block_adapter.pipe
 
@@ -248,33 +242,6 @@ class CachedAdapter:
 
         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,
-        )
-
-        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,
@@ -392,3 +359,159 @@ class CachedAdapter:
             total_cached_blocks.append(cached_blocks_bind_context)
 
         return total_cached_blocks
+
+    @classmethod
+    def apply_params_hooks(
+        cls,
+        block_adapter: BlockAdapter,
+        contexts_kwargs: List[Dict],
+    ):
+        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 apply_stats_hooks(
+        cls,
+        block_adapter: BlockAdapter,
+    ):
+        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 maybe_release_hooks(
+        cls,
+        pipe_or_adapter: Union[
+            DiffusionPipeline,
+            BlockAdapter,
+        ],
+    ):
+        # release model hooks
+        def _release_blocks_hooks(blocks):
+            return
+
+        def _release_transformer_hooks(transformer):
+            if hasattr(transformer, "_original_forward"):
+                original_forward = transformer._original_forward
+                transformer.forward = original_forward.__get__(transformer)
+                del transformer._original_forward
+            if hasattr(transformer, "_is_cached"):
+                del transformer._is_cached
+
+        def _release_pipeline_hooks(pipe):
+            if hasattr(pipe, "_original_call"):
+                original_call = pipe.__class__._original_call
+                pipe.__class__.__call__ = original_call
+                del pipe.__class__._original_call
+            if hasattr(pipe, "_cache_manager"):
+                cache_manager = pipe._cache_manager
+                if isinstance(cache_manager, CachedContextManager):
+                    cache_manager.clear_contexts()
+                del pipe._cache_manager
+            if hasattr(pipe, "_is_cached"):
+                del pipe.__class__._is_cached
+
+        cls.release_hooks(
+            pipe_or_adapter,
+            _release_blocks_hooks,
+            _release_transformer_hooks,
+            _release_pipeline_hooks,
+        )
+
+        # release params hooks
+        def _release_blocks_params(blocks):
+            if hasattr(blocks, "_forward_pattern"):
+                del blocks._forward_pattern
+            if hasattr(blocks, "_cache_context_kwargs"):
+                del blocks._cache_context_kwargs
+
+        def _release_transformer_params(transformer):
+            if hasattr(transformer, "_forward_pattern"):
+                del transformer._forward_pattern
+            if hasattr(transformer, "_has_separate_cfg"):
+                del transformer._has_separate_cfg
+            if hasattr(transformer, "_cache_context_kwargs"):
+                del transformer._cache_context_kwargs
+            for blocks in BlockAdapter.find_blocks(transformer):
+                _release_blocks_params(blocks)
+
+        def _release_pipeline_params(pipe):
+            if hasattr(pipe, "_cache_context_kwargs"):
+                del pipe._cache_context_kwargs
+
+        cls.release_hooks(
+            pipe_or_adapter,
+            _release_blocks_params,
+            _release_transformer_params,
+            _release_pipeline_params,
+        )
+
+        # release stats hooks
+        cls.release_hooks(
+            pipe_or_adapter,
+            remove_cached_stats,
+            remove_cached_stats,
+            remove_cached_stats,
+        )
+
+    @classmethod
+    def release_hooks(
+        cls,
+        pipe_or_adapter: Union[
+            DiffusionPipeline,
+            BlockAdapter,
+        ],
+        _release_blocks: Callable,
+        _release_transformer: Callable,
+        _release_pipeline: Callable,
+    ):
+        if isinstance(pipe_or_adapter, DiffusionPipeline):
+            pipe = pipe_or_adapter
+            _release_pipeline(pipe)
+            if hasattr(pipe, "transformer"):
+                _release_transformer(pipe.transformer)
+            if hasattr(pipe, "transformer_2"):  # Wan 2.2
+                _release_transformer(pipe.transformer_2)
+        elif isinstance(pipe_or_adapter, BlockAdapter):
+            adapter = pipe_or_adapter
+            BlockAdapter.assert_normalized(adapter)
+            _release_pipeline(adapter.pipe)
+            for transformer in BlockAdapter.flatten(adapter.transformer):
+                _release_transformer(transformer)
+            for blocks in BlockAdapter.flatten(adapter.blocks):
+                _release_blocks(blocks)

cache_dit/cache_factory/cache_blocks/pattern_3_4_5.py

@@ -1,5 +1,6 @@
 import torch
 
+from typing import Dict, Any
 from cache_dit.cache_factory import ForwardPattern
 from cache_dit.cache_factory.cache_blocks.pattern_base import (
     CachedBlocks_Pattern_Base,
@@ -31,7 +32,7 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
         # Call first `n` blocks to process the hidden states for
         # more stable diff calculation.
         # encoder_hidden_states: None Pattern 3, else 4, 5
-        hidden_states, encoder_hidden_states = self.call_Fn_blocks(
+        hidden_states, new_encoder_hidden_states = self.call_Fn_blocks(
             hidden_states,
             *args,
             **kwargs,
@@ -60,11 +61,10 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
         if can_use_cache:
             self.cache_manager.add_cached_step()
             del Fn_hidden_states_residual
-            hidden_states, encoder_hidden_states = (
+            hidden_states, new_encoder_hidden_states = (
                 self.cache_manager.apply_cache(
                     hidden_states,
-                    # None Pattern 3, else 4, 5
-                    encoder_hidden_states,
+                    new_encoder_hidden_states,  # encoder_hidden_states not use cache
                     prefix=(
                         f"{self.cache_prefix}_Bn_residual"
                         if self.cache_manager.is_cache_residual()
@@ -80,12 +80,12 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
             torch._dynamo.graph_break()
             # Call last `n` blocks to further process the hidden states
             # for higher precision.
-            hidden_states, encoder_hidden_states = self.call_Bn_blocks(
-                hidden_states,
-                encoder_hidden_states,
-                *args,
-                **kwargs,
-            )
+            if self.cache_manager.Bn_compute_blocks() > 0:
+                hidden_states, new_encoder_hidden_states = self.call_Bn_blocks(
+                    hidden_states,
+                    *args,
+                    **kwargs,
+                )
         else:
             self.cache_manager.set_Fn_buffer(
                 Fn_hidden_states_residual,
@@ -99,19 +99,20 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
                 )
             del Fn_hidden_states_residual
             torch._dynamo.graph_break()
+            old_encoder_hidden_states = new_encoder_hidden_states
             (
                 hidden_states,
-                encoder_hidden_states,
+                new_encoder_hidden_states,
                 hidden_states_residual,
-                # None Pattern 3, else 4, 5
-                encoder_hidden_states_residual,
             ) = self.call_Mn_blocks(  # middle
                 hidden_states,
-                # None Pattern 3, else 4, 5
-                encoder_hidden_states,
                 *args,
                 **kwargs,
             )
+            if new_encoder_hidden_states is not None:
+                new_encoder_hidden_states_residual = (
+                    new_encoder_hidden_states - old_encoder_hidden_states
+                )
             torch._dynamo.graph_break()
             if self.cache_manager.is_cache_residual():
                 self.cache_manager.set_Bn_buffer(
@@ -119,34 +120,32 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
                     prefix=f"{self.cache_prefix}_Bn_residual",
                 )
             else:
-                # TaylorSeer
                 self.cache_manager.set_Bn_buffer(
                     hidden_states,
                     prefix=f"{self.cache_prefix}_Bn_hidden_states",
                 )
+
             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.cache_prefix}_Bn_residual",
-                )
+                if new_encoder_hidden_states is not None:
+                    self.cache_manager.set_Bn_encoder_buffer(
+                        new_encoder_hidden_states_residual,
+                        prefix=f"{self.cache_prefix}_Bn_residual",
+                    )
             else:
-                # TaylorSeer
-                self.cache_manager.set_Bn_encoder_buffer(
-                    # None Pattern 3, else 4, 5
-                    encoder_hidden_states,
-                    prefix=f"{self.cache_prefix}_Bn_hidden_states",
-                )
+                if new_encoder_hidden_states is not None:
+                    self.cache_manager.set_Bn_encoder_buffer(
+                        new_encoder_hidden_states_residual,
+                        prefix=f"{self.cache_prefix}_Bn_hidden_states",
+                    )
             torch._dynamo.graph_break()
             # Call last `n` blocks to further process the hidden states
             # for higher precision.
-            hidden_states, encoder_hidden_states = self.call_Bn_blocks(
-                hidden_states,
-                # None Pattern 3, else 4, 5
-                encoder_hidden_states,
-                *args,
-                **kwargs,
-            )
+            if self.cache_manager.Bn_compute_blocks() > 0:
+                hidden_states, new_encoder_hidden_states = self.call_Bn_blocks(
+                    hidden_states,
+                    *args,
+                    **kwargs,
+                )
 
         torch._dynamo.graph_break()
 
@@ -154,12 +153,21 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
             hidden_states
             if self.forward_pattern.Return_H_Only
             else (
-                (hidden_states, encoder_hidden_states)
+                (hidden_states, new_encoder_hidden_states)
                 if self.forward_pattern.Return_H_First
-                else (encoder_hidden_states, hidden_states)
+                else (new_encoder_hidden_states, hidden_states)
             )
         )
 
+    @torch.compiler.disable
+    def maybe_update_kwargs(
+        self, encoder_hidden_states, kwargs: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        # if "encoder_hidden_states" in kwargs:
+        #     kwargs["encoder_hidden_states"] = encoder_hidden_states
+        # return kwargs
+        return kwargs
+
     def call_Fn_blocks(
         self,
         hidden_states: torch.Tensor,
@@ -172,7 +180,7 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
             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
+        new_encoder_hidden_states = None
         for block in self._Fn_blocks():
             hidden_states = block(
                 hidden_states,
@@ -180,25 +188,27 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
                 **kwargs,
             )
             if not isinstance(hidden_states, torch.Tensor):  # Pattern 4, 5
-                hidden_states, encoder_hidden_states = hidden_states
+                hidden_states, new_encoder_hidden_states = hidden_states
                 if not self.forward_pattern.Return_H_First:
-                    hidden_states, encoder_hidden_states = (
-                        encoder_hidden_states,
+                    hidden_states, new_encoder_hidden_states = (
+                        new_encoder_hidden_states,
                         hidden_states,
                     )
+            kwargs = self.maybe_update_kwargs(
+                new_encoder_hidden_states,
+                kwargs,
+            )
 
-        return hidden_states, encoder_hidden_states
+        return hidden_states, new_encoder_hidden_states
 
     def call_Mn_blocks(
         self,
         hidden_states: torch.Tensor,
-        # None Pattern 3, else 4, 5
-        encoder_hidden_states: torch.Tensor | None,
         *args,
         **kwargs,
     ):
         original_hidden_states = hidden_states
-        original_encoder_hidden_states = encoder_hidden_states
+        new_encoder_hidden_states = None
         for block in self._Mn_blocks():
             hidden_states = block(
                 hidden_states,
@@ -206,44 +216,33 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
                 **kwargs,
             )
             if not isinstance(hidden_states, torch.Tensor):  # Pattern 4, 5
-                hidden_states, encoder_hidden_states = hidden_states
+                hidden_states, new_encoder_hidden_states = hidden_states
                 if not self.forward_pattern.Return_H_First:
-                    hidden_states, encoder_hidden_states = (
-                        encoder_hidden_states,
+                    hidden_states, new_encoder_hidden_states = (
+                        new_encoder_hidden_states,
                         hidden_states,
                     )
+            kwargs = self.maybe_update_kwargs(
+                new_encoder_hidden_states,
+                kwargs,
+            )
 
         # compute hidden_states residual
         hidden_states = hidden_states.contiguous()
         hidden_states_residual = hidden_states - original_hidden_states
-        if (
-            original_encoder_hidden_states is not None
-            and encoder_hidden_states is not None
-        ):  # Pattern 4, 5
-            encoder_hidden_states_residual = (
-                encoder_hidden_states - original_encoder_hidden_states
-            )
-        else:
-            encoder_hidden_states_residual = None  # Pattern 3
 
         return (
             hidden_states,
-            encoder_hidden_states,
+            new_encoder_hidden_states,
             hidden_states_residual,
-            encoder_hidden_states_residual,
         )
 
     def call_Bn_blocks(
         self,
         hidden_states: torch.Tensor,
-        # None Pattern 3, else 4, 5
-        encoder_hidden_states: torch.Tensor | None,
         *args,
         **kwargs,
     ):
-        if self.cache_manager.Bn_compute_blocks() == 0:
-            return hidden_states, encoder_hidden_states
-
         assert self.cache_manager.Bn_compute_blocks() <= len(
             self.transformer_blocks
         ), (
@@ -264,11 +263,15 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
                     **kwargs,
                 )
                 if not isinstance(hidden_states, torch.Tensor):  # Pattern 4,5
-                    hidden_states, encoder_hidden_states = hidden_states
+                    hidden_states, new_encoder_hidden_states = hidden_states
                     if not self.forward_pattern.Return_H_First:
-                        hidden_states, encoder_hidden_states = (
-                            encoder_hidden_states,
+                        hidden_states, new_encoder_hidden_states = (
+                            new_encoder_hidden_states,
                             hidden_states,
                         )
+                kwargs = self.maybe_update_kwargs(
+                    new_encoder_hidden_states,
+                    kwargs,
+                )
 
-        return hidden_states, encoder_hidden_states
+        return hidden_states, new_encoder_hidden_states

cache_dit/cache_factory/cache_blocks/utils.py

@@ -23,3 +23,19 @@ def patch_cached_stats(
     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()
+
+
+def remove_cached_stats(
+    module: torch.nn.Module | Any,
+):
+    if module is None:
+        return
+
+    if hasattr(module, "_cached_steps"):
+        del module._cached_steps
+    if hasattr(module, "_residual_diffs"):
+        del module._residual_diffs
+    if hasattr(module, "_cfg_cached_steps"):
+        del module._cfg_cached_steps
+    if hasattr(module, "_cfg_residual_diffs"):
+        del module._cfg_residual_diffs