Gemma3 doesn't support passing past_key_values
Author: Patchwork53
Pull request: no
Labels: name: bug

### System Info

OS: Ubuntu 24
Python: 3.11.10
Pytorch: 2.1.2
Transformers: 4.51.3

I tried to write a manual `my_generate()` function to emulate `model.generate()`. However, passing the `past_key_values` argument throws an out-of-bounds error.

I tried the same code with other models (`Qwen2_5_VL`) and those worked.

## Code
```
from transformers import AutoProcessor, Gemma3ForConditionalGeneration, Qwen2_5_VLForConditionalGeneration
import torch

# model_id = "Qwen/Qwen2.5-VL-7B-Instruct"
model_id = "google/gemma-3-12b-it"

model = Gemma3ForConditionalGeneration.from_pretrained(
    model_id, 
    device_map="cpu",
).eval()

processor = AutoProcessor.from_pretrained(model_id)

messages = [
    {
        "role": "system",
        "content": [{"type": "text", "text": "You are a helpful assistant."}]
    },
    {
        "role": "user",
        "content": [{"type": "text", "text": "What is the capital of France?"}]
    }
]

inputs = processor.apply_chat_template(
    messages, add_generation_prompt=True, tokenize=True,
    return_dict=True, return_tensors="pt"
).to(model.device)

input_ids = inputs["input_ids"]
device = input_ids.device

with torch.inference_mode():
    outputs = model.forward(
        input_ids=input_ids,
        past_key_values=None,
        use_cache=True
    )

next_token = outputs.logits[:, -1, :].argmax(dim=-1, keepdim=True)
print("1st token:", processor.decode(next_token[0]))

with torch.inference_mode():
    outputs = model.forward(
        input_ids=next_token,
        past_key_values=outputs.past_key_values,
        use_cache=True
    )

next_token = outputs.logits[:, -1, :].argmax(dim=-1, keepdim=True)
print("2nd token:", processor.decode(next_token[0]))
```

