Upload qandc/cache.py

qandc/cache.py

@@ -0,0 +1,58 @@
+"""
+Feature Caching for DiT models (FORA-style).
+Caches self-attention and MLP outputs, reusing for N-1 steps after each full computation.
+"""
+import torch
+import torch.nn as nn
+
+class CachedDiTBlock(nn.Module):
+    """Wraps a DiT transformer block with static feature caching."""
+    def __init__(self, block, cache_interval=2):
+        super().__init__()
+        self.block = block
+        self.cache_interval = cache_interval
+        self.cached_output = None
+        self.step_count = 0
+        self.caching_enabled = True
+
+    def __getattr__(self, name):
+        """Delegate attribute access to the wrapped block for transparency."""
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(self.block, name)
+
+    def reset_cache(self):
+        self.cached_output = None
+        self.step_count = 0
+
+    def forward(self, hidden_states, *args, **kwargs):
+        if not self.caching_enabled:
+            return self.block(hidden_states, *args, **kwargs)
+        if self.step_count % self.cache_interval == 0:
+            output = self.block(hidden_states, *args, **kwargs)
+            if isinstance(output, tuple):
+                self.cached_output = (output[0] - hidden_states).detach()
+            else:
+                self.cached_output = (output - hidden_states).detach()
+            self.step_count += 1
+            return output
+        else:
+            self.step_count += 1
+            if isinstance(self.cached_output, tuple):
+                return (hidden_states + self.cached_output[0],) + self.cached_output[1:]
+            else:
+                return hidden_states + self.cached_output
+
+def apply_cache_to_dit(transformer, cache_interval=2):
+    if hasattr(transformer, 'transformer_blocks'):
+        blocks = transformer.transformer_blocks
+        for i in range(len(blocks)):
+            blocks[i] = CachedDiTBlock(blocks[i], cache_interval=cache_interval)
+        print(f"Applied cache (interval={cache_interval}) to {len(blocks)} transformer blocks")
+    return transformer
+
+def reset_all_caches(transformer):
+    for module in transformer.modules():
+        if isinstance(module, CachedDiTBlock):
+            module.reset_cache()