Gemma 2 crashes with CUDA error on long context input (again)
Author: akkiss
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.48.0
- Platform: Linux-4.18.0-553.6.1.el8.x86_64-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.27.1
- Safetensors version: 0.4.5
- Accelerate version: 1.2.1
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.5.0a0+e000cf0ad9.nv24.10 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: Tesla V100-SXM2-32GB


### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

When I try to generate text with long prompt using gemma-2-2b-it, it crashes with CUDA assertion error. Short code to reproduce the error:
```
import os
os.environ["CUDA_LAUNCH_BLOCKING"]="1"

import torch
from transformers import pipeline

model_name = "google/gemma-2-2b-it"
pipe = pipeline("text-generation", model=model_name, device="cuda", torch_dtype=torch.float16)

messages = [
	{"role": "user", "content": 1000*"Tell me something interesting. "}
]
prompt = pipe.tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
pipe(prompt, max_new_tokens=100, pad_token_id=pipe.tokenizer.eos_token_id, eos_token_id=[pipe.tokenizer.eos_token_id], do_sample=True)
```
This prints assertion failures:

> /opt/pytorch/pytorch/aten/src/ATen/native/cuda/IndexKernel.cu:93: operator(): block: [8310,0,0], thread: [0,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
>/opt/pytorch/pytorch/aten/src/ATen/native/cuda/IndexKernel.cu:93: operator(): block: [8310,0,0], thread: [1,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
>...
>  File /usr/local/lib/python3.10/dist-packages/transformers/models/gemma2/modeling_gemma2.py:224 in forward
>    key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
>
>  File /usr/local/lib/python3.10/dist-packages/transformers/cache_utils.py:1721 in update
>    return update_fn(
>
>  File /usr/local/lib/python3.10/dist-packages/transformers/cache_utils.py:1694 in _static_update
>    k_out[:, :, cache_position] = key_states
>
>RuntimeError: CUDA error: device-side assert triggered
>Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

The problem seems to be that in [Gemma2Attention.forward](https://github.com/huggingface/transformers/blob/12dfd99007dd803b3dfab5fe2ec66260b8007a23/src/transformers/models/gemma2/modeling_gemma2.py#L223) when [HybridCache.update](https://github.com/huggingface/transformers/blob/12dfd99007dd803b3dfab5fe2ec66260b8007a23/src/transformers/cache_utils.py#L1709) is called, the "sliding_window" parameter is not forwarded in the cache_kwargs argument and therefore the _static_update method is called instead of the _sliding update.

This issue seem similar to #31848 but it suggest extracting sliding_window from cache_kwargs as currently done but the caller does not add this value to cache_kwargs.

Also, I thought this model has at least 8k context window though no clear indication on model card, and yet sliding_window size is 4096.


### Expected behavior

If I replace
```
            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
```
with
```
            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position, "sliding_window": self.sliding_window}
```
text generation works.
If someone familiar with the codebase says it's a legit fix, I'm happy to create a pull request but don't know the codebase well enough. @ArthurZucker, do you think it is a good fix?

Comments:
cc @Cyrilvallez who's working on fixing generation with hybrid cache! But your fix is most probably legit! 🤗 

Indeed @akkiss your fix is legit, but it's not the only issue 🥲 see https://github.com/huggingface/transformers/pull/35681 for details!

Great, I see #35681 contains this modification, can't wait for it to be merged 😉