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

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