bleam 0.0.10 → 0.0.11

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()
+    }
+}

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

@@ -0,0 +1,111 @@
+import Foundation
+import MLX
+import MLXNN
+
+public final class QwenHybridLinearAttention: Module, @unchecked Sendable {
+    private let numValueHeads: Int
+    private let numKeyHeads: Int
+    private let keyHeadDim: Int
+    private let valueHeadDim: Int
+    private let keyDimension: Int
+    private let valueDimension: Int
+    private let convolutionDimension: Int
+    private let kernelSize: Int
+    private let keyNormWeight: MLXArray
+
+    @ModuleInfo public var in_proj_qkv: Linear
+    @ModuleInfo public var in_proj_z: Linear
+    @ModuleInfo public var in_proj_b: Linear
+    @ModuleInfo public var in_proj_a: Linear
+    @ModuleInfo public var out_proj: Linear
+    @ModuleInfo public var norm: QwenHybridGatedRMSNorm
+    @ModuleInfo public var conv1d: QwenHybridDepthwiseConv1D
+    @ParameterInfo public var dt_bias: MLXArray
+    @ParameterInfo public var A_log: MLXArray
+
+    public init(config: any QwenHybridConfiguration) {
+        numValueHeads = config.linearNumValueHeads
+        numKeyHeads = config.linearNumKeyHeads
+        keyHeadDim = config.linearKeyHeadDim
+        valueHeadDim = config.linearValueHeadDim
+        keyDimension = config.linearNumKeyHeads * config.linearKeyHeadDim
+        valueDimension = config.linearNumValueHeads * config.linearValueHeadDim
+        convolutionDimension = keyDimension * 2 + valueDimension
+        kernelSize = config.linearConvKernelDim
+        keyNormWeight = MLX.ones([config.linearKeyHeadDim])
+        _in_proj_qkv = ModuleInfo(wrappedValue: Linear(config.hiddenSize, convolutionDimension, bias: false))
+        _in_proj_z = ModuleInfo(wrappedValue: Linear(config.hiddenSize, valueDimension, bias: false))
+        _in_proj_b = ModuleInfo(wrappedValue: Linear(config.hiddenSize, numValueHeads, bias: false))
+        _in_proj_a = ModuleInfo(wrappedValue: Linear(config.hiddenSize, numValueHeads, bias: false))
+        _out_proj = ModuleInfo(wrappedValue: Linear(valueDimension, config.hiddenSize, bias: false))
+        _norm = ModuleInfo(wrappedValue: QwenHybridGatedRMSNorm(dimensions: valueHeadDim, epsilon: config.rmsNormEps))
+        _conv1d = ModuleInfo(wrappedValue: QwenHybridDepthwiseConv1D(channels: convolutionDimension, kernelSize: kernelSize))
+        _dt_bias = ParameterInfo(wrappedValue: MLX.ones([numValueHeads]))
+        _A_log = ParameterInfo(wrappedValue: MLX.zeros([numValueHeads]))
+        super.init()
+    }
+
+    public func callAsFunction(_ input: MLXArray, cache: QwenHybridLayerCache) -> MLXArray {
+        let batch = input.dim(0)
+        let length = input.dim(1)
+        let mixed = causalConvolution(in_proj_qkv(input), cache: cache)
+        let pieces = mixed.split(indices: [keyDimension, keyDimension * 2], axis: -1)
+        var query = pieces[0].reshaped([batch, length, numKeyHeads, keyHeadDim])
+        var key = pieces[1].reshaped([batch, length, numKeyHeads, keyHeadDim])
+        let value = pieces[2].reshaped([batch, length, numValueHeads, valueHeadDim]).asType(.float32)
+        query = MLXFast.rmsNorm(query.asType(.float32), weight: keyNormWeight, eps: 1e-6) / Float(keyHeadDim)
+        key = MLXFast.rmsNorm(key.asType(.float32), weight: keyNormWeight, eps: 1e-6) / sqrt(Float(keyHeadDim))
+        let gate = MLX.exp(-MLX.exp(A_log.asType(.float32)) * MLXNN.softplus(in_proj_a(input).asType(.float32) + dt_bias))
+        let beta = MLX.sigmoid(in_proj_b(input)).asType(.float32)
+        let initialState = cache.recurrentState ?? MLX.zeros([batch, numValueHeads, valueHeadDim, keyHeadDim], dtype: .float32)
+        let updated = QwenHybridGatedDelta.update(query: query, key: key, value: value, gate: gate, beta: beta, state: initialState)
+        cache.recurrentState = updated.state
+        let normalized = norm(
+            updated.output.asType(input.dtype).reshaped([-1, valueHeadDim]),
+            gate: in_proj_z(input).reshaped([-1, valueHeadDim])
+        )
+        return out_proj(normalized.reshaped([batch, length, valueDimension]))
+    }
+
+    private func causalConvolution(_ input: MLXArray, cache: QwenHybridLayerCache) -> MLXArray {
+        let length = input.dim(1)
+        if let previous = cache.convolutionState {
+            let combined = MLX.concatenated([previous, input], axis: 1)
+            cache.convolutionState = combined[0..., (combined.dim(1) - kernelSize + 1)..., 0...]
+            return MLXNN.silu(MLX.conv1d(combined, conv1d.weight, padding: 0, groups: convolutionDimension))
+        }
+        let result = MLXNN.silu(MLX.conv1d(input, conv1d.weight, padding: kernelSize - 1, groups: convolutionDimension))
+        let paddedInput = padded(input, widths: [[0, 0], [max(0, kernelSize - 1 - length), 0], [0, 0]])
+        cache.convolutionState = paddedInput[0..., (paddedInput.dim(1) - kernelSize + 1)..., 0...]
+        return result[0..., 0..<length, 0...]
+    }
+}
+
+public enum QwenHybridGatedDelta {
+    public static func update(
+        query: MLXArray,
+        key: MLXArray,
+        value: MLXArray,
+        gate: MLXArray,
+        beta: MLXArray,
+        state initialState: MLXArray
+    ) -> (output: MLXArray, state: MLXArray) {
+        let repeats = value.dim(2) / key.dim(2)
+        let query = repeats == 1 ? query : MLX.repeated(query, count: repeats, axis: 2)
+        let key = repeats == 1 ? key : MLX.repeated(key, count: repeats, axis: 2)
+        var state = initialState
+        var outputs: [MLXArray] = []
+        outputs.reserveCapacity(query.dim(1))
+        for index in 0..<query.dim(1) {
+            let q = query[0..., index, 0..., 0...]
+            let k = key[0..., index, 0..., 0...]
+            let v = value[0..., index, 0..., 0...]
+            state = state * gate[0..., index, 0...].expandedDimensions(axes: [2, 3])
+            let remembered = (state * k.expandedDimensions(axis: 2)).sum(axis: -1)
+            let delta = (v - remembered) * beta[0..., index, 0...].expandedDimensions(axis: -1)
+            state = state + k.expandedDimensions(axis: 2) * delta.expandedDimensions(axis: -1)
+            outputs.append((state * q.expandedDimensions(axis: 2)).sum(axis: -1))
+        }
+        return (MLX.stacked(outputs, axis: 1), state)
+    }
+}