Update gating condition to include check for supporting GPUs for flash attention (#2184)

* update gating condition for flash attention

* fix test

* update utils

* fix tests

* add on t4 on deny list

* address review comments

* address comments

Author: divyashreepathihalli

Diff for: keras_hub/src/models/gemma/gemma_attention.py

@@ -6,7 +6,9 @@
 
 from keras_hub.src.layers.modeling.rotary_embedding import RotaryEmbedding
 from keras_hub.src.utils.keras_utils import clone_initializer
-from keras_hub.src.utils.keras_utils import has_flash_attention_support
+from keras_hub.src.utils.keras_utils import fused_attention_op_available
+from keras_hub.src.utils.keras_utils import gpu_supports_fused_attention_op
+from keras_hub.src.utils.keras_utils import running_on_gpu
 from keras_hub.src.utils.keras_utils import running_on_tpu
 
 
@@ -106,17 +108,22 @@ def _apply_rope(self, x, start_index):
         )
         return x
 
-    def _can_use_flash_attention(self):
-        if not has_flash_attention_support():
+    def _use_fused_attention_op(self):
+        if not fused_attention_op_available():
             return False
         if self.dropout > 0.0:
             return False
-        if self.logit_soft_cap is None:
-            return True
-        sig = inspect.signature(ops.dot_product_attention)
-        # We can currently only run soft capped attention for keras >= 3.10
-        # and only on TPU.
-        return running_on_tpu() and "attn_logits_soft_cap" in sig.parameters
+        if running_on_gpu():
+            # GPU never supports softcap in the fused op.
+            if self.logit_soft_cap is not None:
+                return False
+            return gpu_supports_fused_attention_op()
+        elif running_on_tpu():
+            # TPU supports softcap with on keras >= 3.10.
+            sig = inspect.signature(ops.dot_product_attention)
+            return "attn_logits_soft_cap" in sig.parameters
+        else:
+            return False
 
     def _compute_attention(
         self,
@@ -140,7 +147,7 @@ def _compute_attention(
                 cache_update_index=cache_update_index,
             )
 
-        if self._can_use_flash_attention():
+        if self._use_fused_attention_op():
             if attention_mask is not None:
                 attention_mask = ops.expand_dims(attention_mask, axis=1)
                 attention_mask = ops.cast(attention_mask, dtype="bool")

Diff for: keras_hub/src/models/gemma/gemma_causal_lm_test.py

@@ -12,7 +12,7 @@
 )
 from keras_hub.src.models.gemma.gemma_tokenizer import GemmaTokenizer
 from keras_hub.src.tests.test_case import TestCase
-from keras_hub.src.utils.keras_utils import has_flash_attention_support
+from keras_hub.src.utils.keras_utils import fused_attention_op_available
 from keras_hub.src.utils.keras_utils import running_on_gpu
 
 
@@ -98,7 +98,10 @@ def test_generate(self):
         )
 
     def test_flash_attention_call(self):
-        if keras.config.backend() != "jax" or not has_flash_attention_support():
+        if (
+            keras.config.backend() != "jax"
+            or not fused_attention_op_available()
+        ):
             self.skipTest("`flash_attention` testing requires the Jax backend.")
 
         with patch("keras.src.backend.nn.dot_product_attention") as mock_func:

Diff for: keras_hub/src/models/gemma3/gemma3_attention.py

@@ -7,7 +7,7 @@
 from keras_hub.src.layers.modeling.rotary_embedding import RotaryEmbedding
 from keras_hub.src.models.gemma.rms_normalization import RMSNormalization
 from keras_hub.src.utils.keras_utils import clone_initializer
-from keras_hub.src.utils.keras_utils import has_flash_attention_support
+from keras_hub.src.utils.keras_utils import fused_attention_op_available
 from keras_hub.src.utils.keras_utils import running_on_tpu
 
 
@@ -140,7 +140,7 @@ def _apply_rope(self, x, start_index):
         return x
 
     def _can_use_flash_attention(self):
