Wrong KV cache update for sliding-window attention (SWA) layers when total sequence length reaches window size
Author: plienhar
Pull request: no
Labels: name: bug

### System Info

```text
- `transformers` version: 4.51.2
- Platform: Linux-6.8.0-1021-aws-x86_64-with-glibc2.35
- Python version: 3.12.8
- Huggingface_hub version: 0.30.2
- Safetensors version: 0.5.3
- Accelerate version: 1.5.2
- Accelerate config:    not found
- DeepSpeed version: not installed
- PyTorch version (GPU?): 2.6.0+cu124 (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?: No
- GPU type: NVIDIA L40S
```

### Who can help?

@ArthurZucker 

### Reproduction

```python
from types import SimpleNamespace

import numpy as np
import torch
import torch.nn.functional as F
from transformers.cache_utils import HybridCache

def update_2d_attention_mask(attention_mask_2d: torch.LongTensor, padding_side: str) -> torch.LongTensor:
    # From tranformers.utils.GenerationMixin._update_model_kwargs_for_generation
    batch_size, _ = attention_mask_2d.shape
    if padding_side == "left":
        attention_mask_2d = torch.cat([attention_mask_2d, attention_mask_2d.new_ones((batch_size, 1))], dim=1)
    else:
        attention_mask_2d = torch.cat([attention_mask_2d.new_ones((batch_size, 1)), attention_mask_2d], dim=1)
    return attention_mask_2d

def create_cache_positions(attention_mask_2d: torch.LongTensor, is_prefill: bool) -> torch.LongTensor:
    # From tranformers.utils.GenerationMixin._get_initial_cache_position
    cache_position = torch.ones_like(attention_mask_2d[0, :], dtype=torch.int64).cumsum(0) - 1
    if is_prefill:
        return cache_position
    else:
        return cache_position[-1:]

config = SimpleNamespace(
    num_hidden_layers=1,
    num_key_value_heads=1,
    num_attention_heads=1,
    head_dim=2,
    hidden_size=2,
    sliding_window=8,
    sliding_window_pattern=4
)

max_batch_size = 1
max_cache_length = 12

kv_cache_manager = HybridCache(
    config=config,
    max_batch_size=max_batch_size,
    max_cache_len=max_cache_length
)

# ** INPUTS **
# 2D attention mask of shape (batch_size, max_input_length) from the tokenizer output
attention_mask_2d = torch.tensor(
    [[1, 1, 1, 1, 1, 1]], 
    dtype=torch.int32
    )

batch_size, input_sequence_length = attention_mask_2d.shape

# Prefill key update of shape (max_batch_size, num_key_value_heads, max_input_length, head_dim)
key_states = torch.arange(1, input_sequence_length+1)[None, None, :, None].to(dtype=torch.float32).expand(-1, -1, -1, config.head_dim)
# Prefill cache positions of shape (max_input_length,). Tensor of absolute cache positions where the update will be written to.
cache_position = create_cache_positions(attention_mask_2d=attention_mask_2d, is_prefill=True)

print("---- INITIAL CACHE STATE ----")
print("  - k_cache:")
print(np.array2string(kv_cache_manager.key_cache[0].transpose(3, 2).numpy(), precision=4, suppress_small=True))

# ** PREFILL ** 
print("\n---- PREFILL ----")
print("> k_cache state before cache update:")
print(np.array2string(kv_cache_manager.key_cache[0].transpose(3, 2).numpy(), precision=4, suppress_small=True))
print("> Cache update inputs:")
print("  - cache_position:")
print(np.array2string(cache_position.numpy(), precision=4, suppress_small=True))
print("  - cache update (keys):")
print(np.array2string(key_states.transpose(3, 2).numpy(), precision=4, suppress_small=True))
k_cache, v_cache = kv_cache_manager.update(
    key_states=key_states,
    value_states=key_states+1,
    layer_idx=0,
    cache_kwargs={
        "cache_position": cache_position,
        "sliding_window": config.sliding_window,
    }
)
print("> k_cache state after cache update:")
print(np.array2string(kv_cache_manager.key_cache[0].transpose(3, 2).numpy(), precision=4, suppress_small=True))

attention_mask_2d = update_2d_attention_mask(attention_mask_2d=attention_mask_2d, padding_side="left")

# ** TOKEN GENERATION **
for token_gen_step_idx in range (max_cache_length - input_sequence_length):
    print(f"\n---- TOKEN GENERATION STEP #{token_gen_step_idx+1} ----")
    batch_size, input_sequence_length = attention_mask_2d.shape
    # Token generation key update of shape (max_batch_size, num_key_value_heads, 1, head_dim)
    key_states = torch.tensor([[input_sequence_length]])[:, None, :, None].to(dtype=torch.float32).expand(-1, -1, -1, config.head_dim)
    # Token generation cache_positions of shape (1,). Its unique value is the absolute cache position the update will be written to.
    cache_position = create_cache_positions(attention_mask_2d=attention_mask_2d, is_prefill=False)
    print("> k_cache state before cache update:")
    print(np.array2string(kv_cache_manager.key_cache[0].transpose(3, 2).numpy(), precision=4, suppress_small=True))
    print("> Cache update inputs:")
    print("  - cache_position:")
    print(np.array2string(cache_position.numpy(), precision=4, suppress_small=True))
    print("  - cache update (keys):")
    print(np.array2string(key_states.transpose(3, 2).numpy(), precision=4, suppress_small=True))

    k_cache, v_cache = kv_cache_manager.update(
        key_states=key_states,
        value_states=key_states+1,
        layer_idx=0,
        cache_kwargs={
            "cache_position": cache_position,
            "sliding_window": config.sliding_window,
        }
    )
    attention_mask_2d = update_2d_attention_mask(attention_mask_2d=attention_mask_2d, padding_side="left")

    print("> k_cache state after cache update:")
    print(np.array2string(kv_cache_manager.key_cache[0].transpose(3, 2).numpy(), precision=4, suppress_small=True))
```

