bleam 0.0.9 → 0.0.11

package/templates/native/ios/Bleam/AI/TextGeneration/Qwen3/Qwen3CheckpointLoader.swift

@@ -0,0 +1,161 @@
+import Foundation
+import MLX
+import MLXNN
+
+public enum Qwen3CheckpointError: LocalizedError {
+    case missingFile(String)
+    case unsupportedArchitecture(String)
+    case unsupportedQuantization(bits: Int?, groupSize: Int?)
+    case invalidConfiguration(String)
+    case invalidCheckpoint(String)
+
+    public var errorDescription: String? {
+        switch self {
+        case .missingFile(let file): "Qwen3 checkpoint is missing \(file)"
+        case .unsupportedArchitecture(let architecture): "Unsupported Qwen3 architecture \(architecture)"
+        case .unsupportedQuantization(let bits, let groupSize):
+            "Bonsai Qwen3 requires 2-bit/group-128 affine weights, got bits=\(bits.map(String.init) ?? "nil") group=\(groupSize.map(String.init) ?? "nil")"
+        case .invalidConfiguration(let message): "Invalid Qwen3 configuration: \(message)"
+        case .invalidCheckpoint(let message): "Invalid Qwen3 checkpoint: \(message)"
+        }
+    }
+}
+
+public enum Qwen3CheckpointLoader {
+    private struct SafetensorsIndex: Decodable {
+        let weightMap: [String: String]
+
+        enum CodingKeys: String, CodingKey {
+            case weightMap = "weight_map"
+        }
+    }
+
+    public static func load(from modelURL: URL) throws -> Qwen3CausalLM {
+        let (configuration, quantization, tensors) = try loadSnapshot(from: modelURL, requireLanguageHead: true)
+
+        let model = Qwen3CausalLM(configuration: configuration)
+        MLXNN.quantize(
+            model: model,
+            groupSize: quantization.groupSize,
+            bits: quantization.bits,
+            mode: .affine
+        )
+        do {
+            try model.update(
+                parameters: ModuleParameters.unflattened(tensors),
+                verify: [.noUnusedKeys, .shapeMismatch]
+            )
+        } catch {
+            throw Qwen3CheckpointError.invalidCheckpoint(error.localizedDescription)
+        }
+        model.usesTiedOutputHead = configuration.tieWordEmbeddings && tensors["lm_head.weight"] == nil
+        MLX.eval(model)
+        MLX.Memory.clearCache()
+        return model
+    }
+
+    public static func loadTransformer(from modelURL: URL) throws -> Qwen3TransformerModel {
+        let (configuration, quantization, tensors) = try loadSnapshot(from: modelURL, requireLanguageHead: false)
+        let model = Qwen3TransformerModel(configuration: configuration)
+        MLXNN.quantize(
+            model: model,
+            groupSize: quantization.groupSize,
+            bits: quantization.bits,
+            mode: .affine
+        )
+        let transformerTensors: [String: MLXArray] = Dictionary(uniqueKeysWithValues: tensors.compactMap { name, value in
+            guard name.hasPrefix("model.") else { return nil }
+            return (String(name.dropFirst("model.".count)), value)
+        })
+        do {
+            try model.update(
+                parameters: ModuleParameters.unflattened(transformerTensors),
+                verify: [.noUnusedKeys, .shapeMismatch]
+            )
+        } catch {
+            throw Qwen3CheckpointError.invalidCheckpoint(error.localizedDescription)
+        }
+        MLX.eval(model)
+        MLX.Memory.clearCache()
+        return model
+    }
+
+    private static func loadSnapshot(
+        from modelURL: URL,
+        requireLanguageHead: Bool
+    ) throws -> (Qwen3ModelConfiguration, Qwen3ModelConfiguration.Quantization, [String: MLXArray]) {
+        let configuration = try Qwen3ModelConfiguration.load(from: modelURL)
+        try configuration.validate()
+        guard let quantization = configuration.quantization,
+              quantization.bits == 2, quantization.groupSize == 128 else {
+            throw Qwen3CheckpointError.unsupportedQuantization(
+                bits: configuration.quantization?.bits,
+                groupSize: configuration.quantization?.groupSize
+            )
+        }
+        let files = try checkpointFiles(at: modelURL)
+        var tensors: [String: MLXArray] = [:]
+        for file in files {
+            for (name, value) in try loadArrays(url: modelURL.appendingPathComponent(file)) {
+                tensors[name] = value
+            }
+        }
+        try validateTensors(
+            tensors,
+            configuration: configuration,
+            requireLanguageHead: requireLanguageHead
+        )
+        return (configuration, quantization, tensors)
+    }
+
+    private static func checkpointFiles(at modelURL: URL) throws -> [String] {
+        let fileManager = FileManager.default
+        let indexURL = modelURL.appendingPathComponent("model.safetensors.index.json")
+        if fileManager.fileExists(atPath: indexURL.path) {
+            let index = try JSONDecoder().decode(SafetensorsIndex.self, from: Data(contentsOf: indexURL))
+            let files = Set(index.weightMap.values).sorted()
+            for file in files where !fileManager.fileExists(atPath: modelURL.appendingPathComponent(file).path) {
+                throw Qwen3CheckpointError.missingFile(file)
+            }
+            return files
+        }
+        let file = "model.safetensors"
+        guard fileManager.fileExists(atPath: modelURL.appendingPathComponent(file).path) else {
+            throw Qwen3CheckpointError.missingFile(file)
+        }
+        return [file]
+    }
+
+    private static func validateTensors(
+        _ tensors: [String: MLXArray],
+        configuration: Qwen3ModelConfiguration,
+        requireLanguageHead: Bool
+    ) throws {
+        var requiredModules = ["model.embed_tokens"]
+        for index in 0..<configuration.hiddenLayerCount {
+            let prefix = "model.layers.\(index)"
+            requiredModules += [
+                "\(prefix).self_attn.q_proj", "\(prefix).self_attn.k_proj",
+                "\(prefix).self_attn.v_proj", "\(prefix).self_attn.o_proj",
+                "\(prefix).mlp.gate_proj", "\(prefix).mlp.up_proj", "\(prefix).mlp.down_proj",
+            ]
+            for norm in ["input_layernorm", "post_attention_layernorm", "self_attn.q_norm", "self_attn.k_norm"] {
+                let key = "\(prefix).\(norm).weight"
+                guard tensors[key] != nil else { throw Qwen3CheckpointError.invalidCheckpoint("missing tensor \(key)") }
+            }
+        }
+        if requireLanguageHead, !configuration.tieWordEmbeddings { requiredModules.append("lm_head") }
+        for module in requiredModules {
+            for suffix in ["weight", "scales", "biases"] {
+                let key = "\(module).\(suffix)"
+                guard tensors[key] != nil else { throw Qwen3CheckpointError.invalidCheckpoint("missing tensor \(key)") }
+            }
+            guard tensors["\(module).weight"]?.dtype == .uint32 else {
+                throw Qwen3CheckpointError.invalidCheckpoint("\(module).weight is not packed uint32")
+            }
+        }
+        guard tensors["model.norm.weight"] != nil else {
+            throw Qwen3CheckpointError.invalidCheckpoint("missing tensor model.norm.weight")
+        }
+    }
+}

