Generate: New `Cache` abstraction and Attention Sinks support
Author: tomaarsen
Pull request: yes

Closes #26553

Hello!

# What does this PR do?
I had a few hours on Saturday to work up a draft version of the updated KV caching mechanism as discussed in #26553. Ideally, this should allow Attention Sinks (https://github.com/tomaarsen/attention_sinks) / StreamingLLM (https://arxiv.org/abs/2309.17453) to be easily implemented in a third-party or in transformers directly. 

The implementation doesn't work well yet, as the VRAM usage quickly shoots up after generating even just 8 tokens. This is probably some bug that I haven't had time for yet. There's a few other comments that I have on specific sections of code, so I'll write some comments below.

## Goal for this draft
The intention for this draft is to continue discussion about whether this is moving in the right direction, and to determine the scope (e.g. do we want to include this updated `Cache` for all architectures that use KV caching?). 

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?

## Who can review?

@patrickvonplaten 
@gante
@LysandreJik 
@Guangxuan-Xiao

- Tom Aarsen


Comments:
@patrickvonplaten @gante 
Based on your feedback, I intend to make the following changes (when I have a bit more time). 

1. move `layer_idx` to `update(key_states, value_states, layer_idx)` rather than storing it as a class attribute on `Cache`. This also involves adding `layer_idx` as class attributes to e.g. `LlamaAttention` and `LlamaDecoderLayer`. This also removes the `set_layer_index` magic from `Cache`.
2. convert `past_key_values` to Cache instance (if not already) at the start of LlamaAttention.forward. This avoids all `isinstance(...)` calls in `LlamaAttention`, and removes the need for the black magicky `__bool__`.
3. convert back to tuple of tuples when returning if some `use_legacy_cache` flag is True. (should this flag be propagated all the way up to `LlamaModel`?)
4. use separate `key_cache` and `value_cache` dicts in the `Cache` instance for efficiency (removes a `torch.cat`) and simplicity.

Regarding https://github.com/huggingface/transformers/pull/26681#discussion_r1350368080 I would have to experiment. There might be options, but they'll probably be slower than necessary.

`past_key_values` and caching will be an incredible feature.

Two PRs should be made irrelevant if this merges: https://github.com/huggingface/transformers/pull/25086  https://github.com/huggingface/transformers/pull/17574

