Repository navigation
Expand file tree
/
Copy pathohmypi_refactor_DynamicCache
More file actions
175 lines (160 loc) · 9.64 KB
/
Copy pathohmypi_refactor_DynamicCache
File metadata and controls
175 lines (160 loc) · 9.64 KB
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
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
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