diff --git a/zig/src/families/dsv41/checkpoint.zig b/zig/src/families/dsv41/checkpoint.zig index 30089f9150..2c9dcb77a2 100644 --- a/zig/src/families/dsv41/checkpoint.zig +++ b/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; diff --git a/zig/src/families/dsv41/checkpoint_test.zig b/zig/src/families/dsv41/checkpoint_test.zig new file mode 100644 index 0000000000..4b202c49c7 --- /dev/null +++ b/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)); +} 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/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/families/dsv41/weights.zig b/zig/src/families/dsv41/weights.zig index 1f8643359f..bc81cf5470 100644 --- a/zig/src/families/dsv41/weights.zig +++ b/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 }); 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 }, };