### Expected behavior
Let's take a simple example assuming batch size 1, window size 8 and input sequence length of 6. The cache is initialized as an tensor 
of shape (batch_size, window_size) full of zeros.
 
* Initial cache state:
```text
[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
```
1. Assuming keys computed during the prefill phase are `[10.0, 10.1, 10.2, 10.3, 10.4, 10.5]`, final sequence length is 6 and 
the cache state after update is: 
```text
[10.0, 10.1, 10.2, 10.3, 10.4, 10.5, 0.0, 0.0]
```
2. Assuming the keys computed during token generation step 1 are `[10.6]`, final sequence length is 7 and the cache state after update is:
```text
[10.0, 10.1, 10.2, 10.3, 10.4, 10.5, 10.6, 0.0]
```
3. Assuming the keys computed during token generation step 2 are `[10.7]`, final sequence length is 8 and the cache state after update is:
```text
[10.0, 10.1, 10.2, 10.3, 10.4, 10.5, 10.6, 10.7]
```
4. Assuming the keys computed during token generation step 3 are `[10.8]`, final sequence length is 9 and now exceeds the window size. The cache must therefore be rolled before being updated and the cache state after update is:
```text
[10.1, 10.2, 10.3, 10.4, 10.5, 10.6, 10.7, 10.8]
```
5. Assuming the keys computed during token generation step 4 are `[10.9]`, final sequence length is 10 and exceeds the window size. The cache must therefore be rolled before being updated and the cache state after update is:
```text
[10.2, 10.3, 10.4, 10.5, 10.6, 10.7, 10.8, 10.9]
```

### Actual behavior
The actual behavior matches the expected behavior in the following cases:
* If the initial prompt length is larger than the window size.
* If the initial prompt length is smaller than the window size & total sequence length remain below window size - 1.

However, for sequences that start smaller than the window size and which grow larger than the window size during generation, 
we observe the following diverging behavior (cf. above Python code for reproduction).

* Initial cache state:
```text
[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
```

