This PR proposes a fix for the flash attention path in the multi-head attention module.
The flash attention block doesn't support the left padding mask, so we apply it upstream.
https://github.com/Dao-AILab/flash-attention/issues/649
We apply the key_pad_mask to the values contained in the KV-cache at the first step only, for all scenari (standard, flash, sdpa mechanisms).
I think it would be great to explain in this PR when and how the key_pad_mask needs to be used, and being clear in the different scenarii (standard, flash, sdpa mechanisms)
This PR proposes a fix for the flash attention path in the multi-head attention module. The flash attention block doesn't support the left padding mask, so we apply it upstream. https://github.com/Dao-AILab/flash-attention/issues/649
We apply the
key_pad_mask
to the values contained in the KV-cache at the first step only, for all scenari (standard, flash, sdpa mechanisms).