Upload LCKVLlamaForCausalLM

README.md

@@ -0,0 +1,199 @@
+---
+library_name: transformers
+tags: []
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]

cache_utils.py

@@ -0,0 +1,550 @@
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+
+from transformers.cache_utils import Cache, DynamicCache, SinkCache
+
+from .utils import LayerTypeParser
+
+
+class IndexedCache(Cache):
+    """
+    Similar to the `DynamicCache` class, but with the ability to index the cache by layer index. DynamicCache
+    assumes that all layers compute KVs, while IndexedCache allows for a more flexible cache structure.
+    """
+    build_position_ids_based_on_cache = False
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.key_cache: Dict[int, torch.Tensor] = {}
+        self.value_cache: Dict[int, torch.Tensor] = {}
+        self._seen_tokens = 0  # Used in `generate` to keep tally of how many tokens the cache has seen
+        self._update = True # to prevent the cache from updating when inference with iterations
+
+    def __getitem__(self, layer_idx: int) -> List[Tuple[torch.Tensor]]:
+        """
+        Support for backwards-compatible `past_key_value` indexing, e.g. `past_key_value[0][0].shape[2]` to get the
+        sequence length.
+        """
+        if layer_idx in self.key_cache:
+            return (self.key_cache[layer_idx], self.value_cache[layer_idx])
+        else:
+            raise KeyError(f"Cache only has {len(self)} layers, attempted to access layer with index {layer_idx}")
+
+    def __iter__(self):
+        """
+        Support for backwards-compatible `past_key_value` iteration, e.g. `for x in past_key_value:` to iterate over
+        keys and values
+        """
+        for layer_idx in sorted(self.key_cache.keys()):
+            yield (self.key_cache[layer_idx], self.value_cache[layer_idx])
+
+    def __len__(self):
+        """
+        Support for backwards-compatible `past_key_value` length, e.g. `len(past_key_value)`. This value corresponds
+        to the number of layers that compute KVs in the model.
+        """
+        return len(self.key_cache)
+
+    @property
+    def min_layer(self) -> int:
+        return min(self.key_cache.keys()) if len(self.key_cache) > 0 else None
+
+    def is_min_layer(self, layer_idx: int) -> bool:
+        return self.min_layer is None or self.min_layer == layer_idx
+
+    def update(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        layer_idx: int,
+        cache_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`.
+
+        Parameters:
+            key_states (`torch.Tensor`):
+                The new key states to cache.
+            value_states (`torch.Tensor`):
+                The new value states to cache.
+            layer_idx (`int`):
+                The index of the layer to cache the states for.
+            cache_kwargs (`Dict[str, Any]`, `optional`):
+                Additional arguments for the cache subclass. No additional arguments are used in `DynamicCache`.
+
+        Return:
+            A tuple containing the updated key and value states.
+        """
+        # Update the number of seen tokens
+        if self.is_min_layer(layer_idx):
+            self._seen_tokens += key_states.shape[-2]
+
+        # Retrieve the cache
+        if layer_idx not in self.key_cache:
+            new_key_states = key_states
+            new_value_states = value_states
+        else:
+            new_key_states = torch.cat([self.key_cache[layer_idx], key_states], dim=-2)
+            new_value_states = torch.cat([self.value_cache[layer_idx], value_states], dim=-2)
+
+        # Update the cache
+        if self._update:
+            self.key_cache[layer_idx] = new_key_states
+            self.value_cache[layer_idx] = new_value_states
+
+        return new_key_states, new_value_states
+
+    def get_seq_length(self, layer_idx: Optional[int] = None) -> int:
+        """Returns the sequence length of the cached states. A layer index can be optionally passed."""
+        if layer_idx is None:
+            layer_idx = self.min_layer
+
+        # TODO: deprecate this function in favor of `cache_position`
+        is_empty_layer = (
+            (len(self.key_cache) == 0)  # no cache in any layer
+            or (layer_idx not in self.key_cache)  # skipped `layer_idx` and hasn't run a layer with cache after it
+        )
+        layer_seq_length = self.key_cache[layer_idx].shape[-2] if not is_empty_layer else 0
+        return layer_seq_length
+
+    def get_max_length(self) -> Optional[int]:
+        """Returns the maximum sequence length of the cached states. IndexedCache does not have a maximum length."""
+        return None
+
+    @classmethod
+    def from_cache(cls, dynamic_cache: DynamicCache, *args, **kwargs) -> "IndexedCache":
+        """Converts a dynamic cache into an equivalent `IndexedCache`."""
+        cache = cls(*args, **kwargs)
+
+        cache._seen_tokens = dynamic_cache._seen_tokens
+        for layer_idx in range(len(dynamic_cache.key_cache)):
+            key_states, value_states = dynamic_cache[layer_idx]
+            cache.update(key_states, value_states, layer_idx)
+
+        return cache
+
+
+class IndexedSinkCache(Cache):
+    """
+    This is a fix to the SinkCache class in the transformers library. It also allows for the cache to be indexed by
+    layer index, similar to the `IndexedCache` class.
+    """
+    build_position_ids_based_on_cache = True
+
+    def __init__(self, window_length: int = None, num_sink_tokens: int = None) -> None:
+        super().__init__()
+        self.key_cache: Dict[int, torch.Tensor] = {}
+        self.value_cache: Dict[int, torch.Tensor] = {}
+        self.window_length = window_length
+        self.num_sink_tokens = num_sink_tokens
+        self.cos_sin_rerotation_cache = {}
+        self._cos_cache = None
+        self._sin_cache = None
+        self._seen_tokens = 0  # Used in `generate` to keep tally of how many tokens the cache has seen
+        self._update = True  # to prevent the cache from updating when inference with iterations
+
+    @staticmethod
+    def _rotate_half(x):
+        x1 = x[..., : x.shape[-1] // 2]
+        x2 = x[..., x.shape[-1] // 2 :]
+        return torch.cat((-x2, x1), dim=-1)
+
+    def _apply_key_rotary_pos_emb(
+        self, key_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+    ) -> torch.Tensor:
+        rotated_key_states = (key_states * cos) + (self._rotate_half(key_states) * sin)
+        return rotated_key_states
+
+    def _get_rerotation_cos_sin(
+        self, offset: int, dtype: torch.dtype, cos: torch.Tensor, sin: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if offset not in self.cos_sin_rerotation_cache:
+            # Upcast to float32 temporarily for better accuracy
+            cos = cos.to(torch.float32)
+            sin = sin.to(torch.float32)
+
+            # Compute the cos and sin required for back- and forward-rotating to one position earlier in the sequence
+            original_cos = cos[self.num_sink_tokens + offset :]
+            shifted_cos = cos[self.num_sink_tokens : -offset]
+            original_sin = sin[self.num_sink_tokens + offset :]
+            shifted_sin = sin[self.num_sink_tokens : -offset]
+            rerotation_cos = original_cos * shifted_cos + original_sin * shifted_sin
+            rerotation_sin = -original_sin * shifted_cos + original_cos * shifted_sin
+
+            self.cos_sin_rerotation_cache[offset] = (
+                rerotation_cos.to(dtype).unsqueeze(0),
+                rerotation_sin.to(dtype).unsqueeze(0),
+            )
+        return self.cos_sin_rerotation_cache[offset]
+
+    @property
+    def min_layer(self) -> int:
+        return min(self.key_cache.keys()) if len(self.key_cache) > 0 else None
+
+    def is_min_layer(self, layer_idx: int) -> bool:
+        return self.min_layer is None or self.min_layer == layer_idx
+
+    def get_seq_length(self, layer_idx: Optional[int] = None) -> int:
+        """Returns the sequence length of the cached states. A layer index can be optionally passed."""
+        # TODO: deprecate this function in favor of `cache_position`
+        # Workaround to make 'key_states.shape[-2] + past_key_value.get_seq_length(self.layer_idx)' <= window_length
+        if layer_idx is None:
+            layer_idx = self.min_layer
+
+        if layer_idx not in self.key_cache:
+            return 0
+
+        return self.key_cache[layer_idx].shape[-2]
+
+    def get_max_length(self) -> Optional[int]:
+        """Returns the maximum sequence length of the cached states."""
+        return self.window_length
+
+    def update(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        layer_idx: int,
+        cache_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`.
+
+        Parameters:
+            key_states (`torch.Tensor`):
+                The new key states to cache.
+            value_states (`torch.Tensor`):
+                The new value states to cache.
+            layer_idx (`int`):
+                The index of the layer to cache the states for.
+            cache_kwargs (`Dict[str, Any]`, `optional`):
+                Additional arguments for the cache subclass. The following arguments can be used in `SinkCache`: `sin`,
+                `cos` and `partial_rotation_size`. These arguments are used with models using RoPE, to recompute the
+                rotation as the tokens are shifted.
+
+        Return:
+            A tuple containing the updated key and value states.
+        """
+        # Optional kwargs for `SinkCache` -- needed on models using RoPE. `partial_rotation_size` is used on models
+        # with partially rotated position embeddings, like Phi or Persimmon.
+        sin = cache_kwargs.get("sin")
+        cos = cache_kwargs.get("cos")
+        partial_rotation_size = cache_kwargs.get("partial_rotation_size")
+        using_rope = cos is not None and sin is not None
+
+        # Update the number of seen tokens
+        if self.is_min_layer(layer_idx):
+            self._seen_tokens += key_states.shape[-2]
+
+        # Update the sin/cos cache, which holds sin/cos values for all possible positions
+        if using_rope and self.is_min_layer(layer_idx):
+            # BC: some models still pass `sin`/`cos` with 2 dims. In those models, they are the full sin/cos. Remove
+            # after all RoPE models have a llama-like cache utilization.
+            if cos.dim() == 2:
+                self._cos_cache = cos
+                self._sin_cache = sin
+            else:
+                if self._cos_cache is None:
+                    self._cos_cache = cos[0, ...]
+                    self._sin_cache = sin[0, ...]
+                elif self._cos_cache.shape[0] < self.window_length + key_states.shape[-2]:
+                    self._cos_cache = torch.cat([self._cos_cache[: self.window_length], cos[0, ...]], dim=0)
+                    self._sin_cache = torch.cat([self._sin_cache[: self.window_length], sin[0, ...]], dim=0)
+
+        # [bsz, num_heads, seq_len, head_dim]
+        if layer_idx not in self.key_cache:
+            # Empty cache
+            new_key_states = key_states
+            new_value_states = value_states
+
+        else:
+            # Growing cache
+            new_key_states = torch.cat([self.key_cache[layer_idx], key_states], dim=-2)
+            new_value_states = torch.cat([self.value_cache[layer_idx], value_states], dim=-2)
+
+        if self._update:
+            self.key_cache[layer_idx] = new_key_states
+            self.value_cache[layer_idx] = new_value_states
+
+        # If the cache is full, we need to shift the cache
+        if (seq_length := self.get_seq_length(layer_idx)) > self.window_length:
+            # Shifting cache
+            keys_to_keep = self.key_cache[layer_idx][:, :, -self.window_length + self.num_sink_tokens :]
+
+            # On RoPE models, we need to recompute the Key rotation as the tokens are shifted
+            if using_rope:
+                rerotation_cos, rerotation_sin = self._get_rerotation_cos_sin(
+                    seq_length - self.window_length,
+                    key_states.dtype,
+                    self._cos_cache[:seq_length],
+                    self._sin_cache[:seq_length],
+                )
+                if partial_rotation_size is not None:
+                    keys_to_keep, keys_pass = (
+                        keys_to_keep[..., :partial_rotation_size],
+                        keys_to_keep[..., partial_rotation_size:],
+                    )
+                keys_to_keep = self._apply_key_rotary_pos_emb(keys_to_keep, rerotation_cos, rerotation_sin)
+                if partial_rotation_size is not None:
+                    keys_to_keep = torch.cat((keys_to_keep, keys_pass), dim=-1)
+
+            # Concatenate sink tokens, shifted & rotated tokens (if needed), and new tokens
+            sink_keys = self.key_cache[layer_idx][:, :, : self.num_sink_tokens]
+            self.key_cache[layer_idx] = torch.cat([sink_keys, keys_to_keep], dim=-2)
+
+            sink_values = self.value_cache[layer_idx][:, :, : self.num_sink_tokens]
+            values_to_keep = self.value_cache[layer_idx][:, :, -self.window_length + self.num_sink_tokens :]
+            self.value_cache[layer_idx] = torch.cat([sink_values, values_to_keep], dim=-2)
+
+        return new_key_states, new_value_states
+
+    @classmethod
+    def from_cache(cls, sink_cache: SinkCache, *args, **kwargs) -> "IndexedSinkCache":
+        """Converts a dynamic cache into an equivalent `IndexedCache`."""
+        cache = cls(*args, **kwargs)
+
+        cache.window_length = sink_cache.window_length
+        cache.num_sink_tokens = sink_cache.num_sink_tokens
+        cache._seen_tokens = sink_cache._seen_tokens
+        cache._cos_cache = sink_cache._cos_cache
+        cache._sin_cache = sink_cache._sin_cache
+        cache.cos_sin_rerotation_cache = sink_cache.cos_sin_rerotation_cache
+        for layer_idx in range(len(sink_cache.key_cache)):
+            cache.key_cache[layer_idx] = sink_cache.key_cache[layer_idx]
+            cache.value_cache[layer_idx] = sink_cache.value_cache[layer_idx]
+
+        return cache
+
+
+class IndexedSlidingWindowCache(IndexedCache):
+    """
+    Similar to the `SlidingWindowCache` class, but with the ability to index the cache by layer index. It is no longer
+    a subclass of `StaticCache` as it is dynamic.
+    """
+    build_position_ids_based_on_cache = False
+
+    def __init__(self, sliding_window: int = None) -> None:
+        super().__init__()
+        self.sliding_window = sliding_window
+
+    def update(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        layer_idx: int,
+        cache_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[torch.Tensor]:
+        # Update the number of seen tokens
+        if self.is_min_layer(layer_idx):
+            self._seen_tokens += key_states.shape[-2]
+
+        # [bsz, num_heads, seq_len, head_dim]
+        if layer_idx not in self.key_cache:
+            # Empty cache
+            new_key_states = key_states
+            new_value_states = value_states
+
+        else:
+            # Growing cache
+            new_key_states = torch.cat([self.key_cache[layer_idx], key_states], dim=-2)
+            new_value_states = torch.cat([self.value_cache[layer_idx], value_states], dim=-2)
+
+        if self._update:
+            self.key_cache[layer_idx] = new_key_states
+            self.value_cache[layer_idx] = new_value_states
+
+        # If the cache is full, we need to shift the cache
+        if self.get_seq_length(layer_idx) > self.sliding_window:
+            self.key_cache[layer_idx] = self.key_cache[layer_idx][:, :, -self.sliding_window :]
+            self.value_cache[layer_idx] = self.value_cache[layer_idx][:, :, -self.sliding_window :]
+
+        return new_key_states, new_value_states
+
+    def get_max_length(self) -> Optional[int]:
+        return self.sliding_window
+
+    @classmethod
+    def from_cache(cls, sliding_window_cache: "IndexedSlidingWindowCache", *args, **kwargs) -> "IndexedSlidingWindowCache":
+        """This is to override the `from_cache` method in the `IndexedCache` class."""
+        cache = cls(*args, **kwargs)
+
+        cache._seen_tokens = sliding_window_cache._seen_tokens
+        cache.sliding_window = sliding_window_cache.sliding_window
+        for layer_idx in range(len(sliding_window_cache.key_cache)):
+            cache.key_cache[layer_idx] = sliding_window_cache.key_cache[layer_idx]
+            cache.value_cache[layer_idx] = sliding_window_cache.value_cache[layer_idx]
+
+        return cache
+
+
+class IndexedHybridCache(IndexedSlidingWindowCache, IndexedCache):
+    """
+    Hybrid Cache class to be used for models that alternate between a local sliding window attention and global
+    attention in every other layer. Under the hood, Hybrid Cache leverages ["IndexedSlidingWindowCache"] for
+    sliding window attention and ["IndexedCache"] for global attention.
+    """
+    build_position_ids_based_on_cache = False
+
+    def __init__(self, parser: LayerTypeParser = None, sliding_window: int = None) -> None:
+        super().__init__(sliding_window=sliding_window)
+        self.parser = parser
+
+    def update(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        layer_idx: int,
+        cache_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[torch.Tensor]:
+        if self.parser[layer_idx].use_sliding_window:
+            return IndexedSlidingWindowCache.update(self, key_states, value_states, layer_idx, cache_kwargs)
+        else:
+            return IndexedCache.update(self, key_states, value_states, layer_idx, cache_kwargs)
+
+    def get_max_length(self) -> Optional[int]:
+        return IndexedCache.get_max_length(self)
+
+    @classmethod
+    def from_cache(cls, hybrid_cache: "IndexedHybridCache", *args, **kwargs) -> "IndexedHybridCache":
+        """This is to override the `from_cache` method in the `IndexedSlidingWindowCache` class."""
+        cache = cls(*args, **kwargs)
+
+        cache._seen_tokens = hybrid_cache._seen_tokens
+        cache.sliding_window = hybrid_cache.sliding_window
+        cache.parser = hybrid_cache.parser
+        for layer_idx in range(len(hybrid_cache.key_cache)):
+            cache.key_cache[layer_idx] = hybrid_cache.key_cache[layer_idx]
+            cache.value_cache[layer_idx] = hybrid_cache.value_cache[layer_idx]
+
+        return cache
+
+
+class LayerCache(torch.nn.Module):
+    """
+    A cache for storing the key-value pairs for layers.
+    """
+    def __init__(self) -> None:
+        """
+        The placeholder is used to expand the key-value pairs if the layer attends to the top layers.
+        Size: (batch_size, num_key_value_heads, 1, head_dim)
+        """
+        super().__init__()
+        self.key_layer_cache: Dict[int, torch.Tensor] = {}
+        self.value_layer_cache: Dict[int, torch.Tensor] = {}
+        self.layer_type = None
+        self.placeholder = None
+
+    def setup(self, placeholder: torch.Tensor):
+        """setup the cache, calling this function is necessary if there is a layer that attends to the top layers"""
+        self.placeholder = placeholder
+
+    def initialize(self, parser: LayerTypeParser, sequence_length: int):
+        """initialize the cache"""
+        layers_to_init = {parser[idx].attends_to for idx in range(len(parser)) if parser[idx].attends_top}
+
+        if layers_to_init:
+            b, h, _, d = self.placeholder.size()
+            init_kvs = self.placeholder.new_zeros((b, h, sequence_length, d))
+
+            for layer_idx in layers_to_init:
+                self.layer_append(layer_idx, init_kvs, init_kvs)
+
+    def layer_get(self, layer_idx: int, zerofill: bool = False) -> Tuple[torch.Tensor, torch.Tensor]:
+        key_states = self.key_layer_cache.get(layer_idx, None)
+        value_states = self.value_layer_cache.get(layer_idx, None)
+
+        if zerofill:
+            if key_states is None:
+                key_states = self.placeholder
+                value_states = self.placeholder
+            else:
+                key_states = torch.cat([self.placeholder, key_states], dim=2)
+                value_states = torch.cat([self.placeholder, value_states], dim=2)
+
+        return key_states, value_states
+
+    def layer_set(self, layer_idx: int, key: torch.Tensor, value: torch.Tensor):
+        self.key_layer_cache[layer_idx] = key
+        self.value_layer_cache[layer_idx] = value
+
+    def layer_append(self, layer_idx: int, key: torch.Tensor, value: torch.Tensor):
+        if layer_idx not in self.key_layer_cache:
+            self.key_layer_cache[layer_idx] = key
+            self.value_layer_cache[layer_idx] = value
+        else:
+            self.key_layer_cache[layer_idx] = torch.cat([self.key_layer_cache[layer_idx], key], dim=2)
+            self.value_layer_cache[layer_idx] = torch.cat([self.value_layer_cache[layer_idx], value], dim=2)
+
+
+class LayerIndexedCache(LayerCache, IndexedCache):
+    """
+    A cache for storing the key-value pairs for layers, in combination with the ability of standard KV cache.
+    """
+    def __init__(self) -> None:
+        LayerCache.__init__(self)
+        IndexedCache.__init__(self)
+
+
+class LayerIndexedSinkCache(LayerCache, IndexedSinkCache):
+    """
+    A cache for storing the key-value pairs for layers, in combination with the ability of sink KV cache.
+    """
+    def __init__(self) -> None:
+        LayerCache.__init__(self)
+        IndexedSinkCache.__init__(self)
+
+
+class LayerIndexedSlidingWindowCache(LayerCache, IndexedSlidingWindowCache):
+    """
+    A cache for storing the key-value pairs for layers, in combination with the ability of sliding window KV cache.
+    """
+    def __init__(self) -> None:
+        LayerCache.__init__(self)
+        IndexedSlidingWindowCache.__init__(self)
+
+
+class LayerIndexedHybridCache(LayerCache, IndexedHybridCache):
+    """
+    A cache for storing the key-value pairs for layers, in combination with the ability of hybrid KV cache.
+    """
+    def __init__(self) -> None:
+        LayerCache.__init__(self)
+        IndexedHybridCache.__init__(self)
+
+
+class AutoLayerCache(torch.nn.Module):
+    """
+    AutoLayerCache is a module that automatically creates a cache from an existing cache.
+    """
+    CACHE_MAPPING = {
+        DynamicCache: LayerIndexedCache,
+        SinkCache: LayerIndexedSinkCache,
+        IndexedSlidingWindowCache: LayerIndexedSlidingWindowCache,
+        IndexedHybridCache: LayerIndexedHybridCache,
+    }
+
+    def __init__(self, *args, **kwargs):
+        raise RuntimeError(
+            f"{self.__class__.__name__} is designed to be instantiated "
+            f"using the `{self.__class__.__name__}.from_cache(cache)` method."
+        )
+
+    @classmethod
+    def from_cache(cls, cache: Cache, *args, **kwargs):
+        """
+        Create a new cache from an existing cache. The new cache will have the same type as the original cache.
+        """
+        cache_type = type(cache)
+        if cache_type not in cls.CACHE_MAPPING:
+            raise ValueError(f"Cache type {cache_type} is not supported by {cls.__name__}.")
+
+        cache_class = cls.CACHE_MAPPING[cache_type]
+
+        if hasattr(cache_class, "from_cache"):
+            return cache_class.from_cache(cache, *args, **kwargs)
+        else:
+            # we init an empty cache and copy the attributes
+            new_cache = cache_class(*args, **kwargs)
+            new_cache.__dict__.update(cache.__dict__)
+            return new_cache

config.json

@@ -0,0 +1,41 @@
+{
+  "_name_or_path": "outputs/tinyllama-lckv-w10-ft-250b",
+  "architectures": [
+    "LCKVLlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_lckv.LCKVLlamaConfig",
+    "AutoModelForCausalLM": "modeling_lckv.LCKVLlamaForCausalLM"
+  },
+  "backward_passes": 2,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "force_nodiag": true,
+  "forward_passes": 7,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intermediate_size": 5632,
+  "layer_types": "0_1_2_3_4_16_16_16_16_16_16_16_16_16_16_16_16_17_18_19_20_21",
+  "max_position_embeddings": 2048,
+  "mlp_bias": false,
+  "model_type": "lckv-llama",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 22,
+  "num_key_value_heads": 4,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": null,
+  "rope_theta": 10000.0,
+  "sliding_window": 4096,
+  "tie_word_embeddings": false,
+  "tokenizer_class": "LlamaTokenizer",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.45.2",
+  "use_cache": true,
+  "use_sequential": false,
+  "vocab_size": 32000
+}

configuration_lckv.py

@@ -0,0 +1,81 @@
+# coding=utf-8
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" LCKV LLaMA model configuration"""
+from transformers.models.llama.configuration_llama import LlamaConfig
+
+from .utils import LayerTypeParser
+
+
+class LCKVLlamaConfig(LlamaConfig):
+
+    model_type = "lckv-llama"
+
+    def __init__(
+        self,
+        layer_types: str = None,
+        forward_passes: int = 7,
+        backward_passes: int = 2,
+        sliding_window: int = 4096,
+        use_sequential: bool = False,
+        force_nodiag: bool = False,
+        **kwargs,
+    ):
+        """
+        Initialize a LCKV LLaMA configuration. Instantiating a configuration with the defaults
+        will yield a similar configuration to that of the LLaMA-7B with the standard transformer
+        training scheme.
+
+        Args:
+            layer_types (`str`, *optional*):
+                A string of integers separated by underscores. The i-th integer means the layer
+                will use the key-value pair in the i-th layer as the kv cache. Special characters
+                may be placed after the integers:
+                - `s` means the layer will use sliding window attention.
+                The default value is "0_1_2_..." till the number of layers in the current config.
+            forward_passes (`int`, *optional*, defaults to 7):
+                The number of forward passes during training and prompt encoding. Equivlent
+                to `m` in the paper.
+            backward_passes (`int`, *optional*, defaults to 2):
+                The number of backward passes during training and prompt encoding. Equivlent
+                to `b` in the paper.
+            sliding_window (`int`, *optional*, defaults to 4096):
+                Sliding window attention window size. If not specified, will default to `4096`.
+                It will only be effective if the corresponding layer uses sliding window attention.
+            use_sequential (`bool`, *optional*, defaults to False):
+                Whether to do forwarding sequentially, token by token. Useful for testing purpose
+                for models with cyclic dependency. Also can be used for sequential training.
+            force_nodiag (`bool`, *optional*, defaults to False):
+                Whether to force the model to not use the diagonal attention. By default, the model
+                will mask the diagonal attention only in layers necessary. If set to `True`, the model
+                will never use the diagonal attention in any layer. This is mainly for backward compatibility.
+        """
+        super().__init__(**kwargs)
+        self.layer_types = layer_types
+        self.forward_passes = forward_passes
+        self.backward_passes = backward_passes
+        self.sliding_window = sliding_window
+        self.use_sequential = use_sequential
+        self.force_nodiag = force_nodiag
+
+        if self.layer_types is None:
+            self.layer_types = "_".join(map(str, range(self.num_hidden_layers)))
+
+        # post check
+        LayerTypeParser(self.layer_types).check(self.num_hidden_layers)

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "transformers_version": "4.45.2"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6d07eabdbaa1683be0a28d1af9a6495ee3f4201b9d04bd9ab9ac4e831188c9b2
+size 2177048648

modeling_lckv.py

@@ -0,0 +1,807 @@
+# coding=utf-8
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" PyTorch LLaMA model."""
+import copy
+import math
+from typing import List, Optional, Tuple, Union
+
+import torch
+from torch import nn
+
+from transformers.cache_utils import Cache, StaticCache
+from transformers.modeling_attn_mask_utils import AttentionMaskConverter
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from transformers.models.llama.modeling_llama import (
+    LLAMA_INPUTS_DOCSTRING,
+    LlamaAttention,
+    LlamaDecoderLayer,
+    LlamaForCausalLM,
+    LlamaModel,
+    LlamaPreTrainedModel,
+    _prepare_4d_causal_attention_mask_with_cache_position,
+    logger,
+    repeat_kv,
+    rotate_half,
+)
+from transformers.utils import add_start_docstrings_to_model_forward, is_flash_attn_greater_or_equal_2_10
+
+from .cache_utils import AutoLayerCache, LayerCache
+from .configuration_lckv import LCKVLlamaConfig
+from .utils import IterStep, LayerTypeParser, flash_attention_forward
+
+
+def apply_rotary(q, cos, sin, unsqueeze_dim=1):
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    return q_embed
+
+
+class LCKVLlamaAttention(LlamaAttention):
+    """
+    LCKV Attention may not need to initialize weights for the key and value projections.
+    """
+
+    def __init__(self, config: LCKVLlamaConfig, layer_idx: Optional[int] = None):
+        super().__init__(config, layer_idx)
+        self.layer_type = LayerTypeParser(config.layer_types)[layer_idx]
+        self.sliding_window = config.sliding_window if self.layer_type.use_sliding_window else None
+
+        # Some layers may not need to compute key-value pairs
+        if not self.layer_type.computes_kv:
+            del self.k_proj
+            del self.v_proj
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # will become mandatory in v4.46
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        bsz, q_len, _ = hidden_states.size()
+        cos, sin = position_embeddings
+
+        query_states = self.q_proj(hidden_states)
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        query_states = apply_rotary(query_states, cos, sin)
+
+        # compute key and value states
+        if self.layer_type.computes_kv:
+            key_states = self.k_proj(hidden_states)
+            value_states = self.v_proj(hidden_states)
+
+            key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+            value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+            key_states = apply_rotary(key_states, cos, sin)
+
+            if isinstance(past_key_value, Cache):
+                # sin and cos are specific to RoPE models; cache_position needed for the static cache
+                cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+                key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+            past_key_value.layer_set(self.layer_idx, key_states, value_states)
+
+        # get the cached key and value states
+        # if the layer attends to the top layers, there are two cases:
+        # 1. the query length is 1, in which case we will not do iterative updates. Therefore, the kv lacks the current
+        #    query length and we need to fill it with zeros.
+        # 2. the query length is greater than 1, in which case we will do iterative updates and the kv will have the
+        #    correct query length.
+        key_states, value_states = past_key_value.layer_get(
+            self.layer_type.attends_to,
+            zerofill=self.layer_type.attends_top and q_len == 1,
+        )
+
+        # handle GQA
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+
+        if attention_mask is not None:  # no matter the length, we just slice it
+            causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+            attn_weights = attn_weights + causal_mask
+
+        # diagonal mask from the right bottom corner
+        if self.config.force_nodiag or self.layer_type.attends_top:
+            kv_len = key_states.size(2)
+            mask = attn_weights.new_full((q_len, kv_len), torch.finfo(attn_weights.dtype).min)
+            mask = mask.tril(diagonal=kv_len - q_len).triu(diagonal=kv_len - q_len)
+            attn_weights = attn_weights + mask
+
+        # sliding window mask
+        if self.sliding_window:
+            kv_len = key_states.size(2)
+            mask = attn_weights.new_full((q_len, kv_len), torch.finfo(attn_weights.dtype).min)
+            mask = mask.tril(diagonal=kv_len - q_len - self.sliding_window)
+            attn_weights = attn_weights + mask
+
+        # upcast attention to fp32
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, -1)
+        attn_output = self.o_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+
+class LCKVLlamaFlashAttention2(LCKVLlamaAttention):
+    """
+    LCKV Attention may not need to initialize weights for the key and value projections.
+    """
+
+    def __init__(self, config: LCKVLlamaConfig, layer_idx: Optional[int] = None):
+        super().__init__(config, layer_idx)
+
+        # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
+        # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
+        # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
+        self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[LayerCache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # will become mandatory in v4.46
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+
+        output_attentions = False
+
+        bsz, q_len, _ = hidden_states.size()
+        cos, sin = position_embeddings
+
+        # Flash attention requires the input to have the shape
+        # batch_size x seq_length x head_dim x hidden_dim
+        # therefore we just need to keep the original shape
+        query_states = self.q_proj(hidden_states)
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        query_states = apply_rotary(query_states, cos, sin)
+
+        # compute key and value states
+        if self.layer_type.computes_kv:
+            key_states = self.k_proj(hidden_states)
+            value_states = self.v_proj(hidden_states)
+
+            key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+            value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+            key_states = apply_rotary(key_states, cos, sin)
+
+            if isinstance(past_key_value, Cache):
+                # sin and cos are specific to RoPE models; cache_position needed for the static cache
+                cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+                key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+            past_key_value.layer_set(self.layer_idx, key_states, value_states)
+
+        # get the cached key and value states
+        # if the layer attends to the top layers, there are two cases:
+        # 1. the query length is 1, in which case we will not do iterative updates. Therefore, the kv lacks the current
+        #    query length and we need to fill it with zeros.
+        # 2. the query length is greater than 1, in which case we will do iterative updates and the kv will have the
+        #    correct query length.
+        key_states, value_states = past_key_value.layer_get(
+            self.layer_type.attends_to,
+            zerofill=self.layer_type.attends_top and q_len == 1,
+        )
+
+        # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
+        # to be able to avoid many of these transpose/reshape/view.
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        dropout_rate = self.attention_dropout if self.training else 0.0
+
+        # In PEFT, usually we cast the layer norms in float32 for training stability reasons
+        # therefore the input hidden states gets silently casted in float32. Hence, we need
+        # cast them back in the correct dtype just to be sure everything works as expected.
+        # This might slowdown training & inference so it is recommended to not cast the LayerNorms
+        # in fp32. (LlamaRMSNorm handles it correctly)
+
+        input_dtype = query_states.dtype
+        if input_dtype == torch.float32:
+            if torch.is_autocast_enabled():
+                target_dtype = torch.get_autocast_gpu_dtype()
+            # Handle the case where the model is quantized
+            elif hasattr(self.config, "_pre_quantization_dtype"):
+                target_dtype = self.config._pre_quantization_dtype
+            else:
+                target_dtype = self.q_proj.weight.dtype
+
+            logger.warning_once(
+                f"The input hidden states seems to be silently casted in float32, this might be related to"
+                f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
+                f" {target_dtype}."
+            )
+
+            query_states = query_states.to(target_dtype)
+            key_states = key_states.to(target_dtype)
+            value_states = value_states.to(target_dtype)
+
+        attn_output = flash_attention_forward(
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            q_len,
+            position_ids=position_ids,
+            dropout=dropout_rate,
+            sliding_window=self.sliding_window,
+            use_top_left_mask=self._flash_attn_uses_top_left_mask,
+            is_causal=self.is_causal,
+            no_diag=(self.config.force_nodiag or self.layer_type.attends_top),
+        )
+
+        attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+
+LCKV_LLAMA_ATTENTION_CLASSES = {
+    "eager": LCKVLlamaAttention,
+    "flash_attention_2": LCKVLlamaFlashAttention2,
+}
+
+
+class LCKVLlamaDecoderLayer(LlamaDecoderLayer):
+    def __init__(self, config: LCKVLlamaConfig, layer_idx: int):
+        super().__init__(config, layer_idx)
+        self.self_attn = LCKV_LLAMA_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx)
+
+
+class LCKVLlamaPreTrainedModel(LlamaPreTrainedModel):
+    config_class = LCKVLlamaConfig
+    supports_gradient_checkpointing = False # not tested yet
+    _no_split_modules = ["LCKVLlamaDecoderLayer"]
+    _supports_flash_attn_2 = True
+    _supports_sdpa = False
+
+
+class LCKVLlamaModel(LCKVLlamaPreTrainedModel, LlamaModel):
+    def __init__(self, config: LCKVLlamaConfig):
+        LCKVLlamaPreTrainedModel.__init__(self, config)
+        LlamaModel.__init__(self, copy.deepcopy(config)) # copy config to avoid modifying the original
+        self.layers = nn.ModuleList([LCKVLlamaDecoderLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)])
+        self.parser = LayerTypeParser(config.layer_types)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[LayerCache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
+            )
+
+        if self.gradient_checkpointing and self.training and use_cache:
+            logger.warning_once(
+                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
+            )
+            use_cache = False
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        # build the cache object
+        if not isinstance(past_key_values, LayerCache):
+            placeholder = inputs_embeds.new_zeros(
+                inputs_embeds.shape[0],
+                self.config.num_key_value_heads,
+                1,
+                getattr(self.config, "head_dim", self.config.hidden_size // self.config.num_attention_heads)
+            )
+
+            if past_key_values is None:
+                past_key_values = LayerCache()
+            elif isinstance(past_key_values, Cache):
+                past_key_values = AutoLayerCache.from_cache(past_key_values)
+            else:
+                raise NotImplementedError("Only DynamicCache is supported for now.")
+
+            past_key_values.setup(placeholder)
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if isinstance(past_key_values, Cache) else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        causal_mask = self._update_causal_mask(
+            attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
+        )
+        hidden_states = inputs_embeds
+
+        # create position embeddings to be shared across the decoder layers
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # whether to forward sequentially
+        use_sequential = (
+            self.config.use_sequential
+            or inputs_embeds.shape[1] <= self.config.forward_passes + self.config.backward_passes
+            and self.parser.attends_top()
+        )
+
+        if use_sequential:
+
+            iteration_outputs = self._modeling_sequential(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                output_hidden_states=output_hidden_states,
+            )
+
+        else:
+
+            # initialize the cache
+            past_key_values.initialize(self.parser, inputs_embeds.shape[1])
+
+            # we need to do forward passes based on a plan if the input is a prompt
+            plan = self.parser.iteration_plan(self.config.forward_passes, self.config.backward_passes)
+
+            iteration_outputs = self._modeling_with_plan(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                output_hidden_states=output_hidden_states,
+                modeling_plan=plan,
+            )
+
+        hidden_states = iteration_outputs.last_hidden_state
+        all_hidden_states = iteration_outputs.hidden_states
+        all_self_attns = iteration_outputs.attentions
+        next_decoder_cache = iteration_outputs.past_key_values
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = next_decoder_cache if use_cache else None
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+    def _iterate_layers(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[LayerCache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        output_hidden_states: Optional[bool] = False,
+        layer_slice: Optional[slice] = None,
+    ) -> BaseModelOutputWithPast:
+        """
+        Iterates over the layers of the model, calling each layer in turn.
+        """
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        next_decoder_cache = None
+
+        # layers to compute
+        if layer_slice is None:
+            layer_slice = slice(None)
+
+        for decoder_layer in self.layers[layer_slice]:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    decoder_layer.__call__,
+                    hidden_states,
+                    attention_mask,
+                    position_ids,
+                    past_key_values,
+                    output_attentions,
+                    use_cache,
+                    cache_position,
+                    position_embeddings,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_values,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                    cache_position=cache_position,
+                    position_embeddings=position_embeddings,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache = layer_outputs[2 if output_attentions else 1]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        next_cache = next_decoder_cache if use_cache else None
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+    def _modeling_with_plan(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[LayerCache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        output_hidden_states: Optional[bool] = False,
+        modeling_plan: List[IterStep] = None,
+    ) -> BaseModelOutputWithPast:
+        """
+        Given a plan, iteratively update the hidden states.
+        """
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        next_decoder_cache = None
+
+        for step in modeling_plan:
+            end = len(self.layers) if step.layer_slice.stop is None else step.layer_slice.stop
+            iteration_func = self._iterate_layers if step.requires_grad else torch.no_grad()(self._iterate_layers)
+
+            if isinstance(past_key_values, Cache):
+                past_key_values._update = step.update
+
+            iteration_outputs = iteration_func(
+                hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                output_hidden_states=output_hidden_states,
+                layer_slice=step.layer_slice
+            )
+
+            # Update the hidden states cache
+            if step.update:
+                hidden_states = iteration_outputs.last_hidden_state
+
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states[:end] + iteration_outputs.hidden_states
+
+            if output_attentions:
+                all_self_attns = all_self_attns[:end] + iteration_outputs.attentions
+
+            if use_cache:
+                next_decoder_cache = iteration_outputs.past_key_values
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_decoder_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+    def _modeling_sequential(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[LayerCache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        output_hidden_states: Optional[bool] = False,
+    ) -> BaseModelOutputWithPast:
+        """
+        Sequentially update the hidden states, token by token.
+        """
+        seq_len = hidden_states.shape[1]
+        last_hidden_state = []
+        all_hidden_states = []
+        all_self_attns = []
+
+        for i in range(seq_len):
+            m_hidden_states = hidden_states[:, i:i+1]
+            m_attention_mask = (
+                (attention_mask[:, : i + 1] if attention_mask.ndim == 2 else attention_mask[:, :, i : i + 1])
+                if attention_mask is not None
+                else None
+            )
+            m_position_ids = position_ids[:, i:i+1] if position_ids is not None else None
+            m_cache_position = cache_position[i:i+1] if cache_position is not None else None
+            m_position_embeddings = (
+                position_embeddings[0][:, i:i+1],
+                position_embeddings[1][:, i:i+1]
+            )
+
+            outputs = self._iterate_layers(
+                m_hidden_states,
+                attention_mask=m_attention_mask,
+                position_ids=m_position_ids,
+                past_key_values=past_key_values,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cache_position=m_cache_position,
+                position_embeddings=m_position_embeddings,
+                output_hidden_states=output_hidden_states
+            )
+
+            last_hidden_state.append(outputs.last_hidden_state)
+
+            if output_hidden_states:
+                all_hidden_states.append(outputs.hidden_states)
+
+            if output_attentions:
+                all_self_attns.append(outputs.attentions)
+
+            if use_cache:
+                past_key_values = outputs.past_key_values
+
+        last_hidden_state = torch.cat(last_hidden_state, dim=1)
+
+        if output_hidden_states:
+            all_hidden_states = [
+                torch.cat([hs[i] for hs in all_hidden_states], dim=1) for i in range(len(all_hidden_states[0]))
+            ]
+
+        if output_attentions:
+            # TODO: deal with attention outputs for non-flash-attention implmentations
+            all_self_attns = all_self_attns[-1]
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=last_hidden_state,
+            past_key_values=past_key_values,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+    def _update_causal_mask(
+        self,
+        attention_mask: torch.Tensor,
+        input_tensor: torch.Tensor,
+        cache_position: torch.Tensor,
+        past_key_values: Cache,
+        output_attentions: bool,
+    ):
+        """fix this function to handle layer cache"""
+        if self.config._attn_implementation == "flash_attention_2":
+            if attention_mask is not None and 0.0 in attention_mask:
+                return attention_mask
+            return None
+
+        # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in
+        # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail
+        # to infer the attention mask.
+        past_seen_tokens = past_key_values.get_seq_length() if isinstance(past_key_values, Cache) else 0
+        using_static_cache = isinstance(past_key_values, StaticCache)
+
+        # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
+        if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions:
+            if AttentionMaskConverter._ignore_causal_mask_sdpa(
+                attention_mask,
+                inputs_embeds=input_tensor,
+                past_key_values_length=past_seen_tokens,
+                is_training=self.training,
+            ):
+                return None
+
+        dtype, device = input_tensor.dtype, input_tensor.device
+        min_dtype = torch.finfo(dtype).min
+        sequence_length = input_tensor.shape[1]
+        if using_static_cache:
+            target_length = past_key_values.get_max_length()
+        else:
+            target_length = (
+                attention_mask.shape[-1]
+                if isinstance(attention_mask, torch.Tensor)
+                else past_seen_tokens + sequence_length + 1
+            )
+
+        # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).
+        causal_mask = _prepare_4d_causal_attention_mask_with_cache_position(
+            attention_mask,
+            sequence_length=sequence_length,
+            target_length=target_length,
+            dtype=dtype,
+            device=device,
+            min_dtype=min_dtype,
+            cache_position=cache_position,
+            batch_size=input_tensor.shape[0],
+        )
+
+        if (
+            self.config._attn_implementation == "sdpa"
+            and attention_mask is not None
+            and attention_mask.device.type == "cuda"
+            and not output_attentions
+        ):
+            # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when
+            # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
+            # Details: https://github.com/pytorch/pytorch/issues/110213
+            causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)
+
+        return causal_mask
+
+
+class LCKVLlamaForCausalLM(LCKVLlamaPreTrainedModel, LlamaForCausalLM):
+    def __init__(self, config):
+        LCKVLlamaPreTrainedModel.__init__(self, config)
+        LlamaForCausalLM.__init__(self, copy.deepcopy(config)) # copy config to avoid modifying the original
+        self.model = LCKVLlamaModel(config)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        cache_position=None,
+        position_ids=None,
+        use_cache=True,
+        num_logits_to_keep=None,
+        **kwargs,
+    ):
+        """fix this function to handle sink cache"""
+        # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens
+        # Exception 1: when passing input_embeds, input_ids may be missing entries
+        # Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here
+        if isinstance(past_key_values, Cache):
+            if inputs_embeds is not None:  # Exception 1
+                input_ids = input_ids[:, -cache_position.shape[0] :]
+            elif input_ids.shape[1] != cache_position.shape[0]:  # Default case (the "else", a no op, is Exception 2)
+                input_ids = input_ids[:, cache_position]
+
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if isinstance(past_key_values, Cache):
+
+                if getattr(past_key_values, "build_position_ids_based_on_cache", False):
+                    cur_cache_length = past_key_values.get_seq_length()
+                    position_ids = position_ids[:, cur_cache_length :cur_cache_length + input_ids.shape[1]]
+                else:
+                    position_ids = position_ids[:, -input_ids.shape[1] :]
+
+                # This `clone` call is needed to avoid recapturing cuda graphs with `torch.compile`'s  `mode="reduce-overhead`, as otherwise the input `position_ids` would have various stride during the decoding. Here, simply using `.contiguous()` is not sufficient as in the batch size = 1 case, `position_ids` is already contiguous but with varying stride which retriggers a capture.
+                position_ids = position_ids.clone(memory_format=torch.contiguous_format)
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and cache_position[0] == 0:
+            model_inputs = {"inputs_embeds": inputs_embeds, "input_ids": None}
+        else:
+            # The clone here is for the same reason as for `position_ids`.
+            model_inputs = {"input_ids": input_ids.clone(memory_format=torch.contiguous_format), "inputs_embeds": None}
+
+        if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2:
+            if model_inputs["inputs_embeds"] is not None:
+                batch_size, sequence_length, _ = model_inputs["inputs_embeds"].shape
+                device = model_inputs["inputs_embeds"].device
+            else:
+                batch_size, sequence_length = model_inputs["input_ids"].shape
+                device = model_inputs["input_ids"].device
+
+            dtype = self.lm_head.weight.dtype
+            min_dtype = torch.finfo(dtype).min
+
+            attention_mask = _prepare_4d_causal_attention_mask_with_cache_position(
+                attention_mask,
+                sequence_length=sequence_length,
+                target_length=past_key_values.get_max_length(),
+                dtype=dtype,
+                device=device,
+                min_dtype=min_dtype,
+                cache_position=cache_position,
+                batch_size=batch_size,
+            )
+
+        if num_logits_to_keep is not None:
+            model_inputs["num_logits_to_keep"] = num_logits_to_keep
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "cache_position": cache_position,
+                "past_key_values": past_key_values,
+                "use_cache": use_cache,
+                "attention_mask": attention_mask,
+            }
+        )
+        return model_inputs

