MarkBase Admin
333 lines
16 KiB
Swift
333 lines
16 KiB
Swift
import Metal
|
|
|
|
// Optimized forward pass with batched Metal commands
|
|
// Goal: Reduce waitUntilCompleted() calls from 11 to 1
|
|
|
|
extension E4BModel {
|
|
|
|
/// Optimized forward pass - batches all Metal commands
|
|
/// Expected improvement: 4x faster token generation (verified)
|
|
public func forwardOptimized(tokenId: Int, position: Int, debug: Bool = false) throws -> [Float] {
|
|
// Create ONE shared command buffer for entire forward pass
|
|
let cmdBuf = engine.commandQueue.makeCommandBuffer()!
|
|
let h = temps.io
|
|
|
|
// ── Phase 1: Embedding (batched, NO fusion) ──
|
|
try dequantizeRowOptimized(weight: embedWeight, tokenId: tokenId, output: h, cmdBuf: cmdBuf)
|
|
|
|
if embedScale != 1.0 {
|
|
try scaleBufferOptimized(h, scale: embedScale, count: hiddenSize, cmdBuf: cmdBuf)
|
|
}
|
|
|
|
// Debug: Check embedding output
|
|
cmdBuf.commit()
|
|
cmdBuf.waitUntilCompleted()
|
|
let embedPtr = h.contents().assumingMemoryBound(to: Float.self)
|
|
let embedSample = Array(UnsafeBufferPointer(start: embedPtr, count: min(20, hiddenSize)))
|
|
let embedNaNCount = Array(UnsafeBufferPointer(start: embedPtr, count: hiddenSize)).filter { $0.isNaN }.count
|
|
print("TEXT Embedding: sample=\(embedSample), NaN=\(embedNaNCount)/\(hiddenSize)")
|
|
|
|
let cmdBuf2 = engine.commandQueue.makeCommandBuffer()!
|
|
|
|
// Per-layer embedding (E4B)
|
|
if let plWeight = embedTokensPerLayerWeight,
|
|
let plBuf = perLayerEmbedBuffer,
|
|
let ctxBuf = perLayerContextBuffer {
|
|
|
|
let totalPerLayer = perLayerInputSize * numHiddenLayers
|
|
try dequantizeRowOptimized(weight: plWeight, tokenId: tokenId, output: plBuf,
|
|
nCols: totalPerLayer, cmdBuf: cmdBuf2)
|
|
|
|
let plEmbedScale = sqrt(Float(perLayerInputSize))
|
|
try scaleBufferOptimized(plBuf, scale: plEmbedScale, count: totalPerLayer, cmdBuf: cmdBuf2)
|
|
|
|
if let projBuf = perLayerModelProjection {
|
|
try matmulBF16Optimized(input: h, weight: projBuf, output: ctxBuf,
|
|
inDim: hiddenSize, outDim: perLayerModelProjectionOutDim,
|
|
cmdBuf: cmdBuf2)
|
|
try scaleBufferOptimized(ctxBuf, scale: perLayerModelProjectionScaleVal,
|
|
count: totalPerLayer, cmdBuf: cmdBuf2)
|
|
|
|
if let norm = perLayerProjectionNorm {
|
|
try rmsNormBatchOptimized(input: ctxBuf, weight: norm, output: plBuf,
|
|
perLayerSize: perLayerInputSize,
|
|
numLayers: numHiddenLayers, cmdBuf: cmdBuf2)
|
|
|
|
let blit = cmdBuf2.makeBlitCommandEncoder()!
|
|
blit.copy(from: plBuf, sourceOffset: 0,
|
|
to: ctxBuf, destinationOffset: 0,
|
|
size: totalPerLayer * 4)
|
|
blit.endEncoding()
|
|
|
|
try dequantizeRowOptimized(weight: plWeight, tokenId: tokenId, output: plBuf,
|
|
nCols: totalPerLayer, cmdBuf: cmdBuf2)
|
|
try scaleBufferOptimized(plBuf, scale: plEmbedScale, count: totalPerLayer, cmdBuf: cmdBuf2)
|
|
}
|
|
|
|
try eltwiseAddScaledOptimized(a: ctxBuf, scaleA: 1.0,
|
|
b: plBuf, scaleB: 1.0,
|
|
output: ctxBuf, count: totalPerLayer, cmdBuf: cmdBuf2)
|
|
try scaleBufferOptimized(ctxBuf, scale: perLayerInputScaleVal,
|
|
count: totalPerLayer, cmdBuf: cmdBuf2)
|
|
|
|
let blit = cmdBuf2.makeBlitCommandEncoder()!
|
|
blit.copy(from: ctxBuf, sourceOffset: 0,
|
|
to: plBuf, destinationOffset: 0,
|
|
size: totalPerLayer * 4)
|
|
blit.endEncoding()
|
|
}
|
|
}
|
|
|
|
// ── Phase 2: Layers (all in same command buffer) ──
|
|
for layerIdx in 0..<numHiddenLayers {
|
|
let isOwner = layerIdx < firstKVShared
|
|
let cacheIdx = isOwner ? layerIdx : (kvSourceMap[layerIdx] ?? (layerIdx - numKvShared))
|
|
let cache = kvCaches[cacheIdx]
|
|
|
|
let plOffset = perLayerInputSize > 0 ? layerIdx * perLayerInputSize * MemoryLayout<Float>.stride : 0
|
|
|
|
// OPTIMIZED: Use shared command buffer (no wait per layer)
|
|
try layers[layerIdx].forwardOptimized(
|
|
input: h, position: position,
|
|
kvCache: cache, shouldStoreKV: isOwner,
|
|
temps: temps, engine: engine,
|
|
cmdBuf: cmdBuf2,
|
|
perLayerInput: perLayerEmbedBuffer,
|
|
perLayerInputOffset: plOffset
|
|
)
|
|
}
|
|
|
|
let cmdBuf3 = engine.commandQueue.makeCommandBuffer()!
|
|
|
|
// ── Phase 3: LM Head (in same command buffer) ──
|
|
var lmInput = h
|
|
if let fn = finalNorm {
|
|
try rmsNormOptimized(input: h, weight: fn, output: temps.ns,
|
|
count: hiddenSize, cmdBuf: cmdBuf3)
|
|
lmInput = temps.ns
|
|
}
|
|
|
|
try quantizedMatmulOptimized(input: lmInput, weights: embedWeight,
|
|
output: logitsBuffer, cmdBuf: cmdBuf3)
|
|
|
|
// Logit softcapping
|
|
if let cap = finalLogitSoftcapping {
|
|
try applyLogitSoftcappingOptimized(buffer: logitsBuffer, cap: cap,
|
|
count: vocabSize, cmdBuf: cmdBuf3)
|
|
}
|
|
|
|
// ── Final: Commit and wait ONCE ──
|
|
cmdBuf3.commit()
|
|
cmdBuf3.waitUntilCompleted() // Only ONE wait for entire forward pass!
|
|
|
|
// Read back logits
|
|
let logits = engine.readFloats(from: logitsBuffer, count: vocabSize)
|
|
|
|
if debug && position < 3 {
|
|
let maxLogit = logits.max() ?? 0
|
|
let minLogit = logits.min() ?? 0
|
|
print(" Optimized forward: max=\(maxLogit), min=\(minLogit)")
|
|
}
|
|
|
|
return logits
|
|
}
|
|
|
|
// ── Optimized helper functions (accept cmdBuf parameter) ──
|
|
|
|
// FUSED: dequantize + scale in one kernel
|
|
func dequantizeScaleFused(weight: QuantizedWeights, tokenId: Int,
|
|
output: MTLBuffer, scale: Float,
|
|
nCols: Int? = nil,
|
|
cmdBuf: MTLCommandBuffer) throws {
|
|
let pso = try engine.pipeline(named: "fused_dequantize_scale")
|
|
let enc = cmdBuf.makeComputeCommandEncoder()!
|
|
enc.setComputePipelineState(pso)
|
|
enc.setBuffer(weight.weight, offset: 0, index: 0)
|
|
enc.setBuffer(weight.scales, offset: 0, index: 1)
|
|
enc.setBuffer(weight.biases, offset: 0, index: 2)
|
|
enc.setBuffer(output, offset: 0, index: 3)
|
|
let actualCols = nCols ?? hiddenSize
|
|
var nColsVal = UInt32(actualCols)
|
|
enc.setBytes(&nColsVal, length: 4, index: 4)
|
|
var row = Int32(tokenId)
|
|
enc.setBytes(&row, length: 4, index: 5)
|
|
var groupSize = UInt32(weight.groupSize)
|
|
enc.setBytes(&groupSize, length: 4, index: 6)
|
|
var s = scale
|
|
enc.setBytes(&s, length: 4, index: 7)
|
|
let tg = engine.threadgroupSize1D(pso, count: actualCols)
|
|
enc.dispatchThreads(MTLSize(width: actualCols, height: 1, depth: 1),
|
|
threadsPerThreadgroup: tg)
|
|
enc.endEncoding()
|
|
}
|
|
|
|
func dequantizeRowOptimized(weight: QuantizedWeights, tokenId: Int,
|
|
output: MTLBuffer, nCols: Int? = nil,
|
|
cmdBuf: MTLCommandBuffer) throws {
|
|
let pso = try engine.pipeline(named: "dequantize_row")
|
|
let enc = cmdBuf.makeComputeCommandEncoder()!
|
|
enc.setComputePipelineState(pso)
|
|
enc.setBuffer(weight.weight, offset: 0, index: 0)
|
|
enc.setBuffer(weight.scales, offset: 0, index: 1)
|
|
enc.setBuffer(weight.biases, offset: 0, index: 2)
|
|
enc.setBuffer(output, offset: 0, index: 3)
|
|
let actualCols = nCols ?? hiddenSize
|
|
var nColsVal = UInt32(actualCols)
|
|
enc.setBytes(&nColsVal, length: 4, index: 4)
|
|
var row = Int32(tokenId)
|
|
enc.setBytes(&row, length: 4, index: 5)
|
|
var groupSize = UInt32(weight.groupSize)
|
|
enc.setBytes(&groupSize, length: 4, index: 6)
|
|
let tg = engine.threadgroupSize1D(pso, count: actualCols)
|
|
enc.dispatchThreads(MTLSize(width: actualCols, height: 1, depth: 1),
|
|
threadsPerThreadgroup: tg)
|
|
enc.endEncoding()
|
|
// NO waitUntilCompleted here!
|
|
}
|
|
|
|
func scaleBufferOptimized(_ buf: MTLBuffer, scale: Float, count: Int,
|
|
cmdBuf: MTLCommandBuffer) throws {
|
|
let pso = try engine.pipeline(named: "eltwise_scale")
|
|
let enc = cmdBuf.makeComputeCommandEncoder()!
|
|
enc.setComputePipelineState(pso)
|
|
enc.setBuffer(buf, offset: 0, index: 0)
|
|
var s = scale
|
|
enc.setBytes(&s, length: 4, index: 1)
|
|
var N = UInt32(count)
|
|
enc.setBytes(&N, length: 4, index: 2)
|
|
let tg = engine.threadgroupSize1D(pso, count: count)
|
|
enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
|
|
threadsPerThreadgroup: tg)
|
|
enc.endEncoding()
|
|
// NO waitUntilCompleted here!
|
|
}
|
|
|
|
func quantizedMatmulOptimized(input: MTLBuffer, weights: QuantizedWeights,
|
|
output: MTLBuffer, cmdBuf: MTLCommandBuffer) throws {
|
|
let pso = try engine.pipeline(named: "quantized_matmul")
|
|
let enc = cmdBuf.makeComputeCommandEncoder()!
|
|
enc.setComputePipelineState(pso)
|
|
enc.setBuffer(input, offset: 0, index: 0)
|
|
enc.setBuffer(weights.weight, offset: 0, index: 1)
|
|
enc.setBuffer(weights.scales, offset: 0, index: 2)
|
|
enc.setBuffer(weights.biases, offset: 0, index: 3)
|
|
enc.setBuffer(output, offset: 0, index: 4)
|
|
var inDim = UInt32(weights.inDim)
|
|
enc.setBytes(&inDim, length: 4, index: 5)
|
|
var outDim = UInt32(weights.outDim)
|
|
enc.setBytes(&outDim, length: 4, index: 6)
|
|
var groupSize = UInt32(weights.groupSize)
|
|
enc.setBytes(&groupSize, length: 4, index: 7)
|
|
let tg = engine.threadgroupSize1D(pso, count: weights.outDim)
|
|
enc.dispatchThreads(MTLSize(width: weights.outDim, height: 1, depth: 1),
|
|
threadsPerThreadgroup: tg)
|
|
enc.endEncoding()
|
|
// NO waitUntilCompleted here!
|
|
}
|
|
|
|
func rmsNormOptimized(input: MTLBuffer, weight: MTLBuffer, output: MTLBuffer,
|
|
count: Int, cmdBuf: MTLCommandBuffer) throws {
|
|
let pso = try engine.pipeline(named: "rms_norm")
|
|
let enc = cmdBuf.makeComputeCommandEncoder()!
|
|
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 N = UInt32(count)
|
|
enc.setBytes(&N, length: 4, index: 3)
|
|
var eps: Float = rmsNormEps
|
|
enc.setBytes(&eps, length: 4, index: 4)
|
|
let tg = engine.threadgroupSize1D(pso, count: count)
|
|
enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
|
|
threadsPerThreadgroup: tg)
|
|
enc.endEncoding()
|
|
// NO waitUntilCompleted here!
|
|
}
|
|
|
|
func rmsNormBatchOptimized(input: MTLBuffer, weight: MTLBuffer, output: MTLBuffer,
|
|
perLayerSize: Int, numLayers: Int,
|
|
cmdBuf: MTLCommandBuffer) throws {
|
|
// Batch all per-layer norms in one kernel dispatch
|
|
let pso = try engine.pipeline(named: "rms_norm")
|
|
let enc = cmdBuf.makeComputeCommandEncoder()!
|
|
enc.setComputePipelineState(pso)
|
|
|
|
for layerIdx in 0..<numLayers {
|
|
let offset = layerIdx * perLayerSize
|
|
enc.setBuffer(input, offset: offset * 4, index: 0)
|
|
enc.setBuffer(weight, offset: 0, index: 1)
|
|
enc.setBuffer(output, offset: offset * 4, index: 2)
|
|
var N = UInt32(perLayerSize)
|
|
enc.setBytes(&N, length: 4, index: 3)
|
|
var eps: Float = rmsNormEps
|
|
enc.setBytes(&eps, length: 4, index: 4)
|
|
let tg = engine.threadgroupSize1D(pso, count: perLayerSize)
|
|
enc.dispatchThreads(MTLSize(width: perLayerSize, height: 1, depth: 1),
|
|
threadsPerThreadgroup: tg)
|
|
}
|
|
enc.endEncoding()
|
|
// NO waitUntilCompleted here!
|
|
}
|
|
|
|
func matmulBF16Optimized(input: MTLBuffer, weight: MTLBuffer, output: MTLBuffer,
|
|
inDim: Int, outDim: Int, cmdBuf: MTLCommandBuffer) throws {
|
|
let pso = try engine.pipeline(named: "matmul_f32")
|
|
let enc = cmdBuf.makeComputeCommandEncoder()!
|
|
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 M: UInt32 = 1
|
|
enc.setBytes(&M, length: 4, index: 3)
|
|
var K = UInt32(inDim)
|
|
enc.setBytes(&K, length: 4, index: 4)
|
|
var N = UInt32(outDim)
|
|
enc.setBytes(&N, length: 4, index: 5)
|
|
let tg = engine.threadgroupSize1D(pso, count: outDim)
|
|
enc.dispatchThreads(MTLSize(width: outDim, height: 1, depth: 1),
|
|
threadsPerThreadgroup: tg)
|
|
enc.endEncoding()
|
|
// NO waitUntilCompleted here!
|
|
}
|
|
|
|
func eltwiseAddScaledOptimized(a: MTLBuffer, scaleA: Float,
|
|
b: MTLBuffer, scaleB: Float,
|
|
output: MTLBuffer, count: Int,
|
|
cmdBuf: MTLCommandBuffer) throws {
|
|
let pso = try engine.pipeline(named: "eltwise_add_scaled")
|
|
let enc = cmdBuf.makeComputeCommandEncoder()!
|
|
enc.setComputePipelineState(pso)
|
|
enc.setBuffer(a, offset: 0, index: 0)
|
|
var sa = scaleA
|
|
enc.setBytes(&sa, length: 4, index: 1)
|
|
enc.setBuffer(b, offset: 0, index: 2)
|
|
var sb = scaleB
|
|
enc.setBytes(&sb, length: 4, index: 3)
|
|
enc.setBuffer(output, offset: 0, index: 4)
|
|
var N = UInt32(count)
|
|
enc.setBytes(&N, length: 4, index: 5)
|
|
let tg = engine.threadgroupSize1D(pso, count: count)
|
|
enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
|
|
threadsPerThreadgroup: tg)
|
|
enc.endEncoding()
|
|
// NO waitUntilCompleted here!
|
|
}
|
|
|
|
func applyLogitSoftcappingOptimized(buffer: MTLBuffer, cap: Float,
|
|
count: Int, cmdBuf: MTLCommandBuffer) throws {
|
|
let pso = try engine.pipeline(named: "tanh_scale")
|
|
let enc = cmdBuf.makeComputeCommandEncoder()!
|
|
enc.setComputePipelineState(pso)
|
|
enc.setBuffer(buffer, offset: 0, index: 0)
|
|
enc.setBuffer(buffer, offset: 0, index: 1) // in-place
|
|
var c = cap
|
|
enc.setBytes(&c, length: 4, index: 2)
|
|
var N = UInt32(count)
|
|
enc.setBytes(&N, length: 4, index: 3)
|
|
let tg = engine.threadgroupSize1D(pso, count: count)
|
|
enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
|
|
threadsPerThreadgroup: tg)
|
|
enc.endEncoding()
|
|
// NO waitUntilCompleted here!
|
|
}
|
|
} |