package/templates/native/ios/Bleam/AI/TextGeneration/Qwen3/Qwen3Configuration.swift

@@ -0,0 +1,119 @@
+import Foundation
+
+public struct Qwen3ModelConfiguration: Decodable, Sendable {
+    public struct RopeScaling: Decodable, Sendable {
+        public let ropeType: String
+        public let factor: Float
+        public let originalMaximumPositions: Int
+        public let betaFast: Float
+        public let betaSlow: Float
+
+        enum CodingKeys: String, CodingKey {
+            case ropeType = "rope_type"
+            case type
+            case factor
+            case originalMaximumPositions = "original_max_position_embeddings"
+            case betaFast = "beta_fast"
+            case betaSlow = "beta_slow"
+        }
+
+        public init(from decoder: Decoder) throws {
+            let values = try decoder.container(keyedBy: CodingKeys.self)
+            ropeType = try values.decodeIfPresent(String.self, forKey: .ropeType)
+                ?? values.decodeIfPresent(String.self, forKey: .type)
+                ?? "yarn"
+            factor = try values.decode(Float.self, forKey: .factor)
+            originalMaximumPositions = try values.decode(Int.self, forKey: .originalMaximumPositions)
+            betaFast = try values.decodeIfPresent(Float.self, forKey: .betaFast) ?? 32
+            betaSlow = try values.decodeIfPresent(Float.self, forKey: .betaSlow) ?? 1
+        }
+    }
+
+    public struct Quantization: Decodable, Sendable {
+        public let bits: Int
+        public let groupSize: Int
+
+        enum CodingKeys: String, CodingKey {
+            case bits
+            case groupSize = "group_size"
+        }
+    }
+
+    public let architectures: [String]
+    public let modelType: String
+    public let vocabularySize: Int
+    public let hiddenSize: Int
+    public let intermediateSize: Int
+    public let hiddenLayerCount: Int
+    public let attentionHeadCount: Int
+    public let keyValueHeadCount: Int
+    public let maximumPositionCount: Int
+    public let rmsNormEpsilon: Float
+    public let ropeTheta: Float
+    public let tieWordEmbeddings: Bool
+    public let attentionBias: Bool
+    public let headDimension: Int
+    public let ropeScaling: RopeScaling?
+    public let quantization: Quantization?
+
+    enum CodingKeys: String, CodingKey {
+        case architectures
+        case modelType = "model_type"
+        case vocabularySize = "vocab_size"
+        case hiddenSize = "hidden_size"
+        case intermediateSize = "intermediate_size"
+        case hiddenLayerCount = "num_hidden_layers"
+        case attentionHeadCount = "num_attention_heads"
+        case keyValueHeadCount = "num_key_value_heads"
+        case maximumPositionCount = "max_position_embeddings"
+        case rmsNormEpsilon = "rms_norm_eps"
+        case ropeTheta = "rope_theta"
+        case tieWordEmbeddings = "tie_word_embeddings"
+        case attentionBias = "attention_bias"
+        case headDimension = "head_dim"
+        case ropeScaling = "rope_scaling"
+        case quantization
+    }
+
+    public init(from decoder: Decoder) throws {
+        let values = try decoder.container(keyedBy: CodingKeys.self)
+        architectures = try values.decodeIfPresent([String].self, forKey: .architectures) ?? []
+        modelType = try values.decodeIfPresent(String.self, forKey: .modelType) ?? "qwen3"
+        vocabularySize = try values.decode(Int.self, forKey: .vocabularySize)
+        hiddenSize = try values.decode(Int.self, forKey: .hiddenSize)
+        intermediateSize = try values.decode(Int.self, forKey: .intermediateSize)
+        hiddenLayerCount = try values.decode(Int.self, forKey: .hiddenLayerCount)
+        attentionHeadCount = try values.decode(Int.self, forKey: .attentionHeadCount)
+        keyValueHeadCount = try values.decode(Int.self, forKey: .keyValueHeadCount)
+        maximumPositionCount = try values.decode(Int.self, forKey: .maximumPositionCount)
+        rmsNormEpsilon = try values.decodeIfPresent(Float.self, forKey: .rmsNormEpsilon) ?? 1e-6
+        ropeTheta = try values.decodeIfPresent(Float.self, forKey: .ropeTheta) ?? 1_000_000
+        tieWordEmbeddings = try values.decodeIfPresent(Bool.self, forKey: .tieWordEmbeddings) ?? true
+        attentionBias = try values.decodeIfPresent(Bool.self, forKey: .attentionBias) ?? false
+        headDimension = try values.decodeIfPresent(Int.self, forKey: .headDimension)
+            ?? hiddenSize / attentionHeadCount
+        ropeScaling = try values.decodeIfPresent(RopeScaling.self, forKey: .ropeScaling)
+        quantization = try values.decodeIfPresent(Quantization.self, forKey: .quantization)
+    }
+
+    public static func load(from modelURL: URL) throws -> Qwen3ModelConfiguration {
+        try JSONDecoder().decode(
+            Qwen3ModelConfiguration.self,
+            from: Data(contentsOf: modelURL.appendingPathComponent("config.json"))
+        )
+    }
+
+    public func validate() throws {
+        guard modelType == "qwen3", architectures.isEmpty || architectures.contains("Qwen3ForCausalLM") else {
+            throw Qwen3CheckpointError.unsupportedArchitecture(architectures.first ?? modelType)
+        }
+        guard hiddenSize > 0, intermediateSize > 0, hiddenLayerCount > 0,
+              attentionHeadCount > 0, keyValueHeadCount > 0,
+              attentionHeadCount.isMultiple(of: keyValueHeadCount), headDimension > 0 else {
+            throw Qwen3CheckpointError.invalidConfiguration("invalid model dimensions")
+        }
+        if let ropeScaling, ropeScaling.ropeType != "yarn" {
+            throw Qwen3CheckpointError.invalidConfiguration("unsupported RoPE scaling type \(ropeScaling.ropeType)")
+        }
+    }
+}

