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

v2: add EmbeddingServer binary, EmbeddingGemma model (WIP forward pass)

Browse Source
This commit is contained in:
MarkBase Admin
2026-07-06 10:13:39 +08:00
parent e7a94b3203
commit dbec6b20ea
3 changed files with 104 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift
+15 -9
View File
@@ -284,15 +284,21 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
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)
enc.dispatchThreads(MTLSize(width: m * n, height: 1, depth: 1),
threadsPerThreadgroup: MTLSize(width: min(256, m * n), height: 1, depth: 1))
// Note: matmul_f32 only handles M=1, so we loop over positions
// For production: implement a proper batch matmul kernel
for i in 0..<m {
let inputOffset = i * k * 4
let outputOffset = i * n * 4
enc.setBuffer(input, offset: inputOffset, index: 0)
enc.setBuffer(weight, offset: 0, index: 1)
enc.setBuffer(output, offset: outputOffset, index: 2)
var mm: UInt32 = 1, 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)
enc.dispatchThreads(MTLSize(width: n, height: 1, depth: 1),
threadsPerThreadgroup: MTLSize(width: min(256, n), height: 1, depth: 1))
}
enc.endEncoding()
}