[refactor] apply qk norm in attention processors

[a-r-r-o-w]

What does this PR do?

• Applies QK-norm in some of the attention processors that currently support it, and newly in the default AttentionProcessor2_0 and FusedAttentionProcessor2_0 (for you know what 👀).
• I think we should also refactor some of the attention processors and determine if we need these many.
• We should also probably re-order these in alphabetical order to make it easier to read, and group the normal and fused variants together.

Happy to take it up in this PR if we're okay with it.

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@sayakpaul @yiyixuxu

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

[a-r-r-o-w]

Additionally, since you can pre-QK norm (before reshaping to batch_size, num_heads, seq_length, head_dim) or post-QK norm (after reshape), I propose we add a parameter like pre_qk_norm=False to Attention to determine if the intent is pre/post-norm and encapsulate the reshape + norm logic into a single function:

def prepare_attention_qkv(self, query, key, value):
    batch_size, seq_length, embedding_dim = query.shape
    head_dim = embedding_dim // self.heads
    
    query = self.norm_q(query) if self.norm_q and self.pre_qk_norm else query
    key = self.norm_k(key) if self.norm_k and self.pre_qk_norm else key

    query = query.reshape(batch_size, -1, self.heads, head_dim).transpose(1, 2)
    key = key.reshape(batch_size, -1, self.heads, head_dim).transpose(1, 2)
    value = value.reshape(batch_size, -1, self.heads, head_dim).transpose(1, 2)

    query = self.norm_q(query) if self.norm_q and not self.pre_qk_norm else query
    key = self.norm_k(key) if self.norm_k and not self.pre_qk_norm else key

# Or a more readable implementation of above

I think this should cover a few cases of QK norm attention processors, and possibly help us remove a few processors (Hunyuan, for example) - if we also handle rotary embeds well in a clean way. But if we introduce too many code branches, it could get difficult to understand code flow so let me know what you think is best.

[yiyixuxu]

let's not add a method and just apply the qk_norm inside the attention processor

[sayakpaul]

Agree with @yiyixuxu regarding not having the method.

Should we try to have a test for this as well?

[sayakpaul]

It's just four lines of code that we're wrapping as a method. It's not doing anything complex, either. So, agree with @yiyixuxu that we should just keep it here.

[a-r-r-o-w]

I see, I was hoping to get in a refactor based on my comment above to have some attention processors removed/reused. But we can come back to this later. Should be good to merge now? Let's do this as soon as possible to perform rebase

[sayakpaul]

Will let @yiyixuxu to give the final approval.