glm/cuda: KV pages by default from two ranks, and the plan funds the windows the room holds before the cache

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).

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" {

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));
     }

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",

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;

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) };