markbaseengine/Sources/MarkBase/Layers/LayerOptimized.swift

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