diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 851f919480..eaa977f265 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -7,19 +7,12 @@ const nemotron = @import("nemotron"); const Allocator = std.mem.Allocator; const budget = @import("cuda_memory.zig"); const lane_backend = @import("cuda_bound.zig"); +const snaps = @import("cuda_snaps.zig"); const nemotron_slide = @import("nemotron_slide.zig"); /// CUDA families provide metadata, open their lane backend and explain their refusals. const registry = .{nemotron.native}; -/// Nemotron's prompt-cache copies. A file-level value so the store's pointer stays valid for the process. -const nemotron_snaps: api.prompt_cache.Snapshots.VTable = .{ - .bytes = nemotron.native.snap_bytes, - .save = nemotron.native.snap_save, - .restore = nemotron.native.snap_restore, - .drop = nemotron.native.snap_drop, -}; - pub const backends: []const []const u8 = &.{"cuda"}; pub const families: []const api.Family = blk: { var out: [registry.len]api.Family = undefined; @@ -322,13 +315,18 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O }, }; const before = try pool(a, io, &g.ctx, problem) orelse return null; - const weights = weightBytes(io, o.dir); + // a family split over hosts holds only its rank's share (GLM-5.3-Flash's TP) + const weights = if (@hasDecl(F, "weightShare")) F.weightShare(weightBytes(io, o.dir)) else weightBytes(io, o.dir); const held0 = cuda.usage(false).device; if (weights > before.room(held0)) { problem.* = try std.fmt.allocPrint(a, "the checkpoint's {d:.1} GiB of weights do not fit the {d:.1} GiB the CUDA memory budget grants ({d:.1} GiB free less a {d:.1} GiB reserve{s}); free device memory or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB", .{ toGib(weights), toGib(before.room(held0)), toGib(before.free), toGib(before.reserve), if (before.limit != null) ", under TENSORFOLD_CUDA_MEMORY_LIMIT_GB" else "" }); return null; } const loaded = F.open(gpa, io, &g.ctx, o.dir, kernels, .{ .context = @intCast(window), .drafts = o.drafts, .segments = segments, .slide = o.slide }) catch |e| { + if (@hasDecl(F, "explains_open")) if (F.explain(null, e)) |text| { + problem.* = text; + return null; + }; problem.* = try std.fmt.allocPrint(a, "the native CUDA engine cannot load {s} with {s} glue kernels ({s})", .{ o.dir, kernels orelse "its own", @errorName(e) }); return null; }; @@ -347,24 +345,24 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O const h = try gpa.create(Host); errdefer gpa.destroy(h); h.* = .{ .gpa = gpa, .gpu = g, .family = loaded.ctx, .release = loaded.deinit, .inner = loaded.backend, .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined, .lone = loaded.lone }; - const cache = budget.cacheBytes(o.prompt_cache_gib, room, loaded.stream_bytes * (streams -| free)); - h.startup = try std.fmt.allocPrint(gpa, "CUDA sm_{d} device {d} ({s}{s}): model {d:.2} GiB{s}; {d} stream{s} at once, {d:.2} GiB each at a {d}-token window, of {d:.1} GiB left after a {d:.1} GiB reserve; prompts in passes of up to {d} rows{s}{s}; {s}", .{ + const cache = if (@hasDecl(F, "snap_save")) budget.cacheBytes(o.prompt_cache_gib, room, loaded.stream_bytes * (streams -| free)) else 0; + h.startup = try std.fmt.allocPrint(gpa, "CUDA sm_{d} device {d} ({s}{s}): model {d:.2} GiB{s}; {d} stream{s} at once, {d:.2} GiB each at a {d}-token window, of {d:.1} GiB left after a {d:.1} GiB reserve; prompts in {d}-row chunks{s}{s}; {s}", .{ capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (loaded.learn_bytes > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(loaded.learn_bytes)}) else "", streams, if (streams == 1) "" else "s", toGib(loaded.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", if (cache > 0) try std.fmt.allocPrint(a, ", a {d:.1} GiB prompt cache", .{toGib(cache)}) else ", no prompt cache", if (kernels) |dir| try std.fmt.allocPrint(a, "glue kernels captured at {s}", .{dir}) else "own glue kernels", }); errdefer gpa.free(h.startup); - if (cache > 0) { + if (@hasDecl(F, "snap_save") and cache > 0) { const pc = api.prompt_cache; const store = try gpa.create(pc.Store); - // states keep at any row at no prompt call (the head reads one token on); replies prefill when it is quiet - store.* = pc.Store.init(gpa, .{ .ptr = loaded.target, .vtable = &nemotron_snaps }, .{ .lookahead = 1, .min_prompt = 0, .warm = .waits }, cache - try nemotron.native.snap_ready(loaded.target, cache)); + // rules and budget from the family's snap_* declarations (cuda_snaps.zig) + store.* = pc.Store.init(gpa, .{ .ptr = loaded.target, .vtable = &snaps.Snaps(F).vtable }, snaps.cacheRules(F), try snaps.cacheBudget(F, loaded.target, cache)); h.store = store; } errdefer if (h.store) |store| { store.deinit(); gpa.destroy(store); }; - if (loaded.learns != null) { + if (@hasDecl(F, "slide")) if (loaded.learns != null) { const Slide = nemotron_slide.Adapter(F.slide); h.dir = try gpa.dupe(u8, o.dir); const s = try gpa.create(Slide); @@ -378,13 +376,13 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O } }.f; h.learn = s.hook(); - } + }; errdefer if (h.slide) |s| { h.slide_free.?(gpa, s); gpa.free(h.dir); }; // the family cuts its own prompt grid from position 0, as `tensorfold run` does: prefill_step 0 - try h.serve(io, loaded.facts, loaded.rows, .{ .lanes = streams, .context_window = @intCast(window), .startup = h.startup, .prompt_cache = h.store != null, .warm_turns = h.store != null, .logprobs = loaded.backend.vtable.first_row != null, .prompt_logprobs = loaded.backend.vtable.prompt_rows, .structures = loaded.backend.vtable.first_masked != null }, .{ .ctx = loaded.ctx, .text = F.explain }); + try h.serve(io, loaded.facts, loaded.rows, .{ .lanes = streams, .context_window = @intCast(window), .startup = h.startup, .prompt_cache = h.store != null, .warm_turns = h.store != null and @hasDecl(F, "warm_quiet"), .logprobs = loaded.backend.vtable.first_row != null, .prompt_logprobs = loaded.backend.vtable.prompt_rows, .structures = loaded.backend.vtable.first_masked != null }, .{ .ctx = loaded.ctx, .text = F.explain }); opened = true; return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; } @@ -397,6 +395,7 @@ test { _ = budget; // cuda_memory.zig's own tests, which no other test binary reaches _ = refusals; _ = lane_backend; + _ = snaps; } test "the lane core sees a family's prefill_step through the lane backend, so a served prompt can step" { diff --git a/zig/src/native/cuda_snaps.zig b/zig/src/native/cuda_snaps.zig new file mode 100644 index 0000000000..961ab866ac --- /dev/null +++ b/zig/src/native/cuda_snaps.zig @@ -0,0 +1,62 @@ +//! A CUDA family's prompt cache in the native server: rules, budget and copies from the family's `snap_*` decls. +const std = @import("std"); +const api = @import("engine_api"); + +/// Any row, the head one token on (Nemotron's); `snap_grid`'s chunk starts and `warm_quiet`'s warm turns if declared. +pub fn cacheRules(comptime F: type) api.prompt_cache.Rules { + var rules: api.prompt_cache.Rules = .{ .lookahead = 1, .min_prompt = 0 }; + if (@hasDecl(F, "warm_quiet")) rules.warm = .waits; + if (@hasDecl(F, "snap_grid")) { + rules.planned = true; + rules.grid = F.snap_grid; + } + return rules; +} + +/// The cache's bytes less what the family sets aside for its states (`snap_ready`, as Nemotron's spare buffers). +pub fn cacheBudget(comptime F: type, target: anytype, cache: u64) !u64 { + return if (@hasDecl(F, "snap_ready")) cache - try F.snap_ready(target, cache) else cache; +} + +/// The family's `snap_*` copies (none: no prompt cache), a container-level value so the store's pointer stays valid. +pub fn Snaps(comptime F: type) type { + return struct { + pub const vtable: api.prompt_cache.Snapshots.VTable = .{ + .bytes = F.snap_bytes, + .save = F.snap_save, + .restore = F.snap_restore, + .drop = F.snap_drop, + }; + }; +} + +test "a family with snap_grid keeps on its chunk starts, any other at any row as Nemotron does" { + const Glm = struct { + pub const snap_grid: u32 = 2048; + }; + const Other = struct {}; + const Warm = struct { + pub const warm_quiet: u64 = 1; + }; + const any = cacheRules(Other); + try std.testing.expect(!any.planned and any.min_prompt == 0 and any.lookahead == 1 and any.warm == .off); + try std.testing.expect(cacheRules(Warm).warm == .waits and cacheRules(Glm).warm == .off); + const Stub = struct { + 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 unused: u8 = 0; + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const prompt = try arena.allocator().alloc(u32, 3547); + for (prompt, 0..) |*t, i| t.* = @intCast(i); + var store = api.prompt_cache.Store.init(std.testing.allocator, .{ .ptr = &unused, .vtable = &.{ .bytes = Stub.bytes, .save = Stub.save, .restore = Stub.restore, .drop = Stub.drop } }, cacheRules(Glm), 1 << 30); + defer store.deinit(); + try std.testing.expectEqualSlices(u32, &[_]u32{2048}, try store.marks(arena.allocator(), prompt, 0, 3540, &.{}, &.{}, &.{})); +}