nemotron metal: size the Mamba pool to active streams and budget the prompt cache from wired footprint

zig/src/families/nemotron/backend.zig

@@ -85,7 +85,7 @@ pub const Metal = struct {
             .gpa = gpa,
             .m = m,
             .o = o,
-            .pool = try st.Pool.init(gpa, m.device, m.config, rows + o.streams + 1, o.streams + 4),
+            .pool = try st.Pool.init(gpa, m.device, m.config, st.poolSlots(o.streams), o.streams + 2),
             .scratch = try st.Scratch.init(m.device, m.config, rows, o.capacity),
             .tokens = try m.device.buffer(ring * 4, opts),
             .prompt = try m.device.buffer(@max(o.capacity, 16) * 4, opts),
@@ -379,6 +379,7 @@ pub const Metal = struct {
         const out = self.take(1);
         if (c.len + 1 > c.kv[0].capacity) return error.ContextFull;
         const state = try self.slotFor(c, .full);
+        errdefer self.pool.give(@intCast(state));
         const Step = struct {
             in: u64,
             out: u64,
@@ -414,6 +415,8 @@ pub const Metal = struct {
         var fills: [st.max_rows]Fill = undefined;
         if (windows.len > segs.len) return error.WindowTooWide;
         var total: usize = 0;
+        var taken: usize = 0; // windows whose fresh slot segs holds until hold() hands it to their cache
+        errdefer for (segs[0..taken]) |g| if (g.slot >= 0) self.pool.give(@intCast(g.slot));
         for (windows, 0..) |w, i| {
             const c = try self.cacheOf(w.stream);
             self.settle(c);
@@ -428,7 +431,9 @@ pub const Metal = struct {
             staged[stage % ring] = w.pending;
             for (w.tokens, 0..) |t, r| staged[(stage + 1 + r) % ring] = t;
             fills[i] = .{ .stage = stage, .held = w.held, .n = w.tokens.len };
-            segs[i] = .{ .rows = rows, .cache = c, .store = .lag, .slot = try self.slotFor(c, .lag), .parents = w.parents, .draw = drawOf(w.stream, w.positions[0]), .mask_row = try grammar.stage(self, w, total) };
+            const mask_row = try grammar.stage(self, w, total);
+            segs[i] = .{ .rows = rows, .cache = c, .store = .lag, .slot = try self.slotFor(c, .lag), .parents = w.parents, .draw = drawOf(w.stream, w.positions[0]), .mask_row = mask_row };
+            taken = i + 1;
             c.start = total;
             c.rows = rows;
             total += rows;

zig/src/families/nemotron/backend_slots_test.zig

@@ -0,0 +1,66 @@
+//! The Metal backend's Mamba slots: a verify refused after a window took its fresh slot gives it back (backend.zig).
+const std = @import("std");
+const mtl = @import("metal");
+const lanes = @import("lanes");
+const cfg = @import("config.zig");
+const st = @import("state.zig");
+const Metal = @import("backend.zig").Metal;
+const Masks = @import("../../core/grammar_mask.zig").Masks;
+
+const be = lanes.backend;
+const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked;
+
+test "a verify refused after a window took its Mamba slot gives the slot back to the pool" {
+    const gpa = std.testing.allocator;
+    var why: cfg.Why = .{};
+    const c = try cfg.parse(gpa, cfg.lightning, &why);
+    var device = mtl.Device.init() catch return error.SkipZigTest;
+    defer device.deinit();
+    // only what a verify reads before its forward: the pool, the token ring, the caches and grammar bits (no weights)
+    var b: Metal = .{
+        .gpa = gpa,
+        .m = undefined,
+        .o = .{ .capacity = 4096, .streams = 2 },
+        .pool = try st.Pool.init(gpa, device, c, st.poolSlots(2), 4),
+        .scratch = undefined,
+        .tokens = try device.buffer((1 << 16) * 4, opts),
+        .prompt = undefined,
+        .fence = undefined,
+        .masks = try Masks.init(device, st.max_rows, c.vocab),
+    };
+    defer {
+        b.masks.?.deinit();
+        b.caches.deinit(gpa);
+        b.tokens.deinit();
+        b.pool.deinit();
+    }
+    const back = b.backend();
+    var streams: [2]lanes.Stream = undefined;
+    for (&streams) |*s| {
+        s.* = try lanes.Stream.init(gpa, .{ .id = "s", .prompt = &.{ 1, 2 }, .max_new = 8 });
+        const cache = try gpa.create(st.Cache);
+        cache.* = try st.Cache.init(device, c, 4096, false);
+        cache.replay = 1; // the last verify's kept row replays first, so the next verify stores to a fresh slot
+        try b.caches.put(gpa, s, cache);
+    }
+    defer for (&streams) |*s| {
+        back.release(s);
+        s.deinit(gpa);
+    };
+    const free = b.pool.free.items.len;
+    var sampled: [2]u32 = undefined;
+    var drafts: [0]u32 = .{};
+    var out = [_]be.Verified{ .{ .sampled = sampled[0..1], .drafts = &drafts }, .{ .sampled = sampled[1..2], .drafts = &drafts } };
+    const first: be.Window = .{ .stream = &streams[0], .pending = 1, .held = 0, .tokens = &.{}, .parents = null, .positions = &.{2} };
+
+    // a later window is refused: the first one's slot is already taken
+    const broken: be.Window = .{ .stream = &streams[1], .pending = 1, .held = 0, .tokens = &.{5}, .parents = null, .positions = &.{ 2, 9 } };
+    try std.testing.expectError(error.PositionsNotAChain, back.verify(&.{ first, broken }, &out));
+    try std.testing.expectEqual(free, b.pool.free.items.len);
+
+    // the window's own grammar bits are refused
+    var masked = first;
+    masked.masks = &.{0};
+    try std.testing.expectError(error.InvalidGrammarMasks, back.verify(&.{masked}, out[0..1]));
+    try std.testing.expectEqual(free, b.pool.free.items.len);
+}

zig/src/families/nemotron/nemotron.zig

@@ -33,6 +33,8 @@ pub const learned_prompts = @import("learned_prompts.zig");
 pub const Model = model.Model;
 
 test {
+    _ = @import("state.zig");
+    _ = @import("backend_slots_test.zig");
     _ = @import("train_check.zig");
     _ = @import("model.zig");
     _ = @import("timing.zig");

zig/src/families/nemotron/state.zig

@@ -19,6 +19,11 @@ pub const tree_saves = 64;
 /// Head levels whose best tokens a stream keeps.
 pub const max_levels = 72;
 
+/// Mamba slots for streams: slot 0, two per stream (state and pending forward), a lone round's third, and a spare.
+pub fn poolSlots(streams: usize) usize {
+    return 2 * streams + 3;
+}
+
 /// Mamba states by slot (slot 0: the zero state, never written), and by replay id a stream's last rows' raw inputs.
 pub const Pool = struct {
     conv: [cfg.max_layers]mtl.Buffer = undefined, // [slots, KC-1, CD] bf16, by Mamba index
@@ -36,13 +41,17 @@ pub const Pool = struct {
         var p = Pool{ .layers = c.count(.mamba), .slots = slots, .gpa = gpa };
         const conv_bytes = slots * (c.conv_kernel - 1) * c.convDim() * 2;
         const ssm_bytes = slots * c.mamba_heads * c.mamba_head_dim * c.state * 4;
+        const raw_bytes = streams * 2 * replay_rows * c.convDim() * 2;
+        const dtraw_bytes = streams * 2 * replay_rows * c.mamba_heads * 2;
         for (0..p.layers) |m| {
             p.conv[m] = try device.buffer(conv_bytes, opts);
             p.ssm[m] = try device.buffer(ssm_bytes, opts);
-            p.raw[m] = try device.buffer(streams * 2 * replay_rows * c.convDim() * 2, opts);
-            p.dtraw[m] = try device.buffer(streams * 2 * replay_rows * c.mamba_heads * 2, opts);
-            @memset(p.conv[m].contents()[0 .. conv_bytes / slots], 0);
-            @memset(p.ssm[m].contents()[0 .. ssm_bytes / slots], 0);
+            p.raw[m] = try device.buffer(raw_bytes, opts);
+            p.dtraw[m] = try device.buffer(dtraw_bytes, opts);
+            @memset(p.conv[m].contents()[0..conv_bytes], 0);
+            @memset(p.ssm[m].contents()[0..ssm_bytes], 0);
+            @memset(p.raw[m].contents()[0..raw_bytes], 0);
+            @memset(p.dtraw[m].contents()[0..dtraw_bytes], 0);
         }
         p.tree = try device.buffer(tree_saves * c.mamba_heads * c.mamba_head_dim * c.state * 4, opts);
         var s = slots;
@@ -148,6 +157,12 @@ pub const Cache = struct {
         return s;
     }
 
+    /// Bytes one stream's KV caches and token windows take on its first prompt.
+    pub fn bytes(c: cfg.Config, capacity: usize, head: bool) u64 {
+        const kv: u64 = (c.count(.attention) + @intFromBool(head)) * 2 * c.kv_heads * capacity * c.head_dim * 2;
+        return kv + 3 * max_rows * 4 + 2 * max_levels * 4 * 4;
+    }
+
     /// After a prompt chunk of `rows` rows whose last row's Mamba states went to `slot` (nothing left to replay).
     pub fn advance(self: *Cache, pool: *Pool, rows: usize, slot: u32) void {
         std.debug.assert(self.replay == 0);
@@ -236,7 +251,10 @@ pub const Scratch = struct {
         const chunks = (capacity + 511) / 512;
         const B = struct {
             fn make(d: mtl.Device, bytes: usize) !mtl.Buffer {
-                return d.buffer(@max(bytes, 64), opts);
+                const n = @max(bytes, 64);
+                const b = try d.buffer(n, opts);
+                @memset(b.contents()[0..n], 0);
+                return b;
             }
         };
         return .{
@@ -299,3 +317,34 @@ fn partCols(c: cfg.Config) usize {
     for (shapes) |nk| most = @max(most, fwd.splitK(nk[0], nk[1]) * nk[0]);
     return most;
 }
+
+test "Cache.bytes matches the Metal buffers Cache.init allocates" {
+    var why: cfg.Why = .{};
+    const c = try cfg.parse(std.testing.allocator, cfg.lightning, &why);
+    var device = mtl.Device.init() catch return error.SkipZigTest;
+    defer device.deinit();
+    var pool = try Pool.init(std.testing.allocator, device, c, 3, 2);
+    defer pool.deinit();
+    for ([_]bool{ false, true }) |head| {
+        var cache = try Cache.init(device, c, 4096, head);
+        defer cache.deinit(&pool);
+        var actual: u64 = cache.windows[0].length() + cache.windows[1].length() + cache.follow.length() + cache.topk.length();
+        for (cache.kv[0..cache.attentions]) |kv| actual += kv.k.length() + kv.v.length();
+        if (cache.mtp) |m| actual += m.k.length() + m.v.length();
+        try std.testing.expectEqual(actual, Cache.bytes(c, 4096, head));
+    }
+}
+
+test "a pool of poolSlots holds every stream's state and pending forward, a lone round's third, and one spare" {
+    var why: cfg.Why = .{};
+    const c = try cfg.parse(std.testing.allocator, cfg.lightning, &why);
+    var device = mtl.Device.init() catch return error.SkipZigTest;
+    defer device.deinit();
+    for ([_]usize{ 1, 2, 4, 8 }) |streams| {
+        var pool = try Pool.init(std.testing.allocator, device, c, poolSlots(streams), streams + 2);
+        defer pool.deinit();
+        for (0..2 * streams + 1) |_| _ = try pool.take();
+        _ = try pool.take();
+        try std.testing.expectError(error.StateSlotsFull, pool.take());
+    }
+}

zig/src/native/metal.zig

@@ -230,7 +230,8 @@ pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]c
     h.round = .{ .depth = 8 };
     h.cache = null;
     var why: []const u8 = "";
-    const plan = cache_fit.budget(o.prompt_cache_gib, o.prompt_cache_over_cap, h.m.device.maxWorkingSet() -| h.m.device.allocated() -| (8 << 30), a, &why, "") catch |e| {
+    const later = @as(u64, o.lanes) * nemotron.state.Cache.bytes(h.m.config, h.metal.o.capacity, h.metal.head != null);
+    const plan = cache_fit.budgetBeside(o.prompt_cache_gib, o.prompt_cache_over_cap, h.m.device.maxWorkingSet() -| h.m.device.allocated() -| (8 << 30), later, a, &why, "") catch |e| {
         if (e == error.CacheOverCap) problem.* = why;
         return e;
     };

zig/tests/nemotron_wide.zig

@@ -117,7 +117,7 @@ fn geometry(gpa: std.mem.Allocator, arena: std.mem.Allocator, o: Options) !void
     for (o.rows) |r| most = @max(most, r);
     const m = try nm.Model.load(gpa, io, o.model, true);
     defer m.deinit();
-    const b = try nm.backend.Metal.init(gpa, m, .{ .chunk = 16, .capacity = prompt.len + o.at + most + 128, .drafts = true, .streams = 2, .batch_rows = @max(most, 16) });
+    const b = try nm.backend.Metal.init(gpa, m, .{ .chunk = 16, .capacity = prompt.len + o.at + most + 128, .drafts = true, .streams = @max(most, 2), .batch_rows = @max(most, 16) });
     defer b.deinit();
     var truth: std.ArrayList(u32) = .empty;
     var p = try Probe.init(b);
@@ -172,7 +172,7 @@ fn profile(gpa: std.mem.Allocator, arena: std.mem.Allocator, o: Options) !void {
     const m = try nm.Model.load(gpa, io, o.model, true);
     defer m.deinit();
     const c = m.config;
-    const b = try nm.backend.Metal.init(gpa, m, .{ .chunk = 16, .capacity = prompt.len + o.at + most + 128, .drafts = true, .streams = 2, .batch_rows = @max(most, 16) });
+    const b = try nm.backend.Metal.init(gpa, m, .{ .chunk = 16, .capacity = prompt.len + o.at + most + 128, .drafts = true, .streams = @max(most, 2), .batch_rows = @max(most, 16) });
     defer b.deinit();
     const kept = 3; // a verify's replayed rows: the last window's kept rows
 
@@ -249,7 +249,7 @@ fn curve(gpa: std.mem.Allocator, arena: std.mem.Allocator, o: Options) !void {
     for (o.rows) |r| most = @max(most, r);
     const m = try nm.Model.load(gpa, io, o.model, true);
     defer m.deinit();
-    const b = try nm.backend.Metal.init(gpa, m, .{ .chunk = 16, .capacity = prompt.len + o.at + most + 128, .drafts = true, .streams = 2, .batch_rows = @max(most, 16) });
+    const b = try nm.backend.Metal.init(gpa, m, .{ .chunk = 16, .capacity = prompt.len + o.at + most + 128, .drafts = true, .streams = @max(most, 2), .batch_rows = @max(most, 16) });
     defer b.deinit();
     const kept = 3;
     var truth: std.ArrayList(u32) = .empty;