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

v2: use MPSMatrixMultiplication for matmul in EmbeddingGemma

Browse Source
This commit is contained in:
MarkBase Admin
2026-07-06 14:13:34 +08:00
parent c48983a413
commit 31d5e8adaf
1 changed files with 38 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift
+38 -23
View File
@@ -1,5 +1,6 @@
import Foundation
import Metal
import MetalPerformanceShaders
import Accelerate
/// EmbeddingGemma configuration
@@ -80,15 +81,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"))
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"))
qProjs.append(try loadAndTranspose("\(p).self_attn.q_proj.weight", rows: config.hiddenSize, cols: config.hiddenSize))
kProjs.append(try loadAndTranspose("\(p).self_attn.k_proj.weight", rows: config.numKeyValueHeads * config.headDim, cols: config.hiddenSize))
vProjs.append(try loadAndTranspose("\(p).self_attn.v_proj.weight", rows: config.numKeyValueHeads * config.headDim, cols: config.hiddenSize))
oProjs.append(try loadAndTranspose("\(p).self_attn.o_proj.weight", rows: config.numAttentionHeads * config.headDim, cols: config.hiddenSize))
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"))
gateProjs.append(try loadAndTranspose("\(p).mlp.gate_proj.weight", rows: config.intermediateSize, cols: config.hiddenSize))
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)!
@@ -161,6 +162,18 @@ 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
let buf = engine.device.makeBuffer(length: seqLen * hs * 4)!
@@ -196,22 +209,24 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
}
private func matmulSeq(input: MTLBuffer, weight: MTLBuffer, output: MTLBuffer, m: Int, k: Int, n: Int, cmdBuf: MTLCommandBuffer) throws {
let enc = cmdBuf.makeComputeCommandEncoder()!
defer { enc.endEncoding() }
let pso = try engine.pipeline(named: "matmul_f32")
enc.setComputePipelineState(pso)
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))
}
// 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)
}
private func eltwiseAdd(a: MTLBuffer, b: MTLBuffer, output: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws {