pimono_refactor_DynamicCache

diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py
index ac324ebb62..db19af4d36 100644
--- a/src/transformers/cache_utils.py
+++ b/src/transformers/cache_utils.py
@@ -1199,7 +1199,9 @@ class DynamicCache(Cache):
         config: PreTrainedConfig | None = None,
         offloading: bool = False,
         offload_only_non_sliding: bool = False,
+        is_stale: bool = False,
     ):
+        self.is_stale = is_stale
         layers = []
         # If a config is passed, use it to infer the layer types and initialize accordingly
         if config is not None:
diff --git a/src/transformers/models/cohere2/modeling_cohere2.py b/src/transformers/models/cohere2/modeling_cohere2.py
index f43b2a0ef4..7d9bc23d30 100644
--- a/src/transformers/models/cohere2/modeling_cohere2.py
+++ b/src/transformers/models/cohere2/modeling_cohere2.py
@@ -244,7 +244,15 @@ class Cohere2Attention(nn.Module):
             query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
 
         if past_key_values is not None:
-            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+            if isinstance(past_key_values, DynamicCache) and past_key_values.is_stale:
+                import warnings
+
+                warnings.warn(
+                    f"The cache is stale for layer {self.layer_idx}. Bypassing the internal cache update.",
+                    UserWarning,
+                )
+            else:
+                key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
 
         attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
             self.config._attn_implementation, eager_attention_forward
diff --git a/src/transformers/models/cohere2/modular_cohere2.py b/src/transformers/models/cohere2/modular_cohere2.py
index d19055a1b7..1f02ad26ac 100644
--- a/src/transformers/models/cohere2/modular_cohere2.py
+++ b/src/transformers/models/cohere2/modular_cohere2.py
@@ -194,7 +194,15 @@ class Cohere2Attention(CohereAttention):
             query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
 
         if past_key_values is not None:
-            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+            if isinstance(past_key_values, DynamicCache) and past_key_values.is_stale:
+                import warnings
+
+                warnings.warn(
+                    f"The cache is stale for layer {self.layer_idx}. Bypassing the internal cache update.",
+                    UserWarning,
+                )
+            else:
+                key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
 
         attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
             self.config._attn_implementation, eager_attention_forward
diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py
index ab998cc99c..8a6cd79bff 100644
--- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py
+++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py
@@ -451,7 +451,15 @@ class DeepseekV3Attention(nn.Module):
         key_states = torch.cat((k_pass, k_rot), dim=-1)
 
         if past_key_values is not None:
-            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+            if isinstance(past_key_values, DynamicCache) and past_key_values.is_stale:
+                import warnings
+
+                warnings.warn(
+                    f"The cache is stale for layer {self.layer_idx}. Bypassing the internal cache update.",
+                    UserWarning,
+                )
+            else:
+                key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
 
         if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim:
             value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim])
diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py
index 3c62a564a3..7ed3bd34a3 100644
--- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py
+++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py
@@ -6,7 +6,7 @@ import torch.nn.functional as F
 from torch import nn
 
 from ... import initialization as init
-from ...cache_utils import Cache
+from ...cache_utils import Cache, DynamicCache
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
 from ...modeling_layers import GenericForSequenceClassification, GenericForTokenClassification
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
@@ -256,7 +256,15 @@ class DeepseekV3Attention(nn.Module):
         key_states = torch.cat((k_pass, k_rot), dim=-1)
 
         if past_key_values is not None:
-            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+            if isinstance(past_key_values, DynamicCache) and past_key_values.is_stale:
+                import warnings
+
+                warnings.warn(
+                    f"The cache is stale for layer {self.layer_idx}. Bypassing the internal cache update.",
+                    UserWarning,
+                )
+            else:
+                key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
 
         if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim:
             value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim])
diff --git a/src/transformers/models/gemma4/modeling_gemma4.py b/src/transformers/models/gemma4/modeling_gemma4.py
index f690c0425c..8bcf501926 100644
--- a/src/transformers/models/gemma4/modeling_gemma4.py
+++ b/src/transformers/models/gemma4/modeling_gemma4.py
@@ -1212,12 +1212,20 @@ class Gemma4TextAttention(nn.Module):
             value_states = value_states.transpose(1, 2)
 
         if past_key_values is not None:
-            if not self.is_kv_shared_layer:
-                key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
-            if self.store_full_length_kv:
-                if not hasattr(past_key_values, "shared_layers"):
-                    past_key_values.shared_layers = {}
-                past_key_values.shared_layers[self.layer_idx] = key_states, value_states
+            if isinstance(past_key_values, DynamicCache) and past_key_values.is_stale:
+                import warnings
+
+                warnings.warn(
+                    f"The cache is stale for layer {self.layer_idx}. Bypassing the internal cache update.",
+                    UserWarning,
+                )
+            else:
+                if not self.is_kv_shared_layer:
+                    key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+                if self.store_full_length_kv:
+                    if not hasattr(past_key_values, "shared_layers"):
+                        past_key_values.shared_layers = {}
+                    past_key_values.shared_layers[self.layer_idx] = key_states, value_states
 
         attention_interface: Callable = eager_attention_forward
         if self.config._attn_implementation != "eager":
