Use fused attention for SmolLM3

keras_hub/src/models/smollm3/smollm3_backbone_test.py

@@ -1,7 +1,10 @@
+from unittest.mock import patch
+
 import pytest
 from keras import ops
 
 from keras_hub.src.models.smollm3.smollm3_backbone import SmolLM3Backbone
+from keras_hub.src.models.smollm3.smollm3_layers import SmolLM3Attention
 from keras_hub.src.tests.test_case import TestCase
 
 
@@ -29,6 +32,18 @@ def setUp(self):
             "padding_mask": ops.ones((2, 5), dtype="int32"),
         }
 
+    def _make_attention(self):
+        return SmolLM3Attention(
+            hidden_size=8,
+            num_attention_heads=2,
+            num_key_value_heads=2,
+            attention_bias=False,
+            attention_dropout=0.0,
+            rope_layer_enabled_list=[True],
+            layer_types=["attention"],
+            layer_idx=0,
+        )
+
     def test_backbone_basics(self):
         self.run_backbone_test(
             cls=SmolLM3Backbone,
@@ -51,3 +66,93 @@ def test_num_parameters(self):
         model = SmolLM3Backbone(**self.init_kwargs)
         # Reference value calculated from the model architecture
         self.assertEqual(model.count_params(), 80464)
+
+    def test_fused_attention(self):
+        attention = self._make_attention()
+        query = ops.ones((1, 4, 2, 4))
+        key = ops.ones((1, 4, 2, 4))
+        value = ops.ones((1, 4, 2, 4))
+        expected = ops.zeros_like(query)
+
+        for attention_mask in (
+            None,
+            ops.ones((1, 4, 4), dtype="int32"),
+        ):
+            with (
+                self.subTest(attention_mask=attention_mask),
+                patch(
+                    "keras_hub.src.models.smollm3.smollm3_layers."
+                    "fused_attention_op_available",
+                    return_value=True,
+                ),
+                patch.object(
+                    ops, "dot_product_attention", return_value=expected
+                ) as mock_attention,
+            ):
+                output = attention._compute_attention(
+                    query, key, value, attention_mask
+                )
+
+                self.assertIs(output, expected)
+                mock_attention.assert_called_once()
+                args, kwargs = mock_attention.call_args
+                self.assertIs(args[0], query)
+                self.assertIs(args[1], key)
+                self.assertIs(args[2], value)
+                expected_mask = (
+                    None
+                    if attention_mask is None
+                    else ops.expand_dims(
+                        ops.cast(attention_mask, "bool"), axis=1
+                    )
+                )
+                self.assertAllEqual(kwargs["mask"], expected_mask)
+                self.assertEqual(kwargs["scale"], attention._inv_norm_factor)
+
+    def test_fused_attention_fallback(self):
+        attention = self._make_attention()
+        query = ops.ones((1, 4, 2, 4))
+        key = ops.ones((1, 4, 2, 4))
+        value = ops.ones((1, 4, 2, 4))
+        attention_mask = ops.ones((1, 4, 4), dtype="bool")
+
+        with (
+            patch(
+                "keras_hub.src.models.smollm3.smollm3_layers."
+                "fused_attention_op_available",
+                return_value=False,
+            ),
+            patch.object(ops, "dot_product_attention") as mock_attention,
+        ):
+            output = attention._compute_attention(
+                query, key, value, attention_mask
+            )
+
+        mock_attention.assert_not_called()
+        self.assertAllClose(output, ops.ones_like(output))
+
+    def test_fused_attention_matches_fallback(self):
+        attention = self._make_attention()
+        query = ops.reshape(ops.arange(32, dtype="float32"), (1, 4, 2, 4))
+        key = ops.flip(query, axis=1)
+        value = ops.cos(query)
+        attention_mask = ops.tril(ops.ones((1, 4, 4), dtype="bool"))
+
+        with patch(
+            "keras_hub.src.models.smollm3.smollm3_layers."
+            "fused_attention_op_available",
+            return_value=False,
+        ):
+            fallback_output = attention._compute_attention(
+                query, key, value, attention_mask
+            )
+        with patch(
+            "keras_hub.src.models.smollm3.smollm3_layers."
+            "fused_attention_op_available",
+            return_value=True,
+        ):
+            fused_output = attention._compute_attention(
+                query, key, value, attention_mask
+            )
+
+        self.assertAllClose(fused_output, fallback_output, atol=1e-5)

keras_hub/src/models/smollm3/smollm3_layers.py

@@ -13,6 +13,7 @@
 )
 from keras_hub.src.models.smollm3.smollm3_utils import apply_rotary_pos_emb
 from keras_hub.src.models.smollm3.smollm3_utils import rope_init
+from keras_hub.src.utils.keras_utils import fused_attention_op_available
 
 
 class SmolLM3Attention(layers.Layer):
@@ -311,6 +312,18 @@ def _compute_attention(
         Returns:
             attention_output: Output tensor after applying attention.
         """
+        if fused_attention_op_available():
+            if attention_mask is not None:
+                attention_mask = ops.expand_dims(attention_mask, axis=1)
+                attention_mask = ops.cast(attention_mask, dtype="bool")
+            return ops.dot_product_attention(
+                query,
+                key,
+                value,
+                mask=attention_mask,
+                scale=self._inv_norm_factor,
+            )
+
         attention_scores = ops.einsum(self._dot_product_equation, query, key)
 
         attention_scores = ops.multiply(