package/templates/native/ios/Bleam/AI/TextGeneration/Qwen3/Qwen3InferenceSession.swift

@@ -0,0 +1,84 @@
+import MLX
+
+public final class Qwen3InferenceSession: @unchecked Sendable {
+    private let model: Qwen3CausalLM
+    private let endTokenIDs: Set<Int>
+
+    public init(model: Qwen3CausalLM, endTokenIDs: Set<Int> = [151643, 151645]) {
+        self.model = model
+        self.endTokenIDs = endTokenIDs
+    }
+
+    public func generate(
+        inputIDs: [Int],
+        parameters: GenerationParameters = .balanced,
+        isCancelled: (() -> Bool)? = nil,
+        onToken: ((Int) -> Bool)? = nil
+    ) throws -> TokenGenerationResult {
+        try TokenSampler.validate(parameters)
+        guard !inputIDs.isEmpty else { throw TextGenerationError.emptyPrompt }
+        for token in inputIDs where token < 0 || token >= model.vocabularySize {
+            throw TextGenerationError.invalidTokenID(token)
+        }
+        let requestedLength = inputIDs.count + parameters.maximumNewTokens
+        guard requestedLength <= model.maximumContextLength else {
+            throw TextGenerationError.contextExceeded(
+                requested: requestedLength,
+                maximum: model.maximumContextLength
+            )
+        }
+        guard parameters.maximumNewTokens > 0 else {
+            return TokenGenerationResult(promptTokenCount: inputIDs.count, tokenIDs: [], finishReason: .length)
+        }
+        if isCancelled?() == true {
+            return TokenGenerationResult(promptTokenCount: inputIDs.count, tokenIDs: [], finishReason: .cancelled)
+        }
+        if let seed = parameters.seed { MLXRandom.seed(seed) }
+
+        let cache = model.makeCache()
+        defer {
+            cache.forEach { $0.clear() }
+            MLX.Memory.clearCache()
+        }
+        var logits = model.logits(
+            for: MLXArray(inputIDs.map(Int32.init)).reshaped([1, inputIDs.count]),
+            cache: cache
+        )
+        var generated: [Int] = []
+        var history = inputIDs
+        var finishReason: GenerationFinishReason = .length
+
+        for _ in 0..<parameters.maximumNewTokens {
+            if isCancelled?() == true {
+                finishReason = .cancelled
+                break
+            }
+            let sampled = try TokenSampler.sample(
+                logits: logits,
+                parameters: parameters,
+                tokenHistory: history
+            )
+            MLX.eval(sampled)
+            let token = Int(sampled.item(Int32.self))
+            if endTokenIDs.contains(token) {
+                finishReason = .endOfSequence
+                break
+            }
+            generated.append(token)
+            history.append(token)
+            if onToken?(token) == false {
+                finishReason = .callbackStopped
+                break
+            }
+            logits = model.logits(
+                for: MLXArray([Int32(token)]).reshaped([1, 1]),
+                cache: cache
+            )
+        }
+        return TokenGenerationResult(
+            promptTokenCount: inputIDs.count,
+            tokenIDs: generated,
+            finishReason: finishReason
+        )
+    }
+}

