Fix mistral generate for long prompt / response
Author: lorabit110
Pull request: yes

# What does this PR do?

Fix the below issue:
When use mistral model to generate texts, if prompt + max_tokens > 4095 and use_cache=True, you would get the below error. 
ValueError: past key much have a shape of (`batch_size, num_heads, self.config.sliding_window-1, head_dim`), got torch.Size([1, 8, 2989, 128]).

This PR fixes the logic that determine which part of the cached key and value should be used for predicting future tokens. 

Fixes https://github.com/huggingface/transformers/issues/27682

## Before submitting
- [X] 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.
- [X] 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).
- [X] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@Bam4d 


Comments:
cc @gante @ArthurZucker 

Hey @lorabit110 👋 

I'm not sure if I follow the need for a fix here :) On my end, the test case you added passes:
```py
from transformers import AutoModelForCausalLM
import torch

EXPECTED_OUTPUT_TOKEN_IDS = [306, 338]

input_ids = [1] + [306, 338] * 2048
model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1", device_map="auto", torch_dtype=torch.float16, use_flash_attention_2=True)
input_ids = torch.tensor([input_ids]).to(model.model.embed_tokens.weight.device)
print(input_ids.shape)
generated_ids = model.generate(input_ids, max_new_tokens=2)
print(generated_ids[0][-2:].tolist() == EXPECTED_OUTPUT_TOKEN_IDS)
# True
```

Can you share a short reproducible script that results in the failure?

> Hey @lorabit110 👋
> 
> I'm not sure if I follow the need for a fix here :) On my end, the test case you added passes:
> 
> ```python
> from transformers import AutoModelForCausalLM
> import torch
> 
> EXPECTED_OUTPUT_TOKEN_IDS = [306, 338]
> 
> input_ids = [1] + [306, 338] * 2048
> model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1", device_map="auto", torch_dtype=torch.float16, use_flash_attention_2=True)
> input_ids = torch.tensor([input_ids]).to(model.model.embed_tokens.weight.device)
> print(input_ids.shape)
> generated_ids = model.generate(input_ids, max_new_tokens=2)
> print(generated_ids[0][-2:].tolist() == EXPECTED_OUTPUT_TOKEN_IDS)
> # True
> ```
> 
> Can you share a short reproducible script that results in the failure?

@gante 

Actually, we need to set max_new_tokens to at least 3 to trigger the error. The below code can reproduce the issue: 

```python
from transformers import AutoModelForCausalLM
import torch

EXPECTED_OUTPUT_TOKEN_IDS = [306, 338]

model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1", device_map="auto", torch_dtype=torch.float16, use_flash_attention_2=True)
input_ids = [1] + [306, 338] * 2048
input_ids = torch.tensor([input_ids]).to(model.model.embed_tokens.weight.device)
generated_ids = model.generate(input_ids, max_new_tokens=4)
print(generated_ids[0][-2:].tolist() == EXPECTED_OUTPUT_TOKEN_IDS)
```

<img width="1494" alt="Screenshot 2023-11-17 at 11 14 35 AM" src="https://github.com/huggingface/transformers/assets/108375850/41249363-9b85-40e9-82f7-a63d9bc1e1f0">


Can anyone take a look?

@younesbelkada I see. I'm missing one variable here, which is how many tokens the model has seen so far through its cache -- does any of the inputs to `prepare_inputs_for_generation` have this information? 

Without it, I'm not sure how to properly slice `input_ids` without the setting the assumption that only one token can be consumed at a time (which would disable the use of `mistral` with assisted generation and the [newly added ability](https://huggingface.co/docs/transformers/main/en/llm_tutorial_optimization#321-multi-round-conversation) to pass `past_key_values` across generations)

Let's make sure the models are quantized for testing

Thanks everyone for the discussion. I will need to spend sometime to figure out how to implement assisted decoding test case. 

All comments have been addressed. I need a maintainer to approve and run a workflow in order to merge the PR.

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

Thanks @lorabit110 for the fix 😉