refactor: replace fa3 wrapper with original fa3 in attention backends registry

[badayvedat]

What does this PR do?

The _wrapped_flash_attn_3 function unconditionally unpacks both out and lse from the return value:

out, lse, *_ = flash_attn_3_func(...)

However, it was not passing return_attn_probs=True to request the tuple return. Since Dao-AILab/flash-attention@203b9b3, flash_attn_func returns only out by default, causing:

 File "/root/flash-attention/diffusers/.venv/lib/python3.12/site-packages/torch/_compile.py", line 51, in inner
    return disable_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/flash-attention/diffusers/.venv/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py", line 838, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/root/flash-attention/diffusers/.venv/lib/python3.12/site-packages/torch/_library/custom_ops.py", line 367, in wrapped_fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/root/flash-attention/diffusers/src/diffusers/models/attention_dispatch.py", line 643, in _wrapped_flash_attn_3
    out, lse, *_ = flash_attn_3_func(
    ^^^^^^^^^^^^
ValueError: not enough values to unpack (expected at least 2, got 1)

How does this pr fixes it

Adds return_attn_probs=True to the flash_attn_3_func call, consistent with how _flash_attention_3_hub handles.

Reproduction

# requirements.txt
git+github.com/huggingface/diffusers@5e48f466b9c0d257f2650e8feec378a0022f2402"
torch==2.7.1
transformers
accelerate
--extra-index-url=https://download.pytorch.org/whl/cu128

and bring your own flash attention build, to repro this i built it from source @ Dao-AILab/flash-attention@ac9b5f1

import torch
from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    torch_dtype=torch.bfloat16
).to("cuda")

pipe.transformer.set_attention_backend("_flash_3")

# ValueError: not enough values to unpack (expected at least 2, got 1)
pipe("a photo of a cat", num_inference_steps=1)

Alternative

The wrapper seems to exist to support fa3 as custom op. However, fa3 now has native torch.compile support as of Dao-AILab/flash-attention@c7697bb. This might be making _wrapped_flash_attn_3 redundant, tho i dont really know if that is the only reason.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

• General functionalities: @sayakpaul @yiyixuxu @DN6

[sayakpaul]

Thanks for your PR. Since torch.compile support has been merged, would you be interested in refactoring and cleaning up the FA3 backend in attention_dispatch.py?

[badayvedat]

Is there any downstream callers of this function that I also need to test?

[sayakpaul]

Thanks! Just a single question.

[sayakpaul]

Do you think we should version-guard this to keep it backwards-compatible?

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.