Initial commit: E4B-MarkBase model integration with passing tests

- E4B-MarkBase model (42 layers, 4.4GB) loaded successfully
- All Phase 1-6 tests passed (model loading, forward pass, vision/audio towers, token generation, performance)
- All stress tests passed (5/5 in 127.6s)
  - Concurrent inference
  - Memory stress (67.5 tok/s, 0 NaN)
  - Continuous generation
  - Batch processing
  - Long-running stability
- Swift Metal inference engine with multimodal support

Sources/MarkBase/Weights/DType.swift

@@ -0,0 +1,35 @@
+/// Supported tensor data types in SafeTensors files.
+public enum TensorDType: String, Codable, Sendable {
+    case bf16 = "BF16"
+    case f32 = "F32"
+    case f16 = "F16"
+    case u32 = "U32"
+    case i32 = "I32"
+    case i64 = "I64"
+
+    /// Case-insensitive lookup: matches both "BF16" and "bfloat16", etc.
+    public static func from(dtype str: String) -> TensorDType? {
+        switch str.lowercased() {
+        case "bf16", "bfloat16": return .bf16
+        case "f32", "float32": return .f32
+        case "f16", "float16": return .f16
+        case "u32", "uint32": return .u32
+        case "i32", "int32": return .i32
+        case "i64", "int64": return .i64
+        default: return nil
+        }
+    }
+
+    /// Number of bytes per element for this dtype.
+    public var byteSize: Int {
+        switch self {
+        case .bf16, .f16: 2
+        case .f32, .u32, .i32: 4
+        case .i64: 8
+        }
+    }
+
+    /// Whether this dtype holds quantized (packed) weight data.
+    /// Quantized weights have separate scales+biases tensors.
+    public var isQuantized: Bool { self == .u32 }
+}

Sources/MarkBase/Weights/ModelConfig.swift

