AudioFlamingo3ForConditionalGeneration: `RuntimeError: Input type (float) and bias type (struct c10::BFloat16) should be the same`
Author: mingyi456
Pull request: no
Labels: name: bug

### System Info

(The command `transformers env` does not work for some reason. I am using Python 3.11.9, and my GPU is the RTX 4090.)
```
>nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2025 NVIDIA Corporation
Built on Fri_Feb_21_20:42:46_Pacific_Standard_Time_2025
Cuda compilation tools, release 12.8, V12.8.93
Build cuda_12.8.r12.8/compiler.35583870_0
```
```
>pip show transformers torch
Name: transformers
Version: 5.0.0.dev0
Summary: Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.
Home-page: https://github.com/huggingface/transformers
Author: The Hugging Face team (past and future) with the help of all our contributors (https://github.com/huggingface/transformers/graphs/contributors)
Author-email: transformers@huggingface.co
License: Apache 2.0 License
Location: C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages
Requires: filelock, huggingface-hub, numpy, packaging, pyyaml, regex, requests, safetensors, tokenizers, tqdm, typer-slim
Required-by: dfloat11, peft
---
Name: torch
Version: 2.7.1+cu128
Summary: Tensors and Dynamic neural networks in Python with strong GPU acceleration
Home-page: https://pytorch.org/
Author: PyTorch Team
Author-email: packages@pytorch.org
License: BSD-3-Clause
Location: C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages
Requires: filelock, fsspec, jinja2, networkx, sympy, typing-extensions
Required-by: accelerate, autoawq_kernels, bitsandbytes, deepspeed, flash-attn, open_clip_torch, openai-whisper, peft, s3tokenizer, timm, torchaudio, torchvision
```

### Who can help?

@zucchini-nlp @eustlb 

### Information

- [x] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

This is the code I used. I modified the `model_id` and audio path to point to local files, and I tried playing around with the `dtype` and `torch_dtype` options:
```python
from transformers import AudioFlamingo3ForConditionalGeneration, AutoProcessor
import torch

model_id = r"..\models\audio-flamingo-3-hf"
processor = AutoProcessor.from_pretrained(model_id)
model = AudioFlamingo3ForConditionalGeneration.from_pretrained(model_id, device_map="cuda", dtype = torch.bfloat16, torch_dtype = torch.bfloat16, attn_implementation = "flash_attention_2").eval()

conversation = [
    {
        "role": "user",
        "content": [
            {"type": "text", "text": "Transcribe the input speech."},
            {"type": "audio", "path": r".\WhDJDIviAOg_120_10.mp3"},
        ],
    }
]

inputs = processor.apply_chat_template(
    conversation,
    tokenize=True,
    add_generation_prompt=True,
    return_dict=True,
).to(model.device)

outputs = model.generate(**inputs, max_new_tokens=500)

decoded_outputs = processor.batch_decode(outputs[:, inputs.input_ids.shape[1]:], skip_special_tokens=True)
print(decoded_outputs)
```
This is the error I get:
```python
Traceback (most recent call last):
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\audio-flamingo-3\inference.py", line 25, in <module>
    outputs = model.generate(**inputs, max_new_tokens=500)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\torch\utils\_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\transformers\generation\utils.py", line 2684, in generate
    result = decoding_method(
             ^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\transformers\generation\utils.py", line 2877, in _sample
    outputs = self._prefill(input_ids, generation_config, model_kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\transformers\generation\utils.py", line 3853, in _prefill
    return self(**model_inputs, return_dict=True)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\torch\nn\modules\module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\torch\nn\modules\module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\transformers\utils\generic.py", line 768, in wrapper
    output = func(self, *args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\transformers\models\audioflamingo3\modeling_audioflamingo3.py", line 563, in forward
    audio_embeds = self.get_audio_features(input_features, input_features_mask)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\transformers\models\audioflamingo3\modeling_audioflamingo3.py", line 470, in get_audio_features
    encoder_output = self.audio_tower(input_features, input_features_mask=input_features_mask)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\torch\nn\modules\module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\torch\nn\modules\module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\transformers\utils\generic.py", line 768, in wrapper
    output = func(self, *args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\transformers\models\audioflamingo3\modeling_audioflamingo3.py", line 355, in forward
    inputs_embeds = nn.functional.gelu(self.conv1(input_features))
                                       ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\torch\nn\modules\module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\torch\nn\modules\module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\torch\nn\modules\conv.py", line 375, in forward
    return self._conv_forward(input, self.weight, self.bias)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\torch\nn\modules\conv.py", line 370, in _conv_forward
    return F.conv1d(
           ^^^^^^^^^
RuntimeError: Input type (float) and bias type (struct c10::BFloat16) should be the same
```

### Expected behavior

I expected that the code should work with BF16, or FP16, without having to resort to FP32.

Comments:
You need to cast inputs to model's dtype probably - `inputs = inputs.to(device=model.device, dtype=torch.bfloat16)`

