175 lines
9.6 KiB
Text
175 lines
9.6 KiB
Text
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/modeling_cohere2.py b/src/transformers/models/cohere2/modeling_cohere2.py
|
|
index f43b2a0ef4..bc89faa454 100644
|
|
--- a/src/transformers/models/cohere2/modeling_cohere2.py
|
|
+++ b/src/transformers/models/cohere2/modeling_cohere2.py
|
|
@@ -244,7 +244,14 @@ 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(
|
|
+ "The provided DynamicCache has been marked as stale. Bypassing the cache update for this layer.",
|
|
+ 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..0ce462c311 100644
|
|
--- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py
|
|
+++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py
|
|
@@ -451,7 +451,14 @@ 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 provided DynamicCache has been marked as stale. Bypassing the cache update for this layer.",
|
|
+ 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..ed9b5a6c96 100644
|
|
--- a/src/transformers/models/gemma4/modular_gemma4.py
|
|
+++ b/src/transformers/models/gemma4/modular_gemma4.py
|
|
@@ -985,8 +985,16 @@ 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 isinstance(past_key_values, DynamicCache) and past_key_values.is_stale:
|
|
+ import warnings
|
|
+
|
|
+ warnings.warn(
|
|
+ "The provided DynamicCache has been marked as stale. Bypassing the cache update for this layer.",
|
|
+ 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 = {}
|
|
diff --git a/src/transformers/models/llama4/modeling_llama4.py b/src/transformers/models/llama4/modeling_llama4.py
|
|
index 08d50bd63f..7e6d1ff804 100644
|
|
--- a/src/transformers/models/llama4/modeling_llama4.py
|
|
+++ b/src/transformers/models/llama4/modeling_llama4.py
|
|
@@ -389,7 +389,14 @@ 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 provided DynamicCache has been marked as stale. Bypassing the cache update for this layer.",
|
|
+ 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..006f4d2d33 100644
|
|
--- a/src/transformers/models/ministral3/modeling_ministral3.py
|
|
+++ b/src/transformers/models/ministral3/modeling_ministral3.py
|
|
@@ -150,7 +150,14 @@ 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(
|
|
+ "The provided DynamicCache has been marked as stale. Bypassing the cache update for this layer.",
|
|
+ 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..05cf8ba273 100644
|
|
--- a/src/transformers/models/mistral4/modeling_mistral4.py
|
|
+++ b/src/transformers/models/mistral4/modeling_mistral4.py
|
|
@@ -459,7 +459,14 @@ 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(
|
|
+ "The provided DynamicCache has been marked as stale. Bypassing the cache update for this layer.",
|
|
+ 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..319d4458de 100644
|
|
--- a/src/transformers/models/olmo3/modeling_olmo3.py
|
|
+++ b/src/transformers/models/olmo3/modeling_olmo3.py
|
|
@@ -185,7 +185,14 @@ 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(
|
|
+ "The provided DynamicCache has been marked as stale. Bypassing the cache update for this layer.",
|
|
+ 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..3398e9b879 100644
|
|
--- a/src/transformers/models/qwen3/modeling_qwen3.py
|
|
+++ b/src/transformers/models/qwen3/modeling_qwen3.py
|
|
@@ -268,7 +268,14 @@ 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(
|
|
+ "The provided DynamicCache has been marked as stale. Bypassing the cache update for this layer.",
|
|
+ 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
|