Fix static generation when compiling! 
Author: ArthurZucker
Pull request: yes

# What does this PR do?
Fixes the static cache generation. Comes with #27931 

thanks @OlivierDehaene for the insight 

https://gist.github.com/ArthurZucker/ae0a86ef8f841c0ef69aaa52ccbc0b03 benchmark

- fixes issue with FlashAttention: when the cache is padded you need the full attention mask, otherwise generations will be wrong with `generate` because the first forward will be fully causal. 
- fixes graph runs: the cache positions have to be stateless, they are otherwise ignored by the model and the compiled generation are random
- fixes potential BC by guarding the use of cache positions




FA2 potential fix if compiled worked: 
```python 
            # we slice the states for static kv cache to be supported in FA2. Not sure it's a must as compile fails
            if (cache_position is not None):  
                key_states = key_states[:, :, : cache_position[-1] + 1, :]
                value_states = value_states[:, :, : cache_position[-1] + 1, :]
```

but I have slowdowns: 
Slicing
![image](https://github.com/huggingface/transformers/assets/48595927/26d4ee49-266e-4289-9e01-88b372138bc6)
vs no Slicing
![image](https://github.com/huggingface/transformers/assets/48595927/397d43c9-6c24-45c9-a2fe-9e7cc7df5a98)


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

I'm not sure adding a new argument `cache_position` to the forward call of the model is strictly backwards compatible. Here's an example to motivate this.

The following works on `transformers==4.37.2`:
```python
import torch
from transformers import AutoModelForCausalLM, LlamaTokenizer

model = AutoModelForCausalLM.from_pretrained("trl-internal-testing/tiny-random-LlamaForCausalLM", attn_implementation="eager")
tokenizer = LlamaTokenizer.from_pretrained("trl-internal-testing/tiny-random-LlamaForCausalLM")

# random input id
inputs = tokenizer("Hey there", return_tensors="pt", return_attention_mask=True)

position_ids = inputs.attention_mask.long().cumsum(-1) - 1
position_ids.masked_fill_(inputs.attention_mask == 0, 1)

with torch.no_grad():
    logits = model.forward(**inputs, position_ids=position_ids).logits
```

If we run the same code on this PR, we get the following error:
```
  File "/Users/sanchitgandhi/transformers/src/transformers/models/llama/modeling_llama.py", line 352, in forward
    attn_weights = attn_weights + causal_mask
                   ~~~~~~~~~~~~~^~~~~~~~~~~~~
RuntimeError: The size of tensor a (3) must match the size of tensor b (2048) at non-singleton dimension 4
```

<details>

<summary> Full traceback: </summary>

```
  File "/Users/sanchitgandhi/transformers/debug_llama.py", line 14, in <module>
    logits = model.forward(**inputs, position_ids=position_ids).logits
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/sanchitgandhi/transformers/src/transformers/models/llama/modeling_llama.py", line 1106, in forward
    outputs = self.model(
              ^^^^^^^^^^^
  File "/Users/sanchitgandhi/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/sanchitgandhi/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/sanchitgandhi/transformers/src/transformers/models/llama/modeling_llama.py", line 950, in forward
    layer_outputs = decoder_layer(
                    ^^^^^^^^^^^^^^
  File "/Users/sanchitgandhi/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/sanchitgandhi/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/sanchitgandhi/transformers/src/transformers/models/llama/modeling_llama.py", line 694, in forward
    hidden_states, self_attn_weights, present_key_value = self.self_attn(
                                                          ^^^^^^^^^^^^^^^
  File "/Users/sanchitgandhi/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/sanchitgandhi/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/sanchitgandhi/transformers/src/transformers/models/llama/modeling_llama.py", line 352, in forward
    attn_weights = attn_weights + causal_mask
                   ~~~~~~~~~~~~~^~~~~~~~~~~~~
RuntimeError: The size of tensor a (3) must match the size of tensor b (2048) at non-singleton dimension 4
```

</details>

This is because `cache_positions` is not specified to the forward call, and so defaults to `None`. When we do our reshape in the attention layer:
https://github.com/huggingface/transformers/blob/56768a028b52290ff55f0ad3902679f6dafc568e/src/transformers/models/llama/modeling_llama.py#L352-L353
instead of reshaping to `[  :, :, cache_position, : key_states.shape[-2]]`, we reshape to `[ :, :, None, : key_states.shape[-2]]`. So instead of slicing, we insert an extra dimension! This gives the size mismatch when we add the attention mask to the weights. The user needs to specify `cache_position` as an argument to the forward call in order for this to work.

Overall, I think we should avoid adding extra arguments that require code changes from the user, especially to the top-level modules which are already highly-used. What about a design more like Flax where we keep track of the `cache_position` internally in the `StaticCache` abstraction? This then requires no changes from the user

We can make it BC! this PR is not ready yet, but generate should check the past key value class and if signature can take cache_position, give them. Something like that. 

I'll work on making it BC! :) 

Thanks, merging asap 

<img width="1795" alt="image" src="https://github.com/huggingface/transformers/assets/48595927/78692611-934d-4035-99f6-3399cbbfd90a">

Slow tests are happy

Example of a breaking behaviour that I introduced while working on FA2: https://github.com/huggingface/transformers/pull/25598#issuecomment-1743338628 so we should be careful when adding new args in our modules

Hey @ArthurZucker, I discovered that this change actually breaks TPU...

Now, TPU training with FSDPv2 will produce loss with NaN. I haven't looked into your PR so I'm not sure why. Just bisecting til this change.

Mmm this might be a ROPE issue? #29109 might also play

Hi @ArthurZucker  I run your [benchmark script](https://gist.github.com/ArthurZucker/ae0a86ef8f841c0ef69aaa52ccbc0b03) with both transformers 4.38.0 and 4.38.2 but got error:

```
Traceback (most recent call last):
  File "/home/best_benchmark.py", line 99, in <module>
    generated_ids[:, cache_position] = input_ids.to(device).to(torch.int)
RuntimeError: shape mismatch: value tensor of shape [1686] cannot be broadcast to indexing result of shape [1, 2048]
```

It is probably out of date! I'll update it 

We'll actually push a full benchmark in `transformers` to make sur we always track this!