DynamicCache results in too many torch recompiles after 4.51
Author: flishwang
Pull request: no
Labels: name: bug

### System Info

accelerate=1.6.0, OS=ubuntu 22.04, numpy=1.26.4, torch=2.6.0+cu124, python=3.10

### Who can help?

@ArthurZucker 


### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [x] My own task or dataset (give details below)

### Reproduction


```
import torch
from transformers import AutoModelForCausalLM
sequence_label_locs = torch.as_tensor([,-40,-20,-10,-1]).cuda()
level_ids_gpu = torch.as_tensor([76,77,78,79,80]).cuda()
cache = None
cls_ids_gpu = torch.as_tensor([103,130,1166,1366,3366]).cuda()
model = AutoModelForCausalLM.from_pretrained('Qwen/Qwen2-1.5B-Base').half().cuda()
model = torch.compile(model)
batch = torch.randint(low=0,high=100000,size=(8,4096)).cuda()
mask = torch.ones((8,4096),dtype=torch.int64).cuda()

loc_stride=256
@torch.no_grad()
def model_predict(model, batch, mask,sequence_label_locs, level_ids_gpu, cls_ids_gpu, cache):
    model_predict_argmax = None
    past_kv_cache = cache
    predicts = []
    non_zero_loc = mask.argmax(-1).min() // loc_stride * loc_stride
    batch=batch[:,non_zero_loc:]
    mask=mask[:,non_zero_loc:]
    sequence_label_locs = sequence_label_locs - non_zero_loc

    for idx in range(len(sequence_label_locs)):
        if idx == 0:
            input_ids = batch[:,:sequence_label_locs[idx]]
        else:
            input_ids = torch.cat([model_predict_argmax,
                                   batch[:,sequence_label_locs[idx-1]+1:sequence_label_locs[idx]]],1)
        input_mask = mask[:,:sequence_label_locs[idx]]
        outputs = model(input_ids,attention_mask = input_mask,past_key_values = past_kv_cache, use_cache=True,
                        logits_to_keep = 1,
                        return_dict = True)
        last_predicts = outputs.logits[:,-1:,:]
        predicts.append(last_predicts)
        past_kv_cache = outputs.past_key_values
        if idx == 0:
            model_predict_argmax = cls_ids_gpu[last_predicts[:,:,cls_ids_gpu].argmax(-1)]
        else:
            model_predict_argmax = level_ids_gpu[last_predicts[:, :, level_ids_gpu].argmax(-1)]
    return torch.cat(predicts,1)

for i in range(0,4000,10):
    mask[:,:i]=0
    outs = model_predict(model, batch, mask,sequence_label_locs, level_ids_gpu, cls_ids_gpu, cache)

```


### Expected behavior


Run the upper scripts with TORCH_LOGS=recompiles
If I use transformer==4.49, only several recompiles happes.
But If I use transformers>=4.51, re-compiling always happens if I changed the input sequences' length.

I notice that it seems to be related to the following code:
https://github.com/huggingface/transformers/blob/7a3e208892c06a5e278144eaf38c8599a42f53e7/src/transformers/cache_utils.py#L442
```
                not self.key_cache[layer_idx].numel()  # prefers not t.numel() to len(t) == 0 to export the model
```
which in transformers<=4.49 is 
```
                not len(self.key_cache[layer_idx])
```


If I modified the code to
```
                not self.key_cache[layer_idx].shape[0]
```
, the re-compiling times could also be reduced.

I'm not expert for torch/transformers, and I'm not sure which project (pytorch or transformers?) this bug belongs to.
I'm also not sure if the proposed modification would break the model export procedure.

Comments:
Hey @flishwang !

`DynamicCache` should not be used when compiling a model. For compilations we have [`StaticCache`](https://huggingface.co/docs/transformers/v4.49.0/kv_cache#static-cache) which pre allocates maximum length for keys and values. 

You can take a look at [this section](https://huggingface.co/docs/transformers/v4.49.0/en/llm_optims#static-kv-cache-and-torchcompile) in docs for how to use static cache with compile. The easiest way would be to call `generate(**inputs, cache_implementation="static")` but it can be configured for specific use cases as well

@zucchini-nlp one issue is static cache is not supported with speculative decoding yet. Are you aware of any efforts to implement static cache + assisted decoding? Thx

@mklasby not that of I know. We had a community contribution in the past and realized that supporting static cache in the assistant model would require us to handle a lot of egde cases. cc @gante if I missed anythin

Speculative + Static Cache is still an open issue, yes! If some contributor wants to take a stab, I'd be happy to support them :D 

Otherwise, I'll pick it up in a few weeks: I'll have bandwidth to explore throughput optimizations 😎 

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.

Hi @gante , I'd like to take this up despite being a beginner. I believe that with some guidance and my readiness to contribute full-time, I should be able to get it done. Do you have any preliminary resources or tips that could help me get started?

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.