v2: Initial clean branch with unit tests + CI/CD pipeline
- Started from ac75faa (initial E4B-MarkBase integration)
- Kept Sources/ (all engine code) + Package.swift + .gitignore
- Removed all ad-hoc tests, documentation, scripts, Python files
- Added Tests/00_Unit/ (MathTest, TokenizerTest, SamplerTest)
- Added .gitea/workflows/ci.yaml (build + unit tests + lint)
- Added Scripts/check_resources.sh (memory-aware test runner)
- Added Tests/Manifest.json (resource requirements for all tests)
- Focus: 4-bit quantized models only
This commit is contained in:
MarkBase Admin
@@ -0,0 +1,339 @@
|
||||
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)
|
||||
|
||||
// Logits scaling (if needed)
|
||||
if embedWeight.groupSize == 32 && embedWeight.inDim == hiddenSize {
|
||||
let logitsScale = Float(30.0 / 116.23 / sqrt(Float(hiddenSize)))
|
||||
try scaleBufferOptimized(logitsBuffer, scale: logitsScale, count: vocabSize, 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!
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user