1. Cache state after prefill update: 
```text
[10.0, 10.1, 10.2, 10.3, 10.4, 10.5, 0.0, 0.0]
```
2. Cache state after token generation step 1 update: 
```text
[10.0, 10.1, 10.2, 10.3, 10.4, 10.5, 10.6, 0.0]
```
3. Cache state after token generation step 2 update: 
```text
[10.1, 10.2, 10.3, 10.4, 10.5, 10.6, 0.0, 10.7]
```
4. Cache state after token generation step 3 update: 
```text
[10.2, 10.3, 10.4, 10.5, 10.6, 0.0, 10.7, 10.8]
```
5. Cache state after token generation step 4 update: 
```text
[10.3, 10.4, 10.5, 10.6, 0.0, 10.7, 10.8, 10.9]
```

Essentially, the current cache update implementation is rolling the cache one step too early.
Within the code source, this means that the rolling condition should be `to_shift = cache_position > max_cache_len - 1` 
instead of `to_shift = cache_position >= max_cache_len - 1`.

Updating one step too early wrongfully:
* Shift by one position all tokens with absolute position lower or equal to window size - 1.
* Add a zero key & value at position window size.

Impact is long lasting as we only recover the expected cache state when total sequence length reaches two times the window size. 

### Possible fix
In our understanding, the above behavior can be fixed by changing two lines in `HybridCache._sliding_update` and `SlidingWindowCache._sliding_update`.

```python
class FixedHybridCache(HybridCache):

    def _sliding_update(self, 
                        cache_position: torch.LongTensor, 
                        layer_idx: int, 
                        key_states: torch.FloatTensor, 
                        value_states: torch.FloatTensor, 
                        k_out: torch.FloatTensor, 
                        v_out: torch.FloatTensor, 
                        max_cache_len: int) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
        if cache_position.shape[0] > max_cache_len:
            k_out = key_states[:, :, -max_cache_len:, :]
            v_out = value_states[:, :, -max_cache_len:, :]
            # Assumption: caches are all zeros at this point, `+=` is equivalent to `=` but compile-friendly
            self.key_cache[layer_idx] += k_out
            self.value_cache[layer_idx] += v_out
            # we should return the whole states instead of k_out, v_out to take the whole prompt
            # into consideration when building kv cache instead of just throwing away tokens outside of the window
            return key_states, value_states

        slicing = torch.ones(max_cache_len, dtype=torch.long, device=value_states.device).cumsum(0)
        # >>>>>>>>>>>>>> START OF FIX >>>>>>>>>>>>>>
        # cache_position = cache_position.clamp(0, max_cache_len - 1)
        # to_shift = cache_position >= max_cache_len - 1
        to_shift = cache_position > max_cache_len - 1
        cache_position = cache_position.clamp(0, max_cache_len - 1)
        # <<<<<<<<<<<<<< END OF FIX <<<<<<<<<<<<<<<<
        indices = (slicing + to_shift[-1].int() - 1) % max_cache_len
        
        k_out = k_out[:, :, indices]
        v_out = v_out[:, :, indices]

        k_out[:, :, cache_position] = key_states
        v_out[:, :, cache_position] = value_states
        # `_.zero()` followed by `+=` is equivalent `=`, but compile-friendly (without graph breaks due to assignment)
        self.key_cache[layer_idx].zero_()
        self.value_cache[layer_idx].zero_()

        self.key_cache[layer_idx] += k_out
        self.value_cache[layer_idx] += v_out
        return k_out, v_out
```

Comments:
cc @gante 

I may have unintentionally fixed this issue while writing sanity-check tests for all caches: https://github.com/huggingface/transformers/pull/37873

(@manueldeprada your tests will still be useful here, to prevent regressions)

The PR should definitely fix the bug and prevent any regression. thanks a lot @plienhar for the super detailed report!! :)

This was not properly handled in the PR mentioned above: proper fix is https://github.com/huggingface/transformers/pull/38046! Thanks a lot for reporting the issue! Great catch!

Actual generation do not look like it changed quite a lot after the fix (I mean from the state before https://github.com/huggingface/transformers/pull/37972 which broke it more, unentotionally), but they still appear better, especially with regards to special tokens, at least with Gemma2