v2: Initial clean branch with unit tests + CI/CD pipeline

- Started from ac75faa (initial E4B-MarkBase integration) - Kept Sources/ (all engine code) + Package.swift + .gitignore - Removed all ad-hoc tests, documentation, scripts, Python files - Added Tests/00_Unit/ (MathTest, TokenizerTest, SamplerTest) - Added .gitea/workflows/ci.yaml (build + unit tests + lint) - Added Scripts/check_resources.sh (memory-aware test runner) - Added Tests/Manifest.json (resource requirements for all tests) - Focus: 4-bit quantized models only
2026-07-05 13:29:25 +08:00
commit 8a66b9086a
90 changed files with 22252 additions and 0 deletions
@@ -0,0 +1,190 @@
+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)
+    }
+}
+