[cache] Allow sliding window layers to be roll-backed for speculative decoding

src/transformers/cache_utils.py

@@ -200,6 +200,14 @@ def __init__(self, sliding_window: int, **kwargs):
         self.sliding_window = sliding_window
         self.cumulative_length = 0
         self._sliding_window_tensor = torch.tensor(self.sliding_window, dtype=torch.long)
+        self.record_past = False
+
+    def activate_past_recording(self):
+        """
+        Calling this function will activate past state recording, meaning that a call to `update` will wait for a call to `crop`
+        before restricting the size of the `k/v_states` to `sliding_window`, to be able to retrieve previous full states.
+        """
+        self.record_past = True
 
     def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tensor) -> None:
         super().lazy_initialization(key_states, value_states)
@@ -228,8 +236,13 @@ def update(
         full_key_states = torch.cat([self.keys, key_states], dim=-2)
         full_value_states = torch.cat([self.values, value_states], dim=-2)
         # Only cache the last `self.sliding_window - 1` tokens (or all of them if lower than that)
-        self.keys = full_key_states[:, :, -self.sliding_window + 1 :, :]
-        self.values = full_value_states[:, :, -self.sliding_window + 1 :, :]
+        if not self.record_past:
+            self.keys = full_key_states[:, :, -self.sliding_window + 1 :, :]
+            self.values = full_value_states[:, :, -self.sliding_window + 1 :, :]
+        # If we record the past, we keep them all for now, and they'll be restricted to the window size in `crop`
+        else:
+            self.keys = full_key_states
+            self.values = full_value_states
 
         # Return the full states
         return full_key_states, full_value_states
@@ -256,16 +269,31 @@ def get_max_length(self) -> int:
 
     def crop(self, max_length: int) -> None:
         """
-        Crop the past key values up to a new `max_length` in terms of tokens. `max_length` can also be
-        negative to remove `max_length` tokens.
+        Crop the past key values up to a new `max_length` in terms of tokens. `max_length` can also be negative to remove `max_length`
+        tokens.
         """
+        # If we are beyond the sliding window, we need to be more careful
         if self.get_seq_length() >= self.sliding_window:
-            raise ValueError(
-                "Cannot `crop` a `DynamicSlidingWindowLayer` after it has seen more tokens than its"
-                "sliding window (otherwise some states are lost)"
-            )
-        super().crop(max_length)
-        self.cumulative_length = self.keys.shape[-2]
+            if not self.record_past:
+                raise RuntimeError(
+                    "`crop` was called, but the current layer does not track past states, and the sliding window size was already "
+                    "reached. Call `activate_past_recording` before `crop` to be able to rollback the cache."
+                )
+            if max_length > 0:
+                raise RuntimeError(
+                    "Once the sliding window size has been reached, `DynamicSlidingWindowLayer` can only be cropped by passing a "
+                    "negative int, to specify how many tokens to remove"
+                )
+            tokens_to_remove = abs(max_length)
+            # We crop, and restrict the size back to the sliding window if still larger
+            self.keys = self.keys[:, :, -self.sliding_window + 1 - tokens_to_remove : -tokens_to_remove, :]
+            self.values = self.values[:, :, -self.sliding_window + 1 - tokens_to_remove : -tokens_to_remove, :]
+            self.cumulative_length = self.cumulative_length - tokens_to_remove
+
+        # If we did not reach the sliding window, we can do the same as for a full attention layer
+        else:
+            super().crop(max_length)
+            self.cumulative_length = self.keys.shape[-2]
 
 
 class DynamicIndexedLayer(DynamicLayer):

src/transformers/generation/candidate_generator.py

@@ -188,9 +188,6 @@ def __init__(
             processor for processor in self.logits_processor if not isinstance(processor, MinLengthLogitsProcessor)
         ]
 
-        # We need to roll back the cache in assisted generation, only DynamicCache is supported
-        self.generation_config.cache_implementation = "dynamic_full"
-
         if (
             is_sklearn_available()
             and self.assistant_generation_config.assistant_confidence_threshold
@@ -295,8 +292,8 @@ def _update_past_and_masks(
         """Update past key values and attention masks for subsequent generation rounds."""
         has_past_key_values = self.assistant_kwargs.get("past_key_values", None) is not None
         if has_past_key_values:
-            new_cache_size = input_ids.shape[-1] - 1 - remove_from_pkv
-            self.assistant_kwargs["past_key_values"].crop(new_cache_size - num_added_tokens)
+            tokens_to_remove = remove_from_pkv + num_added_tokens
+            self.assistant_kwargs["past_key_values"].crop(-tokens_to_remove)
             self.assistant_kwargs = _prepare_attention_mask(
                 self.assistant_kwargs, input_ids.shape[-1], self.assistant_model.config.is_encoder_decoder
             )
@@ -305,10 +302,6 @@ def _update_past_and_masks(
             )
             self.assistant_kwargs = _prepare_token_type_ids(self.assistant_kwargs, input_ids.shape[-1])
 
-            # This unsets `dynamic_full`, needed to initialize a new cache for the assistant. After the first forward
-            # pass on each generation, we reuse the cache instead.
-            self.generation_config.cache_implementation = None
-
         return has_past_key_values
 
     def _prepare_generation_args(self, input_ids: torch.LongTensor, min_new_tokens: int, max_new_tokens: int) -> dict:
@@ -1434,13 +1427,7 @@ def __init__(
         model_kwargs: dict[str, Any],
         logits_processor: Optional["LogitsProcessorList"] = None,
     ):
-        from ..cache_utils import (
-            DynamicLayer,
-            DynamicSlidingWindowLayer,
-            LinearAttentionAndFullAttentionLayer,
-            LinearAttentionAndSlidingWindowAttentionLayer,
-            MtpCache,
-        )
+        from ..cache_utils import MtpCache
         from ..modeling_layers import MtpModel
 
         self.num_mtp_layers = getattr(main_model.config.get_text_config(), "num_mtp_layers", None)
@@ -1456,18 +1443,7 @@ def __init__(
         self.mtp_model = MtpModel.from_pretrained(main_model, device_map={"": self.device})
 
         # Create the mtp cache and allow it to keep its past before we crop it
-        mtp_cache = MtpCache(config=main_model.config.get_mtp_config())
-        # Use full attention for now on sliding layers, same as main_model cache
-        mtp_cache.layers = [
-            DynamicLayer() if type(layer) is DynamicSlidingWindowLayer else layer for layer in mtp_cache.layers
-        ]
-        mtp_cache.layers = [
-            LinearAttentionAndFullAttentionLayer(number_of_states=layer.number_of_states)
-            if type(layer) is LinearAttentionAndSlidingWindowAttentionLayer
-            else layer
-            for layer in mtp_cache.layers
-        ]
-        self.mtp_cache = mtp_cache
+        self.mtp_cache = MtpCache(config=main_model.config.get_mtp_config())
         self.mtp_cache.activate_past_recording()
 
         # Save those to know how to decode mtp tokens

src/transformers/generation/configuration_utils.py

@@ -44,7 +44,7 @@
 logger = logging.get_logger(__name__)
 METADATA_FIELDS = ("_from_model_config", "_commit_hash", "_original_object_hash", "transformers_version")
 STATIC_CACHE_IMPLEMENTATIONS = ("static", "offloaded_static")
-DYNAMIC_CACHE_IMPLEMENTATIONS = ("dynamic", "dynamic_full", "offloaded", "quantized")
+DYNAMIC_CACHE_IMPLEMENTATIONS = ("dynamic", "offloaded", "quantized")
 # All the following are redundant and deprecated, but kept for BC
 DEPRECATED_STATIC_CACHE_IMPLEMENTATIONS = (
     "sliding_window",

src/transformers/generation/utils.py

@@ -1955,20 +1955,7 @@ def _prepare_cache_for_generation(
             return
 
         # Otherwise we NEED to prepare a cache, based on `generation_config.cache_implementation`
-
-        # Assisted decoding and contrastive search require cache rollback, which is incompatible with sliding layers.
-        # To handle this, we skip passing the model config to DynamicCache (forcing a full-layer cache).
-        # The "dynamic_full" option is a shortcut for generate() users to avoid sliding layers on their own.
-        if generation_mode in (GenerationMode.ASSISTED_GENERATION, GenerationMode.CONTRASTIVE_SEARCH):
-            if generation_config.cache_implementation is not None:
-                logger.warning_once(
-                    "An assistant model is provided, using a dynamic cache instead of a cache of type="
-                    f"'{generation_config.cache_implementation}'."
-                )
-            generation_config.cache_implementation = "dynamic_full"
-
-        dynamic_cache_kwargs = {}
-        dynamic_cache_kwargs["config"] = self.config.get_text_config(decoder=True)
+        dynamic_cache_kwargs = {"config": self.config.get_text_config(decoder=True)}
 
         if generation_config.cache_implementation == "offloaded":
             dynamic_cache_kwargs["offloading"] = True
@@ -2005,30 +1992,9 @@ def _prepare_cache_for_generation(
             cache_config.setdefault("config", self.config.get_text_config(decoder=True))
             backend = cache_config.pop("backend", "quanto")
             model_kwargs[cache_name] = QuantizedCache(backend=backend, **cache_config)
-        # i.e. `cache_implementation` in [None, "dynamic", "offloaded", "dynamic_full"]
-        # TODO: prepare linear cache from a single API, instead of creating in modeling code
+        # i.e. `cache_implementation` in [None, "dynamic", "offloaded"]
         else:
-            cache = DynamicCache(**dynamic_cache_kwargs)
-            # Replace sliding by full
-            if generation_config.cache_implementation == "dynamic_full":
-                from ..cache_utils import (
-                    DynamicLayer,
-                    DynamicSlidingWindowLayer,
-                    LinearAttentionAndFullAttentionLayer,
-                    LinearAttentionAndSlidingWindowAttentionLayer,
-                )
-
-                cache.layers = [
-                    DynamicLayer() if type(layer) is DynamicSlidingWindowLayer else layer for layer in cache.layers
-                ]
-                cache.layers = [
-                    LinearAttentionAndFullAttentionLayer(number_of_states=layer.number_of_states)
-                    if type(layer) is LinearAttentionAndSlidingWindowAttentionLayer
-                    else layer
-                    for layer in cache.layers
-                ]
-
-            model_kwargs[cache_name] = cache
+            model_kwargs[cache_name] = DynamicCache(**dynamic_cache_kwargs)
 
         if (
             self.config.is_encoder_decoder
@@ -2040,6 +2006,10 @@ def _prepare_cache_for_generation(
                 DynamicCache(**dynamic_cache_kwargs),  # cross-attention cache
             )
 
+        # If we just created a cache for an assistant model, mark it for past recording, as we will need to rollback it
+        if generation_config.is_assistant:
+            model_kwargs[cache_name].activate_past_recording()
+
     def _supports_logits_to_keep(self: "GenerativePreTrainedModel") -> bool:
         """
         Return True if the current model supports the keyword argument `logits_to_keep` in forward()

src/transformers/models/diffusion_gemma/generation_diffusion_gemma.py

@@ -925,9 +925,7 @@ def _prepare_cache_for_generation(
 
         # Dynamic Caches
         else:
-            dynamic_cache_kwargs = {}
-            if generation_config.cache_implementation != "dynamic_full":
-                dynamic_cache_kwargs["config"] = self.config.get_text_config(decoder=True)
+            dynamic_cache_kwargs = {"config": self.config.get_text_config(decoder=True)}
             if generation_config.cache_implementation == "offloaded":
                 dynamic_cache_kwargs["offloading"] = True
             past_key_values = DynamicCache(**dynamic_cache_kwargs)

src/transformers/models/minimax_m3_vl/modeling_minimax_m3_vl.py

@@ -79,11 +79,12 @@ def batch_select_indices(self, indices: torch.Tensor) -> None:
             self.idx_keys = self.idx_keys[indices, ...]
 
     def crop(self, max_length: int) -> None:
-        super().crop(max_length)
+        # Important to get the seq_len before the call to `super`, as it will be changed inside otherwise
         if max_length < 0:
             max_length = self.get_seq_length() - abs(max_length)
         if self.idx_keys is not None and self.idx_keys.shape[-2] > max_length:
             self.idx_keys = self.idx_keys[..., :max_length, :]
+        super().crop(max_length)
 
 
 class MiniMaxM3VLSparseStaticCacheLayer(StaticLayer):

src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py

@@ -257,11 +257,12 @@ def batch_select_indices(self, indices: torch.Tensor) -> None:
             self.idx_keys = self.idx_keys[indices, ...]
 
     def crop(self, max_length: int) -> None:
-        super().crop(max_length)
+        # Important to get the seq_len before the call to `super`, as it will be changed inside otherwise
         if max_length < 0:
             max_length = self.get_seq_length() - abs(max_length)
         if self.idx_keys is not None and self.idx_keys.shape[-2] > max_length:
             self.idx_keys = self.idx_keys[..., :max_length, :]
+        super().crop(max_length)
 
 
 class MiniMaxM3VLSparseStaticCacheLayer(StaticLayer):

tests/generation/test_utils.py

@@ -699,8 +699,6 @@ def test_assisted_decoding_matches_greedy_search(self, assistant_type):
             #    the assistant model is correct
             # c) there are at least two forward passes in the main model, to ensure the input preparation of
             #    the main model is correct
-            # d) use a cache type compatible with rollbacks (only dynamic cache atm). Otherwise, there may be
-            #     differences vs model-specific default cache
             generation_kwargs = {
                 "eos_token_id": -1,  # see a)
                 "max_new_tokens": 4,  # see c)
@@ -712,7 +710,6 @@ def test_assisted_decoding_matches_greedy_search(self, assistant_type):
                 "output_attentions": self.has_attentions,
                 "return_dict_in_generate": True,
                 "use_cache": True,
-                "cache_implementation": "dynamic_full",  # see d)
             }
             logits_processor_kwargs = self._get_logits_processor_kwargs(config=model.config)
 
@@ -797,8 +794,6 @@ def test_prompt_lookup_decoding_matches_greedy_search(self):
             #    prompt lookup is correct
             # c) there are at least two forward passes in the main model, to ensure the input preparation of
             #    the main model is correct
-            # d) use a cache type compatible with rollbacks (only dynamic cache atm). Otherwise, there may be
-            #     differences vs model-specific default cache
             generation_kwargs = {
                 "eos_token_id": -1,  # see a)
                 "max_new_tokens": 4,  # see c)
@@ -810,7 +805,6 @@ def test_prompt_lookup_decoding_matches_greedy_search(self):
                 "output_attentions": self.has_attentions,
                 "return_dict_in_generate": True,
                 "use_cache": True,
-                "cache_implementation": "dynamic_full",  # see d)
             }
             logits_processor_kwargs = self._get_logits_processor_kwargs(config=model.config)
 
@@ -872,8 +866,6 @@ def test_assisted_decoding_sample(self):
             #    the assistant model is correct
             # c) there are at least two forward passes in the main model, to ensure the input preparation of
             #    the main model is correct
-            # d) use a cache type compatible with rollbacks (only dynamic cache atm). Otherwise, there may be
-            #     differences vs model-specific default cache
             assistant_model = model
             assistant_model.generation_config.num_assistant_tokens = 2  # see b)
             assistant_model.generation_config.num_assistant_tokens_schedule = "constant"  # see b)
@@ -889,7 +881,6 @@ def test_assisted_decoding_sample(self):
                 "output_attentions": self.has_attentions,
                 "return_dict_in_generate": True,
                 "use_cache": True,
-                "cache_implementation": "dynamic_full",  # see d)
             }
             logits_processor_kwargs = self._get_logits_processor_kwargs(config=model.config)
             output_assisted = model.generate(**generation_kwargs, **inputs_dict, **logits_processor_kwargs)