Skip to content

prompt cache: extension links instead of a whole-store scan per candidate #1363

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
jschmied wants to merge 1 commit into from
Closed

prompt cache: extension links instead of a whole-store scan per candidate #1363

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 13 additions & 11 deletions zig/src/core/prompt_cache.zig
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 };
Expand Down Expand Up @@ -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;

Expand All @@ -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);
Expand Down Expand Up @@ -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;
}

Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down
49 changes: 34 additions & 15 deletions zig/src/core/prompt_cache_policy.zig
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
};
}
126 changes: 126 additions & 0 deletions zig/src/core/prompt_cache_scan_test.zig
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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);
}
1 change: 1 addition & 0 deletions zig/src/core/rows_read.zig
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Loading