## Trace
```
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[1], line 44
     41 print("1st token:", processor.decode(next_token[0]))
     43 with torch.inference_mode():
---> 44     outputs = model.forward(
     45         input_ids=next_token,
     46         past_key_values=outputs.past_key_values,
     47         use_cache=True
     48     )
     50 next_token = outputs.logits[:, -1, :].argmax(dim=-1, keepdim=True)
     51 print("2nd token:", processor.decode(next_token[0]))

File ~/miniconda3/envs/lmm/lib/python3.11/site-packages/transformers/utils/generic.py:965, in can_return_tuple.<locals>.wrapper(self, *args, **kwargs)
    962     set_attribute_for_modules(self, "_is_top_level_module", False)
    964 try:
--> 965     output = func(self, *args, **kwargs)
    966     if is_requested_to_return_tuple or (is_configured_to_return_tuple and is_top_level_module):
    967         output = output.to_tuple()

File ~/miniconda3/envs/lmm/lib/python3.11/site-packages/transformers/utils/deprecation.py:172, in deprecate_kwarg.<locals>.wrapper.<locals>.wrapped_func(*args, **kwargs)
    168 elif minimum_action in (Action.NOTIFY, Action.NOTIFY_ALWAYS) and not is_torchdynamo_compiling():
    169     # DeprecationWarning is ignored by default, so we use FutureWarning instead
    170     warnings.warn(message, FutureWarning, stacklevel=2)
--> 172 return func(*args, **kwargs)

File ~/miniconda3/envs/lmm/lib/python3.11/site-packages/transformers/models/gemma3/modeling_gemma3.py:1326, in Gemma3ForConditionalGeneration.forward(self, input_ids, pixel_values, attention_mask, position_ids, past_key_values, token_type_ids, cache_position, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, logits_to_keep, **lm_kwargs)
   1321     labels = torch.where(input_ids == self.pad_token_id, self.config.ignore_index, labels)
   1323 causal_mask = self._update_causal_mask(
   1324     attention_mask, token_type_ids, past_key_values, cache_position, inputs_embeds, is_training
   1325 )
-> 1326 outputs: CausalLMOutputWithPast = self.language_model(
   1327     attention_mask=causal_mask,
   1328     position_ids=position_ids,
   1329     past_key_values=past_key_values,
   1330     inputs_embeds=inputs_embeds,
   1331     use_cache=use_cache,
   1332     output_attentions=output_attentions,
   1333     output_hidden_states=output_hidden_states,
   1334     cache_position=cache_position,
   1335     logits_to_keep=logits_to_keep,
   1336     **lm_kwargs,
   1337 )
   1339 logits = outputs.logits
   1340 loss = None

File ~/miniconda3/envs/lmm/lib/python3.11/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1516     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File ~/miniconda3/envs/lmm/lib/python3.11/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   1522 # If we don't have any hooks, we want to skip the rest of the logic in
   1523 # this function, and just call forward.
   1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1525         or _global_backward_pre_hooks or _global_backward_hooks
   1526         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527     return forward_call(*args, **kwargs)
   1529 try:
   1530     result = None

File ~/miniconda3/envs/lmm/lib/python3.11/site-packages/transformers/utils/generic.py:965, in can_return_tuple.<locals>.wrapper(self, *args, **kwargs)
    962     set_attribute_for_modules(self, "_is_top_level_module", False)
    964 try:
--> 965     output = func(self, *args, **kwargs)
    966     if is_requested_to_return_tuple or (is_configured_to_return_tuple and is_top_level_module):
    967         output = output.to_tuple()

File ~/miniconda3/envs/lmm/lib/python3.11/site-packages/transformers/utils/deprecation.py:172, in deprecate_kwarg.<locals>.wrapper.<locals>.wrapped_func(*args, **kwargs)
    168 elif minimum_action in (Action.NOTIFY, Action.NOTIFY_ALWAYS) and not is_torchdynamo_compiling():
    169     # DeprecationWarning is ignored by default, so we use FutureWarning instead
    170     warnings.warn(message, FutureWarning, stacklevel=2)
--> 172 return func(*args, **kwargs)

File ~/miniconda3/envs/lmm/lib/python3.11/site-packages/transformers/models/gemma3/modeling_gemma3.py:942, in Gemma3ForCausalLM.forward(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, cache_position, logits_to_keep, **loss_kwargs)
    938 output_hidden_states = (
    939     output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    940 )
    941 # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
--> 942 outputs: BaseModelOutputWithPast = self.model(
    943     input_ids=input_ids,
    944     attention_mask=attention_mask,
    945     position_ids=position_ids,
    946     past_key_values=past_key_values,
    947     inputs_embeds=inputs_embeds,
    948     use_cache=use_cache,
    949     output_attentions=output_attentions,
    950     output_hidden_states=output_hidden_states,
    951     cache_position=cache_position,
    952     **loss_kwargs,
    953 )
    955 hidden_states = outputs.last_hidden_state
    956 # Only compute necessary logits, and do not upcast them to float if we are not computing the loss

File ~/miniconda3/envs/lmm/lib/python3.11/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1516     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File ~/miniconda3/envs/lmm/lib/python3.11/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   1522 # If we don't have any hooks, we want to skip the rest of the logic in
   1523 # this function, and just call forward.
   1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1525         or _global_backward_pre_hooks or _global_backward_hooks
   1526         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527     return forward_call(*args, **kwargs)
   1529 try:
   1530     result = None

File ~/miniconda3/envs/lmm/lib/python3.11/site-packages/transformers/utils/generic.py:965, in can_return_tuple.<locals>.wrapper(self, *args, **kwargs)
    962     set_attribute_for_modules(self, "_is_top_level_module", False)
    964 try:
--> 965     output = func(self, *args, **kwargs)
    966     if is_requested_to_return_tuple or (is_configured_to_return_tuple and is_top_level_module):
    967         output = output.to_tuple()

File ~/miniconda3/envs/lmm/lib/python3.11/site-packages/transformers/models/gemma3/modeling_gemma3.py:722, in Gemma3TextModel.forward(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, use_cache, output_attentions, output_hidden_states, cache_position, last_cache_position, **flash_attn_kwargs)
    708     layer_outputs = self._gradient_checkpointing_func(
    709         partial(decoder_layer.__call__, **flash_attn_kwargs),
    710         hidden_states,
   (...)
    719         last_cache_position,
    720     )
    721 else:
--> 722     layer_outputs = decoder_layer(
    723         hidden_states,
    724         position_embeddings_global=position_embeddings_global,
    725         position_embeddings_local=position_embeddings_local,
    726         attention_mask=causal_mask,
    727         position_ids=position_ids,
    728         past_key_value=past_key_values,
    729         output_attentions=output_attentions,
    730         use_cache=use_cache,
    731         cache_position=cache_position,
    732         last_cache_position=last_cache_position,
    733         **flash_attn_kwargs,
    734     )
    736 hidden_states = layer_outputs[0]
    738 if output_attentions:

File ~/miniconda3/envs/lmm/lib/python3.11/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1516     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File ~/miniconda3/envs/lmm/lib/python3.11/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   1522 # If we don't have any hooks, we want to skip the rest of the logic in
   1523 # this function, and just call forward.
   1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1525         or _global_backward_pre_hooks or _global_backward_hooks
   1526         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527     return forward_call(*args, **kwargs)
   1529 try:
   1530     result = None

File ~/miniconda3/envs/lmm/lib/python3.11/site-packages/transformers/models/gemma3/modeling_gemma3.py:420, in Gemma3DecoderLayer.forward(self, hidden_states, position_embeddings_global, position_embeddings_local, attention_mask, position_ids, past_key_value, output_attentions, use_cache, cache_position, last_cache_position, **kwargs)
    417 else:
    418     position_embeddings = position_embeddings_global
--> 420 hidden_states, self_attn_weights = self.self_attn(
    421     hidden_states=hidden_states,
    422     position_embeddings=position_embeddings,
    423     attention_mask=attention_mask,
    424     position_ids=position_ids,
    425     past_key_value=past_key_value,
    426     output_attentions=output_attentions,
    427     use_cache=use_cache,
    428     cache_position=cache_position,
    429     **kwargs,
    430 )
    431 hidden_states = self.post_attention_layernorm(hidden_states)
    432 hidden_states = residual + hidden_states

File ~/miniconda3/envs/lmm/lib/python3.11/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1516     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File ~/miniconda3/envs/lmm/lib/python3.11/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   1522 # If we don't have any hooks, we want to skip the rest of the logic in
   1523 # this function, and just call forward.
   1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1525         or _global_backward_pre_hooks or _global_backward_hooks
   1526         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527     return forward_call(*args, **kwargs)
   1529 try:
   1530     result = None

File ~/miniconda3/envs/lmm/lib/python3.11/site-packages/transformers/models/gemma3/modeling_gemma3.py:322, in Gemma3Attention.forward(self, hidden_states, position_embeddings, attention_mask, past_key_value, cache_position, **kwargs)
    314 if past_key_value is not None:
    315     # sin and cos are specific to RoPE models; cache_position needed for the static cache
    316     cache_kwargs = {
    317         "sin": sin,
    318         "cos": cos,
    319         "cache_position": cache_position,
    320         "sliding_window": self.sliding_window,
    321     }
--> 322     key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
    324     # Here we need to slice as we use a static cache by default, but FA2 does not support it
    325     if attention_mask is not None and self.config._attn_implementation == "flash_attention_2":

File ~/miniconda3/envs/lmm/lib/python3.11/site-packages/transformers/cache_utils.py:1782, in HybridCache.update(self, key_states, value_states, layer_idx, cache_kwargs)
   1779 else:
   1780     update_fn = self._static_update
-> 1782 return update_fn(
   1783     cache_position,
   1784     layer_idx,
   1785     key_states,
   1786     value_states,
   1787     k_out,
   1788     v_out,
   1789     k_out.shape[2],
   1790 )

File ~/miniconda3/envs/lmm/lib/python3.11/site-packages/transformers/cache_utils.py:1746, in HybridCache._static_update(self, cache_position, layer_idx, key_states, value_states, k_out, v_out, max_cache_len)
   1745 def _static_update(self, cache_position, layer_idx, key_states, value_states, k_out, v_out, max_cache_len):
-> 1746     k_out[:, :, cache_position] = key_states
   1747     v_out[:, :, cache_position] = value_states
   1749     self.key_cache[layer_idx] = k_out

IndexError: index 23 is out of bounds for dimension 0 with size 23
```

`_static_update` tries to savethe KV of the current token position. I believe `k_out` and `v_out` was supposed to be expanded in some previous step before `cache_position` can be accessed.





### Who can help?

@amyeroberts, @qubvel

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

Run the code.

### Expected behavior

Output should be:
```
1st token: The
2nd token:  capital
```
with no Errors.

Comments:
@Patchwork53 It is becase Gemma3 requires specific cache and the users are expected to prepare the cache in advance, in contrast to other models. 

You have to prepare [Hybrid Cache](https://huggingface.co/docs/transformers/main/en/internal/generation_utils#transformers.HybridCache) with the maximum sequence length and pass it as `past_key_values`