diff --git a/CHANGELOG.md b/CHANGELOG.md index ac2b0fc8ef..72471ab285 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- LFM2.5-24B-A2B on Metal can hand an external drafter the rows of the layers it reads (its taps), copied on the GPU by a new kernel, `lfm2_tap`, that the kernel contract checks bit for bit with two planted controls (#1390, part of #1374). A drafting stream resumes a kept prompt state only when it also holds the drafter's context. Nothing uses the taps yet: without a drafter none are kept and every forward runs as before. Thanks to @BlivionIaG. + - Nemotron on CUDA: the MoE residual's norm, Mamba's group norm, the MTP input norms and the bf16 token rows run on the shared norm and embedding kernels instead of Nemotron's own copies, which are removed. The shared kernels keep Nemotron's summation order, so replies are the same tokens as before, and one fewer CUDA image loads (#1353, step 1 of #548). Thanks to @jschmied. - Qwen3.8-27B on two Macs, first step: the pair seam's Metal kernels, which exchange whole output rows between the two Macs as bf16 copies, so every sum stays on one Mac and no partial sum crosses the link, plus a one-Mac loopback peer and six kernel-contract controls that each go red as planted (#1432). Nothing serves through it yet, and one Mac's projections keep the same reduction order, so replies are byte-identical. Thanks to @akol1. diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index 8a16d7432e..efa248606d 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -462,6 +462,8 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s b.step("test-nemotron-round-depth-host", "Priced head depth, nonlanding levels and fixed-depth controls without devices").dependOn(&round_depth.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/core/contract/lfm2_records.zig b/zig/src/core/contract/lfm2_records.zig index 355559b1ff..e3f7179595 100644 --- a/zig/src/core/contract/lfm2_records.zig +++ b/zig/src/core/contract/lfm2_records.zig @@ -7,6 +7,7 @@ pub const records = [_]M{ .{ .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 }, + .{ .name = "metal lfm2 tap lfm2_tap", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, }; test "the LFM2 launches fit every Apple family" { 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 4cc2deba22..59afec1320 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -90,6 +90,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_controls.zig b/zig/tests/contract/metal_lfm2_controls.zig index 655a97d4fc..4d875b0fe8 100644 --- a/zig/tests/contract/metal_lfm2_controls.zig +++ b/zig/tests/contract/metal_lfm2_controls.zig @@ -9,6 +9,8 @@ pub const Mutation = enum { argmax_no_flag, conv_first_state, attention_first_cache, + tap_packed_rows, + tap_clobbers_column, pub fn name(m: Mutation) []const u8 { return switch (m) { @@ -17,6 +19,8 @@ pub const Mutation = enum { .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", + .tap_packed_rows => "control: metal lfm2 tap copy, rows packed at the hidden size, not the taps' row", + .tap_clobbers_column => "control: metal lfm2 tap copy, a column of another tap is overwritten", }; } @@ -38,6 +42,8 @@ pub const Mutation = enum { .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," }, + .tap_packed_rows => .{ .from = "(taps + size_t(i.y) * at.x + at.y))[i.x]", .to = "(taps + size_t(i.y) * D + at.y))[i.x]" }, + .tap_clobbers_column => .{ .from = "[i.x] = from[i.x];\n}\n", .to = "[i.x] = from[i.x];\n *(device bfloat*)(taps + size_t(i.y) * at.x) = 0;\n}\n" }, }; } }; diff --git a/zig/tests/contract/metal_lfm2_tap.zig b/zig/tests/contract/metal_lfm2_tap.zig new file mode 100644 index 0000000000..021baffc9e --- /dev/null +++ b/zig/tests/contract/metal_lfm2_tap.zig @@ -0,0 +1,137 @@ +//! 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 controls = @import("metal_lfm2_controls.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, + mutation: controls.Mutation, + 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, 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_tap"), .mutation = mutation }; + } + + 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 { + var s = spec(); + if (p.mutation != .none) s.name = p.mutation.name(); + return .{ .ctx = p, .vt = &vt, .spec = s, .control = p.mutation.gate() }; + } + + 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; + } +}; diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 34567cef0f..1cf2a29ab8 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -28,6 +28,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 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]"; @@ -193,6 +194,22 @@ pub fn main(init: std.process.Init) !u8 { try kernels.append(a, p.kernel()); } } + var tap_list: std.ArrayList(*lfm2_tap.Probe) = .empty; + defer { + for (tap_list.items) |p| { + p.deinit(); + a.destroy(p); + } + tap_list.deinit(a); + } + for ([_]lfm2_controls.Mutation{ .none, .tap_packed_rows, .tap_clobbers_column }) |m| { + if (!wanted(o, if (m == .none) lfm2_tap.name else m.name())) continue; + const p = try a.create(lfm2_tap.Probe); + errdefer a.destroy(p); + p.* = try lfm2_tap.Probe.init(a, device, queue, m); + try tap_list.append(a, 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();