KV cache size is not updating?
Author: tianhaoz95
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.50.0.dev0
- Platform: macOS-15.3.2-arm64-arm-64bit
- Python version: 3.12.9
- Huggingface_hub version: 0.29.3
- Safetensors version: 0.5.3
- Accelerate version: 1.5.2
- Accelerate config:    not found
- DeepSpeed version: not installed
- PyTorch version (GPU?): 2.6.0 (False)
- 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>

### Who can help?

I'm experimenting with using kv cache explicitly (this is for prefill decode disaggregation) with the following code:

```python
def setup_model(device_map="cpu"):
    tokenizer = AutoTokenizer.from_pretrained("google/gemma-3-1b-it")
    model = AutoModelForCausalLM.from_pretrained(
        "google/gemma-3-1b-it",
        torch_dtype=torch.float16,
        device_map=device_map,
    )
    return model, tokenizer

model, tokenizer = setup_model(device_map="cpu")
prompt = "Write a short poem about coding:"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model(
        input_ids=inputs.input_ids,
        use_cache=True,
        return_dict=True
    )
print(f"Prefill kv cache, kv cache size: {len(outputs.past_key_values.key_cache)}, and shape: {outputs.past_key_values.key_cache[0].shape}")
past_key_values = outputs.past_key_values
generated_tokens = []

next_token_logits = outputs.logits[:, -1, :]

for _ in range(50):
    with torch.no_grad():
        next_token_id = torch.argmax(next_token_logits, dim=-1)
        
        if next_token_id.item() == tokenizer.eos_token_id:
            break

        generated_tokens.append(next_token_id.item())

        outputs = model(
            input_ids=next_token_id.unsqueeze(0),
            use_cache=True,
            past_key_values=past_key_values,
            return_dict=True
        )
        print(f"Update kv cache, kv cache size: {len(outputs.past_key_values.key_cache)}, and shape: {outputs.past_key_values.key_cache[0].shape}")
        past_key_values = outputs.past_key_values
        next_token_logits = outputs.logits[:, -1, :]

all_tokens = torch.cat([inputs.input_ids[0], torch.tensor(generated_tokens)])
generated_text = tokenizer.decode(all_tokens)
print(f"\nGenerated text:\n{generated_text}")
```

I found the output is:

```
Prefill kv cache, kv cache size: 26, and shape: torch.Size([1, 1, 8, 256])
Traceback (most recent call last):
  File "/Users/tianhaoz/Downloads/medium/pd_disaggregation/inspect_kv_cache.py", line 79, in <module>
    main()
  File "/Users/tianhaoz/Downloads/medium/pd_disaggregation/inspect_kv_cache.py", line 63, in main
    outputs = model(
              ^^^^^^
  File "/opt/anaconda3/envs/llm/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/llm/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/llm/lib/python3.12/site-packages/transformers/utils/deprecation.py", line 172, in wrapped_func
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/llm/lib/python3.12/site-packages/transformers/models/gemma3/modeling_gemma3.py", line 976, in forward
    outputs = self.model(
              ^^^^^^^^^^^
  File "/opt/anaconda3/envs/llm/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/llm/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/llm/lib/python3.12/site-packages/transformers/models/gemma3/modeling_gemma3.py", line 754, in forward
    layer_outputs = decoder_layer(
                    ^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/llm/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/llm/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/llm/lib/python3.12/site-packages/transformers/models/gemma3/modeling_gemma3.py", line 443, in forward
    hidden_states, self_attn_weights = self.self_attn(
                                       ^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/llm/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/llm/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/llm/lib/python3.12/site-packages/transformers/models/gemma3/modeling_gemma3.py", line 347, in forward
    key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/llm/lib/python3.12/site-packages/transformers/cache_utils.py", line 1721, in update
    return update_fn(
           ^^^^^^^^^^
  File "/opt/anaconda3/envs/llm/lib/python3.12/site-packages/transformers/cache_utils.py", line 1687, in _static_update
    k_out[:, :, cache_position] = key_states
    ~~~~~^^^^^^^^^^^^^^^^^^^^^^
IndexError: index 8 is out of bounds for dimension 0 with size 8
```

does this mean in the decode steps, the sequence length of kv cache is not updating? looks like it might be forever staying at torch.Size([1, 1, 8, 256]) after prefill. Am I doing something wrong?

### 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 example code in the last section.

### Expected behavior

It should generate result and print updated kv cache size each time.

Comments:
Hey @tianhaoz95 !

Gemma3 uses a static-shaped `HybridCache` instead of dynamic cache as most other models. Thus when you do a forward pass the first time, a default cache is initialized with max length equal to input length. The next iterations of forward fail to go further the indicated max length

If you want to generate by manually looping as in the example, you need to explicitly pass a [`HybridCache`](https://huggingface.co/docs/transformers/v4.50.0/en/internal/generation_utils#transformers.HybridCache) with the max length big enough to hold all new tokens

Got it. Very interesting choice in gemma 3. Thanks!