bleam 0.0.9 → 0.0.11

package/templates/native/ios/Bleam/AI/TextGeneration/Qwen35/Qwen35Checkpoint.swift

@@ -0,0 +1,102 @@
+import Foundation
+import MLX
+import MLXNN
+
+public enum Qwen35Checkpoint {
+    private struct Index: Decodable {
+        let weightMap: [String: String]
+
+        private enum CodingKeys: String, CodingKey { case weightMap = "weight_map" }
+    }
+
+    public static func load(
+        from directory: URL,
+        expectedQuantization: Qwen35Quantization? = nil
+    ) throws -> Qwen35ForCausalLM {
+        let config = try Qwen35Configuration.load(from: directory)
+        try config.validate(expectedQuantization: expectedQuantization)
+        let index = try checkpointIndex(in: directory)
+        try validateFiles(index: index, directory: directory, config: config)
+        let model = Qwen35ForCausalLM(config: config)
+        guard config.quantization.mode.lowercased() == "affine" else {
+            throw Qwen35Error.invalidCheckpoint("unsupported quantization mode \(config.quantization.mode)")
+        }
+        MLXNN.quantize(model: model, groupSize: config.quantization.groupSize, bits: config.quantization.bits, mode: .affine)
+
+        var weights: [String: MLXArray] = [:]
+        var usesUpstreamNormOffsets = false
+        for file in Set(index.weightMap.values) {
+            let arrays = try MLX.loadArrays(url: directory.appendingPathComponent(file))
+            usesUpstreamNormOffsets = usesUpstreamNormOffsets || arrays.contains { name, value in
+                name.hasSuffix("conv1d.weight") && value.ndim == 3 && value.dim(2) != 1
+            }
+            for (name, var value) in arrays {
+                guard let key = modelKey(from: name) else { continue }
+                if key.hasSuffix("conv1d.weight"), value.ndim == 3, value.dim(2) != 1 {
+                    value = value.transposed(0, 2, 1)
+                }
+                weights[key] = value
+            }
+        }
+        if usesUpstreamNormOffsets {
+            for key in Array(weights.keys) where isOffsetNorm(key) && weights[key]?.ndim == 1 {
+                weights[key] = weights[key]! + 1
+            }
+        }
+        guard !weights.isEmpty else { throw Qwen35Error.invalidCheckpoint("no language-model tensors") }
+        do {
+            try model.update(parameters: ModuleParameters.unflattened(weights), verify: [.noUnusedKeys, .shapeMismatch])
+        } catch {
+            throw Qwen35Error.invalidCheckpoint(error.localizedDescription)
+        }
+        MLX.eval(model)
+        MLX.Memory.clearCache()
+        return model
+    }
+
+    private static func checkpointIndex(in directory: URL) throws -> Index {
+        let indexURL = directory.appendingPathComponent("model.safetensors.index.json")
+        if FileManager.default.fileExists(atPath: indexURL.path) {
+            return try JSONDecoder().decode(Index.self, from: Data(contentsOf: indexURL))
+        }
+        let single = directory.appendingPathComponent("model.safetensors")
+        guard FileManager.default.fileExists(atPath: single.path) else { throw Qwen35Error.missingFile(single.lastPathComponent) }
+        let names = try MLX.loadArrays(url: single).keys
+        return Index(weightMap: Dictionary(uniqueKeysWithValues: names.map { ($0, single.lastPathComponent) }))
+    }
+
+    private static func validateFiles(index: Index, directory: URL, config: Qwen35Configuration) throws {
+        var required = [
+            "language_model.model.embed_tokens.weight",
+            "language_model.model.layers.0.linear_attn.in_proj_qkv.weight",
+            "language_model.model.layers.0.mlp.gate_proj.weight",
+            "language_model.model.norm.weight",
+        ]
+        if let fullLayer = config.layerTypes.firstIndex(of: "full_attention") {
+            required.append("language_model.model.layers.\(fullLayer).self_attn.q_proj.weight")
+        }
+        if !config.tieWordEmbeddings { required.append("language_model.lm_head.weight") }
+        for key in required where index.weightMap[key] == nil {
+            throw Qwen35Error.invalidCheckpoint("missing tensor \(key)")
+        }
+        for file in Set(index.weightMap.values) where !FileManager.default.fileExists(atPath: directory.appendingPathComponent(file).path) {
+            throw Qwen35Error.missingFile(file)
+        }
+    }
+
+    private static func modelKey(from checkpointKey: String) -> String? {
+        let modelPrefix = "language_model.model."
+        if checkpointKey.hasPrefix(modelPrefix) { return "model." + checkpointKey.dropFirst(modelPrefix.count) }
+        let headPrefix = "language_model.lm_head."
+        if checkpointKey.hasPrefix(headPrefix) { return "lm_head." + checkpointKey.dropFirst(headPrefix.count) }
+        return nil
+    }
+
+    private static func isOffsetNorm(_ key: String) -> Bool {
+        key.hasSuffix(".input_layernorm.weight")
+            || key.hasSuffix(".post_attention_layernorm.weight")
+            || key.hasSuffix(".q_norm.weight")
+            || key.hasSuffix(".k_norm.weight")
+            || key == "model.norm.weight"
+    }
+}

