optimize: use mx.fast.scaled_dot_product_attention

model.py

@@ -25,7 +25,7 @@ Usage:
 
     tokenizer = Tokenizer.from_file("tokenizer.json")
 
-    texts = ["Find the most relevant code snippet given the following query:\\nprint hello world"]
+    texts = ["Find the most relevant code snippet given the following query:\nprint hello world"]
     embeddings = model.encode(texts, tokenizer)
 """
 
@@ -101,22 +101,14 @@ class Attention(nn.Module):
         keys = keys.reshape(B, L, self.n_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
         values = values.reshape(B, L, self.n_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
 
-        # RoPE - rotate_half convention (traditional=False)
+        # RoPE via mx.fast
         queries = mx.fast.rope(queries, self.head_dim, traditional=False, base=self.rope_theta, scale=1.0, offset=0)
         keys = mx.fast.rope(keys, self.head_dim, traditional=False, base=self.rope_theta, scale=1.0, offset=0)
 
-        # GQA: repeat KV heads
-        if self.n_kv_heads != self.n_heads:
-            n_rep = self.n_heads // self.n_kv_heads
-            keys = mx.repeat(keys, n_rep, axis=1)
-            values = mx.repeat(values, n_rep, axis=1)
-
-        # Compute in float32 to avoid fp16 overflow
-        scores = (queries.astype(mx.float32) @ keys.astype(mx.float32).transpose(0, 1, 3, 2)) * self.scale
-        if mask is not None:
-            scores = scores + mask.astype(mx.float32)
-        attn = mx.softmax(scores, axis=-1)
-        output = attn.astype(values.dtype) @ values
+        # Scaled dot-product attention (handles GQA, precision, and masking internally)
+        output = mx.fast.scaled_dot_product_attention(
+            queries, keys, values, mask=mask.astype(queries.dtype) if mask is not None else None, scale=self.scale
+        )
 
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
         return self.o_proj(output)
@@ -188,7 +180,7 @@ class JinaCodeEmbeddingModel(nn.Module):
     ):
         batch_size, seq_len = input_ids.shape
 
-        # Causal mask
+        # Causal mask for SDPA: [1, 1, seq_len, seq_len]
         causal_mask = mx.tril(mx.ones((seq_len, seq_len)))
         causal_mask = mx.where(causal_mask == 0, -1e4, 0.0)
         causal_mask = causal_mask[None, None, :, :]