-        if not has_flash_attention_support():
+        if not fused_attention_op_available():
             return False
         if self.dropout > 0.0:
             return False

Diff for: keras_hub/src/models/gemma3/gemma3_causal_lm_test.py

@@ -12,7 +12,7 @@
 )
 from keras_hub.src.models.gemma3.gemma3_tokenizer import Gemma3Tokenizer
 from keras_hub.src.tests.test_case import TestCase
-from keras_hub.src.utils.keras_utils import has_flash_attention_support
+from keras_hub.src.utils.keras_utils import fused_attention_op_available
 from keras_hub.src.utils.keras_utils import running_on_gpu
 
 
@@ -77,7 +77,10 @@ def test_text_causal_lm_basics(self):
         )
 
     def test_text_flash_attention_call(self):
-        if keras.config.backend() != "jax" or not has_flash_attention_support():
+        if (
+            keras.config.backend() != "jax"
+            or not fused_attention_op_available()
+        ):
             self.skipTest("`flash_attention` testing requires the Jax backend.")
 
         with patch("keras.src.backend.nn.dot_product_attention") as mock_func:

Diff for: keras_hub/src/models/gpt_neo_x/gpt_neo_x_attention.py

@@ -5,7 +5,7 @@
 
 from keras_hub.src.layers.modeling.rotary_embedding import RotaryEmbedding
 from keras_hub.src.utils.keras_utils import clone_initializer
-from keras_hub.src.utils.keras_utils import has_flash_attention_support
+from keras_hub.src.utils.keras_utils import fused_attention_op_available
 
 
 class GPTNeoXAttention(keras.layers.Layer):
@@ -125,7 +125,7 @@ def _masked_softmax(self, attention_scores, attention_mask=None):
     def _compute_attention(
         self, query, key, value, attention_mask=None, training=None
     ):
-        if has_flash_attention_support() and self.dropout == 0:
+        if fused_attention_op_available() and self.dropout == 0:
             # Use `dot_product_attention` with Flash Attention support if
             # available.
             if attention_mask is not None:

Diff for: keras_hub/src/models/llama/llama_attention.py

@@ -5,7 +5,7 @@
 
 from keras_hub.src.layers.modeling.rotary_embedding import RotaryEmbedding
 from keras_hub.src.utils.keras_utils import clone_initializer
-from keras_hub.src.utils.keras_utils import has_flash_attention_support
+from keras_hub.src.utils.keras_utils import fused_attention_op_available
 
 
 class LlamaAttention(keras.layers.Layer):
@@ -185,7 +185,7 @@ def _masked_softmax(self, attention_scores, attention_mask=None):
         return self._softmax(attention_scores)
 
     def _compute_attention(self, query, key, value, attention_mask=None):
-        if has_flash_attention_support():
+        if fused_attention_op_available():
             # Use `dot_product_attention` with Flash Attention support if
             # available.
             if attention_mask is not None:

Diff for: keras_hub/src/models/mistral/mistral_attention.py

@@ -5,7 +5,7 @@
 
 from keras_hub.src.layers.modeling.rotary_embedding import RotaryEmbedding
 from keras_hub.src.utils.keras_utils import clone_initializer
-from keras_hub.src.utils.keras_utils import has_flash_attention_support
+from keras_hub.src.utils.keras_utils import fused_attention_op_available
 
 
 # This is just a self-attention layer in Mistral. But it can be generalized
@@ -196,7 +196,7 @@ def _masked_softmax(self, attention_scores, attention_mask=None):
         return self._softmax(attention_scores)
 
     def _compute_attention(self, query, key, value, attention_mask=None):
-        if has_flash_attention_support():
+        if fused_attention_op_available():
             # Use `dot_product_attention` with Flash Attention support if
             # available.
             if attention_mask is not None:

Diff for: keras_hub/src/models/phi3/phi3_attention.py

@@ -8,7 +8,7 @@
     Phi3SuScaledRotaryEmbedding,
 )
 from keras_hub.src.utils.keras_utils import clone_initializer
