diff --git a/zig/src/families/dsv41/dsv41.zig b/zig/src/families/dsv41/dsv41.zig index 38ba0e44d8..04b13a25a7 100644 --- a/zig/src/families/dsv41/dsv41.zig +++ b/zig/src/families/dsv41/dsv41.zig @@ -98,4 +98,5 @@ test { refAll(sampling); refAll(model); refAll(lanes); + refAll(@import("kept_order.zig")); } diff --git a/zig/src/families/dsv41/exl3_linear.zig b/zig/src/families/dsv41/exl3_linear.zig index 25a9241925..dcce543ad1 100644 --- a/zig/src/families/dsv41/exl3_linear.zig +++ b/zig/src/families/dsv41/exl3_linear.zig @@ -164,7 +164,7 @@ pub fn rotMany(k: *const Kernels, stream: cuda.Stream, calls: []const Call, m: u var a: RArgs = .{ .n = @intCast(calls.len), .M = @intCast(m) }; var warps: usize = 0; for (calls, 0..) |c, i| { - a.l[i] = .{ .x = c.x, .suh = c.layer.suh, .xh = c.xh, .ldx = c.ldx, .x_dtype = @intFromEnum(c.x_dtype), .K = @intCast(c.layer.k), .first = @intCast(warps) }; + a.l[i] = .{ .x = c.x, .suh = c.layer.suh, .xh = c.xh, .ldx = c.ldx, .x_dtype = @backingInt(c.x_dtype), .K = @intCast(c.layer.k), .first = @intCast(warps) }; warps += m * (c.layer.k / 128); } var args: cuda.Args = .{}; @@ -189,7 +189,7 @@ pub fn glinear(k: *const Kernels, stream: cuda.Stream, calls: []const Call, m: u done[j] = true; const s = strides(c.layer); const n: usize = @intCast(a.n); - a.l[n] = .{ .xh = c.xh, .ldx = @intCast(c.layer.k), .T = c.layer.words, .stride_k = s[0], .stride_nb = s[1], .svh = c.layer.svh, .y = c.y, .y_dtype = @intFromEnum(c.y_dtype), .ldy = c.ldy, .K = @intCast(c.layer.k), .N = @intCast(c.layer.n), .SK = @intCast(pj.sk), .counters = c.counters, .first = @intCast(blocks) }; + a.l[n] = .{ .xh = c.xh, .ldx = @intCast(c.layer.k), .T = c.layer.words, .stride_k = s[0], .stride_nb = s[1], .svh = c.layer.svh, .y = c.y, .y_dtype = @backingInt(c.y_dtype), .ldy = c.ldy, .K = @intCast(c.layer.k), .N = @intCast(c.layer.n), .SK = @intCast(pj.sk), .counters = c.counters, .first = @intCast(blocks) }; if (c.rot) |r| { a.l[n].rsuh = r.suh; a.l[n].rxh = r.xh; @@ -243,12 +243,12 @@ test "plans as linear.py makes them" { test "plans and strides equal every recorded glinear layer of the served build" { // (K, N, SK, WK, K2, stride_k, stride_nb) of each distinct layer the recording's glinear calls passed const seen = [_][7]u32{ - .{ 512, 128, 1, 4, 16, 512, 16384 }, .{ 1280, 4096, 2, 4, 10, 320, 25600 }, .{ 1280, 16384, 1, 8, 8, 256, 20480 }, - .{ 1280, 16384, 1, 8, 10, 320, 25600 }, .{ 3072, 25600, 1, 8, 8, 256, 49152 }, .{ 3072, 25600, 1, 8, 10, 320, 61440 }, - .{ 4096, 1024, 8, 4, 8, 256, 65536 }, .{ 4096, 1024, 8, 4, 10, 320, 81920 }, .{ 4096, 5120, 8, 4, 8, 256, 65536 }, - .{ 4096, 5120, 8, 4, 10, 320, 81920 }, .{ 5120, 512, 8, 4, 8, 256, 81920 }, .{ 5120, 512, 8, 4, 10, 320, 102400 }, - .{ 5120, 512, 8, 4, 12, 384, 122880 }, .{ 5120, 1280, 8, 4, 8, 256, 81920 }, .{ 5120, 1280, 8, 4, 10, 320, 102400 }, - .{ 5120, 1280, 8, 4, 12, 384, 122880 }, .{ 5120, 64640, 1, 8, 12, 384, 122880 }, .{ 15360, 5120, 8, 4, 8, 256, 245760 }, + .{ 512, 128, 1, 4, 16, 512, 16384 }, .{ 1280, 4096, 2, 4, 10, 320, 25600 }, .{ 1280, 16384, 1, 8, 8, 256, 20480 }, + .{ 1280, 16384, 1, 8, 10, 320, 25600 }, .{ 3072, 25600, 1, 8, 8, 256, 49152 }, .{ 3072, 25600, 1, 8, 10, 320, 61440 }, + .{ 4096, 1024, 8, 4, 8, 256, 65536 }, .{ 4096, 1024, 8, 4, 10, 320, 81920 }, .{ 4096, 5120, 8, 4, 8, 256, 65536 }, + .{ 4096, 5120, 8, 4, 10, 320, 81920 }, .{ 5120, 512, 8, 4, 8, 256, 81920 }, .{ 5120, 512, 8, 4, 10, 320, 102400 }, + .{ 5120, 512, 8, 4, 12, 384, 122880 }, .{ 5120, 1280, 8, 4, 8, 256, 81920 }, .{ 5120, 1280, 8, 4, 10, 320, 102400 }, + .{ 5120, 1280, 8, 4, 12, 384, 122880 }, .{ 5120, 64640, 1, 8, 12, 384, 122880 }, .{ 15360, 5120, 8, 4, 8, 256, 245760 }, }; for (seen) |r| { const p = plan(r[0], r[1]); diff --git a/zig/src/families/dsv41/kept_order.zig b/zig/src/families/dsv41/kept_order.zig new file mode 100644 index 0000000000..1240367657 --- /dev/null +++ b/zig/src/families/dsv41/kept_order.zig @@ -0,0 +1,102 @@ +//! Kept prompt state order on the ranks (lanes.zig): rank 0 copies before the frame; a failed follower restore stops. +const std = @import("std"); + +/// Saves kept state `sn`: room in `list`, this rank's copy, the frame, then `sn` into `list` (no entry on a failure). +pub fn save(comptime T: type, gpa: std.mem.Allocator, list: *std.ArrayList(*T), sn: *T, x: anytype) !void { + try list.ensureUnusedCapacity(gpa, 1); + try x.copy(); + try x.send(); + list.appendAssumeCapacity(sn); +} + +/// Restores a kept state on rank 0: its copies (`x.copy`) first, then the frame; a failed copy sends no frame. +pub fn restore(x: anytype) !void { + try x.copy(); + try x.send(); +} + +/// Restores kept state `id` on a follower with `x.copy`. An error stops the follower (rank 0 did this restore). +pub fn follow(comptime B: type, kept: *const std.AutoHashMapUnmanaged(u64, B), id: u64, x: anytype) !void { + const b = kept.get(id) orelse return error.UnknownSnapshot; + try x.copy(b); +} + +/// A test's copy and frame: the order of the calls, and the call that fails. +const Fake = struct { + log: *std.ArrayList(u8), + fail: u8 = 0, // 'c': the copy fails, 's': the frame fails + + fn copy(x: Fake) !void { + try x.log.append(std.testing.allocator, 'c'); + if (x.fail == 'c') return error.CopyFailed; + } + + fn send(x: Fake) !void { + try x.log.append(std.testing.allocator, 's'); + if (x.fail == 's') return error.LinkFailed; + } +}; + +test "a save whose copy fails sends no frame and keeps no state" { + const a = std.testing.allocator; + var log: std.ArrayList(u8) = .empty; + defer log.deinit(a); + var list: std.ArrayList(*u32) = .empty; + defer list.deinit(a); + var sn: u32 = 7; + try std.testing.expectError(error.CopyFailed, save(u32, a, &list, &sn, Fake{ .log = &log, .fail = 'c' })); + try std.testing.expectEqualStrings("c", log.items); + try std.testing.expectEqual(@as(usize, 0), list.items.len); +} + +test "a save whose frame fails keeps no state" { + const a = std.testing.allocator; + var log: std.ArrayList(u8) = .empty; + defer log.deinit(a); + var list: std.ArrayList(*u32) = .empty; + defer list.deinit(a); + var sn: u32 = 7; + try std.testing.expectError(error.LinkFailed, save(u32, a, &list, &sn, Fake{ .log = &log, .fail = 's' })); + try std.testing.expectEqual(@as(usize, 0), list.items.len); +} + +test "a save copies, then sends, then keeps the state" { + const a = std.testing.allocator; + var log: std.ArrayList(u8) = .empty; + defer log.deinit(a); + var list: std.ArrayList(*u32) = .empty; + defer list.deinit(a); + var sn: u32 = 7; + try save(u32, a, &list, &sn, Fake{ .log = &log }); + try std.testing.expectEqualStrings("cs", log.items); + try std.testing.expectEqual(@as(usize, 1), list.items.len); + try std.testing.expectEqual(&sn, list.items[0]); +} + +test "a restore whose copy fails on rank 0 sends no frame" { + const a = std.testing.allocator; + var log: std.ArrayList(u8) = .empty; + defer log.deinit(a); + try std.testing.expectError(error.CopyFailed, restore(Fake{ .log = &log, .fail = 'c' })); + try std.testing.expectEqualStrings("c", log.items); +} + +/// A follower's copy that fails. +const FailCopy = struct { + fn copy(_: FailCopy, _: u32) !void { + return error.CopyFailed; + } +}; + +test "a follower whose restore copy fails stops" { + const a = std.testing.allocator; + var kept: std.AutoHashMapUnmanaged(u64, u32) = .empty; + defer kept.deinit(a); + try kept.put(a, 3, 0); + try std.testing.expectError(error.CopyFailed, follow(u32, &kept, 3, FailCopy{})); +} + +test "a follower without the kept state stops" { + const kept: std.AutoHashMapUnmanaged(u64, u32) = .empty; + try std.testing.expectError(error.UnknownSnapshot, follow(u32, &kept, 3, FailCopy{})); +} diff --git a/zig/src/families/dsv41/lanes.zig b/zig/src/families/dsv41/lanes.zig index 3e01638f55..c8ef625c2c 100644 --- a/zig/src/families/dsv41/lanes.zig +++ b/zig/src/families/dsv41/lanes.zig @@ -20,6 +20,7 @@ const round = @import("round.zig"); const draft = @import("draft.zig"); const sampling = @import("sampling.zig"); const link = @import("link.zig"); +const kept_order = @import("kept_order.zig"); const cuda = @import("cuda"); const be = lanes.backend; @@ -520,16 +521,31 @@ pub const Lanes = struct { const replay = len -| m.cfg.window; var start: usize = 0; if (kept) |sn| { - if (try self.begin(.snap_restore)) { - try self.wire.int(gpa, u64, sn.id); - try self.wire.int(gpa, u32, @intCast(slot)); - try self.wire.int(gpa, u64, sn.base); - try self.wire.int(gpa, u64, base); - try self.wire.int(gpa, u64, sn.at); - try self.flush(); - } - if (base != sn.base) try m.copyExtent(sn.base, base, sn.at); - try m.snapCopy(slot, sn.buf.ptr, false); + // this rank's copies, then the frame (kept_order.zig): a failed copy here sends rank 1 no restore + const Restore = struct { + self: *Lanes, + sn: *Snap, + slot: usize, + base: usize, + + pub fn copy(x: @This()) !void { + if (x.base != x.sn.base) try x.self.m.copyExtent(x.sn.base, x.base, x.sn.at); + try x.self.m.snapCopy(x.slot, x.sn.buf.ptr, false); + } + + pub fn send(x: @This()) !void { + const g = x.self.gpa; + if (try x.self.begin(.snap_restore)) { + try x.self.wire.int(g, u64, x.sn.id); + try x.self.wire.int(g, u32, @intCast(x.slot)); + try x.self.wire.int(g, u64, x.sn.base); + try x.self.wire.int(g, u64, x.base); + try x.self.wire.int(g, u64, x.sn.at); + try x.self.flush(); + } + } + }; + try kept_order.restore(Restore{ .self = self, .sn = sn, .slot = slot, .base = base }); start = sn.at; s.cached = @intCast(sn.at); } @@ -823,14 +839,26 @@ pub const Lanes = struct { errdefer self.gpa.destroy(sn); sn.* = .{ .id = self.next_snap, .at = at, .base = l.base, .end = l.end, .buf = try cuda.DeviceBuffer.alloc(self.m.ctx.d, self.m.snapBytes()) }; errdefer sn.buf.free(); - try self.snaps.append(self.gpa, sn); - self.next_snap += 1; - if (try self.begin(.snap_save)) { - try self.wire.int(self.gpa, u64, sn.id); - try self.wire.int(self.gpa, u32, @intCast(l.slot)); - try self.flush(); - } - try self.m.snapCopy(l.slot, sn.buf.ptr, true); + self.next_snap += 1; // an id is never used again, also after a failed save + // copy, frame, list (kept_order.zig): a failed step leaves no freed state in self.snaps nor a save on rank 1 + const Save = struct { + self: *Lanes, + sn: *Snap, + slot: usize, + + pub fn copy(x: @This()) !void { + try x.self.m.snapCopy(x.slot, x.sn.buf.ptr, true); + } + + pub fn send(x: @This()) !void { + if (try x.self.begin(.snap_save)) { + try x.self.wire.int(x.self.gpa, u64, x.sn.id); + try x.self.wire.int(x.self.gpa, u32, @intCast(x.slot)); + try x.self.flush(); + } + } + }; + try kept_order.save(Snap, self.gpa, &self.snaps, sn, Save{ .self = self, .sn = sn, .slot = l.slot }); return sn; } @@ -1011,9 +1039,20 @@ pub fn follow(gpa: std.mem.Allocator, m: *Model) !void { const from: usize = @intCast(try r.int(u64)); const to: usize = @intCast(try r.int(u64)); const at: usize = @intCast(try r.int(u64)); - const b = kept.get(id) orelse return error.UnknownSnapshot; // rank 0 restores only states both ranks keep - if (from != to) m.copyExtent(from, to, at) catch |err| refused(err); - m.snapCopy(s, b.ptr, false) catch |err| refused(err); + // a failed copy stops this rank (kept_order.zig): its state would differ from rank 0's + const Follow = struct { + m: *Model, + s: usize, + from: usize, + to: usize, + at: usize, + + pub fn copy(x: @This(), b: cuda.DeviceBuffer) !void { + if (x.from != x.to) try x.m.copyExtent(x.from, x.to, x.at); + try x.m.snapCopy(x.s, b.ptr, false); + } + }; + try kept_order.follow(cuda.DeviceBuffer, &kept, id, Follow{ .m = m, .s = s, .from = from, .to = to, .at = at }); }, .snap_drop => { const id = try r.int(u64); diff --git a/zig/src/families/dsv41/round.zig b/zig/src/families/dsv41/round.zig index c674aa0c3f..59dac0c8de 100644 --- a/zig/src/families/dsv41/round.zig +++ b/zig/src/families/dsv41/round.zig @@ -446,9 +446,9 @@ pub const Round = struct { fn mark(rd: *Round, e: *const Engine, p: PhaseTimer.Phase) !void { if (rd.clk != 0 and rd.clk_marks and (rd.clk_ratio == null or rd.clk_ratio.? == rd.cur_ratio)) { const now = cuda.launch.count; - rd.launches[@intFromEnum(p)] += now -% rd.launch_mark; + rd.launches[@backingInt(p)] += now -% rd.launch_mark; rd.launch_mark = now; - try rd.clock(e, @intFromEnum(p)); + try rd.clock(e, @backingInt(p)); } if (rd.timer) |t| try t.mark(e.s, p); } @@ -801,7 +801,6 @@ fn linearBytes(l: weights.Linear) [2]u64 { return .{ l.words, @as(u64, l.k) * l.n * l.k2 / 16 }; } - /// A measurement's L2 prefetches left out (TF_DS_L2_SKIP, read at Model.open: bit 1 the "attn" fork after q, 2 the "moe" /// fork after wo_b, 4 the "next" fork after the experts); 0 in service. pub var l2_skip: u8 = 0; diff --git a/zig/src/families/dsv41/split_bench_main.zig b/zig/src/families/dsv41/split_bench_main.zig index fc8952dd2a..7980a7e5ef 100644 --- a/zig/src/families/dsv41/split_bench_main.zig +++ b/zig/src/families/dsv41/split_bench_main.zig @@ -320,10 +320,23 @@ pub fn main(init: std.process.Init) !u8 { try scales.upload(0, hs); } const run: Run = .{ - .t = t, .o = &o, .s = stream, .top = top, .n = n, .rb = rb, + .t = t, + .o = &o, + .s = stream, + .top = top, + .n = n, + .rb = rb, .kx = .{ .fp4 = .{ .codes = codes.ptr, .scales = scales.ptr } }, - .q = q.ptr, .w = w.ptr, .vis = vis.ptr, .nvis = nvis.ptr, .keys = keys.ptr, .tmax = tmax.ptr, .bufs = bufs, - .parts = parts.ptr, .merged = merged.ptr, .vblk = vblk.ptr, + .q = q.ptr, + .w = w.ptr, + .vis = vis.ptr, + .nvis = nvis.ptr, + .keys = keys.ptr, + .tmax = tmax.ptr, + .bufs = bufs, + .parts = parts.ptr, + .merged = merged.ptr, + .vblk = vblk.ptr, }; // every output byte, both splits, both patterns; the served keys past vis diff --git a/zig/src/families/dsv41/tri_basic.zig b/zig/src/families/dsv41/tri_basic.zig index 66cb3470a5..34acadee40 100644 --- a/zig/src/families/dsv41/tri_basic.zig +++ b/zig/src/families/dsv41/tri_basic.zig @@ -19,9 +19,9 @@ fn ci(name: []const u8, v: usize) aot.Const { pub fn hcPre(t: Tri, h: u64, fn_w: u64, scale: u64, base: u64, pre_in: u64, norm_w: u64, eps: f32, hc_eps: f32, iters: usize, out: u64, pre_out: u64, post: u64, comb: u64, part: u64, rows: usize, d: usize) !void { try t.run("_hc_partial", .{ u(rows), hc_blocks, 1 }, &.{ p("X", "*bf16", h), p("FN", "*fp32", fn_w), p("PART", "*fp32", part) }, &.{ ci("WIDE", 4 * d), ci("NB", hc_blocks), ci("SUB", 128) }); try t.run("_hc_finish", .{ u(rows), 1, 1 }, &.{ - p("X", "*bf16", h), p("PART", "*fp32", part), p("BASE", "*fp32", base), p("SCALE", "*fp32", scale), - p("PRE_IN", "*fp32", pre_in), p("NW", "*bf16", norm_w), p("OUT", "*bf16", out), p("PRE_OUT", "*fp32", pre_out), - p("POST", "*fp32", post), p("COMB", "*fp32", comb), aot.float("eps", eps), aot.float("hc_eps", hc_eps), + p("X", "*bf16", h), p("PART", "*fp32", part), p("BASE", "*fp32", base), p("SCALE", "*fp32", scale), + p("PRE_IN", "*fp32", pre_in), p("NW", "*bf16", norm_w), p("OUT", "*bf16", out), p("PRE_OUT", "*fp32", pre_out), + p("POST", "*fp32", post), p("COMB", "*fp32", comb), aot.float("eps", eps), aot.float("hc_eps", hc_eps), }, &.{ ci("D", d), ci("NB", hc_blocks), ci("ITERS", iters), ci("BLOCK", 1024) }); } diff --git a/zig/src/server/effort_values_test.zig b/zig/src/server/effort_values_test.zig index f8b01f8ae9..e3ec47b9ef 100644 --- a/zig/src/server/effort_values_test.zig +++ b/zig/src/server/effort_values_test.zig @@ -10,7 +10,7 @@ const Conn = @import("http_conn.zig").Conn; /// DeepSeek-V4.1's numbers (families/dsv41/text.zig; the served lane's EFFORT_VALUES). const numbers = [_]api.Effort{ - .{ .name = "minimal", .value = 25 }, .{ .name = "low", .value = 50 }, .{ .name = "medium", .value = 62 }, + .{ .name = "minimal", .value = 25 }, .{ .name = "low", .value = 50 }, .{ .name = "medium", .value = 62 }, .{ .name = "high", .value = 75 }, .{ .name = "xhigh", .value = 88 }, .{ .name = "max", .value = 100 }, };