diff --git a/zig/src/core/kv_pool.zig b/zig/src/core/kv_pool.zig index cc439d93e9..55569b2310 100644 --- a/zig/src/core/kv_pool.zig +++ b/zig/src/core/kv_pool.zig @@ -54,7 +54,7 @@ pub const Budget = struct { snaps: ?usize = null, /// Rows one stream holds at most (0: no bound): pages past what the paged streams hold at once stay unmade. window: u64 = 0, - /// The engine's own rows are pages too: a window of them comes first, so one stream can still reach it. + /// The engine's own rows are pages too: a window of them first, then the cache after the windows the room holds. engine_paged: bool = false, }; @@ -78,7 +78,7 @@ pub fn plan(l: Layout, b: Budget) error{ NoStream, TooMany }!Plan { const fit = @min(extra, left / each); if (fit < extra and b.fixed) return error.TooMany; left -= fit * l.fixed_bytes; - const c = split(@min(b.cache, left - fit * (1 + l.runPages()) * l.pageBytes()), l.snap_bytes, b.snaps); + const c = split(@min(b.cache, left -| held(l, b, fit)), l.snap_bytes, b.snaps); left -= c.snaps * l.snap_bytes + c.copied; const most: u64 = if (b.window > 0) (fit + @intFromBool(b.engine_paged)) * l.pages(b.window) + fit * l.runPages() else std.math.maxInt(u64); return .{ .lanes = @intCast(1 + fit), .snaps = c.snaps, .copied = c.copied, .pages = @intCast(@min((left + own) / l.pageBytes(), most)) }; @@ -90,7 +90,20 @@ pub fn holds(l: Layout, b: Budget) bool { if (b.engine + own > b.room or b.lanes == 0) return false; const left = b.room - b.engine - own; const each = l.fixed_bytes + (1 + l.runPages()) * l.pageBytes(); - return left - @min(@as(u64, b.lanes - 1), left / each) * each >= b.cache; + const fit = @min(@as(u64, b.lanes - 1), left / each); + return left -| (fit * l.fixed_bytes + held(l, b, fit)) >= b.cache; +} + +/// The pages the cache leaves the other streams: a page and a run each, or the windows the room would hold if more. +fn held(l: Layout, b: Budget, fit: u64) u64 { + return @max(fit * (1 + l.runPages()), (windows(l, b) - 1) * l.pages(b.window)) * l.pageBytes(); +} + +/// Streams a window each would serve in this room, the engine's own among them (1 unless its rows are pages). +fn windows(l: Layout, b: Budget) u64 { + if (!b.engine_paged or b.window == 0) return 1; + const own = ownPages(l, b); + return @min(b.lanes, 1 + (b.room - own) / (own + l.fixed_bytes)); } /// The bytes of the engine's own window as pages, set aside first when its rows are pages (0 when it has a window). diff --git a/zig/src/core/kv_pool_test.zig b/zig/src/core/kv_pool_test.zig index 20b60dd873..bd79db531e 100644 --- a/zig/src/core/kv_pool_test.zig +++ b/zig/src/core/kv_pool_test.zig @@ -73,10 +73,11 @@ test "the engine's own sequence on pages: a window of them comes first, then the const window: u64 = 1_048_576 + 8; const floor = glm.pages(window) * glm.pageBytes(); try expectEqual(@as(usize, 2049), glm.pages(window)); // one request at --context - const p = try kv.plan(glm, .{ .room = 40 * g, .engine = 0, .lanes = 4, .cache = 16 * g, .window = window, .engine_paged = true }); + // a room under two windows' pages: only the engine's own window is set aside before the cache + const p = try kv.plan(glm, .{ .room = 24 * g, .engine = 0, .lanes = 4, .cache = 8 * g, .window = window, .engine_paged = true }); try expectEqual(@as(u32, 4), p.lanes); - try expectEqual(16 * g, p.copied); - try expectEqual(@as(usize, @intCast((24 * g - 3 * 180 * mib) / glm.pageBytes())), p.pages); + try expectEqual(8 * g, p.copied); + try expectEqual(@as(usize, @intCast((16 * g - 3 * 180 * mib) / glm.pageBytes())), p.pages); // the window's pages before the cache: the cache gets what they and a page a stream leave const each = 180 * mib + glm.pageBytes(); const tight = try kv.plan(glm, .{ .room = floor + 3 * each + g, .engine = 0, .lanes = 4, .cache = 16 * g, .window = window, .engine_paged = true }); @@ -92,11 +93,41 @@ test "the engine's own sequence on pages: a window of them comes first, then the const bounded = try kv.plan(glm, .{ .room = 100 * g, .engine = 0, .lanes = 2, .window = 8192 + 8, .engine_paged = true }); try expectEqual(@as(usize, 2 * 17), bounded.pages); // what holds is what plan gives the cache whole, past the window's pages - const room = floor + 3 * each + 16 * g; - try expect(kv.holds(glm, .{ .room = room, .engine = 0, .lanes = 4, .cache = 16 * g, .window = window, .engine_paged = true })); - try expect(!kv.holds(glm, .{ .room = room - 1, .engine = 0, .lanes = 4, .cache = 16 * g, .window = window, .engine_paged = true })); + const room = floor + 3 * each + 8 * g; + try expect(kv.holds(glm, .{ .room = room, .engine = 0, .lanes = 4, .cache = 8 * g, .window = window, .engine_paged = true })); + try expect(!kv.holds(glm, .{ .room = room - 1, .engine = 0, .lanes = 4, .cache = 8 * g, .window = window, .engine_paged = true })); try expect(!kv.holds(glm, .{ .room = floor - 1, .engine = 0, .lanes = 1, .window = window, .engine_paged = true })); - try expectEqual(16 * g, (try kv.plan(glm, .{ .room = room, .engine = 0, .lanes = 4, .cache = 16 * g, .window = window, .engine_paged = true })).copied); + try expectEqual(8 * g, (try kv.plan(glm, .{ .room = room, .engine = 0, .lanes = 4, .cache = 8 * g, .window = window, .engine_paged = true })).copied); +} + +test "the engine's own sequence on pages: the windows the room would hold come before the cache" { + const g: u64 = 1 << 30; + // GLM-5.3-Flash on CUDA at TP=2 (180 MiB fixed) and TP=3 (131 MiB); the cache asked is the default 16 GiB + const tp2: kv.Layout = .{ .fixed_bytes = 180 * mib, .token_bytes = 13_056, .page_tokens = 512, .snap_bytes = 0 }; + const tp3: kv.Layout = .{ .fixed_bytes = 131 * mib, .token_bytes = 13_056, .page_tokens = 512, .snap_bytes = 0 }; + const page = tp2.pageBytes(); + // a 128k window, four streams: the room holds four windows, so the pool holds their rows and the cache the rest + const short: u64 = 131_072 + 8; + const room2 = 9 * g + g / 2; + const p = try kv.plan(tp2, .{ .room = room2, .engine = 0, .lanes = 4, .cache = 16 * g, .window = short, .engine_paged = true }); + try expectEqual(@as(u32, 4), p.lanes); + try expectEqual(@as(usize, 4 * 257), p.pages); + try expectEqual(room2 - 4 * 257 * page - 3 * 180 * mib, p.copied); + try expect(kv.holds(tp2, .{ .room = room2, .engine = 0, .lanes = 4, .cache = p.copied, .window = short, .engine_paged = true })); + try expect(!kv.holds(tp2, .{ .room = room2, .engine = 0, .lanes = 4, .cache = p.copied + 1, .window = short, .engine_paged = true })); + // a 1M window, eight streams, a room for one window: as before, the window's pages, a page a stream, the cache + const long: u64 = 1_048_576 + 8; + const each = 180 * mib + page; + const one = try kv.plan(tp2, .{ .room = 2049 * page + 7 * each + g, .engine = 0, .lanes = 8, .cache = 16 * g, .window = long, .engine_paged = true }); + try expectEqual(@as(u32, 8), one.lanes); + try expectEqual(@as(usize, 2049 + 7), one.pages); + try expectEqual(g, one.copied); + // TP=3, a 1M window, eight streams: the room holds two windows, so two windows' pages come before the cache + const room3 = 38 * g; + const t = try kv.plan(tp3, .{ .room = room3, .engine = 0, .lanes = 8, .cache = 16 * g, .window = long, .engine_paged = true }); + try expectEqual(@as(u32, 8), t.lanes); + try expectEqual(@as(usize, 2 * 2049), t.pages); + try expectEqual(room3 - 2 * 2049 * page - 7 * 131 * mib, t.copied); } test "a cache holds when the room past the engine and every stream's fixed part and page is at least its bytes" { diff --git a/zig/src/families/glm/cuda_native.zig b/zig/src/families/glm/cuda_native.zig index 99822f26c2..4a6c2f9b44 100644 --- a/zig/src/families/glm/cuda_native.zig +++ b/zig/src/families/glm/cuda_native.zig @@ -433,12 +433,10 @@ test "a checkpoint refused for its format, at the entry or by the loader, is exp for ([_]anyerror{ error.NotModelOpt, error.UnsupportedQuant, error.UnreadableFormat }) |e| try std.testing.expect(explain(null, e) != null); } -test "TF_KV_PAGES: auto pages from three ranks, off and 0 give windows, all and other values page like on" { - for ([_]usize{ 1, 2 }) |w| { - try std.testing.expect(!try paging(null, w, false)); - try std.testing.expect(!try paging("auto", w, false)); - } - for ([_]usize{ 3, 16 }) |w| { +test "TF_KV_PAGES: auto pages from two ranks, off and 0 give windows, all and other values page like on" { + try std.testing.expect(!try paging(null, 1, false)); + try std.testing.expect(!try paging("auto", 1, false)); + for ([_]usize{ 2, 3, 16 }) |w| { try std.testing.expect(try paging(null, w, false)); try std.testing.expect(try paging("auto", w, false)); } diff --git a/zig/src/families/glm/cuda_ranks.zig b/zig/src/families/glm/cuda_ranks.zig index 7300deb65a..e6c4517beb 100644 --- a/zig/src/families/glm/cuda_ranks.zig +++ b/zig/src/families/glm/cuda_ranks.zig @@ -304,7 +304,7 @@ fn explain(f: Fingerprint.Field) []const u8 { .rank => "a rank number given twice or out of order", .world => "the ranks were started with different world sizes", .kv8 => "--kv-cache fp8", - .pages => "TF_KV_PAGES (unset or auto pages from three ranks)", + .pages => "TF_KV_PAGES (unset or auto pages from two ranks)", .prompt_fixed_k => "TF_GLM_PROMPT_K", .tiled => "TF_GLM_MM", .split => "TF_GLM_PREFILL_REDUCE", diff --git a/zig/src/families/glm/cuda_start.zig b/zig/src/families/glm/cuda_start.zig index 4d237182fe..cf5fb99b81 100644 --- a/zig/src/families/glm/cuda_start.zig +++ b/zig/src/families/glm/cuda_start.zig @@ -19,7 +19,7 @@ fn resolved(p: anytype, text: ?[]const u8, initial: anytype, args: anytype, comp /// TF_KV_PAGES keeps FP8 on whole windows; explicit pages with FP8 are refused on every rank. pub fn paging(text: ?[]const u8, world: usize, kv8: bool) !bool { const v = text orelse "auto"; - if (std.mem.eql(u8, v, "auto")) return world >= 3 and !kv8; + if (std.mem.eql(u8, v, "auto")) return world >= 2 and !kv8; if (std.mem.eql(u8, v, "0") or std.mem.eql(u8, v, "off")) return false; if (kv8) return error.Fp8OnPages; return true; diff --git a/zig/src/native/cuda_kv_pool.zig b/zig/src/native/cuda_kv_pool.zig index db1d9c5301..4237edf07c 100644 --- a/zig/src/native/cuda_kv_pool.zig +++ b/zig/src/native/cuda_kv_pool.zig @@ -216,11 +216,11 @@ test "kv pages: an engine whose own rows are pages plans a window of them first, try std.testing.expect(std.mem.indexOf(u8, problem, "8192-token window") != null); const shown = (try text(Owned, a, 180 << 20, l, p, true, 8192, .{ .snaps = 0, .copied = 4 << 30, .snap_bytes = 0 }, false))[0]; try std.testing.expectEqualStrings("the engine's own and 3 more taking 68 pages of 512 rows (0.42 GiB) as they need them (enough for one at a 8192-token window), 180 MiB fixed each", shown); - // the asked cache holds past the window's pages and a fixed part and page for each other stream, not a byte less - const each = (180 << 20) + l.pageBytes(); + // the asked cache holds past the four windows' pages this room holds and a fixed part each, not a byte less + const room = 4 * floor + 3 * (180 << 20) + (4 << 30); const loaded = .{ .layout = @as(?kv_pool.Layout, l), .engine_paged = true, .free_streams = @as(u32, 1) }; - try std.testing.expect(fits(Owned, loaded, o, floor + 3 * each + (4 << 30), 8192, 0)); - try std.testing.expect(!fits(Owned, loaded, o, floor + 3 * each + (4 << 30) - 1, 8192, 0)); + try std.testing.expect(fits(Owned, loaded, o, room, 8192, 0)); + try std.testing.expect(!fits(Owned, loaded, o, room - 1, 8192, 0)); } test "kv merge: host armed-state staging is reserved once before planning the page pool" { const loaded = .{ .host_snaps = true, .snap_bytes = @as(u64, 8 << 20) };