diff --git a/src/transformers/models/gemma4/modular_gemma4.py b/src/transformers/models/gemma4/modular_gemma4.py
index a972738022..2b70f31332 100644
--- a/src/transformers/models/gemma4/modular_gemma4.py
+++ b/src/transformers/models/gemma4/modular_gemma4.py
@@ -985,12 +985,20 @@ class Gemma4TextAttention(nn.Module):
             value_states = value_states.transpose(1, 2)
 
         if past_key_values is not None:
-            if not self.is_kv_shared_layer:
-                key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
-            if self.store_full_length_kv:
-                if not hasattr(past_key_values, "shared_layers"):
-                    past_key_values.shared_layers = {}
-                past_key_values.shared_layers[self.layer_idx] = key_states, value_states
+            if isinstance(past_key_values, DynamicCache) and past_key_values.is_stale:
+                import warnings
+
+                warnings.warn(
+                    f"The cache is stale for layer {self.layer_idx}. Bypassing the internal cache update.",
+                    UserWarning,
+                )
+            else:
+                if not self.is_kv_shared_layer:
+                    key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+                if self.store_full_length_kv:
+                    if not hasattr(past_key_values, "shared_layers"):
+                        past_key_values.shared_layers = {}
+                    past_key_values.shared_layers[self.layer_idx] = key_states, value_states
 
         attention_interface: Callable = eager_attention_forward
         if self.config._attn_implementation != "eager":
diff --git a/src/transformers/models/llama4/modeling_llama4.py b/src/transformers/models/llama4/modeling_llama4.py
index 08d50bd63f..ba642d5d83 100644
--- a/src/transformers/models/llama4/modeling_llama4.py
+++ b/src/transformers/models/llama4/modeling_llama4.py
@@ -389,7 +389,15 @@ class Llama4TextAttention(nn.Module):
         key_states = key_states.transpose(1, 2)
 
         if past_key_values is not None:
-            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+            if isinstance(past_key_values, DynamicCache) and past_key_values.is_stale:
+                import warnings
+
+                warnings.warn(
+                    f"The cache is stale for layer {self.layer_idx}. Bypassing the internal cache update.",
+                    UserWarning,
+                )
+            else:
+                key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
 
         attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
             self.config._attn_implementation, eager_attention_forward
diff --git a/src/transformers/models/ministral3/modeling_ministral3.py b/src/transformers/models/ministral3/modeling_ministral3.py
index 6aacf4c8ce..4701859cbb 100644
--- a/src/transformers/models/ministral3/modeling_ministral3.py
+++ b/src/transformers/models/ministral3/modeling_ministral3.py
@@ -150,7 +150,15 @@ class Ministral3Attention(nn.Module):
         ).to(query_states.dtype)
 
         if past_key_values is not None:
-            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+            if isinstance(past_key_values, DynamicCache) and past_key_values.is_stale:
+                import warnings
+
+                warnings.warn(
+                    f"The cache is stale for layer {self.layer_idx}. Bypassing the internal cache update.",
+                    UserWarning,
+                )
+            else:
+                key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
 
         attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
             self.config._attn_implementation, eager_attention_forward
diff --git a/src/transformers/models/ministral3/modular_ministral3.py b/src/transformers/models/ministral3/modular_ministral3.py
index 0a925f5d40..0a2bde879a 100644
--- a/src/transformers/models/ministral3/modular_ministral3.py
+++ b/src/transformers/models/ministral3/modular_ministral3.py
@@ -2,7 +2,7 @@ from collections.abc import Callable
 
 import torch
 
-from ...cache_utils import Cache
+from ...cache_utils import Cache, DynamicCache
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
 from ...modeling_layers import (
     GenericForQuestionAnswering,
@@ -57,7 +57,15 @@ class Ministral3Attention(MistralAttention):
         ).to(query_states.dtype)
 
         if past_key_values is not None:
-            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+            if isinstance(past_key_values, DynamicCache) and past_key_values.is_stale:
+                import warnings
+
+                warnings.warn(
+                    f"The cache is stale for layer {self.layer_idx}. Bypassing the internal cache update.",
+                    UserWarning,
+                )
+            else:
+                key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
 
         attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
             self.config._attn_implementation, eager_attention_forward
diff --git a/src/transformers/models/mistral4/modeling_mistral4.py b/src/transformers/models/mistral4/modeling_mistral4.py
index 006ddad187..c261b803ac 100644
--- a/src/transformers/models/mistral4/modeling_mistral4.py
+++ b/src/transformers/models/mistral4/modeling_mistral4.py
@@ -459,7 +459,15 @@ class Mistral4Attention(nn.Module):
         ).to(query_states.dtype)
 
         if past_key_values is not None:
