v2: EmbeddingGemma multi-language working, 500-test stability issue pending

Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift

@@ -1,6 +1,5 @@
 import Foundation
 import Metal
-import MetalPerformanceShaders
 import Accelerate
 
 /// EmbeddingGemma configuration
@@ -42,6 +41,7 @@ public struct EmbeddingGemmaConfig: Codable {
     }
 }
 
+/// EmbeddingGemma - Google's 300M parameter embedding model
 public final class EmbeddingGemmaModel: @unchecked Sendable {
     public let config: EmbeddingGemmaConfig
     public let engine: MarkBaseEngine
@@ -67,12 +67,12 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         self.tokenizer = try TokenizerFactory.load(modelDir: modelDir)
         self.reader = try SafeTensorsReader(path: modelDir + "/model.safetensors")
         try loadWeights()
-        print("✓ EmbeddingGemma loaded (\(config.numHiddenLayers) layers, hidden=\(config.hiddenSize))")
     }
 
     private func loadWeights() throws {
         let embedData = try readTensor("embed_tokens.weight")
         embedTokens = engine.device.makeBuffer(bytes: embedData, length: embedData.count * 4)!
+
         for i in 0..<config.numHiddenLayers {
             let p = "layers.\(i)"
             layerNorms.append([
@@ -91,44 +91,39 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
             upProjs.append(try loadAndTranspose("\(p).mlp.up_proj.weight", rows: config.intermediateSize, cols: config.hiddenSize))
             downProjs.append(try loadAndTranspose("\(p).mlp.down_proj.weight", rows: config.hiddenSize, cols: config.intermediateSize))
         }
+
         let fnData = try readTensor("norm.weight")
         finalNorm = engine.device.makeBuffer(bytes: fnData, length: fnData.count * 4)!
     }
 
+    private func loadAndTranspose(_ name: String, rows: Int, cols: Int) throws -> MTLBuffer {
+        let data = try readTensor(name)
+        var transposed = [Float](repeating: 0, count: data.count)
+        for r in 0..<rows {
+            for c in 0..<cols {
+                transposed[c * rows + r] = data[r * cols + c]
+            }
+        }
+        return engine.device.makeBuffer(bytes: transposed, length: transposed.count * 4)!
+    }
+
     public func embed(text: String, maxLen: Int = 2048) throws -> [Float] {
         var tokens = tokenizer.encode(text: text)
         if tokens.count > maxLen { tokens = Array(tokens.prefix(maxLen)) }
         guard !tokens.isEmpty else { return [] }
 
         let seqLen = tokens.count, hs = config.hiddenSize
-
-        // Test 1: Embedding lookup only
-        print("  TEST: Embedding lookup...")
-        let cmdBuf1 = engine.commandQueue.makeCommandBuffer()!
-        let inputBuf = try lookupEmbeddings(tokens: tokens, cmdBuf: cmdBuf1)
-        cmdBuf1.commit()
-        cmdBuf1.waitUntilCompleted()
-        print("  TEST: Embedding lookup OK")
-
-        // Test 2: Single layer forward (layer 0 only)
-        print("  TEST: Layer 0 forward...")
-        let cmdBuf2 = engine.commandQueue.makeCommandBuffer()!
-        var hidden = try forwardLayerDebug(hidden: inputBuf, layerIdx: 0, seqLen: seqLen, cmdBuf: cmdBuf2)
-        cmdBuf2.commit()
-        cmdBuf2.waitUntilCompleted()
-        print("  TEST: Layer 0 OK")
-
-        // Full forward pass
-        print("  TEST: Full forward pass...")
         let cmdBuf = engine.commandQueue.makeCommandBuffer()!
-        hidden = inputBuf
+
+        let inputBuf = try lookupEmbeddings(tokens: tokens, cmdBuf: cmdBuf)
+
+        var hidden = inputBuf
         for idx in 0..<config.numHiddenLayers {
-            if idx % 6 == 0 { print("  Layer \(idx)/\(config.numHiddenLayers)") }
-            hidden = try forwardLayerDebug(hidden: hidden, layerIdx: idx, seqLen: seqLen, cmdBuf: cmdBuf)
+            hidden = try forwardLayer(hidden: hidden, layerIdx: idx, seqLen: seqLen, cmdBuf: cmdBuf)
         }
-        print("  TEST: All layers OK")
 
         let output = try applyRmsNorm(input: hidden, weight: finalNorm, count: seqLen * hs, cmdBuf: cmdBuf)
+
         cmdBuf.commit()
         cmdBuf.waitUntilCompleted()
 
@@ -140,10 +135,12 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         }
         let n = Float(seqLen)
         for i in 0..<hs { embedding[i] /= n }
+
         var norm: Float = 0
         for i in 0..<hs { norm += embedding[i] * embedding[i] }
         norm = sqrt(norm)
         if norm > 0 { for i in 0..<hs { embedding[i] /= norm } }
+
         return embedding
     }
 
@@ -162,39 +159,19 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         return engine.device.makeBuffer(bytes: data, length: data.count * 4)!
     }
 
