Repository: huggingface/transformers
Pull request: handle flash slow tests

fixing broken slow / flash tests as we cannot pass the regular 4dim attention mask for flash

_update_causal_mask does not work with the patch cross attention , handling patch attn mask update for flash in the _prepare function

Changed files:
filename: src/transformers/models/blt/modeling_blt.py
patch: @@ -426,8 +426,8 @@ class BltPreTrainedModel(PreTrainedModel):
     _no_split_modules = ["BltTransformerLayer"]
     _can_compile_fullgraph = False  # static cache cannot have different shapes for each layer
     _supports_sdpa = True
-    _supports_flash_attn = True
-    _supports_flex_attn = True
+    _supports_flash_attn = False
+    _supports_flex_attn = False
     _supports_attention_backend = False
     _can_record_outputs = {
         "hidden_states": OutputRecorder(BltTransformerLayer, index=0, layer_name="local_decoder"),
@@ -1116,7 +1116,12 @@ def forward(
         )
 
         cross_attn_mask_enc = _prepare_patch_cross_attention_mask(
-            patch_ids, patch_lengths.shape[1], sequence_length, True, self.config.cross_attn_k, encoder_embeds.dtype
+            patch_ids=patch_ids,
+            num_patches=patch_lengths.shape[1],
+            sequence_length=sequence_length,
+            patches_as_queries=True,
+            cross_attn_k=self.config.cross_attn_k,
+            dtype=encoder_embeds.dtype,
         )
         encoder_hidden_states, encoder_cross_states = self.local_encoder(
             input_ids=input_ids,
@@ -1148,12 +1153,12 @@ def forward(
         )
         decoder_patch_ids = self._patch_ids_from_lengths(patch_lengths[:, 1:], sequence_length)
         cross_attn_mask_dec = _prepare_patch_cross_attention_mask(
-            decoder_patch_ids,
-            patch_lengths.shape[1],
-            sequence_length,
-            False,
-            self.config.cross_attn_k,
-            encoder_embeds.dtype,
+            patch_ids=decoder_patch_ids,
+            num_patches=patch_lengths.shape[1],
+            sequence_length=sequence_length,
+            patches_as_queries=False,
+            cross_attn_k=self.config.cross_attn_k,
+            dtype=encoder_embeds.dtype,
         )
         output = self.local_decoder(
             input_ids=input_ids,

filename: src/transformers/models/blt/modular_blt.py
patch: @@ -376,6 +376,8 @@ def forward(
 class BltPreTrainedModel(MllamaPreTrainedModel):
     config: BltConfig
     _supports_attention_backend = False
+    _supports_flash_attn = False
+    _supports_flex_attn = False
     _no_split_modules = ["BltTransformerLayer"]
     _can_record_outputs = {
         "hidden_states": OutputRecorder(BltTransformerLayer, index=0, layer_name="local_decoder"),
@@ -866,7 +868,12 @@ def forward(
         )
 
         cross_attn_mask_enc = _prepare_patch_cross_attention_mask(
-            patch_ids, patch_lengths.shape[1], sequence_length, True, self.config.cross_attn_k, encoder_embeds.dtype
+            patch_ids=patch_ids,
+            num_patches=patch_lengths.shape[1],
+            sequence_length=sequence_length,
+            patches_as_queries=True,
+            cross_attn_k=self.config.cross_attn_k,
+            dtype=encoder_embeds.dtype,
         )
         encoder_hidden_states, encoder_cross_states = self.local_encoder(
             input_ids=input_ids,
@@ -898,12 +905,12 @@ def forward(
         )
         decoder_patch_ids = self._patch_ids_from_lengths(patch_lengths[:, 1:], sequence_length)
         cross_attn_mask_dec = _prepare_patch_cross_attention_mask(
-            decoder_patch_ids,
-            patch_lengths.shape[1],
-            sequence_length,
-            False,
-            self.config.cross_attn_k,
-            encoder_embeds.dtype,
+            patch_ids=decoder_patch_ids,
+            num_patches=patch_lengths.shape[1],
+            sequence_length=sequence_length,
+            patches_as_queries=False,
+            cross_attn_k=self.config.cross_attn_k,
+            dtype=encoder_embeds.dtype,
         )
         output = self.local_decoder(
             input_ids=input_ids,

filename: tests/models/blt/test_modeling_blt.py
patch: @@ -284,10 +284,6 @@ def test_model_rope_scaling_from_config(self, scaling_type):
 
         self.assertFalse(torch.allclose(original_long_output, scaled_long_output, atol=1e-5))
 
-    @unittest.skip(reason="Decoder cannot keep gradients")
-    def test_flex_attention_with_grads():
-        pass
-
 
 @require_torch_accelerator
 class BltIntegrationTest(unittest.TestCase):