diff --git a/Package.swift b/Package.swift index e677e77..3a4cd16 100644 --- a/Package.swift +++ b/Package.swift @@ -17,7 +17,7 @@ let package = Package( targets: [ .target( name: "MarkBase", - exclude: ["Metal/MetalKernels.metal", "Metal/OptimizedKernels.metal", "Metal/FusionKernels.metal", "Metal/MetalKernels.metallib", "Metal/metallib"], + exclude: ["Metal/MetalKernels.metal", "Metal/OptimizedKernels.metal", "Metal/FusionKernels.metal", "Metal/EmbeddingKernels.metal", "Metal/MetalKernels.metallib", "Metal/metallib"], linkerSettings: [ .linkedFramework("Metal"), .linkedFramework("Foundation"), diff --git a/Sources/MarkBase/MetalKernels.swift b/Sources/MarkBase/MetalKernels.swift index c5c0aa3..4f3fb09 100644 --- a/Sources/MarkBase/MetalKernels.swift +++ b/Sources/MarkBase/MetalKernels.swift @@ -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 \n", with: "") + .replacingOccurrences(of: "using namespace metal;\n", with: "") + result += "\n" + embStripped + return result } -} \ No newline at end of file + + /// 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) + } +}