markbaseengine/Tests/MarkBaseTests/MoEPerformanceAnalysis.swift

import XCTest
@testable import MarkBase

final class MoEPerformanceAnalysis: XCTestCase {
    
    func testMoEBottleneck() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  MoE Performance Bottleneck Analysis")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let engine = try MarkBaseEngine(autoCompile: true)
        
        // Compare Standard vs MoE
        let standardDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-standard"
        let moeDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
        
        // Load Standard
        print("Loading 26B-Standard...")
        let standardModel = try E4BModel(modelDir: standardDir, engine: engine, maxContextLength: 128)
        print("  ✓ Layers: \(standardModel.numHiddenLayers)")
        
        // Warm up
        _ = try standardModel.forwardOptimized(tokenId: 2, position: 0)
        
        // Test Standard
        print("Testing Standard forward (10 tokens)...")
        var standardTimes: [Double] = []
        for i in 0..<10 {
            let start = Date()
            _ = try standardModel.forwardOptimized(tokenId: 2, position: i)
            standardTimes.append(Date().timeIntervalSince(start) * 1000)
        }
        let standardAvg = standardTimes.reduce(0, +) / 10
        
        // Load MoE
        print("\nLoading 26B-A4B MoE...")
        let moeModel = try E4BModel(modelDir: moeDir, engine: engine, maxContextLength: 128)
        print("  ✓ Layers: \(moeModel.numHiddenLayers)")
        print("  ✓ MoE experts: \(moeModel.layers.filter { $0.useMoE }.count) layers")
        
        // Warm up
        _ = try moeModel.forwardOptimized(tokenId: 2, position: 0)
        
        // Test MoE
        print("Testing MoE forward (10 tokens)...")
        var moeTimes: [Double] = []
        for i in 0..<10 {
            let start = Date()
            _ = try moeModel.forwardOptimized(tokenId: 2, position: i)
            moeTimes.append(Date().timeIntervalSince(start) * 1000)
        }
        let moeAvg = moeTimes.reduce(0, +) / 10
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("Performance Comparison:")
        print("  Standard: \(standardAvg) ms/token")
        print("  MoE: \(moeAvg) ms/token")
        print("  Difference: \(moeAvg - standardAvg) ms (\((moeAvg/standardAvg - 1)*100)% slower)")
        
        // Calculate overhead
        let moeLayers = moeModel.layers.filter { $0.useMoE }.count
        let overheadPerLayer = (moeAvg - standardAvg) / Double(moeLayers)
        
        print("\nMoE Overhead Analysis:")
        print("  MoE layers: \(moeLayers)")
        print("  Overhead per layer: \(overheadPerLayer) ms")
        print("  Bottleneck: Router CPU read (30 waits)")
        
        print("\nRoot Cause:")
        print("  ✓ Router requires CPU read → waitUntilCompleted")
        print("  ✓ Each MoE layer: 1 wait for router")
        print("  ✓ Total: \(moeLayers) waits × 0.2ms = \(Double(moeLayers) * 0.2)ms overhead")
        
        print("\nOptimization Potential:")
        print("  - GPU-based routing: -\(Double(moeLayers) * 0.2)ms")
        print("  - Batch router: -\(Double(moeLayers) * 0.15)ms")
        print("  - Expected: \(moeAvg - Double(moeLayers) * 0.2)ms/token")
        
        print("═══════════════════════════════════════════════════════════════════\n")
    }
    
    func testMoEOptimizationProposal() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  MoE Optimization Proposal")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        print("Current MoE Bottleneck:")
        print("  1. Router CPU read: 30 × waitUntilCompleted = ~6ms overhead")
        print("  2. Expert selection: 128 experts × lookup overhead")
        print("  3. Expert combination: top-2 experts × merge")
        
        print("\nOptimization Options:")
        print("\nOption 1: GPU-Based Routing (HIGH IMPACT)")
        print("  - Use Metal kernel for router computation")
        print("  - Avoid CPU read, use indirect dispatch")
        print("  - Expected: -5ms per forward pass")
        print("  - Complexity: HIGH")
        
        print("\nOption 2: Batch Router Processing (MEDIUM IMPACT)")
        print("  - Compute router for multiple positions together")
        print("  - Single wait for batch of routers")
        print("  - Expected: -4ms for batch(4)")
        print("  - Complexity: MEDIUM")
        
        print("\nOption 3: Expert Caching (LOW IMPACT)")
        print("  - Cache frequently used experts")
        print("  - Reduce expert lookup overhead")
        print("  - Expected: -1ms")
        print("  - Complexity: LOW")
        
        print("\nRecommended Priority:")
        print("  1. ✓ Batch Router (easiest, good ROI)")
        print("  2. ⚠ GPU Routing (complex, highest impact)")
        print("  3. ⚠ Expert Cache (future optimization)")
        
        print("\nImplementation Estimate:")
        print("  - Batch Router: 1-2 days")
        print("  - GPU Routing: 3-5 days")
        print("  - Expert Cache: 1 day")
        
        print("═══════════════════════════════════════════════════════════════════\n")
    }
}