package/templates/native/ios/Bleam/AI/TextGeneration/Qwen35/Qwen35Configuration.swift

@@ -0,0 +1,140 @@
+import Foundation
+
+public struct Qwen35Quantization: Decodable, Equatable, Sendable {
+    public let groupSize: Int
+    public let bits: Int
+    public let mode: String
+
+    private enum CodingKeys: String, CodingKey {
+        case groupSize = "group_size"
+        case bits
+        case mode
+    }
+
+    public init(groupSize: Int, bits: Int, mode: String = "affine") {
+        self.groupSize = groupSize
+        self.bits = bits
+        self.mode = mode
+    }
+
+    public init(from decoder: Decoder) throws {
+        let container = try decoder.container(keyedBy: CodingKeys.self)
+        groupSize = try container.decode(Int.self, forKey: .groupSize)
+        bits = try container.decode(Int.self, forKey: .bits)
+        mode = try container.decodeIfPresent(String.self, forKey: .mode) ?? "affine"
+    }
+}
+
+public struct Qwen35Configuration: Decodable, Sendable, QwenHybridConfiguration {
+    public let modelType: String
+    public let hiddenSize: Int
+    public let intermediateSize: Int
+    public let numHiddenLayers: Int
+    public let numAttentionHeads: Int
+    public let numKeyValueHeads: Int
+    public let headDim: Int
+    public let vocabSize: Int
+    public let maxPositionEmbeddings: Int
+    public let rmsNormEps: Float
+    public let hiddenAct: String
+    public let attentionBias: Bool
+    public let tieWordEmbeddings: Bool
+    public let layerTypes: [String]
+    public let fullAttentionInterval: Int
+    public let linearConvKernelDim: Int
+    public let linearKeyHeadDim: Int
+    public let linearValueHeadDim: Int
+    public let linearNumKeyHeads: Int
+    public let linearNumValueHeads: Int
+    public let ropeTheta: Float
+    public let partialRotaryFactor: Float
+    public let quantization: Qwen35Quantization
+
+    private enum RootKeys: String, CodingKey {
+        case modelType = "model_type"
+        case textConfig = "text_config"
+        case quantization
+        case quantizationConfig = "quantization_config"
+    }
+    private enum TextKeys: String, CodingKey {
+        case hiddenSize = "hidden_size", intermediateSize = "intermediate_size", numHiddenLayers = "num_hidden_layers"
+        case numAttentionHeads = "num_attention_heads", numKeyValueHeads = "num_key_value_heads", headDim = "head_dim"
+        case vocabSize = "vocab_size", maxPositionEmbeddings = "max_position_embeddings", rmsNormEps = "rms_norm_eps"
+        case hiddenAct = "hidden_act", attentionBias = "attention_bias", tieWordEmbeddings = "tie_word_embeddings"
+        case layerTypes = "layer_types", fullAttentionInterval = "full_attention_interval"
+        case linearConvKernelDim = "linear_conv_kernel_dim", linearKeyHeadDim = "linear_key_head_dim"
+        case linearValueHeadDim = "linear_value_head_dim", linearNumKeyHeads = "linear_num_key_heads"
+        case linearNumValueHeads = "linear_num_value_heads", ropeParameters = "rope_parameters"
+    }
+    private enum RopeKeys: String, CodingKey {
+        case ropeTheta = "rope_theta"
+        case partialRotaryFactor = "partial_rotary_factor"
+    }
+
+    public init(from decoder: Decoder) throws {
+        let root = try decoder.container(keyedBy: RootKeys.self)
+        modelType = try root.decode(String.self, forKey: .modelType)
+        quantization = try root.decodeIfPresent(Qwen35Quantization.self, forKey: .quantization)
+            ?? root.decode(Qwen35Quantization.self, forKey: .quantizationConfig)
+        let text = try root.nestedContainer(keyedBy: TextKeys.self, forKey: .textConfig)
+        hiddenSize = try text.decode(Int.self, forKey: .hiddenSize)
+        intermediateSize = try text.decode(Int.self, forKey: .intermediateSize)
+        numHiddenLayers = try text.decode(Int.self, forKey: .numHiddenLayers)
+        numAttentionHeads = try text.decode(Int.self, forKey: .numAttentionHeads)
+        numKeyValueHeads = try text.decode(Int.self, forKey: .numKeyValueHeads)
+        headDim = try text.decode(Int.self, forKey: .headDim)
+        vocabSize = try text.decode(Int.self, forKey: .vocabSize)
+        maxPositionEmbeddings = try text.decode(Int.self, forKey: .maxPositionEmbeddings)
+        rmsNormEps = try text.decode(Float.self, forKey: .rmsNormEps)
+        hiddenAct = try text.decode(String.self, forKey: .hiddenAct)
+        attentionBias = try text.decode(Bool.self, forKey: .attentionBias)
+        tieWordEmbeddings = try text.decode(Bool.self, forKey: .tieWordEmbeddings)
+        layerTypes = try text.decode([String].self, forKey: .layerTypes)
+        fullAttentionInterval = try text.decode(Int.self, forKey: .fullAttentionInterval)
+        linearConvKernelDim = try text.decode(Int.self, forKey: .linearConvKernelDim)
+        linearKeyHeadDim = try text.decode(Int.self, forKey: .linearKeyHeadDim)
+        linearValueHeadDim = try text.decode(Int.self, forKey: .linearValueHeadDim)
+        linearNumKeyHeads = try text.decode(Int.self, forKey: .linearNumKeyHeads)
+        linearNumValueHeads = try text.decode(Int.self, forKey: .linearNumValueHeads)
+        let rope = try text.nestedContainer(keyedBy: RopeKeys.self, forKey: .ropeParameters)
+        ropeTheta = try rope.decode(Float.self, forKey: .ropeTheta)
+        partialRotaryFactor = try rope.decode(Float.self, forKey: .partialRotaryFactor)
+    }
+
+    public static func load(from directory: URL) throws -> Self {
+        try JSONDecoder().decode(Self.self, from: Data(contentsOf: directory.appendingPathComponent("config.json")))
+    }
+
+    public func validate(expectedQuantization: Qwen35Quantization? = nil) throws {
+        guard modelType == "qwen3_5" else { throw Qwen35Error.invalidConfiguration("model_type is \(modelType)") }
+        guard numHiddenLayers > 0, layerTypes.count == numHiddenLayers else { throw Qwen35Error.invalidConfiguration("layer_types count does not match num_hidden_layers") }
+        guard fullAttentionInterval > 0 else { throw Qwen35Error.invalidConfiguration("full_attention_interval must be positive") }
+        guard layerTypes.indices.allSatisfy({ layerTypes[$0] == (($0 + 1).isMultiple(of: fullAttentionInterval) ? "full_attention" : "linear_attention") }) else {
+            throw Qwen35Error.invalidConfiguration("layer_types does not match full_attention_interval")
+        }
+        guard hiddenAct == "silu" else { throw Qwen35Error.invalidConfiguration("hidden_act must be silu") }
+        if let expectedQuantization, quantization != expectedQuantization {
+            throw Qwen35Error.invalidConfiguration("checkpoint quantization does not match the requested runtime policy")
+        }
+    }
+}
+
+public enum Qwen35Error: LocalizedError {
+    case invalidConfiguration(String)
+    case invalidCheckpoint(String)
+    case missingFile(String)
+    case emptyPrompt
+    case contextExceeded(requested: Int, maximum: Int)
+    case decodeBeforePrefill
+
+    public var errorDescription: String? {
+        switch self {
+        case .invalidConfiguration(let message): "Invalid Qwen3.5 configuration: \(message)"
+        case .invalidCheckpoint(let message): "Invalid Qwen3.5 checkpoint: \(message)"
+        case .missingFile(let file): "Missing Qwen3.5 file: \(file)"
+        case .emptyPrompt: "Qwen3.5 requires at least one input token"
+        case .contextExceeded(let requested, let maximum): "Qwen3.5 context length \(requested) exceeds \(maximum)"
+        case .decodeBeforePrefill: "Qwen3.5 decode called before prefill"
+        }
+    }
+}

