diff --git a/README.md b/README.md index e12f385..49f8ab2 100644 --- a/README.md +++ b/README.md @@ -638,6 +638,7 @@ curl http://localhost:5413/v1/chat/completions \ | `--gpu-layers` | `model_default`| Restrict the amount of layers allocated to GPU hardware | | `--stream-experts` | `false` | Enable SSD expert streaming for MoE models (10x speedup) | | `--turbo-kv` | `false` | Enable TurboQuant 3-bit KV cache compression (activates after 2048 tokens, server-wide) | +| `--prompt-cache-entries` | `1` | Prompt cache entries kept in memory (LRU; values below 1 are raised to 1 with a warning). More entries let interleaved sessions each keep their prefix, at the cost of one KV copy per entry. A newer prompt that extends a saved one replaces it instead of taking a second slot; for attention-only models the same goes for a saved prompt it differs from only in its last few tokens (a re-rendered generation prompt, an edited last message), except once a sliding-window layer has outgrown its window. Hybrid (recurrent) models replace only on an exact extension. Under macOS memory pressure the cache keeps only its most recent entry (warning) or is emptied (critical), so with the default of 1 only a critical event drops it | | `--draft-model` | (none) | Draft model path/ID for speculative decoding. When used with `--stream-experts`, `--num-draft-tokens` is auto-capped to 1 to minimise SSD I/O fan-out (see performance note above). | | `--num-draft-tokens` | `4` | Tokens per speculation round. Auto-capped to 1 when combined with `--stream-experts`. | | `--dflash` | `false` | Enable DFlash block-diffusion speculative decoding. Requires a compatible DFlash draft model | diff --git a/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift index b22472e..1820ca6 100644 --- a/Sources/SwiftLM/Server.swift +++ b/Sources/SwiftLM/Server.swift @@ -558,6 +558,9 @@ struct MLXServer: AsyncParsableCommand { @Option(name: .long, help: "Reject prompts longer than this many tokens with an OpenAI-style 400 (context_length_exceeded) before prefill. Default: no limit") var maxPromptTokens: Int? + @Option(name: .long, help: "Prompt cache entries kept in memory (LRU; values below 1 are raised to 1). More entries let interleaved sessions each keep their prefix, at the cost of one KV copy per entry. Under memory pressure the cache keeps only its most recent entry (warning) or is emptied (critical)") + var promptCacheEntries: Int = 1 + @Flag(name: .long, help: "Enable thinking/reasoning mode (Qwen3.5 etc). Default: disabled") var thinking: Bool = false @@ -675,6 +678,11 @@ struct MLXServer: AsyncParsableCommand { setenv("SWIFTLM_MTP_ENABLE", "1", 1) } + if self.promptCacheEntries < 1 { + print("[SwiftLM] Warning: --prompt-cache-entries must be at least 1, got \(self.promptCacheEntries). Using 1.") + self.promptCacheEntries = 1 + } + // Register SwiftLM-owned DFlash model types before any model loading. await registerDFlashModelTypes() @@ -1375,6 +1383,7 @@ struct MLXServer: AsyncParsableCommand { repeatPenalty: self.repeatPenalty, thinking: self.thinking, maxPromptTokens: self.maxPromptTokens, + promptCacheEntries: self.promptCacheEntries, tokenEcho: !self.noTokenEcho, isVision: loadedAsVision, prefillSize: self.prefillSize, @@ -1413,7 +1422,7 @@ struct MLXServer: AsyncParsableCommand { let ssdStr = self.streamExperts ? "enabled" : "disabled" let turboKVStr = config.turboKV ? "enabled" : "disabled" let mtpStr = config.mtp ? "enabled (\(config.numMtpTokens) tokens/round)" : "disabled" - print("[SwiftLM] Config: ctx_size=\(ctxSizeStr), temp=\(config.temp), top_p=\(config.topP), top_k=\(topKStr), min_p=\(minPStr), repeat_penalty=\(penaltyStr), parallel=\(parallelSlots), cors=\(corsStr), mem_limit=\(memLimitStr), auth=\(authStr), thinking=\(thinkingStr), ssd_stream=\(ssdStr), turbo_kv=\(turboKVStr), mtp=\(mtpStr)") + print("[SwiftLM] Config: ctx_size=\(ctxSizeStr), temp=\(config.temp), top_p=\(config.topP), top_k=\(topKStr), min_p=\(minPStr), repeat_penalty=\(penaltyStr), parallel=\(parallelSlots), cors=\(corsStr), mem_limit=\(memLimitStr), auth=\(authStr), thinking=\(thinkingStr), ssd_stream=\(ssdStr), turbo_kv=\(turboKVStr), mtp=\(mtpStr), prompt_cache_entries=\(config.promptCacheEntries)") if config.turboKV, let ctx = config.ctxSize { print("[SwiftLM] ⚠️ --turbo-kv has no effect with --ctx-size \(ctx): a bounded context gives the attention layers a RotatingKVCache, and TurboKV only compresses KVCacheSimple. Drop --ctx-size to use --turbo-kv.") } @@ -1486,7 +1495,28 @@ struct MLXServer: AsyncParsableCommand { } // Chat completions — handler extracted to avoid type-checker timeout - let promptCache = PromptCache() + let promptCache = PromptCache(maxEntries: config.promptCacheEntries) + // Cached KV is the largest evictable allocation; give it back before the OS has to + // compress or kill. A warning drops every entry but the most recent one (the live + // conversation), so with the default --prompt-cache-entries 1 it drops nothing; a + // critical event drops all of them, including that last one. Keep the source for + // the server's lifetime. + let memoryPressure = DispatchSource.makeMemoryPressureSource(eventMask: [.warning, .critical]) + memoryPressure.setEventHandler { + let critical = memoryPressure.data.contains(.critical) + Task { + let dropped = await promptCache.evict(keepMostRecent: !critical) + guard dropped > 0 else { return } + // The evicted KV buffers go back to MLX's allocator cache, not to the OS, so + // without this the eviction frees almost no memory. Safe while a request is + // generating: it only releases buffers the allocator holds unused. It takes + // MLX's global eval lock and so may wait for an in-flight eval, which is why + // it runs here and not on the GCD handler thread. + Memory.clearCache() + print("[SwiftLM] Memory pressure (\(critical ? "critical" : "warning")): evicted \(dropped) prompt cache \(dropped == 1 ? "entry" : "entries")") + } + } + memoryPressure.resume() router.post("/v1/chat/completions") { request, _ -> Response in do { let bodyData = try await collectBody(request) @@ -1671,6 +1701,8 @@ struct ServerConfig: Sendable { let thinking: Bool /// `--max-prompt-tokens`: longer prompts are rejected before prefill. nil = no limit. let maxPromptTokens: Int? + /// `--prompt-cache-entries`: LRU capacity of the prompt cache. + let promptCacheEntries: Int /// Echo each generated chunk to stdout (`--no-token-echo` turns it off). let tokenEcho: Bool let isVision: Bool @@ -1802,10 +1834,89 @@ actor PromptCache { let metaStates: [[String]] // Per-layer metadata } - private var cached: CachedState? + /// Most recently used first. Holds at most `maxEntries`. + private var entries: [CachedState] = [] + private let maxEntries: Int private var hits: Int = 0 private var misses: Int = 0 + init(maxEntries: Int = 1) { self.maxEntries = max(1, maxEntries) } + + /// How many trailing tokens of a saved prompt a newer prompt may drop or rewrite and + /// still replace it (see `isSuperseded`). + /// + /// The generic save runs after the first decode token, so the prompt it saves ends with + /// the generation prompt (`<|im_start|>assistant\n\n`). Templates that re-render + /// history (Qwen3 reasoning templates drop or rewrite it) and edits or regenerations of + /// the last message make the next turn diverge from the previous prompt only in its + /// last few tokens, so the old prompt is no longer an exact prefix even though it is + /// the same conversation. 16 covers those generation prompts with room to spare. + /// + /// For a cache that can be rewound to any shorter prefix, replacing is bounded: the new + /// entry holds the KV for the shared prefix, so any prompt that could have reused the + /// old entry loses at most this many tokens of reuse, and the slot goes to something + /// else instead of a near-duplicate. A state that cannot be rewound that way (recurrent + /// layers, a sliding window that has already dropped tokens) has no such bound and is + /// only replaced by an exact extension; `save` picks the rule. + static let supersedeSlack = 16 + + /// Length of the shared prefix of two token sequences. + private static func commonPrefixLength(_ a: [Int], _ b: [Int]) -> Int { + var n = 0 + for (x, y) in zip(a, b) { + guard x == y else { break } + n += 1 + } + return n + } + + /// Whether saving `new` makes the saved `old` redundant. + /// + /// An exact prefix always does. With `exactPrefixOnly` false (a cache that can be rewound + /// to any shorter prefix) so does a near-prefix: `old` diverges from `new` in at most + /// `supersedeSlack` trailing tokens, and what it shares with `new` outweighs what it does + /// not. The second condition keeps entries that merely start alike apart: a zero-match + /// entry is never removable, and a short one (a shared BOS or chat-template header) is + /// not replaced by an unrelated prompt just because the whole entry is shorter than the + /// slack. Two conversations that share only a system prompt both stay as long as each + /// carries more than `supersedeSlack` tokens of its own. + /// + /// `exactPrefixOnly` is set when the new state cannot stand in for the old one on a + /// partial match (see `canRewindToAnyPrefix`): recurrent (hybrid) snapshots are valid + /// only at exactly `tokens.count` and consumed by `restoreExactPrefix`, and a wrapped + /// ring buffer is rejected by `usableMatch` for any rewind past one token, so for them a + /// near-prefix replacement would lose the whole hit, not a few tokens. + static func isSuperseded(_ old: [Int], by new: [Int], exactPrefixOnly: Bool) -> Bool { + if new.starts(with: old) { return true } + if exactPrefixOnly { return false } + let shared = commonPrefixLength(old, new) + let stale = old.count - shared + return stale <= supersedeSlack && shared > stale + } + + /// Whether a sliding-window ring buffer has already dropped tokens (offset beyond its + /// window). Such a ring can only be rewound by the one slot a full-match replay + /// overwrites; anything deeper would leave the window pointing at evicted keys. + private static func ringHasDroppedTokens(metaState meta: [String]) -> Bool { + // RotatingKVCache.metaState: keep, maxSize, step, offset, idx, ... + guard meta.count > 3, let maxSize = Int(meta[1]), let offset = Int(meta[3]) else { + return false + } + return offset > maxSize + } + + /// Whether a state saved from `cache` can be restored at any shorter prefix, so that it + /// can stand in for a near-prefix entry it replaces: no recurrent layer, and no ring + /// buffer that has dropped tokens. This is the same condition `usableMatch` applies when + /// restoring. + private static func canRewindToAnyPrefix(cache: [KVCache], metaStates: [[String]]) -> Bool { + for (layer, meta) in zip(cache, metaStates) { + if layer is MambaCache { return false } + if layer is RotatingKVCache, ringHasDroppedTokens(metaState: meta) { return false } + } + return true + } + /// Save the full prompt token sequence and its KV state. /// IMPORTANT: We must eval() the state arrays immediately. The state getter may /// produce lazy computation graphs (e.g. TurboKV decode → reshape → concatenate). @@ -1820,6 +1931,8 @@ actor PromptCache { if !allowRecurrent, cache.contains(where: { $0 is MambaCache }) { return } + // Nothing can restore an empty prompt, and it would only take a slot. + guard !tokens.isEmpty else { return } let P = tokens.count // For attention KVCacheSimple layers, the state tensor is [B, H, T, D] with a // pre-allocated T that can exceed the actual prompt length P. If we store the @@ -1846,10 +1959,30 @@ actor PromptCache { if !allArrays.isEmpty { eval(allArrays) } - cached = CachedState(tokens: tokens, states: states, metaStates: metaStates) + // An entry the new one dominates goes: a conversation replaces its own previous turn + // instead of filling the cache. Entries are homogeneous within a run (the generic + // path refuses recurrent layers), so the new state's kind also describes the entries + // being compared. The near-prefix rule needs a state that restores at any prefix. + let exactPrefixOnly = allowRecurrent + || !Self.canRewindToAnyPrefix(cache: cache, metaStates: metaStates) + entries.removeAll { + Self.isSuperseded($0.tokens, by: tokens, exactPrefixOnly: exactPrefixOnly) + } + entries.insert(CachedState(tokens: tokens, states: states, metaStates: metaStates), at: 0) + if entries.count > maxEntries { entries.removeLast(entries.count - maxEntries) } + } + + /// Drop cached entries under memory pressure: all of them, or all but the most recent. + /// Returns how many entries were dropped. + @discardableResult + func evict(keepMostRecent: Bool) -> Int { + let keep = keepMostRecent ? 1 : 0 + let dropped = max(0, entries.count - keep) + if dropped > 0 { entries.removeLast(dropped) } + return dropped } - /// Find the longest common prefix between `newTokens` and the cached sequence. + /// Find the cached entry with the longest usable common prefix with `newTokens`. /// Restores matched KV state, trims any excess — mirrors llama-server behaviour. /// Returns the number of matched tokens, or nil on a complete miss. func restore(newTokens: [Int], into cache: [KVCache]) -> Int? { @@ -1861,10 +1994,6 @@ actor PromptCache { return nil } - guard let cached, !cached.tokens.isEmpty else { - misses += 1 - return nil - } // ── Recurrent-layer safety gate ── // MambaCache (and other recurrent caches) store a 2-D hidden state with no // T dimension, so the dim(2) read below would crash. Hybrid Mamba/attention @@ -1879,73 +2008,87 @@ actor PromptCache { misses += 1 return nil } - // Token-by-token longest common prefix scan - var matchLen = 0 - for (a, b) in zip(cached.tokens, newTokens) { - guard a == b else { break } - matchLen += 1 + var best: (index: Int, matchLen: Int)? + for (i, entry) in entries.enumerated() { + if let m = usableMatch(entry, newTokens: newTokens, cache: cache), + m > (best?.matchLen ?? 0) + { + best = (i, m) + } } - guard matchLen > 0 else { + guard let (index, matchLen) = best else { misses += 1 return nil } + let cached = entries.remove(at: index) + entries.insert(cached, at: 0) + let excess = cached.tokens.count - matchLen + // Safe to restore: trim won't corrupt any layer. Detach so decode steps on the + // live cache cannot write through into the saved snapshot. + for i in 0.. 0 { + for layer in cache { layer.trim(excess) } + } + hits += 1 + print("[SwiftLM] \u{1F5C2} Prompt cache HIT: \(matchLen)/\(newTokens.count) tokens reused (\(excess > 0 ? "partial" : "full") match)") + return matchLen + } + + /// Length of the common prefix between `entry` and `newTokens`, or nil when restoring + /// that entry into `cache` and trimming the excess would corrupt a layer. + private func usableMatch(_ entry: CachedState, newTokens: [Int], cache: [KVCache]) -> Int? { + guard !entry.tokens.isEmpty else { return nil } + let matchLen = Self.commonPrefixLength(entry.tokens, newTokens) + guard matchLen > 0 else { return nil } // Pre-flight safety check: compute the minimum sequence length across // all cached layers. Sliding-window layers (RotatingKVCache) store far // fewer tokens than the full prompt (e.g. 1440 vs 5537). If the trim // would zero-out any layer, bail BEFORE touching the live cache. - let excess = cached.tokens.count - matchLen + let excess = entry.tokens.count - matchLen if excess > 0 { // The state getter stores keys as the first element: [B, H, T, D] // dim(2) = T = the number of cached tokens for that layer. - let minCachedSeqLen = cached.states.map { arrays -> Int in - guard let firstArray = arrays.first else { return 0 } - guard firstArray.ndim >= 3 else { return 0 } - return firstArray.dim(2) // T dimension + let minCachedSeqLen = entry.states.map { arrays -> Int in + guard let firstArray = arrays.first, firstArray.ndim >= 3 else { return 0 } + return firstArray.dim(2) }.min() ?? 0 - if excess >= minCachedSeqLen { - // Trim would empty or corrupt at least one layer → treat as miss - misses += 1 - return nil - } + if excess >= minCachedSeqLen { return nil } } // A ring buffer that has evicted tokens (offset > maxSize) cannot be rewound by // more than the one slot a full-match replay overwrites: trim(n >= 2) would leave // the window pointing at evicted keys. Treat that as a miss. let replayExtra = matchLen >= newTokens.count ? 1 : 0 - for i in 0.. 3, let maxSize = Int(meta[1]), let offset = Int(meta[3]), - offset > maxSize, excess + replayExtra > 1 - { - misses += 1 + for i in 0.. 1 { return nil } } - // Safe to restore: trim won't corrupt any layer. Detach so decode steps on the - // live cache cannot write through into the saved snapshot. - for i in 0.. 0 { - for layer in cache { layer.trim(excess) } - } - hits += 1 - print("[SwiftLM] \u{1F5C2} Prompt cache HIT: \(matchLen)/\(newTokens.count) tokens reused (\(excess > 0 ? "partial" : "full") match)") return matchLen } - /// Hybrid-model restore. Recurrent state cannot be trimmed, only resumed, so the - /// cached sequence is reusable only if it is an exact prefix of `newTokens` no longer - /// than `limit`. Returns the cached length (tokens now in `cache`), or nil on a miss. + /// 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 + /// `cache`), or nil on a miss. func restoreExactPrefix(newTokens: [Int], limit: Int, into cache: [KVCache]) -> Int? { - guard let cached, !cached.tokens.isEmpty, cached.tokens.count <= limit, - cached.states.count == cache.count, newTokens.starts(with: cached.tokens) + // 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 { misses += 1 return nil } + let cached = entries.remove(at: index) + entries.insert(cached, at: 0) for (i, layer) in cache.enumerated() { var layer = layer layer.state = cached.states[i].map(detachedArray) diff --git a/tests/SwiftLMTests/PromptCacheLRUTests.swift b/tests/SwiftLMTests/PromptCacheLRUTests.swift new file mode 100644 index 0000000..9a66449 --- /dev/null +++ b/tests/SwiftLMTests/PromptCacheLRUTests.swift @@ -0,0 +1,635 @@ +import XCTest +import MLX +import MLXLMCommon +@testable import SwiftLM + +/// The prompt cache keeps several entries so interleaved sessions (a side request between +/// two agent turns) don't evict each other's prefix. +final class PromptCacheLRUTests: XCTestCase { + + private func makeCache(seqLen T: Int, fill: Float) -> [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) + return [attn] + } + + private func makeHybrid(seqLen T: Int, fill: Float) -> [any KVCache] { + let mamba = MambaCache() + mamba.state = [MLXArray.ones([1, 3, 8]) * fill, MLXArray.ones([1, 2, 4, 4]) * fill] + return makeCache(seqLen: T, fill: fill) + [mamba] + } + + private func hit(_ pc: PromptCache, _ tokens: [Int]) async -> Int? { + await pc.restore(newTokens: tokens, into: [KVCacheSimple()]) + } + + /// `count` consecutive token ids from `start`. + private func seq(_ start: Int, _ count: Int) -> [Int] { Array(start ..< start + count) } + + /// Save `tokens` with a KV state whose every element is `fill`. + private func put(_ pc: PromptCache, _ tokens: [Int], fill: Float = 1) async { + await pc.save(tokens: tokens, cache: makeCache(seqLen: tokens.count, fill: fill)) + } + + /// Entries still held (destructive: empties the cache). + private func drain(_ pc: PromptCache) async -> Int { + await pc.evict(keepMostRecent: false) + } + + /// Sum of every key element, i.e. `fill * 8 * T` for a `makeCache` layer. + private func keySum(_ layer: any KVCache) -> Float { + layer.state[0].asType(.float32).sum().item(Float.self) + } + + /// A ring fed `n` single tokens; values encode the token index. + private func makeRing(maxSize: Int, tokens n: Int) -> RotatingKVCache { + let ring = RotatingKVCache(maxSize: maxSize, keep: 0, step: 4) + for t in 0 ..< n { + let k = MLXArray([Float(t)]).reshaped([1, 1, 1, 1]) + _ = ring.update(keys: k, values: k) + } + return ring + } + + func testTwoUnrelatedPromptsBothHitWithTwoEntries() async { + let pc = PromptCache(maxEntries: 2) + await pc.save(tokens: [1, 2, 3, 4], cache: makeCache(seqLen: 4, fill: 1)) + await pc.save(tokens: [9, 8, 7, 6], cache: makeCache(seqLen: 4, fill: 2)) + let a = await hit(pc, [1, 2, 3, 4, 5]) + let b = await hit(pc, [9, 8, 7, 6, 5]) + XCTAssertEqual(a, 4) + XCTAssertEqual(b, 4) + } + + func testSingleEntryKeepsTodaysBehaviour() async { + let pc = PromptCache() + await pc.save(tokens: [1, 2, 3, 4], cache: makeCache(seqLen: 4, fill: 1)) + await pc.save(tokens: [9, 8, 7, 6], cache: makeCache(seqLen: 4, fill: 2)) + let first = await hit(pc, [1, 2, 3, 4, 5]) + XCTAssertNil(first) + } + + func testHitRefreshesRecencyAndLeastRecentIsDropped() async { + let pc = PromptCache(maxEntries: 2) + await pc.save(tokens: [1, 2, 3, 4], cache: makeCache(seqLen: 4, fill: 1)) + await pc.save(tokens: [9, 8, 7, 6], cache: makeCache(seqLen: 4, fill: 2)) + _ = await hit(pc, [1, 2, 3, 4, 5]) // [1..] is now most recent + await pc.save(tokens: [5, 5, 5, 5], cache: makeCache(seqLen: 4, fill: 3)) // evicts [9..] + let kept = await hit(pc, [1, 2, 3, 4, 5]) + let dropped = await hit(pc, [9, 8, 7, 6, 5]) + XCTAssertEqual(kept, 4) + XCTAssertNil(dropped) + } + + func testExtensionReplacesItsPrefixEntry() async { + let pc = PromptCache(maxEntries: 2) + await pc.save(tokens: [1, 2, 3], cache: makeCache(seqLen: 3, fill: 1)) + await pc.save(tokens: [9, 9, 9], cache: makeCache(seqLen: 3, fill: 2)) + await pc.save(tokens: [1, 2, 3, 4, 5], cache: makeCache(seqLen: 5, fill: 1)) + // The other session must survive: the linear conversation replaced its own entry. + let other = await hit(pc, [9, 9, 9, 1]) + let conv = await hit(pc, [1, 2, 3, 4, 5, 6]) + XCTAssertEqual(other, 3) + XCTAssertEqual(conv, 5) + } + + func testRestorePicksLongestMatch() async { + let pc = PromptCache(maxEntries: 3) + await pc.save(tokens: [1, 2, 3, 4, 5, 6], cache: makeCache(seqLen: 6, fill: 1)) + await pc.save(tokens: [1, 2, 9, 9, 9, 9], cache: makeCache(seqLen: 6, fill: 2)) + let n = await hit(pc, [1, 2, 3, 4, 7, 8]) + XCTAssertEqual(n, 4) + } + + func testHybridExactPrefixPicksLongestEntry() async { + let pc = PromptCache(maxEntries: 3) + await pc.save(tokens: [1, 2, 3], cache: makeHybrid(seqLen: 3, fill: 1), allowRecurrent: true) + await pc.save(tokens: [7, 7, 7, 7], cache: makeHybrid(seqLen: 4, fill: 3), allowRecurrent: true) + await pc.save(tokens: [1, 2, 3, 4, 5], cache: makeHybrid(seqLen: 5, fill: 2), allowRecurrent: true) + // The 5-token entry replaced the 3-token one, so [1,2,3,4,5,...] reuses all 5. + let fresh: [any KVCache] = [KVCacheSimple(), MambaCache()] + let n = await pc.restoreExactPrefix(newTokens: [1, 2, 3, 4, 5, 6], limit: 5, into: fresh) + XCTAssertEqual(n, 5) + let other = await pc.restoreExactPrefix(newTokens: [7, 7, 7, 7, 1], limit: 4, + into: [KVCacheSimple(), MambaCache()]) + XCTAssertEqual(other, 4, "the other session's snapshot survives") + } + + func testEvictKeepMostRecentAndEvictAll() async { + let pc = PromptCache(maxEntries: 3) + await pc.save(tokens: [1, 2, 3, 4], cache: makeCache(seqLen: 4, fill: 1)) + await pc.save(tokens: [9, 8, 7, 6], cache: makeCache(seqLen: 4, fill: 2)) + await pc.evict(keepMostRecent: true) + let old = await hit(pc, [1, 2, 3, 4, 5]) + let recent = await hit(pc, [9, 8, 7, 6, 5]) + XCTAssertNil(old) + XCTAssertEqual(recent, 4) + await pc.evict(keepMostRecent: false) + let none = await hit(pc, [9, 8, 7, 6, 5]) + XCTAssertNil(none) + } + + // MARK: - Superseding a stale entry of the same conversation + + /// The prompt saved for turn N is not always an exact prefix of turn N+1: templates + /// that re-render history drop or rewrite the generation-prompt tokens it ended with + /// (Qwen3 `` handling), and edits or regenerations change the last message. That + /// entry must be replaced, not left to push another conversation out of the cache. + func testNearPrefixReplacesItsStaleEntryInsteadOfEvictingTheOtherSession() async { + let pc = PromptCache(maxEntries: 2) + await put(pc, [9, 9, 9, 9], fill: 1) + await put(pc, [1, 2, 3, 4, 5, 6], fill: 2) + let partial = await hit(pc, [1, 2, 3, 4, 7, 8, 9]) + XCTAssertEqual(partial, 4) + await put(pc, [1, 2, 3, 4, 7, 8, 9], fill: 3) + + let other = await hit(pc, [9, 9, 9, 9, 0]) + XCTAssertEqual(other, 4, "the other session survives") + let stale = await hit(pc, [1, 2, 3, 4, 5, 6, 0]) + XCTAssertEqual(stale, 4, "the stale entry is gone: only the replacement's 4-token prefix matches") + let held = await drain(pc) + XCTAssertEqual(held, 2) + } + + func testPingPongWithReRenderedHistoryKeepsTheOtherConversation() async { + let b1 = seq(5000, 40) + let history = seq(0, 60) + // `<|im_start|>assistant\n\n\n\n\n`: 7 tokens that the next turn's + // re-rendered history does not repeat. This pins the slack above that real tail. + let a1 = history + seq(900, 7) + let a2 = history + seq(700, 25) + let pc = PromptCache(maxEntries: 2) + await put(pc, b1) + await put(pc, a1) + let partial = await hit(pc, a2) + XCTAssertEqual(partial, history.count) + await put(pc, a2) + + let b = await hit(pc, b1 + [1]) + let a = await hit(pc, a2 + [1]) + let staleA1 = await hit(pc, a1 + [1]) + XCTAssertEqual(b, b1.count, "B1 survives the A1 -> A2 turn") + XCTAssertEqual(a, a2.count) + XCTAssertEqual(staleA1, history.count, "A1's tail is gone") + let held = await drain(pc) + XCTAssertEqual(held, 2) + } + + func testSupersedesWhenTheStaleTailIsExactlyTheSlack() async { + let slack = PromptCache.supersedeSlack + let old = seq(0, 3 * slack) + let other = seq(1000, 3 * slack) + let new = seq(0, 2 * slack) + seq(2000, 2 * slack) // drops the last `slack` tokens of old + let pc = PromptCache(maxEntries: 2) + await put(pc, other) // least recent: it is what a leftover `old` would push out + await put(pc, old) + await put(pc, new) + + let oldTail = await hit(pc, old + [5]) + let otherHit = await hit(pc, other + [5]) + let newHit = await hit(pc, new + [5]) + XCTAssertEqual(oldTail, 2 * slack, "old was replaced; only the shared prefix remains reusable") + XCTAssertEqual(otherHit, other.count) + XCTAssertEqual(newHit, new.count) + let held = await drain(pc) + XCTAssertEqual(held, 2) + } + + func testKeepsBothWhenTheStaleTailExceedsTheSlack() async { + let slack = PromptCache.supersedeSlack + let old = seq(0, 3 * slack) + let other = seq(1000, 3 * slack) + let new = seq(0, 2 * slack - 1) + seq(2000, 2 * slack) // old's tail is slack + 1 + let pc = PromptCache(maxEntries: 3) + await put(pc, old) + await put(pc, other) + await put(pc, new) + + let oldHit = await hit(pc, old + [5]) + let otherHit = await hit(pc, other + [5]) + let newHit = await hit(pc, new + [5]) + XCTAssertEqual(oldHit, old.count) + XCTAssertEqual(otherHit, other.count) + XCTAssertEqual(newHit, new.count) + let held = await drain(pc) + XCTAssertEqual(held, 3) + } + + /// Two conversations that share a long system prompt and then diverge are different + /// conversations. A request that only restored the shared system prompt from one must + /// not make saving the other evict it. + func testConversationsSharingOnlyASystemPromptBothStay() async { + // Fixed sizes on purpose: 20-token tails against a 32-token shared system prompt. + // The shared part outweighs each tail, so only the slack keeps them apart. + XCTAssertLessThan(PromptCache.supersedeSlack, 20, "the tails must exceed the slack for this test to mean anything") + let system = seq(0, 32) + let a = system + seq(1000, 20) + let b = system + seq(2000, 20) + let pc = PromptCache(maxEntries: 2) + await put(pc, a) + let shared = await hit(pc, b) + XCTAssertEqual(shared, system.count, "B's request reuses only the shared system prompt") + await put(pc, b) + + let aHit = await hit(pc, a + [7]) + let bHit = await hit(pc, b + [7]) + XCTAssertEqual(aHit, a.count) + XCTAssertEqual(bHit, b.count) + let held = await drain(pc) + XCTAssertEqual(held, 2) + } + + /// Entries that only start alike (a shared BOS or template header) and are shorter than + /// the slack must not be replaced by an unrelated prompt. + func testShortEntriesSharingOnlyAHeaderAreNotSuperseded() async { + let pc = PromptCache(maxEntries: 2) + await put(pc, [0, 5, 5, 5, 5]) + await put(pc, [0, 6, 6, 6, 6]) + let first = await hit(pc, [0, 5, 5, 5, 5, 1]) + let second = await hit(pc, [0, 6, 6, 6, 6, 1]) + XCTAssertEqual(first, 5) + XCTAssertEqual(second, 5) + let held = await drain(pc) + XCTAssertEqual(held, 2) + } + + func testSavingTheSamePromptAgainReplacesIt() async { + let pc = PromptCache(maxEntries: 2) + await put(pc, [1, 2, 3, 4]) + await put(pc, [1, 2, 3, 4]) + let held = await drain(pc) + XCTAssertEqual(held, 1) + } + + func testEmptyPromptIsNeverSaved() async { + let pc = PromptCache(maxEntries: 1) + await put(pc, [1, 2, 3]) + await pc.save(tokens: [], cache: makeCache(seqLen: 2, fill: 2)) + let kept = await hit(pc, [1, 2, 3, 4]) + XCTAssertEqual(kept, 3, "an unrestorable empty entry must not push a real one out") + let held = await drain(pc) + XCTAssertEqual(held, 1) + } + + /// Recurrent snapshots are valid only at exactly `tokens.count`, so only an exact prefix + /// replaces one. The generic rule would collapse this pair. + func testHybridSavesKeepTheStrictPrefixRule() async { + let older = seq(0, 6) + let newer = seq(0, 4) + [7, 8, 9] // diverges from `older` in its last 2 tokens + let hybrid = PromptCache(maxEntries: 2) + await hybrid.save(tokens: older, cache: makeHybrid(seqLen: 6, fill: 1), allowRecurrent: true) + await hybrid.save(tokens: newer, cache: makeHybrid(seqLen: 7, fill: 2), allowRecurrent: true) + let olderHit = await hybrid.restoreExactPrefix( + newTokens: older + [10], limit: 6, into: [KVCacheSimple(), MambaCache()]) + let newerHit = await hybrid.restoreExactPrefix( + newTokens: newer + [10], limit: 7, into: [KVCacheSimple(), MambaCache()]) + XCTAssertEqual(olderHit, 6) + XCTAssertEqual(newerHit, 7) + let hybridHeld = await drain(hybrid) + XCTAssertEqual(hybridHeld, 2) + + let generic = PromptCache(maxEntries: 2) + await put(generic, older) + await put(generic, newer) + let genericHeld = await drain(generic) + XCTAssertEqual(genericHeld, 1, "the same pair collapses on the generic path") + } + + /// The pure rule behind the tests above. + func testIsSupersededRule() { + let slack = PromptCache.supersedeSlack + let old = seq(0, 3 * slack) + func superseded(_ new: [Int], exactOnly: Bool = false, saved: [Int]? = nil) -> Bool { + PromptCache.isSuperseded(saved ?? old, by: new, exactPrefixOnly: exactOnly) + } + XCTAssertTrue(superseded(old), "identical") + XCTAssertTrue(superseded(old + [1]), "extension") + XCTAssertTrue(superseded(seq(0, 2 * slack) + seq(900, 5)), "stale tail == slack") + XCTAssertFalse(superseded(seq(0, 2 * slack - 1) + seq(900, 5)), "stale tail == slack + 1") + XCTAssertFalse(superseded(seq(900, 40)), "no overlap is never superseded") + XCTAssertFalse(superseded([]), "nothing shared") + // Shared part must outweigh the stale part, whatever the slack. + XCTAssertTrue(superseded([1, 2, 3, 9], saved: [1, 2, 3, 4]), "3 shared, 1 stale") + XCTAssertFalse(superseded([1, 2, 9], saved: [1, 2, 3, 4]), "2 shared, 2 stale") + XCTAssertFalse(superseded([1, 2, 9, 9, 9, 9], saved: [1, 2, 3, 4, 5, 6]), "short entry, minority shared") + XCTAssertTrue(superseded([1, 2, 3, 4, 7, 8, 9], saved: [1, 2, 3, 4, 5, 6])) + // Recurrent snapshots: exact prefix only. + XCTAssertTrue(superseded(old + [1], exactOnly: true)) + XCTAssertTrue(superseded(old, exactOnly: true)) + XCTAssertFalse(superseded(seq(0, 2 * slack) + seq(900, 5), exactOnly: true)) + XCTAssertFalse(superseded(Array(old.dropLast()), exactOnly: true), "shorter prompt does not replace a longer snapshot") + } + + /// One save replaces every entry it dominates, including ones that do not dominate each + /// other. + func testOneSaveSupersedesEveryDominatedEntry() async { + let prefix = seq(0, 30) + let old1 = prefix + seq(1000, 20) + let old2 = prefix + seq(2000, 8) + let new = Array(old1.prefix(45)) + seq(3000, 10) + + let pair = PromptCache(maxEntries: 3) + await put(pair, old1) + await put(pair, old2) + let pairHeld = await drain(pair) + XCTAssertEqual(pairHeld, 2, "setup: old1 and old2 differ in more than the slack, so they coexist") + + let pc = PromptCache(maxEntries: 3) + await put(pc, old1) + await put(pc, old2) + await put(pc, new) // shares 45 of old1 and 30 of old2: both are within the slack + let held = await drain(pc) + XCTAssertEqual(held, 1) + } + + // MARK: - Sliding-window (ring buffer) caches + + /// A ring that has dropped tokens (offset beyond its window) can be rewound by one slot + /// at most, so it cannot stand in for the entry a near-prefix save would replace: that + /// would turn a full hit into a miss, not into a few tokens of lost reuse. Wrapped rings + /// keep the exact-prefix rule. Two conversations that share a system prompt and have + /// short tails (a short first user message) must both stay. + func testWrappedRingsSharingASystemPromptBothStay() async { + let tail = 10 + XCTAssertLessThanOrEqual(tail, PromptCache.supersedeSlack, "the tails must be within the slack for this test to mean anything") + let system = seq(0, 200) + let a = system + seq(1000, tail) + let b = system + seq(2000, tail) + let pc = PromptCache(maxEntries: 2) + await pc.save(tokens: a, cache: [makeRing(maxSize: 32, tokens: a.count)]) + await pc.save(tokens: b, cache: [makeRing(maxSize: 32, tokens: b.count)]) + + // A's next turn extends A exactly: a full hit on A, not a miss. + let fresh = RotatingKVCache(maxSize: 32, keep: 0, step: 4) + let n = await pc.restore(newTokens: a + seq(3000, 30), into: [fresh]) + XCTAssertEqual(n, a.count) + XCTAssertEqual(fresh.offset, a.count) + let held = await drain(pc) + XCTAssertEqual(held, 2) + } + + /// Gemma-style cache: full-attention layers beside sliding-window ones. One wrapped ring + /// is enough, whichever layer it is. + func testMixedCacheWithAWrappedRingKeepsTheExactPrefixRule() async { + func mixed(_ n: Int) -> [any KVCache] { + makeCache(seqLen: n, fill: 1) + [makeRing(maxSize: 32, tokens: n)] + } + let system = seq(0, 200) + let a = system + seq(1000, 10) + let b = system + seq(2000, 10) + let pc = PromptCache(maxEntries: 2) + await pc.save(tokens: a, cache: mixed(a.count)) + await pc.save(tokens: b, cache: mixed(b.count)) + + let n = await pc.restore(newTokens: a + seq(3000, 30), + into: [KVCacheSimple(), RotatingKVCache(maxSize: 32, keep: 0, step: 4)]) + XCTAssertEqual(n, a.count) + let held = await drain(pc) + XCTAssertEqual(held, 2) + } + + /// The exact-prefix rule still applies to a wrapped ring: a turn that extends it + /// replaces it. + func testWrappedRingIsStillReplacedByAnExactExtension() async { + let turn1 = seq(0, 80) + let pc = PromptCache(maxEntries: 2) + await pc.save(tokens: turn1, cache: [makeRing(maxSize: 32, tokens: 80)]) + await pc.save(tokens: turn1 + seq(500, 20), cache: [makeRing(maxSize: 32, tokens: 100)]) + let held = await drain(pc) + XCTAssertEqual(held, 1) + } + + /// A ring that has not dropped tokens rewinds like any attention cache, so the + /// near-prefix rule applies to it (the guard is the wrapped state, not the cache type). + func testUnwrappedRingsReplaceANearPrefix() async { + let old = seq(0, 30) + let new = seq(0, 25) + seq(900, 5) + let pc = PromptCache(maxEntries: 2) + await pc.save(tokens: old, cache: [makeRing(maxSize: 64, tokens: old.count)]) + await pc.save(tokens: new, cache: [makeRing(maxSize: 64, tokens: new.count)]) + + let fresh = RotatingKVCache(maxSize: 64, keep: 0, step: 4) + let n = await pc.restore(newTokens: old + [7], into: [fresh]) + XCTAssertEqual(n, 25, "old was replaced; the replacement still restores the shared prefix") + XCTAssertEqual(fresh.offset, 25) + let held = await drain(pc) + XCTAssertEqual(held, 1) + } + + /// How far a wrapped ring may be rewound on restore: the one slot a full-match replay + /// overwrites, nothing deeper. + func testWrappedRingRewindBoundaries() async { + let saved = seq(0, 80) + func restoreInto(_ newTokens: [Int]) async -> Int? { + let pc = PromptCache(maxEntries: 1) + await pc.save(tokens: saved, cache: [makeRing(maxSize: 32, tokens: 80)]) + return await pc.restore(newTokens: newTokens, + into: [RotatingKVCache(maxSize: 32, keep: 0, step: 4)]) + } + let extended = await restoreInto(saved + [1]) + let replayed = await restoreInto(saved) // excess 0, replays one token + let oneBack = await restoreInto(seq(0, 79) + [500]) // excess 1 + let oneBackReplayed = await restoreInto(seq(0, 79)) // excess 1 plus the replay + let twoBack = await restoreInto(seq(0, 78) + [500]) // excess 2 + XCTAssertEqual(extended, 80) + XCTAssertEqual(replayed, 80) + XCTAssertEqual(oneBack, 79) + XCTAssertNil(oneBackReplayed) + XCTAssertNil(twoBack) + } + + /// A ring that is exactly full has evicted nothing, so it may be rewound as deep as a + /// plain cache. + func testFullButNotWrappedRingRewindsDeeply() async { + let pc = PromptCache(maxEntries: 1) + await pc.save(tokens: seq(0, 32), cache: [makeRing(maxSize: 32, tokens: 32)]) + let fresh = RotatingKVCache(maxSize: 32, keep: 0, step: 4) + let n = await pc.restore(newTokens: seq(0, 20) + [500], into: [fresh]) + XCTAssertEqual(n, 20) + XCTAssertEqual(fresh.offset, 20) + } + + /// The cached state must hold more tokens than are trimmed off it: an excess equal to the + /// length is rejected, one less is not. The 4-token state under a 10-token prompt stands + /// in for a sliding-window layer that kept fewer tokens than the prompt. + func testExcessEqualToTheCachedLengthIsAMissOneLessIsAHit() async { + func restoreInto(_ newTokens: [Int]) async -> Int? { + let pc = PromptCache(maxEntries: 1) + await pc.save(tokens: seq(0, 10), cache: makeCache(seqLen: 4, fill: 1)) + return await pc.restore(newTokens: newTokens, into: [KVCacheSimple()]) + } + let equal = await restoreInto(seq(0, 6) + [500]) // excess 4 == 4 cached + let oneLess = await restoreInto(seq(0, 7) + [500]) // excess 3 < 4 + XCTAssertNil(equal) + XCTAssertEqual(oneLess, 7) + } + + // MARK: - Choosing between several entries + + /// The longest usable match wins, not the most recent entry. + func testRestorePicksTheLongestMatchRegardlessOfRecency() async { + let slack = PromptCache.supersedeSlack + let long = seq(0, 3 * slack) + let short = seq(0, slack) // a prefix of `long`, saved after it, and the most recent + let pc = PromptCache(maxEntries: 2) + await put(pc, long, fill: 1) + await put(pc, short, fill: 2) + + let fresh = KVCacheSimple() + let n = await pc.restore(newTokens: long + [999], into: [fresh]) + XCTAssertEqual(n, long.count) + XCTAssertEqual(fresh.offset, long.count) + XCTAssertEqual(keySum(fresh), 8 * Float(long.count) * 1, "restored from the long entry's KV") + } + + func testEqualMatchesPreferTheMostRecentEntry() async { + let slack = PromptCache.supersedeSlack + let common = seq(0, 2 * slack) + let a = common + seq(1000, slack + 4) + let b = common + seq(2000, slack + 4) + let pc = PromptCache(maxEntries: 2) + await put(pc, a, fill: 1) + await put(pc, b, fill: 2) + + let fresh = KVCacheSimple() + let n = await pc.restore(newTokens: common + seq(3000, 5), into: [fresh]) + XCTAssertEqual(n, common.count) + XCTAssertEqual(keySum(fresh), 8 * Float(common.count) * 2, "tie goes to the most recent entry (b)") + } + + /// A diverging tail restores the shared prefix and trims the rest. + func testTailDivergenceRestoresTheSharedPrefixAndTrimsTheRest() async { + let pc = PromptCache(maxEntries: 2) + await put(pc, seq(0, 30)) + let fresh = KVCacheSimple() + let n = await pc.restore(newTokens: seq(0, 27) + [900, 901, 902], into: [fresh]) + XCTAssertEqual(n, 27) + XCTAssertEqual(fresh.offset, 27) + } + + /// A candidate whose restore would corrupt a layer is skipped, and the next one is used, + /// even when the rejected one matched longer and is the most recent. + func testWrappedRingCandidateIsSkippedForAUsableEntry() async { + let pc = PromptCache(maxEntries: 2) + // Unwrapped ring (offset 30 <= maxSize 32): any trim is exact. + await pc.save(tokens: seq(0, 10) + seq(700, 20), cache: [makeRing(maxSize: 32, tokens: 30)]) + // Wrapped ring (offset 80 > maxSize 32): rewinding it by 4 would point the window at + // evicted keys. Saved last, so it is the first candidate, and the longer match. + await pc.save(tokens: seq(0, 80), cache: [makeRing(maxSize: 32, tokens: 80)]) + + let fresh = RotatingKVCache(maxSize: 32, keep: 0, step: 4) + let n = await pc.restore(newTokens: seq(0, 76) + [500, 501], into: [fresh]) + XCTAssertEqual(n, 10) + XCTAssertEqual(fresh.offset, 10) + let stats = await pc.stats() + XCTAssertEqual(stats.hits, 1) + XCTAssertEqual(stats.misses, 0) + } + + func testEveryCandidateRejectedIsAMiss() async { + let pc = PromptCache(maxEntries: 2) + await pc.save(tokens: seq(0, 80), cache: [makeRing(maxSize: 32, tokens: 80)]) + await pc.save(tokens: seq(1000, 80), cache: [makeRing(maxSize: 32, tokens: 80)]) + let fresh = RotatingKVCache(maxSize: 32, keep: 0, step: 4) + let n = await pc.restore(newTokens: seq(0, 76) + [500, 501], into: [fresh]) + XCTAssertNil(n) + XCTAssertEqual(fresh.offset, 0, "the live cache is untouched") + let stats = await pc.stats() + XCTAssertEqual(stats.misses, 1) + } + + /// An entry whose trim would empty a layer is skipped. The 4-token state under a + /// 10-token prompt stands in for a sliding-window layer that kept fewer tokens than + /// the prompt it was saved with. + func testEntryThatWouldBeTrimmedToNothingIsSkippedForAUsableEntry() async { + let pc = PromptCache(maxEntries: 2) + await put(pc, seq(0, 2) + seq(100, 20), fill: 2) // usable + await pc.save(tokens: seq(0, 10), cache: makeCache(seqLen: 4, fill: 1)) // excess 8 >= 4 + + let fresh = KVCacheSimple() + let n = await pc.restore(newTokens: seq(0, 2) + seq(300, 5), into: [fresh]) + XCTAssertEqual(n, 2) + XCTAssertEqual(fresh.offset, 2) + XCTAssertEqual(keySum(fresh), 8 * 2 * 2, "restored from the usable entry, not the rejected one") + } + + /// Hybrid restore takes the longest exact prefix within the limit, not the most recent. + func testHybridExactPrefixPicksTheLongestEntryNotTheMostRecent() async { + let long = seq(0, 50) + let short = seq(0, 20) + let pc = PromptCache(maxEntries: 2) + await pc.save(tokens: long, cache: makeHybrid(seqLen: 50, fill: 1), allowRecurrent: true) + await pc.save(tokens: short, cache: makeHybrid(seqLen: 20, fill: 2), allowRecurrent: true) + + let fresh: [any KVCache] = [KVCacheSimple(), MambaCache()] + let n = await pc.restoreExactPrefix(newTokens: long + [99], limit: 50, into: fresh) + XCTAssertEqual(n, 50) + XCTAssertEqual(fresh[1].state[1].sum().item(Float.self), 1 * 32, "the long snapshot's recurrent state") + + let capped: [any KVCache] = [KVCacheSimple(), MambaCache()] + let m = await pc.restoreExactPrefix(newTokens: long + [99], limit: 30, into: capped) + XCTAssertEqual(m, 20, "the long snapshot is past the limit, so the short one is used") + XCTAssertEqual(capped[1].state[1].sum().item(Float.self), 2 * 32) + } + + /// A hybrid hit refreshes recency too, so the next save evicts the other snapshot. + func testHybridExactPrefixHitRefreshesRecency() async { + let h1 = seq(0, 10) + let h2 = seq(100, 10) + let h3 = seq(200, 10) + let pc = PromptCache(maxEntries: 2) + await pc.save(tokens: h1, cache: makeHybrid(seqLen: 10, fill: 1), allowRecurrent: true) + await pc.save(tokens: h2, cache: makeHybrid(seqLen: 10, fill: 2), allowRecurrent: true) + let touched = await pc.restoreExactPrefix(newTokens: h1 + [9], limit: 10, + into: [KVCacheSimple(), MambaCache()]) + XCTAssertEqual(touched, 10) + await pc.save(tokens: h3, cache: makeHybrid(seqLen: 10, fill: 3), allowRecurrent: true) // evicts h2 + + let kept = await pc.restoreExactPrefix(newTokens: h1 + [9], limit: 10, + into: [KVCacheSimple(), MambaCache()]) + let dropped = await pc.restoreExactPrefix(newTokens: h2 + [9], limit: 10, + into: [KVCacheSimple(), MambaCache()]) + XCTAssertEqual(kept, 10) + XCTAssertNil(dropped) + } + + // MARK: - Eviction + + func testEvictReportsHowManyEntriesItDropped() async { + let pc = PromptCache(maxEntries: 3) + let emptyKeep = await pc.evict(keepMostRecent: true) + let emptyAll = await pc.evict(keepMostRecent: false) + XCTAssertEqual(emptyKeep, 0) + XCTAssertEqual(emptyAll, 0) + + await put(pc, [1, 1, 1, 1]) + await put(pc, [2, 2, 2, 2]) + await put(pc, [3, 3, 3, 3]) + let keepRecent = await pc.evict(keepMostRecent: true) + let again = await pc.evict(keepMostRecent: true) + let last = await pc.evict(keepMostRecent: false) + let none = await pc.evict(keepMostRecent: false) + XCTAssertEqual(keepRecent, 2) + XCTAssertEqual(again, 0, "nothing left to drop but the most recent") + XCTAssertEqual(last, 1) + XCTAssertEqual(none, 0) + } + + func testCapacityBelowOneIsRaisedToOne() async { + for requested in [0, -1] { + let pc = PromptCache(maxEntries: requested) + await put(pc, [1, 2, 3, 4]) + await put(pc, [9, 8, 7, 6]) + let old = await hit(pc, [1, 2, 3, 4, 5]) + let recent = await hit(pc, [9, 8, 7, 6, 5]) + XCTAssertNil(old, "maxEntries \(requested) still holds one entry, not none or many") + XCTAssertEqual(recent, 4) + } + } + + func testFlagParses() throws { + XCTAssertEqual(try MLXServer.parse(["--model", "m"]).promptCacheEntries, 1) + XCTAssertEqual(try MLXServer.parse(["--model", "m", "--prompt-cache-entries", "4"]).promptCacheEntries, 4) + } +}