dsv41 Zig: dense bf16 linears of a mixed EXL3 pack (EAQ v1.1 indexer path)

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;

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 `<prefix>.weight` and no `<prefix>.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).

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 {

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

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,16 +139,51 @@ 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 `<prefix>.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;
         try b.out.append(b.a, .{ .key = try b.a.dupe(u8, try plainKey(&buf, tensor, dtype)), .dtype = dtype });
     }
 
     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();

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

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 <prefix>.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