Skip to content

Add typed KV cache configuration and runtime reporting #453

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
aleroot wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Add typed KV cache configuration and runtime reporting #453

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
24a4ff1
Add typed KV cache configuration and runtime reporting
aleroot Jul 22, 2026
2419b1c
Harden KV cache recovery paths
aleroot Jul 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 23 additions & 15 deletions Libraries/MLXGuidedGeneration/GuidedGenerationLoop.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand All @@ -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

Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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)
Expand Down
8 changes: 6 additions & 2 deletions Libraries/MLXLLM/Models/FalconH1.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
}
Expand Down
121 changes: 96 additions & 25 deletions Libraries/MLXLMCommon/ChatSession.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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])
}

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
}

Expand All @@ -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
Expand Down Expand Up @@ -740,19 +782,22 @@ 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!

let iterator = try SpeculativeTokenIterator(
input: input,
mainModel: model,
draftModel: draftModel,
mainCache: kvCache,
draftCache: draftCache,
mainCacheStorage: kvCache,
draftCacheStorage: draftCache,
parameters: generateParameters,
numDraftTokens: speculativeDecoding.numDraftTokens
)
Expand Down Expand Up @@ -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()
}
Expand Down Expand Up @@ -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.
Expand All @@ -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)
}
Expand All @@ -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
}
Expand All @@ -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."
}
}
}
Loading
Loading