Experiment with static past key/value buffer
Author: fxmarty
Pull request: yes

This PR is just to see if this could reside in transformers or not.

## Motivation

We suspect [the concatenations of the key/value buffer](https://github.com/huggingface/transformers/blob/ba695c1efd55091e394eb59c90fb33ac3f9f0d41/src/transformers/models/gpt2/modeling_gpt2.py#L321-L322) at each generation step to be expensive. The idea is to modify the buffer in place after a unique allocation.

FasterTransformer [does preallocate the kv cache](https://github.com/NVIDIA/FasterTransformer/blob/c6e8f60ec40da218804a60e6aa986903e7fa8594/src/fastertransformer/models/decoding/Decoding.cc#L83C9-L84) (among others). Preallocating may also help `torch.compile` according to @Chillee, although I don't quite get why yet (there are still dynamic shapes in the model itself, so why care about model I/O static shapes?).

For reference: https://huggingface.slack.com/archives/C055NT312LW/p1683038064467109 & https://huggingface.slack.com/archives/C055NT312LW/p1685570357532069

## Current (ugly) workflow

Very temporary

```python
cache_size = max_new_tokens

past_key_values = tuple([
    tuple([torch.empty(
        batch_size,
        model.config.n_head,
        cache_size,
        model.config.n_embd // model.config.n_head,  # head dimension
        dtype=dtype,
        device=device
    ) for _ in range(2)])

model.enable_static_kv_cache()
res = model.generate(
    **inputs,
    num_beams=1,
    min_new_tokens=max_new_tokens,
    max_new_tokens=max_new_tokens,
    past_key_values=past_key_values
)
```

## Results

Some optimizations may still be missing - right now in small model / small batch size setting this is not interesting.

Results are with PyTorch 2.0.1 eager.

Script: https://github.com/fxmarty/transformers-preallocate-kv-cache/blob/main/run.py
Raw results: https://docs.google.com/spreadsheets/d/15P1o9vDcXOSeLAwLUWiqXQatHxtQYBa_xcyLZanOjQQ/edit?usp=sharing

![image](https://github.com/fxmarty/transformers-preallocate-kv-cache/assets/9808326/19db7654-819d-4939-8306-e0b3e5691af5)
![image](https://github.com/fxmarty/transformers-preallocate-kv-cache/assets/9808326/0dd50f85-a3a0-4571-aa80-2a748b9c9942)
![image](https://github.com/fxmarty/transformers-preallocate-kv-cache/assets/9808326/3a1e6efb-2358-40f4-bf8e-d09261296fd5)
![image](https://github.com/fxmarty/transformers-preallocate-kv-cache/assets/9808326/adc4ec1c-529c-44b6-8e77-105b788e222a)
![image](https://github.com/fxmarty/transformers-preallocate-kv-cache/assets/9808326/e41bc4ca-1d3b-4625-9ec6-283b1442420b)
![image](https://github.com/fxmarty/transformers-preallocate-kv-cache/assets/9808326/48294201-9798-4083-9776-3ef4dbd6ca12)

## Misc

The current implementation of `valid_past_index` being `Optional[int]` is quite debatable and may hamper readability.

I'm also not sure whether adding methods in `PreTrainedModel` that are specific to decoders is OK?

Missing test right now is to generate a sequence length shorter than the KV buffer.

Some todos:
- [ ] `past_key_values` buffer should be initialized in the background, instead of requiring the user to initialize it himself and requiring to pass `generate(**inputs, past_key_values=past_key_values)` (as currently).
- [ ] Current implementation is likely to break with `accelerate` with naive pipeline parallelism, as the buffer is currently initialized on a single device
- [ ] Preallocated kv cache still does not help with small models / batch size.
- [ ] Support an iterative buffer, e.g. that auto-extends each 512 tokens, instead of initializing a buffer of size `max_new_tokens`. This may help reducing memory usage (and speed? probably not).
- [ ] Implement tests
- [ ] Should this be in optimum or transformers?
- [ ] Support all (is it possible?) decoding strategies, instead of currently only `greedy_search`
- [ ] Have it work with cross-attention
- [ ] Preallocate `attention_mask` as well
- [ ] Test on CPU as well

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

In order to efficiently work with the PyTorch team here and to figure out exactly what needs to be done for a super fast generate method, I'd suggest to open a new benchmark repo that includes this PR for say three important LLM models:
- GPT2
- Llama
- starcoder

The repo "maybe just called `transformers-generate-benchmark` should have:
- copied the modeling files of GPT2, Llama, Starcoder
- Use the SDPA attention for all models (just like `betterTransformers` does - you could copy it from here: https://github.com/huggingface/optimum/blob/main/optimum/bettertransformer/models/attention.py)
- A super small generate function that we write ourselves (just greedy without any logit processor)

Then in this repo we run all the benchmarking and also make it easy for the PyTorch team to reproduce the benchmarking numbers.

Sounds good - I'll continue on this PR for easy diff.