From 5b72fe76e1f3de2e182dbd848f9234410c6975bb Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski <24755+wojo@users.noreply.github.com> Date: Sun, 25 Oct 2026 14:02:49 -0400 Subject: [PATCH 1/7] GLM on CUDA: a kept state's disk image on one rank, its parts and where each byte sits in its lease The rows store's view of a GLM kept state, without the store: the words (format, position, head rows, ring bases, KDA copy bits), the KDA prefix, each DSA layer's latent rows, scales and whole pools by position, its rings and pools in progress as the state's own, the head's caches at a whole verify window (rows past the head's own read as zeros), then the tail. Copies run a page of the lease at a time. Nothing is wired to --learn-conversations yet. --- zig/src/families/glm/cuda.zig | 1 + zig/src/families/glm/cuda_disk_image.zig | 189 ++++++++++++++++++ zig/src/families/glm/cuda_disk_image_test.zig | 184 +++++++++++++++++ 3 files changed, 374 insertions(+) create mode 100644 zig/src/families/glm/cuda_disk_image.zig create mode 100644 zig/src/families/glm/cuda_disk_image_test.zig diff --git a/zig/src/families/glm/cuda.zig b/zig/src/families/glm/cuda.zig index bf1f50e8c8..5ea204949e 100644 --- a/zig/src/families/glm/cuda.zig +++ b/zig/src/families/glm/cuda.zig @@ -47,6 +47,7 @@ test { _ = @import("cuda_images_interleave_test.zig"); _ = @import("cuda_pages_test.zig"); _ = @import("cuda_ranks_test.zig"); + _ = @import("cuda_disk_image.zig"); _ = experts; _ = exl3; _ = vision; diff --git a/zig/src/families/glm/cuda_disk_image.zig b/zig/src/families/glm/cuda_disk_image.zig new file mode 100644 index 0000000000..071d300661 --- /dev/null +++ b/zig/src/families/glm/cuda_disk_image.zig @@ -0,0 +1,189 @@ +//! A GLM kept state's disk image on one rank: its parts at a position, and where each image byte sits in its lease. +const std = @import("std"); +const cuda = @import("cuda"); +const st_mod = @import("cuda_state.zig"); +const sparse = @import("cuda_sparse.zig"); +const Config = @import("config.zig").Config; +const pages = @import("cuda_snapshot_pages.zig"); +const State = st_mod.State; +const window_rows = st_mod.window_rows; + +/// The image's first bytes: the format, the position, the head's rows, both ring bases and the KDA copy bits. +pub const Words = struct { + at: u64, + mtp_len: u64, + base: u64, + mtp_base: u64, + cur: u64, + + pub const bytes: u64 = 48; + const format: u64 = 1; + + pub fn encode(w: Words) [6]u64 { + return .{ format, w.at, w.mtp_len, w.base, w.mtp_base, w.cur }; + } + + /// Words read back: refused unless the format is this one and the state stands at `at`. + pub fn decode(x: [6]u64, at: u32) !Words { + if (x[0] != format or x[1] != at) return error.DiskWords; + return .{ .at = x[1], .mtp_len = x[2], .base = x[3], .mtp_base = x[4], .cur = x[5] }; + } +}; + +/// A part as the rows store takes it: `row` bytes for each `per` tokens from position 0, or the state's own (row 0). +pub const Part = struct { bytes: u64, row: u32 = 0, per: u32 = 1 }; + +/// `len` image bytes from `img`: the words, `lease` bytes into the state's lease, or zeros the lease does not hold. +pub const Seg = struct { img: u64, len: u64, from: union(enum) { words, lease: u64, zeros } }; + +/// What a rank's kept state holds, read from a sequence of the engine (its addresses are not used past here). +pub const Shape = struct { + prefix: u64, // the sequence's own fields before its caches: KDA conv and states, a window's saved rows + tail: u64, // its fields after them: the head's position and first-draft row + layers: u32, // DSA layers, whose rows are shared by position + latent: u32, // a latent row's bytes + scales: bool, // fp8 latent rows: a 4-byte scale each + ring_rows: u32, + ring_row: u32, + pk_rows: u32, + pooled_row: u32, + capacity: u32, + head: bool, // the MTP head's caches are kept (drafting) + + pub fn of(st: *const State, c: *const Config, capacity: usize, head: bool) !Shape { + if (st.rings.len == 0) return error.SnapshotShape; + const base = st.arena.buf.ptr; + const first = if (st.table != 0) st.rings[0] else if (st.caches.len > 0) st.caches[0].lc else st.mtp_caches.lc; + const end = if (st.table != 0) st.table else base + st.arena.buf.len; + if (first < base or st.mtp_pos_dev < first or end < st.mtp_pos_dev) return error.SnapshotShape; + const fp8 = st.mtp_caches.ls != 0; + return .{ + .prefix = first - base, + .tail = end - st.mtp_pos_dev, + .layers = @intCast(st.caches.len), + .latent = c.kv_lora * @as(u32, if (fp8) 1 else 2), + .scales = fp8, + .ring_rows = @intCast(st.ring_rows), + .ring_row = @intCast(st.ring_row), + .pk_rows = @intCast(capacity / sparse.pool + 2), + .pooled_row = c.i_dim * 2, + .capacity = @intCast(capacity), + .head = head, + }; + } + + /// Rows the image keeps of the head's caches at `at`: a verify window past it, as the cache charges a state. + pub fn headRows(sh: Shape, at: u32) u32 { + return @min(sh.capacity, at + window_rows); + } +}; + +/// One cache's parts in the lease's order, for `n` rows: latent, scales, the two index rings, the pooled keys. +fn cache(sh: Shape, n: u64) [5]u64 { + return .{ + n * sh.latent, + if (sh.scales) n * 4 else 0, + @min(n, sh.ring_rows) * sh.ring_row, + @min(n, sh.ring_rows) * sh.ring_row, + @min(n / sparse.pool + 2, sh.pk_rows) * sh.pooled_row, + }; +} + +/// The image at `at` (a chunk start): the words, the prefix, each DSA layer's rows, the head's caches, the tail. +pub fn parts(sh: Shape, at: u32, out: *std.ArrayList(Part), a: std.mem.Allocator) !void { + if (at == 0 or at % sparse.pool != 0 or at > sh.capacity) return error.SnapshotShape; + try out.appendSlice(a, &.{ .{ .bytes = Words.bytes }, .{ .bytes = sh.prefix } }); + const pooled: u64 = at / sparse.pool; + for (0..sh.layers) |_| { + const c = cache(sh, at); + try out.append(a, .{ .bytes = c[0], .row = sh.latent }); + if (sh.scales) try out.append(a, .{ .bytes = c[1], .row = 4 }); + try out.appendSlice(a, &.{ .{ .bytes = c[2] }, .{ .bytes = c[3] } }); + // whole pools are rows by position; the pool past them and the spare row stay the state's own + try out.appendSlice(a, &.{ .{ .bytes = pooled * sh.pooled_row, .row = sh.pooled_row, .per = sparse.pool }, .{ .bytes = c[4] - pooled * sh.pooled_row } }); + } + if (sh.head) for (cache(sh, sh.headRows(at))) |n| if (n > 0) try out.append(a, .{ .bytes = n }); + try out.append(a, .{ .bytes = sh.tail }); +} + +/// The image's segments in order for a state at `at` whose head holds `mtp_len` rows (0 without a head). +pub fn segments(sh: Shape, at: u32, mtp_len: u64, out: *std.ArrayList(Seg), a: std.mem.Allocator) !void { + if (sh.head and mtp_len > sh.headRows(at)) return error.SnapshotShape; + var img: u64 = Words.bytes; + var lease: u64 = 0; + try out.append(a, .{ .img = 0, .len = Words.bytes, .from = .words }); + var front: u64 = sh.prefix; // the prefix and every DSA layer's parts sit in the lease as in the image + for (0..sh.layers) |_| for (cache(sh, at)) |n| { + front += n; + }; + try out.append(a, .{ .img = img, .len = front, .from = .{ .lease = 0 } }); + img += front; + lease += front; + if (sh.head) for (cache(sh, sh.headRows(at)), cache(sh, mtp_len)) |full, held| { + if (held > 0) try out.append(a, .{ .img = img, .len = held, .from = .{ .lease = lease } }); + if (full > held) try out.append(a, .{ .img = img + held, .len = full - held, .from = .zeros }); + img += full; + lease += held; + }; + if (sh.tail > 0) try out.append(a, .{ .img = img, .len = sh.tail, .from = .{ .lease = lease } }); +} + +/// The lease bytes of a state at `at` whose head holds `mtp_len` rows (layout.Plan's length for it). +pub fn leaseBytes(sh: Shape, at: u32, mtp_len: u64) u64 { + var n: u64 = sh.prefix + sh.tail; + for (0..sh.layers) |_| for (cache(sh, at)) |x| { + n += x; + }; + if (sh.head) for (cache(sh, mtp_len)) |x| { + n += x; + }; + return n; +} + +/// Copies between `host` and image bytes [off, off + host.len) of a lease: out (lease to host) or in. +pub const Copy = struct { + d: *const cuda.Driver, + pool: *const pages.Pool, + lease: *const pages.Lease, + segs: []const Seg, + words: *[6]u64, + + pub fn run(c: Copy, off: u64, host: []u8, out: bool) !void { + const end = off + host.len; + for (c.segs) |s| { + const lo = @max(off, s.img); + const hi = @min(end, s.img + s.len); + if (lo >= hi) continue; + const h = host[@intCast(lo - off)..][0..@intCast(hi - lo)]; + switch (s.from) { + .words => { + const w = std.mem.asBytes(c.words)[@intCast(lo - s.img)..][0..h.len]; + if (out) @memcpy(h, w) else @memcpy(w, h); + }, + .zeros => if (out) @memset(h, 0), + .lease => |at| try c.leased(at + (lo - s.img), h, out), + } + } + } + + /// Lease bytes from `at`, a page at a time (each page is a run of the snapshot arena). + fn leased(c: Copy, at: u64, host: []u8, out: bool) !void { + const arena = c.pool.arena orelse return error.NoSnapshotRoom; + if (at + host.len > c.lease.ids.items.len * pages.page) return error.SnapshotSize; + var done: usize = 0; + while (done < host.len) { + const x = at + done; + const n = @min(host.len - done, pages.page - x % pages.page); + const dev = arena.ptr + @as(u64, c.lease.ids.items[@intCast(x / pages.page)]) * pages.page + x % pages.page; + const p = host[done..][0..n]; + if (out) { + try c.d.check(c.d.api.cuMemcpyDtoH_v2(p.ptr, dev, n), "cuMemcpyDtoH"); + } else try c.d.check(c.d.api.cuMemcpyHtoD_v2(dev, p.ptr, n), "cuMemcpyHtoD"); + done += n; + } + } +}; + +test { + _ = @import("cuda_disk_image_test.zig"); +} diff --git a/zig/src/families/glm/cuda_disk_image_test.zig b/zig/src/families/glm/cuda_disk_image_test.zig new file mode 100644 index 0000000000..30d1d84b65 --- /dev/null +++ b/zig/src/families/glm/cuda_disk_image_test.zig @@ -0,0 +1,184 @@ +//! GLM's disk image on a driver whose device memory is host memory: sizes against the snapshot plan, a round trip. +const std = @import("std"); +const cuda = @import("cuda"); +const weights = @import("cuda_weights.zig"); +const shard = @import("cuda_shard.zig"); +const image = @import("cuda_disk_image.zig"); +const State = @import("cuda_state.zig").State; +const Plan = @import("cuda_snapshot_layout.zig").Plan; +const pages = @import("cuda_snapshot_pages.zig"); +const expect = std.testing.expect; +const expectEqual = std.testing.expectEqual; +const expectError = std.testing.expectError; + +/// Device addresses are host addresses: allocations from the page allocator, copies and fills as memcpy and memset. +const H = struct { + var held: [16]struct { ptr: u64 = 0, len: usize = 0 } = @splat(.{}); + fn bytes(p: u64, n: usize) []u8 { + return @as([*]u8, @ptrFromInt(p))[0..n]; + } + fn alloc(out: *u64, n: usize) callconv(.c) cuda.abi.Result { + const m = std.heap.page_allocator.alloc(u8, n) catch return 2; + for (&held) |*h| if (h.ptr == 0) { + h.* = .{ .ptr = @intFromPtr(m.ptr), .len = n }; + out.* = h.ptr; + return 0; + }; + std.heap.page_allocator.free(m); + return 2; + } + fn free(p: u64) callconv(.c) cuda.abi.Result { + for (&held) |*h| if (h.ptr == p) { + std.heap.page_allocator.free(bytes(h.ptr, h.len)); + h.* = .{}; + return 0; + }; + return 1; + } + fn set(p: u64, v: u8, n: usize) callconv(.c) cuda.abi.Result { + @memset(bytes(p, n), v); + return 0; + } + fn dtod(dst: u64, src: u64, n: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + @memcpy(bytes(dst, n), bytes(src, n)); + return 0; + } + fn dtoh(dst: ?*anyopaque, src: u64, n: usize) callconv(.c) cuda.abi.Result { + @memcpy(@as([*]u8, @ptrCast(dst.?))[0..n], bytes(src, n)); + return 0; + } + fn htod(dst: u64, src: ?*const anyopaque, n: usize) callconv(.c) cuda.abi.Result { + @memcpy(bytes(dst, n), @as([*]const u8, @ptrCast(src.?))[0..n]); + return 0; + } + fn wait(_: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + return 0; + } + fn driver() cuda.Driver { + var d: cuda.Driver = undefined; + d.api.cuMemAlloc_v2 = alloc; + d.api.cuMemFree_v2 = free; + d.api.cuMemsetD8_v2 = set; + d.api.cuMemcpyDtoDAsync_v2 = dtod; + d.api.cuMemcpyDtoH_v2 = dtoh; + d.api.cuMemcpyHtoD_v2 = htod; + d.api.cuStreamSynchronize = wait; + return d; + } +}; + +const capacity = 6144; // past the index rings' 4,096 rows, so a state at the end holds a full ring + +/// TP=2 rank 0's shapes on KDA, DSA, KDA, DSA layers. +const Model = struct { + layers: [4]weights.Layer, + w: weights.Weights, + + fn init(m: *Model) !void { + for (&m.layers, [_]bool{ false, true, false, true }) |*l, mla| l.mixer = if (mla) .{ .mla = undefined } else .{ .kda = undefined }; + m.w = undefined; + m.w.config = .{}; + m.w.share = try shard.Share.of(&m.w.config, 0, 2); + m.w.layers = &m.layers; + } +}; + +test "GLM disk image: the lease bytes the image counts are the snapshot plan's, at each position and head length" { + var m: Model = undefined; + try m.init(); + var d = H.driver(); + for ([_]bool{ false, true }) |kv8| { + var st = try State.initFormat(std.testing.allocator, &d, &m.w, capacity, kv8); + defer st.deinit(); + for ([_]bool{ false, true }) |head| { + const sh = try image.Shape.of(&st, &m.w.config, capacity, head); + for ([_]u32{ 2048, 4096, 6144 }) |at| { + for ([_]u64{ at - 1, at, sh.headRows(at) }) |mtp| { + const plan = try Plan.init(&st, &m.w.config, at, if (head) mtp else null); + try expectEqual(@as(u64, plan.len), image.leaseBytes(sh, at, if (head) mtp else 0)); + } + } + } + } +} + +test "GLM disk image: the store's rows are by position from 0 in whole blocks; rings, pools in progress and the head are its own" { + var m: Model = undefined; + try m.init(); + var d = H.driver(); + var st = try State.initFormat(std.testing.allocator, &d, &m.w, capacity, true); + defer st.deinit(); + const sh = try image.Shape.of(&st, &m.w.config, capacity, true); + const a = std.testing.allocator; + var list: std.ArrayList(image.Part) = .empty; + defer list.deinit(a); + var segs: std.ArrayList(image.Seg) = .empty; + defer segs.deinit(a); + const at: u32 = 6144; + try image.parts(sh, at, &list, a); + try image.segments(sh, at, at, &segs, a); + var total: u64 = 0; + var rows: usize = 0; + for (list.items) |p| { + total += p.bytes; + if (p.row == 0) continue; + rows += 1; + try expect(256 % p.per == 0); + try expectEqual(@as(u64, at / p.per) * p.row, p.bytes); // a pooled part counts the whole pools alone + } + try expectEqual(@as(usize, 3 * 2), rows); // each DSA layer's latent rows, scales and whole pools + const last = segs.items[segs.items.len - 1]; + try expectEqual(total, last.img + last.len); + try expectEqual(@as(u64, 2 * 256), list.items[7].bytes); // layer 0's pools past `at`: the spare rows + try expectError(error.SnapshotShape, image.segments(sh, at, sh.headRows(at) + 1, &segs, a)); +} + +test "GLM disk image: a kept state's image read out of its lease restores every planned byte into another sequence" { + var m: Model = undefined; + try m.init(); + var d = H.driver(); + const a = std.testing.allocator; + const c = &m.w.config; + var st = try State.initFormat(a, &d, &m.w, capacity, false); + defer st.deinit(); + for (H.bytes(st.arena.buf.ptr, st.arena.buf.len), 0..) |*b, i| b.* = @truncate(i *% 31 +% 7); + var pool = try pages.Pool.init(a, &d, 64 << 20); + defer pool.deinit(); + const stream: cuda.Stream = .{ .d = &d, .handle = null }; + const at: u32 = 2048; + const mtp: u64 = at - 1; // a head one row behind: its last image row is not in the lease + const sh = try image.Shape.of(&st, c, capacity, true); + const plan = try Plan.init(&st, c, at, mtp); + var kept = try pool.take(plan.len); + defer pool.give(&kept); + try pool.move(stream, &kept, &plan, true); + var segs: std.ArrayList(image.Seg) = .empty; + defer segs.deinit(a); + try image.segments(sh, at, mtp, &segs, a); + const last = segs.items[segs.items.len - 1]; + const img = try a.alloc(u8, @intCast(last.img + last.len)); + defer a.free(img); + var words = (image.Words{ .at = at, .mtp_len = mtp, .base = 0, .mtp_base = 0, .cur = 0b10 }).encode(); + const out: image.Copy = .{ .d = &d, .pool = &pool, .lease = &kept, .segs = segs.items, .words = &words }; + const cut = [_]u64{ 0, 1000, plan.len / 2, img.len }; // the store asks for ranges, not the whole image + for (cut[0..3], cut[1..]) |lo, hi| try out.run(lo, img[@intCast(lo)..@intCast(hi)], true); + for (segs.items) |s| if (s.from == .zeros) for (img[@intCast(s.img)..][0..@intCast(s.len)]) |b| try expectEqual(@as(u8, 0), b); + // a blank state as a disk read takes it (room for the head's whole window), filled from the image + var blank = try pool.take(@intCast(image.leaseBytes(sh, at, sh.headRows(at)))); + defer pool.give(&blank); + var back: [6]u64 = undefined; + const in: image.Copy = .{ .d = &d, .pool = &pool, .lease = &blank, .segs = segs.items, .words = &back }; + try in.run(0, img, false); + const w = try image.Words.decode(back, at); + try expectEqual(mtp, w.mtp_len); + try expectEqual(@as(u64, 0b10), w.cur); + try expectError(error.DiskWords, image.Words.decode(back, at + 2048)); + var fresh = try State.initFormat(a, &d, &m.w, capacity, false); + defer fresh.deinit(); + const into = try Plan.init(&fresh, c, at, w.mtp_len); + try expectEqual(plan.len, into.len); + try pool.move(stream, &blank, &into, false); + for (plan.parts[0..plan.count], into.parts[0..into.count]) |x, y| { + try std.testing.expectEqualSlices(u8, H.bytes(x.live, x.len), H.bytes(y.live, y.len)); + } +} From a3c043cca95d585461c9988c61218c0f64613548 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski <24755+wojo@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:03:01 -0400 Subject: [PATCH 2/7] GLM on CUDA: --learn-conversations keeps every rank's part of a kept state on disk Rank 0's engine thread leases, fills, reads and seals each rank's copy through the follow link (cuda_disk_ranks: blank, put, take, seal, prefixes). The store's image is rank 0's image followed by each other rank's prefix, the one part a rank holds alone; the rest are the same bytes on every rank, checked at start by the identity probe. A blank state is refused by restore until its seal. --- docs/api.md | 3 +- zig/src/families/glm/cuda.zig | 2 + zig/src/families/glm/cuda_disk_image.zig | 58 +++++- zig/src/families/glm/cuda_disk_image_test.zig | 86 ++++++++ zig/src/families/glm/cuda_disk_ranks.zig | 152 ++++++++++++++ zig/src/families/glm/cuda_engine.zig | 3 + zig/src/families/glm/cuda_follow.zig | 21 +- zig/src/families/glm/cuda_native.zig | 1 + zig/src/families/glm/cuda_snapshots.zig | 27 +++ zig/src/native/cuda.zig | 4 +- zig/src/native/cuda_glm_disk.zig | 185 ++++++++++++++++++ zig/src/native/cuda_refusal.zig | 5 +- 12 files changed, 534 insertions(+), 13 deletions(-) create mode 100644 zig/src/families/glm/cuda_disk_ranks.zig create mode 100644 zig/src/native/cuda_glm_disk.zig diff --git a/docs/api.md b/docs/api.md index 419e3c8d54..9ce5d3280f 100644 --- a/docs/api.md +++ b/docs/api.md @@ -602,7 +602,8 @@ store for each model and build. That store takes the place of `--learn`'s store with the flag neither reads nor writes the prefix files a plain `--learn` server keeps, and they stay on disk until the cap below needs their room. After a restart, a crash or an eviction from RAM, the next turn reads the latest state back instead of reading the conversation again from its first token. Qwen3.8-27B, Nemotron, Flash Next and GLM-5.3-Flash on -one Mac and Nemotron on CUDA keep them; other engines refuse the flag at startup. +one Mac, and Nemotron and GLM-5.3-Flash on CUDA keep them; other engines refuse the flag at startup. GLM-5.3-Flash on +CUDA across ranks writes every rank's part to rank 0's disk, so a state is read back on every rank or on none. The engine stages an owned copy between rounds, and an I/O thread writes and fsyncs it. Staging has a 1 GiB byte permit and at most eight jobs; when either is full, the save is skipped. The copy waits for its GPU producers before diff --git a/zig/src/families/glm/cuda.zig b/zig/src/families/glm/cuda.zig index 5ea204949e..329386900a 100644 --- a/zig/src/families/glm/cuda.zig +++ b/zig/src/families/glm/cuda.zig @@ -37,6 +37,8 @@ pub const follow = @import("cuda_follow.zig"); pub const lanes = @import("cuda_lanes.zig"); pub const native = @import("cuda_native.zig"); pub const snapshots = @import("cuda_snapshots.zig"); +pub const disk_image = @import("cuda_disk_image.zig"); +pub const disk_ranks = @import("cuda_disk_ranks.zig"); test { _ = @import("cuda_tile.zig"); diff --git a/zig/src/families/glm/cuda_disk_image.zig b/zig/src/families/glm/cuda_disk_image.zig index 071d300661..614a085328 100644 --- a/zig/src/families/glm/cuda_disk_image.zig +++ b/zig/src/families/glm/cuda_disk_image.zig @@ -112,13 +112,10 @@ pub fn segments(sh: Shape, at: u32, mtp_len: u64, out: *std.ArrayList(Seg), a: s var img: u64 = Words.bytes; var lease: u64 = 0; try out.append(a, .{ .img = 0, .len = Words.bytes, .from = .words }); - var front: u64 = sh.prefix; // the prefix and every DSA layer's parts sit in the lease as in the image - for (0..sh.layers) |_| for (cache(sh, at)) |n| { - front += n; - }; - try out.append(a, .{ .img = img, .len = front, .from = .{ .lease = 0 } }); - img += front; - lease += front; + const f = front(sh, at); + try out.append(a, .{ .img = img, .len = f, .from = .{ .lease = 0 } }); + img += f; + lease += f; if (sh.head) for (cache(sh, sh.headRows(at)), cache(sh, mtp_len)) |full, held| { if (held > 0) try out.append(a, .{ .img = img, .len = held, .from = .{ .lease = lease } }); if (full > held) try out.append(a, .{ .img = img + held, .len = full - held, .from = .zeros }); @@ -128,6 +125,15 @@ pub fn segments(sh: Shape, at: u32, mtp_len: u64, out: *std.ArrayList(Seg), a: s if (sh.tail > 0) try out.append(a, .{ .img = img, .len = sh.tail, .from = .{ .lease = lease } }); } +/// The image bytes after the words that sit in the lease as in the image: the prefix and every DSA layer's parts. +pub fn front(sh: Shape, at: u32) u64 { + var n: u64 = sh.prefix; + for (0..sh.layers) |_| for (cache(sh, at)) |x| { + n += x; + }; + return n; +} + /// The lease bytes of a state at `at` whose head holds `mtp_len` rows (layout.Plan's length for it). pub fn leaseBytes(sh: Shape, at: u32, mtp_len: u64) u64 { var n: u64 = sh.prefix + sh.tail; @@ -140,6 +146,44 @@ pub fn leaseBytes(sh: Shape, at: u32, mtp_len: u64) u64 { return n; } +/// Rank 0's image bytes at `at` (its head counted at a verify window past `at`). +pub fn imageBytes(sh: Shape, at: u32) u64 { + var n: u64 = Words.bytes + front(sh, at) + sh.tail; + if (sh.head) for (cache(sh, sh.headRows(at))) |x| { + n += x; + }; + return n; +} + +/// `len` bytes at `off` of rank `rank`'s image; `shared`: past rank 0's prefix, the same bytes every rank holds. +pub const Piece = struct { rank: u32, off: u64, len: u64, shared: bool }; + +/// The store's image at `at`: rank 0's, then each other rank's prefix (rank 0's bytes stand for all the rest). +pub fn storeParts(sh: Shape, prefixes: []const u64, at: u32, out: *std.ArrayList(Part), a: std.mem.Allocator) !void { + try parts(sh, at, out, a); + for (prefixes[1..]) |p| try out.append(a, .{ .bytes = p }); +} + +/// Store image bytes [off, off + len) split by rank: rank 0's image, then rank r's prefix after its words. +pub fn route(sh: Shape, prefixes: []const u64, at: u32, off: u64, len: u64, out: *std.ArrayList(Piece), a: std.mem.Allocator) !void { + const own = Words.bytes + sh.prefix; // rank 0's bytes no other rank holds + var start: u64 = 0; // the section's first store byte + for (prefixes, 0..) |p, r| { + const size = if (r == 0) imageBytes(sh, at) else p; + defer start += size; + const lo = @max(off, start); + const hi = @min(off + len, start + size); + if (lo >= hi) continue; + if (r > 0) { + try out.append(a, .{ .rank = @intCast(r), .off = Words.bytes + lo - start, .len = hi - lo, .shared = false }); + continue; + } + if (lo < own) try out.append(a, .{ .rank = 0, .off = lo, .len = @min(hi, own) - lo, .shared = false }); + if (hi > own) try out.append(a, .{ .rank = 0, .off = @max(lo, own), .len = hi - @max(lo, own), .shared = true }); + } + if (off + len > start) return error.SnapshotShape; +} + /// Copies between `host` and image bytes [off, off + host.len) of a lease: out (lease to host) or in. pub const Copy = struct { d: *const cuda.Driver, diff --git a/zig/src/families/glm/cuda_disk_image_test.zig b/zig/src/families/glm/cuda_disk_image_test.zig index 30d1d84b65..fea37bd00c 100644 --- a/zig/src/families/glm/cuda_disk_image_test.zig +++ b/zig/src/families/glm/cuda_disk_image_test.zig @@ -182,3 +182,89 @@ test "GLM disk image: a kept state's image read out of its lease restores every try std.testing.expectEqualSlices(u8, H.bytes(x.live, x.len), H.bytes(y.live, y.len)); } } + +test "GLM disk image: the store's image is rank 0's then each other rank's prefix, and routes every byte once" { + var m: Model = undefined; + try m.init(); + var d = H.driver(); + const a = std.testing.allocator; + var st = try State.initFormat(a, &d, &m.w, capacity, true); + defer st.deinit(); + const sh = try image.Shape.of(&st, &m.w.config, capacity, true); + const at: u32 = 4096; + const prefixes = [_]u64{ sh.prefix, sh.prefix - 4096, sh.prefix + 512 }; // TP=3: heads one by one + var list: std.ArrayList(image.Part) = .empty; + defer list.deinit(a); + try image.storeParts(sh, &prefixes, at, &list, a); + var total: u64 = 0; + for (list.items) |p| total += p.bytes; + try expectEqual(image.imageBytes(sh, at) + prefixes[1] + prefixes[2], total); + var pieces: std.ArrayList(image.Piece) = .empty; + defer pieces.deinit(a); + var seen: [3]u64 = .{ 0, 0, 0 }; + var shared: u64 = 0; + const cuts = [_]u64{ 0, 40, image.Words.bytes + sh.prefix - 8, image.imageBytes(sh, at) + 100, total }; // ranges across the seams + for (cuts[0 .. cuts.len - 1], cuts[1..]) |lo, hi| { + pieces.clearRetainingCapacity(); + try image.route(sh, &prefixes, at, lo, hi - lo, &pieces, a); + var n: u64 = 0; + for (pieces.items) |p| { + n += p.len; + seen[p.rank] += p.len; + if (p.shared) shared += p.len; + if (p.rank > 0) try expect(p.off >= image.Words.bytes and p.off + p.len <= image.Words.bytes + prefixes[p.rank] and !p.shared); + if (p.shared) try expect(p.rank == 0 and p.off >= image.Words.bytes + sh.prefix); + } + try expectEqual(hi - lo, n); + } + try expectEqual(image.imageBytes(sh, at), seen[0]); + try expectEqual(prefixes[1], seen[1]); + try expectEqual(prefixes[2], seen[2]); + try expectEqual(image.imageBytes(sh, at) - image.Words.bytes - sh.prefix, shared); + try expectError(error.SnapshotShape, image.route(sh, &prefixes, at, total - 4, 8, &pieces, a)); +} + +/// The fields of an engine the kept-state functions read, over a host-memory driver. +const Fake = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + stream: cuda.Stream, + c: @import("config.zig").Config, + cur: struct { st: State, in_head: usize = 0 }, + capacity: usize = capacity, + mtp: bool = true, + snaps: std.AutoHashMapUnmanaged(u64, @import("cuda_engine_types.zig").Snap) = .empty, + snap_pages: ?pages.Pool = null, +}; + +test "GLM disk image: a disk read's blank state is refused until sealed, and its seal checks the words" { + const snaps = @import("cuda_snapshots.zig"); + var m: Model = undefined; + try m.init(); + var d = H.driver(); + const a = std.testing.allocator; + var e: Fake = .{ .gpa = a, .d = &d, .stream = .{ .d = &d, .handle = null }, .c = m.w.config, .cur = .{ .st = try State.initFormat(a, &d, &m.w, capacity, false) } }; + defer e.cur.st.deinit(); + e.snap_pages = try pages.Pool.init(a, &d, 256 << 20); + defer e.snap_pages.?.deinit(); + try e.snaps.ensureTotalCapacity(a, 8); + defer e.snaps.deinit(a); + const at: u32 = 2048; + try snaps.blankAs(&e, 7, at); + defer { + var s = e.snaps.get(7).?; + snaps.free(&e, &s); + } + try expectError(error.StateKept, snaps.blankAs(&e, 7, at)); + try expectError(error.UnsealedState, snaps.restore(&e, 7)); + try expectError(error.DiskWords, snaps.seal(&e, 7, at + 9, .{ 0, 0 }, 0)); // a head past its verify window + try expectError(error.DiskWords, snaps.seal(&e, 7, at, .{ 0, 0 }, @as(u64, 1) << @intCast(e.cur.st.cur.len))); // a copy bit no layer has + try snaps.seal(&e, 7, at - 1, .{ 3, 5 }, 0b01); + const s = e.snaps.get(7).?; + const sh = try image.Shape.of(&e.cur.st, &e.c, capacity, true); + try expectEqual(image.leaseBytes(sh, at, at - 1), @as(u64, s.len)); + try expectEqual(@as(usize, at - 1), s.mtp_len); + try expectEqual([2]usize{ 3, 5 }, s.bases); + try std.testing.expectEqualSlices(u1, &.{ 1, 0 }, s.cur); + try expectError(error.StateKept, snaps.seal(&e, 7, at, .{ 0, 0 }, 0)); +} diff --git a/zig/src/families/glm/cuda_disk_ranks.zig b/zig/src/families/glm/cuda_disk_ranks.zig new file mode 100644 index 0000000000..975b025d18 --- /dev/null +++ b/zig/src/families/glm/cuda_disk_ranks.zig @@ -0,0 +1,152 @@ +//! A GLM kept state's disk image across ranks: rank 0's engine thread leases, fills, reads and seals each rank's copy. +const std = @import("std"); +const follow = @import("cuda_follow.zig"); +const image = @import("cuda_disk_image.zig"); +const Engine = @import("cuda_engine.zig").Engine; + +/// A disk step's op (Header.flags); its id is `a`, its image offset `seed`, its bytes `depth`, its rank `top_k`. +pub const Op = enum(u32) { blank = 1, put, take, seal, prefixes }; +const kind = @backingInt(follow.Kind.disk); // each op crosses the follow link as a disk step, its bytes as the ids + +/// This rank's image shape (every sequence of the engine has the same). +pub fn shape(e: *const Engine) !image.Shape { + return image.Shape.of(&e.cur.st, &e.c, e.capacity, e.mtp); +} + +/// The words of a kept state on this rank: its position, head rows, ring bases and KDA copy bits. +pub fn wordsOf(e: *const Engine, id: u64) !image.Words { + const s = e.snaps.get(id) orelse return error.NoKeptState; + if (s.cur.len > 64) return error.SnapshotShape; + var cur: u64 = 0; + for (s.cur, 0..) |b, i| cur |= @as(u64, b) << @intCast(i); + return .{ .at = s.pos, .mtp_len = s.mtp_len, .base = s.bases[0], .mtp_base = s.bases[1], .cur = cur }; +} + +/// Bytes [off, off + host.len) of this rank's image of `id`, out (lease to host) or in; unsealed: before its head only. +pub fn copy(e: *Engine, id: u64, off: u64, host: []u8, out: bool) !void { + const s = e.snaps.get(id) orelse return error.NoKeptState; + const sh = try shape(e); + const at: u32 = @intCast(s.pos); + if (s.len == 0 and off + host.len > image.Words.bytes + image.front(sh, at)) return error.DiskOrder; + var segs: std.ArrayList(image.Seg) = .empty; + defer segs.deinit(e.gpa); + try image.segments(sh, at, s.mtp_len, &segs, e.gpa); + var w = (try wordsOf(e, id)).encode(); + const pool = if (e.snap_pages) |*p| p else return error.NoSnapshotCache; + try (image.Copy{ .d = e.d, .pool = pool, .lease = &s.lease, .segs = segs.items, .words = &w }).run(off, host, out); +} + +/// Rank 0's link at TP, else null (one rank: every op is local). +fn leader(e: *Engine) ?follow.Leader { + const l = e.link orelse return null; + return .{ .e = e, .link = @ptrCast(@alignCast(l)) }; +} + +/// Rank 0: whether the other ranks sleep on their lines (no stream), and putting them back to sleep after disk steps. +pub fn asleep(e: *Engine) bool { + return if (leader(e)) |x| x.link.asleep else false; +} + +pub fn rest(e: *Engine) !void { + if (leader(e)) |x| try x.idle(); +} + +/// Rank 0: a new id leased at `at` on every rank (rank 0 first: the ranks' pools match, so theirs cannot fail apart). +pub fn blank(e: *Engine, at: u32) !u64 { + if (e.next_snap > std.math.maxInt(u32)) return error.SnapshotIdLimit; + const id = e.next_snap; + try e.blankAs(id, at); + e.next_snap += 1; + if (leader(e)) |x| try x.disk(.{ .kind = kind, .flags = @backingInt(Op.blank), .a = id, .seed = at }, &.{}); + return id; +} + +/// Rank 0: `bytes` into rank `rank`'s image of `id` at `off` (rank 0's own lease, or through the link in pieces). +pub fn put(e: *Engine, rank: u32, id: u64, off: u64, bytes: []const u8) !void { + if (rank == 0) return copy(e, id, off, @constCast(bytes), false); + const x = leader(e) orelse return error.NoRank; + const room = x.link.ids.len * 4; + var done: usize = 0; + while (done < bytes.len) { + const n = @min(room, bytes.len - done); + const words = (n + 3) / 4; + @memcpy(std.mem.sliceAsBytes(x.link.ids[0..words])[0..n], bytes[done..][0..n]); + try x.disk(.{ .kind = kind, .flags = @backingInt(Op.put), .a = id, .seed = off + done, .depth = @intCast(n), .top_k = rank }, x.link.ids[0..words]); + done += n; + } +} + +/// Rank 0: rank `rank`'s image bytes of `id` from `off` into `out`. +pub fn take(e: *Engine, rank: u32, id: u64, off: u64, out: []u8) !void { + if (rank == 0) return copy(e, id, off, out, true); + const x = leader(e) orelse return error.NoRank; + const room = x.link.ids.len * 4; + var done: usize = 0; + while (done < out.len) { + const n = @min(room, out.len - done); + try x.disk(.{ .kind = kind, .flags = @backingInt(Op.take), .a = id, .seed = off + done, .depth = @intCast(n), .top_k = rank }, &.{}); + try x.link.from(rank, n); + try x.link.buf.download(64, out[done..][0..n]); + done += n; + } +} + +/// Rank 0: every rank's blank `id` sealed with `w` (its head's rows, ring bases and KDA copy bits). +pub fn seal(e: *Engine, id: u64, w: image.Words) !void { + try e.sealState(id, w.mtp_len, .{ w.base, w.mtp_base }, w.cur); + if (leader(e)) |x| { + const words = w.encode(); + try x.disk(.{ .kind = kind, .flags = @backingInt(Op.seal), .a = id }, std.mem.bytesAsSlice(u32, std.mem.sliceAsBytes(&words))); + } +} + +/// Rank 0: every rank's prefix bytes in rank order (the part of a state each rank holds alone). +pub fn prefixes(e: *Engine, out: []u64) !void { + const mine = (try shape(e)).prefix; + const x = leader(e) orelse { + if (out.len != 1) return error.NoRank; + out[0] = mine; + return; + }; + try x.disk(.{ .kind = kind, .flags = @backingInt(Op.prefixes) }, &.{}); + try gather(x.link, e, mine, out); +} + +fn gather(l: *follow.Link, e: *const Engine, mine: u64, out: []u64) !void { + try @import("cuda_ranks.zig").gather(u64, l.d, l.lib, l.comm, l.stream, e.share.rank, mine, out); +} + +/// A rank other than 0: rank 0's disk step on this rank's copy of its state. +pub fn step(e: *Engine, l: *follow.Link, h: follow.Header, ids: []const u32) !void { + if (h.flags < @backingInt(Op.blank) or h.flags > @backingInt(Op.prefixes)) return error.BadStep; + const op: Op = @fromBackingInt(h.flags); + const rank: u32 = @intCast(e.share.rank); + switch (op) { + .blank => try e.blankAs(h.a, @intCast(h.seed)), + .seal => { + if (ids.len != 12) return error.BadStep; + var x: [6]u64 = undefined; + @memcpy(std.mem.sliceAsBytes(&x), std.mem.sliceAsBytes(ids)); + const w = try image.Words.decode(x, @intCast((e.snaps.get(h.a) orelse return error.NoKeptState).pos)); + try e.sealState(h.a, w.mtp_len, .{ w.base, w.mtp_base }, w.cur); + }, + .put => if (h.top_k == rank) { + if (h.depth > ids.len * 4) return error.BadStep; + try copy(e, h.a, h.seed, @constCast(std.mem.sliceAsBytes(ids))[0..h.depth], false); + }, + .take => { + if (h.depth > l.ids.len * 4) return error.BadStep; + if (h.top_k == rank) { + const host = std.mem.sliceAsBytes(l.ids)[0..h.depth]; + try copy(e, h.a, h.seed, host, true); + try l.buf.uploadAsync(64, host, l.stream.handle); + } + try l.from(h.top_k, h.depth); + }, + .prefixes => { + var out: [64]u64 = undefined; + if (e.share.world > out.len) return error.BadStep; + try gather(l, e, (try shape(e)).prefix, out[0..e.share.world]); + }, + } +} diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig index 69c309e40b..3ebd6cc76c 100644 --- a/zig/src/families/glm/cuda_engine.zig +++ b/zig/src/families/glm/cuda_engine.zig @@ -84,6 +84,7 @@ pub const Engine = struct { snaps: std.AutoHashMapUnmanaged(u64, Snap) = .empty, // kept prompt states by id (the same ids on every rank) snap_pages: ?snapshots.pages.Pool = null, snap_setup: ?snapshots.Setup = null, + link: ?*anyopaque = null, // rank 0 at TP: its follow.Link, for the disk steps (cuda_disk_ranks) next_snap: u64 = 1, // rank 0: the next id announce: u64 = 0, // rank 0: a state kept at this chunk's mark, for the other ranks to copy at its gate dropped: std.ArrayList(u32) = .empty, // rank 0: ids dropped since the other ranks last heard (cuda_follow) @@ -418,6 +419,8 @@ pub const Engine = struct { pub const prepareSnapshots = snapshots.prepare; pub const keepState = snapshots.keep; pub const keepAs = snapshots.keepAs; + pub const blankAs = snapshots.blankAs; + pub const sealState = snapshots.seal; pub const restoreState = snapshots.restore; pub const holds = snapshots.holds; pub const dropState = snapshots.drop; diff --git a/zig/src/families/glm/cuda_follow.zig b/zig/src/families/glm/cuda_follow.zig index 6e5436dfb1..9e5f326e93 100644 --- a/zig/src/families/glm/cuda_follow.zig +++ b/zig/src/families/glm/cuda_follow.zig @@ -12,7 +12,7 @@ const pages_mod = @import("cuda_pages.zig"); const sparse = @import("cuda_sparse.zig"); const Engine = engine.Engine; -pub const Kind = enum(u32) { begin = 1, prefill, verify, keep, draft, end, drop, release, chunk, shared, image, block, many, redraw, idle, snapshots, pool }; +pub const Kind = enum(u32) { begin = 1, prefill, verify, keep, draft, end, drop, release, chunk, shared, image, block, many, redraw, idle, snapshots, pool, disk }; /// One step as it crosses: 64 bytes, then `n` u32 ids. pub const Header = extern struct { @@ -106,6 +106,14 @@ pub const Link = struct { try l.lib.check(l.lib.api.ncclBroadcast(ptr, ptr, bytes, .u8, 0, l.comm, l.stream.handle), "ncclBroadcast"); } + /// `bytes` after the header from rank `root` into every rank's buffer (a disk step's take: cuda_disk_ranks). + pub fn from(l: *Link, root: u32, bytes: usize) !void { + if (bytes == 0) return; + if (64 + bytes > l.buf.len) return error.TooManyIds; + try l.lib.check(l.lib.api.ncclBroadcast(l.buf.ptr + 64, l.buf.ptr + 64, bytes, .u8, @intCast(root), l.comm, l.stream.handle), "ncclBroadcast"); + try l.stream.synchronize(); + } + /// Other ranks: the next step and its ids (valid until the next receive). pub fn receive(l: *Link) !struct { h: Header, ids: []const u32 } { try l.bcast(64); @@ -291,6 +299,14 @@ pub const Leader = struct { try x.e.draftBlock(reqs); } + /// A kept state's disk step (cuda_disk_ranks): its op and ids to every rank. + pub fn disk(x: Leader, h: Header, ids: []const u32) !void { + try x.flush(); + var hh = h; + hh.kind = @backingInt(Kind.disk); + try x.send(hh, ids); + } + /// The other ranks leave their loops. pub fn end(x: Leader) !void { try x.send(.{ .kind = @backingInt(Kind.end) }, &.{}); @@ -324,7 +340,7 @@ pub fn follow(e: *Engine, link: *Link) !void { const step = try link.receive(); const h = step.h; const kind: Kind = @fromBackingInt(@intCast(h.kind)); - if (kind != .end and kind != .drop and kind != .shared and kind != .block and kind != .many and kind != .idle and kind != .snapshots and kind != .pool) try e.bind(h.slot); + if (kind != .end and kind != .drop and kind != .shared and kind != .block and kind != .many and kind != .idle and kind != .snapshots and kind != .pool and kind != .disk) try e.bind(h.slot); e.prompt_logprobs = if (promptRows(h)) .{} else null; e.structured = structuredStep(h); // rank 0 alone masks them switch (kind) { @@ -398,6 +414,7 @@ pub fn follow(e: *Engine, link: *Link) !void { try link.rows(v.rows.ptr, v.rowCount() * vcfg.out_hidden * 2); }, .snapshots => try @import("cuda_snapshots.zig").agree(e, h.a), + .disk => try @import("cuda_disk_ranks.zig").step(e, link, h, step.ids), .drop => for (step.ids) |id| e.dropState(id), .release => e.release(h.slot), .idle => cuda.tp_watch.sleep(e.io), diff --git a/zig/src/families/glm/cuda_native.zig b/zig/src/families/glm/cuda_native.zig index 620f2c01e9..b7f8d38d8b 100644 --- a/zig/src/families/glm/cuda_native.zig +++ b/zig/src/families/glm/cuda_native.zig @@ -269,6 +269,7 @@ fn finishRank(a: OpenArgs) !Loaded { own.e.group = &own.group; own.link = try follow.Link.init(gpa, ctx.d, &own.group.lib, own.group.comm, own.e.stream, try linkCapacity(cap, o.vision)); errdefer own.link.deinit(); + if (p.rank == 0 and p.world > 1) own.e.link = &own.link; // its disk steps (cuda_disk_ranks) cross on the link own.e.snap_setup = .{ .ptr = own, .run = readySnapshots }; // Rank 0 admits by the smallest room of the ranks; its sizes are the largest (the first ranks take the remainder). const rooms = try gpa.alloc(u64, p.world); diff --git a/zig/src/families/glm/cuda_snapshots.zig b/zig/src/families/glm/cuda_snapshots.zig index 08101ab6c1..5640051d41 100644 --- a/zig/src/families/glm/cuda_snapshots.zig +++ b/zig/src/families/glm/cuda_snapshots.zig @@ -58,8 +58,35 @@ pub fn keepAs(e: anytype, id: u64) !void { e.snaps.putAssumeCapacityNoClobber(id, .{ .lease = lease, .len = plan.len, .pos = e.cur.st.pos, .mtp_len = e.cur.st.mtp_len, .cur = cur, .bases = .{ e.cur.st.base, e.cur.st.mtp_base } }); } +/// A disk read's state: `at`'s pages leased under rank 0's id on each rank, refused by `restore` until `seal`. +pub fn blankAs(e: anytype, id: u64, at: usize) !void { + if (e.snaps.contains(id)) return error.StateKept; + const pool = if (e.snap_pages) |*p| p else return error.NoSnapshotCache; + var lease = try pool.take(bytes(e, at)); + errdefer pool.give(&lease); + const cur = try e.gpa.alloc(u1, e.cur.st.cur.len); + @memset(cur, 0); + e.snaps.putAssumeCapacityNoClobber(id, .{ .lease = lease, .len = 0, .pos = at, .mtp_len = 0, .cur = cur, .bases = .{ 0, 0 } }); +} + +/// A blank state made whole from its image's words (its bytes are in or follow): a restore takes it from here. +pub fn seal(e: anytype, id: u64, mtp_len: usize, bases: [2]usize, cur: u64) !void { + const s = e.snaps.getPtr(id) orelse return error.NoKeptState; + if (s.len != 0) return error.StateKept; + if (s.cur.len > 64 or (s.cur.len < 64 and cur >> @intCast(s.cur.len) != 0)) return error.DiskWords; + const head = if (e.mtp) @min(e.capacity, s.pos + @import("cuda_state.zig").window_rows) else 0; + if (mtp_len > head) return error.DiskWords; + const n = try layout.Plan.size(&e.cur.st, &e.c, e.capacity, s.pos, if (e.mtp) mtp_len else null); + if (n == 0 or n > s.lease.ids.items.len * pages.page) return error.SnapshotSize; + for (s.cur, 0..) |*b, i| b.* = @intCast(cur >> @intCast(i) & 1); + s.len = n; + s.mtp_len = mtp_len; + s.bases = bases; +} + pub fn restore(e: anytype, id: u64) !void { const s = e.snaps.getPtr(id) orelse return error.NoKeptState; + if (s.len == 0) return error.UnsealedState; const pool = if (e.snap_pages) |*p| p else return error.NoSnapshotCache; const plan = try layout.Plan.init(&e.cur.st, &e.c, s.pos, if (e.mtp) s.mtp_len else null); if (plan.len != s.len) return error.SnapshotSize; diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index ab6a6e9443..ab48893c06 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -24,6 +24,7 @@ const lane_backend = @import("cuda_bound.zig"); const snaps = @import("cuda_snaps.zig"); const slide_adapter = @import("slide_adapter.zig"); const cuda_disk = @import("cuda_disk.zig"); +const glm_disk = @import("cuda_glm_disk.zig"); const vision_offer = @import("cuda_vision_offer.zig"); const drafted = @import("cuda_drafted.zig"); const kv_pages = @import("cuda_kv_pool.zig"); @@ -410,7 +411,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O const store = try gpa.create(pc.Store); errdefer gpa.destroy(store); // rules and budget from the family's snap_* declarations (cuda_snaps.zig) - const image: ?*const pc.rows.Image = if (F == nemotron.native) &cuda_disk.image else if (comptime cuda_disk.keptDisk(F)) |K| &K.image else null; + const image: ?*const pc.rows.Image = if (F == nemotron.native) &cuda_disk.image else if (F == glm.native) &glm_disk.image else if (comptime cuda_disk.keptDisk(F)) |K| &K.image else null; const copies: pc.Snapshots = .{ .ptr = loaded.target, .vtable = &snaps.Snaps(F, image).vtable }; store.* = pc.Store.init(gpa, copies, snaps.passRules(F, @import("cuda_startup.zig").rows(F, loaded)), if (@hasDecl(F, "snap_slots")) cache else try snaps.cacheBudget(F, loaded.target, cache)); errdefer { @@ -423,6 +424,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O h.store = store; if (o.learn_conversations) if (o.learn) |root| { if (F == nemotron.native) try cuda_disk.attach(a, io, &g.ctx, o, root, kernels, loaded, store); + if (F == glm.native) try glm_disk.attach(a, io, &g.ctx, o, root, kernels, loaded, store); if (comptime cuda_disk.keptDisk(F)) |K| try cuda_disk.attachKept(K, F.model_type, a, io, &g.ctx, o, root, kernels, loaded, store); }; } diff --git a/zig/src/native/cuda_glm_disk.zig b/zig/src/native/cuda_glm_disk.zig new file mode 100644 index 0000000000..8ddc5ba09a --- /dev/null +++ b/zig/src/native/cuda_glm_disk.zig @@ -0,0 +1,185 @@ +//! GLM-5.3-Flash's CUDA kept states on disk (--learn-conversations): every rank's part in rank 0's files. +const std = @import("std"); +const cuda = @import("cuda"); +const api = @import("engine_api"); +const lanes = @import("lanes"); +const glm = @import("glm_cuda"); +const cuda_disk = @import("cuda_disk.zig"); +const img = glm.disk_image; +const ranks = glm.disk_ranks; +const rows = api.prompt_cache.rows; +const Engine = glm.engine.Engine; +const Kept = glm.engine.Kept; +const Allocator = std.mem.Allocator; + +pub const image: rows.Image = .{ .parts = parts, .gather = gather, .blank = blank, .scatter = scatter, .whole = whole }; + +/// Every rank's prefix bytes, read at attach (one GLM engine a process; rank 0 alone serves). +var prefix_buf: [8]u64 = undefined; +var prefixes: []const u64 = &.{}; + +fn engineOf(ptr: *anyopaque) *Engine { + return @ptrCast(@alignCast(ptr)); +} + +fn parts(ptr: *anyopaque, at: u32, out: *std.ArrayList(rows.image.Part), a: Allocator) anyerror!void { + var list: std.ArrayList(img.Part) = .empty; + defer list.deinit(a); + try img.storeParts(try ranks.shape(engineOf(ptr)), prefixes, at, &list, a); + for (list.items) |p| try out.append(a, .{ .bytes = p.bytes, .row = p.row, .per = p.per }); +} + +/// The other ranks asleep before a disk call go back to sleep after it (their lines, not NCCL, wait meanwhile). +const Rest = struct { + e: *Engine, + asleep: bool, + + fn of(e: *Engine) Rest { + return .{ .e = e, .asleep = ranks.asleep(e) }; + } + + fn done(r: Rest) void { + if (r.asleep) ranks.rest(r.e) catch |err| std.log.warn("glm: the ranks did not sleep after a disk step ({s})", .{@errorName(err)}); + } +}; + +fn gather(ptr: *anyopaque, saved: api.prompt_cache.Saved, ranges: []const rows.image.Range, out: []u8) anyerror!void { + const e = engineOf(ptr); + const k: *Kept = @ptrCast(@alignCast(saved)); + const rest = Rest.of(e); + defer rest.done(); + var pieces: std.ArrayList(img.Piece) = .empty; + defer pieces.deinit(e.gpa); + const sh = try ranks.shape(e); + var o: usize = 0; + for (ranges) |g| { + pieces.clearRetainingCapacity(); + try img.route(sh, prefixes, k.at, g.off, g.len, &pieces, e.gpa); + for (pieces.items) |p| { + const n: usize = @intCast(p.len); + try ranks.take(e, p.rank, k.id, p.off, out[o..][0..n]); + o += n; + } + } +} + +/// A new id leased at `at` on every rank, filled by `scatter` and sealed by the words in its last one. +fn blank(ptr: *anyopaque, at: u32) anyerror!api.prompt_cache.Saved { + const e = engineOf(ptr); + const rest = Rest.of(e); + defer rest.done(); + const k = try e.gpa.create(Kept); + errdefer e.gpa.destroy(k); + k.* = .{ .id = try ranks.blank(e, at), .at = at }; + return @ptrCast(k); +} + +/// The last scatter brings the state's own file, words first: every rank seals it then (the words place the head). +fn scatter(ptr: *anyopaque, saved: api.prompt_cache.Saved, ranges: []const rows.image.Range, in: []const u8) anyerror!void { + const e = engineOf(ptr); + const k: *Kept = @ptrCast(@alignCast(saved)); + const rest = Rest.of(e); + defer rest.done(); + if (ranges.len > 0 and ranges[0].off == 0) { + if (ranges[0].len < img.Words.bytes) return error.DiskWords; + var x: [6]u64 = undefined; + @memcpy(std.mem.sliceAsBytes(&x), in[0..img.Words.bytes]); + try ranks.seal(e, k.id, try img.Words.decode(x, k.at)); + } + var pieces: std.ArrayList(img.Piece) = .empty; + defer pieces.deinit(e.gpa); + const sh = try ranks.shape(e); + var o: usize = 0; + for (ranges) |g| { + pieces.clearRetainingCapacity(); + try img.route(sh, prefixes, k.at, g.off, g.len, &pieces, e.gpa); + for (pieces.items) |p| { + const n: usize = @intCast(p.len); + const bytes = in[o..][0..n]; + o += n; + if (p.rank == 0 and p.off < img.Words.bytes) continue; // the words went with the seal + try ranks.put(e, p.rank, k.id, p.off, bytes); + if (p.shared) for (prefixes[1..], 1..) |q, r| try ranks.put(e, @intCast(r), k.id, p.off - prefixes[0] + q, bytes); + } + } +} + +fn whole(ptr: *anyopaque, saved: api.prompt_cache.Saved) bool { + const k: *Kept = @ptrCast(@alignCast(saved)); + const s = engineOf(ptr).snaps.get(k.id) orelse return false; + return s.len != 0; +} + +/// A fixed prompt's store image hashed, and each other rank's shared bytes checked equal to rank 0's (taken as theirs). +fn probe(a: Allocator, io: std.Io, backend: lanes.backend.Backend, e: *Engine) !struct { hash: u64, rate: f64 } { + const mark: u32 = glm.native.prompt_rows; + var toks: [glm.native.prompt_rows + 40]u32 = undefined; + for (&toks, 0..) |*x, i| x.* = @intCast(1000 + i * 7919 % 50000); + const Hook = struct { + e: *Engine, + saved: ?*anyopaque = null, + failed: ?anyerror = null, + fn at(p: *anyopaque, _: *lanes.Stream, m: u32) void { + const h: *@This() = @ptrCast(@alignCast(p)); + h.saved = glm.native.snap_save(h.e, null, m) catch |err| { + h.failed = err; + return; + }; + } + }; + var hook: Hook = .{ .e = e }; + const marks = [_]u32{mark}; + var s = try lanes.Stream.init(a, .{ .id = "probe", .prompt = &toks, .max_new = 1, .drafts = e.mtp, .reuse = .{ .marks = &marks, .hook = .{ .ptr = &hook, .at = Hook.at } } }); + defer s.deinit(a); + const t0 = std.Io.Clock.awake.now(io).toNanoseconds(); + backend.vtable.prefill(backend.ptr, &s) catch |err| { + backend.vtable.release(backend.ptr, &s); + return err; + }; + const ns = std.Io.Clock.awake.now(io).toNanoseconds() - t0; + backend.vtable.release(backend.ptr, &s); + const saved = hook.saved orelse return hook.failed orelse error.ProbeUnmarked; + defer glm.native.snap_drop(e, saved); + const k: *Kept = @ptrCast(@alignCast(saved)); + const sh = try ranks.shape(e); + const own = img.imageBytes(sh, mark); + var total = own; + for (prefixes[1..]) |p| total += p; + const bytes = try a.alloc(u8, @intCast(total)); + try gather(e, saved, &.{.{ .off = 0, .len = total }}, bytes); + const shared = bytes[@intCast(img.Words.bytes + sh.prefix)..@intCast(own)]; + const theirs = try a.alloc(u8, shared.len); + for (prefixes[1..], 1..) |q, r| { + try ranks.take(e, @intCast(r), k.id, img.Words.bytes + q, theirs); + if (!std.mem.eql(u8, shared, theirs)) return error.RanksDiffer; + } + try ranks.rest(e); + return .{ .hash = std.hash.Wyhash.hash(0x676c6d, bytes), .rate = @as(f64, @floatFromInt(toks.len)) / (@as(f64, @floatFromInt(@max(ns, 1))) / 1e9) }; +} + +/// --learn-conversations: every rank's prefix read, the identity probed on this thread, the store attached. +pub fn attach(a: Allocator, io: std.Io, ctx: *const cuda.Context, o: api.Open, root: []const u8, kernels: ?[]const u8, loaded: anytype, store: *api.prompt_cache.Store) !void { + const e: *Engine = loaded.target; + if (e.share.world > prefix_buf.len) return error.TooManyRanks; + try ranks.prefixes(e, prefix_buf[0..e.share.world]); + prefixes = prefix_buf[0..e.share.world]; + const sh = try ranks.shape(e); + var facts: std.ArrayList(u64) = .empty; + try facts.appendSlice(a, &.{ e.share.world, sh.capacity, sh.ring_rows, sh.tail, @intFromBool(sh.head), @intFromBool(sh.scales), @intFromBool(e.drafter != null) }); + try facts.appendSlice(a, prefixes); + const p = try probe(a, io, loaded.backend, e); + const id = try cuda_disk.identityOf(a, io, ctx, o.dir, kernels, facts.items, p.hash); + const gibs = struct { + fn bytes(g: f64) u64 { + return @intFromFloat(g * (1 << 30)); + } + }.bytes; + const r = try rows.Rows.open(store.gpa, store.family, .{ .root = root, .identity = id, .cap = gibs(o.learn_gib), .floor = gibs(o.learn_min_free_gib), .daily = gibs(o.learn_daily_gib), .lookahead = store.rules.lookahead, .planned = store.rules.planned, .grid = store.rules.grid }); + r.prefill_rate = p.rate; + store.disk = r; + std.log.info("glm: conversation states on disk in {s} ({d} states, {d} MiB of {d} MiB; {d} ranks' parts)", .{ r.dir, r.states.items.len, r.held() >> 20, r.cap >> 20, prefixes.len }); +} + +test { + std.testing.refAllDecls(@This()); +} diff --git a/zig/src/native/cuda_refusal.zig b/zig/src/native/cuda_refusal.zig index 70c2342065..70db25c1e9 100644 --- a/zig/src/native/cuda_refusal.zig +++ b/zig/src/native/cuda_refusal.zig @@ -47,7 +47,7 @@ pub fn family(comptime F: type, o: api.Open) ?[]const u8 { if (@hasDecl(F, "drafter_refusal")) return F.drafter_refusal; // a family without an MTP head says so itself return "--drafter loads a separate drafter, and this CUDA family drafts with its own MTP head: serve without --drafter"; } - if (o.learn_conversations and F != nemotron.native and @import("cuda_disk.zig").keptDisk(F) == null) return "--learn-conversations keeps conversation states on disk for Nemotron, the Qwen hybrids and Flash Next only on CUDA: serve this model without it"; + if (o.learn_conversations and F != nemotron.native and F != @import("glm_cuda").native and @import("cuda_disk.zig").keptDisk(F) == null) return "--learn-conversations keeps conversation states on disk for Nemotron, GLM-5.3-Flash, the Qwen hybrids and Flash Next only on CUDA: serve this model without it"; if (o.learn_conversations and o.tp > 1) if (@import("cuda_disk.zig").keptDisk(F)) |K| if (@hasDecl(K, "one_rank") and K.one_rank) return "--learn-conversations keeps one rank's states on disk for this CUDA family: serve without --tp or without --learn-conversations"; if (o.kv8 and !(@hasDecl(F, "Options") and @hasField(F.Options, "kv8"))) return "--kv-cache fp8: this CUDA family keeps bf16 caches; serve without it"; if (o.prompt_int8 and !(@hasDecl(F, "Options") and @hasField(F.Options, "prompt_int8"))) return "--prompt-precision int8: this CUDA family's prompt matmuls take bf16 inputs; serve without it"; @@ -135,7 +135,7 @@ test "a family with no head of its own words its --drafter refusal" { try std.testing.expectEqualStrings("--drafter: none yet", family(Headless, .{ .dir = "d", .model_type = "m", .drafter = "x" }) orelse ""); } -test "--learn-conversations opens Nemotron and a family whose states declare a disk image, no other" { +test "--learn-conversations opens Nemotron, GLM-5.3-Flash and a family whose states declare a disk image, no other" { const Other = struct {}; const Imaged = struct { pub const kept_disk = struct {}; @@ -160,6 +160,7 @@ test "--learn-conversations opens Nemotron and a family whose states declare a d try std.testing.expect(family(AllRanks, .{ .dir = "d", .model_type = "m", .tp = 2, .learn_conversations = true }) == null); try std.testing.expectStringStartsWith(family(Other, .{ .dir = "d", .model_type = "m", .learn_conversations = true }) orelse "", "--learn-conversations keeps"); try std.testing.expect(family(nemotron.native, .{ .dir = "d", .model_type = "nemotron_h", .learn_conversations = true }) == null); + try std.testing.expect(family(@import("glm_cuda").native, .{ .dir = "d", .model_type = "glm5_next", .learn_conversations = true }) == null); } test "CUDA MTP controls respect window and probability bounds" { From 52c1e382f1d35dafe4efd70d009d161b1273eb8e Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski <24755+wojo@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:31:02 -0400 Subject: [PATCH 3/7] rows: a pool that reuses spare storage without trimming it reads a disk state into that storage A family whose kept states live in preallocated storage it reuses (GLM's snapshot pages on CUDA) and cannot trim never had a disk read chosen: the selection counted its spare storage as fixed and its kept states as not reclaimable, so no state ever fit. Such a family's room for a read is now its spare storage plus what evicting unborrowed states frees, and a read is ready once the family says its spare storage takes the state, as a keep's admission already asks. Families with trim, or without reused storage, are unchanged. --- zig/src/core/rows_pool_select_test.zig | 50 ++++++++++++++++++++++++++ zig/src/core/rows_read.zig | 8 +++-- zig/src/core/rows_select.zig | 7 ++-- 3 files changed, 60 insertions(+), 5 deletions(-) diff --git a/zig/src/core/rows_pool_select_test.zig b/zig/src/core/rows_pool_select_test.zig index b15e756c8a..499c69976f 100644 --- a/zig/src/core/rows_pool_select_test.zig +++ b/zig/src/core/rows_pool_select_test.zig @@ -86,3 +86,53 @@ test "restore selection: pooled buffers trim only when supported and never steal s.release(borrowed); } } + +test "restore selection: a pool without trim reads into the spare storage an eviction leaves, as a keep would" { + var tmp = t.tmpDir(.{}); + defer tmp.cleanup(); + const dir = try tmp.dir.realPathFileAlloc(t.io, ".", a); + defer a.free(dir); + var prompt: [64]u32 = undefined; + for (&prompt, 0..) |*token, i| token.* = @intCast(100 + i); + const rules: pc.Rules = .{ .min_prompt = 0, .min_gap = 1, .lookahead = 1 }; + { + var f: Fake = .{ .gpa = a, .at = 20, .sum = fresh(prompt[0..20]) }; + var table = f.pooled().vtable.*; + table.rows = &Image.api; + const family: pc.Snapshots = .{ .ptr = &f, .vtable = &table }; + const r = try rows.Rows.open(a, family, options(dir)); + defer r.deinit(); + var s = pc.Store.init(a, family, rules, 1000); + defer s.deinit(); + s.disk = r; + try t.expect(s.keep(&prompt, 20, null, &.{}, &.{})); + while (rows.write.step(r)) {} + } + var f: Fake = .{ .gpa = a, .at = 5, .sum = fresh(prompt[0..5]) }; + var table = f.pooled().vtable.*; + table.rows = &Image.api; + table.trim = null; + const family: pc.Snapshots = .{ .ptr = &f, .vtable = &table }; + const r = try rows.Rows.open(a, family, options(dir)); + defer r.deinit(); + var s = pc.Store.init(a, family, rules, 240); + defer s.deinit(); + try t.expect(s.keep(&prompt, 5, null, &.{}, &.{})); + const borrowed = s.entries.items[0]; + borrowed.borrowers = 1; + prompt[0] = 999; + f.at = 30; // another conversation's state: its storage, once evicted, can take the read (120 >= 110) + f.sum = fresh(prompt[0..30]); + try t.expect(s.keep(&prompt, 30, null, &.{}, &.{})); + prompt[0] = 100; + s.disk = r; + try t.expectEqual(.disk, rows.read.resumable(&s, &prompt, &.{}, &.{}).from); + const restored = rows.read.recall(&s, &prompt, 64, &.{}, &.{}, borrowed).?; + try t.expectEqual(@as(u32, 20), restored.at); + try t.expectEqual(@as(u64, 1), r.counts.reads); + try t.expectEqual(@as(u64, 1), s.counts.evicted); + try t.expect(s.held + s.spare() <= s.budget); + try family.vtable.restore(family.ptr, null, restored.saved); + try t.expectEqual(fresh(&prompt), f.pass(&s, &prompt, .{ .from = restored.at })); + s.release(borrowed); +} diff --git a/zig/src/core/rows_read.zig b/zig/src/core/rows_read.zig index 93d0522664..fd80035952 100644 --- a/zig/src/core/rows_read.zig +++ b/zig/src/core/rows_read.zig @@ -91,9 +91,10 @@ pub fn recall(s: *pc.Store, prompt: []const u32, born: u32, starts: []const u32, const name = r.states.items[chosen].name; const bytes = s.family.vtable.bytes(s.family.ptr, r.states.items[chosen].at); if (bytes > s.budget or s.most == 0) return have; - if (!ready(s, bytes)) present(r, chosen) catch return have; // nothing goes to make room for a state whose file is gone + const at = r.states.items[chosen].at; + if (!ready(s, at, bytes)) present(r, chosen) catch return have; // nothing goes to make room for a state whose file is gone if (r.spill != .off) r.touch(chosen); // the newest: the spills that make room for it push older states off first - while (!ready(s, bytes)) { + while (!ready(s, at, bytes)) { if (s.family.vtable.trim) |trim| { const spare = s.spare(); trim(s.family.ptr, s.budget -| s.held -| bytes); @@ -128,8 +129,9 @@ pub fn recall(s: *pc.Store, prompt: []const u32, born: u32, starts: []const u32, return e; } -fn ready(s: *const pc.Store, bytes: u64) bool { +fn ready(s: *const pc.Store, at: u32, bytes: u64) bool { if (s.held > s.budget or s.entries.items.len >= s.most) return false; + if (s.family.vtable.reuses) |f| if (f(s.family.ptr, at)) return true; // spare storage takes it, as a keep's would const room = s.budget - s.held; const spare = s.spare(); return spare <= room and bytes <= room - spare; diff --git a/zig/src/core/rows_select.zig b/zig/src/core/rows_select.zig index 7477620dd6..a3f1529ef6 100644 --- a/zig/src/core/rows_select.zig +++ b/zig/src/core/rows_select.zig @@ -39,12 +39,15 @@ fn capacity(s: *const pc.Store, have: ?*pc.Entry) ?u64 { var protected: usize = 0; const f = s.family.vtable; for (s.entries.items) |e| { - if (e == have or e.borrowers > 0) protected += 1 else if (f.spare == null or f.trim != null) { + if (e == have or e.borrowers > 0) protected += 1 else if (f.spare == null or f.trim != null or f.reuses != null) { reclaimable +|= e.bytes; } } + if (protected >= s.most) return null; + // a pool that reuses spare storage without trimming it (GLM's snapshot pages): a read takes spare or evicted pages + if (f.reuses != null and f.trim == null) return s.spare() +| reclaimable; const fixed = (s.held -| reclaimable) +| if (f.trim == null) s.spare() else 0; - if (protected >= s.most or fixed > s.budget) return null; + if (fixed > s.budget) return null; return s.budget - fixed; } From 44ccb6adc9fd102598a597b67b7432cc60703e85 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski <24755+wojo@users.noreply.github.com> Date: Tue, 27 Oct 2026 12:06:24 -0400 Subject: [PATCH 4/7] rows: a pool that also trims reads a disk state only within its budget, as before The disk read's early admission for a family whose spare storage takes the state now applies only to a family without trim, as the selection's capacity already does. A family with trim, such as Flash Next's or Qwen3.8-27B's pools on Metal, makes room for a read by trimming its spare storage and evicting, exactly as it did before, and a new test pins that its read stays within the budget when its reuses answers yes. --- zig/src/core/rows_pool_select_test.zig | 62 ++++++++++++++++++++++++++ zig/src/core/rows_read.zig | 4 +- 2 files changed, 65 insertions(+), 1 deletion(-) diff --git a/zig/src/core/rows_pool_select_test.zig b/zig/src/core/rows_pool_select_test.zig index 499c69976f..beab09be11 100644 --- a/zig/src/core/rows_pool_select_test.zig +++ b/zig/src/core/rows_pool_select_test.zig @@ -136,3 +136,65 @@ test "restore selection: a pool without trim reads into the spare storage an evi try t.expectEqual(fresh(&prompt), f.pass(&s, &prompt, .{ .from = restored.at })); s.release(borrowed); } + +/// A disk read into new storage, whatever spare storage the pool holds. +const Fresh = struct { + const api: rows.Image = .{ .parts = Image.parts, .gather = Image.gather, .blank = blank, .scatter = Image.scatter }; + fn blank(ptr: *anyopaque, at: u32) !pc.Saved { + const f: *Fake = @ptrCast(@alignCast(ptr)); + const previous = f.at; + f.at = at; + defer f.at = previous; + return f.snapshots().vtable.save(ptr, null, at); + } +}; + +test "restore selection: a pool that also trims reads a disk state only within its budget, whatever it reuses" { + var tmp = t.tmpDir(.{}); + defer tmp.cleanup(); + const dir = try tmp.dir.realPathFileAlloc(t.io, ".", a); + defer a.free(dir); + var prompt: [64]u32 = undefined; + for (&prompt, 0..) |*token, i| token.* = @intCast(100 + i); + const rules: pc.Rules = .{ .min_prompt = 0, .min_gap = 1, .lookahead = 1 }; + { + var f: Fake = .{ .gpa = a, .at = 20, .sum = fresh(prompt[0..20]) }; + var table = f.pooled().vtable.*; + table.rows = &Image.api; + const family: pc.Snapshots = .{ .ptr = &f, .vtable = &table }; + const r = try rows.Rows.open(a, family, options(dir)); + defer r.deinit(); + var s = pc.Store.init(a, family, rules, 1000); + defer s.deinit(); + s.disk = r; + try t.expect(s.keep(&prompt, 20, null, &.{}, &.{})); + while (rows.write.step(r)) {} + } + var f: Fake = .{ .gpa = a, .at = 5, .sum = fresh(prompt[0..5]) }; + var table = f.pooled().vtable.*; + table.rows = &Fresh.api; + try t.expect(table.reuses != null and table.trim != null); + const family: pc.Snapshots = .{ .ptr = &f, .vtable = &table }; + const r = try rows.Rows.open(a, family, options(dir)); + defer r.deinit(); + var s = pc.Store.init(a, family, rules, 240); + defer s.deinit(); + try t.expect(s.keep(&prompt, 5, null, &.{}, &.{})); + const borrowed = s.entries.items[0]; + borrowed.borrowers = 1; + prompt[0] = 999; + f.at = 30; // another conversation's state: once evicted, its spare storage is one the family says takes the read + f.sum = fresh(prompt[0..30]); + try t.expect(s.keep(&prompt, 30, null, &.{}, &.{})); + prompt[0] = 100; + s.disk = r; + try t.expectEqual(.disk, rows.read.resumable(&s, &prompt, &.{}, &.{}).from); + const restored = rows.read.recall(&s, &prompt, 64, &.{}, &.{}, borrowed).?; + try t.expectEqual(@as(u32, 20), restored.at); + try t.expectEqual(@as(u64, 1), r.counts.reads); + try t.expectEqual(@as(u64, 1), s.counts.evicted); + try t.expect(s.held + s.spare() <= s.budget); // the spare storage was trimmed for the new storage, not left over it + try family.vtable.restore(family.ptr, null, restored.saved); + try t.expectEqual(fresh(&prompt), f.pass(&s, &prompt, .{ .from = restored.at })); + s.release(borrowed); +} diff --git a/zig/src/core/rows_read.zig b/zig/src/core/rows_read.zig index fd80035952..df15800079 100644 --- a/zig/src/core/rows_read.zig +++ b/zig/src/core/rows_read.zig @@ -131,7 +131,9 @@ pub fn recall(s: *pc.Store, prompt: []const u32, born: u32, starts: []const u32, fn ready(s: *const pc.Store, at: u32, bytes: u64) bool { if (s.held > s.budget or s.entries.items.len >= s.most) return false; - if (s.family.vtable.reuses) |f| if (f(s.family.ptr, at)) return true; // spare storage takes it, as a keep's would + const f = s.family.vtable; + // a pool without trim reads into spare storage that takes the state, as a keep would + if (f.trim == null) if (f.reuses) |reuses| if (reuses(s.family.ptr, at)) return true; const room = s.budget - s.held; const spare = s.spare(); return spare <= room and bytes <= room - spare; From acfd5a4c48d74951e0bb1c667ad6c11073944527 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski <24755+wojo@users.noreply.github.com> Date: Tue, 27 Oct 2026 12:12:18 -0400 Subject: [PATCH 5/7] GLM on CUDA: a disk read's seal refuses ring bases that do not hold the state's rows An index ring based at `base` holds the rows from `base` up to the state's position, and the image keeps at most a ring's rows of each. The seal now checks the words' main base against the state's position and the head's base against its rows, so an image whose words name a ring that starts past them, or too far back to be held, is refused before the blank state is made whole. --- zig/src/families/glm/cuda_disk_image_test.zig | 18 ++++++++++++++++++ zig/src/families/glm/cuda_snapshots.zig | 7 +++++++ 2 files changed, 25 insertions(+) diff --git a/zig/src/families/glm/cuda_disk_image_test.zig b/zig/src/families/glm/cuda_disk_image_test.zig index fea37bd00c..847dbd4d67 100644 --- a/zig/src/families/glm/cuda_disk_image_test.zig +++ b/zig/src/families/glm/cuda_disk_image_test.zig @@ -259,6 +259,8 @@ test "GLM disk image: a disk read's blank state is refused until sealed, and its try expectError(error.UnsealedState, snaps.restore(&e, 7)); try expectError(error.DiskWords, snaps.seal(&e, 7, at + 9, .{ 0, 0 }, 0)); // a head past its verify window try expectError(error.DiskWords, snaps.seal(&e, 7, at, .{ 0, 0 }, @as(u64, 1) << @intCast(e.cur.st.cur.len))); // a copy bit no layer has + try expectError(error.DiskWords, snaps.seal(&e, 7, at - 1, .{ at + 1, 0 }, 0)); // a main ring based past the state + try expectError(error.DiskWords, snaps.seal(&e, 7, at - 1, .{ 0, at }, 0)); // a head ring based past its rows try snaps.seal(&e, 7, at - 1, .{ 3, 5 }, 0b01); const s = e.snaps.get(7).?; const sh = try image.Shape.of(&e.cur.st, &e.c, capacity, true); @@ -267,4 +269,20 @@ test "GLM disk image: a disk read's blank state is refused until sealed, and its try expectEqual([2]usize{ 3, 5 }, s.bases); try std.testing.expectEqualSlices(u1, &.{ 1, 0 }, s.cur); try expectError(error.StateKept, snaps.seal(&e, 7, at, .{ 0, 0 }, 0)); + const far: u32 = capacity - 1024; // a state past the rings' rows + const back: u32 = far - @as(u32, @intCast(e.cur.st.ring_rows)); // the least base whose ring still holds its rows + try snaps.blankAs(&e, 8, far); + defer { + var held = e.snaps.get(8).?; + snaps.free(&e, &held); + } + try expectError(error.DiskWords, snaps.seal(&e, 8, far, .{ back - 1, back }, 0)); // a main ring too far back + try expectError(error.DiskWords, snaps.seal(&e, 8, far, .{ back, back - 1 }, 0)); // a head ring too far back + try snaps.seal(&e, 8, far, .{ back, back }, 0); + try snaps.blankAs(&e, 9, far); + defer { + var held = e.snaps.get(9).?; + snaps.free(&e, &held); + } + try snaps.seal(&e, 9, 0, .{ back, 0 }, 0); // no head rows yet (as without drafts): the main ring holds the state's rows } diff --git a/zig/src/families/glm/cuda_snapshots.zig b/zig/src/families/glm/cuda_snapshots.zig index 5640051d41..28e83f5319 100644 --- a/zig/src/families/glm/cuda_snapshots.zig +++ b/zig/src/families/glm/cuda_snapshots.zig @@ -69,6 +69,11 @@ pub fn blankAs(e: anytype, id: u64, at: usize) !void { e.snaps.putAssumeCapacityNoClobber(id, .{ .lease = lease, .len = 0, .pos = at, .mtp_len = 0, .cur = cur, .bases = .{ 0, 0 } }); } +/// Whether an index ring of `rows` rows based at `base` holds every row up to `len`, as the image's ring bytes do. +fn ringHolds(base: usize, len: usize, rows: usize) bool { + return base <= len and len - base <= rows; +} + /// A blank state made whole from its image's words (its bytes are in or follow): a restore takes it from here. pub fn seal(e: anytype, id: u64, mtp_len: usize, bases: [2]usize, cur: u64) !void { const s = e.snaps.getPtr(id) orelse return error.NoKeptState; @@ -76,6 +81,8 @@ pub fn seal(e: anytype, id: u64, mtp_len: usize, bases: [2]usize, cur: u64) !voi if (s.cur.len > 64 or (s.cur.len < 64 and cur >> @intCast(s.cur.len) != 0)) return error.DiskWords; const head = if (e.mtp) @min(e.capacity, s.pos + @import("cuda_state.zig").window_rows) else 0; if (mtp_len > head) return error.DiskWords; + const rows = e.cur.st.ring_rows; + if (!ringHolds(bases[0], s.pos, rows) or !ringHolds(bases[1], mtp_len, rows)) return error.DiskWords; const n = try layout.Plan.size(&e.cur.st, &e.c, e.capacity, s.pos, if (e.mtp) mtp_len else null); if (n == 0 or n > s.lease.ids.items.len * pages.page) return error.SnapshotSize; for (s.cur, 0..) |*b, i| b.* = @intCast(cur >> @intCast(i) & 1); From 226c0f632dac54f87af767414e0eb2392bc99bbd Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski <24755+wojo@users.noreply.github.com> Date: Tue, 27 Oct 2026 17:36:42 -0400 Subject: [PATCH 6/7] rows: the grid test runs on both writers, the test's thread and the I/O worker GLM's store uses The write-behind gave GLM's store its own I/O worker, while the grid test wrote on the test's thread (sync_writes). It now runs on each, drained by write.drain, so a worker path that loses the grid's chunk starts fails the test. --- zig/src/core/rows_test.zig | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/zig/src/core/rows_test.zig b/zig/src/core/rows_test.zig index e809569c9c..d14a8c2b7b 100644 --- a/zig/src/core/rows_test.zig +++ b/zig/src/core/rows_test.zig @@ -247,16 +247,18 @@ fn gridRequest(a: Allocator, s: *pc.Store, f: *Fake, p: []const u32, history: u3 return .{ .from = plan.from, .sum = f.sum }; } -test "a grid family's requests name no chunk starts: its states go to disk on the grid and resume after a restart" { +/// A grid family's states through one writer (`sync`: the test's thread, else the I/O worker), read after a restart. +fn gridRoundTrip(sync: bool, tag: []const u8) !void { const gpa = std.testing.allocator; var arena: std.heap.ArenaAllocator = .init(gpa); defer arena.deinit(); const a = arena.allocator(); var buf: [128]u8 = undefined; - const dir = root(&buf, "grid"); + const dir = root(&buf, tag); defer rmTree(dir); const grid: pc.Rules = .{ .lookahead = 1, .min_prompt = 0, .min_gap = 1, .planned = true, .grid = 512 }; var o = options(dir); + o.sync_writes = sync; o.planned = true; o.grid = 512; const p = try prompt(a, 0, 5, 1300); // the engine cuts its chunks on the grid: 512 and 1024 @@ -268,7 +270,7 @@ test "a grid family's requests name no chunk starts: its states go to disk on th defer s.deinit(); s.disk = r; _ = try gridRequest(a, &s, &f, p, 1200); - drain(r); + _ = rows.write.drain(r, rows.write.nowNs() + 3 * std.time.ns_per_s); // the worker publishes the states try std.testing.expect(r.states.items.len > 0); for (r.states.items) |st| { try std.testing.expectEqual(@as(u32, 0), st.at % 512); @@ -289,6 +291,14 @@ test "a grid family's requests name no chunk starts: its states go to disk on th try std.testing.expectEqual(@as(u64, 1), r.counts.reads); } +test "a grid family's requests name no chunk starts: its states go to disk on the grid and resume after a restart" { + try gridRoundTrip(true, "grid"); +} + +test "a grid family's chunk starts reach the I/O worker's writes as well, the writer GLM's store uses" { + try gridRoundTrip(false, "grid-worker"); +} + test "an entry evicted before its write leaves the queue; once written, an eviction costs nothing" { const gpa = std.testing.allocator; var arena: std.heap.ArenaAllocator = .init(gpa); From a2066d29a6b3a895bc4286f8583b9ca26a4fe802 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski <24755+wojo@users.noreply.github.com> Date: Wed, 28 Oct 2026 09:56:19 -0400 Subject: [PATCH 7/7] GLM on CUDA: --spill-gib and --checkpoint-slots are refused, as --learn-conversations alone keeps its states Since #1413 CUDA serves --spill-gib and --checkpoint-slots, and --learn-conversations admits GLM, so --learn-conversations --spill-gib (or --checkpoint-slots) opened GLM and ignored the flag: its disk store writes every rank's part between passes and stages nothing at eviction. GLM now refuses both by name. Its store's spill stays off, since an eviction would gather on rank 0 mid-pass while the other ranks wait in the gate's all-reduce. docs/api.md lists every CUDA family that keeps states on disk. --- docs/api.md | 5 +++-- zig/src/native/cuda_glm_disk.zig | 1 + zig/src/native/cuda_refusal.zig | 13 +++++++++++-- 3 files changed, 15 insertions(+), 4 deletions(-) diff --git a/docs/api.md b/docs/api.md index 9ce5d3280f..a30bb1e985 100644 --- a/docs/api.md +++ b/docs/api.md @@ -602,8 +602,9 @@ store for each model and build. That store takes the place of `--learn`'s store with the flag neither reads nor writes the prefix files a plain `--learn` server keeps, and they stay on disk until the cap below needs their room. After a restart, a crash or an eviction from RAM, the next turn reads the latest state back instead of reading the conversation again from its first token. Qwen3.8-27B, Nemotron, Flash Next and GLM-5.3-Flash on -one Mac, and Nemotron and GLM-5.3-Flash on CUDA keep them; other engines refuse the flag at startup. GLM-5.3-Flash on -CUDA across ranks writes every rank's part to rank 0's disk, so a state is read back on every rank or on none. +one Mac, and Nemotron, the Qwen3.5 hybrids, Flash Next on one rank and GLM-5.3-Flash on CUDA keep them; other engines +refuse the flag at startup. GLM-5.3-Flash on CUDA across ranks writes every rank's part to rank 0's disk, so a state is +read back on every rank or on none; it takes neither `--spill-gib` nor `--checkpoint-slots`. The engine stages an owned copy between rounds, and an I/O thread writes and fsyncs it. Staging has a 1 GiB byte permit and at most eight jobs; when either is full, the save is skipped. The copy waits for its GPU producers before diff --git a/zig/src/native/cuda_glm_disk.zig b/zig/src/native/cuda_glm_disk.zig index 8ddc5ba09a..33e0f8fa06 100644 --- a/zig/src/native/cuda_glm_disk.zig +++ b/zig/src/native/cuda_glm_disk.zig @@ -174,6 +174,7 @@ pub fn attach(a: Allocator, io: std.Io, ctx: *const cuda.Context, o: api.Open, r return @intFromFloat(g * (1 << 30)); } }.bytes; + // spill stays off: an eviction would gather on rank 0 mid-pass while the other ranks wait in the gate's all-reduce const r = try rows.Rows.open(store.gpa, store.family, .{ .root = root, .identity = id, .cap = gibs(o.learn_gib), .floor = gibs(o.learn_min_free_gib), .daily = gibs(o.learn_daily_gib), .lookahead = store.rules.lookahead, .planned = store.rules.planned, .grid = store.rules.grid }); r.prefill_rate = p.rate; store.disk = r; diff --git a/zig/src/native/cuda_refusal.zig b/zig/src/native/cuda_refusal.zig index 70db25c1e9..af00049559 100644 --- a/zig/src/native/cuda_refusal.zig +++ b/zig/src/native/cuda_refusal.zig @@ -31,7 +31,7 @@ pub fn lfm2(o: api.Open) ?[]const u8 { const lfm2_vision = "--vision: LFM2 reads the checkpoint's images through its own tower without the flag; " ++ "serve it without --vision"; -/// A family takes --drafter only if it says so, and --learn-conversations only as Nemotron; null: `o` may open `F`. +/// A family takes --drafter only if it says so, and --learn-conversations only as Nemotron, GLM or a kept_disk family. pub fn family(comptime F: type, o: api.Open) ?[]const u8 { if (o.vision != null and !(@hasDecl(F, "Options") and @hasField(F.Options, "vision"))) return "--vision: this CUDA family reads no images; serve it without --vision"; @@ -43,6 +43,8 @@ pub fn family(comptime F: type, o: api.Open) ?[]const u8 { } // ahead of the drafter's and --learn-conversations' refusals (--spill-gib sets that flag), on a line of its own if (o.spill == .evicted and F != nemotron.native and @import("cuda_disk.zig").keptDisk(F) == null) return "--spill-gib keeps evicted states on disk for Nemotron, the Qwen hybrids and Flash Next only on CUDA: serve this model without it"; + // GLM's disk store writes every rank's part between passes and stages nothing at eviction (cuda_glm_disk.attach) + if (F == @import("glm_cuda").native and (o.spill != .off or o.keep != null)) return "--spill-gib and --checkpoint-slots: GLM-5.3-Flash on CUDA keeps states on disk with --learn-conversations alone; serve it without them"; if (o.drafter != null and !(@hasDecl(F, "takes_drafter") and F.takes_drafter)) { if (@hasDecl(F, "drafter_refusal")) return F.drafter_refusal; // a family without an MTP head says so itself return "--drafter loads a separate drafter, and this CUDA family drafts with its own MTP head: serve without --drafter"; @@ -202,13 +204,20 @@ test "prompt precision: int8 is refused unless a CUDA family consumes the option try std.testing.expect(family(Bf16, .{ .dir = "d", .model_type = "m" }) == null); } -test "--spill-gib and --checkpoint-slots open the families with a disk image; the rest refuse them by name" { +test "--spill-gib and --checkpoint-slots open Nemotron and the kept_disk families; GLM and the rest refuse them by name" { const spilled: api.Open = .{ .dir = "d", .model_type = "m", .learn = "/s", .learn_conversations = true, .spill = .evicted, .keep = 2 }; try std.testing.expect(backend(spilled) == null); try std.testing.expect(family(nemotron.native, spilled) == null); try std.testing.expect(family(@import("qwen35").native, spilled) == null); try std.testing.expect(family(@import("qwen35").native_moe, spilled) == null); try std.testing.expect(family(@import("flashnext").native, spilled) == null); + const glm = @import("glm_cuda").native; // its states go to disk with --learn-conversations alone + try std.testing.expectStringStartsWith(family(glm, spilled) orelse "", "--spill-gib keeps evicted states"); + const finish: api.Open = .{ .dir = "d", .model_type = "glm5_next", .learn = "/s", .learn_conversations = true, .spill = .finish }; + try std.testing.expectStringStartsWith(family(glm, finish) orelse "", "--spill-gib and --checkpoint-slots: GLM"); + const slots: api.Open = .{ .dir = "d", .model_type = "glm5_next", .learn = "/s", .learn_conversations = true, .keep = 2 }; + try std.testing.expectStringStartsWith(family(glm, slots) orelse "", "--spill-gib and --checkpoint-slots: GLM"); + try std.testing.expect(family(glm, .{ .dir = "d", .model_type = "glm5_next", .learn = "/s", .learn_conversations = true }) == null); const Other = struct {}; try std.testing.expectStringStartsWith(family(Other, spilled) orelse "", "--spill-gib keeps evicted states"); try std.testing.expectStringStartsWith(backend(.{ .dir = "d", .model_type = "m", .keep = 2 }) orelse "", "--checkpoint-slots sets how many");