Repository: huggingface/transformers
Pull request: [`Attn`] Allow dynamic causality in SDPA via Kwargs

As per title, it's
- 1. Nice for power users
- 2. Apparently some models do need to switch like CLIP (depending on their input, see https://github.com/huggingface/transformers/pull/41673#discussion_r2438734215)

This allows us to rely on the set module's attribute per default but overwrite with kwarg if given. cc @zucchini-nlp 

Changed files:
filename: src/transformers/integrations/flash_attention.py
patch: @@ -34,6 +34,7 @@ def flash_attention_forward(
     scaling: Optional[float] = None,
     sliding_window: Optional[int] = None,
     softcap: Optional[float] = None,
+    is_causal: Optional[bool] = None,
     **kwargs,
 ) -> tuple[torch.Tensor, None]:
     if kwargs.get("output_attentions", False):
@@ -64,9 +65,7 @@ def flash_attention_forward(
     target_dtype = get_target_dtype(query, module)
 
     # Instead of relying on the value set in the module directly, we use the is_causal passed in kwargs if it is presented
-    is_causal = kwargs.pop("is_causal", None)
-    if is_causal is None:
-        is_causal = module.is_causal
+    is_causal = is_causal if is_causal is not None else module.is_causal
 
     attn_output = _flash_attention_forward(
         query,

filename: src/transformers/integrations/sdpa_attention.py
patch: @@ -68,13 +68,20 @@ def sdpa_attention_forward(
     if attention_mask is not None and attention_mask.ndim == 4:
         attention_mask = attention_mask[:, :, :, : key.shape[-2]]
 
-    # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment
-    # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.
-    # Note that it is important to check first for the shape, otherwise compile will fail with `argument 'is_causal' must be bool, not SymBool`
-    if is_causal is None:
-        # The last condition is for encoder (decoder) models which specify this by passing their own `is_causal` flag
-        # This is mainly due to those models having mixed implementations for encoder, decoder, and encoder-decoder attns
-        is_causal = query.shape[2] > 1 and attention_mask is None and getattr(module, "is_causal", True)
+    # Instead of relying on the value set in the module directly, we use the is_causal passed in kwargs if it is presented
+    is_causal = is_causal if is_causal is not None else getattr(module, "is_causal", True)
+
+    # SDPA's Flash Attention (and cuDNN) kernels rely on the `is_causal` flag. However, there are certain conditions:
+    # - Not in decoding phase (otherwise we want full attention on the single query token)
+    # - Attention mask is not to be provided (even if it is a causal pattern)
+    # - Internally, we marked this as compatible with causal, i.e. it is a decoder attention type
+    #
+    # Quirks on the conditionals:
+    # - We avoid inline passing this to the SDPA function directly to support both torch.compile's dynamic shapes and
+    #   full graph options. Otherwise, dynamic shapes are prevented from compiling.
+    # - It is important to check first for the shape, otherwise compile will fail with
+    #   `argument 'is_causal' must be bool, not SymBool`.
+    is_causal = query.shape[2] > 1 and attention_mask is None and is_causal
 
     # Shapes (e.g. query.shape[2]) are tensors during jit tracing, resulting in `is_causal` being a tensor.
     # We convert it to a bool for the SDPA kernel that only accepts bools.