@@ -0,0 +1,176 @@
+import Foundation
+
+/// Model architecture configuration parsed from config.json.
+/// Uses JSONSerialization (instead of Codable) to tolerate unknown keys.
+public struct ModelConfig: Sendable {
+    public let modelType: String?
+    public let hiddenSize: Int?
+    public let intermediateSize: Int?
+    public let numAttentionHeads: Int?
+    public let numHiddenLayers: Int?
+    public let numKeyValueHeads: Int?
+    public let vocabSize: Int?
+    public let maxPositionEmbeddings: Int?
+    public let rmsNormEps: Float?
+    public let ropeTheta: Float?
+
+    // Gemma 4 specific
+    public let slidingWindow: Int?
+    public let headDim: Int?
+    public let globalHeadDim: Int?
+    public let slidingHeadDim: Int?
+    public let numKvSharedLayers: Int?
+    public let hiddenSizePerLayerInput: Int?
+    public let slidingWindowPattern: Int?
+    public let finalLogitSoftcapping: Float?
+    public let tieWordEmbeddings: Bool?
+    public let perLayerInputScale: Float?
+    public let perLayerProjectionScale: Float?
+    public let embedScale: Float?
+    /// Per-layer attention type: "full_attention" or "sliding_attention"
+    public let layerTypes: [String]?
+
+    // Global KV heads (for full attention layers)
+    public let numGlobalKeyValueHeads: Int?
+
+    // K=V sharing (Gemma 4 full attention layers)
+    public let attentionKEqualsV: Bool?
+
+    // MoE
+    public let enableMoEBlock: Bool?
+    public let numExperts: Int?
+    public let topKExperts: Int?
+    public let moeIntermediateSize: Int?
+
+    public init(
+        modelType: String? = nil,
+        hiddenSize: Int? = nil,
+        intermediateSize: Int? = nil,
+        numAttentionHeads: Int? = nil,
+        numHiddenLayers: Int? = nil,
+        numKeyValueHeads: Int? = nil,
+        vocabSize: Int? = nil,
+        maxPositionEmbeddings: Int? = nil,
+        rmsNormEps: Float? = nil,
+        ropeTheta: Float? = nil,
+        slidingWindow: Int? = nil,
+        headDim: Int? = nil,
+        globalHeadDim: Int? = nil,
+        slidingHeadDim: Int? = nil,
+        numKvSharedLayers: Int? = nil,
+        hiddenSizePerLayerInput: Int? = nil,
+        slidingWindowPattern: Int? = nil,
+        finalLogitSoftcapping: Float? = nil,
+        tieWordEmbeddings: Bool? = nil,
+        perLayerInputScale: Float? = nil,
+        perLayerProjectionScale: Float? = nil,
+        embedScale: Float? = nil,
+        layerTypes: [String]? = nil,
+        numGlobalKeyValueHeads: Int? = nil,
+        enableMoEBlock: Bool? = nil,
+        numExperts: Int? = nil,
+        topKExperts: Int? = nil,
+        moeIntermediateSize: Int? = nil,
+        attentionKEqualsV: Bool? = nil
+    ) {
+        self.modelType = modelType
+        self.hiddenSize = hiddenSize
+        self.intermediateSize = intermediateSize
+        self.numAttentionHeads = numAttentionHeads
+        self.numHiddenLayers = numHiddenLayers
+        self.numKeyValueHeads = numKeyValueHeads
+        self.vocabSize = vocabSize
+        self.maxPositionEmbeddings = maxPositionEmbeddings
+        self.rmsNormEps = rmsNormEps
+        self.ropeTheta = ropeTheta
+        self.slidingWindow = slidingWindow
+        self.headDim = headDim
+        self.globalHeadDim = globalHeadDim
+        self.slidingHeadDim = slidingHeadDim
+        self.numKvSharedLayers = numKvSharedLayers
+        self.hiddenSizePerLayerInput = hiddenSizePerLayerInput
+        self.slidingWindowPattern = slidingWindowPattern
+        self.layerTypes = layerTypes
+        self.finalLogitSoftcapping = finalLogitSoftcapping
+        self.tieWordEmbeddings = tieWordEmbeddings
+        self.perLayerInputScale = perLayerInputScale
+        self.perLayerProjectionScale = perLayerProjectionScale
+        self.embedScale = embedScale
+        self.numGlobalKeyValueHeads = numGlobalKeyValueHeads
+        self.enableMoEBlock = enableMoEBlock
+        self.numExperts = numExperts
+        self.topKExperts = topKExperts
+        self.moeIntermediateSize = moeIntermediateSize
+        self.attentionKEqualsV = attentionKEqualsV
+    }
+
+    /// Load config from a model directory (config.json).
+    /// Uses JSONSerialization to tolerate extra keys.
+    public static func load(from directory: String) throws -> ModelConfig {
+        let url = URL(fileURLWithPath: directory).appendingPathComponent("config.json")
+        let data = try Data(contentsOf: url)
+        guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] else {
+            throw WeightError.invalidHeader("config.json is not a dictionary")
+        }
+
+        // Some configs nest text params in "text_config"
+        let tc = json["text_config"] as? [String: Any] ?? [:]
+
+        return ModelConfig(
+            modelType: json.string("model_type"),
+            hiddenSize: json.int("hidden_size") ?? tc.int("hidden_size"),
+            intermediateSize: json.int("intermediate_size") ?? tc.int("intermediate_size"),
+            numAttentionHeads: json.int("num_attention_heads") ?? tc.int("num_attention_heads"),
+            numHiddenLayers: json.int("num_hidden_layers") ?? tc.int("num_hidden_layers"),
+            numKeyValueHeads: json.int("num_key_value_heads") ?? tc.int("num_key_value_heads"),
+            vocabSize: json.int("vocab_size") ?? tc.int("vocab_size"),
+            maxPositionEmbeddings: json.int("max_position_embeddings") ?? tc.int("max_position_embeddings"),
+            rmsNormEps: json.float("rms_norm_eps") ?? tc.float("rms_norm_eps"),
+            ropeTheta: json.float("rope_theta") ?? tc.float("rope_theta"),
+            slidingWindow: json.int("sliding_window") ?? tc.int("sliding_window"),
+            headDim: json.int("head_dim") ?? tc.int("head_dim"),
+            globalHeadDim: json.int("global_head_dim") ?? tc.int("global_head_dim"),
+            slidingHeadDim: json.int("sliding_head_dim") ?? tc.int("sliding_head_dim"),
+            numKvSharedLayers: json.int("num_kv_shared_layers") ?? tc.int("num_kv_shared_layers"),
+            hiddenSizePerLayerInput: json.int("hidden_size_per_layer_input") ?? tc.int("hidden_size_per_layer_input"),
+            slidingWindowPattern: json.int("sliding_window_pattern") ?? tc.int("sliding_window_pattern"),
+            finalLogitSoftcapping: json.float("final_logit_softcapping") ?? tc.float("final_logit_softcapping"),
+            tieWordEmbeddings: json.bool("tie_word_embeddings") ?? tc.bool("tie_word_embeddings"),
+            perLayerInputScale: json.float("per_layer_input_scale") ?? tc.float("per_layer_input_scale"),
+            perLayerProjectionScale: json.float("per_layer_projection_scale") ?? tc.float("per_layer_projection_scale"),
+            embedScale: json.float("embed_scale") ?? tc.float("embed_scale"),
+            layerTypes: tc.strings("layer_types"),
+            numGlobalKeyValueHeads: json.int("num_global_key_value_heads") ?? tc.int("num_global_key_value_heads"),
+            enableMoEBlock: tc.bool("enable_moe_block"),
+            numExperts: tc.int("num_experts"),
+            topKExperts: tc.int("top_k_experts"),
+            moeIntermediateSize: tc.int("moe_intermediate_size"),
+            attentionKEqualsV: json.bool("attention_k_eq_v") ?? tc.bool("attention_k_eq_v")
+        )
+    }
+}
+
+// ── JSON helpers ──────────────────────────────────────
+
+extension Dictionary where Key == String, Value == Any {
+    func string(_ key: String) -> String? { self[key] as? String }
+    func int(_ key: String) -> Int? {
+        if let v = self[key] as? Int { return v }
+        if let v = self[key] as? Double { return Int(v) }
+        if let v = self[key] as? NSNumber { return v.intValue }
+        return nil
+    }
+    func float(_ key: String) -> Float? {
+        if let v = self[key] as? Float { return v }
+        if let v = self[key] as? Double { return Float(v) }
+        if let v = self[key] as? NSNumber { return v.floatValue }
+        return nil
+    }
+    func bool(_ key: String) -> Bool? {
+        if let v = self[key] as? Bool { return v }
+        return (self[key] as? NSNumber)?.boolValue
+    }
+    func strings(_ key: String) -> [String]? {
+        self[key] as? [String]
+    }
+}