-from keras_hub.src.utils.keras_utils import has_flash_attention_support
+from keras_hub.src.utils.keras_utils import fused_attention_op_available
 
 
 class Phi3Attention(keras.layers.Layer):
@@ -217,7 +217,7 @@ def _masked_softmax(self, attention_scores, attention_mask=None):
         return self.softmax(attention_scores)
 
     def _compute_attention(self, query, key, value, attention_mask=None):
-        if has_flash_attention_support():
+        if fused_attention_op_available():
             # Use `dot_product_attention` with Flash Attention support if
             # available.
             if attention_mask is not None:

Diff for: keras_hub/src/models/qwen/qwen_attention.py

@@ -5,7 +5,7 @@
 
 from keras_hub.src.layers.modeling.rotary_embedding import RotaryEmbedding
 from keras_hub.src.utils.keras_utils import clone_initializer
-from keras_hub.src.utils.keras_utils import has_flash_attention_support
+from keras_hub.src.utils.keras_utils import fused_attention_op_available
 
 
 class QwenAttention(keras.layers.Layer):
@@ -263,7 +263,7 @@ def _compute_attention(
         Returns:
             attention_output: Output tensor after applying attention.
         """
-        if has_flash_attention_support():
+        if fused_attention_op_available():
             # Use `dot_product_attention` with Flash Attention support if
             # available.
             if attention_mask is not None:

Diff for: keras_hub/src/models/stable_diffusion_3/mmdit.py

@@ -6,8 +6,8 @@
 
 from keras_hub.src.layers.modeling.position_embedding import PositionEmbedding
 from keras_hub.src.models.backbone import Backbone
+from keras_hub.src.utils.keras_utils import fused_attention_op_available
 from keras_hub.src.utils.keras_utils import gelu_approximate
-from keras_hub.src.utils.keras_utils import has_flash_attention_support
 from keras_hub.src.utils.keras_utils import standardize_data_format
 
 
@@ -771,7 +771,7 @@ def build(self, inputs_shape, context_shape, timestep_embedding_shape):
     def _compute_attention(self, query, key, value):
         batch_size = ops.shape(query)[0]
 
-        if has_flash_attention_support():
+        if fused_attention_op_available():
             # Use `dot_product_attention` with Flash Attention support if
             # available.
             encoded = ops.dot_product_attention(

Diff for: keras_hub/src/utils/keras_utils.py

@@ -55,7 +55,7 @@ def standardize_data_format(data_format):
     return data_format
 
 
-def has_flash_attention_support():
+def fused_attention_op_available():
     if (
         hasattr(keras.config, "is_flash_attention_enabled")
         and keras.config.backend() == "jax"
@@ -104,3 +104,46 @@ def running_on_gpu():
         import torch
 
         return torch.cuda.is_available()
+
+
+def gpu_supports_fused_attention_op():
+    deny_list = ["T4"]
+    for denied_gpu in deny_list:
+        if any(denied_gpu in gpu.upper() for gpu in get_gpu_names()):
+            return False
+    return True
+
+
+def get_gpu_names():
+    """Detects and returns the names of available GPUs based on the backend.
+
+    Note:
+        The format and content of the returned GPU names are **not normalized**
+        and vary significantly depending on the active backend. This function
+        provides the names as reported by the respective backend's API."
+    """
+    backend = keras.config.backend()
+    if backend == "jax":
+        import jax
+
+        devices = jax.devices()
+
+        return [getattr(d, "device_kind", "") for d in devices]
+
+    elif backend == "tensorflow":
+        import tensorflow as tf
+
+        gpus = tf.config.list_physical_devices("GPU")
+        return [
+            tf.config.experimental.get_device_details(gpu)["device_name"]
+            for gpu in gpus
+        ]
+    elif backend == "torch":
+        import torch
+
+        return [
+            torch.cuda.get_device_name(i)
+            for i in range(torch.cuda.device_count())
+        ]
+    else:
+        return [""]