feat: migrate on-device Qwen3 ASR to CoreML for background Flow dictation
Replace MLX GPU inference with CoreML bundles so transcription continues while the host app is backgrounded. Adds model download and warm-up, vendored Qwen3Speech, and updates onboarding, settings, and copy for the ~1.6 GB CoreML package (iOS 18+).
This commit is contained in:
Rocky
@@ -0,0 +1,107 @@
|
||||
import Foundation
|
||||
|
||||
/// Standalone token sampler for text generation.
|
||||
///
|
||||
/// Supports temperature scaling, top-K filtering, top-P (nucleus) filtering,
|
||||
/// and repetition penalty.
|
||||
enum ChatSampler {
|
||||
|
||||
/// Sample a token from logits using the given sampling config.
|
||||
///
|
||||
/// - Parameters:
|
||||
/// - logits: Raw logits array of size vocab_size
|
||||
/// - config: Sampling parameters (temperature, topK, topP, repetitionPenalty)
|
||||
/// - previousTokens: Recently generated tokens for repetition penalty
|
||||
/// - Returns: Sampled token index
|
||||
static func sample(
|
||||
logits: [Float],
|
||||
config: ChatSamplingConfig,
|
||||
previousTokens: [Int] = []
|
||||
) -> Int {
|
||||
var logits = logits
|
||||
|
||||
// Repetition penalty
|
||||
if config.repetitionPenalty > 1.0 {
|
||||
let seen = Set(previousTokens.suffix(64))
|
||||
for tokenId in seen {
|
||||
if tokenId < logits.count {
|
||||
if logits[tokenId] > 0 {
|
||||
logits[tokenId] /= config.repetitionPenalty
|
||||
} else {
|
||||
logits[tokenId] *= config.repetitionPenalty
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Greedy (argmax) when temperature is 0
|
||||
if config.temperature <= 0 {
|
||||
var maxIdx = 0
|
||||
var maxVal = logits[0]
|
||||
for i in 1..<logits.count {
|
||||
if logits[i] > maxVal {
|
||||
maxVal = logits[i]
|
||||
maxIdx = i
|
||||
}
|
||||
}
|
||||
return maxIdx
|
||||
}
|
||||
|
||||
// Temperature scaling
|
||||
if config.temperature != 1.0 {
|
||||
for i in 0..<logits.count {
|
||||
logits[i] /= config.temperature
|
||||
}
|
||||
}
|
||||
|
||||
// Softmax
|
||||
let maxLogit = logits.max() ?? 0
|
||||
var probs = logits.map { exp($0 - maxLogit) }
|
||||
let sum = probs.reduce(0, +)
|
||||
probs = probs.map { $0 / sum }
|
||||
|
||||
// Top-K filtering
|
||||
if config.topK > 0 && config.topK < probs.count {
|
||||
let indexed = probs.enumerated().sorted { $0.element > $1.element }
|
||||
let topK = Array(indexed.prefix(config.topK))
|
||||
var filtered = [Float](repeating: 0, count: probs.count)
|
||||
for (idx, prob) in topK {
|
||||
filtered[idx] = prob
|
||||
}
|
||||
let filteredSum = filtered.reduce(0, +)
|
||||
if filteredSum > 0 {
|
||||
probs = filtered.map { $0 / filteredSum }
|
||||
}
|
||||
}
|
||||
|
||||
// Top-P (nucleus) filtering
|
||||
if config.topP < 1.0 {
|
||||
let indexed = probs.enumerated().sorted { $0.element > $1.element }
|
||||
var cumProb: Float = 0
|
||||
var mask = [Bool](repeating: false, count: probs.count)
|
||||
for (idx, prob) in indexed {
|
||||
cumProb += prob
|
||||
mask[idx] = true
|
||||
if cumProb >= config.topP { break }
|
||||
}
|
||||
for i in 0..<probs.count {
|
||||
if !mask[i] { probs[i] = 0 }
|
||||
}
|
||||
let filteredSum = probs.reduce(0, +)
|
||||
if filteredSum > 0 {
|
||||
probs = probs.map { $0 / filteredSum }
|
||||
}
|
||||
}
|
||||
|
||||
// Sample from distribution
|
||||
let r = Float.random(in: 0..<1)
|
||||
var cumulative: Float = 0
|
||||
for (i, p) in probs.enumerated() {
|
||||
cumulative += p
|
||||
if cumulative >= r {
|
||||
return i
|
||||
}
|
||||
}
|
||||
return probs.count - 1
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,104 @@
|
||||
import Foundation
|
||||
|
||||
/// Chat message for Qwen3.5 models.
|
||||
public struct ChatMessage: Sendable {
|
||||
public enum Role: String, Sendable {
|
||||
case system
|
||||
case user
|
||||
case assistant
|
||||
}
|
||||
|
||||
public let role: Role
|
||||
public let content: String
|
||||
|
||||
public init(role: Role, content: String) {
|
||||
self.role = role
|
||||
self.content = content
|
||||
}
|
||||
}
|
||||
|
||||
/// Formats messages into Qwen3.5 chat template tokens.
|
||||
///
|
||||
/// ```
|
||||
/// <|im_start|>system
|
||||
/// {system_message}<|im_end|>
|
||||
/// <|im_start|>user
|
||||
/// {user_message}<|im_end|>
|
||||
/// <|im_start|>assistant
|
||||
/// ```
|
||||
enum ChatTemplate {
|
||||
// Qwen3.5 special token IDs (248K vocab)
|
||||
static let imStartId = 248045 // <|im_start|>
|
||||
static let imEndId = 248046 // <|im_end|>
|
||||
static let endOfTextId = 248044 // <|endoftext|>
|
||||
static let thinkStartId = 248068 // <think>
|
||||
static let thinkEndId = 248069 // </think>
|
||||
static let newlineId = 198 // \n
|
||||
|
||||
/// Strip thinking block from generated tokens.
|
||||
///
|
||||
/// Removes tokens from `<think>` through `</think>` (inclusive)
|
||||
/// and any trailing newlines, returning only the response content.
|
||||
static func stripThinking(from tokens: [Int]) -> [Int] {
|
||||
let thinkTokens: Set<Int> = [thinkStartId, thinkEndId]
|
||||
let newlines: Set<Int> = [newlineId, 271] // 198 = \n, 271 = \n\n
|
||||
|
||||
guard let startIdx = tokens.firstIndex(where: { thinkTokens.contains($0) && $0 == thinkStartId }) else {
|
||||
// No <think> — strip any leading </think> + newlines
|
||||
// (happens when non-thinking template causes model to echo end-think)
|
||||
var i = 0
|
||||
while i < tokens.count && (tokens[i] == thinkEndId || newlines.contains(tokens[i])) {
|
||||
i += 1
|
||||
}
|
||||
return i > 0 ? Array(tokens[i...]) : tokens
|
||||
}
|
||||
if let endIdx = tokens[startIdx...].firstIndex(of: thinkEndId) {
|
||||
var afterThink = endIdx + 1
|
||||
while afterThink < tokens.count && newlines.contains(tokens[afterThink]) {
|
||||
afterThink += 1
|
||||
}
|
||||
return Array(tokens[0..<startIdx]) + Array(tokens[afterThink...])
|
||||
}
|
||||
return Array(tokens[0..<startIdx])
|
||||
}
|
||||
|
||||
/// Encode a conversation into token IDs using Qwen3.5 chat template.
|
||||
///
|
||||
/// - Parameters:
|
||||
/// - config: Model config (for future extensibility)
|
||||
/// - enableThinking: If false, injects empty think block to skip reasoning
|
||||
static func encode(
|
||||
messages: [ChatMessage],
|
||||
tokenizer: ChatTokenizer,
|
||||
config: Qwen3ChatConfig? = nil,
|
||||
addGenerationPrompt: Bool = true,
|
||||
enableThinking: Bool = true
|
||||
) -> [Int] {
|
||||
var tokens: [Int] = []
|
||||
|
||||
for message in messages {
|
||||
tokens.append(imStartId)
|
||||
tokens.append(contentsOf: tokenizer.encode(message.role.rawValue))
|
||||
tokens.append(newlineId)
|
||||
tokens.append(contentsOf: tokenizer.encode(message.content))
|
||||
tokens.append(imEndId)
|
||||
tokens.append(newlineId)
|
||||
}
|
||||
|
||||
if addGenerationPrompt {
|
||||
tokens.append(imStartId)
|
||||
tokens.append(contentsOf: tokenizer.encode("assistant"))
|
||||
tokens.append(newlineId)
|
||||
|
||||
if !enableThinking {
|
||||
let doubleNewline = tokenizer.encode("\n\n")
|
||||
tokens.append(thinkStartId)
|
||||
tokens.append(contentsOf: doubleNewline)
|
||||
tokens.append(thinkEndId)
|
||||
tokens.append(contentsOf: doubleNewline)
|
||||
}
|
||||
}
|
||||
|
||||
return tokens
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,272 @@
|
||||
import Foundation
|
||||
|
||||
/// Tokenizer for Qwen3 chat models.
|
||||
///
|
||||
/// Loads vocabulary from HuggingFace tokenizer files and provides
|
||||
/// encode/decode functionality for chat text.
|
||||
public final class ChatTokenizer: @unchecked Sendable {
|
||||
private var idToToken: [Int: String] = [:]
|
||||
private var tokenToId: [String: Int] = [:]
|
||||
private var bpeMerges: [(String, String)] = []
|
||||
private var bpeMergeRanks: [String: Int] = [:]
|
||||
private var addedTokens: [String: Int] = [:]
|
||||
|
||||
public var eosTokenId: Int = 248046 // <|im_end|>
|
||||
public var vocabSize: Int { idToToken.count }
|
||||
|
||||
public init() {}
|
||||
|
||||
/// Load tokenizer from a directory.
|
||||
///
|
||||
/// Supports two formats:
|
||||
/// 1. `tokenizer.json` (HuggingFace format, preferred) — contains vocab, merges, and added tokens
|
||||
/// 2. `vocab.json` + `merges.txt` (legacy) — separate files
|
||||
public func load(from directory: URL) throws {
|
||||
let tokenizerJsonURL = directory.appendingPathComponent("tokenizer.json")
|
||||
let vocabURL = directory.appendingPathComponent("vocab.json")
|
||||
|
||||
if FileManager.default.fileExists(atPath: tokenizerJsonURL.path) {
|
||||
try loadFromTokenizerJson(from: tokenizerJsonURL)
|
||||
} else {
|
||||
try loadVocab(from: vocabURL)
|
||||
|
||||
let mergesURL = directory.appendingPathComponent("merges.txt")
|
||||
if FileManager.default.fileExists(atPath: mergesURL.path) {
|
||||
try loadMerges(from: mergesURL)
|
||||
}
|
||||
}
|
||||
|
||||
let configURL = directory.appendingPathComponent("tokenizer_config.json")
|
||||
if FileManager.default.fileExists(atPath: configURL.path) {
|
||||
try loadAddedTokens(from: configURL)
|
||||
}
|
||||
}
|
||||
|
||||
/// Load from HuggingFace tokenizer.json (contains vocab + merges + added tokens).
|
||||
private func loadFromTokenizerJson(from url: URL) throws {
|
||||
let data = try Data(contentsOf: url)
|
||||
guard let root = try JSONSerialization.jsonObject(with: data) as? [String: Any],
|
||||
let model = root["model"] as? [String: Any],
|
||||
let vocab = model["vocab"] as? [String: Int] else {
|
||||
throw ChatModelError.tokenizerLoadFailed("Invalid tokenizer.json format")
|
||||
}
|
||||
|
||||
tokenToId = vocab
|
||||
idToToken = Dictionary(uniqueKeysWithValues: vocab.map { ($1, $0) })
|
||||
|
||||
// Load merges
|
||||
if let merges = model["merges"] as? [[String]] {
|
||||
for (i, pair) in merges.enumerated() {
|
||||
guard pair.count == 2 else { continue }
|
||||
bpeMerges.append((pair[0], pair[1]))
|
||||
bpeMergeRanks["\(pair[0]) \(pair[1])"] = i
|
||||
}
|
||||
} else if let merges = model["merges"] as? [String] {
|
||||
// Alternative format: each merge as "a b" string
|
||||
for (i, merge) in merges.enumerated() {
|
||||
let parts = merge.split(separator: " ", maxSplits: 1)
|
||||
guard parts.count == 2 else { continue }
|
||||
let pair = (String(parts[0]), String(parts[1]))
|
||||
bpeMerges.append(pair)
|
||||
bpeMergeRanks["\(pair.0) \(pair.1)"] = i
|
||||
}
|
||||
}
|
||||
|
||||
// Load added tokens
|
||||
if let addedList = root["added_tokens"] as? [[String: Any]] {
|
||||
for entry in addedList {
|
||||
guard let content = entry["content"] as? String,
|
||||
let id = entry["id"] as? Int else { continue }
|
||||
addedTokens[content] = id
|
||||
tokenToId[content] = id
|
||||
idToToken[id] = content
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private func loadVocab(from url: URL) throws {
|
||||
let data = try Data(contentsOf: url)
|
||||
guard let vocab = try JSONSerialization.jsonObject(with: data) as? [String: Int] else {
|
||||
throw ChatModelError.tokenizerLoadFailed("Invalid vocab.json format")
|
||||
}
|
||||
tokenToId = vocab
|
||||
idToToken = Dictionary(uniqueKeysWithValues: vocab.map { ($1, $0) })
|
||||
}
|
||||
|
||||
private func loadMerges(from url: URL) throws {
|
||||
let content = try String(contentsOf: url, encoding: .utf8)
|
||||
let lines = content.components(separatedBy: "\n")
|
||||
for (i, line) in lines.enumerated() {
|
||||
if line.hasPrefix("#") || line.isEmpty { continue }
|
||||
let parts = line.split(separator: " ", maxSplits: 1)
|
||||
if parts.count == 2 {
|
||||
let pair = (String(parts[0]), String(parts[1]))
|
||||
bpeMerges.append(pair)
|
||||
bpeMergeRanks["\(pair.0) \(pair.1)"] = i
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private func loadAddedTokens(from url: URL) throws {
|
||||
let data = try Data(contentsOf: url)
|
||||
guard let config = try JSONSerialization.jsonObject(with: data) as? [String: Any] else {
|
||||
return
|
||||
}
|
||||
if let added = config["added_tokens_decoder"] as? [String: Any] {
|
||||
for (idStr, value) in added {
|
||||
guard let id = Int(idStr),
|
||||
let info = value as? [String: Any],
|
||||
let content = info["content"] as? String else { continue }
|
||||
addedTokens[content] = id
|
||||
tokenToId[content] = id
|
||||
idToToken[id] = content
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Encode
|
||||
|
||||
/// Encode text to token IDs using BPE.
|
||||
public func encode(_ text: String) -> [Int] {
|
||||
if text.isEmpty { return [] }
|
||||
|
||||
// Check if it's a special/added token
|
||||
if let id = addedTokens[text] ?? tokenToId[text] {
|
||||
return [id]
|
||||
}
|
||||
|
||||
// Simple BPE encoding
|
||||
var words = tokenizeToWords(text)
|
||||
var allTokens: [Int] = []
|
||||
|
||||
for word in words {
|
||||
let wordTokens = bpeEncode(word)
|
||||
allTokens.append(contentsOf: wordTokens)
|
||||
}
|
||||
|
||||
return allTokens
|
||||
}
|
||||
|
||||
/// Split text into BPE-ready words (GPT-2/Qwen style).
|
||||
private func tokenizeToWords(_ text: String) -> [String] {
|
||||
// Simplified: split on spaces, prefix non-first words with Ġ (space marker)
|
||||
var words: [String] = []
|
||||
var isFirst = true
|
||||
for part in text.components(separatedBy: " ") {
|
||||
if part.isEmpty { continue }
|
||||
if isFirst {
|
||||
words.append(part)
|
||||
isFirst = false
|
||||
} else {
|
||||
words.append("Ġ" + part)
|
||||
}
|
||||
}
|
||||
return words
|
||||
}
|
||||
|
||||
/// BPE encode a single word.
|
||||
private func bpeEncode(_ word: String) -> [Int] {
|
||||
if let id = tokenToId[word] {
|
||||
return [id]
|
||||
}
|
||||
|
||||
var symbols = word.map { String($0) }
|
||||
if symbols.isEmpty { return [] }
|
||||
|
||||
while symbols.count > 1 {
|
||||
// Find best merge
|
||||
var bestRank = Int.max
|
||||
var bestIdx = -1
|
||||
for i in 0..<(symbols.count - 1) {
|
||||
let pair = "\(symbols[i]) \(symbols[i + 1])"
|
||||
if let rank = bpeMergeRanks[pair], rank < bestRank {
|
||||
bestRank = rank
|
||||
bestIdx = i
|
||||
}
|
||||
}
|
||||
if bestIdx < 0 { break }
|
||||
|
||||
// Apply merge
|
||||
let merged = symbols[bestIdx] + symbols[bestIdx + 1]
|
||||
symbols.replaceSubrange(bestIdx...bestIdx + 1, with: [merged])
|
||||
}
|
||||
|
||||
// Look up token IDs
|
||||
return symbols.compactMap { tokenToId[$0] }
|
||||
}
|
||||
|
||||
// MARK: - Decode
|
||||
|
||||
/// Decode token IDs to text.
|
||||
///
|
||||
/// Uses byte-level BPE decoding: token strings are mapped back to bytes
|
||||
/// via the GPT-2 byte-to-unicode table, then assembled into UTF-8 text.
|
||||
public func decode(_ tokenIds: [Int]) -> String {
|
||||
let pieces = tokenIds.compactMap { idToToken[$0] }
|
||||
let joined = pieces.joined()
|
||||
return decodeBPEString(joined)
|
||||
}
|
||||
|
||||
/// Decode a single token ID.
|
||||
public func decodeToken(_ tokenId: Int) -> String? {
|
||||
guard let piece = idToToken[tokenId] else { return nil }
|
||||
return decodeBPEString(piece)
|
||||
}
|
||||
|
||||
/// Convert a BPE token string to UTF-8 text.
|
||||
///
|
||||
/// GPT-2/Qwen byte-level BPE represents each byte as a specific Unicode
|
||||
/// character. This reverses that mapping and decodes the bytes as UTF-8.
|
||||
private func decodeBPEString(_ bpeString: String) -> String {
|
||||
var bytes: [UInt8] = []
|
||||
for char in bpeString {
|
||||
if let byte = Self.unicodeToByte[char] {
|
||||
bytes.append(byte)
|
||||
}
|
||||
}
|
||||
return String(bytes: bytes, encoding: .utf8) ?? bpeString
|
||||
}
|
||||
|
||||
/// GPT-2 byte-to-unicode mapping (reversed for decoding).
|
||||
///
|
||||
/// Maps Unicode characters back to the byte values they represent in
|
||||
/// GPT-2/Qwen byte-level BPE vocabulary.
|
||||
private static let unicodeToByte: [Character: UInt8] = {
|
||||
// Build the standard GPT-2 bytes_to_unicode table
|
||||
var byteToUnicode: [UInt8: Character] = [:]
|
||||
var n = 0
|
||||
|
||||
// Printable ASCII + Latin supplement ranges that map to themselves
|
||||
let ranges: [ClosedRange<UInt8>] = [
|
||||
33...126, // ! through ~
|
||||
161...172, // ¡ through ¬
|
||||
174...255, // ® through ÿ
|
||||
]
|
||||
for range in ranges {
|
||||
for b in range {
|
||||
byteToUnicode[b] = Character(Unicode.Scalar(UInt32(b))!)
|
||||
}
|
||||
}
|
||||
|
||||
// Remaining bytes (0-32, 127-160, 173) map to 256+n
|
||||
for b: UInt16 in 0...255 {
|
||||
if byteToUnicode[UInt8(b)] == nil {
|
||||
byteToUnicode[UInt8(b)] = Character(Unicode.Scalar(256 + UInt32(n))!)
|
||||
n += 1
|
||||
}
|
||||
}
|
||||
|
||||
// Reverse the mapping: unicode char → byte value
|
||||
var result: [Character: UInt8] = [:]
|
||||
for (byte, char) in byteToUnicode {
|
||||
result[char] = byte
|
||||
}
|
||||
return result
|
||||
}()
|
||||
|
||||
/// Check if a token ID is a special token (should not appear in output).
|
||||
public func isSpecialToken(_ tokenId: Int) -> Bool {
|
||||
guard let token = idToToken[tokenId] else { return false }
|
||||
return token.hasPrefix("<|") && token.hasSuffix("|>")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,425 @@
|
||||
import Foundation
|
||||
import MLXCommon
|
||||
import MLX
|
||||
import MLXNN
|
||||
import AudioCommon
|
||||
import os.log
|
||||
|
||||
private let log = OSLog(subsystem: "com.soniqo.qwen3chat", category: "MLX")
|
||||
|
||||
// MARK: - MLX Generator for Qwen3.5 Chat
|
||||
|
||||
/// MLX-based text generator for Qwen3.5-0.8B hybrid model.
|
||||
///
|
||||
/// Uses MLX for GPU inference
|
||||
/// on Apple Silicon GPUs. The hybrid DeltaNet + GatedAttention architecture
|
||||
/// requires managing two types of state:
|
||||
/// 1. DeltaNet recurrent states (carried across all tokens, O(1) per layer)
|
||||
/// 2. GatedAttention KV caches (grow with sequence length, only 6 layers)
|
||||
///
|
||||
/// Usage:
|
||||
/// ```swift
|
||||
/// let model = try await Qwen35MLXChat.fromPretrained()
|
||||
/// let response = try model.generate(messages: [
|
||||
/// ChatMessage(role: .user, content: "Hello!")
|
||||
/// ])
|
||||
/// ```
|
||||
public final class Qwen35MLXChat: @unchecked Sendable {
|
||||
public static let defaultModelId = "aufklarer/Qwen3.5-0.8B-Chat-MLX"
|
||||
|
||||
public let config: Qwen3ChatConfig
|
||||
public let tokenizer: ChatTokenizer
|
||||
let model: Qwen35MLXModel
|
||||
var state: Qwen35MLXModel.InferenceState
|
||||
var _isLoaded = true
|
||||
|
||||
// MARK: - Metrics
|
||||
|
||||
/// Generation metrics for performance tracking.
|
||||
public struct Metrics {
|
||||
public var prefillTimeMs: Double = 0
|
||||
public var prefillTokens: Int = 0
|
||||
public var decodeTimeMs: Double = 0
|
||||
public var decodeTokens: Int = 0
|
||||
|
||||
public var tokensPerSecond: Double {
|
||||
guard decodeTimeMs > 0 else { return 0 }
|
||||
return Double(decodeTokens) / (decodeTimeMs / 1000.0)
|
||||
}
|
||||
|
||||
public var msPerToken: Double {
|
||||
guard decodeTokens > 0 else { return 0 }
|
||||
return decodeTimeMs / Double(decodeTokens)
|
||||
}
|
||||
|
||||
public var prefillTokensPerSecond: Double {
|
||||
guard prefillTimeMs > 0 else { return 0 }
|
||||
return Double(prefillTokens) / (prefillTimeMs / 1000.0)
|
||||
}
|
||||
}
|
||||
|
||||
private(set) var metrics = Metrics()
|
||||
|
||||
/// Latest generation metrics.
|
||||
public var lastMetrics: (tokensPerSec: Double, prefillMs: Double, decodeMs: Double, msPerToken: Double) {
|
||||
(metrics.tokensPerSecond, metrics.prefillTimeMs, metrics.decodeTimeMs, metrics.msPerToken)
|
||||
}
|
||||
|
||||
// MARK: - Init
|
||||
|
||||
private init(config: Qwen3ChatConfig, tokenizer: ChatTokenizer, model: Qwen35MLXModel) {
|
||||
self.config = config
|
||||
self.tokenizer = tokenizer
|
||||
self.model = model
|
||||
self.state = .initial(config: config)
|
||||
}
|
||||
|
||||
// MARK: - Factory
|
||||
|
||||
/// Quantization variant.
|
||||
public enum Quantization: String {
|
||||
case int4
|
||||
case int8
|
||||
}
|
||||
|
||||
/// Load a pre-trained Qwen3.5 chat model from HuggingFace.
|
||||
///
|
||||
/// Downloads quantized safetensors and tokenizer on first use.
|
||||
/// Model is loaded into MLX for GPU inference on Apple Silicon.
|
||||
///
|
||||
/// - Parameters:
|
||||
/// - modelId: HuggingFace model ID (repo with int4/ and int8/ subdirs)
|
||||
/// - quantization: INT4 (404 MB) or INT8 (763 MB)
|
||||
/// - progressHandler: Optional callback for download/load progress
|
||||
public static func fromPretrained(
|
||||
modelId: String = defaultModelId,
|
||||
quantization: Quantization = .int4,
|
||||
cacheDir: URL? = nil,
|
||||
offlineMode: Bool = false,
|
||||
progressHandler: ((Double, String) -> Void)? = nil
|
||||
) async throws -> Qwen35MLXChat {
|
||||
let cacheDir = try cacheDir ?? HuggingFaceDownloader.getCacheDirectory(for: modelId)
|
||||
let variant = quantization.rawValue
|
||||
|
||||
// Download model files from variant subdirectory (int4/ or int8/)
|
||||
progressHandler?(0.05, "Downloading \(variant) model...")
|
||||
try await HuggingFaceDownloader.downloadWeights(
|
||||
modelId: modelId,
|
||||
to: cacheDir,
|
||||
additionalFiles: [
|
||||
"\(variant)/model.safetensors",
|
||||
"\(variant)/config.json",
|
||||
"\(variant)/tokenizer.json",
|
||||
"\(variant)/tokenizer_config.json",
|
||||
],
|
||||
offlineMode: offlineMode,
|
||||
progressHandler: { progress in
|
||||
progressHandler?(progress * 0.5, "Downloading...")
|
||||
}
|
||||
)
|
||||
|
||||
// Variant files are in a subdirectory
|
||||
let variantDir = cacheDir.appendingPathComponent(variant)
|
||||
|
||||
// Load config
|
||||
progressHandler?(0.5, "Loading config...")
|
||||
let config: Qwen3ChatConfig
|
||||
let configURL = variantDir.appendingPathComponent("config.json")
|
||||
if FileManager.default.fileExists(atPath: configURL.path) {
|
||||
config = try Qwen3ChatConfig.load(from: configURL)
|
||||
} else {
|
||||
config = .qwen35_08B
|
||||
}
|
||||
|
||||
// Load tokenizer
|
||||
progressHandler?(0.55, "Loading tokenizer...")
|
||||
let tokenizer = ChatTokenizer()
|
||||
try tokenizer.load(from: variantDir)
|
||||
|
||||
// Create model
|
||||
progressHandler?(0.6, "Creating model...")
|
||||
let model = Qwen35MLXModel(config: config)
|
||||
|
||||
// Load weights
|
||||
progressHandler?(0.65, "Loading weights...")
|
||||
try Qwen35WeightLoader.loadWeights(
|
||||
into: model, from: variantDir,
|
||||
progressHandler: { pct, msg in
|
||||
progressHandler?(0.65 + pct * 0.3, msg)
|
||||
})
|
||||
|
||||
progressHandler?(1.0, "Ready")
|
||||
return Qwen35MLXChat(config: config, tokenizer: tokenizer, model: model)
|
||||
}
|
||||
|
||||
/// Download tokenizer + weight files only — does not load MLX/Metal.
|
||||
public static func downloadWeightsOnly(
|
||||
modelId: String = defaultModelId,
|
||||
quantization: Quantization = .int4,
|
||||
cacheDir: URL? = nil,
|
||||
registry: ModelRegistry = .huggingFace(),
|
||||
progressHandler: ((Double, String) -> Void)? = nil
|
||||
) async throws {
|
||||
let cacheDir = try cacheDir ?? HuggingFaceDownloader.getCacheDirectory(for: modelId)
|
||||
let variant = quantization.rawValue
|
||||
progressHandler?(0.0, "Downloading \(variant) model...")
|
||||
let files = [
|
||||
"\(variant)/model.safetensors",
|
||||
"\(variant)/config.json",
|
||||
"\(variant)/tokenizer.json",
|
||||
"\(variant)/tokenizer_config.json",
|
||||
]
|
||||
switch registry {
|
||||
case .huggingFace(let hubEndpoint):
|
||||
try await HuggingFaceDownloader.downloadFiles(
|
||||
modelId: modelId,
|
||||
to: cacheDir,
|
||||
files: files,
|
||||
hubEndpoint: hubEndpoint,
|
||||
progressHandler: { progress in
|
||||
progressHandler?(progress, "Downloading...")
|
||||
}
|
||||
)
|
||||
case .modelScope(let baseURL, let revision):
|
||||
try await ModelScopeDownloader.downloadFiles(
|
||||
modelId: modelId,
|
||||
to: cacheDir,
|
||||
files: files,
|
||||
baseURL: baseURL,
|
||||
revision: revision,
|
||||
progressHandler: { progress in
|
||||
progressHandler?(progress, "Downloading...")
|
||||
}
|
||||
)
|
||||
}
|
||||
progressHandler?(1.0, "Downloaded")
|
||||
}
|
||||
|
||||
/// Load from a local directory (no HuggingFace download).
|
||||
public static func fromLocal(
|
||||
directory: URL,
|
||||
progressHandler: ((Double, String) -> Void)? = nil
|
||||
) async throws -> Qwen35MLXChat {
|
||||
let config: Qwen3ChatConfig
|
||||
let configURL = directory.appendingPathComponent("config.json")
|
||||
if FileManager.default.fileExists(atPath: configURL.path) {
|
||||
config = try Qwen3ChatConfig.load(from: configURL)
|
||||
} else {
|
||||
config = .qwen35_08B
|
||||
}
|
||||
|
||||
let tokenizer = ChatTokenizer()
|
||||
try tokenizer.load(from: directory)
|
||||
|
||||
let model = Qwen35MLXModel(config: config)
|
||||
try Qwen35WeightLoader.loadWeights(
|
||||
into: model, from: directory,
|
||||
progressHandler: progressHandler)
|
||||
|
||||
return Qwen35MLXChat(config: config, tokenizer: tokenizer, model: model)
|
||||
}
|
||||
|
||||
// MARK: - State Management
|
||||
|
||||
/// Reset all inference state for a new conversation.
|
||||
public func resetState() {
|
||||
state = .initial(config: config)
|
||||
metrics = Metrics()
|
||||
}
|
||||
|
||||
// MARK: - Generation
|
||||
|
||||
/// Generate a response from chat messages.
|
||||
///
|
||||
/// Encodes the messages using the chat template, prefills the prompt,
|
||||
/// then generates tokens autoregressively until EOS or max tokens.
|
||||
public func generate(
|
||||
messages: [ChatMessage],
|
||||
sampling: ChatSamplingConfig = .default
|
||||
) throws -> String {
|
||||
resetState()
|
||||
|
||||
let promptTokens = ChatTemplate.encode(
|
||||
messages: messages,
|
||||
tokenizer: tokenizer,
|
||||
config: config,
|
||||
enableThinking: false)
|
||||
|
||||
// Prefill
|
||||
let prefillStart = CFAbsoluteTimeGetCurrent()
|
||||
let promptArray = MLXArray(promptTokens.map { Int32($0) })
|
||||
.expandedDimensions(axis: 0)
|
||||
let (prefillLogits, prefillState) = model.forward(inputIds: promptArray, state: state)
|
||||
eval(prefillLogits)
|
||||
state = prefillState
|
||||
|
||||
let prefillMs = (CFAbsoluteTimeGetCurrent() - prefillStart) * 1000
|
||||
metrics.prefillTimeMs = prefillMs
|
||||
metrics.prefillTokens = promptTokens.count
|
||||
|
||||
// Extract last-position logits and sample first token
|
||||
var logits = extractLastPositionLogits(prefillLogits)
|
||||
var generatedTokens: [Int] = []
|
||||
var inThinking = false
|
||||
let thinkBudget = 100
|
||||
|
||||
// Decode loop
|
||||
let decodeStart = CFAbsoluteTimeGetCurrent()
|
||||
|
||||
for _ in 0..<(sampling.maxTokens + thinkBudget) {
|
||||
let nextToken = ChatSampler.sample(
|
||||
logits: logits,
|
||||
config: sampling,
|
||||
previousTokens: promptTokens + generatedTokens)
|
||||
|
||||
if nextToken == config.eosTokenId { break }
|
||||
if nextToken == ChatTemplate.imEndId { break }
|
||||
|
||||
generatedTokens.append(nextToken)
|
||||
|
||||
// Thinking token tracking (handle both Qwen3 and Qwen3.5 token IDs)
|
||||
if nextToken == ChatTemplate.thinkStartId {
|
||||
inThinking = true
|
||||
} else if nextToken == ChatTemplate.thinkEndId {
|
||||
inThinking = false
|
||||
}
|
||||
|
||||
if inThinking && generatedTokens.count > thinkBudget {
|
||||
let thinkEnd = ChatTemplate.thinkEndId
|
||||
generatedTokens.append(thinkEnd)
|
||||
let tokenArr = MLXArray([Int32(thinkEnd)])
|
||||
.expandedDimensions(axis: 0)
|
||||
let (stepLogits, newState) = model.forward(inputIds: tokenArr, state: state)
|
||||
eval(stepLogits)
|
||||
state = newState
|
||||
logits = extractLastPositionLogits(stepLogits)
|
||||
inThinking = false
|
||||
continue
|
||||
}
|
||||
|
||||
let thinkTokens: Set<Int> = [
|
||||
ChatTemplate.thinkStartId, ChatTemplate.thinkEndId
|
||||
]
|
||||
let responseCount = generatedTokens.filter { !thinkTokens.contains($0) }.count
|
||||
if !inThinking && responseCount >= sampling.maxTokens { break }
|
||||
|
||||
// Decode one step
|
||||
let tokenArr = MLXArray([Int32(nextToken)]).expandedDimensions(axis: 0)
|
||||
let (stepLogits, newState) = model.forward(inputIds: tokenArr, state: state)
|
||||
eval(stepLogits)
|
||||
state = newState
|
||||
logits = extractLastPositionLogits(stepLogits)
|
||||
}
|
||||
|
||||
let decodeMs = (CFAbsoluteTimeGetCurrent() - decodeStart) * 1000
|
||||
metrics.decodeTimeMs = decodeMs
|
||||
metrics.decodeTokens = generatedTokens.count
|
||||
|
||||
var memInfo = mach_task_basic_info()
|
||||
var memCount = mach_msg_type_number_t(MemoryLayout<mach_task_basic_info>.size) / 4
|
||||
_ = withUnsafeMutablePointer(to: &memInfo) {
|
||||
$0.withMemoryRebound(to: integer_t.self, capacity: Int(memCount)) {
|
||||
task_info(mach_task_self_, task_flavor_t(MACH_TASK_BASIC_INFO), $0, &memCount)
|
||||
}
|
||||
}
|
||||
let memMB = Double(memInfo.resident_size) / 1024 / 1024
|
||||
let tps = decodeMs > 0 ? Double(generatedTokens.count) / (decodeMs / 1000.0) : 0
|
||||
os_log(.info, log: log,
|
||||
"Generate done: prefill=%.0fms (%d tokens), decode=%.0fms (%d tokens, %.1f tok/s), memory=%.0f MB",
|
||||
metrics.prefillTimeMs, promptTokens.count, decodeMs, generatedTokens.count, tps, memMB)
|
||||
|
||||
let responseTokens = ChatTemplate.stripThinking(from: generatedTokens)
|
||||
return tokenizer.decode(responseTokens)
|
||||
}
|
||||
|
||||
/// Generate a streaming response.
|
||||
public func generateStream(
|
||||
messages: [ChatMessage],
|
||||
sampling: ChatSamplingConfig = .default
|
||||
) -> AsyncThrowingStream<String, Error> {
|
||||
AsyncThrowingStream { continuation in
|
||||
Task {
|
||||
self.resetState()
|
||||
|
||||
let promptTokens = ChatTemplate.encode(
|
||||
messages: messages,
|
||||
tokenizer: self.tokenizer,
|
||||
config: self.config,
|
||||
enableThinking: false)
|
||||
|
||||
let promptArray = MLXArray(promptTokens.map { Int32($0) })
|
||||
.expandedDimensions(axis: 0)
|
||||
let (prefillLogits, prefillState) = self.model.forward(
|
||||
inputIds: promptArray, state: self.state)
|
||||
eval(prefillLogits)
|
||||
self.state = prefillState
|
||||
|
||||
var logits = self.extractLastPositionLogits(prefillLogits)
|
||||
var generatedTokens: [Int] = []
|
||||
var inThinking = false
|
||||
|
||||
for _ in 0..<sampling.maxTokens {
|
||||
let nextToken = ChatSampler.sample(
|
||||
logits: logits,
|
||||
config: sampling,
|
||||
previousTokens: promptTokens + generatedTokens)
|
||||
|
||||
if nextToken == self.config.eosTokenId { break }
|
||||
if nextToken == ChatTemplate.imEndId { break }
|
||||
|
||||
generatedTokens.append(nextToken)
|
||||
|
||||
if nextToken == ChatTemplate.thinkStartId {
|
||||
inThinking = true
|
||||
} else if nextToken == ChatTemplate.thinkEndId {
|
||||
inThinking = false
|
||||
} else if !inThinking,
|
||||
let text = self.tokenizer.decodeToken(nextToken),
|
||||
!self.tokenizer.isSpecialToken(nextToken) {
|
||||
continuation.yield(text)
|
||||
}
|
||||
|
||||
let tokenArr = MLXArray([Int32(nextToken)])
|
||||
.expandedDimensions(axis: 0)
|
||||
let (stepLogits, newState) = self.model.forward(
|
||||
inputIds: tokenArr, state: self.state)
|
||||
eval(stepLogits)
|
||||
self.state = newState
|
||||
logits = self.extractLastPositionLogits(stepLogits)
|
||||
}
|
||||
|
||||
continuation.finish()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Helpers
|
||||
|
||||
/// Extract logits for the last sequence position as a Float array.
|
||||
private func extractLastPositionLogits(_ logits: MLXArray) -> [Float] {
|
||||
let t = logits.dim(1)
|
||||
let lastPos = logits[0, t - 1].asType(.float32) // [vocabSize]
|
||||
eval(lastPos)
|
||||
// Bulk extract all floats at once — do NOT use per-element .item() (248K syncs)
|
||||
let all: [Float] = lastPos.asArray(Float.self)
|
||||
return Array(all.prefix(config.vocabSize))
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Memory Management
|
||||
|
||||
extension Qwen35MLXChat: ModelMemoryManageable {
|
||||
public var isLoaded: Bool { _isLoaded }
|
||||
|
||||
public func unload() {
|
||||
guard _isLoaded else { return }
|
||||
model.clearParameters()
|
||||
state = .initial(config: config)
|
||||
_isLoaded = false
|
||||
}
|
||||
|
||||
public var memoryFootprint: Int {
|
||||
guard _isLoaded else { return 0 }
|
||||
return model.parameterMemoryBytes()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,381 @@
|
||||
import CoreML
|
||||
import Foundation
|
||||
import AudioCommon
|
||||
import os.log
|
||||
|
||||
private let log = OSLog(subsystem: "com.soniqo.qwen3chat", category: "CoreML")
|
||||
|
||||
/// CoreML-based Qwen3.5-0.8B chat for iOS Neural Engine.
|
||||
///
|
||||
/// Uses two CoreML models:
|
||||
/// - `embedding.mlmodelc` — token ID → embedding vector
|
||||
/// - `decoder.mlmodelc` — autoregressive transformer with MLState
|
||||
///
|
||||
/// All DeltaNet recurrent states and GatedAttention KV caches are managed
|
||||
/// by CoreML's MLState API — no manual cache tracking needed.
|
||||
public final class Qwen35CoreMLChat: @unchecked Sendable {
|
||||
public static let defaultModelId = "aufklarer/Qwen3.5-0.8B-Chat-CoreML"
|
||||
|
||||
private let embeddingModel: MLModel
|
||||
private let decoderModel: MLModel
|
||||
private var decoderState: MLState
|
||||
public let config: Qwen3ChatConfig
|
||||
public let tokenizer: ChatTokenizer
|
||||
private var position: Int = 0
|
||||
private let maxSeqLen: Int
|
||||
|
||||
/// Quantization variant. Only INT8 available (INT4 removed — CoreML dequantization issues).
|
||||
public enum Quantization: String { case int8 }
|
||||
|
||||
// MARK: - Metrics
|
||||
|
||||
private var _prefillMs: Double = 0
|
||||
private var _decodeMs: Double = 0
|
||||
private var _decodeTokens: Int = 0
|
||||
|
||||
public var lastMetrics: (tokensPerSec: Double, prefillMs: Double, decodeMs: Double, msPerToken: Double) {
|
||||
let tps = _decodeMs > 0 ? Double(_decodeTokens) / (_decodeMs / 1000.0) : 0
|
||||
let mpt = _decodeTokens > 0 ? _decodeMs / Double(_decodeTokens) : 0
|
||||
return (tps, _prefillMs, _decodeMs, mpt)
|
||||
}
|
||||
|
||||
/// Current process memory in MB.
|
||||
private static var memoryMB: Double {
|
||||
var info = mach_task_basic_info()
|
||||
var count = mach_msg_type_number_t(MemoryLayout<mach_task_basic_info>.size) / 4
|
||||
let result = withUnsafeMutablePointer(to: &info) {
|
||||
$0.withMemoryRebound(to: integer_t.self, capacity: Int(count)) {
|
||||
task_info(mach_task_self_, task_flavor_t(MACH_TASK_BASIC_INFO), $0, &count)
|
||||
}
|
||||
}
|
||||
guard result == KERN_SUCCESS else { return 0 }
|
||||
return Double(info.resident_size) / 1024 / 1024
|
||||
}
|
||||
|
||||
// MARK: - Init
|
||||
|
||||
private init(embedding: MLModel, decoder: MLModel, state: MLState,
|
||||
config: Qwen3ChatConfig, tokenizer: ChatTokenizer, maxSeqLen: Int) {
|
||||
self.embeddingModel = embedding
|
||||
self.decoderModel = decoder
|
||||
self.decoderState = state
|
||||
self.config = config
|
||||
self.tokenizer = tokenizer
|
||||
self.maxSeqLen = maxSeqLen
|
||||
}
|
||||
|
||||
// MARK: - Factory
|
||||
|
||||
/// Load from HuggingFace.
|
||||
public static func fromPretrained(
|
||||
modelId: String = defaultModelId,
|
||||
quantization: Quantization = .int8,
|
||||
computeUnits: MLComputeUnits = .cpuAndNeuralEngine,
|
||||
cacheDir: URL? = nil,
|
||||
offlineMode: Bool = false,
|
||||
progressHandler: ((Double, String) -> Void)? = nil
|
||||
) async throws -> Qwen35CoreMLChat {
|
||||
let cacheDir = try cacheDir ?? HuggingFaceDownloader.getCacheDirectory(for: modelId)
|
||||
let variant = quantization.rawValue
|
||||
|
||||
progressHandler?(0.05, "Downloading \(variant) model...")
|
||||
// Fetch the pre-compiled ``.mlmodelc`` bundle only. On-device
|
||||
// ``MLModel.compileModel`` drifts per runtime, and the legacy
|
||||
// ``.mlpackage`` internals (``*.mlmodel`` / ``Manifest.json``) would
|
||||
// force that code path. Users with a stale ``.mlpackage`` cache
|
||||
// transparently re-download because ``.mlmodelc`` is missing.
|
||||
try await HuggingFaceDownloader.downloadWeights(
|
||||
modelId: modelId,
|
||||
to: cacheDir,
|
||||
additionalFiles: [
|
||||
"\(variant)/*.json",
|
||||
"\(variant)/embedding.mlmodelc/**",
|
||||
"\(variant)/decoder.mlmodelc/**",
|
||||
],
|
||||
offlineMode: offlineMode,
|
||||
progressHandler: { p in progressHandler?(p * 0.5, "Downloading...") }
|
||||
)
|
||||
|
||||
let variantDir = cacheDir.appendingPathComponent(variant)
|
||||
return try await fromLocal(directory: variantDir, computeUnits: computeUnits,
|
||||
progressHandler: progressHandler)
|
||||
}
|
||||
|
||||
/// Load from a local directory.
|
||||
public static func fromLocal(
|
||||
directory: URL,
|
||||
computeUnits: MLComputeUnits = .cpuAndNeuralEngine,
|
||||
progressHandler: ((Double, String) -> Void)? = nil
|
||||
) async throws -> Qwen35CoreMLChat {
|
||||
progressHandler?(0.5, "Loading config...")
|
||||
|
||||
// Debug: list directory contents to diagnose missing file issues
|
||||
os_log(.info, log: log, "Loading from directory: %{public}@", directory.path)
|
||||
if let contents = try? FileManager.default.contentsOfDirectory(at: directory, includingPropertiesForKeys: nil) {
|
||||
for item in contents {
|
||||
os_log(.info, log: log, " %{public}@", item.lastPathComponent)
|
||||
}
|
||||
} else {
|
||||
os_log(.error, log: log, "Cannot list directory: %{public}@", directory.path)
|
||||
}
|
||||
|
||||
// Use built-in config — chat_config.json from CoreML conversion has a different schema
|
||||
let config = Qwen3ChatConfig.qwen35_08B
|
||||
os_log(.info, log: log, "Using built-in Qwen3.5-0.8B config")
|
||||
|
||||
progressHandler?(0.55, "Loading tokenizer...")
|
||||
os_log(.info, log: log, "Loading tokenizer from: %{public}@", directory.resolvingSymlinksInPath().path)
|
||||
let tokenizer = ChatTokenizer()
|
||||
try tokenizer.load(from: directory.resolvingSymlinksInPath())
|
||||
|
||||
let memBefore = memoryMB
|
||||
os_log(.info, log: log, "Loading CoreML models, memory before: %.0f MB", memBefore)
|
||||
|
||||
progressHandler?(0.6, "Compiling embedding...")
|
||||
let embModel: MLModel
|
||||
do {
|
||||
embModel = try await loadModel(named: "embedding", from: directory, computeUnits: computeUnits)
|
||||
os_log(.info, log: log, "Embedding loaded, memory: %.0f MB", memoryMB)
|
||||
} catch {
|
||||
os_log(.error, log: log, "Embedding load FAILED: %{public}@", error.localizedDescription)
|
||||
throw error
|
||||
}
|
||||
|
||||
progressHandler?(0.75, "Compiling decoder...")
|
||||
let decModel: MLModel
|
||||
do {
|
||||
decModel = try await loadModel(named: "decoder", from: directory, computeUnits: computeUnits)
|
||||
os_log(.info, log: log, "Decoder loaded, memory: %.0f MB", memoryMB)
|
||||
} catch {
|
||||
os_log(.error, log: log, "Decoder load FAILED: %{public}@", error.localizedDescription)
|
||||
throw error
|
||||
}
|
||||
|
||||
let state = decModel.makeState()
|
||||
let maxSeq = config.maxSeqLen
|
||||
|
||||
os_log(.info, log: log, "Model ready, total memory: %.0f MB (delta: +%.0f MB)",
|
||||
memoryMB, memoryMB - memBefore)
|
||||
progressHandler?(1.0, "Ready")
|
||||
return Qwen35CoreMLChat(
|
||||
embedding: embModel, decoder: decModel, state: state,
|
||||
config: config, tokenizer: tokenizer, maxSeqLen: maxSeq)
|
||||
}
|
||||
|
||||
// MARK: - Model Loading Helpers
|
||||
|
||||
private static func loadModel(
|
||||
named name: String, from dir: URL, computeUnits: MLComputeUnits
|
||||
) async throws -> MLModel {
|
||||
let compiledURL = dir.appendingPathComponent("\(name).mlmodelc")
|
||||
guard FileManager.default.fileExists(atPath: compiledURL.path) else {
|
||||
os_log(.error, log: log,
|
||||
"Model not found: %{public}@.mlmodelc in %{public}@",
|
||||
name, dir.path)
|
||||
throw ChatModelError.modelNotFound(dir)
|
||||
}
|
||||
|
||||
let mlConfig = MLModelConfiguration()
|
||||
mlConfig.computeUnits = CoreMLComputeUnitsResolver.resolved(default: computeUnits)
|
||||
return try await MLModel.load(contentsOf: compiledURL, configuration: mlConfig)
|
||||
}
|
||||
|
||||
// MARK: - Generation
|
||||
|
||||
/// Reset state for a new conversation.
|
||||
public func resetState() {
|
||||
decoderState = decoderModel.makeState()
|
||||
position = 0
|
||||
_prefillMs = 0; _decodeMs = 0; _decodeTokens = 0
|
||||
}
|
||||
|
||||
/// Generate a response from chat messages.
|
||||
public func generate(
|
||||
messages: [ChatMessage],
|
||||
sampling: ChatSamplingConfig = .default
|
||||
) throws -> String {
|
||||
resetState()
|
||||
|
||||
let promptTokens = ChatTemplate.encode(
|
||||
messages: messages, tokenizer: tokenizer,
|
||||
config: config, enableThinking: false)
|
||||
|
||||
// Prefill: feed all prompt tokens one at a time
|
||||
let prefillStart = CFAbsoluteTimeGetCurrent()
|
||||
var lastLogits: [Float] = []
|
||||
|
||||
for token in promptTokens {
|
||||
lastLogits = try forwardStep(tokenId: token)
|
||||
}
|
||||
_prefillMs = (CFAbsoluteTimeGetCurrent() - prefillStart) * 1000
|
||||
|
||||
// Decode loop
|
||||
let decodeStart = CFAbsoluteTimeGetCurrent()
|
||||
var generatedTokens: [Int] = []
|
||||
|
||||
for _ in 0..<sampling.maxTokens {
|
||||
let nextToken = ChatSampler.sample(
|
||||
logits: lastLogits, config: sampling,
|
||||
previousTokens: promptTokens + generatedTokens)
|
||||
|
||||
if nextToken == config.eosTokenId { break }
|
||||
if nextToken == ChatTemplate.imEndId { break }
|
||||
|
||||
generatedTokens.append(nextToken)
|
||||
if generatedTokens.count <= 10 {
|
||||
print("[CoreML] token[\(generatedTokens.count-1)]: \(nextToken) '\(tokenizer.decodeToken(nextToken) ?? "?")'")
|
||||
}
|
||||
lastLogits = try forwardStep(tokenId: nextToken)
|
||||
}
|
||||
|
||||
_decodeMs = (CFAbsoluteTimeGetCurrent() - decodeStart) * 1000
|
||||
_decodeTokens = generatedTokens.count
|
||||
|
||||
let tps = _decodeMs > 0 ? Double(_decodeTokens) / (_decodeMs / 1000.0) : 0
|
||||
os_log(.info, log: log,
|
||||
"Generate done: prefill=%.0fms (%d tokens), decode=%.0fms (%d tokens, %.1f tok/s), memory=%.0f MB",
|
||||
_prefillMs, promptTokens.count, _decodeMs, _decodeTokens, tps, Self.memoryMB)
|
||||
|
||||
let responseTokens = ChatTemplate.stripThinking(from: generatedTokens)
|
||||
let responseText = tokenizer.decode(responseTokens)
|
||||
os_log(.info, log: log, "Response (%d tokens → %d after strip): '%{public}@'",
|
||||
generatedTokens.count, responseTokens.count, responseText)
|
||||
return responseText
|
||||
}
|
||||
|
||||
/// Generate a streaming response.
|
||||
public func generateStream(
|
||||
messages: [ChatMessage],
|
||||
sampling: ChatSamplingConfig = .default
|
||||
) -> AsyncThrowingStream<String, Error> {
|
||||
AsyncThrowingStream { continuation in
|
||||
Task {
|
||||
do {
|
||||
self.resetState()
|
||||
let promptTokens = ChatTemplate.encode(
|
||||
messages: messages, tokenizer: self.tokenizer,
|
||||
config: self.config, enableThinking: false)
|
||||
|
||||
var lastLogits: [Float] = []
|
||||
for token in promptTokens {
|
||||
lastLogits = try self.forwardStep(tokenId: token)
|
||||
}
|
||||
|
||||
var generatedTokens: [Int] = []
|
||||
var inThinking = false
|
||||
let thinkBudget = 100
|
||||
let thinkTokens: Set<Int> = [
|
||||
ChatTemplate.thinkStartId, ChatTemplate.thinkEndId
|
||||
]
|
||||
|
||||
for _ in 0..<(sampling.maxTokens + thinkBudget) {
|
||||
let nextToken = ChatSampler.sample(
|
||||
logits: lastLogits, config: sampling,
|
||||
previousTokens: promptTokens + generatedTokens)
|
||||
|
||||
if nextToken == self.config.eosTokenId { break }
|
||||
if nextToken == ChatTemplate.imEndId { break }
|
||||
|
||||
generatedTokens.append(nextToken)
|
||||
|
||||
if nextToken == ChatTemplate.thinkStartId { inThinking = true }
|
||||
else if nextToken == ChatTemplate.thinkEndId { inThinking = false }
|
||||
else if !inThinking,
|
||||
let text = self.tokenizer.decodeToken(nextToken),
|
||||
!self.tokenizer.isSpecialToken(nextToken) {
|
||||
continuation.yield(text)
|
||||
}
|
||||
|
||||
// Force-end thinking if budget exceeded
|
||||
if inThinking && generatedTokens.count > thinkBudget {
|
||||
generatedTokens.append(ChatTemplate.thinkEndId)
|
||||
lastLogits = try self.forwardStep(tokenId: ChatTemplate.thinkEndId)
|
||||
inThinking = false
|
||||
continue
|
||||
}
|
||||
|
||||
// Only count non-thinking tokens against maxTokens
|
||||
let responseCount = generatedTokens.filter { !thinkTokens.contains($0) }.count
|
||||
if !inThinking && responseCount >= sampling.maxTokens { break }
|
||||
|
||||
lastLogits = try self.forwardStep(tokenId: nextToken)
|
||||
}
|
||||
continuation.finish()
|
||||
} catch {
|
||||
continuation.finish(throwing: error)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Single Step
|
||||
|
||||
/// Run one decoder step: token ID → logits.
|
||||
private func forwardStep(tokenId: Int) throws -> [Float] {
|
||||
// Embedding lookup
|
||||
let tokenInput = try MLMultiArray(shape: [1, 1], dataType: .int32)
|
||||
tokenInput[0] = NSNumber(value: Int32(tokenId))
|
||||
|
||||
let embFeatures = try MLDictionaryFeatureProvider(dictionary: [
|
||||
"token_id": MLFeatureValue(multiArray: tokenInput)
|
||||
])
|
||||
let embResult = try embeddingModel.prediction(from: embFeatures)
|
||||
guard let embedding = embResult.featureValue(for: "embedding")?.multiArrayValue else {
|
||||
throw ChatModelError.inferenceFailed("Embedding output missing")
|
||||
}
|
||||
|
||||
// Build attention mask: 0 for positions ≤ current, -FLT_MAX for future
|
||||
let mask = try MLMultiArray(shape: [1, 1, 1, maxSeqLen as NSNumber], dataType: .float32)
|
||||
let maskPtr = mask.dataPointer.bindMemory(to: Float.self, capacity: maxSeqLen)
|
||||
for i in 0..<maxSeqLen {
|
||||
maskPtr[i] = i <= position ? 0 : -Float.greatestFiniteMagnitude
|
||||
}
|
||||
|
||||
// Position
|
||||
let posArray = try MLMultiArray(shape: [1], dataType: .int32)
|
||||
posArray[0] = NSNumber(value: Int32(position))
|
||||
|
||||
// Decoder forward
|
||||
let decoderInput = try MLDictionaryFeatureProvider(dictionary: [
|
||||
"input_embeds": MLFeatureValue(multiArray: embedding),
|
||||
"position": MLFeatureValue(multiArray: posArray),
|
||||
"attention_mask": MLFeatureValue(multiArray: mask),
|
||||
])
|
||||
|
||||
let result = try decoderModel.prediction(
|
||||
from: decoderInput, using: decoderState)
|
||||
|
||||
position += 1
|
||||
|
||||
// Extract logits
|
||||
guard let logitsArray = result.featureValue(for: "logits")?.multiArrayValue else {
|
||||
throw ChatModelError.inferenceFailed("Decoder output missing")
|
||||
}
|
||||
|
||||
let vocabSize = config.vocabSize
|
||||
|
||||
// Debug: log dtype and shape on first few calls
|
||||
if position <= 3 {
|
||||
print("[CoreML] pos=\(position) logits shape=\(logitsArray.shape) dtype=\(logitsArray.dataType.rawValue) count=\(logitsArray.count)")
|
||||
}
|
||||
|
||||
// Handle both Float16 and Float32 output
|
||||
let logits: [Float]
|
||||
if logitsArray.dataType == .float32 {
|
||||
let ptr = logitsArray.dataPointer.bindMemory(to: Float.self, capacity: vocabSize)
|
||||
logits = Array(UnsafeBufferPointer(start: ptr, count: vocabSize))
|
||||
} else {
|
||||
let ptr = logitsArray.dataPointer.bindMemory(to: Float16.self, capacity: vocabSize)
|
||||
logits = (0..<vocabSize).map { Float(ptr[$0]) }
|
||||
}
|
||||
return logits
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Memory Management
|
||||
|
||||
extension Qwen35CoreMLChat: ModelMemoryManageable {
|
||||
public var isLoaded: Bool { true }
|
||||
public func unload() { /* CoreML manages its own memory */ }
|
||||
public var memoryFootprint: Int { 500 * 1024 * 1024 } // ~500 MB estimate
|
||||
}
|
||||
@@ -0,0 +1,710 @@
|
||||
import Foundation
|
||||
import MLXCommon
|
||||
import MLX
|
||||
import MLXNN
|
||||
import MLXFast
|
||||
|
||||
// MARK: - DeltaNet Linear Attention
|
||||
|
||||
/// DeltaNet linear attention layer for Qwen3.5 hybrid model.
|
||||
///
|
||||
/// Uses linear attention (no softmax) with a recurrent state matrix S of shape [B, H, D, D].
|
||||
/// The state evolves per-token: S = alpha * S + beta * (v outer k), where alpha/beta
|
||||
/// are learned per-head scalar gates derived from the input via softplus/sigmoid.
|
||||
///
|
||||
/// A causal conv1d (kernel=4) provides short-range local context before the attention.
|
||||
/// Output is gated: `o_proj(attention_output * silu(z))` where both are 2*hiddenSize = numHeads*headDim.
|
||||
///
|
||||
/// Weight shapes (HuggingFace safetensors):
|
||||
/// - in_proj_qkv.weight: [6144, 1024] (3 * 16 * 128)
|
||||
/// - in_proj_z.weight: [2048, 1024] (2 * hiddenSize, for gate)
|
||||
/// - in_proj_b.weight: [16, 1024] (beta gate, per head)
|
||||
/// - in_proj_a.weight: [16, 1024] (alpha gate, per head)
|
||||
/// - conv1d.weight: [6144, 4, 1] (depthwise causal conv, MLX [C, K, 1] format)
|
||||
/// - dt_bias: [16] (time-step bias)
|
||||
/// - A_log: [16] (log of decay rate)
|
||||
/// - norm.weight: [128] (per-head RMSNorm)
|
||||
/// - out_proj.weight: [1024, 2048] (gated output projection)
|
||||
public final class DeltaNetLayer: Module {
|
||||
let numHeads: Int
|
||||
let headDim: Int
|
||||
let hiddenSize: Int
|
||||
let convKernel: Int
|
||||
let qkvDim: Int
|
||||
|
||||
@ModuleInfo(key: "in_proj_qkv") var inProjQKV: QuantizedLinear
|
||||
@ModuleInfo(key: "in_proj_z") var inProjZ: QuantizedLinear
|
||||
@ModuleInfo(key: "in_proj_b") var inProjB: QuantizedLinear
|
||||
@ModuleInfo(key: "in_proj_a") var inProjA: QuantizedLinear
|
||||
|
||||
/// Conv1d weight: [C, 1, K] depthwise convolution applied to QKV before attention.
|
||||
/// Stored under a flat key to avoid nested key path issues in MLXNN module traversal.
|
||||
/// Weight loading applies this directly via `layer.convWeight = ...`.
|
||||
@ParameterInfo(key: "conv1d_weight") var convWeight: MLXArray
|
||||
|
||||
@ParameterInfo(key: "dt_bias") var dtBias: MLXArray
|
||||
@ParameterInfo(key: "A_log") var aLog: MLXArray
|
||||
|
||||
@ModuleInfo var norm: RMSNorm
|
||||
|
||||
@ModuleInfo(key: "out_proj") var outProj: QuantizedLinear
|
||||
|
||||
public init(config: Qwen3ChatConfig) {
|
||||
self.numHeads = config.linearNumKeyHeads ?? 16
|
||||
self.headDim = config.linearKeyHeadDim ?? 128
|
||||
self.hiddenSize = config.hiddenSize
|
||||
self.convKernel = config.linearConvKernelDim ?? 4
|
||||
self.qkvDim = 3 * numHeads * headDim
|
||||
|
||||
let groupSize = 64
|
||||
let bits = 4
|
||||
self._inProjQKV = ModuleInfo(wrappedValue: QuantizedLinear(hiddenSize, qkvDim, bias: false, groupSize: groupSize, bits: bits))
|
||||
self._inProjZ = ModuleInfo(wrappedValue: QuantizedLinear(hiddenSize, 2 * hiddenSize, bias: false, groupSize: groupSize, bits: bits))
|
||||
self._inProjB = ModuleInfo(wrappedValue: QuantizedLinear(hiddenSize, numHeads, bias: false, groupSize: groupSize, bits: bits))
|
||||
self._inProjA = ModuleInfo(wrappedValue: QuantizedLinear(hiddenSize, numHeads, bias: false, groupSize: groupSize, bits: bits))
|
||||
|
||||
self._convWeight = ParameterInfo(
|
||||
wrappedValue: MLXArray.zeros([qkvDim, convKernel, 1]))
|
||||
|
||||
self._dtBias = ParameterInfo(wrappedValue: MLXArray.zeros([numHeads]))
|
||||
self._aLog = ParameterInfo(wrappedValue: MLXArray.zeros([numHeads]))
|
||||
|
||||
self._norm = ModuleInfo(
|
||||
wrappedValue: RMSNorm(dimensions: headDim, eps: Float(config.rmsNormEps)))
|
||||
|
||||
self._outProj = ModuleInfo(wrappedValue: QuantizedLinear(2 * hiddenSize, hiddenSize, bias: false, groupSize: groupSize, bits: bits))
|
||||
|
||||
super.init()
|
||||
}
|
||||
|
||||
/// Recurrent state for a DeltaNet layer.
|
||||
public struct State {
|
||||
/// Recurrent state matrix [B, H, D, D]
|
||||
var s: MLXArray
|
||||
/// Conv1d ring buffer [B, C, K-1] storing last K-1 inputs
|
||||
var convState: MLXArray
|
||||
|
||||
public static func initial(
|
||||
batchSize: Int, numHeads: Int, headDim: Int,
|
||||
qkvDim: Int, convKernel: Int, dtype: DType = .float32
|
||||
) -> State {
|
||||
State(
|
||||
s: MLXArray.zeros([batchSize, numHeads, headDim, headDim], dtype: dtype),
|
||||
convState: MLXArray.zeros([batchSize, qkvDim, convKernel - 1], dtype: dtype)
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Forward pass processing a sequence of tokens.
|
||||
///
|
||||
/// Implements the gated delta rule recurrence (reference: mlx-lm/gated_delta.py):
|
||||
/// 1. Decay: S = g * S
|
||||
/// 2. Error: kv_mem = (S * k).sum(-1); delta = (v - kv_mem) * beta
|
||||
/// 3. Update: S = S + k * delta
|
||||
/// 4. Output: y = (S * q).sum(-1)
|
||||
///
|
||||
/// - Parameters:
|
||||
/// - x: Input hidden states [B, T, hiddenSize]
|
||||
/// - state: Previous recurrent state (nil for first call)
|
||||
/// - Returns: (output [B, T, hiddenSize], updated state)
|
||||
public func callAsFunction(_ x: MLXArray, state: State? = nil) -> (MLXArray, State) {
|
||||
let b = x.dim(0)
|
||||
let t = x.dim(1)
|
||||
|
||||
// Project inputs (separate projections matching HuggingFace weight format)
|
||||
let qkvRaw = inProjQKV(x) // [B, T, 3*H*D=6144]
|
||||
let zRaw = inProjZ(x) // [B, T, 2*hiddenSize=2048]
|
||||
let bRaw = inProjB(x) // [B, T, H=16]
|
||||
let aRaw = inProjA(x) // [B, T, H=16]
|
||||
|
||||
|
||||
// Causal conv1d on QKV only (not Z, B, A)
|
||||
let prevConvState: MLXArray
|
||||
if let s = state {
|
||||
prevConvState = s.convState
|
||||
} else {
|
||||
prevConvState = MLXArray.zeros([b, qkvDim, convKernel - 1], dtype: x.dtype)
|
||||
}
|
||||
|
||||
let qkvTransposed = qkvRaw.transposed(0, 2, 1) // [B, C, T]
|
||||
let padded = concatenated([prevConvState, qkvTransposed], axis: 2) // [B, C, T+K-1]
|
||||
|
||||
// Save new conv state (last K-1 columns)
|
||||
let totalLen = padded.dim(2)
|
||||
let newConvState = padded[0..., 0..., (totalLen - convKernel + 1)...]
|
||||
|
||||
// Apply depthwise causal conv1d + SiLU
|
||||
let qkvConv = depthwiseConv1dCausal(padded, outputLen: t)
|
||||
let qkvActivated = silu(qkvConv.transposed(0, 2, 1)) // [B, T, C]
|
||||
|
||||
// Split into Q, K, V — each [B, T, H, D]
|
||||
let hd = numHeads * headDim
|
||||
var q = qkvActivated[0..., 0..., ..<hd].reshaped(b, t, numHeads, headDim)
|
||||
var k = qkvActivated[0..., 0..., hd..<(2 * hd)].reshaped(b, t, numHeads, headDim)
|
||||
let v = qkvActivated[0..., 0..., (2 * hd)...].reshaped(b, t, numHeads, headDim)
|
||||
|
||||
// Q/K normalization (reference: inv_scale * rms_norm, different scaling for Q and K)
|
||||
// rms_norm(x, None, eps) = x / sqrt(mean(x^2) + eps)
|
||||
// q = inv_scale^2 * rms_norm(q) where inv_scale = head_dim^(-0.5)
|
||||
// k = inv_scale * rms_norm(k)
|
||||
let invScale = Float(1.0) / sqrt(Float(headDim))
|
||||
q = MLXArray(invScale * invScale) * rmsNormNoWeight(q)
|
||||
k = MLXArray(invScale) * rmsNormNoWeight(k)
|
||||
|
||||
// Compute gating: g = exp(-exp(A_log) * softplus(a + dt_bias))
|
||||
let g = computeDecayGate(aRaw: aRaw) // [B, T, H]
|
||||
// beta = sigmoid(b_raw) (independent learned gate, NOT 1-alpha)
|
||||
let beta = sigmoid(bRaw) // [B, T, H]
|
||||
|
||||
// Sequential gated delta rule recurrence
|
||||
var currentS: MLXArray
|
||||
if let s = state {
|
||||
currentS = s.s
|
||||
} else {
|
||||
currentS = MLXArray.zeros([b, numHeads, headDim, headDim], dtype: x.dtype)
|
||||
}
|
||||
|
||||
var outputSteps: [MLXArray] = []
|
||||
outputSteps.reserveCapacity(t)
|
||||
|
||||
for step in 0..<t {
|
||||
// Extract step: [B, H, D] or [B, H]
|
||||
let qStep = q[0..., step..<(step + 1), 0..., 0...].squeezed(axis: 1) // [B, H, D]
|
||||
let kStep = k[0..., step..<(step + 1), 0..., 0...].squeezed(axis: 1) // [B, H, D]
|
||||
let vStep = v[0..., step..<(step + 1), 0..., 0...].squeezed(axis: 1) // [B, H, D]
|
||||
let gStep = g[0..., step..<(step + 1), 0...].squeezed(axis: 1) // [B, H]
|
||||
let betaStep = beta[0..., step..<(step + 1), 0...].squeezed(axis: 1) // [B, H]
|
||||
|
||||
// 1. Decay: S = g * S (g is scalar per-head: [B, H, 1, 1])
|
||||
let decay = gStep.reshaped(b, numHeads, 1, 1)
|
||||
currentS = currentS * decay
|
||||
|
||||
// 2. Error correction:
|
||||
// kv_mem = (S * k[..., None, :]).sum(-1) → [B, H, Dv]
|
||||
// delta = (v - kv_mem) * beta[..., None] → [B, H, Dv]
|
||||
let kExpanded = kStep.expandedDimensions(axis: -2) // [B, H, 1, Dk]
|
||||
let kvMem = (currentS * kExpanded).sum(axis: -1) // [B, H, Dv]
|
||||
let delta = (vStep - kvMem) * betaStep.expandedDimensions(axis: -1) // [B, H, Dv]
|
||||
|
||||
// 3. Update: S = S + k[..., None, :] * delta[..., None]
|
||||
// k: [B, H, Dk] → [B, H, 1, Dk], delta: [B, H, Dv] → [B, H, Dv, 1]
|
||||
currentS = currentS + kExpanded * delta.expandedDimensions(axis: -1)
|
||||
|
||||
// 4. Output: y = (S * q[..., None, :]).sum(-1) → [B, H, Dv]
|
||||
let qExpanded = qStep.expandedDimensions(axis: -2) // [B, H, 1, Dk]
|
||||
let oStep = (currentS * qExpanded).sum(axis: -1) // [B, H, Dv]
|
||||
outputSteps.append(oStep)
|
||||
}
|
||||
|
||||
// Stack: [B, T, H, D]
|
||||
let output = stacked(outputSteps, axis: 1)
|
||||
|
||||
// RMSNormGated: norm(output) * silu(z)
|
||||
// z has shape [B, T, 2*hiddenSize=2048], reshape to [B, T, H, D] for per-head norm
|
||||
let zReshaped = zRaw.reshaped(b, t, numHeads, headDim)
|
||||
let normedOutput = norm(output) // per-head RMSNorm, [B, T, H, D]
|
||||
let gated = normedOutput * silu(zReshaped) // [B, T, H, D]
|
||||
|
||||
// Reshape to [B, T, H*D=2048] and project to hiddenSize
|
||||
let result = outProj(gated.reshaped(b, t, numHeads * headDim)) // [B, T, 1024]
|
||||
|
||||
return (result, State(s: currentS, convState: newConvState))
|
||||
}
|
||||
|
||||
/// Compute decay gate: g = exp(-exp(A_log) * softplus(a + dt_bias))
|
||||
private func computeDecayGate(aRaw: MLXArray) -> MLXArray {
|
||||
let a = aRaw + dtBias.reshaped(1, 1, numHeads)
|
||||
let dt = softplus(a)
|
||||
let negExpA = -exp(aLog.asType(.float32)).reshaped(1, 1, numHeads)
|
||||
return exp(negExpA * dt.asType(.float32)).asType(aRaw.dtype)
|
||||
}
|
||||
|
||||
/// RMS normalization without learnable weight (used for Q/K normalization).
|
||||
private func rmsNormNoWeight(_ x: MLXArray) -> MLXArray {
|
||||
let meanSq = (x * x).mean(axis: -1, keepDims: true)
|
||||
return x * rsqrt(meanSq + MLXArray(Float(1e-6)))
|
||||
}
|
||||
|
||||
// MARK: - Depthwise Conv1d
|
||||
|
||||
/// Depthwise causal conv1d via unfolding + element-wise multiply + sum.
|
||||
/// - Parameter input: [B, C, T+K-1] (pre-padded with conv state)
|
||||
/// - Parameter outputLen: number of output time steps T
|
||||
/// - Returns: [B, C, T]
|
||||
private func depthwiseConv1dCausal(_ input: MLXArray, outputLen: Int) -> MLXArray {
|
||||
let c = input.dim(1)
|
||||
let k = convKernel
|
||||
|
||||
// Unfold: gather windows of size K for each output position
|
||||
var windows: [MLXArray] = []
|
||||
windows.reserveCapacity(outputLen)
|
||||
for t in 0..<outputLen {
|
||||
windows.append(input[0..., 0..., t..<(t + k)]) // [B, C, K]
|
||||
}
|
||||
let unfolded = stacked(windows, axis: 2) // [B, C, T, K]
|
||||
|
||||
// Kernel: [C, K, 1] -> squeeze axis 2 -> [C, K] -> [1, C, 1, K]
|
||||
let kernelBcast = convWeight.squeezed(axis: 2).reshaped(1, c, 1, k)
|
||||
|
||||
return (unfolded * kernelBcast).sum(axis: -1) // [B, C, T]
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Softplus
|
||||
|
||||
private func softplus(_ x: MLXArray) -> MLXArray {
|
||||
// Numerically stable: for large x, softplus(x) ~ x
|
||||
MLX.where(x .> MLXArray(Float(20.0)), x, log(1 + exp(x)))
|
||||
}
|
||||
|
||||
// MARK: - GatedAttention (Full Attention)
|
||||
|
||||
/// GatedAttention layer for Qwen3.5 hybrid model.
|
||||
///
|
||||
/// Standard multi-head attention with:
|
||||
/// - GQA: 8 query heads, 2 KV heads, head_dim=256
|
||||
/// - Partial RoPE: only first 25% of head_dim (64 dims) get rotary encoding
|
||||
/// - QK norm: RMSNorm applied per-head to Q and K before RoPE
|
||||
/// - Gated output: q_proj produces [Q; gate], both of dim numQHeads*headDim.
|
||||
/// After attention, output is element-wise multiplied with silu(gate),
|
||||
/// then projected through o_proj.
|
||||
///
|
||||
/// Weight shapes:
|
||||
/// - q_proj: [4096, 1024] = [2 * numQHeads * headDim, hiddenSize] (Q + gate)
|
||||
/// - k_proj: [512, 1024] = [numKVHeads * headDim, hiddenSize]
|
||||
/// - v_proj: [512, 1024] = [numKVHeads * headDim, hiddenSize]
|
||||
/// - o_proj: [1024, 2048] = [hiddenSize, numQHeads * headDim]
|
||||
/// - q_norm: [256] = [headDim]
|
||||
/// - k_norm: [256] = [headDim]
|
||||
public final class GatedAttentionLayer: Module {
|
||||
let numQHeads: Int
|
||||
let numKVHeads: Int
|
||||
let headDim: Int
|
||||
let hiddenSize: Int
|
||||
let scale: Float
|
||||
let ropeDims: Int // partial RoPE dimensions
|
||||
|
||||
@ModuleInfo(key: "q_proj") var qProj: QuantizedLinear
|
||||
@ModuleInfo(key: "k_proj") var kProj: QuantizedLinear
|
||||
@ModuleInfo(key: "v_proj") var vProj: QuantizedLinear
|
||||
@ModuleInfo(key: "o_proj") var oProj: QuantizedLinear
|
||||
|
||||
@ModuleInfo(key: "q_norm") var qNorm: RMSNorm
|
||||
@ModuleInfo(key: "k_norm") var kNorm: RMSNorm
|
||||
|
||||
let rope: MLXNN.RoPE
|
||||
|
||||
public init(config: Qwen3ChatConfig) {
|
||||
self.numQHeads = config.numAttentionHeads // 8
|
||||
self.numKVHeads = config.numKeyValueHeads // 2
|
||||
self.headDim = config.headDim // 256
|
||||
self.hiddenSize = config.hiddenSize // 1024
|
||||
self.scale = 1.0 / sqrt(Float(headDim))
|
||||
|
||||
let factor = config.partialRotaryFactor ?? 0.25
|
||||
self.ropeDims = Int(Double(headDim) * factor) // 64
|
||||
|
||||
let groupSize = 64
|
||||
let bits = 4
|
||||
let qDim = numQHeads * headDim // 2048
|
||||
|
||||
// q_proj outputs 2 * qDim (Q + gate)
|
||||
self._qProj = ModuleInfo(wrappedValue: QuantizedLinear(
|
||||
hiddenSize, 2 * qDim, bias: false,
|
||||
groupSize: groupSize, bits: bits))
|
||||
self._kProj = ModuleInfo(wrappedValue: QuantizedLinear(
|
||||
hiddenSize, numKVHeads * headDim, bias: false,
|
||||
groupSize: groupSize, bits: bits))
|
||||
self._vProj = ModuleInfo(wrappedValue: QuantizedLinear(
|
||||
hiddenSize, numKVHeads * headDim, bias: false,
|
||||
groupSize: groupSize, bits: bits))
|
||||
// o_proj: qDim -> hiddenSize (after gating reduces 2*qDim to qDim)
|
||||
self._oProj = ModuleInfo(wrappedValue: QuantizedLinear(
|
||||
qDim, hiddenSize, bias: false,
|
||||
groupSize: groupSize, bits: bits))
|
||||
|
||||
self._qNorm = ModuleInfo(wrappedValue: RMSNorm(dimensions: headDim, eps: Float(config.rmsNormEps)))
|
||||
self._kNorm = ModuleInfo(wrappedValue: RMSNorm(dimensions: headDim, eps: Float(config.rmsNormEps)))
|
||||
|
||||
// Partial RoPE: only rotates first `ropeDims` of each head
|
||||
self.rope = MLXNN.RoPE(
|
||||
dimensions: ropeDims,
|
||||
traditional: false,
|
||||
base: Float(config.ropeTheta))
|
||||
|
||||
super.init()
|
||||
}
|
||||
|
||||
/// Forward pass.
|
||||
///
|
||||
/// - Parameters:
|
||||
/// - hiddenStates: [B, T, hiddenSize]
|
||||
/// - cache: Optional (keys, values) from previous steps, each [B, H_kv, S, D]
|
||||
/// - offset: RoPE position offset (used when cache is nil, e.g. first call)
|
||||
/// - Returns: (output [B, T, hiddenSize], updated KV cache)
|
||||
public func callAsFunction(
|
||||
_ hiddenStates: MLXArray,
|
||||
cache: (MLXArray, MLXArray)? = nil,
|
||||
offset: Int = 0
|
||||
) -> (MLXArray, (MLXArray, MLXArray)) {
|
||||
let b = hiddenStates.dim(0)
|
||||
let seqLen = hiddenStates.dim(1)
|
||||
let qDim = numQHeads * headDim
|
||||
|
||||
// Q projection: [B, T, 2*qDim] → reshape to [B, T, H, 2*D] → split Q/gate INTERLEAVED per head
|
||||
// CRITICAL: Must reshape BEFORE split (per Python reference).
|
||||
// Interleaved format: for each head, first D dims are Q, next D are gate.
|
||||
let qProjOut = qProj(hiddenStates) // [B, T, 4096 = 2*numQHeads*headDim]
|
||||
let qProjReshaped = qProjOut.reshaped(b, seqLen, numQHeads, 2 * headDim)
|
||||
let qgSplit = qProjReshaped.split(parts: 2, axis: -1)
|
||||
var queries = qgSplit[0] // [B, T, H, D=256]
|
||||
let gateSignal = qgSplit[1].reshaped(b, seqLen, qDim) // [B, T, 2048]
|
||||
|
||||
var keys = kProj(hiddenStates) // [B, T, numKVHeads * headDim]
|
||||
var values = vProj(hiddenStates) // [B, T, numKVHeads * headDim]
|
||||
|
||||
// Reshape K/V to multi-head: [B, T, H, D]
|
||||
keys = keys.reshaped(b, seqLen, numKVHeads, headDim)
|
||||
values = values.reshaped(b, seqLen, numKVHeads, headDim)
|
||||
|
||||
// QK norm (per-head)
|
||||
queries = qNorm(queries)
|
||||
keys = kNorm(keys)
|
||||
|
||||
// Transpose to [B, H, T, D]
|
||||
queries = queries.transposed(0, 2, 1, 3)
|
||||
keys = keys.transposed(0, 2, 1, 3)
|
||||
values = values.transposed(0, 2, 1, 3)
|
||||
|
||||
// Partial RoPE: MLXNN.RoPE with dimensions=ropeDims only rotates first ropeDims
|
||||
let ropeOffset = cache?.0.dim(2) ?? offset
|
||||
queries = rope(queries, offset: ropeOffset)
|
||||
keys = rope(keys, offset: ropeOffset)
|
||||
|
||||
// Update KV cache
|
||||
var cachedKeys = keys
|
||||
var cachedValues = values
|
||||
if let (prevK, prevV) = cache {
|
||||
cachedKeys = concatenated([prevK, keys], axis: 2)
|
||||
cachedValues = concatenated([prevV, values], axis: 2)
|
||||
}
|
||||
|
||||
// Causal mask
|
||||
let mask: MLXFast.ScaledDotProductAttentionMaskMode
|
||||
if seqLen <= 1 && (cache != nil || offset > 0) {
|
||||
mask = .none
|
||||
} else {
|
||||
let kvLen = cachedKeys.dim(2)
|
||||
let pastLen = kvLen - seqLen
|
||||
let causal = MLXArray.tri(seqLen, m: kvLen, k: pastLen, type: Float.self) - 1
|
||||
let additiveMask = causal * Float.greatestFiniteMagnitude // 0 for attended, -FLT_MAX for masked
|
||||
mask = .array(additiveMask.reshaped(1, 1, seqLen, kvLen).asType(queries.dtype))
|
||||
}
|
||||
|
||||
// SDPA (handles GQA natively)
|
||||
let attnOut = SDPA.attendAndMerge(
|
||||
qHeads: queries, kHeads: cachedKeys, vHeads: cachedValues,
|
||||
scale: scale, mask: mask)
|
||||
|
||||
// Gated output: attn_out * sigmoid(gate), then o_proj
|
||||
// Reference: self.o_proj(output * mx.sigmoid(gate))
|
||||
let gated = attnOut * sigmoid(gateSignal) // [B, T, qDim=2048]
|
||||
let output = oProj(gated) // [B, T, hiddenSize=1024]
|
||||
|
||||
return (output, (cachedKeys, cachedValues))
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Qwen3.5 Transformer Layer
|
||||
|
||||
/// A single transformer layer in the Qwen3.5 hybrid model.
|
||||
///
|
||||
/// Either a DeltaNet (linear_attention) or GatedAttention (full_attention) layer,
|
||||
/// both sharing the same pre-norm structure and SwiGLU MLP.
|
||||
///
|
||||
/// The attention submodule is stored as the base `Module` type and registered
|
||||
/// under the key `"self_attn"` via `@ModuleInfo`. This ensures the MLX Module
|
||||
/// system discovers it for parameter traversal (`eval`, `clearParameters`, etc.)
|
||||
/// and weight loading maps to the correct key path.
|
||||
public final class Qwen35TransformerLayer: Module {
|
||||
public let layerType: String
|
||||
|
||||
@ModuleInfo(key: "input_layernorm") var inputLayerNorm: RMSNorm
|
||||
@ModuleInfo(key: "post_attention_layernorm") var postAttentionLayerNorm: RMSNorm
|
||||
@ModuleInfo var mlp: Qwen35MLP
|
||||
|
||||
/// The attention submodule — either DeltaNetLayer or GatedAttentionLayer.
|
||||
/// Key is "linear_attn" for DeltaNet, "self_attn" for GatedAttention (HuggingFace convention).
|
||||
@ModuleInfo var attn: Module
|
||||
|
||||
public init(config: Qwen3ChatConfig, layerType: String) {
|
||||
self.layerType = layerType
|
||||
|
||||
self._inputLayerNorm = ModuleInfo(wrappedValue: RMSNorm(dimensions: config.hiddenSize, eps: Float(config.rmsNormEps)))
|
||||
self._postAttentionLayerNorm = ModuleInfo(wrappedValue: RMSNorm(dimensions: config.hiddenSize, eps: Float(config.rmsNormEps)))
|
||||
self._mlp = ModuleInfo(
|
||||
wrappedValue: Qwen35MLP(config: config))
|
||||
|
||||
if layerType == "linear_attention" {
|
||||
self._attn = ModuleInfo(
|
||||
wrappedValue: DeltaNetLayer(config: config),
|
||||
key: "linear_attn")
|
||||
} else {
|
||||
self._attn = ModuleInfo(
|
||||
wrappedValue: GatedAttentionLayer(config: config),
|
||||
key: "self_attn")
|
||||
}
|
||||
|
||||
super.init()
|
||||
}
|
||||
|
||||
/// Access the DeltaNet submodule (only valid for linear_attention layers).
|
||||
public var deltaNet: DeltaNetLayer? { attn as? DeltaNetLayer }
|
||||
|
||||
/// Access the GatedAttention submodule (only valid for full_attention layers).
|
||||
public var gatedAttn: GatedAttentionLayer? { attn as? GatedAttentionLayer }
|
||||
|
||||
/// Forward for DeltaNet (linear attention) layer.
|
||||
public func forwardDeltaNet(
|
||||
_ x: MLXArray,
|
||||
state: DeltaNetLayer.State?
|
||||
) -> (MLXArray, DeltaNetLayer.State) {
|
||||
guard let dn = deltaNet else {
|
||||
fatalError("forwardDeltaNet called on full_attention layer")
|
||||
}
|
||||
let normed = inputLayerNorm(x)
|
||||
let (attnOut, newState) = dn(normed, state: state)
|
||||
var h = x + attnOut
|
||||
h = h + mlp(postAttentionLayerNorm(h))
|
||||
return (h, newState)
|
||||
}
|
||||
|
||||
/// Forward for GatedAttention (full attention) layer.
|
||||
public func forwardGatedAttention(
|
||||
_ x: MLXArray,
|
||||
cache: (MLXArray, MLXArray)?,
|
||||
offset: Int
|
||||
) -> (MLXArray, (MLXArray, MLXArray)) {
|
||||
guard let ga = gatedAttn else {
|
||||
fatalError("forwardGatedAttention called on linear_attention layer")
|
||||
}
|
||||
let normed = inputLayerNorm(x)
|
||||
let (attnOut, newCache) = ga(normed, cache: cache, offset: offset)
|
||||
var h = x + attnOut
|
||||
|
||||
h = h + mlp(postAttentionLayerNorm(h))
|
||||
return (h, newCache)
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - SwiGLU MLP
|
||||
|
||||
/// SwiGLU MLP for Qwen3.5 (quantized INT4).
|
||||
public final class Qwen35MLP: Module {
|
||||
@ModuleInfo(key: "gate_proj") var gateProj: QuantizedLinear
|
||||
@ModuleInfo(key: "up_proj") var upProj: QuantizedLinear
|
||||
@ModuleInfo(key: "down_proj") var downProj: QuantizedLinear
|
||||
|
||||
public init(config: Qwen3ChatConfig) {
|
||||
let hs = config.hiddenSize
|
||||
let is_ = config.intermediateSize
|
||||
let gs = 64, bits = 4
|
||||
|
||||
self._gateProj = ModuleInfo(
|
||||
wrappedValue: QuantizedLinear(hs, is_, bias: false, groupSize: gs, bits: bits),
|
||||
key: "gate_proj")
|
||||
self._upProj = ModuleInfo(
|
||||
wrappedValue: QuantizedLinear(hs, is_, bias: false, groupSize: gs, bits: bits),
|
||||
key: "up_proj")
|
||||
self._downProj = ModuleInfo(
|
||||
wrappedValue: QuantizedLinear(is_, hs, bias: false, groupSize: gs, bits: bits),
|
||||
key: "down_proj")
|
||||
|
||||
super.init()
|
||||
}
|
||||
|
||||
public func callAsFunction(_ x: MLXArray) -> MLXArray {
|
||||
downProj(silu(gateProj(x)) * upProj(x))
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Qwen3.5 Full Model
|
||||
|
||||
/// Qwen3.5-0.8B hybrid transformer with DeltaNet linear attention and GatedAttention.
|
||||
///
|
||||
/// Architecture: 24 layers in pattern [3x DeltaNet, 1x GatedAttention] x 6.
|
||||
/// - DeltaNet layers (18 of 24): O(1) memory per step via recurrent state, no KV cache.
|
||||
/// - GatedAttention layers (6 of 24): standard SDPA with KV cache, partial RoPE (25%).
|
||||
/// - Tied embeddings: lm_head reuses embed_tokens weights (PreQuantizedEmbedding.asLinear).
|
||||
///
|
||||
/// This gives a favorable memory/compute tradeoff: recurrent DeltaNet layers handle
|
||||
/// most computation with fixed memory, while sparse full attention layers provide
|
||||
/// global context at every 4th layer.
|
||||
public final class Qwen35MLXModel: Module {
|
||||
public let config: Qwen3ChatConfig
|
||||
public let layerTypes: [String]
|
||||
public let fullAttentionIndices: [Int]
|
||||
|
||||
@ModuleInfo(key: "embed_tokens") var embedTokens: PreQuantizedEmbedding
|
||||
@ModuleInfo var layers: [Qwen35TransformerLayer]
|
||||
@ModuleInfo var norm: RMSNorm
|
||||
|
||||
public init(config: Qwen3ChatConfig) {
|
||||
self.config = config
|
||||
|
||||
let types = config.layerTypes ?? Array(
|
||||
repeating: "full_attention", count: config.numHiddenLayers)
|
||||
self.layerTypes = types
|
||||
self.fullAttentionIndices = types.enumerated().compactMap {
|
||||
$0.element == "full_attention" ? $0.offset : nil
|
||||
}
|
||||
|
||||
self._embedTokens = ModuleInfo(wrappedValue: PreQuantizedEmbedding(
|
||||
embeddingCount: config.vocabSize,
|
||||
dimensions: config.hiddenSize,
|
||||
groupSize: 64, bits: 4))
|
||||
|
||||
self._layers = ModuleInfo(
|
||||
wrappedValue: types.map { Qwen35TransformerLayer(config: config, layerType: $0) })
|
||||
|
||||
self._norm = ModuleInfo(
|
||||
wrappedValue: RMSNorm(dimensions: config.hiddenSize, eps: Float(config.rmsNormEps)))
|
||||
|
||||
super.init()
|
||||
}
|
||||
|
||||
// MARK: - Inference State
|
||||
|
||||
/// Combined inference state: DeltaNet recurrent states + GatedAttention KV caches.
|
||||
public struct InferenceState {
|
||||
/// Per-layer DeltaNet state (nil for full_attention layers).
|
||||
public var deltaNetStates: [DeltaNetLayer.State?]
|
||||
/// Per-layer KV cache (nil for linear_attention layers, and for full_attention
|
||||
/// layers before any tokens have been processed).
|
||||
public var kvCaches: [(MLXArray, MLXArray)?]
|
||||
/// Current sequence position (for RoPE offset in GatedAttention layers).
|
||||
public var position: Int
|
||||
|
||||
public static func initial(config: Qwen3ChatConfig, batchSize: Int = 1) -> InferenceState {
|
||||
let types = config.layerTypes ?? Array(
|
||||
repeating: "full_attention", count: config.numHiddenLayers)
|
||||
let numHeads = config.linearNumKeyHeads ?? 16
|
||||
let headDim = config.linearKeyHeadDim ?? 128
|
||||
let qkvDim = 3 * numHeads * headDim
|
||||
let convKernel = config.linearConvKernelDim ?? 4
|
||||
|
||||
return InferenceState(
|
||||
deltaNetStates: types.map { type in
|
||||
type == "linear_attention"
|
||||
? DeltaNetLayer.State.initial(
|
||||
batchSize: batchSize, numHeads: numHeads, headDim: headDim,
|
||||
qkvDim: qkvDim, convKernel: convKernel)
|
||||
: nil
|
||||
},
|
||||
kvCaches: types.map { _ in nil },
|
||||
position: 0
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Forward Pass
|
||||
|
||||
/// Forward pass through the full model.
|
||||
///
|
||||
/// - Parameters:
|
||||
/// - inputIds: Token IDs [B, T]
|
||||
/// - state: Inference state
|
||||
/// - Returns: (logits [B, T, vocabSize], updated state)
|
||||
public func forward(
|
||||
inputIds: MLXArray,
|
||||
state: InferenceState
|
||||
) -> (MLXArray, InferenceState) {
|
||||
let seqLen = inputIds.dim(1)
|
||||
var hidden = embedTokens(inputIds) // [B, T, hiddenSize]
|
||||
|
||||
var newDeltaStates = state.deltaNetStates
|
||||
var newKVCaches = state.kvCaches
|
||||
|
||||
for (i, layer) in layers.enumerated() {
|
||||
if layerTypes[i] == "linear_attention" {
|
||||
let (h, newState) = layer.forwardDeltaNet(hidden, state: state.deltaNetStates[i])
|
||||
hidden = h
|
||||
newDeltaStates[i] = newState
|
||||
} else {
|
||||
let (h, newCache) = layer.forwardGatedAttention(
|
||||
hidden, cache: state.kvCaches[i], offset: state.position)
|
||||
hidden = h
|
||||
newKVCaches[i] = newCache
|
||||
}
|
||||
}
|
||||
|
||||
hidden = norm(hidden)
|
||||
|
||||
// Tied LM head
|
||||
let logits = embedTokens.asLinear(hidden)
|
||||
|
||||
let newState = InferenceState(
|
||||
deltaNetStates: newDeltaStates,
|
||||
kvCaches: newKVCaches,
|
||||
position: state.position + seqLen)
|
||||
|
||||
return (logits, newState)
|
||||
}
|
||||
|
||||
// MARK: - Text Generation
|
||||
|
||||
/// Generate text tokens autoregressively.
|
||||
///
|
||||
/// - Parameters:
|
||||
/// - promptIds: Prompt token IDs
|
||||
/// - sampling: Sampling configuration
|
||||
/// - Returns: Generated token IDs (excluding prompt)
|
||||
public func generate(
|
||||
promptIds: [Int],
|
||||
sampling: ChatSamplingConfig = .default
|
||||
) -> [Int] {
|
||||
var state = InferenceState.initial(config: config)
|
||||
|
||||
// Prefill
|
||||
let prompt = MLXArray(promptIds.map { Int32($0) }).expandedDimensions(axis: 0)
|
||||
let (prefillLogits, prefillState) = forward(inputIds: prompt, state: state)
|
||||
state = prefillState
|
||||
eval(prefillLogits)
|
||||
|
||||
// Sample first token
|
||||
var token = sampleFromLogits(prefillLogits, at: promptIds.count - 1,
|
||||
config: sampling, history: promptIds)
|
||||
if token == config.eosTokenId { return [] }
|
||||
|
||||
var generated = [token]
|
||||
|
||||
// Decode loop
|
||||
for _ in 1..<sampling.maxTokens {
|
||||
let input = MLXArray([Int32(token)]).expandedDimensions(axis: 0)
|
||||
let (logits, newState) = forward(inputIds: input, state: state)
|
||||
state = newState
|
||||
eval(logits)
|
||||
|
||||
token = sampleFromLogits(logits, at: 0,
|
||||
config: sampling, history: promptIds + generated)
|
||||
if token == config.eosTokenId { break }
|
||||
generated.append(token)
|
||||
}
|
||||
|
||||
return generated
|
||||
}
|
||||
|
||||
// MARK: - Helpers
|
||||
|
||||
private func sampleFromLogits(
|
||||
_ logits: MLXArray, at position: Int,
|
||||
config: ChatSamplingConfig, history: [Int]
|
||||
) -> Int {
|
||||
let posLogits = logits[0, position] // [vocabSize]
|
||||
let f32 = posLogits.asType(.float32)
|
||||
eval(f32)
|
||||
let count = self.config.vocabSize
|
||||
let floats: [Float] = f32.asArray(Float.self)
|
||||
return ChatSampler.sample(logits: Array(floats.prefix(count)), config: config, previousTokens: history)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
import Foundation
|
||||
import AudioCommon
|
||||
|
||||
/// Common interface for Qwen3.5 chat backends (MLX or CoreML).
|
||||
public protocol Qwen35ChatBackend: AnyObject {
|
||||
var tokenizer: ChatTokenizer { get }
|
||||
var config: Qwen3ChatConfig { get }
|
||||
func generateStream(messages: [ChatMessage], sampling: ChatSamplingConfig)
|
||||
-> AsyncThrowingStream<String, Error>
|
||||
func resetState()
|
||||
}
|
||||
|
||||
extension Qwen35MLXChat: Qwen35ChatBackend {}
|
||||
extension Qwen35CoreMLChat: Qwen35ChatBackend {}
|
||||
|
||||
/// Bridges any Qwen3.5 backend to VoicePipeline's PipelineLLM protocol.
|
||||
public final class Qwen35PipelineLLM: PipelineLLM {
|
||||
private let model: any Qwen35ChatBackend
|
||||
private let systemPrompt: String
|
||||
private let sampling: ChatSamplingConfig
|
||||
private var cancelled = false
|
||||
|
||||
public var onToken: ((String) -> Void)?
|
||||
|
||||
public init(
|
||||
model: any Qwen35ChatBackend,
|
||||
systemPrompt: String = "Your name is Tama. Give short direct answers. Do not explain your reasoning.",
|
||||
sampling: ChatSamplingConfig = .default
|
||||
) {
|
||||
self.model = model
|
||||
self.systemPrompt = systemPrompt
|
||||
self.sampling = sampling
|
||||
}
|
||||
|
||||
public func chat(
|
||||
messages: [(role: MessageRole, content: String)],
|
||||
onToken: @escaping (String, Bool) -> Void
|
||||
) {
|
||||
cancelled = false
|
||||
|
||||
let chatMessages = messages.compactMap { msg -> ChatMessage? in
|
||||
switch msg.role {
|
||||
case .system: return ChatMessage(role: .system, content: msg.content)
|
||||
case .user: return ChatMessage(role: .user, content: msg.content)
|
||||
case .assistant: return ChatMessage(role: .assistant, content: msg.content)
|
||||
default: return nil
|
||||
}
|
||||
}
|
||||
|
||||
var fullMessages = [ChatMessage(role: .system, content: systemPrompt)]
|
||||
fullMessages.append(contentsOf: chatMessages)
|
||||
|
||||
let stream = model.generateStream(messages: fullMessages, sampling: sampling)
|
||||
let semaphore = DispatchSemaphore(value: 0)
|
||||
var fullResponse = ""
|
||||
|
||||
Task {
|
||||
do {
|
||||
for try await chunk in stream {
|
||||
guard !self.cancelled else { break }
|
||||
fullResponse += chunk
|
||||
self.onToken?(chunk)
|
||||
onToken(chunk, false)
|
||||
}
|
||||
} catch { }
|
||||
|
||||
if !fullResponse.isEmpty {
|
||||
onToken("", true)
|
||||
}
|
||||
semaphore.signal()
|
||||
}
|
||||
|
||||
semaphore.wait()
|
||||
}
|
||||
|
||||
public func cancel() {
|
||||
cancelled = true
|
||||
}
|
||||
}
|
||||
+226
@@ -0,0 +1,226 @@
|
||||
import Foundation
|
||||
import MLXCommon
|
||||
import MLX
|
||||
import MLXNN
|
||||
|
||||
// MARK: - Weight Loading for Qwen3.5-0.8B MLX Model
|
||||
|
||||
/// Loads quantized safetensors weights into the Qwen3.5 MLX model.
|
||||
///
|
||||
/// Expected weight key structure (HuggingFace / mlx-community format):
|
||||
///
|
||||
/// Keys may have `model.` or `language_model.model.` prefix — both are stripped.
|
||||
///
|
||||
/// - `embed_tokens.*` -> embed_tokens (PreQuantizedEmbedding)
|
||||
/// - `layers.{i}.linear_attn.*` -> DeltaNet (linear_attention layers)
|
||||
/// - `layers.{i}.self_attn.*` -> GatedAttention (full_attention layers)
|
||||
/// - `layers.{i}.mlp.*` -> SwiGLU MLP
|
||||
/// - `layers.{i}.input_layernorm.*`
|
||||
/// - `layers.{i}.post_attention_layernorm.*`
|
||||
/// - `norm.*` -> final RMSNorm
|
||||
///
|
||||
/// DeltaNet (linear_attention) weights under `linear_attn.`:
|
||||
/// - `in_proj_qkv.{weight,scales,biases}`: quantized [6144, 1024]
|
||||
/// - `in_proj_z.{weight,scales,biases}`: quantized [2048, 1024]
|
||||
/// - `in_proj_b.{weight,scales,biases}`: quantized [16, 1024]
|
||||
/// - `in_proj_a.{weight,scales,biases}`: quantized [16, 1024]
|
||||
/// - `conv1d.weight`: [6144, 1, 4]
|
||||
/// - `dt_bias`: [16]
|
||||
/// - `A_log`: [16]
|
||||
/// - `norm.weight`: [128]
|
||||
/// - `out_proj.{weight,scales,biases}`: quantized [1024, 2048]
|
||||
///
|
||||
/// GatedAttention (full_attention) weights under `self_attn.`:
|
||||
/// - `q_proj.{weight,scales,biases}`: quantized [4096, 1024]
|
||||
/// - `k_proj.{weight,scales,biases}`: quantized [512, 1024]
|
||||
/// - `v_proj.{weight,scales,biases}`: quantized [512, 1024]
|
||||
/// - `o_proj.{weight,scales,biases}`: quantized [1024, 2048]
|
||||
/// - `q_norm.weight`: [256]
|
||||
/// - `k_norm.weight`: [256]
|
||||
///
|
||||
/// MLP weights under `mlp.`:
|
||||
/// - `gate_proj.{weight,scales,biases}`: quantized [3584, 1024]
|
||||
/// - `up_proj.{weight,scales,biases}`: quantized [3584, 1024]
|
||||
/// - `down_proj.{weight,scales,biases}`: quantized [1024, 3584]
|
||||
public enum Qwen35WeightLoader {
|
||||
|
||||
/// Load weights from a directory containing safetensors files.
|
||||
///
|
||||
/// - Parameters:
|
||||
/// - model: The Qwen3.5 MLX model to load weights into
|
||||
/// - directory: Directory containing safetensors files
|
||||
/// - progressHandler: Optional progress callback
|
||||
public static func loadWeights(
|
||||
into model: Qwen35MLXModel,
|
||||
from directory: URL,
|
||||
progressHandler: ((Double, String) -> Void)? = nil
|
||||
) throws {
|
||||
progressHandler?(0.05, "Loading weight files...")
|
||||
|
||||
// Load all safetensors files from the directory
|
||||
let allWeights = try CommonWeightLoader.loadAllSafetensors(from: directory)
|
||||
progressHandler?(0.3, "Loaded \(allWeights.count) tensors")
|
||||
|
||||
// Strip prefix from keys. Handles two formats:
|
||||
// - Our format: "model.layers.0.*"
|
||||
// - mlx-community VLM: "language_model.model.layers.0.*" (also has vision_tower.* which we skip)
|
||||
var modelWeights: [String: MLXArray] = [:]
|
||||
for (key, value) in allWeights {
|
||||
if key.hasPrefix("language_model.model.") {
|
||||
modelWeights[String(key.dropFirst("language_model.model.".count))] = value
|
||||
} else if key.hasPrefix("model.") {
|
||||
modelWeights[String(key.dropFirst("model.".count))] = value
|
||||
} else if key.hasPrefix("lm_head.") || key.hasPrefix("vision_tower.") {
|
||||
// Skip — lm_head is tied to embed_tokens, vision_tower not needed
|
||||
continue
|
||||
} else {
|
||||
modelWeights[key] = value
|
||||
}
|
||||
}
|
||||
|
||||
progressHandler?(0.4, "Applying embedding weights...")
|
||||
|
||||
// Load embed_tokens (PreQuantizedEmbedding)
|
||||
CommonWeightLoader.applyQuantizedEmbeddingWeights(
|
||||
to: model.embedTokens,
|
||||
prefix: "embed_tokens",
|
||||
from: modelWeights)
|
||||
|
||||
// Load final norm
|
||||
CommonWeightLoader.applyRMSNormWeights(
|
||||
to: model.norm, prefix: "norm", from: modelWeights)
|
||||
|
||||
progressHandler?(0.5, "Loading transformer layers...")
|
||||
|
||||
// Load each layer
|
||||
let numLayers = model.config.numHiddenLayers
|
||||
for i in 0..<numLayers {
|
||||
let prefix = "layers.\(i)"
|
||||
let layer = model.layers[i]
|
||||
|
||||
// Layer norms
|
||||
CommonWeightLoader.applyRMSNormWeights(
|
||||
to: layer.inputLayerNorm,
|
||||
prefix: "\(prefix).input_layernorm",
|
||||
from: modelWeights)
|
||||
CommonWeightLoader.applyRMSNormWeights(
|
||||
to: layer.postAttentionLayerNorm,
|
||||
prefix: "\(prefix).post_attention_layernorm",
|
||||
from: modelWeights)
|
||||
|
||||
// MLP
|
||||
applyQuantizedMLPWeights(
|
||||
to: layer.mlp,
|
||||
prefix: "\(prefix).mlp",
|
||||
from: modelWeights)
|
||||
|
||||
// Attention (type-specific, different key prefix per HuggingFace convention)
|
||||
if layer.layerType == "linear_attention" {
|
||||
try applyDeltaNetWeights(
|
||||
to: layer.deltaNet!,
|
||||
prefix: "\(prefix).linear_attn",
|
||||
from: modelWeights)
|
||||
} else {
|
||||
applyGatedAttentionWeights(
|
||||
to: layer.gatedAttn!,
|
||||
prefix: "\(prefix).self_attn",
|
||||
from: modelWeights)
|
||||
}
|
||||
|
||||
let pct = 0.5 + 0.45 * Double(i + 1) / Double(numLayers)
|
||||
progressHandler?(pct, "Layer \(i + 1)/\(numLayers)")
|
||||
}
|
||||
|
||||
// Evaluate all parameters
|
||||
eval(model)
|
||||
progressHandler?(1.0, "Weights loaded")
|
||||
}
|
||||
|
||||
// MARK: - DeltaNet Weight Loading
|
||||
|
||||
/// Apply weights to a DeltaNet (linear attention) layer.
|
||||
///
|
||||
/// All DeltaNet projections are quantized INT4 (matching mlx-community format).
|
||||
/// The conv1d weight and scalar parameters (dt_bias, A_log) are loaded directly.
|
||||
private static func applyDeltaNetWeights(
|
||||
to layer: DeltaNetLayer,
|
||||
prefix: String,
|
||||
from weights: [String: MLXArray]
|
||||
) throws {
|
||||
// Quantized projections
|
||||
CommonWeightLoader.applyQuantizedLinearWeights(
|
||||
to: layer.inProjQKV, prefix: "\(prefix).in_proj_qkv", from: weights)
|
||||
CommonWeightLoader.applyQuantizedLinearWeights(
|
||||
to: layer.inProjZ, prefix: "\(prefix).in_proj_z", from: weights)
|
||||
CommonWeightLoader.applyQuantizedLinearWeights(
|
||||
to: layer.inProjB, prefix: "\(prefix).in_proj_b", from: weights)
|
||||
CommonWeightLoader.applyQuantizedLinearWeights(
|
||||
to: layer.inProjA, prefix: "\(prefix).in_proj_a", from: weights)
|
||||
|
||||
// Raw parameters: conv1d weight, dt_bias, A_log
|
||||
// Note: ParameterInfo key from property declaration is lost when reassigning
|
||||
// in init(), so use Swift property names (convWeight, dtBias, aLog).
|
||||
var rawParams: [String: NestedItem<String, MLXArray>] = [:]
|
||||
if let w = weights["\(prefix).conv1d.weight"] {
|
||||
rawParams["convWeight"] = .value(w)
|
||||
}
|
||||
if let dtb = weights["\(prefix).dt_bias"] {
|
||||
rawParams["dtBias"] = .value(dtb)
|
||||
}
|
||||
if let alog = weights["\(prefix).A_log"] {
|
||||
rawParams["aLog"] = .value(alog)
|
||||
}
|
||||
if !rawParams.isEmpty {
|
||||
layer.update(parameters: ModuleParameters(values: rawParams))
|
||||
}
|
||||
|
||||
// Per-head norm
|
||||
CommonWeightLoader.applyRMSNormWeights(
|
||||
to: layer.norm, prefix: "\(prefix).norm", from: weights)
|
||||
|
||||
// Output projection (quantized)
|
||||
CommonWeightLoader.applyQuantizedLinearWeights(
|
||||
to: layer.outProj, prefix: "\(prefix).out_proj", from: weights)
|
||||
}
|
||||
|
||||
// MARK: - GatedAttention Weight Loading
|
||||
|
||||
/// Apply weights to a GatedAttention (full attention) layer.
|
||||
///
|
||||
/// All projections are quantized (INT4 with group_size=64).
|
||||
private static func applyGatedAttentionWeights(
|
||||
to layer: GatedAttentionLayer,
|
||||
prefix: String,
|
||||
from weights: [String: MLXArray]
|
||||
) {
|
||||
CommonWeightLoader.applyQuantizedLinearWeights(
|
||||
to: layer.qProj, prefix: "\(prefix).q_proj", from: weights)
|
||||
CommonWeightLoader.applyQuantizedLinearWeights(
|
||||
to: layer.kProj, prefix: "\(prefix).k_proj", from: weights)
|
||||
CommonWeightLoader.applyQuantizedLinearWeights(
|
||||
to: layer.vProj, prefix: "\(prefix).v_proj", from: weights)
|
||||
CommonWeightLoader.applyQuantizedLinearWeights(
|
||||
to: layer.oProj, prefix: "\(prefix).o_proj", from: weights)
|
||||
|
||||
CommonWeightLoader.applyRMSNormWeights(
|
||||
to: layer.qNorm, prefix: "\(prefix).q_norm", from: weights)
|
||||
CommonWeightLoader.applyRMSNormWeights(
|
||||
to: layer.kNorm, prefix: "\(prefix).k_norm", from: weights)
|
||||
}
|
||||
|
||||
// MARK: - MLP Weight Loading
|
||||
|
||||
/// Apply quantized MLP weights (SwiGLU: gate_proj, up_proj, down_proj).
|
||||
private static func applyQuantizedMLPWeights(
|
||||
to mlp: Qwen35MLP,
|
||||
prefix: String,
|
||||
from weights: [String: MLXArray]
|
||||
) {
|
||||
CommonWeightLoader.applyQuantizedLinearWeights(
|
||||
to: mlp.gateProj, prefix: "\(prefix).gate_proj", from: weights)
|
||||
CommonWeightLoader.applyQuantizedLinearWeights(
|
||||
to: mlp.upProj, prefix: "\(prefix).up_proj", from: weights)
|
||||
CommonWeightLoader.applyQuantizedLinearWeights(
|
||||
to: mlp.downProj, prefix: "\(prefix).down_proj", from: weights)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
import Foundation
|
||||
|
||||
/// Model architecture type.
|
||||
public enum ChatModelArch: String, Codable, Sendable {
|
||||
/// Qwen3.5 hybrid (DeltaNet linear attention + GatedAttention)
|
||||
case qwen35 = "qwen3_5_text"
|
||||
}
|
||||
|
||||
/// Configuration for Qwen3.5 chat model.
|
||||
public struct Qwen3ChatConfig: Codable, Sendable {
|
||||
public let hiddenSize: Int
|
||||
public let numHiddenLayers: Int
|
||||
public let numAttentionHeads: Int
|
||||
public let numKeyValueHeads: Int
|
||||
public let headDim: Int
|
||||
public let intermediateSize: Int
|
||||
public let vocabSize: Int
|
||||
public let maxSeqLen: Int
|
||||
public let ropeTheta: Double
|
||||
public let rmsNormEps: Double
|
||||
public let eosTokenId: Int
|
||||
public let padTokenId: Int
|
||||
public let quantization: String
|
||||
|
||||
// Qwen3.5-specific fields
|
||||
public let modelType: ChatModelArch?
|
||||
/// Per-layer type: "linear_attention" (DeltaNet) or "full_attention" (GatedAttention)
|
||||
public let layerTypes: [String]?
|
||||
/// How often a full_attention layer appears (e.g., 4 = every 4th layer)
|
||||
public let fullAttentionInterval: Int?
|
||||
/// DeltaNet linear attention head config
|
||||
public let linearNumKeyHeads: Int?
|
||||
public let linearKeyHeadDim: Int?
|
||||
public let linearNumValueHeads: Int?
|
||||
public let linearValueHeadDim: Int?
|
||||
/// Causal conv1d kernel size for DeltaNet
|
||||
public let linearConvKernelDim: Int?
|
||||
/// Partial RoPE factor for GatedAttention (e.g., 0.25)
|
||||
public let partialRotaryFactor: Double?
|
||||
/// Whether embeddings are tied (lm_head = embed_tokens)
|
||||
public let tieWordEmbeddings: Bool?
|
||||
|
||||
enum CodingKeys: String, CodingKey {
|
||||
case hiddenSize = "hidden_size"
|
||||
case numHiddenLayers = "num_hidden_layers"
|
||||
case numAttentionHeads = "num_attention_heads"
|
||||
case numKeyValueHeads = "num_key_value_heads"
|
||||
case headDim = "head_dim"
|
||||
case intermediateSize = "intermediate_size"
|
||||
case vocabSize = "vocab_size"
|
||||
case maxSeqLen = "max_seq_len"
|
||||
case ropeTheta = "rope_theta"
|
||||
case rmsNormEps = "rms_norm_eps"
|
||||
case eosTokenId = "eos_token_id"
|
||||
case padTokenId = "pad_token_id"
|
||||
case quantization
|
||||
case modelType = "model_type"
|
||||
case layerTypes = "layer_types"
|
||||
case fullAttentionInterval = "full_attention_interval"
|
||||
case linearNumKeyHeads = "linear_num_key_heads"
|
||||
case linearKeyHeadDim = "linear_key_head_dim"
|
||||
case linearNumValueHeads = "linear_num_value_heads"
|
||||
case linearValueHeadDim = "linear_value_head_dim"
|
||||
case linearConvKernelDim = "linear_conv_kernel_dim"
|
||||
case partialRotaryFactor = "partial_rotary_factor"
|
||||
case tieWordEmbeddings = "tie_word_embeddings"
|
||||
}
|
||||
|
||||
/// Whether this is a Qwen3.5 hybrid model.
|
||||
public var isQwen35: Bool {
|
||||
modelType == .qwen35 || layerTypes != nil
|
||||
}
|
||||
|
||||
/// Number of full-attention layers (that need KV cache).
|
||||
public var numFullAttentionLayers: Int {
|
||||
guard let types = layerTypes else { return numHiddenLayers }
|
||||
return types.filter { $0 == "full_attention" }.count
|
||||
}
|
||||
|
||||
/// Default config for Qwen3.5-0.8B.
|
||||
public static let qwen35_08B = Qwen3ChatConfig(
|
||||
hiddenSize: 1024,
|
||||
numHiddenLayers: 24,
|
||||
numAttentionHeads: 8,
|
||||
numKeyValueHeads: 2,
|
||||
headDim: 256,
|
||||
intermediateSize: 3584,
|
||||
vocabSize: 248320,
|
||||
maxSeqLen: 2048,
|
||||
ropeTheta: 10_000_000.0,
|
||||
rmsNormEps: 1e-6,
|
||||
eosTokenId: 248046, // <|im_end|> — stops generation at end of assistant turn
|
||||
padTokenId: 248044, // <|endoftext|>
|
||||
quantization: "int4",
|
||||
modelType: .qwen35,
|
||||
layerTypes: [
|
||||
"linear_attention", "linear_attention", "linear_attention", "full_attention",
|
||||
"linear_attention", "linear_attention", "linear_attention", "full_attention",
|
||||
"linear_attention", "linear_attention", "linear_attention", "full_attention",
|
||||
"linear_attention", "linear_attention", "linear_attention", "full_attention",
|
||||
"linear_attention", "linear_attention", "linear_attention", "full_attention",
|
||||
"linear_attention", "linear_attention", "linear_attention", "full_attention",
|
||||
],
|
||||
fullAttentionInterval: 4,
|
||||
linearNumKeyHeads: 16,
|
||||
linearKeyHeadDim: 128,
|
||||
linearNumValueHeads: 16,
|
||||
linearValueHeadDim: 128,
|
||||
linearConvKernelDim: 4,
|
||||
partialRotaryFactor: 0.25,
|
||||
tieWordEmbeddings: true
|
||||
)
|
||||
|
||||
/// Load config from a JSON file.
|
||||
public static func load(from url: URL) throws -> Qwen3ChatConfig {
|
||||
let data = try Data(contentsOf: url)
|
||||
return try JSONDecoder().decode(Qwen3ChatConfig.self, from: data)
|
||||
}
|
||||
}
|
||||
|
||||
/// Sampling parameters for text generation.
|
||||
public struct ChatSamplingConfig: Sendable {
|
||||
public var temperature: Float
|
||||
public var topK: Int
|
||||
public var topP: Float
|
||||
public var maxTokens: Int
|
||||
public var repetitionPenalty: Float
|
||||
|
||||
public init(
|
||||
temperature: Float = 0.7,
|
||||
topK: Int = 50,
|
||||
topP: Float = 0.9,
|
||||
maxTokens: Int = 256,
|
||||
repetitionPenalty: Float = 1.1
|
||||
) {
|
||||
self.temperature = temperature
|
||||
self.topK = topK
|
||||
self.topP = topP
|
||||
self.maxTokens = maxTokens
|
||||
self.repetitionPenalty = repetitionPenalty
|
||||
}
|
||||
|
||||
public static let `default` = ChatSamplingConfig()
|
||||
public static let creative = ChatSamplingConfig(temperature: 0.9, topP: 0.95)
|
||||
public static let precise = ChatSamplingConfig(temperature: 0.3, topK: 20, topP: 0.8)
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
import Foundation
|
||||
|
||||
/// Errors for Qwen3.5 chat model operations.
|
||||
public enum ChatModelError: LocalizedError {
|
||||
case modelLoadFailed(String)
|
||||
case tokenizerLoadFailed(String)
|
||||
case inferenceFailed(String)
|
||||
case configNotFound(URL)
|
||||
case modelNotFound(URL)
|
||||
|
||||
public var errorDescription: String? {
|
||||
switch self {
|
||||
case .modelLoadFailed(let reason):
|
||||
"Failed to load chat model: \(reason)"
|
||||
case .tokenizerLoadFailed(let reason):
|
||||
"Failed to load tokenizer: \(reason)"
|
||||
case .inferenceFailed(let reason):
|
||||
"Inference failed: \(reason)"
|
||||
case .configNotFound(let url):
|
||||
"Config not found at \(url.path)"
|
||||
case .modelNotFound(let url):
|
||||
"Model not found at \(url.path)"
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user