lmcache-cli 0.4.5.dev0__py3-none-any.whl

lmcache/v1/compute/attention/flash_attn.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import TYPE_CHECKING
+
+# Third Party
+from vllm.attention import Attention
+from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl
+from vllm.vllm_flash_attn import flash_attn_varlen_func, get_scheduler_metadata
+import torch
+
+# First Party
+from lmcache.v1.compute.attention.abstract import AttentionInterface
+from lmcache.v1.compute.attention.metadata import LMCFlashAttnMetadata
+
+if TYPE_CHECKING:
+    # First Party
+    from lmcache.v1.compute.attention.metadata import LMCAttnMetadata
+
+
+class LMCFlashAttnBackend(AttentionInterface):
+    """
+    FlashAttention backend for LMCache.
+    This backend uses the FlashAttention implementation
+    for efficient attention computation.
+    """
+
+    def __init__(
+        self,
+        vllm_attn: Attention,
+    ):
+        self.vllm_attn = vllm_attn
+        self.vllm_attn_impl: FlashAttentionImpl = vllm_attn.impl
+
+        # TODO(Jiayi): remove this hardcode
+        self.aot_schedule = False
+
+        idx = torch.cuda.current_device()
+        self.device = torch.device(f"cuda:{idx}")
+
+    def forward_contiguous(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        output: torch.Tensor,
+        attn_metadata: "LMCAttnMetadata",
+        **kwargs,
+    ) -> torch.Tensor:
+        assert isinstance(attn_metadata, LMCFlashAttnMetadata)
+
+        cu_seqlens_q = attn_metadata.query_start_loc
+        seqused_k = attn_metadata.seq_lens
+        cu_seqlens_k = attn_metadata.cu_seqlens_k
+        max_seqlen_q = attn_metadata.max_query_len
+        max_seqlen_k = attn_metadata.max_seq_len
+
+        descale_shape = (cu_seqlens_q.shape[0] - 1, key.shape[1])
+
+        # TODO(Jiayi): Figure out how to use aot_schedule.
+        scheduler_metadata = self._schedule(
+            batch_size=1,  # NOTE(Jiayi): Assuming batch size is 1,
+            # since we are processing request by request.
+            cu_query_lens=cu_seqlens_q,
+            max_query_len=max_seqlen_q,
+            seqlens=seqused_k,
+            max_seq_len=max_seqlen_k,
+            causal=True,  # Assuming causal attention
+        )
+
+        flash_attn_varlen_func(
+            q=query,  # contiguous
+            k=key,  # contiguous
+            v=value,  # contiguous
+            out=output,
+            cu_seqlens_q=cu_seqlens_q,
+            max_seqlen_q=max_seqlen_q,
+            cu_seqlens_k=cu_seqlens_k,
+            # seqused_k=seqused_k,
+            max_seqlen_k=max_seqlen_k,
+            softmax_scale=self.vllm_attn_impl.scale,
+            causal=True,
+            alibi_slopes=self.vllm_attn_impl.alibi_slopes,
+            window_size=self.vllm_attn_impl.sliding_window,
+            block_table=None,
+            softcap=self.vllm_attn_impl.logits_soft_cap,
+            scheduler_metadata=scheduler_metadata,
+            fa_version=self.vllm_attn_impl.vllm_flash_attn_version,
+            q_descale=self.vllm_attn._q_scale.expand(descale_shape),
+            k_descale=self.vllm_attn._k_scale.expand(descale_shape),
+            v_descale=self.vllm_attn._v_scale.expand(descale_shape),
+        )
+
+        return output
+
+    def _schedule(
+        self, batch_size, cu_query_lens, max_query_len, seqlens, max_seq_len, causal
+    ):
+        if self.aot_schedule:
+            return get_scheduler_metadata(
+                batch_size=batch_size,
+                max_seqlen_q=max_query_len,
+                max_seqlen_k=max_seq_len,
+                cache_seqlens=seqlens,
+                num_heads_q=self.vllm_attn_impl.num_heads_q,
+                num_heads_kv=self.vllm_attn_impl.num_heads_kv,
+                headdim=self.vllm_attn_impl.headdim,
+                page_size=self.vllm_attn_impl.block_size,
+                cu_seqlens_q=cu_query_lens,
+                causal=causal,
+                window_size=self.vllm_attn_impl.aot_sliding_window,
+            )
+        return None
+
+    def init_attn_metadata(
+        self,
+        input_ids: torch.tensor,
+        **kwargs,
+    ) -> LMCFlashAttnMetadata:
+        seq_len = input_ids.shape[0]
+        device = input_ids.device
+        return LMCFlashAttnMetadata(
+            query_start_loc=torch.tensor(
+                [0, seq_len], dtype=torch.int32, device=device
+            ),
+            seq_lens=torch.tensor([seq_len], device=device),
+            cu_seqlens_k=torch.tensor([0, seq_len], dtype=torch.int32, device=device),
+            max_query_len=seq_len,
+            max_seq_len=seq_len,
+        )

