Repository: huggingface/transformers
Pull request: New cache tests and modular Hybrid Cache
## What does this PR do?

1. Refactor out the cache update logic for static and sliding window attention mechanisms.
Extracts the core update logic into two new utility functions: `_static_cache_update_logic` and `_sliding_cache_update_logic`. This way, there is only one implementation for StaticCache, SlidingWindowCache, and HybridCache.
@ArthurZucker @gante this is a first step towards per-layer modular cache definitions.
2. Added new synthetic tests for caches. Fixes #37574 and should catch similar bugs.

This is preliminary work for #38077

Problem statement:
Wrong KV cache update for sliding-window attention (SWA) layers when total sequence length reaches window size
### 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
```

Hints from issue discussion:
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!! :)

Commit message: Update src/transformers/cache_utils.py

Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>

Patch to review:
diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py
index 85a09f03de22..6cf78906df3e 100644
--- a/src/transformers/cache_utils.py
+++ b/src/transformers/cache_utils.py
@@ -21,6 +21,104 @@
 logger = logging.get_logger(__name__)
 
 
+# Utility functions for static/sliding cache update logic
+def _static_cache_update(
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    key_states: torch.Tensor,
+    value_states: torch.Tensor,
+    cache_position: Optional[torch.LongTensor],
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Updates the static cache tensors in place.
+
+    Args:
+        k_cache (`torch.Tensor`): The key cache tensor to update.
+        v_cache (`torch.Tensor`): The value cache tensor to update.
+        key_states (`torch.Tensor`): The new key states to add.
+        value_states (`torch.Tensor`): The new value states to add.
+        cache_position (`Optional[torch.LongTensor]`): The position indices where the new states should be inserted.
+                                                       If None, the entire cache is overwritten (prefill).
+
+    Returns:
+        Tuple[`torch.Tensor`, `torch.Tensor`]: The updated key and value cache tensors (modified in-place).
+    """
+    if cache_position is None:
+        # Prefill phase where seq_len potentially equals max_cache_len. Directly copy.
+        k_cache.copy_(key_states)
+        v_cache.copy_(value_states)
+    else:
+        # Generation phase. Update specific positions.
+        # Use index_copy_ for in-place update (compile-friendly).
+        try:
+            k_cache.index_copy_(2, cache_position, key_states)
+            v_cache.index_copy_(2, cache_position, value_states)
+        except NotImplementedError:
+            # Fallback for devices like MPS where index_copy_ might not be supported.
+            k_cache[:, :, cache_position] = key_states
+            v_cache[:, :, cache_position] = value_states
+    return k_cache, v_cache
+
+
+def _sliding_cache_update_logic(
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    key_states: torch.Tensor,
+    value_states: torch.Tensor,
+    cache_position: torch.LongTensor,
+    max_cache_len: int,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Updates the sliding window cache tensors, returning the potentially modified tensors.
+
+    Args:
+        k_cache (`torch.Tensor`): The key cache tensor to update.
+        v_cache (`torch.Tensor`): The value cache tensor to update.
+        key_states (`torch.Tensor`): The new key states to add.
+        value_states (`torch.Tensor`): The new value states to add.
+        cache_position (`torch.LongTensor`): The position indices where the new states should be inserted.
+        max_cache_len (`int`): The maximum length of the sliding window cache.
+
+    Returns:
+        Tuple[`torch.Tensor`, `torch.Tensor`]: The key and value tensors representing the cache state after the update.
+                                               For prefill > window, these are the full input states.
+                                               Otherwise, they are the updated cache tensors.
+    """
+    # Handle prefill phase when prompt length > sliding_window_size
+    if cache_position.shape[0] > max_cache_len:
+        new_k = key_states[:, :, -max_cache_len:, :]
+        new_v = value_states[:, :, -max_cache_len:, :]
+        k_cache.copy_(new_k)
+        v_cache.copy_(new_v)
+        return key_states, value_states
+
+    # Sliding window logic for generation phase or prefill < window
+    slicing = torch.arange(max_cache_len, device=value_states.device)
+    current_seq_len = cache_position[-1] + 1  # Use last position to determine current length
+    should_shift = current_seq_len > max_cache_len
+    indices = (slicing + should_shift.sum()) % max_cache_len
+
+    k_out_shifted = k_cache[:, :, indices]
+    v_out_shifted = v_cache[:, :, indices]
+
+    # Clamp cache_position to determine the *target index* within the shifted cache view
+    update_position = cache_position.clamp(min=0, max=max_cache_len - 1)
+
+    try:
+        k_out_updated = k_out_shifted.index_copy(2, update_position, key_states)
+        v_out_updated = v_out_shifted.index_copy(2, update_position, value_states)
+    except NotImplementedError:
+        # Fallback for MPS: clone and modify the clone
+        k_out_updated = k_out_shifted.clone()
+        v_out_updated = v_out_shifted.clone()
+        k_out_updated[:, :, update_position] = key_states
+        v_out_updated[:, :, update_position] = value_states
+
+    k_cache.copy_(k_out_updated)
+    v_cache.copy_(v_out_updated)
+    return k_out_updated, v_out_updated
+
+
 class Cache:
     """
     Base, abstract class for all caches. The actual data structure is specific to each subclass.
@@ -1262,28 +1360,16 @@ def update(
         """
         if cache_kwargs is None:
             cache_kwargs = {}
-        cache_position = cache_kwargs.get("cache_position")
-        k_out = self.key_cache[layer_idx]
-        v_out = self.value_cache[layer_idx]
-        key_states = key_states.to(k_out.dtype)
-        value_states = value_states.to(v_out.dtype)
 
-        if cache_position is None:
-            k_out.copy_(key_states)
-            v_out.copy_(value_states)
-        else:
-            # Note: here we use `tensor.index_copy_(dim, index, tensor)` that is equivalent to
-            # `tensor[:, :, index] = tensor`, but the first one is compile-friendly and it does explicitly an in-place
-            # operation, that avoids copies and uses less memory.
-            try:
-                k_out.index_copy_(2, cache_position, key_states)
-                v_out.index_copy_(2, cache_position, value_states)
-            except NotImplementedError:
-                # The operator 'aten::index_copy.out' is not currently implemented for the MPS device.
-                k_out[:, :, cache_position] = key_states
-                v_out[:, :, cache_position] = value_states
-
-        return k_out, v_out
+        key_states = key_states.to(self.key_cache[layer_idx].dtype)
+        value_states = value_states.to(self.value_cache[layer_idx].dtype)
+        return _static_cache_update_logic(
+            self.key_cache[layer_idx],
+            self.value_cache[layer_idx],
+            key_states,
+            value_states,
+            cache_kwargs.get("cache_position"),
+        )
 
     def get_seq_length(self, layer_idx: Optional[int] = 0) -> int:
         """Returns the sequence length of the cached states that were seen by the model."""
@@ -1310,9 +1396,9 @@ class SlidingWindowCache(StaticCache):
     if true(which means the cache can not hold all the old key value states and new states together because of the sliding window constraint),
     we need to do a cycle shift based on `indices` to replace the oldest states by the new key value states passed in.
 
-    The `to_shift` is only true once we are above sliding_window. Thus with `sliding_window==64`:
+    The `should_shift` is only true once we are above sliding_window. Thus with `sliding_window==64`:
 
-    indices = (slicing + to_shift[-1].int()-1) % self.config.sliding_window
+    indices = (slicing + should_shift[-1].sum()-1) % self.config.sliding_window
     tensor([ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17, 18,
         19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36,
         37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54,
@@ -1396,46 +1482,21 @@ def update(
         if cache_kwargs is None:
             cache_kwargs = {}
         cache_position = cache_kwargs.get("cache_position")
-        k_out = self.key_cache[layer_idx]
-        v_out = self.value_cache[layer_idx]
-        key_states = key_states.to(k_out.dtype)
-        value_states = value_states.to(v_out.dtype)
-
-        # assume this only happens in prefill phase when prompt length > sliding_window_size (= max_cache_len)
-        if cache_position.shape[0] > self.max_cache_len:
-            k_out = key_states[:, :, -self.max_cache_len :, :]
-            v_out = value_states[:, :, -self.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(self.max_cache_len, dtype=torch.long, device=value_states.device).cumsum(0)
-        cache_position = cache_position.clamp(0, self.max_cache_len - 1)
-        to_shift = cache_position > self.max_cache_len - 1
-        indices = (slicing + to_shift[-1].int() - 1) % self.max_cache_len
-
-        k_out = k_out[:, :, indices]
-        v_out = v_out[:, :, indices]
 
-        try:
-            k_out.index_copy_(2, cache_position, key_states)
-            v_out.index_copy_(2, cache_position, value_states)
-        except NotImplementedError:
-            # The operator 'aten::index_copy.out' is not currently implemented for the MPS device.
-            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_()
+        if cache_position is None:
+            raise ValueError("`cache_position` must be provided for SlidingWindowCache.")
 
-        self.key_cache[layer_idx] += k_out
-        self.value_cache[layer_idx] += v_out
+        key_states = key_states.to(self.key_cache[layer_idx].dtype)
+        value_states = value_states.to(self.value_cache[layer_idx].dtype)
 
-        return k_out, v_out
+        return _sliding_cache_update_logic(
+            self.key_cache[layer_idx],
+            self.value_cache[layer_idx],
+            key_states,
+            value_states,
+            cache_position,
+            self.max_cache_len,
+        )
 
     def get_max_cache_shape(self) -> Optional[int]:
         return self.max_cache_len
@@ -1678,12 +1739,13 @@ def __init__(
         super().__init__()
         if not hasattr(config, "sliding_window") or config.sliding_window is None:
             raise ValueError(
-                "Setting `cache_implementation` to 'sliding_window' requires the model config supporting "
+                "Setting `cache_implementation` to 'hybrid' requires the model config supporting "
                 "sliding window attention, please check if there is a `sliding_window` field in the model "
                 "config and it's not set to None."
             )
-        self.max_cache_len = max_cache_len
-        self._sliding_window_max_len = min(config.sliding_window, max_cache_len)
+        self.max_cache_len = max_cache_len if max_cache_len is not None else config.max_position_embeddings
+        # Sliding layers can't be larger than the overall max cache len
+        self.sliding_window_len = min(config.sliding_window, self.max_cache_len)
         self.max_batch_size = max_batch_size
         # Some model define a custom `head_dim` != config.hidden_size // config.num_attention_heads
         self.head_dim = (
@@ -1696,18 +1758,11 @@ def __init__(
         )
 
         layer_switch = config.sliding_window_pattern if hasattr(config, "sliding_window_pattern") else 2  # 2 is for BC
-        self.is_sliding = torch.tensor(
-            [bool((i + 1) % layer_switch) for i in range(config.num_hidden_layers)], dtype=torch.bool
-        )
+        self.is_sliding_list = [bool((i + 1) % layer_switch) for i in range(config.num_hidden_layers)]
         self.key_cache: List[torch.Tensor] = []
         self.value_cache: List[torch.Tensor] = []
-        global_cache_shape = (self.max_batch_size, self.num_key_value_heads, max_cache_len, self.head_dim)
-        sliding_cache_shape = (
-            self.max_batch_size,
-            self.num_key_value_heads,
-            self._sliding_window_max_len,
-            self.head_dim,
-        )
+        global_cache_shape = (self.max_batch_size, self.num_key_value_heads, self.max_cache_len, self.head_dim)
+        sliding_cache_shape = (self.max_batch_size, self.num_key_value_heads, self.sliding_window_len, self.head_dim)
         device = torch.device(device) if device is not None else None
         for i in range(config.num_hidden_layers):
             if layer_device_map is not None:
@@ -1716,7 +1771,7 @@ def __init__(
                 layer_device = device
             # Note: `mark_static_address` is used to tag the cache as an fixed data pointer, preventing cuda graph
             # breaks when updating the cache.
-            cache_shape = global_cache_shape if not self.is_sliding[i] else sliding_cache_shape
+            cache_shape = sliding_cache_shape if self.is_sliding_list[i] else global_cache_shape
             new_layer_key_cache = torch.zeros(cache_shape, dtype=self._dtype, device=layer_device)
             new_layer_value_cache = torch.zeros(cache_shape, dtype=self._dtype, device=layer_device)
             torch._dynamo.mark_static_address(new_layer_key_cache)
@@ -1724,42 +1779,6 @@ def __init__(
             self.key_cache.append(new_layer_key_cache)
             self.value_cache.append(new_layer_value_cache)
 
-    def _sliding_update(self, cache_position, layer_idx, key_states, value_states, k_out, v_out, max_cache_len):
-        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)
-        cache_position = cache_position.clamp(0, max_cache_len - 1)
-        to_shift = cache_position > max_cache_len - 1
-        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
-
-    def _static_update(self, cache_position, layer_idx, key_states, value_states, k_out, v_out, max_cache_len):
-        k_out[:, :, cache_position] = key_states
-        v_out[:, :, cache_position] = value_states
-
-        self.key_cache[layer_idx] = k_out
-        self.value_cache[layer_idx] = v_out
-        return k_out, v_out
-
     def update(
         self,
         key_states: torch.Tensor,
@@ -1770,7 +1789,10 @@ def update(
         if cache_kwargs is None:
             cache_kwargs = {}
         cache_position = cache_kwargs.get("cache_position")
-        sliding_window = cache_kwargs.get("sliding_window")
+        if cache_position is None:
+            raise ValueError("`cache_position` must be provided for HybridCache.")
+
+        is_sliding_layer = self.is_sliding_list[layer_idx]
 
         # These two `if` blocks are only reached in multigpu and if `layer_device_map` is not passed. They are used
         # when the cache is initialized in the forward pass (e.g. Gemma2)
@@ -1779,25 +1801,22 @@ def update(
         if self.value_cache[layer_idx].device != value_states.device:
             self.value_cache[layer_idx] = self.value_cache[layer_idx].to(value_states.device)
 
-        k_out = self.key_cache[layer_idx]
-        v_out = self.value_cache[layer_idx]
-        key_states = key_states.to(k_out.dtype)
-        value_states = value_states.to(v_out.dtype)
-
-        if sliding_window:
-            update_fn = self._sliding_update
+        k_cache = self.key_cache[layer_idx]
+        v_cache = self.value_cache[layer_idx]
+        key_states = key_states.to(k_cache.dtype)
+        value_states = value_states.to(v_cache.dtype)
+
+        if is_sliding_layer:
+            return _sliding_cache_update_logic(
+                k_cache,
+                v_cache,
+                key_states,
+                value_states,
+                cache_position,
+                k_cache.shape[2],  # Use actual cache dim as max cache len
+            )
         else:
-            update_fn = self._static_update
-
-        return update_fn(
-            cache_position,
-            layer_idx,
-            key_states,
-            value_states,
-            k_out,
-            v_out,
-            k_out.shape[2],
-        )
+            return _static_cache_update_logic(k_cache, v_cache, key_states, value_states, cache_position)
 
     def get_max_cache_shape(self) -> Optional[int]:
         return self.max_cache_len
diff --git a/tests/utils/test_cache_utils.py b/tests/utils/test_cache_utils.py
index 691c5aa53578..73e86f03d70a 100644
--- a/tests/utils/test_cache_utils.py
+++ b/tests/utils/test_cache_utils.py
@@ -44,8 +44,11 @@
         Cache,
         ClvpForCausalLM,
         DynamicCache,
+        Gemma2Config,
         GenerationConfig,
+        HybridCache,
         LlamaConfig,
+        SlidingWindowCache,
         StaticCache,
         convert_and_export_with_cache,
     )
@@ -713,3 +716,352 @@ def test_hybrid_cache_exportability(self):
             dynamic_shapes=dynamic_shapes,
             strict=False,
         )
+
+
+class SyntheticCacheTest(unittest.TestCase):
+    """Tests cache behavior with simple dummy data."""
+
+    def setUp(self):
+        """Set up common configuration for all tests."""
+        self.window_size = 4
+        self.max_cache_len = 4
+        self.config = Gemma2Config(
+            num_hidden_layers=1,
+            num_key_value_heads=1,
+            num_attention_heads=1,
+            head_dim=1,
+            hidden_size=1,
+            sliding_window=self.window_size,
+            sliding_window_pattern=2,
+        )
+
+    def test_static_cache_within_bounds(self):
+        """
+        Test StaticCache preserves values at written positions within bounds.
+        Example (max_cache_len=4):
+        step 1 (pos 0): [1.0, 0.0, 0.0, 0.0]
+        step 2 (pos 1): [1.0, 2.0, 0.0, 0.0]
+        step 3 (pos 2): [1.0, 2.0, 3.0, 0.0]
+        step 4 (pos 3): [1.0, 2.0, 3.0, 4.0]
+        """
+        static_cache = StaticCache(config=self.config, max_batch_size=1, max_cache_len=self.max_cache_len)
+        expected_state = [0.0] * self.max_cache_len
+
+        for step in range(1, self.max_cache_len + 1):  # Test up to cache capacity
+            pos_idx = step - 1
+            value = float(step)
+
+            static_cache.update(
+                key_states=torch.tensor([[[[value]]]]),
+                value_states=torch.tensor([[[[value]]]]),  # Use same value for simplicity
+                layer_idx=0,
+                cache_kwargs={"cache_position": torch.tensor([pos_idx])},
+            )
+
+            expected_state[pos_idx] = value
+            k = static_cache.key_cache[0][0, 0, :, 0].tolist()
+            v = static_cache.value_cache[0][0, 0, :, 0].tolist()
+
+            self.assertEqual(k, expected_state, f"Static Key cache failed at step {step}")
+            self.assertEqual(v, expected_state, f"Static Value cache failed at step {step}")
+
+    def test_static_cache_out_of_bounds(self):
+        """Test StaticCache raises IndexError for out-of-bounds positions."""
+        static_cache = StaticCache(config=self.config, max_batch_size=1, max_cache_len=self.max_cache_len)
+        pos_out_of_bounds = torch.tensor([self.max_cache_len])  # Position >= max_cache_len
+
+        with self.assertRaises(IndexError):
+            static_cache.update(
+                key_states=torch.tensor([[[[1.0]]]]),
+                value_states=torch.tensor([[[[1.0]]]]),
+                layer_idx=0,
+                cache_kwargs={"cache_position": pos_out_of_bounds},
+            )
+
+    def test_sliding_window_cache(self):
+        """
+        Test SlidingWindowCache accumulates then slides.
+        Example (window_size=4):
+        step 4 (pos 3): [1.0, 2.0, 3.0, 4.0]
+        step 5 (pos 4): [2.0, 3.0, 4.0, 5.0] (shift happens as pos > window_size-1)
+        step 6 (pos 5): [3.0, 4.0, 5.0, 6.0]
+        """
+        sliding_cache = SlidingWindowCache(config=self.config, max_batch_size=1, max_cache_len=self.max_cache_len)
+        expected_state = [0.0] * self.window_size
+
+        for step in range(1, self.window_size * 2):  # Test beyond window size
+            pos_idx = step - 1
+            value = float(step)
+
+            sliding_cache.update(
+                key_states=torch.tensor([[[[value]]]]),
+                value_states=torch.tensor([[[[value]]]]),
+                layer_idx=0,
+                cache_kwargs={"cache_position": torch.tensor([pos_idx]), "sliding_window": self.window_size},
+            )
+
+            # Calculate expected state based on corrected sliding logic
+            clamped_pos = min(pos_idx, self.window_size - 1)
+            to_shift = pos_idx > self.window_size - 1
+
+            if to_shift:
+                expected_state = expected_state[1:] + [0.0]  # Shift left, add placeholder
+
+            expected_state[clamped_pos] = value  # Insert new value
+
+            # Only verify key cache for simplicity, assuming value is symmetrical
+            k = sliding_cache.key_cache[0][0, 0, :, 0].tolist()
+            self.assertEqual(k, expected_state, f"SlidingWindowCache failed at step {step}")
+
+    def test_hybrid_cache_static_mode(self):
+        """
+        Test HybridCache acts like StaticCache in the static layers.
+        Example (max_cache_len=4):
+        step 4 (pos 3): [1.0, 2.0, 3.0, 4.0]
+        """
+        config = copy.deepcopy(self.config)
+        config.sliding_window_pattern = 1  # Force layer 0 to be 1 % 1 = 0 (static)
+        hybrid_cache = HybridCache(config=config, max_batch_size=1, max_cache_len=self.max_cache_len)
+        expected_state = [0.0] * self.max_cache_len
+
+        for step in range(0, self.max_cache_len):  # Test up to cache size
+            hybrid_cache.update(
+                key_states=torch.tensor([[[[float(step + 1)]]]]),
+                value_states=torch.tensor([[[[float(step + 1)]]]]),
+                layer_idx=0,
+                cache_kwargs={"cache_position": torch.tensor([step])},
+            )
+
+            expected_state[step] = float(step + 1)
+            k = hybrid_cache.key_cache[0][0, 0, :, 0].tolist()
+            self.assertEqual(k, expected_state, f"HybridCache (static) failed at step {step}")
+
+    def test_hybrid_cache_sliding_mode(self):
+        """
+        Test HybridCache acts like SlidingWindowCache when 'sliding_window' is present.
+        Example (window_size=4):
+        step 4 (pos 3): [1.0, 2.0, 3.0, 4.0]
+        step 5 (pos 4): [2.0, 3.0, 4.0, 5.0]
+        step 6 (pos 5): [3.0, 4.0, 5.0, 6.0]
+        """
+        hybrid_cache = HybridCache(config=self.config, max_batch_size=1, max_cache_len=self.max_cache_len)
+        expected_state = [0.0] * self.window_size
+
+        for step in range(1, self.window_size * 2):  # Test beyond window size
+            pos_idx = step - 1
+            value = float(step)
+
+            hybrid_cache.update(
+                key_states=torch.tensor([[[[value]]]]),
+                value_states=torch.tensor([[[[value]]]]),
+                layer_idx=0,
+                cache_kwargs={"cache_position": torch.tensor([pos_idx]), "sliding_window": self.window_size},
+            )
+
+            clamped_pos = min(pos_idx, self.window_size - 1)
+            to_shift = pos_idx > self.window_size - 1
+            if to_shift:
+                expected_state = expected_state[1:] + [0.0]
+            expected_state[clamped_pos] = value
+
+            k = hybrid_cache.key_cache[0][0, 0, :, 0].tolist()
+            self.assertEqual(k, expected_state, f"HybridCache (sliding) failed at step {step}")
+
+    def test_sliding_window_cache_prompt_longer_than_max_cache_len(self):
+        """Test SlidingWindowCache when prompt length > max_cache_len (should keep only last max_cache_len tokens)."""
+        sliding_cache = SlidingWindowCache(config=self.config, max_batch_size=1, max_cache_len=self.max_cache_len)
+        prompt_len = self.max_cache_len + 2  # e.g., 6 if max_cache_len=4
+        values = [float(i + 1) for i in range(prompt_len)]
+        key_states = torch.tensor([[[[v] for v in values]]])  # shape (1,1,6,1)
+        value_states = torch.tensor([[[[v] for v in values]]])  # shape (1,1,6,1)
+        cache_position = torch.arange(prompt_len)
+        sliding_cache.update(
+            key_states=key_states,
+            value_states=value_states,
+            layer_idx=0,
+            cache_kwargs={"cache_position": cache_position, "sliding_window": self.window_size},
+        )
+        k = sliding_cache.key_cache[0][0, 0, :, 0].tolist()
+        v = sliding_cache.value_cache[0][0, 0, :, 0].tolist()
+        self.assertEqual(k, values[-self.window_size :], "SlidingWindowCache did not keep last window tokens")
+        self.assertEqual(v, values[-self.window_size :], "SlidingWindowCache did not keep last window tokens")
+
+    def test_hybrid_cache_prompt_longer_than_max_cache_len(self):
+        """Test HybridCache when prompt length > max_cache_len (should keep only last max_cache_len tokens in sliding mode)."""
+        hybrid_cache = HybridCache(config=self.config, max_batch_size=1, max_cache_len=self.max_cache_len)
+        prompt_len = self.max_cache_len + 2  # e.g., 6 if max_cache_len=4
+        values = [float(i + 1) for i in range(prompt_len)]
+        key_states = torch.tensor([[[[v] for v in values]]])  # shape (1,1,6,1)
+        value_states = torch.tensor([[[[v] for v in values]]])  # shape (1,1,6,1)
+        cache_position = torch.arange(prompt_len)
+        # Use sliding mode
+        hybrid_cache.update(
+            key_states=key_states,
+            value_states=value_states,
+            layer_idx=0,
+            cache_kwargs={"cache_position": cache_position, "sliding_window": self.window_size},
+        )
+        k = hybrid_cache.key_cache[0][0, 0, :, 0].tolist()
+        v = hybrid_cache.value_cache[0][0, 0, :, 0].tolist()
+        self.assertEqual(k, values[-self.window_size :], "HybridCache (sliding) did not keep last window tokens")
+        self.assertEqual(v, values[-self.window_size :], "HybridCache (sliding) did not keep last window tokens")
+
+    def test_static_cache_hardcoded(self):
+        """Test StaticCache with manually prefilled states and hardcoded assertions."""
+        # Scenario 1: Fill up to near capacity
+        static_cache = StaticCache(config=self.config, max_batch_size=1, max_cache_len=self.max_cache_len)
+        prefill = torch.tensor([[[[v] for v in [1.0, 2.0, 0.0, 0.0]]]], dtype=torch.float)
+        static_cache.key_cache[0].copy_(prefill)
+        static_cache.value_cache[0].copy_(prefill)
+        static_cache.update(
+            key_states=torch.tensor([[[[3.0]]]]),
+            value_states=torch.tensor([[[[3.0]]]]),
+            layer_idx=0,
+            cache_kwargs={"cache_position": torch.tensor([2])},
+        )
+        self.assertEqual(
+            static_cache.key_cache[0][0, 0, :, 0].tolist(), [1.0, 2.0, 3.0, 0.0], "StaticCache Scenario 1 failed"
+        )
+
+        # Scenario 2: Fill to capacity
+        static_cache = StaticCache(config=self.config, max_batch_size=1, max_cache_len=self.max_cache_len)
+        prefill = torch.tensor([[[[v] for v in [1.0, 2.0, 3.0, 0.0]]]], dtype=torch.float)
+        static_cache.key_cache[0].copy_(prefill)
+        static_cache.value_cache[0].copy_(prefill)
+        static_cache.update(
+            key_states=torch.tensor([[[[4.0]]]]),
+            value_states=torch.tensor([[[[4.0]]]]),
+            layer_idx=0,
+            cache_kwargs={"cache_position": torch.tensor([3])},
+        )
+        self.assertEqual(
+            static_cache.key_cache[0][0, 0, :, 0].tolist(), [1.0, 2.0, 3.0, 4.0], "StaticCache Scenario 2 failed"
+        )
+
+    def test_sliding_window_cache_hardcode(self):
+        """Test SlidingWindowCache with manually prefilled states and hardcoded assertions."""
+        # Scenario 1: Update within window, no slide yet
+        sliding_cache = SlidingWindowCache(config=self.config, max_batch_size=1, max_cache_len=self.max_cache_len)
+        prefill = torch.tensor([[[[v] for v in [1.0, 2.0, 0.0, 0.0]]]], dtype=torch.float)
+        sliding_cache.key_cache[0].copy_(prefill)
+        sliding_cache.value_cache[0].copy_(prefill)
+        sliding_cache.update(
+            key_states=torch.tensor([[[[3.0]]]]),
+            value_states=torch.tensor([[[[3.0]]]]),
+            layer_idx=0,
+            cache_kwargs={"cache_position": torch.tensor([2]), "sliding_window": self.window_size},
+        )
+        self.assertEqual(
+            sliding_cache.key_cache[0][0, 0, :, 0].tolist(),
+            [1.0, 2.0, 3.0, 0.0],
+            "SlidingWindowCache Scenario 1 failed",
+        )
+
+        # Scenario 2: Update causing slide
+        sliding_cache = SlidingWindowCache(config=self.config, max_batch_size=1, max_cache_len=self.max_cache_len)
+        prefill = torch.tensor([[[[v] for v in [1.0, 2.0, 3.0, 4.0]]]], dtype=torch.float)
+        sliding_cache.key_cache[0].copy_(prefill)
+        sliding_cache.value_cache[0].copy_(prefill)
+        sliding_cache.update(
+            key_states=torch.tensor([[[[5.0]]]]),
+            value_states=torch.tensor([[[[5.0]]]]),
+            layer_idx=0,
+            cache_kwargs={"cache_position": torch.tensor([4]), "sliding_window": self.window_size},
+        )
+        self.assertEqual(
+            sliding_cache.key_cache[0][0, 0, :, 0].tolist(),
+            [2.0, 3.0, 4.0, 5.0],
+            "SlidingWindowCache Scenario 2 failed",
+        )
+
+    def test_hybrid_cache_static_mode_hardcoded(self):
+        """Test HybridCache in static mode with hardcoded assertions."""
+        config = copy.deepcopy(self.config)
+        config.sliding_window_pattern = 1  # Layer 0 is static (1 % 1 == 0)
+
+        # Scenario 1
+        hybrid_cache = HybridCache(config=config, max_batch_size=1, max_cache_len=self.max_cache_len)
+        prefill = torch.tensor([[[[v] for v in [1.0, 2.0, 0.0, 0.0]]]], dtype=torch.float)
+        hybrid_cache.key_cache[0].copy_(prefill)
+        hybrid_cache.value_cache[0].copy_(prefill)
+        hybrid_cache.update(
+            key_states=torch.tensor([[[[3.0]]]]),
+            value_states=torch.tensor([[[[3.0]]]]),
+            layer_idx=0,
+            cache_kwargs={"cache_position": torch.tensor([2])},
+        )
+        self.assertEqual(
+            hybrid_cache.key_cache[0][0, 0, :, 0].tolist(),
+            [1.0, 2.0, 3.0, 0.0],
+            "HybridCache Static Scenario 1 failed",
+        )
+
+        # Scenario 2
+        hybrid_cache = HybridCache(config=config, max_batch_size=1, max_cache_len=self.max_cache_len)
+        prefill = torch.tensor([[[[v] for v in [1.0, 2.0, 3.0, 0.0]]]], dtype=torch.float)
+        hybrid_cache.key_cache[0].copy_(prefill)
+        hybrid_cache.value_cache[0].copy_(prefill)
+        hybrid_cache.update(
+            key_states=torch.tensor([[[[4.0]]]]),
+            value_states=torch.tensor([[[[4.0]]]]),
+            layer_idx=0,
+            cache_kwargs={"cache_position": torch.tensor([3])},
+        )
+        self.assertEqual(
+            hybrid_cache.key_cache[0][0, 0, :, 0].tolist(),
+            [1.0, 2.0, 3.0, 4.0],
+            "HybridCache Static Scenario 2 failed",
+        )
+
+    def test_hybrid_cache_sliding_mode_hardcoded(self):
+        """Test HybridCache in sliding mode with hardcoded assertions."""
+        # Scenario 1: Update within window, no slide yet
+        hybrid_cache = HybridCache(config=self.config, max_batch_size=1, max_cache_len=self.max_cache_len)
+        prefill = torch.tensor([[[[v] for v in [1.0, 2.0, 0.0, 0.0]]]], dtype=torch.float)
+        hybrid_cache.key_cache[0].copy_(prefill)
+        hybrid_cache.value_cache[0].copy_(prefill)
+        hybrid_cache.update(
+            key_states=torch.tensor([[[[3.0]]]]),
+            value_states=torch.tensor([[[[3.0]]]]),
+            layer_idx=0,
+            cache_kwargs={"cache_position": torch.tensor([2]), "sliding_window": self.window_size},
+        )
+        self.assertEqual(
+            hybrid_cache.key_cache[0][0, 0, :, 0].tolist(),
+            [1.0, 2.0, 3.0, 0.0],
+            "HybridCache Sliding Scenario 1 failed",
+        )
+
+        # Scenario 2: Update causing first slide
+        hybrid_cache = HybridCache(config=self.config, max_batch_size=1, max_cache_len=self.max_cache_len)
+        prefill = torch.tensor([[[[v] for v in [1.0, 2.0, 3.0, 4.0]]]], dtype=torch.float)
+        hybrid_cache.key_cache[0].copy_(prefill)
+        hybrid_cache.value_cache[0].copy_(prefill)
+        hybrid_cache.update(
+            key_states=torch.tensor([[[[5.0]]]]),
+            value_states=torch.tensor([[[[5.0]]]]),
+            layer_idx=0,
+            cache_kwargs={"cache_position": torch.tensor([4]), "sliding_window": self.window_size},
+        )
+        self.assertEqual(
+            hybrid_cache.key_cache[0][0, 0, :, 0].tolist(),
+            [2.0, 3.0, 4.0, 5.0],
+            "HybridCache Sliding Scenario 2 failed",
+        )
+
+        # Scenario 3: Update causing subsequent slide
+        hybrid_cache = HybridCache(config=self.config, max_batch_size=1, max_cache_len=self.max_cache_len)
+        prefill = torch.tensor([[[[v] for v in [2.0, 3.0, 4.0, 5.0]]]], dtype=torch.float)
+        hybrid_cache.key_cache[0].copy_(prefill)
+        hybrid_cache.value_cache[0].copy_(prefill)
+        hybrid_cache.update(
+            key_states=torch.tensor([[[[6.0]]]]),
+            value_states=torch.tensor([[[[6.0]]]]),
+            layer_idx=0,
+            cache_kwargs={"cache_position": torch.tensor([5]), "sliding_window": self.window_size},
+        )
+        self.assertEqual(
+            hybrid_cache.key_cache[0][0, 0, :, 0].tolist(),
+            [3.0, 4.0, 5.0, 6.0],
+            "HybridCache Sliding Scenario 3 failed",
+        )


Code review comments:
diff_hunk: @@ -916,10 +901,167 @@ def test_hybrid_cache_prompt_longer_than_max_cache_len(self):
             layer_idx=0,
             cache_kwargs={"cache_position": cache_position, "sliding_window": self.window_size},
         )
-        stored_keys, stored_values = self._extract_both_caches(hybrid_cache)
+        k = hybrid_cache.key_cache[0][0, 0, :, 0].tolist()
+        v = hybrid_cache.value_cache[0][0, 0, :, 0].tolist()
+        self.assertEqual(k, values[-self.window_size :], "HybridCache (sliding) did not keep last window tokens")
+        self.assertEqual(v, values[-self.window_size :], "HybridCache (sliding) did not keep last window tokens")
+
+    def test_static_cache_hardcoded(self):
+        """Test StaticCache with manually prefilled states and hardcoded assertions."""
+        # Scenario 1: Fill up to near capacity
+        static_cache = StaticCache(config=self.config, max_batch_size=1, max_cache_len=self.max_cache_len)
+        prefill = torch.tensor([[[[v] for v in [1.0, 2.0, 0.0, 0.0]]]], dtype=torch.float)
path: tests/utils/test_cache_utils.py
text: @user1:
perhaps something like `torch.tensor([1.0, 2.0, 0.0, 0.0])[None, None, :, None]`? (more in line with how we expand dimensions throughout our codebase)

diff_hunk: @@ -916,10 +901,167 @@ def test_hybrid_cache_prompt_longer_than_max_cache_len(self):
             layer_idx=0,
             cache_kwargs={"cache_position": cache_position, "sliding_window": self.window_size},
         )
-        stored_keys, stored_values = self._extract_both_caches(hybrid_cache)
+        k = hybrid_cache.key_cache[0][0, 0, :, 0].tolist()
+        v = hybrid_cache.value_cache[0][0, 0, :, 0].tolist()
+        self.assertEqual(k, values[-self.window_size :], "HybridCache (sliding) did not keep last window tokens")
+        self.assertEqual(v, values[-self.window_size :], "HybridCache (sliding) did not keep last window tokens")
+
+    def test_static_cache_hardcoded(self):
+        """Test StaticCache with manually prefilled states and hardcoded assertions."""
+        # Scenario 1: Fill up to near capacity
+        static_cache = StaticCache(config=self.config, max_batch_size=1, max_cache_len=self.max_cache_len)
+        prefill = torch.tensor([[[[v] for v in [1.0, 2.0, 0.0, 0.0]]]], dtype=torch.float)
+        static_cache.key_cache[0].copy_(prefill)
+        static_cache.value_cache[0].copy_(prefill)
path: tests/utils/test_cache_utils.py
text: @user1:
let's use `static_cache.update()` instead for prefill too -> then we implicitly test that prefill is working as expected

diff_hunk: @@ -916,10 +901,167 @@ def test_hybrid_cache_prompt_longer_than_max_cache_len(self):
             layer_idx=0,
             cache_kwargs={"cache_position": cache_position, "sliding_window": self.window_size},
         )
-        stored_keys, stored_values = self._extract_both_caches(hybrid_cache)
+        k = hybrid_cache.key_cache[0][0, 0, :, 0].tolist()
+        v = hybrid_cache.value_cache[0][0, 0, :, 0].tolist()
+        self.assertEqual(k, values[-self.window_size :], "HybridCache (sliding) did not keep last window tokens")
+        self.assertEqual(v, values[-self.window_size :], "HybridCache (sliding) did not keep last window tokens")
+
+    def test_static_cache_hardcoded(self):
+        """Test StaticCache with manually prefilled states and hardcoded assertions."""
+        # Scenario 1: Fill up to near capacity
+        static_cache = StaticCache(config=self.config, max_batch_size=1, max_cache_len=self.max_cache_len)
+        prefill = torch.tensor([[[[v] for v in [1.0, 2.0, 0.0, 0.0]]]], dtype=torch.float)
+        static_cache.key_cache[0].copy_(prefill)
+        static_cache.value_cache[0].copy_(prefill)
+        static_cache.update(
+            key_states=torch.tensor([[[[3.0]]]]),
+            value_states=torch.tensor([[[[3.0]]]]),
+            layer_idx=0,
+            cache_kwargs={"cache_position": torch.tensor([2])},
+        )
+        self.assertEqual(
+            static_cache.key_cache[0][0, 0, :, 0].tolist(), [1.0, 2.0, 3.0, 0.0], "StaticCache Scenario 1 failed"
+        )
+
+        # Scenario 2: Fill to capacity
+        static_cache = StaticCache(config=self.config, max_batch_size=1, max_cache_len=self.max_cache_len)
+        prefill = torch.tensor([[[[v] for v in [1.0, 2.0, 3.0, 0.0]]]], dtype=torch.float)
+        static_cache.key_cache[0].copy_(prefill)
+        static_cache.value_cache[0].copy_(prefill)
path: tests/utils/test_cache_utils.py
text: @user1:
Reinitializing the cache (these 4 lines) can be removed, correct? 

(same for the other tests)