dsv41 Zig: load the EXL3 2.8 bpw pack (bf16 router weights and confidence head, one source a tensor)

zig/src/families/dsv41/checkpoint.zig

@@ -49,7 +49,19 @@ pub const Shards = struct {
             sh.files[n] = try st.File.open(gpa, io, p);
             n += 1;
             var names = sh.files[n - 1].names.iterator();
-            while (names.next()) |kv| try sh.names.put(gpa, kv.key_ptr.*, @intCast(n - 1));
+            while (names.next()) |kv| {
+                const slot = try sh.names.getOrPut(gpa, kv.key_ptr.*);
+                if (slot.found_existing) {
+                    // same dtype and shape (1017 Engram tables): the later shard wins; else no single source
+                    const was = sh.files[slot.value_ptr.*].get(kv.key_ptr.*).?;
+                    const now = sh.files[n - 1].get(kv.key_ptr.*).?;
+                    if (was.dtype != now.dtype or was.rank != now.rank or !std.mem.eql(usize, was.shape[0..was.rank], now.shape[0..now.rank])) {
+                        std.log.warn("tensor {s} is in two shards with another dtype or shape: {s} and {s}", .{ kv.key_ptr.*, paths.items[slot.value_ptr.*], p });
+                        return error.DuplicateTensor;
+                    }
+                }
+                slot.value_ptr.* = @intCast(n - 1);
+            }
         }
         return sh;
     }
@@ -200,6 +212,8 @@ pub fn readFrom(sh: *const Shards, key: []const u8, bytes: ?[]const u8, out: []u
         // the 1017 quants hold the Engram q_weight / k_weight tables in bf16; the engine reads them as f32, and the
         // widening is exact (bf16 is an f32's top 16 bits), as the Python loader's .to(torch.float32) is.
         if (from == .bf16 and k.dtype.? == .f32) return widenBfloat(t.bytes, out);
+        // the 2.8 bpw pack's bf16 router weights read as f16, as Python's .to(torch.float16) (round to nearest even)
+        if (from == .bf16 and k.dtype.? == .f16) return narrowBfloat(t.bytes, out);
         return error.UnsupportedConversion;
     }
     const es = (try dtypeOf(t.dtype)).size();
@@ -371,6 +385,10 @@ test "a pack with the engine's own names reads them as before: the converter's n
     try std.testing.expectEqualSlices(u8, body[16..24], &f32_got);
 }
 