package/templates/native/ios/Bleam/AI/TextGeneration/Qwen3/Qwen3Layers.swift

@@ -0,0 +1,197 @@
+import Foundation
+import MLX
+import MLXNN
+
+public final class Qwen3RMSNorm: Module, UnaryLayer {
+    @ParameterInfo public var weight: MLXArray
+    public let epsilon: Float
+
+    public init(dimensions: Int, epsilon: Float) {
+        _weight = ParameterInfo(wrappedValue: MLXArray.ones([dimensions]))
+        self.epsilon = epsilon
+        super.init()
+    }
+
+    public func callAsFunction(_ x: MLXArray) -> MLXArray {
+        MLXFast.rmsNorm(x, weight: weight, eps: epsilon)
+    }
+}
+
+public final class Qwen3RotaryEmbedding: Module {
+    private let dimensions: Int
+    private let base: Float
+    private let frequencies: MLXArray?
+    private let magnitudeScale: Float
+
+    public init(configuration: Qwen3ModelConfiguration) {
+        dimensions = configuration.headDimension
+        base = configuration.ropeTheta
+        if let scaling = configuration.ropeScaling {
+            precondition(configuration.headDimension.isMultiple(of: 2))
+            func correctionDimension(_ rotations: Float) -> Float {
+                Float(configuration.headDimension)
+                    * log(Float(scaling.originalMaximumPositions) / (rotations * 2 * .pi))
+                    / (2 * log(configuration.ropeTheta))
+            }
+            let low = max(Int(floor(correctionDimension(scaling.betaFast))), 0)
+            let high = min(Int(ceil(correctionDimension(scaling.betaSlow))), configuration.headDimension - 1)
+            let safeHigh = low == high ? Float(high) + 0.001 : Float(high)
+            let ramp = clip(
+                (MLXArray(0..<(configuration.headDimension / 2)).asType(.float32) - Float(low))
+                    / (safeHigh - Float(low)),
+                min: 0,
+                max: 1
+            )
+            let exponent = MLXArray(stride(from: 0, to: configuration.headDimension, by: 2))
+                .asType(.float32) / Float(configuration.headDimension)
+            let unscaled = pow(configuration.ropeTheta, exponent)
+            let interpolated = scaling.factor * unscaled
+            let mask = 1 - ramp
+            frequencies = (interpolated * unscaled) / (interpolated * mask + unscaled * (1 - mask))
+            magnitudeScale = scaling.factor <= 1 ? 1 : 0.1 * log(scaling.factor) + 1
+        } else {
+            frequencies = nil
+            magnitudeScale = 1
+        }
+        super.init()
+    }
+
+    public func callAsFunction(_ x: MLXArray, offset: Int) -> MLXArray {
+        MLXFast.RoPE(
+            magnitudeScale == 1 ? x : x * magnitudeScale,
+            dimensions: dimensions,
+            traditional: false,
+            base: frequencies == nil ? base : nil,
+            scale: 1,
+            offset: offset,
+            freqs: frequencies
+        )
+    }
+}
+
+public final class Qwen3SelfAttention: Module {
+    private let headCount: Int
+    private let keyValueHeadCount: Int
+    private let headDimension: Int
+    private let scale: Float
+
+    @ModuleInfo public var q_proj: Linear
+    @ModuleInfo public var k_proj: Linear
+    @ModuleInfo public var v_proj: Linear
+    @ModuleInfo public var o_proj: Linear
+    @ModuleInfo public var q_norm: Qwen3RMSNorm
+    @ModuleInfo public var k_norm: Qwen3RMSNorm
+    @ModuleInfo public var rope: Qwen3RotaryEmbedding
+
+    public init(configuration: Qwen3ModelConfiguration) {
+        headCount = configuration.attentionHeadCount
+        keyValueHeadCount = configuration.keyValueHeadCount
+        headDimension = configuration.headDimension
+        scale = 1 / sqrt(Float(configuration.headDimension))
+        _q_proj = ModuleInfo(wrappedValue: Linear(
+            configuration.hiddenSize,
+            configuration.attentionHeadCount * configuration.headDimension,
+            bias: configuration.attentionBias
+        ))
+        _k_proj = ModuleInfo(wrappedValue: Linear(
+            configuration.hiddenSize,
+            configuration.keyValueHeadCount * configuration.headDimension,
+            bias: configuration.attentionBias
+        ))
+        _v_proj = ModuleInfo(wrappedValue: Linear(
+            configuration.hiddenSize,
+            configuration.keyValueHeadCount * configuration.headDimension,
+            bias: configuration.attentionBias
+        ))
+        _o_proj = ModuleInfo(wrappedValue: Linear(
+            configuration.attentionHeadCount * configuration.headDimension,
+            configuration.hiddenSize,
+            bias: configuration.attentionBias
+        ))
+        _q_norm = ModuleInfo(wrappedValue: Qwen3RMSNorm(
+            dimensions: configuration.headDimension,
+            epsilon: configuration.rmsNormEpsilon
+        ))
+        _k_norm = ModuleInfo(wrappedValue: Qwen3RMSNorm(
+            dimensions: configuration.headDimension,
+            epsilon: configuration.rmsNormEpsilon
+        ))
+        _rope = ModuleInfo(wrappedValue: Qwen3RotaryEmbedding(configuration: configuration))
+        super.init()
+    }
+
+    public func callAsFunction(_ hiddenStates: MLXArray, mask: MLXArray?, cache: TextKVCache?) -> MLXArray {
+        let batchSize = hiddenStates.dim(0)
+        let sequenceLength = hiddenStates.dim(1)
+        var queries = q_proj(hiddenStates).reshaped([batchSize, sequenceLength, headCount, headDimension])
+        var keys = k_proj(hiddenStates).reshaped([batchSize, sequenceLength, keyValueHeadCount, headDimension])
+        var values = v_proj(hiddenStates).reshaped([batchSize, sequenceLength, keyValueHeadCount, headDimension])
+        queries = q_norm(queries).transposed(0, 2, 1, 3)
+        keys = k_norm(keys).transposed(0, 2, 1, 3)
+        values = values.transposed(0, 2, 1, 3)
+        let offset = cache?.length ?? 0
+        queries = rope(queries, offset: offset)
+        keys = rope(keys, offset: offset)
+        if let cache {
+            (keys, values) = cache.update(keys: keys, values: values)
+        }
+        let attended = MLXFast.scaledDotProductAttention(
+            queries: queries,
+            keys: keys,
+            values: values,
+            scale: scale,
+            mask: mask
+        )
+        return o_proj(attended.transposed(0, 2, 1, 3).reshaped([
+            batchSize, sequenceLength, headCount * headDimension,
+        ]))
+    }
+}
+
+public final class Qwen3FeedForward: Module, @unchecked Sendable {
+    @ModuleInfo public var gate_proj: Linear
+    @ModuleInfo public var up_proj: Linear
+    @ModuleInfo public var down_proj: Linear
+
+    nonisolated(unsafe) private static let activate: (MLXArray, MLXArray) -> MLXArray = compile {
+        silu($0) * $1
+    }
+
+    public init(configuration: Qwen3ModelConfiguration) {
+        _gate_proj = ModuleInfo(wrappedValue: Linear(configuration.hiddenSize, configuration.intermediateSize, bias: false))
+        _up_proj = ModuleInfo(wrappedValue: Linear(configuration.hiddenSize, configuration.intermediateSize, bias: false))
+        _down_proj = ModuleInfo(wrappedValue: Linear(configuration.intermediateSize, configuration.hiddenSize, bias: false))
+        super.init()
+    }
+
+    public func callAsFunction(_ x: MLXArray) -> MLXArray {
+        down_proj(Self.activate(gate_proj(x), up_proj(x)))
+    }
+}
+
+public final class Qwen3DecoderBlock: Module {
+    @ModuleInfo public var self_attn: Qwen3SelfAttention
+    @ModuleInfo public var mlp: Qwen3FeedForward
+    @ModuleInfo public var input_layernorm: Qwen3RMSNorm
+    @ModuleInfo public var post_attention_layernorm: Qwen3RMSNorm
+
+    public init(configuration: Qwen3ModelConfiguration) {
+        _self_attn = ModuleInfo(wrappedValue: Qwen3SelfAttention(configuration: configuration))
+        _mlp = ModuleInfo(wrappedValue: Qwen3FeedForward(configuration: configuration))
+        _input_layernorm = ModuleInfo(wrappedValue: Qwen3RMSNorm(
+            dimensions: configuration.hiddenSize,
+            epsilon: configuration.rmsNormEpsilon
+        ))
+        _post_attention_layernorm = ModuleInfo(wrappedValue: Qwen3RMSNorm(
+            dimensions: configuration.hiddenSize,
+            epsilon: configuration.rmsNormEpsilon
+        ))
+        super.init()
+    }
+
+    public func callAsFunction(_ hiddenStates: MLXArray, mask: MLXArray?, cache: TextKVCache?) -> MLXArray {
+        var hidden = hiddenStates + self_attn(input_layernorm(hiddenStates), mask: mask, cache: cache)
+        hidden = hidden + mlp(post_attention_layernorm(hidden))
+        return hidden
+    }
+}