package/templates/native/ios/Bleam/AI/TextGeneration/Qwen35/Qwen35Model.swift

@@ -0,0 +1,112 @@
+import MLX
+import MLXNN
+
+public final class Qwen35MLP: Module, @unchecked Sendable {
+    @ModuleInfo public var gate_proj: Linear
+    @ModuleInfo public var up_proj: Linear
+    @ModuleInfo public var down_proj: Linear
+
+    public init(config: Qwen35Configuration) {
+        _gate_proj = ModuleInfo(wrappedValue: Linear(config.hiddenSize, config.intermediateSize, bias: false))
+        _up_proj = ModuleInfo(wrappedValue: Linear(config.hiddenSize, config.intermediateSize, bias: false))
+        _down_proj = ModuleInfo(wrappedValue: Linear(config.intermediateSize, config.hiddenSize, bias: false))
+        super.init()
+    }
+
+    public func callAsFunction(_ input: MLXArray) -> MLXArray {
+        down_proj(MLXNN.silu(gate_proj(input)) * up_proj(input))
+    }
+}
+
+public final class Qwen35DecoderLayer: Module, @unchecked Sendable {
+    @ModuleInfo public var self_attn: QwenHybridFullAttention?
+    @ModuleInfo public var linear_attn: QwenHybridLinearAttention?
+    @ModuleInfo public var mlp: Qwen35MLP
+    @ModuleInfo public var input_layernorm: QwenHybridRMSNorm
+    @ModuleInfo public var post_attention_layernorm: QwenHybridRMSNorm
+
+    public init(config: Qwen35Configuration, layerIndex: Int) {
+        let isFullAttention = config.layerTypes[layerIndex] == "full_attention"
+        _self_attn = ModuleInfo(wrappedValue: isFullAttention ? QwenHybridFullAttention(config: config) : nil)
+        _linear_attn = ModuleInfo(wrappedValue: isFullAttention ? nil : QwenHybridLinearAttention(config: config))
+        _mlp = ModuleInfo(wrappedValue: Qwen35MLP(config: config))
+        _input_layernorm = ModuleInfo(wrappedValue: QwenHybridRMSNorm(dimensions: config.hiddenSize, epsilon: config.rmsNormEps))
+        _post_attention_layernorm = ModuleInfo(wrappedValue: QwenHybridRMSNorm(dimensions: config.hiddenSize, epsilon: config.rmsNormEps))
+        super.init()
+    }
+
+    public func callAsFunction(_ input: MLXArray, mask: MLXArray?, cache: QwenHybridLayerCache) -> MLXArray {
+        let normalized = input_layernorm(input)
+        let mixed = linear_attn.map { $0(normalized, cache: cache) } ?? self_attn!(normalized, mask: mask, cache: cache)
+        let residual = input + mixed
+        return residual + mlp(post_attention_layernorm(residual))
+    }
+}
+
+public final class Qwen35Model: Module, @unchecked Sendable {
+    public let config: Qwen35Configuration
+    @ModuleInfo public var embed_tokens: Embedding
+    public var layers: [Qwen35DecoderLayer]
+    @ModuleInfo public var norm: QwenHybridRMSNorm
+
+    public init(config: Qwen35Configuration) {
+        self.config = config
+        _embed_tokens = ModuleInfo(wrappedValue: Embedding(embeddingCount: config.vocabSize, dimensions: config.hiddenSize))
+        layers = config.layerTypes.indices.map { Qwen35DecoderLayer(config: config, layerIndex: $0) }
+        _norm = ModuleInfo(wrappedValue: QwenHybridRMSNorm(dimensions: config.hiddenSize, epsilon: config.rmsNormEps))
+        super.init()
+    }
+
+    public func callAsFunction(_ tokenIDs: MLXArray, cache: [QwenHybridLayerCache]) -> MLXArray {
+        var hidden = embed_tokens(tokenIDs)
+        let length = tokenIDs.dim(1)
+        let offset = cache.first?.position ?? 0
+        let rows = MLXArray.arange(length, dtype: .float32).expandedDimensions(axis: 1)
+        let columns = MLXArray.arange(length + offset, dtype: .float32).expandedDimensions(axis: 0)
+        let mask: MLXArray? = length == 1 ? nil : MLX.where(
+            columns .<= rows + Float(offset),
+            MLXArray(Float(0)),
+            MLXArray(-Float.infinity)
+        ).reshaped([1, 1, length, length + offset])
+        for index in layers.indices {
+            hidden = layers[index](hidden, mask: mask, cache: cache[index])
+        }
+        cache.forEach { $0.position += length }
+        return norm(hidden)
+    }
+}
+
+public final class Qwen35ForCausalLM: Module, @unchecked Sendable {
+    public let config: Qwen35Configuration
+    public var model: Qwen35Model
+    @ModuleInfo public var lm_head: Linear?
+
+    public init(config: Qwen35Configuration) {
+        self.config = config
+        model = Qwen35Model(config: config)
+        _lm_head = ModuleInfo(wrappedValue: config.tieWordEmbeddings ? nil : Linear(config.hiddenSize, config.vocabSize, bias: false))
+        super.init()
+    }
+
+    public func callAsFunction(_ tokenIDs: MLXArray, cache: [QwenHybridLayerCache]) -> MLXArray {
+        let hidden = model(tokenIDs, cache: cache)
+        if let lm_head { return lm_head(hidden) }
+        if let embedding = model.embed_tokens as? QuantizedEmbedding {
+            return MLX.quantizedMM(
+                hidden,
+                embedding.weight,
+                scales: embedding.scales,
+                biases: embedding.biases,
+                transpose: true,
+                groupSize: embedding.groupSize,
+                bits: embedding.bits,
+                mode: embedding.mode
+            )
+        }
+        return MLX.matmul(hidden, model.embed_tokens.weight.T)
+    }
+
+    public func makeCache() -> [QwenHybridLayerCache] {
+        config.layerTypes.map { _ in QwenHybridLayerCache() }
+    }
+}

