[WIP] Sink cache: fix implementation to shift key states
Author: why-in-Shanghaitech
Pull request: yes

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes #33691

This PR shifts the new key states along with those already in the cache. It could also solve some issues mentioned in #31537, #32315, where the initial input key states are longer than the context window or shorter than the attn sink.

I hope it could better explain the issue in #33691.

It does not solve the issue that the new key states have large positional encodings (See #33691 the last few lines). So it still does not work with `model.generate`. Once this is fixed, the model will be able to generate tokens very well.

Also, it will be better if we can come up with some good tests. E.g. re-design key state so that it takes the actual positional encoding as its attributes in the test?

## 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
- vision models: @amyeroberts, @qubvel
- speech models: @ylacombe, @eustlb
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Rocketknight1
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @SunMarc
- chat templates: @Rocketknight1

Integrations:

- deepspeed: HF Trainer/Accelerate: @muellerzr
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc

Documentation: @stevhliu

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 @zucchini-nlp @gante
I didn't finish up reading all the previous discussions (it's too long)... But I know you should be very familiar about this.

Comments:
If we make some ad-hoc changes:
https://github.com/huggingface/transformers/blob/52daf4ec768fb9ffe84a0c373834172a7c54aecc/src/transformers/models/llama/modeling_llama.py#L962
https://github.com/huggingface/transformers/blob/52daf4ec768fb9ffe84a0c373834172a7c54aecc/src/transformers/models/llama/modeling_llama.py#L967

Make this two lines always true. Then it can output the almost same texts as in [the original streamingllm implementation](https://github.com/mit-han-lab/streaming-llm/tree/main).

It might be related to the changes in #32315, but this PR will report errors from my end.

https://github.com/huggingface/transformers/pull/33692/commits/45c70a70c1588ae839a097088d8b36e4a688d859 This commit tries to solve the problem above. I am not sure whether this will affect codes in other places...

> @why-in-Shanghaitech yes, SinkCache currently is very heavily bugged and the tests you mentioned are failing. We know when it started bugging, when we modified attention mask preparation stages. Prev the rerotations worked because attention mask (and consequently position ids) never went beyond the `window_length`
> 
> Currently my idea was to bring back the cropping as keeping cos/sin until the maximum possible sequence length seemed a bit inefficient in terms of memory. But I see it is also a possible solution. Since I paused work on SinkCache for a few weeks for focus on other tasks, tagging @gante
> 
> @gante will say what is better for sink cache, he is the point of contact for anything in generation

Thank you for the reply!

Yes... I think cropping is necessary since the attention mask grows w/o looking at the past key values.

https://github.com/huggingface/transformers/blob/f2c388e3f946862f657acc1e21b272ec946fc66c/src/transformers/generation/utils.py#L645-L649

Maybe we can crop here instead? The past key values are known in `model_kwargs`. Or maybe we should process the attention mask in the Cache classes?

Also, since we cannot modify the positional encoding of the queries, it seems impossible to keep the cos/sin within the window length. And I think we don't need to do this? This PR only caches cos/sin with window length + the longest prompt length (not the entire generated sequence length).

https://github.com/why-in-Shanghaitech/transformers/blob/1550f8f126946570c748131552c6b97fd33ba874/src/transformers/cache_utils.py#L969-L971


Is this stale?

AFAIK this is not needed anymore, @gante was going to remove SinkCache so it won;t be supported anymore

> AFAIK this is not needed anymore, @gante was going to remove SinkCache so it won;t be supported anymore

ok thanks, I'll close it.