[`Core generation`] Adds support for static KV cache
Author: ArthurZucker
Pull request: yes

# ~4x speedups with cuda graphs! 🥳 

Currently getting  ~4x speedups compare to dynamic cache with torch.compile for a single forward pass (agnostic to batch but faster for smaller batch) 

Forward is very very very fast, but materializing the input costs a bit!
~10ms / forward is what we get to! 

- [x] Refactors the way we deal with attention mask:
    - causal and padding are separated 
    - does not rely on the `past_key_values`
    - merged in 2 line. No attention mask utils are needed, no extra complicated logic all explicit
    - LlamaAttention is not self contained, this added 20% overhead in a simple forward
    - Gets rid of the entire `mask_attn_utils` 😄 
- [x] Save the cache class in the generation config
- [x] Init the cache with the batch size (from the generate call) and the `max_length` from the generation config (taking max_new_tokens) into account
- [x] torch.compile

## Benchmark using [af097af](https://github.com/huggingface/transformers/pull/27931/commits/af097af7582de01fcf73856e7fa37be972e96456)


## Use it in generate:
Use this: EDIT: TO COME


## Failing test left for @gante
Related to the fact that I don't return `past_key_values` / is None so the `test_new_cache_format` fails. I don't want to dive in this. 


fixes #28075 , fixes #28610, fixes #28190

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_27931). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

If I understand correctly, this PR should close the existing gap between inference with transformers + AutoGPTQ and inference with ExLlama, as the VRAM usage would become much more controlled. I'm rooting for it :)

Thanks! 🤗

Exciting PR! 

1. This was some debug / tests but removed hehe
2. The idea is to be able to save this and track usage + just make the api seamless. You can pass the cache instance but it should be able to be deduced. It's not a generation kwargs for me but a generation_config argument rather than that 
3. Attention mask is mostly what is left to adress TBH, because the control flow introduced are non trivial that is why I have a seperate PR that tried to take inspiration for JAX. Still working on it! Jax does update the attention mask with the cache see:

https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_flax_llama.py#L225

Main attention challenge is that we have different attention classes, which need different kind of masks. My goal is to only pass the 2d attention mask, and have each attention layers use it. 

How do I use this PR on npu device, here is my demo with errors.
```python
import torch
import torch_npu
from transformers import LlamaForCausalLM, LlamaTokenizer, TextStreamer



tokenizer = LlamaTokenizer.from_pretrained(
    "../Llama-2-7b-chat-hf",
    device_map="npu:2"
)
res = tokenizer("hello")

llama_model = LlamaForCausalLM.from_pretrained(
    "../Llama-2-7b-chat-hf",
    device_map="npu:2"
)
streamer = TextStreamer(tokenizer)


print("load successful")
while True:
    ins = input("user: ")
    res = tokenizer.encode(ins, return_tensors="pt").to("npu:2")
    print(res.dtype)
    outputs = llama_model.generate(
        inputs=res,
        streamer=streamer,
        max_new_tokens=10,
    )
```

