diff --git a/Sources/MereRunCore/LFM2/LFM2Generator.swift b/Sources/MereRunCore/LFM2/LFM2Generator.swift index 645dfcef..23836030 100644 --- a/Sources/MereRunCore/LFM2/LFM2Generator.swift +++ b/Sources/MereRunCore/LFM2/LFM2Generator.swift @@ -11,8 +11,36 @@ private struct LFM2DecodeResult { let decodeSeconds: Double } +private struct LFM2PrefixKVCacheKey: Hashable { + let modelPath: String + let tokens: [Int] +} + +private struct LFM2PrefixKVCacheEntry { + let caches: [LFM2LayerCache?] + let logits: MLXArray + let priority: RuntimePrefixCacheEntryPriority + var lastAccess: Date +} + public actor LFM2Generator: ChatGenerator { private static let prefillChunkSize = 512 + private static let prefixKVCacheMaxEntries = 4 + + /// Opt-in (MERERUN_LFM2_PREFIX_KV_CACHE=1) in-memory prompt-prefix reuse, + /// mirroring the Qwen-family implementation: forked layer caches (both + /// attention KV and conv states support forking) are stored at prefill + /// chunk boundaries and the longest matching token prefix seeds the next + /// request. Chunk-boundary checkpoints only for now — the Gemma4-style + /// semantic chat-template checkpoints are not yet derived for LFM2. + private static let prefixKVCacheEnabled: Bool = + ProcessInfo.processInfo.environment["MERERUN_LFM2_PREFIX_KV_CACHE"] == "1" + + private var prefixKVCache: [LFM2PrefixKVCacheKey: LFM2PrefixKVCacheEntry] = [:] + private var prefixKVCacheHits = 0 + private var prefixKVCacheMisses = 0 + private var prefixKVCacheStores = 0 + private var prefixKVCacheReusedTokens = 0 private var model: LFM2Model? private var tokenizerAndTemplate: LFM2TokenizerAndTemplate? @@ -168,11 +196,22 @@ public actor LFM2Generator: ChatGenerator { repetitionContextSize: 64 ) - let layerCaches = makeLayerCaches(config: loadedConfig) + var layerCaches = makeLayerCaches(config: loadedConfig) + var prefillStartIndex = 0 + var prefillExistingLogits: MLXArray? + if let seed = prefixKVCacheSeed(modelPath: loadedModelPath, promptTokens: promptTokens) { + layerCaches = seed.caches + prefillStartIndex = seed.tokenCount + prefillExistingLogits = seed.logits + progressHandler?(ChatProgress(stage: .encoding, message: "Reusing \(seed.tokenCount) prompt KV tokens")) + } let prefillOutput = try await chunkedPrefill( model: model, promptTokens: promptTokens, cache: layerCaches, + modelPath: loadedModelPath, + startIndex: prefillStartIndex, + existingLogits: prefillExistingLogits, progressHandler: progressHandler ) let prefillSeconds = Date().timeIntervalSince(prefillStart) @@ -280,14 +319,18 @@ public actor LFM2Generator: ChatGenerator { model: LFM2Model, promptTokens: [Int], cache: [LFM2LayerCache?], + modelPath: String? = nil, + startIndex: Int = 0, + existingLogits: MLXArray? = nil, progressHandler: (@Sendable (ChatProgress) -> Void)? ) async throws -> LFM2PrefillOutput { guard !promptTokens.isEmpty else { throw LFM2Error.generationFailed("Prompt tokenization produced no tokens.") } - var offset = 0 - var lastOutput: LFM2ForwardOutput? + var offset = startIndex + var logits = existingLogits + var hidden: MLXArray? while offset < promptTokens.count { try Task.checkCancellation() let end = min(promptTokens.count, offset + Self.prefillChunkSize) @@ -296,8 +339,18 @@ public actor LFM2Generator: ChatGenerator { let output = model.forwardPrefill(input, cache: cache) MLX.eval(output.logits) MLX.eval(output.hidden) - lastOutput = output + logits = output.logits + hidden = output.hidden offset = end + if let modelPath { + storePrefixKVCache( + modelPath: modelPath, + promptTokens: promptTokens, + tokenCount: end, + cache: cache, + logits: output.logits + ) + } progressHandler?(ChatProgress( stage: .encoding, message: "Prefilled \(offset)/\(promptTokens.count) tokens" @@ -305,10 +358,87 @@ public actor LFM2Generator: ChatGenerator { await Task.yield() } - guard let lastOutput else { + guard let logits else { throw LFM2Error.generationFailed("LFM2 prefill produced no logits.") } - return LFM2PrefillOutput(logits: lastOutput.logits, hidden: lastOutput.hidden) + return LFM2PrefillOutput(logits: logits, hidden: hidden ?? logits) + } + + // MARK: - Prefix KV cache + + public func prefixKVCacheStats() -> PrefixKVCacheStats { + PrefixKVCacheStats( + enabled: Self.prefixKVCacheEnabled, + entries: prefixKVCache.count, + maxEntries: Self.prefixKVCacheMaxEntries, + hits: prefixKVCacheHits, + misses: prefixKVCacheMisses, + storedPrefixes: prefixKVCacheStores, + reusedTokens: prefixKVCacheReusedTokens, + storedTokens: prefixKVCache.keys.reduce(0) { $0 + $1.tokens.count } + ) + } + + private func prefixKVCacheSeed( + modelPath: String?, + promptTokens: [Int] + ) -> (tokenCount: Int, caches: [LFM2LayerCache?], logits: MLXArray)? { + guard Self.prefixKVCacheEnabled, let modelPath else { return nil } + let matchingKey = prefixKVCache.keys + .filter { key in + key.modelPath == modelPath + && key.tokens.count <= promptTokens.count + && promptTokens.starts(with: key.tokens) + } + .max { $0.tokens.count < $1.tokens.count } + + guard let matchingKey, var entry = prefixKVCache[matchingKey] else { + prefixKVCacheMisses += 1 + return nil + } + + entry.lastAccess = Date() + prefixKVCache[matchingKey] = entry + prefixKVCacheHits += 1 + prefixKVCacheReusedTokens += matchingKey.tokens.count + return ( + matchingKey.tokens.count, + entry.caches.map { $0?.fork() }, + entry.logits + ) + } + + private func storePrefixKVCache( + modelPath: String, + promptTokens: [Int], + tokenCount: Int, + cache: [LFM2LayerCache?], + logits: MLXArray + ) { + guard Self.prefixKVCacheEnabled, tokenCount > 0 else { return } + let key = LFM2PrefixKVCacheKey( + modelPath: modelPath, + tokens: Array(promptTokens.prefix(tokenCount)) + ) + prefixKVCache[key] = LFM2PrefixKVCacheEntry( + caches: cache.map { $0?.fork() }, + logits: logits, + priority: .chunk, + lastAccess: Date() + ) + prefixKVCacheStores += 1 + while prefixKVCache.count > Self.prefixKVCacheMaxEntries { + let metadata = prefixKVCache.mapValues { + RuntimePrefixCacheRetentionMetadata( + priority: $0.priority, + lastAccess: $0.lastAccess + ) + } + guard let oldest = RuntimePrefixCacheRetentionPlanner.keyToPrune(entries: metadata) else { + return + } + prefixKVCache.removeValue(forKey: oldest) + } } private func resolveModelRoot( diff --git a/Sources/MereRunCore/ZImageTurbo/Model/TextEncoder/LLMGeneration/KVCache.swift b/Sources/MereRunCore/ZImageTurbo/Model/TextEncoder/LLMGeneration/KVCache.swift index 498d187f..ffd1feec 100644 --- a/Sources/MereRunCore/ZImageTurbo/Model/TextEncoder/LLMGeneration/KVCache.swift +++ b/Sources/MereRunCore/ZImageTurbo/Model/TextEncoder/LLMGeneration/KVCache.swift @@ -117,8 +117,15 @@ public class KVCacheSimple: KVCache { public func fork() -> KVCache { let copy = KVCacheSimple(step: step) - copy.keys = keys - copy.values = values + // `update` writes new tokens with subscript assignment, which rebinds + // the SAME MLXArray wrapper in place (`_updateInternal`). Sharing the + // wrapper objects with a fork means every later write on the parent + // mutates the fork too — prefix-KV snapshots stored mid-request were + // silently corrupted by the request's remaining prefill and decode. + // Fresh wrappers over the current (immutable) arrays isolate the fork; + // the parent's rebinds can no longer reach it. + copy.keys = keys.map { $0.asType($0.dtype) } + copy.values = values.map { $0.asType($0.dtype) } copy.offset = offset return copy } diff --git a/Tests/MereRunCoreTests/KVCacheForkTests.swift b/Tests/MereRunCoreTests/KVCacheForkTests.swift new file mode 100644 index 00000000..004e5410 --- /dev/null +++ b/Tests/MereRunCoreTests/KVCacheForkTests.swift @@ -0,0 +1,71 @@ +import Foundation +import MLX +import XCTest +@testable import MereRunCore + +/// `KVCacheSimple.update` writes new tokens with subscript assignment, which +/// rebinds the same `MLXArray` wrapper in place. A fork that shares the +/// wrapper objects therefore sees every later write on the parent — which +/// silently corrupted prefix-KV snapshots stored mid-request (the request's +/// remaining prefill and decode kept writing into the stored copy). These +/// tests pin the isolation contract. +final class KVCacheForkTests: XCTestCase { + override class func setUp() { + super.setUp() + MLXTestSupport.ensureMetalLibraryAvailable() + } + + private func makeKV(_ value: Float, tokens: Int) -> (MLXArray, MLXArray) { + ( + MLXArray.full([1, 2, tokens, 4], values: MLXArray(value)), + MLXArray.full([1, 2, tokens, 4], values: MLXArray(value)) + ) + } + + func testForkIsIsolatedFromLaterParentWrites() throws { + let parent = KVCacheSimple(step: 4) + let (k1, v1) = makeKV(1.0, tokens: 3) + _ = parent.update(keys: k1, values: v1) + + let fork = try XCTUnwrap(parent.fork() as? KVCacheSimple) + XCTAssertEqual(fork.offset, 3, "fork offset must be frozen at fork time") + + // Parent keeps decoding: writes MORE tokens into its buffers. With a + // wrapper-sharing fork these writes land in the fork's arrays too. + let (k2, v2) = makeKV(2.0, tokens: 1) + _ = parent.update(keys: k2, values: v2) + let (k3, v3) = makeKV(3.0, tokens: 1) + _ = parent.update(keys: k3, values: v3) + + // Read the fork's snapshot by appending one marker token: the first + // three tokens it returns must still be the pre-fork 1.0s, not the + // parent's later 2.0/3.0 writes. + let (kM, vM) = makeKV(9.0, tokens: 1) + let forkView = fork.update(keys: kM, values: vM) + MLX.eval(forkView.0) + let snapshot = forkView.0[0..., 0..., 0..<3, 0...].asArray(Float.self) + XCTAssertEqual(snapshot.count, 1 * 2 * 3 * 4) + XCTAssertTrue( + snapshot.allSatisfy { $0 == 1.0 }, + "fork observed the parent's post-fork writes (max=\(snapshot.max() ?? 0)) — snapshot corrupted" + ) + } + + func testForkedCacheCanDivergeIndependently() throws { + let parent = KVCacheSimple(step: 4) + let (k1, v1) = makeKV(1.0, tokens: 2) + _ = parent.update(keys: k1, values: v1) + + let fork = try XCTUnwrap(parent.fork() as? KVCacheSimple) + let (kF, vF) = makeKV(5.0, tokens: 1) + let forkView = fork.update(keys: kF, values: vF) + let (kP, vP) = makeKV(7.0, tokens: 1) + let parentView = parent.update(keys: kP, values: vP) + + MLX.eval(forkView.0, parentView.0) + let forkTail = forkView.0[0..., 0..., 2..<3, 0...].asArray(Float.self) + let parentTail = parentView.0[0..., 0..., 2..<3, 0...].asArray(Float.self) + XCTAssertTrue(forkTail.allSatisfy { $0 == 5.0 }) + XCTAssertTrue(parentTail.allSatisfy { $0 == 7.0 }) + } +} diff --git a/docs/configuration.md b/docs/configuration.md index 0af23ce9..7a494460 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -71,6 +71,18 @@ keeps semantic checkpoints ahead of ordinary chunk checkpoints when pruning. Continuous batching and SSD KV cache are not enabled by this flag. +### `MERERUN_LFM2_PREFIX_KV_CACHE` + +Opt-in (`1`): in-memory prompt-prefix reuse for the LFM2 chat runtime, +mirroring the Qwen-family implementation. Forked layer caches (attention KV +and short-conv states both support forking) are stored at prefill chunk +boundaries and the longest matching token prefix seeds later requests, so a +repeated or extended prompt re-prefills only its tail. Chunk-boundary +checkpoints only — Gemma4-style semantic chat-template checkpoints are not +yet derived for LFM2. Bounded to 4 entries with the shared retention planner. +Measured on a ~2.9k-token prompt: repeat requests drop from 11.7s to 0.3s +end to end. + ### `MERERUN_GEMMA4_MTP` Gemma 4 12B MTP is enabled by default when the managed