Skip to content

Instantly share code, notes, and snippets.

@Blaizzy

Blaizzy/BaseConfiguration.swift

Created February 23, 2026 21:36
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save Blaizzy/7e56017eee617852d403d6b643899844 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save Blaizzy/7e56017eee617852d403d6b643899844 to your computer and use it in GitHub Desktop.
Download ZIP
Get started with MLX-Swift with Qwen3 port
Raw
BaseConfiguration.swift
//
// BaseConfiguration.swift
// mlx-test
//
// Created by Prince Canuma on 29/12/25.
//
import Foundation
import MLX
/// Base ``LanguageModel`` configuration -- provides `modelType`
/// and `quantization` (used in loading the model).
///
/// This is used by ``ModelFactory/load(hub:configuration:progressHandler:)``
/// to determine the type of model to load.
public struct BaseConfiguration: Codable, Sendable {
public let modelType: String
public struct Quantization: Codable, Sendable, Equatable {
public init(groupSize: Int, bits: Int) {
self.groupSize = groupSize
self.bits = bits
}
public let groupSize: Int
public let bits: Int
private var _mode: QuantizationMode? = nil
public var mode: QuantizationMode { _mode ?? .affine }
public var asTuple: (Int, Int, QuantizationMode) { (groupSize, bits, mode) }
enum CodingKeys: String, CodingKey {
case groupSize = "group_size"
case bits = "bits"
case _mode = "mode"
}
}
/// handling instructions for ``PerLayerQuantization``
public enum QuantizationOption: Sendable {
case skip
case quantize(Quantization)
}
/// Per-layer ``Quantization`` values with optional default.
public struct PerLayerQuantization: Sendable {
public var quantization: Quantization? = nil
public var perLayerQuantization: [String: QuantizationOption]
public init(
quantization: BaseConfiguration.Quantization? = nil,
perLayerQuantization: [String: BaseConfiguration.QuantizationOption]
) {
self.quantization = quantization
self.perLayerQuantization = perLayerQuantization
}
/// The quantization to apply for the given layer name or nil for no quantization.
public func quantization(layer: String) -> Quantization? {
if let perLayer = perLayerQuantization[layer] {
switch perLayer {
case .skip:
return nil
case .quantize(let quantization):
return quantization
}
} else {
return quantization
}
}
}
/// Special codable to support a mixed key: Int / key: Quantization
/// structure for hereogenous quantization, e.g.
///
/// ```
/// "quantization": {
/// "group_size": 64,
/// "bits": 4,
/// "model.embed_tokens": {
/// "group_size": 32,
/// "bits": 4
/// },
/// "model.layers.0.self_attn.q_norm": false,
/// ```
///
/// This mixed type structure requires manual decoding.
struct QuantizationContainer: Codable, Sendable {
var quantization: Quantization
var perLayerQuantization: PerLayerQuantization
// based on Dictionary's coding key
internal struct _DictionaryCodingKey: CodingKey {
internal let stringValue: String
internal let intValue: Int?
internal init(stringValue: String) {
self.stringValue = stringValue
self.intValue = Int(stringValue)
}
internal init(intValue: Int) {
self.stringValue = "\(intValue)"
self.intValue = intValue
}
}
init(from decoder: any Decoder) throws {
// handle the embedded Quantization
self.quantization = try Quantization(from: decoder)
// and the interleaved per-layer values
var perLayerQuantization = [String: QuantizationOption]()
let container = try decoder.container(keyedBy: _DictionaryCodingKey.self)
for key in container.allKeys {
switch key.stringValue {
case Quantization.CodingKeys.groupSize.rawValue: continue
case Quantization.CodingKeys.bits.rawValue: continue
case Quantization.CodingKeys._mode.rawValue: continue
// additional keys that are not layer instructions, see
// mlx-community/bitnet-b1.58-2B-4T-4bit
case "quant_method", "linear_class", "quantization_mode": continue
default:
if let f = try? container.decode(Bool.self, forKey: key) {
if !f {
perLayerQuantization[key.stringValue] = .skip
}
} else {
perLayerQuantization[key.stringValue] = .quantize(
try container.decode(Quantization.self, forKey: key))
}
}
}
self.perLayerQuantization = PerLayerQuantization(
quantization: quantization, perLayerQuantization: perLayerQuantization)
}
func encode(to encoder: any Encoder) throws {
try quantization.encode(to: encoder)
var container = encoder.container(keyedBy: _DictionaryCodingKey.self)
for (key, value) in perLayerQuantization.perLayerQuantization {
switch value {
case .skip:
try container.encode(false, forKey: .init(stringValue: key))
case .quantize(let q):
try container.encode(q, forKey: .init(stringValue: key))
}
}
}
}
var quantizationContainer: QuantizationContainer?
@available(*, deprecated, message: "Please use perLayerQuantization instead")
public var quantization: Quantization? {
quantizationContainer?.quantization
}
public var perLayerQuantization: PerLayerQuantization? {
quantizationContainer?.perLayerQuantization
}
enum CodingKeys: String, CodingKey {
case modelType = "model_type"
case quantizationContainer = "quantization"
}
}
Raw
KVCache.swift
//
// KVCache.swift
// mlx-test
//
// Created by Prince Canuma on 27/12/25.
//
import Foundation
import MLX
import MLXNN
// MARK: - Model Protocols and Parameters
/// Protocol for language models that support KV caching
public protocol LanguageModel {
/// Create a new cache for this model
func newCache(parameters: GenerateParameters?) -> [KVCache]
}
/// Parameters for text generation
public struct GenerateParameters: Sendable {
public var maxKVSize: Int?
public var maxTokens: Int?
public var temperature: Float
public var topP: Float
public var repetitionPenalty: Float?
public var repetitionContextSize: Int
public init(
maxKVSize: Int? = nil,
maxTokens: Int? = nil,
temperature: Float = 0.7,
topP: Float = 1.0,
repetitionPenalty: Float? = nil,
repetitionContextSize: Int = 20
) {
self.maxKVSize = maxKVSize
self.maxTokens = maxTokens
self.temperature = temperature
self.topP = topP
self.repetitionPenalty = repetitionPenalty
self.repetitionContextSize = repetitionContextSize
}
}
/// Protocol for models that can provide KV cache dimensions
public protocol KVCacheDimensionProvider {
/// Number of layers in the model
var numLayers: Int { get }
}
/// Default implementation for models that conform to KVCacheDimensionProvider
extension LanguageModel where Self: KVCacheDimensionProvider {
public func newCache(parameters: GenerateParameters?) -> [KVCache] {
if let maxKVSize = parameters?.maxKVSize {
return (0..<numLayers).map { _ in
RotatingKVCache(maxSize: maxKVSize, keep: 4)
}
} else {
return (0..<numLayers).map { _ in KVCacheSimple() }
}
}
}
/// Implementation of KV cache functionality for MLX Swift
///
///
/// ## Quantized Cache Usage
///
/// **Standard caches:**
/// ```swift
/// let cache = KVCacheSimple()
/// let (keys, values) = cache.update(keys: keys, values: values)
/// let output = MLXFast.scaledDotProductAttention(queries: q, keys: keys, values: values, ...)
/// ```
///
/// **Quantized cache:**
/// ```swift
/// let quantizedCache = QuantizedKVCache(groupSize: 64, bits: 4)
/// let (qKeys, qValues) = quantizedCache.updateQuantized(keys: keys, values: values)
///
/// let output = quantizedScaledDotProductAttention(
/// queries: queries,
/// quantizedKeys: qKeys,
/// quantizedValues: qValues,
/// scale: scale,
/// mask: mask,
/// groupSize: quantizedCache.groupSize,
/// bits: quantizedCache.bits
/// )
/// ```
///
/// Interface for Key/Value cache for LLMs.
///
/// See ``LanguageModel/newCache(parameters:)``
public protocol KVCache: Evaluatable {
/// get the current offset
var offset: Int { get }
/// get the maximum size (if any)
var maxSize: Int? { get }
/// update the cache with new keys and values and return all keys/values
func update(keys: MLXArray, values: MLXArray) -> (MLXArray, MLXArray)
/// get the current state for serialization
var state: [MLXArray] { get set }
/// get/set metadata state as string array for serialization
var metaState: [String] { get set }
/// whether this cache can be trimmed
var isTrimmable: Bool { get }
/// trim n tokens from the cache, returning actual number trimmed
@discardableResult
func trim(_ n: Int) -> Int
/// Create an attention mask for this cache
///
/// This method encapsulates cache-specific mask creation logic. Implementations should handle offset capping, window size logic,
/// and optimization decisions (symbolic vs array masks).
///
/// - Parameters:
/// - n: The sequence length for the new tokens
/// - windowSize: Optional sliding window size
/// - returnArray: Force return of array mask instead of symbolic
/// - Returns: Attention mask mode for scaled dot product attention
func makeMask(
n: Int, windowSize: Int?, returnArray: Bool
) -> MLXFast.ScaledDotProductAttentionMaskMode
}
/// Protocol for caches that support efficient quantized operations
///
/// **Usage Example:**
/// ```swift
/// // Efficient quantized path
/// if let quantizedCache = cache as? QuantizedKVCacheProtocol {
/// let (qKeys, qValues) = quantizedCache.updateQuantized(keys: k, values: v)
/// // Use native quantized operations
/// let scores = quantizedMatmul(queries, w: qKeys.0, scales: qKeys.1, biases: qKeys.2, ...)
/// } else {
/// // Regular path
/// let (k, v) = cache.update(keys: k, values: v)
/// let output = MLXFast.scaledDotProductAttention(queries: q, keys: k, values: v, ...)
/// }
/// ```
public protocol QuantizedKVCacheProtocol: KVCache {
/// The quantization group size used
var groupSize: Int { get }
/// The number of quantization bits used
var bits: Int { get }
/// Quantization mode
var mode: QuantizationMode { get }
/// Update cache and return quantized tuples for maximum efficiency
///
/// - Parameters:
/// - keys: New key data to add to cache
/// - values: New value data to add to cache
/// - Returns: Quantized tuples (keys, values) as ((weight, scales, biases), (weight, scales, biases))
func updateQuantized(keys: MLXArray, values: MLXArray) -> (
(MLXArray, MLXArray, MLXArray?), (MLXArray, MLXArray, MLXArray?)
)
/// Get current quantized state without updating
///
/// Useful for accessing cached data without adding new tokens.
/// - Returns: Current quantized state, or nil if cache is empty
func getQuantizedState() -> ((MLXArray, MLXArray, MLXArray?), (MLXArray, MLXArray, MLXArray?))?
}
/// Base cache implementation providing default behaviors
open class BaseKVCache: KVCache {
public var offset: Int = 0
public var maxSize: Int? { nil }
public func innerState() -> [MLXArray] { [] }
open func update(keys: MLXArray, values: MLXArray) -> (MLXArray, MLXArray) {
fatalError("update(keys:values:) must be implemented by subclass")
}
open var state: [MLXArray] {
get { [] }
set {
if !newValue.isEmpty {
fatalError("This cache has no state but a state was set.")
}
}
}
open var metaState: [String] {
get {
// Python base class returns empty string, but we return empty array for Swift compatibility
// This is handled in the save/load functions
[]
}
set {
if !newValue.isEmpty {
fatalError("This cache has no meta_state but a meta_state was set.")
}
}
}
open var isTrimmable: Bool { false }
@discardableResult
open func trim(_ n: Int) -> Int { 0 }
/// Default implementation for caches without special mask requirements
open func makeMask(
n: Int, windowSize: Int?, returnArray: Bool
) -> MLXFast.ScaledDotProductAttentionMaskMode {
// For single token, no mask needed
if n == 1 {
return .none
}
// For multi-token sequences
if returnArray || (windowSize != nil && n > windowSize!) {
return .array(createCausalMask(n: n, offset: offset, windowSize: windowSize))
}
return .causal
}
}
public func createCausalMask(
n: Int,
offset: Int,
windowSize: Int? = nil,
lengths: MLXArray? = nil
) -> MLXArray {
var rinds = MLXArray(Int32(0) ..< Int32(offset + n))
var linds = offset != 0 ? MLXArray(Int32(offset) ..< Int32(offset + n)) : rinds
linds = linds[0..., .newAxis]
rinds = rinds[.newAxis]
var mask = linds .>= rinds
if let windowSize {
mask = mask & (linds .< rinds + windowSize)
}
if var lengths {
lengths = lengths[0..., .newAxis, .newAxis, .newAxis]
mask = mask & (rinds .< lengths)
}
return mask
}
/// Create an attention mask using the parameters from the KVCache.
///
/// See also ``MultiHeadAttention/createAdditiveCausalMask(_:dtype:)`` -- same idea
/// but doesn't honor the cache offset.
@_disfavoredOverload
public func createAttentionMask(h: MLXArray, cache: [KVCache]?) -> MLXArray? {
let t = h.dim(1)
if t > 1 {
var offset = 0
if let c = cache?.first {
offset = c.offset
}
return createCausalMask(n: t, offset: offset)
}
return nil
}
@available(
*, deprecated,
message: "Use createAttentionMask(h:cache:windowSize:returnArray:) with a single cache instead"
)
public func createAttentionMask(h: MLXArray, cache: [KVCache]?, returnArray: Bool = false)
-> MLXFast.ScaledDotProductAttentionMaskMode
{
let t = h.dim(1)
if t > 1 {
var returnArray = returnArray
var offset = 0
var windowSize: Int? = nil
if let c = cache?.first {
offset = c.offset
if let maxSize = c.maxSize {
windowSize = maxSize
offset = min(maxSize - 1, offset)
if !returnArray {
returnArray = offset + t > maxSize
}
}
}
if returnArray {
return .array(createCausalMask(n: t, offset: offset, windowSize: windowSize))
} else {
return .causal
}
}
return .none
}
/// Create an attention mask with explicit window size parameter.
///
/// - Parameters:
/// - h: The input array (used to determine sequence length)
/// - cache: Optional single KV cache
/// - windowSize: Optional sliding window size (if provided, creates windowed attention)
/// - returnArray: Force return of array mask instead of symbolic "causal"
/// - Returns: Attention mask mode for scaled dot product attention
public func createAttentionMask(
h: MLXArray,
cache: KVCache?,
windowSize: Int? = nil,
returnArray: Bool = false
) -> MLXFast.ScaledDotProductAttentionMaskMode {
let n = h.dim(1)
// Delegate to cache's makeMask if available
if let cache = cache {
return cache.makeMask(n: n, windowSize: windowSize, returnArray: returnArray)
}
// Fallback for no cache
if n == 1 {
return .none
}
if returnArray || (windowSize != nil && n > windowSize!) {
return .array(createCausalMask(n: n, offset: 0, windowSize: windowSize))
}
return .causal
}
public func createSSMMask(h: MLXArray, cache: MambaCache?) -> MLXArray? {
if let cache {
return cache.makeMask(N: h.dim(1))
}
return nil
}
/// Standard KV cache implementation based on Python's KVCache
/// See https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/models/base.py#L11
public class KVCacheSimple: BaseKVCache, CustomDebugStringConvertible {
internal var keys: MLXArray?
internal var values: MLXArray?
public var step = 256
public override init() {
super.init()
}
public override func innerState() -> [MLXArray] {
[self.keys, self.values].compactMap { $0 }
}
public override func update(keys: MLXArray, values: MLXArray) -> (MLXArray, MLXArray) {
let previous = self.offset
let reset =
if let currentKeys = self.keys, (previous + keys.dim(2)) > currentKeys.dim(2) {
true
} else {
self.keys == nil
}
if reset {
let B = keys.dim(0)
let kvHeads = keys.dim(1)
let kHeadDim = keys.dim(3)
let vHeadDim = values.dim(3)
let nSteps = (step + keys.dim(2) - 1) / step
let kShape = [B, kvHeads, nSteps * step, kHeadDim]
let vShape = [B, kvHeads, nSteps * step, vHeadDim]
let newK = MLXArray.zeros(kShape, dtype: keys.dtype)
let newV = MLXArray.zeros(vShape, dtype: values.dtype)
if var currentKeys = self.keys, var currentValues = self.values {
if previous % step != 0 {
currentKeys = currentKeys[.ellipsis, ..<previous, 0...]
currentValues = currentValues[.ellipsis, ..<previous, 0...]
}
self.keys = concatenated([currentKeys, newK], axis: 2)
self.values = concatenated([currentValues, newV], axis: 2)
} else {
self.keys = newK
self.values = newV
}
}
self.offset += keys.dim(2)
self.keys?[.ellipsis, previous ..< self.offset, 0...] = keys
self.values?[.ellipsis, previous ..< self.offset, 0...] = values
let returnedKeys = self.keys![.ellipsis, ..<self.offset, 0...]
let returnedValues = self.values![.ellipsis, ..<self.offset, 0...]
return (returnedKeys, returnedValues)
}
public override var state: [MLXArray] {
get {
guard let keys = self.keys, let values = self.values else { return [] }
if offset == keys.dim(2) {
return [keys, values]
} else {
return [
keys[.ellipsis, ..<offset, 0...],
values[.ellipsis, ..<offset, 0...],
]
}
}
set {
guard newValue.count == 2 else {
fatalError("KVCacheSimple state must have exactly 2 arrays (keys, values)")
}
self.keys = newValue[0]
self.values = newValue[1]
self.offset = self.keys!.dim(2)
}
}
public override var metaState: [String] {
get { [] }
set {
if !newValue.isEmpty {
fatalError("KVCacheSimple should not have metaState.")
}
}
}
public override var isTrimmable: Bool { true }
@discardableResult
public override func trim(_ n: Int) -> Int {
let trimmed = min(offset, n)
offset -= trimmed
return trimmed
}
/// Convert to quantized cache for maximum efficiency
///
/// Use `updateQuantized()` and `quantizedScaledDotProductAttention()` for zero-overhead operation.
public func toQuantized(groupSize: Int = 64, bits: Int = 4) -> QuantizedKVCache {
let quantizedCache = QuantizedKVCache(groupSize: groupSize, bits: bits)
quantizedCache.offset = self.offset
if let keys = self.keys, let values = self.values {
// Quantize the current keys and values
let currentKeys = keys[.ellipsis, ..<offset, 0...]
let currentValues = values[.ellipsis, ..<offset, 0...]
let quantizedKeys = quantized(currentKeys, groupSize: groupSize, bits: bits)
let quantizedValues = quantized(currentValues, groupSize: groupSize, bits: bits)
// Set the quantized state
quantizedCache.state = [
quantizedKeys.wq, quantizedKeys.scales, quantizedKeys.biases,
quantizedValues.wq, quantizedValues.scales, quantizedValues.biases,
].compactMap { $0 }
}
return quantizedCache
}
public var debugDescription: String {
"\(String(describing: Self.self)) \(Unmanaged.passUnretained(self).toOpaque()), offset: \(offset), step: \(step), keys: \(keys?.shape.description ?? "-"), values: \(values?.shape.description ?? "-")"
}
}
/// Rotating KV cache for sliding window attention
public class RotatingKVCache: BaseKVCache, CustomDebugStringConvertible {
private var keep: Int
private var keys: MLXArray?
private var values: MLXArray?
private var maxCacheSize: Int
private var step: Int
private var idx: Int = 0
public override var maxSize: Int? { maxCacheSize }
public init(maxSize: Int, keep: Int = 0, step: Int = 256) {
self.maxCacheSize = maxSize
self.keep = keep
self.step = step
super.init()
}
public override func innerState() -> [MLXArray] {
[self.keys, self.values].compactMap { $0 }
}
private func trim(trimSize: Int, _ array: MLXArray, append: MLXArray? = nil) -> MLXArray {
var toCat: [MLXArray] = []
if trimSize > 0 {
toCat = [
array[.ellipsis, ..<keep, 0...],
array[.ellipsis, (trimSize + keep)..., 0...],
]
} else {
toCat = [array]
}
if let append {
toCat.append(append)
}
return concatenated(toCat, axis: 2)
}
private func temporalOrder(_ array: MLXArray) -> MLXArray {
// Rearrange the cache into temporal order, slicing off the end if unused
if idx == array.dim(2) {
return array
} else if idx < offset {
return concatenated(
[
array[.ellipsis, ..<keep, 0...],
array[.ellipsis, idx..., 0...],
array[.ellipsis, keep ..< idx, 0...],
], axis: 2)
} else {
return array[.ellipsis, ..<idx, 0...]
}
}
private func updateConcat(keys: MLXArray, values: MLXArray) -> (MLXArray, MLXArray) {
if self.keys == nil {
self.keys = keys
self.values = values
} else {
// Put the keys/values in temporal order to preserve context
self.keys = temporalOrder(self.keys!)
self.values = temporalOrder(self.values!)
idx = self.keys!.dim(2)
// Allow temporary cache growth during multi-token processing (e.g., prompt prefill).
// The largest size is maxCacheSize + S - 1 to ensure
// every token gets at least maxCacheSize context
let trimSize = idx - maxCacheSize + 1
self.keys = trim(trimSize: trimSize, self.keys!, append: keys)
self.values = trim(trimSize: trimSize, self.values!, append: values)
}
offset += keys.dim(2)
idx = self.keys!.dim(2)
return (self.keys!, self.values!)
}
private func updateInPlace(keys: MLXArray, values: MLXArray) -> (MLXArray, MLXArray) {
let B = keys.dim(0)
let nKVHeads = keys.dim(1)
let S = keys.dim(2)
let kHeadDim = keys.dim(3)
let vHeadDim = values.dim(3)
let prev = offset
// May not have hit the max size yet, so potentially keep growing the cache
if self.keys == nil
|| (prev >= self.keys!.dim(2) && self.keys!.dim(2) < maxCacheSize)
{
let newSize = min(step, maxCacheSize - prev)
let kShape = [B, nKVHeads, newSize, kHeadDim]
let vShape = [B, nKVHeads, newSize, vHeadDim]
let newK = MLXArray.zeros(kShape, dtype: keys.dtype)
let newV = MLXArray.zeros(vShape, dtype: values.dtype)
if let currentKeys = self.keys, let currentValues = self.values {
self.keys = concatenated([currentKeys, newK], axis: 2)
self.values = concatenated([currentValues, newV], axis: 2)
} else {
self.keys = newK
self.values = newV
}
idx = prev
}
// Trim if needed
let trimSize = self.keys!.dim(2) - maxCacheSize
if trimSize > 0 {
self.keys = trim(trimSize: trimSize, self.keys!)
self.values = trim(trimSize: trimSize, self.values!)
idx = maxCacheSize
}
// Rotate if we've hit the end
if idx == maxCacheSize {
idx = keep
}
// Assign
self.keys![.ellipsis, idx ..< (idx + S), 0...] = keys
self.values![.ellipsis, idx ..< (idx + S), 0...] = values
offset += S
idx += S
// Return the appropriate cache slice
if offset < maxCacheSize {
return (
self.keys![.ellipsis, ..<offset, 0...],
self.values![.ellipsis, ..<offset, 0...]
)
}
return (self.keys!, self.values!)
}
public override func update(keys: MLXArray, values: MLXArray) -> (MLXArray, MLXArray) {
let result =
if keys.dim(2) == 1 {
updateInPlace(keys: keys, values: values)
} else {
updateConcat(keys: keys, values: values)
}
return result
}
public override var state: [MLXArray] {
get {
guard let keys = self.keys, let values = self.values else { return [] }
if offset < keys.dim(2) {
return [
keys[.ellipsis, ..<offset, 0...],
values[.ellipsis, ..<offset, 0...],
]
} else {
return [keys, values]
}
}
set {
guard newValue.count == 2 else {
fatalError("RotatingKVCache state must have exactly 2 arrays")
}
self.keys = newValue[0]
self.values = newValue[1]
// Note: RotatingKVCache doesn't set offset from keys like KVCache does
// The offset is managed through meta_state
}
}
public override var metaState: [String] {
get {
return [String(keep), String(maxCacheSize), String(step), String(offset), String(idx)]
}
set {
guard newValue.count == 5 else {
fatalError("RotatingKVCache metaState must have exactly 5 values")
}
guard let keepVal = Int(newValue[0]),
let stepVal = Int(newValue[2]),
let offsetVal = Int(newValue[3]),
let idxVal = Int(newValue[4])
else {
fatalError("Failed to convert metaState values to integers")
}
if newValue[1] == "None" {
fatalError(
"RotatingKVCache requires a non-nil maxSize. Cannot load cache with maxSize=None."
)
}
guard let maxSizeVal = Int(newValue[1]) else {
fatalError("Failed to convert maxCacheSize '\(newValue[1])' to integer")
}
self.keep = keepVal
self.maxCacheSize = maxSizeVal
self.step = stepVal
self.offset = offsetVal
self.idx = idxVal
}
}
public override var isTrimmable: Bool {
return offset < maxCacheSize
}
@discardableResult
public override func trim(_ n: Int) -> Int {
let trimmed = min(offset, n)
offset -= trimmed
idx -= trimmed
return trimmed
}
/// Optimized mask creation for rotating cache with offset capping
public override func makeMask(
n: Int, windowSize: Int?, returnArray: Bool
) -> MLXFast.ScaledDotProductAttentionMaskMode {
if n > 1 {
// Multi-token case
let actualWindowSize = windowSize ?? maxCacheSize
let cappedOffset = min(maxCacheSize - 1, offset)
// Decide if we need an array mask
if cappedOffset + n > actualWindowSize || returnArray {
return .array(
createCausalMask(n: n, offset: cappedOffset, windowSize: actualWindowSize))
}
return .causal
} else {
// Single token case (n == 1)
guard let windowSize = windowSize else {
return .none
}
// May need a mask when window_size < max_size and cache has wrapped
if offset >= windowSize, maxCacheSize > windowSize {
var currentIdx = idx
if currentIdx >= maxCacheSize {
currentIdx = 0
}
let maskSize = offset < maxCacheSize ? offset + 1 : maxCacheSize
let mask = MLXArray(0 ..< Int32(maskSize)) .>= Int32(maskSize - windowSize)
// Roll the mask to account for rotation
let rolledMask = roll(mask, shift: currentIdx + 1)
return .array(rolledMask)
}
return .none
}
}
public var debugDescription: String {
"\(String(describing: Self.self)) offset: \(offset), maxSize: \(maxCacheSize.description), keep: \(keep), idx: \(idx)"
}
/// Convert to quantized cache
/// Note: This is complex due to the rotating nature and temporal ordering
public func toQuantized(groupSize: Int = 64, bits: Int = 4) -> QuantizedKVCache {
// For now, throw an error like the Python version does
// A full implementation would need to handle the temporal ordering correctly
fatalError(
"RotatingKVCache quantization not yet implemented - temporal ordering makes this complex"
)
// Future implementation would need to:
// 1. Put keys/values in temporal order using temporalOrder()
// 2. Quantize the temporally ordered arrays
// 3. Store metadata about rotation state
// 4. Implement corresponding dequantization with rotation restoration
}
}
/// Quantized KV cache for memory efficiency using MLX quantization
public class QuantizedKVCache: BaseKVCache, QuantizedKVCacheProtocol {
private var keys: (MLXArray, MLXArray, MLXArray?)?
private var values: (MLXArray, MLXArray, MLXArray?)?
private let step: Int
public let groupSize: Int
public let bits: Int
public let mode: QuantizationMode
public init(groupSize: Int = 64, bits: Int = 8, mode: QuantizationMode = .affine) {
self.groupSize = groupSize
self.bits = bits
self.step = 256
self.mode = mode
super.init()
}
public override func innerState() -> [MLXArray] {
var arrays: [MLXArray] = []
if let keys = keys {
arrays.append(contentsOf: [keys.0, keys.1, keys.2].compactMap { $0 })
}
if let values = values {
arrays.append(contentsOf: [values.0, values.1, values.2].compactMap { $0 })
}
return arrays
}
/// Tree map equivalent for applying function to tuple elements
private func treeMap<T>(_ transform: (MLXArray) -> T, _ tuple: (MLXArray, MLXArray, MLXArray?))
-> (T, T, T?)
{
if let biases = tuple.2 {
return (transform(tuple.0), transform(tuple.1), transform(biases))
} else {
return (transform(tuple.0), transform(tuple.1), nil)
}
}
/// Tree map for two tuples (like Python's tree_map over (keys, values))
private func treeMapPair<T>(
_ transform: (MLXArray) -> T, _ tuple1: (MLXArray, MLXArray, MLXArray?),
_ tuple2: (MLXArray, MLXArray, MLXArray?)
) -> ((T, T, T?), (T, T, T?)) {
return (treeMap(transform, tuple1), treeMap(transform, tuple2))
}
/// Create initial quantized tuples (like Python's init_quant)
private func initQuant(dim: Int, shape: [Int], dtype: DType) -> (MLXArray, MLXArray, MLXArray?)
{
// Create temporary zero arrays and quantize them using native MLX Swift
let tempArray = MLXArray.zeros(shape + [dim], dtype: dtype)
let quantized = quantized(tempArray, groupSize: groupSize, bits: bits)
return (quantized.wq, quantized.scales, quantized.biases)
}
/// Expand quantized tuple
private func expandQuant(_ quantTuple: (MLXArray, MLXArray, MLXArray?), newShape: [Int]) -> (
MLXArray, MLXArray, MLXArray?
) {
return treeMap(
{ array in
let newArray = MLXArray.zeros(newShape + [array.dim(-1)], dtype: array.dtype)
return concatenated([array, newArray], axis: -2)
}, quantTuple)
}
/// Get current quantized keys and values as tuples (efficient access)
/// - Returns: Tuple of ((keyWeight, keyScales, keyBiases), (valueWeight, valueScales, valueBiases))
public func getQuantizedState() -> (
(MLXArray, MLXArray, MLXArray?), (MLXArray, MLXArray, MLXArray?)
)? {
guard let keys = keys, let values = values else { return nil }
let trimmedKeys = treeMap({ $0[.ellipsis, ..<offset, 0...] }, keys)
let trimmedValues = treeMap({ $0[.ellipsis, ..<offset, 0...] }, values)
return (trimmedKeys, trimmedValues)
}
/// Update cache and return quantized tuples (Python's update_and_fetch)
/// This is needed because `update` in Swift must return `(MLXArray, MLXArray)`
///
/// - Parameters:
/// - keys: New key data to add to cache
/// - values: New value data to add to cache
/// - Returns: Quantized tuples (keys, values) as ((weight, scales, biases), (weight, scales, biases))
public func updateQuantized(keys: MLXArray, values: MLXArray) -> (
(MLXArray, MLXArray, MLXArray?), (MLXArray, MLXArray, MLXArray?)
) {
let B = keys.dim(0)
let nKVHeads = keys.dim(1)
let numSteps = keys.dim(2)
let kHeadDim = keys.dim(3)
let vHeadDim = values.dim(3)
let prev = offset
// Check if we need to expand the cache
if self.keys == nil || (prev + numSteps) > self.keys!.0.dim(-2) {
let newSteps = ((step + numSteps - 1) / step) * step
let shape = [B, nKVHeads, newSteps]
if let existingKeys = self.keys, let existingValues = self.values {
// Trim if needed
if prev % step != 0 {
// Use tree_map equivalent to trim both keys and values
let (trimmedKeys, trimmedValues) = treeMapPair(
{ array in
array[.ellipsis, ..<prev, 0...]
}, existingKeys, existingValues)
self.keys = trimmedKeys
self.values = trimmedValues
}
// Expand using tree_map equivalent (Python's tree_map(expand_quant, ...))
self.keys = expandQuant(self.keys!, newShape: shape)
self.values = expandQuant(self.values!, newShape: shape)
} else {
// Initialize new quantized cache
self.keys = initQuant(dim: kHeadDim, shape: shape, dtype: keys.dtype)
self.values = initQuant(dim: vHeadDim, shape: shape, dtype: keys.dtype)
}
}
offset += numSteps
let quantizedKeys = quantized(keys, groupSize: groupSize, bits: bits)
let quantizedValues = quantized(values, groupSize: groupSize, bits: bits)
// Convert named tuples to positional tuples
let qKeys = (quantizedKeys.wq, quantizedKeys.scales, quantizedKeys.biases)
let qValues = (quantizedValues.wq, quantizedValues.scales, quantizedValues.biases)
// Assign to storage
guard let currentKeys = self.keys, let currentValues = self.values else {
fatalError("Quantized cache not properly initialized")
}
// Update each component of the quantized tuples
currentKeys.0[.ellipsis, prev ..< offset, 0...] = qKeys.0
currentKeys.1[.ellipsis, prev ..< offset, 0...] = qKeys.1
if let qKeysBiases = qKeys.2 {
currentKeys.2![.ellipsis, prev ..< offset, 0...] = qKeysBiases
}
currentValues.0[.ellipsis, prev ..< offset, 0...] = qValues.0
currentValues.1[.ellipsis, prev ..< offset, 0...] = qValues.1
if let qValuesBiases = qValues.2 {
currentValues.2![.ellipsis, prev ..< offset, 0...] = qValuesBiases
}
self.keys = currentKeys
self.values = currentValues
// Return quantized tuples
let trimmedKeys = treeMap({ $0[.ellipsis, ..<offset, 0...] }, currentKeys)
let trimmedValues = treeMap({ $0[.ellipsis, ..<offset, 0...] }, currentValues)
return (trimmedKeys, trimmedValues)
}
/// This method is required by the KVCache protocol, but it is not intended to be used with QuantizedKVCache.
/// Use `updateQuantized` instead.
public override func update(keys: MLXArray, values: MLXArray) -> (MLXArray, MLXArray) {
fatalError(
"`update` was called on `QuantizedKVCache`. Use `updateQuantized` instead."
)
}
/// Array of keys and values -- this will have either 6 elements or 4 elements (if biases are nil).
public override var state: [MLXArray] {
get {
guard let keys = keys, let values = values else { return [] }
if offset < keys.0.dim(2) {
// Trim to current offset using tree_map
let trimmedKeys = treeMap({ $0[.ellipsis, ..<offset, 0...] }, keys)
let trimmedValues = treeMap({ $0[.ellipsis, ..<offset, 0...] }, values)
// Flatten tuples to array for serialization
return [
trimmedKeys.0, trimmedKeys.1, trimmedKeys.2, trimmedValues.0, trimmedValues.1,
trimmedValues.2,
].compactMap { $0 }
} else {
// Flatten tuples to array for serialization
return [keys.0, keys.1, keys.2, values.0, values.1, values.2].compactMap { $0 }
}
}
set {
switch newValue.count {
case 4:
// nil biases case
keys = (newValue[0], newValue[1], nil)
values = (newValue[2], newValue[3], nil)
case 6:
keys = (newValue[0], newValue[1], newValue[2])
values = (newValue[3], newValue[4], newValue[5])
default:
fatalError(
"QuantizedKVCache state must have exactly 6 or 4 arrays (3/2 for keys, 3/2 for values)"
)
}
}
}
public override var metaState: [String] {
get { [String(step), String(offset), String(groupSize), String(bits)] }
set {
guard newValue.count == 4 else {
fatalError("QuantizedKVCache metaState must have exactly 4 values")
}
self.offset = Int(newValue[1]) ?? 0
}
}
public override var isTrimmable: Bool { true }
@discardableResult
public override func trim(_ n: Int) -> Int {
let trimmed = min(offset, n)
offset -= trimmed
return trimmed
}
/// Convert to unquantized cache
public func toUnquantized() -> KVCacheSimple {
let simpleCache = KVCacheSimple()
simpleCache.offset = self.offset
if let keys = keys, let values = values {
// Dequantize the current state using tree_map approach
let currentKeys = treeMap({ $0[.ellipsis, ..<offset, 0...] }, keys)
let currentValues = treeMap({ $0[.ellipsis, ..<offset, 0...] }, values)
let dequantizedKeys = dequantized(
currentKeys.0, scales: currentKeys.1, biases: currentKeys.2,
groupSize: groupSize, bits: bits, mode: mode)
let dequantizedValues = dequantized(
currentValues.0, scales: currentValues.1, biases: currentValues.2,
groupSize: groupSize, bits: bits, mode: mode)
// Set the unquantized state
simpleCache.state = [dequantizedKeys, dequantizedValues]
}
return simpleCache
}
}
/// Chunked KV cache for processing large contexts in chunks
public class ChunkedKVCache: KVCacheSimple {
private var chunkSize: Int?
private var startPosition: Int = 0
public init(chunkSize: Int? = nil) {
self.chunkSize = chunkSize
super.init()
}
public func maybeTrimFront() {
guard let keys = self.keys,
let chunkSize = chunkSize,
keys.dim(2) >= chunkSize
else { return }
startPosition += keys.dim(2) - chunkSize
self.keys = keys[.ellipsis, (-chunkSize)..., 0...]
self.values = values?[.ellipsis, (-chunkSize)..., 0...]
}
public override func update(keys: MLXArray, values: MLXArray) -> (MLXArray, MLXArray) {
let prev = offset - startPosition
if self.keys == nil || (prev + keys.dim(2)) > self.keys!.dim(2) {
let B = keys.dim(0)
let kvHeads = keys.dim(1)
let kHeadDim = keys.dim(3)
let vHeadDim = values.dim(3)
let nSteps = (step + keys.dim(2) - 1) / step
let kShape = [B, kvHeads, nSteps * step, kHeadDim]
let vShape = [B, kvHeads, nSteps * step, vHeadDim]
let newK = MLXArray.zeros(kShape, dtype: keys.dtype)
let newV = MLXArray.zeros(vShape, dtype: values.dtype)
if var currentKeys = self.keys, var currentValues = self.values {
if prev % step != 0 {
currentKeys = currentKeys[.ellipsis, ..<prev, 0...]
currentValues = currentValues[.ellipsis, ..<prev, 0...]
}
self.keys = concatenated([currentKeys, newK], axis: 2)
self.values = concatenated([currentValues, newV], axis: 2)
} else {
self.keys = newK
self.values = newV
}
}
offset += keys.dim(2)
let end = offset - startPosition
self.keys![.ellipsis, prev ..< end, 0...] = keys
self.values![.ellipsis, prev ..< end, 0...] = values
return (self.keys![.ellipsis, ..<end, 0...], self.values![.ellipsis, ..<end, 0...])
}
@discardableResult
public override func trim(_ n: Int) -> Int {
let trimmed = min(offset - startPosition, n)
offset -= trimmed
return trimmed
}
public override var metaState: [String] {
get {
let chunkSizeStr = chunkSize?.description ?? "None"
return [chunkSizeStr, String(startPosition)]
}
set {
guard newValue.count == 2 else {
fatalError("ChunkedKVCache metaState must have exactly 2 values")
}
if newValue[0] == "None" {
self.chunkSize = nil
} else {
self.chunkSize = Int(newValue[0])
}
self.startPosition = Int(newValue[1]) ?? 0
}
}
}
/// Base cache for array-based state storage
public class ArraysCache: BaseKVCache {
private var cache: [MLXArray?]
private var leftPadding: MLXArray?
public init(size: Int, leftPadding: [Int]? = nil) {
self.cache = Array(repeating: nil, count: size)
self.leftPadding = leftPadding.map { MLXArray($0) }
super.init()
}
public override func innerState() -> [MLXArray] {
cache.compactMap { $0 }
}
public subscript(index: Int) -> MLXArray? {
get { cache[index] }
set { cache[index] = newValue }
}
public override var state: [MLXArray] {
get {
return cache.compactMap { $0 }
}
set {
cache = newValue.map { $0 as MLXArray? }
}
}
/// In-place filter to keep just the given indices in the cache
public func filter(batchIndices: MLXArray) {
cache = cache.map { c in
c?[batchIndices]
}
leftPadding = nil
}
/// In-place extend this cache with the other cache
public func extend(other: ArraysCache) {
cache = zip(cache, other.cache).map { (c, o) in
if let c = c, let o = o {
return MLX.concatenated([c, o])
}
return c ?? o
}
leftPadding = nil
}
/// Create attention mask based on left padding
public func makeMask(N: Int) -> MLXArray? {
if cache[0] == nil, let leftPadding = leftPadding {
return MLXArray(0 ..< N) .>= leftPadding[0..., .newAxis]
} else {
return nil
}
}
}
/// Simple cache for Mamba-style state space models
public class MambaCache: ArraysCache {
public init(leftPadding: [Int]? = nil) {
super.init(size: 2, leftPadding: leftPadding)
}
}
/// Composite cache that manages multiple sub-caches
public class CacheList: BaseKVCache {
private var caches: [KVCache]
public init(_ caches: KVCache...) {
self.caches = caches
super.init()
}
public override func innerState() -> [MLXArray] {
caches.flatMap { $0.innerState() }
}
public subscript(index: Int) -> KVCache {
return caches[index]
}
public override func update(keys: MLXArray, values: MLXArray) -> (MLXArray, MLXArray) {
fatalError("CacheList should not use update(keys:values:) - use subscript access instead")
}
public override var state: [MLXArray] {
get { caches.flatMap { $0.state } }
set {
let stateLengths = caches.map { $0.state.count }
var start = 0
for i in 0 ..< caches.count {
let length = stateLengths[i]
caches[i].state = Array(newValue[start ..< (start + length)])
start += length
}
}
}
public override var isTrimmable: Bool {
caches.allSatisfy { $0.isTrimmable }
}
@discardableResult
public override func trim(_ n: Int) -> Int {
var result = 0
for cache in caches {
result = cache.trim(n)
}
return result
}
}
// MARK: - Error Types
struct KVCacheError: Error {
let message: String
}
// MARK: - Utility Functions
/// Save a pre-computed prompt cache to a file.
///
/// - Parameters:
/// - url: The URL to the `.safetensors` file
/// - cache: The model cache state
/// - metadata: Optional metadata to save along with cache state
public func savePromptCache(
url: URL,
cache: [KVCache],
metadata: [String: String] = [:]
) throws {
let cacheData = cache.map { $0.state }
let cacheInfo = cache.map { $0.metaState }
// Use Python-compatible class names for cross-platform compatibility
let cacheClasses = cache.map { cache -> String in
switch cache {
case is KVCacheSimple:
return "KVCache" // Python uses "KVCache" for the basic cache
case is RotatingKVCache:
return "RotatingKVCache"
case is QuantizedKVCache:
return "QuantizedKVCache"
case is ChunkedKVCache:
return "ChunkedKVCache"
case is MambaCache:
return "MambaCache"
case is CacheList:
return "CacheList"
default:
return "KVCache" // Default fallback
}
}
// Flatten cache data using tree_flatten compatible structure: "i.j" format
var flattenedData: [String: MLXArray] = [:]
for (i, arrays) in cacheData.enumerated() {
for (j, array) in arrays.enumerated() {
flattenedData["\(i).\(j)"] = array
}
}
// Create cache_metadata structure compatible with Python: [cache_info, metadata, cache_classes]
var flattenedMetadata: [String: String] = [:]
// Flatten cache_info as "0.i.j" (first element of cache_metadata)
for (i, info) in cacheInfo.enumerated() {
for (j, metaValue) in info.enumerated() {
flattenedMetadata["0.\(i).\(j)"] = metaValue
}
}
// Flatten user metadata as "1.key" (second element of cache_metadata)
for (key, value) in metadata {
flattenedMetadata["1.\(key)"] = value
}
// Flatten cache_classes as "2.i" (third element of cache_metadata)
for (i, className) in cacheClasses.enumerated() {
flattenedMetadata["2.\(i)"] = className
}
try save(arrays: flattenedData, metadata: flattenedMetadata, url: url)
}
/// Load a prompt cache from a file.
///
/// - Parameters:
/// - url: The URL to the `.safetensors` file
/// - Returns: The prompt cache and the metadata
public func loadPromptCache(
url: URL
) throws -> ([KVCache], [String: String]?) {
let (arrays, metadata) = try loadArraysAndMetadata(url: url)
// Unflatten arrays using tree_unflatten compatible logic
let cacheData = unflattenArrays(arrays)
// Unflatten metadata using tree_unflatten compatible logic
let unflattenedMetadata = unflattenMetadata(metadata)
// Extract cache_info, user_metadata, and cache_classes from unflattened structure
// Structure: [cache_info, user_metadata, cache_classes]
guard unflattenedMetadata.count >= 3 else {
throw KVCacheError(message: "Invalid cache metadata format")
}
let cacheInfo = unflattenedMetadata[0] as? [[String]] ?? []
let userMetadata = unflattenedMetadata[1] as? [String: String] ?? [:]
let cacheClasses = unflattenedMetadata[2] as? [String] ?? []
guard cacheData.count == cacheInfo.count && cacheData.count == cacheClasses.count else {
throw KVCacheError(message: "Mismatch in cache counts")
}
// Reconstruct cache instances
var caches: [KVCache] = []
for i in 0 ..< cacheData.count {
let className = cacheClasses[i]
var cache: KVCache
switch className {
case "KVCache", "KVCacheSimple": // Handle both Python and Swift names
cache = KVCacheSimple()
case "RotatingKVCache":
// Parse metaState first to get maxSize, then create cache
let info = i < cacheInfo.count ? cacheInfo[i] : []
guard info.count >= 5 else {
throw KVCacheError(message: "Invalid RotatingKVCache metaState - expected 5 values")
}
if info[1] == "None" {
throw KVCacheError(
message:
"RotatingKVCache with maxSize=None is not supported. This cache was created with invalid parameters."
)
}
guard let maxSize = Int(info[1]) else {
throw KVCacheError(
message: "Failed to parse RotatingKVCache maxSize from: \(info[1])")
}
cache = RotatingKVCache(maxSize: maxSize) // Create with parsed maxSize
case "QuantizedKVCache":
cache = QuantizedKVCache()
case "ChunkedKVCache":
cache = ChunkedKVCache()
case "MambaCache":
cache = MambaCache()
case "CacheList":
// Note: CacheList requires special handling as it contains sub-caches
// For now, create an empty CacheList - this may not work correctly
// for complex cache hierarchies loaded from Python
cache = CacheList()
print("Warning: CacheList loading may not preserve sub-cache structure correctly")
default:
throw KVCacheError(message: "Unknown cache class: \(className)")
}
cache.state = cacheData[i]
if i < cacheInfo.count {
cache.metaState = cacheInfo[i]
}
caches.append(cache)
}
return (caches, userMetadata)
}
/// Unflatten arrays from tree_flatten format (e.g., "0.1", "1.0") to nested structure
private func unflattenArrays(_ flatArrays: [String: MLXArray]) -> [[MLXArray]] {
var arrayMap: [Int: [Int: MLXArray]] = [:]
// Parse all keys and organize by indices
for (key, array) in flatArrays {
let components = key.split(separator: ".")
if components.count >= 2,
let i = Int(components[0]),
let j = Int(components[1])
{
if arrayMap[i] == nil {
arrayMap[i] = [:]
}
arrayMap[i]![j] = array
}
}
// Convert to ordered array structure
var result: [[MLXArray]] = []
let maxI = arrayMap.keys.max() ?? -1
for i in 0 ... maxI {
if let innerMap = arrayMap[i] {
let maxJ = innerMap.keys.max() ?? -1
var innerArray: [MLXArray] = []
for j in 0 ... maxJ {
if let array = innerMap[j] {
innerArray.append(array)
}
}
result.append(innerArray)
} else {
result.append([])
}
}
return result
}
/// Unflatten metadata from tree_flatten format to nested structure
private func unflattenMetadata(_ flatMetadata: [String: String]) -> [Any] {
var cacheInfo: [[String]] = []
var userMetadata: [String: String] = [:]
var cacheClasses: [String] = []
for (key, value) in flatMetadata {
let components = key.split(separator: ".")
if components.count >= 3 && components[0] == "0" {
// Cache info: "0.i.j" format
if let i = Int(components[1]), let j = Int(components[2]) {
// Ensure cacheInfo is large enough
while cacheInfo.count <= i {
cacheInfo.append([])
}
// Ensure inner array is large enough
while cacheInfo[i].count <= j {
cacheInfo[i].append("")
}
cacheInfo[i][j] = value
}
} else if components.count >= 2 && components[0] == "1" {
// User metadata: "1.key" format
let metaKey = components.dropFirst().joined(separator: ".")
userMetadata[metaKey] = value
} else if components.count >= 2 && components[0] == "2" {
// Cache classes: "2.i" format
if let i = Int(components[1]) {
// Ensure cacheClasses is large enough
while cacheClasses.count <= i {
cacheClasses.append("")
}
cacheClasses[i] = value
}
}
}
return [cacheInfo, userMetadata, cacheClasses]
}
/// Construct the model's cache for use when generating.
///
/// This function will defer the cache construction to the model if it has a
/// `newCache` method, otherwise it will make a default KV cache.
public func makePromptCache(
model: any LanguageModel,
parameters: GenerateParameters? = nil
) -> [KVCache] {
// The model already conforms to LanguageModel which has newCache
// If it also conforms to KVCacheDimensionProvider, the extension will provide the implementation
return model.newCache(parameters: parameters)
}
/// Legacy function for backwards compatibility
public func makePromptCache(
model: any LanguageModel,
maxKVSize: Int? = nil
) -> [KVCache] {
let parameters = maxKVSize.map { GenerateParameters(maxKVSize: $0) }
return makePromptCache(model: model, parameters: parameters)
}
/// Fallback function to create cache when layer count is known
///
/// This function creates a default cache structure when the number of layers is known.
/// Use this when `makePromptCache` cannot determine the layer count automatically.
public func makePromptCacheWithLayerCount(
numLayers: Int,
maxKVSize: Int? = nil
) -> [KVCache] {
if let maxKVSize = maxKVSize {
return (0 ..< numLayers).map { _ in
RotatingKVCache(maxSize: maxKVSize, keep: 4)
}
} else {
return (0 ..< numLayers).map { _ in KVCacheSimple() }
}
}
/// Check if model's cache can be trimmed.
public func canTrimPromptCache(_ cache: [KVCache]) -> Bool {
return cache.allSatisfy { $0.isTrimmable }
}
/// Trim the model's cache by the given number of tokens.
///
/// This function will trim the cache if possible (in-place) and return the
/// number of tokens that were trimmed.
@discardableResult
public func trimPromptCache(_ cache: [KVCache], numTokens: Int) -> Int {
guard canTrimPromptCache(cache), !cache.isEmpty else { return 0 }
return cache.first?.trim(numTokens) ?? 0
}
// MARK: - Type Aliases
/// Standard KV cache - alias to KVCacheSimple for compatibility
public typealias StandardKVCache = KVCacheSimple
// MARK: - Quantized Attention Operations
public func quantizedScaledDotProductAttention(
queries: MLXArray,
quantizedKeys: (MLXArray, MLXArray, MLXArray?),
quantizedValues: (MLXArray, MLXArray, MLXArray?),
scale: Float,
mask: MLXFast.ScaledDotProductAttentionMaskMode = .none,
groupSize: Int = 64,
bits: Int = 8,
mode: QuantizationMode = .affine
) -> MLXArray {
let (B, nQHeads, L, D) = (queries.dim(0), queries.dim(1), queries.dim(2), queries.dim(3))
let nKVHeads = quantizedKeys.0.dim(-3)
let nRepeats = nQHeads / nKVHeads
// Scale queries
var scaledQueries = queries * scale
// Handle GQA (Grouped Query Attention)
var qKeys = quantizedKeys
var qValues = quantizedValues
if nRepeats > 1 {
scaledQueries = scaledQueries.reshaped([B, nKVHeads, nRepeats, L, D])
qKeys = (
expandedDimensions(qKeys.0, axis: -3),
expandedDimensions(qKeys.1, axis: -3),
qKeys.2 == nil ? nil : expandedDimensions(qKeys.2!, axis: -3)
)
qValues = (
expandedDimensions(qValues.0, axis: -3),
expandedDimensions(qValues.1, axis: -3),
qValues.2 == nil ? nil : expandedDimensions(qValues.2!, axis: -3)
)
}
// Compute attention scores using quantized matmul
var scores = quantizedMatmul(
scaledQueries, qKeys.0, scales: qKeys.1, biases: qKeys.2,
transpose: true, groupSize: groupSize, bits: bits,
mode: mode
)
// Apply mask
switch mask {
case .causal:
let (qL, kL) = (scores.dim(-2), scores.dim(-1))
let qIndices = MLXArray(0 ..< qL) + MLXArray(kL - qL)
let kIndices = MLXArray(0 ..< kL)
let causalMask = greaterEqual(
expandedDimensions(qIndices, axis: -1), expandedDimensions(kIndices, axis: -2))
scores = MLX.where(causalMask, scores, MLXArray(Float.leastNormalMagnitude))
case .array(let maskArray):
if maskArray.dtype == .bool {
scores = MLX.where(maskArray, scores, MLXArray(Float.leastNormalMagnitude))
} else {
scores = scores + maskArray
}
case .arrays(let maskArrays):
// Handle multiple mask arrays - just use the first one for simplicity
if let maskArray = maskArrays.first {
if maskArray.dtype == .bool {
scores = MLX.where(maskArray, scores, MLXArray(Float.leastNormalMagnitude))
} else {
scores = scores + maskArray
}
}
case .none:
break
}
let attentionWeights = softmax(scores, axis: -1)
// Compute output using quantized matmul
var output = quantizedMatmul(
attentionWeights, qValues.0, scales: qValues.1, biases: qValues.2,
transpose: false, groupSize: groupSize, bits: bits,
mode: mode
)
// Reshape output for GQA
if nRepeats > 1 {
output = output.reshaped([B, nQHeads, L, D])
}
return output
}
// MARK: - Dynamic Cache Quantization
/// Dynamically quantize KV caches during generation if conditions are met
///
/// Converts regular caches to quantized caches when:
/// - kvBits is specified
/// - The cache is not already quantized
/// - The cache offset is greater than quantizedKVStart
///
/// - Parameters:
/// - cache: Array of KV caches to potentially quantize
/// - kvBits: Number of bits for quantization (nil = no quantization)
/// - kvGroupSize: Group size for quantization
/// - quantizedKVStart: Token count threshold to begin quantizing
public func maybeQuantizeKVCache(
cache: inout [KVCache],
kvBits: Int?,
kvGroupSize: Int = 64,
quantizedKVStart: Int = 0
) {
guard let kvBits = kvBits,
!cache.isEmpty,
!(cache[0] is QuantizedKVCache),
cache[0].offset > quantizedKVStart
else {
return
}
for i in 0 ..< cache.count {
// Handle cache types that support quantization
if let simpleCache = cache[i] as? KVCacheSimple {
cache[i] = simpleCache.toQuantized(groupSize: kvGroupSize, bits: kvBits)
}
// TODO: RotatingKVCache.toQuantized() is not implemented yet, like in Python.
// When implemented, add: else if let rotatingCache = cache[i] as? RotatingKVCache { ... }
// MambaCache and CacheList don't use traditional KV quantization
}
}
Raw
main.swift
//
// main.swift
// mlx-test
//
// Created by Prince Canuma on 27/12/25.
//
import Foundation
import MLX
import HuggingFace
import Tokenizers
import MLXFast
import MLXNN
import Combine
// MARK: - StringOrNumber Helper
public enum StringOrNumber: Codable, Sendable, Equatable {
case string(String)
case int(Int)
case float(Float)
case double(Double)
public init(from decoder: Swift.Decoder) throws {
let container = try decoder.singleValueContainer()
if let stringValue = try? container.decode(String.self) {
self = .string(stringValue)
} else if let intValue = try? container.decode(Int.self) {
self = .int(intValue)
} else if let doubleValue = try? container.decode(Double.self) {
self = .double(doubleValue)
} else if let floatValue = try? container.decode(Float.self) {
self = .float(floatValue)
} else {
throw DecodingError.typeMismatch(
StringOrNumber.self,
DecodingError.Context(
codingPath: container.codingPath,
debugDescription: "Expected String, Int, Float, or Double"
)
)
}
}
public func encode(to encoder: Swift.Encoder) throws {
var container = encoder.singleValueContainer()
switch self {
case .string(let value):
try container.encode(value)
case .int(let value):
try container.encode(value)
case .float(let value):
try container.encode(value)
case .double(let value):
try container.encode(value)
}
}
public func asFloat() -> Float? {
switch self {
case .float(let value):
return value
case .double(let value):
return Float(value)
case .int(let value):
return Float(value)
case .string:
return nil
}
}
public func asString() -> String? {
switch self {
case .string(let value):
return value
default:
return nil
}
}
}
// MARK: - Configuration
public struct Qwen3Configuration: Codable, Sendable {
var hiddenSize: Int
var hiddenLayers: Int
var intermediateSize: Int
var eosTokenId: Int
var attentionHeads: Int
var rmsNormEps: Float
var vocabularySize: Int
var quantization: BaseConfiguration.Quantization?
var perLayerQuantization: BaseConfiguration.PerLayerQuantization?
var kvHeads: Int
var ropeTheta: Float = 1_000_000
var headDim: Int
var ropeScaling: [String: StringOrNumber]? = nil
var tieWordEmbeddings = false
var maxPositionEmbeddings: Int = 32768
enum CodingKeys: String, CodingKey {
case hiddenSize = "hidden_size"
case eosTokenId = "eos_token_id"
case hiddenLayers = "num_hidden_layers"
case intermediateSize = "intermediate_size"
case attentionHeads = "num_attention_heads"
case rmsNormEps = "rms_norm_eps"
case vocabularySize = "vocab_size"
case quantization = "quantization"
case quantization_config = "quantization_config"
case kvHeads = "num_key_value_heads"
case ropeTheta = "rope_theta"
case headDim = "head_dim"
case ropeScaling = "rope_scaling"
case tieWordEmbeddings = "tie_word_embeddings"
case maxPositionEmbeddings = "max_position_embeddings"
}
public init(from decoder: Swift.Decoder) throws {
// custom implementation to handle optional keys with required values
let container: KeyedDecodingContainer<Qwen3Configuration.CodingKeys> =
try decoder.container(
keyedBy: Qwen3Configuration.CodingKeys.self)
self.hiddenSize = try container.decode(
Int.self, forKey: Qwen3Configuration.CodingKeys.hiddenSize)
self.eosTokenId = try container.decode(
Int.self, forKey: Qwen3Configuration.CodingKeys.eosTokenId)
self.hiddenLayers = try container.decode(
Int.self, forKey: Qwen3Configuration.CodingKeys.hiddenLayers)
self.intermediateSize = try container.decode(
Int.self, forKey: Qwen3Configuration.CodingKeys.intermediateSize)
self.attentionHeads = try container.decode(
Int.self, forKey: Qwen3Configuration.CodingKeys.attentionHeads)
self.rmsNormEps = try container.decode(
Float.self, forKey: Qwen3Configuration.CodingKeys.rmsNormEps)
self.vocabularySize = try container.decode(
Int.self, forKey: Qwen3Configuration.CodingKeys.vocabularySize)
self.kvHeads = try container.decode(Int.self, forKey: Qwen3Configuration.CodingKeys.kvHeads)
self.ropeTheta =
try container.decodeIfPresent(
Float.self, forKey: Qwen3Configuration.CodingKeys.ropeTheta)
?? 1_000_000
self.headDim = try container.decode(
Int.self, forKey: Qwen3Configuration.CodingKeys.headDim)
self.ropeScaling = try container.decodeIfPresent(
[String: StringOrNumber].self, forKey: Qwen3Configuration.CodingKeys.ropeScaling)
self.tieWordEmbeddings =
try container.decodeIfPresent(Bool.self, forKey: .tieWordEmbeddings) ?? false
self.maxPositionEmbeddings =
try container.decodeIfPresent(Int.self, forKey: .maxPositionEmbeddings) ?? 32768
// MARK: - Decode Quantization
let baseConfig = try? BaseConfiguration(from: decoder)
self.quantization = baseConfig?.quantization
self.perLayerQuantization = baseConfig?.perLayerQuantization
}
public func encode(to encoder: Swift.Encoder) throws {
var container = encoder.container(keyedBy: CodingKeys.self)
try container.encode(hiddenSize, forKey: .hiddenSize)
try container.encode(eosTokenId, forKey: .eosTokenId)
try container.encode(hiddenLayers, forKey: .hiddenLayers)
try container.encode(intermediateSize, forKey: .intermediateSize)
try container.encode(attentionHeads, forKey: .attentionHeads)
try container.encode(rmsNormEps, forKey: .rmsNormEps)
try container.encode(vocabularySize, forKey: .vocabularySize)
try container.encode(kvHeads, forKey: .kvHeads)
try container.encode(ropeTheta, forKey: .ropeTheta)
try container.encode(headDim, forKey: .headDim)
try container.encodeIfPresent(ropeScaling, forKey: .ropeScaling)
try container.encode(tieWordEmbeddings, forKey: .tieWordEmbeddings)
try container.encode(maxPositionEmbeddings, forKey: .maxPositionEmbeddings)
try container.encodeIfPresent(quantization, forKey: .quantization)
// Note: perLayerQuantization doesn't have a direct coding key, so we skip it
}
}
// MARK: - Attention
public class Attention: Module {
let args: Qwen3Configuration
let scale: Float
@ModuleInfo(key: "q_proj") var wq: Linear
@ModuleInfo(key: "k_proj") var wk: Linear
@ModuleInfo(key: "v_proj") var wv: Linear
@ModuleInfo(key: "o_proj") var wo: Linear
@ModuleInfo(key: "q_norm") var qNorm: RMSNorm
@ModuleInfo(key: "k_norm") var kNorm: RMSNorm
let rope: RoPE
public init(_ args: Qwen3Configuration) {
self.args = args
let dim = args.hiddenSize
let heads = args.attentionHeads
let kvHeads = args.kvHeads
let headDim = args.headDim
self.scale = pow(Float(headDim), -0.5)
self._wq.wrappedValue = Linear(dim, heads * headDim, bias: false)
self._wk.wrappedValue = Linear(dim, kvHeads * headDim, bias: false)
self._wv.wrappedValue = Linear(dim, kvHeads * headDim, bias: false)
self._wo.wrappedValue = Linear(heads * headDim, dim, bias: false)
self._qNorm.wrappedValue = RMSNorm(dimensions: headDim, eps: args.rmsNormEps)
self._kNorm.wrappedValue = RMSNorm(dimensions: headDim, eps: args.rmsNormEps)
let ropeScale: Float
if let ropeScaling = args.ropeScaling, ropeScaling["type"] == .string("linear"),
let factor = ropeScaling["factor"]
{
if let v = factor.asFloat() {
ropeScale = 1 / v
} else {
fatalError("ropeScaling.factor must be a Float")
}
} else {
ropeScale = 1
}
self.rope = RoPE(
dimensions: headDim, traditional: false, base: args.ropeTheta, scale: ropeScale
)
}
public func callAsFunction(
_ x: MLXArray, mask: MLXFast.ScaledDotProductAttentionMaskMode, cache: KVCache?
) -> MLXArray {
let (B, L) = (x.dim(0), x.dim(1))
var queries = wq(x)
var keys = wk(x)
var values = wv(x)
queries = qNorm(queries.reshaped(B, L, args.attentionHeads, -1)).transposed(0, 2, 1, 3)
keys = kNorm(keys.reshaped(B, L, args.kvHeads, -1)).transposed(0, 2, 1, 3)
values = values.reshaped(B, L, args.kvHeads, -1).transposed(0, 2, 1, 3)
if let cache {
queries = rope(queries, offset: cache.offset)
keys = rope(keys, offset: cache.offset)
} else {
queries = rope(queries)
keys = rope(keys)
}
let output = MLXFast.scaledDotProductAttention(
queries: queries,
keys: keys,
values: values,
scale: scale,
mask: mask
).transposed(0, 2, 1, 3).reshaped(B, L, -1)
return wo(
output
)
}
}
// MARK: - MLP
private class MLP: Module {
@ModuleInfo(key: "gate_proj") var gate: Linear
@ModuleInfo(key: "down_proj") var down: Linear
@ModuleInfo(key: "up_proj") var up: Linear
public init(dimensions: Int, hiddenDimensions: Int) {
self._gate.wrappedValue = Linear(dimensions, hiddenDimensions, bias: false)
self._down.wrappedValue = Linear(hiddenDimensions, dimensions, bias: false)
self._up.wrappedValue = Linear(dimensions, hiddenDimensions, bias: false)
}
public func callAsFunction(_ x: MLXArray) -> MLXArray {
return down(silu(gate(x)) * up(x))
}
}
private class TransformerBlock: Module {
@ModuleInfo(key: "self_attn") var attention: Attention
let mlp: MLP
@ModuleInfo(key: "input_layernorm") var inputLayerNorm: RMSNorm
@ModuleInfo(key: "post_attention_layernorm") var postAttentionLayerNorm: RMSNorm
public init(_ args: Qwen3Configuration) {
self._attention.wrappedValue = Attention(args)
self.mlp = MLP(dimensions: args.hiddenSize, hiddenDimensions: args.intermediateSize)
self._inputLayerNorm.wrappedValue = RMSNorm(dimensions: args.hiddenSize, eps: args.rmsNormEps)
self._postAttentionLayerNorm.wrappedValue = RMSNorm(dimensions: args.hiddenSize, eps: args.rmsNormEps)
}
public func callAsFunction(
_ x: MLXArray, mask: MLXFast.ScaledDotProductAttentionMaskMode, cache: KVCache?
) -> MLXArray {
var r = attention(inputLayerNorm(x), mask: mask, cache: cache)
let h = x + r
r = mlp(postAttentionLayerNorm(h))
return h + r
}
}
private class Qwen3ModelInner: Module {
@ModuleInfo(key: "embed_tokens") var embedTokens: Embedding
fileprivate let layers: [TransformerBlock]
let norm: RMSNorm
public init(_ args: Qwen3Configuration) {
precondition(args.vocabularySize > 0)
self._embedTokens.wrappedValue = Embedding(
embeddingCount: args.vocabularySize,
dimensions: args.hiddenSize
)
self.layers = (0..<args.hiddenLayers)
.map { _ in TransformerBlock(args) }
self.norm = RMSNorm(dimensions: args.hiddenSize, eps: args.rmsNormEps)
}
public func callAsFunction(_ inputs: MLXArray, cache: [KVCache]? = nil) -> MLXArray {
var h = embedTokens(inputs)
let mask = createAttentionMask(h: h, cache: cache?.first)
for (i, layer) in layers.enumerated() {
h = layer(h, mask: mask, cache: cache?[i])
}
return norm(h)
}
}
public class Qwen3Model: Module, LanguageModel, KVCacheDimensionProvider {
public let vocabularySize: Int
public let kvHeads: [Int]
private let model: Qwen3ModelInner
let configuration: Qwen3Configuration
@ModuleInfo(key: "lm_head") var lmHead: Linear?
// KVCacheDimensionProvider conformance
public var numLayers: Int {
configuration.hiddenLayers
}
public init(_ args: Qwen3Configuration){
self.configuration = args
self.vocabularySize = args.vocabularySize
self.kvHeads = (0..<args.hiddenLayers).map {_ in args.kvHeads}
self.model = Qwen3ModelInner(args)
if !args.tieWordEmbeddings {
self._lmHead.wrappedValue = Linear(args.hiddenSize, args.vocabularySize, bias: false)
}
}
public func callAsFunction(_ inputs: MLXArray, cache: [KVCache]? = nil) -> MLXArray {
var out = model(inputs, cache: cache)
if let lmHead {
out = lmHead(out)
} else {
out = model.embedTokens.asLinear(out)
}
return out
}
public func sanitize(weights: [String: MLXArray]) -> [String: MLXArray] {
var weights = weights
if configuration.tieWordEmbeddings {
weights["lm_head.weight"] = nil
}
return weights
}
public static func fromPretrained(_ model_repo: String = "Qwen/Qwen3-0.6") async throws -> Qwen3Model {
let client = HubClient.default
let snapshotDir = FileManager.default.temporaryDirectory
let progress = Progress(totalUnitCount: 0)
Task {
for await value in progress.publisher(for: \.fractionCompleted).values {
print("Snapshot download progress: \(value * 100)%")
}
}
guard let repoID = Repo.ID(rawValue: model_repo) else {
throw NSError(domain: "Qwen3Model", code: 1, userInfo: [NSLocalizedDescriptionKey: "Invalid repository ID: \(model_repo)"])
}
let modelDir = try await client.downloadSnapshot(
of: repoID,
kind: .model,
to: snapshotDir,
revision: "main",
progressHandler: { progress in
// Accurate progress per file
print("\(progress.completedUnitCount)/\(progress.totalUnitCount) files")
})
let configPath = snapshotDir.appendingPathComponent("config.json")
let configData = try Data(contentsOf: configPath)
let config = try JSONDecoder().decode(Qwen3Configuration.self, from: configData)
let perLayerQuantization = config.perLayerQuantization
let model = Qwen3Model(config)
// Load weights from safetensors
let weights = try loadWeights(from: modelDir)
let sanitizedWeights = model.sanitize(weights: weights)
// Quantize if needed
if perLayerQuantization != nil {
print("Applying quantizaiton from config...")
if let perLayerQuant = perLayerQuantization {
print(" Per-layer: \(perLayerQuant)")
}
quantize(model: model) { path, module in
// Only quantize if scales exist for this layer
if weights["\(path).scales"] != nil {
return perLayerQuantization?.quantization(layer: path)?.asTuple
} else {
return nil
}
}
}
try model.update(parameters: ModuleParameters.unflattened(sanitizedWeights), verify: [.all])
return model
}
}
func loadWeights(from directory: URL) throws -> [String: MLXArray] {
let fileManager = FileManager.default
let files = try fileManager.contentsOfDirectory(at: directory, includingPropertiesForKeys: nil)
let safetensorFiles = files.filter { $0.pathExtension == "safetensors" }
var weights: [String: MLXArray] = [:]
for file in safetensorFiles {
let fileWeights = try MLX.loadArrays(url: file)
weights.merge(fileWeights) { _, new in new }
}
return weights
}
// Main entry point
Task {
do {
let modelRepo = "mlx-community/Qwen3-0.6B-4bit"
// let modelRepo = "Qwen/Qwen3-0.6b"
let model = try await Qwen3Model.fromPretrained(modelRepo)
let configuration = model.configuration
let eosTokenId = configuration.eosTokenId
print("Model loaded successfully!")
let tokenizer = try await AutoTokenizer.from(pretrained: modelRepo)
let tokens = tokenizer.encode(text: "<|im_start|>user\nHi there!<|im_end|>\n<|im_start|>assistant\n<think>")
var inputIds = MLXArray(tokens).expandedDimensions(axis: 0)
print(inputIds)
print("Tokenizer loaded successfully!")
let cache = (0..<configuration.hiddenLayers).map { _ in RotatingKVCache(maxSize: 2048, keep: 4)}
print("Cache initialized")
var output: MLXArray!
var nextToken: MLXArray!
// MARK: Generate
for i in 0..<500 {
output = model(inputIds, cache: cache as [any KVCache])
// Get logits for the last token in the sequence: shape [batch, vocab_size]
let lastTokenLogits = output[0..., -1, 0...]
// Get the token with highest probability: shape [batch]
nextToken = argMax(lastTokenLogits, axis: -1)
// Expand dimensions to match inputIds shape [batch, 1]
nextToken = nextToken.expandedDimensions(axis: 1)
// Eval y
eval(nextToken!)
if nextToken.squeezed().item(Int.self) == eosTokenId {
break
}
// Decode the token
let decodedText = tokenizer.decode(tokens: [nextToken.squeezed().item(Int.self)])
print(decodedText, terminator: "")
// Concatenate along sequence dimension
inputIds = concatenated([inputIds, nextToken], axis: 1)
if i != 4 {
GPU.clearCache()
}
}
print() // Final newline
print("DONE !")
} catch {
print("Failed to load model: \(error)")
}
}
// Keep the program running to allow async tasks to complete
RunLoop.main.run()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment