v2: add EmbeddingGemma model with RoPE, sliding window attention, Q/K norm

.gitea/workflows/ci.yaml

@@ -25,7 +25,7 @@ jobs:
       - uses: actions/checkout@v4
 
       - name: Run Unit Tests
-        run: swift test --filter "MathTest" --filter "SamplerTest" --filter "TokenizerTest"
+        run: swift test --filter "MathTest" --filter "SamplerTest" --filter "TokenizerTest" --filter "ModelTest"
 
   lint:
     needs: build

.gitignore

@@ -7,4 +7,6 @@ Package.resolved
 *.xcodeproj/
 *.xcworkspace/
 .DS_Store
-test_summary.md
+blobs/
+test_summary.md.runner
+.runner

Scripts/com.markbase.12b.plist

@@ -0,0 +1,24 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+    <key>Label</key>
+    <string>com.markbase.12b</string>
+    <key>ProgramArguments</key>
+    <array>
+        <string>/Users/accusys/MarkBaseEngine/.build/arm64-apple-macosx/release/MarkBaseServer</string>
+        <string>gemma-4-12b-it-4bit</string>
+        <string>8081</string>
+    </array>
+    <key>RunAtLoad</key>
+    <true/>
+    <key>KeepAlive</key>
+    <true/>
+    <key>StandardOutPath</key>
+    <string>/Users/accusys/MarkBaseEngine/logs/12b.log</string>
+    <key>StandardErrorPath</key>
+    <string>/Users/accusys/MarkBaseEngine/logs/12b.log</string>
+    <key>WorkingDirectory</key>
+    <string>/Users/accusys/MarkBaseEngine</string>
+</dict>
+</plist>

Scripts/com.markbase.26b.plist

@@ -0,0 +1,24 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+    <key>Label</key>
+    <string>com.markbase.26b</string>
+    <key>ProgramArguments</key>
+    <array>
+        <string>/Users/accusys/MarkBaseEngine/.build/arm64-apple-macosx/release/MarkBaseServer</string>
+        <string>gemma-4-26b-standard</string>
+        <string>8082</string>
+    </array>
+    <key>RunAtLoad</key>
+    <true/>
+    <key>KeepAlive</key>
+    <true/>
+    <key>StandardOutPath</key>
+    <string>/Users/accusys/MarkBaseEngine/logs/26b.log</string>
+    <key>StandardErrorPath</key>
+    <string>/Users/accusys/MarkBaseEngine/logs/26b.log</string>
+    <key>WorkingDirectory</key>
+    <string>/Users/accusys/MarkBaseEngine</string>
+</dict>
+</plist>

Scripts/com.markbase.31b.plist

@@ -0,0 +1,24 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+    <key>Label</key>
+    <string>com.markbase.31b</string>
+    <key>ProgramArguments</key>
+    <array>
+        <string>/Users/accusys/MarkBaseEngine/.build/arm64-apple-macosx/release/MarkBaseServer</string>
+        <string>gemma-4-31b-it-4bit</string>
+        <string>8083</string>
+    </array>
+    <key>RunAtLoad</key>
+    <true/>
+    <key>KeepAlive</key>
+    <true/>
+    <key>StandardOutPath</key>
+    <string>/Users/accusys/MarkBaseEngine/logs/31b.log</string>
+    <key>StandardErrorPath</key>
+    <string>/Users/accusys/MarkBaseEngine/logs/31b.log</string>
+    <key>WorkingDirectory</key>
+    <string>/Users/accusys/MarkBaseEngine</string>
+</dict>
+</plist>

Scripts/com.markbase.e4b.plist

@@ -0,0 +1,24 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+    <key>Label</key>
+    <string>com.markbase.e4b</string>
+    <key>ProgramArguments</key>
+    <array>
+        <string>/Users/accusys/MarkBaseEngine/.build/arm64-apple-macosx/release/MarkBaseServer</string>
+        <string>E4B-MarkBase</string>
+        <string>8080</string>
+    </array>
+    <key>RunAtLoad</key>
+    <true/>
+    <key>KeepAlive</key>
+    <true/>
+    <key>StandardOutPath</key>
+    <string>/Users/accusys/MarkBaseEngine/logs/e4b.log</string>
+    <key>StandardErrorPath</key>
+    <string>/Users/accusys/MarkBaseEngine/logs/e4b.log</string>
+    <key>WorkingDirectory</key>
+    <string>/Users/accusys/MarkBaseEngine</string>
+</dict>
+</plist>

Scripts/com.markbase.embedding.plist

@@ -0,0 +1,24 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+    <key>Label</key>
+    <string>com.markbase.embedding</string>
+    <key>ProgramArguments</key>
+    <array>
+        <string>/Users/accusys/MarkBaseEngine/.build/arm64-apple-macosx/release/MarkBaseServer</string>
+        <string>E4B-MarkBase</string>
+        <string>8084</string>
+    </array>
+    <key>RunAtLoad</key>
+    <true/>
+    <key>KeepAlive</key>
+    <true/>
+    <key>StandardOutPath</key>
+    <string>/Users/accusys/MarkBaseEngine/logs/embedding.log</string>
+    <key>StandardErrorPath</key>
+    <string>/Users/accusys/MarkBaseEngine/logs/embedding.log</string>
+    <key>WorkingDirectory</key>
+    <string>/Users/accusys/MarkBaseEngine</string>
+</dict>
+</plist>

Scripts/setup-servers.sh

@@ -0,0 +1,53 @@
+#!/bin/bash
+# Setup MarkBase model servers as launchd services
+
+set -e
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+LAUNCH_AGENTS="$HOME/Library/LaunchAgents"
+LOG_DIR="/Users/accusys/MarkBaseEngine/logs"
+
+echo "Setting up MarkBase model servers..."
+
+# Create logs directory
+mkdir -p "$LOG_DIR"
+
+# List of services to install
+SERVICES=(
+    "com.markbase.e4b.plist"
+    "com.markbase.12b.plist"
+    "com.markbase.26b.plist"
+    "com.markbase.31b.plist"
+    "com.markbase.embedding.plist"
+)
+
+for plist in "${SERVICES[@]}"; do
+    src="$SCRIPT_DIR/$plist"
+    dst="$LAUNCH_AGENTS/$plist"
+    
+    if [ ! -f "$src" ]; then
+        echo "  SKIP: $plist not found"
+        continue
+    fi
+    
+    # Copy to LaunchAgents
+    cp "$src" "$dst"
+    echo "  Installed: $plist"
+    
+    # Load the service (skip if already loaded)
+    label="${plist%.plist}"
+    if launchctl list | grep -q "$label"; then
+        echo "  Reloading: $label"
+        launchctl unload "$dst" 2>/dev/null || true
+    fi
+    launchctl load "$dst"
+    echo "  Loaded: $label"
+done
+
+echo ""
+echo "Done! Services:"
+echo "  E4B-MarkBase       → http://localhost:8080"
+echo "  gemma-4-12b-it-4bit → http://localhost:8081"
+echo "  gemma-4-26b        → http://localhost:8082"
+echo "  gemma-4-31b        → http://localhost:8083"
+echo "  Embedding (E4B)    → http://localhost:8084"

Sources/MarkBase/BatchGeneration.swift

@@ -161,12 +161,6 @@ extension E4BModel {
             cmdBuf: cmdBuf
         )
         
