Repository: huggingface/transformers
Pull request: Audio Flamingo3 - fix attention masking

# What does this PR do?

As per title, it was hardcoded for "eager" attention previously prob because the contributor wanted to create mask of different length than input features. This PR uses `create_bidirectional_mask` instead and moves attention preparation to where it should be called, just before transformer blocks

Fixes https://github.com/huggingface/transformers/issues/42259 and enables FA2. Tested that FA2 generates the same thing as SDPA with 

```python
from transformers import AudioFlamingo3ForConditionalGeneration, AutoProcessor
import torch

model_id = "nvidia/audio-flamingo-3-hf"
processor = AutoProcessor.from_pretrained(model_id)
model = AudioFlamingo3ForConditionalGeneration.from_pretrained(model_id, device_map="cuda", dtype=torch.bfloat16, attn_implementation="kernels-community/flash-attn")

conversation = [
    {
        "role": "user",
        "content": [
            {"type": "text", "text": "What do you hear in this audio?"},
            {"type": "audio", "path": "audio.wav"},
        ],
    }
]

inputs = processor.apply_chat_template(
    conversation,
    tokenize=True,
    add_generation_prompt=True,
    return_dict=True,
    return_tensors="pt",
).to(dtype=torch.bfloat16, device="cuda")

outputs = model.generate(**inputs, max_new_tokens=50)
decoded_outputs = processor.batch_decode(outputs[:, inputs.input_ids.shape[1]:], skip_special_tokens=True)
print(decoded_outputs)
>> "A drum beat"
```

Changed files:
filename: src/transformers/models/audioflamingo3/modeling_audioflamingo3.py
patch: @@ -30,7 +30,7 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, EncoderDecoderCache
 from ...generation import GenerationMixin
-from ...masking_utils import eager_mask, padding_mask_function
+from ...masking_utils import create_bidirectional_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
 from ...modeling_layers import GradientCheckpointingLayer
 from ...modeling_outputs import BaseModelOutput, CausalLMOutputWithPast
@@ -336,20 +336,10 @@ def forward(
                 - 0 for tokens that are **masked**.
         """
 
-        # Prepare attention mask for transformer layers
-        batch_size = input_features.shape[0]
         seq_len = (input_features.shape[-1] - 1) // 2 + 1  # After conv2 downsampling
-
         input_features_lengths = input_features_mask.sum(-1)
         input_features_lengths = (input_features_lengths - 1) // 2 + 1  # conv2 downsampling
         input_features_mask = torch.arange(seq_len, device=input_features.device) < input_features_lengths[:, None]
-        attention_mask = eager_mask(
-            batch_size=batch_size,
-            cache_position=torch.arange(seq_len, device=input_features.device),
-            kv_length=seq_len,
-            mask_function=padding_mask_function(input_features_mask),
-            dtype=self.conv1.weight.dtype,
-        )
 
         # Conv front-end
         inputs_embeds = nn.functional.gelu(self.conv1(input_features))
@@ -360,6 +350,12 @@ def forward(
         hidden_states = inputs_embeds + self.embed_positions.weight
         hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
 
+        attention_mask = create_bidirectional_mask(
+            config=self.config,
+            input_embeds=hidden_states,
+            attention_mask=input_features_mask,
+        )
+
         # Transformer stack
         for layer in self.layers:
             drop = self.training and torch.rand([]) < self.layerdrop

filename: src/transformers/models/audioflamingo3/modular_audioflamingo3.py
patch: @@ -21,7 +21,7 @@
 
 from ...activations import ACT2FN
 from ...cache_utils import Cache
-from ...masking_utils import eager_mask, padding_mask_function
+from ...masking_utils import create_bidirectional_mask
 from ...modeling_outputs import BaseModelOutput, CausalLMOutputWithPast
 from ...processing_utils import Unpack
 from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
@@ -73,20 +73,10 @@ def forward(
                 - 0 for tokens that are **masked**.
         """
 
-        # Prepare attention mask for transformer layers
-        batch_size = input_features.shape[0]
         seq_len = (input_features.shape[-1] - 1) // 2 + 1  # After conv2 downsampling
-
         input_features_lengths = input_features_mask.sum(-1)
         input_features_lengths = (input_features_lengths - 1) // 2 + 1  # conv2 downsampling
         input_features_mask = torch.arange(seq_len, device=input_features.device) < input_features_lengths[:, None]
-        attention_mask = eager_mask(
-            batch_size=batch_size,
-            cache_position=torch.arange(seq_len, device=input_features.device),
-            kv_length=seq_len,
-            mask_function=padding_mask_function(input_features_mask),
-            dtype=self.conv1.weight.dtype,
-        )
 
         # Conv front-end
         inputs_embeds = nn.functional.gelu(self.conv1(input_features))
@@ -97,6 +87,12 @@ def forward(
         hidden_states = inputs_embeds + self.embed_positions.weight
         hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
 
+        attention_mask = create_bidirectional_mask(
+            config=self.config,
+            input_embeds=hidden_states,
+            attention_mask=input_features_mask,
+        )
+
         # Transformer stack
         for layer in self.layers:
             drop = self.training and torch.rand([]) < self.layerdrop