Fix incorrect attention mask truncate in WhisperFlashAttention2

[OliBomby]

What does this PR do?

Fixes incorrect attention calculation when training Whisper with Flash Attention 2 and passing decoder_attention_mask with some values set to False.

This error might've been made when copying the same truncating code from the other attention implementations. The problem is that in WhisperFlashAttention2 the dimensions have been transposed.

cc @sanchit-gandhi @ylacombe

[github-actions]

Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. When it is ready for review, please click the Ready for review button (at the bottom of the PR page).

[Rocketknight1]

cc @eustlb

[vasqu]

I think this is still a pretty big bug since it highly affects the attention computations in fa2 and addresses #36585 (which has been closed under the assumption of this fix)

cc @eustlb

[ronansgd]

It would be great to merge this fix! cc @eustlb

[greg2451]

Also affected by the issue, thanks for the fix @OliBomby.

@eustlb would be really great to get this merged 🙏🏼

[eustlb]

LGTM! Great catch, thanks a lot 🤗

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