XLA optimized Implementation of StaticCache with `Tensor Indexing API`
Author: huzama
Pull request: yes

Use the `Index_copy` method to update the static cache inplace and avoid recompilation during each iteration in XLA

# What does this PR do?
The PR focuses on avoiding repeated recompilation during each iteration in XLA by performing in-place updates on the static cache. Specifically, it replaces direct tensor indexing assignments with `index_copy_` method calls. This technique ensures that cache updates are executed without triggering full tensor recompilations, improving runtime performance.

### Code Changes

**Original Code:**
```python
k_out[:, :, cache_position] = key_states
v_out[:, :, cache_position] = value_states
```

**Updated Code:**
```python
k_out.index_copy_(2, cache_position, key_states)
v_out.index_copy_(2, cache_position, value_states)
```




## Before submitting
- [X] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?

 @zucchini-nlp @gante 



Edit: 

- #### Discussion: [31126](https://github.com/huggingface/transformers/issues/31126)
- The code was further updated to create a new list instead of performing in-place updates to avoid recompilation in XLA.




Comments:
cc @ArthurZucker as well

### Summary of Changes

@tengomucho, I conducted experiments with both `.index_copy_` and Python slicing methods on the GPU and found no significant performance differences. Consequently, I have reverted the `StaticCache` code to its original form.

> In LLaMA, the KV-cache tensor slices are updated in-place; this leads to recompilation events every time a token is generated. To address this issue, we use index tensors and tensor.index_copy() ops to replace the in-place slice updates. Attention masks and output sequences also benefit from the same optimization. [[Ref]
](https://pytorch.org/blog/path-achieve-low-inference-latency/)

Based on @ArthurZucker's suggestions, and above-mentioned guide, I created a new class, `StaticCacheXLA`, with several key differences from the simple `StaticCache` class:

1. **Updating Cache with Out-of-Place Operations with `index_copy`**
   _XLA lazy tensors perform well with out-of-place operations:_
   https://github.com/huggingface/transformers/blob/71efd4c1f7a2f3411d19aff4a7f811f331c7330a/src/transformers/cache_utils.py#L1007-L1014

2. **Get `seq_len` Out-of-Place with `index_select`**
   https://github.com/huggingface/transformers/blob/71efd4c1f7a2f3411d19aff4a7f811f331c7330a/src/transformers/cache_utils.py#L1027-L1030

### Performance Improvements

#### Architecture Design

Using SPMD on LLAMA 3 following the guide from [PyTorch's high-performance LLAMA 2 blog](https://pytorch.org/blog/high-performance-llama-2/) and a simple generate function with greedy decoding, I achieved the following results:

#### Results

These experiments were run on a TPU v3-8. The input sequence length was 256, and the maximum number of new tokens was 512. I can upload numbers on TPU v4-128 later.

As a result of these changes and optimizations, the new `StaticCacheXLA` class achieved a generation rate of 6.35 iterations per second with TPU utilization at 31%. In comparison, the original `StaticCache` achieved 5.5 iterations per second with TPU utilization at 27%.

- **StaticCacheXLA**: 6.35 it/s, TPU utilization: 3.1%
- **StaticCache**: 5.5 it/s, TPU utilization: 2.7%
- **Without Cache**: 3 s/it, TPU utilization: 79%

Edit:
The significant disparity in TPU utilization arises from the difference between single token generation during cache-enabled runs and the full input processing during non-cache runs. Despite the low TPU utilization, cache-enabled generation is 10 times faster. Additionally, when employing in-place cache updates or Python slicing, the model recompiles the graph in proportion to the number of layers it contains. A similar pattern is observed with XLA on GPUs. I can provide further details, including the relevant code and performance metrics if needed.

@ArthurZucker Thank you for the feedback. If I understand correctly, you are suggesting that we should dynamically map to the appropriate cache implementation based on whether XLA is available. Would the following update reflect your suggestion?

```python
class StaticCacheDefault(Cache):
    """Default static cache implementation."""
    pass

class StaticCacheXLA(Cache):
    """XLA-optimized static cache implementation."""
    pass

# Determine which StaticCache implementation to use based on the availability of XLA.
StaticCache = StaticCacheXLA if is_torch_xla_available() else StaticCacheDefault
```

Hello @gante,

A quick test on GPUs (as my TPUs are running jobs) with the following script revealed it can be used and it produces the same results as the original `StaticCache`.

```python
from transformers import LlamaForCausalLM, AutoTokenizer
from transformers.cache_utils import StaticCacheXLA, StaticCache
import torch

model = LlamaForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B", torch_dtype=torch.bfloat16
).to(1)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B")

xla_cache = StaticCacheXLA(model.config, 1, 128, dtype=torch.bfloat16, device=1)
cache = StaticCache(model.config, 1, 128, dtype=torch.bfloat16, device=1)

input_ids = tokenizer("This is a test", return_tensors="pt").to(1)

out = model.generate(
    **input_ids, max_new_tokens=32, do_sample=False, past_key_values=cache
)

out_cache = model.generate(
    **input_ids, max_new_tokens=32, do_sample=False, past_key_values=xla_cache
)

torch.all(out == out_cache) 
# True
```


@huzama awesome! Can we add it as a test? 💛 After the test is added, I'm more than happy to approve the PR!

@gante I even wonder if it would be possible just to patch the original StaticCache class, if they behave the same? 

@tengomucho if all that's needed is to replace slicing by `.index_copy_`, I'm happy with changing `StaticCache` itself -- as long as we add a comment like "`k_out.index_copy_(2, cache_position, key_states)` is equivalent to `k_out[:, :, cache_position] = key_states`, but with better generalized support", for readability

If that's the case, no new tests would be needed 🤗 @huzama 

@gante I agree that combining both classes into a single implementation is a better approach. I'll add comments to explain the code and will push the changes with the merged implementation.

@huzama you'll need to run `make fixup` on your `transformers` root dir and push the changes to make our CI happy 🤗 

(unrelated CI failure also observed on fresh PRs from `main`)

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

Hey @huzama are you planning to clean this up to get it merged? Otherwise let me know, I will be happy to take over to make `StaticCache` more efficient on XLA!

@tengomucho Thank you for taking over. I have been running some other experiments, and I lost track of this pull request. I appreciate your help in making StaticCache more efficient on XLA!

I opened another PR so I could push on the branch. Closing this!