utils.py

@@ -0,0 +1,190 @@
+import re
+from dataclasses import dataclass
+from typing import List, Optional
+
+import torch
+
+from transformers.modeling_flash_attention_utils import _flash_attention_forward
+
+
+@dataclass
+class IterStep:
+    """A helper class for the iteration plan"""
+    layer_slice: slice = slice(None)
+    requires_grad: bool = True
+    update: bool = True
+
+@dataclass
+class LayerType:
+    """A helper class to collect the layer type information"""
+    layer_idx: int
+    use_sliding_window: bool
+    attends_to: int
+    attends_top: bool
+    computes_kv: bool
+
+class LayerTypeParser:
+    """
+    A helper class to parse the layer type string and provide some useful methods.
+
+    Arguments:
+        layer_type (str): A string of integers separated by underscores. The i-th integer
+            means the layer will use the key-value pair in the i-th layer as the kv cache.
+            Special characters may be placed after the integers:
+            - `s` means the layer will use sliding window attention.
+
+    >>> layer_type = LayerTypeParser("0_0_0_5s_5s_5s_8_8_8")[3]
+    >>> layer_type.attends_to
+    5
+    >>> layer_type.attends_top
+    True
+    >>> layer_type.use_sliding_window
+    True
+    """
+    def __init__(self, layer_type: str):
+        self._layer_type = layer_type
+
+        # parse the layer type
+        self.layer_indices = []
+        self.sliding_window = []
+        for s in layer_type.split("_"):
+            layer_idx, sliding_window = re.match(r"^(\d+)(s)?$", s).groups()
+            self.layer_indices.append(int(layer_idx))
+            self.sliding_window.append(bool(sliding_window))
+
+    def __len__(self):
+        return len(self.layer_indices)
+
+    def __getitem__(self, layer_idx: int) -> LayerType:
+        """return the layer type information for the given layer index"""
+        return LayerType(
+            layer_idx=layer_idx,
+            use_sliding_window=self.sliding_window[layer_idx],
+            attends_to=self.layer_indices[layer_idx],
+            attends_top=self.layer_indices[layer_idx] > layer_idx,
+            computes_kv=layer_idx in self.layer_indices,
+        )
+
+    def use_sliding_window(self) -> bool:
+        """whether there exists a layer that uses sliding window attention"""
+        return any(self.sliding_window)
+
+    def attends_top(self) -> bool:
+        """whether there exists a layer that attends to layers above it"""
+        return any(self.layer_indices[i] > i for i in range(len(self)))
+
+    def iteration_plan(self, forward_passes: int = 7, backward_passes: int = 2) -> List[IterStep]:
+        """
+        Return a iteration plan for the layer types. The plan is a list of IterStep objects.
+        """
+        # if there is no cyclic dependency, return the default plan
+        if not self.attends_top():
+            return [IterStep()]
+
+        # otherwise, return the plan for the cyclic dependency
+        plan = []
+        i = 0
+        while i < len(self):
+
+            # if the layer attends to top layers, resolve the cyclic dependency
+            if self[i].attends_top:
+
+                # find the top layer in the cyclic dependency
+                top = self[i].attends_to
+                while top < max(self.layer_indices[i: top + 1]):
+                    top = max(self.layer_indices[i: top + 1])
+                top += 1
+
+                # create iteration plan for this group
+                layer_slice = slice(i, top)
+                plan.extend([
+                    *forward_passes * [IterStep(layer_slice, requires_grad=False, update=False)],
+                    *(backward_passes - 1) * [IterStep(layer_slice, update=False)],
+                    IterStep(layer_slice)
+                ])
+
+            # otherwise, create a default plan
+            else:
+
+                top = i + 1
+                while top < len(self) and not self[top].attends_top:
+                    top += 1
+                plan.append(IterStep(slice(i, top)))
+
+            # update the index
+            i = top
+
+        return plan
+
+    def check(self, num_hidden_layers: int):
+        """Check if the layer type is valid"""
+        if len(self.layer_indices) != num_hidden_layers:
+            raise ValueError("The number of layer types should be equal to the number of hidden layers.")
+        for i in range(num_hidden_layers):
+            if self.layer_indices[i] not in range(num_hidden_layers):
+                raise ValueError("The layer type should be in the range of the number of hidden layers.")
+
+
+def flash_attention_forward(
+    query_states: torch.Tensor,
+    key_states: torch.Tensor,
+    value_states: torch.Tensor,
+    attention_mask: torch.Tensor,
+    query_length: int,
+    is_causal: bool,
+    dropout: float = 0.0,
+    position_ids: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    sliding_window: Optional[int] = None,
+    use_top_left_mask: bool = False,
+    softcap: Optional[float] = None,
+    deterministic: bool = None,
+    no_diag: bool = False,
+):
+    """
+    This function is a wrapper around the _flash_attention_forward function in the
+    transformers library. It adds support to mask the diagonal elements of the attention
+    matrix. The diagonal mask is used to resolve the cyclic dependencies in the LCKV model.
+    """
+    prune_query = False
+    if no_diag:
+        if key_states.size(1) == 1:
+            b, l, _, d = value_states.size()
+            _, _, h, _ = query_states.size()
+            return value_states.new_zeros((b, l, h, d))
+
+        if key_states.size(1) == query_states.size(1):
+            prune_query = True
+            query_states = query_states[:, 1:, :, :]
+            query_length -= 1
+
+            if attention_mask is not None:
+                attention_mask = attention_mask[:, 1:]
+
+        key_states = key_states[:, :-1, :, :]
+        value_states = value_states[:, :-1, :, :]
+
+        if sliding_window is not None:
+            sliding_window = sliding_window - 1
+
+    result: torch.Tensor = _flash_attention_forward(
+        query_states=query_states,
+        key_states=key_states,
+        value_states=value_states,
+        attention_mask=attention_mask,
+        query_length=query_length,
+        is_causal=is_causal,
+        dropout=dropout,
+        position_ids=position_ids,
+        softmax_scale=softmax_scale,
+        sliding_window=sliding_window,
+        use_top_left_mask=use_top_left_mask,
+        softcap=softcap,
+        deterministic=deterministic,
+    )
+
+    if prune_query:
+        b, _, h, d = result.size()
+        result = torch.cat([result.new_zeros((b, 1, h, d)), result], dim=1)
+
+    return result