v2: add EmbeddingKernels to Metal compilation pipeline

2026-07-06 12:11:31 +08:00
parent 5e060c7aea
commit f122d854dc
2 changed files with 17 additions and 2 deletions
@@ -109,6 +109,21 @@ public enum MetalKernels {
            .replacingOccurrences(of: "using namespace metal;\n", with: "")
        result += "\n" + fusedStripped
        
+        // Strip preamble from embedding kernels source
+        let embStripped = embeddingKernelsSource
+            .replacingOccurrences(of: "#include <metal_stdlib>\n", with: "")
+            .replacingOccurrences(of: "using namespace metal;\n", with: "")
+        result += "\n" + embStripped
+        
        return result
    }
-}
+
+    /// Embedding kernel source for EmbeddingGemma.
+    /// Includes RoPE, bidirectional sliding window attention, Q/K norm, and GELU.
+    public static var embeddingKernelsSource: String {
+        let url = URL(fileURLWithPath: #filePath)
+            .deletingLastPathComponent()
+            .appendingPathComponent("Metal/EmbeddingKernels.metal")
+        return try! String(contentsOf: url, encoding: .utf8)
+    }
+}