cache-dit 1.0.2__py3-none-any.whl → 1.0.4__py3-none-any.whl

cache_dit/cache_factory/cache_blocks/__init__.py

@@ -1,28 +1,33 @@
 import torch
 
 from cache_dit.cache_factory import ForwardPattern
+from cache_dit.cache_factory.cache_types import CacheType
 from cache_dit.cache_factory.cache_contexts.cache_context import CachedContext
+from cache_dit.cache_factory.cache_contexts.prune_context import PrunedContext
 from cache_dit.cache_factory.cache_contexts.cache_manager import (
     CachedContextManager,
 )
+from cache_dit.cache_factory.cache_contexts.prune_manager import (
+    PrunedContextManager,
+)
 
 from cache_dit.cache_factory.cache_blocks.pattern_0_1_2 import (
     CachedBlocks_Pattern_0_1_2,
+    PrunedBlocks_Pattern_0_1_2,
 )
 from cache_dit.cache_factory.cache_blocks.pattern_3_4_5 import (
     CachedBlocks_Pattern_3_4_5,
+    PrunedBlocks_Pattern_3_4_5,
 )
-from cache_dit.cache_factory.cache_blocks.pattern_utils import (
-    patch_cached_stats,
-    remove_cached_stats,
-)
+from cache_dit.cache_factory.cache_blocks.pattern_utils import apply_stats
+from cache_dit.cache_factory.cache_blocks.pattern_utils import remove_stats
 
 from cache_dit.logger import init_logger
 
 logger = init_logger(__name__)
 
 