Traceback (most recent call last):
  File "test.py", line 27, in <module>
    streamer=streamer,
  File "/usr/local/conda/envs/PyTorch2.1.0/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/usr/local/conda/envs/PyTorch2.1.0/lib/python3.8/site-packages/transformers/generation/utils.py", line 1790, in generate
    return self.sample(
  File "/usr/local/conda/envs/PyTorch2.1.0/lib/python3.8/site-packages/transformers/generation/utils.py", line 2887, in sample
    outputs = self(
  File "/usr/local/conda/envs/PyTorch2.1.0/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/conda/envs/PyTorch2.1.0/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/conda/envs/PyTorch2.1.0/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 1206, in forward
    outputs = self.model(
  File "/usr/local/conda/envs/PyTorch2.1.0/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/conda/envs/PyTorch2.1.0/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/conda/envs/PyTorch2.1.0/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 1093, in forward
    layer_outputs = decoder_layer(
  File "/usr/local/conda/envs/PyTorch2.1.0/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/conda/envs/PyTorch2.1.0/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/conda/envs/PyTorch2.1.0/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 820, in forward
    hidden_states, self_attn_weights, present_key_value = self.self_attn(
  File "/usr/local/conda/envs/PyTorch2.1.0/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/conda/envs/PyTorch2.1.0/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/conda/envs/PyTorch2.1.0/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 438, in forward
    raise ValueError(
ValueError: Attention mask should be of size (1, 1, 2, 2), but is torch.Size([1, 2])

This PR is still very much in draft mode. A script can be found in the tests, but does not support everything yet! Would recommend you to wait a tad bit, should be in a better state by the end of the week

cc @oobabooga the PR is more or less ready, would love to get your feedback from using it in `oobabooga`! (this needs a bit of custom work to compile only the decoding step, but should already help!) 

I had a look at them, the main pain point here is: 
- the overhead on a generation is ~2 token /s 
- the overhead of not using the `position_ids` tensor in compiled mode is more than 40 tokens / s
s

Alternative for ROPE that does not break cuda graphs + faster (5%) than vanilla recomputing. Issue is you save these tensor so loading will be wrong for sure.
```python
class LlamaRotaryEmbedding(nn.Module):
    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
        super().__init__()

        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        # Build here to make `torch.jit.trace` work.
        self._set_cos_sin_cache(
            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
        )

        
    def _set_cos_sin_cache(self, seq_len, device, dtype):
        self.max_seq_len_cached = seq_len
        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)

        freqs = torch.outer(t, self.inv_freq)
        # Different from paper, but it uses a different permutation in order to obtain the same calculation
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)

        
    def forward(self, x, position_ids, seq_len=None):
        # x: [bs, num_attention_heads, seq_len, head_size]
        if seq_len > self.max_seq_len_cached:
            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
        cos = F.embedding(position_ids, self.cos_cached)[0,:,:].to(dtype=x.dtype)
        sin = F.embedding(position_ids, self.sin_cached)[0,:,:].to(dtype=x.dtype)
        return cos, sin
```
<img width="896" alt="image" src="https://github.com/huggingface/transformers/assets/48595927/0e37d453-2712-462b-9e33-e6183e3bd983">


@ArthurZucker when will the next minor release (4.37.3?) of transformers come out? Will this PR be included in 4.37.3? 

@tsengalb99 v4.37.3 would be a micro release. These are reserved for patch fixes which resolve newly introduced breaking changes in the repo or fixing regressions from the last release.

The next minor release is v4.38. We release on a roughly monthly schedule, so would be in around two weeks. 

If you want to use this feature immediately, you can [install from source](https://huggingface.co/docs/transformers/installation#install-from-source). 

Is this working in general or only for some specific hardware, etc? 
I tried using fresh installs with WSL2, ubuntu, Cuda 11.8 (and again with 12.1) on an RTX A5000

 I installed from source today and tried the test script (https://gist.github.com/ArthurZucker/2dd607c4333ac4c489af30f54a1d8a2d)

 but the model.generate for compiled forward throws runtimeError: attn_bias is not correctly aligned (strideH). attn_bias.stride(1) = 21, and should be a multiple of 8.   



What is the recommended way to use cuda graphs with this PR? The torch cuda graph wrapper does not appear to be working iirc due to this line `attention_mask is not None and not torch.all(attention_mask[..., 0] == 1) and q_len != 1` / the equivalent in all model files.

This line is probably gonne be removed in #28937 
I'update once this is ready.
I use torch nightly so torch 2.3! 

@ArthurZucker Great to see the PR has been merged. I tried with the latest nightlies on H100, and running into errors (full [logs](https://gist.github.com/chauhang/a0dabe50c9958af5c92a3b1a23d1e888))

with sdpa attention: 
`RuntimeError: Event device type CUDA does not match blocking stream's device type CPU.`

with flash_attention_2 (fails in import itself):
`flash_attn_2_cuda.cpython-310-x86_64-linux-gnu.so: undefined symbol: _ZN3c104cuda9SetDeviceEi`

Nice, could you share a full stack trace? 🤗 


@chauhang 

> with flash_attention_2 (fails in import itself):

This means that the your local torch version does not match the torch version flash-attn was built against. You would need to uninstall & rebuild flash-attn.

Could you help clarify the removal of comment?
![image](https://github.com/huggingface/transformers/assets/5203025/f5579c0a-362f-4e6c-a28a-52f4c3128293)
It's probably removed by mistake as attn_weights is still not supported in FA2.
We should also add a warning before setting it to False, just like the SDPA counterpart.

Hi @ArthurZucker It seems that the increase in VRAM could potentially lead to out of memory (OOM) [comment1](https://github.com/huggingface/transformers/pull/29285#issuecomment-1979879993) [comment2](https://github.com/huggingface/transformers/pull/29285#issuecomment-1980165355), as pointed out in this PR by @danielhanchen 

> It seems like a change was made in another PR which allocates a causal mask of size (16384, 16384) https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L940
> 
> The triu causes the causal mask to upcast to float32, using 16384^2 * 4bytes = 1GB of extra VRAM. We have n^2 * 4 / 1024 / 1024 = 37.25GB in your screenshot, so I'm assuming you're also doing RoPE Scaling to 100K context length? So ie a (100K, 100K) matrix was trying to be created.

Could you please take a look into it?

@paulcx your issue is related to this one (https://github.com/huggingface/transformers/issues/29484) -- let's keep the discussion there! :)

Yes! And @paulcx I'm sorry this broke for you

See https://github.com/huggingface/transformers/pull/29241 which alleviates but does not fix the issue @paulcx 

Does this work for llava? from my testing it doesnt

> See https://github.com/huggingface/transformers/pull/29241 which alleviates but does not fix the issue @paulcx 

does this temporary fix work for 200K?

No, you have to update the `max_position_embedding`. It is allocating 200k because you set it to 200K while you machine does not support 200K input

`max_position_embedding` is only use for the `causal_mask` now, while previously it was used for `sin` and `cos`. In both cases it was to `cache` the maximum number of positions that will be passed to the model. If you have 200K context length, that does not mean you can do inference / training with it!

We can also just remove it, but then you need to allocate the causal_mask at each forward pass

@ArthurZucker If I understand correctly, my thought is to lower the max_position_embedding, such as 4096, because it only affects the initial position embedding and caching for 200K. But during inference, lengths approaching 200K will still be calculated, just slower. This workaround, it can ensure normal training and inference for non-200K cases. Is my understanding correct?

This was fixed by #29753 ! Sorry @paulcx for the inconvenience. For static / compile cache you should still reduce the max position embedding or it will OOM 😉 

> static / compile cache

Thank you and great work on new release @ArthurZucker.

Would you mind clarifying the use case of "static / compile cache" in release note? I'm not sure if I understand correctly.

It is mostly this: https://gist.github.com/ArthurZucker/af34221def212259b43d55a2811d2dbb, you can get x4 generation speed in transformers with torch compile and static cache!

> It is mostly this: https://gist.github.com/ArthurZucker/af34221def212259b43d55a2811d2dbb, you can get x4 generation speed in transformers with torch compile and static cache!

Is this expected to work with llava-next?

I believe so yes, if not we can add support for it

Feel free to open an issue if it doesnt work

I have tried and it dont work because the vision tower changes the shape of inputs after encoding to patches. Also, it doesnt work for bnb 4 bits

bnb is a different issue, torch.compile might not support this (int8 yes). 
For the encoder part cc @NielsRogge could be nice

we are using NF4 for bnb

> It is mostly this: https://gist.github.com/ArthurZucker/af34221def212259b43d55a2811d2dbb, you can get x4 generation speed in transformers with torch compile and static cache!

@ArthurZucker  Just checking have you added this to `model.generate`, or we still have to follow your script there to use static KV cache?

@gante is working on this here #29374