From 471193c6af891f1893fd157efa446e099f0aa359 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 01:22:21 +0100 Subject: [PATCH 01/10] lfm2: a verify's greedy rows take the GPU's pick of their logits The host drew every verify row by scanning its 128,000 bf16 logits, after the forward and with the GPU idle. lfm2_argmax picks each head row's first largest logit on the GPU and flags a row with a nonfinite logit. verifyFn takes the pick for greedy rows without a mask; sampled, masked and flagged rows draw on the host as before. The kernel is registered with the contract ("metal lfm2 argmax": the host draw's pick, ties to the lower id, nonfinite rows flagged), and tf-lfm2-exact checks every window's picks against the host draw. --- zig/kernels/metal/lfm2/lfm2.metal | 57 +++++++++++ zig/src/families/lfm2_moe/backend.zig | 7 ++ zig/src/families/lfm2_moe/forward.zig | 12 ++- zig/src/families/lfm2_moe/kernels.zig | 1 + zig/src/families/lfm2_moe/state.zig | 3 + zig/tests/contract/metal_lfm2_argmax.zig | 125 +++++++++++++++++++++++ zig/tests/contract/metal_main.zig | 8 ++ zig/tests/lfm2_exact.zig | 12 +++ 8 files changed, 222 insertions(+), 3 deletions(-) create mode 100644 zig/tests/contract/metal_lfm2_argmax.zig diff --git a/zig/kernels/metal/lfm2/lfm2.metal b/zig/kernels/metal/lfm2/lfm2.metal index 0d17e093ae..ab417e1f62 100644 --- a/zig/kernels/metal/lfm2/lfm2.metal +++ b/zig/kernels/metal/lfm2/lfm2.metal @@ -494,3 +494,60 @@ constant constexpr uint SORT_CHUNKS = 16; // thread (e, c) counts expert e's pi } r[size_t(at.y) * D + at.x] = static_cast(acc); } + +// Row r's greedy pick over `vocab` bf16 logits: the lowest id of the largest, then 1 if any logit is not finite. +[[kernel]] void lfm2_argmax(const device ushort* logits [[buffer(0)]], constant uint& vocab [[buffer(1)]], + device uint* best [[buffer(2)]], uint row [[threadgroup_position_in_grid]], + uint t [[thread_index_in_threadgroup]], uint sg [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + threadgroup float tv[32]; + threadgroup uint ti[32], tn[32]; + const device ushort* x = logits + size_t(row) * vocab; + float bv = -INFINITY; + uint bi = 0xffffffffu, bad = 0; + // Eight logits a load, ids rising along a thread: a strictly larger value moves the pick (vocab % 8 == 0). + for (uint v = 8 * t; v < vocab; v += 8 * 1024) { + const uint4 w = *(const device uint4*)(x + v); + const uint words[4] = {w.x, w.y, w.z, w.w}; + for (uint j = 0; j < 8; ++j) { + const float f = as_type((j % 2) ? (words[j / 2] & 0xFFFF0000u) : (words[j / 2] << 16)); + if (!isfinite(f)) bad = 1; + else if (f > bv) { + bv = f; + bi = v + j; + } + } + } + for (uint off = 16; off > 0; off >>= 1) { + const float ov = simd_shuffle_xor(bv, off); + const uint oi = simd_shuffle_xor(bi, off); + if (ov > bv || (ov == bv && oi < bi)) { + bv = ov; + bi = oi; + } + } + bad = simd_max(bad); + if (lane == 0) { + tv[sg] = bv; + ti[sg] = bi; + tn[sg] = bad; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg != 0) return; + bv = tv[lane]; + bi = ti[lane]; + bad = tn[lane]; + for (uint off = 16; off > 0; off >>= 1) { + const float ov = simd_shuffle_xor(bv, off); + const uint oi = simd_shuffle_xor(bi, off); + if (ov > bv || (ov == bv && oi < bi)) { + bv = ov; + bi = oi; + } + } + bad = simd_max(bad); + if (lane == 0) { + best[2 * row] = bi; + best[2 * row + 1] = bad; + } +} diff --git a/zig/src/families/lfm2_moe/backend.zig b/zig/src/families/lfm2_moe/backend.zig index 670b9dacee..f83935333c 100644 --- a/zig/src/families/lfm2_moe/backend.zig +++ b/zig/src/families/lfm2_moe/backend.zig @@ -193,11 +193,18 @@ pub const Metal = struct { } try fwd.run(self.model, &self.scratch, segments[0..windows.len], ids[0..total], true, .all); const logits = self.scratch.logits.slice(u16, total * c.vocab); + const best = self.scratch.best.slice(u32, 2 * total); var base: usize = 0; for (windows, out) |w, o| { if (o.sampled.len != w.rows() or o.drafts.len != w.tokens.len) return error.InvalidLfmOutputs; for (o.sampled, w.positions, 0..) |*t, p, r| { const mask = if (w.masks.len > 0) w.masks[r * mask_words ..][0..mask_words] else &.{}; + // A greedy row takes the GPU's pick, the host's draw on the same logits; a nonfinite row errors there. + const greedy = if (w.stream.sampling) |s| s.temperature <= 0 else true; + if (greedy and mask.len == 0 and best[2 * (base + r) + 1] == 0) { + t.* = best[2 * (base + r)]; + continue; + } t.* = try draw(self.gpa, logits[(base + r) * c.vocab ..][0..c.vocab], w.stream.sampling, p, mask); } @memcpy(o.drafts, w.tokens); diff --git a/zig/src/families/lfm2_moe/forward.zig b/zig/src/families/lfm2_moe/forward.zig index 0bbfc3ef02..574e536839 100644 --- a/zig/src/families/lfm2_moe/forward.zig +++ b/zig/src/families/lfm2_moe/forward.zig @@ -12,6 +12,7 @@ const Buffer = mtl.Buffer; pub const Segment = struct { cache: *st.Cache, rows: usize }; /// GPU seconds of every call so far (the bench reports a chunk's GPU time beside its wall time). pub var gpu_seconds: f64 = 0; +/// `all`: every row's logits and greedy pick (Scratch.best); `last`: the last row's logits alone. pub const Head = enum { last, all }; /// Rows from which a call takes the tiles (the twin's bits, each weight read once for every row). const tile_rows = 33; @@ -157,9 +158,14 @@ pub fn run(m: *Model, s: *st.Scratch, segments: []const Segment, tokens: []const } } enc.norm(s, m.weights.norm, false, total, combine); - if (head == .all) enc.projection(embedding, s.x, 0, s.logits, 0, total) else { - enc.projection(embedding, s.x, (total - 1) * c.hidden * 2, s.logits, 0, 1); - } + if (head == .all) { + enc.projection(embedding, s.x, 0, s.logits, 0, total); + enc.pipe("lfm2_argmax"); + e.setBuffer(s.logits, 0, 0); + e.setValue(@as(u32, c.vocab), 1); + e.setBuffer(s.best, 0, 2); + enc.groups(.{ total, 1, 1 }, .{ 1024, 1, 1 }); + } else enc.projection(embedding, s.x, (total - 1) * c.hidden * 2, s.logits, 0, 1); e.end(); cb.commit(); cb.wait(); diff --git a/zig/src/families/lfm2_moe/kernels.zig b/zig/src/families/lfm2_moe/kernels.zig index 2cbf18b018..48da29433f 100644 --- a/zig/src/families/lfm2_moe/kernels.zig +++ b/zig/src/families/lfm2_moe/kernels.zig @@ -22,6 +22,7 @@ const names = [_][:0]const u8{ "lfm2_combine_norm", "lfm2_sort", "lfm2_combine_sorted", + "lfm2_argmax", }; pub const Kernels = struct { diff --git a/zig/src/families/lfm2_moe/state.zig b/zig/src/families/lfm2_moe/state.zig index 955110fb1e..646510581f 100644 --- a/zig/src/families/lfm2_moe/state.zig +++ b/zig/src/families/lfm2_moe/state.zig @@ -141,6 +141,8 @@ pub const Scratch = struct { /// A projection's input rows' group sums for the tiles and their twin (f32 [rows, k / 64]). sums: mtl.Buffer, logits: mtl.Buffer, + /// Each head row's greedy pick and nonfinite flag (u32 pairs), from a forward whose head takes every row. + best: mtl.Buffer, snapshots: [c.layers]?mtl.Buffer = @splat(null), pub fn init(gpa: std.mem.Allocator, device: mtl.Device, attention: *const [c.layers]bool, rows: usize) !Scratch { @@ -179,6 +181,7 @@ pub const Scratch = struct { out.sums = try mem.alloc(@max(rows * (c.intermediate / 64), slots * (c.hidden / 64)) * 4); const logit_rows: usize = @min(rows, batch_rows); out.logits = try mem.alloc(logit_rows * c.vocab * 2); + out.best = try mem.alloc(logit_rows * 2 * 4); for (&out.snapshots, attention) |*s, a| if (!a) { // a recorded window's rows, or a prompt call's last row s.* = try mem.alloc(logit_rows * conv_bytes); }; diff --git a/zig/tests/contract/metal_lfm2_argmax.zig b/zig/tests/contract/metal_lfm2_argmax.zig new file mode 100644 index 0000000000..a60e4e273f --- /dev/null +++ b/zig/tests/contract/metal_lfm2_argmax.zig @@ -0,0 +1,125 @@ +//! LFM2's greedy pick of every verify row (lfm2_argmax): the lowest id of a row's largest bf16 logit, and a flag. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); +const c = contract.spec; +const g = contract.gates; +const rig_ = @import("rig.zig"); + +const vocab = 128000; +pub const name = "metal lfm2 argmax"; + +/// rows: a verify's head rows, one to a shared round's 32. +const shapes = [_]c.Shape{ .{ .rows = 1, .n = vocab, .k = 1 }, .{ .rows = 9, .n = vocab, .k = 1 }, .{ .rows = 32, .n = vocab, .k = 1 } }; + +fn none(_: c.Shape) usize { + return 0; +} + +pub fn spec() c.Spec { + return .{ .name = name, .accum = .none, .depth = none, .out = .raw, .served = &shapes, .nan = .{ .stated = "a row with a NaN or an infinity is flagged nonfinite" }, .sums = .{ .f16 = 0, .bf16 = 0, .why = "picks compare values: no sums" } }; +} + +/// A bf16 logit for element `i`: few distinct values (ties in every row); the nonfinite fill plants NaN and infinities. +fn value(i: usize, seed: u64, fill: g.Fill) u16 { + const h: u64 = (@as(u64, i) *% 0x9e3779b97f4a7c15 +% seed *% 0xbf58476d1ce4e5b9) >> 40; + if (fill == .nonfinite and h % 50021 < 3) return ([3]u16{ 0x7fc0, 0x7f80, 0xff80 })[h % 50021]; + const v = @as(f32, @floatFromInt(h % 61)) / 8 - 4; // -4 to 3.5 in eighths, signed zeros among them + return if (v == 0 and (h >> 20) & 1 == 1) 0x8000 else c.toBf16(v); +} + +pub const Probe = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + rows: usize = 0, + x: []u16 = &.{}, + args: [2]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue) !Probe { + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .pipe = try rig_.pipeline(device, sources.lfm2, "lfm2_argmax") }; + } + + fn drop(p: *Probe) void { + p.a.free(p.x); + p.x = &.{}; + p.rig.reset(); + } + + pub fn deinit(p: *Probe) void { + p.drop(); + p.rig.deinit(); + p.pipe.deinit(); + } + + const vt = g.hooks(Probe); + pub fn kernel(p: *Probe) g.Kernel { + return .{ .ctx = p, .vt = &vt, .spec = spec() }; + } + + pub fn prepare(p: *Probe, s: c.Shape, f: g.Fill, seed: u64) !bool { + p.drop(); + if (s.n != vocab or s.k != 1 or s.rows == 0 or s.rows > 32) return false; + p.rows = s.rows; + p.x = try p.a.alloc(u16, s.rows * vocab); + for (p.x, 0..) |*v, i| v.* = value(i, seed, f); + // Row 0's best value sits at two ids: the pick is the lower one. + p.x[0 * vocab + 70001] = c.toBf16(5); + p.x[0 * vocab + 1234] = c.toBf16(5); + p.args[0] = try p.rig.input(std.mem.sliceAsBytes(p.x)); + const n: u32 = vocab; + p.args[1] = try p.rig.input(std.mem.asBytes(&n)); + p.out = try p.rig.output(p.count()); + return true; + } + + pub fn count(p: *Probe) usize { + return p.rows * 8; + } + + pub fn run(p: *Probe) !void { + const pass = p.rig.begin(); + const e = pass.e; + e.setPipeline(p.pipe); + e.setBuffer(p.args[0].buf, p.args[0].off, 0); + e.setBuffer(p.args[1].buf, p.args[1].off, 1); + e.setBuffer(p.out.buf, p.out.off, 2); + e.dispatchGroups(mtl.Size.of(p.rows, 1, 1), mtl.Size.of(1024, 1, 1)); + try p.rig.finish(pass); + } + + pub fn reference(_: *Probe, want: []f64, mag: []f64) !void { + @memset(want, 0); // the picks check themselves (check) + @memset(mag, 0); + } + + /// Each row's pick against the host's greedy draw: the first id of the largest value, or the row flagged. + pub fn check(p: *Probe, _: g.Gate) !?[]const u8 { + const out = p.rig.bytes(p.out, p.count()); + for (0..p.rows) |r| { + const token = std.mem.readInt(u32, out[8 * r ..][0..4], .little); + const flag = std.mem.readInt(u32, out[8 * r + 4 ..][0..4], .little); + var best: ?usize = null; + var bad = false; + for (p.x[r * vocab ..][0..vocab], 0..) |w, id| { + const v = c.fromBf16(w); + if (!std.math.isFinite(v)) bad = true else if (best == null or v > c.fromBf16(p.x[r * vocab + best.?])) best = id; + } + if (bad) { + if (flag != 1) return "a row with a nonfinite logit is not flagged"; + continue; + } + if (flag != 0) return "a finite row is flagged nonfinite"; + if (token != best.?) return "a pick differs from the first id of the row's largest logit"; + } + return null; + } +}; + +/// Gate 5 for the launch on this GPU: 1,024 threads, 32 simdgroups' best values, ids and flags. +pub fn limits(a: std.mem.Allocator, device: mtl.Device, only: ?[]const u8) !usize { + if (only) |s| if (std.mem.indexOf(u8, name, s) == null) return 0; + return @import("metal_limits.zig").onThisGpu(a, device, .{ .decl = .{ .name = name, .threads = 1024, .tg_memory = 384, .offline = false }, .function = "lfm2_argmax", .here = sources.lfm2, .pre_m5 = "" }, false); +} diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index f266ffcd2d..97a3f9153a 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -23,6 +23,7 @@ const flash = @import("metal_flashnext.zig"); const gemma = @import("metal_gemma4.zig"); const gemma_embed = @import("metal_gemma4_embed.zig"); const gemma_mixed = @import("metal_gemma4_mixed.zig"); +const lfm2_argmax = @import("metal_lfm2_argmax.zig"); const g = contract.gates; const usage = "usage: tf-kernel-contract [--only TEXT] [--no-gates] [--no-limits] [--offline DIR] [--seed N] [--repro] [--rotary-long] [--rotary-accuracy FILE]"; @@ -151,6 +152,12 @@ pub fn main(init: std.process.Init) !u8 { if (!tensor) return error.TensorUnitsRequired; try kernels.append(a, mixed.kernel()); } + var argmax: ?lfm2_argmax.Probe = null; + defer if (argmax) |*p| p.deinit(); + if (wanted(o, lfm2_argmax.name)) { + argmax = try lfm2_argmax.Probe.init(a, device, queue); + try kernels.append(a, argmax.?.kernel()); + } var rows_list: std.ArrayList(*rows_k.Row) = .empty; var proj_list: std.ArrayList(*affine_k.Proj) = .empty; defer for (rows_list.items) |r| r.deinit(); @@ -290,6 +297,7 @@ pub fn main(init: std.process.Init) !u8 { lim += try gemma.limits(a, device, o.only); lim += try gemma_embed.limits(a, device, o.only); lim += try gemma_mixed.limits(a, device, o.only); + lim += try lfm2_argmax.limits(a, device, o.only); } const t2 = std.Io.Clock.awake.now(init.io); const ms = struct { diff --git a/zig/tests/lfm2_exact.zig b/zig/tests/lfm2_exact.zig index f9a8d6b66c..6fe9a3980b 100644 --- a/zig/tests/lfm2_exact.zig +++ b/zig/tests/lfm2_exact.zig @@ -33,6 +33,16 @@ fn prompt(m: *q.Model, scratch: *st.Scratch, cache: *st.Cache, ids: []const u32) } } +/// Each of a forward's `rows` head rows: the GPU's greedy pick equals the host's draw on its logits. +fn picks(m: *q.Model, s: *st.Scratch, rows: usize) !void { + const best = s.best.slice(u32, 2 * rows); + for (0..rows) |r| { + const want = try q.backend.draw(m.gpa, s.logits.slice(u16, rows * c.vocab)[r * c.vocab ..][0..c.vocab], null, 0, &.{}); + try std.testing.expectEqual(want, best[2 * r]); + try std.testing.expectEqual(@as(u32, 0), best[2 * r + 1]); + } +} + /// A prompt of `prefix` rows, then plain decode to `len`: what a window's kept rows equal. fn decoded(m: *q.Model, solo: *st.Scratch, cache: *st.Cache, ids: []const u32, prefix: usize, len: usize) !void { try prompt(m, solo, cache, ids[0..prefix]); @@ -48,6 +58,7 @@ fn window(m: *q.Model, wide: *st.Scratch, solo: *st.Scratch, ids: []const u32, p try prompt(m, solo, &b, ids[0..prefix]); try equal(&a, &b); try fwd.run(m, wide, &.{.{ .cache = &a, .rows = rows }}, ids[prefix..][0..rows], true, .all); + try picks(m, wide, rows); for (0..rows) |r| { try fwd.run(m, solo, &.{.{ .cache = &b, .rows = 1 }}, ids[prefix + r ..][0..1], false, .all); const want = solo.logits.contents()[0 .. c.vocab * 2]; @@ -124,6 +135,7 @@ fn round( total += r; } try fwd.run(m, wide, segs[0..n], tokens[0..total], true, .all); + try picks(m, wide, total); var base: usize = 0; for (caches, prompts, lens, rows) |*cache, start, len, r| { if (r == 0) continue; From 886d4173ee47c1536e52e869dbcefc2d989d3c90 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 01:23:05 +0100 Subject: [PATCH 02/10] lfm2: the conv, qk, attention and merge kernels' bodies as inline functions lfm2_conv, lfm2_qk, lfm2_attn_mma and lfm2_attn_merge each call an inline body (lfm2_conv_row, lfm2_qk_row, lfm2_attn_rows, lfm2_merge_row) with the same arguments, so another kernel can run a row's arithmetic unchanged. The kernels' bits do not change. --- zig/kernels/metal/lfm2/lfm2.metal | 63 ++++++++++++++++++++----------- 1 file changed, 40 insertions(+), 23 deletions(-) diff --git a/zig/kernels/metal/lfm2/lfm2.metal b/zig/kernels/metal/lfm2/lfm2.metal index ab417e1f62..af6a2a7372 100644 --- a/zig/kernels/metal/lfm2/lfm2.metal +++ b/zig/kernels/metal/lfm2/lfm2.metal @@ -105,11 +105,8 @@ inline void lfm2_rms(const thread float* v, const device bfloat* w, float eps, d } // A segment's short conv: Bx = B * x, 3 taps over (state [2, D], Bx), y = C * conv; SNAP: rows' states from FIRST on. -[[kernel]] void lfm2_conv(const device bfloat* bcx [[buffer(0)]], const device bfloat* state [[buffer(1)]], - const device bfloat* w [[buffer(2)]], device bfloat* y [[buffer(3)]], device bfloat* snap [[buffer(4)]], - constant uint& first [[buffer(5)]], uint2 at [[thread_position_in_grid]]) { - const uint ch = at.x, r = at.y; - if (ch >= D) return; +inline void lfm2_conv_row(const device bfloat* bcx, const device bfloat* state, const device bfloat* w, device bfloat* y, + device bfloat* snap, uint first, uint ch, uint r) { bfloat bx[3]; for (int tap = 0; tap < 3; ++tap) { const int src = int(r) + tap - 2; @@ -127,6 +124,11 @@ inline void lfm2_rms(const thread float* v, const device bfloat* w, float eps, d snap[(size_t(r - first) * 2) * D + ch] = bx[1]; snap[(size_t(r - first) * 2 + 1) * D + ch] = bx[2]; } +[[kernel]] void lfm2_conv(const device bfloat* bcx [[buffer(0)]], const device bfloat* state [[buffer(1)]], + const device bfloat* w [[buffer(2)]], device bfloat* y [[buffer(3)]], device bfloat* snap [[buffer(4)]], + constant uint& first [[buffer(5)]], uint2 at [[thread_position_in_grid]]) { + if (at.x < D) lfm2_conv_row(bcx, state, w, y, snap, first, at.x, at.y); +} // One head of one row: RMS norm over its 64 values (lane l: l and l + 32), then RoPE on the pair at position `pos`. inline void lfm2_head(const device bfloat* src, const device bfloat* gain, float eps, float log2_theta, uint pos, uint lane, @@ -143,11 +145,10 @@ inline void lfm2_head(const device bfloat* src, const device bfloat* gain, float // A simdgroup a (head, row): queries normed and rotated, keys and values into the [KVH, cap, HD] cache at len + r. struct lfm2_qk_args { uint len, cap; float eps, log2_theta; }; -[[kernel]] void lfm2_qk(const device bfloat* q [[buffer(0)]], const device bfloat* k [[buffer(1)]], const device bfloat* v [[buffer(2)]], - const device bfloat* qn [[buffer(3)]], const device bfloat* kn [[buffer(4)]], constant lfm2_qk_args& p [[buffer(5)]], - device bfloat* queries [[buffer(6)]], device bfloat* keys [[buffer(7)]], device bfloat* values [[buffer(8)]], - uint2 tg [[threadgroup_position_in_grid]], uint lane [[thread_index_in_simdgroup]]) { - const uint head = tg.x, r = tg.y, pos = p.len + r; +inline void lfm2_qk_row(const device bfloat* q, const device bfloat* k, const device bfloat* v, const device bfloat* qn, + const device bfloat* kn, lfm2_qk_args p, device bfloat* queries, device bfloat* keys, + device bfloat* values, uint head, uint r, uint lane) { + const uint pos = p.len + r; bfloat a, b; if (head < QH) { lfm2_head(q + (size_t(r) * QH + head) * HD, qn, p.eps, p.log2_theta, pos, lane, a, b); @@ -164,6 +165,12 @@ struct lfm2_qk_args { uint len, cap; float eps, log2_theta; }; values[at + lane] = v[(size_t(r) * KVH + kh) * HD + lane]; values[at + lane + 32] = v[(size_t(r) * KVH + kh) * HD + lane + 32]; } +[[kernel]] void lfm2_qk(const device bfloat* q [[buffer(0)]], const device bfloat* k [[buffer(1)]], const device bfloat* v [[buffer(2)]], + const device bfloat* qn [[buffer(3)]], const device bfloat* kn [[buffer(4)]], constant lfm2_qk_args& p [[buffer(5)]], + device bfloat* queries [[buffer(6)]], device bfloat* keys [[buffer(7)]], device bfloat* values [[buffer(8)]], + uint2 tg [[threadgroup_position_in_grid]], uint lane [[thread_index_in_simdgroup]]) { + lfm2_qk_row(q, k, v, qn, kn, p, queries, keys, values, tg.x, tg.y, lane); +} // Keys a chunk: a row's keys split at fixed positions, so prompt tiles and decode rows cut them the same way. constant constexpr uint PER = 4096; @@ -226,14 +233,10 @@ inline void lfm2_stage(const thread uint4* pre, threadgroup uint* ks, threadgrou // Every row's attention partials (one kernel for prompt and decode): a threadgroup a (chunk, kv head, 8 * RB rows). template // a simdgroup's RB row blocks share each staged K and V fragment, each block's arithmetic its own -[[kernel]] void lfm2_attn_mma(const device bfloat* queries [[buffer(0)]], const device bfloat* keys [[buffer(1)]], - const device bfloat* values [[buffer(2)]], constant lfm2_attn_args& p [[buffer(3)]], - device float* pm [[buffer(4)]], device float* pl [[buffer(5)]], device float* po [[buffer(6)]], - constant uint& rows [[buffer(7)]], uint3 tg [[threadgroup_position_in_grid]], - uint tid [[thread_index_in_threadgroup]], uint sg [[simdgroup_index_in_threadgroup]], - uint lane [[thread_index_in_simdgroup]]) { - threadgroup uint ks[2][HD * (KT / 2)]; - threadgroup bfloat vs[2][KB * KP]; +inline void lfm2_attn_rows(const device bfloat* queries, const device bfloat* keys, const device bfloat* values, + lfm2_attn_args p, device float* pm, device float* pl, device float* po, uint rows, uint3 tg, + uint tid, uint sg, uint lane, threadgroup uint (*ks)[HD * (KT / 2)], + threadgroup bfloat (*vs)[KB * KP]) { // Row tiles first in the grid: the threadgroups in flight walk the same keys together (one chunk, one kv head). const uint chunk = tg.z / p.subs, sub0 = p.subs == 1 ? 0 : tg.z % SUBS, sub1 = p.subs == 1 ? SUBS : sub0 + 1; const uint kh = tg.y, r0 = tg.x * 8 * RB, head = kh * (QH / KVH) + sg; @@ -337,15 +340,23 @@ template // a simdgroup's RB row blocks share each staged K and V frag for (int d = 0; d < 8; ++d) *(device float2*)(po + at * HD + 8 * d + fn) = O[b][d]; } } +template +[[kernel]] void lfm2_attn_mma(const device bfloat* queries [[buffer(0)]], const device bfloat* keys [[buffer(1)]], + const device bfloat* values [[buffer(2)]], constant lfm2_attn_args& p [[buffer(3)]], + device float* pm [[buffer(4)]], device float* pl [[buffer(5)]], device float* po [[buffer(6)]], + constant uint& rows [[buffer(7)]], uint3 tg [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], uint sg [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + threadgroup uint ks[2][HD * (KT / 2)]; + threadgroup bfloat vs[2][KB * KP]; + lfm2_attn_rows(queries, keys, values, p, pm, pl, po, rows, tg, tid, sg, lane, ks, vs); +} template [[host_name("lfm2_attn_mma")]] [[kernel]] decltype(lfm2_attn_mma<1>) lfm2_attn_mma<1>; template [[host_name("lfm2_attn_mma_r2")]] [[kernel]] decltype(lfm2_attn_mma<2>) lfm2_attn_mma<2>; // The chunks merged in chunk order for (head, row): Y [rows, QH*HD] bf16. -[[kernel]] void lfm2_attn_merge(const device float* pm [[buffer(0)]], const device float* pl [[buffer(1)]], - const device float* po [[buffer(2)]], constant lfm2_attn_args& p [[buffer(3)]], - device bfloat* y [[buffer(4)]], uint2 tg [[threadgroup_position_in_grid]], - uint lane [[thread_index_in_simdgroup]]) { - const uint head = tg.x, r = tg.y; +inline void lfm2_merge_row(const device float* pm, const device float* pl, const device float* po, lfm2_attn_args p, + device bfloat* y, uint head, uint r, uint lane) { const uint count = p.len + r + 1; const uint per = PER; const uint chunks = (count + per - 1) / per; @@ -380,6 +391,12 @@ template [[host_name("lfm2_attn_mma_r2")]] [[kernel]] decltype(lfm2_attn_mma<2>) out[lane] = static_cast(o0 / l); out[lane + 32] = static_cast(o1 / l); } +[[kernel]] void lfm2_attn_merge(const device float* pm [[buffer(0)]], const device float* pl [[buffer(1)]], + const device float* po [[buffer(2)]], constant lfm2_attn_args& p [[buffer(3)]], + device bfloat* y [[buffer(4)]], uint2 tg [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]]) { + lfm2_merge_row(pm, pl, po, p, y, tg.x, tg.y, lane); +} // Router logits [rows, E] in fp32: a simdgroup an expert, lane l over inputs 4l..4l+3 of each 128, then simd_sum. [[kernel]] void lfm2_route_logits(const device bfloat* x [[buffer(0)]], const device float* router [[buffer(1)]], From 793b68e2b66bb77cd38c5715bbf19b8d3d4f77dd Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 10:58:48 +0100 Subject: [PATCH 03/10] build: test-lfm2 runs the LFM2-MoE family's tests alone The family's tests (config, draws, state) ran only inside `zig build test`. test-lfm2 runs the engine's test binary filtered to lfm2_moe, so a change to the family can run them without the whole suite. --- zig/build/metal_targets.zig | 2 ++ 1 file changed, 2 insertions(+) diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index d5ab369dd7..7f1c531bab 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -455,6 +455,8 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s }); test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = helpers })).step); test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = engine })).step); + const lfm2_tests = b.addRunArtifact(b.addTest(.{ .root_module = engine, .filters = &.{"lfm2_moe"} })); + b.step("test-lfm2", "The LFM2-MoE family's tests: config, draws, state and its shared rounds on this Mac's GPU").dependOn(&lfm2_tests.step); test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = lanes })).step); // the fake fabric maps its windows with mmap: the Macs' RDMA, not Windows if (target.result.os.tag != .windows) test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = fabric })).step); From 807aa4d6270b3e01a09898b8c7905b8f414ff002 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 01:27:07 +0100 Subject: [PATCH 04/10] lfm2: a shared round runs each layer's conv and attention once over all its streams A recorded call of several streams launched each stream's own short conv (30 layers) and its qk, attention and merge (10 layers): at 16 one-row streams 960 launches a forward for little work each. lfm2_segments.metal runs each of them once a layer over every row, from a table of the streams' cache addresses, row offsets and key counts (segments.zig); each row's arithmetic is the per-stream kernel's, called through the same inline bodies. A stream of more than 8 rows folds a chunk's sub-spans in its 8-row tiles, as its own 16-row tiles do. Prompt calls and one-stream windows keep their launches. Every cache buffer the table reaches is declared to the encoder: a new stream's caches join the server's residency set only at its next flush, so a kernel that reaches them through the table alone may read them before they are resident. The scratch's attention partials now hold each stream's apart (partialCount: a 32-row scratch's grow fourfold; the served scratch of 2,048 rows already holds them). The contract registers "metal lfm2 stream attention" and "metal lfm2 stream conv" as bit-equal replacements of the per-stream kernels, with float64 attention and a host replay of the conv steps. test-lfm2 checks the table, the row map, the tiles and the declared buffers, and the partials' count at 1, 32 and 2,048 rows. --- zig/kernels/metal/lfm2/lfm2_segments.metal | 66 +++ zig/kernels/metal/sources.zig | 2 +- zig/src/families/lfm2_moe/backend.zig | 4 + zig/src/families/lfm2_moe/forward.zig | 59 ++- zig/src/families/lfm2_moe/kernels.zig | 4 + zig/src/families/lfm2_moe/segments.zig | 103 +++++ zig/src/families/lfm2_moe/state.zig | 25 +- zig/tests/contract/metal_lfm2_streams.zig | 488 +++++++++++++++++++++ zig/tests/contract/metal_main.zig | 11 + 9 files changed, 757 insertions(+), 5 deletions(-) create mode 100644 zig/kernels/metal/lfm2/lfm2_segments.metal create mode 100644 zig/src/families/lfm2_moe/segments.zig create mode 100644 zig/tests/contract/metal_lfm2_streams.zig diff --git a/zig/kernels/metal/lfm2/lfm2_segments.metal b/zig/kernels/metal/lfm2/lfm2_segments.metal new file mode 100644 index 0000000000..4bf25135c7 --- /dev/null +++ b/zig/kernels/metal/lfm2/lfm2_segments.metal @@ -0,0 +1,66 @@ + +// A shared round's streams in one dispatch a kernel (after lfm2.metal): each row runs its own stream's arithmetic. + +// A stream of a recorded call at one layer: its cache rows, its rows from `base`, its key count and sub-spans. +struct lfm2_seg { + device bfloat* keys; + device bfloat* values; + device bfloat* conv; + uint base, rows, len, cap, subs, pad; +}; +// Partials a row owns: every stream's at its first row times the most sub-spans, so no two streams' overlap. +constant constexpr size_t SEG_PART = size_t(QH) * CHUNKS * SUBS; + +// lfm2_conv for every row of a recorded call: row y of stream rowseg[y], its state and snapshot rows. +[[kernel]] void lfm2_conv_segs(const device bfloat* bcx [[buffer(0)]], const device lfm2_seg* segs [[buffer(1)]], + const device bfloat* w [[buffer(2)]], device bfloat* y [[buffer(3)]], + device bfloat* snap [[buffer(4)]], const device uint* rowseg [[buffer(5)]], + uint2 at [[thread_position_in_grid]]) { + if (at.x >= D) return; + const lfm2_seg s = segs[rowseg[at.y]]; + lfm2_conv_row(bcx + size_t(s.base) * 3 * D, s.conv, w, y + size_t(s.base) * D, snap + size_t(s.base) * 2 * D, 0, at.x, + at.y - s.base); +} + +// lfm2_qk for every row of a recorded call, each stream's keys and values into its own cache. +[[kernel]] void lfm2_qk_segs(const device bfloat* q [[buffer(0)]], const device bfloat* k [[buffer(1)]], + const device bfloat* v [[buffer(2)]], const device bfloat* qn [[buffer(3)]], + const device bfloat* kn [[buffer(4)]], constant lfm2_qk_args& a [[buffer(5)]], + device bfloat* queries [[buffer(6)]], const device lfm2_seg* segs [[buffer(7)]], + const device uint* rowseg [[buffer(8)]], uint2 tg [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]]) { + const lfm2_seg s = segs[rowseg[tg.y]]; + const lfm2_qk_args p = {s.len, s.cap, a.eps, a.log2_theta}; + const size_t b = s.base; + lfm2_qk_row(q + b * QH * HD, k + b * KVH * HD, v + b * KVH * HD, qn, kn, p, queries + b * QH * HD, s.keys, s.values, + tg.x, tg.y - s.base, lane); +} + +// lfm2_attn_mma over the streams' 8-row tiles: tile tg.x is tiles[tg.x] = (stream << 16) | the stream's tile. +[[kernel]] void lfm2_attn_segs(const device bfloat* queries [[buffer(0)]], const device lfm2_seg* segs [[buffer(1)]], + const device uint* tiles [[buffer(2)]], constant lfm2_attn_args& a [[buffer(3)]], + device float* pm [[buffer(4)]], device float* pl [[buffer(5)]], + device float* po [[buffer(6)]], uint3 tg [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], uint sg [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + threadgroup uint ks[2][HD * (KT / 2)]; + threadgroup bfloat vs[2][KB * KP]; + const uint item = tiles[tg.x]; + const lfm2_seg s = segs[item >> 16]; + const lfm2_attn_args p = {s.len, s.cap, a.scale, s.subs}; + const size_t part = size_t(s.base) * SEG_PART; + lfm2_attn_rows<1>(queries + size_t(s.base) * QH * HD, s.keys, s.values, p, pm + part, pl + part, po + part * HD, + s.rows, uint3(item & 0xFFFFu, tg.y, tg.z), tid, sg, lane, ks, vs); +} + +// lfm2_attn_merge for every row of a recorded call, from its stream's partials. +[[kernel]] void lfm2_attn_merge_segs(const device float* pm [[buffer(0)]], const device float* pl [[buffer(1)]], + const device float* po [[buffer(2)]], constant lfm2_attn_args& a [[buffer(3)]], + device bfloat* y [[buffer(4)]], const device lfm2_seg* segs [[buffer(5)]], + const device uint* rowseg [[buffer(6)]], uint2 tg [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]]) { + const lfm2_seg s = segs[rowseg[tg.y]]; + const lfm2_attn_args p = {s.len, s.cap, a.scale, s.subs}; + const size_t part = size_t(s.base) * SEG_PART; + lfm2_merge_row(pm + part, pl + part, po + part * HD, p, y + size_t(s.base) * QH * HD, tg.x, tg.y - s.base, lane); +} diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index 7caee6ee67..18662768d7 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -19,7 +19,7 @@ pub const gemma4_prompt_attention = @embedFile("gemma4/prompt_attention.metal"); pub const gemma4_vision = @embedFile("prefill/gemm_nax.metal") ++ @embedFile("gemma4/vision.metal"); /// LFM2-MoE: embedding, norms, the short conv, attention, the route and the experts' glue. -pub const lfm2 = @embedFile("lfm2/lfm2.metal"); +pub const lfm2 = @embedFile("lfm2/lfm2.metal") ++ @embedFile("lfm2/lfm2_segments.metal"); /// LFM2's prompt attention on the tensor units (nax.h inlined at load, after lfm2.metal's text). pub const lfm2_attn_nax = @embedFile("lfm2/lfm2_attn_nax.metal"); diff --git a/zig/src/families/lfm2_moe/backend.zig b/zig/src/families/lfm2_moe/backend.zig index f83935333c..d713c9ec24 100644 --- a/zig/src/families/lfm2_moe/backend.zig +++ b/zig/src/families/lfm2_moe/backend.zig @@ -313,6 +313,10 @@ fn factsFor(width: u32, streams: u32, costs: []const lanes.config.Cost) lanes.Mo }; } +test { + _ = @import("segments.zig"); // a shared round's table, with the family's tests +} + test "after a failed wider-lanes check every forward is one stream's one row, so no window or by-expert twin runs" { const gpa = std.testing.allocator; var one = try lanes.Config.init(gpa, factsFor(1, 8, &.{}), st.window_rows, st.window_rows - 1); diff --git a/zig/src/families/lfm2_moe/forward.zig b/zig/src/families/lfm2_moe/forward.zig index 574e536839..212b5978bc 100644 --- a/zig/src/families/lfm2_moe/forward.zig +++ b/zig/src/families/lfm2_moe/forward.zig @@ -8,6 +8,7 @@ const c = @import("config.zig"); const row = @import("../../core/row_projection.zig"); const tile = @import("../../core/tile_qmm.zig"); const am = @import("../../core/affine_mm.zig"); +const segs = @import("segments.zig"); const Buffer = mtl.Buffer; pub const Segment = struct { cache: *st.Cache, rows: usize }; /// GPU seconds of every call so far (the bench reports a chunk's GPU time beside its wall time). @@ -30,6 +31,8 @@ const Encoder = struct { prompt: bool, /// A recorded window: every row's conv state is kept for a later keep. record: bool, + /// A recorded window of several streams: their conv and attention kernels run once a layer over every row. + plan: ?segs.Plan = null, fn pipe(self: Encoder, comptime key: []const u8) void { self.e.setPipeline(self.m.kernels.get(key)); @@ -129,8 +132,9 @@ pub fn run(m: *Model, s: *st.Scratch, segments: []const Segment, tokens: []const defer pool.pop(); const cb = m.queue.commandBuffer(); const prompt = (!record and head == .last) or total >= tile_rows; - const enc = Encoder{ .e = cb.compute(.serial), .m = m, .s = s, .prompt = prompt, .record = record }; - const e = enc.e; + const e = cb.compute(.serial); + const plan = if (segs.batched(record, segments.len)) segs.prepare(e, s, segments) else null; + const enc = Encoder{ .e = e, .m = m, .s = s, .prompt = prompt, .record = record, .plan = plan }; const embedding = m.weights.embedding; e.setPipeline(m.kernels.embed(embedding.bits)); e.setBuffer(s.ids, 0, 0); @@ -186,6 +190,17 @@ pub fn run(m: *Model, s: *st.Scratch, segments: []const Segment, tokens: []const fn conv(enc: Encoder, s: *st.Scratch, cv: wts.Conv, layer: usize, segments: []const Segment, total: usize) void { const e = enc.e; enc.projection(cv.in, s.x, 0, s.bcx, 0, total); + if (enc.plan != null) { + enc.pipe("lfm2_conv_segs"); + e.setBuffer(s.bcx, 0, 0); + e.setBuffer(s.segs, layer * segs.table_bytes, 1); + enc.tensor(cv.weight, 2); + e.setBuffer(s.y, 0, 3); + e.setBuffer(s.snapshots[layer].?, 0, 4); + e.setBuffer(s.rowseg, 0, 5); + enc.threads(.{ c.hidden, total, 1 }, .{ 256, 1, 1 }); + return enc.projection(cv.out, s.y, 0, s.r, 0, total); + } var base: usize = 0; for (segments) |seg| { enc.pipe("lfm2_conv"); @@ -206,6 +221,10 @@ fn attention(enc: Encoder, s: *st.Scratch, a: wts.Attention, config: *const c.Co enc.projection(a.q, s.x, 0, s.q, 0, total); enc.projection(a.k, s.x, 0, s.k, 0, total); enc.projection(a.v, s.x, 0, s.v, 0, total); + if (enc.plan) |plan| { + attentionSegs(enc, s, a, config, layer, total, plan); + return enc.projection(a.out, s.y, 0, s.r, 0, total); + } var base: usize = 0; for (segments) |seg| { const cache = seg.cache.blocks[layer].attention; @@ -259,6 +278,42 @@ fn attention(enc: Encoder, s: *st.Scratch, a: wts.Attention, config: *const c.Co enc.projection(a.out, s.y, 0, s.r, 0, total); } +/// A recorded call's attention over every stream in one dispatch a kernel: qk, the 8-row tiles, the merge. +fn attentionSegs(enc: Encoder, s: *st.Scratch, a: wts.Attention, config: *const c.Config, layer: usize, total: usize, plan: segs.Plan) void { + const e = enc.e; + const table = layer * segs.table_bytes; + enc.pipe("lfm2_qk_segs"); + e.setBuffer(s.q, 0, 0); + e.setBuffer(s.k, 0, 1); + e.setBuffer(s.v, 0, 2); + enc.tensor(a.q_norm, 3); + enc.tensor(a.k_norm, 4); + e.setBytes(std.mem.asBytes(&QkArgs{ .len = 0, .cap = 0, .eps = c.eps, .log2_theta = std.math.log2(config.theta) }), 5); + e.setBuffer(s.queries, 0, 6); + e.setBuffer(s.segs, table, 7); + e.setBuffer(s.rowseg, 0, 8); + enc.groups(.{ c.query_heads + c.kv_heads, total, 1 }, .{ 32, 1, 1 }); + const args = AttnArgs{ .len = 0, .cap = 0, .scale = 1.0 / @sqrt(@as(f32, c.head_dim)), .subs = 0 }; + enc.pipe("lfm2_attn_segs"); + e.setBuffer(s.queries, 0, 0); + e.setBuffer(s.segs, table, 1); + e.setBuffer(s.tiles, 0, 2); + e.setBytes(std.mem.asBytes(&args), 3); + e.setBuffer(s.pm, 0, 4); + e.setBuffer(s.pl, 0, 5); + e.setBuffer(s.po, 0, 6); + enc.groups(.{ plan.tiles, c.kv_heads, plan.used }, .{ 128, 1, 1 }); + enc.pipe("lfm2_attn_merge_segs"); + e.setBuffer(s.pm, 0, 0); + e.setBuffer(s.pl, 0, 1); + e.setBuffer(s.po, 0, 2); + e.setBytes(std.mem.asBytes(&args), 3); + e.setBuffer(s.y, 0, 4); + e.setBuffer(s.segs, table, 5); + e.setBuffer(s.rowseg, 0, 6); + enc.groups(.{ c.query_heads, total, 1 }, .{ 32, 1, 1 }); +} + /// True when the picks' combine is left to the next norm (decode rows). fn moe(enc: Encoder, s: *st.Scratch, x: wts.Moe, scale: f32, total: usize) bool { const e = enc.e; diff --git a/zig/src/families/lfm2_moe/kernels.zig b/zig/src/families/lfm2_moe/kernels.zig index 48da29433f..aa1e798fc0 100644 --- a/zig/src/families/lfm2_moe/kernels.zig +++ b/zig/src/families/lfm2_moe/kernels.zig @@ -23,6 +23,10 @@ const names = [_][:0]const u8{ "lfm2_sort", "lfm2_combine_sorted", "lfm2_argmax", + "lfm2_conv_segs", + "lfm2_qk_segs", + "lfm2_attn_segs", + "lfm2_attn_merge_segs", }; pub const Kernels = struct { diff --git a/zig/src/families/lfm2_moe/segments.zig b/zig/src/families/lfm2_moe/segments.zig new file mode 100644 index 0000000000..aa2a7b810f --- /dev/null +++ b/zig/src/families/lfm2_moe/segments.zig @@ -0,0 +1,103 @@ +//! A recorded call's streams in one dispatch a kernel and layer: each layer's stream table, the rows' streams, tiles. +const std = @import("std"); +const mtl = @import("metal"); +const st = @import("state.zig"); +const c = @import("config.zig"); + +/// lfm2_segments.metal's lfm2_seg: a stream's cache at one layer, its rows from `base`, key count and sub-spans. +pub const Seg = extern struct { keys: u64 = 0, values: u64 = 0, conv: u64 = 0, base: u32, rows: u32, len: u32, cap: u32, subs: u32, pad: u32 = 0 }; + +/// Bytes of one layer's table in Scratch.segs. +pub const table_bytes = st.batch_rows * @sizeOf(Seg); + +/// The attention's 8-row tiles over every stream and the grid depth (chunks times sub-spans) the deepest needs. +pub const Plan = struct { tiles: u32 = 0, used: u32 = 0 }; + +/// Whether a call takes these kernels: a recorded window of more than one stream. +pub fn batched(record: bool, streams: usize) bool { + return record and streams > 1; +} + +/// The streams' tables, row map and tiles into the scratch, and every cache buffer the kernels reach declared on `e`. +pub fn prepare(e: anytype, s: *st.Scratch, streams: anytype) Plan { + var plan: Plan = .{}; + const rowseg = s.rowseg.slice(u32, st.batch_rows); + const tiles = s.tiles.slice(u32, st.batch_rows); + var base: usize = 0; + for (streams, 0..) |seg, j| { + // A stream of more than split_rows rows folds a chunk's sub-spans in its tiles, as its own call would. + const subs: u32 = if (seg.rows > st.split_rows) 1 else st.subs; + for (0..(seg.rows + 7) / 8) |t| { + tiles[plan.tiles] = @intCast((j << 16) | t); + plan.tiles += 1; + } + const used: u32 = @intCast((seg.cache.len + seg.rows + st.chunk_keys - 1) / st.chunk_keys * subs); + plan.used = @max(plan.used, used); + @memset(rowseg[base..][0..seg.rows], @intCast(j)); + for (seg.cache.blocks, 0..) |block, layer| { + var entry: Seg = .{ .base = @intCast(base), .rows = @intCast(seg.rows), .len = @intCast(seg.cache.len), .cap = @intCast(seg.cache.capacity), .subs = subs }; + switch (block) { + .conv => |b| entry.conv = b.gpuAddress(), + .attention => |a| { + entry.keys = a.keys.gpuAddress(); + entry.values = a.values.gpuAddress(); + }, + } + s.segs.slice(Seg, c.layers * st.batch_rows)[layer * st.batch_rows + j] = entry; + } + // The kernels reach the caches through the table: each buffer is declared, in a residency set or not. + for (seg.cache.memory.buffers.items) |b| e.useResource(b, mtl.ResourceUsage.read | mtl.ResourceUsage.write); + base += seg.rows; + } + return plan; +} + +/// Buffers a call declares to its encoder (useResource), in order. +const Declared = struct { + ids: [8 * c.layers]?*anyopaque = undefined, + count: usize = 0, + pub fn useResource(self: *Declared, b: mtl.Buffer, _: usize) void { + self.ids[self.count] = @ptrCast(b.id); + self.count += 1; + } + fn has(self: *const Declared, b: mtl.Buffer) bool { + return std.mem.indexOfScalar(?*anyopaque, self.ids[0..self.count], @ptrCast(b.id)) != null; + } +}; + +test "lfm2_moe: a shared round's table reaches only cache buffers it declares, and maps its rows and tiles" { + const gpa = std.testing.allocator; + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + var attention: [c.layers]bool = @splat(false); + for (&attention, 0..) |*a, i| a.* = i % 4 == 2; + var s = try st.Scratch.init(gpa, device, &attention, st.batch_rows); + defer s.deinit(); + var caches: [2]st.Cache = undefined; + for (&caches) |*cache| cache.* = try st.Cache.init(gpa, device, &attention, 64); + defer for (&caches) |*cache| cache.deinit(); + caches[0].len = 5; + caches[1].len = 40; + const streams = [_]struct { cache: *st.Cache, rows: usize }{ .{ .cache = &caches[0], .rows = 9 }, .{ .cache = &caches[1], .rows = 1 } }; + var declared: Declared = .{}; + const plan = prepare(&declared, &s, &streams); + // Each cache's every buffer: a buffer the kernels reach through the table only is not resident unless declared. + var want: usize = 0; + for (&caches) |*cache| for (cache.memory.buffers.items) |b| { + try std.testing.expect(declared.has(b)); + want += 1; + }; + try std.testing.expectEqual(want, declared.count); + try std.testing.expectEqual(@as(u32, 3), plan.tiles); + try std.testing.expectEqual(@as(u32, st.subs), plan.used); + try std.testing.expectEqualSlices(u32, &.{ 0, 1, 1 << 16 }, s.tiles.slice(u32, 3)); + try std.testing.expectEqualSlices(u32, &.{ 0, 0, 0, 0, 0, 0, 0, 0, 0, 1 }, s.rowseg.slice(u32, 10)); + const table = s.segs.slice(Seg, c.layers * st.batch_rows); + for (attention, 0..) |a, layer| for (&caches, 0..) |*cache, j| { + const e = table[layer * st.batch_rows + j]; + try std.testing.expectEqual(@as(u32, if (j == 0) 0 else 9), e.base); + try std.testing.expectEqual(@as(u32, if (j == 0) 1 else st.subs), e.subs); + try std.testing.expectEqual(@as(u32, @intCast(cache.len)), e.len); + if (a) try std.testing.expectEqual(cache.blocks[layer].attention.keys.gpuAddress(), e.keys) else try std.testing.expectEqual(cache.blocks[layer].conv.gpuAddress(), e.conv); + }; +} diff --git a/zig/src/families/lfm2_moe/state.zig b/zig/src/families/lfm2_moe/state.zig index 646510581f..3242d8bf26 100644 --- a/zig/src/families/lfm2_moe/state.zig +++ b/zig/src/families/lfm2_moe/state.zig @@ -144,6 +144,10 @@ pub const Scratch = struct { /// Each head row's greedy pick and nonfinite flag (u32 pairs), from a forward whose head takes every row. best: mtl.Buffer, snapshots: [c.layers]?mtl.Buffer = @splat(null), + /// A recorded call's streams at every layer, each row's stream and the attention's row tiles (segments.zig). + segs: mtl.Buffer, + rowseg: mtl.Buffer, + tiles: mtl.Buffer, pub fn init(gpa: std.mem.Allocator, device: mtl.Device, attention: *const [c.layers]bool, rows: usize) !Scratch { var out: Scratch = undefined; @@ -159,8 +163,7 @@ pub const Scratch = struct { out.bcx = try mem.alloc(rows * 3 * c.hidden * 2); out.k = try mem.alloc(rows * c.kv_heads * c.head_dim * 2); out.v = try mem.alloc(rows * c.kv_heads * c.head_dim * 2); - const slice_rows: usize = @min(rows, attn_slice); - const partials = @max(slice_rows, split_rows * subs) * c.query_heads * chunks; + const partials = partialCount(rows); out.pm = try mem.alloc(partials * 4); out.pl = try mem.alloc(partials * 4); out.po = try mem.alloc(partials * c.head_dim * 4); @@ -182,6 +185,9 @@ pub const Scratch = struct { const logit_rows: usize = @min(rows, batch_rows); out.logits = try mem.alloc(logit_rows * c.vocab * 2); out.best = try mem.alloc(logit_rows * 2 * 4); + out.segs = try mem.alloc(c.layers * batch_rows * @sizeOf(@import("segments.zig").Seg)); + out.rowseg = try mem.alloc(batch_rows * 4); + out.tiles = try mem.alloc(batch_rows * 4); for (&out.snapshots, attention) |*s, a| if (!a) { // a recorded window's rows, or a prompt call's last row s.* = try mem.alloc(logit_rows * conv_bytes); }; @@ -193,6 +199,14 @@ pub const Scratch = struct { } }; +/// Attention partials a scratch of `rows` holds: a prompt slice's, and a recorded call's streams kept apart. +pub fn partialCount(rows: usize) usize { + const slice_rows: usize = @min(rows, attn_slice); + // usize: @min with a constant narrows its type, and `most * subs` would overflow it. + const most: usize = @min(rows, batch_rows); + return @max(slice_rows, split_rows * subs, most * subs) * c.query_heads * chunks; +} + /// The served window: --context, else the checkpoint's, never past `max_context`; null when --context is outside it. pub fn window(checkpoint: usize, requested: ?i64) ?usize { const most = @min(checkpoint, max_context); @@ -200,6 +214,13 @@ pub fn window(checkpoint: usize, requested: ?i64) ?usize { return if (w >= 1 and w <= most) @as(usize, @intCast(w)) else null; } +test "a shared round's 32 rows keep their attention partials apart" { + // 32 rows of 8 sub-spans each, every row's 32 heads over 32 chunks; a prompt chunk's rows hold more. + try std.testing.expectEqual(@as(usize, 32 * subs * c.query_heads * chunks), partialCount(32)); + try std.testing.expectEqual(@as(usize, split_rows * subs * c.query_heads * chunks), partialCount(1)); + try std.testing.expectEqual(@as(usize, 2048 * c.query_heads * chunks), partialCount(2048)); +} + test "the served window stays inside the attention's 131,072 positions" { try std.testing.expectEqual(@as(usize, 131072), max_context); try std.testing.expectEqual(@as(?usize, 131072), window(131072, null)); diff --git a/zig/tests/contract/metal_lfm2_streams.zig b/zig/tests/contract/metal_lfm2_streams.zig new file mode 100644 index 0000000000..0a46f7a526 --- /dev/null +++ b/zig/tests/contract/metal_lfm2_streams.zig @@ -0,0 +1,488 @@ +//! LFM2's shared-round kernels (lfm2_segments.metal): every stream's conv or attention rows in one dispatch a kernel. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); +const c = contract.spec; +const g = contract.gates; +const rig_ = @import("rig.zig"); + +const d = 2048; +const qh = 32; +const kvh = 8; +const hd = 64; +const subs = 8; // sub-spans a 4,096-key chunk for a stream of at most 8 rows (state.subs) +const chunks = 32; +const chunk_keys = 4096; +const most = 32; // a shared round's rows (state.batch_rows) +const part = qh * chunks * subs; // partials a row owns (the kernels' SEG_PART) +pub const Which = enum { attention, conv }; +const names = [_][]const u8{ "metal lfm2 stream attention", "metal lfm2 stream conv" }; + +/// rows: a round's rows over its streams (sizes from `sizes` in turn); k: the first stream's cached keys. +const attn_shapes = [_]c.Shape{ .{ .rows = 2, .n = qh * hd, .k = 5 }, .{ .rows = 13, .n = qh * hd, .k = 300 }, .{ .rows = 32, .n = qh * hd, .k = 4100 } }; +const conv_shapes = [_]c.Shape{ .{ .rows = 2, .n = d, .k = 3 }, .{ .rows = 13, .n = d, .k = 3 }, .{ .rows = 32, .n = d, .k = 3 } }; +const sizes = [_]usize{ 1, 9, 3, 16, 2, 8, 5 }; + +fn none(_: c.Shape) usize { + return 0; +} + +pub fn spec(w: Which) c.Spec { + return switch (w) { + .attention => .{ .name = names[0], .accum = .none, .depth = none, .out = .bf16, .served = &attn_shapes, .replaces = "lfm2_qk, lfm2_attn_mma and lfm2_attn_merge a stream", .nan = .{ .stated = "a NaN key of one stream's kv head makes that stream's heads of it NaN, every other output stays finite" }, .sums = .{ .f16 = 0, .bf16 = 0, .why = "the per-stream kernels' bits; the check bounds outputs by float64 attention" } }, + .conv => .{ .name = names[1], .accum = .none, .depth = none, .out = .bf16, .served = &conv_shapes, .replaces = "lfm2_conv a stream", .nan = .{ .stated = "outputs equal the host's replay of the kernel's bf16 steps, NaN for NaN" }, .sums = .{ .f16 = 0, .bf16 = 0, .why = "three fp32 products a channel, replayed exactly on the host" } }, + }; +} + +/// lfm2_segments.metal's lfm2_seg. +const Seg = extern struct { keys: u64 = 0, values: u64 = 0, conv: u64 = 0, base: u32, rows: u32, len: u32, cap: u32, subs: u32, pad: u32 = 0 }; + +fn value(i: usize, seed: u64, scale: f32) u16 { + const h: u64 = (@as(u64, i) *% 0x9e3779b97f4a7c15 +% seed *% 0xbf58476d1ce4e5b9) >> 40; + return c.toBf16((@as(f32, @floatFromInt(h % 2048)) / 1024 - 1) * scale); +} + +pub const Probe = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + which: Which, + pipes: [8]mtl.Pipeline, + made: usize = 0, + rows: usize = 0, + nan: bool = false, + first_len: usize = 0, + streams: [most]usize = undefined, // rows a stream + count_streams: usize = 0, + host: []u16 = &.{}, // attention: the cached keys and values as prepared; conv: bcx, then the states + weights: []u16 = &.{}, + args: [6]rig_.Ref = undefined, + caches: [2][most]rig_.Ref = undefined, // attention: each stream's [keys | values], batched and per stream + y: rig_.Ref = undefined, + queries: rig_.Ref = undefined, + work: [9]rig_.Ref = undefined, // segs, rowseg, tiles, pm, pl, po; the per-stream queries, y and partials' base + golden: []u8 = &.{}, + snap: rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, w: Which) !Probe { + const list: []const [:0]const u8 = switch (w) { + .attention => &.{ "lfm2_qk_segs", "lfm2_attn_segs", "lfm2_attn_merge_segs", "lfm2_qk", "lfm2_attn_mma", "lfm2_attn_mma_r2", "lfm2_attn_merge" }, + .conv => &.{ "lfm2_conv_segs", "lfm2_conv" }, + }; + var p: Probe = .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .which = w, .pipes = undefined }; + errdefer for (p.pipes[0..p.made]) |x| x.deinit(); + for (list) |n| { + p.pipes[p.made] = try rig_.pipeline(device, sources.lfm2, n); + p.made += 1; + } + return p; + } + + fn drop(p: *Probe) void { + p.a.free(p.golden); + p.golden = &.{}; + p.a.free(p.host); + p.a.free(p.weights); + p.host = &.{}; + p.weights = &.{}; + p.rig.reset(); + } + + pub fn deinit(p: *Probe) void { + p.drop(); + p.rig.deinit(); + for (p.pipes[0..p.made]) |x| x.deinit(); + } + + const vt = g.hooks(Probe); + pub fn kernel(p: *Probe) g.Kernel { + return .{ .ctx = p, .vt = &vt, .spec = spec(p.which) }; + } + + fn len(p: *const Probe, j: usize) usize { + return p.first_len + 37 * j; + } + + fn cap(p: *const Probe, j: usize) usize { + return (p.len(j) + p.streams[j] + 63) / 64 * 64 + 32; + } + + fn base(p: *const Probe, j: usize) usize { + var b: usize = 0; + for (p.streams[0..j]) |r| b += r; + return b; + } + + pub fn prepare(p: *Probe, s: c.Shape, f: g.Fill, seed: u64) !bool { + if (s.rows < 2 or s.rows > most or s.n != @as(usize, if (p.which == .attention) qh * hd else d)) return false; + if (p.which == .attention and (s.k == 0 or s.k + 37 * most + most > chunks * chunk_keys)) return false; + p.drop(); + p.rows = s.rows; + p.first_len = s.k; + p.nan = f == .nonfinite; + p.count_streams = 0; + var left = s.rows; + while (left > 0) : (p.count_streams += 1) { + const r = @min(left, sizes[p.count_streams % sizes.len]); + p.streams[p.count_streams] = r; + left -= r; + } + return if (p.which == .attention) p.prepareAttention(f, seed) else p.prepareConv(f, seed); + } + + fn prepareAttention(p: *Probe, f: g.Fill, seed: u64) !bool { + const n = p.rows; + var q = try p.a.alloc(u16, n * (qh + 2 * kvh) * hd); + defer p.a.free(q); + for (q, 0..) |*x, i| x.* = value(i, seed, if (f == .stress) 4 else 1); + p.weights = try p.a.alloc(u16, 2 * hd); + for (p.weights, 0..) |*x, i| x.* = value(i, seed +% 9, 1.5); + p.args[0] = try p.rig.input(std.mem.sliceAsBytes(q[0 .. n * qh * hd])); + p.args[1] = try p.rig.input(std.mem.sliceAsBytes(q[n * qh * hd ..][0 .. n * kvh * hd])); + p.args[2] = try p.rig.input(std.mem.sliceAsBytes(q[n * (qh + kvh) * hd ..][0 .. n * kvh * hd])); + p.args[3] = try p.rig.input(std.mem.sliceAsBytes(p.weights)); + var total: usize = 0; + for (0..p.count_streams) |j| total += 2 * kvh * p.cap(j) * hd; + p.host = try p.a.alloc(u16, total); + var at: usize = 0; + for (0..p.count_streams) |j| { + const kv = p.host[at..][0 .. 2 * kvh * p.cap(j) * hd]; + for (kv, 0..) |*x, i| x.* = value(i, seed +% 3 +% j, 1); + if (p.nan and j == 0) kv[(0 * p.cap(j) + p.len(j) / 2) * hd + 7] = 0x7fc0; // stream 0, kv head 0 + for (0..2) |copy| p.caches[copy][j] = try p.rig.scratch(kv.len * 2); + at += kv.len; + } + p.y = try p.rig.output(n * d * 2); + p.queries = try p.rig.scratch(n * d * 2); + p.work[0] = try p.rig.scratch(most * @sizeOf(Seg)); + p.work[1] = try p.rig.scratch(most * 4); + p.work[2] = try p.rig.scratch(most * 4); + p.work[3] = try p.rig.scratch(n * part * 4); + p.work[4] = try p.rig.scratch(n * part * 4); + p.work[5] = try p.rig.scratch(n * part * hd * 4); + p.work[6] = try p.rig.scratch(n * d * 2); + p.work[7] = try p.rig.scratch(n * d * 2); + return true; + } + + fn prepareConv(p: *Probe, f: g.Fill, seed: u64) !bool { + const n = p.rows; + p.host = try p.a.alloc(u16, n * 3 * d + p.count_streams * 2 * d); + for (p.host, 0..) |*x, i| x.* = value(i, seed, if (f == .stress) 8 else 1); + if (p.nan) { + p.host[1 * 3 * d + 11] = 0x7fc0; // row 1's B at channel 11 + p.host[n * 3 * d + 2 * d + 5] = 0x7f80; // stream 1's state at channel 5 + } + p.weights = try p.a.alloc(u16, d * 3); + for (p.weights, 0..) |*x, i| x.* = value(i, seed +% 5, 1); + p.args[0] = try p.rig.input(std.mem.sliceAsBytes(p.host[0 .. n * 3 * d])); + p.args[1] = try p.rig.input(std.mem.sliceAsBytes(p.weights)); + for (0..p.count_streams) |j| p.caches[0][j] = try p.rig.input(std.mem.sliceAsBytes(p.host[n * 3 * d + j * 2 * d ..][0 .. 2 * d])); + p.y = try p.rig.output(n * d * 2); + p.snap = try p.rig.output(n * 2 * d * 2); + return true; + } + + pub fn count(p: *Probe) usize { + return p.rows * d * @as(usize, if (p.which == .conv) 3 else 1); + } + + /// Each stream's cache as prepared (its cached keys before the call's). + fn fill(p: *Probe, copy: usize) void { + var at: usize = 0; + for (0..p.count_streams) |j| { + const n = 2 * kvh * p.cap(j) * hd; + @memcpy(std.mem.bytesAsSlice(u16, p.rig.bytes(p.caches[copy][j], n * 2)), p.host[at..][0..n]); + at += n; + } + } + + fn table(p: *Probe, copy: usize) ![]Seg { + const t = try p.a.alloc(Seg, most); + for (0..p.count_streams) |j| { + const r = p.caches[copy][j]; + const addr = r.buf.gpuAddress() + r.off; + const wide = p.streams[j] > 8; + t[j] = .{ .base = @intCast(p.base(j)), .rows = @intCast(p.streams[j]), .len = @intCast(p.len(j)), .cap = @intCast(p.cap(j)), .subs = if (wide) 1 else subs }; + if (p.which == .attention) { + t[j].keys = addr; + t[j].values = addr + kvh * p.cap(j) * hd * 2; + } else t[j].conv = addr; + } + return t; + } + + pub fn run(p: *Probe) !void { + if (p.which == .conv) return p.runConv(p.y, p.snap, false); + p.fill(0); + const n = p.rows; + const t = try p.table(0); + defer p.a.free(t); + var rowseg: [most]u32 = undefined; + var tiles: [most]u32 = undefined; + var count_tiles: u32 = 0; + var used: u32 = 0; + for (0..p.count_streams) |j| { + @memset(rowseg[p.base(j)..][0..p.streams[j]], @intCast(j)); + for (0..(p.streams[j] + 7) / 8) |i| { + tiles[count_tiles] = @intCast((j << 16) | i); + count_tiles += 1; + } + const s: usize = if (p.streams[j] > 8) 1 else subs; + used = @max(used, @as(u32, @intCast((p.len(j) + p.streams[j] + chunk_keys - 1) / chunk_keys * s))); + } + const segs = p.work[0]; + @memcpy(p.rig.bytes(segs, most * @sizeOf(Seg)), std.mem.sliceAsBytes(t)); + const rs = p.work[1]; + @memcpy(p.rig.bytes(rs, most * 4), std.mem.asBytes(&rowseg)); + const ts = p.work[2]; + @memcpy(p.rig.bytes(ts, most * 4), std.mem.asBytes(&tiles)); + const queries = p.queries; + const pm = p.work[3]; + const pl = p.work[4]; + const po = p.work[5]; + const pass = p.rig.begin(); + const e = pass.e; + for (p.caches[0][0..p.count_streams]) |r| e.useResource(r.buf, mtl.ResourceUsage.read | mtl.ResourceUsage.write); + e.setPipeline(p.pipes[0]); + for (p.args[0..3], 0..) |r, i| e.setBuffer(r.buf, r.off, i); + e.setBuffer(p.args[3].buf, p.args[3].off, 3); + e.setBuffer(p.args[3].buf, p.args[3].off + hd * 2, 4); + const qk = [4]u32{ 0, 0, @bitCast(@as(f32, 1e-5)), @bitCast(std.math.log2(@as(f32, 1e6))) }; + e.setBytes(std.mem.asBytes(&qk), 5); + e.setBuffer(queries.buf, queries.off, 6); + e.setBuffer(segs.buf, segs.off, 7); + e.setBuffer(rs.buf, rs.off, 8); + e.dispatchGroups(mtl.Size.of(qh + kvh, n, 1), mtl.Size.of(32, 1, 1)); + const args = [4]u32{ 0, 0, @bitCast(1.0 / @sqrt(@as(f32, hd))), 0 }; + e.setPipeline(p.pipes[1]); + e.setBuffer(queries.buf, queries.off, 0); + e.setBuffer(segs.buf, segs.off, 1); + e.setBuffer(ts.buf, ts.off, 2); + e.setBytes(std.mem.asBytes(&args), 3); + for ([_]rig_.Ref{ pm, pl, po }, 4..) |r, i| e.setBuffer(r.buf, r.off, i); + e.dispatchGroups(mtl.Size.of(count_tiles, kvh, used), mtl.Size.of(128, 1, 1)); + e.setPipeline(p.pipes[2]); + for ([_]rig_.Ref{ pm, pl, po }, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + e.setBytes(std.mem.asBytes(&args), 3); + e.setBuffer(p.y.buf, p.y.off, 4); + e.setBuffer(segs.buf, segs.off, 5); + e.setBuffer(rs.buf, rs.off, 6); + e.dispatchGroups(mtl.Size.of(qh, n, 1), mtl.Size.of(32, 1, 1)); + try p.rig.finish(pass); + } + + /// conv: the batched kernel, or (`alone`) lfm2_conv a stream, into y and snap. + fn runConv(p: *Probe, y: rig_.Ref, snap: rig_.Ref, alone: bool) !void { + const pass = p.rig.begin(); + const e = pass.e; + if (alone) { + e.setPipeline(p.pipes[1]); + for (0..p.count_streams) |j| { + const b = p.base(j); + e.setBuffer(p.args[0].buf, p.args[0].off + b * 3 * d * 2, 0); + e.setBuffer(p.caches[0][j].buf, p.caches[0][j].off, 1); + e.setBuffer(p.args[1].buf, p.args[1].off, 2); + e.setBuffer(y.buf, y.off + b * d * 2, 3); + e.setBuffer(snap.buf, snap.off + b * 2 * d * 2, 4); + e.setValue(@as(u32, 0), 5); + e.dispatchThreads(mtl.Size.of(d, p.streams[j], 1), mtl.Size.of(256, 1, 1)); + } + return p.rig.finish(pass); + } + const t = try p.table(0); + defer p.a.free(t); + var rowseg: [most]u32 = undefined; + for (0..p.count_streams) |j| @memset(rowseg[p.base(j)..][0..p.streams[j]], @intCast(j)); + const segs = try p.rig.scratch(most * @sizeOf(Seg)); + @memcpy(p.rig.bytes(segs, most * @sizeOf(Seg)), std.mem.sliceAsBytes(t)); + const rs = try p.rig.scratch(most * 4); + @memcpy(p.rig.bytes(rs, most * 4), std.mem.asBytes(&rowseg)); + for (p.caches[0][0..p.count_streams]) |r| e.useResource(r.buf, mtl.ResourceUsage.read); + e.setPipeline(p.pipes[0]); + e.setBuffer(p.args[0].buf, p.args[0].off, 0); + e.setBuffer(segs.buf, segs.off, 1); + e.setBuffer(p.args[1].buf, p.args[1].off, 2); + e.setBuffer(y.buf, y.off, 3); + e.setBuffer(snap.buf, snap.off, 4); + e.setBuffer(rs.buf, rs.off, 5); + e.dispatchThreads(mtl.Size.of(d, p.rows, 1), mtl.Size.of(256, 1, 1)); + try p.rig.finish(pass); + } + + /// The per-stream kernels' bytes on the same inputs. + pub fn replaced(p: *Probe) ![]const u8 { + const n = p.rows; + if (p.which == .conv) { + const y = try p.rig.scratch(n * d * 2); + const snap = try p.rig.scratch(n * 2 * d * 2); + try p.runConv(y, snap, true); + const out = try p.a.alloc(u8, n * 3 * d * 2); + @memcpy(out[0 .. n * d * 2], p.rig.bytes(y, n * d * 2)); + @memcpy(out[n * d * 2 ..], p.rig.bytes(snap, n * 2 * d * 2)); + return p.keep(out); + } + p.fill(1); + const queries = p.work[6]; + const y = p.work[7]; + const pm = p.work[3]; + const pl = p.work[4]; + const po = p.work[5]; + const pass = p.rig.begin(); + const e = pass.e; + for (0..p.count_streams) |j| { + const b = p.base(j); + const rows = p.streams[j]; + const cache = p.caches[1][j]; + const kv_bytes = kvh * p.cap(j) * hd * 2; + e.setPipeline(p.pipes[3]); + e.setBuffer(p.args[0].buf, p.args[0].off + b * qh * hd * 2, 0); + e.setBuffer(p.args[1].buf, p.args[1].off + b * kvh * hd * 2, 1); + e.setBuffer(p.args[2].buf, p.args[2].off + b * kvh * hd * 2, 2); + e.setBuffer(p.args[3].buf, p.args[3].off, 3); + e.setBuffer(p.args[3].buf, p.args[3].off + hd * 2, 4); + const qk = [4]u32{ @intCast(p.len(j)), @intCast(p.cap(j)), @bitCast(@as(f32, 1e-5)), @bitCast(std.math.log2(@as(f32, 1e6))) }; + e.setBytes(std.mem.asBytes(&qk), 5); + e.setBuffer(queries.buf, queries.off + b * d * 2, 6); + e.setBuffer(cache.buf, cache.off, 7); + e.setBuffer(cache.buf, cache.off + kv_bytes, 8); + e.dispatchGroups(mtl.Size.of(qh + kvh, rows, 1), mtl.Size.of(32, 1, 1)); + const wide = rows > 8; + const s: u32 = if (wide) 1 else subs; + const args = [4]u32{ @intCast(p.len(j)), @intCast(p.cap(j)), @bitCast(1.0 / @sqrt(@as(f32, hd))), s }; + e.setPipeline(p.pipes[if (wide) 5 else 4]); + e.setBuffer(queries.buf, queries.off + b * d * 2, 0); + e.setBuffer(cache.buf, cache.off, 1); + e.setBuffer(cache.buf, cache.off + kv_bytes, 2); + e.setBytes(std.mem.asBytes(&args), 3); + for ([_]rig_.Ref{ pm, pl, po }, 4..) |r, i| e.setBuffer(r.buf, r.off, i); + e.setValue(@as(u32, @intCast(rows)), 7); + const per: usize = if (wide) 16 else 8; + const used = (p.len(j) + rows + chunk_keys - 1) / chunk_keys * s; + e.dispatchGroups(mtl.Size.of((rows + per - 1) / per, kvh, used), mtl.Size.of(128, 1, 1)); + e.setPipeline(p.pipes[6]); + for ([_]rig_.Ref{ pm, pl, po }, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + e.setBytes(std.mem.asBytes(&args), 3); + e.setBuffer(y.buf, y.off + b * d * 2, 4); + e.dispatchGroups(mtl.Size.of(qh, rows, 1), mtl.Size.of(32, 1, 1)); + } + try p.rig.finish(pass); + return p.keep(try p.a.dupe(u8, p.rig.bytes(y, n * d * 2))); + } + + fn keep(p: *Probe, bytes: []u8) []const u8 { + p.a.free(p.golden); + p.golden = bytes; + return bytes; + } + + pub fn reference(_: *Probe, want: []f64, mag: []f64) !void { + @memset(want, 0); // the check computes its own references + @memset(mag, 0); + } + + pub fn check(p: *Probe, _: g.Gate) !?[]const u8 { + return if (p.which == .conv) p.checkConv() else p.checkAttention(); + } + + /// Each output against float64 attention over its stream's keys, from the queries and cache the call left. + fn checkAttention(p: *Probe) !?[]const u8 { + const out = std.mem.bytesAsSlice(u16, p.rig.bytes(p.y, p.rows * d * 2)); + const queries = std.mem.bytesAsSlice(u16, p.rig.bytes(p.queries, p.rows * d * 2)); + const w = try p.a.alloc(f64, chunks * chunk_keys); + defer p.a.free(w); + for (0..p.count_streams) |j| { + const cp = p.cap(j); + const kv = std.mem.bytesAsSlice(u16, p.rig.bytes(p.caches[0][j], 2 * kvh * cp * hd * 2)); + for (0..p.streams[j]) |r| for (0..qh) |h| { + const row = p.base(j) + r; + const kh = h / (qh / kvh); + const keys = p.len(j) + r + 1; + var top: f64 = -std.math.inf(f64); + for (0..keys) |k| { + var dot: f64 = 0; + for (0..hd) |i| dot += c.fromBf16(queries[(row * qh + h) * hd + i]) * c.fromBf16(kv[(kh * cp + k) * hd + i]); + w[k] = dot / @sqrt(@as(f64, hd)); + top = @max(top, w[k]); + } + var sum: f64 = 0; + for (w[0..keys]) |*x| { + x.* = @exp(x.* - top); + sum += x.*; + } + for (0..hd) |i| { + var o: f64 = 0; + var big: f64 = 0; + for (0..keys) |k| { + const v: f64 = c.fromBf16(kv[((kvh + kh) * cp + k) * hd + i]); + o += w[k] / sum * v; + big = @max(big, @abs(v)); + } + const got: f64 = c.fromBf16(out[(row * qh + h) * hd + i]); + if (p.nan and j == 0 and kh == 0) { + if (!std.math.isNan(got)) return "a head behind its stream's NaN key gave a number"; + continue; + } + if (!std.math.isFinite(got)) return "a finite head gave a non-finite output"; + if (@abs(got - o) > big * 0x1p-7 + @abs(o) * 0x1p-8) return "an output is past its float64 bound"; + } + }; + } + return null; + } + + /// Each output against the host's replay of lfm2_conv's steps (bf16 products, fp32 taps, the bf16 gate). + fn checkConv(p: *Probe) !?[]const u8 { + const n = p.rows; + const y = std.mem.bytesAsSlice(u16, p.rig.bytes(p.y, n * d * 2)); + const snap = std.mem.bytesAsSlice(u16, p.rig.bytes(p.snap, n * 2 * d * 2)); + for (0..p.count_streams) |j| { + const b = p.base(j); + const state = p.host[n * 3 * d + j * 2 * d ..][0 .. 2 * d]; + for (0..p.streams[j]) |r| for (0..d) |ch| { + var bx: [3]u16 = undefined; + for (0..3) |tap| { + const src = @as(isize, @intCast(r + tap)) - 2; + if (src >= 0) { + const row = p.host[(b + @as(usize, @intCast(src))) * 3 * d ..]; + bx[tap] = c.toBf16(c.fromBf16(row[ch]) * c.fromBf16(row[2 * d + ch])); + } else bx[tap] = state[@as(usize, @intCast(src + 2)) * d + ch]; + } + // The taps' sum as products then adds, or as fused steps: the compiler may contract either way. + var acc: f32 = 0; + var fused: f32 = 0; + for (0..3) |tap| { + acc += c.fromBf16(bx[tap]) * c.fromBf16(p.weights[ch * 3 + tap]); + fused = @mulAdd(f32, c.fromBf16(bx[tap]), c.fromBf16(p.weights[ch * 3 + tap]), fused); + } + const gate = c.fromBf16(p.host[(b + r) * 3 * d + d + ch]); + const got = y[(b + r) * d + ch]; + if (!same(got, c.toBf16(gate * c.fromBf16(c.toBf16(acc)))) and !same(got, c.toBf16(gate * c.fromBf16(c.toBf16(fused))))) + return "a conv output differs from the host's replay"; + if (!same(snap[((b + r) * 2) * d + ch], bx[1]) or !same(snap[((b + r) * 2 + 1) * d + ch], bx[2])) return "a kept state differs from the host's replay"; + }; + } + return null; + } +}; + +fn same(x: u16, y: u16) bool { + const nx = std.math.isNan(c.fromBf16(x)); + return if (nx) std.math.isNan(c.fromBf16(y)) else x == y; +} + +/// Gate 5 for the shared-round launches on this GPU. +pub fn limits(a: std.mem.Allocator, device: mtl.Device, only: ?[]const u8) !usize { + var problems: usize = 0; + const launches = [_]struct { []const u8, []const u8, u32, u32 }{ + .{ names[1], "lfm2_conv_segs", 256, 0 }, + .{ names[0] ++ " qk", "lfm2_qk_segs", 32, 0 }, + .{ names[0], "lfm2_attn_segs", 128, 18688 }, + .{ names[0] ++ " merge", "lfm2_attn_merge_segs", 32, 0 }, + }; + for (launches) |l| { + if (only) |s| if (std.mem.indexOf(u8, l[0], s) == null) continue; + problems += try @import("metal_limits.zig").onThisGpu(a, device, .{ .decl = .{ .name = l[0], .threads = l[2], .tg_memory = l[3], .offline = false }, .function = l[1], .here = sources.lfm2, .pre_m5 = "" }, false); + } + return problems; +} diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 97a3f9153a..c0c20308e4 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -24,6 +24,7 @@ const gemma = @import("metal_gemma4.zig"); const gemma_embed = @import("metal_gemma4_embed.zig"); const gemma_mixed = @import("metal_gemma4_mixed.zig"); const lfm2_argmax = @import("metal_lfm2_argmax.zig"); +const lfm2_streams = @import("metal_lfm2_streams.zig"); const g = contract.gates; const usage = "usage: tf-kernel-contract [--only TEXT] [--no-gates] [--no-limits] [--offline DIR] [--seed N] [--repro] [--rotary-long] [--rotary-accuracy FILE]"; @@ -158,6 +159,15 @@ pub fn main(init: std.process.Init) !u8 { argmax = try lfm2_argmax.Probe.init(a, device, queue); try kernels.append(a, argmax.?.kernel()); } + var streams_list: [2]lfm2_streams.Probe = undefined; + var streams_made: usize = 0; + defer for (streams_list[0..streams_made]) |*p| p.deinit(); + for (std.enums.values(lfm2_streams.Which)) |w| { + if (!wanted(o, lfm2_streams.spec(w).name)) continue; + streams_list[streams_made] = try lfm2_streams.Probe.init(a, device, queue, w); + streams_made += 1; + } + for (streams_list[0..streams_made]) |*p| try kernels.append(a, p.kernel()); var rows_list: std.ArrayList(*rows_k.Row) = .empty; var proj_list: std.ArrayList(*affine_k.Proj) = .empty; defer for (rows_list.items) |r| r.deinit(); @@ -298,6 +308,7 @@ pub fn main(init: std.process.Init) !u8 { lim += try gemma_embed.limits(a, device, o.only); lim += try gemma_mixed.limits(a, device, o.only); lim += try lfm2_argmax.limits(a, device, o.only); + lim += try lfm2_streams.limits(a, device, o.only); } const t2 = std.Io.Clock.awake.now(init.io); const ms = struct { From 520d61792b41ed27b0e0f2af12bb3905c92d442f Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 01:27:38 +0100 Subject: [PATCH 05/10] tile_qmm: the twin's MMA window only for matrices of 4,096 outputs or more From 8 rows the dense twin took its MMA window kernel, a threadgroup for 16 outputs over every row. A 2,048-wide matrix then has 128 threadgroups and runs slower than the scalar twin's 4-row passes; shared rounds of 8 to 32 rows call such matrices at every layer. The 6,144-wide short conv input and the head keep the window. Both kernels give the tiles' bits (the load check runs each matrix's twin at every width); the contract's twin now also serves 16 rows of 6,144 from 2,048. --- zig/src/core/tile_qmm.zig | 4 +++- zig/tests/contract/metal_tile.zig | 3 ++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/zig/src/core/tile_qmm.zig b/zig/src/core/tile_qmm.zig index 64bb1a4971..030bd097e1 100644 --- a/zig/src/core/tile_qmm.zig +++ b/zig/src/core/tile_qmm.zig @@ -125,6 +125,8 @@ pub const twin_rows = 32; pub const window_slots = 4 * twin_rows; /// Rows from which a dense twin call takes the MMA window kernel over the scalar twin. pub const window_rows = 8; +/// Outputs from which a matrix takes that window kernel: narrower ones keep the scalar passes' threadgroups. +pub const window_n = 4096; /// Window rows from which the by-expert twin (each picked expert read once) beats a threadgroup row a slot. pub const group_rows = 16; /// Columns a scalar twin threadgroup (2 simdgroups of 4) and a window threadgroup (2 simdgroups of 8, one MMA block). @@ -168,7 +170,7 @@ pub fn encodeTwin( ) !void { const i = try width(m); if (rows == 0 or rows > twin_rows) return error.NoRows; - const window = rows >= window_rows; + const window = rows >= window_rows and m.n >= window_n; e.setPipeline(if (window) p.rows[i] else if (rows > 1) p.twin_few[i] else p.twin[i]); bind(e, m, .{ .k = @intCast(m.k), .n = @intCast(m.n), .r = @intCast(rows) }, x, x_off, y, y_off); if (window) return e.dispatchGroups(.{ .width = m.n / window_columns }, .{ .width = 64 }); diff --git a/zig/tests/contract/metal_tile.zig b/zig/tests/contract/metal_tile.zig index 98fb49ae1d..5acdb678c6 100644 --- a/zig/tests/contract/metal_tile.zig +++ b/zig/tests/contract/metal_tile.zig @@ -19,7 +19,8 @@ fn tileDepth(s: c.Shape) usize { /// LFM2.5-24B-A2B's expert projections, which #822 will run on the tiles: gate or up, 1,536 from 2,048, and down. const prompt_rows = [_]c.Shape{ .{ .rows = 64, .n = 1536, .k = 2048 }, .{ .rows = 33, .n = 2048, .k = 1536 } }; -const decode_rows = [_]c.Shape{ .{ .rows = 32, .n = 1536, .k = 2048 }, .{ .rows = 5, .n = 2048, .k = 1536 } }; +/// The twin's rows: those projections, and the short conv's 6,144 from 2,048 at 16 rows (the twin's MMA window). +const decode_rows = [_]c.Shape{ .{ .rows = 32, .n = 1536, .k = 2048 }, .{ .rows = 5, .n = 2048, .k = 1536 }, .{ .rows = 16, .n = 6144, .k = 2048 } }; pub fn spec(kind: Kind, name: []const u8) c.Spec { const k_tile = c.Tile{ .size = 128, .step = 64, .min = 64 }; From 9df2cd982b93f13c94cbe9be5954407488aeaedd Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 10:59:07 +0100 Subject: [PATCH 06/10] docs: the LFM2.5-24B-A2B recipe, its shared rounds and their M5 Ultra receipt README pointed at no native recipe for LFM2.5-24B-A2B. The page gives the checkpoint, the serve command, the shared-round kernels and their tests, and the measured M5 Ultra decode rates before and after them. --- README.md | 2 +- docs/recipes/README.md | 1 + docs/recipes/lfm2.5-24b-a2b.md | 43 ++++++++++++++++++++++++++++++++++ 3 files changed, 45 insertions(+), 1 deletion(-) create mode 100644 docs/recipes/lfm2.5-24b-a2b.md diff --git a/README.md b/README.md index 8e60137708..1c121be219 100644 --- a/README.md +++ b/README.md @@ -173,7 +173,7 @@ Conversions are explicitly marked: their upstream weights need the linked conver | GLM-5.3-Flash, FP4-Local | `tensorfold pull nvidia/GLM-5.3-Flash-FP4-Local@76348ddb2be2e35e39ecfc00e75dc07703f4f3d5` | Two GB10s at TP=2, or three at TP=3 over a ring of direct links: @Bizuayeu loaded it as stored and served it at both sizes. 165.3 GiB of weights, 13% less than NVFP4's 190.4 GiB. | [Native CUDA](docs/recipes/glm-5.3-flash-cuda.md), in 1.1.5 (#1050). NVIDIA's mixed-precision checkpoint: NVFP4 routed experts with compressed block scales, MXFP8 and 128×128-block FP8 for most of the rest. As close to BF16 as NVFP4 at its smaller size: at TP=2 its NLL is 1.11% above BF16's against NVFP4's 1.37% ([contributed receipt](docs/receipts/glm-fp4local-80517fc3/README.md) on #1050's head, not rerun on main), a difference the per-document intervals in [XSFP4's card](https://huggingface.co/Bizuayeu/GLM-5.3-Flash-XSFP4#measured) do not resolve. | | GLM-5.3-Flash, XSFP4 | `tensorfold pull Bizuayeu/GLM-5.3-Flash-XSFP4@88736f916f8939607fde8a4b881dcae749f7551d` | Two GB10s at TP=2, or three at TP=3 over a ring of direct links: on main `8086c205c` @Bizuayeu loaded it in 370.7 s at TP=2 (rank 0's model 96.97 GiB at a 524,288-token window) and in 313.8 s at TP=3 (63.60 GiB at 1,048,576), and both served. 165.7 GiB of weights, 0.4 GiB more than FP4-Local's. | [Native CUDA](docs/recipes/glm-5.3-flash-cuda.md), in 1.1.5 through FP4-Local's path. FP4-Local with its routed experts' NVFP4 blocks fitted again from the BF16 weights, a smoothing folded into them, so it runs on the same kernels. Closer to BF16 than FP4-Local and NVFP4 at TP=2 by all three measures in [its card](https://huggingface.co/Bizuayeu/GLM-5.3-Flash-XSFP4#measured), on #1050's code: top-1 agreement 93.26%, NLL 0.85% above and a KL bound of 0.0431, against FP4-Local's 92.82%, 1.11% and 0.0510. | | GLM-5.3, EXLR8 `home` | `tensorfold pull kindlingai/glm-5.3-exlr8-k2-k3-k4@f459f5d0272a529403dd9b3fd32f771f4250a0d4 --exlr8 home --node 0/2 --dir ~/models/glm-5.3-exlr8-home` | Two M5 Ultra Macs, 256 GB each; about 160 GB resident per Mac at 8k context / four streams. | [Two-Mac recipe](docs/glm-5.3-exlr8.md). Use `--node 1/2` on the other Mac; Thunderbolt 5 and MCDMA required. No CUDA backend. | -| LFM2.5-24B-A2B, MLX 4 / 6 / 8-bit | `tensorfold pull LiquidAI/LFM2.5-24B-A2B-MLX-4bit@667408d494bde191c2cb9dd4f4549827816b9275`
`tensorfold pull LiquidAI/LFM2.5-24B-A2B-MLX-6bit@1deec53eb49571e21cdcbf82fd5839248f07ed88`
`tensorfold pull LiquidAI/LFM2.5-24B-A2B-MLX-8bit@67ab417dd02d0c9545802a726df668620eae049d` | One Mac; M4 Pro checked. | [Native operator recipe](docs/recipes/README.md), [release history](CHANGELOG.md). No separate native model recipe yet. | +| LFM2.5-24B-A2B, MLX 4 / 6 / 8-bit | `tensorfold pull LiquidAI/LFM2.5-24B-A2B-MLX-4bit@667408d494bde191c2cb9dd4f4549827816b9275`
`tensorfold pull LiquidAI/LFM2.5-24B-A2B-MLX-6bit@1deec53eb49571e21cdcbf82fd5839248f07ed88`
`tensorfold pull LiquidAI/LFM2.5-24B-A2B-MLX-8bit@67ab417dd02d0c9545802a726df668620eae049d` | One Mac; M4 Pro checked. | [Native recipe](docs/recipes/lfm2.5-24b-a2b.md), [release history](CHANGELOG.md). | | Qwen3.6-35B-A3B, MLX with MTP | `tensorfold pull TensorFold/Qwen3.6-35B-A3B-MLX-4bit-MTP@81169a9bc511a27c1b4eedb77a2cd98ced431847` | One GB10; 20.9 GB checkpoint. | [Qwen3.6](docs/recipes/qwen3.6-moe.md). Native CUDA; the recipe's Metal path belongs to Python. | | Gemma 4 31B, MLX 4-bit | `tensorfold pull mlx-community/gemma-4-31B-it-4bit@696d436c404745a59f30e4939a658162b0a9e57f`
`tensorfold pull z-lab/gemma-4-31B-it-DFlash@eabd648301ce28583cc14757912e5e0f84e152e1` | One M5 Mac; M5 Ultra checked. M1-M4 refused. | [Native 31B](docs/recipes/gemma-4-native.md). Name the drafter with `--drafter`; images with `--vision`. | | Gemma 4 12B, converted MLX 4-bit | `hf download google/gemma-4-12B-it --revision 707f0a3b8a3c7ad586ed01e27eafbad8a27dd0f7 --local-dir ~/models/gemma-4-12b-hf`
`hf download google/gemma-4-12B-it-assistant --revision 46d4c6f13f0ac0ad827b915669b8df9b81c64c51 --local-dir ~/models/gemma-4-12b-assistant` | One M5 Mac; converted text weights 6.3 GB, image weights 100 MB, plus assistant/caches. M1-M4 refused. | [Convert and serve](docs/recipes/gemma-4-12b-native.md). Upstream download needs `hf`; conversion needs Python, mlx and mlx-lm. Serving is native. | diff --git a/docs/recipes/README.md b/docs/recipes/README.md index 204bb8980e..fa160dc631 100644 --- a/docs/recipes/README.md +++ b/docs/recipes/README.md @@ -21,6 +21,7 @@ For agent-assisted contributions, read the [TensorFold contributor skill](SKILL. | Qwen3.6-35B-A3B | [MLX with MTP drafts, one-GPU CUDA](qwen3.6-moe.md) | | Step-5-Preview | [Native Metal on one 256 GB Mac: 2-bit experts, MTP drafts, images](step-5-preview.md) | | Kolibri-1 | [Native CUDA, block FP8, research checkpoint](kolibri-1.md) | +| LFM2.5-24B-A2B | [Native Metal, MLX 4-bit, shared rounds](lfm2.5-24b-a2b.md) | | Spark-X2.6-4B | [Native CUDA](spark-x2.6-4b-cuda.md), [native Metal, the bf16 release, contexts to 1M](spark-x2.6-4b-metal.md) | ## Capability floor diff --git a/docs/recipes/lfm2.5-24b-a2b.md b/docs/recipes/lfm2.5-24b-a2b.md new file mode 100644 index 0000000000..7b8b2cdf4b --- /dev/null +++ b/docs/recipes/lfm2.5-24b-a2b.md @@ -0,0 +1,43 @@ +# LFM2.5-24B-A2B, native Metal + +This recipe covers [`LiquidAI/LFM2.5-24B-A2B-MLX-4bit`](https://huggingface.co/LiquidAI/LFM2.5-24B-A2B-MLX-4bit) +at revision `667408d494bde191c2cb9dd4f4549827816b9275`. The numbers below are measured on this 4-bit export only. + +```bash +hf download LiquidAI/LFM2.5-24B-A2B-MLX-4bit \ + --revision 667408d494bde191c2cb9dd4f4549827816b9275 --local-dir "$HOME/models/lfm2.5-24b-a2b-4bit" +zig build native -Dcpu=apple_m1 +zig-out/native/bin/tensorfold-native serve "$HOME/models/lfm2.5-24b-a2b-4bit" --port 8080 --parallel 16 +``` + +A Mac with less memory may not hold 16 streams at the checkpoint's 131,072-token window; the server then names how +many fit. `--context 40960` holds 16 on a 48 GB M4 Pro. + +## Shared rounds + +Every stream's tokens go through the lane rounds. A round of several streams is one forward over all their rows. +In it each layer's short conv (30 layers) and each attention layer's qk, attention and merge (10 layers) run once +over every row (`lfm2_segments.metal`), from a table of each stream's cache addresses, row offsets and key counts +(`segments.zig`). Each row's arithmetic is the per-stream kernel's, so a reply equals its solo run. Prompt calls and +one-stream windows keep the per-stream launches. A greedy row takes the GPU's pick of its logits (`lfm2_argmax`); +sampled and masked rows draw on the host. The dense twin takes its MMA window only for matrices of 4,096 outputs or +more (the short conv's input and the head); narrower ones keep the scalar passes. + +`zig build test-lfm2` runs the family's tests. `tf-kernel-contract --only lfm2` checks the shared-round kernels and +the pick, and `tf-lfm2-exact` compares windows and shared rounds with one-row decode on the real weights. + +## Measured + +Greedy, 256 tokens, 16 public prompts, a different one a stream, 3 reps in each of 3 alternated rounds on an M5 Ultra; +aggregate tok/s, the range of the rounds' medians. Before is main at 0d24ecb3 (each stream's own launches), after adds +the shared-round kernels on it. The M4 Pro is not yet measured on a quiet machine. + +| Streams | before | after | +| --- | --- | --- | +| 1 | 258.5-259.2 | 266.5-267.7 | +| 4 | 483.5-486.5 | 565.2-566.1 | +| 8 | 551.5-551.9 | 732.2-735.6 | +| 16 | 660.7-661.9 | 922.5-923.5 | + +Every reply was token for token equal before and after, and equal to its solo run and to `"draft": false`. +Cold prefill at 2k, 8k and 32k tokens and peak memory did not change. From 970f10d728b280c42eee4d0de351aca44840e435 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Sun, 25 Oct 2026 01:13:08 +0200 Subject: [PATCH 07/10] lfm2: keep a drafter's tapped rows on every forward The lane core can put an external drafter in front of a target. The drafter reads the target's hidden rows after some layers (its taps). LFM2-MoE had no taps, so it could draft only with the core's own heads. prepare_features names the layers. From then on, each forward copies each row's output of those layers into the scratch (lfm2_tap, a plain copy of bf16 values). features() gives a stream's rows of the latest verify or one-token round. prefill_features hands each prompt chunk's rows to the drafter before the chunk's end is kept for reuse. The prompt's last row is kept beside the stream's caches for its first draft. A drafting stream resumes only a kept state that also holds its drafter's context (Snap.draft). Without a drafter nothing changes: no taps are kept and every forward runs as before. lfm2_tap registers with the kernel contract (metal_lfm2_tap.zig): the third of five taps' columns at 1, 10, 32 and 2,048 rows must hold the layer's bits and the other columns what the call found. Gate 5 checks its launch against every Apple chip we serve and this GPU. test-lfm2-dspark checks the same copy through the family's own kernels. --- zig/build/metal_targets.zig | 2 + zig/kernels/metal/lfm2/lfm2.metal | 8 ++ zig/src/families/lfm2_moe/backend.zig | 82 ++++++++++++- zig/src/families/lfm2_moe/dspark_test.zig | 77 ++++++++++++ zig/src/families/lfm2_moe/forward.zig | 14 +++ zig/src/families/lfm2_moe/kernels.zig | 1 + zig/src/families/lfm2_moe/snapshot.zig | 3 +- zig/src/families/lfm2_moe/state.zig | 12 ++ zig/src/tensorfold.zig | 1 + zig/tests/contract/metal_lfm2_tap.zig | 143 ++++++++++++++++++++++ zig/tests/contract/metal_main.zig | 8 ++ 11 files changed, 345 insertions(+), 6 deletions(-) create mode 100644 zig/src/families/lfm2_moe/dspark_test.zig create mode 100644 zig/tests/contract/metal_lfm2_tap.zig diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index 7f1c531bab..6db7154689 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -457,6 +457,8 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = engine })).step); const lfm2_tests = b.addRunArtifact(b.addTest(.{ .root_module = engine, .filters = &.{"lfm2_moe"} })); b.step("test-lfm2", "The LFM2-MoE family's tests: config, draws, state and its shared rounds on this Mac's GPU").dependOn(&lfm2_tests.step); + const lfm2_dspark = b.addRunArtifact(b.addTest(.{ .root_module = engine, .filters = &.{"lfm2 dspark:"} })); + b.step("test-lfm2-dspark", "LFM2's taps for a drafter: the tap kernel on this Mac's GPU").dependOn(&lfm2_dspark.step); test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = lanes })).step); // the fake fabric maps its windows with mmap: the Macs' RDMA, not Windows if (target.result.os.tag != .windows) test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = fabric })).step); diff --git a/zig/kernels/metal/lfm2/lfm2.metal b/zig/kernels/metal/lfm2/lfm2.metal index af6a2a7372..b2772cf64e 100644 --- a/zig/kernels/metal/lfm2/lfm2.metal +++ b/zig/kernels/metal/lfm2/lfm2.metal @@ -568,3 +568,11 @@ constant constexpr uint SORT_CHUNKS = 16; // thread (e, c) counts expert e's pi best[2 * row + 1] = bad; } } + +// A drafter's tap: H's rows [rows, D] copied into TAPS at column `at.y` of rows `at.x` values apart, 8 a thread. +[[kernel]] void lfm2_tap(const device bfloat* h [[buffer(0)]], device bfloat* taps [[buffer(1)]], constant uint2& at [[buffer(2)]], + uint2 i [[thread_position_in_grid]]) { + if (i.x >= D / 8) return; + const device vec* from = (const device vec*)(h + size_t(i.y) * D); + ((device vec*)(taps + size_t(i.y) * at.x + at.y))[i.x] = from[i.x]; +} diff --git a/zig/src/families/lfm2_moe/backend.zig b/zig/src/families/lfm2_moe/backend.zig index d713c9ec24..44427d435a 100644 --- a/zig/src/families/lfm2_moe/backend.zig +++ b/zig/src/families/lfm2_moe/backend.zig @@ -23,6 +23,8 @@ pub const Metal = struct { /// Forward ms by window width, timed at load: the lane core weighs drafts against them. costs: [st.window_rows]lanes.config.Cost = undefined, timed: usize = 0, + /// Forwards run so far: a stream's tapped rows in the scratch are current while their gen matches. + gen: u64 = 0, pub fn init(gpa: std.mem.Allocator, model: *Model, options: Options) !*Metal { const out = try gpa.create(Metal); @@ -88,7 +90,7 @@ pub const Metal = struct { } pub fn backend(self: *Metal) be.Backend { - return .{ .ptr = self, .vtable = &.{ .prefill = prefillFn, .first = firstFn, .first_masked = firstMaskedFn, .queue = queueFn, .read = readFn, .verify = verifyFn, .keep = keepFn, .draft = draftFn, .release = releaseFn } }; + return .{ .ptr = self, .vtable = &.{ .prefill = prefillFn, .prepare_features = prepareFn, .features = featuresFn, .prefill_features = prefillFeaturesFn, .first = firstFn, .first_masked = firstMaskedFn, .queue = queueFn, .read = readFn, .verify = verifyFn, .keep = keepFn, .draft = draftFn, .release = releaseFn } }; } fn cast(ptr: *anyopaque) *Metal { @@ -107,8 +109,65 @@ pub const Metal = struct { return cache; } + /// One forward on the scratch; a stream's tapped rows from earlier forwards are stale after it. + fn run(self: *Metal, segments: []const fwd.Segment, tokens: []const u32, record: bool, head: fwd.Head) !void { + self.gen += 1; + try fwd.run(self.model, &self.scratch, segments, tokens, record, head); + } + fn prefillFn(ptr: *anyopaque, stream: *lanes.Stream) !void { + return cast(ptr).prefillWith(stream, null); + } + + /// The prompt pass, `sink` taking each chunk's tapped rows before the chunk's end is kept for reuse. + fn prefillFeaturesFn(ptr: *anyopaque, stream: *lanes.Stream, layers: []const u32, sink: be.FeatureSink) !void { + const self = cast(ptr); + const t = self.scratch.taps orelse return error.UnpreparedLfmTaps; + if (layers.len != t.count) return error.UnpreparedLfmTaps; + return self.prefillWith(stream, sink); + } + + /// A tapped row's bytes: every tap's hidden row, bf16. + fn tapBytes(self: *const Metal) usize { + return if (self.scratch.taps) |t| t.count * c.hidden * 2 else 0; + } + + /// The drafter's layers: every forward from now on keeps their outputs for each row. + fn prepareFn(ptr: *anyopaque, layers: []const u32) !void { const self = cast(ptr); + if (layers.len == 0 or layers.len > st.max_taps) return error.BadLfmTaps; + for (layers) |l| if (l >= c.layers) return error.BadLfmTaps; + if (self.scratch.taps) |t| return if (std.mem.eql(u32, t.layers[0..t.count], layers)) {} else error.BadLfmTaps; + var t: st.Taps = .{ .count = @intCast(layers.len), .buf = undefined }; + @memcpy(t.layers[0..layers.len], layers); + t.buf = try self.scratch.memory.alloc(self.scratch.rows * layers.len * c.hidden * 2); + self.scratch.taps = t; + } + + /// Cache rows [start, start + count) of the stream's latest tapped rows, on the device. + fn featuresFn(ptr: *anyopaque, stream: *lanes.Stream, layers: []const u32, start: u64, count: u32) !be.Features { + const self = cast(ptr); + const t = self.scratch.taps orelse return error.UnpreparedLfmTaps; + if (layers.len != t.count) return error.UnpreparedLfmTaps; + const cache = try self.cacheOf(stream); + const held = cache.tap orelse return error.LfmTapsNotHeld; + if (count == 0 or start < held.start or start + count > held.start + held.rows) return error.LfmTapsNotHeld; + const width = self.tapBytes(); + const buf = if (held.gen == 0) cache.last_tap.? else if (held.gen == self.gen) t.buf else return error.StaleLfmTaps; + const row = held.row + @as(usize, @intCast(start - held.start)); + return .{ .buffer = @intFromPtr(buf.id), .offset = row * width, .rows = count, .row_bytes = @intCast(width), .space = .device, .dtype = .bf16 }; + } + + /// The prompt's last row's taps, copied beside the stream's caches for its first draft (gen 0: kept there). + fn keepLastTap(self: *Metal, cache: *st.Cache, row: usize, position: usize) !void { + const width = self.tapBytes(); + if (width == 0) return; + if (cache.last_tap == null) cache.last_tap = try cache.memory.alloc(width); + @memcpy(cache.last_tap.?.contents()[0..width], (self.scratch.taps.?.buf.contents() + row * width)[0..width]); + cache.tap = .{ .start = position, .row = 0, .rows = 1, .gen = 0 }; + } + + fn prefillWith(self: *Metal, stream: *lanes.Stream, sink: ?be.FeatureSink) !void { const ids = stream.prompt(); if (ids.len == 0 or ids.len + stream.max_new > self.options.capacity or self.caches.contains(stream)) return error.InvalidLfmPrompt; const capacity = @min(self.options.capacity, ids.len + stream.max_new + st.window_rows); @@ -121,7 +180,9 @@ pub const Metal = struct { stream.cached = 0; if (stream.reuse.saved) |saved| { // a kept state of this prompt's prefix: the pass starts there const snap: *snapshot.Snap = @ptrCast(@alignCast(saved)); - if (snap.at < ids.len) if (snapshot.restore(cache, snap)) |_| { + // a drafting stream resumes only a state kept with its drafter's context + const drafter_less = sink != null and snap.draft == null; + if (snap.at < ids.len and !drafter_less) if (snapshot.restore(cache, snap)) |_| { at = snap.at; } else |_| {}; stream.cached = @intCast(at); @@ -136,7 +197,13 @@ pub const Metal = struct { }; while (mark < stream.reuse.marks.len and stream.reuse.marks[mark] <= at) mark += 1; if (mark < stream.reuse.marks.len) end = @min(end, stream.reuse.marks[mark]); - try fwd.run(self.model, &self.scratch, &.{.{ .cache = cache, .rows = end - at }}, ids[at..end], false, .last); + try self.run(&.{.{ .cache = cache, .rows = end - at }}, ids[at..end], false, .last); + if (sink) |k| { + const width = self.tapBytes(); + const f: be.Features = .{ .buffer = @intFromPtr(self.scratch.taps.?.buf.id), .rows = @intCast(end - at), .row_bytes = @intCast(width), .space = .device, .dtype = .bf16 }; + try k.take(k.ptr, stream, f, at); + if (end == ids.len) try self.keepLastTap(cache, end - 1 - at, end - 1); + } at = end; if (stream.reuse.hook) |hk| if (std.mem.indexOfScalar(u32, stream.reuse.marks, @intCast(at)) != null) hk.at(hk.ptr, stream, @intCast(at)); @@ -165,7 +232,8 @@ pub const Metal = struct { .value => |v| v, .handle => |h| @intCast(h), }; - try fwd.run(self.model, &self.scratch, &.{.{ .cache = cache, .rows = 1 }}, &.{token}, false, .all); + try self.run(&.{.{ .cache = cache, .rows = 1 }}, &.{token}, false, .all); + if (self.scratch.taps != null) cache.tap = .{ .start = cache.len - 1, .row = 0, .rows = 1, .gen = self.gen }; return try draw(self.gpa, cache.logits.slice(u16, c.vocab), stream.sampling, position, &.{}); } @@ -191,7 +259,11 @@ pub const Metal = struct { @memcpy(ids[total + 1 ..][0..w.tokens.len], w.tokens); total += rows; } - try fwd.run(self.model, &self.scratch, segments[0..windows.len], ids[0..total], true, .all); + try self.run(segments[0..windows.len], ids[0..total], true, .all); + if (self.scratch.taps != null) for (segments[0..windows.len]) |seg| { + const last = seg.cache.last.?; + seg.cache.tap = .{ .start = last.start, .row = last.base, .rows = last.rows, .gen = self.gen }; + }; const logits = self.scratch.logits.slice(u16, total * c.vocab); const best = self.scratch.best.slice(u32, 2 * total); var base: usize = 0; diff --git a/zig/src/families/lfm2_moe/dspark_test.zig b/zig/src/families/lfm2_moe/dspark_test.zig new file mode 100644 index 0000000000..d220b9ea66 --- /dev/null +++ b/zig/src/families/lfm2_moe/dspark_test.zig @@ -0,0 +1,77 @@ +//! The DSpark drafter's kernels on the GPU against host references: the target's taps. +const std = @import("std"); +const mtl = @import("metal"); +const c = @import("config.zig"); +const kernels = @import("kernels.zig"); + +const shared = mtl.ResourceOptions.shared; + +const Rig = struct { + device: mtl.Device, + queue: mtl.Queue, + k: kernels.Kernels, + + fn init() !Rig { + const device = mtl.Device.init() catch return error.SkipZigTest; + errdefer device.deinit(); + const queue = try device.queue(); + errdefer queue.deinit(); + return .{ .device = device, .queue = queue, .k = try kernels.load(device) }; + } + + fn deinit(r: *Rig) void { + r.k.deinit(); + r.queue.deinit(); + r.device.deinit(); + } + + fn buffer(r: *Rig, comptime T: type, n: usize) !mtl.Buffer { + const b = try r.device.buffer(n * @sizeOf(T), shared); + @memset(b.slice(u8, n * @sizeOf(T)), 0); + return b; + } + + fn run(r: *Rig, comptime key: []const u8, bufs: []const mtl.Buffer, value: anytype, at: usize, grid: [2]usize) !void { + const cb = r.queue.commandBuffer(); + const e = cb.compute(.serial); + e.setPipeline(r.k.get(key)); + for (bufs, 0..) |b, i| e.setBuffer(b, 0, if (i >= at) i + 1 else i); + e.setValue(value, at); + e.dispatchThreads(mtl.Size.of(grid[0], grid[1], 1), mtl.Size.of(@min(grid[0], 256), 1, 1)); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |_| return error.GpuFailed; + } +}; + +/// Round to nearest even, as MLX casts. +fn toBf16(f: f32) u16 { + const b: u32 = @bitCast(f); + return @truncate((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +fn noise(seed: u64, out: []u16, scale: f32) void { + var prng = std.Random.DefaultPrng.init(seed); + for (out) |*v| v.* = toBf16((prng.random().float(f32) * 2 - 1) * scale); +} + +test "lfm2 dspark: the tap kernel copies each row's hidden values into its tap's columns, the rest untouched" { + var r = try Rig.init(); + defer r.deinit(); + const rows = 5; + const taps = 3; + const h = try r.buffer(u16, rows * c.hidden); + defer h.deinit(); + const out = try r.buffer(u16, rows * taps * c.hidden); + defer out.deinit(); + noise(7, h.slice(u16, rows * c.hidden), 4); + try r.run("lfm2_tap", &.{ h, out }, [2]u32{ taps * c.hidden, c.hidden }, 2, .{ c.hidden / 8, rows }); + const want = h.slice(u16, rows * c.hidden); + const got = out.slice(u16, rows * taps * c.hidden); + for (0..rows) |row| { + try std.testing.expectEqualSlices(u16, want[row * c.hidden ..][0..c.hidden], got[(row * taps + 1) * c.hidden ..][0..c.hidden]); + for (got[row * taps * c.hidden ..][0..c.hidden]) |v| try std.testing.expectEqual(@as(u16, 0), v); + for (got[(row * taps + 2) * c.hidden ..][0..c.hidden]) |v| try std.testing.expectEqual(@as(u16, 0), v); + } +} diff --git a/zig/src/families/lfm2_moe/forward.zig b/zig/src/families/lfm2_moe/forward.zig index 212b5978bc..7b63833192 100644 --- a/zig/src/families/lfm2_moe/forward.zig +++ b/zig/src/families/lfm2_moe/forward.zig @@ -146,6 +146,7 @@ pub fn run(m: *Model, s: *st.Scratch, segments: []const Segment, tokens: []const for (m.weights.blocks, 0..) |block, i| { enc.norm(s, block.operator_norm, i == 0, total, combine); combine = false; + if (i > 0) tap(enc, s, i - 1, total); switch (block.mixer) { .conv => |cv| conv(enc, s, cv, i, segments, total), .attention => |a| attention(enc, s, a, &m.config, segments, i, total), @@ -162,6 +163,7 @@ pub fn run(m: *Model, s: *st.Scratch, segments: []const Segment, tokens: []const } } enc.norm(s, m.weights.norm, false, total, combine); + tap(enc, s, c.layers - 1, total); if (head == .all) { enc.projection(embedding, s.x, 0, s.logits, 0, total); enc.pipe("lfm2_argmax"); @@ -187,6 +189,18 @@ pub fn run(m: *Model, s: *st.Scratch, segments: []const Segment, tokens: []const } } +/// Layer `layer`'s output rows (h, just summed by the next norm) into each tap of it, when a drafter reads them. +fn tap(enc: Encoder, s: *st.Scratch, layer: usize, total: usize) void { + const taps = s.taps orelse return; + for (taps.layers[0..taps.count], 0..) |l, t| if (l == layer) { + enc.pipe("lfm2_tap"); + enc.e.setBuffer(s.h, 0, 0); + enc.e.setBuffer(taps.buf, 0, 1); + enc.e.setValue([2]u32{ taps.count * c.hidden, @intCast(t * c.hidden) }, 2); + enc.threads(.{ c.hidden / 8, total, 1 }, .{ 256, 1, 1 }); + }; +} + fn conv(enc: Encoder, s: *st.Scratch, cv: wts.Conv, layer: usize, segments: []const Segment, total: usize) void { const e = enc.e; enc.projection(cv.in, s.x, 0, s.bcx, 0, total); diff --git a/zig/src/families/lfm2_moe/kernels.zig b/zig/src/families/lfm2_moe/kernels.zig index aa1e798fc0..61af45f4c8 100644 --- a/zig/src/families/lfm2_moe/kernels.zig +++ b/zig/src/families/lfm2_moe/kernels.zig @@ -27,6 +27,7 @@ const names = [_][:0]const u8{ "lfm2_qk_segs", "lfm2_attn_segs", "lfm2_attn_merge_segs", + "lfm2_tap", }; pub const Kernels = struct { diff --git a/zig/src/families/lfm2_moe/snapshot.zig b/zig/src/families/lfm2_moe/snapshot.zig index d1c1a8b3ae..d7289331b7 100644 --- a/zig/src/families/lfm2_moe/snapshot.zig +++ b/zig/src/families/lfm2_moe/snapshot.zig @@ -4,7 +4,7 @@ const c = @import("config.zig"); const st = @import("state.zig"); /// One state: conv states by conv layer, then keys and values [kv head][at][head dim] by attention layer. -pub const Snap = struct { at: u32, bytes: []u8 }; +pub const Snap = struct { at: u32, bytes: []u8, draft: ?[]u8 = null }; /// A row's keys or values for every kv head of one attention layer. const row_bytes = c.kv_heads * c.head_dim * 2; @@ -56,6 +56,7 @@ pub fn restore(cache: *st.Cache, snap: *const Snap) !void { } pub fn drop(gpa: std.mem.Allocator, snap: *Snap) void { + if (snap.draft) |d| gpa.free(d); // the drafter's context at `at`, when one was kept with the state gpa.free(snap.bytes); gpa.destroy(snap); } diff --git a/zig/src/families/lfm2_moe/state.zig b/zig/src/families/lfm2_moe/state.zig index 3242d8bf26..1bc749ae95 100644 --- a/zig/src/families/lfm2_moe/state.zig +++ b/zig/src/families/lfm2_moe/state.zig @@ -51,6 +51,12 @@ pub const Storage = struct { }; pub const Attention = struct { keys: mtl.Buffer, values: mtl.Buffer }; +/// Layers a drafter reads at most (each layer's output, hidden wide). +pub const max_taps = 8; +/// The drafter's layers in its order and every forward row's outputs of them, [rows][count][hidden] bf16. +pub const Taps = struct { layers: [max_taps]u32 = @splat(0), count: u32 = 0, buf: mtl.Buffer }; +/// Where a stream's latest tapped rows sit in the scratch: cache rows [start, start + rows) from scratch row `row`. +pub const TapRows = struct { start: usize, row: usize, rows: usize, gen: u64 }; pub const Mixer = union(enum) { conv: mtl.Buffer, attention: Attention }; pub const Cache = struct { @@ -60,6 +66,9 @@ pub const Cache = struct { last: ?struct { start: usize, base: usize, rows: usize } = null, blocks: [c.layers]Mixer, logits: mtl.Buffer, + /// The stream's latest tapped rows in the scratch, and its prompt's last row's taps (kept for the first draft). + tap: ?TapRows = null, + last_tap: ?mtl.Buffer = null, pub fn init(gpa: std.mem.Allocator, device: mtl.Device, attention: *const [c.layers]bool, capacity: usize) !Cache { var out: Cache = .{ .memory = .{ .gpa = gpa, .device = device }, .capacity = capacity, .blocks = undefined, .logits = undefined }; @@ -148,6 +157,8 @@ pub const Scratch = struct { segs: mtl.Buffer, rowseg: mtl.Buffer, tiles: mtl.Buffer, + /// A drafter's taps of every row, once the backend is told its layers (null: none kept). + taps: ?Taps = null, pub fn init(gpa: std.mem.Allocator, device: mtl.Device, attention: *const [c.layers]bool, rows: usize) !Scratch { var out: Scratch = undefined; @@ -155,6 +166,7 @@ pub const Scratch = struct { errdefer out.memory.deinit(); out.rows = rows; out.snapshots = @splat(null); + out.taps = null; const mem = &out.memory; const slots = rows * c.topk; out.ids = try mem.alloc(rows * 4); diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 26b68d7925..7797558179 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -84,6 +84,7 @@ test { _ = @import("families/lfm2_moe/config.zig"); _ = @import("families/lfm2_moe/backend.zig"); _ = @import("families/lfm2_moe/state.zig"); + _ = @import("families/lfm2_moe/dspark_test.zig"); } pub const lane_projection = @import("core/lane_projection.zig"); diff --git a/zig/tests/contract/metal_lfm2_tap.zig b/zig/tests/contract/metal_lfm2_tap.zig new file mode 100644 index 0000000000..871ea8a8a2 --- /dev/null +++ b/zig/tests/contract/metal_lfm2_tap.zig @@ -0,0 +1,143 @@ +//! LFM2's tap copy for an external drafter (lfm2_tap): a layer's rows into their tap's columns, the others kept. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); +const c = contract.spec; +const g = contract.gates; +const rig_ = @import("rig.zig"); + +const hidden = 2048; +const taps = 5; +pub const name = "metal lfm2 tap copy"; + +/// rows: one decode row, a verify window, a shared round's 32 and a 2,048-row prompt chunk. +const shapes = [_]c.Shape{ .{ .rows = 1, .n = taps * hidden, .k = 1 }, .{ .rows = 10, .n = taps * hidden, .k = 1 }, .{ .rows = 32, .n = taps * hidden, .k = 1 }, .{ .rows = 2048, .n = taps * hidden, .k = 1 } }; + +fn none(_: c.Shape) usize { + return 0; +} + +pub fn spec() c.Spec { + return .{ .name = name, .accum = .none, .depth = none, .out = .bf16, .served = &shapes, .sums = .{ .f16 = 0, .bf16 = 0, .why = "a copy moves values: no sums" } }; +} + +/// A bf16 value for element `i` over signs and binades; the nonfinite fill plants NaN and infinities. +fn value(i: usize, seed: u64, fill: g.Fill) u16 { + if (fill == .nonfinite) switch ((@as(u64, i) *% 2654435761 +% seed) % 9973) { + 0 => return 0x7fc0, + 1 => return 0x7f80, + 2 => return 0xff80, + else => {}, + }; + const h: u64 = (@as(u64, i) *% 0x9e3779b97f4a7c15 +% seed *% 0xbf58476d1ce4e5b9) >> 40; + const mant: f32 = @as(f32, @floatFromInt(h % 1024)) / 1024 + 1; + const exp: i32 = @as(i32, @intCast((h >> 10) % 8)) - 4; + const v = std.math.ldexp(mant, exp) * @as(f32, if (((h >> 15) & 1) == 0) 1 else -1); + return c.toBf16(v); +} + +pub const Probe = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + rows: usize = 0, + src: []u16 = &.{}, // H rows + prior: []u16 = &.{}, // the taps rows as the call finds them (the other taps' columns must stay) + args: [2]rig_.Ref = undefined, + y: rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue) !Probe { + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .pipe = try rig_.pipeline(device, sources.lfm2, "lfm2_tap") }; + } + + fn drop(p: *Probe) void { + p.a.free(p.src); + p.a.free(p.prior); + p.src = &.{}; + p.prior = &.{}; + p.rig.reset(); + } + + pub fn deinit(p: *Probe) void { + p.drop(); + p.rig.deinit(); + p.pipe.deinit(); + } + + const vt = g.hooks(Probe); + pub fn kernel(p: *Probe) g.Kernel { + return .{ .ctx = p, .vt = &vt, .spec = spec() }; + } + + fn fill(out: []u16, seed: u64, f: g.Fill) void { + for (out, 0..) |*v, i| v.* = value(i, seed, f); + } + + pub fn prepare(p: *Probe, s: c.Shape, f: g.Fill, seed: u64) !bool { + p.drop(); + if (s.n != taps * hidden or s.rows == 0) return false; + p.rows = s.rows; + p.src = try p.a.alloc(u16, s.rows * hidden); + fill(p.src, seed, f); + p.prior = try p.a.alloc(u16, s.rows * taps * hidden); + fill(p.prior, seed +% 3, .random); + p.args[0] = try p.rig.input(std.mem.sliceAsBytes(p.src)); + const at = [2]u32{ taps * hidden, 2 * hidden }; // the third tap's columns + p.args[1] = try p.rig.input(std.mem.asBytes(&at)); + p.y = try p.rig.output(p.count() * 2); + return true; + } + + pub fn count(p: *Probe) usize { + return p.rows * taps * hidden; + } + + pub fn run(p: *Probe) !void { + @memcpy(p.rig.bytes(p.y, p.count() * 2), std.mem.sliceAsBytes(p.prior)); + const pass = p.rig.begin(); + const e = pass.e; + e.setPipeline(p.pipe); + e.setBuffer(p.args[0].buf, p.args[0].off, 0); + e.setBuffer(p.y.buf, p.y.off, 1); + e.setBuffer(p.args[1].buf, p.args[1].off, 2); + e.dispatchThreads(mtl.Size.of(hidden / 8, p.rows, 1), mtl.Size.of(256, 1, 1)); + return p.rig.finish(pass); + } + + pub fn reference(p: *Probe, want: []f64, mag: []f64) !void { + for (want, mag, p.prior) |*w, *m, v| { // the other taps' columns keep what the call found + w.* = c.fromBf16(v); + m.* = @abs(w.*); + } + for (0..p.rows) |r| for (0..hidden) |i| { + const v = c.fromBf16(p.src[r * hidden + i]); + want[r * taps * hidden + 2 * hidden + i] = v; + mag[r * taps * hidden + 2 * hidden + i] = @abs(v); + }; + } + + /// The copy is exact: the third tap's columns hold H's bits, the others what the call found. + pub fn check(p: *Probe, _: g.Gate) !?[]const u8 { + const out = std.mem.bytesAsSlice(u16, p.rig.bytes(p.y, p.count() * 2)); + for (out, p.prior, 0..) |got, before, i| { + const r = i / (taps * hidden); + const col = i % (taps * hidden); + const want = if (col >= 2 * hidden and col < 3 * hidden) p.src[r * hidden + col - 2 * hidden] else before; + if (got != want) return "an output differs from the copy's bits"; + } + return null; + } +}; + +/// Gate 5 for the tap's launch of 256 threads and no threadgroup memory: every Apple chip we serve, then this GPU. +pub fn limits(a: std.mem.Allocator, device: mtl.Device, only: ?[]const u8) !usize { + if (only) |s| if (std.mem.indexOf(u8, name, s) == null) return 0; + const d = contract.limits.MetalLaunch{ .name = name, .threads = 256, .tg_memory = 0, .offline = false }; + var notes: std.ArrayList(u8) = .empty; + defer notes.deinit(a); + var problems: usize = 0; + for (contract.limits.apple) |chip| problems += try contract.limits.metal(d, chip, ¬es, a); + std.debug.print("{s}", .{notes.items}); + return problems + try @import("metal_limits.zig").onThisGpu(a, device, .{ .decl = d, .function = "lfm2_tap", .here = sources.lfm2, .pre_m5 = "" }, false); +} diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index c0c20308e4..70bedb69dc 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -25,6 +25,7 @@ const gemma_embed = @import("metal_gemma4_embed.zig"); const gemma_mixed = @import("metal_gemma4_mixed.zig"); const lfm2_argmax = @import("metal_lfm2_argmax.zig"); const lfm2_streams = @import("metal_lfm2_streams.zig"); +const lfm2_tap = @import("metal_lfm2_tap.zig"); const g = contract.gates; const usage = "usage: tf-kernel-contract [--only TEXT] [--no-gates] [--no-limits] [--offline DIR] [--seed N] [--repro] [--rotary-long] [--rotary-accuracy FILE]"; @@ -168,6 +169,12 @@ pub fn main(init: std.process.Init) !u8 { streams_made += 1; } for (streams_list[0..streams_made]) |*p| try kernels.append(a, p.kernel()); + var tap: ?lfm2_tap.Probe = null; + defer if (tap) |*p| p.deinit(); + if (wanted(o, lfm2_tap.name)) { + tap = try lfm2_tap.Probe.init(a, device, queue); + try kernels.append(a, tap.?.kernel()); + } var rows_list: std.ArrayList(*rows_k.Row) = .empty; var proj_list: std.ArrayList(*affine_k.Proj) = .empty; defer for (rows_list.items) |r| r.deinit(); @@ -309,6 +316,7 @@ pub fn main(init: std.process.Init) !u8 { lim += try gemma_mixed.limits(a, device, o.only); lim += try lfm2_argmax.limits(a, device, o.only); lim += try lfm2_streams.limits(a, device, o.only); + lim += try lfm2_tap.limits(a, device, o.only); } const t2 = std.Io.Clock.awake.now(init.io); const ms = struct { From 9f74d9175e2ac51119ee6b982dd9130027b751b9 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 18:47:34 +0100 Subject: [PATCH 08/10] kernel contract: LFM2's shared-round launches carry recorded M1/M2 limits and four planted controls The five launches are in lfm2_records.zig, with the limits that --offline compiled, not this GPU's. lfm2_attn_segs reserves its 128 threads in its source, as M1 and M2 cannot build it for another count. Four controls edit the source: a tie to the highest id, a nonfinite logit not flagged, and every stream reading or writing the first stream's conv state or kv cache. Each goes red at its gate. --- docs/kernel-contract.md | 5 +++ zig/kernels/metal/lfm2/lfm2_segments.metal | 2 +- zig/src/core/contract/lfm2_records.zig | 14 ++++++ zig/src/core/contract/registry.zig | 2 +- zig/tests/contract/metal_lfm2_argmax.zig | 18 ++++---- zig/tests/contract/metal_lfm2_controls.zig | 51 ++++++++++++++++++++++ zig/tests/contract/metal_lfm2_streams.zig | 30 +++++-------- zig/tests/contract/metal_limits.zig | 1 + zig/tests/contract/metal_main.zig | 50 ++++++++++++++------- 9 files changed, 127 insertions(+), 46 deletions(-) create mode 100644 zig/src/core/contract/lfm2_records.zig create mode 100644 zig/tests/contract/metal_lfm2_controls.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 0d8241d5c2..52d9014565 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -131,6 +131,11 @@ Metal: softmax, RMS norm, bf16 top-k, the row projection (3, 4, 6 and 8 bits), a formats on the M5's tensor units, where its reg kernel must equal the cooperative one bit for bit), and tile_qmm's prompt tiles and their decode twin (4, 6 and 8 bits; the twin bit-equal to the tiles). Gate 5 also holds the rest of tile_qmm's pipelines and GLM-5.3-Flash's hc_expand_10. +LFM2.5-24B-A2B's shared rounds: the greedy pick (`lfm2_argmax`) and the conv and attention kernels that run every +stream's rows in one dispatch, each bit-equal to the per-stream kernels it replaces. Four controls edit the source: a tie +to the highest id, a nonfinite logit not flagged, and every stream reading or writing the first stream's state or cache. +Their launches are in `lfm2_records.zig` with the M1 and M2 limits from `--offline`; `lfm2_attn_segs` reserves its 128 +threads in its source, as Metal cannot build it for M1 or M2 at another count. CUDA: the RMS norm rows (16-byte form and the fp32-out tree), the bf16 residual's norm tails (MLX's rounded rows and the once-rounded ones, each checked on its residual and its output, with four controls each dropping a rounding point), Nemotron's router logits (and their 128-row tiles, bit-equal to them) and top-k, the latent top-k's decode form (bit-equal to `tf_super_topk`), the latent combine and diff --git a/zig/kernels/metal/lfm2/lfm2_segments.metal b/zig/kernels/metal/lfm2/lfm2_segments.metal index 4bf25135c7..09997c53e8 100644 --- a/zig/kernels/metal/lfm2/lfm2_segments.metal +++ b/zig/kernels/metal/lfm2/lfm2_segments.metal @@ -37,7 +37,7 @@ constant constexpr size_t SEG_PART = size_t(QH) * CHUNKS * SUBS; } // lfm2_attn_mma over the streams' 8-row tiles: tile tg.x is tiles[tg.x] = (stream << 16) | the stream's tile. -[[kernel]] void lfm2_attn_segs(const device bfloat* queries [[buffer(0)]], const device lfm2_seg* segs [[buffer(1)]], +[[max_total_threads_per_threadgroup(128)]] [[kernel]] void lfm2_attn_segs(const device bfloat* queries [[buffer(0)]], const device lfm2_seg* segs [[buffer(1)]], const device uint* tiles [[buffer(2)]], constant lfm2_attn_args& a [[buffer(3)]], device float* pm [[buffer(4)]], device float* pl [[buffer(5)]], device float* po [[buffer(6)]], uint3 tg [[threadgroup_position_in_grid]], diff --git a/zig/src/core/contract/lfm2_records.zig b/zig/src/core/contract/lfm2_records.zig new file mode 100644 index 0000000000..355559b1ff --- /dev/null +++ b/zig/src/core/contract/lfm2_records.zig @@ -0,0 +1,14 @@ +//! LFM2-MoE's shared-round and greedy-pick launches (lfm2.metal, lfm2_segments.metal), as the encoders launch them. +const std = @import("std"); +const M = @import("limits.zig").MetalLaunch; +pub const records = [_]M{ + .{ .name = "metal lfm2 argmax lfm2_argmax", .threads = 1024, .tg_memory = 384, .m1 = 1024, .m2 = 1024 }, + .{ .name = "metal lfm2 stream conv lfm2_conv_segs", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "metal lfm2 stream attention qk lfm2_qk_segs", .threads = 32, .tg_memory = 0, .m1 = 832, .m2 = 832 }, + .{ .name = "metal lfm2 stream attention lfm2_attn_segs", .threads = 128, .tg_memory = 18688, .reserved = 128 }, + .{ .name = "metal lfm2 stream attention merge lfm2_attn_merge_segs", .threads = 32, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, +}; + +test "the LFM2 launches fit every Apple family" { + try std.testing.expectEqual(@as(usize, 0), try @import("registry.zig").checkMetal(std.testing.allocator, &records, true)); +} diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 90f556180e..4637d84ddc 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -566,7 +566,7 @@ test { _ = paging; } /// The Metal collector keeps measured slots and the restored vision declarations. -pub const all_metal = metal ++ @import("delta_slots_records.zig").records ++ @import("flashnext_pending.zig").records ++ @import("simd_pairs_records.zig").records ++ @import("glm_rows_records.zig").records; +pub const all_metal = metal ++ @import("delta_slots_records.zig").records ++ @import("flashnext_pending.zig").records ++ @import("simd_pairs_records.zig").records ++ @import("glm_rows_records.zig").records ++ @import("lfm2_records.zig").records; const learning = @import("registry_learning.zig"); pub const gradient_pending = learning.gradient_pending; pub const qwen_learning_pending = learning.qwen_learning_pending; diff --git a/zig/tests/contract/metal_lfm2_argmax.zig b/zig/tests/contract/metal_lfm2_argmax.zig index a60e4e273f..423919447f 100644 --- a/zig/tests/contract/metal_lfm2_argmax.zig +++ b/zig/tests/contract/metal_lfm2_argmax.zig @@ -6,6 +6,7 @@ const sources = @import("kernel_sources"); const c = contract.spec; const g = contract.gates; const rig_ = @import("rig.zig"); +const controls = @import("metal_lfm2_controls.zig"); const vocab = 128000; pub const name = "metal lfm2 argmax"; @@ -33,13 +34,16 @@ pub const Probe = struct { a: std.mem.Allocator, rig: rig_.Rig, pipe: mtl.Pipeline, + mutation: controls.Mutation, rows: usize = 0, x: []u16 = &.{}, args: [2]rig_.Ref = undefined, out: rig_.Ref = undefined, - pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue) !Probe { - return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .pipe = try rig_.pipeline(device, sources.lfm2, "lfm2_argmax") }; + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, mutation: controls.Mutation) !Probe { + const text = try controls.source(a, mutation); + defer if (mutation != .none) a.free(text); + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .pipe = try rig_.pipeline(device, text, "lfm2_argmax"), .mutation = mutation }; } fn drop(p: *Probe) void { @@ -56,7 +60,9 @@ pub const Probe = struct { const vt = g.hooks(Probe); pub fn kernel(p: *Probe) g.Kernel { - return .{ .ctx = p, .vt = &vt, .spec = spec() }; + var s = spec(); + if (p.mutation != .none) s.name = p.mutation.name(); + return .{ .ctx = p, .vt = &vt, .spec = s, .control = p.mutation.gate() }; } pub fn prepare(p: *Probe, s: c.Shape, f: g.Fill, seed: u64) !bool { @@ -117,9 +123,3 @@ pub const Probe = struct { return null; } }; - -/// Gate 5 for the launch on this GPU: 1,024 threads, 32 simdgroups' best values, ids and flags. -pub fn limits(a: std.mem.Allocator, device: mtl.Device, only: ?[]const u8) !usize { - if (only) |s| if (std.mem.indexOf(u8, name, s) == null) return 0; - return @import("metal_limits.zig").onThisGpu(a, device, .{ .decl = .{ .name = name, .threads = 1024, .tg_memory = 384, .offline = false }, .function = "lfm2_argmax", .here = sources.lfm2, .pre_m5 = "" }, false); -} diff --git a/zig/tests/contract/metal_lfm2_controls.zig b/zig/tests/contract/metal_lfm2_controls.zig new file mode 100644 index 0000000000..655a97d4fc --- /dev/null +++ b/zig/tests/contract/metal_lfm2_controls.zig @@ -0,0 +1,51 @@ +//! Planted controls for LFM2's shared-round kernels: lfm2.metal's text with one fault, which one gate must report. +const std = @import("std"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); + +pub const Mutation = enum { + none, + argmax_last_tie, + argmax_no_flag, + conv_first_state, + attention_first_cache, + + pub fn name(m: Mutation) []const u8 { + return switch (m) { + .none => "", + .argmax_last_tie => "control: metal lfm2 argmax, a tie goes to the highest id", + .argmax_no_flag => "control: metal lfm2 argmax, a nonfinite logit is not flagged", + .conv_first_state => "control: metal lfm2 stream conv, every stream reads the first stream's state", + .attention_first_cache => "control: metal lfm2 stream attention, every stream writes the first stream's cache", + }; + } + + /// The gate that must go red. + pub fn gate(m: Mutation) ?contract.gates.Gate { + return switch (m) { + .none => null, + .argmax_no_flag => .nonfinite, + else => .shapes, + }; + } + + const Edit = struct { from: []const u8, to: []const u8 }; + + fn edit(m: Mutation) Edit { + return switch (m) { + .none => unreachable, + .argmax_last_tie => .{ .from = "(ov == bv && oi < bi)", .to = "(ov == bv && oi > bi)" }, + .argmax_no_flag => .{ .from = "if (!isfinite(f)) bad = 1;", .to = "if (false) bad = 1;" }, + .conv_first_state => .{ .from = "lfm2_conv_row(bcx + size_t(s.base) * 3 * D, s.conv,", .to = "lfm2_conv_row(bcx + size_t(s.base) * 3 * D, segs[0].conv," }, + .attention_first_cache => .{ .from = "queries + b * QH * HD, s.keys, s.values,", .to = "queries + b * QH * HD, segs[0].keys, segs[0].values," }, + }; + } +}; + +/// lfm2.metal's text as built for the kernel under test; a control's one edit must change it, or it proves nothing. +pub fn source(a: std.mem.Allocator, m: Mutation) ![]const u8 { + if (m == .none) return sources.lfm2; + const e = m.edit(); + if (std.mem.indexOf(u8, sources.lfm2, e.from) == null) return error.ControlEditMissing; + return std.mem.replaceOwned(u8, a, sources.lfm2, e.from, e.to); +} diff --git a/zig/tests/contract/metal_lfm2_streams.zig b/zig/tests/contract/metal_lfm2_streams.zig index 0a46f7a526..e9a59ef4f8 100644 --- a/zig/tests/contract/metal_lfm2_streams.zig +++ b/zig/tests/contract/metal_lfm2_streams.zig @@ -6,6 +6,7 @@ const sources = @import("kernel_sources"); const c = contract.spec; const g = contract.gates; const rig_ = @import("rig.zig"); +const controls = @import("metal_lfm2_controls.zig"); const d = 2048; const qh = 32; @@ -47,6 +48,7 @@ pub const Probe = struct { a: std.mem.Allocator, rig: rig_.Rig, which: Which, + mutation: controls.Mutation, pipes: [8]mtl.Pipeline, made: usize = 0, rows: usize = 0, @@ -64,15 +66,17 @@ pub const Probe = struct { golden: []u8 = &.{}, snap: rig_.Ref = undefined, - pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, w: Which) !Probe { + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, w: Which, mutation: controls.Mutation) !Probe { const list: []const [:0]const u8 = switch (w) { .attention => &.{ "lfm2_qk_segs", "lfm2_attn_segs", "lfm2_attn_merge_segs", "lfm2_qk", "lfm2_attn_mma", "lfm2_attn_mma_r2", "lfm2_attn_merge" }, .conv => &.{ "lfm2_conv_segs", "lfm2_conv" }, }; - var p: Probe = .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .which = w, .pipes = undefined }; + var p: Probe = .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .which = w, .mutation = mutation, .pipes = undefined }; + const text = try controls.source(a, mutation); + defer if (mutation != .none) a.free(text); errdefer for (p.pipes[0..p.made]) |x| x.deinit(); for (list) |n| { - p.pipes[p.made] = try rig_.pipeline(device, sources.lfm2, n); + p.pipes[p.made] = try rig_.pipeline(device, text, n); p.made += 1; } return p; @@ -96,7 +100,9 @@ pub const Probe = struct { const vt = g.hooks(Probe); pub fn kernel(p: *Probe) g.Kernel { - return .{ .ctx = p, .vt = &vt, .spec = spec(p.which) }; + var s = spec(p.which); + if (p.mutation != .none) s.name = p.mutation.name(); + return .{ .ctx = p, .vt = &vt, .spec = s, .control = p.mutation.gate() }; } fn len(p: *const Probe, j: usize) usize { @@ -470,19 +476,3 @@ fn same(x: u16, y: u16) bool { const nx = std.math.isNan(c.fromBf16(x)); return if (nx) std.math.isNan(c.fromBf16(y)) else x == y; } - -/// Gate 5 for the shared-round launches on this GPU. -pub fn limits(a: std.mem.Allocator, device: mtl.Device, only: ?[]const u8) !usize { - var problems: usize = 0; - const launches = [_]struct { []const u8, []const u8, u32, u32 }{ - .{ names[1], "lfm2_conv_segs", 256, 0 }, - .{ names[0] ++ " qk", "lfm2_qk_segs", 32, 0 }, - .{ names[0], "lfm2_attn_segs", 128, 18688 }, - .{ names[0] ++ " merge", "lfm2_attn_merge_segs", 32, 0 }, - }; - for (launches) |l| { - if (only) |s| if (std.mem.indexOf(u8, l[0], s) == null) continue; - problems += try @import("metal_limits.zig").onThisGpu(a, device, .{ .decl = .{ .name = l[0], .threads = l[2], .tg_memory = l[3], .offline = false }, .function = l[1], .here = sources.lfm2, .pre_m5 = "" }, false); - } - return problems; -} diff --git a/zig/tests/contract/metal_limits.zig b/zig/tests/contract/metal_limits.zig index a5a8fc143f..dbe7a12df3 100644 --- a/zig/tests/contract/metal_limits.zig +++ b/zig/tests/contract/metal_limits.zig @@ -36,6 +36,7 @@ fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch) !La text = tf.flashnext_vision_contract.source(w); return .{ .decl = d, .function = fun, .here = text, .pre_m5 = text }; } + if (std.mem.startsWith(u8, d.name, "metal lfm2 ")) return .{ .decl = d, .function = fun, .here = sources.lfm2, .pre_m5 = sources.lfm2 }; if (std.mem.startsWith(u8, d.name, "drama3 ")) { const text_d3 = try @import("metal_drama3.zig").source(a, d.name); return .{ .decl = d, .function = fun, .here = text_d3, .pre_m5 = text_d3 }; diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index c0c20308e4..86e8a786fd 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -25,6 +25,7 @@ const gemma_embed = @import("metal_gemma4_embed.zig"); const gemma_mixed = @import("metal_gemma4_mixed.zig"); const lfm2_argmax = @import("metal_lfm2_argmax.zig"); const lfm2_streams = @import("metal_lfm2_streams.zig"); +const lfm2_controls = @import("metal_lfm2_controls.zig"); const g = contract.gates; const usage = "usage: tf-kernel-contract [--only TEXT] [--no-gates] [--no-limits] [--offline DIR] [--seed N] [--repro] [--rotary-long] [--rotary-accuracy FILE]"; @@ -153,21 +154,42 @@ pub fn main(init: std.process.Init) !u8 { if (!tensor) return error.TensorUnitsRequired; try kernels.append(a, mixed.kernel()); } - var argmax: ?lfm2_argmax.Probe = null; - defer if (argmax) |*p| p.deinit(); - if (wanted(o, lfm2_argmax.name)) { - argmax = try lfm2_argmax.Probe.init(a, device, queue); - try kernels.append(a, argmax.?.kernel()); + var lfm2_list: std.ArrayList(*lfm2_argmax.Probe) = .empty; + var lfm2_streams_list: std.ArrayList(*lfm2_streams.Probe) = .empty; + defer { + for (lfm2_list.items) |p| { + p.deinit(); + a.destroy(p); + } + for (lfm2_streams_list.items) |p| { + p.deinit(); + a.destroy(p); + } + lfm2_list.deinit(a); + lfm2_streams_list.deinit(a); } - var streams_list: [2]lfm2_streams.Probe = undefined; - var streams_made: usize = 0; - defer for (streams_list[0..streams_made]) |*p| p.deinit(); - for (std.enums.values(lfm2_streams.Which)) |w| { - if (!wanted(o, lfm2_streams.spec(w).name)) continue; - streams_list[streams_made] = try lfm2_streams.Probe.init(a, device, queue, w); - streams_made += 1; + for (std.enums.values(lfm2_controls.Mutation)) |m| { + const argmax_side = m == .none or m == .argmax_last_tie or m == .argmax_no_flag; + if (argmax_side) { + if (!wanted(o, if (m == .none) lfm2_argmax.name else m.name())) continue; + const p = try a.create(lfm2_argmax.Probe); + errdefer a.destroy(p); + p.* = try lfm2_argmax.Probe.init(a, device, queue, m); + try lfm2_list.append(a, p); + try kernels.append(a, p.kernel()); + if (m != .none) continue; + } + for (std.enums.values(lfm2_streams.Which)) |w| { + const control_for = (m == .conv_first_state and w == .conv) or (m == .attention_first_cache and w == .attention); + if (m != .none and !control_for) continue; + if (!wanted(o, if (m == .none) lfm2_streams.spec(w).name else m.name())) continue; + const p = try a.create(lfm2_streams.Probe); + errdefer a.destroy(p); + p.* = try lfm2_streams.Probe.init(a, device, queue, w, m); + try lfm2_streams_list.append(a, p); + try kernels.append(a, p.kernel()); + } } - for (streams_list[0..streams_made]) |*p| try kernels.append(a, p.kernel()); var rows_list: std.ArrayList(*rows_k.Row) = .empty; var proj_list: std.ArrayList(*affine_k.Proj) = .empty; defer for (rows_list.items) |r| r.deinit(); @@ -307,8 +329,6 @@ pub fn main(init: std.process.Init) !u8 { lim += try gemma.limits(a, device, o.only); lim += try gemma_embed.limits(a, device, o.only); lim += try gemma_mixed.limits(a, device, o.only); - lim += try lfm2_argmax.limits(a, device, o.only); - lim += try lfm2_streams.limits(a, device, o.only); } const t2 = std.Io.Clock.awake.now(init.io); const ms = struct { From 05d83e852fb7142a1c959c8a92d9dc582185ddc9 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 19:34:55 +0100 Subject: [PATCH 09/10] lfm2: serve only GPUs from the M3 on, and say an M3 is untested M1 and M2 cap a pipeline's threads as its registers rise, and no run covers them. The server refuses them by name at open. An M3 is served, with one log line that it is not yet tested with shared rounds. --- zig/src/native/lfm2_moe.zig | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/zig/src/native/lfm2_moe.zig b/zig/src/native/lfm2_moe.zig index 5340d807de..1512eaef4c 100644 --- a/zig/src/native/lfm2_moe.zig +++ b/zig/src/native/lfm2_moe.zig @@ -35,6 +35,19 @@ const Host = struct { const cannot = "the native LFM2 engine cannot load this checkpoint ({s})"; +/// The first GPU generation served: applegpu_g15 (M3); M1 is 13 and M2 is 14. +const first_generation = 15; + +/// Whether a GPU of this generation (0: not reported) serves; on M1 and M2 a pipeline's threads fall as registers rise. +fn generationServed(generation: u32) bool { + return generation == 0 or generation >= first_generation; +} + +test "M1 and M2 are refused, M3 and later are served" { + try std.testing.expect(!generationServed(13) and !generationServed(14)); + try std.testing.expect(generationServed(15) and generationServed(16) and generationServed(17) and generationServed(0)); +} + pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { const pool = mtl.objc.Pool.push(); defer pool.pop(); @@ -57,6 +70,11 @@ pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Ope return null; }; errdefer model.deinit(); + if (!generationServed(model.device.generation())) { + problem.* = "LFM2.5-24B-A2B needs an M3 or later GPU: M1 and M2 are not qualified for its shared rounds"; + return null; + } + if (model.device.generation() == first_generation) std.log.info("lfm2: an M3 is not yet tested with shared rounds", .{}); const h = try gpa.create(Host); errdefer gpa.destroy(h); h.gpa = gpa; From 70f74ca3464b4234b60b61691eb004f2c3e2cbcc Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 20:40:13 +0100 Subject: [PATCH 10/10] lfm2: M1 and M2 are served again The shared-round kernels carry recorded M1 and M2 launch limits in the kernel contract, and no run covers those chips. The server no longer refuses them. --- zig/src/native/lfm2_moe.zig | 18 ------------------ 1 file changed, 18 deletions(-) diff --git a/zig/src/native/lfm2_moe.zig b/zig/src/native/lfm2_moe.zig index 1512eaef4c..5340d807de 100644 --- a/zig/src/native/lfm2_moe.zig +++ b/zig/src/native/lfm2_moe.zig @@ -35,19 +35,6 @@ const Host = struct { const cannot = "the native LFM2 engine cannot load this checkpoint ({s})"; -/// The first GPU generation served: applegpu_g15 (M3); M1 is 13 and M2 is 14. -const first_generation = 15; - -/// Whether a GPU of this generation (0: not reported) serves; on M1 and M2 a pipeline's threads fall as registers rise. -fn generationServed(generation: u32) bool { - return generation == 0 or generation >= first_generation; -} - -test "M1 and M2 are refused, M3 and later are served" { - try std.testing.expect(!generationServed(13) and !generationServed(14)); - try std.testing.expect(generationServed(15) and generationServed(16) and generationServed(17) and generationServed(0)); -} - pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { const pool = mtl.objc.Pool.push(); defer pool.pop(); @@ -70,11 +57,6 @@ pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Ope return null; }; errdefer model.deinit(); - if (!generationServed(model.device.generation())) { - problem.* = "LFM2.5-24B-A2B needs an M3 or later GPU: M1 and M2 are not qualified for its shared rounds"; - return null; - } - if (model.device.generation() == first_generation) std.log.info("lfm2: an M3 is not yet tested with shared rounds", .{}); const h = try gpa.create(Host); errdefer gpa.destroy(h); h.gpa = gpa;