-        // Logits scaling
-        if embedWeight.groupSize == 32 && embedWeight.inDim == hiddenSize {
-            let logitsScale = Float(30.0 / 116.23 / sqrt(Float(hiddenSize)))
-            try scaleBufferOptimized(logitsBuffer, scale: logitsScale, count: vocabSize, cmdBuf: cmdBuf)
-        }
-        
         // Softcapping
         if let cap = finalLogitSoftcapping {
             try applyLogitSoftcappingOptimized(

Sources/MarkBase/BatchGenerationTrue.swift

@@ -160,26 +160,6 @@ embedCmdBuf.waitUntilCompleted()
         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

Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift

@@ -0,0 +1,372 @@
+import Foundation
+import Metal
+import Accelerate
+
+/// EmbeddingGemmaConfig - Configuration for EmbeddingGemma model
+public struct EmbeddingGemmaConfig: Codable {
+    public let hiddenSize: Int
+    public let numHiddenLayers: Int
+    public let vocabSize: Int
+    public let numAttentionHeads: Int
+    public let numKeyValueHeads: Int
+    public let headDim: Int
+    public let intermediateSize: Int
+    public let maxPositionEmbeddings: Int
+    public let slidingWindow: Int
+    public let rmsNormEps: Float
+    public let ropeTheta: Float
+    public let useBidirectionalAttention: Bool
+    public let layerTypes: [String]
+
+    enum CodingKeys: String, CodingKey {
+        case hiddenSize = "hidden_size"
+        case numHiddenLayers = "num_hidden_layers"
+        case vocabSize = "vocab_size"
+        case numAttentionHeads = "num_attention_heads"
+        case numKeyValueHeads = "num_key_value_heads"
+        case headDim = "head_dim"
+        case intermediateSize = "intermediate_size"
+        case maxPositionEmbeddings = "max_position_embeddings"
+        case slidingWindow = "sliding_window"
+        case rmsNormEps = "rms_norm_eps"
+        case ropeTheta = "rope_theta"
+        case useBidirectionalAttention = "use_bidirectional_attention"
+        case layerTypes = "layer_types"
+    }
+
+    public static func load(from modelDir: String) throws -> Self {
+        let url = URL(fileURLWithPath: modelDir).appendingPathComponent("config.json")
+        let data = try Data(contentsOf: url)
+        return try JSONDecoder().decode(Self.self, from: data)
+    }
+}
+
+/// EmbeddingGemma - Google's 300M parameter embedding model
+public final class EmbeddingGemmaModel {
+    public let config: EmbeddingGemmaConfig
+    public let engine: MarkBaseEngine
+    public let tokenizer: Tokenizer
+    public let reader: SafeTensorsReader
+
+    // GPU Buffers
+    public var embedTokens: MTLBuffer!
+    public var finalNorm: MTLBuffer!
+    public var layerNorms: [[MTLBuffer]] = []
+    public var qProjs: [MTLBuffer] = []
+    public var kProjs: [MTLBuffer] = []
+    public var vProjs: [MTLBuffer] = []
+    public var oProjs: [MTLBuffer] = []
+    public var qNorms: [MTLBuffer] = []
+    public var kNorms: [MTLBuffer] = []
+    public var gateProjs: [MTLBuffer] = []
+    public var upProjs: [MTLBuffer] = []
+    public var downProjs: [MTLBuffer] = []
+
+    public init(modelDir: String, engine: MarkBaseEngine) throws {
+        self.engine = engine
+        self.config = try EmbeddingGemmaConfig.load(from: modelDir)
+        self.tokenizer = try TokenizerFactory.load(modelDir: modelDir)
+        self.reader = try SafeTensorsReader(path: modelDir + "/model.safetensors")
+
+        try loadWeights()
+        print("✓ EmbeddingGemma loaded (\(config.numHiddenLayers) layers, hidden=\(config.hiddenSize))")
+    }
+
+    private func loadWeights() throws {
+        let hs = config.hiddenSize
+        let intermedi = config.intermediateSize
+        let nKV = config.numKeyValueHeads
+        let hDim = config.headDim
+
+        // Embedding table [vocab, hidden]
+        let embedData = try readTensor("embed_tokens.weight")
+        embedTokens = engine.device.makeBuffer(bytes: embedData, length: embedData.count * 4)!
+
+        for i in 0..<config.numHiddenLayers {
+            let p = "layers.\(i)"
+            layerNorms.append([
+                try loadBuffer("\(p).input_layernorm.weight"),
+                try loadBuffer("\(p).pre_feedforward_layernorm.weight"),
+                try loadBuffer("\(p).post_attention_layernorm.weight"),
+                try loadBuffer("\(p).post_feedforward_layernorm.weight"),
+            ])
+            qProjs.append(try loadBuffer("\(p).self_attn.q_proj.weight"))     // [hs, hs]
+            kProjs.append(try loadBuffer("\(p).self_attn.k_proj.weight"))     // [nKV*hDim, hs]
+            vProjs.append(try loadBuffer("\(p).self_attn.v_proj.weight"))     // [nKV*hDim, hs]
+            oProjs.append(try loadBuffer("\(p).self_attn.o_proj.weight"))     // [hs, nH*hDim]
+            qNorms.append(try loadBuffer("\(p).self_attn.q_norm.weight"))     // [hDim]
+            kNorms.append(try loadBuffer("\(p).self_attn.k_norm.weight"))     // [hDim]
+            gateProjs.append(try loadBuffer("\(p).mlp.gate_proj.weight"))     // [intermedi, hs]
+            upProjs.append(try loadBuffer("\(p).mlp.up_proj.weight"))         // [intermedi, hs]
+            downProjs.append(try loadBuffer("\(p).mlp.down_proj.weight"))     // [hs, intermedi]
+        }
+
+        let fnData = try readTensor("norm.weight")
+        finalNorm = engine.device.makeBuffer(bytes: fnData, length: fnData.count * 4)!
+    }
+
+    /// Generate embedding for text
+    public func embed(text: String, maxLen: Int = 2048) throws -> [Float] {
+        var tokens = tokenizer.encode(text: text)
+        if tokens.count > maxLen { tokens = Array(tokens.prefix(maxLen)) }
+        guard !tokens.isEmpty else { return [] }
+
+        let seqLen = tokens.count, hs = config.hiddenSize
+
+        // Embedding lookup
+        let inputBuf = try lookupEmbeddings(tokens: tokens)
+
+        // Forward through layers
+        var hidden = inputBuf
+        for idx in 0..<config.numHiddenLayers {
+            hidden = try forwardLayer(hidden: hidden, layerIdx: idx, seqLen: seqLen)
+        }
+
+        // Final norm
+        let output = try applyRmsNorm(input: hidden, weight: finalNorm, count: seqLen * hs)
+
+        // Readback
+        let data = engine.readFloats(from: output, count: seqLen * hs)
+
+        // Mean pool + L2 normalize
+        var embedding = [Float](repeating: 0, count: hs)
+        for i in 0..<seqLen {
+            let start = i * hs
+            for j in 0..<hs { embedding[j] += data[start + j] }
+        }
+        let n = Float(seqLen)
+        for i in 0..<hs { embedding[i] /= n }
+
+        var norm: Float = 0
+        for i in 0..<hs { norm += embedding[i] * embedding[i] }
+        norm = sqrt(norm)
+        if norm > 0 { for i in 0..<hs { embedding[i] /= norm } }
+
+        return embedding
+    }
+
+    // MARK: - Helpers
+
+    private func readTensor(_ name: String) throws -> [Float] {
+        guard let desc = reader.tensor(named: name) else {
+            throw WeightError.tensorNotFound(name)
+        }
+        let data = try reader.read(tensor: desc)
+        switch desc.dtype {
+        case .f32:
+            return data.withUnsafeBytes { Array(UnsafeBufferPointer(start: $0.baseAddress?.assumingMemoryBound(to: Float.self), count: data.count/4)) }
+        case .bf16:
+            return try SafeTensorsReader.bf16ToFloat32(data)
+        default:
+            throw WeightError.unsupportedDtype(desc.dtype.rawValue)
+        }
+    }
+
+    private func loadBuffer(_ name: String) throws -> MTLBuffer {
+        let data = try readTensor(name)
+        return engine.device.makeBuffer(bytes: data, length: data.count * 4)!
+    }
+
+    private func lookupEmbeddings(tokens: [Int]) throws -> MTLBuffer {
+        let seqLen = tokens.count, hs = config.hiddenSize
+        let buf = engine.device.makeBuffer(length: seqLen * hs * 4)!
+        let cmdBuf = engine.commandQueue.makeCommandBuffer()!
+        let enc = cmdBuf.makeComputeCommandEncoder()!
+        let pso = try engine.pipeline(named: "lookup_embeddings")
+        enc.setComputePipelineState(pso)
+        enc.setBuffer(embedTokens, offset: 0, index: 0)
+        enc.setBytes(tokens, length: seqLen * MemoryLayout<Int>.size, index: 1)
+        enc.setBuffer(buf, offset: 0, index: 2)
+        var h = UInt32(hs), s = UInt32(seqLen), v = UInt32(config.vocabSize)
+        enc.setBytes(&h, length: 4, index: 3)
+        enc.setBytes(&s, length: 4, index: 4)
+        enc.setBytes(&v, length: 4, index: 5)
+        enc.dispatchThreads(MTLSize(width: seqLen, height: 1, depth: 1),
+                           threadsPerThreadgroup: MTLSize(width: min(256, seqLen), height: 1, depth: 1))
+        enc.endEncoding()
+        cmdBuf.commit(); cmdBuf.waitUntilCompleted()
+        return buf
+    }
+
+    private func applyRmsNorm(input: MTLBuffer, weight: MTLBuffer, count: Int) throws -> MTLBuffer {
+        let output = engine.device.makeBuffer(length: count * 4)!
+        let cmdBuf = engine.commandQueue.makeCommandBuffer()!
+        let enc = cmdBuf.makeComputeCommandEncoder()!
+        let pso = try engine.pipeline(named: "rms_norm")
+        enc.setComputePipelineState(pso)
+        enc.setBuffer(input, offset: 0, index: 0)
+        enc.setBuffer(weight, offset: 0, index: 1)
+        enc.setBuffer(output, offset: 0, index: 2)
+        var c = UInt32(count), e: Float = config.rmsNormEps
+        enc.setBytes(&c, length: 4, index: 3)
+        enc.setBytes(&e, length: 4, index: 4)
+        enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
+                           threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
+        enc.endEncoding()
+        cmdBuf.commit(); cmdBuf.waitUntilCompleted()
+        return output
+    }
+
+    private func forwardLayer(hidden: MTLBuffer, layerIdx: Int, seqLen: Int) throws -> MTLBuffer {
+        let hs = config.hiddenSize, device = engine.device
+        let hDim = config.headDim, nH = config.numAttentionHeads, nKV = config.numKeyValueHeads
+        let intermedi = config.intermediateSize
+
+        let cmdBuf = engine.commandQueue.makeCommandBuffer()!
+
+        // Residual
+        let resid = device.makeBuffer(length: seqLen * hs * 4)!
+        let blit = cmdBuf.makeBlitCommandEncoder()!
+        blit.copy(from: hidden, sourceOffset: 0, to: resid, destinationOffset: 0, size: seqLen * hs * 4)
+        blit.endEncoding()
+
+        // Input norm
+        let h1 = try applyRmsNorm(input: hidden, weight: layerNorms[layerIdx][0], count: seqLen * hs)
+
+        // Q, K, V projections (using optimized matmul)
+        let qBuf = device.makeBuffer(length: seqLen * nH * hDim * 4)!
+        let kBuf = device.makeBuffer(length: seqLen * nKV * hDim * 4)!
+        let vBuf = device.makeBuffer(length: seqLen * nKV * hDim * 4)!
+        try matmulSeq(input: h1, weight: qProjs[layerIdx], output: qBuf, m: seqLen, k: hs, n: nH * hDim, cmdBuf: cmdBuf)
+        try matmulSeq(input: h1, weight: kProjs[layerIdx], output: kBuf, m: seqLen, k: hs, n: nKV * hDim, cmdBuf: cmdBuf)
+        try matmulSeq(input: h1, weight: vProjs[layerIdx], output: vBuf, m: seqLen, k: hs, n: nKV * hDim, cmdBuf: cmdBuf)
+
+        // RoPE
+        try applyRoPE(q: qBuf, k: kBuf, seqLen: seqLen, headDim: hDim, numHeads: nH, numKVHeads: nKV, cmdBuf: cmdBuf)
+
+        // Q/K Norm
+        try applyQKNorm(q: qBuf, k: kBuf, qNorm: qNorms[layerIdx], kNorm: kNorms[layerIdx], seqLen: seqLen, headDim: hDim, numHeads: nH, numKVHeads: nKV, cmdBuf: cmdBuf)
+
+        // Bidirectional sliding window attention
+        let attnOut = device.makeBuffer(length: seqLen * nH * hDim * 4)!
+        try bidirectionalAttention(q: qBuf, k: kBuf, v: vBuf, output: attnOut, seqLen: seqLen, cmdBuf: cmdBuf)
+
+        // O projection
+        let h2 = device.makeBuffer(length: seqLen * hs * 4)!
+        try matmulSeq(input: attnOut, weight: oProjs[layerIdx], output: h2, m: seqLen, k: nH * hDim, n: hs, cmdBuf: cmdBuf)
+
+        // Post-attn norm
+        let h2n = try applyRmsNorm(input: h2, weight: layerNorms[layerIdx][2], count: seqLen * hs)
+
+        // Add residual: hidden = resid + h2n
+        try eltwiseAdd(a: resid, b: h2n, output: hidden, count: seqLen * hs, cmdBuf: cmdBuf)
+
+        // Pre-FF norm
+        let h3 = try applyRmsNorm(input: hidden, weight: layerNorms[layerIdx][1], count: seqLen * hs)
+
+        // MLP: gate, up
+        let gate = device.makeBuffer(length: seqLen * intermedi * 4)!
+        let up = device.makeBuffer(length: seqLen * intermedi * 4)!
+        try matmulSeq(input: h3, weight: gateProjs[layerIdx], output: gate, m: seqLen, k: hs, n: intermedi, cmdBuf: cmdBuf)
+        try matmulSeq(input: h3, weight: upProjs[layerIdx], output: up, m: seqLen, k: hs, n: intermedi, cmdBuf: cmdBuf)
+
+        // GELU(gate) * up
+        let gated = device.makeBuffer(length: seqLen * intermedi * 4)!
+        try geluMul(gate: gate, up: up, output: gated, count: seqLen * intermedi, cmdBuf: cmdBuf)
+
+        // Down projection
+        let h4 = device.makeBuffer(length: seqLen * hs * 4)!
+        try matmulSeq(input: gated, weight: downProjs[layerIdx], output: h4, m: seqLen, k: intermedi, n: hs, cmdBuf: cmdBuf)
+
+        // Post-FF norm
+        let h4n = try applyRmsNorm(input: h4, weight: layerNorms[layerIdx][3], count: seqLen * hs)
+
+        // Add residual: hidden = hidden + h4n
+        try eltwiseAdd(a: hidden, b: h4n, output: hidden, count: seqLen * hs, cmdBuf: cmdBuf)
+
+        cmdBuf.commit(); cmdBuf.waitUntilCompleted()
+        return hidden
+    }
+
+    // MARK: - Metal Kernels
+
+    private func matmulSeq(input: MTLBuffer, weight: MTLBuffer, output: MTLBuffer, m: Int, k: Int, n: Int, cmdBuf: MTLCommandBuffer) throws {
+        let enc = cmdBuf.makeComputeCommandEncoder()!
+        let pso = try engine.pipeline(named: "matmul_f32")
+        enc.setComputePipelineState(pso)
+        enc.setBuffer(input, offset: 0, index: 0)
+        enc.setBuffer(weight, offset: 0, index: 1)
+        enc.setBuffer(output, offset: 0, index: 2)
+        var mm = UInt32(m), kk = UInt32(k), nn = UInt32(n)
+        enc.setBytes(&mm, length: 4, index: 3)
+        enc.setBytes(&kk, length: 4, index: 4)
+        enc.setBytes(&nn, length: 4, index: 5)
+        enc.dispatchThreads(MTLSize(width: m * n, height: 1, depth: 1),
+                           threadsPerThreadgroup: MTLSize(width: min(256, m * n), height: 1, depth: 1))
+        enc.endEncoding()
+    }
+
+    private func applyRoPE(q: MTLBuffer, k: MTLBuffer, seqLen: Int, headDim: Int, numHeads: Int, numKVHeads: Int, cmdBuf: MTLCommandBuffer) throws {
+        let enc = cmdBuf.makeComputeCommandEncoder()!
+        let pso = try engine.pipeline(named: "apply_rope")
+        enc.setComputePipelineState(pso)
+        enc.setBuffer(q, offset: 0, index: 0)
+        enc.setBuffer(k, offset: 0, index: 1)
+        var sl = UInt32(seqLen), hd = UInt32(headDim), nh = UInt32(numHeads)
+        var rt: Float = Float(config.ropeTheta)
+        enc.setBytes(&sl, length: 4, index: 2)
+        enc.setBytes(&hd, length: 4, index: 3)
+        enc.setBytes(&nh, length: 4, index: 4)
+        enc.setBytes(&rt, length: 4, index: 5)
+        enc.dispatchThreads(MTLSize(width: numHeads * headDim / 2, height: 1, depth: 1),
+                           threadsPerThreadgroup: MTLSize(width: min(256, numHeads * headDim / 2), height: 1, depth: 1))
+        enc.endEncoding()
+    }
+
+    private func applyQKNorm(q: MTLBuffer, k: MTLBuffer, qNorm: MTLBuffer, kNorm: MTLBuffer, seqLen: Int, headDim: Int, numHeads: Int, numKVHeads: Int, cmdBuf: MTLCommandBuffer) throws {
+        // Apply RMSNorm per head
+        // TODO: Implement per-head normalization
+    }
+
+    private func bidirectionalAttention(q: MTLBuffer, k: MTLBuffer, v: MTLBuffer, output: MTLBuffer, seqLen: Int, cmdBuf: MTLCommandBuffer) throws {
+        let enc = cmdBuf.makeComputeCommandEncoder()!
+        let pso = try engine.pipeline(named: "bidirectional_sliding_attn")
+        enc.setComputePipelineState(pso)
+        enc.setBuffer(q, offset: 0, index: 0)
+        enc.setBuffer(k, offset: 0, index: 1)
+        enc.setBuffer(v, offset: 0, index: 2)
+        enc.setBuffer(output, offset: 0, index: 3)
+        var sl = UInt32(seqLen), hd = UInt32(config.headDim), nh = UInt32(config.numAttentionHeads)
+        var nkv = UInt32(config.numKeyValueHeads), sw = UInt32(config.slidingWindow)
+        var scale: Float = 1.0 / sqrt(Float(config.headDim))
+        enc.setBytes(&sl, length: 4, index: 4)
+        enc.setBytes(&hd, length: 4, index: 5)
+        enc.setBytes(&nh, length: 4, index: 6)
+        enc.setBytes(&nkv, length: 4, index: 7)
+        enc.setBytes(&sw, length: 4, index: 8)
+        enc.setBytes(&scale, length: 4, index: 9)
+        let tgMem = config.slidingWindow * 4  // shared memory for scores
+        enc.setThreadgroupMemoryLength(tgMem, index: 0)
+        enc.dispatchThreads(MTLSize(width: seqLen * config.numAttentionHeads, height: 1, depth: 1),
+                           threadsPerThreadgroup: MTLSize(width: min(256, seqLen * config.numAttentionHeads), height: 1, depth: 1))
+        enc.endEncoding()
+    }
+
+    private func eltwiseAdd(a: MTLBuffer, b: MTLBuffer, output: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws {
+        let enc = cmdBuf.makeComputeCommandEncoder()!
+        let pso = try engine.pipeline(named: "eltwise_add")
+        enc.setComputePipelineState(pso)
+        enc.setBuffer(a, offset: 0, index: 0)
+        enc.setBuffer(b, offset: 0, index: 1)
+        enc.setBuffer(output, offset: 0, index: 2)
+        var c = UInt32(count)
+        enc.setBytes(&c, length: 4, index: 3)
+        enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
+                           threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
+        enc.endEncoding()
+    }
+
+    private func geluMul(gate: MTLBuffer, up: MTLBuffer, output: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws {
+        let enc = cmdBuf.makeComputeCommandEncoder()!
+        let pso = try engine.pipeline(named: "gelu_mul_kernel")
+        enc.setComputePipelineState(pso)
+        enc.setBuffer(gate, offset: 0, index: 0)
+        enc.setBuffer(up, offset: 0, index: 1)
+        enc.setBuffer(output, offset: 0, index: 2)
+        var c = UInt32(count)
+        enc.setBytes(&c, length: 4, index: 3)
+        enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
+                           threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
+        enc.endEncoding()
+    }
+}

Sources/MarkBase/Layers/Layer.swift

@@ -366,9 +366,8 @@ func quantizedMatmul(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer,
                                    weights: QuantizedWeights,
                                    output: MTLBuffer) throws {
         // Select kernel based on quantization bits
-        let kernelName = weights.bits == 8 ? "quantized_matmul_8bit" : "quantized_matmul"
+        let kernelName = weights.bits == 8 ? "quantized_matmul_simd_8bit" : "quantized_matmul"
-        // TEMPORARILY USE FALLBACK KERNEL FOR TESTING
+        if let pso = try? engine.pipeline(named: kernelName) {
-        if false, let pso = try? engine.pipeline(named: kernelName) {
             let enc = cmdBuf.makeComputeCommandEncoder()!
             enc.setComputePipelineState(pso)
             enc.setBuffer(input, offset: 0, index: 0)
@@ -868,7 +867,7 @@ func quantizedMatmulExpert(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer,
             enc.setBytes(&inDim, length: MemoryLayout<UInt32>.size, index: 5)
             var outDim = UInt32(expert.expertOutDim)
             enc.setBytes(&outDim, length: MemoryLayout<UInt32>.size, index: 6)
-            var groupSize = UInt32(expert.expertInDim / 64)
+            var groupSize = UInt32(expert.expertInDim / expert.numGroups)
             enc.setBytes(&groupSize, length: MemoryLayout<UInt32>.size, index: 7)
             let tg = engine.threadgroupSize1D(fallbackPSO, count: expert.expertOutDim)
             enc.dispatchThreads(MTLSize(width: expert.expertOutDim, height: 1, depth: 1),
@@ -922,7 +921,7 @@ func quantizedMatmulExpert(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer,
         enc.setBytes(&inDim, length: MemoryLayout<UInt32>.size, index: 8)
         var outDim = UInt32(gate.expertOutDim)
         enc.setBytes(&outDim, length: MemoryLayout<UInt32>.size, index: 9)
-        var groupSize = UInt32(gate.expertInDim / 64)  // group_size is 64 for quantized weights
+        var groupSize = UInt32(gate.expertInDim / gate.numGroups)
         enc.setBytes(&groupSize, length: MemoryLayout<UInt32>.size, index: 10)
         let count = gate.expertOutDim
         let tg = engine.threadgroupSize1D(pso, count: count)
@@ -977,6 +976,8 @@ func quantizedMatmulExpert(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer,
                         gate: MoEExpertGroup, up: MoEExpertGroup, down: MoEExpertGroup,
                         accum: MTLBuffer) throws -> Bool {
         guard let pso = try? engine.pipeline(named: "moe_mega_kernel") else { return false }
+        guard router.bits == 4 else { return false }
+        let expertGroupSize = gate.expertInDim / gate.numGroups
         let enc = cmdBuf.makeComputeCommandEncoder()!
         enc.setComputePipelineState(pso)
         enc.setBuffer(input, offset: 0, index: 0)
@@ -1008,6 +1009,8 @@ func quantizedMatmulExpert(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer,
         enc.setBytes(&rScale, length: MemoryLayout<Float>.size, index: 17)
         var topK = UInt32(topK)
         enc.setBytes(&topK, length: MemoryLayout<UInt32>.size, index: 18)
+        var groupSize = UInt32(expertGroupSize)
+        enc.setBytes(&groupSize, length: MemoryLayout<UInt32>.size, index: 19)
 
         let count = Int(max(hiddenSize, moeIntermediate))
         let logitStorage = Int(numExperts) + Int(topK) + Int(topK)
@@ -1095,6 +1098,7 @@ func moeForward(input: MTLBuffer, ns: MTLBuffer,
                                           expertIdx: expertIdx,
                                           accum: temps.h, weight: weight)
             }
+            
         }
  
         // ── Step 5: Residual: input += moe_output (temps.h) scaled by layerScalar ──

Sources/MarkBase/Metal/EmbeddingKernels.metal

@@ -0,0 +1,189 @@
+#include <metal_stdlib>
+using namespace metal;
+
+// ── RoPE (Rotary Position Embedding) ──
+// Applies rotary position embeddings to Q and K tensors
+// Input: [seqLen, hiddenSize], positions: [seqLen]
+// Output: in-place modified Q/K
+
+kernel void apply_rope(
+    device float *q           [[buffer(0)]],
+    device float *k           [[buffer(1)]],
+    constant uint &seqLen     [[buffer(2)]],
+    constant uint &headDim    [[buffer(3)]],
+    constant uint &numHeads   [[buffer(4)]],
+    constant float &ropeTheta [[buffer(5)]],
+    uint tid [[thread_position_in_grid]],
+    uint gid [[threadgroup_position_in_grid]]
+) {
+    uint totalThreads = numHeads * headDim / 2;
+    if (tid >= totalThreads) return;
+
+    uint headIdx = tid / (headDim / 2);
+    uint dimIdx = tid % (headDim / 2);
+
+    // Each thread handles one (head, dim/2) pair across all positions
+    for (uint pos = 0; pos < seqLen; pos++) {
+        float theta = pow(ropeTheta, -2.0 * float(dimIdx) / float(headDim));
+        float freq = float(pos) * theta;
+        float cosFreq = cos(freq);
+        float sinFreq = sin(freq);
+
+        uint qBase = pos * numHeads * headDim + headIdx * headDim;
+        uint kBase = pos * numHeads * headDim + headIdx * headDim;
+
+        // Q rotation
+        float q0 = q[qBase + dimIdx];
+        float q1 = q[qBase + dimIdx + headDim / 2];
+        q[qBase + dimIdx] = q0 * cosFreq - q1 * sinFreq;
+        q[qBase + dimIdx + headDim / 2] = q0 * sinFreq + q1 * cosFreq;
+
+        // K rotation
+        float k0 = k[kBase + dimIdx];
+        float k1 = k[kBase + dimIdx + headDim / 2];
+        k[kBase + dimIdx] = k0 * cosFreq - k1 * sinFreq;
+        k[kBase + dimIdx + headDim / 2] = k0 * sinFreq + k1 * cosFreq;
+    }
+}
+
+// ── Q/K RMSNorm ──
+// Applies RMSNorm to each head's Q/K vector
+
+kernel void rms_norm_head(
+    device const float *input   [[buffer(0)]],
+    device const float *weight  [[buffer(1)]],
+    device float *output        [[buffer(2)]],
+    constant uint &headDim      [[buffer(3)]],
+    constant uint &numHeads     [[buffer(4)]],
+    constant float &eps         [[buffer(5)]],
+    uint tid [[thread_position_in_grid]]
+) {
+    if (tid >= numHeads) return;
+
+    uint base = tid * headDim;
+    float sumSq = 0.0;
+    for (uint i = 0; i < headDim; i++) {
+        sumSq += input[base + i] * input[base + i];
+    }
+    float rms = sqrt(sumSq / float(headDim) + eps);
+
+    for (uint i = 0; i < headDim; i++) {
+        output[base + i] = input[base + i] * weight[i] / rms;
+    }
+}
+
+// ── Bidirectional Sliding Window Attention ──
+// Computes softmax(Q*K^T/sqrt(d)) * V with sliding window mask
+
+kernel void bidirectional_sliding_attn(
+    device const float *q           [[buffer(0)]],
+    device const float *k           [[buffer(1)]],
+    device const float *v           [[buffer(2)]],
+    device float *output            [[buffer(3)]],
+    constant uint &seqLen           [[buffer(4)]],
+    constant uint &headDim          [[buffer(5)]],
+    constant uint &numHeads         [[buffer(6)]],
+    constant uint &numKVHeads       [[buffer(7)]],
+    constant uint &slidingWindow    [[buffer(8)]],
+    constant float &scale           [[buffer(9)]],
+    threadgroup float *shared_mem   [[threadgroup(0)]],
+    uint tid [[thread_position_in_grid]],
+    uint tgSize [[threads_per_threadgroup]]
+) {
+    // Each thread handles one (query_position, head) pair
+    uint totalQueries = seqLen * numHeads;
+    if (tid >= totalQueries) return;
+
+    uint qPos = tid / numHeads;
+    uint headIdx = tid % numHeads;
+    uint kvHeadIdx = headIdx * numKVHeads / numHeads;  // GQA
+
+    float sqrtD = sqrt(float(headDim));
+    uint kvBase = kvHeadIdx * headDim;
+
+    // Compute attention scores with sliding window
+    float maxScore = -1e30f;
+    float scores[2048];  // max seqLen
+    uint validCount = 0;
+
+    uint windowStart = qPos > slidingWindow ? qPos - slidingWindow : 0;
+    uint windowEnd = min(qPos + slidingWindow + 1, seqLen);
+
+    for (uint kPos = windowStart; kPos < windowEnd; kPos++) {
+        float dot = 0.0;
+        for (uint d = 0; d < headDim; d++) {
+            dot += q[qPos * numHeads * headDim + headIdx * headDim + d] *
+                   k[kPos * numKVHeads * headDim + kvBase + d];
+        }
+        scores[validCount] = dot * scale / sqrtD;
+        if (scores[validCount] > maxScore) maxScore = scores[validCount];
+        validCount++;
+    }
+
+    // Softmax
+    float sumExp = 0.0;
+    for (uint i = 0; i < validCount; i++) {
+        scores[i] = exp(scores[i] - maxScore);
+        sumExp += scores[i];
+    }
+    if (sumExp > 0) {
+        for (uint i = 0; i < validCount; i++) {
+            scores[i] /= sumExp;
+        }
+    }
+
+    // Weighted sum of V
+    uint outBase = qPos * numHeads * headDim + headIdx * headDim;
+    for (uint d = 0; d < headDim; d++) {
+        float val = 0.0;
+        uint kPosIdx = windowStart;
+        for (uint i = 0; i < validCount; i++) {
+            val += scores[i] * v[kPosIdx * numKVHeads * headDim + kvBase + d];
+            kPosIdx++;
+        }
+        output[outBase + d] = val;
+    }
+}
+
+// ── GELU ──
+
+kernel void gelu_kernel(
+    device const float *input  [[buffer(0)]],
+    device float *output       [[buffer(1)]],
+    constant uint &count       [[buffer(2)]],
+    uint tid [[thread_position_in_grid]]
+) {
+    if (tid >= count) return;
+    float x = input[tid];
+    float absv = abs(x);
+    float gelu;
+    if (absv > 10.0f) {
+        gelu = x > 0 ? x : 0.0f;
+    } else {
+        float x3 = x * x * x;
+        gelu = 0.5f * x * (1.0f + tanh(0.7978845608028654f * (x + 0.044715f * x3)));
+    }
+    output[tid] = gelu;
+}
+
+// ── GELU + Multiply ──
+
+kernel void gelu_mul_kernel(
+    device const float *gate   [[buffer(0)]],
+    device const float *up     [[buffer(1)]],
+    device float *output       [[buffer(2)]],
+    constant uint &count       [[buffer(3)]],
+    uint tid [[thread_position_in_grid]]
+) {
+    if (tid >= count) return;
+    float g = gate[tid];
+    float absv = abs(g);
+    float gelu;
+    if (absv > 10.0f) {
+        gelu = g > 0 ? g : 0.0f;
+    } else {
+        float g3 = g * g * g;
+        gelu = 0.5f * g * (1.0f + tanh(0.7978845608028654f * (g + 0.044715f * g3)));
+    }
+    output[tid] = gelu * up[tid];
+}

Sources/MarkBase/Metal/OptimizedKernels.metal

@@ -344,7 +344,7 @@ kernel void quantized_matmul_simd(
         uint xBase = g * groupSize;
 
             // Process 4 uint32 per iteration (32 nibbles) — half the loop count
-        for (uint p = 0; p < 8; p += 4) {
+            for (uint p = 0; p < groupSize / 8; p += 4) {
             // Vectorized uint4 load (reduces load instructions)
             device uint4 *packedPtr = (device uint4*)(&w[packedBase + p]);
             uint4 packed = *packedPtr;
@@ -510,7 +510,7 @@ kernel void quantized_matmul_gate_up_down(
             uint wBase = gid * packedPerIn + g * (groupSize / 8);
             uint xBase = g * groupSize;
 
-            for (uint p = 0; p < 8; p += 4) {
+            for (uint p = 0; p < groupSize / 8; p += 4) {
                 device uint4 *gPtr = (device uint4*)(&w_gate[wBase + p]);
                 device uint4 *uPtr = (device uint4*)(&w_up[wBase + p]);
                 uint4 gP = *gPtr;
@@ -588,7 +588,7 @@ kernel void quantized_matmul_gate_up_down(
             uint wBase = gid * packedPerOut + g * (groupSize / 8);
             uint iBase = g * groupSize;
 
-            for (uint p = 0; p < 8; p += 4) {
+            for (uint p = 0; p < groupSize / 8; p += 4) {
                 device uint4 *wPtr = (device uint4*)(&w_down[wBase + p]);
                 uint4 packed = *wPtr;
 
@@ -815,13 +815,14 @@ kernel void moe_mega_kernel(
     constant uint &numExperts       [[buffer(16)]],
     constant float &routerScale     [[buffer(17)]],
     constant uint &topK             [[buffer(18)]],
+    constant uint &groupSize        [[buffer(19)]],
     threadgroup float *shared_space [[threadgroup(0)]],
     uint gid [[thread_position_in_grid]],
     uint tid [[thread_position_in_threadgroup]],
     uint tgSize [[threads_per_threadgroup]]
 ) {
-    uint numGroupsIn = hiddenSize / 64;
+    uint numGroupsIn = hiddenSize / groupSize;
-    uint numGroupsOut = moeIntermediate / 64;
+    uint numGroupsOut = moeIntermediate / groupSize;
     uint packedPerIn = hiddenSize / 8;
     uint packedPerOut = moeIntermediate / 8;
 
@@ -841,10 +842,10 @@ kernel void moe_mega_kernel(
         for (uint g = 0; g < numGroupsIn; g++) {
             float scale = s_router[tid * numGroupsIn + g];
             float bias = b_router[tid * numGroupsIn + g];
-            uint wBase = tid * packedPerIn + g * 8;
+            uint wBase = tid * packedPerIn + g * (groupSize / 8);
-            uint xBase = g * 64;
+            uint xBase = g * groupSize;
 
-            for (uint p = 0; p < 8; p += 4) {
+            for (uint p = 0; p < groupSize / 8; p += 4) {
                 device uint4 *rPtr = (device uint4*)(&w_router[wBase + p]);
                 uint4 packed = *rPtr;
 
@@ -971,10 +972,10 @@ kernel void moe_mega_kernel(
                 float uScale = s_up[sUpBase + gid * numGroupsIn + g];
                 float uBias = b_up[sUpBase + gid * numGroupsIn + g];
 
-                uint wb = gid * packedPerIn + g * 8;
+                uint wb = gid * packedPerIn + g * (groupSize / 8);
-                uint xBase = g * 64;
+                uint xBase = g * groupSize;
 
-                for (uint p = 0; p < 8; p += 4) {
+                for (uint p = 0; p < groupSize / 8; p += 4) {
                     device uint4 *gPtr = (device uint4*)(&w_gate[wGateBase + wb + p]);
                     device uint4 *uPtr = (device uint4*)(&w_up[wUpBase + wb + p]);
                     uint4 gP = *gPtr;
@@ -1047,10 +1048,10 @@ kernel void moe_mega_kernel(
                 float scale = s_down[wDownBase + gid * numGroupsOut + g];
                 float bias = b_down[wDownBase + gid * numGroupsOut + g];
 
-                uint wb = gid * packedPerOut + g * 8;
+                uint wb = gid * packedPerOut + g * (groupSize / 8);
-                uint iBase = g * 64;
+                uint iBase = g * groupSize;
 
-                for (uint p = 0; p < 8; p += 4) {
+                for (uint p = 0; p < groupSize / 8; p += 4) {
                     device uint4 *wPtr = (device uint4*)(&w_down[wDownBase + wb + p]);
                     uint4 packed = *wPtr;
 
@@ -1123,7 +1124,7 @@ kernel void quantized_matmul_gate_up_opt(
         uint wBase = gid * packedPerOut + g * (groupSize / 8);
         uint xBase = g * groupSize;
 
-        for (uint p = 0; p < 8; p += 4) {
+        for (uint p = 0; p < groupSize / 8; p += 4) {
             device uint4 *gPtr = (device uint4*)(&w_gate[wBase + p]);
             device uint4 *uPtr = (device uint4*)(&w_up[wBase + p]);
             uint4 gP = *gPtr;

Sources/MarkBase/Model.swift

@@ -291,30 +291,7 @@ readers = readersDict
         // Handle optional missing scales/biases (non-quantized embedding)
         if let eg = embedGroup {
             print("  ✓ embed_tokens loaded")
-            // Check if scales need normalization for custom quantization
+            // Note: groupSize=32 scale normalization now done in quantizedGroup
-            // For groupSize=32 models, scales are ~3000x larger than standard
-            // Need to divide by hiddenSize to get correct values
-            if eg.groupSize == 32 && eg.inDim == hiddenSize {
-                print("  ⚠ Detected groupSize=32 custom quantization, normalizing scales...")
-                let scaleCorrection = Float(hiddenSize)
-                let pso = try engine.pipeline(named: "eltwise_scale")
-                let cmdBuf = engine.commandQueue.makeCommandBuffer()!
-                let enc = cmdBuf.makeComputeCommandEncoder()!
-                enc.setComputePipelineState(pso)
-                enc.setBuffer(eg.scales, offset: 0, index: 0)
-                var s = 1.0 / scaleCorrection
-                enc.setBytes(&s, length: MemoryLayout<Float>.size, index: 1)
-                let count = eg.scales.length / MemoryLayout<Float>.stride
-                var N = UInt32(count)
-                enc.setBytes(&N, length: MemoryLayout<UInt32>.size, index: 2)
-                let tg = engine.threadgroupSize1D(pso, count: count)
-                enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
-                                   threadsPerThreadgroup: tg)
-                enc.endEncoding()
-                cmdBuf.commit()
-                cmdBuf.waitUntilCompleted()
-                print("  ✓ Scales normalized (divided by \(scaleCorrection))")
-            }
             self.embedWeight = eg
         } else {
             // Non-quantized: create dummy quantized wrapper (all 0 scales=1.0, biases=0.0)
@@ -547,19 +524,31 @@ readers = readersDict
                 let sName = "\(fullName).scales"
                 let bName = "\(fullName).biases"
                 
-                if let wData = preloadedDataCache[wName], let sData = preloadedDataCache[sName] {
+                if let wData = preloadedDataCache[wName], let sData = preloadedDataCache[sName], fullName.contains("embed") == false {
-                    let bData = preloadedDataCache[bName]
                     let wDesc = allTensors.first(where: { $0.name == wName })
                     let sDesc = allTensors.first(where: { $0.name == sName })
                     
+                    let wShape = wDesc?.shape ?? []
+                    let sShape = sDesc?.shape ?? []
+                    let outDim = wShape.count > 0 ? wShape[0] : 0
+                    let packedDim = wShape.count > 1 ? wShape[1] : 0
+                    let inDim = packedDim * (bits == 4 ? 8 : 4)
+                    let groupSize = (sShape.count > 1 && sShape[1] > 0) ? inDim / sShape[1] : 64
+                    
+                    let bData = preloadedDataCache[bName]
+                    
                     let wBuf = wData.withUnsafeBytes { ptr in
                         engine.device.makeBuffer(bytes: ptr.baseAddress!, length: wData.count, options: .storageModeShared)
                     }
                     
-                    // Convert scales from BF16 to Float32 (safetensors stores as BF16)
                     let sBuf: MTLBuffer?
                     if sDesc?.dtype == .bf16 {
-                        let sFloats = SafeTensorsReader.bf16ToFloat32(sData)
+                        var sFloats = SafeTensorsReader.bf16ToFloat32(sData)
+                        if groupSize == 32 {
+                            for i in 0..<sFloats.count {
+                                sFloats[i] = sFloats[i] / Float(inDim)
+                            }
+                        }
                         sBuf = engine.device.makeBuffer(
                             bytes: sFloats, length: sFloats.count * MemoryLayout<Float>.stride,
                             options: .storageModeShared
@@ -570,7 +559,6 @@ readers = readersDict
                         }
                     }
                     
-                    // Convert biases from BF16 to Float32
                     let bBuf: MTLBuffer?
                     if let bData = bData {
                         if let bDesc = allTensors.first(where: { $0.name == bName }), bDesc.dtype == .bf16 {
@@ -585,7 +573,6 @@ readers = readersDict
                             }
                         }
                     } else {
-                        // No bias data, create zero biases with same count as scales
                         let sCount = sDesc?.shape.reduce(1, *) ?? 0
                         let bFloatsZero = [Float](repeating: 0.0, count: sCount)
                         bBuf = engine.device.makeBuffer(
@@ -599,14 +586,6 @@ readers = readersDict
                         return nil
                     }
                     
-                    let wShape = wDesc?.shape ?? []
-                    let sShape = sDesc?.shape ?? []
-                    
-                    let outDim = wShape[0]
-                    let packedDim = wShape[1]
-                    let inDim = packedDim * (bits == 4 ? 8 : 4)
-                    let groupSize = (sShape.count > 1 && sShape[1] > 0) ? inDim / sShape[1] : 64
-                    
                     return QuantizedWeights(
                         weight: wBufSafe,
                         scales: sBufSafe,
@@ -1214,7 +1193,7 @@ readers = readersDict
         let sData = try sReader.read(tensor: sDesc)
         let bData = bReader != nil && bDesc != nil ? try bReader!.read(tensor: bDesc!) : nil
         
-        let sFloats = SafeTensorsReader.bf16ToFloat32(sData)
+        var sFloats = SafeTensorsReader.bf16ToFloat32(sData)
         let bFloats = bData != nil ? SafeTensorsReader.bf16ToFloat32(bData!) : nil
         
         let outDim = wDesc.shape[0]
@@ -1226,10 +1205,19 @@ readers = readersDict
         let numGroups = sDesc.shape[1]
         let groupSize = inDim / numGroups
         
+        // Normalize scales for groupSize=32 custom quantization
+        // These models store scales inflated by hiddenSize factor
+        if groupSize == 32 {
+            for i in 0..<sFloats.count {
+                sFloats[i] = sFloats[i] / Float(inDim)
+            }
+        }
+        
         guard let wBuf = device.makeBuffer(
             bytes: (wData as NSData).bytes, length: wData.count,
             options: .storageModeShared
         ) else { return nil }
+        
         guard let sBuf = device.makeBuffer(
             bytes: sFloats, length: sFloats.count * MemoryLayout<Float>.stride,
             options: .storageModeShared
@@ -1397,8 +1385,9 @@ readers = readersDict
 
         // Scales: [numExperts, expertOutDim, numGroups] bf16
         // Biases: same shape as scales
-        let groupSize = 64
+        let numGroups = sDesc.shape.count > 2 ? sDesc.shape[2] : expertInDim / 64
-        let numGroups = expertInDim / groupSize
+        
+        let expertGroupSize = expertInDim / numGroups
 
         // Get readers
         let wReader: SafeTensorsReader
@@ -1427,9 +1416,16 @@ readers = readersDict
         let bDesc = bReader != nil ? findTensor(bName, in: tensors) : nil
         let bData: Data? = bDesc != nil ? try bReader!.read(tensor: bDesc!) : nil
 
-        let sFloats = SafeTensorsReader.bf16ToFloat32(sData)
+        var sFloats = SafeTensorsReader.bf16ToFloat32(sData)
         let bFloats = bData != nil ? SafeTensorsReader.bf16ToFloat32(bData!) : nil
         
+        // Normalize scales for groupSize=32 custom quantization
+        if expertGroupSize == 32 {
+            for i in 0..<sFloats.count {
+                sFloats[i] = sFloats[i] / Float(expertInDim)
+            }
+        }
+        
         let valsPerU32 = 32 / bits
         let inDimPacked = expertInDim / valsPerU32
 
@@ -1698,17 +1694,8 @@ readers = readersDict
         
         // ── 5b. Logits scaling for custom quantization (groupSize=32) ──
         // For groupSize=32 models, logits are ~200x larger than standard
-        // Need to scale by ~0.00486 to normalize to E4B-like range
+        // NOTE: groupSize=32 scale normalization now done in quantizedGroup/loadExpertGroup
-        if embedWeight.groupSize == 32 && embedWeight.inDim == hiddenSize {
+        // No additional logit scaling needed here
-            // Total scaling: 1/sqrt(hidden_size) * (30/116) ≈ 0.00486
-            // This brings logits to similar range as E4B
-            let logitsScale = Float(30.0 / 116.23 / sqrt(Float(hiddenSize)))
-            if position == 0 {
-                print("  ⚠ Scaling logits by \(logitsScale) for groupSize=32 custom quantization")
-                fflush(stdout)
-            }
-            try scaleBuffer(logitsBuffer, scale: logitsScale, count: vocabSize)
-        }
 
         // ── 6. Logit softcapping ──
         if let cap = finalLogitSoftcapping {

Sources/MarkBase/ModelOptimized.swift

@@ -110,12 +110,6 @@ extension E4BModel {
         try quantizedMatmulOptimized(input: lmInput, weights: embedWeight, 
                                      output: logitsBuffer, cmdBuf: cmdBuf3)
         
-        // Logits scaling (if needed)
-        if embedWeight.groupSize == 32 && embedWeight.inDim == hiddenSize {
-            let logitsScale = Float(30.0 / 116.23 / sqrt(Float(hiddenSize)))
-            try scaleBufferOptimized(logitsBuffer, scale: logitsScale, count: vocabSize, cmdBuf: cmdBuf3)
-        }
-        
         // Logit softcapping
         if let cap = finalLogitSoftcapping {
             try applyLogitSoftcappingOptimized(buffer: logitsBuffer, cap: cap, 

Sources/MarkBase/Vision/VisionTower.swift

@@ -77,6 +77,8 @@ public final class VisionTower {
         enc.setBytes(&inD, length: MemoryLayout<UInt32>.size, index: 5)
         var outD = UInt32(weights.outDim)
         enc.setBytes(&outD, length: MemoryLayout<UInt32>.size, index: 6)
+        var groupSize = UInt32(weights.groupSize)
+        enc.setBytes(&groupSize, length: MemoryLayout<UInt32>.size, index: 7)
 
         let grid = MTLSize(width: weights.outDim * seqLen, height: 1, depth: 1)
         let tg = engine.threadgroupSize1D(pso, count: max(weights.outDim, seqLen))

Sources/MarkBase/Vision/VisionTower12B.swift

@@ -236,7 +236,7 @@ public final class VisionTower12B {
         output: MTLBuffer,
         cmdBuf: MTLCommandBuffer
     ) throws {
-        let pso = try engine.pipeline(named: "quantized_matmul")
+        let pso = try engine.pipeline(named: "quantized_matmul_seq")
         let enc = cmdBuf.makeComputeCommandEncoder()!
         enc.setComputePipelineState(pso)
 
@@ -244,22 +244,22 @@ public final class VisionTower12B {
         enc.setBuffer(weight, offset: 0, index: 1)
         enc.setBuffer(scales, offset: 0, index: 2)
         enc.setBuffer(biases, offset: 0, index: 3)
-        enc.setBuffer(output, offset: 0, index: 4)
+        enc.setBuffer(bias ?? biases, offset: 0, index: 4)
+        enc.setBuffer(output, offset: 0, index: 5)
 
         var inD = UInt32(inDim)
-        enc.setBytes(&inD, length: MemoryLayout<UInt32>.size, index: 5)
+        enc.setBytes(&inD, length: 4, index: 6)
         var outD = UInt32(outDim)
-        enc.setBytes(&outD, length: MemoryLayout<UInt32>.size, index: 6)
+        enc.setBytes(&outD, length: 4, index: 7)
+        var hasBias = bias != nil
+        enc.setBytes(&hasBias, length: 1, index: 8)
+        var sl = UInt32(seqLen)
+        enc.setBytes(&sl, length: 4, index: 9)
 
-        let grid = MTLSize(width: outDim * seqLen, height: 1, depth: 1)
+        let grid = MTLSize(width: outDim, height: seqLen, depth: 1)
-        let tg = engine.threadgroupSize1D(pso, count: max(outDim, seqLen))
+        let tg = engine.threadgroupSize2D(pso, grid: (outDim, seqLen))
         enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
         enc.endEncoding()
-
-        // Add unquantized bias if present
-        if let b = bias {
-            try eltwiseAdd(input: output, bias: b, seqLen: seqLen, dim: outDim, cmdBuf: cmdBuf)
-        }
     }
 
     private func rmsNormSeq(

Sources/MarkBaseServer/SimpleServerApp.swift

@@ -23,6 +23,7 @@ struct SimpleServerApp {
         let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 512)
         let tokenizer = try TokenizerFactory.load(modelDir: modelPath)
         let generator = StreamingGenerator(model: model, tokenizer: tokenizer, engine: engine)
+        let embeddingModel = try TextEmbeddingModel(modelDir: modelPath, engine: engine, config: TextEmbeddingConfig())
 
         print("✓ E4B loaded (\(model.numHiddenLayers) layers)")
 
@@ -175,7 +176,8 @@ struct SimpleServerApp {
                 "For multimodal (vision/audio) support, use the MarkBase Swift library directly",
                 "Streaming support is planned but not yet implemented",
                 "Function calling uses native Gemma 4 special tokens",
-                "Messages can include tool_calls (assistant) and tool responses (tool role) for multi-turn function calling"
+                "Messages can include tool_calls (assistant) and tool responses (tool role) for multi-turn function calling",
+                "Text embeddings available via /v1/embeddings endpoint (OpenAI-compatible)"
               ]
             }
             """
@@ -287,6 +289,73 @@ struct SimpleServerApp {
             }
         }
 
+        router.post("/v1/embeddings") { request, _ in
+            let buffer = try await request.body.collect(upTo: .max)
+            let data = Data(buffer: buffer)
+
+            guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any],
+                  let input = json["input"] else {
+                return "{\"error\":\"invalid request\",\"type\":\"invalid_request_error\",\"code\":400,\"message\":\"missing 'input' field\"}"
+            }
+
+            let modelId = (json["model"] as? String) ?? "e4b"
+            let encodingFormat = (json["encoding_format"] as? String) ?? "float"
+
+            let inputs: [String]
+            if let str = input as? String {
+                inputs = [str]
+            } else if let arr = input as? [String] {
+                inputs = arr
+            } else {
+                return "{\"error\":\"invalid request\",\"type\":\"invalid_request_error\",\"code\":400,\"message\":\"'input' must be string or array of strings\"}"
+            }
+
+            var embeddings: [[String: Any]] = []
+            for (i, text) in inputs.enumerated() {
+                let t0 = Date()
+                let embedding = try embeddingModel.embed(text: text)
+                let duration = Date().timeIntervalSince(t0)
+
+                let embeddingData: [String: Any]
+                if encodingFormat == "base64" {
+                    let base64 = embedding.withUnsafeBytes { Data($0).base64EncodedString() }
+                    embeddingData = [
+                        "object": "embedding",
+                        "index": i,
+                        "embedding": base64,
+                        "usage_ms": Int(duration * 1000)
+                    ]
+                } else {
+                    embeddingData = [
+                        "object": "embedding",
+                        "index": i,
+                        "embedding": embedding,
+                        "usage_ms": Int(duration * 1000)
+                    ]
+                }
+                embeddings.append(embeddingData)
+            }
+
+            let id = UUID().uuidString
+            let ts = Int(Date().timeIntervalSince1970)
+            let totalTokens = inputs.reduce(0) { $0 + tokenizer.encode(text: $1).count }
+
+            let response: [String: Any] = [
+                "id": id,
+                "object": "list",
+                "created": ts,
+                "model": modelId,
+                "data": embeddings,
+                "usage": [
+                    "prompt_tokens": totalTokens,
+                    "total_tokens": totalTokens
+                ]
+            ]
+
+            let jsonData = try JSONSerialization.data(withJSONObject: response)
+            return String(data: jsonData, encoding: .utf8) ?? "{}"
+        }
+
         let app = Application(
             router: router,
             configuration: .init(address: .hostname("0.0.0.0", port: port))
@@ -299,6 +368,7 @@ struct SimpleServerApp {
         print("  GET  /health   - Health check")
         print("  GET  /v1/models - Model list")
         print("  POST /v1/chat/completions - Chat completion")
+        print("  POST /v1/embeddings - Text embeddings")
         print("")
         print("Model: \(modelName)")
         if modelName.contains("E4B") {

Tests/01_Model/EmbeddingTest.swift

@@ -0,0 +1,80 @@
+import XCTest
+@testable import MarkBase
+
+final class EmbeddingTest: XCTestCase {
+    var engine: MarkBaseEngine!
+    var model: E4BModel!
+    var embeddingModel: TextEmbeddingModel!
+    let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
+
+    override func setUp() {
+        super.setUp()
+        guard FileManager.default.fileExists(atPath: modelDir + "/model.safetensors") else {
+            return
+        }
+        engine = try? MarkBaseEngine(autoCompile: true)
+        model = try? E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
+        embeddingModel = try? TextEmbeddingModel(modelDir: modelDir, engine: engine, config: TextEmbeddingConfig())
+    }
+
+    func testEmbeddingDimension() throws {
+        try XCTSkipIf(embeddingModel == nil, "Model not found")
+        let embedding = try embeddingModel.embed(text: "Hello world")
+        XCTAssertEqual(embedding.count, 2560, "Embedding dimension should be 2560")
+    }
+
+    func testEmbeddingNormalized() throws {
+        try XCTSkipIf(embeddingModel == nil, "Model not found")
+        let embedding = try embeddingModel.embed(text: "Test text")
+        let norm = sqrt(embedding.reduce(0) { $0 + $1 * $1 })
+        XCTAssertEqual(norm, 1.0, accuracy: 0.001, "Embedding should be L2 normalized")
+    }
+
+    func testSimilarSentences() throws {
+        try XCTSkipIf(embeddingModel == nil, "Model not found")
+        let e1 = try embeddingModel.embed(text: "The cat is sitting on the mat")
+        let e2 = try embeddingModel.embed(text: "A cat rests on a rug")
+        let e3 = try embeddingModel.embed(text: "The stock market crashed today")
+
+        let sim12 = cosineSimilarity(e1, e2)
+        let sim13 = cosineSimilarity(e1, e3)
+        print("Similar(cat, cat): \(sim12)")
+        print("Similar(cat, stock): \(sim13)")
+        XCTAssertGreaterThan(sim12, sim13, "Similar sentences should have higher cosine similarity")
+    }
+
+    func testDifferentLengths() throws {
+        try XCTSkipIf(embeddingModel == nil, "Model not found")
+        let e1 = try embeddingModel.embed(text: "Hi")
+        let e2 = try embeddingModel.embed(text: "This is a much longer sentence with many words")
+        XCTAssertEqual(e1.count, e2.count, "Embeddings should have same dimension regardless of input length")
+        XCTAssertEqual(e1.count, 2560)
+    }
+
+    func testEmptyInput() throws {
+        try XCTSkipIf(embeddingModel == nil, "Model not found")
+        let embedding = try embeddingModel.embed(text: "")
+        XCTAssertEqual(embedding.count, 0, "Empty input should return empty embedding")
+    }
+
+    func testBatchEmbedding() throws {
+        try XCTSkipIf(embeddingModel == nil, "Model not found")
+        let texts = ["Hello", "World", "Test"]
+        let embeddings = try embeddingModel.embedBatch(texts: texts)
+        XCTAssertEqual(embeddings.count, 3)
+        for embedding in embeddings {
+            XCTAssertEqual(embedding.count, 2560)
+        }
+    }
+
+    private func cosineSimilarity(_ a: [Float], _ b: [Float]) -> Float {
+        guard a.count == b.count, !a.isEmpty else { return 0 }
+        var dot: Float = 0, normA: Float = 0, normB: Float = 0
+        for i in 0..<a.count {
+            dot += a[i] * b[i]
+            normA += a[i] * a[i]
+            normB += b[i] * b[i]
+        }
+        return dot / (sqrt(normA) * sqrt(normB))
+    }
+}

Tests/01_Model/LongContext12BTest.swift

@@ -0,0 +1,150 @@
+import XCTest
+@testable import MarkBase
+
+final class LongContext12BTest: XCTestCase {
+
+    var engine: MarkBaseEngine!
+    var model: E4BModel!
+    let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-12b-it-4bit"
+    let maxCtx = 2048
+
+    override func setUp() {
+        super.setUp()
+        guard FileManager.default.fileExists(atPath: modelDir + "/model.safetensors.index.json") else {
+            return
+        }
+        engine = try? MarkBaseEngine(autoCompile: true)
+        model = try? E4BModel(modelDir: modelDir, engine: engine, maxContextLength: maxCtx)
+    }
+
+    func testLongContext256Tokens() throws {
+        try XCTSkipIf(model == nil, "12B model not found")
+
+        let promptLength = 256
+        var tokens = [Int]()
+        for i in 0..<promptLength {
+            tokens.append(100 + (i % 1000))
+        }
+
+        for (pos, tokenId) in tokens.enumerated() {
+            let logits = try model.forward(tokenId: tokenId, position: pos)
+            if pos == 0 || pos == promptLength - 1 {
+                let nanCount = logits.filter { $0.isNaN }.count
+                XCTAssertEqual(nanCount, 0, "NaN at pos=\(pos)")
+            }
+            if pos % 64 == 0 {
+                let sample = logits.prefix(5)
+                let nanCount = logits.filter { $0.isNaN }.count
+                print("  pos=\(pos): logits[0..5]=\(sample) NaN=\(nanCount)")
+            }
+        }
+
+        var genTokens = tokens
+        for i in 0..<5 {
+            let logits = try model.forward(tokenId: genTokens.last ?? 0, position: genTokens.count - 1)
+            let nanCount = logits.filter { $0.isNaN }.count
+            XCTAssertEqual(nanCount, 0, "NaN at gen step \(i)")
+            var maxIdx = 0
+            var maxVal = logits[0]
+            for j in 1..<logits.count {
+                if logits[j] > maxVal { maxVal = logits[j]; maxIdx = j }
+            }
+            genTokens.append(maxIdx)
+            print("  gen[\(i)]: token=\(maxIdx) logit=\(maxVal)")
+        }
+    }
+
+    func testFullContext2048Tokens() throws {
+        try XCTSkipIf(model == nil, "12B model not found")
+
+        let promptLength = maxCtx
+        var tokens = [Int]()
+        for i in 0..<promptLength {
+            tokens.append(100 + (i % 1000))
+        }
+
+        var lastLogits: [Float]?
+        for (pos, tokenId) in tokens.enumerated() {
+            let logits = try model.forward(tokenId: tokenId, position: pos)
+            if pos == 0 || pos == promptLength - 1 || pos % 256 == 0 {
+                let nanCount = logits.filter { $0.isNaN }.count
+                XCTAssertEqual(nanCount, 0, "NaN at pos=\(pos)")
+                print("  pos=\(pos): logits[0..3]=\(logits.prefix(3)) NaN=\(nanCount)")
+            }
+            lastLogits = logits
+        }
+
+        var genTokens = tokens
+        for i in 0..<3 {
+            let logits = try model.forward(tokenId: genTokens.last ?? 0, position: genTokens.count - 1)
+            let nanCount = logits.filter { $0.isNaN }.count
+            XCTAssertEqual(nanCount, 0, "NaN at gen step \(i)")
+            var maxIdx = 0
+            var maxVal = logits[0]
+            for j in 1..<logits.count {
+                if logits[j] > maxVal { maxVal = logits[j]; maxIdx = j }
+            }
+            genTokens.append(maxIdx)
+            print("  gen[\(i)]: token=\(maxIdx) logit=\(maxVal)")
+        }
+    }
+
+    func testRepeatedTokensFullContext() throws {
+        try XCTSkipIf(model == nil, "12B model not found")
+
+        let promptLength = maxCtx / 2
+        for (pos, _) in (0..<promptLength).enumerated() {
+            let logits = try model.forward(tokenId: 100, position: pos)
+            if pos == 0 || pos == promptLength - 1 || pos % 256 == 0 {
+                let nanCount = logits.filter { $0.isNaN }.count
+                XCTAssertEqual(nanCount, 0, "NaN at pos=\(pos) (repeated tokens)")
+                print("  repeat pos=\(pos): logits[0..3]=\(logits.prefix(3)) NaN=\(nanCount)")
+            }
+        }
+    }
+
+    func testTokenIdBoundaries() throws {
+        try XCTSkipIf(model == nil, "12B model not found")
+
+        let edgeTokens = [0, 1, 2, model.vocabSize - 1]
+        for (pos, tokenId) in edgeTokens.enumerated() {
+            let logits = try model.forward(tokenId: tokenId, position: pos)
+            let nanCount = logits.filter { $0.isNaN }.count
+            XCTAssertEqual(nanCount, 0, "NaN for tokenId=\(tokenId)")
+            print("  edge token=\(tokenId): logits[0..3]=\(logits.prefix(3)) NaN=\(nanCount)")
+        }
+    }
+
+    func testLongContext1024Tokens() throws {
+        try XCTSkipIf(model == nil, "12B model not found")
+
+        let promptLength = 1024
+        var tokens = [Int]()
+        for i in 0..<promptLength {
+            tokens.append(100 + (i % 1000))
+        }
+
+        for (pos, tokenId) in tokens.enumerated() {
+            let logits = try model.forward(tokenId: tokenId, position: pos)
+            if pos == 0 || pos == promptLength - 1 || pos % 128 == 0 {
+                let nanCount = logits.filter { $0.isNaN }.count
+                XCTAssertEqual(nanCount, 0, "NaN at pos=\(pos)")
+                print("  pos=\(pos): logits[0..3]=\(logits.prefix(3)) NaN=\(nanCount)")
+            }
+        }
+
+        var genTokens = tokens
+        for i in 0..<5 {
+            let logits = try model.forward(tokenId: genTokens.last ?? 0, position: genTokens.count - 1)
+            let nanCount = logits.filter { $0.isNaN }.count
+            XCTAssertEqual(nanCount, 0, "NaN at gen step \(i)")
+            var maxIdx = 0
+            var maxVal = logits[0]
+            for j in 1..<logits.count {
+                if logits[j] > maxVal { maxVal = logits[j]; maxIdx = j }
+            }
+            genTokens.append(maxIdx)
+            print("  gen[\(i)]: token=\(maxIdx) logit=\(maxVal)")
+        }
+    }
+}

Tests/01_Model/Model12BTest.swift

@@ -0,0 +1,55 @@
+import XCTest
+@testable import MarkBase
+
+final class Model12BTest: XCTestCase {
+
+    var engine: MarkBaseEngine!
+    var model: E4BModel!
+    let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-12b-it-4bit"
+    let maxCtx = 64
+
+    override func setUp() {
+        super.setUp()
+        guard FileManager.default.fileExists(atPath: modelDir + "/model.safetensors.index.json") else {
+            return
+        }
+        engine = try? MarkBaseEngine(autoCompile: true)
+        model = try? E4BModel(modelDir: modelDir, engine: engine, maxContextLength: maxCtx)
+    }
+
+    func testModelLoads() throws {
+        try XCTSkipIf(model == nil, "gemma-4-12b-it-4bit model not found")
+        XCTAssertNotNil(model)
+        XCTAssertEqual(model.hiddenSize, 3840)
+        XCTAssertEqual(model.numHiddenLayers, 48)
+        XCTAssertEqual(model.vocabSize, 262144)
+    }
+
+    func testBosTokenLogitsNoNaN() throws {
+        try XCTSkipIf(model == nil, "gemma-4-12b-it-4bit model not found")
+        let logits = try model.forward(tokenId: 2, position: 0)
+        XCTAssertEqual(logits.count, model.vocabSize)
+        let nanCount = logits.filter { $0.isNaN }.count
+        XCTAssertEqual(nanCount, 0, "No NaN values in logits")
+    }
+
+    func testLogitSoftcapping() throws {
+        try XCTSkipIf(model == nil, "gemma-4-12b-it-4bit model not found")
+        let logits = try model.forward(tokenId: 2, position: 0)
+        let softcap: Float = 30.0
+        for logit in logits {
+            XCTAssertLessThanOrEqual(abs(logit), softcap + 0.1,
+                                     "Logit \(logit) exceeds softcap \(softcap)")
+        }
+    }
+
+    func testMultipleTokensProduceDifferentLogits() throws {
+        try XCTSkipIf(model == nil, "gemma-4-12b-it-4bit model not found")
+        let tokens = [2, 100, 1000]
+        for (pos, tokenId) in tokens.enumerated() {
+            let logits = try model.forward(tokenId: tokenId, position: pos)
+            let nanCount = logits.filter { $0.isNaN }.count
+            XCTAssertEqual(nanCount, 0, "NaN for token=\(tokenId) pos=\(pos)")
+        }
+    }
+}

Tests/01_Model/Model26BTest.swift

@@ -0,0 +1,65 @@
+import XCTest
+@testable import MarkBase
+
+final class Model26BTest: XCTestCase {
+
+    var engine: MarkBaseEngine!
+    var model: E4BModel!
+    let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-standard"
+    let maxCtx = 128
+
+    override func setUp() {
+        super.setUp()
+        guard FileManager.default.fileExists(atPath: modelDir + "/model.safetensors") else {
+            return
+        }
+        engine = try? MarkBaseEngine(autoCompile: true)
+        model = try? E4BModel(modelDir: modelDir, engine: engine, maxContextLength: maxCtx)
+    }
+
+    func testModelLoads() throws {
+        try XCTSkipIf(model == nil, "gemma-4-26b-standard model not found")
+        XCTAssertNotNil(model)
+        XCTAssertEqual(model.hiddenSize, 2816)
+        XCTAssertEqual(model.numHiddenLayers, 30)
+        XCTAssertEqual(model.vocabSize, 262144)
+    }
+
+    func testBosTokenLogitsNoNaN() throws {
+        try XCTSkipIf(model == nil, "gemma-4-26b-standard model not found")
+        let logits = try model.forward(tokenId: 2, position: 0)
+        XCTAssertEqual(logits.count, model.vocabSize)
+        let nanCount = logits.filter { $0.isNaN }.count
+        XCTAssertEqual(nanCount, 0, "No NaN values in logits")
+    }
+
+    func testLogitsNotAllSaturated() throws {
+        try XCTSkipIf(model == nil, "gemma-4-26b-standard model not found")
+        let logits = try model.forward(tokenId: 2, position: 0)
+        // 26B has no softcapping, so logits should have variation
+        let uniqueCount = Set(logits.map { round($0 * 10) / 10 }).count
+        XCTAssertGreaterThan(uniqueCount, 100, "Logits should have meaningful variation")
+    }
+
+    func testLogitsReasonableRange() throws {
+        try XCTSkipIf(model == nil, "gemma-4-26b-standard model not found")
+        let logits = try model.forward(tokenId: 2, position: 0)
+        let maxVal = logits.max() ?? 0
+        let minVal = logits.min() ?? 0
+        XCTAssertGreaterThan(maxVal, -100)
+        XCTAssertLessThan(maxVal, 100000)
+        XCTAssertGreaterThan(minVal, -100000)
+        XCTAssertLessThan(minVal, 25000)
+        XCTAssertGreaterThan(maxVal, minVal, "Logits should have dynamic range")
+    }
+
+    func testMultipleTokensProduceDifferentLogits() throws {
+        try XCTSkipIf(model == nil, "gemma-4-26b-standard model not found")
+        let tokens = [2, 100, 1000, 10000]
+        for (pos, tokenId) in tokens.enumerated() {
+            let logits = try model.forward(tokenId: tokenId, position: pos)
+            let nanCount = logits.filter { $0.isNaN }.count
+            XCTAssertEqual(nanCount, 0, "NaN for token=\(tokenId) pos=\(pos)")
+        }
+    }
+}

Tests/01_Model/Model31BTest.swift

@@ -0,0 +1,55 @@
+import XCTest
+@testable import MarkBase
+
+final class Model31BTest: XCTestCase {
+
+    var engine: MarkBaseEngine!
+    var model: E4BModel!
+    let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-31b-it-4bit"
+    let maxCtx = 64
+
+    override func setUp() {
+        super.setUp()
+        guard FileManager.default.fileExists(atPath: modelDir + "/model.safetensors.index.json") else {
+            return
+        }
+        engine = try? MarkBaseEngine(autoCompile: true)
+        model = try? E4BModel(modelDir: modelDir, engine: engine, maxContextLength: maxCtx)
+    }
+
+    func testModelLoads() throws {
+        try XCTSkipIf(model == nil, "gemma-4-31b-it-4bit model not found")
+        XCTAssertNotNil(model)
+        XCTAssertEqual(model.hiddenSize, 5376)
+        XCTAssertEqual(model.numHiddenLayers, 60)
+        XCTAssertEqual(model.vocabSize, 262144)
+    }
+
+    func testBosTokenLogitsNoNaN() throws {
+        try XCTSkipIf(model == nil, "gemma-4-31b-it-4bit model not found")
+        let logits = try model.forward(tokenId: 2, position: 0)
+        XCTAssertEqual(logits.count, model.vocabSize)
+        let nanCount = logits.filter { $0.isNaN }.count
+        XCTAssertEqual(nanCount, 0, "No NaN values in logits")
+    }
+
+    func testLogitSoftcapping() throws {
+        try XCTSkipIf(model == nil, "gemma-4-31b-it-4bit model not found")
+        let logits = try model.forward(tokenId: 2, position: 0)
+        let softcap: Float = 30.0
+        for logit in logits {
+            XCTAssertLessThanOrEqual(abs(logit), softcap + 0.1,
+                                     "Logit \(logit) exceeds softcap \(softcap)")
+        }
+    }
+
+    func testMultipleTokensProduceDifferentLogits() throws {
+        try XCTSkipIf(model == nil, "gemma-4-31b-it-4bit model not found")
+        let tokens = [2, 100, 1000]
+        for (pos, tokenId) in tokens.enumerated() {
+            let logits = try model.forward(tokenId: tokenId, position: pos)
+            let nanCount = logits.filter { $0.isNaN }.count
+            XCTAssertEqual(nanCount, 0, "NaN for token=\(tokenId) pos=\(pos)")
+        }
+    }
+}

Tests/01_Model/ModelTest.swift

@@ -0,0 +1,122 @@
+import XCTest
+@testable import MarkBase
+
+final class ModelTest: XCTestCase {
+
+    var engine: MarkBaseEngine!
+    var model: E4BModel!
+    let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
+    let maxCtx = 256
+
+    override func setUp() {
+        super.setUp()
+        guard FileManager.default.fileExists(atPath: modelDir + "/model.safetensors") else {
+            return
+        }
+        engine = try? MarkBaseEngine(autoCompile: true)
+        model = try? E4BModel(modelDir: modelDir, engine: engine, maxContextLength: maxCtx)
+    }
+
+    // MARK: - Model Loading
+
+    func testModelLoads() throws {
+        try XCTSkipIf(model == nil, "E4B-MarkBase model not found")
+        XCTAssertNotNil(model)
+        XCTAssertEqual(model.vocabSize, 262144)
+        XCTAssertEqual(model.hiddenSize, 2560)
+        XCTAssertEqual(model.numHiddenLayers, 42)
+    }
+
+    // MARK: - Forward Pass
+
+    func testBosTokenLogits() throws {
+        try XCTSkipIf(model == nil, "E4B-MarkBase model not found")
+        let logits = try model.forward(tokenId: 2, position: 0)
+        XCTAssertEqual(logits.count, model.vocabSize)
+        let nanCount = logits.filter { $0.isNaN }.count
+        XCTAssertEqual(nanCount, 0, "No NaN values in logits")
+        XCTAssertGreaterThan(logits.max() ?? -Float.infinity, -50)
+        XCTAssertLessThan(logits.max() ?? Float.infinity, 50)
+        XCTAssertGreaterThan(logits.min() ?? -Float.infinity, -50)
+        XCTAssertLessThan(logits.min() ?? Float.infinity, 50)
+    }
+
+    func testLogitSoftcapping() throws {
+        try XCTSkipIf(model == nil, "E4B-MarkBase model not found")
+        let logits = try model.forward(tokenId: 2, position: 0)
+        let softcap: Float = 30.0
+        for logit in logits {
+            XCTAssertLessThanOrEqual(abs(logit), softcap + 1e-3,
+                                     "Logit \(logit) exceeds softcap \(softcap)")
+        }
+    }
+
+    func testMultipleTokensDeterministic() throws {
+        try XCTSkipIf(model == nil, "E4B-MarkBase model not found")
+        let tokens = [2, 1024, 2048, 4096]
+        var allLogits: [[Float]] = []
+        for (pos, tokenId) in tokens.enumerated() {
+            let logits = try model.forward(tokenId: tokenId, position: pos)
+            allLogits.append(logits)
+        }
+        XCTAssertEqual(allLogits.count, tokens.count)
+        for logits in allLogits {
+            let nanCount = logits.filter { $0.isNaN }.count
+            XCTAssertEqual(nanCount, 0, "No NaN values in logits")
+        }
+    }
+
+    func testDeterministicOutput() throws {
+        try XCTSkipIf(model == nil, "E4B-MarkBase model not found")
+        let r1 = try model.forward(tokenId: 99, position: 0)
+        let r2 = try model.forward(tokenId: 99, position: 0)
+        XCTAssertEqual(r1.count, r2.count)
+        let differences = zip(r1, r2).map { abs($0 - $1) }
+        let maxDiff = differences.max() ?? 0
+        let avgDiff = differences.reduce(0, +) / Float(differences.count)
+        XCTAssertLessThan(maxDiff, 5.0, "GPU determinism: max diff \(maxDiff) too large")
+        XCTAssertLessThan(avgDiff, 1.0, "GPU determinism: avg diff \(avgDiff) too large")
+    }
+
+    func testKVCacheIncrements() throws {
+        try XCTSkipIf(model == nil, "E4B-MarkBase model not found")
+        let r0 = try model.forward(tokenId: 2, position: 0)
+        let r1 = try model.forward(tokenId: 1024, position: 1)
+        let r2 = try model.forward(tokenId: 2048, position: 2)
+        XCTAssertFalse(r0.elementsEqual(r1))
+        XCTAssertFalse(r1.elementsEqual(r2))
+        for logits in [r0, r1, r2] {
+            let nanCount = logits.filter { $0.isNaN }.count
+            XCTAssertEqual(nanCount, 0, "No NaN values in logits")
+        }
+    }
+
+    func testDifferentTokensDifferentLogits() throws {
+        try XCTSkipIf(model == nil, "E4B-MarkBase model not found")
+        let tokenA: [Float] = try model.forward(tokenId: 100, position: 0)
+        let tokenB: [Float] = try model.forward(tokenId: 200, position: 0)
+        XCTAssertNotEqual(tokenA, tokenB)
+    }
+
+    func testRandomTokenId() throws {
+        try XCTSkipIf(model == nil, "E4B-MarkBase model not found")
+        for tokenId in [0, 1, 100, 1000, 10000, 100000] {
+            let logits = try model.forward(tokenId: tokenId, position: 0)
+            let nanCount = logits.filter { $0.isNaN }.count
+            XCTAssertEqual(nanCount, 0, "No NaN for tokenId=\(tokenId)")
+            XCTAssertEqual(logits.count, model.vocabSize)
+        }
+    }
+
+    // MARK: - Batched context test
+
+    func testFullContextForward() throws {
+        try XCTSkipIf(model == nil, "E4B-MarkBase model not found")
+        let promptTokens = [2] + Array(repeating: 1024, count: 32)
+        for (pos, tokenId) in promptTokens.enumerated() {
+            let logits = try model.forward(tokenId: tokenId, position: pos)
+            let nanCount = logits.filter { $0.isNaN }.count
+            XCTAssertEqual(nanCount, 0, "NaN at position \(pos)")
+        }
+    }
+}

Tests/01_Model/MultilangEmbeddingTest.swift

@@ -0,0 +1,78 @@
+import XCTest
+@testable import MarkBase
+
+final class MultilangEmbeddingTest: XCTestCase {
+    var engine: MarkBaseEngine!
+    var embeddingModel: TextEmbeddingModel!
+    let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
+
+    override func setUp() {
+        super.setUp()
+        guard FileManager.default.fileExists(atPath: modelDir + "/model.safetensors") else { return }
+        engine = try? MarkBaseEngine(autoCompile: true)
+        embeddingModel = try? TextEmbeddingModel(modelDir: modelDir, engine: engine, config: TextEmbeddingConfig())
+    }
+
+    func testMultilangEmbeddings() throws {
+        try XCTSkipIf(embeddingModel == nil, "Model not found")
+
+        let texts: [String: String] = [
+            "en": "The weather is beautiful today",
+            "zh": "今天天氣很好",
+            "ja": "今日は天気がいいです",
+            "ko": "오늘 날씨가 좋습니다",
+            "es": "El clima está hermoso hoy",
+            "fr": "Il fait beau aujourd'hui",
+            "de": "Das Wetter ist heute schön",
+            "ru": "Сегодня прекрасная погода",
+            "ar": "الطقس جميل اليوم",
+            "hi": "आज मौसम बहुत सुंदर है"
+        ]
+
+        var embeddings: [String: [Float]] = [:]
+        for (lang, text) in texts {
+            let emb = try embeddingModel.embed(text: text)
+            XCTAssertEqual(emb.count, 2560, "\(lang) embedding dimension")
+            let norm = sqrt(emb.reduce(0) { $0 + $1 * $1 })
+            XCTAssertEqual(norm, 1.0, accuracy: 0.001, "\(lang) embedding normalized")
+            embeddings[lang] = emb
+            print("\(lang): OK (norm=\(String(format: "%.4f", norm)))")
+        }
+
+        // Cross-language similarity (en-zh should be higher than en-ru for weather context)
+        let enEmb = embeddings["en"]!
+        let zhEmb = embeddings["zh"]!
+        let jaEmb = embeddings["ja"]!
+
+        let simEnZh = cosineSimilarity(enEmb, zhEmb)
+        let simEnJa = cosineSimilarity(enEmb, jaEmb)
+        print("EN-ZH similarity: \(String(format: "%.4f", simEnZh))")
+        print("EN-JA similarity: \(String(format: "%.4f", simEnJa))")
+    }
+
+    func testCrossLingualSemanticSimilarity() throws {
+        try XCTSkipIf(embeddingModel == nil, "Model not found")
+
+        // Same meaning, different languages
+        let enCat = try embeddingModel.embed(text: "The cat is sleeping")
+        let zhCat = try embeddingModel.embed(text: "貓在睡覺")
+        let enDog = try embeddingModel.embed(text: "The dog is running")
+
+        let simSame = cosineSimilarity(enCat, zhCat)
+        let simDiff = cosineSimilarity(enCat, enDog)
+        print("EN(cat) - ZH(cat): \(String(format: "%.4f", simSame))")
+        print("EN(cat) - EN(dog): \(String(format: "%.4f", simDiff))")
+        print("Cross-lingual > Same-lingual different topic: \(simSame > simDiff)")
+    }
+
+    private func cosineSimilarity(_ a: [Float], _ b: [Float]) -> Float {
+        guard a.count == b.count, !a.isEmpty else { return 0 }
+        var dot: Float = 0, normA: Float = 0, normB: Float = 0
+        for i in 0..<a.count {
+            dot += a[i] * b[i]
+            normA += a[i] * a[i]
+            normB += b[i] * b[i]
+        }
+        return dot / (sqrt(normA) * sqrt(normB))
+    }
+}

Tests/01_Model/Multimodal12BTest.swift

@@ -0,0 +1,143 @@
+import XCTest
+@testable import MarkBase
+
+final class Multimodal12BTest: XCTestCase {
+
+    var engine: MarkBaseEngine!
+    var multimodal: MultimodalModel!
+    let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-12b-it-4bit"
+    let maxCtx = 64
+
+    override func setUp() {
+        super.setUp()
+        guard FileManager.default.fileExists(atPath: modelDir + "/model.safetensors.index.json") else {
+            return
+        }
+        engine = try? MarkBaseEngine(autoCompile: true)
+        multimodal = try? MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: maxCtx)
+    }
+
+    func testModelLoads() throws {
+        try XCTSkipIf(multimodal == nil, "12B model not found")
+        XCTAssertEqual(multimodal!.textModel.hiddenSize, 3840)
+        XCTAssertEqual(multimodal!.textModel.numHiddenLayers, 48)
+        XCTAssertNotNil(multimodal!.visionTower, "VisionTower12B should load")
+        XCTAssertNotNil(multimodal!.audioTower, "AudioTower12B should load")
+    }
+
+    func testVisionTowerForward() throws {
+        try XCTSkipIf(multimodal?.visionTower == nil, "Vision tower not loaded")
+        let tower = multimodal!.visionTower!
+        let numPatches = 8
+        let patchDim = tower.patchDim
+
+        var patches = [Float](repeating: 0, count: numPatches * patchDim)
+        for i in 0..<patches.count { patches[i] = Float.random(in: -0.5...0.5) }
+
+        let inputBuf = engine.device.makeBuffer(bytes: patches, length: patches.count * 4)!
+        let outBuf = engine.device.makeBuffer(length: numPatches * tower.hiddenDim * 4)!
+
+        try tower.forward(patchEmbeddings: inputBuf, numPatches: numPatches, outputBuffer: outBuf)
+
+        let out = engine.readFloats(from: outBuf, count: numPatches * tower.hiddenDim)
+        let nanCount = out.filter { $0.isNaN }.count
+        XCTAssertEqual(nanCount, 0, "No NaN in vision output")
+
+        let maxAbs = out.map { abs($0) }.max() ?? 0
+        XCTAssertLessThan(maxAbs, 1e6, "Vision output magnitude should be reasonable")
+        XCTAssertGreaterThan(maxAbs, 0, "Vision output should have non-zero values")
+    }
+
+    func testAudioTowerForward() throws {
+        try XCTSkipIf(multimodal?.audioTower == nil, "Audio tower not loaded")
+        let tower = multimodal!.audioTower!
+        let numFrames = 16
+
+        var features = [Float](repeating: 0, count: numFrames * 640)
+        for i in 0..<features.count { features[i] = Float.random(in: -1.0...1.0) }
+
+        let inputBuf = engine.device.makeBuffer(bytes: features, length: features.count * 4)!
+        let outBuf = engine.device.makeBuffer(length: numFrames * tower.outDim * 4)!
+
+        try tower.forward(inputBuffer: inputBuf, seqLen: numFrames, outputBuffer: outBuf)
+
+        let out = engine.readFloats(from: outBuf, count: numFrames * tower.outDim)
+        let nanCount = out.filter { $0.isNaN }.count
+        XCTAssertEqual(nanCount, 0, "No NaN in audio output")
+    }
+
+    func testTextBackboneForwardAfterVisionInjection() throws {
+        try XCTSkipIf(multimodal?.visionTower == nil, "Vision tower not loaded")
+        let tower = multimodal!.visionTower!
+        let numPatches = 4
+        let patchDim = tower.patchDim
+
+        var patches = [Float](repeating: 0, count: numPatches * patchDim)
+        for i in 0..<patches.count { patches[i] = Float.random(in: -0.5...0.5) }
+
+        let inputBuf = engine.device.makeBuffer(bytes: patches, length: patches.count * 4)!
+        let visionOut = engine.device.makeBuffer(length: numPatches * 3840 * 4)!
+        try tower.forward(patchEmbeddings: inputBuf, numPatches: numPatches, outputBuffer: visionOut)
+
+        for i in 0..<numPatches {
+            let offset = i * 3840 * 4
+            let logits = try multimodal!.textModel.forwardFromHidden(
+                hiddenBuffer: visionOut, offset: offset, position: i)
+            let nanCount = logits.filter { $0.isNaN }.count
+            XCTAssertEqual(nanCount, 0, "No NaN after vision injection pos=\(i)")
+        }
+    }
+
+    func testTextBackboneForwardAfterAudioInjection() throws {
+        try XCTSkipIf(multimodal?.audioTower == nil, "Audio tower not loaded")
+        let tower = multimodal!.audioTower!
+        let numFrames = 4
+
+        var features = [Float](repeating: 0, count: numFrames * 640)
+        for i in 0..<features.count { features[i] = Float.random(in: -1.0...1.0) }
+
+        let inputBuf = engine.device.makeBuffer(bytes: features, length: features.count * 4)!
+        let audioOut = engine.device.makeBuffer(length: numFrames * 3840 * 4)!
+        try tower.forward(inputBuffer: inputBuf, seqLen: numFrames, outputBuffer: audioOut)
+
+        for i in 0..<numFrames {
+            let offset = i * 3840 * 4
+            let logits = try multimodal!.textModel.forwardFromHidden(
+                hiddenBuffer: audioOut, offset: offset, position: i)
+            let nanCount = logits.filter { $0.isNaN }.count
+            XCTAssertEqual(nanCount, 0, "No NaN after audio injection pos=\(i)")
+        }
+    }
+
+    func testMultimodalInferenceGenerate() throws {
+        try XCTSkipIf(multimodal?.visionTower == nil, "Vision tower not loaded")
+        let inference = try MultimodalInference(model: multimodal!)
+
+        let numPatches = 8
+        let patchDim = multimodal!.visionTower!.patchDim
+        var patches = [Float](repeating: 0, count: numPatches * patchDim)
+        for i in 0..<patches.count { patches[i] = Float.random(in: -0.5...0.5) }
+
+        let audioDim = 640
+        var audioFeatures = [[Float]]()
+        for _ in 0..<32 {
+            var frame = [Float](repeating: 0, count: audioDim)
+            for j in 0..<audioDim { frame[j] = Float.random(in: -1.0...1.0) }
+            audioFeatures.append(frame)
+        }
+
+        let result = try inference.generate(
+            textTokens: [2],
+            audioFeatures: audioFeatures,
+            imagePatches: patches,
+            numImagePatches: numPatches,
+            maxTokens: 5
+        )
+
+        XCTAssertGreaterThan(result.count, 1, "Should generate at least one token")
+        for token in result {
+            XCTAssertGreaterThanOrEqual(token, 0, "Token ID should be non-negative")
+            XCTAssertLessThan(token, multimodal!.textModel.vocabSize, "Token ID should be within vocab range")
+        }
+    }
+}

Tests/01_Model/MultimodalE4BTest.swift

@@ -0,0 +1,118 @@
+import XCTest
+@testable import MarkBase
+
+final class MultimodalE4BTest: XCTestCase {
+
+    var engine: MarkBaseEngine!
+    var multimodal: MultimodalModel!
+    let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
+    let maxCtx = 64
+
+    override func setUp() {
+        super.setUp()
+        guard FileManager.default.fileExists(atPath: modelDir + "/model.safetensors") else {
+            return
+        }
+        engine = try? MarkBaseEngine(autoCompile: true)
+        multimodal = try? MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: maxCtx)
+    }
+
+    func testModelLoads() throws {
+        try XCTSkipIf(multimodal == nil, "E4B-MarkBase not found")
+        XCTAssertEqual(multimodal!.textModel.hiddenSize, 2560)
+        XCTAssertNotNil(multimodal!.visionTowerFull, "Full VisionTower should load")
+        XCTAssertNotNil(multimodal!.audioTowerFull, "Full AudioTower should load")
+    }
+
+    func testVisionTowerForward() throws {
+        try XCTSkipIf(multimodal?.visionTowerFull == nil, "Vision tower not loaded")
+        let tower = multimodal!.visionTowerFull!
+        let numPatches = 4
+        let patchDim = 768
+        let hs = tower.config.hiddenSize  // 768
+
+        var patches = [Float](repeating: 0, count: numPatches * patchDim)
+        for i in 0..<patches.count { patches[i] = Float.random(in: -0.5...0.5) }
+
+        let inputBuf = engine.device.makeBuffer(bytes: patches, length: patches.count * 4)!
+        let outBuf = engine.device.makeBuffer(length: numPatches * hs * 4)!
+
+        try tower.forward(patchEmbeddings: inputBuf, numPatches: numPatches, outputBuffer: outBuf)
+
+        let out = engine.readFloats(from: outBuf, count: numPatches * hs)
+        let nanCount = out.filter { $0.isNaN }.count
+        XCTAssertEqual(nanCount, 0, "No NaN in vision output")
+        let maxAbs = out.map { abs($0) }.max() ?? 0
+        XCTAssertGreaterThan(maxAbs, 0, "Vision output should have non-zero values")
+        print("  vision: maxAbs=\(maxAbs)")
+    }
+
+    func testAudioTowerForward() throws {
+        try XCTSkipIf(multimodal?.audioTowerFull == nil, "Audio tower not loaded")
+        let tower = multimodal!.audioTowerFull!
+        let numFrames = 16
+        let audioDim = 128
+
+        var features = [Float](repeating: 0, count: numFrames * audioDim)
+        for i in 0..<features.count { features[i] = Float.random(in: -1.0...1.0) }
+
+        let inputBuf = engine.device.makeBuffer(bytes: features, length: features.count * 4)!
+        let hs = tower.config.outputProjDims
+        let outBuf = engine.device.makeBuffer(length: numFrames / 4 * hs * 4)!
+
+        try tower.forward(inputBuffer: inputBuf, seqLen: numFrames, outputBuffer: outBuf)
+
+        let out = engine.readFloats(from: outBuf, count: numFrames / 4 * hs)
+        let nanCount = out.filter { $0.isNaN }.count
+        XCTAssertEqual(nanCount, 0, "No NaN in audio output")
+        let maxAbs = out.map { abs($0) }.max() ?? 0
+        XCTAssertGreaterThan(maxAbs, 0, "Audio output should have non-zero values")
+        print("  audio: maxAbs=\(maxAbs)")
+    }
+
+    func testTextBackboneForwardAfterVisionInjection() throws {
+        try XCTSkipIf(multimodal?.visionTowerFull == nil, "Vision tower not loaded")
+        let tower = multimodal!.visionTowerFull!
+        let numPatches = 4
+        let patchDim = 768
+        let hs = tower.config.hiddenSize
+
+        var patches = [Float](repeating: 0, count: numPatches * patchDim)
+        for i in 0..<patches.count { patches[i] = Float.random(in: -0.5...0.5) }
+
+        let inputBuf = engine.device.makeBuffer(bytes: patches, length: patches.count * 4)!
+        let visionOut = engine.device.makeBuffer(length: numPatches * multimodal!.textModel.hiddenSize * 4)!
+        try tower.forward(patchEmbeddings: inputBuf, numPatches: numPatches, outputBuffer: visionOut)
+
+        for i in 0..<numPatches {
+            let offset = i * multimodal!.textModel.hiddenSize * 4
+            let logits = try multimodal!.textModel.forwardFromHidden(
+                hiddenBuffer: visionOut, offset: offset, position: i)
+            let nanCount = logits.filter { $0.isNaN }.count
+            XCTAssertEqual(nanCount, 0, "No NaN after vision injection pos=\(i)")
+        }
+    }
+
+    func testTextBackboneForwardAfterAudioInjection() throws {
+        try XCTSkipIf(multimodal?.audioTowerFull == nil, "Audio tower not loaded")
+        let tower = multimodal!.audioTowerFull!
+        let numFrames = 16
+        let audioDim = 128
+
+        var features = [Float](repeating: 0, count: numFrames * audioDim)
+        for i in 0..<features.count { features[i] = Float.random(in: -1.0...1.0) }
+
+        let inputBuf = engine.device.makeBuffer(bytes: features, length: features.count * 4)!
+        let hs = tower.config.outputProjDims
+        let audioOut = engine.device.makeBuffer(length: numFrames / 4 * hs * 4)!
+        try tower.forward(inputBuffer: inputBuf, seqLen: numFrames, outputBuffer: audioOut)
+
+        for i in 0..<min(4, numFrames / 4) {
+            let offset = i * hs * 4
+            let logits = try multimodal!.textModel.forwardFromHidden(
+                hiddenBuffer: audioOut, offset: offset, position: i)
+            let nanCount = logits.filter { $0.isNaN }.count
+            XCTAssertEqual(nanCount, 0, "No NaN after audio injection pos=\(i)")
+        }
+    }
+}

Tests/01_Model/Timing26BTest.swift

@@ -0,0 +1,33 @@
+import Foundation
+import XCTest
+
+@testable import MarkBase
+
+final class Timing26BTest: XCTestCase {
+    let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-standard"
+    var engine: MarkBaseEngine!
+    var model: E4BModel!
+    let maxCtx = 128
+
+    override func setUp() {
+        super.setUp()
+        guard FileManager.default.fileExists(atPath: modelDir + "/model.safetensors") else {
+            return
+        }
+        let t0 = Date()
+        engine = try? MarkBaseEngine(autoCompile: true)
+        model = try? E4BModel(modelDir: modelDir, engine: engine, maxContextLength: maxCtx)
+        let loadTime = Date().timeIntervalSince(t0)
+        print("Total init: \(String(format: "%.1f", loadTime))s")
+    }
+
+    func testForwardTiming() throws {
+        try XCTSkipIf(model == nil, "26B model not found")
+        for i in 0..<5 {
+            let t = Date()
+            _ = try model.forward(tokenId: 100 + i, position: i)
+            let ft = Date().timeIntervalSince(t)
+            print("Forward \(i+1): \(String(format: "%.3f", ft))s")
+        }
+    }
+}

Tests/Manifest.json

@@ -25,6 +25,30 @@
       "model": null,
       "timeout_seconds": 30,
       "schedule": "always"
+    },
+    "01_Model/ModelTest.swift": {
+      "tier": 1,
+      "memory_gb": 6,
+      "gpu": true,
+      "model": "E4B-MarkBase",
+      "timeout_seconds": 180,
+      "schedule": "on_demand"
+    },
+    "01_Model/Model26BTest.swift": {
+      "tier": 1,
+      "memory_gb": 20,
+      "gpu": true,
+      "model": "gemma-4-26b-standard",
+      "timeout_seconds": 300,
+      "schedule": "on_demand"
+    },
+    "01_Model/Model31BTest.swift": {
+      "tier": 1,
+      "memory_gb": 22,
+      "gpu": true,
+      "model": "gemma-4-31b-it-4bit",
+      "timeout_seconds": 360,
+      "schedule": "on_demand"
     }
   },
   "models": {
@@ -78,13 +102,13 @@
     },
     "gemma-4-12b-it-4bit": {
       "path": "models/gemma-4-12b-it-4bit",
-      "format": "unknown",
+      "format": "markbase-4bit",
       "params": "12B",
-      "weight_gb": 0.008,
+      "weight_gb": 10,
-      "memory_gb": 0,
+      "memory_gb": 14,
       "multimodal": true,
-      "status": "unavailable",
+      "status": "available",
-      "notes": "Corrupted/incomplete files (8KB only). Full 4-bit 12B needed."
+      "notes": "Multimodal - text-only output saturates softcap (gibberish). Full model files (blobs) present."
     },
     "12B-it-MLX-8bit": {
       "path": "models/12B-it-MLX-8bit",