package/templates/native/ios/Bleam/AI/TextGeneration/Qwen35/Qwen35Session.swift

@@ -0,0 +1,121 @@
+import Foundation
+import MLX
+
+public final class Qwen35Session: @unchecked Sendable {
+    public let model: Qwen35ForCausalLM
+    public let cache: [QwenHybridLayerCache]
+    private var nextLogits: MLXArray?
+
+    public var contextLength: Int { cache.first?.position ?? 0 }
+    public var maximumContextLength: Int { model.config.maxPositionEmbeddings }
+
+    public init(model: Qwen35ForCausalLM) {
+        self.model = model
+        cache = model.makeCache()
+    }
+
+    public func reset() {
+        cache.forEach { $0.clear() }
+        nextLogits = nil
+        MLX.Memory.clearCache()
+    }
+
+    public func prefill(
+        tokenIDs: [Int],
+        chunkSize: Int = 512,
+        isCancelled: (@Sendable () -> Bool)? = nil
+    ) async throws {
+        guard !tokenIDs.isEmpty else { throw Qwen35Error.emptyPrompt }
+        guard contextLength + tokenIDs.count <= maximumContextLength else {
+            throw Qwen35Error.contextExceeded(requested: contextLength + tokenIDs.count, maximum: maximumContextLength)
+        }
+        var offset = 0
+        while offset < tokenIDs.count {
+            if isCancelled?() == true { throw CancellationError() }
+            try Task.checkCancellation()
+            let end = min(offset + max(1, chunkSize), tokenIDs.count)
+            let input = MLXArray(tokenIDs[offset..<end].map(Int32.init)).reshaped([1, end - offset])
+            let logits = model(input, cache: cache)[0..., -1, 0...]
+            MLX.eval(logits)
+            nextLogits = logits
+            offset = end
+            if offset < tokenIDs.count { MLX.Memory.clearCache() }
+            await Task.yield()
+        }
+    }
+
+    public func decodeNext(parameters: GenerationParameters = .balanced, tokenHistory: [Int] = []) throws -> Int {
+        try Task.checkCancellation()
+        guard let nextLogits else { throw Qwen35Error.decodeBeforePrefill }
+        guard contextLength < maximumContextLength else {
+            throw Qwen35Error.contextExceeded(requested: contextLength + 1, maximum: maximumContextLength)
+        }
+        let tokenArray = try TokenSampler.sample(logits: nextLogits, parameters: parameters, tokenHistory: tokenHistory)
+        MLX.eval(tokenArray)
+        let token = Int(tokenArray.item(Int32.self))
+        let input = MLXArray([Int32(token)]).reshaped([1, 1])
+        let logits = model(input, cache: cache)[0..., -1, 0...]
+        MLX.asyncEval(logits)
+        self.nextLogits = logits
+        return token
+    }
+
+    public func generate(
+        inputIDs: [Int],
+        parameters: GenerationParameters = .balanced,
+        endTokenIDs: Set<Int> = [151643, 151645],
+        chunkSize: Int = 512,
+        isCancelled: (@Sendable () -> Bool)? = nil,
+        onToken: (@Sendable (Int) -> Bool)? = nil
+    ) async throws -> TokenGenerationResult {
+        try TokenSampler.validate(parameters)
+        guard !inputIDs.isEmpty else { throw TextGenerationError.emptyPrompt }
+        for token in inputIDs where token < 0 || token >= model.config.vocabSize {
+            throw TextGenerationError.invalidTokenID(token)
+        }
+        let requestedLength = inputIDs.count + parameters.maximumNewTokens
+        guard requestedLength <= maximumContextLength else {
+            throw TextGenerationError.contextExceeded(requested: requestedLength, maximum: maximumContextLength)
+        }
+        if let seed = parameters.seed { MLXRandom.seed(seed) }
+        reset()
+        if parameters.maximumNewTokens == 0 {
+            return TokenGenerationResult(promptTokenCount: inputIDs.count, tokenIDs: [], finishReason: .length)
+        }
+        if isCancelled?() == true || Task.isCancelled {
+            return TokenGenerationResult(promptTokenCount: inputIDs.count, tokenIDs: [], finishReason: .cancelled)
+        }
+        do {
+            try await prefill(tokenIDs: inputIDs, chunkSize: chunkSize, isCancelled: isCancelled)
+        } catch is CancellationError {
+            return TokenGenerationResult(promptTokenCount: inputIDs.count, tokenIDs: [], finishReason: .cancelled)
+        }
+        var generated: [Int] = []
+        var history = inputIDs
+        var finishReason: GenerationFinishReason = .length
+        generated.reserveCapacity(parameters.maximumNewTokens)
+        for _ in 0..<parameters.maximumNewTokens {
+            if isCancelled?() == true || Task.isCancelled {
+                finishReason = .cancelled
+                break
+            }
+            let token = try decodeNext(parameters: parameters, tokenHistory: history)
+            if endTokenIDs.contains(token) {
+                finishReason = .endOfSequence
+                break
+            }
+            generated.append(token)
+            history.append(token)
+            if onToken?(token) == false {
+                finishReason = .callbackStopped
+                break
+            }
+            await Task.yield()
+        }
+        return TokenGenerationResult(promptTokenCount: inputIDs.count, tokenIDs: generated, finishReason: finishReason)
+    }
+}
+
+// CausalLanguageModel currently requires [TextKVCache], which cannot carry the
+// convolution and recurrent states needed by hybrid layers. Qwen35Session still
+// uses the generic generation parameters, sampler, result, and error types.