Ok, now it works if I change `attn_implementation` to "sdpa". If I use flash attention I get this error:
```python
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [64,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [65,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [66,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [67,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [68,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [69,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [70,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [71,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [72,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [73,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [74,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [75,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [76,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [77,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [78,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [79,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [80,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [81,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [82,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [83,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [84,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [85,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [86,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [87,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [88,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [89,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [90,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [91,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [92,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [93,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [94,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\native\cuda\IndexKernel.cu:93: block: [2549,0,0], thread: [95,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
Traceback (most recent call last):
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\audio-flamingo-3\inference.py", line 28, in <module>
    outputs = model.generate(**inputs, max_new_tokens=500)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\torch\utils\_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\transformers\generation\utils.py", line 2684, in generate
    result = decoding_method(
             ^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\transformers\generation\utils.py", line 2877, in _sample
    outputs = self._prefill(input_ids, generation_config, model_kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\transformers\generation\utils.py", line 3853, in _prefill
    return self(**model_inputs, return_dict=True)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\torch\nn\modules\module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\torch\nn\modules\module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\transformers\utils\generic.py", line 768, in wrapper
    output = func(self, *args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\transformers\models\audioflamingo3\modeling_audioflamingo3.py", line 563, in forward
    audio_embeds = self.get_audio_features(input_features, input_features_mask)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\transformers\models\audioflamingo3\modeling_audioflamingo3.py", line 470, in get_audio_features
    encoder_output = self.audio_tower(input_features, input_features_mask=input_features_mask)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\torch\nn\modules\module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\torch\nn\modules\module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\transformers\utils\generic.py", line 768, in wrapper
    output = func(self, *args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\transformers\models\audioflamingo3\modeling_audioflamingo3.py", line 367, in forward
    hidden_states = layer(hidden_states, attention_mask)[0]
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\transformers\modeling_layers.py", line 94, in __call__
    return super().__call__(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\torch\nn\modules\module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\torch\nn\modules\module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\transformers\models\audioflamingo3\modeling_audioflamingo3.py", line 233, in forward
    hidden_states, attn_weights = self.self_attn(
                                  ^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\torch\nn\modules\module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\torch\nn\modules\module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\transformers\models\audioflamingo3\modeling_audioflamingo3.py", line 179, in forward
    attn_output, attn_weights = attention_interface(
                                ^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\transformers\integrations\flash_attention.py", line 70, in flash_attention_forward
    attn_output = _flash_attention_forward(
                  ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\transformers\modeling_flash_attention_utils.py", line 607, in _flash_attention_forward
    q, k, v, indices_q, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = _upad_input(
                                                                                       ^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\transformers\modeling_flash_attention_utils.py", line 290, in _upad_input
    query_layer = _index_first_axis(query_layer, indices_k)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\mingy\OneDrive\Desktop\Transformers\xformers-venv\Lib\site-packages\transformers\modeling_flash_attention_utils.py", line 154, in _index_first_axis
    return reshaped_tensor[indices]
           ~~~~~~~~~~~~~~~^^^^^^^^^
RuntimeError: CUDA error: device-side assert triggered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
```

By the way, I also run Qwen2.5-VL, ~~Qwen2.5-Omni~~, and Qwen3-VL with the same venv, and I verified that they work with flash attention.

cc @ebezzam 

Will close this, see https://github.com/huggingface/transformers/issues/42259#issuecomment-3547809952

It is expected that FA doesn't work with fp32 and only half precision. The difference between something like qwen vl and audio flamingo is that qwen uses input ids only while flamingo uses extra input features. The type of input ids is less relevant but the features cannot be fp32.

@vasqu No, after I followed the tip about casting the inputs to BF16 before passing them, to `model.generate()`, I still get the cuda assert error unless I disable FA. To be clear, this is the modified code that I used, with the inputs casted to BF16 accordingly, that had the error with FA:
```python
from transformers import AudioFlamingo3ForConditionalGeneration, AutoProcessor
import torch

model_id = r"..\models\audio-flamingo-3-hf"
processor = AutoProcessor.from_pretrained(model_id)
model = AudioFlamingo3ForConditionalGeneration.from_pretrained(model_id, device_map="cuda", dtype = torch.bfloat16, torch_dtype = torch.bfloat16, attn_implementation = "flash_attention_2")

conversation = [
    {
        "role": "user",
        "content": [
            {"type": "text", "text": "Transcribe the input speech."},
            {"type": "audio", "path": r".\WhDJDIviAOg_120_10.mp3"},
        ],
    }
]

inputs = processor.apply_chat_template(
    conversation,
    tokenize=True,
    add_generation_prompt=True,
    return_dict=True,
)

inputs = inputs.to(torch.bfloat16).to("cuda")
outputs = model.generate(**inputs, max_new_tokens=500)

decoded_outputs = processor.batch_decode(outputs[:, inputs.input_ids.shape[1]:], skip_special_tokens=True)
print(decoded_outputs)
```

Re-opening as this indeed seems to be the issue. Audio Flamingo has hardcoded `eager_mask` in its modeling code therefore FA2 can't be supported with 4D attention mask

No idea if the hardcoding was done for a reason or not, so keeping it open until then. We need to use `create_bidirectional_mask` instead and check if FA2 produces legible output. If anyone wants to take a stab, feel free 

Update: actually can make a PR myself soon 😄 

Sorry about that @mingyi456 (totally on me, acted too fast) and thx @zucchini-nlp 

We should indeed use the new mask api for this (if it works). Lmk, I can also take a stab 

Good timing, just submitted a PR.Can you review? @vasqu