device mismatch error when using `SlidingWindowLayer`.
Author: nuxlear
Pull request: no
Labels: name: bug

name: Cache

### System Info

- `transformers` version: 4.54.0
- Platform: Linux-5.15.0-1083-gcp-x86_64-with-glibc2.31
- Python version: 3.12.10
- Huggingface_hub version: 0.34.1
- Safetensors version: 0.5.3
- Accelerate version: 1.6.0
- Accelerate config:    not found
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.7.1+cu126 (CUDA)
- 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?: yes (`device_map="auto"` with accelerate)
- Using GPU in script?: yes
- GPU type: NVIDIA A100-SXM4-40GB

### Who can help?

@ArthurZucker 

### Information

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

### Tasks

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

### Reproduction

1. Run quickstart example on EXAONE-4.0 using multiple GPUs (we use 2 x A100-40GB)
  ```python
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "LGAI-EXAONE/EXAONE-4.0-32B"

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="bfloat16",
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# choose your prompt
prompt = "Explain how wonderful you are"
prompt = "Explica lo increíble que eres"
prompt = "너가 얼마나 대단한지 설명해 봐"

messages = [
    {"role": "user", "content": prompt}
]
input_ids = tokenizer.apply_chat_template(
    messages,
    tokenize=True,
    add_generation_prompt=True,
    return_tensors="pt"
)

output = model.generate(
    input_ids.to(model.device),
    max_new_tokens=128,
    do_sample=False,
)
print(tokenizer.decode(output[0]))
  ```
2. It raises an error on indexing KV cache on `SlidingWindowLayer`.
```
Traceback (most recent call last):
  File "/home/junwon_hwang/transformers_release/quickstart.py", line 27, in <module>
    output = model.generate(
             ^^^^^^^^^^^^^^^
  File "/home/junwon_hwang/uv_env/exaone4/.venv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/junwon_hwang/uv_env/exaone4/.venv/lib/python3.12/site-packages/transformers/generation/utils.py", line 2633, in generate
    result = self._sample(
             ^^^^^^^^^^^^^
  File "/home/junwon_hwang/uv_env/exaone4/.venv/lib/python3.12/site-packages/transformers/generation/utils.py", line 3614, in _sample
    outputs = self(**model_inputs, return_dict=True)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/junwon_hwang/uv_env/exaone4/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/junwon_hwang/uv_env/exaone4/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/junwon_hwang/uv_env/exaone4/.venv/lib/python3.12/site-packages/accelerate/hooks.py", line 176, in new_forward
    output = module._old_forward(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/junwon_hwang/uv_env/exaone4/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 961, in wrapper
    output = func(self, *args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/junwon_hwang/uv_env/exaone4/.venv/lib/python3.12/site-packages/transformers/models/exaone4/modeling_exaone4.py", line 491, in forward
    outputs: BaseModelOutputWithPast = self.model(
                                       ^^^^^^^^^^^
  File "/home/junwon_hwang/uv_env/exaone4/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/junwon_hwang/uv_env/exaone4/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/junwon_hwang/uv_env/exaone4/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 1069, in wrapper
    outputs = func(self, *args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/junwon_hwang/uv_env/exaone4/.venv/lib/python3.12/site-packages/transformers/models/exaone4/modeling_exaone4.py", line 404, in forward
    hidden_states = decoder_layer(
                    ^^^^^^^^^^^^^^
  File "/home/junwon_hwang/uv_env/exaone4/.venv/lib/python3.12/site-packages/transformers/modeling_layers.py", line 94, in __call__
    return super().__call__(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/junwon_hwang/uv_env/exaone4/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/junwon_hwang/uv_env/exaone4/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/junwon_hwang/uv_env/exaone4/.venv/lib/python3.12/site-packages/accelerate/hooks.py", line 176, in new_forward
    output = module._old_forward(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/junwon_hwang/uv_env/exaone4/.venv/lib/python3.12/site-packages/transformers/models/exaone4/modeling_exaone4.py", line 291, in forward
    hidden_states, _ = self.self_attn(
                       ^^^^^^^^^^^^^^^
  File "/home/junwon_hwang/uv_env/exaone4/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/junwon_hwang/uv_env/exaone4/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/junwon_hwang/uv_env/exaone4/.venv/lib/python3.12/site-packages/accelerate/hooks.py", line 176, in new_forward
    output = module._old_forward(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/junwon_hwang/uv_env/exaone4/.venv/lib/python3.12/site-packages/transformers/models/exaone4/modeling_exaone4.py", line 230, in forward
    key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/junwon_hwang/uv_env/exaone4/.venv/lib/python3.12/site-packages/transformers/cache_utils.py", line 967, in _wrapped_update
    key_tensors, value_tensors = fn(self, key_states, value_states, layer_idx, cache_kwargs)
                                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/junwon_hwang/uv_env/exaone4/.venv/lib/python3.12/site-packages/transformers/cache_utils.py", line 1176, in update
    return self.layers[layer_idx].update(key_states, value_states, cache_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/junwon_hwang/uv_env/exaone4/.venv/lib/python3.12/site-packages/transformers/cache_utils.py", line 362, in update
    k_out_shifted = self.keys[:, :, indices]
                    ~~~~~~~~~^^^^^^^^^^^^^^^
RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cuda:1)
```

### Expected behavior

It should be run successfully when the model is deployed on multiple devices.
I guess matching the devices before indexing can be resolve this issue. 

Comments:
cc @Cyrilvallez and @manueldeprada 

(Thanks for reporting btw!)