Max cache length issue with Gemma 3
Author: mitchelldehaven
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.54.0
- Platform: Linux-6.11.0-1011-nvidia-x86_64-with-glibc2.39
- Python version: 3.11.13
- Huggingface_hub version: 0.34.1
- Safetensors version: 0.5.3
- Accelerate version: 1.9.0
- Accelerate config:    not found
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.7.1+cu128 (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?: no
- Using GPU in script?: yes
- GPU type: NVIDIA RTX PRO 6000 Blackwell Workstation Edition

### Who can help?

@gante 

### 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

Trying to generate with `AutoModelForCausalLM` with `google/gemma3-1b-it` is running into the following error when using `eager` attention as specified to do from logging messages (rather than using `sdpa`).

```
ValueError: Max cache length is not consistent across layers: [512, 512, 512, 512, 512, 741, 512, 512, 512, 512, 512, 741, 512, 512, 512, 512, 512, 741, 512, 512, 512, 512, 512, 741, 512, 512]
```

The offending code seems to be here in `transformers/cache_utils.py`
```
    @property
    def max_cache_len(self) -> int:
        """Return the maximum cache length of the cache"""
        values = [layer.max_cache_len for layer in self.layers]
        if len(set(values)) > 1:
            raise ValueError(f"Max cache length is not consistent across layers: {values}")
        return values[0]
```

This check seems to be inconsistent with Gemma3's layer structure, where 5 layers use sliding attention with size `512` and 6th layer uses full casual attention.

The gitblame shows this was changed recently from the commit: https://github.com/huggingface/transformers/commit/c338fd43b0be2c7f5d73e693fa6fb1b5e7a0bdc2

This worked fine in the previous transformers version I was running, but I needed to update my version recently and this error started occurring. I'm happy to post a PR for a fix if this is determined to be a bug.

### Expected behavior

That the model generates text as expected.

Comments:
For what its worth, this simple change seems to fix the issue:
```
    @property
    def max_cache_len(self) -> int:
        """Return the maximum cache length of the cache"""
        values = [layer.max_cache_len for layer in self.layers]
        # if len(set(values)) > 1:
        #     raise ValueError(f"Max cache length is not consistent across layers: {values}")
        return max(values)
```

We ran into the same bug with gemma-2-2b

Hey, can you share a small code snippet reproducing the issue?

cc @manueldeprada as well

@zucchini-nlp 
```
from transformers import Gemma3ForConditionalGeneration, AutoTokenizer, AutoModelForCausalLM, GenerationConfig
import datasets
import json
import torch
from tqdm import tqdm


def tools_format_python(tool_dicts):
    functions_formatted = []
    tab_s = "    "
    for tool_dict in tool_dicts:
        function_name = tool_dict["name"]
        function_args = [f"{k}: {v['type']}" for k, v in tool_dict["parameters"].items()]
        function_formatted = (
            f"def {function_name}({', '.join(function_args)})\n" +
            f"{tab_s}{tool_dict['description']}\n\n" +
            f"{tab_s}Args:\n" +
            "\n".join(f"{tab_s*2}{k}: {v['description']}" for k, v in tool_dict["parameters"].items())
        )
        functions_formatted.append(function_formatted)
    tools_prompt = (
        "```python\n" +
        "\n\n".join(functions_formatted) + 
        "\n```\n\n"
    )
    return tools_prompt

llm_model_name = "google/gemma-3-1b-it"
device = "cuda" if torch.cuda.is_available() else "cpu"
model = AutoModelForCausalLM.from_pretrained(llm_model_name, attn_implementation="eager").to(device)
tokenizer = AutoTokenizer.from_pretrained(llm_model_name)
dataset = datasets.load_dataset("Salesforce/xlam-function-calling-60k")["train"]
system_message = "You are a helpful function calling AI assistant. Below are the Python functions accessible to you:\n"
for sample in tqdm(dataset):
    sample_tools = json.loads(sample["tools"]) * 10
    tools_prompt = tools_format_python(sample_tools)
    gemma_prompt_format = (
        "<bos><start_of_turn>user\n" +
        system_message + 
        tools_prompt + 
        "User: " + sample["query"] + "<end_of_turn>\n<start_of_turn>model\n```tool_call\n"
    )
    tokenized_prompt = tokenizer(gemma_prompt_format, return_tensors="pt", add_special_tokens=False)["input_ids"].to(device)
    with torch.no_grad(), torch.amp.autocast(device, dtype=torch.bfloat16):
        outputs = model.generate(
            tokenized_prompt,
            max_new_tokens=768,
            use_cache=True,
        )
```

This will trigger the error.

Indeed, this is consistent with previous behaviour. Thanks for reporting! Opening a PR