Update FlashAttention Hub integration and guard CUDA-only tests

[sywangyi]

version=1 only has <=torch2.12 support and in v3, return_lse is not the input arg(https://huggingface.co/kernels/kernels-community/flash-attn2/blob/v3/build/torch-stable-abi210-cu132-x86_64-linux/flash_attn_interface.py#L164) and do not support dropout>0

[sywangyi]

@sayakpaul please help review, verified in A100

[sayakpaul]

So this means return_softmax is always False?

[sywangyi]

see https://huggingface.co/kernels/kernels-community/flash-attn2/blob/v3/build/torch-stable-abi210-cu132-x86_64-linux/flash_attn_interface.py#L142, return_softmax will impact the third output(dmask) which is not used in the code.

[sayakpaul]

My question is what parameter does that False value correspond to? Is it not return_softmax?

[sywangyi]

it's return_softmax in v3, and default is False

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

Hi @sywangyi, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice.

Once the PR links an issue (or gets the no-issue-needed label), you can ignore this message — it stays here as a comment, but it no longer applies.

[sayakpaul]

Thanks! @DN6 I think this also fixes the issue you recently faced.