diff --git a/Libraries/MLXGuidedGeneration/GuidedGenerationLoop.swift b/Libraries/MLXGuidedGeneration/GuidedGenerationLoop.swift index 7d3b158d5..d44e6a96b 100644 --- a/Libraries/MLXGuidedGeneration/GuidedGenerationLoop.swift +++ b/Libraries/MLXGuidedGeneration/GuidedGenerationLoop.swift @@ -57,6 +57,8 @@ public enum GuidedGenerationLoop { /// - vocabSize: Number of tokens in the grammar's vocabulary. May differ /// from the model's logit dimension (e.g. added special tokens beyond /// the embedding size). Used to correctly interpret the grammar bitmask. + /// - kvCache: Typed KV-cache capacity and compression configuration. Prefer + /// this over the legacy scalar quantization parameters. /// - kvBits: Bit width for KV-cache quantization, or nil to disable. When /// set, the KV cache is quantized after each forward pass to reduce /// memory use, mirroring the unconstrained `TokenIterator`. nil is a @@ -87,6 +89,7 @@ public enum GuidedGenerationLoop { constraint: GrammarConstraint, maxTokens: Int, vocabSize: Int, + kvCache: KVCacheConfiguration? = nil, kvBits: Int? = nil, kvGroupSize: Int = 64, quantizedKVStart: Int = 0, @@ -99,7 +102,15 @@ public enum GuidedGenerationLoop { emit: (String) -> Bool ) throws -> Int { let model = context.model - var cache = model.newCache(parameters: nil) + let generationParameters = GenerateParameters( + kvCache: kvCache, + kvBits: kvBits, + kvGroupSize: kvGroupSize, + quantizedKVStart: quantizedKVStart) + let kvCachePlan = try generationParameters.kvCachePlan() + let cacheStorage = try kvCachePlan.validated( + KVCacheStorage( + model.newCache(parameters: generationParameters), plan: kvCachePlan)) var modelState: LMOutput.State? // Build EOS token set @@ -110,9 +121,12 @@ public enum GuidedGenerationLoop { // Prefill prompt and get first set of logits var logits: MLXArray - switch try model.prepare(input, cache: cache, state: nil, windowSize: 512) { + switch try model.prepare( + input, cache: cacheStorage.cache, state: nil, windowSize: 512) + { case .tokens(let tokens): - let result = model(tokens[text: .newAxis], cache: cache, state: nil) + let result = model( + tokens[text: .newAxis], cache: cacheStorage.cache, state: nil) modelState = result.state logits = result.logits @@ -121,6 +135,8 @@ public enum GuidedGenerationLoop { logits = result.logits } + try kvCachePlan.applyAndValidate(to: cacheStorage) + var detokenizer = NaiveStreamingDetokenizer(tokenizer: context.tokenizer) var tokenCount = 0 var grammarStopped = false @@ -361,7 +377,7 @@ public enum GuidedGenerationLoop { let tokenInput = LMInput.Text(tokens: MLXArray([ffToken])) let result = model( tokenInput[text: .newAxis], - cache: cache.isEmpty ? nil : cache, + cache: cacheStorage.cache.isEmpty ? nil : cacheStorage.cache, state: modelState ) modelState = result.state @@ -371,11 +387,7 @@ public enum GuidedGenerationLoop { } } - // Quantize the KV cache after the forward pass(es), matching the - // unconstrained TokenIterator. No-op unless `kvBits` is set. - maybeQuantizeKVCache( - cache: &cache, kvBits: kvBits, kvGroupSize: kvGroupSize, - quantizedKVStart: quantizedKVStart) + kvCachePlan.apply(to: cacheStorage) // Kick off GPU computation asynchronously asyncEval(logits) @@ -392,17 +404,13 @@ public enum GuidedGenerationLoop { let nextInput = LMInput.Text(tokens: MLXArray([Int32(token)])) let result = model( nextInput[text: .newAxis], - cache: cache.isEmpty ? nil : cache, + cache: cacheStorage.cache.isEmpty ? nil : cacheStorage.cache, state: modelState ) modelState = result.state logits = result.logits - // Quantize the KV cache after the forward pass, matching the - // unconstrained TokenIterator. No-op unless `kvBits` is set. - maybeQuantizeKVCache( - cache: &cache, kvBits: kvBits, kvGroupSize: kvGroupSize, - quantizedKVStart: quantizedKVStart) + kvCachePlan.apply(to: cacheStorage) // Kick off GPU computation asynchronously asyncEval(logits) diff --git a/Libraries/MLXLLM/Models/FalconH1.swift b/Libraries/MLXLLM/Models/FalconH1.swift index da0fd8d25..cf01943db 100644 --- a/Libraries/MLXLLM/Models/FalconH1.swift +++ b/Libraries/MLXLLM/Models/FalconH1.swift @@ -806,8 +806,8 @@ public class FalconH1Model: Module, LLMModel, KVCacheDimensionProvider { private func makeAttentionCache(parameters: GenerateParameters?) -> any KVCache { // Sliding-window attention: only the KV attention cache is bounded. The Mamba // recurrent state retains its full history because it cannot be safely windowed. - if let maxKVSize = parameters?.maxKVSize { - return RotatingKVCache(maxSize: maxKVSize, keep: 4) + if let capacity = parameters?.effectiveKVCacheCapacity { + return capacity.makeRotatingCache() } // Quantized attention cache. We create it eagerly when quantization starts at @@ -825,6 +825,10 @@ public class FalconH1Model: Module, LLMModel, KVCacheDimensionProvider { private func resolveKVQuantizationParameters(_ parameters: GenerateParameters?) -> (bits: Int, groupSize: Int)? { + if case .affine(let configuration) = parameters?.kvCache?.strategy.storage { + guard configuration.compressionStart == 0 else { return nil } + return (configuration.bits, configuration.groupSize) + } if let scheme = parameters?.kvScheme, let resolved = resolveAffineScheme(scheme) { return resolved } diff --git a/Libraries/MLXLMCommon/ChatSession.swift b/Libraries/MLXLMCommon/ChatSession.swift index 3b7963d9e..2e0d94e22 100644 --- a/Libraries/MLXLMCommon/ChatSession.swift +++ b/Libraries/MLXLMCommon/ChatSession.swift @@ -147,13 +147,41 @@ public struct SpeculativeDecodingConfig: Sendable { /// model operations. public final class ChatSession { + struct RealizedCache { + let main: KVCacheStorage + var draft: KVCacheStorage? = nil + var state: LMOutput.State? = nil + + init( + cache: consuming [KVCache], draft: consuming [KVCache]? = nil, + state: LMOutput.State? = nil, plan: KVCachePlan + ) { + self.main = KVCacheStorage(cache, plan: plan) + self.draft = draft.map { KVCacheStorage($0, plan: plan) } + self.state = state + } + + init(main: KVCacheStorage, draft: KVCacheStorage? = nil, state: LMOutput.State? = nil) { + self.main = main + self.draft = draft + self.state = state + } + + func requirePlan(_ requested: KVCachePlan) throws { + guard main.plan == requested else { + throw ChatSessionError.kvCacheConfigurationChanged( + previous: main.plan.configuration, requested: requested.configuration) + } + } + } + enum Cache { /// `state` is the per-call model state (e.g. M-RoPE rope deltas) /// from the last prefill against this cache. It must survive across /// turns: without it, a model that anchors positions on carried /// state re-derives them from a cold start on the next turn. case empty - case kvcache([KVCache], draftKVCache: [KVCache]?, state: LMOutput.State?) + case kvcache(RealizedCache) case history([Chat.Message]) } @@ -346,7 +374,11 @@ public final class ChatSession { ) { self.model = model self.instructions = instructions - self.cache = .init(.kvcache(cache, draftKVCache: nil, state: nil)) + self.cache = .init( + .kvcache( + .init( + cache: cache, + plan: (try? generateParameters.kvCachePlan()) ?? .disabled))) self.loadedDraftModel = .init(speculativeDecoding?.draftModel) self.processing = processing self.generateParameters = generateParameters @@ -392,7 +424,11 @@ public final class ChatSession { ) { self.model = ModelContainer(context: model) self.instructions = instructions - self.cache = .init(.kvcache(cache, draftKVCache: nil, state: nil)) + self.cache = .init( + .kvcache( + .init( + cache: cache, + plan: (try? generateParameters.kvCachePlan()) ?? .disabled))) self.loadedDraftModel = .init(speculativeDecoding?.draftModel) self.processing = processing self.generateParameters = generateParameters @@ -616,27 +652,33 @@ public final class ChatSession { let model = await model.perform { context in SendableBox(context.model) }.consume() + let kvCachePlan = try generateParameters.kvCachePlan() - var kvCache: [KVCache] - var draftKVCache: [KVCache]? + let kvCache: KVCacheStorage + var draftKVCache: KVCacheStorage? // Per-call model state (e.g. M-RoPE rope deltas) carried // across turns alongside the KV cache; updated after each // prefill and stored back at the end of the turn. var lmState: LMOutput.State? switch cache { case .empty: - kvCache = model.newCache(parameters: generateParameters) - cache = .kvcache(kvCache, draftKVCache: nil, state: nil) + kvCache = KVCacheStorage( + model.newCache(parameters: generateParameters), plan: kvCachePlan) + cache = .kvcache( + .init(main: kvCache)) - case .kvcache(let array, let storedDraftCache, let storedState): - kvCache = array - draftKVCache = storedDraftCache - lmState = storedState + case .kvcache(let stored): + try stored.requirePlan(kvCachePlan) + kvCache = stored.main + draftKVCache = stored.draft + lmState = stored.state case .history(let history): // the KVCache is represented by a chat history - kvCache = model.newCache(parameters: generateParameters) - cache = .kvcache(kvCache, draftKVCache: nil, state: nil) + kvCache = KVCacheStorage( + model.newCache(parameters: generateParameters), plan: kvCachePlan) + cache = .kvcache( + .init(main: kvCache)) messages.append(contentsOf: history) } @@ -663,7 +705,7 @@ public final class ChatSession { // change during decode) so the next turn — or the // next tool restart — anchors correctly. let iterator = try TokenIterator( - input: input, model: model, cache: kvCache, + input: input, model: model, cacheStorage: kvCache, state: lmState, parameters: generateParameters) lmState = iterator.state @@ -740,10 +782,13 @@ public final class ChatSession { // Allocate the draft KV cache once and reuse it across turns, // exactly like the main model's KV cache. if draftKVCache == nil { - draftKVCache = draftModel.newCache( - parameters: generateParameters) + draftKVCache = KVCacheStorage( + draftModel.newCache(parameters: generateParameters), + plan: kvCachePlan) cache = .kvcache( - kvCache, draftKVCache: draftKVCache, state: lmState) + .init( + main: kvCache, draft: draftKVCache, + state: lmState)) } let draftCache = draftKVCache! @@ -751,8 +796,8 @@ public final class ChatSession { input: input, mainModel: model, draftModel: draftModel, - mainCache: kvCache, - draftCache: draftCache, + mainCacheStorage: kvCache, + draftCacheStorage: draftCache, parameters: generateParameters, numDraftTokens: speculativeDecoding.numDraftTokens ) @@ -819,7 +864,9 @@ public final class ChatSession { // Store the carried state back alongside the KV cache so // the next turn resumes with correct position anchoring. - cache = .kvcache(kvCache, draftKVCache: draftKVCache, state: lmState) + cache = .kvcache( + .init( + main: kvCache, draft: draftKVCache, state: lmState)) continuation.finish() } @@ -872,6 +919,20 @@ public final class ChatSession { await cache.read { _ in } } + /// Return the effective per-layer state of the configured KV-cache strategy. + /// + /// The report is `nil` until a typed or legacy cache configuration exists, + /// or when the session currently stores history rather than a realized cache. + public func kvCacheRuntimeReport() async throws -> KVCacheRuntimeReport? { + let kvCachePlan = try generateParameters.kvCachePlan() + return try await cache.read { cache in + guard case .kvcache(let stored) = cache else { return nil } + try stored.requirePlan(kvCachePlan) + _ = try stored.main.plan.validated(stored.main.cache) + return stored.main.plan.report(for: stored.main.cache) + } + } + /// Visit the current cache value, if realized as a `[KVCache]`. /// /// This method is meant for test support. @@ -880,8 +941,8 @@ public final class ChatSession { { try await cache.read { cache in switch cache { - case .kvcache(let cache, _, _): - return try await body(cache) + case .kvcache(let stored): + return try await body(stored.main.cache) default: return try await body(nil) } @@ -899,8 +960,8 @@ public final class ChatSession { public func saveCache(to url: URL) async throws { try await cache.read { cache in switch cache { - case .kvcache(let cache, _, _): - try savePromptCache(url: url, cache: cache) + case .kvcache(let stored): + try savePromptCache(url: url, cache: stored.main.cache) default: throw ChatSessionError.noCacheAvailable } @@ -912,8 +973,18 @@ public final class ChatSession { public enum ChatSessionError: LocalizedError { /// ``ChatSession/saveCache(to:)`` was called before any generation occurred. case noCacheAvailable + /// The cache was realized under a different KV-cache configuration. + case kvCacheConfigurationChanged( + previous: KVCacheConfiguration?, + requested: KVCacheConfiguration? + ) public var errorDescription: String? { - "No KV cache is available. Call respond() or streamResponse() before saveCache(to:)." + switch self { + case .noCacheAvailable: + "No KV cache is available. Call respond() or streamResponse() before saveCache(to:)." + case .kvCacheConfigurationChanged: + "KV-cache configuration changed after the session cache was realized. Call clear() before continuing with the new configuration." + } } } diff --git a/Libraries/MLXLMCommon/Documentation.docc/kv-cache-quantization.md b/Libraries/MLXLMCommon/Documentation.docc/kv-cache-quantization.md index 5338e932b..543a166bf 100644 --- a/Libraries/MLXLMCommon/Documentation.docc/kv-cache-quantization.md +++ b/Libraries/MLXLMCommon/Documentation.docc/kv-cache-quantization.md @@ -5,27 +5,70 @@ key/value cache. ## Overview -At long context lengths the KV cache, not the weights, dominates memory. Two -mechanisms are available through ``GenerateParameters``: - -- **Affine quantization**: `kvBits`, `kvGroupSize`, and `quantizedKVStart` - quantize both K and V with MLX's affine scheme. The named schemes - `"affine4"` and `"affine8"` are shorthands for the same path. -- **TurboQuant schemes via `kvScheme`**: vectors are rotated with a - Walsh-Hadamard transform and quantized against a Lloyd-Max codebook. +At long context lengths the KV cache, not the weights, dominates memory. Use +``KVCacheConfiguration`` to select capacity, compression, and compatibility as +one validated value: + +```swift +let parameters = GenerateParameters( + kvCache: KVCacheConfiguration( + strategy: .turboQuant(.balanced))) +``` + +The strategy is opaque so new cache implementations can be added without +turning a public enum into an exhaustive client-side switch. The current +strategies are: + +- **Affine quantization**: ``KVCacheConfiguration/Strategy/affine(_:)`` + quantizes both K and V with MLX's affine scheme. +- **TurboQuant**: ``KVCacheConfiguration/Strategy/turboQuant(_:)`` rotates + vectors with a Walsh-Hadamard transform and quantizes them against a + Lloyd-Max codebook. TurboQuant schemes are asymmetric. Keys and values can use different precision because attention quality is far more sensitive to key error (softmax amplifies it) than to value error (linear averaging smooths it). -```swift -var parameters = GenerateParameters() -parameters.kvScheme = "turbo8v3" // recommended default -``` +The older `maxKVSize`, `kvBits`, `kvGroupSize`, `quantizedKVStart`, and +`kvScheme` fields remain available as a compatibility adapter. Do not combine +them with `kvCache`; unknown legacy scheme strings are rejected when generation +starts. + +The standalone legacy ``maybeQuantizeKVCache(cache:kvBits:kvGroupSize:quantizedKVStart:kvScheme:)`` +hook leaves custom schemes unchanged; generation APIs reject them. Prefill is unaffected: the cache stores raw fp16 during prompt processing, compresses on the first decode step, and encodes each new token incrementally afterwards. Compressed caches round-trip through prompt-cache save/restore. +## Capacity and compatibility + +A capacity creates rotating caches for model cache factories that support the +generic bounded-cache path. Rotating caches are not compressible, so combining a +capacity with compression can leave no eligible layers. Select the failure +semantics explicitly. Typed configuration defaults to +`.requireAtLeastOneLayer`, preventing a compression request from silently +becoming an all-fp16 no-op: + +- `.allowPartial` compresses eligible global layers and retains rotating layers + as fp16. +- `.requireAtLeastOneLayer` rejects an all-rotating no-op configuration. +- `.requireAllLayers` rejects any uncompressed attention layer. + +Compatibility is checked again after prefill, so unsupported realized head +shapes fail according to the selected policy. + +Applications that need a hard total-context cap and compression should omit +cache capacity, use `.requireAtLeastOneLayer`, and enforce the total token budget +before inference. A bounded compressed ring cache is not currently implemented. + +``ChatSession/kvCacheRuntimeReport()`` reports the requested configuration and +each realized layer's state, resolved strategy, and skip reason. Its aggregate +counts let an application display compressed, pending, and skipped layer counts +without assuming that a request took effect. + +A realized `ChatSession` cache is bound to its configuration. Call +``ChatSession/clear()`` before changing its capacity or strategy. + ## Scheme reference Scheme names read `turbov`; `0` means keys stay fp16. @@ -103,7 +146,7 @@ when exact fp16-parity decode matters more than footprint. actually grows, do compress, and a one-time notice lists the layers that kept fp16 rotating caches. Hybrid recurrent layers are likewise left untouched. -- Unrecognized scheme strings are ignored. +- A bounded compressed rotating cache is not implemented. - For memory estimation with wired limits see ; effective bytes per element follow from the table above (for example `turbo8v3` is about 0.73 bytes per K/V element pair average against 4 for fp16). diff --git a/Libraries/MLXLMCommon/Evaluate.swift b/Libraries/MLXLMCommon/Evaluate.swift index df34922e1..3eb387e25 100644 --- a/Libraries/MLXLMCommon/Evaluate.swift +++ b/Libraries/MLXLMCommon/Evaluate.swift @@ -64,6 +64,13 @@ public struct GenerateParameters: Sendable { /// When set, uses ``RotatingKVCache`` instead of ``KVCacheSimple`` public var maxKVSize: Int? + /// Typed key-value cache configuration. + /// + /// When set, this is the canonical cache configuration. Do not combine it + /// with `maxKVSize`, `kvBits`, `kvGroupSize`, `quantizedKVStart`, or + /// `kvScheme`. + public var kvCache: KVCacheConfiguration? + /// Number of bits to use for KV cache quantization. nil implies no cache quantization. public var kvBits: Int? @@ -88,7 +95,8 @@ public struct GenerateParameters: Sendable { /// maximum compression; K sensitivity varies by model family, so /// validate on your model (asym is the recommended starting point). /// - /// Unrecognized schemes are ignored. See `resolveTurboScheme`. + /// Unrecognized schemes are rejected when generation starts. Prefer + /// ``kvCache`` for compile-time-safe configuration. public var kvScheme: String? /// Sampling temperature @@ -131,6 +139,7 @@ public struct GenerateParameters: Sendable { public init( maxTokens: Int? = nil, maxKVSize: Int? = nil, + kvCache: KVCacheConfiguration? = nil, kvBits: Int? = nil, kvGroupSize: Int = 64, quantizedKVStart: Int = 0, @@ -150,6 +159,7 @@ public struct GenerateParameters: Sendable { ) { self.maxTokens = maxTokens self.maxKVSize = maxKVSize + self.kvCache = kvCache self.kvBits = kvBits self.kvGroupSize = kvGroupSize self.quantizedKVStart = quantizedKVStart @@ -577,18 +587,18 @@ public struct TokenIterator: TokenIteratorProtocol { public internal(set) var state: LMOutput.State? var y: LMInput.Text - var cache: [KVCache] + let cacheStorage: KVCacheStorage + var cache: [KVCache] { + get { cacheStorage.cache } + set { cacheStorage.replace(with: newValue) } + } var processor: LogitProcessor? let sampler: LogitSampler public var tokenCount = 0 public let maxTokens: Int? - // Cache quantization parameters - let kvBits: Int? - let kvGroupSize: Int - let quantizedKVStart: Int - let kvScheme: String? + var kvCachePlan: KVCachePlan { cacheStorage.plan } // Internal metrics public var promptPrefillTime: TimeInterval = 0.0 @@ -606,22 +616,12 @@ public struct TokenIterator: TokenIteratorProtocol { prompt: MLXArray, model: any LanguageModel, cache: [KVCache]? = nil, parameters: GenerateParameters ) throws { - self.model = model - self.y = .init(tokens: prompt) - self.cache = cache ?? model.newCache(parameters: parameters) - - self.processor = parameters.processor() - self.sampler = parameters.sampler() - self.maxTokens = parameters.maxTokens - - self.kvBits = parameters.kvBits - self.kvGroupSize = parameters.kvGroupSize - self.quantizedKVStart = parameters.quantizedKVStart - self.kvScheme = parameters.kvScheme - - self.promptPrefillTime = try measure { - try prepare(input: .init(text: y), windowSize: parameters.prefillStepSize) - } + let plan = try parameters.kvCachePlan() + try self.init( + input: .init(text: .init(tokens: prompt)), model: model, + cacheStorage: KVCacheStorage( + cache ?? model.newCache(parameters: parameters), plan: plan), + parameters: parameters) } /// Initialize a `TokenIterator` with the given input. @@ -643,20 +643,32 @@ public struct TokenIterator: TokenIteratorProtocol { state: LMOutput.State? = nil, parameters: GenerateParameters ) throws { + let plan = try parameters.kvCachePlan() + try self.init( + input: input, model: model, + cacheStorage: KVCacheStorage( + cache ?? model.newCache(parameters: parameters), plan: plan), + state: state, parameters: parameters) + } + + package init( + input: LMInput, model: any LanguageModel, + cacheStorage: KVCacheStorage, + state: LMOutput.State? = nil, + parameters: GenerateParameters + ) throws { + let kvCachePlan = cacheStorage.plan + let cacheStorage = try kvCachePlan.validated(cacheStorage) + self.model = model self.state = state self.y = input.text - self.cache = cache ?? model.newCache(parameters: parameters) + self.cacheStorage = cacheStorage self.processor = parameters.processor() self.sampler = parameters.sampler() self.maxTokens = parameters.maxTokens - self.kvBits = parameters.kvBits - self.kvGroupSize = parameters.kvGroupSize - self.quantizedKVStart = parameters.quantizedKVStart - self.kvScheme = parameters.kvScheme - self.promptPrefillTime = try measure { try prepare(input: input, windowSize: parameters.prefillStepSize) } @@ -683,18 +695,13 @@ public struct TokenIterator: TokenIteratorProtocol { self.model = model self.state = state self.y = input.text - self.cache = cache ?? model.newCache(parameters: nil) + self.cacheStorage = KVCacheStorage( + cache ?? model.newCache(parameters: nil), plan: .disabled) self.processor = processor self.sampler = sampler self.maxTokens = maxTokens - // No cache quantization for this direct initialization - self.kvBits = nil - self.kvGroupSize = 64 - self.quantizedKVStart = 0 - self.kvScheme = nil - self.promptPrefillTime = try measure { try prepare(input: input, windowSize: prefillStepSize) } @@ -720,6 +727,8 @@ public struct TokenIterator: TokenIteratorProtocol { break } + + try kvCachePlan.applyAndValidate(to: cacheStorage) } mutating func convertToToken(logits: MLXArray) -> MLXArray { @@ -743,13 +752,7 @@ public struct TokenIterator: TokenIteratorProtocol { self.state = result.state // Apply dynamic cache quantization after each step - maybeQuantizeKVCache( - cache: &cache, - kvBits: kvBits, - kvGroupSize: kvGroupSize, - quantizedKVStart: quantizedKVStart, - kvScheme: kvScheme - ) + kvCachePlan.apply(to: cacheStorage) return convertToToken(logits: result.logits) } @@ -825,9 +828,17 @@ public struct SpeculativeTokenIterator: TokenIteratorProtocol { let draftModel: any LanguageModel var mainState: LMOutput.State? - var mainCache: [KVCache] - var draftCache: [KVCache] - let quantizeKVCache: (inout [KVCache]) -> Void + let mainCacheStorage: KVCacheStorage + let draftCacheStorage: KVCacheStorage + var mainCache: [KVCache] { + get { mainCacheStorage.cache } + set { mainCacheStorage.replace(with: newValue) } + } + var draftCache: [KVCache] { + get { draftCacheStorage.cache } + set { draftCacheStorage.replace(with: newValue) } + } + var kvCachePlan: KVCachePlan { mainCacheStorage.plan } var processor: LogitProcessor? let sampler: LogitSampler @@ -870,16 +881,45 @@ public struct SpeculativeTokenIterator: TokenIteratorProtocol { parameters: GenerateParameters, numDraftTokens: Int ) throws { + let plan = try parameters.kvCachePlan() + try self.init( + input: input, mainModel: mainModel, draftModel: draftModel, + mainCacheStorage: KVCacheStorage( + mainCache ?? mainModel.newCache(parameters: parameters), plan: plan), + draftCacheStorage: KVCacheStorage( + draftCache ?? draftModel.newCache(parameters: parameters), plan: plan), + parameters: parameters, numDraftTokens: numDraftTokens) + } + + package init( + input: LMInput, + mainModel: any LanguageModel, + draftModel: any LanguageModel, + mainCacheStorage: KVCacheStorage, + draftCacheStorage: KVCacheStorage, + parameters: GenerateParameters, + numDraftTokens: Int + ) throws { + let kvCachePlan = mainCacheStorage.plan + precondition( + draftCacheStorage.plan == kvCachePlan, + "Speculative caches must use the same KV-cache plan") + let mainCacheStorage = try kvCachePlan.validated(mainCacheStorage) + let draftCacheStorage = try kvCachePlan.validated(draftCacheStorage) + guard + canTrimPromptCache(mainCacheStorage.cache), + canTrimPromptCache(draftCacheStorage.cache) + else { + throw KVCacheError(message: "Speculative decoding requires trimmable KV caches.") + } + self.y = input.text self.draftY = input.text self.mainModel = mainModel self.draftModel = draftModel - self.mainCache = mainCache ?? mainModel.newCache(parameters: parameters) - self.draftCache = draftCache ?? draftModel.newCache(parameters: parameters) - guard canTrimPromptCache(self.mainCache), canTrimPromptCache(self.draftCache) else { - throw KVCacheError(message: "Speculative decoding requires trimmable KV caches.") - } + self.mainCacheStorage = mainCacheStorage + self.draftCacheStorage = draftCacheStorage self.sampler = parameters.sampler() self.processor = parameters.processor() @@ -887,16 +927,6 @@ public struct SpeculativeTokenIterator: TokenIteratorProtocol { self.maxTokens = parameters.maxTokens self.numDraftTokens = numDraftTokens - self.quantizeKVCache = { cache in - maybeQuantizeKVCache( - cache: &cache, - kvBits: parameters.kvBits, - kvGroupSize: parameters.kvGroupSize, - quantizedKVStart: parameters.quantizedKVStart, - kvScheme: parameters.kvScheme - ) - } - self.promptPrefillTime = try measure { try prepare(input: input, windowSize: parameters.prefillStepSize) } @@ -933,6 +963,9 @@ public struct SpeculativeTokenIterator: TokenIteratorProtocol { draftY = .init(tokens: token) asyncEval(draftY.tokens) } + + try kvCachePlan.applyAndValidate(to: mainCacheStorage) + try kvCachePlan.applyAndValidate(to: draftCacheStorage) } /// Run one round of speculative decoding: draft, verify, accept/reject @@ -1018,8 +1051,8 @@ public struct SpeculativeTokenIterator: TokenIteratorProtocol { trimPromptCache(draftCache, numTokens: Swift.max(numDraft - accepted - 1, 0)) // Apply dynamic cache quantization after rewind - quantizeKVCache(&mainCache) - quantizeKVCache(&draftCache) + kvCachePlan.apply(to: mainCacheStorage) + kvCachePlan.apply(to: draftCacheStorage) // Set y/draftY for the next round y = .init(tokens: finalToken) diff --git a/Libraries/MLXLMCommon/KVCache.swift b/Libraries/MLXLMCommon/KVCache.swift index 6c56e92ca..f88070fbe 100644 --- a/Libraries/MLXLMCommon/KVCache.swift +++ b/Libraries/MLXLMCommon/KVCache.swift @@ -462,10 +462,13 @@ public class KVCacheSimple: BaseKVCache, CustomDebugStringConvertible { return trimmed } - /// Convert to quantized cache for maximum efficiency + /// Convert to a quantized cache for maximum efficiency. /// /// Use `updateQuantized()` and `quantizedScaledDotProductAttention()` for zero-overhead operation. - public func toQuantized(groupSize: Int = 64, bits: Int = 4) -> QuantizedKVCache { + /// + /// - Throws: If neither the requested group size nor another supported group size can + /// represent both the key and value head dimensions. + public func toQuantized(groupSize: Int = 64, bits: Int = 4) throws -> QuantizedKVCache { if let keys = self.keys, let values = self.values { // Quantize the current keys and values let currentKeys = keys[.ellipsis, .. 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" + /// Convert to a quantized cache. + /// + /// Rotating-cache quantization needs a representation that preserves temporal ordering and + /// rotation metadata. Until that representation exists, callers can recover by retaining the + /// full-precision rotating cache. + /// + /// - Throws: Always, because rotating-cache quantization is not implemented. + public func toQuantized(groupSize: Int = 64, bits: Int = 4) throws -> QuantizedKVCache { + throw KVCacheError( + message: + "RotatingKVCache quantization is not implemented because its temporal ordering requires dedicated handling." ) - - // 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 } } @@ -1474,6 +1498,22 @@ public class CacheList: BaseKVCache { } } + /// Rewrite every non-composite child while preserving its stable tree path. + func rewriteLeaves( + path: [Int], + using transform: (KVCacheLeaf) -> KVCache + ) { + caches = caches.enumerated().map { index, child in + let childPath = path + [index] + if let list = child as? CacheList { + list.rewriteLeaves(path: childPath, using: transform) + return list + } + let leaf = KVCacheLeaf(path: childPath, cache: child) + return transform(leaf) + } + } + public override func prepare(lengths: [Int]?) { caches.forEach { $0.prepare(lengths: lengths) } } @@ -1567,8 +1607,10 @@ public class CacheList: BaseKVCache { // MARK: - Error Types -struct KVCacheError: Error { +struct KVCacheError: Error, LocalizedError { let message: String + + var errorDescription: String? { message } } // MARK: - Utility Functions @@ -1644,9 +1686,6 @@ public func loadPromptCache( ) 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) @@ -1660,10 +1699,13 @@ public func loadPromptCache( let userMetadata = unflattenedMetadata[1] as? [String: String] ?? [:] let cacheClasses = unflattenedMetadata[2] as? [String] ?? [] - guard cacheData.count == cacheInfo.count && cacheData.count == cacheClasses.count else { + guard cacheInfo.count == cacheClasses.count else { throw KVCacheError(message: "Mismatch in cache counts") } + // Metadata carries the cache count even when one or more valid caches have no arrays. + let cacheData = try unflattenArrays(arrays, cacheCount: cacheClasses.count) + // Reconstruct cache instances var caches: [KVCache] = [] for i in 0 ..< cacheData.count { @@ -1689,39 +1731,71 @@ private func restoreCacheFromMetaState( ) throws -> KVCache { switch className { case "KVCache", "KVCacheSimple": + try validatePromptCache( + className: "KVCacheSimple", state: state, stateCounts: [0, 2], + metadata: metaState, metadataCounts: [1]) + guard metaState == [""] else { + throw KVCacheError( + message: + "Corrupt prompt cache: KVCacheSimple metadata must contain its single empty placeholder." + ) + } let cache = KVCacheSimple() - cache.state = state - cache.metaState = metaState + if !state.isEmpty { + cache.state = state + } return cache case "RotatingKVCache": - guard metaState.count >= 5 else { - throw KVCacheError( - message: "Invalid RotatingKVCache metaState - expected 5 values") - } - if metaState[1] == "None" { + try validatePromptCache( + className: className, state: state, stateCounts: [0, 2], + metadata: metaState, metadataCounts: [5, 6]) + let values = try promptCacheIntegers(metaState.prefix(5), className: className) + if metaState.count == 6, + RotatingKVCache.CapacityOrigin(rawValue: metaState[5]) == nil + { throw KVCacheError( message: - "RotatingKVCache with maxSize=None is not supported.") + "Corrupt prompt cache: invalid RotatingKVCache capacity origin '\(metaState[5])'." + ) } - guard let maxSize = Int(metaState[1]) else { - throw KVCacheError( - message: "Failed to parse RotatingKVCache maxSize from: \(metaState[1])") + + let cache = RotatingKVCache(maxSize: values[1]) + if !state.isEmpty { + cache.state = state } - let cache = RotatingKVCache(maxSize: maxSize) - cache.state = state cache.metaState = metaState return cache case "QuantizedKVCache": - let cache = QuantizedKVCache() - cache.state = state + try validatePromptCache( + className: className, state: state, stateCounts: [0, 4, 6], + metadata: metaState, metadataCounts: [4]) + let values = try promptCacheIntegers(metaState, className: className) + let cache = QuantizedKVCache(groupSize: values[2], bits: values[3]) + if !state.isEmpty { + cache.state = state + } cache.metaState = metaState return cache case "ChunkedKVCache": - let cache = ChunkedKVCache() - cache.state = state + try validatePromptCache( + className: className, state: state, stateCounts: [0, 2], + metadata: metaState, metadataCounts: [2]) + + let chunkSize: Int? = + if metaState[0] == "None" { + nil + } else { + try promptCacheInteger(metaState[0], className: className) + } + _ = try promptCacheInteger(metaState[1], className: className) + + let cache = ChunkedKVCache(chunkSize: chunkSize) + if !state.isEmpty { + cache.state = state + } cache.metaState = metaState return cache @@ -1759,44 +1833,74 @@ private func restoreCacheFromMetaState( } } +private func validatePromptCache( + className: String, + state: [MLXArray], + stateCounts: Set, + metadata: [String], + metadataCounts: Set +) throws { + guard stateCounts.contains(state.count), metadataCounts.contains(metadata.count), + state.allSatisfy({ $0.ndim == 4 }) + else { + throw KVCacheError( + message: "Corrupt prompt cache: invalid \(className) state or metadata shape." + ) + } +} + +private func promptCacheInteger(_ value: String, className: String) throws -> Int { + guard let value = Int(value) else { + throw KVCacheError( + message: "Corrupt prompt cache: \(className) metadata must contain integers." + ) + } + return value +} + +private func promptCacheIntegers( + _ metadata: some Collection, className: String +) throws -> [Int] { + try metadata.map { try promptCacheInteger($0, className: className) } +} + /// Unflatten arrays from tree_flatten format (e.g., "0.1", "1.0") to nested structure -private func unflattenArrays(_ flatArrays: [String: MLXArray]) -> [[MLXArray]] { +private func unflattenArrays( + _ flatArrays: [String: MLXArray], + cacheCount: Int +) throws -> [[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, + guard components.count == 2, let i = Int(components[0]), - let j = Int(components[1]) - { - if arrayMap[i] == nil { - arrayMap[i] = [:] - } - arrayMap[i]![j] = array + let j = Int(components[1]), + (0 ..< cacheCount).contains(i), + j >= 0 + else { + throw KVCacheError( + message: "Corrupt prompt cache: invalid array key '\(key)'.") } + arrayMap[i, default: [:]][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) - } + return try (0 ..< cacheCount).map { cacheIndex in + guard let arrays = arrayMap[cacheIndex], !arrays.isEmpty else { return [] } + var result: [MLXArray] = [] + result.reserveCapacity(arrays.count) + for arrayIndex in 0 ..< arrays.count { + guard let array = arrays[arrayIndex] else { + throw KVCacheError( + message: + "Corrupt prompt cache: cache \(cacheIndex) has non-contiguous array indices." + ) } - result.append(innerArray) - } else { - result.append([]) + result.append(array) } + return result } - - return result } /// Unflatten metadata from tree_flatten format to nested structure @@ -2035,76 +2139,40 @@ public func maybeQuantizeKVCache( quantizedKVStart: Int = 0, kvScheme: String? = nil ) { - // TurboQuant schemes convert eligible layers to TurboQuantKVCache - // (handled in TurboQuantKVCache.swift to keep this file scheme-agnostic). - if let scheme = kvScheme, let turbo = resolveTurboScheme(scheme) { - maybeTurboQuantizeKVCache( - cache: &cache, - keyBits: turbo.keyBits, - valueBits: turbo.valueBits, - quantizedKVStart: quantizedKVStart - ) - return - } - - // Resolve effective bits: kvScheme overrides kvBits. - let effectiveBits: Int - let effectiveGroupSize: Int - if let scheme = kvScheme, let resolved = resolveAffineScheme(scheme) { - effectiveBits = resolved.bits - effectiveGroupSize = resolved.groupSize - } else if let kvBits { - effectiveBits = kvBits - effectiveGroupSize = kvGroupSize - } else { - return - } - - /// Recursively decide whether a cache (or any of its children) is eligible for - /// quantization: it must be a plain ``KVCacheSimple`` that is not already - /// quantized and whose offset has crossed the requested start threshold. - func isQuantizable(_ cache: KVCache) -> Bool { - if let list = cache as? CacheList { - return list.children.contains(where: isQuantizable) - } - return cache is KVCacheSimple - && !(cache is QuantizedKVCache) - && cache.offset > quantizedKVStart - } - - guard cache.contains(where: isQuantizable) else { + if let kvScheme, + resolveAffineScheme(kvScheme) == nil, + resolveTurboScheme(kvScheme) == nil + { return } + let parameters = GenerateParameters( + kvBits: kvBits, + kvGroupSize: kvGroupSize, + quantizedKVStart: quantizedKVStart, + kvScheme: kvScheme) + guard let plan = try? parameters.kvCachePlan() else { return } + plan.apply(to: &cache) +} - /// Attempt to convert a single cache to its quantized form. Returns the - /// original cache if conversion is not possible for this entry. - func quantize(_ cache: KVCacheSimple) -> KVCache { - let state = cache.state - if state.count == 2 { - let keyHeadDim = state[0].dim(3) - let valueHeadDim = state[1].dim(3) - guard - resolvedKVQuantizationGroupSize( - requested: effectiveGroupSize, - keyHeadDim: keyHeadDim, - valueHeadDim: valueHeadDim - ) != nil - else { - return cache - } +@discardableResult +func maybeAffineQuantizeKVCache( + cache: inout [KVCache], + bits: Int, + groupSize: Int, + compressionStart: Int +) -> Bool { + var awaitsCompressionStart = false + KVCacheTree.rewrite(&cache) { leaf in + guard case .simple(let simple) = leaf.kind else { return leaf.cache } + guard simple.offset > compressionStart else { + awaitsCompressionStart = true + return simple } - return cache.toQuantized(groupSize: effectiveGroupSize, bits: effectiveBits) - } - for i in 0 ..< cache.count { - if let list = cache[i] as? CacheList { - list.mapChildren { child in - guard let simpleCache = child as? KVCacheSimple else { return child } - return quantize(simpleCache) - } - } else if let simpleCache = cache[i] as? KVCacheSimple { - cache[i] = quantize(simpleCache) + guard let quantized = try? simple.toQuantized(groupSize: groupSize, bits: bits) else { + return simple } - // TODO: RotatingKVCache.toQuantized() is not implemented yet, like in Python. + return quantized } + return !awaitsCompressionStart } diff --git a/Libraries/MLXLMCommon/KVCacheConfiguration.swift b/Libraries/MLXLMCommon/KVCacheConfiguration.swift new file mode 100644 index 000000000..2c1702e5a --- /dev/null +++ b/Libraries/MLXLMCommon/KVCacheConfiguration.swift @@ -0,0 +1,343 @@ +// Copyright © 2026 Apple Inc. + +import Foundation + +/// A complete request for key-value cache storage during generation. +/// +/// This keeps cache capacity, compression, and compatibility semantics in one +/// value so every generation path forwards the same configuration. +public struct KVCacheConfiguration: Sendable, Hashable { + public var capacity: Capacity? + public var strategy: Strategy + public var compatibility: CompatibilityPolicy + + public init( + capacity: Capacity? = nil, + strategy: Strategy = .fullPrecision, + compatibility: CompatibilityPolicy = .requireAtLeastOneLayer + ) { + self.capacity = capacity + self.strategy = strategy + self.compatibility = compatibility + } + + /// Maximum resident size for caches whose capacity is caller-configurable. + /// + /// Model-native sliding-window layers retain their architecture-defined + /// window and prefix behavior. This value bounds the remaining attention + /// layers created by the model's `newCache(parameters:)` implementation. + public struct Capacity: Sendable, Hashable { + public let maxTokens: Int + public let preservedPrefixTokens: Int + + public init(maxTokens: Int, preservedPrefixTokens: Int = 4) throws { + guard maxTokens > 0 else { + throw KVCacheConfigurationError.invalidCapacity(maxTokens) + } + guard preservedPrefixTokens >= 0, preservedPrefixTokens < maxTokens else { + throw KVCacheConfigurationError.invalidPreservedPrefix( + preservedPrefixTokens, capacity: maxTokens) + } + self.maxTokens = maxTokens + self.preservedPrefixTokens = preservedPrefixTokens + } + + package init(uncheckedMaxTokens maxTokens: Int, preservedPrefixTokens: Int = 4) { + self.maxTokens = maxTokens + self.preservedPrefixTokens = preservedPrefixTokens + } + } + + /// Behavior when only part of a model's cache topology supports the strategy. + public enum CompatibilityPolicy: Sendable, Hashable { + /// Apply the strategy where supported and retain other layers unchanged. + case allowPartial + /// Reject the request unless at least one attention layer can use it. + case requireAtLeastOneLayer + /// Reject the request unless every attention layer can use it. + case requireAllLayers + } + + /// Opaque cache strategy. New strategies can be added without introducing + /// public enum cases that clients must exhaustively switch over. + public struct Strategy: Sendable, Hashable { + package enum Storage: Sendable, Hashable { + case fullPrecision + case affine(AffineKVCacheConfiguration) + case turboQuant(TurboQuantKVCacheConfiguration) + } + + package let storage: Storage + + private init(storage: Storage) { + self.storage = storage + } + + public static let fullPrecision = Strategy(storage: .fullPrecision) + + public static func affine(_ configuration: AffineKVCacheConfiguration) -> Strategy { + Strategy(storage: .affine(configuration)) + } + + public static func turboQuant( + _ configuration: TurboQuantKVCacheConfiguration + ) -> Strategy { + Strategy(storage: .turboQuant(configuration)) + } + + /// Stable algorithm identity for diagnostics and persistence adapters. + public var identifier: KVCacheStrategyIdentifier { + switch storage { + case .fullPrecision: .fullPrecision + case .affine: .affine + case .turboQuant: .turboQuant + } + } + + package var compressionStart: Int { + switch storage { + case .fullPrecision: 0 + case .affine(let configuration): configuration.compressionStart + case .turboQuant(let configuration): configuration.compressionStart + } + } + } +} + +/// Stable, open-ended identity for a resolved cache strategy. +public struct KVCacheStrategyIdentifier: Sendable, Hashable, CustomStringConvertible { + public let rawValue: String + + private init(_ rawValue: String) { + self.rawValue = rawValue + } + + public static let fullPrecision = KVCacheStrategyIdentifier("full-precision") + public static let affine = KVCacheStrategyIdentifier("affine") + public static let turboQuant = KVCacheStrategyIdentifier("turbo-quant") + + public var description: String { rawValue } +} + +/// Affine cache compression settings. +public struct AffineKVCacheConfiguration: Sendable, Hashable { + private static let supportedBitWidths: Set = [2, 3, 4, 5, 6, 8] + + public let bits: Int + public let groupSize: Int + public let compressionStart: Int + + public init(bits: Int, groupSize: Int = 64, compressionStart: Int = 0) throws { + guard Self.supportedBitWidths.contains(bits) else { + throw KVCacheConfigurationError.invalidAffineBits(bits) + } + guard groupSize > 0 else { + throw KVCacheConfigurationError.invalidGroupSize(groupSize) + } + guard compressionStart >= 0 else { + throw KVCacheConfigurationError.invalidCompressionStart(compressionStart) + } + self.bits = bits + self.groupSize = groupSize + self.compressionStart = compressionStart + } + + package init(uncheckedBits bits: Int, groupSize: Int, compressionStart: Int) { + self.bits = bits + self.groupSize = groupSize + self.compressionStart = compressionStart + } + + public static let fourBit = AffineKVCacheConfiguration( + uncheckedBits: 4, groupSize: 64, compressionStart: 0) + public static let eightBit = AffineKVCacheConfiguration( + uncheckedBits: 8, groupSize: 64, compressionStart: 0) +} + +/// TurboQuant cache compression settings. +public struct TurboQuantKVCacheConfiguration: Sendable, Hashable { + /// Opaque key encoding precision. Additional encodings can be introduced + /// without adding public enum cases that clients must exhaustively handle. + public struct KeyPrecision: Sendable, Hashable { + package let bitWidth: Int + + private init(bitWidth: Int) { + self.bitWidth = bitWidth + } + + /// Keep keys in FP16; only values are compressed. + public static let fp16 = KeyPrecision(bitWidth: 0) + /// TurboQuant-compressed keys. + public static let twoBit = KeyPrecision(bitWidth: 2) + public static let threeBit = KeyPrecision(bitWidth: 3) + public static let fourBit = KeyPrecision(bitWidth: 4) + /// Affine 8-bit keys with TurboQuant-compressed values. + public static let affineEightBit = KeyPrecision(bitWidth: 8) + + package init?(legacyBitWidth: Int) { + switch legacyBitWidth { + case 0: self = .fp16 + case 2: self = .twoBit + case 3: self = .threeBit + case 4: self = .fourBit + case 8: self = .affineEightBit + default: return nil + } + } + } + + /// Opaque value encoding precision. + public struct ValuePrecision: Sendable, Hashable { + package let bitWidth: Int + + private init(bitWidth: Int) { + self.bitWidth = bitWidth + } + + public static let twoBit = ValuePrecision(bitWidth: 2) + public static let threeBit = ValuePrecision(bitWidth: 3) + public static let fourBit = ValuePrecision(bitWidth: 4) + + package init?(legacyBitWidth: Int) { + switch legacyBitWidth { + case 2: self = .twoBit + case 3: self = .threeBit + case 4: self = .fourBit + default: return nil + } + } + } + + public let keyPrecision: KeyPrecision + public let valuePrecision: ValuePrecision + public let compressionStart: Int + + public init( + keyPrecision: KeyPrecision, + valuePrecision: ValuePrecision, + compressionStart: Int = 0 + ) throws { + guard compressionStart >= 0 else { + throw KVCacheConfigurationError.invalidCompressionStart(compressionStart) + } + self.keyPrecision = keyPrecision + self.valuePrecision = valuePrecision + self.compressionStart = compressionStart + } + + private init( + uncheckedKeyPrecision keyPrecision: KeyPrecision, + valuePrecision: ValuePrecision, + compressionStart: Int + ) { + self.keyPrecision = keyPrecision + self.valuePrecision = valuePrecision + self.compressionStart = compressionStart + } + + /// FP16 keys and 4-bit values; the conservative quality-first preset. + public static let qualityFirst = TurboQuantKVCacheConfiguration( + uncheckedKeyPrecision: .fp16, valuePrecision: .fourBit, compressionStart: 0) + + /// Affine 8-bit keys and 3-bit values; the recommended general preset. + public static let balanced = TurboQuantKVCacheConfiguration( + uncheckedKeyPrecision: .affineEightBit, + valuePrecision: .threeBit, + compressionStart: 0) + + /// Affine 8-bit keys and 2-bit values for memory-bound long contexts. + public static let memoryFirst = TurboQuantKVCacheConfiguration( + uncheckedKeyPrecision: .affineEightBit, + valuePrecision: .twoBit, + compressionStart: 0) +} + +/// Validation failures for typed or legacy KV-cache configuration. +public enum KVCacheConfigurationError: Error, Sendable, Equatable, LocalizedError { + case conflictingLegacyConfiguration + case invalidCapacity(Int) + case invalidPreservedPrefix(Int, capacity: Int) + case invalidAffineBits(Int) + case invalidGroupSize(Int) + case invalidCompressionStart(Int) + case unsupportedLegacyScheme(String) + case incompatibleCapacity(expected: Int, count: Int) + case noCompatibleLayers(strategy: KVCacheStrategyIdentifier) + case incompatibleLayers(strategy: KVCacheStrategyIdentifier, count: Int) + + public var errorDescription: String? { + switch self { + case .conflictingLegacyConfiguration: + "Set either GenerateParameters.kvCache or the legacy KV-cache fields, not both." + case .invalidCapacity(let value): + "KV-cache capacity must be positive; received \(value)." + case .invalidPreservedPrefix(let value, let capacity): + "Preserved prefix \(value) must be non-negative and smaller than capacity \(capacity)." + case .invalidAffineBits(let value): + "Affine KV-cache bit width must be one of 2, 3, 4, 5, 6, or 8; received \(value)." + case .invalidGroupSize(let value): + "KV-cache group size must be positive; received \(value)." + case .invalidCompressionStart(let value): + "KV-cache compression start must be non-negative; received \(value)." + case .unsupportedLegacyScheme(let value): + "Unsupported legacy KV-cache scheme: \(value)." + case .incompatibleCapacity(let expected, let count): + "KV-cache capacity \(expected) is not realized by \(count) attention layer(s)." + case .noCompatibleLayers(let strategy): + "KV-cache strategy \(strategy) is not supported by any attention layer." + case .incompatibleLayers(let strategy, let count): + "KV-cache strategy \(strategy) is unsupported by \(count) attention layer(s)." + } + } +} + +/// Effective per-layer state for a configured KV cache. +public struct KVCacheRuntimeReport: Sendable, Hashable { + public struct Layer: Sendable, Hashable { + public enum State: Sendable, Hashable { + case active + case pending + case skipped + case notApplicable + } + + public enum Reason: Sendable, Hashable { + case awaitingCompressionStart + case boundaryProtection + case slidingWindow + case unsupportedShape + case differentStrategy + case nonAttentionState + } + + public let path: [Int] + public let state: State + public let resolvedStrategy: KVCacheStrategyIdentifier? + public let reason: Reason? + } + + public let requestedConfiguration: KVCacheConfiguration + public let layers: [Layer] + + public var compressedLayerCount: Int { + layers.filter { + $0.state == .active && $0.resolvedStrategy != .fullPrecision + }.count + } + + public var pendingLayerCount: Int { + layers.filter { $0.state == .pending }.count + } + + public var skippedLayerCount: Int { + layers.filter { $0.state == .skipped }.count + } +} + +/// The observed outcome of applying a cache configuration to a realized cache. +public struct KVCacheApplicationResult: Sendable, Hashable { + public let convertedLayerCount: Int + public let alreadyCompatibleLayerCount: Int + public let pendingLayerCount: Int + public let skipped: [KVCacheRuntimeReport.Layer] +} diff --git a/Libraries/MLXLMCommon/KVCachePlan.swift b/Libraries/MLXLMCommon/KVCachePlan.swift new file mode 100644 index 000000000..9c5c9b309 --- /dev/null +++ b/Libraries/MLXLMCommon/KVCachePlan.swift @@ -0,0 +1,197 @@ +// Copyright © 2026 Apple Inc. + +/// Immutable, resolved cache behavior for one generation request. +/// +/// A disabled plan is a first-class value, so generation paths do not need to +/// repeat optional checks around validation and dynamic compression. +package struct KVCachePlan: Sendable, Equatable { + package static let disabled = KVCachePlan(configuration: nil) + + package let configuration: KVCacheConfiguration? + + package init(configuration: KVCacheConfiguration?) { + self.configuration = configuration + } + + package func validated(_ cache: [KVCache]) throws -> [KVCache] { + if let configuration { + try validateKVCacheCompatibility(cache, configuration: configuration) + } + return cache + } + + package func validated(_ storage: KVCacheStorage) throws -> KVCacheStorage { + precondition(storage.plan == self, "KVCacheStorage used with a different plan") + storage.cache = try validated(storage.cache) + return storage + } + + package func apply(to cache: inout [KVCache]) { + guard let configuration else { return } + _ = applyKVCacheConfigurationFast(cache: &cache, configuration: configuration) + } + + /// Apply dynamic conversion to shared cache storage. + /// + /// Once every eligible layer has either converted or reached a terminal + /// unsupported state, later decode steps only compare the completed plan. + package func apply(to storage: KVCacheStorage) { + precondition(storage.plan == self, "KVCacheStorage used with a different plan") + guard !storage.isApplicationTerminal else { return } + guard let configuration else { + storage.isApplicationTerminal = true + return + } + + if applyKVCacheConfigurationFast( + cache: &storage.cache, configuration: configuration) + { + storage.isApplicationTerminal = true + } + } + + @discardableResult + package func applyAndValidate( + to cache: inout [KVCache] + ) throws -> KVCacheApplicationResult? { + guard let configuration else { return nil } + return try applyKVCacheConfiguration(cache: &cache, configuration: configuration) + } + + @discardableResult + package func applyAndValidate( + to storage: KVCacheStorage + ) throws -> KVCacheApplicationResult? { + precondition(storage.plan == self, "KVCacheStorage used with a different plan") + guard !storage.isApplicationTerminal else { return nil } + guard let configuration else { + storage.isApplicationTerminal = true + return nil + } + + let application = try applyKVCacheConfigurationValidated( + cache: &storage.cache, configuration: configuration) + if application.isTerminal { + storage.isApplicationTerminal = true + } + return application.result + } + + package func report(for cache: [KVCache]) -> KVCacheRuntimeReport? { + configuration.map { kvCacheRuntimeReport(cache: cache, configuration: $0) } + } +} + +/// Shared ownership for a realized cache and its dynamic application state. +/// +/// Cache arrays have value semantics, but dynamic compression replaces array +/// elements. Sharing this storage keeps sessions and iterators on the same +/// realized array while the cache objects themselves remain reference types. +/// Access is externally serialized by generation/session ownership. +package final class KVCacheStorage: @unchecked Sendable { + package var cache: [KVCache] { + didSet { isApplicationTerminal = false } + } + package let plan: KVCachePlan + package fileprivate(set) var isApplicationTerminal = false + + package init(_ cache: [KVCache], plan: KVCachePlan) { + self.cache = cache + self.plan = plan + } + + package func replace(with cache: [KVCache]) { + self.cache = cache + } +} + +extension KVCacheConfiguration.Capacity { + /// Construct the bounded cache represented by this value. + package func makeRotatingCache() -> RotatingKVCache { + let cache = RotatingKVCache(maxSize: maxTokens, keep: preservedPrefixTokens) + cache.capacityOrigin = .requested + return cache + } +} + +extension GenerateParameters { + package var effectiveKVCacheCapacity: KVCacheConfiguration.Capacity? { + if let capacity = kvCache?.capacity { + return capacity + } + guard let maxKVSize else { return nil } + return .init( + uncheckedMaxTokens: maxKVSize, + preservedPrefixTokens: min(4, max(0, maxKVSize - 1))) + } + + package func kvCachePlan() throws -> KVCachePlan { + KVCachePlan(configuration: try resolvedKVCacheConfiguration()) + } + + package func resolvedKVCacheConfiguration() throws -> KVCacheConfiguration? { + if let kvCache { + guard !hasLegacyKVCacheOverrides else { + throw KVCacheConfigurationError.conflictingLegacyConfiguration + } + return kvCache + } + + let capacity = try legacyKVCacheCapacity() + let strategy = try legacyKVCacheStrategy() + guard capacity != nil || strategy != .fullPrecision else { return nil } + return KVCacheConfiguration( + capacity: capacity, + strategy: strategy, + compatibility: .allowPartial) + } + + private var hasLegacyKVCacheOverrides: Bool { + maxKVSize != nil || kvBits != nil || kvScheme != nil || kvGroupSize != 64 + || quantizedKVStart != 0 + } + + private func legacyKVCacheCapacity() throws -> KVCacheConfiguration.Capacity? { + guard let maxKVSize else { return nil } + return try .init( + maxTokens: maxKVSize, + preservedPrefixTokens: min(4, max(0, maxKVSize - 1))) + } + + private func legacyKVCacheStrategy() throws -> KVCacheConfiguration.Strategy { + if let kvScheme { + return try strategy(forLegacyScheme: kvScheme) + } + guard let kvBits else { return .fullPrecision } + return .affine( + try .init( + bits: kvBits, + groupSize: kvGroupSize, + compressionStart: quantizedKVStart)) + } + + private func strategy( + forLegacyScheme scheme: String + ) throws -> KVCacheConfiguration.Strategy { + if let affine = resolveAffineScheme(scheme) { + return .affine( + try .init( + bits: affine.bits, + groupSize: affine.groupSize, + compressionStart: quantizedKVStart)) + } + if let turbo = resolveTurboScheme(scheme), + let keyPrecision = TurboQuantKVCacheConfiguration.KeyPrecision( + legacyBitWidth: turbo.keyBits), + let valuePrecision = TurboQuantKVCacheConfiguration.ValuePrecision( + legacyBitWidth: turbo.valueBits) + { + return .turboQuant( + try .init( + keyPrecision: keyPrecision, + valuePrecision: valuePrecision, + compressionStart: quantizedKVStart)) + } + throw KVCacheConfigurationError.unsupportedLegacyScheme(scheme) + } +} diff --git a/Libraries/MLXLMCommon/KVCacheRuntime.swift b/Libraries/MLXLMCommon/KVCacheRuntime.swift new file mode 100644 index 000000000..a2dae2896 --- /dev/null +++ b/Libraries/MLXLMCommon/KVCacheRuntime.swift @@ -0,0 +1,231 @@ +// Copyright © 2026 Apple Inc. + +/// Apply a typed cache strategy to every eligible cache layer. +/// +/// This is the single runtime dispatch point used by autoregressive, +/// speculative, guided, and diagnostic generation paths. +public func applyKVCacheConfiguration( + cache: inout [KVCache], + configuration: KVCacheConfiguration +) throws -> KVCacheApplicationResult { + try applyKVCacheConfigurationValidated( + cache: &cache, configuration: configuration + ).result +} + +/// Validate before mutation, then produce the diagnostic application result. +/// This keeps throwing application transactional for all validation failures. +package func applyKVCacheConfigurationValidated( + cache: inout [KVCache], + configuration: KVCacheConfiguration +) throws -> (result: KVCacheApplicationResult, isTerminal: Bool) { + try validateKVCacheCompatibility(cache, configuration: configuration) + let before = KVCacheTree.leaves(in: cache) + let isTerminal = applyKVCacheConfigurationFast( + cache: &cache, configuration: configuration) + + let after = KVCacheTree.leaves(in: cache) + let report = kvCacheRuntimeReport(cache: cache, configuration: configuration) + let converted = zip(before, after).count { original, current in + guard case .simple = original.kind else { return false } + return current.isCompressed + } + let protectedPaths = protectedPaths(for: before, configuration: configuration) + let alreadyCompatible = before.count { leaf in + if case .simple = leaf.kind { + return configuration.strategy.identifier == .fullPrecision + } else { + return leaf.supports(configuration, protectedPaths: protectedPaths) == true + } + } + return ( + KVCacheApplicationResult( + convertedLayerCount: converted, + alreadyCompatibleLayerCount: alreadyCompatible, + pendingLayerCount: report.pendingLayerCount, + skipped: report.layers.filter { $0.state == .skipped }), + isTerminal + ) +} + +/// Minimal decode-loop application. Returns true once no simple layer can +/// become newly eligible on a later token. +@discardableResult +package func applyKVCacheConfigurationFast( + cache: inout [KVCache], + configuration: KVCacheConfiguration +) -> Bool { + switch configuration.strategy.storage { + case .fullPrecision: + true + case .affine(let affine): + maybeAffineQuantizeKVCache( + cache: &cache, + bits: affine.bits, + groupSize: affine.groupSize, + compressionStart: affine.compressionStart) + case .turboQuant(let turbo): + maybeTurboQuantizeKVCache( + cache: &cache, + keyBits: turbo.keyPrecision.bitWidth, + valueBits: turbo.valuePrecision.bitWidth, + quantizedKVStart: turbo.compressionStart) + } +} + +/// Validate the requested compatibility policy against a concrete model cache. +/// +/// Full-precision and non-attention state entries do not participate. Affine +/// caches are accepted for TurboQuant because aggressive configurations use +/// them as an intentional boundary-layer fallback. +public func validateKVCacheCompatibility( + _ cache: [KVCache], + configuration: KVCacheConfiguration +) throws { + let leaves = KVCacheTree.leaves(in: cache) + + if let capacity = configuration.capacity { + let incompatibleCapacityCount = leaves.count { leaf in + switch leaf.kind { + case .recurrent: + return false + case .rotating(let rotating) where rotating.capacityOrigin == .modelNative: + // Architecture-defined sliding windows are independent of the + // caller-configurable capacity used for global attention. + return false + case .rotating(let rotating): + return rotating.maxSize != capacity.maxTokens + || rotating.preservedPrefixTokens != capacity.preservedPrefixTokens + default: + return true + } + } + guard incompatibleCapacityCount == 0 else { + throw KVCacheConfigurationError.incompatibleCapacity( + expected: capacity.maxTokens, count: incompatibleCapacityCount) + } + } + + if configuration.strategy.identifier == .fullPrecision { + let compressedCount = leaves.count { $0.isCompressed } + guard compressedCount == 0 else { + throw KVCacheConfigurationError.incompatibleLayers( + strategy: .fullPrecision, count: compressedCount) + } + return + } + + let protectedPaths = protectedPaths(for: leaves, configuration: configuration) + let support = leaves.compactMap { + $0.supports(configuration, protectedPaths: protectedPaths) + } + let compatibleCount = support.count { $0 } + let incompatibleCount = support.count { !$0 } + + switch configuration.compatibility { + case .allowPartial: + return + case .requireAtLeastOneLayer: + guard compatibleCount > 0 else { + throw KVCacheConfigurationError.noCompatibleLayers( + strategy: configuration.strategy.identifier) + } + case .requireAllLayers: + guard compatibleCount > 0 else { + throw KVCacheConfigurationError.noCompatibleLayers( + strategy: configuration.strategy.identifier) + } + guard incompatibleCount == 0 else { + throw KVCacheConfigurationError.incompatibleLayers( + strategy: configuration.strategy.identifier, + count: incompatibleCount) + } + } +} + +/// Describe the effective strategy of every cache entry without mutating it. +public func kvCacheRuntimeReport( + cache: [KVCache], + configuration: KVCacheConfiguration +) -> KVCacheRuntimeReport { + let leaves = KVCacheTree.leaves(in: cache) + let protectedPaths = protectedPaths(for: leaves, configuration: configuration) + return KVCacheRuntimeReport( + requestedConfiguration: configuration, + layers: leaves.map { + $0.runtimeLayer( + configuration: configuration, + protectedPaths: protectedPaths) + }) +} + +private func protectedPaths( + for leaves: [KVCacheLeaf], + configuration: KVCacheConfiguration +) -> Set<[Int]> { + guard case .turboQuant(let turbo) = configuration.strategy.storage else { return [] } + return KVCacheTree.turboQuantProtectedPaths( + in: leaves, + keyBits: turbo.keyPrecision.bitWidth, + valueBits: turbo.valuePrecision.bitWidth) +} + +extension KVCacheLeaf { + fileprivate func runtimeLayer( + configuration: KVCacheConfiguration, + protectedPaths: Set<[Int]> + ) -> KVCacheRuntimeReport.Layer { + let requested = configuration.strategy.identifier + switch kind { + case .recurrent: + return .init( + path: path, state: .notApplicable, resolvedStrategy: nil, + reason: .nonAttentionState) + case .turboQuant(let turbo): + let matches = requested == .turboQuant + return .init( + path: path, + state: matches ? (turbo.isCompressed ? .active : .pending) : .skipped, + resolvedStrategy: .turboQuant, + reason: matches + ? (turbo.isCompressed ? nil : .awaitingCompressionStart) + : .differentStrategy) + case .affine: + let isBoundaryProtection = + requested == .turboQuant && protectedPaths.contains(path) + let matches = requested == .affine || isBoundaryProtection + return .init( + path: path, + state: matches ? .active : .skipped, + resolvedStrategy: .affine, + reason: isBoundaryProtection + ? .boundaryProtection : (matches ? nil : .differentStrategy)) + case .rotating: + return .init( + path: path, + state: requested == .fullPrecision ? .active : .skipped, + resolvedStrategy: .fullPrecision, + reason: requested == .fullPrecision ? nil : .slidingWindow) + case .simple where requested == .fullPrecision: + return .init( + path: path, state: .active, resolvedStrategy: .fullPrecision, reason: nil) + case .simple: + guard supports(configuration, protectedPaths: protectedPaths) == true else { + return .init( + path: path, + state: .skipped, + resolvedStrategy: .fullPrecision, + reason: .unsupportedShape) + } + let compressionPending = cache.offset <= configuration.strategy.compressionStart + return .init( + path: path, + state: compressionPending ? .pending : .skipped, + resolvedStrategy: .fullPrecision, + reason: compressionPending ? .awaitingCompressionStart : .unsupportedShape) + case .unsupported: + return .init( + path: path, state: .skipped, resolvedStrategy: nil, reason: .unsupportedShape) + } + } +} diff --git a/Libraries/MLXLMCommon/KVCacheTree.swift b/Libraries/MLXLMCommon/KVCacheTree.swift new file mode 100644 index 000000000..2395de6ad --- /dev/null +++ b/Libraries/MLXLMCommon/KVCacheTree.swift @@ -0,0 +1,132 @@ +// Copyright © 2026 Apple Inc. + +/// A classified leaf in a potentially nested cache topology. +struct KVCacheLeaf { + enum Kind { + case recurrent + case rotating(RotatingKVCache) + case simple(KVCacheSimple) + case affine(QuantizedKVCache) + case turboQuant(TurboQuantKVCache) + case unsupported + } + + let path: [Int] + let cache: KVCache + + var kind: Kind { + switch cache { + case is MambaCache, is ArraysCache: .recurrent + case let cache as RotatingKVCache: .rotating(cache) + case let cache as TurboQuantKVCache: .turboQuant(cache) + case let cache as QuantizedKVCache: .affine(cache) + case let cache as KVCacheSimple: .simple(cache) + default: .unsupported + } + } + + var isAttentionCache: Bool { + if case .recurrent = kind { return false } + return true + } + + var isCompressed: Bool { + switch kind { + case .affine, .turboQuant: true + default: false + } + } + + /// Whether this leaf can satisfy a compression strategy. `nil` identifies + /// state that is not an attention cache and does not participate. + func supports( + _ configuration: KVCacheConfiguration, + protectedPaths: Set<[Int]> + ) -> Bool? { + let strategy = configuration.strategy.identifier + switch kind { + case .recurrent: + return nil + case .rotating, .unsupported: + return strategy == .fullPrecision + case .simple(let simple): + let groupSize: Int + switch configuration.strategy.storage { + case .fullPrecision: + return true + case .affine(let affine): + groupSize = affine.groupSize + case .turboQuant(let turbo): + guard protectedPaths.contains(path) || turbo.keyPrecision == .affineEightBit + else { return true } + groupSize = 64 + } + + let state = simple.innerState() + guard state.count >= 2 else { return true } + return resolvedKVQuantizationGroupSize( + requested: groupSize, + keyHeadDim: state[0].dim(3), + valueHeadDim: state[1].dim(3)) != nil + case .affine: + return strategy == .affine + || (strategy == .turboQuant && protectedPaths.contains(path)) + case .turboQuant: + return strategy == .turboQuant + } + } +} + +/// The sole traversal and mutation boundary for nested cache topologies. +enum KVCacheTree { + static func leaves(in cache: [KVCache]) -> [KVCacheLeaf] { + var leaves = [KVCacheLeaf]() + for (index, entry) in cache.enumerated() { + appendLeaves(from: entry, path: [index], to: &leaves) + } + return leaves + } + + static func rewrite( + _ cache: inout [KVCache], + using transform: (KVCacheLeaf) -> KVCache + ) { + for index in cache.indices { + if let list = cache[index] as? CacheList { + list.rewriteLeaves(path: [index], using: transform) + } else { + let leaf = KVCacheLeaf(path: [index], cache: cache[index]) + cache[index] = transform(leaf) + } + } + } + + static func turboQuantProtectedPaths( + in leaves: [KVCacheLeaf], + keyBits: Int, + valueBits: Int + ) -> Set<[Int]> { + let fragile = (keyBits > 0 && keyBits < 8) || valueBits <= 2 + guard fragile else { return [] } + + let attentionPaths = leaves.filter(\.isAttentionCache).map(\.path) + let boundaryCount = min(2, attentionPaths.count / 2) + guard boundaryCount > 0 else { return [] } + return Set(attentionPaths.prefix(boundaryCount)) + .union(attentionPaths.suffix(boundaryCount)) + } + + private static func appendLeaves( + from cache: KVCache, + path: [Int], + to leaves: inout [KVCacheLeaf] + ) { + if let list = cache as? CacheList { + for (index, child) in list.children.enumerated() { + appendLeaves(from: child, path: path + [index], to: &leaves) + } + } else { + leaves.append(KVCacheLeaf(path: path, cache: cache)) + } + } +} diff --git a/Libraries/MLXLMCommon/LanguageModel.swift b/Libraries/MLXLMCommon/LanguageModel.swift index 147b07cca..c69c56e83 100644 --- a/Libraries/MLXLMCommon/LanguageModel.swift +++ b/Libraries/MLXLMCommon/LanguageModel.swift @@ -290,10 +290,10 @@ extension LanguageModel where Self: KVCacheDimensionProvider { // The number of heads per layer (kvHeads[i]) is not used for cache creation let numLayers = kvHeads.count - // Follow Python logic: use RotatingKVCache if maxKVSize is provided - if let maxKVSize = parameters?.maxKVSize { + // Follow Python logic: use RotatingKVCache if a capacity is provided. + if let capacity = parameters?.effectiveKVCacheCapacity { return (0 ..< numLayers).map { _ in - RotatingKVCache(maxSize: maxKVSize, keep: 4) + capacity.makeRotatingCache() } } else { return (0 ..< numLayers).map { _ in KVCacheSimple() } diff --git a/Libraries/MLXLMCommon/MTPSpeculativeTokenIterator.swift b/Libraries/MLXLMCommon/MTPSpeculativeTokenIterator.swift index b0a75fdae..701e29a99 100644 --- a/Libraries/MLXLMCommon/MTPSpeculativeTokenIterator.swift +++ b/Libraries/MLXLMCommon/MTPSpeculativeTokenIterator.swift @@ -38,8 +38,12 @@ public struct MTPSpeculativeTokenIterator: TokenIteratorProtocol { let drafter: any MTPDrafterModel var mainState: LMOutput.State? - var mainCache: [KVCache] - let quantizeKVCache: (inout [KVCache]) -> Void + let mainCacheStorage: KVCacheStorage + var mainCache: [KVCache] { + get { mainCacheStorage.cache } + set { mainCacheStorage.replace(with: newValue) } + } + var kvCachePlan: KVCachePlan { mainCacheStorage.plan } var processor: LogitProcessor? let sampler: LogitSampler @@ -107,15 +111,19 @@ public struct MTPSpeculativeTokenIterator: TokenIteratorProtocol { blockSize >= 2, "MTPSpeculativeTokenIterator requires blockSize >= 2 (1 bonus + K-1 drafted)") + let kvCachePlan = try parameters.kvCachePlan() + let mainCache = try kvCachePlan.validated( + mainCache ?? mainModel.newCache(parameters: parameters)) + guard canTrimPromptCache(mainCache) else { + throw KVCacheError( + message: "MTP speculative decoding requires a trimmable main KV cache.") + } + self.y = input.text self.mainModel = mainModel self.drafter = drafter - self.mainCache = mainCache ?? mainModel.newCache(parameters: parameters) - guard canTrimPromptCache(self.mainCache) else { - throw KVCacheError( - message: "MTP speculative decoding requires a trimmable main KV cache.") - } + self.mainCacheStorage = KVCacheStorage(mainCache, plan: kvCachePlan) self.sampler = parameters.sampler() self.processor = parameters.processor() @@ -123,15 +131,6 @@ public struct MTPSpeculativeTokenIterator: TokenIteratorProtocol { self.maxTokens = parameters.maxTokens self.blockSize = blockSize - self.quantizeKVCache = { cache in - maybeQuantizeKVCache( - cache: &cache, - kvBits: parameters.kvBits, - kvGroupSize: parameters.kvGroupSize, - quantizedKVStart: parameters.quantizedKVStart - ) - } - let prefillStart = Date.timeIntervalSinceReferenceDate try prepare(input: input, windowSize: parameters.prefillStepSize) self.promptPrefillTime = Date.timeIntervalSinceReferenceDate - prefillStart @@ -208,6 +207,8 @@ public struct MTPSpeculativeTokenIterator: TokenIteratorProtocol { pendingTokens.append(token.item(Int.self)) } } + + try kvCachePlan.applyAndValidate(to: mainCacheStorage) } /// Single round: draft `blockSize - 1` tokens, verify with main, accept @@ -363,7 +364,7 @@ public struct MTPSpeculativeTokenIterator: TokenIteratorProtocol { // Dynamic cache quantization may convert `.regular` K/V to `.quantized`, // at which point the target's emit-hook returns sharedKV: nil and the // next round transitions to passthrough. - quantizeKVCache(&mainCache) + kvCachePlan.apply(to: mainCacheStorage) y = .init(tokens: finalToken) } @@ -395,7 +396,7 @@ public struct MTPSpeculativeTokenIterator: TokenIteratorProtocol { eval(token) let tokenInt = token.item(Int.self) y = .init(tokens: token) - quantizeKVCache(&mainCache) + kvCachePlan.apply(to: mainCacheStorage) return tokenInt } diff --git a/Libraries/MLXLMCommon/TurboQuantKVCache.swift b/Libraries/MLXLMCommon/TurboQuantKVCache.swift index 9f8ff9af7..b4d5ea7ca 100644 --- a/Libraries/MLXLMCommon/TurboQuantKVCache.swift +++ b/Libraries/MLXLMCommon/TurboQuantKVCache.swift @@ -1647,7 +1647,7 @@ private enum RotatingSkipNotice { } /// Log, at most once per set of `RotatingKVCache` instances, that a -/// requested TurboQuant kvScheme is leaving sliding-window layers at fp16. +/// requested TurboQuant strategy is leaving sliding-window layers at fp16. /// /// TurboQuant only compresses `KVCacheSimple` layers (see /// `maybeTurboQuantizeKVCache` below); `RotatingKVCache` (Gemma-style @@ -1656,9 +1656,12 @@ private enum RotatingSkipNotice { /// sequential-append compression path, so they are intentionally left /// untouched. The notice tells a user which layers kept fp16 rotating /// caches so a partially engaged scheme is visible. -private func logRotatingKVCacheSkipOnce(cache: [KVCache]) { - let rotatingIndices = cache.indices.filter { cache[$0] is RotatingKVCache } - guard !rotatingIndices.isEmpty else { return } +private func logRotatingKVCacheSkipOnce(leaves: [KVCacheLeaf]) { + let paths = leaves.compactMap { leaf -> [Int]? in + guard case .rotating = leaf.kind else { return nil } + return leaf.path + } + guard !paths.isEmpty else { return } RotatingSkipNotice.lock.lock() defer { RotatingSkipNotice.lock.unlock() } @@ -1666,9 +1669,10 @@ private func logRotatingKVCacheSkipOnce(cache: [KVCache]) { guard !RotatingSkipNotice.logged else { return } RotatingSkipNotice.logged = true - let indexList = rotatingIndices.map(String.init).joined(separator: ", ") + let indexList = paths.map { $0.map(String.init).joined(separator: ".") } + .joined(separator: ", ") print( - "[TurboQuant] kvScheme requested KV compression, but layer(s) at index \(indexList) " + "[TurboQuant] KV compression was requested, but layer(s) at index \(indexList) " + "use RotatingKVCache (sliding-window) and will stay fp16. TurboQuant only " + "compresses non-rotating (global) KV cache layers." ) @@ -1682,54 +1686,45 @@ private func logRotatingKVCacheSkipOnce(cache: [KVCache]) { /// their window size and their rotating storage layout does not fit the /// sequential-append compression path. Called from `maybeQuantizeKVCache` /// when `kvScheme` names a TurboQuant scheme. +@discardableResult public func maybeTurboQuantizeKVCache( cache: inout [KVCache], keyBits: Int, valueBits: Int, quantizedKVStart: Int -) { - logRotatingKVCacheSkipOnce(cache: cache) - - guard - cache.contains(where: { $0 is KVCacheSimple && $0.offset > quantizedKVStart }) - else { return } +) -> Bool { + let leaves = KVCacheTree.leaves(in: cache) + logRotatingKVCacheSkipOnce(leaves: leaves) // Boundary layer protection: the first and last attention layers are // disproportionately sensitive to KV quantization, keeping 2 on each // end at FP16 recovers 37-91% of the quality gap at minimal // compression cost. Non-attention layers (Mamba/wrapper caches) are // excluded from the rank count so hybrids protect the right layers. - let kvLayerIndices = cache.indices.filter { cache[$0] is KVCacheSimple } - // Only engage when the scheme is actually fragile: turbo-quantized keys - // (keyBits 2-4) or 2-bit values. Raw/affine-8 keys with >=3-bit values - // are near-lossless everywhere and protection just costs compression. - let fragile = (keyBits > 0 && keyBits < 8) || valueBits <= 2 - let boundaryLayers = fragile ? min(2, kvLayerIndices.count / 2) : 0 - let protected = Set( - kvLayerIndices.prefix(boundaryLayers) + kvLayerIndices.suffix(boundaryLayers)) - - for i in 0 ..< cache.count { - guard let simple = cache[i] as? KVCacheSimple, cache[i].offset > quantizedKVStart, - !(cache[i] is TurboQuantKVCache) - else { continue } - - let state = cache[i].innerState() + let protectedPaths = KVCacheTree.turboQuantProtectedPaths( + in: leaves, keyBits: keyBits, valueBits: valueBits) + + var awaitsCompressionStart = false + KVCacheTree.rewrite(&cache) { leaf in + guard case .simple(let simple) = leaf.kind else { return leaf.cache } + guard simple.offset > quantizedKVStart else { + awaitsCompressionStart = true + return simple + } + + let state = simple.innerState() let headDims: (key: Int, value: Int)? = state.count >= 2 ? (state[0].dim(3), state[1].dim(3)) : nil - if protected.contains(i) { + if protectedPaths.contains(leaf.path) { // Boundary layers use 8-bit affine instead of the turbo scheme: // near-lossless protection for the quantization-sensitive first - // and last layers at a quarter of the fp16 cost. Skip the - // conversion (leave fp16) rather than crash when neither head - // dimension divides one of the supported affine group sizes. - guard - let headDims, - resolvedKVQuantizationGroupSize( - requested: 64, keyHeadDim: headDims.key, valueHeadDim: headDims.value) != nil - else { continue } - cache[i] = simple.toQuantized(groupSize: 64, bits: 8) - continue + // and last layers at a quarter of the fp16 cost. An unsupported + // realized shape remains in full precision. + guard let quantized = try? simple.toQuantized(groupSize: 64, bits: 8) else { + return simple + } + return quantized } // Affine-K mode (keyBits == 8) quantizes keys in groups, resolve the @@ -1745,7 +1740,7 @@ public func maybeTurboQuantizeKVCache( let headDims, let resolved = resolvedKVQuantizationGroupSize( requested: 64, keyHeadDim: headDims.key, valueHeadDim: headDims.value) - else { continue } + else { return simple } resolvedKeyGroupSize = resolved } @@ -1754,12 +1749,13 @@ public func maybeTurboQuantizeKVCache( keyGroupSize: resolvedKeyGroupSize) // Transfer existing KV data, trimmed to the live offset (the simple // cache over-allocates in steps). - let offset = cache[i].offset + let offset = simple.offset if state.count >= 2, offset > 0 { let keys = state[0][.ellipsis, .. [KVCache] { - var cache = model.newCache(parameters: parameters) + let kvCachePlan = try parameters.kvCachePlan() + var cache = try kvCachePlan.validated(model.newCache(parameters: parameters)) switch try model.prepare( input, cache: cache, state: nil, windowSize: parameters.prefillStepSize) @@ -104,23 +105,13 @@ public enum WiredMemoryUtils { cache: cache.isEmpty ? nil : cache, state: nil ) - maybeQuantizeKVCache( - cache: &cache, - kvBits: parameters.kvBits, - kvGroupSize: parameters.kvGroupSize, - quantizedKVStart: parameters.quantizedKVStart - ) eval(result.logits) case .logits(let result): - maybeQuantizeKVCache( - cache: &cache, - kvBits: parameters.kvBits, - kvGroupSize: parameters.kvGroupSize, - quantizedKVStart: parameters.quantizedKVStart - ) eval(result.logits) } + try kvCachePlan.applyAndValidate(to: &cache) + return cache } diff --git a/Libraries/MLXVLM/Models/Mistral3.swift b/Libraries/MLXVLM/Models/Mistral3.swift index b517ae652..8fedd1138 100644 --- a/Libraries/MLXVLM/Models/Mistral3.swift +++ b/Libraries/MLXVLM/Models/Mistral3.swift @@ -579,8 +579,8 @@ private enum Language { return layerTypes.map { layerType in if layerType == "sliding_attention", let slidingWindow = config.slidingWindow { return RotatingKVCache(maxSize: slidingWindow) - } else if let maxKVSize = parameters?.maxKVSize { - return RotatingKVCache(maxSize: maxKVSize, keep: 4) + } else if let capacity = parameters?.effectiveKVCacheCapacity { + return capacity.makeRotatingCache() } else { return KVCacheSimple() } diff --git a/Libraries/MLXVLM/Models/Pixtral.swift b/Libraries/MLXVLM/Models/Pixtral.swift index 9770279fa..3f65ccc77 100644 --- a/Libraries/MLXVLM/Models/Pixtral.swift +++ b/Libraries/MLXVLM/Models/Pixtral.swift @@ -714,8 +714,8 @@ private enum PixtralLanguage { func newCache(parameters: GenerateParameters?) -> [KVCache] { (0 ..< config.numHiddenLayers).map { _ in - if let maxKVSize = parameters?.maxKVSize { - return RotatingKVCache(maxSize: maxKVSize, keep: 4) + if let capacity = parameters?.effectiveKVCacheCapacity { + return capacity.makeRotatingCache() } else { return KVCacheSimple() } diff --git a/Libraries/MLXVLM/Models/Qwen35.swift b/Libraries/MLXVLM/Models/Qwen35.swift index 79660e56f..87a7bdd2e 100644 --- a/Libraries/MLXVLM/Models/Qwen35.swift +++ b/Libraries/MLXVLM/Models/Qwen35.swift @@ -872,13 +872,13 @@ enum Qwen35Language { return LMOutput(logits: out, state: state) } - func makeCache(maxKVSize: Int?) -> [KVCache] { + func makeCache(capacity: KVCacheConfiguration.Capacity?) -> [KVCache] { model.layers.map { layer in if layer.isLinear { return MambaCache() } - if let maxKVSize { - return RotatingKVCache(maxSize: maxKVSize, keep: 4) + if let capacity { + return capacity.makeRotatingCache() } return KVCacheSimple() } @@ -908,7 +908,7 @@ public class Qwen35: Module, VLMModel { } public func newCache(parameters: GenerateParameters?) -> [KVCache] { - languageModel.makeCache(maxKVSize: parameters?.maxKVSize) + languageModel.makeCache(capacity: parameters?.effectiveKVCacheCapacity) } private func mergeInputIdsWithImageFeatures( diff --git a/Tests/MLXLMTests/ChatSessionTests.swift b/Tests/MLXLMTests/ChatSessionTests.swift index bc03afda3..01adafd5a 100644 --- a/Tests/MLXLMTests/ChatSessionTests.swift +++ b/Tests/MLXLMTests/ChatSessionTests.swift @@ -514,6 +514,55 @@ public class ChatSessionTests: XCTestCase { } } + func testChangingKVCacheConfigurationRequiresClearingRealizedCache() async throws { + let initialConfiguration = KVCacheConfiguration( + strategy: .turboQuant(.qualityFirst)) + let session = ChatSession( + model(), + cache: [KVCacheSimple()], + generateParameters: GenerateParameters(kvCache: initialConfiguration)) + + session.generateParameters.kvCache = KVCacheConfiguration(strategy: .fullPrecision) + + do { + _ = try await session.kvCacheRuntimeReport() + XCTFail("Expected a realized cache to reject configuration changes") + } catch ChatSessionError.kvCacheConfigurationChanged( + let previous, let requested) + { + XCTAssertEqual(previous, initialConfiguration) + XCTAssertEqual(requested, KVCacheConfiguration(strategy: .fullPrecision)) + } + } + + func testSessionRetainsCacheReplacementTriggeredDuringDecode() async throws { + let affine = try AffineKVCacheConfiguration( + bits: 4, groupSize: 64, compressionStart: 8) + let session = ChatSession( + model(), + generateParameters: GenerateParameters( + maxTokens: 3, + kvCache: KVCacheConfiguration( + strategy: .affine(affine), compatibility: .allowPartial))) + + _ = try await session.respond(to: "hello") + let observedFirstOffset = try await session.withCache { cache in + let cache = try XCTUnwrap(cache) + let quantized = try XCTUnwrap( + cache.first { $0 is QuantizedKVCache } as? QuantizedKVCache) + return quantized.offset + } + let firstOffset = try XCTUnwrap(observedFirstOffset) + + _ = try await session.respond(to: "hello again") + try await session.withCache { cache in + let cache = try XCTUnwrap(cache) + let quantized = try XCTUnwrap( + cache.first { $0 is QuantizedKVCache } as? QuantizedKVCache) + XCTAssertGreaterThan(quantized.offset, firstOffset) + } + } + /// something that looks like a view model @MainActor class ChatModel { let session: ChatSession diff --git a/Tests/MLXLMTests/KVCacheConfigurationTests.swift b/Tests/MLXLMTests/KVCacheConfigurationTests.swift new file mode 100644 index 000000000..3f022c18d --- /dev/null +++ b/Tests/MLXLMTests/KVCacheConfigurationTests.swift @@ -0,0 +1,383 @@ +// Copyright © 2026 Apple Inc. + +import Foundation +import MLX +import Testing + +@testable import MLXLMCommon + +@Suite("KV-cache configuration") +struct KVCacheConfigurationTests { + @Test func typedPresetsExposeStableAlgorithmIdentity() { + let quality = KVCacheConfiguration( + strategy: .turboQuant(.qualityFirst)) + let balanced = KVCacheConfiguration( + strategy: .turboQuant(.balanced)) + let affine = KVCacheConfiguration(strategy: .affine(.fourBit)) + + #expect(quality.strategy.identifier == .turboQuant) + #expect(balanced.strategy.identifier == .turboQuant) + #expect(affine.strategy.identifier == .affine) + #expect(TurboQuantKVCacheConfiguration.qualityFirst.keyPrecision == .fp16) + #expect(TurboQuantKVCacheConfiguration.qualityFirst.valuePrecision == .fourBit) + #expect(TurboQuantKVCacheConfiguration.balanced.keyPrecision == .affineEightBit) + #expect(TurboQuantKVCacheConfiguration.balanced.valuePrecision == .threeBit) + } + + @Test func invalidTypedValuesAreRejectedAtConstruction() { + #expect(throws: KVCacheConfigurationError.invalidCapacity(0)) { + try KVCacheConfiguration.Capacity(maxTokens: 0) + } + #expect(throws: KVCacheConfigurationError.invalidPreservedPrefix(8, capacity: 8)) { + try KVCacheConfiguration.Capacity(maxTokens: 8, preservedPrefixTokens: 8) + } + for bits in [1, 7, 9] { + #expect(throws: KVCacheConfigurationError.invalidAffineBits(bits)) { + try AffineKVCacheConfiguration(bits: bits) + } + } + #expect(throws: KVCacheConfigurationError.invalidGroupSize(0)) { + try AffineKVCacheConfiguration(bits: 4, groupSize: 0) + } + } + + @Test func legacyTurboSchemeResolvesToTypedConfiguration() throws { + let parameters = GenerateParameters( + kvGroupSize: 64, + quantizedKVStart: 128, + kvScheme: "turbo8v3") + + let optionalConfiguration = try parameters.resolvedKVCacheConfiguration() + let resolved = try #require(optionalConfiguration) + #expect(resolved.strategy.identifier == .turboQuant) + guard case .turboQuant(let turbo) = resolved.strategy.storage else { + Issue.record("Expected a TurboQuant strategy") + return + } + #expect(turbo.keyPrecision == .affineEightBit) + #expect(turbo.valuePrecision == .threeBit) + #expect(turbo.compressionStart == 128) + } + + @Test func unknownLegacySchemeIsRejected() { + let parameters = GenerateParameters(kvScheme: "future-unregistered-scheme") + #expect( + throws: KVCacheConfigurationError.unsupportedLegacyScheme( + "future-unregistered-scheme") + ) { + try parameters.resolvedKVCacheConfiguration() + } + } + + @Test func typedAndLegacyConfigurationCannotBeMixed() { + let parameters = GenerateParameters( + kvCache: .init(strategy: .turboQuant(.balanced)), + kvBits: 4) + #expect(throws: KVCacheConfigurationError.conflictingLegacyConfiguration) { + try parameters.resolvedKVCacheConfiguration() + } + } + + @Test func capacityPreservesCallerSelectedPrefix() throws { + let capacity = try KVCacheConfiguration.Capacity( + maxTokens: 4_096, preservedPrefixTokens: 16) + let parameters = GenerateParameters(kvCache: .init(capacity: capacity)) + let resolved = try parameters.resolvedKVCacheConfiguration() + + #expect(parameters.effectiveKVCacheCapacity == capacity) + #expect(resolved?.capacity == capacity) + } + + @Test func typedConfigurationRejectsAllRotatingCacheByDefault() { + let configuration = KVCacheConfiguration( + strategy: .turboQuant(.balanced)) + let cache: [KVCache] = [ + RotatingKVCache(maxSize: 128), + RotatingKVCache(maxSize: 128), + ] + + #expect( + throws: KVCacheConfigurationError.noCompatibleLayers(strategy: .turboQuant) + ) { + try validateKVCacheCompatibility(cache, configuration: configuration) + } + } + + @Test func requireAllLayersRejectsMixedRotatingTopology() { + let configuration = KVCacheConfiguration( + strategy: .affine(.fourBit), + compatibility: .requireAllLayers) + let cache: [KVCache] = [ + KVCacheSimple(), + RotatingKVCache(maxSize: 128), + ] + + #expect( + throws: KVCacheConfigurationError.incompatibleLayers( + strategy: .affine, count: 1) + ) { + try validateKVCacheCompatibility(cache, configuration: configuration) + } + } + + @Test func runtimeReportRecursesThroughCompositeCaches() { + let configuration = KVCacheConfiguration(strategy: .affine(.fourBit)) + let cache: [KVCache] = [ + CacheList( + MambaCache(), + QuantizedKVCache(groupSize: 64, bits: 4)), + RotatingKVCache(maxSize: 128), + ] + + let report = kvCacheRuntimeReport(cache: cache, configuration: configuration) + + #expect(report.compressedLayerCount == 1) + #expect(report.pendingLayerCount == 0) + #expect(report.skippedLayerCount == 1) + #expect(report.layers.count == 3) + #expect(report.layers[0].path == [0, 0]) + #expect(report.layers[0].state == .notApplicable) + #expect(report.layers[1].path == [0, 1]) + #expect(report.layers[1].resolvedStrategy == .affine) + #expect(report.layers[2].reason == .slidingWindow) + } + + @Test func typedTurboQuantDispatchRewritesNestedAttentionCache() throws { + let simple = KVCacheSimple() + simple.offset = 8 + let list = CacheList(MambaCache(), simple) + var cache: [KVCache] = [list] + let configuration = KVCacheConfiguration( + strategy: .turboQuant(.qualityFirst), + compatibility: .requireAllLayers) + + try validateKVCacheCompatibility(cache, configuration: configuration) + _ = try applyKVCacheConfiguration(cache: &cache, configuration: configuration) + + #expect(list[0] is MambaCache) + #expect(list[1] is TurboQuantKVCache) + let report = kvCacheRuntimeReport(cache: cache, configuration: configuration) + #expect(report.pendingLayerCount == 1) + #expect(report.skippedLayerCount == 0) + } + + @Test func realizedUnsupportedShapeFailsClosed() throws { + var cache: [KVCache] = [populatedSimpleCache(headDimension: 48)] + let configuration = KVCacheConfiguration( + strategy: .affine(.fourBit), + compatibility: .requireAtLeastOneLayer) + + #expect( + throws: KVCacheConfigurationError.noCompatibleLayers(strategy: .affine) + ) { + _ = try applyKVCacheConfiguration(cache: &cache, configuration: configuration) + } + #expect(cache[0] is KVCacheSimple) + } + + @Test func mixedRealizedShapesReportConversionAndSkip() throws { + var cache = mixedShapeCache() + let configuration = KVCacheConfiguration( + strategy: .affine(.fourBit), + compatibility: .allowPartial) + + let result = try applyKVCacheConfiguration( + cache: &cache, configuration: configuration) + + #expect(result.convertedLayerCount == 1) + #expect(result.alreadyCompatibleLayerCount == 0) + #expect(result.skipped.count == 1) + #expect(result.skipped[0].path == [1]) + #expect(result.skipped[0].reason == .unsupportedShape) + #expect(cache[0] is QuantizedKVCache) + #expect(cache[1] is KVCacheSimple) + } + + @Test func requireAllLayersRejectsMixedRealizedShapesWithoutMutation() throws { + var cache = mixedShapeCache() + let configuration = KVCacheConfiguration( + strategy: .affine(.fourBit), + compatibility: .requireAllLayers) + + #expect( + throws: KVCacheConfigurationError.incompatibleLayers( + strategy: .affine, count: 1) + ) { + _ = try applyKVCacheConfiguration(cache: &cache, configuration: configuration) + } + #expect(cache[0] is KVCacheSimple) + #expect(cache[1] is KVCacheSimple) + } + + @Test func explicitFullPrecisionRejectsCompressedCache() { + let configuration = KVCacheConfiguration(strategy: .fullPrecision) + let cache: [KVCache] = [QuantizedKVCache(groupSize: 64, bits: 4)] + + #expect( + throws: KVCacheConfigurationError.incompatibleLayers( + strategy: .fullPrecision, count: 1) + ) { + try validateKVCacheCompatibility(cache, configuration: configuration) + } + } + + @Test func explicitFullPrecisionReportsSimpleCachesAsCompatible() throws { + var cache: [KVCache] = [KVCacheSimple(), KVCacheSimple()] + + let result = try applyKVCacheConfiguration( + cache: &cache, + configuration: KVCacheConfiguration(strategy: .fullPrecision)) + + #expect(result.convertedLayerCount == 0) + #expect(result.alreadyCompatibleLayerCount == 2) + #expect(result.pendingLayerCount == 0) + #expect(result.skipped.isEmpty) + } + + @Test func configuredCacheMustRealizeRequestedCapacity() throws { + let configuration = KVCacheConfiguration( + capacity: try .init(maxTokens: 128, preservedPrefixTokens: 4)) + let suppliedCapacity = try KVCacheConfiguration.Capacity( + maxTokens: 64, preservedPrefixTokens: 4) + let cache: [KVCache] = [suppliedCapacity.makeRotatingCache()] + + #expect( + throws: KVCacheConfigurationError.incompatibleCapacity(expected: 128, count: 1) + ) { + try validateKVCacheCompatibility(cache, configuration: configuration) + } + } + + @Test func modelNativeSlidingWindowDoesNotConflictWithRequestedCapacity() throws { + let configuration = KVCacheConfiguration( + capacity: try .init(maxTokens: 128, preservedPrefixTokens: 4)) + let nativeSlidingCache = RotatingKVCache(maxSize: 512, keep: 0) + let configuredCache = try KVCacheConfiguration.Capacity( + maxTokens: 128, preservedPrefixTokens: 4 + ).makeRotatingCache() + + try validateKVCacheCompatibility( + [nativeSlidingCache, configuredCache], configuration: configuration) + + #expect(nativeSlidingCache.preservedPrefixTokens == 0) + #expect(configuredCache.preservedPrefixTokens == 4) + } + + @Test func requestedCapacityOriginSurvivesPromptCacheRoundTrip() throws { + let requested = try KVCacheConfiguration.Capacity( + maxTokens: 64, preservedPrefixTokens: 4) + let cache = requested.makeRotatingCache() + _ = cache.update( + keys: MLXArray.ones([1, 1, 8, 16]), + values: MLXArray.ones([1, 1, 8, 16])) + + let url = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString) + .appendingPathExtension("safetensors") + defer { try? FileManager.default.removeItem(at: url) } + + try savePromptCache(url: url, cache: [cache]) + let (loaded, _) = try loadPromptCache(url: url) + let restored = try #require(loaded.first as? RotatingKVCache) + + #expect(restored.capacityOrigin == .requested) + try validateKVCacheCompatibility( + loaded, configuration: KVCacheConfiguration(capacity: requested)) + + let incompatible = KVCacheConfiguration( + capacity: try .init(maxTokens: 128, preservedPrefixTokens: 4)) + #expect( + throws: KVCacheConfigurationError.incompatibleCapacity(expected: 128, count: 1) + ) { + try validateKVCacheCompatibility(loaded, configuration: incompatible) + } + } + + @Test func legacyRotatingMetadataDefaultsToModelNativeCapacity() { + let cache = RotatingKVCache(maxSize: 64) + cache.capacityOrigin = .requested + + cache.metaState = ["0", "64", "256", "0", "0"] + + #expect(cache.capacityOrigin == .modelNative) + } + + @Test func sharedRealizationPublishesConversionAndBecomesTerminal() throws { + let affine = try AffineKVCacheConfiguration( + bits: 4, groupSize: 64, compressionStart: 1) + let plan = KVCachePlan( + configuration: KVCacheConfiguration( + strategy: .affine(affine), compatibility: .allowPartial)) + let simple = populatedSimpleCache(headDimension: 64) + simple.offset = 1 + let storage = KVCacheStorage([simple], plan: plan) + let observer = storage + + _ = try plan.applyAndValidate(to: storage) + #expect(!storage.isApplicationTerminal) + + simple.offset = 2 + plan.apply(to: storage) + + #expect(storage.isApplicationTerminal) + #expect(storage.cache[0] is QuantizedKVCache) + #expect(observer.cache[0] is QuantizedKVCache) + + plan.apply(to: storage) + #expect(storage.isApplicationTerminal) + #expect(storage.cache[0] is QuantizedKVCache) + } + + @Test func turboRuntimeReportMarksDifferentRequestedStrategy() { + let configuration = KVCacheConfiguration(strategy: .affine(.fourBit)) + let cache: [KVCache] = [ + TurboQuantKVCache(bits: 4, keyBits: 0, valueBits: 4) + ] + + let layer = kvCacheRuntimeReport(cache: cache, configuration: configuration).layers[0] + #expect(layer.state == .skipped) + #expect(layer.resolvedStrategy == .turboQuant) + #expect(layer.reason == .differentStrategy) + } + + @Test func turboBoundaryProtectionUsesStableAttentionPaths() throws { + let turbo = try TurboQuantKVCacheConfiguration( + keyPrecision: .twoBit, + valuePrecision: .fourBit) + let configuration = KVCacheConfiguration( + strategy: .turboQuant(turbo), + compatibility: .allowPartial) + var cache: [KVCache] = (0 ..< 6).map { + [2, 5].contains($0) ? populatedSimpleCache(headDimension: 64) : KVCacheSimple() + } + + _ = try applyKVCacheConfiguration(cache: &cache, configuration: configuration) + #expect(cache[2] is TurboQuantKVCache) + #expect(cache[5] is QuantizedKVCache) + + for index in [0, 1, 3, 4] { + cache[index].state = populatedSimpleCache(headDimension: 64).state + } + _ = try applyKVCacheConfiguration(cache: &cache, configuration: configuration) + + for index in [0, 1, 4, 5] { #expect(cache[index] is QuantizedKVCache) } + for index in [2, 3] { #expect(cache[index] is TurboQuantKVCache) } + + let report = kvCacheRuntimeReport(cache: cache, configuration: configuration) + #expect(report.layers[0].reason == .boundaryProtection) + #expect(report.layers[3].reason == .awaitingCompressionStart) + } + + private func mixedShapeCache() -> [KVCache] { + [populatedSimpleCache(headDimension: 64), populatedSimpleCache(headDimension: 48)] + } + + private func populatedSimpleCache(headDimension: Int) -> KVCacheSimple { + let cache = KVCacheSimple() + cache.state = [ + MLXArray.zeros([1, 1, 1, headDimension]), + MLXArray.zeros([1, 1, 1, headDimension]), + ] + return cache + } +} diff --git a/Tests/MLXLMTests/KVCachePlanBenchmark.swift b/Tests/MLXLMTests/KVCachePlanBenchmark.swift new file mode 100644 index 000000000..aa98e033f --- /dev/null +++ b/Tests/MLXLMTests/KVCachePlanBenchmark.swift @@ -0,0 +1,62 @@ +// Copyright © 2026 Apple Inc. + +import Foundation +import Testing + +@testable import MLXLMCommon + +/// Microbenchmark for dynamic KV-cache application on the decode hot path. +/// +/// The traversal measurement intentionally runs only the minimal conversion +/// walk. The previous implementation also collected before/after leaves, +/// built a runtime report, counted outcomes, and allocated a result, so this +/// is a conservative lower bound for the work removed by terminal plans. +@Suite(.serialized) +struct KVCachePlanBenchmark { + private func nanosecondsPerCall(_ elapsed: Duration, iterations: Int) -> Double { + let components = elapsed.components + let nanoseconds = + Double(components.seconds) * 1_000_000_000 + + Double(components.attoseconds) / 1_000_000_000 + return nanoseconds / Double(iterations) + } + + @Test func completedApplicationCost() { + let configuration = KVCacheConfiguration( + strategy: .affine(.fourBit), compatibility: .allowPartial) + let plan = KVCachePlan(configuration: configuration) + let caches: [KVCache] = (0 ..< 32).map { _ in + QuantizedKVCache(groupSize: 64, bits: 4) + } + let storage = KVCacheStorage(caches, plan: plan) + plan.apply(to: storage) + #expect(storage.isApplicationTerminal) + + let clock = ContinuousClock() + let warmup = 1_000 + let iterations = 100_000 + + for _ in 0 ..< warmup { plan.apply(to: storage) } + var start = clock.now + for _ in 0 ..< iterations { plan.apply(to: storage) } + let terminal = nanosecondsPerCall(clock.now - start, iterations: iterations) + + var traversed = caches + for _ in 0 ..< warmup { + _ = applyKVCacheConfigurationFast( + cache: &traversed, configuration: configuration) + } + start = clock.now + for _ in 0 ..< iterations { + _ = applyKVCacheConfigurationFast( + cache: &traversed, configuration: configuration) + } + let traversal = nanosecondsPerCall(clock.now - start, iterations: iterations) + + print( + String( + format: + "[KVCACHEBENCH] terminal %.1f ns/call | traversal floor %.1f ns/call | %.1fx faster", + terminal, traversal, traversal / terminal)) + } +} diff --git a/Tests/MLXLMTests/KVCacheTests.swift b/Tests/MLXLMTests/KVCacheTests.swift index 3d64f84d4..b3df57f57 100644 --- a/Tests/MLXLMTests/KVCacheTests.swift +++ b/Tests/MLXLMTests/KVCacheTests.swift @@ -21,6 +21,42 @@ private func tempURL() -> URL { .appendingPathExtension("safetensors") } +private struct SerializedCacheFixture { + let className: String + let arrayCount: Int + let metadata: [String] + var arrayShape = [1, 1, 1, 1] +} + +private func writePromptCacheFixture(_ fixture: SerializedCacheFixture) throws -> URL { + let url = tempURL() + let arrays: [String: MLXArray] = Dictionary( + uniqueKeysWithValues: (0 ..< fixture.arrayCount).map { index in + ( + "0.\(index)", + MLXArray.zeros(fixture.arrayShape, dtype: .float32, stream: .cpu) + ) + }) + var metadata: [String: String] = Dictionary( + uniqueKeysWithValues: fixture.metadata.enumerated().map { + ("0.0.\($0.offset)", $0.element) + }) + metadata["2.0"] = fixture.className + try save(arrays: arrays, metadata: metadata, url: url, stream: .cpu) + return url +} + +private func expectPromptCacheLoadToFail(_ fixtures: [SerializedCacheFixture]) throws { + for fixture in fixtures { + let url = try writePromptCacheFixture(fixture) + defer { try? FileManager.default.removeItem(at: url) } + + #expect(throws: KVCacheError.self) { + _ = try loadPromptCache(url: url) + } + } +} + /// Assert two arrays of MLXArray are element-wise close private func assertArraysClose(_ lhs: [MLXArray], _ rhs: [MLXArray], label: String = "") { #expect(lhs.count == rhs.count, "state count mismatch \(label)") @@ -134,11 +170,11 @@ func testCacheSerialization(creator: (() -> any KVCache)) async throws { #expect(copied.metaState == cache.metaState) } -@Test func testEmptyKVCacheSimpleToQuantizedPreservesRequestedQuantizationMetadata() { +@Test func testEmptyKVCacheSimpleToQuantizedPreservesRequestedQuantizationMetadata() throws { let cache = KVCacheSimple() cache.offset = 7 - let quantized = cache.toQuantized(groupSize: 128, bits: 4) + let quantized = try cache.toQuantized(groupSize: 128, bits: 4) #expect(quantized.offset == 7) #expect(quantized.groupSize == 128) @@ -146,6 +182,104 @@ func testCacheSerialization(creator: (() -> any KVCache)) async throws { #expect(quantized.metaState == ["256", "7", "128", "4"]) } +@Test func testDirectKVCacheQuantizationFailuresAreRecoverable() throws { + let incompatible = KVCacheSimple() + incompatible.state = [ + MLXArray.zeros([1, 2, 4, 5], dtype: .float32, stream: .cpu), + MLXArray.zeros([1, 2, 4, 5], dtype: .float32, stream: .cpu), + ] + + #expect(throws: KVCacheError.self) { + _ = try incompatible.toQuantized(groupSize: 64, bits: 4) + } + #expect(throws: KVCacheError.self) { + _ = try RotatingKVCache(maxSize: 32).toQuantized() + } +} + +@Test func testPromptCacheRestorationRejectsIncompleteState() throws { + try expectPromptCacheLoadToFail([ + .init(className: "KVCacheSimple", arrayCount: 1, metadata: [""]), + .init( + className: "RotatingKVCache", arrayCount: 1, + metadata: ["0", "32", "256", "0", "0", "modelNative"]), + .init( + className: "QuantizedKVCache", arrayCount: 3, + metadata: ["256", "0", "64", "4"]), + .init(className: "ChunkedKVCache", arrayCount: 1, metadata: ["None", "0"]), + ]) +} + +@Test func testPromptCacheRestorationRejectsInvalidMetadata() throws { + try expectPromptCacheLoadToFail([ + .init(className: "KVCacheSimple", arrayCount: 0, metadata: ["unexpected"]), + .init( + className: "RotatingKVCache", arrayCount: 0, + metadata: ["0", "32", "invalid", "0", "0", "modelNative"]), + .init( + className: "RotatingKVCache", arrayCount: 0, + metadata: ["0", "32", "256", "0", "0", "invalid-origin"]), + .init( + className: "QuantizedKVCache", arrayCount: 0, + metadata: ["256", "invalid", "64", "4"]), + .init(className: "ChunkedKVCache", arrayCount: 0, metadata: ["invalid", "0"]), + ]) +} + +@Test func testPromptCacheRestorationRejectsInvalidStateRank() throws { + try expectPromptCacheLoadToFail([ + .init( + className: "KVCacheSimple", arrayCount: 2, metadata: [""], + arrayShape: [1, 1]) + ]) +} + +@Test func testPromptCacheRestorationAcceptsLegacyRotatingMetadata() throws { + let fixture = SerializedCacheFixture( + className: "RotatingKVCache", + arrayCount: 2, + metadata: ["0", "32", "256", "1", "1"]) + let url = try writePromptCacheFixture(fixture) + defer { try? FileManager.default.removeItem(at: url) } + + let (restored, _) = try loadPromptCache(url: url) + let rotating = try #require(restored.first as? RotatingKVCache) + + #expect(rotating.capacityOrigin == .modelNative) + #expect(rotating.metaState == ["0", "32", "256", "1", "1", "modelNative"]) +} + +@Test func testPromptCacheRoundTripPreservesEmptyCaches() throws { + let populated = KVCacheSimple() + populated.state = [ + MLXArray.ones([1, 1, 1, 32], stream: .cpu), + MLXArray.ones([1, 1, 1, 32], stream: .cpu), + ] + let caches: [KVCache] = [ + KVCacheSimple(), + populated, + RotatingKVCache(maxSize: 32), + QuantizedKVCache(), + ChunkedKVCache(chunkSize: 16), + ] + let url = tempURL() + defer { try? FileManager.default.removeItem(at: url) } + + try savePromptCache(url: url, cache: caches) + let (restored, _) = try loadPromptCache(url: url) + + #expect(restored.count == caches.count) + #expect(restored[0] is KVCacheSimple) + #expect(restored[0].state.isEmpty) + #expect(restored[1].state.count == 2) + #expect(restored[2] is RotatingKVCache) + #expect(restored[2].state.isEmpty) + #expect(restored[3] is QuantizedKVCache) + #expect(restored[3].state.isEmpty) + #expect(restored[4] is ChunkedKVCache) + #expect(restored[4].state.isEmpty) +} + // MARK: - ArraysCache sparse slot round-trip @Test func testArraysCacheSparseSlots() throws { diff --git a/Tests/MLXLMTests/MTPSpeculativeTokenIteratorTests.swift b/Tests/MLXLMTests/MTPSpeculativeTokenIteratorTests.swift index 2afb59384..b2cb3a19d 100644 --- a/Tests/MLXLMTests/MTPSpeculativeTokenIteratorTests.swift +++ b/Tests/MLXLMTests/MTPSpeculativeTokenIteratorTests.swift @@ -183,6 +183,28 @@ private final class CountingKVCache: KVCache { // MARK: - Smallest-unit-of-work smoke test +@Suite("MTP KV-cache configuration") +struct MTPKVCacheConfigurationTests { + @Test func legacyTurboSchemeUsesTypedDispatcher() throws { + let main = MockMainModel(nextLogitTokens: [0, 0, 7]) + let drafter = MockDrafter(draftedTokenValue: 7) + let input = LMInput(tokens: MLXArray([Int32(1), 2, 3])) + let iterator = try MTPSpeculativeTokenIterator( + input: input, + mainModel: main, + drafter: drafter, + parameters: GenerateParameters(kvScheme: "turbo0v4"), + blockSize: 2) + let simple = KVCacheSimple() + simple.offset = 1 + var cache: [KVCache] = [simple] + + iterator.kvCachePlan.apply(to: &cache) + + #expect(cache[0] is TurboQuantKVCache) + } +} + @Test func testMTPSpeculateRoundSmokeWithSynthetics() throws { // Plan: prompt of 3 tokens [1, 2, 3], main model is rigged so that it