> @patrickvonplaten @gante Based on your feedback, I intend to make the following changes (when I have a bit more time).
> 
> 1. move `layer_idx` to `update(key_states, value_states, layer_idx)` rather than storing it as a class attribute on `Cache`. This also involves adding `layer_idx` as class attributes to e.g. `LlamaAttention` and `LlamaDecoderLayer`. This also removes the `set_layer_index` magic from `Cache`.
> 2. convert `past_key_values` to Cache instance (if not already) at the start of LlamaAttention.forward. This avoids all `isinstance(...)` calls in `LlamaAttention`, and removes the need for the black magicky `__bool__`.
> 3. convert back to tuple of tuples when returning if some `use_legacy_cache` flag is True. (should this flag be propagated all the way up to `LlamaModel`?)
> 4. use separate `key_cache` and `value_cache` dicts in the `Cache` instance for efficiency (removes a `torch.cat`) and simplicity.
> 
> Regarding [#26681 (comment)](https://github.com/huggingface/transformers/pull/26681#discussion_r1350368080) I would have to experiment. There might be options, but they'll probably be slower than necessary.

Sounds great!

Addressed the various comments. Beyond that, I also made `past_key_values.update(key, value, idx)` returns `key, value` as you'd expect. Manual generation (i.e. repeated calling of a `LlamaForCausalLM` instance with the `past_key_values`) works well, I even see a ~2% speedup, but don't quote me on that speedup. `model.generate` doesn't work yet because `use_legacy_flag` defaults to None, i.e. Falsey, and `model.generate` doesn't work with the cache yet. @patrickvonplaten @gante Should we go for:
1. Immediately update `model.generate` to work with Cache instances or,
2. Insert `use_legacy_cache=True` as default for `model.generate` until this can be removed in some later PR?

This is all under the assumption that the PR is heading in the right direction 😄 

As a heads up, the sink cache does not work yet. I still need to do some experiments to see if I can store rotated keys and back-rotate + forward-rotate them when the cached keys are requested, rather than storing non-rotated keys. That is what I'll work on next. 

That leaves me with an additional question: each architecture requires slightly different key rotations. I'd like to implement this to be sufficiently adaptable, e.g. allowing architecture-specific functionality in `src/transformers/models/<architecture>/cache_utils.py`. However, what the relation to classes or functions in this file should be to the `src/transformers/cache_utils.py` cache classes is still unclear to me. In short: how do we allow an architecture to e.g. have slightly different implementations for the `SinkCache` or the `DynamicCache`?

- Tom Aarsen

Hey @tomaarsen 👋 

### `generate` compatibility
Re `generate` compatibility: usually, I'm pro small PRs, tackling one problem at a time. However, since `generate` is the main gate to LLMs in `transformers`, and caching is only needed for auto-regressive generation, day 0 support is important -- at least `gready_search` and `sample` should be operational by the time this PR can be merged. Otherwise, we might get ourselves in a position where we realize we need to rewrite a significant part of `Cache`/`generate` to enable them together, which is very undesirable. However, I'm pro having this PR enabling the new cache on a single model (Llama) and taking care of the other models later as needed 🤗 I can also give a hand on the `generate` side, when we are happy with the state of the PR with the exception of `generate`.

### model-specific cache changes
I do expect some models to require different code to convert back and forth from the legacy cache format (not all models have the same cache format, there are ~5% of them have slight differences e.g. [see this in BLOOM](https://github.com/huggingface/transformers/blob/288bf5c1d2844e89a2a87aeee90033532335e2e6/src/transformers/models/bloom/modeling_bloom.py#L504)). You are also writing that RoPE-based models may also need custom logic. 

The model-specific cache modification part is an important design decision that will have an impact over many versions of `transformers` to come, so I'd also like to hear your thoughts @tomaarsen @patrickvonplaten 

EDIT: what follows bellow is my "plan B" suggestion, read my next comment for the "plan A" :D
To me, that suggests five things:
1) a model may need to subclass the base `Cache` to implement its own nuances regarding the cache format
2) despite the above, the cache operations and output format are standardized for each type of cache, so we would benefit from a strong base class (as we do in the config files or in the base pretrained class).
3) I suspect the model-specific changes for the cache would be small, so we could keep them in the modeling file. We can always move to a new file if needed :)
4) Because of 1), each model would have to implement a thin wrapper class to the instantiable cache classes, even if the base class is fully compatible. This also goes in line with our philosophy in `tramsformers` where each model implements its own model-specific operations.
5) Because of 4), in practice, each model defines its own cache. This means users are free to write their own custom caches -- power to the users 💪 



Upon some shower thoughts, I've come across an alternative plan for the model-specific cache modification problem -- easier to implement and that would result in more readable code. 

Instead of the base `Cache` holding the code to convert to and from the legacy format (which then requires subclassing at each model, if it has a different legacy cache format), the conversion logic could be held in model-specific functions to convert to and from the legacy format. In other words, each model would implement a `to_legacy_cache` and a `from_legacy_cache`, and the different types of `Cache` would be immediately available to the model with no further modifications!

> lding the code to convert to and from the legacy format (which then requires subclassing at each model, if it has a different legacy cache format)

Agree here! Think it would be better if we only have a few selected cache classes that work for all models. The functions `from_legacy_cache` and `to_legacy_cache` indeed need to be model specific, so I think they can be just stand-alone functions in each `modeling_....py` file. 

=> Thus, I think we can:
- a) Have all `Cache` classes in a general `generation/cache.py` file. These general cache implementations should be identical for all models
- b) For backwards compatibility each model needs specific `from_legacy_cache` and `to_legacy_cache` functions, but they don't need to be part of the Cache class IMO, they can just be stand-alone functions in each model class so that they are easier to deprecate going forward.

I just noticed that Joao's message has been edited, and I missed Patrick's message, so my response doesn't make much sense anymore - I deleted it.

I also prefer the `from_legacy_cache` and `to_legacy_cache` implementations. I'll get on it.

@gante  @tomaarsen  This is really a good abstraction of kv_cache to enable window context and be compatible to the legacy kv_cache. But for the memory footprint, the 'torch.cat' are still needed when update the cache and reorder_cache using 'index_select' is also there with beam search.  
Do you have a plan to use the pre-allocate buffer to avoid 'torch.cat'?  The pre-allocated buffer should be compatible with semantic of your Cache.  For the cache.update operation,  token slots from the pre-allocated buffer are needed to store key/value token states. 

