Reduce by 2 the memory requirement in `generate()` 🔥🔥🔥
Author: Cyrilvallez
Pull request: yes
Labels: name: run-slow

# What does this PR do?

Change the data structure and implementation of `DynamicCache` to **halve** the memory requirement in `generate()` at no noticeable speed degradation.

## Reason

I was working on precise memory estimation of `generate()` and noticed that the expected memory peak and the one observed are different (observed one is much higher). I was able to track down the problem to the current implementation of `DynamicCache`.

Since `torch.cat()` creates a copy, and that during `generate` we loop over each token using the previous cache in the inputs (thus the old cache is referenced in the inputs so cannot be garbage collected), the current implementation basically has 2 copies of the **full** cache in every iteration. By changing the data structure from `Tuple[Tuple[Tensor]]` to `Tuple[Tuple[List[Tensor]]]`, I was able to avoid this copy and reduce the memory footprint of `generate()` by 2. 

However, that means that potentially thousands of tensors must be `cat()`'ed when `cache.update()` is called to feed the correct tensor to the `AttentionLayer`. This results in a speed degradation. I was able to mitigate this issue by periodically `cat()`'ing the tensors in the cache when there are more than `N` (I used 50 as of now), which incurs a negligible memory increase, as the cache for a sequence length of `N=50` is usually completely negligible compared to the size of the full cache (hundreds, or even more frequently thousands). This strategy almost completely removes the speeds degradation, allowing to get the best of both worlds. `N` could even be chosen dynamically in `generate()` depending on the input length and max new tokens, but 50 seemed like a good heuristic to start with.

Basically, at a very small performance penalty that is visible only for very large sequence length, we reduce the memory footprint by 2, which by itself allows to increase the batch size by 2 (so should be able to actually speed up the process as passing a sequence of 2 times the batch at lower speed should still be faster then runnings 2 loops at faster speed).

Of course, the best would still be to use a `StaticCache` as I saw that you started implementing, but a DynamicCache is still very much useful, and should not imply to double its effective memory footprint when used in loop referencing itself.

## Benchmark

Here you can see the benchmark I ran using a RTX 4090 (24 GB) and Mistral 7B, Llama2 7B and Llama3 8B.

### Fix batch size of 1 and input length of 300, variable new token number

