From 74be69cd3b106fbfb69272e53297646de4526ac9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sun, 25 Oct 2026 17:00:19 +0100 Subject: [PATCH] prompt cache: extension links instead of a whole-store scan per candidate victimBut asked extended() of every candidate, and extended() compared the candidate's tokens with every other entry's: O(E^2 L) a victim, and eviction loops pick several. Each entry now keeps the entries that extend it (longer tokens with them as a prefix, born later), linked once when an entry joins the store and unlinked when it leaves; tokens and born never change, and expired is read when asked, so extended() answers as the scan did. victimBut also compares each entry's tokens with the prompt once, for both of its passes, instead of again in every pairwise order of the oldest scan. If a link allocation fails, extended() falls back to the scan. keepBorn copies the prompt once: the state's tokens are a prefix of the kept `last`, which a later resume replaces as before. prompt_cache_scan_test.zig checks every victim and extended() against the old scans on random many-entry stores whose prompts share long prefixes. --- zig/src/core/prompt_cache.zig | 24 ++--- zig/src/core/prompt_cache_policy.zig | 49 ++++++--- zig/src/core/prompt_cache_scan_test.zig | 126 ++++++++++++++++++++++++ zig/src/core/rows_read.zig | 1 + 4 files changed, 174 insertions(+), 26 deletions(-) create mode 100644 zig/src/core/prompt_cache_scan_test.zig diff --git a/zig/src/core/prompt_cache.zig b/zig/src/core/prompt_cache.zig index d99bd21c3c..8cee1671d0 100644 --- a/zig/src/core/prompt_cache.zig +++ b/zig/src/core/prompt_cache.zig @@ -115,6 +115,8 @@ pub const Entry = struct { branched: bool = false, // resumed again after a later prompt kept a longer state extending it: never superseded expired: bool = false, // hidden from lookups immediately, reclaimed only at safe host boundaries borrowers: u32 = 0, // outstanding lookups: released by resumed or release, never evicted while borrowed + ext: std.ArrayList(*Entry) = .empty, // entries extending its tokens, born later (policy.link): `extended` reads them + own: []u32 = &.{}, // the kept prompt's one copy, which `tokens` and at first `last` slice (empty: they own theirs) }; pub const Counts = struct { hits: u64 = 0, misses: u64 = 0, kept: u64 = 0, evicted: u64 = 0, refused: u64 = 0, failed: u64 = 0 }; @@ -143,6 +145,7 @@ pub const Store = struct { imprint: ?*imprint.Imprint = null, // --learn: shared cuts' states on disk, read back by later sessions and servers disk: ?*rows.Rows = null, // --learn-conversations: every kept state written behind to disk, read back where it pays active_rewind_at: u32 = 0, // the last lookup's rewind state (0: none); hosts keep a pass's states before the next lookup + links_lost: bool = false, // an allocation failure left the extension links incomplete: `extended` scans const SHARED_KEYS = 64; @@ -160,15 +163,17 @@ pub const Store = struct { fn free(s: *Store, e: *Entry) void { if (s.disk) |r| rows.write.dropped(r, e); s.family.vtable.drop(s.family.ptr, e.saved); - s.gpa.free(e.tokens); + s.gpa.free(if (e.own.len > 0) e.own else e.tokens); s.gpa.free(e.decode_spans); - s.gpa.free(e.last); + if (e.own.len == 0 or e.last.ptr != e.own.ptr) s.gpa.free(e.last); + e.ext.deinit(s.gpa); s.gpa.destroy(e); } pub fn remove(s: *Store, i: usize) void { std.debug.assert(s.entries.items[i].borrowers == 0); const e = s.entries.orderedRemove(i); + policy.unlink(s, e); s.held -= e.bytes; if (s.disk) |r| if (!e.expired) rows.write.evicted(r, e); // --spill-gib: an evicted state is written before it goes s.free(e); @@ -396,7 +401,7 @@ pub const Store = struct { e.branched = e.branched or s.extended(e); // a turn several continuations share (a re-sent, edited or forked one) s.counts.hits += 1; const last = s.gpa.dupe(u32, prompt) catch return; - s.gpa.free(e.last); + if (e.own.len == 0 or e.last.ptr != e.own.ptr) s.gpa.free(e.last); // the first `last` is the tokens' copy e.last = last; } @@ -479,28 +484,23 @@ pub const Store = struct { } if (s.entries.items.len == 0) if (s.family.vtable.trim) |f| f(s.family.ptr, s.budget -| bytes); const e = s.gpa.create(Entry) catch return s.fail(at, error.OutOfMemory); - const tokens = s.gpa.dupe(u32, prompt[0..n]) catch { - s.gpa.destroy(e); - return s.fail(at, error.OutOfMemory); - }; - const last = s.gpa.dupe(u32, prompt) catch { - s.gpa.free(tokens); + const last = s.gpa.dupe(u32, prompt) catch { // one copy: the state's tokens are its prefix s.gpa.destroy(e); return s.fail(at, error.OutOfMemory); }; const saved = s.family.vtable.save(s.family.ptr, owner, at) catch |err| { s.gpa.free(last); - s.gpa.free(tokens); s.gpa.destroy(e); return s.fail(at, err); }; const charged = if (s.family.vtable.charged) |f| f(s.family.ptr, saved) else bytes; - e.* = .{ .tokens = tokens, .decode_spans = decoded, .at = at, .saved = saved, .bytes = charged, .born = born, .used = s.clock, .last = last, .shared = shared }; + e.* = .{ .tokens = last[0..n], .decode_spans = decoded, .at = at, .saved = saved, .bytes = charged, .born = born, .used = s.clock, .last = last, .shared = shared, .own = last }; owns_decoded = false; s.entries.append(s.gpa, e) catch { s.free(e); return s.fail(at, error.OutOfMemory); }; + policy.link(s, e); s.held += charged; s.counts.kept += 1; if (shared) learned.learn(s, e, starts); @@ -533,6 +533,7 @@ pub const Store = struct { }; im.touch(key); s.entries.append(s.gpa, e) catch return s.drop(e, have, true); + policy.link(s, e); s.held += bytes; if (!@import("builtin").is_test) std.log.info("prompt cache: read {d} learned tokens from disk", .{m.at}); return e; @@ -578,6 +579,7 @@ test { _ = @import("prompt_cache_branch_test.zig"); _ = @import("prompt_cache_give_test.zig"); _ = @import("prompt_cache_lifetime_test.zig"); + _ = @import("prompt_cache_scan_test.zig"); _ = @import("learned_fault_test.zig"); _ = modes; _ = rows; diff --git a/zig/src/core/prompt_cache_policy.zig b/zig/src/core/prompt_cache_policy.zig index fb07e975bc..745b64ba6a 100644 --- a/zig/src/core/prompt_cache_policy.zig +++ b/zig/src/core/prompt_cache_policy.zig @@ -4,32 +4,51 @@ const pc = @import("prompt_cache.zig"); pub fn victimBut(s: *const pc.Store, skip: ?*pc.Entry, prompt: []const u32, born: u32) ?usize { var best: ?usize = null; - for (s.entries.items, 0..) |e, i| { - if (e.borrowers > 0 or e.shared or e.branched or e == skip or isActiveRewind(s, e, prompt)) continue; - const now = born > e.born and prompt.len > e.tokens.len and std.mem.eql(u32, prompt[0..e.tokens.len], e.tokens); // a longer request's tokens extend it + var oldest: ?usize = null; + var oldest_rewind = false; + for (s.entries.items, 0..) |e, i| { // one prefix compare an entry: the superseded scan and the oldest scan share it + const on = e.tokens.len <= prompt.len and std.mem.eql(u32, prompt[0..e.tokens.len], e.tokens); + const rewind = !e.expired and e.at == s.active_rewind_at and on; // the active rewind state goes last + if (e.borrowers == 0 and e != skip and (oldest == null or (if (rewind != oldest_rewind) oldest_rewind else e.used < s.entries.items[oldest.?].used))) { + oldest = i; + oldest_rewind = rewind; + } + if (e.borrowers > 0 or e.shared or e.branched or e == skip or rewind) continue; + const now = born > e.born and prompt.len > e.tokens.len and on; // a longer request's tokens extend it if ((now or extended(s, e)) and (best == null or e.used < s.entries.items[best.?].used)) best = i; } - if (best) |i| return i; - var oldest: ?usize = null; - for (s.entries.items, 0..) |e, i| if (e.borrowers == 0 and e != skip and (oldest == null or before(s, e, s.entries.items[oldest.?], prompt))) { - oldest = i; - }; - return oldest; + return best orelse oldest; } /// Only live continuations make a prefix a branch point or a superseded state. pub fn extended(s: *const pc.Store, e: *const pc.Entry) bool { + if (s.links_lost) return extendedScan(s, e); + for (e.ext.items) |o| if (!o.expired) return true; + return false; +} + +/// `extended` over every entry, for a store whose links an allocation failure left incomplete. +pub fn extendedScan(s: *const pc.Store, e: *const pc.Entry) bool { for (s.entries.items) |o| if (!o.expired and o != e and o.born > e.born and o.tokens.len > e.tokens.len and std.mem.eql(u32, o.tokens[0..e.tokens.len], e.tokens)) return true; return false; } -fn before(s: *const pc.Store, e: *const pc.Entry, o: *const pc.Entry, prompt: []const u32) bool { - const e_rewind = isActiveRewind(s, e, prompt); - const o_rewind = isActiveRewind(s, o, prompt); - return if (e_rewind != o_rewind) o_rewind else e.used < o.used; +/// A new entry linked to those it extends and those extending it, later born (tokens and born never change). +pub fn link(s: *pc.Store, x: *pc.Entry) void { + for (s.entries.items) |o| { + if (o == x or o.born == x.born or o.tokens.len == x.tokens.len) continue; + const short, const long = if (o.tokens.len < x.tokens.len) .{ o, x } else .{ x, o }; + if (long.born < short.born or !std.mem.eql(u32, long.tokens[0..short.tokens.len], short.tokens)) continue; + short.ext.append(s.gpa, long) catch { + s.links_lost = true; + }; + } } -fn isActiveRewind(s: *const pc.Store, e: *const pc.Entry, prompt: []const u32) bool { - return !e.expired and e.at == s.active_rewind_at and e.tokens.len <= prompt.len and std.mem.eql(u32, e.tokens, prompt[0..e.tokens.len]); +/// An entry leaving the store leaves every link to it. +pub fn unlink(s: *pc.Store, x: *pc.Entry) void { + for (s.entries.items) |o| if (std.mem.indexOfScalar(*pc.Entry, o.ext.items, x)) |i| { + _ = o.ext.swapRemove(i); + }; } diff --git a/zig/src/core/prompt_cache_scan_test.zig b/zig/src/core/prompt_cache_scan_test.zig new file mode 100644 index 0000000000..9cf560142e --- /dev/null +++ b/zig/src/core/prompt_cache_scan_test.zig @@ -0,0 +1,126 @@ +//! Eviction choices with the extension links against the whole-store scans they replace, on random long-prefix stores. +const std = @import("std"); +const pc = @import("prompt_cache.zig"); +const policy = @import("prompt_cache_policy.zig"); +const Fake = @import("prompt_cache_fake.zig").Fake; +const gpa = std.testing.allocator; + +/// The victim as the scans chose it: every entry's prefix compared again in each pass, and in each pairwise order. +fn victimScan(s: *const pc.Store, skip: ?*pc.Entry, prompt: []const u32, born: u32) ?usize { + var best: ?usize = null; + for (s.entries.items, 0..) |e, i| { + if (e.borrowers > 0 or e.shared or e.branched or e == skip or rewound(s, e, prompt)) continue; + const now = born > e.born and prompt.len > e.tokens.len and std.mem.eql(u32, prompt[0..e.tokens.len], e.tokens); + if ((now or policy.extendedScan(s, e)) and (best == null or e.used < s.entries.items[best.?].used)) best = i; + } + if (best) |i| return i; + var oldest: ?usize = null; + for (s.entries.items, 0..) |e, i| if (e.borrowers == 0 and e != skip and (oldest == null or before(s, e, s.entries.items[oldest.?], prompt))) { + oldest = i; + }; + return oldest; +} + +fn before(s: *const pc.Store, e: *const pc.Entry, o: *const pc.Entry, prompt: []const u32) bool { + const e_rewind = rewound(s, e, prompt); + const o_rewind = rewound(s, o, prompt); + return if (e_rewind != o_rewind) o_rewind else e.used < o.used; +} + +fn rewound(s: *const pc.Store, e: *const pc.Entry, prompt: []const u32) bool { + return !e.expired and e.at == s.active_rewind_at and e.tokens.len <= prompt.len and std.mem.eql(u32, e.tokens, prompt[0..e.tokens.len]); +} + +/// A conversation tree: each prompt a stored one's prefix plus new tokens, so entries share long prefixes. +const Prompts = struct { + all: std.ArrayList([]u32) = .empty, + fn deinit(p: *Prompts) void { + for (p.all.items) |x| gpa.free(x); + p.all.deinit(gpa); + } + fn next(p: *Prompts, r: std.Random) ![]const u32 { + const from: []const u32 = if (p.all.items.len == 0 or r.uintLessThan(u8, 8) == 0) &.{} else blk: { + const x = p.all.items[r.uintLessThan(usize, p.all.items.len)]; + break :blk x[0 .. x.len - r.uintLessThan(usize, @min(x.len, 40))]; + }; + const out = try gpa.alloc(u32, from.len + 8 + r.uintLessThan(usize, 120)); + @memcpy(out[0..from.len], from); + for (out[from.len..]) |*t| t.* = r.uintLessThan(u32, 3); // a small vocabulary: prompts often agree by chance + try p.all.append(gpa, out); + return out; + } +}; + +fn agree(s: *pc.Store, r: std.Random, prompts: *Prompts) !void { + for (0..@min(8, s.entries.items.len)) |_| { + const e = s.entries.items[r.uintLessThan(usize, s.entries.items.len)]; + try std.testing.expectEqual(policy.extendedScan(s, e), policy.extended(s, e)); + } + for (0..4) |_| { + const prompt = prompts.all.items[r.uintLessThan(usize, prompts.all.items.len)]; + const skip: ?*pc.Entry = if (s.entries.items.len > 0 and r.boolean()) s.entries.items[r.uintLessThan(usize, s.entries.items.len)] else null; + const born: u32 = r.uintLessThan(u32, @intCast(prompt.len + 2)); + const saved = s.active_rewind_at; + if (s.entries.items.len > 0 and r.boolean()) s.active_rewind_at = s.entries.items[r.uintLessThan(usize, s.entries.items.len)].at; + try std.testing.expectEqual(victimScan(s, skip, prompt, born), s.victimBut(skip, prompt, born)); + s.active_rewind_at = saved; + } +} + +fn run(seed: u64, most: usize) !u64 { + var prng = std.Random.DefaultPrng.init(seed); + const r = prng.random(); + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + var f: Fake = .{ .gpa = gpa }; + var s = pc.Store.init(gpa, f.snapshots(), .{ .lookahead = 1, .min_prompt = 0, .min_gap = 4 }, 1 << 40); + s.most = most; + defer s.deinit(); + var prompts: Prompts = .{}; + defer prompts.deinit(); + for (0..400) |_| { + const prompt = try prompts.next(r); + switch (r.uintLessThan(u8, 6)) { + 0, 1, 2 => { // a pass keeps a few states, some born as a shorter request (a reply's end) + for (0..1 + r.uintLessThan(usize, 3)) |_| { + const at: u32 = 1 + r.uintLessThan(u32, @intCast(prompt.len - 1)); + f.at = at; + _ = s.keepBorn(prompt, if (r.boolean()) @intCast(prompt.len) else at + 1, at, null, &.{}, &.{}); + } + }, + 3, 4 => { // a lookup, resumed or failed + const history: u32 = r.uintLessThan(u32, @intCast(prompt.len)); + const l = try s.lookupRewind(arena.allocator(), prompt, @intCast(prompt.len), history, r.uintLessThan(u32, @intCast(prompt.len)), &.{}, &.{}, &.{}); + if (l.entry) |e| s.resumed(e, prompt, r.uintLessThan(u8, 5) != 0); + }, + else => if (s.entries.items.len > 0) { // a state the family drops + s.expired(s.entries.items[r.uintLessThan(usize, s.entries.items.len)].saved); + if (r.boolean()) s.collectExpired(); + }, + } + try agree(&s, r, &prompts); + } + return s.counts.evicted; +} + +test "eviction with extension links picks every victim the whole-store scans picked" { + var evicted: u64 = 0; + for (0..24) |seed| evicted += try run(seed, if (seed % 4 == 3) 200 else 24); + try std.testing.expect(evicted > 1000); // the stores were full: the choices were exercised +} + +test "a state's tokens are its kept prompt's prefix, held once" { + var f: Fake = .{ .gpa = gpa }; + var s = pc.Store.init(gpa, f.snapshots(), .{ .lookahead = 1, .min_prompt = 0 }, 1 << 20); + defer s.deinit(); + const prompt = [_]u32{ 1, 2, 3, 4, 5, 6 }; + f.at = 3; + try std.testing.expect(s.keep(&prompt, 3, null, &.{}, &.{})); + const e = s.entries.items[0]; + try std.testing.expectEqualSlices(u32, prompt[0..4], e.tokens); + try std.testing.expectEqualSlices(u32, &prompt, e.last); + try std.testing.expect(e.tokens.ptr == e.last.ptr); + s.resumed(e, &.{ 1, 2, 3, 4, 9 }, true); // a later prompt takes `last`'s place; the tokens keep their copy + try std.testing.expectEqualSlices(u32, prompt[0..4], e.tokens); + try std.testing.expectEqualSlices(u32, &.{ 1, 2, 3, 4, 9 }, e.last); +} diff --git a/zig/src/core/rows_read.zig b/zig/src/core/rows_read.zig index aa45720681..0c05a23ae0 100644 --- a/zig/src/core/rows_read.zig +++ b/zig/src/core/rows_read.zig @@ -127,6 +127,7 @@ pub fn recall(s: *pc.Store, prompt: []const u32, born: u32, starts: []const u32, s.family.vtable.drop(s.family.ptr, e.saved); return have; }; + @import("prompt_cache_policy.zig").link(s, e); ok = true; s.held += e.bytes; return e;