MarkBase Admin
5a94501f95
- Add FloatWeights fields to E4BLayer (qProjFloat, kProjFloat, etc.) - Add matmulFloat and matmulAny helpers for float matmul operations - Update Layer.swift forward pass to use matmulAny (bf16 or quantized) - Update LayerOptimized.swift and LayerBatch.swift for bf16 weights - Modify Model.swift to load bf16 layer weights via fw() helper - Add guards in LayerBatch.swift for quantized-only batch operations - Fix test files for optional QuantizedWeights handling - bf16 model loading uses preloaded cache for weight conversion Tested: E4B bf16 model forward pass works (5.5 tok/s, no NaN/Inf) Tested: 4-bit models still work correctly after changes
456 lines
18 KiB
Swift
456 lines
18 KiB
Swift
import Metal
|
|
|
|
// ══════════════════════════════════════════════════════════════════
|
|
// Batch Layer Processing - TRUE parallel layer forward pass
|
|
// Process multiple tokens through entire layer simultaneously
|
|
// ══════════════════════════════════════════════════════════════════
|
|
|
|
extension E4BLayer {
|
|
|
|
/// Batch forward pass - process N tokens through entire layer in parallel
|
|
/// Expected: 8-15x speedup for batch inference
|
|
public func forwardBatchTrue(
|
|
batchInput: MTLBuffer, // [batchSize, hiddenSize]
|
|
positions: [Int],
|
|
batchSize: Int,
|
|
kvCache: KVCache,
|
|
shouldStoreKV: Bool,
|
|
temps: ForwardTemps,
|
|
batchTemps: BatchTemps, // Batch-specific buffers
|
|
engine: MarkBaseEngine,
|
|
cmdBuf: MTLCommandBuffer
|
|
) throws {
|
|
// Note: This is a simplified implementation focusing on FFN batch processing
|
|
// Attention still needs sequential KV cache updates
|
|
|
|
// ── Phase 1: Batch Input Norm ──
|
|
guard let inputLN = inputLayernorm else {
|
|
throw NSError(domain: "LayerBatch", code: -1,
|
|
userInfo: [NSLocalizedDescriptionKey: "inputLayernorm required for batch processing"])
|
|
}
|
|
|
|
try batchLayerRMSNorm(
|
|
batchInput: batchInput,
|
|
weights: inputLN,
|
|
batchOutput: batchTemps.hBatch,
|
|
hiddenSize: config.hiddenSize,
|
|
batchSize: batchSize,
|
|
cmdBuf: cmdBuf,
|
|
engine: engine
|
|
)
|
|
|
|
// ── Phase 2: Batch Attention (Sequential for KV cache) ──
|
|
// Note: Attention needs per-token KV cache updates, so we process sequentially
|
|
// But we can batch Q/K/V projections
|
|
|
|
guard let qp = qProj else {
|
|
throw NSError(domain: "LayerBatch", code: -3,
|
|
userInfo: [NSLocalizedDescriptionKey: "Quantized weights required for batch processing"])
|
|
}
|
|
|
|
try batchQuantizedMatmul(
|
|
batchInput: batchTemps.hBatch,
|
|
weights: qp,
|
|
batchOutput: batchTemps.qBatch,
|
|
batchSize: batchSize,
|
|
cmdBuf: cmdBuf,
|
|
engine: engine
|
|
)
|
|
|
|
// Batch grouped norm for Q
|
|
guard let qN = qNorm else {
|
|
throw NSError(domain: "LayerBatch", code: -2,
|
|
userInfo: [NSLocalizedDescriptionKey: "qNorm required for batch processing"])
|
|
}
|
|
|
|
try batchGroupedRMSNorm(
|
|
batchInput: batchTemps.qBatch,
|
|
weights: qN,
|
|
batchOutput: batchTemps.nsBatch,
|
|
count: config.nHeads * config.headDim,
|
|
groupSize: config.headDim,
|
|
batchSize: batchSize,
|
|
cmdBuf: cmdBuf,
|
|
engine: engine
|
|
)
|
|
|
|
// Sequential RoPE and attention (KV cache dependency)
|
|
for i in 0..<batchSize {
|
|
let pos = positions[i]
|
|
let offset = i * config.nHeads * config.headDim
|
|
|
|
// Get Q for this token
|
|
let qToken = engine.device.makeBuffer(
|
|
bytes: batchTemps.nsBatch.contents() + offset * 4,
|
|
length: config.nHeads * config.headDim * 4,
|
|
options: .storageModeShared
|
|
)!
|
|
|
|
// Apply RoPE
|
|
try applyRoPEQ(engine: engine, cmdBuf: cmdBuf, q: qToken, position: pos)
|
|
|
|
// K/V projections (batched, but we need per-token results)
|
|
let hToken = engine.device.makeBuffer(
|
|
bytes: batchTemps.hBatch.contents() + i * config.hiddenSize * 4,
|
|
length: config.hiddenSize * 4,
|
|
options: .storageModeShared
|
|
)!
|
|
|
|
try matmulAny(engine: engine, cmdBuf: cmdBuf, input: hToken, weightsQ: kProj, weightsF: kProjFloat, output: temps.k)
|
|
if let vp = vProj, let vpF = vProjFloat {
|
|
if vp != nil || vpF != nil {
|
|
try matmulAny(engine: engine, cmdBuf: cmdBuf, input: hToken, weightsQ: vp, weightsF: vpF, output: temps.v)
|
|
}
|
|
}
|
|
|
|
// K/V norms
|
|
try groupedRmsNorm(engine: engine, cmdBuf: cmdBuf, input: temps.k, weight: kNorm, output: temps.up,
|
|
count: config.nKvHeads * config.headDim, groupSize: config.headDim, eps: rmsNormEps)
|
|
|
|
// RoPE K
|
|
try applyRoPEK(engine: engine, cmdBuf: cmdBuf, k: temps.up, position: pos)
|
|
|
|
// Store KV
|
|
if shouldStoreKV {
|
|
let valueBuf = vNorm != nil ? temps.gate : temps.v
|
|
kvCache.store(key: temps.up, keySrcOffset: 0, value: valueBuf, valueSrcOffset: 0,
|
|
position: pos, commandBuffer: cmdBuf)
|
|
}
|
|
|
|
// Attention
|
|
let curK = temps.up
|
|
let curV = vNorm != nil ? temps.gate : temps.v
|
|
if config.isSliding {
|
|
if shouldStoreKV {
|
|
try slidingAttention(engine: engine, cmdBuf: cmdBuf, q: qToken, cache: kvCache, position: pos)
|
|
} else {
|
|
try slidingAttentionWithCurrent(engine: engine, cmdBuf: cmdBuf, q: qToken, cache: kvCache,
|
|
curK: curK, curV: curV, position: pos)
|
|
}
|
|
} else {
|
|
if shouldStoreKV {
|
|
try fullAttention(engine: engine, cmdBuf: cmdBuf, q: qToken, cache: kvCache, position: pos)
|
|
} else {
|
|
try fullAttentionWithCurrent(engine: engine, cmdBuf: cmdBuf, q: qToken, cache: kvCache,
|
|
curK: curK, curV: curV, position: pos)
|
|
}
|
|
}
|
|
|
|
// O projection (write back to batch buffer)
|
|
try matmulAny(engine: engine, cmdBuf: cmdBuf, input: temps.attn, weightsQ: oProj, weightsF: oProjFloat, output: temps.h)
|
|
|
|
// Copy to batch position
|
|
let batchOffset = i * config.hiddenSize * 4
|
|
memcpy(batchInput.contents() + batchOffset, temps.h.contents(), config.hiddenSize * 4)
|
|
}
|
|
|
|
// ── Phase 3: Batch FFN (TRUE batch processing) ──
|
|
// This is where we get the big speedup
|
|
|
|
// Post-attention norm (batched)
|
|
guard let postAttnLN = postAttentionLayernorm else {
|
|
throw NSError(domain: "LayerBatch", code: -3,
|
|
userInfo: [NSLocalizedDescriptionKey: "postAttentionLayernorm required"])
|
|
}
|
|
|
|
try batchLayerRMSNorm(
|
|
batchInput: batchInput,
|
|
weights: postAttnLN,
|
|
batchOutput: batchTemps.hBatch,
|
|
hiddenSize: config.hiddenSize,
|
|
batchSize: batchSize,
|
|
cmdBuf: cmdBuf,
|
|
engine: engine
|
|
)
|
|
|
|
// Pre-FFN norm (batched)
|
|
guard let preFFNLN = preFeedforwardLayernorm else {
|
|
throw NSError(domain: "LayerBatch", code: -4,
|
|
userInfo: [NSLocalizedDescriptionKey: "preFeedforwardLayernorm required"])
|
|
}
|
|
|
|
try batchLayerRMSNorm(
|
|
batchInput: batchTemps.hBatch,
|
|
weights: preFFNLN,
|
|
batchOutput: batchTemps.nsBatch,
|
|
hiddenSize: config.hiddenSize,
|
|
batchSize: batchSize,
|
|
cmdBuf: cmdBuf,
|
|
engine: engine
|
|
)
|
|
|
|
// Batch FFN: Gate + Up (fused)
|
|
guard let gp = gateProj, let up = upProj else {
|
|
throw NSError(domain: "LayerBatch", code: -4,
|
|
userInfo: [NSLocalizedDescriptionKey: "Quantized weights required for batch FFN"])
|
|
}
|
|
|
|
try batchFusedGateUp(
|
|
batchInput: batchTemps.nsBatch,
|
|
gateWeights: gp,
|
|
upWeights: up,
|
|
batchOutput: batchTemps.interBatch,
|
|
batchSize: batchSize,
|
|
cmdBuf: cmdBuf,
|
|
engine: engine
|
|
)
|
|
|
|
// Batch Down projection
|
|
guard let dp = downProj else {
|
|
throw NSError(domain: "LayerBatch", code: -5,
|
|
userInfo: [NSLocalizedDescriptionKey: "Quantized weights required for batch down projection"])
|
|
}
|
|
|
|
try batchDownProjection(
|
|
batchInter: batchTemps.interBatch,
|
|
downWeights: dp,
|
|
batchOutput: batchTemps.hBatch,
|
|
batchSize: batchSize,
|
|
cmdBuf: cmdBuf,
|
|
engine: engine
|
|
)
|
|
|
|
// Batch residual add
|
|
try batchEltwiseAdd(
|
|
batchA: batchInput,
|
|
batchB: batchTemps.hBatch,
|
|
batchOutput: batchInput,
|
|
size: config.hiddenSize,
|
|
batchSize: batchSize,
|
|
cmdBuf: cmdBuf,
|
|
engine: engine
|
|
)
|
|
|
|
// Layer scalar (if needed)
|
|
if layerScalar != 1.0 {
|
|
try batchScaleBuffer(
|
|
batchBuffer: batchInput,
|
|
scale: layerScalar,
|
|
size: config.hiddenSize,
|
|
batchSize: batchSize,
|
|
cmdBuf: cmdBuf,
|
|
engine: engine
|
|
)
|
|
}
|
|
}
|
|
|
|
// ── Batch Layer Helper Functions ──
|
|
|
|
private func batchLayerRMSNorm(
|
|
batchInput: MTLBuffer,
|
|
weights: MTLBuffer,
|
|
batchOutput: MTLBuffer,
|
|
hiddenSize: Int,
|
|
batchSize: Int,
|
|
cmdBuf: MTLCommandBuffer,
|
|
engine: MarkBaseEngine
|
|
) throws {
|
|
let pso = try engine.pipeline(named: "batch_layer_rms_norm")
|
|
let enc = cmdBuf.makeComputeCommandEncoder()!
|
|
enc.setComputePipelineState(pso)
|
|
|
|
enc.setBuffer(batchInput, offset: 0, index: 0)
|
|
enc.setBuffer(weights, offset: 0, index: 1)
|
|
enc.setBuffer(batchOutput, offset: 0, index: 2)
|
|
|
|
var hs = UInt32(hiddenSize)
|
|
enc.setBytes(&hs, length: 4, index: 3)
|
|
var eps: Float = rmsNormEps
|
|
enc.setBytes(&eps, length: 4, index: 4)
|
|
var batch = UInt32(batchSize)
|
|
enc.setBytes(&batch, length: 4, index: 5)
|
|
|
|
let tg = MTLSize(width: 256, height: 1, depth: 1)
|
|
let grid = MTLSize(width: batchSize, height: hiddenSize, depth: 1)
|
|
enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
|
|
enc.endEncoding()
|
|
}
|
|
|
|
private func batchQuantizedMatmul(
|
|
batchInput: MTLBuffer,
|
|
weights: QuantizedWeights,
|
|
batchOutput: MTLBuffer,
|
|
batchSize: Int,
|
|
cmdBuf: MTLCommandBuffer,
|
|
engine: MarkBaseEngine
|
|
) throws {
|
|
let pso = try engine.pipeline(named: "batch_layer_quantized_matmul")
|
|
let enc = cmdBuf.makeComputeCommandEncoder()!
|
|
enc.setComputePipelineState(pso)
|
|
|
|
enc.setBuffer(batchInput, 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(batchOutput, 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)
|
|
var batch = UInt32(batchSize)
|
|
enc.setBytes(&batch, length: 4, index: 8)
|
|
|
|
let tg = MTLSize(width: 256, height: 1, depth: 1)
|
|
let grid = MTLSize(width: batchSize, height: Int(weights.outDim), depth: 1)
|
|
enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
|
|
enc.endEncoding()
|
|
}
|
|
|
|
private func batchGroupedRMSNorm(
|
|
batchInput: MTLBuffer,
|
|
weights: MTLBuffer,
|
|
batchOutput: MTLBuffer,
|
|
count: Int,
|
|
groupSize: Int,
|
|
batchSize: Int,
|
|
cmdBuf: MTLCommandBuffer,
|
|
engine: MarkBaseEngine
|
|
) throws {
|
|
// Use existing grouped_rms_norm kernel with batch iteration
|
|
// For now, process sequentially (can optimize later)
|
|
let inputPtr = batchInput.contents().assumingMemoryBound(to: Float.self)
|
|
let outputPtr = batchOutput.contents().assumingMemoryBound(to: Float.self)
|
|
|
|
for i in 0..<batchSize {
|
|
let offset = i * count
|
|
let tokenInput = engine.device.makeBuffer(
|
|
bytes: inputPtr + offset,
|
|
length: count * 4,
|
|
options: .storageModeShared
|
|
)!
|
|
let tokenOutput = engine.device.makeBuffer(
|
|
bytes: outputPtr + offset,
|
|
length: count * 4,
|
|
options: .storageModeShared
|
|
)!
|
|
|
|
try groupedRmsNorm(engine: engine, cmdBuf: cmdBuf, input: tokenInput, weight: weights,
|
|
output: tokenOutput, count: count, groupSize: groupSize, eps: rmsNormEps)
|
|
}
|
|
}
|
|
|
|
private func batchFusedGateUp(
|
|
batchInput: MTLBuffer,
|
|
gateWeights: QuantizedWeights,
|
|
upWeights: QuantizedWeights,
|
|
batchOutput: MTLBuffer,
|
|
batchSize: Int,
|
|
cmdBuf: MTLCommandBuffer,
|
|
engine: MarkBaseEngine
|
|
) throws {
|
|
let pso = try engine.pipeline(named: "batch_fused_gate_up")
|
|
let enc = cmdBuf.makeComputeCommandEncoder()!
|
|
enc.setComputePipelineState(pso)
|
|
|
|
enc.setBuffer(batchInput, offset: 0, index: 0)
|
|
enc.setBuffer(gateWeights.weight, offset: 0, index: 1)
|
|
enc.setBuffer(gateWeights.scales, offset: 0, index: 2)
|
|
enc.setBuffer(gateWeights.biases, offset: 0, index: 3)
|
|
enc.setBuffer(upWeights.weight, offset: 0, index: 4)
|
|
enc.setBuffer(upWeights.scales, offset: 0, index: 5)
|
|
enc.setBuffer(upWeights.biases, offset: 0, index: 6)
|
|
enc.setBuffer(batchOutput, offset: 0, index: 7)
|
|
|
|
var hiddenSize = UInt32(gateWeights.inDim)
|
|
enc.setBytes(&hiddenSize, length: 4, index: 8)
|
|
var intermediateSize = UInt32(gateWeights.outDim)
|
|
enc.setBytes(&intermediateSize, length: 4, index: 9)
|
|
var groupSize = UInt32(gateWeights.groupSize)
|
|
enc.setBytes(&groupSize, length: 4, index: 10)
|
|
var batch = UInt32(batchSize)
|
|
enc.setBytes(&batch, length: 4, index: 11)
|
|
|
|
let tg = MTLSize(width: 256, height: 1, depth: 1)
|
|
let grid = MTLSize(width: batchSize, height: Int(gateWeights.outDim), depth: 1)
|
|
enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
|
|
enc.endEncoding()
|
|
}
|
|
|
|
private func batchDownProjection(
|
|
batchInter: MTLBuffer,
|
|
downWeights: QuantizedWeights,
|
|
batchOutput: MTLBuffer,
|
|
batchSize: Int,
|
|
cmdBuf: MTLCommandBuffer,
|
|
engine: MarkBaseEngine
|
|
) throws {
|
|
let pso = try engine.pipeline(named: "batch_down_projection")
|
|
let enc = cmdBuf.makeComputeCommandEncoder()!
|
|
enc.setComputePipelineState(pso)
|
|
|
|
enc.setBuffer(batchInter, offset: 0, index: 0)
|
|
enc.setBuffer(downWeights.weight, offset: 0, index: 1)
|
|
enc.setBuffer(downWeights.scales, offset: 0, index: 2)
|
|
enc.setBuffer(downWeights.biases, offset: 0, index: 3)
|
|
enc.setBuffer(batchOutput, offset: 0, index: 4)
|
|
|
|
var hiddenSize = UInt32(downWeights.outDim)
|
|
enc.setBytes(&hiddenSize, length: 4, index: 5)
|
|
var intermediateSize = UInt32(downWeights.inDim)
|
|
enc.setBytes(&intermediateSize, length: 4, index: 6)
|
|
var groupSize = UInt32(downWeights.groupSize)
|
|
enc.setBytes(&groupSize, length: 4, index: 7)
|
|
var batch = UInt32(batchSize)
|
|
enc.setBytes(&batch, length: 4, index: 8)
|
|
|
|
let tg = MTLSize(width: 256, height: 1, depth: 1)
|
|
let grid = MTLSize(width: batchSize, height: Int(downWeights.outDim), depth: 1)
|
|
enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
|
|
enc.endEncoding()
|
|
}
|
|
|
|
private func batchEltwiseAdd(
|
|
batchA: MTLBuffer,
|
|
batchB: MTLBuffer,
|
|
batchOutput: MTLBuffer,
|
|
size: Int,
|
|
batchSize: Int,
|
|
cmdBuf: MTLCommandBuffer,
|
|
engine: MarkBaseEngine
|
|
) throws {
|
|
let pso = try engine.pipeline(named: "batch_eltwise_add")
|
|
let enc = cmdBuf.makeComputeCommandEncoder()!
|
|
enc.setComputePipelineState(pso)
|
|
|
|
enc.setBuffer(batchA, offset: 0, index: 0)
|
|
enc.setBuffer(batchB, offset: 0, index: 1)
|
|
enc.setBuffer(batchOutput, offset: 0, index: 2)
|
|
|
|
var s = UInt32(size)
|
|
enc.setBytes(&s, length: 4, index: 3)
|
|
var batch = UInt32(batchSize)
|
|
enc.setBytes(&batch, length: 4, index: 4)
|
|
|
|
let tg = MTLSize(width: 256, height: 1, depth: 1)
|
|
let grid = MTLSize(width: batchSize * size, height: 1, depth: 1)
|
|
enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
|
|
enc.endEncoding()
|
|
}
|
|
|
|
private func batchScaleBuffer(
|
|
batchBuffer: MTLBuffer,
|
|
scale: Float,
|
|
size: Int,
|
|
batchSize: Int,
|
|
cmdBuf: MTLCommandBuffer,
|
|
engine: MarkBaseEngine
|
|
) throws {
|
|
let pso = try engine.pipeline(named: "eltwise_scale")
|
|
let enc = cmdBuf.makeComputeCommandEncoder()!
|
|
enc.setComputePipelineState(pso)
|
|
|
|
enc.setBuffer(batchBuffer, offset: 0, index: 0)
|
|
var s = scale
|
|
enc.setBytes(&s, length: 4, index: 1)
|
|
var count = UInt32(batchSize * size)
|
|
enc.setBytes(&count, length: 4, index: 2)
|
|
|
|
let tg = MTLSize(width: 256, height: 1, depth: 1)
|
|
let grid = MTLSize(width: batchSize * size, height: 1, depth: 1)
|
|
enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
|
|
enc.endEncoding()
|
|
}
|
|
} |