Sources/MarkBase/Weights/SafeTensors.swift

@@ -0,0 +1,126 @@
+import Foundation
+
+/// SafeTensors file reader. Handles single-file and sharded (index) formats,
+/// BF16→Float32 conversion, and quantized tensor grouping.
+public final class SafeTensorsReader {
+    public let fileURL: URL
+    private let headerSize: Int
+    private let rawHeader: [String: Any]
+    private let fileHandle: FileHandle  // kept open for fast repeated reads
+    private let lock = NSLock()  // thread-safe access to fileHandle
+
+    // ── Init ──────────────────────────────────────────
+
+    /// Open a single .safetensors file and parse its header.
+    public init(path: String) throws {
+        self.fileURL = URL(fileURLWithPath: path)
+        let handle = try FileHandle(forReadingFrom: fileURL)
+
+        let lenData = handle.readData(ofLength: 8)
+        headerSize = Int(UInt64(littleEndian: lenData.withUnsafeBytes { $0.load(as: UInt64.self) }))
+
+        let jsonData = handle.readData(ofLength: headerSize)
+        guard let json = try JSONSerialization.jsonObject(with: jsonData) as? [String: Any] else {
+            try? handle.close()
+            throw WeightError.invalidHeader("Top-level JSON is not a dictionary")
+        }
+        self.rawHeader = json
+        self.fileHandle = handle
+    }
+
+    deinit {
+        try? fileHandle.close()
+    }
+
+    // ── Tensor listing ────────────────────────────────
+
+    /// All tensor descriptors in this file.
+    public var allTensors: [TensorDescriptor] {
+        rawHeader.compactMap { name, value in
+            guard let info = value as? [String: Any],
+                  let dtypeStr = info["dtype"] as? String,
+                  let dtype = TensorDType.from(dtype: dtypeStr),
+                  let shape = info["shape"] as? [Int],
+                  let offsets = info["data_offsets"] as? [Int],
+                  offsets.count == 2
+            else { return nil }
+            return TensorDescriptor(
+                name: name, dtype: dtype, shape: shape,
+                dataOffset: headerSize + 8 + offsets[0],
+                dataSize: offsets[1] - offsets[0]
+            )
+        }
+    }
+
+    /// All tensor descriptors (convenience).
+    public func allDescriptors() -> [TensorDescriptor] { allTensors }
+
+    /// Look up a specific tensor by name.
+    public func tensor(named name: String) -> TensorDescriptor? {
+        allTensors.first { $0.name == name }
+    }
+
+    // ── Reading raw data ──────────────────────────────
+
+    /// Read raw bytes for a tensor.
+    public func read(tensor: TensorDescriptor) throws -> Data {
+        lock.lock()
+        defer { lock.unlock() }
+        try fileHandle.seek(toOffset: UInt64(tensor.dataOffset))
+        return fileHandle.readData(ofLength: tensor.dataSize)
+    }
+
+    /// Read a specific tensor by name.
+    public func read(named name: String) throws -> Data {
+        guard let desc = tensor(named: name) else {
+            throw WeightError.tensorNotFound(name)
+        }
+        return try read(tensor: desc)
+    }
+    
+    /// Read raw bytes for a tensor as uint32 array
+    public func readUint32(named name: String) throws -> [UInt32] {
+        guard let desc = tensor(named: name) else {
+            throw WeightError.tensorNotFound(name)
+        }
+        let data = try read(tensor: desc)
+        return data.withUnsafeBytes { ptr in
+            let uint32Ptr = ptr.bindMemory(to: UInt32.self)
+            return Array(uint32Ptr)
+        }
+    }
+
+    // ── BF16 → Float32 conversion ─────────────────────
+
+    /// Convert BF16 binary data to Float32 array.
+    public static func bf16ToFloat32(_ data: Data) -> [Float] {
+        data.withUnsafeBytes { ptr in
+            let bf16 = ptr.assumingMemoryBound(to: UInt16.self)
+            return (0..<data.count / 2).map { i in
+                Float(bitPattern: UInt32(bf16[i]) << 16)
+            }
+        }
+    }
+}
+
+// ── Errors ────────────────────────────────────────────
+
+public enum WeightError: Error, LocalizedError {
+    case invalidHeader(String)
+    case tensorNotFound(String)
+    case unsupportedDtype(String)
+    case fileNotFound(String)
+    case readFailed(String)
+    case bufferCreationFailed(String)
+
+    public var errorDescription: String? {
+        switch self {
+        case .invalidHeader(let detail): return "Invalid SafeTensors header: \(detail)"
+        case .tensorNotFound(let name): return "Tensor '\(name)' not found"
+        case .unsupportedDtype(let dtype): return "Unsupported dtype: \(dtype)"
+        case .fileNotFound(let path): return "File not found: \(path)"
+        case .readFailed(let detail): return "Read failed: \(detail)"
+        case .bufferCreationFailed(let name): return "Failed to create Metal buffer: \(name)"
+        }
+    }
+}