-class CachedBlocks:
+class UnifiedBlocks:
     def __new__(
         cls,
         # 0. Transformer blocks configuration
@@ -36,16 +41,13 @@ class CachedBlocks:
         # '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,
+        cache_context: CachedContext | PrunedContext | str = None,
+        context_manager: CachedContextManager | PrunedContextManager = None,
+        cache_type: CacheType = CacheType.DBCache,
         **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(
+        if cache_type == CacheType.DBCache:
+            return CachedBlocks(
                 # 0. Transformer blocks configuration
                 transformer_blocks,
                 transformer=transformer,
@@ -55,11 +57,12 @@ class CachedBlocks:
                 # 1. Cache context configuration
                 cache_prefix=cache_prefix,
                 cache_context=cache_context,
-                cache_manager=cache_manager,
+                context_manager=context_manager,
+                cache_type=cache_type,
                 **kwargs,
             )
-        elif forward_pattern in CachedBlocks_Pattern_3_4_5._supported_patterns:
-            return CachedBlocks_Pattern_3_4_5(
+        elif cache_type == CacheType.DBPrune:
+            return PrunedBlocks(
                 # 0. Transformer blocks configuration
                 transformer_blocks,
                 transformer=transformer,
@@ -69,8 +72,155 @@ class CachedBlocks:
                 # 1. Cache context configuration
                 cache_prefix=cache_prefix,
                 cache_context=cache_context,
-                cache_manager=cache_manager,
+                context_manager=context_manager,
+                cache_type=cache_type,
                 **kwargs,
             )
+        else:
+            raise ValueError(f"Cache type {cache_type} is not supported now!")
+
+
+class CachedBlocks:
+    def __new__(
+        cls,
+        # 0. Transformer blocks configuration
+        transformer_blocks: torch.nn.ModuleList,
+        transformer: torch.nn.Module = None,
+        forward_pattern: ForwardPattern = None,
+        check_forward_pattern: bool = True,
+        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 | PrunedContext | str = None,
+        context_manager: CachedContextManager | PrunedContextManager = None,
+        cache_type: CacheType = CacheType.DBCache,
+        **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 context_manager is not None, "context_manager can't be None."
+        if forward_pattern in CachedBlocks_Pattern_0_1_2._supported_patterns:
+            if cache_type == CacheType.DBCache:
+                assert isinstance(
+                    context_manager, CachedContextManager
+                ), "context_manager must be CachedContextManager for DBCache."
+                return CachedBlocks_Pattern_0_1_2(
+                    # 0. Transformer blocks configuration
+                    transformer_blocks,
+                    transformer=transformer,
+                    forward_pattern=forward_pattern,
+                    check_forward_pattern=check_forward_pattern,
+                    check_num_outputs=check_num_outputs,
+                    # 1. Cache context configuration
+                    cache_prefix=cache_prefix,
+                    cache_context=cache_context,
+                    context_manager=context_manager,
+                    cache_type=cache_type,
+                    **kwargs,
+                )
+            else:
+                raise ValueError(
+                    f"Cache type {cache_type} is not supported now!"
+                )
+        elif forward_pattern in CachedBlocks_Pattern_3_4_5._supported_patterns:
+            if cache_type == CacheType.DBCache:
+                assert isinstance(
+                    context_manager, CachedContextManager
+                ), "context_manager must be CachedContextManager for DBCache."
+                return CachedBlocks_Pattern_3_4_5(
+                    # 0. Transformer blocks configuration
+                    transformer_blocks,
+                    transformer=transformer,
+                    forward_pattern=forward_pattern,
+                    check_forward_pattern=check_forward_pattern,
+                    check_num_outputs=check_num_outputs,
+                    # 1. Cache context configuration
+                    cache_prefix=cache_prefix,
+                    cache_context=cache_context,
+                    context_manager=context_manager,
+                    cache_type=cache_type,
+                    **kwargs,
+                )
+            else:
+                raise ValueError(
+                    f"Cache type {cache_type} is not supported now!"
+                )
+        else:
+            raise ValueError(f"Pattern {forward_pattern} is not supported now!")
+
+
+class PrunedBlocks:
+    def __new__(
+        cls,
+        # 0. Transformer blocks configuration
+        transformer_blocks: torch.nn.ModuleList,
+        transformer: torch.nn.Module = None,
+        forward_pattern: ForwardPattern = None,
+        check_forward_pattern: bool = True,
+        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 | PrunedContext | str = None,
+        context_manager: CachedContextManager | PrunedContextManager = None,
+        cache_type: CacheType = CacheType.DBCache,
+        **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 context_manager is not None, "context_manager can't be None."
+        if forward_pattern in PrunedBlocks_Pattern_0_1_2._supported_patterns:
+            if cache_type == CacheType.DBPrune:
+                assert isinstance(
+                    context_manager, PrunedContextManager
+                ), "context_manager must be PrunedContextManager for DBPrune."
+                return PrunedBlocks_Pattern_0_1_2(
+                    # 0. Transformer blocks configuration
+                    transformer_blocks,
+                    transformer=transformer,
+                    forward_pattern=forward_pattern,
+                    check_forward_pattern=check_forward_pattern,
+                    check_num_outputs=check_num_outputs,
+                    # 1. Cache context configuration
+                    cache_prefix=cache_prefix,
+                    cache_context=cache_context,
+                    context_manager=context_manager,
+                    cache_type=cache_type,
+                    **kwargs,
+                )
+            else:
+                raise ValueError(
+                    f"Cache type {cache_type} is not supported now!"
+                )
+        elif forward_pattern in PrunedBlocks_Pattern_3_4_5._supported_patterns:
+            if cache_type == CacheType.DBPrune:
+                assert isinstance(
+                    context_manager, PrunedContextManager
+                ), "context_manager must be PrunedContextManager for DBPrune."
+                return PrunedBlocks_Pattern_3_4_5(
+                    # 0. Transformer blocks configuration
+                    transformer_blocks,
+                    transformer=transformer,
+                    forward_pattern=forward_pattern,
+                    check_forward_pattern=check_forward_pattern,
+                    check_num_outputs=check_num_outputs,
+                    # 1. Cache context configuration
+                    cache_prefix=cache_prefix,
+                    cache_context=cache_context,
+                    context_manager=context_manager,
+                    cache_type=cache_type,
+                    **kwargs,
+                )
+            else:
+                raise ValueError(
+                    f"Cache type {cache_type} is not supported now!"
+                )
         else:
             raise ValueError(f"Pattern {forward_pattern} is not supported now!")

cache_dit/cache_factory/cache_blocks/pattern_0_1_2.py

@@ -1,6 +1,7 @@
 from cache_dit.cache_factory import ForwardPattern
 from cache_dit.cache_factory.cache_blocks.pattern_base import (
     CachedBlocks_Pattern_Base,
+    PrunedBlocks_Pattern_Base,
 )
 from cache_dit.logger import init_logger
 
@@ -14,3 +15,12 @@ class CachedBlocks_Pattern_0_1_2(CachedBlocks_Pattern_Base):
         ForwardPattern.Pattern_2,
     ]
     ...
+
+
+class PrunedBlocks_Pattern_0_1_2(PrunedBlocks_Pattern_Base):
+    _supported_patterns = [
+        ForwardPattern.Pattern_0,
+        ForwardPattern.Pattern_1,
+        ForwardPattern.Pattern_2,
+    ]
+    ...