dsv41: kept prompt states: rank 0 copies before the frame, a failed follower restore stops

zig/src/families/dsv41/dsv41.zig

@@ -98,4 +98,5 @@ test {
     refAll(sampling);
     refAll(model);
     refAll(lanes);
+    refAll(@import("kept_order.zig"));
 }

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]);

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{}));
+}

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);

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;

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

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) });
 }
 

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 },
 };