From f9fc1a786d51e289411989d808295d2895f9c272 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 15 Oct 2026 08:53:05 +0200 Subject: [PATCH 01/34] glm: group-limited top-k routing (Ling's 8 groups, keep 4) and a per-layer shared SwiGLU limit, GLM's bytes unchanged tf_glm_topk takes n_group and topk_group: with more than one group it keeps the topk_group groups of the best top-2 biased sums (lower group on ties) before the top-k. GLM's one group skips that path. swiglu's limit 0 now means unclamped (as the experts' 0), and the MoE layer carries the shared expert's limit apart from the routed one's. Tests: glm-ref adds group routing vs float64 and the swiglu limit; glm-route-old compares GLM's select and swiglu with 6380c8da's fatbin; glm-route-ling checks captured Ling layer-4 routing against HF and exllamav3. --- zig/kernels/cuda/glm_glue_moe.cuh | 47 +++- zig/src/families/glm/cuda_glue.zig | 19 +- zig/src/families/glm/cuda_moe.zig | 6 +- zig/src/families/glm/cuda_weights.zig | 1 + zig/tests/cuda/glm_ref_route.zig | 356 ++++++++++++++++++++++++++ zig/tests/cuda/glm_ref_tests.zig | 4 +- zig/tests/cuda/main.zig | 5 + 7 files changed, 428 insertions(+), 10 deletions(-) create mode 100644 zig/tests/cuda/glm_ref_route.zig diff --git a/zig/kernels/cuda/glm_glue_moe.cuh b/zig/kernels/cuda/glm_glue_moe.cuh index 7f2a27d6b0..d9cd649159 100644 --- a/zig/kernels/cuda/glm_glue_moe.cuh +++ b/zig/kernels/cuda/glm_glue_moe.cuh @@ -23,6 +23,48 @@ __device__ __forceinline__ float swiglu1(float g, float u, float limit) { return rbf(__fmul_rn(u, s)); } +// Group-limited routing: -inf on every expert outside the `keep` groups of the best top-2 sums (ties to lower group). +template +__device__ __forceinline__ void keep_groups(float* choice, int lane, int ne, int gs, int ng, int keep) { + float gscore = -INFINITY; // lane g: group g's top-2 sum + for (int g = 0; g < ng; ++g) { + float a = -INFINITY, b = -INFINITY; // this lane's best two in group g, a >= b + for (int i = 0; i < MAXE; ++i) { + const int e = i * 32 + lane; + if (e < ne && e / gs == g) { + const float c = choice[i]; + if (c > a) { + b = a; + a = c; + } else if (c > b) { + b = c; + } + } + } + for (int o = 16; o > 0; o >>= 1) { + const float oa = __shfl_xor_sync(0xffffffffu, a, o), ob = __shfl_xor_sync(0xffffffffu, b, o); + b = fmaxf(fminf(a, oa), fmaxf(b, ob)); + a = fmaxf(a, oa); + } + if (lane == g) gscore = a + b; + } + bool live = lane < ng; + unsigned kept = 0u; + for (int k = 0; k < keep; ++k) { + float mx = live ? gscore : -INFINITY; + for (int o = 16; o > 0; o >>= 1) mx = fmaxf(mx, __shfl_xor_sync(0xffffffffu, mx, o)); + const unsigned at = __ballot_sync(0xffffffffu, live && gscore == mx); + if (at == 0u) break; + const int id = __ffs(at) - 1; + kept |= 1u << id; + if (lane == id) live = false; + } + for (int i = 0; i < MAXE; ++i) { + const int e = i * 32 + lane; + if (e < ne && !((kept >> (e / gs)) & 1u)) choice[i] = -INFINITY; + } +} + } // namespace glm_moe // glue._swiglu with BLOCK 512 (128 threads, 4 values each, 64-group sums over 16 lanes), column block blockIdx.y. @@ -66,10 +108,10 @@ extern "C" __global__ void __launch_bounds__(128) if ((threadIdx.x & 7) == 0 && g < kg) xs[m * kg + g] = s; } -// glue._topk, a warp a row: top_k by sigmoid + bias (ties to lower id), normed if `norm`, x scale; then shared expert. +// glue._topk, a warp a row: top_k of sigmoid + bias (lower id on ties) in the kept groups, normed, x scale; + shared. extern "C" __global__ void __launch_bounds__(32) tf_glm_topk(const float* logits, const float* bias, int* pick, float* wts, float scale, int ne, int top_k, - int norm) { + int norm, int n_group, int topk_group) { constexpr int MAXE = 16; // 512 experts a row const int64_t r = blockIdx.x; const int lane = threadIdx.x; @@ -79,6 +121,7 @@ extern "C" __global__ void __launch_bounds__(32) score[i] = e < ne ? div_full(1.0f, exp_approx(0.0f - logits[r * ne + e]) + 1.0f) : 0.0f; choice[i] = e < ne ? score[i] + bias[e] : -INFINITY; } + if (n_group > 1) glm_moe::keep_groups(choice, lane, ne, ne / n_group, n_group, topk_group); float total = 0.0f, mine = 0.0f; int mine_id = 0; for (int k = 0; k < top_k; ++k) { diff --git a/zig/src/families/glm/cuda_glue.zig b/zig/src/families/glm/cuda_glue.zig index 733edcdb27..abb53038d0 100644 --- a/zig/src/families/glm/cuda_glue.zig +++ b/zig/src/families/glm/cuda_glue.zig @@ -108,6 +108,11 @@ fn cdiv(a: usize, b: usize) u32 { return @intCast((a + b - 1) / b); } +/// The kernel's clamp for a limit: 0 (unclamped, as the experts' 0) as +inf, whose min and clip are exact. +pub fn clampOf(limit: f32) f32 { + return if (limit == 0) std.math.inf(f32) else limit; +} + /// swiglu's column block: 512 where it divides the width, else 64 (TP=3's 704 and 640, the latter 128 in glue.py). fn swigluBlock(w: usize) !usize { if (w % 512 == 0) return 512; @@ -115,12 +120,12 @@ fn swigluBlock(w: usize) !usize { return error.NoSwigluBlock; } -/// glue.swiglu: [gate | up] -> bf16(bf16(silu(min(g, L))) * clip(u, -L, L)), 64-sums in `xs`. +/// glue.swiglu: [gate | up] -> bf16(bf16(silu(min(g, L))) * clip(u, -L, L)), 64-sums in `xs`; L per layer, 0 for none. pub fn swiglu(set: Set, stream: cuda.Stream, gu: u64, out: u64, xs: u64, limit: f32, rows: usize, w: usize) !void { const block = try swigluBlock(w); const f, const threads: u32 = if (block == 512) .{ set.glue.swiglu_512, 128 } else .{ set.glue.swiglu_64, 64 }; var a: cuda.Args = .{}; - inline for (.{ gu, out, xs, limit, i64of(w) }) |v| a.add(v); + inline for (.{ gu, out, xs, clampOf(limit), i64of(w) }) |v| a.add(v); try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(w / block) }, .block = .{ .x = threads } }, stream, &a); } @@ -148,11 +153,15 @@ pub fn router(set: Set, stream: cuda.Stream, x: u64, x_stride: usize, w: u64, ou try sliceSum(set, stream, part, out, m * ne, router_ks, true); } -/// glue.select: top_k by sigmoid + bias (lower id on ties), normed if `norm`, times `scale`, then the shared expert. -pub fn select(set: Set, stream: cuda.Stream, logits: u64, bias: u64, pick: u64, wts: u64, rows: usize, top_k: usize, experts: usize, scale: f32, norm: bool) !void { +/// Group-limited routing (Ling's n_group 8, topk_group 4): GLM's one group keeps every expert. +pub const Groups = struct { n: usize = 1, keep: usize = 1 }; + +/// glue.select: top_k by sigmoid + bias (lower id on ties) in the kept groups, normed if `norm`, x `scale`, + shared. +pub fn select(set: Set, stream: cuda.Stream, logits: u64, bias: u64, pick: u64, wts: u64, rows: usize, top_k: usize, experts: usize, scale: f32, norm: bool, groups: Groups) !void { if (experts > 512 or top_k >= 32) return error.TopkTooWide; + if (groups.n == 0 or groups.n > 32 or experts % groups.n != 0 or groups.keep == 0 or groups.keep > groups.n or top_k > groups.keep * (experts / groups.n)) return error.TopkGroups; var a: cuda.Args = .{}; - inline for (.{ logits, bias, pick, wts, scale, cint(experts), cint(top_k), @as(c_int, @intFromBool(norm)) }) |v| a.add(v); + inline for (.{ logits, bias, pick, wts, scale, cint(experts), cint(top_k), @as(c_int, @intFromBool(norm)), cint(groups.n), cint(groups.keep) }) |v| a.add(v); try cuda.launch.launch(set.glue.topk, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 32 } }, stream, &a); } diff --git a/zig/src/families/glm/cuda_moe.zig b/zig/src/families/glm/cuda_moe.zig index ec860a188e..d608597722 100644 --- a/zig/src/families/glm/cuda_moe.zig +++ b/zig/src/families/glm/cuda_moe.zig @@ -15,9 +15,11 @@ pub const Layer = struct { shared_gu: u64, shared_down: u64, shared: usize, + shared_limit: f32, // the shared expert's SwiGLU clamp, 0 for none (Ling's differs from the routed one's) top_k: usize, scale: f32, norm: bool, + groups: glue.Groups = .{}, }; /// The block's scratch for up to `rows` rows (forward.Buffers' MoE part). @@ -51,7 +53,7 @@ pub fn body(k: Kernels, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_str const slots = l.top_k + 1; const shared_id: i32 = @intCast(l.count); try glue.router(k.set, stream, x, x_stride, l.router, s.mlog, s.router_part, rows, d, l.count); - try glue.select(k.set, stream, s.mlog, l.bias, s.pick, s.wts, rows, l.top_k, l.count, l.scale, l.norm); + try glue.select(k.set, stream, s.mlog, l.bias, s.pick, s.wts, rows, l.top_k, l.count, l.scale, l.norm, l.groups); const plan: experts.Plan = .{ .route = s.plan, .slots = slots, .experts = l.count + 1, .prompt = prompt }; try experts.route(k.experts, stream, s.pick, rows * slots, plan); try experts.run(k.experts, stream, .gate_up, l.routed, plan, x, x_stride, rows, s.eact, shared_id); @@ -65,7 +67,7 @@ pub fn fill(k: Kernels, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_str const n = hi - lo; const eb: usize = if (prompt) 2 else 4; // the slot rows: bf16 in a prompt chunk try b16.matmul(k.set, stream, x + lo * x_stride * 2, x_stride, l.shared_gu, s.sgu + lo * 2 * l.shared * 2, s.mm_part, n, 2 * l.shared, d, false); - try glue.swiglu(k.set, stream, s.sgu + lo * 2 * l.shared * 2, s.sact + lo * l.shared * 2, s.sxs + lo * (l.shared / 64) * 4, l.routed.limit, n, l.shared); + try glue.swiglu(k.set, stream, s.sgu + lo * 2 * l.shared * 2, s.sact + lo * l.shared * 2, s.sxs + lo * (l.shared / 64) * 4, l.shared_limit, n, l.shared); try b16.matmul(k.set, stream, s.sact + lo * l.shared * 2, l.shared, l.shared_down, s.sy + lo * d * 4, s.mm_part, n, d, l.shared, true); var a: cuda.Args = .{}; a.add(s.sy + lo * d * 4); diff --git a/zig/src/families/glm/cuda_weights.zig b/zig/src/families/glm/cuda_weights.zig index 80bc44a8e0..58f89293cb 100644 --- a/zig/src/families/glm/cuda_weights.zig +++ b/zig/src/families/glm/cuda_weights.zig @@ -396,6 +396,7 @@ const Loader = struct { .shared_gu = gu, .shared_down = sd, .shared = sh.len(), + .shared_limit = c.swiglu_limit, .top_k = c.topk, .scale = c.routed_scale, .norm = c.norm_topk, diff --git a/zig/tests/cuda/glm_ref_route.zig b/zig/tests/cuda/glm_ref_route.zig new file mode 100644 index 0000000000..ede4237915 --- /dev/null +++ b/zig/tests/cuda/glm_ref_route.zig @@ -0,0 +1,356 @@ +//! glue.select's group-limited routing (Ling's 512 / 8 groups / keep 4 / top 8) vs float64; GLM's bytes vs 6380c8da. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const qr = @import("quant_ref_tests.zig"); +const glm = @import("glm_cuda"); +const Gpu = check.Gpu; +const Rig = qr.RigT; +const bf = qr.helpers.bfOf; + +const glue = glm.glue; + +/// One routing problem: fp32 logits [rows][ne], fp32 bias, top_k of the `keep` best of `ng` groups, times `scale`. +const Route = struct { rows: usize, ne: usize, top_k: usize, ng: usize, keep: usize, scale: f32 }; + +/// The kernels' near ties: two choices (or group sums) this close in float64 may rank either way after fp32 rounding. +const tie_expert = std.math.pow(f64, 2, -20); +const tie_group = std.math.pow(f64, 2, -19); + +fn sigmoid(x: f32) f64 { + return 1 / (1 + @exp(-@as(f64, x))); +} + +/// float64 rule: `keep` best top-2 group sums (or `forced`), then top_k in them, lower ids on ties; returns kept bits. +fn refPick(choice: []const f64, q: Route, forced: ?u32, order: []usize, gs_out: []f64) u32 { + const gs = q.ne / q.ng; + var kept: u32 = if (q.ng == 1) 1 else 0; + if (q.ng > 1) { + for (0..q.ng) |g| { + var a: f64 = -std.math.inf(f64); + var b: f64 = -std.math.inf(f64); + for (choice[g * gs ..][0..gs]) |c| { + if (c > a) { + b = a; + a = c; + } else if (c > b) b = c; + } + gs_out[g] = a + b; + } + if (forced) |f| kept = f else for (0..q.keep) |_| { + var best: usize = q.ng; + for (0..q.ng) |g| if (kept >> @intCast(g) & 1 == 0 and (best == q.ng or gs_out[g] > gs_out[best])) { + best = g; + }; + kept |= @as(u32, 1) << @intCast(best); + } + } + var taken: [512]bool = @splat(false); + for (order[0..q.top_k]) |*o| { + var best: usize = q.ne; + for (0..q.ne) |e| { + if (taken[e] or kept >> @intCast(e / gs) & 1 == 0) continue; + if (best == q.ne or choice[e] > choice[best]) best = e; + } + o.* = best; + taken[best] = true; + } + return kept; +} + +/// The k-th and (k+1)-th best group sums' gap: below tie_group the kernel's kept set may legitimately differ. +fn boundaryGap(gsc: []const f64, keep: usize) f64 { + var s: [32]f64 = undefined; + @memcpy(s[0..gsc.len], gsc); + std.mem.sort(f64, s[0..gsc.len], {}, std.sort.desc(f64)); + return s[keep - 1] - s[keep]; +} + +const Tally = struct { rows: usize = 0, group_ties: usize = 0, expert_ties: usize = 0, widest: f64 = 0, worst: f64 = 0 }; + +/// Kernel picks and weights for every row against the float64 rule; near ties counted, anything else fails. +fn against(r: *Rig, q: Route, logits: []const f32, bias: []const f32, got_pick: []const i32, got_wts: []const f32) !Tally { + var t: Tally = .{ .rows = q.rows }; + const choice = try r.a.alloc(f64, q.ne); + const score = try r.a.alloc(f64, q.ne); + const order = try r.a.alloc(usize, q.top_k); + var gsc: [32]f64 = undefined; + const slots = q.top_k + 1; + for (0..q.rows) |row| { + for (score, choice, logits[row * q.ne ..][0..q.ne], bias) |*s, *c, l, b| { + s.* = sigmoid(l); + c.* = s.* + b; + } + const kept = refPick(choice, q, null, order, gsc[0..q.ng]); + const gs = q.ne / q.ng; + var mine: u32 = 0; + for (got_pick[row * slots ..][0..q.top_k]) |id| mine |= @as(u32, 1) << @intCast(@as(usize, @intCast(id)) / gs); + if (mine & ~kept != 0) { // a pick outside the float64 groups: only a near tie at the group boundary may do that + const gap = boundaryGap(gsc[0..q.ng], q.keep); + try check.expect(gap <= tie_group, "row {d}: groups {b} picked where float64 keeps {b} (boundary gap {e})", .{ row, mine, kept, gap }); + t.group_ties += 1; + t.widest = @max(t.widest, gap); + var f = mine; + var g: usize = 0; + while (@popCount(f) < q.keep) : (g += 1) if (kept >> @intCast(g) & 1 != 0) { + f |= @as(u32, 1) << @intCast(g); + }; + _ = refPick(choice, q, f, order, gsc[0..q.ng]); + } + var total: f64 = 0; + for (got_pick[row * slots ..][0..q.top_k]) |id| total += score[@intCast(id)]; + for (0..slots) |i| { + const id: usize = @intCast(got_pick[row * slots + i]); + const wt: f64 = got_wts[row * slots + i]; + if (i == q.top_k) { + try check.expect(id == q.ne and wt == 1, "row {d}: the shared expert last at weight 1", .{row}); + continue; + } + if (id != order[i]) { + const gap = @abs(choice[id] - choice[order[i]]); + try check.expect(gap <= tie_expert, "row {d} pick {d}: {d} where float64 ranks {d} (gap {e})", .{ row, i, id, order[i], gap }); + t.expert_ties += 1; + t.widest = @max(t.widest, gap); + } + const want = score[id] / total * q.scale; + t.worst = @max(t.worst, @abs(wt - want) / (@abs(want) * std.math.pow(f64, 2, -16))); + } + } + return t; +} + +fn f32s(r: *Rig, n: usize, lo: f32, hi: f32) ![]f32 { + const out = try r.a.alloc(f32, n); + for (out) |*v| v.* = lo + r.rng.float(f32) * (hi - lo); + return out; +} + +const Out = struct { pick: []const i32, wts: []const f32 }; + +fn runSelect(r: *Rig, set: glue.Set, q: Route, logits: []const f32, bias: []const f32, norm: bool) !Out { + const n = q.rows * (q.top_k + 1); + const pick = try r.zeros(n * 4); + const wts = try r.zeros(n * 4); + try glue.select(set, r.s, try r.dev(std.mem.sliceAsBytes(logits)), try r.dev(std.mem.sliceAsBytes(bias)), pick, wts, q.rows, q.top_k, q.ne, q.scale, norm, .{ .n = q.ng, .keep = q.keep }); + return .{ .pick = std.mem.bytesAsSlice(i32, @as([]align(4) u8, @alignCast(try r.back(pick, n * 4)))), .wts = std.mem.bytesAsSlice(f32, @as([]align(4) u8, @alignCast(try r.back(wts, n * 4)))) }; +} + +/// Rows built for exact ties: every group a copy of group 0 (all group sums tie), and a group whose top two are equal. +fn tieRows(r: *Rig, q: Route, logits: []f32, bias: []f32) void { + const gs = q.ne / q.ng; + @memset(bias, 0); + for (0..q.rows) |row| { + const l = logits[row * q.ne ..][0..q.ne]; + switch (row % 3) { + 0 => for (0..q.ne) |e| { + l[e] = l[e % gs]; + }, + 1 => { // equal pairs at 2.0 (one lane's in 6, two lanes' in 7) beat 2.2 + 0.0 only if the duplicate counts + for (l) |*v| v.* = -6 + r.rng.float(f32); + for (0..q.ng) |g| { + l[g * gs + 3] = if (g >= 6) 2.0 else 2.2; + l[g * gs + 35] = if (g == 6) 2.0 else 0.0; + if (g == 7) l[g * gs + 7] = 2.0; + } + }, + else => for (0..q.ne) |e| { + l[e] = if (e % 7 == 0) 1.0 else l[e]; + }, + } + } +} + +/// Synthetic Ling-shaped routing (and a few odd shapes) against float64, plus exact-tie rows that must match exactly. +pub fn groups(r: *Rig, set: glue.Set) !void { + const shapes = [_]Route{ + .{ .rows = 300, .ne = 512, .top_k = 8, .ng = 8, .keep = 4, .scale = 2.5 }, + .{ .rows = 300, .ne = 256, .top_k = 8, .ng = 8, .keep = 4, .scale = 2.5 }, // DeepSeek-V3's + .{ .rows = 100, .ne = 480, .top_k = 6, .ng = 12, .keep = 3, .scale = 1.0 }, // groups of 40 straddle the lanes + .{ .rows = 100, .ne = 512, .top_k = 8, .ng = 32, .keep = 1, .scale = 2.5 }, + }; + for (shapes) |q| for ([_]f32{ 0.001, 0.05 }) |bspread| { + const logits = try f32s(r, q.rows * q.ne, -6, 6); + const bias = try f32s(r, q.ne, -bspread, bspread); + const o = try runSelect(r, set, q, logits, bias, true); + const t = try against(r, q, logits, bias, o.pick, o.wts); + r.worst = @max(r.worst, t.worst); + std.debug.print("RESULT group select {d} rows, top {d} of {d} in {d} of {d} groups (bias +-{d}): weights at {e:.3} of 2^-16 relative, {d} group and {d} expert near ties (widest gap {e:.2})\n", .{ q.rows, q.top_k, q.ne, q.keep, q.ng, bspread, t.worst, t.group_ties, t.expert_ties, t.widest }); + try check.expect(t.worst <= 1, "group select weights within 2^-16 of float64", .{}); + }; + const q: Route = .{ .rows = 30, .ne = 512, .top_k = 8, .ng = 8, .keep = 4, .scale = 2.5 }; + const logits = try f32s(r, q.rows * q.ne, -6, 6); + const bias = try r.a.alloc(f32, q.ne); + tieRows(r, q, logits, bias); + const o = try runSelect(r, set, q, logits, bias, true); + const t = try against(r, q, logits, bias, o.pick, o.wts); + try check.expect(t.group_ties == 0 and t.expert_ties == 0, "exact-tie rows: {d} group and {d} expert swaps, want none", .{ t.group_ties, t.expert_ties }); + check.pass("group select: picks follow the float64 rule (exact ties to the lower id, near ties counted), weights within 2^-16", .{}); +} + +/// swiglu at limit 0 (Ling's unclamped layers) against float64 unclamped, and limit 10 differing where it clamps. +pub fn swigluLimit(r: *Rig, set: glue.Set) !void { + for ([_]usize{ 1024, 704 }) |w| { + const rows = 40; + const gu = try r.bfs(rows * 2 * w, 16.0); + const dgu = try r.dev(std.mem.sliceAsBytes(gu)); + const out0 = try r.zeros(rows * w * 2); + const out10 = try r.zeros(rows * w * 2); + const xs = try r.zeros(rows * w / 64 * 4); + try glue.swiglu(set, r.s, dgu, out0, xs, 0, rows, w); + try glue.swiglu(set, r.s, dgu, out10, xs, 10, rows, w); + const raw0 = try r.back(out0, rows * w * 2); + const raw10 = try r.back(out10, rows * w * 2); + var most: f64 = 0; + var clamped: usize = 0; + for (0..rows) |m| for (0..w) |d| { + const g = bf(gu[m * 2 * w + d]); + const u = bf(gu[m * 2 * w + w + d]); + const want = g / (1 + @exp(-g)) * u; + const got = bf(std.mem.bytesToValue(u16, raw0[2 * (m * w + d) ..][0..2])); + most = @max(most, @abs(got - want) / (@abs(want) * std.math.pow(f64, 2, -7) + 1e-30)); + if (!std.mem.eql(u8, raw0[2 * (m * w + d) ..][0..2], raw10[2 * (m * w + d) ..][0..2])) clamped += 1; + }; + r.worst = @max(r.worst, most); + std.debug.print("RESULT swiglu width {d}, limit 0: unclamped float64 at {e:.3} of 2^-7 relative; limit 10 changes {d} of {d}\n", .{ w, most, clamped, rows * w }); + try check.expect(most <= 1 and clamped > 0, "swiglu limit 0 is unclamped, limit 10 clamps", .{}); + } + check.pass("swiglu: the limit is per call, 0 runs unclamped (as the experts' 0), 10 clamps", .{}); +} + +const Old = struct { topk: cuda.Function, swiglu_512: cuda.Function, swiglu_64: cuda.Function }; + +/// GLM's calls through the new glue against 6380c8da's kernels (its glm_glue fatbin) on the same inputs: every byte. +fn glmBytes(r: *Rig, set: glue.Set, prev: Old) !void { + for ([_]usize{ 288, 256, 160 }) |ne| for ([_]bool{ true, false }) |norm| { + const q: Route = .{ .rows = 300, .ne = ne, .top_k = 8, .ng = 1, .keep = 1, .scale = 2.5 }; + const logits = try f32s(r, q.rows * ne, -8, 8); + for (logits[0..ne], 0..) |*v, e| v.* = if (e % 5 == 0) 1.0 else v.*; // exact ties in row 0 + const bias = try f32s(r, ne, -0.05, 0.05); + const n = q.rows * 9; + const dl = try r.dev(std.mem.sliceAsBytes(logits)); + const db = try r.dev(std.mem.sliceAsBytes(bias)); + const pick = try r.zeros(n * 4); + const wts = try r.zeros(n * 4); + var a: cuda.Args = .{}; + inline for (.{ dl, db, pick, wts, q.scale, @as(c_int, @intCast(ne)), @as(c_int, 8), @as(c_int, @intFromBool(norm)) }) |v| a.add(v); + try cuda.launch.launch(prev.topk, .{ .grid = .{ .x = @intCast(q.rows) }, .block = .{ .x = 32 } }, r.s, &a); + const pick_new = try r.zeros(n * 4); + const wts_new = try r.zeros(n * 4); + try glue.select(set, r.s, dl, db, pick_new, wts_new, q.rows, 8, ne, q.scale, norm, .{}); + try check.sameBytes("GLM select picks vs 6380c8da", try r.back(pick_new, n * 4), try r.back(pick, n * 4)); + try check.sameBytes("GLM select weights vs 6380c8da", try r.back(wts_new, n * 4), try r.back(wts, n * 4)); + }; + for ([_]usize{ 2048, 704 }) |w| { + const rows = 40; + const dgu = try r.dev(std.mem.sliceAsBytes(try r.bfs(rows * 2 * w, 16.0))); + const outs = [2]u64{ try r.zeros(rows * w * 2), try r.zeros(rows * w * 2) }; + const xss = [2]u64{ try r.zeros(rows * w / 64 * 4), try r.zeros(rows * w / 64 * 4) }; + const block: usize = if (w % 512 == 0) 512 else 64; + var a: cuda.Args = .{}; + inline for (.{ dgu, outs[0], xss[0], @as(f32, 10.0), @as(i64, @intCast(w)) }) |v| a.add(v); + try cuda.launch.launch(if (block == 512) prev.swiglu_512 else prev.swiglu_64, .{ .grid = .{ .x = rows, .y = @intCast(w / block) }, .block = .{ .x = if (block == 512) 128 else 64 } }, r.s, &a); + try glue.swiglu(set, r.s, dgu, outs[1], xss[1], 10, rows, w); + try check.sameBytes("GLM swiglu vs 6380c8da", try r.back(outs[1], rows * w * 2), try r.back(outs[0], rows * w * 2)); + try check.sameBytes("GLM swiglu sums vs 6380c8da", try r.back(xss[1], rows * w / 64 * 4), try r.back(xss[0], rows * w / 64 * 4)); + } + check.pass("GLM select (288, 256, 160 experts, normed and not) and swiglu (limit 10): every byte equals 6380c8da's kernels", .{}); +} + +const Env = struct { + arena: std.heap.ArenaAllocator, + s: cuda.Stream, + mod: cuda.Module, + ks: glue.Kernels, + prng: std.Random.DefaultPrng, + r: Rig, + + fn deinit(e: *Env) void { + for (e.r.bufs.items) |*b| b.free(); + e.r.bufs.deinit(e.r.gpu.gpa); + e.mod.unload(); + e.s.deinit(); + e.arena.deinit(); + } +}; + +fn env(gpu: Gpu, e: *Env) !void { + e.arena = .init(gpu.gpa); + e.s = try cuda.Stream.init(gpu.d, false); + e.mod = try cuda.Module.load(gpu.d, cuda.kernels.glm_glue); + e.ks = try glue.Kernels.load(e.mod); + e.prng = .init(804); + e.r = .{ .gpu = gpu, .a = e.arena.allocator(), .s = e.s, .rng = e.prng.random() }; +} + +/// `glm-route-old `: GLM's select and swiglu against 6380c8da's glm_glue fatbin, byte for byte. +pub fn old(gpu: Gpu, path: []const u8) !void { + var e: Env = undefined; + try env(gpu, &e); + defer e.deinit(); + const bytes = try std.Io.Dir.cwd().readFileAllocOptions(gpu.io, path, e.r.a, .limited(1 << 26), .@"64", null); + var m = try cuda.Module.load(gpu.d, bytes); + defer m.unload(); + try glmBytes(&e.r, .{ .glue = &e.ks }, .{ .topk = try m.function("tf_glm_topk"), .swiglu_512 = try m.function("tf_glm_swiglu_512"), .swiglu_64 = try m.function("tf_glm_swiglu_64") }); +} + +/// `glm-route-ling `: Ling layer 4's rows (HF's fp32 logits) vs float64, HF's routing and exllamav3's. +pub fn ling(gpu: Gpu, path: []const u8) !void { + var e: Env = undefined; + try env(gpu, &e); + defer e.deinit(); + const r = &e.r; + const raw = try std.Io.Dir.cwd().readFileAllocOptions(gpu.io, path, r.a, .limited(1 << 26), .@"4", null); + const hdr = std.mem.bytesAsSlice(i32, raw[0..20]); + const q: Route = .{ .rows = @intCast(hdr[0]), .ne = @intCast(hdr[1]), .top_k = @intCast(hdr[2]), .ng = @intCast(hdr[3]), .keep = @intCast(hdr[4]), .scale = std.mem.bytesToValue(f32, raw[20..24]) }; + var at: usize = 24; + const R = q.rows; + const K = q.top_k; + const parts = .{ q.rows * q.ne, q.ne, R * K, R * K, R * K, R * K }; + var sl: [6][]align(4) u8 = undefined; + inline for (parts, 0..) |n, i| { + sl[i] = @alignCast(raw[at..][0 .. n * 4]); + at += n * 4; + } + try check.expect(at == raw.len, "fixture: {d} bytes, read {d}", .{ raw.len, at }); + const logits = std.mem.bytesAsSlice(f32, sl[0]); + const bias = std.mem.bytesAsSlice(f32, sl[1]); + const hf_pick = std.mem.bytesAsSlice(i32, sl[2]); + const hf_w = std.mem.bytesAsSlice(f32, sl[3]); + const ex_pick = std.mem.bytesAsSlice(i32, sl[4]); + const ex_w = std.mem.bytesAsSlice(f32, sl[5]); + const o = try runSelect(r, .{ .glue = &e.ks }, q, logits, bias, true); + const t = try against(r, q, logits, bias, o.pick, o.wts); + std.debug.print("RESULT Ling L4 vs float64: {d} rows, weights at {e:.3} of 2^-16 relative, {d} group and {d} expert near ties (widest gap {e:.2})\n", .{ R, t.worst, t.group_ties, t.expert_ties, t.widest }); + try check.expect(t.worst <= 1, "Ling weights within 2^-16 of float64", .{}); + var hf_same: usize = 0; + var ex_same: usize = 0; + var d2: f64 = 0; + var n2: f64 = 0; + var hf_most: f64 = 0; + var ex_most: f64 = 0; + for (0..R) |row| { + const mine = o.pick[row * (K + 1) ..][0..K]; + const mw = o.wts[row * (K + 1) ..][0..K]; + if (std.mem.eql(i32, mine, hf_pick[row * K ..][0..K])) { + hf_same += 1; + for (mw, hf_w[row * K ..][0..K]) |a, b| { + d2 += (@as(f64, a) - b) * (@as(f64, a) - b); + n2 += @as(f64, b) * b; + hf_most = @max(hf_most, @abs(@as(f64, a) - b) / @abs(@as(f64, b))); + } + } else std.debug.print(" row {d}: ours {any} HF {any}\n", .{ row, mine, hf_pick[row * K ..][0..K] }); + var all = true; + for (ex_pick[row * K ..][0..K], ex_w[row * K ..][0..K]) |id, w| { + const j = std.mem.indexOfScalar(i32, mine, id) orelse { + all = false; + continue; + }; + ex_most = @max(ex_most, @abs(@as(f64, mw[j]) - w) / @abs(@as(f64, mw[j]))); + } + if (all) ex_same += 1 else std.debug.print(" row {d}: ours {any} exllamav3 {any}\n", .{ row, mine, ex_pick[row * K ..][0..K] }); + } + std.debug.print("RESULT Ling L4 vs HF BailingMoeV3Gate: {d}/{d} rows pick the same experts in the same order; weights rel L2 {e:.3}, worst rel {e:.3}\n", .{ hf_same, R, @sqrt(d2 / n2), hf_most }); + std.debug.print("RESULT Ling L4 vs exllamav3 routing_ds3: {d}/{d} rows pick the same set; weights (half) worst rel {e:.3}\n", .{ ex_same, R, ex_most }); + try check.expect(hf_same == R and ex_same == R and hf_most <= std.math.pow(f64, 2, -16) and ex_most <= std.math.pow(f64, 2, -10), "Ling routing equals HF's and exllamav3's", .{}); + check.pass("Ling layer 4 routing: picks equal HF's and exllamav3's on all {d} rows, weights within float64 tolerance", .{R}); +} diff --git a/zig/tests/cuda/glm_ref_tests.zig b/zig/tests/cuda/glm_ref_tests.zig index b0e7598990..724cec58f4 100644 --- a/zig/tests/cuda/glm_ref_tests.zig +++ b/zig/tests/cuda/glm_ref_tests.zig @@ -134,7 +134,7 @@ fn routers(r: *Rig, set: glue.Set) !void { for (bias) |*b| b.* = (r.rng.float(f32) * 2 - 1) * 0.05; const pick = try r.zeros(m * (top_k + 1) * 4); const wts = try r.zeros(m * (top_k + 1) * 4); - try glue.select(set, r.s, full, try r.dev(std.mem.sliceAsBytes(bias)), pick, wts, m, top_k, ne, scale, true); + try glue.select(set, r.s, full, try r.dev(std.mem.sliceAsBytes(bias)), pick, wts, m, top_k, ne, scale, true, .{}); const logits = try r.back(full, m * ne * 4); const got_pick = try r.back(pick, m * (top_k + 1) * 4); const got_wts = try r.back(wts, m * (top_k + 1) * 4); @@ -196,6 +196,8 @@ pub fn run(gpu: Gpu) !void { const set: glue.Set = .{ .glue = &ks }; for ([_]Shape{ .{ .n = 512, .k = 4096 }, .{ .n = 1280, .k = 1536 }, .{ .n = 1024, .k = 128 } }) |sh| try matmuls(&r, set, sh); try routers(&r, set); + try @import("glm_ref_route.zig").groups(&r, set); + try @import("glm_ref_route.zig").swigluLimit(&r, set); try @import("glm_ref_attn.zig").dense(&r, set); try @import("glm_ref_attn.zig").sparse(&r, set); try @import("glm_ref_hc.zig").mixing(&r, set); diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index fe98682c86..b6fabfc13e 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -38,6 +38,7 @@ const lfm2_vision_tests = @import("lfm2_vision_tests.zig"); const radio_tests = @import("radio_tests.zig"); const glm_lanes_tests = @import("glm_lanes_tests.zig"); const glm_synth = @import("glm_synth.zig"); +const glm_ref_route = @import("glm_ref_route.zig"); const usage = \\usage: tf-cuda-test @@ -94,6 +95,8 @@ const usage = \\ glm-lanes-cover copy drafts, cancels, discards, eos and max_new on one lanes.Engine across requests \\ glm-lanes-concurrent concurrent streams on one lanes.Engine, each against its solo run \\ glm-shared-check shared verify windows against lone twins, logits and states byte for byte + \\ glm-route-old GLM's select and swiglu against an older glm_glue fatbin (6380c8da), byte for byte + \\ glm-route-ling Ling's group-limited routing on captured layer-4 logits vs float64, HF and exllamav3 \\ glm-synth a synthetic GLM-5.3-Flash checkpoint (2 layers, 16 experts) and its dc.json for the lane checks \\ ; @@ -191,6 +194,8 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "train-pack")) return train_pack.run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) try std.fmt.parseInt(usize, rest[2], 10) else 24); if (std.mem.eql(u8, cmd, "glm-lanes-cover")) return glm_lanes_tests.coverCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); if (std.mem.eql(u8, cmd, "glm-shared-check")) return glm_lanes_tests.sharedCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); + if (std.mem.eql(u8, cmd, "glm-route-old")) return glm_ref_route.old(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "glm-route-ling")) return glm_ref_route.ling(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "glm-synth")) return glm_synth.run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "glm-lanes-concurrent")) return glm_lanes_tests.concurrentCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); std.debug.print("{s}", .{usage}); From a6559571fb7cb6c31f8a92ac752f6dcf365ccba9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 15 Oct 2026 09:19:44 +0200 Subject: [PATCH 02/34] ling: Ling-3.1-flash's config loader and weight map, its routing and SwiGLU limits as GLM's kernels take them config.zig parses config.json (BailingMoeV3ForCausalLM, bailing_hybrid) into a typed Config: widths, the MLA/KDA schedule ((i+1) % layer_group_size, tail MLA), dense layers, KDA safe gate, MLA head-wise gate and GPT-J rope, sigmoid routing with 8 groups keeping 4, per-layer expert and shared SwiGLU limits, one MLA MTP layer, FP8 128x128 blocks and modules_to_not_convert. It refuses what the engine does not run, each with a one-sentence reason. names.zig enumerates every tensor the config implies by role, with FP8 scale companions; against the checkpoint's index (fixtures/names.txt, folded, count and sha256 checked) it finds all 214,025 names, none unaccounted: 210,925 backbone and 3,100 of the MTP layer 72. cuda.zig gives select()'s Groups and each layer's limits for moe.Layer. Host tests: zig build test-ling (also under test-cuda-host). --- zig/build/cuda.zig | 6 + zig/src/families/ling/config.zig | 361 +++++++++++++++++++++ zig/src/families/ling/config_test.zig | 156 +++++++++ zig/src/families/ling/cuda.zig | 68 ++++ zig/src/families/ling/fixtures/config.json | 254 +++++++++++++++ zig/src/families/ling/fixtures/names.txt | 61 ++++ zig/src/families/ling/names.zig | 220 +++++++++++++ zig/src/families/ling/names_test.zig | 150 +++++++++ 8 files changed, 1276 insertions(+) create mode 100644 zig/src/families/ling/config.zig create mode 100644 zig/src/families/ling/config_test.zig create mode 100644 zig/src/families/ling/cuda.zig create mode 100644 zig/src/families/ling/fixtures/config.json create mode 100644 zig/src/families/ling/fixtures/names.txt create mode 100644 zig/src/families/ling/names.zig create mode 100644 zig/src/families/ling/names_test.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 684d6c3cc9..4c241edf7b 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -272,6 +272,12 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st // the GLM-5.3-Flash CUDA family's host side const glm = glmCuda(b, host, .debug, cuda, mods.core, mods.lanes, mods.heat); step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = glm })).step); + // Ling-3.1-flash's config loader and weight map, and its routing and limits as the GLM kernels take them + const ling = b.createModule(.{ .root_source_file = b.path("zig/src/families/ling/cuda.zig"), .target = host, .optimize = .debug, .link_libc = true }); + ling.addImport("glm_cuda", glm); + const ling_run = b.addRunArtifact(b.addTest(.{ .root_module = ling })); + step.dependOn(&ling_run.step); + b.step("test-ling", "Ling-3.1-flash's config loader and weight map (no GPU work)").dependOn(&ling_run.step); // the shares a GLM-5.3-Flash rank takes (cluster/canon.zig against the Python family's) const shares = b.createModule(.{ .root_source_file = b.path("zig/src/cluster/canon_glm_test.zig"), .target = host, .optimize = .debug }); step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = shares })).step); diff --git a/zig/src/families/ling/config.zig b/zig/src/families/ling/config.zig new file mode 100644 index 0000000000..40f2f877d1 --- /dev/null +++ b/zig/src/families/ling/config.zig @@ -0,0 +1,361 @@ +//! Ling-3.1-flash's shape from config.json (bailing_hybrid), refusing what the engine does not run. +const std = @import("std"); + +/// Backbone plus MTP layers the per-layer tables hold (Ling-3.1-flash: 72 + 1). +pub const max_layers = 80; + +pub const Kind = enum { kda, mla }; + +/// MLA's rope pairs: rope_interleave true is exllamav3's RopeStyle.GPTJ (adjacent pairs), false NeoX (halves). +pub const RopeStyle = enum { gptj, neox }; + +/// The checkpoint's linear weights: bf16, or FP8 e4m3 in 128x128 blocks with a weight_scale_inv each. +pub const Quant = enum { bf16, fp8_block128 }; + +/// Why parse refused the config, one sentence for the caller's log line (tests read it too). +pub const Why = struct { + buf: [320]u8 = undefined, + len: usize = 0, + + pub fn set(self: *Why, comptime fmt: []const u8, args: anytype) void { + var w: std.Io.Writer = .fixed(&self.buf); + w.print(fmt, args) catch {}; + self.len = w.end; + } + + pub fn text(self: *const Why) []const u8 { + return self.buf[0..self.len]; + } +}; + +pub const Config = struct { + hidden: u32, + layers: u32, // backbone layers; the MTP layers follow at [layers, layers + mtp) + vocab: u32, + eps: f32, + rope_theta: f64, + rope_dim: u32, // qk_rope_head_dim: the modeling's rotary embedding takes it whole (partial factor 1) + rope_style: RopeStyle, + group: u32, // layer_group_size: MLA on each group's last layer and on the tail past the last full group + dense_layers: u32, // first_k_dense_replace + dense_inter: u32, + kda_heads: u32, + kda_dim: u32, + conv: u32, + lower_bound: f32, // kda_safe_gate's kda_lower_bound + mla_heads: u32, + q_lora: u32, // 0: no q_lora_rank, one q_proj + kv_lora: u32, + nope: u32, + v_dim: u32, + head_gate: bool, // gated_attention_proj_granularity_type head_wise: g_proj gives a sigmoid gate per head + experts: u32, + topk: u32, + n_group: u32, + topk_group: u32, + routed_scale: f32, + norm_topk: bool, + expert_bias: bool, // moe_router_enable_expert_bias: gate.expert_bias, added for the pick only + moe_inter: u32, + shared_inter: u32, // moe_shared_expert_intermediate_size x num_shared_experts, 0 without a shared expert + expert_limit: [max_layers]f32, // expert_swiglu_limit_list by layer, 0 for none + shared_limit: [max_layers]f32, // share_expert_swiglu_limit_list by layer, 0 for none + mtp: u32, + tie: bool, + eos: [4]u32, + eos_n: u32, + bos: ?u32, + quant: Quant, + plain_text: [1024]u8, // modules_to_not_convert, packed; plain_ends[i] ends entry i + plain_ends: [32]u16, + plain_n: u32, + + /// Layer i's mixer: the modeling's schedule for the backbone; the MTP layers are MLA (mtp_use_kda false). + pub fn kind(c: *const Config, i: u32) Kind { + if (i >= c.layers) return .mla; + const full = c.layers / c.group * c.group; + return if ((i + 1) % c.group == 0 or i >= full) .mla else .kda; + } + + /// Layer i has routed experts (the MTP layers always do); the first dense_layers run a dense MLP. + pub fn isMoe(c: *const Config, i: u32) bool { + return i >= c.dense_layers; + } + + pub fn countKind(c: *const Config, k: Kind) u32 { + var n: u32 = 0; + for (0..c.layers) |i| n += @intFromBool(c.kind(@intCast(i)) == k); + return n; + } + + /// MLA's query and key head width: nope then rope. + pub fn qkDim(c: *const Config) u32 { + return c.nope + c.rope_dim; + } + + /// MLA's softmax scale, qk_head_dim^-0.5 (no rope_scaling mscale: parse refuses rope_scaling). + pub fn smScale(c: *const Config) f32 { + return @floatCast(1.0 / @sqrt(@as(f64, @floatFromInt(c.qkDim())))); + } + + pub fn plain(c: *const Config, i: u32) []const u8 { + const lo: usize = if (i == 0) 0 else c.plain_ends[i - 1]; + return c.plain_text[lo..c.plain_ends[i]]; + } + + /// A linear module stored without FP8 (modules_to_not_convert names it, or ends it after a dot). + pub fn unconverted(c: *const Config, module: []const u8) bool { + for (0..c.plain_n) |i| { + const p = c.plain(@intCast(i)); + if (std.mem.eql(u8, module, p)) return true; + if (module.len > p.len and std.mem.endsWith(u8, module, p) and module[module.len - p.len - 1] == '.') return true; + } + return false; + } + + /// A linear module carries an FP8 weight_scale_inv beside its weight. + pub fn scaled(c: *const Config, module: []const u8) bool { + return c.quant == .fp8_block128 and !c.unconverted(module); + } +}; + +fn get(o: std.json.ObjectMap, key: []const u8, why: *Why) !std.json.Value { + return o.get(key) orelse { + why.set("config.json has no {s}", .{key}); + return error.BadConfig; + }; +} + +fn int(o: std.json.ObjectMap, key: []const u8, why: *Why) !u32 { + const v = try get(o, key, why); + if (v != .integer or v.integer < 0 or v.integer > std.math.maxInt(u32)) { + why.set("config.json's {s} is not a whole number", .{key}); + return error.BadConfig; + } + return @intCast(v.integer); +} + +fn number(v: std.json.Value) ?f64 { + return switch (v) { + .float => |f| f, + .integer => |i| @floatFromInt(i), + else => null, + }; +} + +fn float(o: std.json.ObjectMap, key: []const u8, why: *Why) !f64 { + return number(try get(o, key, why)) orelse { + why.set("config.json's {s} is not a number", .{key}); + return error.BadConfig; + }; +} + +fn flag(o: std.json.ObjectMap, key: []const u8, default: bool, why: *Why) !bool { + const v = o.get(key) orelse return default; + if (v != .bool) { + why.set("config.json's {s} is not true or false", .{key}); + return error.BadConfig; + } + return v.bool; +} + +fn text(o: std.json.ObjectMap, key: []const u8) ?[]const u8 { + const v = o.get(key) orelse return null; + return if (v == .string) v.string else null; +} + +/// A string key that, when present and not null, must read `want`. +fn must(o: std.json.ObjectMap, key: []const u8, want: []const u8, err: anyerror, why: *Why) !void { + const v = o.get(key) orelse return; + if (v == .null) return; + if (v == .string and std.mem.eql(u8, v.string, want)) return; + why.set("{s} must be \"{s}\": the engine runs no other", .{ key, want }); + return err; +} + +fn limits(o: std.json.ObjectMap, key: []const u8, layers: u32, out: *[max_layers]f32, why: *Why) !void { + out.* = @splat(0); + const v = o.get(key) orelse return; // absent: no clamp anywhere, as the reference reads it + if (v != .array or v.array.items.len != layers) { + why.set("{s} has {d} entries for {d} layers", .{ key, if (v == .array) v.array.items.len else 0, layers }); + return error.LimitList; + } + for (v.array.items, 0..) |e, i| { + const x = if (e == .null) 0 else number(e) orelse -1; + if (x < 0 or !std.math.isFinite(x)) { + why.set("{s}[{d}] is not a limit (0 for none, else positive)", .{ key, i }); + return error.LimitList; + } + out[i] = @floatCast(x); + } +} + +fn quantization(o: std.json.ObjectMap, c: *Config, why: *Why) !void { + c.quant = .bf16; + c.plain_n = 0; + const qv = o.get("quantization_config") orelse return; + if (qv == .null) return; + if (qv != .object) { + why.set("quantization_config is not an object", .{}); + return error.Quant; + } + const q = qv.object; + const method = text(q, "quant_method") orelse ""; + const block = q.get("weight_block_size"); + const b128 = block != null and block.? == .array and block.?.array.items.len == 2 and + block.?.array.items[0] == .integer and block.?.array.items[0].integer == 128 and + block.?.array.items[1] == .integer and block.?.array.items[1].integer == 128; + if (!std.mem.eql(u8, method, "fp8") or !b128 or !std.mem.eql(u8, text(q, "fmt") orelse "e4m3", "e4m3") or + !std.mem.eql(u8, text(q, "activation_scheme") orelse "dynamic", "dynamic")) + { + why.set("quantization {s}: the loader reads bf16 or FP8 e4m3 in 128x128 blocks", .{if (method.len > 0) method else "without quant_method"}); + return error.Quant; + } + c.quant = .fp8_block128; + var at: usize = 0; + if (q.get("modules_to_not_convert")) |list| if (list == .array) for (list.array.items) |m| { + if (m != .string or c.plain_n == c.plain_ends.len or at + m.string.len > c.plain_text.len) { + why.set("modules_to_not_convert holds more than the loader keeps", .{}); + return error.Quant; + } + @memcpy(c.plain_text[at..][0..m.string.len], m.string); + at += m.string.len; + c.plain_ends[c.plain_n] = @intCast(at); + c.plain_n += 1; + }; +} + +fn tokens(o: std.json.ObjectMap, c: *Config, why: *Why) !void { + c.eos_n = 0; + const v = try get(o, "eos_token_id", why); + const items: []const std.json.Value = if (v == .array) v.array.items else &.{v}; + if (items.len == 0 or items.len > c.eos.len) { + why.set("eos_token_id holds {d} ids, the engine stops on 1 to {d}", .{ items.len, c.eos.len }); + return error.BadConfig; + } + for (items) |e| { + if (e != .integer or e.integer < 0 or e.integer >= c.vocab) { + why.set("eos_token_id is not a token of the {d}-token vocabulary", .{c.vocab}); + return error.BadConfig; + } + c.eos[c.eos_n] = @intCast(e.integer); + c.eos_n += 1; + } + c.bos = null; + if (o.get("bos_token_id")) |b| if (b == .integer and b.integer >= 0 and b.integer < c.vocab) { + c.bos = @intCast(b.integer); + }; +} + +/// config.json as the reference (configuration_bailing_moe_v3.py, exllamav3's bailing_hybrid) reads it. +pub fn parse(gpa: std.mem.Allocator, json: []const u8, why: *Why) !Config { + const parsed = std.json.parseFromSlice(std.json.Value, gpa, json, .{}) catch { + why.set("config.json is not JSON", .{}); + return error.BadConfig; + }; + defer parsed.deinit(); + if (parsed.value != .object) return error.BadConfig; + const o = parsed.value.object; + const arch = o.get("architectures"); + const ling = arch != null and arch.? == .array and arch.?.array.items.len == 1 and arch.?.array.items[0] == .string and + std.mem.eql(u8, arch.?.array.items[0].string, "BailingMoeV3ForCausalLM"); + if (!ling or !std.mem.eql(u8, text(o, "model_type") orelse "", "bailing_hybrid")) { + why.set("not a Ling-3.1 checkpoint: architectures BailingMoeV3ForCausalLM, model_type bailing_hybrid", .{}); + return error.NotLing; + } + var c = std.mem.zeroes(Config); + c.hidden = try int(o, "hidden_size", why); + c.layers = try int(o, "num_hidden_layers", why); + c.vocab = try int(o, "vocab_size", why); + c.eps = @floatCast(try float(o, "rms_norm_eps", why)); + c.group = try int(o, "layer_group_size", why); + c.dense_layers = try int(o, "first_k_dense_replace", why); + c.dense_inter = try int(o, "intermediate_size", why); + c.mtp = if (o.get("num_nextn_predict_layers")) |_| try int(o, "num_nextn_predict_layers", why) else 0; + if (c.layers == 0 or c.group == 0 or c.dense_layers > c.layers or c.hidden == 0 or c.vocab == 0) { + why.set("{d} layers in groups of {d}, {d} dense: not a layer schedule", .{ c.layers, c.group, c.dense_layers }); + return error.BadConfig; + } + if (c.layers + c.mtp > max_layers) { + why.set("{d} layers and {d} MTP layers: the loader holds {d}", .{ c.layers, c.mtp, max_layers }); + return error.TooManyLayers; + } + if (c.mtp > 1 or try flag(o, "mtp_use_kda", false, why)) { + why.set("MTP: the engine runs one MLA MTP layer", .{}); + return error.Mtp; + } + try must(o, "hidden_act", "silu", error.Activation, why); + if (try flag(o, "use_qkv_bias", false, why) or try flag(o, "use_bias", false, why)) { + why.set("use_qkv_bias/use_bias: the engine's projections have no bias", .{}); + return error.Bias; + } + // KDA: one head geometry for q, k and v, full-rank f and g, the safe gate's lower bound + if (!try flag(o, "no_kda_lora", false, why) or try flag(o, "use_kda_lora", false, why)) { + why.set("no_kda_lora must be true: the engine runs full-rank f_proj and g_proj", .{}); + return error.KdaLora; + } + c.kda_heads = try int(o, "num_attention_heads", why); + c.kda_dim = try int(o, "head_dim", why); + c.conv = try int(o, "short_conv_kernel_size", why); + if (!try flag(o, "kda_safe_gate", false, why) or o.get("kda_lower_bound") == null) { + why.set("kda_safe_gate with a kda_lower_bound: the engine's KDA gate is the safe one", .{}); + return error.KdaGate; + } + c.lower_bound = @floatCast(try float(o, "kda_lower_bound", why)); + // MLA + c.mla_heads = c.kda_heads; + const ql = o.get("q_lora_rank") orelse std.json.Value.null; + c.q_lora = if (ql == .null) 0 else try int(o, "q_lora_rank", why); + c.kv_lora = try int(o, "kv_lora_rank", why); + c.nope = try int(o, "qk_nope_head_dim", why); + c.rope_dim = try int(o, "qk_rope_head_dim", why); + c.v_dim = try int(o, "v_head_dim", why); + c.rope_theta = try float(o, "rope_theta", why); + c.rope_style = if (try flag(o, "rope_interleave", true, why)) .gptj else .neox; + if (o.get("rope_scaling")) |r| if (r != .null) { + why.set("rope_scaling: the engine's MLA has plain rope only", .{}); + return error.RopeScaling; + }; + const gate = o.get("gated_attention_proj_granularity_type") orelse std.json.Value.null; + c.head_gate = gate == .string and std.mem.eql(u8, gate.string, "head_wise"); + if (gate != .null and !c.head_gate) { + why.set("gated_attention_proj_granularity_type {s}: the engine gates MLA per head or not at all", .{if (gate == .string) gate.string else "?"}); + return error.GateGranularity; + } + // MoE: DeepSeek-V3's sigmoid router with group-limited top-k and a bias for the pick + try must(o, "score_function", "sigmoid", error.ScoreFunction, why); + try must(o, "scoring_func", "sigmoid", error.ScoreFunction, why); + try must(o, "topk_method", "noaux_tc", error.Routing, why); + c.experts = try int(o, "num_experts", why); + c.topk = try int(o, "num_experts_per_tok", why); + c.n_group = if (o.get("n_group")) |_| try int(o, "n_group", why) else 1; + c.topk_group = if (o.get("topk_group")) |_| try int(o, "topk_group", why) else 1; + c.routed_scale = @floatCast(if (o.get("routed_scaling_factor")) |_| try float(o, "routed_scaling_factor", why) else 1.0); + c.norm_topk = try flag(o, "norm_topk_prob", true, why); + if (!c.norm_topk) { + why.set("norm_topk_prob must be true: the routed weights are renormalized over the top-k", .{}); + return error.Routing; + } + c.expert_bias = try flag(o, "moe_router_enable_expert_bias", true, why); + // glue.select's bounds: at most 512 experts, top-k under 32, up to 32 equal groups that hold the top-k + const g_ok = c.n_group >= 1 and c.n_group <= 32 and c.experts % c.n_group == 0 and c.topk_group >= 1 and + c.topk_group <= c.n_group and c.topk <= c.topk_group * (c.experts / c.n_group); + if (c.experts > 512 or c.topk == 0 or c.topk >= 32 or c.topk > c.experts or !g_ok) { + why.set("routing {d} of {d} experts in {d} groups keeping {d}: past glue.select's bounds", .{ c.topk, c.experts, c.n_group, c.topk_group }); + return error.Routing; + } + c.moe_inter = try int(o, "moe_intermediate_size", why); + const shared_n = if (o.get("num_shared_experts")) |v| (if (v == .null) 0 else try int(o, "num_shared_experts", why)) else 0; + const shared_w = if (o.get("moe_shared_expert_intermediate_size")) |_| try int(o, "moe_shared_expert_intermediate_size", why) else c.moe_inter; + c.shared_inter = shared_n * shared_w; + try limits(o, "expert_swiglu_limit_list", c.layers, &c.expert_limit, why); + try limits(o, "share_expert_swiglu_limit_list", c.layers, &c.shared_limit, why); + for (0..c.dense_layers) |i| if (c.expert_limit[i] != 0 or c.shared_limit[i] != 0) { + why.set("layer {d} is dense with a SwiGLU limit: neither list says which a dense MLP takes", .{i}); + return error.LimitList; + }; + c.tie = try flag(o, "tie_word_embeddings", false, why); + try tokens(o, &c, why); + try quantization(o, &c, why); + return c; +} diff --git a/zig/src/families/ling/config_test.zig b/zig/src/families/ling/config_test.zig new file mode 100644 index 0000000000..f68eaac905 --- /dev/null +++ b/zig/src/families/ling/config_test.zig @@ -0,0 +1,156 @@ +//! Ling-3.1-flash's own config.json (fixtures/config.json) parsed, and each refusal on a one-edit copy. +const std = @import("std"); +const config = @import("config.zig"); + +const real = @embedFile("fixtures/config.json"); +const t = std.testing; + +fn load() !config.Config { + var why: config.Why = .{}; + return config.parse(t.allocator, real, &why); +} + +/// The fixture with `from` (which must occur exactly once) replaced by `to`. +fn edit(from: []const u8, to: []const u8) ![]u8 { + if (std.mem.count(u8, real, from) != 1) return error.EditAnchor; + return std.mem.replaceOwned(u8, t.allocator, real, from, to); +} + +fn refused(from: []const u8, to: []const u8, want: anyerror, says: []const u8) !void { + const json = try edit(from, to); + defer t.allocator.free(json); + var why: config.Why = .{}; + try t.expectError(want, config.parse(t.allocator, json, &why)); + try t.expect(std.mem.indexOf(u8, why.text(), says) != null); +} + +test "the checkpoint's config.json parses to Ling-3.1-flash's shape" { + const c = try load(); + try t.expectEqual(@as(u32, 4096), c.hidden); + try t.expectEqual(@as(u32, 72), c.layers); + try t.expectEqual(@as(u32, 157184), c.vocab); + try t.expectEqual(@as(f32, 1e-6), c.eps); + try t.expectEqual(@as(f64, 6000000), c.rope_theta); + try t.expectEqual(@as(u32, 64), c.rope_dim); + try t.expectEqual(config.RopeStyle.gptj, c.rope_style); + try t.expectEqual(@as(u32, 8), c.group); + try t.expectEqual(@as(u32, 4), c.dense_layers); + try t.expectEqual(@as(u32, 10240), c.dense_inter); + try t.expectEqual(@as(u32, 64), c.kda_heads); + try t.expectEqual(@as(u32, 128), c.kda_dim); + try t.expectEqual(@as(u32, 4), c.conv); + try t.expectEqual(@as(f32, -5.0), c.lower_bound); + try t.expectEqual(@as(u32, 64), c.mla_heads); + try t.expectEqual(@as(u32, 1536), c.q_lora); + try t.expectEqual(@as(u32, 512), c.kv_lora); + try t.expectEqual(@as(u32, 128), c.nope); + try t.expectEqual(@as(u32, 128), c.v_dim); + try t.expectEqual(@as(u32, 192), c.qkDim()); + try t.expectEqual(@as(f32, @floatCast(1.0 / @sqrt(192.0))), c.smScale()); + try t.expect(c.head_gate); + try t.expectEqual(@as(u32, 512), c.experts); + try t.expectEqual(@as(u32, 8), c.topk); + try t.expectEqual(@as(u32, 8), c.n_group); + try t.expectEqual(@as(u32, 4), c.topk_group); + try t.expectEqual(@as(f32, 2.5), c.routed_scale); + try t.expect(c.norm_topk and c.expert_bias); + try t.expectEqual(@as(u32, 1280), c.moe_inter); + try t.expectEqual(@as(u32, 1280), c.shared_inter); + try t.expectEqual(@as(u32, 1), c.mtp); + try t.expect(!c.tie); + try t.expectEqual(@as(u32, 1), c.eos_n); + try t.expectEqual(@as(u32, 156895), c.eos[0]); + try t.expect(c.bos == null); // config.json names no bos_token_id; the tokenizer adds none either + try t.expectEqual(config.Quant.fp8_block128, c.quant); + try t.expectEqual(@as(u32, 14), c.plain_n); +} + +test "the schedule: MLA exactly on 7, 15, ..., 71 (9 layers), KDA on the other 63, dense 0-3, the MTP layer MLA + MoE" { + const c = try load(); + var mla: [9]u32 = undefined; + var n: usize = 0; + for (0..c.layers) |i| if (c.kind(@intCast(i)) == .mla) { + mla[n] = @intCast(i); + n += 1; + }; + try t.expectEqual(@as(usize, 9), n); + try t.expectEqualSlices(u32, &.{ 7, 15, 23, 31, 39, 47, 55, 63, 71 }, &mla); + try t.expectEqual(@as(u32, 9), c.countKind(.mla)); + try t.expectEqual(@as(u32, 63), c.countKind(.kda)); + for (0..c.layers) |i| try t.expectEqual(i >= 4, c.isMoe(@intCast(i))); + try t.expectEqual(config.Kind.mla, c.kind(72)); + try t.expect(c.isMoe(72)); +} + +test "a tail past the last full group is MLA, as the modeling's schedule says" { + var c = try load(); + c.layers = 10; + c.group = 4; + for ([_]config.Kind{ .kda, .kda, .kda, .mla, .kda, .kda, .kda, .mla, .mla, .mla }, 0..) |k, i| try t.expectEqual(k, c.kind(@intCast(i))); +} + +test "SwiGLU limits from the file: 10 on routed experts of layers 65-71 and shared ones of 64-71, else 0" { + const c = try load(); + for (0..c.layers) |i| { + try t.expectEqual(@as(f32, if (i >= 65) 10 else 0), c.expert_limit[i]); + try t.expectEqual(@as(f32, if (i >= 64) 10 else 0), c.shared_limit[i]); + } +} + +test "modules_to_not_convert matches whole dotted tails: mlp.gate is the router, not gate_proj" { + const c = try load(); + try t.expect(c.unconverted("model.layers.4.mlp.gate")); + try t.expect(!c.unconverted("model.layers.4.mlp.gate_proj")); + try t.expect(!c.unconverted("model.layers.4.mlp.shared_experts.gate_proj")); + try t.expect(c.unconverted("model.layers.0.attention.b_proj")); + try t.expect(!c.unconverted("model.layers.7.attention.q_b_proj")); + try t.expect(c.unconverted("model.layers.7.attention.g_proj")); + try t.expect(!c.unconverted("model.layers.72.attention.g_proj")); + try t.expect(c.unconverted("model.layers.72.eh_proj")); + try t.expect(c.unconverted("lm_head") and !c.scaled("lm_head")); + try t.expect(c.scaled("model.layers.0.attention.q_proj")); +} + +test "refusals: one edit each, with the reason" { + try refused("\"BailingMoeV3ForCausalLM\"", "\"BailingMoeV2ForCausalLM\"", error.NotLing, "not a Ling-3.1"); + try refused("\"score_function\": \"sigmoid\"", "\"score_function\": \"softmax\"", error.ScoreFunction, "score_function"); + try refused("\"scoring_func\": \"sigmoid\"", "\"scoring_func\": \"softmax\"", error.ScoreFunction, "scoring_func"); + try refused("\"no_kda_lora\": true", "\"no_kda_lora\": false", error.KdaLora, "no_kda_lora"); + try refused("\"use_kda_lora\": false", "\"use_kda_lora\": true", error.KdaLora, "no_kda_lora"); + try refused("\"kda_safe_gate\": true", "\"kda_safe_gate\": false", error.KdaGate, "kda_safe_gate"); + try refused("\"head_wise\"", "\"element_wise\"", error.GateGranularity, "element_wise"); + try refused("\"rope_scaling\": null", "\"rope_scaling\": {\"type\": \"yarn\", \"factor\": 4}", error.RopeScaling, "rope_scaling"); + try refused("\"n_group\": 8", "\"n_group\": 7", error.Routing, "groups"); + try refused("\"topk_group\": 4", "\"topk_group\": 9", error.Routing, "groups"); + try refused("\"num_experts\": 512", "\"num_experts\": 1024", error.Routing, "glue.select"); + try refused("\"norm_topk_prob\": true", "\"norm_topk_prob\": false", error.Routing, "norm_topk_prob"); + try refused("\"expert_swiglu_limit_list\": [\n 0,\n", "\"expert_swiglu_limit_list\": [\n", error.LimitList, "71 entries for 72"); + try refused("\"share_expert_swiglu_limit_list\": [\n 0,\n", "\"share_expert_swiglu_limit_list\": [\n 0,\n 0,\n", error.LimitList, "73 entries"); + try refused("\"share_expert_swiglu_limit_list\": [\n 0,\n", "\"share_expert_swiglu_limit_list\": [\n -1,\n", error.LimitList, "[0]"); + try refused("\"expert_swiglu_limit_list\": [\n 0,\n", "\"expert_swiglu_limit_list\": [\n 7,\n", error.LimitList, "layer 0 is dense"); + try refused("\"quant_method\": \"fp8\"", "\"quant_method\": \"gptq\"", error.Quant, "gptq"); + try refused("\"num_nextn_predict_layers\": 1", "\"num_nextn_predict_layers\": 2", error.Mtp, "MTP"); + try refused("\"mtp_use_kda\": false", "\"mtp_use_kda\": true", error.Mtp, "MTP"); + try refused("\"hidden_act\": \"silu\"", "\"hidden_act\": \"gelu\"", error.Activation, "hidden_act"); + try refused("\"use_qkv_bias\": false", "\"use_qkv_bias\": true", error.Bias, "bias"); + try refused("\"eos_token_id\": 156895", "\"eos_token_id\": 157184", error.BadConfig, "eos_token_id"); + try refused("\"num_hidden_layers\": 72", "\"num_hidden_layers\": 73", error.LimitList, "72 entries for 73"); + try refused("\"hidden_size\": 4096", "\"hidden_sz\": 4096", error.BadConfig, "hidden_size"); +} + +test "optional keys: no quantization reads bf16, a null q_lora_rank one q_proj, rope_interleave false NeoX" { + const a = try edit("\"q_lora_rank\": 1536", "\"q_lora_rank\": null"); + defer t.allocator.free(a); + const b = try std.mem.replaceOwned(u8, t.allocator, a, "\"rope_interleave\": true", "\"rope_interleave\": false"); + defer t.allocator.free(b); + const at = std.mem.indexOf(u8, b, "\"quantization_config\"").?; + const end = std.mem.indexOf(u8, b[at..], " },\n").? + at + 5; + const c_json = try std.mem.concat(t.allocator, u8, &.{ b[0..at], b[end..] }); + defer t.allocator.free(c_json); + var why: config.Why = .{}; + const c = try config.parse(t.allocator, c_json, &why); + try t.expectEqual(config.Quant.bf16, c.quant); + try t.expectEqual(@as(u32, 0), c.q_lora); + try t.expectEqual(config.RopeStyle.neox, c.rope_style); + try t.expect(!c.scaled("model.layers.0.attention.q_proj")); +} diff --git a/zig/src/families/ling/cuda.zig b/zig/src/families/ling/cuda.zig new file mode 100644 index 0000000000..615e36b02b --- /dev/null +++ b/zig/src/families/ling/cuda.zig @@ -0,0 +1,68 @@ +//! Ling-3.1-flash on CUDA: its config as the GLM-5.3-Flash kernels it reuses take it (routing groups, SwiGLU limits). +const std = @import("std"); +const glm = @import("glm_cuda"); +pub const config = @import("config.zig"); +pub const names = @import("names.zig"); + +const Config = config.Config; + +/// glue.select's group-limited routing: n_group groups, the best topk_group kept (parse holds select's bounds). +pub fn groups(c: *const Config) glm.glue.Groups { + return .{ .n = c.n_group, .keep = c.topk_group }; +} + +/// Layer i's SwiGLU clamps, 0 for none: experts.Layer.limit, moe.Layer.shared_limit, mlp.Layer.limit. +pub const Limits = struct { routed: f32, shared: f32, dense: f32 }; + +/// The lists cover the backbone; the MTP layer and the dense layers run unclamped (parse refuses a dense limit). +pub fn limits(c: *const Config, i: u32) Limits { + if (i >= c.layers or !c.isMoe(i)) return .{ .routed = 0, .shared = 0, .dense = 0 }; + return .{ .routed = c.expert_limit[i], .shared = c.shared_limit[i], .dense = 0 }; +} + +/// The MoE layer's routing fields as moe.Layer holds them (pointers and widths come from the weight loader). +pub fn route(c: *const Config, i: u32, l: *glm.moe.Layer) void { + const lim = limits(c, i); + l.top_k = c.topk; + l.scale = c.routed_scale; + l.norm = c.norm_topk; + l.groups = groups(c); + l.shared_limit = lim.shared; + l.routed.limit = lim.routed; +} + +test { + _ = @import("config_test.zig"); + _ = @import("names_test.zig"); +} + +test "Ling's routing as glue.select takes it: 8 groups keeping 4, and select's own bounds hold" { + var why: config.Why = .{}; + const c = try config.parse(std.testing.allocator, @embedFile("fixtures/config.json"), &why); + const g = groups(&c); + try std.testing.expectEqual(glm.glue.Groups{ .n = 8, .keep = 4 }, g); + try std.testing.expect(c.experts % g.n == 0 and g.keep <= g.n and c.topk <= g.keep * (c.experts / g.n) and g.n <= 32); + try std.testing.expect(!std.meta.eql(g, glm.glue.Groups{})); // GLM's one group would skip Ling's group pass +} + +test "per-layer limits on the MoE layers: routed 10 from 65, shared 10 from 64, none on dense or MTP layers" { + var why: config.Why = .{}; + const c = try config.parse(std.testing.allocator, @embedFile("fixtures/config.json"), &why); + for (0..c.layers + c.mtp) |i| { + const l = limits(&c, @intCast(i)); + const moe = i >= 4 and i < 72; + try std.testing.expectEqual(@as(f32, if (moe and i >= 65) 10 else 0), l.routed); + try std.testing.expectEqual(@as(f32, if (moe and i >= 64) 10 else 0), l.shared); + try std.testing.expectEqual(@as(f32, 0), l.dense); + } + var l: glm.moe.Layer = undefined; + route(&c, 64, &l); + try std.testing.expectEqual(@as(f32, 0), l.routed.limit); + try std.testing.expectEqual(@as(f32, 10), l.shared_limit); + route(&c, 65, &l); + try std.testing.expectEqual(@as(f32, 10), l.routed.limit); + try std.testing.expectEqual(glm.glue.Groups{ .n = 8, .keep = 4 }, l.groups); + try std.testing.expectEqual(@as(usize, 8), l.top_k); + try std.testing.expectEqual(@as(f32, 2.5), l.scale); + try std.testing.expect(l.norm); +} diff --git a/zig/src/families/ling/fixtures/config.json b/zig/src/families/ling/fixtures/config.json new file mode 100644 index 0000000000..bbdaf010c8 --- /dev/null +++ b/zig/src/families/ling/fixtures/config.json @@ -0,0 +1,254 @@ +{ + "architectures": [ + "BailingMoeV3ForCausalLM" + ], + "attention_dropout": 0.0, + "auto_map": { + "AutoConfig": "configuration_bailing_moe_v3.BailingMoeV3Config", + "AutoModel": "modeling_bailing_moe_v3.BailingMoeV3Model", + "AutoModelForCausalLM": "modeling_bailing_moe_v3.BailingMoeV3ForCausalLM" + }, + "embedding_dropout": 0.0, + "eos_token_id": 156895, + "expert_swiglu_limit_list": [ + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 10, + 10, + 10, + 10, + 10, + 10, + 10 + ], + "first_k_dense_replace": 4, + "gated_attention_proj_granularity_type": "head_wise", + "group_norm_size": 1, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 10240, + "kda_lower_bound": -5.0, + "kda_safe_gate": true, + "kv_lora_rank": 512, + "layer_group_size": 8, + "linear_silu": true, + "max_position_embeddings": 262144, + "max_window_layers": 20, + "moe_intermediate_size": 1280, + "moe_router_enable_expert_bias": true, + "moe_shared_expert_intermediate_size": 1280, + "mtp_loss_scaling_factor": 0, + "mtp_use_kda": false, + "n_group": 8, + "no_kda_lora": true, + "norm_topk_prob": true, + "num_attention_heads": 64, + "num_experts": 512, + "num_experts_per_tok": 8, + "num_hidden_layers": 72, + "num_key_value_heads": 64, + "num_kv_heads_for_linear_attn": 0, + "num_nextn_predict_layers": 1, + "num_shared_experts": 1, + "output_dropout": 0.0, + "output_router_logits": false, + "pad_token_id": 156892, + "partial_rotary_factor": 0.5, + "q_lora_rank": 1536, + "qk_head_dim": 192, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 64, + "quantization_config": { + "activation_scheme": "dynamic", + "fmt": "e4m3", + "quant_method": "fp8", + "modules_to_not_convert": [ + "lm_head", + "model.word_embeddings", + "b_proj", + "mlp.gate", + "model.layers.7.attention.g_proj", + "model.layers.15.attention.g_proj", + "model.layers.23.attention.g_proj", + "model.layers.31.attention.g_proj", + "model.layers.39.attention.g_proj", + "model.layers.47.attention.g_proj", + "model.layers.55.attention.g_proj", + "model.layers.63.attention.g_proj", + "model.layers.71.attention.g_proj", + "model.layers.72.eh_proj" + ], + "weight_block_size": [ + 128, + 128 + ] + }, + "rms_norm_eps": 1e-06, + "rope_interleave": true, + "rope_scaling": null, + "rope_theta": 6000000, + "rotary_dim": 64, + "routed_scaling_factor": 2.5, + "router_dtype": "fp32", + "scale_router_input": false, + "score_function": "sigmoid", + "scoring_func": "sigmoid", + "seq_aux": true, + "share_expert_swiglu_limit_list": [ + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 10, + 10, + 10, + 10, + 10, + 10, + 10, + 10 + ], + "short_conv_kernel_size": 4, + "tie_word_embeddings": false, + "topk_group": 4, + "topk_method": "noaux_tc", + "transformers_version": "4.45.0", + "up_proj_norm": false, + "use_bias": false, + "use_cache": true, + "use_kda_lora": false, + "use_mla_nope": false, + "use_nGPT": false, + "use_qk_norm": true, + "use_qkv_bias": false, + "v_head_dim": 128, + "value_norm": false, + "vocab_size": 157184, + "model_type": "bailing_hybrid", + "torch_dtype": "bfloat16" +} \ No newline at end of file diff --git a/zig/src/families/ling/fixtures/names.txt b/zig/src/families/ling/fixtures/names.txt new file mode 100644 index 0000000000..b2502f4427 --- /dev/null +++ b/zig/src/families/ling/fixtures/names.txt @@ -0,0 +1,61 @@ +ling31-names 214025 16f5c51dd86af196318448fc16a4177efbcbe83cbd34ab863409c57fc7f028a7 +lm_head.weight +model.layers.{L}.attention.A_log 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.b_proj.weight 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.dense.weight 7,15,23,31,39,47,55,63,71-72 +model.layers.{L}.attention.dense.weight_scale_inv 7,15,23,31,39,47,55,63,71-72 +model.layers.{L}.attention.dt_bias 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.f_proj.weight 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.f_proj.weight_scale_inv 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.g_proj.weight 0-72 +model.layers.{L}.attention.g_proj.weight_scale_inv 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70,72 +model.layers.{L}.attention.k_conv1d.weight 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.k_proj.weight 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.k_proj.weight_scale_inv 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.kv_a_layernorm.weight 7,15,23,31,39,47,55,63,71-72 +model.layers.{L}.attention.kv_a_proj_with_mqa.weight 7,15,23,31,39,47,55,63,71-72 +model.layers.{L}.attention.kv_a_proj_with_mqa.weight_scale_inv 7,15,23,31,39,47,55,63,71-72 +model.layers.{L}.attention.kv_b_proj.weight 7,15,23,31,39,47,55,63,71-72 +model.layers.{L}.attention.kv_b_proj.weight_scale_inv 7,15,23,31,39,47,55,63,71-72 +model.layers.{L}.attention.o_norm.weight 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.o_proj.weight 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.o_proj.weight_scale_inv 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.q_a_layernorm.weight 7,15,23,31,39,47,55,63,71-72 +model.layers.{L}.attention.q_a_proj.weight 7,15,23,31,39,47,55,63,71-72 +model.layers.{L}.attention.q_a_proj.weight_scale_inv 7,15,23,31,39,47,55,63,71-72 +model.layers.{L}.attention.q_b_proj.weight 7,15,23,31,39,47,55,63,71-72 +model.layers.{L}.attention.q_b_proj.weight_scale_inv 7,15,23,31,39,47,55,63,71-72 +model.layers.{L}.attention.q_conv1d.weight 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.q_proj.weight 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.q_proj.weight_scale_inv 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.v_conv1d.weight 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.v_proj.weight 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.v_proj.weight_scale_inv 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.eh_proj.weight 72 +model.layers.{L}.enorm.weight 72 +model.layers.{L}.final_layernorm.weight 72 +model.layers.{L}.hnorm.weight 72 +model.layers.{L}.input_layernorm.weight 0-72 +model.layers.{L}.mlp.down_proj.weight 0-3 +model.layers.{L}.mlp.down_proj.weight_scale_inv 0-3 +model.layers.{L}.mlp.experts.{E}.down_proj.weight 4-72 512 +model.layers.{L}.mlp.experts.{E}.down_proj.weight_scale_inv 4-72 512 +model.layers.{L}.mlp.experts.{E}.gate_proj.weight 4-72 512 +model.layers.{L}.mlp.experts.{E}.gate_proj.weight_scale_inv 4-72 512 +model.layers.{L}.mlp.experts.{E}.up_proj.weight 4-72 512 +model.layers.{L}.mlp.experts.{E}.up_proj.weight_scale_inv 4-72 512 +model.layers.{L}.mlp.gate.expert_bias 4-72 +model.layers.{L}.mlp.gate.weight 4-72 +model.layers.{L}.mlp.gate_proj.weight 0-3 +model.layers.{L}.mlp.gate_proj.weight_scale_inv 0-3 +model.layers.{L}.mlp.shared_experts.down_proj.weight 4-72 +model.layers.{L}.mlp.shared_experts.down_proj.weight_scale_inv 4-72 +model.layers.{L}.mlp.shared_experts.gate_proj.weight 4-72 +model.layers.{L}.mlp.shared_experts.gate_proj.weight_scale_inv 4-72 +model.layers.{L}.mlp.shared_experts.up_proj.weight 4-72 +model.layers.{L}.mlp.shared_experts.up_proj.weight_scale_inv 4-72 +model.layers.{L}.mlp.up_proj.weight 0-3 +model.layers.{L}.mlp.up_proj.weight_scale_inv 0-3 +model.layers.{L}.post_attention_layernorm.weight 0-72 +model.norm.weight +model.word_embeddings.weight diff --git a/zig/src/families/ling/names.zig b/zig/src/families/ling/names.zig new file mode 100644 index 0000000000..20c4b1e0b6 --- /dev/null +++ b/zig/src/families/ling/names.zig @@ -0,0 +1,220 @@ +//! Ling-3.1-flash's weight map: each tensor name the config implies, by role, checked against an index. +const std = @import("std"); +const config = @import("config.zig"); + +const Config = config.Config; + +pub const Role = enum { + embed, + final_norm, + lm_head, + attn_norm, + mlp_norm, + kda_q, + kda_k, + kda_v, + kda_q_conv, + kda_k_conv, + kda_v_conv, + kda_b, + kda_f, + kda_g, + kda_a_log, + kda_dt_bias, + kda_o_norm, + kda_o, + mla_q, + mla_q_a, + mla_q_a_norm, + mla_q_b, + mla_kv_a, + mla_kv_a_norm, + mla_kv_b, + mla_gate, + mla_o, + dense_gate, + dense_up, + dense_down, + router, + router_bias, + expert_gate, + expert_up, + expert_down, + shared_gate, + shared_up, + shared_down, + mtp_enorm, + mtp_hnorm, + mtp_eh, + mtp_final_norm, +}; + +/// One stored tensor: its role, layer (null for the model's own), routed expert, and whether it is an FP8 scale. +pub const Tensor = struct { role: Role, layer: ?u32 = null, expert: ?u32 = null, scale: bool = false }; + +/// linear: `.weight`, plus `.weight_scale_inv` when FP8; weight: `.weight` only; bare: the module's own name. +const Form = enum { linear, weight, bare }; + +const Part = struct { role: Role, path: []const u8, form: Form }; + +const norms = [_]Part{ + .{ .role = .attn_norm, .path = "input_layernorm", .form = .weight }, + .{ .role = .mlp_norm, .path = "post_attention_layernorm", .form = .weight }, +}; + +const kda = [_]Part{ + .{ .role = .kda_q, .path = "attention.q_proj", .form = .linear }, + .{ .role = .kda_k, .path = "attention.k_proj", .form = .linear }, + .{ .role = .kda_v, .path = "attention.v_proj", .form = .linear }, + .{ .role = .kda_q_conv, .path = "attention.q_conv1d", .form = .weight }, + .{ .role = .kda_k_conv, .path = "attention.k_conv1d", .form = .weight }, + .{ .role = .kda_v_conv, .path = "attention.v_conv1d", .form = .weight }, + .{ .role = .kda_b, .path = "attention.b_proj", .form = .linear }, + .{ .role = .kda_f, .path = "attention.f_proj", .form = .linear }, + .{ .role = .kda_g, .path = "attention.g_proj", .form = .linear }, + .{ .role = .kda_a_log, .path = "attention.A_log", .form = .bare }, + .{ .role = .kda_dt_bias, .path = "attention.dt_bias", .form = .bare }, + .{ .role = .kda_o_norm, .path = "attention.o_norm", .form = .weight }, + .{ .role = .kda_o, .path = "attention.o_proj", .form = .linear }, +}; + +const mla_q_lora = [_]Part{ + .{ .role = .mla_q_a, .path = "attention.q_a_proj", .form = .linear }, + .{ .role = .mla_q_a_norm, .path = "attention.q_a_layernorm", .form = .weight }, + .{ .role = .mla_q_b, .path = "attention.q_b_proj", .form = .linear }, +}; + +const mla_q_whole = [_]Part{.{ .role = .mla_q, .path = "attention.q_proj", .form = .linear }}; + +const mla = [_]Part{ + .{ .role = .mla_kv_a, .path = "attention.kv_a_proj_with_mqa", .form = .linear }, + .{ .role = .mla_kv_a_norm, .path = "attention.kv_a_layernorm", .form = .weight }, + .{ .role = .mla_kv_b, .path = "attention.kv_b_proj", .form = .linear }, + .{ .role = .mla_o, .path = "attention.dense", .form = .linear }, +}; + +const mla_gate = [_]Part{.{ .role = .mla_gate, .path = "attention.g_proj", .form = .linear }}; + +const dense = [_]Part{ + .{ .role = .dense_gate, .path = "mlp.gate_proj", .form = .linear }, + .{ .role = .dense_up, .path = "mlp.up_proj", .form = .linear }, + .{ .role = .dense_down, .path = "mlp.down_proj", .form = .linear }, +}; + +const router = [_]Part{.{ .role = .router, .path = "mlp.gate", .form = .linear }}; +const router_bias = [_]Part{.{ .role = .router_bias, .path = "mlp.gate.expert_bias", .form = .bare }}; + +const shared = [_]Part{ + .{ .role = .shared_gate, .path = "mlp.shared_experts.gate_proj", .form = .linear }, + .{ .role = .shared_up, .path = "mlp.shared_experts.up_proj", .form = .linear }, + .{ .role = .shared_down, .path = "mlp.shared_experts.down_proj", .form = .linear }, +}; + +const expert = [_]Part{ + .{ .role = .expert_gate, .path = "gate_proj", .form = .linear }, + .{ .role = .expert_up, .path = "up_proj", .form = .linear }, + .{ .role = .expert_down, .path = "down_proj", .form = .linear }, +}; + +const mtp = [_]Part{ + .{ .role = .mtp_enorm, .path = "enorm", .form = .weight }, + .{ .role = .mtp_hnorm, .path = "hnorm", .form = .weight }, + .{ .role = .mtp_eh, .path = "eh_proj", .form = .linear }, + .{ .role = .mtp_final_norm, .path = "final_layernorm", .form = .weight }, +}; + +/// The part's tensors under `prefix`: the name built in `buf`, handed to `sink.put(name, tensor)`. +fn emit(c: *const Config, sink: anytype, prefix: []const u8, p: Part, t: Tensor) !void { + var buf: [256]u8 = undefined; + const module = try std.fmt.bufPrint(&buf, "{s}{s}", .{ prefix, p.path }); + var name: [256]u8 = undefined; + var u = t; + u.role = p.role; + switch (p.form) { + .bare => try sink.put(module, u), + .weight, .linear => try sink.put(try std.fmt.bufPrint(&name, "{s}.weight", .{module}), u), + } + if (p.form == .linear and c.scaled(module)) { + u.scale = true; + try sink.put(try std.fmt.bufPrint(&name, "{s}.weight_scale_inv", .{module}), u); + } +} + +fn emitAll(c: *const Config, sink: anytype, prefix: []const u8, parts: []const Part, t: Tensor) !void { + for (parts) |p| try emit(c, sink, prefix, p, t); +} + +/// Layer i's tensors (an MTP layer when i >= c.layers): norms, the mixer its kind names, the MLP or MoE. +pub fn layer(c: *const Config, i: u32, sink: anytype) !void { + var buf: [64]u8 = undefined; + const prefix = try std.fmt.bufPrint(&buf, "model.layers.{d}.", .{i}); + const t: Tensor = .{ .role = .attn_norm, .layer = i }; + try emitAll(c, sink, prefix, &norms, t); + switch (c.kind(i)) { + .kda => try emitAll(c, sink, prefix, &kda, t), + .mla => { + try emitAll(c, sink, prefix, if (c.q_lora > 0) &mla_q_lora else &mla_q_whole, t); + try emitAll(c, sink, prefix, &mla, t); + if (c.head_gate) try emitAll(c, sink, prefix, &mla_gate, t); + }, + } + if (!c.isMoe(i)) return emitAll(c, sink, prefix, &dense, t); + try emitAll(c, sink, prefix, &router, t); + if (c.expert_bias) try emitAll(c, sink, prefix, &router_bias, t); + if (c.shared_inter > 0) try emitAll(c, sink, prefix, &shared, t); + if (i >= c.layers) try emitAll(c, sink, prefix, &mtp, t); + var ebuf: [96]u8 = undefined; + for (0..c.experts) |e| { + const ep = try std.fmt.bufPrint(&ebuf, "{s}mlp.experts.{d}.", .{ prefix, e }); + try emitAll(c, sink, ep, &expert, .{ .role = .expert_gate, .layer = i, .expert = @intCast(e) }); + } +} + +/// The model's own tensors: the embeddings, the final norm and, unless tied, the head. +pub fn globals(c: *const Config, sink: anytype) !void { + try emit(c, sink, "model.", .{ .role = .embed, .path = "word_embeddings", .form = .weight }, .{ .role = .embed }); + try emit(c, sink, "model.", .{ .role = .final_norm, .path = "norm", .form = .weight }, .{ .role = .final_norm }); + if (!c.tie) try emit(c, sink, "", .{ .role = .lm_head, .path = "lm_head", .form = .linear }, .{ .role = .lm_head }); +} + +/// Every tensor the config implies, backbone then MTP layers. +pub fn all(c: *const Config, sink: anytype) !void { + try globals(c, sink); + for (0..c.layers + c.mtp) |i| try layer(c, @intCast(i), sink); +} + +/// The map against an index's names: expected tensors (MTP ones counted apart), missing ones and unaccounted ones. +pub const Report = struct { + expected: usize = 0, + mtp: usize = 0, + missing: std.ArrayList([]const u8) = .empty, + unknown: std.ArrayList([]const u8) = .empty, +}; + +const Checker = struct { + arena: std.mem.Allocator, + seen: *std.StringHashMapUnmanaged(bool), + c: *const Config, + r: *Report, + + pub fn put(k: *Checker, name: []const u8, t: Tensor) !void { + k.r.expected += 1; + if (t.layer != null and t.layer.? >= k.c.layers) k.r.mtp += 1; + if (k.seen.getPtr(name)) |hit| { + if (hit.*) return error.DuplicateName; + hit.* = true; + } else try k.r.missing.append(k.arena, try k.arena.dupe(u8, name)); + } +}; + +/// Every name the map expects is in `index` and every index name is one the map expects (lists in `arena`). +pub fn check(arena: std.mem.Allocator, c: *const Config, index: []const []const u8) !Report { + var seen: std.StringHashMapUnmanaged(bool) = .empty; + try seen.ensureTotalCapacity(arena, @intCast(index.len)); + for (index) |n| if (seen.fetchPutAssumeCapacity(n, false) != null) return error.DuplicateName; + var r: Report = .{}; + var k: Checker = .{ .arena = arena, .seen = &seen, .c = c, .r = &r }; + try all(c, &k); + for (index) |n| if (!seen.get(n).?) try r.unknown.append(arena, n); + return r; +} diff --git a/zig/src/families/ling/names_test.zig b/zig/src/families/ling/names_test.zig new file mode 100644 index 0000000000..03bb485141 --- /dev/null +++ b/zig/src/families/ling/names_test.zig @@ -0,0 +1,150 @@ +//! The weight map against the checkpoint's index: fixtures/names.txt folds its 214,025 names. +const std = @import("std"); +const config = @import("config.zig"); +const names = @import("names.zig"); + +const t = std.testing; +const real = @embedFile("fixtures/config.json"); +const folded = @embedFile("fixtures/names.txt"); + +fn load() !config.Config { + var why: config.Why = .{}; + return config.parse(t.allocator, real, &why); +} + +fn lessThan(_: void, a: []const u8, b: []const u8) bool { + return std.mem.lessThan(u8, a, b); +} + +/// Each `{L}` line over its layer ranges, each `{E}` line over its expert count; the header holds count and sha256. +fn unfold(arena: std.mem.Allocator) ![]const []const u8 { + var out: std.ArrayList([]const u8) = .empty; + var lines = std.mem.tokenizeScalar(u8, folded, '\n'); + var head = std.mem.tokenizeScalar(u8, lines.next().?, ' '); + if (!std.mem.eql(u8, head.next().?, "ling31-names")) return error.BadFixture; + const count = try std.fmt.parseInt(usize, head.next().?, 10); + const digest = head.next().?; + while (lines.next()) |line| { + var f = std.mem.tokenizeScalar(u8, line, ' '); + const pat = f.next().?; + const layers = f.next() orelse { + try out.append(arena, pat); + continue; + }; + const experts = if (f.next()) |e| try std.fmt.parseInt(u32, e, 10) else 1; + var spans = std.mem.tokenizeScalar(u8, layers, ','); + while (spans.next()) |span| { + const dash = std.mem.indexOfScalar(u8, span, '-'); + const lo = try std.fmt.parseInt(u32, if (dash) |d| span[0..d] else span, 10); + const hi = if (dash) |d| try std.fmt.parseInt(u32, span[d + 1 ..], 10) else lo; + for (lo..hi + 1) |l| for (0..experts) |e| { + var lb: [8]u8 = undefined; + var eb: [8]u8 = undefined; + const a = try std.mem.replaceOwned(u8, arena, pat, "{L}", try std.fmt.bufPrint(&lb, "{d}", .{l})); + try out.append(arena, try std.mem.replaceOwned(u8, arena, a, "{E}", try std.fmt.bufPrint(&eb, "{d}", .{e}))); + }; + } + } + std.mem.sort([]const u8, out.items, {}, lessThan); + var h = std.crypto.hash.sha2.Sha256.init(.{}); + for (out.items, 0..) |n, i| { + if (i > 0) h.update("\n"); + h.update(n); + } + var d: [32]u8 = undefined; + h.final(&d); + if (out.items.len != count or !std.mem.eql(u8, &std.fmt.bytesToHex(d, .lower), digest)) return error.FixtureNotIndex; + return out.items; +} + +test "the folded fixture unfolds to the index's exact name list (count and sha256 of the sorted names)" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + try t.expectEqual(@as(usize, 214025), (try unfold(arena.allocator())).len); +} + +test "the map accounts for every index name and finds each it expects: 210,925 backbone, 3,100 MTP layer 72" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const c = try load(); + const r = try names.check(arena.allocator(), &c, try unfold(arena.allocator())); + for (r.missing.items) |n| std.debug.print("missing {s}\n", .{n}); + for (r.unknown.items) |n| std.debug.print("unaccounted {s}\n", .{n}); + try t.expectEqual(@as(usize, 0), r.missing.items.len); + try t.expectEqual(@as(usize, 0), r.unknown.items.len); + try t.expectEqual(@as(usize, 214025), r.expected); + try t.expectEqual(@as(usize, 3100), r.mtp); +} + +const Count = struct { + n: usize = 0, + scales: usize = 0, + roles: std.EnumSet(names.Role) = .empty, + + pub fn put(k: *Count, name: []const u8, x: names.Tensor) !void { + _ = name; + k.n += 1; + k.scales += @intFromBool(x.scale); + k.roles.insert(x.role); + } +}; + +test "per layer: KDA dense 0, KDA MoE 4, MLA MoE 7, the MTP layer 72" { + const c = try load(); + var k0: Count = .{}; + try names.layer(&c, 0, &k0); + try t.expectEqual(@as(usize, 2 + 13 + 6 + 6), k0.n); // norms, KDA parts, their 6 FP8 scales (b_proj is bf16), dense MLP + try t.expect(k0.roles.contains(.kda_a_log) and k0.roles.contains(.dense_down) and !k0.roles.contains(.router)); + var k4: Count = .{}; + try names.layer(&c, 4, &k4); + try t.expectEqual(@as(usize, 2 + 19 + 1 + 1 + 6 + 512 * 6), k4.n); // the router is bf16 and has a bias + var k7: Count = .{}; + try names.layer(&c, 7, &k7); + try t.expectEqual(@as(usize, 2 + 3 + 2 + 4 + 3 + 1 + 2 + 6 + 512 * 6), k7.n); // g_proj of 7 is bf16 + try t.expect(k7.roles.contains(.mla_gate) and !k7.roles.contains(.kda_q)); + var k72: Count = .{}; + try names.layer(&c, 72, &k72); + try t.expectEqual(@as(usize, 3100), k72.n); + try t.expect(k72.roles.contains(.mtp_eh) and k72.roles.contains(.mtp_final_norm)); +} + +test "a config that disagrees with the index is caught: the router taken as FP8 misses 69 scales" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const json = try std.mem.replaceOwned(u8, t.allocator, real, "\"mlp.gate\",", "\"mlp.gate_x\","); + defer t.allocator.free(json); + var why: config.Why = .{}; + const c = try config.parse(t.allocator, json, &why); + const r = try names.check(arena.allocator(), &c, try unfold(arena.allocator())); + try t.expectEqual(@as(usize, 69), r.missing.items.len); + try t.expectEqualStrings("model.layers.4.mlp.gate.weight_scale_inv", r.missing.items[0]); + try t.expectEqual(@as(usize, 0), r.unknown.items.len); +} + +test "tied embeddings drop lm_head, bf16 drops every scale, no head gate drops MLA's g_proj" { + var c = try load(); + c.tie = true; + c.quant = .bf16; + c.head_gate = false; + var g: Count = .{}; + try names.globals(&c, &g); + try t.expectEqual(@as(usize, 2), g.n); + var k7: Count = .{}; + try names.layer(&c, 7, &k7); + try t.expectEqual(@as(usize, 0), k7.scales); + try t.expect(!k7.roles.contains(.mla_gate)); +} + +test "the real index (TF_LING_INDEX=) against the map" { + const path = std.mem.span(std.c.getenv("TF_LING_INDEX") orelse return error.SkipZigTest); + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const bytes = try std.Io.Dir.cwd().readFileAlloc(t.io, path, a, .unlimited); + const doc = try std.json.parseFromSliceLeaky(std.json.Value, a, bytes, .{}); + const map = doc.object.get("weight_map").?.object; + const c = try load(); + const r = try names.check(a, &c, map.keys()); + try t.expectEqual(@as(usize, 0), r.missing.items.len + r.unknown.items.len); + try t.expectEqual(@as(usize, map.count()), r.expected); +} From d46e37234fc9aeb70a689f2b10f101035b97a4ca Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 15 Oct 2026 10:01:40 +0200 Subject: [PATCH 03/34] ling: weights in the checkpoint's own format (bf16, FP8 blocks, ModelOpt NVFP4) at TP=1, synthetic checkpoints in each config.zig: Quant gains nvfp4. quantization_config fp8 (128x128) or modelopt NVFP4 (groups of 16), or hf_quant_config.json's block (parseWith); ignore / exclude_modules / modules_to_not_convert as names, dotted tails or * globs; format(module) gives each linear's format; exl3 refused naming #832. names.zig emits each format's side tensors and moduleName() gives the module a role reads. weights.zig / weights_proj.zig: one Layer per layer with a qlinear.Weight per projection (bf16, fp8g, nvfp4 through modelopt.Packer and fp8.packScales), routed experts as experts.Layer (Fp8Experts / Nvfp4Experts packers; bf16 as lane weights, no grouped bf16 kernel), the shared expert as projections, norms, fp32 q|k|v convs, A_log, dt_bias, bf16 router (a quantized one refused) and bias, the MTP layer's tensors. A file tensor whose dtype or shape disagrees with the config is refused. Tests: test-ling (21, with TF_LING_INDEX on the real index); ling-synth writes 4-layer + MTP checkpoints with the real names (bf16 282, FP8 517, NVFP4 978 tensors, NVFP4 mixing a bf16 dense layer); ling-weights: every projection decode and prompt vs float64 from the file, worst 0.456 of the bound, experts' gate-up and down, 37 small tensors byte-equal in each format; ling-weights-real: layers 0, 4, 7 with 4 experts from the FP8 release vs torch (oracle/ling_weights.py), 35 products, worst 0.414 of the bound, rel L2 <= 1.7e-3. Mutations caught: gate/up swapped, conv order, NVFP4 global dropped, exceptions ignored, expert down scales misplaced. --- zig/build/cuda.zig | 5 + zig/src/families/ling/config.zig | 147 +++++++++-- zig/src/families/ling/config_test.zig | 47 ++++ zig/src/families/ling/cuda.zig | 5 + zig/src/families/ling/names.zig | 36 ++- zig/src/families/ling/names_test.zig | 14 + zig/src/families/ling/synth.zig | 99 +++++++ zig/src/families/ling/weights.zig | 235 +++++++++++++++++ zig/src/families/ling/weights_proj.zig | 333 ++++++++++++++++++++++++ zig/tests/cuda/ling_synth.zig | 154 +++++++++++ zig/tests/cuda/ling_weights_real.zig | 201 +++++++++++++++ zig/tests/cuda/ling_weights_tests.zig | 340 +++++++++++++++++++++++++ zig/tests/cuda/main.zig | 9 + zig/tests/cuda/oracle/ling_weights.py | 62 +++++ 14 files changed, 1654 insertions(+), 33 deletions(-) create mode 100644 zig/src/families/ling/synth.zig create mode 100644 zig/src/families/ling/weights.zig create mode 100644 zig/src/families/ling/weights_proj.zig create mode 100644 zig/tests/cuda/ling_synth.zig create mode 100644 zig/tests/cuda/ling_weights_real.zig create mode 100644 zig/tests/cuda/ling_weights_tests.zig create mode 100644 zig/tests/cuda/oracle/ling_weights.py diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 4c241edf7b..90accb6244 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -154,6 +154,9 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu for (exlr8_build.denseModules(b, target, optimize, exlr8, exlr8_tests[0])) |m| runner.addImport(m.name, m.module); const glm_cuda = glmCuda(b, target, optimize, cuda, mods.core, mods.lanes, mods.heat); runner.addImport("glm_cuda", glm_cuda); + const ling = b.createModule(.{ .root_source_file = b.path("zig/src/families/ling/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true }); + for ([_]struct { []const u8, *std.Build.Module }{ .{ "glm_cuda", glm_cuda }, .{ "cuda", cuda }, .{ "core", mods.core } }) |m| ling.addImport(m[0], m[1]); + runner.addImport("ling_cuda", ling); const runner_exe = b.addExecutable(.{ .name = "tf-cuda-test", .root_module = runner }); b.installArtifact(runner_exe); const controls = if (nvcc) |tool| fatbin(b, tool, version.?, .{ .name = "contract_controls", .flags = &.{"-O3"} }, sms) else null; @@ -275,6 +278,8 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st // Ling-3.1-flash's config loader and weight map, and its routing and limits as the GLM kernels take them const ling = b.createModule(.{ .root_source_file = b.path("zig/src/families/ling/cuda.zig"), .target = host, .optimize = .debug, .link_libc = true }); ling.addImport("glm_cuda", glm); + ling.addImport("cuda", cuda); + ling.addImport("core", mods.core); const ling_run = b.addRunArtifact(b.addTest(.{ .root_module = ling })); step.dependOn(&ling_run.step); b.step("test-ling", "Ling-3.1-flash's config loader and weight map (no GPU work)").dependOn(&ling_run.step); diff --git a/zig/src/families/ling/config.zig b/zig/src/families/ling/config.zig index 40f2f877d1..ba835602c5 100644 --- a/zig/src/families/ling/config.zig +++ b/zig/src/families/ling/config.zig @@ -9,8 +9,8 @@ pub const Kind = enum { kda, mla }; /// MLA's rope pairs: rope_interleave true is exllamav3's RopeStyle.GPTJ (adjacent pairs), false NeoX (halves). pub const RopeStyle = enum { gptj, neox }; -/// The checkpoint's linear weights: bf16, or FP8 e4m3 in 128x128 blocks with a weight_scale_inv each. -pub const Quant = enum { bf16, fp8_block128 }; +/// A linear's format: bf16, FP8 128x128 blocks, ModelOpt NVFP4; EXL3 (#832) adds a variant and a weights_proj reader. +pub const Quant = enum { bf16, fp8_block128, nvfp4 }; /// Why parse refused the config, one sentence for the caller's log line (tests read it too). pub const Why = struct { @@ -66,8 +66,8 @@ pub const Config = struct { eos_n: u32, bos: ?u32, quant: Quant, - plain_text: [1024]u8, // modules_to_not_convert, packed; plain_ends[i] ends entry i - plain_ends: [32]u16, + plain_text: [8192]u8, // modules_to_not_convert (FP8) or ignore / exclude_modules (ModelOpt), packed + plain_ends: [512]u16, // plain_ends[i] ends entry i plain_n: u32, /// Layer i's mixer: the modeling's schedule for the backbone; the MTP layers are MLA (mtp_use_kda false). @@ -103,22 +103,41 @@ pub const Config = struct { return c.plain_text[lo..c.plain_ends[i]]; } - /// A linear module stored without FP8 (modules_to_not_convert names it, or ends it after a dot). + /// A linear module the exception list leaves in bf16: named whole, ended after a dot, or matched by a `*` glob. pub fn unconverted(c: *const Config, module: []const u8) bool { for (0..c.plain_n) |i| { const p = c.plain(@intCast(i)); + if (std.mem.indexOfScalar(u8, p, '*') != null) { + if (glob(p, module)) return true; + continue; + } if (std.mem.eql(u8, module, p)) return true; if (module.len > p.len and std.mem.endsWith(u8, module, p) and module[module.len - p.len - 1] == '.') return true; } return false; } - /// A linear module carries an FP8 weight_scale_inv beside its weight. + /// The stored format of linear `module`: the checkpoint's quantization unless its exception list names it. + pub fn format(c: *const Config, module: []const u8) Quant { + return if (c.unconverted(module)) .bf16 else c.quant; + } + + /// A linear module carries quantization tensors beside its weight. pub fn scaled(c: *const Config, module: []const u8) bool { - return c.quant == .fp8_block128 and !c.unconverted(module); + return c.format(module) != .bf16; } }; +/// fnmatch with `*` only (ModelOpt's ignore patterns): `*` takes any run of characters, dots included. +pub fn glob(pattern: []const u8, name: []const u8) bool { + const star = std.mem.indexOfScalar(u8, pattern, '*') orelse return std.mem.eql(u8, pattern, name); + if (!std.mem.startsWith(u8, name, pattern[0..star])) return false; + const rest = pattern[star + 1 ..]; + var at = star; + while (at <= name.len) : (at += 1) if (glob(rest, name[at..])) return true; + return false; +} + fn get(o: std.json.ObjectMap, key: []const u8, why: *Why) !std.json.Value { return o.get(key) orelse { why.set("config.json has no {s}", .{key}); @@ -190,39 +209,88 @@ fn limits(o: std.json.ObjectMap, key: []const u8, layers: u32, out: *[max_layers } } -fn quantization(o: std.json.ObjectMap, c: *Config, why: *Why) !void { - c.quant = .bf16; - c.plain_n = 0; - const qv = o.get("quantization_config") orelse return; - if (qv == .null) return; - if (qv != .object) { - why.set("quantization_config is not an object", .{}); - return error.Quant; +/// An exception list's names, packed into the config (modules left in bf16). +fn exceptions(v: ?std.json.Value, c: *Config, why: *Why) !void { + const list = v orelse return; + if (list != .array) return; + var at: usize = if (c.plain_n == 0) 0 else c.plain_ends[c.plain_n - 1]; + for (list.array.items) |m| { + if (m != .string or c.plain_n == c.plain_ends.len or at + m.string.len > c.plain_text.len) { + why.set("the quantization's exception list holds more than the loader keeps", .{}); + return error.Quant; + } + @memcpy(c.plain_text[at..][0..m.string.len], m.string); + at += m.string.len; + c.plain_ends[c.plain_n] = @intCast(at); + c.plain_n += 1; } - const q = qv.object; - const method = text(q, "quant_method") orelse ""; +} + +fn fp8Block(q: std.json.ObjectMap, c: *Config, why: *Why) !void { const block = q.get("weight_block_size"); const b128 = block != null and block.? == .array and block.?.array.items.len == 2 and block.?.array.items[0] == .integer and block.?.array.items[0].integer == 128 and block.?.array.items[1] == .integer and block.?.array.items[1].integer == 128; - if (!std.mem.eql(u8, method, "fp8") or !b128 or !std.mem.eql(u8, text(q, "fmt") orelse "e4m3", "e4m3") or + if (!b128 or !std.mem.eql(u8, text(q, "fmt") orelse "e4m3", "e4m3") or !std.mem.eql(u8, text(q, "activation_scheme") orelse "dynamic", "dynamic")) { - why.set("quantization {s}: the loader reads bf16 or FP8 e4m3 in 128x128 blocks", .{if (method.len > 0) method else "without quant_method"}); + why.set("quantization fp8: the loader reads FP8 e4m3 in 128x128 blocks with dynamic activations", .{}); return error.Quant; } c.quant = .fp8_block128; - var at: usize = 0; - if (q.get("modules_to_not_convert")) |list| if (list == .array) for (list.array.items) |m| { - if (m != .string or c.plain_n == c.plain_ends.len or at + m.string.len > c.plain_text.len) { - why.set("modules_to_not_convert holds more than the loader keeps", .{}); - return error.Quant; + try exceptions(q.get("modules_to_not_convert"), c, why); +} + +fn isInt(v: std.json.Value, want: i64) bool { + return v == .integer and v.integer == want; +} + +fn groups16(why: *Why) error{Quant} { + why.set("ModelOpt NVFP4 with weights other than 4-bit in groups of 16", .{}); + return error.Quant; +} + +/// ModelOpt NVFP4 from config.json's quantization_config or hf_quant_config.json's "quantization" block. +fn modelopt(q: std.json.ObjectMap, c: *Config, why: *Why) !void { + const algo = text(q, "quant_algo") orelse ""; + if (!std.ascii.eqlIgnoreCase(algo, "NVFP4")) { + why.set("ModelOpt {s}: the loader reads ModelOpt NVFP4 (e2m1 in groups of 16)", .{if (algo.len > 0) algo else "without quant_algo"}); + return error.Quant; + } + if (q.get("group_size")) |v| if (!isInt(v, 16)) return groups16(why); + if (q.get("config_groups")) |g| if (g == .object) { + var it = g.object.iterator(); + while (it.next()) |e| { + const w = if (e.value_ptr.* == .object) e.value_ptr.object.get("weights") else null; + if (w == null or w.? != .object) return groups16(why); + const bits = w.?.object.get("num_bits") orelse return groups16(why); + const size = w.?.object.get("group_size") orelse return groups16(why); + if (!isInt(bits, 4) or !isInt(size, 16)) return groups16(why); } - @memcpy(c.plain_text[at..][0..m.string.len], m.string); - at += m.string.len; - c.plain_ends[c.plain_n] = @intCast(at); - c.plain_n += 1; }; + c.quant = .nvfp4; + inline for (.{ "ignore", "exclude_modules", "modules_to_not_convert" }) |key| try exceptions(q.get(key), c, why); +} + +fn quantization(o: std.json.ObjectMap, c: *Config, why: *Why) !void { + c.quant = .bf16; + c.plain_n = 0; + const qv = o.get("quantization_config") orelse return; + if (qv == .null) return; + if (qv != .object) { + why.set("quantization_config is not an object", .{}); + return error.Quant; + } + const q = qv.object; + const method = text(q, "quant_method") orelse ""; + if (std.mem.eql(u8, method, "fp8")) return fp8Block(q, c, why); + if (std.ascii.eqlIgnoreCase(method, "modelopt")) return modelopt(q, c, why); + if (std.ascii.eqlIgnoreCase(method, "exl3")) { + why.set("EXL3: Ling's EXL3 weights come with ashhart's EXL3 module (tracker #832), not this loader yet", .{}); + return error.Exl3; + } + why.set("quantization {s}: the loader reads bf16, FP8 128x128 blocks or ModelOpt NVFP4", .{if (method.len > 0) method else "without quant_method"}); + return error.Quant; } fn tokens(o: std.json.ObjectMap, c: *Config, why: *Why) !void { @@ -249,6 +317,29 @@ fn tokens(o: std.json.ObjectMap, c: *Config, why: *Why) !void { /// config.json as the reference (configuration_bailing_moe_v3.py, exllamav3's bailing_hybrid) reads it. pub fn parse(gpa: std.mem.Allocator, json: []const u8, why: *Why) !Config { + return parseWith(gpa, json, null, why); +} + +/// parse() with ModelOpt's hf_quant_config.json beside it: its NVFP4 block when config.json has no quantization_config. +pub fn parseWith(gpa: std.mem.Allocator, json: []const u8, hf_quant: ?[]const u8, why: *Why) !Config { + var c = try parseConfig(gpa, json, why); + const extra = hf_quant orelse return c; + if (c.quant != .bf16 or c.plain_n != 0) return c; + const parsed = std.json.parseFromSlice(std.json.Value, gpa, extra, .{}) catch { + why.set("hf_quant_config.json is not JSON", .{}); + return error.Quant; + }; + defer parsed.deinit(); + const q = if (parsed.value == .object) parsed.value.object.get("quantization") else null; + if (q == null or q.? != .object) { + why.set("hf_quant_config.json has no quantization block", .{}); + return error.Quant; + } + try modelopt(q.?.object, &c, why); + return c; +} + +fn parseConfig(gpa: std.mem.Allocator, json: []const u8, why: *Why) !Config { const parsed = std.json.parseFromSlice(std.json.Value, gpa, json, .{}) catch { why.set("config.json is not JSON", .{}); return error.BadConfig; diff --git a/zig/src/families/ling/config_test.zig b/zig/src/families/ling/config_test.zig index f68eaac905..16d18777fb 100644 --- a/zig/src/families/ling/config_test.zig +++ b/zig/src/families/ling/config_test.zig @@ -154,3 +154,50 @@ test "optional keys: no quantization reads bf16, a null q_lora_rank one q_proj, try t.expectEqual(config.RopeStyle.neox, c.rope_style); try t.expect(!c.scaled("model.layers.0.attention.q_proj")); } + +/// The fixture without its quantization_config block. +fn plainJson() ![]u8 { + const at = std.mem.indexOf(u8, real, "\"quantization_config\"").?; + const end = std.mem.indexOf(u8, real[at..], " },\n").? + at + 5; + return std.mem.concat(t.allocator, u8, &.{ real[0..at], real[end..] }); +} + +test "ModelOpt NVFP4 from quantization_config: ignore globs and names pick the bf16 modules, the rest is nvfp4" { + const json = try edit("\"quant_method\": \"fp8\",", "\"quant_method\": \"modelopt\", \"quant_algo\": \"NVFP4\", \"ignore\": [\"*.attention.b_proj\", \"model.layers.7.mlp*\"],"); + defer t.allocator.free(json); + var why: config.Why = .{}; + const c = try config.parse(t.allocator, json, &why); + try t.expectEqual(config.Quant.nvfp4, c.quant); + try t.expectEqual(config.Quant.nvfp4, c.format("model.layers.0.attention.q_proj")); + try t.expectEqual(config.Quant.bf16, c.format("model.layers.30.attention.b_proj")); + try t.expectEqual(config.Quant.bf16, c.format("model.layers.7.mlp.experts.3.up_proj")); + try t.expectEqual(config.Quant.nvfp4, c.format("model.layers.17.mlp.experts.3.up_proj")); + try t.expectEqual(config.Quant.bf16, c.format("lm_head")); // modules_to_not_convert stays an exception list too +} + +test "hf_quant_config.json's NVFP4 block when config.json has none; groups other than 16 and other algorithms refused" { + const json = try plainJson(); + defer t.allocator.free(json); + var why: config.Why = .{}; + const hf = "{\"producer\": {\"name\": \"modelopt\"}, \"quantization\": {\"quant_algo\": \"NVFP4\", \"group_size\": 16, \"exclude_modules\": [\"lm_head\", \"model.layers.*.mlp.gate\"]}}"; + const c = try config.parseWith(t.allocator, json, hf, &why); + try t.expectEqual(config.Quant.nvfp4, c.quant); + try t.expectEqual(config.Quant.bf16, c.format("model.layers.40.mlp.gate")); + try t.expectEqual(config.Quant.nvfp4, c.format("model.layers.40.mlp.gate_proj")); + const g32 = "{\"quantization\": {\"quant_algo\": \"NVFP4\", \"group_size\": 32}}"; + try t.expectError(error.Quant, config.parseWith(t.allocator, json, g32, &why)); + try t.expect(std.mem.indexOf(u8, why.text(), "groups of 16") != null); + const fp8 = "{\"quantization\": {\"quant_algo\": \"FP8\"}}"; + try t.expectError(error.Quant, config.parseWith(t.allocator, json, fp8, &why)); + try t.expect(std.mem.indexOf(u8, why.text(), "ModelOpt FP8") != null); + try t.expectEqual(config.Quant.bf16, (try config.parseWith(t.allocator, json, null, &why)).quant); +} + +test "EXL3 is refused naming #832; a glob matches whole names only" { + try refused("\"quant_method\": \"fp8\"", "\"quant_method\": \"exl3\"", error.Exl3, "#832"); + try t.expect(config.glob("*.mlp.gate", "model.layers.4.mlp.gate")); + try t.expect(!config.glob("*.mlp.gate", "model.layers.4.mlp.gate_proj")); + try t.expect(config.glob("model.layers.*.attention*", "model.layers.12.attention.q_proj")); + try t.expect(!config.glob("model.layers.1*", "model.layers.0.mlp")); + try t.expect(config.glob("a*b*c", "a.b.x.c") and !config.glob("a*b*c", "a.c")); +} diff --git a/zig/src/families/ling/cuda.zig b/zig/src/families/ling/cuda.zig index 615e36b02b..5d534ed0b0 100644 --- a/zig/src/families/ling/cuda.zig +++ b/zig/src/families/ling/cuda.zig @@ -3,6 +3,8 @@ const std = @import("std"); const glm = @import("glm_cuda"); pub const config = @import("config.zig"); pub const names = @import("names.zig"); +pub const weights = @import("weights.zig"); +pub const synth = @import("synth.zig"); const Config = config.Config; @@ -34,6 +36,9 @@ pub fn route(c: *const Config, i: u32, l: *glm.moe.Layer) void { test { _ = @import("config_test.zig"); _ = @import("names_test.zig"); + _ = @import("synth.zig"); + std.testing.refAllDecls(weights); // the loader compiles on hosts without a GPU + std.testing.refAllDecls(@import("weights_proj.zig")); } test "Ling's routing as glue.select takes it: 8 groups keeping 4, and select's own bounds hold" { diff --git a/zig/src/families/ling/names.zig b/zig/src/families/ling/names.zig index 20c4b1e0b6..b6f3786acb 100644 --- a/zig/src/families/ling/names.zig +++ b/zig/src/families/ling/names.zig @@ -49,10 +49,10 @@ pub const Role = enum { mtp_final_norm, }; -/// One stored tensor: its role, layer (null for the model's own), routed expert, and whether it is an FP8 scale. +/// One stored tensor: its role, layer (null for the model's own), routed expert, and whether it is a scale tensor. pub const Tensor = struct { role: Role, layer: ?u32 = null, expert: ?u32 = null, scale: bool = false }; -/// linear: `.weight`, plus `.weight_scale_inv` when FP8; weight: `.weight` only; bare: the module's own name. +/// linear: `.weight` plus its format's side tensors (side()); weight: `.weight` only; bare: the module's own name. const Form = enum { linear, weight, bare }; const Part = struct { role: Role, path: []const u8, form: Form }; @@ -134,10 +134,36 @@ fn emit(c: *const Config, sink: anytype, prefix: []const u8, p: Part, t: Tensor) .bare => try sink.put(module, u), .weight, .linear => try sink.put(try std.fmt.bufPrint(&name, "{s}.weight", .{module}), u), } - if (p.form == .linear and c.scaled(module)) { - u.scale = true; - try sink.put(try std.fmt.bufPrint(&name, "{s}.weight_scale_inv", .{module}), u); + if (p.form != .linear) return; + u.scale = true; + for (side(c.format(module))) |tail| try sink.put(try std.fmt.bufPrint(&name, "{s}.{s}", .{ module, tail }), u); +} + +/// The tensors a linear stores beside `.weight`: FP8's block scales, ModelOpt NVFP4's group, global and input scales. +pub fn side(f: config.Quant) []const []const u8 { + return switch (f) { + .bf16 => &.{}, + .fp8_block128 => &.{"weight_scale_inv"}, + .nvfp4 => &.{ "weight_scale", "weight_scale_2", "input_scale" }, + }; +} + +const tables = [_][]const Part{ &norms, &kda, &mla_q_lora, &mla_q_whole, &mla, &mla_gate, &dense, &router, &router_bias, &shared, &expert, &mtp }; + +/// Role's module in layer i (`e` for routed experts) as the map names it: the name before `.weight`. +pub fn moduleName(buf: []u8, i: u32, role: Role, e: ?u32) ![]const u8 { + switch (role) { + .embed => return "model.word_embeddings", + .final_norm => return "model.norm", + .lm_head => return "lm_head", + .expert_gate, .expert_up, .expert_down => { + const p = for (expert) |x| (if (x.role == role) break x) else unreachable; + return std.fmt.bufPrint(buf, "model.layers.{d}.mlp.experts.{d}.{s}", .{ i, e orelse return error.NoExpert, p.path }); + }, + else => {}, } + for (tables) |t| for (t) |p| if (p.role == role) return std.fmt.bufPrint(buf, "model.layers.{d}.{s}", .{ i, p.path }); + unreachable; } fn emitAll(c: *const Config, sink: anytype, prefix: []const u8, parts: []const Part, t: Tensor) !void { diff --git a/zig/src/families/ling/names_test.zig b/zig/src/families/ling/names_test.zig index 03bb485141..17aef45f9b 100644 --- a/zig/src/families/ling/names_test.zig +++ b/zig/src/families/ling/names_test.zig @@ -148,3 +148,17 @@ test "the real index (TF_LING_INDEX=) against the try t.expectEqual(@as(usize, 0), r.missing.items.len + r.unknown.items.len); try t.expectEqual(@as(usize, map.count()), r.expected); } + +test "NVFP4 linears store weight, weight_scale, weight_scale_2 and input_scale; the module names the loader reads" { + var c = try load(); + c.quant = .nvfp4; + var k0: Count = .{}; + try names.layer(&c, 0, &k0); + try t.expectEqual(@as(usize, 2 + 13 + 6 * 3 + 3 + 3 * 3), k0.n); // norms, KDA, 3 side tensors a quantized linear, dense + var b: [192]u8 = undefined; + try t.expectEqualStrings("model.layers.4.mlp.experts.511.down_proj", try names.moduleName(&b, 4, .expert_down, 511)); + try t.expectEqualStrings("model.layers.7.attention.dense", try names.moduleName(&b, 7, .mla_o, null)); + try t.expectEqualStrings("model.layers.0.attention.A_log", try names.moduleName(&b, 0, .kda_a_log, null)); + try t.expectEqualStrings("model.layers.72.eh_proj", try names.moduleName(&b, 72, .mtp_eh, null)); + try t.expectEqualStrings("lm_head", try names.moduleName(&b, 0, .lm_head, null)); +} diff --git a/zig/src/families/ling/synth.zig b/zig/src/families/ling/synth.zig new file mode 100644 index 0000000000..593f8ba916 --- /dev/null +++ b/zig/src/families/ling/synth.zig @@ -0,0 +1,99 @@ +//! A tiny Ling-3.1 config: the real config.json's schema with fewer heads, experts and dims, in any loader format. +const std = @import("std"); +const config = @import("config.zig"); + +pub const real_config = @embedFile("fixtures/config.json"); + +/// The shrunk widths; KDA's head dim 128, MLA's 512 latent, 64 rope and 128-wide heads stay as the kernels take them. +pub const shape = .{ + .hidden = 512, + .heads = 4, + .dense_inter = 256, + .moe_inter = 128, + .experts = 16, + .topk = 4, + .n_group = 4, + .topk_group = 2, + .q_lora = 256, + .vocab = 1024, + .group = 2, // KDA, MLA, KDA, MLA, ... +}; + +/// The tiny config for `layers` backbone layers (at least 2: layer 0 KDA + dense, 1 MLA + MoE) and one MTP layer. +pub fn configJson(gpa: std.mem.Allocator, format: config.Quant, layers: u32) ![]u8 { + if (layers < 2 or layers > config.max_layers - 1) return error.BadLayerCount; + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var root = try std.json.parseFromSliceLeaky(std.json.Value, a, real_config, .{}); + const o = &root.object; + const ints = .{ + .{ "hidden_size", shape.hidden }, .{ "num_attention_heads", shape.heads }, .{ "num_key_value_heads", shape.heads }, + .{ "intermediate_size", shape.dense_inter }, .{ "moe_intermediate_size", shape.moe_inter }, .{ "moe_shared_expert_intermediate_size", shape.moe_inter }, + .{ "num_experts", shape.experts }, .{ "num_experts_per_tok", shape.topk }, .{ "n_group", shape.n_group }, + .{ "topk_group", shape.topk_group }, .{ "q_lora_rank", shape.q_lora }, .{ "vocab_size", shape.vocab }, + .{ "eos_token_id", shape.vocab - 1 }, .{ "pad_token_id", shape.vocab - 2 }, .{ "layer_group_size", shape.group }, + .{ "first_k_dense_replace", 1 }, .{ "num_hidden_layers", layers }, .{ "num_nextn_predict_layers", 1 }, + }; + inline for (ints) |kv| try o.put(a, kv[0], .{ .integer = kv[1] }); + // clamps on the last layer only: routed and shared, as layers 65-71 carry them in the real file + inline for (.{ "expert_swiglu_limit_list", "share_expert_swiglu_limit_list" }) |key| { + var list = std.json.Array.init(a); + for (0..layers) |i| try list.append(.{ .integer = if (i + 1 == layers) 10 else 0 }); + try o.put(a, key, .{ .array = list }); + } + switch (format) { + .bf16 => _ = o.orderedRemove("quantization_config"), + .fp8_block128 => try o.put(a, "quantization_config", try quant(a, fp8_block, layers)), + .nvfp4 => try o.put(a, "quantization_config", try quant(a, modelopt_block, layers)), + } + return std.json.Stringify.valueAlloc(gpa, root, .{ .whitespace = .indent_2 }); +} + +const fp8_block = + \\{"activation_scheme": "dynamic", "fmt": "e4m3", "quant_method": "fp8", "weight_block_size": [128, 128], + \\ "modules_to_not_convert": ["lm_head", "model.word_embeddings", "b_proj", "mlp.gate"]} +; + +/// ModelOpt NVFP4 as nvidia's checkpoints write it; `*` globs, and layer 0's dense MLP left in bf16 to mix formats. +const modelopt_block = + \\{"quant_method": "modelopt", "quant_algo": "NVFP4", "config_groups": {"group_0": {"targets": ["Linear"], + \\ "weights": {"dynamic": false, "num_bits": 4, "type": "float", "group_size": 16}, + \\ "input_activations": {"dynamic": false, "num_bits": 4, "type": "float", "group_size": 16}}}, + \\ "ignore": ["lm_head", "model.word_embeddings", "*.b_proj", "*.mlp.gate", "model.layers.0.mlp*"]} +; + +/// The block with the real file's per-layer exceptions added: each MLA layer's g_proj and the MTP layer's eh_proj. +fn quant(a: std.mem.Allocator, block: []const u8, layers: u32) !std.json.Value { + var q = try std.json.parseFromSliceLeaky(std.json.Value, a, block, .{}); + const key = if (q.object.contains("ignore")) "ignore" else "modules_to_not_convert"; + const list = &q.object.getPtr(key).?.array; + var i: u32 = 1; + while (i < layers) : (i += shape.group) try list.append(.{ .string = try std.fmt.allocPrint(a, "model.layers.{d}.attention.g_proj", .{i}) }); + try list.append(.{ .string = try std.fmt.allocPrint(a, "model.layers.{d}.eh_proj", .{layers}) }); + return q; +} + +test "the tiny configs parse in every format: KDA dense 0, MLA MoE 1, two groups, the MTP layer, the exceptions" { + const gpa = std.testing.allocator; + for ([_]config.Quant{ .bf16, .fp8_block128, .nvfp4 }) |f| { + const json = try configJson(gpa, f, 4); + defer gpa.free(json); + var why: config.Why = .{}; + const c = config.parse(gpa, json, &why) catch |e| { + std.debug.print("{t}: {s}\n", .{ f, why.text() }); + return e; + }; + try std.testing.expectEqual(f, c.quant); + try std.testing.expect(c.kind(0) == .kda and c.kind(1) == .mla and c.kind(2) == .kda and c.kind(4) == .mla); + try std.testing.expect(!c.isMoe(0) and c.isMoe(1) and c.mtp == 1 and c.n_group == 4 and c.topk_group == 2); + try std.testing.expectEqual(@as(f32, 10), c.expert_limit[3]); + try std.testing.expectEqual(@as(u32, 512), c.kv_lora); + try std.testing.expectEqual(config.Quant.bf16, c.format("model.layers.3.attention.g_proj")); + try std.testing.expectEqual(config.Quant.bf16, c.format("model.layers.2.attention.b_proj")); + try std.testing.expectEqual(config.Quant.bf16, c.format("model.layers.4.eh_proj")); + try std.testing.expectEqual(f, c.format("model.layers.2.attention.g_proj")); + try std.testing.expectEqual(f, c.format("model.layers.1.mlp.experts.3.down_proj")); + try std.testing.expectEqual(if (f == .nvfp4) .bf16 else f, c.format("model.layers.0.mlp.up_proj")); + } +} diff --git a/zig/src/families/ling/weights.zig b/zig/src/families/ling/weights.zig new file mode 100644 index 0000000000..32ab240336 --- /dev/null +++ b/zig/src/families/ling/weights.zig @@ -0,0 +1,235 @@ +//! Ling-3.1-flash's weights on CUDA at TP=1, each tensor in the format its checkpoint stores, layer by layer. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const config = @import("config.zig"); +const names = @import("names.zig"); +pub const proj = @import("weights_proj.zig"); + +const Config = config.Config; +const Role = names.Role; +pub const Proj = proj.Proj; +pub const Mlp = proj.Mlp; +pub const Routed = proj.Routed; +pub const Store = proj.Store; + +/// KDA: full-rank projections, the q|k|v short convs as #804's chain reads them, the decay and output norm tensors. +pub const Kda = struct { + q: Proj, + k: Proj, + v: Proj, + b: Proj, // beta, a row per head + f: Proj, // the decay's full-rank projection + g: Proj, // the output gate's full-rank projection + o: Proj, + conv: u64, // fp32 [3 heads dk][taps]: q_conv1d | k_conv1d | v_conv1d + a_log: u64, // fp32 [heads] + dt_bias: u64, // fp32 [heads dk] + o_norm: u64, // bf16 [dk] +}; + +/// MLA: q through the q_lora pair (or one q_proj), the latent kv pair, the head-wise gate, the output (`dense`). +pub const Mla = struct { + q: ?Proj = null, // without q_lora_rank + q_a: ?Proj = null, + q_a_norm: u64 = 0, // bf16 [q_lora] + q_b: ?Proj = null, + kv_a: Proj, // [kv_lora + rope, hidden] + kv_a_norm: u64, // bf16 [kv_lora] + kv_b: Proj, // [heads (nope + v), kv_lora] + gate: ?Proj = null, // head_wise: [heads, hidden] + o: Proj, +}; + +pub const Mixer = union(config.Kind) { kda: Kda, mla: Mla }; + +/// A MoE block: bf16 router and its pick bias, routed experts, the shared expert as projections, SwiGLU limits. +pub const Moe = struct { + router: u64, // bf16 [experts][hidden], as GLM's router kernel reads it + bias: u64, // fp32 [experts], 0 without moe_router_enable_expert_bias + routed: Routed, + shared: ?Mlp, + experts: usize, // loaded experts (all but in a test's partial load) + shared_limit: f32, +}; + +pub const Ffn = union(enum) { dense: Mlp, moe: Moe }; + +pub const Layer = struct { index: u32, attn_norm: u64, mlp_norm: u64, mixer: Mixer, ffn: Ffn }; + +/// The MTP layer's own tensors beside its MLA + MoE layer: the two input norms, eh_proj, its final norm. +pub const Mtp = struct { enorm: u64, hnorm: u64, eh: Proj, norm: u64 }; + +pub const Options = struct { + globals: bool = true, // embedding, final norm and head + experts: ?usize = null, // the first n routed experts of each MoE layer only (tests on a partial checkpoint) +}; + +pub const Weights = struct { + config: Config, + store: Store, + embed: u64 = 0, // bf16 [vocab][hidden] + norm: u64 = 0, + head: ?Proj = null, // null when tied: the embedding's rows + layers: []Layer = &.{}, + mtp: ?Mtp = null, + + pub fn deinit(w: *Weights) void { + w.store.gpa.free(w.layers); + w.store.deinit(); + w.* = undefined; + } + + /// Layer `i` among the loaded ones. + pub fn layer(w: *const Weights, i: u32) ?*const Layer { + for (w.layers) |*l| if (l.index == i) return l; + return null; + } +}; + +const Loader = struct { + r: proj.Reader, + c: *const Config, + opts: Options, + buf: [192]u8 = undefined, + + fn module(L: *Loader, i: u32, role: Role) ![]const u8 { + return names.moduleName(&L.buf, i, role, null); + } + + fn lin(L: *Loader, i: u32, role: Role, n: usize, k: usize) !Proj { + return L.r.linear(try L.module(i, role), n, k); + } + + /// A norm (`.weight`) or bare tensor (A_log, dt_bias, expert_bias) of `numel` values, as `to`. + fn small(L: *Loader, i: u32, role: Role, numel: usize, to: core.checkpoint.DType) !u64 { + var b: [200]u8 = undefined; + const m = try L.module(i, role); + const bare = role == .kda_a_log or role == .kda_dt_bias or role == .router_bias; + return L.r.small(if (bare) m else try std.fmt.bufPrint(&b, "{s}.weight", .{m}), numel, to); + } + + fn kda(L: *Loader, i: u32) !Kda { + const c = L.c; + const n: usize = c.kda_heads * c.kda_dim; + const d: usize = c.hidden; + var b: [64]u8 = undefined; + var cb: [96]u8 = undefined; + return .{ + .q = try L.lin(i, .kda_q, n, d), + .k = try L.lin(i, .kda_k, n, d), + .v = try L.lin(i, .kda_v, n, d), + .b = try L.lin(i, .kda_b, c.kda_heads, d), + .f = try L.lin(i, .kda_f, n, d), + .g = try L.lin(i, .kda_g, n, d), + .o = try L.lin(i, .kda_o, d, n), + .conv = try L.r.convs(try std.fmt.bufPrint(&b, "model.layers.{d}.attention.", .{i}), n, try std.fmt.bufPrint(&cb, "model.layers.{d}.attention.conv", .{i})), + .a_log = try L.small(i, .kda_a_log, c.kda_heads, .f32), + .dt_bias = try L.small(i, .kda_dt_bias, n, .f32), + .o_norm = try L.small(i, .kda_o_norm, c.kda_dim, .bf16), + }; + } + + fn mla(L: *Loader, i: u32) !Mla { + const c = L.c; + const d: usize = c.hidden; + const h: usize = c.mla_heads; + var m: Mla = .{ + .kv_a = try L.lin(i, .mla_kv_a, c.kv_lora + c.rope_dim, d), + .kv_a_norm = try L.small(i, .mla_kv_a_norm, c.kv_lora, .bf16), + .kv_b = try L.lin(i, .mla_kv_b, h * (c.nope + c.v_dim), c.kv_lora), + .o = try L.lin(i, .mla_o, d, h * c.v_dim), + }; + if (c.q_lora > 0) { + m.q_a = try L.lin(i, .mla_q_a, c.q_lora, d); + m.q_a_norm = try L.small(i, .mla_q_a_norm, c.q_lora, .bf16); + m.q_b = try L.lin(i, .mla_q_b, h * c.qkDim(), c.q_lora); + } else m.q = try L.lin(i, .mla_q, h * c.qkDim(), d); + if (c.head_gate) m.gate = try L.lin(i, .mla_gate, h, d); + return m; + } + + fn moe(L: *Loader, i: u32) !Moe { + const c = L.c; + const d: usize = c.hidden; + const route = try L.module(i, .router); + if (c.format(route) != .bf16) { + std.log.err("{s} is {t}: the routing kernels read a bf16 router", .{ route, c.format(route) }); + return error.QuantizedRouter; + } + const count = @min(c.experts, L.opts.experts orelse c.experts); + const lim = @import("cuda.zig").limits(c, i); + var b: [96]u8 = undefined; + return .{ + .router = try L.small(i, .router, c.experts * d, .bf16), + .bias = if (c.expert_bias) try L.small(i, .router_bias, c.experts, .f32) else 0, + .routed = try L.r.experts(i, count, c.moe_inter, d, lim.routed), + .shared = if (c.shared_inter > 0) try L.r.mlp(try std.fmt.bufPrint(&b, "model.layers.{d}.mlp.shared_experts.", .{i}), c.shared_inter, d) else null, + .experts = count, + .shared_limit = lim.shared, + }; + } + + fn layer(L: *Loader, i: u32) !Layer { + const c = L.c; + var b: [64]u8 = undefined; + return .{ + .index = i, + .attn_norm = try L.small(i, .attn_norm, c.hidden, .bf16), + .mlp_norm = try L.small(i, .mlp_norm, c.hidden, .bf16), + .mixer = switch (c.kind(i)) { + .kda => .{ .kda = try L.kda(i) }, + .mla => .{ .mla = try L.mla(i) }, + }, + .ffn = if (c.isMoe(i)) .{ .moe = try L.moe(i) } else .{ .dense = try L.r.mlp(try std.fmt.bufPrint(&b, "model.layers.{d}.mlp.", .{i}), c.dense_inter, c.hidden) }, + }; + } + + fn mtp(L: *Loader, i: u32) !Mtp { + const d: usize = L.c.hidden; + return .{ + .enorm = try L.small(i, .mtp_enorm, d, .bf16), + .hnorm = try L.small(i, .mtp_hnorm, d, .bf16), + .eh = try L.lin(i, .mtp_eh, d, 2 * d), + .norm = try L.small(i, .mtp_final_norm, d, .bf16), + }; + } +}; + +/// Layers `layers` (index c.layers is the MTP layer) from `ck`, each tensor in its stored format; `stream` blocking. +pub fn load(gpa: std.mem.Allocator, d: *const cuda.Driver, stream: cuda.Stream, pack: *const cuda.modelopt.Packer, ck: *core.Checkpoint, c: Config, layers: []const u32, opts: Options) !Weights { + var w: Weights = .{ .config = c, .store = .{ .gpa = gpa, .d = d } }; + errdefer w.store.deinit(); + var L: Loader = .{ .r = .{ .gpa = gpa, .stream = stream, .pack = pack, .ck = ck, .c = &w.config, .store = &w.store }, .c = &w.config, .opts = opts }; + const out = try gpa.alloc(Layer, layers.len); + errdefer gpa.free(out); + for (layers, out) |i, *l| { + if (i >= c.layers + c.mtp) return error.BadLayer; + l.* = try L.layer(i); + if (i >= c.layers) w.mtp = try L.mtp(i); + ck.release(); // a layer's pages out of the cache before the next: the GPU shares this memory + } + w.layers = out; + if (!opts.globals) return w; + const e = try ck.get("model.word_embeddings.weight"); + if (!e.is(.bf16, &.{ c.vocab, c.hidden })) { + std.log.err("model.word_embeddings: {t} {any}, expected bf16 [{d}, {d}]", .{ e.dtype, e.shape[0..e.rank], c.vocab, c.hidden }); + return error.UnexpectedTensor; + } + w.embed = try w.store.keep(e.bytes, "model.word_embeddings.weight"); + w.norm = try L.r.small("model.norm.weight", c.hidden, .bf16); + if (!c.tie) w.head = try L.r.linear("lm_head", c.vocab, c.hidden); + return w; +} + +/// load() from a checkpoint directory: config.json (and hf_quant_config.json, if any) read there too. +pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, why: *config.Why) !struct { ck: core.Checkpoint, c: Config } { + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + const cwd = std.Io.Dir.cwd(); + const json = try cwd.readFileAlloc(io, try std.fs.path.join(a, &.{ dir, "config.json" }), a, .limited(1 << 22)); + const hf = cwd.readFileAlloc(io, try std.fs.path.join(a, &.{ dir, "hf_quant_config.json" }), a, .limited(1 << 22)) catch null; + const c = try config.parseWith(gpa, json, hf, why); + return .{ .ck = try core.Checkpoint.openModel(gpa, io, dir), .c = c }; +} diff --git a/zig/src/families/ling/weights_proj.zig b/zig/src/families/ling/weights_proj.zig new file mode 100644 index 0000000000..a1b5dbd8f7 --- /dev/null +++ b/zig/src/families/ling/weights_proj.zig @@ -0,0 +1,333 @@ +//! Ling's linears and routed experts as stored (bf16, FP8 blocks, ModelOpt NVFP4) into qlinear and experts views. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const config = @import("config.zig"); + +const Quant = config.Quant; +const Tensor = core.checkpoint.Tensor; +const X8 = cuda.experts.Fp8Experts; +const X4 = cuda.experts.Nvfp4Experts; + +/// One projection in the format its checkpoint stores: a lane matmul view (bf16, fp8g, nvfp4). +pub const Proj = cuda.qlinear.Weight; + +/// Gate, up and down: the dense MLP, the shared expert, and one bf16 routed expert. +pub const Mlp = struct { gate: Proj, up: Proj, down: Proj }; + +/// Routed experts: the grouped kernels' layout (fp8g, nvfp4), or bf16 lane weights (no grouped bf16 kernel). +pub const Routed = union(enum) { grouped: cuda.experts.Layer, lanes: []Mlp }; + +/// One upload, by the checkpoint name it came from (tests find small tensors by it). +pub const Named = struct { name: []u8, ptr: u64, len: usize }; + +/// Every device buffer a load made, freed together. +pub const Store = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + buffers: std.ArrayList(cuda.DeviceBuffer) = .empty, + named: std.ArrayList(Named) = .empty, + lanes: std.ArrayList([]Mlp) = .empty, + + pub fn deinit(s: *Store) void { + for (s.buffers.items) |*b| b.free(); + for (s.named.items) |n| s.gpa.free(n.name); + for (s.lanes.items) |l| s.gpa.free(l); + s.buffers.deinit(s.gpa); + s.named.deinit(s.gpa); + s.lanes.deinit(s.gpa); + } + + /// A device buffer of `len` bytes the store owns, named `name`. + pub fn alloc(s: *Store, len: usize, name: []const u8) !u64 { + const b = try cuda.DeviceBuffer.alloc(s.d, @max(len, 16)); + errdefer { + var x = b; + x.free(); + } + try s.buffers.append(s.gpa, b); + try s.named.append(s.gpa, .{ .name = try s.gpa.dupe(u8, name), .ptr = b.ptr, .len = len }); + return b.ptr; + } + + /// `bytes` uploaded into a buffer the store owns. + pub fn keep(s: *Store, bytes: []const u8, name: []const u8) !u64 { + const p = try s.alloc(bytes.len, name); + try s.buffers.items[s.buffers.items.len - 1].upload(0, bytes); + return p; + } + + pub fn find(s: *const Store, name: []const u8) ?Named { + for (s.named.items) |n| if (std.mem.eql(u8, n.name, name)) return n; + return null; + } +}; + +/// Host bf16 from fp32: round to nearest even, NaN kept quiet. +pub fn toBf16(v: f32) u16 { + const b: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((b >> 16) | 0x40); + return @intCast((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +/// Element i of a bf16 or fp32 tensor. +pub fn value(t: Tensor, i: usize) !f32 { + return switch (t.dtype) { + .bf16 => @bitCast(@as(u32, std.mem.readInt(u16, t.bytes[2 * i ..][0..2], .little)) << 16), + .f32 => @bitCast(std.mem.readInt(u32, t.bytes[4 * i ..][0..4], .little)), + else => error.UnexpectedTensor, + }; +} + +/// What reading a layer needs: the checkpoint, the config, the GPU packers, the store the uploads land in. +pub const Reader = struct { + gpa: std.mem.Allocator, + stream: cuda.Stream, // blocking: the packers wait for the uploads + pack: *const cuda.modelopt.Packer, + ck: *core.Checkpoint, + c: *const config.Config, + store: *Store, + buf: [256]u8 = undefined, + + fn name(r: *Reader, module: []const u8, tail: []const u8) ![]const u8 { + return std.fmt.bufPrint(&r.buf, "{s}.{s}", .{ module, tail }); + } + + /// The config gives `module` format `f`; the file disagrees (dtype or shape): refused, both named. + fn wrong(module: []const u8, f: Quant, t: Tensor, n: usize, k: usize) error{UnexpectedTensor} { + std.log.err("{s}: the config makes it {t} [{d}, {d}], the file holds {t} {any}", .{ module, f, n, k, t.dtype, t.shape[0..t.rank] }); + return error.UnexpectedTensor; + } + + /// Linear `module` [n, k] in the format its config gives it; a file tensor of another format or shape is refused. + pub fn linear(r: *Reader, module: []const u8, n: usize, k: usize) !Proj { + if (k % cuda.qmmf.group != 0) return error.UnexpectedTensor; + return switch (r.c.format(module)) { + .bf16 => r.bf16(module, n, k), + .fp8_block128 => r.fp8(module, n, k), + .nvfp4 => r.fp4(module, n, k), + }; + } + + /// A temporary device copy of `bytes` for a packer, freed by the caller after the stream drains. + fn staged(r: *Reader, bytes: []const u8) !cuda.DeviceBuffer { + return cuda.DeviceBuffer.fromHost(r.store.d, bytes); + } + + fn bf16(r: *Reader, module: []const u8, n: usize, k: usize) !Proj { + const t = try r.ck.get(try r.name(module, "weight")); + if (!t.is(.bf16, &.{ n, k })) return wrong(module, .bf16, t, n, k); + var src = try r.staged(t.bytes); + defer src.free(); + const out = try r.store.alloc(cuda.modelopt.padded(n) * k * 2, module); + try r.pack.bf16Lane(r.stream, src.ptr, out, n, k); + try r.stream.synchronize(); + return .{ .bf16 = cuda.modelopt.bf16Weight(out, n, k) }; + } + + /// e4m3 [n, k] and fp32 scale_inv [n/128 up, k/128] checked against the config's shape. + fn fp8Pair(r: *Reader, module: []const u8, n: usize, k: usize) !struct { w: Tensor, s: Tensor } { + const w = try r.ck.get(try r.name(module, "weight")); + if (!w.is(.f8_e4m3, &.{ n, k }) or k % cuda.fp8.block != 0) return wrong(module, .fp8_block128, w, n, k); + const s = try r.ck.get(try r.name(module, "weight_scale_inv")); + if (!s.is(.f32, &.{ (n + 127) / 128, k / 128 })) return wrong(module, .fp8_block128, s, n, k); + return .{ .w = w, .s = s }; + } + + fn fp8(r: *Reader, module: []const u8, n: usize, k: usize) !Proj { + const p = try r.fp8Pair(module, n, k); + var src = try r.staged(p.w.bytes); + defer src.free(); + const codes = try r.store.alloc(cuda.fp8.padded(n) * k, module); + try r.pack.fp8Lane(r.stream, src.ptr, codes, n, k); + const inv = try r.gpa.alloc(f32, p.s.numel()); + defer r.gpa.free(inv); + @memcpy(std.mem.sliceAsBytes(inv), p.s.bytes); + const tiles = try r.gpa.alloc(f32, cuda.fp8.padded(n) * (k / cuda.fp8.group)); + defer r.gpa.free(tiles); + cuda.fp8.packScales(tiles, inv, n, k); + const scales = try r.store.keep(std.mem.sliceAsBytes(tiles), try r.name(module, "weight_scale_inv")); + try r.stream.synchronize(); + return .{ .fp8g = cuda.fp8.weight(codes, scales, n, k) }; + } + + /// ModelOpt NVFP4 [n, k]: a weight stored without its weight_scale is refused before core.nvfp4 reads it. + fn fp4Read(r: *Reader, module: []const u8, n: usize, k: usize) !core.nvfp4.Linear { + if (!r.ck.has(try r.name(module, "weight_scale"))) return wrong(module, .nvfp4, try r.ck.get(try r.name(module, "weight")), n, k); + const l = try core.nvfp4.Linear.read(r.ck, module); + if (l.n != n or l.k != k) return wrong(module, .nvfp4, try r.ck.get(try r.name(module, "weight")), n, k); + return l; + } + + fn fp4(r: *Reader, module: []const u8, n: usize, k: usize) !Proj { + const l = try r.fp4Read(module, n, k); + var codes = try r.staged(l.codes); + defer codes.free(); + var scales = try r.staged(l.scales); + defer scales.free(); + const np = cuda.nvfp4.padded(n); + const words = try r.store.alloc(np / 64 * (k / 64) * 512 * 4, module); + const tiles = try r.store.alloc(np * (k / 16), try r.name(module, "weight_scale")); + try r.pack.fp4Lane(r.stream, codes.ptr, scales.ptr, words, tiles, n, k); + try r.stream.synchronize(); + return .{ .nvfp4 = cuda.nvfp4.weight(words, tiles, l.global, n, k) }; + } + + /// Gate, up [width, dims] and down [dims, width] under `prefix` (`prefix`gate_proj and so on). + pub fn mlp(r: *Reader, prefix: []const u8, width: usize, dims: usize) !Mlp { + var m: [3]Proj = undefined; + inline for (.{ "gate_proj", "up_proj", "down_proj" }, 0..) |p, j| { + var b: [192]u8 = undefined; + const module = try std.fmt.bufPrint(&b, "{s}{s}", .{ prefix, p }); + m[j] = if (j < 2) try r.linear(module, width, dims) else try r.linear(module, dims, width); + } + return .{ .gate = m[0], .up = m[1], .down = m[2] }; + } + + /// Small tensor `tensor` (bf16 or fp32, `numel` long) as `to`, uploaded under its own name. + pub fn small(r: *Reader, tensor: []const u8, numel: usize, to: core.checkpoint.DType) !u64 { + const t = try r.ck.get(tensor); + if (t.numel() != numel or (t.dtype != .bf16 and t.dtype != .f32)) { + std.log.err("{s}: {t} {any}, expected {d} bf16 or fp32 values", .{ tensor, t.dtype, t.shape[0..t.rank], numel }); + return error.UnexpectedTensor; + } + var host: std.ArrayList(u8) = .empty; + defer host.deinit(r.gpa); + try append(r.gpa, &host, t, to); + return r.store.keep(host.items, tensor); + } + + /// The three short-conv weights [n, 1, taps] as one fp32 [3n][taps]: #804 KDA's q|k|v channels. + pub fn convs(r: *Reader, prefix: []const u8, n: usize, out_name: []const u8) !u64 { + var host: std.ArrayList(u8) = .empty; + defer host.deinit(r.gpa); + inline for (.{ "q", "k", "v" }) |x| { + var b: [192]u8 = undefined; + const t = try r.ck.get(try std.fmt.bufPrint(&b, "{s}" ++ x ++ "_conv1d.weight", .{prefix})); + if (!t.is(t.dtype, &.{ n, 1, r.c.conv }) or (t.dtype != .bf16 and t.dtype != .f32)) { + std.log.err("{s}" ++ x ++ "_conv1d: {t} {any}, expected [{d}, 1, {d}]", .{ prefix, t.dtype, t.shape[0..t.rank], n, r.c.conv }); + return error.UnexpectedTensor; + } + try append(r.gpa, &host, t, .f32); + } + return r.store.keep(host.items, out_name); + } + + /// Layer i's routed experts 0..count [width, dims]: every projection of every expert in one format, or refused. + pub fn experts(r: *Reader, i: u32, count: usize, width: usize, dims: usize, limit: f32) !Routed { + const names = @import("names.zig"); + var b: [192]u8 = undefined; + const f = r.c.format(try names.moduleName(&b, i, .expert_gate, 0)); + for (0..count) |e| for ([_]names.Role{ .expert_gate, .expert_up, .expert_down }) |role| { + const g = r.c.format(try names.moduleName(&b, i, role, @intCast(e))); + if (g != f) { + std.log.err("{s} is {t}, expert 0's gate {t}: the grouped kernels take one format a layer", .{ try names.moduleName(&b, i, role, @intCast(e)), g, f }); + return error.MixedExperts; + } + }; + var prefix: [96]u8 = undefined; + switch (f) { + .bf16 => { + const lanes = try r.gpa.alloc(Mlp, count); + errdefer r.gpa.free(lanes); + for (lanes, 0..) |*m, e| m.* = try r.mlp(try std.fmt.bufPrint(&prefix, "model.layers.{d}.mlp.experts.{d}.", .{ i, e }), width, dims); + try r.store.lanes.append(r.gpa, lanes); + return .{ .lanes = lanes }; + }, + .fp8_block128 => return .{ .grouped = try r.fp8Experts(i, count, width, dims, limit) }, + .nvfp4 => return .{ .grouped = try r.fp4Experts(i, count, width, dims, limit) }, + } + } + + fn expertModule(b: []u8, i: u32, e: usize, comptime proj: []const u8) ![]const u8 { + return std.fmt.bufPrint(b, "model.layers.{d}.mlp.experts.{d}." ++ proj, .{ i, e }); + } + + /// A view of `ptr` the store owns, for uploads at an offset. + fn view(r: *Reader, ptr: u64, len: usize) cuda.DeviceBuffer { + return .{ .d = r.store.d, .ptr = ptr, .len = len }; + } + + fn label(b: []u8, i: u32, what: []const u8) ![]const u8 { + return std.fmt.bufPrint(b, "model.layers.{d}.mlp.experts.{s}", .{ i, what }); + } + + /// experts.Layer's fp8g form: Fp8Experts.packGateUp and packOne per expert, scale_inv blocks as stored. + fn fp8Experts(r: *Reader, i: u32, count: usize, ni: usize, d: usize, limit: f32) !cuda.experts.Layer { + if (ni % X8.block != 0 or d % X8.block != 0) return error.UnexpectedTensor; + const up_n = 2 * ni * d / 4; // u32 words of one expert's gate and up + const down_n = d * ni / 4; + const inv_n = (ni / 128) * (d / 128); // one matrix's scale_inv blocks + var lb: [96]u8 = undefined; + const up = try r.store.alloc(count * up_n * 4, try label(&lb, i, "up")); + const down = try r.store.alloc(count * down_n * 4, try label(&lb, i, "down")); + const up_scale = try r.store.alloc(count * 2 * inv_n * 4, try label(&lb, i, "up_scale")); + const down_scale = try r.store.alloc(count * inv_n * 4, try label(&lb, i, "down_scale")); + const hu = try r.gpa.alloc(u32, up_n); + defer r.gpa.free(hu); + const hd = try r.gpa.alloc(u32, down_n); + defer r.gpa.free(hd); + var b: [192]u8 = undefined; + for (0..count) |e| { + const g = try r.fp8Pair(try expertModule(&b, i, e, "gate_proj"), ni, d); + const u = try r.fp8Pair(try expertModule(&b, i, e, "up_proj"), ni, d); + const dn = try r.fp8Pair(try expertModule(&b, i, e, "down_proj"), d, ni); + try X8.packGateUp(r.gpa, hu, g.w.bytes, u.w.bytes, ni, d); + X8.packOne(hd, dn.w.bytes, d, ni); + try r.view(up, count * up_n * 4).upload(e * up_n * 4, std.mem.sliceAsBytes(hu)); + try r.view(down, count * down_n * 4).upload(e * down_n * 4, std.mem.sliceAsBytes(hd)); + const us = r.view(up_scale, count * 2 * inv_n * 4); + try us.upload(2 * e * inv_n * 4, g.s.bytes); + try us.upload((2 * e + 1) * inv_n * 4, u.s.bytes); + try r.view(down_scale, count * inv_n * 4).upload(e * inv_n * 4, dn.s.bytes); + } + return .{ .format = .fp8g, .up = up, .down = down, .up_scale = up_scale, .down_scale = down_scale, .width = ni, .dims = d, .experts = count, .limit = limit }; + } + + /// experts.Layer's nvfp4 form: Nvfp4Experts.packGateUp and packOne per expert, weight_scale_2 as the globals. + fn fp4Experts(r: *Reader, i: u32, count: usize, ni: usize, d: usize, limit: f32) !cuda.experts.Layer { + if (ni % X4.cols != 0 or d % 32 != 0) return error.UnexpectedTensor; + const up_n = X4.bytes(1, ni, d, 2) / 4; + const down_n = X4.bytes(1, d, ni, 1) / 4; + var lb: [96]u8 = undefined; + const up = try r.store.alloc(count * up_n * 4, try label(&lb, i, "up")); + const down = try r.store.alloc(count * down_n * 4, try label(&lb, i, "down")); + const hu = try r.gpa.alloc(u32, up_n); + defer r.gpa.free(hu); + const hd = try r.gpa.alloc(u32, down_n); + defer r.gpa.free(hd); + const gu = try r.gpa.alloc(f32, 2 * count); + defer r.gpa.free(gu); + const gd = try r.gpa.alloc(f32, count); + defer r.gpa.free(gd); + var b: [192]u8 = undefined; + for (0..count) |e| { + const g = try r.fp4Read(try expertModule(&b, i, e, "gate_proj"), ni, d); + const u = try r.fp4Read(try expertModule(&b, i, e, "up_proj"), ni, d); + const dn = try r.fp4Read(try expertModule(&b, i, e, "down_proj"), d, ni); + X4.packGateUp(hu, .{ g.codes, g.scales }, .{ u.codes, u.scales }, ni, d); + X4.packOne(hd, dn.codes, dn.scales, d, ni); + try r.view(up, count * up_n * 4).upload(e * up_n * 4, std.mem.sliceAsBytes(hu)); + try r.view(down, count * down_n * 4).upload(e * down_n * 4, std.mem.sliceAsBytes(hd)); + gu[2 * e] = g.global; + gu[2 * e + 1] = u.global; + gd[e] = dn.global; + } + const up_scale = try r.store.keep(std.mem.sliceAsBytes(gu), try label(&lb, i, "up_scale")); + const down_scale = try r.store.keep(std.mem.sliceAsBytes(gd), try label(&lb, i, "down_scale")); + return .{ .format = .nvfp4, .up = up, .down = down, .up_scale = up_scale, .down_scale = down_scale, .width = ni, .dims = d, .experts = count, .limit = limit }; + } +}; + +/// `t`'s elements appended to `host` as bf16 or fp32 (bf16 from fp32 rounds to nearest even). +pub fn append(gpa: std.mem.Allocator, host: *std.ArrayList(u8), t: Tensor, to: core.checkpoint.DType) !void { + if (t.dtype == to) return host.appendSlice(gpa, t.bytes); + for (0..t.numel()) |i| { + const v = try value(t, i); + switch (to) { + .f32 => try host.appendSlice(gpa, std.mem.asBytes(&v)), + .bf16 => try host.appendSlice(gpa, std.mem.asBytes(&toBf16(v))), + else => return error.UnexpectedTensor, + } + } +} diff --git a/zig/tests/cuda/ling_synth.zig b/zig/tests/cuda/ling_synth.zig new file mode 100644 index 0000000000..eef50607df --- /dev/null +++ b/zig/tests/cuda/ling_synth.zig @@ -0,0 +1,154 @@ +//! `ling-synth DIR [--format bf16|fp8|nvfp4] [--layers N]`: a tiny Ling-3.1 checkpoint, every name from the weight map. +const std = @import("std"); +const core = @import("core"); +const ling = @import("ling_cuda"); +const check = @import("check.zig"); + +const st = core.safetensors; +const New = core.shard_edit.New; +const Config = ling.config.Config; +const Quant = ling.config.Quant; +const Role = ling.names.Role; + +/// A linear role's stored [n, k], from the config: the writer's own table, apart from the loader's. +fn linear(c: *const Config, role: Role) ?[2]usize { + const d: usize = c.hidden; + const n: usize = c.kda_heads * c.kda_dim; + const h: usize = c.mla_heads; + return switch (role) { + .kda_q, .kda_k, .kda_v, .kda_f, .kda_g => .{ n, d }, + .kda_b => .{ c.kda_heads, d }, + .kda_o => .{ d, n }, + .mla_q => .{ h * c.qkDim(), d }, + .mla_q_a => .{ c.q_lora, d }, + .mla_q_b => .{ h * c.qkDim(), c.q_lora }, + .mla_kv_a => .{ c.kv_lora + c.rope_dim, d }, + .mla_kv_b => .{ h * (c.nope + c.v_dim), c.kv_lora }, + .mla_gate => .{ h, d }, + .mla_o => .{ d, h * c.v_dim }, + .dense_gate, .dense_up => .{ c.dense_inter, d }, + .dense_down => .{ d, c.dense_inter }, + .router => .{ c.experts, d }, + .expert_gate, .expert_up => .{ c.moe_inter, d }, + .expert_down => .{ d, c.moe_inter }, + .shared_gate, .shared_up => .{ c.shared_inter, d }, + .shared_down => .{ d, c.shared_inter }, + .mtp_eh => .{ d, 2 * d }, + .lm_head, .embed => .{ c.vocab, d }, + else => null, + }; +} + +const Out = struct { + a: std.mem.Allocator, + rng: std.Random, + c: *const Config, + list: std.ArrayList(New) = .empty, + bytes: usize = 0, + + fn add(o: *Out, name: []const u8, dtype: st.DType, shape: []const usize, bytes: []const u8) !void { + try o.list.append(o.a, .{ .name = try o.a.dupe(u8, name), .dtype = dtype, .shape = try o.a.dupe(usize, shape), .bytes = bytes }); + o.bytes += bytes.len; + } + + fn bf16s(o: *Out, count: usize, lo: f32, hi: f32) ![]u8 { + const v = try o.a.alloc(u16, count); + for (v) |*x| x.* = ling.weights.proj.toBf16(lo + o.rng.float(f32) * (hi - lo)); + return std.mem.sliceAsBytes(v); + } + + fn f32s(o: *Out, count: usize, lo: f32, hi: f32) ![]u8 { + const v = try o.a.alloc(f32, count); + for (v) |*x| x.* = lo + o.rng.float(f32) * (hi - lo); + return std.mem.sliceAsBytes(v); + } + + /// One name the map emits: a weight in its module's format, a side tensor by its suffix, or a small tensor by role. + pub fn put(o: *Out, name: []const u8, t: ling.names.Tensor) !void { + const c = o.c; + if (linear(c, t.role)) |nk| { + const n, const k = nk; + const module = name[0 .. std.mem.lastIndexOfScalar(u8, name, '.') orelse return error.BadName]; + const f = if (t.role == .embed) .bf16 else c.format(module); + const s = 1.0 / @sqrt(@as(f32, @floatFromInt(k))); + if (!t.scale) return switch (f) { + .bf16 => o.add(name, .bf16, &.{ n, k }, try o.bf16s(n * k, -s, s)), + .fp8_block128 => o.add(name, .f8_e4m3, &.{ n, k }, try o.codes(n * k, .e4m3)), + .nvfp4 => o.add(name, .u8, &.{ n, k / 2 }, try o.codes(n * k / 2, .any)), + }; + const tail = name[module.len + 1 ..]; + if (std.mem.eql(u8, tail, "weight_scale_inv")) return o.add(name, .f32, &.{ (n + 127) / 128, k / 128 }, try o.f32s((n + 127) / 128 * (k / 128), 0.001, 0.011)); + if (std.mem.eql(u8, tail, "weight_scale")) return o.add(name, .f8_e4m3, &.{ n, k / 16 }, try o.codes(n * k / 16, .scale)); + if (std.mem.eql(u8, tail, "weight_scale_2")) return o.add(name, .f32, &.{}, try o.f32s(1, s / 8, s / 4)); + if (std.mem.eql(u8, tail, "input_scale")) return o.add(name, .f32, &.{}, try o.f32s(1, 1, 1)); + return error.BadName; + } + const d: usize = c.hidden; + const n: usize = c.kda_heads * c.kda_dim; + return switch (t.role) { + .attn_norm, .mlp_norm, .final_norm, .mtp_enorm, .mtp_hnorm, .mtp_final_norm => o.add(name, .bf16, &.{d}, try o.bf16s(d, 0.5, 1.5)), + .kda_q_conv, .kda_k_conv, .kda_v_conv => o.add(name, .bf16, &.{ n, 1, c.conv }, try o.bf16s(n * c.conv, -0.5, 0.5)), + .kda_a_log => o.add(name, .f32, &.{c.kda_heads}, try o.f32s(c.kda_heads, 0, 2.5)), + .kda_dt_bias => o.add(name, .f32, &.{n}, try o.f32s(n, -0.5, 0.5)), + .kda_o_norm => o.add(name, .bf16, &.{c.kda_dim}, try o.bf16s(c.kda_dim, 0.5, 1.5)), + .mla_q_a_norm => o.add(name, .bf16, &.{c.q_lora}, try o.bf16s(c.q_lora, 0.5, 1.5)), + .mla_kv_a_norm => o.add(name, .bf16, &.{c.kv_lora}, try o.bf16s(c.kv_lora, 0.5, 1.5)), + .router_bias => o.add(name, .f32, &.{c.experts}, try o.f32s(c.experts, -0.1, 0.1)), + else => error.BadName, + }; + } + + /// Random bytes: any; e4m3 codes without NaN; positive e4m3 scales 2^-3 to 2^1. + fn codes(o: *Out, count: usize, what: enum { any, e4m3, scale }) ![]u8 { + const v = try o.a.alloc(u8, count); + for (v) |*b| b.* = switch (what) { + .any => o.rng.int(u8), + .e4m3 => blk: { + const x = o.rng.int(u8); + break :blk if (x & 0x7f == 0x7f) x ^ 1 else x; + }, + .scale => 0x20 + o.rng.uintLessThan(u8, 0x28), + }; + return v; + } +}; + +fn formatOf(text: []const u8) !Quant { + if (std.mem.eql(u8, text, "bf16")) return .bf16; + if (std.mem.eql(u8, text, "fp8")) return .fp8_block128; + if (std.mem.eql(u8, text, "nvfp4")) return .nvfp4; + std.debug.print("--format takes bf16, fp8 or nvfp4, not {s}\n", .{text}); + return error.BadArgument; +} + +pub fn run(gpu: check.Gpu, rest: []const [:0]const u8) !void { + if (rest.len == 0) return error.MissingArgument; + var format: Quant = .fp8_block128; + var layers: u32 = 4; + var i: usize = 1; + while (i + 1 < rest.len) : (i += 2) { + if (std.mem.eql(u8, rest[i], "--format")) format = try formatOf(rest[i + 1]) else if (std.mem.eql(u8, rest[i], "--layers")) layers = try std.fmt.parseInt(u32, rest[i + 1], 10) else return error.BadArgument; + } + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const json = try ling.synth.configJson(a, format, layers); + var why: ling.config.Why = .{}; + const c = ling.config.parse(a, json, &why) catch |e| { + std.debug.print("the tiny config does not parse: {s}\n", .{why.text()}); + return e; + }; + var prng = std.Random.DefaultPrng.init(0x11a6_3100 + @as(u64, @backingInt(format))); + var o: Out = .{ .a = a, .rng = prng.random(), .c = &c }; + try ling.names.all(&c, &o); + const dir = rest[0]; + try std.Io.Dir.cwd().createDirPath(gpu.io, dir); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = try std.fs.path.join(a, &.{ dir, "config.json" }), .data = json }); + try core.shard_edit.create(gpu.gpa, try std.fs.path.join(a, &.{ dir, "model.safetensors" }), &.{.{ "format", "pt" }}, o.list.items); + // the file's names against the map: every name expected, nothing extra + const all = try a.alloc([]const u8, o.list.items.len); + for (o.list.items, all) |t, *n| n.* = t.name; + const r = try ling.names.check(a, &c, all); + try check.expect(r.missing.items.len == 0 and r.unknown.items.len == 0, "ling-synth names: {d} missing, {d} unaccounted", .{ r.missing.items.len, r.unknown.items.len }); + check.pass("ling-synth {t}: {d} layers + MTP, {d} tensors, {d} KiB in {s}", .{ format, layers, o.list.items.len, o.bytes >> 10, dir }); +} diff --git a/zig/tests/cuda/ling_weights_real.zig b/zig/tests/cuda/ling_weights_real.zig new file mode 100644 index 0000000000..1110020206 --- /dev/null +++ b/zig/tests/cuda/ling_weights_real.zig @@ -0,0 +1,201 @@ +//! `ling-weights-real DIR REFDIR`: real FP8 layers 0, 4, 7 (experts 0-3) vs float64 and torch (oracle/ling_weights.py). +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const ling = @import("ling_cuda"); +const check = @import("check.zig"); +const qr = @import("quant_ref_tests.zig"); +const lw = @import("ling_weights_tests.zig"); + +const Rig = qr.RigT; +const h = qr.helpers; +const W = ling.weights; +const rows = 8; + +/// Torch's rows for a module: x [8][k] bf16, y and sum |x w| [8][n] float64. +const Ref = struct { x: []const u16, y: []const f64, abs: []const f64 }; + +fn readRef(r: *Rig, dir: []const u8, module: []const u8, k: usize, n: usize) !Ref { + const io = r.gpu.io; + var b: [256]u8 = undefined; + const x = try std.Io.Dir.cwd().readFileAlloc(io, try std.fmt.bufPrint(&b, "{s}/x{d}.bin", .{ dir, k }), r.a, .unlimited); + const y = try std.Io.Dir.cwd().readFileAlloc(io, try std.fmt.bufPrint(&b, "{s}/{s}.y", .{ dir, module }), r.a, .unlimited); + const abs = try std.Io.Dir.cwd().readFileAlloc(io, try std.fmt.bufPrint(&b, "{s}/{s}.abs", .{ dir, module }), r.a, .unlimited); + if (x.len != rows * k * 2 or y.len != rows * n * 8 or abs.len != y.len) return error.BadReference; + const xs = try r.a.alloc(u16, rows * k); + @memcpy(std.mem.sliceAsBytes(xs), x); + const ys = try r.a.alloc(f64, rows * n); + @memcpy(std.mem.sliceAsBytes(ys), y); + const as = try r.a.alloc(f64, rows * n); + @memcpy(std.mem.sliceAsBytes(as), abs); + return .{ .x = xs, .y = ys, .abs = as }; +} + +/// Agreement with torch: the largest difference as a fraction of the bound, and the relative L2 error. +const Stat = struct { + worst: f64 = 0, + rel: f64 = 0, // the largest relative L2 over the checks + n: usize = 0, + + fn add(s: *Stat, name: []const u8, got: []const f64, want: []const f64, tol: []const f64) !void { + var most: f64 = 0; + var e2: f64 = 0; + var y2: f64 = 0; + for (got, want, tol) |g, w, t| { + most = @max(most, @abs(g - w) / t); + e2 += (g - w) * (g - w); + y2 += w * w; + } + const rel = @sqrt(e2 / @max(y2, 1e-300)); + std.debug.print("RESULT torch {s}: {d:.3} of its bound, rel L2 {e:.3}\n", .{ name, most, rel }); + s.worst = @max(s.worst, most); + s.rel = @max(s.rel, rel); + s.n += 1; + try check.expect(most <= 1, "{s} within its bound against torch", .{name}); + } +}; + +/// The bf16 (or fp32) output at `ptr` as float64. +fn values(r: *Rig, ptr: u64, count: usize, f32_out: bool) ![]f64 { + const raw = try r.back(ptr, count * @as(usize, if (f32_out) 4 else 2)); + const out = try r.a.alloc(f64, count); + for (out, 0..) |*v, i| v.* = if (f32_out) std.mem.bytesToValue(f32, raw[4 * i ..][0..4]) else h.bfOf(std.mem.bytesToValue(u16, raw[2 * i ..][0..2])); + return out; +} + +/// The projection named `module` among the loaded layers, or null (routed experts and the router are not projections). +fn find(w: *const W.Weights, module: []const u8) ?W.Proj { + var b: [192]u8 = undefined; + for (w.layers) |*l| { + var list: [16]lw.Item = undefined; + for (list[0..lw.projections(l, &list)]) |e| if (std.mem.eql(u8, ling.names.moduleName(&b, l.index, e[0], null) catch continue, module)) return e[1]; + } + return null; +} + +fn bound(ref: Ref, i: usize, fp32: bool) f64 { + return (if (fp32) 0 else @abs(ref.y[i]) * std.math.pow(f64, 2, -7)) + ref.abs[i] * std.math.pow(f64, 2, -12) + 1e-30; +} + +/// A projection's 8 decode rows against torch's. +fn projection(r: *Rig, kit: *const lw.Kit, st: *Stat, dir: []const u8, module: []const u8, p: W.Proj) !void { + const n = p.outputs(); + const k = p.inputs(); + const ref = try readRef(r, dir, module, k, n); + const out = try r.zeros(rows * n * 2); + try kit.linear().decode(r.s, p, .{ .x = try r.dev(std.mem.sliceAsBytes(ref.x)), .ldx = k, .part = try r.zeros(8 * rows * n * 4) }, out, rows); + const tol = try r.a.alloc(f64, rows * n); + for (tol, 0..) |*t, i| t.* = bound(ref, i, false); + try st.add(module, try values(r, out, rows * n, false), ref.y, tol); +} + +/// The router as loaded (bf16 [experts][hidden]) times torch's rows in float64, against torch's. +fn router(r: *Rig, st: *Stat, dir: []const u8, module: []const u8, l: *const W.Layer, c: *const ling.config.Config) !void { + const n: usize = c.experts; + const k: usize = c.hidden; + const ref = try readRef(r, dir, module, k, n); + const w = try r.back(l.ffn.moe.router, n * k * 2); + const got = try r.a.alloc(f64, rows * n); + for (0..rows) |row| for (0..n) |j| { + var s: f64 = 0; + for (0..k) |i| s += h.bfOf(ref.x[row * k + i]) * h.bfOf(std.mem.bytesToValue(u16, w[2 * (j * k + i) ..][0..2])); + got[row * n + j] = s; + }; + const tol = try r.a.alloc(f64, rows * n); + for (tol, ref.abs) |*t, a| t.* = a * std.math.pow(f64, 2, -40) + 1e-30; // float64 sums both sides + try st.add(module, got, ref.y, tol); +} + +/// Experts 0-3 of layer i: gate-up on torch's rows 0-1 of x (SwiGLU of torch's gate and up), down on its 8 pair rows. +fn expertsVsTorch(r: *Rig, kit: *const lw.Kit, st: *Stat, dir: []const u8, i: u32, l: cuda.experts.Layer) !void { + const slots = 4; + const pairs = 2 * slots; + const ni = l.width; + const d = l.dims; + const pk = try r.a.alloc(i32, pairs); + for (pk, 0..) |*e, p| e.* = lw.pick(p, slots, l.experts); + const items = cuda.grouped.maxItems(pairs, l.experts, 16); + const plan: cuda.grouped.Plan = .{ .members = try r.zeros(pairs * 4), .items = try r.zeros(items * 12), .counts = try r.zeros(8), .rank = try r.zeros(4), .hist = try r.zeros(4) }; + try kit.router.route(r.s, try r.dev(std.mem.sliceAsBytes(pk)), pairs, l.experts, 16, plan); + var nolimit = l; + nolimit.limit = 0; + var b: [192]u8 = undefined; + const want = try r.a.alloc(f64, pairs * ni); + const tol = try r.a.alloc(f64, pairs * ni); + const dwant = try r.a.alloc(f64, pairs * d); + const dtol = try r.a.alloc(f64, pairs * d); + var x: []const u16 = &.{}; + var xd: []const u16 = &.{}; + for (0..pairs) |p| { + const e: u32 = @intCast(pk[p]); + const g = try readRef(r, dir, try ling.names.moduleName(&b, i, .expert_gate, e), d, ni); + const u = try readRef(r, dir, try ling.names.moduleName(&b, i, .expert_up, e), d, ni); + const dn = try readRef(r, dir, try ling.names.moduleName(&b, i, .expert_down, e), ni, d); + x = g.x; + xd = dn.x; + const row = p / slots; + for (0..ni) |j| { + const gy = g.y[row * ni + j]; + want[p * ni + j] = lw.swiglu(gy, u.y[row * ni + j]); + const gt = bound(g, row * ni + j, false); + const ut = bound(u, row * ni + j, false); + tol[p * ni + j] = @abs(want[p * ni + j]) * std.math.pow(f64, 2, -6) + 1.1 * @abs(u.y[row * ni + j]) * gt + @abs(gy / (1 + @exp(-gy))) * ut + 1e-30; + } + for (0..d) |j| { + dwant[p * d + j] = dn.y[p * d + j]; + dtol[p * d + j] = bound(dn, p * d + j, true); + } + } + const act = try r.zeros(pairs * ni * 2); + try kit.grouped().decode(r.s, .up, .{ .x = try r.dev(std.mem.sliceAsBytes(x[0 .. 2 * d])), .stride = d, .slots = slots }, nolimit, plan, act, items, -1); + const y = try r.zeros(pairs * d * 4); + try kit.grouped().decode(r.s, .down_f32, .{ .x = try r.dev(std.mem.sliceAsBytes(xd)), .stride = ni }, nolimit, plan, y, items, -1); + var name: [96]u8 = undefined; + try st.add(try std.fmt.bufPrint(&name, "layer {d} experts 0-3 gate-up (SwiGLU)", .{i}), try values(r, act, pairs * ni, false), want, tol); + try st.add(try std.fmt.bufPrint(&name, "layer {d} experts 0-3 down", .{i}), try values(r, y, pairs * d, true), dwant, dtol); +} + +pub fn run(gpu: check.Gpu, dir: []const u8, refdir: []const u8) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(0x11a6_0047); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*x| x.free(); + r.bufs.deinit(gpu.gpa); + } + var kit: lw.Kit = undefined; + try kit.init(gpu); + defer kit.deinit(); + var why: ling.config.Why = .{}; + var o = W.open(gpu.gpa, gpu.io, dir, &why) catch |e| { + std.debug.print("{s}: {s}\n", .{ dir, why.text() }); + return e; + }; + defer o.ck.close(); + const t0 = check.now(gpu.io); + var w = try W.load(gpu.gpa, gpu.d, s, &kit.pack, &o.ck, o.c, &.{ 0, 4, 7 }, .{ .globals = false, .experts = 4 }); + defer w.deinit(); + const secs = @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e9; + var bytes: usize = 0; + for (w.store.named.items) |n| bytes += n.len; + std.debug.print("RESULT loaded layers 0, 4, 7 (4 experts each) in {d:.1} s: {d} MiB on the device\n", .{ secs, bytes >> 20 }); + const n = try lw.layers(&r, &kit, &o.ck, &w); + var st: Stat = .{}; + const manifest = try std.Io.Dir.cwd().readFileAlloc(gpu.io, try std.fmt.allocPrint(r.a, "{s}/manifest.txt", .{refdir}), r.a, .unlimited); + var lines = std.mem.tokenizeScalar(u8, manifest, '\n'); + while (lines.next()) |line| { + const module = line[0 .. std.mem.indexOfScalar(u8, line, ' ') orelse return error.BadReference]; + if (std.mem.indexOf(u8, module, ".experts.") != null) continue; + if (std.mem.endsWith(u8, module, ".mlp.gate")) { + const i = try std.fmt.parseInt(u32, module["model.layers.".len..std.mem.indexOfScalarPos(u8, module, "model.layers.".len, '.').?], 10); + try router(&r, &st, refdir, module, w.layer(i).?, &w.config); + continue; + } + try projection(&r, &kit, &st, refdir, module, find(&w, module) orelse return error.NotLoaded); + } + for ([_]u32{ 4, 7 }) |i| try expertsVsTorch(&r, &kit, &st, refdir, i, w.layer(i).?.ffn.moe.routed.grouped); + check.pass("ling-weights-real: {d} projections within float64 bounds from the file (worst {d:.3}), {d} small tensors byte-equal; {d} checks against torch, worst {d:.3} of the bound, rel L2 at most {e:.3}", .{ n.proj, r.worst, n.small, st.n, st.worst, st.rel }); +} diff --git a/zig/tests/cuda/ling_weights_tests.zig b/zig/tests/cuda/ling_weights_tests.zig new file mode 100644 index 0000000000..185ecb6bac --- /dev/null +++ b/zig/tests/cuda/ling_weights_tests.zig @@ -0,0 +1,340 @@ +//! The Ling loader against its checkpoint: each projection's y = W x vs float64 from the file, small tensors as bytes. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const ling = @import("ling_cuda"); +const check = @import("check.zig"); +const qr = @import("quant_ref_tests.zig"); + +const Rig = qr.RigT; +const h = qr.helpers; +const W = ling.weights; +const Config = ling.config.Config; +const Quant = ling.config.Quant; +const Role = ling.names.Role; + +/// The kernels a check runs: lane matmul and prompt GEMM, the packers, both grouped expert formats and the plan. +pub const Kit = struct { + lane: cuda.qmmf.Lane, + gemm: cuda.nvfp4.Prompt, + pack: cuda.modelopt.Packer, + x8: cuda.experts.Fp8Experts, + x4: cuda.experts.Nvfp4Experts, + mod: cuda.Module, + router: cuda.grouped.Router, + + pub fn init(k: *Kit, gpu: check.Gpu) !void { + const sms: usize = @intCast(try gpu.ctx.attribute(.multiprocessor_count)); + k.lane = try cuda.qmmf.Lane.load(gpu.d, (try gpu.ctx.capability()) / 10); + k.gemm = try cuda.nvfp4.Prompt.load(gpu.d); + k.pack = try cuda.modelopt.Packer.load(gpu.d); + k.x8 = try cuda.experts.Fp8Experts.load(gpu.d, sms); + k.x4 = try cuda.experts.Nvfp4Experts.load(gpu.d, sms); + k.mod = try cuda.Module.load(gpu.d, cuda.kernels.experts); + k.router = try cuda.grouped.Router.resolve(k.mod); + } + + pub fn deinit(k: *Kit) void { + k.lane.unload(); + k.gemm.unload(); + k.pack.unload(); + k.x8.unload(); + k.x4.unload(); + k.mod.unload(); + } + + pub fn linear(k: *const Kit) cuda.qlinear.Linear { + return .{ .lane = &k.lane, .nvfp4_prompt = &k.gemm }; + } + + pub fn grouped(k: *const Kit) cuda.experts.Grouped { + return .{ .fp8g = &k.x8, .nvfp4 = &k.x4 }; + } +}; + +/// A linear as its file stores it, read back to float64 by (row, col): the host's dequantisation. +pub const Mat = struct { + f: Quant, + n: usize, + k: usize, + w: []const u8, + s: []const u8 = &.{}, + global: f64 = 1, + + pub fn read(ck: *core.Checkpoint, c: *const Config, module: []const u8, n: usize, k: usize) !Mat { + var b: [256]u8 = undefined; + const f = c.format(module); + const w = try ck.get(try std.fmt.bufPrint(&b, "{s}.weight", .{module})); + return switch (f) { + .bf16 => .{ .f = f, .n = n, .k = k, .w = w.bytes }, + .fp8_block128 => .{ .f = f, .n = n, .k = k, .w = w.bytes, .s = (try ck.get(try std.fmt.bufPrint(&b, "{s}.weight_scale_inv", .{module}))).bytes }, + .nvfp4 => blk: { + const l = try core.nvfp4.Linear.read(ck, module); + break :blk .{ .f = f, .n = n, .k = k, .w = l.codes, .s = l.scales, .global = l.global }; + }, + }; + } + + pub fn at(m: Mat, row: usize, col: usize) f64 { + return switch (m.f) { + .bf16 => h.bfOf(std.mem.readInt(u16, m.w[2 * (row * m.k + col) ..][0..2], .little)), + .fp8_block128 => h.e4m3Of(m.w[row * m.k + col]) * @as(f64, @as(f32, @bitCast(std.mem.readInt(u32, m.s[4 * ((row / 128) * (m.k / 128) + col / 128) ..][0..4], .little)))), + .nvfp4 => blk: { + const byte = m.w[row * (m.k / 2) + col / 2]; + break :blk h.e2m1Of(if (col & 1 == 0) byte & 0xF else byte >> 4) * h.e4m3Of(m.s[row * (m.k / 16) + col / 16]) * m.global; + }, + }; + } +}; + +/// y = x W^T in float64 for `rows` bf16 rows, and the bound: 2^-7 of |y| (bf16 out) and 2^-12 of sum |x w| (fp32 sums). +pub fn reference(a: std.mem.Allocator, m: Mat, x: []const u16, rows: usize) !struct { want: []f64, tol: []f64 } { + const want = try a.alloc(f64, rows * m.n); + const tol = try a.alloc(f64, rows * m.n); + @memset(want, 0); + @memset(tol, 0); + for (0..m.n) |j| for (0..m.k) |i| { + const v = m.at(j, i); + for (0..rows) |r| { + const t = h.bfOf(x[r * m.k + i]) * v; + want[r * m.n + j] += t; + tol[r * m.n + j] += @abs(t); + } + }; + for (want, tol) |y, *t| t.* = @abs(y) * std.math.pow(f64, 2, -7) + t.* * std.math.pow(f64, 2, -12) + 1e-30; + return .{ .want = want, .tol = tol }; +} + +/// The view's tag is the format the config gives the module. +fn tagged(w: W.Proj, f: Quant) bool { + return switch (w) { + .bf16 => f == .bf16, + .fp8g => f == .fp8_block128, + .nvfp4 => f == .nvfp4, + else => false, + }; +} + +/// One projection: its tag, then decode and prompt rows against the file's float64 product. +pub fn linear(r: *Rig, kit: *const Kit, ck: *core.Checkpoint, c: *const Config, module: []const u8, w: W.Proj, rows: usize) !void { + const n = w.outputs(); + const k = w.inputs(); + try check.expect(tagged(w, c.format(module)), "{s}: loaded as {s}, the config says {t}", .{ module, @tagName(w), c.format(module) }); + const m = try Mat.read(ck, c, module, n, k); + const x = try r.bfs(rows * k, 1.0); + const ref = try reference(r.a, m, x, rows); + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const part = try r.zeros(8 * rows * n * 4); + var name: [256]u8 = undefined; + for ([_]bool{ false, true }) |prompt| { + const out = try r.zeros(rows * n * 2); + const in: cuda.qlinear.Rows = .{ .x = dx, .ldx = k, .part = part }; + if (prompt) try kit.linear().prompt(r.s, w, in, out, rows) else try kit.linear().decode(r.s, w, in, out, rows); + try r.close(try std.fmt.bufPrint(&name, "{s} {t} {s} [{d}, {d}]", .{ module, c.format(module), if (prompt) "prompt" else "decode", n, k }), out, false, ref.want, ref.tol); + } +} + +pub const Item = struct { Role, W.Proj }; + +fn put(out: []Item, at: *usize, role: Role, p: ?W.Proj) void { + const x = p orelse return; + out[at.*] = .{ role, x }; + at.* += 1; +} + +/// Every projection a layer holds, by role (routed experts apart). +pub fn projections(l: *const W.Layer, out: []Item) usize { + var n: usize = 0; + switch (l.mixer) { + .kda => |a| for ([_]Role{ .kda_q, .kda_k, .kda_v, .kda_b, .kda_f, .kda_g, .kda_o }, [_]W.Proj{ a.q, a.k, a.v, a.b, a.f, a.g, a.o }) |role, p| put(out, &n, role, p), + .mla => |a| for ([_]Role{ .mla_q, .mla_q_a, .mla_q_b, .mla_kv_a, .mla_kv_b, .mla_gate, .mla_o }, [_]?W.Proj{ a.q, a.q_a, a.q_b, a.kv_a, a.kv_b, a.gate, a.o }) |role, p| put(out, &n, role, p), + } + var roles = [3]Role{ .dense_gate, .dense_up, .dense_down }; + var m: ?W.Mlp = null; + switch (l.ffn) { + .dense => |d| m = d, + .moe => |e| { + m = e.shared; + roles = .{ .shared_gate, .shared_up, .shared_down }; + }, + } + if (m) |x| for (roles, [_]W.Proj{ x.gate, x.up, x.down }) |role, p| put(out, &n, role, p); + return n; +} + +/// The bytes at `ptr` equal the file's tensors (concatenated) as `to`. +fn sameAs(r: *Rig, ck: *core.Checkpoint, ptr: u64, tensors: []const []const u8, to: core.checkpoint.DType) !void { + var want: std.ArrayList(u8) = .empty; + for (tensors) |t| try W.proj.append(r.a, &want, try ck.get(t), to); + try check.sameBytes(tensors[0], try r.back(ptr, want.items.len), want.items); +} + +/// A layer's small tensors byte for byte: norms, the KDA convs (fp32 q|k|v), A_log, dt_bias, router and bias. +pub fn smalls(r: *Rig, ck: *core.Checkpoint, c: *const Config, l: *const W.Layer, w: *const W.Weights) !usize { + var b: [8][192]u8 = undefined; + const i = l.index; + const p = "model.layers."; + var n: usize = 2; + try sameAs(r, ck, l.attn_norm, &.{try std.fmt.bufPrint(&b[0], p ++ "{d}.input_layernorm.weight", .{i})}, .bf16); + try sameAs(r, ck, l.mlp_norm, &.{try std.fmt.bufPrint(&b[0], p ++ "{d}.post_attention_layernorm.weight", .{i})}, .bf16); + switch (l.mixer) { + .kda => |a| { + var convs: [3][]const u8 = undefined; + for (&convs, [_][]const u8{ "q", "k", "v" }, 0..) |*t, x, j| t.* = try std.fmt.bufPrint(&b[j], p ++ "{d}.attention.{s}_conv1d.weight", .{ i, x }); + try sameAs(r, ck, a.conv, &convs, .f32); + try sameAs(r, ck, a.a_log, &.{try std.fmt.bufPrint(&b[3], p ++ "{d}.attention.A_log", .{i})}, .f32); + try sameAs(r, ck, a.dt_bias, &.{try std.fmt.bufPrint(&b[3], p ++ "{d}.attention.dt_bias", .{i})}, .f32); + try sameAs(r, ck, a.o_norm, &.{try std.fmt.bufPrint(&b[3], p ++ "{d}.attention.o_norm.weight", .{i})}, .bf16); + n += 4; + }, + .mla => |a| { + try sameAs(r, ck, a.kv_a_norm, &.{try std.fmt.bufPrint(&b[0], p ++ "{d}.attention.kv_a_layernorm.weight", .{i})}, .bf16); + n += 1; + if (c.q_lora > 0) { + try sameAs(r, ck, a.q_a_norm, &.{try std.fmt.bufPrint(&b[0], p ++ "{d}.attention.q_a_layernorm.weight", .{i})}, .bf16); + n += 1; + } + }, + } + if (l.ffn == .moe) { + try sameAs(r, ck, l.ffn.moe.router, &.{try std.fmt.bufPrint(&b[0], p ++ "{d}.mlp.gate.weight", .{i})}, .bf16); + if (c.expert_bias) try sameAs(r, ck, l.ffn.moe.bias, &.{try std.fmt.bufPrint(&b[0], p ++ "{d}.mlp.gate.expert_bias", .{i})}, .f32); + n += 2; + } + if (i >= c.layers) if (w.mtp) |m| { + for ([_]u64{ m.enorm, m.hnorm, m.norm }, [_][]const u8{ "enorm", "hnorm", "final_layernorm" }) |ptr, x| try sameAs(r, ck, ptr, &.{try std.fmt.bufPrint(&b[0], p ++ "{d}.{s}.weight", .{ i, x })}, .bf16); + n += 3; + }; + return n; +} + +/// SwiGLU as the expert kernels round it: gate and up to bf16, silu to bf16, the product to bf16 (limit 0). +pub fn swiglu(g: f64, u: f64) f64 { + const gv = h.bfOf(h.bf16Of(@floatCast(g))); + const uv = h.bfOf(h.bf16Of(@floatCast(u))); + return h.bfOf(h.bf16Of(@floatCast(gv / (1 + @exp(-gv))))) * uv; +} + +/// Pair p's expert: row p / slots takes `slots` distinct experts, rotated a row. +pub fn pick(p: usize, slots: usize, count: usize) i32 { + return @intCast((p % slots + p / slots) % count); +} + +/// The grouped experts' gate-up (SwiGLU) and down on their own inputs, against float64 from the file. +pub fn experts(r: *Rig, kit: *const Kit, ck: *core.Checkpoint, c: *const Config, i: u32, l: cuda.experts.Layer, rows: usize) !void { + const slots = @min(4, l.experts); + const pairs = rows * slots; + const ni = l.width; + const d = l.dims; + const pk = try r.a.alloc(i32, pairs); + for (pk, 0..) |*e, p| e.* = pick(p, slots, l.experts); + const plan: cuda.grouped.Plan = .{ .members = try r.zeros(pairs * 4), .items = try r.zeros(cuda.grouped.maxItems(pairs, l.experts, 16) * 12), .counts = try r.zeros(8), .rank = try r.zeros(4), .hist = try r.zeros(4) }; + try kit.router.route(r.s, try r.dev(std.mem.sliceAsBytes(pk)), pairs, l.experts, 16, plan); + const items = cuda.grouped.maxItems(pairs, l.experts, 16); + try check.expect(l.limit == ling.limits(c, i).routed, "layer {d}: routed limit {d}, the config's {d}", .{ i, l.limit, ling.limits(c, i).routed }); + var nolimit = l; + nolimit.limit = 0; + const x = try r.bfs(rows * d, 1.0); + const act = try r.zeros(pairs * ni * 2); + try kit.grouped().decode(r.s, .up, .{ .x = try r.dev(std.mem.sliceAsBytes(x)), .stride = d, .slots = slots }, nolimit, plan, act, items, -1); + const xd = try r.bfs(pairs * ni, 1.0); + const y = try r.zeros(pairs * d * 4); + try kit.grouped().decode(r.s, .down_f32, .{ .x = try r.dev(std.mem.sliceAsBytes(xd)), .stride = ni }, nolimit, plan, y, items, -1); + var b: [3][192]u8 = undefined; + const want = try r.a.alloc(f64, pairs * ni); + const tol = try r.a.alloc(f64, pairs * ni); + const dwant = try r.a.alloc(f64, pairs * d); + const dtol = try r.a.alloc(f64, pairs * d); + for (0..pairs) |p| { + const e: u32 = @intCast(pk[p]); + const g = try Mat.read(ck, c, try ling.names.moduleName(&b[0], i, .expert_gate, e), ni, d); + const u = try Mat.read(ck, c, try ling.names.moduleName(&b[1], i, .expert_up, e), ni, d); + const dn = try Mat.read(ck, c, try ling.names.moduleName(&b[2], i, .expert_down, e), d, ni); + const gr = try reference(r.a, g, x[(p / slots) * d ..][0..d], 1); + const ur = try reference(r.a, u, x[(p / slots) * d ..][0..d], 1); + for (0..ni) |j| { + const g0 = gr.want[j]; + want[p * ni + j] = swiglu(g0, ur.want[j]); + // the bf16 steps (2^-6 of |a|), then each product's bound through silu(g) u: silu' stays under 1.1 + tol[p * ni + j] = @abs(want[p * ni + j]) * std.math.pow(f64, 2, -6) + 1.1 * @abs(ur.want[j]) * gr.tol[j] + @abs(g0 / (1 + @exp(-g0))) * ur.tol[j] + 1e-30; + } + const dr = try reference(r.a, dn, xd[p * ni ..][0..ni], 1); + for (0..d) |j| { + dwant[p * d + j] = dr.want[j]; + dtol[p * d + j] = dr.tol[j] - @abs(dr.want[j]) * std.math.pow(f64, 2, -7); // fp32 out: no bf16 step + } + } + var name: [128]u8 = undefined; + try r.close(try std.fmt.bufPrint(&name, "layer {d} experts {t} gate-up, {d} pairs of {d}", .{ i, l.format, pairs, l.experts }), act, false, want, tol); + try r.close(try std.fmt.bufPrint(&name, "layer {d} experts {t} down, {d} pairs", .{ i, l.format, pairs }), y, true, dwant, dtol); +} + +/// Every loaded layer's projections, experts and small tensors; the globals when loaded. +pub fn layers(r: *Rig, kit: *const Kit, ck: *core.Checkpoint, w: *const W.Weights) !struct { proj: usize, small: usize } { + const c = &w.config; + var np: usize = 0; + var ns: usize = 0; + var b: [192]u8 = undefined; + for (w.layers) |*l| { + var list: [16]Item = undefined; + for (list[0..projections(l, &list)]) |e| { + try linear(r, kit, ck, c, try ling.names.moduleName(&b, l.index, e[0], null), e[1], 3); + np += 1; + } + if (l.ffn == .moe) switch (l.ffn.moe.routed) { + .grouped => |g| { + try experts(r, kit, ck, c, l.index, g, 2); + np += 3; + }, + .lanes => |ls| for (ls[0..@min(ls.len, 3)], 0..) |m, e| { + for ([_]Role{ .expert_gate, .expert_up, .expert_down }, [_]W.Proj{ m.gate, m.up, m.down }) |role, p| try linear(r, kit, ck, c, try ling.names.moduleName(&b, l.index, role, @intCast(e)), p, 3); + np += 3; + }, + }; + if (l.index >= c.layers) if (w.mtp) |m| { + try linear(r, kit, ck, c, try ling.names.moduleName(&b, l.index, .mtp_eh, null), m.eh, 3); + np += 1; + }; + ns += try smalls(r, ck, c, l, w); + } + if (w.embed != 0) { + try sameAs(r, ck, w.embed, &.{"model.word_embeddings.weight"}, .bf16); + try sameAs(r, ck, w.norm, &.{"model.norm.weight"}, .bf16); + ns += 2; + if (w.head) |hd| { + try linear(r, kit, ck, c, "lm_head", hd, 3); + np += 1; + } + } + return .{ .proj = np, .small = ns }; +} + +/// `ling-weights DIR`: every layer of a (synthetic) checkpoint loaded and checked. +pub fn run(gpu: check.Gpu, dir: []const u8) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(0x11a6); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*x| x.free(); + r.bufs.deinit(gpu.gpa); + } + var kit: Kit = undefined; + try kit.init(gpu); + defer kit.deinit(); + var why: ling.config.Why = .{}; + var o = W.open(gpu.gpa, gpu.io, dir, &why) catch |e| { + std.debug.print("{s}: {s}\n", .{ dir, why.text() }); + return e; + }; + defer o.ck.close(); + const all = try r.a.alloc(u32, o.c.layers + o.c.mtp); + for (all, 0..) |*i, j| i.* = @intCast(j); + var w = try W.load(gpu.gpa, gpu.d, s, &kit.pack, &o.ck, o.c, all, .{}); + defer w.deinit(); + try check.expect(o.ck.unused() == 0, "{s}: every tensor read", .{dir}); + const n = try layers(&r, &kit, &o.ck, &w); + check.pass("ling-weights {t}: {d} layers + MTP, {d} projections within float64 bounds (worst {d:.3}), {d} small tensors byte-equal", .{ o.c.quant, o.c.layers, n.proj, r.worst, n.small }); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index b6fabfc13e..aaf69b0300 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -39,6 +39,9 @@ const radio_tests = @import("radio_tests.zig"); const glm_lanes_tests = @import("glm_lanes_tests.zig"); const glm_synth = @import("glm_synth.zig"); const glm_ref_route = @import("glm_ref_route.zig"); +const ling_synth = @import("ling_synth.zig"); +const ling_weights = @import("ling_weights_tests.zig"); +const ling_real = @import("ling_weights_real.zig"); const usage = \\usage: tf-cuda-test @@ -98,6 +101,9 @@ const usage = \\ glm-route-old GLM's select and swiglu against an older glm_glue fatbin (6380c8da), byte for byte \\ glm-route-ling Ling's group-limited routing on captured layer-4 logits vs float64, HF and exllamav3 \\ glm-synth a synthetic GLM-5.3-Flash checkpoint (2 layers, 16 experts) and its dc.json for the lane checks + \\ ling-synth [--format bf16|fp8|nvfp4] [--layers N] a tiny Ling-3.1 checkpoint with the real names and config schema + \\ ling-weights every layer of a Ling checkpoint loaded: projections vs float64, small tensors byte for byte + \\ ling-weights-real real Ling-3.1 layers 0, 4, 7 (4 experts) vs float64 and torch (oracle/ling_weights.py) \\ ; @@ -197,6 +203,9 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "glm-route-old")) return glm_ref_route.old(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "glm-route-ling")) return glm_ref_route.ling(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "glm-synth")) return glm_synth.run(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "ling-synth")) return ling_synth.run(gpu, rest); + if (std.mem.eql(u8, cmd, "ling-weights")) return ling_weights.run(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "ling-weights-real")) return ling_real.run(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "glm-lanes-concurrent")) return glm_lanes_tests.concurrentCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); std.debug.print("{s}", .{usage}); return error.UnknownCommand; diff --git a/zig/tests/cuda/oracle/ling_weights.py b/zig/tests/cuda/oracle/ling_weights.py new file mode 100644 index 0000000000..9ce13cad9b --- /dev/null +++ b/zig/tests/cuda/oracle/ling_weights.py @@ -0,0 +1,62 @@ +"""For `tf-cuda-test ling-weights-real`: torch's float64 y = x W^T of real Ling-3.1 projections on fixed rows.""" +import argparse +import json +import os + +import torch +from safetensors import safe_open + +ROWS = 8 # decode rows, and the down projection's pair rows (2 tokens x 4 experts) +EXPERTS = 4 +KDA = ("q_proj", "k_proj", "v_proj", "b_proj", "f_proj", "g_proj", "o_proj") +MLA = ("q_a_proj", "q_b_proj", "kv_a_proj_with_mqa", "kv_b_proj", "g_proj", "dense") +MLP = ("gate_proj", "up_proj", "down_proj") + + +def modules() -> list[str]: + """Layer 0 (KDA, dense MLP), 4 (KDA, MoE) and 7 (MLA, MoE): every projection, the router, experts 0-3.""" + out = [f"model.layers.0.attention.{m}" for m in KDA] + [f"model.layers.0.mlp.{m}" for m in MLP] + for layer, mixer in ((4, KDA), (7, MLA)): + out += [f"model.layers.{layer}.attention.{m}" for m in mixer] + [f"model.layers.{layer}.mlp.gate"] + out += [f"model.layers.{layer}.mlp.shared_experts.{m}" for m in MLP] + out += [f"model.layers.{layer}.mlp.experts.{e}.{m}" for e in range(EXPERTS) for m in MLP] + return out + + +def rows(k: int) -> torch.Tensor: + """The bf16 rows of width k both sides read: uniform in [-1, 1), seeded by k.""" + g = torch.Generator().manual_seed(k) + return (torch.rand((ROWS, k), generator=g) * 2 - 1).to(torch.bfloat16) + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("src", help="the FP8 checkpoint (model.safetensors.index.json and its shards)") + p.add_argument("out", help="where x.bin, .y, .abs and manifest.txt go") + a = p.parse_args() + os.makedirs(a.out, exist_ok=True) + weight_map = json.load(open(os.path.join(a.src, "model.safetensors.index.json")))["weight_map"] + files, xs, lines = {}, {}, [] + for m in modules(): + name = weight_map[m + ".weight"] + f = files.setdefault(name, safe_open(os.path.join(a.src, name), "pt")) + w = f.get_tensor(m + ".weight") + n, k = w.shape + if w.dtype == torch.float8_e4m3fn: + s = weight_map[m + ".weight_scale_inv"] + inv = files.setdefault(s, safe_open(os.path.join(a.src, s), "pt")).get_tensor(m + ".weight_scale_inv") + w = w.double() * inv.double().repeat_interleave(128, 0)[:n].repeat_interleave(128, 1)[:, :k] + w = w.double() + if k not in xs: + xs[k] = rows(k) + xs[k].view(torch.int16).numpy().tofile(os.path.join(a.out, f"x{k}.bin")) + x = xs[k].double() + (x @ w.T).numpy().tofile(os.path.join(a.out, m + ".y")) + (x.abs() @ w.abs().T).numpy().tofile(os.path.join(a.out, m + ".abs")) + lines.append(f"{m} {n} {k}") + print(m, n, k, flush=True) + open(os.path.join(a.out, "manifest.txt"), "w").write("\n".join(lines) + "\n") + + +if __name__ == "__main__": + main() From faf7af1cd4c5b1f3b81f94f5c87ccd10ae161c3a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 15 Oct 2026 10:01:40 +0200 Subject: [PATCH 04/34] docs: Ling-3.1-flash recipe: what exists, what is missing, testing without the checkpoint, memory per format --- docs/recipes/README.md | 1 + docs/recipes/ling-3.1.md | 122 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 123 insertions(+) create mode 100644 docs/recipes/ling-3.1.md diff --git a/docs/recipes/README.md b/docs/recipes/README.md index f3e82f73b4..fabe4435f6 100644 --- a/docs/recipes/README.md +++ b/docs/recipes/README.md @@ -15,6 +15,7 @@ the native engine, or the Python engine frozen at 0.6.6. | Gemma 4 26B-A4B | [MLX, fused one-row decode](gemma-4.md) | | DeepSeek-V4-Flash | [MLX on a 256 GB Mac, DSpark and MTP drafts](deepseek-v4-flash.md) | | Qwen3.6-35B-A3B | [MLX with MTP drafts, one-GPU CUDA](qwen3.6-moe.md) | +| Ling-3.1-flash | [Not served yet: CUDA weight loader, synthetic checkpoints](ling-3.1.md) | ## Capability floor diff --git a/docs/recipes/ling-3.1.md b/docs/recipes/ling-3.1.md new file mode 100644 index 0000000000..44938e3fde --- /dev/null +++ b/docs/recipes/ling-3.1.md @@ -0,0 +1,122 @@ +# Ling-3.1-flash + +Work in progress on the Zig engine's CUDA side: the model is **not served yet**. This page lists what exists, what is +missing, and how to test without the 572 GB checkpoint. Everything below was measured on one GB10. + +Ling-3.1-flash (`BailingMoeV3ForCausalLM`, `bailing_hybrid`) has 72 layers plus one MTP layer: Kimi delta attention +(KDA) on layers with `(i + 1) % 8 != 0`, MLA on every eighth, dense MLPs on layers 0-3 and 512 routed experts (top 8, +8 groups keeping 4, sigmoid scores with a pick bias) plus one shared expert from layer 4 on. The released checkpoint is +block FP8 (e4m3 in 128 x 128 blocks with `weight_scale_inv`). + +## What exists + +| Part | Where | Checked by | +| --- | --- | --- | +| Group-limited top-k routing, per-layer shared SwiGLU limit (GLM's bytes unchanged) | `ling-family`, commit 51c01573 | `glm-ref`, `glm-route-old`, `glm-route-ling` | +| MLA at Ling's shape: 64 roped key columns, 128-wide heads, GPT-J rope, head-wise gate | branch `ling-mla`, f8e594f6 | `ling-mla`, `ling-mla-real`, `glm-mla-same` | +| KDA on Ling's layer 0 (q/k/v convs as one q\|k\|v conv), opt-in chunked prompt path | branch `ling-kda`, b55aa13d | `glm-kda-ling`, `glm-kda-chunk` | +| Config loader and weight map (all 214,025 names of the real index) | `zig/src/families/ling/config.zig`, `names.zig` | `zig build test-ling` | +| Weight loader at TP=1: bf16, FP8 block, ModelOpt NVFP4 | `zig/src/families/ling/weights.zig`, `weights_proj.zig` | `ling-weights`, `ling-weights-real` | +| Synthetic checkpoints in all three formats | `tf-cuda-test ling-synth` | its own name check | + +The MLA and KDA commits live on their own branches and are not merged into `ling-family`. + +### Formats + +The format of every tensor comes from the checkpoint, never from the layer graph: + +- `quantization_config.quant_method: "fp8"` with `weight_block_size [128, 128]`: FP8 block. `modules_to_not_convert` + names the bf16 linears (a whole name, or a tail after a dot: `mlp.gate` is the router, not `gate_proj`). +- `quant_method: "modelopt"`, `quant_algo: "NVFP4"`, 4-bit weights in groups of 16, or the same block in + `hf_quant_config.json`: ModelOpt NVFP4. `ignore`, `exclude_modules` and `modules_to_not_convert` take names and + `*` globs. +- No `quantization_config`: bf16. +- `quant_method: "exl3"` is refused with a message naming tracker #832. The extension point is a variant in + `config.Quant`, its side tensors in `names.side()` and a reader in `weights_proj.Reader.linear()`. + +A file tensor whose dtype or shape disagrees with the format the config gives it is refused, naming the module and +both formats. Each projection loads as a `qlinear.Weight` (`bf16`, `fp8g` or `nvfp4`) through main's packers +(`modelopt.Packer`, `fp8.packScales`); routed experts load as `experts.Layer` (`fp8g` or `nvfp4`, packed by +`Fp8Experts`/`Nvfp4Experts`). There is no grouped bf16 expert kernel, so bf16 routed experts load as one lane weight +per projection. The router must be bf16 (GLM's router kernel reads bf16); a quantized router is refused. + +## What is missing + +- The layer graph: `weights.Layer` holds one `qlinear.Weight` per projection, while #804's GLM blocks read raw bf16 + pointers (KDA's fused `[q | k | v | f_a | g_a | b]`, MLA's latent projections). Ling's blocks have to call + `qlinear.Linear` per projection, or fuse same-format rows. +- Lanes hooks (drafted, concurrent and cut-prompt windows) for Ling. +- MTP drafting: the MTP layer's tensors load (`weights.Mtp`), nothing drafts with them. +- EXL3, through ashhart's EXL3 module (tracker #832). +- Tensor parallelism: the loader is TP=1; shares and transport would come from #804's `cuda_shard.zig` and #855. + +## Testing without the full model + +```bash +zig build test-ling # config, name map, tiny configs (host only) +TF_LING_INDEX=/path/to/model.safetensors.index.json zig build test-ling # adds the real index +zig build -Doptimize=ReleaseFast -Dnvcc=/usr/local/cuda/bin/nvcc +zig-out/bin/tf-cuda-test ling-synth /tmp/ling-fp8 --format fp8 # bf16 | fp8 | nvfp4, --layers N (default 4) +zig-out/bin/tf-cuda-test ling-weights /tmp/ling-fp8 +``` + +`ling-synth` writes `config.json` (the real file's schema) and `model.safetensors` with the real key names: hidden 512, +4 heads, 16 experts in 4 groups keeping 2, `q_lora_rank` 256, vocabulary 1,024, layer groups of 2 (KDA, MLA, KDA, +MLA), layer 0 dense, the last layer clamped, plus the MTP layer. KDA's head dim 128 and MLA's 512 latent, 64 rope and +128-wide heads stay as the kernels take them. The FP8 file carries the real exception list's pattern; the NVFP4 file +uses globs and leaves layer 0's dense MLP in bf16, so one checkpoint mixes formats. It checks its names against the +weight map before it passes. Sizes: bf16 282 tensors, 43 MiB; FP8 517, 23 MiB; NVFP4 978, 15 MiB. + +`ling-weights` loads every layer and checks every projection (decode and prompt rows, 3 rows) against float64 from +the file, the grouped experts' gate-up (SwiGLU) and down, and each small tensor (norms, convs as fp32 q|k|v, `A_log`, +`dt_bias`, router, bias) byte for byte. Measured: bf16 85 projections, FP8 and NVFP4 61 each, worst 0.456 of the bound +(2^-7 of |y| plus 2^-12 of the sum of |x w|); 37 small tensors equal. + +With the real checkpoint: + +```bash +python zig/tests/cuda/oracle/ling_weights.py /path/to/Ling-3.1-flash /tmp/ling-ref # needs torch, safetensors +zig-out/bin/tf-cuda-test ling-weights-real /path/to/Ling-3.1-flash /tmp/ling-ref +``` + +It loads layers 0, 4 and 7 with experts 0-3 (769 MiB on the device) and compares 35 products with torch's float64 +on the same 8 rows: worst 0.414 of the bound, relative L2 at most 1.7e-3 for bf16 outputs and 6.9e-8 for the +experts' fp32 down rows. + +## Memory + +Computed from `config.json` (572.1 B parameters; routed experts are 97.1 %). Quantized formats keep the FP8 release's +bf16 modules (embedding, head, routers, KDA `b_proj`, MLA `g_proj`, MTP `eh_proj`). The FP8 figure equals the +checkpoint's shard total. EXL3 assumes every other linear at the given rate. + +| Format | Weights | Sparks, 128 GB each | Sparks, 100 GB of weights each | +| --- | --- | --- | --- | +| bf16 | 1,144 GB | 9 | 12 | +| FP8 block (the release) | 574 GB | 5 | 6 | +| ModelOpt NVFP4 | 324 GB | 3 | 4 | +| EXL3 2.5 bpw | 181 GB | 2 | 2 | +| EXL3 1.25 bpw | 92 GB | 1 | 1 | + +The first count is weights alone; the second leaves room for KV cache, activations and the system. Loaded, the FP8 +lane weights add 0.94 GB of fp32 scale tiles (one per 64 inputs) on the non-expert linears; NVFP4 loads at its stored +size. + +## How to help + +Each item stands alone. + +1. **KDA block on `qlinear`.** Write Ling's KDA block over `weights.Kda`: seven `qlinear.Linear` calls (or one fused + `[q | k | v]` per format) feeding #804's chain with `conv`, `a_log`, `dt_bias`, `o_norm` as loaded. Check it with + `ling-synth` plus `glm-kda-ling`'s fixture. +2. **MLA block on `qlinear`.** The same for `weights.Mla` on `ling-mla`'s kernels (`q_a`, `q_b`, `kv_a`, `kv_b`, + head gate, `dense`). +3. **MoE block.** Route with GLM's `moe.Layer` (`ling.route()` fills top-k, groups and limits), run `experts.Grouped` + on `Moe.routed`, the shared expert as three projections. +4. **Bf16 routed experts.** Add a grouped bf16 expert kernel, or quantize at load, and drop the per-expert lanes. +5. **TP.** Cut `weights.load` by heads and widths as `glm/cuda_shard.zig` does; FP8 scale blocks span 128 rows and + 128 inputs, so FP8 cuts fall on multiples of 128. +6. **EXL3.** Wire #832's reader behind the `config.Quant` extension point and lift the refusal. +7. **An NVFP4 checkpoint.** Quantize the release with ModelOpt, keep the routers in bf16, and run `ling-weights` on + it; the loader expects `weight`, `weight_scale`, `weight_scale_2` and `input_scale` per linear. +8. **GPU packing of FP8 experts.** They pack on the host today (`Fp8Experts.packGateUp`); `modelopt.Packer` has the + NVFP4 experts' GPU packer to follow. From 7d184cb7e92885f9fabc9d1e14ef0a61dda1960c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 15 Oct 2026 09:07:06 +0200 Subject: [PATCH 05/34] glm: MLA kernels take Ling-3.1's shape: 64 roped key columns, 128-wide heads, head-wise gate chunks_wide: q and cache rows LW + RW wide, scores over all of them, values from the 512 latent; tf_glm_chunks_r16 is the 576-wide 16-head form (RW 0 builds GLM's kernels to the same SASS). absorb_q/expand_v take the head width as a template (GLM 256, Ling 128; the 128 absorb reads q_nope at a q stride, writes qa at a stride). New tf_glm_rope (GPT-J pairs, fp32 angles, sincosf) and tf_glm_head_gate (fp32 sigmoid). Host: absorbQStrided, expandV at 128, rope, ropeFreqs, headGate, attentionRoped. Checks: ling-mla (float64 bounds; decode rows, a verify window and a cut prompt across key 1536 equal the uncut call), ling-mla-real (layer 7 vs HF and float64; 400 + 32 decode rows, an 8-row window and a cut at 300 equal the prompt), glm-mla-same (GLM's MLA launches against a 6380c8da glm_glue fatbin, byte for byte). --- zig/kernels/cuda/glm_glue_attn.cuh | 51 ++--- zig/kernels/cuda/glm_glue_latent.cuh | 80 ++++++-- zig/src/families/glm/cuda_glue.zig | 11 ++ zig/src/families/glm/cuda_latent.zig | 72 ++++++-- zig/tests/cuda/glm_mla_same.zig | 116 ++++++++++++ zig/tests/cuda/ling_mla_real.zig | 139 ++++++++++++++ zig/tests/cuda/ling_ref_mla.zig | 267 +++++++++++++++++++++++++++ zig/tests/cuda/main.zig | 6 + 8 files changed, 693 insertions(+), 49 deletions(-) create mode 100644 zig/tests/cuda/glm_mla_same.zig create mode 100644 zig/tests/cuda/ling_mla_real.zig create mode 100644 zig/tests/cuda/ling_ref_mla.zig diff --git a/zig/kernels/cuda/glm_glue_attn.cuh b/zig/kernels/cuda/glm_glue_attn.cuh index db481c0d9e..2c2bcaed86 100644 --- a/zig/kernels/cuda/glm_glue_attn.cuh +++ b/zig/kernels/cuda/glm_glue_attn.cuh @@ -39,26 +39,28 @@ __device__ __forceinline__ void ldm4(uint32_t* r, const void* p) { : "r"(a)); } -// A 512-wide bf16 row's 16-byte unit u in shared memory, swizzled by the row so ldmatrix's eight rows hit eight banks. -__device__ __forceinline__ int wide(int row, int u) { return row * LW + ((u ^ (row & 7)) << 3); } +// A KW-wide bf16 row's 16-byte unit u in shared memory, swizzled by the row so ldmatrix's eight rows hit eight banks. +template +__device__ __forceinline__ int wide(int row, int u) { return row * KW + ((u ^ (row & 7)) << 3); } // p's [32][32] rows (64 bytes): unit u of row r. __device__ __forceinline__ int narrow(int row, int u) { return row * KT + ((u ^ ((row >> 1) & 3)) << 3); } // latent._dense_chunks/_sparse_chunks (sparse 0/1), block (row, HB heads, chunk), 8 warps: online softmax by 32 keys. -template +template constexpr int wide_smem() { - return (HB * LW + 2 * KT * LW + HB * KT) * 2 + (HB * 4 + 2 * HB) * 4; + return (HB * (LW + RW) + 2 * KT * (LW + RW) + HB * KT) * 2 + (HB * 4 + 2 * HB) * 4; } -template +// RW rope columns (Ling's 64) follow the latent in q and keys: scores over LW + RW, values from the latent alone. +template __device__ void chunks_wide(const __nv_bfloat16* qa, const __nv_bfloat16* lc, const int* pos, const int* tok, const int* cnt, float* po, float* pm, float* pl, int64_t rows, int64_t heads, int64_t w, float scale, int sparse) { - constexpr int MT = HB / 16; + constexpr int MT = HB / 16, KW = LW + RW; extern __shared__ __align__(16) unsigned char smem[]; __nv_bfloat16* qs = reinterpret_cast<__nv_bfloat16*>(smem); - __nv_bfloat16* kvs = qs + HB * LW; - __nv_bfloat16* ps = kvs + 2 * KT * LW; + __nv_bfloat16* kvs = qs + HB * KW; + __nv_bfloat16* ps = kvs + 2 * KT * KW; float* red = reinterpret_cast(ps + HB * KT); float* stat = red + HB * 4; const int t = threadIdx.x, lane = t & 31, warp = t >> 5, g = lane >> 2, q = lane & 3; @@ -67,14 +69,14 @@ __device__ void chunks_wide(const __nv_bfloat16* qa, const __nv_bfloat16* lc, co const int64_t n = sparse ? cnt[r] : 0, limit = sparse ? 0 : (int64_t)pos[0] + r, start = c * CH; const bool live = sparse ? start < n : start <= limit; auto load_kv = [&](int tile) { - __nv_bfloat16* dst = kvs + (tile & 1) * KT * LW; + __nv_bfloat16* dst = kvs + (tile & 1) * KT * KW; const int64_t first = start + tile * KT; - for (int idx = t; idx < KT * (LW / 8); idx += blockDim.x) { - const int key = idx / (LW / 8), u = idx % (LW / 8); + for (int idx = t; idx < KT * (KW / 8); idx += blockDim.x) { + const int key = idx / (KW / 8), u = idx % (KW / 8); const int64_t ki = first + key; const bool ok = sparse ? ki < n : ki <= limit; const int64_t src = ok ? (sparse ? (int64_t)tok[r * w + ki] : ki) : 0; - cp16(dst + wide(key, u), lc + src * LW + u * 8, ok); + cp16(dst + wide(key, u), lc + src * KW + u * 8, ok); } }; float m[MT][2], l[MT][2], o[MT][8][4]; @@ -85,10 +87,10 @@ __device__ void chunks_wide(const __nv_bfloat16* qa, const __nv_bfloat16* lc, co for (int e = 0; e < 4; ++e) o[i][j][e] = 0.0f; constexpr int tiles = CH / KT; if (live) { - for (int idx = t; idx < HB * (LW / 8); idx += blockDim.x) { - const int hd = idx / (LW / 8), u = idx % (LW / 8); + for (int idx = t; idx < HB * (KW / 8); idx += blockDim.x) { + const int hd = idx / (KW / 8), u = idx % (KW / 8); const bool ok = h0 + hd < heads; - cp16(qs + wide(hd, u), qa + (r * heads + (ok ? h0 + hd : 0)) * LW + u * 8, ok); + cp16(qs + wide(hd, u), qa + (r * heads + (ok ? h0 + hd : 0)) * KW + u * 8, ok); } load_kv(0); commit(); @@ -101,17 +103,17 @@ __device__ void chunks_wide(const __nv_bfloat16* qa, const __nv_bfloat16* lc, co else wait<0>(); __syncthreads(); - const __nv_bfloat16* kv = kvs + (tile & 1) * KT * LW; + const __nv_bfloat16* kv = kvs + (tile & 1) * KT * KW; const int64_t first = start + tile * KT; float sc[4]; if (hh < MT) { float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; #pragma unroll 4 - for (int kk = 0; kk < LW; kk += 32) { + for (int kk = 0; kk < KW; kk += 32) { uint32_t a0[4], a1[4], b[4]; - ldm4(a0, qs + wide(hh * 16 + (lane & 15), (kk >> 3) + (lane >> 4))); - ldm4(a1, qs + wide(hh * 16 + (lane & 15), ((kk + 16) >> 3) + (lane >> 4))); - ldm4(b, kv + wide(kg * 8 + (lane & 7), (kk >> 3) + (lane >> 3))); + ldm4(a0, qs + wide(hh * 16 + (lane & 15), (kk >> 3) + (lane >> 4))); + ldm4(a1, qs + wide(hh * 16 + (lane & 15), ((kk + 16) >> 3) + (lane >> 4))); + ldm4(b, kv + wide(kg * 8 + (lane & 7), (kk >> 3) + (lane >> 3))); glm_router::mma16816(acc, a0, b); glm_router::mma16816(acc, a1, b + 2); } @@ -172,7 +174,7 @@ __device__ void chunks_wide(const __nv_bfloat16* qa, const __nv_bfloat16* lc, co #pragma unroll for (int j = 0; j < 8; j += 2) { uint32_t b[4]; - ldm4(b, kv + wide(k16 + (lane & 15), ((warp * 64 + j * 8) >> 3) + (lane >> 4))); + ldm4(b, kv + wide(k16 + (lane & 15), ((warp * 64 + j * 8) >> 3) + (lane >> 4))); for (int i = 0; i < MT; ++i) { glm_router::mma16816(o[i][j], a[i], b); glm_router::mma16816(o[i][j + 1], a[i], b + 2); @@ -213,6 +215,13 @@ extern "C" __global__ void __launch_bounds__(256) glm_attn::chunks_wide<32>(qa, lc, pos, tok, cnt, po, pm, pl, rows, heads, w, scale, sparse); } +// Ling's 16-head form: 576-wide q and keys (latent, then 64 roped columns), in wide_smem<16, 64>() bytes. +extern "C" __global__ void __launch_bounds__(256) + tf_glm_chunks_r16(const __nv_bfloat16* qa, const __nv_bfloat16* lc, const int* pos, const int* tok, const int* cnt, + float* po, float* pm, float* pl, int64_t rows, int64_t heads, int64_t w, float scale, int sparse) { + glm_attn::chunks_wide<16, 64>(qa, lc, pos, tok, cnt, po, pm, pl, rows, heads, w, scale, sparse); +} + // latent._merge, block (row, head), 128 threads: chunks merged in order, o / l by div.full; skips sparse count-0 rows. extern "C" __global__ void __launch_bounds__(128) tf_glm_merge(const float* po, const float* pm, const float* pl, __nv_bfloat16* out, const int* cnt, int64_t rows, diff --git a/zig/kernels/cuda/glm_glue_latent.cuh b/zig/kernels/cuda/glm_glue_latent.cuh index 12afc2daa3..ff29d6a74b 100644 --- a/zig/kernels/cuda/glm_glue_latent.cuh +++ b/zig/kernels/cuda/glm_glue_latent.cuh @@ -1,4 +1,4 @@ -// latent.py's BF16 absorb, expand and cache write by hand (the MLA layers' 512-wide latent, 256-wide heads). +// latent.py's BF16 absorb, expand and cache write by hand (512-wide latent; GLM's 256-wide heads, Ling's 128). #pragma once #include "glm_triton.cuh" @@ -6,23 +6,23 @@ using namespace glm_triton; // latent._absorb_q(_rows), block (head, 32 latent columns, row block) of 128 threads; rows per block change no bits. -extern "C" __global__ void __launch_bounds__(128) - tf_glm_absorb_q(const __nv_bfloat16* q, const __nv_bfloat16* wk, __nv_bfloat16* qa, int64_t rows, int64_t heads, - int64_t rb) { - constexpr int D = 256, LW = 512; +template // q_nope D wide (GLM 256, Ling 128), q rows qs apart per head, qa rows qaw apart +__device__ __forceinline__ void absorb_q(const __nv_bfloat16* q, const __nv_bfloat16* wk, __nv_bfloat16* qa, + int64_t rows, int64_t heads, int64_t rb, int64_t qs, int64_t qaw) { + constexpr int LW = 512, R = D / 32; __shared__ float smem[128]; const int t = threadIdx.x, k0 = t >> 2, nq = t & 3, warp = t >> 5; const int64_t h = blockIdx.x, n0 = (int64_t)blockIdx.y * 32 + nq * 8, r0 = blockIdx.z * rb; - float w[8][8]; - for (int rep = 0; rep < 8; ++rep) + float w[R][8]; + for (int rep = 0; rep < R; ++rep) for (int j = 0; j < 8; ++j) w[rep][j] = bf(wk[(h * D + k0 + 32 * rep) * LW + n0 + j]); for (int64_t r = r0; r < r0 + rb && r < rows; ++r) { - float qv[8]; - for (int rep = 0; rep < 8; ++rep) qv[rep] = bf(q[(r * heads + h) * D + k0 + 32 * rep]); + float qv[R]; + for (int rep = 0; rep < R; ++rep) qv[rep] = bf(q[(r * heads + h) * qs + k0 + 32 * rep]); for (int j = 0; j < 8; ++j) { float a = fmul(w[1][j], qv[1]); a = __fmaf_rn(w[0][j], qv[0], a); - for (int rep = 2; rep < 8; ++rep) a = __fmaf_rn(w[rep][j], qv[rep], a); + for (int rep = 2; rep < R; ++rep) a = __fmaf_rn(w[rep][j], qv[rep], a); for (int o = 16; o >= 4; o >>= 1) a = a + __shfl_xor_sync(0xffffffffu, a, o); if ((t & 28) == 0) smem[(nq * 8 + j) * 4 + warp] = a; } @@ -30,16 +30,29 @@ extern "C" __global__ void __launch_bounds__(128) float v = smem[t]; v = v + __shfl_xor_sync(0xffffffffu, v, 2); v = v + __shfl_xor_sync(0xffffffffu, v, 1); - if ((t & 3) == 0) qa[(r * heads + h) * LW + (int64_t)blockIdx.y * 32 + t / 4] = __float2bfloat16_rn(v); + if ((t & 3) == 0) qa[(r * heads + h) * qaw + (int64_t)blockIdx.y * 32 + t / 4] = __float2bfloat16_rn(v); __syncthreads(); } } -// latent._expand_v(_rows), block (head, 16 outputs, row block) of 128 threads: o . wv by fma, lanes xor 16 .. 1. extern "C" __global__ void __launch_bounds__(128) - tf_glm_expand_v(const __nv_bfloat16* ol, const __nv_bfloat16* wv, __nv_bfloat16* out, int64_t rows, - int64_t heads, int64_t rb) { - constexpr int DV = 256, LW = 512; + tf_glm_absorb_q(const __nv_bfloat16* q, const __nv_bfloat16* wk, __nv_bfloat16* qa, int64_t rows, int64_t heads, + int64_t rb) { + absorb_q<256>(q, wk, qa, rows, heads, rb, 256, 512); +} + +// Ling: q_nope the first 128 of each head's qs columns, qa rows qaw (576: the latent, then the roped q_rope). +extern "C" __global__ void __launch_bounds__(128) + tf_glm_absorb_q_128(const __nv_bfloat16* q, const __nv_bfloat16* wk, __nv_bfloat16* qa, int64_t rows, + int64_t heads, int64_t rb, int64_t qs, int64_t qaw) { + absorb_q<128>(q, wk, qa, rows, heads, rb, qs, qaw); +} + +// latent._expand_v(_rows), block (head, 16 outputs, row block) of 128 threads: o . wv by fma, lanes xor 16 .. 1. +template +__device__ __forceinline__ void expand_v(const __nv_bfloat16* ol, const __nv_bfloat16* wv, __nv_bfloat16* out, + int64_t rows, int64_t heads, int64_t rb) { + constexpr int LW = 512; __shared__ float smem[32]; const int t = threadIdx.x, lane = t & 31, warp = t >> 5, kb = (warp & 1) * 256 + lane * 8, nb = warp >> 1; const int64_t h = blockIdx.x, n0 = (int64_t)blockIdx.y * 16, r0 = blockIdx.z * rb; @@ -66,9 +79,46 @@ extern "C" __global__ void __launch_bounds__(128) } } +extern "C" __global__ void __launch_bounds__(128) + tf_glm_expand_v(const __nv_bfloat16* ol, const __nv_bfloat16* wv, __nv_bfloat16* out, int64_t rows, + int64_t heads, int64_t rb) { + expand_v<256>(ol, wv, out, rows, heads, rb); +} + +// Ling's 128-wide v heads. +extern "C" __global__ void __launch_bounds__(128) + tf_glm_expand_v_128(const __nv_bfloat16* ol, const __nv_bfloat16* wv, __nv_bfloat16* out, int64_t rows, + int64_t heads, int64_t rb) { + expand_v<128>(ol, wv, out, rows, heads, rb); +} + // latent._lat_write (BF16 cache): row r of `lat` into cache slot pos + r. extern "C" __global__ void tf_glm_lat_write(const __nv_bfloat16* lat, int64_t lat_stride, __nv_bfloat16* cache, const int* pos, int64_t lw) { const int64_t r = blockIdx.x, at = (int64_t)pos[0] + r; for (int64_t k = threadIdx.x; k < lw; k += blockDim.x) cache[at * lw + k] = lat[r * lat_stride + k]; } + +// Ling's rope (exl3's GPT-J pairs), block (row, head) of 32: x + r xs + h xh into out + r os + h oh, in place ok. +extern "C" __global__ void __launch_bounds__(32) + tf_glm_rope(const __nv_bfloat16* x, int64_t xs, int64_t xh, __nv_bfloat16* out, int64_t os, int64_t oh, + const int* pos, const float* inv_freq) { + const int64_t r = blockIdx.x, h = blockIdx.y; + const int i = threadIdx.x; // pair i at angle fp32 inv_freq[i] (pos + r); explicit rn ops, no fma contraction + const __nv_bfloat16* src = x + r * xs + h * xh + 2 * i; + const float x0 = bf(src[0]), x1 = bf(src[1]); + float s, c; + sincosf(__fmul_rn(inv_freq[i], (float)((int64_t)pos[0] + r)), &s, &c); + __nv_bfloat16* dst = out + r * os + h * oh + 2 * i; + dst[0] = __float2bfloat16_rn(__fsub_rn(__fmul_rn(x0, c), __fmul_rn(x1, s))); + dst[1] = __float2bfloat16_rn(__fadd_rn(__fmul_rn(x1, c), __fmul_rn(x0, s))); +} + +// Ling's head-wise output gate, block (row, head) of dv threads: o *= sigmoid(g[r gs + h]) in fp32, rounded to bf16. +extern "C" __global__ void __launch_bounds__(256) + tf_glm_head_gate(__nv_bfloat16* o, int64_t heads, const float* g, int64_t gs) { + const int64_t r = blockIdx.x, h = blockIdx.y, dv = blockDim.x; + const float s = __fdiv_rn(1.0f, __fadd_rn(1.0f, expf(-g[r * gs + h]))); + __nv_bfloat16* p = o + (r * heads + h) * dv + threadIdx.x; + *p = __float2bfloat16_rn(__fmul_rn(bf(*p), s)); +} diff --git a/zig/src/families/glm/cuda_glue.zig b/zig/src/families/glm/cuda_glue.zig index abb53038d0..dbc52cc201 100644 --- a/zig/src/families/glm/cuda_glue.zig +++ b/zig/src/families/glm/cuda_glue.zig @@ -31,7 +31,11 @@ pub const Kernels = struct { hc_partial8: cuda.Function, hc_finish: cuda.Function, absorb_q: cuda.Function, + absorb_q_128: cuda.Function, expand_v: cuda.Function, + expand_v_128: cuda.Function, + rope: cuda.Function, + head_gate: cuda.Function, lat_write: cuda.Function, index_write: cuda.Function, pool_keys: cuda.Function, @@ -46,6 +50,7 @@ pub const Kernels = struct { bmm_d64s2: cuda.Function, chunks_16: cuda.Function, chunks_32: cuda.Function, + chunks_r16: cuda.Function, merge: cuda.Function, scores: cuda.Function, @@ -56,6 +61,7 @@ pub const Kernels = struct { try k.bmm_tile1.allowDynamicShared(bmm_tile1_smem[1]); try k.chunks_16.allowDynamicShared(chunksSmem(16)); try k.chunks_32.allowDynamicShared(chunksSmem(32)); + try k.chunks_r16.allowDynamicShared(chunksSmemRoped(16)); return k; } }; @@ -138,6 +144,11 @@ pub fn chunksSmem(comptime hb: u32) u32 { return (hb * 512 + 2 * 32 * 512 + hb * 32) * 2 + (hb * 4 + 2 * hb) * 4; } +/// wide_smem(): as above with q and key rows 576 wide (Ling's roped columns after the latent). +pub fn chunksSmemRoped(comptime hb: u32) u32 { + return (hb * 576 + 2 * 32 * 576 + hb * 32) * 2 + (hb * 4 + 2 * hb) * 4; +} + /// K slices of the router's dot (glue.ROUTER_KS): fixed by the shape, as the sliced and the single pass differ in bits. pub const router_ks = 8; diff --git a/zig/src/families/glm/cuda_latent.zig b/zig/src/families/glm/cuda_latent.zig index a41381ef3b..6d2b08ad31 100644 --- a/zig/src/families/glm/cuda_latent.zig +++ b/zig/src/families/glm/cuda_latent.zig @@ -1,4 +1,4 @@ -//! GLM-5.3-Flash's MLA on the 512-wide latent cache (DSA layers): the Python family's latent on BF16, row by row. +//! GLM-5.3-Flash's MLA on the 512-wide latent cache (DSA layers), and Ling-3.1's with 64 roped columns after it. const std = @import("std"); const cuda = @import("cuda"); @@ -29,21 +29,60 @@ fn cdiv(a: usize, b: usize) u32 { /// latent.absorb_q (BF16 kv_b): out [rows][heads][lw] bf16 = q [rows][heads][d] . wk [heads][d][lw], fp32 sums. pub fn absorbQ(set: glue.Set, stream: cuda.Stream, q: u64, wk: u64, out: u64, rows: usize, heads: usize, d: usize, lw: usize) !void { if (d != 256 or lw != 512) return error.LatentShape; - try rowBlocks(set.glue.absorb_q, stream, q, wk, out, rows, heads, lw / 32); + try rowBlocks(set.glue.absorb_q, stream, q, wk, out, rows, heads, lw / 32, .{}); +} + +/// Ling's absorb: q_nope the first d = 128 of each head's `q_stride` columns, out rows `out_stride` apart per head. +pub fn absorbQStrided(set: glue.Set, stream: cuda.Stream, q: u64, q_stride: usize, wk: u64, out: u64, out_stride: usize, rows: usize, heads: usize, d: usize, lw: usize) !void { + if (d != 128 or lw != 512 or q_stride < d or out_stride < lw) return error.LatentShape; + try rowBlocks(set.glue.absorb_q_128, stream, q, wk, out, rows, heads, lw / 32, .{ @as(i64, @intCast(q_stride)), @as(i64, @intCast(out_stride)) }); } /// The absorb and expand launches: rows in one block in decode windows, prompt_rb a block in prompt chunks. -fn rowBlocks(f: cuda.Function, stream: cuda.Stream, in: u64, w: u64, out: u64, rows: usize, heads: usize, cols: usize) !void { +fn rowBlocks(f: cuda.Function, stream: cuda.Stream, in: u64, w: u64, out: u64, rows: usize, heads: usize, cols: usize, extra: anytype) !void { const rb = if (rows >= prompt_rows) prompt_rb else rows; var a: cuda.Args = .{}; - inline for (.{ in, w, out, @as(i64, @intCast(rows)), @as(i64, @intCast(heads)), @as(i64, @intCast(rb)) }) |v| a.add(v); + inline for (.{ in, w, out, @as(i64, @intCast(rows)), @as(i64, @intCast(heads)), @as(i64, @intCast(rb)) } ++ extra) |v| a.add(v); try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(heads), .y = @intCast(cols), .z = cdiv(rows, rb) }, .block = .{ .x = 128 } }, stream, &a); } -/// latent.expand_v (BF16 kv_b): out [rows][heads][dv] bf16 = o_lat [rows][heads][lw] . wv [heads][dv][lw]. +/// latent.expand_v (BF16 kv_b): out [rows][heads][dv] bf16 = o_lat [rows][heads][lw] . wv [heads][dv][lw], dv 256|128. pub fn expandV(set: glue.Set, stream: cuda.Stream, ol: u64, wv: u64, out: u64, rows: usize, heads: usize, dv: usize, lw: usize) !void { - if (dv != 256 or lw != 512) return error.LatentShape; - try rowBlocks(set.glue.expand_v, stream, ol, wv, out, rows, heads, dv / 16); + if (lw != 512) return error.LatentShape; + const f = switch (dv) { + 256 => set.glue.expand_v, + 128 => set.glue.expand_v_128, + else => return error.LatentShape, + }; + try rowBlocks(f, stream, ol, wv, out, rows, heads, dv / 16, .{}); +} + +/// Ling's 64 rope columns of `heads` heads a row: in at x + r x_row + h x_head, out likewise (bf16 elements). +pub const RopeRows = struct { x: u64, x_row: usize, x_head: usize, out: u64, o_row: usize, o_head: usize }; + +/// The rope (GPT-J pairs) at positions pos .. pos + rows - 1 (pos an i32 on device), inv_freq 32 fp32 on device. +pub fn rope(set: glue.Set, stream: cuda.Stream, v: RopeRows, rows: usize, heads: usize, pos: u64, inv_freq: u64) !void { + var a: cuda.Args = .{}; + inline for (.{ v.x, @as(i64, @intCast(v.x_row)), @as(i64, @intCast(v.x_head)), v.out, @as(i64, @intCast(v.o_row)), @as(i64, @intCast(v.o_head)), pos, inv_freq }) |x| a.add(x); + try cuda.launch.launch(set.glue.rope, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(heads) }, .block = .{ .x = rope_width / 2 } }, stream, &a); +} + +/// Ling's qk_rope_head_dim: the roped columns after the latent in q and in the cache. +pub const rope_width = 64; + +/// The rope's inv_freq: theta^(-2i / 64) for i < 32, from float64 rounded to fp32. +pub fn ropeFreqs(theta: f64) [rope_width / 2]f32 { + var f: [rope_width / 2]f32 = undefined; + for (&f, 0..) |*x, i| x.* = @floatCast(1.0 / std.math.pow(f64, theta, @as(f64, @floatFromInt(2 * i)) / rope_width)); + return f; +} + +/// Ling's head-wise gate: o [rows][heads][dv] bf16 *= sigmoid(g fp32 [rows][g_stride], column h), in place. +pub fn headGate(set: glue.Set, stream: cuda.Stream, o: u64, g: u64, g_stride: usize, rows: usize, heads: usize, dv: usize) !void { + if (dv == 0 or dv > 256) return error.LatentShape; + var a: cuda.Args = .{}; + inline for (.{ o, @as(i64, @intCast(heads)), g, @as(i64, @intCast(g_stride)) }) |x| a.add(x); + try cuda.launch.launch(set.glue.head_gate, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(heads) }, .block = .{ .x = @intCast(dv) } }, stream, &a); } /// latent.latent_write: rows of `lat` (stride `lat_stride`) into slots pos .. pos + rows - 1 (pos an i32 on device). @@ -58,28 +97,35 @@ pub const Partials = struct { po: u64, pm: u64, pl: u64 }; /// latent.attention: causal qa over the cache in `nch` 512-key chunks into out bf16; `dummy`: the merge's count i32. pub fn attention(set: glue.Set, stream: cuda.Stream, qa: u64, cache: u64, pos: u64, p: Partials, out: u64, dummy: u64, rows: usize, heads: usize, lw: usize, nch: usize, scale: f32, hb: usize) !void { - try chunks(set, stream, qa, cache, pos, 0, 0, p, rows, heads, lw, nch, 0, scale, hb, false); + try chunks(set, stream, qa, cache, pos, 0, 0, p, rows, heads, lw, 0, nch, 0, scale, hb, false); + try merge(set, stream, p, out, dummy, rows, heads, lw, nch, false); +} + +/// Ling's latent.attention: qa and cache rows lw + rope_width wide (roped columns last), 16-head blocks; out lw wide. +pub fn attentionRoped(set: glue.Set, stream: cuda.Stream, qa: u64, cache: u64, pos: u64, p: Partials, out: u64, dummy: u64, rows: usize, heads: usize, lw: usize, nch: usize, scale: f32) !void { + try chunks(set, stream, qa, cache, pos, 0, 0, p, rows, heads, lw, rope_width, nch, 0, scale, heads_tile, false); try merge(set, stream, p, out, dummy, rows, heads, lw, nch, false); } /// latent._dense_chunks (sparse false: keys through pos + r) and _sparse_chunks (row r's `counts` tokens of `tokens`). -fn chunks(set: glue.Set, stream: cuda.Stream, qa: u64, cache: u64, pos: u64, tokens: u64, counts: u64, p: Partials, rows: usize, heads: usize, lw: usize, nch: usize, w: usize, scale: f32, hb: usize, sparse: bool) !void { - if (lw != 512) return error.LatentShape; - const f = switch (hb) { +fn chunks(set: glue.Set, stream: cuda.Stream, qa: u64, cache: u64, pos: u64, tokens: u64, counts: u64, p: Partials, rows: usize, heads: usize, lw: usize, rw: usize, nch: usize, w: usize, scale: f32, hb: usize, sparse: bool) !void { + if (lw != 512 or (rw != 0 and rw != rope_width)) return error.LatentShape; + if (rw != 0 and hb != heads_tile) return error.HeadBlock; + const f = if (rw != 0) set.glue.chunks_r16 else switch (hb) { heads_tile => set.glue.chunks_16, heads_tile_wide => set.glue.chunks_32, else => return error.HeadBlock, }; var a: cuda.Args = .{}; inline for (.{ qa, cache, pos, tokens, counts, p.po, p.pm, p.pl, @as(i64, @intCast(rows)), @as(i64, @intCast(heads)), @as(i64, @intCast(w)), scale, @as(c_int, @intFromBool(sparse)) }) |v| a.add(v); - const shared = if (hb == heads_tile_wide) glue.chunksSmem(heads_tile_wide) else glue.chunksSmem(heads_tile); + const shared = if (rw != 0) glue.chunksSmemRoped(heads_tile) else if (hb == heads_tile_wide) glue.chunksSmem(heads_tile_wide) else glue.chunksSmem(heads_tile); try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows), .y = cdiv(heads, hb), .z = @intCast(nch) }, .block = .{ .x = 256 }, .shared = shared }, stream, &a); } /// latent.sparse_attention: rows with counts > 0 over their selected tokens (i32 [rows][w]); others left alone. pub fn sparseAttention(set: glue.Set, stream: cuda.Stream, qa: u64, cache: u64, tokens: u64, counts: u64, w: usize, p: Partials, out: u64, rows: usize, heads: usize, lw: usize, scale: f32) !void { const nch = chunksFor(w); - try chunks(set, stream, qa, cache, 0, tokens, counts, p, rows, heads, lw, nch, w, scale, headBlock(rows), true); + try chunks(set, stream, qa, cache, 0, tokens, counts, p, rows, heads, lw, 0, nch, w, scale, headBlock(rows), true); try merge(set, stream, p, out, counts, rows, heads, lw, nch, true); } diff --git a/zig/tests/cuda/glm_mla_same.zig b/zig/tests/cuda/glm_mla_same.zig new file mode 100644 index 0000000000..f5bf6b4afa --- /dev/null +++ b/zig/tests/cuda/glm_mla_same.zig @@ -0,0 +1,116 @@ +//! GLM's MLA launches on this build's glm_glue against a 6380c8da glm_glue fatbin: the same bytes for GLM's shapes. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const qr = @import("quant_ref_tests.zig"); +const glm = @import("glm_cuda"); +const Gpu = check.Gpu; +const Rig = qr.RigT; +const latent = glm.latent; +const Set = glm.glue.Set; + +const heads = 32; +const lw = 512; +const d = 256; + +/// The kernels Ling added, absent from the old image: taken from the new one (GLM's launches never reach them). +const added = [_][]const u8{ "absorb_q_128", "expand_v_128", "rope", "head_gate", "chunks_r16" }; + +fn oldKernels(old: cuda.Module, new: glm.glue.Kernels) !glm.glue.Kernels { + var k = new; + inline for (@typeInfo(glm.glue.Kernels).@"struct".field_names) |name| { + const fresh = for (added) |a| { + if (std.mem.eql(u8, a, name)) break true; + } else false; + if (!fresh) @field(k, name) = try old.function("tf_glm_" ++ name); + } + try k.chunks_16.allowDynamicShared(glm.glue.chunksSmem(16)); + try k.chunks_32.allowDynamicShared(glm.glue.chunksSmem(32)); + return k; +} + +fn partials(r: *Rig, nch: usize, rows: usize) !latent.Partials { + return .{ .po = try r.zeros(nch * rows * heads * lw * 4), .pm = try r.zeros(nch * rows * heads * 4), .pl = try r.zeros(nch * rows * heads * 4) }; +} + +const Inputs = struct { q: u64, wk: u64, ol: u64, wv: u64, cache: u64, lat: u64, tok: u64, cnt: u64 }; + +/// Every GLM MLA launch (absorb, expand, cache write, dense 16/32-head and sparse attention) on `set`: output bytes. +fn outputs(r: *Rig, set: Set, in: Inputs) ![]const []const u8 { + var list: std.ArrayList([]const u8) = .empty; + for ([_]usize{ 1, 5, 70 }) |rows| { + const qa = try r.zeros(rows * heads * lw * 2); + try latent.absorbQ(set, r.s, in.q, in.wk, qa, rows, heads, d, lw); + try list.append(r.a, try r.back(qa, rows * heads * lw * 2)); + const vn = try r.zeros(rows * heads * d * 2); + try latent.expandV(set, r.s, in.ol, in.wv, vn, rows, heads, d, lw); + try list.append(r.a, try r.back(vn, rows * heads * d * 2)); + } + const c = try r.zeros(1500 * lw * 2); + const p0: i32 = 1400; + try latent.latentWrite(set, r.s, in.lat, lw + 64, c, try r.dev(std.mem.asBytes(&p0)), 100, lw); + try list.append(r.a, try r.back(c, 1500 * lw * 2)); + for ([_][3]usize{ .{ 1430, 70, 32 }, .{ 1430, 70, 16 }, .{ 1499, 1, 16 }, .{ 0, 64, 32 } }) |case| { + const pos, const rows, const hb = case; + const nch = latent.chunksFor(pos + rows); + const out = try r.zeros(rows * heads * lw * 2); + const at: i32 = @intCast(pos); + try latent.attention(set, r.s, in.q, in.cache, try r.dev(std.mem.asBytes(&at)), try partials(r, nch, rows), out, try r.zeros(4), rows, heads, lw, nch, 0.125, hb); + try list.append(r.a, try r.back(out, rows * heads * lw * 2)); + } + for ([_]usize{ 70, 1 }) |rows| { + const out = try r.zeros(rows * heads * lw * 2); + try latent.sparseAttention(set, r.s, in.q, in.cache, in.tok, in.cnt, 1100, try partials(r, latent.chunksFor(1100), rows), out, rows, heads, lw, 0.125); + try list.append(r.a, try r.back(out, rows * heads * lw * 2)); + } + return list.items; +} + +/// `glm-mla-same `: the old and new images' GLM MLA outputs, byte for byte, on seeded inputs. +pub fn run(gpu: Gpu, old_path: []const u8) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(804); + var r: Rig = .{ .gpu = gpu, .a = a, .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + } + const raw = try std.Io.Dir.cwd().readFileAlloc(gpu.io, old_path, a, .limited(64 << 20)); + const image = try a.alignedAlloc(u8, .@"16", raw.len); + @memcpy(image, raw); + var new_mod = try cuda.Module.load(gpu.d, cuda.kernels.glm_glue); + defer new_mod.unload(); + var old_mod = try cuda.Module.load(gpu.d, image); + defer old_mod.unload(); + const new = try glm.glue.Kernels.load(new_mod); + const old = try oldKernels(old_mod, new); + const tok = try a.alloc(i32, 70 * 1100); + const cnt = try a.alloc(i32, 70); + for (cnt, 0..) |*n, row| { + n.* = if (row == 5) 0 else @intCast(1 + r.rng.uintLessThan(usize, 1100)); + for (tok[row * 1100 ..][0..1100], 0..) |*t, i| t.* = if (i < n.*) @intCast(r.rng.uintLessThan(u32, 6000)) else -1; + } + const in: Inputs = .{ + .q = try r.dev(std.mem.sliceAsBytes(try r.bfs(70 * heads * lw, 1.0))), + .wk = try r.dev(std.mem.sliceAsBytes(try r.bfs(heads * d * lw, 0.05))), + .ol = try r.dev(std.mem.sliceAsBytes(try r.bfs(70 * heads * lw, 1.0))), + .wv = try r.dev(std.mem.sliceAsBytes(try r.bfs(heads * d * lw, 0.05))), + .cache = try r.dev(std.mem.sliceAsBytes(try r.bfs(6000 * lw, 1.0))), + .lat = try r.dev(std.mem.sliceAsBytes(try r.bfs(100 * (lw + 64), 1.0))), + .tok = try r.dev(std.mem.sliceAsBytes(tok)), + .cnt = try r.dev(std.mem.sliceAsBytes(cnt)), + }; + const want = try outputs(&r, .{ .glue = &old }, in); + const got = try outputs(&r, .{ .glue = &new }, in); + const names = [_][]const u8{ "absorb_q 1 row", "expand_v 1 row", "absorb_q 5 rows", "expand_v 5 rows", "absorb_q 70 rows", "expand_v 70 rows", "lat_write 100 rows", "dense 70 rows at 1430, 32-head", "dense 70 rows at 1430, 16-head", "dense 1 row at 1499", "dense 64 rows at 0, 32-head", "sparse 70 rows", "sparse 1 row" }; + var bytes: usize = 0; + for (names, want, got) |name, w, g| { + try check.sameBytes(name, g, w); + bytes += g.len; + } + check.pass("glm-mla-same: {d} GLM MLA launches, {d} output bytes equal 6380c8da's", .{ names.len, bytes }); +} diff --git a/zig/tests/cuda/ling_mla_real.zig b/zig/tests/cuda/ling_mla_real.zig new file mode 100644 index 0000000000..baa0994999 --- /dev/null +++ b/zig/tests/cuda/ling_mla_real.zig @@ -0,0 +1,139 @@ +//! Ling layer 7's MLA on real inputs (an HF fixture): absorb, rope, cache, attention, expand, gate vs HF and float64. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const qr = @import("quant_ref_tests.zig"); +const glm = @import("glm_cuda"); +const Gpu = check.Gpu; +const Rig = qr.RigT; +const bf = qr.helpers.bfOf; +const latent = glm.latent; +const Set = glm.glue.Set; + +const lw = 512; +const nope = 128; +const dv = 128; + +/// The fixture's dims and device buffers: q [t][h][192], lat [t][576] (raw k_rope), wk/wv [h][128][512], g fp32 [t][h]. +const Layer = struct { t: usize, h: usize, qk: usize, kw: usize, q: u64, lat: u64, wk: u64, wv: u64, g: u64, freq: u64 }; + +/// Device buffers one pass writes: the roped cache, qa, latent output, expanded and gated rows. +const Work = struct { cache: u64, qa: u64, ol: u64, vn: u64 }; + +fn take(comptime T: type, a: std.mem.Allocator, bytes: []const u8, at: *usize, n: usize) ![]T { + const out = try a.alloc(T, n); + @memcpy(std.mem.sliceAsBytes(out), bytes[at.*..][0 .. n * @sizeOf(T)]); + at.* += n * @sizeOf(T); + return out; +} + +/// Rows lo .. hi through absorb, rope, cache write, attention and expand at positions lo ..; out at vn row lo. +fn pass(r: *Rig, set: Set, l: Layer, w: Work, lo: usize, hi: usize) !void { + const n = hi - lo; + const p0: i32 = @intCast(lo); + const pos = try r.dev(std.mem.asBytes(&p0)); + const qa = w.qa + lo * l.h * l.kw * 2; + try latent.absorbQStrided(set, r.s, l.q + lo * l.h * l.qk * 2, l.qk, l.wk, qa, l.kw, n, l.h, nope, lw); + try latent.rope(set, r.s, .{ .x = l.q + (lo * l.h * l.qk + nope) * 2, .x_row = l.h * l.qk, .x_head = l.qk, .out = qa + lw * 2, .o_row = l.h * l.kw, .o_head = l.kw }, n, l.h, pos, l.freq); + const k = try r.zeros(n * l.kw * 2); + const zero: i32 = 0; + try latent.latentWrite(set, r.s, l.lat + lo * l.kw * 2, l.kw, k, try r.dev(std.mem.asBytes(&zero)), n, l.kw); + try latent.rope(set, r.s, .{ .x = k + lw * 2, .x_row = l.kw, .x_head = 0, .out = k + lw * 2, .o_row = l.kw, .o_head = 0 }, n, 1, pos, l.freq); + try latent.latentWrite(set, r.s, k, l.kw, w.cache, pos, n, l.kw); + const nch = latent.chunksFor(hi); + const p: latent.Partials = .{ .po = try r.zeros(nch * n * l.h * lw * 4), .pm = try r.zeros(nch * n * l.h * 4), .pl = try r.zeros(nch * n * l.h * 4) }; + const scale: f32 = @floatCast(1 / @sqrt(@as(f64, @floatFromInt(l.qk)))); + try latent.attentionRoped(set, r.s, qa, w.cache, pos, p, w.ol + lo * l.h * lw * 2, try r.zeros(4), n, l.h, lw, nch, scale); + try latent.expandV(set, r.s, w.ol + lo * l.h * lw * 2, l.wv, w.vn + lo * l.h * dv * 2, n, l.h, dv, lw); +} + +fn work(r: *Rig, l: Layer) !Work { + return .{ .cache = try r.zeros(l.t * l.kw * 2), .qa = try r.zeros(l.t * l.h * l.kw * 2), .ol = try r.zeros(l.t * l.h * lw * 2), .vn = try r.zeros(l.t * l.h * dv * 2) }; +} + +/// Relative L2 of bf16 `got` against fp32 `want`, and the largest difference over want's rms. +fn distance(got: []const u8, want: []const f32) [2]f64 { + var dd: f64 = 0; + var ww: f64 = 0; + var most: f64 = 0; + for (want, 0..) |wv, i| { + const g = bf(std.mem.bytesToValue(u16, got[2 * i ..][0..2])); + dd += (g - wv) * (g - wv); + ww += @as(f64, wv) * wv; + most = @max(most, @abs(g - wv)); + } + return .{ @sqrt(dd / ww), most / @sqrt(ww / @as(f64, @floatFromInt(want.len))) }; +} + +fn report(name: []const u8, got: []const u8, want: []const f32, bound: f64) !void { + const dist = distance(got, want); + std.debug.print("RESULT {s}: rel L2 {e:.3}, max |diff| {e:.3} of rms\n", .{ name, dist[0], dist[1] }); + try check.expect(dist[0] <= bound, "{s}: rel L2 within {e}", .{ name, bound }); +} + +/// `ling-mla-real `: 512-token prompt vs HF and float64; decode rows, a window and a cut equal its bytes. +pub fn run(gpu: Gpu, path: []const u8) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(7); + var r: Rig = .{ .gpu = gpu, .a = a, .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + } + var mod = try cuda.Module.load(gpu.d, cuda.kernels.glm_glue); + defer mod.unload(); + const ks = try glm.glue.Kernels.load(mod); + const set: Set = .{ .glue = &ks }; + const bytes = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, a, .limited(256 << 20)); + var at: usize = 0; + const dims = try take(i32, a, bytes, &at, 4); + const t: usize = @intCast(dims[0]); + const h: usize = @intCast(dims[1]); + const qk: usize = @intCast(dims[2]); + const kw: usize = @intCast(dims[3]); + const q = try take(u16, a, bytes, &at, t * h * qk); + const lat = try take(u16, a, bytes, &at, t * kw); + const wk = try take(u16, a, bytes, &at, h * nope * lw); + const wv = try take(u16, a, bytes, &at, h * dv * lw); + const g = try take(f32, a, bytes, &at, t * h); + const freq = try take(f32, a, bytes, &at, latent.rope_width / 2); + var refs: [4][]f32 = undefined; + for (&refs) |*x| x.* = try take(f32, a, bytes, &at, t * h * dv); + try check.expect(at == bytes.len, "fixture: {d} bytes read of {d}", .{ at, bytes.len }); + const ours = latent.ropeFreqs(6000000); + var fd: f64 = 0; + for (ours, freq) |o, f| fd = @max(fd, @abs(@as(f64, o) - f) / f); + std.debug.print("RESULT rope inv_freq: ropeFreqs vs HF's fp32 table, worst relative {e:.3}\n", .{fd}); + const l: Layer = .{ + .t = t, .h = h, .qk = qk, .kw = kw, + .q = try r.dev(std.mem.sliceAsBytes(q)), .lat = try r.dev(std.mem.sliceAsBytes(lat)), + .wk = try r.dev(std.mem.sliceAsBytes(wk)), .wv = try r.dev(std.mem.sliceAsBytes(wv)), + .g = try r.dev(std.mem.sliceAsBytes(g)), .freq = try r.dev(std.mem.sliceAsBytes(freq)), + }; + const full = try work(&r, l); + try pass(&r, set, l, full, 0, t); + const pre = try r.back(full.vn, t * h * dv * 2); + try latent.headGate(set, r.s, full.vn, l.g, h, t, h, dv); + const post = try r.back(full.vn, t * h * dv * 2); + try report("Ling L7 attention pre-gate vs HF bf16 (eager)", pre, refs[0], 2e-2); + try report("Ling L7 attention post-gate (dense input) vs HF bf16 (eager)", post, refs[1], 2e-2); + try report("Ling L7 attention pre-gate vs float64 on the kernel inputs", pre, refs[2], 1e-2); + try report("Ling L7 attention post-gate vs float64 on the kernel inputs", post, refs[3], 1e-2); + const row = h * dv * 2; + const dec = try work(&r, l); + try pass(&r, set, l, dec, 0, 400); + for (400..432) |i| try pass(&r, set, l, dec, i, i + 1); + try pass(&r, set, l, dec, 432, 440); + try latent.headGate(set, r.s, dec.vn, l.g, h, 440, h, dv); + try check.sameBytes("Ling L7: 400-row prompt, 32 decode rows, an 8-row window equal the 512-row prompt", try r.back(dec.vn, 440 * row), post[0 .. 440 * row]); + const cut = try work(&r, l); + try pass(&r, set, l, cut, 0, 300); + try pass(&r, set, l, cut, 300, t); + try latent.headGate(set, r.s, cut.vn, l.g, h, t, h, dv); + try check.sameBytes("Ling L7: the prompt cut at 300 equals uncut", try r.back(cut.vn, t * row), post); + check.pass("ling-mla-real: layer 7 within rel L2 of HF and float64; decode, verify window and cut prompt equal the prompt", .{}); +} diff --git a/zig/tests/cuda/ling_ref_mla.zig b/zig/tests/cuda/ling_ref_mla.zig new file mode 100644 index 0000000000..fd855256be --- /dev/null +++ b/zig/tests/cuda/ling_ref_mla.zig @@ -0,0 +1,267 @@ +//! Ling's MLA on GLM's kernels vs float64: 128-wide absorb and expand, rope, head gate, the 576-wide roped attention. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const qr = @import("quant_ref_tests.zig"); +const glm = @import("glm_cuda"); +const Gpu = check.Gpu; +const Rig = qr.RigT; +const bf = qr.helpers.bfOf; +const latent = glm.latent; + +const heads = 64; +const lw = 512; +const rw = latent.rope_width; +const kw = lw + rw; +const nope = 128; +const qk = nope + rw; +const dv = 128; +const u = std.math.pow(f64, 2, -24); + +/// Worst |got - want| / tol: `want` is `width` of every `stride` bf16 elements of `raw` from `off`. +fn within(r: *Rig, name: []const u8, raw: []const u8, stride: usize, off: usize, width: usize, want: []const f64, tol: []const f64) !void { + var most: f64 = 0; + for (want, tol, 0..) |w, t, i| { + const at = (i / width) * stride + off + i % width; + const q = @abs(bf(std.mem.bytesToValue(u16, raw[2 * at ..][0..2])) - w) / t; + most = if (std.math.isNan(q)) std.math.inf(f64) else @max(most, q); + } + r.worst = @max(r.worst, most); + std.debug.print("RESULT {s}: worst {e:.3} of its float64 tolerance\n", .{ name, most }); + try check.expect(most <= 1, "{s} within its float64 tolerance", .{name}); +} + +fn f64s(r: *Rig, x: []const u16) ![]f64 { + const out = try r.a.alloc(f64, x.len); + for (out, x) |*o, v| o.* = bf(v); + return out; +} + +/// The bytes of rows lo .. hi of a [rows][n] bf16 output equal `got`'s. +fn sameSpan(name: []const u8, all: []const u8, got: []const u8, lo: usize, n: usize) !void { + try check.sameBytes(name, got, all[lo * n * 2 ..][0..got.len]); +} + +/// absorbQStrided (192-wide q heads into 576-wide qa rows) and expandV at dv 128 vs float64; row counts' bits. +fn projections(r: *Rig, set: glm.glue.Set) !void { + const rows = 70; + const q = try r.bfs(rows * heads * qk, 1.0); + const wk = try r.bfs(heads * nope * lw, 0.05); + const dq = try r.dev(std.mem.sliceAsBytes(q)); + const dwk = try r.dev(std.mem.sliceAsBytes(wk)); + const qa = try r.zeros(rows * heads * kw * 2); + try latent.absorbQStrided(set, r.s, dq, qk, dwk, qa, kw, rows, heads, nope, lw); + var want = try r.a.alloc(f64, rows * heads * lw); + var tol = try r.a.alloc(f64, rows * heads * lw); + const qd = try f64s(r, q); + const wkd = try f64s(r, wk); + for (0..rows) |row| for (0..heads) |h| for (0..lw) |c| { + var s: f64 = 0; + var abs: f64 = 0; + for (0..nope) |k| { + const p = qd[(row * heads + h) * qk + k] * wkd[(h * nope + k) * lw + c]; + s += p; + abs += @abs(p); + } + want[(row * heads + h) * lw + c] = s; + tol[(row * heads + h) * lw + c] = @abs(s) * std.math.pow(f64, 2, -8) + abs * (nope + 8) * u + 1e-30; + }; + const all = try r.back(qa, rows * heads * kw * 2); + try within(r, "absorb_q 128 (Ling): 70 rows x 64 heads, q stride 192, qa stride 576", all, kw, 0, lw, want, tol); + for ([_][2]usize{ .{ 0, 1 }, .{ 37, 42 }, .{ 3, 67 } }) |span| { + const n = span[1] - span[0]; + const o = try r.zeros(n * heads * kw * 2); + try latent.absorbQStrided(set, r.s, dq + span[0] * heads * qk * 2, qk, dwk, o, kw, n, heads, nope, lw); + const got = try r.back(o, n * heads * kw * 2); + for (0..n * heads) |i| try check.sameBytes("absorb_q 128: rows of a narrower call equal the 70-row call's", got[i * kw * 2 ..][0 .. lw * 2], all[(span[0] * heads + i) * kw * 2 ..][0 .. lw * 2]); + } + const ol = try r.bfs(rows * heads * lw, 1.0); + const wv = try r.bfs(heads * dv * lw, 0.05); + const dol = try r.dev(std.mem.sliceAsBytes(ol)); + const dwv = try r.dev(std.mem.sliceAsBytes(wv)); + const vn = try r.zeros(rows * heads * dv * 2); + try latent.expandV(set, r.s, dol, dwv, vn, rows, heads, dv, lw); + want = try r.a.alloc(f64, rows * heads * dv); + tol = try r.a.alloc(f64, rows * heads * dv); + const old = try f64s(r, ol); + const wvd = try f64s(r, wv); + for (0..rows) |row| for (0..heads) |h| for (0..dv) |n| { + var s: f64 = 0; + var abs: f64 = 0; + for (old[(row * heads + h) * lw ..][0..lw], wvd[(h * dv + n) * lw ..][0..lw]) |a, b| { + s += a * b; + abs += @abs(a * b); + } + want[(row * heads + h) * dv + n] = s; + tol[(row * heads + h) * dv + n] = @abs(s) * std.math.pow(f64, 2, -8) + abs * (lw + 8) * u + 1e-30; + }; + const vall = try r.back(vn, rows * heads * dv * 2); + try within(r, "expand_v 128 (Ling): 70 rows x 64 heads", vall, dv, 0, dv, want, tol); + for ([_][2]usize{ .{ 0, 1 }, .{ 37, 42 }, .{ 3, 67 } }) |span| { + const n = span[1] - span[0]; + const o = try r.zeros(n * heads * dv * 2); + try latent.expandV(set, r.s, dol + span[0] * heads * lw * 2, dwv, o, n, heads, dv, lw); + try sameSpan("expand_v 128: rows of a narrower call equal the 70-row call's", vall, try r.back(o, n * heads * dv * 2), span[0], heads * dv); + } + check.pass("absorb_q / expand_v at 128-wide heads: within float64 tolerance, every row count's bits equal", .{}); +} + +/// rope (q heads into qa, cache rows in place) at two positions vs float64 on fp32 angles; head gate vs float64. +fn ropeGate(r: *Rig, set: glm.glue.Set) !void { + const rows = 40; + const freqs = latent.ropeFreqs(6000000); + const dfreq = try r.dev(std.mem.asBytes(&freqs)); + const q = try r.bfs(rows * heads * qk, 1.0); + const dq = try r.dev(std.mem.sliceAsBytes(q)); + const qd = try f64s(r, q); + var name: [128]u8 = undefined; + for ([_]i32{ 1000, 261000 }) |p0| { + const at = try r.dev(std.mem.asBytes(&p0)); + const qa = try r.zeros(rows * heads * kw * 2); + const v: latent.RopeRows = .{ .x = dq + nope * 2, .x_row = heads * qk, .x_head = qk, .out = qa + lw * 2, .o_row = heads * kw, .o_head = kw }; + try latent.rope(set, r.s, v, rows, heads, at, dfreq); + const want = try r.a.alloc(f64, rows * heads * rw); + const tol = try r.a.alloc(f64, rows * heads * rw); + for (0..rows) |row| for (0..heads) |h| for (0..rw / 2) |i| { + const ang: f64 = freqs[i] * @as(f32, @floatFromInt(p0 + @as(i32, @intCast(row)))); + const x0 = qd[(row * heads + h) * qk + nope + 2 * i]; + const x1 = qd[(row * heads + h) * qk + nope + 2 * i + 1]; + const c = @cos(ang); + const s = @sin(ang); + const at2 = (row * heads + h) * rw + 2 * i; + want[at2] = x0 * c - x1 * s; + want[at2 + 1] = x1 * c + x0 * s; + for (0..2) |e| tol[at2 + e] = @abs(want[at2 + e]) * std.math.pow(f64, 2, -8) + (@abs(x0) + @abs(x1)) * 8 * u + 1e-30; + }; + const all = try r.back(qa, rows * heads * kw * 2); + try within(r, try std.fmt.bufPrint(&name, "rope: 40 rows x 64 heads at {d} (GPT-J pairs, fp32 angles)", .{p0}), all, kw, lw, rw, want, tol); + const p1: i32 = p0 + 23; + const one = try r.zeros(heads * kw * 2); + const v1: latent.RopeRows = .{ .x = dq + (23 * heads * qk + nope) * 2, .x_row = heads * qk, .x_head = qk, .out = one + lw * 2, .o_row = heads * kw, .o_head = kw }; + try latent.rope(set, r.s, v1, 1, heads, try r.dev(std.mem.asBytes(&p1)), dfreq); + const got = try r.back(one, heads * kw * 2); + for (0..heads) |h| try check.sameBytes("rope: row 23 alone equals its row", got[(h * kw + lw) * 2 ..][0 .. rw * 2], all[((23 * heads + h) * kw + lw) * 2 ..][0 .. rw * 2]); + // a cache row's k_rope in place equals q's rope of the same values at the same position + const k = try r.dev(all[(23 * heads * kw) * 2 ..][0 .. kw * 2]); + @memcpy(all[(23 * heads * kw + lw) * 2 ..][0 .. rw * 2], std.mem.sliceAsBytes(q[(23 * heads) * qk + nope ..][0..rw])); + const kin = try r.dev(all[(23 * heads * kw) * 2 ..][0 .. kw * 2]); + try latent.rope(set, r.s, .{ .x = kin + lw * 2, .x_row = kw, .x_head = 0, .out = kin + lw * 2, .o_row = kw, .o_head = 0 }, 1, 1, try r.dev(std.mem.asBytes(&p1)), dfreq); + try check.sameBytes("rope: a cache row in place equals the q rope", try r.back(kin, kw * 2), try r.back(k, kw * 2)); + } + const o = try r.bfs(rows * heads * dv, 1.0); + const g = try r.a.alloc(f32, rows * (heads + 3)); + for (g) |*x| x.* = (r.rng.float(f32) * 2 - 1) * 12; + const dov = try r.dev(std.mem.sliceAsBytes(o)); + try latent.headGate(set, r.s, dov, try r.dev(std.mem.sliceAsBytes(g)), heads + 3, rows, heads, dv); + const want = try r.a.alloc(f64, rows * heads * dv); + const tol = try r.a.alloc(f64, rows * heads * dv); + for (want, tol, 0..) |*w, *t, i| { + const sg = 1 / (1 + @exp(-@as(f64, g[(i / (heads * dv)) * (heads + 3) + (i / dv) % heads]))); + w.* = bf(o[i]) * sg; + t.* = @abs(w.*) * std.math.pow(f64, 2, -8) + @abs(bf(o[i])) * 8 * u + 1e-30; + } + try within(r, "head gate: 40 rows x 64 heads x 128, g stride 67", try r.back(dov, rows * heads * dv * 2), dv, 0, dv, want, tol); + check.pass("rope and head gate: within float64 tolerance, a row alone and a cache row in place equal", .{}); +} + +/// float64 attention over 576-wide rows (scores) and their first 512 (values); bound as glm_ref_attn's, 576 + 8 terms. +fn reference(r: *Rig, q: []const f64, cd: []const f64, pos: usize, rows: usize, scale: f64) !struct { want: []f64, tol: []f64 } { + const want = try r.a.alloc(f64, rows * heads * lw); + const tol = try r.a.alloc(f64, rows * heads * lw); + const s = try r.a.alloc(f64, pos + rows); + const ds = try r.a.alloc(f64, pos + rows); + const o = try r.a.alloc(f64, lw); + const spread = try r.a.alloc(f64, lw); + for (0..rows) |row| for (0..heads) |h| { + const qr_ = q[(row * heads + h) * kw ..][0..kw]; + const n = pos + row + 1; + var most = -std.math.inf(f64); + for (0..n) |key| { + var dot: f64 = 0; + var abs: f64 = 0; + for (qr_, cd[key * kw ..][0..kw]) |a, b| { + dot += a * b; + abs += @abs(a * b); + } + s[key] = dot * scale; + ds[key] = abs * scale * (kw + 8) * u; + most = @max(most, s[key]); + } + @memset(o, 0); + @memset(spread, 0); + var l: f64 = 0; + for (0..n) |key| { + const p = @exp(s[key] - most); + l += p; + for (o, spread, cd[key * kw ..][0..lw]) |*od, *sd, v| { + od.* += p * v; + sd.* += p * @abs(v) * (std.math.pow(f64, 2, -8) + ds[key]); + } + } + const base = (row * heads + h) * lw; + for (0..lw) |d| { + want[base + d] = o[d] / l; + tol[base + d] = @abs(o[d] / l) * std.math.pow(f64, 2, -8) + spread[d] / l + 1e-30; + } + }; + return .{ .want = want, .tol = tol }; +} + +/// attentionRoped of `rows` rows at `pos` (rows of qa from `first`) into a fresh [rows][heads][lw] output. +fn attend(r: *Rig, set: glm.glue.Set, qa: u64, cache: u64, pos: usize, first: usize, rows: usize, scale: f32) ![]u8 { + const nch = latent.chunksFor(pos + rows); + const out = try r.zeros(rows * heads * lw * 2); + const p0: i32 = @intCast(pos); + const p: latent.Partials = .{ .po = try r.zeros(nch * rows * heads * lw * 4), .pm = try r.zeros(nch * rows * heads * 4), .pl = try r.zeros(nch * rows * heads * 4) }; + try latent.attentionRoped(set, r.s, qa + first * heads * kw * 2, cache, try r.dev(std.mem.asBytes(&p0)), p, out, try r.zeros(4), rows, heads, lw, nch, scale); + return r.back(out, rows * heads * lw * 2); +} + +/// Roped attention: 60 rows at 1,500 (keys cross 1,536) vs float64; one-row decode, a verify window and a cut prompt. +fn attention(r: *Rig, set: glm.glue.Set) !void { + const pos = 1500; + const rows = 60; + const scale: f32 = @floatCast(1 / @sqrt(@as(f64, qk))); + const q = try r.bfs(rows * heads * kw, 1.0); + const c = try r.bfs((pos + rows) * kw, 1.0); + const dq = try r.dev(std.mem.sliceAsBytes(q)); + const dc = try r.dev(std.mem.sliceAsBytes(c)); + const all = try attend(r, set, dq, dc, pos, 0, rows, scale); + const ref64 = try reference(r, try f64s(r, q), try f64s(r, c), pos, rows, scale); + try within(r, "roped attention (Ling): 60 rows at 1500, 64 heads, 576-wide keys, 16-head blocks", all, lw, 0, lw, ref64.want, ref64.tol); + const n = heads * lw; + var name: [128]u8 = undefined; + for ([_]usize{ 0, 35, 36, 59 }) |row| { + const one = try attend(r, set, dq, dc, pos + row, row, 1, scale); + try sameSpan(try std.fmt.bufPrint(&name, "roped attention: row {d} alone (decode) equals its row", .{row}), all, one, row, n); + } + const win = try attend(r, set, dq, dc, pos + 30, 30, 8, scale); + try sameSpan("roped attention: an 8-row verify window across 1,536 equals serial rows", all, win, 30, n); + const head = try attend(r, set, dq, dc, pos, 0, 36, scale); + const tail = try attend(r, set, dq, dc, pos + 36, 36, rows - 36, scale); + try sameSpan("roped attention: a prompt cut at 36 (key 1,536): first part equals uncut", all, head, 0, n); + try sameSpan("roped attention: a prompt cut at 36 (key 1,536): second part equals uncut", all, tail, 36, n); + check.pass("roped attention over {d} keys: within float64 tolerance; decode rows, verify window, cut prompt equal", .{pos + rows}); +} + +/// `ling-mla`: the checks above on synthetic inputs (seeded). +pub fn run(gpu: Gpu) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(731); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + } + var mod = try cuda.Module.load(gpu.d, cuda.kernels.glm_glue); + defer mod.unload(); + const ks = try glm.glue.Kernels.load(mod); + const set: glm.glue.Set = .{ .glue = &ks }; + try projections(&r, set); + try ropeGate(&r, set); + try attention(&r, set); + check.pass("ling-mla: every check within {e:.3} of its float64 tolerance", .{r.worst}); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index aaf69b0300..da8e589563 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -95,6 +95,9 @@ const usage = \\ kernel-contract the registered kernels through the contract's gates, each planted control red (needs -Dnvcc) \\ kernel-contract-limits DIR each registered launch against sm_80/86/89/120/121 from the fatbins in DIR \\ glm-ref GLM projections (matmuls, router, attention, indexer, hc) against float64 (synthetic) + \\ ling-mla Ling's MLA (128-wide absorb/expand, rope, gate, 576-wide attention) vs float64 (synthetic) + \\ ling-mla-real Ling layer 7's MLA on real inputs vs HF and float64; decode, window and cut bytes + \\ glm-mla-same GLM's MLA launches on this build vs a 6380c8da glm_glue fatbin, byte for byte \\ glm-lanes-cover copy drafts, cancels, discards, eos and max_new on one lanes.Engine across requests \\ glm-lanes-concurrent concurrent streams on one lanes.Engine, each against its solo run \\ glm-shared-check shared verify windows against lone twins, logits and states byte for byte @@ -166,6 +169,9 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "nvfp4-experts")) return nvfp4_experts_tests.experts(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "quant-ref")) return quant_ref_tests.run(gpu); if (std.mem.eql(u8, cmd, "glm-ref")) return glm_ref_tests.run(gpu); + if (std.mem.eql(u8, cmd, "ling-mla")) return @import("ling_ref_mla.zig").run(gpu); + if (std.mem.eql(u8, cmd, "ling-mla-real")) return @import("ling_mla_real.zig").run(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "glm-mla-same")) return @import("glm_mla_same.zig").run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "modelopt")) return modelopt_tests.run(gpu); if (std.mem.eql(u8, cmd, "exlr8")) return exlr8_tests.run(gpu); if (std.mem.eql(u8, cmd, "exlr8-dense")) return exlr8_dense_tests.run(gpu, if (rest.len > 0) rest[0] else null); From ef8f12f35a1a28ab185f74ae83efeab8d8b560bd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 15 Oct 2026 09:01:50 +0200 Subject: [PATCH 06/34] glm: KDA serves Ling-3.1 as is (real layer 0 vs fla and float64); an opt-in 16-row WY chunked prompt path - glm-kda-ling: glm_kda's fused chain, 16-row windows, wide prompt and chunked prompt on Ling-3.1 layer 0 (q/k/v convs as one q|k|v conv): rel L2 3.6e-3 vs fla, within float64's bound; replays == chain bits. - glm_kda_chunk.cu: new entry tf_glm_kda_chunk over prep's saves (kda.chunked, kda.chunkReplay). Opt-in only: GLM stays on the row path; glm_kda.cu and its fatbin are unchanged. Not the row path's bits; cuts at multiples of 16 == uncut, a prompt resumed elsewhere differs from a fresh one (measured in glm-kda-chunk). - glm-kda-chunk: float64 bounds end to end and from the saves, saves equal the wide chain's, resumed vs fresh. - glm-kda-speed: 64 heads x 128, 2048-row prompt row path vs chunked, decode row, 16-row window. --- zig/build/cuda.zig | 1 + zig/kernels/cuda/glm_kda_chunk.cu | 237 +++++++++++++++++++++++++ zig/src/cuda/kernels.zig | 1 + zig/src/families/glm/cuda_kda.zig | 84 ++++++--- zig/tests/cuda/glm_kda_chunk_tests.zig | 124 +++++++++++++ zig/tests/cuda/glm_kda_f64.zig | 201 +++++++++++++++++++++ zig/tests/cuda/glm_kda_ling.zig | 216 ++++++++++++++++++++++ zig/tests/cuda/glm_kda_rig.zig | 117 ++++++++++++ zig/tests/cuda/glm_kda_speed.zig | 76 ++++++++ zig/tests/cuda/main.zig | 9 + 10 files changed, 1045 insertions(+), 21 deletions(-) create mode 100644 zig/kernels/cuda/glm_kda_chunk.cu create mode 100644 zig/tests/cuda/glm_kda_chunk_tests.zig create mode 100644 zig/tests/cuda/glm_kda_f64.zig create mode 100644 zig/tests/cuda/glm_kda_ling.zig create mode 100644 zig/tests/cuda/glm_kda_rig.zig create mode 100644 zig/tests/cuda/glm_kda_speed.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 90accb6244..88b8744638 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -49,6 +49,7 @@ const kernels = [_]Kernel{ .{ .name = "prompt16", .flags = &.{"-O3"} }, // nvfp4/prompt.cu's FP4 and FP8 tile 4, tensorfold_nvfp4_prompt_v1 .{ .name = "modelopt_pack", .flags = &.{"-O3"} }, // ours: ModelOpt tensors into the lane, prompt and expert layouts .{ .name = "glm_kda", .flags = &.{ "-O3", "--fmad=false" } }, // GLM: the Python family's glm5_next/cuda/kda.cu + .{ .name = "glm_kda_chunk", .flags = &.{ "-O3", "--fmad=false" } }, // ours: glm_kda's prompt recurrence in 16-row WY chunks .{ .name = "glm_glue", .flags = &.{"-O3"} }, // ours: the Python family's Triton glue by hand, its PTX op for op .{ .name = "glm_ops", .flags = &.{"-O3"} }, // ours: GLM's integer work the Python family did in torch (the DSA token lists) .{ .name = "glm_quant", .flags = torch_ops }, // ours: GLM's draft-only weight quantizers, torch at load in the Python family diff --git a/zig/kernels/cuda/glm_kda_chunk.cu b/zig/kernels/cuda/glm_kda_chunk.cu new file mode 100644 index 0000000000..8f705777cb --- /dev/null +++ b/zig/kernels/cuda/glm_kda_chunk.cu @@ -0,0 +1,237 @@ +// GLM's KDA prompt recurrence in 16-row WY chunks over glm_kda.cu's prep saves; opt-in, -O3 --fmad=false, _rn ops. + +// A block a head, chunks from the call's first row: calls cut at multiples of 16 rows keep one call's bits. + +// Not the serial row path's bits (other sums): a prompt resumed off a multiple of 16 differs from a fresh one. + +// Thread t holds S rows 4 (t / 8) .. + 3 (value), keys 4 (t % 8) + 32 e + 0..3; every sum runs in one fixed order. + +#include +#include +#include +namespace { + +constexpr int D = 128; // key and value dims a head +constexpr int C = 16; // rows a chunk +constexpr int NT = 256; // threads a block +constexpr int PARTS = 8; // threads sharing a value row (key parts) +constexpr int R = 4; // value rows a thread +constexpr int KG = 4; // float4 key groups a thread +constexpr unsigned FULL = 0xffffffffu; + +// N per-lane partials summed across the warp in one fixed order; lanes 32 j / N .. 32 (j + 1) / N - 1 hold total j. +template +__device__ __forceinline__ float spread(float (&v)[N], int lane) { +#pragma unroll + for (int m = 16, h = N / 2; m > 0; m >>= 1, h >>= 1) { + if (h >= 1) { + const bool up = lane & m; +#pragma unroll + for (int i = 0; i < h; ++i) { + const float give = up ? v[i] : v[h + i], keep = up ? v[h + i] : v[i]; + v[i] = __fadd_rn(keep, __shfl_xor_sync(FULL, give, m)); + } + } else { + v[0] = __fadd_rn(v[0], __shfl_xor_sync(FULL, v[0], m)); + } + } + return v[0]; +} + +struct Chunk { + float q[C][D]; // q, then its decay from the chunk's start times q + float k[C][D]; // k, then its decay to the chunk's end times k + float g[C][D]; // decay factors, then (decayed q) S + float w[C][D]; // beta (decay from the start) k, then W = T of that + float v[C][D]; // v, then beta v, then U = T (beta v), then U - W S + float l[C][C]; // beta_i k_i . (decay j+1 .. i) k_j below the diagonal + float m[C][C]; // q_i . (decay j+1 .. i) k_j on and below it + float gl[D]; // the chunk's whole decay + float beta[C]; +}; + +// out[i][v] (sub: minus) x_i . S row v: a thread's keys in order, 8 parts by xor butterfly (same bits in every lane). +__device__ __forceinline__ void rows_dot(const float (*x)[D], const float (&s)[R][KG][4], int p, int vg, bool sub, float (*out)[D]) { +#pragma unroll + for (int i = 0; i < C; ++i) { + float acc[R] = {0.f, 0.f, 0.f, 0.f}; +#pragma unroll + for (int e = 0; e < KG; ++e) { + const float4 a = *reinterpret_cast(&x[i][4 * p + 32 * e]); +#pragma unroll + for (int r = 0; r < R; ++r) { + acc[r] = __fmaf_rn(a.x, s[r][e][0], acc[r]); + acc[r] = __fmaf_rn(a.y, s[r][e][1], acc[r]); + acc[r] = __fmaf_rn(a.z, s[r][e][2], acc[r]); + acc[r] = __fmaf_rn(a.w, s[r][e][3], acc[r]); + } + } +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int m = 1; m < PARTS; m <<= 1) acc[r] = __fadd_rn(acc[r], __shfl_xor_sync(FULL, acc[r], m)); + float mine = acc[0]; +#pragma unroll + for (int r = 1; r < R; ++r) mine = p == r ? acc[r] : mine; + if (p < R) out[i][R * vg + p] = sub ? __fsub_rn(out[i][R * vg + p], mine) : mine; + } +} + +// Block (head); q_in null: the state alone (a chunked replay of saved rows, the same state bits), y_out unwritten. +__device__ __forceinline__ void chunk_body(int H, const float* __restrict__ state_in, const float* __restrict__ q_in, + const float* __restrict__ k_save, const __nv_bfloat16* __restrict__ v_save, + const float* __restrict__ g_save, const float* __restrict__ b_save, int rows, + __nv_bfloat16* __restrict__ y_out, float* __restrict__ state_out) { + __shared__ __align__(16) Chunk sh; + const int t = threadIdx.x, warp = t >> 5, lane = t & 31; + const int h = blockIdx.x; + const bool out = q_in != nullptr; + const int vg = t / PARTS, p = t % PARTS; + const size_t sbase = (size_t)h * D * D + (size_t)R * vg * D + 4 * p; + float s[R][KG][4]; +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int e = 0; e < KG; ++e) { + const float4 x = *reinterpret_cast(state_in + sbase + r * D + 32 * e); + s[r][e][0] = x.x, s[r][e][1] = x.y, s[r][e][2] = x.z, s[r][e][3] = x.w; + } + for (int c0 = 0; c0 < rows; c0 += C) { + const int n = min(C, rows - c0); + constexpr int LD = C * D / NT; + float lk[LD], lq[LD], lg[LD], lv[LD]; // every load of the chunk issued before any store +#pragma unroll + for (int e = 0; e < LD; ++e) { // the chunk's saved rows; rows past the call: k, q, v 0, decay 1, beta 0 + const int x = t + e * NT, i = x / D; + const size_t at = ((size_t)(c0 + i) * H + h) * D + x % D; + lk[e] = i < n ? k_save[at] : 0.f; + lq[e] = i < n && out ? q_in[at] : 0.f; + lg[e] = i < n ? g_save[at] : 1.f; + lv[e] = i < n ? __bfloat162float(v_save[at]) : 0.f; + } +#pragma unroll + for (int e = 0; e < LD; ++e) { + const int i = (t + e * NT) / D, d = (t + e * NT) % D; + sh.k[i][d] = lk[e], sh.q[i][d] = lq[e], sh.g[i][d] = lg[e], sh.v[i][d] = lv[e]; + } + if (t < C) sh.beta[t] = t < n ? b_save[(c0 + t) * H + h] : 0.f; + { + const int i = t / C, j = t % C; + if (j > i) sh.l[i][j] = 0.f, sh.m[i][j] = 0.f; + if (j == i) sh.l[i][j] = 0.f; + } + __syncthreads(); +#pragma unroll + for (int jj = 0; jj < 2; ++jj) { // columns j = warp and 15 - warp: e = prod a_(j+1 .. i), a running product per channel + const int j = jj == 0 ? warp : C - 1 - warp; + float kj[4], ej[4], part[2 * C]; +#pragma unroll + for (int e = 0; e < 4; ++e) kj[e] = sh.k[j][lane + 32 * e], ej[e] = 1.f; +#pragma unroll + for (int i = 0; i < C; ++i) { + float a = 0.f, m = 0.f; + if (i >= j) { +#pragma unroll + for (int e = 0; e < 4; ++e) { + const int d = lane + 32 * e; + if (i > j) ej[e] = __fmul_rn(ej[e], sh.g[i][d]); + const float ke = __fmul_rn(kj[e], ej[e]); + a = __fmaf_rn(sh.k[i][d], ke, a); + m = __fmaf_rn(sh.q[i][d], ke, m); + } + } + part[i] = a, part[C + i] = m; + } + const float tot = spread<2 * C>(part, lane); // lane i: row i's a, lane 16 + i: row i's m + const int i = lane % C; + if (lane < C && j < i) sh.l[i][j] = __fmul_rn(sh.beta[i], tot); + if (lane >= C && j <= i) sh.m[i][j] = tot; + } + __syncthreads(); + if (t < D) { // a channel's prefix products (q, beta k) and suffix products (the keys' decay to the chunk's end) + float pre = 1.f; +#pragma unroll + for (int i = 0; i < C; ++i) { + pre = __fmul_rn(pre, sh.g[i][t]); + sh.w[i][t] = __fmul_rn(sh.beta[i], __fmul_rn(pre, sh.k[i][t])); + sh.q[i][t] = __fmul_rn(pre, sh.q[i][t]); + } + sh.gl[t] = pre; + float suf = 1.f; +#pragma unroll + for (int i = C - 1; i >= 0; --i) { + sh.k[i][t] = __fmul_rn(sh.k[i][t], suf); + suf = __fmul_rn(suf, sh.g[i][t]); + } + } else { +#pragma unroll + for (int i = 0; i < C; ++i) sh.v[i][t - D] = __fmul_rn(sh.beta[i], sh.v[i][t - D]); + } + __syncthreads(); + { // (I + L) X = R by forward substitution, a thread a column of [W | U] + float* col = t < D ? &sh.w[0][t] : &sh.v[0][t - D]; + float x[C]; +#pragma unroll + for (int i = 0; i < C; ++i) { + float acc = col[i * D]; +#pragma unroll + for (int j = 0; j < i; ++j) acc = __fmaf_rn(-sh.l[i][j], x[j], acc); + x[i] = acc; + } +#pragma unroll + for (int i = 0; i < C; ++i) col[i * D] = x[i]; + } + __syncthreads(); + rows_dot(sh.w, s, p, vg, true, sh.v); // v: U - W S + if (out) rows_dot(sh.q, s, p, vg, false, sh.g); // g: (decayed q) S + __syncthreads(); + if (out) { // row i's read-out (decayed q) S + sum_j<=i M_ij (U - W S)_j: value row t % 128, rows 8 (t / 128) .. + 7 + const int v = t % D; +#pragma unroll + for (int ii = 0; ii < C / 2; ++ii) { + const int i = C / 2 * (t / D) + ii; + float o = sh.g[i][v]; + for (int j = 0; j <= i; ++j) o = __fmaf_rn(sh.m[i][j], sh.v[j][v], o); + if (i < n) y_out[((size_t)(c0 + i) * H + h) * D + v] = __float2bfloat16_rn(o); + } + } +#pragma unroll + for (int e = 0; e < KG; ++e) { // S = S diag(chunk decay) + sum_j (U - W S)_j (decayed k_j)^T + const int k = 4 * p + 32 * e; + const float4 g4 = *reinterpret_cast(&sh.gl[k]); + float x[R][4]; +#pragma unroll + for (int r = 0; r < R; ++r) + x[r][0] = __fmul_rn(s[r][e][0], g4.x), x[r][1] = __fmul_rn(s[r][e][1], g4.y), x[r][2] = __fmul_rn(s[r][e][2], g4.z), x[r][3] = __fmul_rn(s[r][e][3], g4.w); +#pragma unroll + for (int j = 0; j < C; ++j) { + const float4 kd = *reinterpret_cast(&sh.k[j][k]); + const float4 vn = *reinterpret_cast(&sh.v[j][R * vg]); + const float u[R] = {vn.x, vn.y, vn.z, vn.w}; +#pragma unroll + for (int r = 0; r < R; ++r) + x[r][0] = __fmaf_rn(kd.x, u[r], x[r][0]), x[r][1] = __fmaf_rn(kd.y, u[r], x[r][1]), x[r][2] = __fmaf_rn(kd.z, u[r], x[r][2]), x[r][3] = __fmaf_rn(kd.w, u[r], x[r][3]); + } +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int c = 0; c < 4; ++c) s[r][e][c] = x[r][c]; + } + __syncthreads(); + } +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int e = 0; e < KG; ++e) + *reinterpret_cast(state_out + sbase + r * D + 32 * e) = make_float4(s[r][e][0], s[r][e][1], s[r][e][2], s[r][e][3]); +} + +} // namespace + +// glm_kda.cu's step entry, chunked: q, k, v, decay, beta rows [rows][H][128] from prep; y bf16, the state after. +extern "C" __global__ void __launch_bounds__(NT, 2) tf_glm_kda_chunk( + int H, const float* __restrict__ state_in, const float* __restrict__ q_in, const float* __restrict__ k_save, + const __nv_bfloat16* __restrict__ v_save, const float* __restrict__ g_save, const float* __restrict__ b_save, + int rows, __nv_bfloat16* __restrict__ y_out, float* __restrict__ state_out) { + chunk_body(H, state_in, q_in, k_save, v_save, g_save, b_save, rows, y_out, state_out); +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 18542b8cb4..28e1811a38 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -40,6 +40,7 @@ pub const glm_glue: []const u8 = if (available) &Blob("fatbin_glm_glue").bytes e pub const glm_ops: []const u8 = if (available) &Blob("fatbin_glm_ops").bytes else &.{}; pub const glm_quant: []const u8 = if (available) &Blob("fatbin_glm_quant").bytes else &.{}; pub const glm_kda: []const u8 = if (available) &Blob("fatbin_glm_kda").bytes else &.{}; +pub const glm_kda_chunk: []const u8 = if (available) &Blob("fatbin_glm_kda_chunk").bytes else &.{}; pub const torch_argmax: []const u8 = if (available) &Blob("fatbin_torch_argmax").bytes else &.{}; pub const torch_topk: []const u8 = if (available) &Blob("fatbin_torch_topk").bytes else &.{}; pub const torch_pointwise: []const u8 = if (available) &Blob("fatbin_torch_pointwise").bytes else &.{}; diff --git a/zig/src/families/glm/cuda_kda.zig b/zig/src/families/glm/cuda_kda.zig index cf2388e640..c5b58d28d1 100644 --- a/zig/src/families/glm/cuda_kda.zig +++ b/zig/src/families/glm/cuda_kda.zig @@ -54,18 +54,7 @@ pub const Wide = struct { q: u64, y: u64 }; pub fn chain(ks: *const Kernels, stream: cuda.Stream, c: Chain, rows: usize, s: Saves, wide: ?Wide) !void { const h: c_int = @intCast(c.heads); if (wide) |w| { - var a: cuda.Args = .{}; - a.add(h); - a.add(c.p); - a.add(@as(c_int, @intCast(c.p_stride))); - a.add(@as(c_int, @intCast(c.b_off))); - a.add(c.a); - a.add(@as(c_int, @intCast(c.a_stride))); - for ([_]u64{ c.conv_state, c.conv_w, c.a_log, c.dt_bias }) |v| a.add(v); - a.add(c.lower); - for ([_]u64{ w.q, s.k, s.v, s.g, s.b }) |v| a.add(v); - const grid: cuda.launch.Dim3 = .{ .x = @intCast(rows), .y = @intCast(c.heads) }; - try cuda.launch.launch(ks.prep, .{ .grid = grid, .block = .{ .x = 512 } }, stream, &a); + try prep(ks, stream, c, rows, s, w); var b: cuda.Args = .{}; b.add(h); for ([_]u64{ c.state_in, w.q, s.k, s.v, s.g, s.b }) |v| b.add(v); @@ -73,15 +62,7 @@ pub fn chain(ks: *const Kernels, stream: cuda.Stream, c: Chain, rows: usize, s: b.add(w.y); b.add(s.state_out); try cuda.launch.launch(ks.step, .{ .grid = .{ .x = @intCast(c.heads), .y = dv / 4 / 8 }, .block = .{ .x = 8 * 32 } }, stream, &b); - var o: cuda.Args = .{}; - o.add(h); - o.add(w.y); - o.add(c.g); - o.add(@as(c_int, @intCast(c.g_stride))); - o.add(c.norm_w); - o.add(c.eps); - o.add(s.out); - return cuda.launch.launch(ks.out, .{ .grid = grid, .block = .{ .x = dv } }, stream, &o); + return out(ks, stream, c, rows, s, w); } var a: cuda.Args = .{}; a.add(h); @@ -100,6 +81,67 @@ pub fn chain(ks: *const Kernels, stream: cuda.Stream, c: Chain, rows: usize, s: try cuda.launch.launch(ks.chain, .{ .grid = .{ .x = @intCast(c.heads) }, .block = .{ .x = 1024 } }, stream, &a); } +/// The wide chain's first kernel: conv, decay, beta and q/k norms of `rows` rows into w.q and the saves. +pub fn prep(ks: *const Kernels, stream: cuda.Stream, c: Chain, rows: usize, s: Saves, w: Wide) !void { + var a: cuda.Args = .{}; + a.add(@as(c_int, @intCast(c.heads))); + a.add(c.p); + a.add(@as(c_int, @intCast(c.p_stride))); + a.add(@as(c_int, @intCast(c.b_off))); + a.add(c.a); + a.add(@as(c_int, @intCast(c.a_stride))); + for ([_]u64{ c.conv_state, c.conv_w, c.a_log, c.dt_bias }) |v| a.add(v); + a.add(c.lower); + for ([_]u64{ w.q, s.k, s.v, s.g, s.b }) |v| a.add(v); + const grid: cuda.launch.Dim3 = .{ .x = @intCast(rows), .y = @intCast(c.heads) }; + try cuda.launch.launch(ks.prep, .{ .grid = grid, .block = .{ .x = 512 } }, stream, &a); +} + +/// The wide chain's last kernel: the gated RMSNorm of w.y into s.out. +pub fn out(ks: *const Kernels, stream: cuda.Stream, c: Chain, rows: usize, s: Saves, w: Wide) !void { + var o: cuda.Args = .{}; + o.add(@as(c_int, @intCast(c.heads))); + o.add(w.y); + o.add(c.g); + o.add(@as(c_int, @intCast(c.g_stride))); + o.add(c.norm_w); + o.add(c.eps); + o.add(s.out); + const grid: cuda.launch.Dim3 = .{ .x = @intCast(rows), .y = @intCast(c.heads) }; + return cuda.launch.launch(ks.out, .{ .grid = grid, .block = .{ .x = dv } }, stream, &o); +} + +/// glm_kda_chunk.cu's entry: the prompt recurrence in 16-row WY chunks over prep's saves. +pub const Chunked = struct { + f: cuda.Function, + + pub fn load(m: cuda.Module) !Chunked { + return .{ .f = try m.function("tf_glm_kda_chunk") }; + } +}; + +fn chunkLaunch(ck: Chunked, stream: cuda.Stream, heads: usize, state_in: u64, q: u64, k: u64, v: u64, g: u64, b: u64, rows: usize, y: u64, state_out: u64) !void { + var a: cuda.Args = .{}; + a.add(@as(c_int, @intCast(heads))); + for ([_]u64{ state_in, q, k, v, g, b }) |x| a.add(x); + a.add(@as(c_int, @intCast(rows))); + a.add(y); + a.add(state_out); + try cuda.launch.launch(ck.f, .{ .grid = .{ .x = @intCast(heads) }, .block = .{ .x = 2 * dk } }, stream, &a); +} + +/// Opt-in: the wide chain with 16-row WY chunks; float64's bound, not the row path's bits; cuts at 16 k == uncut. +pub fn chunked(ks: *const Kernels, ck: Chunked, stream: cuda.Stream, c: Chain, rows: usize, s: Saves, w: Wide) !void { + try prep(ks, stream, c, rows, s, w); + try chunkLaunch(ck, stream, c.heads, c.state_in, w.q, s.k, s.v, s.g, s.b, rows, w.y, s.state_out); + try out(ks, stream, c, rows, s, w); +} + +/// The chunked recurrence over `rows` saved rows, state only: `chunked`'s state bits for a call of exactly those rows. +pub fn chunkReplay(ck: Chunked, stream: cuda.Stream, state_in: u64, k: u64, v: u64, g: u64, b: u64, rows: usize, heads: usize, state_out: u64) !void { + try chunkLaunch(ck, stream, heads, state_in, 0, k, v, g, b, rows, 0, state_out); +} + /// kda.replay: the state after the first `rows` saved rows, from state_in. pub fn replay(ks: *const Kernels, stream: cuda.Stream, state_in: u64, k: u64, v: u64, g: u64, b: u64, rows: usize, heads: usize, state_out: u64) !void { var a: cuda.Args = .{}; diff --git a/zig/tests/cuda/glm_kda_chunk_tests.zig b/zig/tests/cuda/glm_kda_chunk_tests.zig new file mode 100644 index 0000000000..8cfec7d58c --- /dev/null +++ b/zig/tests/cuda/glm_kda_chunk_tests.zig @@ -0,0 +1,124 @@ +//! glm_kda_chunk.cu on synthetic GLM rows: float64 bounds, the wide chain's saves, its replay, resumed vs fresh. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const f64r = @import("glm_kda_f64.zig"); +const rig = @import("glm_kda_rig.zig"); +const ref = @import("glm_ref_tests.zig"); +const Rig = rig.Rig; +const kda = rig.kda; +const dk = rig.dk; + +const heads = 8; +const n = heads * dk; +const p_stride = 3 * n + 64; +const layout: f64r.Layout = .{ .heads = heads, .p_stride = p_stride, .b_off = 3 * n, .lower = -5, .eps = 1e-6 }; + +fn uniform(r: *Rig, count: usize, lo: f32, hi: f32) ![]f32 { + const out = try r.a.alloc(f32, count); + for (out) |*v| v.* = lo + r.rng.float(f32) * (hi - lo); + return out; +} + +/// glm_ref_kda's synthetic inputs (its scales) at `heads` heads, and the chain over them on the device. +const Case = struct { + in: f64r.Inputs, + c: kda.Chain, + + fn make(r: *Rig, rows: usize) !Case { + const in: f64r.Inputs = .{ .p = try r.bfs(rows * p_stride, 1.0), .a = try r.bfs(rows * n, 2.0), .g = try r.bfs(rows * n, 2.0), .cs = try r.bfs(3 * 3 * n, 1.0), .cw = try uniform(r, 3 * n * 4, -0.5, 0.5), .state = try uniform(r, n * dk, -0.05, 0.05), .a_log = try uniform(r, heads, -1, 1), .dt_bias = try uniform(r, n, -0.5, 0.5), .norm = try r.bfs(dk, 1.0) }; + const B = std.mem.sliceAsBytes; + const c: kda.Chain = .{ .p = try r.dev(B(in.p)), .p_stride = p_stride, .b_off = 3 * n, .a = try r.dev(B(in.a)), .a_stride = n, .g = try r.dev(B(in.g)), .g_stride = n, .conv_state = try r.dev(B(in.cs)), .conv_w = try r.dev(B(in.cw)), .state_in = try r.dev(B(in.state)), .a_log = try r.dev(B(in.a_log)), .dt_bias = try r.dev(B(in.dt_bias)), .norm_w = try r.dev(B(in.norm)), .eps = 1e-6, .lower = -5, .heads = heads }; + return .{ .in = in, .c = c }; + } + + /// The chain from row `lo` on: its rows, and the conv window of rows lo - 3 .. lo - 1 (lo >= 3) on the device. + fn from(k: Case, r: *Rig, lo: usize, state_in: u64) !kda.Chain { + const cs = try r.a.alloc(u16, 3 * 3 * n); + for (0..3) |j| @memcpy(cs[j * 3 * n ..][0 .. 3 * n], k.in.p[(lo - 3 + j) * p_stride ..][0 .. 3 * n]); + var c = k.c; + c.p += lo * p_stride * 2; + c.a += lo * n * 2; + c.g += lo * n * 2; + c.conv_state = try r.dev(std.mem.sliceAsBytes(cs)); + c.state_in = state_in; + return c; + } +}; + +/// Worst fractions of the end-to-end bound (out, state) and of the recurrence's from the path's saves (y, state). +fn bounds(r: *Rig, k: Case, o: rig.Out, rows: usize, e2e: f64r.Want) ![4]f64 { + const rec = try f64r.recurrence(r.a, heads, k.in.state, try rig.saved(r, o, heads), rows); + const out = try rig.wide64(r, o.out, rows * n, true); + const st = try rig.wide64(r, o.state, n * dk, false); + const y = try rig.wide64(r, o.y, rows * n, true); + return .{ f64r.fraction(out, e2e.out, e2e.out_tol), f64r.fraction(st, e2e.state, e2e.state_tol), f64r.fraction(y, rec.out, rec.out_tol), f64r.fraction(st, rec.state, rec.state_tol) }; +} + +fn rowsCheck(r: *Rig, m: *const rig.Mods, rows: usize) !f64 { + const k = try Case.make(r, rows); + const wide = try rig.Out.alloc(r, heads, rows); + const chunk = try rig.Out.alloc(r, heads, rows); + try rig.run(m, r.s, k.c, rows, wide, .wide); + try rig.run(m, r.s, k.c, rows, chunk, .chunked); + var name: [128]u8 = undefined; + const P = std.fmt.bufPrint; + for ([_]u64{ wide.q, wide.k, wide.v, wide.g, wide.b }, [_]u64{ chunk.q, chunk.k, chunk.v, chunk.g, chunk.b }, [_]usize{ 4 * n, 4 * n, 2 * n, 4 * n, 4 * heads }, [_][]const u8{ "q", "k", "v", "decay", "beta" }) |a, b, size, what| + try ref.sameRows(r, try P(&name, "KDA chunked {d} rows: saved {s} equal the wide chain's", .{ rows, what }), a, b, size * rows); + const replayed = try r.zeros(n * dk * 4); + try kda.chunkReplay(m.ck, r.s, k.c.state_in, chunk.k, chunk.v, chunk.g, chunk.b, rows, heads, replayed); + try ref.sameRows(r, try P(&name, "KDA chunked {d} rows: the chunked replay rebuilds its state", .{rows}), chunk.state, replayed, n * dk * 4); + const e2e = try f64r.chain(r.a, layout, k.in, rows); + const bc = try bounds(r, k, chunk, rows, e2e); + const bw = try bounds(r, k, wide, rows, e2e); + const dout = rig.diff(try rig.wide64(r, chunk.out, rows * n, true), try rig.wide64(r, wide.out, rows * n, true)); + const dst = rig.diff(try rig.wide64(r, chunk.state, n * dk, false), try rig.wide64(r, wide.state, n * dk, false)); + std.debug.print("RESULT KDA {d} rows x {d} heads, worst fraction of float64's bound, chunked (wide): end to end out {d:.3} ({d:.3}), state {d:.4} ({d:.4}); " ++ + "recurrence from the saves y {d:.3} ({d:.3}), state {d:.3} ({d:.3})\n", .{ rows, heads, bc[0], bw[0], bc[1], bw[1], bc[2], bw[2], bc[3], bw[3] }); + std.debug.print("RESULT KDA {d} rows: chunked vs wide, out {d} of {d} bf16 differ (max {e:.2} of max |out|), state {d} of {d} fp32 differ (max {e:.2} of max |S|)\n", .{ rows, dout.count, rows * n, dout.rel, dst.count, n * dk, dst.rel }); + var worst: f64 = 0; + for (bc ++ bw) |x| worst = @max(worst, x); + try check.expect(worst <= 1, "KDA {d} rows: chunked and wide within float64's bounds (worst {d:.3})", .{ rows, worst }); + return worst; +} + +/// 300 rows resumed at `at` (second call from the first's state and conv window) against one call: equal bytes? +fn cut(r: *Rig, m: *const rig.Mods, at: usize) !bool { + const rows = 300; + const k = try Case.make(r, rows); + const whole = try rig.Out.alloc(r, heads, rows); + try rig.run(m, r.s, k.c, rows, whole, .chunked); + const parts = try rig.Out.alloc(r, heads, rows); + try rig.run(m, r.s, k.c, at, parts, .chunked); + var second = parts.from(heads, at); + second.state = try r.zeros(n * dk * 4); + try rig.run(m, r.s, try k.from(r, at, parts.state), rows - at, second, .chunked); + const out_same = std.mem.eql(u8, try r.back(whole.out, rows * n * 2), try r.back(parts.out, rows * n * 2)); + const st_same = std.mem.eql(u8, try r.back(whole.state, n * dk * 4), try r.back(second.state, n * dk * 4)); + const d = rig.diff(try rig.wide64(r, second.state, n * dk, false), try rig.wide64(r, whole.state, n * dk, false)); + const dy = rig.diff(try rig.wide64(r, parts.out, rows * n, true), try rig.wide64(r, whole.out, rows * n, true)); + std.debug.print("RESULT KDA chunked 300 rows resumed at {d} vs fresh: outputs equal {} ({d} bf16 differ, max {e:.2} of max |out|), state equal {} ({d} fp32 differ, max {e:.2} of max |S|)\n", .{ at, out_same, dy.count, dy.rel, st_same, d.count, d.rel }); + return out_same and st_same; +} + +/// `glm-kda-chunk`. +pub fn run(gpu: check.Gpu) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(804); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + } + var m = try rig.Mods.load(gpu.d); + defer m.unload(); + var worst: f64 = 0; + for ([_]usize{ 1, 15, 16, 17, 70, 300 }) |rows| worst = @max(worst, try rowsCheck(&r, &m, rows)); + for ([_]usize{ 16, 48, 160 }) |at| try check.expect(try cut(&r, &m, at), "KDA chunked: a cut at {d} rows keeps one call's bits", .{at}); + var off = false; + for ([_]usize{ 37, 100, 255 }) |at| off = off or try cut(&r, &m, at); + check.pass("KDA chunked {d} heads: within float64's bounds (worst {d:.3}), saves equal the wide chain's, chunked replay rebuilds the state, cuts at 16/48/160 == uncut (resumed at 37/100/255: {s})", .{ heads, worst, if (off) "equal somewhere" else "differ, as expected" }); +} diff --git a/zig/tests/cuda/glm_kda_f64.zig b/zig/tests/cuda/glm_kda_f64.zig new file mode 100644 index 0000000000..b0db69748a --- /dev/null +++ b/zig/tests/cuda/glm_kda_f64.zig @@ -0,0 +1,201 @@ +//! glm_kda.cu's chain in float64 at any head count, end to end and the recurrence alone (glm-kda-chunk, glm-kda-ling). +const std = @import("std"); + +pub const dk = 128; + +/// x rounded to the nearest bf16, ties to even (finite normal values). +pub fn bf16(x: f64) f64 { + const u: u64 = @bitCast(x); + const lsb = (u >> 45) & 1; + return @bitCast((u + ((@as(u64, 1) << 44) - 1) + lsb) & ~((@as(u64, 1) << 45) - 1)); +} + +pub fn bf(u: u16) f64 { + return @as(f32, @bitCast(@as(u32, u) << 16)); +} + +fn sigmoid(x: f64) f64 { + return 1 / (1 + @exp(-x)); +} + +const p7 = std.math.pow(f64, 2, -7); +const p8 = std.math.pow(f64, 2, -8); +/// The recurrence's bound from the saves: fp32 sums of up to 144 terms a row, contracting rows, 2^4 to spare. +const p12 = std.math.pow(f64, 2, -12); + +/// The call's shape: p rows hold q | k | v (3 H 128 channels) and beta's logits at b_off; a and g rows are H 128 wide. +pub const Layout = struct { heads: usize, p_stride: usize, b_off: usize, lower: f64, eps: f64 }; + +/// Raw rows and weights as the chain takes them: cs [3][3 H 128] oldest first, taps [3 H 128][4], state [H][v][k]. +pub const Inputs = struct { p: []const u16, a: []const u16, g: []const u16, cs: []const u16, cw: []const f32, state: []const f32, a_log: []const f32, dt_bias: []const f32, norm: []const u16 }; + +/// prep's rows [rows][H][128]: q (scaled), k, v bf16, decay factors, and beta [rows][H]. +pub const Saved = struct { q: []const f32, k: []const f32, v: []const u16, g: []const f32, b: []const f32 }; + +/// Outputs [rows][H 128] and the final state [H][128][128], each with a per-element bound. +pub const Want = struct { out: []f64, out_tol: []f64, state: []f64, state_tol: []f64 }; + +fn want(a: std.mem.Allocator, heads: usize, rows: usize) !Want { + const n = heads * dk; + return .{ .out = try a.alloc(f64, rows * n), .out_tol = try a.alloc(f64, rows * n), .state = try a.alloc(f64, n * dk), .state_tol = try a.alloc(f64, n * dk) }; +} + +/// A delta step: decay, u = beta (v - S k), S += k u, o = S q; mag bounds |S|; o_mag kda_ref's or (fp32) sum mag |q|. +fn step(s: []f64, mag: []f64, q: []const f64, k: []const f64, v: []const f64, g: []const f64, beta: f64, o: []f64, o_mag: []f64, fp32: bool) void { + var kq: f64 = 0; + for (q, k) |qq, kk| kq += @abs(qq * kk); + for (0..dk) |j| { + const row = s[j * dk ..][0..dk]; + const mrow = mag[j * dk ..][0..dk]; + var kv: f64 = 0; + var km: f64 = 0; + for (row, mrow, k, g) |*x, *m, kk, gg| { + x.* *= gg; + m.* *= gg; + kv += x.* * kk; + km += m.* * @abs(kk); + } + const u = beta * (v[j] - kv); + const um = beta * (@abs(v[j]) + km); + var acc: f64 = 0; + var am: f64 = 0; + var aq: f64 = 0; + for (row, mrow, k, q) |*x, *m, kk, qq| { + x.* += kk * u; + m.* += @abs(kk) * um; + acc += x.* * qq; + am += m.* * @abs(qq); + aq += @abs(x.* * qq); + } + o[j] = acc; + o_mag[j] = if (fp32) am else aq + beta * @abs(v[j]) * kq; + } +} + +/// kda_ref.gatedNorm: o / rms * w * sigmoid(gate); tol: o's bound in, y's bound out (first order) + y's bf16 rounding. +fn gatedNorm(o: []const f64, gate: []const f64, w: []const f64, eps: f64, y: []f64, tol: []f64) void { + var ss: f64 = 0; + var cross: f64 = 0; + for (o, tol) |x, t| { + ss += x * x; + cross += @abs(x) * t; + } + const nn: f64 = @floatFromInt(o.len); + const rstd = 1 / @sqrt(ss / nn + eps); + const moved = cross / nn / (ss / nn + eps); + for (y, tol, o, gate, w) |*out, *t, x, z, ww| { + out.* = x * rstd * ww * sigmoid(z); + t.* = @abs(ww * sigmoid(z)) * rstd * (t.* + @abs(x) * moved) + p8 * @abs(out.*); + } +} + +fn l2norm(x: []f64, scale: f64) void { + var ss: f64 = 0; + for (x) |e| ss += e * e; + const f = scale / @sqrt(ss + 1e-6); + for (x) |*e| e.* *= f; +} + +/// End to end: conv+SiLU, q/k L2 norms, safe gate, bf16 beta, delta rule, bf16(S q), gated norm; kda_ref's 2^-7 bounds. +pub fn chain(a: std.mem.Allocator, l: Layout, in: Inputs, rows: usize) !Want { + const n = l.heads * dk; + const cc = 3 * n; + var w = try want(a, l.heads, rows); + const mag = try a.alloc(f64, dk * dk); + var qkv: [3][dk]f64 = undefined; + var g: [dk]f64 = undefined; + var o: [dk]f64 = undefined; + var om: [dk]f64 = undefined; + var gate: [dk]f64 = undefined; + var nw: [dk]f64 = undefined; + for (&nw, in.norm[0..dk]) |*d, x| d.* = bf(x); + for (0..l.heads) |h| { + const s = w.state[h * dk * dk ..][0 .. dk * dk]; + for (s, in.state[h * dk * dk ..][0 .. dk * dk]) |*d, x| d.* = x; + for (mag, s) |*m, x| m.* = @abs(x); + const rate = @exp(@as(f64, in.a_log[h])); + for (0..rows) |row| { + for (0..3) |part| for (0..dk) |i| { + const c = part * n + h * dk + i; + var acc: f64 = 0; + for (0..4) |tap| { + const at = row + tap; + const x = if (at < 3) bf(in.cs[at * cc + c]) else bf(in.p[(at - 3) * l.p_stride + c]); + acc += @as(f64, in.cw[c * 4 + tap]) * x; + } + qkv[part][i] = bf16(acc * sigmoid(acc)); + }; + l2norm(&qkv[0], 1 / @sqrt(@as(f64, dk))); + l2norm(&qkv[1], 1); + for (&g, 0..) |*e, i| e.* = @exp(l.lower * sigmoid(rate * (bf(in.a[row * n + h * dk + i]) + in.dt_bias[h * dk + i]))); + const beta = bf16(sigmoid(bf(in.p[row * l.p_stride + l.b_off + h]))); + step(s, mag, &qkv[0], &qkv[1], &qkv[2], &g, beta, &o, &om, false); + const tol = w.out_tol[row * n + h * dk ..][0..dk]; + for (&o, tol, om) |*e, *t, m| { + t.* = p7 * m; + e.* = bf16(e.*); + } + for (&gate, 0..) |*z, t| z.* = bf(in.g[row * n + h * dk + t]); + gatedNorm(&o, &gate, &nw, l.eps, w.out[row * n + h * dk ..][0..dk], tol); + } + for (w.state_tol[h * dk * dk ..][0 .. dk * dk], mag) |*t, m| t.* = p7 * m + 1e-30; + } + return w; +} + +/// The recurrence alone from prep's saves: S q (tol 2^-8 |y| + 2^-12 sum |mag q|), state (tol 2^-12 of magnitudes). +pub fn recurrence(a: std.mem.Allocator, heads: usize, state0: []const f32, sv: Saved, rows: usize) !Want { + const n = heads * dk; + var w = try want(a, heads, rows); + const mag = try a.alloc(f64, dk * dk); + var q: [dk]f64 = undefined; + var k: [dk]f64 = undefined; + var v: [dk]f64 = undefined; + var g: [dk]f64 = undefined; + var om: [dk]f64 = undefined; + for (0..heads) |h| { + const s = w.state[h * dk * dk ..][0 .. dk * dk]; + for (s, state0[h * dk * dk ..][0 .. dk * dk]) |*d, x| d.* = x; + for (mag, s) |*m, x| m.* = @abs(x); + for (0..rows) |row| { + const at = row * n + h * dk; + for (&q, &k, &v, &g, sv.q[at..][0..dk], sv.k[at..][0..dk], sv.v[at..][0..dk], sv.g[at..][0..dk]) |*qq, *kk, *vv, *gg, a0, a1, a2, a3| { + qq.* = a0; + kk.* = a1; + vv.* = bf(a2); + gg.* = a3; + } + const o = w.out[at..][0..dk]; + step(s, mag, &q, &k, &v, &g, sv.b[row * heads + h], o, &om, true); + for (w.out_tol[at..][0..dk], o, om) |*t, e, m| t.* = p8 * @abs(e) + p12 * m + 1e-30; + } + for (w.state_tol[h * dk * dk ..][0 .. dk * dk], mag) |*t, m| t.* = p12 * m + 1e-30; + } + return w; +} + +/// The worst |got - want| / tol. +pub fn fraction(got: []const f64, w: []const f64, tol: []const f64) f64 { + var most: f64 = 0; + for (got, w, tol) |x, e, t| { + const q = @abs(x - e) / t; + if (!(q <= most)) most = if (std.math.isNan(q)) std.math.inf(f64) else q; + } + return most; +} + +/// Relative L2 distance of got from want. +pub fn relL2(got: []const f64, w: []const f64) f64 { + var dd: f64 = 0; + var ww: f64 = 0; + for (got, w) |x, e| { + dd += (x - e) * (x - e); + ww += e * e; + } + return @sqrt(dd / ww); +} + +test "bf16 rounding ties to even" { + try std.testing.expectEqual(@as(f64, 1.0), bf16(1.0 + 1.0 / 256.0)); + try std.testing.expectEqual(@as(f64, 1.0 + 1.0 / 64.0), bf16(1.0 + 3.0 / 256.0)); +} diff --git a/zig/tests/cuda/glm_kda_ling.zig b/zig/tests/cuda/glm_kda_ling.zig new file mode 100644 index 0000000000..d357a1f586 --- /dev/null +++ b/zig/tests/cuda/glm_kda_ling.zig @@ -0,0 +1,216 @@ +//! glm_kda's paths on Ling-3.1 layer 0 (kda-real's fla fixture, q/k/v convs as one q|k|v): vs fla and float64. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const f64r = @import("glm_kda_f64.zig"); +const rig = @import("glm_kda_rig.zig"); +const ref = @import("glm_ref_tests.zig"); +const Rig = rig.Rig; +const kda = rig.kda; +const dk = rig.dk; + +fn take(comptime T: type, a: std.mem.Allocator, bytes: []const u8, at: *usize, count: usize) ![]T { + const out = try a.alloc(T, count); + @memcpy(std.mem.sliceAsBytes(out), bytes[at.*..][0 .. count * @sizeOf(T)]); + at.* += count * @sizeOf(T); + return out; +} + +/// The layer's rows as GLM's chain takes them, fla's outputs and states, and per-row conv windows on the device. +const Layer = struct { + P: usize, + T: usize, + H: usize, + l: f64r.Layout, + in: f64r.Inputs, + c: kda.Chain, + cs_all: u64, // [T][3][3 H 128]: row t's window, rows t - 3 .. t - 1 (zeros before the first) + fla_y: []f64, + fla_sp: []f64, + fla_st: []f64, + + /// The chain over rows lo .. from state_in with row lo's conv window. + fn at(l: Layer, lo: usize, state_in: u64) kda.Chain { + const n = l.H * dk; + var c = l.c; + c.p += lo * l.l.p_stride * 2; + c.a += lo * n * 2; + c.g += lo * n * 2; + c.conv_state = l.cs_all + lo * 9 * n * 2; + c.state_in = state_in; + return c; + } +}; + +fn load(r: *Rig, path: []const u8) !?Layer { + const a = r.a; + const bytes = std.Io.Dir.cwd().readFileAlloc(r.gpu.io, path, a, .limited(1 << 30)) catch |e| switch (e) { + error.FileNotFound => return null, + else => return e, + }; + if (bytes.len < 32 or !std.mem.eql(u8, bytes[0..8], "TFKDA001")) return error.BadFixture; + const P: usize = std.mem.readInt(u32, bytes[8..12], .little); + const N: usize = std.mem.readInt(u32, bytes[12..16], .little); + const H: usize = std.mem.readInt(u32, bytes[16..20], .little); + if (std.mem.readInt(u32, bytes[20..24], .little) != dk) return error.BadFixture; + const lower: f32 = @bitCast(std.mem.readInt(u32, bytes[24..28], .little)); + const eps: f32 = @bitCast(std.mem.readInt(u32, bytes[28..32], .little)); + const T = P + N; + const n = H * dk; + var o: usize = 32; + const qkv = try take(u16, a, bytes, &o, T * 3 * n); + const f = try take(u16, a, bytes, &o, T * n); + const b = try take(u16, a, bytes, &o, T * H); + const gate = try take(u16, a, bytes, &o, T * n); + const conv = try take(f32, a, bytes, &o, 3 * n * 4); // Ling's q_conv1d | k_conv1d | v_conv1d: GLM's q|k|v channels + const a_log = try take(f32, a, bytes, &o, H); + const dt_bias = try take(f32, a, bytes, &o, n); + const norm32 = try take(f32, a, bytes, &o, dk); + const fy = try take(u16, a, bytes, &o, T * n); + const fsp = try take(f32, a, bytes, &o, n * dk); + const fst = try take(f32, a, bytes, &o, n * dk); + if (o != bytes.len) return error.BadFixture; + const p_stride = 3 * n + H; // [q | k | v | beta logits] + const p = try a.alloc(u16, T * p_stride); + for (0..T) |t| { + @memcpy(p[t * p_stride ..][0 .. 3 * n], qkv[t * 3 * n ..][0 .. 3 * n]); + @memcpy(p[t * p_stride + 3 * n ..][0..H], b[t * H ..][0..H]); + } + const norm = try a.alloc(u16, dk); // the o_norm weight is bf16 in the checkpoint: exact + for (norm, norm32) |*d, x| { + d.* = @truncate(@as(u32, @bitCast(x)) >> 16); + if (f64r.bf(d.*) != x) return error.NormNotBf16; + } + const zero = try a.alloc(f32, n * dk); + @memset(zero, 0); + const cs = try a.alloc(u16, T * 9 * n); + @memset(cs, 0); + for (0..T) |t| for (0..3) |j| if (t + j >= 3) @memcpy(cs[(t * 3 + j) * 3 * n ..][0 .. 3 * n], qkv[(t + j - 3) * 3 * n ..][0 .. 3 * n]); + const in: f64r.Inputs = .{ .p = p, .a = f, .g = gate, .cs = cs[0 .. 9 * n], .cw = conv, .state = zero, .a_log = a_log, .dt_bias = dt_bias, .norm = norm }; + const B = std.mem.sliceAsBytes; + const c: kda.Chain = .{ .p = try r.dev(B(p)), .p_stride = p_stride, .b_off = 3 * n, .a = try r.dev(B(f)), .a_stride = n, .g = try r.dev(B(gate)), .g_stride = n, .conv_state = 0, .conv_w = try r.dev(B(conv)), .state_in = try r.dev(B(zero)), .a_log = try r.dev(B(a_log)), .dt_bias = try r.dev(B(dt_bias)), .norm_w = try r.dev(B(norm)), .eps = eps, .lower = lower, .heads = H }; + const w64 = struct { + fn of(al: std.mem.Allocator, x: anytype) ![]f64 { + const out = try al.alloc(f64, x.len); + for (out, x) |*d, v| d.* = if (@TypeOf(v) == u16) f64r.bf(v) else v; + return out; + } + }; + return .{ .P = P, .T = T, .H = H, .l = .{ .heads = H, .p_stride = p_stride, .b_off = 3 * n, .lower = lower, .eps = eps }, .in = in, .c = c, .cs_all = try r.dev(B(cs)), .fla_y = try w64.of(a, fy), .fla_sp = try w64.of(a, fsp), .fla_st = try w64.of(a, fst) }; +} + +/// One run's outputs [T][H 128] and its states after the prompt (P rows) and at the end. +const Result = struct { y: []f64, sp: u64, st: u64 }; + +/// Rows lo .. T as decode rows (fused chain, one row a call) from `state`; each row's state into states[t] when given. +fn decode(r: *Rig, m: *const rig.Mods, l: Layer, lo: usize, state: u64, y: []f64, states: ?[]u64) !u64 { + const n = l.H * dk; + var s = state; + for (lo..l.T) |t| { + const o = try rig.Out.alloc(r, l.H, 1); + try rig.run(m, r.s, l.at(t, s), 1, o, .fused); + @memcpy(y[t * n ..][0..n], try rig.wide64(r, o.out, n, true)); + s = o.state; + if (states) |ss| ss[t + 1] = s; + } + return s; +} + +/// A prompt of P rows through `path`, then decode rows; `saves` returns the prompt's buffers. +fn promptThenDecode(r: *Rig, m: *const rig.Mods, l: Layer, path: rig.Path, saves: *rig.Out) !Result { + const n = l.H * dk; + const y = try r.a.alloc(f64, l.T * n); + saves.* = try rig.Out.alloc(r, l.H, l.P); + try rig.run(m, r.s, l.at(0, l.c.state_in), l.P, saves.*, path); + @memcpy(y[0 .. l.P * n], try rig.wide64(r, saves.out, l.P * n, true)); + return .{ .y = y, .sp = saves.state, .st = try decode(r, m, l, l.P, saves.state, y, null) }; +} + +fn fracs(r: *Rig, l: Layer, res: Result, e2e_p: f64r.Want, e2e: f64r.Want) ![4]f64 { + const n = l.H * dk; + const sp = try rig.wide64(r, res.sp, n * dk, false); + const st = try rig.wide64(r, res.st, n * dk, false); + return .{ f64r.fraction(res.y, e2e.out, e2e.out_tol), f64r.fraction(sp, e2e_p.state, e2e_p.state_tol), f64r.fraction(st, e2e.state, e2e.state_tol), @max(f64r.relL2(sp, l.fla_sp), f64r.relL2(st, l.fla_st)) }; +} + +/// `glm-kda-ling FIXTURE`. +pub fn run(gpu: check.Gpu, path: []const u8) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(1); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + } + const l = (try load(&r, path)) orelse return std.debug.print("SKIP glm-kda-ling: no fixture at {s}\n", .{path}); + var m = try rig.Mods.load(gpu.d); + defer m.unload(); + const n = l.H * dk; + const T = l.T; + var name: [128]u8 = undefined; + const P = std.fmt.bufPrint; + // every row a decode row, each row's state kept + const states = try r.a.alloc(u64, T + 1); + states[0] = l.c.state_in; + const yd = try r.a.alloc(f64, T * n); + _ = try decode(&r, &m, l, 0, l.c.state_in, yd, states); + const dec: Result = .{ .y = yd, .sp = states[l.P], .st = states[T] }; + // 16-row verify windows; each replayed whole (== its state) and over its first 5 rows (== the decode rows' state) + const yw = try r.a.alloc(f64, T * n); + var ws = l.c.state_in; + var w0: usize = 0; + while (w0 < T) : (w0 += 16) { + const rows = @min(16, T - w0); + const o = try rig.Out.alloc(&r, l.H, rows); + try rig.run(&m, r.s, l.at(w0, ws), rows, o, .fused); + @memcpy(yw[w0 * n ..][0 .. rows * n], try rig.wide64(&r, o.out, rows * n, true)); + const full = try r.zeros(n * dk * 4); + const part = try r.zeros(n * dk * 4); + try kda.replay(&m.ks, r.s, ws, o.k, o.v, o.g, o.b, rows, l.H, full); + try kda.replay(&m.ks, r.s, ws, o.k, o.v, o.g, o.b, @min(5, rows), l.H, part); + try check.sameBytes(try P(&name, "Ling window at {d}: replay == the window's state", .{w0}), try r.back(full, n * dk * 4), try r.back(o.state, n * dk * 4)); + try check.sameBytes(try P(&name, "Ling window at {d}: replay of 5 rows == 5 decode rows' state", .{w0}), try r.back(part, n * dk * 4), try r.back(states[w0 + @min(5, rows)], n * dk * 4)); + ws = o.state; + } + try check.sameBytes("Ling: 16-row windows' state == decode rows' state", try r.back(ws, n * dk * 4), try r.back(states[T], n * dk * 4)); + for (yw, yd) |a, b| if (a != b) return check.expect(false, "Ling: 16-row windows' outputs == decode rows' outputs", .{}); + const win: Result = .{ .y = yw, .sp = states[l.P], .st = ws }; + var wsv: rig.Out = undefined; + var csv: rig.Out = undefined; + const wide = try promptThenDecode(&r, &m, l, .wide, &wsv); + const chunk = try promptThenDecode(&r, &m, l, .chunked, &csv); + for (wide.y, yd) |a, b| if (a != b) return check.expect(false, "Ling: the wide prompt's and decode outputs == every row a decode row's", .{}); + try ref.sameRows(&r, "Ling: the wide prompt's state == P decode rows' state", wide.sp, states[l.P], n * dk * 4); + const crep = try r.zeros(n * dk * 4); + try kda.chunkReplay(m.ck, r.s, l.c.state_in, csv.k, csv.v, csv.g, csv.b, l.P, l.H, crep); + try ref.sameRows(&r, "Ling: the chunked replay == the chunked prompt's state", csv.state, crep, n * dk * 4); + const e2e_p = try f64r.chain(r.a, l.l, l.in, l.P); + const e2e = try f64r.chain(r.a, l.l, l.in, T); + const rec = try f64r.recurrence(r.a, l.H, l.in.state, try rig.saved(&r, csv, l.H), l.P); + const rc = [2]f64{ f64r.fraction(try rig.wide64(&r, csv.y, l.P * n, true), rec.out, rec.out_tol), f64r.fraction(try rig.wide64(&r, csv.state, n * dk, false), rec.state, rec.state_tol) }; + const recw = [2]f64{ f64r.fraction(try rig.wide64(&r, wsv.y, l.P * n, true), rec.out, rec.out_tol), f64r.fraction(try rig.wide64(&r, wsv.state, n * dk, false), rec.state, rec.state_tol) }; + var worst: f64 = @max(@max(rc[0], rc[1]), @max(recw[0], recw[1])); + var most_l2: f64 = 0; + std.debug.print("RESULT glm-kda-ling: Ling-3.1 layer 0, {d} heads, {d} prompt + {d} decode rows; fla rel L2 of outputs (states after the prompt / end), worst fraction of float64's bound (out, state after prompt, at end)\n", .{ l.H, l.P, T - l.P }); + for ([_]Result{ dec, win, wide, chunk }, [_][]const u8{ "chain, every row a decode row", "chain, 16-row windows", "wide prompt + decode", "chunked prompt + decode" }) |res, label| { + const f = try fracs(&r, l, res, e2e_p, e2e); + const sp = try rig.wide64(&r, res.sp, n * dk, false); + const st = try rig.wide64(&r, res.st, n * dk, false); + const yl2 = f64r.relL2(res.y, l.fla_y); + std.debug.print("RESULT {s}: rel L2 {e:.3} (states {e:.3} / {e:.3}); float64 {d:.3}, {d:.4}, {d:.4}\n", .{ label, yl2, f64r.relL2(sp, l.fla_sp), f64r.relL2(st, l.fla_st), f[0], f[1], f[2] }); + worst = @max(worst, @max(f[0], @max(f[1], f[2]))); + most_l2 = @max(most_l2, @max(yl2, f[3])); + } + const fla_f = f64r.fraction(l.fla_y, e2e.out, e2e.out_tol); + const dy = rig.diff(chunk.y[0 .. l.P * n], wide.y[0 .. l.P * n]); + const dsp = rig.diff(try rig.wide64(&r, chunk.sp, n * dk, false), try rig.wide64(&r, wide.sp, n * dk, false)); + const dall = rig.diff(chunk.y, wide.y); + std.debug.print("RESULT fla's own outputs: float64 {d:.3}; recurrence from the chunked saves, chunked y {d:.3} state {d:.3} (wide {d:.3}, {d:.3})\n", .{ fla_f, rc[0], rc[1], recw[0], recw[1] }); + std.debug.print("RESULT chunked vs wide prompt: out {d} of {d} bf16 differ (max {e:.2} of max |out|), state {d} of {d} fp32 differ (max {e:.2} of max |S|); all {d} rows incl. decode: {d} differ\n", .{ dy.count, l.P * n, dy.rel, dsp.count, n * dk, dsp.rel, T, dall.count }); + try check.expect(worst <= 1, "glm-kda-ling: every path within float64's bound (worst {d:.3})", .{worst}); + try check.expect(most_l2 < 1e-2, "glm-kda-ling: every path within 1 % of fla (worst rel L2 {e:.3})", .{most_l2}); + check.pass("glm-kda-ling: glm_kda serves Ling layer 0 as is; chain, windows, wide and chunked within float64's bound (worst {d:.3}) and fla (rel L2 <= {e:.3}); replays == chain bits", .{ worst, most_l2 }); +} diff --git a/zig/tests/cuda/glm_kda_rig.zig b/zig/tests/cuda/glm_kda_rig.zig new file mode 100644 index 0000000000..e7d22615d9 --- /dev/null +++ b/zig/tests/cuda/glm_kda_rig.zig @@ -0,0 +1,117 @@ +//! glm-kda-chunk's and glm-kda-ling's shared rig: a call's inputs on the device, one path's outputs, downloads to f64. +const std = @import("std"); +const cuda = @import("cuda"); +const qr = @import("quant_ref_tests.zig"); +const glm = @import("glm_cuda"); +const f64r = @import("glm_kda_f64.zig"); +pub const Rig = qr.RigT; + +pub const kda = glm.kda; +pub const dk = kda.dk; + +/// The two modules: glm_kda.cu (chain, replay, wide) and glm_kda_chunk.cu. +pub const Mods = struct { + kda_mod: cuda.Module, + chunk_mod: cuda.Module, + ks: kda.Kernels, + ck: kda.Chunked, + + pub fn load(d: *const cuda.Driver) !Mods { + var m: Mods = undefined; + m.kda_mod = try cuda.Module.load(d, cuda.kernels.glm_kda); + m.chunk_mod = try cuda.Module.load(d, cuda.kernels.glm_kda_chunk); + m.ks = try kda.Kernels.load(m.kda_mod); + m.ck = try kda.Chunked.load(m.chunk_mod); + return m; + } + + pub fn unload(m: *Mods) void { + m.chunk_mod.unload(); + m.kda_mod.unload(); + } +}; + +/// How a call runs: the fused chain, the three-kernel (serial wide) chain, or the chunked prompt path. +pub const Path = enum { fused, wide, chunked }; + +/// One path's buffers for `rows` rows: out, state, saves (k, v, decay, beta), wide scratch (q, y). +pub const Out = struct { + out: u64, + state: u64, + k: u64, + v: u64, + g: u64, + b: u64, + q: u64, + y: u64, + rows: usize, + + pub fn alloc(r: *Rig, heads: usize, rows: usize) !Out { + const n = heads * dk; + return .{ .out = try r.zeros(rows * n * 2), .state = try r.zeros(n * dk * 4), .k = try r.zeros(rows * n * 4), .v = try r.zeros(rows * n * 2), .g = try r.zeros(rows * n * 4), .b = try r.zeros(rows * heads * 4), .q = try r.zeros(rows * n * 4), .y = try r.zeros(rows * n * 2), .rows = rows }; + } + + pub fn saves(o: Out) kda.Saves { + return .{ .out = o.out, .state_out = o.state, .k = o.k, .v = o.v, .g = o.g, .b = o.b }; + } + + /// The same buffers from row `lo` on (state unchanged). + pub fn from(o: Out, heads: usize, lo: usize) Out { + const n = heads * dk; + return .{ .out = o.out + lo * n * 2, .state = o.state, .k = o.k + lo * n * 4, .v = o.v + lo * n * 2, .g = o.g + lo * n * 4, .b = o.b + lo * heads * 4, .q = o.q + lo * n * 4, .y = o.y + lo * n * 2, .rows = o.rows - lo }; + } +}; + +/// `rows` rows of c through `path` into o. +pub fn run(m: *const Mods, s: cuda.Stream, c: kda.Chain, rows: usize, o: Out, path: Path) !void { + const w: kda.Wide = .{ .q = o.q, .y = o.y }; + switch (path) { + .fused => try kda.chain(&m.ks, s, c, rows, o.saves(), null), + .wide => try kda.chain(&m.ks, s, c, rows, o.saves(), w), + .chunked => try kda.chunked(&m.ks, m.ck, s, c, rows, o.saves(), w), + } +} + +pub fn floats(r: *Rig, ptr: u64, count: usize) ![]f32 { + const raw = try r.back(ptr, count * 4); + const out = try r.a.alloc(f32, count); + @memcpy(std.mem.sliceAsBytes(out), raw); + return out; +} + +pub fn halves(r: *Rig, ptr: u64, count: usize) ![]u16 { + const raw = try r.back(ptr, count * 2); + const out = try r.a.alloc(u16, count); + @memcpy(std.mem.sliceAsBytes(out), raw); + return out; +} + +/// fp32 or bf16 device values as f64. +pub fn wide64(r: *Rig, ptr: u64, count: usize, comptime bf16: bool) ![]f64 { + const out = try r.a.alloc(f64, count); + if (bf16) { + for (out, try halves(r, ptr, count)) |*d, x| d.* = f64r.bf(x); + } else for (out, try floats(r, ptr, count)) |*d, x| { + d.* = x; + } + return out; +} + +/// prep's saves of o, downloaded. +pub fn saved(r: *Rig, o: Out, heads: usize) !f64r.Saved { + const n = o.rows * heads * dk; + return .{ .q = try floats(r, o.q, n), .k = try floats(r, o.k, n), .v = try halves(r, o.v, n), .g = try floats(r, o.g, n), .b = try floats(r, o.b, o.rows * heads) }; +} + +/// Largest |a - b| over the largest |b|, and how many elements differ at all. +pub fn diff(a: []const f64, b: []const f64) struct { rel: f64, count: usize } { + var most: f64 = 0; + var big: f64 = 0; + var count: usize = 0; + for (a, b) |x, y| { + most = @max(most, @abs(x - y)); + big = @max(big, @abs(y)); + if (x != y) count += 1; + } + return .{ .rel = most / big, .count = count }; +} diff --git a/zig/tests/cuda/glm_kda_speed.zig b/zig/tests/cuda/glm_kda_speed.zig new file mode 100644 index 0000000000..2cfc97e885 --- /dev/null +++ b/zig/tests/cuda/glm_kda_speed.zig @@ -0,0 +1,76 @@ +//! glm-kda-speed: 64 heads x 128, a 2048-row prompt (row path vs chunked), a decode row, a 16-row window; median ms. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const rig = @import("glm_kda_rig.zig"); +const Rig = rig.Rig; +const kda = rig.kda; +const dk = rig.dk; + +const heads = 64; +const n = heads * dk; +const prompt = 2048; + +const What = enum { path, replay, prep_out }; + +fn time(r: *Rig, m: *const rig.Mods, c: kda.Chain, rows: usize, o: rig.Out, path: rig.Path, what: What) !f64 { + var start = try cuda.Event.init(r.gpu.d, true); + defer start.deinit(); + var end = try cuda.Event.init(r.gpu.d, true); + defer end.deinit(); + var ms: [9]f64 = undefined; + for (0..ms.len + 1) |i| { + try start.record(r.s); + switch (what) { + .path => try rig.run(m, r.s, c, rows, o, path), + .replay => if (path == .chunked) try kda.chunkReplay(m.ck, r.s, c.state_in, o.k, o.v, o.g, o.b, rows, heads, o.state) else try kda.replay(&m.ks, r.s, c.state_in, o.k, o.v, o.g, o.b, rows, heads, o.state), + .prep_out => { + try kda.prep(&m.ks, r.s, c, rows, o.saves(), .{ .q = o.q, .y = o.y }); + try kda.out(&m.ks, r.s, c, rows, o.saves(), .{ .q = o.q, .y = o.y }); + }, + } + try end.record(r.s); + try r.s.synchronize(); + if (i > 0) ms[i - 1] = try start.elapsedMs(end); + } + std.mem.sort(f64, &ms, {}, std.sort.asc(f64)); + return ms[ms.len / 2]; +} + +pub fn run(gpu: check.Gpu) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(64); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + } + var m = try rig.Mods.load(gpu.d); + defer m.unload(); + const p_stride = 3 * n + heads; + const state = try r.a.alloc(f32, n * dk); + for (state) |*v| v.* = (r.rng.float(f32) * 2 - 1) * 0.05; + const cw = try r.a.alloc(f32, 3 * n * 4); + for (cw) |*v| v.* = r.rng.float(f32) - 0.5; + const small = try r.a.alloc(f32, n); + for (small) |*v| v.* = r.rng.float(f32) - 0.5; + const B = std.mem.sliceAsBytes; + const c: kda.Chain = .{ .p = try r.dev(B(try r.bfs(prompt * p_stride, 1.0))), .p_stride = p_stride, .b_off = 3 * n, .a = try r.dev(B(try r.bfs(prompt * n, 2.0))), .a_stride = n, .g = try r.dev(B(try r.bfs(prompt * n, 2.0))), .g_stride = n, .conv_state = try r.dev(B(try r.bfs(9 * n, 1.0))), .conv_w = try r.dev(B(cw)), .state_in = try r.dev(B(state)), .a_log = try r.dev(B(small[0..heads])), .dt_bias = try r.dev(B(small)), .norm_w = try r.dev(B(try r.bfs(dk, 1.0))), .eps = 1e-6, .lower = -5, .heads = heads }; + const o = try rig.Out.alloc(&r, heads, prompt); + const wide = try time(&r, &m, c, prompt, o, .wide, .path); + const chunk = try time(&r, &m, c, prompt, o, .chunked, .path); + const rw = try time(&r, &m, c, prompt, o, .wide, .replay); + const rc = try time(&r, &m, c, prompt, o, .chunked, .replay); + const po = try time(&r, &m, c, prompt, o, .wide, .prep_out); + const row = try time(&r, &m, c, 1, o, .fused, .path); + const win = try time(&r, &m, c, 16, o, .fused, .path); + const win_wide = try time(&r, &m, c, 16, o, .wide, .path); + const win_chunk = try time(&r, &m, c, 16, o, .chunked, .path); + std.debug.print("RESULT glm-kda-speed {d} heads x 128, {d}-row prompt: wide (prep, serial step, out) {d:.3} ms, chunked (prep, WY chunks, out) {d:.3} ms ({d:.2}x); " ++ + "prep + out alone {d:.3} ms; state-only: serial replay {d:.3} ms, chunked replay {d:.3} ms\n", .{ heads, prompt, wide, chunk, wide / chunk, po, rw, rc }); + std.debug.print("RESULT glm-kda-speed decode row (fused chain) {d:.4} ms; 16-row window: fused chain {d:.4} ms, wide {d:.4} ms, chunked {d:.4} ms\n", .{ row, win, win_wide, win_chunk }); + check.pass("glm-kda-speed: timed", .{}); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index da8e589563..ff056558c8 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -42,6 +42,9 @@ const glm_ref_route = @import("glm_ref_route.zig"); const ling_synth = @import("ling_synth.zig"); const ling_weights = @import("ling_weights_tests.zig"); const ling_real = @import("ling_weights_real.zig"); +const glm_kda_chunk_tests = @import("glm_kda_chunk_tests.zig"); +const glm_kda_ling = @import("glm_kda_ling.zig"); +const glm_kda_speed = @import("glm_kda_speed.zig"); const usage = \\usage: tf-cuda-test @@ -98,6 +101,9 @@ const usage = \\ ling-mla Ling's MLA (128-wide absorb/expand, rope, gate, 576-wide attention) vs float64 (synthetic) \\ ling-mla-real Ling layer 7's MLA on real inputs vs HF and float64; decode, window and cut bytes \\ glm-mla-same GLM's MLA launches on this build vs a 6380c8da glm_glue fatbin, byte for byte + \\ glm-kda-chunk GLM KDA's chunked prompt path against float64, the wide chain's saves, cuts (synthetic) + \\ glm-kda-ling FIXTURE GLM KDA's paths on Ling-3.1 layer 0 (kda-real's fixture) against fla and float64 + \\ glm-kda-speed GLM KDA at 64 heads: a 2048-row prompt wide vs chunked, a decode row, a 16-row window \\ glm-lanes-cover copy drafts, cancels, discards, eos and max_new on one lanes.Engine across requests \\ glm-lanes-concurrent concurrent streams on one lanes.Engine, each against its solo run \\ glm-shared-check shared verify windows against lone twins, logits and states byte for byte @@ -172,6 +178,9 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "ling-mla")) return @import("ling_ref_mla.zig").run(gpu); if (std.mem.eql(u8, cmd, "ling-mla-real")) return @import("ling_mla_real.zig").run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "glm-mla-same")) return @import("glm_mla_same.zig").run(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "glm-kda-chunk")) return glm_kda_chunk_tests.run(gpu); + if (std.mem.eql(u8, cmd, "glm-kda-ling")) return glm_kda_ling.run(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "glm-kda-speed")) return glm_kda_speed.run(gpu); if (std.mem.eql(u8, cmd, "modelopt")) return modelopt_tests.run(gpu); if (std.mem.eql(u8, cmd, "exlr8")) return exlr8_tests.run(gpu); if (std.mem.eql(u8, cmd, "exlr8-dense")) return exlr8_dense_tests.run(gpu, if (rest.len > 0) rest[0] else null); From 6f583baf91ab3bf06ae20045ef32fa95f604d382 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 15 Oct 2026 10:56:47 +0200 Subject: [PATCH 07/34] ling: the TP=1 layer graph in any loaded format: KDA and MLA mixers, dense and MoE MLPs, head, windows and commits Every projection runs as the qlinear.Weight and every routed-expert layer as the experts.Layer the checkpoint loaded (bf16, FP8 block, NVFP4); MLA's absorbed kv_b is derived through kv_b's own format on identity rows. ling_glue.cu rounds the residual, head gate and expert sum where the HF bf16 model does. Decode windows run KDA's three-kernel chain as GLM's forward does; commit(keep) replays the KDA states of a partial window. ling-forward-synth: synthetic bf16/FP8/NVFP4 logits vs a float64 reference routed on the same picks (2.1-2.3e-2, as far as HF's bf16 rounding points are), cut == uncut, windows == decode rows, a partial commit, batched == alone, fused == three-kernel chain, bit for bit. ling-layer-real: real FP8 layers 0, 4, 7, 71 vs HF 3.0-3.7e-3 without routing flips. GLM files untouched. --- docs/recipes/ling-3.1.md | 107 +++++++--- zig/build/cuda.zig | 1 + zig/kernels/cuda/ling_glue.cu | 39 ++++ zig/src/cuda/kernels.zig | 1 + zig/src/families/ling/block_ffn.zig | 61 ++++++ zig/src/families/ling/block_kda.zig | 91 ++++++++ zig/src/families/ling/block_mla.zig | 59 ++++++ zig/src/families/ling/cuda.zig | 5 + zig/src/families/ling/forward.zig | 118 +++++++++++ zig/src/families/ling/graph.zig | 230 +++++++++++++++++++++ zig/src/families/ling/state.zig | 286 ++++++++++++++++++++++++++ zig/src/families/ling/synth.zig | 2 +- zig/tests/cuda/ling_forward_rig.zig | 130 ++++++++++++ zig/tests/cuda/ling_forward_synth.zig | 178 ++++++++++++++++ zig/tests/cuda/ling_layer_real.zig | 149 ++++++++++++++ zig/tests/cuda/ling_synth.zig | 9 +- zig/tests/cuda/main.zig | 4 + zig/tests/cuda/oracle/ling_forward.py | 186 +++++++++++++++++ zig/tests/cuda/oracle/ling_layer.py | 26 +++ 19 files changed, 1652 insertions(+), 30 deletions(-) create mode 100644 zig/kernels/cuda/ling_glue.cu create mode 100644 zig/src/families/ling/block_ffn.zig create mode 100644 zig/src/families/ling/block_kda.zig create mode 100644 zig/src/families/ling/block_mla.zig create mode 100644 zig/src/families/ling/forward.zig create mode 100644 zig/src/families/ling/graph.zig create mode 100644 zig/src/families/ling/state.zig create mode 100644 zig/tests/cuda/ling_forward_rig.zig create mode 100644 zig/tests/cuda/ling_forward_synth.zig create mode 100644 zig/tests/cuda/ling_layer_real.zig create mode 100644 zig/tests/cuda/oracle/ling_forward.py create mode 100644 zig/tests/cuda/oracle/ling_layer.py diff --git a/docs/recipes/ling-3.1.md b/docs/recipes/ling-3.1.md index 44938e3fde..58d8f9c44f 100644 --- a/docs/recipes/ling-3.1.md +++ b/docs/recipes/ling-3.1.md @@ -13,13 +13,67 @@ block FP8 (e4m3 in 128 x 128 blocks with `weight_scale_inv`). | Part | Where | Checked by | | --- | --- | --- | | Group-limited top-k routing, per-layer shared SwiGLU limit (GLM's bytes unchanged) | `ling-family`, commit 51c01573 | `glm-ref`, `glm-route-old`, `glm-route-ling` | -| MLA at Ling's shape: 64 roped key columns, 128-wide heads, GPT-J rope, head-wise gate | branch `ling-mla`, f8e594f6 | `ling-mla`, `ling-mla-real`, `glm-mla-same` | -| KDA on Ling's layer 0 (q/k/v convs as one q\|k\|v conv), opt-in chunked prompt path | branch `ling-kda`, b55aa13d | `glm-kda-ling`, `glm-kda-chunk` | +| MLA at Ling's shape: 64 roped key columns, 128-wide heads, GPT-J rope, head-wise gate | f8e594f6, merged | `ling-mla`, `ling-mla-real`, `glm-mla-same` | +| KDA on Ling's layer 0 (q/k/v convs as one q\|k\|v conv), opt-in chunked prompt path | b55aa13d, merged | `glm-kda-ling`, `glm-kda-chunk` | | Config loader and weight map (all 214,025 names of the real index) | `zig/src/families/ling/config.zig`, `names.zig` | `zig build test-ling` | | Weight loader at TP=1: bf16, FP8 block, ModelOpt NVFP4 | `zig/src/families/ling/weights.zig`, `weights_proj.zig` | `ling-weights`, `ling-weights-real` | | Synthetic checkpoints in all three formats | `tf-cuda-test ling-synth` | its own name check | - -The MLA and KDA commits live on their own branches and are not merged into `ling-family`. +| The TP=1 layer graph: embedding, KDA/MLA mixers, dense MLP and MoE, final norm, head; windows and commits | `zig/src/families/ling/forward.zig`, `graph.zig`, `state.zig`, `block_*.zig`, `kernels/cuda/ling_glue.cu` | `ling-forward-synth`, `ling-layer-real` | + +## The layer graph + +`forward.Forward` runs rows of one or more sequences (`Seg`: a `State`, its first row, its row count) through the +embedding, every loaded layer and the head; `stack()` runs the layers alone on `Buffers.x`, `layer()` one of them. +`Buffers` are prompt rows or decode rows; each sequence's rows stay pending until `commit(state, keep)`. A prompt +keeps every row; a decode window (at most `window_rows`, 16) keeps any prefix: the KDA states are replayed from the +window's saves (`replay_layers`) and the conv windows shifted, as #804's GLM commit does. The API follows GLM's +`cuda_forward.zig` (`Rank`/`Seg`, `Buffers`/`State`, `commit`) so lanes hooks can follow. + +Every projection is a `qlinear.Weight` and every routed-expert layer an `experts.Layer` in the format the checkpoint +stored; the graph only calls `qlinear.Linear.decode`/`prompt` and `experts.Grouped.decode`/`prompt`. MLA's absorbed +`kv_b` (bf16 `wk`/`wv` per head) is derived at load by running `kv_b` in its own format on identity rows, so any format +`qlinear` serves works. bf16 routed experts (no grouped bf16 kernel) run every expert on every row and the combine +reads each row's picks: correct, for tests, not for speed. + +Per layer, in the HF model's order: input RMSNorm; KDA (q, k, v, beta into one `[q|k|v|b]` row, f and g rows, GLM's +chain with the safe gate, gated output norm, `o_proj`) or MLA (`q_a`, `q_a_layernorm`, `q_b`; `kv_a_proj_with_mqa`, +`kv_a_layernorm` on the 512 latent, the 64 key columns roped into the cache row; absorbed queries, roped attention +over the 576-wide cache, `expand_v`, sigmoid head gate from `g_proj`, `dense`); residual; post-attention RMSNorm; +dense SwiGLU (layers 0-3) or MoE (bf16 router dot in fp32, sigmoid + `expert_bias` with 8 groups keeping 4, +normalized weights x 2.5, grouped experts with the layer's limit, shared expert with its own limit); residual. The +residual stream, every linear's output and the MoE sums are bf16 at the points where the HF bf16 model rounds them +(`ling_glue.cu`: residual add, head gate, expert sum + shared). Decode windows run KDA's three-kernel chain, as +GLM's forward does by default (bit-equal to the fused chain, checked). + +Measured on synthetic checkpoints (4 layers: KDA dense, MLA MoE, KDA MoE, MLA MoE; 80 prompt rows, 8 decode rows), +logits against `oracle/ling_forward.py` routed on the graph's own picks (with free routing, 15-21 of 264 picks +differ from float64's: group near-ties, the worst group sums 0.0014-0.008 apart): + +| Format | vs float64 | vs the bf16 rounding points | bf16 points vs float64 | argmax vs float64 | +| --- | --- | --- | --- | --- | +| bf16 | 2.17e-2 | 1.31e-2 | 2.18e-2 | 86/88 | +| FP8 block | 2.12e-2 | 1.22e-2 | 2.12e-2 | 85/88 | +| NVFP4 | 2.29e-2 | 1.40e-2 | 2.25e-2 | 85/88 | + +(relative L2 of the logits). The graph is as far from float64 as the HF model's own bf16 rounding is. Bit for bit, in +every format: a prompt cut at 48 equals the uncut prompt; an 8-row verify window equals 8 decode rows; keeping 3 of 8 +and running the next 5 leaves every KDA state, conv window and MLA cache row equal to 8 decode rows; two sequences in +one prompt chunk and in one decode window equal each alone; the fused and three-kernel KDA chains agree. + +Real layers of the FP8 release (`ling-layer-real`, the HF bf16 chain's captured input and output, 512 tokens): + +| Layer | Kind | rel L2 vs HF | rows with a different top-8 | rel L2 without them | decode rows vs HF | device peak | +| --- | --- | --- | --- | --- | --- | --- | +| 0 | KDA, dense | 3.27e-3 | - | 3.27e-3 | 3.12e-3 | 0.9 GB | +| 4 | KDA, MoE | 5.88e-3 | 21 | 3.13e-3 | 4.28e-3 | 8.5 GB | +| 7 | MLA, MoE | 4.89e-3 | 17 | 3.05e-3 | 3.57e-3 | 8.3 GB | +| 71 | MLA, MoE, limits 10 | 4.13e-3 | 7 | 3.75e-3 | 3.89e-3 | 8.3 GB | + +exllamav3's port measured 0.4-1.1e-3 without flips (4 flips each on layers 4 and 7) on the same chain: it keeps fp16 +activations, the captured HF chain carries an fp32 residual from layer 3's output on, and this graph rounds the stream +to bf16 as the bf16 model does. On every +real layer a prompt cut at 300 equals the uncut prompt and a 16-row window equals 16 decode rows; for FP8 decode rows +equal the prompt's rows (both run the lane). Loading one MoE layer's 512 experts from the sshfs mount takes ~150 s. ### Formats @@ -42,11 +96,15 @@ per projection. The router must be bf16 (GLM's router kernel reads bf16); a quan ## What is missing -- The layer graph: `weights.Layer` holds one `qlinear.Weight` per projection, while #804's GLM blocks read raw bf16 - pointers (KDA's fused `[q | k | v | f_a | g_a | b]`, MLA's latent projections). Ling's blocks have to call - `qlinear.Linear` per projection, or fuse same-format rows. -- Lanes hooks (drafted, concurrent and cut-prompt windows) for Ling. -- MTP drafting: the MTP layer's tensors load (`weights.Mtp`), nothing drafts with them. +- Lanes hooks (drafted, concurrent and cut-prompt windows) for Ling: the graph has the pieces (segments, windows, + `commit(keep)`), nothing schedules them yet. Sampling, the server and a full-checkpoint run are not wired. +- MTP drafting: the MTP layer's tensors load (`weights.Mtp`); the graph leaves it unwired. It would be + `enorm(embed(next)) | hnorm(hidden)` through `eh_proj`, one MLA + MoE layer on its own cache, the final norm and + the shared head. +- The SwiGLU limit clamps the gate before SiLU (`silu(min(g, 10))`, the grouped expert kernels' and GLM's rule); + the Ling reference clamps after (`min(silu(g), 10)`). They differ by at most 4.5e-4 on a clamped element. +- Speed: nothing is tuned. Projections land in their columns through a byte copy, bf16 routed experts run densely, + and the RMS norms take hidden and `q_lora_rank` widths 512, 1536 and 4096 only (`glue.rmsnorm`). - EXL3, through ashhart's EXL3 module (tracker #832). - Tensor parallelism: the loader is TP=1; shares and transport would come from #804's `cuda_shard.zig` and #855. @@ -58,18 +116,23 @@ TF_LING_INDEX=/path/to/model.safetensors.index.json zig build test-ling # adds zig build -Doptimize=ReleaseFast -Dnvcc=/usr/local/cuda/bin/nvcc zig-out/bin/tf-cuda-test ling-synth /tmp/ling-fp8 --format fp8 # bf16 | fp8 | nvfp4, --layers N (default 4) zig-out/bin/tf-cuda-test ling-weights /tmp/ling-fp8 +python zig/tests/cuda/oracle/ling_forward.py /tmp/ling-fp8 /tmp/ref # float64 reference, free routing +zig-out/bin/tf-cuda-test ling-forward-synth /tmp/ling-fp8 /tmp/ref /tmp/picks # exactness; writes the picks +python zig/tests/cuda/oracle/ling_forward.py /tmp/ling-fp8 /tmp/fref --picks /tmp/picks +zig-out/bin/tf-cuda-test ling-forward-synth /tmp/ling-fp8 /tmp/fref # same picks: logits bounded ``` `ling-synth` writes `config.json` (the real file's schema) and `model.safetensors` with the real key names: hidden 512, -4 heads, 16 experts in 4 groups keeping 2, `q_lora_rank` 256, vocabulary 1,024, layer groups of 2 (KDA, MLA, KDA, -MLA), layer 0 dense, the last layer clamped, plus the MTP layer. KDA's head dim 128 and MLA's 512 latent, 64 rope and -128-wide heads stay as the kernels take them. The FP8 file carries the real exception list's pattern; the NVFP4 file -uses globs and leaves layer 0's dense MLP in bf16, so one checkpoint mixes formats. It checks its names against the -weight map before it passes. Sizes: bf16 282 tensors, 43 MiB; FP8 517, 23 MiB; NVFP4 978, 15 MiB. +4 heads, 16 experts in 4 groups keeping 2, `q_lora_rank` 512 (a width the norm kernel takes), vocabulary 1,024, layer +groups of 2 (KDA, MLA, KDA, MLA), layer 0 dense, the last layer clamped, plus the MTP layer. KDA's head dim 128 and +MLA's 512 latent, 64 rope and 128-wide heads stay as the kernels take them. The FP8 file carries the real exception +list's pattern; the NVFP4 file uses globs and leaves layer 0's dense MLP in bf16, so one checkpoint mixes formats. It +checks its names against the weight map before it passes. FP8 codes stay within 2^-3 to 3.75 so a forward stays in +range. Sizes: bf16 282 tensors, 45 MiB; FP8 517, 24 MiB; NVFP4 978, 16 MiB. `ling-weights` loads every layer and checks every projection (decode and prompt rows, 3 rows) against float64 from the file, the grouped experts' gate-up (SwiGLU) and down, and each small tensor (norms, convs as fp32 q|k|v, `A_log`, -`dt_bias`, router, bias) byte for byte. Measured: bf16 85 projections, FP8 and NVFP4 61 each, worst 0.456 of the bound +`dt_bias`, router, bias) byte for byte. Measured: bf16 85 projections, FP8 and NVFP4 61 each, worst 0.433 of the bound (2^-7 of |y| plus 2^-12 of the sum of |x w|); 37 small tensors equal. With the real checkpoint: @@ -79,7 +142,8 @@ python zig/tests/cuda/oracle/ling_weights.py /path/to/Ling-3.1-flash /tmp/ling-r zig-out/bin/tf-cuda-test ling-weights-real /path/to/Ling-3.1-flash /tmp/ling-ref ``` -It loads layers 0, 4 and 7 with experts 0-3 (769 MiB on the device) and compares 35 products with torch's float64 +`ling-layer-real DIR FIXTURE` runs one whole layer (fixture from `oracle/ling_layer.py CHAIN L OUT`, the captured +HF chain's `h.pt`, `h.pt`, `r.pt`). `ling-weights-real` loads layers 0, 4 and 7 with experts 0-3 (769 MiB on the device) and compares 35 products with torch's float64 on the same 8 rows: worst 0.414 of the bound, relative L2 at most 1.7e-3 for bf16 outputs and 6.9e-8 for the experts' fp32 down rows. @@ -105,13 +169,10 @@ size. Each item stands alone. -1. **KDA block on `qlinear`.** Write Ling's KDA block over `weights.Kda`: seven `qlinear.Linear` calls (or one fused - `[q | k | v]` per format) feeding #804's chain with `conv`, `a_log`, `dt_bias`, `o_norm` as loaded. Check it with - `ling-synth` plus `glm-kda-ling`'s fixture. -2. **MLA block on `qlinear`.** The same for `weights.Mla` on `ling-mla`'s kernels (`q_a`, `q_b`, `kv_a`, `kv_b`, - head gate, `dense`). -3. **MoE block.** Route with GLM's `moe.Layer` (`ling.route()` fills top-k, groups and limits), run `experts.Grouped` - on `Moe.routed`, the shared expert as three projections. +1. **Lanes hooks.** Drive `Forward` from #804's `cuda_lanes.zig` pattern: drafted windows, concurrent streams as + segments, cut prompts; `ling-forward-synth` already checks the exactness each needs. +2. **MTP.** Wire `weights.Mtp` as a drafter (see "What is missing"). +3. **Fused projections.** Same-format q|k|v (and gate|up) rows as one weight per format, without the column copies. 4. **Bf16 routed experts.** Add a grouped bf16 expert kernel, or quantize at load, and drop the per-expert lanes. 5. **TP.** Cut `weights.load` by heads and widths as `glm/cuda_shard.zig` does; FP8 scale blocks span 128 rows and 128 inputs, so FP8 cuts fall on multiples of 128. diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 88b8744638..b263d9b4b4 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -53,6 +53,7 @@ const kernels = [_]Kernel{ .{ .name = "glm_glue", .flags = &.{"-O3"} }, // ours: the Python family's Triton glue by hand, its PTX op for op .{ .name = "glm_ops", .flags = &.{"-O3"} }, // ours: GLM's integer work the Python family did in torch (the DSA token lists) .{ .name = "glm_quant", .flags = torch_ops }, // ours: GLM's draft-only weight quantizers, torch at load in the Python family + .{ .name = "ling_glue", .flags = glue }, // ours: Ling-3.1's bf16 glue at the HF model's rounding points .{ .name = "torch_argmax", .src = "torch_ops/argmax", .flags = torch_ops }, .{ .name = "torch_topk", .src = "torch_ops/topk", .flags = torch_ops }, .{ .name = "torch_pointwise", .src = "torch_ops/pointwise", .flags = torch_ops }, diff --git a/zig/kernels/cuda/ling_glue.cu b/zig/kernels/cuda/ling_glue.cu new file mode 100644 index 0000000000..34a6917017 --- /dev/null +++ b/zig/kernels/cuda/ling_glue.cu @@ -0,0 +1,39 @@ +// Ling-3.1-flash's bf16 glue at the HF model's rounding points: the residual add, the head-wise gate, the MoE sum. +#include +#include + +namespace ling_glue { +__device__ __forceinline__ float bf(__nv_bfloat16 x) { return __bfloat162float(x); } +__device__ __forceinline__ float rbf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } +} // namespace ling_glue + +// torch's bf16 x + y: the fp32 sum rounded once; out may alias x or y. +extern "C" __global__ void tf_ling_add(const __nv_bfloat16* x, const __nv_bfloat16* y, __nv_bfloat16* out, int64_t n) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < n) out[i] = __float2bfloat16_rn(__fadd_rn(ling_glue::bf(x[i]), ling_glue::bf(y[i]))); +} + +// The head-wise gate: o [rows][heads][dv] times bf16(sigmoid(float(g [rows][gs], column h))), rounded to bf16. +extern "C" __global__ void tf_ling_gate(__nv_bfloat16* o, int64_t heads, const __nv_bfloat16* g, int64_t gs) { + const int64_t r = blockIdx.x, h = blockIdx.y, dv = blockDim.x; + const float z = ling_glue::bf(g[r * gs + h]); + const float s = ling_glue::rbf(__fdiv_rn(1.0f, __fadd_rn(1.0f, expf(0.0f - z)))); + __nv_bfloat16* p = o + (r * heads + h) * dv + threadIdx.x; + *p = __float2bfloat16_rn(__fmul_rn(ling_glue::bf(*p), s)); +} + +// moe_infer's sum, bf16(sum of float(y) w over k < used) + shared, y [rows][slots][d] or by_pick [experts][rows][d]. +extern "C" __global__ void tf_ling_combine(const __nv_bfloat16* y, const int* pick, const float* wts, + const __nv_bfloat16* shared, __nv_bfloat16* out, int64_t rows, int64_t d, + int slots, int used, int by_pick) { + const int64_t r = blockIdx.x, c = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; + if (c >= d) return; + float a = 0.0f; + for (int k = 0; k < used; ++k) { + const int64_t at = by_pick ? (static_cast(pick[r * slots + k]) * rows + r) * d + c : (r * slots + k) * d + c; + a = __fadd_rn(a, __fmul_rn(ling_glue::bf(y[at]), wts[r * slots + k])); + } + float v = ling_glue::rbf(a); + if (shared != nullptr) v = __fadd_rn(v, ling_glue::bf(shared[r * d + c])); + out[r * d + c] = __float2bfloat16_rn(v); +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 28e1811a38..34deca60cc 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -40,6 +40,7 @@ pub const glm_glue: []const u8 = if (available) &Blob("fatbin_glm_glue").bytes e pub const glm_ops: []const u8 = if (available) &Blob("fatbin_glm_ops").bytes else &.{}; pub const glm_quant: []const u8 = if (available) &Blob("fatbin_glm_quant").bytes else &.{}; pub const glm_kda: []const u8 = if (available) &Blob("fatbin_glm_kda").bytes else &.{}; +pub const ling_glue: []const u8 = if (available) &Blob("fatbin_ling_glue").bytes else &.{}; pub const glm_kda_chunk: []const u8 = if (available) &Blob("fatbin_glm_kda_chunk").bytes else &.{}; pub const torch_argmax: []const u8 = if (available) &Blob("fatbin_torch_argmax").bytes else &.{}; pub const torch_topk: []const u8 = if (available) &Blob("fatbin_torch_topk").bytes else &.{}; diff --git a/zig/src/families/ling/block_ffn.zig b/zig/src/families/ling/block_ffn.zig new file mode 100644 index 0000000000..23171fbbcf --- /dev/null +++ b/zig/src/families/ling/block_ffn.zig @@ -0,0 +1,61 @@ +//! Ling's MLPs on CUDA: SwiGLU over projections in their loaded formats, and the MoE block (router, groups, experts). +const std = @import("std"); +const cuda = @import("cuda"); +const graph = @import("graph.zig"); +const weights = @import("weights.zig"); +const state = @import("state.zig"); +const ling = @import("cuda.zig"); + +const glue = graph.glue; +const Ops = graph.Ops; +const Config = @import("config.zig").Config; + +/// Gate and up into [gate | up] rows, SwiGLU clamped at `limit` (0: none), down into `out` [rows][hidden]. +pub fn mlp(o: Ops, b: *const state.Buffers, m: weights.Mlp, x: u64, ldx: usize, rows: usize, limit: f32, out: u64) !void { + const w = m.gate.outputs(); + try o.mmCols(m.gate, x, ldx, rows, b.tmp, b.gu, 2 * w, 0); + try o.mmCols(m.up, x, ldx, rows, b.tmp, b.gu, 2 * w, w); + try glue.swiglu(o.set(), o.s, b.gu, b.act, b.xs, limit, rows, w); + try o.mm(m.down, b.act, w, rows, out); +} + +/// The MoE block on b.normed into b.y: routed experts' slots summed (weights fp32), plus the shared expert. +pub fn moe(o: Ops, b: *const state.Buffers, e: weights.Moe, g: *const graph.Graph, layer: u32, rows: usize, picks: ?u64) !void { + const c = &g.w.config; + const D: usize = c.hidden; + const E: usize = c.experts; + if (e.experts != E) return error.PartialExperts; // the router picks among every expert + const top: usize = c.topk; + const slots = top + 1; + const lim = ling.limits(c, layer); + try glue.router(o.set(), o.s, b.normed, D, e.router, b.mlog, b.rpart, rows, D, E); + try glue.select(o.set(), o.s, b.mlog, if (e.bias != 0) e.bias else g.zero_bias, b.pick, b.wts, rows, top, E, c.routed_scale, c.norm_topk, ling.groups(c)); + if (picks) |p| try o.copy(b.pick, slots * 4, p, slots * 4, rows, slots * 4); + switch (e.routed) { + .grouped => |l| try grouped(o, b, l, rows, slots, E), + .lanes => |ls| for (ls, 0..) |x, j| try mlp(o, b, x, b.normed, D, rows, lim.routed, b.ey + j * rows * D * 2), + } + if (e.shared) |s| try mlp(o, b, s, b.normed, D, rows, e.shared_limit, b.sh); + var a: cuda.Args = .{}; + inline for (.{ b.ey, b.pick, b.wts, if (e.shared != null) b.sh else @as(u64, 0), b.y, @as(i64, @intCast(rows)), @as(i64, @intCast(D)), @as(c_int, @intCast(slots)), @as(c_int, @intCast(top)), @as(c_int, @intFromBool(e.routed == .lanes)) }) |v| a.add(v); + try cuda.launch.launch(o.m.combine, .{ .grid = .{ .x = @intCast(rows), .y = @intCast((D + 255) / 256) }, .block = .{ .x = 256 } }, o.s, &a); +} + +/// Grouped experts on every row's top-k slots (the shared slot skipped): gate-up SwiGLU, then down to bf16 slot rows. +fn grouped(o: Ops, b: *const state.Buffers, l: cuda.experts.Layer, rows: usize, slots: usize, count: usize) !void { + const pairs = rows * slots; + const tile: usize = if (o.prompt) cuda.experts.promptTile(l.format) else 16; + const skip: i32 = @intCast(count); + try o.m.router.route(o.s, b.pick, pairs, count + 1, tile, b.plan); + const items = cuda.grouped.maxItems(pairs, count + 1, tile); + const g = o.m.grouped(); + const up: cuda.experts.Rows = .{ .x = b.normed, .stride = l.dims, .slots = slots }; + const down: cuda.experts.Rows = .{ .x = b.eact, .stride = l.width }; + if (o.prompt) { + try g.prompt(o.s, .up, up, l, b.plan, b.eact, items, skip); + try g.prompt(o.s, .down_bf16, down, l, b.plan, b.ey, items, skip); + } else { + try g.decode(o.s, .up, up, l, b.plan, b.eact, items, skip); + try g.decode(o.s, .down_bf16, down, l, b.plan, b.ey, items, skip); + } +} diff --git a/zig/src/families/ling/block_kda.zig b/zig/src/families/ling/block_kda.zig new file mode 100644 index 0000000000..706de08fdf --- /dev/null +++ b/zig/src/families/ling/block_kda.zig @@ -0,0 +1,91 @@ +//! Ling's KDA mixer on CUDA: seven projections in their loaded formats feeding GLM's chain, then o_proj. +const std = @import("std"); +const cuda = @import("cuda"); +const graph = @import("graph.zig"); +const weights = @import("weights.zig"); +const state = @import("state.zig"); + +const kda = graph.kda; +const Ops = graph.Ops; + +/// One sequence's rows of a forward: rows row0 .. row0 + rows on `st` from its position. +pub const Seg = struct { st: *state.State, row0: usize, rows: usize }; + +/// Where segment `sg`'s projection rows and saves go: the State's window rows, or the Buffers' prompt rows. +fn place(b: *const state.Buffers, sg: Seg, li: usize) struct { p: u64, k: u64, v: u64, g: u64, b: u64 } { + const st = sg.st; + if (!b.prompt) { + const s = st.saves(li); + return .{ .p = s.p, .k = s.k, .v = s.v, .g = s.g, .b = s.b }; + } + const n = st.heads * kda.dk; + const r = sg.row0; + return .{ .p = b.kp + r * st.stride * 2, .k = b.saves.k + r * n * 4, .v = b.saves.v + r * n * 2, .g = b.saves.g + r * n * 4, .b = b.saves.b + r * st.heads * 4 }; +} + +/// The block on b.normed into b.y: layer `li` among the KDA layers; prompt rows commit their states and conv windows. +pub fn block(o: Ops, b: *const state.Buffers, a: weights.Kda, c: *const @import("config.zig").Config, li: usize, segs: []const Seg, rows: usize) !void { + const D: usize = c.hidden; + const H: usize = c.kda_heads; + const n = H * kda.dk; + const ps = segs[0].st.stride; + for ([_]weights.Proj{ a.q, a.k, a.v, a.b }, [_]usize{ 0, n, 2 * n, 3 * n }) |w, col0| { + try o.mm(w, b.normed, D, rows, b.tmp); + const w_n = w.outputs(); + for (segs) |sg| try o.copy(b.tmp + sg.row0 * w_n * 2, w_n * 2, place(b, sg, li).p + col0 * 2, ps * 2, sg.rows, w_n * 2); + } + try o.mm(a.f, b.normed, D, rows, b.ka); + try o.mm(a.g, b.normed, D, rows, b.kg); + for (segs) |sg| { + const st = sg.st; + const at = place(b, sg, li); + const cur = st.cur[li]; + const ch: kda.Chain = .{ + .p = at.p, + .p_stride = ps, + .b_off = 3 * n, + .a = b.ka + sg.row0 * n * 2, + .a_stride = n, + .g = b.kg + sg.row0 * n * 2, + .g_stride = n, + .conv_state = st.convAt(li), + .conv_w = a.conv, + .state_in = st.recAt(cur, li), + .a_log = a.a_log, + .dt_bias = a.dt_bias, + .norm_w = a.o_norm, + .eps = c.eps, + .lower = c.lower_bound, + .heads = H, + }; + const sv: kda.Saves = .{ .out = b.saves.out + sg.row0 * n * 2, .state_out = st.recAt(1 - cur, li), .k = at.k, .v = at.v, .g = at.g, .b = at.b }; + const wide: ?kda.Wide = if (if (b.prompt) sg.rows >= kda.wide_rows else o.kda_wide) .{ .q = b.wide.q + sg.row0 * n * 4, .y = b.wide.y + sg.row0 * n * 2 } else null; + try kda.chain(&o.m.kda, o.s, ch, sg.rows, sv, wide); + if (!b.prompt) continue; + try graph.glue.convShift(o.set(), o.s, st.convAt(li), at.p, sg.rows, sg.rows * ps, ps, 1, 3 * n); + st.cur[li] = 1 - cur; // a prompt keeps every row: the layer commits now + } + try o.mm(a.o, b.saves.out, n, rows, b.y); +} + +/// A window's commit of `keep` rows on every KDA layer: the states replayed when rows are dropped, the conv windows. +pub fn commit(o: Ops, st: *state.State, keep: usize) !void { + if (st.kda_layers == 0) return; + const cur = st.cur[0]; + for (st.cur) |x| if (x != cur) return error.KdaLayersOutOfStep; + const h = st.heads; + if (keep < st.pending) { + var a: cuda.Args = .{}; + a.add(@as(c_int, @intCast(h))); + a.add(st.recAt(cur, 0)); + a.add(@as(usize, st.rec_layer / 4)); + for ([_]u64{ st.save_k, st.save_v, st.save_g, st.save_b }) |x| a.add(x); + a.add(@as(usize, state.window_rows * h * kda.dk)); + a.add(@as(usize, state.window_rows * h)); + a.add(@as(c_int, @intCast(keep))); + a.add(st.recAt(1 - cur, 0)); + try cuda.launch.launch(o.m.kda.replay_layers, .{ .grid = .{ .x = @intCast(st.kda_layers * h) }, .block = .{ .x = 1024 } }, o.s, &a); + } + for (st.cur) |*x| x.* = 1 - x.*; + try graph.glue.convShift(o.set(), o.s, st.conv, st.proj, keep, state.window_rows * st.stride, st.stride, st.kda_layers, st.convWidth()); +} diff --git a/zig/src/families/ling/block_mla.zig b/zig/src/families/ling/block_mla.zig new file mode 100644 index 0000000000..bf76aaf2f1 --- /dev/null +++ b/zig/src/families/ling/block_mla.zig @@ -0,0 +1,59 @@ +//! Ling's MLA mixer on CUDA, absorbed: projections in their loaded formats, roped latent cache, head gate, `dense`. +const std = @import("std"); +const graph = @import("graph.zig"); +const weights = @import("weights.zig"); +const state = @import("state.zig"); +const block_kda = @import("block_kda.zig"); + +const glue = graph.glue; +const latent = graph.latent; +const Ops = graph.Ops; +const Config = @import("config.zig").Config; +const kw = state.cache_width; +const lw = 512; + +/// The block on b.normed into b.y: q (lora pair or q_proj), latent and roped key into the cache, attention, gate. +pub fn block(o: Ops, b: *const state.Buffers, a: weights.Mla, ab: graph.Absorb, g: *const graph.Graph, mi: usize, segs: []const block_kda.Seg, rows: usize) !void { + const c = &g.w.config; + const D: usize = c.hidden; + const H: usize = c.mla_heads; + const qk: usize = c.qkDim(); + const nope: usize = c.nope; + const dv: usize = c.v_dim; + if (c.kv_lora != lw or nope != 128 or dv != 128 or c.rope_dim != latent.rope_width) return error.LatentShape; + if (a.q_a) |qa| { + const ql: usize = c.q_lora; + try o.mm(qa, b.normed, D, rows, b.qa_in); + try glue.rmsnorm(o.set(), o.s, b.qa_in, ql, a.q_a_norm, b.qr, ql, null, c.eps, rows, ql); + try o.mm(a.q_b.?, b.qr, ql, rows, b.q); + } else try o.mm(a.q.?, b.normed, D, rows, b.q); + try o.mm(a.kv_a, b.normed, D, rows, b.kv); + try glue.rmsnorm(o.set(), o.s, b.kv, kw, a.kv_a_norm, b.lat, kw, null, c.eps, rows, lw); + if (a.gate) |gp| try o.mm(gp, b.normed, D, rows, b.g); + try latent.absorbQStrided(o.set(), o.s, b.q, qk, ab.wk, b.qa, kw, rows, H, nope, lw); + const scale = c.smScale(); + for (segs) |sg| { + const st = sg.st; + const r = sg.row0; + if (st.pos + sg.rows > st.capacity) return error.CacheFull; + try latent.rope(o.set(), o.s, .{ .x = b.q + (r * H * qk + nope) * 2, .x_row = H * qk, .x_head = qk, .out = b.qa + (r * H * kw + lw) * 2, .o_row = H * kw, .o_head = kw }, sg.rows, H, st.pos_dev, g.freq); + try latent.rope(o.set(), o.s, .{ .x = b.kv + (r * kw + lw) * 2, .x_row = kw, .x_head = 0, .out = b.lat + (r * kw + lw) * 2, .o_row = kw, .o_head = 0 }, sg.rows, 1, st.pos_dev, g.freq); + try latent.latentWrite(o.set(), o.s, b.lat + r * kw * 2, kw, st.cacheAt(mi), st.pos_dev, sg.rows, kw); + var r0: usize = 0; + while (r0 < sg.rows) : (r0 += state.att_rows) { + const m = @min(state.att_rows, sg.rows - r0); + var at = st.pos_dev; + if (r0 > 0) { + try o.addI32(st.pos_dev, r0, b.at); + at = b.at; + } + const nch = latent.chunksFor(st.pos + r0 + m); + if (nch > b.nch) return error.CacheFull; + const row = r + r0; + try latent.attentionRoped(o.set(), o.s, b.qa + row * H * kw * 2, st.cacheAt(mi), at, b.partials, b.ol + row * H * lw * 2, b.dummy, m, H, lw, nch, scale); + } + } + try latent.expandV(o.set(), o.s, b.ol, ab.wv, b.vn, rows, H, dv, lw); + if (a.gate != null) try o.gate(b.vn, b.g, H, rows, H, dv); + try o.mm(a.o, b.vn, H * dv, rows, b.y); +} diff --git a/zig/src/families/ling/cuda.zig b/zig/src/families/ling/cuda.zig index 5d534ed0b0..095810d32a 100644 --- a/zig/src/families/ling/cuda.zig +++ b/zig/src/families/ling/cuda.zig @@ -5,6 +5,9 @@ pub const config = @import("config.zig"); pub const names = @import("names.zig"); pub const weights = @import("weights.zig"); pub const synth = @import("synth.zig"); +pub const graph = @import("graph.zig"); +pub const state = @import("state.zig"); +pub const forward = @import("forward.zig"); const Config = config.Config; @@ -39,6 +42,8 @@ test { _ = @import("synth.zig"); std.testing.refAllDecls(weights); // the loader compiles on hosts without a GPU std.testing.refAllDecls(@import("weights_proj.zig")); + std.testing.refAllDecls(forward); // the layer graph compiles on hosts without a GPU + std.testing.refAllDecls(forward.Forward); } test "Ling's routing as glue.select takes it: 8 groups keeping 4, and select's own bounds hold" { diff --git a/zig/src/families/ling/forward.zig b/zig/src/families/ling/forward.zig new file mode 100644 index 0000000000..a10072b41c --- /dev/null +++ b/zig/src/families/ling/forward.zig @@ -0,0 +1,118 @@ +//! Ling-3.1-flash's TP=1 forward on CUDA: embedding, the layer stack, final norm and head, and a window's commit. +const std = @import("std"); +const cuda = @import("cuda"); +const graph = @import("graph.zig"); +const weights = @import("weights.zig"); +const state = @import("state.zig"); +const block_kda = @import("block_kda.zig"); +const block_mla = @import("block_mla.zig"); +const block_ffn = @import("block_ffn.zig"); + +const glue = graph.glue; +pub const Seg = block_kda.Seg; +pub const Modules = graph.Modules; +pub const Graph = graph.Graph; +pub const Buffers = state.Buffers; +pub const State = state.State; +pub const window_rows = state.window_rows; +pub const cache_width = state.cache_width; + +pub const Forward = struct { + m: *const graph.Modules, + g: *const graph.Graph, + s: cuda.Stream, + b: *const state.Buffers, + picks: ?u64 = null, // i32 [MoE layers][rows][top_k + 1]: every MoE layer's picks, for checks + kda_wide: bool = true, // GLM's TF_GLM_KDA_DECODE_WIDE default: decode windows on the three-kernel chain + + pub fn ops(f: *const Forward) graph.Ops { + return .{ .m = f.m, .s = f.s, .prompt = f.b.prompt, .part = f.b.part, .kda_wide = f.kda_wide }; + } + + /// Segments in row order from row 0, each on its own committed sequence; a window's at most window_rows rows. + fn rowsOf(f: *const Forward, segs: []const Seg) !usize { + var rows: usize = 0; + for (segs, 0..) |sg, i| { + if (sg.row0 != rows or sg.rows == 0 or sg.st.pending != 0) return error.BadSegments; + if (!f.b.prompt and sg.rows > window_rows) return error.WindowTooWide; + for (segs[0..i]) |o| if (o.st == sg.st) return error.BadSegments; + rows += sg.rows; + } + if (rows == 0 or rows > f.b.rows) return error.TooManyRows; + return rows; + } + + /// `ids` (one per row of `segs`) through the embedding, every loaded layer, the final norm and the head. + pub fn run(f: *const Forward, segs: []const Seg, ids: []const i32) !void { + const rows = try f.rowsOf(segs); + if (ids.len != rows) return error.TooManyRows; + try f.embed(ids); + try f.stack(segs); + try f.head(rows); + } + + /// The tokens' bf16 embedding rows into b.x. + pub fn embed(f: *const Forward, ids: []const i32) !void { + const w = f.g.w; + const d = f.b.arena.buf.d; + if (ids.len > f.b.rows) return error.TooManyRows; + for (ids) |t| if (t < 0 or t >= w.config.vocab) return error.BadToken; + try d.check(d.api.cuMemcpyHtoDAsync_v2(f.b.ids, ids.ptr, ids.len * 4, f.s.handle), "cuMemcpyHtoDAsync"); + try glue.embedB16(f.ops().set(), f.s, f.b.ids, w.embed, f.b.x, ids.len, w.config.hidden, 1); + } + + /// Every loaded layer on b.x in place; each sequence then holds its rows pending until commit(). + pub fn stack(f: *const Forward, segs: []const Seg) !void { + const rows = try f.rowsOf(segs); + var moe: usize = 0; + for (f.g.w.layers, 0..) |*l, j| { + const slots: usize = f.g.w.config.topk + 1; + const p: ?u64 = if (f.picks) |x| x + moe * rows * slots * 4 else null; + try f.layer(segs, j, rows, p); + moe += @intFromBool(l.ffn == .moe); + } + for (segs) |sg| { + sg.st.pending = sg.rows; + sg.st.prompt = f.b.prompt; + } + } + + /// Loaded layer j: input norm, mixer, residual; post-attention norm, MLP or MoE, residual. + pub fn layer(f: *const Forward, segs: []const Seg, j: usize, rows: usize, picks: ?u64) !void { + const o = f.ops(); + const b = f.b; + const g = f.g; + const c = &g.w.config; + const D: usize = c.hidden; + const l = g.w.layers[j]; + try glue.rmsnorm(o.set(), o.s, b.x, D, l.attn_norm, b.normed, D, null, c.eps, rows, D); + switch (l.mixer) { + .kda => |a| try block_kda.block(o, b, a, c, g.slot[j], segs, rows), + .mla => |a| try block_mla.block(o, b, a, g.absorb[g.slot[j]], g, g.slot[j], segs, rows), + } + try o.add(b.x, b.y, b.x, rows * D); + try glue.rmsnorm(o.set(), o.s, b.x, D, l.mlp_norm, b.normed, D, null, c.eps, rows, D); + switch (l.ffn) { + .dense => |m| try block_ffn.mlp(o, b, m, b.normed, D, rows, 0, b.y), + .moe => |e| try block_ffn.moe(o, b, e, g, l.index, rows, picks), + } + try o.add(b.x, b.y, b.x, rows * D); + } + + /// The final norm and lm_head of b.x's rows into b.logits (bf16 [rows][vocab]). + pub fn head(f: *const Forward, rows: usize) !void { + const o = f.ops(); + const w = f.g.w; + const D: usize = w.config.hidden; + try glue.rmsnorm(o.set(), o.s, f.b.x, D, w.norm, f.b.normed, D, null, w.config.eps, rows, D); + try o.mm(w.head orelse return error.TiedHeadNotBuilt, f.b.normed, D, rows, f.b.logits); + } + + /// Keep the first `keep` rows of the sequence's last forward (a prompt keeps them all): KDA states, position. + pub fn commit(f: *const Forward, st: *state.State, keep: usize) !void { + if (keep == 0 or keep > st.pending or (st.prompt and keep != st.pending)) return error.BadCommit; + if (!st.prompt) try block_kda.commit(f.ops(), st, keep); + st.pending = 0; + try st.setPos(f.s, st.pos + keep); + } +}; diff --git a/zig/src/families/ling/graph.zig b/zig/src/families/ling/graph.zig new file mode 100644 index 0000000000..3296fe18fe --- /dev/null +++ b/zig/src/families/ling/graph.zig @@ -0,0 +1,230 @@ +//! Ling-3.1-flash's TP=1 layer graph on CUDA: the kernels it launches, the launches its blocks share, derived weights. +const std = @import("std"); +const cuda = @import("cuda"); +const glm = @import("glm_cuda"); +const weights = @import("weights.zig"); + +pub const glue = glm.glue; +pub const kda = glm.kda; +pub const latent = glm.latent; +const Proj = weights.Proj; +const Config = @import("config.zig").Config; + +/// Every module the graph launches from: GLM's glue, KDA and ops, the lane formats, the grouped experts, our glue. +pub const Modules = struct { + mods: [5]cuda.Module, // glm_glue, glm_kda, glm_ops, experts (the plan kernels), ling_glue + glue: glue.Kernels, + kda: kda.Kernels, + lane: cuda.qmmf.Lane, + prompt: cuda.nvfp4.Prompt, + x8: cuda.experts.Fp8Experts, + x4: cuda.experts.Nvfp4Experts, + router: cuda.grouped.Router, + copy_rows: cuda.Function, + add_i32: cuda.Function, + add: cuda.Function, + gate: cuda.Function, + combine: cuda.Function, + + pub fn load(d: *const cuda.Driver, ctx: *const cuda.Context) !Modules { + const k = cuda.kernels; + var m: Modules = undefined; + var n: usize = 0; + errdefer for (m.mods[0..n]) |*x| x.unload(); + for (&m.mods, [_][]const u8{ k.glm_glue, k.glm_kda, k.glm_ops, k.experts, k.ling_glue }) |*x, image| { + x.* = try cuda.Module.load(d, image); + n += 1; + } + const sms: usize = @intCast(try ctx.attribute(.multiprocessor_count)); + m.glue = try glue.Kernels.load(m.mods[0]); + m.kda = try kda.Kernels.load(m.mods[1]); + m.copy_rows = try m.mods[2].function("tf_glm_copy_rows"); + m.add_i32 = try m.mods[2].function("tf_glm_add_i32"); + m.router = try cuda.grouped.Router.resolve(m.mods[3]); + m.add = try m.mods[4].function("tf_ling_add"); + m.gate = try m.mods[4].function("tf_ling_gate"); + m.combine = try m.mods[4].function("tf_ling_combine"); + m.lane = try cuda.qmmf.Lane.load(d, @intCast(try ctx.attribute(.compute_capability_major))); + errdefer m.lane.unload(); + m.prompt = try cuda.nvfp4.Prompt.load(d); + errdefer m.prompt.unload(); + m.x8 = try cuda.experts.Fp8Experts.load(d, sms); + errdefer m.x8.unload(); + m.x4 = try cuda.experts.Nvfp4Experts.load(d, sms); + return m; + } + + pub fn deinit(m: *Modules) void { + for (&m.mods) |*x| x.unload(); + m.lane.unload(); + m.prompt.unload(); + m.x8.unload(); + m.x4.unload(); + m.* = undefined; + } + + pub fn set(m: *const Modules) glue.Set { + return .{ .glue = &m.glue }; + } + + /// Every lane format's decode rows; nvfp4 prompt rows on the prompt GEMM. + pub fn linear(m: *const Modules) cuda.qlinear.Linear { + return .{ .lane = &m.lane, .nvfp4_prompt = &m.prompt }; + } + + pub fn grouped(m: *const Modules) cuda.experts.Grouped { + return .{ .fp8g = &m.x8, .nvfp4 = &m.x4 }; + } +}; + +fn cdiv(a: usize, b: usize) u32 { + return @intCast((a + b - 1) / b); +} + +/// One forward's launches: the modules, the stream, prompt or decode kernels, the lane's split-K partials. +pub const Ops = struct { + m: *const Modules, + s: cuda.Stream, + prompt: bool, + part: u64, + kda_wide: bool = true, // decode windows on KDA's three-kernel chain (bit-equal to the fused one, faster) + + pub fn set(o: Ops) glue.Set { + return o.m.set(); + } + + /// out [rows][n] bf16 = x rows (`ldx` apart) through `w` in whatever format it loaded. + pub fn mm(o: Ops, w: Proj, x: u64, ldx: usize, rows: usize, out: u64) !void { + const in: cuda.qlinear.Rows = .{ .x = x, .ldx = ldx, .part = o.part }; + if (o.prompt) return o.m.linear().prompt(o.s, w, in, out, rows); + return o.m.linear().decode(o.s, w, in, out, rows); + } + + /// Bytes [rows][row_bytes] from rows `src_pitch` apart into rows `dst_pitch` apart, as they are. + pub fn copy(o: Ops, src: u64, src_pitch: usize, dst: u64, dst_pitch: usize, rows: usize, row_bytes: usize) !void { + var a: cuda.Args = .{}; + inline for (.{ src, @as(c_longlong, @intCast(src_pitch)), dst, @as(c_longlong, @intCast(dst_pitch)), @as(c_int, @intCast(rows)), @as(c_int, @intCast(row_bytes)) }) |v| a.add(v); + try cuda.launch.launch(o.m.copy_rows, .{ .grid = .{ .x = cdiv(rows * row_bytes, 256) }, .block = .{ .x = 256 } }, o.s, &a); + } + + /// A projection into columns col0 .. col0 + n of rows `stride` wide (bf16), through `tmp`. + pub fn mmCols(o: Ops, w: Proj, x: u64, ldx: usize, rows: usize, tmp: u64, dst: u64, stride: usize, col0: usize) !void { + const n = w.outputs(); + try o.mm(w, x, ldx, rows, tmp); + try o.copy(tmp, n * 2, dst + col0 * 2, stride * 2, rows, n * 2); + } + + /// out = bf16(x + y), n bf16 elements; out may alias either. + pub fn add(o: Ops, x: u64, y: u64, out: u64, n: usize) !void { + var a: cuda.Args = .{}; + inline for (.{ x, y, out, @as(i64, @intCast(n)) }) |v| a.add(v); + try cuda.launch.launch(o.m.add, .{ .grid = .{ .x = cdiv(n, 256) }, .block = .{ .x = 256 } }, o.s, &a); + } + + /// The head-wise gate: v [rows][heads][dv] times bf16 sigmoid of g [rows][g_stride] column h. + pub fn gate(o: Ops, v: u64, g: u64, g_stride: usize, rows: usize, heads: usize, dv: usize) !void { + var a: cuda.Args = .{}; + inline for (.{ v, @as(i64, @intCast(heads)), g, @as(i64, @intCast(g_stride)) }) |x| a.add(x); + try cuda.launch.launch(o.m.gate, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(heads) }, .block = .{ .x = @intCast(dv) } }, o.s, &a); + } + + /// The device i32 at `pos` plus `r` into `out`. + pub fn addI32(o: Ops, pos: u64, r: usize, out: u64) !void { + var a: cuda.Args = .{}; + inline for (.{ pos, @as(c_int, @intCast(r)), out }) |v| a.add(v); + try cuda.launch.launch(o.m.add_i32, .{ .grid = .{ .x = 1 }, .block = .{ .x = 32 } }, o.s, &a); + } +}; + +/// An MLA layer's kv_b as the absorbed path reads it, bf16: keys wk [heads][nope][lw] and values wv [heads][v][lw]. +pub const Absorb = struct { wk: u64, wv: u64 }; + +/// What the graph derives from the loaded weights: absorbed kv_b, each layer's slot, the rope table, a zero bias. +pub const Graph = struct { + w: *const weights.Weights, + store: weights.Store, + slot: []u32, // loaded layer j's index among the loaded KDA layers, or among the MLA layers + absorb: []Absorb, // by MLA slot + kda_layers: usize, + mla_layers: usize, + freq: u64, // fp32 [32]: the rope's inv_freq + zero_bias: u64, // fp32 [experts]: select's bias without moe_router_enable_expert_bias + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, m: *const Modules, s: cuda.Stream, w: *const weights.Weights) !Graph { + const c = &w.config; + var g: Graph = .{ .w = w, .store = .{ .gpa = gpa, .d = d }, .slot = try gpa.alloc(u32, w.layers.len), .absorb = &.{}, .kda_layers = 0, .mla_layers = 0, .freq = 0, .zero_bias = 0 }; + errdefer g.deinit(); + for (w.layers, g.slot) |l, *x| switch (l.mixer) { + .kda => { + x.* = @intCast(g.kda_layers); + g.kda_layers += 1; + }, + .mla => { + x.* = @intCast(g.mla_layers); + g.mla_layers += 1; + }, + }; + if (c.rope_dim != latent.rope_width or c.rope_style != .gptj) return error.RopeNotBuilt; + const f = latent.ropeFreqs(c.rope_theta); + g.freq = try g.store.keep(std.mem.sliceAsBytes(&f), "ling.rope.inv_freq"); + const zero = try gpa.alloc(u8, @as(usize, c.experts) * 4); + defer gpa.free(zero); + @memset(zero, 0); + g.zero_bias = try g.store.keep(zero, "ling.zero_bias"); + g.absorb = try gpa.alloc(Absorb, g.mla_layers); + for (w.layers, g.slot) |l, x| if (l.mixer == .mla) { + g.absorb[x] = try absorbOf(gpa, &g.store, m, s, c, l.mixer.mla.kv_b, l.index); + }; + return g; + } + + pub fn deinit(g: *Graph) void { + const gpa = g.store.gpa; + gpa.free(g.slot); + gpa.free(g.absorb); + g.store.deinit(); + g.* = undefined; + } +}; + +/// kv_b in its loaded format through the lane on identity rows: W^T [lw][heads (nope + v)], regrouped by head. +fn absorbOf(gpa: std.mem.Allocator, store: *weights.Store, m: *const Modules, s: cuda.Stream, c: *const Config, kv_b: Proj, index: u32) !Absorb { + const lw: usize = c.kv_lora; + const h: usize = c.mla_heads; + const nope: usize = c.nope; + const dv: usize = c.v_dim; + const per = nope + dv; + const n = h * per; + if (kv_b.outputs() != n or kv_b.inputs() != lw) return error.UnexpectedTensor; + const eye = try gpa.alloc(u16, lw * lw); + defer gpa.free(eye); + @memset(eye, 0); + for (0..lw) |i| eye[i * lw + i] = 0x3f80; // bf16 1.0 + var x = try cuda.DeviceBuffer.fromHost(store.d, std.mem.sliceAsBytes(eye)); + defer x.free(); + var y = try cuda.DeviceBuffer.alloc(store.d, lw * n * 2); + defer y.free(); + var part = try cuda.DeviceBuffer.alloc(store.d, @max(16, m.lane.partBytes(lw, switch (kv_b) { + .fp8g, .nvfp4, .fp8, .bf16 => |q| q, + else => return error.FormatNotBuilt, + }))); + defer part.free(); + try m.linear().decode(s, kv_b, .{ .x = x.ptr, .ldx = lw, .part = part.ptr }, y.ptr, lw); + try s.synchronize(); + const t = try gpa.alloc(u16, lw * n); + defer gpa.free(t); + try y.download(0, std.mem.sliceAsBytes(t)); + const wk = try gpa.alloc(u16, h * nope * lw); + defer gpa.free(wk); + const wv = try gpa.alloc(u16, h * dv * lw); + defer gpa.free(wv); + for (0..h) |hd| for (0..per) |j| for (0..lw) |col| { + const v = t[col * n + hd * per + j]; + if (j < nope) wk[(hd * nope + j) * lw + col] = v else wv[(hd * dv + j - nope) * lw + col] = v; + }; + var b: [64]u8 = undefined; + return .{ + .wk = try store.keep(std.mem.sliceAsBytes(wk), try std.fmt.bufPrint(&b, "ling.{d}.absorb.wk", .{index})), + .wv = try store.keep(std.mem.sliceAsBytes(wv), try std.fmt.bufPrint(&b, "ling.{d}.absorb.wv", .{index})), + }; +} diff --git a/zig/src/families/ling/state.zig b/zig/src/families/ling/state.zig new file mode 100644 index 0000000000..50505e696d --- /dev/null +++ b/zig/src/families/ling/state.zig @@ -0,0 +1,286 @@ +//! Ling-3.1-flash's CUDA memory at TP=1: a forward's Buffers (any sequences) and one sequence's State, an arena each. +const std = @import("std"); +const cuda = @import("cuda"); +const graph = @import("graph.zig"); +const weights = @import("weights.zig"); + +const kda = graph.kda; +const latent = graph.latent; + +/// A verify window's most rows on one sequence: the KDA saves and projection rows a commit replays. +pub const window_rows = 16; +/// The most rows of one roped attention call (a prompt chunk's rows go in blocks of these). +pub const att_rows = 128; +/// MLA's cache row: the 512-wide latent, then the 64 roped key columns. +pub const cache_width = 512 + latent.rope_width; + +/// Sub-allocations of one zeroed device allocation, 256-byte aligned. +const Arena = struct { + buf: cuda.DeviceBuffer, + + fn init(d: *const cuda.Driver, sizes: []const usize, out: []u64) !Arena { + var total: usize = 0; + for (sizes) |n| total += std.mem.alignForward(usize, n, 256); + const buf = try cuda.DeviceBuffer.alloc(d, @max(total, 256)); + errdefer { + var b = buf; + b.free(); + } + try buf.fill8(0, null); + var at: usize = 0; + for (sizes, out) |n, *p| { + p.* = buf.ptr + at; + at += std.mem.alignForward(usize, n, 256); + } + return .{ .buf = buf }; + } +}; + +/// The KDA projection row: q | k | v (3 heads 128), then beta's logits, padded to 8 columns. +pub fn kdaStride(w: *const weights.Weights) usize { + const h: usize = w.config.kda_heads; + return 3 * h * kda.dk + std.mem.alignForward(usize, h, 8); +} + +/// The largest split-K partials any loaded projection takes at 1 .. rows rows. +fn partBytes(m: *const graph.Modules, w: *const weights.Weights, rows: usize) usize { + var most: usize = 16; + const one = struct { + fn of(mo: *const graph.Modules, p: weights.Proj, r: usize, best: *usize) void { + const q = switch (p) { + .fp8g, .nvfp4, .fp8, .bf16 => |x| x, + else => return, + }; + for (1..r + 1) |k| best.* = @max(best.*, mo.lane.partBytes(k, q)); + } + }; + for (w.layers) |l| { + switch (l.mixer) { + .kda => |a| for ([_]weights.Proj{ a.q, a.k, a.v, a.b, a.f, a.g, a.o }) |p| one.of(m, p, rows, &most), + .mla => |a| for ([_]?weights.Proj{ a.q, a.q_a, a.q_b, a.kv_a, a.gate, a.o }) |p| if (p) |x| one.of(m, x, rows, &most), + } + const mlps: [2]?weights.Mlp = switch (l.ffn) { + .dense => |x| .{ x, null }, + .moe => |e| .{ e.shared, if (e.routed == .lanes and e.routed.lanes.len > 0) e.routed.lanes[0] else null }, + }; + for (mlps) |x| if (x) |p| for ([_]weights.Proj{ p.gate, p.up, p.down }) |q| one.of(m, q, rows, &most); + } + if (w.head) |hd| one.of(m, hd, rows, &most); + return most; +} + +/// A forward's scratch for up to `rows` rows on any sequences: prompt rows (`prompt`) or decode rows. +pub const Buffers = struct { + arena: Arena, + rows: usize, + prompt: bool, + nch: usize, // the most 512-key chunks one attention call reads + ids: u64, // i32 [rows] + x: u64, // bf16 [rows][hidden]: the residual stream + normed: u64, + y: u64, // a block's output rows + sh: u64, // the shared expert's rows + tmp: u64, // one projection's rows before they go into their columns + part: u64, // the lane's split-K partials + kp: u64, // bf16 [rows][kdaStride]: a prompt's KDA projection rows (a window's stay in the State) + ka: u64, // bf16 [rows][kda n]: the decay's projection + kg: u64, // bf16 [rows][kda n]: the output gate's projection + saves: kda.Saves, // out bf16; a prompt's k, v, g, b (a window's stay in the State); state_out set per layer + wide: kda.Wide, + qa_in: u64, // bf16 [rows][q_lora] + qr: u64, // bf16 [rows][q_lora]: q_a_layernorm's rows + q: u64, // bf16 [rows][heads 192] + kv: u64, // bf16 [rows][576]: kv_a_proj_with_mqa's rows + lat: u64, // bf16 [rows][576]: normed latent, roped key columns + qa: u64, // bf16 [rows][heads][576]: absorbed then roped queries + ol: u64, // bf16 [rows][heads][512] + vn: u64, // bf16 [rows][heads 128] + g: u64, // bf16 [rows][heads]: g_proj's rows + partials: latent.Partials, + dummy: u64, + at: u64, // i32: an attention block's first position + gu: u64, // bf16 [rows][2 width]: [gate | up] + act: u64, // bf16 [rows][width] + xs: u64, // fp32 [rows][width / 64] + mlog: u64, // fp32 [rows][experts] + rpart: u64, // fp32 [router_ks][rows][experts] + pick: u64, // i32 [rows][top_k + 1] + wts: u64, // fp32 [rows][top_k + 1] + plan: cuda.grouped.Plan, + eact: u64, // bf16 [pairs][moe_inter] + ey: u64, // bf16 [pairs][hidden], or every expert's rows [experts][rows][hidden] for per-expert lanes + logits: u64, // bf16 [rows][vocab] + + pub fn init(d: *const cuda.Driver, m: *const graph.Modules, w: *const weights.Weights, rows: usize, prompt: bool, capacity: usize) !Buffers { + const c = &w.config; + const D: usize = c.hidden; + const H: usize = c.kda_heads; + const n = H * kda.dk; + const HM: usize = c.mla_heads; + const ql: usize = @max(c.q_lora, 1); + const E: usize = c.experts; + const slots: usize = c.topk + 1; + const pairs = rows * slots; + const width: usize = @max(@max(c.dense_inter, c.shared_inter), c.moe_inter); + var lanes = false; + for (w.layers) |l| if (l.ffn == .moe and l.ffn.moe.routed == .lanes) { + lanes = true; + }; + const ar = @min(rows, att_rows); + const nch = latent.chunksFor(capacity); + const hist = if (pairs > cuda.grouped.small) (pairs + cuda.grouped.small - 1) / cuda.grouped.small * (E + 1) else 1; + const sizes = [_]usize{ + rows * 4, rows * D * 2, rows * D * 2, rows * D * 2, + rows * D * 2, rows * @max(n, width) * 2, partBytes(m, w, rows), if (prompt) rows * kdaStride(w) * 2 else 16, + rows * n * 2, rows * n * 2, rows * n * 2, rows * n * 4, + rows * n * 2, rows * n * 4, rows * H * 4, rows * n * 4, + rows * n * 2, rows * ql * 2, rows * ql * 2, rows * HM * c.qkDim() * 2, + rows * cache_width * 2, rows * cache_width * 2, rows * HM * cache_width * 2, rows * HM * 512 * 2, + rows * HM * c.v_dim * 2, rows * HM * 2, nch * ar * HM * 512 * 4, nch * ar * HM * 4, + nch * ar * HM * 4, 4, 4, rows * 2 * width * 2, + rows * width * 2, rows * (width / 64 + 1) * 4, rows * E * 4, glue_router_ks * rows * E * 4, + pairs * 4, pairs * 4, pairs * 4, cuda.grouped.maxItems(pairs, E + 1, 16) * 12, + 8, pairs * 4, hist * 4, pairs * c.moe_inter * 2, + @max(pairs, if (lanes) E * rows else 0) * D * 2, rows * @as(usize, c.vocab) * 2, + }; + var p: [sizes.len]u64 = undefined; + const a = try Arena.init(d, &sizes, &p); + return .{ + .arena = a, + .rows = rows, + .prompt = prompt, + .nch = nch, + .ids = p[0], + .x = p[1], + .normed = p[2], + .y = p[3], + .sh = p[4], + .tmp = p[5], + .part = p[6], + .kp = p[7], + .ka = p[8], + .kg = p[9], + .saves = .{ .out = p[10], .state_out = 0, .k = p[11], .v = p[12], .g = p[13], .b = p[14] }, + .wide = .{ .q = p[15], .y = p[16] }, + .qa_in = p[17], + .qr = p[18], + .q = p[19], + .kv = p[20], + .lat = p[21], + .qa = p[22], + .ol = p[23], + .vn = p[24], + .g = p[25], + .partials = .{ .po = p[26], .pm = p[27], .pl = p[28] }, + .dummy = p[29], + .at = p[30], + .gu = p[31], + .act = p[32], + .xs = p[33], + .mlog = p[34], + .rpart = p[35], + .pick = p[36], + .wts = p[37], + .plan = .{ .members = p[38], .items = p[39], .counts = p[40], .rank = p[41], .hist = p[42] }, + .eact = p[43], + .ey = p[44], + .logits = p[45], + }; + } + + pub fn deinit(b: *Buffers) void { + b.arena.buf.free(); + b.* = undefined; + } +}; + +const glue_router_ks = graph.glue.router_ks; + +/// One sequence's committed caches: KDA conv windows and states (two copies), a window's saves, MLA latent rows. +pub const State = struct { + arena: Arena, + gpa: std.mem.Allocator, + capacity: usize, + pos: usize = 0, + pending: usize = 0, // rows of the last forward until its commit + prompt: bool = false, // whether those rows were prompt rows (their KDA states already committed) + pos_dev: u64, // i32 + conv: u64, // bf16 [kda layers][3][3 heads 128] + rec: u64, // fp32 [2][kda layers][heads][128][128] + proj: u64, // bf16 [kda layers][window_rows][kdaStride] + save_k: u64, // fp32 [kda layers][window_rows][heads 128] + save_v: u64, // bf16, as save_k + save_g: u64, // fp32, as save_k + save_b: u64, // fp32 [kda layers][window_rows][heads] + cache: u64, // bf16 [mla layers][capacity][576] + cur: []u1, // each KDA layer's committed copy of rec + kda_layers: usize, + mla_layers: usize, + heads: usize, + stride: usize, // kdaStride + rec_layer: usize, // bytes of one layer's state + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, g: *const graph.Graph, capacity: usize) !State { + const c = &g.w.config; + const nk = g.kda_layers; + const nm = g.mla_layers; + const H: usize = c.kda_heads; + const n = H * kda.dk; + const ps = kdaStride(g.w); + const W = window_rows; + const sizes = [_]usize{ 4, nk * 3 * 3 * n * 2, 2 * nk * n * kda.dv * 4, nk * W * ps * 2, nk * W * n * 4, nk * W * n * 2, nk * W * n * 4, nk * W * H * 4, nm * capacity * cache_width * 2 }; + var p: [sizes.len]u64 = undefined; + var a = try Arena.init(d, &sizes, &p); + errdefer a.buf.free(); + const cur = try gpa.alloc(u1, nk); + @memset(cur, 0); + return .{ .arena = a, .gpa = gpa, .capacity = capacity, .pos_dev = p[0], .conv = p[1], .rec = p[2], .proj = p[3], .save_k = p[4], .save_v = p[5], .save_g = p[6], .save_b = p[7], .cache = p[8], .cur = cur, .kda_layers = nk, .mla_layers = nm, .heads = H, .stride = ps, .rec_layer = n * kda.dv * 4 }; + } + + pub fn deinit(st: *State) void { + st.arena.buf.free(); + st.gpa.free(st.cur); + st.* = undefined; + } + + /// A fresh sequence: zero conv windows and states, position 0 (cache rows past it are never read). + pub fn reset(st: *State, s: cuda.Stream) !void { + const d = st.arena.buf.d; + try d.check(d.api.cuMemsetD8Async(st.conv, 0, st.kda_layers * 3 * st.convWidth() * 2, s.handle), "cuMemsetD8Async"); + try d.check(d.api.cuMemsetD8Async(st.rec, 0, 2 * st.kda_layers * st.rec_layer, s.handle), "cuMemsetD8Async"); + @memset(st.cur, 0); + st.pending = 0; + st.prompt = false; + try st.setPos(s, 0); + } + + pub fn setPos(st: *State, s: cuda.Stream, pos: usize) !void { + st.pos = pos; + const v: i32 = @intCast(pos); + try st.arena.buf.d.check(st.arena.buf.d.api.cuMemsetD32Async(st.pos_dev, @bitCast(v), 1, s.handle), "cuMemsetD32Async"); + } + + /// The conv window's channels: q | k | v. + pub fn convWidth(st: *const State) usize { + return 3 * st.heads * kda.dk; + } + + pub fn recAt(st: *const State, copy: u1, li: usize) u64 { + return st.rec + (@as(usize, copy) * st.kda_layers + li) * st.rec_layer; + } + + pub fn convAt(st: *const State, li: usize) u64 { + return st.conv + li * 3 * st.convWidth() * 2; + } + + pub fn cacheAt(st: *const State, mi: usize) u64 { + return st.cache + mi * st.capacity * cache_width * 2; + } + + /// Layer li's window saves: projection rows, k, v, decay, beta. + pub fn saves(st: *const State, li: usize) struct { p: u64, k: u64, v: u64, g: u64, b: u64 } { + const n = st.heads * kda.dk; + const kv = li * window_rows * n; + return .{ .p = st.proj + li * window_rows * st.stride * 2, .k = st.save_k + kv * 4, .v = st.save_v + kv * 2, .g = st.save_g + kv * 4, .b = st.save_b + li * window_rows * st.heads * 4 }; + } +}; diff --git a/zig/src/families/ling/synth.zig b/zig/src/families/ling/synth.zig index 593f8ba916..69cdd3fef1 100644 --- a/zig/src/families/ling/synth.zig +++ b/zig/src/families/ling/synth.zig @@ -14,7 +14,7 @@ pub const shape = .{ .topk = 4, .n_group = 4, .topk_group = 2, - .q_lora = 256, + .q_lora = 512, // a width glue.rmsnorm takes (512, 1536, 4096) .vocab = 1024, .group = 2, // KDA, MLA, KDA, MLA, ... }; diff --git a/zig/tests/cuda/ling_forward_rig.zig b/zig/tests/cuda/ling_forward_rig.zig new file mode 100644 index 0000000000..37927d3dc8 --- /dev/null +++ b/zig/tests/cuda/ling_forward_rig.zig @@ -0,0 +1,130 @@ +//! What the Ling forward checks share: the graph's kit on one stream, runs that return logits, distances and flips. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const ling = @import("ling_cuda"); + +pub const F = ling.forward; +pub const W = ling.weights; + +pub fn bfOf(u: u16) f64 { + return @as(f32, @bitCast(@as(u32, u) << 16)); +} + +pub const Rig = struct { + gpu: check.Gpu, + a: std.mem.Allocator, + s: cuda.Stream, + m: *const F.Modules, + g: *const F.Graph, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + kda_wide: bool = true, // decode windows on KDA's three-kernel chain (false: the fused one) + + pub fn deinit(r: *Rig) void { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(r.gpu.gpa); + } + + pub fn zeros(r: *Rig, len: usize) !u64 { + const b = try cuda.DeviceBuffer.alloc(r.gpu.d, @max(len, 16)); + try r.bufs.append(r.gpu.gpa, b); + try b.fill8(0, null); + return b.ptr; + } + + pub fn back(r: *Rig, ptr: u64, len: usize) ![]u8 { + const out = try r.a.alloc(u8, len); + try r.s.synchronize(); + const b: cuda.DeviceBuffer = .{ .d = r.gpu.d, .ptr = ptr, .len = len }; + try b.download(0, out); + return out; + } + + pub fn fwd(r: *const Rig, b: *const F.Buffers, picks: ?u64) F.Forward { + return .{ .m = r.m, .g = r.g, .s = r.s, .b = b, .picks = picks, .kda_wide = r.kda_wide }; + } + + /// One forward of `segs` on `b`; its logits rows (bf16 bits) and, with `picks`, every MoE layer's picks. + pub fn run(r: *Rig, b: *const F.Buffers, segs: []const F.Seg, ids: []const i32, picks: ?u64) ![]u16 { + try r.fwd(b, picks).run(segs, ids); + const v: usize = r.g.w.config.vocab; + const raw = try r.back(b.logits, ids.len * v * 2); + const out = try r.a.alloc(u16, ids.len * v); + @memcpy(std.mem.sliceAsBytes(out), raw); + return out; + } + + /// A fresh sequence of `capacity` positions. + pub fn state(r: *Rig, capacity: usize) !*F.State { + const st = try r.a.create(F.State); + st.* = try F.State.init(r.gpu.gpa, r.gpu.d, r.g, capacity); + try st.reset(r.s); + return st; + } +}; + +/// Picks [layers][rows][top_k + 1] (i32) as each row's top-k sorted, [layers][rows][top_k]. +pub fn sortedPicks(a: std.mem.Allocator, raw: []const u8, layers: usize, rows: usize, top: usize) ![]i32 { + const all = std.mem.bytesAsSlice(i32, raw); + const out = try a.alloc(i32, layers * rows * top); + for (0..layers * rows) |i| { + const dst = out[i * top ..][0..top]; + for (dst, 0..) |*d, k| d.* = all[i * (top + 1) + k]; + std.mem.sort(i32, dst, {}, std.sort.asc(i32)); + } + return out; +} + +/// Rows whose sorted picks differ from the reference's in any layer (`want` rows from `w0` of its `wrows`). +pub fn flips(got: []const i32, want: []const i32, layers: usize, rows: usize, w0: usize, wrows: usize, top: usize, out: []bool) usize { + var n: usize = 0; + for (0..layers) |l| for (0..rows) |t| { + const g = got[(l * rows + t) * top ..][0..top]; + const e = want[(l * wrows + w0 + t) * top ..][0..top]; + if (!std.mem.eql(i32, g, e)) { + n += 1; + out[t] = true; + } + }; + return n; +} + +/// bf16 rows against float64 rows: relative L2, the same excluding `skip` rows, the worst row's, argmax agreement. +pub const Dist = struct { rel: f64, rel_kept: f64, worst_row: f64, argmax: usize }; + +pub fn distance(got: []const u16, want: []const f64, rows: usize, width: usize, skip: []const bool) Dist { + var dd: f64 = 0; + var ww: f64 = 0; + var kd: f64 = 0; + var kw: f64 = 0; + var worst: f64 = 0; + var agree: usize = 0; + for (0..rows) |t| { + var rd: f64 = 0; + var rw: f64 = 0; + var ga: usize = 0; + var wa: usize = 0; + for (0..width) |j| { + const g = bfOf(got[t * width + j]); + const e = want[t * width + j]; + rd += (g - e) * (g - e); + rw += e * e; + if (g > bfOf(got[t * width + ga])) ga = j; + if (e > want[t * width + wa]) wa = j; + } + dd += rd; + ww += rw; + if (!skip[t]) { + kd += rd; + kw += rw; + } + worst = @max(worst, @sqrt(rd / rw)); + agree += @intFromBool(ga == wa); + } + return .{ .rel = @sqrt(dd / ww), .rel_kept = if (kw > 0) @sqrt(kd / kw) else 0, .worst_row = worst, .argmax = agree }; +} + +/// Two runs' logits rows equal, bit for bit. +pub fn same(what: []const u8, got: []const u16, want: []const u16) !void { + try check.sameBytes(what, std.mem.sliceAsBytes(got), std.mem.sliceAsBytes(want)); +} diff --git a/zig/tests/cuda/ling_forward_synth.zig b/zig/tests/cuda/ling_forward_synth.zig new file mode 100644 index 0000000000..dcad3c8f0b --- /dev/null +++ b/zig/tests/cuda/ling_forward_synth.zig @@ -0,0 +1,178 @@ +//! `ling-forward-synth DIR REF`: a synthetic checkpoint's logits vs ling_forward.py's; cuts, windows, batches exact. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const ling = @import("ling_cuda"); +const rig = @import("ling_forward_rig.zig"); + +const F = rig.F; +const Rig = rig.Rig; +const P = 80; // prompt rows: KDA's prompt takes the wide chain (64 rows or more) + +/// ling_forward.py's file: ids, float64 and bf16-point logits [T][V], the picks it routed on [layers][T][top]. +const Ref = struct { T: usize, V: usize, layers: usize, top: usize, forced: bool, gap: f64, ids: []i32, l64: []f64, lbf: []f64, picks: []i32 }; + +fn take(comptime T: type, a: std.mem.Allocator, bytes: []const u8, at: *usize, n: usize) ![]T { + const out = try a.alloc(T, n); + @memcpy(std.mem.sliceAsBytes(out), bytes[at.*..][0 .. n * @sizeOf(T)]); + at.* += n * @sizeOf(T); + return out; +} + +fn readRef(a: std.mem.Allocator, io: std.Io, path: []const u8) !Ref { + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 30)); + if (bytes.len < 36 or !std.mem.eql(u8, bytes[0..8], "TFLINGF2")) return error.BadFixture; + var at: usize = 8; + const dims = try take(u32, a, bytes, &at, 5); + const gap = (try take(f64, a, bytes, &at, 1))[0]; + const T: usize = dims[0]; + const V: usize = dims[1]; + const L: usize = dims[2]; + const K: usize = dims[3]; + const r: Ref = .{ .T = T, .V = V, .layers = L, .top = K, .forced = dims[4] != 0, .gap = gap, .ids = try take(i32, a, bytes, &at, T), .l64 = try take(f64, a, bytes, &at, T * V), .lbf = try take(f64, a, bytes, &at, T * V), .picks = try take(i32, a, bytes, &at, L * T * K) }; + if (at != bytes.len) return error.BadFixture; + return r; +} + +/// The serial run: the prompt in one chunk, then one decode row at a time; logits rows and sorted picks. +fn serial(r: *Rig, pb: *const F.Buffers, db: *const F.Buffers, st: *F.State, ids: []const i32, nm: usize) !struct { logits: []u16, picks: []i32 } { + const c = &r.g.w.config; + const V: usize = c.vocab; + const T = ids.len; + const top: usize = c.topk; + const logits = try r.a.alloc(u16, T * V); + const picks = try r.a.alloc(i32, nm * T * top); + const pp = try r.zeros(nm * P * (top + 1) * 4); + @memcpy(logits[0 .. P * V], try r.run(pb, &.{.{ .st = st, .row0 = 0, .rows = P }}, ids[0..P], pp)); + const sp = try rig.sortedPicks(r.a, try r.back(pp, nm * P * (top + 1) * 4), nm, P, top); + for (0..nm) |l| @memcpy(picks[l * T * top ..][0 .. P * top], sp[l * P * top ..][0 .. P * top]); + try r.fwd(pb, null).commit(st, P); + const dp = try r.zeros(nm * (top + 1) * 4); + for (P..T) |t| { + @memcpy(logits[t * V ..][0..V], try r.run(db, &.{.{ .st = st, .row0 = 0, .rows = 1 }}, ids[t .. t + 1], dp)); + const one = try rig.sortedPicks(r.a, try r.back(dp, nm * (top + 1) * 4), nm, 1, top); + for (0..nm) |l| @memcpy(picks[(l * T + t) * top ..][0..top], one[l * top ..][0..top]); + try r.fwd(db, null).commit(st, 1); + } + return .{ .logits = logits, .picks = picks }; +} + +/// A committed sequence's caches (KDA states and conv windows, MLA rows below pos) equal another's, byte for byte. +fn sameState(r: *Rig, what: []const u8, a: *const F.State, b: *const F.State) !void { + var name: [96]u8 = undefined; + try check.expect(a.pos == b.pos, "{s}: positions {d} and {d}", .{ what, a.pos, b.pos }); + for (0..a.kda_layers) |li| { + try check.sameBytes(try std.fmt.bufPrint(&name, "{s}: KDA layer {d} state", .{ what, li }), try r.back(a.recAt(a.cur[li], li), a.rec_layer), try r.back(b.recAt(b.cur[li], li), b.rec_layer)); + try check.sameBytes(try std.fmt.bufPrint(&name, "{s}: KDA layer {d} conv window", .{ what, li }), try r.back(a.convAt(li), 3 * a.convWidth() * 2), try r.back(b.convAt(li), 3 * b.convWidth() * 2)); + } + for (0..a.mla_layers) |mi| try check.sameBytes(try std.fmt.bufPrint(&name, "{s}: MLA layer {d} cache rows", .{ what, mi }), try r.back(a.cacheAt(mi), a.pos * F.cache_width * 2), try r.back(b.cacheAt(mi), b.pos * F.cache_width * 2)); +} + +/// `ling-forward-synth DIR REF [PICKS]`: PICKS gets the serial run's picks, for ling_forward.py --picks. +pub fn run(gpu: check.Gpu, dir: []const u8, ref_path: []const u8, picks_out: ?[]const u8) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + const ref = try readRef(a, gpu.io, ref_path); + var why: ling.config.Why = .{}; + var o = rig.W.open(gpu.gpa, gpu.io, dir, &why) catch |e| { + std.debug.print("{s}: {s}\n", .{ dir, why.text() }); + return e; + }; + defer o.ck.close(); + var m = try F.Modules.load(gpu.d, gpu.ctx); + defer m.deinit(); + const backbone = try a.alloc(u32, o.c.layers); + for (backbone, 0..) |*i, j| i.* = @intCast(j); + var pack = try cuda.modelopt.Packer.load(gpu.d); + defer pack.unload(); + var w = try rig.W.load(gpu.gpa, gpu.d, s, &pack, &o.ck, o.c, backbone, .{}); + defer w.deinit(); + var g = try F.Graph.init(gpu.gpa, gpu.d, &m, s, &w); + defer g.deinit(); + const c = &w.config; + const V: usize = c.vocab; + const T = ref.T; + try check.expect(ref.V == V and T > P + 8 - 1 and ref.top == c.topk, "reference: {d} tokens of vocabulary {d}, top {d}", .{ T, ref.V, ref.top }); + var r: Rig = .{ .gpu = gpu, .a = a, .s = s, .m = &m, .g = &g }; + defer r.deinit(); + var pb = try F.Buffers.init(gpu.d, &m, &w, 2 * P, true, 128); + defer pb.deinit(); + var db = try F.Buffers.init(gpu.d, &m, &w, 2 * F.window_rows, false, 128); + defer db.deinit(); + var nm: usize = 0; + for (w.layers) |l| nm += @intFromBool(l.ffn == .moe); + try check.expect(nm == ref.layers, "MoE layers {d}, the reference's {d}", .{ nm, ref.layers }); + const ids = ref.ids; + // A: the serial run, against the references + const sa = try r.state(128); + const A = try serial(&r, &pb, &db, sa, ids, nm); + const flipped = try a.alloc(bool, T); + @memset(flipped, false); + const nf = rig.flips(A.picks, ref.picks, nm, T, 0, T, c.topk, flipped); + var rows_flipped: usize = 0; + for (flipped) |x| rows_flipped += @intFromBool(x); + const d64 = rig.distance(A.logits, ref.l64, T, V, flipped); + const dbf = rig.distance(A.logits, ref.lbf, T, V, flipped); + if (picks_out) |path| try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = path, .data = std.mem.sliceAsBytes(A.picks) }); + const how = if (ref.forced) "the reference routed on these picks" else "the reference's own float64 picks"; + std.debug.print("RESULT ling-forward-synth {t}: {d} layers, {d} prompt + {d} decode rows; picks vs {s}: {d} of {d} differ ({d} rows)\n", .{ c.quant, c.layers, P, T - P, how, nf, nm * T, rows_flipped }); + if (ref.forced) std.debug.print("RESULT float64's own top-k vs these: the worst pick ranks {e:.3} below its k-th biased score\n", .{ref.gap}); + std.debug.print("RESULT logits vs float64: rel L2 {e:.3} ({e:.3} without flipped rows), worst row {e:.3}, argmax {d}/{d}\n", .{ d64.rel, d64.rel_kept, d64.worst_row, d64.argmax, T }); + std.debug.print("RESULT logits vs bf16 points: rel L2 {e:.3} ({e:.3} without flipped rows), worst row {e:.3}, argmax {d}/{d}\n", .{ dbf.rel, dbf.rel_kept, dbf.worst_row, dbf.argmax, T }); + var bd: f64 = 0; + var bw: f64 = 0; + for (ref.lbf, ref.l64) |x, e| { + bd += (x - e) * (x - e); + bw += e * e; + } + const base = @sqrt(bd / bw); // HF's bf16 rounding points alone, against float64 + std.debug.print("RESULT the bf16-point reference itself vs float64: rel L2 {e:.3}\n", .{base}); + if (ref.forced) try check.expect(nf == 0 and d64.rel <= 1.5 * base and dbf.rel <= base, "on the same picks: logits within 1.5x the bf16 points' rel L2 of float64, and within it of the bf16 points", .{}); + // B: the prompt cut at 48 + const sb = try r.state(128); + try rig.same("prompt cut at 48 == uncut (rows 0-47)", try r.run(&pb, &.{.{ .st = sb, .row0 = 0, .rows = 48 }}, ids[0..48], null), A.logits[0 .. 48 * V]); + try r.fwd(&pb, null).commit(sb, 48); + try rig.same("prompt cut at 48 == uncut (rows 48-79)", try r.run(&pb, &.{.{ .st = sb, .row0 = 0, .rows = P - 48 }}, ids[48..P], null), A.logits[48 * V .. P * V]); + try r.fwd(&pb, null).commit(sb, P - 48); + // C: one 8-row verify window == 8 decode rows + const sc = try r.state(128); + _ = try r.run(&pb, &.{.{ .st = sc, .row0 = 0, .rows = P }}, ids[0..P], null); + try r.fwd(&pb, null).commit(sc, P); + try rig.same("8-row verify window == 8 decode rows", try r.run(&db, &.{.{ .st = sc, .row0 = 0, .rows = 8 }}, ids[P .. P + 8], null), A.logits[P * V .. (P + 8) * V]); + try r.fwd(&db, null).commit(sc, 3); + try rig.same("after keeping 3 of 8: the next 5-row window == decode rows", try r.run(&db, &.{.{ .st = sc, .row0 = 0, .rows = 5 }}, ids[P + 3 .. P + 8], null), A.logits[(P + 3) * V .. (P + 8) * V]); + try r.fwd(&db, null).commit(sc, 5); + try sameState(&r, "8-row window, keep 3, 5-row window == 8 decode rows", sc, sa); + // the fused KDA chain on the same window: bit-equal to the three-kernel one + const sf = try r.state(128); + _ = try r.run(&pb, &.{.{ .st = sf, .row0 = 0, .rows = P }}, ids[0..P], null); + try r.fwd(&pb, null).commit(sf, P); + r.kda_wide = false; + try rig.same("8-row window on the fused KDA chain == on the three-kernel chain", try r.run(&db, &.{.{ .st = sf, .row0 = 0, .rows = 8 }}, ids[P .. P + 8], null), A.logits[P * V .. (P + 8) * V]); + r.kda_wide = true; + sf.deinit(); + // D and E: two sequences batched, in a prompt chunk and in a decode window + const sd = try r.state(128); + const se = try r.state(128); + const both = try a.alloc(i32, P + 48); + @memcpy(both[0..P], ids[0..P]); + @memcpy(both[P..], ids[0..48]); + const pl = try r.run(&pb, &.{ .{ .st = sd, .row0 = 0, .rows = P }, .{ .st = se, .row0 = P, .rows = 48 } }, both, null); + try rig.same("batched prompt: sequence D == alone", pl[0 .. P * V], A.logits[0 .. P * V]); + try rig.same("batched prompt: sequence E (48 rows) == alone", pl[P * V ..], A.logits[0 .. 48 * V]); + try r.fwd(&pb, null).commit(sd, P); + try r.fwd(&pb, null).commit(se, 48); + _ = try r.run(&pb, &.{.{ .st = se, .row0 = 0, .rows = P - 48 }}, ids[48..P], null); + try r.fwd(&pb, null).commit(se, P - 48); + const win = try a.alloc(i32, 10); + @memcpy(win[0..4], ids[P .. P + 4]); + @memcpy(win[4..], ids[P .. P + 6]); + const wl = try r.run(&db, &.{ .{ .st = sd, .row0 = 0, .rows = 4 }, .{ .st = se, .row0 = 4, .rows = 6 } }, win, null); + try rig.same("batched window: sequence D (4 rows) == alone", wl[0 .. 4 * V], A.logits[P * V .. (P + 4) * V]); + try rig.same("batched window: sequence E (6 rows) == alone", wl[4 * V ..], A.logits[P * V .. (P + 6) * V]); + for ([_]*F.State{ sa, sb, sc, sd, se }) |x| x.deinit(); + check.pass("ling-forward-synth {t}: logits rel L2 {e:.3} of float64 ({s}); cut == uncut, windows == decode rows, a partial commit replays, batched == alone", .{ c.quant, d64.rel, if (ref.forced) "same picks" else "free routing" }); +} diff --git a/zig/tests/cuda/ling_layer_real.zig b/zig/tests/cuda/ling_layer_real.zig new file mode 100644 index 0000000000..1c83490a2f --- /dev/null +++ b/zig/tests/cuda/ling_layer_real.zig @@ -0,0 +1,149 @@ +//! `ling-layer-real DIR FIXTURE`: one real Ling-3.1 layer through the graph vs HF's captured output; exactness on it. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const ling = @import("ling_cuda"); +const rig = @import("ling_forward_rig.zig"); + +const F = rig.F; +const Rig = rig.Rig; + +fn take(comptime T: type, a: std.mem.Allocator, bytes: []const u8, at: *usize, n: usize) ![]T { + const out = try a.alloc(T, n); + @memcpy(std.mem.sliceAsBytes(out), bytes[at.*..][0 .. n * @sizeOf(T)]); + at.* += n * @sizeOf(T); + return out; +} + +/// Rows lo .. hi of x through the stack on `st` (prompt or decode buffers), committed; the output rows. +fn pass(r: *Rig, b: *const F.Buffers, st: *F.State, x: []const u16, D: usize, lo: usize, hi: usize, picks: ?u64) ![]u16 { + const d = r.gpu.d; + const n = hi - lo; + try d.check(d.api.cuMemcpyHtoDAsync_v2(b.x, x[lo * D ..].ptr, n * D * 2, r.s.handle), "cuMemcpyHtoDAsync"); + const f = r.fwd(b, picks); + try f.stack(&.{.{ .st = st, .row0 = 0, .rows = n }}); + const out = try r.a.alloc(u16, n * D); + @memcpy(std.mem.sliceAsBytes(out), try r.back(b.x, n * D * 2)); + try f.commit(st, n); + return out; +} + +/// Relative L2 of bf16 rows against fp32 rows, all and without `skip` rows, and the worst row's. +fn dist(got: []const u16, want: []const f32, rows: usize, D: usize, skip: []const bool) [3]f64 { + var dd: f64 = 0; + var ww: f64 = 0; + var kd: f64 = 0; + var kw: f64 = 0; + var worst: f64 = 0; + for (0..rows) |t| { + var rd: f64 = 0; + var rw: f64 = 0; + for (0..D) |j| { + const e: f64 = want[t * D + j]; + const g = rig.bfOf(got[t * D + j]); + rd += (g - e) * (g - e); + rw += e * e; + } + dd += rd; + ww += rw; + if (!skip[t]) { + kd += rd; + kw += rw; + worst = @max(worst, @sqrt(rd / rw)); + } + } + return .{ @sqrt(dd / ww), @sqrt(kd / kw), worst }; +} + +/// `ling-layer-real DIR FIXTURE` (oracle/ling_layer.py's file for layer L of the checkpoint at DIR). +pub fn run(gpu: check.Gpu, dir: []const u8, path: []const u8) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + const bytes = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, a, .limited(1 << 30)); + if (bytes.len < 24 or !std.mem.eql(u8, bytes[0..8], "TFLINGL1")) return error.BadFixture; + var at: usize = 8; + const dims = try take(u32, a, bytes, &at, 4); + const L = dims[0]; + const T: usize = dims[1]; + const D: usize = dims[2]; + const K: usize = dims[3]; + const x = try take(u16, a, bytes, &at, T * D); + const want = try take(f32, a, bytes, &at, T * D); + const topk = try take(i32, a, bytes, &at, T * K); + if (at != bytes.len) return error.BadFixture; + var why: ling.config.Why = .{}; + var o = rig.W.open(gpu.gpa, gpu.io, dir, &why) catch |e| { + std.debug.print("{s}: {s}\n", .{ dir, why.text() }); + return e; + }; + defer o.ck.close(); + var m = try F.Modules.load(gpu.d, gpu.ctx); + defer m.deinit(); + var pack = try cuda.modelopt.Packer.load(gpu.d); + defer pack.unload(); + const t0 = check.now(gpu.io); + var w = try rig.W.load(gpu.gpa, gpu.d, s, &pack, &o.ck, o.c, &.{L}, .{ .globals = false }); + defer w.deinit(); + const load_s = @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e9; + var g = try F.Graph.init(gpu.gpa, gpu.d, &m, s, &w); + defer g.deinit(); + const c = &w.config; + try check.expect(D == c.hidden and K == (if (c.isMoe(L)) @as(usize, c.topk) else 0), "fixture: layer {d}, {d} wide, top {d}", .{ L, D, K }); + var r: Rig = .{ .gpu = gpu, .a = a, .s = s, .m = &m, .g = &g }; + defer r.deinit(); + const cap = T + 32; + var pb = try F.Buffers.init(gpu.d, &m, &w, T, true, cap); + defer pb.deinit(); + var db = try F.Buffers.init(gpu.d, &m, &w, F.window_rows, false, cap); + defer db.deinit(); + // the whole prompt against HF, picks against HF's + const top: usize = c.topk; + const pk = try r.zeros(T * (top + 1) * 4); + const st = try r.state(cap); + defer st.deinit(); + const got = try pass(&r, &pb, st, x, D, 0, T, if (K > 0) pk else null); + const skip = try a.alloc(bool, T); + @memset(skip, false); + var nf: usize = 0; + if (K > 0) { + const ours = try rig.sortedPicks(a, try r.back(pk, T * (top + 1) * 4), 1, T, top); + const theirs = try a.dupe(i32, topk); + for (0..T) |t| std.mem.sort(i32, theirs[t * K ..][0..K], {}, std.sort.asc(i32)); + nf = rig.flips(ours, theirs, 1, T, 0, T, top, skip); + } + const dv = dist(got, want, T, D, skip); + std.debug.print("RESULT ling-layer-real layer {d} ({t} {s}, {t}): {d} tokens, loaded in {d:.1} s\n", .{ L, c.kind(L), if (K > 0) "MoE" else "dense", c.quant, T, load_s }); + std.debug.print("RESULT vs HF: rel L2 {e:.3}, {e:.3} without the {d} rows whose top-{d} differs, worst kept row {e:.3}\n", .{ dv[0], dv[1], nf, top, dv[2] }); + // cut prompt, decode rows, a verify window + const sc = try r.state(cap); + defer sc.deinit(); + const cut0 = try pass(&r, &pb, sc, x, D, 0, 300, null); + const cut1 = try pass(&r, &pb, sc, x, D, 300, T, null); + try check.sameBytes("prompt cut at 300 == uncut", std.mem.sliceAsBytes(cut0), std.mem.sliceAsBytes(got[0 .. 300 * D])); + try check.sameBytes("prompt cut at 300 == uncut (the rest)", std.mem.sliceAsBytes(cut1), std.mem.sliceAsBytes(got[300 * D ..])); + const P = 400; + const N = F.window_rows; + const sd = try r.state(cap); + defer sd.deinit(); + const sw = try r.state(cap); + defer sw.deinit(); + _ = try pass(&r, &pb, sd, x, D, 0, P, null); + _ = try pass(&r, &pb, sw, x, D, 0, P, null); + const dec = try a.alloc(u16, N * D); + for (0..N) |i| @memcpy(dec[i * D ..][0..D], try pass(&r, &db, sd, x, D, P + i, P + i + 1, null)); + try check.sameBytes("16-row verify window == 16 decode rows", std.mem.sliceAsBytes(try pass(&r, &db, sw, x, D, P, P + N, null)), std.mem.sliceAsBytes(dec)); + const dd = dist(dec, want[P * D ..], N, D, skip[P..]); + const dp = dist(dec, @as([]const f32, try widen(a, got[P * D .. (P + N) * D])), N, D, skip[P..]); + std.debug.print("RESULT decode rows {d}-{d}: rel L2 {e:.3} vs HF, {e:.3} vs the prompt's rows\n", .{ P, P + N - 1, dd[0], dp[0] }); + try check.expect(dv[1] <= 1e-2 and dd[1] <= 1e-2, "layer {d} within rel L2 1e-2 of HF (rows whose picks agree)", .{L}); + check.pass("ling-layer-real layer {d}: rel L2 {e:.3} of HF ({d} flipped rows apart); cut == uncut, window == decode rows; peak {d} MiB on the device", .{ L, dv[1], nf, cuda.usage(false).peak >> 20 }); +} + +fn widen(a: std.mem.Allocator, x: []const u16) ![]f32 { + const out = try a.alloc(f32, x.len); + for (out, x) |*o, v| o.* = @floatCast(rig.bfOf(v)); + return out; +} diff --git a/zig/tests/cuda/ling_synth.zig b/zig/tests/cuda/ling_synth.zig index eef50607df..7dde1b698d 100644 --- a/zig/tests/cuda/ling_synth.zig +++ b/zig/tests/cuda/ling_synth.zig @@ -77,7 +77,7 @@ const Out = struct { .nvfp4 => o.add(name, .u8, &.{ n, k / 2 }, try o.codes(n * k / 2, .any)), }; const tail = name[module.len + 1 ..]; - if (std.mem.eql(u8, tail, "weight_scale_inv")) return o.add(name, .f32, &.{ (n + 127) / 128, k / 128 }, try o.f32s((n + 127) / 128 * (k / 128), 0.001, 0.011)); + if (std.mem.eql(u8, tail, "weight_scale_inv")) return o.add(name, .f32, &.{ (n + 127) / 128, k / 128 }, try o.f32s((n + 127) / 128 * (k / 128), s / 4, s / 2)); if (std.mem.eql(u8, tail, "weight_scale")) return o.add(name, .f8_e4m3, &.{ n, k / 16 }, try o.codes(n * k / 16, .scale)); if (std.mem.eql(u8, tail, "weight_scale_2")) return o.add(name, .f32, &.{}, try o.f32s(1, s / 8, s / 4)); if (std.mem.eql(u8, tail, "input_scale")) return o.add(name, .f32, &.{}, try o.f32s(1, 1, 1)); @@ -98,15 +98,12 @@ const Out = struct { }; } - /// Random bytes: any; e4m3 codes without NaN; positive e4m3 scales 2^-3 to 2^1. + /// Random bytes: any; e4m3 codes of magnitude 2^-3 to 3.75 (a forward stays in range); e4m3 scales 2^-3 to 2^1. fn codes(o: *Out, count: usize, what: enum { any, e4m3, scale }) ![]u8 { const v = try o.a.alloc(u8, count); for (v) |*b| b.* = switch (what) { .any => o.rng.int(u8), - .e4m3 => blk: { - const x = o.rng.int(u8); - break :blk if (x & 0x7f == 0x7f) x ^ 1 else x; - }, + .e4m3 => (o.rng.int(u8) & 0x87) | (@as(u8, 4) + o.rng.uintLessThan(u8, 5)) << 3, .scale => 0x20 + o.rng.uintLessThan(u8, 0x28), }; return v; diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index ff056558c8..9ecf832425 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -112,6 +112,8 @@ const usage = \\ glm-synth a synthetic GLM-5.3-Flash checkpoint (2 layers, 16 experts) and its dc.json for the lane checks \\ ling-synth [--format bf16|fp8|nvfp4] [--layers N] a tiny Ling-3.1 checkpoint with the real names and config schema \\ ling-weights every layer of a Ling checkpoint loaded: projections vs float64, small tensors byte for byte + \\ ling-forward-synth [picks] a ling-synth checkpoint's forward vs oracle/ling_forward.py; cuts, windows, batches + \\ ling-layer-real one real Ling-3.1 layer vs HF's captured output (oracle/ling_layer.py); exactness \\ ling-weights-real real Ling-3.1 layers 0, 4, 7 (4 experts) vs float64 and torch (oracle/ling_weights.py) \\ ; @@ -220,6 +222,8 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "glm-synth")) return glm_synth.run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "ling-synth")) return ling_synth.run(gpu, rest); if (std.mem.eql(u8, cmd, "ling-weights")) return ling_weights.run(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "ling-forward-synth")) return @import("ling_forward_synth.zig").run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) rest[2] else null); + if (std.mem.eql(u8, cmd, "ling-layer-real")) return @import("ling_layer_real.zig").run(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "ling-weights-real")) return ling_real.run(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "glm-lanes-concurrent")) return glm_lanes_tests.concurrentCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); std.debug.print("{s}", .{usage}); diff --git a/zig/tests/cuda/oracle/ling_forward.py b/zig/tests/cuda/oracle/ling_forward.py new file mode 100644 index 0000000000..f63ca768ce --- /dev/null +++ b/zig/tests/cuda/oracle/ling_forward.py @@ -0,0 +1,186 @@ +"""For `tf-cuda-test ling-forward-synth`: a Ling-3.1 checkpoint's logits in float64, from its own file and config.""" +import argparse +import json +import math +import os +import struct + +import torch +from safetensors import safe_open + +E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float64) + + +class Ckpt: + """Every tensor of model.safetensors; linears dequantized to float64 by the side tensors beside them.""" + + def __init__(self, path): + self.f = safe_open(path, "pt") + self.keys = set(self.f.keys()) + self.cache = {} + + def t(self, name): + return self.f.get_tensor(name) + + def lin(self, module): + if module not in self.cache: + self.cache[module] = self.dequant(module) + return self.cache[module] + + def dequant(self, module): + w = self.t(module + ".weight") + if module + ".weight_scale_inv" in self.keys: + s = self.t(module + ".weight_scale_inv").double() + n, k = w.shape + return w.double() * s.repeat_interleave(128, 0)[:n].repeat_interleave(128, 1)[:, :k] + if module + ".weight_scale_2" in self.keys: + n = w.shape[0] + codes = torch.stack((w & 0xF, w >> 4), -1).reshape(n, -1).long() + v = E2M1[codes & 7] * torch.where(codes >= 8, -1.0, 1.0).double() + scale = self.t(module + ".weight_scale").double().repeat_interleave(16, 1) + return v * scale * self.t(module + ".weight_scale_2").double() + return w.double() + + +def bf16(x): + return x.to(torch.bfloat16).double() + + +class Model: + """The decoder in float64 (`r` = identity) or at the HF bf16 model's rounding points (`r` = bf16).""" + + def __init__(self, ck, c, r, forced=None): + self.ck, self.c, self.r, self.forced = ck, c, r, forced + self.eps = c["rms_norm_eps"] + self.picks = [] + self.gap = 0.0 + + def norm(self, x, w): + y = x / torch.sqrt((x * x).mean(-1, keepdim=True) + self.eps) + return self.r(self.r(y) * w.double()) + + def lin(self, x, module): + return self.r(x @ self.ck.lin(module).T) + + def kda(self, x, p): + c, r, ck = self.c, self.r, self.ck + T, H, dk = x.shape[0], c["num_attention_heads"], c["head_dim"] + qkv = torch.cat([self.lin(x, p + n) for n in ("q_proj", "k_proj", "v_proj")], -1) + cw = torch.cat([ck.t(p + n + "_conv1d.weight").double()[:, 0, :] for n in ("q", "k", "v")], 0) + xp = torch.cat([torch.zeros(3, qkv.shape[1], dtype=torch.float64), qkv], 0) + acc = sum(cw[:, j] * xp[j:j + T] for j in range(4)) + qkv = r(acc * torch.sigmoid(acc)).view(T, 3, H, dk) + q = qkv[:, 0] / torch.sqrt((qkv[:, 0] ** 2).sum(-1, keepdim=True) + 1e-6) / math.sqrt(dk) + k = qkv[:, 1] / torch.sqrt((qkv[:, 1] ** 2).sum(-1, keepdim=True) + 1e-6) + v = qkv[:, 2] + rate = torch.exp(ck.t(p + "A_log").double()) + f = self.lin(x, p + "f_proj").view(T, H, dk) + ck.t(p + "dt_bias").double().view(H, dk) + g = torch.exp(c["kda_lower_bound"] * torch.sigmoid(rate[None, :, None] * f)) + beta = r(torch.sigmoid(self.lin(x, p + "b_proj"))) + gate = self.lin(x, p + "g_proj").view(T, H, dk) + s = torch.zeros(H, dk, dk, dtype=torch.float64) + out = torch.zeros(T, H, dk, dtype=torch.float64) + for t in range(T): + s = s * g[t][:, None, :] + u = beta[t][:, None] * (v[t] - torch.einsum("hji,hi->hj", s, k[t])) + s = s + u[:, :, None] * k[t][:, None, :] + out[t] = r(torch.einsum("hji,hi->hj", s, q[t])) + w = ck.t(p + "o_norm.weight").double() + y = out / torch.sqrt((out * out).mean(-1, keepdim=True) + self.eps) * w * torch.sigmoid(gate) + return self.lin(r(y).reshape(T, H * dk), p + "o_proj") + + def mla(self, x, p): + c, r, ck = self.c, self.r, self.ck + T, H = x.shape[0], c["num_attention_heads"] + nope, rope, dv, lw = c["qk_nope_head_dim"], c["qk_rope_head_dim"], c["v_head_dim"], c["kv_lora_rank"] + q = self.lin(self.norm(self.lin(x, p + "q_a_proj"), ck.t(p + "q_a_layernorm.weight")), p + "q_b_proj").view(T, H, -1) + kv = self.lin(x, p + "kv_a_proj_with_mqa") + lat = self.norm(kv[:, :lw], ck.t(p + "kv_a_layernorm.weight")) + kvb = self.lin(lat, p + "kv_b_proj").view(T, H, nope + dv) + inv = 1.0 / c["rope_theta"] ** (torch.arange(0, rope, 2, dtype=torch.float64) / rope) + ang = torch.arange(T, dtype=torch.float64)[:, None] * inv + cos, sin = torch.cos(ang), torch.sin(ang) + + def rot(z): # GPT-J pairs (rope_interleave), out in halves as HF's interleave helper writes them + a, b = z[..., 0::2], z[..., 1::2] + cs, sn = (cos[:, None], sin[:, None]) if z.dim() == 3 else (cos, sin) + return r(torch.cat((a * cs - b * sn, b * cs + a * sn), -1)) + + qr, kr = rot(q[..., nope:]), rot(kv[:, lw:]) + sc = (torch.einsum("thd,shd->hts", q[..., :nope], kvb[..., :nope]) + torch.einsum("thd,sd->hts", qr, kr)) * (nope + rope) ** -0.5 + sc = sc.masked_fill(torch.triu(torch.ones(T, T, dtype=torch.bool), 1), float("-inf")) + o = r(torch.einsum("hts,shd->thd", torch.softmax(sc, -1), kvb[..., nope:])) + o = r(o * r(torch.sigmoid(self.lin(x, p + "g_proj")))[..., None]) + return self.lin(o.reshape(T, H * dv), p + "dense") + + def mlp(self, x, p, limit): + g, u = self.lin(x, p + "gate_proj"), self.lin(x, p + "up_proj") + if limit: + g, u = g.clamp(max=limit), u.clamp(-limit, limit) + return self.lin(self.r(self.r(g * torch.sigmoid(g)) * u), p + "down_proj") + + def moe(self, x, p, i): + c, r, ck = self.c, self.r, self.ck + ne, top, ng, kg = c["num_experts"], c["num_experts_per_tok"], c["n_group"], c["topk_group"] + s = torch.sigmoid(x @ ck.t(p + "gate.weight").double().T) + choice = s + ck.t(p + "gate.expert_bias").double() + gs = choice.view(-1, ng, ne // ng).topk(2, -1)[0].sum(-1) + keep = torch.zeros_like(gs).scatter_(1, gs.topk(kg, -1)[1], 1).repeat_interleave(ne // ng, 1) + masked = choice.masked_fill(keep == 0, float("-inf")) + idx = masked.topk(top, -1)[1] + self.picks.append(idx.sort(-1)[0].int()) + if self.forced is not None: # the graph's picks; gap: how far the worst one ranks below the top-k + idx = self.forced[len(self.picks) - 1].long() + kth = masked.topk(top, -1)[0][:, -1] + self.gap = max(self.gap, (kth - choice.gather(1, idx).min(-1)[0]).max().item()) + w = s.gather(1, idx) + w = w / (w.sum(-1, keepdim=True) + 1e-20) * c["routed_scaling_factor"] + le, ls = c["expert_swiglu_limit_list"][i], c["share_expert_swiglu_limit_list"][i] + y = torch.zeros_like(x) + for t in range(x.shape[0]): + for j in range(top): + y[t] += self.mlp(x[t:t + 1], f"{p}experts.{int(idx[t, j])}.", le)[0] * w[t, j] + return r(r(y) + self.mlp(x, p + "shared_experts.", ls)) + + def logits(self, ids): + c, ck = self.c, self.ck + x = ck.t("model.word_embeddings.weight").double()[ids] + for i in range(c["num_hidden_layers"]): + p = f"model.layers.{i}." + h = self.norm(x, ck.t(p + "input_layernorm.weight")) + g = c["layer_group_size"] + mla = (i + 1) % g == 0 or i >= c["num_hidden_layers"] // g * g + x = self.r(x + (self.mla(h, p + "attention.") if mla else self.kda(h, p + "attention."))) + h = self.norm(x, ck.t(p + "post_attention_layernorm.weight")) + x = self.r(x + (self.moe(h, p + "mlp.", i) if i >= c["first_k_dense_replace"] else self.mlp(h, p + "mlp.", 0))) + return self.norm(x, ck.t("model.norm.weight")) @ ck.lin("lm_head").T + + +def main(): + a = argparse.ArgumentParser(description=__doc__) + a.add_argument("ckpt", help="a ling-synth directory (config.json, model.safetensors)") + a.add_argument("out", help="the reference file: ids, float64 and bf16-point logits, the picks it routed on") + a.add_argument("--tokens", type=int, default=88) + a.add_argument("--picks", help="route on these picks (i32 [MoE layers][tokens][top_k], the graph's) instead") + o = a.parse_args() + c = json.load(open(os.path.join(o.ckpt, "config.json"))) + ck = Ckpt(os.path.join(o.ckpt, "model.safetensors")) + T, V = o.tokens, c["vocab_size"] + ids = torch.tensor([(37 * t + 11) % (V - 2) for t in range(T)]) + forced = None + if o.picks: + nm = sum(1 for i in range(c["num_hidden_layers"]) if i >= c["first_k_dense_replace"]) + forced = torch.frombuffer(bytearray(open(o.picks, "rb").read()), dtype=torch.int32).view(nm, T, -1) + with torch.no_grad(): + exact = Model(ck, c, lambda z: z, forced) + l64 = exact.logits(ids) + lbf = Model(ck, c, bf16, forced).logits(ids) + picks = forced if forced is not None else torch.stack(exact.picks) + with open(o.out, "wb") as f: + f.write(b"TFLINGF2" + struct.pack("<4I", T, V, picks.shape[0], picks.shape[2]) + struct.pack(" None: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("chain", help="the captured chain: h.pt (layer L's input), h.pt (its output), r.pt") + p.add_argument("layer", type=int) + p.add_argument("out", help="the fixture: x bf16 [T][D], HF's output fp32 [T][D], HF's top-k i32 [T][K] (K 0: dense)") + a = p.parse_args() + x = torch.load(f"{a.chain}/h{a.layer}.pt")[0].to(torch.bfloat16) + y = torch.load(f"{a.chain}/h{a.layer + 1}.pt")[0].float() + r = torch.load(f"{a.chain}/r{a.layer}.pt") + topk = r["topk"][0].int() if "topk" in r else torch.zeros(x.shape[0], 0, dtype=torch.int32) + T, D = x.shape + with open(a.out, "wb") as f: + f.write(b"TFLINGL1" + struct.pack("<4I", a.layer, T, D, topk.shape[1])) + f.write(x.view(torch.int16).numpy().tobytes() + y.numpy().tobytes() + topk.numpy().tobytes()) + print(a.layer, r["type"], "tokens", T, "top", topk.shape[1], "clamps", r["LE"], r["LS"]) + + +if __name__ == "__main__": + main() From bfb44ddc505e1e3022065e3d50c262f31feaf271 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 15 Oct 2026 12:10:06 +0200 Subject: [PATCH 08/34] ling: SwiGLU clamp in Ling's order, min(silu(g), L) * clip(u, -L, L) Ling's reference (inclusionAI's SwigluStepAndMul, SGLang's BailingMoeV3MLP) clamps silu's output; GLM's glue and the grouped expert kernels clamp the gate before silu. Ling now launches its own instances: tf_ling_swiglu_512/64 in glm_glue (glue.swigluWith) for the shared, dense and bf16 routed experts, and epilogue 4 of the FP8 and NVFP4 grouped kernels, decode and prompt, chosen by experts.Layer.after (set by the Ling loader, refused by affine4). Every pre-existing kernel keeps its SASS and resource usage (cuobjdump over all 40 fatbins vs a33dde0c; 37 fatbins byte-identical, the other 3 only add the new instances). ling-clamp-old runs the kept FP8 (12) and NVFP4 (13) instance calls against a33dde0c's fatbins byte for byte; glm-route-old and glm-mla-same still equal 6380c8da. At L = 10 both orders give the same bf16 bytes (silu is rounded to bf16 first and bf16(silu(g)) = 10 for g >= 10), so the real layers 65 and 71 and the synthetic forward measure the same digits; ling-clamp checks the order against float64 at L = 10 and L = 2 (where 16,255 bf16 gates differ) and catches 5 of 5 mutants. The float64 oracle clamps after silu too. --- docs/recipes/ling-3.1.md | 26 ++- zig/kernels/cuda/experts.cuh | 14 ++ zig/kernels/cuda/fp8_experts.cu | 4 +- zig/kernels/cuda/glm_glue_moe.cuh | 67 +++--- zig/kernels/cuda/nvfp4_experts.cu | 6 +- zig/src/cuda/experts.zig | 44 ++-- zig/src/families/glm/cuda_glue.zig | 7 +- zig/src/families/ling/block_ffn.zig | 4 +- zig/src/families/ling/graph.zig | 2 + zig/src/families/ling/weights_proj.zig | 4 +- zig/tests/cuda/ling_clamp.zig | 269 +++++++++++++++++++++++++ zig/tests/cuda/main.zig | 4 + zig/tests/cuda/oracle/ling_forward.py | 7 +- 13 files changed, 401 insertions(+), 57 deletions(-) create mode 100644 zig/tests/cuda/ling_clamp.zig diff --git a/docs/recipes/ling-3.1.md b/docs/recipes/ling-3.1.md index 58d8f9c44f..3b4b6da25f 100644 --- a/docs/recipes/ling-3.1.md +++ b/docs/recipes/ling-3.1.md @@ -19,6 +19,7 @@ block FP8 (e4m3 in 128 x 128 blocks with `weight_scale_inv`). | Weight loader at TP=1: bf16, FP8 block, ModelOpt NVFP4 | `zig/src/families/ling/weights.zig`, `weights_proj.zig` | `ling-weights`, `ling-weights-real` | | Synthetic checkpoints in all three formats | `tf-cuda-test ling-synth` | its own name check | | The TP=1 layer graph: embedding, KDA/MLA mixers, dense MLP and MoE, final norm, head; windows and commits | `zig/src/families/ling/forward.zig`, `graph.zig`, `state.zig`, `block_*.zig`, `kernels/cuda/ling_glue.cu` | `ling-forward-synth`, `ling-layer-real` | +| SwiGLU clamped in Ling's order, `min(silu(g), L) * clip(u, -L, L)`, in the glue and the FP8/NVFP4 grouped experts (every older kernel's SASS unchanged) | `glm_glue_moe.cuh` (`tf_ling_swiglu_*`), `experts.cuh` (epilogue 4), `experts.Layer.after` | `ling-clamp`, `ling-clamp-old`, `glm-route-old` | ## The layer graph @@ -67,6 +68,7 @@ Real layers of the FP8 release (`ling-layer-real`, the HF bf16 chain's captured | 0 | KDA, dense | 3.27e-3 | - | 3.27e-3 | 3.12e-3 | 0.9 GB | | 4 | KDA, MoE | 5.88e-3 | 21 | 3.13e-3 | 4.28e-3 | 8.5 GB | | 7 | MLA, MoE | 4.89e-3 | 17 | 3.05e-3 | 3.57e-3 | 8.3 GB | +| 65 | KDA, MoE, limits 10 | 4.85e-3 | 14 | 3.35e-3 | 3.33e-3 | 8.5 GB | | 71 | MLA, MoE, limits 10 | 4.13e-3 | 7 | 3.75e-3 | 3.89e-3 | 8.3 GB | exllamav3's port measured 0.4-1.1e-3 without flips (4 flips each on layers 4 and 7) on the same chain: it keeps fp16 @@ -94,6 +96,26 @@ both formats. Each projection loads as a `qlinear.Weight` (`bf16`, `fp8g` or `nv `Fp8Experts`/`Nvfp4Experts`). There is no grouped bf16 expert kernel, so bf16 routed experts load as one lane weight per projection. The router must be bf16 (GLM's router kernel reads bf16); a quantized router is refused. +### The SwiGLU clamp + +Layers 64-71 clamp the shared expert and layers 65-71 the routed experts at L = 10 as inclusionAI's vLLM +(`SwigluStepAndMul`) and SGLang (`BailingMoeV3MLP`) do: `min(silu(g), L) * clip(u, -L, L)`, the limit on SiLU's output. +GLM's glue and the grouped expert kernels clamp the gate before SiLU, so Ling launches its own instances: the glue's +`tf_ling_swiglu_512`/`_64` (`glue.swigluWith`) for the shared expert, dense and bf16 routed experts, and epilogue 4 of +the FP8 and NVFP4 grouped kernels (decode and prompt), chosen by `experts.Layer.after`, which the Ling loader sets. +Every kernel that existed before has the same SASS and resource usage as at a33dde0c (`cuobjdump`, all 40 fatbins); +`ling-clamp-old` runs the kept FP8 (12 calls) and NVFP4 (13 calls) instances against a33dde0c's fatbins byte for byte. + +At L = 10 the two orders give the same bf16 bytes: the kernels (and HF's bf16 `act_fn`) round SiLU to bf16 first, and +bf16(silu(g)) is 10 for every bf16 g >= 10. They differ at small limits (L = 2: 16,255 of the 65,536 bf16 gates). +`ling-clamp DIR` checks Ling's order against float64 at L = 10 and L = 2 on gates and ups straddling L: the glue at +widths 1024 and 704 (worst 0.93 of 2^-7 relative; 0 bytes differ from GLM's order at 10, 6,794 at 2) and the clamped +layer's grouped gate-up, decode and prompt rows (FP8 worst 0.090, NVFP4 0.102 of the bound). a33dde0c's binary and +this one print the same digits on the real layers (71: rel L2 4.125e-3, 3.747e-3 without 7 flipped rows, decode rows +3.893e-3; 65: 4.851e-3, 3.347e-3 without 14, 3.325e-3) and on the synthetic forward in all three formats, whose +float64 references do not change either (the synthetic clamped layer never reaches g > 10). HF's clamp itself moves +layer 71's output by 3.7e-3 (rel L2 against the unclamped layer) and layer 65's by 0 on these 512 tokens. + ## What is missing - Lanes hooks (drafted, concurrent and cut-prompt windows) for Ling: the graph has the pieces (segments, windows, @@ -101,8 +123,6 @@ per projection. The router must be bf16 (GLM's router kernel reads bf16); a quan - MTP drafting: the MTP layer's tensors load (`weights.Mtp`); the graph leaves it unwired. It would be `enorm(embed(next)) | hnorm(hidden)` through `eh_proj`, one MLA + MoE layer on its own cache, the final norm and the shared head. -- The SwiGLU limit clamps the gate before SiLU (`silu(min(g, 10))`, the grouped expert kernels' and GLM's rule); - the Ling reference clamps after (`min(silu(g), 10)`). They differ by at most 4.5e-4 on a clamped element. - Speed: nothing is tuned. Projections land in their columns through a byte copy, bf16 routed experts run densely, and the RMS norms take hidden and `q_lora_rank` widths 512, 1536 and 4096 only (`glue.rmsnorm`). - EXL3, through ashhart's EXL3 module (tracker #832). @@ -120,6 +140,8 @@ python zig/tests/cuda/oracle/ling_forward.py /tmp/ling-fp8 /tmp/ref zig-out/bin/tf-cuda-test ling-forward-synth /tmp/ling-fp8 /tmp/ref /tmp/picks # exactness; writes the picks python zig/tests/cuda/oracle/ling_forward.py /tmp/ling-fp8 /tmp/fref --picks /tmp/picks zig-out/bin/tf-cuda-test ling-forward-synth /tmp/ling-fp8 /tmp/fref # same picks: logits bounded +zig-out/bin/tf-cuda-test ling-clamp /tmp/ling-fp8 # Ling's SwiGLU order vs float64, L 10 and 2 +zig-out/bin/tf-cuda-test ling-clamp-old /tmp/ling-fp8 OLD/fatbin # kept expert instances vs an older build's bytes ``` `ling-synth` writes `config.json` (the real file's schema) and `model.safetensors` with the real key names: hidden 512, diff --git a/zig/kernels/cuda/experts.cuh b/zig/kernels/cuda/experts.cuh index 3396e42a40..cb5dfdca37 100644 --- a/zig/kernels/cuda/experts.cuh +++ b/zig/kernels/cuda/experts.cuh @@ -62,6 +62,17 @@ __device__ __forceinline__ float swiglu(float g, float u, float limit) { return bf(gv / (1.f + expf(-gv))) * uv; } +// Ling's order (EPI 4): the limit on silu's output, min(silu(g), L) * clip(u, -L, L). +__device__ __forceinline__ float swiglu_after(float g, float u, float limit) { + const float gv = bf(g); + float uv = bf(u), s = bf(gv / (1.f + expf(-gv))); + if (limit > 0.f) { + s = fminf(s, limit); + uv = fminf(fmaxf(uv, -limit), limit); + } + return s * uv; +} + __device__ __forceinline__ float relu2(float a) { const float u = fmaxf(bf(a), 0.f); return u * u; @@ -88,6 +99,9 @@ __device__ __forceinline__ void epilogue(const float (&acc)[M][RT][NTW][4], int } else if constexpr (EPI == 2) { o0 = swiglu(a0, acc[M - 1][r][j][2 * h], limit); o1 = swiglu(a1, acc[M - 1][r][j][2 * h + 1], limit); + } else if constexpr (EPI == 4) { + o0 = swiglu_after(a0, acc[M - 1][r][j][2 * h], limit); + o1 = swiglu_after(a1, acc[M - 1][r][j][2 * h + 1], limit); } else { o0 = a0; o1 = a1; diff --git a/zig/kernels/cuda/fp8_experts.cu b/zig/kernels/cuda/fp8_experts.cu index e90d4e7180..1e4302ab4e 100644 --- a/zig/kernels/cuda/fp8_experts.cu +++ b/zig/kernels/cuda/fp8_experts.cu @@ -267,11 +267,13 @@ __global__ void __launch_bounds__(PWARPS * 32, PMIN) } } -// The instances fp8_experts_cuda and fp8_experts_prompt_cuda launch: SwiGLU gate-up, fp32 and bf16 down. +// The instances: SwiGLU gate-up (EPI 4 in Ling's clamp order), fp32 and bf16 down; decode and prompt forms. template __global__ void fp8_expert_kernel<2, 2, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); +template __global__ void fp8_expert_kernel<2, 4, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); template __global__ void fp8_expert_kernel<1, 0, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); template __global__ void fp8_expert_kernel<1, 3, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); template __global__ void fp8_expert_prompt_kernel<2, 2, 2, 2>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); +template __global__ void fp8_expert_prompt_kernel<2, 4, 2, 2>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); template __global__ void fp8_expert_prompt_kernel<1, 0, 4, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); template __global__ void fp8_expert_prompt_kernel<1, 3, 4, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); } diff --git a/zig/kernels/cuda/glm_glue_moe.cuh b/zig/kernels/cuda/glm_glue_moe.cuh index d9cd649159..edb23b9ce5 100644 --- a/zig/kernels/cuda/glm_glue_moe.cuh +++ b/zig/kernels/cuda/glm_glue_moe.cuh @@ -16,10 +16,13 @@ __device__ __forceinline__ float group64(const float* q) { return g; } +// AFTER false: GLM's silu(min(g, L)); true: Ling's min(silu(g), L) (inclusionAI's SwigluStepAndMul). Both clip u. +template __device__ __forceinline__ float swiglu1(float g, float u, float limit) { - g = fminf(g, limit); + if (!AFTER) g = fminf(g, limit); u = fminf(fmaxf(u, 0.0f - limit), limit); - const float s = rbf(div_full(g, exp_approx(0.0f - g) + 1.0f)); + float s = rbf(div_full(g, exp_approx(0.0f - g) + 1.0f)); + if (AFTER) s = fminf(s, limit); return rbf(__fmul_rn(u, s)); } @@ -68,35 +71,43 @@ __device__ __forceinline__ void keep_groups(float* choice, int lane, int ne, int } // namespace glm_moe // glue._swiglu with BLOCK 512 (128 threads, 4 values each, 64-group sums over 16 lanes), column block blockIdx.y. -extern "C" __global__ void __launch_bounds__(128) - tf_glm_swiglu_512(const __nv_bfloat16* gu, __nv_bfloat16* out, float* xs, float limit, int64_t w) { - const int64_t r = blockIdx.x, d0 = blockIdx.y * 512 + threadIdx.x * 4; - const uint2 gv = *reinterpret_cast(gu + r * 2 * w + d0); - const uint2 uv = *reinterpret_cast(gu + r * 2 * w + w + d0); - const __nv_bfloat16* gh = reinterpret_cast(&gv); - const __nv_bfloat16* uh = reinterpret_cast(&uv); - float a[4]; - uint2 ov; - __nv_bfloat16* oh = reinterpret_cast<__nv_bfloat16*>(&ov); - for (int j = 0; j < 4; ++j) { - a[j] = glm_moe::swiglu1(bf(gh[j]), bf(uh[j]), limit); - oh[j] = __float2bfloat16_rn(a[j]); +#define TF_GLM_SWIGLU_512(NAME, AFTER) \ + extern "C" __global__ void __launch_bounds__(128) \ + NAME(const __nv_bfloat16* gu, __nv_bfloat16* out, float* xs, float limit, int64_t w) { \ + const int64_t r = blockIdx.x, d0 = blockIdx.y * 512 + threadIdx.x * 4; \ + const uint2 gv = *reinterpret_cast(gu + r * 2 * w + d0); \ + const uint2 uv = *reinterpret_cast(gu + r * 2 * w + w + d0); \ + const __nv_bfloat16* gh = reinterpret_cast(&gv); \ + const __nv_bfloat16* uh = reinterpret_cast(&uv); \ + float a[4]; \ + uint2 ov; \ + __nv_bfloat16* oh = reinterpret_cast<__nv_bfloat16*>(&ov); \ + for (int j = 0; j < 4; ++j) { \ + a[j] = glm_moe::swiglu1(bf(gh[j]), bf(uh[j]), limit); \ + oh[j] = __float2bfloat16_rn(a[j]); \ + } \ + *reinterpret_cast(out + r * w + d0) = ov; \ + const float g = glm_moe::group64<4>(a); \ + if ((threadIdx.x & 15) == 0) xs[r * (w / 64) + d0 / 64] = g; \ } - *reinterpret_cast(out + r * w + d0) = ov; - const float g = glm_moe::group64<4>(a); - if ((threadIdx.x & 15) == 0) xs[r * (w / 64) + d0 / 64] = g; -} // glue._swiglu with BLOCK 64 (the TP=3 shared expert's 704): one value a thread, the warp's butterfly, then 2 warps. -extern "C" __global__ void __launch_bounds__(64) - tf_glm_swiglu_64(const __nv_bfloat16* gu, __nv_bfloat16* out, float* xs, float limit, int64_t w) { - __shared__ float smem[2]; - const int64_t r = blockIdx.x, d = blockIdx.y * 64 + threadIdx.x; - const float a = glm_moe::swiglu1(bf(gu[r * 2 * w + d]), bf(gu[r * 2 * w + w + d]), limit); - out[r * w + d] = __float2bfloat16_rn(a); - const float g = block_sum<2>(a, smem); - if (threadIdx.x == 0) xs[r * (w / 64) + blockIdx.y] = g; -} +#define TF_GLM_SWIGLU_64(NAME, AFTER) \ + extern "C" __global__ void __launch_bounds__(64) \ + NAME(const __nv_bfloat16* gu, __nv_bfloat16* out, float* xs, float limit, int64_t w) { \ + __shared__ float smem[2]; \ + const int64_t r = blockIdx.x, d = blockIdx.y * 64 + threadIdx.x; \ + const float a = glm_moe::swiglu1(bf(gu[r * 2 * w + d]), bf(gu[r * 2 * w + w + d]), limit); \ + out[r * w + d] = __float2bfloat16_rn(a); \ + const float g = block_sum<2>(a, smem); \ + if (threadIdx.x == 0) xs[r * (w / 64) + blockIdx.y] = g; \ + } + +TF_GLM_SWIGLU_512(tf_glm_swiglu_512, false) +TF_GLM_SWIGLU_64(tf_glm_swiglu_64, false) +// Ling's order, the limit on silu's output: min(silu(g), L) * clip(u, -L, L). +TF_GLM_SWIGLU_512(tf_ling_swiglu_512, true) +TF_GLM_SWIGLU_64(tf_ling_swiglu_64, true) // qmm._group_sums: row blockIdx.x's 64-input sums, 8 lanes a group of 8 consecutive inputs each, 16 groups a block. extern "C" __global__ void __launch_bounds__(128) diff --git a/zig/kernels/cuda/nvfp4_experts.cu b/zig/kernels/cuda/nvfp4_experts.cu index be724dead8..611d010bde 100644 --- a/zig/kernels/cuda/nvfp4_experts.cu +++ b/zig/kernels/cuda/nvfp4_experts.cu @@ -310,14 +310,15 @@ __device__ __forceinline__ void prompt_body(const __nv_bfloat16* __restrict__ X, #undef TF_TILES } -// The instances: SwiGLU gate-up, fp32 and bf16 down, and relu^2 up for experts without a gate (Nemotron). +// The instances: SwiGLU gate-up (EPI 4 in Ling's order), fp32 and bf16 down, relu^2 up without a gate (Nemotron). template __global__ void nvfp4_expert_kernel<2, 2, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); +template __global__ void nvfp4_expert_kernel<2, 4, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); template __global__ void nvfp4_expert_kernel<1, 0, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); template __global__ void nvfp4_expert_kernel<1, 3, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); template __global__ void nvfp4_expert_kernel<1, 1, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); } -// The prompt form's entries, 4 row tiles a warp in one row of WN warps: relu^2 up, SwiGLU gate-up, bf16 sums down. +// The prompt form's entries, 4 row tiles a warp in a row of WN warps: relu^2 up, SwiGLU gate-up (2 orders), down. #define TF_NVFP4_PROMPT(NAME, M, EPI, WN) \ extern "C" __global__ void __launch_bounds__(WN * 32) NAME( \ const __nv_bfloat16* __restrict__ X, int x_stride, int slots, const uint4* __restrict__ W, \ @@ -329,3 +330,4 @@ template __global__ void nvfp4_expert_kernel<1, 1, 4>(const __nv_bfloat16* __res TF_NVFP4_PROMPT(tf_nvfp4_prompt_relu2, 1, 1, 8) TF_NVFP4_PROMPT(tf_nvfp4_prompt_swiglu, 2, 2, 4) TF_NVFP4_PROMPT(tf_nvfp4_prompt_down, 1, 3, 8) +TF_NVFP4_PROMPT(tf_nvfp4_prompt_swiglu_after, 2, 4, 4) diff --git a/zig/src/cuda/experts.zig b/zig/src/cuda/experts.zig index 8a6664cbd6..8f6d950457 100644 --- a/zig/src/cuda/experts.zig +++ b/zig/src/cuda/experts.zig @@ -26,6 +26,7 @@ pub const Layer = struct { dims: usize, // D experts: usize, limit: f32 = 0.0, // fp8g, nvfp4: the SwiGLU clamp, 0 for none + after: bool = false, // fp8g, nvfp4: clamp silu's output (Ling: min(silu(g), L)), not its input (silu(min(g, L))) relu2: bool = false, // nvfp4: up is one projection under relu^2 (Nemotron), not SwiGLU's gate and up }; @@ -47,7 +48,7 @@ pub const Grouped = struct { .affine4 => { const a = g.affine4 orelse return error.FormatNotLoaded; if (epi == .down_bf16) return error.EpilogueNotBuilt; - if (skip >= 0 or l.limit != 0) return error.OptionNotBuilt; + if (skip >= 0 or l.limit != 0 or l.after) return error.OptionNotBuilt; const up = epi == .up; const nb = (if (up) l.width else l.dims) / 32; const w = Affine4Experts.view(l, up); @@ -64,7 +65,7 @@ pub const Grouped = struct { .affine4 => { const a = g.affine4 orelse return error.FormatNotLoaded; if (epi == .down_f32) return error.EpilogueNotBuilt; - if (skip >= 0 or l.limit != 0) return error.OptionNotBuilt; + if (skip >= 0 or l.limit != 0 or l.after) return error.OptionNotBuilt; const up = epi == .up; const nb = (if (up) l.width else l.dims) / 32; const w = Affine4Experts.view(l, up); @@ -179,6 +180,8 @@ pub const Fp8Experts = struct { mod: Module, fns: [2][3]Function, // decode, prompt; by Epi resident: [2][3]usize, // blocks the grid may hold: per SM times SMs + after: [2]Function, // decode, prompt: gate-up in Ling's clamp order (Layer.after) + after_resident: [2]usize, pub const cols = 32; // output columns a warp pub const block = 128; // the checkpoint's scale block, both ways @@ -190,18 +193,24 @@ pub const Fp8Experts = struct { pub const p_gate_up = "_ZN14tf_fp8_experts24fp8_expert_prompt_kernelILi2ELi2ELi2ELi2EEEvPK13__nv_bfloat16iiPK5uint4PKfiiPKiSA_SA_Pvifi"; pub const p_down_f32 = "_ZN14tf_fp8_experts24fp8_expert_prompt_kernelILi1ELi0ELi4ELi4EEEvPK13__nv_bfloat16iiPK5uint4PKfiiPKiSA_SA_Pvifi"; pub const p_down_bf16 = "_ZN14tf_fp8_experts24fp8_expert_prompt_kernelILi1ELi3ELi4ELi4EEEvPK13__nv_bfloat16iiPK5uint4PKfiiPKiSA_SA_Pvifi"; + pub const gate_up_after = "_ZN14tf_fp8_experts17fp8_expert_kernelILi2ELi4ELi4EEEvPK13__nv_bfloat16iiPK5uint4PKfiiPKiSA_SA_Pvifi"; + pub const p_gate_up_after = "_ZN14tf_fp8_experts24fp8_expert_prompt_kernelILi2ELi4ELi2ELi2EEEvPK13__nv_bfloat16iiPK5uint4PKfiiPKiSA_SA_Pvifi"; }; pub fn load(d: *const Driver, sms: usize) !Fp8Experts { if (!kernels.available) return error.BuiltWithoutKernels; var mod = try Module.load(d, kernels.fp8_experts); errdefer mod.unload(); - var e: Fp8Experts = .{ .mod = mod, .fns = undefined, .resident = undefined }; + var e: Fp8Experts = .{ .mod = mod, .fns = undefined, .resident = undefined, .after = undefined, .after_resident = undefined }; const names = [2][3][:0]const u8{ .{ symbols.gate_up, symbols.down_f32, symbols.down_bf16 }, .{ symbols.p_gate_up, symbols.p_down_f32, symbols.p_down_bf16 } }; for (names, 0..) |set, i| for (set, 0..) |name, j| { e.fns[i][j] = try mod.function(name); e.resident[i][j] = @max(1, try e.fns[i][j].occupancy(128, 0)) * sms; }; + for ([2][:0]const u8{ symbols.gate_up_after, symbols.p_gate_up_after }, 0..) |name, i| { + e.after[i] = try mod.function(name); + e.after_resident[i] = @max(1, try e.after[i].occupancy(128, 0)) * sms; + } return e; } @@ -217,10 +226,12 @@ pub const Fp8Experts = struct { const n = if (gate) l.width else l.dims; const pi: usize = @intFromBool(prompt); const ei: usize = @backingInt(epi); + const after = gate and l.after; + const resident = if (after) e.after_resident[pi] else e.resident[pi][ei]; const grid: usize = if (prompt) blk: { const cw: usize = if (gate) 2 else 4; - break :blk @min(items * (nb / cw), e.resident[pi][ei]); - } else @min((items * nb + 3) / 4, e.resident[pi][ei]); + break :blk @min(items * (nb / cw), resident); + } else @min((items * nb + 3) / 4, resident); if (grid < 1) return; var a: launch_.Args = .{}; a.add(in.x); @@ -234,7 +245,7 @@ pub const Fp8Experts = struct { a.add(@as(i32, @intCast(n))); a.add(if (gate) l.limit else @as(f32, 0.0)); a.add(skip); - try launch_.launch(e.fns[pi][ei], .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = 128, .y = 1, .z = 1 } }, s, &a); + try launch_.launch(if (after) e.after[pi] else e.fns[pi][ei], .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = 128, .y = 1, .z = 1 } }, s, &a); } /// e4m3 [n, k] of one expert as int32 [n/32][k/32][32][2][4]: lane (gq, t)'s half h, tile j (fp8/experts.py pack). @@ -270,9 +281,9 @@ pub const Fp8Experts = struct { /// NVFP4 experts (nvfp4_experts.cu): decode rows 16 pairs at a time; prompt rows staged, 64 pairs a CTA. pub const Nvfp4Experts = struct { mod: Module, - fns: [4]Function, // by Epi, then relu^2 up - resident: [4]usize, // blocks the grid may hold: per SM times SMs - prompt_fns: [3]Function, // relu^2 up, SwiGLU gate-up, bf16 sums down + fns: [5]Function, // by Epi, then relu^2 up, then gate-up in Ling's clamp order (Layer.after) + resident: [5]usize, // blocks the grid may hold: per SM times SMs + prompt_fns: [4]Function, // relu^2 up, SwiGLU gate-up, bf16 sums down, gate-up in Ling's order pub const cols = 32; // output columns a block pub const words = 144; // int32 a (32 columns, 32 inputs) block: 128 code words, then 16 of e4m3 scales @@ -282,11 +293,12 @@ pub const Nvfp4Experts = struct { pub const down_f32 = "_ZN16tf_nvfp4_experts19nvfp4_expert_kernelILi1ELi0ELi4EEEvPK13__nv_bfloat16iiPK5uint4PKfiiPKiSA_SA_Pvifi"; pub const down_bf16 = "_ZN16tf_nvfp4_experts19nvfp4_expert_kernelILi1ELi3ELi4EEEvPK13__nv_bfloat16iiPK5uint4PKfiiPKiSA_SA_Pvifi"; pub const up_relu2 = "_ZN16tf_nvfp4_experts19nvfp4_expert_kernelILi1ELi1ELi4EEEvPK13__nv_bfloat16iiPK5uint4PKfiiPKiSA_SA_Pvifi"; - pub const prompt = [3][:0]const u8{ "tf_nvfp4_prompt_relu2", "tf_nvfp4_prompt_swiglu", "tf_nvfp4_prompt_down" }; + pub const gate_up_after = "_ZN16tf_nvfp4_experts19nvfp4_expert_kernelILi2ELi4ELi4EEEvPK13__nv_bfloat16iiPK5uint4PKfiiPKiSA_SA_Pvifi"; + pub const prompt = [4][:0]const u8{ "tf_nvfp4_prompt_relu2", "tf_nvfp4_prompt_swiglu", "tf_nvfp4_prompt_down", "tf_nvfp4_prompt_swiglu_after" }; }; - /// The prompt entries' warps (32-column blocks) a CTA: relu^2 up, SwiGLU gate-up, down (GB10's fastest tiles). - pub const prompt_wn = [3]usize{ 8, 4, 8 }; + /// The prompt entries' warps (32-column blocks) a CTA: relu^2 up, SwiGLU gate-up, down, Ling's gate-up. + pub const prompt_wn = [4]usize{ 8, 4, 8, 4 }; /// Prompt4's dynamic shared bytes: stages of 64 rows' 64 inputs, then WN x 2 blocks of `parts`. pub fn promptSmem(parts: usize, wn: usize) u32 { @@ -300,11 +312,11 @@ pub const Nvfp4Experts = struct { var mod = try Module.load(d, kernels.nvfp4_experts); errdefer mod.unload(); var e: Nvfp4Experts = .{ .mod = mod, .fns = undefined, .resident = undefined, .prompt_fns = undefined }; - for ([_][:0]const u8{ symbols.gate_up, symbols.down_f32, symbols.down_bf16, symbols.up_relu2 }, 0..) |name, i| { + for ([_][:0]const u8{ symbols.gate_up, symbols.down_f32, symbols.down_bf16, symbols.up_relu2, symbols.gate_up_after }, 0..) |name, i| { e.fns[i] = try mod.function(name); e.resident[i] = @max(1, try e.fns[i].occupancy(128, 0)) * sms; } - for (symbols.prompt, &e.prompt_fns, [3]usize{ 1, 2, 1 }, prompt_wn) |name, *f, parts, wn| { + for (symbols.prompt, &e.prompt_fns, [4]usize{ 1, 2, 1, 2 }, prompt_wn) |name, *f, parts, wn| { f.* = try mod.function(name); try f.allowDynamicShared(promptSmem(parts, wn)); } @@ -320,7 +332,7 @@ pub const Nvfp4Experts = struct { const gate = epi == .up; const kg = (if (gate) l.dims else l.width) / 32; const nb = (if (gate) l.width else l.dims) / cols; - const ei: usize = if (gate and l.relu2) 3 else @backingInt(epi); + const ei: usize = if (gate and l.relu2) 3 else if (gate and l.after) 4 else @backingInt(epi); const grid = @min((items * nb + 3) / 4, e.resident[ei]); if (grid < 1) return; var a: launch_.Args = .{}; @@ -357,7 +369,7 @@ pub const Nvfp4Experts = struct { a.add(@as(i32, @intCast(if (gate) l.width else l.dims))); a.add(if (gate) l.limit else @as(f32, 0.0)); a.add(skip); - const fi: usize = if (!gate) 2 else if (l.relu2) 0 else 1; + const fi: usize = if (!gate) 2 else if (l.relu2) 0 else if (l.after) 3 else 1; const cfg: launch_.Config = .{ .grid = .{ .x = @intCast(items * ((nb + prompt_wn[fi] - 1) / prompt_wn[fi])), .y = 1, .z = 1 }, .block = .{ .x = @intCast(prompt_wn[fi] * 32), .y = 1, .z = 1 }, .shared = promptSmem(parts, prompt_wn[fi]) }; try launch_.launch(e.prompt_fns[fi], cfg, s, &a); } diff --git a/zig/src/families/glm/cuda_glue.zig b/zig/src/families/glm/cuda_glue.zig index dbc52cc201..01c71a3f73 100644 --- a/zig/src/families/glm/cuda_glue.zig +++ b/zig/src/families/glm/cuda_glue.zig @@ -128,8 +128,13 @@ fn swigluBlock(w: usize) !usize { /// glue.swiglu: [gate | up] -> bf16(bf16(silu(min(g, L))) * clip(u, -L, L)), 64-sums in `xs`; L per layer, 0 for none. pub fn swiglu(set: Set, stream: cuda.Stream, gu: u64, out: u64, xs: u64, limit: f32, rows: usize, w: usize) !void { + return swigluWith(.{ set.glue.swiglu_512, set.glue.swiglu_64 }, stream, gu, out, xs, limit, rows, w); +} + +/// swiglu on a pair of glm_glue_moe.cuh's 512 and 64 kernels (tf_ling_swiglu_*: the limit on silu's output, Ling). +pub fn swigluWith(fs: [2]cuda.Function, stream: cuda.Stream, gu: u64, out: u64, xs: u64, limit: f32, rows: usize, w: usize) !void { const block = try swigluBlock(w); - const f, const threads: u32 = if (block == 512) .{ set.glue.swiglu_512, 128 } else .{ set.glue.swiglu_64, 64 }; + const f, const threads: u32 = if (block == 512) .{ fs[0], 128 } else .{ fs[1], 64 }; var a: cuda.Args = .{}; inline for (.{ gu, out, xs, clampOf(limit), i64of(w) }) |v| a.add(v); try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(w / block) }, .block = .{ .x = threads } }, stream, &a); diff --git a/zig/src/families/ling/block_ffn.zig b/zig/src/families/ling/block_ffn.zig index 23171fbbcf..e6f0c4d0fa 100644 --- a/zig/src/families/ling/block_ffn.zig +++ b/zig/src/families/ling/block_ffn.zig @@ -10,12 +10,12 @@ const glue = graph.glue; const Ops = graph.Ops; const Config = @import("config.zig").Config; -/// Gate and up into [gate | up] rows, SwiGLU clamped at `limit` (0: none), down into `out` [rows][hidden]. +/// Gate and up into [gate | up] rows, SwiGLU with min(silu(g), limit) (0: none), down into `out` [rows][hidden]. pub fn mlp(o: Ops, b: *const state.Buffers, m: weights.Mlp, x: u64, ldx: usize, rows: usize, limit: f32, out: u64) !void { const w = m.gate.outputs(); try o.mmCols(m.gate, x, ldx, rows, b.tmp, b.gu, 2 * w, 0); try o.mmCols(m.up, x, ldx, rows, b.tmp, b.gu, 2 * w, w); - try glue.swiglu(o.set(), o.s, b.gu, b.act, b.xs, limit, rows, w); + try glue.swigluWith(o.m.swiglu, o.s, b.gu, b.act, b.xs, limit, rows, w); try o.mm(m.down, b.act, w, rows, out); } diff --git a/zig/src/families/ling/graph.zig b/zig/src/families/ling/graph.zig index 3296fe18fe..9a2a645265 100644 --- a/zig/src/families/ling/graph.zig +++ b/zig/src/families/ling/graph.zig @@ -25,6 +25,7 @@ pub const Modules = struct { add: cuda.Function, gate: cuda.Function, combine: cuda.Function, + swiglu: [2]cuda.Function, // glm_glue's 512 and 64 SwiGLU in Ling's order, min(silu(g), L) * clip(u, -L, L) pub fn load(d: *const cuda.Driver, ctx: *const cuda.Context) !Modules { const k = cuda.kernels; @@ -44,6 +45,7 @@ pub const Modules = struct { m.add = try m.mods[4].function("tf_ling_add"); m.gate = try m.mods[4].function("tf_ling_gate"); m.combine = try m.mods[4].function("tf_ling_combine"); + m.swiglu = .{ try m.mods[0].function("tf_ling_swiglu_512"), try m.mods[0].function("tf_ling_swiglu_64") }; m.lane = try cuda.qmmf.Lane.load(d, @intCast(try ctx.attribute(.compute_capability_major))); errdefer m.lane.unload(); m.prompt = try cuda.nvfp4.Prompt.load(d); diff --git a/zig/src/families/ling/weights_proj.zig b/zig/src/families/ling/weights_proj.zig index a1b5dbd8f7..e0e372261c 100644 --- a/zig/src/families/ling/weights_proj.zig +++ b/zig/src/families/ling/weights_proj.zig @@ -281,7 +281,7 @@ pub const Reader = struct { try us.upload((2 * e + 1) * inv_n * 4, u.s.bytes); try r.view(down_scale, count * inv_n * 4).upload(e * inv_n * 4, dn.s.bytes); } - return .{ .format = .fp8g, .up = up, .down = down, .up_scale = up_scale, .down_scale = down_scale, .width = ni, .dims = d, .experts = count, .limit = limit }; + return .{ .format = .fp8g, .up = up, .down = down, .up_scale = up_scale, .down_scale = down_scale, .width = ni, .dims = d, .experts = count, .limit = limit, .after = true }; } /// experts.Layer's nvfp4 form: Nvfp4Experts.packGateUp and packOne per expert, weight_scale_2 as the globals. @@ -315,7 +315,7 @@ pub const Reader = struct { } const up_scale = try r.store.keep(std.mem.sliceAsBytes(gu), try label(&lb, i, "up_scale")); const down_scale = try r.store.keep(std.mem.sliceAsBytes(gd), try label(&lb, i, "down_scale")); - return .{ .format = .nvfp4, .up = up, .down = down, .up_scale = up_scale, .down_scale = down_scale, .width = ni, .dims = d, .experts = count, .limit = limit }; + return .{ .format = .nvfp4, .up = up, .down = down, .up_scale = up_scale, .down_scale = down_scale, .width = ni, .dims = d, .experts = count, .limit = limit, .after = true }; } }; diff --git a/zig/tests/cuda/ling_clamp.zig b/zig/tests/cuda/ling_clamp.zig new file mode 100644 index 0000000000..e2b3b24ee9 --- /dev/null +++ b/zig/tests/cuda/ling_clamp.zig @@ -0,0 +1,269 @@ +//! Ling's SwiGLU order, min(silu(g), L) * clip(u, -L, L): glue and grouped experts vs float64, old kernels' bytes. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const glm = @import("glm_cuda"); +const ling = @import("ling_cuda"); +const check = @import("check.zig"); +const qr = @import("quant_ref_tests.zig"); +const lw = @import("ling_weights_tests.zig"); + +const Rig = qr.RigT; +const h = qr.helpers; +const W = ling.weights; +const ex = cuda.experts; +const glue = glm.glue; +const Config = ling.config.Config; +const two7 = std.math.pow(f64, 2, -7); + +/// Ling's activation in float64 on bf16 values: silu(g) as `s`, then min(s, L) and clip(u, -L, L); L 0 for none. +fn act(s: f64, u: f64, limit: f64) f64 { + if (limit == 0) return s * u; + return @min(s, limit) * std.math.clamp(u, -limit, limit); +} + +fn silu(g: f64) f64 { + return g / (1 + @exp(-g)); +} + +/// The glue's tf_ling_swiglu at L 10 and 2 on gates and ups straddling L vs float64; GLM's order beside it. +fn glueOrder(r: *Rig, fs: [2]cuda.Function, set: glue.Set) !void { + for ([_]f32{ 10, 2 }) |lim| for ([_]usize{ 1024, 704 }) |w| { + const rows = 40; + const gu = try r.bfs(rows * 2 * w, 1.5 * lim); + const dgu = try r.dev(std.mem.sliceAsBytes(gu)); + const outs = [2]u64{ try r.zeros(rows * w * 2), try r.zeros(rows * w * 2) }; + const xs = try r.zeros(rows * w / 64 * 4); + try glue.swigluWith(fs, r.s, dgu, outs[0], xs, lim, rows, w); + try glue.swiglu(set, r.s, dgu, outs[1], xs, lim, rows, w); + const got = try r.back(outs[0], rows * w * 2); + const before = try r.back(outs[1], rows * w * 2); + var most: f64 = 0; + var above: usize = 0; + var wide: usize = 0; + var differ: usize = 0; + for (0..rows) |m| for (0..w) |d| { + const g = h.bfOf(gu[m * 2 * w + d]); + const u = h.bfOf(gu[m * 2 * w + w + d]); + const want = act(silu(g), u, lim); + const i = m * w + d; + const v = h.bfOf(std.mem.bytesToValue(u16, got[2 * i ..][0..2])); + most = @max(most, @abs(v - want) / (@abs(want) * two7 + 1e-30)); + above += @intFromBool(silu(g) > lim); + wide += @intFromBool(@abs(u) > lim); + differ += @intFromBool(!std.mem.eql(u8, got[2 * i ..][0..2], before[2 * i ..][0..2])); + }; + r.worst = @max(r.worst, most); + std.debug.print("RESULT glue L {d}, width {d}: float64 at {d:.3} of 2^-7; silu(g) > L on {d}, |u| > L on {d} of {d}; {d} differ from GLM's order\n", .{ lim, w, most, above, wide, rows * w, differ }); + try check.expect(most <= 1 and above > 0 and wide > 0, "tf_ling_swiglu at L {d}: Ling's order within 2^-7 of float64", .{lim}); + // bf16(silu(g)) is 10 for every bf16 g >= 10, so the two orders agree on bytes at 10 and part at 2 + try check.expect(if (lim == 10) differ == 0 else differ > 0, "L {d}: {d} bytes differ from GLM's order", .{ lim, differ }); + }; +} + +/// A plan for `pairs` pairs, pair p on expert lw.pick(p), at the given tile. +fn plan(r: *Rig, kit: *const lw.Kit, pairs: usize, slots: usize, count: usize, tile: usize) !cuda.grouped.Plan { + const pk = try r.a.alloc(i32, pairs); + for (pk, 0..) |*e, p| e.* = lw.pick(p, slots, count); + const p: cuda.grouped.Plan = .{ .members = try r.zeros(pairs * 4), .items = try r.zeros(cuda.grouped.maxItems(pairs, count, 16) * 12), .counts = try r.zeros(8), .rank = try r.zeros(4), .hist = try r.zeros(4) }; + try kit.router.route(r.s, try r.dev(std.mem.sliceAsBytes(pk)), pairs, count, tile, p); + return p; +} + +/// The grouped gate-up with Layer.after at L 10 and 2, decode and prompt rows, vs float64 from the file. +fn expertsOrder(r: *Rig, kit: *const lw.Kit, ck: *core.Checkpoint, c: *const Config, i: u32, l: ex.Layer) !void { + const slots = @min(4, l.experts); + const ni = l.width; + const d = l.dims; + var b: [2][192]u8 = undefined; + const e0: u32 = @intCast(lw.pick(0, slots, l.experts)); + const g0 = try lw.Mat.read(ck, c, try ling.names.moduleName(&b[0], i, .expert_gate, e0), ni, d); + const probe = try lw.reference(r.a, g0, try r.bfs(d, 1), 1); + var ss: f64 = 0; + for (probe.want) |v| ss += v * v; + const rms = @sqrt(ss / @as(f64, @floatFromInt(ni))); + for ([_]bool{ false, true }) |prompt| for ([_]f32{ 10, 2 }) |lim| { + const rows: usize = if (prompt) 24 else 3; + const pairs = rows * slots; + const tile: usize = if (prompt) ex.promptTile(l.format) else 16; + const p = try plan(r, kit, pairs, slots, l.experts, tile); + const x = try r.bfs(rows * d, @floatCast(std.math.pow(f64, 2, @round(std.math.log2(1.2 * lim / rms))))); + var lc = l; + lc.limit = lim; + lc.after = true; + const out = try r.zeros(pairs * ni * 2); + const in: ex.Rows = .{ .x = try r.dev(std.mem.sliceAsBytes(x)), .stride = d, .slots = slots }; + const items = cuda.grouped.maxItems(pairs, l.experts, tile); + if (prompt) try kit.grouped().prompt(r.s, .up, in, lc, p, out, items, -1) else try kit.grouped().decode(r.s, .up, in, lc, p, out, items, -1); + const want = try r.a.alloc(f64, pairs * ni); + const tol = try r.a.alloc(f64, pairs * ni); + var above: usize = 0; + for (0..pairs) |q| { + const e: u32 = @intCast(lw.pick(q, slots, l.experts)); + const g = try lw.Mat.read(ck, c, try ling.names.moduleName(&b[0], i, .expert_gate, e), ni, d); + const u = try lw.Mat.read(ck, c, try ling.names.moduleName(&b[1], i, .expert_up, e), ni, d); + const gr = try lw.reference(r.a, g, x[(q / slots) * d ..][0..d], 1); + const ur = try lw.reference(r.a, u, x[(q / slots) * d ..][0..d], 1); + for (0..ni) |j| { + const gv = h.bfOf(h.bf16Of(@floatCast(gr.want[j]))); + const uv = h.bfOf(h.bf16Of(@floatCast(ur.want[j]))); + const s = h.bfOf(h.bf16Of(@floatCast(silu(gv)))); + want[q * ni + j] = act(s, uv, lim); + above += @intFromBool(s > lim); + // the bf16 steps (2^-6), then each product's bound: silu' under 1.1, min and clip 1-Lipschitz + tol[q * ni + j] = @abs(want[q * ni + j]) * std.math.pow(f64, 2, -6) + 1.1 * @min(@abs(uv), lim) * gr.tol[j] + @min(@abs(s), lim) * ur.tol[j] + 1e-30; + } + } + var name: [128]u8 = undefined; + try r.close(try std.fmt.bufPrint(&name, "layer {d} {t} gate-up L {d} ({s}), silu(g) > L on {d} of {d}", .{ i, l.format, lim, if (prompt) "prompt" else "decode", above, pairs * ni }), out, false, want, tol); + try check.expect(above > 0, "some gates past L", .{}); + }; +} + +/// `ling-clamp DIR`: the glue's Ling order vs float64, then the clamped MoE layer's grouped experts (FP8 and NVFP4). +pub fn run(gpu: check.Gpu, dir: []const u8) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(0xc1a); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*x| x.free(); + r.bufs.deinit(gpu.gpa); + } + var gm = try cuda.Module.load(gpu.d, cuda.kernels.glm_glue); + defer gm.unload(); + const ks = try glue.Kernels.load(gm); + try glueOrder(&r, .{ try gm.function("tf_ling_swiglu_512"), try gm.function("tf_ling_swiglu_64") }, .{ .glue = &ks }); + var kit: lw.Kit = undefined; + try kit.init(gpu); + defer kit.deinit(); + var why: ling.config.Why = .{}; + var o = W.open(gpu.gpa, gpu.io, dir, &why) catch |e| { + std.debug.print("{s}: {s}\n", .{ dir, why.text() }); + return e; + }; + defer o.ck.close(); + const last: u32 = o.c.layers - 1; + var w = try W.load(gpu.gpa, gpu.d, s, &kit.pack, &o.ck, o.c, &.{last}, .{ .globals = false }); + defer w.deinit(); + const moe = w.layers[0].ffn.moe; + try check.expect(moe.shared_limit == ling.limits(&w.config, last).shared, "shared limit from the config", .{}); + switch (moe.routed) { + .grouped => |l| { + try check.expect(l.after and l.limit == ling.limits(&w.config, last).routed, "loader: Ling's order at the config's limit", .{}); + try expertsOrder(&r, &kit, &o.ck, &w.config, last, l); + }, + .lanes => std.debug.print("RESULT {t}: routed experts run as lanes, through the glue checked above\n", .{o.c.quant}), + } + check.pass("ling-clamp {t}: min(silu(g), L) * clip(u, -L, L) within float64 bounds (worst {d:.3}), GLM's order beside it", .{ o.c.quant, r.worst }); +} + +/// An old fatbin's instance, launched with the arguments and grid the host gave it before Layer.after existed. +const Old = struct { f: cuda.Function, threads: u32, smem: u32 = 0, per: usize, cap: usize }; + +fn launchOld(r: *Rig, k: Old, in: ex.Rows, l: ex.Layer, gate: bool, p: cuda.grouped.Plan, out: u64, units: usize) !void { + var a: cuda.Args = .{}; + a.add(in.x); + a.add(@as(i32, @intCast(in.stride))); + a.add(@as(i32, @intCast(if (gate) in.slots else 0))); + a.add(if (gate) l.up else l.down); + a.add(if (gate) l.up_scale else l.down_scale); + a.add(@as(i32, @intCast((if (gate) l.dims else l.width) / 32))); + a.add(@as(i32, @intCast((if (gate) l.width else l.dims) / 32))); + for ([_]u64{ p.items, p.counts, p.members, out }) |v| a.add(v); + a.add(@as(i32, @intCast(if (gate) l.width else l.dims))); + a.add(if (gate) l.limit else @as(f32, 0.0)); + a.add(@as(i32, -1)); + const grid = @min((units + k.per - 1) / k.per, k.cap); + try cuda.launch.launch(k.f, .{ .grid = .{ .x = @intCast(grid) }, .block = .{ .x = k.threads }, .shared = k.smem }, r.s, &a); +} + +fn oldFn(m: cuda.Module, name: [:0]const u8, threads: u32, smem: u32, per: usize, sms: usize, resident: bool) !Old { + const f = try m.function(name); + if (smem > 0) try f.allowDynamicShared(smem); + return .{ .f = f, .threads = threads, .smem = smem, .per = per, .cap = if (resident) @max(1, try f.occupancy(128, 0)) * sms else std.math.maxInt(usize) }; +} + +/// `ling-clamp-old DIR FATBINS`: each grouped instance this build kept, via experts.Grouped, vs FATBINS' kernel. +pub fn old(gpu: check.Gpu, dir: []const u8, fatbins: []const u8) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(0x01d); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*x| x.free(); + r.bufs.deinit(gpu.gpa); + } + var kit: lw.Kit = undefined; + try kit.init(gpu); + defer kit.deinit(); + var why: ling.config.Why = .{}; + var o = try W.open(gpu.gpa, gpu.io, dir, &why); + defer o.ck.close(); + const last: u32 = o.c.layers - 1; + var w = try W.load(gpu.gpa, gpu.d, s, &kit.pack, &o.ck, o.c, &.{last}, .{ .globals = false }); + defer w.deinit(); + const l = w.layers[0].ffn.moe.routed.grouped; + const fp8 = l.format == .fp8g; + const path = try std.fs.path.join(r.a, &.{ fatbins, if (fp8) "fp8_experts.fatbin" else "nvfp4_experts.fatbin" }); + const bytes = try std.Io.Dir.cwd().readFileAllocOptions(gpu.io, path, r.a, .limited(1 << 26), .@"64", null); + var m = try cuda.Module.load(gpu.d, bytes); + defer m.unload(); + const sms: usize = @intCast(try gpu.ctx.attribute(.multiprocessor_count)); + const F8 = ex.Fp8Experts.symbols; + const F4 = ex.Nvfp4Experts.symbols; + const p4 = ex.Nvfp4Experts.promptSmem; + // [decode, prompt] x [gate-up, down_f32, down_bf16, relu^2]; null where the format has no such instance + const ks: [2][4]?Old = if (fp8) .{ + .{ try oldFn(m, F8.gate_up, 128, 0, 4, sms, true), try oldFn(m, F8.down_f32, 128, 0, 4, sms, true), try oldFn(m, F8.down_bf16, 128, 0, 4, sms, true), null }, + .{ try oldFn(m, F8.p_gate_up, 128, 0, 1, sms, true), try oldFn(m, F8.p_down_f32, 128, 0, 1, sms, true), try oldFn(m, F8.p_down_bf16, 128, 0, 1, sms, true), null }, + } else .{ + .{ try oldFn(m, F4.gate_up, 128, 0, 4, sms, true), try oldFn(m, F4.down_f32, 128, 0, 4, sms, true), try oldFn(m, F4.down_bf16, 128, 0, 4, sms, true), try oldFn(m, F4.up_relu2, 128, 0, 4, sms, true) }, + .{ try oldFn(m, F4.prompt[1], 128, p4(2, 4), 1, sms, false), null, try oldFn(m, F4.prompt[2], 256, p4(1, 8), 1, sms, false), try oldFn(m, F4.prompt[0], 256, p4(1, 8), 1, sms, false) }, + }; + const slots = @min(4, l.experts); + var n: usize = 0; + var same_at_10: usize = 0; + for ([_]bool{ false, true }, 0..) |prompt, pi| { + const rows: usize = if (prompt) 40 else 3; + const pairs = rows * slots; + const tile: usize = if (prompt) ex.promptTile(l.format) else 16; + const p = try plan(&r, &kit, pairs, slots, l.experts, tile); + const items = cuda.grouped.maxItems(pairs, l.experts, tile); + const up_in: ex.Rows = .{ .x = try r.dev(std.mem.sliceAsBytes(try r.bfs(rows * l.dims, 4))), .stride = l.dims, .slots = slots }; + const down_in: ex.Rows = .{ .x = try r.dev(std.mem.sliceAsBytes(try r.bfs(pairs * l.width, 4))), .stride = l.width }; + for (ks[pi], 0..) |maybe, ki| if (maybe) |k| for ([_]f32{ 0, 10, 2 }) |lim| for ([_]bool{ false, true }) |after| { + if (ki != 0 and (lim != 0 or after)) continue; + var lc = l; + lc.limit = if (ki == 0) lim else 0; + lc.after = after; + lc.relu2 = ki == 3; + const gate = ki == 0 or ki == 3; + const epi: ex.Epi = if (gate) .up else if (ki == 1) .down_f32 else .down_bf16; + const n_out = if (gate) l.width else l.dims; + const len = pairs * n_out * @as(usize, if (ki == 1) 4 else 2); + const outs = [2]u64{ try r.zeros(len), try r.zeros(len) }; + const in = if (gate) up_in else down_in; + if (prompt) try kit.grouped().prompt(r.s, epi, in, lc, p, outs[0], items, -1) else try kit.grouped().decode(r.s, epi, in, lc, p, outs[0], items, -1); + var lo = lc; + lo.after = false; + const nb = n_out / 32; + const units = if (!prompt) items * nb else if (fp8) items * (nb / @as(usize, if (gate) 2 else 4)) else items * ((nb + @as(usize, if (ki == 0) 4 else 8) - 1) / @as(usize, if (ki == 0) 4 else 8)); + try launchOld(&r, k, in, lo, gate, p, outs[1], units); + const got = try r.back(outs[0], len); + const want = try r.back(outs[1], len); + if (!after or lim == 0) { + try check.sameBytes("grouped experts vs the old fatbin", got, want); + n += 1; + } else if (lim == 10) { + try check.sameBytes("Ling's order at L 10 vs the old order", got, want); + same_at_10 += 1; + } else try check.expect(!std.mem.eql(u8, got, want), "Ling's order at L 2 differs from the old order", .{}); + }; + } + check.pass("ling-clamp-old {t}: {d} calls of the kept instances equal {s}'s bytes; Ling's order at L 10 equals them {d} times, at L 2 differs", .{ l.format, n, fatbins, same_at_10 }); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 9ecf832425..b3d13c7f1a 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -114,6 +114,8 @@ const usage = \\ ling-weights every layer of a Ling checkpoint loaded: projections vs float64, small tensors byte for byte \\ ling-forward-synth [picks] a ling-synth checkpoint's forward vs oracle/ling_forward.py; cuts, windows, batches \\ ling-layer-real one real Ling-3.1 layer vs HF's captured output (oracle/ling_layer.py); exactness + \\ ling-clamp Ling's SwiGLU order min(silu(g), L): glue and the dir's grouped experts vs float64 + \\ ling-clamp-old the grouped expert instances vs an older build's fatbins, byte for byte \\ ling-weights-real real Ling-3.1 layers 0, 4, 7 (4 experts) vs float64 and torch (oracle/ling_weights.py) \\ ; @@ -224,6 +226,8 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "ling-weights")) return ling_weights.run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "ling-forward-synth")) return @import("ling_forward_synth.zig").run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) rest[2] else null); if (std.mem.eql(u8, cmd, "ling-layer-real")) return @import("ling_layer_real.zig").run(gpu, try arg(rest, 0), try arg(rest, 1)); + if (std.mem.eql(u8, cmd, "ling-clamp")) return @import("ling_clamp.zig").run(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "ling-clamp-old")) return @import("ling_clamp.zig").old(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "ling-weights-real")) return ling_real.run(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "glm-lanes-concurrent")) return glm_lanes_tests.concurrentCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); std.debug.print("{s}", .{usage}); diff --git a/zig/tests/cuda/oracle/ling_forward.py b/zig/tests/cuda/oracle/ling_forward.py index f63ca768ce..5eb34d1915 100644 --- a/zig/tests/cuda/oracle/ling_forward.py +++ b/zig/tests/cuda/oracle/ling_forward.py @@ -115,9 +115,10 @@ def rot(z): # GPT-J pairs (rope_interleave), out in halves as HF's interleave h def mlp(self, x, p, limit): g, u = self.lin(x, p + "gate_proj"), self.lin(x, p + "up_proj") - if limit: - g, u = g.clamp(max=limit), u.clamp(-limit, limit) - return self.lin(self.r(self.r(g * torch.sigmoid(g)) * u), p + "down_proj") + s = self.r(g * torch.sigmoid(g)) + if limit: # inclusionAI's SwigluStepAndMul: the limit on silu's output, the up clipped both ways + s, u = s.clamp(max=limit), u.clamp(-limit, limit) + return self.lin(self.r(s * u), p + "down_proj") def moe(self, x, p, i): c, r, ck = self.c, self.r, self.ck From 2a5e1b65a062741697fabc192f26a5cb912e9eb5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 15 Oct 2026 12:12:47 +0200 Subject: [PATCH 09/34] tests: zig fmt ling_mla_real.zig --- zig/tests/cuda/ling_mla_real.zig | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/zig/tests/cuda/ling_mla_real.zig b/zig/tests/cuda/ling_mla_real.zig index baa0994999..5ce184cd44 100644 --- a/zig/tests/cuda/ling_mla_real.zig +++ b/zig/tests/cuda/ling_mla_real.zig @@ -109,10 +109,16 @@ pub fn run(gpu: Gpu, path: []const u8) !void { for (ours, freq) |o, f| fd = @max(fd, @abs(@as(f64, o) - f) / f); std.debug.print("RESULT rope inv_freq: ropeFreqs vs HF's fp32 table, worst relative {e:.3}\n", .{fd}); const l: Layer = .{ - .t = t, .h = h, .qk = qk, .kw = kw, - .q = try r.dev(std.mem.sliceAsBytes(q)), .lat = try r.dev(std.mem.sliceAsBytes(lat)), - .wk = try r.dev(std.mem.sliceAsBytes(wk)), .wv = try r.dev(std.mem.sliceAsBytes(wv)), - .g = try r.dev(std.mem.sliceAsBytes(g)), .freq = try r.dev(std.mem.sliceAsBytes(freq)), + .t = t, + .h = h, + .qk = qk, + .kw = kw, + .q = try r.dev(std.mem.sliceAsBytes(q)), + .lat = try r.dev(std.mem.sliceAsBytes(lat)), + .wk = try r.dev(std.mem.sliceAsBytes(wk)), + .wv = try r.dev(std.mem.sliceAsBytes(wv)), + .g = try r.dev(std.mem.sliceAsBytes(g)), + .freq = try r.dev(std.mem.sliceAsBytes(freq)), }; const full = try work(&r, l); try pass(&r, set, l, full, 0, t); From a3f14d30117cf597522b0515121beba2390d82a8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 15 Oct 2026 16:07:30 +0200 Subject: [PATCH 10/34] ling: served through the lane core at TP=1: engine, lane backend, server hooks and kept prompt states Engine: a sequence per slot, prompt chunks cut at any row (marks included), verify windows of up to 16 rows, shared windows of up to 8 streams, greedy (tf_glm_top1) and sampled (sample.cu) draws, kept states as arena copies. Lanes: GLM's cuda_lanes.zig without the ranks. native.zig registers bailing_hybrid with the CUDA server. ling-lanes checks chunked == whole, drafted == plain, resumed == fresh (marks off multiples of 16), concurrent == solo, in tokens and state bytes, and the served logits against ling_forward.py's float64 on the same picks. --- docs/recipes/ling-3.1.md | 56 ++++- zig/build/cuda.zig | 20 +- zig/src/cli/hub.zig | 1 + zig/src/families/ling/cuda.zig | 6 + zig/src/families/ling/engine.zig | 344 ++++++++++++++++++++++++++ zig/src/families/ling/forward.zig | 13 +- zig/src/families/ling/lanes.zig | 255 +++++++++++++++++++ zig/src/families/ling/native.zig | 115 +++++++++ zig/src/native/cuda.zig | 3 +- zig/tests/cuda/ling_forward_synth.zig | 4 +- zig/tests/cuda/ling_lanes.zig | 178 +++++++++++++ zig/tests/cuda/ling_lanes_bytes.zig | 198 +++++++++++++++ zig/tests/cuda/ling_lanes_rig.zig | 218 ++++++++++++++++ zig/tests/cuda/main.zig | 2 + zig/tests/cuda/oracle/ling_forward.py | 39 ++- 15 files changed, 1430 insertions(+), 22 deletions(-) create mode 100644 zig/src/families/ling/engine.zig create mode 100644 zig/src/families/ling/lanes.zig create mode 100644 zig/src/families/ling/native.zig create mode 100644 zig/tests/cuda/ling_lanes.zig create mode 100644 zig/tests/cuda/ling_lanes_bytes.zig create mode 100644 zig/tests/cuda/ling_lanes_rig.zig diff --git a/docs/recipes/ling-3.1.md b/docs/recipes/ling-3.1.md index 3b4b6da25f..4b354da37f 100644 --- a/docs/recipes/ling-3.1.md +++ b/docs/recipes/ling-3.1.md @@ -1,6 +1,8 @@ # Ling-3.1-flash -Work in progress on the Zig engine's CUDA side: the model is **not served yet**. This page lists what exists, what is +Work in progress on the Zig engine's CUDA side. The family is registered with the native CUDA server +(`model_type` `bailing_hybrid`) and serves requests through the lane core at TP=1, but the whole model does not fit +one GB10, so serving is verified on synthetic checkpoints and a thin real model. This page lists what exists, what is missing, and how to test without the 572 GB checkpoint. Everything below was measured on one GB10. Ling-3.1-flash (`BailingMoeV3ForCausalLM`, `bailing_hybrid`) has 72 layers plus one MTP layer: Kimi delta attention @@ -19,6 +21,7 @@ block FP8 (e4m3 in 128 x 128 blocks with `weight_scale_inv`). | Weight loader at TP=1: bf16, FP8 block, ModelOpt NVFP4 | `zig/src/families/ling/weights.zig`, `weights_proj.zig` | `ling-weights`, `ling-weights-real` | | Synthetic checkpoints in all three formats | `tf-cuda-test ling-synth` | its own name check | | The TP=1 layer graph: embedding, KDA/MLA mixers, dense MLP and MoE, final norm, head; windows and commits | `zig/src/families/ling/forward.zig`, `graph.zig`, `state.zig`, `block_*.zig`, `kernels/cuda/ling_glue.cu` | `ling-forward-synth`, `ling-layer-real` | +| The engine and lane backend: a sequence per slot, prompt chunks cut anywhere, verify windows of up to 16 rows, shared windows of up to 8 streams (64 rows), kept prompt states, the server hooks | `engine.zig`, `lanes.zig`, `native.zig`; `native/cuda.zig` registers it | `ling-lanes` | | SwiGLU clamped in Ling's order, `min(silu(g), L) * clip(u, -L, L)`, in the glue and the FP8/NVFP4 grouped experts (every older kernel's SASS unchanged) | `glm_glue_moe.cuh` (`tf_ling_swiglu_*`), `experts.cuh` (epilogue 4), `experts.Layer.after` | `ling-clamp`, `ling-clamp-old`, `glm-route-old` | ## The layer graph @@ -116,10 +119,51 @@ this one print the same digits on the real layers (71: rel L2 4.125e-3, 3.747e-3 float64 references do not change either (the synthetic clamped layer never reaches g > 10). HF's clamp itself moves layer 71's output by 3.7e-3 (rel L2 against the unclamped layer) and layer 65's by 0 on these 512 tokens. +## Serving + +`engine.Engine` owns the weights (every backbone layer, or `Options.layers` for a thin model), the graph, three +`Buffers` (prompt chunks, one stream's window, shared windows) and one `State` per slot. `Lanes` is the lane backend +(GLM's `cuda_lanes.zig` without the ranks): `prefill`/`prefill_step` (a chunk at a time, other streams' rounds +between), `first`, `verify` (one stream's window, or every stream's in one forward through `verifyShared`), `keep` +(a window's prefix, the KDA states replayed), `release`. Greedy draws take GLM's `tf_glm_top1`, sampled ones +upstream's `sample.cu` through GLM's sampler at TP=1. `native.zig` gives the server `open`, the facts (exact windows of +16 rows, shared rounds unless `TF_LING_SHARED=0`; `TF_LING_PROMPT_STEPS=0` passes prompts whole) and the prompt +cache's `snap_*` hooks: a kept state is the sequence's arena copied whole at a mark, and the lane backend restores it +inside its own prompt pass. Ling's prompt arithmetic does not depend on where a chunk ends, so a mark may fall on any +row (no `snap_grid`): the pass cuts its chunk there. There is no recorder hook (`features`): the drafter recorder +(#848) is not on this branch. + +`ling-lanes DIR OUT [REF] [--layers a,b]` drives the backend under upstream's round loop as the server's lane host +does, prompts in 37-row chunks; every check compares tokens, and the engine-level ones bytes: + +- chunked == whole: chunks of 512, 37 and 16 rows and the whole `prefill` give the same reply; the states after + 37- and 16-row chunks equal the whole pass's byte for byte (KDA states, conv windows, MLA rows). +- drafted == plain: drafts from the plain reply (2 to 15 a round, most rounds one of them wrong, so windows of 3 to + 16 rows are cut anywhere), greedy and sampled; a 9-row window kept to 4 then a 5-row window equal 9 single rows in + tokens, logits and state bytes. +- resumed == fresh: states kept at marks 37, 101, 150, 263, 349 (off the chunk grid, not multiples of 16), each + resumed pass equals the fresh one, and a different tail after a kept 263; resumed states equal the whole pass's. +- concurrent == solo: nine cases of three streams (at once, staggered, reversed, the same prompt twice, sampled, + one discarded, prompts a chunk at a time between the others' rounds, drafted or not) equal their solo replies; + shared windows of 2 streams over 8 rounds (1 to 16 rows each) equal lone twins in logits, tokens and states. +- cancels in the prompt pass, discards between rounds, eos and `max_new` 1-3. +- the reply equals the serial forward's argmax at every row; OUT.ids and OUT.picks feed + `ling_forward.py --ids OUT.ids --picks OUT.picks`, and with REF the served logits are checked against float64: + +| Format | served vs float64 (same picks) | vs the bf16 points | bf16 points vs float64 | argmax vs float64 | +| --- | --- | --- | --- | --- | +| bf16 | 2.24e-2 | 1.27e-2 | 2.25e-2 | 301/309 | +| FP8 block | 2.09e-2 | 1.18e-2 | 2.10e-2 | 302/309 | +| NVFP4 | 2.21e-2 | 1.38e-2 | 2.22e-2 | 294/309 | + +92 checks pass in each format. On the real FP8 release, `--layers 0,1,2,3` (the real embedding, layers 0-3, final +norm and head; KDA and dense MLP only, 6 GB) passes the same 90 exactness checks; it has no float64 reference (an +MLA + MoE layer adds 8 GB). Mutations caught: committing every row on `keep`, a restore that keeps the fresh KDA copy +bits, chunks not cut at marks, shared draws keyed one position early, a lane commit that keeps one row too many. + ## What is missing -- Lanes hooks (drafted, concurrent and cut-prompt windows) for Ling: the graph has the pieces (segments, windows, - `commit(keep)`), nothing schedules them yet. Sampling, the server and a full-checkpoint run are not wired. +- A full-checkpoint run: 574 GB in FP8 needs five Sparks and tensor parallelism. - MTP drafting: the MTP layer's tensors load (`weights.Mtp`); the graph leaves it unwired. It would be `enorm(embed(next)) | hnorm(hidden)` through `eh_proj`, one MLA + MoE layer on its own cache, the final norm and the shared head. @@ -140,6 +184,9 @@ python zig/tests/cuda/oracle/ling_forward.py /tmp/ling-fp8 /tmp/ref zig-out/bin/tf-cuda-test ling-forward-synth /tmp/ling-fp8 /tmp/ref /tmp/picks # exactness; writes the picks python zig/tests/cuda/oracle/ling_forward.py /tmp/ling-fp8 /tmp/fref --picks /tmp/picks zig-out/bin/tf-cuda-test ling-forward-synth /tmp/ling-fp8 /tmp/fref # same picks: logits bounded +zig-out/bin/tf-cuda-test ling-lanes /tmp/ling-fp8 /tmp/served # serving checks; writes .ids/.picks +python zig/tests/cuda/oracle/ling_forward.py /tmp/ling-fp8 /tmp/sref --ids /tmp/served.ids --picks /tmp/served.picks +zig-out/bin/tf-cuda-test ling-lanes /tmp/ling-fp8 /tmp/served /tmp/sref # and the served logits vs float64 zig-out/bin/tf-cuda-test ling-clamp /tmp/ling-fp8 # Ling's SwiGLU order vs float64, L 10 and 2 zig-out/bin/tf-cuda-test ling-clamp-old /tmp/ling-fp8 OLD/fatbin # kept expert instances vs an older build's bytes ``` @@ -191,8 +238,7 @@ size. Each item stands alone. -1. **Lanes hooks.** Drive `Forward` from #804's `cuda_lanes.zig` pattern: drafted windows, concurrent streams as - segments, cut prompts; `ling-forward-synth` already checks the exactness each needs. +1. **Recorder.** Give `Lanes` a `features` hook (the final-norm rows `Engine.verify` already keeps) once #848 lands. 2. **MTP.** Wire `weights.Mtp` as a drafter (see "What is missing"). 3. **Fused projections.** Same-format q|k|v (and gate|up) rows as one weight per format, without the column copies. 4. **Bf16 routed experts.** Add a grouped bf16 expert kernel, or quantize at load, and drop the per-expert lanes. diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index b263d9b4b4..4b1c865ad3 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -156,9 +156,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu for (exlr8_build.denseModules(b, target, optimize, exlr8, exlr8_tests[0])) |m| runner.addImport(m.name, m.module); const glm_cuda = glmCuda(b, target, optimize, cuda, mods.core, mods.lanes, mods.heat); runner.addImport("glm_cuda", glm_cuda); - const ling = b.createModule(.{ .root_source_file = b.path("zig/src/families/ling/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true }); - for ([_]struct { []const u8, *std.Build.Module }{ .{ "glm_cuda", glm_cuda }, .{ "cuda", cuda }, .{ "core", mods.core } }) |m| ling.addImport(m[0], m[1]); - runner.addImport("ling_cuda", ling); + runner.addImport("ling_cuda", lingCuda(b, target, optimize, cuda, mods.core, mods.lanes, glm_cuda)); const runner_exe = b.addExecutable(.{ .name = "tf-cuda-test", .root_module = runner }); b.installArtifact(runner_exe); const controls = if (nvcc) |tool| fatbin(b, tool, version.?, .{ .name = "contract_controls", .flags = &.{"-O3"} }, sms) else null; @@ -188,9 +186,17 @@ fn radioModule(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu return b.createModule(.{ .root_source_file = b.path("zig/src/families/radio/radio.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core }, .{ .name = "lanes", .module = lanes } } }); } +/// The Ling-3.1-flash CUDA family over the runtime, the core, the lane core and the GLM kernels it reuses. +fn lingCuda(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, glm: *std.Build.Module) *std.Build.Module { + const m = b.createModule(.{ .root_source_file = b.path("zig/src/families/ling/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true }); + for ([_]struct { []const u8, *std.Build.Module }{ .{ "glm_cuda", glm }, .{ "cuda", cuda }, .{ "core", core }, .{ "lanes", lanes } }) |x| m.addImport(x[0], x[1]); + return m; +} + /// The CUDA engines a native server opens (native/cuda.zig), over the given runtime and families. fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, heat: *std.Build.Module) struct { api: *std.Build.Module, engines: *std.Build.Module } { const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} }); + const glm = glmCuda(b, target, optimize, cuda, core, lanes, heat); const mod = b.createModule(.{ .root_source_file = b.path("zig/src/native/cuda.zig"), .target = target, @@ -203,7 +209,8 @@ fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builti .{ .name = "nemotron", .module = nemotron }, .{ .name = "lfm2", .module = lfm2Module(b, target, optimize, cuda, core) }, .{ .name = "core", .module = core }, - .{ .name = "glm_cuda", .module = glmCuda(b, target, optimize, cuda, core, lanes, heat) }, + .{ .name = "glm_cuda", .module = glm }, + .{ .name = "ling_cuda", .module = lingCuda(b, target, optimize, cuda, core, lanes, glm) }, }, }); return .{ .api = api, .engines = mod }; @@ -278,10 +285,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st const glm = glmCuda(b, host, .debug, cuda, mods.core, mods.lanes, mods.heat); step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = glm })).step); // Ling-3.1-flash's config loader and weight map, and its routing and limits as the GLM kernels take them - const ling = b.createModule(.{ .root_source_file = b.path("zig/src/families/ling/cuda.zig"), .target = host, .optimize = .debug, .link_libc = true }); - ling.addImport("glm_cuda", glm); - ling.addImport("cuda", cuda); - ling.addImport("core", mods.core); + const ling = lingCuda(b, host, .debug, cuda, mods.core, mods.lanes, glm); const ling_run = b.addRunArtifact(b.addTest(.{ .root_module = ling })); step.dependOn(&ling_run.step); b.step("test-ling", "Ling-3.1-flash's config loader and weight map (no GPU work)").dependOn(&ling_run.step); diff --git a/zig/src/cli/hub.zig b/zig/src/cli/hub.zig index 99d6835a49..318dd88eb6 100644 --- a/zig/src/cli/hub.zig +++ b/zig/src/cli/hub.zig @@ -15,6 +15,7 @@ const titles = [_]struct { model_type: []const u8, title: []const u8 }{ .{ .model_type = "glm5_next", .title = "GLM-5.3-Flash" }, .{ .model_type = "qwen3_5", .title = "Qwen3.5-2B or Qwen3.8-27B" }, .{ .model_type = "lfm2_moe", .title = "LFM2.5-24B-A2B" }, + .{ .model_type = "bailing_hybrid", .title = "Ling-3.1-flash" }, }; /// The family that serves `model_type` from the engine registry, or null when Zig cannot serve it (the 0.6 line may). diff --git a/zig/src/families/ling/cuda.zig b/zig/src/families/ling/cuda.zig index 095810d32a..f578ca89da 100644 --- a/zig/src/families/ling/cuda.zig +++ b/zig/src/families/ling/cuda.zig @@ -8,6 +8,9 @@ pub const synth = @import("synth.zig"); pub const graph = @import("graph.zig"); pub const state = @import("state.zig"); pub const forward = @import("forward.zig"); +pub const engine = @import("engine.zig"); +pub const lanes = @import("lanes.zig"); +pub const native = @import("native.zig"); const Config = config.Config; @@ -44,6 +47,9 @@ test { std.testing.refAllDecls(@import("weights_proj.zig")); std.testing.refAllDecls(forward); // the layer graph compiles on hosts without a GPU std.testing.refAllDecls(forward.Forward); + std.testing.refAllDecls(engine.Engine); // the engine, its lane backend and its server hooks too + std.testing.refAllDecls(lanes.Lanes); + _ = native; } test "Ling's routing as glue.select takes it: 8 groups keeping 4, and select's own bounds hold" { diff --git a/zig/src/families/ling/engine.zig b/zig/src/families/ling/engine.zig new file mode 100644 index 0000000000..f97f160371 --- /dev/null +++ b/zig/src/families/ling/engine.zig @@ -0,0 +1,344 @@ +//! Ling-3.1-flash's TP=1 CUDA engine: a sequence per slot, prompt chunks cut anywhere, verify windows, kept states. +const std = @import("std"); +const cuda = @import("cuda"); +const glm = @import("glm_cuda"); +const config = @import("config.zig"); +const weights = @import("weights.zig"); +const fw = @import("forward.zig"); + +const sampler = glm.sampler; +pub const Sampling = sampler.Sampling; + +/// A prompt pass's chunk rows by default (Options.prompt_rows); marks cut chunks anywhere besides. +pub const prompt_rows = 1024; +/// A verify window's most rows: the pending token and up to 15 drafts (forward.window_rows). +pub const window_rows = fw.window_rows; +/// A shared verify's most rows and streams. +pub const shared_rows = 64; +pub const max_shared = 8; + +pub const Options = struct { + capacity: usize, // tokens a sequence holds + prompt_rows: usize = prompt_rows, // rows a prompt chunk takes at most + layers: ?[]const u32 = null, // a thin model's backbone layers, ascending (checks on a real checkpoint); null: all +}; + +/// Rank 0's word on its request (cancelled or not), asked before each chunk. +pub const Cancel = struct { ptr: ?*anyopaque, check: *const fn (?*anyopaque) bool }; +/// One stream's window of a shared verify: its slot, its ids, and its drawn tokens. +pub const Part = struct { slot: u32, ids: []const u32, out: []u32 }; +/// A kept prompt state as the prompt cache holds it: its id and position. +pub const Kept = struct { id: u64, at: u32 }; +/// Where a prompt pass resumes: kept state `id` at `at`. +pub const Resume = struct { id: u64, at: usize }; +/// The pass's marks: it cuts a chunk at each and calls `call` standing there. +pub const Marks = struct { at: []const u32, ctx: *anyopaque, call: *const fn (ctx: *anyopaque, at: u32) void }; +/// One chunk's outcome: where the pass stands, and the first token once the prompt is done. +pub const Step = struct { at: usize, first: ?u32 }; + +/// A stream's sequence: the layers' caches and its sampling rule. +const Seq = struct { + st: fw.State, + sampling: ?Sampling = null, // null: greedy +}; + +/// A kept state: the sequence's arena copied whole, and the host fields beside it. +const Snap = struct { buf: cuda.DeviceBuffer, pos: usize, cur: []u1 }; + +pub const Engine = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + c: config.Config, + m: fw.Modules, + stream: cuda.Stream, + w: weights.Weights, // every loaded layer + main: weights.Weights, // a view of w: the backbone layers alone + g: fw.Graph, + bp: fw.Buffers, // prompt chunks + bd: fw.Buffers, // one stream's verify window + bs: fw.Buffers, // shared verify windows + smp: sampler.Sampler, + pairs: cuda.DeviceBuffer, // (value, id) a row of greedy draws + seqs: std.ArrayList(?*Seq) = .empty, // by slot; slot 0 made at init + cur: *Seq = undefined, + slot: u32 = 0, + capacity: usize, + chunk: usize, + snaps: std.AutoHashMapUnmanaged(u64, Snap) = .empty, + next_snap: u64 = 1, + + /// The checkpoint at `dir` in its own format. + pub fn init(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, ctx: *const cuda.Context, dir: []const u8, o: Options) !*Engine { + if (o.prompt_rows == 0 or o.capacity < window_rows + 2) return error.BadOptions; + var why: config.Why = .{}; + var opened = weights.open(gpa, io, dir, &why) catch |err| { + std.log.err("{s}: {s}", .{ dir, why.text() }); + return err; + }; + defer opened.ck.close(); + const c = opened.c; + const e = try gpa.create(Engine); + errdefer gpa.destroy(e); + e.* = .{ .gpa = gpa, .d = d, .c = c, .m = undefined, .stream = undefined, .w = undefined, .main = undefined, .g = undefined, .bp = undefined, .bd = undefined, .bs = undefined, .smp = undefined, .pairs = undefined, .capacity = o.capacity, .chunk = o.prompt_rows }; + e.m = try fw.Modules.load(d, ctx); + errdefer e.m.deinit(); + e.stream = try cuda.Stream.init(d, false); + errdefer e.stream.deinit(); + const layers = try backbone(gpa, &c, o.layers); + defer gpa.free(layers); + { + var pack = try cuda.modelopt.Packer.load(d); + defer pack.unload(); + e.w = try weights.load(gpa, d, e.stream, &pack, &opened.ck, c, layers, .{}); + } + errdefer e.w.deinit(); + e.main = e.w; + e.main.layers = e.w.layers[0..layers.len]; + e.g = try fw.Graph.init(gpa, d, &e.m, e.stream, &e.main); + errdefer e.g.deinit(); + e.bp = try fw.Buffers.init(d, &e.m, &e.w, o.prompt_rows, true, o.capacity); + errdefer e.bp.deinit(); + e.bd = try fw.Buffers.init(d, &e.m, &e.w, window_rows, false, o.capacity); + errdefer e.bd.deinit(); + e.bs = try fw.Buffers.init(d, &e.m, &e.w, shared_rows, false, o.capacity); + errdefer e.bs.deinit(); + e.smp = try sampler.Sampler.init(d, c.vocab, window_rows); + errdefer e.smp.deinit(); + e.pairs = try cuda.DeviceBuffer.alloc(d, 8 * window_rows); + errdefer e.pairs.free(); + try e.bind(0); + try e.stream.synchronize(); + return e; + } + + /// The backbone layers to load: `thin` (ascending, each below c.layers), or all of them. + fn backbone(gpa: std.mem.Allocator, c: *const config.Config, thin: ?[]const u32) ![]u32 { + if (thin) |t| { + for (t, 0..) |i, k| if (i >= c.layers or (k > 0 and i <= t[k - 1])) return error.BadLayers; + return gpa.dupe(u32, t); + } + const all = try gpa.alloc(u32, c.layers); + for (all, 0..) |*l, i| l.* = @intCast(i); + return all; + } + + pub fn deinit(e: *Engine) void { + e.stream.synchronize() catch {}; + var it = e.snaps.valueIterator(); + while (it.next()) |s| e.freeSnap(s); + e.snaps.deinit(e.gpa); + for (0..e.seqs.items.len) |i| e.free(i); + e.seqs.deinit(e.gpa); + e.pairs.free(); + e.smp.deinit(); + e.bs.deinit(); + e.bd.deinit(); + e.bp.deinit(); + e.g.deinit(); + e.w.deinit(); + e.stream.deinit(); + e.m.deinit(); + const gpa = e.gpa; + gpa.destroy(e); + } + + pub fn forward(e: *const Engine, b: *const fw.Buffers) fw.Forward { + return .{ .m = &e.m, .g = &e.g, .s = e.stream, .b = b }; + } + + /// Binds slot `slot`'s sequence for the steps that follow, making it (zeroed) on first use. + pub fn bind(e: *Engine, slot: u32) !void { + if (slot >= e.seqs.items.len) try e.seqs.appendNTimes(e.gpa, null, slot + 1 - e.seqs.items.len); + if (e.seqs.items[slot] == null) { + const q = try e.gpa.create(Seq); + errdefer e.gpa.destroy(q); + q.* = .{ .st = try fw.State.init(e.gpa, e.d, &e.g, e.capacity) }; + e.seqs.items[slot] = q; + } + e.cur = e.seqs.items[slot].?; + e.slot = slot; + } + + /// Frees slot `slot`'s sequence (slot 0 stays) and binds slot 0. + pub fn release(e: *Engine, slot: u32) void { + if (slot != 0) e.free(slot); + e.cur = e.seqs.items[0].?; + e.slot = 0; + } + + fn free(e: *Engine, slot: usize) void { + const q = e.seqs.items[slot] orelse return; + e.stream.synchronize() catch {}; + q.st.deinit(); + e.gpa.destroy(q); + e.seqs.items[slot] = null; + } + + /// Bytes one more stream's sequence takes: its arena. + pub fn seqBytes(e: *const Engine) u64 { + return e.seqs.items[0].?.st.arena.buf.len; + } + + /// The committed length: the pending token sits here, a window's row r draws position pos + 1 + r. + pub fn pos(e: *const Engine) usize { + return e.cur.st.pos; + } + + /// A fresh sequence with `sampling` (null or temperature 0: greedy). + pub fn begin(e: *Engine, sampling: ?Sampling) !void { + e.cur.sampling = sampler.check(sampling); + try e.cur.st.reset(e.stream); + } + + /// The whole prompt pass: chunks, then the first token from the last row; `from` restores a kept state. + pub fn prefill(e: *Engine, ids: []const u32, from: ?Resume, cancel: ?Cancel, marks: ?Marks) !u32 { + var at = try e.prefillStart(ids.len, from); + while (true) { + const s = try e.prefillChunk(ids, at, cancel, marks); + if (s.first) |t| return t; + at = s.at; + } + } + + /// Where a pass over `len` ids starts: 0, or `from`'s position with its kept state restored. + pub fn prefillStart(e: *Engine, len: usize, from: ?Resume) !usize { + if (len == 0 or len + window_rows >= e.capacity) return error.PromptTooLong; + const r = from orelse return 0; + if (r.at == 0 or r.at >= len) return error.BadResume; + try e.restoreState(r.id); + if (e.cur.st.pos != r.at) return error.BadResume; + return r.at; + } + + /// One chunk of `ids` from `at` (the sequence's position): up to the chunk rows, cut at the next mark. + pub fn prefillChunk(e: *Engine, ids: []const u32, at: usize, cancel: ?Cancel, marks: ?Marks) !Step { + const len = ids.len; + if (at >= len or at != e.cur.st.pos) return error.BadChunk; + if (marks) |m| if (at > 0 and std.mem.indexOfScalar(u32, m.at, @intCast(at)) != null) m.call(m.ctx, @intCast(at)); + if (cancel) |c| if (c.check(c.ptr)) return error.Cancelled; + var end = @min(len, at + e.chunk); + if (marks) |m| for (m.at) |x| if (x > at and x < end) { + end = x; + break; + }; + const rows = end - at; + const f = e.forward(&e.bp); + const seg = [_]fw.Seg{.{ .st = &e.cur.st, .row0 = 0, .rows = rows }}; + try f.embed(@ptrCast(ids[at..end])); + try f.stack(&seg); + try f.final(rows); + var first: ?u32 = null; + if (end == len) { + try f.logitsOf(rows - 1, 1); + var t: [1]u32 = undefined; + try e.draw(e.bp.logits, 1, len, &t); + first = t[0]; + } + try f.commit(&e.cur.st, rows); + return .{ .at = end, .first = first }; + } + + /// A verify window: `ids` (pending token, then drafts) at the committed position; `out[r]` drawn at pos + 1 + r. + pub fn verify(e: *Engine, ids: []const u32, out: []u32) !void { + if (ids.len == 0 or ids.len > window_rows) return error.WindowTooWide; + const f = e.forward(&e.bd); + const at = e.cur.st.pos; + try f.run(&.{.{ .st = &e.cur.st, .row0 = 0, .rows = ids.len }}, @ptrCast(ids)); + try e.draw(e.bd.logits, ids.len, at + 1, out); + } + + /// Several streams' windows in one forward, each stream's rows on its own sequence. + pub fn verifyShared(e: *Engine, parts: []const Part) !void { + if (parts.len == 0 or parts.len > max_shared) return error.WindowTooWide; + var segs: [max_shared]fw.Seg = undefined; + var ids: [shared_rows]u32 = undefined; + var rows: usize = 0; + for (parts, 0..) |p, i| { + if (p.ids.len == 0 or p.ids.len > window_rows or rows + p.ids.len > shared_rows) return error.WindowTooWide; + for (parts[0..i]) |q| if (q.slot == p.slot) return error.SlotTwice; + try e.bind(p.slot); + segs[i] = .{ .st = &e.cur.st, .row0 = rows, .rows = p.ids.len }; + @memcpy(ids[rows..][0..p.ids.len], p.ids); + rows += p.ids.len; + } + try e.forward(&e.bs).run(segs[0..parts.len], @ptrCast(ids[0..rows])); + for (parts, segs[0..parts.len]) |p, sg| { + try e.bind(p.slot); + try e.draw(e.bs.logits + sg.row0 * e.c.vocab * 2, sg.rows, sg.st.pos + 1, p.out); + } + } + + /// The last window's first `n` rows committed (the KDA states replayed when fewer than it ran). + pub fn keep(e: *Engine, n: usize) !void { + try e.forward(&e.bd).commit(&e.cur.st, n); + } + + /// `rows` logits rows -> the tokens at positions first, first + 1, ... + pub fn draw(e: *Engine, logits: u64, rows: usize, first: u64, out: []u32) !void { + if (e.cur.sampling) |s| { + const span = [_]glm.shard.Span{.{ .lo = 0, .hi = e.c.vocab }}; + return e.smp.draw(e.stream, .local, &span, &.{.{ .rank = 0, .logits = logits }}, rows, first, sampler.rule(s), out); + } + const V: usize = e.c.vocab; + for (0..rows) |r| { + var a: cuda.Args = .{}; + inline for (.{ logits + r * V * 2, @as(i64, @intCast(V)), @as(u32, 0), e.pairs.ptr + r * 8 }) |v| a.add(v); + try cuda.launch.launch(e.m.glue.top1, .{ .grid = .{ .x = 1 }, .block = .{ .x = 1024 } }, e.stream, &a); + } + try e.stream.synchronize(); + var got: [window_rows][2]u32 = undefined; + try e.pairs.download(0, std.mem.sliceAsBytes(got[0..rows])); + for (out[0..rows], got[0..rows]) |*o, p| o.* = p[1]; + } + + /// Bytes one kept state takes: the sequence's arenas. + pub fn stateBytes(e: *const Engine) u64 { + return e.seqBytes(); + } + + /// A copy of the bound sequence's state under a new id (the pass standing at a mark). + pub fn keepState(e: *Engine) !u64 { + const q = e.cur; + if (q.st.pending != 0) return error.NotCommitted; + var buf = try cuda.DeviceBuffer.alloc(e.d, q.st.arena.buf.len); + errdefer buf.free(); + try e.copy(buf.ptr, q.st.arena.buf.ptr, buf.len); + try e.stream.synchronize(); + const cur = try e.gpa.dupe(u1, q.st.cur); + errdefer e.gpa.free(cur); + const id = e.next_snap; + try e.snaps.put(e.gpa, id, .{ .buf = buf, .pos = q.st.pos, .cur = cur }); + e.next_snap += 1; + return id; + } + + /// The bound sequence made kept state `id`'s. + fn restoreState(e: *Engine, id: u64) !void { + const s = e.snaps.getPtr(id) orelse return error.NoKeptState; + const q = e.cur; + try e.copy(q.st.arena.buf.ptr, s.buf.ptr, s.buf.len); + @memcpy(q.st.cur, s.cur); + q.st.pending = 0; + q.st.prompt = false; + try q.st.setPos(e.stream, s.pos); + try e.stream.synchronize(); + } + + pub fn holds(e: *const Engine, id: u64) bool { + return e.snaps.contains(id); + } + + pub fn dropState(e: *Engine, id: u64) void { + var kv = e.snaps.fetchRemove(id) orelse return; + e.freeSnap(&kv.value); + } + + fn freeSnap(e: *Engine, s: *Snap) void { + s.buf.free(); + e.gpa.free(s.cur); + } + + pub fn copy(e: *const Engine, dst: u64, src: u64, bytes: usize) !void { + try e.d.check(e.d.api.cuMemcpyDtoDAsync_v2(dst, src, bytes, e.stream.handle), "cuMemcpyDtoDAsync"); + } +}; diff --git a/zig/src/families/ling/forward.zig b/zig/src/families/ling/forward.zig index a10072b41c..6cb0cdd337 100644 --- a/zig/src/families/ling/forward.zig +++ b/zig/src/families/ling/forward.zig @@ -101,11 +101,22 @@ pub const Forward = struct { /// The final norm and lm_head of b.x's rows into b.logits (bf16 [rows][vocab]). pub fn head(f: *const Forward, rows: usize) !void { + try f.final(rows); + try f.logitsOf(0, rows); + } + + /// The final norm of b.x's rows into b.normed: the rows lm_head and the MTP layer read. + pub fn final(f: *const Forward, rows: usize) !void { const o = f.ops(); const w = f.g.w; const D: usize = w.config.hidden; try glue.rmsnorm(o.set(), o.s, f.b.x, D, w.norm, f.b.normed, D, null, w.config.eps, rows, D); - try o.mm(w.head orelse return error.TiedHeadNotBuilt, f.b.normed, D, rows, f.b.logits); + } + + /// lm_head of b.normed's rows row0 .. row0 + n into b.logits from its first row. + pub fn logitsOf(f: *const Forward, row0: usize, n: usize) !void { + const D: usize = f.g.w.config.hidden; + try f.ops().mm(f.g.w.head orelse return error.TiedHeadNotBuilt, f.b.normed + row0 * D * 2, D, n, f.b.logits); } /// Keep the first `keep` rows of the sequence's last forward (a prompt keeps them all): KDA states, position. diff --git a/zig/src/families/ling/lanes.zig b/zig/src/families/ling/lanes.zig new file mode 100644 index 0000000000..9cab5286fd --- /dev/null +++ b/zig/src/families/ling/lanes.zig @@ -0,0 +1,255 @@ +//! Ling-3.1-flash's CUDA lane backend (as GLM's cuda_lanes.zig, one process): a slot per stream, shared verify. +const std = @import("std"); +const lanes = @import("lanes"); +const engine = @import("engine.zig"); + +const be = lanes.backend; +const Engine = engine.Engine; + +/// First tokens a handle names. +const ring = 1024; + +/// The widest verify window: the pending token and 15 drafts (copies or the MTP layer's chain). +pub const max_rows = engine.window_rows; + +/// A stream's place: its sequence's slot, the rows its last verify left uncommitted, where its prompt pass stands. +const Lane = struct { + slot: u32, + pending_rows: ?usize = null, + fill: ?usize = null, // a prompt pass going a chunk at a time: where its next chunk starts +}; + +pub const Lanes = struct { + e: *Engine, + whole: bool = false, // prompt passes run whole: no other stream's round between chunks + alone: bool = false, // each window in its own forward + shared: u64 = 0, // forwards that held several streams' windows (checks read it) + lanes: std.AutoHashMapUnmanaged(*const lanes.Stream, Lane) = .empty, + drawn: [ring]u32 = undefined, + next: u64 = 0, + + pub fn deinit(self: *Lanes) void { + self.lanes.deinit(self.e.gpa); + } + + pub fn backend(self: *Lanes) be.Backend { + return .{ .ptr = self, .vtable = if (self.whole) &whole_vt else &steps_vt }; + } + + const whole_vt: be.Backend.VTable = .{ + .prefill = prefillFn, + .first = firstFn, + .queue = queueFn, + .read = readFn, + .verify = verifyFn, + .keep = keepFn, + .draft = draftFn, + .release = releaseFn, + }; + const steps_vt: be.Backend.VTable = vt: { + var v = whole_vt; + v.prefill_step = prefillStepFn; + break :vt v; + }; + + /// Round loop facts: exact windows of max_rows (copies), shared windows unless `alone`. + pub fn facts(self: *const Lanes) lanes.Model { + const sharing = !self.alone; + return .{ + .exact_width = max_rows, + .gpu_tokens = false, + .speculate_early = false, + .hidden_rows = sharing, + .batch_rows = if (sharing) engine.shared_rows else 0, + .max_streams = if (sharing) engine.max_shared else 1, + }; + } + + fn of(ptr: *anyopaque) *Lanes { + return @ptrCast(@alignCast(ptr)); + } + + fn take(self: *Lanes, token: u32) u64 { + const h = self.next; + self.drawn[h % ring] = token; + self.next += 1; + return h; + } + + /// The lowest slot no stream holds. + fn freeSlot(self: *const Lanes) u32 { + var slot: u32 = 0; + while (true) : (slot += 1) { + var it = self.lanes.valueIterator(); + const taken = while (it.next()) |l| { + if (l.slot == slot) break true; + } else false; + if (!taken) return slot; + } + } + + /// The stream's sequence bound, its last verify committed (`kept` rows, else all of them). + fn bind(self: *Lanes, s: *const lanes.Stream, kept: ?usize) !*Lane { + const l = self.lanes.getPtr(s) orelse return error.UnknownStream; + try self.e.bind(l.slot); + if (l.pending_rows) |rows| try self.e.keep(kept orelse rows); + l.pending_rows = null; + return l; + } + + fn cancelled(ptr: ?*anyopaque) bool { + return @as(*lanes.Stream, @ptrCast(@alignCast(ptr.?))).isCancelled(); + } + + fn cancelOf(s: *lanes.Stream) engine.Cancel { + return .{ .ptr = s, .check = cancelled }; + } + + fn prefillFn(ptr: *anyopaque, s: *lanes.Stream) anyerror!void { + const self = of(ptr); + const from = try self.opening(s); + _ = self.take(try self.e.prefill(s.prompt(), from, cancelOf(s), marksOf(s))); + } + + /// One chunk of the stream's prompt pass (fillOne runs other rounds between); true once the first token is drawn. + fn prefillStepFn(ptr: *anyopaque, s: *lanes.Stream) anyerror!bool { + const self = of(ptr); + const going: ?usize = if (self.lanes.getPtr(s)) |l| l.fill else null; + const at = going orelse try self.e.prefillStart(s.prompt().len, try self.opening(s)); + const l = self.lanes.getPtr(s).?; + try self.e.bind(l.slot); + errdefer l.fill = null; + const step = try self.e.prefillChunk(s.prompt(), at, cancelOf(s), marksOf(s)); + if (step.first) |first| { + _ = self.take(first); + l.fill = null; + return true; + } + l.fill = step.at; + return false; + } + + /// A prompt pass begins: the stream's slot bound and its sequence fresh; where the pass resumes (a kept state). + fn opening(self: *Lanes, s: *lanes.Stream) !?engine.Resume { + const ids = s.prompt(); + if (ids.len + s.max_new + max_rows >= self.e.capacity) return error.PromptTooLong; + const gop = try self.lanes.getOrPut(self.e.gpa, s); + if (!gop.found_existing) gop.value_ptr.* = .{ .slot = self.freeSlot() }; + gop.value_ptr.pending_rows = null; + gop.value_ptr.fill = null; + try self.e.bind(gop.value_ptr.slot); + try self.e.begin(s.sampling); + s.cached = 0; + return self.resumeFrom(s, ids.len); + } + + fn marksOf(s: *lanes.Stream) ?engine.Marks { + return if (s.reuse.marks.len == 0) null else .{ .at = s.reuse.marks, .ctx = s, .call = keptAt }; + } + + /// The stream's kept state, where its prompt pass starts (null: none, or no longer held). + fn resumeFrom(self: *Lanes, s: *lanes.Stream, len: usize) ?engine.Resume { + const saved = s.reuse.saved orelse return null; + const k: *const engine.Kept = @ptrCast(@alignCast(saved)); + if (k.at != s.reuse.at or k.at == 0 or k.at >= len or !self.e.holds(k.id)) { + s.reuse_failed = true; + return null; + } + s.cached = k.at; + return .{ .id = k.id, .at = k.at }; + } + + /// The prompt pass stands at a mark: the lane host keeps the state there (native.snap_save). + fn keptAt(ctx: *anyopaque, at: u32) void { + const s: *lanes.Stream = @ptrCast(@alignCast(ctx)); + if (s.reuse.hook) |k| k.at(k.ptr, s, at); + } + + fn firstFn(ptr: *anyopaque, s: *lanes.Stream, position: u64) anyerror!u64 { + const self = of(ptr); + if (position != s.prompt_len) return error.PositionMismatch; + return self.next - 1; + } + + fn queueFn(ptr: *anyopaque, s: *lanes.Stream, feed: be.Feed, position: u64) anyerror!u64 { + _ = ptr; + _ = s; + _ = feed; + _ = position; + return error.NotPipelined; + } + + fn readFn(ptr: *anyopaque, handle: u64) anyerror!u32 { + const self = of(ptr); + if (handle >= self.next or self.next - handle > ring) return error.NoSuchToken; + return self.drawn[handle % ring]; + } + + /// A window's ids (pending token, then its drafts) checked against the bound sequence's position. + fn windowIds(self: *Lanes, w: be.Window, ids: *[max_rows]u32) ![]const u32 { + if (w.parents != null) return error.TreesNotBuilt; + const rows = w.rows(); + if (rows > max_rows) return error.WindowTooWide; + if (w.held > 0) return error.NoHeldDrafts; + const pos = self.e.pos(); + for (w.positions, 0..) |p, r| if (p != pos + 1 + r) return error.PositionMismatch; + ids[0] = w.pending; + @memcpy(ids[1..][0..w.tokens.len], w.tokens); + return ids[0..rows]; + } + + /// Each window in its own forward on its stream's sequence (the core sends one when no forward is shared). + fn verifyFn(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { + const self = of(ptr); + if (windows.len > 1) return self.verifyShared(windows, out); + for (windows, out) |w, *o| { + const l = try self.bind(w.stream, null); + var ids: [max_rows]u32 = undefined; + const got = try self.windowIds(w, &ids); + try self.e.verify(got, o.sampled[0..got.len]); + @memcpy(o.drafts[0 .. got.len - 1], got[1..]); + l.pending_rows = got.len; + } + } + + /// A round's windows in one forward (Engine.verifyShared), each stream's rows on its own sequence. + fn verifyShared(self: *Lanes, windows: []const be.Window, out: []be.Verified) !void { + if (windows.len > engine.max_shared) return error.WindowTooWide; + var parts: [engine.max_shared]engine.Part = undefined; + var ids: [engine.max_shared][max_rows]u32 = undefined; + var lane: [engine.max_shared]*Lane = undefined; + for (windows, out[0..windows.len], 0..) |w, o, k| { + lane[k] = try self.bind(w.stream, null); + const got = try self.windowIds(w, &ids[k]); + parts[k] = .{ .slot = lane[k].slot, .ids = got, .out = o.sampled[0..got.len] }; + } + try self.e.verifyShared(parts[0..windows.len]); + self.shared += 1; + for (out[0..windows.len], parts[0..windows.len], lane[0..windows.len]) |*o, p, l| { + @memcpy(o.drafts[0 .. p.ids.len - 1], p.ids[1..]); + l.pending_rows = p.ids.len; + } + } + + fn keepFn(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) anyerror!void { + const self = of(ptr); + for (windows, paths) |w, path| { + if (path.len == 0) return error.EmptyPath; + for (path, 0..) |r, i| if (r != i) return error.TreesNotBuilt; + _ = try self.bind(w.stream, path.len); + } + } + + fn draftFn(ptr: *anyopaque, requests: []const be.DraftRequest) anyerror!void { + _ = ptr; + _ = requests; + return error.NoDraftHead; + } + + /// The stream left the rounds: its sequence freed (slot 0 stays for the next stream). + fn releaseFn(ptr: *anyopaque, s: *lanes.Stream) void { + const self = of(ptr); + const kv = self.lanes.fetchRemove(s) orelse return; + self.e.release(kv.value.slot); + } +}; diff --git a/zig/src/families/ling/native.zig b/zig/src/families/ling/native.zig new file mode 100644 index 0000000000..f26780d967 --- /dev/null +++ b/zig/src/families/ling/native.zig @@ -0,0 +1,115 @@ +//! Ling-3.1-flash on the native CUDA server at TP=1: the engine, its lane backend, the prompt cache's copies. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const engine = @import("engine.zig"); +const ling_lanes = @import("lanes.zig"); + +pub const model_type = "bailing_hybrid"; +pub const formats: []const []const u8 = &.{ "bf16", "fp8-block", "modelopt-nvfp4" }; +pub const default_context: i64 = 8192; +pub const max_segments: u32 = 1; +pub const prompt_rows: u32 = engine.prompt_rows; + +pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false }; + +pub const LoneRun = *const fn (ctx: *anyopaque, s: *lanes.Stream, hooks: *anyopaque, committed: *const fn (*anyopaque) void, yield: *const fn (*anyopaque) bool, resumed: bool) anyerror!bool; + +pub const Loaded = struct { + backend: lanes.backend.Backend, + facts: lanes.Model, + rows: u32, + stream_bytes: usize, + free_streams: u32 = 1, + learn_bytes: usize = 0, + ctx: *anyopaque, + deinit: *const fn (*anyopaque) void, + lone: ?LoneRun = null, + target: *engine.Engine, // the prompt cache keeps this engine's states (snap_*) +}; + +/// Positions a sequence holds: `context` plus a verify window, a multiple of 64. +pub fn capacity(context: usize) usize { + return std.mem.alignForward(usize, context + engine.window_rows, 64); +} + +const Owned = struct { gpa: std.mem.Allocator, e: *engine.Engine, lanes: ling_lanes.Lanes }; + +/// TF_LING_PROMPT_STEPS=0 passes prompts whole, TF_LING_SHARED=0 runs windows alone. +fn off(name: [*:0]const u8) bool { + return if (std.c.getenv(name)) |v| std.mem.eql(u8, std.mem.span(v), "0") else false; +} + +pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { + _ = kernels; // our own kernels: a captured set is not used + const own = try gpa.create(Owned); + errdefer gpa.destroy(own); + own.gpa = gpa; + own.e = try engine.Engine.init(gpa, io, ctx.d, ctx, dir, .{ .capacity = capacity(o.context) }); + errdefer own.e.deinit(); + own.lanes = .{ .e = own.e, .whole = off("TF_LING_PROMPT_STEPS"), .alone = off("TF_LING_SHARED") }; + return .{ + .backend = own.lanes.backend(), + .facts = own.lanes.facts(), + .rows = ling_lanes.max_rows, + .stream_bytes = own.e.seqBytes(), + .ctx = own, + .deinit = close, + .target = own.e, + }; +} + +fn close(ptr: *anyopaque) void { + const own: *Owned = @ptrCast(@alignCast(ptr)); + own.lanes.deinit(); + own.e.deinit(); + own.gpa.destroy(own); +} + +/// Bytes a state takes at any position: the sequence's arena. +pub fn snap_bytes(ptr: *anyopaque, at: u32) u64 { + _ = at; + const e: *engine.Engine = @ptrCast(@alignCast(ptr)); + return e.stateBytes(); +} + +/// The prompt pass stands at `at` (a mark): copy the state there under a new id. +pub fn snap_save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { + _ = owner; + const e: *engine.Engine = @ptrCast(@alignCast(ptr)); + if (e.pos() != at) return error.NotAtMark; + const k = try e.gpa.create(engine.Kept); + errdefer e.gpa.destroy(k); + k.* = .{ .id = try e.keepState(), .at = at }; + return k; +} + +/// The lane backend restores inside its own prompt pass (Lanes.opening), never from the store's `begin`. +pub fn snap_restore(ptr: *anyopaque, owner: ?*anyopaque, saved: *anyopaque) anyerror!void { + _ = ptr; + _ = owner; + _ = saved; + return error.BackendRestores; +} + +pub fn snap_drop(ptr: *anyopaque, saved: *anyopaque) void { + const e: *engine.Engine = @ptrCast(@alignCast(ptr)); + const k: *engine.Kept = @ptrCast(@alignCast(saved)); + e.dropState(k.id); + e.gpa.destroy(k); +} + +/// The server's words for this family's refusals. +pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { + return switch (err) { + error.Quant => "Ling-3.1 on CUDA reads bf16, block-FP8 and ModelOpt NVFP4 checkpoints (the log names what config.json has)", + error.QuantizedRouter => "Ling-3.1's routing kernels read a bf16 router (mlp.gate): this checkpoint quantized it", + error.PromptTooLong => "the prompt and its reply exceed the window: raise --context", + else => null, + }; +} + +test "capacity keeps a verify window past the context and a multiple of 64" { + try std.testing.expectEqual(@as(usize, 8256), capacity(8192)); + try std.testing.expect(capacity(1000) % 64 == 0 and capacity(1000) >= 1016); +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 254a8799a9..dcc5c02ee9 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -6,6 +6,7 @@ const lanes = @import("lanes"); const nemotron = @import("nemotron"); const lfm2 = @import("lfm2.zig"); const glm = @import("glm_cuda"); +const ling = @import("ling_cuda"); const Allocator = std.mem.Allocator; const budget = @import("cuda_memory.zig"); const lane_backend = @import("cuda_bound.zig"); @@ -15,7 +16,7 @@ const cuda_disk = @import("cuda_disk.zig"); const vision_offer = @import("cuda_vision_offer.zig"); /// CUDA families provide metadata, open their lane backend and explain their refusals. -const registry = .{ nemotron.native, nemotron.omni, glm.native }; +const registry = .{ nemotron.native, nemotron.omni, glm.native, ling.native }; /// --record opens here (engines.zig refuses it for a backend without this). pub const records = true; diff --git a/zig/tests/cuda/ling_forward_synth.zig b/zig/tests/cuda/ling_forward_synth.zig index dcad3c8f0b..73a927ad02 100644 --- a/zig/tests/cuda/ling_forward_synth.zig +++ b/zig/tests/cuda/ling_forward_synth.zig @@ -10,7 +10,7 @@ const Rig = rig.Rig; const P = 80; // prompt rows: KDA's prompt takes the wide chain (64 rows or more) /// ling_forward.py's file: ids, float64 and bf16-point logits [T][V], the picks it routed on [layers][T][top]. -const Ref = struct { T: usize, V: usize, layers: usize, top: usize, forced: bool, gap: f64, ids: []i32, l64: []f64, lbf: []f64, picks: []i32 }; +pub const Ref = struct { T: usize, V: usize, layers: usize, top: usize, forced: bool, gap: f64, ids: []i32, l64: []f64, lbf: []f64, picks: []i32 }; fn take(comptime T: type, a: std.mem.Allocator, bytes: []const u8, at: *usize, n: usize) ![]T { const out = try a.alloc(T, n); @@ -19,7 +19,7 @@ fn take(comptime T: type, a: std.mem.Allocator, bytes: []const u8, at: *usize, n return out; } -fn readRef(a: std.mem.Allocator, io: std.Io, path: []const u8) !Ref { +pub fn readRef(a: std.mem.Allocator, io: std.Io, path: []const u8) !Ref { const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 30)); if (bytes.len < 36 or !std.mem.eql(u8, bytes[0..8], "TFLINGF2")) return error.BadFixture; var at: usize = 8; diff --git a/zig/tests/cuda/ling_lanes.zig b/zig/tests/cuda/ling_lanes.zig new file mode 100644 index 0000000000..4bac636973 --- /dev/null +++ b/zig/tests/cuda/ling_lanes.zig @@ -0,0 +1,178 @@ +//! `ling-lanes DIR OUT [REF] [--layers a,b,..]`: Ling's lane backend under upstream's round loop, no server. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const check = @import("check.zig"); +const ling = @import("ling_cuda"); +const rig = @import("ling_lanes_rig.zig"); +const bytes = @import("ling_lanes_bytes.zig"); + +const Cover = rig.Cover; +const Engine = rig.Engine; + +/// The prompt's tokens: a fixed walk over the vocabulary, as ling_forward.py's. +pub fn walk(a: std.mem.Allocator, n: usize, vocab: usize, salt: usize) ![]u32 { + const p = try a.alloc(u32, n); + for (p, 0..) |*t, i| t.* = @intCast((37 * i + 11 + 101 * salt) % (vocab - 2)); + return p; +} + +/// Rows a prompt chunk takes in these checks: several chunks a prompt, and not a multiple of 16. +const chunk = 37; + +pub fn run(gpu: check.Gpu, rest: []const [:0]const u8) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + if (rest.len < 2) return error.MissingArgument; + const dir = rest[0]; + const out = rest[1]; + var ref: ?[]const u8 = null; + var thin: ?[]u32 = null; + var i: usize = 2; + while (i < rest.len) : (i += 1) { + if (std.mem.eql(u8, rest[i], "--layers") and i + 1 < rest.len) { + var list: std.ArrayList(u32) = .empty; + var it = std.mem.tokenizeScalar(u8, rest[i + 1], ','); + while (it.next()) |x| try list.append(a, try std.fmt.parseInt(u32, x, 10)); + thin = list.items; + i += 1; + } else ref = rest[i]; + } + const e = try Engine.init(gpu.gpa, gpu.io, gpu.d, gpu.ctx, dir, .{ .capacity = 1024, .prompt_rows = 512, .layers = thin }); + defer e.deinit(); + e.chunk = chunk; + var ll: ling.lanes.Lanes = .{ .e = e }; + defer ll.deinit(); + var cfg = try lanes.Config.init(gpu.gpa, ll.facts(), ling.lanes.max_rows, ling.lanes.max_rows - 1); + defer cfg.deinit(gpu.gpa); + var wall: lanes.backend.WallClock = .{ .io = gpu.io }; + var le = lanes.Engine.init(gpu.gpa, &cfg, ll.backend(), wall.clock()); + defer le.deinit(); + var c: Cover = .{ .gpa = gpu.gpa, .a = a, .le = &le }; + const V: usize = e.c.vocab; + std.debug.print("RESULT ling-lanes {t}: {d} layers loaded of {d}, vocabulary {d}, prompt chunks of {d} rows\n", .{ e.c.quant, e.main.layers.len, e.c.layers, V, chunk }); + + const p = try walk(a, 150, V, 0); + // W, what every contract below holds to: the plain greedy pass, its prompt in 37-row chunks + const w = (try c.run(.{ .prompt = p, .max_new = 160, .drafts = false })).ids; + const copy_prompt = try std.mem.concat(a, u32, &.{ p, w[0..100], p }); + + // chunked == whole: the prompt in one chunk, in 16-row chunks, and through `prefill` (no steps) + for ([_]usize{ 512, 16 }) |rows| { + e.chunk = rows; + try c.same(try std.fmt.allocPrint(a, "prompt chunks of {d} rows, against chunks of {d}", .{ rows, chunk }), try c.run(.{ .prompt = p, .max_new = 160, .drafts = false }), w, .length); + } + e.chunk = chunk; + ll.whole = true; + le.backend = ll.backend(); + try c.same("the prompt pass whole (prefill, no steps)", try c.run(.{ .prompt = p, .max_new = 160, .drafts = false }), w, .length); + ll.whole = false; + le.backend = ll.backend(); + + // drafted == plain: windows of 3 to 16 rows from a known reply, most with a wrong row; greedy and sampled + const keyed: lanes.Sampling = .{ .seed = lanes.sampling.seedFor(copy_prompt, 0) }; + const plain_g = try c.run(.{ .prompt = copy_prompt, .max_new = 200, .drafts = false }); + const plain_s = try c.run(.{ .prompt = copy_prompt, .max_new = 200, .drafts = false, .sampling = keyed }); + for ([_]Cover.Got{ plain_g, plain_s }, [_]?lanes.Sampling{ null, keyed }) |plain, sampling| { + const drafted = try c.run(.{ .prompt = copy_prompt, .max_new = 200, .sampling = sampling, .truth = plain.ids }); + try c.same(if (sampling != null) "drafted, sampled, against plain" else "drafted, greedy, against plain", drafted, plain.ids, .length); + try c.expect(drafted.drafted > 0 and drafted.accepted > 0 and drafted.accepted < drafted.drafted, "drafted{s}: {d} drafts in {d} rounds, {d} kept", .{ if (sampling != null) ", sampled" else "", drafted.drafted, drafted.rounds, drafted.accepted }); + } + const g = plain_g.ids; + const differ = std.mem.indexOfDiff(u32, plain_s.ids, g); + try c.expect(differ != null, "the sampled reply leaves the greedy one at token {?d} (sampling draws)", .{differ}); + + // resumed == fresh: states kept at marks off the chunk grid and off multiples of 16, then passes resumed there + { + var keeper: rig.Keeper = .{ .e = e }; + defer keeper.drop(); + const marks = [_]u32{ 37, 101, 150, 263, 349 }; + const marked = try c.run(.{ .prompt = copy_prompt, .max_new = 64, .truth = g, .reuse = .{ .marks = &marks, .hook = keeper.hook() } }); + try c.same("a drafted pass cut at its marks, against the plain fresh pass", marked, g[0..64], .length); + try c.expect(keeper.n == marks.len, "the pass kept {d} states at its marks ({d} asked)", .{ keeper.n, marks.len }); + for (keeper.kept[0..keeper.n], keeper.at[0..keeper.n]) |kept, at| { + const got = try c.run(.{ .prompt = copy_prompt, .max_new = 64, .truth = g, .reuse = .{ .saved = kept, .at = at } }); + try c.expect(got.cached == at, "the pass resumed at {d} ({d} restored)", .{ at, got.cached }); + try c.same(try std.fmt.allocPrint(a, "resumed at {d}, drafted, against the plain fresh pass", .{at}), got, g[0..64], .length); + } + const other = try std.mem.concat(a, u32, &.{ copy_prompt[0..263], try walk(a, 60, V, 3) }); + const other_fresh = try c.run(.{ .prompt = other, .max_new = 64, .drafts = false }); + const got = try c.run(.{ .prompt = other, .max_new = 64, .drafts = false, .reuse = .{ .saved = keeper.kept[3], .at = 263 } }); + try c.same("another tail after the kept 263, against its fresh pass", got, other_fresh.ids, .length); + } + + // cancels in the prompt pass, discards between rounds, eos, max_new at the window's edges + for ([_]u32{ 0, 1, 3 }) |at| { + var cut: rig.Cut = .{ .at = at }; + const r = c.run(.{ .prompt = copy_prompt, .max_new = 64, .cancel = &cut }); + try c.expect(if (r) |_| false else |err| err == error.Cancelled, "a cancel at ask {d} of the prompt pass ended it with error.Cancelled", .{at}); + try c.same(try std.fmt.allocPrint(a, "the request after a cancel at prompt ask {d}", .{at}), try c.run(.{ .prompt = p, .max_new = 64, .drafts = false }), w[0..64], .length); + } + for ([_]u32{ 1, 8 }) |rounds| { + try c.expect(try c.discardAfter(copy_prompt, g, rounds), "a drafted stream was discarded after {d} rounds", .{rounds}); + try c.same(try std.fmt.allocPrint(a, "the request after a discard after {d} rounds", .{rounds}), try c.run(.{ .prompt = p, .max_new = 64, .drafts = false }), w[0..64], .length); + } + for ([_]usize{ 0, 9, 33 }) |from| { + const j = for (from..g.len) |k| { + if (std.mem.indexOfScalar(u32, g[0..k], g[k]) == null) break k; + } else continue; + try c.same(try std.fmt.allocPrint(a, "eos at token {d}, drafted", .{j}), try c.run(.{ .prompt = copy_prompt, .max_new = 512, .eos = &.{g[j]}, .truth = g }), g[0 .. j + 1], .stop); + } + for (1..4) |n| try c.same(try std.fmt.allocPrint(a, "max_new {d}, drafted", .{n}), try c.run(.{ .prompt = copy_prompt, .max_new = @intCast(n), .truth = g }), g[0..n], .length); + + // concurrent == solo + try concurrent(&c, &ll, a, V); + // the engine's own steps, byte for byte: cuts, resumes, windows and shared forwards against their twins + try bytes.engineChecks(&c, e, copy_prompt, w); + // the plain pass's ids and picks for ling_forward.py; with REF, its logits against float64 + try bytes.serial(&c, gpu, e, p, w, out, ref); + if (c.failed) return error.TestFailed; + check.pass("ling-lanes {t}: {d} checks", .{ e.c.quant, c.checks }); +} + +/// Each concurrent stream on one lanes.Engine equals its solo run. +fn concurrent(c: *Cover, ll: *ling.lanes.Lanes, a: std.mem.Allocator, V: usize) !void { + const n = 96; + var ps: [3][]const u32 = undefined; + for (&ps, 0..) |*p, k| { + const base = try walk(a, 60 + 23 * k, V, k + 1); + p.* = try std.mem.concat(a, u32, &.{ base, base[0..40], base }); + } + const seeds = [_]lanes.Sampling{ .{ .seed = 11 }, .{ .seed = 22 }, .{ .seed = 33 } }; + var solo: [3][]const u32 = undefined; + var solo_s: [3][]const u32 = undefined; + for (ps, &solo, &solo_s, 0..) |p, *g, *gs, k| { + g.* = (try c.run(.{ .prompt = p, .max_new = n, .drafts = false })).ids; + gs.* = (try c.run(.{ .prompt = p, .max_new = n, .drafts = false, .sampling = seeds[k] })).ids; + } + var sampled_seen = false; // the drafts' truth: the greedy solo replies, or the sampled ones in sampled cases + const cases = [_]struct { name: []const u8, joins: []const rig.Join }{ + .{ .name = "all at once", .joins = &.{ .{ .p = 0 }, .{ .p = 1 }, .{ .p = 2 } } }, + .{ .name = "all at once, drafted", .joins = &.{ .{ .p = 0, .drafts = true }, .{ .p = 1, .drafts = true }, .{ .p = 2, .drafts = true } } }, + .{ .name = "staggered by 5 rounds, drafted", .joins = &.{ .{ .p = 0, .drafts = true }, .{ .p = 1, .at = 5, .drafts = true }, .{ .p = 2, .at = 10 } } }, + .{ .name = "in reverse", .joins = &.{ .{ .p = 2 }, .{ .p = 1, .drafts = true }, .{ .p = 0 } } }, + .{ .name = "the same prompt twice", .joins = &.{ .{ .p = 0, .drafts = true }, .{ .p = 0 } } }, + .{ .name = "sampled, drafted", .joins = &.{ .{ .p = 0, .sampled = true, .drafts = true }, .{ .p = 1, .sampled = true }, .{ .p = 2, .sampled = true, .drafts = true } } }, + .{ .name = "the second discarded after 6 rounds", .joins = &.{ .{ .p = 0 }, .{ .p = 1, .discard_at = 6 }, .{ .p = 2, .drafts = true } } }, + .{ .name = "prompts a chunk at a time while the others decode", .joins = &.{ .{ .p = 0, .drafts = true }, .{ .p = 1, .at = 2, .filled = true }, .{ .p = 2, .at = 4, .filled = true, .drafts = true } } }, + .{ .name = "sampled, prompts a chunk at a time", .joins = &.{ .{ .p = 0, .sampled = true }, .{ .p = 1, .at = 2, .sampled = true, .filled = true }, .{ .p = 2, .at = 4, .sampled = true, .filled = true, .drafts = true } } }, + }; + const before = ll.shared; + for (cases) |cs| { + sampled_seen = cs.joins[0].sampled; + const got = try c.together(cs.joins, &ps, if (sampled_seen) &solo_s else &solo, &seeds, n); + for (cs.joins, got.ids, 0..) |j, ids, k| { + if (j.discard_at != null) continue; + try c.same(try std.fmt.allocPrint(a, "{s}: stream {d} (prompt {d}{s}) against its solo run", .{ cs.name, k, j.p, if (j.sampled) ", sampled" else "" }), .{ .ids = ids, .reason = .length, .rounds = 0, .cached = 0, .drafted = 0, .accepted = 0 }, if (j.sampled) solo_s[j.p] else solo[j.p], .length); + } + var drafting = false; + for (cs.joins) |j| drafting = drafting or j.drafts; + try c.expect(got.most >= 2 and (got.drafted > 0) == drafting, "{s}: {d} streams at most in one round, {d} rounds with two or more, {d} drafts", .{ cs.name, got.most, got.shared, got.drafted }); + for (cs.joins) |j| if (j.filled) { + try c.expect(got.between > 0, "{s}: {d} rounds of the others between a prompt's chunks", .{ cs.name, got.between }); + break; + }; + } + try c.expect(ll.shared > before, "{d} forwards held several streams' windows (Engine.verifyShared)", .{ll.shared - before}); +} diff --git a/zig/tests/cuda/ling_lanes_bytes.zig b/zig/tests/cuda/ling_lanes_bytes.zig new file mode 100644 index 0000000000..85e0147765 --- /dev/null +++ b/zig/tests/cuda/ling_lanes_bytes.zig @@ -0,0 +1,198 @@ +//! The Ling engine's steps against their twins byte for byte, and the plain pass's logits against float64. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const ling = @import("ling_cuda"); +const rig = @import("ling_lanes_rig.zig"); +const frig = @import("ling_forward_rig.zig"); +const synth = @import("ling_forward_synth.zig"); + +const Cover = rig.Cover; +const Engine = rig.Engine; +const F = ling.forward; + +/// Kept-state ids by mark, as a direct engine pass keeps them. +const Marked = struct { + e: *Engine, + ids: [8]u64 = undefined, + n: usize = 0, + fn at(ptr: *anyopaque, pos: u32) void { + _ = pos; + const m: *Marked = @ptrCast(@alignCast(ptr)); + m.ids[m.n] = m.e.keepState() catch return; + m.n += 1; + } +}; + +fn logitsRows(e: *Engine, a: std.mem.Allocator, ptr: u64, rows: usize) ![]u8 { + const out = try a.alloc(u8, rows * e.c.vocab * 2); + try e.stream.synchronize(); + const b: cuda.DeviceBuffer = .{ .d = e.d, .ptr = ptr, .len = out.len }; + try b.download(0, out); + return out; +} + +/// A fresh pass of `ids` on `slot` with `rows`-row chunks (and `marks`), or resumed from `from`; its first token. +fn pass(e: *Engine, slot: u32, ids: []const u32, rows: usize, from: ?ling.engine.Resume, marks: ?ling.engine.Marks) !u32 { + try e.bind(slot); + try e.begin(null); + e.chunk = rows; + return e.prefill(ids, from, null, marks); +} + +pub fn engineChecks(c: *Cover, e: *Engine, prompt: []const u32, w: []const u32) !void { + const a = c.a; + const keep_chunk = e.chunk; + defer e.chunk = keep_chunk; + const t10 = try pass(e, 10, prompt, 512, null, null); + const t11 = try pass(e, 11, prompt, 37, null, null); + const t12 = try pass(e, 12, prompt, 16, null, null); + try c.expect(t10 == t11 and t10 == t12, "first tokens: whole {d}, 37-row chunks {d}, 16-row chunks {d}", .{ t10, t11, t12 }); + try rig.sameSeq(c, e, "prompt in 37-row chunks == whole", 11, 10); + try rig.sameSeq(c, e, "prompt in 16-row chunks == whole", 12, 10); + var mk: Marked = .{ .e = e }; + const marks = [_]u32{ 101, 263, 349 }; + _ = try pass(e, 13, prompt, 48, null, .{ .at = &marks, .ctx = &mk, .call = Marked.at }); + try c.expect(mk.n == marks.len, "a pass kept {d} states at marks 101, 263, 349", .{mk.n}); + for (mk.ids[0..mk.n], marks[0..mk.n]) |id, at| { + const t = try pass(e, 14, prompt, 37, .{ .id = id, .at = at }, null); + try c.expect(t == t10, "resumed at {d}: first token {d}, the whole pass's {d}", .{ at, t, t10 }); + try rig.sameSeq(c, e, try std.fmt.allocPrint(a, "resumed at {d} == whole", .{at}), 14, 10); + e.dropState(id); + } + // a 9-row window, keep 4, a 5-row window == 9 single rows on the twin (tokens, logits, states) + var g: [10]u32 = undefined; + g[0] = t11; + const lone = try a.alloc([]u8, 9); + try e.bind(11); + for (0..9) |r| { + try e.verify(g[r .. r + 1], g[r + 1 .. r + 2]); + lone[r] = try logitsRows(e, a, e.bd.logits, 1); + try e.keep(1); + } + try e.bind(10); + var got: [9]u32 = undefined; + try e.verify(g[0..9], &got); + const win = try logitsRows(e, a, e.bd.logits, 9); + var same_rows = std.mem.eql(u32, &got, g[1..10]); + for (0..9) |r| same_rows = same_rows and std.mem.eql(u8, win[r * lone[r].len ..][0..lone[r].len], lone[r]); + try e.keep(4); + try e.verify(g[4..9], got[0..5]); + const win2 = try logitsRows(e, a, e.bd.logits, 5); + same_rows = same_rows and std.mem.eql(u32, got[0..5], g[5..10]); + for (0..5) |r| same_rows = same_rows and std.mem.eql(u8, win2[r * lone[4 + r].len ..][0..lone[4 + r].len], lone[4 + r]); + try e.keep(5); + try c.expect(same_rows, "a 9-row window and, after keeping 4, a 5-row window: tokens and logits == 9 single rows", .{}); + try rig.sameSeq(c, e, "9-row window, keep 4, 5-row window == 9 single rows", 10, 11); + try sharedChecks(c, e, prompt, w); + for ([_]u32{ 10, 11, 12, 13, 14 }) |s| e.release(s); +} + +/// Two streams' windows in one shared forward against lone twins over 8 rounds: logits, tokens, then states. +fn sharedChecks(c: *Cover, e: *Engine, prompt: []const u32, w: []const u32) !void { + const a = c.a; + const pa = prompt[0..120]; + const pb = prompt[50..250]; + var pend: [2]u32 = undefined; + for ([_]u32{ 20, 21, 22, 23 }, 0..) |slot, k| pend[k % 2] = try pass(e, slot, if (k % 2 == 0) pa else pb, 37, null, null); + var ok = true; + var rows_seen: usize = 0; + for (0..8) |round| { + var ids: [2][16]u32 = undefined; + var outs: [2][16]u32 = undefined; + var twin: [2][16]u32 = undefined; + var parts: [2]ling.engine.Part = undefined; + for (0..2) |k| { + const rows = 1 + (round * 5 + k * 7) % 16; + ids[k][0] = pend[k]; + for (1..rows) |r| ids[k][r] = w[(round * 13 + k * 31 + r) % w.len]; + parts[k] = .{ .slot = @intCast(20 + k), .ids = ids[k][0..rows], .out = outs[k][0..rows] }; + } + try e.verifyShared(&parts); + const sh = try logitsRows(e, a, e.bs.logits, parts[0].ids.len + parts[1].ids.len); + var row0: usize = 0; + for (0..2) |k| { + const rows = parts[k].ids.len; + try e.bind(@intCast(22 + k)); + try e.verify(parts[k].ids, twin[k][0..rows]); + const lo = try logitsRows(e, a, e.bd.logits, rows); + const V2 = e.c.vocab * 2; + if (!std.mem.eql(u8, sh[row0 * V2 ..][0 .. rows * V2], lo) or !std.mem.eql(u32, outs[k][0..rows], twin[k][0..rows])) ok = false; + row0 += rows; + rows_seen += rows; + const keep = 1 + (round * 3 + k) % rows; + for ([_]u32{ @intCast(20 + k), @intCast(22 + k) }) |s| { + try e.bind(s); + try e.keep(keep); + } + pend[k] = outs[k][keep - 1]; + } + } + try c.expect(ok, "shared windows: {d} rows of 2 streams over 8 rounds (1 to 16 rows each) == lone twins' logits and tokens", .{rows_seen}); + try rig.sameSeq(c, e, "shared stream 0 after 8 rounds == its lone twin", 20, 22); + try rig.sameSeq(c, e, "shared stream 1 after 8 rounds == its lone twin", 21, 23); + for ([_]u32{ 20, 21, 22, 23 }) |s| e.release(s); +} + +/// The plain pass serially (whole prompt, then one row at a time) with its picks; files for ling_forward.py --ids. +pub fn serial(c: *Cover, gpu: check.Gpu, e: *Engine, p: []const u32, w: []const u32, out: []const u8, ref_path: ?[]const u8) !void { + const a = c.a; + const ids = try a.alloc(i32, p.len + w.len - 1); + for (ids, 0..) |*x, k| x.* = @intCast(if (k < p.len) p[k] else w[k - p.len]); + const T = ids.len; + const P = p.len; + const V: usize = e.c.vocab; + const top: usize = e.c.topk; + var nm: usize = 0; + for (e.main.layers) |l| nm += @intFromBool(l.ffn == .moe); + var r: frig.Rig = .{ .gpu = gpu, .a = a, .s = e.stream, .m = &e.m, .g = &e.g }; + defer r.deinit(); + try e.bind(30); + try e.begin(null); + const st = &e.cur.st; + const logits = try a.alloc(u16, T * V); + const picks = try a.alloc(i32, nm * T * top); + const pp = try r.zeros(nm * P * (top + 1) * 4); + @memcpy(logits[0 .. P * V], try r.run(&e.bp, &.{.{ .st = st, .row0 = 0, .rows = P }}, ids[0..P], pp)); + const sp = try frig.sortedPicks(a, try r.back(pp, nm * P * (top + 1) * 4), nm, P, top); + for (0..nm) |l| @memcpy(picks[l * T * top ..][0 .. P * top], sp[l * P * top ..][0 .. P * top]); + try r.fwd(&e.bp, null).commit(st, P); + const dp = try r.zeros(nm * (top + 1) * 4); + for (P..T) |t| { + @memcpy(logits[t * V ..][0..V], try r.run(&e.bd, &.{.{ .st = st, .row0 = 0, .rows = 1 }}, ids[t .. t + 1], dp)); + const one = try frig.sortedPicks(a, try r.back(dp, nm * (top + 1) * 4), nm, 1, top); + for (0..nm) |l| @memcpy(picks[(l * T + t) * top ..][0..top], one[l * top ..][0..top]); + try r.fwd(&e.bd, null).commit(st, 1); + } + e.release(30); + var agree: usize = 0; + for (P - 1..T) |t| { + var best: usize = 0; + for (0..V) |j| if (frig.bfOf(logits[t * V + j]) > frig.bfOf(logits[t * V + best])) { + best = j; + }; + const next: u32 = if (t + 1 < T) @intCast(ids[t + 1]) else w[w.len - 1]; + agree += @intFromBool(best == next); + } + try c.expect(agree == T - P + 1, "the lanes' greedy reply == the serial forward's argmax at all {d} of {d} rows", .{ agree, T - P + 1 }); + var name: [512]u8 = undefined; + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = try std.fmt.bufPrint(&name, "{s}.ids", .{out}), .data = std.mem.sliceAsBytes(ids) }); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = try std.fmt.bufPrint(&name, "{s}.picks", .{out}), .data = std.mem.sliceAsBytes(picks) }); + const path = ref_path orelse return; + const ref = try synth.readRef(a, gpu.io, path); + try c.expect(ref.T == T and ref.V == V and ref.forced and std.mem.eql(i32, ref.ids, ids), "reference: {d} tokens routed on these picks, the same ids", .{ref.T}); + const flipped = try a.alloc(bool, T); + @memset(flipped, false); + const nf = frig.flips(picks, ref.picks, nm, T, 0, T, top, flipped); + const d64 = frig.distance(logits, ref.l64, T, V, flipped); + const dbf = frig.distance(logits, ref.lbf, T, V, flipped); + var bd: f64 = 0; + var bw: f64 = 0; + for (ref.lbf, ref.l64) |x, y| { + bd += (x - y) * (x - y); + bw += y * y; + } + const base = @sqrt(bd / bw); + std.debug.print("RESULT served logits vs float64 (same picks): rel L2 {e:.3}, worst row {e:.3}, argmax {d}/{d}; vs bf16 points {e:.3}; bf16 points vs float64 {e:.3}\n", .{ d64.rel, d64.worst_row, d64.argmax, T, dbf.rel, base }); + try c.expect(nf == 0 and d64.rel <= 1.5 * base and dbf.rel <= base, "served logits within 1.5x the bf16 points' rel L2 of float64 ({e:.3} <= 1.5 x {e:.3}), within it of the bf16 points", .{ d64.rel, base }); +} diff --git a/zig/tests/cuda/ling_lanes_rig.zig b/zig/tests/cuda/ling_lanes_rig.zig new file mode 100644 index 0000000000..9c9dc4c0c6 --- /dev/null +++ b/zig/tests/cuda/ling_lanes_rig.zig @@ -0,0 +1,218 @@ +//! What the Ling lane checks share: requests through one lanes.Engine, joined streams, kept states, byte comparisons. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const check = @import("check.zig"); +const ling = @import("ling_cuda"); + +pub const Engine = ling.engine.Engine; + +/// A mark hook that keeps the state as the lane host's prompt cache does (native.snap_save). +pub const Keeper = struct { + e: *Engine, + kept: [8]*anyopaque = undefined, + at: [8]u32 = undefined, + n: usize = 0, + + pub fn hook(k: *Keeper) lanes.stream.MarkHook { + return .{ .ptr = k, .at = keep }; + } + + fn keep(ptr: *anyopaque, s: *lanes.Stream, pos: u32) void { + _ = s; + const k: *Keeper = @ptrCast(@alignCast(ptr)); + k.kept[k.n] = ling.native.snap_save(k.e, null, pos) catch |err| { + std.debug.print("FAIL keeping the state at {d}: {s}\n", .{ pos, @errorName(err) }); + return; + }; + k.at[k.n] = pos; + k.n += 1; + } + + pub fn drop(k: *Keeper) void { + for (k.kept[0..k.n]) |x| ling.native.snap_drop(k.e, x); + k.n = 0; + } +}; + +/// A cancel hook that answers yes from its `at`-th ask on. +pub const Cut = struct { + at: u32, + asked: u32 = 0, + + fn ask(ptr: *anyopaque) bool { + const cut: *Cut = @ptrCast(@alignCast(ptr)); + cut.asked += 1; + return cut.asked > cut.at; + } +}; + +/// Drafts from a known reply (the plain run's): 2 to 15 a round, most rounds one of them wrong, to reject anywhere. +pub const Oracle = struct { + prompt_len: usize, + truth: []const u32, + out: [16]u32 = undefined, + rounds: u32 = 0, + + pub fn proposer(o: *Oracle) lanes.proposer.Proposer { + return .{ .ptr = o, .vtable = &.{ .propose = propose, .last_match = lastMatch } }; + } + + fn propose(ptr: *anyopaque, context: []const u32, max_draft: i64) anyerror![]const u32 { + const o: *Oracle = @ptrCast(@alignCast(ptr)); + o.rounds += 1; + const k = context.len - o.prompt_len; // the reply so far, its last token pending: drafts follow it + if (max_draft < 2 or k >= o.truth.len) return &.{}; + const n = @min(@as(usize, @intCast(max_draft)), o.truth.len - k, 2 + o.rounds % 14); + if (n < 2) return &.{}; + @memcpy(o.out[0..n], o.truth[k..][0..n]); + if (o.rounds % 4 != 0) o.out[(o.rounds * 5) % n] ^= 1; // a wrong token at a moving row + return o.out[0..n]; + } + + fn lastMatch(_: *anyopaque) i64 { + return 64; + } +}; + +/// A concurrent case's stream: prompt `p` joining at round `at`; `filled`: its prompt a chunk at a time between rounds. +pub const Join = struct { p: usize, at: u32 = 0, sampled: bool = false, discard_at: ?u32 = null, filled: bool = false, drafts: bool = false }; + +pub const Cover = struct { + gpa: std.mem.Allocator, + a: std.mem.Allocator, + le: *lanes.Engine, + failed: bool = false, + checks: usize = 0, + + pub const Ask = struct { + prompt: []const u32, + max_new: u32, + eos: []const u32 = &.{}, + drafts: bool = true, + sampling: ?lanes.Sampling = null, + truth: []const u32 = &.{}, // drafts from this reply (Oracle); empty: none + cancel: ?*Cut = null, + reuse: lanes.stream.Reuse = .{}, + }; + pub const Got = struct { ids: []const u32, reason: lanes.stream.Reason, rounds: u64, cached: u32, drafted: u64, accepted: u64 }; + + /// One request through the shared lanes.Engine, as the server's lane host runs it. + pub fn run(c: *Cover, ask: Ask) !Got { + var o: Oracle = .{ .prompt_len = ask.prompt.len, .truth = ask.truth }; + var s = try lanes.Stream.init(c.gpa, .{ .id = "cover", .prompt = ask.prompt, .max_new = ask.max_new, .eos = ask.eos, .drafts = ask.drafts, .sampling = ask.sampling, .proposer = if (ask.truth.len > 0) o.proposer() else null, .cancel_check = if (ask.cancel) |cut| .{ .ptr = cut, .check = Cut.ask } else null, .reuse = ask.reuse }); + defer s.deinit(c.gpa); + try c.le.addStream(&s); + while (c.le.activeCount() > 0) try c.le.step(); + return .{ .ids = try c.a.dupe(u32, s.emitted()), .reason = s.reason, .rounds = s.rounds, .cached = s.cached, .drafted = s.drafted, .accepted = s.accepted }; + } + + pub const Together = struct { ids: []const []const u32, most: usize, shared: u32, between: u32, drafted: u64 }; + + /// The streams of `joins` through the rounds together, each joining at its round; how many were live at once. + pub fn together(c: *Cover, joins: []const Join, ps: []const []const u32, truths: []const []const u32, seeds: []const lanes.Sampling, max_new: u32) !Together { + const ss = try c.a.alloc(lanes.Stream, joins.len); + const looks = try c.a.alloc(Oracle, joins.len); + const made = try c.a.alloc(bool, joins.len); + @memset(made, false); + defer for (ss, made) |*s, m| if (m) s.deinit(c.gpa); + var most: usize = 0; + var shared: u32 = 0; + var between: u32 = 0; + var round: u32 = 0; + var filling: std.ArrayList(*lanes.Stream) = .empty; // the lane host's fillOne: the oldest's next chunk a round + defer filling.deinit(c.gpa); + var began = false; + while (true) : (round += 1) { + for (joins, ss, looks, made) |j, *s, *l, *m| if (!m.* and j.at <= round) { + l.* = .{ .prompt_len = ps[j.p].len, .truth = truths[j.p] }; + s.* = try lanes.Stream.init(c.gpa, .{ .id = "joined", .prompt = ps[j.p], .max_new = max_new, .sampling = if (j.sampled) seeds[j.p % seeds.len] else null, .proposer = if (j.drafts) l.proposer() else null }); + m.* = true; + if (j.filled) try filling.append(c.gpa, s) else try c.le.addStream(s); + }; + if (filling.items.len > 0) { + if (try c.le.fillStream(filling.items[0], !began)) { + _ = filling.orderedRemove(0); + began = false; + } else began = true; + } + for (joins, ss, made) |j, *s, m| if (m and j.discard_at == round and !s.finished) c.le.discard(s); + const live = c.le.activeCount(); + most = @max(most, live); + shared += @intFromBool(live >= 2); + if (live == 0) { + if (std.mem.indexOfScalar(bool, made, false) == null and filling.items.len == 0) break; + continue; + } + between += @intFromBool(began); + try c.le.step(); + } + const ids = try c.a.alloc([]const u32, joins.len); + var drafted: u64 = 0; + for (ss, ids) |*s, *i| { + i.* = try c.a.dupe(u32, s.emitted()); + drafted += s.drafted; + } + return .{ .ids = ids, .most = most, .shared = shared, .between = between, .drafted = drafted }; + } + + /// A drafted stream over `prompt` (drafts from `truth`), discarded after `rounds`; false if it ended first. + pub fn discardAfter(c: *Cover, prompt: []const u32, truth: []const u32, rounds: u32) !bool { + var o: Oracle = .{ .prompt_len = prompt.len, .truth = truth }; + var s = try lanes.Stream.init(c.gpa, .{ .id = "discarded", .prompt = prompt, .max_new = 512, .proposer = o.proposer() }); + defer s.deinit(c.gpa); + try c.le.addStream(&s); + var n: u32 = 0; + while (c.le.activeCount() > 0 and n < rounds) : (n += 1) try c.le.step(); + if (c.le.activeCount() == 0) return false; + c.le.discard(&s); + return true; + } + + pub fn same(c: *Cover, what: []const u8, got: Got, want: []const u32, reason: lanes.stream.Reason) !void { + c.checks += 1; + const n = @min(got.ids.len, want.len); + const first = for (got.ids[0..n], want[0..n], 0..) |x, y, i| { + if (x != y) break i; + } else if (got.ids.len != want.len) n else null; + if (first == null and got.reason == reason) { + check.pass("{s}: {d} ids equal, finished by {s} ({d} rounds, {d} of {d} drafts kept)", .{ what, want.len, reason.name(), got.rounds, got.accepted, got.drafted }); + return; + } + c.failed = true; + if (first) |i| { + std.debug.print("FAIL {s}: token {d} of {d} is {d}, want {d} (emitted {d}, {s})\n", .{ what, i, want.len, if (i < got.ids.len) got.ids[i] else 0, if (i < want.len) want[i] else 0, got.ids.len, got.reason.name() }); + } else std.debug.print("FAIL {s}: finished by {s}, want {s}\n", .{ what, got.reason.name(), reason.name() }); + } + + pub fn expect(c: *Cover, ok: bool, comptime fmt: []const u8, args: anytype) !void { + c.checks += 1; + if (ok) return check.pass(fmt, args); + c.failed = true; + std.debug.print("FAIL " ++ fmt ++ "\n", args); + } +}; + +fn back(e: *Engine, a: std.mem.Allocator, ptr: u64, len: usize) ![]u8 { + const out = try a.alloc(u8, len); + try e.stream.synchronize(); + const b: cuda.DeviceBuffer = .{ .d = e.d, .ptr = ptr, .len = len }; + try b.download(0, out); + return out; +} + +/// Slots `x` and `y` hold the same committed sequence: position, KDA states and conv windows, MLA rows below it. +pub fn sameSeq(c: *Cover, e: *Engine, what: []const u8, x: u32, y: u32) !void { + const p = &e.seqs.items[x].?.st; + const q = &e.seqs.items[y].?.st; + var bad: ?[]const u8 = null; + if (p.pos != q.pos) bad = "positions"; + for (0..p.kda_layers) |li| { + if (!std.mem.eql(u8, try back(e, c.a, p.recAt(p.cur[li], li), p.rec_layer), try back(e, c.a, q.recAt(q.cur[li], li), q.rec_layer))) bad = "a KDA state"; + if (!std.mem.eql(u8, try back(e, c.a, p.convAt(li), 3 * p.convWidth() * 2), try back(e, c.a, q.convAt(li), 3 * q.convWidth() * 2))) bad = "a conv window"; + } + for (0..p.mla_layers) |mi| if (!std.mem.eql(u8, try back(e, c.a, p.cacheAt(mi), p.pos * ling.forward.cache_width * 2), try back(e, c.a, q.cacheAt(mi), p.pos * ling.forward.cache_width * 2))) { + bad = "MLA cache rows"; + }; + try c.expect(bad == null, "{s}: slots {d} and {d} at {d}, byte for byte{s}{s}", .{ what, x, y, p.pos, if (bad != null) ", but differ in " else "", bad orelse "" }); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index b3d13c7f1a..5533ea1488 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -116,6 +116,7 @@ const usage = \\ ling-layer-real one real Ling-3.1 layer vs HF's captured output (oracle/ling_layer.py); exactness \\ ling-clamp Ling's SwiGLU order min(silu(g), L): glue and the dir's grouped experts vs float64 \\ ling-clamp-old the grouped expert instances vs an older build's fatbins, byte for byte + \\ ling-lanes [ref] [--layers a,b] Ling's lane backend: chunked, drafted, resumed, concurrent == plain \\ ling-weights-real real Ling-3.1 layers 0, 4, 7 (4 experts) vs float64 and torch (oracle/ling_weights.py) \\ ; @@ -228,6 +229,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "ling-layer-real")) return @import("ling_layer_real.zig").run(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "ling-clamp")) return @import("ling_clamp.zig").run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "ling-clamp-old")) return @import("ling_clamp.zig").old(gpu, try arg(rest, 0), try arg(rest, 1)); + if (std.mem.eql(u8, cmd, "ling-lanes")) return @import("ling_lanes.zig").run(gpu, rest); if (std.mem.eql(u8, cmd, "ling-weights-real")) return ling_real.run(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "glm-lanes-concurrent")) return glm_lanes_tests.concurrentCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); std.debug.print("{s}", .{usage}); diff --git a/zig/tests/cuda/oracle/ling_forward.py b/zig/tests/cuda/oracle/ling_forward.py index 5eb34d1915..5e76e09176 100644 --- a/zig/tests/cuda/oracle/ling_forward.py +++ b/zig/tests/cuda/oracle/ling_forward.py @@ -130,7 +130,7 @@ def moe(self, x, p, i): masked = choice.masked_fill(keep == 0, float("-inf")) idx = masked.topk(top, -1)[1] self.picks.append(idx.sort(-1)[0].int()) - if self.forced is not None: # the graph's picks; gap: how far the worst one ranks below the top-k + if self.forced is not None and len(self.picks) <= self.forced.shape[0]: # the graph's picks; gap: how far the worst one ranks below the top-k idx = self.forced[len(self.picks) - 1].long() kth = masked.topk(top, -1)[0][:, -1] self.gap = max(self.gap, (kth - choice.gather(1, idx).min(-1)[0]).max().item()) @@ -154,7 +154,20 @@ def logits(self, ids): x = self.r(x + (self.mla(h, p + "attention.") if mla else self.kda(h, p + "attention."))) h = self.norm(x, ck.t(p + "post_attention_layernorm.weight")) x = self.r(x + (self.moe(h, p + "mlp.", i) if i >= c["first_k_dense_replace"] else self.mlp(h, p + "mlp.", 0))) - return self.norm(x, ck.t("model.norm.weight")) @ ck.lin("lm_head").T + self.hidden = self.norm(x, ck.t("model.norm.weight")) + return self.hidden @ ck.lin("lm_head").T + + def mtp_logits(self, ids): + """The MTP layer's logits at rows 0 .. T - 2: hidden row t (after logits()) with token t + 1's embedding.""" + c, ck = self.c, self.ck + L = c["num_hidden_layers"] + p = f"model.layers.{L}." + e = self.norm(ck.t("model.word_embeddings.weight").double()[ids[1:]], ck.t(p + "enorm.weight")) + h = self.norm(self.hidden[:-1], ck.t(p + "hnorm.weight")) + x = self.lin(torch.cat([e, h], -1), p + "eh_proj") + x = self.r(x + self.mla(self.norm(x, ck.t(p + "input_layernorm.weight")), p + "attention.")) + x = self.r(x + self.moe(self.norm(x, ck.t(p + "post_attention_layernorm.weight")), p + "mlp.", L)) + return self.norm(x, ck.t(p + "final_layernorm.weight")) @ ck.lin("lm_head").T def main(): @@ -163,11 +176,17 @@ def main(): a.add_argument("out", help="the reference file: ids, float64 and bf16-point logits, the picks it routed on") a.add_argument("--tokens", type=int, default=88) a.add_argument("--picks", help="route on these picks (i32 [MoE layers][tokens][top_k], the graph's) instead") + a.add_argument("--mtp", help="also the MTP layer's logits (float64, bf16 points) and its own picks, to this file") + a.add_argument("--ids", help="these tokens (i32, ling-lanes' OUT.ids) instead of the fixed walk; --tokens follows") o = a.parse_args() c = json.load(open(os.path.join(o.ckpt, "config.json"))) ck = Ckpt(os.path.join(o.ckpt, "model.safetensors")) - T, V = o.tokens, c["vocab_size"] - ids = torch.tensor([(37 * t + 11) % (V - 2) for t in range(T)]) + V = c["vocab_size"] + if o.ids: + ids = torch.frombuffer(bytearray(open(o.ids, "rb").read()), dtype=torch.int32).long() + else: + ids = torch.tensor([(37 * t + 11) % (V - 2) for t in range(o.tokens)]) + T = ids.shape[0] forced = None if o.picks: nm = sum(1 for i in range(c["num_hidden_layers"]) if i >= c["first_k_dense_replace"]) @@ -176,7 +195,17 @@ def main(): exact = Model(ck, c, lambda z: z, forced) l64 = exact.logits(ids) lbf = Model(ck, c, bf16, forced).logits(ids) - picks = forced if forced is not None else torch.stack(exact.picks) + if o.mtp: # the bf16-point model's own picks route its MTP layer; float64's are written for the flips + rounded = Model(ck, c, bf16, forced) + rounded.logits(ids) + m64, mbf = exact.mtp_logits(ids), rounded.mtp_logits(ids) + mp = exact.picks[-1] + with open(o.mtp, "wb") as f: + f.write(b"TFLINGM2" + struct.pack("<3I", T - 1, V, mp.shape[1])) + f.write(m64.numpy().tobytes() + mbf.numpy().tobytes() + mp.numpy().tobytes()) + print("mtp rows", T - 1, "logit rms", m64.pow(2).mean().sqrt().item()) + nm = sum(1 for i in range(c["num_hidden_layers"]) if i >= c["first_k_dense_replace"]) + picks = forced if forced is not None else torch.stack(exact.picks[:nm]) with open(o.out, "wb") as f: f.write(b"TFLINGF2" + struct.pack("<4I", T, V, picks.shape[0], picks.shape[2]) + struct.pack(" Date: Thu, 15 Oct 2026 16:59:29 +0200 Subject: [PATCH 11/34] ling: the MTP layer (model.layers.72) drafts: enorm | hnorm through eh_proj, MLA + MoE on its own cache, chains of 3 mtp.Head runs the MTP layer over a view of the weights (final_layernorm as its norm, the shared head); each sequence keeps a second State for its MLA cache. Prompt passes feed every row whose next token is known, Engine.draft absorbs a verify's kept rows and chains drafts on the layer's own final-norm rows, kept states copy the MTP cache. ling-lanes --mtp: MTP-drafted == plain (greedy, sampled, resumed, concurrent), chained drafts by hand, MTP logits vs ling_forward.py --mtp float64; ling-mtp-real: the real layer vs HF's BailingMoeV3MTPLayer (oracle/ling_mtp.py). --- docs/recipes/ling-3.1.md | 42 +++++++- zig/src/families/ling/cuda.zig | 2 + zig/src/families/ling/engine.zig | 104 ++++++++++++++++-- zig/src/families/ling/lanes.zig | 43 ++++++-- zig/src/families/ling/mtp.zig | 70 ++++++++++++ zig/src/families/ling/native.zig | 2 +- zig/tests/cuda/ling_lanes.zig | 25 ++++- zig/tests/cuda/ling_lanes_bytes.zig | 120 +++++++++++++++++++++ zig/tests/cuda/ling_lanes_rig.zig | 5 + zig/tests/cuda/ling_mtp_real.zig | 150 ++++++++++++++++++++++++++ zig/tests/cuda/main.zig | 4 +- zig/tests/cuda/oracle/ling_forward.py | 3 +- zig/tests/cuda/oracle/ling_mtp.py | 82 ++++++++++++++ 13 files changed, 626 insertions(+), 26 deletions(-) create mode 100644 zig/src/families/ling/mtp.zig create mode 100644 zig/tests/cuda/ling_mtp_real.zig create mode 100644 zig/tests/cuda/oracle/ling_mtp.py diff --git a/docs/recipes/ling-3.1.md b/docs/recipes/ling-3.1.md index 4b354da37f..ffe1686dd0 100644 --- a/docs/recipes/ling-3.1.md +++ b/docs/recipes/ling-3.1.md @@ -21,6 +21,7 @@ block FP8 (e4m3 in 128 x 128 blocks with `weight_scale_inv`). | Weight loader at TP=1: bf16, FP8 block, ModelOpt NVFP4 | `zig/src/families/ling/weights.zig`, `weights_proj.zig` | `ling-weights`, `ling-weights-real` | | Synthetic checkpoints in all three formats | `tf-cuda-test ling-synth` | its own name check | | The TP=1 layer graph: embedding, KDA/MLA mixers, dense MLP and MoE, final norm, head; windows and commits | `zig/src/families/ling/forward.zig`, `graph.zig`, `state.zig`, `block_*.zig`, `kernels/cuda/ling_glue.cu` | `ling-forward-synth`, `ling-layer-real` | +| The MTP layer (`model.layers.72`) as the drafter: `enorm(embed(next)) \| hnorm(hidden)` through `eh_proj`, MLA + MoE on its own cache, `final_layernorm`, the shared head; chains of up to 3 | `mtp.zig`, `Engine.draft` | `ling-lanes --mtp`, `ling-mtp-real` | | The engine and lane backend: a sequence per slot, prompt chunks cut anywhere, verify windows of up to 16 rows, shared windows of up to 8 streams (64 rows), kept prompt states, the server hooks | `engine.zig`, `lanes.zig`, `native.zig`; `native/cuda.zig` registers it | `ling-lanes` | | SwiGLU clamped in Ling's order, `min(silu(g), L) * clip(u, -L, L)`, in the glue and the FP8/NVFP4 grouped experts (every older kernel's SASS unchanged) | `glm_glue_moe.cuh` (`tf_ling_swiglu_*`), `experts.cuh` (epilogue 4), `experts.Layer.after` | `ling-clamp`, `ling-clamp-old`, `glm-route-old` | @@ -161,12 +162,41 @@ norm and head; KDA and dense MLP only, 6 GB) passes the same 90 exactness checks MLA + MoE layer adds 8 GB). Mutations caught: committing every row on `keep`, a restore that keeps the fresh KDA copy bits, chunks not cut at marks, shared draws keyed one position early, a lane commit that keeps one row too many. +## MTP drafting + +The HF model (`BailingMoeV3MTPLayer`) feeds the MTP layer the backbone's rows after `model.norm` and the next +token's embedding: `eh_proj(enorm(embed(t+1)) | hnorm(h_t))`, then input norm, MLA, residual, post-attention norm, +MoE, residual, `final_layernorm`, and `lm_head` reads that. The layer runs unclamped (the SwiGLU limit lists cover +the backbone). `mtp.Head` holds a view of the weights (the MTP layer alone, `final_layernorm` as its `norm`), its +own graph and the concatenated input rows; each sequence gets a second `State` (one MLA cache, its own position). +A prompt pass feeds every row whose next token is known; `Engine.draft` absorbs a verify's kept rows (their +final-norm rows and the tokens after them), then chains up to 3 drafts, each step's hidden row the MTP layer's own +`final_layernorm` output, as GLM's `decode.draft` does; the chain's rows are dropped before the next absorb. Kept +prompt states copy the MTP cache too. The server loads it unless `--no-drafts`; the lane facts then say `mtp`. + +Checked (`ling-lanes ... --mtp`, 104 checks a format; acceptance on random weights says nothing about quality and is +not reported as such): MTP-drafted replies, greedy and sampled, equal the plain ones; so do MTP-drafted passes +resumed at 53 and 211 and every concurrent case; a prompt pass's MTP cache in 37-row chunks equals the serial MTP +rows byte for byte; twins drafting chains of 3 and of 1 agree on tokens and first drafts over 6 windows, each +chain's second draft equals one computed by hand, and after one more window both states are equal byte for byte. +The MTP logits at every prompt row against `ling_forward.py --mtp` (float64 from the file, picks free): + +| Format | vs float64, rows routed alike | rows whose MTP picks differ | the bf16 points vs float64 on those rows | argmax | +| --- | --- | --- | --- | --- | +| bf16 | 1.72e-2 | 14 of 308 | 2.47e-2 | 293/308 | +| FP8 block | 1.65e-2 | 5 | 2.38e-2 | 300/308 | +| NVFP4 | 3.20e-2 | 11 | 7.81e-2 | 283/308 | + +The real layer (`ling-mtp-real`, fixture from `oracle/ling_mtp.py`: HF's `BailingMoeV3MTPLayer` in bf16 on 511 real +tokens' embeddings and `model.norm` of layer 0's output): `final_layernorm` rows at rel L2 7.28e-3 of HF on the 488 +rows whose top-8 agrees (4.56e-2 with the 23 others), worst row 1.61e-2, decode rows 6.74e-3; cut at 300 == uncut and +a 16-row window == 16 decode rows byte for byte; 8.4 GB peak, 153 s to load from the mount. Mutations caught: no +rollback of the chain's rows, `hnorm` with `enorm`'s weight, prompt rows absorbing their own token, kept states +without the MTP cache, the chain's hidden row taken from row 0. + ## What is missing - A full-checkpoint run: 574 GB in FP8 needs five Sparks and tensor parallelism. -- MTP drafting: the MTP layer's tensors load (`weights.Mtp`); the graph leaves it unwired. It would be - `enorm(embed(next)) | hnorm(hidden)` through `eh_proj`, one MLA + MoE layer on its own cache, the final norm and - the shared head. - Speed: nothing is tuned. Projections land in their columns through a byte copy, bf16 routed experts run densely, and the RMS norms take hidden and `q_lora_rank` widths 512, 1536 and 4096 only (`glue.rmsnorm`). - EXL3, through ashhart's EXL3 module (tracker #832). @@ -187,6 +217,9 @@ zig-out/bin/tf-cuda-test ling-forward-synth /tmp/ling-fp8 /tmp/fref zig-out/bin/tf-cuda-test ling-lanes /tmp/ling-fp8 /tmp/served # serving checks; writes .ids/.picks python zig/tests/cuda/oracle/ling_forward.py /tmp/ling-fp8 /tmp/sref --ids /tmp/served.ids --picks /tmp/served.picks zig-out/bin/tf-cuda-test ling-lanes /tmp/ling-fp8 /tmp/served /tmp/sref # and the served logits vs float64 +zig-out/bin/tf-cuda-test ling-lanes /tmp/ling-fp8 /tmp/m --mtp # the same with MTP drafts +python zig/tests/cuda/oracle/ling_forward.py /tmp/ling-fp8 /tmp/mref --ids /tmp/m.ids --picks /tmp/m.picks --mtp /tmp/mref.mtp +zig-out/bin/tf-cuda-test ling-lanes /tmp/ling-fp8 /tmp/m /tmp/mref --mtp # and the MTP logits vs float64 zig-out/bin/tf-cuda-test ling-clamp /tmp/ling-fp8 # Ling's SwiGLU order vs float64, L 10 and 2 zig-out/bin/tf-cuda-test ling-clamp-old /tmp/ling-fp8 OLD/fatbin # kept expert instances vs an older build's bytes ``` @@ -211,6 +244,7 @@ python zig/tests/cuda/oracle/ling_weights.py /path/to/Ling-3.1-flash /tmp/ling-r zig-out/bin/tf-cuda-test ling-weights-real /path/to/Ling-3.1-flash /tmp/ling-ref ``` +`ling-mtp-real DIR FIXTURE` runs the MTP layer (fixture from `oracle/ling_mtp.py SRC MODELING_PKG ROWS.pt OUT`). `ling-layer-real DIR FIXTURE` runs one whole layer (fixture from `oracle/ling_layer.py CHAIN L OUT`, the captured HF chain's `h.pt`, `h.pt`, `r.pt`). `ling-weights-real` loads layers 0, 4 and 7 with experts 0-3 (769 MiB on the device) and compares 35 products with torch's float64 on the same 8 rows: worst 0.414 of the bound, relative L2 at most 1.7e-3 for bf16 outputs and 6.9e-8 for the @@ -239,7 +273,7 @@ size. Each item stands alone. 1. **Recorder.** Give `Lanes` a `features` hook (the final-norm rows `Engine.verify` already keeps) once #848 lands. -2. **MTP.** Wire `weights.Mtp` as a drafter (see "What is missing"). +2. **MTP acceptance.** Measure the MTP layer's acceptance on real text once a whole model runs (TP). 3. **Fused projections.** Same-format q|k|v (and gate|up) rows as one weight per format, without the column copies. 4. **Bf16 routed experts.** Add a grouped bf16 expert kernel, or quantize at load, and drop the per-expert lanes. 5. **TP.** Cut `weights.load` by heads and widths as `glm/cuda_shard.zig` does; FP8 scale blocks span 128 rows and diff --git a/zig/src/families/ling/cuda.zig b/zig/src/families/ling/cuda.zig index f578ca89da..d00fa55f6e 100644 --- a/zig/src/families/ling/cuda.zig +++ b/zig/src/families/ling/cuda.zig @@ -8,6 +8,7 @@ pub const synth = @import("synth.zig"); pub const graph = @import("graph.zig"); pub const state = @import("state.zig"); pub const forward = @import("forward.zig"); +pub const mtp = @import("mtp.zig"); pub const engine = @import("engine.zig"); pub const lanes = @import("lanes.zig"); pub const native = @import("native.zig"); @@ -49,6 +50,7 @@ test { std.testing.refAllDecls(forward.Forward); std.testing.refAllDecls(engine.Engine); // the engine, its lane backend and its server hooks too std.testing.refAllDecls(lanes.Lanes); + std.testing.refAllDecls(mtp.Head); _ = native; } diff --git a/zig/src/families/ling/engine.zig b/zig/src/families/ling/engine.zig index f97f160371..8b4efd7df8 100644 --- a/zig/src/families/ling/engine.zig +++ b/zig/src/families/ling/engine.zig @@ -5,6 +5,7 @@ const glm = @import("glm_cuda"); const config = @import("config.zig"); const weights = @import("weights.zig"); const fw = @import("forward.zig"); +const head = @import("mtp.zig"); const sampler = glm.sampler; pub const Sampling = sampler.Sampling; @@ -16,10 +17,13 @@ pub const window_rows = fw.window_rows; /// A shared verify's most rows and streams. pub const shared_rows = 64; pub const max_shared = 8; +/// The MTP layer's chain at most. +pub const max_drafts = head.max_drafts; pub const Options = struct { capacity: usize, // tokens a sequence holds prompt_rows: usize = prompt_rows, // rows a prompt chunk takes at most + mtp: bool = false, // load the MTP layer, if the checkpoint has one, and draft with it layers: ?[]const u32 = null, // a thin model's backbone layers, ascending (checks on a real checkpoint); null: all }; @@ -36,14 +40,18 @@ pub const Marks = struct { at: []const u32, ctx: *anyopaque, call: *const fn (ct /// One chunk's outcome: where the pass stands, and the first token once the prompt is done. pub const Step = struct { at: usize, first: ?u32 }; -/// A stream's sequence: the layers' caches and its sampling rule. +/// A stream's sequence: the layers' caches, the MTP layer's, its sampling rule, the rows its next draft reads. const Seq = struct { st: fw.State, + hd: ?fw.State = null, // the MTP layer's MLA cache at its own position + hold: ?cuda.DeviceBuffer = null, // the prompt's last final-norm row: the first draft's hidden row sampling: ?Sampling = null, // null: greedy + in_head: usize = 0, // chained draft rows the MTP cache holds past the absorbed ones + rows_at: u64 = 0, // the last verify's final-norm rows (bf16 [rows][hidden]): the next draft's hidden rows }; /// A kept state: the sequence's arena copied whole, and the host fields beside it. -const Snap = struct { buf: cuda.DeviceBuffer, pos: usize, cur: []u1 }; +const Snap = struct { buf: cuda.DeviceBuffer, hd: ?cuda.DeviceBuffer, pos: usize, hd_pos: usize, cur: []u1 }; pub const Engine = struct { gpa: std.mem.Allocator, @@ -54,6 +62,7 @@ pub const Engine = struct { w: weights.Weights, // every loaded layer main: weights.Weights, // a view of w: the backbone layers alone g: fw.Graph, + mtp: ?head.Head = null, bp: fw.Buffers, // prompt chunks bd: fw.Buffers, // one stream's verify window bs: fw.Buffers, // shared verify windows @@ -77,6 +86,7 @@ pub const Engine = struct { }; defer opened.ck.close(); const c = opened.c; + const mtp = o.mtp and c.mtp > 0; const e = try gpa.create(Engine); errdefer gpa.destroy(e); e.* = .{ .gpa = gpa, .d = d, .c = c, .m = undefined, .stream = undefined, .w = undefined, .main = undefined, .g = undefined, .bp = undefined, .bd = undefined, .bs = undefined, .smp = undefined, .pairs = undefined, .capacity = o.capacity, .chunk = o.prompt_rows }; @@ -87,15 +97,25 @@ pub const Engine = struct { const layers = try backbone(gpa, &c, o.layers); defer gpa.free(layers); { + const all = try std.mem.concat(gpa, u32, &.{ layers, if (mtp) &.{c.layers} else &.{} }); + defer gpa.free(all); var pack = try cuda.modelopt.Packer.load(d); defer pack.unload(); - e.w = try weights.load(gpa, d, e.stream, &pack, &opened.ck, c, layers, .{}); + e.w = try weights.load(gpa, d, e.stream, &pack, &opened.ck, c, all, .{}); } errdefer e.w.deinit(); e.main = e.w; e.main.layers = e.w.layers[0..layers.len]; e.g = try fw.Graph.init(gpa, d, &e.m, e.stream, &e.main); errdefer e.g.deinit(); + if (mtp) { + e.mtp = @as(head.Head, undefined); + e.mtp.?.init(gpa, d, &e.m, e.stream, &e.w, o.prompt_rows, o.capacity) catch |err| { + e.mtp = null; + return err; + }; + } + errdefer if (e.mtp) |*h| h.deinit(); e.bp = try fw.Buffers.init(d, &e.m, &e.w, o.prompt_rows, true, o.capacity); errdefer e.bp.deinit(); e.bd = try fw.Buffers.init(d, &e.m, &e.w, window_rows, false, o.capacity); @@ -134,6 +154,7 @@ pub const Engine = struct { e.bs.deinit(); e.bd.deinit(); e.bp.deinit(); + if (e.mtp) |*h| h.deinit(); e.g.deinit(); e.w.deinit(); e.stream.deinit(); @@ -153,6 +174,11 @@ pub const Engine = struct { const q = try e.gpa.create(Seq); errdefer e.gpa.destroy(q); q.* = .{ .st = try fw.State.init(e.gpa, e.d, &e.g, e.capacity) }; + errdefer e.freeSeq(q); + if (e.mtp) |*h| { + q.hd = try fw.State.init(e.gpa, e.d, &h.g, e.capacity); + q.hold = try cuda.DeviceBuffer.alloc(e.d, @as(usize, e.c.hidden) * 2); + } e.seqs.items[slot] = q; } e.cur = e.seqs.items[slot].?; @@ -169,14 +195,24 @@ pub const Engine = struct { fn free(e: *Engine, slot: usize) void { const q = e.seqs.items[slot] orelse return; e.stream.synchronize() catch {}; - q.st.deinit(); + e.freeSeq(q); e.gpa.destroy(q); e.seqs.items[slot] = null; } - /// Bytes one more stream's sequence takes: its arena. + fn freeSeq(e: *Engine, q: *Seq) void { + _ = e; + q.st.deinit(); + if (q.hd) |*h| h.deinit(); + if (q.hold) |*b| b.free(); + q.hd = null; + q.hold = null; + } + + /// Bytes one more stream's sequence takes: its arenas. pub fn seqBytes(e: *const Engine) u64 { - return e.seqs.items[0].?.st.arena.buf.len; + const q = e.seqs.items[0].?; + return q.st.arena.buf.len + if (q.hd) |h| h.arena.buf.len + @as(usize, e.c.hidden) * 2 else 0; } /// The committed length: the pending token sits here, a window's row r draws position pos + 1 + r. @@ -187,7 +223,9 @@ pub const Engine = struct { /// A fresh sequence with `sampling` (null or temperature 0: greedy). pub fn begin(e: *Engine, sampling: ?Sampling) !void { e.cur.sampling = sampler.check(sampling); + e.cur.in_head = 0; try e.cur.st.reset(e.stream); + if (e.cur.hd) |*h| try h.reset(e.stream); } /// The whole prompt pass: chunks, then the first token from the last row; `from` restores a kept state. @@ -234,6 +272,12 @@ pub const Engine = struct { try e.draw(e.bp.logits, 1, len, &t); first = t[0]; } + if (e.mtp) |*h| { // each row whose next token is known through the MTP layer (it overwrites b.normed) + const q = e.cur; + if (end == len) try e.copy(q.hold.?.ptr, e.bp.normed + (rows - 1) * e.c.hidden * 2, e.c.hidden * 2); + const known = @min(rows, len - 1 - at); + if (known > 0) try h.absorb(&e.m, e.stream, &e.bp, &q.hd.?, @ptrCast(ids[at + 1 ..][0..known]), e.bp.normed); + } try f.commit(&e.cur.st, rows); return .{ .at = end, .first = first }; } @@ -244,6 +288,7 @@ pub const Engine = struct { const f = e.forward(&e.bd); const at = e.cur.st.pos; try f.run(&.{.{ .st = &e.cur.st, .row0 = 0, .rows = ids.len }}, @ptrCast(ids)); + e.cur.rows_at = e.bd.normed; try e.draw(e.bd.logits, ids.len, at + 1, out); } @@ -264,6 +309,7 @@ pub const Engine = struct { try e.forward(&e.bs).run(segs[0..parts.len], @ptrCast(ids[0..rows])); for (parts, segs[0..parts.len]) |p, sg| { try e.bind(p.slot); + e.cur.rows_at = e.bs.normed + sg.row0 * e.c.hidden * 2; try e.draw(e.bs.logits + sg.row0 * e.c.vocab * 2, sg.rows, sg.st.pos + 1, p.out); } } @@ -273,6 +319,36 @@ pub const Engine = struct { try e.forward(&e.bd).commit(&e.cur.st, n); } + /// The MTP layer absorbs the kept rows with `follow` (the token after each), then chains `depth` drafts. + pub fn draft(e: *Engine, from_prompt: bool, follow: []const u32, depth: usize, out: []u32) !void { + const h = if (e.mtp) |*x| x else return error.NoDraftHead; + if (depth > max_drafts or out.len < depth) return error.TooDeep; + if (follow.len == 0) return error.NoFollow; + const q = e.cur; + const hd = &q.hd.?; + const D: usize = e.c.hidden; + try hd.setPos(e.stream, hd.pos - q.in_head); // the last chain's rows go + q.in_head = 0; + const hidden = if (from_prompt) q.hold.?.ptr else q.rows_at; + var s0: usize = 0; + while (s0 < follow.len) : (s0 += window_rows) { + const part = follow[s0..@min(follow.len, s0 + window_rows)]; + try h.absorb(&e.m, e.stream, &h.bm, hd, @ptrCast(part), hidden + s0 * D * 2); + } + if (hd.pos != q.st.pos) return error.MtpOutOfStep; + const f = h.forward(&e.m, e.stream, &h.bm); + var last = (follow.len - 1) % window_rows; // h.bm's row of the newest absorbed token + for (0..depth) |j| { + try f.final(last + 1); + try f.logitsOf(last, 1); + try e.draw(h.bm.logits, 1, q.st.pos + 1 + j, out[j..][0..1]); + if (j + 1 == depth) break; + try h.absorb(&e.m, e.stream, &h.bm, hd, @ptrCast(out[j..][0..1]), h.bm.normed + last * D * 2); + q.in_head += 1; + last = 0; + } + } + /// `rows` logits rows -> the tokens at positions first, first + 1, ... pub fn draw(e: *Engine, logits: u64, rows: usize, first: u64, out: []u32) !void { if (e.cur.sampling) |s| { @@ -303,11 +379,18 @@ pub const Engine = struct { var buf = try cuda.DeviceBuffer.alloc(e.d, q.st.arena.buf.len); errdefer buf.free(); try e.copy(buf.ptr, q.st.arena.buf.ptr, buf.len); + var hd: ?cuda.DeviceBuffer = null; + errdefer if (hd) |*x| x.free(); + if (q.hd) |h| { + if (q.in_head != 0) return error.NotCommitted; + hd = try cuda.DeviceBuffer.alloc(e.d, h.arena.buf.len); + try e.copy(hd.?.ptr, h.arena.buf.ptr, h.arena.buf.len); + } try e.stream.synchronize(); const cur = try e.gpa.dupe(u1, q.st.cur); errdefer e.gpa.free(cur); const id = e.next_snap; - try e.snaps.put(e.gpa, id, .{ .buf = buf, .pos = q.st.pos, .cur = cur }); + try e.snaps.put(e.gpa, id, .{ .buf = buf, .hd = hd, .pos = q.st.pos, .hd_pos = if (q.hd) |h| h.pos else 0, .cur = cur }); e.next_snap += 1; return id; } @@ -321,6 +404,12 @@ pub const Engine = struct { q.st.pending = 0; q.st.prompt = false; try q.st.setPos(e.stream, s.pos); + if (q.hd) |*h| { + try e.copy(h.arena.buf.ptr, (s.hd orelse return error.NoKeptState).ptr, h.arena.buf.len); + h.pending = 0; + try h.setPos(e.stream, s.hd_pos); + } + q.in_head = 0; try e.stream.synchronize(); } @@ -335,6 +424,7 @@ pub const Engine = struct { fn freeSnap(e: *Engine, s: *Snap) void { s.buf.free(); + if (s.hd) |*x| x.free(); e.gpa.free(s.cur); } diff --git a/zig/src/families/ling/lanes.zig b/zig/src/families/ling/lanes.zig index 9cab5286fd..91ed5ce646 100644 --- a/zig/src/families/ling/lanes.zig +++ b/zig/src/families/ling/lanes.zig @@ -16,6 +16,8 @@ pub const max_rows = engine.window_rows; const Lane = struct { slot: u32, pending_rows: ?usize = null, + held: [engine.max_drafts]u32 = undefined, // the MTP layer's drafts for the next window + held_n: usize = 0, fill: ?usize = null, // a prompt pass going a chunk at a time: where its next chunk starts }; @@ -52,12 +54,16 @@ pub const Lanes = struct { break :vt v; }; - /// Round loop facts: exact windows of max_rows (copies), shared windows unless `alone`. + /// Round loop facts: exact windows of max_rows (copies, MTP chains), shared windows unless `alone`. pub fn facts(self: *const Lanes) lanes.Model { const sharing = !self.alone; + const drafting = self.e.mtp != null; return .{ .exact_width = max_rows, .gpu_tokens = false, + .mtp = drafting, + .speculate = drafting, + .drafts = engine.max_drafts, .speculate_early = false, .hidden_rows = sharing, .batch_rows = if (sharing) engine.shared_rows else 0, @@ -136,6 +142,7 @@ pub const Lanes = struct { const gop = try self.lanes.getOrPut(self.e.gpa, s); if (!gop.found_existing) gop.value_ptr.* = .{ .slot = self.freeSlot() }; gop.value_ptr.pending_rows = null; + gop.value_ptr.held_n = 0; gop.value_ptr.fill = null; try self.e.bind(gop.value_ptr.slot); try self.e.begin(s.sampling); @@ -186,15 +193,16 @@ pub const Lanes = struct { } /// A window's ids (pending token, then its drafts) checked against the bound sequence's position. - fn windowIds(self: *Lanes, w: be.Window, ids: *[max_rows]u32) ![]const u32 { + fn windowIds(self: *Lanes, l: *const Lane, w: be.Window, ids: *[max_rows]u32) ![]const u32 { if (w.parents != null) return error.TreesNotBuilt; const rows = w.rows(); if (rows > max_rows) return error.WindowTooWide; - if (w.held > 0) return error.NoHeldDrafts; + if (w.held > l.held_n) return error.NoHeldDrafts; const pos = self.e.pos(); for (w.positions, 0..) |p, r| if (p != pos + 1 + r) return error.PositionMismatch; + const drafts = if (w.held > 0) l.held[0..w.held] else w.tokens; ids[0] = w.pending; - @memcpy(ids[1..][0..w.tokens.len], w.tokens); + @memcpy(ids[1..][0..drafts.len], drafts); return ids[0..rows]; } @@ -205,7 +213,7 @@ pub const Lanes = struct { for (windows, out) |w, *o| { const l = try self.bind(w.stream, null); var ids: [max_rows]u32 = undefined; - const got = try self.windowIds(w, &ids); + const got = try self.windowIds(l, w, &ids); try self.e.verify(got, o.sampled[0..got.len]); @memcpy(o.drafts[0 .. got.len - 1], got[1..]); l.pending_rows = got.len; @@ -220,7 +228,7 @@ pub const Lanes = struct { var lane: [engine.max_shared]*Lane = undefined; for (windows, out[0..windows.len], 0..) |w, o, k| { lane[k] = try self.bind(w.stream, null); - const got = try self.windowIds(w, &ids[k]); + const got = try self.windowIds(lane[k], w, &ids[k]); parts[k] = .{ .slot = lane[k].slot, .ids = got, .out = o.sampled[0..got.len] }; } try self.e.verifyShared(parts[0..windows.len]); @@ -240,10 +248,27 @@ pub const Lanes = struct { } } + /// The MTP layer absorbs each stream's kept rows (its last verify committed first) and chains its drafts. fn draftFn(ptr: *anyopaque, requests: []const be.DraftRequest) anyerror!void { - _ = ptr; - _ = requests; - return error.NoDraftHead; + const self = of(ptr); + if (self.e.mtp == null) return error.NoDraftHead; + for (requests) |r| { + const l = try self.bind(r.stream, if (r.rows) |rows| rows.len else null); + if (r.position != self.e.pos() + 1) return error.PositionMismatch; + if (r.lanes != null) return error.TreesNotBuilt; + const depth = @min(r.depth, engine.max_drafts); + if (r.rows) |rows| { + for (rows, 0..) |row, i| if (row != i) return error.TreesNotBuilt; + try self.e.draft(false, r.follow[0..rows.len], depth, &l.held); + } else { + const first = switch (r.first orelse return error.NoFirstToken) { + .handle => |h| self.drawn[h % ring], + .value => |v| v, + }; + try self.e.draft(true, &.{first}, depth, &l.held); + } + l.held_n = depth; + } } /// The stream left the rounds: its sequence freed (slot 0 stays for the next stream). diff --git a/zig/src/families/ling/mtp.zig b/zig/src/families/ling/mtp.zig new file mode 100644 index 0000000000..a94a4a7d83 --- /dev/null +++ b/zig/src/families/ling/mtp.zig @@ -0,0 +1,70 @@ +//! Ling-3.1's MTP layer (model.layers.72) as the family's drafter: enorm | hnorm through eh_proj, MLA + MoE, its norm. +const std = @import("std"); +const cuda = @import("cuda"); +const weights = @import("weights.zig"); +const fw = @import("forward.zig"); + +const glue = @import("graph.zig").glue; + +/// The most drafts one chain holds. +pub const max_drafts = 3; + +/// The MTP layer's graph over a view of the weights (its layer alone, its final norm as `norm`) and its scratch. +pub const Head = struct { + view: weights.Weights, + mtp: weights.Mtp, + g: fw.Graph, + bm: fw.Buffers, // chained draft rows and a verify's kept rows + cat: cuda.DeviceBuffer, // bf16 [rows][2 hidden]: enorm(embed(next)) | hnorm(hidden) + + /// In place (the graph points at `view`); `w` holds the MTP layer last; cat rows for `rows` (the prompt chunk). + pub fn init(h: *Head, gpa: std.mem.Allocator, d: *const cuda.Driver, m: *const fw.Modules, s: cuda.Stream, w: *const weights.Weights, rows: usize, capacity: usize) !void { + const c = &w.config; + const mtp = w.mtp orelse return error.NoMtpLayer; + const last = w.layers[w.layers.len - 1]; + if (last.index != c.layers or last.mixer != .mla or last.ffn != .moe) return error.NoMtpLayer; + h.mtp = mtp; + h.view = w.*; + h.view.layers = w.layers[w.layers.len - 1 ..]; + h.view.norm = mtp.norm; // final_layernorm before the shared lm_head + h.g = try fw.Graph.init(gpa, d, m, s, &h.view); + errdefer h.g.deinit(); + h.bm = try fw.Buffers.init(d, m, w, fw.window_rows, false, capacity); + errdefer h.bm.deinit(); + h.cat = try cuda.DeviceBuffer.alloc(d, @max(rows, fw.window_rows) * 2 * @as(usize, c.hidden) * 2); + } + + pub fn deinit(h: *Head) void { + h.cat.free(); + h.bm.deinit(); + h.g.deinit(); + h.* = undefined; + } + + pub fn forward(h: *const Head, m: *const fw.Modules, s: cuda.Stream, b: *const fw.Buffers) fw.Forward { + return .{ .m = m, .g = &h.g, .s = s, .b = b }; + } + + /// The layer on b.x's `rows` embedding rows and `hidden`'s rows (bf16 [rows][hidden]), committed on `hd`. + pub fn step(h: *const Head, f: fw.Forward, hd: *fw.State, rows: usize, hidden: u64, picks: ?u64) !void { + const b = f.b; + const c = &h.view.config; + const D: usize = c.hidden; + const o = f.ops(); + if (rows == 0 or rows > b.rows or hd.pending != 0) return error.BadRows; + try glue.rmsnorm(o.set(), o.s, b.x, D, h.mtp.enorm, h.cat.ptr, 2 * D, null, c.eps, rows, D); + try glue.rmsnorm(o.set(), o.s, hidden, D, h.mtp.hnorm, h.cat.ptr + D * 2, 2 * D, null, c.eps, rows, D); + try o.mm(h.mtp.eh, h.cat.ptr, 2 * D, rows, b.x); + try f.layer(&.{.{ .st = hd, .row0 = 0, .rows = rows }}, 0, rows, picks); + hd.pending = rows; + hd.prompt = b.prompt; + try f.commit(hd, rows); + } + + /// `next` (the token after each row) with `hidden`'s rows through the layer on `b`, committed on `hd`. + pub fn absorb(h: *const Head, m: *const fw.Modules, s: cuda.Stream, b: *const fw.Buffers, hd: *fw.State, next: []const i32, hidden: u64) !void { + const f = h.forward(m, s, b); + try f.embed(next); + try h.step(f, hd, next.len, hidden, null); + } +}; diff --git a/zig/src/families/ling/native.zig b/zig/src/families/ling/native.zig index f26780d967..bb8bdfe31c 100644 --- a/zig/src/families/ling/native.zig +++ b/zig/src/families/ling/native.zig @@ -45,7 +45,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ const own = try gpa.create(Owned); errdefer gpa.destroy(own); own.gpa = gpa; - own.e = try engine.Engine.init(gpa, io, ctx.d, ctx, dir, .{ .capacity = capacity(o.context) }); + own.e = try engine.Engine.init(gpa, io, ctx.d, ctx, dir, .{ .capacity = capacity(o.context), .mtp = o.drafts }); errdefer own.e.deinit(); own.lanes = .{ .e = own.e, .whole = off("TF_LING_PROMPT_STEPS"), .alone = off("TF_LING_SHARED") }; return .{ diff --git a/zig/tests/cuda/ling_lanes.zig b/zig/tests/cuda/ling_lanes.zig index 4bac636973..1040288ce8 100644 --- a/zig/tests/cuda/ling_lanes.zig +++ b/zig/tests/cuda/ling_lanes.zig @@ -29,6 +29,7 @@ pub fn run(gpu: check.Gpu, rest: []const [:0]const u8) !void { const out = rest[1]; var ref: ?[]const u8 = null; var thin: ?[]u32 = null; + var mtp = false; var i: usize = 2; while (i < rest.len) : (i += 1) { if (std.mem.eql(u8, rest[i], "--layers") and i + 1 < rest.len) { @@ -37,9 +38,11 @@ pub fn run(gpu: check.Gpu, rest: []const [:0]const u8) !void { while (it.next()) |x| try list.append(a, try std.fmt.parseInt(u32, x, 10)); thin = list.items; i += 1; + } else if (std.mem.eql(u8, rest[i], "--mtp")) { + mtp = true; } else ref = rest[i]; } - const e = try Engine.init(gpu.gpa, gpu.io, gpu.d, gpu.ctx, dir, .{ .capacity = 1024, .prompt_rows = 512, .layers = thin }); + const e = try Engine.init(gpu.gpa, gpu.io, gpu.d, gpu.ctx, dir, .{ .capacity = 1024, .prompt_rows = 512, .layers = thin, .mtp = mtp }); defer e.deinit(); e.chunk = chunk; var ll: ling.lanes.Lanes = .{ .e = e }; @@ -51,7 +54,7 @@ pub fn run(gpu: check.Gpu, rest: []const [:0]const u8) !void { defer le.deinit(); var c: Cover = .{ .gpa = gpu.gpa, .a = a, .le = &le }; const V: usize = e.c.vocab; - std.debug.print("RESULT ling-lanes {t}: {d} layers loaded of {d}, vocabulary {d}, prompt chunks of {d} rows\n", .{ e.c.quant, e.main.layers.len, e.c.layers, V, chunk }); + std.debug.print("RESULT ling-lanes {t}: {d} layers loaded of {d}, vocabulary {d}, prompt chunks of {d} rows{s}\n", .{ e.c.quant, e.main.layers.len, e.c.layers, V, chunk, if (e.mtp != null) ", drafts from the MTP layer" else "" }); const p = try walk(a, 150, V, 0); // W, what every contract below holds to: the plain greedy pass, its prompt in 37-row chunks @@ -83,6 +86,22 @@ pub fn run(gpu: check.Gpu, rest: []const [:0]const u8) !void { const differ = std.mem.indexOfDiff(u32, plain_s.ids, g); try c.expect(differ != null, "the sampled reply leaves the greedy one at token {?d} (sampling draws)", .{differ}); + if (e.mtp != null) { // the MTP layer's drafts against the plain replies (acceptance on random weights says nothing) + for ([_]Cover.Got{ plain_g, plain_s }, [_]?lanes.Sampling{ null, keyed }) |plain, sampling| { + const got = try c.run(.{ .prompt = copy_prompt, .max_new = 200, .sampling = sampling }); + try c.same(if (sampling != null) "MTP drafts, sampled, against plain" else "MTP drafts, greedy, against plain", got, plain.ids, .length); + try c.expect(got.drafted > 0, "MTP drafts{s}: {d} drafted in {d} rounds, {d} kept", .{ if (sampling != null) ", sampled" else "", got.drafted, got.rounds, got.accepted }); + } + var keeper: rig.Keeper = .{ .e = e }; + defer keeper.drop(); + const marks = [_]u32{ 53, 211 }; + _ = try c.run(.{ .prompt = copy_prompt, .max_new = 8, .reuse = .{ .marks = &marks, .hook = keeper.hook() } }); + for (keeper.kept[0..keeper.n], keeper.at[0..keeper.n]) |kept, at| { + const got = try c.run(.{ .prompt = copy_prompt, .max_new = 200, .reuse = .{ .saved = kept, .at = at } }); + try c.same(try std.fmt.allocPrint(a, "MTP drafts resumed at {d}, against plain", .{at}), got, plain_g.ids, .length); + } + } + // resumed == fresh: states kept at marks off the chunk grid and off multiples of 16, then passes resumed there { var keeper: rig.Keeper = .{ .e = e }; @@ -166,7 +185,7 @@ fn concurrent(c: *Cover, ll: *ling.lanes.Lanes, a: std.mem.Allocator, V: usize) if (j.discard_at != null) continue; try c.same(try std.fmt.allocPrint(a, "{s}: stream {d} (prompt {d}{s}) against its solo run", .{ cs.name, k, j.p, if (j.sampled) ", sampled" else "" }), .{ .ids = ids, .reason = .length, .rounds = 0, .cached = 0, .drafted = 0, .accepted = 0 }, if (j.sampled) solo_s[j.p] else solo[j.p], .length); } - var drafting = false; + var drafting = ll.e.mtp != null; // every joined stream drafts with the MTP layer when it is loaded for (cs.joins) |j| drafting = drafting or j.drafts; try c.expect(got.most >= 2 and (got.drafted > 0) == drafting, "{s}: {d} streams at most in one round, {d} rounds with two or more, {d} drafts", .{ cs.name, got.most, got.shared, got.drafted }); for (cs.joins) |j| if (j.filled) { diff --git a/zig/tests/cuda/ling_lanes_bytes.zig b/zig/tests/cuda/ling_lanes_bytes.zig index 85e0147765..d03e9b560f 100644 --- a/zig/tests/cuda/ling_lanes_bytes.zig +++ b/zig/tests/cuda/ling_lanes_bytes.zig @@ -85,9 +85,66 @@ pub fn engineChecks(c: *Cover, e: *Engine, prompt: []const u32, w: []const u32) try c.expect(same_rows, "a 9-row window and, after keeping 4, a 5-row window: tokens and logits == 9 single rows", .{}); try rig.sameSeq(c, e, "9-row window, keep 4, 5-row window == 9 single rows", 10, 11); try sharedChecks(c, e, prompt, w); + if (e.mtp != null) try mtpChain(c, e, prompt, w); for ([_]u32{ 10, 11, 12, 13, 14 }) |s| e.release(s); } +/// MTP chains of 3 and of 1 on twins: the same first draft, and after the same windows the chain's rows rolled back. +fn mtpChain(c: *Cover, e: *Engine, prompt: []const u32, w: []const u32) !void { + var t: [2]u32 = undefined; + var d: [2][3]u32 = undefined; + for (0..2) |k| { + t[k] = try pass(e, @intCast(40 + k), prompt[0..200], 37, null, null); + try e.draft(true, &.{t[k]}, if (k == 0) 3 else 1, &d[k]); + } + var ok = t[0] == t[1] and d[0][0] == d[1][0]; + var pend = t[0]; + var chained: usize = 0; + for (0..6) |round| { + var ids: [8]u32 = undefined; + var out: [2][8]u32 = undefined; + const rows = 2 + round % 6; + ids[0] = pend; + for (1..rows) |r| ids[r] = w[(round * 11 + r) % w.len]; + const keep = 1 + (round * 3) % rows; + for (0..2) |k| { + try e.bind(@intCast(40 + k)); + try e.verify(ids[0..rows], out[k][0..rows]); + try e.keep(keep); + try e.draft(false, out[k][0..keep], if (k == 0) 3 else 1, &d[k]); + } + ok = ok and std.mem.eql(u32, out[0][0..rows], out[1][0..rows]) and d[0][0] == d[1][0]; + { // the chain's second draft by hand on the twin: d0 with the newest kept row's MTP output as its hidden row + const h = &e.mtp.?; + const hd = &e.seqs.items[41].?.hd.?; + const D: usize = e.c.hidden; + const f = h.forward(&e.m, e.stream, &h.bm); + try h.absorb(&e.m, e.stream, &h.bm, hd, @ptrCast(d[1][0..1]), h.bm.normed + (keep - 1) * D * 2); + try f.final(1); + try f.logitsOf(0, 1); + var second: [1]u32 = undefined; + try e.draw(h.bm.logits, 1, hd.pos + 1, &second); + try hd.setPos(e.stream, hd.pos - 1); + ok = ok and second[0] == d[0][1]; + chained += 1; + } + pend = out[0][keep - 1]; + } + try c.expect(ok, "MTP chains of 3 and of 1 over 6 windows (kept 1 to 7 rows): the same tokens and first drafts; {d} second drafts == by hand", .{chained}); + const sa = e.seqs.items[40].?; + const sb = e.seqs.items[41].?; + try c.expect(sa.hd.?.pos == sa.st.pos + 2 and sb.hd.?.pos == sb.st.pos, "the chain of 3 holds its 2 chained rows past the kept ones (MTP at {d}, the sequence at {d})", .{ sa.hd.?.pos, sa.st.pos }); + for (0..2) |k| { // one more window, both drafting 1: the chain's rows rolled back, the twins byte for byte + var out: [2]u32 = undefined; + try e.bind(@intCast(40 + k)); + try e.verify(&.{ pend, w[3] }, &out); + try e.keep(2); + try e.draft(false, &out, 1, &d[k]); + } + try rig.sameSeq(c, e, "after chains of 3 and of 1, a window and a chain of 1 each", 40, 41); + for (0..2) |k| e.release(@intCast(40 + k)); +} + /// Two streams' windows in one shared forward against lone twins over 8 rounds: logits, tokens, then states. fn sharedChecks(c: *Cover, e: *Engine, prompt: []const u32, w: []const u32) !void { const a = c.a; @@ -194,5 +251,68 @@ pub fn serial(c: *Cover, gpu: check.Gpu, e: *Engine, p: []const u32, w: []const } const base = @sqrt(bd / bw); std.debug.print("RESULT served logits vs float64 (same picks): rel L2 {e:.3}, worst row {e:.3}, argmax {d}/{d}; vs bf16 points {e:.3}; bf16 points vs float64 {e:.3}\n", .{ d64.rel, d64.worst_row, d64.argmax, T, dbf.rel, base }); + if (e.mtp != null) try mtpLogits(c, gpu, e, ids, path); try c.expect(nf == 0 and d64.rel <= 1.5 * base and dbf.rel <= base, "served logits within 1.5x the bf16 points' rel L2 of float64 ({e:.3} <= 1.5 x {e:.3}), within it of the bf16 points", .{ d64.rel, base }); } + +/// The MTP layer's logits at every row but the last (hidden row t, token t + 1) against ling_forward.py --mtp REF.mtp. +fn mtpLogits(c: *Cover, gpu: check.Gpu, e: *Engine, ids: []const i32, ref_path: []const u8) !void { + const a = c.a; + var name: [512]u8 = undefined; + const bytes = try std.Io.Dir.cwd().readFileAlloc(gpu.io, try std.fmt.bufPrint(&name, "{s}.mtp", .{ref_path}), a, .limited(1 << 30)); + if (bytes.len < 20 or !std.mem.eql(u8, bytes[0..8], "TFLINGM2")) return error.BadFixture; + const dims = std.mem.bytesAsSlice(u32, bytes[8..20]); + const R: usize = dims[0]; + const V: usize = dims[1]; + const K: usize = dims[2]; + const T = ids.len; + try c.expect(R == T - 1 and V == e.c.vocab and K == e.c.topk, "MTP reference: {d} rows", .{R}); + const m64 = try a.alloc(f64, R * V); + const mbf = try a.alloc(f64, R * V); + const theirs = try a.alloc(i32, R * K); + var at: usize = 20; + for ([_][]u8{ std.mem.sliceAsBytes(m64), std.mem.sliceAsBytes(mbf), std.mem.sliceAsBytes(theirs) }) |dst| { + @memcpy(dst, bytes[at..][0..dst.len]); + at += dst.len; + } + const h = &e.mtp.?; + var r: frig.Rig = .{ .gpu = gpu, .a = a, .s = e.stream, .m = &e.m, .g = &e.g }; + defer r.deinit(); + try e.bind(31); + try e.begin(null); + const q = e.seqs.items[31].?; + const main = e.forward(&e.bp); + const seg = [_]F.Seg{.{ .st = &q.st, .row0 = 0, .rows = T }}; + try main.embed(ids); + try main.stack(&seg); + try main.final(T); // e.bp.normed: the hidden rows, as a prompt pass leaves them + const pk = try r.zeros(R * (K + 1) * 4); + const f = h.forward(&e.m, e.stream, &e.bp); + try f.embed(ids[1..]); + try h.step(f, &q.hd.?, R, e.bp.normed, pk); + try f.final(R); + try f.logitsOf(0, R); + const raw = try r.back(e.bp.logits, R * V * 2); + const got = try a.alloc(u16, R * V); + @memcpy(std.mem.sliceAsBytes(got), raw); + try main.commit(&q.st, T); + _ = try pass(e, 32, @ptrCast(ids), 37, null, null); // a prompt pass in 37-row chunks: its MTP rows the same + try rig.sameSeq(c, e, "a prompt pass's MTP rows in 37-row chunks == the serial MTP rows", 32, 31); + e.release(31); + e.release(32); + const ours = try frig.sortedPicks(a, try r.back(pk, R * (K + 1) * 4), 1, R, K); + const flipped = try a.alloc(bool, R); + @memset(flipped, false); + const nf = frig.flips(ours, theirs, 1, R, 0, R, K, flipped); + const d64 = frig.distance(got, m64, R, V, flipped); + const dbf = frig.distance(got, mbf, R, V, flipped); + var bd: f64 = 0; + var bw: f64 = 0; + for (0..R) |t| if (!flipped[t]) for (mbf[t * V ..][0..V], m64[t * V ..][0..V]) |x, y| { + bd += (x - y) * (x - y); + bw += y * y; + }; + const base = @sqrt(bd / bw); // the bf16 points' own distance on the same rows (their routing may differ there too) + std.debug.print("RESULT MTP logits vs float64: rel L2 {e:.3} ({e:.3} without {d} rows whose picks differ), argmax {d}/{d}; vs bf16 points {e:.3}; bf16 points vs float64 on those rows {e:.3}\n", .{ d64.rel, d64.rel_kept, nf, d64.argmax, R, dbf.rel, base }); + try c.expect(d64.rel_kept <= 1.5 * base, "MTP logits within 1.5x the bf16 points' rel L2 of float64 on rows routed alike ({e:.3} <= 1.5 x {e:.3})", .{ d64.rel_kept, base }); +} diff --git a/zig/tests/cuda/ling_lanes_rig.zig b/zig/tests/cuda/ling_lanes_rig.zig index 9c9dc4c0c6..e54b4bcecc 100644 --- a/zig/tests/cuda/ling_lanes_rig.zig +++ b/zig/tests/cuda/ling_lanes_rig.zig @@ -214,5 +214,10 @@ pub fn sameSeq(c: *Cover, e: *Engine, what: []const u8, x: u32, y: u32) !void { for (0..p.mla_layers) |mi| if (!std.mem.eql(u8, try back(e, c.a, p.cacheAt(mi), p.pos * ling.forward.cache_width * 2), try back(e, c.a, q.cacheAt(mi), p.pos * ling.forward.cache_width * 2))) { bad = "MLA cache rows"; }; + if (e.seqs.items[x].?.hd) |*ph| { + const qh = &e.seqs.items[y].?.hd.?; + if (ph.pos != qh.pos) bad = "MTP positions"; + if (!std.mem.eql(u8, try back(e, c.a, ph.cacheAt(0), ph.pos * ling.forward.cache_width * 2), try back(e, c.a, qh.cacheAt(0), ph.pos * ling.forward.cache_width * 2))) bad = "MTP cache rows"; + } try c.expect(bad == null, "{s}: slots {d} and {d} at {d}, byte for byte{s}{s}", .{ what, x, y, p.pos, if (bad != null) ", but differ in " else "", bad orelse "" }); } diff --git a/zig/tests/cuda/ling_mtp_real.zig b/zig/tests/cuda/ling_mtp_real.zig new file mode 100644 index 0000000000..92dcfc9352 --- /dev/null +++ b/zig/tests/cuda/ling_mtp_real.zig @@ -0,0 +1,150 @@ +//! `ling-mtp-real DIR FIXTURE`: the real MTP layer (model.layers.72) vs HF's BailingMoeV3MTPLayer; exactness on it. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const ling = @import("ling_cuda"); +const rig = @import("ling_forward_rig.zig"); + +const F = rig.F; + +fn take(comptime T: type, a: std.mem.Allocator, bytes: []const u8, at: *usize, n: usize) ![]T { + const out = try a.alloc(T, n); + @memcpy(std.mem.sliceAsBytes(out), bytes[at.*..][0 .. n * @sizeOf(T)]); + at.* += n * @sizeOf(T); + return out; +} + +const Kit = struct { + r: *rig.Rig, + h: *const ling.mtp.Head, + emb: []const u16, + hid: []const u16, + hd_dev: u64, // the hidden rows on the device + D: usize, + + /// Rows lo .. hi through the MTP layer on `hd` (prompt or decode buffers), committed; its final-norm rows. + fn pass(k: Kit, b: *const F.Buffers, hd: *F.State, lo: usize, hi: usize, picks: ?u64) ![]u16 { + const d = k.r.gpu.d; + const n = hi - lo; + try d.check(d.api.cuMemcpyHtoDAsync_v2(b.x, k.emb[lo * k.D ..].ptr, n * k.D * 2, k.r.s.handle), "cuMemcpyHtoDAsync"); + const f = k.h.forward(k.r.m, k.r.s, b); + try k.h.step(f, hd, n, k.hd_dev + lo * k.D * 2, picks); + try f.final(n); + const out = try k.r.a.alloc(u16, n * k.D); + @memcpy(std.mem.sliceAsBytes(out), try k.r.back(b.normed, n * k.D * 2)); + return out; + } +}; + +pub fn run(gpu: check.Gpu, dir: []const u8, path: []const u8) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + const bytes = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, a, .limited(1 << 30)); + if (bytes.len < 20 or !std.mem.eql(u8, bytes[0..8], "TFLINGM1")) return error.BadFixture; + var at: usize = 8; + const dims = try take(u32, a, bytes, &at, 3); + const T: usize = dims[0]; + const D: usize = dims[1]; + const K: usize = dims[2]; + const emb = try take(u16, a, bytes, &at, T * D); + const hid = try take(u16, a, bytes, &at, T * D); + const want = try take(f32, a, bytes, &at, T * D); + const topk = try take(i32, a, bytes, &at, T * K); + if (at != bytes.len) return error.BadFixture; + var why: ling.config.Why = .{}; + var o = rig.W.open(gpu.gpa, gpu.io, dir, &why) catch |e| { + std.debug.print("{s}: {s}\n", .{ dir, why.text() }); + return e; + }; + defer o.ck.close(); + var m = try F.Modules.load(gpu.d, gpu.ctx); + defer m.deinit(); + var pack = try cuda.modelopt.Packer.load(gpu.d); + defer pack.unload(); + const L = o.c.layers; + const t0 = check.now(gpu.io); + var w = try rig.W.load(gpu.gpa, gpu.d, s, &pack, &o.ck, o.c, &.{L}, .{ .globals = false }); + defer w.deinit(); + const load_s = @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e9; + const cap = T + 32; + var h: ling.mtp.Head = undefined; + try h.init(gpu.gpa, gpu.d, &m, s, &w, T, cap); + defer h.deinit(); + const c = &w.config; + try check.expect(D == c.hidden and K == c.topk, "fixture: {d} rows {d} wide, top {d}", .{ T, D, K }); + var r: rig.Rig = .{ .gpu = gpu, .a = a, .s = s, .m = &m, .g = &h.g }; + defer r.deinit(); + var pb = try F.Buffers.init(gpu.d, &m, &w, T, true, cap); + defer pb.deinit(); + var db = try F.Buffers.init(gpu.d, &m, &w, F.window_rows, false, cap); + defer db.deinit(); + const hd_dev = try r.zeros(T * D * 2); + try gpu.d.check(gpu.d.api.cuMemcpyHtoDAsync_v2(hd_dev, hid.ptr, T * D * 2, s.handle), "cuMemcpyHtoDAsync"); + const k: Kit = .{ .r = &r, .h = &h, .emb = emb, .hid = hid, .hd_dev = hd_dev, .D = D }; + // the whole prompt against HF, picks against HF's + const pk = try r.zeros(T * (K + 1) * 4); + const st = try r.state(cap); + defer st.deinit(); + const got = try k.pass(&pb, st, 0, T, pk); + const ours = try rig.sortedPicks(a, try r.back(pk, T * (K + 1) * 4), 1, T, K); + const theirs = try a.dupe(i32, topk); + for (0..T) |t| std.mem.sort(i32, theirs[t * K ..][0..K], {}, std.sort.asc(i32)); + const skip = try a.alloc(bool, T); + @memset(skip, false); + const nf = rig.flips(ours, theirs, 1, T, 0, T, K, skip); + const dv = dist(got, want, T, D, skip); + std.debug.print("RESULT ling-mtp-real layer {d} ({t}): {d} rows, loaded in {d:.1} s\n", .{ L, c.quant, T, load_s }); + std.debug.print("RESULT final-norm rows vs HF: rel L2 {e:.3}, {e:.3} without the {d} rows whose top-{d} differs, worst kept row {e:.3}\n", .{ dv[0], dv[1], nf, K, dv[2] }); + // cut prompt, decode rows, a verify window + const sc = try r.state(cap); + defer sc.deinit(); + const cut0 = try k.pass(&pb, sc, 0, 300, null); + const cut1 = try k.pass(&pb, sc, 300, T, null); + try check.sameBytes("MTP rows cut at 300 == uncut", std.mem.sliceAsBytes(cut0), std.mem.sliceAsBytes(got[0 .. 300 * D])); + try check.sameBytes("MTP rows cut at 300 == uncut (the rest)", std.mem.sliceAsBytes(cut1), std.mem.sliceAsBytes(got[300 * D ..])); + const P = 400; + const N = F.window_rows; + const sd = try r.state(cap); + defer sd.deinit(); + const sw = try r.state(cap); + defer sw.deinit(); + _ = try k.pass(&pb, sd, 0, P, null); + _ = try k.pass(&pb, sw, 0, P, null); + const dec = try a.alloc(u16, N * D); + for (0..N) |i| @memcpy(dec[i * D ..][0..D], try k.pass(&db, sd, P + i, P + i + 1, null)); + try check.sameBytes("16-row MTP window == 16 decode rows", std.mem.sliceAsBytes(try k.pass(&db, sw, P, P + N, null)), std.mem.sliceAsBytes(dec)); + const dd = dist(dec, want[P * D ..], N, D, skip[P..]); + std.debug.print("RESULT decode rows {d}-{d}: rel L2 {e:.3} vs HF\n", .{ P, P + N - 1, dd[0] }); + try check.expect(dv[1] <= 1e-2 and dd[1] <= 1e-2, "the MTP layer within rel L2 1e-2 of HF (rows whose picks agree)", .{}); + check.pass("ling-mtp-real: rel L2 {e:.3} of HF ({d} flipped rows apart); cut == uncut, window == decode rows; peak {d} MiB on the device", .{ dv[1], nf, cuda.usage(false).peak >> 20 }); +} + +/// Relative L2 of bf16 rows against fp32 rows, all and without `skip` rows, and the worst kept row's. +fn dist(got: []const u16, want: []const f32, rows: usize, D: usize, skip: []const bool) [3]f64 { + var dd: f64 = 0; + var ww: f64 = 0; + var kd: f64 = 0; + var kw: f64 = 0; + var worst: f64 = 0; + for (0..rows) |t| { + var rd: f64 = 0; + var rw: f64 = 0; + for (0..D) |j| { + const e: f64 = want[t * D + j]; + const g = rig.bfOf(got[t * D + j]); + rd += (g - e) * (g - e); + rw += e * e; + } + dd += rd; + ww += rw; + if (!skip[t]) { + kd += rd; + kw += rw; + worst = @max(worst, @sqrt(rd / rw)); + } + } + return .{ @sqrt(dd / ww), @sqrt(kd / kw), worst }; +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 5533ea1488..368c1e6f82 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -116,7 +116,8 @@ const usage = \\ ling-layer-real one real Ling-3.1 layer vs HF's captured output (oracle/ling_layer.py); exactness \\ ling-clamp Ling's SwiGLU order min(silu(g), L): glue and the dir's grouped experts vs float64 \\ ling-clamp-old the grouped expert instances vs an older build's fatbins, byte for byte - \\ ling-lanes [ref] [--layers a,b] Ling's lane backend: chunked, drafted, resumed, concurrent == plain + \\ ling-lanes [ref] [--layers a,b] [--mtp] Ling's lane backend: chunked, drafted, resumed, concurrent == plain + \\ ling-mtp-real the real MTP layer vs HF's (oracle/ling_mtp.py); cut, window == decode rows \\ ling-weights-real real Ling-3.1 layers 0, 4, 7 (4 experts) vs float64 and torch (oracle/ling_weights.py) \\ ; @@ -229,6 +230,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "ling-layer-real")) return @import("ling_layer_real.zig").run(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "ling-clamp")) return @import("ling_clamp.zig").run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "ling-clamp-old")) return @import("ling_clamp.zig").old(gpu, try arg(rest, 0), try arg(rest, 1)); + if (std.mem.eql(u8, cmd, "ling-mtp-real")) return @import("ling_mtp_real.zig").run(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "ling-lanes")) return @import("ling_lanes.zig").run(gpu, rest); if (std.mem.eql(u8, cmd, "ling-weights-real")) return ling_real.run(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "glm-lanes-concurrent")) return glm_lanes_tests.concurrentCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); diff --git a/zig/tests/cuda/oracle/ling_forward.py b/zig/tests/cuda/oracle/ling_forward.py index 5e76e09176..d2507a39cf 100644 --- a/zig/tests/cuda/oracle/ling_forward.py +++ b/zig/tests/cuda/oracle/ling_forward.py @@ -136,7 +136,8 @@ def moe(self, x, p, i): self.gap = max(self.gap, (kth - choice.gather(1, idx).min(-1)[0]).max().item()) w = s.gather(1, idx) w = w / (w.sum(-1, keepdim=True) + 1e-20) * c["routed_scaling_factor"] - le, ls = c["expert_swiglu_limit_list"][i], c["share_expert_swiglu_limit_list"][i] + lists = c["expert_swiglu_limit_list"], c["share_expert_swiglu_limit_list"] + le, ls = (x[i] if i < len(x) else 0 for x in lists) # the lists cover the backbone: the MTP layer runs unclamped y = torch.zeros_like(x) for t in range(x.shape[0]): for j in range(top): diff --git a/zig/tests/cuda/oracle/ling_mtp.py b/zig/tests/cuda/oracle/ling_mtp.py new file mode 100644 index 0000000000..dcc0a5f1c9 --- /dev/null +++ b/zig/tests/cuda/oracle/ling_mtp.py @@ -0,0 +1,82 @@ +"""For `tf-cuda-test ling-mtp-real`: HF's MTP layer (BailingMoeV3MTPLayer) on real rows, as a flat file.""" +import argparse +import importlib +import json +import os +import struct +import sys + +import torch +from safetensors import safe_open + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("src", help="the checkpoint: config.json, model.safetensors.index.json, its shards (FP8 blocks)") + p.add_argument("modeling", help="a package folder (with __init__.py) holding the HF modeling and configuration files") + p.add_argument("rows", help="a .pt with 'ids' [1][T] and 'out' [1][T][D] (a layer's output, e.g. ref0.pt)") + p.add_argument("out", help="the fixture: embeddings, hidden rows (bf16), HF's output (fp32), HF's top-k (i32)") + a = p.parse_args() + pkg = os.path.basename(os.path.normpath(a.modeling)) + sys.path.insert(0, os.path.dirname(os.path.normpath(a.modeling))) + import transformers.utils.import_utils as iu + if not hasattr(iu, "is_torch_fx_available"): # newer transformers dropped it; the modeling only imports it + iu.is_torch_fx_available = lambda: False + cfg_mod = importlib.import_module(pkg + ".configuration_bailing_moe_v3") + mod = importlib.import_module(pkg + ".modeling_bailing_moe_v3") + raw = json.load(open(a.src + "/config.json")) + cfg = cfg_mod.BailingMoeV3Config(**{k: v for k, v in raw.items() if k not in ("auto_map", "architectures", "quantization_config")}) + cfg._attn_implementation = "eager" + if raw.get("rope_scaling") is None: # newer transformers fill in a default dict the modeling reads as YaRN's + cfg.rope_scaling = None + wmap = json.load(open(a.src + "/model.safetensors.index.json"))["weight_map"] + files = {} + + def get(k): + fn = wmap[k] + if fn not in files: + files[fn] = safe_open(a.src + "/" + fn, "pt", device="cuda") + w = files[fn].get_tensor(k) + if w.dtype == torch.float8_e4m3fn: # 128 x 128 blocks of weight_scale_inv + s = get(k + "_scale_inv").float() + w = (w.float() * s.repeat_interleave(128, 0)[: w.shape[0]].repeat_interleave(128, 1)[:, : w.shape[1]]).to(torch.bfloat16) + return w + + L = cfg.num_hidden_layers + with torch.device("meta"): + layer = mod.BailingMoeV3MTPLayer(cfg, L) + sd = {} + for name, prm in list(layer.named_parameters()) + list(layer.named_buffers()): + t = get(f"model.layers.{L}." + name) + sd[name] = t.float() if prm.dtype == torch.float32 else t.to(torch.bfloat16) + layer.load_state_dict(sd, strict=True, assign=True) + layer = layer.cuda().eval() + src = torch.load(a.rows) + ids, h = src["ids"][0].cuda(), src["out"].cuda().to(torch.bfloat16) + T = ids.shape[0] - 1 # row t: the hidden row t and the embedding of token t + 1 + emb = get("model.word_embeddings.weight")[ids[1:]][None] + norm = get("model.norm.weight") + hf = h[:, :T].float() + hidden = (hf * torch.rsqrt(hf.pow(2).mean(-1, keepdim=True) + cfg.rms_norm_eps)).to(torch.bfloat16) * norm # model.norm's rows + pos = torch.arange(T, device="cuda")[None] + rd = cfg.qk_rope_head_dim # BailingMoeV3RotaryEmbedding's default rope (its init table moved in newer transformers) + inv = 1.0 / cfg.rope_theta ** (torch.arange(0, rd, 2, device="cuda").float() / rd) + ang = torch.cat([pos[0, :, None].float() * inv] * 2, -1)[None] + cos_sin = (ang.cos().to(torch.bfloat16), ang.sin().to(torch.bfloat16)) + mask = torch.full((T, T), torch.finfo(torch.bfloat16).min, device="cuda").triu(1)[None, None].to(torch.bfloat16) + picks = [] + layer.mlp.gate.register_forward_hook(lambda m, i, o: picks.append(o[0])) + with torch.no_grad(): + out = layer(emb, hidden, attention_mask=mask, position_ids=pos, position_embeddings=cos_sin)[0] + topk = picks[0].view(T, -1).int().cpu() + D = emb.shape[-1] + with open(a.out, "wb") as f: + f.write(b"TFLINGM1" + struct.pack("<3I", T, D, topk.shape[1])) + for t in (emb[0], hidden[0]): + f.write(t.cpu().view(torch.int16).numpy().tobytes()) + f.write(out[0].float().cpu().numpy().tobytes() + topk.numpy().tobytes()) + print("mtp layer", L, "tokens", T, "rms", out.float().pow(2).mean().sqrt().item(), "finite", torch.isfinite(out).all().item()) + + +if __name__ == "__main__": + main() From 7bd7a40a3f4aefe2243008d27cbbbc0eed0e59fc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 15 Oct 2026 15:17:07 +0200 Subject: [PATCH 12/34] tests: ling-pack times one MoE layer's expert load (read, pack, cold) and checks its bytes against the host packers; ling-layer-real can dump its rows --- zig/tests/cuda/ling_layer_real.zig | 10 +- zig/tests/cuda/ling_pack.zig | 160 +++++++++++++++++++++++++++++ zig/tests/cuda/main.zig | 6 +- 3 files changed, 172 insertions(+), 4 deletions(-) create mode 100644 zig/tests/cuda/ling_pack.zig diff --git a/zig/tests/cuda/ling_layer_real.zig b/zig/tests/cuda/ling_layer_real.zig index 1c83490a2f..1de72f6e28 100644 --- a/zig/tests/cuda/ling_layer_real.zig +++ b/zig/tests/cuda/ling_layer_real.zig @@ -1,4 +1,4 @@ -//! `ling-layer-real DIR FIXTURE`: one real Ling-3.1 layer through the graph vs HF's captured output; exactness on it. +//! `ling-layer-real DIR FIXTURE [DUMP]`: one real Ling-3.1 layer through the graph vs HF's output; exactness on it. const std = @import("std"); const cuda = @import("cuda"); const check = @import("check.zig"); @@ -56,7 +56,7 @@ fn dist(got: []const u16, want: []const f32, rows: usize, D: usize, skip: []cons } /// `ling-layer-real DIR FIXTURE` (oracle/ling_layer.py's file for layer L of the checkpoint at DIR). -pub fn run(gpu: check.Gpu, dir: []const u8, path: []const u8) !void { +pub fn run(gpu: check.Gpu, dir: []const u8, path: []const u8, dump: ?[]const u8) !void { var arena = std.heap.ArenaAllocator.init(gpu.gpa); defer arena.deinit(); const a = arena.allocator(); @@ -138,6 +138,12 @@ pub fn run(gpu: check.Gpu, dir: []const u8, path: []const u8) !void { const dd = dist(dec, want[P * D ..], N, D, skip[P..]); const dp = dist(dec, @as([]const f32, try widen(a, got[P * D .. (P + N) * D])), N, D, skip[P..]); std.debug.print("RESULT decode rows {d}-{d}: rel L2 {e:.3} vs HF, {e:.3} vs the prompt's rows\n", .{ P, P + N - 1, dd[0], dp[0] }); + if (dump) |d| { // the prompt's and the decode rows' bytes, for comparing two builds + const all = try a.alloc(u16, got.len + dec.len); + @memcpy(all[0..got.len], got); + @memcpy(all[got.len..], dec); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = d, .data = std.mem.sliceAsBytes(all) }); + } try check.expect(dv[1] <= 1e-2 and dd[1] <= 1e-2, "layer {d} within rel L2 1e-2 of HF (rows whose picks agree)", .{L}); check.pass("ling-layer-real layer {d}: rel L2 {e:.3} of HF ({d} flipped rows apart); cut == uncut, window == decode rows; peak {d} MiB on the device", .{ L, dv[1], nf, cuda.usage(false).peak >> 20 }); } diff --git a/zig/tests/cuda/ling_pack.zig b/zig/tests/cuda/ling_pack.zig new file mode 100644 index 0000000000..12d75ee844 --- /dev/null +++ b/zig/tests/cuda/ling_pack.zig @@ -0,0 +1,160 @@ +//! `ling-pack DIR LAYER [EXPERTS]`: a MoE layer's experts loaded and timed, bytes vs the host packers the loader used. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const ling = @import("ling_cuda"); +const check = @import("check.zig"); +const lw = @import("ling_weights_tests.zig"); + +const W = ling.weights; +const X8 = cuda.experts.Fp8Experts; +const X4 = cuda.experts.Nvfp4Experts; + +/// Host memory from /proc/self/status: the peak resident set since the last reset, and anonymous memory now (KiB). +const Host = struct { + hwm: usize = 0, + anon: usize = 0, + + fn read(io: std.Io) Host { + var buf: [8192]u8 = undefined; + const text = std.Io.Dir.cwd().readFile(io, "/proc/self/status", &buf) catch return .{}; + var h: Host = .{}; + var lines = std.mem.tokenizeScalar(u8, text, '\n'); + while (lines.next()) |line| { + inline for (.{ .{ "VmHWM:", "hwm" }, .{ "RssAnon:", "anon" } }) |kv| if (std.mem.startsWith(u8, line, kv[0])) { + var it = std.mem.tokenizeAny(u8, line[kv[0].len..], " \tkB"); + @field(h, kv[1]) = std.fmt.parseInt(usize, it.next() orelse "0", 10) catch 0; + }; + } + return h; + } + + /// Resets VmHWM to the current resident set (proc(5) clear_refs 5). + fn reset(io: std.Io) void { + std.Io.Dir.cwd().writeFile(io, .{ .sub_path = "/proc/self/clear_refs", .data = "5" }) catch {}; + } +}; + +fn secs(io: std.Io, t0: i96) f64 { + return @as(f64, @floatFromInt(check.now(io) - t0)) / 1e9; +} + +/// Every byte of layer i's routed experts 0..count read once from the mapping (the read alone, no device work). +fn touch(ck: *core.Checkpoint, i: u32, count: usize) !usize { + var sum: u64 = 0; + var bytes: usize = 0; + var b: [64]u8 = undefined; + const prefix = try std.fmt.bufPrint(&b, "model.layers.{d}.mlp.experts.", .{i}); + for (ck.files.items) |*f| for (f.names.keys()) |k| { + if (!std.mem.startsWith(u8, k, prefix)) continue; + const rest = k[prefix.len..]; + const e = std.fmt.parseInt(usize, rest[0 .. std.mem.indexOfScalar(u8, rest, '.') orelse continue], 10) catch continue; + if (e >= count) continue; + const t = f.get(k).?; + for (std.mem.bytesAsSlice(u64, @as([]align(1) const u8, t.bytes[0 .. t.bytes.len / 8 * 8]))) |w| sum +%= w; + bytes += t.bytes.len; + }; + std.mem.doNotOptimizeAway(sum); + return bytes; +} + +/// One expert's packed device words against the host packers' bytes from the same file tensors. +fn compare(gpu: check.Gpu, ck: *core.Checkpoint, i: u32, l: cuda.experts.Layer, a: std.mem.Allocator) !void { + const ni = l.width; + const d = l.dims; + const nvfp4 = l.format == .nvfp4; + const up_n = if (nvfp4) X4.bytes(1, ni, d, 2) / 4 else 2 * ni * d / 4; + const down_n = if (nvfp4) X4.bytes(1, d, ni, 1) / 4 else d * ni / 4; + const inv_n = (ni / 128) * (d / 128); + const hu = try a.alloc(u32, up_n); + const hd = try a.alloc(u32, down_n); + const gu = try a.alloc(u32, up_n); + const gd = try a.alloc(u32, down_n); + const scale_up = try a.alloc(u8, l.experts * (if (nvfp4) 2 * 4 else 2 * inv_n * 4)); + const scale_down = try a.alloc(u8, l.experts * (if (nvfp4) 4 else inv_n * 4)); + try (cuda.DeviceBuffer{ .d = gpu.d, .ptr = l.up_scale, .len = scale_up.len }).download(0, scale_up); + try (cuda.DeviceBuffer{ .d = gpu.d, .ptr = l.down_scale, .len = scale_down.len }).download(0, scale_down); + var b: [3][192]u8 = undefined; + for (0..l.experts) |e| { + var m: [3][]const u8 = undefined; + for ([_][]const u8{ "gate_proj", "up_proj", "down_proj" }, 0..) |p, j| m[j] = try std.fmt.bufPrint(&b[j], "model.layers.{d}.mlp.experts.{d}.{s}", .{ i, e, p }); + if (nvfp4) { + var x: [3]core.nvfp4.Linear = undefined; + for (&x, m) |*y, name| y.* = try core.nvfp4.Linear.read(ck, name); + X4.packGateUp(hu, .{ x[0].codes, x[0].scales }, .{ x[1].codes, x[1].scales }, ni, d); + X4.packOne(hd, x[2].codes, x[2].scales, d, ni); + const globals = [3]f32{ x[0].global, x[1].global, x[2].global }; + try check.sameBytes("nvfp4 gate and up globals", scale_up[8 * e ..][0..8], std.mem.sliceAsBytes(globals[0..2])); + try check.sameBytes("nvfp4 down global", scale_down[4 * e ..][0..4], std.mem.sliceAsBytes(globals[2..3])); + } else { + var t: [3]core.checkpoint.Tensor = undefined; + var s: [3]core.checkpoint.Tensor = undefined; + for (&t, &s, m) |*w, *sc, name| { + var nb: [200]u8 = undefined; + w.* = try ck.get(try std.fmt.bufPrint(&nb, "{s}.weight", .{name})); + sc.* = try ck.get(try std.fmt.bufPrint(&nb, "{s}.weight_scale_inv", .{name})); + } + try X8.packGateUp(a, hu, t[0].bytes, t[1].bytes, ni, d); + X8.packOne(hd, t[2].bytes, d, ni); + try check.sameBytes("fp8 gate scale_inv", scale_up[(2 * e) * inv_n * 4 ..][0 .. inv_n * 4], s[0].bytes); + try check.sameBytes("fp8 up scale_inv", scale_up[(2 * e + 1) * inv_n * 4 ..][0 .. inv_n * 4], s[1].bytes); + try check.sameBytes("fp8 down scale_inv", scale_down[e * inv_n * 4 ..][0 .. inv_n * 4], s[2].bytes); + } + try (cuda.DeviceBuffer{ .d = gpu.d, .ptr = l.up, .len = l.experts * up_n * 4 }).download(e * up_n * 4, std.mem.sliceAsBytes(gu)); + try (cuda.DeviceBuffer{ .d = gpu.d, .ptr = l.down, .len = l.experts * down_n * 4 }).download(e * down_n * 4, std.mem.sliceAsBytes(gd)); + var nb: [64]u8 = undefined; + try check.sameBytes(try std.fmt.bufPrint(&nb, "expert {d} gate-up words", .{e}), std.mem.sliceAsBytes(gu), std.mem.sliceAsBytes(hu)); + try check.sameBytes(try std.fmt.bufPrint(&nb, "expert {d} down words", .{e}), std.mem.sliceAsBytes(gd), std.mem.sliceAsBytes(hd)); + } +} + +pub fn run(gpu: check.Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 2) return error.MissingArgument; + const dir = rest[0]; + const i = try std.fmt.parseInt(u32, rest[1], 10); + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const io = gpu.io; + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var kit: lw.Kit = undefined; + try kit.init(gpu); + defer kit.deinit(); + var why: ling.config.Why = .{}; + var o = W.open(gpu.gpa, io, dir, &why) catch |e| { + std.debug.print("{s}: {s}\n", .{ dir, why.text() }); + return e; + }; + defer o.ck.close(); + if (!o.c.isMoe(i)) return error.NotMoeLayer; + const count = if (rest.len > 2) try std.fmt.parseInt(usize, rest[2], 10) else o.c.experts; + const opts: W.Options = .{ .globals = false, .experts = count }; + // the read alone from a cold cache, then a load on hot pages (pack and upload), then a cold load (both) + o.ck.release(); + var t0 = check.now(io); + const bytes = try touch(&o.ck, i, count); + const read_s = secs(io, t0); + var hot_s: f64 = 0; + var cold_s: f64 = 0; + var peak: [2]usize = .{ 0, 0 }; + var host: [2]Host = undefined; + for (0..2) |pass| { + if (pass == 1) o.ck.release(); + Host.reset(io); + const base = cuda.usage(true).device; + t0 = check.now(io); + var w = try W.load(gpu.gpa, gpu.d, s, &kit.pack, &o.ck, o.c, &.{i}, opts); + defer w.deinit(); + if (pass == 0) hot_s = secs(io, t0) else cold_s = secs(io, t0); + peak[pass] = cuda.usage(false).peak - base; + host[pass] = Host.read(io); + if (pass == 1) { + const l = w.layer(i).?.ffn.moe.routed.grouped; + try compare(gpu, &o.ck, i, l, arena.allocator()); + std.debug.print("RESULT layer {d}: {d} experts {t}, {d:.2} GiB of expert tensors\n", .{ i, count, l.format, @as(f64, @floatFromInt(bytes)) / (1 << 30) }); + } + } + std.debug.print("RESULT read alone {d:.2} s ({d:.0} MiB/s); load on hot pages {d:.2} s; cold load {d:.2} s\n", .{ read_s, @as(f64, @floatFromInt(bytes >> 20)) / read_s, hot_s, cold_s }); + std.debug.print("RESULT device peak above the start {d} MiB (hot), {d} MiB (cold); host VmHWM {d} MiB (hot), {d} MiB (cold), RssAnon {d} MiB\n", .{ peak[0] >> 20, peak[1] >> 20, host[0].hwm >> 10, host[1].hwm >> 10, host[1].anon >> 10 }); + check.pass("ling-pack layer {d}: every expert's packed words and scales equal the host packers' bytes", .{i}); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 368c1e6f82..2233fdc502 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -113,12 +113,13 @@ const usage = \\ ling-synth [--format bf16|fp8|nvfp4] [--layers N] a tiny Ling-3.1 checkpoint with the real names and config schema \\ ling-weights every layer of a Ling checkpoint loaded: projections vs float64, small tensors byte for byte \\ ling-forward-synth [picks] a ling-synth checkpoint's forward vs oracle/ling_forward.py; cuts, windows, batches - \\ ling-layer-real one real Ling-3.1 layer vs HF's captured output (oracle/ling_layer.py); exactness + \\ ling-layer-real [dump] one real Ling-3.1 layer vs HF's captured output (oracle/ling_layer.py); exactness \\ ling-clamp Ling's SwiGLU order min(silu(g), L): glue and the dir's grouped experts vs float64 \\ ling-clamp-old the grouped expert instances vs an older build's fatbins, byte for byte \\ ling-lanes [ref] [--layers a,b] [--mtp] Ling's lane backend: chunked, drafted, resumed, concurrent == plain \\ ling-mtp-real the real MTP layer vs HF's (oracle/ling_mtp.py); cut, window == decode rows \\ ling-weights-real real Ling-3.1 layers 0, 4, 7 (4 experts) vs float64 and torch (oracle/ling_weights.py) + \\ ling-pack [experts] one MoE layer's experts: read, pack and cold load timed; bytes vs the host packers \\ ; @@ -227,12 +228,13 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "ling-synth")) return ling_synth.run(gpu, rest); if (std.mem.eql(u8, cmd, "ling-weights")) return ling_weights.run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "ling-forward-synth")) return @import("ling_forward_synth.zig").run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) rest[2] else null); - if (std.mem.eql(u8, cmd, "ling-layer-real")) return @import("ling_layer_real.zig").run(gpu, try arg(rest, 0), try arg(rest, 1)); + if (std.mem.eql(u8, cmd, "ling-layer-real")) return @import("ling_layer_real.zig").run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) rest[2] else null); if (std.mem.eql(u8, cmd, "ling-clamp")) return @import("ling_clamp.zig").run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "ling-clamp-old")) return @import("ling_clamp.zig").old(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "ling-mtp-real")) return @import("ling_mtp_real.zig").run(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "ling-lanes")) return @import("ling_lanes.zig").run(gpu, rest); if (std.mem.eql(u8, cmd, "ling-weights-real")) return ling_real.run(gpu, try arg(rest, 0), try arg(rest, 1)); + if (std.mem.eql(u8, cmd, "ling-pack")) return @import("ling_pack.zig").run(gpu, rest); if (std.mem.eql(u8, cmd, "glm-lanes-concurrent")) return glm_lanes_tests.concurrentCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); std.debug.print("{s}", .{usage}); return error.UnknownCommand; From bbe360f9df38280b271ff3dd58758845ca784c1a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 15 Oct 2026 16:09:41 +0200 Subject: [PATCH 13/34] ling: routed experts staged in batches through pinned memory and packed on the GPU (Packer.fp8Experts, fp4Experts), the file's pages dropped per batch --- docs/recipes/ling-3.1.md | 22 +++- zig/kernels/cuda/modelopt_pack.cu | 9 ++ zig/src/cuda/modelopt.zig | 15 ++- zig/src/families/ling/weights.zig | 3 +- zig/src/families/ling/weights_proj.zig | 155 +++++++++++++++++-------- zig/tests/cuda/ling_pack.zig | 4 +- zig/tests/cuda/main.zig | 2 +- 7 files changed, 155 insertions(+), 55 deletions(-) diff --git a/docs/recipes/ling-3.1.md b/docs/recipes/ling-3.1.md index ffe1686dd0..05505341ca 100644 --- a/docs/recipes/ling-3.1.md +++ b/docs/recipes/ling-3.1.md @@ -79,7 +79,8 @@ exllamav3's port measured 0.4-1.1e-3 without flips (4 flips each on layers 4 and activations, the captured HF chain carries an fp32 residual from layer 3's output on, and this graph rounds the stream to bf16 as the bf16 model does. On every real layer a prompt cut at 300 equals the uncut prompt and a 16-row window equals 16 decode rows; for FP8 decode rows -equal the prompt's rows (both run the lane). Loading one MoE layer's 512 experts from the sshfs mount takes ~150 s. +equal the prompt's rows (both run the lane). Loading one MoE layer's 512 experts takes 6.1 s from +local disk and 122 s from the sshfs mount (its read alone is 92 s); see Loading below. ### Formats @@ -97,9 +98,26 @@ The format of every tensor comes from the checkpoint, never from the layer graph A file tensor whose dtype or shape disagrees with the format the config gives it is refused, naming the module and both formats. Each projection loads as a `qlinear.Weight` (`bf16`, `fp8g` or `nvfp4`) through main's packers (`modelopt.Packer`, `fp8.packScales`); routed experts load as `experts.Layer` (`fp8g` or `nvfp4`, packed by -`Fp8Experts`/`Nvfp4Experts`). There is no grouped bf16 expert kernel, so bf16 routed experts load as one lane weight +`Fp8Experts`/`Nvfp4Experts`'s layouts on the GPU: `Packer.fp8Experts` and `Packer.fp4Experts`). There is no grouped bf16 expert kernel, so bf16 routed experts load as one lane weight per projection. The router must be bf16 (GLM's router kernel reads bf16); a quantized router is refused. +### Loading + +Routed experts load in batches of `weights.Options.stage` bytes (512 MiB, 32 FP8 experts): eight threads copy the batch's +raw tensors from the mapped file into pinned memory, one upload moves them to a device stage of the same size, and the +packers lay them out in the layer's buffers; the batch's file pages are then dropped. Scales are collected on the host +and uploaded once a layer. `ling-pack DIR LAYER [EXPERTS [STAGE]]` times it and checks every packed byte against the +host packers (`Fp8Experts`/`Nvfp4Experts.packGateUp`, `packOne`). Real layer 4 (512 FP8 experts, 7.5 GiB), from local +disk / the sshfs mount: + +| Loader | cold load | host peak (VmHWM) | device peak | +| --- | --- | --- | --- | +| host packers per expert (c48be499) | 14.3 s / ~157 s | 8.0 GiB (the layer's mapped pages) | 7.9 GiB | +| batches packed on the GPU | 6.1 s / 122 s | 1.3 GiB | 8.4 GiB (+ one stage) | + +From local disk the load runs at the read's rate (5-7 s to read the layer once single-threaded); from the mount the +network bounds it (86 MiB/s here: ~1.9 h for the 572 GB release). + ### The SwiGLU clamp Layers 64-71 clamp the shared expert and layers 65-71 the routed experts at L = 10 as inclusionAI's vLLM diff --git a/zig/kernels/cuda/modelopt_pack.cu b/zig/kernels/cuda/modelopt_pack.cu index 000c834358..c7df5e1ea7 100644 --- a/zig/kernels/cuda/modelopt_pack.cu +++ b/zig/kernels/cuda/modelopt_pack.cu @@ -106,3 +106,12 @@ extern "C" __global__ void __launch_bounds__(144) tf_pack_fp4_experts(const uint dst[w] = word; } } + +// Experts' e4m3 [E][n][k] as part m of [E][n/32][k/32][parts][32][2][4] words (Fp8Experts.packOne's order). +extern "C" __global__ void __launch_bounds__(256) tf_pack_fp8_experts(const uint32_t* __restrict__ src, uint32_t* __restrict__ out, + int n, int k, int parts, int m) { + const int g = blockIdx.x, cb = blockIdx.y, e = blockIdx.z, w = threadIdx.x, kg = k / 32; + const int j = w % 4, h = (w / 4) % 2, t = (w / 8) % 4, gq = w / 32; + const long long row = static_cast(e) * n + cb * 32 + j * 8 + gq; + out[(((static_cast(e) * (n / 32) + cb) * kg + g) * parts + m) * 256 + w] = src[row * (k / 4) + g * 8 + h * 4 + t]; +} diff --git a/zig/src/cuda/modelopt.zig b/zig/src/cuda/modelopt.zig index 4338ef4db9..c489b7ede0 100644 --- a/zig/src/cuda/modelopt.zig +++ b/zig/src/cuda/modelopt.zig @@ -16,7 +16,7 @@ const Stream = stream_.Stream; /// Rows padded to whole 128-row blocks, as every lane format pads them. pub const padded = nvfp4.padded; -/// modelopt_pack.cu's kernels; each matches a host packer (fp8.packCodes, qmmf.packBf16, nvfp4.pack*, Nvfp4Experts). +/// modelopt_pack.cu's kernels; each matches a host packer (fp8.packCodes, qmmf.packBf16, nvfp4.pack*, the experts'). pub const Packer = struct { mod: Module, fp8_lane: Function, @@ -24,6 +24,7 @@ pub const Packer = struct { fp4_words: Function, fp4_scales: Function, fp4_experts: Function, + fp8_experts: Function, pub fn load(d: *const Driver) !Packer { if (!kernels.available) return error.BuiltWithoutKernels; @@ -36,6 +37,7 @@ pub const Packer = struct { .fp4_words = try mod.function("tf_pack_fp4_words"), .fp4_scales = try mod.function("tf_pack_fp4_scales"), .fp4_experts = try mod.function("tf_pack_fp4_experts"), + .fp8_experts = try mod.function("tf_pack_fp8_experts"), }; } @@ -79,6 +81,17 @@ pub const Packer = struct { for ([_]usize{ n, ks, k0, k, parts, m }) |v| a.add(@as(i32, @intCast(v))); try launch_.launch(p.fp4_experts, .{ .grid = .{ .x = @intCast(k / 32), .y = @intCast(n / 32), .z = @intCast(count) }, .block = .{ .x = 144 } }, s, &a); } + + /// `count` experts' e4m3 [E][n][k] as part m of `parts` per (column block, k32 group) (Fp8Experts.packGateUp). + pub fn fp8Experts(p: *const Packer, s: Stream, src: u64, out: u64, count: usize, n: usize, k: usize, parts: usize, m: usize) !void { + if (n % 32 != 0 or k % 32 != 0 or m >= parts) return error.UnexpectedTensor; + if (count == 0) return; + var a: launch_.Args = .{}; + a.add(src); + a.add(out); + for ([_]usize{ n, k, parts, m }) |v| a.add(@as(i32, @intCast(v))); + try launch_.launch(p.fp8_experts, .{ .grid = .{ .x = @intCast(k / 32), .y = @intCast(n / 32), .z = @intCast(count) }, .block = .{ .x = 256 } }, s, &a); + } }; /// The lane matmul's view of a packed per-tensor FP8 projection. diff --git a/zig/src/families/ling/weights.zig b/zig/src/families/ling/weights.zig index 32ab240336..f2339a276d 100644 --- a/zig/src/families/ling/weights.zig +++ b/zig/src/families/ling/weights.zig @@ -63,6 +63,7 @@ pub const Mtp = struct { enorm: u64, hnorm: u64, eh: Proj, norm: u64 }; pub const Options = struct { globals: bool = true, // embedding, final norm and head experts: ?usize = null, // the first n routed experts of each MoE layer only (tests on a partial checkpoint) + stage: usize = 512 << 20, // raw expert bytes staged at a time (pinned and on the device), at least one expert }; pub const Weights = struct { @@ -200,7 +201,7 @@ const Loader = struct { pub fn load(gpa: std.mem.Allocator, d: *const cuda.Driver, stream: cuda.Stream, pack: *const cuda.modelopt.Packer, ck: *core.Checkpoint, c: Config, layers: []const u32, opts: Options) !Weights { var w: Weights = .{ .config = c, .store = .{ .gpa = gpa, .d = d } }; errdefer w.store.deinit(); - var L: Loader = .{ .r = .{ .gpa = gpa, .stream = stream, .pack = pack, .ck = ck, .c = &w.config, .store = &w.store }, .c = &w.config, .opts = opts }; + var L: Loader = .{ .r = .{ .gpa = gpa, .stream = stream, .pack = pack, .ck = ck, .c = &w.config, .store = &w.store, .stage = opts.stage }, .c = &w.config, .opts = opts }; const out = try gpa.alloc(Layer, layers.len); errdefer gpa.free(out); for (layers, out) |i, *l| { diff --git a/zig/src/families/ling/weights_proj.zig b/zig/src/families/ling/weights_proj.zig index e0e372261c..3a3404015b 100644 --- a/zig/src/families/ling/weights_proj.zig +++ b/zig/src/families/ling/weights_proj.zig @@ -87,6 +87,7 @@ pub const Reader = struct { ck: *core.Checkpoint, c: *const config.Config, store: *Store, + stage: usize = 512 << 20, // raw expert bytes staged at a time (weights.Options.stage) buf: [256]u8 = undefined, fn name(r: *Reader, module: []const u8, tail: []const u8) ![]const u8 { @@ -243,75 +244,133 @@ pub const Reader = struct { return std.fmt.bufPrint(b, "model.layers.{d}.mlp.experts.{d}." ++ proj, .{ i, e }); } - /// A view of `ptr` the store owns, for uploads at an offset. - fn view(r: *Reader, ptr: u64, len: usize) cuda.DeviceBuffer { - return .{ .d = r.store.d, .ptr = ptr, .len = len }; - } - fn label(b: []u8, i: u32, what: []const u8) ![]const u8 { return std.fmt.bufPrint(b, "model.layers.{d}.mlp.experts.{s}", .{ i, what }); } - /// experts.Layer's fp8g form: Fp8Experts.packGateUp and packOne per expert, scale_inv blocks as stored. + /// Experts a stage holds when one expert's raw tensors take `one` bytes. + fn batch(r: *const Reader, count: usize, one: usize) usize { + return @max(1, @min(count, r.stage / one)); + } + + /// A batch's raw tensors: copied from the file into pinned memory by `threads` threads, uploaded, packed there. + const Stage = struct { + host: cuda.HostBuffer, + dev: cuda.DeviceBuffer, + jobs: std.ArrayList(Job) = .empty, + + const Job = struct { at: usize, bytes: []const u8 }; + const threads = 8; // page faults in flight: one reader leaves the disk (or the mount) idle between them + + fn init(r: *Reader, len: usize) !Stage { + var host = try cuda.HostBuffer.alloc(r.store.d, len); + errdefer host.free(); + return .{ .host = host, .dev = try cuda.DeviceBuffer.alloc(r.store.d, len) }; + } + + fn deinit(st: *Stage, r: *Reader) void { + st.jobs.deinit(r.gpa); + st.host.free(); + st.dev.free(); + } + + fn put(st: *Stage, r: *Reader, at: usize, bytes: []const u8) !void { + try st.jobs.append(r.gpa, .{ .at = at, .bytes = bytes }); + } + + fn copy(st: *const Stage, first: usize) void { + var j = first; + while (j < st.jobs.items.len) : (j += threads) { + const job = st.jobs.items[j]; + @memcpy(st.host.bytes[job.at..][0..job.bytes.len], job.bytes); + } + } + + /// The batch on the device; the packers run after it on `r.stream`. + fn upload(st: *Stage, r: *Reader) !void { + var pool: [threads - 1]?std.Thread = @splat(null); + for (&pool, 1..) |*t, first| t.* = std.Thread.spawn(.{}, copy, .{ st, first }) catch null; + for (pool, 1..) |t, first| if (t == null) st.copy(first); // no thread: this one copies its share + st.copy(0); + for (pool) |t| if (t) |x| x.join(); + st.jobs.clearRetainingCapacity(); + try st.dev.uploadAsync(0, st.host.bytes, r.stream.handle); + } + + /// The packers done (the stage is reused) and the batch's file pages dropped: the GPU shares this memory. + fn drained(_: *Stage, r: *Reader) !void { + try r.stream.synchronize(); + r.ck.release(); + } + }; + + /// experts.Layer's fp8g form: batches of e4m3 bytes staged on the device, packed there (Packer.fp8Experts). fn fp8Experts(r: *Reader, i: u32, count: usize, ni: usize, d: usize, limit: f32) !cuda.experts.Layer { if (ni % X8.block != 0 or d % X8.block != 0) return error.UnexpectedTensor; - const up_n = 2 * ni * d / 4; // u32 words of one expert's gate and up - const down_n = d * ni / 4; + const one = ni * d; // bytes of one e4m3 matrix const inv_n = (ni / 128) * (d / 128); // one matrix's scale_inv blocks var lb: [96]u8 = undefined; - const up = try r.store.alloc(count * up_n * 4, try label(&lb, i, "up")); - const down = try r.store.alloc(count * down_n * 4, try label(&lb, i, "down")); - const up_scale = try r.store.alloc(count * 2 * inv_n * 4, try label(&lb, i, "up_scale")); - const down_scale = try r.store.alloc(count * inv_n * 4, try label(&lb, i, "down_scale")); - const hu = try r.gpa.alloc(u32, up_n); - defer r.gpa.free(hu); - const hd = try r.gpa.alloc(u32, down_n); - defer r.gpa.free(hd); + const up = try r.store.alloc(count * 2 * one, try label(&lb, i, "up")); + const down = try r.store.alloc(count * one, try label(&lb, i, "down")); + const inv = try r.gpa.alloc(f32, 3 * count * inv_n); // [E][gate, up] then [E][down], as stored + defer r.gpa.free(inv); + const nb = r.batch(count, 3 * one); + var stage = try Stage.init(r, 3 * nb * one); // [gate, up, down][nb][n][k] + defer stage.deinit(r); var b: [192]u8 = undefined; - for (0..count) |e| { - const g = try r.fp8Pair(try expertModule(&b, i, e, "gate_proj"), ni, d); - const u = try r.fp8Pair(try expertModule(&b, i, e, "up_proj"), ni, d); - const dn = try r.fp8Pair(try expertModule(&b, i, e, "down_proj"), d, ni); - try X8.packGateUp(r.gpa, hu, g.w.bytes, u.w.bytes, ni, d); - X8.packOne(hd, dn.w.bytes, d, ni); - try r.view(up, count * up_n * 4).upload(e * up_n * 4, std.mem.sliceAsBytes(hu)); - try r.view(down, count * down_n * 4).upload(e * down_n * 4, std.mem.sliceAsBytes(hd)); - const us = r.view(up_scale, count * 2 * inv_n * 4); - try us.upload(2 * e * inv_n * 4, g.s.bytes); - try us.upload((2 * e + 1) * inv_n * 4, u.s.bytes); - try r.view(down_scale, count * inv_n * 4).upload(e * inv_n * 4, dn.s.bytes); + var e0: usize = 0; + while (e0 < count) : (e0 += nb) { + const n = @min(nb, count - e0); + for (0..n) |j| inline for (.{ "gate_proj", "up_proj", "down_proj" }, 0..) |proj, m| { + const e = e0 + j; + const x = if (m < 2) try r.fp8Pair(try expertModule(&b, i, e, proj), ni, d) else try r.fp8Pair(try expertModule(&b, i, e, proj), d, ni); + try stage.put(r, (m * nb + j) * one, x.w.bytes); + @memcpy(std.mem.sliceAsBytes(inv[(if (m < 2) 2 * e + m else 2 * count + e) * inv_n ..][0..inv_n]), x.s.bytes); + }; + try stage.upload(r); + for (0..2) |m| try r.pack.fp8Experts(r.stream, stage.dev.ptr + m * nb * one, up + e0 * 2 * one, n, ni, d, 2, m); + try r.pack.fp8Experts(r.stream, stage.dev.ptr + 2 * nb * one, down + e0 * one, n, d, ni, 1, 0); + try stage.drained(r); } + const up_scale = try r.store.keep(std.mem.sliceAsBytes(inv[0 .. 2 * count * inv_n]), try label(&lb, i, "up_scale")); + const down_scale = try r.store.keep(std.mem.sliceAsBytes(inv[2 * count * inv_n ..]), try label(&lb, i, "down_scale")); return .{ .format = .fp8g, .up = up, .down = down, .up_scale = up_scale, .down_scale = down_scale, .width = ni, .dims = d, .experts = count, .limit = limit, .after = true }; } - /// experts.Layer's nvfp4 form: Nvfp4Experts.packGateUp and packOne per expert, weight_scale_2 as the globals. + /// experts.Layer's nvfp4 form: batches of codes and scales staged on the device, packed there (Packer.fp4Experts). fn fp4Experts(r: *Reader, i: u32, count: usize, ni: usize, d: usize, limit: f32) !cuda.experts.Layer { if (ni % X4.cols != 0 or d % 32 != 0) return error.UnexpectedTensor; - const up_n = X4.bytes(1, ni, d, 2) / 4; - const down_n = X4.bytes(1, d, ni, 1) / 4; + const up_n = X4.bytes(1, ni, d, 2); + const down_n = X4.bytes(1, d, ni, 1); + const cb = ni * d / 2; // code bytes of one matrix + const sb = ni * d / 16; // its e4m3 scale bytes var lb: [96]u8 = undefined; - const up = try r.store.alloc(count * up_n * 4, try label(&lb, i, "up")); - const down = try r.store.alloc(count * down_n * 4, try label(&lb, i, "down")); - const hu = try r.gpa.alloc(u32, up_n); - defer r.gpa.free(hu); - const hd = try r.gpa.alloc(u32, down_n); - defer r.gpa.free(hd); + const up = try r.store.alloc(count * up_n, try label(&lb, i, "up")); + const down = try r.store.alloc(count * down_n, try label(&lb, i, "down")); const gu = try r.gpa.alloc(f32, 2 * count); defer r.gpa.free(gu); const gd = try r.gpa.alloc(f32, count); defer r.gpa.free(gd); + const nb = r.batch(count, 3 * (cb + sb)); + var stage = try Stage.init(r, 3 * nb * (cb + sb)); // codes [3][nb], then scales [3][nb] + defer stage.deinit(r); + const codes = stage.dev.ptr; + const scales = stage.dev.ptr + 3 * nb * cb; var b: [192]u8 = undefined; - for (0..count) |e| { - const g = try r.fp4Read(try expertModule(&b, i, e, "gate_proj"), ni, d); - const u = try r.fp4Read(try expertModule(&b, i, e, "up_proj"), ni, d); - const dn = try r.fp4Read(try expertModule(&b, i, e, "down_proj"), d, ni); - X4.packGateUp(hu, .{ g.codes, g.scales }, .{ u.codes, u.scales }, ni, d); - X4.packOne(hd, dn.codes, dn.scales, d, ni); - try r.view(up, count * up_n * 4).upload(e * up_n * 4, std.mem.sliceAsBytes(hu)); - try r.view(down, count * down_n * 4).upload(e * down_n * 4, std.mem.sliceAsBytes(hd)); - gu[2 * e] = g.global; - gu[2 * e + 1] = u.global; - gd[e] = dn.global; + var e0: usize = 0; + while (e0 < count) : (e0 += nb) { + const n = @min(nb, count - e0); + for (0..n) |j| inline for (.{ "gate_proj", "up_proj", "down_proj" }, 0..) |proj, m| { + const e = e0 + j; + const x = if (m < 2) try r.fp4Read(try expertModule(&b, i, e, proj), ni, d) else try r.fp4Read(try expertModule(&b, i, e, proj), d, ni); + try stage.put(r, (m * nb + j) * cb, x.codes); + try stage.put(r, 3 * nb * cb + (m * nb + j) * sb, x.scales); + if (m < 2) gu[2 * e + m] = x.global else gd[e] = x.global; + }; + try stage.upload(r); + for (0..2) |m| try r.pack.fp4Experts(r.stream, codes + m * nb * cb, scales + m * nb * sb, up + e0 * up_n, n, ni, d, 0, d, 2, m); + try r.pack.fp4Experts(r.stream, codes + 2 * nb * cb, scales + 2 * nb * sb, down + e0 * down_n, n, d, ni, 0, ni, 1, 0); + try stage.drained(r); } const up_scale = try r.store.keep(std.mem.sliceAsBytes(gu), try label(&lb, i, "up_scale")); const down_scale = try r.store.keep(std.mem.sliceAsBytes(gd), try label(&lb, i, "down_scale")); diff --git a/zig/tests/cuda/ling_pack.zig b/zig/tests/cuda/ling_pack.zig index 12d75ee844..eb4747c2ef 100644 --- a/zig/tests/cuda/ling_pack.zig +++ b/zig/tests/cuda/ling_pack.zig @@ -1,4 +1,4 @@ -//! `ling-pack DIR LAYER [EXPERTS]`: a MoE layer's experts loaded and timed, bytes vs the host packers the loader used. +//! `ling-pack DIR LAYER [EXPERTS [STAGE]]`: MoE experts loaded and timed, bytes vs the host packers the loader used. const std = @import("std"); const cuda = @import("cuda"); const core = @import("core"); @@ -128,7 +128,7 @@ pub fn run(gpu: check.Gpu, rest: []const [:0]const u8) !void { defer o.ck.close(); if (!o.c.isMoe(i)) return error.NotMoeLayer; const count = if (rest.len > 2) try std.fmt.parseInt(usize, rest[2], 10) else o.c.experts; - const opts: W.Options = .{ .globals = false, .experts = count }; + const opts: W.Options = .{ .globals = false, .experts = count, .stage = if (rest.len > 3) try std.fmt.parseInt(usize, rest[3], 10) else 512 << 20 }; // the read alone from a cold cache, then a load on hot pages (pack and upload), then a cold load (both) o.ck.release(); var t0 = check.now(io); diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 2233fdc502..1247d3c0e8 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -119,7 +119,7 @@ const usage = \\ ling-lanes [ref] [--layers a,b] [--mtp] Ling's lane backend: chunked, drafted, resumed, concurrent == plain \\ ling-mtp-real the real MTP layer vs HF's (oracle/ling_mtp.py); cut, window == decode rows \\ ling-weights-real real Ling-3.1 layers 0, 4, 7 (4 experts) vs float64 and torch (oracle/ling_weights.py) - \\ ling-pack [experts] one MoE layer's experts: read, pack and cold load timed; bytes vs the host packers + \\ ling-pack [experts [stage]] one MoE layer's experts: read, pack and cold load timed; bytes vs the host packers \\ ; From 412f66fc69c192f2d04e7355ada3c8dd72d23def Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 15 Oct 2026 15:45:53 +0200 Subject: [PATCH 14/34] tests: ling-layer-speed times one real layer whole and op by op per phase; ling-layer-bytes writes every path's bytes for cmp --- zig/tests/cuda/ling_layer_bytes.zig | 119 ++++++++++++ zig/tests/cuda/ling_layer_speed.zig | 275 ++++++++++++++++++++++++++++ zig/tests/cuda/main.zig | 4 + 3 files changed, 398 insertions(+) create mode 100644 zig/tests/cuda/ling_layer_bytes.zig create mode 100644 zig/tests/cuda/ling_layer_speed.zig diff --git a/zig/tests/cuda/ling_layer_bytes.zig b/zig/tests/cuda/ling_layer_bytes.zig new file mode 100644 index 0000000000..7414298946 --- /dev/null +++ b/zig/tests/cuda/ling_layer_bytes.zig @@ -0,0 +1,119 @@ +//! `ling-layer-bytes DIR FIXTURE OUT`: one real Ling-3.1 layer's outputs and states on every path, written for cmp. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const ling = @import("ling_cuda"); +const rig = @import("ling_forward_rig.zig"); + +const F = rig.F; +const Rig = rig.Rig; +const Opened = @typeInfo(@typeInfo(@TypeOf(rig.W.open)).@"fn".return_type.?).error_union.payload; + +/// The loaded layer, its graph, a prompt and a decode Buffers, the fixture's rows tiled to `rows`. +pub const Kit = struct { + o: Opened, + m: F.Modules, + w: rig.W.Weights, + g: F.Graph, + r: Rig, + pb: F.Buffers, + db: F.Buffers, + x: []u16, + D: usize, + L: u32, + cap: usize, + + /// Loads FIXTURE's layer from DIR; prompt buffers for `rows` rows, states of `cap` positions. + pub fn init(k: *Kit, gpu: check.Gpu, a: std.mem.Allocator, s: cuda.Stream, dir: []const u8, path: []const u8, rows: usize, cap: usize) !void { + const bytes = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, a, .limited(1 << 30)); + if (bytes.len < 24 or !std.mem.eql(u8, bytes[0..8], "TFLINGL1")) return error.BadFixture; + const dims = std.mem.bytesAsSlice(u32, bytes[8..24]); + k.L = dims[0]; + const T: usize = dims[1]; + k.D = dims[2]; + k.x = try a.alloc(u16, rows * k.D); + for (0..rows) |t| @memcpy(std.mem.sliceAsBytes(k.x[t * k.D ..][0..k.D]), bytes[24 + (t % T) * k.D * 2 ..][0 .. k.D * 2]); + var why: ling.config.Why = .{}; + k.o = rig.W.open(gpu.gpa, gpu.io, dir, &why) catch |e| { + std.debug.print("{s}: {s}\n", .{ dir, why.text() }); + return e; + }; + k.m = try F.Modules.load(gpu.d, gpu.ctx); + var pack = try cuda.modelopt.Packer.load(gpu.d); + defer pack.unload(); + k.w = try rig.W.load(gpu.gpa, gpu.d, s, &pack, &k.o.ck, k.o.c, &.{k.L}, .{ .globals = false }); + k.g = try F.Graph.init(gpu.gpa, gpu.d, &k.m, s, &k.w); + k.r = .{ .gpu = gpu, .a = a, .s = s, .m = &k.m, .g = &k.g }; + k.cap = cap; + k.pb = try F.Buffers.init(gpu.d, &k.m, &k.w, rows, true, cap); + k.db = try F.Buffers.init(gpu.d, &k.m, &k.w, F.window_rows, false, cap); + } + + pub fn deinit(k: *Kit) void { + k.pb.deinit(); + k.db.deinit(); + k.r.deinit(); + k.g.deinit(); + k.w.deinit(); + k.m.deinit(); + k.o.ck.close(); + } + + /// Rows lo .. hi through the layer on `st`, committing `keep` of them (0: all); the output rows' bytes. + pub fn pass(k: *Kit, b: *const F.Buffers, st: *F.State, lo: usize, hi: usize, keep: usize) ![]u8 { + const d = k.r.gpu.d; + const n = hi - lo; + try d.check(d.api.cuMemcpyHtoDAsync_v2(b.x, k.x[lo * k.D ..].ptr, n * k.D * 2, k.r.s.handle), "cuMemcpyHtoDAsync"); + const f = k.r.fwd(b, null); + try f.stack(&.{.{ .st = st, .row0 = 0, .rows = n }}); + const out = try k.r.back(b.x, n * k.D * 2); + try f.commit(st, if (keep == 0) n else keep); + return out; + } +}; + +fn put(list: *std.ArrayList(u8), a: std.mem.Allocator, bytes: []const u8) !void { + try list.appendSlice(a, bytes); +} + +/// Every path's bytes: prompts of 512 and 2048 rows, 16 decode rows, windows of 2..16 with partial commits, states. +pub fn run(gpu: check.Gpu, dir: []const u8, path: []const u8, out: []const u8) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var k: Kit = undefined; + try k.init(gpu, a, s, dir, path, 2048, 2600); + defer k.deinit(); + var all: std.ArrayList(u8) = .empty; + const P = 400; + for ([_]usize{ 512, 2048 }) |T| { + const st = try k.r.state(k.cap); + defer st.deinit(); + try put(&all, a, try k.pass(&k.pb, st, 0, T, 0)); + try put(&all, a, try k.r.back(st.arena.buf.ptr, st.arena.buf.len)); + } + for ([_]bool{ true, false }) |wide| { + k.r.kda_wide = wide; + const st = try k.r.state(k.cap); + defer st.deinit(); + _ = try k.pass(&k.pb, st, 0, P, 0); + for (0..F.window_rows) |i| try put(&all, a, try k.pass(&k.db, st, P + i, P + i + 1, 0)); + try put(&all, a, try k.r.back(st.arena.buf.ptr, st.arena.buf.len)); + for ([_]usize{ 2, 3, 4, 8, 16 }) |w| { + const sw = try k.r.state(k.cap); + defer sw.deinit(); + _ = try k.pass(&k.pb, sw, 0, P, 0); + const keep = (w + 1) / 2; + try put(&all, a, try k.pass(&k.db, sw, P, P + w, keep)); + try put(&all, a, try k.pass(&k.db, sw, P + keep, P + keep + w, 0)); + try put(&all, a, try k.r.back(sw.arena.buf.ptr, sw.arena.buf.len)); + } + } + k.r.kda_wide = true; + var h = std.hash.Wyhash.init(0); + h.update(all.items); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = out, .data = all.items }); + check.pass("ling-layer-bytes layer {d}: {d} bytes into {s}, wyhash {x:0>16}", .{ k.L, all.items.len, out, h.final() }); +} diff --git a/zig/tests/cuda/ling_layer_speed.zig b/zig/tests/cuda/ling_layer_speed.zig new file mode 100644 index 0000000000..4ec40a2d8a --- /dev/null +++ b/zig/tests/cuda/ling_layer_speed.zig @@ -0,0 +1,275 @@ +//! `ling-layer-speed DIR FIXTURE [REPS]`: one real Ling-3.1 layer timed by CUDA events, whole and op by op, per phase. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const ling = @import("ling_cuda"); +const rig = @import("ling_forward_rig.zig"); +const bytes = @import("ling_layer_bytes.zig"); + +const F = rig.F; +const G = ling.graph; +const W = rig.W; +const kda = G.kda; +const latent = G.latent; +const glue = G.glue; + +/// Events between op groups; each group's milliseconds by name, medians over reps. +const Timer = struct { + ev: [48]cuda.Event, + names: [48][]const u8 = undefined, + n: usize = 0, + s: cuda.Stream, + on: bool, + + fn mark(t: *Timer, name: []const u8) !void { + if (!t.on) return; + try t.ev[t.n].record(t.s); + t.names[t.n] = name; + t.n += 1; + } +}; + +const Seg = F.Seg; + +fn kdaBlock(o: G.Ops, b: *const F.Buffers, a: W.Kda, c: *const ling.config.Config, li: usize, sg: Seg, rows: usize, t: *Timer) !void { + const D: usize = c.hidden; + const H: usize = c.kda_heads; + const n = H * kda.dk; + const st = sg.st; + const ps = st.stride; + const sv = st.saves(li); + const p = if (b.prompt) b.kp else sv.p; + for ([_]W.Proj{ a.q, a.k, a.v, a.b }, [_]usize{ 0, n, 2 * n, 3 * n }, [_][]const u8{ "proj q", "proj k", "proj v", "proj b (bf16)" }) |w, col0, name| { + try o.mm(w, b.normed, D, rows, b.tmp); + try t.mark(name); + const w_n = w.outputs(); + try o.copy(b.tmp, w_n * 2, p + col0 * 2, ps * 2, rows, w_n * 2); + try t.mark("copies into p"); + } + try o.mm(a.f, b.normed, D, rows, b.ka); + try t.mark("proj f"); + try o.mm(a.g, b.normed, D, rows, b.kg); + try t.mark("proj g"); + const cur = st.cur[li]; + const ch: kda.Chain = .{ .p = p, .p_stride = ps, .b_off = 3 * n, .a = b.ka, .a_stride = n, .g = b.kg, .g_stride = n, .conv_state = st.convAt(li), .conv_w = a.conv, .state_in = st.recAt(cur, li), .a_log = a.a_log, .dt_bias = a.dt_bias, .norm_w = a.o_norm, .eps = c.eps, .lower = c.lower_bound, .heads = H }; + const k = if (b.prompt) b.saves.k else sv.k; + const v = if (b.prompt) b.saves.v else sv.v; + const g = if (b.prompt) b.saves.g else sv.g; + const bb = if (b.prompt) b.saves.b else sv.b; + const s: kda.Saves = .{ .out = b.saves.out, .state_out = st.recAt(1 - cur, li), .k = k, .v = v, .g = g, .b = bb }; + const wide: ?kda.Wide = if (if (b.prompt) rows >= kda.wide_rows else o.kda_wide) .{ .q = b.wide.q, .y = b.wide.y } else null; + try kda.chain(&o.m.kda, o.s, ch, rows, s, wide); + try t.mark("KDA chain"); + if (b.prompt) { + try glue.convShift(o.set(), o.s, st.convAt(li), p, rows, rows * ps, ps, 1, 3 * n); + try t.mark("conv shift"); + } + try o.mm(a.o, b.saves.out, n, rows, b.y); + try t.mark("proj o"); +} + +fn mlaBlock(o: G.Ops, b: *const F.Buffers, a: W.Mla, ab: G.Absorb, g: *const F.Graph, sg: Seg, rows: usize, t: *Timer) !void { + const c = &g.w.config; + const D: usize = c.hidden; + const H: usize = c.mla_heads; + const qk: usize = c.qkDim(); + const nope: usize = c.nope; + const kw = F.cache_width; + const lw = 512; + const ql: usize = c.q_lora; + try o.mm(a.q_a.?, b.normed, D, rows, b.qa_in); + try t.mark("proj q_a"); + try glue.rmsnorm(o.set(), o.s, b.qa_in, ql, a.q_a_norm, b.qr, ql, null, c.eps, rows, ql); + try t.mark("norms q_a, kv"); + try o.mm(a.q_b.?, b.qr, ql, rows, b.q); + try t.mark("proj q_b"); + try o.mm(a.kv_a, b.normed, D, rows, b.kv); + try t.mark("proj kv_a"); + try glue.rmsnorm(o.set(), o.s, b.kv, kw, a.kv_a_norm, b.lat, kw, null, c.eps, rows, lw); + try t.mark("norms q_a, kv"); + try o.mm(a.gate.?, b.normed, D, rows, b.g); + try t.mark("proj gate (bf16)"); + try latent.absorbQStrided(o.set(), o.s, b.q, qk, ab.wk, b.qa, kw, rows, H, nope, lw); + try t.mark("absorb q"); + const st = sg.st; + try latent.rope(o.set(), o.s, .{ .x = b.q + nope * 2, .x_row = H * qk, .x_head = qk, .out = b.qa + lw * 2, .o_row = H * kw, .o_head = kw }, rows, H, st.pos_dev, g.freq); + try latent.rope(o.set(), o.s, .{ .x = b.kv + lw * 2, .x_row = kw, .x_head = 0, .out = b.lat + lw * 2, .o_row = kw, .o_head = 0 }, rows, 1, st.pos_dev, g.freq); + try latent.latentWrite(o.set(), o.s, b.lat, kw, st.cacheAt(0), st.pos_dev, rows, kw); + try t.mark("rope, cache write"); + var r0: usize = 0; + while (r0 < rows) : (r0 += ling.state.att_rows) { + const m = @min(ling.state.att_rows, rows - r0); + var at = st.pos_dev; + if (r0 > 0) { + try o.addI32(st.pos_dev, r0, b.at); + at = b.at; + } + const nch = latent.chunksFor(st.pos + r0 + m); + try latent.attentionRoped(o.set(), o.s, b.qa + r0 * H * kw * 2, st.cacheAt(0), at, b.partials, b.ol + r0 * H * lw * 2, b.dummy, m, H, lw, nch, c.smScale()); + } + try t.mark("latent attention"); + try latent.expandV(o.set(), o.s, b.ol, ab.wv, b.vn, rows, H, c.v_dim, lw); + try t.mark("expand v"); + try o.gate(b.vn, b.g, H, rows, H, c.v_dim); + try t.mark("head gate"); + try o.mm(a.o, b.vn, H * c.v_dim, rows, b.y); + try t.mark("proj o"); +} + +fn mlp(o: G.Ops, b: *const F.Buffers, m: W.Mlp, rows: usize, limit: f32, out: u64, t: *Timer, comptime tag: []const u8) !void { + const D = m.gate.inputs(); + const w = m.gate.outputs(); + for ([_]W.Proj{ m.gate, m.up }, [_]usize{ 0, w }) |p, col0| { + try o.mm(p, b.normed, D, rows, b.tmp); + try t.mark(tag ++ " gate, up"); + try o.copy(b.tmp, w * 2, b.gu + col0 * 2, 4 * w, rows, w * 2); + try t.mark(tag ++ " copies"); + } + try glue.swigluWith(o.m.swiglu, o.s, b.gu, b.act, b.xs, limit, rows, w); + try t.mark(tag ++ " swiglu"); + try o.mm(m.down, b.act, w, rows, out); + try t.mark(tag ++ " down"); +} + +fn moe(o: G.Ops, b: *const F.Buffers, e: W.Moe, g: *const F.Graph, rows: usize, t: *Timer) !void { + const c = &g.w.config; + const D: usize = c.hidden; + const E: usize = c.experts; + const top: usize = c.topk; + const slots = top + 1; + try glue.router(o.set(), o.s, b.normed, D, e.router, b.mlog, b.rpart, rows, D, E); + try t.mark("router"); + try glue.select(o.set(), o.s, b.mlog, if (e.bias != 0) e.bias else g.zero_bias, b.pick, b.wts, rows, top, E, c.routed_scale, c.norm_topk, ling.groups(c)); + try t.mark("select"); + const l = e.routed.grouped; + const tile: usize = if (o.prompt) cuda.experts.promptTile(l.format) else 16; + try o.m.router.route(o.s, b.pick, rows * slots, E + 1, tile, b.plan); + try t.mark("expert plan"); + const items = cuda.grouped.maxItems(rows * slots, E + 1, tile); + const gr = o.m.grouped(); + const up: cuda.experts.Rows = .{ .x = b.normed, .stride = l.dims, .slots = slots }; + const down: cuda.experts.Rows = .{ .x = b.eact, .stride = l.width }; + const skip: i32 = @intCast(E); + if (o.prompt) try gr.prompt(o.s, .up, up, l, b.plan, b.eact, items, skip) else try gr.decode(o.s, .up, up, l, b.plan, b.eact, items, skip); + try t.mark("experts gate-up"); + if (o.prompt) try gr.prompt(o.s, .down_bf16, down, l, b.plan, b.ey, items, skip) else try gr.decode(o.s, .down_bf16, down, l, b.plan, b.ey, items, skip); + try t.mark("experts down"); + if (o.prompt) { // the decode kernels on 16-pair items over the same rows, for comparison + try o.m.router.route(o.s, b.pick, rows * slots, E + 1, 16, b.plan); + try t.mark("alt plan 16"); + const it16 = cuda.grouped.maxItems(rows * slots, E + 1, 16); + try gr.decode(o.s, .up, up, l, b.plan, b.eact, it16, skip); + try t.mark("alt gate-up decode kernel"); + try gr.decode(o.s, .down_bf16, down, l, b.plan, b.ey, it16, skip); + try t.mark("alt down decode kernel"); + } + try mlp(o, b, e.shared.?, rows, e.shared_limit, b.sh, t, "shared"); + var a: cuda.Args = .{}; + inline for (.{ b.ey, b.pick, b.wts, b.sh, b.y, @as(i64, @intCast(rows)), @as(i64, @intCast(D)), @as(c_int, @intCast(slots)), @as(c_int, @intCast(top)), @as(c_int, 0) }) |v| a.add(v); + try cuda.launch.launch(o.m.combine, .{ .grid = .{ .x = @intCast(rows), .y = @intCast((D + 255) / 256) }, .block = .{ .x = 256 } }, o.s, &a); + try t.mark("combine"); +} + +/// The layer as forward.layer runs it, op groups marked. +fn layerOps(f: *const F.Forward, sg: Seg, rows: usize, t: *Timer) !void { + const o = f.ops(); + const b = f.b; + const g = f.g; + const c = &g.w.config; + const D: usize = c.hidden; + const l = g.w.layers[0]; + try t.mark("start"); + try glue.rmsnorm(o.set(), o.s, b.x, D, l.attn_norm, b.normed, D, null, c.eps, rows, D); + try t.mark("norms + residual adds"); + switch (l.mixer) { + .kda => |a| try kdaBlock(o, b, a, c, 0, sg, rows, t), + .mla => |a| try mlaBlock(o, b, a, g.absorb[0], g, sg, rows, t), + } + try o.add(b.x, b.y, b.x, rows * D); + try glue.rmsnorm(o.set(), o.s, b.x, D, l.mlp_norm, b.normed, D, null, c.eps, rows, D); + try t.mark("norms + residual adds"); + switch (l.ffn) { + .dense => |m| try mlp(o, b, m, rows, 0, b.y, t, "dense"), + .moe => |e| try moe(o, b, e, g, rows, t), + } + try o.add(b.x, b.y, b.x, rows * D); + try t.mark("norms + residual adds"); +} + +const Row = struct { name: []const u8, ms: [32]f64 = undefined, n: usize = 0 }; + +fn median(xs: []f64) f64 { + std.mem.sort(f64, xs, {}, std.sort.asc(f64)); + return xs[xs.len / 2]; +} + +/// `ling-layer-speed DIR FIXTURE [REPS]`: decode at position 2048, windows 2..16, prompts of 512 and 2048 rows. +pub fn run(gpu: check.Gpu, dir: []const u8, path: []const u8, reps_arg: ?[]const u8) !void { + const reps: usize = if (reps_arg) |x| try std.fmt.parseInt(usize, x, 10) else 15; + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var k: bytes.Kit = undefined; + try k.init(gpu, a, s, dir, path, 2048, 4200); + defer k.deinit(); + var t: Timer = .{ .ev = undefined, .s = s, .on = true }; + for (&t.ev) |*e| e.* = try cuda.Event.init(gpu.d, true); + defer for (&t.ev) |*e| e.deinit(); + const st = try k.r.state(k.cap); + defer st.deinit(); + const c = &k.w.config; + std.debug.print("RESULT ling-layer-speed layer {d} ({t}, {s}, {t}), {d} reps, median and least us (other GPU contexts' time slices land in the medians)\n", .{ k.L, c.kind(k.L), if (c.isMoe(k.L)) "MoE" else "dense", c.quant, reps }); + for ([_]usize{ 1, 2, 4, 8, 16, 512, 2048 }) |rows| { + const b = if (rows > F.window_rows) &k.pb else &k.db; + const pos: usize = if (rows > F.window_rows) 0 else 2048; + try st.setPos(s, pos); + const d = gpu.d; + try d.check(d.api.cuMemcpyHtoDAsync_v2(b.x, k.x.ptr, rows * k.D * 2, s.handle), "cuMemcpyHtoDAsync"); + const f = k.r.fwd(b, null); + const sg: Seg = .{ .st = st, .row0 = 0, .rows = rows }; + var whole: [64]f64 = undefined; + var table: [32]Row = undefined; + var nt: usize = 0; + for (0..reps + 2) |i| { + // whole layer through forward.layer, then the same launches op by op + try t.ev[46].record(s); + try f.layer(&.{sg}, 0, rows, null); + try t.ev[47].record(s); + try s.synchronize(); + if (i >= 2) whole[i - 2] = try t.ev[46].elapsedMs(t.ev[47]); + t.n = 0; + try layerOps(&f, sg, rows, &t); + try s.synchronize(); + if (i < 2) continue; + var acc: [32]f64 = @splat(0); + for (1..t.n) |j| { + const name = t.names[j]; + var at: usize = 0; + while (at < nt and !std.mem.eql(u8, table[at].name, name)) at += 1; + if (at == nt) { + table[nt] = .{ .name = name }; + nt += 1; + } + acc[at] += try t.ev[j - 1].elapsedMs(t.ev[j]); + } + for (table[0..nt], acc[0..nt]) |*r, v| { + r.ms[r.n] = v; + r.n += 1; + } + } + var sum: f64 = 0; + var low: f64 = 0; + std.debug.print("RESULT rows {d:>4}: layer {d:.1} {d:.1}", .{ rows, median(whole[0..reps]) * 1000, std.mem.min(f64, whole[0..reps]) * 1000 }); + for (table[0..nt]) |*r| { + const lo = std.mem.min(f64, r.ms[0..r.n]) * 1000; + const m = median(r.ms[0..r.n]) * 1000; + sum += m; + low += lo; + std.debug.print(" | {s} {d:.1} {d:.1}", .{ r.name, m, lo }); + } + std.debug.print(" | sum of ops {d:.1} {d:.1}\n", .{ sum, low }); + } + check.pass("ling-layer-speed layer {d}: timed", .{k.L}); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 1247d3c0e8..2e4616fcb2 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -114,6 +114,8 @@ const usage = \\ ling-weights every layer of a Ling checkpoint loaded: projections vs float64, small tensors byte for byte \\ ling-forward-synth [picks] a ling-synth checkpoint's forward vs oracle/ling_forward.py; cuts, windows, batches \\ ling-layer-real [dump] one real Ling-3.1 layer vs HF's captured output (oracle/ling_layer.py); exactness + \\ ling-layer-speed [reps] one real Ling-3.1 layer timed whole and op by op: decode, windows, prompts + \\ ling-layer-bytes one real Ling-3.1 layer's outputs and states on every path into a file (cmp) \\ ling-clamp Ling's SwiGLU order min(silu(g), L): glue and the dir's grouped experts vs float64 \\ ling-clamp-old the grouped expert instances vs an older build's fatbins, byte for byte \\ ling-lanes [ref] [--layers a,b] [--mtp] Ling's lane backend: chunked, drafted, resumed, concurrent == plain @@ -229,6 +231,8 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "ling-weights")) return ling_weights.run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "ling-forward-synth")) return @import("ling_forward_synth.zig").run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) rest[2] else null); if (std.mem.eql(u8, cmd, "ling-layer-real")) return @import("ling_layer_real.zig").run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) rest[2] else null); + if (std.mem.eql(u8, cmd, "ling-layer-speed")) return @import("ling_layer_speed.zig").run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) rest[2] else null); + if (std.mem.eql(u8, cmd, "ling-layer-bytes")) return @import("ling_layer_bytes.zig").run(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2)); if (std.mem.eql(u8, cmd, "ling-clamp")) return @import("ling_clamp.zig").run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "ling-clamp-old")) return @import("ling_clamp.zig").old(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "ling-mtp-real")) return @import("ling_mtp_real.zig").run(gpu, try arg(rest, 0), try arg(rest, 1)); From 99652babf2cb72485361a39d2a98c96322068564 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 15 Oct 2026 15:45:53 +0200 Subject: [PATCH 15/34] ling: row copies 16 bytes a thread on aligned rows (prompt KDA and MLP column copies -50 %), the same bytes --- zig/kernels/cuda/ling_glue.cu | 18 ++++++++++++++++++ zig/src/families/ling/graph.zig | 4 ++-- 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/zig/kernels/cuda/ling_glue.cu b/zig/kernels/cuda/ling_glue.cu index 34a6917017..83373f50f0 100644 --- a/zig/kernels/cuda/ling_glue.cu +++ b/zig/kernels/cuda/ling_glue.cu @@ -37,3 +37,21 @@ extern "C" __global__ void tf_ling_combine(const __nv_bfloat16* y, const int* pi if (shared != nullptr) v = __fadd_rn(v, ling_glue::bf(shared[r * d + c])); out[r * d + c] = __float2bfloat16_rn(v); } + +// Bytes [rows][row_bytes] between pitched rows, as they are: 16 bytes a thread where every address is 16-aligned. +extern "C" __global__ void tf_ling_copy_rows(const unsigned char* __restrict__ src, long long src_pitch, + unsigned char* __restrict__ dst, long long dst_pitch, int rows, + int row_bytes) { + const unsigned long long a = reinterpret_cast(src) | reinterpret_cast(dst); + const bool wide = ((a | static_cast(src_pitch | dst_pitch | row_bytes)) & 15) == 0; + const int step = gridDim.x * blockDim.x, at = blockIdx.x * blockDim.x + threadIdx.x; + for (int r = blockIdx.y; r < rows; r += gridDim.y) { + const unsigned char* s = src + r * src_pitch; + unsigned char* d = dst + r * dst_pitch; + if (wide) { + for (int c = at; c < row_bytes / 16; c += step) reinterpret_cast(d)[c] = reinterpret_cast(s)[c]; + } else { + for (int c = at; c < row_bytes; c += step) d[c] = s[c]; + } + } +} diff --git a/zig/src/families/ling/graph.zig b/zig/src/families/ling/graph.zig index 9a2a645265..c202b8b966 100644 --- a/zig/src/families/ling/graph.zig +++ b/zig/src/families/ling/graph.zig @@ -39,7 +39,7 @@ pub const Modules = struct { const sms: usize = @intCast(try ctx.attribute(.multiprocessor_count)); m.glue = try glue.Kernels.load(m.mods[0]); m.kda = try kda.Kernels.load(m.mods[1]); - m.copy_rows = try m.mods[2].function("tf_glm_copy_rows"); + m.copy_rows = try m.mods[4].function("tf_ling_copy_rows"); m.add_i32 = try m.mods[2].function("tf_glm_add_i32"); m.router = try cuda.grouped.Router.resolve(m.mods[3]); m.add = try m.mods[4].function("tf_ling_add"); @@ -106,7 +106,7 @@ pub const Ops = struct { pub fn copy(o: Ops, src: u64, src_pitch: usize, dst: u64, dst_pitch: usize, rows: usize, row_bytes: usize) !void { var a: cuda.Args = .{}; inline for (.{ src, @as(c_longlong, @intCast(src_pitch)), dst, @as(c_longlong, @intCast(dst_pitch)), @as(c_int, @intCast(rows)), @as(c_int, @intCast(row_bytes)) }) |v| a.add(v); - try cuda.launch.launch(o.m.copy_rows, .{ .grid = .{ .x = cdiv(rows * row_bytes, 256) }, .block = .{ .x = 256 } }, o.s, &a); + try cuda.launch.launch(o.m.copy_rows, .{ .grid = .{ .x = cdiv(row_bytes, 256 * 16), .y = @intCast(@max(1, @min(rows, 65535))) }, .block = .{ .x = 256 } }, o.s, &a); } /// A projection into columns col0 .. col0 + n of rows `stride` wide (bf16), through `tmp`. From b3faea69d9b63e8713d58e468b5683e6bc061684 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 15 Oct 2026 16:59:37 +0200 Subject: [PATCH 16/34] ling: a single decode row's projections write their columns in place, no copies (4 KDA, 2 per MLP), the same bytes --- zig/src/families/ling/block_ffn.zig | 9 +++++++-- zig/src/families/ling/block_kda.zig | 4 ++++ 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/zig/src/families/ling/block_ffn.zig b/zig/src/families/ling/block_ffn.zig index e6f0c4d0fa..5b95112700 100644 --- a/zig/src/families/ling/block_ffn.zig +++ b/zig/src/families/ling/block_ffn.zig @@ -13,8 +13,13 @@ const Config = @import("config.zig").Config; /// Gate and up into [gate | up] rows, SwiGLU with min(silu(g), limit) (0: none), down into `out` [rows][hidden]. pub fn mlp(o: Ops, b: *const state.Buffers, m: weights.Mlp, x: u64, ldx: usize, rows: usize, limit: f32, out: u64) !void { const w = m.gate.outputs(); - try o.mmCols(m.gate, x, ldx, rows, b.tmp, b.gu, 2 * w, 0); - try o.mmCols(m.up, x, ldx, rows, b.tmp, b.gu, 2 * w, w); + if (rows == 1) { // one row is its own columns: no copies + try o.mm(m.gate, x, ldx, 1, b.gu); + try o.mm(m.up, x, ldx, 1, b.gu + w * 2); + } else { + try o.mmCols(m.gate, x, ldx, rows, b.tmp, b.gu, 2 * w, 0); + try o.mmCols(m.up, x, ldx, rows, b.tmp, b.gu, 2 * w, w); + } try glue.swigluWith(o.m.swiglu, o.s, b.gu, b.act, b.xs, limit, rows, w); try o.mm(m.down, b.act, w, rows, out); } diff --git a/zig/src/families/ling/block_kda.zig b/zig/src/families/ling/block_kda.zig index 706de08fdf..3817a88911 100644 --- a/zig/src/families/ling/block_kda.zig +++ b/zig/src/families/ling/block_kda.zig @@ -30,6 +30,10 @@ pub fn block(o: Ops, b: *const state.Buffers, a: weights.Kda, c: *const @import( const n = H * kda.dk; const ps = segs[0].st.stride; for ([_]weights.Proj{ a.q, a.k, a.v, a.b }, [_]usize{ 0, n, 2 * n, 3 * n }) |w, col0| { + if (rows == 1) { // one row is its own columns: no copy + try o.mm(w, b.normed, D, 1, place(b, segs[0], li).p + col0 * 2); + continue; + } try o.mm(w, b.normed, D, rows, b.tmp); const w_n = w.outputs(); for (segs) |sg| try o.copy(b.tmp + sg.row0 * w_n * 2, w_n * 2, place(b, sg, li).p + col0 * 2, ps * 2, sg.rows, w_n * 2); From 5b006370b0fa8b490d034763a4a8e571842595df Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 15 Oct 2026 16:59:37 +0200 Subject: [PATCH 17/34] ling: FP8 prompt chunks under 16 routed pairs an expert run the decode expert kernels (same bits, -7 to -27 % expert time up to 768 rows) --- zig/src/families/ling/block_ffn.zig | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/zig/src/families/ling/block_ffn.zig b/zig/src/families/ling/block_ffn.zig index 5b95112700..d20010e7ef 100644 --- a/zig/src/families/ling/block_ffn.zig +++ b/zig/src/families/ling/block_ffn.zig @@ -46,17 +46,21 @@ pub fn moe(o: Ops, b: *const state.Buffers, e: weights.Moe, g: *const graph.Grap try cuda.launch.launch(o.m.combine, .{ .grid = .{ .x = @intCast(rows), .y = @intCast((D + 255) / 256) }, .block = .{ .x = 256 } }, o.s, &a); } +/// Prompt chunks under this many routed pairs an expert stream FP8 experts faster on the decode kernels (same bits). +const decode_pairs = 16; + /// Grouped experts on every row's top-k slots (the shared slot skipped): gate-up SwiGLU, then down to bf16 slot rows. fn grouped(o: Ops, b: *const state.Buffers, l: cuda.experts.Layer, rows: usize, slots: usize, count: usize) !void { const pairs = rows * slots; - const tile: usize = if (o.prompt) cuda.experts.promptTile(l.format) else 16; + const prompt = o.prompt and !(l.format == .fp8g and rows * (slots - 1) < decode_pairs * count); + const tile: usize = if (prompt) cuda.experts.promptTile(l.format) else 16; const skip: i32 = @intCast(count); try o.m.router.route(o.s, b.pick, pairs, count + 1, tile, b.plan); const items = cuda.grouped.maxItems(pairs, count + 1, tile); const g = o.m.grouped(); const up: cuda.experts.Rows = .{ .x = b.normed, .stride = l.dims, .slots = slots }; const down: cuda.experts.Rows = .{ .x = b.eact, .stride = l.width }; - if (o.prompt) { + if (prompt) { try g.prompt(o.s, .up, up, l, b.plan, b.eact, items, skip); try g.prompt(o.s, .down_bf16, down, l, b.plan, b.ey, items, skip); } else { From e78d23bb0dcfa53ecf1fc6307d34d3c888eb9163 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 15 Oct 2026 16:59:37 +0200 Subject: [PATCH 18/34] glm: group-limited select without divisions where groups are whole lane columns (Ling's 64), same picks and weights; ling-route-old --- zig/kernels/cuda/glm_glue_moe.cuh | 44 ++++++++++++++++++++++++++++++- zig/tests/cuda/glm_ref_route.zig | 37 ++++++++++++++++++++++++++ zig/tests/cuda/main.zig | 6 +++-- 3 files changed, 84 insertions(+), 3 deletions(-) diff --git a/zig/kernels/cuda/glm_glue_moe.cuh b/zig/kernels/cuda/glm_glue_moe.cuh index edb23b9ce5..d156092022 100644 --- a/zig/kernels/cuda/glm_glue_moe.cuh +++ b/zig/kernels/cuda/glm_glue_moe.cuh @@ -30,7 +30,35 @@ __device__ __forceinline__ float swiglu1(float g, float u, float limit) { template __device__ __forceinline__ void keep_groups(float* choice, int lane, int ne, int gs, int ng, int keep) { float gscore = -INFINITY; // lane g: group g's top-2 sum - for (int g = 0; g < ng; ++g) { + const bool whole = gs % 32 == 0 && ne % 32 == 0; // groups of whole i: one pass, no divisions (Ling's 64) + if (whole) { + const int per = gs / 32, ni = ne / 32; + float a = -INFINITY, b = -INFINITY; + int g = 0, left = per; +#pragma unroll + for (int i = 0; i < MAXE; ++i) { + if (i >= ni) break; + const float c = choice[i]; + if (c > a) { + b = a; + a = c; + } else if (c > b) { + b = c; + } + if (--left == 0) { + for (int o = 16; o > 0; o >>= 1) { + const float oa = __shfl_xor_sync(0xffffffffu, a, o), ob = __shfl_xor_sync(0xffffffffu, b, o); + b = fmaxf(fminf(a, oa), fmaxf(b, ob)); + a = fmaxf(a, oa); + } + if (lane == g) gscore = a + b; + a = b = -INFINITY; + ++g; + left = per; + } + } + } + for (int g = 0; g < (whole ? 0 : ng); ++g) { float a = -INFINITY, b = -INFINITY; // this lane's best two in group g, a >= b for (int i = 0; i < MAXE; ++i) { const int e = i * 32 + lane; @@ -62,6 +90,20 @@ __device__ __forceinline__ void keep_groups(float* choice, int lane, int ne, int kept |= 1u << id; if (lane == id) live = false; } + if (whole) { + const int per = gs / 32, ni = ne / 32; + int g = 0, left = per; +#pragma unroll + for (int i = 0; i < MAXE; ++i) { + if (i >= ni) break; + if (!((kept >> g) & 1u)) choice[i] = -INFINITY; + if (--left == 0) { + ++g; + left = per; + } + } + return; + } for (int i = 0; i < MAXE; ++i) { const int e = i * 32 + lane; if (e < ne && !((kept >> (e / gs)) & 1u)) choice[i] = -INFINITY; diff --git a/zig/tests/cuda/glm_ref_route.zig b/zig/tests/cuda/glm_ref_route.zig index ede4237915..798c8068ac 100644 --- a/zig/tests/cuda/glm_ref_route.zig +++ b/zig/tests/cuda/glm_ref_route.zig @@ -354,3 +354,40 @@ pub fn ling(gpu: Gpu, path: []const u8) !void { try check.expect(hf_same == R and ex_same == R and hf_most <= std.math.pow(f64, 2, -16) and ex_most <= std.math.pow(f64, 2, -10), "Ling routing equals HF's and exllamav3's", .{}); check.pass("Ling layer 4 routing: picks equal HF's and exllamav3's on all {d} rows, weights within float64 tolerance", .{R}); } + +/// `ling-route-old `: group-limited select (Ling's shapes, odd groups, tie rows) vs an older glm_glue, byte for byte. +pub fn lingOld(gpu: Gpu, path: []const u8) !void { + var e: Env = undefined; + try env(gpu, &e); + defer e.deinit(); + const r = &e.r; + const bytes = try std.Io.Dir.cwd().readFileAllocOptions(gpu.io, path, r.a, .limited(1 << 26), .@"64", null); + var m = try cuda.Module.load(gpu.d, bytes); + defer m.unload(); + const prev = try m.function("tf_glm_topk"); + const shapes = [_]Route{ + .{ .rows = 300, .ne = 512, .top_k = 8, .ng = 8, .keep = 4, .scale = 2.5 }, + .{ .rows = 300, .ne = 256, .top_k = 8, .ng = 8, .keep = 4, .scale = 2.5 }, + .{ .rows = 100, .ne = 480, .top_k = 6, .ng = 12, .keep = 3, .scale = 1.0 }, + .{ .rows = 100, .ne = 512, .top_k = 8, .ng = 32, .keep = 1, .scale = 2.5 }, + .{ .rows = 100, .ne = 384, .top_k = 8, .ng = 4, .keep = 2, .scale = 2.5 }, + }; + var rows: usize = 0; + for (shapes) |q| for ([_]bool{ true, false }) |norm| for ([_]bool{ false, true }) |ties| { + const logits = try f32s(r, q.rows * q.ne, -6, 6); + const bias = try f32s(r, q.ne, -0.05, 0.05); + if (ties) tieRows(r, q, logits, bias); + const n = q.rows * (q.top_k + 1); + const dl = try r.dev(std.mem.sliceAsBytes(logits)); + const db = try r.dev(std.mem.sliceAsBytes(bias)); + const outs = [4]u64{ try r.zeros(n * 4), try r.zeros(n * 4), try r.zeros(n * 4), try r.zeros(n * 4) }; + var a: cuda.Args = .{}; + inline for (.{ dl, db, outs[0], outs[1], q.scale, @as(c_int, @intCast(q.ne)), @as(c_int, @intCast(q.top_k)), @as(c_int, @intFromBool(norm)), @as(c_int, @intCast(q.ng)), @as(c_int, @intCast(q.keep)) }) |v| a.add(v); + try cuda.launch.launch(prev, .{ .grid = .{ .x = @intCast(q.rows) }, .block = .{ .x = 32 } }, r.s, &a); + try glue.select(.{ .glue = &e.ks }, r.s, dl, db, outs[2], outs[3], q.rows, q.top_k, q.ne, q.scale, norm, .{ .n = q.ng, .keep = q.keep }); + try check.sameBytes("group-limited select picks vs the older build", try r.back(outs[2], n * 4), try r.back(outs[0], n * 4)); + try check.sameBytes("group-limited select weights vs the older build", try r.back(outs[3], n * 4), try r.back(outs[1], n * 4)); + rows += q.rows; + }; + check.pass("ling-route-old: {d} rows of group-limited select (5 shapes, normed and not, tie rows) equal the older kernel's bytes", .{rows}); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 2e4616fcb2..84584ef6f3 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -108,13 +108,14 @@ const usage = \\ glm-lanes-concurrent concurrent streams on one lanes.Engine, each against its solo run \\ glm-shared-check shared verify windows against lone twins, logits and states byte for byte \\ glm-route-old GLM's select and swiglu against an older glm_glue fatbin (6380c8da), byte for byte + \\ ling-route-old group-limited select (Ling's shapes, odd groups, ties) vs an older glm_glue fatbin, byte for byte \\ glm-route-ling Ling's group-limited routing on captured layer-4 logits vs float64, HF and exllamav3 \\ glm-synth a synthetic GLM-5.3-Flash checkpoint (2 layers, 16 experts) and its dc.json for the lane checks \\ ling-synth [--format bf16|fp8|nvfp4] [--layers N] a tiny Ling-3.1 checkpoint with the real names and config schema \\ ling-weights every layer of a Ling checkpoint loaded: projections vs float64, small tensors byte for byte \\ ling-forward-synth [picks] a ling-synth checkpoint's forward vs oracle/ling_forward.py; cuts, windows, batches \\ ling-layer-real [dump] one real Ling-3.1 layer vs HF's captured output (oracle/ling_layer.py); exactness - \\ ling-layer-speed [reps] one real Ling-3.1 layer timed whole and op by op: decode, windows, prompts + \\ ling-layer-speed [reps [rows,..]] one real Ling-3.1 layer timed whole and op by op: decode, windows, prompts \\ ling-layer-bytes one real Ling-3.1 layer's outputs and states on every path into a file (cmp) \\ ling-clamp Ling's SwiGLU order min(silu(g), L): glue and the dir's grouped experts vs float64 \\ ling-clamp-old the grouped expert instances vs an older build's fatbins, byte for byte @@ -225,13 +226,14 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "glm-lanes-cover")) return glm_lanes_tests.coverCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); if (std.mem.eql(u8, cmd, "glm-shared-check")) return glm_lanes_tests.sharedCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); if (std.mem.eql(u8, cmd, "glm-route-old")) return glm_ref_route.old(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "ling-route-old")) return glm_ref_route.lingOld(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "glm-route-ling")) return glm_ref_route.ling(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "glm-synth")) return glm_synth.run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "ling-synth")) return ling_synth.run(gpu, rest); if (std.mem.eql(u8, cmd, "ling-weights")) return ling_weights.run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "ling-forward-synth")) return @import("ling_forward_synth.zig").run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) rest[2] else null); if (std.mem.eql(u8, cmd, "ling-layer-real")) return @import("ling_layer_real.zig").run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) rest[2] else null); - if (std.mem.eql(u8, cmd, "ling-layer-speed")) return @import("ling_layer_speed.zig").run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) rest[2] else null); + if (std.mem.eql(u8, cmd, "ling-layer-speed")) return @import("ling_layer_speed.zig").run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) rest[2] else null, if (rest.len > 3) rest[3] else null); if (std.mem.eql(u8, cmd, "ling-layer-bytes")) return @import("ling_layer_bytes.zig").run(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2)); if (std.mem.eql(u8, cmd, "ling-clamp")) return @import("ling_clamp.zig").run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "ling-clamp-old")) return @import("ling_clamp.zig").old(gpu, try arg(rest, 0), try arg(rest, 1)); From 0ffce4f7f4a7ab5ec265ea3624b3ae434282418d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 15 Oct 2026 17:10:15 +0200 Subject: [PATCH 19/34] tests: ling-layer-speed takes a rows list, times prompt ops from an idle GPU, the expert kernels and plan tiles side by side --- zig/tests/cuda/glm_ref_route.zig | 2 +- zig/tests/cuda/ling_layer_speed.zig | 78 ++++++++++++++++++----------- 2 files changed, 51 insertions(+), 29 deletions(-) diff --git a/zig/tests/cuda/glm_ref_route.zig b/zig/tests/cuda/glm_ref_route.zig index 798c8068ac..cc8e07e3a0 100644 --- a/zig/tests/cuda/glm_ref_route.zig +++ b/zig/tests/cuda/glm_ref_route.zig @@ -355,7 +355,7 @@ pub fn ling(gpu: Gpu, path: []const u8) !void { check.pass("Ling layer 4 routing: picks equal HF's and exllamav3's on all {d} rows, weights within float64 tolerance", .{R}); } -/// `ling-route-old `: group-limited select (Ling's shapes, odd groups, tie rows) vs an older glm_glue, byte for byte. +/// `ling-route-old `: group-limited select (Ling's shapes, odd groups, ties) vs an older glm_glue, bytes. pub fn lingOld(gpu: Gpu, path: []const u8) !void { var e: Env = undefined; try env(gpu, &e); diff --git a/zig/tests/cuda/ling_layer_speed.zig b/zig/tests/cuda/ling_layer_speed.zig index 4ec40a2d8a..8902a7cbb7 100644 --- a/zig/tests/cuda/ling_layer_speed.zig +++ b/zig/tests/cuda/ling_layer_speed.zig @@ -15,17 +15,19 @@ const glue = G.glue; /// Events between op groups; each group's milliseconds by name, medians over reps. const Timer = struct { - ev: [48]cuda.Event, - names: [48][]const u8 = undefined, + ev: [64]cuda.Event, + names: [64][]const u8 = undefined, n: usize = 0, s: cuda.Stream, on: bool, + isolate: bool = false, fn mark(t: *Timer, name: []const u8) !void { if (!t.on) return; try t.ev[t.n].record(t.s); t.names[t.n] = name; t.n += 1; + if (t.isolate) try t.s.synchronize(); // each op group starts on an idle GPU: another context's slice lands between } }; @@ -40,6 +42,11 @@ fn kdaBlock(o: G.Ops, b: *const F.Buffers, a: W.Kda, c: *const ling.config.Confi const sv = st.saves(li); const p = if (b.prompt) b.kp else sv.p; for ([_]W.Proj{ a.q, a.k, a.v, a.b }, [_]usize{ 0, n, 2 * n, 3 * n }, [_][]const u8{ "proj q", "proj k", "proj v", "proj b (bf16)" }) |w, col0, name| { + if (rows == 1) { + try o.mm(w, b.normed, D, 1, p + col0 * 2); + try t.mark(name); + continue; + } try o.mm(w, b.normed, D, rows, b.tmp); try t.mark(name); const w_n = w.outputs(); @@ -120,6 +127,11 @@ fn mlp(o: G.Ops, b: *const F.Buffers, m: W.Mlp, rows: usize, limit: f32, out: u6 const D = m.gate.inputs(); const w = m.gate.outputs(); for ([_]W.Proj{ m.gate, m.up }, [_]usize{ 0, w }) |p, col0| { + if (rows == 1) { + try o.mm(p, b.normed, D, 1, b.gu + col0 * 2); + try t.mark(tag ++ " gate, up"); + continue; + } try o.mm(p, b.normed, D, rows, b.tmp); try t.mark(tag ++ " gate, up"); try o.copy(b.tmp, w * 2, b.gu + col0 * 2, 4 * w, rows, w * 2); @@ -142,26 +154,29 @@ fn moe(o: G.Ops, b: *const F.Buffers, e: W.Moe, g: *const F.Graph, rows: usize, try glue.select(o.set(), o.s, b.mlog, if (e.bias != 0) e.bias else g.zero_bias, b.pick, b.wts, rows, top, E, c.routed_scale, c.norm_topk, ling.groups(c)); try t.mark("select"); const l = e.routed.grouped; - const tile: usize = if (o.prompt) cuda.experts.promptTile(l.format) else 16; - try o.m.router.route(o.s, b.pick, rows * slots, E + 1, tile, b.plan); - try t.mark("expert plan"); - const items = cuda.grouped.maxItems(rows * slots, E + 1, tile); const gr = o.m.grouped(); const up: cuda.experts.Rows = .{ .x = b.normed, .stride = l.dims, .slots = slots }; const down: cuda.experts.Rows = .{ .x = b.eact, .stride = l.width }; const skip: i32 = @intCast(E); - if (o.prompt) try gr.prompt(o.s, .up, up, l, b.plan, b.eact, items, skip) else try gr.decode(o.s, .up, up, l, b.plan, b.eact, items, skip); - try t.mark("experts gate-up"); - if (o.prompt) try gr.prompt(o.s, .down_bf16, down, l, b.plan, b.ey, items, skip) else try gr.decode(o.s, .down_bf16, down, l, b.plan, b.ey, items, skip); - try t.mark("experts down"); - if (o.prompt) { // the decode kernels on 16-pair items over the same rows, for comparison - try o.m.router.route(o.s, b.pick, rows * slots, E + 1, 16, b.plan); - try t.mark("alt plan 16"); - const it16 = cuda.grouped.maxItems(rows * slots, E + 1, 16); - try gr.decode(o.s, .up, up, l, b.plan, b.eact, it16, skip); - try t.mark("alt gate-up decode kernel"); - try gr.decode(o.s, .down_bf16, down, l, b.plan, b.ey, it16, skip); - try t.mark("alt down decode kernel"); + // block_ffn's choice first, then (prompt rows) every kernel and plan tile on the same rows for comparison + const pk = o.prompt and !(l.format == .fp8g and rows * top < 16 * E); + const Way = struct { prompt: bool, tile: usize, up: []const u8, down: []const u8 }; + const ways = [_]Way{ + .{ .prompt = pk, .tile = if (pk) 64 else 16, .up = "experts gate-up", .down = "experts down" }, + .{ .prompt = true, .tile = 64, .up = "alt gate-up, prompt kernel, tile 64", .down = "alt down, prompt kernel, tile 64" }, + .{ .prompt = false, .tile = 16, .up = "alt gate-up, decode kernel, tile 16", .down = "alt down, decode kernel, tile 16" }, + .{ .prompt = false, .tile = 32, .up = "alt gate-up, decode kernel, tile 32", .down = "alt down, decode kernel, tile 32" }, + .{ .prompt = false, .tile = 64, .up = "alt gate-up, decode kernel, tile 64", .down = "alt down, decode kernel, tile 64" }, + }; + for (ways, 0..) |w, i| { + if (i > 0 and !o.prompt) break; + try o.m.router.route(o.s, b.pick, rows * slots, E + 1, w.tile, b.plan); + try t.mark(if (i == 0) "expert plan" else "alt plans"); + const items = cuda.grouped.maxItems(rows * slots, E + 1, w.tile); + if (w.prompt) try gr.prompt(o.s, .up, up, l, b.plan, b.eact, items, skip) else try gr.decode(o.s, .up, up, l, b.plan, b.eact, items, skip); + try t.mark(w.up); + if (w.prompt) try gr.prompt(o.s, .down_bf16, down, l, b.plan, b.ey, items, skip) else try gr.decode(o.s, .down_bf16, down, l, b.plan, b.ey, items, skip); + try t.mark(w.down); } try mlp(o, b, e.shared.?, rows, e.shared_limit, b.sh, t, "shared"); var a: cuda.Args = .{}; @@ -203,9 +218,13 @@ fn median(xs: []f64) f64 { return xs[xs.len / 2]; } -/// `ling-layer-speed DIR FIXTURE [REPS]`: decode at position 2048, windows 2..16, prompts of 512 and 2048 rows. -pub fn run(gpu: check.Gpu, dir: []const u8, path: []const u8, reps_arg: ?[]const u8) !void { - const reps: usize = if (reps_arg) |x| try std.fmt.parseInt(usize, x, 10) else 15; +/// `ling-layer-speed DIR FIXTURE [REPS [ROWS,..]]`: decode at position 2048, windows, prompts (default 512, 2048). +pub fn run(gpu: check.Gpu, dir: []const u8, path: []const u8, reps_arg: ?[]const u8, rows_arg: ?[]const u8) !void { + const reps: usize = @min(30, if (reps_arg) |x| try std.fmt.parseInt(usize, x, 10) else 15); + var phases: std.ArrayList(usize) = .empty; + var it = std.mem.tokenizeScalar(u8, rows_arg orelse "1,2,4,8,16,512,2048", ','); + while (it.next()) |x| try phases.append(gpu.gpa, try std.fmt.parseInt(usize, x, 10)); + defer phases.deinit(gpu.gpa); var arena = std.heap.ArenaAllocator.init(gpu.gpa); defer arena.deinit(); const a = arena.allocator(); @@ -220,28 +239,31 @@ pub fn run(gpu: check.Gpu, dir: []const u8, path: []const u8, reps_arg: ?[]const const st = try k.r.state(k.cap); defer st.deinit(); const c = &k.w.config; - std.debug.print("RESULT ling-layer-speed layer {d} ({t}, {s}, {t}), {d} reps, median and least us (other GPU contexts' time slices land in the medians)\n", .{ k.L, c.kind(k.L), if (c.isMoe(k.L)) "MoE" else "dense", c.quant, reps }); - for ([_]usize{ 1, 2, 4, 8, 16, 512, 2048 }) |rows| { + std.debug.print("RESULT ling-layer-speed layer {d} ({t}, {s}, {t}), {d} reps, median and least us; prompt ops each from an idle GPU (other contexts' time slices land in medians)\n", .{ k.L, c.kind(k.L), if (c.isMoe(k.L)) "MoE" else "dense", c.quant, reps }); + for (phases.items) |rows| { + if (rows == 0 or rows > 2048) return error.TooManyRows; const b = if (rows > F.window_rows) &k.pb else &k.db; const pos: usize = if (rows > F.window_rows) 0 else 2048; try st.setPos(s, pos); const d = gpu.d; - try d.check(d.api.cuMemcpyHtoDAsync_v2(b.x, k.x.ptr, rows * k.D * 2, s.handle), "cuMemcpyHtoDAsync"); const f = k.r.fwd(b, null); const sg: Seg = .{ .st = st, .row0 = 0, .rows = rows }; var whole: [64]f64 = undefined; var table: [32]Row = undefined; var nt: usize = 0; for (0..reps + 2) |i| { - // whole layer through forward.layer, then the same launches op by op - try t.ev[46].record(s); + // whole layer through forward.layer, then the same launches op by op, each rep on the fixture's rows + try d.check(d.api.cuMemcpyHtoDAsync_v2(b.x, k.x.ptr, rows * k.D * 2, s.handle), "cuMemcpyHtoDAsync"); + try f.layer(&.{sg}, 0, rows, null); // untimed: the GPU busy before the timed launches + try t.ev[62].record(s); try f.layer(&.{sg}, 0, rows, null); - try t.ev[47].record(s); + try t.ev[63].record(s); try s.synchronize(); - if (i >= 2) whole[i - 2] = try t.ev[46].elapsedMs(t.ev[47]); t.n = 0; + t.isolate = rows > F.window_rows; // prompt ops one at a time; a decode layer's run back to back try layerOps(&f, sg, rows, &t); try s.synchronize(); + if (i >= 2) whole[i - 2] = try t.ev[62].elapsedMs(t.ev[63]); if (i < 2) continue; var acc: [32]f64 = @splat(0); for (1..t.n) |j| { From bd7c23e3fe4c21bef54cd6e97de84060ce3c4ad2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 15 Oct 2026 17:25:42 +0200 Subject: [PATCH 20/34] tests: ling-layer-speed starts each rep after a random spin (tf_probe_spin), so other contexts' time slices miss the least --- zig/kernels/cuda/probe.cu | 9 ++++++++ zig/tests/cuda/ling_layer_speed.zig | 33 +++++++++++++++++++---------- 2 files changed, 31 insertions(+), 11 deletions(-) diff --git a/zig/kernels/cuda/probe.cu b/zig/kernels/cuda/probe.cu index ca8f345e7b..a3edf6de24 100644 --- a/zig/kernels/cuda/probe.cu +++ b/zig/kernels/cuda/probe.cu @@ -60,3 +60,12 @@ extern "C" __global__ void tf_probe_pdl_step(unsigned long long* counter, unsign asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); #endif } + +// One thread busy for `ns` of the global timer: a stream kept busy while the host queues what follows it. +extern "C" __global__ void tf_probe_spin(long long ns) { + long long t0, t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + do { + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + } while (t - t0 < ns); +} diff --git a/zig/tests/cuda/ling_layer_speed.zig b/zig/tests/cuda/ling_layer_speed.zig index 8902a7cbb7..6b36204dfc 100644 --- a/zig/tests/cuda/ling_layer_speed.zig +++ b/zig/tests/cuda/ling_layer_speed.zig @@ -13,21 +13,30 @@ const kda = G.kda; const latent = G.latent; const glue = G.glue; -/// Events between op groups; each group's milliseconds by name, medians over reps. +/// Events between op groups; each group's milliseconds by name, medians and least over reps. const Timer = struct { ev: [64]cuda.Event, names: [64][]const u8 = undefined, n: usize = 0, s: cuda.Stream, on: bool, - isolate: bool = false, + spin: cuda.Function, fn mark(t: *Timer, name: []const u8) !void { if (!t.on) return; try t.ev[t.n].record(t.s); t.names[t.n] = name; t.n += 1; - if (t.isolate) try t.s.synchronize(); // each op group starts on an idle GPU: another context's slice lands between + } + + fn spinFor(t: *Timer, ns: i64) !void { + var a: cuda.Args = .{}; + a.add(@as(c_longlong, ns)); + try cuda.launch.launch(t.spin, .{ .grid = .{}, .block = .{} }, t.s, &a); + } + + fn ms(t: *const Timer, j: usize) !f32 { + return t.ev[j - 1].elapsedMs(t.ev[j]); } }; @@ -233,13 +242,16 @@ pub fn run(gpu: check.Gpu, dir: []const u8, path: []const u8, reps_arg: ?[]const var k: bytes.Kit = undefined; try k.init(gpu, a, s, dir, path, 2048, 4200); defer k.deinit(); - var t: Timer = .{ .ev = undefined, .s = s, .on = true }; + var probe = try cuda.Module.load(gpu.d, cuda.kernels.probe); + defer probe.unload(); + var t: Timer = .{ .ev = undefined, .s = s, .on = true, .spin = try probe.function("tf_probe_spin") }; for (&t.ev) |*e| e.* = try cuda.Event.init(gpu.d, true); defer for (&t.ev) |*e| e.deinit(); + var prng = std.Random.DefaultPrng.init(0x5eed); const st = try k.r.state(k.cap); defer st.deinit(); const c = &k.w.config; - std.debug.print("RESULT ling-layer-speed layer {d} ({t}, {s}, {t}), {d} reps, median and least us; prompt ops each from an idle GPU (other contexts' time slices land in medians)\n", .{ k.L, c.kind(k.L), if (c.isMoe(k.L)) "MoE" else "dense", c.quant, reps }); + std.debug.print("RESULT ling-layer-speed layer {d} ({t}, {s}, {t}), {d} reps, median and least us (each rep after a random spin: other contexts' slices land at random)\n", .{ k.L, c.kind(k.L), if (c.isMoe(k.L)) "MoE" else "dense", c.quant, reps }); for (phases.items) |rows| { if (rows == 0 or rows > 2048) return error.TooManyRows; const b = if (rows > F.window_rows) &k.pb else &k.db; @@ -252,15 +264,14 @@ pub fn run(gpu: check.Gpu, dir: []const u8, path: []const u8, reps_arg: ?[]const var table: [32]Row = undefined; var nt: usize = 0; for (0..reps + 2) |i| { - // whole layer through forward.layer, then the same launches op by op, each rep on the fixture's rows + // a spin of random length first: other contexts' time slices land at a new phase each rep, least is clean try d.check(d.api.cuMemcpyHtoDAsync_v2(b.x, k.x.ptr, rows * k.D * 2, s.handle), "cuMemcpyHtoDAsync"); - try f.layer(&.{sg}, 0, rows, null); // untimed: the GPU busy before the timed launches + try t.spinFor(prng.random().intRangeLessThan(i64, 0, 4_000_000)); + try f.layer(&.{sg}, 0, rows, null); // untimed: the GPU warm before the timed launches try t.ev[62].record(s); try f.layer(&.{sg}, 0, rows, null); try t.ev[63].record(s); - try s.synchronize(); - t.n = 0; - t.isolate = rows > F.window_rows; // prompt ops one at a time; a decode layer's run back to back + t.n = 0; // then the same launches op by op, back to back try layerOps(&f, sg, rows, &t); try s.synchronize(); if (i >= 2) whole[i - 2] = try t.ev[62].elapsedMs(t.ev[63]); @@ -274,7 +285,7 @@ pub fn run(gpu: check.Gpu, dir: []const u8, path: []const u8, reps_arg: ?[]const table[nt] = .{ .name = name }; nt += 1; } - acc[at] += try t.ev[j - 1].elapsedMs(t.ev[j]); + acc[at] += try t.ms(j); } for (table[0..nt], acc[0..nt]) |*r, v| { r.ms[r.n] = v; From 32055e9e1f1bec39964f7d9120858bc9ab3b1b45 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 15 Oct 2026 20:44:48 +0200 Subject: [PATCH 21/34] glm: swiglu keeps GLM's limit as given; only Ling's swigluWith maps 0 to unclamped At 6380c8da GLM's glue swiglu (like its Metal glm_rows_glu) clamps at the limit it is given, 0 included. 51c01573 made the glue map 0 to +inf for every caller; that mapping now applies to Ling's kernels only. - glm-route-old: GLM's swiglu equals 6380c8da's kernels at limits 10 and 0 (0 was not compared before). - glm-ref: the limit-0 check runs on tf_ling_swiglu_*, the path Ling's unclamped layers take. --- zig/src/families/glm/cuda_glue.zig | 12 ++++++++---- zig/tests/cuda/glm_ref_route.zig | 20 ++++++++++---------- zig/tests/cuda/glm_ref_tests.zig | 2 +- 3 files changed, 19 insertions(+), 15 deletions(-) diff --git a/zig/src/families/glm/cuda_glue.zig b/zig/src/families/glm/cuda_glue.zig index 01c71a3f73..7dc0cb854a 100644 --- a/zig/src/families/glm/cuda_glue.zig +++ b/zig/src/families/glm/cuda_glue.zig @@ -126,17 +126,21 @@ fn swigluBlock(w: usize) !usize { return error.NoSwigluBlock; } -/// glue.swiglu: [gate | up] -> bf16(bf16(silu(min(g, L))) * clip(u, -L, L)), 64-sums in `xs`; L per layer, 0 for none. +/// glue.swiglu: [gate | up] -> bf16(bf16(silu(min(g, L))) * clip(u, -L, L)), 64-sums in `xs`; L as given (GLM's). pub fn swiglu(set: Set, stream: cuda.Stream, gu: u64, out: u64, xs: u64, limit: f32, rows: usize, w: usize) !void { - return swigluWith(.{ set.glue.swiglu_512, set.glue.swiglu_64 }, stream, gu, out, xs, limit, rows, w); + return swigluLaunch(.{ set.glue.swiglu_512, set.glue.swiglu_64 }, stream, gu, out, xs, limit, rows, w); } -/// swiglu on a pair of glm_glue_moe.cuh's 512 and 64 kernels (tf_ling_swiglu_*: the limit on silu's output, Ling). +/// swiglu on a pair of glm_glue_moe.cuh's 512 and 64 kernels (tf_ling_swiglu_*: the limit on silu's output, Ling); 0 for none. pub fn swigluWith(fs: [2]cuda.Function, stream: cuda.Stream, gu: u64, out: u64, xs: u64, limit: f32, rows: usize, w: usize) !void { + return swigluLaunch(fs, stream, gu, out, xs, clampOf(limit), rows, w); +} + +fn swigluLaunch(fs: [2]cuda.Function, stream: cuda.Stream, gu: u64, out: u64, xs: u64, clamp: f32, rows: usize, w: usize) !void { const block = try swigluBlock(w); const f, const threads: u32 = if (block == 512) .{ fs[0], 128 } else .{ fs[1], 64 }; var a: cuda.Args = .{}; - inline for (.{ gu, out, xs, clampOf(limit), i64of(w) }) |v| a.add(v); + inline for (.{ gu, out, xs, clamp, i64of(w) }) |v| a.add(v); try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(w / block) }, .block = .{ .x = threads } }, stream, &a); } diff --git a/zig/tests/cuda/glm_ref_route.zig b/zig/tests/cuda/glm_ref_route.zig index cc8e07e3a0..c1a4b9439a 100644 --- a/zig/tests/cuda/glm_ref_route.zig +++ b/zig/tests/cuda/glm_ref_route.zig @@ -187,8 +187,8 @@ pub fn groups(r: *Rig, set: glue.Set) !void { check.pass("group select: picks follow the float64 rule (exact ties to the lower id, near ties counted), weights within 2^-16", .{}); } -/// swiglu at limit 0 (Ling's unclamped layers) against float64 unclamped, and limit 10 differing where it clamps. -pub fn swigluLimit(r: *Rig, set: glue.Set) !void { +/// Ling's swiglu at limit 0 (its unclamped layers) against float64 unclamped, and limit 10 differing where it clamps. +pub fn swigluLimit(r: *Rig, fs: [2]cuda.Function) !void { for ([_]usize{ 1024, 704 }) |w| { const rows = 40; const gu = try r.bfs(rows * 2 * w, 16.0); @@ -196,8 +196,8 @@ pub fn swigluLimit(r: *Rig, set: glue.Set) !void { const out0 = try r.zeros(rows * w * 2); const out10 = try r.zeros(rows * w * 2); const xs = try r.zeros(rows * w / 64 * 4); - try glue.swiglu(set, r.s, dgu, out0, xs, 0, rows, w); - try glue.swiglu(set, r.s, dgu, out10, xs, 10, rows, w); + try glue.swigluWith(fs, r.s, dgu, out0, xs, 0, rows, w); + try glue.swigluWith(fs, r.s, dgu, out10, xs, 10, rows, w); const raw0 = try r.back(out0, rows * w * 2); const raw10 = try r.back(out10, rows * w * 2); var most: f64 = 0; @@ -214,7 +214,7 @@ pub fn swigluLimit(r: *Rig, set: glue.Set) !void { std.debug.print("RESULT swiglu width {d}, limit 0: unclamped float64 at {e:.3} of 2^-7 relative; limit 10 changes {d} of {d}\n", .{ w, most, clamped, rows * w }); try check.expect(most <= 1 and clamped > 0, "swiglu limit 0 is unclamped, limit 10 clamps", .{}); } - check.pass("swiglu: the limit is per call, 0 runs unclamped (as the experts' 0), 10 clamps", .{}); + check.pass("Ling swiglu: the limit is per call, 0 runs unclamped (as the experts' 0), 10 clamps", .{}); } const Old = struct { topk: cuda.Function, swiglu_512: cuda.Function, swiglu_64: cuda.Function }; @@ -240,20 +240,20 @@ fn glmBytes(r: *Rig, set: glue.Set, prev: Old) !void { try check.sameBytes("GLM select picks vs 6380c8da", try r.back(pick_new, n * 4), try r.back(pick, n * 4)); try check.sameBytes("GLM select weights vs 6380c8da", try r.back(wts_new, n * 4), try r.back(wts, n * 4)); }; - for ([_]usize{ 2048, 704 }) |w| { + for ([_]f32{ 10, 0 }) |lim| for ([_]usize{ 2048, 704 }) |w| { const rows = 40; const dgu = try r.dev(std.mem.sliceAsBytes(try r.bfs(rows * 2 * w, 16.0))); const outs = [2]u64{ try r.zeros(rows * w * 2), try r.zeros(rows * w * 2) }; const xss = [2]u64{ try r.zeros(rows * w / 64 * 4), try r.zeros(rows * w / 64 * 4) }; const block: usize = if (w % 512 == 0) 512 else 64; var a: cuda.Args = .{}; - inline for (.{ dgu, outs[0], xss[0], @as(f32, 10.0), @as(i64, @intCast(w)) }) |v| a.add(v); + inline for (.{ dgu, outs[0], xss[0], lim, @as(i64, @intCast(w)) }) |v| a.add(v); try cuda.launch.launch(if (block == 512) prev.swiglu_512 else prev.swiglu_64, .{ .grid = .{ .x = rows, .y = @intCast(w / block) }, .block = .{ .x = if (block == 512) 128 else 64 } }, r.s, &a); - try glue.swiglu(set, r.s, dgu, outs[1], xss[1], 10, rows, w); + try glue.swiglu(set, r.s, dgu, outs[1], xss[1], lim, rows, w); try check.sameBytes("GLM swiglu vs 6380c8da", try r.back(outs[1], rows * w * 2), try r.back(outs[0], rows * w * 2)); try check.sameBytes("GLM swiglu sums vs 6380c8da", try r.back(xss[1], rows * w / 64 * 4), try r.back(xss[0], rows * w / 64 * 4)); - } - check.pass("GLM select (288, 256, 160 experts, normed and not) and swiglu (limit 10): every byte equals 6380c8da's kernels", .{}); + }; + check.pass("GLM select (288, 256, 160 experts, normed and not) and swiglu (limits 10 and 0): every byte equals 6380c8da's kernels", .{}); } const Env = struct { diff --git a/zig/tests/cuda/glm_ref_tests.zig b/zig/tests/cuda/glm_ref_tests.zig index 724cec58f4..ad8b7e8e2a 100644 --- a/zig/tests/cuda/glm_ref_tests.zig +++ b/zig/tests/cuda/glm_ref_tests.zig @@ -197,7 +197,7 @@ pub fn run(gpu: Gpu) !void { for ([_]Shape{ .{ .n = 512, .k = 4096 }, .{ .n = 1280, .k = 1536 }, .{ .n = 1024, .k = 128 } }) |sh| try matmuls(&r, set, sh); try routers(&r, set); try @import("glm_ref_route.zig").groups(&r, set); - try @import("glm_ref_route.zig").swigluLimit(&r, set); + try @import("glm_ref_route.zig").swigluLimit(&r, .{ try mod.function("tf_ling_swiglu_512"), try mod.function("tf_ling_swiglu_64") }); try @import("glm_ref_attn.zig").dense(&r, set); try @import("glm_ref_attn.zig").sparse(&r, set); try @import("glm_ref_hc.zig").mixing(&r, set); From 991383f39ec73e4ed7b976be92dabb419ba92a3b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 15 Oct 2026 21:34:43 +0200 Subject: [PATCH 22/34] glm: swigluWith's doc comment within 120 columns --- zig/src/families/glm/cuda_glue.zig | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zig/src/families/glm/cuda_glue.zig b/zig/src/families/glm/cuda_glue.zig index 7dc0cb854a..24da99eea3 100644 --- a/zig/src/families/glm/cuda_glue.zig +++ b/zig/src/families/glm/cuda_glue.zig @@ -131,7 +131,7 @@ pub fn swiglu(set: Set, stream: cuda.Stream, gu: u64, out: u64, xs: u64, limit: return swigluLaunch(.{ set.glue.swiglu_512, set.glue.swiglu_64 }, stream, gu, out, xs, limit, rows, w); } -/// swiglu on a pair of glm_glue_moe.cuh's 512 and 64 kernels (tf_ling_swiglu_*: the limit on silu's output, Ling); 0 for none. +/// swiglu on a pair of glm_glue_moe.cuh's 512 and 64 kernels (tf_ling_swiglu_*, the limit on silu's output); 0: none. pub fn swigluWith(fs: [2]cuda.Function, stream: cuda.Stream, gu: u64, out: u64, xs: u64, limit: f32, rows: usize, w: usize) !void { return swigluLaunch(fs, stream, gu, out, xs, clampOf(limit), rows, w); } From 03171bf55a3035983b455bf9cd8afec637519531 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 15 Oct 2026 21:34:43 +0200 Subject: [PATCH 23/34] glm: drop the opt-in 16-row chunked KDA prompt path; GLM's cuda_kda.zig is 6380c8da's again Nothing served called it: GLM and Ling run the row path, which resumes at any row. Removed with it: glm_kda_chunk.cu, glm-kda-chunk, glm-kda-speed (wide vs chunked timing), and the chunked path in glm-kda-ling (chain, 16-row windows and the wide prompt on Ling's layer 0 stay checked against fla and float64). --- docs/recipes/ling-3.1.md | 2 +- zig/build/cuda.zig | 1 - zig/kernels/cuda/glm_kda_chunk.cu | 237 ------------------------- zig/src/cuda/kernels.zig | 1 - zig/src/families/glm/cuda_kda.zig | 84 +++------ zig/tests/cuda/glm_kda_chunk_tests.zig | 124 ------------- zig/tests/cuda/glm_kda_f64.zig | 2 +- zig/tests/cuda/glm_kda_ling.zig | 20 +-- zig/tests/cuda/glm_kda_rig.zig | 14 +- zig/tests/cuda/glm_kda_speed.zig | 76 -------- zig/tests/cuda/main.zig | 6 - 11 files changed, 32 insertions(+), 535 deletions(-) delete mode 100644 zig/kernels/cuda/glm_kda_chunk.cu delete mode 100644 zig/tests/cuda/glm_kda_chunk_tests.zig delete mode 100644 zig/tests/cuda/glm_kda_speed.zig diff --git a/docs/recipes/ling-3.1.md b/docs/recipes/ling-3.1.md index 05505341ca..1fe89c41f6 100644 --- a/docs/recipes/ling-3.1.md +++ b/docs/recipes/ling-3.1.md @@ -16,7 +16,7 @@ block FP8 (e4m3 in 128 x 128 blocks with `weight_scale_inv`). | --- | --- | --- | | Group-limited top-k routing, per-layer shared SwiGLU limit (GLM's bytes unchanged) | `ling-family`, commit 51c01573 | `glm-ref`, `glm-route-old`, `glm-route-ling` | | MLA at Ling's shape: 64 roped key columns, 128-wide heads, GPT-J rope, head-wise gate | f8e594f6, merged | `ling-mla`, `ling-mla-real`, `glm-mla-same` | -| KDA on Ling's layer 0 (q/k/v convs as one q\|k\|v conv), opt-in chunked prompt path | b55aa13d, merged | `glm-kda-ling`, `glm-kda-chunk` | +| KDA on Ling's layer 0 (q/k/v convs as one q\|k\|v conv), GLM's row path as is | b55aa13d, merged | `glm-kda-ling` | | Config loader and weight map (all 214,025 names of the real index) | `zig/src/families/ling/config.zig`, `names.zig` | `zig build test-ling` | | Weight loader at TP=1: bf16, FP8 block, ModelOpt NVFP4 | `zig/src/families/ling/weights.zig`, `weights_proj.zig` | `ling-weights`, `ling-weights-real` | | Synthetic checkpoints in all three formats | `tf-cuda-test ling-synth` | its own name check | diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 4b1c865ad3..b4a0cf9c4f 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -49,7 +49,6 @@ const kernels = [_]Kernel{ .{ .name = "prompt16", .flags = &.{"-O3"} }, // nvfp4/prompt.cu's FP4 and FP8 tile 4, tensorfold_nvfp4_prompt_v1 .{ .name = "modelopt_pack", .flags = &.{"-O3"} }, // ours: ModelOpt tensors into the lane, prompt and expert layouts .{ .name = "glm_kda", .flags = &.{ "-O3", "--fmad=false" } }, // GLM: the Python family's glm5_next/cuda/kda.cu - .{ .name = "glm_kda_chunk", .flags = &.{ "-O3", "--fmad=false" } }, // ours: glm_kda's prompt recurrence in 16-row WY chunks .{ .name = "glm_glue", .flags = &.{"-O3"} }, // ours: the Python family's Triton glue by hand, its PTX op for op .{ .name = "glm_ops", .flags = &.{"-O3"} }, // ours: GLM's integer work the Python family did in torch (the DSA token lists) .{ .name = "glm_quant", .flags = torch_ops }, // ours: GLM's draft-only weight quantizers, torch at load in the Python family diff --git a/zig/kernels/cuda/glm_kda_chunk.cu b/zig/kernels/cuda/glm_kda_chunk.cu deleted file mode 100644 index 8f705777cb..0000000000 --- a/zig/kernels/cuda/glm_kda_chunk.cu +++ /dev/null @@ -1,237 +0,0 @@ -// GLM's KDA prompt recurrence in 16-row WY chunks over glm_kda.cu's prep saves; opt-in, -O3 --fmad=false, _rn ops. - -// A block a head, chunks from the call's first row: calls cut at multiples of 16 rows keep one call's bits. - -// Not the serial row path's bits (other sums): a prompt resumed off a multiple of 16 differs from a fresh one. - -// Thread t holds S rows 4 (t / 8) .. + 3 (value), keys 4 (t % 8) + 32 e + 0..3; every sum runs in one fixed order. - -#include -#include -#include -namespace { - -constexpr int D = 128; // key and value dims a head -constexpr int C = 16; // rows a chunk -constexpr int NT = 256; // threads a block -constexpr int PARTS = 8; // threads sharing a value row (key parts) -constexpr int R = 4; // value rows a thread -constexpr int KG = 4; // float4 key groups a thread -constexpr unsigned FULL = 0xffffffffu; - -// N per-lane partials summed across the warp in one fixed order; lanes 32 j / N .. 32 (j + 1) / N - 1 hold total j. -template -__device__ __forceinline__ float spread(float (&v)[N], int lane) { -#pragma unroll - for (int m = 16, h = N / 2; m > 0; m >>= 1, h >>= 1) { - if (h >= 1) { - const bool up = lane & m; -#pragma unroll - for (int i = 0; i < h; ++i) { - const float give = up ? v[i] : v[h + i], keep = up ? v[h + i] : v[i]; - v[i] = __fadd_rn(keep, __shfl_xor_sync(FULL, give, m)); - } - } else { - v[0] = __fadd_rn(v[0], __shfl_xor_sync(FULL, v[0], m)); - } - } - return v[0]; -} - -struct Chunk { - float q[C][D]; // q, then its decay from the chunk's start times q - float k[C][D]; // k, then its decay to the chunk's end times k - float g[C][D]; // decay factors, then (decayed q) S - float w[C][D]; // beta (decay from the start) k, then W = T of that - float v[C][D]; // v, then beta v, then U = T (beta v), then U - W S - float l[C][C]; // beta_i k_i . (decay j+1 .. i) k_j below the diagonal - float m[C][C]; // q_i . (decay j+1 .. i) k_j on and below it - float gl[D]; // the chunk's whole decay - float beta[C]; -}; - -// out[i][v] (sub: minus) x_i . S row v: a thread's keys in order, 8 parts by xor butterfly (same bits in every lane). -__device__ __forceinline__ void rows_dot(const float (*x)[D], const float (&s)[R][KG][4], int p, int vg, bool sub, float (*out)[D]) { -#pragma unroll - for (int i = 0; i < C; ++i) { - float acc[R] = {0.f, 0.f, 0.f, 0.f}; -#pragma unroll - for (int e = 0; e < KG; ++e) { - const float4 a = *reinterpret_cast(&x[i][4 * p + 32 * e]); -#pragma unroll - for (int r = 0; r < R; ++r) { - acc[r] = __fmaf_rn(a.x, s[r][e][0], acc[r]); - acc[r] = __fmaf_rn(a.y, s[r][e][1], acc[r]); - acc[r] = __fmaf_rn(a.z, s[r][e][2], acc[r]); - acc[r] = __fmaf_rn(a.w, s[r][e][3], acc[r]); - } - } -#pragma unroll - for (int r = 0; r < R; ++r) -#pragma unroll - for (int m = 1; m < PARTS; m <<= 1) acc[r] = __fadd_rn(acc[r], __shfl_xor_sync(FULL, acc[r], m)); - float mine = acc[0]; -#pragma unroll - for (int r = 1; r < R; ++r) mine = p == r ? acc[r] : mine; - if (p < R) out[i][R * vg + p] = sub ? __fsub_rn(out[i][R * vg + p], mine) : mine; - } -} - -// Block (head); q_in null: the state alone (a chunked replay of saved rows, the same state bits), y_out unwritten. -__device__ __forceinline__ void chunk_body(int H, const float* __restrict__ state_in, const float* __restrict__ q_in, - const float* __restrict__ k_save, const __nv_bfloat16* __restrict__ v_save, - const float* __restrict__ g_save, const float* __restrict__ b_save, int rows, - __nv_bfloat16* __restrict__ y_out, float* __restrict__ state_out) { - __shared__ __align__(16) Chunk sh; - const int t = threadIdx.x, warp = t >> 5, lane = t & 31; - const int h = blockIdx.x; - const bool out = q_in != nullptr; - const int vg = t / PARTS, p = t % PARTS; - const size_t sbase = (size_t)h * D * D + (size_t)R * vg * D + 4 * p; - float s[R][KG][4]; -#pragma unroll - for (int r = 0; r < R; ++r) -#pragma unroll - for (int e = 0; e < KG; ++e) { - const float4 x = *reinterpret_cast(state_in + sbase + r * D + 32 * e); - s[r][e][0] = x.x, s[r][e][1] = x.y, s[r][e][2] = x.z, s[r][e][3] = x.w; - } - for (int c0 = 0; c0 < rows; c0 += C) { - const int n = min(C, rows - c0); - constexpr int LD = C * D / NT; - float lk[LD], lq[LD], lg[LD], lv[LD]; // every load of the chunk issued before any store -#pragma unroll - for (int e = 0; e < LD; ++e) { // the chunk's saved rows; rows past the call: k, q, v 0, decay 1, beta 0 - const int x = t + e * NT, i = x / D; - const size_t at = ((size_t)(c0 + i) * H + h) * D + x % D; - lk[e] = i < n ? k_save[at] : 0.f; - lq[e] = i < n && out ? q_in[at] : 0.f; - lg[e] = i < n ? g_save[at] : 1.f; - lv[e] = i < n ? __bfloat162float(v_save[at]) : 0.f; - } -#pragma unroll - for (int e = 0; e < LD; ++e) { - const int i = (t + e * NT) / D, d = (t + e * NT) % D; - sh.k[i][d] = lk[e], sh.q[i][d] = lq[e], sh.g[i][d] = lg[e], sh.v[i][d] = lv[e]; - } - if (t < C) sh.beta[t] = t < n ? b_save[(c0 + t) * H + h] : 0.f; - { - const int i = t / C, j = t % C; - if (j > i) sh.l[i][j] = 0.f, sh.m[i][j] = 0.f; - if (j == i) sh.l[i][j] = 0.f; - } - __syncthreads(); -#pragma unroll - for (int jj = 0; jj < 2; ++jj) { // columns j = warp and 15 - warp: e = prod a_(j+1 .. i), a running product per channel - const int j = jj == 0 ? warp : C - 1 - warp; - float kj[4], ej[4], part[2 * C]; -#pragma unroll - for (int e = 0; e < 4; ++e) kj[e] = sh.k[j][lane + 32 * e], ej[e] = 1.f; -#pragma unroll - for (int i = 0; i < C; ++i) { - float a = 0.f, m = 0.f; - if (i >= j) { -#pragma unroll - for (int e = 0; e < 4; ++e) { - const int d = lane + 32 * e; - if (i > j) ej[e] = __fmul_rn(ej[e], sh.g[i][d]); - const float ke = __fmul_rn(kj[e], ej[e]); - a = __fmaf_rn(sh.k[i][d], ke, a); - m = __fmaf_rn(sh.q[i][d], ke, m); - } - } - part[i] = a, part[C + i] = m; - } - const float tot = spread<2 * C>(part, lane); // lane i: row i's a, lane 16 + i: row i's m - const int i = lane % C; - if (lane < C && j < i) sh.l[i][j] = __fmul_rn(sh.beta[i], tot); - if (lane >= C && j <= i) sh.m[i][j] = tot; - } - __syncthreads(); - if (t < D) { // a channel's prefix products (q, beta k) and suffix products (the keys' decay to the chunk's end) - float pre = 1.f; -#pragma unroll - for (int i = 0; i < C; ++i) { - pre = __fmul_rn(pre, sh.g[i][t]); - sh.w[i][t] = __fmul_rn(sh.beta[i], __fmul_rn(pre, sh.k[i][t])); - sh.q[i][t] = __fmul_rn(pre, sh.q[i][t]); - } - sh.gl[t] = pre; - float suf = 1.f; -#pragma unroll - for (int i = C - 1; i >= 0; --i) { - sh.k[i][t] = __fmul_rn(sh.k[i][t], suf); - suf = __fmul_rn(suf, sh.g[i][t]); - } - } else { -#pragma unroll - for (int i = 0; i < C; ++i) sh.v[i][t - D] = __fmul_rn(sh.beta[i], sh.v[i][t - D]); - } - __syncthreads(); - { // (I + L) X = R by forward substitution, a thread a column of [W | U] - float* col = t < D ? &sh.w[0][t] : &sh.v[0][t - D]; - float x[C]; -#pragma unroll - for (int i = 0; i < C; ++i) { - float acc = col[i * D]; -#pragma unroll - for (int j = 0; j < i; ++j) acc = __fmaf_rn(-sh.l[i][j], x[j], acc); - x[i] = acc; - } -#pragma unroll - for (int i = 0; i < C; ++i) col[i * D] = x[i]; - } - __syncthreads(); - rows_dot(sh.w, s, p, vg, true, sh.v); // v: U - W S - if (out) rows_dot(sh.q, s, p, vg, false, sh.g); // g: (decayed q) S - __syncthreads(); - if (out) { // row i's read-out (decayed q) S + sum_j<=i M_ij (U - W S)_j: value row t % 128, rows 8 (t / 128) .. + 7 - const int v = t % D; -#pragma unroll - for (int ii = 0; ii < C / 2; ++ii) { - const int i = C / 2 * (t / D) + ii; - float o = sh.g[i][v]; - for (int j = 0; j <= i; ++j) o = __fmaf_rn(sh.m[i][j], sh.v[j][v], o); - if (i < n) y_out[((size_t)(c0 + i) * H + h) * D + v] = __float2bfloat16_rn(o); - } - } -#pragma unroll - for (int e = 0; e < KG; ++e) { // S = S diag(chunk decay) + sum_j (U - W S)_j (decayed k_j)^T - const int k = 4 * p + 32 * e; - const float4 g4 = *reinterpret_cast(&sh.gl[k]); - float x[R][4]; -#pragma unroll - for (int r = 0; r < R; ++r) - x[r][0] = __fmul_rn(s[r][e][0], g4.x), x[r][1] = __fmul_rn(s[r][e][1], g4.y), x[r][2] = __fmul_rn(s[r][e][2], g4.z), x[r][3] = __fmul_rn(s[r][e][3], g4.w); -#pragma unroll - for (int j = 0; j < C; ++j) { - const float4 kd = *reinterpret_cast(&sh.k[j][k]); - const float4 vn = *reinterpret_cast(&sh.v[j][R * vg]); - const float u[R] = {vn.x, vn.y, vn.z, vn.w}; -#pragma unroll - for (int r = 0; r < R; ++r) - x[r][0] = __fmaf_rn(kd.x, u[r], x[r][0]), x[r][1] = __fmaf_rn(kd.y, u[r], x[r][1]), x[r][2] = __fmaf_rn(kd.z, u[r], x[r][2]), x[r][3] = __fmaf_rn(kd.w, u[r], x[r][3]); - } -#pragma unroll - for (int r = 0; r < R; ++r) -#pragma unroll - for (int c = 0; c < 4; ++c) s[r][e][c] = x[r][c]; - } - __syncthreads(); - } -#pragma unroll - for (int r = 0; r < R; ++r) -#pragma unroll - for (int e = 0; e < KG; ++e) - *reinterpret_cast(state_out + sbase + r * D + 32 * e) = make_float4(s[r][e][0], s[r][e][1], s[r][e][2], s[r][e][3]); -} - -} // namespace - -// glm_kda.cu's step entry, chunked: q, k, v, decay, beta rows [rows][H][128] from prep; y bf16, the state after. -extern "C" __global__ void __launch_bounds__(NT, 2) tf_glm_kda_chunk( - int H, const float* __restrict__ state_in, const float* __restrict__ q_in, const float* __restrict__ k_save, - const __nv_bfloat16* __restrict__ v_save, const float* __restrict__ g_save, const float* __restrict__ b_save, - int rows, __nv_bfloat16* __restrict__ y_out, float* __restrict__ state_out) { - chunk_body(H, state_in, q_in, k_save, v_save, g_save, b_save, rows, y_out, state_out); -} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 34deca60cc..dfb0f74ed4 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -41,7 +41,6 @@ pub const glm_ops: []const u8 = if (available) &Blob("fatbin_glm_ops").bytes els pub const glm_quant: []const u8 = if (available) &Blob("fatbin_glm_quant").bytes else &.{}; pub const glm_kda: []const u8 = if (available) &Blob("fatbin_glm_kda").bytes else &.{}; pub const ling_glue: []const u8 = if (available) &Blob("fatbin_ling_glue").bytes else &.{}; -pub const glm_kda_chunk: []const u8 = if (available) &Blob("fatbin_glm_kda_chunk").bytes else &.{}; pub const torch_argmax: []const u8 = if (available) &Blob("fatbin_torch_argmax").bytes else &.{}; pub const torch_topk: []const u8 = if (available) &Blob("fatbin_torch_topk").bytes else &.{}; pub const torch_pointwise: []const u8 = if (available) &Blob("fatbin_torch_pointwise").bytes else &.{}; diff --git a/zig/src/families/glm/cuda_kda.zig b/zig/src/families/glm/cuda_kda.zig index c5b58d28d1..cf2388e640 100644 --- a/zig/src/families/glm/cuda_kda.zig +++ b/zig/src/families/glm/cuda_kda.zig @@ -54,7 +54,18 @@ pub const Wide = struct { q: u64, y: u64 }; pub fn chain(ks: *const Kernels, stream: cuda.Stream, c: Chain, rows: usize, s: Saves, wide: ?Wide) !void { const h: c_int = @intCast(c.heads); if (wide) |w| { - try prep(ks, stream, c, rows, s, w); + var a: cuda.Args = .{}; + a.add(h); + a.add(c.p); + a.add(@as(c_int, @intCast(c.p_stride))); + a.add(@as(c_int, @intCast(c.b_off))); + a.add(c.a); + a.add(@as(c_int, @intCast(c.a_stride))); + for ([_]u64{ c.conv_state, c.conv_w, c.a_log, c.dt_bias }) |v| a.add(v); + a.add(c.lower); + for ([_]u64{ w.q, s.k, s.v, s.g, s.b }) |v| a.add(v); + const grid: cuda.launch.Dim3 = .{ .x = @intCast(rows), .y = @intCast(c.heads) }; + try cuda.launch.launch(ks.prep, .{ .grid = grid, .block = .{ .x = 512 } }, stream, &a); var b: cuda.Args = .{}; b.add(h); for ([_]u64{ c.state_in, w.q, s.k, s.v, s.g, s.b }) |v| b.add(v); @@ -62,7 +73,15 @@ pub fn chain(ks: *const Kernels, stream: cuda.Stream, c: Chain, rows: usize, s: b.add(w.y); b.add(s.state_out); try cuda.launch.launch(ks.step, .{ .grid = .{ .x = @intCast(c.heads), .y = dv / 4 / 8 }, .block = .{ .x = 8 * 32 } }, stream, &b); - return out(ks, stream, c, rows, s, w); + var o: cuda.Args = .{}; + o.add(h); + o.add(w.y); + o.add(c.g); + o.add(@as(c_int, @intCast(c.g_stride))); + o.add(c.norm_w); + o.add(c.eps); + o.add(s.out); + return cuda.launch.launch(ks.out, .{ .grid = grid, .block = .{ .x = dv } }, stream, &o); } var a: cuda.Args = .{}; a.add(h); @@ -81,67 +100,6 @@ pub fn chain(ks: *const Kernels, stream: cuda.Stream, c: Chain, rows: usize, s: try cuda.launch.launch(ks.chain, .{ .grid = .{ .x = @intCast(c.heads) }, .block = .{ .x = 1024 } }, stream, &a); } -/// The wide chain's first kernel: conv, decay, beta and q/k norms of `rows` rows into w.q and the saves. -pub fn prep(ks: *const Kernels, stream: cuda.Stream, c: Chain, rows: usize, s: Saves, w: Wide) !void { - var a: cuda.Args = .{}; - a.add(@as(c_int, @intCast(c.heads))); - a.add(c.p); - a.add(@as(c_int, @intCast(c.p_stride))); - a.add(@as(c_int, @intCast(c.b_off))); - a.add(c.a); - a.add(@as(c_int, @intCast(c.a_stride))); - for ([_]u64{ c.conv_state, c.conv_w, c.a_log, c.dt_bias }) |v| a.add(v); - a.add(c.lower); - for ([_]u64{ w.q, s.k, s.v, s.g, s.b }) |v| a.add(v); - const grid: cuda.launch.Dim3 = .{ .x = @intCast(rows), .y = @intCast(c.heads) }; - try cuda.launch.launch(ks.prep, .{ .grid = grid, .block = .{ .x = 512 } }, stream, &a); -} - -/// The wide chain's last kernel: the gated RMSNorm of w.y into s.out. -pub fn out(ks: *const Kernels, stream: cuda.Stream, c: Chain, rows: usize, s: Saves, w: Wide) !void { - var o: cuda.Args = .{}; - o.add(@as(c_int, @intCast(c.heads))); - o.add(w.y); - o.add(c.g); - o.add(@as(c_int, @intCast(c.g_stride))); - o.add(c.norm_w); - o.add(c.eps); - o.add(s.out); - const grid: cuda.launch.Dim3 = .{ .x = @intCast(rows), .y = @intCast(c.heads) }; - return cuda.launch.launch(ks.out, .{ .grid = grid, .block = .{ .x = dv } }, stream, &o); -} - -/// glm_kda_chunk.cu's entry: the prompt recurrence in 16-row WY chunks over prep's saves. -pub const Chunked = struct { - f: cuda.Function, - - pub fn load(m: cuda.Module) !Chunked { - return .{ .f = try m.function("tf_glm_kda_chunk") }; - } -}; - -fn chunkLaunch(ck: Chunked, stream: cuda.Stream, heads: usize, state_in: u64, q: u64, k: u64, v: u64, g: u64, b: u64, rows: usize, y: u64, state_out: u64) !void { - var a: cuda.Args = .{}; - a.add(@as(c_int, @intCast(heads))); - for ([_]u64{ state_in, q, k, v, g, b }) |x| a.add(x); - a.add(@as(c_int, @intCast(rows))); - a.add(y); - a.add(state_out); - try cuda.launch.launch(ck.f, .{ .grid = .{ .x = @intCast(heads) }, .block = .{ .x = 2 * dk } }, stream, &a); -} - -/// Opt-in: the wide chain with 16-row WY chunks; float64's bound, not the row path's bits; cuts at 16 k == uncut. -pub fn chunked(ks: *const Kernels, ck: Chunked, stream: cuda.Stream, c: Chain, rows: usize, s: Saves, w: Wide) !void { - try prep(ks, stream, c, rows, s, w); - try chunkLaunch(ck, stream, c.heads, c.state_in, w.q, s.k, s.v, s.g, s.b, rows, w.y, s.state_out); - try out(ks, stream, c, rows, s, w); -} - -/// The chunked recurrence over `rows` saved rows, state only: `chunked`'s state bits for a call of exactly those rows. -pub fn chunkReplay(ck: Chunked, stream: cuda.Stream, state_in: u64, k: u64, v: u64, g: u64, b: u64, rows: usize, heads: usize, state_out: u64) !void { - try chunkLaunch(ck, stream, heads, state_in, 0, k, v, g, b, rows, 0, state_out); -} - /// kda.replay: the state after the first `rows` saved rows, from state_in. pub fn replay(ks: *const Kernels, stream: cuda.Stream, state_in: u64, k: u64, v: u64, g: u64, b: u64, rows: usize, heads: usize, state_out: u64) !void { var a: cuda.Args = .{}; diff --git a/zig/tests/cuda/glm_kda_chunk_tests.zig b/zig/tests/cuda/glm_kda_chunk_tests.zig deleted file mode 100644 index 8cfec7d58c..0000000000 --- a/zig/tests/cuda/glm_kda_chunk_tests.zig +++ /dev/null @@ -1,124 +0,0 @@ -//! glm_kda_chunk.cu on synthetic GLM rows: float64 bounds, the wide chain's saves, its replay, resumed vs fresh. -const std = @import("std"); -const cuda = @import("cuda"); -const check = @import("check.zig"); -const f64r = @import("glm_kda_f64.zig"); -const rig = @import("glm_kda_rig.zig"); -const ref = @import("glm_ref_tests.zig"); -const Rig = rig.Rig; -const kda = rig.kda; -const dk = rig.dk; - -const heads = 8; -const n = heads * dk; -const p_stride = 3 * n + 64; -const layout: f64r.Layout = .{ .heads = heads, .p_stride = p_stride, .b_off = 3 * n, .lower = -5, .eps = 1e-6 }; - -fn uniform(r: *Rig, count: usize, lo: f32, hi: f32) ![]f32 { - const out = try r.a.alloc(f32, count); - for (out) |*v| v.* = lo + r.rng.float(f32) * (hi - lo); - return out; -} - -/// glm_ref_kda's synthetic inputs (its scales) at `heads` heads, and the chain over them on the device. -const Case = struct { - in: f64r.Inputs, - c: kda.Chain, - - fn make(r: *Rig, rows: usize) !Case { - const in: f64r.Inputs = .{ .p = try r.bfs(rows * p_stride, 1.0), .a = try r.bfs(rows * n, 2.0), .g = try r.bfs(rows * n, 2.0), .cs = try r.bfs(3 * 3 * n, 1.0), .cw = try uniform(r, 3 * n * 4, -0.5, 0.5), .state = try uniform(r, n * dk, -0.05, 0.05), .a_log = try uniform(r, heads, -1, 1), .dt_bias = try uniform(r, n, -0.5, 0.5), .norm = try r.bfs(dk, 1.0) }; - const B = std.mem.sliceAsBytes; - const c: kda.Chain = .{ .p = try r.dev(B(in.p)), .p_stride = p_stride, .b_off = 3 * n, .a = try r.dev(B(in.a)), .a_stride = n, .g = try r.dev(B(in.g)), .g_stride = n, .conv_state = try r.dev(B(in.cs)), .conv_w = try r.dev(B(in.cw)), .state_in = try r.dev(B(in.state)), .a_log = try r.dev(B(in.a_log)), .dt_bias = try r.dev(B(in.dt_bias)), .norm_w = try r.dev(B(in.norm)), .eps = 1e-6, .lower = -5, .heads = heads }; - return .{ .in = in, .c = c }; - } - - /// The chain from row `lo` on: its rows, and the conv window of rows lo - 3 .. lo - 1 (lo >= 3) on the device. - fn from(k: Case, r: *Rig, lo: usize, state_in: u64) !kda.Chain { - const cs = try r.a.alloc(u16, 3 * 3 * n); - for (0..3) |j| @memcpy(cs[j * 3 * n ..][0 .. 3 * n], k.in.p[(lo - 3 + j) * p_stride ..][0 .. 3 * n]); - var c = k.c; - c.p += lo * p_stride * 2; - c.a += lo * n * 2; - c.g += lo * n * 2; - c.conv_state = try r.dev(std.mem.sliceAsBytes(cs)); - c.state_in = state_in; - return c; - } -}; - -/// Worst fractions of the end-to-end bound (out, state) and of the recurrence's from the path's saves (y, state). -fn bounds(r: *Rig, k: Case, o: rig.Out, rows: usize, e2e: f64r.Want) ![4]f64 { - const rec = try f64r.recurrence(r.a, heads, k.in.state, try rig.saved(r, o, heads), rows); - const out = try rig.wide64(r, o.out, rows * n, true); - const st = try rig.wide64(r, o.state, n * dk, false); - const y = try rig.wide64(r, o.y, rows * n, true); - return .{ f64r.fraction(out, e2e.out, e2e.out_tol), f64r.fraction(st, e2e.state, e2e.state_tol), f64r.fraction(y, rec.out, rec.out_tol), f64r.fraction(st, rec.state, rec.state_tol) }; -} - -fn rowsCheck(r: *Rig, m: *const rig.Mods, rows: usize) !f64 { - const k = try Case.make(r, rows); - const wide = try rig.Out.alloc(r, heads, rows); - const chunk = try rig.Out.alloc(r, heads, rows); - try rig.run(m, r.s, k.c, rows, wide, .wide); - try rig.run(m, r.s, k.c, rows, chunk, .chunked); - var name: [128]u8 = undefined; - const P = std.fmt.bufPrint; - for ([_]u64{ wide.q, wide.k, wide.v, wide.g, wide.b }, [_]u64{ chunk.q, chunk.k, chunk.v, chunk.g, chunk.b }, [_]usize{ 4 * n, 4 * n, 2 * n, 4 * n, 4 * heads }, [_][]const u8{ "q", "k", "v", "decay", "beta" }) |a, b, size, what| - try ref.sameRows(r, try P(&name, "KDA chunked {d} rows: saved {s} equal the wide chain's", .{ rows, what }), a, b, size * rows); - const replayed = try r.zeros(n * dk * 4); - try kda.chunkReplay(m.ck, r.s, k.c.state_in, chunk.k, chunk.v, chunk.g, chunk.b, rows, heads, replayed); - try ref.sameRows(r, try P(&name, "KDA chunked {d} rows: the chunked replay rebuilds its state", .{rows}), chunk.state, replayed, n * dk * 4); - const e2e = try f64r.chain(r.a, layout, k.in, rows); - const bc = try bounds(r, k, chunk, rows, e2e); - const bw = try bounds(r, k, wide, rows, e2e); - const dout = rig.diff(try rig.wide64(r, chunk.out, rows * n, true), try rig.wide64(r, wide.out, rows * n, true)); - const dst = rig.diff(try rig.wide64(r, chunk.state, n * dk, false), try rig.wide64(r, wide.state, n * dk, false)); - std.debug.print("RESULT KDA {d} rows x {d} heads, worst fraction of float64's bound, chunked (wide): end to end out {d:.3} ({d:.3}), state {d:.4} ({d:.4}); " ++ - "recurrence from the saves y {d:.3} ({d:.3}), state {d:.3} ({d:.3})\n", .{ rows, heads, bc[0], bw[0], bc[1], bw[1], bc[2], bw[2], bc[3], bw[3] }); - std.debug.print("RESULT KDA {d} rows: chunked vs wide, out {d} of {d} bf16 differ (max {e:.2} of max |out|), state {d} of {d} fp32 differ (max {e:.2} of max |S|)\n", .{ rows, dout.count, rows * n, dout.rel, dst.count, n * dk, dst.rel }); - var worst: f64 = 0; - for (bc ++ bw) |x| worst = @max(worst, x); - try check.expect(worst <= 1, "KDA {d} rows: chunked and wide within float64's bounds (worst {d:.3})", .{ rows, worst }); - return worst; -} - -/// 300 rows resumed at `at` (second call from the first's state and conv window) against one call: equal bytes? -fn cut(r: *Rig, m: *const rig.Mods, at: usize) !bool { - const rows = 300; - const k = try Case.make(r, rows); - const whole = try rig.Out.alloc(r, heads, rows); - try rig.run(m, r.s, k.c, rows, whole, .chunked); - const parts = try rig.Out.alloc(r, heads, rows); - try rig.run(m, r.s, k.c, at, parts, .chunked); - var second = parts.from(heads, at); - second.state = try r.zeros(n * dk * 4); - try rig.run(m, r.s, try k.from(r, at, parts.state), rows - at, second, .chunked); - const out_same = std.mem.eql(u8, try r.back(whole.out, rows * n * 2), try r.back(parts.out, rows * n * 2)); - const st_same = std.mem.eql(u8, try r.back(whole.state, n * dk * 4), try r.back(second.state, n * dk * 4)); - const d = rig.diff(try rig.wide64(r, second.state, n * dk, false), try rig.wide64(r, whole.state, n * dk, false)); - const dy = rig.diff(try rig.wide64(r, parts.out, rows * n, true), try rig.wide64(r, whole.out, rows * n, true)); - std.debug.print("RESULT KDA chunked 300 rows resumed at {d} vs fresh: outputs equal {} ({d} bf16 differ, max {e:.2} of max |out|), state equal {} ({d} fp32 differ, max {e:.2} of max |S|)\n", .{ at, out_same, dy.count, dy.rel, st_same, d.count, d.rel }); - return out_same and st_same; -} - -/// `glm-kda-chunk`. -pub fn run(gpu: check.Gpu) !void { - var arena: std.heap.ArenaAllocator = .init(gpu.gpa); - defer arena.deinit(); - var s = try cuda.Stream.init(gpu.d, false); - defer s.deinit(); - var prng = std.Random.DefaultPrng.init(804); - var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .rng = prng.random() }; - defer { - for (r.bufs.items) |*b| b.free(); - r.bufs.deinit(gpu.gpa); - } - var m = try rig.Mods.load(gpu.d); - defer m.unload(); - var worst: f64 = 0; - for ([_]usize{ 1, 15, 16, 17, 70, 300 }) |rows| worst = @max(worst, try rowsCheck(&r, &m, rows)); - for ([_]usize{ 16, 48, 160 }) |at| try check.expect(try cut(&r, &m, at), "KDA chunked: a cut at {d} rows keeps one call's bits", .{at}); - var off = false; - for ([_]usize{ 37, 100, 255 }) |at| off = off or try cut(&r, &m, at); - check.pass("KDA chunked {d} heads: within float64's bounds (worst {d:.3}), saves equal the wide chain's, chunked replay rebuilds the state, cuts at 16/48/160 == uncut (resumed at 37/100/255: {s})", .{ heads, worst, if (off) "equal somewhere" else "differ, as expected" }); -} diff --git a/zig/tests/cuda/glm_kda_f64.zig b/zig/tests/cuda/glm_kda_f64.zig index b0db69748a..319839a5b6 100644 --- a/zig/tests/cuda/glm_kda_f64.zig +++ b/zig/tests/cuda/glm_kda_f64.zig @@ -1,4 +1,4 @@ -//! glm_kda.cu's chain in float64 at any head count, end to end and the recurrence alone (glm-kda-chunk, glm-kda-ling). +//! glm_kda.cu's chain in float64 at any head count, end to end and the recurrence alone (glm-kda-ling). const std = @import("std"); pub const dk = 128; diff --git a/zig/tests/cuda/glm_kda_ling.zig b/zig/tests/cuda/glm_kda_ling.zig index d357a1f586..0316fd4361 100644 --- a/zig/tests/cuda/glm_kda_ling.zig +++ b/zig/tests/cuda/glm_kda_ling.zig @@ -179,23 +179,17 @@ pub fn run(gpu: check.Gpu, path: []const u8) !void { for (yw, yd) |a, b| if (a != b) return check.expect(false, "Ling: 16-row windows' outputs == decode rows' outputs", .{}); const win: Result = .{ .y = yw, .sp = states[l.P], .st = ws }; var wsv: rig.Out = undefined; - var csv: rig.Out = undefined; const wide = try promptThenDecode(&r, &m, l, .wide, &wsv); - const chunk = try promptThenDecode(&r, &m, l, .chunked, &csv); for (wide.y, yd) |a, b| if (a != b) return check.expect(false, "Ling: the wide prompt's and decode outputs == every row a decode row's", .{}); try ref.sameRows(&r, "Ling: the wide prompt's state == P decode rows' state", wide.sp, states[l.P], n * dk * 4); - const crep = try r.zeros(n * dk * 4); - try kda.chunkReplay(m.ck, r.s, l.c.state_in, csv.k, csv.v, csv.g, csv.b, l.P, l.H, crep); - try ref.sameRows(&r, "Ling: the chunked replay == the chunked prompt's state", csv.state, crep, n * dk * 4); const e2e_p = try f64r.chain(r.a, l.l, l.in, l.P); const e2e = try f64r.chain(r.a, l.l, l.in, T); - const rec = try f64r.recurrence(r.a, l.H, l.in.state, try rig.saved(&r, csv, l.H), l.P); - const rc = [2]f64{ f64r.fraction(try rig.wide64(&r, csv.y, l.P * n, true), rec.out, rec.out_tol), f64r.fraction(try rig.wide64(&r, csv.state, n * dk, false), rec.state, rec.state_tol) }; + const rec = try f64r.recurrence(r.a, l.H, l.in.state, try rig.saved(&r, wsv, l.H), l.P); const recw = [2]f64{ f64r.fraction(try rig.wide64(&r, wsv.y, l.P * n, true), rec.out, rec.out_tol), f64r.fraction(try rig.wide64(&r, wsv.state, n * dk, false), rec.state, rec.state_tol) }; - var worst: f64 = @max(@max(rc[0], rc[1]), @max(recw[0], recw[1])); + var worst: f64 = @max(recw[0], recw[1]); var most_l2: f64 = 0; std.debug.print("RESULT glm-kda-ling: Ling-3.1 layer 0, {d} heads, {d} prompt + {d} decode rows; fla rel L2 of outputs (states after the prompt / end), worst fraction of float64's bound (out, state after prompt, at end)\n", .{ l.H, l.P, T - l.P }); - for ([_]Result{ dec, win, wide, chunk }, [_][]const u8{ "chain, every row a decode row", "chain, 16-row windows", "wide prompt + decode", "chunked prompt + decode" }) |res, label| { + for ([_]Result{ dec, win, wide }, [_][]const u8{ "chain, every row a decode row", "chain, 16-row windows", "wide prompt + decode" }) |res, label| { const f = try fracs(&r, l, res, e2e_p, e2e); const sp = try rig.wide64(&r, res.sp, n * dk, false); const st = try rig.wide64(&r, res.st, n * dk, false); @@ -205,12 +199,8 @@ pub fn run(gpu: check.Gpu, path: []const u8) !void { most_l2 = @max(most_l2, @max(yl2, f[3])); } const fla_f = f64r.fraction(l.fla_y, e2e.out, e2e.out_tol); - const dy = rig.diff(chunk.y[0 .. l.P * n], wide.y[0 .. l.P * n]); - const dsp = rig.diff(try rig.wide64(&r, chunk.sp, n * dk, false), try rig.wide64(&r, wide.sp, n * dk, false)); - const dall = rig.diff(chunk.y, wide.y); - std.debug.print("RESULT fla's own outputs: float64 {d:.3}; recurrence from the chunked saves, chunked y {d:.3} state {d:.3} (wide {d:.3}, {d:.3})\n", .{ fla_f, rc[0], rc[1], recw[0], recw[1] }); - std.debug.print("RESULT chunked vs wide prompt: out {d} of {d} bf16 differ (max {e:.2} of max |out|), state {d} of {d} fp32 differ (max {e:.2} of max |S|); all {d} rows incl. decode: {d} differ\n", .{ dy.count, l.P * n, dy.rel, dsp.count, n * dk, dsp.rel, T, dall.count }); + std.debug.print("RESULT fla's own outputs: float64 {d:.3}; recurrence from the wide prompt's saves, y {d:.3} state {d:.3}\n", .{ fla_f, recw[0], recw[1] }); try check.expect(worst <= 1, "glm-kda-ling: every path within float64's bound (worst {d:.3})", .{worst}); try check.expect(most_l2 < 1e-2, "glm-kda-ling: every path within 1 % of fla (worst rel L2 {e:.3})", .{most_l2}); - check.pass("glm-kda-ling: glm_kda serves Ling layer 0 as is; chain, windows, wide and chunked within float64's bound (worst {d:.3}) and fla (rel L2 <= {e:.3}); replays == chain bits", .{ worst, most_l2 }); + check.pass("glm-kda-ling: glm_kda serves Ling layer 0 as is; chain, windows and wide within float64's bound (worst {d:.3}) and fla (rel L2 <= {e:.3}); replays == chain bits", .{ worst, most_l2 }); } diff --git a/zig/tests/cuda/glm_kda_rig.zig b/zig/tests/cuda/glm_kda_rig.zig index e7d22615d9..a39b5385db 100644 --- a/zig/tests/cuda/glm_kda_rig.zig +++ b/zig/tests/cuda/glm_kda_rig.zig @@ -1,4 +1,4 @@ -//! glm-kda-chunk's and glm-kda-ling's shared rig: a call's inputs on the device, one path's outputs, downloads to f64. +//! glm-kda-ling's rig: a call's inputs on the device, one path's outputs, downloads to f64. const std = @import("std"); const cuda = @import("cuda"); const qr = @import("quant_ref_tests.zig"); @@ -9,30 +9,25 @@ pub const Rig = qr.RigT; pub const kda = glm.kda; pub const dk = kda.dk; -/// The two modules: glm_kda.cu (chain, replay, wide) and glm_kda_chunk.cu. +/// glm_kda.cu's module (chain, replay, wide). pub const Mods = struct { kda_mod: cuda.Module, - chunk_mod: cuda.Module, ks: kda.Kernels, - ck: kda.Chunked, pub fn load(d: *const cuda.Driver) !Mods { var m: Mods = undefined; m.kda_mod = try cuda.Module.load(d, cuda.kernels.glm_kda); - m.chunk_mod = try cuda.Module.load(d, cuda.kernels.glm_kda_chunk); m.ks = try kda.Kernels.load(m.kda_mod); - m.ck = try kda.Chunked.load(m.chunk_mod); return m; } pub fn unload(m: *Mods) void { - m.chunk_mod.unload(); m.kda_mod.unload(); } }; -/// How a call runs: the fused chain, the three-kernel (serial wide) chain, or the chunked prompt path. -pub const Path = enum { fused, wide, chunked }; +/// How a call runs: the fused chain or the three-kernel (serial wide) chain. +pub const Path = enum { fused, wide }; /// One path's buffers for `rows` rows: out, state, saves (k, v, decay, beta), wide scratch (q, y). pub const Out = struct { @@ -68,7 +63,6 @@ pub fn run(m: *const Mods, s: cuda.Stream, c: kda.Chain, rows: usize, o: Out, pa switch (path) { .fused => try kda.chain(&m.ks, s, c, rows, o.saves(), null), .wide => try kda.chain(&m.ks, s, c, rows, o.saves(), w), - .chunked => try kda.chunked(&m.ks, m.ck, s, c, rows, o.saves(), w), } } diff --git a/zig/tests/cuda/glm_kda_speed.zig b/zig/tests/cuda/glm_kda_speed.zig deleted file mode 100644 index 2cfc97e885..0000000000 --- a/zig/tests/cuda/glm_kda_speed.zig +++ /dev/null @@ -1,76 +0,0 @@ -//! glm-kda-speed: 64 heads x 128, a 2048-row prompt (row path vs chunked), a decode row, a 16-row window; median ms. -const std = @import("std"); -const cuda = @import("cuda"); -const check = @import("check.zig"); -const rig = @import("glm_kda_rig.zig"); -const Rig = rig.Rig; -const kda = rig.kda; -const dk = rig.dk; - -const heads = 64; -const n = heads * dk; -const prompt = 2048; - -const What = enum { path, replay, prep_out }; - -fn time(r: *Rig, m: *const rig.Mods, c: kda.Chain, rows: usize, o: rig.Out, path: rig.Path, what: What) !f64 { - var start = try cuda.Event.init(r.gpu.d, true); - defer start.deinit(); - var end = try cuda.Event.init(r.gpu.d, true); - defer end.deinit(); - var ms: [9]f64 = undefined; - for (0..ms.len + 1) |i| { - try start.record(r.s); - switch (what) { - .path => try rig.run(m, r.s, c, rows, o, path), - .replay => if (path == .chunked) try kda.chunkReplay(m.ck, r.s, c.state_in, o.k, o.v, o.g, o.b, rows, heads, o.state) else try kda.replay(&m.ks, r.s, c.state_in, o.k, o.v, o.g, o.b, rows, heads, o.state), - .prep_out => { - try kda.prep(&m.ks, r.s, c, rows, o.saves(), .{ .q = o.q, .y = o.y }); - try kda.out(&m.ks, r.s, c, rows, o.saves(), .{ .q = o.q, .y = o.y }); - }, - } - try end.record(r.s); - try r.s.synchronize(); - if (i > 0) ms[i - 1] = try start.elapsedMs(end); - } - std.mem.sort(f64, &ms, {}, std.sort.asc(f64)); - return ms[ms.len / 2]; -} - -pub fn run(gpu: check.Gpu) !void { - var arena: std.heap.ArenaAllocator = .init(gpu.gpa); - defer arena.deinit(); - var s = try cuda.Stream.init(gpu.d, false); - defer s.deinit(); - var prng = std.Random.DefaultPrng.init(64); - var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .rng = prng.random() }; - defer { - for (r.bufs.items) |*b| b.free(); - r.bufs.deinit(gpu.gpa); - } - var m = try rig.Mods.load(gpu.d); - defer m.unload(); - const p_stride = 3 * n + heads; - const state = try r.a.alloc(f32, n * dk); - for (state) |*v| v.* = (r.rng.float(f32) * 2 - 1) * 0.05; - const cw = try r.a.alloc(f32, 3 * n * 4); - for (cw) |*v| v.* = r.rng.float(f32) - 0.5; - const small = try r.a.alloc(f32, n); - for (small) |*v| v.* = r.rng.float(f32) - 0.5; - const B = std.mem.sliceAsBytes; - const c: kda.Chain = .{ .p = try r.dev(B(try r.bfs(prompt * p_stride, 1.0))), .p_stride = p_stride, .b_off = 3 * n, .a = try r.dev(B(try r.bfs(prompt * n, 2.0))), .a_stride = n, .g = try r.dev(B(try r.bfs(prompt * n, 2.0))), .g_stride = n, .conv_state = try r.dev(B(try r.bfs(9 * n, 1.0))), .conv_w = try r.dev(B(cw)), .state_in = try r.dev(B(state)), .a_log = try r.dev(B(small[0..heads])), .dt_bias = try r.dev(B(small)), .norm_w = try r.dev(B(try r.bfs(dk, 1.0))), .eps = 1e-6, .lower = -5, .heads = heads }; - const o = try rig.Out.alloc(&r, heads, prompt); - const wide = try time(&r, &m, c, prompt, o, .wide, .path); - const chunk = try time(&r, &m, c, prompt, o, .chunked, .path); - const rw = try time(&r, &m, c, prompt, o, .wide, .replay); - const rc = try time(&r, &m, c, prompt, o, .chunked, .replay); - const po = try time(&r, &m, c, prompt, o, .wide, .prep_out); - const row = try time(&r, &m, c, 1, o, .fused, .path); - const win = try time(&r, &m, c, 16, o, .fused, .path); - const win_wide = try time(&r, &m, c, 16, o, .wide, .path); - const win_chunk = try time(&r, &m, c, 16, o, .chunked, .path); - std.debug.print("RESULT glm-kda-speed {d} heads x 128, {d}-row prompt: wide (prep, serial step, out) {d:.3} ms, chunked (prep, WY chunks, out) {d:.3} ms ({d:.2}x); " ++ - "prep + out alone {d:.3} ms; state-only: serial replay {d:.3} ms, chunked replay {d:.3} ms\n", .{ heads, prompt, wide, chunk, wide / chunk, po, rw, rc }); - std.debug.print("RESULT glm-kda-speed decode row (fused chain) {d:.4} ms; 16-row window: fused chain {d:.4} ms, wide {d:.4} ms, chunked {d:.4} ms\n", .{ row, win, win_wide, win_chunk }); - check.pass("glm-kda-speed: timed", .{}); -} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 84584ef6f3..f14dbcc650 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -42,9 +42,7 @@ const glm_ref_route = @import("glm_ref_route.zig"); const ling_synth = @import("ling_synth.zig"); const ling_weights = @import("ling_weights_tests.zig"); const ling_real = @import("ling_weights_real.zig"); -const glm_kda_chunk_tests = @import("glm_kda_chunk_tests.zig"); const glm_kda_ling = @import("glm_kda_ling.zig"); -const glm_kda_speed = @import("glm_kda_speed.zig"); const usage = \\usage: tf-cuda-test @@ -101,9 +99,7 @@ const usage = \\ ling-mla Ling's MLA (128-wide absorb/expand, rope, gate, 576-wide attention) vs float64 (synthetic) \\ ling-mla-real Ling layer 7's MLA on real inputs vs HF and float64; decode, window and cut bytes \\ glm-mla-same GLM's MLA launches on this build vs a 6380c8da glm_glue fatbin, byte for byte - \\ glm-kda-chunk GLM KDA's chunked prompt path against float64, the wide chain's saves, cuts (synthetic) \\ glm-kda-ling FIXTURE GLM KDA's paths on Ling-3.1 layer 0 (kda-real's fixture) against fla and float64 - \\ glm-kda-speed GLM KDA at 64 heads: a 2048-row prompt wide vs chunked, a decode row, a 16-row window \\ glm-lanes-cover copy drafts, cancels, discards, eos and max_new on one lanes.Engine across requests \\ glm-lanes-concurrent concurrent streams on one lanes.Engine, each against its solo run \\ glm-shared-check shared verify windows against lone twins, logits and states byte for byte @@ -188,9 +184,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "ling-mla")) return @import("ling_ref_mla.zig").run(gpu); if (std.mem.eql(u8, cmd, "ling-mla-real")) return @import("ling_mla_real.zig").run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "glm-mla-same")) return @import("glm_mla_same.zig").run(gpu, try arg(rest, 0)); - if (std.mem.eql(u8, cmd, "glm-kda-chunk")) return glm_kda_chunk_tests.run(gpu); if (std.mem.eql(u8, cmd, "glm-kda-ling")) return glm_kda_ling.run(gpu, try arg(rest, 0)); - if (std.mem.eql(u8, cmd, "glm-kda-speed")) return glm_kda_speed.run(gpu); if (std.mem.eql(u8, cmd, "modelopt")) return modelopt_tests.run(gpu); if (std.mem.eql(u8, cmd, "exlr8")) return exlr8_tests.run(gpu); if (std.mem.eql(u8, cmd, "exlr8-dense")) return exlr8_dense_tests.run(gpu, if (rest.len > 0) rest[0] else null); From 66579cbadea2571c048d3ed26e28af8496a46f29 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 16 Oct 2026 07:06:27 +0200 Subject: [PATCH 24/34] ling: config refuses every geometry the kernels don't take, at load, with a named error parse ran KDA head widths, conv tap counts and the MLA widths straight into kernels built for one shape. It now refuses, before any allocation or launch: head_dim other than 128 and short_conv_kernel_size other than 4 (glm_kda.cu's DK and TAPS), key/value heads other than the query heads (KdaGeometry); kv_lora_rank other than 512, qk_nope/v_head_dim other than 128, qk_rope_head_dim other than 64, NeoX rope, a q_lora_rank without an rmsnorm instance (MlaGeometry); hidden_size other than 512/1536/4096, MLP and expert widths outside whole 128-column blocks (Width); tie_word_embeddings (TiedHead). - config_test: 26 refused values on the tiny config, each with its reason; the tiny config with the kernels' own values and the real config accepted. cuda.zig's test holds the fixed widths to glm.kda, glm.latent and fp8.block. --- docs/recipes/ling-3.1.md | 9 +++ zig/src/families/ling/config.zig | 65 +++++++++++++++++++++ zig/src/families/ling/config_test.zig | 81 +++++++++++++++++++++++++-- zig/src/families/ling/cuda.zig | 7 +++ 4 files changed, 157 insertions(+), 5 deletions(-) diff --git a/docs/recipes/ling-3.1.md b/docs/recipes/ling-3.1.md index 1fe89c41f6..e548852104 100644 --- a/docs/recipes/ling-3.1.md +++ b/docs/recipes/ling-3.1.md @@ -101,6 +101,15 @@ both formats. Each projection loads as a `qlinear.Weight` (`bf16`, `fp8g` or `nv `Fp8Experts`/`Nvfp4Experts`'s layouts on the GPU: `Packer.fp8Experts` and `Packer.fp4Experts`). There is no grouped bf16 expert kernel, so bf16 routed experts load as one lane weight per projection. The router must be bf16 (GLM's router kernel reads bf16); a quantized router is refused. +### Shapes + +`config.parse` refuses, with a named error and the reason, every shape the kernels don't take, before anything is +allocated or launched: KDA heads other than 128 wide, short convs other than 4 taps, key/value heads other than the +query heads (`KdaGeometry`); a latent other than 512, nope or value heads other than 128, rope columns other than 64, +NeoX rope pairs, a `q_lora_rank` without a norm instance (`MlaGeometry`); a hidden width other than 512, 1536 or 4096, +MLP and expert widths outside whole 128-column blocks (`Width`); a tied head (`TiedHead`); routing past `glue.select`'s +bounds (`Routing`). + ### Loading Routed experts load in batches of `weights.Options.stage` bytes (512 MiB, 32 FP8 experts): eight threads copy the batch's diff --git a/zig/src/families/ling/config.zig b/zig/src/families/ling/config.zig index ba835602c5..18b0a86b68 100644 --- a/zig/src/families/ling/config.zig +++ b/zig/src/families/ling/config.zig @@ -447,6 +447,71 @@ fn parseConfig(gpa: std.mem.Allocator, json: []const u8, why: *Why) !Config { }; c.tie = try flag(o, "tie_word_embeddings", false, why); try tokens(o, &c, why); + try geometry(o, &c, why); try quantization(o, &c, why); return c; } + +/// The widths the kernels are built for (cuda.zig's tests hold them to glm_kda, glm_glue and the latent cache). +pub const fixed = .{ .kda_dim = 128, .conv = 4, .kv_lora = 512, .nope = 128, .v_dim = 128, .rope = 64, .block = 128 }; +/// glue.rmsnorm's instances: the hidden width and q_lora_rank take one of these. +pub const norm_widths = [_]u32{ 512, 1536, 4096 }; + +fn normWidth(w: u32) bool { + return std.mem.indexOfScalar(u32, &norm_widths, w) != null; +} + +/// Every shape the kernels take fixed, refused at load before anything is allocated or launched. +fn geometry(o: std.json.ObjectMap, c: *const Config, why: *Why) !void { + // KDA: glm_kda.cu's 128-wide heads and 4 taps (a 3-row conv window); one key and value head per query head + if (c.kda_dim != fixed.kda_dim) { + why.set("head_dim {d}: the KDA kernels take {d}-wide heads", .{ c.kda_dim, fixed.kda_dim }); + return error.KdaGeometry; + } + if (c.conv != fixed.conv) { + why.set("short_conv_kernel_size {d}: the KDA kernels take {d} taps", .{ c.conv, fixed.conv }); + return error.KdaGeometry; + } + const kv_heads = if (o.get("num_key_value_heads")) |_| try int(o, "num_key_value_heads", why) else c.kda_heads; + if (c.kda_heads == 0 or kv_heads != c.kda_heads) { + why.set("{d} attention heads, {d} key/value heads: the engine runs one key/value head per query head", .{ c.kda_heads, kv_heads }); + return error.KdaGeometry; + } + // MLA: the 512-wide latent, 128-wide nope and value heads, 64 GPT-J roped columns + const mla = [_]struct { []const u8, u32, u32 }{ + .{ "kv_lora_rank", c.kv_lora, fixed.kv_lora }, + .{ "qk_nope_head_dim", c.nope, fixed.nope }, + .{ "v_head_dim", c.v_dim, fixed.v_dim }, + .{ "qk_rope_head_dim", c.rope_dim, fixed.rope }, + }; + for (mla) |m| if (m[1] != m[2]) { + why.set("{s} {d}: the MLA kernels take {d}", .{ m[0], m[1], m[2] }); + return error.MlaGeometry; + }; + if (c.rope_style != .gptj) { + why.set("rope_interleave false: the MLA rope kernel pairs adjacent columns (GPT-J) only", .{}); + return error.MlaGeometry; + } + if (c.q_lora != 0 and !normWidth(c.q_lora)) { + why.set("q_lora_rank {d}: q_a_layernorm runs at 512, 1536 or 4096", .{c.q_lora}); + return error.MlaGeometry; + } + // widths: the norms' instances, and every linear's inputs in whole 128-column blocks (FP8's, the lanes' 64) + if (!normWidth(c.hidden)) { + why.set("hidden_size {d}: the norms and the router run at 512, 1536 or 4096", .{c.hidden}); + return error.Width; + } + const widths = [_]struct { []const u8, u32, bool }{ + .{ "intermediate_size", c.dense_inter, c.dense_layers > 0 }, + .{ "moe_intermediate_size", c.moe_inter, true }, + .{ "the shared expert's width", c.shared_inter, c.shared_inter > 0 }, + }; + for (widths) |x| if (x[2] and (x[1] == 0 or x[1] % fixed.block != 0)) { + why.set("{s} {d}: the projections and experts take whole blocks of {d}", .{ x[0], x[1], fixed.block }); + return error.Width; + }; + if (c.tie) { + why.set("tie_word_embeddings: the engine reads its own lm_head", .{}); + return error.TiedHead; + } +} diff --git a/zig/src/families/ling/config_test.zig b/zig/src/families/ling/config_test.zig index 16d18777fb..a0d1f38516 100644 --- a/zig/src/families/ling/config_test.zig +++ b/zig/src/families/ling/config_test.zig @@ -138,10 +138,8 @@ test "refusals: one edit each, with the reason" { try refused("\"hidden_size\": 4096", "\"hidden_sz\": 4096", error.BadConfig, "hidden_size"); } -test "optional keys: no quantization reads bf16, a null q_lora_rank one q_proj, rope_interleave false NeoX" { - const a = try edit("\"q_lora_rank\": 1536", "\"q_lora_rank\": null"); - defer t.allocator.free(a); - const b = try std.mem.replaceOwned(u8, t.allocator, a, "\"rope_interleave\": true", "\"rope_interleave\": false"); +test "optional keys: no quantization reads bf16, a null q_lora_rank one q_proj" { + const b = try edit("\"q_lora_rank\": 1536", "\"q_lora_rank\": null"); defer t.allocator.free(b); const at = std.mem.indexOf(u8, b, "\"quantization_config\"").?; const end = std.mem.indexOf(u8, b[at..], " },\n").? + at + 5; @@ -151,7 +149,7 @@ test "optional keys: no quantization reads bf16, a null q_lora_rank one q_proj, const c = try config.parse(t.allocator, c_json, &why); try t.expectEqual(config.Quant.bf16, c.quant); try t.expectEqual(@as(u32, 0), c.q_lora); - try t.expectEqual(config.RopeStyle.neox, c.rope_style); + try t.expectEqual(config.RopeStyle.gptj, c.rope_style); try t.expect(!c.scaled("model.layers.0.attention.q_proj")); } @@ -201,3 +199,76 @@ test "EXL3 is refused naming #832; a glob matches whole names only" { try t.expect(!config.glob("model.layers.1*", "model.layers.0.mlp")); try t.expect(config.glob("a*b*c", "a.b.x.c") and !config.glob("a*b*c", "a.c")); } + +const synth = @import("synth.zig"); + +/// The tiny config (synth.zig, FP8) with `key` set to `value` (JSON text) and `also`'s pair, parsed. +fn tiny(key: []const u8, value: []const u8, also: ?[2][]const u8, why: *config.Why) !config.Config { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var root = try std.json.parseFromSliceLeaky(std.json.Value, a, try synth.configJson(a, .fp8_block128, 4), .{}); + try root.object.put(a, key, try std.json.parseFromSliceLeaky(std.json.Value, a, value, .{})); + if (also) |kv| try root.object.put(a, kv[0], try std.json.parseFromSliceLeaky(std.json.Value, a, kv[1], .{})); + return config.parse(t.allocator, try std.json.Stringify.valueAlloc(a, root, .{}), why); +} + +const Case = struct { key: []const u8, value: []const u8, err: anyerror, says: []const u8, also: ?[2][]const u8 = null }; + +test "geometry the kernels do not take is refused at load, one value at a time on the tiny config, with the reason" { + const cases = [_]Case{ + .{ .key = "head_dim", .value = "64", .err = error.KdaGeometry, .says = "head_dim 64" }, + .{ .key = "head_dim", .value = "256", .err = error.KdaGeometry, .says = "128-wide heads" }, + .{ .key = "short_conv_kernel_size", .value = "3", .err = error.KdaGeometry, .says = "short_conv_kernel_size 3" }, + .{ .key = "short_conv_kernel_size", .value = "5", .err = error.KdaGeometry, .says = "4 taps" }, + .{ .key = "num_attention_heads", .value = "0", .err = error.KdaGeometry, .says = "0 attention heads" }, + .{ .key = "num_attention_heads", .value = "0", .err = error.KdaGeometry, .says = "0 attention heads, 0", .also = .{ "num_key_value_heads", "0" } }, + .{ .key = "num_key_value_heads", .value = "2", .err = error.KdaGeometry, .says = "2 key/value heads" }, + .{ .key = "kv_lora_rank", .value = "256", .err = error.MlaGeometry, .says = "kv_lora_rank 256" }, + .{ .key = "kv_lora_rank", .value = "1024", .err = error.MlaGeometry, .says = "take 512" }, + .{ .key = "qk_nope_head_dim", .value = "64", .err = error.MlaGeometry, .says = "qk_nope_head_dim 64" }, + .{ .key = "qk_nope_head_dim", .value = "192", .err = error.MlaGeometry, .says = "qk_nope_head_dim 192" }, + .{ .key = "v_head_dim", .value = "256", .err = error.MlaGeometry, .says = "v_head_dim 256" }, + .{ .key = "v_head_dim", .value = "64", .err = error.MlaGeometry, .says = "v_head_dim 64" }, + .{ .key = "qk_rope_head_dim", .value = "32", .err = error.MlaGeometry, .says = "qk_rope_head_dim 32" }, + .{ .key = "qk_rope_head_dim", .value = "128", .err = error.MlaGeometry, .says = "qk_rope_head_dim 128" }, + .{ .key = "rope_interleave", .value = "false", .err = error.MlaGeometry, .says = "GPT-J" }, + .{ .key = "q_lora_rank", .value = "768", .err = error.MlaGeometry, .says = "q_lora_rank 768" }, + .{ .key = "q_lora_rank", .value = "1024", .err = error.MlaGeometry, .says = "q_lora_rank 1024" }, + .{ .key = "hidden_size", .value = "1024", .err = error.Width, .says = "hidden_size 1024" }, + .{ .key = "hidden_size", .value = "640", .err = error.Width, .says = "hidden_size 640" }, + .{ .key = "intermediate_size", .value = "200", .err = error.Width, .says = "intermediate_size 200" }, + .{ .key = "intermediate_size", .value = "0", .err = error.Width, .says = "intermediate_size 0" }, + .{ .key = "moe_intermediate_size", .value = "96", .err = error.Width, .says = "moe_intermediate_size 96" }, + .{ .key = "moe_intermediate_size", .value = "192", .err = error.Width, .says = "moe_intermediate_size 192" }, + .{ .key = "moe_shared_expert_intermediate_size", .value = "64", .err = error.Width, .says = "shared expert's width 64" }, + .{ .key = "tie_word_embeddings", .value = "true", .err = error.TiedHead, .says = "lm_head" }, + }; + for (cases) |x| { + var why: config.Why = .{}; + const got = tiny(x.key, x.value, x.also, &why); + if (got) |_| { + std.debug.print("{s} {s}: accepted\n", .{ x.key, x.value }); + return error.TestUnexpectedResult; + } else |e| if (e != x.err or std.mem.indexOf(u8, why.text(), x.says) == null) { + std.debug.print("{s} {s}: {t}, \"{s}\"\n", .{ x.key, x.value, e, why.text() }); + return error.TestUnexpectedResult; + } + } +} + +test "the shapes the kernels take are accepted: the tiny config as written and with these values, the real one" { + var why: config.Why = .{}; + const ok = [_][2][]const u8{ + .{ "head_dim", "128" }, .{ "short_conv_kernel_size", "4" }, .{ "num_key_value_heads", "4" }, + .{ "q_lora_rank", "1536" }, .{ "q_lora_rank", "null" }, .{ "hidden_size", "4096" }, + .{ "hidden_size", "1536" }, .{ "moe_intermediate_size", "1280" }, .{ "intermediate_size", "10240" }, + .{ "rope_interleave", "true" }, + }; + for (ok) |kv| _ = tiny(kv[0], kv[1], null, &why) catch |e| { + std.debug.print("{s} {s}: {t}, \"{s}\"\n", .{ kv[0], kv[1], e, why.text() }); + return e; + }; + const c = try load(); + try t.expect(c.kda_dim == config.fixed.kda_dim and c.conv == config.fixed.conv and c.kv_lora == config.fixed.kv_lora); +} diff --git a/zig/src/families/ling/cuda.zig b/zig/src/families/ling/cuda.zig index d00fa55f6e..314eb30103 100644 --- a/zig/src/families/ling/cuda.zig +++ b/zig/src/families/ling/cuda.zig @@ -54,6 +54,13 @@ test { _ = native; } +test "config's fixed widths are the kernels' own: KDA heads, the latent cache, the rope" { + const f = config.fixed; + try std.testing.expect(f.kda_dim == glm.kda.dk and f.kda_dim == glm.kda.dv); + try std.testing.expect(f.rope == glm.latent.rope_width and f.kv_lora + f.rope == state.cache_width); + try std.testing.expect(f.block == @import("cuda").fp8.block); +} + test "Ling's routing as glue.select takes it: 8 groups keeping 4, and select's own bounds hold" { var why: config.Why = .{}; const c = try config.parse(std.testing.allocator, @embedFile("fixtures/config.json"), &why); From 92afed6143bc3f4ddc670eefceb23f0576a904c5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 16 Oct 2026 07:12:55 +0200 Subject: [PATCH 25/34] ling: routing groups of one and a single pick refused at load; one group and every group kept checked vs float64 HF's group_limited_topk scores a group by its top two (topk(2) raises on a group of one), and with one pick it skips the renormalization glue.select applies. Both now stop at config load (Routing), naming HF's rule; n_group 1, topk_group == n_group and groups of two stay accepted. - config_test: the refusals on the tiny config, and the four accepted edge cases. - glm-ref: glue.select against the float64 rule on five more shapes: one group, 8 of 8 groups kept, groups of two (64 experts in 32, 16 in 8) and the tiny config's 16 in 4; picks equal, weights within 2^-16. --- zig/src/families/ling/config.zig | 9 +++++++ zig/src/families/ling/config_test.zig | 34 +++++++++++++++++++++++++++ zig/tests/cuda/glm_ref_route.zig | 5 ++++ 3 files changed, 48 insertions(+) diff --git a/zig/src/families/ling/config.zig b/zig/src/families/ling/config.zig index 18b0a86b68..8939a8f4a7 100644 --- a/zig/src/families/ling/config.zig +++ b/zig/src/families/ling/config.zig @@ -435,6 +435,15 @@ fn parseConfig(gpa: std.mem.Allocator, json: []const u8, why: *Why) !Config { why.set("routing {d} of {d} experts in {d} groups keeping {d}: past glue.select's bounds", .{ c.topk, c.experts, c.n_group, c.topk_group }); return error.Routing; } + // HF's rule: a group scores its top two (topk(2) fails on a group of one); one pick keeps its weight unnormalized + if (c.experts / c.n_group < 2) { + why.set("{d} groups of {d} experts: HF scores a group by its best two, so groups hold two or more", .{ c.n_group, c.experts / c.n_group }); + return error.Routing; + } + if (c.topk == 1) { + why.set("num_experts_per_tok 1: HF leaves one pick's weight unnormalized, the kernel normalizes it", .{}); + return error.Routing; + } c.moe_inter = try int(o, "moe_intermediate_size", why); const shared_n = if (o.get("num_shared_experts")) |v| (if (v == .null) 0 else try int(o, "num_shared_experts", why)) else 0; const shared_w = if (o.get("moe_shared_expert_intermediate_size")) |_| try int(o, "moe_shared_expert_intermediate_size", why) else c.moe_inter; diff --git a/zig/src/families/ling/config_test.zig b/zig/src/families/ling/config_test.zig index a0d1f38516..f8a08e87c7 100644 --- a/zig/src/families/ling/config_test.zig +++ b/zig/src/families/ling/config_test.zig @@ -272,3 +272,37 @@ test "the shapes the kernels take are accepted: the tiny config as written and w const c = try load(); try t.expect(c.kda_dim == config.fixed.kda_dim and c.conv == config.fixed.conv and c.kv_lora == config.fixed.kv_lora); } + +test "routing groups: groups of one and a single pick refused (HF's rule differs there); one group, all groups kept, groups of two accepted" { + const refusals = [_]Case{ + .{ .key = "n_group", .value = "16", .err = error.Routing, .says = "16 groups of 1 experts", .also = .{ "topk_group", "4" } }, + .{ .key = "n_group", .value = "16", .err = error.Routing, .says = "best two", .also = .{ "topk_group", "16" } }, + .{ .key = "num_experts_per_tok", .value = "1", .err = error.Routing, .says = "num_experts_per_tok 1" }, + .{ .key = "n_group", .value = "1", .err = error.Routing, .says = "past glue.select's bounds" }, // keeping 2 of 1 + }; + for (refusals) |x| { + var why: config.Why = .{}; + const got = tiny(x.key, x.value, x.also, &why); + if (got) |_| { + std.debug.print("{s} {s}: accepted\n", .{ x.key, x.value }); + return error.TestUnexpectedResult; + } else |e| if (e != x.err or std.mem.indexOf(u8, why.text(), x.says) == null) { + std.debug.print("{s} {s}: {t}, \"{s}\"\n", .{ x.key, x.value, e, why.text() }); + return error.TestUnexpectedResult; + } + } + const ok = [_][4][]const u8{ + .{ "n_group", "1", "topk_group", "1" }, // one group: glue.select skips the group pass, HF keeps every expert + .{ "n_group", "4", "topk_group", "4" }, // every group kept + .{ "n_group", "8", "topk_group", "2" }, // groups of two + .{ "n_group", "8", "topk_group", "8" }, + }; + for (ok) |kv| { + var why: config.Why = .{}; + const c = tiny(kv[0], kv[1], .{ kv[2], kv[3] }, &why) catch |e| { + std.debug.print("{s} {s}, {s} {s}: {t}, \"{s}\"\n", .{ kv[0], kv[1], kv[2], kv[3], e, why.text() }); + return e; + }; + try t.expect(c.experts / c.n_group >= 2 and c.topk_group <= c.n_group); + } +} diff --git a/zig/tests/cuda/glm_ref_route.zig b/zig/tests/cuda/glm_ref_route.zig index c1a4b9439a..b125ac37b6 100644 --- a/zig/tests/cuda/glm_ref_route.zig +++ b/zig/tests/cuda/glm_ref_route.zig @@ -167,6 +167,11 @@ pub fn groups(r: *Rig, set: glue.Set) !void { .{ .rows = 300, .ne = 256, .top_k = 8, .ng = 8, .keep = 4, .scale = 2.5 }, // DeepSeek-V3's .{ .rows = 100, .ne = 480, .top_k = 6, .ng = 12, .keep = 3, .scale = 1.0 }, // groups of 40 straddle the lanes .{ .rows = 100, .ne = 512, .top_k = 8, .ng = 32, .keep = 1, .scale = 2.5 }, + .{ .rows = 100, .ne = 512, .top_k = 8, .ng = 1, .keep = 1, .scale = 2.5 }, // one group: no group pass + .{ .rows = 100, .ne = 512, .top_k = 8, .ng = 8, .keep = 8, .scale = 2.5 }, // every group kept + .{ .rows = 100, .ne = 64, .top_k = 8, .ng = 32, .keep = 4, .scale = 2.5 }, // groups of two + .{ .rows = 100, .ne = 16, .top_k = 4, .ng = 8, .keep = 3, .scale = 1.0 }, // groups of two, the tiny config's experts + .{ .rows = 100, .ne = 16, .top_k = 4, .ng = 4, .keep = 2, .scale = 2.5 }, // the tiny config's routing }; for (shapes) |q| for ([_]f32{ 0.001, 0.05 }) |bspread| { const logits = try f32s(r, q.rows * q.ne, -6, 6); From 674163facacd1ce357c13963ffb98b85fe1c7b45 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 16 Oct 2026 07:15:59 +0200 Subject: [PATCH 26/34] ling: Store.alloc owns a buffer only once it is in both lists; a fake driver to fail every call in turn A failed name copy or named.append used to free the device buffer already in buffers (Store.deinit freed it again) and lose the copied name. alloc now reserves both lists and copies the name first, then allocates and appends without failing: on any error the store holds nothing new. - fake_driver.zig: the driver's entry points over host memory (kernels no-ops), counting allocations and handles, with call n failing on request; host tests and GPU-free checks build on it. - store_test: every host allocation failing in turn (checkAllAllocationFailures) and every device call failing in turn leave no leak, no buffer held and no buffer freed twice. The old alloc fails both. --- zig/src/families/ling/cuda.zig | 3 + zig/src/families/ling/fake_driver.zig | 343 +++++++++++++++++++++++++ zig/src/families/ling/store_test.zig | 43 ++++ zig/src/families/ling/weights_proj.zig | 14 +- 4 files changed, 396 insertions(+), 7 deletions(-) create mode 100644 zig/src/families/ling/fake_driver.zig create mode 100644 zig/src/families/ling/store_test.zig diff --git a/zig/src/families/ling/cuda.zig b/zig/src/families/ling/cuda.zig index 314eb30103..33c42efd23 100644 --- a/zig/src/families/ling/cuda.zig +++ b/zig/src/families/ling/cuda.zig @@ -12,6 +12,7 @@ pub const mtp = @import("mtp.zig"); pub const engine = @import("engine.zig"); pub const lanes = @import("lanes.zig"); pub const native = @import("native.zig"); +pub const fake_driver = @import("fake_driver.zig"); const Config = config.Config; @@ -52,6 +53,8 @@ test { std.testing.refAllDecls(lanes.Lanes); std.testing.refAllDecls(mtp.Head); _ = native; + _ = @import("fake_driver.zig"); + _ = @import("store_test.zig"); } test "config's fixed widths are the kernels' own: KDA heads, the latent cache, the rope" { diff --git a/zig/src/families/ling/fake_driver.zig b/zig/src/families/ling/fake_driver.zig new file mode 100644 index 0000000000..f5656a899c --- /dev/null +++ b/zig/src/families/ling/fake_driver.zig @@ -0,0 +1,343 @@ +//! A host stand-in for the CUDA driver in checks: device memory in host memory, no-op kernels, a failure on cue. +const std = @import("std"); +const cuda = @import("cuda"); + +const abi = cuda.abi; +const R = abi.Result; +const ok = abi.success; +const backing = std.heap.page_allocator; + +/// What the entry points share (they are C functions, so process-wide): allocations, handles, the countdown. +var held: std.AutoHashMapUnmanaged(u64, usize) = .empty; // device and pinned host allocations by address +var handles: isize = 0; // modules, streams, events, graphs and execs made and not yet released +var calls: u64 = 0; // fallible calls since reset +var fail_at: ?u64 = null; // the call that fails +var code: R = 2; // what it returns (2: out of memory; NOT_READY reaches the caller without a log line) +var bad: usize = 0; // frees of something not held, releases of nothing +var next: usize = 0x1000; +var globals: [1 << 16]u8 align(256) = undefined; + +/// From now on, fallible call number `at` (1-based) returns `with`; null fails none. Bad releases start at 0. +pub fn reset(at: ?u64, with: R) void { + calls = 0; + bad = 0; + fail_at = at; + code = with; +} + +/// Fallible calls since reset, and whether the failure was handed out. +pub fn count() u64 { + return calls; +} + +pub fn fired() bool { + return if (fail_at) |n| calls >= n else false; +} + +/// Allocations and handles held now; frees and releases of things never made. +pub fn live() usize { + return held.count() + @as(usize, @intCast(@max(handles, 0))); +} + +pub fn badReleases() usize { + return bad + @as(usize, @intCast(@max(-handles, 0))); +} + +/// A Driver whose entry points are these (its library is never opened or closed). +pub fn driver() cuda.Driver { + var api: abi.Api = undefined; + inline for (@typeInfo(abi.Api).@"struct".field_names) |name| @field(api, name) = @ptrCast(&passes); + inline for (comptime std.meta.declarations(table)) |name| @field(api, name) = &@field(table, name); + return .{ .lib = undefined, .api = api }; +} + +/// A context on `d` (device 0) that nothing retains. +pub fn context(d: *const cuda.Driver) cuda.Context { + return .{ .d = d, .device = 0, .handle = @ptrFromInt(0x10) }; +} + +/// The countdown: each fallible call takes one step, and the chosen one fails. +fn step() R { + calls += 1; + return if (fail_at != null and calls == fail_at.?) code else ok; +} + +/// Every entry point without out-parameters: a step, nothing else (the C ABI lets a caller pass arguments it ignores). +fn passes() callconv(.c) R { + return step(); +} + +fn handle() ?*anyopaque { + next += 16; + handles += 1; + return @ptrFromInt(next); +} + +fn release() R { + handles -= 1; + return ok; +} + +fn take(n: usize) ?u64 { + const m = backing.alignedAlloc(u8, .fromByteUnits(256), @max(n, 1)) catch return null; + held.put(backing, @intFromPtr(m.ptr), @max(n, 1)) catch { + backing.free(m); + return null; + }; + return @intFromPtr(m.ptr); +} + +fn give(p: u64) R { + const kv = held.fetchRemove(p) orelse { + bad += 1; + return 1; + }; + backing.free(@as([*]align(256) u8, @ptrFromInt(p))[0..kv.value]); + return ok; +} + +fn bytes(p: u64, n: usize) []u8 { + return @as([*]u8, @ptrFromInt(p))[0..n]; +} + +fn cbytes(p: ?*const anyopaque, n: usize) []const u8 { + return @as([*]const u8, @ptrCast(p.?))[0..n]; +} + +const table = struct { + pub fn cuDriverGetVersion(v: *c_int) callconv(.c) R { + v.* = 13000; + return ok; + } + pub fn cuDeviceGet(d: *abi.Device, _: c_int) callconv(.c) R { + d.* = 0; + return ok; + } + pub fn cuDeviceGetCount(n: *c_int) callconv(.c) R { + n.* = 1; + return ok; + } + pub fn cuDeviceGetName(out: [*]u8, len: c_int, _: abi.Device) callconv(.c) R { + const name = "fake"; + @memcpy(out[0..@min(name.len, @as(usize, @intCast(len)))], name[0..@min(name.len, @as(usize, @intCast(len)))]); + if (len > name.len) out[name.len] = 0; + return ok; + } + pub fn cuDeviceGetAttribute(v: *c_int, a: abi.DeviceAttribute, _: abi.Device) callconv(.c) R { + v.* = switch (a) { + .max_threads_per_block => 1024, + .warp_size => 32, + .clock_rate => 1_000_000, + .multiprocessor_count => 48, + .integrated => 1, + .l2_cache_size => 24 << 20, + .compute_capability_major => 12, + .compute_capability_minor => 1, + .max_shared_memory_per_multiprocessor, .max_shared_memory_per_block_optin => 101376, + }; + return step(); + } + pub fn cuDeviceTotalMem_v2(n: *usize, _: abi.Device) callconv(.c) R { + n.* = 128 << 30; + return step(); + } + pub fn cuDevicePrimaryCtxRetain(c: *abi.Context, _: abi.Device) callconv(.c) R { + c.* = @ptrFromInt(0x10); + return step(); + } + pub fn cuDevicePrimaryCtxRelease_v2(_: abi.Device) callconv(.c) R { + return ok; + } + pub fn cuCtxGetCurrent(c: *abi.Context) callconv(.c) R { + c.* = @ptrFromInt(0x10); + return ok; + } + pub fn cuMemGetInfo_v2(free: *usize, total: *usize) callconv(.c) R { + free.* = 100 << 30; + total.* = 128 << 30; + return step(); + } + pub fn cuMemAlloc_v2(p: *abi.DevicePtr, n: usize) callconv(.c) R { + if (step() != ok) return code; + p.* = take(n) orelse return 2; + return ok; + } + pub fn cuMemFree_v2(p: abi.DevicePtr) callconv(.c) R { + return give(p); + } + pub fn cuMemHostAlloc(p: *?*anyopaque, n: usize, _: c_uint) callconv(.c) R { + if (step() != ok) return code; + p.* = @ptrFromInt(take(n) orelse return 2); + return ok; + } + pub fn cuMemHostGetDevicePointer_v2(p: *abi.DevicePtr, host: ?*anyopaque, _: c_uint) callconv(.c) R { + p.* = @intFromPtr(host); + return step(); + } + pub fn cuMemFreeHost(p: ?*anyopaque) callconv(.c) R { + return give(@intFromPtr(p)); + } + pub fn cuMemHostUnregister(_: ?*anyopaque) callconv(.c) R { + return ok; + } + pub fn cuMemcpyHtoD_v2(dst: abi.DevicePtr, src: ?*const anyopaque, n: usize) callconv(.c) R { + if (step() != ok) return code; + @memcpy(bytes(dst, n), cbytes(src, n)); + return ok; + } + pub fn cuMemcpyDtoH_v2(dst: ?*anyopaque, src: abi.DevicePtr, n: usize) callconv(.c) R { + if (step() != ok) return code; + @memcpy(bytes(@intFromPtr(dst), n), bytes(src, n)); + return ok; + } + pub fn cuMemcpyDtoD_v2(dst: abi.DevicePtr, src: abi.DevicePtr, n: usize) callconv(.c) R { + if (step() != ok) return code; + @memmove(bytes(dst, n), bytes(src, n)); + return ok; + } + pub fn cuMemcpyHtoDAsync_v2(dst: abi.DevicePtr, src: ?*const anyopaque, n: usize, _: abi.Stream) callconv(.c) R { + return cuMemcpyHtoD_v2(dst, src, n); + } + pub fn cuMemcpyDtoHAsync_v2(dst: ?*anyopaque, src: abi.DevicePtr, n: usize, _: abi.Stream) callconv(.c) R { + return cuMemcpyDtoH_v2(dst, src, n); + } + pub fn cuMemcpyDtoDAsync_v2(dst: abi.DevicePtr, src: abi.DevicePtr, n: usize, _: abi.Stream) callconv(.c) R { + return cuMemcpyDtoD_v2(dst, src, n); + } + pub fn cuMemcpy2DAsync_v2(m: *const abi.Memcpy2D, _: abi.Stream) callconv(.c) R { + if (step() != ok) return code; + for (0..m.height) |r| @memmove(bytes(m.dst + r * m.dst_pitch, m.width), bytes(m.src + r * m.src_pitch, m.width)); + return ok; + } + pub fn cuMemsetD8_v2(p: abi.DevicePtr, v: u8, n: usize) callconv(.c) R { + if (step() != ok) return code; + @memset(bytes(p, n), v); + return ok; + } + pub fn cuMemsetD32_v2(p: abi.DevicePtr, v: c_uint, n: usize) callconv(.c) R { + if (step() != ok) return code; + @memset(@as([*]align(1) u32, @ptrFromInt(p))[0..n], v); + return ok; + } + pub fn cuMemsetD8Async(p: abi.DevicePtr, v: u8, n: usize, _: abi.Stream) callconv(.c) R { + return cuMemsetD8_v2(p, v, n); + } + pub fn cuMemsetD32Async(p: abi.DevicePtr, v: c_uint, n: usize, _: abi.Stream) callconv(.c) R { + return cuMemsetD32_v2(p, v, n); + } + pub fn cuStreamCreate(s: *abi.Stream, _: c_uint) callconv(.c) R { + if (step() != ok) return code; + s.* = @ptrCast(handle()); + return ok; + } + pub fn cuStreamDestroy_v2(_: abi.Stream) callconv(.c) R { + return release(); + } + pub fn cuStreamIsCapturing(_: abi.Stream, s: *abi.CaptureStatus) callconv(.c) R { + s.* = .none; + return step(); + } + pub fn cuEventCreate(e: *abi.Event, _: c_uint) callconv(.c) R { + if (step() != ok) return code; + e.* = @ptrCast(handle()); + return ok; + } + pub fn cuEventDestroy_v2(_: abi.Event) callconv(.c) R { + return release(); + } + pub fn cuEventElapsedTime(ms: *f32, _: abi.Event, _: abi.Event) callconv(.c) R { + ms.* = 0; + return step(); + } + pub fn cuModuleLoadData(m: *abi.Module, _: ?*const anyopaque) callconv(.c) R { + if (step() != ok) return code; + m.* = @ptrCast(handle()); + return ok; + } + pub fn cuModuleLoadDataEx(m: *abi.Module, image: ?*const anyopaque, _: c_uint, _: ?[*]c_int, _: ?[*]?*anyopaque) callconv(.c) R { + return cuModuleLoadData(m, image); + } + pub fn cuModuleUnload(_: abi.Module) callconv(.c) R { + return release(); + } + pub fn cuModuleGetFunction(f: *abi.Function, _: abi.Module, _: [*:0]const u8) callconv(.c) R { + f.* = @ptrFromInt(0x20); + return step(); + } + pub fn cuModuleGetGlobal_v2(p: *abi.DevicePtr, n: *usize, _: abi.Module, _: [*:0]const u8) callconv(.c) R { + p.* = @intFromPtr(&globals); + n.* = globals.len; + return step(); + } + pub fn cuFuncGetAttribute(v: *c_int, a: abi.FunctionAttribute, _: abi.Function) callconv(.c) R { + v.* = switch (a) { + .max_threads_per_block => 1024, + .num_regs => 32, + .ptx_version, .binary_version => 121, + else => 0, + }; + return step(); + } + pub fn cuOccupancyMaxActiveBlocksPerMultiprocessor(n: *c_int, _: abi.Function, _: c_int, _: usize) callconv(.c) R { + n.* = 4; + return step(); + } + pub fn cuStreamEndCapture(_: abi.Stream, g: *abi.Graph) callconv(.c) R { + if (step() != ok) return code; + g.* = @ptrCast(handle()); + return ok; + } + pub fn cuGraphCreate(g: *abi.Graph, _: c_uint) callconv(.c) R { + if (step() != ok) return code; + g.* = @ptrCast(handle()); + return ok; + } + pub fn cuGraphDestroy(_: abi.Graph) callconv(.c) R { + return release(); + } + pub fn cuGraphAddKernelNode_v2(n: *abi.GraphNode, _: abi.Graph, _: ?[*]const abi.GraphNode, _: usize, _: *const abi.KernelNodeParams) callconv(.c) R { + n.* = @ptrFromInt(0x30); + return step(); + } + pub fn cuGraphGetNodes(_: abi.Graph, _: ?[*]abi.GraphNode, n: *usize) callconv(.c) R { + n.* = 0; + return step(); + } + pub fn cuGraphInstantiateWithFlags(x: *abi.GraphExec, _: abi.Graph, _: u64) callconv(.c) R { + if (step() != ok) return code; + x.* = @ptrCast(handle()); + return ok; + } + pub fn cuGraphExecDestroy(_: abi.GraphExec) callconv(.c) R { + return release(); + } + pub fn cuGetErrorName(_: R, s: *?[*:0]const u8) callconv(.c) R { + s.* = "CUDA_ERROR_INJECTED"; + return ok; + } + pub fn cuGetErrorString(_: R, s: *?[*:0]const u8) callconv(.c) R { + s.* = "a failure the fake driver was told to give"; + return ok; + } +}; + +test "the fake driver: buffers live in host memory, the chosen call fails, frees of strangers are counted" { + const d = driver(); + reset(null, ok); + var b = try cuda.DeviceBuffer.fromHost(&d, "abcdefgh"); + var back: [8]u8 = undefined; + try b.download(0, &back); + try std.testing.expectEqualStrings("abcdefgh", &back); + try std.testing.expectEqual(@as(usize, 1), live()); + b.free(); + try std.testing.expectEqual(@as(usize, 0), live()); + reset(2, abi.error_not_ready); + var s = try cuda.Stream.init(&d, false); + try std.testing.expectError(error.NotReady, cuda.DeviceBuffer.alloc(&d, 64)); + try std.testing.expect(fired()); + s.deinit(); + try std.testing.expectEqual(@as(usize, 0), live()); + _ = d.api.cuMemFree_v2(0x1234); + try std.testing.expectEqual(@as(usize, 1), badReleases()); + reset(null, ok); +} diff --git a/zig/src/families/ling/store_test.zig b/zig/src/families/ling/store_test.zig new file mode 100644 index 0000000000..80ee8d1179 --- /dev/null +++ b/zig/src/families/ling/store_test.zig @@ -0,0 +1,43 @@ +//! weights.Store on the fake driver: each host or device allocation failing in turn leaves nothing lost or freed twice. +const std = @import("std"); +const cuda = @import("cuda"); +const proj = @import("weights_proj.zig"); +const fake = @import("fake_driver.zig"); + +const t = std.testing; + +/// Twelve buffers (past the lists' first growth), one upload, a lookup; the store freed on every path. +fn fill(gpa: std.mem.Allocator, d: *const cuda.Driver) !void { + var s: proj.Store = .{ .gpa = gpa, .d = d }; + defer s.deinit(); + for (0..12) |i| { + var name: [32]u8 = undefined; + _ = try s.alloc(64 + i, try std.fmt.bufPrint(&name, "buffer.{d}", .{i})); + } + _ = try s.keep("bytes", "kept"); + try t.expect(s.find("buffer.3") != null and s.find("kept").?.len == 5); +} + +test "Store.alloc: each host allocation failing in turn leaks nothing and frees no buffer twice" { + const d = fake.driver(); + fake.reset(null, cuda.abi.success); + try t.checkAllAllocationFailures(t.allocator, fill, .{&d}); + try t.expectEqual(@as(usize, 0), fake.live()); + try t.expectEqual(@as(usize, 0), fake.badReleases()); +} + +test "Store.alloc: each device call failing in turn (allocations, the upload) leaves nothing held or freed twice" { + const d = fake.driver(); + fake.reset(null, cuda.abi.success); + try fill(t.allocator, &d); + const n = fake.count(); + try t.expect(n >= 13); + for (1..n + 1) |k| { + fake.reset(k, cuda.abi.error_not_ready); + try t.expectError(error.NotReady, fill(t.allocator, &d)); + try t.expect(fake.fired()); + try t.expectEqual(@as(usize, 0), fake.live()); + try t.expectEqual(@as(usize, 0), fake.badReleases()); + } + fake.reset(null, cuda.abi.success); +} diff --git a/zig/src/families/ling/weights_proj.zig b/zig/src/families/ling/weights_proj.zig index 3a3404015b..1478695c85 100644 --- a/zig/src/families/ling/weights_proj.zig +++ b/zig/src/families/ling/weights_proj.zig @@ -38,15 +38,15 @@ pub const Store = struct { s.lanes.deinit(s.gpa); } - /// A device buffer of `len` bytes the store owns, named `name`. + /// A device buffer of `len` bytes the store owns, named `name`; on failure the store holds nothing new. pub fn alloc(s: *Store, len: usize, name: []const u8) !u64 { + try s.buffers.ensureUnusedCapacity(s.gpa, 1); + try s.named.ensureUnusedCapacity(s.gpa, 1); + const owned = try s.gpa.dupe(u8, name); + errdefer s.gpa.free(owned); const b = try cuda.DeviceBuffer.alloc(s.d, @max(len, 16)); - errdefer { - var x = b; - x.free(); - } - try s.buffers.append(s.gpa, b); - try s.named.append(s.gpa, .{ .name = try s.gpa.dupe(u8, name), .ptr = b.ptr, .len = len }); + s.buffers.appendAssumeCapacity(b); + s.named.appendAssumeCapacity(.{ .name = owned, .ptr = b.ptr, .len = len }); return b.ptr; } From 0ba510adb3e3e663121cf6380219e0a7a0132c9a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 16 Oct 2026 07:20:37 +0200 Subject: [PATCH 27/34] ling: a new stream's slot is chosen before its map entry exists (lanes.zig read the entry's unset slot) opening() did getOrPut and then freeSlot(), which walks every entry, the new one included, before its value was set. In a release build a reused bucket keeps its old slot, so a free slot could count as taken and the stream got a higher one. Lanes.claim picks the slot first, then puts the entry. - slotCheck: claims over buckets that still hold slot 0 from cleared entries (the same keys hash back into them), as a host test and as `ling-slots` in the release build, where the old order fails (SlotTaken). --- zig/src/families/ling/cuda.zig | 1 + zig/src/families/ling/lanes.zig | 40 ++++++++++++++++++++++++++++----- zig/tests/cuda/main.zig | 7 ++++++ 3 files changed, 42 insertions(+), 6 deletions(-) diff --git a/zig/src/families/ling/cuda.zig b/zig/src/families/ling/cuda.zig index 33c42efd23..0743957aa8 100644 --- a/zig/src/families/ling/cuda.zig +++ b/zig/src/families/ling/cuda.zig @@ -55,6 +55,7 @@ test { _ = native; _ = @import("fake_driver.zig"); _ = @import("store_test.zig"); + _ = lanes; } test "config's fixed widths are the kernels' own: KDA heads, the latent cache, the rope" { diff --git a/zig/src/families/ling/lanes.zig b/zig/src/families/ling/lanes.zig index 91ed5ce646..bbcbc51c1c 100644 --- a/zig/src/families/ling/lanes.zig +++ b/zig/src/families/ling/lanes.zig @@ -94,6 +94,14 @@ pub const Lanes = struct { } } + /// The stream's lane; a new one takes the lowest free slot, chosen before the map holds its (unset) entry. + pub fn claim(self: *Lanes, gpa: std.mem.Allocator, s: *const lanes.Stream) !*Lane { + if (self.lanes.getPtr(s)) |l| return l; + const slot = self.freeSlot(); + try self.lanes.put(gpa, s, .{ .slot = slot }); + return self.lanes.getPtr(s).?; + } + /// The stream's sequence bound, its last verify committed (`kept` rows, else all of them). fn bind(self: *Lanes, s: *const lanes.Stream, kept: ?usize) !*Lane { const l = self.lanes.getPtr(s) orelse return error.UnknownStream; @@ -139,12 +147,11 @@ pub const Lanes = struct { fn opening(self: *Lanes, s: *lanes.Stream) !?engine.Resume { const ids = s.prompt(); if (ids.len + s.max_new + max_rows >= self.e.capacity) return error.PromptTooLong; - const gop = try self.lanes.getOrPut(self.e.gpa, s); - if (!gop.found_existing) gop.value_ptr.* = .{ .slot = self.freeSlot() }; - gop.value_ptr.pending_rows = null; - gop.value_ptr.held_n = 0; - gop.value_ptr.fill = null; - try self.e.bind(gop.value_ptr.slot); + const l = try self.claim(self.e.gpa, s); + l.pending_rows = null; + l.held_n = 0; + l.fill = null; + try self.e.bind(l.slot); try self.e.begin(s.sampling); s.cached = 0; return self.resumeFrom(s, ids.len); @@ -278,3 +285,24 @@ pub const Lanes = struct { self.e.release(kv.value.slot); } }; + +/// A stand-in stream key: claim hashes the address and never reads through it. +fn key(i: usize) *const lanes.Stream { + return @ptrFromInt(i * 4096); +} + +/// Slots claimed over map buckets still holding slot 0 from cleared entries (a release build leaves a new value as is). +pub fn slotCheck(gpa: std.mem.Allocator) !void { + var ll: Lanes = .{ .e = undefined }; + defer ll.lanes.deinit(gpa); + for (1..9) |i| try ll.lanes.put(gpa, key(i), .{ .slot = 0 }); + ll.lanes.clearRetainingCapacity(); + const want = [_]struct { k: usize, slot: u32 }{ .{ .k = 1, .slot = 0 }, .{ .k = 2, .slot = 1 }, .{ .k = 1, .slot = 0 }, .{ .k = 3, .slot = 2 } }; + for (want) |w| if ((try ll.claim(gpa, key(w.k))).slot != w.slot) return error.SlotTaken; + if (!ll.lanes.remove(key(1))) return error.SlotTaken; + if ((try ll.claim(gpa, key(4))).slot != 0 or (try ll.claim(gpa, key(5))).slot != 3) return error.SlotTaken; +} + +test "a new stream takes the lowest free slot, a known one keeps its own, a released one's goes to the next" { + try slotCheck(std.testing.allocator); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index f14dbcc650..4b682437dc 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -116,6 +116,7 @@ const usage = \\ ling-clamp Ling's SwiGLU order min(silu(g), L): glue and the dir's grouped experts vs float64 \\ ling-clamp-old the grouped expert instances vs an older build's fatbins, byte for byte \\ ling-lanes [ref] [--layers a,b] [--mtp] Ling's lane backend: chunked, drafted, resumed, concurrent == plain + \\ ling-slots the lane backend's slots over reused map buckets, in this build's mode (no GPU work) \\ ling-mtp-real the real MTP layer vs HF's (oracle/ling_mtp.py); cut, window == decode rows \\ ling-weights-real real Ling-3.1 layers 0, 4, 7 (4 experts) vs float64 and torch (oracle/ling_weights.py) \\ ling-pack [experts [stage]] one MoE layer's experts: read, pack and cold load timed; bytes vs the host packers @@ -233,6 +234,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "ling-clamp-old")) return @import("ling_clamp.zig").old(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "ling-mtp-real")) return @import("ling_mtp_real.zig").run(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "ling-lanes")) return @import("ling_lanes.zig").run(gpu, rest); + if (std.mem.eql(u8, cmd, "ling-slots")) return lingSlots(gpu); if (std.mem.eql(u8, cmd, "ling-weights-real")) return ling_real.run(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "ling-pack")) return @import("ling_pack.zig").run(gpu, rest); if (std.mem.eql(u8, cmd, "glm-lanes-concurrent")) return glm_lanes_tests.concurrentCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); @@ -256,3 +258,8 @@ fn hasFlag(flags: []const [:0]const u8, name: []const u8) bool { for (flags) |f| if (std.mem.eql(u8, f, name)) return true; return false; } + +fn lingSlots(gpu: check.Gpu) !void { + try @import("ling_cuda").lanes.slotCheck(gpu.gpa); + check.pass("ling-slots: new streams take the lowest free slot over buckets holding stale slots ({t})", .{@import("builtin").mode}); +} From 82b3817bbf56bddb91f6c6afecaa75446585695b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 16 Oct 2026 07:29:46 +0200 Subject: [PATCH 28/34] ling: rows past a sequence's capacity are refused before any layer runs, not at the first MLA layer The capacity check sat in block_mla, after the KDA layers before it had run: a prompt chunk past the capacity had already committed those layers' states and conv windows (and a window had written its saves) when CacheFull came back. Forward.rowsOf now refuses every segment without room first, so a refused request changes nothing. The MTP layer, one MLA block, already refused before its cache write. - ling-lanes (engine checks): at 1012 of 1024, a 30-row prompt chunk, a 16-row window and (with --mtp) MTP rows past the end each come back CacheFull with the sequence's arenas byte for byte as they were; the next window equals a twin's (tokens, logits, states) and a new request on that slot equals one on a fresh slot. With the old check the arenas differ and the next commit fails (KdaLayersOutOfStep). --- zig/src/families/ling/forward.zig | 3 +- zig/tests/cuda/ling_lanes_bytes.zig | 69 +++++++++++++++++++++++++++++ zig/tests/cuda/ling_lanes_rig.zig | 2 +- 3 files changed, 72 insertions(+), 2 deletions(-) diff --git a/zig/src/families/ling/forward.zig b/zig/src/families/ling/forward.zig index 6cb0cdd337..f31edfb54d 100644 --- a/zig/src/families/ling/forward.zig +++ b/zig/src/families/ling/forward.zig @@ -29,11 +29,12 @@ pub const Forward = struct { return .{ .m = f.m, .s = f.s, .prompt = f.b.prompt, .part = f.b.part, .kda_wide = f.kda_wide }; } - /// Segments in row order from row 0, each on its own committed sequence; a window's at most window_rows rows. + /// Segments in row order from row 0, each on its own committed sequence with room for its rows; windows <= 16 rows. fn rowsOf(f: *const Forward, segs: []const Seg) !usize { var rows: usize = 0; for (segs, 0..) |sg, i| { if (sg.row0 != rows or sg.rows == 0 or sg.st.pending != 0) return error.BadSegments; + if (sg.st.pos + sg.rows > sg.st.capacity) return error.CacheFull; // before a prompt's KDA layers commit if (!f.b.prompt and sg.rows > window_rows) return error.WindowTooWide; for (segs[0..i]) |o| if (o.st == sg.st) return error.BadSegments; rows += sg.rows; diff --git a/zig/tests/cuda/ling_lanes_bytes.zig b/zig/tests/cuda/ling_lanes_bytes.zig index d03e9b560f..922c7eb231 100644 --- a/zig/tests/cuda/ling_lanes_bytes.zig +++ b/zig/tests/cuda/ling_lanes_bytes.zig @@ -86,9 +86,78 @@ pub fn engineChecks(c: *Cover, e: *Engine, prompt: []const u32, w: []const u32) try rig.sameSeq(c, e, "9-row window, keep 4, 5-row window == 9 single rows", 10, 11); try sharedChecks(c, e, prompt, w); if (e.mtp != null) try mtpChain(c, e, prompt, w); + try capacityRefusal(c, e, prompt, w); for ([_]u32{ 10, 11, 12, 13, 14 }) |s| e.release(s); } +/// A sequence's arenas as they stand: the layers' caches and states, the MTP layer's cache. +fn arenas(c: *Cover, e: *Engine, slot: u32) ![2][]u8 { + const q = e.seqs.items[slot].?; + const st = &q.st.arena.buf; + const hd: []u8 = if (q.hd) |*h| try rig.back(e, c.a, h.arena.buf.ptr, h.arena.buf.len) else &.{}; + return .{ try rig.back(e, c.a, st.ptr, st.len), hd }; +} + +/// Rows past the capacity refused with the arenas unchanged; the next window and a new request then equal twins'. +fn capacityRefusal(c: *Cover, e: *Engine, prompt: []const u32, w: []const u32) !void { + const a = c.a; + const cap = e.capacity; + const long = try a.alloc(u32, cap + 40); + for (long, 0..) |*x, i| x.* = prompt[i % prompt.len]; + var pend: [2]u32 = undefined; + for (0..2) |k| { // twins at cap - 12: a pass of cap - 24 rows, a 12-row window kept whole + const t = try pass(e, @intCast(50 + k), long[0 .. cap - 24], 512, null, null); + var ids: [12]u32 = undefined; + var out: [12]u32 = undefined; + ids[0] = t; + @memcpy(ids[1..], w[0..11]); + try e.verify(&ids, &out); + try e.keep(12); + pend[k] = out[11]; + } + try e.bind(50); + const before = try arenas(c, e, 50); + const q = e.seqs.items[50].?; + const cur = try a.dupe(u1, q.st.cur); + const pos = q.st.pos; + var refused: usize = 0; + // a prompt chunk of 30 rows (its KDA layers would commit in place), then a 16-row window + e.chunk = 512; + if (e.prefillChunk(long[0 .. cap + 18], pos, null, null)) |_| {} else |err| refused += @intFromBool(err == error.CacheFull); + var ids: [16]u32 = undefined; + var out: [16]u32 = undefined; + ids[0] = pend[0]; + @memcpy(ids[1..], w[20..35]); + if (e.verify(&ids, &out)) |_| {} else |err| refused += @intFromBool(err == error.CacheFull); + var asked: usize = 2; + if (e.mtp) |*h| { // one row past the capacity into the MTP layer's cache (its MLA block refuses before writing) + asked += 1; + const n = cap + 1 - q.hd.?.pos; + if (h.absorb(&e.m, e.stream, &e.bp, &q.hd.?, @ptrCast(long[0..n]), e.bp.normed)) |_| {} else |err| refused += @intFromBool(err == error.CacheFull); + } + const after = try arenas(c, e, 50); + const same = std.mem.eql(u8, before[0], after[0]) and std.mem.eql(u8, before[1], after[1]) and + std.mem.eql(u1, cur, q.st.cur) and q.st.pos == pos and q.st.pending == 0; + try c.expect(refused == asked and same, "{d} of {d} requests past the capacity ({d} at {d}) refused with CacheFull, the arenas ({d} and {d} bytes) and position as they were", .{ refused, asked, cap, pos, before[0].len, before[1].len }); + // the next window on both twins: tokens, logits, then states + var got: [2][4]u32 = undefined; + var lo: [2][]u8 = undefined; + for (0..2) |k| { + try e.bind(@intCast(50 + k)); + var win: [4]u32 = .{ pend[k], w[40], w[41], w[42] }; + try e.verify(&win, &got[k]); + lo[k] = try logitsRows(e, a, e.bd.logits, 4); + try e.keep(4); + } + try c.expect(std.mem.eql(u32, &got[0], &got[1]) and std.mem.eql(u8, lo[0], lo[1]), "the window after the refusals: tokens and logits == the twin's", .{}); + try rig.sameSeq(c, e, "after the refusals and a window, the twin", 50, 51); + const t50 = try pass(e, 50, prompt, 37, null, null); + const t52 = try pass(e, 52, prompt, 37, null, null); + try c.expect(t50 == t52, "a new request on the refused slot: first token {d}, on a fresh slot {d}", .{ t50, t52 }); + try rig.sameSeq(c, e, "a new request on the refused slot == on a fresh slot", 50, 52); + for ([_]u32{ 50, 51, 52 }) |s| e.release(s); +} + /// MTP chains of 3 and of 1 on twins: the same first draft, and after the same windows the chain's rows rolled back. fn mtpChain(c: *Cover, e: *Engine, prompt: []const u32, w: []const u32) !void { var t: [2]u32 = undefined; diff --git a/zig/tests/cuda/ling_lanes_rig.zig b/zig/tests/cuda/ling_lanes_rig.zig index e54b4bcecc..1e3b44b9a7 100644 --- a/zig/tests/cuda/ling_lanes_rig.zig +++ b/zig/tests/cuda/ling_lanes_rig.zig @@ -193,7 +193,7 @@ pub const Cover = struct { } }; -fn back(e: *Engine, a: std.mem.Allocator, ptr: u64, len: usize) ![]u8 { +pub fn back(e: *Engine, a: std.mem.Allocator, ptr: u64, len: usize) ![]u8 { const out = try a.alloc(u8, len); try e.stream.synchronize(); const b: cuda.DeviceBuffer = .{ .d = e.d, .ptr = ptr, .len = len }; From 0ba2adfb1aba4946deb428c26577a93177acccfa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 16 Oct 2026 07:53:05 +0200 Subject: [PATCH 29/34] ling: Engine.init frees its sequences when binding slot 0 or the last sync fails; every init failure checked The one gap in init's errdefer chain: sequence 0 (its arena, MTP cache and hold row) and the slot list leaked when bind(0) or the closing synchronize failed. Both are freed now. - ling-init-faults DIR [--mtp]: native.open (Engine.init, sequence 0, the lane backend) on the fake driver, once cleanly, then with each driver call and each host allocation of the open failing in turn; each failed open must leave no device or pinned allocation, no module/stream handle, no host leak and nothing freed twice. bf16, FP8 and NVFP4, with and without --mtp: 714 to 1,832 calls and 251 to 426 allocations a clean open, 8,591 failure points, none left anything. Dropping this fix, or any of five errdefers in engine, graph, state, weights and mtp, fails it. --- docs/recipes/ling-3.1.md | 5 ++ zig/src/families/ling/engine.zig | 4 ++ zig/tests/cuda/ling_faults.zig | 82 ++++++++++++++++++++++++++++++++ zig/tests/cuda/main.zig | 2 + 4 files changed, 93 insertions(+) create mode 100644 zig/tests/cuda/ling_faults.zig diff --git a/docs/recipes/ling-3.1.md b/docs/recipes/ling-3.1.md index e548852104..ab39cdfe26 100644 --- a/docs/recipes/ling-3.1.md +++ b/docs/recipes/ling-3.1.md @@ -264,6 +264,11 @@ the file, the grouped experts' gate-up (SwiGLU) and down, and each small tensor `dt_bias`, router, bias) byte for byte. Measured: bf16 85 projections, FP8 and NVFP4 61 each, worst 0.433 of the bound (2^-7 of |y| plus 2^-12 of the sum of |x w|); 37 small tensors equal. +`ling-init-faults DIR [--mtp]` opens the backend (engine, sequence 0, lane backend) on `fake_driver.zig`, a host +stand-in for the driver, once cleanly and then with each of its driver calls, and each of its host allocations, +failing in turn: every failed open must leave nothing allocated, held or freed twice (no GPU work; FP8 with `--mtp` +is 869 calls and 394 allocations). `ling-slots` checks the lane backend's slot choice in the release build. + With the real checkpoint: ```bash diff --git a/zig/src/families/ling/engine.zig b/zig/src/families/ling/engine.zig index 8b4efd7df8..3e64f0554c 100644 --- a/zig/src/families/ling/engine.zig +++ b/zig/src/families/ling/engine.zig @@ -126,6 +126,10 @@ pub const Engine = struct { errdefer e.smp.deinit(); e.pairs = try cuda.DeviceBuffer.alloc(d, 8 * window_rows); errdefer e.pairs.free(); + errdefer { + for (0..e.seqs.items.len) |i| e.free(i); + e.seqs.deinit(gpa); + } try e.bind(0); try e.stream.synchronize(); return e; diff --git a/zig/tests/cuda/ling_faults.zig b/zig/tests/cuda/ling_faults.zig new file mode 100644 index 0000000000..35efebc760 --- /dev/null +++ b/zig/tests/cuda/ling_faults.zig @@ -0,0 +1,82 @@ +//! `ling-init-faults DIR [--mtp] [--every N]`: the backend opened on the fake driver, each call and allocation failing. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const ling = @import("ling_cuda"); + +const fake = ling.fake_driver; +const Dbg = std.heap.DebugAllocator(.{ .safety = true }); // leak checks in release builds too + +/// Opened (the engine, its sequence, the lane backend) and closed again. +fn once(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, mtp: bool) !void { + const l = try open(gpa, io, ctx, dir, mtp); + l.deinit(l.ctx); +} + +fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, mtp: bool) !ling.native.Loaded { + return ling.native.open(gpa, io, ctx, dir, null, .{ .context = 256, .drafts = mtp }); +} + +/// What one failed open left: the error, host leaks, the fake driver's live and bad releases. +const Left = struct { err: ?anyerror, leaked: bool, live: usize, bad: usize }; + +fn failed(io: std.Io, ctx: *const cuda.Context, dir: []const u8, mtp: bool, call: ?u64, alloc: ?usize) Left { + fake.reset(call, cuda.abi.error_not_ready); + var dbg: Dbg = .init; + var fa = std.testing.FailingAllocator.init(dbg.allocator(), .{ .fail_index = alloc orelse std.math.maxInt(usize) }); + const err: ?anyerror = if (once(fa.allocator(), io, ctx, dir, mtp)) |_| null else |e| e; + const live = fake.live(); + const bad = fake.badReleases(); + const leaked = dbg.deinit() == .leak; + return .{ .err = err, .leaked = leaked, .live = live, .bad = bad }; +} + +pub fn run(gpu: check.Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 1) return error.MissingArgument; + const dir = rest[0]; + var mtp = false; + var every: usize = 1; + var i: usize = 1; + while (i < rest.len) : (i += 1) { + if (std.mem.eql(u8, rest[i], "--mtp")) mtp = true; + if (std.mem.eql(u8, rest[i], "--every") and i + 1 < rest.len) { + every = try std.fmt.parseInt(usize, rest[i + 1], 10); + i += 1; + } + } + const d = fake.driver(); + const ctx = fake.context(&d); + // a clean open first: its calls and host allocations are the failure points + fake.reset(null, cuda.abi.success); + var dbg: Dbg = .init; + var counter = std.testing.FailingAllocator.init(dbg.allocator(), .{}); + const clean = try open(counter.allocator(), gpu.io, &ctx, dir, mtp); + const calls = fake.count(); // the open's own calls and allocations: close's are not failure points + const allocs = counter.allocations; + clean.deinit(clean.ctx); + try check.expect(dbg.deinit() == .ok and fake.live() == 0 and fake.badReleases() == 0, "a clean open and close leaves nothing", .{}); + var bad: usize = 0; + var tried: usize = 0; + var k: u64 = 1; + while (k <= calls) : (k += every) { + const l = failed(gpu.io, &ctx, dir, mtp, k, null); + tried += 1; + if (l.err == null or !fake.fired() or l.leaked or l.live != 0 or l.bad != 0) { + bad += 1; + if (bad <= 8) std.debug.print(" driver call {d} failing: {?t}, host leak {}, {d} held, {d} bad releases\n", .{ k, l.err, l.leaked, l.live, l.bad }); + } + } + var a: usize = 0; + while (a < allocs) : (a += every) { + const l = failed(gpu.io, &ctx, dir, mtp, null, a); + tried += 1; + if (l.err == null or l.leaked or l.live != 0 or l.bad != 0) { + bad += 1; + if (bad <= 16) std.debug.print(" host allocation {d} failing: {?t}, host leak {}, {d} held, {d} bad releases\n", .{ a, l.err, l.leaked, l.live, l.bad }); + } + } + fake.reset(null, cuda.abi.success); + std.debug.print("RESULT ling-init-faults{s}: {d} driver calls and {d} host allocations in a clean open; {d} failure points tried, {d} left something behind\n", .{ if (mtp) " --mtp" else "", calls, allocs, tried, bad }); + try check.expect(bad == 0, "every failing driver call and host allocation in the open is cleaned up", .{}); + check.pass("ling-init-faults{s}: {d} failure points, each open refused with nothing left held, leaked or freed twice", .{ if (mtp) " --mtp" else "", tried }); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 4b682437dc..5cb0aaa1e0 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -117,6 +117,7 @@ const usage = \\ ling-clamp-old the grouped expert instances vs an older build's fatbins, byte for byte \\ ling-lanes [ref] [--layers a,b] [--mtp] Ling's lane backend: chunked, drafted, resumed, concurrent == plain \\ ling-slots the lane backend's slots over reused map buckets, in this build's mode (no GPU work) + \\ ling-init-faults [--mtp] [--every N] the backend opened with each driver call, then each host allocation, failing \\ ling-mtp-real the real MTP layer vs HF's (oracle/ling_mtp.py); cut, window == decode rows \\ ling-weights-real real Ling-3.1 layers 0, 4, 7 (4 experts) vs float64 and torch (oracle/ling_weights.py) \\ ling-pack [experts [stage]] one MoE layer's experts: read, pack and cold load timed; bytes vs the host packers @@ -235,6 +236,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "ling-mtp-real")) return @import("ling_mtp_real.zig").run(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "ling-lanes")) return @import("ling_lanes.zig").run(gpu, rest); if (std.mem.eql(u8, cmd, "ling-slots")) return lingSlots(gpu); + if (std.mem.eql(u8, cmd, "ling-init-faults")) return @import("ling_faults.zig").run(gpu, rest); if (std.mem.eql(u8, cmd, "ling-weights-real")) return ling_real.run(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "ling-pack")) return @import("ling_pack.zig").run(gpu, rest); if (std.mem.eql(u8, cmd, "glm-lanes-concurrent")) return glm_lanes_tests.concurrentCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); From 5649c63b66c0ccbf18416ee72c07599a49e12f94 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 16 Oct 2026 08:06:43 +0200 Subject: [PATCH 30/34] ling: every tensor of the full checkpoint checked against the loader's names, dtypes and shapes, from headers alone The loader's shapes now come from one table, names.dims (each role's [n, k], value count or conv shape); weights.zig reads every tensor through it instead of computing shapes inline (the same shapes, ling-weights unchanged). - ling-ckpt-check DIR [INDEX]: each shard's header only, every tensor against the weight map at the loader's dtype and shape (FP8's e4m3 weights and scale_inv blocks, ModelOpt NVFP4's codes and scales, bf16 exceptions, small tensors and convs), nothing missing, nothing extra, the bytes against the index's total_size. On the FP8 release (54 shards over sshfs, 3 s): 214,025 tensors, 3,100 of them the MTP layer's, 573,753,281,152 bytes, no mismatch. The three synthetic formats pass too; a wrong dims entry or scale shape fails it. --- docs/recipes/ling-3.1.md | 5 ++ zig/src/families/ling/names.zig | 48 +++++++++++++++ zig/src/families/ling/weights.zig | 97 ++++++++++++++++--------------- zig/tests/cuda/ling_ckpt.zig | 95 ++++++++++++++++++++++++++++++ zig/tests/cuda/main.zig | 2 + 5 files changed, 201 insertions(+), 46 deletions(-) create mode 100644 zig/tests/cuda/ling_ckpt.zig diff --git a/docs/recipes/ling-3.1.md b/docs/recipes/ling-3.1.md index ab39cdfe26..1f374fb2b8 100644 --- a/docs/recipes/ling-3.1.md +++ b/docs/recipes/ling-3.1.md @@ -276,6 +276,11 @@ python zig/tests/cuda/oracle/ling_weights.py /path/to/Ling-3.1-flash /tmp/ling-r zig-out/bin/tf-cuda-test ling-weights-real /path/to/Ling-3.1-flash /tmp/ling-ref ``` +`ling-ckpt-check DIR [INDEX]` reads every shard's header and nothing else: each of the release's 214,025 tensors is +one the weight map names, at the dtype and shape the loader reads (`names.dims`, the loader's own table, and each +format's side tensors), none missing, and their bytes add up to the index's `total_size` (573,753,281,152). Over a +network mount it takes seconds. + `ling-mtp-real DIR FIXTURE` runs the MTP layer (fixture from `oracle/ling_mtp.py SRC MODELING_PKG ROWS.pt OUT`). `ling-layer-real DIR FIXTURE` runs one whole layer (fixture from `oracle/ling_layer.py CHAIN L OUT`, the captured HF chain's `h.pt`, `h.pt`, `r.pt`). `ling-weights-real` loads layers 0, 4 and 7 with experts 0-3 (769 MiB on the device) and compares 35 products with torch's float64 diff --git a/zig/src/families/ling/names.zig b/zig/src/families/ling/names.zig index b6f3786acb..f59d867ba7 100644 --- a/zig/src/families/ling/names.zig +++ b/zig/src/families/ling/names.zig @@ -148,6 +148,54 @@ pub fn side(f: config.Quant) []const []const u8 { }; } +/// What the loader reads for a role: a linear [n, k] (`.weight` and its side tensors), `values` values, a conv. +pub const Dims = union(enum) { linear: [2]usize, values: usize, conv: [3]usize }; + +/// The loader's shape for each role (weights.zig reads every tensor through it). +pub fn dims(c: *const Config, role: Role) Dims { + const d: usize = c.hidden; + const n: usize = @as(usize, c.kda_heads) * c.kda_dim; + const h: usize = c.mla_heads; + return switch (role) { + .embed, .lm_head => .{ .linear = .{ c.vocab, d } }, + .final_norm, .attn_norm, .mlp_norm, .mtp_enorm, .mtp_hnorm, .mtp_final_norm => .{ .values = d }, + .kda_q, .kda_k, .kda_v, .kda_f, .kda_g => .{ .linear = .{ n, d } }, + .kda_b => .{ .linear = .{ c.kda_heads, d } }, + .kda_o => .{ .linear = .{ d, n } }, + .kda_q_conv, .kda_k_conv, .kda_v_conv => .{ .conv = .{ n, 1, c.conv } }, + .kda_a_log => .{ .values = c.kda_heads }, + .kda_dt_bias => .{ .values = n }, + .kda_o_norm => .{ .values = c.kda_dim }, + .mla_q => .{ .linear = .{ h * c.qkDim(), d } }, + .mla_q_a => .{ .linear = .{ c.q_lora, d } }, + .mla_q_a_norm => .{ .values = c.q_lora }, + .mla_q_b => .{ .linear = .{ h * c.qkDim(), c.q_lora } }, + .mla_kv_a => .{ .linear = .{ c.kv_lora + c.rope_dim, d } }, + .mla_kv_a_norm => .{ .values = c.kv_lora }, + .mla_kv_b => .{ .linear = .{ h * (c.nope + c.v_dim), c.kv_lora } }, + .mla_gate => .{ .linear = .{ h, d } }, + .mla_o => .{ .linear = .{ d, h * c.v_dim } }, + .dense_gate, .dense_up => .{ .linear = .{ c.dense_inter, d } }, + .dense_down => .{ .linear = .{ d, c.dense_inter } }, + .router => .{ .linear = .{ c.experts, d } }, + .router_bias => .{ .values = c.experts }, + .expert_gate, .expert_up => .{ .linear = .{ c.moe_inter, d } }, + .expert_down => .{ .linear = .{ d, c.moe_inter } }, + .shared_gate, .shared_up => .{ .linear = .{ c.shared_inter, d } }, + .shared_down => .{ .linear = .{ d, c.shared_inter } }, + .mtp_eh => .{ .linear = .{ d, 2 * d } }, + }; +} + +/// A role's value count: a linear's n k, a small tensor's values, a conv's n taps. +pub fn numel(c: *const Config, role: Role) usize { + return switch (dims(c, role)) { + .linear => |x| x[0] * x[1], + .values => |x| x, + .conv => |x| x[0] * x[1] * x[2], + }; +} + const tables = [_][]const Part{ &norms, &kda, &mla_q_lora, &mla_q_whole, &mla, &mla_gate, &dense, &router, &router_bias, &shared, &expert, &mtp }; /// Role's module in layer i (`e` for routed experts) as the map names it: the name before `.weight`. diff --git a/zig/src/families/ling/weights.zig b/zig/src/families/ling/weights.zig index f2339a276d..e1b7b2f16a 100644 --- a/zig/src/families/ling/weights.zig +++ b/zig/src/families/ling/weights.zig @@ -98,61 +98,64 @@ const Loader = struct { return names.moduleName(&L.buf, i, role, null); } - fn lin(L: *Loader, i: u32, role: Role, n: usize, k: usize) !Proj { - return L.r.linear(try L.module(i, role), n, k); + /// Linear `role` of layer i at names.dims' [n, k]. + fn lin(L: *Loader, i: u32, role: Role) !Proj { + const nk = names.dims(L.c, role).linear; + return L.r.linear(try L.module(i, role), nk[0], nk[1]); } - /// A norm (`.weight`) or bare tensor (A_log, dt_bias, expert_bias) of `numel` values, as `to`. - fn small(L: *Loader, i: u32, role: Role, numel: usize, to: core.checkpoint.DType) !u64 { + /// A norm (`.weight`) or bare tensor (A_log, dt_bias, expert_bias) of names.numel values, as `to`. + fn small(L: *Loader, i: u32, role: Role, to: core.checkpoint.DType) !u64 { var b: [200]u8 = undefined; const m = try L.module(i, role); const bare = role == .kda_a_log or role == .kda_dt_bias or role == .router_bias; - return L.r.small(if (bare) m else try std.fmt.bufPrint(&b, "{s}.weight", .{m}), numel, to); + return L.r.small(if (bare) m else try std.fmt.bufPrint(&b, "{s}.weight", .{m}), names.numel(L.c, role), to); + } + + /// The MLP of gate, up and down `role`s under `prefix`, at names.dims' widths. + fn mlp(L: *Loader, prefix: []const u8, gate: Role) !Mlp { + const nk = names.dims(L.c, gate).linear; + return L.r.mlp(prefix, nk[0], nk[1]); } fn kda(L: *Loader, i: u32) !Kda { - const c = L.c; - const n: usize = c.kda_heads * c.kda_dim; - const d: usize = c.hidden; + const n = names.dims(L.c, .kda_q_conv).conv[0]; var b: [64]u8 = undefined; var cb: [96]u8 = undefined; return .{ - .q = try L.lin(i, .kda_q, n, d), - .k = try L.lin(i, .kda_k, n, d), - .v = try L.lin(i, .kda_v, n, d), - .b = try L.lin(i, .kda_b, c.kda_heads, d), - .f = try L.lin(i, .kda_f, n, d), - .g = try L.lin(i, .kda_g, n, d), - .o = try L.lin(i, .kda_o, d, n), + .q = try L.lin(i, .kda_q), + .k = try L.lin(i, .kda_k), + .v = try L.lin(i, .kda_v), + .b = try L.lin(i, .kda_b), + .f = try L.lin(i, .kda_f), + .g = try L.lin(i, .kda_g), + .o = try L.lin(i, .kda_o), .conv = try L.r.convs(try std.fmt.bufPrint(&b, "model.layers.{d}.attention.", .{i}), n, try std.fmt.bufPrint(&cb, "model.layers.{d}.attention.conv", .{i})), - .a_log = try L.small(i, .kda_a_log, c.kda_heads, .f32), - .dt_bias = try L.small(i, .kda_dt_bias, n, .f32), - .o_norm = try L.small(i, .kda_o_norm, c.kda_dim, .bf16), + .a_log = try L.small(i, .kda_a_log, .f32), + .dt_bias = try L.small(i, .kda_dt_bias, .f32), + .o_norm = try L.small(i, .kda_o_norm, .bf16), }; } fn mla(L: *Loader, i: u32) !Mla { const c = L.c; - const d: usize = c.hidden; - const h: usize = c.mla_heads; var m: Mla = .{ - .kv_a = try L.lin(i, .mla_kv_a, c.kv_lora + c.rope_dim, d), - .kv_a_norm = try L.small(i, .mla_kv_a_norm, c.kv_lora, .bf16), - .kv_b = try L.lin(i, .mla_kv_b, h * (c.nope + c.v_dim), c.kv_lora), - .o = try L.lin(i, .mla_o, d, h * c.v_dim), + .kv_a = try L.lin(i, .mla_kv_a), + .kv_a_norm = try L.small(i, .mla_kv_a_norm, .bf16), + .kv_b = try L.lin(i, .mla_kv_b), + .o = try L.lin(i, .mla_o), }; if (c.q_lora > 0) { - m.q_a = try L.lin(i, .mla_q_a, c.q_lora, d); - m.q_a_norm = try L.small(i, .mla_q_a_norm, c.q_lora, .bf16); - m.q_b = try L.lin(i, .mla_q_b, h * c.qkDim(), c.q_lora); - } else m.q = try L.lin(i, .mla_q, h * c.qkDim(), d); - if (c.head_gate) m.gate = try L.lin(i, .mla_gate, h, d); + m.q_a = try L.lin(i, .mla_q_a); + m.q_a_norm = try L.small(i, .mla_q_a_norm, .bf16); + m.q_b = try L.lin(i, .mla_q_b); + } else m.q = try L.lin(i, .mla_q); + if (c.head_gate) m.gate = try L.lin(i, .mla_gate); return m; } fn moe(L: *Loader, i: u32) !Moe { const c = L.c; - const d: usize = c.hidden; const route = try L.module(i, .router); if (c.format(route) != .bf16) { std.log.err("{s} is {t}: the routing kernels read a bf16 router", .{ route, c.format(route) }); @@ -161,11 +164,12 @@ const Loader = struct { const count = @min(c.experts, L.opts.experts orelse c.experts); const lim = @import("cuda.zig").limits(c, i); var b: [96]u8 = undefined; + const x = names.dims(c, .expert_gate).linear; return .{ - .router = try L.small(i, .router, c.experts * d, .bf16), - .bias = if (c.expert_bias) try L.small(i, .router_bias, c.experts, .f32) else 0, - .routed = try L.r.experts(i, count, c.moe_inter, d, lim.routed), - .shared = if (c.shared_inter > 0) try L.r.mlp(try std.fmt.bufPrint(&b, "model.layers.{d}.mlp.shared_experts.", .{i}), c.shared_inter, d) else null, + .router = try L.small(i, .router, .bf16), + .bias = if (c.expert_bias) try L.small(i, .router_bias, .f32) else 0, + .routed = try L.r.experts(i, count, x[0], x[1], lim.routed), + .shared = if (c.shared_inter > 0) try L.mlp(try std.fmt.bufPrint(&b, "model.layers.{d}.mlp.shared_experts.", .{i}), .shared_gate) else null, .experts = count, .shared_limit = lim.shared, }; @@ -176,23 +180,22 @@ const Loader = struct { var b: [64]u8 = undefined; return .{ .index = i, - .attn_norm = try L.small(i, .attn_norm, c.hidden, .bf16), - .mlp_norm = try L.small(i, .mlp_norm, c.hidden, .bf16), + .attn_norm = try L.small(i, .attn_norm, .bf16), + .mlp_norm = try L.small(i, .mlp_norm, .bf16), .mixer = switch (c.kind(i)) { .kda => .{ .kda = try L.kda(i) }, .mla => .{ .mla = try L.mla(i) }, }, - .ffn = if (c.isMoe(i)) .{ .moe = try L.moe(i) } else .{ .dense = try L.r.mlp(try std.fmt.bufPrint(&b, "model.layers.{d}.mlp.", .{i}), c.dense_inter, c.hidden) }, + .ffn = if (c.isMoe(i)) .{ .moe = try L.moe(i) } else .{ .dense = try L.mlp(try std.fmt.bufPrint(&b, "model.layers.{d}.mlp.", .{i}), .dense_gate) }, }; } fn mtp(L: *Loader, i: u32) !Mtp { - const d: usize = L.c.hidden; return .{ - .enorm = try L.small(i, .mtp_enorm, d, .bf16), - .hnorm = try L.small(i, .mtp_hnorm, d, .bf16), - .eh = try L.lin(i, .mtp_eh, d, 2 * d), - .norm = try L.small(i, .mtp_final_norm, d, .bf16), + .enorm = try L.small(i, .mtp_enorm, .bf16), + .hnorm = try L.small(i, .mtp_hnorm, .bf16), + .eh = try L.lin(i, .mtp_eh), + .norm = try L.small(i, .mtp_final_norm, .bf16), }; } }; @@ -213,13 +216,15 @@ pub fn load(gpa: std.mem.Allocator, d: *const cuda.Driver, stream: cuda.Stream, w.layers = out; if (!opts.globals) return w; const e = try ck.get("model.word_embeddings.weight"); - if (!e.is(.bf16, &.{ c.vocab, c.hidden })) { - std.log.err("model.word_embeddings: {t} {any}, expected bf16 [{d}, {d}]", .{ e.dtype, e.shape[0..e.rank], c.vocab, c.hidden }); + const vd = names.dims(&w.config, .embed).linear; + if (!e.is(.bf16, &vd)) { + std.log.err("model.word_embeddings: {t} {any}, expected bf16 {any}", .{ e.dtype, e.shape[0..e.rank], vd }); return error.UnexpectedTensor; } w.embed = try w.store.keep(e.bytes, "model.word_embeddings.weight"); - w.norm = try L.r.small("model.norm.weight", c.hidden, .bf16); - if (!c.tie) w.head = try L.r.linear("lm_head", c.vocab, c.hidden); + w.norm = try L.r.small("model.norm.weight", names.numel(&w.config, .final_norm), .bf16); + const hd = names.dims(&w.config, .lm_head).linear; + if (!c.tie) w.head = try L.r.linear("lm_head", hd[0], hd[1]); return w; } diff --git a/zig/tests/cuda/ling_ckpt.zig b/zig/tests/cuda/ling_ckpt.zig new file mode 100644 index 0000000000..c52b368d89 --- /dev/null +++ b/zig/tests/cuda/ling_ckpt.zig @@ -0,0 +1,95 @@ +//! `ling-ckpt-check DIR [INDEX]`: a Ling checkpoint against the loader's names, dtypes and shapes, from headers only. +const std = @import("std"); +const core = @import("core"); +const check = @import("check.zig"); +const ling = @import("ling_cuda"); + +const names = ling.names; +const Config = ling.config.Config; +const Tensor = core.checkpoint.Tensor; + +/// A tensor as the loader takes it: names.dims' shape, its format's dtype; side tensors as weights_proj reads them. +fn fits(c: *const Config, name: []const u8, t: names.Tensor, x: Tensor) bool { + const tail = name[std.mem.lastIndexOfScalar(u8, name, '.').? + 1 ..]; + const small = x.dtype == .bf16 or x.dtype == .f32; // Reader.small and Reader.convs take either + switch (names.dims(c, t.role)) { + .values => |n| return small and x.is(x.dtype, &.{n}), + .conv => |s| return small and x.is(x.dtype, &s), + .linear => |nk| { + const n = nk[0]; + const k = nk[1]; + if (t.role == .embed) return x.is(.bf16, &nk); + if (t.role == .router) return c.format(name[0 .. name.len - tail.len - 1]) == .bf16 and small and x.is(x.dtype, &nk); + const f = c.format(name[0 .. name.len - tail.len - 1]); + if (!t.scale) return switch (f) { + .bf16 => x.is(.bf16, &nk), + .fp8_block128 => x.is(.f8_e4m3, &nk) and k % 128 == 0, + .nvfp4 => x.is(.u8, &.{ n, k / 2 }) and k % 16 == 0, + }; + return switch (f) { + .bf16 => false, + .fp8_block128 => x.is(.f32, &.{ (n + 127) / 128, k / 128 }), + .nvfp4 => if (std.mem.eql(u8, tail, "weight_scale")) x.is(.f8_e4m3, &.{ n, k / 16 }) else x.dtype == .f32 and x.numel() == 1, + }; + }, + } +} + +const Tally = struct { + ck: *core.Checkpoint, + c: *const Config, + expected: usize = 0, + mtp: usize = 0, + missing: usize = 0, + wrong: usize = 0, + bytes: u64 = 0, + kinds: [3]usize = @splat(0), // weights, scale tensors, small tensors and convs + + pub fn put(k: *Tally, name: []const u8, t: names.Tensor) !void { + k.expected += 1; + if (t.layer != null and t.layer.? >= k.c.layers) k.mtp += 1; + if (!k.ck.has(name)) { + k.missing += 1; + if (k.missing <= 10) std.debug.print(" missing {s}\n", .{name}); + return; + } + const x = try k.ck.get(name); + k.bytes += x.bytes.len; + k.kinds[if (names.dims(k.c, t.role) != .linear) 2 else if (t.scale) 1 else 0] += 1; + if (fits(k.c, name, t, x)) return; + k.wrong += 1; + if (k.wrong <= 10) std.debug.print(" {s}: {t} {any}, not what the loader takes\n", .{ name, x.dtype, x.shape[0..x.rank] }); + } +}; + +pub fn run(gpu: check.Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 1) return error.MissingArgument; + const dir = rest[0]; + var why: ling.config.Why = .{}; + var opened = ling.weights.open(gpu.gpa, gpu.io, dir, &why) catch |err| { + std.debug.print("FAIL {s}: {s}\n", .{ dir, why.text() }); + return err; + }; + defer opened.ck.close(); + const c = &opened.c; + var stored: usize = 0; + for (opened.ck.files.items) |f| stored += f.names.count(); + var k: Tally = .{ .ck = &opened.ck, .c = c }; + try names.all(c, &k); + const unknown = opened.ck.unused(); + // the index's total_size, or without an index the files' data regions: every byte belongs to a checked tensor + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const index = if (rest.len > 1) rest[1] else try std.fs.path.join(arena.allocator(), &.{ dir, "model.safetensors.index.json" }); + var total: u64 = 0; + if (std.Io.Dir.cwd().readFileAlloc(gpu.io, index, arena.allocator(), .limited(1 << 26))) |text| { + const parsed = try std.json.parseFromSliceLeaky(std.json.Value, arena.allocator(), text, .{}); + total = @intCast(parsed.object.get("metadata").?.object.get("total_size").?.integer); + } else |_| for (opened.ck.files.items) |f| { + total += f.map.memory.len - f.data; + } + std.debug.print("RESULT ling-ckpt-check {t}: {d} shards, {d} tensors stored, {d} expected ({d} of the MTP layer): {d} weights, {d} scale tensors, {d} small tensors and convs\n", .{ c.quant, opened.ck.files.items.len, stored, k.expected, k.mtp, k.kinds[0], k.kinds[1], k.kinds[2] }); + std.debug.print("RESULT {d} missing, {d} stored but not expected, {d} of another dtype or shape; {d} bytes against {d} (the index's total_size, else the files' data regions)\n", .{ k.missing, unknown, k.wrong, k.bytes, total }); + try check.expect(k.missing == 0 and unknown == 0 and k.wrong == 0 and stored == k.expected and k.bytes == total, "every stored tensor is one the loader reads, at its dtype and shape, and nothing is missing", .{}); + check.pass("ling-ckpt-check: all {d} tensors' names, dtypes and shapes as the loader takes them, {d} bytes (headers only)", .{ stored, k.bytes }); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 5cb0aaa1e0..811f130d3c 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -117,6 +117,7 @@ const usage = \\ ling-clamp-old the grouped expert instances vs an older build's fatbins, byte for byte \\ ling-lanes [ref] [--layers a,b] [--mtp] Ling's lane backend: chunked, drafted, resumed, concurrent == plain \\ ling-slots the lane backend's slots over reused map buckets, in this build's mode (no GPU work) + \\ ling-ckpt-check [index] every tensor's name, dtype and shape against the loader's, from the headers alone (no GPU work) \\ ling-init-faults [--mtp] [--every N] the backend opened with each driver call, then each host allocation, failing \\ ling-mtp-real the real MTP layer vs HF's (oracle/ling_mtp.py); cut, window == decode rows \\ ling-weights-real real Ling-3.1 layers 0, 4, 7 (4 experts) vs float64 and torch (oracle/ling_weights.py) @@ -236,6 +237,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "ling-mtp-real")) return @import("ling_mtp_real.zig").run(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "ling-lanes")) return @import("ling_lanes.zig").run(gpu, rest); if (std.mem.eql(u8, cmd, "ling-slots")) return lingSlots(gpu); + if (std.mem.eql(u8, cmd, "ling-ckpt-check")) return @import("ling_ckpt.zig").run(gpu, rest); if (std.mem.eql(u8, cmd, "ling-init-faults")) return @import("ling_faults.zig").run(gpu, rest); if (std.mem.eql(u8, cmd, "ling-weights-real")) return ling_real.run(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "ling-pack")) return @import("ling_pack.zig").run(gpu, rest); From e4c9bcc5088d6b8be59781a6874925be80ac243c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 16 Oct 2026 08:08:32 +0200 Subject: [PATCH 31/34] oracle: the Ling oracles check their HF source against inclusionAI/Ling-3.1-flash@a6943f4b by sha256 ling_hf.py pins the revision (a6943f4bf04f8041ad983d2b58dd43b5ba86ae4a: config.json as of #2; the modeling, configuration and index files are unchanged since the upload) and the sha256 of modeling_bailing_moe_v3.py, configuration_bailing_moe_v3.py, config.json and model.safetensors.index.json. ling_mtp.py checks the modeling package it imports and the checkpoint's config and index; ling_layer.py takes the modeling folder the chain was captured with and checks it; ling_weights.py checks the index. Any other copy stops the oracle naming both hashes. - test-ling holds fixtures/config.json to the pinned sha256. - The local copies match HF at that revision byte for byte; ling_layer.py rebuilds the layer-4 fixture identically. --- docs/recipes/ling-3.1.md | 7 +++++-- zig/src/families/ling/config_test.zig | 6 ++++++ zig/tests/cuda/oracle/ling_hf.py | 24 ++++++++++++++++++++++++ zig/tests/cuda/oracle/ling_layer.py | 4 ++++ zig/tests/cuda/oracle/ling_mtp.py | 4 ++++ zig/tests/cuda/oracle/ling_weights.py | 3 +++ 6 files changed, 46 insertions(+), 2 deletions(-) create mode 100644 zig/tests/cuda/oracle/ling_hf.py diff --git a/docs/recipes/ling-3.1.md b/docs/recipes/ling-3.1.md index 1f374fb2b8..a94fff9aef 100644 --- a/docs/recipes/ling-3.1.md +++ b/docs/recipes/ling-3.1.md @@ -282,8 +282,11 @@ format's side tensors), none missing, and their bytes add up to the index's `tot network mount it takes seconds. `ling-mtp-real DIR FIXTURE` runs the MTP layer (fixture from `oracle/ling_mtp.py SRC MODELING_PKG ROWS.pt OUT`). -`ling-layer-real DIR FIXTURE` runs one whole layer (fixture from `oracle/ling_layer.py CHAIN L OUT`, the captured -HF chain's `h.pt`, `h.pt`, `r.pt`). `ling-weights-real` loads layers 0, 4 and 7 with experts 0-3 (769 MiB on the device) and compares 35 products with torch's float64 +`ling-layer-real DIR FIXTURE` runs one whole layer (fixture from `oracle/ling_layer.py CHAIN L OUT MODELING`, the captured +HF chain's `h.pt`, `h.pt`, `r.pt`). The HF reference is pinned: `oracle/ling_hf.py` names +`inclusionAI/Ling-3.1-flash` at `a6943f4bf04f8041ad983d2b58dd43b5ba86ae4a` and the sha256 of its +`modeling_bailing_moe_v3.py`, `configuration_bailing_moe_v3.py`, `config.json` and index; `ling_mtp.py`, +`ling_layer.py` and `ling_weights.py` stop on any other copy, and `fixtures/config.json` is held to it in `test-ling`. `ling-weights-real` loads layers 0, 4 and 7 with experts 0-3 (769 MiB on the device) and compares 35 products with torch's float64 on the same 8 rows: worst 0.414 of the bound, relative L2 at most 1.7e-3 for bf16 outputs and 6.9e-8 for the experts' fp32 down rows. diff --git a/zig/src/families/ling/config_test.zig b/zig/src/families/ling/config_test.zig index f8a08e87c7..880d8ddd03 100644 --- a/zig/src/families/ling/config_test.zig +++ b/zig/src/families/ling/config_test.zig @@ -306,3 +306,9 @@ test "routing groups: groups of one and a single pick refused (HF's rule differs try t.expect(c.experts / c.n_group >= 2 and c.topk_group <= c.n_group); } } + +test "fixtures/config.json is inclusionAI/Ling-3.1-flash's at a6943f4b, byte for byte (oracle/ling_hf.py pins the same)" { + var h: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(real, &h, .{}); + try t.expectEqualStrings("6f3773e52a12509e772e698aacacaa156d2402ebb8250812a1da52a03d92f084", &std.fmt.bytesToHex(h, .lower)); +} diff --git a/zig/tests/cuda/oracle/ling_hf.py b/zig/tests/cuda/oracle/ling_hf.py new file mode 100644 index 0000000000..a256eff52d --- /dev/null +++ b/zig/tests/cuda/oracle/ling_hf.py @@ -0,0 +1,24 @@ +"""The HF reference of the Ling oracles: inclusionAI/Ling-3.1-flash at one revision, its files by sha256.""" +import hashlib +import os + +REPO = "inclusionAI/Ling-3.1-flash" +# config.json is this revision's (#2); the modeling, configuration and index files are unchanged since the upload +REVISION = "a6943f4bf04f8041ad983d2b58dd43b5ba86ae4a" +SHA256 = { + "modeling_bailing_moe_v3.py": "c2509bf7ac580c262e2581d34d6403aa21682d2e10beb9ad85ad8820a7e33a40", + "configuration_bailing_moe_v3.py": "f2c048966aec8a2f042cfeb1351f74d51a28589b409c55baae7d24e841c1f6c4", + "config.json": "6f3773e52a12509e772e698aacacaa156d2402ebb8250812a1da52a03d92f084", + "model.safetensors.index.json": "f01006f7a65eb4b32b66ba60868ef7cc1642c66afe658cb7e6faf80eb246b468", +} +MODELING = ("modeling_bailing_moe_v3.py", "configuration_bailing_moe_v3.py") + + +def pinned(folder: str, files=MODELING) -> None: + """Stops unless each of `files` in `folder` is the pinned revision's, byte for byte.""" + for f in files: + path = os.path.join(folder, f) + with open(path, "rb") as fh: + got = hashlib.sha256(fh.read()).hexdigest() + if got != SHA256[f]: + raise SystemExit(f"{path}: sha256 {got}, not {REPO}@{REVISION[:12]}'s {SHA256[f]}") diff --git a/zig/tests/cuda/oracle/ling_layer.py b/zig/tests/cuda/oracle/ling_layer.py index 1b6f15915b..5bd9cf19ea 100644 --- a/zig/tests/cuda/oracle/ling_layer.py +++ b/zig/tests/cuda/oracle/ling_layer.py @@ -4,13 +4,17 @@ import torch +import ling_hf + def main() -> None: p = argparse.ArgumentParser(description=__doc__) p.add_argument("chain", help="the captured chain: h.pt (layer L's input), h.pt (its output), r.pt") p.add_argument("layer", type=int) p.add_argument("out", help="the fixture: x bf16 [T][D], HF's output fp32 [T][D], HF's top-k i32 [T][K] (K 0: dense)") + p.add_argument("modeling", help="the folder of the HF modeling and configuration files the chain was captured with") a = p.parse_args() + ling_hf.pinned(a.modeling) x = torch.load(f"{a.chain}/h{a.layer}.pt")[0].to(torch.bfloat16) y = torch.load(f"{a.chain}/h{a.layer + 1}.pt")[0].float() r = torch.load(f"{a.chain}/r{a.layer}.pt") diff --git a/zig/tests/cuda/oracle/ling_mtp.py b/zig/tests/cuda/oracle/ling_mtp.py index dcc0a5f1c9..82676f82a7 100644 --- a/zig/tests/cuda/oracle/ling_mtp.py +++ b/zig/tests/cuda/oracle/ling_mtp.py @@ -9,6 +9,8 @@ import torch from safetensors import safe_open +import ling_hf + def main() -> None: p = argparse.ArgumentParser(description=__doc__) @@ -17,6 +19,8 @@ def main() -> None: p.add_argument("rows", help="a .pt with 'ids' [1][T] and 'out' [1][T][D] (a layer's output, e.g. ref0.pt)") p.add_argument("out", help="the fixture: embeddings, hidden rows (bf16), HF's output (fp32), HF's top-k (i32)") a = p.parse_args() + ling_hf.pinned(a.modeling) + ling_hf.pinned(a.src, ("config.json", "model.safetensors.index.json")) pkg = os.path.basename(os.path.normpath(a.modeling)) sys.path.insert(0, os.path.dirname(os.path.normpath(a.modeling))) import transformers.utils.import_utils as iu diff --git a/zig/tests/cuda/oracle/ling_weights.py b/zig/tests/cuda/oracle/ling_weights.py index 9ce13cad9b..de7d842abe 100644 --- a/zig/tests/cuda/oracle/ling_weights.py +++ b/zig/tests/cuda/oracle/ling_weights.py @@ -6,6 +6,8 @@ import torch from safetensors import safe_open +import ling_hf + ROWS = 8 # decode rows, and the down projection's pair rows (2 tokens x 4 experts) EXPERTS = 4 KDA = ("q_proj", "k_proj", "v_proj", "b_proj", "f_proj", "g_proj", "o_proj") @@ -34,6 +36,7 @@ def main() -> None: p.add_argument("src", help="the FP8 checkpoint (model.safetensors.index.json and its shards)") p.add_argument("out", help="where x.bin, .y, .abs and manifest.txt go") a = p.parse_args() + ling_hf.pinned(a.src, ("model.safetensors.index.json",)) os.makedirs(a.out, exist_ok=True) weight_map = json.load(open(os.path.join(a.src, "model.safetensors.index.json")))["weight_map"] files, xs, lines = {}, {}, [] From 5e931c3a70d65a43e0d8fb445173d45be20338d0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 16 Oct 2026 11:31:51 +0200 Subject: [PATCH 32/34] ling: a segment the scratch cannot hold is refused before any layer runs, not at the first MLA layer Forward.rowsOf checked the sequence's capacity only; the scratch's attention partials (Buffers.nch, sized from the capacity the Buffers were made for) were checked inside block_mla, after the KDA layers before it had committed and after the segment's latent row was written. With scratch smaller than the state a prompt chunk or window that fit the state still changed it before CacheFull came back. block_mla.room now checks both; rowsOf calls it for every segment when the graph has MLA layers, and block_mla calls it for every segment before its first cache write (the MTP layer, which runs the block without rowsOf). - ling-lanes (engine checks): Buffers made for 512 positions against sequences of 1024, twins at 500: a 30-row prompt chunk, a 16-row window and (with --mtp) 16 MTP rows each come back CacheFull with the arenas byte for byte as they were; the next window equals the twin's (tokens, logits, states) and a new request on that slot equals one on a fresh slot. With the preflight reverted the arenas differ and the next commit fails (KdaLayersOutOfStep); with only block_mla's order reverted the MTP case alone fails. --- zig/src/families/ling/block_mla.zig | 8 +++- zig/src/families/ling/forward.zig | 3 +- zig/tests/cuda/ling_lanes_bytes.zig | 57 ++++++++++++++++++++++++----- 3 files changed, 56 insertions(+), 12 deletions(-) diff --git a/zig/src/families/ling/block_mla.zig b/zig/src/families/ling/block_mla.zig index bf76aaf2f1..02a5e5beb7 100644 --- a/zig/src/families/ling/block_mla.zig +++ b/zig/src/families/ling/block_mla.zig @@ -12,6 +12,11 @@ const Config = @import("config.zig").Config; const kw = state.cache_width; const lw = 512; +/// Room for a segment's rows: the sequence's cache and the scratch's attention partials (512-key chunks). +pub fn room(b: *const state.Buffers, st: *const state.State, rows: usize) !void { + if (st.pos + rows > st.capacity or latent.chunksFor(st.pos + rows) > b.nch) return error.CacheFull; +} + /// The block on b.normed into b.y: q (lora pair or q_proj), latent and roped key into the cache, attention, gate. pub fn block(o: Ops, b: *const state.Buffers, a: weights.Mla, ab: graph.Absorb, g: *const graph.Graph, mi: usize, segs: []const block_kda.Seg, rows: usize) !void { const c = &g.w.config; @@ -32,10 +37,10 @@ pub fn block(o: Ops, b: *const state.Buffers, a: weights.Mla, ab: graph.Absorb, if (a.gate) |gp| try o.mm(gp, b.normed, D, rows, b.g); try latent.absorbQStrided(o.set(), o.s, b.q, qk, ab.wk, b.qa, kw, rows, H, nope, lw); const scale = c.smScale(); + for (segs) |sg| try room(b, sg.st, sg.rows); // every segment before the first cache write for (segs) |sg| { const st = sg.st; const r = sg.row0; - if (st.pos + sg.rows > st.capacity) return error.CacheFull; try latent.rope(o.set(), o.s, .{ .x = b.q + (r * H * qk + nope) * 2, .x_row = H * qk, .x_head = qk, .out = b.qa + (r * H * kw + lw) * 2, .o_row = H * kw, .o_head = kw }, sg.rows, H, st.pos_dev, g.freq); try latent.rope(o.set(), o.s, .{ .x = b.kv + (r * kw + lw) * 2, .x_row = kw, .x_head = 0, .out = b.lat + (r * kw + lw) * 2, .o_row = kw, .o_head = 0 }, sg.rows, 1, st.pos_dev, g.freq); try latent.latentWrite(o.set(), o.s, b.lat + r * kw * 2, kw, st.cacheAt(mi), st.pos_dev, sg.rows, kw); @@ -48,7 +53,6 @@ pub fn block(o: Ops, b: *const state.Buffers, a: weights.Mla, ab: graph.Absorb, at = b.at; } const nch = latent.chunksFor(st.pos + r0 + m); - if (nch > b.nch) return error.CacheFull; const row = r + r0; try latent.attentionRoped(o.set(), o.s, b.qa + row * H * kw * 2, st.cacheAt(mi), at, b.partials, b.ol + row * H * lw * 2, b.dummy, m, H, lw, nch, scale); } diff --git a/zig/src/families/ling/forward.zig b/zig/src/families/ling/forward.zig index f31edfb54d..61bc7befc5 100644 --- a/zig/src/families/ling/forward.zig +++ b/zig/src/families/ling/forward.zig @@ -29,12 +29,13 @@ pub const Forward = struct { return .{ .m = f.m, .s = f.s, .prompt = f.b.prompt, .part = f.b.part, .kda_wide = f.kda_wide }; } - /// Segments in row order from row 0, each on its own committed sequence with room for its rows; windows <= 16 rows. + /// Segments in row order from row 0, each on its own committed sequence with room (and scratch) for its rows. fn rowsOf(f: *const Forward, segs: []const Seg) !usize { var rows: usize = 0; for (segs, 0..) |sg, i| { if (sg.row0 != rows or sg.rows == 0 or sg.st.pending != 0) return error.BadSegments; if (sg.st.pos + sg.rows > sg.st.capacity) return error.CacheFull; // before a prompt's KDA layers commit + if (f.g.mla_layers > 0) try block_mla.room(f.b, sg.st, sg.rows); // and the scratch's attention partials if (!f.b.prompt and sg.rows > window_rows) return error.WindowTooWide; for (segs[0..i]) |o| if (o.st == sg.st) return error.BadSegments; rows += sg.rows; diff --git a/zig/tests/cuda/ling_lanes_bytes.zig b/zig/tests/cuda/ling_lanes_bytes.zig index 922c7eb231..4716e3c60f 100644 --- a/zig/tests/cuda/ling_lanes_bytes.zig +++ b/zig/tests/cuda/ling_lanes_bytes.zig @@ -87,6 +87,7 @@ pub fn engineChecks(c: *Cover, e: *Engine, prompt: []const u32, w: []const u32) try sharedChecks(c, e, prompt, w); if (e.mtp != null) try mtpChain(c, e, prompt, w); try capacityRefusal(c, e, prompt, w); + try scratchRefusal(c, e, prompt, w); for ([_]u32{ 10, 11, 12, 13, 14 }) |s| e.release(s); } @@ -139,23 +140,61 @@ fn capacityRefusal(c: *Cover, e: *Engine, prompt: []const u32, w: []const u32) ! const same = std.mem.eql(u8, before[0], after[0]) and std.mem.eql(u8, before[1], after[1]) and std.mem.eql(u1, cur, q.st.cur) and q.st.pos == pos and q.st.pending == 0; try c.expect(refused == asked and same, "{d} of {d} requests past the capacity ({d} at {d}) refused with CacheFull, the arenas ({d} and {d} bytes) and position as they were", .{ refused, asked, cap, pos, before[0].len, before[1].len }); - // the next window on both twins: tokens, logits, then states + try afterRefusal(c, e, prompt, w, pend, 50); +} + +/// The next window on twins `s` (refused) and s + 1, then a new request on `s` and on a fresh slot s + 2. +fn afterRefusal(c: *Cover, e: *Engine, prompt: []const u32, w: []const u32, pend: [2]u32, s: u32) !void { var got: [2][4]u32 = undefined; var lo: [2][]u8 = undefined; for (0..2) |k| { - try e.bind(@intCast(50 + k)); + try e.bind(s + @as(u32, @intCast(k))); var win: [4]u32 = .{ pend[k], w[40], w[41], w[42] }; try e.verify(&win, &got[k]); - lo[k] = try logitsRows(e, a, e.bd.logits, 4); + lo[k] = try logitsRows(e, c.a, e.bd.logits, 4); try e.keep(4); } try c.expect(std.mem.eql(u32, &got[0], &got[1]) and std.mem.eql(u8, lo[0], lo[1]), "the window after the refusals: tokens and logits == the twin's", .{}); - try rig.sameSeq(c, e, "after the refusals and a window, the twin", 50, 51); - const t50 = try pass(e, 50, prompt, 37, null, null); - const t52 = try pass(e, 52, prompt, 37, null, null); - try c.expect(t50 == t52, "a new request on the refused slot: first token {d}, on a fresh slot {d}", .{ t50, t52 }); - try rig.sameSeq(c, e, "a new request on the refused slot == on a fresh slot", 50, 52); - for ([_]u32{ 50, 51, 52 }) |s| e.release(s); + try rig.sameSeq(c, e, "after the refusals and a window, the twin", s, s + 1); + const t0 = try pass(e, s, prompt, 37, null, null); + const t2 = try pass(e, s + 2, prompt, 37, null, null); + try c.expect(t0 == t2, "a new request on the refused slot: first token {d}, on a fresh slot {d}", .{ t0, t2 }); + try rig.sameSeq(c, e, "a new request on the refused slot == on a fresh slot", s, s + 2); + for ([_]u32{ s, s + 1, s + 2 }) |x| e.release(x); +} + +/// Scratch for half the sequences' capacity: a prompt chunk, a window and MTP rows past it refused, nothing changed. +fn scratchRefusal(c: *Cover, e: *Engine, prompt: []const u32, w: []const u32) !void { + const a = c.a; + const half = e.capacity / 2; + const long = try a.alloc(u32, half + 40); + for (long, 0..) |*x, i| x.* = prompt[i % prompt.len]; + var pend: [2]u32 = undefined; + for (0..2) |k| pend[k] = try pass(e, @intCast(60 + k), long[0 .. half - 12], 512, null, null); // twins at half - 12 + var bp = try F.Buffers.init(e.d, &e.m, &e.w, 64, true, half); + defer bp.deinit(); + var bd = try F.Buffers.init(e.d, &e.m, &e.w, F.window_rows, false, half); + defer bd.deinit(); + try e.bind(60); + const q = e.seqs.items[60].?; + const before = try arenas(c, e, 60); + const cur = try a.dupe(u1, q.st.cur); + const pos = q.st.pos; + var refused: usize = 0; + const seg = [_]F.Seg{.{ .st = &q.st, .row0 = 0, .rows = 30 }}; // fits the sequence, not the scratch + if (e.forward(&bp).run(&seg, @ptrCast(long[pos..][0..30]))) |_| {} else |err| refused += @intFromBool(err == error.CacheFull); + const win = [_]F.Seg{.{ .st = &q.st, .row0 = 0, .rows = 16 }}; + if (e.forward(&bd).run(&win, @ptrCast(w[20..36]))) |_| {} else |err| refused += @intFromBool(err == error.CacheFull); + var asked: usize = 2; + if (e.mtp) |*h| { // 16 rows through the MTP layer from its position, past the scratch's + asked += 1; + if (h.absorb(&e.m, e.stream, &bd, &q.hd.?, @ptrCast(long[0..16]), e.bd.normed)) |_| {} else |err| refused += @intFromBool(err == error.CacheFull); + } + const after = try arenas(c, e, 60); + const same = std.mem.eql(u8, before[0], after[0]) and std.mem.eql(u8, before[1], after[1]) and + std.mem.eql(u1, cur, q.st.cur) and q.st.pos == pos and q.st.pending == 0 and (q.hd == null or q.hd.?.pending == 0); + try c.expect(refused == asked and same, "{d} of {d} requests past the scratch ({d} of {d}, at {d}) refused with CacheFull, the arenas and position as they were", .{ refused, asked, half, e.capacity, pos }); + try afterRefusal(c, e, prompt, w, pend, 60); } /// MTP chains of 3 and of 1 on twins: the same first draft, and after the same windows the chain's rows rolled back. From 1395e2678a516679c7681979f35e4abc3cfebbe1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 16 Oct 2026 11:31:59 +0200 Subject: [PATCH 33/34] ling: staging and loaded weights freed only after the stream drains; a failed drain leaks them on purpose A failed packer launch, upload or copy returned through defers that freed what queued work could still read or write: the expert stage (pinned host and device, weights_proj Stage.deinit), the lane staging of bf16, FP8 and NVFP4 projections, the store of packed outputs (weights.load), absorbOf's identity rows and results, the engine's weights and graph when a later init step failed, and keepState's copy target. The Reader now tracks whether an upload or launch is queued (settled() drains it first); the other paths drain before freeing. A drain that fails leaves the memory allocated: Store.drop, Weights.drop and Graph.drop free the host lists and leak the device buffers. - fake_driver.zig: async copies, memsets and launches stay pending until a successful synchronize; frees with work pending are counted, failures can persist from call N on (drains included), and leaked memory can be swept. - ling-init-faults: every driver call and host allocation of the open and a kept state fails in turn with transfers pending, then every call from N on; no free may come while work is pending, and memory may stay held only after a failed drain. Over bf16, FP8 and NVFP4 with and without MTP: 15,180 points (1,804 failing with a transfer pending), all clean. Each fix reverted alone fails it (Stage.deinit, lane staging, NVFP4 lane staging, load's store, absorbOf, the engine's weights, keepState). --- zig/src/families/ling/engine.zig | 8 +-- zig/src/families/ling/fake_driver.zig | 83 +++++++++++++++++++++++--- zig/src/families/ling/graph.zig | 23 +++++-- zig/src/families/ling/weights.zig | 10 +++- zig/src/families/ling/weights_proj.zig | 48 ++++++++++++--- zig/tests/cuda/ling_faults.zig | 71 ++++++++++++++++------ 6 files changed, 197 insertions(+), 46 deletions(-) diff --git a/zig/src/families/ling/engine.zig b/zig/src/families/ling/engine.zig index 3e64f0554c..4d3b4ae1da 100644 --- a/zig/src/families/ling/engine.zig +++ b/zig/src/families/ling/engine.zig @@ -103,11 +103,11 @@ pub const Engine = struct { defer pack.unload(); e.w = try weights.load(gpa, d, e.stream, &pack, &opened.ck, c, all, .{}); } - errdefer e.w.deinit(); + errdefer e.w.drop(weights.drains(e.stream)); // leaked on a failed drain: a queued launch may read them e.main = e.w; e.main.layers = e.w.layers[0..layers.len]; e.g = try fw.Graph.init(gpa, d, &e.m, e.stream, &e.main); - errdefer e.g.deinit(); + errdefer e.g.drop(weights.drains(e.stream)); // as e.w if (mtp) { e.mtp = @as(head.Head, undefined); e.mtp.?.init(gpa, d, &e.m, e.stream, &e.w, o.prompt_rows, o.capacity) catch |err| { @@ -381,10 +381,10 @@ pub const Engine = struct { const q = e.cur; if (q.st.pending != 0) return error.NotCommitted; var buf = try cuda.DeviceBuffer.alloc(e.d, q.st.arena.buf.len); - errdefer buf.free(); + errdefer if (weights.drains(e.stream)) buf.free(); // else leaked: the copy into it may be queued try e.copy(buf.ptr, q.st.arena.buf.ptr, buf.len); var hd: ?cuda.DeviceBuffer = null; - errdefer if (hd) |*x| x.free(); + errdefer if (hd) |*x| if (weights.drains(e.stream)) x.free(); if (q.hd) |h| { if (q.in_head != 0) return error.NotCommitted; hd = try cuda.DeviceBuffer.alloc(e.d, h.arena.buf.len); diff --git a/zig/src/families/ling/fake_driver.zig b/zig/src/families/ling/fake_driver.zig index f5656a899c..7a8efc8ad5 100644 --- a/zig/src/families/ling/fake_driver.zig +++ b/zig/src/families/ling/fake_driver.zig @@ -14,6 +14,12 @@ var calls: u64 = 0; // fallible calls since reset var fail_at: ?u64 = null; // the call that fails var code: R = 2; // what it returns (2: out of memory; NOT_READY reaches the caller without a log line) var bad: usize = 0; // frees of something not held, releases of nothing +var lazy = false; // async copies, memsets and launches stay pending until a stream or context drains +var sticky = false; // every call from the failing one on fails too, drains included +var pending: usize = 0; // async work queued since the last drain (lazy only) +var early: usize = 0; // frees while work was pending +var stuck: usize = 0; // drains that failed +var hit = false; // the failing call came with work pending var next: usize = 0x1000; var globals: [1 << 16]u8 align(256) = undefined; @@ -21,10 +27,46 @@ var globals: [1 << 16]u8 align(256) = undefined; pub fn reset(at: ?u64, with: R) void { calls = 0; bad = 0; + pending = 0; + early = 0; + stuck = 0; + hit = false; fail_at = at; code = with; } +/// From now on async work stays pending until a successful drain (`queue`); failures persist (`stick`). +pub fn mode(queue: bool, stick: bool) void { + lazy = queue; + sticky = stick; +} + +/// Frees of device or pinned memory while async work was pending, since reset; work pending now; failed drains. +pub fn freedEarly() usize { + return early; +} + +pub fn queued() usize { + return pending; +} + +pub fn failedDrains() usize { + return stuck; +} + +pub fn failedPending() bool { + return hit; +} + +/// Everything still held freed (what a failed drain leaked on purpose); returns how many. +pub fn sweep() usize { + const n = held.count(); + var it = held.iterator(); + while (it.next()) |kv| backing.free(@as([*]align(256) u8, @ptrFromInt(kv.key_ptr.*))[0..kv.value_ptr.*]); + held.clearRetainingCapacity(); + return n; +} + /// Fallible calls since reset, and whether the failure was handed out. pub fn count() u64 { return calls; @@ -59,7 +101,21 @@ pub fn context(d: *const cuda.Driver) cuda.Context { /// The countdown: each fallible call takes one step, and the chosen one fails. fn step() R { calls += 1; - return if (fail_at != null and calls == fail_at.?) code else ok; + const n = fail_at orelse return ok; + if (calls == n) hit = pending > 0; + return if (calls == n or (sticky and calls > n)) code else ok; +} + +/// An async call's result: queued work pending until a drain when lazy. +fn later(r: R) R { + if (r == ok and lazy) pending += 1; + return r; +} + +fn drain() R { + const r = step(); + if (r == ok) pending = 0 else stuck += 1; + return r; } /// Every entry point without out-parameters: a step, nothing else (the C ABI lets a caller pass arguments it ignores). @@ -88,6 +144,7 @@ fn take(n: usize) ?u64 { } fn give(p: u64) R { + if (pending > 0) early += 1; const kv = held.fetchRemove(p) orelse { bad += 1; return 1; @@ -196,18 +253,18 @@ const table = struct { return ok; } pub fn cuMemcpyHtoDAsync_v2(dst: abi.DevicePtr, src: ?*const anyopaque, n: usize, _: abi.Stream) callconv(.c) R { - return cuMemcpyHtoD_v2(dst, src, n); + return later(cuMemcpyHtoD_v2(dst, src, n)); } pub fn cuMemcpyDtoHAsync_v2(dst: ?*anyopaque, src: abi.DevicePtr, n: usize, _: abi.Stream) callconv(.c) R { - return cuMemcpyDtoH_v2(dst, src, n); + return later(cuMemcpyDtoH_v2(dst, src, n)); } pub fn cuMemcpyDtoDAsync_v2(dst: abi.DevicePtr, src: abi.DevicePtr, n: usize, _: abi.Stream) callconv(.c) R { - return cuMemcpyDtoD_v2(dst, src, n); + return later(cuMemcpyDtoD_v2(dst, src, n)); } pub fn cuMemcpy2DAsync_v2(m: *const abi.Memcpy2D, _: abi.Stream) callconv(.c) R { if (step() != ok) return code; for (0..m.height) |r| @memmove(bytes(m.dst + r * m.dst_pitch, m.width), bytes(m.src + r * m.src_pitch, m.width)); - return ok; + return later(ok); } pub fn cuMemsetD8_v2(p: abi.DevicePtr, v: u8, n: usize) callconv(.c) R { if (step() != ok) return code; @@ -220,10 +277,22 @@ const table = struct { return ok; } pub fn cuMemsetD8Async(p: abi.DevicePtr, v: u8, n: usize, _: abi.Stream) callconv(.c) R { - return cuMemsetD8_v2(p, v, n); + return later(cuMemsetD8_v2(p, v, n)); } pub fn cuMemsetD32Async(p: abi.DevicePtr, v: c_uint, n: usize, _: abi.Stream) callconv(.c) R { - return cuMemsetD32_v2(p, v, n); + return later(cuMemsetD32_v2(p, v, n)); + } + pub fn cuLaunchKernel(_: abi.Function, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: abi.Stream, _: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) R { + return later(step()); + } + pub fn cuLaunchKernelEx(_: *const abi.LaunchConfig, _: abi.Function, _: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) R { + return later(step()); + } + pub fn cuStreamSynchronize(_: abi.Stream) callconv(.c) R { + return drain(); + } + pub fn cuCtxSynchronize() callconv(.c) R { + return drain(); } pub fn cuStreamCreate(s: *abi.Stream, _: c_uint) callconv(.c) R { if (step() != ok) return code; diff --git a/zig/src/families/ling/graph.zig b/zig/src/families/ling/graph.zig index c202b8b966..e0f055b80c 100644 --- a/zig/src/families/ling/graph.zig +++ b/zig/src/families/ling/graph.zig @@ -155,7 +155,7 @@ pub const Graph = struct { pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, m: *const Modules, s: cuda.Stream, w: *const weights.Weights) !Graph { const c = &w.config; var g: Graph = .{ .w = w, .store = .{ .gpa = gpa, .d = d }, .slot = try gpa.alloc(u32, w.layers.len), .absorb = &.{}, .kda_layers = 0, .mla_layers = 0, .freq = 0, .zero_bias = 0 }; - errdefer g.deinit(); + errdefer g.drop(weights.drains(s)); // a failed drain leaks the store: a queued decode may still use it for (w.layers, g.slot) |l, *x| switch (l.mixer) { .kda => { x.* = @intCast(g.kda_layers); @@ -181,14 +181,24 @@ pub const Graph = struct { } pub fn deinit(g: *Graph) void { + g.drop(true); + } + + /// deinit, but the store's device buffers leaked unless `settled` (nothing queued can still use them). + pub fn drop(g: *Graph, settled: bool) void { const gpa = g.store.gpa; gpa.free(g.slot); gpa.free(g.absorb); - g.store.deinit(); + g.store.drop(settled); g.* = undefined; } }; +/// `b` freed, after draining `s` if work is `pending` on it; a failed drain leaks it (the device may still read it). +fn freeDrained(s: cuda.Stream, b: *cuda.DeviceBuffer, pending: bool) void { + if (!pending or weights.drains(s)) b.free(); +} + /// kv_b in its loaded format through the lane on identity rows: W^T [lw][heads (nope + v)], regrouped by head. fn absorbOf(gpa: std.mem.Allocator, store: *weights.Store, m: *const Modules, s: cuda.Stream, c: *const Config, kv_b: Proj, index: u32) !Absorb { const lw: usize = c.kv_lora; @@ -202,17 +212,20 @@ fn absorbOf(gpa: std.mem.Allocator, store: *weights.Store, m: *const Modules, s: defer gpa.free(eye); @memset(eye, 0); for (0..lw) |i| eye[i * lw + i] = 0x3f80; // bf16 1.0 + var pending = false; // the decode queued and not yet drained var x = try cuda.DeviceBuffer.fromHost(store.d, std.mem.sliceAsBytes(eye)); - defer x.free(); + defer freeDrained(s, &x, pending); var y = try cuda.DeviceBuffer.alloc(store.d, lw * n * 2); - defer y.free(); + defer freeDrained(s, &y, pending); var part = try cuda.DeviceBuffer.alloc(store.d, @max(16, m.lane.partBytes(lw, switch (kv_b) { .fp8g, .nvfp4, .fp8, .bf16 => |q| q, else => return error.FormatNotBuilt, }))); - defer part.free(); + defer freeDrained(s, &part, pending); + pending = true; try m.linear().decode(s, kv_b, .{ .x = x.ptr, .ldx = lw, .part = part.ptr }, y.ptr, lw); try s.synchronize(); + pending = false; const t = try gpa.alloc(u16, lw * n); defer gpa.free(t); try y.download(0, std.mem.sliceAsBytes(t)); diff --git a/zig/src/families/ling/weights.zig b/zig/src/families/ling/weights.zig index e1b7b2f16a..1ccecba653 100644 --- a/zig/src/families/ling/weights.zig +++ b/zig/src/families/ling/weights.zig @@ -12,6 +12,7 @@ pub const Proj = proj.Proj; pub const Mlp = proj.Mlp; pub const Routed = proj.Routed; pub const Store = proj.Store; +pub const drains = proj.drains; /// KDA: full-rank projections, the q|k|v short convs as #804's chain reads them, the decay and output norm tensors. pub const Kda = struct { @@ -76,8 +77,13 @@ pub const Weights = struct { mtp: ?Mtp = null, pub fn deinit(w: *Weights) void { + w.drop(true); + } + + /// deinit, but the device buffers leaked unless `settled` (nothing queued can still read them). + pub fn drop(w: *Weights, settled: bool) void { w.store.gpa.free(w.layers); - w.store.deinit(); + w.store.drop(settled); w.* = undefined; } @@ -203,8 +209,8 @@ const Loader = struct { /// Layers `layers` (index c.layers is the MTP layer) from `ck`, each tensor in its stored format; `stream` blocking. pub fn load(gpa: std.mem.Allocator, d: *const cuda.Driver, stream: cuda.Stream, pack: *const cuda.modelopt.Packer, ck: *core.Checkpoint, c: Config, layers: []const u32, opts: Options) !Weights { var w: Weights = .{ .config = c, .store = .{ .gpa = gpa, .d = d } }; - errdefer w.store.deinit(); var L: Loader = .{ .r = .{ .gpa = gpa, .stream = stream, .pack = pack, .ck = ck, .c = &w.config, .store = &w.store, .stage = opts.stage }, .c = &w.config, .opts = opts }; + errdefer w.store.drop(L.r.settled()); // the packers' outputs freed only once nothing queued writes them const out = try gpa.alloc(Layer, layers.len); errdefer gpa.free(out); for (layers, out) |i, *l| { diff --git a/zig/src/families/ling/weights_proj.zig b/zig/src/families/ling/weights_proj.zig index 1478695c85..0c7123e952 100644 --- a/zig/src/families/ling/weights_proj.zig +++ b/zig/src/families/ling/weights_proj.zig @@ -57,12 +57,24 @@ pub const Store = struct { return p; } + /// deinit, or when the device may still write them (a failed drain) the host lists alone: buffers leaked. + pub fn drop(s: *Store, settled: bool) void { + if (!settled) s.buffers.clearRetainingCapacity(); + s.deinit(); + } + pub fn find(s: *const Store, name: []const u8) ?Named { for (s.named.items) |n| if (std.mem.eql(u8, n.name, name)) return n; return null; } }; +/// Whether `s` drained: false when its synchronize fails (work queued on it may still touch memory). +pub fn drains(s: cuda.Stream) bool { + s.synchronize() catch return false; + return true; +} + /// Host bf16 from fp32: round to nearest even, NaN kept quiet. pub fn toBf16(v: f32) u16 { const b: u32 = @bitCast(v); @@ -88,6 +100,7 @@ pub const Reader = struct { c: *const config.Config, store: *Store, stage: usize = 512 << 20, // raw expert bytes staged at a time (weights.Options.stage) + queued: bool = false, // an upload or packer launch not yet drained buf: [256]u8 = undefined, fn name(r: *Reader, module: []const u8, tail: []const u8) ![]const u8 { @@ -110,19 +123,32 @@ pub const Reader = struct { }; } - /// A temporary device copy of `bytes` for a packer, freed by the caller after the stream drains. + /// A temporary device copy of `bytes` for a packer, freed by the caller once settled(). fn staged(r: *Reader, bytes: []const u8) !cuda.DeviceBuffer { return cuda.DeviceBuffer.fromHost(r.store.d, bytes); } + /// Nothing queued can still read or write memory: at once, or after a drain; false when the drain fails. + pub fn settled(r: *Reader) bool { + if (r.queued and !drains(r.stream)) return false; + r.queued = false; + return true; + } + + fn drain(r: *Reader) !void { + try r.stream.synchronize(); + r.queued = false; + } + fn bf16(r: *Reader, module: []const u8, n: usize, k: usize) !Proj { const t = try r.ck.get(try r.name(module, "weight")); if (!t.is(.bf16, &.{ n, k })) return wrong(module, .bf16, t, n, k); var src = try r.staged(t.bytes); - defer src.free(); + defer if (r.settled()) src.free(); // else leaked: a packer may still read it const out = try r.store.alloc(cuda.modelopt.padded(n) * k * 2, module); + r.queued = true; try r.pack.bf16Lane(r.stream, src.ptr, out, n, k); - try r.stream.synchronize(); + try r.drain(); return .{ .bf16 = cuda.modelopt.bf16Weight(out, n, k) }; } @@ -138,8 +164,9 @@ pub const Reader = struct { fn fp8(r: *Reader, module: []const u8, n: usize, k: usize) !Proj { const p = try r.fp8Pair(module, n, k); var src = try r.staged(p.w.bytes); - defer src.free(); + defer if (r.settled()) src.free(); // else leaked: a packer may still read it const codes = try r.store.alloc(cuda.fp8.padded(n) * k, module); + r.queued = true; try r.pack.fp8Lane(r.stream, src.ptr, codes, n, k); const inv = try r.gpa.alloc(f32, p.s.numel()); defer r.gpa.free(inv); @@ -148,7 +175,7 @@ pub const Reader = struct { defer r.gpa.free(tiles); cuda.fp8.packScales(tiles, inv, n, k); const scales = try r.store.keep(std.mem.sliceAsBytes(tiles), try r.name(module, "weight_scale_inv")); - try r.stream.synchronize(); + try r.drain(); return .{ .fp8g = cuda.fp8.weight(codes, scales, n, k) }; } @@ -163,14 +190,15 @@ pub const Reader = struct { fn fp4(r: *Reader, module: []const u8, n: usize, k: usize) !Proj { const l = try r.fp4Read(module, n, k); var codes = try r.staged(l.codes); - defer codes.free(); + defer if (r.settled()) codes.free(); // else both leaked var scales = try r.staged(l.scales); - defer scales.free(); + defer if (r.settled()) scales.free(); const np = cuda.nvfp4.padded(n); const words = try r.store.alloc(np / 64 * (k / 64) * 512 * 4, module); const tiles = try r.store.alloc(np * (k / 16), try r.name(module, "weight_scale")); + r.queued = true; try r.pack.fp4Lane(r.stream, codes.ptr, scales.ptr, words, tiles, n, k); - try r.stream.synchronize(); + try r.drain(); return .{ .nvfp4 = cuda.nvfp4.weight(words, tiles, l.global, n, k) }; } @@ -270,6 +298,7 @@ pub const Reader = struct { fn deinit(st: *Stage, r: *Reader) void { st.jobs.deinit(r.gpa); + if (!r.settled()) return; // both leaked on purpose: the upload or a packer may still read them st.host.free(); st.dev.free(); } @@ -294,12 +323,13 @@ pub const Reader = struct { st.copy(0); for (pool) |t| if (t) |x| x.join(); st.jobs.clearRetainingCapacity(); + r.queued = true; try st.dev.uploadAsync(0, st.host.bytes, r.stream.handle); } /// The packers done (the stage is reused) and the batch's file pages dropped: the GPU shares this memory. fn drained(_: *Stage, r: *Reader) !void { - try r.stream.synchronize(); + try r.drain(); r.ck.release(); } }; diff --git a/zig/tests/cuda/ling_faults.zig b/zig/tests/cuda/ling_faults.zig index 35efebc760..b380ed0cad 100644 --- a/zig/tests/cuda/ling_faults.zig +++ b/zig/tests/cuda/ling_faults.zig @@ -7,28 +7,52 @@ const ling = @import("ling_cuda"); const fake = ling.fake_driver; const Dbg = std.heap.DebugAllocator(.{ .safety = true }); // leak checks in release builds too -/// Opened (the engine, its sequence, the lane backend) and closed again. -fn once(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, mtp: bool) !void { +/// Opened (the engine, its sequence, the lane backend), with `keep` a state kept and dropped, and closed again. +fn once(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, mtp: bool, keep: bool) !void { const l = try open(gpa, io, ctx, dir, mtp); - l.deinit(l.ctx); + defer l.deinit(l.ctx); + if (keep) l.target.dropState(try l.target.keepState()); } fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, mtp: bool) !ling.native.Loaded { return ling.native.open(gpa, io, ctx, dir, null, .{ .context = 256, .drafts = mtp }); } -/// What one failed open left: the error, host leaks, the fake driver's live and bad releases. -const Left = struct { err: ?anyerror, leaked: bool, live: usize, bad: usize }; +/// What one failed open left: the error, host leaks, the fake driver's live, bad releases, frees with work pending. +const Left = struct { err: ?anyerror, leaked: bool, live: usize, bad: usize, early: usize, stuck: usize, pending: bool }; -fn failed(io: std.Io, ctx: *const cuda.Context, dir: []const u8, mtp: bool, call: ?u64, alloc: ?usize) Left { +fn failed(io: std.Io, ctx: *const cuda.Context, dir: []const u8, mtp: bool, keep: bool, call: ?u64, alloc: ?usize) Left { fake.reset(call, cuda.abi.error_not_ready); var dbg: Dbg = .init; var fa = std.testing.FailingAllocator.init(dbg.allocator(), .{ .fail_index = alloc orelse std.math.maxInt(usize) }); - const err: ?anyerror = if (once(fa.allocator(), io, ctx, dir, mtp)) |_| null else |e| e; + const err: ?anyerror = if (once(fa.allocator(), io, ctx, dir, mtp, keep)) |_| null else |e| e; const live = fake.live(); const bad = fake.badReleases(); const leaked = dbg.deinit() == .leak; - return .{ .err = err, .leaked = leaked, .live = live, .bad = bad }; + const out: Left = .{ .err = err, .leaked = leaked, .live = live, .bad = bad, .early = fake.freedEarly(), .stuck = fake.failedDrains(), .pending = fake.failedPending() }; + _ = fake.sweep(); // what this point left held does not count against the next + return out; +} + +const Sticky = struct { tried: usize = 0, bad: usize = 0, held: usize = 0, pending: usize = 0 }; + +/// Call k and every one after it failing (the open alone): no free with work pending, leaks only past a failed drain. +fn sticky(io: std.Io, ctx: *const cuda.Context, dir: []const u8, mtp: bool, calls: u64, every: usize) Sticky { + fake.mode(true, true); + defer fake.mode(true, false); + var out: Sticky = .{}; + var k: u64 = 1; + while (k <= calls) : (k += every) { + const l = failed(io, ctx, dir, mtp, false, k, null); + out.tried += 1; + out.held += @intFromBool(l.live != 0); + out.pending += @intFromBool(l.pending); + if (l.err == null or l.leaked or l.bad != 0 or l.early != 0 or (l.live != 0 and l.stuck == 0)) { + out.bad += 1; + if (out.bad <= 8) std.debug.print(" driver calls {d}.. failing: {?t}, host leak {}, {d} held ({d} failed drains), {d} bad releases, {d} freed while pending\n", .{ k, l.err, l.leaked, l.live, l.stuck, l.bad, l.early }); + } + } + return out; } pub fn run(gpu: check.Gpu, rest: []const [:0]const u8) !void { @@ -46,37 +70,46 @@ pub fn run(gpu: check.Gpu, rest: []const [:0]const u8) !void { } const d = fake.driver(); const ctx = fake.context(&d); - // a clean open first: its calls and host allocations are the failure points + fake.mode(true, false); // async copies and launches pending until a drain: no free may come before one, ever + defer fake.mode(false, false); + // a clean open and kept state first: their calls and host allocations are the failure points fake.reset(null, cuda.abi.success); var dbg: Dbg = .init; var counter = std.testing.FailingAllocator.init(dbg.allocator(), .{}); const clean = try open(counter.allocator(), gpu.io, &ctx, dir, mtp); - const calls = fake.count(); // the open's own calls and allocations: close's are not failure points + const opened = fake.count(); // the open's own calls: close's are not failure points + clean.target.dropState(try clean.target.keepState()); + const calls = fake.count(); const allocs = counter.allocations; clean.deinit(clean.ctx); - try check.expect(dbg.deinit() == .ok and fake.live() == 0 and fake.badReleases() == 0, "a clean open and close leaves nothing", .{}); + try check.expect(dbg.deinit() == .ok and fake.live() == 0 and fake.badReleases() == 0 and fake.freedEarly() == 0, "a clean open and close leaves nothing", .{}); var bad: usize = 0; var tried: usize = 0; + var pending: usize = 0; var k: u64 = 1; while (k <= calls) : (k += every) { - const l = failed(gpu.io, &ctx, dir, mtp, k, null); + const l = failed(gpu.io, &ctx, dir, mtp, true, k, null); tried += 1; - if (l.err == null or !fake.fired() or l.leaked or l.live != 0 or l.bad != 0) { + pending += @intFromBool(l.pending); + if (l.err == null or !fake.fired() or l.leaked or l.live != 0 or l.bad != 0 or l.early != 0) { bad += 1; - if (bad <= 8) std.debug.print(" driver call {d} failing: {?t}, host leak {}, {d} held, {d} bad releases\n", .{ k, l.err, l.leaked, l.live, l.bad }); + if (bad <= 8) std.debug.print(" driver call {d} failing: {?t}, host leak {}, {d} held, {d} bad releases, {d} freed while pending\n", .{ k, l.err, l.leaked, l.live, l.bad, l.early }); } } var a: usize = 0; while (a < allocs) : (a += every) { - const l = failed(gpu.io, &ctx, dir, mtp, null, a); + const l = failed(gpu.io, &ctx, dir, mtp, true, null, a); tried += 1; - if (l.err == null or l.leaked or l.live != 0 or l.bad != 0) { + if (l.err == null or l.leaked or l.live != 0 or l.bad != 0 or l.early != 0) { bad += 1; if (bad <= 16) std.debug.print(" host allocation {d} failing: {?t}, host leak {}, {d} held, {d} bad releases\n", .{ a, l.err, l.leaked, l.live, l.bad }); } } + const st = sticky(gpu.io, &ctx, dir, mtp, opened, every); + bad += st.bad; + tried += st.tried; fake.reset(null, cuda.abi.success); - std.debug.print("RESULT ling-init-faults{s}: {d} driver calls and {d} host allocations in a clean open; {d} failure points tried, {d} left something behind\n", .{ if (mtp) " --mtp" else "", calls, allocs, tried, bad }); - try check.expect(bad == 0, "every failing driver call and host allocation in the open is cleaned up", .{}); - check.pass("ling-init-faults{s}: {d} failure points, each open refused with nothing left held, leaked or freed twice", .{ if (mtp) " --mtp" else "", tried }); + std.debug.print("RESULT ling-init-faults{s}: {d} driver calls ({d} the open's) and {d} host allocations in a clean open and kept state; {d} failure points tried ({d} with a transfer pending), {d} left something behind; failing from a call on: {d} points ({d} with a transfer pending), {d} leaking after a failed drain\n", .{ if (mtp) " --mtp" else "", calls, opened, allocs, tried, pending + st.pending, bad, st.tried, st.pending, st.held }); + try check.expect(bad == 0, "every failing driver call and host allocation in the open is cleaned up, nothing freed with work pending", .{}); + check.pass("ling-init-faults{s}: {d} failure points, each refused with nothing freed before a drain, held, leaked or freed twice", .{ if (mtp) " --mtp" else "", tried }); } From e99e6505fc3c632577c19fe70638f9e8a6aa09ac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 16 Oct 2026 21:19:08 +0200 Subject: [PATCH 34/34] docs/recipes/ling-3.1: file names for hashes the rebase changed, current check counts, the recorder on main --- docs/recipes/ling-3.1.md | 29 ++++++++++++++--------------- 1 file changed, 14 insertions(+), 15 deletions(-) diff --git a/docs/recipes/ling-3.1.md b/docs/recipes/ling-3.1.md index a94fff9aef..97af269852 100644 --- a/docs/recipes/ling-3.1.md +++ b/docs/recipes/ling-3.1.md @@ -14,9 +14,9 @@ block FP8 (e4m3 in 128 x 128 blocks with `weight_scale_inv`). | Part | Where | Checked by | | --- | --- | --- | -| Group-limited top-k routing, per-layer shared SwiGLU limit (GLM's bytes unchanged) | `ling-family`, commit 51c01573 | `glm-ref`, `glm-route-old`, `glm-route-ling` | -| MLA at Ling's shape: 64 roped key columns, 128-wide heads, GPT-J rope, head-wise gate | f8e594f6, merged | `ling-mla`, `ling-mla-real`, `glm-mla-same` | -| KDA on Ling's layer 0 (q/k/v convs as one q\|k\|v conv), GLM's row path as is | b55aa13d, merged | `glm-kda-ling` | +| Group-limited top-k routing, per-layer shared SwiGLU limit (GLM's bytes unchanged) | `glm_glue_moe.cuh` (`tf_glm_topk`) | `glm-ref`, `glm-route-old`, `glm-route-ling` | +| MLA at Ling's shape: 64 roped key columns, 128-wide heads, GPT-J rope, head-wise gate | `glm_glue_attn.cuh`, `glm_glue_latent.cuh`, `block_mla.zig` | `ling-mla`, `ling-mla-real`, `glm-mla-same` | +| KDA on Ling's layer 0 (q/k/v convs as one q\|k\|v conv), GLM's row path as is | `glm_kda.cu`, `block_kda.zig` | `glm-kda-ling` | | Config loader and weight map (all 214,025 names of the real index) | `zig/src/families/ling/config.zig`, `names.zig` | `zig build test-ling` | | Weight loader at TP=1: bf16, FP8 block, ModelOpt NVFP4 | `zig/src/families/ling/weights.zig`, `weights_proj.zig` | `ling-weights`, `ling-weights-real` | | Synthetic checkpoints in all three formats | `tf-cuda-test ling-synth` | its own name check | @@ -121,7 +121,7 @@ disk / the sshfs mount: | Loader | cold load | host peak (VmHWM) | device peak | | --- | --- | --- | --- | -| host packers per expert (c48be499) | 14.3 s / ~157 s | 8.0 GiB (the layer's mapped pages) | 7.9 GiB | +| host packers per expert (before GPU packing) | 14.3 s / ~157 s | 8.0 GiB (the layer's mapped pages) | 7.9 GiB | | batches packed on the GPU | 6.1 s / 122 s | 1.3 GiB | 8.4 GiB (+ one stage) | From local disk the load runs at the read's rate (5-7 s to read the layer once single-threaded); from the mount the @@ -134,14 +134,15 @@ Layers 64-71 clamp the shared expert and layers 65-71 the routed experts at L = GLM's glue and the grouped expert kernels clamp the gate before SiLU, so Ling launches its own instances: the glue's `tf_ling_swiglu_512`/`_64` (`glue.swigluWith`) for the shared expert, dense and bf16 routed experts, and epilogue 4 of the FP8 and NVFP4 grouped kernels (decode and prompt), chosen by `experts.Layer.after`, which the Ling loader sets. -Every kernel that existed before has the same SASS and resource usage as at a33dde0c (`cuobjdump`, all 40 fatbins); -`ling-clamp-old` runs the kept FP8 (12 calls) and NVFP4 (13 calls) instances against a33dde0c's fatbins byte for byte. +Every kernel that existed before has the same SASS and resource usage as at 6f583baf, the layer graph's commit +(`cuobjdump`, all 40 fatbins); `ling-clamp-old` runs the kept FP8 (12 calls) and NVFP4 (13 calls) instances against +6f583baf's fatbins byte for byte. At L = 10 the two orders give the same bf16 bytes: the kernels (and HF's bf16 `act_fn`) round SiLU to bf16 first, and bf16(silu(g)) is 10 for every bf16 g >= 10. They differ at small limits (L = 2: 16,255 of the 65,536 bf16 gates). `ling-clamp DIR` checks Ling's order against float64 at L = 10 and L = 2 on gates and ups straddling L: the glue at widths 1024 and 704 (worst 0.93 of 2^-7 relative; 0 bytes differ from GLM's order at 10, 6,794 at 2) and the clamped -layer's grouped gate-up, decode and prompt rows (FP8 worst 0.090, NVFP4 0.102 of the bound). a33dde0c's binary and +layer's grouped gate-up, decode and prompt rows (FP8 worst 0.090, NVFP4 0.102 of the bound). 6f583baf's binary and this one print the same digits on the real layers (71: rel L2 4.125e-3, 3.747e-3 without 7 flipped rows, decode rows 3.893e-3; 65: 4.851e-3, 3.347e-3 without 14, 3.325e-3) and on the synthetic forward in all three formats, whose float64 references do not change either (the synthetic clamped layer never reaches g > 10). HF's clamp itself moves @@ -158,8 +159,8 @@ upstream's `sample.cu` through GLM's sampler at TP=1. `native.zig` gives the ser 16 rows, shared rounds unless `TF_LING_SHARED=0`; `TF_LING_PROMPT_STEPS=0` passes prompts whole) and the prompt cache's `snap_*` hooks: a kept state is the sequence's arena copied whole at a mark, and the lane backend restores it inside its own prompt pass. Ling's prompt arithmetic does not depend on where a chunk ends, so a mark may fall on any -row (no `snap_grid`): the pass cuts its chunk there. There is no recorder hook (`features`): the drafter recorder -(#848) is not on this branch. +row (no `snap_grid`): the pass cuts its chunk there. There is no recorder hook (`features`) yet, so `--record` (#848) +is refused for `bailing_hybrid`. `ling-lanes DIR OUT [REF] [--layers a,b]` drives the backend under upstream's round loop as the server's lane host does, prompts in 37-row chunks; every check compares tokens, and the engine-level ones bytes: @@ -184,8 +185,8 @@ does, prompts in 37-row chunks; every check compares tokens, and the engine-leve | FP8 block | 2.09e-2 | 1.18e-2 | 2.10e-2 | 302/309 | | NVFP4 | 2.21e-2 | 1.38e-2 | 2.22e-2 | 294/309 | -92 checks pass in each format. On the real FP8 release, `--layers 0,1,2,3` (the real embedding, layers 0-3, final -norm and head; KDA and dense MLP only, 6 GB) passes the same 90 exactness checks; it has no float64 reference (an +100 checks pass in each format. On the real FP8 release, `--layers 0,1,2,3` (the real embedding, layers 0-3, final +norm and head; KDA and dense MLP only, 6 GB) passes the same exactness checks; it has no float64 reference (an MLA + MoE layer adds 8 GB). Mutations caught: committing every row on `keep`, a restore that keeps the fresh KDA copy bits, chunks not cut at marks, shared draws keyed one position early, a lane commit that keeps one row too many. @@ -201,7 +202,7 @@ final-norm rows and the tokens after them), then chains up to 3 drafts, each ste `final_layernorm` output, as GLM's `decode.draft` does; the chain's rows are dropped before the next absorb. Kept prompt states copy the MTP cache too. The server loads it unless `--no-drafts`; the lane facts then say `mtp`. -Checked (`ling-lanes ... --mtp`, 104 checks a format; acceptance on random weights says nothing about quality and is +Checked (`ling-lanes ... --mtp`, 109 checks a format; acceptance on random weights says nothing about quality and is not reported as such): MTP-drafted replies, greedy and sampled, equal the plain ones; so do MTP-drafted passes resumed at 53 and 211 and every concurrent case; a prompt pass's MTP cache in 37-row chunks equals the serial MTP rows byte for byte; twins drafting chains of 3 and of 1 agree on tokens and first drafts over 6 windows, each @@ -312,7 +313,7 @@ size. Each item stands alone. -1. **Recorder.** Give `Lanes` a `features` hook (the final-norm rows `Engine.verify` already keeps) once #848 lands. +1. **Recorder.** Give `Lanes` a `features` hook (the final-norm rows `Engine.verify` already keeps) for #848's recorder. 2. **MTP acceptance.** Measure the MTP layer's acceptance on real text once a whole model runs (TP). 3. **Fused projections.** Same-format q|k|v (and gate|up) rows as one weight per format, without the column copies. 4. **Bf16 routed experts.** Add a grouped bf16 expert kernel, or quantize at load, and drop the per-expert lanes. @@ -321,5 +322,3 @@ Each item stands alone. 6. **EXL3.** Wire #832's reader behind the `config.Quant` extension point and lift the refusal. 7. **An NVFP4 checkpoint.** Quantize the release with ModelOpt, keep the routers in bf16, and run `ling-weights` on it; the loader expects `weight`, `weight_scale`, `weight_scale_2` and `input_scale` per linear. -8. **GPU packing of FP8 experts.** They pack on the host today (`Fp8Experts.packGateUp`); `modelopt.Packer` has the - NVFP4 experts' GPU packer to follow.