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
2026-06-23 18:12:35 +08:00
commit ac75faa0cc
301 changed files with 63426 additions and 0 deletions
@@ -0,0 +1,143 @@
+import XCTest
+@testable import MarkBase
+
+final class BatchEmbeddingOptimizationTest: XCTestCase {
+    
+    func testBatchEmbeddingPerformance() throws {
+        print("\n═══════════════════════════════════════════════════════════════════")
+        print("  Batch Embedding Optimization Test")
+        print("═══════════════════════════════════════════════════════════════════\n")
+        
+        let engine = try MarkBaseEngine(autoCompile: true)
+        
+        // Test with E4B-MarkBase (smaller, faster to load)
+        let modelPath = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
+        
+print("Loading E4B-MarkBase...")
+let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 512)
+print("✓ Model loaded: \(model.layers.count) layers\n")
+
+// Create batch context (using model's helper method)
+let maxBatchSize = 8
+let context = model.createBatchContext(maxBatchSize: maxBatchSize)
+        
+        print("Testing batch embedding performance:")
+        print("  Batch sizes: 1, 2, 4, 8\n")
+        
+        for batchSize in [1, 2, 4, 8] {
+            print("Testing batch(\(batchSize))...")
+            
+            // Generate random token IDs
+            let tokenIds = (0..<batchSize).map { _ in Int.random(in: 0..<model.vocabSize) }
+            let positions = Array(0..<batchSize)
+            
+            // Test 5 forward passes
+            var times: [Double] = []
+            for i in 0..<5 {
+                let start = Date()
+                let logits = try model.forwardBatchTrue(
+                    tokenIds: tokenIds,
+                    positions: positions,
+                    context: context
+                )
+                let elapsed = Date().timeIntervalSince(start) * 1000
+                times.append(elapsed)
+                
+                // Verify output
+                XCTAssertEqual(logits.count, batchSize, "Batch size mismatch")
+                XCTAssertEqual(logits[0].count, model.vocabSize, "Vocab size mismatch")
+                
+                let nanCount = logits.flatMap { $0 }.filter { $0.isNaN }.count
+                XCTAssertEqual(nanCount, 0, "NaN detected in batch output")
+            }
+            
+            let avgTime = times.reduce(0, +) / Double(times.count)
+            let perToken = avgTime / Double(batchSize)
+            
+            print("  ✓ Batch(\(batchSize)): avg \(String(format: "%.1f", avgTime))ms total, \(String(format: "%.1f", perToken))ms/token")
+        }
+        
+        print("\n═══════════════════════════════════════════════════════════════════")
+        print("  Batch Embedding Optimization COMPLETE ✓✓✓")
+        print("  Single GPU dispatch for entire batch (eliminated sequential waits)")
+        print("═══════════════════════════════════════════════════════════════════\n")
+    }
+    
+    func test31BBatchPerformance() throws {
+        print("\n═══════════════════════════════════════════════════════════════════")
+        print("  31B Batch Performance Test")
+        print("═══════════════════════════════════════════════════════════════════\n")
+        
+        let engine = try MarkBaseEngine(autoCompile: true)
+        let modelPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-31b-it-4bit"
+        
+        print("Loading 31B model...")
+        let startLoad = Date()
+        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 512)
+        let loadTime = Date().timeIntervalSince(startLoad) * 1000
+print("✓ Model loaded in \(String(format: "%.1f", loadTime))ms")
+print("  Layers: \(model.layers.count)")
+print("  Hidden: \(model.hiddenSize)\n")
+
+// Create batch context (using model's helper method)
+let maxBatchSize = 4
+let context = model.createBatchContext(maxBatchSize: maxBatchSize)
+        
+        print("Testing 31B batch performance:")
+        
+        // Single token baseline
+        print("\nSingle token baseline...")
+        var singleTimes: [Double] = []
+        for i in 0..<3 {
+            let start = Date()
+            let logits = try model.forwardOptimized(tokenId: 2, position: i)
+            let elapsed = Date().timeIntervalSince(start) * 1000
+            singleTimes.append(elapsed)
+            
+            let nanCount = logits.filter { $0.isNaN }.count
+            XCTAssertEqual(nanCount, 0, "NaN in single forward")
+        }
+        let singleAvg = singleTimes.reduce(0, +) / Double(singleTimes.count)
+        print("  ✓ Single: \(String(format: "%.1f", singleAvg))ms/token")
+        
+        // Batch(4) test
+        print("\nBatch(4) test...")
+        let tokenIds = [2, 2, 2, 2]
+        let positions = [0, 1, 2, 3]
+        
+        var batchTimes: [Double] = []
+        for i in 0..<3 {
+            let start = Date()
+            let logits = try model.forwardBatchTrue(
+                tokenIds: tokenIds,
+                positions: positions,
+                context: context
+            )
+            let elapsed = Date().timeIntervalSince(start) * 1000
+            batchTimes.append(elapsed)
+            
+            XCTAssertEqual(logits.count, 4, "Batch size mismatch")
+            let nanCount = logits.flatMap { $0 }.filter { $0.isNaN }.count
+            XCTAssertEqual(nanCount, 0, "NaN in batch forward")
+        }
+        
+        let batchAvg = batchTimes.reduce(0, +) / Double(batchTimes.count)
+        let batchPerToken = batchAvg / 4.0
+        let speedup = singleAvg / batchPerToken
+        
+        print("  ✓ Batch(4): \(String(format: "%.1f", batchAvg))ms total, \(String(format: "%.1f", batchPerToken))ms/token")
+        print("  ✓ Speedup: \(String(format: "%.1f", speedup))x")
+        
+        // Target check
+        if batchPerToken < 50.0 {
+            print("\n✓✓✓ TARGET MET! 31B batch <50ms/token")
+        } else {
+            print("\n⚠ Target not met: \(String(format: "%.1f", batchPerToken))ms/token (target <50ms)")
+            print("  Further optimization needed")
+        }
+        
+        print("\n═══════════════════════════════════════════════════════════════════")
+        print("  31B Batch Test Complete")
+        print("═══════════════════════════════════════════════════════════════════\n")
+    }
+}