feat(cache): StaticCache uses index_copy_ to avoid useless copy
Author: tengomucho
Pull request: yes
Labels: name: run-slow

# What does this PR do?

Using index_copy_ allows for explicit in-place change of the tensor. Some backends (XLA) will otherwise copy the tensor, making the code slower and using more memory.

Proposed implementation will end up using less memory and on XLA will result in less compilation, but the change is also quite generic, making no change whatsoever on CUDA or CPU backend.

This was discussed in https://github.com/huggingface/transformers/pull/31129 that I squashed/rebased. 

## Who can review?


Models:

- text models: @ArthurZucker @gante

Library:

- generate: @gante (all others)


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

cc @huzama

@ArthurZucker I applied the change to sliding window too. 
What is the easiest way to trigger slow test run on Llama?

Hi @tengomucho,

Thank you for taking over. I read the code for the sliding window, and I think we can optimize the following lines of code with `index_select` operations as well:

https://github.com/huggingface/transformers/blob/53e99d185cbfaae32788281a7214acfd82403277/src/transformers/cache_utils.py#L955-L956

https://github.com/huggingface/transformers/blob/53e99d185cbfaae32788281a7214acfd82403277/src/transformers/cache_utils.py#L969-L970


Additionally, IMO 

https://github.com/huggingface/transformers/blob/53e99d185cbfaae32788281a7214acfd82403277/src/transformers/cache_utils.py#L975-L983

This can be replaced with `self.key_cache[layer_idx].index_copy_` 

you should do a commit with `[run-slow] llama` inside to trigger them! 

@huzama thanks for the suggestion, though while the very first change allowed to avoid a tensor copy and improve performance on XLA, I am a little concerned about the other suggestions, because:
- Using `index_select` will result in a tensor copy anyway, it's just a different syntax
- For the latest one, it might be good, but I have some concerns about touching that part of code, given the comment about the CUDA graph. I do not have the time now to verify if that works better or not on CUDA, so I will leave it like that
- I had to add a try/catch to allow the static and sliding window cache to keep on working on MPS!

`git commit -m "[run slow] llama"` should have triggered it

Umh I see an issue in the slow test `LlamaIntegrationTest.test_compile_static_cache`, the message says:
```python
>               k_out.index_copy_(2, cache_position, key_states)
E               RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument for argument index in method wrapper_CUDA_index_copy_)
```

I have different options here:
- add a condition to check if `index_copy_` inputs are on same device, otherwise use the legacy `tensor[:, :, index] = tensor` syntax. Or a try-catch.
- move the `index_copy_` inputs to the tensor device, but I am not sure of side effects of that. I wonder is the legacy syntax does that implicitly.
- Give up on this change, as it affects multi-gpus (but it will penalize XLA).

Any suggestion?

Can you try moving the cache position to the correct device? (IMo this is what’s bound to be moved around)

Do you guys think we can merge this now?

Yeah if you can fix the ci!