Enable flash attention

[divyashreepathihalli]

This PR

• refactors the MHA layer so that its _compute_attention method would just call ops.dot_production_attention
• Adds a global toggle keras.config.enable_flash_attention and keras.config.is_flash_attention_enabled

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 72.72727% with 6 lines in your changes missing coverage. Please review.

Project coverage is 76.84%. Comparing base (c052cea) to head (3a47c53).
Report is 2 commits behind head on master.

Files with missing lines	Patch %	Lines
keras/src/layers/attention/multi_head_attention.py	63.63%	3 Missing and 1 partial ⚠️
keras/api/_tf_keras/keras/config/__init__.py	0.00%	2 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (c052cea) and HEAD (3a47c53). Click for more details.

HEAD has 2 uploads less than BASE

Flag	BASE (c052cea)	HEAD (3a47c53)
keras	4	3
keras-jax	1	0

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #20448      +/-   ##
==========================================
- Coverage   82.01%   76.84%   -5.18%     
==========================================
  Files         514      514              
  Lines       47194    47261      +67     
  Branches     7408     7417       +9     
==========================================
- Hits        38706    36317    -2389     
- Misses       6698     9220    +2522     
+ Partials     1790     1724      -66     

Flag	Coverage Δ
keras	76.72% <72.72%> (-5.15%)	⬇️
keras-jax	?
keras-numpy	59.86% <72.72%> (-0.01%)	⬇️
keras-tensorflow	65.90% <72.72%> (-0.02%)	⬇️
keras-torch	64.85% <72.72%> (-0.02%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[fchollet]

Thanks for the PR!

[fchollet]

This is needs to be threadlocal. Instead of doing it like this, use set_global_attribute/get_global_attribute from keras.src.backend.common.global_state. See how other global flags are implemented.

[divyashreepathihalli]

importing global_state in config created circular import error so I have moved the configs to attention.py

[fchollet]

Add a numerical correctness test with and without FA.