LFM2: opt-in prefix-KV reuse; KVCacheSimple fork isolation + tests

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,19 +339,106 @@ 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"
             ))
             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(

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
     }

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 })
+    }
+}

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