Whisper Static Cache 
Author: huseinzol05
Pull request: yes
Labels: name: Audio

name: Cache

# Static Cache for Whisper

This enable to use `torch.compile` for Whisper generation to enable faster generation, example https://gist.github.com/huseinzol05/9aff34ec1427ee8c92240cb4f3cc0c88

Compiled static cache able to achieve 186.26it/s while non-compiled got 150.20it/s .

## Still work in progress

1. Current forked only work to use static cache, need to follow caching steps as Llama.
2. There are so many conditions need to fulfill first.
3. Only worked on Pytorch `2.4.0.dev20240508+cu121` version, not yet released as stable for custom function `reduce-overhead` torch compile.


Comments:
Thank you very much @huseinzol05 for the work. 
Here's a version with HQQ 4-bit using the torchao backend. As expected there's a good speed-up with the static cache and fullgraph compilation: https://gist.github.com/mobicham/ecfe09a48efb11e4014386901a5c6cce

```
GPU: 4090
orig - no compile : 48 it/sec
orig + compiled   : 227 it/sec

hqq - no compile  : 42 it/sec
hqq + compile     : 308 it/sec
```


Will it be merged? @younesbelkada

cc @sanchit-gandhi 

@kadirnar , this PR is not ready to merge, or you can continue to work on it to fulfill no 1, 2 and 3. But if you want to use it, you have to split the audio into 30s chunks with overlap and feed into encoder-decoder process, feel free to add `temperature` and `top_k` like https://github.com/pytorch-labs/gpt-fast/blob/main/generate.py#L52

Will it work if I run these codes?  Also, should I make any changes to the gpt-fast library?

https://gist.github.com/huseinzol05/9aff34ec1427ee8c92240cb4f3cc0c88

yeah it should work, i use it in my prod, but dont forget to warmup the static cache multiple time first

> yeah it should work, i use it in my prod, but dont forget to warmup the static cache multiple time first

I ran the notebook file. It gives this error.
```bash
File /usr/local/lib/python3.10/dist-packages/transformers/cache_utils.py:484, in WhisperStaticCache.__init__(self, config, dtype, device, existing_cache, batch_size)
    482 torch._dynamo.mark_static_address(e_key_cache)
    483 torch._dynamo.mark_static_address(e_value_cache)
--> 484 e_key_cache[:, :, :, :] = existing_cache[k][2].clone()
    485 e_value_cache[:, :, :, :] = existing_cache[k][3].clone()
    486 self.key_cache.append(new_layer_key_cache)File /usr/local/lib/python3.10/dist-packages/transformers/cache_utils.py:484, in WhisperStaticCache.__init__(self, config, dtype, device, existing_cache, batch_size)
    482 torch._dynamo.mark_static_address(e_key_cache)
    483 torch._dynamo.mark_static_address(e_value_cache)
--> 484 e_key_cache[:, :, :, :] = existing_cache[k][2].clone()
    485 e_value_cache[:, :, :, :] = existing_cache[k][3].clone()
    486 self.key_cache.append(new_layer_key_cache)

IndexError: tuple index out of range
```

I just reran and no issue, super weird, which line is that you the error?

We did a comparison of the performance of the torch compiled version with static cache and its HQQ variants (4,3,2 and 1.58 bits) on both short-form audio (open_asr_eval) and long-form audio (internal test benchmark). 

Here is the link to the blog post: https://mobiusml.github.io/whisper-static-cache-blog/
Colab Notebook: https://colab.research.google.com/drive/18Zs-oG1Ztco3cfnNexcHDi-Zn9vk2RJ5?usp=sharing

I think the speech community can benefit a lot from this speed-up once integrated into transformers 🤗 !

Any progress on this folks? Is there a timeline for a general static support in transformers? We are very excited to see this officially supported in transformers!

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.

Will you merge this pull request? @sanchit-gandhi 

Closing this PR: `whisper` + compilation had a few sensible design decisions, as shown in the discussion above, so we took charge of adding static caches to whisper ([PR](https://github.com/huggingface/transformers/pull/31166))

Thank you for kickstarting the process and for the discussion 🤗