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;