cuda: serve --spill-gib and --checkpoint-slots for Nemotron, the Qwen hybrids and Flash Next (#840, split from #1359)

docs/api.md

@@ -637,14 +637,24 @@ Older queued saves cannot evict states saved earlier in that close. States that
 daily budget or free-disk floor are skipped; unfinished writes are cancelled at the deadline.
 Past N GiB other models' spilled sets go first, then the state this server used least recently. A damaged file is
 deleted and its tokens are prefilled, and a write cut off midway leaves nothing a later turn or a restart reads.
-`--checkpoint-slots N` sets how many states stay in RAM beside the store, so the next one sends the oldest to disk; on a
-Mac it needs `--spill-gib` or `--learn-conversations`. With `--learn-conversations` too, every state is still written
+`--checkpoint-slots N` sets how many states stay in RAM beside the store, so the next one sends the oldest to disk; it
+needs `--spill-gib` or `--learn-conversations`. With `--learn-conversations` too, every state is still written
 behind, a staged write survives eviction of its RAM entry, and the cap is the smaller of `--learn-gib` and
 `--spill-gib`. `--learn-min-free-gib` and `--learn-daily-gib` apply to spilled states as well. The same four Mac
-families take both flags and the others refuse them at startup; without either flag nothing changes, and the routes
-below stay off. `TENSORFOLD_ROWS_SYNC=1` selects the previous synchronous writer for comparisons; its block and state
-files have the same bytes. That comparison writer retains blocking I/O; the bounded disk wait uses the threaded writer.
-`TF_IMPRINT_VERIFY=1` reads threaded writes back against their owned staging before publication.
+families take both flags, and so do three CUDA ones; the others refuse them at startup. Without either flag nothing
+changes, and the routes below stay off. `TENSORFOLD_ROWS_SYNC=1` selects the previous synchronous writer for
+comparisons; its block and state files have the same bytes. That comparison writer retains blocking I/O; the bounded
+disk wait uses the threaded writer. `TF_IMPRINT_VERIFY=1` reads threaded writes back against their owned staging before
+publication.
+
+On CUDA, Nemotron, the Qwen3.5 hybrids and Flash Next take `--spill-gib` and `--checkpoint-slots` with the same store
+and cap as a Mac. At an eviction the engine thread stages a copy of the state into host memory, waiting for the GPU
+work that produced it, then releases its RAM entry and goes on; the I/O thread writes and fsyncs the copy while the
+engine decodes. Staging holds at most 1 GiB and eight jobs: a state whose new bytes (rows not already on disk) do not
+fit what is free is not saved, so one past 1 GiB never is, and a later request for it prefills. When the engine closes,
+at a planned stop (SIGTERM or SIGINT) or an idle unload, it saves the states still in RAM, newest first, within the same
+disk cap and 20-second stop budget. While other streams run, a queued request whose state is on disk has it read ahead
+on an I/O thread, as described above.
 
 Three routes let a client such as the Mac app ask about a conversation. Each takes the body of a chat completion
 (`messages`, `tools`, `chat_template_kwargs`), renders it as that request would be rendered, and answers on the engine's

zig/build/cuda.zig

@@ -330,6 +330,9 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st
     for ([_]*std.Build.Module{ native, mods.nemotron }) |m| safety.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m, .filters = &.{"kv safety:"} })).step);
     b.step("test-cache-capacity-host", "Startup cache capacity facts and warnings without a GPU").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = native, .filters = &.{"cache capacity:"} })).step);
     b.step("test-snap-fork-host", "CUDA prompt-cache rules: no fork marks unless a family sets snap_fork, its marks when it does").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = native, .filters = &.{"snap fork:"} })).step);
