warren/markbaseengine
1
0
Fork 0
Code Issues Pull Requests Actions 15 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 Foundation
import Metal import Metal
import MetalPerformanceShaders
import Accelerate import Accelerate
/// EmbeddingGemma configuration /// EmbeddingGemma configuration
@@ -80,15 +81,15 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
try loadBuffer("\(p).post_attention_layernorm.weight"), try loadBuffer("\(p).post_attention_layernorm.weight"),
try loadBuffer("\(p).post_feedforward_layernorm.weight"), try loadBuffer("\(p).post_feedforward_layernorm.weight"),
]) ])
qProjs.append(try loadBuffer("\(p).self_attn.q_proj.weight")) qProjs.append(try loadAndTranspose("\(p).self_attn.q_proj.weight", rows: config.hiddenSize, cols: config.hiddenSize))
kProjs.append(try loadBuffer("\(p).self_attn.k_proj.weight")) kProjs.append(try loadAndTranspose("\(p).self_attn.k_proj.weight", rows: config.numKeyValueHeads * config.headDim, cols: config.hiddenSize))
vProjs.append(try loadBuffer("\(p).self_attn.v_proj.weight")) vProjs.append(try loadAndTranspose("\(p).self_attn.v_proj.weight", rows: config.numKeyValueHeads * config.headDim, cols: config.hiddenSize))
oProjs.append(try loadBuffer("\(p).self_attn.o_proj.weight")) 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")) qNorms.append(try loadBuffer("\(p).self_attn.q_norm.weight"))
kNorms.append(try loadBuffer("\(p).self_attn.k_norm.weight")) kNorms.append(try loadBuffer("\(p).self_attn.k_norm.weight"))
gateProjs.append(try loadBuffer("\(p).mlp.gate_proj.weight")) gateProjs.append(try loadAndTranspose("\(p).mlp.gate_proj.weight", rows: config.intermediateSize, cols: config.hiddenSize))
upProjs.append(try loadBuffer("\(p).mlp.up_proj.weight")) upProjs.append(try loadAndTranspose("\(p).mlp.up_proj.weight", rows: config.intermediateSize, cols: config.hiddenSize))
downProjs.append(try loadBuffer("\(p).mlp.down_proj.weight")) downProjs.append(try loadAndTranspose("\(p).mlp.down_proj.weight", rows: config.hiddenSize, cols: config.intermediateSize))
} }
let fnData = try readTensor("norm.weight") let fnData = try readTensor("norm.weight")
finalNorm = engine.device.makeBuffer(bytes: fnData, length: fnData.count * 4)! 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)! 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 { 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)! 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 { private func matmulSeq(input: MTLBuffer, weight: MTLBuffer, output: MTLBuffer, m: Int, k: Int, n: Int, cmdBuf: MTLCommandBuffer) throws {
let enc = cmdBuf.makeComputeCommandEncoder()! // Use MPS for optimized matrix multiplication on Apple Silicon
defer { enc.endEncoding() } // Weight is stored transposed [k, n] for C = A × B
let pso = try engine.pipeline(named: "matmul_f32") let descA = MPSMatrixDescriptor(rows: m, columns: k, rowBytes: k * 4, dataType: .float32)
enc.setComputePipelineState(pso) let descB = MPSMatrixDescriptor(rows: k, columns: n, rowBytes: n * 4, dataType: .float32)
for i in 0..<m { let descC = MPSMatrixDescriptor(rows: m, columns: n, rowBytes: n * 4, dataType: .float32)
let inputOffset = i * k * 4 let matA = MPSMatrix(buffer: input, descriptor: descA)
let outputOffset = i * n * 4 let matB = MPSMatrix(buffer: weight, descriptor: descB)
enc.setBuffer(input, offset: inputOffset, index: 0) let matC = MPSMatrix(buffer: output, descriptor: descC)
enc.setBuffer(weight, offset: 0, index: 1)
enc.setBuffer(output, offset: outputOffset, index: 2) let matMul = MPSMatrixMultiplication(device: engine.device,
var mm: UInt32 = 1, kk = UInt32(k), nn = UInt32(n) transposeLeft: false,
enc.setBytes(&mm, length: 4, index: 3) transposeRight: false,
enc.setBytes(&kk, length: 4, index: 4) resultRows: m,
enc.setBytes(&nn, length: 4, index: 5) resultColumns: n,
enc.dispatchThreads(MTLSize(width: n, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, n), height: 1, depth: 1)) 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 { private func eltwiseAdd(a: MTLBuffer, b: MTLBuffer, output: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws {