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

v2: add defer enc.endEncoding() to all Metal encoders

Browse Source
This commit is contained in:
MarkBase Admin
2026-07-06 13:05:30 +08:00
parent f122d854dc
commit 492b779634
1 changed files with 15 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift
+15 -7
View File
@@ -106,24 +106,32 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
guard !tokens.isEmpty else { return [] }
let seqLen = tokens.count, hs = config.hiddenSize
print(" embed: seqLen=\(seqLen)")
// Use single command buffer for entire forward pass
let cmdBuf = engine.commandQueue.makeCommandBuffer()!
// Embedding lookup
print(" embed: lookupEmbeddings")
let inputBuf = try lookupEmbeddings(tokens: tokens, cmdBuf: cmdBuf)
print(" embed: lookupEmbeddings done")
// Forward through layers
var hidden = inputBuf
for idx in 0..<config.numHiddenLayers {
if idx % 6 == 0 { print(" embed: layer \(idx)/\(config.numHiddenLayers)") }
hidden = try forwardLayer(hidden: hidden, layerIdx: idx, seqLen: seqLen, cmdBuf: cmdBuf)
}
print(" embed: all layers done")
// Final norm
print(" embed: final norm")
let output = try applyRmsNorm(input: hidden, weight: finalNorm, count: seqLen * hs, cmdBuf: cmdBuf)
print(" embed: committing")
cmdBuf.commit()
cmdBuf.waitUntilCompleted()
print(" embed: completed")
// Readback
let data = engine.readFloats(from: output, count: seqLen * hs)
@@ -171,6 +179,7 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
let seqLen = tokens.count, hs = config.hiddenSize
let buf = engine.device.makeBuffer(length: seqLen * hs * 4)!
let enc = cmdBuf.makeComputeCommandEncoder()!
defer { enc.endEncoding() }
let pso = try engine.pipeline(named: "lookup_embeddings")
enc.setComputePipelineState(pso)
enc.setBuffer(embedTokens, offset: 0, index: 0)
@@ -182,13 +191,13 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
enc.setBytes(&v, length: 4, index: 5)
enc.dispatchThreads(MTLSize(width: seqLen, height: 1, depth: 1),
threadsPerThreadgroup: MTLSize(width: min(256, seqLen), height: 1, depth: 1))
enc.endEncoding()
return buf
}
private func applyRmsNorm(input: MTLBuffer, weight: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
let output = engine.device.makeBuffer(length: count * 4)!
let enc = cmdBuf.makeComputeCommandEncoder()!
defer { enc.endEncoding() }
let pso = try engine.pipeline(named: "rms_norm")
enc.setComputePipelineState(pso)
enc.setBuffer(input, offset: 0, index: 0)
@@ -199,12 +208,12 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
enc.setBytes(&e, length: 4, index: 4)
enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
enc.endEncoding()
return output
}
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 {
@@ -220,11 +229,11 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
enc.dispatchThreads(MTLSize(width: n, height: 1, depth: 1),
threadsPerThreadgroup: MTLSize(width: min(256, n), height: 1, depth: 1))
}
enc.endEncoding()
}
private func eltwiseAdd(a: MTLBuffer, b: MTLBuffer, output: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws {
let enc = cmdBuf.makeComputeCommandEncoder()!
defer { enc.endEncoding() }
let pso = try engine.pipeline(named: "eltwise_add")
enc.setComputePipelineState(pso)
enc.setBuffer(a, offset: 0, index: 0)
@@ -234,11 +243,11 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
enc.setBytes(&c, length: 4, index: 3)
enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
enc.endEncoding()
}
private func geluMul(gate: MTLBuffer, up: MTLBuffer, output: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws {
let enc = cmdBuf.makeComputeCommandEncoder()!
defer { enc.endEncoding() }
let pso = try engine.pipeline(named: "gelu_mul_kernel")
enc.setComputePipelineState(pso)
enc.setBuffer(gate, offset: 0, index: 0)
@@ -248,11 +257,11 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
enc.setBytes(&c, length: 4, index: 3)
enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
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()!
defer { enc.endEncoding() }
let pso = try engine.pipeline(named: "apply_rope")
enc.setComputePipelineState(pso)
enc.setBuffer(q, offset: 0, index: 0)
@@ -265,11 +274,11 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
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()!
defer { enc.endEncoding() }
let pso = try engine.pipeline(named: "bidirectional_sliding_attn")
enc.setComputePipelineState(pso)
enc.setBuffer(q, offset: 0, index: 0)
@@ -289,7 +298,6 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
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 {