portugalfuturista/replica-omnisciente
2
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 73
replica-omnisciente/dirac/evals/roo/roo_code_refactor_DynamicCache

226 lines
12 KiB
Text
Raw Blame History

diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py
index ac324ebb62..2b8b3f4f4e 100644
--- a/src/transformers/cache_utils.py
+++ b/src/transformers/cache_utils.py
@@ -1200,6 +1200,7 @@ class DynamicCache(Cache):
offloading: bool = False,
offload_only_non_sliding: bool = False,
):
+ self.is_stale = False
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/modular_cohere2.py b/src/transformers/models/cohere2/modular_cohere2.py
index d19055a1b7..0be1bdc155 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(
+ "The cache is stale and will not be updated. Please reset the cache or provide a new one.",
+ 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/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py
index 3c62a564a3..72e8c31900 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(
+ "The cache is stale and will not be updated. Please reset the cache or provide a new one.",
+ 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/modular_gemma4.py b/src/transformers/models/gemma4/modular_gemma4.py
index a972738022..ab1f736d35 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(
+ "The cache is stale and will not be updated. Please reset the cache or provide a new one.",
+ 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..c8c8f7416f 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(
+ "The cache is stale and will not be updated. Please reset the cache or provide a new one.",
+ 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..2ad3c65d4d 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(
+ "The cache is stale and will not be updated. Please reset the cache or provide a new one.",
+ 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/modular_mistral4.py b/src/transformers/models/mistral4/modular_mistral4.py
index acd9f1f601..856d2ac3a8 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(
+ "The cache is stale and will not be updated. Please reset the cache or provide a new one.",
+ 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/modular_olmo3.py b/src/transformers/models/olmo3/modular_olmo3.py
index 81325db7cc..ca35600efc 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(
+ "The cache is stale and will not be updated. Please reset the cache or provide a new one.",
+ 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..179a7bf4e8 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(
+ "The cache is stale and will not be updated. Please reset the cache or provide a new one.",
+ 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
Reference in a new issue View git blame Copy permalink