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

v2: EmbeddingGemma - single cmdBuf fix, needs Metal kernel compilation

Browse Source
This commit is contained in:
MarkBase Admin
2026-07-06 11:37:10 +08:00
parent dbec6b20ea
commit 5e060c7aea
2 changed files with 157 additions and 147 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift
+128 -147
View File
@@ -2,7 +2,7 @@ import Foundation
import Metal
import Accelerate
/// EmbeddingGemmaConfig - Configuration for EmbeddingGemma model
/// EmbeddingGemma configuration
public struct EmbeddingGemmaConfig: Codable {
public let hiddenSize: Int
public let numHiddenLayers: Int
@@ -73,12 +73,6 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
}
private func loadWeights() throws {
let hs = config.hiddenSize
let intermedi = config.intermediateSize
let nKV = config.numKeyValueHeads
let hDim = config.headDim
// Embedding table [vocab, hidden]
let embedData = try readTensor("embed_tokens.weight")
embedTokens = engine.device.makeBuffer(bytes: embedData, length: embedData.count * 4)!
@@ -90,15 +84,15 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
try loadBuffer("\(p).post_attention_layernorm.weight"),
try loadBuffer("\(p).post_feedforward_layernorm.weight"),
])
qProjs.append(try loadBuffer("\(p).self_attn.q_proj.weight")) // [hs, hs]
kProjs.append(try loadBuffer("\(p).self_attn.k_proj.weight")) // [nKV*hDim, hs]
vProjs.append(try loadBuffer("\(p).self_attn.v_proj.weight")) // [nKV*hDim, hs]
oProjs.append(try loadBuffer("\(p).self_attn.o_proj.weight")) // [hs, nH*hDim]
qNorms.append(try loadBuffer("\(p).self_attn.q_norm.weight")) // [hDim]
kNorms.append(try loadBuffer("\(p).self_attn.k_norm.weight")) // [hDim]
gateProjs.append(try loadBuffer("\(p).mlp.gate_proj.weight")) // [intermedi, hs]
upProjs.append(try loadBuffer("\(p).mlp.up_proj.weight")) // [intermedi, hs]
downProjs.append(try loadBuffer("\(p).mlp.down_proj.weight")) // [hs, intermedi]
qProjs.append(try loadBuffer("\(p).self_attn.q_proj.weight"))
kProjs.append(try loadBuffer("\(p).self_attn.k_proj.weight"))
vProjs.append(try loadBuffer("\(p).self_attn.v_proj.weight"))
oProjs.append(try loadBuffer("\(p).self_attn.o_proj.weight"))
qNorms.append(try loadBuffer("\(p).self_attn.q_norm.weight"))
kNorms.append(try loadBuffer("\(p).self_attn.k_norm.weight"))
gateProjs.append(try loadBuffer("\(p).mlp.gate_proj.weight"))
upProjs.append(try loadBuffer("\(p).mlp.up_proj.weight"))
downProjs.append(try loadBuffer("\(p).mlp.down_proj.weight"))
}
let fnData = try readTensor("norm.weight")
@@ -113,17 +107,23 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
let seqLen = tokens.count, hs = config.hiddenSize
// Use single command buffer for entire forward pass
let cmdBuf = engine.commandQueue.makeCommandBuffer()!
// Embedding lookup
let inputBuf = try lookupEmbeddings(tokens: tokens)
let inputBuf = try lookupEmbeddings(tokens: tokens, cmdBuf: cmdBuf)
// Forward through layers
var hidden = inputBuf
for idx in 0..<config.numHiddenLayers {
hidden = try forwardLayer(hidden: hidden, layerIdx: idx, seqLen: seqLen)
hidden = try forwardLayer(hidden: hidden, layerIdx: idx, seqLen: seqLen, cmdBuf: cmdBuf)
}
// Final norm
let output = try applyRmsNorm(input: hidden, weight: finalNorm, count: seqLen * hs)
let output = try applyRmsNorm(input: hidden, weight: finalNorm, count: seqLen * hs, cmdBuf: cmdBuf)
cmdBuf.commit()
cmdBuf.waitUntilCompleted()
// Readback
let data = engine.readFloats(from: output, count: seqLen * hs)
@@ -145,7 +145,7 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
return embedding
}
// MARK: - Helpers
// MARK: - Helpers (all use shared cmdBuf)
private func readTensor(_ name: String) throws -> [Float] {
guard let desc = reader.tensor(named: name) else {
@@ -167,10 +167,9 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
return engine.device.makeBuffer(bytes: data, length: data.count * 4)!
}
private func lookupEmbeddings(tokens: [Int]) throws -> MTLBuffer {
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 cmdBuf = engine.commandQueue.makeCommandBuffer()!
let enc = cmdBuf.makeComputeCommandEncoder()!
let pso = try engine.pipeline(named: "lookup_embeddings")
enc.setComputePipelineState(pso)
@@ -184,13 +183,11 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
enc.dispatchThreads(MTLSize(width: seqLen, height: 1, depth: 1),
threadsPerThreadgroup: MTLSize(width: min(256, seqLen), height: 1, depth: 1))
enc.endEncoding()
cmdBuf.commit(); cmdBuf.waitUntilCompleted()
return buf
}
private func applyRmsNorm(input: MTLBuffer, weight: MTLBuffer, count: Int) throws -> MTLBuffer {
private func applyRmsNorm(input: MTLBuffer, weight: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
let output = engine.device.makeBuffer(length: count * 4)!
let cmdBuf = engine.commandQueue.makeCommandBuffer()!
let enc = cmdBuf.makeComputeCommandEncoder()!
let pso = try engine.pipeline(named: "rms_norm")
enc.setComputePipelineState(pso)
@@ -203,89 +200,13 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
enc.endEncoding()
cmdBuf.commit(); cmdBuf.waitUntilCompleted()
return output
}
private func forwardLayer(hidden: MTLBuffer, layerIdx: Int, seqLen: Int) throws -> MTLBuffer {
let hs = config.hiddenSize, device = engine.device
let hDim = config.headDim, nH = config.numAttentionHeads, nKV = config.numKeyValueHeads
let intermedi = config.intermediateSize
let cmdBuf = engine.commandQueue.makeCommandBuffer()!
// Residual
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()
// Input norm
let h1 = try applyRmsNorm(input: hidden, weight: layerNorms[layerIdx][0], count: seqLen * hs)
// Q, K, V projections (using optimized matmul)
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)
try matmulSeq(input: h1, weight: kProjs[layerIdx], output: kBuf, m: seqLen, k: hs, n: nKV * hDim, cmdBuf: cmdBuf)
try matmulSeq(input: h1, weight: vProjs[layerIdx], output: vBuf, m: seqLen, k: hs, n: nKV * hDim, cmdBuf: cmdBuf)
// RoPE
try applyRoPE(q: qBuf, k: kBuf, seqLen: seqLen, headDim: hDim, numHeads: nH, numKVHeads: nKV, cmdBuf: cmdBuf)
// Q/K Norm
try applyQKNorm(q: qBuf, k: kBuf, qNorm: qNorms[layerIdx], kNorm: kNorms[layerIdx], seqLen: seqLen, headDim: hDim, numHeads: nH, numKVHeads: nKV, cmdBuf: cmdBuf)
// Bidirectional sliding window attention
let attnOut = device.makeBuffer(length: seqLen * nH * hDim * 4)!
try bidirectionalAttention(q: qBuf, k: kBuf, v: vBuf, output: attnOut, seqLen: seqLen, cmdBuf: cmdBuf)
// 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)
// Post-attn norm
let h2n = try applyRmsNorm(input: h2, weight: layerNorms[layerIdx][2], count: seqLen * hs)
// Add residual: hidden = resid + h2n
try eltwiseAdd(a: resid, b: h2n, output: hidden, count: seqLen * hs, cmdBuf: cmdBuf)
// Pre-FF norm
let h3 = try applyRmsNorm(input: hidden, weight: layerNorms[layerIdx][1], count: seqLen * hs)
// 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)
try matmulSeq(input: h3, weight: upProjs[layerIdx], output: up, m: seqLen, k: hs, n: intermedi, cmdBuf: cmdBuf)
// GELU(gate) * up
let gated = device.makeBuffer(length: seqLen * intermedi * 4)!
try geluMul(gate: gate, up: up, output: gated, count: seqLen * intermedi, cmdBuf: cmdBuf)
// 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)
// Post-FF norm
let h4n = try applyRmsNorm(input: h4, weight: layerNorms[layerIdx][3], count: seqLen * hs)
// Add residual: hidden = hidden + h4n
try eltwiseAdd(a: hidden, b: h4n, output: hidden, count: seqLen * hs, cmdBuf: cmdBuf)
cmdBuf.commit(); cmdBuf.waitUntilCompleted()
return hidden
}
// MARK: - Metal Kernels
private func matmulSeq(input: MTLBuffer, weight: MTLBuffer, output: MTLBuffer, m: Int, k: Int, n: Int, cmdBuf: MTLCommandBuffer) throws {
let enc = cmdBuf.makeComputeCommandEncoder()!
let pso = try engine.pipeline(named: "matmul_f32")
enc.setComputePipelineState(pso)
// 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
@@ -302,52 +223,6 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
enc.endEncoding()
}
private func applyRoPE(q: MTLBuffer, k: MTLBuffer, seqLen: Int, headDim: Int, numHeads: Int, numKVHeads: Int, cmdBuf: MTLCommandBuffer) throws {
let enc = cmdBuf.makeComputeCommandEncoder()!
let pso = try engine.pipeline(named: "apply_rope")
enc.setComputePipelineState(pso)
enc.setBuffer(q, offset: 0, index: 0)
enc.setBuffer(k, offset: 0, index: 1)
var sl = UInt32(seqLen), hd = UInt32(headDim), nh = UInt32(numHeads)
var rt: Float = Float(config.ropeTheta)
enc.setBytes(&sl, length: 4, index: 2)
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.endEncoding()
}
private func applyQKNorm(q: MTLBuffer, k: MTLBuffer, qNorm: MTLBuffer, kNorm: MTLBuffer, seqLen: Int, headDim: Int, numHeads: Int, numKVHeads: Int, cmdBuf: MTLCommandBuffer) throws {
// Apply RMSNorm per head
// TODO: Implement per-head normalization
}
private func bidirectionalAttention(q: MTLBuffer, k: MTLBuffer, v: MTLBuffer, output: MTLBuffer, seqLen: Int, cmdBuf: MTLCommandBuffer) throws {
let enc = cmdBuf.makeComputeCommandEncoder()!
let pso = try engine.pipeline(named: "bidirectional_sliding_attn")
enc.setComputePipelineState(pso)
enc.setBuffer(q, offset: 0, index: 0)
enc.setBuffer(k, offset: 0, index: 1)
enc.setBuffer(v, offset: 0, index: 2)
enc.setBuffer(output, offset: 0, index: 3)
var sl = UInt32(seqLen), hd = UInt32(config.headDim), nh = UInt32(config.numAttentionHeads)
var nkv = UInt32(config.numKeyValueHeads), sw = UInt32(config.slidingWindow)
var scale: Float = 1.0 / sqrt(Float(config.headDim))
enc.setBytes(&sl, length: 4, index: 4)
enc.setBytes(&hd, length: 4, index: 5)
enc.setBytes(&nh, length: 4, index: 6)
enc.setBytes(&nkv, length: 4, index: 7)
enc.setBytes(&sw, length: 4, index: 8)
enc.setBytes(&scale, length: 4, index: 9)
let tgMem = config.slidingWindow * 4 // shared memory for scores
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.endEncoding()
}
private func eltwiseAdd(a: MTLBuffer, b: MTLBuffer, output: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws {
let enc = cmdBuf.makeComputeCommandEncoder()!
let pso = try engine.pipeline(named: "eltwise_add")
@@ -375,4 +250,110 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
enc.endEncoding()
}
private func applyRoPE(q: MTLBuffer, k: MTLBuffer, seqLen: Int, headDim: Int, numHeads: Int, cmdBuf: MTLCommandBuffer) throws {
let enc = cmdBuf.makeComputeCommandEncoder()!
let pso = try engine.pipeline(named: "apply_rope")
enc.setComputePipelineState(pso)
enc.setBuffer(q, offset: 0, index: 0)
enc.setBuffer(k, offset: 0, index: 1)
var sl = UInt32(seqLen), hd = UInt32(headDim), nh = UInt32(numHeads)
var rt: Float = Float(config.ropeTheta)
enc.setBytes(&sl, length: 4, index: 2)
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.endEncoding()
}
private func bidirectionalAttention(q: MTLBuffer, k: MTLBuffer, v: MTLBuffer, output: MTLBuffer, seqLen: Int, cmdBuf: MTLCommandBuffer) throws {
let enc = cmdBuf.makeComputeCommandEncoder()!
let pso = try engine.pipeline(named: "bidirectional_sliding_attn")
enc.setComputePipelineState(pso)
enc.setBuffer(q, offset: 0, index: 0)
enc.setBuffer(k, offset: 0, index: 1)
enc.setBuffer(v, offset: 0, index: 2)
enc.setBuffer(output, offset: 0, index: 3)
var sl = UInt32(seqLen), hd = UInt32(config.headDim), nh = UInt32(config.numAttentionHeads)
var nkv = UInt32(config.numKeyValueHeads), sw = UInt32(config.slidingWindow)
var scale: Float = 1.0 / sqrt(Float(config.headDim))
enc.setBytes(&sl, length: 4, index: 4)
enc.setBytes(&hd, length: 4, index: 5)
enc.setBytes(&nh, length: 4, index: 6)
enc.setBytes(&nkv, length: 4, index: 7)
enc.setBytes(&sw, length: 4, index: 8)
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.endEncoding()
}
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
// Residual
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()
// Input norm
let h1 = try applyRmsNorm(input: hidden, weight: layerNorms[layerIdx][0], count: seqLen * hs, cmdBuf: cmdBuf)
// Q, K, V 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)
try matmulSeq(input: h1, weight: kProjs[layerIdx], output: kBuf, m: seqLen, k: hs, n: nKV * hDim, cmdBuf: cmdBuf)
try matmulSeq(input: h1, weight: vProjs[layerIdx], output: vBuf, m: seqLen, k: hs, n: nKV * hDim, cmdBuf: cmdBuf)
// RoPE
try applyRoPE(q: qBuf, k: kBuf, seqLen: seqLen, headDim: hDim, numHeads: nH, cmdBuf: cmdBuf)
// Bidirectional sliding window attention
let attnOut = device.makeBuffer(length: seqLen * nH * hDim * 4)!
try bidirectionalAttention(q: qBuf, k: kBuf, v: vBuf, output: attnOut, seqLen: seqLen, cmdBuf: cmdBuf)
// 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)
// Post-attn norm
let h2n = try applyRmsNorm(input: h2, weight: layerNorms[layerIdx][2], count: seqLen * hs, cmdBuf: cmdBuf)
// Add residual: hidden = resid + h2n
try eltwiseAdd(a: resid, b: h2n, output: hidden, count: seqLen * hs, cmdBuf: cmdBuf)
// Pre-FF norm
let h3 = try applyRmsNorm(input: hidden, weight: layerNorms[layerIdx][1], count: seqLen * hs, cmdBuf: cmdBuf)
// 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)
try matmulSeq(input: h3, weight: upProjs[layerIdx], output: up, m: seqLen, k: hs, n: intermedi, cmdBuf: cmdBuf)
// GELU(gate) * up
let gated = device.makeBuffer(length: seqLen * intermedi * 4)!
try geluMul(gate: gate, up: up, output: gated, count: seqLen * intermedi, cmdBuf: cmdBuf)
// 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)
// Post-FF norm
let h4n = try applyRmsNorm(input: h4, weight: layerNorms[layerIdx][3], count: seqLen * hs, cmdBuf: cmdBuf)
// Add residual: hidden = hidden + h4n
try eltwiseAdd(a: hidden, b: h4n, output: hidden, count: seqLen * hs, cmdBuf: cmdBuf)
return hidden
}
}
Tests/01_Model/EmbeddingGemmaTest.swift
+29
View File
@@ -0,0 +1,29 @@
import XCTest
@testable import MarkBase
final class EmbeddingGemmaTest: XCTestCase {
func testLoadAndEmbed() throws {
let modelDir = "/Users/accusys/MarkBaseEngine/models/embeddinggemma-300m"
guard FileManager.default.fileExists(atPath: modelDir + "/model.safetensors") else {
XCTFail("Model not found"); return
}
let engine = try MarkBaseEngine(autoCompile: true)
let model = try EmbeddingGemmaModel(modelDir: modelDir, engine: engine)
XCTAssertEqual(model.config.numHiddenLayers, 24)
XCTAssertEqual(model.config.hiddenSize, 768)
let emb = try model.embed(text: "Hello world")
XCTAssertEqual(emb.count, 768)
// Check L2 norm is 1.0
var norm: Float = 0
for v in emb { norm += v * v }
norm = sqrt(norm)
XCTAssertEqual(norm, 1.0, accuracy: 0.001, "L2 norm should be 1.0")
// Check no NaN
XCTAssertFalse(emb.contains { $0.isNaN }, "No NaN values")
print("✓ Embedding dim=\(emb.count), norm=\(String(format: "%.4f", norm))")
}
}