diff --git a/zig/src/families/nemotron/backend.zig b/zig/src/families/nemotron/backend.zig index 61abb9d12b..577b0a1c27 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), @@ -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 54ff2c9ffc..e6d29f0764 100644 --- a/zig/src/families/nemotron/nemotron.zig +++ b/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"); 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;