From 467eb62d760789a0a8800e9b21d925d60d1ff1c7 Mon Sep 17 00:00:00 2001 From: Jishnu Venugopal <50513172+jishnuvenugopal@users.noreply.github.com> Date: Sat, 17 Oct 2026 19:03:24 +0800 Subject: [PATCH 1/2] nemotron metal: size the Mamba pool to active streams and budget the prompt cache from wired footprint Pool.init and Scratch.init allocated their Metal buffers without touching the pages, so when the window-cost table was read from disk the process footprint at cache_fit.budget was 3.5 GiB below what the first forward wired, while a fresh timing run wired those pages first and left 0.0 GiB for the prompt cache. backend.Metal.init also sized Pool.init for rows + streams + 1 Mamba slots (67 slots at 2 streams, 3.10 GiB) when each active stream holds at most two slots plus slot 0 and a spare. Size Pool.init to 2 * o.streams + 3 Mamba slots and o.streams + 2 replay ids, zero-fill Pool.init and Scratch.init buffers on creation so their pages enter the process footprint before cache_fit runs, and pass the lazy per-stream Cache.bytes total to cache_fit.budgetBeside. --- zig/src/families/nemotron/backend.zig | 2 +- zig/src/families/nemotron/nemotron.zig | 1 + zig/src/families/nemotron/state.zig | 59 +++++++++++++++++++++++--- zig/src/native/metal.zig | 3 +- zig/tests/nemotron_wide.zig | 6 +-- 5 files changed, 61 insertions(+), 10 deletions(-) diff --git a/zig/src/families/nemotron/backend.zig b/zig/src/families/nemotron/backend.zig index 61abb9d12b..d02396d7a8 100644 --- a/zig/src/families/nemotron/backend.zig +++ b/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), diff --git a/zig/src/families/nemotron/nemotron.zig b/zig/src/families/nemotron/nemotron.zig index 54ff2c9ffc..5a636fb07b 100644 --- a/zig/src/families/nemotron/nemotron.zig +++ b/zig/src/families/nemotron/nemotron.zig @@ -33,6 +33,7 @@ pub const learned_prompts = @import("learned_prompts.zig"); pub const Model = model.Model; test { + _ = @import("state.zig"); _ = @import("train_check.zig"); _ = @import("model.zig"); _ = @import("timing.zig"); diff --git a/zig/src/families/nemotron/state.zig b/zig/src/families/nemotron/state.zig index eec5bab7a1..944181ccbd 100644 --- a/zig/src/families/nemotron/state.zig +++ b/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()); + } +} diff --git a/zig/src/native/metal.zig b/zig/src/native/metal.zig index c778274196..71dc15bcbf 100644 --- a/zig/src/native/metal.zig +++ b/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; }; diff --git a/zig/tests/nemotron_wide.zig b/zig/tests/nemotron_wide.zig index 7c32bcb3e4..94f47faf4c 100644 --- a/zig/tests/nemotron_wide.zig +++ b/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; From ad79aa89b15a3580104e93d10a8c4016e1ede082 Mon Sep 17 00:00:00 2001 From: Jishnu Venugopal <50513172+jishnuvenugopal@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:49:25 +0800 Subject: [PATCH 2/2] nemotron metal: a refused queue or verify gives its fresh Mamba slot back to the pool queueFn and verifyFn took a fresh Mamba slot from the pool and only handed it to the stream's cache through hold() after the command buffer landed, so an error in between (a later window refused, grammar bits refused, a failed command buffer) left the slot out of the pool for the life of the process while the lane host failed the round and served on. The pool now holds one spare beyond what every active stream needs, so each such round removes headroom a later lone round needs. Return the slot with errdefer in queueFn, and in verifyFn for every window whose slot is already in segs; stage a window's grammar bits before its slot is taken so a refused mask has nothing to return. --- zig/src/families/nemotron/backend.zig | 7 +- .../families/nemotron/backend_slots_test.zig | 66 +++++++++++++++++++ zig/src/families/nemotron/nemotron.zig | 1 + 3 files changed, 73 insertions(+), 1 deletion(-) create mode 100644 zig/src/families/nemotron/backend_slots_test.zig diff --git a/zig/src/families/nemotron/backend.zig b/zig/src/families/nemotron/backend.zig index d02396d7a8..577b0a1c27 100644 --- a/zig/src/families/nemotron/backend.zig +++ b/zig/src/families/nemotron/backend.zig @@ -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; diff --git a/zig/src/families/nemotron/backend_slots_test.zig b/zig/src/families/nemotron/backend_slots_test.zig new file mode 100644 index 0000000000..df9b51ac01 --- /dev/null +++ b/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); +} diff --git a/zig/src/families/nemotron/nemotron.zig b/zig/src/families/nemotron/nemotron.zig index 5a636fb07b..e6d29f0764 100644 --- a/zig/src/families/nemotron/nemotron.zig +++ b/zig/src/families/nemotron/nemotron.zig @@ -34,6 +34,7 @@ pub const Model = model.Model; test { _ = @import("state.zig"); + _ = @import("backend_slots_test.zig"); _ = @import("train_check.zig"); _ = @import("model.zig"); _ = @import("timing.zig");