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
• Modify attention mask to match _masked_softmax implementation

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 73.68421% with 10 lines in your changes missing coverage. Please review.

Project coverage is 82.02%. Comparing base (30a6b87) to head (9a5200e).

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

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #20448      +/-   ##
==========================================
- Coverage   82.03%   82.02%   -0.02%     
==========================================
  Files         515      515              
  Lines       47346    47383      +37     
  Branches     7427     7435       +8     
==========================================
+ Hits        38842    38865      +23     
- Misses       6705     6714       +9     
- Partials     1799     1804       +5     

Flag	Coverage Δ
keras	81.87% <73.68%> (-0.02%)	⬇️
keras-jax	64.94% <68.42%> (-0.01%)	⬇️
keras-numpy	59.89% <50.00%> (-0.02%)	⬇️
keras-tensorflow	65.96% <60.52%> (-0.02%)	⬇️
keras-torch	64.86% <71.05%> (-0.01%)	⬇️

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!

[SamanehSaadat]

@divyashreepathihalli shouldn't return_attention_scores have a default value here?
There might be many instances of _compute_attention call that don't pass return_attention_scores value and I think this change can break them. There is an example here. So I was wondering if it's possible to set a default value here so that other references of _compute_attention work as before.
I think this change is probably the reason that the test is failing in keras-team/keras-hub#1977

[divyashreepathihalli]

This is tricky, right now the value passed to the call method is passed to _compute_attention. If we add a default value here and users don't pass the arg value from call it might override the call arg. and that could cause discrepancies.

[SamanehSaadat]

Isn't there any other way to check and make sure the users pass the arg value (rather than lack of default value here)?

[divyashreepathihalli]

I think a work around would be to add a self._return_attention_scores and then set it in the call method and use it in _compute_attention. wdyt?