package/templates/native/ios/Bleam/AI/TextGeneration/QwenHybrid/QwenHybridCache.swift

@@ -0,0 +1,42 @@
+import MLX
+
+public final class QwenHybridKVCache: @unchecked Sendable {
+    public private(set) var keys: MLXArray?
+    public private(set) var values: MLXArray?
+
+    public var length: Int { keys?.dim(2) ?? 0 }
+
+    public init() {}
+
+    public func update(keys newKeys: MLXArray, values newValues: MLXArray) -> (MLXArray, MLXArray) {
+        if let keys, let values {
+            self.keys = MLX.concatenated([keys, newKeys], axis: 2)
+            self.values = MLX.concatenated([values, newValues], axis: 2)
+        } else {
+            keys = newKeys
+            values = newValues
+        }
+        return (keys!, values!)
+    }
+
+    public func clear() {
+        keys = nil
+        values = nil
+    }
+}
+
+public final class QwenHybridLayerCache: @unchecked Sendable {
+    public let kv = QwenHybridKVCache()
+    public var convolutionState: MLXArray?
+    public var recurrentState: MLXArray?
+    public var position = 0
+
+    public init() {}
+
+    public func clear() {
+        kv.clear()
+        convolutionState = nil
+        recurrentState = nil
+        position = 0
+    }
+}

