KV Cache Size Issue during Inference
Author: gopikrishnajha
Pull request: no

This is w.r.t inference using models like mistral7b or llama. 

In my understanding, KV cache size should grow as we process more tokens, however I see in the code that it shrinks as more tokens are processed. For example, in transformers/src/transformers/models/mistral/modeling_mistral.py, see the following code.

```key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)```

The cache ```past_key_value``` instead of growing, shrinks in size. Initial size is ```[batch_size, num_heads, seq_len, head_dim]``` and with increasing iterations, while ```batch_size, num_heads and head_dim``` remain the same, ```seq_len``` decreases. This results in a shrinking cache.

Can anyone explain why is the cache shrinking instead of growing in size?

Comments:
cc @ArthurZucker 

Hey! That's the normal behavior for auto regressive transformers. The key value states's shape decreases, while the actual cache increases in size. 
I don't know how you debugged, but the prefill (first forward) will be big, then each new forward will only add 1 to the sequence length dimension. 
Have never seen this leading to a "decrease" shape of the actual key_states and value_states

> Hey! That's the normal behavior for auto regressive transformers. The key value states's shape decreases, while the actual cache increases in size. I don't know how you debugged, but the prefill (first forward) will be big, then each new forward will only add 1 to the sequence length dimension. Have never seen this leading to a "decrease" shape of the actual key_states and value_states

That's what I was expecting too. But the size of ```past_key_value``` variable is decreasing in the sequence length dimension with every new forward call. 

```past_key_value``` variable should grow in size with every forward call, right? Or is ```past_key_value``` not the kv cache?

I don't know what you are checking, but the `key_states` used to compute the attention score, so after you update the `past_key_value` should grow. That is exactly why you cannot compile it: the size of the key and values grows

> I don't know what you are checking, but the `key_states` used to compute the attention score, so after you update the `past_key_value` should grow. That is exactly why you cannot compile it: the size of the key and values grows

In the ```modeling_mistral.py``` I am adding the following prints (after  ```past_key_value.update```) for knowing the current size.

```
if past_key_value is not None:
            cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
if (self.layer_idx == 0):
        print(len(past_key_value)) #always prints 1
        print(past_key_value[0][0].shape, past_key_value[0][1].shape)
        print(key_states.shape, value_states.shape)
```
The output of first few forward calls is following.

1
torch.Size([10, 8, 75, 128]) torch.Size([10, 8, 75, 128])
torch.Size([10, 8, 75, 128]) torch.Size([10, 8, 75, 128])
1
torch.Size([10, 8, 58, 128]) torch.Size([10, 8, 58, 128])
torch.Size([10, 8, 58, 128]) torch.Size([10, 8, 58, 128])
1
torch.Size([10, 8, 51, 128]) torch.Size([10, 8, 51, 128])
torch.Size([10, 8, 51, 128]) torch.Size([10, 8, 51, 128])
1
torch.Size([10, 8, 45, 128]) torch.Size([10, 8, 45, 128])
torch.Size([10, 8, 45, 128]) torch.Size([10, 8, 45, 128])
1
torch.Size([10, 8, 43, 128]) torch.Size([10, 8, 43, 128])
torch.Size([10, 8, 43, 128]) torch.Size([10, 8, 43, 128])

You can see that the cache size is decreasing from the ```seq_len``` dimension.

Also, when I am confirming the DRAM usage using ```htop```, it seems to be decreasing with each iteration.

Well this can also be due to beam search. After a few iterations you have less beams for example. There are various factors but recommending you to use greedy generation 

> Well this can also be due to beam search. After a few iterations you have less beams for example. There are various factors but recommending you to use greedy generation

I am already using greedy generation.

If you don't provide a reproduce there's no way I can know that 🤗 could you share a small snippet of how you are calling the model? 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.