Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
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
2 changes: 2 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel

- The prompt cache can let a family count repeated image or keyed prompts as repeats, under the cache's own key. No family uses it yet (#1929, part of #1927). Thanks to @laudney.

- The prompt cache can keep a bounded ladder of intermediate checkpoints under a long prompt's history mark, so a later, compacted prompt resumes at the longest checkpoint inside its retained prefix instead of the shared cut, and the checkpoints stay across append turns. Host tests cover the plan, admission, retention and eviction; the Flash Next on CUDA opt-in (`TF_PROMPT_INTERMEDIATE=1`) is not yet measured on a GPU (#1909).

- HIP Qwen3.5-family servers price shared draft rounds at the row bucket they actually run, so Qwen3.6-35B-A3B with MTP at four streams no longer drops into a slower drafting mode, and models load faster. Tokens are unchanged (#1830, fixes #1695). Thanks to @AbrahamCete.

- The Living Weights view reads a shard's header with core's safetensors reader, so a malformed header, length word or offset is refused with an error instead of a panic or a 64 GiB allocation, and its sums and row counts are checked (#1923, part of #1872). Thanks to @jschmied.
Expand Down
15 changes: 11 additions & 4 deletions zig/src/core/prompt_cache.zig
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ const lifetime = @import("prompt_cache_lifetime.zig");
const policy = @import("prompt_cache_policy.zig");
const cuts = @import("prompt_cuts.zig");
const interior = @import("prompt_interior.zig");
const intermediate = @import("prompt_intermediate.zig");
const index = @import("prompt_cache_index.zig");
const forks = @import("prompt_cache_fork.zig");
pub const rows = @import("rows_store.zig");
Expand Down Expand Up @@ -92,6 +93,7 @@ pub const Entry = struct {
last: []u32, // the last prompt that kept or resumed it
shared: bool = false, // kept at a shared system block's cut: other conversations resume it, so its own never supersedes it
branched: bool = false, // resumed again after a later prompt kept a longer state extending it: never superseded
intermediate: bool = false, // a ladder checkpoint under the history: kept across append turns, 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
Expand Down Expand Up @@ -296,13 +298,17 @@ pub const Store = struct {
a.free(ordinary);
return err;
};
const marked = intermediate.add(s, a, prompt, from, history_len, starts, spans, marks_, endpoint, endpoint_bytes, e) catch |err| {
a.free(marks_);
return err;
};
for (shared) |w| { // the shared cuts this pass keeps, and a short prompt's mark inside a shared block: other conversations resume them
const at = if (s.rules.planned) s.startAt(starts, w) else w;
for (marks_) |m| if (m == at or (m < at and !s.rules.planned and prompt.len < s.rules.min_prompt)) s.noteShared(prompt[0 .. m + s.rules.lookahead]);
for (marked) |m| if (m == at or (m < at and !s.rules.planned and prompt.len < s.rules.min_prompt)) s.noteShared(prompt[0 .. m + s.rules.lookahead]);
}
s.reserve(prompt, born, e, marks_, spans);
s.reserve(prompt, born, e, marked, spans);
if (e) |entry| entry.borrowers += 1;
return .{ .entry = e, .marks = marks_ };
return .{ .entry = e, .marks = marked };
}

/// Room for every state the pass keeps, made before it: a peer that mirrors the keeps hears of these evictions with the request.
Expand Down Expand Up @@ -474,7 +480,7 @@ pub const Store = struct {
return s.fail(at, err);
};
const charged = if (s.family.vtable.charged) |f| f(s.family.ptr, saved) else bytes;
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 };
e.* = .{ .tokens = last[0..n], .decode_spans = decoded, .at = at, .saved = saved, .bytes = charged, .born = born, .used = s.clock, .last = last, .shared = shared, .intermediate = intermediate.is(s.rules, prompt.len, at), .own = last };
owns_decoded = false;
index.add(s, e) catch {
s.free(e);
Expand Down Expand Up @@ -557,6 +563,7 @@ test {
_ = @import("prompt_cache_learn_test.zig");
_ = @import("prompt_cache_rewind_test.zig");
_ = @import("prompt_cache_interior_test.zig");
_ = @import("prompt_cache_intermediate_test.zig");
_ = @import("prompt_cache_branch_test.zig");
_ = @import("prompt_cache_fork_test.zig");
_ = @import("prompt_cache_give_test.zig");
Expand Down
149 changes: 149 additions & 0 deletions zig/src/core/prompt_cache_intermediate_test.zig
Original file line number Diff line number Diff line change
@@ -0,0 +1,149 @@
//! The bounded ladder of intermediate checkpoints: planned, admitted beside the endpoint, resumed when compacted.
const std = @import("std");
const pc = @import("prompt_cache.zig");
const im = @import("prompt_intermediate.zig");
const Fake = @import("prompt_cache_fake.zig").Fake;
const fresh = @import("prompt_cache_test.zig").fresh;

const ladder: pc.Rules = .{ .chunk = 2048, .intermediate = 8192, .intermediate_most = 4 };

test "ladder takes the largest powers-of-two positions under the history" {
const gpa = std.testing.allocator;
const cases = [_]struct { rules: pc.Rules, prompt: usize, history: u32, from: u32, marks: []const u32, want: []const u32 }{
.{ .rules = ladder, .prompt = 160 * 1024, .history = 160 * 1024 - 8, .from = 0, .marks = &.{}, .want = &.{ 32768, 65536, 131072 } },
.{ .rules = ladder, .prompt = 40 * 1024, .history = 40 * 1024, .from = 0, .marks = &.{}, .want = &.{32768} },
// An existing mark and the resume point drop their positions; a short prompt or the default rules plan none.
.{ .rules = ladder, .prompt = 40 * 1024, .history = 40 * 1024, .from = 0, .marks = &.{32768}, .want = &.{} },
.{ .rules = ladder, .prompt = 40 * 1024, .history = 40 * 1024, .from = 32768, .marks = &.{}, .want = &.{} },
.{ .rules = ladder, .prompt = 40 * 1024, .history = 40 * 1024, .from = 0, .marks = &.{40000}, .want = &.{32768} },
.{ .rules = ladder, .prompt = 8192, .history = 8192, .from = 0, .marks = &.{}, .want = &.{} },
.{ .rules = .{ .chunk = 2048 }, .prompt = 160 * 1024, .history = 160 * 1024, .from = 0, .marks = &.{}, .want = &.{} },
};
for (cases) |c| {
const got = try im.cuts(gpa, c.rules, c.prompt, c.history, c.from, c.marks);
defer gpa.free(got);
try std.testing.expectEqualSlices(u32, c.want, got);
}
}

test "a compacted long history resumes its longest ladder checkpoint exactly" {
const gpa = std.testing.allocator;
var arena: std.heap.ArenaAllocator = .init(gpa);
defer arena.deinit();
const a = arena.allocator();
const prompt = try a.alloc(u32, 160 * 1024);
for (prompt, 0..) |*t, i| t.* = @intCast(i % 65521);
var f: Fake = .{ .gpa = gpa };
var s = pc.Store.init(gpa, f.snapshots(), ladder, 1 << 20);
defer s.deinit();
const history: u32 = @intCast(prompt.len - 8);
const p = try s.beginUser(a, prompt, history, 0, 0, &.{4096}, &.{}, null, &.{});
try std.testing.expectEqualSlices(u32, &.{ 4096, 32768, 65536, 131072, history }, p.marks);
try std.testing.expectEqual(fresh(prompt), f.pass(&s, prompt, p));
const compacted = try a.dupe(u32, prompt[0 .. 120 * 1024]);
for (compacted[100 * 1024 ..]) |*t| t.* = @intCast((t.* + 1) % 65521);
const q = try s.beginUser(a, compacted, @intCast(compacted.len - 8), 0, 0, &.{4096}, &.{}, null, &.{});
try std.testing.expectEqual(@as(u32, 65536), q.from);
try std.testing.expectEqual(fresh(compacted), f.pass(&s, compacted, q));
try std.testing.expectEqual(@as(u64, 1), s.counts.hits);
}

test "without the ladder a compacted long history resumes only its shared cut" {
const gpa = std.testing.allocator;
var arena: std.heap.ArenaAllocator = .init(gpa);
defer arena.deinit();
const a = arena.allocator();
const prompt = try a.alloc(u32, 160 * 1024);
for (prompt, 0..) |*t, i| t.* = @intCast(i % 65521);
var f: Fake = .{ .gpa = gpa };
var s = pc.Store.init(gpa, f.snapshots(), .{ .chunk = 2048 }, 1 << 20);
defer s.deinit();
const history: u32 = @intCast(prompt.len - 8);
const p = try s.beginUser(a, prompt, history, 0, 0, &.{4096}, &.{}, null, &.{});
try std.testing.expectEqual(fresh(prompt), f.pass(&s, prompt, p));
const compacted = try a.dupe(u32, prompt[0 .. 120 * 1024]);
for (compacted[100 * 1024 ..]) |*t| t.* = @intCast((t.* + 1) % 65521);
const q = try s.beginUser(a, compacted, @intCast(compacted.len - 8), 0, 0, &.{4096}, &.{}, null, &.{});
try std.testing.expectEqual(@as(u32, 4096), q.from);
}

test "ladder checkpoints stay across an append turn and are evictable" {
const gpa = std.testing.allocator;
var arena: std.heap.ArenaAllocator = .init(gpa);
defer arena.deinit();
const a = arena.allocator();
const prompt = try a.alloc(u32, 160 * 1024);
for (prompt, 0..) |*t, i| t.* = @intCast(i % 65521);
var f: Fake = .{ .gpa = gpa };
var s = pc.Store.init(gpa, f.snapshots(), ladder, 1 << 20);
defer s.deinit();
const p = try s.beginUser(a, prompt, @intCast(prompt.len - 8), 0, 0, &.{}, &.{}, null, &.{});
try std.testing.expectEqual(fresh(prompt), f.pass(&s, prompt, p));
const rung = for (s.entries.items) |e| {
if (e.at == 65536) break e;
} else null;
try std.testing.expect(rung != null and rung.?.intermediate);
// The next turn extends the prompt; the checkpoint is not superseded by the later state.
const grown = try a.alloc(u32, 200 * 1024);
@memcpy(grown[0..prompt.len], prompt);
for (grown[prompt.len..], 0..) |*t, i| t.* = @intCast(i % 7);
const q = try s.beginUser(a, grown, @intCast(grown.len - 8), 0, 0, &.{}, &.{}, null, &.{});
try std.testing.expectEqual(@as(u32, @intCast(prompt.len - 8)), q.from);
try std.testing.expectEqual(fresh(grown), f.pass(&s, grown, q));
const compacted = try a.dupe(u32, grown[0 .. 120 * 1024]);
for (compacted[100 * 1024 ..]) |*t| t.* = @intCast((t.* + 1) % 7);
const r = try s.beginUser(a, compacted, @intCast(compacted.len - 8), 0, 0, &.{}, &.{}, null, &.{});
try std.testing.expectEqual(@as(u32, 65536), r.from);
try std.testing.expectEqual(fresh(compacted), f.pass(&s, compacted, r));
// Under pressure the same checkpoint is no immortal state: evicting picks it once better victims are gone.
var picked = false;
while (s.entries.items.len > 0) {
const victim = s.victimBut(null, grown, @intCast(grown.len)) orelse break;
picked = picked or s.entries.items[victim].intermediate;
s.remove(victim);
}
try std.testing.expect(picked);
}

test "ladder admission keeps the endpoint ahead of the checkpoints" {
const gpa = std.testing.allocator;
var arena: std.heap.ArenaAllocator = .init(gpa);
defer arena.deinit();
const a = arena.allocator();
const prompt = try a.alloc(u32, 160 * 1024);
@memset(prompt, 5);
const history: u32 = @intCast(prompt.len - 8);
const endpoint: u64 = 100 + history;
for ([_]u64{ 0, endpoint - 1, endpoint, endpoint + 100 + 32768, 1 << 20 }) |budget| for ([_]usize{ 1, 2, 3, 6 }) |slots| {
var f: Fake = .{ .gpa = gpa };
var s = pc.Store.init(gpa, f.snapshots(), ladder, budget);
defer s.deinit();
s.most = slots;
const p = try s.beginUser(a, prompt, history, 0, 0, &.{}, &.{}, null, &.{});
const rungs = for (p.marks) |m| {
if (im.is(ladder, prompt.len, m)) break true;
} else false;
try std.testing.expectEqual(budget >= endpoint and slots >= 1, std.mem.indexOfScalar(u32, p.marks, history) != null);
try std.testing.expectEqual(budget >= endpoint + 100 + 32768 and slots >= 2, rungs);
try std.testing.expectEqual(fresh(prompt), f.pass(&s, prompt, p));
try std.testing.expect(s.held + s.spare() <= budget);
try std.testing.expect(s.entries.items.len <= slots);
};
}

test "a short prompt or a cache-off store plans no ladder" {
const gpa = std.testing.allocator;
var arena: std.heap.ArenaAllocator = .init(gpa);
defer arena.deinit();
const a = arena.allocator();
const prompt = try a.alloc(u32, 8000);
@memset(prompt, 7);
for ([_]pc.Rules{ ladder, .{ .chunk = 2048 } }) |rules| {
var f: Fake = .{ .gpa = gpa };
var s = pc.Store.init(gpa, f.snapshots(), rules, 1 << 20);
defer s.deinit();
const p = try s.beginUser(a, prompt, 7992, 16, 7992, &.{}, &.{}, null, &.{});
for (p.marks) |m| try std.testing.expect(!im.is(rules, prompt.len, m));
try std.testing.expectEqual(fresh(prompt), f.pass(&s, prompt, p));
}
}
2 changes: 1 addition & 1 deletion zig/src/core/prompt_cache_policy.zig
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ pub fn victimBut(s: *const pc.Store, skip: ?*pc.Entry, prompt: []const u32, born
oldest = i;
oldest_rewind = rewind;
}
if (e.borrowers > 0 or e.shared or e.branched or e == skip or rewind) continue;
if (e.borrowers > 0 or e.shared or e.branched or e.intermediate 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;
}
Expand Down
109 changes: 109 additions & 0 deletions zig/src/core/prompt_intermediate.zig
Original file line number Diff line number Diff line change
@@ -0,0 +1,109 @@
//! A bounded ladder of intermediate checkpoints under the history: a later, compacted prompt resumes inside its prefix.
const std = @import("std");
const pc = @import("prompt_cache.zig");
const modes = @import("cache_modes.zig");
const Allocator = std.mem.Allocator;

/// Whether `at` belongs to the ladder: a power-of-two multiple of its spacing, under the prompt's history.
pub fn is(r: pc.Rules, prompt_len: usize, at: u32) bool {
if (r.intermediate == 0 or r.intermediate_most == 0 or prompt_len < r.intermediate_min or at < r.intermediate_min) return false;
if (at >= prompt_len or at % r.intermediate != 0) return false;
return std.math.isPowerOfTwo(at / r.intermediate);
}

/// The largest `intermediate_most` ladder positions below `history`, ascending, away from `from` and `marks`.
pub fn cuts(a: Allocator, r: pc.Rules, prompt_len: usize, history: u32, from: u32, marks: []const u32) ![]u32 {
if (r.intermediate == 0 or r.intermediate_most == 0 or prompt_len < r.intermediate_min or history == 0) return a.alloc(u32, 0);
var all: std.ArrayList(u32) = .empty;
defer all.deinit(a);
var at = r.intermediate;
while (at < history) {
const outside = at < r.intermediate_min or at <= from or at + r.lookahead > prompt_len or history - at < r.min_gap;
var skip = outside;
if (!skip) skip = std.mem.indexOfScalar(u32, marks, at) != null;
if (!skip) skip = for (marks) |m| {
if (@max(m, at) - @min(m, at) < r.min_gap) break true;
} else false;
if (!skip) try all.append(a, at);
if (at > history / 2) break;
at *= 2;
}
if (all.items.len <= @as(usize, r.intermediate_most)) return all.toOwnedSlice(a);
const first = all.items.len - r.intermediate_most;
const out = try a.dupe(u32, all.items[first..]);
return out;
}

/// Append the ladder's marks that fit beside the endpoint, rewind and existing marks; a new slice in `a`.
pub fn add(
s: *pc.Store,
a: Allocator,
prompt: []const u32,
from: u32,
history: u32,
starts: []const u32,
spans: []const modes.Span,
marks: []const u32,
endpoint: u32,
endpoint_bytes: u64,
restored: ?*pc.Entry,
) ![]const u32 {
// The endpoint keeps its priority: no ladder beside a state that does not fit.
if (endpoint_bytes > 0 and endpoint_bytes > s.budget) return marks;
const candidates = try cuts(a, s.rules, prompt.len, history, from, marks);
defer a.free(candidates);
if (candidates.len == 0) return marks;
var required: std.ArrayList(u32) = .empty;
defer required.deinit(a);
try required.appendSlice(a, marks);
for ([_]u32{ if (endpoint_bytes > 0 and endpoint_bytes <= s.budget) endpoint else 0, s.active_rewind_at }) |p| {
if (p > 0 and std.mem.indexOfScalar(u32, required.items, p) == null) try required.append(a, p);
}
var used: u64 = 0;
var slots: usize = required.items.len;
for (required.items) |p| {
var bytes = s.family.vtable.bytes(s.family.ptr, p);
for (s.entries.items) |e| if (same(e, prompt, p, s.rules.lookahead, spans)) {
bytes = @max(bytes, e.bytes);
break;
};
used +|= bytes;
}
// Other in-flight restores cannot be evicted. The current restore has not been borrowed by lookup yet.
for (s.entries.items) |e| if (e.borrowers > 0 or e == restored) {
const included = for (required.items) |p| {
if (same(e, prompt, p, s.rules.lookahead, spans)) break true;
} else false;
if (!included) {
used +|= e.bytes;
slots +|= 1;
}
};
// A family without trim owns its spare storage until it reuses it; conservative admission leaves it charged.
if (s.family.vtable.trim == null) used +|= s.spare();
var added: usize = 0;
var out: std.ArrayList(u32) = .empty;
defer out.deinit(a);
try out.appendSlice(a, marks);
for (candidates) |at| {
if (added == @as(usize, s.rules.intermediate_most) or slots >= s.most) break;
if (!s.usable(at, starts)) continue;
const bytes = s.family.vtable.bytes(s.family.ptr, at);
if (used > s.budget or bytes > s.budget - used) continue;
try out.append(a, at);
used +|= bytes;
slots +|= 1;
added += 1;
}
if (added == 0) return marks;
std.mem.sort(u32, out.items, {}, std.sort.asc(u32));
const kept = try out.toOwnedSlice(a);
a.free(marks);
return kept;
}

fn same(e: *const pc.Entry, prompt: []const u32, at: u32, lookahead: u32, spans: []const modes.Span) bool {
return !e.expired and e.at == at and e.tokens.len == @as(usize, at) + lookahead and
e.tokens.len <= prompt.len and std.mem.eql(u32, e.tokens, prompt[0..e.tokens.len]) and
modes.equal(e.decode_spans, spans, at);
}
6 changes: 6 additions & 0 deletions zig/src/core/prompt_rules.zig
Original file line number Diff line number Diff line change
Expand Up @@ -28,4 +28,10 @@ pub const Rules = struct {
repeat_keyed: bool = false,
/// Keep one before where the prompt parts from the one it resumed (a miss: the longest kept one), however near.
fork: bool = false,
/// Spacing of a bounded ladder of checkpoints under the history (0: off; only chunk-invariant families opt in).
intermediate: u32 = 0,
/// Checkpoints kept from the ladder at most, the largest positions first.
intermediate_most: u8 = 0,
/// Prompts shorter than this, and ladder positions before it, keep none.
intermediate_min: u32 = 32768,
};
3 changes: 3 additions & 0 deletions zig/src/families/flashnext/cuda_lanes.zig
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,9 @@ pub fn Namespace(comptime Tri: type) type {
/// The prompt cache's copies of a stream's state (cuda_snap.zig), called with `Loaded.target`.
pub const snap_bytes = snaps.bytes;
pub const snap_repeat_min: u32 = 16;
/// The prompt cache's ladder of intermediate checkpoints under a long prompt's history mark.
pub const snap_intermediate: u32 = 8192;
pub const snap_intermediate_most: u8 = 4;
pub const snap_save = snaps.save;
pub const snap_restore = snaps.restore;
pub const snap_drop = snaps.drop;
Expand Down
Loading
Loading