package/templates/native/ios/Bleam/AI/TextGeneration/QwenHybrid/QwenHybridFullAttention.swift

@@ -0,0 +1,57 @@
+import Foundation
+import MLX
+import MLXNN
+
+public final class QwenHybridFullAttention: Module, @unchecked Sendable {
+    private let numHeads: Int
+    private let numKVHeads: Int
+    private let headDim: Int
+    private let rotaryDimensions: Int
+    private let ropeBase: Float
+    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: QwenHybridRMSNorm
+    @ModuleInfo public var k_norm: QwenHybridRMSNorm
+
+    public init(config: any QwenHybridConfiguration) {
+        numHeads = config.numAttentionHeads
+        numKVHeads = config.numKeyValueHeads
+        headDim = config.headDim
+        rotaryDimensions = Int(Float(config.headDim) * config.partialRotaryFactor)
+        ropeBase = config.ropeTheta
+        scale = 1 / sqrt(Float(config.headDim))
+        _q_proj = ModuleInfo(wrappedValue: Linear(config.hiddenSize, config.numAttentionHeads * config.headDim * 2, bias: config.attentionBias))
+        _k_proj = ModuleInfo(wrappedValue: Linear(config.hiddenSize, config.numKeyValueHeads * config.headDim, bias: config.attentionBias))
+        _v_proj = ModuleInfo(wrappedValue: Linear(config.hiddenSize, config.numKeyValueHeads * config.headDim, bias: config.attentionBias))
+        _o_proj = ModuleInfo(wrappedValue: Linear(config.numAttentionHeads * config.headDim, config.hiddenSize, bias: config.attentionBias))
+        _q_norm = ModuleInfo(wrappedValue: QwenHybridRMSNorm(dimensions: config.headDim, epsilon: config.rmsNormEps))
+        _k_norm = ModuleInfo(wrappedValue: QwenHybridRMSNorm(dimensions: config.headDim, epsilon: config.rmsNormEps))
+        super.init()
+    }
+
+    public func callAsFunction(_ input: MLXArray, mask: MLXArray?, cache: QwenHybridLayerCache) -> MLXArray {
+        let batch = input.dim(0)
+        let length = input.dim(1)
+        let projected = q_proj(input).reshaped([batch, length, numHeads, headDim * 2]).split(parts: 2, axis: -1)
+        var query = q_norm(projected[0]).transposed(0, 2, 1, 3)
+        let outputGate = projected[1].reshaped([batch, length, numHeads * headDim])
+        var key = k_norm(k_proj(input).reshaped([batch, length, numKVHeads, headDim])).transposed(0, 2, 1, 3)
+        var value = v_proj(input).reshaped([batch, length, numKVHeads, headDim]).transposed(0, 2, 1, 3)
+        query = MLXFast.RoPE(query, dimensions: rotaryDimensions, traditional: false, base: ropeBase, scale: 1, offset: cache.position)
+        key = MLXFast.RoPE(key, dimensions: rotaryDimensions, traditional: false, base: ropeBase, scale: 1, offset: cache.position)
+        (key, value) = cache.kv.update(keys: key, values: value)
+        let attended = MLXFast.scaledDotProductAttention(
+            queries: query,
+            keys: key,
+            values: value,
+            scale: scale,
+            mask: mask?.asType(query.dtype)
+        )
+        let merged = attended.transposed(0, 2, 1, 3).reshaped([batch, length, numHeads * headDim])
+        return o_proj(merged * MLXNN.sigmoid(outputGate))
+    }
+}

