diff --git a/zig/src/families/dsv41/check_main.zig b/zig/src/families/dsv41/check_main.zig index 6ac51fafcf..e6b2ba1091 100644 --- a/zig/src/families/dsv41/check_main.zig +++ b/zig/src/families/dsv41/check_main.zig @@ -37,7 +37,8 @@ pub fn main(init: std.process.Init) !u8 { defer ix.deinit(); // four ranks: the exact 2D split (node g = TP2 rank g % 2 of pair g / 2); the cache file is still rank g of 4 const split: dsv41.plan.Split = if (world == 4) .{ .rank = rank % 2, .world = 2, .pair = rank / 2, .parity = parity } else .{ .rank = rank, .world = world }; - const want = try dsv41.plan.wants(a, cfg, split, dspark); + // a mixed pack's dense linears: planned as the cache holds them + const want = try dsv41.plan.wantsDense(a, cfg, split, dspark, dsv41.plan.Dense.inCache(&ix)); var out_buf: [4096]u8 = undefined; var out = std.Io.File.stdout().writer(io, &out_buf); const w = &out.interface; diff --git a/zig/src/families/dsv41/checkpoint.zig b/zig/src/families/dsv41/checkpoint.zig index 383f701592..883ee48290 100644 --- a/zig/src/families/dsv41/checkpoint.zig +++ b/zig/src/families/dsv41/checkpoint.zig @@ -79,6 +79,20 @@ pub const Shards = struct { const i = sh.names.get(name) orelse return null; return sh.files[i].get(name); } + + /// The linears this checkpoint holds dense (plan.Dense): a `.weight` and no `.trellis`. + pub fn dense(sh: *const Shards) plan.Dense { + return .{ .ctx = sh, .has = hasDense }; + } + + fn hasDense(ctx: *const anyopaque, prefix: []const u8) bool { + const sh: *const Shards = @ptrCast(@alignCast(ctx)); + var buf: [256]u8 = undefined; + const tr = std.fmt.bufPrint(&buf, "{s}.trellis", .{prefix}) catch return false; + if (sh.names.contains(tr)) return false; + const w = std.fmt.bufPrint(&buf, "{s}.weight", .{prefix}) catch return false; + return sh.names.contains(w); + } }; /// A loader key split into its fields: "name|dtype" (plain) or "prefix|cols|rows|part" (an EXL3 part). diff --git a/zig/src/families/dsv41/cublas.zig b/zig/src/families/dsv41/cublas.zig index 0411372578..f5e47e31be 100644 --- a/zig/src/families/dsv41/cublas.zig +++ b/zig/src/families/dsv41/cublas.zig @@ -73,6 +73,15 @@ pub const Blas = struct { try b.check(b.api.cublasGemmEx(b.handle, .t, .n, @intCast(n), @intCast(rows), @intCast(k), &one, w, bf16_type, @intCast(k), x, bf16_type, @intCast(k), &zero, out, bf16_type, @intCast(n), 68, 99)); } + /// A dense linear (weights.Linear.dense): x bf16 [rows, k] @ w bf16 [n, k]^T into bf16 or fp32 out, as linearBf16. + pub fn linearDense(b: *const Blas, x: u64, ldx: usize, w: u64, out: u64, ldo: usize, out_f32: bool, rows: usize, k: usize, n: usize) Error!void { + const one: f32 = 1; + const zero: f32 = 0; + const bf16_type: c_int = 14; // CUDA_R_16BF + const out_type: c_int = if (out_f32) 0 else bf16_type; // CUDA_R_32F + try b.check(b.api.cublasGemmEx(b.handle, .t, .n, @intCast(n), @intCast(rows), @intCast(k), &one, w, bf16_type, @intCast(k), x, bf16_type, @intCast(ldx), &zero, out, out_type, @intCast(ldo), 68, 99)); + } + /// torch.matmul of x [rows, k] fp32 (row-major) and w [n, k] fp32 transposed: out [rows, n] fp32, as torch issues /// it (column-major C^T = W x^T: transa T, transb N, m = n, n = rows, lda = ldb = k, ldc = n, alpha 1, beta 0). pub fn xwT(b: *const Blas, x: u64, w: u64, out: u64, rows: usize, k: usize, n: usize) Error!void { diff --git a/zig/src/families/dsv41/exl3_linear.zig b/zig/src/families/dsv41/exl3_linear.zig index 25a9241925..47b7d75538 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; @@ -224,6 +224,7 @@ pub fn glinear(k: *const Kernels, stream: cuda.Stream, calls: []const Call, m: u pub fn zFloats(layers: []const weights.Linear, m: usize) usize { var n: usize = 0; for (layers) |l| { + if (l.isDense()) continue; // runs through cuBLAS (prompt.denseRows) const p = plan(l.k, planN(l)); if (p.sk > 1) n += p.sk * m * l.n; } @@ -243,12 +244,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/plan.zig b/zig/src/families/dsv41/plan.zig index 0b045244cf..69741e570c 100644 --- a/zig/src/families/dsv41/plan.zig +++ b/zig/src/families/dsv41/plan.zig @@ -7,7 +7,8 @@ //! and keeps pair p's whole 128-column blocks of each of its rank's split outputs; DSpark stays TP2 inside the pair. const std = @import("std"); const Config = @import("config.zig").Config; -const DType = @import("rank_cache.zig").DType; +const rank_cache = @import("rank_cache.zig"); +const DType = rank_cache.DType; /// What a split means for one rank. pub const Split = struct { @@ -138,9 +139,38 @@ pub fn partKey(buf: []u8, prefix: []const u8, cols: ?[2]usize, rows: ?[2]usize, return std.fmt.bufPrint(buf, "{s}|{s}|{s}|{s}", .{ prefix, cs, rs, part }); } +/// A dense linear's key: its plain `.weight` in the checkpoint's own dtype (bf16 [n, k]). +pub fn denseKey(buf: []u8, prefix: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}.weight|None", .{prefix}); +} + +/// Whether a dense linear has a run path: not wo_a's slices (Round.init reads their suh; decode runs groupedRotated). +pub fn denseRunnable(prefix: []const u8) bool { + return std.mem.indexOf(u8, prefix, ".attn.wo_a.") == null; +} + +/// Which whole linears a mixed pack holds dense (`has(ctx, prefix)`); a split (2D / TP part) dense linear is refused. +pub const Dense = struct { + ctx: *const anyopaque, + has: *const fn (ctx: *const anyopaque, prefix: []const u8) bool, + + /// The linears a rank cache holds dense: a dense entry and no trellis part. + pub fn inCache(ix: *const rank_cache.Index) Dense { + return .{ .ctx = ix, .has = cacheHas }; + } + + fn cacheHas(ctx: *const anyopaque, prefix: []const u8) bool { + const ix: *const rank_cache.Index = @ptrCast(@alignCast(ctx)); + var buf: [256]u8 = undefined; + if (ix.get(partKey(&buf, prefix, null, null, "tr") catch return false) != null) return false; + return ix.get(denseKey(&buf, prefix) catch return false) != null; + } +}; + const Builder = struct { a: std.mem.Allocator, out: std.ArrayList(Want) = .empty, + dense: ?Dense = null, fn plain(b: *Builder, tensor: []const u8, dtype: ?DType) !void { var buf: [256]u8 = undefined; @@ -148,6 +178,12 @@ const Builder = struct { } fn exl3(b: *Builder, prefix: []const u8, cols: ?[2]usize, rows: ?[2]usize, k: ?usize, n: ?usize) !void { + if (b.dense) |d| if (d.has(d.ctx, prefix)) { + if (cols != null or rows != null) return error.DenseLinearSplit; + if (!denseRunnable(prefix)) return error.DenseLinearUnsupported; + var buf: [256]u8 = undefined; + return b.out.append(b.a, .{ .key = try b.a.dupe(u8, try denseKey(&buf, prefix)), .dtype = .bf16 }); + }; const k16: ?usize = if (rows) |r| (r[1] - r[0]) / 16 else if (k) |x| x / 16 else null; const n16: ?usize = if (cols) |c| (c[1] - c[0]) / 16 else if (n) |x| x / 16 else null; for ([_][]const u8{ "tr", "suh", "svh" }) |part| { @@ -220,7 +256,12 @@ fn block(b: *Builder, c: Config, s: Split, p: []const u8, i: usize, experts: usi /// Every tensor rank `s` loads, in load order (`dspark`: the DSpark stages and heads too). Keys live in `a`. pub fn wants(a: std.mem.Allocator, c: Config, s: Split, dspark: bool) ![]Want { - var b: Builder = .{ .a = a }; + return wantsDense(a, c, s, dspark, null); +} + +/// wants, the linears `dense` names planned as their dense entries (a mixed pack). +pub fn wantsDense(a: std.mem.Allocator, c: Config, s: Split, dspark: bool, dense: ?Dense) ![]Want { + var b: Builder = .{ .a = a, .dense = dense }; try b.plain("embed.weight", null); try b.plain("norm.weight", null); try b.exl3("head", s.headCols(c), null, c.hidden, null); @@ -266,6 +307,43 @@ test "a rank's tensors in load order: the split's ranges, keys as Python prints try std.testing.expectEqualStrings("mtp.2.confidence_head.proj.weight|None", w[w.len - 1].key); } +test "a mixed pack's dense linears: one bf16 .weight entry for a whole linear, a split one refused" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var why: @import("config.zig").Why = .{}; + const c = try @import("config.zig").parse(a, @import("config.zig").test_config, &why); + const names = struct { + fn has(_: *const anyopaque, prefix: []const u8) bool { + return std.mem.eql(u8, prefix, "layers.2.attn.wq_a") or std.mem.eql(u8, prefix, "layers.2.attn.indexer.wq_b"); + } + fn split(_: *const anyopaque, prefix: []const u8) bool { + return std.mem.eql(u8, prefix, "layers.0.attn.wq_b"); + } + fn woA(_: *const anyopaque, prefix: []const u8) bool { + return std.mem.eql(u8, prefix, "layers.0.attn.wo_a.slice.0"); + } + }; + const ctx: u8 = 0; + const w = try wantsDense(a, c, .{ .rank = 0, .world = 2 }, true, .{ .ctx = &ctx, .has = names.has }); + const plain = try wants(a, c, .{ .rank = 0, .world = 2 }, true); + // each dense linear: its three parts give way to one entry + try std.testing.expectEqual(plain.len - 2 * 2, w.len); + var seen: usize = 0; + for (w) |x| { + if (std.mem.startsWith(u8, x.key, "layers.2.attn.wq_a|") or std.mem.startsWith(u8, x.key, "layers.2.attn.indexer.wq_b|")) return error.PartsLeft; + if (std.mem.eql(u8, x.key, "layers.2.attn.wq_a.weight|None") or std.mem.eql(u8, x.key, "layers.2.attn.indexer.wq_b.weight|None")) { + try std.testing.expectEqual(@as(?DType, .bf16), x.dtype); + seen += 1; + } + } + try std.testing.expectEqual(@as(usize, 2), seen); + // wq_b is a column slice of the rank's heads: no dense form of it + try std.testing.expectError(error.DenseLinearSplit, wantsDense(a, c, .{ .rank = 0, .world = 2 }, true, .{ .ctx = &ctx, .has = names.split })); + // a whole wo_a slice has no dense run path yet: refused at plan time, not at Round.init + try std.testing.expectError(error.DenseLinearUnsupported, wantsDense(a, c, .{ .rank = 0, .world = 2 }, true, .{ .ctx = &ctx, .has = names.woA })); +} + test "each split tensor's ranges: TP2's slices, a 2D node's pair parts of them" { var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); diff --git a/zig/src/families/dsv41/prompt.zig b/zig/src/families/dsv41/prompt.zig index 2fb961c169..084edd0e8d 100644 --- a/zig/src/families/dsv41/prompt.zig +++ b/zig/src/families/dsv41/prompt.zig @@ -498,6 +498,7 @@ fn layerLinears(e: *const Engine, ch: *Chunk, a: *Arena, lay: weights.Layer, max /// One linear alone at 128 rows in the chunk's grouped-linear sizes, and its split-K counters (int32 [8 N / 128], zeros, /// left zero by each launch). fn addLinear(e: *const Engine, ch: *Chunk, a: *Arena, l: weights.Linear, max_gk: *usize, max_gz: *usize) !void { + if (l.isDense()) return; // cuBLAS: no rotated rows, partials or counters max_gk.* = @max(max_gk.*, l.k); max_gz.* = @max(max_gz.*, exl3_linear.zFloats(&[_]weights.Linear{l}, 128)); if (ch.n_counters == ch.counters.len) return error.TooManyLinears; @@ -555,19 +556,38 @@ pub fn counterOf(ch: *const Chunk, l: weights.Linear) !u64 { pub fn grouped(e: *const Engine, ch: *const Chunk, m: usize, ls: []const weights.Linear, xs: []const u64, ldxs: []const usize, outs: []const u64, ldos: []const usize, types: []const exl3_linear.DType) !void { if (m == 0 or m > 128 or ls.len > 8) return error.NotADecodeWindow; var calls: [8]exl3_linear.Call = undefined; + var n: usize = 0; var off: usize = 0; for (ls, 0..) |l, i| { - calls[i] = .{ .layer = l, .x = xs[i], .ldx = @intCast(ldxs[i]), .x_dtype = .bf16, .xh = ch.gxh + off * 2, .y = outs[i], .ldy = @intCast(ldos[i]), .y_dtype = types[i], .counters = try counterOf(ch, l) }; + // a mixed pack's dense linears through cuBLAS, the group's EXL3 ones as before + if (l.isDense()) { + try denseRows(e, m, l, xs[i], ldxs[i], outs[i], ldos[i], types[i]); + continue; + } + calls[n] = .{ .layer = l, .x = xs[i], .ldx = @intCast(ldxs[i]), .x_dtype = .bf16, .xh = ch.gxh + off * 2, .y = outs[i], .ldy = @intCast(ldos[i]), .y_dtype = types[i], .counters = try counterOf(ch, l) }; off += m * l.k; + n += 1; } - try exl3_linear.rotMany(e.lin, e.s, calls[0..ls.len], m, true); - try exl3_linear.glinear(e.lin, e.s, calls[0..ls.len], m, ch.gz, true, true); + if (n == 0) return; + try exl3_linear.rotMany(e.lin, e.s, calls[0..n], m, true); + try exl3_linear.glinear(e.lin, e.s, calls[0..n], m, ch.gz, true, true); +} + +/// A dense linear (weights.Linear.dense) on m rows: x bf16 (row stride ldx) into out (row stride ldo, bf16 or fp32). +pub fn denseRows(e: *const Engine, m: usize, l: weights.Linear, x: u64, ldx: usize, out: u64, ldo: usize, dt: exl3_linear.DType) !void { + if (!l.isDense()) return error.NotDense; + const f32_out = switch (dt) { + .bf16 => false, + .f32 => true, + .f16 => return error.UnsupportedDType, + }; + try e.blas.linearDense(x, ldx, l.dense, out, ldo, f32_out, m, l.k, l.n); } -/// Exl3Group.rotated on m rows already rotated into each call's xh (a kernel that folds the rotation in wrote them): -/// the glinear launches alone, each call's counters and folds as given. +/// Exl3Group.rotated on m rows already rotated into each call's xh: the glinear launches alone (dense: denseRows). pub fn groupedRotated(e: *const Engine, ch: *const Chunk, m: usize, calls: []exl3_linear.Call) !void { if (m == 0 or m > 128 or calls.len > exl3_linear.gmax) return error.NotADecodeWindow; + for (calls) |c| if (c.layer.isDense()) return error.DenseNotRotated; for (calls) |*c| c.counters = try counterOf(ch, c.layer); try exl3_linear.glinear(e.lin, e.s, calls, m, ch.gz, true, true); } @@ -711,11 +731,12 @@ pub fn attnMixes(e: *const Engine, ch: *Chunk, li: usize) !void { /// mm(layer, x, out_dtype) of m rows: the prompt GEMM beyond 128 rows, the layer's own group (GROUPED) up to 128. fn mmRows(e: *const Engine, ch: *const Chunk, m: usize, l: weights.Linear, x: u64, ldx: usize, out: u64, out_type: tri_markov.OutType, os: usize) !void { - if (m > 128) return exl3_prefill.matmul(e.pf, e.t, ch.ws, l, x, .bf16, ldx, out, out_type, os, m); const dt: exl3_linear.DType = switch (out_type) { .bf16 => .bf16, .fp32 => .f32, }; + if (l.isDense()) return denseRows(e, m, l, x, ldx, out, os, dt); + if (m > 128) return exl3_prefill.matmul(e.pf, e.t, ch.ws, l, x, .bf16, ldx, out, out_type, os, m); try grouped(e, ch, m, &.{l}, &.{x}, &.{ldx}, &.{out}, &.{os}, &.{dt}); } diff --git a/zig/src/families/dsv41/rank_cache_main.zig b/zig/src/families/dsv41/rank_cache_main.zig index 879be3ccdc..2a1a3e4be2 100644 --- a/zig/src/families/dsv41/rank_cache_main.zig +++ b/zig/src/families/dsv41/rank_cache_main.zig @@ -49,9 +49,10 @@ pub fn main(init: std.process.Init) !u8 { // four ranks: the exact 2D split (node g = TP2 rank g % 2 of pair g / 2); the file is still rank g of 4 const split: dsv41.plan.Split = if (world == 4) .{ .rank = rank % 2, .world = 2, .pair = rank / 2, .parity = parity } else .{ .rank = rank, .world = world }; const cfg = try dsv41.Config.read(a, io, args[1]); - const want = try dsv41.plan.wants(a, cfg, split, dspark); var sh = try dsv41.checkpoint.Shards.open(gpa, io, args[1]); defer sh.close(io); + // a mixed pack's dense linears (a plain .weight, no trellis): one bf16 entry each + const want = try dsv41.plan.wantsDense(a, cfg, split, dspark, sh.dense()); // every entry's layout first, and from the layouts the index: a missing tensor or a layout the split does not // expect stops before a byte is written diff --git a/zig/src/families/dsv41/round.zig b/zig/src/families/dsv41/round.zig index c674aa0c3f..ba8f4ec1fd 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); } @@ -590,6 +590,7 @@ fn attention(e: *const Engine, rd: *Round, ch: *const Chunk, cs: *const Caches, const outs = [_]u64{ rd.q, rd.iq }; for (qls) |ol| { const l = ol orelse continue; + if (l.isDense()) continue; // a mixed pack's dense linear: from q's normed rows below const xh = rd.xhq + n * R * c.q_lora * 2; rot[n] = .{ .suh = l.suh, .h = xh }; calls[n] = .{ .layer = l, .x = 0, .ldx = 0, .x_dtype = .bf16, .xh = xh, .y = outs[n], .ldy = @intCast(l.n), .y_dtype = .bf16, .counters = 0 }; @@ -598,6 +599,10 @@ fn attention(e: *const Engine, rd: *Round, ch: *const Chunk, cs: *const Caches, } try tri_norm.qKvNorm(t, rd.qa, lay.wq_a.n, lay.q_norm, c.eps, rot[0..n], rd.ykv, lay.kv_norm, rope.cos, rope.sin, pos, ring, ring_rows, rd.g(.wslot), true, rd_dim, rd.qr, rd.kout, R, lay.wq_a.n, hd); try prompt.groupedRotated(e, ch, R, calls[0..n]); + for (qls, outs) |ol, out| { + const l = ol orelse continue; + if (l.isDense()) try prompt.denseRows(e, R, l, rd.qr, lay.wq_a.n, out, l.n, .bf16); + } } // the prefetch: this layer's wo_a and wo_b behind q's linears { @@ -798,10 +803,10 @@ fn moe(e: *const Engine, rd: *Round, ch: *const Chunk, li: usize) !void { /// An EXL3 linear's words: their first byte and length (k * n values of k2 / 2 bits). fn linearBytes(l: weights.Linear) [2]u64 { + if (l.isDense()) return .{ l.dense, @as(u64, l.k) * l.n * 2 }; 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 3997ca531d..c25d23c84a 100644 --- a/zig/src/families/dsv41/weights.zig +++ b/zig/src/families/dsv41/weights.zig @@ -14,10 +14,23 @@ const max_span = 4 << 30; /// The size of one direct read within such a stretch. const span_piece = 64 << 20; -/// An EXL3 linear in the strips layout: words int32 [N/128, K/16, 8, 8 * bits], fp16 suh [K] and svh [N]. -/// plan_n: the N its decode plan is made for, when not n: a 2D node's column part of a TP2 slice keeps that slice's -/// plan (linear.py pinned plans), so its columns sum in the slice's K order. -pub const Linear = struct { words: u64, suh: u64, svh: u64, k: u32, n: u32, k2: u32, plan_n: u32 = 0 }; +/// An EXL3 linear in the strips layout: words int32 [N/128, K/16, 8, 8 * bits], fp16 suh [K] and svh [N]; or dense. +pub const Linear = struct { + words: u64, + suh: u64, + svh: u64, + k: u32, + n: u32, + k2: u32, + /// The N of its decode plan when not n: a 2D column part keeps its TP2 slice's plan (linear.py pinned plans). + plan_n: u32 = 0, + /// A mixed pack's plain bf16 `.weight` [n, k] row-major (EAQ v1.1); words, suh, svh and k2 are then 0. + dense: u64 = 0, + + pub fn isDense(l: Linear) bool { + return l.dense != 0; + } +}; /// A layer's routed experts and its shared expert (the last entry), as exl3_experts.prepare lays them out. On a 2D node /// (exl3/experts2d.py prepare) gate / up hold its intermediate blocks (width of the half's down_k) and down its @@ -305,10 +318,14 @@ const Loader = struct { return L.upload(tensor, L.aux.items); } - /// linear(prefix, cols, rows): Exl3Linear.from_tensors in strips. + /// linear(prefix, cols, rows): Exl3Linear.from_tensors in strips, or a whole linear the pack holds dense. fn linear(L: *Loader, prefix: []const u8, cols: ?[2]usize, rows: ?[2]usize) !Linear { var kb: [256]u8 = undefined; var nb: [256]u8 = undefined; + if (cols == null and rows == null) if (try L.denseLinear(prefix)) |d| { + if (!plan.denseRunnable(prefix)) return error.DenseLinearUnsupported; + return d; + }; const tr = try L.entry(try plan.partKey(&kb, prefix, cols, rows, "tr")); const t = try L.shape(try plan.partKey(&kb, prefix, cols, rows, "tr")); const kt = t[0]; @@ -323,6 +340,40 @@ const Loader = struct { return .{ .words = words, .suh = suh, .svh = svh, .k = @intCast(kt * 16), .n = @intCast(nt * 16), .k2 = @intCast(w16 / 8) }; } + /// A whole linear the source holds as a plain bf16 `.weight` [n, k], uploaded as it is; null: a trellis. + fn denseLinear(L: *Loader, prefix: []const u8) !?Linear { + var kb: [256]u8 = undefined; + var nb: [256]u8 = undefined; + if (L.has(try plan.partKey(&kb, prefix, null, null, "tr"))) return null; + const key = try plan.denseKey(&kb, prefix); + if (!L.has(key)) return null; + const s = try L.dims(key); + if (s.dtype != .bf16 or s.rank != 2 or s.n % 8 != 0 or s.k % 8 != 0) return error.UnexpectedTensor; + const w = try L.upload(try std.fmt.bufPrint(&nb, "{s}.weight", .{prefix}), try L.entry(key)); + return .{ .words = 0, .suh = 0, .svh = 0, .k = @intCast(s.k), .n = @intCast(s.n), .k2 = 0, .dense = w }; + } + + fn has(L: *Loader, key: []const u8) bool { + return switch (L.src) { + .cache => |c| c.index.get(key) != null, + .shards => |sh| if (checkpoint.describe(sh, key)) |_| true else |_| false, + }; + } + + /// An entry's dtype, rank and first two dimensions ([n, k] for a dense linear). + fn dims(L: *Loader, key: []const u8) !struct { dtype: rank_cache.DType, rank: u8, n: usize, k: usize } { + return switch (L.src) { + .cache => |c| blk: { + const e = c.index.get(key) orelse return error.MissingEntry; + break :blk .{ .dtype = e.dtype, .rank = e.rank, .n = e.shape[0], .k = e.shape[1] }; + }, + .shards => |sh| blk: { + const s = try checkpoint.describe(sh, key); + break :blk .{ .dtype = s.dtype, .rank = s.rank, .n = s.shape[0], .k = s.shape[1] }; + }, + }; + } + /// A trellis entry's [kt, nt, w16]. fn shape(L: *Loader, key: []const u8) ![3]usize { return switch (L.src) { 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 }, };