StaticCache crashes when the batch-size changes
Author: mobicham
Pull request: no
Labels: name: bug

### System Info

```
- `transformers` version: 4.57.1
- Platform: Linux-6.8.0-1043-aws-x86_64-with-glibc2.35
- Python version: 3.11.7
- Huggingface_hub version: 0.36.0
- Safetensors version: 0.6.2
- Accelerate version: 1.12.0
- Accelerate config:    not found
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.8.0+cu128 (CUDA)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA L40S
```

### Who can help?

@gante @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

Hi folks! I was doing some testing with static cache and noticed that it breaks when the batch-size changes:

```
RuntimeError: index_copy_(): Source/destination tensor must have same slice shapes. Destination slice shape: 4 20 64 at dimension 2 and source slice shape: 8 20 64 at dimension 0.
```

Is this expected behavior or a known limitation? The issue is that, in order to use static cache with cudaraphs, the cache needs to be pre-allocated. If it crashes like this, we need a separate cache for each batch-size which quickly blows up the VRAM usage...

```Python
import torch
from transformers import WhisperForConditionalGeneration, AutoProcessor
from transformers.cache_utils import StaticCache, EncoderDecoderCache
import math
import numpy as np

device = 'cuda:0'
torch_dtype = torch.float16
model_id = "openai/whisper-large-v3-turbo"

model = WhisperForConditionalGeneration.from_pretrained(model_id, torch_dtype=torch_dtype, attn_implementation="sdpa", device_map=device) 
model.generation_config.cache_implementation = "static"

@torch.no_grad()
def run_encoder(model, labels, encoder_outputs, past_key_values, cache_position=None):
    out_decoder = model.model.decoder(
        labels, 
        encoder_hidden_states=encoder_outputs,
        past_key_values = past_key_values,
        cache_position=cache_position,
        use_cache = True,
        return_dict=True,
    ) 

    cur_token = model.proj_out(out_decoder.last_hidden_state[:,-1:]).argmax(axis=-1) 
    past_key_values = out_decoder.past_key_values

    return cur_token, past_key_values

max_batch_size = 32
max_cache_len = 256
enc_len = 1500
decoder = model.model.decoder

# Cache
self_cache = StaticCache(
    config=decoder.config,
    max_batch_size=max_batch_size,
    max_cache_len=max_cache_len,
    device=device,
    dtype=torch_dtype,
)

cross_cache = StaticCache(
    config=decoder.config,
    max_batch_size=max_batch_size,
    max_cache_len=enc_len,
    device=device,
    dtype=torch_dtype,
)

past_key_values = EncoderDecoderCache(self_cache, cross_cache)

for bs in [8, 4, 2, 1]:
	past_key_values.reset()
	assert bs <= max_batch_size, "batch_size should be <= max_batch_size"
	seq_length = 3
	labels = torch.tensor([[50258, 50259, 50360]] * bs, device=device, dtype=torch.int64)
	cache_position = torch.arange(seq_length, device=device, dtype=torch.int64)
	encoder_outputs = torch.randn([bs, enc_len, 1280], device=device, dtype=torch_dtype)

	cur_token, past_key_values_out = run_encoder(model, labels, encoder_outputs, past_key_values, cache_position=cache_position)

# RuntimeError: index_copy_(): Source/destination tensor must have same slice shapes. Destination slice shape: 4 20 64 at dimension 2 and source slice shape: 8 20 64 at dimension 0.
```

### Expected behavior

A single static cache instance should work across multiple batch-sizes as long as the batch-size <= `max_batch_size`

Comments:
Hmm, I don't think this is supported! The cache is intended to be initialized and used with a single static batch size. Cache creation should be fast, though - maybe just make a new one each time your batch size shifts?

If I create a new one, it would break cudagraphs. How do you recommend changing the example above by creating a new cache everytime but still avoiding torch.compile recompilation?

Btw, we had a PR for that which unfortunately didn't make it to the release (https://github.com/huggingface/transformers/pull/37394). I agree that having a `max_batch_size` will be a nice feat and I can work on it next week


@zucchini-nlp can I work on this.

Thank you! I think that's the most important feature to actually use static cache in a deployment not just with toy examples.
 
I guess the update could simply use `k/v_cache[:batch_size, ....]` which shouldn't create a copy and keep it compatible with cudagraphs 🤔 

Yeah I've implemented the list slicing fix should I submit a PR for it ? 
It does work on the reproduction script . 
@mobicham  @zucchini-nlp 

@i3hz sure, if you have a working solution, feel free to make a PR and tag me. I think the general design will be similar to the prev PR we had

You can also run a benchmark by creating a big static cache and using several batch sizes smaller than cache size, so we know that there's smth wrong if the performance doesn't match the performance on `main` branch

@i3hz Thank you! Have you checked if it works with cudagraphs via `torch.compile` ?
```Python
run_encoder = torch.compile(run_encoder, mode='reduce-overhead', fullgraph=True)

for bs in [8, 4, 2, 1]:
    past_key_values.reset()
    assert bs <= max_batch_size, "batch_size should be <= max_batch_size"
    seq_length = 3
    labels = torch.tensor([[50258, 50259, 50360]] * bs, device=device, dtype=torch.int64)
    
    encoder_outputs = torch.randn([bs, enc_len, 1280], device=device, dtype=torch_dtype)

    cache_position = torch.arange(seq_length, device=device)
    cur_token, past_key_values_out = run_encoder(model, labels, encoder_outputs, past_key_values, cache_position=cache_position)

    cache_position = torch.tensor(seq_length, device=device)
    cur_token, past_key_values_out = run_encoder(model, cur_token, encoder_outputs, past_key_values_out, cache_position=cache_position)
```

I did a benchmark run with the slicing method against the baseline (where we create a new one every time ) .
```
[old] Batch= 1 | Cache_Max= 1 | Latency: 4.05 ms
[new ] Batch= 1 | Cache_Max=64 | Latency: 3.91 ms
[old] Batch= 8 | Cache_Max= 8 | Latency: 3.99 ms
[new ] Batch= 8 | Cache_Max=64 | Latency: 4.26 ms
[old] Batch=32 | Cache_Max=32 | Latency: 5.57 ms
[new ] Batch=32 | Cache_Max=64 | Latency: 5.53 ms
[old] Batch=64 | Cache_Max=64 | Latency: 11.51 ms
[new ] Batch=64 | Cache_Max=64 | Latency: 11.45 ms
```
the latency is pretty close 

How about cudagraphs @i3hz ? It's the main reason to use static cache

yeah @mobicham  the thing it it's working with gpt 2 but it does crash when im using whisper and I can't really figure out why

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.