I removed the commits regarding changing the order of the caches based on @gante's recommendation. He is working on more changes for this PR here: https://github.com/tomaarsen/transformers/pull/1

@liangan1 both are issues in the near-future roadmap, with this cache abstraction being a requirement! 🙌 

First we will work on pre-allocated buffers (a subclass of `Cache`), which result in significant speedups (especially with `torch.compile`). Then we will simplify beam search into an XLA-friendly method, just like in our TF and JAX implementations.

> @liangan1 both are issues in the near-future roadmap, with this cache abstraction being a requirement! 🙌
> 
> First we will work on pre-allocated buffers (a subclass of `Cache`), which result in significant speedups (especially with `torch.compile`). Then we will simplify beam search into an XLA-friendly method, just like in our TF and JAX implementations.

Wow, cool. Can you share more details about your roadmap? e.g., the pre-allocate buffers. In fact, we also have implemented a indirect access kv_cache in [Intel-extension-for-pytorch optimization for LLM](https://intel.github.io/intel-extension-for-pytorch/cpu/latest/tutorials/llm.html#indirect-access-kv-cache)  which can be used for both greedy and beam and we are pleasure to contribute code if I can know your design about the pre-allocated buffer.

> Should we roll out the new cache feature for the most important models and aim to get this PR in soon ?

Yes! @gante has picked up the work in https://github.com/tomaarsen/transformers/pull/1 to get `generate()` working, and has proposed to implement the cache for Llama, Mistral and Persimmon + some tests. Once he has completed that work, I can merge it into this PR, and we can get it merged into `main`.

Does that sound like a plan?

This PR is now generate-compatible, and has tests to ensure the two formats are interchangeable 💪 

The suggested review order is `cache_utils.py` -> `modeling_llama.py` -> `test_utils.py` -> others

@patrickvonplaten in the added test, we confirm that the two formats are interchangeable: they can be converted back and forth and be manipulated the same way. However, under the hood, the updated models use the new cache internally. It would be best to add a long slow test to key models (Llama and Mistral here) before merging this PR, correct?

@tomaarsen can you confirm that the attention sinks are working as expected?

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_26681). All of your documentation changes will be reflected on that endpoint.

@patrickvonplaten applied the suggested changes 👍  I agree they the full abstraction will be better in the long run.

Two minor modifications over your suggestions:
1. Replaced the `__init__` `warn` by `warning_once`, otherwise the user would get `n_layers` repeated warnings
2. In the beam methods, to check against `Cache`, we needed to perform the check before calling `_reorder_cache`

Will review tomorrow ! 

@ArthurZucker depending on your answer to [this comment](https://github.com/huggingface/transformers/pull/26681#discussion_r1411079517), the PR will move very differently -- I'll wait for your answer before continuing :)

Ready for a new round of reviews. IMO, it's now much better than before: no new flags, `generate` compatibility for any `Cache` subclass, among others.

@patrickvonplaten @ArthurZucker:
- There is no more `use_legacy_cache` flag. If `past_key_values` is a `Cache` instance (which can be an empty cache), the model sticks to that format. No flags to deprecate in the future :)
- Because of the point above, if we pass  `past_key_values=some_cache_instance` to `generate`, it will use that cache instance (and, therefore, format). `generate` is ready to receive static caches! See the generation test for an example.
- The caches now accept a `cache_kwargs` dictionary in `update` (its is not a `**dict`). It is the responsibility of each subclass to process this dictionary. It should gives us plenty of room (and flexibility) to grow without exploding the signature.
- `Cache` has a `reorder_cache` method, as suggested. This also means we can start thinking about standardizing cache formats and removing per-model functions;
- Added tons of docstrings

@tomaarsen:
- Added support for partial rotation in `SinkCache`. An extra set of eyes there would be appreciated! :)

~⚠️ There is a bug with beam search + legacy cache format. I suspect it is a subtle dimension-related bug, but the PR can be reviewed while I chase it.~ EDIT: bug sorted

I've ran perplexity tests for Attention Sinks with Persimmon, Llama and Mistral, and they all look good in terms of performance, VRAM and latency. I plan to do some more tests with Phi tomorrow, but I assume it's all good other than the above 3 comments.
I'll also try to use SinkCache with `generate`.

