metal nemotron: TF_KV_SHARE keeps prompt states as core pool pages (stacked on #1837)

CHANGELOG.md

@@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel
 
 ## Unreleased
 
+- Nemotron on Metal: with `TF_NEMOTRON_METAL_KV_PAGES=1`, `TF_KV_SHARE=n` keeps prompt states in the core `prompt_radix` store over the KV pool, as CUDA does; `TF_KV_SHARE` acts only on the paged pool, and unset or 0 keeps the host prompt cache (#1850). Thanks to @jschmied.
+
 - GLM CUDA at TP=2 and up: each rank names its startup phases (NCCL init, peer warm-up, weights, calibration, ready), prints NCCL's init device, version and log path, and warns once a minute while NCCL init or the warm-up is unfinished, with the transport seen so far. It only reports; transfers are unchanged (#2354, part of #2347).
 
 - Docs: the serve flags match the CLI: `--metrics-open` opens `/v1/metrics` too, `--learn-conversations` on CUDA serves the Qwen3.5 hybrids and Flash Next as well as Nemotron, Flash Next shares sessions under an explicit `--parallel` above one and refuses `--recall` only on a Mac, and the Metal Qwen3.5-2B keeps prompt states (#2367).

zig/build/drama3_metal.zig

@@ -4,15 +4,15 @@ const std = @import("std");
 const dir = "zig/kernels/metal/drama3";
 const sources = [_][]const u8{ "dense.metal", "gdn.metal", "attn.metal", "moe.metal", "nvfp4.metal", "guidance.metal", "sample.metal", "codec.metal" };
 
-/// Drama in the Metal engines: its `core` is the engine's safetensors, the server's tokenizer and frame_rounds.
+/// Drama in the Metal engines: its `core` is the engine's safetensors and KV pool, the tokenizer and frame_rounds.
 pub fn addNative(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, metal: *std.Build.Module, native: *std.Build.Module, metal_std: []const u8) void {
     const tokenizer = native.import_table.get("tokenizer") orelse tok: {
         const t = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true });
         native.addImport("tokenizer", t);
         break :tok t;
     };
     const rounds = b.createModule(.{ .root_source_file = b.path("zig/src/core/frame_rounds.zig"), .target = target, .optimize = optimize });
-    const shim = b.addWriteFiles().add("drama3_core.zig", "pub const safetensors = @import(\"tensorfold\").safetensors;\npub const tokenizer = @import(\"tokenizer\");\npub const frame_rounds = @import(\"frame_rounds\");\n");
+    const shim = b.addWriteFiles().add("drama3_core.zig", "pub const safetensors = @import(\"tensorfold\").safetensors;\npub const kv_pool = @import(\"tensorfold\").kv_pool;\npub const kv_prompts = @import(\"tensorfold\").kv_prompts;\npub const tokenizer = @import(\"tokenizer\");\npub const frame_rounds = @import(\"frame_rounds\");\n");
     const core = b.createModule(.{ .root_source_file = shim, .target = target, .optimize = optimize, .imports = &.{
         .{ .name = "tensorfold", .module = native.import_table.get("tensorfold").? },
         .{ .name = "tokenizer", .module = tokenizer },

zig/src/core/metal_kv_pages.zig

@@ -179,6 +179,17 @@ pub const Pages = struct {
         p.tables.set(lane, i, id);
     }
 
+    /// Lane `lane`'s table starts with a kept prompt's pages, held once more (TF_KV_SHARE); a failure releases them.
+    pub fn adopt(p: *Pages, lane: usize, held: []const u32) !void {
+        const old = p.leases[lane].ids.items.len;
+        if (old + held.len > p.tables.most) {
+            for (held) |id| p.pool.release(id);
+            return error.PastTable;
+        }
+        try p.pool.adopt(&p.leases[lane], held);
+        for (held, old..) |id, i| p.tables.set(lane, i, id);
+    }
+
     /// Lane `lane` keeps the pages of its first `tokens` rows; those past them go back to its promise.
     pub fn trim(p: *Pages, lane: usize, tokens: u64) void {
         p.pool.trim(&p.leases[lane], p.layout.pages(tokens));
@@ -314,3 +325,22 @@ test "the kernels' address helper pages by TF_PAGE_TOKENS" {
 test {
     _ = @import("metal_kv_pages_test.zig");
 }
+
+test "a kept prompt's pages head a new lane's table, held by both, and stay the store's when the stream ends" {
+    const p = try Pages.init(std.testing.allocator, null, testLayout(), 2, 6, 4 * page_tokens, 0);
+    defer p.deinit();
+    var kept: [2]u32 = undefined;
+    try std.testing.expect(p.pool.hold(&kept)); // the store's two pages
+    for (kept) |id| p.pool.retain(id); // a resume holds them once more for the stream
+    try p.admit(1, 4 - kept.len);
+    try p.adopt(1, &kept);
+    try p.cover(1, 3 * page_tokens); // the stream's own pages past them, from its promise
+    try std.testing.expectEqualSlices(u32, &kept, p.tables.entries(1, 2));
+    try std.testing.expectEqualSlices(u32, p.ids(1), p.tables.entries(1, 3));
+    try std.testing.expectEqual(@as(u32, 2), p.pool.holders(kept[0]));
+    p.finish(1);
+    try std.testing.expect(p.idle()); // kept pages are the store's, not a lane's
+    try std.testing.expectEqual(@as(usize, 2), p.pool.census().kept);
+    for (kept) |id| try p.pool.drop(id);
+    try std.testing.expectEqual(@as(usize, 6), p.pool.census().free);
+}

zig/src/families/nemotron/backend.zig

@@ -9,7 +9,7 @@ const st = @import("state.zig");
 const kv_pages = @import("kv_pages.zig");
 const kvp = @import("../../core/metal_kv_pages.zig");
 const lesson_route = @import("living").lesson_route;
-const layers = @import("layers.zig");
+const prompt_pages = @import("prompt_pages.zig");
 const mtp = @import("mtp.zig");
 const timing = @import("timing.zig");
 const tree = @import("tree.zig");
@@ -82,6 +82,8 @@ pub const Metal = struct {
     kv_tables: ?mtl.Buffer = null, // and every lane's page table
     kv_pages: ?*kvp.Pages = null, // the core pool over them: a stream leases a lane's pages when it opens
     kv_pipes: ?*kv_pages.Pipelines = null, // the paged twins (the model's kernels point at them)
+    prompts: ?@import("../../core/kv_prompts.zig").Prompts = null, // TF_KV_SHARE: the store keeping prompts as pages
+    pending: ?*const snapshot.Snap = null, // the store's resume, restored once the stream's lane adopts its pages
 
     pub fn init(gpa: std.mem.Allocator, m: *Model, o: Options) !*Metal {
         const b = try gpa.create(Metal);
@@ -240,7 +242,10 @@ pub const Metal = struct {
             self.gpa.destroy(old);
             _ = self.caches.remove(s);
         }
-        const c = try @import("residency.zig").openCache(self, s, rows_needed);
+        var pass = try prompt_pages.Pass.begin(self, s);
+        var ok = false;
+        defer pass.end(self, s, ok);
+        const c = try @import("residency.zig").openCache(self, s, rows_needed, pass.held());
         @memcpy(self.prompt.slice(u32, ids.len), ids);
         const Chunks = struct {
             n: usize,
@@ -280,7 +285,7 @@ pub const Metal = struct {
             }
         };
         s.cached = 0;
-        var at: usize = 0;
+        var at = try pass.resume_(self, s, c); // TF_KV_SHARE: the store's kept pages and state
         if (s.reuse.saved) |saved| { // a kept state of this prompt's prefix: the pass starts there
             const snap: *snapshot.Snap = @ptrCast(@alignCast(saved));
             if (snap.at < ids.len) if (snapshot.restore(self, c, snap)) |_| {
@@ -294,31 +299,18 @@ pub const Metal = struct {
         var k: usize = 0;
         while (at < ids.len) {
             const decision: ?usize = if (c.decide == lesson_route.none) null else @intCast(c.decide);
-            const rows = @import("prompt_flight.zig").rows(s, at, self.o.chunk, decision, &k);
+            const rows = @import("prompt_flight.zig").rows(s, at, self.o.chunk, decision, pass.marks, &k); // TF_KV_SHARE: and its marks
             if (s.isCancelled()) return error.Cancelled; // release() waits for the chunk still in flight
             try kv_pages.grow(c, at + rows);
             try @import("prompt_flight.zig").submit(self, s, Chunks{ .n = ids.len, .c = c, .at = at, .rows = rows });
             at += rows;
             // landed tokens only: the newest chunk is still queued on the GPU
             s.fill_at.store(@intCast(at - rows), .monotonic);
             if (s.reuse.hook) |hk| if (std.mem.indexOfScalar(u32, s.reuse.marks, @intCast(at)) != null) hk.at(hk.ptr, s, @intCast(at));
+            pass.mark(self, s, c, at);
         }
-        if (self.trace and c.decide != lesson_route.none) try self.traceLessons(c);
-    }
-
-    /// Dev trace (TF_SLIDE_TRACE): the stream's router scores at its decision row and the lessons acting on its reply.
-    fn traceLessons(self: *Metal, c: *const st.Cache) !void {
-        const ad = layers.firstAdapter(&self.m.weights) orelse return;
-        try self.drain();
-        const n = lesson_route.max_blocks;
-        const at: usize = @intCast(@max(c.rid, 0));
-        const score = ad.score.slice(f32, lesson_route.max_slots * n)[at * n ..][0..n];
-        const flags = ad.latch.slice(u32, lesson_route.max_slots * n)[at * n ..][0..n];
-        const need = ad.need.slice(f32, n);
-        var line: [2048]u8 = undefined;
-        var w: std.Io.Writer = .fixed(&line);
-        for (0..ad.rank / 16) |b| w.print(" {d:.2}/{d:.2}{s}", .{ score[b], need[b], if (flags[b] != 0) "*" else "" }) catch break;
-        std.log.info("slide trace: decision row {d}: scores{s} (* acting)", .{ c.decide, w.buffered() });
+        ok = true;
+        if (self.trace and c.decide != lesson_route.none) try @import("slide.zig").traceLessons(self, c);
     }
 
     fn firstFn(ptr: *anyopaque, s: *lanes.Stream, position: u64) anyerror!u64 {

zig/src/families/nemotron/kv_pages.zig

@@ -64,12 +64,12 @@ pub inline fn holds(kv: st.Kv, end: usize) void {
     std.debug.assert(end <= p.ids().len * kvp.page_tokens);
 }
 
-/// With no stream open every page must be free: the census says so, or warns.
+/// With no stream open every page is free or the store's (TF_KV_SHARE): the census says so, or warns.
 pub fn census(b: anytype) void {
     const p = b.kv_pages orelse return;
     const c = p.pool.census();
-    if (p.idle() and c.free == c.total) {
-        std.log.info("nemotron: KV pages idle, {d} of {d} free", .{ c.free, c.total });
+    if (p.idle() and (b.prompts != null or c.free == c.total)) {
+        std.log.info("nemotron: KV pages idle, {d} of {d} free, {d} kept", .{ c.free, c.total, c.kept });
     } else std.log.warn("nemotron: no stream open, yet {d} of {d} KV pages are not free", .{ c.total - c.free, c.total });
 }
 

zig/src/families/nemotron/learned_prompts.zig

@@ -77,7 +77,7 @@ pub fn probe(b: *Metal) !u64 {
                 p.failed = e;
                 return;
             };
-            p.snap = snapshot.save(p.b, c, mark) catch |e| {
+            p.snap = snapshot.save(p.b, c, mark, false) catch |e| {
                 p.failed = e;
                 return;
             };
@@ -116,7 +116,7 @@ pub fn snapRead(ptr: *anyopaque, dir: [:0]const u8, key: u64, at: u32) anyerror!
     var buf: [1100]u8 = undefined;
     const file = try path(&buf, dir, key);
     const head = b.head != null;
-    const n = snapshot.bytes(b.m.config, at, head);
+    const n = snapshot.bytes(b.m.config, at, head, false);
     const snap = try b.gpa.create(Snap);
     errdefer b.gpa.destroy(snap);
     snap.* = .{ .at = at, .head = head, .buf = try b.m.device.buffer(n, opts), .bytes = n };

zig/src/families/nemotron/prompt_flight.zig

@@ -3,13 +3,16 @@ const lanes = @import("lanes");
 
 pub const Prompt = struct { stream: *lanes.Stream, rows: usize };
 
-/// The original prompt cuts: requested boundaries, the decision row, and the configured row cap.
-pub inline fn rows(s: *const lanes.Stream, at: usize, cap: usize, decision: ?usize, next: *usize) usize {
+/// The prompt cuts: requested boundaries, the decision row, TF_KV_SHARE's store marks, and the configured row cap.
+pub inline fn rows(s: *const lanes.Stream, at: usize, cap: usize, decision: ?usize, marks: []const u32, next: *usize) usize {
     while (next.* < s.chunks.len and s.chunks[next.*] <= at) next.* += 1;
     var end = if (next.* < s.chunks.len) @min(s.chunks[next.*], s.prompt_len) else s.prompt_len;
     if (decision) |d| if (d >= at and d + 1 < end) {
         end = d + 1;
     };
+    for (marks) |m| if (m > at) {
+        end = @min(end, m); // the store keeps a state only where a chunk ends
+    };
     return @min(cap, end - at);
 }
 

zig/src/families/nemotron/prompt_flight_test.zig

@@ -50,10 +50,14 @@ const Fake = struct {
 };
 
 fn fill(f: *Fake, s: *lanes.Stream, cap: usize, decision: ?usize) !void {
+    try fillAt(f, s, cap, decision, &.{});
+}
+
+fn fillAt(f: *Fake, s: *lanes.Stream, cap: usize, decision: ?usize, marks: []const u32) !void {
     var at: usize = 0;
     var next: usize = 0;
     while (at < s.prompt_len) {
-        const n = workflow.rows(s, at, cap, decision, &next);
+        const n = workflow.rows(s, at, cap, decision, marks, &next);
         try workflow.submit(f, s, .{ .rows = n });
         at += n;
     }
@@ -79,6 +83,24 @@ test "Nemotron prefill widths: actual requested cuts, decision row and cap land
     try t.expectEqual(@as(f64, 12), f.gpu_ms);
 }
 
+test "Nemotron prefill widths: a chunk ends at every TF_KV_SHARE store mark, in any order, beside the other cuts" {
+    const ids: [29]u32 = @splat(7);
+    var s = try lanes.Stream.init(a, .{ .id = "marks", .prompt = &ids, .max_new = 1, .chunks = &.{ 3, 20 } });
+    defer s.deinit(a);
+    var f: Fake = .{};
+    defer f.deinit();
+    try fillAt(&f, &s, 8, 12, &.{ 24, 16, 29, 40 });
+    try f.drain();
+    try t.expectEqualSlices(u32, &.{ 3, 8, 2, 3, 4, 4, 5 }, s.prefill_widths.items);
+    var ends: [8]u32 = undefined;
+    var end: u32 = 0;
+    for (s.prefill_widths.items, 0..) |w, i| {
+        end += w;
+        ends[i] = end;
+    }
+    for ([_]u32{ 16, 24, 29 }) |m| try t.expect(std.mem.indexOfScalar(u32, ends[0..s.prefill_widths.items.len], m) != null);
+}
+
 test "Nemotron prefill widths: one pass waits once; verify and draft buffers never count as prompt passes" {
     var s = try lanes.Stream.init(a, .{ .id = "one", .prompt = &.{ 1, 2, 3, 4, 5, 6, 7 }, .max_new = 1 });
     defer s.deinit(a);

zig/src/families/nemotron/prompt_pages.zig

@@ -0,0 +1,81 @@
+//! TF_KV_SHARE on Metal: a prompt pass resumes on the store's kept pages and keeps its own pages at the store's marks.
+const std = @import("std");
+const lanes = @import("lanes");
+const st = @import("state.zig");
+const snapshot = @import("snapshot.zig");
+const kvp = @import("../../core/metal_kv_pages.zig");
+const Metal = @import("backend.zig").Metal;
+
+/// One prompt pass's plan from the store: the pages it adopts, where they end, and the marks it keeps at.
+pub const Pass = struct {
+    on: bool = false,
+    adopt: std.ArrayList(u32) = .empty, // held once more for the stream until its lane adopts them
+    from: u32 = 0,
+    marks: []const u32 = &.{}, // the store's allocation, freed by `end`
+    kept: u64 = 0,
+
+    /// The store's deepest kept state of `s`'s prompt (none without TF_KV_SHARE): its pages held, its state pending.
+    pub fn begin(b: *Metal, s: *lanes.Stream) !Pass {
+        const pr = b.prompts orelse return .{};
+        b.pending = null;
+        var p: Pass = .{ .on = true };
+        errdefer p.adopt.deinit(b.gpa);
+        const got = try pr.begin(s.prompt(), s.history_len, s.shared_prefixes, @ptrCast(s), &p.adopt);
+        p.from = got.from;
+        p.marks = got.marks;
+        return p;
+    }
+
+    /// Pages the lane holds from the store: its promise is that many fewer.
+    pub fn held(p: *const Pass) usize {
+        return p.adopt.items.len;
+    }
+
+    /// The lane's table starts with the adopted pages and the pending Mamba states go in; the pass resumes there.
+    pub fn resume_(p: *Pass, b: *Metal, s: *lanes.Stream, c: *st.Cache) !usize {
+        if (!p.on) return 0;
+        if (p.adopt.items.len > 0) {
+            const pg = c.kv[0].paged orelse return error.NotPaged;
+            const ids = p.adopt.items;
+            p.adopt.items.len = 0; // the lane's now: adopt releases them if it fails
+            try pg.lease.pages.adopt(pg.lease.lane, ids);
+        }
+        if (p.from == 0) return 0;
+        const k = b.pending orelse return error.KeptStateMismatch;
+        b.pending = null;
+        if (k.at != p.from or !k.pages) return error.KeptStateMismatch;
+        try snapshot.restore(b, c, k);
+        s.cached = p.from;
+        return p.from;
+    }
+
+    /// At one of the store's marks the store takes the pages it lacks and the state; pages it held replace the lane's.
+    pub fn mark(p: *Pass, b: *Metal, s: *lanes.Stream, c: *st.Cache, at: usize) void {
+        if (!p.on or std.mem.indexOfScalar(u32, p.marks, @intCast(at)) == null) return;
+        p.keep(b, s, c, at) catch |e| std.log.warn("nemotron: prompt pages not kept: {s}", .{@errorName(e)});
+    }
+
+    fn keep(p: *Pass, b: *Metal, s: *lanes.Stream, c: *st.Cache, at: usize) !void {
+        const pg = c.kv[0].paged orelse return error.NotPaged;
+        const n = at / kvp.page_tokens;
+        const mine = try b.gpa.dupe(u32, pg.ids()[0..n]);
+        defer b.gpa.free(mine);
+        const path = try b.gpa.alloc(u32, n);
+        defer b.gpa.free(path);
+        const shared = b.prompts.?.keep(s.prompt(), @intCast(at), @ptrCast(s), mine, path);
+        if (shared == n) p.kept += 1;
+        if (std.mem.eql(u32, mine[0..shared], path[0..shared])) return;
+        try b.drain(); // the chunk in flight writes the lane's pages: none goes back while it may
+        for (path[0..shared], mine[0..shared], 0..) |id, own, i| if (id != own) pg.lease.pages.share(pg.lease.lane, i, id);
+    }
+
+    /// The pass's line, and the plan's memory back.
+    pub fn end(p: *Pass, b: *Metal, s: *lanes.Stream, ok: bool) void {
+        if (!p.on) return;
+        if (ok) b.prompts.?.report(s.prompt().len, s.cached, p.kept);
+        for (p.adopt.items) |id| b.kv_pages.?.pool.release(id); // pages a failed pass never adopted
+        p.adopt.deinit(b.gpa);
+        b.prompts.?.free(p.marks);
+        p.* = .{};
+    }
+};

zig/src/families/nemotron/residency.zig

@@ -3,11 +3,11 @@ const std = @import("std");
 const mtl = @import("metal");
 const fields = mtl.resident.fields;
 
-pub fn openCache(b: anytype, stream: *@import("lanes").Stream, capacity: usize) !*@import("state.zig").Cache {
+pub fn openCache(b: anytype, stream: *@import("lanes").Stream, capacity: usize, held: usize) !*@import("state.zig").Cache {
     const state = @import("state.zig");
     const cache = try b.gpa.create(state.Cache);
     errdefer b.gpa.destroy(cache);
-    const lease = try leasePages(b, capacity);
+    const lease = try leasePages(b, capacity, held);
     errdefer if (lease) |l| l.pages.finish(l.lane);
     cache.* = try state.Cache.initPaged(b.m.device, b.m.config, capacity, b.head != null, lease);
     cache.setResident(b.residency);
@@ -22,13 +22,13 @@ pub fn openCache(b: anytype, stream: *@import("lanes").Stream, capacity: usize)
     return cache;
 }
 
-/// TF_NEMOTRON_METAL_KV_PAGES: a free lane promised the pages of `capacity` rows; LaneMemory: none now.
-fn leasePages(b: anytype, capacity: usize) !?@import("state.zig").Lease {
+/// TF_NEMOTRON_METAL_KV_PAGES: a free lane promised the pages of `capacity` rows less `held` adopted; LaneMemory: none.
+fn leasePages(b: anytype, capacity: usize, held: usize) !?@import("state.zig").Lease {
     const pages = b.kv_pages orelse return null;
     const lane = for (pages.leases, 0..) |l, i| {
         if (l.promised == 0 and l.ids.items.len == 0) break i;
     } else return error.LaneMemory;
-    try pages.admit(lane, pages.layout.pages(capacity)); // its pages are taken as its rows grow (kv_pages.grow)
+    try pages.admit(lane, pages.layout.pages(capacity) -| held); // its pages are taken as its rows grow (kv_pages.grow)
     return .{ .pages = pages, .lane = lane };
 }