From 4fe232f9eab82844c6b62236c88dd14a8db7341e Mon Sep 17 00:00:00 2001 From: CodeAndCanvas728 Date: Mon, 5 Oct 2026 22:46:28 +0200 Subject: [PATCH] feat: resume a diverged hybrid prompt from a recurrent-state checkpoint Hybrid (recurrent + attention) models could only reuse a cached prompt that was an exact prefix of the new one, so editing earlier history re-prefilled the whole conversation. Snapshot the recurrent layers at turn boundaries while prefilling (the turn after the system prompt, then every 2048+ tokens, at most 4 per entry) and, when a prompt diverges, resume from the newest checkpoint before the divergence: recurrent layers are set from it, attention layers are trimmed back to it (plain KV always; a sliding-window ring only until it fills), and only the rest is prefilled. Co-Authored-By: Claude Sonnet 5.5 --- Sources/SwiftLM/Server.swift | 207 ++++++++++++++++-- .../HybridCacheCheckpointTests.swift | 200 +++++++++++++++++ 2 files changed, 388 insertions(+), 19 deletions(-) create mode 100644 tests/SwiftLMTests/HybridCacheCheckpointTests.swift diff --git a/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift index 243065f..52da5c3 100644 --- a/Sources/SwiftLM/Server.swift +++ b/Sources/SwiftLM/Server.swift @@ -1827,13 +1827,27 @@ actor ServerStats { +/// Recurrent (MambaCache) layer state captured after `position` prompt tokens, keyed by +/// layer index. Recurrent state cannot be rewound, so a hybrid prompt that diverges from a +/// cached one resumes from the last checkpoint before the divergence instead of from zero. +struct RecurrentCheckpoint { + let position: Int + let layers: [Int: [MLXArray]] +} + actor PromptCache { struct CachedState { let tokens: [Int] // Full token sequence that generated this KV state let states: [[MLXArray]] // Per-layer KV state arrays let metaStates: [[String]] // Per-layer metadata + /// Hybrid entries only: earlier recurrent states, ascending by position, each + /// strictly before `tokens.count` (the entry's own state covers that one). + var checkpoints: [RecurrentCheckpoint] = [] } + /// Most checkpoints kept per entry; each holds one small recurrent state per layer. + static let maxCheckpoints = 4 + /// Most recently used first. Holds at most `maxEntries`. private var entries: [CachedState] = [] private let maxEntries: Int @@ -1927,7 +1941,14 @@ actor PromptCache { /// state is only valid if captured at exactly `tokens.count`, which the generic /// post-first-token save cannot guarantee. Only the hybrid path (which snapshots at /// a known boundary) passes true. - func save(tokens: [Int], cache: [KVCache], allowRecurrent: Bool = false) { + /// + /// `checkpoints` are the recurrent states captured while prefilling this prompt; + /// `resumedFrom` is how many tokens a restore supplied, so the entry that supplied them + /// passes its own checkpoints (and, for an exact hit, its end state) on to this one. + func save( + tokens: [Int], cache: [KVCache], allowRecurrent: Bool = false, + checkpoints: [RecurrentCheckpoint] = [], resumedFrom: Int = 0 + ) { if !allowRecurrent, cache.contains(where: { $0 is MambaCache }) { return } @@ -1965,13 +1986,64 @@ actor PromptCache { // being compared. The near-prefix rule needs a state that restores at any prefix. let exactPrefixOnly = allowRecurrent || !Self.canRewindToAnyPrefix(cache: cache, metaStates: metaStates) + // Read the resumed-from entry before the removal below drops it. + let carried = allowRecurrent + ? inheritedCheckpoints(tokens: tokens, resumedFrom: resumedFrom, cache: cache) : [] + let merged = Self.mergedCheckpoints(carried + checkpoints, before: P) entries.removeAll { Self.isSuperseded($0.tokens, by: tokens, exactPrefixOnly: exactPrefixOnly) } - entries.insert(CachedState(tokens: tokens, states: states, metaStates: metaStates), at: 0) + entries.insert( + CachedState(tokens: tokens, states: states, metaStates: metaStates, checkpoints: merged), + at: 0) if entries.count > maxEntries { entries.removeLast(entries.count - maxEntries) } } + /// Checkpoints of the entry a restore resumed from, up to `resumedFrom`, plus that + /// entry's end state when the restore took all of it. Both are valid for `tokens` + /// because they lie inside the prefix it shares with the entry. + private func inheritedCheckpoints( + tokens: [Int], resumedFrom: Int, cache: [KVCache] + ) -> [RecurrentCheckpoint] { + guard resumedFrom > 0, + let source = entries.first(where: { + Self.commonPrefixLength($0.tokens, tokens) >= resumedFrom + && $0.states.count == cache.count + }) + else { return [] } + var inherited = source.checkpoints.filter { $0.position <= resumedFrom } + if source.tokens.count == resumedFrom { + var layers: [Int: [MLXArray]] = [:] + for (i, layer) in cache.enumerated() where layer is MambaCache { + layers[i] = source.states[i] + } + inherited.append(RecurrentCheckpoint(position: resumedFrom, layers: layers)) + } + return inherited + } + + /// Ascending, deduplicated checkpoints strictly before `end`, thinned to + /// `maxCheckpoints`. The first (the anchor after the system prompt) and the newest are + /// kept; otherwise the checkpoint closest to a neighbour goes, so the spread stays even. + static func mergedCheckpoints(_ all: [RecurrentCheckpoint], before end: Int) -> [RecurrentCheckpoint] { + var out: [RecurrentCheckpoint] = [] + for cp in all.sorted(by: { $0.position < $1.position }) + where cp.position > 0 && cp.position < end && cp.position != out.last?.position { + out.append(cp) + } + while out.count > maxCheckpoints { + let drop = (1.. Int { + min(cps[i].position - cps[i - 1].position, cps[i + 1].position - cps[i].position) + } + /// Drop cached entries under memory pressure: all of them, or all but the most recent. /// Returns how many entries were dropped. @discardableResult @@ -2072,18 +2144,19 @@ actor PromptCache { /// Hybrid-model restore. Recurrent state cannot be trimmed, only resumed, so a cached /// sequence is reusable only if it is an exact prefix of `newTokens` no longer than - /// `limit`; the longest such entry wins. Returns the cached length (tokens now in + /// `limit`, or (when the prompt diverges from it) if it holds a recurrent checkpoint at + /// or before the divergence and the attention layers can be rewound to that point. The + /// entry that resumes the most tokens wins. Returns the resumed length (tokens now in /// `cache`), or nil on a miss. func restoreExactPrefix(newTokens: [Int], limit: Int, into cache: [KVCache]) -> Int? { - // Longest exact-prefix entry wins. - guard let index = entries.indices - .filter({ i in - let e = entries[i] - return !e.tokens.isEmpty && e.tokens.count <= limit - && e.states.count == cache.count && newTokens.starts(with: e.tokens) - }) - .max(by: { entries[$0].tokens.count < entries[$1].tokens.count }) - else { + var best: (index: Int, length: Int, checkpoint: RecurrentCheckpoint?)? + for (i, e) in entries.enumerated() { + guard let r = hybridResumePoint(e, newTokens: newTokens, limit: limit, cache: cache), + r.length > (best?.length ?? 0) + else { continue } + best = (i, r.length, r.checkpoint) + } + guard let (index, length, checkpoint) = best else { misses += 1 return nil } @@ -2095,9 +2168,53 @@ actor PromptCache { // ArraysCache's metaState setter traps; it carries no metadata to restore. if !(layer is MambaCache) { layer.metaState = cached.metaStates[i] } } + if let checkpoint { + // Recurrent layers jump back to the checkpoint; attention layers drop the + // positions after it (`hybridResumePoint` checked they can). + for (i, state) in checkpoint.layers { + var layer = cache[i] + layer.state = state.map(detachedArray) + } + for layer in cache where !(layer is MambaCache) { + layer.trim(cached.tokens.count - length) + } + } hits += 1 - print("[SwiftLM] \u{1F5C2} Prompt cache HIT (hybrid): \(cached.tokens.count)/\(newTokens.count) tokens reused") - return cached.tokens.count + print("[SwiftLM] \u{1F5C2} Prompt cache HIT (hybrid): \(length)/\(newTokens.count) tokens reused" + + (checkpoint == nil ? "" : " (checkpoint of \(cached.tokens.count))")) + return length + } + + /// How many tokens of `entry` a hybrid restore can resume, and the checkpoint to + /// resume from when that is short of the whole entry. Nil when it cannot. + private func hybridResumePoint( + _ entry: CachedState, newTokens: [Int], limit: Int, cache: [KVCache] + ) -> (length: Int, checkpoint: RecurrentCheckpoint?)? { + guard !entry.tokens.isEmpty, entry.states.count == cache.count else { return nil } + if entry.tokens.count <= limit, newTokens.starts(with: entry.tokens) { + return (entry.tokens.count, nil) + } + // Diverged (or longer than `limit`): the newest checkpoint inside the shared prefix. + let reach = min(Self.commonPrefixLength(entry.tokens, newTokens), limit) + guard let cp = entry.checkpoints.last(where: { $0.position <= reach }), + canRewindAttention(entry, cache: cache) + else { return nil } + return (cp.position, cp) + } + + /// Whether every non-recurrent layer of `entry` can be trimmed exactly: a plain KV + /// cache always can; a sliding-window ring only until it fills (the same predicate as + /// `RotatingKVCache.isTrimmable(after:)`, read from the saved metaState because the + /// live cache has not been restored yet). + private func canRewindAttention(_ entry: CachedState, cache: [KVCache]) -> Bool { + for (i, layer) in cache.enumerated() { + if layer is MambaCache || type(of: layer) == KVCacheSimple.self { continue } + guard layer is RotatingKVCache, entry.metaStates[i].count > 3, + let maxSize = Int(entry.metaStates[i][1]), let offset = Int(entry.metaStates[i][3]), + offset < maxSize + else { return false } + } + return true } func stats() -> (hits: Int, misses: Int) { (hits, misses) } @@ -2126,6 +2243,41 @@ func hybridCacheBoundary(promptTokens: [Int], imStartId: Int?, cache: [KVCache]) return boundary } +/// Prompt positions at which a hybrid prefill from `start` to `boundary` should snapshot +/// its recurrent state: turn starts, so a later prompt that edits history shares a prefix +/// that ends on a turn boundary. A cold prefill (`start == 0`) first takes the turn after +/// the system prompt (the anchor, which survives every later edit to the conversation); +/// after that, a turn start only once `minGap` tokens have passed since the last snapshot +/// (the resume point counts as one), which bounds how much an edit has to re-prefill. +func hybridCheckpointPositions( + promptTokens: [Int], imStartId: Int?, start: Int, boundary: Int, minGap: Int = 2048 +) -> [Int] { + guard let imStartId else { return [] } + let turnStarts = promptTokens.indices.filter { promptTokens[$0] == imStartId } + var positions: [Int] = [] + var last = start + if start == 0, turnStarts.count > 1, turnStarts[1] < boundary { + positions.append(turnStarts[1]) + last = turnStarts[1] + } + for i in turnStarts where i > last && i < boundary && i - last >= minGap { + positions.append(i) + last = i + } + return positions +} + +/// Detached, evaluated copy of every recurrent layer's state, keyed by layer index. The +/// live arrays are overwritten by the next prefill segment, so they must not be shared. +func recurrentSnapshot(_ cache: [KVCache]) -> [Int: [MLXArray]] { + var layers: [Int: [MLXArray]] = [:] + for (i, layer) in cache.enumerated() where layer is MambaCache { + layers[i] = layer.state.map(detachedArray) + } + eval(layers.values.flatMap { $0 }) + return layers +} + /// Copy-on-write snapshot of a cache state array. RotatingKVCache (and the KVCacheSimple /// setter) update their buffers in place, which would otherwise write through into a /// saved or restored snapshot that shares them. @@ -2501,12 +2653,29 @@ func handleChatCompletion( if start < boundary { // TokenIterator.init prefills its input into `cache` (same chunking and // SSD-streaming error handling as a cold prefill), then samples one token - // without feeding it back, so the cache ends at exactly `boundary`. - _ = try TokenIterator( - input: LMInput(text: sliceText(lmInput.text, from: start, to: boundary)), - model: context.model, cache: cache, parameters: params) + // without feeding it back, so the cache ends at exactly its input's end. + // Prefill in segments, snapshotting the recurrent layers between them, so a + // later prompt that edits earlier history can resume from the last snapshot + // before the edit instead of from zero. + let cuts = hybridCheckpointPositions( + promptTokens: promptTokens, + imStartId: context.tokenizer.convertTokenToId("<|im_start|>"), + start: start, boundary: boundary) + var checkpoints: [RecurrentCheckpoint] = [] + var from = start + for stop in cuts + [boundary] { + _ = try TokenIterator( + input: LMInput(text: sliceText(lmInput.text, from: from, to: stop)), + model: context.model, cache: cache, parameters: params) + if stop < boundary { + checkpoints.append( + RecurrentCheckpoint(position: stop, layers: recurrentSnapshot(cache))) + } + from = stop + } await promptCache.save( - tokens: Array(promptTokens[.. [any KVCache] { + let attn = KVCacheSimple() + _ = attn.update(keys: MLXArray.ones([1, 2, T, 4], dtype: .float16) * fill, + values: MLXArray.ones([1, 2, T, 4], dtype: .float16) * fill) + let mamba = MambaCache() + mamba.state = [MLXArray.ones([1, 3, 8]) * fill, MLXArray.ones([1, 2, 4, 4]) * fill] + return [attn, mamba] + } + + /// Recurrent state of layer 1 as the cache's second array (32 elements) filled with `fill`. + private func checkpoint(at position: Int, fill: Float) -> RecurrentCheckpoint { + RecurrentCheckpoint(position: position, layers: [ + 1: [MLXArray.ones([1, 3, 8]) * fill, MLXArray.ones([1, 2, 4, 4]) * fill], + ]) + } + + private func recurrentSum(_ cache: [any KVCache]) -> Float { + cache[1].state[1].sum().item(Float.self) + } + + // MARK: - Choosing checkpoint positions + + func testColdPrefillTakesAnchorThenSpacedTurnStarts() { + // Turn starts at 0, 3, 6, 12, 20 and boundary 20. + var tokens = Array(repeating: 1, count: 25) + for i in [0, 3, 6, 12, 20] { tokens[i] = imStart } + let cuts = hybridCheckpointPositions(promptTokens: tokens, imStartId: imStart, + start: 0, boundary: 20, minGap: 5) + XCTAssertEqual(cuts, [3, 12], "anchor after the system prompt, then the first start 5+ tokens on") + } + + func testResumedPrefillSkipsAnchorAndCountsGapFromResumePoint() { + var tokens = Array(repeating: 1, count: 30) + for i in [0, 3, 8, 14, 25] { tokens[i] = imStart } + let cuts = hybridCheckpointPositions(promptTokens: tokens, imStartId: imStart, + start: 6, boundary: 25, minGap: 5) + XCTAssertEqual(cuts, [14], "8 is only 2 past the resume point; 14 is 8 past it") + } + + func testNoCheckpointsWithoutImStartOrBeforeAnchor() { + XCTAssertEqual(hybridCheckpointPositions(promptTokens: [1, 2, 3], imStartId: nil, + start: 0, boundary: 3), []) + // Only the boundary turn follows the system prompt: nothing to checkpoint. + XCTAssertEqual(hybridCheckpointPositions(promptTokens: [imStart, 1, imStart, 2], imStartId: imStart, + start: 0, boundary: 2), []) + } + + // MARK: - Thinning + + func testThinningKeepsAnchorAndNewestAndDropsTheClosestPair() { + let all = [100, 2000, 2100, 4500, 7000, 9500].map { checkpoint(at: $0, fill: 1) } + let kept = PromptCache.mergedCheckpoints(all, before: 20_000).map(\.position) + XCTAssertEqual(kept.count, PromptCache.maxCheckpoints) + XCTAssertEqual(kept.first, 100, "anchor stays") + XCTAssertEqual(kept.last, 9500, "newest stays") + XCTAssertFalse(kept.contains(2000) && kept.contains(2100), "the 100-token gap is thinned first") + } + + func testMergeDropsPositionsAtOrPastEntryEndAndDuplicates() { + let all = [checkpoint(at: 5, fill: 1), checkpoint(at: 5, fill: 2), checkpoint(at: 9, fill: 3), + checkpoint(at: 10, fill: 4)] + XCTAssertEqual(PromptCache.mergedCheckpoints(all, before: 10).map(\.position), [5, 9]) + } + + // MARK: - Restoring from a checkpoint + + func testDivergentPromptResumesFromLatestCheckpointBeforeTheEdit() async { + let pc = PromptCache() + await pc.save(tokens: [1, 2, 3, 4, 5, 6], cache: makeHybridCache(seqLen: 6, fill: 9), + allowRecurrent: true, + checkpoints: [checkpoint(at: 2, fill: 2), checkpoint(at: 4, fill: 4)]) + + // Edit at index 4: tokens 0..<4 are shared. + let fresh: [any KVCache] = [KVCacheSimple(), MambaCache()] + let n = await pc.restoreExactPrefix(newTokens: [1, 2, 3, 4, 8, 8, 8], limit: 7, into: fresh) + + XCTAssertEqual(n, 4) + XCTAssertEqual(fresh[0].offset, 4, "attention layer rewound to the checkpoint") + XCTAssertEqual(recurrentSum(fresh), 4 * 32, "recurrent state is the checkpoint's, not the entry's") + } + + func testCheckpointAfterTheDivergenceIsNotUsed() async { + let pc = PromptCache() + await pc.save(tokens: [1, 2, 3, 4, 5, 6], cache: makeHybridCache(seqLen: 6, fill: 9), + allowRecurrent: true, checkpoints: [checkpoint(at: 5, fill: 5)]) + let n = await pc.restoreExactPrefix(newTokens: [1, 2, 3, 9, 9, 9, 9], limit: 7, + into: [KVCacheSimple(), MambaCache()]) + XCTAssertNil(n, "the only checkpoint is past where the prompts diverge") + } + + func testExactPrefixStillBeatsAnEarlierCheckpoint() async { + let pc = PromptCache() + await pc.save(tokens: [1, 2, 3, 4], cache: makeHybridCache(seqLen: 4, fill: 4), + allowRecurrent: true, checkpoints: [checkpoint(at: 2, fill: 2)]) + let fresh: [any KVCache] = [KVCacheSimple(), MambaCache()] + let n = await pc.restoreExactPrefix(newTokens: [1, 2, 3, 4, 5, 6], limit: 6, into: fresh) + XCTAssertEqual(n, 4) + XCTAssertEqual(recurrentSum(fresh), 4 * 32) + } + + func testEntryWithLongestResumePointWinsAcrossEntries() async { + let pc = PromptCache(maxEntries: 2) + await pc.save(tokens: [1, 2, 3, 4, 5], cache: makeHybridCache(seqLen: 5, fill: 1), + allowRecurrent: true, checkpoints: [checkpoint(at: 2, fill: 2)]) + await pc.save(tokens: [1, 2, 3, 4, 6, 6], cache: makeHybridCache(seqLen: 6, fill: 1), + allowRecurrent: true, checkpoints: [checkpoint(at: 3, fill: 3)]) + let fresh: [any KVCache] = [KVCacheSimple(), MambaCache()] + let n = await pc.restoreExactPrefix(newTokens: [1, 2, 3, 4, 9, 9], limit: 6, into: fresh) + XCTAssertEqual(n, 3) + XCTAssertEqual(recurrentSum(fresh), 3 * 32) + } + + func testWrappedRingRefusesCheckpointRestore() async { + // A ring that has dropped tokens can't be rewound exactly. + let ring = RotatingKVCache(maxSize: 4, keep: 0, step: 4) + for t in 0 ..< 6 { + let k = MLXArray([Float(t)]).reshaped([1, 1, 1, 1]) + _ = ring.update(keys: k, values: k) + } + let mamba = MambaCache() + mamba.state = [MLXArray.ones([1, 3, 8]), MLXArray.ones([1, 2, 4, 4])] + let pc = PromptCache() + await pc.save(tokens: [1, 2, 3, 4, 5, 6], cache: [ring, mamba], allowRecurrent: true, + checkpoints: [checkpoint(at: 2, fill: 2)]) + let n = await pc.restoreExactPrefix(newTokens: [1, 2, 9, 9], limit: 4, + into: [RotatingKVCache(maxSize: 4, keep: 0, step: 4), MambaCache()]) + XCTAssertNil(n) + } + + func testUnwrappedRingRestoresFromCheckpoint() async { + let ring = RotatingKVCache(maxSize: 16, keep: 0, step: 4) + for t in 0 ..< 6 { + let k = MLXArray([Float(t)]).reshaped([1, 1, 1, 1]) + _ = ring.update(keys: k, values: k) + } + let mamba = MambaCache() + mamba.state = [MLXArray.ones([1, 3, 8]), MLXArray.ones([1, 2, 4, 4])] + let pc = PromptCache() + await pc.save(tokens: [1, 2, 3, 4, 5, 6], cache: [ring, mamba], allowRecurrent: true, + checkpoints: [checkpoint(at: 2, fill: 2)]) + let fresh: [any KVCache] = [RotatingKVCache(maxSize: 16, keep: 0, step: 4), MambaCache()] + let n = await pc.restoreExactPrefix(newTokens: [1, 2, 9, 9], limit: 4, into: fresh) + XCTAssertEqual(n, 2) + XCTAssertEqual(fresh[0].offset, 2) + XCTAssertEqual(recurrentSum(fresh), 2 * 32) + } + + // MARK: - Checkpoints carry across turns + + func testExtensionInheritsCheckpointsAndKeepsThePreviousEndState() async { + let pc = PromptCache() + await pc.save(tokens: [1, 2, 3, 4], cache: makeHybridCache(seqLen: 4, fill: 4), + allowRecurrent: true, checkpoints: [checkpoint(at: 2, fill: 2)]) + // The next turn resumes from the whole entry, prefills on, and saves. + let live: [any KVCache] = [KVCacheSimple(), MambaCache()] + let resumed = await pc.restoreExactPrefix(newTokens: [1, 2, 3, 4, 5, 6, 7], limit: 7, into: live) + XCTAssertEqual(resumed, 4) + await pc.save(tokens: [1, 2, 3, 4, 5, 6, 7], cache: makeHybridCache(seqLen: 7, fill: 7), + allowRecurrent: true, checkpoints: [], resumedFrom: 4) + + // An edit at index 5 can now resume from the old end state at 4 ... + let a: [any KVCache] = [KVCacheSimple(), MambaCache()] + let nA = await pc.restoreExactPrefix(newTokens: [1, 2, 3, 4, 5, 8, 8], limit: 7, into: a) + XCTAssertEqual(nA, 4) + XCTAssertEqual(recurrentSum(a), 4 * 32) + // ... and an edit at index 2 from the inherited checkpoint at 2. + let b: [any KVCache] = [KVCacheSimple(), MambaCache()] + let nB = await pc.restoreExactPrefix(newTokens: [1, 2, 8, 8, 8], limit: 5, into: b) + XCTAssertEqual(nB, 2) + XCTAssertEqual(recurrentSum(b), 2 * 32) + } + + func testCheckpointRestoreThenSaveInheritsTheCheckpointsBeforeTheEdit() async { + let pc = PromptCache() + await pc.save(tokens: [1, 2, 3, 4, 5, 6], cache: makeHybridCache(seqLen: 6, fill: 6), + allowRecurrent: true, + checkpoints: [checkpoint(at: 2, fill: 2), checkpoint(at: 4, fill: 4)]) + let live: [any KVCache] = [KVCacheSimple(), MambaCache()] + let resumed = await pc.restoreExactPrefix(newTokens: [1, 2, 3, 4, 8, 8], limit: 6, into: live) + XCTAssertEqual(resumed, 4) + // Single-entry cache: the edited prompt replaces the original. + await pc.save(tokens: [1, 2, 3, 4, 8, 8], cache: makeHybridCache(seqLen: 6, fill: 8), + allowRecurrent: true, checkpoints: [], resumedFrom: 4) + let fresh: [any KVCache] = [KVCacheSimple(), MambaCache()] + let n = await pc.restoreExactPrefix(newTokens: [1, 2, 7, 7], limit: 4, into: fresh) + XCTAssertEqual(n, 2, "the checkpoint at 2 survived the replacement") + } +}