-    private func loadAndTranspose(_ name: String, rows: Int, cols: Int) throws -> MTLBuffer {
-        // Load [rows, cols] and transpose to [cols, rows] for MPS matmul C = A × B
-        let data = try readTensor(name)
-        var transposed = [Float](repeating: 0, count: data.count)
-        for r in 0..<rows {
-            for c in 0..<cols {
-                transposed[c * rows + r] = data[r * cols + c]
-            }
-        }
-        return engine.device.makeBuffer(bytes: transposed, length: transposed.count * 4)!
-    }
-
     private func lookupEmbeddings(tokens: [Int], cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
         let seqLen = tokens.count, hs = config.hiddenSize
-        print("    lookupEmbeddings: seqLen=\(seqLen), hs=\(hs)")
-        // Read embedding table to CPU
+        // CPU-based embedding lookup
         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 dstStart = i * hs
             let srcStart = token * hs
-            embedData[start..<start+hs] = embedArray[srcStart..<srcStart+hs]
+            for j in 0..<hs {
+                embedData[dstStart + j] = embedPtr[srcStart + j]
+            }
         }
-        print("    Looked up \(seqLen) tokens")
-        
-        let buf = engine.device.makeBuffer(bytes: embedData, length: embedData.count * 4)!
-        print("    Created MTLBuffer")
-        return buf
+        return engine.device.makeBuffer(bytes: embedData, length: embedData.count * 4)!
     }
 
     private func applyRmsNorm(input: MTLBuffer, weight: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
@@ -209,29 +186,26 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         var c = UInt32(count), e: Float = config.rmsNormEps
         enc.setBytes(&c, length: 4, index: 3)
         enc.setBytes(&e, length: 4, index: 4)
-        enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
+        enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
+                           threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
         return output
     }
 
     private func matmulSeq(input: MTLBuffer, weight: MTLBuffer, output: MTLBuffer, m: Int, k: Int, n: Int, cmdBuf: MTLCommandBuffer) throws {
-        // Use MPS for optimized matrix multiplication on Apple Silicon
-        // Weight is stored transposed [k, n] for C = A × B
-        let descA = MPSMatrixDescriptor(rows: m, columns: k, rowBytes: k * 4, dataType: .float32)
-        let descB = MPSMatrixDescriptor(rows: k, columns: n, rowBytes: n * 4, dataType: .float32)
-        let descC = MPSMatrixDescriptor(rows: m, columns: n, rowBytes: n * 4, dataType: .float32)
-        let matA = MPSMatrix(buffer: input, descriptor: descA)
-        let matB = MPSMatrix(buffer: weight, descriptor: descB)
-        let matC = MPSMatrix(buffer: output, descriptor: descC)
-        
-        let matMul = MPSMatrixMultiplication(device: engine.device,
-                                              transposeLeft: false,
-                                              transposeRight: false,
-                                              resultRows: m,
-                                              resultColumns: n,
-                                              interiorColumns: k,
-                                              alpha: 1.0,
-                                              beta: 0.0)
-        matMul.encode(commandBuffer: cmdBuf, leftMatrix: matA, rightMatrix: matB, resultMatrix: matC)
+        let enc = cmdBuf.makeComputeCommandEncoder()!
+        let pso = try engine.pipeline(named: "matmul_f32")
+        enc.setComputePipelineState(pso)
+        enc.setBuffer(input, offset: 0, index: 0)
+        enc.setBuffer(weight, offset: 0, index: 1)
+        enc.setBuffer(output, offset: 0, index: 2)
+        var mm = UInt32(m), kk = UInt32(k), nn = UInt32(n)
+        enc.setBytes(&mm, length: 4, index: 3)
+        enc.setBytes(&kk, length: 4, index: 4)
+        enc.setBytes(&nn, length: 4, index: 5)
+        let total = m * n
+        enc.dispatchThreads(MTLSize(width: total, height: 1, depth: 1),
+                           threadsPerThreadgroup: MTLSize(width: min(256, total), height: 1, depth: 1))
+        enc.endEncoding()
     }
 
     private func eltwiseAdd(a: MTLBuffer, b: MTLBuffer, output: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws {
@@ -244,7 +218,8 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         enc.setBuffer(output, offset: 0, index: 2)
         var c = UInt32(count)
         enc.setBytes(&c, length: 4, index: 3)
-        enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
+        enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
+                           threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
     }
 
     private func geluMul(gate: MTLBuffer, up: MTLBuffer, output: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws {
@@ -257,7 +232,8 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         enc.setBuffer(output, offset: 0, index: 2)
         var c = UInt32(count)
         enc.setBytes(&c, length: 4, index: 3)
-        enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
+        enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
+                           threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
     }
 
     private func applyRoPE(q: MTLBuffer, k: MTLBuffer, seqLen: Int, headDim: Int, numHeads: Int, cmdBuf: MTLCommandBuffer) throws {
@@ -273,7 +249,8 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         enc.setBytes(&hd, length: 4, index: 3)
         enc.setBytes(&nh, length: 4, index: 4)
         enc.setBytes(&rt, length: 4, index: 5)
-        enc.dispatchThreads(MTLSize(width: numHeads * headDim / 2, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, numHeads * headDim / 2), height: 1, depth: 1))
+        enc.dispatchThreads(MTLSize(width: numHeads * headDim / 2, height: 1, depth: 1),
+                           threadsPerThreadgroup: MTLSize(width: min(256, numHeads * headDim / 2), height: 1, depth: 1))
     }
 
     private func bidirectionalAttention(q: MTLBuffer, k: MTLBuffer, v: MTLBuffer, output: MTLBuffer, seqLen: Int, cmdBuf: MTLCommandBuffer) throws {
@@ -296,81 +273,55 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         enc.setBytes(&scale, length: 4, index: 9)
         let tgMem = config.slidingWindow * 4
         enc.setThreadgroupMemoryLength(tgMem, index: 0)
-        enc.dispatchThreads(MTLSize(width: seqLen * config.numAttentionHeads, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, seqLen * config.numAttentionHeads), height: 1, depth: 1))
+        enc.dispatchThreads(MTLSize(width: seqLen * config.numAttentionHeads, height: 1, depth: 1),
+                           threadsPerThreadgroup: MTLSize(width: min(256, seqLen * config.numAttentionHeads), height: 1, depth: 1))
     }
 
