diff --git a/Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift b/Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift index 6245486..82b7b36 100644 --- a/Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift +++ b/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.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..