warren/markbaseengine
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
MarkBase Admin
2026-07-06 15:38:17 +08:00
parent 31d5e8adaf
commit 6af56f58ea
1 changed files with 18 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift
+18 -13
View File
@@ -176,19 +176,24 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
private func lookupEmbeddings(tokens: [Int], cmdBuf: MTLCommandBuffer) throws -> MTLBuffer { private func lookupEmbeddings(tokens: [Int], cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
let seqLen = tokens.count, hs = config.hiddenSize let seqLen = tokens.count, hs = config.hiddenSize
let buf = engine.device.makeBuffer(length: seqLen * hs * 4)! print(" lookupEmbeddings: seqLen=\(seqLen), hs=\(hs)")
let enc = cmdBuf.makeComputeCommandEncoder()! // Read embedding table to CPU
defer { enc.endEncoding() } let embedPtr = embedTokens.contents().assumingMemoryBound(to: Float.self)
let pso = try engine.pipeline(named: "lookup_embeddings") let embedCount = embedTokens.length / 4
enc.setComputePipelineState(pso) let embedArray = Array(UnsafeBufferPointer(start: embedPtr, count: embedCount))
enc.setBuffer(embedTokens, offset: 0, index: 0) print(" Read \(embedCount) floats from embedTokens")
enc.setBytes(tokens, length: seqLen * MemoryLayout<Int>.size, index: 1)
enc.setBuffer(buf, offset: 0, index: 2) // Lookup embeddings for tokens
var h = UInt32(hs), s = UInt32(seqLen), v = UInt32(config.vocabSize) var embedData = [Float](repeating: 0, count: seqLen * hs)
enc.setBytes(&h, length: 4, index: 3) for (i, token) in tokens.enumerated() {
enc.setBytes(&s, length: 4, index: 4) let start = i * hs
enc.setBytes(&v, length: 4, index: 5) let srcStart = token * hs
enc.dispatchThreads(MTLSize(width: seqLen, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, seqLen), height: 1, depth: 1)) 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 return buf
} }