+    const disk_open = b.addRunArtifact(b.addTest(.{ .root_module = native, .filters = &.{"CUDA disk open:"} }));
+    b.step("test-cuda-disk-open-host", "CUDA's disk store takes --spill-gib's mode and --checkpoint-slots").dependOn(&disk_open.step);
+    step.dependOn(&disk_open.step);
     const merge_checks = b.step("test-kv-merge-host", "KV page integration with landed snapshot modes and budgets");
     for ([_]*std.Build.Module{ native, mods.nemotron }) |m| merge_checks.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m, .filters = &.{"kv merge:"} })).step);
     const nemo_safety = b.step("test-nemotron-learning-host", "Actual Nemotron CUDA close, checked pocket reads and retry under fake driver failures");

zig/src/core/engine_api.zig

@@ -319,7 +319,7 @@ pub const Open = struct {
     learn_conversations: bool = false,
     /// --learn-daily-gib: what learned states may write a day under the root, every model together.
     learn_daily_gib: f64 = 32,
-    /// --spill-gib (Macs): states evicted from RAM go to the store `learn_conversations` opens (off: written behind).
+    /// --spill-gib: states evicted from RAM go to the store `learn_conversations` opens (off: written behind).
     spill: prompt_cache.rows.Spill = .off,
     /// --slide: Sliding Weights learns into the served weights, live (off: learn requests are refused).
     slide: bool = false,

zig/src/native/cuda.zig

@@ -32,7 +32,7 @@ const registry = .{ nemotron.native, nemotron.omni, glm.native, qwen35.native, q
 /// --record opens here (engines.zig refuses it for a backend without this).
 pub const records = true;
 pub const backends: []const []const u8 = &.{"cuda"};
-pub const serves: []const []const u8 = &.{ "--mtp-drafts", "--mtp-confidence", "--policy" };
+pub const serves: []const []const u8 = &.{ "--mtp-drafts", "--mtp-confidence", "--policy", "--checkpoint-slots", "--spill-gib" };
 pub const families: []const api.Family = blk: {
     var out: [registry.len + 4]api.Family = undefined;
     for (registry, 0..) |F, i| out[i] = .{ .model_type = F.model_type, .formats = F.formats };
@@ -536,6 +536,7 @@ test {
     _ = @import("cuda_vision_offer.zig");
     _ = kv_pages;
     _ = drafted;
+    _ = cuda_disk; // its tests run under test-cuda-disk-open-host's filter
 }
 test "Learner ownership: actual native close drains Nemotron before any source, trainer or parent free" {
     try @import("nemotron_trainer_test.zig").run(Host, Host.serve, Host.close, false);

zig/src/native/cuda_disk.zig

@@ -180,14 +180,15 @@ fn open(o: api.Open, root: []const u8, id: u64, rate: f64, store: *api.prompt_ca
             return @intFromFloat(g * (1 << 30));
         }
     }.bytes;
-    const r = try api.prompt_cache.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 });
+    // --spill-gib: evicted states only (or with --learn-conversations, both); --checkpoint-slots bounds RAM beside it
+    const r = try api.prompt_cache.rows.attach(store.gpa, store, .{ .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, .spill = o.spill, .slots = @import("disk_spill.zig").slots(o) });
     r.prefill_rate = rate;
-    store.disk = r;
     std.log.info("{s}: conversation states on disk in {s} ({d} states, {d} MiB of {d} MiB)", .{ name, r.dir, r.states.items.len, r.held() >> 20, r.cap >> 20 });
 }
 
 test {
     std.testing.refAllDecls(@This());
+    _ = @import("disk_spill.zig");
 }
 
 test "a disk module that declares attach opens the store itself" {
@@ -201,3 +202,41 @@ test "a disk module that declares attach opens the store itself" {
     try attachKept(K, "k", std.testing.allocator, std.testing.io, undefined, .{ .dir = "d", .model_type = "m" }, "/r", null, @as(u8, 0), undefined);
     try std.testing.expectEqual(@as(usize, 1), K.calls);
 }
+
+test "CUDA disk open: --spill-gib's mode and --checkpoint-slots reach the store" {
+    const gpa = std.testing.allocator;
+    const io = std.testing.io;
+    // a disk module that opens the store itself, through this file's open, under a fixed identity and no device
+    const K = struct {
+        pub fn attach(_: Allocator, _: std.Io, _: *const cuda.Context, o: api.Open, root: []const u8, _: ?[]const u8, _: anytype, store: *api.prompt_cache.Store) !void {
+            try open(o, root, 7, 1.0, store, "fake");
+        }
+        fn bytes(_: *anyopaque, _: u32) u64 {
+            return 0;
+        }
+        fn save(_: *anyopaque, _: ?*anyopaque, _: u32) anyerror!api.prompt_cache.Saved {
+            return error.Unused;
+        }
+        fn restore(_: *anyopaque, _: ?*anyopaque, _: api.prompt_cache.Saved) anyerror!void {}
+        fn drop(_: *anyopaque, _: api.prompt_cache.Saved) void {}
+    };
+    var tmp = std.testing.tmpDir(.{});
+    defer tmp.cleanup();
+    const root = try tmp.dir.realPathFileAlloc(io, ".", gpa);
+    defer gpa.free(root);
+    var family: u8 = 0;
+    const vt: api.prompt_cache.Snapshots.VTable = .{ .bytes = K.bytes, .save = K.save, .restore = K.restore, .drop = K.drop, .rows = &image };
+    for ([_]api.prompt_cache.rows.Spill{ .evicted, .finish }, [_]?i64{ 3, null }) |mode, keep| {
+        var store = api.prompt_cache.Store.init(gpa, .{ .ptr = &family, .vtable = &vt }, .{ .lookahead = 1 }, 1 << 20);
+        const o: api.Open = .{ .dir = "d", .model_type = "nemotron_h", .keep = keep, .learn = root, .learn_conversations = true, .spill = mode };
+        attachKept(K, "k", gpa, io, undefined, o, root, null, @as(u8, 0), &store) catch |e| {
+            store.deinit();
+            return e;
+        };
+        const r = store.disk.?;
+        defer r.deinit();
+        defer store.deinit(); // the store before its disk, as the rows fixtures close them
+        try std.testing.expectEqual(mode, r.spill);
+        try std.testing.expectEqual(if (keep) |n| @as(usize, @intCast(n)) else std.math.maxInt(usize), store.most);
+    }
+}

zig/src/native/cuda_refusal.zig

@@ -16,6 +16,7 @@ pub fn backend(o: api.Open) ?[]const u8 {
             return "--mtp-drafts: this CUDA family's drafts are not tunable; serve it without --mtp-drafts or --mtp-confidence";
     }
     if (o.policy.len > 0) return "--policy is not supported by this CUDA build";
+    if (@import("disk_spill.zig").refusal(o, "")) |why| return why;
     if (o.learn != null and !o.learn_conversations) return "--learn learns shared prompt prefixes on Macs only; CUDA keeps conversation states on disk with --learn-conversations: serve with it, or without --learn, --learn-dir or --learn-gib";
     if (o.speed_up != null) return "--speed-up links two Macs: serve without --speed-up on CUDA";
     if (o.prompt_cache_over_cap) return "--prompt-cache-over-cap lifts a Mac's 70% cap: on CUDA the prompt cache takes at most the memory the weights and streams leave, so serve without it";
@@ -40,6 +41,8 @@ pub fn family(comptime F: type, o: api.Open) ?[]const u8 {
         if (!F.tp_draft_snapshots and o.tp > 1 and o.drafts and (o.prompt_cache_gib orelse 1) > 0)
             return "Flash Next tensor-rank MTP snapshots are unavailable: use --prompt-cache-gib 0 with --tp and drafting";
     }
+    // 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";
     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";
@@ -197,3 +200,16 @@ test "prompt precision: int8 is refused unless a CUDA family consumes the option
     try std.testing.expect(family(Int8, request) == null);
     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" {
+    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 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");
+    try std.testing.expect(backend(.{ .dir = "d", .model_type = "m" }) == null);
+}