+test {
+    _ = @import("checkpoint_test.zig");
+}
+
 test "column, row and whole slices of a toy trellis and its scales" {
     const a = std.testing.allocator;
     const io = std.testing.io;

zig/src/families/dsv41/checkpoint_test.zig

@@ -0,0 +1,68 @@
+//! checkpoint.zig's 2.8 bpw pack test: bf16 router weights, an f32 router bias, and one name in two shards.
+
+const std = @import("std");
+const checkpoint = @import("checkpoint.zig");
+const DType = @import("rank_cache.zig").DType;
+const Shards = checkpoint.Shards;
+const describe = checkpoint.describe;
+const read = checkpoint.read;
+
+test "a 2.8 bpw pack: bf16 router weights read as f16, f32 router bias as f32, a name in two shards fails" {
+    const a = std.testing.allocator;
+    const io = std.testing.io;
+    var tmp = std.testing.tmpDir(.{ .iterate = true });
+    defer tmp.cleanup();
+    // the router weight in bf16 (0.5, -2.0, 0.0234375, 1.0) and the bias in f32 (0.25, -1.5)
+    const vals = [_]f32{ 0.5, -2.0, 0.0234375, 1.0 };
+    var body: [16]u8 = undefined;
+    for (vals, 0..) |v, i| std.mem.writeInt(u16, body[2 * i ..][0..2], @truncate(@as(u32, @bitCast(v)) >> 16), .little);
+    std.mem.writeInt(u32, body[8..12], @bitCast(@as(f32, 0.25)), .little);
+    std.mem.writeInt(u32, body[12..16], @bitCast(@as(f32, -1.5)), .little);
+    const json = "{\"layers.0.ffn.gate.weight\":{\"dtype\":\"BF16\",\"shape\":[2,2],\"data_offsets\":[0,8]}," ++
+        "\"layers.0.ffn.gate.bias\":{\"dtype\":\"F32\",\"shape\":[2],\"data_offsets\":[8,16]}}";
+    var file: std.ArrayList(u8) = .empty;
+    defer file.deinit(a);
+    var len: [8]u8 = undefined;
+    std.mem.writeInt(u64, &len, json.len, .little);
+    try file.appendSlice(a, &len);
+    try file.appendSlice(a, json);
+    try file.appendSlice(a, &body);
+    try tmp.dir.writeFile(io, .{ .sub_path = "layer-00.safetensors", .data = file.items });
+    const dir = try std.fs.path.join(a, &.{ ".zig-cache", "tmp", &tmp.sub_path });
+    defer a.free(dir);
+    {
+        var sh = try Shards.open(a, io, dir);
+        defer sh.close(io);
+        const w = try describe(&sh, "layers.0.ffn.gate.weight|torch.float16");
+        try std.testing.expectEqual(DType.f16, w.dtype);
+        try std.testing.expectEqual(@as(usize, 8), w.bytes);
+        var half: [8]u8 = undefined;
+        try read(&sh, "layers.0.ffn.gate.weight|torch.float16", &half);
+        for (vals, 0..) |v, i| {
+            const h: f16 = @bitCast(std.mem.readInt(u16, half[2 * i ..][0..2], .little));
+            try std.testing.expectEqual(v, @as(f32, h));
+        }
+        var bias: [8]u8 = undefined;
+        try read(&sh, "layers.0.ffn.gate.bias|torch.float32", &bias);
+        try std.testing.expectEqualSlices(u8, body[8..16], &bias);
+    }
+    // the same names, dtypes and shapes in a second shard (the 1017 packs' repeated Engram tables): the later one wins
+    try tmp.dir.writeFile(io, .{ .sub_path = "layer-01.safetensors", .data = file.items });
+    {
+        var sh = try Shards.open(a, io, dir);
+        defer sh.close(io);
+        var bias: [8]u8 = undefined;
+        try read(&sh, "layers.0.ffn.gate.bias|torch.float32", &bias);
+        try std.testing.expectEqualSlices(u8, body[8..16], &bias);
+    }
+    // the router weight again in f16 (a dtype-override file next to the pack): no single source, so the open fails
+    const json16 = "{\"layers.0.ffn.gate.weight\":{\"dtype\":\"F16\",\"shape\":[2,2],\"data_offsets\":[0,8]}}";
+    var over: std.ArrayList(u8) = .empty;
+    defer over.deinit(a);
+    std.mem.writeInt(u64, &len, json16.len, .little);
+    try over.appendSlice(a, &len);
+    try over.appendSlice(a, json16);
+    try over.appendSlice(a, body[0..8]);
+    try tmp.dir.writeFile(io, .{ .sub_path = "zz-override.safetensors", .data = over.items });
+    try std.testing.expectError(error.DuplicateTensor, Shards.open(a, io, dir));
+}

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/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/families/dsv41/weights.zig

@@ -291,10 +291,8 @@ const Loader = struct {
         return L.upload(tensor, try L.entry(try plan.plainKey(&kb, tensor, dtype)));
     }
 
-    /// The MTP markov head, which the kernel reads as fp16: a pack that holds it in bf16 (the 1017 checkpoint's own
-    /// dtype; the converter casts it to f16, a community quant may not) is converted exactly at load, and an f16 head
-    /// is uploaded as it is. The cache key stays the checkpoint's own dtype, so no rank cache is invalidated.
-    fn markovHead(L: *Loader, tensor: []const u8) !u64 {
+    /// An fp16 kernel input (markov or confidence head): bf16 converts at load; the cache key keeps the pack dtype.
+    fn halfAtLoad(L: *Loader, tensor: []const u8) !u64 {
         var kb: [256]u8 = undefined;
         const key = try plan.plainKey(&kb, tensor, null);
         const src = try L.entry(key);
@@ -638,8 +636,8 @@ pub fn load(gpa: std.mem.Allocator, d: *const cuda.Driver, src: Source, c: Confi
             .main_norm = try L.plain("mtp.0.main_norm.weight", null),
             .norm = try L.plain(try std.fmt.bufPrint(&tb, "{s}.norm.weight", .{last}), null),
             .markov_embed = try L.plain(try std.fmt.bufPrint(&tb, "{s}.markov_head.embed.weight", .{last}), null),
-            .markov_head = try L.markovHead(try std.fmt.bufPrint(&tb, "{s}.markov_head.head.weight", .{last})),
-            .conf = try L.plain(try std.fmt.bufPrint(&tb, "{s}.confidence_head.proj.weight", .{last}), null),
+            .markov_head = try L.halfAtLoad(try std.fmt.bufPrint(&tb, "{s}.markov_head.head.weight", .{last})),
+            .conf = try L.halfAtLoad(try std.fmt.bufPrint(&tb, "{s}.confidence_head.proj.weight", .{last})),
         };
     }
     if (L.span_bytes > 0) std.log.info("rank cache: {d:.1} GiB of expert entries read whole in {d:.1} s", .{ @as(f64, @floatFromInt(L.span_bytes)) / (1 << 30), @as(f64, @floatFromInt(L.span_ns)) / 1e9 });

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