markbaseengine/Sources/MarkBase/AsyncInference.swift

import Foundation
import Metal

// ─────────────────────────────────────────────────────────────
// Async Inference Optimizations
// ─────────────────────────────────────────────────────────────

/// Async inference result
public struct InferenceResult {
    public let logits: [Float]
    public let elapsed: TimeInterval
    public let token: Int
    
    public init(logits: [Float], elapsed: TimeInterval, token: Int = 0) {
        self.logits = logits
        self.elapsed = elapsed
        self.token = token
    }
}

/// Async inference queue for batching requests
public final class AsyncInferenceQueue {
    private let maxBatchSize: Int
    private var pending: [(input: Int, position: Int, completion: (Result<InferenceResult, Error>) -> Void)] = []
    private let lock = NSLock()
    
    public init(maxBatchSize: Int = 8) {
        self.maxBatchSize = maxBatchSize
    }
    
    /// Add inference request
    public func enqueue(input: Int, position: Int, completion: @escaping (Result<InferenceResult, Error>) -> Void) {
        lock.lock()
        pending.append((input, position, completion))
        let shouldProcess = pending.count >= maxBatchSize
        lock.unlock()
        
        if shouldProcess {
            processBatch()
        }
    }
    
    /// Process batch of requests
    private func processBatch() {
        lock.lock()
        let batch = pending.prefix(maxBatchSize)
        pending.removeFirst(min(batch.count, maxBatchSize))
        lock.unlock()
        
        // TODO: Implement actual batch processing
        // This would require modifications to the model to support batch inference
    }
}

/// Async token generator with prefetching
public final class AsyncTokenGenerator: @unchecked Sendable {
    private let model: E4BModel
    private let tokenizer: Tokenizer
    private let engine: MarkBaseEngine
    private let sampler: Sampler
    
    private var currentLogits: [Float]?
    private var currentPosition: Int = 0
    private var generatedTokens: [Int] = []
    
    public init(model: E4BModel, tokenizer: Tokenizer, engine: MarkBaseEngine) {
        self.model = model
        self.tokenizer = tokenizer
        self.engine = engine
        self.sampler = Sampler()
    }
    
    /// Start generation with async streaming
    public func start(prompt: String, config: GenerationConfig) -> AsyncStream<String> {
        return AsyncStream { [weak self] continuation in
            guard let self = self else {
                continuation.finish()
                return
            }
            
            Task.detached {
                do {
                    try await self.generateAsync(prompt: prompt, config: config) { tokenText in
                        continuation.yield(tokenText)
                    }
                    continuation.finish()
                } catch {
                    continuation.finish()
                }
            }
        }
    }
    
    /// Async generation with callback
    private func generateAsync(
        prompt: String,
        config: GenerationConfig,
        onToken: (String) -> Void
    ) async throws {
        let promptTokens = tokenizer.encode(text: prompt)
        
        // Pre-fill KV cache
        var lastLogits: [Float] = []
        for (position, tokenId) in promptTokens.enumerated() {
            lastLogits = try model.forward(tokenId: tokenId, position: position)
        }
        
        currentLogits = lastLogits
        currentPosition = promptTokens.count
        generatedTokens = []
        var streamDecoder = StreamingDecoder(tokenizer: tokenizer)
        
        // Generate tokens
        for _ in 0..<config.maxTokens {
            guard let logits = currentLogits else { break }
            
            // Sample next token
            let nextToken = sampler.sample(
                logits: logits,
                temperature: config.temperature,
                topK: config.topK,
                topP: config.topP
            )
            
            // Check EOS
            if tokenizer.eosTokenIds.contains(nextToken) {
                break
            }
            
            generatedTokens.append(nextToken)
            let tokenText = streamDecoder.consume(tokenId: nextToken)
            if !tokenText.isEmpty {
                onToken(tokenText)
            }
            
            // Forward pass
            let position = currentPosition
            do {
                let newLogits = try model.forward(tokenId: nextToken, position: position)
                currentLogits = newLogits
                currentPosition = position + 1
            } catch {
                break
            }
        }
    }
    
    /// Get current generation state
    public var state: (tokens: [Int], position: Int) {
        return (generatedTokens, currentPosition)
    }
}

/// Pre-fetching helper for overlapping computation
public final class Prefetcher: @unchecked Sendable {
    private var nextToken: Int?
    private var nextPosition: Int?
    private var prefetchedLogits: [Float]?
    private let model: E4BModel
    
    public init(model: E4BModel) {
        self.model = model
    }
    
    /// Start prefetching for next token
    public func startPrefetch(token: Int, position: Int) async {
        nextToken = token
        nextPosition = position
        
        // Prefetch in background
        do {
            let logits = try model.forward(tokenId: token, position: position)
            prefetchedLogits = logits
        } catch {
            // Prefetch failed, will compute on demand
        }
    }
    
    /// Get prefetched result or compute on demand
    public func getOrCompute(token: Int, position: Int) throws -> [Float] {
        if let logits = prefetchedLogits, nextToken == token, nextPosition == position {
            prefetchedLogits = nil
            return logits
        }
        
        // Compute on demand
        return try model.forward(tokenId: token, position: position)
    }
}