From 75218c46e5e31097a57f877048cf215b02a91df0 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski <24755+wojo@users.noreply.github.com> Date: Sun, 25 Oct 2026 13:59:59 -0400 Subject: [PATCH 1/2] core kv_pool: with the engine's own sequence on pages, the windows the room holds come before the cache kv_pool.plan funded one window of pages, a page a stream, then the prompt cache, then the rest. Where the room holds more than one window, that gave the cache rows that a window a stream would have given the streams: at TP=2, a 128k window and four streams, 260 pages against four windows. Now, with engine_paged, the pages of as many windows as the room would hold as windows come before the cache, so the pool never holds fewer rows than windows in the same room. holds() counts the same way. Nemotron (engine_paged false) is unchanged. The engine's-own test keeps its checks at a room under two windows; a new test covers TP=2 at 128k and 1M and TP=3 at 1M. --- zig/src/core/kv_pool.zig | 19 +++++++++++--- zig/src/core/kv_pool_test.zig | 45 ++++++++++++++++++++++++++++----- zig/src/native/cuda_kv_pool.zig | 8 +++--- 3 files changed, 58 insertions(+), 14 deletions(-) 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/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) }; From dbd073dc227453e9d87817f6f89eef10848ec4dc Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski <24755+wojo@users.noreply.github.com> Date: Sun, 25 Oct 2026 14:05:49 -0400 Subject: [PATCH 2/2] glm/cuda: TF_KV_PAGES=auto pages from two ranks With the plan funding the windows the room holds before the cache, pages at two ranks hold at least the rows a window a stream would, and serve --parallel 8 at a 1M window where windows are refused. One GPU and --kv-cache fp8 keep their windows. --- zig/src/families/glm/cuda_native.zig | 10 ++++------ zig/src/families/glm/cuda_ranks.zig | 2 +- zig/src/families/glm/cuda_start.zig | 2 +- 3 files changed, 6 insertions(+), 8 deletions(-) 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;