import Metal // ── Quantized Weights ──────────────────────────── public struct QuantizedWeights { public let weight: MTLBuffer // U32 packed [outDim, inDim/(32/bits)] public let scales: MTLBuffer // Float32 [outDim, inDim/groupSize] public let biases: MTLBuffer // Float32 [outDim, inDim/groupSize] public let inDim: Int public let outDim: Int public let bits: Int // 4 or 8 public let groupSize: Int // Quantization group size (32, 64, etc.) } // ── Layer Configuration ────────────────────────── public struct E4BLayerConfig { public let isSliding: Bool public let headDim: Int public let intermediateSize: Int public let rotatedDim: Int public let ropeTheta: Float public let ropeScale: Float public let windowSize: Int public let nHeads: Int public let nKvHeads: Int public let hiddenSize: Int public static func sliding(hiddenSize: Int, headDim: Int, intermediateSize: Int, nHeads: Int, nKvHeads: Int, windowSize: Int = 512) -> E4BLayerConfig { E4BLayerConfig( isSliding: true, headDim: headDim, intermediateSize: intermediateSize, rotatedDim: headDim / 2, // default RoPE: half of dimensions rotated ropeTheta: 10000.0, ropeScale: 1.0, windowSize: windowSize, nHeads: nHeads, nKvHeads: nKvHeads, hiddenSize: hiddenSize ) } public static func full(hiddenSize: Int, headDim: Int, intermediateSize: Int, nHeads: Int, nKvHeads: Int, maxPosition: Int = 8192) -> E4BLayerConfig { E4BLayerConfig( isSliding: false, headDim: headDim, intermediateSize: intermediateSize, rotatedDim: Int(Float(headDim) * 0.25), ropeTheta: 1000000.0, ropeScale: 1.0, windowSize: maxPosition, nHeads: nHeads, nKvHeads: nKvHeads, hiddenSize: hiddenSize ) } } // ── Temp buffers shared across forward pass ────── public struct ForwardTemps { public let q: MTLBuffer // [maxHeads * maxHeadDim] public let k: MTLBuffer // [maxKvHeads * maxHeadDim] public let v: MTLBuffer // [maxKvHeads * maxHeadDim] public let h: MTLBuffer // [hiddenSize] scratch (FFN专用) public let attnH: MTLBuffer // [hiddenSize] attention专用 (避免覆盖h) public let gate: MTLBuffer // [maxIntermediateSize] public let up: MTLBuffer // [maxIntermediateSize] public let attn: MTLBuffer // [maxHeads * maxHeadDim] public let gating: MTLBuffer // [256] per-layer gating scratch public let ns: MTLBuffer // [max(hiddenSize, nHeads*headDim)] norm scratch public let io: MTLBuffer // [hiddenSize] dedicated layer I/O (separate from scratch) public init(device: MTLDevice, maxHeadDim: Int = 512, maxIntermediate: Int = 20480, hiddenSize: Int = 2560, nHeads: Int = 8, nKvHeads: Int = 2) throws { func buf(_ n: Int) throws -> MTLBuffer { guard let b = device.makeBuffer(length: n * MemoryLayout.stride, options: .storageModeShared) else { throw E4BError.bufferCreationFailed } return b } q = try buf(nHeads * maxHeadDim) k = try buf(nKvHeads * maxHeadDim) v = try buf(nKvHeads * maxHeadDim) h = try buf(hiddenSize) attnH = try buf(hiddenSize) // NEW: attention专用buffer gate = try buf(maxIntermediate) up = try buf(maxIntermediate) attn = try buf(nHeads * maxHeadDim) gating = try buf(256) ns = try buf(max(hiddenSize, nHeads * maxHeadDim)) io = try buf(hiddenSize) } } // ── MoE structures ────────────────────────────── public struct MoEExpert { public let gateProj: QuantizedWeights public let upProj: QuantizedWeights public let downProj: QuantizedWeights public init(gateProj: QuantizedWeights, upProj: QuantizedWeights, downProj: QuantizedWeights) { self.gateProj = gateProj self.upProj = upProj self.downProj = downProj } } /// Expert weights stored as contiguous 3D tensors [numExperts, outDim, inDimPacked] /// and [numExperts, outDim, numGroups] for scales/biases. /// Per-expert access uses byte offsets into the shared buffers. public struct MoEExpertGroup { /// Full 3D weight buffer [numExperts * expertOutDim, expertInDimPacked] uint32 public let weight: MTLBuffer /// Full 3D scales buffer [numExperts * expertOutDim, numGroups] float32 public let scales: MTLBuffer /// Full 3D biases buffer (same shape as scales) public let biases: MTLBuffer /// Per-expert output dimension public let expertOutDim: Int /// Input dimension (hidden size) public let expertInDim: Int /// Number of groups per output row = inDim / groupSize public let numGroups: Int /// Total number of experts public let numExperts: Int /// Quantization bits (4 or 8) public let bits: Int /// Byte stride per expert for weight buffer public var weightStride: Int { expertOutDim * (expertInDim * bits / 32) * 4 } /// Byte stride per expert for scales/biases buffer public var scalesStride: Int { expertOutDim * numGroups * 4 } public init(weight: MTLBuffer, scales: MTLBuffer, biases: MTLBuffer, expertOutDim: Int, expertInDim: Int, numGroups: Int, numExperts: Int, bits: Int = 4) { self.weight = weight self.scales = scales self.biases = biases self.expertOutDim = expertOutDim self.expertInDim = expertInDim self.numGroups = numGroups self.numExperts = numExperts self.bits = bits } } // ── Layer forward pass ─────────────────────────── public final class E4BLayer { let config: E4BLayerConfig let rmsNormEps: Float = 1e-6 let layerIdx: Int // For debug logging // Norm weights (Float32) let inputLayernorm: MTLBuffer? let postAttentionLayernorm: MTLBuffer? let preFeedforwardLayernorm: MTLBuffer? let postFeedforwardLayernorm: MTLBuffer? let postPerLayerInputNorm: MTLBuffer? // after per-layer gating let qNorm: MTLBuffer? let kNorm: MTLBuffer? let vNorm: MTLBuffer? // nil — no-scale variant // Quantized projections let qProj: QuantizedWeights let kProj: QuantizedWeights let vProj: QuantizedWeights? let oProj: QuantizedWeights let gateProj: QuantizedWeights let upProj: QuantizedWeights let downProj: QuantizedWeights let perLayerGate: QuantizedWeights? let perLayerProjection: QuantizedWeights? // MoE let useMoE: Bool let routerProj: QuantizedWeights? let routerScale: Float let perExpertScale: [Float]? let expertGate: MoEExpertGroup? let expertUp: MoEExpertGroup? let expertDown: MoEExpertGroup? let topK: Int // K=V sharing for full attention layers (Gemma 4) let kEqualsV: Bool // Per-layer constants let perLayerInput: MTLBuffer? let perLayerInputScale: Float let perLayerProjectionScale: Float let layerScalar: Float public init(config: E4BLayerConfig, layerIdx: Int = 0, inputLayernorm: MTLBuffer?, postAttentionLayernorm: MTLBuffer?, preFeedforwardLayernorm: MTLBuffer?, postFeedforwardLayernorm: MTLBuffer?, postPerLayerInputNorm: MTLBuffer?, qNorm: MTLBuffer?, kNorm: MTLBuffer?, vNorm: MTLBuffer?, qProj: QuantizedWeights, kProj: QuantizedWeights, vProj: QuantizedWeights?, oProj: QuantizedWeights, gateProj: QuantizedWeights, upProj: QuantizedWeights, downProj: QuantizedWeights, perLayerGate: QuantizedWeights?, perLayerProjection: QuantizedWeights?, perLayerInput: MTLBuffer?, perLayerInputScale: Float, perLayerProjectionScale: Float, layerScalar: Float, useMoE: Bool = false, routerProj: QuantizedWeights? = nil, routerScale: Float = 1.0, perExpertScale: [Float]? = nil, expertGate: MoEExpertGroup? = nil, expertUp: MoEExpertGroup? = nil, expertDown: MoEExpertGroup? = nil, topK: Int = 8, kEqualsV: Bool = false) { self.layerIdx = layerIdx self.config = config self.inputLayernorm = inputLayernorm self.postAttentionLayernorm = postAttentionLayernorm self.preFeedforwardLayernorm = preFeedforwardLayernorm self.postFeedforwardLayernorm = postFeedforwardLayernorm self.postPerLayerInputNorm = postPerLayerInputNorm self.qNorm = qNorm self.kNorm = kNorm self.vNorm = vNorm self.qProj = qProj self.kProj = kProj self.vProj = vProj self.oProj = oProj self.gateProj = gateProj self.upProj = upProj self.downProj = downProj self.perLayerGate = perLayerGate self.perLayerProjection = perLayerProjection self.kEqualsV = kEqualsV self.perLayerInput = perLayerInput self.perLayerInputScale = perLayerInputScale self.perLayerProjectionScale = perLayerProjectionScale self.layerScalar = layerScalar self.useMoE = useMoE self.routerProj = routerProj self.routerScale = routerScale self.perExpertScale = perExpertScale self.expertGate = expertGate self.expertUp = expertUp self.expertDown = expertDown self.topK = topK } // ── Kernel dispatch helpers (optimized versions preferred) ────────────────── func rmsNorm(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer, input: MTLBuffer, weight: MTLBuffer?, output: MTLBuffer, count: Int, eps: Float) 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: MemoryLayout.size, index: 3) var e = eps enc.setBytes(&e, length: MemoryLayout.size, index: 4) let gridSize = MTLSize(width: count, height: 1, depth: 1) let tgSize = min(256, count) let tg = MTLSize(width: tgSize, height: 1, depth: 1) enc.dispatchThreads(gridSize, threadsPerThreadgroup: tg) enc.endEncoding() } func groupedRmsNorm(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer, input: MTLBuffer, weight: MTLBuffer?, output: MTLBuffer, count: Int, groupSize: Int, eps: Float) throws { let pso = try engine.pipeline(named: "rms_norm_grouped") 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: MemoryLayout.size, index: 3) var gs = UInt32(groupSize) enc.setBytes(&gs, length: MemoryLayout.size, index: 4) var e = eps enc.setBytes(&e, length: MemoryLayout.size, index: 5) let tg = engine.threadgroupSize1D(pso, count: count) enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: tg) enc.endEncoding() } func gelu(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer, input: MTLBuffer, output: MTLBuffer, count: Int) throws { let pso = try engine.pipeline(named: "gelu_approx") let enc = cmdBuf.makeComputeCommandEncoder()! enc.setComputePipelineState(pso) enc.setBuffer(input, offset: 0, index: 0) enc.setBuffer(output, offset: 0, index: 1) var N = UInt32(count) enc.setBytes(&N, length: MemoryLayout.size, index: 2) let tg = engine.threadgroupSize1D(pso, count: count) enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: tg) enc.endEncoding() } func quantizedMatmul(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer, input: MTLBuffer, weights: QuantizedWeights, output: MTLBuffer) throws { // Select kernel based on quantization bits let kernelName = weights.bits == 8 ? "quantized_matmul_8bit" : "quantized_matmul" // TEMPORARILY USE FALLBACK KERNEL FOR TESTING if false, let pso = try? engine.pipeline(named: kernelName) { 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: MemoryLayout.size, index: 5) var outDim = UInt32(weights.outDim) enc.setBytes(&outDim, length: MemoryLayout.size, index: 6) var groupSize = UInt32(weights.groupSize) // quantization group size from weights enc.setBytes(&groupSize, length: MemoryLayout.size, index: 7) // Threadgroup memory for input vector cache let tgMemSize = weights.inDim * 4 // Float32 enc.setThreadgroupMemoryLength(tgMemSize, index: 0) let count = weights.outDim let tg = MTLSize(width: 256, height: 1, depth: 1) enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: tg) enc.endEncoding() return } // Fallback to original 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: MemoryLayout.size, index: 5) var outDim = UInt32(weights.outDim) enc.setBytes(&outDim, length: MemoryLayout.size, index: 6) var groupSize = UInt32(weights.groupSize) // FIX: Add groupSize! enc.setBytes(&groupSize, length: MemoryLayout.size, index: 7) let count = weights.outDim let tg = engine.threadgroupSize1D(pso, count: count) enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: tg) enc.endEncoding() } func applyRoPEQ(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer, q: MTLBuffer, position: Int) throws { let pso = try engine.pipeline(named: "apply_rope_q") let enc = cmdBuf.makeComputeCommandEncoder()! enc.setComputePipelineState(pso) enc.setBuffer(q, offset: 0, index: 0) var nHeads = UInt32(config.nHeads) enc.setBytes(&nHeads, length: MemoryLayout.size, index: 1) var headDim = UInt32(config.headDim) enc.setBytes(&headDim, length: MemoryLayout.size, index: 2) var rotatedDim = UInt32(config.rotatedDim) enc.setBytes(&rotatedDim, length: MemoryLayout.size, index: 3) var theta = config.ropeTheta enc.setBytes(&theta, length: MemoryLayout.size, index: 4) var scale = config.ropeScale enc.setBytes(&scale, length: MemoryLayout.size, index: 5) var pos = Int32(position) enc.setBytes(&pos, length: MemoryLayout.size, index: 6) let count = config.nHeads * (config.rotatedDim / 2) let tg = engine.threadgroupSize1D(pso, count: count) enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: tg) enc.endEncoding() } func applyRoPEK(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer, k: MTLBuffer, position: Int) throws { let pso = try engine.pipeline(named: "apply_rope_k") let enc = cmdBuf.makeComputeCommandEncoder()! enc.setComputePipelineState(pso) enc.setBuffer(k, offset: 0, index: 0) var nKvHeads = UInt32(config.nKvHeads) enc.setBytes(&nKvHeads, length: MemoryLayout.size, index: 1) var headDim = UInt32(config.headDim) enc.setBytes(&headDim, length: MemoryLayout.size, index: 2) var rotatedDim = UInt32(config.rotatedDim) enc.setBytes(&rotatedDim, length: MemoryLayout.size, index: 3) var theta = config.ropeTheta enc.setBytes(&theta, length: MemoryLayout.size, index: 4) var scale = config.ropeScale enc.setBytes(&scale, length: MemoryLayout.size, index: 5) var pos = Int32(position) enc.setBytes(&pos, length: MemoryLayout.size, index: 6) let count = config.nKvHeads * (config.rotatedDim / 2) let tg = engine.threadgroupSize1D(pso, count: count) enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: tg) enc.endEncoding() } func slidingAttention(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer, q: MTLBuffer, cache: KVCache, position: Int) throws { // Try optimized SIMD version first (softcapping removed) if let pso = try? engine.pipeline(named: "sliding_attention_simd") { let enc = cmdBuf.makeComputeCommandEncoder()! enc.setComputePipelineState(pso) enc.setBuffer(q, offset: 0, index: 0) enc.setBuffer(cache.buffer, offset: cache.keyBaseOffset, index: 1) enc.setBuffer(cache.buffer, offset: cache.valueBaseOffset, index: 2) enc.setBuffer(attnBuf, offset: 0, index: 3) var nHeads = UInt32(config.nHeads) enc.setBytes(&nHeads, length: MemoryLayout.size, index: 4) var nKvHeads = UInt32(config.nKvHeads) enc.setBytes(&nKvHeads, length: MemoryLayout.size, index: 5) var headDim = UInt32(config.headDim) enc.setBytes(&headDim, length: MemoryLayout.size, index: 6) var windowSize = UInt32(config.windowSize) enc.setBytes(&windowSize, length: MemoryLayout.size, index: 7) var off = Int32(position) enc.setBytes(&off, length: MemoryLayout.size, index: 8) // Threadgroup memory for K/V cache let kvCacheSize = config.windowSize * config.nKvHeads * (config.headDim/4) * 16 // float4 = 16 bytes enc.setThreadgroupMemoryLength(kvCacheSize, index: 0) enc.setThreadgroupMemoryLength(kvCacheSize, index: 1) let grid = MTLSize(width: config.nHeads, height: config.headDim/4, depth: 1) let tg = MTLSize(width: 8, height: 16, depth: 1) // Tune for cache efficiency enc.dispatchThreads(grid, threadsPerThreadgroup: tg) enc.endEncoding() return } // Fallback to original let pso = try engine.pipeline(named: "sliding_attention") let enc = cmdBuf.makeComputeCommandEncoder()! enc.setComputePipelineState(pso) enc.setBuffer(q, offset: 0, index: 0) enc.setBuffer(cache.buffer, offset: cache.keyBaseOffset, index: 1) enc.setBuffer(cache.buffer, offset: cache.valueBaseOffset, index: 2) enc.setBuffer(attnBuf, offset: 0, index: 3) var nHeads = UInt32(config.nHeads) enc.setBytes(&nHeads, length: MemoryLayout.size, index: 4) var nKvHeads = UInt32(config.nKvHeads) enc.setBytes(&nKvHeads, length: MemoryLayout.size, index: 5) var headDim = UInt32(config.headDim) enc.setBytes(&headDim, length: MemoryLayout.size, index: 6) var windowSize = UInt32(config.windowSize) enc.setBytes(&windowSize, length: MemoryLayout.size, index: 7) var off = Int32(position) enc.setBytes(&off, length: MemoryLayout.size, index: 8) let grid = MTLSize(width: config.nHeads, height: config.headDim, depth: 1) let tg = engine.threadgroupSize2D(pso, grid: (config.nHeads, config.headDim)) enc.dispatchThreads(grid, threadsPerThreadgroup: tg) enc.endEncoding() } func slidingAttentionWithCurrent(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer, q: MTLBuffer, cache: KVCache, curK: MTLBuffer, curV: MTLBuffer, position: Int) throws { let pso = try engine.pipeline(named: "sliding_attention_with_current") let enc = cmdBuf.makeComputeCommandEncoder()! enc.setComputePipelineState(pso) enc.setBuffer(q, offset: 0, index: 0) enc.setBuffer(cache.buffer, offset: cache.keyBaseOffset, index: 1) enc.setBuffer(cache.buffer, offset: cache.valueBaseOffset, index: 2) enc.setBuffer(curK, offset: 0, index: 3) enc.setBuffer(curV, offset: 0, index: 4) enc.setBuffer(attnBuf, offset: 0, index: 5) var nHeads = UInt32(config.nHeads) enc.setBytes(&nHeads, length: MemoryLayout.size, index: 6) var nKvHeads = UInt32(config.nKvHeads) enc.setBytes(&nKvHeads, length: MemoryLayout.size, index: 7) var headDim = UInt32(config.headDim) enc.setBytes(&headDim, length: MemoryLayout.size, index: 8) var windowSize = UInt32(cache.maxLength) enc.setBytes(&windowSize, length: MemoryLayout.size, index: 9) // For shared layers: read ALL entries from owner's cache (positions 0..N) // The owner has already stored at all positions up to current position let cacheLen = cache.currentLength var cacheLenVal = UInt32(cacheLen) enc.setBytes(&cacheLenVal, length: MemoryLayout.size, index: 10) var pos = Int32(position) enc.setBytes(&pos, length: MemoryLayout.size, index: 11) let grid = MTLSize(width: config.nHeads, height: config.headDim, depth: 1) let tg = engine.threadgroupSize2D(pso, grid: (config.nHeads, config.headDim)) enc.dispatchThreads(grid, threadsPerThreadgroup: tg) enc.endEncoding() } // Temp buffer for attention output — set externally from ForwardTemps var attnBuf: MTLBuffer! func fullAttention(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer, q: MTLBuffer, cache: KVCache, position: Int) throws { // Try optimized SIMD version first (no softcapping for text models) if let pso = try? engine.pipeline(named: "full_attention_simd") { let enc = cmdBuf.makeComputeCommandEncoder()! enc.setComputePipelineState(pso) enc.setBuffer(q, offset: 0, index: 0) enc.setBuffer(cache.buffer, offset: cache.keyBaseOffset, index: 1) enc.setBuffer(cache.buffer, offset: cache.valueBaseOffset, index: 2) enc.setBuffer(attnBuf, offset: 0, index: 3) var nHeads = UInt32(config.nHeads) enc.setBytes(&nHeads, length: MemoryLayout.size, index: 4) var nKvHeads = UInt32(config.nKvHeads) enc.setBytes(&nKvHeads, length: MemoryLayout.size, index: 5) var headDim = UInt32(config.headDim) enc.setBytes(&headDim, length: MemoryLayout.size, index: 6) var seqLen = UInt32(position + 1) // FIX: use seqLen, not maxPos enc.setBytes(&seqLen, length: MemoryLayout.size, index: 7) // Threadgroup memory for K/V cache (use seqLen, not maxPos) let kvCacheSize = Int(seqLen) * config.nKvHeads * (config.headDim/4) * 16 // float4 = 16 bytes enc.setThreadgroupMemoryLength(kvCacheSize, index: 0) enc.setThreadgroupMemoryLength(kvCacheSize, index: 1) let grid = MTLSize(width: config.nHeads, height: config.headDim/4, depth: 1) let tg = MTLSize(width: 8, height: 16, depth: 1) // Tune for cache efficiency enc.dispatchThreads(grid, threadsPerThreadgroup: tg) enc.endEncoding() return } // Fallback to original let pso = try engine.pipeline(named: "full_attention") let enc = cmdBuf.makeComputeCommandEncoder()! enc.setComputePipelineState(pso) enc.setBuffer(q, offset: 0, index: 0) enc.setBuffer(cache.buffer, offset: cache.keyBaseOffset, index: 1) enc.setBuffer(cache.buffer, offset: cache.valueBaseOffset, index: 2) enc.setBuffer(attnBuf, offset: 0, index: 3) var nHeads = UInt32(config.nHeads) enc.setBytes(&nHeads, length: MemoryLayout.size, index: 4) var nKvHeads = UInt32(config.nKvHeads) enc.setBytes(&nKvHeads, length: MemoryLayout.size, index: 5) var headDim = UInt32(config.headDim) enc.setBytes(&headDim, length: MemoryLayout.size, index: 6) var maxPos = UInt32(cache.maxLength) enc.setBytes(&maxPos, length: MemoryLayout.size, index: 7) var off = Int32(position) enc.setBytes(&off, length: MemoryLayout.size, index: 8) let grid = MTLSize(width: config.nHeads, height: config.headDim, depth: 1) let tg = engine.threadgroupSize2D(pso, grid: (config.nHeads, config.headDim)) enc.dispatchThreads(grid, threadsPerThreadgroup: tg) enc.endEncoding() } func fullAttentionWithCurrent(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer, q: MTLBuffer, cache: KVCache, curK: MTLBuffer, curV: MTLBuffer, position: Int) throws { let pso = try engine.pipeline(named: "full_attention_with_current") let enc = cmdBuf.makeComputeCommandEncoder()! enc.setComputePipelineState(pso) enc.setBuffer(q, offset: 0, index: 0) enc.setBuffer(cache.buffer, offset: cache.keyBaseOffset, index: 1) enc.setBuffer(cache.buffer, offset: cache.valueBaseOffset, index: 2) enc.setBuffer(curK, offset: 0, index: 3) enc.setBuffer(curV, offset: 0, index: 4) enc.setBuffer(attnBuf, offset: 0, index: 5) var nHeads = UInt32(config.nHeads) enc.setBytes(&nHeads, length: MemoryLayout.size, index: 6) var nKvHeads = UInt32(config.nKvHeads) enc.setBytes(&nKvHeads, length: MemoryLayout.size, index: 7) var headDim = UInt32(config.headDim) enc.setBytes(&headDim, length: MemoryLayout.size, index: 8) // For shared layers: read ALL entries from owner's cache (positions 0..N) // The owner has already stored at all positions up to current position let cacheLen = cache.currentLength var cacheLenVal = UInt32(cacheLen) enc.setBytes(&cacheLenVal, length: MemoryLayout.size, index: 9) var pos = Int32(position) enc.setBytes(&pos, length: MemoryLayout.size, index: 10) let grid = MTLSize(width: config.nHeads, height: config.headDim, depth: 1) let tg = engine.threadgroupSize2D(pso, grid: (config.nHeads, config.headDim)) enc.dispatchThreads(grid, threadsPerThreadgroup: tg) enc.endEncoding() } func eltwiseAdd(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer, a: MTLBuffer, b: MTLBuffer, output: MTLBuffer, count: Int) throws { // Try optimized SIMD version first if let pso = try? engine.pipeline(named: "eltwise_add_simd") { let enc = cmdBuf.makeComputeCommandEncoder()! enc.setComputePipelineState(pso) enc.setBuffer(a, offset: 0, index: 0) enc.setBuffer(b, offset: 0, index: 1) enc.setBuffer(output, offset: 0, index: 2) var N = UInt32(count) enc.setBytes(&N, length: MemoryLayout.size, index: 3) let grid = MTLSize(width: (count + 3) / 4, height: 1, depth: 1) // float4 processing let tg = engine.threadgroupSize1D(pso, count: (count + 3) / 4) enc.dispatchThreads(grid, threadsPerThreadgroup: tg) enc.endEncoding() return } // Fallback to original let pso = try engine.pipeline(named: "eltwise_add") let enc = cmdBuf.makeComputeCommandEncoder()! enc.setComputePipelineState(pso) enc.setBuffer(a, offset: 0, index: 0) enc.setBuffer(b, offset: 0, index: 1) enc.setBuffer(output, offset: 0, index: 2) var N = UInt32(count) enc.setBytes(&N, length: MemoryLayout.size, index: 3) let tg = engine.threadgroupSize1D(pso, count: count) enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: tg) enc.endEncoding() } func eltwiseMul(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer, a: MTLBuffer, aOffset: Int = 0, b: MTLBuffer, bOffset: Int = 0, output: MTLBuffer, outputOffset: Int = 0, count: Int) throws { // Try optimized SIMD version first if let pso = try? engine.pipeline(named: "eltwise_mul_simd") { let enc = cmdBuf.makeComputeCommandEncoder()! enc.setComputePipelineState(pso) enc.setBuffer(a, offset: aOffset, index: 0) enc.setBuffer(b, offset: bOffset, index: 1) enc.setBuffer(output, offset: outputOffset, index: 2) var N = UInt32(count) enc.setBytes(&N, length: MemoryLayout.size, index: 3) let grid = MTLSize(width: (count + 3) / 4, height: 1, depth: 1) // float4 processing let tg = engine.threadgroupSize1D(pso, count: (count + 3) / 4) enc.dispatchThreads(grid, threadsPerThreadgroup: tg) enc.endEncoding() return } // Fallback to original let pso = try engine.pipeline(named: "eltwise_mul") let enc = cmdBuf.makeComputeCommandEncoder()! enc.setComputePipelineState(pso) enc.setBuffer(a, offset: 0, index: 0) enc.setBuffer(b, offset: 0, index: 1) enc.setBuffer(output, offset: 0, index: 2) var N = UInt32(count) enc.setBytes(&N, length: MemoryLayout.size, index: 3) let tg = engine.threadgroupSize1D(pso, count: count) enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: tg) enc.endEncoding() } func eltwiseAddScaled(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer, a: MTLBuffer, scaleA: Float, b: MTLBuffer, scaleB: Float, output: MTLBuffer, count: Int) 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: MemoryLayout.size, index: 1) enc.setBuffer(b, offset: 0, index: 2) var sb = scaleB enc.setBytes(&sb, length: MemoryLayout.size, index: 3) enc.setBuffer(output, offset: 0, index: 4) var N = UInt32(count) enc.setBytes(&N, length: MemoryLayout.size, index: 5) let tg = engine.threadgroupSize1D(pso, count: count) enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: tg) enc.endEncoding() } func fusedGateUp(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer, input: MTLBuffer, output: MTLBuffer) throws { let kernelName = gateProj.bits == 8 ? "quantized_matmul_gate_up_opt_8bit" : "quantized_matmul_gate_up_opt" if let pso = try? engine.pipeline(named: kernelName) { // Optimized path: threadgroup-cached input + uint4 loads let enc = cmdBuf.makeComputeCommandEncoder()! enc.setComputePipelineState(pso) enc.setBuffer(input, offset: 0, index: 0) enc.setBuffer(gateProj.weight, offset: 0, index: 1) enc.setBuffer(gateProj.scales, offset: 0, index: 2) enc.setBuffer(gateProj.biases, offset: 0, index: 3) enc.setBuffer(upProj.weight, offset: 0, index: 4) enc.setBuffer(upProj.scales, offset: 0, index: 5) enc.setBuffer(upProj.biases, offset: 0, index: 6) enc.setBuffer(output, offset: 0, index: 7) var inDim = UInt32(gateProj.inDim) enc.setBytes(&inDim, length: MemoryLayout.size, index: 8) var outDim = UInt32(gateProj.outDim) enc.setBytes(&outDim, length: MemoryLayout.size, index: 9) var groupSize = UInt32(gateProj.groupSize) enc.setBytes(&groupSize, length: MemoryLayout.size, index: 10) let tgMemSize = gateProj.inDim * 4 enc.setThreadgroupMemoryLength(tgMemSize, index: 0) let count = gateProj.outDim let tg = MTLSize(width: 256, height: 1, depth: 1) enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: tg) enc.endEncoding() } else { // Fallback to old kernel let fallbackName = gateProj.bits == 8 ? "quantized_matmul_gate_up_8bit" : "quantized_matmul_gate_up" let fallbackPSO = try engine.pipeline(named: fallbackName) let enc = cmdBuf.makeComputeCommandEncoder()! enc.setComputePipelineState(fallbackPSO) enc.setBuffer(input, offset: 0, index: 0) enc.setBuffer(gateProj.weight, offset: 0, index: 1) enc.setBuffer(gateProj.scales, offset: 0, index: 2) enc.setBuffer(gateProj.biases, offset: 0, index: 3) enc.setBuffer(upProj.weight, offset: 0, index: 4) enc.setBuffer(upProj.scales, offset: 0, index: 5) enc.setBuffer(upProj.biases, offset: 0, index: 6) enc.setBuffer(output, offset: 0, index: 7) var inDim = UInt32(gateProj.inDim) enc.setBytes(&inDim, length: MemoryLayout.size, index: 8) var outDim = UInt32(gateProj.outDim) enc.setBytes(&outDim, length: MemoryLayout.size, index: 9) var groupSize = UInt32(gateProj.groupSize) enc.setBytes(&groupSize, length: MemoryLayout.size, index: 10) let count = gateProj.outDim let tg = engine.threadgroupSize1D(fallbackPSO, count: count) enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: tg) enc.endEncoding() } } // ── MoE forward helpers ────────────────────── /// Quantized matmul for a specific expert slice within a 3D tensor. func quantizedMatmulExpert(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer, input: MTLBuffer, expert: MoEExpertGroup, expertIdx: Int, output: MTLBuffer) throws { let kernelName = expert.bits == 8 ? "quantized_matmul_simd_8bit" : "quantized_matmul_simd" guard let pso = try? engine.pipeline(named: kernelName) else { print(" [ERROR] quantizedMatmulExpert: Shader \(kernelName) not found! Falling back to quantized_matmul_seq") let fallbackPSO = try engine.pipeline(named: "quantized_matmul_seq") // Use fallback shader... let enc = cmdBuf.makeComputeCommandEncoder()! enc.setComputePipelineState(fallbackPSO) enc.setBuffer(input, offset: 0, index: 0) enc.setBuffer(expert.weight, offset: expert.weightStride * expertIdx, index: 1) enc.setBuffer(expert.scales, offset: expert.scalesStride * expertIdx, index: 2) enc.setBuffer(expert.biases, offset: expert.scalesStride * expertIdx, index: 3) enc.setBuffer(output, offset: 0, index: 4) var inDim = UInt32(expert.expertInDim) enc.setBytes(&inDim, length: MemoryLayout.size, index: 5) var outDim = UInt32(expert.expertOutDim) enc.setBytes(&outDim, length: MemoryLayout.size, index: 6) var groupSize = UInt32(expert.expertInDim / 64) enc.setBytes(&groupSize, length: MemoryLayout.size, index: 7) let tg = engine.threadgroupSize1D(fallbackPSO, count: expert.expertOutDim) enc.dispatchThreads(MTLSize(width: expert.expertOutDim, height: 1, depth: 1), threadsPerThreadgroup: tg) enc.endEncoding() return } let enc = cmdBuf.makeComputeCommandEncoder()! enc.setComputePipelineState(pso) enc.setBuffer(input, offset: 0, index: 0) enc.setBuffer(expert.weight, offset: expert.weightStride * expertIdx, index: 1) enc.setBuffer(expert.scales, offset: expert.scalesStride * expertIdx, index: 2) enc.setBuffer(expert.biases, offset: expert.scalesStride * expertIdx, index: 3) enc.setBuffer(output, offset: 0, index: 4) var inDim = UInt32(expert.expertInDim) enc.setBytes(&inDim, length: MemoryLayout.size, index: 5) var outDim = UInt32(expert.expertOutDim) enc.setBytes(&outDim, length: MemoryLayout.size, index: 6) var groupSize = UInt32(expert.expertInDim / expert.numGroups) // dynamic group size enc.setBytes(&groupSize, length: MemoryLayout.size, index: 7) let tgMemSize = expert.expertInDim * 4 enc.setThreadgroupMemoryLength(tgMemSize, index: 0) let count = expert.expertOutDim let tg = MTLSize(width: 256, height: 1, depth: 1) enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: tg) enc.endEncoding() } /// Fused gate+up matmul for a specific expert slice. func expertFusedGateUp(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer, input: MTLBuffer, gate: MoEExpertGroup, up: MoEExpertGroup, expertIdx: Int, output: MTLBuffer) throws { let kernelName = gate.bits == 8 ? "quantized_matmul_gate_up_8bit" : "quantized_matmul_gate_up" let pso = try engine.pipeline(named: kernelName) let enc = cmdBuf.makeComputeCommandEncoder()! enc.setComputePipelineState(pso) enc.setBuffer(input, offset: 0, index: 0) enc.setBuffer(gate.weight, offset: gate.weightStride * expertIdx, index: 1) enc.setBuffer(gate.scales, offset: gate.scalesStride * expertIdx, index: 2) enc.setBuffer(gate.biases, offset: gate.scalesStride * expertIdx, index: 3) enc.setBuffer(up.weight, offset: up.weightStride * expertIdx, index: 4) enc.setBuffer(up.scales, offset: up.scalesStride * expertIdx, index: 5) enc.setBuffer(up.biases, offset: up.scalesStride * expertIdx, index: 6) enc.setBuffer(output, offset: 0, index: 7) var inDim = UInt32(gate.expertInDim) enc.setBytes(&inDim, length: MemoryLayout.size, index: 8) var outDim = UInt32(gate.expertOutDim) enc.setBytes(&outDim, length: MemoryLayout.size, index: 9) var groupSize = UInt32(gate.expertInDim / 64) // group_size is 64 for quantized weights enc.setBytes(&groupSize, length: MemoryLayout.size, index: 10) let count = gate.expertOutDim let tg = engine.threadgroupSize1D(pso, count: count) enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: tg) enc.endEncoding() } /// Fused gate+up+down for a specific expert slice. /// Replaces: fusedGateUp + blit + downMatmul + scaledAdd with a single kernel. func expertFusedGateUpDown(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer, input: MTLBuffer, gate: MoEExpertGroup, up: MoEExpertGroup, down: MoEExpertGroup, expertIdx: Int, accum: MTLBuffer, weight: Float) throws { let kernelName = gate.bits == 8 ? "quantized_matmul_gate_up_down_8bit" : "quantized_matmul_gate_up_down" let pso = try engine.pipeline(named: kernelName) let enc = cmdBuf.makeComputeCommandEncoder()! enc.setComputePipelineState(pso) enc.setBuffer(input, offset: 0, index: 0) enc.setBuffer(gate.weight, offset: gate.weightStride * expertIdx, index: 1) enc.setBuffer(gate.scales, offset: gate.scalesStride * expertIdx, index: 2) enc.setBuffer(gate.biases, offset: gate.scalesStride * expertIdx, index: 3) enc.setBuffer(up.weight, offset: up.weightStride * expertIdx, index: 4) enc.setBuffer(up.scales, offset: up.scalesStride * expertIdx, index: 5) enc.setBuffer(up.biases, offset: up.scalesStride * expertIdx, index: 6) enc.setBuffer(down.weight, offset: down.weightStride * expertIdx, index: 7) enc.setBuffer(down.scales, offset: down.scalesStride * expertIdx, index: 8) enc.setBuffer(down.biases, offset: down.scalesStride * expertIdx, index: 9) enc.setBuffer(accum, offset: 0, index: 10) var hiddenSize = UInt32(gate.expertInDim) enc.setBytes(&hiddenSize, length: MemoryLayout.size, index: 11) var moeIntermediate = UInt32(gate.expertOutDim) enc.setBytes(&moeIntermediate, length: MemoryLayout.size, index: 12) var groupSize = UInt32(gate.expertInDim / gate.numGroups) enc.setBytes(&groupSize, length: MemoryLayout.size, index: 13) var w = weight enc.setBytes(&w, length: MemoryLayout.size, index: 14) let count = Int(max(hiddenSize, moeIntermediate)) let tgMemSize = count * 4 enc.setThreadgroupMemoryLength(tgMemSize, index: 0) let tg = MTLSize(width: 256, height: 1, depth: 1) enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: tg) enc.endEncoding() } func moeMegaKernel(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer, input: MTLBuffer, router: QuantizedWeights, gate: MoEExpertGroup, up: MoEExpertGroup, down: MoEExpertGroup, accum: MTLBuffer) throws -> Bool { guard let pso = try? engine.pipeline(named: "moe_mega_kernel") else { return false } let enc = cmdBuf.makeComputeCommandEncoder()! enc.setComputePipelineState(pso) enc.setBuffer(input, offset: 0, index: 0) // Router: full 3D weight buffer [numExperts * outDim, inDimPacked] enc.setBuffer(router.weight, offset: 0, index: 1) enc.setBuffer(router.scales, offset: 0, index: 2) enc.setBuffer(router.biases, offset: 0, index: 3) // Gate: full 3D weight buffer [numExperts * expertOutDim, expertInDimPacked] enc.setBuffer(gate.weight, offset: 0, index: 4) enc.setBuffer(gate.scales, offset: 0, index: 5) enc.setBuffer(gate.biases, offset: 0, index: 6) // Up: full 3D weight buffer enc.setBuffer(up.weight, offset: 0, index: 7) enc.setBuffer(up.scales, offset: 0, index: 8) enc.setBuffer(up.biases, offset: 0, index: 9) // Down: full 3D weight buffer [numExperts * expertOutDim, expertInDimPacked] enc.setBuffer(down.weight, offset: 0, index: 10) enc.setBuffer(down.scales, offset: 0, index: 11) enc.setBuffer(down.biases, offset: 0, index: 12) enc.setBuffer(accum, offset: 0, index: 13) var hiddenSize = UInt32(gate.expertInDim) enc.setBytes(&hiddenSize, length: MemoryLayout.size, index: 14) var moeIntermediate = UInt32(gate.expertOutDim) enc.setBytes(&moeIntermediate, length: MemoryLayout.size, index: 15) var numExperts = UInt32(gate.numExperts) enc.setBytes(&numExperts, length: MemoryLayout.size, index: 16) var rScale = routerScale enc.setBytes(&rScale, length: MemoryLayout.size, index: 17) var topK = UInt32(topK) enc.setBytes(&topK, length: MemoryLayout.size, index: 18) let count = Int(max(hiddenSize, moeIntermediate)) let logitStorage = Int(numExperts) + Int(topK) + Int(topK) let tgMemSize = (count + logitStorage) * 4 enc.setThreadgroupMemoryLength(tgMemSize, index: 0) let tg = MTLSize(width: 256, height: 1, depth: 1) enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: tg) enc.endEncoding() return true } func moeForward(input: MTLBuffer, ns: MTLBuffer, temps: ForwardTemps, cmdBuf: MTLCommandBuffer, engine: MarkBaseEngine) throws { guard let router = routerProj, let eGate = expertGate, let eUp = expertUp, let eDown = expertDown else { return } let hs = config.hiddenSize // ── Step 1: Copy MoE input ns → io (preserve from overwrite by expert output) ── let moeInput = temps.io let blit = cmdBuf.makeBlitCommandEncoder()! blit.copy(from: ns, sourceOffset: 0, to: moeInput, destinationOffset: 0, size: hs * 4) blit.endEncoding() // Zero accumulation buffer temps.h let zeroBlit = cmdBuf.makeBlitCommandEncoder()! zeroBlit.fill(buffer: temps.h, range: 0.. 0 { for i in 0.. $1.element } let topK = indexed.prefix(k) let topKSum = topK.reduce(0) { $0 + $1.element } guard topKSum > 0 else { return } // Compute experts on passed cmdBuf (can be batched with other layers) for (expertIdx, rawWeight) in topK { let weight = rawWeight / topKSum try expertFusedGateUpDown(engine: engine, cmdBuf: cmdBuf, input: moeInput, gate: eGate, up: eUp, down: eDown, expertIdx: expertIdx, accum: temps.h, weight: weight) } } // ── Step 5: Residual: input += moe_output (temps.h) scaled by layerScalar ── if layerScalar != 1.0 { try eltwiseAddScaled(engine: engine, cmdBuf: cmdBuf, a: input, scaleA: 1.0, b: temps.h, scaleB: layerScalar, output: input, count: hs) } else { try eltwiseAdd(engine: engine, cmdBuf: cmdBuf, a: input, b: temps.h, output: input, count: hs) } } // ── Main forward ───────────────────────────── /// Run one E4B layer forward pass. /// /// - Parameters: /// - input: [hiddenSize] — **modified in-place** to become the output /// - position: absolute token position (0-based) /// - kvCache: KV cache for reading attention. Must be non-nil. /// - shouldStoreKV: compute K,V and write into kvCache when true /// - temps: pre-allocated scratch buffers /// - engine: Metal engine public func forward(input: MTLBuffer, position: Int, kvCache: KVCache, shouldStoreKV: Bool, temps: ForwardTemps, engine: MarkBaseEngine, 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) let cmdBuf = engine.commandQueue.makeCommandBuffer()! try attentionForward(input: input, position: position, kvCache: kvCache, shouldStoreKV: shouldStoreKV, temps: temps, engine: engine, cmdBuf: cmdBuf) try moeForward(input: input, ns: temps.ns, temps: temps, cmdBuf: cmdBuf, engine: engine) try postFfnForward(input: input, temps: temps, engine: engine, cmdBuf: cmdBuf, perLayerInput: perLayerInput, perLayerInputOffset: perLayerInputOffset) cmdBuf.commit() cmdBuf.waitUntilCompleted() } else { // Dense path: unified flow for all positions let cmdBuf = engine.commandQueue.makeCommandBuffer()! try attentionForward(input: input, position: position, kvCache: kvCache, shouldStoreKV: shouldStoreKV, temps: temps, engine: engine, cmdBuf: cmdBuf) // FFN: gate+up fused → down → residual (scaled by layerScalar) 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) if layerScalar != 1.0 { try eltwiseAddScaled(engine: engine, cmdBuf: cmdBuf, a: input, scaleA: 1.0, b: temps.h, scaleB: layerScalar, output: input, count: config.hiddenSize) } else { try eltwiseAdd(engine: engine, cmdBuf: cmdBuf, a: input, b: temps.h, output: input, count: config.hiddenSize) } // Per-layer gating for dense path if let pg = perLayerGate, let pp = perLayerProjection, let pl = perLayerInput { try rmsNorm(engine: engine, cmdBuf: cmdBuf, input: input, weight: postFeedforwardLayernorm, output: temps.h, count: config.hiddenSize, eps: rmsNormEps) 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 eltwiseAdd(engine: engine, cmdBuf: cmdBuf, a: input, b: temps.h, output: input, count: config.hiddenSize) } cmdBuf.commit() cmdBuf.waitUntilCompleted() } } // ── Attention forward (steps 1-13) ── private func attentionForward(input: MTLBuffer, position: Int, kvCache: KVCache, shouldStoreKV: Bool, temps: ForwardTemps, engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer) throws { // ── 1. input_layernorm(x) → temps.h ── try rmsNorm(engine: engine, cmdBuf: cmdBuf, input: input, weight: inputLayernorm, output: temps.h, count: config.hiddenSize, eps: rmsNormEps) // ── 2. Q = q_proj(temps.h) → temps.q ── try quantizedMatmul(engine: engine, cmdBuf: cmdBuf, input: temps.h, 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.h, weights: kProj, output: temps.k) if let vp = vProj { try quantizedMatmul(engine: engine, cmdBuf: cmdBuf, input: temps.h, weights: vp, output: temps.v) } else if kEqualsV { let blit = cmdBuf.makeBlitCommandEncoder()! let copyBytes = config.nKvHeads * config.headDim * MemoryLayout.stride blit.copy(from: temps.k, sourceOffset: 0, to: temps.v, destinationOffset: 0, size: copyBytes) blit.endEncoding() } else { try eltwiseAddScaled(engine: engine, cmdBuf: cmdBuf, a: temps.v, scaleA: 0.0, b: temps.v, scaleB: 0.0, output: temps.v, count: config.nKvHeads * config.headDim) } // ── 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.h) // ── 11. Residual 1 (scaled by layerScalar) ── if layerScalar != 1.0 { try eltwiseAddScaled(engine: engine, cmdBuf: cmdBuf, a: input, scaleA: 1.0, b: temps.h, scaleB: layerScalar, output: input, count: config.hiddenSize) } else { try eltwiseAdd(engine: engine, cmdBuf: cmdBuf, a: input, b: temps.h, output: input, count: config.hiddenSize) } // ── 12. post_attention_layernorm → temps.h ── try rmsNorm(engine: engine, cmdBuf: cmdBuf, input: input, weight: postAttentionLayernorm, output: temps.h, count: config.hiddenSize, eps: rmsNormEps) // ── 13. pre_feedforward_layernorm → ns ── try rmsNorm(engine: engine, cmdBuf: cmdBuf, input: temps.h, weight: preFeedforwardLayernorm, output: temps.ns, count: config.hiddenSize, eps: rmsNormEps) } // ── Post-FFN forward (steps 17-19) ── private func postFfnForward(input: MTLBuffer, temps: ForwardTemps, engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer, perLayerInput: MTLBuffer?, perLayerInputOffset: Int) throws { // ── 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 eltwiseAdd(engine: engine, cmdBuf: cmdBuf, a: input, b: temps.h, output: input, count: config.hiddenSize) } else { try eltwiseAdd(engine: engine, cmdBuf: cmdBuf, a: input, b: temps.h, output: input, count: config.hiddenSize) } } }