Fix cross-attention cache layer type for T5Gemma2 long inputs

src/transformers/models/t5gemma2/modeling_t5gemma2.py

@@ -1341,8 +1341,8 @@ def _prepare_cache_for_generation(
         cross_attn_config = copy.deepcopy(self.config.get_text_config(decoder=True))
 
         # cross-attention does not use sliding window
-        del cross_attn_config.sliding_window
-        del cross_attn_config.layer_types
+        cross_attn_config.sliding_window = None
+        cross_attn_config.layer_types = ["full_attention"] * cross_attn_config.num_hidden_layers
 
         cross_attn_cache_kwargs = {
             "config": cross_attn_config,

src/transformers/models/t5gemma2/modular_t5gemma2.py

@@ -1129,8 +1129,8 @@ def _prepare_cache_for_generation(
         cross_attn_config = copy.deepcopy(self.config.get_text_config(decoder=True))
 
         # cross-attention does not use sliding window
-        del cross_attn_config.sliding_window
-        del cross_attn_config.layer_types
+        cross_attn_config.sliding_window = None
+        cross_attn_config.layer_types = ["full_attention"] * cross_attn_config.num_hidden_layers
 
         cross_attn_cache_kwargs = {
             "config": cross_attn_config,

tests/models/t5gemma2/test_modeling_t5gemma2.py

@@ -28,6 +28,7 @@
     is_torch_available,
     is_vision_available,
 )
+from transformers.cache_utils import DynamicLayer, DynamicSlidingWindowLayer, EncoderDecoderCache
 from transformers.testing_utils import (
     Expectations,
     cleanup,
@@ -610,6 +611,48 @@ def create_and_check_generate_with_past_key_values(
         )
         self.parent.assertTrue(torch.all(output_with_past_cache == output_without_past_cache))
 
+    def create_and_check_cross_attention_cache_is_not_sliding(
+        self,
+        config,
+        input_ids,
+        decoder_input_ids,
+        attention_mask,
+        decoder_attention_mask,
+        lm_labels,
+        pixel_values,
+    ):
+        """
+        Regression test for #45521. Checks whether the cross attention cache is correctly handled, i.e. not a SWA cache.
+        This would previously fail on instances where the sliding window < encoder len.
+        """
+        config.decoder.sliding_window = self.encoder_seq_length // 2
+        self.parent.assertGreater(self.encoder_seq_length, config.decoder.sliding_window)
+        model = self.causal_lm_class(config=config).to(torch_device).eval()
+        output = model.generate(
+            input_ids,
+            pixel_values=pixel_values,
+            max_new_tokens=2,
+            do_sample=False,
+            use_cache=True,
+            return_dict_in_generate=True,
+        )
+        self.parent.assertIsInstance(output.past_key_values, EncoderDecoderCache)
+        cross_cache = output.past_key_values.cross_attention_cache
+        for layer_idx, layer in enumerate(cross_cache.layers):
+            self.parent.assertNotIsInstance(
+                layer,
+                DynamicSlidingWindowLayer,
+                msg=(
+                    f"Cross-attention layer {layer_idx} must not be a sliding-window layer "
+                    f"(got {type(layer).__name__}); cross-attention attends to all encoder tokens."
+                ),
+            )
+            self.parent.assertIs(
+                type(layer),
+                DynamicLayer,
+                msg=(f"Cross-attention layer {layer_idx} must be DynamicLayer (got {type(layer).__name__})."),
+            )
+
     def create_and_check_model_fp16_forward(
         self,
         config,
@@ -773,6 +816,10 @@ def test_generate_with_past_key_values(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_generate_with_past_key_values(*config_and_inputs)
 
+    def test_cross_attention_cache_is_not_sliding(self):
+        config_and_inputs = self.model_tester.prepare_config_and_inputs()
+        self.model_tester.create_and_check_cross_attention_cache_is_not_sliding(*config_and_inputs)
+
     @unittest.skipIf(torch_device == "cpu", "Can't do half precision")
     def test_model_fp16_forward(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()