package/templates/native/ios/Bleam/AI/TextGeneration/Qwen3/Qwen3Model.swift

@@ -0,0 +1,130 @@
+import MLX
+import MLXNN
+
+public final class Qwen3TransformerModel: Module, HiddenStateModel {
+    public let configuration: Qwen3ModelConfiguration
+
+    @ModuleInfo public var embed_tokens: Embedding
+    public var layers: [Qwen3DecoderBlock]
+    @ModuleInfo public var norm: Qwen3RMSNorm
+
+    public var vocabularySize: Int { configuration.vocabularySize }
+    public var maximumContextLength: Int { configuration.maximumPositionCount }
+
+    public init(configuration: Qwen3ModelConfiguration) {
+        self.configuration = configuration
+        _embed_tokens = ModuleInfo(wrappedValue: Embedding(
+            embeddingCount: configuration.vocabularySize,
+            dimensions: configuration.hiddenSize
+        ))
+        layers = (0..<configuration.hiddenLayerCount).map { _ in
+            Qwen3DecoderBlock(configuration: configuration)
+        }
+        _norm = ModuleInfo(wrappedValue: Qwen3RMSNorm(
+            dimensions: configuration.hiddenSize,
+            epsilon: configuration.rmsNormEpsilon
+        ))
+        super.init()
+    }
+
+    public func callAsFunction(_ inputIDs: MLXArray, cache: [TextKVCache]? = nil) -> MLXArray {
+        var hidden = embed_tokens(inputIDs)
+        let mask = causalMask(sequenceLength: inputIDs.dim(1), offset: cache?.first?.length ?? 0)
+        for (index, layer) in layers.enumerated() {
+            hidden = layer(hidden, mask: mask, cache: cache?[index])
+        }
+        return norm(hidden)
+    }
+
+    public func hiddenStates(
+        for inputIDs: MLXArray,
+        at layerIndices: Set<Int>,
+        attentionMask: MLXArray? = nil
+    ) -> [Int: MLXArray] {
+        var hidden = embed_tokens(inputIDs)
+        let mask = causalMask(sequenceLength: inputIDs.dim(1), offset: 0, attentionMask: attentionMask)
+        var result: [Int: MLXArray] = [:]
+        if layerIndices.contains(0) {
+            MLX.eval(hidden)
+            result[0] = hidden
+        }
+        for (index, layer) in layers.enumerated() {
+            hidden = layer(hidden, mask: mask, cache: nil)
+            let layerIndex = index + 1
+            if layerIndices.contains(layerIndex) {
+                let state = layerIndex == layers.count ? norm(hidden) : hidden
+                MLX.eval(state)
+                result[layerIndex] = state
+            }
+        }
+        return result
+    }
+
+    private func causalMask(
+        sequenceLength: Int,
+        offset: Int,
+        attentionMask: MLXArray? = nil
+    ) -> MLXArray? {
+        if sequenceLength == 1, attentionMask == nil { return nil }
+        let totalLength = sequenceLength + offset
+        let rows = MLXArray.arange(sequenceLength, dtype: .float32).expandedDimensions(axis: 1)
+        let columns = MLXArray.arange(totalLength, dtype: .float32).expandedDimensions(axis: 0)
+        var mask = MLX.where(
+            columns .<= rows + Float(offset),
+            MLXArray(Float(0)),
+            MLXArray(-Float.infinity)
+        ).reshaped([1, 1, sequenceLength, totalLength])
+        if let attentionMask {
+            let padding = MLX.where(
+                attentionMask .== Int32(1),
+                MLXArray(Float(0)),
+                MLXArray(Float(-1e9))
+            ).reshaped([attentionMask.dim(0), 1, 1, attentionMask.dim(1)])
+            mask = mask + padding
+        }
+        return mask
+    }
+}
+
+public final class Qwen3CausalLM: Module, CausalLanguageModel {
+    public let configuration: Qwen3ModelConfiguration
+    @ModuleInfo public var model: Qwen3TransformerModel
+    @ModuleInfo public var lm_head: Linear
+    public var usesTiedOutputHead: Bool
+
+    public var vocabularySize: Int { configuration.vocabularySize }
+    public var maximumContextLength: Int { configuration.maximumPositionCount }
+
+    public init(configuration: Qwen3ModelConfiguration) {
+        self.configuration = configuration
+        usesTiedOutputHead = configuration.tieWordEmbeddings
+        _model = ModuleInfo(wrappedValue: Qwen3TransformerModel(configuration: configuration))
+        _lm_head = ModuleInfo(wrappedValue: Linear(configuration.hiddenSize, configuration.vocabularySize, bias: false))
+        super.init()
+    }
+
+    public func logits(for inputIDs: MLXArray, cache: [TextKVCache]) -> MLXArray {
+        project(model(inputIDs, cache: cache))
+    }
+
+    public func makeCache() -> [TextKVCache] {
+        (0..<configuration.hiddenLayerCount).map { _ in TextKVCache() }
+    }
+
+    private func project(_ hiddenStates: MLXArray) -> MLXArray {
+        guard usesTiedOutputHead else { return lm_head(hiddenStates) }
+        if let embedding = model.embed_tokens as? QuantizedEmbedding {
+            return MLX.quantizedMM(
+                hiddenStates,
+                embedding.weight,
+                scales: embedding.scales,
+                biases: embedding.biases,
+                transpose: true,
+                groupSize: embedding.groupSize,
+                bits: embedding.bits,
+                mode: .affine
+            )
+        }
+        return matmul(hiddenStates, model.embed_tokens.weight.T)
+    }
+}