Represent query_length in a different way to solve jit issue
Hi @ArthurZucker @younesbelkada  @sgugger 

Thanks for contributing to the MPT model. I found a jit issue with this model.
```python
from transformers import AutoModelForCausalLM, AutoConfig
from optimum.intel.generation.modeling import jit_trace

model = AutoModelForCausalLM.from_pretrained("mosaicml/mpt-7b")

jit_model = jit_trace(model=model, task="text-generation", use_cache=True)
```
When I tried to trace mpt model, I got this error
![image](https://github.com/huggingface/transformers/assets/107918818/67244f76-625a-4080-aa54-34b8b032545a)

This is because float values like seq_length and query_length are detected as tensor in trace mode. When we set `query_length = seq_length` and `seq_length += past_key_value[0].shape[2]`, the `seq_length` is changed too which is unexpected.
![MicrosoftTeams-image (1)](https://github.com/huggingface/transformers/assets/107918818/8a6cff61-3a07-4f1a-96d6-d17a7efe9292)

So I use a more clean way to set `query_length` and it can also avoid the jit issue.
`query_length = seq_length if past_key_value is None else seq_length + past_key_value[0].shape[2]`

Would you please help me review it? Thanks!



_The documentation is not available anymore as the PR was closed or merged._