From bafc2850e5f414930d45968b559e4bbe472c5746 Mon Sep 17 00:00:00 2001 From: Cyril Vallez Date: Tue, 21 Jul 2026 14:15:06 +0900 Subject: [PATCH 1/8] native sliding --- src/transformers/cache_utils.py | 47 ++++++++++++++----- .../generation/candidate_generator.py | 21 +-------- src/transformers/generation/utils.py | 40 ++-------------- 3 files changed, 41 insertions(+), 67 deletions(-) diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py index 5da4688f8f4e..05ed4b497a18 100644 --- a/src/transformers/cache_utils.py +++ b/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) @@ -227,9 +235,11 @@ def update( # Compute the full states 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 :, :] + # Only cache the last `self.sliding_window - 1` tokens (or all of them if lower than that) - if we record the past, we keep + # them all for now, and they'll be restricted to the window size in `crop` + if not self.record_past: + self.keys = full_key_states[:, :, -self.sliding_window + 1 :, :] + self.values = full_value_states[:, :, -self.sliding_window + 1 :, :] # Return the full states return full_key_states, full_value_states @@ -256,16 +266,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): diff --git a/src/transformers/generation/candidate_generator.py b/src/transformers/generation/candidate_generator.py index 5776e2ccec98..cd75d6534500 100644 --- a/src/transformers/generation/candidate_generator.py +++ b/src/transformers/generation/candidate_generator.py @@ -1434,13 +1434,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 +1450,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 diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index 54e70ee16485..8b665bf6838c 100644 --- a/src/transformers/generation/utils.py +++ b/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 From 1d605187d9d780c635d79bf8c4c36554177f4f7b Mon Sep 17 00:00:00 2001 From: Cyril Vallez Date: Tue, 21 Jul 2026 14:23:00 +0900 Subject: [PATCH 2/8] remove all traces of dynamic_full --- src/transformers/generation/candidate_generator.py | 7 ------- src/transformers/generation/configuration_utils.py | 2 +- .../models/diffusion_gemma/generation_diffusion_gemma.py | 4 +--- tests/generation/test_utils.py | 9 --------- 4 files changed, 2 insertions(+), 20 deletions(-) diff --git a/src/transformers/generation/candidate_generator.py b/src/transformers/generation/candidate_generator.py index cd75d6534500..72b0cb225a53 100644 --- a/src/transformers/generation/candidate_generator.py +++ b/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 @@ -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: diff --git a/src/transformers/generation/configuration_utils.py b/src/transformers/generation/configuration_utils.py index ad2847f18eba..2eeb3739797c 100644 --- a/src/transformers/generation/configuration_utils.py +++ b/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", diff --git a/src/transformers/models/diffusion_gemma/generation_diffusion_gemma.py b/src/transformers/models/diffusion_gemma/generation_diffusion_gemma.py index 714802956af8..d0b7d0140b0d 100644 --- a/src/transformers/models/diffusion_gemma/generation_diffusion_gemma.py +++ b/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) diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py index f239e09d29ff..f1469d715090 100644 --- a/tests/generation/test_utils.py +++ b/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) From 68075f8104df3f8845be20772ea7d05e5616de9a Mon Sep 17 00:00:00 2001 From: Cyril Vallez Date: Tue, 21 Jul 2026 14:26:20 +0900 Subject: [PATCH 3/8] record for assistants --- src/transformers/generation/utils.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index 8b665bf6838c..3ac3c3a1dad8 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -2006,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() From 5ec7b73615d32962af83c69c54cf46c752a2a2fa Mon Sep 17 00:00:00 2001 From: Cyril Vallez Date: Tue, 21 Jul 2026 14:37:24 +0900 Subject: [PATCH 4/8] crop with negative always --- src/transformers/generation/candidate_generator.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/transformers/generation/candidate_generator.py b/src/transformers/generation/candidate_generator.py index 72b0cb225a53..db5b5a31baa2 100644 --- a/src/transformers/generation/candidate_generator.py +++ b/src/transformers/generation/candidate_generator.py @@ -292,8 +292,7 @@ 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) + self.assistant_kwargs["past_key_values"].crop(-remove_from_pkv - num_added_tokens) self.assistant_kwargs = _prepare_attention_mask( self.assistant_kwargs, input_ids.shape[-1], self.assistant_model.config.is_encoder_decoder ) From 2cbeaf4385e12d21e7a656e1b0a3d02365c80131 Mon Sep 17 00:00:00 2001 From: Cyril Vallez Date: Tue, 21 Jul 2026 15:16:04 +0900 Subject: [PATCH 5/8] oupsi forgot else --- src/transformers/cache_utils.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py index 05ed4b497a18..439556217bde 100644 --- a/src/transformers/cache_utils.py +++ b/src/transformers/cache_utils.py @@ -235,11 +235,14 @@ def update( # Compute the full states 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) - if we record the past, we keep - # them all for now, and they'll be restricted to the window size in `crop` + # Only cache the last `self.sliding_window - 1` tokens (or all of them if lower than that) 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 From 3a73fccbb29c1b0f1abd0ac2535b041a697b750e Mon Sep 17 00:00:00 2001 From: Cyril Vallez Date: Tue, 21 Jul 2026 16:09:21 +0900 Subject: [PATCH 6/8] small fix --- src/transformers/models/minimax_m3_vl/modeling_minimax_m3_vl.py | 2 +- src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/minimax_m3_vl/modeling_minimax_m3_vl.py b/src/transformers/models/minimax_m3_vl/modeling_minimax_m3_vl.py index b04737837609..6b98e8b9f82e 100644 --- a/src/transformers/models/minimax_m3_vl/modeling_minimax_m3_vl.py +++ b/src/transformers/models/minimax_m3_vl/modeling_minimax_m3_vl.py @@ -81,7 +81,7 @@ def batch_select_indices(self, indices: torch.Tensor) -> None: def crop(self, max_length: int) -> None: super().crop(max_length) if max_length < 0: - max_length = self.get_seq_length() - abs(max_length) + max_length = self.idx_keys.shape[-2] - 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, :] diff --git a/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py b/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py index 438316b593c9..ef8b7b430b10 100644 --- a/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py +++ b/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py @@ -259,7 +259,7 @@ def batch_select_indices(self, indices: torch.Tensor) -> None: def crop(self, max_length: int) -> None: super().crop(max_length) if max_length < 0: - max_length = self.get_seq_length() - abs(max_length) + max_length = self.idx_keys.shape[-2] - 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, :] From 1ece6c23e051d61dfb773b6f9dfaa535b5149650 Mon Sep 17 00:00:00 2001 From: Cyril Vallez Date: Wed, 22 Jul 2026 11:59:57 +0900 Subject: [PATCH 7/8] nit --- src/transformers/generation/candidate_generator.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/transformers/generation/candidate_generator.py b/src/transformers/generation/candidate_generator.py index db5b5a31baa2..71edaa53a393 100644 --- a/src/transformers/generation/candidate_generator.py +++ b/src/transformers/generation/candidate_generator.py @@ -292,7 +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: - self.assistant_kwargs["past_key_values"].crop(-remove_from_pkv - 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 ) From 37ad4d4e178b49cc275b6ea0a9529be9f24ca6de Mon Sep 17 00:00:00 2001 From: Cyril Vallez Date: Wed, 22 Jul 2026 12:05:47 +0900 Subject: [PATCH 8/8] better fix for minimax --- .../models/minimax_m3_vl/modeling_minimax_m3_vl.py | 5 +++-- .../models/minimax_m3_vl/modular_minimax_m3_vl.py | 5 +++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/transformers/models/minimax_m3_vl/modeling_minimax_m3_vl.py b/src/transformers/models/minimax_m3_vl/modeling_minimax_m3_vl.py index 6b98e8b9f82e..a9b38c1ae2e9 100644 --- a/src/transformers/models/minimax_m3_vl/modeling_minimax_m3_vl.py +++ b/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.idx_keys.shape[-2] - abs(max_length) + 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): diff --git a/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py b/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py index ef8b7b430b10..a7d92104c87b 100644 --- a/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py +++ b/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.idx_keys.shape[-2] - abs(max_length) + 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):