refactor rotary embeddings

Author: DavidLandup0

keras_hub/src/models/smollm3/smollm3_backbone.py

@@ -8,7 +8,6 @@
 from keras_hub.src.models.backbone import Backbone
 from keras_hub.src.models.smollm3.smollm3_layers import SmolLM3DecoderLayer
 from keras_hub.src.models.smollm3.smollm3_layers import SmolLM3RotaryEmbedding
-from keras_hub.src.layers.modeling.rotary_embedding import RotaryEmbedding
 
 
 @keras_hub_export(
@@ -70,7 +69,6 @@ def __init__(
         max_position_embeddings,
         rope_theta,
         partial_rotary_factor,
-        rope_scaling=1,
         **kwargs,
     ):
         # === Layers ===
@@ -102,17 +100,12 @@ def __init__(
             name="sequence_output_layernorm",
         )
 
-        #self.rotary_embedding = SmolLM3RotaryEmbedding(
-        #    hidden_size=hidden_dim,
-        #    num_attention_heads=num_attention_heads,
-        #    max_position_embeddings=max_position_embeddings,
-        #    rope_theta=rope_theta,
-        #    partial_rotary_factor=partial_rotary_factor,
-        #)
-        self.rotary_embedding = RotaryEmbedding(
-            max_wavelength=rope_theta,
-            scaling_factor=rope_scaling,
-            dtype=self.token_embedding.dtype_policy
+        self.rotary_embedding = SmolLM3RotaryEmbedding(
+            hidden_size=hidden_dim,
+            num_attention_heads=num_attention_heads,
+            max_position_embeddings=max_position_embeddings,
+            rope_theta=rope_theta,
+            partial_rotary_factor=partial_rotary_factor,
         )
 
         # === Functional Model ===
@@ -124,8 +117,14 @@ def __init__(
             shape=(None,), dtype="int32", name="padding_mask"
         )
 
+        cache_update_index = kwargs.get('self_attention_cache_index')
+
+        start_index = (
+            cache_update_index if cache_update_index is not None else 0
+        )
+
         hidden_states = self.token_embedding(token_id_input)
-        position_embeddings = self.rotary_embedding(hidden_states)
+        position_embeddings = self.rotary_embedding(hidden_states, start_index)
 
         for decoder_layer in self.transformer_layers[:num_layers]:
             hidden_states = decoder_layer(
@@ -161,7 +160,6 @@ def __init__(
         self.max_position_embeddings = max_position_embeddings
         self.rope_theta = rope_theta
         self.partial_rotary_factor = partial_rotary_factor
-        self.rope_scaling = rope_scaling
 
     def get_config(self):
         config = super().get_config()
@@ -182,7 +180,6 @@ def get_config(self):
                 "max_position_embeddings": self.max_position_embeddings,
                 "rope_theta": self.rope_theta,
                 "partial_rotary_factor": self.partial_rotary_factor,
-                "rope_scaling": self.rope_scaling
             }
         )
         return config

keras_hub/src/models/smollm3/smollm3_layers.py

@@ -583,7 +583,7 @@ def build(self, input_shape):
     def call(
         self,
         x,
-        position_ids,
+        start_index=0,
     ):
         """
         Forward pass for SmolLM3RotaryEmbedding.
@@ -596,13 +596,17 @@ def call(
         inv_freq_expanded = ops.expand_dims(
             ops.expand_dims(self.inv_freq, axis=0), axis=-1
         )
+        
+        batch_size = ops.shape(x)[0]
+        seq_len = ops.shape(x)[1]
+        positions = ops.arange(seq_len, dtype="float32")
+        positions + ops.cast(start_index, dtype="float32")
 
-        batch_size = ops.shape(position_ids)[0]
         inv_freq_expanded = ops.broadcast_to(
             inv_freq_expanded, (batch_size, ops.shape(self.inv_freq)[0], 1)
         )
 
-        position_ids_expanded = ops.expand_dims(position_ids, axis=1)
+        position_ids_expanded = ops.expand_dims(positions, axis=1)
 
         freqs = ops.matmul(
             ops.cast(inv_freq_expanded, "float32"),