zig/src/native/disk_spill.zig

@@ -0,0 +1,32 @@
+//! --spill-gib and --checkpoint-slots on any backend: the RAM store's slots beside the disk store, the refusals' words.
+const std = @import("std");
+const api = @import("engine_api");
+
+/// --checkpoint-slots: the prompt states the RAM store keeps at most beside its disk store (null: the family's plan).
+pub fn slots(o: api.Open) ?usize {
+    return if (o.keep) |n| @intCast(@max(n, 0)) else null;
+}
+
+/// --checkpoint-slots bounds what stays in RAM beside the disk store: without one it is refused (`where`: " on Macs").
+pub fn refusal(o: api.Open, comptime where: []const u8) ?[]const u8 {
+    if (o.keep != null and !o.learn_conversations) return "--checkpoint-slots sets how many prompt states stay in RAM beside the disk store" ++ where ++ ": serve with --spill-gib N (or --learn-conversations), or without it";
+    return null;
+}
+
+/// The flag a family without the disk store names in its refusal: --spill-gib when that alone asked for the store.
+pub fn flag(o: api.Open) []const u8 {
+    return if (o.spill == .evicted) "--spill-gib" else "--learn-conversations";
+}
+
+test "--checkpoint-slots takes the disk store's company, and a refusal names the flag that asked" {
+    const alone: api.Open = .{ .dir = "d", .model_type = "nemotron_h", .keep = 2 };
+    try std.testing.expectEqualStrings("--checkpoint-slots sets how many prompt states stay in RAM beside the disk store on Macs: serve with --spill-gib N (or --learn-conversations), or without it", refusal(alone, " on Macs") orelse "");
+    try std.testing.expectStringStartsWith(refusal(alone, "") orelse "", "--checkpoint-slots sets how many prompt states stay in RAM beside the disk store: serve");
+    try std.testing.expect(refusal(.{ .dir = "d", .model_type = "nemotron_h" }, "") == null);
+    const spilled: api.Open = .{ .dir = "d", .model_type = "nemotron_h", .keep = 2, .learn = "/s", .learn_conversations = true, .spill = .evicted };
+    try std.testing.expect(refusal(spilled, "") == null);
+    try std.testing.expectEqual(@as(?usize, 2), slots(spilled));
+    try std.testing.expectEqual(@as(?usize, null), slots(.{ .dir = "d", .model_type = "nemotron_h" }));
+    try std.testing.expectEqualStrings("--spill-gib", flag(spilled));
+    for ([_]api.prompt_cache.rows.Spill{ .off, .finish }) |m| try std.testing.expectEqualStrings("--learn-conversations", flag(.{ .dir = "d", .model_type = "t", .spill = m }));
+}