-            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+            if isinstance(past_key_values, DynamicCache) and past_key_values.is_stale:
+                import warnings
+
+                warnings.warn(
+                    f"The cache is stale for layer {self.layer_idx}. Bypassing the internal cache update.",
+                    UserWarning,
+                )
+            else:
+                key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
 
         if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim:
             value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim])
diff --git a/src/transformers/models/mistral4/modular_mistral4.py b/src/transformers/models/mistral4/modular_mistral4.py
index acd9f1f601..5efbe66e83 100644
--- a/src/transformers/models/mistral4/modular_mistral4.py
+++ b/src/transformers/models/mistral4/modular_mistral4.py
@@ -18,7 +18,7 @@ import torch.nn.functional as F
 from torch import nn
 
 from ... import initialization as init
-from ...cache_utils import Cache
+from ...cache_utils import Cache, DynamicCache
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
 from ...modeling_layers import GenericForSequenceClassification, GenericForTokenClassification
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
@@ -191,7 +191,15 @@ class Mistral4Attention(DeepseekV3Attention):
         ).to(query_states.dtype)
 
         if past_key_values is not None:
-            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+            if isinstance(past_key_values, DynamicCache) and past_key_values.is_stale:
+                import warnings
+
+                warnings.warn(
+                    f"The cache is stale for layer {self.layer_idx}. Bypassing the internal cache update.",
+                    UserWarning,
+                )
+            else:
+                key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
 
         if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim:
             value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim])
diff --git a/src/transformers/models/olmo3/modeling_olmo3.py b/src/transformers/models/olmo3/modeling_olmo3.py
index 5baa8e5f24..40ba42adc2 100644
--- a/src/transformers/models/olmo3/modeling_olmo3.py
+++ b/src/transformers/models/olmo3/modeling_olmo3.py
@@ -185,7 +185,15 @@ class Olmo3Attention(nn.Module):
         query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
 
         if past_key_values is not None:
-            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+            if isinstance(past_key_values, DynamicCache) and past_key_values.is_stale:
+                import warnings
+
+                warnings.warn(
+                    f"The cache is stale for layer {self.layer_idx}. Bypassing the internal cache update.",
+                    UserWarning,
+                )
+            else:
+                key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
 
         attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
             self.config._attn_implementation, eager_attention_forward
diff --git a/src/transformers/models/olmo3/modular_olmo3.py b/src/transformers/models/olmo3/modular_olmo3.py
index 81325db7cc..11fb68b5ea 100644
--- a/src/transformers/models/olmo3/modular_olmo3.py
+++ b/src/transformers/models/olmo3/modular_olmo3.py
@@ -126,7 +126,15 @@ class Olmo3Attention(Olmo2Attention):
         query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
 
         if past_key_values is not None:
-            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+            if isinstance(past_key_values, DynamicCache) and past_key_values.is_stale:
+                import warnings
+
+                warnings.warn(
+                    f"The cache is stale for layer {self.layer_idx}. Bypassing the internal cache update.",
+                    UserWarning,
+                )
+            else:
+                key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
 
         attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
             self.config._attn_implementation, eager_attention_forward
diff --git a/src/transformers/models/qwen3/modeling_qwen3.py b/src/transformers/models/qwen3/modeling_qwen3.py
index 91715a33cf..c62b908c8d 100644
--- a/src/transformers/models/qwen3/modeling_qwen3.py
+++ b/src/transformers/models/qwen3/modeling_qwen3.py
@@ -268,7 +268,15 @@ class Qwen3Attention(nn.Module):
         query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
 
         if past_key_values is not None:
-            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+            if isinstance(past_key_values, DynamicCache) and past_key_values.is_stale:
+                import warnings
+
+                warnings.warn(
+                    f"The cache is stale for layer {self.layer_idx}. Bypassing the internal cache update.",
+                    UserWarning,
+                )
+            else:
+                key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
 
         attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
             self.config._attn_implementation, eager_attention_forward
diff --git a/src/transformers/models/qwen3/modular_qwen3.py b/src/transformers/models/qwen3/modular_qwen3.py
index 73cde6d89a..286d545e0c 100644
--- a/src/transformers/models/qwen3/modular_qwen3.py
+++ b/src/transformers/models/qwen3/modular_qwen3.py
@@ -17,7 +17,7 @@ from collections.abc import Callable
 
 import torch
 
-from ...cache_utils import Cache
+from ...cache_utils import Cache, DynamicCache
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
 from ...modeling_outputs import CausalLMOutputWithPast
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS
@@ -84,7 +84,15 @@ class Qwen3Attention(LlamaAttention):
         query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
 
         if past_key_values is not None:
-            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+            if isinstance(past_key_values, DynamicCache) and past_key_values.is_stale:
+                import warnings
+
+                warnings.warn(
+                    f"The cache is stale for layer {self.layer_idx}. Bypassing the internal cache update.",
+                    UserWarning,
+                )
+            else:
+                key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
 
         attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
             self.config._attn_implementation, eager_attention_forward