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
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,439 @@
|
||||
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
|
||||
|
||||
try batchQuantizedMatmul(
|
||||
batchInput: batchTemps.hBatch,
|
||||
weights: qProj,
|
||||
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 quantizedMatmul(engine: engine, cmdBuf: cmdBuf, input: hToken, weights: kProj, output: temps.k)
|
||||
if let vp = vProj {
|
||||
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf, input: hToken, weights: vp, 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 quantizedMatmul(engine: engine, cmdBuf: cmdBuf, input: temps.attn, weights: oProj, 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)
|
||||
try batchFusedGateUp(
|
||||
batchInput: batchTemps.nsBatch,
|
||||
gateWeights: gateProj,
|
||||
upWeights: upProj,
|
||||
batchOutput: batchTemps.interBatch,
|
||||
batchSize: batchSize,
|
||||
cmdBuf: cmdBuf,
|
||||
engine: engine
|
||||
)
|
||||
|
||||
// Batch Down projection
|
||||
try batchDownProjection(
|
||||
batchInter: batchTemps.interBatch,
|
||||
downWeights: downProj,
|
||||
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()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,246 @@
|
||||
import Metal
|
||||
|
||||
// Optimized E4BLayer forward pass - accepts shared command buffer
|
||||
// Goal: Eliminate per-layer waitUntilCompleted calls
|
||||
|
||||
extension E4BLayer {
|
||||
|
||||
/// Optimized forward pass - batches operations with shared command buffer
|
||||
/// No waitUntilCompleted at end - caller handles that
|
||||
public func forwardOptimized(input: MTLBuffer, position: Int,
|
||||
kvCache: KVCache,
|
||||
shouldStoreKV: Bool,
|
||||
temps: ForwardTemps,
|
||||
engine: MarkBaseEngine,
|
||||
cmdBuf: MTLCommandBuffer,
|
||||
perLayerInput: MTLBuffer? = nil,
|
||||
perLayerInputOffset: Int = 0) throws {
|
||||
self.attnBuf = temps.attn
|
||||
|
||||
if useMoE {
|
||||
// ── MoE path: GPU mega kernel eliminates CPU dependency ──
|
||||
// All operations use shared command buffer (NO waits)
|
||||
|
||||
// Attention + MoE + post-FFN all use shared command buffer
|
||||
try attentionForwardOptimized(input: input, position: position,
|
||||
kvCache: kvCache, shouldStoreKV: shouldStoreKV,
|
||||
temps: temps, engine: engine, cmdBuf: cmdBuf)
|
||||
|
||||
try moeForwardOptimized(input: input, ns: temps.ns, temps: temps,
|
||||
cmdBuf: cmdBuf, engine: engine)
|
||||
|
||||
try postFfnForwardOptimized(input: input, temps: temps, engine: engine,
|
||||
cmdBuf: cmdBuf,
|
||||
perLayerInput: perLayerInput,
|
||||
perLayerInputOffset: perLayerInputOffset)
|
||||
|
||||
if layerScalar != 1.0 {
|
||||
try eltwiseAddScaled(engine: engine, cmdBuf: cmdBuf,
|
||||
a: input, scaleA: layerScalar,
|
||||
b: input, scaleB: 0,
|
||||
output: input, count: config.hiddenSize)
|
||||
}
|
||||
// NO waitUntilCompleted - mega kernel does ALL work on GPU!
|
||||
} else {
|
||||
// ── Dense path: all operations in shared command buffer (NO wait) ──
|
||||
try attentionForwardOptimized(input: input, position: position,
|
||||
kvCache: kvCache, shouldStoreKV: shouldStoreKV,
|
||||
temps: temps, engine: engine, cmdBuf: cmdBuf)
|
||||
|
||||
// FFN: gate+up fused → down → residual
|
||||
try fusedGateUp(engine: engine, cmdBuf: cmdBuf,
|
||||
input: temps.ns, output: temps.gate)
|
||||
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf,
|
||||
input: temps.gate, weights: downProj, output: temps.h)
|
||||
try eltwiseAdd(engine: engine, cmdBuf: cmdBuf,
|
||||
a: input, b: temps.h,
|
||||
output: input, count: config.hiddenSize)
|
||||
|
||||
try postFfnForwardOptimized(input: input, temps: temps, engine: engine,
|
||||
cmdBuf: cmdBuf,
|
||||
perLayerInput: perLayerInput,
|
||||
perLayerInputOffset: perLayerInputOffset)
|
||||
|
||||
if layerScalar != 1.0 {
|
||||
try eltwiseAddScaled(engine: engine, cmdBuf: cmdBuf,
|
||||
a: input, scaleA: layerScalar,
|
||||
b: input, scaleB: 0,
|
||||
output: input, count: config.hiddenSize)
|
||||
}
|
||||
// NO waitUntilCompleted - caller handles that!
|
||||
}
|
||||
}
|
||||
|
||||
// ── Optimized attention forward (reuses existing functions) ──
|
||||
private func attentionForwardOptimized(input: MTLBuffer, position: Int,
|
||||
kvCache: KVCache,
|
||||
shouldStoreKV: Bool,
|
||||
temps: ForwardTemps,
|
||||
engine: MarkBaseEngine,
|
||||
cmdBuf: MTLCommandBuffer) throws {
|
||||
// Same logic as attentionForward, but using passed cmdBuf
|
||||
// Steps 1-13 from original implementation
|
||||
|
||||
// ── 1. input_layernorm(x) → temps.attnH ──
|
||||
try rmsNorm(engine: engine, cmdBuf: cmdBuf,
|
||||
input: input, weight: inputLayernorm,
|
||||
output: temps.attnH, count: config.hiddenSize, eps: rmsNormEps)
|
||||
|
||||
// ── 2. Q = q_proj(temps.attnH) → temps.q ──
|
||||
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf,
|
||||
input: temps.attnH, weights: qProj, output: temps.q)
|
||||
|
||||
// ── 3. Q = q_norm(Q) → ns (per-head RMSNorm) ──
|
||||
try groupedRmsNorm(engine: engine, cmdBuf: cmdBuf,
|
||||
input: temps.q, weight: qNorm,
|
||||
output: temps.ns,
|
||||
count: config.nHeads * config.headDim,
|
||||
groupSize: config.headDim, eps: rmsNormEps)
|
||||
|
||||
// ── 4. RoPE(Q) on ns ──
|
||||
try applyRoPEQ(engine: engine, cmdBuf: cmdBuf,
|
||||
q: temps.ns, position: position)
|
||||
|
||||
// ── 5. K,V projections ──
|
||||
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf,
|
||||
input: temps.attnH, weights: kProj, output: temps.k)
|
||||
if let vp = vProj {
|
||||
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf,
|
||||
input: temps.attnH, weights: vp, output: temps.v)
|
||||
} else if kEqualsV {
|
||||
let blit = cmdBuf.makeBlitCommandEncoder()!
|
||||
let copyBytes = config.nKvHeads * config.headDim * MemoryLayout<Float>.stride
|
||||
blit.copy(from: temps.k, sourceOffset: 0,
|
||||
to: temps.v, destinationOffset: 0,
|
||||
size: copyBytes)
|
||||
blit.endEncoding()
|
||||
}
|
||||
|
||||
// ── 6. 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)
|
||||
if let vn = vNorm {
|
||||
try groupedRmsNorm(engine: engine, cmdBuf: cmdBuf,
|
||||
input: temps.v, weight: vn,
|
||||
output: temps.gate,
|
||||
count: config.nKvHeads * config.headDim,
|
||||
groupSize: config.headDim, eps: rmsNormEps)
|
||||
}
|
||||
|
||||
// ── 7. RoPE(K) ──
|
||||
try applyRoPEK(engine: engine, cmdBuf: cmdBuf,
|
||||
k: temps.up, position: position)
|
||||
|
||||
// ── 8. Store K,V ──
|
||||
if shouldStoreKV {
|
||||
let valueBuf = vNorm != nil ? temps.gate : temps.v
|
||||
kvCache.store(key: temps.up, keySrcOffset: 0,
|
||||
value: valueBuf, valueSrcOffset: 0,
|
||||
position: position, commandBuffer: cmdBuf)
|
||||
}
|
||||
|
||||
// ── 9. 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: temps.ns, cache: kvCache, position: position)
|
||||
} else {
|
||||
try slidingAttentionWithCurrent(engine: engine, cmdBuf: cmdBuf,
|
||||
q: temps.ns, cache: kvCache,
|
||||
curK: curK, curV: curV,
|
||||
position: position)
|
||||
}
|
||||
} else {
|
||||
if shouldStoreKV {
|
||||
try fullAttention(engine: engine, cmdBuf: cmdBuf,
|
||||
q: temps.ns, cache: kvCache, position: position)
|
||||
} else {
|
||||
try fullAttentionWithCurrent(engine: engine, cmdBuf: cmdBuf,
|
||||
q: temps.ns, cache: kvCache,
|
||||
curK: curK, curV: curV,
|
||||
position: position)
|
||||
}
|
||||
}
|
||||
|
||||
// ── 10. O projection ──
|
||||
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf,
|
||||
input: temps.attn, weights: oProj, output: temps.attnH)
|
||||
|
||||
// ── 11. Residual 1 ──
|
||||
try eltwiseAdd(engine: engine, cmdBuf: cmdBuf,
|
||||
a: input, b: temps.attnH,
|
||||
output: input, count: config.hiddenSize)
|
||||
|
||||
// ── 12. post_attention_layernorm → temps.attnH ──
|
||||
try rmsNorm(engine: engine, cmdBuf: cmdBuf,
|
||||
input: input, weight: postAttentionLayernorm,
|
||||
output: temps.attnH, count: config.hiddenSize, eps: rmsNormEps)
|
||||
|
||||
// ── 13. pre_feedforward_layernorm → ns ──
|
||||
try rmsNorm(engine: engine, cmdBuf: cmdBuf,
|
||||
input: temps.attnH, weight: preFeedforwardLayernorm,
|
||||
output: temps.ns, count: config.hiddenSize, eps: rmsNormEps)
|
||||
}
|
||||
|
||||
// ── Optimized MoE forward ──
|
||||
private func moeForwardOptimized(input: MTLBuffer, ns: MTLBuffer, temps: ForwardTemps,
|
||||
cmdBuf: MTLCommandBuffer, engine: MarkBaseEngine) throws {
|
||||
// Call existing moeForward with shared cmdBuf
|
||||
try moeForward(input: input, ns: ns, temps: temps,
|
||||
cmdBuf: cmdBuf, engine: engine)
|
||||
}
|
||||
|
||||
// ── Optimized post-FFN forward ──
|
||||
private func postFfnForwardOptimized(input: MTLBuffer, temps: ForwardTemps,
|
||||
engine: MarkBaseEngine,
|
||||
cmdBuf: MTLCommandBuffer,
|
||||
perLayerInput: MTLBuffer?,
|
||||
perLayerInputOffset: Int) throws {
|
||||
// Duplicate logic from postFfnForward (it's private in E4BLayer)
|
||||
|
||||
// ── 17. post_feedforward_layernorm → temps.h ──
|
||||
try rmsNorm(engine: engine, cmdBuf: cmdBuf,
|
||||
input: input, weight: postFeedforwardLayernorm,
|
||||
output: temps.h, count: config.hiddenSize, eps: rmsNormEps)
|
||||
|
||||
// ── 18. Per-layer gating (optional) ──
|
||||
if let pg = perLayerGate, let pp = perLayerProjection, let pl = perLayerInput {
|
||||
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf,
|
||||
input: temps.h, weights: pg,
|
||||
output: temps.gating)
|
||||
try gelu(engine: engine, cmdBuf: cmdBuf,
|
||||
input: temps.gating, output: temps.gating, count: 256)
|
||||
|
||||
try eltwiseMul(engine: engine, cmdBuf: cmdBuf,
|
||||
a: temps.gating, aOffset: 0,
|
||||
b: pl, bOffset: perLayerInputOffset,
|
||||
output: temps.gating, outputOffset: 0,
|
||||
count: 256)
|
||||
|
||||
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf,
|
||||
input: temps.gating, weights: pp,
|
||||
output: temps.h)
|
||||
|
||||
if let ppn = postPerLayerInputNorm {
|
||||
try rmsNorm(engine: engine, cmdBuf: cmdBuf,
|
||||
input: temps.h, weight: ppn,
|
||||
output: temps.h, count: config.hiddenSize, eps: rmsNormEps)
|
||||
}
|
||||
|
||||
try eltwiseAddScaled(engine: engine, cmdBuf: cmdBuf,
|
||||
a: temps.h, scaleA: 1.0,
|
||||
b: temps.h, scaleB: 0.0,
|
||||
output: input, count: config.hiddenSize)
|
||||
} else {
|
||||
try eltwiseAddScaled(engine: engine, cmdBuf: cmdBuf,
|
||||
a: temps.h, scaleA: 1.0,
|
||||
b: temps.h, scaleB: 0.0,
|
||||
output: input, count: config.hiddenSize)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user