back to storage inside Cache()

Author: manueldeprada

docs/source/en/cache_explanation.md

@@ -89,8 +89,8 @@ Layers can be of different types (e.g. `DynamicLayer`, `StaticLayer`, `SlidingWi
 The simplest is a `DynamicLayer` that grows as more tokens are processed. The sequence length dimension (`seq_len`) increases with each new token:
 
 ```py
-cache.layers[idx].keys = torch.cat([cache.layers[idx].keys, key_states], dim=-2)
-cache.layers[idx].values = torch.cat([cache.layers[idx].values, value_states], dim=-2)
+cache.key_cache[idx] = torch.cat([cache.key_cache[idx], key_states], dim=-2)
+cache.value_cache[idx] = torch.cat([cache.value_cache[idx], value_states], dim=-2)
 ```
 
 Other layers like `StaticLayer` and `SlidingWindowLayer` have a fixed sequence length that is set when the cache is created. This makes them compatible with `torch.compile`. In the case of `SlidingWindowLayer`, existing tokens are shifted out of the cache when a new token is added.

src/transformers/cache_utils.py

@@ -282,8 +282,8 @@ def __init__(self, model: PreTrainedModel):
             dtype=self.model.dtype,
         )
         for i in range(len(self.static_cache)):
-            self.register_buffer(f"key_cache_{i}", self.static_cache.layers[i].keys, persistent=False)
-            self.register_buffer(f"value_cache_{i}", self.static_cache.layers[i].values, persistent=False)
+            self.register_buffer(f"key_cache_{i}", self.static_cache.key_cache[i], persistent=False)
+            self.register_buffer(f"value_cache_{i}", self.static_cache.value_cache[i], persistent=False)
 
     def forward(self, input_ids: torch.Tensor, cache_position: torch.Tensor):
         """
@@ -413,8 +413,8 @@ def __init__(
 
         # Register all key and value cache tensors as buffers
         for i in range(len(self.cache)):
-            self.register_buffer(f"key_cache_{i}", self.cache.layers[i].keys, persistent=False)
-            self.register_buffer(f"value_cache_{i}", self.cache.layers[i].values, persistent=False)
+            self.register_buffer(f"key_cache_{i}", self.cache.key_cache[i], persistent=False)
+            self.register_buffer(f"value_cache_{i}", self.cache.value_cache[i], persistent=False)
 
     def forward(
         self,
@@ -559,8 +559,8 @@ def __init__(self, model, max_static_cache_length, batch_size):
 
         # Register cache buffers to make them exportable
         for i in range(len(self.static_cache)):
-            self.register_buffer(f"key_cache_{i}", self.static_cache.layers[i].keys, persistent=False)
-            self.register_buffer(f"value_cache_{i}", self.static_cache.layers[i].values, persistent=False)
+            self.register_buffer(f"key_cache_{i}", self.static_cache.key_cache[i], persistent=False)
+            self.register_buffer(f"value_cache_{i}", self.static_cache.value_cache[i], persistent=False)
 
     def forward(self, decoder_input_ids, encoder_hidden_states, cache_position):
         # Get outputs from decoder

src/transformers/integrations/executorch.py

@@ -230,8 +230,8 @@ def forward(
         current_states = key_value_states if is_cross_attention else hidden_states
         if is_cross_attention and past_key_value is not None and is_updated:
             # reuse k,v, cross_attentions
-            key_states = curr_past_key_value.layers[self.layer_idx].keys
-            value_states = curr_past_key_value.layers[self.layer_idx].values
+            key_states = curr_past_key_value.key_cache[self.layer_idx]
+            value_states = curr_past_key_value.value_cache[self.layer_idx]
         else:
             key_states = self.k_proj(current_states)
             value_states = self.v_proj(current_states)

src/transformers/models/bart/modeling_bart.py

@@ -1293,8 +1293,8 @@ def forward(
         current_states = key_value_states if is_cross_attention else hidden_states
         if is_cross_attention and past_key_value is not None and is_updated:
             # reuse k,v, cross_attentions
-            key_states = curr_past_key_value.layers[self.layer_idx].keys
-            value_states = curr_past_key_value.layers[self.layer_idx].values
+            key_states = curr_past_key_value.key_cache[self.layer_idx]
+            value_states = curr_past_key_value.value_cache[self.layer_idx]
         else:
             key_states = self.k_proj(current_states)
             value_states = self.v_proj(current_states)

src/transformers/models/bigbird_pegasus/modeling_bigbird_pegasus.py

@@ -207,8 +207,8 @@ def forward(
         current_states = key_value_states if is_cross_attention else hidden_states
         if is_cross_attention and past_key_value is not None and is_updated:
             # reuse k,v, cross_attentions
-            key_states = curr_past_key_value.layers[self.layer_idx].keys
-            value_states = curr_past_key_value.layers[self.layer_idx].values
+            key_states = curr_past_key_value.key_cache[self.layer_idx]
+            value_states = curr_past_key_value.value_cache[self.layer_idx]
         else:
             key_states = self.k_proj(current_states)
             value_states = self.v_proj(current_states)

src/transformers/models/biogpt/modeling_biogpt.py

@@ -229,8 +229,8 @@ def forward(
         current_states = key_value_states if is_cross_attention else hidden_states
         if is_cross_attention and past_key_value is not None and is_updated:
             # reuse k,v, cross_attentions
-            key_states = curr_past_key_value.layers[self.layer_idx].keys
-            value_states = curr_past_key_value.layers[self.layer_idx].values
+            key_states = curr_past_key_value.key_cache[self.layer_idx]
+            value_states = curr_past_key_value.value_cache[self.layer_idx]
         else:
             key_states = self.k_proj(current_states)
             value_states = self.v_proj(current_states)

src/transformers/models/blenderbot/modeling_blenderbot.py

@@ -213,8 +213,8 @@ def forward(
         current_states = key_value_states if is_cross_attention else hidden_states
         if is_cross_attention and past_key_value is not None and is_updated:
             # reuse k,v, cross_attentions
-            key_states = curr_past_key_value.layers[self.layer_idx].keys
-            value_states = curr_past_key_value.layers[self.layer_idx].values
+            key_states = curr_past_key_value.key_cache[self.layer_idx]
+            value_states = curr_past_key_value.value_cache[self.layer_idx]
         else:
             key_states = self.k_proj(current_states)
             value_states = self.v_proj(current_states)

src/transformers/models/blenderbot_small/modeling_blenderbot_small.py

@@ -356,8 +356,8 @@ def forward(
         is_updated = past_key_values.is_updated.get(self.layer_idx) if past_key_values is not None else False
         if past_key_values is not None and is_updated:
             # reuse k,v, cross_attentions
-            key_states = past_key_values.cross_attention_cache.layers[self.layer_idx].keys
-            value_states = past_key_values.cross_attention_cache.layers[self.layer_idx].values
+            key_states = past_key_values.cross_attention_cache.key_cache[self.layer_idx]
+            value_states = past_key_values.cross_attention_cache.value_cache[self.layer_idx]
         else:
             key_states = self.k_proj(cross_attention_states).view(cross_shape).transpose(1, 2)
             value_states = self.v_proj(cross_attention_states).view(cross_shape).transpose(1, 2)

src/transformers/models/dia/modeling_dia.py

@@ -182,8 +182,8 @@ def forward(
         is_updated = past_key_values.is_updated.get(self.layer_idx) if past_key_values is not None else False
         if past_key_values is not None and is_updated:
             # reuse k,v, cross_attentions
-            key_states = past_key_values.cross_attention_cache.layers[self.layer_idx].keys
-            value_states = past_key_values.cross_attention_cache.layers[self.layer_idx].values
+            key_states = past_key_values.cross_attention_cache.key_cache[self.layer_idx]
+            value_states = past_key_values.cross_attention_cache.value_cache[self.layer_idx]
         else:
             key_states = self.k_proj(cross_attention_states).view(cross_shape).transpose(1, 2)
             value_states = self.v_proj(cross_attention_states).view(cross_shape).transpose(1, 2)