![Mistral-7B-v0.1_memory_fix_batch.pdf](https://github.com/huggingface/transformers/files/15149403/Mistral-7B-v0.1_memory_fix_batch.pdf)
![Mistral-7B-v0.1_time_fix_batch.pdf](https://github.com/huggingface/transformers/files/15149405/Mistral-7B-v0.1_time_fix_batch.pdf)
![Llama-2-7b-hf_memory_fix_batch.pdf](https://github.com/huggingface/transformers/files/15149408/Llama-2-7b-hf_memory_fix_batch.pdf)
![Llama-2-7b-hf_time_fix_batch.pdf](https://github.com/huggingface/transformers/files/15149410/Llama-2-7b-hf_time_fix_batch.pdf)
![Meta-Llama-3-8B_memory_fix_batch.pdf](https://github.com/huggingface/transformers/files/15149411/Meta-Llama-3-8B_memory_fix_batch.pdf)
![Meta-Llama-3-8B_time_fix_batch.pdf](https://github.com/huggingface/transformers/files/15149412/Meta-Llama-3-8B_time_fix_batch.pdf)

### Fix new token number of 2000 and input length of 300, variable batch size

![Mistral-7B-v0.1_memory_fix_length.pdf](https://github.com/huggingface/transformers/files/15149416/Mistral-7B-v0.1_memory_fix_length.pdf)
![Mistral-7B-v0.1_time_fix_length.pdf](https://github.com/huggingface/transformers/files/15149418/Mistral-7B-v0.1_time_fix_length.pdf)
![Llama-2-7b-hf_memory_fix_length.pdf](https://github.com/huggingface/transformers/files/15149419/Llama-2-7b-hf_memory_fix_length.pdf)
![Llama-2-7b-hf_time_fix_length.pdf](https://github.com/huggingface/transformers/files/15149420/Llama-2-7b-hf_time_fix_length.pdf)
![Meta-Llama-3-8B_memory_fix_length.pdf](https://github.com/huggingface/transformers/files/15149421/Meta-Llama-3-8B_memory_fix_length.pdf)
![Meta-Llama-3-8B_time_fix_length.pdf](https://github.com/huggingface/transformers/files/15149422/Meta-Llama-3-8B_time_fix_length.pdf)

## Integration in Transformers

As I am changing the data structure of the DynamicCache, I had to modify how it is used in models modeling. For now, I only modified `LlamaForCausalLM` and `MistralForCausalLM` which are using the `DynamicCache` by default to test my implementation. Also, the change of data structure may have impacts elsewhere that I overlooked (if other code rely on attributes of the cache, e.g. a call to `cache[0][0].shape` would be an `AttributeError` now). If so, do not hesitate to point me towards these parts and I can modify it.

I would be happy to help integrate the change in all models if you decide to move forward with the PR.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker and @younesbelkada
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @pacman100

Integrations:

- deepspeed: HF Trainer/Accelerate: @pacman100
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc and @younesbelkada

Documentation: @stevhliu and @MKhalusova

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->

@ArthurZucker and @younesbelkada and @gante


Comments:
Hi @ArthurZucker, thanks for the feedback!

Unfortunately, I don't think we can handle all cases by keeping the same data structure. For example, if someone passes the `past_key_values` to `generate()` to resume generation, there is no way to avoid having the reference around as it will be referenced outside "transformer's scope" (e.g. the global scope). In this case, the current data structure will always create a copy and return the copy while still having the old one in memory because we cannot just dynamically add more memory to an already existing tensor. With the new data structure however, we can simply append to the old cache in-place. I just tested it and it works nicely, without allocating any more memory (except for the memory needed for adding a few tokens to the already existing cache of course)!

I checked the rest of the code base and I agree that it is a big change, but it is very much worth it. I'm definitely down to collaborate on it if you wish!

Hi @ArthurZucker,

I am done with the work. Could you please review it? At this point, it should be 100% backward compatible. The only change is that now `past_key_values` will be returned as a `Tuple[Tuple[List[Tensor]]]`, and be modified in-place if directly supplied to `model.generate()` for any of the model architectures that use `DynamicCache` by default (i.e. Cohere, Dbrx, Gemma, Idefics2, Llama, Mistral, Mixtral, Olmo, Persimmon, Phi, Phi3, Qwen2, Qwen2Moe, StableLM, StarCoder2 as of now). 
Every other model will still use the old format `Tuple[Tuple[Tensor]]` (or their own custom format, e.g. for Bloom) but will NOT benefit from the improved memory performance. However, if you make older models also use `DynamicCache` by default, they will automatically benefit from it without further modifications on your part.

Here is a final benchmark of the improved performances for every decoding strategy available in Transformers. It was still done on one of my RTX 4090, and with Mistral-7B-v0.1:

![contrastive_search.pdf](https://github.com/huggingface/transformers/files/15218488/contrastive_search.pdf)
![greedy.pdf](https://github.com/huggingface/transformers/files/15218490/greedy.pdf)
![sample.pdf](https://github.com/huggingface/transformers/files/15218492/sample.pdf)
![beam_search.pdf](https://github.com/huggingface/transformers/files/15218493/beam_search.pdf)
![beam_sample.pdf](https://github.com/huggingface/transformers/files/15218494/beam_sample.pdf)
![constrained_beam_search.pdf](https://github.com/huggingface/transformers/files/15218495/constrained_beam_search.pdf)
![group_beam_search.pdf](https://github.com/huggingface/transformers/files/15218496/group_beam_search.pdf)
![assisted.pdf](https://github.com/huggingface/transformers/files/15218498/assisted.pdf)

Note that for assisted decoding, for 1000 tokens and more I ran into a weird CUDA error (`[/opt/conda/conda-bld/pytorch_1704987289929/work/aten/src/ATen/native/cuda/Indexing.cu:1237]: indexSelectSmallIndex: block: [29,0,0], thread: [84,0,0] Assertion 'srcIndex < srcSelectDimSize' failed.` and `RuntimeError: CUDA error: device-side assert triggered`) in BOTH the actual Transformers version and my modified script, which is why this particular benchmark stops after 500 tokens. I don't know exactly what is happening here, but it is not due to my changes as it is also happening in the current version.

As you can see, for non-beam decoding methods, the memory is divided by 2x as advertised (minus some small overhead, so around 86% improvement in practice for Mistral; between 91% and 98% for Llama2 7B and Llama3 8B). For beam methods, as I still rely on `index_select()` to reorder the cache which performs a copy, it is less than that, but still very much significant (I measured it as an improvement of 46%, so memory is divided by 1.5).
The speed impact is only noticeable when the number of new tokens generated in one go becomes very large (more than 4000), which in practice should be very rare anyway.

Finally, however, as the input size becomes large, the memory bottleneck may become the first model forward pass instead of the cache size. In settings where the input size is large compared to the number of new tokens generated (say e.g. input size 4000 and new tokens 5), the memory footprint will be dominated by the cost of the first model forward pass instead of the cache size. 
In this case, the memory improvement will still be **at least** 25% for models with small caches such as Mistral, and **usually** around 60-70% for models with bigger caches such as Llama2 from my experiments, because we still save a full copy of the cache just after the first forward.
This is illustrated in the following 2 figures (here the input size is already 1000):

![scaling_example_beam.pdf](https://github.com/huggingface/transformers/files/15223549/scaling_example_beam.pdf)
![scaling_example_non_beam.pdf](https://github.com/huggingface/transformers/files/15223551/scaling_example_non_beam.pdf)

In those figures (the first is for beam-based methods, the second for non beam-based), we can observe "3 zones". First, as the number of new tokens generated is very small compared to the input size, we have a plateau of "minimal memory improvement", then as the number of new tokens increases a sharp increase in memory improvement, until the size of the cache is similar to the memory footprint of the first forward pass. At this point, we then plateau to the zone of « maximum memory improvement », which is roughly 2 times more efficient for all non beam-based methods, and roughly 1.5 more efficient for all beam-based decoding strategies.

I checked and we ALWAYS hit the zone of minimal improved performance, which means in this case (see figure) that ANY call to non beam-based method will yield at least a memory improvement of ~1.22 for Mistral and ~1.62 for Llama2 (even generating only 2 new tokens for any input size), with still even better improved efficiency (up to between 1.85 and 2) for a large range of input size and max new tokens. The same is true for beam-based methods, with minimal improvement ~1.3 for Mistral and more for Llama2 (up to ~1.5 max for a large range of inputs).

I hope this is clear and can be integrated as soon as possible! Don't hesitate to tell me if I am missing something that should be done before merging.

Hi @ArthurZucker,

I was investigating why we observe those "minimal improvements" even with very large input sizes and just 2 new tokens. I found out that the reason was that the logits in every iteration leaked to the next iteration though some hanging references to the outputs dictionary. This was true for every decoding strategy except assisted decoding. **I fixed it in my last 2 commits.**

## Non-beam methods

That means that for any input size, the memory peak was not caused by the forward pass of the first iteration with the large input, but during the 2nd iteration, when the copying of the full cache was happening. At that point, the memory peak consisted of `2*past_KV + first_iteration_logits` (the logits in the first iteration when all tokens are feeded to the model can quickly grow very large). After my first modifications, we were no longer keeping the 2nd cache copy, thus in the second iteration, the peak would be `past_KV + first_iteration_logits`, which is always at least smaller or equal to the memory needed for the forward in the first iteration (this pass is creating those values so need at least that much memory, in practice more). Thus the leaking of logits was not a memory bottleneck after solving the copying of the cache in this case. I still solved it though, because it does not hurt performances and would be weird to keep. In beam methods however, removing it is a big improvement (see below).
Anyway, that means that we can compute the "minimal improvement" as:

```
minimal_improvement = (2*past_KV + first_iteration_logits) / memory(first_forward_pass)
```

Now, both the cache and the logits trivially scale linearly with the input size. Since flash attention, the memory needed for the forward pass also scales linearly with the input size, thus `minimal_improvement` is a constant for all input sizes. This is exactly what I observed in my past comment, with this constant being ~1.22 for Mistral and ~1.62 for Llama2 (so we always divide the memory requirements by at least those numbers). I computed the ratio presented above and obtained the same values, which confirmed my analysis. 

Now, as we generate more tokens after the first iteration, and the cache size exceeds the `memory(first_forward_pass)`, we tend to the `maximal_improvement = 2`, because we removed the perpetual copy of the cache (logits are negligible here because after the first pass, the logits will always be for only 1 token, thus small in size). 

Illustration:
![non_beam-based.pdf](https://github.com/huggingface/transformers/files/15241727/non_beam-based.pdf)


## Beam methods

For beam methods, the story is slightly different. Because reordering the cache allocates a full copy of the cache, the leak of the cache and the logits caused `3*past_KV + first_iteration_logits` to be allocated in memory during the 2nd iteration. After my work, it is now only `2*past_KV` (I believe we could even reorder in a smarter way that would avoid that 2nd copy, but that would be future work). That means, as before, the `minimal_improvement` is:

```
minimal_improvement = (3*past_KV + first_iteration_logits) / max{memory(first_forward_pass), 2*past_KV}
```

and when generating enough new tokens, the `maximal_improvement` is going from `3*past_KV` to `2*past_KV`, thus dividing the memory requirement by 1.5. However, in this case, `minimal_improvement` is always bigger than 1.5, meaning that if we generate few tokens compared to the size of the input we will benefit from memory improvement bigger than 1.5, and if we generate enough new tokens compared to the input size, the improvement will tend to 1.5.

Illustration:
![beam-based.pdf](https://github.com/huggingface/transformers/files/15241729/beam-based.pdf)

## Contrastive search

Contrastive search is a bit of a special case. As a non-beam strategy, it will benefit from 2x memory savings but will not scale in the same way. Because we artificially multiply the size of the cache by `top_k` in each iteration to forward with artificially increased batch size, if we neglect the memory needed for the initial forward pass (which is of the input size, not with batch size increased), we will immediately benefit from a memory improvement of `(2 * top_k + 1) / (top_k + 1)` which is ~2 very quickly. So, independently of the number of new tokens, we will get 2x increased performance. 

Illustration:
![contrastive_top_k.pdf](https://github.com/huggingface/transformers/files/15247303/contrastive_top_k.pdf)

As you see on the image, as `top_k` increases, we quickly tend to 2x for any new tokens number. Note that this was obtained with Mistral which has a small cache (thus cost of initial forward is not completely negligible), which means that other models such as Llama2 would get to 2x much faster (for smaller `top_k`) but I couldn't really show it as I would immediately OOM with the old version.


## Restarting from previous cache

Now, because greedy and sample search never perform any copies of the cache anymore, restarting from previous cache such as: 
```
out = model.generate(inputs, max_new_tokens=100, return_dict_in_generate=True, do_sample=False)
cache = out.past_key_values

# Do whatever here

outputs = model.generate(out.sequences, past_key_values=cache, max_new_tokens=100, do_sample=False)
```

will actually benefit from a **3x** memory improvement (that is, we divide the memory footprint by 3x). This is because before we would always have 2 copies of the cache internally during the second call to `generate` + the copy outside (variable `cache`). Now, no copies are performed and `cache` is modified in-place (we just append to it). This results in a 3x memory improvement. If the number of new tokens in the second call to `generate` becomes very large, then the size of the `cache` variable becomes somewhat negligible, and we are back to our 2x memory savings. So this setting will always yield between 2x and 3x savings.

As most chat applications will actually always use greedy or sample decoding, and will restart every new conversation turn from the previous cache (at least they should most of the time), it means that these chat applications will always save between 2x and 3x memory.

# TL;DR

This PR significantly reduces the memory footprint of `generate()` for all inputs combinations and models.
For decoding strategies that do not rely on beam (including what I think are the most 2 populars, greedy and sample), most models will benefit of at least ~1.6x memory reduction independently of the input size. Those with smaller caches such as Mistral will benefit from roughly at least ~1.2x memory decrease for all input sizes. Moreover, as the `max_new_tokens` generated increases, the previous numbers quickly go up to 2x improvement. In most practical cases, we will actually hit the 2x increase zone (that is dividing the necessary memory by 2x).
For decoding strategies relying on beams, the improvement is always AT LEAST 1.5x with even better performance for small `max_new_tokens`.
Restarting generation from previous cache for greedy and sample decoding will ALWAYS (independently of all input sizes and number of new tokens generated) yield memory savings between 2x and 3x.
This comes at no cost ("free lunch") as the speed penalty is only noticeable when generating more than 6k or 7k new tokens in one go (around 15%-20% performance penalty in the worst case when generating 7k tokens in one-go, which may even still be acceptable), which in practice should (except in extremely niche cases) never happen. Moreover, as the batch size can be largely increased due to the memory savings, the actual processing speed may be even faster.
Nowadays, as the model sizes increase, the memory is (almost) always the bottleneck in all applications. This work tries to alleviate the pressure on those GPUs as much as possible.

> @Cyrilvallez kudos for this really high quality contribution 🤗 🚀 Your in-depth explanations are very useful, and all of it makes sense. I think at some point we did want to change the way the cache was stored, to also remove the two transpose operation that are always needed in the modeling code (which is not the case for jax codes).
> 
> @gante some potential break is the cache format, but I think we can add an `EfficientDynamicCache` class that would be used upon activation. @Cyrilvallez this would be for use better than having two code paths and makes more sense for isolating bugs / maintenance!
> 
> PS: do you have any gist on how to reproduce memory benchmarks and etc?
> 
> WDYT @gante

Thanks for reviewing @ArthurZucker! Here is a link to the benchmarks I ran: https://gist.github.com/Cyrilvallez/ce1adfad1d561c1e8dc92666ab5a9e8c

It is a bit messy but you should find everything. Each benchmark was run on the official version 4.40.1, saving outputs with `legacy` in the filenames, and then on my current branch using `patched` in the filenames. You should manually change the filenames depending on which version you use if you re-run the benchmarks. Also, you may need to modify the actual paths to those filenames when creating the figures. Apart from that, I think the link should contain everything you need to reproduce. 

@Cyrilvallez what do you think about creating a new cache class? 

@ArthurZucker I agree that a new class would be clearer and more maintainable. Not sure what you meant by "that would be used upon activation" however? I would use the new EfficientDynamicCache as the new default for all models actually using DynamicCache by default if not provided. As I cannot think of any backward compatibility issues (even if some people save cache to file to restart from it, the new class will work), doing otherwise would be counter-productive I think (it would be weird to have to pass an argument to activate a feature that does the same as before, but more efficiently in all cases). I don't have your experience maintaining such a widely used codebase however, so if you think people may run into problems, we could raise a warning to point them to how to use the old DynamicCache if they run into any issue?

By default we don't change anything, and just passing "cache_implementaiton="efficient"` would use your implementation! 
We can warn about "you are using the non efficient version of the DynamicCache` for example 🤗 

Ok, then by default we will still benefit from removing the leak of the logits which is already a big gain. I will make the necessary changes next Monday 👌🏻

@Cyrilvallez very cool in-depth exploration 😮🔥 And also very impactful consequences of the suggested changes!

Reading the discussion, from a usage perspective, I agree with Arthur: a separate class would be ideal! With a separate cache class, we:
1. Can keep a simple reference cache class in `DynamicCache`;
2. The default class used and returned in `forward` stays as simple as possible;
3. [after a proper deprecation cycle] Can use the new efficient class by default in `generate`, getting all the benefits showcased in the benchmarks above.

I'll review the full code after it is implemented as a separate class, as we all seem aligned on what should be done 🤗 @Cyrilvallez ping me when it's ready :D 



@Cyrilvallez out of curiosity: have you considered/explored expanding the cache with fixed-size blocks every time we hit the limit, similarly to paged attention? 

@ArthurZucker @gante I realized yesterday that what actually creates the copies is not the current `DynamicCache` itself, but the back and forth `from_legacy_cache` and `to_legacy_cache` calls (that creates tuples that are immutables). Thus, forcing to use proper `DynamicCache` in `generate()` before going into specialized decoding strategy would retain all the above benefits while removing the need for a new class. Something like:

```
if use_cache and self._supports_cache_class:
    if generation_config.cache_implementation is None and not isinstance(model_kwargs.get('past_key_values', None), Cache):
        model_kwargs['past_key_values'] = DynamicCache() # or DynamicCache.from_legacy_cache() depending on model_kwargs['past_key_values'] type
```

Doing so will avoid the back and forth switch from legacy to cache, and thus all references in the decoding strategies functions will point to the same `Cache` object that will be modified in-place for all of them because the caches are List[Tensor]. 

This would greatly simplify the above approach while retaining all benefits (even more as `reorder` would also happen in-place layer after layer, thus not allocating new memory -> beam decoding will use 3x less memory instead of 1.5). The small speed decrease would also be removed. 

@ArthurZucker @gante The work is ready for final review! 

As previously said, `EfficientDynamicCache` was not needed in the end. This makes the changes more natural based on the current state of the library, and will be much easier to maintain. I rebased my branch to remove previous commits and make everything clearer. Models with support for `DynamicCache` will use it from start to end as of now, without back and forth to the old legacy format. However, `generate()` will return the legacy format by default for backward compatibility. This can be controlled with the generation config argument `return_legacy_cache` if this is not wanted (e.g. in chat applications, when restarting from previous cache, it is more efficient to always keep a `DynamicCache` to avoid any copying and get between 3X and 2X memory improvement even for sample and greedy decoding as previously mentioned).

Models without support for `DynamicCache` will still benefit from improvements as well as I removed all the references to leaked logits, and, as much as possible, leaked copies of past_key_values. This is illustrated in the following for greedy search:
![old_models.pdf](https://github.com/huggingface/transformers/files/15332941/old_models.pdf)
As you can see, models with large logits (i.e. large vocabulary size) such as bloom benefit a lot from the removed leaked logits as long as the number of new tokens generated is not too big. This is because we divide memory consumption by the factor `(2*cache + logits_first_iteration) / max{memory_first_forward, 2*cache}` and that usually for large logits the denominator is equal to `memory_first_forward`. For smaller logits, the denominator will be `2*cache` and we "only" save the full logits size.

Moreover, this work paves the way to the full adoption of the new `Cache` format everywhere as all generate functions can now accomodate a `DynamicCache` instance all the way (some functions and/or some keywords were not completely supported before).

Finally, here is the final benchmark of performances:
![beam-based.pdf](https://github.com/huggingface/transformers/files/15324681/beam-based.pdf)
![non_beam-based.pdf](https://github.com/huggingface/transformers/files/15324683/non_beam-based.pdf)
![contrastive.pdf](https://github.com/huggingface/transformers/files/15324684/contrastive.pdf)
![contrastive_top_k.pdf](https://github.com/huggingface/transformers/files/15324685/contrastive_top_k.pdf)


Beam-based methods now benefit from using 3X less memory because even `reorder()` is made in-place. Non beam-based are the same as before (except contrastive search). Contrastive always benefits from at least 2X less memory now. It will be even better for small `top_k`, and then scale to 2X as `top_k` increases, because I changed how the cache is expanded (without temporary copy, thus we only have `top_k` cache replications instead of `top_k +1` which is a nice gain for small `top_k` compared to previous benchmarks -> memory gain factor is `(2*top_k + 1) / top_k` now).

Finally, as we are using `DynamicCache`, we ne longer have any speed degradation. All those memory improvements are FREE LUNCH all the way.


New idea to further improve memory: #30860

Damn that's impressive! Reviewing now! 


@ArthurZucker @gante I applied all changes following your comments! 

Repo consistency and code quality errors do not come from files I modified (I think code quality errors come from ruff version now being 0.4.4, previously 0.1.x).
Same goes for tests_tf, I did not touch any tf files either so it must come from the fact that I just rebased on latest main branch.

@Cyrilvallez I believe CI is red due to external causes, [this PR](https://github.com/huggingface/transformers/pull/30932) should fix it after it gets merged

@ArthurZucker can you give this PR a final check? :)

Ran a few tests locally (slow tests for popular models, slow `generate` tests, slow cache tests), everything is passing 👍 

Having a look! 

I wanted to run the slow tests but did not work 😓 


@ArthurZucker Should I try to rebase once again to see if it solves the issue?

@Cyrilvallez please yes, rebasing should work !

Yep, according to @younes this should not work. I'll review and push another empty commit 

Hi @zucchini-nlp! When rebasing I noticed that in your recent https://github.com/huggingface/transformers/pull/30483 you made QuantizedCache a subclass of DynamicCache. However, some code paths in `generate()` rely on checking `isinstance(cache, DynamicCache)` (this is the case in e.g. `_contrastive_search`). This is likely to crash if a `QuantizedCache` instance go inside one of those paths. 

Before, I had added a few methods (`crop`, `batch_split`, `from_batch_split`) in `DynanicCache` to make the code go through cache methods instead of specialized snippets elsewhere. After I saw your PR, I added `batch_repeat_interleave` and `batch_select_indices` to use in `_contrastive_search` instead of hard-coding what should happen inside this function if the cache is a `DynamicCache`. At this point, I think that there are no more code paths left that hard-code a specific behavior that do not use a `DynamicCache` method.

This means that if `QuantizedCache` overwrites these 5 methods (I think it is likely to be needed, but did not look at `QuantizedCache` in details maybe default implementation in `DynamicCache` is enough?) , everything should work correctly in all current code paths checking `isinstance(cache, DynamicCache)` in `generate`. As of now, I think that at least contrastive search is very likely to fail in e.g. this part https://github.com/huggingface/transformers/blob/main/src/transformers/generation/utils.py#L2146

If this PR gets merged, overwriting the 5 methods mentioned above should make `QuantizedCache` work in any scenario and simplify `generate` code paths for the future.

@ArthurZucker Rebasing is done and all CIs are green!

@Cyrilvallez Right, QuantizedCache stores most of the past kv in a private list, so I think these methods would not work even before your changes. 

Thanks for noticing! I will think if we want to make quantized cache compatible with these generation strategies, otherwise raise an error when user tries to go that path. 

And amazing work in this PR for making generation better! 🔥 

Yes, it will fail in the current state of the library, but adding support will be straight-forward and won't require messing with 'generation.utils' after this is merged!

Reviewing today sorry for the delay

Hi @ArthurZucker, any news? 😁

Super excited to see this get merged, to hopefully resolve https://github.com/huggingface/transformers/issues/30019

Thank you all for your hard work on this  🚀 

@ArthurZucker The modeling update was mostly used in the case `inputs_embeds` were passed, but checking `past_length == 0` is equivalent and cleaner so I updated. I also removed the `TODO: remove if-else check after generate prioritize Cache objects` as this is now the case.

Should be good now!

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_30536). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@ArthurZucker With how things evolved in the end, the docstring for `DynamicCache` is still accurate! 

I am more thinking about explaining that we do internally some bucketing
but if you feel like it's good as is, we can merge

The bucketing part was removed in the end because it added unnecessary complexity to the code, and we could get all memory benefits without it! So I don't think that the docstring needs any more details (the `update()` method still work the same as before)

Okay! Merging then! Congrats and great work 🤗 🚀