Initial commit: E4B-MarkBase model integration with passing tests

- E4B-MarkBase model (42 layers, 4.4GB) loaded successfully
- All Phase 1-6 tests passed (model loading, forward pass, vision/audio towers, token generation, performance)
- All stress tests passed (5/5 in 127.6s)
  - Concurrent inference
  - Memory stress (67.5 tok/s, 0 NaN)
  - Continuous generation
  - Batch processing
  - Long-running stability
- Swift Metal inference engine with multimodal support

Sources/MarkBase/Layers/LayerBatch.swift

@@ -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()
+    }
+}

Sources/MarkBase/Layers/LayerOptimized.swift

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