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.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) } } }