Use Flash Attention if available (#2058)

james77777778 · web-flow · commit 63863ab7bc41 · 2025-01-27T19:10:54.000-08:00
* Use Flash Attention if available

* Torch's `dot_product_attention` doesn't support `bias`.
diff --git a/keras_hub/src/models/falcon/falcon_attention.py b/keras_hub/src/models/falcon/falcon_attention.py
@@ -110,16 +110,19 @@ def call(
 
         attention_scores = ops.einsum("bqnh,bknh->bnqk", query, key)
         attention_scores = ops.add(attention_scores, alibi)
-        attention_scores = (
-            attention_scores * self.inv_norm_factor
-        )  # [batch_size, num_heads, query_length, kv_length]
+        # [batch_size, num_heads, query_length, kv_length]
+        attention_scores = ops.multiply(
+            attention_scores,
+            ops.cast(self.inv_norm_factor, self.compute_dtype),
+        )
         attention_scores = self.softmax(
             attention_scores, ops.expand_dims(attention_mask, 1)
         )
         attention_scores = self.attention_dropout(attention_scores)
         attention_output = ops.einsum(
             "bnqk,bknh->bqnh", attention_scores, value
         )
+
         attention_output = ops.reshape(
             attention_output,
             [batch_size, seq_length, self.num_heads * self.head_dim],
diff --git a/keras_hub/src/models/gpt_neo_x/gpt_neo_x_attention.py b/keras_hub/src/models/gpt_neo_x/gpt_neo_x_attention.py
@@ -1,8 +1,11 @@
+import math
+
 import keras
 from keras import ops
 
 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
 
 
 class GPTNeoXAttention(keras.layers.Layer):
@@ -58,6 +61,8 @@ def __init__(
         self.bias_initializer = keras.initializers.get(bias_initializer)
         self.max_sequence_length = max_sequence_length
 
+        self._inv_norm_factor = 1.0 / math.sqrt(self.attn_head_size)
+
     def build(self, input_shape):
         self._qkv_dense = keras.layers.EinsumDense(
             equation="abc,cde->abde",
@@ -120,14 +125,26 @@ def _masked_softmax(self, attention_scores, attention_mask=None):
     def _compute_attention(
         self, query, key, value, attention_mask=None, training=None
     ):
-        attention_scores = ops.einsum("aecd,abcd->acbe", key, query)
+        if has_flash_attention_support() and self.dropout == 0:
+            # Use `dot_product_attention` with Flash Attention support if
+            # available.
+            if attention_mask is not None:
+                attention_mask = ops.expand_dims(attention_mask, axis=1)
+                attention_mask = ops.cast(attention_mask, dtype="bool")
+            attention_output = ops.dot_product_attention(
+                query,
+                key,
+                value,
+                mask=attention_mask,
+                scale=self._inv_norm_factor,
+            )
+            return attention_output
 
-        norm_factor = ops.sqrt(
-            ops.convert_to_tensor(self.attn_head_size, self.compute_dtype)
+        attention_scores = ops.einsum("aecd,abcd->acbe", key, query)
+        attention_scores = ops.multiply(
+            attention_scores,
+            ops.cast(self._inv_norm_factor, self.compute_dtype),
         )
-
-        attention_scores /= norm_factor
-
         attention_scores = self._masked_softmax(
             attention_scores, attention_mask
         )
diff --git a/keras_hub/src/models/llama/llama_attention.py b/keras_hub/src/models/llama/llama_attention.py
@@ -1,8 +1,11 @@
+import math
+
 import keras
 from keras import ops
 
 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
 
 
 class LlamaAttention(keras.layers.Layer):
@@ -43,7 +46,7 @@ def build(self, inputs_shape):
         # h = head dim
         hidden_dim = inputs_shape[-1]
         head_dim = hidden_dim // self.num_query_heads
-        self._norm_factor = ops.sqrt(ops.cast(head_dim, self.compute_dtype))
+        self._inv_norm_factor = 1.0 / math.sqrt(head_dim)
 
         self._query_dense = keras.layers.EinsumDense(
             equation="bqm,muh->bquh",
@@ -182,9 +185,27 @@ 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():
+            # Use `dot_product_attention` with Flash Attention support if
+            # available.
+            if attention_mask is not None:
+                attention_mask = ops.expand_dims(attention_mask, axis=1)
+                attention_mask = ops.cast(attention_mask, dtype="bool")
+            attention_output = ops.dot_product_attention(
+                query,
+                key,
+                value,
+                mask=attention_mask,
+                scale=self._inv_norm_factor,
+            )
+            return attention_output
+
         attention_scores = ops.einsum(self._dot_product_equation, query, key)
 
-        attention_scores = attention_scores / self._norm_factor
+        attention_scores = ops.multiply(
+            attention_scores,
+            ops.cast(self._inv_norm_factor, self.compute_dtype),
+        )
         attention_scores = self._masked_softmax(
             attention_scores, attention_mask
         )
diff --git a/keras_hub/src/models/mistral/mistral_attention.py b/keras_hub/src/models/mistral/mistral_attention.py
@@ -1,8 +1,11 @@
+import math
+
 import keras
 from keras import ops
 
 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
 
 
 # This is just a self-attention layer in Mistral. But it can be generalized
@@ -52,6 +55,7 @@ def build(self, inputs_shape):
         # h = head dim
         self._hidden_dim = inputs_shape[-1]
         self._head_dim = self._hidden_dim // self._num_query_heads
+        self._inv_norm_factor = 1.0 / math.sqrt(self._head_dim)
 
         self._query_dense = keras.layers.EinsumDense(
             equation="bqm,muh->bquh",
@@ -192,11 +196,26 @@ def _masked_softmax(self, attention_scores, attention_mask=None):
         return self._softmax(attention_scores)
 
     def _compute_attention(self, query, key, value, attention_mask=None):
-        attention_scores = ops.einsum(self._dot_product_equation, query, key)
-
-        norm_factor = ops.sqrt(ops.cast(self._head_dim, self.compute_dtype))
+        if has_flash_attention_support():
+            # Use `dot_product_attention` with Flash Attention support if
+            # available.
+            if attention_mask is not None:
+                attention_mask = ops.expand_dims(attention_mask, axis=1)
+                attention_mask = ops.cast(attention_mask, dtype="bool")
+            attention_output = ops.dot_product_attention(
+                query,
+                key,
+                value,
+                mask=attention_mask,
+                scale=self._inv_norm_factor,
+            )
+            return attention_output
 
-        attention_scores = attention_scores / norm_factor
+        attention_scores = ops.einsum(self._dot_product_equation, query, key)
+        attention_scores = ops.multiply(
+            attention_scores,
+            ops.cast(self._inv_norm_factor, self.compute_dtype),
+        )
         attention_scores = self._masked_softmax(
             attention_scores, attention_mask
         )
diff --git a/keras_hub/src/models/phi3/phi3_attention.py b/keras_hub/src/models/phi3/phi3_attention.py
@@ -1,3 +1,5 @@
+import math
+
 import keras
 from keras import ops
 
@@ -6,6 +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
 
 
 class Phi3Attention(keras.layers.Layer):
@@ -53,7 +56,7 @@ def build(self, inputs_shape):
         # h = head dim
         hidden_dim = inputs_shape[-1]
         head_dim = hidden_dim // self.num_query_heads
-        self._norm_factor = ops.sqrt(ops.cast(head_dim, self.compute_dtype))
+        self._inv_norm_factor = 1.0 / math.sqrt(head_dim)
 
         self.query_dense = keras.layers.EinsumDense(
             equation="bqm,muh->bquh",
@@ -214,8 +217,26 @@ 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():
+            # Use `dot_product_attention` with Flash Attention support if
+            # available.
+            if attention_mask is not None:
+                attention_mask = ops.expand_dims(attention_mask, axis=1)
+                attention_mask = ops.cast(attention_mask, dtype="bool")
+            attention_output = ops.dot_product_attention(
+                query,
+                key,
+                value,
+                mask=attention_mask,
+                scale=self._inv_norm_factor,
+            )
+            return attention_output
+
         attention_scores = ops.einsum("bquh,bkuh->buqk", query, key)
-        attention_scores = attention_scores / self._norm_factor
+        attention_scores = ops.multiply(
+            attention_scores,
+            ops.cast(self._inv_norm_factor, self.compute_dtype),
+        )
         attention_scores = self._masked_softmax(
             attention_scores, attention_mask
         )
diff --git a/keras_hub/src/models/stable_diffusion_3/mmdit.py b/keras_hub/src/models/stable_diffusion_3/mmdit.py
@@ -7,6 +7,7 @@
 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 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
 
 
@@ -770,17 +771,14 @@ def build(self, inputs_shape, context_shape, timestep_embedding_shape):
     def _compute_attention(self, query, key, value):
         batch_size = ops.shape(query)[0]
 
-        # Use the fast path when `ops.dot_product_attention` and flash attention
-        # are available.
-        if hasattr(ops, "dot_product_attention") and hasattr(
-            keras.config, "is_flash_attention_enabled"
-        ):
+        if has_flash_attention_support():
+            # Use `dot_product_attention` with Flash Attention support if
+            # available.
             encoded = ops.dot_product_attention(
                 query,
                 key,
                 value,
                 scale=self._inverse_sqrt_key_dim,
-                flash_attention=keras.config.is_flash_attention_enabled(),
             )
             return ops.reshape(
                 encoded, (batch_size, -1, self.num_heads * self.head_dim)
@@ -793,10 +791,9 @@ def _compute_attention(self, query, key, value):
         probs = self.softmax(logits)
         probs = ops.cast(probs, self.compute_dtype)
         encoded = ops.einsum("BNTS,BSNH->BTNH", probs, value)
-        encoded = ops.reshape(
+        return ops.reshape(
             encoded, (batch_size, -1, self.num_heads * self.head_dim)
         )
-        return encoded
 
     def call(self, inputs, context, timestep_embedding, training=None):
         # Compute pre-attention.
diff --git a/keras_hub/src/utils/keras_utils.py b/keras_hub/src/utils/keras_utils.py
@@ -53,3 +53,10 @@ def standardize_data_format(data_format):
             f"Received: data_format={data_format}"
         )
     return data_format
+
+
+def has_flash_attention_support():
+    if hasattr(keras.config, "is_flash_attention_enabled"):
+        return True
+    else:
+        return False
