RecurrentGemma crashes during inference for inputs longer than sliding window width
Author: assafbk
Pull request: no
Labels: name: bug

### System Info

**System Info:**
- `transformers` version: 4.50.3
- Platform: Linux-6.8.0-52-generic-x86_64-with-glibc2.35
- Python version: 3.10.16
- Huggingface_hub version: 0.30.1
- Safetensors version: 0.5.3
- Accelerate version: 1.6.0
- Accelerate config:    not found
- DeepSpeed version: not installed
- PyTorch version (GPU?): 2.5.1+cu118 (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?: no
- Using GPU in script?: yes
- GPU type: NVIDIA RTX A6000

### Who can help?

@ArthurZucker, @gante 

### Information

- [x] The official example scripts
- [ ] 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

Code snippet for reproduction:
```
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("google/recurrentgemma-9b-it")
model = AutoModelForCausalLM.from_pretrained("google/recurrentgemma-9b-it", device_map="cuda", torch_dtype=torch.float16)

input_text = "Write me a poem about Machine Learning." * 300    # This string is 2402 tokens long, which is larger than 2048, the sliding window attention width
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")

outputs = model.generate(**input_ids, max_new_tokens=20)
print(tokenizer.decode(outputs[0]))
```

Error message:

> Traceback (most recent call last):
>   File "/data2/assaf/tmp/test_rg.py", line 13, in <module>
>     outputs = model.generate(**input_ids, max_new_tokens=20)
>   File "/data2/assaf/conda_envs/recurrent_gemma_tmp/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
>     return func(*args, **kwargs)
>   File "/data2/assaf/conda_envs/recurrent_gemma_tmp/lib/python3.10/site-packages/transformers/generation/utils.py", line 2326, in generate
>     result = self._sample(
>   File "/data2/assaf/conda_envs/recurrent_gemma_tmp/lib/python3.10/site-packages/transformers/generation/utils.py", line 3289, in _sample
>     outputs = model_forward(**model_inputs, return_dict=True)
>   File "/data2/assaf/conda_envs/recurrent_gemma_tmp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
>     return self._call_impl(*args, **kwargs)
>   File "/data2/assaf/conda_envs/recurrent_gemma_tmp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
>     return forward_call(*args, **kwargs)
>   File "/data2/assaf/conda_envs/recurrent_gemma_tmp/lib/python3.10/site-packages/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py", line 852, in forward
>     outputs = self.model(
>   File "/data2/assaf/conda_envs/recurrent_gemma_tmp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
>     return self._call_impl(*args, **kwargs)
>   File "/data2/assaf/conda_envs/recurrent_gemma_tmp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
>     return forward_call(*args, **kwargs)
>   File "/data2/assaf/conda_envs/recurrent_gemma_tmp/lib/python3.10/site-packages/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py", line 717, in forward
>     causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position)
>   File "/data2/assaf/conda_envs/recurrent_gemma_tmp/lib/python3.10/site-packages/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py", line 764, in _update_causal_mask
>     padding_mask = causal_mask[..., :mask_length].eq(0.0) * attention_mask[:, None, None, :].eq(0.0)
> RuntimeError: The size of tensor a (2048) must match the size of tensor b (2402) at non-singleton dimension 3

### Expected behavior

If the sequence is longer than the sliding window width (like it is now in the script) then the script crashes with the error message above.

If the sequence is shorter than the sliding window width (e.g. replace *300 by *200) then the script runs fine.

The bug was seen in transformers version v4.50.3
It does not reproduce on earlier transformers versions, such as v4.42.4



Comments:
Hi @assafbk, can you use [git bisect](https://git-scm.com/docs/git-bisect) to identify the commit where this bug appeared? It'll help a lot in identifying the cause and assigning the right person to fix it!

RecurrentGemma generation was functional until commit 295a90cb40468ddd607577448e674348cff55adc (v4.46.0), which introduced two regressions:

1. Removed the model-specific `prepare_inputs_for_generation`, and in doing so, stopped forwarding `past_key_values` via `**kwargs`. This broke generation entirely.
2. Even if `past_key_values` were passed correctly, the same commit also introduced a bug for long prompts due to incorrect `attention_mask` handling during decoding.

Commit 3d6e55c7e73f425e0fcf0f27cae54ddf8b369291 (v4.48.3) later restored forwarding of `**kwargs`, which incidentally fixed generation. However, the long-context bug persists.

### Root cause
RecurrentGemma assumes that the `attention_mask` received during decoding is already cropped to the model’s sliding attention window (`attention_window_size`). The generic `prepare_inputs_for_generation` does not currently enforce this, leading to shape mismatches in `_update_causal_mask()` when the prompt length exceeds the attention window.

### Proposed fix
Add a one-line safeguard inside `_update_causal_mask()` to truncate `attention_mask` when it's longer than the causal mask, as its was done on the original `prepare_inputs_for_generation`.

Submitting a PR with the fix.