Sources/MarkBase/Weights/SafeTensorsIndex.swift

@@ -0,0 +1,34 @@
+import Foundation
+
+/// Handles sharded SafeTensors models (with model.safetensors.index.json).
+public final class SafeTensorsIndex {
+    public let weightMap: [String: String]
+    public let baseDir: String
+
+    /// Load the index file from a model directory.
+    public init(modelDir: String) throws {
+        let indexURL = URL(fileURLWithPath: modelDir).appendingPathComponent("model.safetensors.index.json")
+        let data = try Data(contentsOf: indexURL)
+        guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any],
+              let weightMap = json["weight_map"] as? [String: String]
+        else {
+            throw WeightError.invalidHeader("Index file missing weight_map")
+        }
+        self.weightMap = weightMap
+        self.baseDir = modelDir
+    }
+
+    /// All unique shard filenames referenced by the index.
+    public var shardFiles: Set<String> {
+        Set(weightMap.values)
+    }
+
+    /// Resolve a tensor name to its shard file path.
+    public func shardPath(for tensor: String) -> String? {
+        guard let shard = weightMap[tensor] else { return nil }
+        return (baseDir as NSString).appendingPathComponent(shard)
+    }
+
+    /// List all tensor names.
+    public var allTensors: [String] { Array(weightMap.keys) }
+}

Sources/MarkBase/Weights/TensorDescriptor.swift

@@ -0,0 +1,38 @@
+/// Metadata for a single tensor stored in a SafeTensors file.
+public struct TensorDescriptor: Sendable, Codable {
+    public let name: String
+    public let dtype: TensorDType
+    public let shape: [Int]
+    /// Byte offset from the start of the safetensors data section.
+    public let dataOffset: Int
+    /// Byte size of the tensor data.
+    public let dataSize: Int
+
+    /// Total number of elements.
+    public var elementCount: Int { shape.reduce(1, *) }
+
+    /// Check if shape is compatible with a given dim count.
+    public func hasRank(_ rank: Int) -> Bool { shape.count == rank }
+
+    /// For quantized tensors: returns the grouping factor (elements per group).
+    /// MLX default: 64 elements per quantization group (for Gemma 4 E4B 4-bit).
+    public var quantizationGroupSize: Int { 64 }
+}
+
+/// Group of tensors that together represent a quantized linear layer.
+///   weight: U32 packed (shape: [outDim, inDim / 32 * 4])
+///   scales: BF16    (shape: [outDim, inDim / 32])
+///   biases: BF16    (shape: [outDim, inDim / 32])
+public struct QuantizedTensorGroup: Sendable {
+    public let name: String
+    public let weight: TensorDescriptor
+    public let scales: TensorDescriptor
+    public let biases: TensorDescriptor
+
+    /// Output dimension.
+    public var outDim: Int { weight.shape[0] }
+    /// Input dimension (pre-quantization).
+    public var inDim: Int { scales.shape[1] * 32 }
+    /// Block size (elements per group).
+    public let groupSize: Int = 64
+}