lmcache/v1/compute/attention/flash_infer_sparse.py

@@ -0,0 +1,284 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import TYPE_CHECKING, Optional, Tuple, Union
+import math
+
+# Third Party
+from flashinfer import VariableBlockSparseAttentionWrapper
+from flashinfer.page import block_sparse_indices_to_vector_sparse_offsets
+from flashinfer.utils import (
+    TensorLayout,
+    _check_pos_encoding_mode,
+    check_shape_dtype_device,
+    device_support_pdl,
+)
+from vllm.attention import Attention
+from vllm.v1.attention.backends.flashinfer import FlashInferImpl
+import flashinfer
+import torch
+
+# First Party
+from lmcache.v1.compute.attention.abstract import AttentionInterface
+from lmcache.v1.compute.attention.metadata import LMCFlashInferSparseMetadata
+
+if TYPE_CHECKING:
+    # First Party
+    from lmcache.v1.compute.attention.metadata import LMCAttnMetadata
+
+
+# NOTE(Jiayi): This flashinfer version is 0.3.1.
+class HackBSAWrapper(VariableBlockSparseAttentionWrapper):
+    def run(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        out: Optional[torch.Tensor] = None,
+        lse: Optional[torch.Tensor] = None,
+        return_lse: bool = False,
+        enable_pdl: Optional[bool] = None,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        """
+        This is adapted from https://github.com/flashinfer-ai/flashinfer/blob/cc5ab77370dd9a489357a47e34315d9a8f3ad5fb/flashinfer/sparse.py#L1073
+
+        Compute block-sparse attention between Q/K/V tensors.
+
+        Parameters
+        ----------
+        q : torch.Tensor
+            The query tensor with shape ``(qo_len, num_qo_heads, head_dim)``.
+        k : torch.Tensor
+            The key tensor with shape ``(kv_len, num_kv_heads, head_dim)``.
+        v : torch.Tensor
+            The value tensor with shape ``(kv_len, num_kv_heads, head_dim)``.
+        out : Optional[torch.Tensor]
+            The output tensor, if not provided, will be allocated internally.
+        lse : Optional[torch.Tensor]
+            The log-sum-exp of attention logits, if not provided, will be
+            allocated internally.
+        return_lse : bool
+            Whether to return the log-sum-exp of attention logits
+        enable_pdl : bool
+            Whether to enable Programmatic Dependent Launch (PDL). See https://docs.nvidia.com/cuda/cuda-c-programming-guide/#programmatic-dependent-launch-and-synchronization
+            Only supported for >= sm90, and currently only for FA2 and CUDA core decode.
+
+        Returns
+        -------
+        Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
+            If :attr:`return_lse` is ``False``, the attention output, shape:
+            ``[M, num_qo_heads, head_dim]``.
+            If :attr:`return_lse` is ``True``, a tuple of two tensors:
+
+            * The attention output, shape: ``[M, num_qo_heads, head_dim]``.
+            * The logsumexp of attention output, shape: ``[M, num_qo_heads]``.
+        """
+
+        if enable_pdl is None:
+            enable_pdl = device_support_pdl(q.device)
+
+        pos_encoding_mode = self._pos_encoding_mode
+        logits_soft_cap = self._logits_soft_cap
+        sm_scale = self._sm_scale
+        rope_scale = self._rope_scale
+        rope_theta = self._rope_theta
+        _check_pos_encoding_mode(pos_encoding_mode)
+        if logits_soft_cap is None:
+            logits_soft_cap = 0.0
+        if sm_scale is None:
+            sm_scale = 1.0 / math.sqrt(q.size(-1))
+        if rope_scale is None:
+            rope_scale = 1.0
+        if rope_theta is None:
+            rope_theta = 1e4
+
+        # 1 denotes page size
+        k = k.reshape(-1, 1, *k.shape[-2:])
+        v = v.reshape(-1, 1, *v.shape[-2:])
+
+        stride_block = k.stride(0)
+        stride_n = k.stride(1)
+
+        if return_lse:
+            if lse is None:
+                lse = torch.empty(
+                    (q.size(0), q.size(1)), dtype=torch.float32, device=q.device
+                )
+            else:
+                check_shape_dtype_device(
+                    lse, (q.size(0), q.size(1)), torch.float32, q.device, "lse"
+                )
+
+        if out is None:
+            out = torch.empty_like(q, dtype=self._o_dtype)
+        else:
+            check_shape_dtype_device(out, q.shape, self._o_dtype, q.device, "out")
+
+        if self._backend == "fa3":
+            if (
+                self._vector_sparse_indices_buffer.numel()
+                <= self._paged_kv_indices_buf.numel()
+            ):
+                raise ValueError(
+                    "_vector_sparse_indices_buffer is not large enough. "
+                    "Please increase the buffer size."
+                )
+
+            sparse_indices = block_sparse_indices_to_vector_sparse_offsets(
+                self._paged_kv_indices_buf,
+                self._paged_kv_indptr_buf,
+                self._vector_sparse_indices_buffer,  # output
+                self._vector_sparse_indptr_buffer,
+                self._kv_lens_buffer,
+                stride_block // stride_n,
+                1,  # stride_n // stride_n
+                1,  # block_size
+            )
+            sparse_indptr = self._vector_sparse_indptr_buffer
+        else:
+            sparse_indices = self._paged_kv_indices_buf
+            sparse_indptr = self._paged_kv_indptr_buf
+
+        self._cached_module.paged_run(
+            self._float_workspace_buffer,
+            self._int_workspace_buffer,
+            self._plan_info,
+            q,
+            k,
+            v,
+            self._qo_indptr,
+            sparse_indptr,
+            sparse_indices,
+            self._paged_kv_last_page_len,
+            out,
+            lse,
+            self._mask_mode,
+            TensorLayout[self._kv_layout].value,
+            -1,  # window_left
+            enable_pdl,
+            # ADDITIONAL_FUNC_PARAMS
+            # Not supported yet
+            None,  # packed_mask_buf
+            None,  # mask_indptr_buf
+            None,  # alibi_slopes_buf
+            None,
+            None,
+            None,
+            logits_soft_cap,
+            sm_scale,
+            None,  # scale_q
+            None,  # scale_k
+            None,  # scale_v
+            rope_scale,
+            rope_theta,
+            0,  # token_pos_in_items_len
+        )
+
+        return (out, lse) if return_lse else out
+
+
+class LMCFlashInferSparseBackend(AttentionInterface):
+    """
+    FlashAttention backend for LMCache.
+    This backend uses the FlashAttention implementation
+    for efficient attention computation.
+    """
+
+    # Workspace buffer size in bytes (128 MiB)
+    _WORKSPACE_BUFFER_SIZE_BYTES = 128 * 1024 * 1024
+
+    def __init__(
+        self,
+        vllm_attn: Attention,
+    ):
+        self.vllm_attn = vllm_attn
+        self.vllm_attn_impl: FlashInferImpl = vllm_attn.impl
+
+        idx = torch.cuda.current_device()
+        self.device = torch.device(f"cuda:{idx}")
+
+        self.workspace_buffer = torch.empty(
+            self._WORKSPACE_BUFFER_SIZE_BYTES, dtype=torch.uint8, device=self.device
+        )
+
+        self.num_qo_heads = self.vllm_attn_impl.num_heads
+        self.num_kv_heads = self.vllm_attn_impl.num_kv_heads
+        self.head_dim = self.vllm_attn_impl.head_size
+
+        self.sm_scale = self.vllm_attn_impl.scale
+        self.window_left = self.vllm_attn_impl.window_left
+        self.logits_soft_cap = self.vllm_attn_impl.logits_soft_cap
+
+        self.k_scale = vllm_attn._k_scale_float
+        self.v_scale = vllm_attn._v_scale_float
+
+    def forward_contiguous(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        output: torch.Tensor,
+        attn_metadata: "LMCAttnMetadata",
+        **kwargs,
+    ) -> torch.Tensor:
+        assert isinstance(attn_metadata, LMCFlashInferSparseMetadata)
+        is_causal = attn_metadata.is_causal
+        if is_causal:
+            output = flashinfer.prefill.single_prefill_with_kv_cache(
+                q=query,
+                k=key,
+                v=value,
+                scale_q=None,
+                scale_k=self.k_scale,
+                scale_v=self.v_scale,
+                o_dtype=query.dtype,
+                custom_mask=None,
+                packed_custom_mask=None,
+                causal=True,
+                kv_layout="NHD",
+                pos_encoding_mode="NONE",
+                use_fp16_qk_reduction=False,
+                sm_scale=self.sm_scale,
+                window_left=self.window_left,
+                logits_soft_cap=self.logits_soft_cap,
+                rope_scale=None,
+                rope_theta=None,
+                backend="auto",
+                return_lse=False,
+            )
+        else:
+            output = attn_metadata.wrapper.run(query, key, value, output)
+        return output
+
+    def init_attn_metadata(
+        self,
+        input_ids: torch.Tensor,
+        **kwargs,
+    ) -> LMCFlashInferSparseMetadata:
+        """
+        Initialize non-sparse attention metadata first.
+        """
+
+        wrapper = HackBSAWrapper(self.workspace_buffer)
+        seq_len = len(input_ids)
+
+        # TODO(Jiayi): remove this hardcode
+        sparse_blk_row_size = 32
+        sparse_blk_col_size = 32
+
+        num_block_col = seq_len // sparse_blk_col_size
+        last_col_len = seq_len % sparse_blk_col_size
+        block_col_sizes = torch.tensor(
+            [sparse_blk_col_size] * num_block_col, device=self.device
+        )
+        block_col_sizes[-1] += last_col_len
+        return LMCFlashInferSparseMetadata(
+            wrapper,
+            seq_len,
+            self.num_qo_heads,
+            self.num_kv_heads,
+            self.head_dim,
+            block_col_sizes,
+            sparse_blk_row_size,
+            sparse_blk_col_size,
+            is_causal=True,
+        )

lmcache/v1/compute/attention/metadata.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from abc import abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+import abc
+
+# Third Party
+import torch
+
+if TYPE_CHECKING:
+    # First Party
+    from lmcache.v1.compute.attention.flash_infer_sparse import HackBSAWrapper
+
+
+@dataclass
+class LMCAttnMetadata(metaclass=abc.ABCMeta):
+    @abstractmethod
+    def update_from_top_indices(self, top_indices: torch.Tensor):
+        raise NotImplementedError("This method should be implemented in subclasses.")
+
+
+@dataclass
+class LMCFlashAttnMetadata(LMCAttnMetadata):
+    query_start_loc: torch.Tensor
+    seq_lens: torch.Tensor
+    cu_seqlens_k: torch.Tensor
+    max_query_len: torch.Tensor
+    max_seq_len: torch.Tensor
+
+    def update_from_top_indices(self, top_indices: torch.Tensor):
+        top_k_num = len(top_indices)
+        self.max_query_len = top_k_num
+        device = self.query_start_loc.device
+        dtype = self.query_start_loc.dtype
+        self.query_start_loc = torch.tensor([0, top_k_num], dtype=dtype, device=device)
+
+
+@dataclass
+class LMCFlashInferSparseMetadata(LMCAttnMetadata):
+    wrapper: "HackBSAWrapper"
+    seq_len: int
+    num_qo_heads: int
+    num_kv_heads: int
+    head_dim: int
+    block_col_sizes: torch.Tensor
+    sparse_blk_row_size: int = 32  # TODO(Jiayi): make this tunable
+    sparse_blk_col_size: int = 32  # TODO(Jiayi): make this tunable
+    is_causal: bool = True
+    q_data_dtype: torch.dtype = torch.bfloat16  # TODO(Jiayi): remove hardcode
+
+    def update_from_top_indices(self, top_indices: torch.Tensor):
+        # self.is_causal = False
+        device = top_indices.device
+        top_k_num = len(top_indices)
+        num_block_row = top_k_num // self.sparse_blk_row_size
+        block_row_sizes = torch.tensor(
+            [self.sparse_blk_row_size] * num_block_row, device=device
+        )
+        block_row_sizes[-1] += top_k_num % self.sparse_blk_row_size
+
+        block_mask_map = torch.zeros(
+            num_block_row, len(self.block_col_sizes), dtype=torch.bool, device=device
+        )
+        cols = torch.arange(block_mask_map.size(1), device=device).expand(
+            block_mask_map.size(0), -1
+        )
+
+        # NOTE(Jiayi): select every `sparse_blk_row_size`-th index from top_indices
+        # to approximate the attention mask at block level.
+        top_indices_slice = top_indices[
+            self.sparse_blk_row_size - 1 :: self.sparse_blk_row_size
+        ]
+        top_indices_slice //= self.sparse_blk_col_size
+        mask = cols < top_indices_slice.unsqueeze(1)
+        block_mask_map[mask] = 1
+        self.wrapper.plan(
+            block_mask_map.expand(self.num_kv_heads, -1, -1),
+            block_row_sizes.expand(self.num_kv_heads, -1),
+            self.block_col_sizes.expand(self.num_kv_heads, -1),
+            self.num_qo_heads,
+            self.num_kv_heads,
+            self.head_dim,
+            q_data_type=self.q_data_dtype,
+        )

lmcache/v1/compute/attention/utils.py

@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: Apache-2.0
+# Local
+from .flash_attn import LMCFlashAttnBackend
+from .flash_infer_sparse import LMCFlashInferSparseBackend
+
+
+def infer_attn_backend_from_vllm(vllm_attn, enable_sparse=False):
+    attn_name = type(vllm_attn.impl).__name__
+    if attn_name == "FlashInferImpl" and enable_sparse:
+        return LMCFlashInferSparseBackend(vllm_attn)
+    elif attn_name == "FlashAttentionImpl" and not enable_sparse:
+        return LMCFlashAttnBackend(vllm_attn)
+    else:
+        raise ValueError(f"Attention backend {attn_name} is not supported in LMCache.")

lmcache/v1/compute/blend/__init__.py

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+# First Party
+from lmcache.v1.compute.blend.utils import LMCBlenderBuilder
+
+__all__ = [
+    "LMCBlenderBuilder",
+]

lmcache/v1/compute/blend/blender.py

@@ -0,0 +1,168 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Optional, Union
+
+# Third Party
+import torch
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.compute.attention.metadata import LMCAttnMetadata
+from lmcache.v1.compute.blend.metadata import LMCBlendCommonMetadata, LMCBlendMetadata
+from lmcache.v1.compute.models.utils import infer_model_from_vllm
+from lmcache.v1.config import LMCacheEngineConfig
+
+logger = init_logger(__name__)
+
+
+class LMCBlender:
+    """
+    Cache-blender backend for LMCache.
+    This backend uses the Blender implementation for efficient blending computation.
+    """
+
+    def __init__(
+        self,
+        cache_engine,
+        gpu_connector,
+        vllm_model,
+        config: LMCacheEngineConfig,
+    ):
+        self.cache_engine = cache_engine
+        self.gpu_connector = gpu_connector
+
+        enable_sparse = False
+        if config.extra_config is not None:
+            enable_sparse = config.extra_config.get("enable_sparse", False)
+
+        self.layerwise_model = infer_model_from_vllm(vllm_model, self, enable_sparse)
+
+        # TODO: remove this hardcode
+        self.num_layers = len(vllm_model.model.layers)
+
+        # TODO(Jiayi): support threshold-based blending
+        # TODO(Jiayi): support different ratios for different layers
+        # TODO(Jiayi): support "skipping blending if hit too short"
+        self.common_metadata = LMCBlendCommonMetadata(
+            check_layers=config.blend_check_layers,
+            recomp_ratios=config.blend_recompute_ratios,
+            thresholds=config.blend_thresholds,
+        )
+
+        # This will be set during the blending process
+        self.metadata = LMCBlendMetadata(
+            imp_indices=None,
+            attn_mask=None,
+            positions=None,
+        )
+
+    def process_qkv(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        residual: torch.Tensor,
+        layer_id: int,
+        attn_output: Optional[torch.Tensor],
+        attn_metadata: LMCAttnMetadata,
+    ):
+        logger.debug(f"Blender is processing KV for layer {layer_id}")
+        old_k, old_v = self.gpu_connector.get_kv(layer_id)
+
+        if attn_output is None:
+            attn_output = torch.empty(
+                q.shape,
+                dtype=q.dtype,
+                device=q.device,
+            )
+
+        # perform positional encoding
+        if self.metadata.positions is None:
+            self.metadata.positions = torch.arange(
+                q.shape[0], device=q.device, dtype=torch.int64
+            )
+        layer = self.layerwise_model.vllm_model.model.layers[layer_id]
+        attn_layer = layer.self_attn
+        q, k = attn_layer.rotary_emb(self.metadata.positions, q, k)
+
+        if layer_id in self.common_metadata.check_layers:
+            diff_k = torch.sum(
+                (k.to(torch.float32) - old_k.to(torch.float32)) ** 2, dim=[1]
+            )
+            total_len = diff_k.shape[0]
+
+            assert self.common_metadata.recomp_ratios is not None
+
+            # TODO(Jiayi): remove `[0]` hardcode
+            topk_num = int(total_len * self.common_metadata.recomp_ratios[0])
+            topk_num = max(topk_num, 1)
+
+            top_indices = torch.topk(diff_k, k=topk_num).indices
+            top_indices, _ = torch.sort(top_indices)
+
+            k, v = k[top_indices], v[top_indices]
+            q = q[top_indices]
+            residual = residual[top_indices]
+
+            logger.debug(f"Number of indices picked: {len(top_indices)}")
+
+            self.metadata.imp_indices = top_indices
+            self.metadata.positions = self.metadata.positions[top_indices]
+            attn_output = attn_output[:topk_num]
+
+            attn_metadata.update_from_top_indices(top_indices)
+
+        if self.metadata.imp_indices is not None:
+            old_k[self.metadata.imp_indices] = k
+            old_v[self.metadata.imp_indices] = v
+            return q, old_k, old_v, residual, attn_output, attn_metadata
+        else:
+            return q, k, v, residual, attn_output, attn_metadata
+
+    # NOTE(Jiayi): Exposing this `blend_layer` interface as we might
+    # want to ochestrate the blending process elsewhere
+    def blend_layer(
+        self,
+        tokens: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        **kwargs,
+    ):
+        """
+        Perform layerwiese retrieve + blending.
+        """
+
+        # TODO(Jiayi): store is currently not included in this function
+
+        layerwise_model_executor = self.layerwise_model.compute_layer(tokens)
+        layerwise_retriever = self.cache_engine.retrieve_layer(tokens, mask, **kwargs)
+
+        next(layerwise_retriever)
+        yield
+
+        for i in range(self.num_layers):
+            next(layerwise_retriever)
+            next(layerwise_model_executor)
+            yield
+
+        next(layerwise_retriever)
+
+        self.metadata.clean()
+        yield
+
+    def blend(
+        self,
+        tokens: Union[torch.Tensor, list[int]],
+        mask: Optional[torch.Tensor] = None,
+        **kwargs,
+    ):
+        """
+        Perform blending for the given tokens.
+        """
+
+        if isinstance(tokens, list):
+            tokens = torch.tensor(tokens).cuda()
+
+        layerwise_blender = self.blend_layer(tokens, mask, **kwargs)
+
+        for i in range(self.num_layers + 2):
+            next(layerwise_blender)