package/templates/native/ios/Bleam/AI/TextGeneration/QwenHybrid/QwenHybridLayers.swift

@@ -0,0 +1,58 @@
+import MLX
+import MLXNN
+
+public protocol QwenHybridConfiguration: Sendable {
+    var hiddenSize: Int { get }
+    var numAttentionHeads: Int { get }
+    var numKeyValueHeads: Int { get }
+    var headDim: Int { get }
+    var partialRotaryFactor: Float { get }
+    var attentionBias: Bool { get }
+    var rmsNormEps: Float { get }
+    var ropeTheta: Float { get }
+    var linearNumValueHeads: Int { get }
+    var linearNumKeyHeads: Int { get }
+    var linearKeyHeadDim: Int { get }
+    var linearValueHeadDim: Int { get }
+    var linearConvKernelDim: Int { get }
+}
+
+public final class QwenHybridRMSNorm: Module, UnaryLayer, @unchecked Sendable {
+    @ParameterInfo public var weight: MLXArray
+    public let epsilon: Float
+
+    public init(dimensions: Int, epsilon: Float) {
+        _weight = ParameterInfo(wrappedValue: MLXArray.zeros([dimensions]))
+        self.epsilon = epsilon
+        super.init()
+    }
+
+    public func callAsFunction(_ input: MLXArray) -> MLXArray {
+        MLXFast.rmsNorm(input, weight: weight, eps: epsilon)
+    }
+}
+
+public final class QwenHybridGatedRMSNorm: Module, @unchecked Sendable {
+    @ParameterInfo public var weight: MLXArray
+    public let epsilon: Float
+
+    public init(dimensions: Int, epsilon: Float) {
+        _weight = ParameterInfo(wrappedValue: MLXArray.zeros([dimensions]))
+        self.epsilon = epsilon
+        super.init()
+    }
+
+    public func callAsFunction(_ input: MLXArray, gate: MLXArray) -> MLXArray {
+        let normalized = MLXFast.rmsNorm(input, weight: weight, eps: epsilon)
+        return (MLXNN.silu(gate.asType(.float32)) * normalized.asType(.float32)).asType(input.dtype)
+    }
+}
+
+public final class QwenHybridDepthwiseConv1D: Module, @unchecked Sendable {
+    @ParameterInfo public var weight: MLXArray
+
+    public init(channels: Int, kernelSize: Int) {
+        _weight = ParameterInfo(wrappedValue: MLX.zeros([channels, kernelSize, 1]))
+        super.init()
+    }
+}