cache-dit 0.3.2__py3-none-any.whl → 1.0.14__py3-none-any.whl

cache_dit/caching/cache_blocks/pattern_base.py

@@ -0,0 +1,748 @@
+import inspect
+import logging
+import torch
+import torch.distributed as dist
+from diffusers.hooks import HookRegistry
+
+try:
+    from diffusers.hooks.context_parallel import ContextParallelSplitHook
+except ImportError:
+    ContextParallelSplitHook = None
+    raise UserWarning(
+        "Context parallelism requires the 'diffusers>=0.36.dev0'."
+        "Please install latest version of diffusers from source: \n"
+        "pip3 install git+https://github.com/huggingface/diffusers.git"
+    )
+from cache_dit.caching.cache_contexts.cache_context import CachedContext
+from cache_dit.caching.cache_contexts.prune_context import PrunedContext
+from cache_dit.caching.cache_contexts.cache_manager import (
+    CachedContextManager,
+    ContextNotExistError,
+)
+from cache_dit.caching.cache_contexts.prune_manager import (
+    PrunedContextManager,
+)
+from cache_dit.caching import ForwardPattern
+from cache_dit.caching.cache_types import CacheType
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class CachedBlocks_Pattern_Base(torch.nn.Module):
+    _supported_patterns = [
+        ForwardPattern.Pattern_0,
+        ForwardPattern.Pattern_1,
+        ForwardPattern.Pattern_2,
+    ]
+
+    def __init__(
+        self,
+        # 0. Transformer blocks configuration
+        transformer_blocks: torch.nn.ModuleList,
+        transformer: torch.nn.Module = None,
+        forward_pattern: ForwardPattern = ForwardPattern.Pattern_0,
+        check_forward_pattern: bool = True,
+        check_num_outputs: bool = True,
+        # 1. Cache context configuration
+        cache_prefix: str = None,  # maybe un-need.
+        cache_context: CachedContext | str = None,
+        context_manager: CachedContextManager = None,
+        cache_type: CacheType = CacheType.DBCache,
+        **kwargs,
+    ):
+        super().__init__()
+
+        # 0. Transformer blocks configuration
+        self.transformer = transformer
+        self.transformer_blocks = transformer_blocks
+        self.forward_pattern = forward_pattern
+        self.check_forward_pattern = check_forward_pattern
+        self.check_num_outputs = check_num_outputs
+        # 1. Cache context configuration
+        self.cache_prefix = cache_prefix
+        self.cache_context = cache_context
+        self.context_manager = context_manager
+        self.cache_type = cache_type
+
+        self._check_forward_pattern()
+        self._check_cache_type()
+        logger.info(
+            f"Match Blocks: {self.__class__.__name__}, for "
+            f"{self.cache_prefix}, cache_context: {self.cache_context}, "
+            f"context_manager: {self.context_manager.name}."
+        )
+
+    def _check_forward_pattern(self):
+        if not self.check_forward_pattern:
+            logger.warning(
+                f"Skipped Forward Pattern Check: {self.forward_pattern}"
+            )
+            return
+
+        assert (
+            self.forward_pattern.Supported
+            and self.forward_pattern in self._supported_patterns
+        ), f"Pattern {self.forward_pattern} is not supported now!"
+
+        if self.transformer_blocks is not None:
+            for block in self.transformer_blocks:
+                # Special case for HiDreamBlock
+                if hasattr(block, "block"):
+                    if isinstance(block.block, torch.nn.Module):
+                        block = block.block
+
+                forward_parameters = set(
+                    inspect.signature(block.forward).parameters.keys()
+                )
+
+                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 (
+                        required_param in forward_parameters
+                    ), f"The input parameters must contains: {required_param}."
+
+    @torch.compiler.disable
+    def _check_cache_type(self):
+        assert (
+            self.cache_type == CacheType.DBCache
+        ), f"Cache type {self.cache_type} is not supported for CachedBlocks."
+
+    @torch.compiler.disable
+    def _check_cache_params(self):
+        self._check_cache_type()
+        assert self.context_manager.Fn_compute_blocks() <= len(
+            self.transformer_blocks
+        ), (
+            f"Fn_compute_blocks {self.context_manager.Fn_compute_blocks()} must be less than "
+            f"the number of transformer blocks {len(self.transformer_blocks)}"
+        )
+        assert self.context_manager.Bn_compute_blocks() <= len(
+            self.transformer_blocks
+        ), (
+            f"Bn_compute_blocks {self.context_manager.Bn_compute_blocks()} must be less than "
+            f"the number of transformer blocks {len(self.transformer_blocks)}"
+        )
+
+    def call_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        # Call all blocks to process the hidden states without cache.
+        for block in self.transformer_blocks:
+            hidden_states = block(
+                hidden_states,
+                encoder_hidden_states,
+                *args,
+                **kwargs,
+            )
+            hidden_states, encoder_hidden_states = self._process_block_outputs(
+                hidden_states, encoder_hidden_states
+            )
+        return hidden_states, encoder_hidden_states
+
+    @torch.compiler.disable
+    def _process_block_outputs(
+        self,
+        hidden_states: torch.Tensor | tuple,
+        encoder_hidden_states: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if not isinstance(hidden_states, torch.Tensor):
+            hidden_states, encoder_hidden_states = hidden_states
+            if not self.forward_pattern.Return_H_First:
+                hidden_states, encoder_hidden_states = (
+                    encoder_hidden_states,
+                    hidden_states,
+                )
+        return hidden_states, encoder_hidden_states
+
+    @torch.compiler.disable
+    def _process_forward_outputs(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None] | torch.Tensor:
+        return (
+            hidden_states
+            if self.forward_pattern.Return_H_Only
+            else (
+                (hidden_states, encoder_hidden_states)
+                if self.forward_pattern.Return_H_First
+                else (encoder_hidden_states, hidden_states)
+            )
+        )
+
+    @torch.compiler.disable
+    def _check_if_context_parallel_enabled(
+        self,
+        module: torch.nn.Module,
+    ) -> bool:
+        if ContextParallelSplitHook is None:
+            return False
+        if hasattr(module, "_diffusers_hook"):
+            _diffusers_hook: HookRegistry = module._diffusers_hook
+            for hook in _diffusers_hook.hooks.values():
+                if isinstance(hook, ContextParallelSplitHook):
+                    return True
+        return False
+
+    def _get_Fn_residual(
+        self,
+        original_hidden_states: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        # NOTE: Make cases compatible with context parallelism while using
+        # block level cp plan, e.g., WanTransformer3DModel. The shape of
+        # `original_hidden_states` and `hidden_states` after Fn maybe
+        # different due to seqlen split in context parallelism.
+        if self._check_if_context_parallel_enabled(
+            self.transformer_blocks[0]
+        ) and (original_hidden_states.shape != hidden_states.shape):
+            # Force use `hidden_states` as the Fn states residual for subsequent
+            # dynamic cache processing if the shape is different.
+            Fn_hidden_states_residual = hidden_states
+            if logger.isEnabledFor(logging.DEBUG):
+                logger.debug(
+                    f"Context parallelism is enabled in Fn blocks, and the shape of "
+                    f"original_hidden_states {original_hidden_states.shape} and "
+                    f"hidden_states {hidden_states.shape} are different after Fn blocks. "
+                    f"Use hidden_states as Fn_hidden_states_residual directly."
+                )
+        else:
+            Fn_hidden_states_residual = (
+                hidden_states - original_hidden_states.to(hidden_states.device)
+            )
+        return Fn_hidden_states_residual
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        # Use it's own cache context.
+        try:
+            self.context_manager.set_context(self.cache_context)
+            self._check_cache_params()
+        except ContextNotExistError as e:
+            logger.warning(f"Cache context not exist: {e}, skip cache.")
+            # Call all blocks to process the hidden states.
+            hidden_states, encoder_hidden_states = self.call_blocks(
+                hidden_states,
+                encoder_hidden_states,
+                *args,
+                **kwargs,
+            )
+            return self._process_forward_outputs(
+                hidden_states,
+                encoder_hidden_states,
+            )
+
+        original_hidden_states = hidden_states
+        # Call first `n` blocks to process the hidden states for
+        # more stable diff calculation.
+        hidden_states, encoder_hidden_states = self.call_Fn_blocks(
+            hidden_states,
+            encoder_hidden_states,
+            *args,
+            **kwargs,
+        )
+
+        Fn_hidden_states_residual = self._get_Fn_residual(
+            original_hidden_states, hidden_states
+        )
+        del original_hidden_states
+
+        self.context_manager.mark_step_begin()
+        # Residual L1 diff or Hidden States L1 diff
+        can_use_cache = self.context_manager.can_cache(
+            (
+                Fn_hidden_states_residual
+                if not self.context_manager.is_l1_diff_enabled()
+                else hidden_states
+            ),
+            parallelized=self._is_parallelized(),
+            prefix=(
+                f"{self.cache_prefix}_Fn_residual"
+                if not self.context_manager.is_l1_diff_enabled()
+                else f"{self.cache_prefix}_Fn_hidden_states"
+            ),
+        )
+
+        torch._dynamo.graph_break()
+        if can_use_cache:
+            self.context_manager.add_cached_step()
+            del Fn_hidden_states_residual
+            hidden_states, encoder_hidden_states = (
+                self.context_manager.apply_cache(
+                    hidden_states,
+                    encoder_hidden_states,
+                    prefix=(
+                        f"{self.cache_prefix}_Bn_residual"
+                        if self.context_manager.is_cache_residual()
+                        else f"{self.cache_prefix}_Bn_hidden_states"
+                    ),
+                    encoder_prefix=(
+                        f"{self.cache_prefix}_Bn_residual"
+                        if self.context_manager.is_encoder_cache_residual()
+                        else 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,
+                encoder_hidden_states,
+                *args,
+                **kwargs,
+            )
+        else:
+            self.context_manager.set_Fn_buffer(
+                Fn_hidden_states_residual,
+                prefix=f"{self.cache_prefix}_Fn_residual",
+            )
+            if self.context_manager.is_l1_diff_enabled():
+                # for hidden states L1 diff
+                self.context_manager.set_Fn_buffer(
+                    hidden_states,
+                    f"{self.cache_prefix}_Fn_hidden_states",
+                )
+            del Fn_hidden_states_residual
+            torch._dynamo.graph_break()
+            (
+                hidden_states,
+                encoder_hidden_states,
+                hidden_states_residual,
+                encoder_hidden_states_residual,
+            ) = self.call_Mn_blocks(  # middle
+                hidden_states,
+                encoder_hidden_states,
+                *args,
+                **kwargs,
+            )
+            torch._dynamo.graph_break()
+            if self.context_manager.is_cache_residual():
+                self.context_manager.set_Bn_buffer(
+                    hidden_states_residual,
+                    prefix=f"{self.cache_prefix}_Bn_residual",
+                )
+            else:
+                self.context_manager.set_Bn_buffer(
+                    hidden_states,
+                    prefix=f"{self.cache_prefix}_Bn_hidden_states",
+                )
+
+            if self.context_manager.is_encoder_cache_residual():
+                self.context_manager.set_Bn_encoder_buffer(
+                    encoder_hidden_states_residual,
+                    prefix=f"{self.cache_prefix}_Bn_residual",
+                )
+            else:
+                self.context_manager.set_Bn_encoder_buffer(
+                    encoder_hidden_states,
+                    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,
+                encoder_hidden_states,
+                *args,
+                **kwargs,
+            )
+
+        # patch cached stats for blocks or remove it.
+        torch._dynamo.graph_break()
+
+        return self._process_forward_outputs(
+            hidden_states,
+            encoder_hidden_states,
+        )
+
+    @torch.compiler.disable
+    def _is_parallelized(self):
+        # Compatible with distributed inference.
+        return any(
+            (
+                all(
+                    (
+                        self.transformer is not None,
+                        getattr(self.transformer, "_is_parallelized", False),
+                    )
+                ),
+                (dist.is_initialized() and dist.get_world_size() > 1),
+            )
+        )
+
+    @torch.compiler.disable
+    def _is_in_cache_step(self):
+        # Check if the current step is in cache steps.
+        # If so, we can skip some Bn blocks and directly
+        # use the cached values.
+        return (
+            self.context_manager.get_current_step()
+            in self.context_manager.get_cached_steps()
+        ) or (
+            self.context_manager.get_current_step()
+            in self.context_manager.get_cfg_cached_steps()
+        )
+
+    @torch.compiler.disable
+    def _Fn_blocks(self):
+        # Select first `n` blocks to process the hidden states for
+        # more stable diff calculation.
+        # Fn: [0,...,n-1]
+        selected_Fn_blocks = self.transformer_blocks[
+            : self.context_manager.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 self.context_manager.Bn_compute_blocks() == 0:  # WARN: x[:-0] = []
+            selected_Mn_blocks = self.transformer_blocks[
+                self.context_manager.Fn_compute_blocks() :
+            ]
+        else:
+            selected_Mn_blocks = self.transformer_blocks[
+                self.context_manager.Fn_compute_blocks() : -self.context_manager.Bn_compute_blocks()
+            ]
+        return selected_Mn_blocks
+
+    @torch.compiler.disable
+    def _Bn_blocks(self):
+        # Bn: transformer_blocks [N-n+1,...,N-1]
+        selected_Bn_blocks = self.transformer_blocks[
+            -self.context_manager.Bn_compute_blocks() :
+        ]
+        return selected_Bn_blocks
+
+    def call_Fn_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        for block in self._Fn_blocks():
+            hidden_states = block(
+                hidden_states,
+                encoder_hidden_states,
+                *args,
+                **kwargs,
+            )
+            hidden_states, encoder_hidden_states = self._process_block_outputs(
+                hidden_states, encoder_hidden_states
+            )
+
+        return hidden_states, encoder_hidden_states
+
+    def call_Mn_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        original_hidden_states = hidden_states
+        original_encoder_hidden_states = encoder_hidden_states
+        for block in self._Mn_blocks():
+            hidden_states = block(
+                hidden_states,
+                encoder_hidden_states,
+                *args,
+                **kwargs,
+            )
+            hidden_states, encoder_hidden_states = self._process_block_outputs(
+                hidden_states, encoder_hidden_states
+            )
+
+        # compute hidden_states residual
+        hidden_states = hidden_states.contiguous()
+
+        hidden_states_residual = hidden_states - original_hidden_states
+
+        if (
+            encoder_hidden_states is not None
+            and original_encoder_hidden_states is not None
+        ):
+            encoder_hidden_states = encoder_hidden_states.contiguous()
+            encoder_hidden_states_residual = (
+                encoder_hidden_states - original_encoder_hidden_states
+            )
+        else:
+            encoder_hidden_states_residual = None
+
+        return (
+            hidden_states,
+            encoder_hidden_states,
+            hidden_states_residual,
+            encoder_hidden_states_residual,
+        )
+
+    def call_Bn_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        if self.context_manager.Bn_compute_blocks() == 0:
+            return hidden_states, encoder_hidden_states
+
+        for block in self._Bn_blocks():
+            hidden_states = block(
+                hidden_states,
+                encoder_hidden_states,
+                *args,
+                **kwargs,
+            )
+            hidden_states, encoder_hidden_states = self._process_block_outputs(
+                hidden_states, encoder_hidden_states
+            )
+
+        return hidden_states, encoder_hidden_states
+
+
+class PrunedBlocks_Pattern_Base(CachedBlocks_Pattern_Base):
+    pruned_blocks_step: int = 0  # number of pruned blocks in current step
+
+    def __init__(
+        self,
+        # 0. Transformer blocks configuration
+        transformer_blocks: torch.nn.ModuleList,
+        transformer: torch.nn.Module = None,
+        forward_pattern: ForwardPattern = ForwardPattern.Pattern_0,
+        check_forward_pattern: bool = True,
+        check_num_outputs: bool = True,
+        # 1. Prune context configuration
+        cache_prefix: str = None,  # maybe un-need.
+        cache_context: PrunedContext | str = None,
+        context_manager: PrunedContextManager = None,
+        cache_type: CacheType = CacheType.DBPrune,
+        **kwargs,
+    ):
+        super().__init__(
+            # 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,
+        )
+        assert isinstance(
+            self.context_manager, PrunedContextManager
+        ), "context_manager must be PrunedContextManager for PrunedBlocks."
+        self.context_manager: PrunedContextManager = (
+            self.context_manager
+        )  # For type hint
+
+    @torch.compiler.disable
+    def _check_cache_type(self):
+        assert (
+            self.cache_type == CacheType.DBPrune
+        ), f"Cache type {self.cache_type} is not supported for PrunedBlocks."
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        self.pruned_blocks_step: int = 0  # reset for each step
+
+        # Use it's own cache context.
+        try:
+            self.context_manager.set_context(self.cache_context)
+            self._check_cache_params()
+        except ContextNotExistError as e:
+            logger.warning(f"Cache context not exist: {e}, skip prune.")
+            # Fallback to call all blocks to process the hidden states w/o prune.
+            hidden_states, encoder_hidden_states = self.call_blocks(
+                hidden_states,
+                encoder_hidden_states,
+                *args,
+                **kwargs,
+            )
+            return self._process_forward_outputs(
+                hidden_states,
+                encoder_hidden_states,
+            )
+
+        self.context_manager.mark_step_begin()
+
+        if self._check_if_context_parallel_enabled(self.transformer_blocks[0]):
+            raise RuntimeError(
+                "Block level Context parallelism is not supported in PrunedBlocks."
+            )
+
+        # Call all blocks with prune strategy to process the hidden states.
+        for i, block in enumerate(self.transformer_blocks):
+            hidden_states, encoder_hidden_states = self.compute_or_prune(
+                i,
+                block,
+                hidden_states,
+                encoder_hidden_states,
+                *args,
+                **kwargs,
+            )
+
+        self.context_manager.add_pruned_block(self.pruned_blocks_step)
+        self.context_manager.add_actual_block(self.num_blocks)
+
+        return self._process_forward_outputs(
+            hidden_states,
+            encoder_hidden_states,
+        )
+
+    @property
+    @torch.compiler.disable
+    def num_blocks(self):
+        return len(self.transformer_blocks)
+
+    @torch.compiler.disable
+    def _skip_prune(self, block_id: int) -> bool:
+        # Wrap for non compiled mode.
+        return block_id in self.context_manager.get_non_prune_blocks_ids(
+            self.num_blocks
+        )
+
+    @torch.compiler.disable
+    def _maybe_prune(
+        self,
+        block_id: int,  # Block index in the transformer blocks
+        hidden_states: torch.Tensor,  # hidden_states or residual
+        prefix: str = "Bn_original",  # prev step name for single blocks
+    ):
+        # Wrap for non compiled mode.
+        can_use_prune = False
+        if not self._skip_prune(block_id):
+            can_use_prune = self.context_manager.can_prune(
+                hidden_states,  # curr step
+                parallelized=self._is_parallelized(),
+                prefix=prefix,  # prev step
+            )
+        self.pruned_blocks_step += int(can_use_prune)
+        return can_use_prune
+
+    def compute_or_prune(
+        self,
+        block_id: int,  # Block index in the transformer blocks
+        # Below are the inputs to the block
+        block,  # The transformer block to be executed
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        original_hidden_states = hidden_states
+        original_encoder_hidden_states = encoder_hidden_states
+
+        can_use_prune = self._maybe_prune(
+            block_id,
+            hidden_states,
+            prefix=f"{self.cache_prefix}_{block_id}_Fn_original",
+        )
+
+        # Prune steps: Prune current block and reuse the cached
+        # residuals for hidden states approximate.
+        torch._dynamo.graph_break()
+        if can_use_prune:
+            self.context_manager.add_pruned_step()
+            hidden_states, encoder_hidden_states = (
+                self.context_manager.apply_prune(
+                    hidden_states,
+                    encoder_hidden_states,
+                    prefix=(
+                        f"{self.cache_prefix}_{block_id}_Bn_residual"
+                        if self.context_manager.is_cache_residual()
+                        else f"{self.cache_prefix}_{block_id}_Bn_hidden_states"
+                    ),
+                    encoder_prefix=(
+                        f"{self.cache_prefix}_{block_id}_Bn_encoder_residual"
+                        if self.context_manager.is_encoder_cache_residual()
+                        else f"{self.cache_prefix}_{block_id}_Bn_encoder_hidden_states"
+                    ),
+                )
+            )
+            torch._dynamo.graph_break()
+        else:
+            # Normal steps: Compute the block and cache the residuals.
+            hidden_states = block(
+                hidden_states,
+                encoder_hidden_states,
+                *args,
+                **kwargs,
+            )
+            hidden_states, encoder_hidden_states = self._process_block_outputs(
+                hidden_states, encoder_hidden_states
+            )
+            if not self._skip_prune(block_id):
+                hidden_states = hidden_states.contiguous()
+                hidden_states_residual = hidden_states - original_hidden_states
+
+                if (
+                    encoder_hidden_states is not None
+                    and original_encoder_hidden_states is not None
+                ):
+                    encoder_hidden_states = encoder_hidden_states.contiguous()
+                    encoder_hidden_states_residual = (
+                        encoder_hidden_states - original_encoder_hidden_states
+                    )
+                else:
+                    encoder_hidden_states_residual = None
+
+                self.context_manager.set_Fn_buffer(
+                    original_hidden_states,
+                    prefix=f"{self.cache_prefix}_{block_id}_Fn_original",
+                )
+                if self.context_manager.is_cache_residual():
+                    self.context_manager.set_Bn_buffer(
+                        hidden_states_residual,
+                        prefix=f"{self.cache_prefix}_{block_id}_Bn_residual",
+                    )
+                else:
+                    self.context_manager.set_Bn_buffer(
+                        hidden_states,
+                        prefix=f"{self.cache_prefix}_{block_id}_Bn_hidden_states",
+                    )
+                if encoder_hidden_states_residual is not None:
+                    if self.context_manager.is_encoder_cache_residual():
+                        self.context_manager.set_Bn_encoder_buffer(
+                            encoder_hidden_states_residual,
+                            prefix=f"{self.cache_prefix}_{block_id}_Bn_encoder_residual",
+                        )
+                    else:
+                        self.context_manager.set_Bn_encoder_buffer(
+                            encoder_hidden_states_residual,
+                            prefix=f"{self.cache_prefix}_{block_id}_Bn_encoder_hidden_states",
+                        )
+            torch._dynamo.graph_break()
+
+        return hidden_states, encoder_hidden_states