Add MultiQueryAttention & GroupedQueryAttention

[awsaf49]

MultiQueryAttention (MQA) [Used in Falcon LLM] and GroupedQueryAttention (GQA) [Used in Llama 2 LLM] are alternatives to MultiHeadAttention (MHA) but they are a lot faster. Here's the speed comparison in my naive implementation,

===================================
          TensorFlow - GPU
===================================
Attention                 : 0.004 sec
Multi Head Attention      : 0.035 sec
Multi Query Attention     : 0.018 sec ( 50.17% faster than MHA )
Grouped Query Attention   : 0.030 sec ( 15.02% faster than MHA )

I think it would be nice to have these layers in keras-core.

Reference Papers:

• MQA: Paper
• GQA: Paper

[mattdangerw]

Probably easiest to just write GroupedQueryAttention, and consider MultiQueryAttention a special case of it. We can expose MultiQueryAttention, as subclass of GroupedQueryAttention that sets a single init value num_key_value_heads=1 on the base class. Somewhat similar to our AdamW class with weight_decay.

This is also some discussion in #18423 for more context. I definitely think adding support here makes sense. And probably clearer to have this standalone from MultiHeadAttention instead of just throwing more parameters at that (already quite complex) class.

Thanks for filing!

[awsaf49]

We can expose MultiQueryAttention, as subclass of GroupedQueryAttention that sets a single init value num_key_value_heads=1 on the base class.

I was thinking the same thing.

[awsaf49]

Should I open a PR for this??

[mattdangerw]

Sounds good! Thank you!

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No