KV cache in qwen2
Author: hvlgo
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.40.2
- Platform: Linux-5.4.0-200-generic-x86_64-with-glibc2.31
- Python version: 3.10.4
- Huggingface_hub version: 0.26.2
- Safetensors version: 0.4.5
- Accelerate version: 1.1.1
- Accelerate config:    not found
- PyTorch version (GPU?): 2.0.1+cu117 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help?

_No response_

### Information

- [x] The official example scripts
- [ ] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Anyway to run Qwen, the issue is more about the code implementation. In [this line](https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2/modeling_qwen2.py#L292), the key_states and value_states are all computed but not using the cache. So, has the KV cache been implemented correctly?

### Expected behavior

I am studying the implementation of KV cache. Can someone help clarify my doubts? Thanks!

Comments:
EDIT: Bug below was fixed in master, just not released to pip yet. It also only affected the Qwen2-VL models and not the Qwen models, apparently.
<details>
<summary>An apparently unrelated issue for Qwen2-VL</summary>
<br>

I was just about to open an issue about this. Good timing. The KV cache for Qwen2-VL models have not been implemented correctly. Here's a slightly modified example:
```python
import copy
import torch
from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, DynamicCache, StaticCache

# Load the model in half-precision on the available device(s)
model = Qwen2VLForConditionalGeneration.from_pretrained(
    "Qwen/Qwen2-VL-2B-Instruct", torch_dtype=torch.float16, device_map="auto"
)

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-VL-2B-Instruct")

prompt_cache = StaticCache(config=model.config, max_batch_size=1, max_cache_len=1024, device="cuda", dtype=torch.float16)
#prompt_cache = DynamicCache() # <-- this also doesn't work

INITIAL_PROMPT = "You are a helpful assistant."
inputs_initial_prompt = tokenizer(INITIAL_PROMPT, return_tensors="pt").to("cuda")
# This is the common prompt cached, we need to run forward without grad to be abel to copy
with torch.no_grad():
     prompt_cache = model(**inputs_initial_prompt, past_key_values = prompt_cache).past_key_values

prompts = [" Help me to write a blogpost about travelling.", " What is the capital of France?"]

responses = []
responses_expected = []

for prompt in prompts:
    new_inputs = tokenizer(INITIAL_PROMPT + prompt, return_tensors="pt").to("cuda")
    past_key_values = copy.deepcopy(prompt_cache)
    
    # behavior with cache
    outputs = model.generate(**new_inputs, past_key_values=past_key_values,max_new_tokens=20)
    response = tokenizer.batch_decode(outputs)[0]
    responses.append(response)
    
    # expected behavior
    outputs_expected = model.generate(**new_inputs,max_new_tokens=20)
    response_expected = tokenizer.batch_decode(outputs_expected)[0]
    responses_expected.append(response_expected)


print(responses)
print(responses_expected)
```
gives:
```
['You are a helpful assistant. Help me to write a blogpost about travelling. I with with with with with with with. with,. with,......', 'You are a helpful assistant. What is the capital of France? I a a a a a a a a a a a a a a a a a a a']
['You are a helpful assistant. Help me to write a blogpost about travelling. I want to write a blogpost about travelling. I want to write a blogpost about travelling.', 'You are a helpful assistant. What is the capital of France? The capital of France is Paris. It is located in the north of the country and is the largest']
```

If I use, say, Mistral-0.3, the results match the expected results.
</details>



Thank you, but the correctness of the above code does not seem to ensure that the KV cache implementation is correct. It appears to store the KV cache but does not actually use it. For example, in [this line](https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py#L545), theoretically, it should load the past K and V from the cache, calculate K and V only for the new input, and compute attention for only the necessary part. What I mean is that the KV cache is not being utilized, and there is no acceleration effect. Am I misunderstanding this, or is there an issue with the implementation?

Hi @hvlgo , has it been resolved? I also see that kvcache doesn't seem to be implemented yet. The `use_cache` is not being used at all.

Same question. is it solved?

To sovle this problem, your custom inference code may split to 2 stage, prefill and incremental, more details could ask ai