-    private func forwardLayerDebug(hidden: MTLBuffer, layerIdx: Int, seqLen: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
+    private func forwardLayer(hidden: MTLBuffer, layerIdx: Int, seqLen: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
         let hs = config.hiddenSize, device = engine.device
         let hDim = config.headDim, nH = config.numAttentionHeads, nKV = config.numKeyValueHeads
         let intermedi = config.intermediateSize
 
-        print("    Residual copy...")
         let resid = device.makeBuffer(length: seqLen * hs * 4)!
         let blit = cmdBuf.makeBlitCommandEncoder()!
         blit.copy(from: hidden, sourceOffset: 0, to: resid, destinationOffset: 0, size: seqLen * hs * 4)
         blit.endEncoding()
 
-        print("    Input norm...")
         let h1 = try applyRmsNorm(input: hidden, weight: layerNorms[layerIdx][0], count: seqLen * hs, cmdBuf: cmdBuf)
 
-        print("    QKV projections...")
         let qBuf = device.makeBuffer(length: seqLen * nH * hDim * 4)!
         let kBuf = device.makeBuffer(length: seqLen * nKV * hDim * 4)!
         let vBuf = device.makeBuffer(length: seqLen * nKV * hDim * 4)!
         try matmulSeq(input: h1, weight: qProjs[layerIdx], output: qBuf, m: seqLen, k: hs, n: nH * hDim, cmdBuf: cmdBuf)
-        print("    Q done")
         try matmulSeq(input: h1, weight: kProjs[layerIdx], output: kBuf, m: seqLen, k: hs, n: nKV * hDim, cmdBuf: cmdBuf)
-        print("    K done")
         try matmulSeq(input: h1, weight: vProjs[layerIdx], output: vBuf, m: seqLen, k: hs, n: nKV * hDim, cmdBuf: cmdBuf)
-        print("    V done")
 
-        print("    RoPE...")
         try applyRoPE(q: qBuf, k: kBuf, seqLen: seqLen, headDim: hDim, numHeads: nH, cmdBuf: cmdBuf)
-        print("    RoPE done")
 
-        print("    Attention...")
         let attnOut = device.makeBuffer(length: seqLen * nH * hDim * 4)!
         try bidirectionalAttention(q: qBuf, k: kBuf, v: vBuf, output: attnOut, seqLen: seqLen, cmdBuf: cmdBuf)
-        print("    Attention done")
 
-        print("    O projection...")
         let h2 = device.makeBuffer(length: seqLen * hs * 4)!
         try matmulSeq(input: attnOut, weight: oProjs[layerIdx], output: h2, m: seqLen, k: nH * hDim, n: hs, cmdBuf: cmdBuf)
-        print("    O done")
 
-        print("    Post-attn norm...")
         let h2n = try applyRmsNorm(input: h2, weight: layerNorms[layerIdx][2], count: seqLen * hs, cmdBuf: cmdBuf)
-
-        print("    Add residual 1...")
         try eltwiseAdd(a: resid, b: h2n, output: hidden, count: seqLen * hs, cmdBuf: cmdBuf)
 
-        print("    Pre-FF norm...")
         let h3 = try applyRmsNorm(input: hidden, weight: layerNorms[layerIdx][1], count: seqLen * hs, cmdBuf: cmdBuf)
 
-        print("    MLP gate/up...")
         let gate = device.makeBuffer(length: seqLen * intermedi * 4)!
         let up = device.makeBuffer(length: seqLen * intermedi * 4)!
         try matmulSeq(input: h3, weight: gateProjs[layerIdx], output: gate, m: seqLen, k: hs, n: intermedi, cmdBuf: cmdBuf)
-        print("    Gate done")
         try matmulSeq(input: h3, weight: upProjs[layerIdx], output: up, m: seqLen, k: hs, n: intermedi, cmdBuf: cmdBuf)
-        print("    Up done")
 
-        print("    GELU mul...")
         let gated = device.makeBuffer(length: seqLen * intermedi * 4)!
         try geluMul(gate: gate, up: up, output: gated, count: seqLen * intermedi, cmdBuf: cmdBuf)
-        print("    GELU done")
 
-        print("    Down projection...")
         let h4 = device.makeBuffer(length: seqLen * hs * 4)!
         try matmulSeq(input: gated, weight: downProjs[layerIdx], output: h4, m: seqLen, k: intermedi, n: hs, cmdBuf: cmdBuf)
-        print("    Down done")
 
-        print("    Post-FF norm...")
         let h4n = try applyRmsNorm(input: h4, weight: layerNorms[layerIdx][3], count: seqLen * hs, cmdBuf: cmdBuf)
-
-        print("    Add residual 2...")
         try eltwiseAdd(a: hidden, b: h4n, output: hidden, count: seqLen * hs, cmdBuf: cmdBuf)
-        print("    Layer \(layerIdx) complete")
 
         return hidden
     }