StaticCache crashes when the batch-size changes

[mobicham]

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
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• 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...

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

[Rocketknight1]

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?

[mobicham]

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?

[zucchini-nlp]

Btw, we had a PR for that which unfortunately didn't make it to the release (#37394). I agree that having a max_batch_size will be a nice feat and I can work on it next week

[Aaraviitkgp]

@zucchini-nlp can I work on this.

[mobicham]

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 🤔

[i3hz]

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

[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

[mobicham]

@i3hz Thank you! Have you checked if it works with cudagraphs via torch.compile ?

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)

[i3hz]

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

[mobicham]

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

[i3hz]

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