markbaseengine/Sources/MarkBase/BatchGenerationTrue.swift

import Metal

// ══════════════════════════════════════════════════════════════════
// TRUE Batch Generation - Using batch Metal kernels
// Expected: 8-15x speedup for batch inference
// ══════════════════════════════════════════════════════════════════

extension E4BModel {
    
    /// TRUE batch forward pass - process multiple tokens with batch kernels
    /// This achieves real parallelism, not sequential processing
    public func forwardBatchTrue(
        tokenIds: [Int],
        positions: [Int],
        context: BatchContext
    ) throws -> [[Float]] {
        guard tokenIds.count == positions.count else { return [] }
        let batchSize = tokenIds.count
        guard batchSize <= context.maxBatchSize else { return [] }
        
if batchSize == 0 { return [] }
if batchSize == 1 {
return [try forwardOptimized(tokenId: tokenIds[0], position: positions[0])]
}

// ── Phase 1: Embedding Lookup (FIXED: Use batch kernel) ──
// Debug: Check embedWeight parameters BEFORE batch embedding
print("BEFORE batch embedding:")
print("  hiddenSize=\(hiddenSize)")
print("  embedWeight.groupSize=\(embedWeight.groupSize)")
print("  embedWeight.weight.length=\(embedWeight.weight.length)")
print("  embedWeight.scales.length=\(embedWeight.scales.length)")
print("  embedWeight.biases.length=\(embedWeight.biases.length)")
print("  embedWeight.inDim=\(embedWeight.inDim)")
print("  embedWeight.outDim=\(embedWeight.outDim)")
print("  vocabSize=\(vocabSize)")
print("  batchSize=\(batchSize)")
print("  embedScale=\(embedScale) (should be ~50.6 for hiddenSize=2560)")
print("  tokenIds=\(tokenIds)")

// Prepare tokenIds array for Metal
let tokenIdsBuffer = engine.device.makeBuffer(
    bytes: tokenIds.map { UInt32($0) },
    length: batchSize * 4,
    options: .storageModeShared
)!

// Use batch embedding kernel
let embedCmdBuf = engine.commandQueue.makeCommandBuffer()!
let pso = try engine.pipeline(named: embedScale != 1.0 ? "dequantize_row_batch_scaled" : "dequantize_row_batch")
let enc = embedCmdBuf.makeComputeCommandEncoder()!
enc.setComputePipelineState(pso)

enc.setBuffer(embedWeight.weight, offset: 0, index: 0)
enc.setBuffer(embedWeight.scales, offset: 0, index: 1)
enc.setBuffer(embedWeight.biases, offset: 0, index: 2)
enc.setBuffer(tokenIdsBuffer, offset: 0, index: 3)
enc.setBuffer(context.batchInputBuffer, offset: 0, index: 4)

var nCols = UInt32(hiddenSize)
var batchSz = UInt32(batchSize)
var groupSz = UInt32(embedWeight.groupSize)
enc.setBytes(&nCols, length: 4, index: 5)
enc.setBytes(&batchSz, length: 4, index: 6)
enc.setBytes(&groupSz, length: 4, index: 7)

if embedScale != 1.0 {
    var scale = embedScale
    enc.setBytes(&scale, length: 4, index: 8)
}

// Calculate threadgroup size (2D grid: batchSize × hiddenSize)
let threadsPerThreadgroup = MTLSize(width: 32, height: 8, depth: 1)
let gridSize = MTLSize(width: batchSize, height: hiddenSize, depth: 1)
enc.dispatchThreads(gridSize, threadsPerThreadgroup: threadsPerThreadgroup)

enc.endEncoding()
embedCmdBuf.commit()
embedCmdBuf.waitUntilCompleted()
        
        // ── Phase 2: Layer Processing with BATCH KERNELS ──
        let layerCmdBuf = engine.commandQueue.makeCommandBuffer()!
        
        // Create batch temps for layer processing
        let batchTemps = try temps.createBatchBuffers(
            device: engine.device,
            batchSize: batchSize,
            hiddenSize: hiddenSize,
            nHeads: layers[0].config.nHeads,
            headDim: layers[0].config.headDim,
            intermediateSize: layers[0].config.intermediateSize
        )
        
        // Process all 42 layers with batch kernels
        for layerIdx in 0..<numHiddenLayers {
            let isOwner = layerIdx < firstKVShared
            let cacheIdx = isOwner ? layerIdx : (kvSourceMap[layerIdx] ?? (layerIdx - numKvShared))
            let cache = kvCaches[cacheIdx]
            
            // Use batch layer processing
            try layers[layerIdx].forwardBatchTrue(
                batchInput: context.batchInputBuffer,
                positions: positions,
                batchSize: batchSize,
                kvCache: cache,
                shouldStoreKV: isOwner,
                temps: temps,
                batchTemps: batchTemps,
                engine: engine,
                cmdBuf: layerCmdBuf
            )
        }
        
        // ── Phase 3: Final Norm + LM Head (batch) ──
        if let fn = finalNorm {
            // Inline batch RMS norm
            let pso = try engine.pipeline(named: "rms_norm_batch")
            let enc = layerCmdBuf.makeComputeCommandEncoder()!
            enc.setComputePipelineState(pso)
            
            enc.setBuffer(context.batchInputBuffer, offset: 0, index: 0)
            enc.setBuffer(fn, offset: 0, index: 1)
            enc.setBuffer(context.batchInputBuffer, offset: 0, index: 2)  // In-place
            
            var N = UInt32(hiddenSize)
            enc.setBytes(&N, 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()
        }
        
        // Batch LM head
        let psoLM = try engine.pipeline(named: "quantized_matmul_batch")
        let encLM = layerCmdBuf.makeComputeCommandEncoder()!
        encLM.setComputePipelineState(psoLM)
        
        encLM.setBuffer(context.batchInputBuffer, offset: 0, index: 0)
        encLM.setBuffer(embedWeight.weight, offset: 0, index: 1)
        encLM.setBuffer(embedWeight.scales, offset: 0, index: 2)
        encLM.setBuffer(embedWeight.biases, offset: 0, index: 3)
        encLM.setBuffer(context.batchOutputBuffer, offset: 0, index: 4)
        
        var inDim = UInt32(embedWeight.inDim)
        encLM.setBytes(&inDim, length: 4, index: 5)
        var outDim = UInt32(embedWeight.outDim)
        encLM.setBytes(&outDim, length: 4, index: 6)
        var groupSize = UInt32(embedWeight.groupSize)
        encLM.setBytes(&groupSize, length: 4, index: 7)
        var batchLM = UInt32(batchSize)
        encLM.setBytes(&batchLM, length: 4, index: 8)
        
        let tgLM = MTLSize(width: 256, height: 1, depth: 1)
        let gridLM = MTLSize(width: batchSize, height: embedWeight.outDim, depth: 1)
        encLM.dispatchThreads(gridLM, threadsPerThreadgroup: tgLM)
        encLM.endEncoding()
        
        // Logits scaling and softcapping (batch)
        if embedWeight.groupSize == 32 {
            let logitsScale = Float(30.0 / 116.23 / sqrt(Float(hiddenSize)))
            // Use eltwise_scale for batch scaling
            let pso = try engine.pipeline(named: "eltwise_scale")
            let enc = layerCmdBuf.makeComputeCommandEncoder()!
            enc.setComputePipelineState(pso)
            
            enc.setBuffer(context.batchOutputBuffer, offset: 0, index: 0)
            var ls = logitsScale
            enc.setBytes(&ls, length: 4, index: 1)
            var total = UInt32(batchSize * vocabSize)
            enc.setBytes(&total, length: 4, index: 2)
            
            let tg = MTLSize(width: 256, height: 1, depth: 1)
            let grid = MTLSize(width: batchSize * vocabSize, height: 1, depth: 1)
            enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
            enc.endEncoding()
        }
        
        // Softcapping (skip if kernel not found)
        if let cap = finalLogitSoftcapping {
            // Try to use tanh_scale kernel
            do {
                let pso = try engine.pipeline(named: "tanh_scale")
                let enc = layerCmdBuf.makeComputeCommandEncoder()!
                enc.setComputePipelineState(pso)
                
                enc.setBuffer(context.batchOutputBuffer, offset: 0, index: 0)
                var c = cap
                enc.setBytes(&c, length: 4, index: 1)
                var total = UInt32(batchSize * vocabSize)
                enc.setBytes(&total, length: 4, index: 2)
                
                let tg = MTLSize(width: 256, height: 1, depth: 1)
                let grid = MTLSize(width: batchSize * vocabSize, height: 1, depth: 1)
                enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
                enc.endEncoding()
            } catch {
                // Skip softcapping if kernel not found
            }
        }
        
        // Single commit for entire batch
        layerCmdBuf.commit()
        layerCmdBuf.waitUntilCompleted()
        
        // Read results
        let outputPtr = context.batchOutputBuffer.contents().assumingMemoryBound(to: Float.self)
        var results: [[Float]] = []
        for i in 0..<batchSize {
            let logits = Array(UnsafeBufferPointer<Float>(
                start: outputPtr + i * vocabSize,
                count: vocabSize
            ))
            results.append(logits)
        }
        
        return results
    }
    
}