|
|
|
|
|
|
|
|
from __future__ import annotations |
|
|
|
|
|
from typing import Any, Dict, List, Optional, Tuple |
|
|
|
|
|
import torch |
|
|
import transformers |
|
|
|
|
|
|
|
|
class Cache(transformers.cache_utils.Cache): |
|
|
""" |
|
|
A cache used for storing hidden states produced by flash linear attention models. |
|
|
|
|
|
It stores the states of each layer as the tensor of shape `[batch_size, key_dim, value_dim]`. |
|
|
""" |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
seen_tokens: int = 0 |
|
|
) -> Cache: |
|
|
|
|
|
self.states: List[torch.Tensor] = [] |
|
|
self._seen_tokens = seen_tokens |
|
|
|
|
|
def __getitem__(self, layer_idx: int) -> torch.Tensor: |
|
|
if layer_idx < len(self): |
|
|
return self.states[layer_idx] |
|
|
else: |
|
|
raise KeyError(f"Cache only has {len(self)} layers, attempted to access layer with index {layer_idx}") |
|
|
|
|
|
def __iter__(self): |
|
|
for state in self.states: |
|
|
yield state |
|
|
|
|
|
def __len__(self): |
|
|
return len(self.states) |
|
|
|
|
|
def update( |
|
|
self, |
|
|
state: Tuple[torch.Tensor], |
|
|
layer_idx: int, |
|
|
offset: Optional[int] = 1, |
|
|
cache_kwargs: Optional[Dict[str, Any]] = None, |
|
|
) -> Tuple[torch.Tensor]: |
|
|
""" |
|
|
Updates the cache with the new `state` for the layer `layer_idx`. |
|
|
|
|
|
Parameters: |
|
|
state (`Tuple[torch.Tensor]`): |
|
|
The new state to cache. |
|
|
layer_idx (`int`): |
|
|
The index of the layer to cache the states for. |
|
|
offset (`int`): |
|
|
The offset of current fed tokens. |
|
|
cache_kwargs (`Dict[str, Any]`, `optional`): |
|
|
Additional arguments for the cache subclass. |
|
|
|
|
|
Return: |
|
|
The updated state. |
|
|
""" |
|
|
|
|
|
if isinstance(state, torch.Tensor): |
|
|
state = (state,) |
|
|
if len(self.states) <= layer_idx: |
|
|
self.states.append(state) |
|
|
else: |
|
|
for i, s in enumerate(state): |
|
|
self.states[layer_idx][i].copy_(s) |
|
|
|
|
|
if layer_idx == len(self) - 1: |
|
|
self._seen_tokens += offset |
|
|
|
|
|
return state |
|
|
|
|
|
def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: |
|
|
"""Returns the sequence length of the cached states. A layer index can be optionally passed.""" |
|
|
if len(self.states) <= layer_idx: |
|
|
return 0 |
|
|
return self._seen_tokens |
|
|
|
|
|
def get_max_length(self) -> Optional[int]: |
|
|
"""Returns the maximum sequence length of the cached states. Cache does not have a maximum length.""" |
|
|
return None |
|
|
|
|
|
def reorder_cache(self, beam_idx: torch.LongTensor): |
|
|
"""Reorders the cache for beam search, given the selected beam indices.""" |
|
|
for layer_idx in range(len(self.states)): |
|
|
device = self.states[layer_idx].device |
|
|
self.states[layer_idx] = self.states[layer_idx].index_select(0, beam_idx.to(device)) |
|
|
|
|
|
def to_legacy_cache(self) -> Tuple[torch.Tensor]: |
|
|
return tuple(self.states) |
|
|
|
|
|
@classmethod |
|
|
def from_legacy_cache( |
|
|
cls, |
|
|
past_key_values: Optional[Tuple[torch.Tensor]] = None, |
|
|
seen_tokens: int = 0 |
|
|
) -> Cache: |
|
|
"""Converts a cache in the legacy cache format into an equivalent `Cache`.""" |
|
|
|
|
|
cache = cls(seen_tokens) |
|
|
if past_key_values is not None: |
|
|
for layer_idx in range(len(past_key_values)): |
|
|
cache.update(past_key_values[layer_idx], layer_idx) |
|
|
return cache |
|
|
|