zig/src/native/metal_spill.zig

@@ -1,16 +1,16 @@
 //! --spill-gib and --checkpoint-slots on Macs: what the disk store's families take, and the refusal without one.
 const std = @import("std");
 const api = @import("engine_api");
+const neutral = @import("disk_spill.zig"); // CUDA's too
 
 /// --checkpoint-slots: the prompt states the RAM store keeps at most beside its disk store (null: the family's plan).
 pub fn slots(o: api.Open) ?usize {
-    return if (o.keep) |n| @intCast(@max(n, 0)) else null;
+    return neutral.slots(o);
 }
 
 /// --checkpoint-slots bounds what stays in RAM beside the disk store, so a Mac refuses it without one.
 pub fn refusal(o: api.Open) ?[]const u8 {
-    if (o.keep != null and !o.learn_conversations) return "--checkpoint-slots sets how many prompt states stay in RAM beside the disk store on Macs: serve with --spill-gib N (or --learn-conversations), or without it";
-    return null;
+    return neutral.refusal(o, " on Macs");
 }
 
 /// Whether a disk store writes states behind (--learn, --learn-conversations), not only those RAM evicts (--spill-gib).
@@ -20,7 +20,7 @@ pub fn writesBehind(o: api.Open) bool {
 
 /// The flag a family without the disk store names in its refusal: --spill-gib when that alone asked for the store.
 pub fn flag(o: api.Open) []const u8 {
-    return if (o.spill == .evicted) "--spill-gib" else "--learn-conversations";
+    return neutral.flag(o);
 }
 
 test "--checkpoint-slots on a Mac takes the disk store's company, and a refusal names the flag that asked" {