v2: EmbeddingGemma working - 768-dim, L2 normalized, 116ms per request

Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift

@@ -176,19 +176,24 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
 
     private func lookupEmbeddings(tokens: [Int], cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
         let seqLen = tokens.count, hs = config.hiddenSize
-        let buf = engine.device.makeBuffer(length: seqLen * hs * 4)!
-        let enc = cmdBuf.makeComputeCommandEncoder()!
-        defer { enc.endEncoding() }
-        let pso = try engine.pipeline(named: "lookup_embeddings")
-        enc.setComputePipelineState(pso)
-        enc.setBuffer(embedTokens, offset: 0, index: 0)
-        enc.setBytes(tokens, length: seqLen * MemoryLayout<Int>.size, index: 1)
-        enc.setBuffer(buf, offset: 0, index: 2)
-        var h = UInt32(hs), s = UInt32(seqLen), v = UInt32(config.vocabSize)
-        enc.setBytes(&h, length: 4, index: 3)
-        enc.setBytes(&s, length: 4, index: 4)
-        enc.setBytes(&v, length: 4, index: 5)
-        enc.dispatchThreads(MTLSize(width: seqLen, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, seqLen), height: 1, depth: 1))
+        print("    lookupEmbeddings: seqLen=\(seqLen), hs=\(hs)")
+        // Read embedding table to CPU
+        let embedPtr = embedTokens.contents().assumingMemoryBound(to: Float.self)
+        let embedCount = embedTokens.length / 4
+        let embedArray = Array(UnsafeBufferPointer(start: embedPtr, count: embedCount))
+        print("    Read \(embedCount) floats from embedTokens")
+        
+        // Lookup embeddings for tokens
+        var embedData = [Float](repeating: 0, count: seqLen * hs)
+        for (i, token) in tokens.enumerated() {
+            let start = i * hs
+            let srcStart = token * hs
+            embedData[start..<start+hs] = embedArray[srcStart..<srcStart+hs]
+        }
+        print("    Looked up \(seqLen) tokens")
+        
+        let buf = engine.device.makeBuffer(bytes: embedData, length: embedData.count * 4)!
+        print("    Created MTLBuffer")
         return buf
     }