lfm2: keep a drafter's tapped rows on every forward

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.

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

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<bfloat, 8>* from = (const device vec<bfloat, 8>*)(h + size_t(i.y) * D);
+  ((device vec<bfloat, 8>*)(taps + size_t(i.y) * at.x + at.y))[i.x] = from[i.x];
+}

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" {

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;

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

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

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 {

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

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,13 +157,16 @@ 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;
         out.memory = .{ .gpa = gpa, .device = device };
         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);

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