Do not drop mask with SDPA for more cases (#30311)

* overlooked * style * cleaner
huggingface · May 14, 2024 · 157bf0a · 157bf0a
1 parent 0fcf8b3
commit 157bf0a
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/src/transformers/modeling_attn_mask_utils.py b/src/transformers/modeling_attn_mask_utils.py
@@ -319,8 +319,12 @@ def _prepare_4d_causal_attention_mask_for_sdpa(
     ignore_causal_mask = False
 
     if attention_mask is None:
-        if sliding_window is None or key_value_length < sliding_window:
-            ignore_causal_mask = not is_tracing
+        if (
+            not is_tracing
+            and (query_length == 1 or key_value_length == query_length)
+            and (sliding_window is None or key_value_length < sliding_window)
+        ):
+            ignore_causal_mask = True
     elif sliding_window is None or key_value_length < sliding_window:
         # 4d mask is passed through
         if len(attention_mask.shape) == 4: