CLIP issues with Flash Attention 3 pt.2
Author: akashpalla
Pull request: no
Labels: name: bug

Follow up to https://github.com/huggingface/transformers/issues/41668

### System Info

- `transformers` version: 4.57.0
- Platform: Linux-6.14.0-33-generic-x86_64-with-glibc2.39
- Python version: 3.12.3
- Huggingface_hub version: 1.0.0.rc4


### Who can help?

@ArthurZucker @vasqu 

### Information

- [ ] 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

```
from diffusers import EulerDiscreteScheduler, StableDiffusionXLPipeline
import torch
import os

model_path =os.getenv('MODEL_PATH')
print(model_path)
pipe = StableDiffusionXLPipeline.from_pretrained(
    model_path,
    variant="fp16",
    torch_dtype=torch.float16
).to("cuda")
pipe.fuse_qkv_projections()

pipe.text_encoder.config._attn_implementation = "flash_attention_3"
pipe.text_encoder_2.config._attn_implementation = "flash_attention_3"

prompt = ["A city at night with people walking around."] * 8
image = pipe(prompt, num_inference_steps=1).images[0]
```


### Expected behavior

When using the default attention implementation of SDPA with CLIP although `is_causal` is set to true, `create_causal_mask` generates a non-null attention mask leading `is_causal` to be set to False by SDPA and not lowering into [SDPA's](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html) own FA implementation. 

https://github.com/huggingface/transformers/blob/37d48bbb48483dd93541bca7c0c689b06e3ba729/src/transformers/integrations/sdpa_attention.py#L82

When setting the attention implementation to FA3, the following error is thrown
`AssertionError: expected size 16==16, stride 177408==59136 at dim=0; expected size 77==77, stride 2304==768 at dim=1
Error in op: torch.ops.flash_attn.flash_attn_func.default`

Setting .contiguous() to k,q,v resolves the error
https://github.com/huggingface/transformers/blob/37d48bbb48483dd93541bca7c0c689b06e3ba729/src/transformers/integrations/flash_attention.py#L56-L58

and is similar to what is done in flash_paged
https://github.com/huggingface/transformers/blob/37d48bbb48483dd93541bca7c0c689b06e3ba729/src/transformers/integrations/flash_paged.py#L88


and is done before calling flash_attn_varlen_func 
https://github.com/huggingface/transformers/blob/37d48bbb48483dd93541bca7c0c689b06e3ba729/src/transformers/modeling_flash_attention_utils.py#L389-L391

but not before calling `flash_varlen_fn`

Comments:
cc @cyrilvallez maybe?

Some initial questions for repro
- Do you have a model I can check this out with? I.e. `model_path =os.getenv('MODEL_PATH')`
- Which version of FA3 are you using, best the respective commit
- Please check out main, #41750 is essential for FA compatibility for text clip

There are multiple things going on here, somewhat even unrelated
- The issue is about FA3 but the first parts focus on SDPA - it is expected that is_causal can't be leveraged most of the time because padding forces us to create a mask. Also, torch SDPA only supports FA2 (while FA4 seems to be worked on internally)
- The assertion error is thrown in torch? The strides indicate what you fix, i.e. non-contiguous memory is not supported. Are you really using FA or SDPA?

I'd like to avoid having to call `.contiguous` tbh as it's a bit expensive op

NVM ignore the fa3 bug, I was just making a dumb mistake on my end. 

So just to confirm theres no way to lower into FA2 from SDPA when running CLIP?

> So just to confirm theres no way to lower into FA2 from SDPA when running CLIP?

Possibly related to #41856. It's a bit hard because we have logic that sometimes forces a mask creation even if it's not necessary (e.g. packed masks) - I would like to avoid these but haven't found a good way yet.