[modelling] remove un-necessary transpose for fa2 attention

[sanchit-gandhi]

What does this PR do?

Fixes #31166 (comment). Previously, we were performing three shape operations on the q_proj tensor:

1. Reshape to (bsz, tgt_len, num_heads, head_dim)
2. Transpose to (bsz, num_heads, tgt_len, head_dim)
3. Transpose to (bsz, tgt_len, num_heads, head_dim)

Clearly 2 and 3 are not required! This PR fixes the modelling code accordingly.

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

[amyeroberts]

Thanks!

Could you check for similar patterns in the library in FA2, it's possible we're extra transposing all over the place

[sanchit-gandhi]

The pattern appears in most LLaMA-style attention modules, but here the back-to-back transpose operations are required, since there's a step in between the transpose ops:

1. Reshape to (bsz, tgt_len, num_heads, head_dim)
2. Transpose to (bsz, num_heads, tgt_len, head_dim)
3. Apply rotary pos embeddings (note that this requires the shape produced from 2)
4. Transpose to (bsz, tgt_len, num_heads, head_dim)

This leaves only two additional models where there are redundant transpose ops: Idefics2 and Jamba, changes for which I've pushed to the PR.

[ArthurZucker]

good find!