Let's get this one in for the release no? cc @ArthurZucker @amyeroberts @LysandreJik 

Yep 


I am going to trigger the CI now (with a reabse on main in order to include the fix #27887

https://github.com/huggingface/transformers/actions/runs/7129508785

@tomaarsen Great work! Looks like the current implementation cannot handle case **input_length > window_length** at least for some models (mistral) 
Do I miss anything? Thanks! 

```
cache = SinkCache(window_length=1024, num_sink_tokens=4)
gen_out = model.generate(**inputs, do_sample=False, max_new_tokens=10, 
                         use_cache=True,
                         past_key_values=cache, 
                         pad_token_id=tokenizer.eos_token_id)
decoded = tokenizer.batch_decode(gen_out, skip_special_tokens=True)
```
Supporse inputs is of torch.Size([1, 15759]). It will raise
`RuntimeError: shape '[-1, 15759]' is invalid for input of size 1024`
when it trys to reshape position_ids

```
File ~/anaconda3/envs/llm/lib/python3.8/site-packages/transformers/models/mistral/modeling_mistral.py:984, in MistralModel.forward(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict)
    982     position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
    983 else:
--> 984     position_ids = position_ids.view(-1, seq_length).long()
```

Not sure it was tested fro mistral! Can you open a separate issue with a full reproducer? 🤗 

@fayejf You might want to look at this paper, too - handling extremely long inputs in a streaming fashion: [Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention](https://arxiv.org/abs/2404.07143).

> 1B model that was fine-tuned on up to 5K sequence length passkey instances solved the 1M length problem.

@ArthurZucker It seems that even with Llama2, passing in SinkCache to `generate` causes errors.

I'm using transformers 4.39.3 and the Llama2 model was loaded using the following code:

```
from transformers import AutoConfig, LlamaForCausalLM, AutoTokenizer

model_id = "meta-llama/Llama-2-7b-hf"
model = LlamaForCausalLM.from_pretrained(
      model_id, low_cpu_mem_usage=True, device_map='auto',
      torch_dtype=torch.bfloat16)
model = model.eval()
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = 'left'
```

The SinkCache was passed to generate as in @fayejf's script. I am not sure if this is the correct way to use SinkCache:

```
from transformers import SinkCache

prefix = 'Hello world!'
inputs = tokenizer(prefix, return_tensors='pt').to(device)

cache = SinkCache(window_length=1024, num_sink_tokens=4)
gen_out = model.generate(**inputs, do_sample=False, max_new_tokens=256,
                         use_cache=True,
                         past_key_values=cache,
                         pad_token_id=tokenizer.pad_token_id)
decoded = tokenizer.batch_decode(gen_out, skip_special_tokens=True)
```

The code caused `TypeError: object of type 'SinkCache' has no len()` as a result of this `DynamicCache.from_legacy_cache` call (see the stack trace below). Looks like you are familiar with the `StaticCache` stuff, any suggestions on how to get around this? Thanks in advance!

```
File /miniconda3/envs/pytorch2/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:977, in LlamaModel.forward(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict, cache_position)
    975 if use_cache:  # kept for BC (cache positions)
    976     if not isinstance(past_key_values, StaticCache):
--> 977         past_key_values = DynamicCache.from_legacy_cache(past_key_values)
    978         past_seen_tokens = past_key_values.get_seq_length()
    980 if cache_position is None:

File /miniconda3/envs/pytorch2/lib/python3.10/site-packages/transformers/cache_utils.py:181, in DynamicCache.from_legacy_cache(cls, past_key_values)
    179 cache = cls()
    180 if past_key_values is not None:
--> 181     for layer_idx in range(len(past_key_values)):
    182         key_states, value_states = past_key_values[layer_idx]
    183         cache.update(key_states, value_states, layer_idx)
```


 

> @fayejf You might want to look at this paper, too - handling extremely long inputs in a streaming fashion: [Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention](https://arxiv.org/abs/2404.07143).
> 
> > 1B model that was fine-tuned on up to 5K sequence length passkey instances solved the 1M length problem.

@MoonRide303 Thanks for sharing! very interesting paper!

@ArthurZucker I was with 4.38.2 and haven't got a chance to test it with latest release. But looks like @explanare have seen similar issues. I also tried with llama model and I've seen legacy cache error as well. 

I'm not sure if I'm using SinkCache correctly. I didn't find any doc for it unfortuantely. :(
I follow the way in test file. 

> @ArthurZucker It seems that even with Llama2, passing in SinkCache to `generate` causes errors.
> 
> I'm using transformers 4.39.3 and the Llama2 model was loaded using the following code:
> 
> ```
> from transformers import AutoConfig, LlamaForCausalLM, AutoTokenizer
> 
> model_id = "meta-llama/Llama-2-7b-hf"
> model = LlamaForCausalLM.from_pretrained(
>       model_id, low_cpu_mem_usage=True, device_map='auto',
>       torch_dtype=torch.bfloat16)
> model = model.eval()
> tokenizer = AutoTokenizer.from_pretrained(model_id)
> tokenizer.pad_token = tokenizer.eos_token
> tokenizer.padding_side = 'left'
> ```
> 
> The SinkCache was passed to generate as in @fayejf's script. I am not sure if this is the correct way to use SinkCache:
> 
> ```
> from transformers import SinkCache
> 
> prefix = 'Hello world!'
> inputs = tokenizer(prefix, return_tensors='pt').to(device)
> 
> cache = SinkCache(window_length=1024, num_sink_tokens=4)
> gen_out = model.generate(**inputs, do_sample=False, max_new_tokens=256,
>                          use_cache=True,
>                          past_key_values=cache,
>                          pad_token_id=tokenizer.pad_token_id)
> decoded = tokenizer.batch_decode(gen_out, skip_special_tokens=True)
> ```
> 
> The code caused `TypeError: object of type 'SinkCache' has no len()` as a result of this `DynamicCache.from_legacy_cache` call (see the stack trace below). Looks like you are familiar with the `StaticCache` stuff, any suggestions on how to get around this? Thanks in advance!
> 
> ```
> File /miniconda3/envs/pytorch2/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:977, in LlamaModel.forward(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict, cache_position)
>     975 if use_cache:  # kept for BC (cache positions)
>     976     if not isinstance(past_key_values, StaticCache):
> --> 977         past_key_values = DynamicCache.from_legacy_cache(past_key_values)
>     978         past_seen_tokens = past_key_values.get_seq_length()
>     980 if cache_position is None:
> 
> File /miniconda3/envs/pytorch2/lib/python3.10/site-packages/transformers/cache_utils.py:181, in DynamicCache.from_legacy_cache(cls, past_key_values)
>     179 cache = cls()
>     180 if past_key_values is not None:
> --> 181     for layer_idx in range(len(past_key_values)):
>     182         key_states, value_states = past_key_values[layer_idx]
>     183         cache.update(key_states, value_states, layer_idx)
> ```

I got the same error with Llama2 using 4.38.2 following code [here](https://colab.research.google.com/drive/1S0oIPaqxAVp0oWEwTadhZXDjhWiTyF12?usp=sharing#scrollTo=fwxG9OFAV4Qv). Also tried `transformers==4.39.0`, got the same error.
```
File "/opt/conda/envs/dn/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 982, in forward
    past_key_values = DynamicCache.from_legacy_cache(past_key_values)
  File "/opt/conda/envs/dn/lib/python3.10/site-packages/transformers/cache_utils.py", line 166, in from_legacy_cache
    for layer_idx in range(len(past_key_values)):
TypeError: object of type 'SinkCache' has no len()
```

cc @gante as well as we broke it 

Will might be fixed by #30476. Only tuples should be going that path 

@fayejf @ys-zong @ArthurZucker 

`SinkCache` is operational on `main` with Mistral and a few other models (especially non-RoPE models). Llama and Llama-dependent models + `SinkCache` is fixed in [this PR](https://github.com/huggingface/transformers/pull/30581), which hopefully will get merged soon 🤗 

Apologies for this temporary issue, we had to break a few eggs (sink cache) to make an omelet (torch.compile support) :D

Edit: merged, let me know if you run into issues!

Hi @tomaarsen I just came across this line https://github.com/huggingface/transformers/blob/main/src/transformers/cache_utils.py#L979 in your PR and wonder why the condition for cache growing is  `key_states.shape[-2] + self.get_seq_length(layer_idx) < self.window_length`?  To my understanding even if the conditon is `<=` the new key states can still safely grow into the window size without need to evict.  Is there any logic that the codition must be `<`?  Thanks!