From 06a5236ce70bd07123d96a018c30905109252cb2 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 08:53:38 +0100 Subject: [PATCH 01/46] server: capabilities --json's Engines struct moves to cli_capabilities.zig cli_capabilities.zig writes the capabilities document, but the Engines struct it reads still sat in cli.zig, which is 598 lines. The struct moves next to its writer and cli.zig re-exports it, so callers do not change. cli.zig drops to 589 lines. capabilities --json prints the same bytes as on main. --- zig/src/server/cli.zig | 13 ++----------- zig/src/server/cli_capabilities.zig | 16 ++++++++++++++-- 2 files changed, 16 insertions(+), 13 deletions(-) diff --git a/zig/src/server/cli.zig b/zig/src/server/cli.zig index 0e2da51172..d36a933aa7 100644 --- a/zig/src/server/cli.zig +++ b/zig/src/server/cli.zig @@ -571,17 +571,8 @@ pub fn request(a: Allocator, dir: []const u8, model_type: []const u8, args: Args }); } -/// What ``capabilities --json`` reports about the engine side: its version, chip, backends and families. -pub const Engines = struct { - version: []const u8, - chip: ?[]const u8 = null, - optional_libraries: []const @import("sys").optional_libraries.Status = &.{}, - backends: []const []const u8 = &.{}, - /// model_type to the weight formats it reads, as gate entries name them. - families: []const api.Family = &.{}, -}; - -/// The native capabilities document: the table's native flags and the honoured variables. +/// ``capabilities --json`` (cli_capabilities.zig): the engine side, then the table's native flags and variables. +pub const Engines = @import("cli_capabilities.zig").Engines; pub const capabilities = @import("cli_capabilities.zig").capabilities; const cli = @This(); diff --git a/zig/src/server/cli_capabilities.zig b/zig/src/server/cli_capabilities.zig index 7ee1fd2283..4c5d8607d1 100644 --- a/zig/src/server/cli_capabilities.zig +++ b/zig/src/server/cli_capabilities.zig @@ -1,8 +1,20 @@ -//! The native capabilities document: the table's native flags and the honoured variables. +//! ``capabilities --json``: the engine's version, chip, backends and families, then the native flags and variables. const std = @import("std"); +const api = @import("engine_api"); const cli = @import("cli.zig"); -pub fn capabilities(w: *std.Io.Writer, e: cli.Engines) !void { +/// What ``capabilities --json`` reports about the engine side: its version, chip, backends and families. +pub const Engines = struct { + version: []const u8, + chip: ?[]const u8 = null, + optional_libraries: []const @import("sys").optional_libraries.Status = &.{}, + backends: []const []const u8 = &.{}, + /// model_type to the weight formats it reads, as gate entries name them. + families: []const api.Family = &.{}, +}; + +/// The native capabilities document: the table's native flags and the honoured variables. +pub fn capabilities(w: *std.Io.Writer, e: Engines) !void { try w.print("{{\"schema\": 1, \"engine\": \"zig\", \"version\": \"{s}\", \"chip\": ", .{e.version}); if (e.chip) |c| try w.print("\"{s}\"", .{c}) else try w.writeAll("null"); try w.writeAll(", \"backends\": ["); From 9a8d654efdab297ab07d5a09b27b4d694895690b Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:01:08 +0100 Subject: [PATCH 02/46] lanes: Drafted passes held windows to a target that reads them itself With Drafted.held_on_target set, Drafted.verify gives the target its windows with Window.held as the round loop made them, so a target that reads held drafts from the drafter on the device skips the host readback. A window of a stream with fillers still gets host tokens. The flag is off by default; without it verify runs as before and the new check is not evaluated. drafted_held_test.zig checks the three cases: held windows pass through, a stream with fillers gets its pending token as host tokens, and with the flag off the drafter's drafts are read back as host tokens. --- zig/src/core/lanes/drafted.zig | 6 +- zig/src/core/lanes/drafted_held_test.zig | 78 ++++++++++++++++++++++++ zig/src/core/lanes/drafted_test.zig | 1 + 3 files changed, 84 insertions(+), 1 deletion(-) create mode 100644 zig/src/core/lanes/drafted_held_test.zig diff --git a/zig/src/core/lanes/drafted.zig b/zig/src/core/lanes/drafted.zig index 999aaeab6a..1aacfa30e3 100644 --- a/zig/src/core/lanes/drafted.zig +++ b/zig/src/core/lanes/drafted.zig @@ -29,6 +29,8 @@ pub const Drafted = struct { absorbs: std.ArrayList(dr.Absorb) = .empty, holds: std.ArrayList(dr.Hold) = .empty, firsts: std.ArrayList(u32) = .empty, + /// The target reads held drafts from the drafter itself (Window.held), so windows without fillers pass as they are. + held_on_target: bool = false, /// The wrapper, with the target told the drafter's taps before any forward. pub fn init(gpa: Allocator, target: be.Backend, drafter: dr.Drafter) !Drafted { @@ -217,11 +219,13 @@ pub const Drafted = struct { lane.tree_first = null; }; var total: usize = 0; + var through = x.held_on_target; for (windows) |w| { if (w.held > 0 and w.tokens.len != 0) return error.MixedDrafts; total += w.held; + if (through and w.held > 0 and (x.lanes.getPtr(w.stream) orelse return error.NotDrafting).filler > 0) through = false; } - if (total == 0) return x.target.verify(x.windows.items, out); + if (total == 0 or through) return x.target.verify(x.windows.items, out); try x.tokens.resize(x.gpa, total); x.streams.clearRetainingCapacity(); x.outs.clearRetainingCapacity(); diff --git a/zig/src/core/lanes/drafted_held_test.zig b/zig/src/core/lanes/drafted_held_test.zig new file mode 100644 index 0000000000..c8a6010744 --- /dev/null +++ b/zig/src/core/lanes/drafted_held_test.zig @@ -0,0 +1,78 @@ +//! Drafted.verify with held_on_target: windows of drafting streams reach the target with Window.held as asked. +const std = @import("std"); +const be = @import("backend.zig"); +const dr = @import("drafter.zig"); +const Stream = @import("stream.zig").Stream; +const Drafted = @import("drafted.zig").Drafted; + +/// A target that records each verify's held rows and host tokens, and a drafter that only counts host readbacks. +const Fake = struct { + seen: u32 = 0, + tokens: [4]u32 = @splat(0), + readbacks: u32 = 0, + + fn cast(ptr: *anyopaque) *Fake { + return @ptrCast(@alignCast(ptr)); + } + fn verify(ptr: *anyopaque, windows: []const be.Window, _: []be.Verified) anyerror!void { + const f = cast(ptr); + f.seen = windows[0].held; + @memcpy(f.tokens[0..windows[0].tokens.len], windows[0].tokens); + } + fn held(ptr: *anyopaque, _: []const *Stream, out: []const []u32) anyerror!void { + cast(ptr).readbacks += 1; + for (out) |o| @memset(o, 5); + } + fn none(_: *anyopaque) []const u32 { + return &.{}; + } + fn facts(_: *anyopaque) dr.Facts { + return .{ .depth = 3, .step_ms = 0, .prior = &.{}, .stream_bytes = 0 }; + } + fn unused(_: *anyopaque, _: *Stream) anyerror!void { + return error.Unused; + } + fn absorb(_: *anyopaque, _: []const dr.Absorb) anyerror!void { + return error.Unused; + } + fn hold(_: *anyopaque, _: []const dr.Hold) anyerror!void { + return error.Unused; + } + fn release(_: *anyopaque, _: *Stream) void {} + fn first(_: *anyopaque, _: *Stream, _: u64) anyerror!u64 { + return error.Unused; + } + fn queue(_: *anyopaque, _: *Stream, _: be.Feed, _: u64) anyerror!u64 { + return error.Unused; + } + fn read(_: *anyopaque, _: u64) anyerror!u32 { + return error.Unused; + } + fn keep(_: *anyopaque, _: []const be.Window, _: []const []const u32) anyerror!void {} + fn draft(_: *anyopaque, _: []const be.DraftRequest) anyerror!void {} +}; + +test "held_on_target passes held windows through, and a stream with fillers still gets host tokens" { + var f: Fake = .{}; + const target: be.Backend = .{ .ptr = &f, .vtable = &.{ .prefill = Fake.unused, .first = Fake.first, .queue = Fake.queue, .read = Fake.read, .verify = Fake.verify, .keep = Fake.keep, .draft = Fake.draft, .release = Fake.release } }; + const drafter: dr.Drafter = .{ .ptr = &f, .vtable = &.{ .taps = Fake.none, .facts = Fake.facts, .open = Fake.unused, .absorb = Fake.absorb, .hold = Fake.hold, .held = Fake.held, .release = Fake.release } }; + var x: Drafted = .{ .gpa = std.testing.allocator, .target = target, .drafter = drafter, .held_on_target = true }; + defer x.deinit(); + var s: Stream = undefined; + try x.lanes.put(std.testing.allocator, &s, .{ .rows = true }); + var sampled: [4]u32 = undefined; + var drafts: [3]u32 = undefined; + const w = [_]be.Window{.{ .stream = &s, .pending = 9, .held = 3, .tokens = &.{}, .parents = null, .positions = &.{ 1, 2, 3, 4 } }}; + var o = [_]be.Verified{.{ .sampled = &sampled, .drafts = &drafts }}; + try x.backend().verify(&w, &o); + try std.testing.expectEqual(@as(u32, 3), f.seen); + try std.testing.expectEqual(@as(u32, 0), f.readbacks); + x.lanes.getPtr(&s).?.filler = 3; // fillers: the pending token as host tokens, no readback + try x.backend().verify(&w, &o); + try std.testing.expectEqual(@as(u32, 0), f.seen); + try std.testing.expectEqualSlices(u32, &.{ 9, 9, 9 }, f.tokens[0..3]); + x.held_on_target = false; // off: the drafter's drafts as host tokens + try x.backend().verify(&w, &o); + try std.testing.expectEqual(@as(u32, 1), f.readbacks); + try std.testing.expectEqualSlices(u32, &.{ 5, 5, 5 }, f.tokens[0..3]); +} diff --git a/zig/src/core/lanes/drafted_test.zig b/zig/src/core/lanes/drafted_test.zig index 1c313bc305..be10023409 100644 --- a/zig/src/core/lanes/drafted_test.zig +++ b/zig/src/core/lanes/drafted_test.zig @@ -594,4 +594,5 @@ test "the drafter wrapper sees the committed callback before its first hold" { test { _ = @import("drafted_role_test.zig"); _ = @import("drafted_step_test.zig"); + _ = @import("drafted_held_test.zig"); } From 47eecb523aa3902616fdcda85d353f7aea7c5acb Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:02:12 +0100 Subject: [PATCH 03/46] lanes: a lone stream's pipelined steps move to pipelined.zig engine.zig is 592 lines and the shared-round guard needs room in it. pipelinedRound and landInflight move unchanged to pipelined.zig as round and land; Engine.familyRound and Result become public for them. No behaviour change. --- zig/src/core/lanes/engine.zig | 67 +++--------------------------- zig/src/core/lanes/pipelined.zig | 70 ++++++++++++++++++++++++++++++++ 2 files changed, 75 insertions(+), 62 deletions(-) create mode 100644 zig/src/core/lanes/pipelined.zig diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index a92b27ae7f..1664587c6d 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -16,6 +16,7 @@ const lone_costs = @import("lone_costs.zig"); const open = @import("open.zig"); const handover = @import("handover.zig"); const trim = @import("trim.zig"); +const pipelined = @import("pipelined.zig"); const Stream = sm.Stream; const Feed = be.Feed; const LogRow = @import("logprob.zig").Row; @@ -24,7 +25,7 @@ const f = ev.f; const str = trail.str; const int = trail.int; const Outcome = struct { got: []const u32, path: []const u32, cut: ?usize, follow: []const u32 }; -const Result = struct { rows: u32, keep: u32, got: []const u32 }; +pub const Result = struct { rows: u32, keep: u32, got: []const u32 }; const Unread = struct { event: usize, handle: u64 }; pub const Engine = struct { const pieces_ = @import("pieces.zig"); @@ -124,7 +125,7 @@ pub const Engine = struct { if (live.items.len > 1 and e.cfg.family_streams) { // a shared round is synchronous: a stream that ran a step ahead lands its queued token first for (live.items) |s| { - if (s.inflight != null) try e.landInflight(s); + if (s.inflight != null) try pipelined.land(e, s); } var kept: usize = 0; for (live.items) |s| { @@ -147,7 +148,7 @@ pub const Engine = struct { for (live.items) |s| { // logprobs and grammars take verify windows: rows' logits are read, or bits made, before a forward const windowed = (e.cfg.family_mtp and s.drafts) or !e.cfg.pipelined or s.logprobs != null or s.grammar != null; - const r = if (windowed) try e.familyRound(s, null) else try e.pipelinedRound(s); + const r = if (windowed) try e.familyRound(s, null) else try pipelined.round(e, s); s.min_rows = if (s.min_rows == 0) r.rows else @min(s.min_rows, r.rows); } } @@ -195,7 +196,7 @@ pub const Engine = struct { } /// One stream's round: verify pending and drafted tokens together, keep drafts to the first mismatch. - fn familyRound(e: *Engine, s: *Stream, copied: ?win.Copy) !Result { + pub fn familyRound(e: *Engine, s: *Stream, copied: ?win.Copy) !Result { e.clock.start(); const priced = s.priced; s.priced = null; @@ -492,64 +493,6 @@ pub const Engine = struct { } } - /// One token, the next forward queued first; a copied continuation ahead switches to verify windows. - fn pipelinedRound(e: *Engine, s: *Stream) !Result { - const before = s.mode; - // an adopted stream holds neither a mode nor a queued step: it starts where a drain step leaves one - const mode: sm.Mode = s.mode orelse if (s.inflight == null) .verify else .pipe; - var result: Result = .{ .rows = 1, .keep = 1, .got = &.{} }; - if (mode == .verify or mode == .exit) { - const verifying = mode == .verify and e.cfg.family_width >= 2; - const proposal: win.Copy = if (verifying) try win.copyProposal(e, s) else .{}; - if (proposal.tokens.len > 0) { - const copied = try e.arena.allocator().dupe(u32, proposal.tokens); - result = try e.familyRound(s, .{ .tokens = copied, .recalled = proposal.recalled }); - if (result.keep == 1 or s.force.items.len > 0) s.mode = .exit; - } else { - // nothing to copy: back to steps queued ahead (this one lands next round) - s.mode = .pipe; - try e.queueNext(s, .{ .value = s.pending.? }); - } - try trail.pipe(e, s, before, result.got); - return result; - } - var current: Feed = .{ .handle = s.inflight.? }; - s.inflight = null; - if (try e.forcedNext(s)) |t| current = .{ .value = t }; - var token: u32 = undefined; - if (mode == .drain) { - token = try e.readFeed(current); // the last queued step: no new one - s.mode = .verify; - } else { - try e.queueNext(s, current); // the GPU starts the next step first - token = try e.readFeed(current); - } - s.rounds += 1; - const landed = try s.commit(e.gpa, &.{token}, &.{}); - s.pending = token; - if (s.recall) |r| r.noteHead(1); // a step lands one token: the bar a copy must beat - if (e.cfg.root_copies) s.root_copy_miss = false; - const piping = e.cfg.family_width >= 2 and (s.mode orelse .pipe) == .pipe and !s.finished and s.drafts; - if (piping and (try win.copyProposal(e, s)).tokens.len > 0) - s.mode = .drain; // a copy window is ahead: land the queued step - result.got = try e.arena.allocator().dupe(u32, s.emitted()[s.emitted().len - landed ..]); - try trail.pipe(e, s, before, result.got); - return result; - } - - /// Land the queued token without queuing another (a shared round is synchronous). - fn landInflight(e: *Engine, s: *Stream) !void { - var current: Feed = .{ .handle = s.inflight.? }; - s.inflight = null; - if (try e.forcedNext(s)) |t| current = .{ .value = t }; - const token = try e.readFeed(current); - s.rounds += 1; - const landed = try s.commit(e.gpa, &.{token}, &.{}); - s.pending = token; - s.mode = .verify; - try trail.event(e, &.{ f("ev", str("land")), f("stream", str(s.id)), f("got", .{ .u32s = s.emitted()[s.emitted().len - landed ..] }) }); - } - /// The first token's draw: among the grammar's allowed tokens when the stream has one that constrains it. pub fn drawFirst(e: *Engine, s: *Stream, position: u64) !u64 { const g = s.grammar orelse return e.backend.first(s, position); diff --git a/zig/src/core/lanes/pipelined.zig b/zig/src/core/lanes/pipelined.zig new file mode 100644 index 0000000000..99de33304b --- /dev/null +++ b/zig/src/core/lanes/pipelined.zig @@ -0,0 +1,70 @@ +//! A lone stream's pipelined steps: the next forward queued before the last one's token is read. +const be = @import("backend.zig"); +const sm = @import("stream.zig"); +const ev = @import("events.zig"); +const win = @import("windows.zig"); +const trail = @import("trail.zig"); +const Engine = @import("engine.zig").Engine; +const Result = @import("engine.zig").Result; +const Stream = sm.Stream; +const Feed = be.Feed; +const f = ev.f; +const str = trail.str; + +/// One token, the next forward queued first; a copied continuation ahead switches to verify windows. +pub fn round(e: *Engine, s: *Stream) !Result { + const before = s.mode; + // an adopted stream holds neither a mode nor a queued step: it starts where a drain step leaves one + const mode: sm.Mode = s.mode orelse if (s.inflight == null) .verify else .pipe; + var result: Result = .{ .rows = 1, .keep = 1, .got = &.{} }; + if (mode == .verify or mode == .exit) { + const verifying = mode == .verify and e.cfg.family_width >= 2; + const proposal: win.Copy = if (verifying) try win.copyProposal(e, s) else .{}; + if (proposal.tokens.len > 0) { + const copied = try e.arena.allocator().dupe(u32, proposal.tokens); + result = try e.familyRound(s, .{ .tokens = copied, .recalled = proposal.recalled }); + if (result.keep == 1 or s.force.items.len > 0) s.mode = .exit; + } else { + // nothing to copy: back to steps queued ahead (this one lands next round) + s.mode = .pipe; + try e.queueNext(s, .{ .value = s.pending.? }); + } + try trail.pipe(e, s, before, result.got); + return result; + } + var current: Feed = .{ .handle = s.inflight.? }; + s.inflight = null; + if (try e.forcedNext(s)) |t| current = .{ .value = t }; + var token: u32 = undefined; + if (mode == .drain) { + token = try e.readFeed(current); // the last queued step: no new one + s.mode = .verify; + } else { + try e.queueNext(s, current); // the GPU starts the next step first + token = try e.readFeed(current); + } + s.rounds += 1; + const landed = try s.commit(e.gpa, &.{token}, &.{}); + s.pending = token; + if (s.recall) |r| r.noteHead(1); // a step lands one token: the bar a copy must beat + if (e.cfg.root_copies) s.root_copy_miss = false; + const piping = e.cfg.family_width >= 2 and (s.mode orelse .pipe) == .pipe and !s.finished and s.drafts; + if (piping and (try win.copyProposal(e, s)).tokens.len > 0) + s.mode = .drain; // a copy window is ahead: land the queued step + result.got = try e.arena.allocator().dupe(u32, s.emitted()[s.emitted().len - landed ..]); + try trail.pipe(e, s, before, result.got); + return result; +} + +/// Land the queued token without queuing another (a shared round is synchronous). +pub fn land(e: *Engine, s: *Stream) !void { + var current: Feed = .{ .handle = s.inflight.? }; + s.inflight = null; + if (try e.forcedNext(s)) |t| current = .{ .value = t }; + const token = try e.readFeed(current); + s.rounds += 1; + const landed = try s.commit(e.gpa, &.{token}, &.{}); + s.pending = token; + s.mode = .verify; + try trail.event(e, &.{ f("ev", str("land")), f("stream", str(s.id)), f("got", .{ .u32s = s.emitted()[s.emitted().len - landed ..] }) }); +} From 5d9fcf68f87dd46ff653c17a2fd937fd09074ce1 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:03:43 +0100 Subject: [PATCH 04/46] lanes: a shared round drafts only while drafting measures faster (Model.shared_guard) A drafter whose hold is a pass of its own costs a shared round more than the depth rule's head steps price: the rule treats the rest of a round's time as overhead that plain rounds pay too. Facts.shared_guard (and Model.shared_guard for a backend) lets the engine compare the two modes directly. shared_guard.zig keeps, for each stream count, committed tokens over the round's wall time in plain stretches (no budgets, so no holds) and in drafted stretches. It starts plain, measures plain rounds for 16 rounds and drafted ones for 4, runs the faster mode for 512 rounds and doubles that, up to 64 times, while the same mode wins again; then it measures the other mode again. Drafted rounds must beat plain ones by 5%. A mode's first round still carries the other mode's drafts or holds and is not measured. An entry is allocated once a stream count; later rounds take no allocation. A plain stretch changes only how many drafts a window carries, never what the target draws, so replies stay the same tokens. Without the flag the engine never consults the guard. Tests: the guard's choice and re-measures, ever rarer re-measures of a losing mode, its skipped first rounds, no allocation after a stream count's first round, and replies equal to one-token decode with the guard on and off. --- zig/src/core/lanes/config.zig | 3 + zig/src/core/lanes/drafted.zig | 1 + zig/src/core/lanes/drafted_guard_test.zig | 51 ++++++++ zig/src/core/lanes/drafter.zig | 1 + zig/src/core/lanes/engine.zig | 9 +- zig/src/core/lanes/lanes.zig | 2 + zig/src/core/lanes/shared_guard.zig | 145 ++++++++++++++++++++++ 7 files changed, 211 insertions(+), 1 deletion(-) create mode 100644 zig/src/core/lanes/drafted_guard_test.zig create mode 100644 zig/src/core/lanes/shared_guard.zig diff --git a/zig/src/core/lanes/config.zig b/zig/src/core/lanes/config.zig index cb2159b03a..a7eac1ea7d 100644 --- a/zig/src/core/lanes/config.zig +++ b/zig/src/core/lanes/config.zig @@ -36,6 +36,7 @@ pub const Model = struct { shared_costs: []const Cost = &.{}, // shared forwards' ms by total rows (timed at load) draft_probabilities: bool = false, // the head gives each drafted node's chance of landing draft_streams: bool = false, // the head drafts every stream of a shared round in one batch + shared_guard: bool = false, // a shared round drafts only while drafted rounds measure faster (shared_guard.zig) head_trees: bool = false, // the head drafts trees of lanes (DraftRequest.lanes) and holds them join_tail: bool = false, // a prompt in pieces leaves its last token for the next shared round's window (Backend.join) piece_runs: bool = false, // a round's pieces go to the backend in one call (Backend.pieces: one forward over them) @@ -68,6 +69,7 @@ pub const Config = struct { fill_match: u32 = 8, // context tokens a suffix-match branch must match (shorter ones are coincidental phrases) fill_floor: f64 = 0.35, // the share of grafted rows a stream must keep to go on reading its chain back every round draft_streams: bool, + shared_guard: bool = false, head_trees: bool, /// a prompt in pieces leaves its last token as the pending row of its first shared round's window join_tail: bool = false, @@ -147,6 +149,7 @@ pub const Config = struct { .batch_streams = if (streams) (if (m.max_streams != 0) m.max_streams else 32) else 1, .node_probabilities = m.draft_probabilities or m.draft_confidence != null, .draft_streams = m.draft_streams, + .shared_guard = m.shared_guard, .head_trees = m.head_trees, .join_tail = m.join_tail, .piece_runs = m.piece_runs, diff --git a/zig/src/core/lanes/drafted.zig b/zig/src/core/lanes/drafted.zig index 999aaeab6a..e814988934 100644 --- a/zig/src/core/lanes/drafted.zig +++ b/zig/src/core/lanes/drafted.zig @@ -103,6 +103,7 @@ pub const Drafted = struct { m.draft_streams = d.batched; m.draft_probabilities = false; m.head_trees = d.trees; + m.shared_guard = d.shared_guard; if (d.tree_nodes > 0) { // the drafter grows its own trees, best first, and prices their nodes m.drafts = @min(d.tree_nodes, if (target.exact_width > 1) target.exact_width - 1 else 0); diff --git a/zig/src/core/lanes/drafted_guard_test.zig b/zig/src/core/lanes/drafted_guard_test.zig new file mode 100644 index 0000000000..e0e38868a4 --- /dev/null +++ b/zig/src/core/lanes/drafted_guard_test.zig @@ -0,0 +1,51 @@ +//! A drafter that asks for the shared-round guard (Facts.shared_guard): which shared rounds draft, replies unchanged. +const std = @import("std"); +const Config = @import("config.zig").Config; +const Model = @import("config.zig").Model; +const Engine = @import("engine.zig").Engine; +const sm = @import("stream.zig"); +const fake = @import("fake.zig"); +const Drafted = @import("drafted.zig").Drafted; +const dt = @import("drafted_test.zig"); + +const gpa = std.testing.allocator; + +const Run = struct { tokens: [3][64]u32 = undefined, lens: [3]usize = @splat(0), holds: u32 = 0 }; + +/// Three streams decoded together: drafts off, or the fake drafter with or without the guard. +fn run(drafts: bool, guard: bool) !Run { + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var head: dt.FakeDrafter = .{ .facts_ = .{ .depth = 4, .step_ms = 0.5, .batched = true, .shared_guard = guard } }; + defer head.deinit(); + var wrapped = try Drafted.init(gpa, target.backend(), head.drafter()); + defer wrapped.deinit(); + const m: Model = .{ .exact_width = 16, .gpu_tokens = true, .hidden_rows = true, .batch_rows = 32, .max_streams = 4 }; + var cfg = try Config.init(gpa, if (drafts) wrapped.facts(m) else m, 16, 15); + defer cfg.deinit(gpa); + var clock: fake.FixedClock = .{}; + var engine = Engine.init(gpa, &cfg, if (drafts) wrapped.backend() else target.backend(), clock.clock()); + defer engine.deinit(); + var streams: [3]sm.Stream = undefined; + for (&streams, 0..) |*s, i| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = dt.long[0 .. 11 + i], .max_new = 48, .eos = &.{96}, .drafts = drafts }); + defer for (&streams) |*s| s.deinit(gpa); + for (&streams) |*s| try engine.addStream(s); + while (engine.activeCount() > 0) try engine.step(); + var out: Run = .{ .holds = head.holds }; + for (&streams, 0..) |*s, i| { + out.lens[i] = s.emitted().len; + @memcpy(out.tokens[i][0..out.lens[i]], s.emitted()); + } + return out; +} + +test "the guard changes only which shared rounds draft: replies equal one-token decode with it on and off" { + const plain = try run(false, false); + const off = try run(true, false); + const on = try run(true, true); + for (0..3) |i| { + try std.testing.expectEqualSlices(u32, plain.tokens[i][0..plain.lens[i]], off.tokens[i][0..off.lens[i]]); + try std.testing.expectEqualSlices(u32, plain.tokens[i][0..plain.lens[i]], on.tokens[i][0..on.lens[i]]); + } + try std.testing.expect(on.holds < off.holds); // its plain stretch holds no drafts +} diff --git a/zig/src/core/lanes/drafter.zig b/zig/src/core/lanes/drafter.zig index 80577e0edf..88492dd650 100644 --- a/zig/src/core/lanes/drafter.zig +++ b/zig/src/core/lanes/drafter.zig @@ -18,6 +18,7 @@ pub const Facts = struct { window_costs: []const Cost = &.{}, // the target's window ms by rows (timed at load), for the depth rule lane_costs: []const Cost = &.{}, // the same for the lane planner (Model.lane_costs) tree_nodes: u32 = 0, // hold() without lanes grows its own tree of up to `depth` nodes, at most this many (tree, chances) + shared_guard: bool = false, // shared rounds draft only while they measure faster than plain ones (Model.shared_guard) }; /// Target rows [start, start + features.rows) and the token after each; `rows`, when given, picks which to absorb. diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index 1664587c6d..ba545f18c4 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -3,6 +3,7 @@ const std = @import("std"); const Allocator = std.mem.Allocator; const Config = @import("config.zig").Config; const depth = @import("depth.zig"); +const shared_guard = @import("shared_guard.zig"); const accept = @import("accept.zig"); const be = @import("backend.zig"); const sm = @import("stream.zig"); @@ -43,6 +44,7 @@ pub const Engine = struct { log: ?*ev.Log = null, trim: ?trim.Trim = null, // a family's row policy in place of the built-in allocation (lanes/trim.zig) rule: depth.Rule, + guard: shared_guard.Guard, live: std.ArrayList(*Stream) = .empty, // admitted streams in admission order (Python `_live`) drafted: u64 = 0, accepted: u64 = 0, @@ -75,6 +77,7 @@ pub const Engine = struct { .batch_rows = cfg.batch_rows, .batched_heads = cfg.draft_streams, }, + .guard = .{ .gpa = gpa }, }; } @@ -82,6 +85,7 @@ pub const Engine = struct { e.live.deinit(e.gpa); e.unread.deinit(e.gpa); e.rule.deinit(); + e.guard.deinit(); for (&e.lone_kept) |*mode| for (mode) |*k| if (k.*) |*x| x.deinit(e.gpa); e.arena.deinit(); } @@ -285,7 +289,9 @@ pub const Engine = struct { for (whos, heads.items) |*w, i| w.* = win.who(plans[i].stream); const budgets = try a.alloc(i64, heads.items.len); const sampled = win.anySampled(plans); - try e.rule.budgets(whos, win.probe(e), sampled, budgets); + // the guard's plain stretches: no stream drafts, so the round costs what it would without a head + const guarded = e.cfg.shared_guard and !try e.guard.drafts(plans.len); + if (guarded) @memset(budgets, 0) else try e.rule.budgets(whos, win.probe(e), sampled, budgets); const ids = try a.alloc(ev.Value, heads.items.len); for (ids, heads.items) |*id, i| id.* = str(plans[i].stream.id); try trail.event(e, &.{ f("ev", str("budgets")), f("streams", .{ .list = ids }), f("depths", .{ .i64s = budgets }) }); @@ -314,6 +320,7 @@ pub const Engine = struct { try trail.event(e, &.{ f("ev", str("overhead")), f("streams", int(plans.len)), f("rows", int(total)), f("ms", ev.bits(ms)) }); try e.rule.observeOverhead(@intCast(plans.len), total, ms, sampled, e.rule.headMs(budgets)); } + if (e.cfg.shared_guard) try e.guard.observe(plans.len, shared_guard.committed(outcomes), e.clock.elapsedMs(.overhead)); if (e.backend.round_observer != null) e.backend.observed(paths); } diff --git a/zig/src/core/lanes/lanes.zig b/zig/src/core/lanes/lanes.zig index ff09b58fb2..9f9ab516a0 100644 --- a/zig/src/core/lanes/lanes.zig +++ b/zig/src/core/lanes/lanes.zig @@ -59,6 +59,8 @@ test { _ = @import("windows_grammar_test.zig"); _ = @import("pricing_test.zig"); _ = @import("depth_probe_test.zig"); + _ = @import("shared_guard.zig"); + _ = @import("drafted_guard_test.zig"); _ = @import("proposal_external_test.zig"); _ = @import("lone_costs_test.zig"); _ = @import("lone_backoff_test.zig"); diff --git a/zig/src/core/lanes/shared_guard.zig b/zig/src/core/lanes/shared_guard.zig new file mode 100644 index 0000000000..8adaee2646 --- /dev/null +++ b/zig/src/core/lanes/shared_guard.zig @@ -0,0 +1,145 @@ +//! Whether a shared round drafts at all: drafted and plain rounds' measured tokens a ms, by stream count. +const std = @import("std"); +const Allocator = std.mem.Allocator; + +pub const Constants = struct { + explore: [2]u32 = .{ 16, 4 }, // measured rounds plain and drafted rounds run when their rate is (re)taken + exploit: u32 = 512, // measured rounds the faster mode runs before the slower one is measured again + longest: u5 = 6, // each time the same mode wins again its run doubles, up to exploit << longest + margin: f64 = 1.05, // drafted rounds must beat plain ones' tokens a ms by this much + blend: f64 = 0.5, // the newest stretch's weight in a mode's rate +}; + +/// One stream count's state: each mode's rate (0 plain, 1 drafted), the mode running and its stretch. +const Entry = struct { + streams: u32, + rate: [2]?f64 = .{ null, null }, + mode: u1 = 0, // plain first: drafted rounds that lose cost more than plain ones, so they are measured short + exploring: bool = true, + left: u32, + tokens: f64 = 0, + ms: f64 = 0, + skip: bool = true, // a mode's first round still carries the other mode's drafts or holds: not measured + best: ?u1 = null, // the mode the last measurement chose + streak: u5 = 0, // how many times in a row it chose that mode +}; + +pub const Guard = struct { + gpa: Allocator, + k: Constants = .{}, + list: std.ArrayList(Entry) = .empty, + + pub fn deinit(g: *Guard) void { + g.list.deinit(g.gpa); + } + + fn entry(g: *Guard, streams: usize) !*Entry { + const n: u32 = @intCast(streams); + for (g.list.items) |*e| if (e.streams == n) return e; + try g.list.append(g.gpa, .{ .streams = n, .left = g.k.explore[0] }); + return &g.list.items[g.list.items.len - 1]; + } + + /// Whether a shared round of `streams` streams may draft (false: every stream runs its row alone). + pub fn drafts(g: *Guard, streams: usize) !bool { + return (try g.entry(streams)).mode == 1; + } + + /// A finished shared round: its committed tokens and wall ms, in the mode drafts() gave it. + pub fn observe(g: *Guard, streams: usize, tokens: u64, ms: f64) !void { + const e = try g.entry(streams); + if (e.skip) { + e.skip = false; + return; + } + e.tokens += @floatFromInt(tokens); + e.ms += ms; + e.left -|= 1; + if (e.left > 0) return; + if (e.ms > 0) { + const now = e.tokens / e.ms; + e.rate[e.mode] = if (e.rate[e.mode]) |r| r + g.k.blend * (now - r) else now; + } + e.tokens = 0; + e.ms = 0; + const before = e.mode; + if (e.exploring and e.rate[e.mode ^ 1] == null) { + e.mode ^= 1; // the other mode has no rate yet: measure it next + } else if (e.exploring) { + e.mode = if (e.rate[1].? >= e.rate[0].? * g.k.margin) 1 else 0; + e.streak = if (e.best == e.mode) @min(e.streak + 1, g.k.longest) else 0; + e.best = e.mode; + e.exploring = false; + } else { + e.mode ^= 1; // the slower mode is measured again, in case the streams' text changed + e.exploring = true; + } + e.left = if (e.exploring) g.k.explore[e.mode] else g.k.exploit << e.streak; + e.skip = e.mode != before; + } +}; + +/// The tokens a round's outcomes committed (each outcome's `got`). +pub fn committed(outcomes: anytype) u64 { + var n: u64 = 0; + for (outcomes) |o| n += o.got.len; + return n; +} + +/// Rounds of a guard at `streams` with rates `plain` and `drafted` tokens a ms: how many it drafted. +fn run(g: *Guard, streams: usize, rounds: usize, plain: f64, drafted: f64) !usize { + var n: usize = 0; + for (0..rounds) |_| { + const d = try g.drafts(streams); + n += @intFromBool(d); + try g.observe(streams, if (d) 3 else 1, if (d) 3 / drafted else 1 / plain); + } + return n; +} + +test "a shared round drafts only while drafted rounds beat plain ones, measuring the slower mode again now and then" { + var g: Guard = .{ .gpa = std.testing.allocator }; + defer g.deinit(); + // drafted rounds slower at four streams: after one stretch of each, plain but for a re-measure a period + const slow = try run(&g, 4, 1000, 0.25, 0.20); + try std.testing.expect(slow <= 4 * (g.k.explore[1] + 1)); // the first measure and three re-measures, short ones + try std.testing.expect(slow > 0); + // the same verdict again and again: the slower mode is measured ever more rarely, under 1% of rounds + const later = try run(&g, 4, 20000, 0.25, 0.20); + try std.testing.expect(later * 500 < 20000); + // faster at one stream: drafted but for the plain re-measures + const fast = try run(&g, 1, 1000, 0.13, 0.15); + try std.testing.expect(fast > 800); + // each stream count keeps its own rates + try std.testing.expect(!try g.drafts(4) or (try g.entry(4)).exploring); + try std.testing.expect(try g.drafts(1) or (try g.entry(1)).exploring); +} + +test "a mode's first round is not measured: it carries the other mode's drafts" { + var g: Guard = .{ .gpa = std.testing.allocator, .k = .{ .explore = .{ 2, 2 }, .exploit = 4 } }; + defer g.deinit(); + try std.testing.expect(!try g.drafts(2)); // plain first + try g.observe(2, 0, 1); // skipped: the stream count's first round + try g.observe(2, 1, 4); + try g.observe(2, 1, 4); + try std.testing.expect(try g.drafts(2)); // plain measured at 0.25 a ms; drafted next + try g.observe(2, 100, 1); // skipped: plain's last held rows, a fantasy rate that would keep drafting + try g.observe(2, 1, 10); + try g.observe(2, 1, 10); + try std.testing.expect(!try g.drafts(2)); // plain at 0.25 beats drafted at 0.1 + try std.testing.expectEqual(@as(?f64, 0.1), (try g.entry(2)).rate[1]); +} + +test "a guard allocates once a stream count: its rounds after that take no allocation" { + var counted = std.testing.FailingAllocator.init(std.testing.allocator, .{}); + var g: Guard = .{ .gpa = counted.allocator(), .k = .{ .explore = .{ 2, 2 }, .exploit = 4 } }; + defer g.deinit(); + for ([_]usize{ 1, 4, 8, 16 }) |streams| _ = try run(&g, streams, 1, 0.25, 0.2); + const first = .{ counted.allocations, counted.resize_index }; + for (0..200) |i| for ([_]usize{ 1, 4, 8, 16 }) |streams| { + _ = try g.drafts(streams); + try g.observe(streams, i % 3 + 1, 4); + }; + try std.testing.expectEqual(first[0], counted.allocations); + try std.testing.expectEqual(first[1], counted.resize_index); +} From 1b57e583855b552aa003004e80db702ac1e6ae4e Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:04:57 +0100 Subject: [PATCH 05/46] lanes: a shared round the guard keeps plain proposes no copies, and a stream joining one holds no drafts With the shared-round guard on, a plain round still verified the lane core's context and Recall copies, which are drafts too, priced from a prior; and a stream joining other streams held its first drafts even while the guard ran plain rounds at that stream count, a hold pass for each join. The engine now marks a shared round plain when the guard runs that stream count plain, and copyProposal proposes nothing for it, so the guard's plain mode is the drafts-off round. Open.firstDepth gives no drafts to a stream that joins live streams while the guard runs their count plus one plain. Both are behind Model.shared_guard; without it plain_round stays false and firstDepth is unchanged. Tests: a guarded plain round proposes no copies; a stream joining plain-guarded rounds holds no first drafts and, without the flag, holds them as before. --- zig/src/core/lanes/drafted_guard_test.zig | 51 +++++++++++++++++++++++ zig/src/core/lanes/engine.zig | 3 ++ zig/src/core/lanes/engine_test.zig | 3 ++ zig/src/core/lanes/open.zig | 3 ++ zig/src/core/lanes/windows.zig | 2 +- 5 files changed, 61 insertions(+), 1 deletion(-) diff --git a/zig/src/core/lanes/drafted_guard_test.zig b/zig/src/core/lanes/drafted_guard_test.zig index e0e38868a4..ca6d6bc568 100644 --- a/zig/src/core/lanes/drafted_guard_test.zig +++ b/zig/src/core/lanes/drafted_guard_test.zig @@ -10,6 +10,57 @@ const dt = @import("drafted_test.zig"); const gpa = std.testing.allocator; +test "a stream joining shared rounds that the guard runs plain holds no first drafts; a lone one does" { + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var head: dt.FakeDrafter = .{ .facts_ = .{ .depth = 4, .step_ms = 0.5, .batched = true, .shared_guard = true } }; + defer head.deinit(); + var wrapped = try Drafted.init(gpa, target.backend(), head.drafter()); + defer wrapped.deinit(); + const m: Model = .{ .exact_width = 16, .gpu_tokens = true, .hidden_rows = true, .batch_rows = 32, .max_streams = 4 }; + var cfg = try Config.init(gpa, wrapped.facts(m), 16, 15); + defer cfg.deinit(gpa); + var clock: fake.FixedClock = .{}; + var engine = Engine.init(gpa, &cfg, wrapped.backend(), clock.clock()); + defer engine.deinit(); + var streams: [2]sm.Stream = undefined; + for (&streams, 0..) |*s, i| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = dt.long[0 .. 12 + i], .max_new = 24, .eos = &.{96} }); + defer for (&streams) |*s| s.deinit(gpa); + try engine.addStream(&streams[0]); + try std.testing.expectEqual(@as(u32, 1), head.holds); // alone: its first drafts are held + for (0..3) |_| try engine.step(); + const before = head.holds; + try std.testing.expect(!try engine.guard.drafts(2)); // the guard measures plain shared rounds first + try engine.addStream(&streams[1]); + try std.testing.expectEqual(before, head.holds); // joining: none held + while (engine.activeCount() > 0) try engine.step(); +} + +test "without the guard's flag a stream joining shared rounds holds its first drafts as before" { + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var head: dt.FakeDrafter = .{ .facts_ = .{ .depth = 4, .step_ms = 0.5, .batched = true } }; + defer head.deinit(); + var wrapped = try Drafted.init(gpa, target.backend(), head.drafter()); + defer wrapped.deinit(); + const m: Model = .{ .exact_width = 16, .gpu_tokens = true, .hidden_rows = true, .batch_rows = 32, .max_streams = 4 }; + var cfg = try Config.init(gpa, wrapped.facts(m), 16, 15); + defer cfg.deinit(gpa); + try std.testing.expect(!cfg.shared_guard); + var clock: fake.FixedClock = .{}; + var engine = Engine.init(gpa, &cfg, wrapped.backend(), clock.clock()); + defer engine.deinit(); + var streams: [2]sm.Stream = undefined; + for (&streams, 0..) |*s, i| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = dt.long[0 .. 12 + i], .max_new = 24, .eos = &.{96} }); + defer for (&streams) |*s| s.deinit(gpa); + try engine.addStream(&streams[0]); + for (0..3) |_| try engine.step(); + const before = head.holds; + try engine.addStream(&streams[1]); + try std.testing.expectEqual(before + 1, head.holds); // joining: its first drafts held + while (engine.activeCount() > 0) try engine.step(); +} + const Run = struct { tokens: [3][64]u32 = undefined, lens: [3]usize = @splat(0), holds: u32 = 0 }; /// Three streams decoded together: drafts off, or the fake drafter with or without the guard. diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index ba545f18c4..5349b38904 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -50,6 +50,7 @@ pub const Engine = struct { accepted: u64 = 0, shared_rounds: u64 = 0, alone: bool = true, // this step has one live stream: its copies may take their ramp + plain_round: bool = false, // a shared round the guard keeps plain: no stream copies or drafts, as with drafts off lone_kept: [2][lone_costs.bands]?lone_costs.Kept = @splat(@splat(null)), // the newest lone calibrations, greedy then sampled, a band each steps: u64 = 0, arena: std.heap.ArenaAllocator, // a round's temporaries @@ -260,6 +261,8 @@ pub const Engine = struct { e.clock.start(); const a = e.arena.allocator(); const plans = try a.alloc(Plan, entries.len); + e.plain_round = e.cfg.shared_guard and !try e.guard.drafts(entries.len); + defer e.plain_round = false; for (entries, plans) |s, *p| p.* = try win.plan(e, s, null); try win.allocate(e, plans); for (plans) |*p| try win.constrain(e, p); diff --git a/zig/src/core/lanes/engine_test.zig b/zig/src/core/lanes/engine_test.zig index 36167e30d6..6995de7dbe 100644 --- a/zig/src/core/lanes/engine_test.zig +++ b/zig/src/core/lanes/engine_test.zig @@ -463,6 +463,9 @@ test "a family that keeps copies to shared rounds proposes none for a lone strea try std.testing.expect((try win.copyProposal(&engine, &s)).tokens.len >= 2); engine.alone = true; try std.testing.expectEqual(lone_copies, (try win.copyProposal(&engine, &s)).tokens.len >= 2); + engine.alone = false; + engine.plain_round = true; // a shared round the guard keeps plain: no copies either, as with drafts off + try std.testing.expectEqual(@as(usize, 0), (try win.copyProposal(&engine, &s)).tokens.len); } } diff --git a/zig/src/core/lanes/open.zig b/zig/src/core/lanes/open.zig index 4c4cb585e5..733d3bb2f8 100644 --- a/zig/src/core/lanes/open.zig +++ b/zig/src/core/lanes/open.zig @@ -135,6 +135,9 @@ fn openedCommitted(e: *Engine, s: *Stream, reach: Reach, hook: ?be.Committed) !v /// The first drafts' depth for `reach`. pub fn firstDepth(e: *Engine, s: *Stream, reach: Reach) !u32 { var who = win.who(s); + // a stream joining shared rounds the guard runs plain holds no drafts, as the guard measured it faster + const others = e.activeCount(); + if (others > 0 and e.cfg.family_streams and e.cfg.shared_guard and !try e.guard.drafts(others + 1)) return 0; if (reach == .head and e.cfg.node_probabilities) return @intCast(try e.rule.headDepth(&who, null, win.probe(e))); return @intCast(try e.rule.depth(who)); } diff --git a/zig/src/core/lanes/windows.zig b/zig/src/core/lanes/windows.zig index e8e48be2ad..b70c2f634d 100644 --- a/zig/src/core/lanes/windows.zig +++ b/zig/src/core/lanes/windows.zig @@ -197,7 +197,7 @@ fn headTrunk(e: *Engine, s: *Stream, p: Plan, levels: usize) !?fill.Trunk { /// A copy from the context backed by `enter_match` ids (Python `_copy_proposal`), else Recall's when it should pay. pub fn copyProposal(e: *Engine, s: *Stream) !Copy { - if (e.cfg.max_copy == 0 or s.force.items.len > 0 or (e.cfg.root_copies and s.root_copy_miss)) return .{}; + if (e.cfg.max_copy == 0 or e.plain_round or s.force.items.len > 0 or (e.cfg.root_copies and s.root_copy_miss)) return .{}; const width: u32 = if (e.alone) (s.copy_width orelse e.cfg.first_copy) else e.cfg.first_copy; const most = @min(@as(i64, width), s.draftRoom() - 1); if (s.proposer) |p| if (!e.alone or e.cfg.lone_copies) { From 47eaf78bd7278d5dcbd58d63d6d59559f19523dc Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:06:32 +0100 Subject: [PATCH 06/46] lanes: an external drafter's chain chances price verify rows, and its holds follow the grants A drafter that gives each held chain draft's chance (Drafter.VTable.chances) may set Facts.chances. Drafted.facts then sets Model.draft_probabilities and Model.granted_depth, so the window allocator prices each stream's verify rows by those chances, as it does a head's node chances. With Facts.confidence_floor it instead verifies the drafts at least that likely (a fixed window when 0). With granted_depth a stream's next hold is as deep as its last grant plus a slack: one more draft when the grant took every held draft, none when it stopped short of them (Stream.grant_short), and after a grant of none a probe of draft_slack drafts at the depth rule's probe interval. Facts.slack sets one slack for all three. A shared round's granted depths are cut so the drafting streams' blocks fit one hold pass of batch_rows. When the shared-round guard starts a drafted stretch, its first round grants every stream one draft, so the stretch measures drafted rounds rather than plain ones left by grants of none. All of it is behind Facts.chances and Model.granted_depth, off by default; windows.allocate records grant_short only for granted depths. Tests: chances price the chain's rows and drafted replies equal one-token decode alone and at 1, 4 and 8 shared streams, with the scheduled verify ragged across streams and verifying fewer rows than a fixed window; the facts reach the lane config; granted depths, their slacks and the one-pass cut; holds as wide as the grants. --- zig/src/core/lanes/config.zig | 11 +- zig/src/core/lanes/depth.zig | 62 ++++- zig/src/core/lanes/drafted.zig | 7 + zig/src/core/lanes/drafted_chances_test.zig | 266 ++++++++++++++++++++ zig/src/core/lanes/drafter.zig | 3 + zig/src/core/lanes/engine.zig | 3 +- zig/src/core/lanes/lanes.zig | 1 + zig/src/core/lanes/shared_guard.zig | 32 ++- zig/src/core/lanes/stream.zig | 2 + zig/src/core/lanes/windows.zig | 7 +- 10 files changed, 387 insertions(+), 7 deletions(-) create mode 100644 zig/src/core/lanes/drafted_chances_test.zig diff --git a/zig/src/core/lanes/config.zig b/zig/src/core/lanes/config.zig index a7eac1ea7d..340788ccfe 100644 --- a/zig/src/core/lanes/config.zig +++ b/zig/src/core/lanes/config.zig @@ -35,6 +35,8 @@ pub const Model = struct { max_streams: u32 = 0, // a shared forward's streams at most (0: 32) shared_costs: []const Cost = &.{}, // shared forwards' ms by total rows (timed at load) draft_probabilities: bool = false, // the head gives each drafted node's chance of landing + granted_depth: bool = false, // with draft_probabilities: every stream drafts its last grant plus a slack + draft_slack: u32 = 0, // drafts past a stream's last grant, short or not (0: the Constants' own) draft_streams: bool = false, // the head drafts every stream of a shared round in one batch shared_guard: bool = false, // a shared round drafts only while drafted rounds measure faster (shared_guard.zig) head_trees: bool = false, // the head drafts trees of lanes (DraftRequest.lanes) and holds them @@ -65,6 +67,7 @@ pub const Config = struct { batch_rows: u32, batch_streams: u32, node_probabilities: bool, + granted_depth: bool = false, // Model.granted_depth fill_lanes: u32 = 0, // rows a window may grow to with suffix-match branches (0: off; lanes/fill.zig) fill_match: u32 = 8, // context tokens a suffix-match branch must match (shorter ones are coincidental phrases) fill_floor: f64 = 0.35, // the share of grafted rows a stream must keep to go on reading its chain back every round @@ -148,6 +151,7 @@ pub const Config = struct { .batch_rows = batch_rows, .batch_streams = if (streams) (if (m.max_streams != 0) m.max_streams else 32) else 1, .node_probabilities = m.draft_probabilities or m.draft_confidence != null, + .granted_depth = m.granted_depth and m.draft_probabilities, .draft_streams = m.draft_streams, .shared_guard = m.shared_guard, .head_trees = m.head_trees, @@ -160,7 +164,12 @@ pub const Config = struct { .shared_costs = try alloc.extendCosts(gpa, timed.items, batch_rows), .lane_costs = try alloc.extendCosts(gpa, m.lane_costs, 64), .mtp_step_ms = m.mtp_step_ms, - .k = .{ .depth_rate = if (m.depth_rate > 0) m.depth_rate else (Constants{}).depth_rate }, + .k = .{ + .depth_rate = if (m.depth_rate > 0) m.depth_rate else (Constants{}).depth_rate, + .draft_slack = if (m.draft_slack > 0) m.draft_slack else (Constants{}).draft_slack, + .grant_slack = if (m.draft_slack > 0) m.draft_slack else (Constants{}).grant_slack, + .short_slack = if (m.draft_slack > 0) m.draft_slack else (Constants{}).short_slack, + }, .draft_step_ms = m.draft_step_ms, }; } diff --git a/zig/src/core/lanes/depth.zig b/zig/src/core/lanes/depth.zig index 9f6c711ac5..3c45e6ee6d 100644 --- a/zig/src/core/lanes/depth.zig +++ b/zig/src/core/lanes/depth.zig @@ -11,6 +11,8 @@ pub const Constants = struct { plain_margin: f64 = 1.05, // a drafted depth must beat plain rounds' tokens a ms by this much plain_wait_most: u64 = 128, // plain rounds between probes at most draft_slack: i64 = 2, // drafts a stream offers past what its last shared round granted it + grant_slack: i64 = 1, // granted depths: drafts past a grant that verified every draft the stream held + short_slack: i64 = 0, // granted depths: drafts past a grant that left some of them unverified }; /// A stream's acceptance by depth and probe counters (Python `_depth_state[stream]`). @@ -34,6 +36,7 @@ pub const Who = struct { finished: bool, forced: bool, granted: ?u32 = null, + grant_short: bool = false, // the last grant left some held drafts unverified (the width did not bind) stream: ?*anyopaque = null, sampled: bool = false, // keyed draws: host work a row that a greedy round does not have }; @@ -58,6 +61,7 @@ pub const Rule = struct { draft_step_ms: f64 = 0, // a shared round's draft past a stream's first (budgets price it) plain_guard: bool, node_probabilities: bool, + granted_depth: bool = false, // node_probabilities: a stream drafts its last grant plus draft_slack, lone ones too confidence_floor: ?f64 = null, batch_rows: u32, batched_heads: bool = false, // one batch drafts every stream of a shared round (its steps: the deepest stream's) @@ -191,9 +195,28 @@ pub const Rule = struct { return true; } + /// Drafts past a grant: none when it stopped short of the held drafts (the width did not bind), else grant_slack. + fn slack(r: *const Rule, who: Who) i64 { + return if (who.grant_short) r.k.short_slack else r.k.grant_slack; + } + + /// The last grant plus its slack (the deepest chain before any); after a grant of none, a probe of draft_slack. + fn grantDepth(r: *Rule, who: Who) !i64 { + const most = @min(@as(i64, r.most_drafts), @max(1, who.draft_room - 1)); + const granted = who.granted orelse return most; + _ = try r.rates(who); + const st = &who.state.*.?; + if (granted == 0) return if (!r.plain_guard or r.probePlain(st)) @min(most, r.k.draft_slack) else 0; + st.plain = 0; + st.wait = r.k.depth_probe_every; + return @min(most, @as(i64, granted) + r.slack(who)); + } + /// Python `_head_depth`: the next round's head drafts for a stream (0 when forced, 1 with a copy ahead). pub fn headDepth(r: *Rule, who: *const Who, budget: ?i64, probe: Probe) !i64 { - const d: i64 = if (who.finished) 0 else if (budget) |b| b else if (r.node_probabilities) + const d: i64 = if (who.finished) 0 else if (budget) |b| b else if (r.granted_depth) + try r.grantDepth(who.*) + else if (r.node_probabilities) @min(@as(i64, r.most_drafts), @max(1, who.draft_room - 1)) else try r.depth(who.*); @@ -202,6 +225,15 @@ pub const Rule = struct { return d; } + /// A shared hold pads every block to its deepest: depths cut so the drafting streams' blocks fit `rows` (one pass). + fn capPass(out: []i64, rows: u32) void { + var n: i64 = 0; + for (out) |d| n += @intFromBool(d > 0); + if (n == 0) return; + const most = @max(1, @divFloor(@as(i64, rows), n)); + for (out) |*d| d.* = @min(d.*, most); + } + /// Head drafts for the streams of a shared round, sized by allocation over their chain chances. pub fn budgets(r: *Rule, whos: []Who, probe: Probe, sampled: bool, out: []i64) !void { if (r.confidence_floor != null) { @@ -212,8 +244,9 @@ pub const Rule = struct { for (whos, out) |*w, *o| { const inner = try r.headDepth(w, null, probe); const granted: i64 = if (w.granted) |g| g else r.most_drafts; - o.* = try r.headDepth(w, @min(inner, granted + r.k.draft_slack), probe); + o.* = try r.headDepth(w, if (r.granted_depth) inner else @min(inner, granted + r.k.draft_slack), probe); } + if (r.granted_depth) capPass(out, r.batch_rows); return; } const most = try r.gpa.alloc(i64, whos.len); @@ -380,3 +413,28 @@ test "a shared round's drafts come out the same whichever depths its overhead wa } try std.testing.expectEqualSlices(i64, &got[0], &got[1]); } + +test "granted depths in a shared round fit one hold pass: every block is as deep as the deepest" { + const empty: Table = .{}; + var rule: Rule = .{ .gpa = std.testing.allocator, .prior = &.{0.8}, .most_drafts = 9, .family_costs = &empty, .shared_costs = &empty, .mtp_step_ms = 0.5, .plain_guard = false, .node_probabilities = true, .granted_depth = true, .batch_rows = 32 }; + defer rule.deinit(); + var states: [8]?State = @splat(null); + defer for (&states) |*st| if (st.*) |*x| x.deinit(std.testing.allocator); + var whos: [8]Who = undefined; + for (&whos, &states, 0..) |*w, *st, i| w.* = .{ .state = st, .draft_room = 100, .finished = false, .forced = false, .granted = @intCast(i) }; + const No = struct { + fn copy(_: *anyopaque, _: *const Who) bool { + return false; + } + }; + var unused: u8 = 0; + var out: [8]i64 = undefined; + try rule.budgets(&whos, .{ .ptr = &unused, .copy = No.copy }, false, &out); + // a grant of none probes two; grants 1..7 that took every held draft, plus one; cut to 32 rows over eight blocks + try std.testing.expectEqualSlices(i64, &.{ 2, 2, 3, 4, 4, 4, 4, 4 }, &out); + try rule.budgets(whos[0..2], .{ .ptr = &unused, .copy = No.copy }, false, out[0..2]); + try std.testing.expectEqualSlices(i64, &.{ 2, 2 }, out[0..2]); // two blocks fit whole + for (whos[0..2]) |*w| w.grant_short = true; + try rule.budgets(whos[0..2], .{ .ptr = &unused, .copy = No.copy }, false, out[0..2]); + try std.testing.expectEqualSlices(i64, &.{ 2, 1 }, out[0..2]); // a grant short of the held drafts: no slack +} diff --git a/zig/src/core/lanes/drafted.zig b/zig/src/core/lanes/drafted.zig index e814988934..50f5739913 100644 --- a/zig/src/core/lanes/drafted.zig +++ b/zig/src/core/lanes/drafted.zig @@ -104,6 +104,13 @@ pub const Drafted = struct { m.draft_probabilities = false; m.head_trees = d.trees; m.shared_guard = d.shared_guard; + if (d.chances and x.drafter.vtable.chances != null) { + // a chain's drafts priced by the drafter's own chances; the next hold as deep as the grant and its slack + m.draft_probabilities = d.confidence_floor == null; + m.draft_confidence = d.confidence_floor; + m.granted_depth = d.confidence_floor == null; + m.draft_slack = d.slack; + } if (d.tree_nodes > 0) { // the drafter grows its own trees, best first, and prices their nodes m.drafts = @min(d.tree_nodes, if (target.exact_width > 1) target.exact_width - 1 else 0); diff --git a/zig/src/core/lanes/drafted_chances_test.zig b/zig/src/core/lanes/drafted_chances_test.zig new file mode 100644 index 0000000000..02c5eb4cb1 --- /dev/null +++ b/zig/src/core/lanes/drafted_chances_test.zig @@ -0,0 +1,266 @@ +//! A drafter that prices its chains (Facts.chances): verify rows go by its chances and replies stay one-token decode. +const std = @import("std"); +const Config = @import("config.zig").Config; +const Model = @import("config.zig").Model; +const Cost = @import("config.zig").Cost; +const Engine = @import("engine.zig").Engine; +const be = @import("backend.zig"); +const sm = @import("stream.zig"); +const fake = @import("fake.zig"); +const dr = @import("drafter.zig"); +const depth = @import("depth.zig"); +const Table = @import("table.zig").Table; +const Drafted = @import("drafted.zig").Drafted; +const dt = @import("drafted_test.zig"); + +const gpa = std.testing.allocator; +const most_streams = 8; + +/// drafted_test's drafter with a chance for each held draft: `sure` where its guess is right, `unsure` where wrong. +const Priced = struct { + head: dt.FakeDrafter, + sure: f64 = 0.7, + unsure: f64 = 0.7, + vtable: dr.Drafter.VTable = undefined, + inner_hold: *const fn (*anyopaque, []const dr.Hold) anyerror!void = undefined, + at: std.AutoHashMapUnmanaged(*sm.Stream, u64) = .empty, // each stream's last hold's position + asked: u32 = 0, + held: u64 = 0, // drafts its holds asked for + + fn of(ptr: *anyopaque) *Priced { + return @fieldParentPtr("head", @as(*dt.FakeDrafter, @ptrCast(@alignCast(ptr)))); + } + + fn hold(ptr: *anyopaque, items: []const dr.Hold) anyerror!void { + const p = of(ptr); + for (items) |it| { + try p.at.put(gpa, it.stream, it.position); + p.held += it.depth; + } + return p.inner_hold(ptr, items); + } + + /// The running product along the chain, as the DSpark confidence head's survivals are. + fn chances(ptr: *anyopaque, s: *sm.Stream, out: []f64) anyerror!bool { + const p = of(ptr); + const lane = p.head.lanes.getPtr(s) orelse return false; + if (out.len > lane.held.items.len) return error.TooFewDrafts; + const position = p.at.get(s) orelse return error.NoHold; + p.asked += 1; + var reach: f64 = 1; + for (out, 0..) |*o, j| { + const at = position + j; + reach *= if ((at *% 2654435761 + j) % 5 == 0) p.unsure else p.sure; // drafted_test's wrong guesses + o.* = reach; + } + return true; + } + + fn drafter(p: *Priced) dr.Drafter { + const base_ = p.head.drafter(); + p.vtable = base_.vtable.*; + p.inner_hold = base_.vtable.hold; + p.vtable.hold = hold; + p.vtable.chances = chances; + return .{ .ptr = &p.head, .vtable = &p.vtable }; + } + + fn deinit(p: *Priced) void { + p.at.deinit(gpa); + p.head.deinit(); + } +}; + +/// Every shared verify's window rows: whether one round gave its streams unequal rows (a ragged verify). +const Seen = struct { + var inner: *const fn (*anyopaque, []const be.Window, []be.Verified) anyerror!void = undefined; + var ragged: u32 = 0; + var rows: u64 = 0; + + fn verify(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { + for (windows) |w| { + rows += w.rows(); + if (w.rows() != windows[0].rows()) { + ragged += 1; + break; + } + } + return inner(ptr, windows, out); + } +}; + +var costs: [16]Cost = undefined; + +fn base() Model { + for (&costs, 1..) |*c, w| c.* = .{ .width = @intCast(w), .ms = 5.0 + 0.8 * @as(f64, @floatFromInt(w)) }; + return .{ .exact_width = 16, .gpu_tokens = true, .window_costs = &costs, .hidden_rows = true, .batch_rows = 32, .max_streams = most_streams }; +} + +const Run = struct { tokens: [most_streams][64]u32 = undefined, lens: [most_streams]usize = @splat(0), drafted: u64 = 0, holds: u32 = 0, asked: u32 = 0, ragged: u32 = 0, rows: u64 = 0, held: u64 = 0 }; + +/// The drafter's Facts.slack for the next runs (0: the depth rule's own; 9: every hold the whole block). +var slack: u32 = 0; + +/// The drafter's chances: equal everywhere, or knowing its wrong guesses (a confidence head that ranks well). +const Head = enum { flat, knows }; + +/// Each prompt decoded together (one stream: alone); `floor` null: the chances price rows, else a fixed window. +fn run(prompts: []const []const u32, drafts: bool, floor: ?f64, head: Head) !Run { + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var p: Priced = .{ .head = .{ .facts_ = .{ .depth = 9, .step_ms = 0.5, .batched = true, .chances = true, .confidence_floor = floor, .slack = slack } } }; + if (head == .knows) { + p.sure = 0.95; + p.unsure = 0.05; + } + defer p.deinit(); + var strict: dt.Strict = .{ .inner = target.backend() }; + defer strict.deinit(); + var seen = strict.backend(); + var table = seen.vtable.*; + Seen.inner = table.verify; + Seen.ragged = 0; + Seen.rows = 0; + table.verify = Seen.verify; + seen.vtable = &table; + var wrapped = try Drafted.init(gpa, seen, p.drafter()); + defer wrapped.deinit(); + var cfg = try Config.init(gpa, if (drafts) wrapped.facts(base()) else base(), 16, 15); + defer cfg.deinit(gpa); + var clock: fake.FixedClock = .{}; + var engine = Engine.init(gpa, &cfg, if (drafts) wrapped.backend() else target.backend(), clock.clock()); + defer engine.deinit(); + var streams: [most_streams]sm.Stream = undefined; + const n = prompts.len; + for (prompts, streams[0..n]) |prompt, *s| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = prompt, .max_new = 48, .eos = &.{96}, .drafts = drafts }); + defer for (streams[0..n]) |*s| s.deinit(gpa); + for (streams[0..n]) |*s| try engine.addStream(s); + while (engine.activeCount() > 0) try engine.step(); + var out: Run = .{ .holds = p.head.holds, .asked = p.asked, .ragged = Seen.ragged, .rows = Seen.rows, .held = p.held }; + for (streams[0..n], 0..) |*s, i| { + out.lens[i] = s.emitted().len; + @memcpy(out.tokens[i][0..out.lens[i]], s.emitted()); + out.drafted += s.drafted; + } + return out; +} + +fn expectSame(a: Run, b: Run, n: usize) !void { + for (0..n) |i| try std.testing.expectEqualSlices(u32, a.tokens[i][0..a.lens[i]], b.tokens[i][0..b.lens[i]]); +} + +/// `n` streams of unequal prompts, each a prefix of drafted_test's long prompt. +fn unequal(buf: *[most_streams][]const u32, n: usize) []const []const u32 { + for (buf[0..n], 0..) |*p, i| p.* = dt.long[0 .. 11 + i]; + return buf[0..n]; +} + +test "a drafter's chances price its chain's rows, and drafted replies equal one-token decode alone and shared" { + const sets = [_][]const []const u32{ &.{&dt.p1}, &.{ &dt.p1, &dt.long, &dt.p1 } }; + for (sets) |ps| { + const plain = try run(ps, false, null, .flat); + const priced = try run(ps, true, null, .flat); + try expectSame(plain, priced, ps.len); + try std.testing.expect(priced.asked > 0 and priced.drafted > 0); // the window allocator read the chances + const fixed = try run(ps, true, 0, .flat); + try expectSame(plain, fixed, ps.len); + try std.testing.expect(fixed.drafted > priced.drafted); // a fixed window verifies every draft the rows fit + } +} + +test "at 1, 4 and 8 streams the scheduled and the fixed verify equal one-token decode; shared rounds are ragged" { + var buf: [most_streams][]const u32 = undefined; + for ([_]usize{ 1, 4, 8 }) |n| { + const ps = unequal(&buf, n); + const plain = try run(ps, false, null, .flat); + const scheduled = try run(ps, true, null, .knows); + try expectSame(plain, scheduled, n); + const fixed = try run(ps, true, 0, .knows); + try expectSame(plain, fixed, n); + try std.testing.expect(scheduled.asked > 0 and scheduled.drafted > 0); + // a confident head's survivals cut each stream's chain where its guess goes wrong: fewer rows verified + try std.testing.expect(scheduled.rows < fixed.rows); + if (n > 1) try std.testing.expect(scheduled.ragged > 0); // streams' verify lengths differ within a round + } +} + +test "the drafter's chances reach the lane facts: rows by chances and holds by grants, or a fixed window" { + var p: Priced = .{ .head = .{ .facts_ = .{ .depth = 9, .step_ms = 0.5, .chances = true } } }; + defer p.deinit(); + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var wrapped: Drafted = .{ .gpa = gpa, .target = target.backend(), .drafter = p.drafter() }; + defer wrapped.deinit(); + const m = wrapped.facts(base()); + try std.testing.expect(m.draft_probabilities and m.granted_depth and m.draft_confidence == null); + var cfg = try Config.init(gpa, m, 16, 15); + defer cfg.deinit(gpa); + try std.testing.expectEqual(@as(i64, 2), cfg.k.draft_slack); // the depth rule's own + try std.testing.expect(cfg.k.grant_slack == 1 and cfg.k.short_slack == 0); // a granted width's own + p.head.facts_.slack = 1; + var one = try Config.init(gpa, wrapped.facts(base()), 16, 15); + defer one.deinit(gpa); + try std.testing.expectEqual(@as(i64, 1), one.k.draft_slack); // a drafter's own slack past its grant + try std.testing.expect(one.k.grant_slack == 1 and one.k.short_slack == 1); // short or not + p.head.facts_.confidence_floor = 0; + const f = wrapped.facts(base()); + try std.testing.expect(!f.draft_probabilities and !f.granted_depth and f.draft_confidence.? == 0); + p.head.facts_.chances = false; + const none = wrapped.facts(base()); + try std.testing.expect(!none.draft_probabilities and none.draft_confidence == null); +} + +test "a granted depth: the deepest chain first, then the grant plus its slack, and after a grant of none a probe" { + const empty: Table = .{}; + var rule: depth.Rule = .{ .gpa = gpa, .prior = &.{0.8}, .most_drafts = 9, .family_costs = &empty, .shared_costs = &empty, .mtp_step_ms = 0.5, .plain_guard = true, .node_probabilities = true, .granted_depth = true, .batch_rows = 32 }; + defer rule.deinit(); + var st: ?depth.State = null; + defer if (st) |*s| s.deinit(gpa); + const No = struct { + fn copy(_: *anyopaque, _: *const depth.Who) bool { + return false; + } + }; + var unused: u8 = 0; + const probe: depth.Probe = .{ .ptr = &unused, .copy = No.copy }; + var who: depth.Who = .{ .state = &st, .draft_room = 100, .finished = false, .forced = false }; + try std.testing.expectEqual(@as(i64, 9), try rule.headDepth(&who, null, probe)); + who.granted = 3; + try std.testing.expectEqual(@as(i64, 4), try rule.headDepth(&who, null, probe)); // every held draft granted: one more + who.grant_short = true; + try std.testing.expectEqual(@as(i64, 3), try rule.headDepth(&who, null, probe)); // drafts left unverified: the grant + who.grant_short = false; + who.granted = 8; + try std.testing.expectEqual(@as(i64, 9), try rule.headDepth(&who, null, probe)); + rule.k.grant_slack = 2; // a drafter's slack sets both + rule.k.short_slack = 2; + who.grant_short = true; + who.granted = 3; + try std.testing.expectEqual(@as(i64, 5), try rule.headDepth(&who, null, probe)); + who.granted = 0; + var drafted: usize = 0; + for (0..8) |_| drafted += @intFromBool(try rule.headDepth(&who, null, probe) > 0); + try std.testing.expectEqual(@as(usize, 1), drafted); // one probe of draft_slack drafts in the first eight rounds +} + +test "holds as wide as the grants: replies equal one-token decode at 1, 4 and 8 streams, fewer drafts than whole blocks" { + var buf: [most_streams][]const u32 = undefined; + defer slack = 0; + for ([_]usize{ 1, 4, 8 }) |n| { + const ps = unequal(&buf, n); + const plain = try run(ps, false, null, .flat); + for ([_]Head{ .flat, .knows }) |head| { + slack = 0; + const narrow = try run(ps, true, null, head); + slack = 2; + const two = try run(ps, true, null, head); + slack = 9; + const whole = try run(ps, true, null, head); + try expectSame(plain, narrow, n); + try expectSame(plain, two, n); + try expectSame(plain, whole, n); + try std.testing.expect(narrow.drafted > 0); + try std.testing.expect(narrow.held < two.held and two.held < whole.held); + } + } +} diff --git a/zig/src/core/lanes/drafter.zig b/zig/src/core/lanes/drafter.zig index 88492dd650..35a79e2e1c 100644 --- a/zig/src/core/lanes/drafter.zig +++ b/zig/src/core/lanes/drafter.zig @@ -19,6 +19,9 @@ pub const Facts = struct { lane_costs: []const Cost = &.{}, // the same for the lane planner (Model.lane_costs) tree_nodes: u32 = 0, // hold() without lanes grows its own tree of up to `depth` nodes, at most this many (tree, chances) shared_guard: bool = false, // shared rounds draft only while they measure faster than plain ones (Model.shared_guard) + chances: bool = false, // chances() gives each held chain draft's chance: verify rows go by them (Model.granted_depth) + confidence_floor: ?f64 = null, // with chances: every round drafts the deepest chain and verifies drafts this likely + slack: u32 = 0, // with chances: drafts a hold offers past the stream's last grant (0: the depth rule's own) }; /// Target rows [start, start + features.rows) and the token after each; `rows`, when given, picks which to absorb. diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index 5349b38904..9ff27b18d5 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -74,6 +74,7 @@ pub const Engine = struct { .draft_step_ms = cfg.draft_step_ms, .plain_guard = cfg.plain_guard, .node_probabilities = cfg.node_probabilities, + .granted_depth = cfg.granted_depth, .confidence_floor = cfg.draft_confidence, .batch_rows = cfg.batch_rows, .batched_heads = cfg.draft_streams, @@ -293,7 +294,7 @@ pub const Engine = struct { const budgets = try a.alloc(i64, heads.items.len); const sampled = win.anySampled(plans); // the guard's plain stretches: no stream drafts, so the round costs what it would without a head - const guarded = e.cfg.shared_guard and !try e.guard.drafts(plans.len); + const guarded = e.cfg.shared_guard and !try e.guard.draftsFresh(plans.len, whos); if (guarded) @memset(budgets, 0) else try e.rule.budgets(whos, win.probe(e), sampled, budgets); const ids = try a.alloc(ev.Value, heads.items.len); for (ids, heads.items) |*id, i| id.* = str(plans[i].stream.id); diff --git a/zig/src/core/lanes/lanes.zig b/zig/src/core/lanes/lanes.zig index 9f9ab516a0..8f8adf1e1b 100644 --- a/zig/src/core/lanes/lanes.zig +++ b/zig/src/core/lanes/lanes.zig @@ -54,6 +54,7 @@ test { _ = @import("recorder_test.zig"); _ = @import("gpu_full_test.zig"); _ = @import("drafted_test.zig"); + _ = @import("drafted_chances_test.zig"); _ = @import("root_copy_test.zig"); _ = @import("drafted_grammar_test.zig"); _ = @import("windows_grammar_test.zig"); diff --git a/zig/src/core/lanes/shared_guard.zig b/zig/src/core/lanes/shared_guard.zig index 8adaee2646..cdaf870736 100644 --- a/zig/src/core/lanes/shared_guard.zig +++ b/zig/src/core/lanes/shared_guard.zig @@ -1,6 +1,7 @@ //! Whether a shared round drafts at all: drafted and plain rounds' measured tokens a ms, by stream count. const std = @import("std"); const Allocator = std.mem.Allocator; +const depth = @import("depth.zig"); pub const Constants = struct { explore: [2]u32 = .{ 16, 4 }, // measured rounds plain and drafted rounds run when their rate is (re)taken @@ -45,6 +46,16 @@ pub const Guard = struct { return (try g.entry(streams)).mode == 1; } + /// drafts(); a drafted stretch's first round grants every stream one draft, so it measures drafted rounds. + pub fn draftsFresh(g: *Guard, streams: usize, whos: []depth.Who) !bool { + const e = try g.entry(streams); + if (e.mode == 1 and e.skip) for (whos) |*w| { + w.granted = 1; + w.grant_short = false; + }; + return e.mode == 1; + } + /// A finished shared round: its committed tokens and wall ms, in the mode drafts() gave it. pub fn observe(g: *Guard, streams: usize, tokens: u64, ms: f64) !void { const e = try g.entry(streams); @@ -130,14 +141,33 @@ test "a mode's first round is not measured: it carries the other mode's drafts" try std.testing.expectEqual(@as(?f64, 0.1), (try g.entry(2)).rate[1]); } +test "a drafted stretch starts from a grant of one: a plain stretch's grants of none would hold nothing to measure" { + var g: Guard = .{ .gpa = std.testing.allocator, .k = .{ .explore = .{ 2, 2 }, .exploit = 4 } }; + defer g.deinit(); + var st: ?depth.State = null; + var whos = [_]depth.Who{.{ .state = &st, .draft_room = 100, .finished = false, .forced = false, .granted = 0 }}; + try std.testing.expect(!try g.draftsFresh(2, &whos)); // plain first: grants as they are + try std.testing.expectEqual(@as(?u32, 0), whos[0].granted); + for (0..3) |_| try g.observe(2, 1, 4); + try std.testing.expect(try g.draftsFresh(2, &whos)); // the drafted stretch's first round + try std.testing.expectEqual(@as(?u32, 1), whos[0].granted); + try g.observe(2, 1, 4); // its first round, skipped + whos[0].granted = 0; + try std.testing.expect(try g.draftsFresh(2, &whos)); // later rounds keep the allocation's grants + try std.testing.expectEqual(@as(?u32, 0), whos[0].granted); +} + test "a guard allocates once a stream count: its rounds after that take no allocation" { var counted = std.testing.FailingAllocator.init(std.testing.allocator, .{}); var g: Guard = .{ .gpa = counted.allocator(), .k = .{ .explore = .{ 2, 2 }, .exploit = 4 } }; defer g.deinit(); for ([_]usize{ 1, 4, 8, 16 }) |streams| _ = try run(&g, streams, 1, 0.25, 0.2); const first = .{ counted.allocations, counted.resize_index }; + var st: [16]?depth.State = @splat(null); + var whos: [16]depth.Who = undefined; + for (&whos, &st) |*w, *s| w.* = .{ .state = s, .draft_room = 100, .finished = false, .forced = false, .granted = 1 }; for (0..200) |i| for ([_]usize{ 1, 4, 8, 16 }) |streams| { - _ = try g.drafts(streams); + _ = try g.draftsFresh(streams, whos[0..streams]); try g.observe(streams, i % 3 + 1, 4); }; try std.testing.expectEqual(first[0], counted.allocations); diff --git a/zig/src/core/lanes/stream.zig b/zig/src/core/lanes/stream.zig index 8c348ff95f..a35c8aaddc 100644 --- a/zig/src/core/lanes/stream.zig +++ b/zig/src/core/lanes/stream.zig @@ -225,6 +225,7 @@ pub const Stream = struct { served: ?u64 = null, proposal_depth: ?u32 = null, // null: legacy; otherwise the last routed request, zero until it succeeds granted: ?u32 = null, + grant_short: bool = false, // the last grant verified fewer drafts than the stream held copy_width: ?u32 = null, pub fn init(gpa: Allocator, spec: Spec) !Stream { @@ -287,6 +288,7 @@ pub const Stream = struct { s.inflight = null; s.served = null; s.granted = null; + s.grant_short = false; s.copy_width = null; s.proposal_depth = null; } diff --git a/zig/src/core/lanes/windows.zig b/zig/src/core/lanes/windows.zig index b70c2f634d..278b761325 100644 --- a/zig/src/core/lanes/windows.zig +++ b/zig/src/core/lanes/windows.zig @@ -226,7 +226,7 @@ pub fn probe(e: *Engine) depth.Probe { } pub fn who(s: *Stream) depth.Who { - return .{ .state = &s.depth, .draft_room = s.draftRoom(), .finished = s.finished, .forced = s.force.items.len > 0, .granted = s.granted, .stream = s, .sampled = s.sampling != null }; + return .{ .state = &s.depth, .draft_room = s.draftRoom(), .finished = s.finished, .forced = s.force.items.len > 0, .granted = s.granted, .grant_short = s.grant_short, .stream = s, .sampled = s.sampling != null }; } /// Whether any of a round's streams draws keyed samples (its rows carry their host work). @@ -282,7 +282,10 @@ pub fn allocate(e: *Engine, plans: []Plan) !void { try trail.event(e, &.{ f("ev", trail.str("lone-price")), f("stream", trail.str(plans[0].stream.id)), f("rows", trail.int(counts[0] + 1)), f("expected", ev.bits(expected)), f("ms", if (ms) |value| ev.bits(value) else .null), f("probe", .{ .bool = trial != null }) }); } for (plans, counts) |*p, c| { - if (p.kind == .head) p.stream.granted = c; + if (p.kind == .head) { + p.stream.granted = c; + if (e.cfg.granted_depth) p.stream.grant_short = c < p.count(); + } if (p.kind == .forced or p.count() == 0 or c == p.count()) continue; try cut.prefix(a, p, c); } From 27089360acf7ff50ce12840577559445cb02d20e Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:07:53 +0100 Subject: [PATCH 07/46] lanes: a drafter whose hold is a pass of its own: no hold beside a copy, and holds kept only while they pay The depth rule asks one head draft for a stream with a copy ahead, and a shared round's granted holds were priced only by their verify rows. A drafter whose hold runs as a pass of its own (Facts.copy_first) pays that pass for the one draft, and the copy round then drops it; Model.copy_first now asks it for none. A drafter may also give its hold pass's timed cost by blocks and depth (Facts.hold_costs, priced by hold_price.zig). With granted depths a shared round then keeps granted holds, most expected drafts a row first, only while each raises the next round's expected tokens a ms over its verify rows and the pass; a probe after a grant of none keeps its hold. A stream's last grant records the drafts it expected to keep (Stream.grant_gain). Both are off by default (copy_first false, hold_costs empty). Tests: with copy_first no hold is asked for a copy round and replies equal one-token decode at 1 and 3 streams; a hold's cost between and past its timed passes; granted holds kept while they pay, none for a dear pass, and a probe's hold kept. --- zig/src/core/lanes/config.zig | 4 + zig/src/core/lanes/depth.zig | 123 ++++++++++++++++++++- zig/src/core/lanes/drafted.zig | 2 + zig/src/core/lanes/drafted_copy_test.zig | 135 +++++++++++++++++++++++ zig/src/core/lanes/drafter.zig | 2 + zig/src/core/lanes/hold_price.zig | 58 ++++++++++ zig/src/core/lanes/lanes.zig | 2 + zig/src/core/lanes/shared_guard.zig | 1 + zig/src/core/lanes/stream.zig | 2 + zig/src/core/lanes/windows.zig | 11 +- 10 files changed, 335 insertions(+), 5 deletions(-) create mode 100644 zig/src/core/lanes/drafted_copy_test.zig create mode 100644 zig/src/core/lanes/hold_price.zig diff --git a/zig/src/core/lanes/config.zig b/zig/src/core/lanes/config.zig index 340788ccfe..a20d219c55 100644 --- a/zig/src/core/lanes/config.zig +++ b/zig/src/core/lanes/config.zig @@ -39,6 +39,8 @@ pub const Model = struct { draft_slack: u32 = 0, // drafts past a stream's last grant, short or not (0: the Constants' own) draft_streams: bool = false, // the head drafts every stream of a shared round in one batch shared_guard: bool = false, // a shared round drafts only while drafted rounds measure faster (shared_guard.zig) + copy_first: bool = false, // a copy ahead takes the round: the head holds no drafts for it (a drafter's own pass) + hold_costs: []const @import("hold_price.zig").HoldCost = &.{}, // a drafter's own hold pass, timed at load head_trees: bool = false, // the head drafts trees of lanes (DraftRequest.lanes) and holds them join_tail: bool = false, // a prompt in pieces leaves its last token for the next shared round's window (Backend.join) piece_runs: bool = false, // a round's pieces go to the backend in one call (Backend.pieces: one forward over them) @@ -169,6 +171,8 @@ pub const Config = struct { .draft_slack = if (m.draft_slack > 0) m.draft_slack else (Constants{}).draft_slack, .grant_slack = if (m.draft_slack > 0) m.draft_slack else (Constants{}).grant_slack, .short_slack = if (m.draft_slack > 0) m.draft_slack else (Constants{}).short_slack, + .copy_first = m.copy_first and !m.context_copies, + .hold_costs = m.hold_costs, }, .draft_step_ms = m.draft_step_ms, }; diff --git a/zig/src/core/lanes/depth.zig b/zig/src/core/lanes/depth.zig index 3c45e6ee6d..02d096a07c 100644 --- a/zig/src/core/lanes/depth.zig +++ b/zig/src/core/lanes/depth.zig @@ -13,8 +13,13 @@ pub const Constants = struct { draft_slack: i64 = 2, // drafts a stream offers past what its last shared round granted it grant_slack: i64 = 1, // granted depths: drafts past a grant that verified every draft the stream held short_slack: i64 = 0, // granted depths: drafts past a grant that left some of them unverified + copy_first: bool = false, // a copy ahead takes the next round whole: the head drafts none for it (else one) + hold_costs: []const HoldCost = &.{}, // a drafter's own hold pass by blocks and depth: granted holds must pay for it }; +pub const HoldCost = @import("hold_price.zig").HoldCost; +const holdMs = @import("hold_price.zig").holdMs; + /// A stream's acceptance by depth and probe counters (Python `_depth_state[stream]`). pub const State = struct { p: []f64, @@ -37,6 +42,7 @@ pub const Who = struct { forced: bool, granted: ?u32 = null, grant_short: bool = false, // the last grant left some held drafts unverified (the width did not bind) + gain: ?f64 = null, // the drafts its last grant expected to keep (null: none priced) stream: ?*anyopaque = null, sampled: bool = false, // keyed draws: host work a row that a greedy round does not have }; @@ -212,7 +218,7 @@ pub const Rule = struct { return @min(most, @as(i64, granted) + r.slack(who)); } - /// Python `_head_depth`: the next round's head drafts for a stream (0 when forced, 1 with a copy ahead). + /// Python `_head_depth`: the next round's head drafts for a stream (0 when forced, 1 or none with a copy ahead). pub fn headDepth(r: *Rule, who: *const Who, budget: ?i64, probe: Probe) !i64 { const d: i64 = if (who.finished) 0 else if (budget) |b| b else if (r.granted_depth) try r.grantDepth(who.*) @@ -221,7 +227,7 @@ pub const Rule = struct { else try r.depth(who.*); if (who.forced) return 0; - if (d != 0 and probe.copy(probe.ptr, who)) return 1; + if (d != 0 and probe.copy(probe.ptr, who)) return @intFromBool(!r.k.copy_first); return d; } @@ -234,6 +240,65 @@ pub const Rule = struct { for (out) |*d| d.* = @min(d.*, most); } + /// Granted holds kept, most expected drafts a row first, while they raise the next round's expected tokens a ms. + fn payHolds(r: *Rule, whos: []const Who, sampled: bool, out: []i64) !void { + const order = try r.gpa.alloc(usize, whos.len); + defer r.gpa.free(order); + var m: usize = 0; + var probes: usize = 0; // a probe after a grant of none keeps its hold: it measures the stream again + var deepest: i64 = 0; + for (whos, out, 0..) |w, d, i| { + if (d <= 0) continue; + if (w.granted == @as(?u32, 0)) { + probes += 1; + deepest = @max(deepest, d); + } else { + order[m] = i; + m += 1; + } + } + if (m == 0) return; + const Ctx = struct { + r: *Rule, + whos: []const Who, + fn less(c: @This(), x: usize, y: usize) bool { + const gx = c.r.gainOf(c.whos[x]) catch 0; + const gy = c.r.gainOf(c.whos[y]) catch 0; + return gx * rowsOf(c.whos[y]) > gy * rowsOf(c.whos[x]) or (gx * rowsOf(c.whos[y]) == gy * rowsOf(c.whos[x]) and x < y); + } + }; + std.mem.sort(usize, order[0..m], Ctx{ .r = r, .whos = whos }, Ctx.less); + const fixed = r.fixedMs(whos.len, sampled); + var rows: f64 = @floatFromInt(whos.len); + var tokens = rows; + const plain = r.shared_costs.get(@intCast(whos.len)) orelse return; + const probe_ms = if (probes > 0) holdMs(r.k.hold_costs, probes, deepest) else 0; + var best = tokens / (plain + fixed + probe_ms); + var keep: usize = 0; + for (order[0..m], 1..) |i, k| { + rows += rowsOf(whos[i]); + tokens += try r.gainOf(whos[i]); + deepest = @max(deepest, out[i]); + const verify = r.shared_costs.get(@intFromFloat(@ceil(rows))) orelse break; + const now = tokens / (verify + fixed + holdMs(r.k.hold_costs, k + probes, deepest)); + if (now > best) { + best = now; + keep = k; + } + } + for (order[keep..m]) |i| out[i] = 0; + } + + /// The drafts a granted stream's next round is expected to keep: its last grant's, or its first draft's rate. + fn gainOf(r: *Rule, w: Who) !f64 { + if (w.gain) |g| if (w.granted != null) return g; + return (try r.rates(w))[0]; + } + + fn rowsOf(w: Who) f64 { + return @floatFromInt(@max(1, w.granted orelse 1)); + } + /// Head drafts for the streams of a shared round, sized by allocation over their chain chances. pub fn budgets(r: *Rule, whos: []Who, probe: Probe, sampled: bool, out: []i64) !void { if (r.confidence_floor != null) { @@ -247,6 +312,7 @@ pub const Rule = struct { o.* = try r.headDepth(w, if (r.granted_depth) inner else @min(inner, granted + r.k.draft_slack), probe); } if (r.granted_depth) capPass(out, r.batch_rows); + if (r.granted_depth and r.k.hold_costs.len > 0) try r.payHolds(whos, sampled, out); return; } const most = try r.gpa.alloc(i64, whos.len); @@ -438,3 +504,56 @@ test "granted depths in a shared round fit one hold pass: every block is as deep try rule.budgets(whos[0..2], .{ .ptr = &unused, .copy = No.copy }, false, out[0..2]); try std.testing.expectEqualSlices(i64, &.{ 2, 1 }, out[0..2]); // a grant short of the held drafts: no slack } + +test "a copy ahead: one head draft beside it, or none when the head's drafts cost a pass of their own" { + const empty: Table = .{}; + const Yes = struct { + fn copy(_: *anyopaque, _: *const Who) bool { + return true; + } + }; + var unused: u8 = 0; + const probe: Probe = .{ .ptr = &unused, .copy = Yes.copy }; + var st: ?State = null; + defer if (st) |*x| x.deinit(std.testing.allocator); + const who: Who = .{ .state = &st, .draft_room = 100, .finished = false, .forced = false, .granted = 3 }; + for ([_]bool{ false, true }) |first| { + var rule: Rule = .{ .gpa = std.testing.allocator, .k = .{ .copy_first = first }, .prior = &.{0.8}, .most_drafts = 9, .family_costs = &empty, .shared_costs = &empty, .mtp_step_ms = 0.5, .plain_guard = false, .node_probabilities = true, .granted_depth = true, .batch_rows = 32 }; + defer rule.deinit(); + try std.testing.expectEqual(@as(i64, @intFromBool(!first)), try rule.headDepth(&who, null, probe)); + var out: [1]i64 = undefined; + var whos = [_]Who{who}; + try rule.budgets(&whos, probe, false, &out); + try std.testing.expectEqual(@as(i64, @intFromBool(!first)), out[0]); + } +} + +test "a granted hold stays only while its stream's expected drafts pay for the verify rows and the hold pass" { + const gpa = std.testing.allocator; + var shared: Table = .{}; + defer shared.deinit(gpa); + for (1..33) |w| try shared.put(gpa, @intCast(w), 10.0 + @as(f64, @floatFromInt(w))); + const empty: Table = .{}; + const cheap = [_]HoldCost{ .{ .streams = 1, .depth = 1, .ms = 1.0 }, .{ .streams = 8, .depth = 1, .ms = 4.5 } }; + var states: [4]?State = @splat(null); + defer for (&states) |*st| if (st.*) |*x| x.deinit(gpa); + const granted = [_]?u32{ 2, 2, 0, 1 }; + const gains = [_]?f64{ 1.8, 0.2, null, 0.9 }; + var whos: [4]Who = undefined; + for (&whos, &states, granted, gains) |*w, *st, g, x| w.* = .{ .state = st, .draft_room = 100, .finished = false, .forced = false, .granted = g, .gain = x }; + var rule: Rule = .{ .gpa = gpa, .k = .{ .hold_costs = &cheap }, .prior = &.{0.8}, .most_drafts = 9, .family_costs = &empty, .shared_costs = &shared, .mtp_step_ms = 0.5, .plain_guard = true, .node_probabilities = true, .granted_depth = true, .batch_rows = 32 }; + defer rule.deinit(); + // 8 ms overhead: plain 4 / 22; stream 0 then 3 raise the rate (5.8 / 25, 6.7 / 26.5), stream 1 would lower it + var out = [_]i64{ 4, 4, 0, 3 }; + try rule.payHolds(&whos, false, &out); + try std.testing.expectEqualSlices(i64, &.{ 4, 0, 0, 3 }, &out); + const dear = [_]HoldCost{.{ .streams = 1, .depth = 1, .ms = 20.0 }}; + rule.k.hold_costs = &dear; + out = .{ 4, 4, 0, 3 }; + try rule.payHolds(&whos, false, &out); + try std.testing.expectEqualSlices(i64, &.{ 0, 0, 0, 0 }, &out); // no gain pays a 20 ms pass + whos[2].granted = 0; + out = .{ 4, 4, 2, 3 }; + try rule.payHolds(&whos, false, &out); + try std.testing.expectEqualSlices(i64, &.{ 0, 0, 2, 0 }, &out); // a probe after a grant of none keeps its hold +} diff --git a/zig/src/core/lanes/drafted.zig b/zig/src/core/lanes/drafted.zig index 50f5739913..211f30112c 100644 --- a/zig/src/core/lanes/drafted.zig +++ b/zig/src/core/lanes/drafted.zig @@ -104,6 +104,8 @@ pub const Drafted = struct { m.draft_probabilities = false; m.head_trees = d.trees; m.shared_guard = d.shared_guard; + m.copy_first = d.copy_first; + m.hold_costs = d.hold_costs; if (d.chances and x.drafter.vtable.chances != null) { // a chain's drafts priced by the drafter's own chances; the next hold as deep as the grant and its slack m.draft_probabilities = d.confidence_floor == null; diff --git a/zig/src/core/lanes/drafted_copy_test.zig b/zig/src/core/lanes/drafted_copy_test.zig new file mode 100644 index 0000000000..329731f69c --- /dev/null +++ b/zig/src/core/lanes/drafted_copy_test.zig @@ -0,0 +1,135 @@ +//! A drafter whose hold costs a pass of its own: no hold for a stream whose next round is a copy. +const std = @import("std"); +const Config = @import("config.zig").Config; +const Engine = @import("engine.zig").Engine; +const be = @import("backend.zig"); +const sm = @import("stream.zig"); +const fake = @import("fake.zig"); +const dr = @import("drafter.zig"); +const Proposer = @import("proposer.zig").Proposer; +const Drafted = @import("drafted.zig").Drafted; +const dt = @import("drafted_test.zig"); + +const gpa = std.testing.allocator; + +/// Copies the fake target's own continuation (it lands whole) whenever the context's length is a multiple of three. +const Copier = struct { + out: [3]u32 = undefined, + match: i64 = 0, + + fn propose(ptr: *anyopaque, context: []const u32, most: i64) anyerror![]const u32 { + const c: *Copier = @ptrCast(@alignCast(ptr)); + c.match = 0; + if (context.len % 3 != 0 or most < 2) return &.{}; + var history: [128]u32 = undefined; + @memcpy(history[0..context.len], context); + const n: usize = @intCast(@min(3, most)); + for (c.out[0..n], 0..) |*t, k| { + t.* = fake.next(history[0 .. context.len + k], null, 0); + history[context.len + k] = t.*; + } + c.match = 100; + return c.out[0..n]; + } + + fn lastMatch(ptr: *anyopaque) i64 { + return @as(*Copier, @ptrCast(@alignCast(ptr))).match; + } + + const vtable: Proposer.VTable = .{ .propose = propose, .last_match = lastMatch }; +}; + +/// Each stream's last hold position, and the copy windows verified at a position a hold was asked for. +const Watch = struct { + var inner_hold: *const fn (*anyopaque, []const dr.Hold) anyerror!void = undefined; + var inner_verify: *const fn (*anyopaque, []const be.Window, []be.Verified) anyerror!void = undefined; + var at: [4]?u64 = @splat(null); + var streams: [4]?*sm.Stream = @splat(null); + var wasted: u32 = 0; + var copies: u32 = 0; + var blocks: u32 = 0; // the streams' blocks held, over every hold + + fn slot(s: *sm.Stream) usize { + for (&streams, 0..) |*x, i| if (x.* == null or x.* == s) { + x.* = s; + return i; + }; + unreachable; + } + + fn hold(ptr: *anyopaque, items: []const dr.Hold) anyerror!void { + blocks += @intCast(items.len); + // a joining stream's first hold comes before its first token is read: no copy can be foreseen for it + for (items) |it| at[slot(it.stream)] = if (it.stream.emitted().len > 0) it.position else null; + return inner_hold(ptr, items); + } + + fn verify(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { + for (windows) |w| if (w.held == 0 and w.tokens.len > 0) { + copies += 1; + if (at[slot(w.stream)] == w.positions[0]) wasted += 1; + }; + return inner_verify(ptr, windows, out); + } +}; + +const Run = struct { tokens: [3][64]u32 = undefined, lens: [3]usize = @splat(0), holds: u32 = 0, wasted: u32 = 0, copies: u32 = 0 }; + +fn run(n: usize, drafts: bool, copy_first: bool) !Run { + Watch.at = @splat(null); + Watch.streams = @splat(null); + Watch.wasted = 0; + Watch.copies = 0; + Watch.blocks = 0; + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var head: dt.FakeDrafter = .{ .facts_ = .{ .depth = 4, .step_ms = 0.5, .batched = true, .copy_first = copy_first } }; + defer head.deinit(); + var hooked = head.drafter(); + var table = hooked.vtable.*; + Watch.inner_hold = table.hold; + table.hold = Watch.hold; + hooked.vtable = &table; + var wrapped = try Drafted.init(gpa, target.backend(), hooked); + defer wrapped.deinit(); + // the round loop's windows, before the wrapper turns held drafts into host tokens + var seen = if (drafts) wrapped.backend() else target.backend(); + var verified = seen.vtable.*; + Watch.inner_verify = verified.verify; + verified.verify = Watch.verify; + seen.vtable = &verified; + const m: @import("config.zig").Model = .{ .exact_width = 16, .gpu_tokens = true, .hidden_rows = true, .batch_rows = 32, .max_streams = 4 }; + var cfg = try Config.init(gpa, if (drafts) wrapped.facts(m) else m, 16, 15); + defer cfg.deinit(gpa); + var clock: fake.FixedClock = .{}; + var engine = Engine.init(gpa, &cfg, seen, clock.clock()); + defer engine.deinit(); + var copiers: [3]Copier = @splat(.{}); + var streams: [3]sm.Stream = undefined; + for (streams[0..n], copiers[0..n], 0..) |*s, *c, i| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = dt.long[0 .. 12 + i], .max_new = 40, .eos = &.{96}, .drafts = drafts, .proposer = .{ .ptr = c, .vtable = &Copier.vtable } }); + defer for (streams[0..n]) |*s| s.deinit(gpa); + for (streams[0..n]) |*s| try engine.addStream(s); + while (engine.activeCount() > 0) try engine.step(); + var out: Run = .{ .holds = Watch.blocks, .wasted = Watch.wasted, .copies = Watch.copies }; + for (streams[0..n], 0..) |*s, i| { + out.lens[i] = s.emitted().len; + @memcpy(out.tokens[i][0..out.lens[i]], s.emitted()); + } + return out; +} + +test "a copy ahead takes the round: with copy_first no hold is asked for it, and replies equal one-token decode" { + for ([_]usize{ 1, 3 }) |n| { + const plain = try run(n, false, false); + const before = try run(n, true, false); + const after = try run(n, true, true); + for (0..n) |i| { + try std.testing.expectEqualSlices(u32, plain.tokens[i][0..plain.lens[i]], before.tokens[i][0..before.lens[i]]); + try std.testing.expectEqualSlices(u32, plain.tokens[i][0..plain.lens[i]], after.tokens[i][0..after.lens[i]]); + } + try std.testing.expect(before.copies > 0 and after.copies > 0); + try std.testing.expect(before.wasted > 0); // a hold of one draft beside each copy, dropped by the copy round + try std.testing.expectEqual(@as(u32, 0), after.wasted); + try std.testing.expect(after.holds < before.holds); + } +} diff --git a/zig/src/core/lanes/drafter.zig b/zig/src/core/lanes/drafter.zig index 35a79e2e1c..3ba9feb671 100644 --- a/zig/src/core/lanes/drafter.zig +++ b/zig/src/core/lanes/drafter.zig @@ -22,6 +22,8 @@ pub const Facts = struct { chances: bool = false, // chances() gives each held chain draft's chance: verify rows go by them (Model.granted_depth) confidence_floor: ?f64 = null, // with chances: every round drafts the deepest chain and verifies drafts this likely slack: u32 = 0, // with chances: drafts a hold offers past the stream's last grant (0: the depth rule's own) + copy_first: bool = false, // a hold costs a pass of its own: none is asked for a stream with a copy ahead + hold_costs: []const @import("hold_price.zig").HoldCost = &.{}, // that pass's ms by blocks and depth (timed at load) }; /// Target rows [start, start + features.rows) and the token after each; `rows`, when given, picks which to absorb. diff --git a/zig/src/core/lanes/hold_price.zig b/zig/src/core/lanes/hold_price.zig new file mode 100644 index 0000000000..e4c854410e --- /dev/null +++ b/zig/src/core/lanes/hold_price.zig @@ -0,0 +1,58 @@ +//! A drafter's hold pass priced from its timed passes: by the blocks it holds and their depth. +const std = @import("std"); + +/// A drafter's hold pass of `streams` blocks, each `depth` drafts deep, in ms (timed at load). +pub const HoldCost = struct { streams: u32, depth: u32, ms: f64 }; + +/// A hold pass's ms: linear in depth at each timed stream count, then in streams (past the widest, by its slope). +pub fn holdMs(costs: []const HoldCost, streams: usize, depth: i64) f64 { + var below: ?u32 = null; + var above: ?u32 = null; + for (costs) |c| { + if (c.streams <= streams and (below == null or c.streams > below.?)) below = c.streams; + if (c.streams >= streams and (above == null or c.streams < above.?)) above = c.streams; + } + const lo = below orelse above.?; + const hi = above orelse lo; + const at_lo = atDepth(costs, lo, depth); + if (hi == lo) { + if (above != null) return at_lo; + // past the widest timed count: the slope from the count below it + var under: ?u32 = null; + for (costs) |c| if (c.streams < lo and (under == null or c.streams > under.?)) { + under = c.streams; + }; + const u = under orelse return at_lo * @as(f64, @floatFromInt(streams)) / @as(f64, @floatFromInt(lo)); + const slope = (at_lo - atDepth(costs, u, depth)) / @as(f64, @floatFromInt(lo - u)); + return at_lo + slope * @as(f64, @floatFromInt(streams - lo)); + } + const t = @as(f64, @floatFromInt(streams - lo)) / @as(f64, @floatFromInt(hi - lo)); + return at_lo + t * (atDepth(costs, hi, depth) - at_lo); +} + +/// One stream count's hold ms at `depth`: linear between its two nearest timed depths (one alone: as timed). +fn atDepth(costs: []const HoldCost, streams: u32, depth: i64) f64 { + var a: ?HoldCost = null; + var b: ?HoldCost = null; + for (costs) |c| { + if (c.streams != streams) continue; + if (a == null or c.depth < a.?.depth) { + b = a; + a = c; + } else if (b == null or c.depth < b.?.depth) b = c; + } + const x = a.?; + const y = b orelse return x.ms; + const t = (@as(f64, @floatFromInt(depth)) - @as(f64, @floatFromInt(x.depth))) / @as(f64, @floatFromInt(y.depth - x.depth)); + return x.ms + t * (y.ms - x.ms); +} + +test "a hold's ms between and past the timed passes" { + const t = [_]HoldCost{ .{ .streams = 1, .depth = 1, .ms = 1.0 }, .{ .streams = 1, .depth = 4, .ms = 2.5 }, .{ .streams = 4, .depth = 1, .ms = 3.0 }, .{ .streams = 4, .depth = 4, .ms = 4.5 }, .{ .streams = 8, .depth = 1, .ms = 4.0 } }; + try std.testing.expectApproxEqAbs(@as(f64, 1.0), holdMs(&t, 1, 1), 1e-12); + try std.testing.expectApproxEqAbs(@as(f64, 2.0), holdMs(&t, 1, 3), 1e-12); // linear in depth + try std.testing.expectApproxEqAbs(@as(f64, 3.0), holdMs(&t, 1, 5), 1e-12); // and past the deepest + try std.testing.expectApproxEqAbs(@as(f64, 1.0 + 2.0 / 3.0), holdMs(&t, 2, 1), 1e-12); // linear in streams + try std.testing.expectApproxEqAbs(@as(f64, 4.0), holdMs(&t, 8, 3), 1e-12); // one depth timed: as timed + try std.testing.expectApproxEqAbs(@as(f64, 5.0), holdMs(&t, 12, 1), 1e-12); // past the widest: the slope below it +} diff --git a/zig/src/core/lanes/lanes.zig b/zig/src/core/lanes/lanes.zig index 8f8adf1e1b..c5cd5f5e24 100644 --- a/zig/src/core/lanes/lanes.zig +++ b/zig/src/core/lanes/lanes.zig @@ -18,6 +18,7 @@ pub const cost_rule = @import("cost_rule.zig"); pub const gpu_rule = @import("gpu_rule.zig"); pub const gpu_full = @import("gpu_full.zig"); pub const depth = @import("depth.zig"); +pub const hold_price = @import("hold_price.zig"); pub const mtp_policy = @import("mtp_policy.zig"); pub const trim = @import("trim.zig"); pub const config = @import("config.zig"); @@ -55,6 +56,7 @@ test { _ = @import("gpu_full_test.zig"); _ = @import("drafted_test.zig"); _ = @import("drafted_chances_test.zig"); + _ = @import("drafted_copy_test.zig"); _ = @import("root_copy_test.zig"); _ = @import("drafted_grammar_test.zig"); _ = @import("windows_grammar_test.zig"); diff --git a/zig/src/core/lanes/shared_guard.zig b/zig/src/core/lanes/shared_guard.zig index cdaf870736..52dc848b95 100644 --- a/zig/src/core/lanes/shared_guard.zig +++ b/zig/src/core/lanes/shared_guard.zig @@ -52,6 +52,7 @@ pub const Guard = struct { if (e.mode == 1 and e.skip) for (whos) |*w| { w.granted = 1; w.grant_short = false; + w.gain = null; // priced by its first draft's rate, not a grant of none }; return e.mode == 1; } diff --git a/zig/src/core/lanes/stream.zig b/zig/src/core/lanes/stream.zig index a35c8aaddc..41ffc9054c 100644 --- a/zig/src/core/lanes/stream.zig +++ b/zig/src/core/lanes/stream.zig @@ -226,6 +226,7 @@ pub const Stream = struct { proposal_depth: ?u32 = null, // null: legacy; otherwise the last routed request, zero until it succeeds granted: ?u32 = null, grant_short: bool = false, // the last grant verified fewer drafts than the stream held + grant_gain: ?f64 = null, // the drafts its last grant expected to keep (the chances it was priced by) copy_width: ?u32 = null, pub fn init(gpa: Allocator, spec: Spec) !Stream { @@ -289,6 +290,7 @@ pub const Stream = struct { s.served = null; s.granted = null; s.grant_short = false; + s.grant_gain = null; s.copy_width = null; s.proposal_depth = null; } diff --git a/zig/src/core/lanes/windows.zig b/zig/src/core/lanes/windows.zig index 278b761325..41f87b769c 100644 --- a/zig/src/core/lanes/windows.zig +++ b/zig/src/core/lanes/windows.zig @@ -226,7 +226,7 @@ pub fn probe(e: *Engine) depth.Probe { } pub fn who(s: *Stream) depth.Who { - return .{ .state = &s.depth, .draft_room = s.draftRoom(), .finished = s.finished, .forced = s.force.items.len > 0, .granted = s.granted, .grant_short = s.grant_short, .stream = s, .sampled = s.sampling != null }; + return .{ .state = &s.depth, .draft_room = s.draftRoom(), .finished = s.finished, .forced = s.force.items.len > 0, .granted = s.granted, .grant_short = s.grant_short, .gain = s.grant_gain, .stream = s, .sampled = s.sampling != null }; } /// Whether any of a round's streams draws keyed samples (its rows carry their host work). @@ -281,10 +281,15 @@ pub fn allocate(e: *Engine, plans: []Plan) !void { const ms = own.get(1 + @as(i64, counts[0])); try trail.event(e, &.{ f("ev", trail.str("lone-price")), f("stream", trail.str(plans[0].stream.id)), f("rows", trail.int(counts[0] + 1)), f("expected", ev.bits(expected)), f("ms", if (ms) |value| ev.bits(value) else .null), f("probe", .{ .bool = trial != null }) }); } - for (plans, counts) |*p, c| { + for (plans, counts, probs) |*p, c, pr| { if (p.kind == .head) { p.stream.granted = c; - if (e.cfg.granted_depth) p.stream.grant_short = c < p.count(); + if (e.cfg.granted_depth) { + p.stream.grant_short = c < p.count(); + var gain: f64 = 0; + for (pr[0..@min(c, pr.len)]) |chance| gain += chance; + p.stream.grant_gain = gain; + } } if (p.kind == .forced or p.count() == 0 or c == p.count()) continue; try cut.prefix(a, p, c); From b66a7629e60d7283d5a837f7b3d2632f03000896 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:09:02 +0100 Subject: [PATCH 08/46] lanes: a lone stream holds drafts only while its drafted rounds measure faster With granted depths a lone stream's verify rows are priced over the shared costs and the one-stream overhead, which is measured on drafted rounds and so holds the drafter's hold time; nothing compared a lone drafted round's tokens a ms with a plain round's. lone_guard.zig measures each lone stream's rounds of held drafts in stretches of 16 against plain rounds' tokens a ms, kept from stream to stream by sampling mode and context band and taken once in a plain stretch when unknown; copy rounds are left out of both. Drafting goes on while it beats plain by 5%; else the stream runs plain rounds and measures drafting again after 16 plain rounds, doubling to 512 while plain keeps winning. A drafted stretch's first hold drafts one and its slack, and a mode's first round is not measured. It runs only with granted depths and the shared-round guard on a timed engine (Facts.chances and Facts.shared_guard); shared rounds keep their own guard. Off by default. Tests: the guard's verdicts, rarer probes and kept plain rate; through the round loop a lone stream whose drafted rounds lose runs mostly plain and one whose drafted rounds win keeps drafting, replies equal to one-token decode. --- zig/src/core/lanes/depth.zig | 29 +++- zig/src/core/lanes/drafted_chances_test.zig | 8 +- zig/src/core/lanes/engine.zig | 2 + zig/src/core/lanes/lanes.zig | 2 + zig/src/core/lanes/lone_guard.zig | 155 ++++++++++++++++++++ zig/src/core/lanes/lone_guard_test.zig | 115 +++++++++++++++ 6 files changed, 305 insertions(+), 6 deletions(-) create mode 100644 zig/src/core/lanes/lone_guard.zig create mode 100644 zig/src/core/lanes/lone_guard_test.zig diff --git a/zig/src/core/lanes/depth.zig b/zig/src/core/lanes/depth.zig index 02d096a07c..e02c1e6240 100644 --- a/zig/src/core/lanes/depth.zig +++ b/zig/src/core/lanes/depth.zig @@ -19,6 +19,7 @@ pub const Constants = struct { pub const HoldCost = @import("hold_price.zig").HoldCost; const holdMs = @import("hold_price.zig").holdMs; +const lone_guard = @import("lone_guard.zig"); /// A stream's acceptance by depth and probe counters (Python `_depth_state[stream]`). pub const State = struct { @@ -27,6 +28,7 @@ pub const State = struct { plain: ?u64 = null, wait: ?u64 = null, ms: ?Table = null, // plain_guard: the stream's own round times + lone: lone_guard.Lone = .{}, // lone_guard: whether its lone holds draft pub fn deinit(st: *State, gpa: Allocator) void { gpa.free(st.p); @@ -45,6 +47,7 @@ pub const Who = struct { gain: ?f64 = null, // the drafts its last grant expected to keep (null: none priced) stream: ?*anyopaque = null, sampled: bool = false, // keyed draws: host work a row that a greedy round does not have + shared: bool = false, // a shared round's budget: the lone guard stays out }; /// Whether a copied continuation is ahead of a stream (the round loop proposes; Python `_copy_proposal`). @@ -68,6 +71,9 @@ pub const Rule = struct { plain_guard: bool, node_probabilities: bool, granted_depth: bool = false, // node_probabilities: a stream drafts its last grant plus draft_slack, lone ones too + lone_guard: bool = false, // granted_depth: a lone stream holds drafts only while its drafted rounds measure faster + lone_k: lone_guard.Constants = .{}, + lone_plain: lone_guard.Plain = .{}, // plain rounds' tokens a ms, every lone stream's confidence_floor: ?f64 = null, batch_rows: u32, batched_heads: bool = false, // one batch drafts every stream of a shared round (its steps: the deepest stream's) @@ -209,15 +215,33 @@ pub const Rule = struct { /// The last grant plus its slack (the deepest chain before any); after a grant of none, a probe of draft_slack. fn grantDepth(r: *Rule, who: Who) !i64 { const most = @min(@as(i64, r.most_drafts), @max(1, who.draft_room - 1)); + if (r.lone_guard and !who.shared) { + const l = &(try r.state(who)).lone; + if (!l.drafts()) return 0; + if (l.fresh) { + l.fresh = false; + return @min(most, 1 + r.k.grant_slack); + } + } const granted = who.granted orelse return most; - _ = try r.rates(who); - const st = &who.state.*.?; + const st = try r.state(who); if (granted == 0) return if (!r.plain_guard or r.probePlain(st)) @min(most, r.k.draft_slack) else 0; st.plain = 0; st.wait = r.k.depth_probe_every; return @min(most, @as(i64, granted) + r.slack(who)); } + fn state(r: *Rule, who: Who) !*State { + _ = try r.rates(who); + return &who.state.*.?; + } + + /// A lone round's committed tokens and wall ms at the stream's context band, for the lone guard. + pub fn observeLone(r: *Rule, who: Who, band: usize, tokens: usize, ms: f64) !void { + if (!r.lone_guard) return; + (try r.state(who)).lone.observe(r.lone_k, r.lone_plain.at(who.sampled, band), tokens, ms); + } + /// Python `_head_depth`: the next round's head drafts for a stream (0 when forced, 1 or none with a copy ahead). pub fn headDepth(r: *Rule, who: *const Who, budget: ?i64, probe: Probe) !i64 { const d: i64 = if (who.finished) 0 else if (budget) |b| b else if (r.granted_depth) @@ -307,6 +331,7 @@ pub const Rule = struct { } if (r.node_probabilities) { for (whos, out) |*w, *o| { + w.shared = true; const inner = try r.headDepth(w, null, probe); const granted: i64 = if (w.granted) |g| g else r.most_drafts; o.* = try r.headDepth(w, if (r.granted_depth) inner else @min(inner, granted + r.k.draft_slack), probe); diff --git a/zig/src/core/lanes/drafted_chances_test.zig b/zig/src/core/lanes/drafted_chances_test.zig index 02c5eb4cb1..d3f248e8da 100644 --- a/zig/src/core/lanes/drafted_chances_test.zig +++ b/zig/src/core/lanes/drafted_chances_test.zig @@ -17,7 +17,7 @@ const gpa = std.testing.allocator; const most_streams = 8; /// drafted_test's drafter with a chance for each held draft: `sure` where its guess is right, `unsure` where wrong. -const Priced = struct { +pub const Priced = struct { head: dt.FakeDrafter, sure: f64 = 0.7, unsure: f64 = 0.7, @@ -56,7 +56,7 @@ const Priced = struct { return true; } - fn drafter(p: *Priced) dr.Drafter { + pub fn drafter(p: *Priced) dr.Drafter { const base_ = p.head.drafter(); p.vtable = base_.vtable.*; p.inner_hold = base_.vtable.hold; @@ -65,7 +65,7 @@ const Priced = struct { return .{ .ptr = &p.head, .vtable = &p.vtable }; } - fn deinit(p: *Priced) void { + pub fn deinit(p: *Priced) void { p.at.deinit(gpa); p.head.deinit(); } @@ -91,7 +91,7 @@ const Seen = struct { var costs: [16]Cost = undefined; -fn base() Model { +pub fn base() Model { for (&costs, 1..) |*c, w| c.* = .{ .width = @intCast(w), .ms = 5.0 + 0.8 * @as(f64, @floatFromInt(w)) }; return .{ .exact_width = 16, .gpu_tokens = true, .window_costs = &costs, .hidden_rows = true, .batch_rows = 32, .max_streams = most_streams }; } diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index 9ff27b18d5..32e46c02e8 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -75,6 +75,7 @@ pub const Engine = struct { .plain_guard = cfg.plain_guard, .node_probabilities = cfg.node_probabilities, .granted_depth = cfg.granted_depth, + .lone_guard = cfg.granted_depth and cfg.shared_guard and cfg.timed, .confidence_floor = cfg.draft_confidence, .batch_rows = cfg.batch_rows, .batched_heads = cfg.draft_streams, @@ -238,6 +239,7 @@ pub const Engine = struct { lone_costs.observe(s, plan.cost_probe_rows, rows, ms); } } + if (e.rule.lone_guard and e.alone and s.drafts and (plan.kind == .head or plan.kind == .none)) try e.rule.observeLone(win.who(s), lone_costs.band(s.context.items.len), o.got.len, e.clock.elapsedMs(.cost)); if (e.cfg.timed and s.graft_room != null and s.rounds > 1) { // the stream's round time beyond its window and head levels, for the lane planner if (e.cfg.lane_costs.get(@intCast(rows))) |w| { diff --git a/zig/src/core/lanes/lanes.zig b/zig/src/core/lanes/lanes.zig index c5cd5f5e24..a3b2587704 100644 --- a/zig/src/core/lanes/lanes.zig +++ b/zig/src/core/lanes/lanes.zig @@ -63,6 +63,8 @@ test { _ = @import("pricing_test.zig"); _ = @import("depth_probe_test.zig"); _ = @import("shared_guard.zig"); + _ = @import("lone_guard.zig"); + _ = @import("lone_guard_test.zig"); _ = @import("drafted_guard_test.zig"); _ = @import("proposal_external_test.zig"); _ = @import("lone_costs_test.zig"); diff --git a/zig/src/core/lanes/lone_guard.zig b/zig/src/core/lanes/lone_guard.zig new file mode 100644 index 0000000000..d861174dc3 --- /dev/null +++ b/zig/src/core/lanes/lone_guard.zig @@ -0,0 +1,155 @@ +//! Whether a lone stream holds drafts: its drafted rounds' tokens a ms against plain rounds', the loser re-measured. +const std = @import("std"); + +pub const Constants = struct { + drafted: u32 = 16, // measured drafted rounds a verdict takes, again and again while drafting wins + plain: u32 = 16, // measured plain rounds that take the plain rate when there is none yet + wait: u32 = 16, // plain rounds before drafting is measured again; doubles each time plain wins again + wait_most: u32 = 512, + margin: f64 = 1.05, // drafted rounds must beat plain ones' tokens a ms by this much + blend: f64 = 0.5, // the newest stretch's weight in a rate +}; + +pub const Mode = enum { drafted, plain }; + +/// Plain rounds' tokens a ms by sampling mode and context band, kept from stream to stream (no draft changes them). +pub const Plain = struct { + rates: [2][bands]?f64 = @splat(@splat(null)), + + pub const bands = 6; + + pub fn at(p: *Plain, sampled: bool, band: usize) *?f64 { + return &p.rates[@intFromBool(sampled)][@min(band, bands - 1)]; + } +}; + +/// One lone stream's guard: the mode its holds follow, the stretch measured and its drafted rate. +pub const Lone = struct { + mode: Mode = .drafted, + left: u32 = 0, + tokens: f64 = 0, + ms: f64 = 0, + skip: bool = true, // a mode's first round still verifies the other mode's hold (or none): not measured + started: bool = false, + pending: bool = false, // a drafted verdict waits for the plain rate this stretch takes + fresh: bool = false, // a drafted stretch's first hold: one draft and its slack, whatever the last grant + drafted: ?f64 = null, + wait: u32 = 0, + + /// Whether the next hold may draft. + pub fn drafts(l: *const Lone) bool { + return l.mode == .drafted; + } + + /// A lone round of held drafts or none (copies left out): its tokens and wall ms; `plain` the band's rate. + pub fn observe(l: *Lone, k: Constants, plain: *?f64, tokens: usize, ms: f64) void { + if (!l.started) { + l.started = true; + l.left = k.drafted; + l.wait = k.wait; + } + if (l.skip) { + l.skip = false; + return; + } + l.tokens += @floatFromInt(tokens); + l.ms += ms; + l.left -|= 1; + if (l.left > 0) return; + const now = if (l.ms > 0) l.tokens / l.ms else 0; + l.tokens = 0; + l.ms = 0; + switch (l.mode) { + .drafted => { + l.drafted = blend(k, l.drafted, now); + if (plain.* == null) { + l.pending = true; + l.switchTo(.plain, k.plain); + } else l.verdict(k, plain.*.?); + }, + .plain => { + plain.* = blend(k, plain.*, now); + if (l.pending) { + l.pending = false; + l.verdict(k, plain.*.?); + } else l.switchTo(.drafted, k.drafted); + }, + } + } + + /// Drafting goes on while its rate beats plain's by the margin, else plain rounds for a wait that doubles. + fn verdict(l: *Lone, k: Constants, plain: f64) void { + if (l.drafted.? >= plain * k.margin) { + l.wait = k.wait; + l.switchTo(.drafted, k.drafted); + return; + } + l.switchTo(.plain, l.wait); + l.wait = @min(l.wait *| 2, k.wait_most); + } + + fn switchTo(l: *Lone, mode: Mode, rounds: u32) void { + if (mode != l.mode) { + l.skip = true; + l.fresh = mode == .drafted; + } + l.mode = mode; + l.left = @max(1, rounds); + } +}; + +fn blend(k: Constants, old: ?f64, now: f64) f64 { + return if (old) |r| r + k.blend * (now - r) else now; +} + +/// Rounds of one stream whose drafted rounds give `drafted` and plain ones `plain` tokens a ms: how many drafted. +fn run(l: *Lone, k: Constants, rate: *?f64, rounds: usize, plain: f64, drafted: f64) usize { + var n: usize = 0; + for (0..rounds) |_| { + const d = l.drafts(); + n += @intFromBool(d); + l.observe(k, rate, if (d) 2 else 1, if (d) 2 / drafted else 1 / plain); + } + return n; +} + +test "a lone stream whose drafted rounds lose runs plain but for probes that grow rarer" { + var l: Lone = .{}; + var rate: ?f64 = null; + const k: Constants = .{}; + const drafted = run(&l, k, &rate, 1000, 0.128, 0.122); + // the first drafted stretch, the plain rate's own, then probes after 16, 32, 64, 128, 256 and 512 plain rounds + try std.testing.expect(drafted <= 7 * (k.drafted + 1)); + try std.testing.expect(drafted > k.drafted); + try std.testing.expect(!l.drafts() or l.left <= k.drafted); + try std.testing.expectApproxEqAbs(@as(f64, 0.128), rate.?, 1e-9); +} + +test "a lone stream whose drafted rounds win keeps drafting; a later stream reuses the plain rate" { + var rate: ?f64 = null; + const k: Constants = .{}; + var first: Lone = .{}; + const n = run(&first, k, &rate, 1000, 0.128, 0.15); + try std.testing.expect(n > 1000 - k.plain - 2); // one plain stretch takes the plain rate + var second: Lone = .{}; + try std.testing.expectEqual(@as(usize, 300), run(&second, k, &rate, 300, 0.128, 0.15)); // no plain stretch again + var third: Lone = .{}; + const lost = run(&third, k, &rate, 300, 0.128, 0.11); + try std.testing.expect(lost < 300 / 3); // plain from its first verdict, without measuring plain again +} + +test "a drafted stretch's first round is not measured and its first hold is fresh" { + var l: Lone = .{}; + var rate: ?f64 = 0.2; + const k: Constants = .{ .drafted = 2, .wait = 2 }; + l.observe(k, &rate, 9, 1); // skipped: the stream's first round + l.observe(k, &rate, 1, 10); + l.observe(k, &rate, 1, 10); + try std.testing.expect(!l.drafts()); // 0.1 against 0.2 + try std.testing.expect(!l.fresh); + l.observe(k, &rate, 9, 1); // skipped: the last drafted round's hold + l.observe(k, &rate, 1, 5); + l.observe(k, &rate, 1, 5); + try std.testing.expect(l.drafts() and l.fresh); // the probe after the wait + try std.testing.expectApproxEqAbs(@as(f64, 0.2), rate.?, 1e-12); +} diff --git a/zig/src/core/lanes/lone_guard_test.zig b/zig/src/core/lanes/lone_guard_test.zig new file mode 100644 index 0000000000..2becd90158 --- /dev/null +++ b/zig/src/core/lanes/lone_guard_test.zig @@ -0,0 +1,115 @@ +//! The lone guard through the round loop: a lone stream drafts while drafting measures faster, else runs plain. +const std = @import("std"); +const Config = @import("config.zig").Config; +const Engine = @import("engine.zig").Engine; +const be = @import("backend.zig"); +const sm = @import("stream.zig"); +const fake = @import("fake.zig"); +const dr = @import("drafter.zig"); +const Drafted = @import("drafted.zig").Drafted; +const dt = @import("drafted_test.zig"); +const ct = @import("drafted_chances_test.zig"); + +const gpa = std.testing.allocator; + +/// A round's ms from its verify's rows and whether a hold ran before it, as a GPU that runs both in turn. +const Costs = struct { row: f64, extra: f64, hold: f64 }; + +const Cost = struct { + var costs: Costs = undefined; + var inner_verify: *const fn (*anyopaque, []const be.Window, []be.Verified) anyerror!void = undefined; + var inner_hold: *const fn (*anyopaque, []const dr.Hold) anyerror!void = undefined; + var rows: usize = 1; + var held: bool = false; // a hold since the last verify + var before: bool = false; // the last verify ran after a hold + var holds: u32 = 0; + var verifies: u32 = 0; + + fn verify(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { + rows = windows[0].rows(); + before = held; + held = false; + verifies += 1; + return inner_verify(ptr, windows, out); + } + + fn hold(ptr: *anyopaque, items: []const dr.Hold) anyerror!void { + for (items) |it| held = held or it.depth > 0; + holds += @intFromBool(held); + return inner_hold(ptr, items); + } + + fn clock() be.Clock { + return .{ .ptr = @ptrCast(&costs), .vtable = &.{ .start = start, .elapsed_ms = elapsed } }; + } + + fn start(_: *anyopaque) void {} + + fn elapsed(_: *anyopaque, _: be.Mark) f64 { + const extra: f64 = @floatFromInt(rows - 1); + return costs.row + costs.extra * extra + (if (before) costs.hold else 0); + } +}; + +const Run = struct { tokens: [400]u32 = undefined, len: usize = 0, holds: u32 = 0, rounds: u32 = 0, drafted: u64 = 0 }; + +/// One lone stream of 300 tokens; `guard` the drafter's guard switch, `costs` null: drafts off. +fn run(costs: ?Costs, guard: bool) !Run { + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var p: ct.Priced = .{ .head = .{ .facts_ = .{ .depth = 9, .step_ms = 0.5, .batched = true, .chances = true, .shared_guard = guard } }, .sure = 0.95, .unsure = 0.05 }; + defer p.deinit(); + var strict: dt.Strict = .{ .inner = target.backend() }; + defer strict.deinit(); + var seen = strict.backend(); + var table = seen.vtable.*; + Cost.inner_verify = table.verify; + table.verify = Cost.verify; + seen.vtable = &table; + var drafter = p.drafter(); + var hold_table = drafter.vtable.*; + Cost.inner_hold = hold_table.hold; + hold_table.hold = Cost.hold; + drafter.vtable = &hold_table; + Cost.costs = costs orelse .{ .row = 8, .extra = 0, .hold = 0 }; + Cost.holds = 0; + Cost.verifies = 0; + Cost.held = false; + var wrapped = try Drafted.init(gpa, seen, drafter); + defer wrapped.deinit(); + const drafts = costs != null; + var cfg = try Config.init(gpa, if (drafts) wrapped.facts(ct.base()) else ct.base(), 16, 15); + defer cfg.deinit(gpa); + var engine = Engine.init(gpa, &cfg, if (drafts) wrapped.backend() else target.backend(), Cost.clock()); + defer engine.deinit(); + var s = try sm.Stream.init(gpa, .{ .id = "s", .prompt = &dt.p1, .max_new = 300, .eos = &.{}, .drafts = drafts }); + defer s.deinit(gpa); + try engine.addStream(&s); + while (engine.activeCount() > 0) try engine.step(); + var out: Run = .{ .len = s.emitted().len, .holds = Cost.holds, .rounds = Cost.verifies, .drafted = s.drafted }; + @memcpy(out.tokens[0..out.len], s.emitted()); + return out; +} + +const slow: Costs = .{ .row = 8, .extra = 4, .hold = 9 }; // a hold and a verify row cost more than drafts save +const fast: Costs = .{ .row = 8, .extra = 1, .hold = 1 }; + +test "a lone stream whose drafted rounds lose runs plain rounds and probes now and then, replies unchanged" { + const plain = try run(null, true); + const unguarded = try run(slow, false); + const guarded = try run(slow, true); + try std.testing.expectEqualSlices(u32, plain.tokens[0..plain.len], unguarded.tokens[0..unguarded.len]); + try std.testing.expectEqualSlices(u32, plain.tokens[0..plain.len], guarded.tokens[0..guarded.len]); + try std.testing.expect(unguarded.holds * 10 > unguarded.rounds * 9); // without the guard nearly every round holds + try std.testing.expect(guarded.holds * 3 < guarded.rounds); // with it most rounds are plain + try std.testing.expect(guarded.holds > 16); // the first drafted stretch and probes +} + +test "a lone stream whose drafted rounds win keeps drafting with the guard on" { + const plain = try run(null, true); + const unguarded = try run(fast, false); + const guarded = try run(fast, true); + try std.testing.expectEqualSlices(u32, plain.tokens[0..plain.len], guarded.tokens[0..guarded.len]); + try std.testing.expect(guarded.holds * 10 > guarded.rounds * 8); // drafting but for the plain rate's one stretch + try std.testing.expect(guarded.rounds * 10 < unguarded.rounds * 12); +} From 1e718458ff686197a8eb70aa82217b2c587f3710 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:11:23 +0100 Subject: [PATCH 09/46] lanes: a copy window may be priced by its stream's own copies (Model.own_copy_rates) A shared round prices every copy window's rows at cfg.copy_rate, a fixed prior. A stream whose copies keep breaking (a different prompt on each stream, say) keeps paying verify rows for them. With Model.own_copy_rates a copy window is priced from the stream's own decayed copy history, by source (context or Recall), starting from four rows of copy_rate's prior; each verified copy window teaches its source's history. Off by default: without the flag allocation prices copies at copy_rate as before and the history is not kept. Test (copied_rate_test.zig): four streams' seven-row copy windows in one shared round, through windows.allocate. With the flag, after its context copies land one row of seven, a stream gets at most one copy row where the prior gave seven, the other streams keep theirs and its Recall copies keep copy_rate; without it every copy keeps seven rows and no history is kept. --- zig/src/core/lanes/config.zig | 3 ++ zig/src/core/lanes/copied_rate_test.zig | 56 +++++++++++++++++++++++++ zig/src/core/lanes/engine.zig | 4 +- zig/src/core/lanes/lanes.zig | 1 + zig/src/core/lanes/stream.zig | 2 + zig/src/core/lanes/windows.zig | 28 ++++++++++++- 6 files changed, 90 insertions(+), 4 deletions(-) create mode 100644 zig/src/core/lanes/copied_rate_test.zig diff --git a/zig/src/core/lanes/config.zig b/zig/src/core/lanes/config.zig index cb2159b03a..a6721fcf57 100644 --- a/zig/src/core/lanes/config.zig +++ b/zig/src/core/lanes/config.zig @@ -14,6 +14,7 @@ pub const default_prior = [_]f64{ 0.85, 0.75, 0.7, 0.65, 0.6, 0.55, 0.5, 0.5, 0. pub const Model = struct { root_copies: bool = false, // verified context copies precede head drafts and skip their proposal pass copy_drafts: bool = true, // global serving switch, independent of shared verification capacity + own_copy_rates: bool = false, // a copy window is priced by its stream's own copies, not copy_rate's prior exact_width: u32 = 1, // the widest window that reproduces one-token steps bit for bit first_copy_rows: u32 = 0, // a lone stream's first copy window (0: the exact width) @@ -52,6 +53,7 @@ pub const Config = struct { root_copies: bool = false, family_width: u32, max_copy: u32, + own_copy_rates: bool = false, // Model.own_copy_rates first_copy: u32, base_width: u32, family_mtp: bool, @@ -134,6 +136,7 @@ pub const Config = struct { .lone_cost_tile = m.lone_cost_tile, .family_width = width, .max_copy = max_copy, + .own_copy_rates = m.own_copy_rates, .first_copy = first_copy, .base_width = @min(width, first_copy + 1), .family_mtp = m.mtp and width >= 2 and m.speculate, diff --git a/zig/src/core/lanes/copied_rate_test.zig b/zig/src/core/lanes/copied_rate_test.zig new file mode 100644 index 0000000000..d037e3bf7c --- /dev/null +++ b/zig/src/core/lanes/copied_rate_test.zig @@ -0,0 +1,56 @@ +//! Model.own_copy_rates: a shared round's copy windows priced by each stream's own copies, else by copy_rate. +const std = @import("std"); +const Config = @import("config.zig").Config; +const Cost = @import("config.zig").Cost; +const Engine = @import("engine.zig").Engine; +const Stream = @import("stream.zig").Stream; +const win = @import("windows.zig"); +const fake = @import("fake.zig"); + +const gpa = std.testing.allocator; +const copied = [_]u32{ 5, 6, 7, 8, 9, 10, 11 }; + +/// The copy rows allocation grants each of four streams' seven-row copy windows in one shared round. +fn granted(e: *Engine, streams: []Stream, out: []usize) !void { + _ = e.arena.reset(.retain_capacity); + var plans: [4]win.Plan = undefined; + for (plans[0..streams.len], streams) |*p, *s| p.* = .{ .stream = s, .position = s.context.items.len, .kind = .copy, .tokens = &copied }; + try win.allocate(e, plans[0..streams.len]); + for (plans[0..streams.len], out) |p, *o| o.* = p.count(); +} + +test "own_copy_rates: copies that keep breaking are cut, by source; without it every copy is priced at copy_rate" { + // plain rows 10 ms, 0.6 ms a row past four: seven copy rows a stream pay at the prior, not landing 1 a copy + var costs: [32]Cost = undefined; + for (&costs, 1..) |*c, r| c.* = .{ .width = @intCast(r), .ms = 10 + 0.6 * @as(f64, @floatFromInt(@max(r, 4) - 4)) }; + for ([_]bool{ false, true }) |own| { + var cfg = try Config.init(gpa, .{ .exact_width = 16, .hidden_rows = true, .window_costs = &costs, .batch_rows = 32, .own_copy_rates = own }, 16, 15); + defer cfg.deinit(gpa); + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: fake.FixedClock = .{}; + var e = Engine.init(gpa, &cfg, target.backend(), clock.clock()); + defer e.deinit(); + var streams: [4]Stream = undefined; + var made: usize = 0; + defer for (streams[0..made]) |*s| s.deinit(gpa); + for (&streams) |*s| { + s.* = try Stream.init(gpa, .{ .id = "copied", .prompt = &.{ 1, 2 }, .max_new = 64 }); + made += 1; + } + var rows: [4]usize = undefined; + try granted(&e, &streams, &rows); + try std.testing.expectEqualSlices(usize, &.{ 7, 7, 7, 7 }, &rows); + for (0..8) |_| win.observeCopied(&e, &streams[0], false, 1, 7); // stream 0's context copies land one row of seven + try granted(&e, &streams, &rows); + if (own) { + try std.testing.expect(rows[0] <= 1); + try std.testing.expectEqualSlices(usize, &.{ 7, 7, 7 }, rows[1..]); + try std.testing.expect(win.copiedRate(&e, &streams[0], false) < 0.6); + try std.testing.expectApproxEqAbs(cfg.copy_rate, win.copiedRate(&e, &streams[0], true), 1e-12); // Recall's keep theirs + } else { + try std.testing.expectEqualSlices(usize, &.{ 7, 7, 7, 7 }, &rows); + try std.testing.expectEqual(@as(f64, -1), streams[0].copied_landed[0]); // no history kept + } + } +} diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index a92b27ae7f..f4902bdc7a 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -344,9 +344,7 @@ pub const Engine = struct { for (path[1..]) |r| own_kept += @intFromBool(r < own_end); const own: usize = own_end - 1; if (p.kind == .copy) { - if (p.recalled) { - if (s.recall) |r| r.observe(@intCast(own), @intCast(own_kept), s.rounds); - } else if (s.proposer) |pr| pr.observe(@intCast(own), @intCast(own_kept)); + win.observeCopied(e, s, p.recalled, own_kept, own); if (e.alone) { const width = s.copy_width orelse e.cfg.first_copy; s.copy_width = if (own_kept == own) @min(e.cfg.max_copy, 2 * width + 1) else e.cfg.first_copy; diff --git a/zig/src/core/lanes/lanes.zig b/zig/src/core/lanes/lanes.zig index ff09b58fb2..deb434e894 100644 --- a/zig/src/core/lanes/lanes.zig +++ b/zig/src/core/lanes/lanes.zig @@ -65,6 +65,7 @@ test { _ = @import("fill_test.zig"); _ = @import("fill_trim_test.zig"); _ = @import("fair_turns_test.zig"); + _ = @import("copied_rate_test.zig"); } pub const learned_disk = @import("learned_disk.zig"); diff --git a/zig/src/core/lanes/stream.zig b/zig/src/core/lanes/stream.zig index 8c348ff95f..4c7315bf7e 100644 --- a/zig/src/core/lanes/stream.zig +++ b/zig/src/core/lanes/stream.zig @@ -200,6 +200,8 @@ pub const Stream = struct { prices_seeded: bool = false, // its own prices started from the engine's kept calibration copy_landed: [2]f64 = .{ -1, -1 }, // a lone context copy's rows that landed, decayed, by short and long match (-1: unprimed) copy_broke: [2]f64 = .{ 0, 0 }, // and its chains that broke before their end: their ratio is a copied row's landing rate + copied_landed: [2]f64 = .{ -1, -1 }, // a copy window's rows that landed, decayed, by source: context, Recall (-1: unprimed) + copied_broke: [2]f64 = .{ 0, 0 }, // and its copies that broke before their end drafted: u64 = 0, accepted: u64 = 0, branch_rows: u64 = 0, // suffix-match branch rows verified (lanes/fill.zig) diff --git a/zig/src/core/lanes/windows.zig b/zig/src/core/lanes/windows.zig index e8e48be2ad..fd7ffec95d 100644 --- a/zig/src/core/lanes/windows.zig +++ b/zig/src/core/lanes/windows.zig @@ -248,7 +248,8 @@ pub fn allocate(e: *Engine, plans: []Plan) !void { pr.* = &.{}; } else if (p.kind == .copy) { fx.* = 1; - pr.* = try alloc.chainProbabilities(a, &.{e.cfg.copy_rate}, n); + const rate = if (e.cfg.own_copy_rates) copiedRate(e, p.stream, p.recalled) else e.cfg.copy_rate; + pr.* = try alloc.chainProbabilities(a, &.{rate}, n); } else { fx.* = 1; const chances = try planChances(e, p, n); @@ -362,6 +363,31 @@ pub fn observeCopy(e: *const Engine, s: *Stream, trusted: bool, kept: u64, offer } } +/// A copy window's history weighs this many rows of cfg.copy_rate's prior before the stream's own copies teach it. +const copied_prior = 4; + +/// A copy window's row landing rate for this stream and source: landed rows over those plus breaks, decayed a copy. +pub fn copiedRate(e: *const Engine, s: *Stream, recalled: bool) f64 { + const i = @intFromBool(recalled); + if (s.copied_landed[i] < 0) { + s.copied_landed[i] = copied_prior * e.cfg.copy_rate; + s.copied_broke[i] = copied_prior * (1 - e.cfg.copy_rate); + } + return s.copied_landed[i] / (s.copied_landed[i] + s.copied_broke[i]); +} + +/// A verified copy window of `offered` rows, `kept` landing: its source's counters, then its own rate's history. +pub fn observeCopied(e: *const Engine, s: *Stream, recalled: bool, kept: usize, offered: usize) void { + if (recalled) { + if (s.recall) |r| r.observe(@intCast(offered), @intCast(kept), s.rounds); + } else if (s.proposer) |pr| pr.observe(@intCast(offered), @intCast(kept)); + if (!e.cfg.own_copy_rates) return; + _ = copiedRate(e, s, recalled); + const i = @intFromBool(recalled); + s.copied_landed[i] = copy_decay * s.copied_landed[i] + @as(f64, @floatFromInt(kept)); + s.copied_broke[i] = copy_decay * s.copied_broke[i] + @as(f64, @floatFromInt(@intFromBool(kept < offered))); +} + /// A context copy's appended rows: each its parent's chance times the stream's copy rate; the drafts keep their own. fn planChances(e: *Engine, p: Plan, n: usize) ![]const f64 { if (!p.context_copy or p.branch_from == 0 or p.parents == null or n < p.branch_from) return nodeChances(e, p.stream, n); 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 10/46] 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 11/46] 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 12/46] 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 13/46] 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 14/46] 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 15/46] 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 b18058e920d4543497c3e316dfeb5d53c3738047 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:30:40 +0100 Subject: [PATCH 16/46] lfm2: draft with LiquidAI's DSpark drafter (--drafter) LiquidAI/LFM2.5-24B-A2B-DSpark is LiquidAI's speculative-decoding drafter for LFM2.5-24B-A2B. `tensorfold serve` now takes it with --drafter DIR for lfm2_moe checkpoints. The lane core verifies every draft, so a drafted reply equals its "draft": false reply. The drafter reads the target's rows after layers 2, 10, 18, 26 and 34 (the taps). fc and hidden_norm turn each row into the context keys and values of five Qwen3 layers. A block of the pending token and mask tokens runs the layers with non-causal attention, then the target's head. The Markov head adds W2 W1[previous token] to each row's logits and takes the argmax, one row after the other. A block is as long as the deepest chain its pass holds, at most the nine it was trained on. The drafter runs on the family's own kernels and core/tile_qmm (on GPUs with tensor units, core/affine_mm for absorb passes past the twin's rows): its bf16 linears are quantized to 8-bit MLX affine codes at load, and the q and k rows are reordered so that the target's half-split RoPE kernel gives the drafter's interleaved RoPE. New kernels: - lfm2_markov_scan and lfm2_markov_pick (lfm2_dspark.metal): a Markov step's W1 gather, W2 dot over 8-bit codes, bias add and argmax, ties to the lowest token, a NaN score flagged; - lfm2_attn_block and lfm2_attn_block_merge: the target's attention and merge kernels with a FULL template flag, every block row seeing the keys [0, len + 1). The target's own instantiations keep their causal count. Both register with the kernel contract (metal_lfm2_dspark.zig, metal_lfm2_attn.zig), gate 5 against every Apple chip we serve and this GPU. Rows reach the drafter in passes: a prompt chunk in one pass of up to 2,048 rows; a round's verified rows are staged on the host (128 a lane) and absorbed by the next hold, so a round without holds does no drafter GPU work. A hold pass takes as many streams as their blocks fit in 32 rows. The drafter sets Facts.shared_guard, so shared rounds draft only while drafted rounds measure faster than plain ones at that stream count; a lone stream keeps the depth rule's own plain guard. The drafter keeps its context with each kept prompt state, so a drafting stream resumes a prompt as before. Its lane's bytes join the streams the load limit admits. --drafter-bits 0 is refused for this drafter, and a refused open frees the model it loaded. Tests: test-lfm2-dspark checks the config's admission and refusals, the quantizer, the RoPE reordering through lfm2_qk, the Markov step and the block attention on the GPU. tf-lfm2-dspark runs the real drafter against an fp32 host reference of the released checkpoint and times its passes. --- zig/build/metal_targets.zig | 3 +- zig/kernels/metal/lfm2/lfm2.metal | 19 +- zig/kernels/metal/lfm2/lfm2_dspark.metal | 99 +++++ zig/kernels/metal/sources.zig | 2 +- zig/src/families/lfm2_moe/dspark.zig | 386 +++++++++++++++++++ zig/src/families/lfm2_moe/dspark_block.zig | 271 +++++++++++++ zig/src/families/lfm2_moe/dspark_config.zig | 136 +++++++ zig/src/families/lfm2_moe/dspark_quant.zig | 118 ++++++ zig/src/families/lfm2_moe/dspark_test.zig | 147 ++++++- zig/src/families/lfm2_moe/dspark_weights.zig | 144 +++++++ zig/src/families/lfm2_moe/kernels.zig | 4 + zig/src/families/lfm2_moe/lfm2_moe.zig | 4 + zig/src/native/lfm2_moe.zig | 80 +++- zig/src/native/metal.zig | 4 +- zig/src/tensorfold.zig | 2 + zig/tests/contract/metal_lfm2_attn.zig | 182 +++++++++ zig/tests/contract/metal_lfm2_dspark.zig | 221 +++++++++++ zig/tests/contract/metal_main.zig | 19 + zig/tests/lfm2_dspark.zig | 301 +++++++++++++++ 19 files changed, 2114 insertions(+), 28 deletions(-) create mode 100644 zig/kernels/metal/lfm2/lfm2_dspark.metal create mode 100644 zig/src/families/lfm2_moe/dspark.zig create mode 100644 zig/src/families/lfm2_moe/dspark_block.zig create mode 100644 zig/src/families/lfm2_moe/dspark_config.zig create mode 100644 zig/src/families/lfm2_moe/dspark_quant.zig create mode 100644 zig/src/families/lfm2_moe/dspark_weights.zig create mode 100644 zig/tests/contract/metal_lfm2_attn.zig create mode 100644 zig/tests/contract/metal_lfm2_dspark.zig create mode 100644 zig/tests/lfm2_dspark.zig diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index 6db7154689..727daf7e2e 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -292,6 +292,7 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s .{ .name = "tf-lfm2-forward", .path = "zig/tests/lfm2_forward.zig", .about = "LFM2-MoE native teacher-forced logits" }, .{ .name = "tf-lfm2-bench", .path = "zig/tests/lfm2_bench.zig", .about = "LFM2-MoE prompt-chunk and one-row decode timing" }, .{ .name = "tf-lfm2-exact", .path = "zig/tests/lfm2_exact.zig", .about = "LFM2-MoE prompt chunks, windows and committed-state exactness" }, + .{ .name = "tf-lfm2-dspark", .path = "zig/tests/lfm2_dspark.zig", .about = "LFM2's DSpark drafter on the GPU against an fp32 host reference of the released drafter" }, .{ .name = "tf-lfm2-fp64", .path = "zig/tests/lfm2_fp64.zig", .about = "LFM2 projection classes and attention on the tensor units, the tiles and the twin, each against float64" }, .{ .name = "tensorfold", .path = "zig/src/main.zig", .about = "The native engine's command line" }, .{ .name = "tf-nemotron-fixtures", .path = "zig/tests/nemotron_fixtures.zig", .about = "Nemotron kernels against the Python engine's captured ops" }, @@ -458,7 +459,7 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s 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); + b.step("test-lfm2-dspark", "LFM2's DSpark drafter and its taps: config, quantizer and kernels 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 b2772cf64e..76c223805d 100644 --- a/zig/kernels/metal/lfm2/lfm2.metal +++ b/zig/kernels/metal/lfm2/lfm2.metal @@ -232,7 +232,7 @@ 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 +template // RB row blocks share each staged K and V fragment; FULL: all see [0, len + 1) 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)], @@ -242,8 +242,8 @@ inline void lfm2_attn_rows(const device bfloat* queries, const device bfloat* ke const uint kh = tg.y, r0 = tg.x * 8 * RB, head = kh * (QH / KVH) + sg; const uint qid = lane / 4, fm = (qid & 4) + ((lane / 2) % 4), fn = (qid & 2) * 2 + (lane % 2) * 2; uint count[RB]; // this lane's rows' visible keys - for (int b = 0; b < RB; ++b) count[b] = p.len + min(r0 + 8 * b + fm, rows - 1) + 1; - const uint top = p.len + min(r0 + 8 * RB - 1, rows - 1) + 1; // the tile's most + for (int b = 0; b < RB; ++b) count[b] = FULL ? p.len + 1 : p.len + min(r0 + 8 * b + fm, rows - 1) + 1; + const uint top = FULL ? p.len + 1 : p.len + min(r0 + 8 * RB - 1, rows - 1) + 1; // the tile's most if (chunk * PER + sub0 * SUB >= top) return; vec q[RB][8]; for (int b = 0; b < RB; ++b) @@ -340,7 +340,7 @@ inline void lfm2_attn_rows(const device bfloat* queries, const device bfloat* ke for (int d = 0; d < 8; ++d) *(device float2*)(po + at * HD + 8 * d + fn) = O[b][d]; } } -template +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)]], @@ -349,15 +349,17 @@ template 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); + 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>; +template [[host_name("lfm2_attn_block")]] [[kernel]] decltype(lfm2_attn_mma<1, true>) lfm2_attn_mma<1, true>; // The chunks merged in chunk order for (head, row): Y [rows, QH*HD] bf16. +template 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 count = FULL ? p.len + 1 : p.len + r + 1; const uint per = PER; const uint chunks = (count + per - 1) / per; const size_t at = (size_t(r) * QH + head) * (CHUNKS * p.subs); @@ -391,12 +393,15 @@ inline void lfm2_merge_row(const device float* pm, const device float* pl, const out[lane] = static_cast(o0 / l); out[lane + 32] = static_cast(o1 / l); } +template [[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); + lfm2_merge_row(pm, pl, po, p, y, tg.x, tg.y, lane); } +template [[host_name("lfm2_attn_merge")]] [[kernel]] decltype(lfm2_attn_merge) lfm2_attn_merge; +template [[host_name("lfm2_attn_block_merge")]] [[kernel]] decltype(lfm2_attn_merge) lfm2_attn_merge; // 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)]], diff --git a/zig/kernels/metal/lfm2/lfm2_dspark.metal b/zig/kernels/metal/lfm2/lfm2_dspark.metal new file mode 100644 index 0000000000..995d63a20a --- /dev/null +++ b/zig/kernels/metal/lfm2/lfm2_dspark.metal @@ -0,0 +1,99 @@ +// DSpark's Markov step, compiled after lfm2.metal's text (its includes and constants). + +// DSpark's Markov scan: a threadgroup's best of 512 tokens of a row, bf16(base + bf16(W2 . W1[prev])), W2 8-bit codes. +struct lfm2_markov_args { uint stride, rank, vocab, prev_stride; }; +constant constexpr uint MARKOV_SPAN = 512; +[[kernel]] void lfm2_markov_scan(const device bfloat* logits [[buffer(0)]], const device bfloat* w1 [[buffer(1)]], + const device uint8_t* w2 [[buffer(2)]], const device uint* prev [[buffer(3)]], + constant lfm2_markov_args& p [[buffer(4)]], device float* best [[buffer(5)]], + device uint* at [[buffer(6)]], const device bfloat* w2s [[buffer(7)]], + const device bfloat* w2b [[buffer(8)]], uint2 tg [[threadgroup_position_in_grid]], + uint2 grid [[threadgroups_per_grid]], uint sg [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + threadgroup float e[256]; + threadgroup float sv[8]; + threadgroup uint si[8]; + const uint s = tg.y, t = sg * 32 + lane, id = prev[size_t(s) * p.prev_stride]; + for (uint i = t; i < p.rank; i += 256) e[i] = id < p.vocab ? float(w1[size_t(id) * p.rank + i]) : 0.0f; + threadgroup_barrier(mem_flags::mem_threadgroup); + const device bfloat* row = logits + size_t(s) * p.stride; + float mv = -INFINITY; + uint mi = 0xffffffffu; + bool nan = false; + const uint first = tg.x * MARKOV_SPAN, last = min(first + MARKOV_SPAN, p.vocab); + for (uint v = first + sg; v < last; v += 8) { + float d = 0.0f; + for (uint i = 8 * lane; i < p.rank; i += 256) { + const size_t g = (size_t(v) * p.rank + i) / 64; + const float sc = w2s[g], bi = w2b[g]; + const uchar4 lo = *(const device uchar4*)(w2 + size_t(v) * p.rank + i); + const uchar4 hi = *(const device uchar4*)(w2 + size_t(v) * p.rank + i + 4); + for (int j = 0; j < 4; ++j) d = fma(fma(sc, float(lo[j]), bi), e[i + j], d); + for (int j = 0; j < 4; ++j) d = fma(fma(sc, float(hi[j]), bi), e[i + 4 + j], d); + } + d = simd_sum(d); + const float sc = float(static_cast(float(row[v]) + float(static_cast(d)))); + nan = nan || isnan(sc); + if (sc > mv || (mi == 0xffffffffu && sc == mv)) { + mv = sc; + mi = v; + } + } + if (lane == 0) { + sv[sg] = nan ? NAN : mv; + si[sg] = mi; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (t != 0) return; + float bv = sv[0]; + uint bi = si[0]; + for (uint g = 1; g < 8; ++g) { + if (isnan(bv) || isnan(sv[g])) bv = NAN; + else if (sv[g] > bv || (sv[g] == bv && si[g] < bi)) { + bv = sv[g]; + bi = si[g]; + } + } + best[size_t(s) * grid.x + tg.x] = bv; + at[size_t(s) * grid.x + tg.x] = bi; +} + +// DSpark's Markov step pick: the scan's best over a stream's threadgroups (ties to the lowest token), as a Pick. +[[kernel]] void lfm2_markov_pick(const device float* best [[buffer(0)]], const device uint* at [[buffer(1)]], + constant uint& groups [[buffer(2)]], device uint* picks [[buffer(3)]], + uint s [[threadgroup_position_in_grid]], uint t [[thread_index_in_threadgroup]]) { + threadgroup float tv[256]; + threadgroup uint ti[256]; + threadgroup uint tn[256]; + float bv = -INFINITY; + uint bi = 0xffffffffu, nan = 0; + for (uint g = t; g < groups; g += 256) { + const float v = best[size_t(s) * groups + g]; + const uint i = at[size_t(s) * groups + g]; + if (isnan(v)) nan = 1; + else if (v > bv || (v == bv && i < bi)) { + bv = v; + bi = i; + } + } + tv[t] = bv; + ti[t] = bi; + tn[t] = nan; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint w = 128; w > 0; w >>= 1) { + if (t < w) { + const float v = tv[t + w]; + const uint i = ti[t + w]; + if (v > tv[t] || (v == tv[t] && i < ti[t])) { + tv[t] = v; + ti[t] = i; + } + tn[t] |= tn[t + w]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (t == 0) { + picks[2 * s] = ti[0]; + picks[2 * s + 1] = tn[0]; + } +} diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index 18662768d7..6d350f26b6 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") ++ @embedFile("lfm2/lfm2_segments.metal"); +pub const lfm2 = @embedFile("lfm2/lfm2.metal") ++ @embedFile("lfm2/lfm2_dspark.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/dspark.zig b/zig/src/families/lfm2_moe/dspark.zig new file mode 100644 index 0000000000..339071ac6c --- /dev/null +++ b/zig/src/families/lfm2_moe/dspark.zig @@ -0,0 +1,386 @@ +//! LiquidAI's DSpark drafter on LFM2's taps: a stream's context keys, a block of rows, a Markov chain of drafts. +const std = @import("std"); +const mtl = @import("metal"); +const lanes = @import("lanes"); +const c = @import("config.zig"); +const dc = @import("dspark_config.zig"); +const dw = @import("dspark_weights.zig"); +const block = @import("dspark_block.zig"); +const st = @import("state.zig"); +const snapshot = @import("snapshot.zig"); +const Model = @import("model.zig").Model; + +const dr = lanes.drafter; +const Stream = lanes.Stream; +const shared = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; + +/// Block rows one hold pass takes (the decode twin's rows), and context rows one absorb pass takes. +pub const pass_rows = 32; +pub const absorb_rows = 2048; +/// Verified rows a lane keeps on the host until its next hold absorbs them (a round's rows for no GPU work). +pub const stage_rows = 128; +const kw = c.kv_heads * c.head_dim; + +/// `bits`: the width the bf16 linears are quantized to (8 keeps the fp32 reference's drafts; 4 does not). +pub const Options = struct { capacity: usize, bits: u8 = 8 }; + +/// One stream's context keys and values a layer ([kv head][cap][head dim], as the target's attention reads them). +pub const Lane = struct { + keys: [8]mtl.Buffer, + values: [8]mtl.Buffer, + cap: usize, // positions each holds: the context and one block past it + end: u64 = 0, // context rows absorbed: positions [0, end) + stage: mtl.Buffer, // tapped rows [end, end + staged), not absorbed yet (shared storage) + staged: u32 = 0, + held: [16]u32 = @splat(0), + held_len: u32 = 0, +}; + +/// The passes' working buffers, one pass at a time on the target's queue. +pub const Scratch = struct { + ids: mtl.Buffer, + h: mtl.Buffer, + x: mtl.Buffer, + r: mtl.Buffer, + y: mtl.Buffer, + q: mtl.Buffer, + k: mtl.Buffer, + v: mtl.Buffer, + queries: mtl.Buffer, + pm: mtl.Buffer, + pl: mtl.Buffer, + po: mtl.Buffer, + gate: mtl.Buffer, + up: mtl.Buffer, + act: mtl.Buffer, + logits: mtl.Buffer, + sums: mtl.Buffer, + fcx: mtl.Buffer, + hctx: mtl.Buffer, + best: mtl.Buffer, // the Markov scan's best score and token of each threadgroup, by stream + at: mtl.Buffer, + picks: mtl.Buffer, +}; + +pub const Drafter = struct { + gpa: std.mem.Allocator, + io: std.Io, + target: *Model, + config: dc.Config, + w: dw.Weights, + options: Options, + s: Scratch, + owned: std.ArrayList(mtl.Buffer) = .empty, + lanes: std.AutoHashMapUnmanaged(*Stream, *Lane) = .empty, + /// One stream's hold wall ms at full depth, and that over the block (the depth rule's cost a draft). + hold_ms: f64 = 0, + step_ms: f64 = 0, + + pub fn init(gpa: std.mem.Allocator, io: std.Io, target: *Model, dir: []const u8, options: Options) !*Drafter { + const d = try gpa.create(Drafter); + errdefer gpa.destroy(d); + d.* = .{ .gpa = gpa, .io = io, .target = target, .config = try dc.Config.read(gpa, io, dir), .w = undefined, .options = options, .s = undefined }; + d.w = try dw.load(gpa, target.device, io, dir, &d.config, options.bits); + errdefer d.w.deinit(gpa); + errdefer for (d.owned.items) |b| b.deinit(); + errdefer d.owned.deinit(gpa); + try d.scratch(); + try d.warm(); + return d; + } + + pub fn deinit(d: *Drafter) void { + var it = d.lanes.valueIterator(); + while (it.next()) |l| d.dropLane(l.*); + d.lanes.deinit(d.gpa); + for (d.owned.items) |b| b.deinit(); + d.owned.deinit(d.gpa); + d.w.deinit(d.gpa); + d.gpa.destroy(d); + } + + fn buffer(d: *Drafter, bytes: usize) !mtl.Buffer { + const b = try d.target.device.buffer(@max(bytes, 16), shared); + errdefer b.deinit(); + try d.owned.append(d.gpa, b); + return b; + } + + fn scratch(d: *Drafter) !void { + const r = pass_rows; + const wide = @max(pass_rows, absorb_rows); + const i = d.config.intermediate; + const partials = @max(st.split_rows, d.config.block) * st.subs * c.query_heads * st.chunks; + const sums = @max(absorb_rows * (d.config.tap_count * c.hidden / 64), r * (i / 64)); + d.s = .{ + .ids = try d.buffer(r * 4), + .h = try d.buffer(r * c.hidden * 2), + .x = try d.buffer(wide * c.hidden * 2), + .r = try d.buffer(r * c.hidden * 2), + .y = try d.buffer(r * c.hidden * 2), + .q = try d.buffer(r * c.hidden * 2), + .k = try d.buffer(wide * kw * 2), + .v = try d.buffer(wide * kw * 2), + .queries = try d.buffer((r + st.attn_rows) * c.hidden * 2), + .pm = try d.buffer(partials * 4), + .pl = try d.buffer(partials * 4), + .po = try d.buffer(partials * c.head_dim * 4), + .gate = try d.buffer(r * i * 2), + .up = try d.buffer(r * i * 2), + .act = try d.buffer(r * i * 2), + .logits = try d.buffer(r * c.vocab * 2), + .sums = try d.buffer(sums * 4), + .fcx = try d.buffer(absorb_rows * c.hidden * 2), + .hctx = try d.buffer(absorb_rows * c.hidden * 2), + .best = try d.buffer(r * block.scan_groups * 4), + .at = try d.buffer(r * block.scan_groups * 4), + .picks = try d.buffer(d.config.block * r * 8), + }; + } + + /// The drafter's buffers that every round reads (weights and scratch), for the server's residency set. + pub fn buffers(d: *const Drafter, gpa: std.mem.Allocator) ![]mtl.Buffer { + return std.mem.concat(gpa, mtl.Buffer, &.{ d.w.buffers(), d.owned.items }); + } + + pub fn drafter(d: *Drafter) dr.Drafter { + return .{ .ptr = d, .vtable = &.{ .taps = tapsFn, .facts = factsFn, .open = openFn, .absorb = absorbFn, .hold = holdFn, .held = heldFn, .release = releaseFn } }; + } + + fn cast(ptr: *anyopaque) *Drafter { + return @ptrCast(@alignCast(ptr)); + } + + fn tapsFn(ptr: *anyopaque) []const u32 { + const d = cast(ptr); + return d.config.taps[0..d.config.tap_count]; + } + + /// Acceptance a depth until a stream has its own: the card's 5.16 tokens a step at block 9 is 0.84 a depth. + const prior: [16]f64 = @splat(0.84); + + fn factsFn(ptr: *anyopaque) dr.Facts { + const d = cast(ptr); + return .{ .depth = d.config.block, .step_ms = d.step_ms, .prior = prior[0..d.config.block], .batched = true, .shared_guard = true, .stream_bytes = d.laneBytes(d.options.capacity) }; + } + + pub fn laneBytes(d: *const Drafter, capacity: usize) usize { + return 2 * d.config.layers * (capacity + d.config.block) * kw * 2 + stage_rows * d.tapBytes(); + } + + /// A tapped row's bytes: every tap's hidden row, bf16. + pub fn tapBytes(d: *const Drafter) usize { + return d.config.tap_count * c.hidden * 2; + } + + /// The lane's staged rows as an absorb piece (null: none staged). + pub fn stagedPiece(l: *Lane) ?block.Piece { + if (l.staged == 0) return null; + return .{ .lane = l, .src = .{ .buf = l.stage }, .rows = l.staged, .p0 = l.end }; + } + + /// The lane's staged rows absorbed now, in a command buffer of their own. + fn flush(d: *Drafter, l: *Lane) !void { + const pc = stagedPiece(l) orelse return; + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const run = block.begin(d); + try block.absorbPass(d, run.e, &.{pc}); + try block.end(run.cb, run.e); + l.end += l.staged; + l.staged = 0; + } + + fn openFn(ptr: *anyopaque, s: *Stream) anyerror!void { + const d = cast(ptr); + if (d.lanes.contains(s)) return error.DraftStreamOpen; + const want = s.prompt_len + s.max_new + 2 * d.config.block; + const l = try d.newLane(@min(d.options.capacity, (want + 63) / 64 * 64)); + errdefer d.dropLane(l); + try d.lanes.put(d.gpa, s, l); + } + + pub fn newLane(d: *Drafter, capacity: usize) !*Lane { + const l = try d.gpa.create(Lane); + errdefer d.gpa.destroy(l); + l.* = .{ .keys = undefined, .values = undefined, .stage = try d.target.device.buffer(stage_rows * d.tapBytes(), shared), .cap = capacity + d.config.block }; + errdefer l.stage.deinit(); + var made: usize = 0; + errdefer for (0..made) |i| { + l.keys[i].deinit(); + l.values[i].deinit(); + }; + for (0..d.config.layers) |i| { + l.keys[i] = try d.target.device.buffer(l.cap * kw * 2, shared); + l.values[i] = d.target.device.buffer(l.cap * kw * 2, shared) catch |e| { + l.keys[i].deinit(); + return e; + }; + made += 1; + } + return l; + } + + pub fn dropLane(d: *Drafter, l: *Lane) void { + l.stage.deinit(); + for (0..d.config.layers) |i| { + l.keys[i].deinit(); + l.values[i].deinit(); + } + d.gpa.destroy(l); + } + + fn releaseFn(ptr: *anyopaque, s: *Stream) void { + const d = cast(ptr); + if (d.lanes.fetchRemove(s)) |e| d.dropLane(e.value); + } + + fn heldFn(ptr: *anyopaque, streams: []const *Stream, out: []const []u32) anyerror!void { + const d = cast(ptr); + for (streams, out) |s, o| { + const l = d.lanes.get(s) orelse return error.DraftStreamUnknown; + if (o.len > l.held_len) return error.DraftHeldFewer; + @memcpy(o, l.held[0..o.len]); + } + } + + /// Bytes a lane's context at `at` takes beside a kept prompt state. + pub fn keptBytes(d: *const Drafter, at: usize) u64 { + return 2 * @as(u64, d.config.layers) * at * kw * 2; + } + + /// Stream `s`'s context copied into `snap` when its lane stands at the snapshot's position (a prompt chunk's end). + pub fn keepLane(d: *Drafter, s: *Stream, snap: *snapshot.Snap) !void { + const l = d.lanes.get(s) orelse return; + try d.flush(l); + if (l.end != snap.at) return; + const bytes = try d.gpa.alloc(u8, @intCast(d.keptBytes(snap.at))); + copyLane(d, l, snap.at, bytes, true); + snap.draft = bytes; + } + + /// The lane takes the context kept with the state its stream's prompt pass resumed at `start`. + fn resumeLane(d: *Drafter, s: *Stream, l: *Lane, start: u64) !void { + const saved = s.reuse.saved orelse return error.DraftContextGap; + const snap: *const snapshot.Snap = @ptrCast(@alignCast(saved)); + const bytes = snap.draft orelse return error.DraftContextGap; + if (s.cached != start or snap.at != start or l.end != 0 or start + d.config.block > l.cap) return error.DraftContextGap; + copyLane(d, l, snap.at, bytes, false); + l.end = start; + } + + /// Rows [0, at) of every layer's keys and values and `bytes` (`out`: the lane's into bytes), kv head by kv head. + fn copyLane(d: *const Drafter, l: *const Lane, at: usize, bytes: []u8, out: bool) void { + const n = at * c.head_dim * 2; + var off: usize = 0; + for (0..d.config.layers) |i| for ([_]mtl.Buffer{ l.keys[i], l.values[i] }) |buf| for (0..c.kv_heads) |h| { + const live = (buf.contents() + h * l.cap * c.head_dim * 2)[0..n]; + if (out) @memcpy(bytes[off..][0..n], live) else @memcpy(live, bytes[off..][0..n]); + off += n; + }; + } + + /// Rows from `start` replace the context there: verified rows wait on the host for the next hold, others run now. + fn absorbFn(ptr: *anyopaque, items: []const dr.Absorb) anyerror!void { + const d = cast(ptr); + const width = d.tapBytes(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + var run: ?@TypeOf(block.begin(d)) = null; + var pieces: [32]block.Piece = undefined; + var count: usize = 0; + var rows: usize = 0; + for (items) |it| { + const l = d.lanes.get(it.stream) orelse return error.DraftStreamUnknown; + const f = it.features; + if (f.space != .device or f.dtype != .bf16 or f.row_bytes != width or it.rows != null) return error.DraftFeaturesUnsupported; + if (it.start > l.end + l.staged) try d.resumeLane(it.stream, l, it.start); // a prompt pass that resumed + if (it.start + f.rows + d.config.block > l.cap) return error.DraftContextGap; + l.held_len = 0; + const src = mtl.Buffer{ .id = @ptrFromInt(f.buffer) }; + // rows from `start` replace what was staged past it; rows before the absorbed end go to the GPU now + const keep: u32 = if (it.start >= l.end) @intCast(it.start - l.end) else 0; + l.staged = @min(l.staged, keep); + if (!it.prompt and it.start >= l.end and l.staged + f.rows <= stage_rows) { + const from = (src.contents() + f.offset)[0 .. f.rows * width]; + @memcpy((l.stage.contents() + l.staged * width)[0 .. f.rows * width], from); + l.staged += f.rows; + continue; + } + if (run == null) run = block.begin(d); + const e = run.?.e; + if (stagedPiece(l)) |pc| { + if (count == pieces.len or rows + pc.rows > absorb_rows) { + try block.absorbPass(d, e, pieces[0..count]); + count = 0; + rows = 0; + } + pieces[count] = pc; + count += 1; + rows += pc.rows; + } + var at: usize = 0; + while (at < f.rows) { + if (rows == absorb_rows or count == pieces.len) { // the pass is full: encode it + try block.absorbPass(d, e, pieces[0..count]); + count = 0; + rows = 0; + } + const n = @min(absorb_rows - rows, f.rows - at); + pieces[count] = .{ .lane = l, .src = .{ .buf = src, .off = f.offset + at * width }, .rows = n, .p0 = it.start + at }; + count += 1; + rows += n; + at += n; + } + l.end = it.start + f.rows; + l.staged = 0; + } + const r = run orelse return; + if (count > 0) try block.absorbPass(d, r.e, pieces[0..count]); + try block.end(r.cb, r.e); + } + + /// Each stream's chain of `depth` drafts after `pending` at `position`, its block's rows in a pass of pass_rows. + fn holdFn(ptr: *anyopaque, items: []const dr.Hold) anyerror!void { + const d = cast(ptr); + var at: usize = 0; + while (at < items.len) { + // as many streams as their blocks' rows fit a pass: each block is as long as the pass's deepest chain + var deepest: usize = 0; + var end = at; + while (end < items.len) : (end += 1) { + const rows = @max(deepest, items[end].depth); + if ((end - at + 1) * rows > pass_rows) break; + deepest = rows; + } + if (end == at) return error.DraftDepth; + var blocks: [pass_rows]block.Block = undefined; + const group = items[at..end]; + at = end; + for (group, blocks[0..group.len]) |it, *b| { + const l = d.lanes.get(it.stream) orelse return error.DraftStreamUnknown; + if (it.depth == 0 or it.depth > d.config.block or it.position == 0) return error.DraftDepth; + const anchor = it.position - 1; + if (l.end + l.staged != anchor or anchor + d.config.block > l.cap) return error.DraftContextGap; + b.* = .{ .lane = l, .pending = it.pending, .anchor = anchor, .depth = it.depth }; + } + try block.hold(d, blocks[0..group.len]); + } + } + + /// Holds on an empty context: the pipelines compiled, then one-draft and full-depth holds timed (best of three). + fn warm(d: *Drafter) !void { + const l = try d.newLane(64); + defer d.dropLane(l); + var best: [2]f64 = @splat(std.math.floatMax(f64)); + for (0..4) |i| for ([_]u32{ 1, d.config.block }, &best) |depth, *b| { + const t0 = std.Io.Clock.awake.now(d.io); + try block.hold(d, &.{.{ .lane = l, .pending = 1, .anchor = 0, .depth = depth }}); + const ms = @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(d.io)).toNanoseconds())) / 1e6; + if (i > 0) b.* = @min(b.*, ms); + }; + d.hold_ms = best[1]; + d.step_ms = best[1] / @as(f64, @floatFromInt(d.config.block)); + std.log.info("lfm2 dspark: a 1-draft hold takes {d:.2} ms, a {d}-draft hold {d:.2} ms", .{ best[0], d.config.block, best[1] }); + } +}; diff --git a/zig/src/families/lfm2_moe/dspark_block.zig b/zig/src/families/lfm2_moe/dspark_block.zig new file mode 100644 index 0000000000..22676dca6e --- /dev/null +++ b/zig/src/families/lfm2_moe/dspark_block.zig @@ -0,0 +1,271 @@ +//! DSpark's passes on the target's queue: context rows into a lane's keys, and a block's rows into Markov drafts. +const std = @import("std"); +const mtl = @import("metal"); +const c = @import("config.zig"); +const st = @import("state.zig"); +const fwd = @import("forward.zig"); +const tile = @import("../../core/tile_qmm.zig"); +const am = @import("../../core/affine_mm.zig"); +const round = @import("tree_round_gpu"); +const ds = @import("dspark.zig"); +const Drafter = ds.Drafter; +const Lane = ds.Lane; + +const kw = c.kv_heads * c.head_dim; +const QkArgs = extern struct { len: u32, cap: u32, eps: f32, log2_theta: f32 }; +const AttnArgs = extern struct { len: u32, cap: u32, scale: f32, subs: u32 }; +const Pick = round.contract.Pick; +const MarkovArgs = extern struct { stride: u32, rank: u32, vocab: u32, prev_stride: u32 }; +/// Tokens a Markov scan threadgroup takes (the kernel's MARKOV_SPAN), and threadgroups a stream's row takes. +pub const scan_span = 512; +pub const scan_groups = (c.vocab + scan_span - 1) / scan_span; + +/// A buffer and byte offset. +pub const Ref = struct { buf: mtl.Buffer, off: usize = 0 }; + +/// One stream's block in a hold pass: rows [pending, mask x (block - 1)] at positions anchor .. anchor + block - 1. +pub const Block = struct { lane: *Lane, pending: u32, anchor: u64, depth: u32 }; + +const Enc = struct { + d: *Drafter, + e: mtl.ComputeEncoder, + + fn pipe(x: Enc, comptime key: []const u8) void { + x.e.setPipeline(x.d.target.kernels.get(key)); + } + + /// y[rows] = x[rows] W^T: the twin up to its rows, past them the M5's tensor-unit kernels, else the tiles. + fn project(x: Enc, w: tile.Matrix, in: Ref, out: Ref, rows: usize) !void { + const t = &x.d.target.kernels.tiles; + if (rows > tile.twin_rows) if (x.d.target.kernels.naxPipes(w.bits)) |p| { + const sums = Ref{ .buf = x.d.s.sums }; + am.rowSums(x.e, p, c.group, in, sums, @intCast(rows), @intCast(w.k)); + const q = .{ .w = Ref{ .buf = w.w, .off = w.w_off }, .s = Ref{ .buf = w.scales, .off = w.s_off }, .b = Ref{ .buf = w.biases, .off = w.b_off }, .n = @as(u32, @intCast(w.n)), .k = @as(u32, @intCast(w.k)) }; + return am.dense(x.e, p, in, sums, q, out, @intCast(rows)); + }; + if (rows <= tile.twin_rows and x.d.target.kernels.twin) return tile.encodeTwin(x.e, t, w, in.buf, in.off, out.buf, out.off, rows); + tile.encodeSums(x.e, t, w.k, in.buf, in.off, x.d.s.sums, rows); + try tile.encode(x.e, t, w, in.buf, in.off, x.d.s.sums, 0, out.buf, out.off, rows); + } + + /// h += r when `add`, then out = norm(h) * gain: LFM2's own norm (the drafter's width and eps are the target's). + fn norm(x: Enc, h: mtl.Buffer, r: mtl.Buffer, gain: mtl.Buffer, add: bool, out: mtl.Buffer, rows: usize) void { + x.pipe("lfm2_norm"); + x.e.setBuffer(h, 0, 0); + x.e.setBuffer(r, 0, 1); + x.e.setBuffer(gain, 0, 2); + x.e.setValue(c.eps, 3); + x.e.setValue(@as(u32, @intFromBool(add)), 4); + x.e.setBuffer(out, 0, 5); + x.e.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(c.hidden / 4, 1, 1)); + } + + /// Rows' queries normed and rotated into `queries`, keys and values into the lane's layer at positions from `len`. + fn qk(x: Enc, q: Ref, k: Ref, v: Ref, layer: usize, lane: *const Lane, len: u64, queries: Ref, rows: usize) void { + const ly = &x.d.w.layers[layer]; + x.pipe("lfm2_qk"); + x.e.setBuffer(q.buf, q.off, 0); + x.e.setBuffer(k.buf, k.off, 1); + x.e.setBuffer(v.buf, v.off, 2); + x.e.setBuffer(ly.q_norm, 0, 3); + x.e.setBuffer(ly.k_norm, 0, 4); + const args = QkArgs{ .len = @intCast(len), .cap = @intCast(lane.cap), .eps = c.eps, .log2_theta = std.math.log2(x.d.config.theta) }; + x.e.setBytes(std.mem.asBytes(&args), 5); + x.e.setBuffer(queries.buf, queries.off, 6); + x.e.setBuffer(lane.keys[layer], 0, 7); + x.e.setBuffer(lane.values[layer], 0, 8); + x.e.dispatchGroups(mtl.Size.of(c.query_heads + c.kv_heads, rows, 1), mtl.Size.of(32, 1, 1)); + } + + /// Query rows [row, row + rows) against the lane's keys [0, keys), all of them visible to every row. + fn attend(x: Enc, layer: usize, lane: *const Lane, row: usize, rows: usize, keys: u64) void { + const s = &x.d.s; + const args = AttnArgs{ .len = @intCast(keys - 1), .cap = @intCast(lane.cap), .scale = 1.0 / @sqrt(@as(f32, c.head_dim)), .subs = st.subs }; + const used = (keys + st.chunk_keys - 1) / st.chunk_keys * st.subs; + x.pipe("lfm2_attn_block"); + x.e.setBuffer(s.queries, row * c.hidden * 2, 0); + x.e.setBuffer(lane.keys[layer], 0, 1); + x.e.setBuffer(lane.values[layer], 0, 2); + x.e.setBytes(std.mem.asBytes(&args), 3); + x.e.setBuffer(s.pm, 0, 4); + x.e.setBuffer(s.pl, 0, 5); + x.e.setBuffer(s.po, 0, 6); + x.e.setValue(@as(u32, @intCast(rows)), 7); + x.e.dispatchGroups(mtl.Size.of((rows + 7) / 8, c.kv_heads, @intCast(used)), mtl.Size.of(128, 1, 1)); + x.pipe("lfm2_attn_block_merge"); + x.e.setBuffer(s.pm, 0, 0); + x.e.setBuffer(s.pl, 0, 1); + x.e.setBuffer(s.po, 0, 2); + x.e.setBytes(std.mem.asBytes(&args), 3); + x.e.setBuffer(s.y, row * c.hidden * 2, 4); + x.e.dispatchGroups(mtl.Size.of(c.query_heads, rows, 1), mtl.Size.of(32, 1, 1)); + } +}; + +fn failed(msg: [*:0]const u8) anyerror { + std.log.err("lfm2 dspark: command buffer failed: {s}", .{msg}); + return error.GpuFailed; +} + +fn finish(cb: mtl.CommandBuffer, e: mtl.ComputeEncoder) !void { + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| return failed(msg); +} + +/// One stream's tapped rows in an absorb pass: `rows` rows at `src`, the context from position `p0`. +pub const Piece = struct { lane: *Lane, src: Ref, rows: usize, p0: u64 }; + +/// Each piece's rows into its lane's keys and values (fc, hidden_norm, k, v, k_norm, RoPE) in passes on `e`. +pub fn absorbPass(d: *Drafter, e: mtl.ComputeEncoder, pieces: []const Piece) !void { + const s = &d.s; + const x = Enc{ .d = d, .e = e }; + var rows: usize = 0; + for (pieces) |pc| { + try x.project(d.w.fc, pc.src, .{ .buf = s.fcx, .off = rows * c.hidden * 2 }, pc.rows); + rows += pc.rows; + } + if (rows == 0 or rows > ds.absorb_rows) return error.DraftContextGap; + x.norm(s.fcx, s.fcx, d.w.hidden_norm, false, s.hctx, rows); + for (0..d.config.layers) |i| { + const ly = &d.w.layers[i]; + try x.project(ly.k, .{ .buf = s.hctx }, .{ .buf = s.k }, rows); + try x.project(ly.v, .{ .buf = s.hctx }, .{ .buf = s.v }, rows); + var at: usize = 0; + for (pieces) |pc| { + // the kernel rotates query heads too: they read the context rows and land in x, which no one reads + const q = Ref{ .buf = s.hctx, .off = at * c.hidden * 2 }; + x.qk(q, .{ .buf = s.k, .off = at * kw * 2 }, .{ .buf = s.v, .off = at * kw * 2 }, i, pc.lane, pc.p0, .{ .buf = s.x }, pc.rows); + at += pc.rows; + } + } +} + +/// A command buffer for absorb passes, and its wait. +pub fn begin(d: *Drafter) struct { cb: mtl.CommandBuffer, e: mtl.ComputeEncoder } { + const cb = d.target.queue.commandBuffer(); + return .{ .cb = cb, .e = cb.compute(.serial) }; +} + +pub fn end(cb: mtl.CommandBuffer, e: mtl.ComputeEncoder) !void { + return finish(cb, e); +} + +/// The blocks' drafts into their lanes' held chains: the backbone over every block's rows, then the Markov chain. +pub fn hold(d: *Drafter, blocks: []const Block) !void { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const s = &d.s; + var depth: u32 = 0; + for (blocks) |b| depth = @max(depth, b.depth); + // a block of the deepest chain's rows, at most the 9 the drafter was trained on + const n: u32 = depth; + const rows = blocks.len * n; + if (rows == 0 or rows > ds.pass_rows) return error.DraftTooManyStreams; + const ids = s.ids.slice(u32, ds.pass_rows); + for (0..blocks.len) |i| { + ids[i * n] = blocks[i].pending; + @memset(ids[i * n + 1 ..][0 .. n - 1], d.config.mask); + } + const steps = depth; + // past the last layer's attention a lone block carries only the rows its chain reads; several keep them all + const tail = if (blocks.len == 1) steps else rows; + const cb = d.target.queue.commandBuffer(); + const x = Enc{ .d = d, .e = cb.compute(.serial) }; + const e = x.e; + // the blocks' staged context rows first, in the same command buffer + var staged: [ds.pass_rows]Piece = undefined; + var count: usize = 0; + var held: usize = 0; + for (blocks) |b| if (ds.Drafter.stagedPiece(b.lane)) |pc| { + if (held + pc.rows > ds.absorb_rows) { + try absorbPass(d, e, staged[0..count]); + count = 0; + held = 0; + } + staged[count] = pc; + count += 1; + held += pc.rows; + }; + if (count > 0) try absorbPass(d, e, staged[0..count]); + for (blocks) |b| { + b.lane.end += b.lane.staged; + b.lane.staged = 0; + } + const t = d.target; + const emb = t.weights.embedding; + e.setPipeline(t.kernels.embed(emb.bits)); + e.setBuffer(s.ids, 0, 0); + e.setBuffer(emb.w, emb.w_off, 1); + e.setBuffer(emb.scales, emb.s_off, 2); + e.setBuffer(emb.biases, emb.b_off, 3); + e.setBuffer(s.h, 0, 4); + const units = c.hidden / @as(usize, if (emb.bits == 6) 4 else 8 / emb.bits); + e.dispatchThreads(mtl.Size.of(units, rows, 1), mtl.Size.of(256, 1, 1)); + for (0..d.config.layers) |i| { + const ly = &d.w.layers[i]; + const out = if (i + 1 == d.config.layers) tail else rows; + x.norm(s.h, s.r, ly.input_norm, i > 0, s.x, rows); + try x.project(ly.q, .{ .buf = s.x }, .{ .buf = s.q }, rows); + try x.project(ly.k, .{ .buf = s.x }, .{ .buf = s.k }, rows); + try x.project(ly.v, .{ .buf = s.x }, .{ .buf = s.v }, rows); + for (blocks, 0..) |b, bi| { + const at = bi * n; + x.qk(.{ .buf = s.q, .off = at * c.hidden * 2 }, .{ .buf = s.k, .off = at * kw * 2 }, .{ .buf = s.v, .off = at * kw * 2 }, i, b.lane, b.anchor, .{ .buf = s.queries, .off = at * c.hidden * 2 }, n); + x.attend(i, b.lane, at, if (blocks.len == 1) out else n, b.anchor + n); + } + try x.project(ly.o, .{ .buf = s.y }, .{ .buf = s.r }, out); + x.norm(s.h, s.r, ly.post_norm, true, s.x, out); + try x.project(ly.gate, .{ .buf = s.x }, .{ .buf = s.gate }, out); + try x.project(ly.up, .{ .buf = s.x }, .{ .buf = s.up }, out); + x.pipe("lfm2_swiglu"); + e.setBuffer(s.gate, 0, 0); + e.setBuffer(s.up, 0, 1); + e.setBuffer(s.act, 0, 2); + e.setValue(@as(u32, @intCast(out * d.config.intermediate)), 3); + e.dispatchThreads(mtl.Size.of(out * d.config.intermediate / 8, 1, 1), mtl.Size.of(256, 1, 1)); + try x.project(ly.down, .{ .buf = s.act }, .{ .buf = s.r }, out); + } + x.norm(s.h, s.r, d.w.norm, true, s.x, tail); + try x.project(fwd.matrix(emb), .{ .buf = s.x }, .{ .buf = s.logits }, tail); + try markov(x, blocks.len, n, steps); + try finish(cb, e); + const picks = s.picks.slice(Pick, d.config.block * ds.pass_rows); + for (blocks, 0..) |b, bi| { + for (b.lane.held[0..steps], 0..) |*h, i| { + const p = picks[i * ds.pass_rows + bi]; + if (p.nonfinite != 0 or p.token >= c.vocab) return error.DraftNonfinite; + h.* = p.token; + } + b.lane.held_len = b.depth; + } +} + +/// Step i of each stream: its row i's logits plus W2 W1[previous draft] (the pending token first), then the argmax. +fn markov(x: Enc, streams: usize, n: u32, depth: u32) !void { + const d = x.d; + const s = &d.s; + const e = x.e; + for (0..depth) |i| { + // step i's previous token: the pending one (each block's first id), then step i - 1's pick + x.pipe("lfm2_markov_scan"); + e.setBuffer(s.logits, i * c.vocab * 2, 0); + e.setBuffer(d.w.w1, 0, 1); + e.setBuffer(d.w.w2.w, d.w.w2.w_off, 2); + if (i == 0) e.setBuffer(s.ids, 0, 3) else e.setBuffer(s.picks, (i - 1) * ds.pass_rows * @sizeOf(Pick), 3); + const args = MarkovArgs{ .stride = n * c.vocab, .rank = d.config.rank, .vocab = c.vocab, .prev_stride = if (i == 0) n else @sizeOf(Pick) / 4 }; + e.setBytes(std.mem.asBytes(&args), 4); + e.setBuffer(s.best, 0, 5); + e.setBuffer(s.at, 0, 6); + e.setBuffer(d.w.w2.scales, d.w.w2.s_off, 7); + e.setBuffer(d.w.w2.biases, d.w.w2.b_off, 8); + e.dispatchGroups(mtl.Size.of(scan_groups, streams, 1), mtl.Size.of(256, 1, 1)); + x.pipe("lfm2_markov_pick"); + e.setBuffer(s.best, 0, 0); + e.setBuffer(s.at, 0, 1); + e.setValue(@as(u32, scan_groups), 2); + e.setBuffer(s.picks, i * ds.pass_rows * @sizeOf(Pick), 3); + e.dispatchGroups(mtl.Size.of(streams, 1, 1), mtl.Size.of(256, 1, 1)); + } +} diff --git a/zig/src/families/lfm2_moe/dspark_config.zig b/zig/src/families/lfm2_moe/dspark_config.zig new file mode 100644 index 0000000000..7b3b628a02 --- /dev/null +++ b/zig/src/families/lfm2_moe/dspark_config.zig @@ -0,0 +1,136 @@ +//! LiquidAI's DSpark drafter for LFM2.5-24B-A2B: its config admitted, any other shape refused with its reason. +const std = @import("std"); +const c = @import("config.zig"); + +/// The geometry the drafter's kernels take: the target's attention shape, five taps at most eight layers. +pub const Config = struct { + taps: [8]u32 = @splat(0), + tap_count: u32, + layers: u32, + intermediate: u32, + /// Rows a block takes: the pending token and block - 1 mask tokens, one draft a row. + block: u32, + mask: u32, + rank: u32, + theta: f32, + + pub fn read(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Config { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 20)); + defer gpa.free(bytes); + return parse(gpa, bytes); + } + + pub fn parse(gpa: std.mem.Allocator, bytes: []const u8) !Config { + var parsed = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + defer parsed.deinit(); + const o = if (parsed.value == .object) parsed.value.object else return error.BadDraftConfig; + const archs = o.get("architectures") orelse return error.BadDraftConfig; + if (archs != .array or archs.array.items.len != 1 or archs.array.items[0] != .string) return error.BadDraftConfig; + if (!std.mem.eql(u8, archs.array.items[0].string, "Lfm2DSparkDraftModel")) return error.NotLfmDSpark; + const d = if (o.get("dflash_config")) |v| (if (v == .object) v.object else return error.BadDraftConfig) else return error.BadDraftConfig; + var out: Config = .{ + .tap_count = 0, + .layers = try small(o, "num_hidden_layers"), + .intermediate = try small(o, "intermediate_size"), + .block = try small(o, "block_size"), + .mask = try small(d, "mask_token_id"), + .rank = try small(o, "markov_rank"), + .theta = @floatCast(try number(o, "rope_theta")), + }; + // the target's widths: the drafter shares its embedding, head and attention kernels + if (try small(o, "hidden_size") != c.hidden or try small(o, "vocab_size") != c.vocab) return error.DraftForAnotherTarget; + if (try small(d, "num_target_layers") != c.layers) return error.DraftForAnotherTarget; + const heads = try small(o, "num_attention_heads"); + const kv = try small(o, "num_key_value_heads"); + if (heads != c.query_heads or kv != c.kv_heads or try small(o, "head_dim") != c.head_dim) return error.UnsupportedDraftShape; + if (@as(f32, @floatCast(try number(o, "rms_norm_eps"))) != c.eps) return error.UnsupportedDraftShape; + const taps = d.get("target_layer_ids") orelse return error.BadDraftConfig; + if (taps != .array or taps.array.items.len == 0 or taps.array.items.len > out.taps.len) return error.BadDraftConfig; + for (taps.array.items, 0..) |t, i| { + if (t != .integer or t.integer < 0) return error.BadDraftConfig; + if (t.integer >= c.layers) return error.DraftForAnotherTarget; + out.taps[i] = @intCast(t.integer); + } + out.tap_count = @intCast(taps.array.items.len); + try kinds(o, out.layers); + // the vendored defaults decide the block: non-causal attention, interleaved RoPE, a plain Markov head + if (o.get("is_causal") != null) return error.UnsupportedDraftMask; + const neox = o.get("rope_is_neox_style") orelse return error.BadDraftConfig; + if (neox != .bool or neox.bool) return error.UnsupportedDraftRope; + const head = o.get("markov_head_type") orelse return error.BadDraftConfig; + if (head != .string or !std.mem.eql(u8, head.string, "vanilla")) return error.UnsupportedMarkovHead; + if (o.get("sample_from_anchor")) |v| if (v != .bool or !v.bool) return error.UnsupportedDraftAnchor; + if (out.layers == 0 or out.layers > 8 or out.block < 2 or out.block > 16) return error.UnsupportedDraftShape; + if (out.intermediate == 0 or out.intermediate % 64 != 0 or out.rank == 0 or out.rank % 64 != 0) return error.UnsupportedDraftShape; + if (out.mask >= c.vocab or !(out.theta > 0)) return error.BadDraftConfig; + return out; + } +}; + +/// Every layer full attention (the block's mask has no sliding window). +fn kinds(o: std.json.ObjectMap, layers: u32) !void { + const v = o.get("layer_types") orelse return; + if (v != .array or v.array.items.len != layers) return error.BadDraftConfig; + for (v.array.items) |k| if (k != .string or !std.mem.eql(u8, k.string, "full_attention")) return error.UnsupportedDraftMask; +} + +fn small(o: std.json.ObjectMap, key: []const u8) !u32 { + const v = o.get(key) orelse return error.BadDraftConfig; + if (v != .integer or v.integer < 0 or v.integer > std.math.maxInt(u32)) return error.BadDraftConfig; + return @intCast(v.integer); +} + +fn number(o: std.json.ObjectMap, key: []const u8) !f64 { + const v = o.get(key) orelse return error.BadDraftConfig; + return switch (v) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => error.BadDraftConfig, + }; +} + +/// LiquidAI/LFM2.5-24B-A2B-DSpark's config.json at e052e5b1dd10, as released. +pub const stock = + \\{"architectures":["Lfm2DSparkDraftModel"],"model_type":"qwen3","dtype":"bfloat16","hidden_size":2048, + \\"num_hidden_layers":5,"num_attention_heads":32,"num_key_value_heads":8,"head_dim":64,"intermediate_size":6144, + \\"hidden_act":"silu","rms_norm_eps":1e-05,"vocab_size":128000,"rope_theta":10000000.0, + \\"max_position_embeddings":131072,"layer_types":["full_attention","full_attention","full_attention", + \\"full_attention","full_attention"],"block_size":9,"dflash_config":{"mask_token_id":125017, + \\"target_layer_ids":[2,10,18,26,34],"num_target_layers":40},"markov_rank":256,"rope_is_neox_style":false, + \\"enable_confidence_head":true,"markov_head_type":"vanilla"} +; + +fn edited(from: []const u8, to: []const u8) !Config { + const a = std.testing.allocator; + const text = try std.mem.replaceOwned(u8, a, stock, from, to); + defer a.free(text); + return Config.parse(a, text); +} + +test "lfm2 dspark: the released config is admitted with its taps, block, mask token and Markov rank" { + const cfg = try Config.parse(std.testing.allocator, stock); + try std.testing.expectEqualSlices(u32, &.{ 2, 10, 18, 26, 34 }, cfg.taps[0..cfg.tap_count]); + try std.testing.expectEqual(@as(u32, 5), cfg.layers); + try std.testing.expectEqual(@as(u32, 9), cfg.block); + try std.testing.expectEqual(@as(u32, 125017), cfg.mask); + try std.testing.expectEqual(@as(u32, 256), cfg.rank); + try std.testing.expectEqual(@as(u32, 6144), cfg.intermediate); + try std.testing.expectEqual(@as(f32, 1e7), cfg.theta); +} + +test "lfm2 dspark: other drafters and shapes are refused with their reason" { + try std.testing.expectError(error.NotLfmDSpark, edited("Lfm2DSparkDraftModel", "DFlashDraftModel")); + try std.testing.expectError(error.DraftForAnotherTarget, edited("\"hidden_size\":2048", "\"hidden_size\":4096")); + try std.testing.expectError(error.DraftForAnotherTarget, edited("\"num_target_layers\":40", "\"num_target_layers\":24")); + try std.testing.expectError(error.DraftForAnotherTarget, edited("26,34]", "26,44]")); + try std.testing.expectError(error.UnsupportedDraftShape, edited("\"head_dim\":64", "\"head_dim\":128")); + try std.testing.expectError(error.UnsupportedDraftRope, edited("\"rope_is_neox_style\":false", "\"rope_is_neox_style\":true")); + try std.testing.expectError(error.UnsupportedMarkovHead, edited("\"vanilla\"", "\"gated\"")); + try std.testing.expectError(error.UnsupportedDraftMask, edited("\"full_attention\"]", "\"sliding_attention\"]")); + try std.testing.expectError(error.UnsupportedDraftMask, edited("\"block_size\"", "\"is_causal\":true,\"block_size\"")); + try std.testing.expectError(error.UnsupportedDraftAnchor, edited("\"block_size\"", "\"sample_from_anchor\":false,\"block_size\"")); + try std.testing.expectError(error.UnsupportedDraftShape, edited("\"block_size\":9", "\"block_size\":17")); + try std.testing.expectError(error.BadDraftConfig, edited("\"mask_token_id\":125017", "\"mask_token_id\":128000")); +} diff --git a/zig/src/families/lfm2_moe/dspark_quant.zig b/zig/src/families/lfm2_moe/dspark_quant.zig new file mode 100644 index 0000000000..a6a5baccd6 --- /dev/null +++ b/zig/src/families/lfm2_moe/dspark_quant.zig @@ -0,0 +1,118 @@ +//! BF16 drafter rows quantized at load to MLX's affine codes in groups of 64, the layout core/tile_qmm reads. +const std = @import("std"); + +/// MLX's affine codes of 64 bf16 values at `bits`: bias the larger edge, scale |edge| / round(|edge| / step). +pub fn group(x: []const u16, bits: u8, codes: []u8, scale: *u16, bias: *u16) void { + var lo: f32 = std.math.inf(f32); + var hi: f32 = -std.math.inf(f32); + for (x) |v| { + lo = @min(lo, toF32(v)); + hi = @max(hi, toF32(v)); + } + const bins: f32 = @floatFromInt((@as(u32, 1) << @intCast(bits)) - 1); + const step = @max(hi - lo, 1e-7) / bins; + const from_min = @abs(lo) > @abs(hi); + const edge = if (from_min) lo else hi; + const q0 = @round(@abs(edge) / step); + const s_f: f32 = if (q0 == 0) step else (if (from_min) @abs(edge) else -@abs(edge)) / q0; + scale.* = toBf16(s_f); + bias.* = toBf16(if (q0 == 0) lo else edge); + const s = toF32(scale.*); + const b = toF32(bias.*); + @memset(codes, 0); + for (x, 0..) |v, i| { + const q: u8 = @intFromFloat(std.math.clamp(@round((toF32(v) - b) / s), 0, bins)); + if (bits == 8) codes[i] = q else codes[i / 2] |= q << @intCast((i % 2) * 4); + } +} + +pub fn toF32(v: u16) f32 { + return @bitCast(@as(u32, v) << 16); +} + +/// Round to nearest even, as MLX casts. +pub fn toBf16(f: f32) u16 { + const b: u32 = @bitCast(f); + if ((b & 0x7f80_0000) == 0x7f80_0000) return @truncate(b >> 16); + return @truncate((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +/// A matrix's codes [n][k * bits / 8], scales and biases [n][k / 64], each written by one thread's rows. +pub const Job = struct { + src: []const u16, + k: usize, + bits: u8, + codes: []u8, + scales: []u16, + biases: []u16, + fn run(j: Job, from: usize, to: usize) void { + const groups = j.k / 64; + const per = 64 * @as(usize, j.bits) / 8; + for (from..to) |r| for (0..groups) |g| { + const at = r * groups + g; + group(j.src[r * j.k + g * 64 ..][0..64], j.bits, j.codes[at * per ..][0..per], &j.scales[at], &j.biases[at]); + }; + } + + /// The job over its rows on every CPU (at most 16 threads). + pub fn all(j: Job, gpa: std.mem.Allocator, rows: usize) !void { + const threads = @min(std.Thread.getCpuCount() catch 1, 16); + const per = (rows + threads - 1) / threads; + const handles = try gpa.alloc(std.Thread, threads); + defer gpa.free(handles); + var spawned: usize = 0; + defer for (handles[0..spawned]) |h| h.join(); + var at: usize = 0; + while (at < rows) : (at += per) { + handles[spawned] = try std.Thread.spawn(.{}, run, .{ j, at, @min(rows, at + per) }); + spawned += 1; + } + } +}; + +/// Head `h`'s interleaved RoPE pairs (2i, 2i + 1) as half-split pairs (i, i + dim / 2): row h * dim + j's source. +pub fn halfSplit(row: usize, dim: usize) usize { + const h = row / dim; + const j = row % dim; + const half = dim / 2; + return h * dim + if (j < half) 2 * j else 2 * (j - half) + 1; +} + +test "lfm2 dspark: 4- and 8-bit codes keep the larger edge exact, as MLX quantizes" { + var x: [64]u16 = undefined; + for (&x, 0..) |*v, i| v.* = toBf16(@as(f32, @floatFromInt(i)) * 0.125 - 2.0); + for ([_]u8{ 4, 8 }) |bits| { + var codes: [64]u8 = undefined; + var s: u16 = undefined; + var b: u16 = undefined; + group(&x, bits, codes[0 .. 64 * bits / 8], &s, &b); + try std.testing.expect(toF32(s) < 0 and toF32(b) == 5.875); // the max is the edge: codes descend from it + const last: u8 = if (bits == 8) codes[63] else codes[31] >> 4; + try std.testing.expectEqual(@as(u8, 0), last); + for (x, 0..) |v, i| { + const q: u8 = if (bits == 8) codes[i] else (codes[i / 2] >> @intCast((i % 2) * 4)) & 0xf; + const back = toF32(s) * @as(f32, @floatFromInt(q)) + toF32(b); + try std.testing.expect(@abs(back - toF32(v)) <= @abs(toF32(s)) * 0.51); + } + } +} + +test "lfm2 dspark: half-split rows rotated by the half-split rule equal interleaved pairs rotated, row for row" { + const dim = 64; + var x: [2 * dim]f32 = undefined; + for (&x, 0..) |*v, i| v.* = @sin(@as(f32, @floatFromInt(i)) * 0.37) * 3; + var moved: [2 * dim]f32 = undefined; + for (&moved, 0..) |*v, r| v.* = x[halfSplit(r, dim)]; + const pos: f32 = 1234; + for (0..2) |h| for (0..dim / 2) |i| { + const angle = pos * std.math.pow(f32, 1e7, -@as(f32, @floatFromInt(i)) / (dim / 2)); + const cs = @cos(angle); + const sn = @sin(angle); + const a = x[h * dim + 2 * i]; + const b = x[h * dim + 2 * i + 1]; + const l = moved[h * dim + i]; + const r = moved[h * dim + i + dim / 2]; + try std.testing.expectEqual(a * cs - b * sn, l * cs - r * sn); + try std.testing.expectEqual(a * sn + b * cs, l * sn + r * cs); + }; +} diff --git a/zig/src/families/lfm2_moe/dspark_test.zig b/zig/src/families/lfm2_moe/dspark_test.zig index d220b9ea66..9ee31d868e 100644 --- a/zig/src/families/lfm2_moe/dspark_test.zig +++ b/zig/src/families/lfm2_moe/dspark_test.zig @@ -1,8 +1,9 @@ -//! The DSpark drafter's kernels on the GPU against host references: the target's taps. +//! The DSpark drafter's kernels on the GPU against host references: taps, the Markov step, half-split RoPE. const std = @import("std"); const mtl = @import("metal"); const c = @import("config.zig"); const kernels = @import("kernels.zig"); +const quant = @import("dspark_quant.zig"); const shared = mtl.ResourceOptions.shared; @@ -45,15 +46,9 @@ const Rig = struct { } }; -/// 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); + for (out) |*v| v.* = quant.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" { @@ -75,3 +70,139 @@ test "lfm2 dspark: the tap kernel copies each row's hidden values into its tap's for (got[(row * taps + 2) * c.hidden ..][0..c.hidden]) |v| try std.testing.expectEqual(@as(u16, 0), v); } } + +test "lfm2 dspark: a Markov step picks the argmax of bf16(base + bf16(W2 . W1[prev])) over every threadgroup" { + var r = try Rig.init(); + defer r.deinit(); + const rank = 256; + const vocab = 1300; // three scan threadgroups, the last one short + const streams = 4; + const groups = (vocab + 511) / 512; + const w1 = try r.buffer(u16, vocab * rank); + defer w1.deinit(); + noise(11, w1.slice(u16, vocab * rank), 1); + const codes = try r.buffer(u8, vocab * rank); + defer codes.deinit(); + for (codes.slice(u8, vocab * rank), 0..) |*q, i| q.* = @truncate(i *% 2654435761 >> 7); + const sc = try r.buffer(u16, vocab * rank / 64); + defer sc.deinit(); + const bi = try r.buffer(u16, vocab * rank / 64); + defer bi.deinit(); + noise(12, sc.slice(u16, vocab * rank / 64), 0.01); + noise(13, bi.slice(u16, vocab * rank / 64), 0.5); + // a token every second word (the picks' layout): inside W1, its last row, past it (a zero row), and a NaN row + const ids = try r.buffer(u32, 2 * streams); + defer ids.deinit(); + @memcpy(ids.slice(u32, 2 * streams), &[_]u32{ 17, 9, vocab - 1, 9, vocab + 5, 0, 3, 3 }); + const logits = try r.buffer(u16, streams * 2 * vocab); + defer logits.deinit(); + noise(14, logits.slice(u16, streams * 2 * vocab), 4); + logits.slice(u16, streams * 2 * vocab)[3 * 2 * vocab + 700] = 0x7fc0; + const best = try r.buffer(f32, streams * groups); + defer best.deinit(); + const at = try r.buffer(u32, streams * groups); + defer at.deinit(); + const picks = try r.buffer(u32, 2 * streams); + defer picks.deinit(); + const cb = r.queue.commandBuffer(); + const e = cb.compute(.serial); + e.setPipeline(r.k.get("lfm2_markov_scan")); + for ([_]mtl.Buffer{ logits, w1, codes, ids }, 0..) |b, i| e.setBuffer(b, 0, i); + e.setValue([4]u32{ 2 * vocab, rank, vocab, 2 }, 4); + for ([_]mtl.Buffer{ best, at, sc, bi }, 5..) |b, i| e.setBuffer(b, 0, i); + e.dispatchGroups(mtl.Size.of(groups, streams, 1), mtl.Size.of(256, 1, 1)); + e.setPipeline(r.k.get("lfm2_markov_pick")); + e.setBuffer(best, 0, 0); + e.setBuffer(at, 0, 1); + e.setValue(@as(u32, groups), 2); + e.setBuffer(picks, 0, 3); + e.dispatchGroups(mtl.Size.of(streams, 1, 1), mtl.Size.of(256, 1, 1)); + e.end(); + cb.commit(); + cb.wait(); + const got = picks.slice(u32, 2 * streams); + const base = logits.slice(u16, streams * 2 * vocab); + const id_list = ids.slice(u32, 2 * streams); + for (0..streams) |s| { + if (s == 3) { // the planted NaN: flagged + try std.testing.expectEqual(@as(u32, 1), got[2 * s + 1]); + continue; + } + try std.testing.expectEqual(@as(u32, 0), got[2 * s + 1]); + const id = id_list[2 * s]; + var top: f64 = -std.math.inf(f64); + var mine: f64 = 0; + for (0..vocab) |v| { + var dot: f64 = 0; + for (0..rank) |i| { + const g = (v * rank + i) / 64; + const w = @as(f64, quant.toF32(sc.slice(u16, vocab * rank / 64)[g])) * @as(f64, @floatFromInt(codes.slice(u8, vocab * rank)[v * rank + i])) + quant.toF32(bi.slice(u16, vocab * rank / 64)[g]); + dot += w * (if (id < vocab) quant.toF32(w1.slice(u16, vocab * rank)[id * rank + i]) else 0); + } + const score: f64 = quant.toF32(quant.toBf16(quant.toF32(base[s * 2 * vocab + v]) + quant.toF32(quant.toBf16(@floatCast(dot))))); + top = @max(top, score); + if (v == got[2 * s]) mine = score; + } + // the pick is the best score, or within two bf16 steps of it where fp32 sums round apart + try std.testing.expect(got[2 * s] < vocab and mine >= top - @abs(top) * 0x1p-7); + } +} + +test "lfm2 dspark: lfm2_qk on half-split rows and gains gives interleaved RoPE's keys, reordered" { + var r = try Rig.init(); + defer r.deinit(); + const hd = c.head_dim; + const pos: u32 = 777; + const theta: f32 = 1e7; + var k_il: [c.kv_heads * hd]u16 = undefined; + var gain_il: [hd]u16 = undefined; + noise(21, &k_il, 2); + noise(22, &gain_il, 1.5); + const q = try r.buffer(u16, c.query_heads * hd); + defer q.deinit(); + const k = try r.buffer(u16, c.kv_heads * hd); + defer k.deinit(); + const v = try r.buffer(u16, c.kv_heads * hd); + defer v.deinit(); + const qn = try r.buffer(u16, hd); + defer qn.deinit(); + const kn = try r.buffer(u16, hd); + defer kn.deinit(); + for (k.slice(u16, c.kv_heads * hd), 0..) |*x, i| x.* = k_il[quant.halfSplit(i, hd)]; + for (kn.slice(u16, hd), 0..) |*x, i| x.* = gain_il[quant.halfSplit(i, hd)]; + @memset(qn.slice(u16, hd), quant.toBf16(1)); + const cap = 1024; + const queries = try r.buffer(u16, c.query_heads * hd); + defer queries.deinit(); + const keys = try r.buffer(u16, c.kv_heads * cap * hd); + defer keys.deinit(); + const values = try r.buffer(u16, c.kv_heads * cap * hd); + defer values.deinit(); + const Args = extern struct { len: u32, cap: u32, eps: f32, log2_theta: f32 }; + const cb = r.queue.commandBuffer(); + const e = cb.compute(.serial); + e.setPipeline(r.k.get("lfm2_qk")); + for ([_]mtl.Buffer{ q, k, v, qn, kn }, 0..) |b, i| e.setBuffer(b, 0, i); + const args = Args{ .len = pos, .cap = cap, .eps = c.eps, .log2_theta = std.math.log2(theta) }; + e.setBytes(std.mem.asBytes(&args), 5); + for ([_]mtl.Buffer{ queries, keys, values }, 6..) |b, i| e.setBuffer(b, 0, i); + e.dispatchGroups(mtl.Size.of(c.query_heads + c.kv_heads, 1, 1), mtl.Size.of(32, 1, 1)); + e.end(); + cb.commit(); + cb.wait(); + const got = keys.slice(u16, c.kv_heads * cap * hd); + for (0..c.kv_heads) |h| { + var ss: f32 = 0; + for (0..hd) |i| ss += quant.toF32(k_il[h * hd + i]) * quant.toF32(k_il[h * hd + i]); + const inv = 1 / @sqrt(ss / hd + c.eps); + for (0..hd / 2) |i| { + const angle = @as(f32, pos) * std.math.pow(f32, theta, -@as(f32, @floatFromInt(i)) / (hd / 2)); + const a = quant.toF32(k_il[h * hd + 2 * i]) * inv * quant.toF32(gain_il[2 * i]); + const b = quant.toF32(k_il[h * hd + 2 * i + 1]) * inv * quant.toF32(gain_il[2 * i + 1]); + const want = [2]f32{ a * @cos(angle) - b * @sin(angle), a * @sin(angle) + b * @cos(angle) }; + const at = (h * cap + pos) * hd; + const have = [2]f32{ quant.toF32(got[at + i]), quant.toF32(got[at + i + hd / 2]) }; + for (want, have) |w, x| try std.testing.expect(@abs(w - x) <= 0.02 * @max(1, @abs(w))); + } + } +} diff --git a/zig/src/families/lfm2_moe/dspark_weights.zig b/zig/src/families/lfm2_moe/dspark_weights.zig new file mode 100644 index 0000000000..81dc727e56 --- /dev/null +++ b/zig/src/families/lfm2_moe/dspark_weights.zig @@ -0,0 +1,144 @@ +//! The DSpark drafter's weights: bf16 linears quantized to MLX affine codes at load, norms and W1 kept as stored. +const std = @import("std"); +const mtl = @import("metal"); +const c = @import("config.zig"); +const dc = @import("dspark_config.zig"); +const quant = @import("dspark_quant.zig"); +const ckpt = @import("../../core/checkpoint_metal.zig"); +const shards = @import("../../core/checkpoint.zig"); +const tile = @import("../../core/tile_qmm.zig"); + +const shared = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; + +/// A bf16 vector in its own buffer (norm gains; RoPE-paired ones in the half-split order the kernels rotate). +pub const Vector = mtl.Buffer; + +pub const Layer = struct { + q: tile.Matrix, + k: tile.Matrix, + v: tile.Matrix, + o: tile.Matrix, + gate: tile.Matrix, + up: tile.Matrix, + down: tile.Matrix, + input_norm: Vector, + post_norm: Vector, + q_norm: Vector, + k_norm: Vector, +}; + +pub const Weights = struct { + fc: tile.Matrix, + hidden_norm: Vector, + norm: Vector, + layers: [8]Layer, + /// The Markov head: W1 [vocab, rank] bf16 as stored (rows gathered by token), W2 [vocab, rank] as 8-bit codes. + w1: mtl.Buffer, + w2: tile.Matrix, + owned: std.ArrayList(mtl.Buffer) = .empty, + + pub fn deinit(w: *Weights, gpa: std.mem.Allocator) void { + for (w.owned.items) |b| b.deinit(); + w.owned.deinit(gpa); + } + + /// Every buffer the weights hold, for a residency set. + pub fn buffers(w: *const Weights) []const mtl.Buffer { + return w.owned.items; + } +}; + +const Loader = struct { + gpa: std.mem.Allocator, + device: mtl.Device, + check: *const ckpt.Checkpoint, + owned: *std.ArrayList(mtl.Buffer), + bits: u8, + + fn buffer(l: Loader, bytes: usize) !mtl.Buffer { + const b = try l.device.buffer(bytes, shared); + errdefer b.deinit(); + try l.owned.append(l.gpa, b); + return b; + } + + fn tensor(l: Loader, name: []const u8, shape: []const usize) !ckpt.Tensor { + const t = try l.check.get(name); + if (t.dtype != .bf16 or t.rank != shape.len or !std.mem.eql(usize, t.shape[0..t.rank], shape)) { + std.log.err("{s}: unexpected DSpark tensor dtype or shape", .{name}); + return error.BadDraftWeight; + } + return t; + } + + /// A bf16 vector copied out, its pairs reordered for half-split RoPE when `paired`. + fn vector(l: Loader, name: []const u8, len: usize, paired: bool) !Vector { + const src = (try l.tensor(name, &.{len})).host(u16); + const b = try l.buffer(len * 2); + const dst = b.slice(u16, len); + for (dst, 0..) |*v, i| v.* = src[if (paired) quant.halfSplit(i, len) else i]; + return b; + } + + /// An [n, k] bf16 matrix as affine codes at the loader's width; rows reordered for half-split RoPE when `paired`. + fn matrix(l: Loader, name: []const u8, n: usize, k: usize, paired: bool) !tile.Matrix { + const t = try l.tensor(name, &.{ n, k }); + var rows = t.host(u16); + var moved: ?[]u16 = null; + defer if (moved) |m| l.gpa.free(m); + if (paired) { + const m = try l.gpa.alloc(u16, n * k); + moved = m; + for (0..n) |r| @memcpy(m[r * k ..][0..k], rows[quant.halfSplit(r, c.head_dim) * k ..][0..k]); + rows = m; + } + const codes = try l.buffer(n * k * l.bits / 8); + const scales = try l.buffer(n * k / 64 * 2); + const biases = try l.buffer(n * k / 64 * 2); + const job: quant.Job = .{ .src = rows, .k = k, .bits = l.bits, .codes = codes.slice(u8, n * k * l.bits / 8), .scales = scales.slice(u16, n * k / 64), .biases = biases.slice(u16, n * k / 64) }; + try job.all(l.gpa, n); + return .{ .w = codes, .scales = scales, .biases = biases, .n = n, .k = k, .bits = l.bits }; + } +}; + +fn part(buf: []u8, layer: usize, name: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "layers.{d}.{s}.weight", .{ layer, name }); +} + +/// The drafter's checkpoint read and quantized at `bits` (4 or 8); the checkpoint's own buffers are released after. +pub fn load(gpa: std.mem.Allocator, device: mtl.Device, io: std.Io, dir: []const u8, cfg: *const dc.Config, bits: u8) !Weights { + if (bits != 4 and bits != 8) return error.UnsupportedDraftBits; + var check = ckpt.Checkpoint.init(gpa); + defer check.deinit(); + const files = try shards.shardFiles(gpa, io, dir); + defer shards.freeShardFiles(gpa, files); + for (files) |path| try check.addFile(device, path, ""); + var w: Weights = undefined; + w.owned = .empty; + errdefer w.deinit(gpa); + const l: Loader = .{ .gpa = gpa, .device = device, .check = &check, .owned = &w.owned, .bits = bits }; + const qw = c.query_heads * c.head_dim; + const kw = c.kv_heads * c.head_dim; + w.fc = try l.matrix("fc.weight", c.hidden, cfg.tap_count * c.hidden, false); + w.hidden_norm = try l.vector("hidden_norm.weight", c.hidden, false); + w.norm = try l.vector("norm.weight", c.hidden, false); + var name: [128]u8 = undefined; + for (w.layers[0..cfg.layers], 0..) |*ly, i| { + ly.q = try l.matrix(try part(&name, i, "self_attn.q_proj"), qw, c.hidden, true); + ly.k = try l.matrix(try part(&name, i, "self_attn.k_proj"), kw, c.hidden, true); + ly.v = try l.matrix(try part(&name, i, "self_attn.v_proj"), kw, c.hidden, false); + ly.o = try l.matrix(try part(&name, i, "self_attn.o_proj"), c.hidden, qw, false); + ly.gate = try l.matrix(try part(&name, i, "mlp.gate_proj"), cfg.intermediate, c.hidden, false); + ly.up = try l.matrix(try part(&name, i, "mlp.up_proj"), cfg.intermediate, c.hidden, false); + ly.down = try l.matrix(try part(&name, i, "mlp.down_proj"), c.hidden, cfg.intermediate, false); + ly.input_norm = try l.vector(try part(&name, i, "input_layernorm"), c.hidden, false); + ly.post_norm = try l.vector(try part(&name, i, "post_attention_layernorm"), c.hidden, false); + ly.q_norm = try l.vector(try part(&name, i, "self_attn.q_norm"), c.head_dim, true); + ly.k_norm = try l.vector(try part(&name, i, "self_attn.k_norm"), c.head_dim, true); + } + const w1 = (try l.tensor("markov_head.markov_w1.weight", &.{ c.vocab, cfg.rank })).host(u16); + w.w1 = try l.buffer(c.vocab * cfg.rank * 2); + @memcpy(w.w1.slice(u16, c.vocab * cfg.rank), w1); + w.w2 = try l.matrix("markov_head.markov_w2.weight", c.vocab, cfg.rank, false); + return w; +} diff --git a/zig/src/families/lfm2_moe/kernels.zig b/zig/src/families/lfm2_moe/kernels.zig index 61af45f4c8..357500c04d 100644 --- a/zig/src/families/lfm2_moe/kernels.zig +++ b/zig/src/families/lfm2_moe/kernels.zig @@ -16,6 +16,8 @@ const names = [_][:0]const u8{ "lfm2_attn_mma", "lfm2_attn_mma_r2", "lfm2_attn_merge", + "lfm2_attn_block", + "lfm2_attn_block_merge", "lfm2_route_logits", "lfm2_route_select", "lfm2_swiglu", @@ -28,6 +30,8 @@ const names = [_][:0]const u8{ "lfm2_attn_segs", "lfm2_attn_merge_segs", "lfm2_tap", + "lfm2_markov_scan", + "lfm2_markov_pick", }; pub const Kernels = struct { diff --git a/zig/src/families/lfm2_moe/lfm2_moe.zig b/zig/src/families/lfm2_moe/lfm2_moe.zig index 97f9cbd5d8..8f861e31b4 100644 --- a/zig/src/families/lfm2_moe/lfm2_moe.zig +++ b/zig/src/families/lfm2_moe/lfm2_moe.zig @@ -7,3 +7,7 @@ pub const state = @import("state.zig"); pub const forward = @import("forward.zig"); pub const backend = @import("backend.zig"); pub const snapshot = @import("snapshot.zig"); +pub const dspark = @import("dspark.zig"); +pub const dspark_config = @import("dspark_config.zig"); +pub const dspark_block = @import("dspark_block.zig"); +pub const dspark_quant = @import("dspark_quant.zig"); diff --git a/zig/src/native/lfm2_moe.zig b/zig/src/native/lfm2_moe.zig index 5340d807de..05f0260dcc 100644 --- a/zig/src/native/lfm2_moe.zig +++ b/zig/src/native/lfm2_moe.zig @@ -1,4 +1,4 @@ -//! LFM2-MoE served by the existing lane host and scheduler. +//! LFM2-MoE served by the existing lane host and scheduler, with LiquidAI's DSpark drafter from --drafter. const std = @import("std"); const mtl = @import("metal"); const api = @import("engine_api"); @@ -8,11 +8,16 @@ const lanes = tf.lanes; const cache_fit = @import("cache_fit.zig"); const metal_admit = @import("metal_admit.zig"); +/// --drafter loads LiquidAI's DSpark drafter (LiquidAI/LFM2.5-24B-A2B-DSpark) in front of the engine. +pub const takes_drafter = true; + const Host = struct { resident: ?mtl.resident.Resident = null, gpa: std.mem.Allocator, model: *q.Model, metal: *q.backend.Metal, + draft: ?*q.dspark.Drafter = null, + wrap: ?lanes.drafted.Drafted = null, warm: mtl.keepalive.Target = undefined, // the model's queue, for the lane host's idle ticker config: lanes.Config, clock: lanes.backend.WallClock, @@ -25,7 +30,9 @@ const Host = struct { h.host.stop(); if (h.cache) |*store| store.deinit(); h.core.deinit(); + if (h.wrap) |*w| w.deinit(); mtl.resident.streams.drop(h); + if (h.draft) |d| d.deinit(); h.config.deinit(h.gpa); h.metal.deinit(); h.model.deinit(); @@ -35,6 +42,13 @@ const Host = struct { const cannot = "the native LFM2 engine cannot load this checkpoint ({s})"; +/// A refused open's backend, model and host, freed before it reports why. +fn release(h: *Host) void { + h.metal.deinit(); + h.model.deinit(); + h.gpa.destroy(h); +} + 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(); @@ -65,18 +79,36 @@ pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Ope const chunk = 2048; h.metal = try q.backend.Metal.init(gpa, model, .{ .capacity = window, .chunk = chunk, .streams = @max(o.lanes, 2) }); errdefer h.metal.deinit(); - const per_stream = q.state.cacheBytes(&model.config.attention, window); + h.draft = null; + h.wrap = null; + if (if (o.drafts) o.drafter else null) |dir| { + if (o.drafter_bits == 0) { + problem.* = "--drafter-bits 0: the LFM2 DSpark drafter runs its bf16 linears as 8-bit codes, quantized at load"; + release(h); + return null; + } + h.draft = q.dspark.Drafter.init(gpa, io, model, dir, .{ .capacity = window }) catch |err| { + problem.* = try std.fmt.allocPrint(a, "the native LFM2 drafter cannot load {s} ({s}): it takes LiquidAI/LFM2.5-24B-A2B-DSpark", .{ dir, @errorName(err) }); + release(h); + return null; + }; + } + errdefer if (h.draft) |d| d.deinit(); + // a stream's caches, and its drafter lane when one drafts + const per_stream = q.state.cacheBytes(&model.config.attention, window) + if (h.draft) |d| d.laneBytes(window) else 0; const served = try metal_admit.streams(a, "lfm2", per_stream, metal_admit.room(model.device), o, window, problem) orelse { - h.metal.deinit(); - model.deinit(); - gpa.destroy(h); + if (h.draft) |d| d.deinit(); + release(h); return null; }; h.metal.options.streams = @max(served, 2); - h.config = try lanes.Config.init(gpa, h.metal.facts(), q.state.window_rows, q.state.window_rows - 1); + if (h.draft) |d| h.wrap = try lanes.drafted.Drafted.init(gpa, h.metal.backend(), d.drafter()); + errdefer if (h.wrap) |*w| w.deinit(); + const facts = if (h.wrap) |*w| w.facts(h.metal.facts()) else h.metal.facts(); + h.config = try lanes.Config.init(gpa, facts, q.state.window_rows, q.state.window_rows - 1); errdefer h.config.deinit(gpa); h.clock = .{ .io = io }; - h.core = lanes.Engine.init(gpa, &h.config, h.metal.backend(), h.clock.clock()); + h.core = lanes.Engine.init(gpa, &h.config, if (h.wrap) |*w| w.backend() else h.metal.backend(), h.clock.clock()); errdefer h.core.deinit(); h.host = api.LaneHost.init(gpa, io, &h.core, .{ .lanes = served, .context_window = @intCast(window), .prefill_step = chunk, .structures = true }); h.cache = null; @@ -97,8 +129,8 @@ pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Ope if (plan.budget_bytes > 0) { const S = q.backend.Snaps; h.cache = api.prompt_cache.Store.init(gpa, .{ - .ptr = h.metal, - .vtable = &.{ .bytes = S.bytes, .save = S.save, .restore = S.restore, .drop = S.drop }, + .ptr = h, + .vtable = &.{ .bytes = Snaps.bytes, .save = Snaps.save, .restore = S.restore, .drop = Snaps.drop }, }, .{ .chunk = chunk }, plan.budget_bytes); h.host.cache = &h.cache.?; h.host.info_.prompt_cache = true; @@ -106,7 +138,37 @@ pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Ope errdefer if (h.cache) |*store| store.deinit(); try mtl.resident.streams.hold(h, io, "lfm2 moe"); errdefer mtl.resident.streams.drop(h); + if (h.draft) |d| if (h.resident) |*r| { + const held = try d.buffers(gpa); + defer gpa.free(held); + r.add(held); + }; h.host.keepalive_target = .{ .ctx = &h.warm, .tick = mtl.keepalive.Target.tick }; try h.host.start(); + if (h.draft) |d| std.log.info("lfm2: DSpark drafts from --drafter, {d} a round at most, {d}-bit codes", .{ d.config.block, d.options.bits }); return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; } + +/// The prompt cache's copies of a stream: the target's state at a mark (backend.zig), and the drafter's context there. +const Snaps = struct { + const S = q.backend.Snaps; + fn host(ptr: *anyopaque) *Host { + return @ptrCast(@alignCast(ptr)); + } + fn bytes(ptr: *anyopaque, at: u32) u64 { + const h = host(ptr); + return S.bytes(h.metal, at) + if (h.draft) |d| d.keptBytes(at) else 0; + } + fn save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { + const h = host(ptr); + const saved = try S.save(h.metal, owner, at); + if (h.draft) |d| d.keepLane(@ptrCast(@alignCast(owner.?)), @ptrCast(@alignCast(saved))) catch |err| { + S.drop(h.metal, saved); + return err; + }; + return saved; + } + fn drop(ptr: *anyopaque, saved: *anyopaque) void { + S.drop(host(ptr).metal, saved); + } +}; diff --git a/zig/src/native/metal.zig b/zig/src/native/metal.zig index 41ab5c453e..33533c3597 100644 --- a/zig/src/native/metal.zig +++ b/zig/src/native/metal.zig @@ -170,8 +170,8 @@ pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]c problem.* = why; return null; } - if (o.drafter != null and !dense and !isGemma4(o.model_type)) { - problem.* = "--drafter is supported for Qwen3.8-27B-class checkpoints (qwen3_5 with its own output head) and Gemma 4 dense"; + if (o.drafter != null and !dense and !isGemma4(o.model_type) and !std.mem.eql(u8, o.model_type, "lfm2_moe")) { + problem.* = "--drafter is supported for Qwen3.8-27B-class checkpoints (qwen3_5 with its own output head), Gemma 4 dense and LFM2.5-24B-A2B"; return null; } if (dense) return openQwen27(a, gpa, io, o, problem); diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 7797558179..4a278138ac 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -84,6 +84,8 @@ test { _ = @import("families/lfm2_moe/config.zig"); _ = @import("families/lfm2_moe/backend.zig"); _ = @import("families/lfm2_moe/state.zig"); + _ = @import("families/lfm2_moe/dspark_config.zig"); + _ = @import("families/lfm2_moe/dspark_quant.zig"); _ = @import("families/lfm2_moe/dspark_test.zig"); } diff --git a/zig/tests/contract/metal_lfm2_attn.zig b/zig/tests/contract/metal_lfm2_attn.zig new file mode 100644 index 0000000000..799b337c59 --- /dev/null +++ b/zig/tests/contract/metal_lfm2_attn.zig @@ -0,0 +1,182 @@ +//! LFM2's block attention (lfm2_attn_block, then lfm2_attn_block_merge): every row of a drafter's block sees all keys. +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 qh = 32; +const kvh = 8; +const hd = 64; +const subs = 8; // the drafter's sub-spans a 4,096-key chunk (state.subs) +const chunks = 32; +const chunk_keys = 4096; +pub const name = "metal lfm2 block attention"; + +/// rows: a block's rows; k: the keys every row sees (a context and the block past it). +const shapes = [_]c.Shape{ .{ .rows = 1, .n = qh * hd, .k = 9 }, .{ .rows = 3, .n = qh * hd, .k = 700 }, .{ .rows = 9, .n = qh * hd, .k = 4500 } }; + +fn depth(s: c.Shape) usize { + return s.k; +} + +pub fn spec() c.Spec { + return .{ .name = name, .accum = .none, .depth = depth, .out = .bf16, .served = &shapes, .nan = .{ .stated = "a NaN key of a kv head makes its query heads' outputs NaN, the others stay finite" }, .sums = .{ .f16 = 0, .bf16 = 0, .why = "the check bounds each output by its bf16 probabilities' rounding, not a sum's bits" } }; +} + +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, + pipes: [2]mtl.Pipeline, + shape: c.Shape = undefined, + cap: usize = 0, + nan: bool = false, + q: []u16 = &.{}, + k: []u16 = &.{}, + v: []u16 = &.{}, + args: [5]rig_.Ref = undefined, + y: rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue) !Probe { + const first = try rig_.pipeline(device, sources.lfm2, "lfm2_attn_block"); + errdefer first.deinit(); + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .pipes = .{ first, try rig_.pipeline(device, sources.lfm2, "lfm2_attn_block_merge") } }; + } + + fn drop(p: *Probe) void { + inline for (.{ "q", "k", "v" }) |f| { + p.a.free(@field(p, f)); + @field(p, f) = &.{}; + } + p.rig.reset(); + } + + pub fn deinit(p: *Probe) void { + p.drop(); + p.rig.deinit(); + for (p.pipes) |x| x.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 { + if (s.n != qh * hd or s.rows == 0 or s.rows > 16 or s.k < s.rows or s.k > chunks * chunk_keys) return false; + p.drop(); + p.shape = s; + p.nan = f == .nonfinite; + p.cap = (s.k + 63) / 64 * 64 + 32; // keys past `k` hold finite values, as a lane's cache does + const padded = (s.rows + 7) / 8 * 8 + 8; // the kernel reads a tile's 8 rows whatever its row count + p.q = try p.a.alloc(u16, padded * qh * hd); + p.k = try p.a.alloc(u16, kvh * p.cap * hd); + p.v = try p.a.alloc(u16, kvh * p.cap * hd); + for (p.q, 0..) |*x, i| x.* = value(i, seed, 1); + for (p.k, 0..) |*x, i| x.* = value(i, seed +% 1, if (f == .stress) 4 else 1); + for (p.v, 0..) |*x, i| x.* = value(i, seed +% 2, 2); + if (p.nan) p.k[(0 * p.cap + s.k / 2) * hd + 5] = 0x7fc0; // kv head 0, a visible key + p.args[0] = try p.rig.input(std.mem.sliceAsBytes(p.q)); + p.args[1] = try p.rig.input(std.mem.sliceAsBytes(p.k)); + p.args[2] = try p.rig.input(std.mem.sliceAsBytes(p.v)); + const Args = extern struct { len: u32, cap: u32, scale: f32, subs: u32 }; + const args = Args{ .len = @intCast(s.k - 1), .cap = @intCast(p.cap), .scale = 1.0 / @sqrt(@as(f32, hd)), .subs = subs }; + p.args[3] = try p.rig.input(std.mem.asBytes(&args)); + const n: u32 = @intCast(s.rows); + p.args[4] = try p.rig.input(std.mem.asBytes(&n)); + p.y = try p.rig.output(p.count() * 2); + return true; + } + + pub fn count(p: *Probe) usize { + return p.shape.rows * qh * hd; + } + + pub fn run(p: *Probe) !void { + const rows = p.shape.rows; + const partials = @max(rows, 8) * subs * qh * chunks; + const pm = try p.rig.scratch(partials * 4); + const pl = try p.rig.scratch(partials * 4); + const po = try p.rig.scratch(partials * hd * 4); + const pass = p.rig.begin(); + const e = pass.e; + e.setPipeline(p.pipes[0]); + for (p.args[0..4], 0..) |r, i| e.setBuffer(r.buf, r.off, i); + for ([_]rig_.Ref{ pm, pl, po, p.args[4] }, 4..) |r, i| e.setBuffer(r.buf, r.off, i); + const used = (p.shape.k + chunk_keys - 1) / chunk_keys * subs; + e.dispatchGroups(mtl.Size.of((rows + 7) / 8, kvh, used), mtl.Size.of(128, 1, 1)); + e.setPipeline(p.pipes[1]); + for ([_]rig_.Ref{ pm, pl, po, p.args[3], p.y }, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + e.dispatchGroups(mtl.Size.of(qh, rows, 1), mtl.Size.of(32, 1, 1)); + try p.rig.finish(pass); + } + + pub fn reference(_: *Probe, want: []f64, mag: []f64) !void { + @memset(want, 0); // the check computes its own references + @memset(mag, 0); + } + + /// Every output against float64 attention: within 2^-7 of its values' largest magnitude and 2^-8 of itself. + pub fn check(p: *Probe, _: g.Gate) !?[]const u8 { + const s = p.shape; + const out = std.mem.bytesAsSlice(u16, p.rig.bytes(p.y, p.count() * 2)); + const w = try p.a.alloc(f64, s.k); + defer p.a.free(w); + for (0..s.rows) |r| for (0..qh) |h| { + const kv = h / (qh / kvh); + const nan_head = p.nan and kv == 0; + var top: f64 = -std.math.inf(f64); + for (0..s.k) |j| { + var dot: f64 = 0; + for (0..hd) |d| dot += c.fromBf16(p.q[(r * qh + h) * hd + d]) * c.fromBf16(p.k[(kv * p.cap + j) * hd + d]); + w[j] = dot / @sqrt(@as(f64, hd)); + top = @max(top, w[j]); + } + var sum: f64 = 0; + for (w) |*x| { + x.* = @exp(x.* - top); + sum += x.*; + } + for (0..hd) |d| { + var o: f64 = 0; + var big: f64 = 0; + for (0..s.k) |j| { + const vv: f64 = c.fromBf16(p.v[(kv * p.cap + j) * hd + d]); + o += w[j] / sum * vv; + big = @max(big, @abs(vv)); + } + const got: f64 = c.fromBf16(out[(r * qh + h) * hd + d]); + if (nan_head) { + if (!std.math.isNan(got)) return "a head behind a 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; + } +}; + +/// Gate 5 for the block attention's two launches on every Apple chip we serve and 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 }{ .{ "metal lfm2 block attention", "lfm2_attn_block", 128, 18688 }, .{ "metal lfm2 block attention merge", "lfm2_attn_block_merge", 32, 0 } }; + for (launches) |l| { + if (only) |s| if (std.mem.indexOf(u8, l[0], s) == null) continue; + const d = contract.limits.MetalLaunch{ .name = l[0], .threads = l[2], .tg_memory = l[3], .offline = false }; + var notes: std.ArrayList(u8) = .empty; + defer notes.deinit(a); + for (contract.limits.apple) |chip| problems += try contract.limits.metal(d, chip, ¬es, a); + std.debug.print("{s}", .{notes.items}); + problems += try @import("metal_limits.zig").onThisGpu(a, device, .{ .decl = d, .function = l[1], .here = sources.lfm2, .pre_m5 = "" }, false); + } + return problems; +} diff --git a/zig/tests/contract/metal_lfm2_dspark.zig b/zig/tests/contract/metal_lfm2_dspark.zig new file mode 100644 index 0000000000..3f2ae29907 --- /dev/null +++ b/zig/tests/contract/metal_lfm2_dspark.zig @@ -0,0 +1,221 @@ +//! LFM2's DSpark Markov step (scan then pick): each block row's draft from its base logits and the previous token. +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"); + +pub const Which = enum { markov_step }; +const rank = 256; +const vocab = 128000; +const span = 512; // tokens a scan threadgroup takes (the kernel's MARKOV_SPAN) +const groups = (vocab + span - 1) / span; +const table = 1000; // W1's rows in the check: ids are inside it or past the vocabulary (a zero row) +const names = [_][]const u8{"metal lfm2 markov step"}; + +const step_shapes = [_]c.Shape{ .{ .rows = 1, .n = vocab, .k = rank }, .{ .rows = 4, .n = vocab, .k = rank }, .{ .rows = 9, .n = vocab, .k = rank } }; + +fn none(_: c.Shape) usize { + return 0; +} + +fn depth(_: c.Shape) usize { + return rank; +} + +pub fn spec(w: Which) c.Spec { + return switch (w) { + .markov_step => .{ .name = names[0], .accum = .none, .depth = depth, .out = .raw, .served = &step_shapes, .nan = .{ .stated = "a row with a NaN score is flagged nonfinite" }, .sums = .{ .f16 = 0, .bf16 = 0, .why = "picks are tokens: the check bounds each pick's score by gamma(256), not sum bits" } }, + }; +} + +/// 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, + which: Which, + pipes: [2]mtl.Pipeline, + rows: usize = 0, + src: []u16 = &.{}, // the base logits rows + w1: []u16 = &.{}, + codes: []u8 = &.{}, + scales: []u16 = &.{}, + biases: []u16 = &.{}, + ids: []u32 = &.{}, + args: [8]rig_.Ref = undefined, + y: rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, w: Which) !Probe { + const kernels: [2][:0]const u8 = switch (w) { + .markov_step => .{ "lfm2_markov_scan", "lfm2_markov_pick" }, + }; + const first = try rig_.pipeline(device, sources.lfm2, kernels[0]); + errdefer first.deinit(); + const second = try rig_.pipeline(device, sources.lfm2, kernels[1]); + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .which = w, .pipes = .{ first, second } }; + } + + fn drop(p: *Probe) void { + inline for (.{ "src", "w1", "codes", "scales", "biases", "ids" }) |f| { + p.a.free(@field(p, f)); + @field(p, f) = &.{}; + } + p.rig.reset(); + } + + pub fn deinit(p: *Probe) void { + p.drop(); + p.rig.deinit(); + for (p.pipes) |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 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(); + p.rows = s.rows; + if (s.n != vocab or s.k != rank or s.rows == 0) return false; + // each stream's step row sits a block (two rows here) apart, as a hold's rows of one step do + p.src = try p.a.alloc(u16, 2 * s.rows * vocab); + fill(p.src, seed, f); + p.w1 = try p.a.alloc(u16, table * rank); + fill(p.w1, seed +% 5, .random); + p.codes = try p.a.alloc(u8, vocab * rank); + for (p.codes, 0..) |*q, i| q.* = @truncate((@as(u64, i) *% 2654435761 +% seed) >> 11); + p.scales = try p.a.alloc(u16, vocab * rank / 64); + p.biases = try p.a.alloc(u16, vocab * rank / 64); + for (p.scales, p.biases, 0..) |*sc, *bi, i| { + sc.* = c.toBf16(@as(f32, @floatFromInt(i % 7 + 1)) / 4096); + bi.* = c.toBf16(@as(f32, @floatFromInt(i % 5)) / 64 - 0.03); + } + // the previous token every second word (the picks' layout): inside W1, its last row, or past the vocabulary + p.ids = try p.a.alloc(u32, 2 * s.rows); + for (p.ids, 0..) |*id, i| id.* = if (i % 2 == 1) 0xdead else switch ((i / 2) % 3) { + 0 => @intCast((@as(u64, i) * 37 + seed) % table), + 1 => table - 1, + else => vocab + @as(u32, @intCast(i)), + }; + p.args[0] = try p.rig.input(std.mem.sliceAsBytes(p.src)); + p.args[1] = try p.rig.input(std.mem.sliceAsBytes(p.w1)); + p.args[2] = try p.rig.input(p.codes); + p.args[3] = try p.rig.input(std.mem.sliceAsBytes(p.ids)); + const args = [4]u32{ 2 * vocab, rank, vocab, 2 }; + p.args[4] = try p.rig.input(std.mem.asBytes(&args)); + p.args[5] = try p.rig.input(std.mem.sliceAsBytes(p.scales)); + p.args[6] = try p.rig.input(std.mem.sliceAsBytes(p.biases)); + const n: u32 = groups; + p.args[7] = try p.rig.input(std.mem.asBytes(&n)); + p.y = try p.rig.output(p.count()); + return true; + } + + pub fn count(p: *Probe) usize { + return switch (p.which) { + .markov_step => p.rows * 8, + }; + } + + pub fn run(p: *Probe) !void { + const pass = p.rig.begin(); + const e = pass.e; + const best = try p.rig.scratch(p.rows * groups * 4); + const at = try p.rig.scratch(p.rows * groups * 4); + e.setPipeline(p.pipes[0]); + for (p.args[0..5], 0..) |r, i| e.setBuffer(r.buf, r.off, i); + e.setBuffer(best.buf, best.off, 5); + e.setBuffer(at.buf, at.off, 6); + e.setBuffer(p.args[5].buf, p.args[5].off, 7); + e.setBuffer(p.args[6].buf, p.args[6].off, 8); + e.dispatchGroups(mtl.Size.of(groups, p.rows, 1), mtl.Size.of(256, 1, 1)); + e.setPipeline(p.pipes[1]); + e.setBuffer(best.buf, best.off, 0); + e.setBuffer(at.buf, at.off, 1); + e.setBuffer(p.args[7].buf, p.args[7].off, 2); + e.setBuffer(p.y.buf, p.y.off, 3); + e.dispatchGroups(mtl.Size.of(p.rows, 1, 1), mtl.Size.of(256, 1, 1)); + try p.rig.finish(pass); + } + + pub fn reference(_: *Probe, want: []f64, mag: []f64) !void { + @memset(mag, 0); + @memset(want, 0); // the step and the ids check themselves (check) + } + + /// Each stream's pick against float64 scores with the kernel's bf16 roundings: the best, but for near ties. + pub fn check(p: *Probe, _: g.Gate) !?[]const u8 { + const out = p.rig.bytes(p.y, p.count()); + for (0..p.rows) |s| { + const token = std.mem.readInt(u32, out[8 * s ..][0..4], .little); + const flag = std.mem.readInt(u32, out[8 * s + 4 ..][0..4], .little); + const id = p.ids[2 * s]; + var top: f64 = -std.math.inf(f64); + var mine: ?f64 = null; + var nan = false; + var bound: f64 = 0; + for (0..vocab) |v| { + var dot: f64 = 0; + var size: f64 = 0; + if (id < table) for (0..rank) |i| { + const k = (v * rank + i) / 64; + const w = @as(f64, c.fromBf16(p.scales[k])) * @as(f64, @floatFromInt(p.codes[v * rank + i])) + c.fromBf16(p.biases[k]); + dot += w * c.fromBf16(p.w1[id * rank + i]); + size += @abs(w * c.fromBf16(p.w1[id * rank + i])); + }; + const base: f32 = c.fromBf16(p.src[2 * s * vocab + v]); + const score: f64 = c.fromBf16(c.toBf16(base + c.fromBf16(c.toBf16(@floatCast(dot))))); + if (std.math.isNan(score)) nan = true; + if (score > top) top = score; + if (v == token) mine = score; + // two bf16 steps of the score and the fp32 sum's own error + bound = @max(bound, @abs(score) * 0x1p-7 + c.gamma(rank, 0x1p-24) * size); + } + if (nan) { + if (flag != 1) return "a row with a NaN score is not flagged"; + continue; + } + if (flag != 0) return "a finite row is flagged nonfinite"; + const got = mine orelse return "the pick is past the vocabulary"; + if (got < top - bound) return "the pick's score is below the row's best by more than its bound"; + } + return null; + } +}; + +/// Gate 5 for the scan and pick on every Apple chip and this GPU: 256 threads each, 1,088 and 3,072 threadgroup bytes. +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 }{ .{ "metal lfm2 markov scan", "lfm2_markov_scan", 1088, 256 }, .{ "metal lfm2 markov pick", "lfm2_markov_pick", 3072, 256 } }; + for (launches) |l| { + if (only) |s| if (std.mem.indexOf(u8, l[0], s) == null) continue; + const d = contract.limits.MetalLaunch{ .name = l[0], .threads = l[3], .tg_memory = l[2], .offline = false }; + var notes: std.ArrayList(u8) = .empty; + defer notes.deinit(a); + for (contract.limits.apple) |chip| problems += try contract.limits.metal(d, chip, ¬es, a); + std.debug.print("{s}", .{notes.items}); + problems += try @import("metal_limits.zig").onThisGpu(a, device, .{ .decl = d, .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 70bedb69dc..6f2e5d76b6 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -26,6 +26,8 @@ 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_dspark = @import("metal_lfm2_dspark.zig"); +const lfm2_attn = @import("metal_lfm2_attn.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]"; @@ -175,6 +177,21 @@ pub fn main(init: std.process.Init) !u8 { tap = try lfm2_tap.Probe.init(a, device, queue); try kernels.append(a, tap.?.kernel()); } + var lfm2_list: [1]lfm2_dspark.Probe = undefined; + var lfm2_made: usize = 0; + defer for (lfm2_list[0..lfm2_made]) |*p| p.deinit(); + for (std.enums.values(lfm2_dspark.Which)) |w| { + if (!wanted(o, lfm2_dspark.spec(w).name)) continue; + lfm2_list[lfm2_made] = try lfm2_dspark.Probe.init(a, device, queue, w); + lfm2_made += 1; + } + for (lfm2_list[0..lfm2_made]) |*p| try kernels.append(a, p.kernel()); + var block_attn: ?lfm2_attn.Probe = null; + defer if (block_attn) |*p| p.deinit(); + if (wanted(o, lfm2_attn.name)) { + block_attn = try lfm2_attn.Probe.init(a, device, queue); + try kernels.append(a, block_attn.?.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(); @@ -317,6 +334,8 @@ pub fn main(init: std.process.Init) !u8 { 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); + lim += try lfm2_dspark.limits(a, device, o.only); + lim += try lfm2_attn.limits(a, device, o.only); } const t2 = std.Io.Clock.awake.now(init.io); const ms = struct { diff --git a/zig/tests/lfm2_dspark.zig b/zig/tests/lfm2_dspark.zig new file mode 100644 index 0000000000..09c7f68a52 --- /dev/null +++ b/zig/tests/lfm2_dspark.zig @@ -0,0 +1,301 @@ +//! LFM2's DSpark drafter on the GPU against an fp32 host reference of the released model, on the real weights. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const q = tf.lfm2_moe; +const c = q.config; +const st = q.state; +const ck = tf.checkpoint; +const quant = q.dspark_quant; + +const hd = c.head_dim; +const kw = c.kv_heads * hd; +const qw = c.query_heads * hd; + +/// The drafter's tensors as fp32 on the host, in the checkpoint's own layout (interleaved RoPE pairs). +const Ref = struct { + gpa: std.mem.Allocator, + check: ck.Checkpoint, + cfg: q.dspark_config.Config, + + fn get(r: *const Ref, name: []const u8) ![]f32 { + const t = try r.check.get(name); + const src = t.host(u16); + const out = try r.gpa.alloc(f32, src.len); + for (out, src) |*o, v| o.* = quant.toF32(v); + return out; + } + + fn layer(r: *const Ref, buf: []u8, i: usize, part: []const u8) ![]f32 { + return r.get(try std.fmt.bufPrint(buf, "layers.{d}.{s}.weight", .{ i, part })); + } +}; + +/// y[n] = W[n, k] x[k] for each of `rows` rows, on every core. +fn matmul(gpa: std.mem.Allocator, w: []const f32, x: []const f32, y: []f32, n: usize, k: usize, rows: usize) !void { + const Job = struct { + fn run(w_: []const f32, x_: []const f32, y_: []f32, n_: usize, k_: usize, rows_: usize, from: usize, to: usize) void { + for (0..rows_) |r| for (from..to) |col| { + var s: f32 = 0; + const wr = w_[col * k_ ..][0..k_]; + const xr = x_[r * k_ ..][0..k_]; + for (wr, xr) |a, b| s += a * b; + y_[r * n_ + col] = s; + }; + } + }; + const threads = @min(std.Thread.getCpuCount() catch 1, 16); + const per = (n + threads - 1) / threads; + const hs = try gpa.alloc(std.Thread, threads); + defer gpa.free(hs); + var made: usize = 0; + var at: usize = 0; + while (at < n) : (at += per) { + hs[made] = try std.Thread.spawn(.{}, Job.run, .{ w, x, y, n, k, rows, at, @min(n, at + per) }); + made += 1; + } + for (hs[0..made]) |h| h.join(); +} + +fn rms(x: []f32, gain: []const f32) void { + var ss: f32 = 0; + for (x) |v| ss += v * v; + const inv = 1 / @sqrt(ss / @as(f32, @floatFromInt(x.len)) + c.eps); + for (x, gain) |*v, g| v.* = v.* * inv * g; +} + +/// Interleaved RoPE at `pos` over each head of `x` (pairs 2i, 2i + 1). +fn rope(x: []f32, pos: u64, theta: f32) void { + var h: usize = 0; + while (h < x.len) : (h += hd) for (0..hd / 2) |i| { + const angle = @as(f32, @floatFromInt(pos)) * std.math.pow(f32, theta, -@as(f32, @floatFromInt(i)) / (hd / 2)); + const a = x[h + 2 * i]; + const b = x[h + 2 * i + 1]; + x[h + 2 * i] = a * @cos(angle) - b * @sin(angle); + x[h + 2 * i + 1] = a * @sin(angle) + b * @cos(angle); + }; +} + +/// The target's embedding row `id`, dequantized from its 6-bit codes (MLX's scale * q + bias). +fn embedRow(m: *q.Model, id: u32, out: []f32) void { + const e = m.weights.embedding; + const codes = (e.w.contents() + e.w_off)[@as(usize, id) * c.hidden * 6 / 8 ..][0 .. c.hidden * 6 / 8]; + const sc = @as([*]const u16, @ptrCast(@alignCast(e.scales.contents() + e.s_off)))[@as(usize, id) * c.hidden / 64 ..][0 .. c.hidden / 64]; + const bi = @as([*]const u16, @ptrCast(@alignCast(e.biases.contents() + e.b_off)))[@as(usize, id) * c.hidden / 64 ..][0 .. c.hidden / 64]; + for (0..c.hidden / 4) |u| { + const w = codes[u * 3 ..][0..3]; + const qs = [4]u32{ w[0] & 0x3f, (w[0] >> 6) | ((@as(u32, w[1]) & 0x0f) << 2), (w[1] >> 4) | ((@as(u32, w[2]) & 0x03) << 4), w[2] >> 2 }; + for (qs, 0..) |qv, j| { + const g = (u * 4 + j) / 64; + out[u * 4 + j] = quant.toF32(sc[g]) * @as(f32, @floatFromInt(qv)) + quant.toF32(bi[g]); + } + } +} + +/// GPU bf16 row against the reference's row + bias: cosine and worst difference. +fn compare(label: []const u8, gpu: []const u16, row: []const f32, bias: []const f32) void { + var dot: f64 = 0; + var gg: f64 = 0; + var rr: f64 = 0; + var worst: f32 = 0; + for (gpu, row, bias) |g, b, m2| { + const x = quant.toF32(g); + const y = b + m2; + dot += x * y; + gg += x * x; + rr += y * y; + worst = @max(worst, @abs(x - y)); + } + std.debug.print("{s}: cosine {d:.5}, worst difference {d:.3}\n", .{ label, dot / @sqrt(gg * rr), worst }); +} + +fn absorb(d: *q.dspark.Drafter, piece: q.dspark_block.Piece) !void { + const run = q.dspark_block.begin(d); + try q.dspark_block.absorbPass(d, run.e, &.{piece}); + try q.dspark_block.end(run.cb, run.e); +} + +pub fn main(init: std.process.Init) !void { + const gpa = init.gpa; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 4 or args.len > 6) return error.ExpectedModelDrafterPromptLengthRowsBits; + const bits: u8 = if (args.len == 6) try std.fmt.parseInt(u8, args[5], 10) else 8; + const length = try std.fmt.parseInt(usize, args[3], 10); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const m = try q.Model.load(gpa, init.io, args[1]); + defer m.deinit(); + const d = try q.dspark.Drafter.init(gpa, init.io, m, args[2], .{ .capacity = length + 64, .bits = bits }); + defer d.deinit(); + const cfg = d.config; + // the compared block's rows (a hold of depth d runs a d-row block) + const block: u32 = if (args.len >= 5) try std.fmt.parseInt(u32, args[4], 10) else cfg.block; + if (block == 0 or block > cfg.block) return error.BadBlockRows; + var scratch = try st.Scratch.init(gpa, m.device, &m.config.attention, length); + defer scratch.deinit(); + const width = cfg.tap_count * c.hidden; + var taps: st.Taps = .{ .count = cfg.tap_count, .buf = try scratch.memory.alloc(length * width * 2) }; + @memcpy(taps.layers[0..cfg.tap_count], cfg.taps[0..cfg.tap_count]); + scratch.taps = taps; + var cache = try st.Cache.init(gpa, m.device, &m.config.attention, length + 16); + defer cache.deinit(); + const ids = try gpa.alloc(u32, length); + defer gpa.free(ids); + for (ids, 0..) |*id, i| id.* = @intCast((i * 7919 + 13) % 120000); + try q.forward.run(m, &scratch, &.{.{ .cache = &cache, .rows = length }}, ids, false, .last); + // the context: rows [0, length - 1); the pending token is the last id, at position length - 1 + const ctx = length - 1; + const lane = try d.newLane(length + 32); + defer d.dropLane(lane); + var at: usize = 0; + while (at < ctx) : (at += q.dspark.absorb_rows) { + const rows = @min(q.dspark.absorb_rows, ctx - at); + try absorb(d, .{ .lane = lane, .src = .{ .buf = taps.buf, .off = at * width * 2 }, .rows = rows, .p0 = at }); + } + lane.end = ctx; + try q.dspark_block.hold(d, &.{.{ .lane = lane, .pending = ids[ctx], .anchor = ctx, .depth = block }}); + var held: [16]u32 = undefined; + @memcpy(held[0..block], lane.held[0..block]); + const gpu = held[0..block]; + for ([_]usize{ 1, 4, 16 }) |rows| { + var best: f64 = std.math.floatMax(f64); + for (0..5) |_| { + const t0 = std.Io.Clock.awake.now(init.io); + try absorb(d, .{ .lane = lane, .src = .{ .buf = taps.buf, .off = (ctx - rows) * width * 2 }, .rows = rows, .p0 = ctx - rows }); + best = @min(best, @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(init.io)).toNanoseconds())) / 1e6); + } + std.debug.print("absorb of {d} rows: {d:.2} ms (best of 5)\n", .{ rows, best }); + } + const depths = [_]u32{ 1, 3, cfg.block }; + for (depths) |depth| { + var best: f64 = std.math.floatMax(f64); + for (0..5) |_| { + const t0 = std.Io.Clock.awake.now(init.io); + try q.dspark_block.hold(d, &.{.{ .lane = lane, .pending = ids[ctx], .anchor = ctx, .depth = depth }}); + best = @min(best, @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(init.io)).toNanoseconds())) / 1e6); + } + std.debug.print("hold at depth {d}: {d:.2} ms (best of 5)\n", .{ depth, best }); + } + try q.dspark_block.hold(d, &.{.{ .lane = lane, .pending = ids[ctx], .anchor = ctx, .depth = block }}); + // the host reference, fp32 throughout + var r: Ref = .{ .gpa = gpa, .check = ck.Checkpoint.init(gpa), .cfg = cfg }; + defer r.check.deinit(); + const path = try std.fs.path.joinZ(gpa, &.{ args[2], "model.safetensors" }); + defer gpa.free(path); + try r.check.addFile(m.device, path, ""); + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + r.gpa = arena.allocator(); + const a = r.gpa; + const n = block; + const tap_rows = taps.buf.slice(u16, length * width); + const feats = try a.alloc(f32, ctx * width); + for (feats, tap_rows[0 .. ctx * width]) |*f, v| f.* = quant.toF32(v); + const hctx = try a.alloc(f32, ctx * c.hidden); + try matmul(gpa, try r.get("fc.weight"), feats, hctx, c.hidden, width, ctx); + const hn = try r.get("hidden_norm.weight"); + for (0..ctx) |j| rms(hctx[j * c.hidden ..][0..c.hidden], hn); + const h = try a.alloc(f32, n * c.hidden); + for (0..n) |i| embedRow(m, if (i == 0) ids[ctx] else cfg.mask, h[i * c.hidden ..][0..c.hidden]); + const keys = ctx + n; + var name: [96]u8 = undefined; + const x = try a.alloc(f32, n * c.hidden); + for (0..cfg.layers) |li| { + const kmat = try r.layer(&name, li, "self_attn.k_proj"); + const vmat = try r.layer(&name, li, "self_attn.v_proj"); + const qn = try r.layer(&name, li, "self_attn.q_norm"); + const kn = try r.layer(&name, li, "self_attn.k_norm"); + const kk = try a.alloc(f32, keys * kw); + const vv = try a.alloc(f32, keys * kw); + try matmul(gpa, kmat, hctx, kk, kw, c.hidden, ctx); + try matmul(gpa, vmat, hctx, vv, kw, c.hidden, ctx); + @memcpy(x, h); + for (0..n) |i| rms(x[i * c.hidden ..][0..c.hidden], try r.layer(&name, li, "input_layernorm")); + const qq = try a.alloc(f32, n * qw); + try matmul(gpa, try r.layer(&name, li, "self_attn.q_proj"), x, qq, qw, c.hidden, n); + try matmul(gpa, kmat, x, kk[ctx * kw ..], kw, c.hidden, n); + try matmul(gpa, vmat, x, vv[ctx * kw ..], kw, c.hidden, n); + for (0..keys) |j| for (0..c.kv_heads) |kh| { + const row = kk[j * kw + kh * hd ..][0..hd]; + rms(row, kn); + rope(row, j, cfg.theta); + }; + for (0..n) |i| for (0..c.query_heads) |qh| { + const row = qq[i * qw + qh * hd ..][0..hd]; + rms(row, qn); + rope(row, ctx + i, cfg.theta); + }; + const o = try a.alloc(f32, n * qw); + const sc = try a.alloc(f32, keys); + for (0..n) |i| for (0..c.query_heads) |qh| { + const kh = qh / (c.query_heads / c.kv_heads); + var top: f32 = -std.math.inf(f32); + for (0..keys) |j| { + var s: f32 = 0; + for (0..hd) |t| s += qq[i * qw + qh * hd + t] * kk[j * kw + kh * hd + t]; + sc[j] = s / @sqrt(@as(f32, hd)); + top = @max(top, sc[j]); + } + var sum: f32 = 0; + for (sc) |*s| { + s.* = @exp(s.* - top); + sum += s.*; + } + for (0..hd) |t| { + var acc: f32 = 0; + for (0..keys) |j| acc += sc[j] * vv[j * kw + kh * hd + t]; + o[i * qw + qh * hd + t] = acc / sum; + } + }; + const y = try a.alloc(f32, n * c.hidden); + try matmul(gpa, try r.layer(&name, li, "self_attn.o_proj"), o, y, c.hidden, qw, n); + for (h, y) |*v, w| v.* += w; + @memcpy(x, h); + for (0..n) |i| rms(x[i * c.hidden ..][0..c.hidden], try r.layer(&name, li, "post_attention_layernorm")); + const gate = try a.alloc(f32, n * cfg.intermediate); + const up = try a.alloc(f32, n * cfg.intermediate); + try matmul(gpa, try r.layer(&name, li, "mlp.gate_proj"), x, gate, cfg.intermediate, c.hidden, n); + try matmul(gpa, try r.layer(&name, li, "mlp.up_proj"), x, up, cfg.intermediate, c.hidden, n); + for (gate, up) |*g, u| g.* = g.* / (1 + @exp(-g.*)) * u; + try matmul(gpa, try r.layer(&name, li, "mlp.down_proj"), gate, y, c.hidden, cfg.intermediate, n); + for (h, y) |*v, w| v.* += w; + } + const fin = try r.get("norm.weight"); + for (0..n) |i| rms(h[i * c.hidden ..][0..c.hidden], fin); + const head = try a.alloc(f32, c.vocab * c.hidden); + for (0..c.vocab) |v| embedRow(m, @intCast(v), head[v * c.hidden ..][0..c.hidden]); + const base = try a.alloc(f32, n * c.vocab); + try matmul(gpa, head, h, base, c.vocab, c.hidden, n); + const w1 = try r.get("markov_head.markov_w1.weight"); + const w2 = try r.get("markov_head.markov_w2.weight"); + const bias = try a.alloc(f32, c.vocab); + var prev = ids[ctx]; + var same: usize = 0; + var worst: f32 = 0; + for (0..n) |i| { + try matmul(gpa, w2, w1[@as(usize, prev) * cfg.rank ..][0..cfg.rank], bias, c.vocab, cfg.rank, 1); + const row = base[i * c.vocab ..][0..c.vocab]; + var best: usize = 0; + var second: usize = 1; + for (row, bias, 0..) |b, m2, v| { + const s = b + m2; + if (s > row[best] + bias[best]) { + second = best; + best = v; + } else if (v != best and s > row[second] + bias[second]) second = v; + } + const margin = (row[best] + bias[best]) - (row[second] + bias[second]); + if (i == 0) { // the GPU keeps the base logits (the Markov step adds its bias on the way to the pick) + const zero = try a.alloc(f32, c.vocab); + @memset(zero, 0); + compare("row 0 base logits", d.s.logits.slice(u16, c.vocab), row, zero); + } + const ok = gpu[i] == best; + same += @intFromBool(ok); + std.debug.print("step {d}: reference {d} (margin {d:.3}, second {d}), GPU {d} {s}\n", .{ i, best, margin, second, gpu[i], if (ok) "equal" else "DIFFERENT" }); + if (!ok) worst = @max(worst, margin); + prev = gpu[i]; // teacher-forced: step i + 1 follows the GPU's pick, so each step is compared alone + } + std.debug.print("the reference's largest margin over a different GPU pick: {d:.3}\n", .{worst}); + std.debug.print("{d} of {d} drafts equal the fp32 reference's picks on the same chain ({d}-token context)\n", .{ same, n, ctx }); +} From 092254c5d00864acecb10cc617868adf0d4fbdfc Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:30:52 +0100 Subject: [PATCH 17/46] lfm2: a DSpark hold is not waited on; the verify copies its drafts on the GPU A hold that fits in one pass signals a GPU event and returns without a wait. The next verify waits on that event on the GPU, and a new kernel, lfm2_draft_ids, copies each window's drafts from the drafter's picks into the verify's ids. The host reads the picks only after the verify, or when a grammar, an absorb, a keep or a release needs them. Multi-pass holds stay synchronous. The LFM2 backend takes held windows (Window.held) from a HeldSource the drafter gives, and Drafted.held_on_target passes windows without fillers to it as they are. lfm2_draft_ids registers with the kernel contract (metal lfm2 draft ids): each stream's first pick, or 0 for a token past the vocabulary, copied exactly; gate 5 against every Apple chip we serve and this GPU. --- zig/kernels/metal/lfm2/lfm2_dspark.metal | 8 +++ zig/src/families/lfm2_moe/backend.zig | 34 ++++++++-- zig/src/families/lfm2_moe/dspark.zig | 73 +++++++++++++++++++++- zig/src/families/lfm2_moe/dspark_block.zig | 27 +++++++- zig/src/families/lfm2_moe/forward.zig | 18 ++++++ zig/src/families/lfm2_moe/kernels.zig | 1 + zig/src/native/lfm2_moe.zig | 7 ++- zig/tests/contract/metal_lfm2_dspark.zig | 44 +++++++++++-- zig/tests/contract/metal_main.zig | 2 +- zig/tests/lfm2_dspark.zig | 6 +- 10 files changed, 201 insertions(+), 19 deletions(-) diff --git a/zig/kernels/metal/lfm2/lfm2_dspark.metal b/zig/kernels/metal/lfm2/lfm2_dspark.metal index 995d63a20a..050fd5e8cc 100644 --- a/zig/kernels/metal/lfm2/lfm2_dspark.metal +++ b/zig/kernels/metal/lfm2/lfm2_dspark.metal @@ -97,3 +97,11 @@ constant constexpr uint MARKOV_SPAN = 512; picks[2 * s + 1] = tn[0]; } } + +// A verify window's held drafts from the drafter's picks: picks[i * stride] into ids[i], a token past the vocabulary 0. +[[kernel]] void lfm2_draft_ids(const device uint* picks [[buffer(0)]], constant uint4& p [[buffer(1)]], + device uint* ids [[buffer(2)]], uint i [[thread_position_in_grid]]) { + if (i >= p.y) return; + const uint token = picks[size_t(i) * p.x]; + ids[i] = token < p.z ? token : 0u; +} diff --git a/zig/src/families/lfm2_moe/backend.zig b/zig/src/families/lfm2_moe/backend.zig index 44427d435a..8af90dc43b 100644 --- a/zig/src/families/lfm2_moe/backend.zig +++ b/zig/src/families/lfm2_moe/backend.zig @@ -10,6 +10,14 @@ const snapshot = @import("snapshot.zig"); const Model = @import("model.zig").Model; const be = lanes.backend; +/// Where a drafter holds a stream's drafts on the device, for verify windows that take them there (Window.held). +pub const HeldSource = struct { + ptr: *anyopaque, + get: *const fn (ptr: *anyopaque, s: *lanes.Stream) ?Device, + host: *const fn (ptr: *anyopaque, s: *lanes.Stream, out: []u32) anyerror!void, // where `get` gives none +}; +pub const Device = struct { src: mtl.Buffer, off: usize, stride: u32, wait: fwd.Wait }; + pub const Options = struct { capacity: usize, chunk: usize = 128, streams: u32 = 8 }; pub const Metal = struct { @@ -25,6 +33,8 @@ pub const Metal = struct { timed: usize = 0, /// Forwards run so far: a stream's tapped rows in the scratch are current while their gen matches. gen: u64 = 0, + /// A drafter whose drafts verify windows read on the device (null: windows carry host tokens only). + held_source: ?HeldSource = null, pub fn init(gpa: std.mem.Allocator, model: *Model, options: Options) !*Metal { const out = try gpa.create(Metal); @@ -246,20 +256,34 @@ pub const Metal = struct { const self = cast(ptr); var segments: [st.batch_rows]fwd.Segment = undefined; var ids: [st.batch_rows]u32 = undefined; + var held: [st.batch_rows]fwd.Held = undefined; + var count: usize = 0; + var wait: ?fwd.Wait = null; var total: usize = 0; if (windows.len > segments.len or windows.len != out.len) return error.InvalidLfmWindows; for (windows, 0..) |w, i| { const rows = w.rows(); - if (w.parents != null or w.held != 0 or rows > st.window_rows or total + rows > st.batch_rows or w.positions.len != rows) return error.UnsupportedLfmWindow; + if (w.held != 0) { // the drafter's picks, copied into the window's rows on the GPU + const source = self.held_source orelse return error.UnsupportedLfmWindow; + if (total + rows > st.batch_rows) return error.UnsupportedLfmWindow; + if (source.get(source.ptr, w.stream)) |dev| { + held[count] = .{ .row = total + 1, .src = dev.src, .off = dev.off, .stride = dev.stride, .count = w.held }; + count += 1; + wait = if (wait) |x| (if (x.value >= dev.wait.value) x else dev.wait) else dev.wait; + @memset(ids[total + 1 ..][0..w.held], 0); + } else try source.host(source.ptr, w.stream, ids[total + 1 ..][0..w.held]); + } + if (w.parents != null or rows > st.window_rows or total + rows > st.batch_rows or w.positions.len != rows) return error.UnsupportedLfmWindow; if (w.masks.len != 0 and w.masks.len != rows * mask_words) return error.InvalidLfmMasks; const cache = try self.cacheOf(w.stream); for (w.positions, 0..) |p, r| if (p != cache.len + r + 1) return error.InvalidLfmPosition; segments[i] = .{ .cache = cache, .rows = rows }; ids[total] = w.pending; - @memcpy(ids[total + 1 ..][0..w.tokens.len], w.tokens); + @memcpy(ids[total + 1 + w.held ..][0..w.tokens.len], w.tokens); total += rows; } - try self.run(segments[0..windows.len], ids[0..total], true, .all); + self.gen += 1; + try fwd.runHeld(self.model, &self.scratch, segments[0..windows.len], ids[0..total], true, .all, held[0..count], wait); 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 }; @@ -268,7 +292,7 @@ pub const Metal = struct { 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; + if (o.sampled.len != w.rows() or o.drafts.len != w.rows() - 1) 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. @@ -279,7 +303,7 @@ pub const Metal = struct { } t.* = try draw(self.gpa, logits[(base + r) * c.vocab ..][0..c.vocab], w.stream.sampling, p, mask); } - @memcpy(o.drafts, w.tokens); + @memcpy(o.drafts, self.scratch.ids.slice(u32, base + w.rows())[base + 1 ..][0 .. w.rows() - 1]); base += w.rows(); } } diff --git a/zig/src/families/lfm2_moe/dspark.zig b/zig/src/families/lfm2_moe/dspark.zig index 339071ac6c..850c1ce34c 100644 --- a/zig/src/families/lfm2_moe/dspark.zig +++ b/zig/src/families/lfm2_moe/dspark.zig @@ -9,6 +9,7 @@ const block = @import("dspark_block.zig"); const st = @import("state.zig"); const snapshot = @import("snapshot.zig"); const Model = @import("model.zig").Model; +const backend = @import("backend.zig"); const dr = lanes.drafter; const Stream = lanes.Stream; @@ -34,6 +35,7 @@ pub const Lane = struct { staged: u32 = 0, held: [16]u32 = @splat(0), held_len: u32 = 0, + slot: ?u32 = null, // its block's place in the last signalled hold, whose picks still hold its drafts }; /// The passes' working buffers, one pass at a time on the target's queue. @@ -75,11 +77,19 @@ pub const Drafter = struct { /// One stream's hold wall ms at full depth, and that over the block (the depth rule's cost a draft). hold_ms: f64 = 0, step_ms: f64 = 0, + event: mtl.Event = undefined, + signalled: u64 = 0, + pending: ?mtl.CommandBuffer = null, + pass: [pass_rows]*Lane = undefined, // the signalled hold's lanes and depths, in its blocks' order + depths: [pass_rows]u32 = undefined, + pass_len: usize = 0, pub fn init(gpa: std.mem.Allocator, io: std.Io, target: *Model, dir: []const u8, options: Options) !*Drafter { const d = try gpa.create(Drafter); errdefer gpa.destroy(d); d.* = .{ .gpa = gpa, .io = io, .target = target, .config = try dc.Config.read(gpa, io, dir), .w = undefined, .options = options, .s = undefined }; + d.event = try target.device.event(); + errdefer d.event.deinit(); d.w = try dw.load(gpa, target.device, io, dir, &d.config, options.bits); errdefer d.w.deinit(gpa); errdefer for (d.owned.items) |b| b.deinit(); @@ -90,6 +100,8 @@ pub const Drafter = struct { } pub fn deinit(d: *Drafter) void { + d.settle() catch {}; + d.event.deinit(); var it = d.lanes.valueIterator(); while (it.next()) |l| d.dropLane(l.*); d.lanes.deinit(d.gpa); @@ -232,11 +244,13 @@ pub const Drafter = struct { fn releaseFn(ptr: *anyopaque, s: *Stream) void { const d = cast(ptr); + d.unpass() catch {}; if (d.lanes.fetchRemove(s)) |e| d.dropLane(e.value); } fn heldFn(ptr: *anyopaque, streams: []const *Stream, out: []const []u32) anyerror!void { const d = cast(ptr); + try d.settle(); for (streams, out) |s, o| { const l = d.lanes.get(s) orelse return error.DraftStreamUnknown; if (o.len > l.held_len) return error.DraftHeldFewer; @@ -252,6 +266,7 @@ pub const Drafter = struct { /// Stream `s`'s context copied into `snap` when its lane stands at the snapshot's position (a prompt chunk's end). pub fn keepLane(d: *Drafter, s: *Stream, snap: *snapshot.Snap) !void { const l = d.lanes.get(s) orelse return; + try d.settle(); try d.flush(l); if (l.end != snap.at) return; const bytes = try d.gpa.alloc(u8, @intCast(d.keptBytes(snap.at))); @@ -283,6 +298,7 @@ pub const Drafter = struct { /// Rows from `start` replace the context there: verified rows wait on the host for the next hold, others run now. fn absorbFn(ptr: *anyopaque, items: []const dr.Absorb) anyerror!void { const d = cast(ptr); + try d.settle(); const width = d.tapBytes(); const pool = mtl.objc.Pool.push(); defer pool.pop(); @@ -297,6 +313,7 @@ pub const Drafter = struct { if (it.start > l.end + l.staged) try d.resumeLane(it.stream, l, it.start); // a prompt pass that resumed if (it.start + f.rows + d.config.block > l.cap) return error.DraftContextGap; l.held_len = 0; + l.slot = null; const src = mtl.Buffer{ .id = @ptrFromInt(f.buffer) }; // rows from `start` replace what was staged past it; rows before the absorbed end go to the GPU now const keep: u32 = if (it.start >= l.end) @intCast(it.start - l.end) else 0; @@ -343,6 +360,7 @@ pub const Drafter = struct { /// Each stream's chain of `depth` drafts after `pending` at `position`, its block's rows in a pass of pass_rows. fn holdFn(ptr: *anyopaque, items: []const dr.Hold) anyerror!void { const d = cast(ptr); + try d.unpass(); var at: usize = 0; while (at < items.len) { // as many streams as their blocks' rows fit a pass: each block is as long as the pass's deepest chain @@ -364,8 +382,59 @@ pub const Drafter = struct { if (l.end + l.staged != anchor or anchor + d.config.block > l.cap) return error.DraftContextGap; b.* = .{ .lane = l, .pending = it.pending, .anchor = anchor, .depth = it.depth }; } - try block.hold(d, blocks[0..group.len]); + if (group.len == items.len) { // one pass: signalled, not waited on + d.signalled += 1; + d.pending = try block.hold(d, blocks[0..group.len], .{ .event = d.event, .value = d.signalled }); + for (blocks[0..group.len], 0..) |b, bi| { + d.pass[bi] = b.lane; + d.depths[bi] = b.depth; + b.lane.slot = @intCast(bi); + b.lane.held_len = 0; + } + d.pass_len = group.len; + continue; + } + _ = try block.hold(d, blocks[0..group.len], null); + } + } + + /// The signalled hold waited on, its picks read into its lanes' held chains (nothing pending: nothing done). + fn settle(d: *Drafter) !void { + const cb = d.pending orelse return; + d.pending = null; + defer cb.release(); + cb.wait(); + if (cb.failure()) |msg| { + std.log.err("lfm2 dspark: a hold failed on the GPU: {s}", .{msg}); + return error.DraftGpuFailed; + } + block.readPicks(d, d.pass[0..d.pass_len], d.depths[0..d.pass_len]); + } + + /// The signalled hold settled and its lanes' picks forgotten (a lane may be dropped, or a hold may overwrite). + fn unpass(d: *Drafter) !void { + defer { + for (d.pass[0..d.pass_len]) |l| l.slot = null; + d.pass_len = 0; } + try d.settle(); + } + + /// Where stream `s`'s drafts from the last signalled hold lie in the picks, for a verify to copy on the GPU. + pub fn device(ptr: *anyopaque, s: *Stream) ?backend.Device { + const d = cast(ptr); + const l = d.lanes.get(s) orelse return null; + const slot = l.slot orelse return null; + return .{ .src = d.s.picks, .off = slot * 8, .stride = pass_rows * 2, .wait = .{ .event = d.event, .value = d.signalled } }; + } + + /// The target's verify reading the held drafts on the device, or as host tokens where it cannot. + pub fn heldSource(d: *Drafter) backend.HeldSource { + return .{ .ptr = d, .get = device, .host = hostHeld }; + } + + fn hostHeld(ptr: *anyopaque, s: *Stream, out: []u32) anyerror!void { + return heldFn(ptr, &.{s}, &.{out}); } /// Holds on an empty context: the pipelines compiled, then one-draft and full-depth holds timed (best of three). @@ -375,7 +444,7 @@ pub const Drafter = struct { var best: [2]f64 = @splat(std.math.floatMax(f64)); for (0..4) |i| for ([_]u32{ 1, d.config.block }, &best) |depth, *b| { const t0 = std.Io.Clock.awake.now(d.io); - try block.hold(d, &.{.{ .lane = l, .pending = 1, .anchor = 0, .depth = depth }}); + _ = try block.hold(d, &.{.{ .lane = l, .pending = 1, .anchor = 0, .depth = depth }}, null); const ms = @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(d.io)).toNanoseconds())) / 1e6; if (i > 0) b.* = @min(b.*, ms); }; diff --git a/zig/src/families/lfm2_moe/dspark_block.zig b/zig/src/families/lfm2_moe/dspark_block.zig index 22676dca6e..d3f6cbabb8 100644 --- a/zig/src/families/lfm2_moe/dspark_block.zig +++ b/zig/src/families/lfm2_moe/dspark_block.zig @@ -152,8 +152,11 @@ pub fn end(cb: mtl.CommandBuffer, e: mtl.ComputeEncoder) !void { return finish(cb, e); } -/// The blocks' drafts into their lanes' held chains: the backbone over every block's rows, then the Markov chain. -pub fn hold(d: *Drafter, blocks: []const Block) !void { +/// A GPU event a hold signals instead of being waited on, and the value it sets. +pub const Signal = struct { event: mtl.Event, value: u64 }; + +/// The blocks' drafts (backbone, then Markov chain) into the lanes' held chains; with `signal`, left for readPicks. +pub fn hold(d: *Drafter, blocks: []const Block, signal: ?Signal) !?mtl.CommandBuffer { const pool = mtl.objc.Pool.push(); defer pool.pop(); const s = &d.s; @@ -230,6 +233,13 @@ pub fn hold(d: *Drafter, blocks: []const Block) !void { x.norm(s.h, s.r, d.w.norm, true, s.x, tail); try x.project(fwd.matrix(emb), .{ .buf = s.x }, .{ .buf = s.logits }, tail); try markov(x, blocks.len, n, steps); + if (signal) |sg| { // committed unwaited, the command buffer returned retained + e.end(); + cb.signalEvent(sg.event, sg.value); + const kept = cb.retain(); + cb.commit(); + return kept; + } try finish(cb, e); const picks = s.picks.slice(Pick, d.config.block * ds.pass_rows); for (blocks, 0..) |b, bi| { @@ -240,6 +250,19 @@ pub fn hold(d: *Drafter, blocks: []const Block) !void { } b.lane.held_len = b.depth; } + return null; +} + +/// A finished signalled hold's picks into the lanes' held chains (a pick past the vocabulary reads as token 0). +pub fn readPicks(d: *Drafter, lanes: []const *Lane, depths: []const u32) void { + const picks = d.s.picks.slice(Pick, d.config.block * ds.pass_rows); + for (lanes, depths, 0..) |l, depth, bi| { + for (l.held[0..depth], 0..) |*h, i| { + const p = picks[i * ds.pass_rows + bi]; + h.* = if (p.token < c.vocab) p.token else 0; + } + l.held_len = depth; + } } /// Step i of each stream: its row i's logits plus W2 W1[previous draft] (the pending token first), then the argmax. diff --git a/zig/src/families/lfm2_moe/forward.zig b/zig/src/families/lfm2_moe/forward.zig index 7b63833192..10040c0d77 100644 --- a/zig/src/families/lfm2_moe/forward.zig +++ b/zig/src/families/lfm2_moe/forward.zig @@ -118,7 +118,17 @@ pub fn matrix(w: row.Weights) tile.Matrix { return .{ .w = w.w, .w_off = w.w_off, .scales = w.scales, .s_off = w.s_off, .biases = w.biases, .b_off = w.b_off, .n = w.n, .k = w.k, .bits = w.bits }; } +/// A window's held drafts on the device: `count` tokens at `src` + `off`, `stride` words apart, into ids from `row`. +pub const Held = struct { row: usize, src: Buffer, off: usize, stride: u32, count: u32 }; +/// A GPU event the forward waits on before it reads held drafts, and the value it waits for. +pub const Wait = struct { event: mtl.Event, value: u64 }; + pub fn run(m: *Model, s: *st.Scratch, segments: []const Segment, tokens: []const u32, record: bool, head: Head) !void { + return runHeld(m, s, segments, tokens, record, head, &.{}, null); +} + +/// `run` with held drafts copied into the rows' ids on the GPU first (after `wait`): their host ids are ignored. +pub fn runHeld(m: *Model, s: *st.Scratch, segments: []const Segment, tokens: []const u32, record: bool, head: Head, held: []const Held, wait: ?Wait) !void { if (tokens.len == 0 or tokens.len > s.rows or (record and tokens.len > st.batch_rows)) return error.LfmRowsExceeded; var total: usize = 0; for (segments) |seg| { @@ -131,10 +141,18 @@ pub fn run(m: *Model, s: *st.Scratch, segments: []const Segment, tokens: []const const pool = mtl.objc.Pool.push(); defer pool.pop(); const cb = m.queue.commandBuffer(); + if (wait) |w| cb.waitForEvent(w.event, w.value); const prompt = (!record and head == .last) or total >= tile_rows; 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 }; + for (held) |h| { + enc.pipe("lfm2_draft_ids"); + e.setBuffer(h.src, h.off, 0); + e.setValue([4]u32{ h.stride, h.count, c.vocab, 0 }, 1); + e.setBuffer(s.ids, h.row * 4, 2); + enc.threads(.{ h.count, 1, 1 }, .{ 32, 1, 1 }); + } const embedding = m.weights.embedding; e.setPipeline(m.kernels.embed(embedding.bits)); e.setBuffer(s.ids, 0, 0); diff --git a/zig/src/families/lfm2_moe/kernels.zig b/zig/src/families/lfm2_moe/kernels.zig index 357500c04d..936ba33197 100644 --- a/zig/src/families/lfm2_moe/kernels.zig +++ b/zig/src/families/lfm2_moe/kernels.zig @@ -32,6 +32,7 @@ const names = [_][:0]const u8{ "lfm2_tap", "lfm2_markov_scan", "lfm2_markov_pick", + "lfm2_draft_ids", }; pub const Kernels = struct { diff --git a/zig/src/native/lfm2_moe.zig b/zig/src/native/lfm2_moe.zig index 05f0260dcc..a78d67d97b 100644 --- a/zig/src/native/lfm2_moe.zig +++ b/zig/src/native/lfm2_moe.zig @@ -102,7 +102,12 @@ pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Ope return null; }; h.metal.options.streams = @max(served, 2); - if (h.draft) |d| h.wrap = try lanes.drafted.Drafted.init(gpa, h.metal.backend(), d.drafter()); + if (h.draft) |d| { + h.wrap = try lanes.drafted.Drafted.init(gpa, h.metal.backend(), d.drafter()); + // the verify copies the drafts on the GPU after the unwaited hold + h.metal.held_source = d.heldSource(); + h.wrap.?.held_on_target = true; + } errdefer if (h.wrap) |*w| w.deinit(); const facts = if (h.wrap) |*w| w.facts(h.metal.facts()) else h.metal.facts(); h.config = try lanes.Config.init(gpa, facts, q.state.window_rows, q.state.window_rows - 1); diff --git a/zig/tests/contract/metal_lfm2_dspark.zig b/zig/tests/contract/metal_lfm2_dspark.zig index 3f2ae29907..63d7bcf590 100644 --- a/zig/tests/contract/metal_lfm2_dspark.zig +++ b/zig/tests/contract/metal_lfm2_dspark.zig @@ -1,4 +1,4 @@ -//! LFM2's DSpark Markov step (scan then pick): each block row's draft from its base logits and the previous token. +//! LFM2's DSpark kernels: the Markov step (scan then pick), and held drafts into a verify's ids. const std = @import("std"); const mtl = @import("metal"); const contract = @import("contract"); @@ -7,14 +7,16 @@ const c = contract.spec; const g = contract.gates; const rig_ = @import("rig.zig"); -pub const Which = enum { markov_step }; +pub const Which = enum { markov_step, draft_ids }; const rank = 256; const vocab = 128000; const span = 512; // tokens a scan threadgroup takes (the kernel's MARKOV_SPAN) const groups = (vocab + span - 1) / span; const table = 1000; // W1's rows in the check: ids are inside it or past the vocabulary (a zero row) -const names = [_][]const u8{"metal lfm2 markov step"}; +const names = [_][]const u8{ "metal lfm2 markov step", "metal lfm2 draft ids" }; +const stride = 64; // words between a stream's picks: a hold pass's 32 blocks, a token and a flag each +const ids_shapes = [_]c.Shape{ .{ .rows = 1, .n = 1, .k = 1 }, .{ .rows = 4, .n = 1, .k = 1 }, .{ .rows = 15, .n = 1, .k = 1 } }; const step_shapes = [_]c.Shape{ .{ .rows = 1, .n = vocab, .k = rank }, .{ .rows = 4, .n = vocab, .k = rank }, .{ .rows = 9, .n = vocab, .k = rank } }; fn none(_: c.Shape) usize { @@ -27,6 +29,7 @@ fn depth(_: c.Shape) usize { pub fn spec(w: Which) c.Spec { return switch (w) { + .draft_ids => .{ .name = names[1], .accum = .none, .depth = none, .out = .raw, .served = &ids_shapes, .sums = .{ .f16 = 0, .bf16 = 0, .why = "a copy moves tokens: no sums" } }, .markov_step => .{ .name = names[0], .accum = .none, .depth = depth, .out = .raw, .served = &step_shapes, .nan = .{ .stated = "a row with a NaN score is flagged nonfinite" }, .sums = .{ .f16 = 0, .bf16 = 0, .why = "picks are tokens: the check bounds each pick's score by gamma(256), not sum bits" } }, }; } @@ -64,6 +67,7 @@ pub const Probe = struct { pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, w: Which) !Probe { const kernels: [2][:0]const u8 = switch (w) { .markov_step => .{ "lfm2_markov_scan", "lfm2_markov_pick" }, + .draft_ids => .{ "lfm2_draft_ids", "lfm2_draft_ids" }, }; const first = try rig_.pipeline(device, sources.lfm2, kernels[0]); errdefer first.deinit(); @@ -97,6 +101,20 @@ pub const Probe = struct { pub fn prepare(p: *Probe, s: c.Shape, f: g.Fill, seed: u64) !bool { p.drop(); p.rows = s.rows; + if (p.which == .draft_ids) { // picks a stride apart: tokens inside the vocabulary, its last, and past it + if (s.n != 1 or s.rows == 0 or s.rows > 16) return false; + p.ids = try p.a.alloc(u32, s.rows * stride); + for (p.ids, 0..) |*id, i| id.* = switch ((i / stride) % 3) { + 0 => @intCast((@as(u64, i) * 7919 + seed) % vocab), + 1 => vocab - 1, + else => vocab + @as(u32, @intCast(i % 5)), + }; + p.args[0] = try p.rig.input(std.mem.sliceAsBytes(p.ids)); + const at = [4]u32{ stride, @intCast(s.rows), vocab, 0 }; + p.args[1] = try p.rig.input(std.mem.asBytes(&at)); + p.y = try p.rig.output(p.count()); + return true; + } if (s.n != vocab or s.k != rank or s.rows == 0) return false; // each stream's step row sits a block (two rows here) apart, as a hold's rows of one step do p.src = try p.a.alloc(u16, 2 * s.rows * vocab); @@ -135,12 +153,21 @@ pub const Probe = struct { pub fn count(p: *Probe) usize { return switch (p.which) { .markov_step => p.rows * 8, + .draft_ids => p.rows * 4, }; } pub fn run(p: *Probe) !void { const pass = p.rig.begin(); const e = pass.e; + if (p.which == .draft_ids) { + e.setPipeline(p.pipes[0]); + 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.y.buf, p.y.off, 2); + e.dispatchThreads(mtl.Size.of(p.rows, 1, 1), mtl.Size.of(32, 1, 1)); + return p.rig.finish(pass); + } const best = try p.rig.scratch(p.rows * groups * 4); const at = try p.rig.scratch(p.rows * groups * 4); e.setPipeline(p.pipes[0]); @@ -167,6 +194,13 @@ pub const Probe = struct { /// Each stream's pick against float64 scores with the kernel's bf16 roundings: the best, but for near ties. pub fn check(p: *Probe, _: g.Gate) !?[]const u8 { const out = p.rig.bytes(p.y, p.count()); + if (p.which == .draft_ids) { // exact: each stream's first pick, or 0 for a token past the vocabulary + for (0..p.rows) |r| { + const token = p.ids[r * stride]; + if (std.mem.readInt(u32, out[4 * r ..][0..4], .little) != (if (token < vocab) token else 0)) return "an id differs from its pick"; + } + return null; + } for (0..p.rows) |s| { const token = std.mem.readInt(u32, out[8 * s ..][0..4], .little); const flag = std.mem.readInt(u32, out[8 * s + 4 ..][0..4], .little); @@ -204,10 +238,10 @@ pub const Probe = struct { } }; -/// Gate 5 for the scan and pick on every Apple chip and this GPU: 256 threads each, 1,088 and 3,072 threadgroup bytes. +/// Gate 5 for the three launches on every Apple chip and this GPU: ids 32 threads, others 256, the scan 1,088 bytes. 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 }{ .{ "metal lfm2 markov scan", "lfm2_markov_scan", 1088, 256 }, .{ "metal lfm2 markov pick", "lfm2_markov_pick", 3072, 256 } }; + const launches = [_]struct { []const u8, []const u8, u32, u32 }{ .{ "metal lfm2 markov scan", "lfm2_markov_scan", 1088, 256 }, .{ "metal lfm2 markov pick", "lfm2_markov_pick", 3072, 256 }, .{ "metal lfm2 draft ids", "lfm2_draft_ids", 0, 32 } }; for (launches) |l| { if (only) |s| if (std.mem.indexOf(u8, l[0], s) == null) continue; const d = contract.limits.MetalLaunch{ .name = l[0], .threads = l[3], .tg_memory = l[2], .offline = false }; diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 6f2e5d76b6..3d42ff32b2 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -177,7 +177,7 @@ pub fn main(init: std.process.Init) !u8 { tap = try lfm2_tap.Probe.init(a, device, queue); try kernels.append(a, tap.?.kernel()); } - var lfm2_list: [1]lfm2_dspark.Probe = undefined; + var lfm2_list: [2]lfm2_dspark.Probe = undefined; var lfm2_made: usize = 0; defer for (lfm2_list[0..lfm2_made]) |*p| p.deinit(); for (std.enums.values(lfm2_dspark.Which)) |w| { diff --git a/zig/tests/lfm2_dspark.zig b/zig/tests/lfm2_dspark.zig index 09c7f68a52..f95ce615e5 100644 --- a/zig/tests/lfm2_dspark.zig +++ b/zig/tests/lfm2_dspark.zig @@ -153,7 +153,7 @@ pub fn main(init: std.process.Init) !void { try absorb(d, .{ .lane = lane, .src = .{ .buf = taps.buf, .off = at * width * 2 }, .rows = rows, .p0 = at }); } lane.end = ctx; - try q.dspark_block.hold(d, &.{.{ .lane = lane, .pending = ids[ctx], .anchor = ctx, .depth = block }}); + _ = try q.dspark_block.hold(d, &.{.{ .lane = lane, .pending = ids[ctx], .anchor = ctx, .depth = block }}, null); var held: [16]u32 = undefined; @memcpy(held[0..block], lane.held[0..block]); const gpu = held[0..block]; @@ -171,12 +171,12 @@ pub fn main(init: std.process.Init) !void { var best: f64 = std.math.floatMax(f64); for (0..5) |_| { const t0 = std.Io.Clock.awake.now(init.io); - try q.dspark_block.hold(d, &.{.{ .lane = lane, .pending = ids[ctx], .anchor = ctx, .depth = depth }}); + _ = try q.dspark_block.hold(d, &.{.{ .lane = lane, .pending = ids[ctx], .anchor = ctx, .depth = depth }}, null); best = @min(best, @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(init.io)).toNanoseconds())) / 1e6); } std.debug.print("hold at depth {d}: {d:.2} ms (best of 5)\n", .{ depth, best }); } - try q.dspark_block.hold(d, &.{.{ .lane = lane, .pending = ids[ctx], .anchor = ctx, .depth = block }}); + _ = try q.dspark_block.hold(d, &.{.{ .lane = lane, .pending = ids[ctx], .anchor = ctx, .depth = block }}, null); // the host reference, fp32 throughout var r: Ref = .{ .gpa = gpa, .check = ck.Checkpoint.init(gpa), .cfg = cfg }; defer r.check.deinit(); 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 18/46] 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 57608634baf00179b90d24b31227d2da0233ad7e Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 10:06:08 +0100 Subject: [PATCH 19/46] docs: the LFM2 recipe's DSpark drafts with LiquidAI's drafter The recipe says where the drafter comes from (LiquidAI/LFM2.5-24B-A2B-DSpark at its revision), that its weights are under the LFM Open License v1.0 and not shipped with TensorFold, how to serve with --drafter, what the drafter runs and which kernels and checks cover it. The README row and the recipe index name the drafts. --- README.md | 2 +- docs/recipes/README.md | 2 +- docs/recipes/lfm2.5-24b-a2b.md | 41 ++++++++++++++++++++++++++++++++++ 3 files changed, 43 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 1c121be219..f013d4160e 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 recipe](docs/recipes/lfm2.5-24b-a2b.md), [release history](CHANGELOG.md). | +| 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). LiquidAI's DSpark drafter with `--drafter`. | | 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 fa160dc631..f0104d747d 100644 --- a/docs/recipes/README.md +++ b/docs/recipes/README.md @@ -21,7 +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) | +| LFM2.5-24B-A2B | [Native Metal, MLX 4-bit, shared rounds, DSpark drafts](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 index 7b8b2cdf4b..f4272b59a5 100644 --- a/docs/recipes/lfm2.5-24b-a2b.md +++ b/docs/recipes/lfm2.5-24b-a2b.md @@ -41,3 +41,44 @@ the shared-round kernels on it. The M4 Pro is not yet measured on a quiet machin 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. + +## DSpark drafts + +LiquidAI's DSpark drafter for this model is +[`LiquidAI/LFM2.5-24B-A2B-DSpark`](https://huggingface.co/LiquidAI/LFM2.5-24B-A2B-DSpark/tree/e052e5b1dd103c38cdd4f8ddc068d7dea4586f55) +at revision `e052e5b1dd103c38cdd4f8ddc068d7dea4586f55` (327.7M bf16 parameters, 655 MB). Its weights are under the +[LFM Open License v1.0](https://huggingface.co/LiquidAI/LFM2.5-24B-A2B-DSpark/blob/e052e5b1dd103c38cdd4f8ddc068d7dea4586f55/LICENSE); +TensorFold does not ship them, so read the license before use. + +```bash +hf download LiquidAI/LFM2.5-24B-A2B-DSpark --revision e052e5b1dd103c38cdd4f8ddc068d7dea4586f55 \ + config.json model.safetensors LICENSE README.md --local-dir "$HOME/models/LFM2.5-24B-A2B-DSpark" +zig-out/native/bin/tensorfold-native serve "$MODEL_DIR" --port 8080 --parallel 16 \ + --drafter "$HOME/models/LFM2.5-24B-A2B-DSpark" +``` + +Without `--drafter` the engine serves as before. `"draft": false` on a request, or `--no-drafts`, turns the drafter +off; the lane core verifies every draft, so a drafted reply equals its `"draft": false` reply token for token. + +The drafter has no embedding and no head of its own: + +- Context: the target's rows after layers 2, 10, 18, 26 and 34 (the backend's taps, `lfm2_tap`), through `fc` and + `hidden_norm`, give each of its five layers' context keys and values (k_norm, then RoPE at the row's position). +- Block: the pending token and mask tokens run five Qwen3 layers with non-causal attention over the context and the + block, then the target's head. A block is as long as the deepest chain its pass holds, at most the nine rows the + drafter was trained on. +- Markov head: row i's logits plus `W2 W1[previous token]` (rank 256); draft i is their argmax. + +Its bf16 linears are quantized at load to 8-bit MLX affine codes in groups of 64 and run on the family's kernels +and `core/tile_qmm`; the q and k rows are reordered so that the target's half-split RoPE kernel gives the +checkpoint's interleaved RoPE. Its own kernels are `lfm2_tap`, `lfm2_attn_block` and `lfm2_attn_block_merge` +(non-causal block attention, the target's attention kernels with a template flag), `lfm2_markov_scan` and +`lfm2_markov_pick` (a Markov step) and `lfm2_draft_ids` (a hold's drafts into the next verify's rows on the GPU, so +the host does not wait for a hold that fits one pass). Each kept prompt state keeps the drafter's context too, so a +drafting stream resumes prompts. Shared rounds take the lane core's shared-round guard: they draft only while +drafted rounds measure faster than plain ones at that stream count. + +`tf-kernel-contract --only lfm2` checks these kernels. `tf-lfm2-dspark MODEL DRAFTER TOKENS [ROWS]` runs the drafter +against an fp32 host reference of the released checkpoint (the same taps and the target's embedding and head), +prints each draft against the reference's pick with the margin between the reference's two best tokens, and times +the drafter's passes; a near tie can pick the other token. From c5d6c2d131751dec9e13485a53b97f0ae4583548 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:37:14 +0100 Subject: [PATCH 20/46] lfm2: a DSpark confidence logit kernel, registered with the kernel contract LiquidAI's DSpark drafter carries a confidence head: one linear over a draft's final block row (the row the target's head reads) and the W1 row of the token before it, 2,048 + 256 inputs, to one logit. lfm2_dspark_confidence (lfm2_dspark.metal) computes that logit for every row of a hold pass: 256 threads a row, fp32 products and sums, the eight simdgroups' partials summed in threadgroup memory. Nothing calls it yet. It registers with the kernel contract (metal_lfm2_confidence.zig): the served pass shapes against a float64 reference, the accumulation gate against reduced-precision sums, and gate 5 against every Apple chip we serve and this GPU. --- zig/kernels/metal/lfm2/lfm2_dspark.metal | 27 +++++ zig/src/families/lfm2_moe/kernels.zig | 1 + zig/tests/contract/metal_lfm2_confidence.zig | 116 +++++++++++++++++++ zig/tests/contract/metal_main.zig | 8 ++ 4 files changed, 152 insertions(+) create mode 100644 zig/tests/contract/metal_lfm2_confidence.zig diff --git a/zig/kernels/metal/lfm2/lfm2_dspark.metal b/zig/kernels/metal/lfm2/lfm2_dspark.metal index 050fd5e8cc..7f613def9b 100644 --- a/zig/kernels/metal/lfm2/lfm2_dspark.metal +++ b/zig/kernels/metal/lfm2/lfm2_dspark.metal @@ -105,3 +105,30 @@ constant constexpr uint MARKOV_SPAN = 512; const uint token = picks[size_t(i) * p.x]; ids[i] = token < p.z ? token : 0u; } + +// DSpark's confidence logit of draft i of stream s: proj . [x row s * n + i, W1[previous token]] in fp32, no bias. +struct lfm2_confidence_args { uint n, pass, vocab, rank, hidden; }; +[[kernel]] void lfm2_dspark_confidence(const device bfloat* x [[buffer(0)]], const device bfloat* w1 [[buffer(1)]], + const device uint* ids [[buffer(2)]], const device uint* picks [[buffer(3)]], + const device bfloat* proj [[buffer(4)]], + constant lfm2_confidence_args& p [[buffer(5)]], device float* out [[buffer(6)]], + uint2 tg [[threadgroup_position_in_grid]], uint t [[thread_index_in_threadgroup]], + uint sg [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + const uint i = tg.x, s = tg.y; + // the previous token: the block's pending id for the first draft, else the Markov step's pick before it + const uint id = i == 0 ? ids[size_t(s) * p.n] : picks[2 * (size_t(i - 1) * p.pass + s)]; + const device bfloat* row = x + (size_t(s) * p.n + i) * p.hidden; + float d = 0.0f; + for (uint k = t; k < p.hidden; k += 256) d = fma(float(row[k]), float(proj[k]), d); + if (id < p.vocab) + for (uint k = t; k < p.rank; k += 256) d = fma(float(w1[size_t(id) * p.rank + k]), float(proj[p.hidden + k]), d); + d = simd_sum(d); + if (lane == 0) part[sg] = d; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (t != 0) return; + float sum = 0.0f; + for (uint g = 0; g < 8; ++g) sum += part[g]; + out[size_t(i) * p.pass + s] = sum; +} diff --git a/zig/src/families/lfm2_moe/kernels.zig b/zig/src/families/lfm2_moe/kernels.zig index 936ba33197..96997a1109 100644 --- a/zig/src/families/lfm2_moe/kernels.zig +++ b/zig/src/families/lfm2_moe/kernels.zig @@ -33,6 +33,7 @@ const names = [_][:0]const u8{ "lfm2_markov_scan", "lfm2_markov_pick", "lfm2_draft_ids", + "lfm2_dspark_confidence", }; pub const Kernels = struct { diff --git a/zig/tests/contract/metal_lfm2_confidence.zig b/zig/tests/contract/metal_lfm2_confidence.zig new file mode 100644 index 0000000000..f1543a183d --- /dev/null +++ b/zig/tests/contract/metal_lfm2_confidence.zig @@ -0,0 +1,116 @@ +//! LFM2's DSpark confidence logit (lfm2_dspark_confidence): a dot of a row and a W1 row with the head's 2,304 weights. +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 rows = contract.rows; +const rig_ = @import("rig.zig"); + +const hidden = 2048; +const rank = 256; +pub const name = "metal lfm2 dspark confidence"; + +/// rows: one stream's drafts (row i reads the block's row i and the W1 row of the token before draft i). +const shapes = [_]c.Shape{ .{ .rows = 1, .n = 1, .k = hidden + rank }, .{ .rows = 4, .n = 1, .k = hidden + rank }, .{ .rows = 9, .n = 1, .k = hidden + rank } }; + +/// A thread's chain over every 256th input, a 32-way simd_sum tree, then the eight simdgroups' partials in order. +fn depth(s: c.Shape) usize { + return (s.k + 255) / 256 + 5 + 8; +} + +pub fn spec() c.Spec { + return .{ .name = name, .accum = .f32, .depth = depth, .out = .f32, .served = &shapes }; +} + +pub const Probe = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + r: ?rows.Rows = null, + args: [6]rig_.Ref = undefined, // x, W1 rows, ids, picks, the head's weights, the shape + 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_dspark_confidence") }; + } + + pub fn deinit(p: *Probe) void { + if (p.r) |r| r.deinit(); + 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() }; + } + + /// Row i's first 2,048 inputs as the block's row i, its last 256 as W1 row i: the token before draft i is i. + pub fn prepare(p: *Probe, s: c.Shape, f: g.Fill, seed: u64) !bool { + if (s.k != hidden + rank or s.rows == 0 or s.rows > 16) return false; + if (p.r) |r| r.deinit(); + p.r = null; + p.rig.reset(); + const r = try rows.Rows.init(p.a, .dot, s); + p.r = r; + r.fill(f, seed); + const x = try p.a.alloc(u16, s.rows * hidden); + defer p.a.free(x); + const w1 = try p.a.alloc(u16, s.rows * rank); + defer p.a.free(w1); + for (0..s.rows) |i| { + @memcpy(x[i * hidden ..][0..hidden], r.x[i * s.k ..][0..hidden]); + @memcpy(w1[i * rank ..][0..rank], r.x[i * s.k + hidden ..][0..rank]); + } + // the block's pending id is token 0; the pick before draft i (i > 0) is token i, a flag word after each + var picks: [32]u32 = @splat(0xdead); + for (1..s.rows) |i| picks[2 * (i - 1)] = @intCast(i); + const ids = [1]u32{0}; + const args = [5]u32{ @intCast(s.rows), 1, @intCast(s.rows), rank, hidden }; + p.args[0] = try p.rig.input(std.mem.sliceAsBytes(x)); + p.args[1] = try p.rig.input(std.mem.sliceAsBytes(w1)); + p.args[2] = try p.rig.input(std.mem.asBytes(&ids)); + p.args[3] = try p.rig.input(std.mem.asBytes(&picks)); + p.args[4] = try p.rig.input(std.mem.sliceAsBytes(r.w)); + p.args[5] = try p.rig.input(std.mem.asBytes(&args)); + p.out = try p.rig.output(s.rows * 4); + return true; + } + + pub fn run(p: *Probe) !void { + const n = p.r.?.shape.rows; + const pass = p.rig.begin(); + const e = pass.e; + e.setPipeline(p.pipe); + for (p.args, 0..) |ref, i| e.setBuffer(ref.buf, ref.off, i); + e.setBuffer(p.out.buf, p.out.off, 6); + e.dispatchGroups(mtl.Size.of(n, 1, 1), mtl.Size.of(256, 1, 1)); + try p.rig.finish(pass); + } + + pub fn count(p: *Probe) usize { + return p.r.?.outputs(); + } + + pub fn reference(p: *Probe, want: []f64, mag: []f64) !void { + return p.r.?.reference(want, mag); + } + + pub fn simulate(p: *Probe, acc: c.Accum, order: g.Order, out: []f64) !void { + return p.r.?.simulate(acc, order, out); + } +}; + +/// Gate 5 for the launch on every Apple chip we serve and this GPU: 256 threads, the eight simdgroups' partials. +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 = 32, .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_dspark_confidence", .here = sources.lfm2, .pre_m5 = "" }, false); +} diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 3d42ff32b2..687f5c83f4 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -28,6 +28,7 @@ const lfm2_streams = @import("metal_lfm2_streams.zig"); const lfm2_tap = @import("metal_lfm2_tap.zig"); const lfm2_dspark = @import("metal_lfm2_dspark.zig"); const lfm2_attn = @import("metal_lfm2_attn.zig"); +const lfm2_conf = @import("metal_lfm2_confidence.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]"; @@ -192,6 +193,12 @@ pub fn main(init: std.process.Init) !u8 { block_attn = try lfm2_attn.Probe.init(a, device, queue); try kernels.append(a, block_attn.?.kernel()); } + var confidence: ?lfm2_conf.Probe = null; + defer if (confidence) |*p| p.deinit(); + if (wanted(o, lfm2_conf.name)) { + confidence = try lfm2_conf.Probe.init(a, device, queue); + try kernels.append(a, confidence.?.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(); @@ -336,6 +343,7 @@ pub fn main(init: std.process.Init) !u8 { lim += try lfm2_tap.limits(a, device, o.only); lim += try lfm2_dspark.limits(a, device, o.only); lim += try lfm2_attn.limits(a, device, o.only); + lim += try lfm2_conf.limits(a, device, o.only); } const t2 = std.Io.Clock.awake.now(init.io); const ms = struct { From 6e1340999e9e584b5ce77ba2411ef1b81f576d9c Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:37:43 +0100 Subject: [PATCH 21/46] lfm2: the DSpark confidence head schedules the verify (--draft-verify, --verify-costs) With the scheduled verify a hold also runs lfm2_dspark_confidence, and each lane keeps sigmoid(logit + bias) for its drafts. chances() gives their running product, the chance that a draft and every draft before it land, as Facts.chances: the lane core's window allocator then picks each stream's verify rows by those chances and the verify costs, and the next hold drafts as deep as the stream's grant. No second scheduler. While a hold still runs, its rows are priced by the lane's last settled confidences (draft j by the last hold's draft j, the deepest one repeated past it), so a verify does not wait on the GPU for them. A lane with none settled yet waits. Two serve flags on macOS, listed by capabilities --json and refused unless an lfm2_moe checkpoint is served with --drafter: - --draft-verify depth|fixed|scheduled: the core's depth rule, the whole block every round (the published fixed window), or the confidence schedule. scheduled is the default. - --verify-costs FILE: a table of verify forward ms by streams and rows a stream, written by tf-lfm2-verify-costs on the serving Mac, in place of the costs timed at load. Shared rounds keep the shared-round guard in every mode but fixed. Tests: test-lfm2-dspark reads a cost table, refuses malformed ones and checks the default mode; the server's parser tests take both flags and refuse an unknown mode. tf-lfm2-dspark checks the head against its fp32 reference. --- zig/build/metal_targets.zig | 1 + zig/src/core/engine_api.zig | 6 ++ zig/src/families/lfm2_moe/dspark.zig | 59 ++++++++++++++++- zig/src/families/lfm2_moe/dspark_block.zig | 28 ++++++++ zig/src/families/lfm2_moe/dspark_test.zig | 5 ++ zig/src/families/lfm2_moe/dspark_weights.zig | 12 ++++ zig/src/families/lfm2_moe/lfm2_moe.zig | 1 + zig/src/families/lfm2_moe/verify_costs.zig | 44 +++++++++++++ zig/src/native/lfm2_moe.zig | 30 ++++++++- zig/src/native/metal.zig | 1 + zig/src/server/cli.zig | 8 ++- zig/src/server/cli_test.zig | 7 ++ zig/src/server/main.zig | 2 +- zig/src/tensorfold.zig | 1 + zig/tests/lfm2_dspark.zig | 15 ++++- zig/tests/lfm2_verify_costs.zig | 67 ++++++++++++++++++++ 16 files changed, 278 insertions(+), 9 deletions(-) create mode 100644 zig/src/families/lfm2_moe/verify_costs.zig create mode 100644 zig/tests/lfm2_verify_costs.zig diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index 727daf7e2e..9896a764f2 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -293,6 +293,7 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s .{ .name = "tf-lfm2-bench", .path = "zig/tests/lfm2_bench.zig", .about = "LFM2-MoE prompt-chunk and one-row decode timing" }, .{ .name = "tf-lfm2-exact", .path = "zig/tests/lfm2_exact.zig", .about = "LFM2-MoE prompt chunks, windows and committed-state exactness" }, .{ .name = "tf-lfm2-dspark", .path = "zig/tests/lfm2_dspark.zig", .about = "LFM2's DSpark drafter on the GPU against an fp32 host reference of the released drafter" }, + .{ .name = "tf-lfm2-verify-costs", .path = "zig/tests/lfm2_verify_costs.zig", .about = "LFM2 verify forward ms by stream count and rows a stream, written as a cost table" }, .{ .name = "tf-lfm2-fp64", .path = "zig/tests/lfm2_fp64.zig", .about = "LFM2 projection classes and attention on the tensor units, the tiles and the twin, each against float64" }, .{ .name = "tensorfold", .path = "zig/src/main.zig", .about = "The native engine's command line" }, .{ .name = "tf-nemotron-fixtures", .path = "zig/tests/nemotron_fixtures.zig", .about = "Nemotron kernels against the Python engine's captured ops" }, diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index a50d0de404..8427cc9e0e 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -259,6 +259,9 @@ pub const Family = struct { model_type: []const u8, formats: []const []const u8, /// The math an NVFP4 checkpoint runs: its own (FP4 rows into NVFP4 layers, FP8 into FP8 ones), or bf16 rows throughout. pub const Precision = enum { checkpoint, full }; +/// An external drafter's verify: the depth rule's rows, its whole block, or its confidence head's survivals. +pub const DraftVerify = enum { depth, fixed, scheduled }; + /// What a server asks of the engine it opens: the checkpoint, and the serve flags an engine reads. pub const Open = struct { dir: []const u8, @@ -277,6 +280,9 @@ pub const Open = struct { /// --drafter: a draft model's directory (Qwen3.8-27B's DFlash2; null: none), and --drafter-bits: 4, or 0 for bf16. drafter: ?[]const u8 = null, drafter_bits: u8 = 4, + /// --draft-verify and --verify-costs (a tf-lfm2-verify-costs table): the LFM2 DSpark drafter's verify rows. + draft_verify: ?DraftVerify = null, + verify_costs: ?[]const u8 = null, /// --record: each stream's committed rows and tapped features written under this folder, for drafter training. record: ?[]const u8 = null, speed_up: ?[]const u8 = null, diff --git a/zig/src/families/lfm2_moe/dspark.zig b/zig/src/families/lfm2_moe/dspark.zig index 850c1ce34c..0263d03f64 100644 --- a/zig/src/families/lfm2_moe/dspark.zig +++ b/zig/src/families/lfm2_moe/dspark.zig @@ -22,8 +22,13 @@ pub const absorb_rows = 2048; pub const stage_rows = 128; const kw = c.kv_heads * c.head_dim; +/// Verify rows by the core's depth rule, the whole block (fixed), or the confidence head's chances (scheduled). +pub const Verify = enum { depth, fixed, scheduled }; +/// The verify when none is asked: the confidence-scheduled one. +pub const default_verify: Verify = .scheduled; + /// `bits`: the width the bf16 linears are quantized to (8 keeps the fp32 reference's drafts; 4 does not). -pub const Options = struct { capacity: usize, bits: u8 = 8 }; +pub const Options = struct { capacity: usize, bits: u8 = 8, verify: Verify = default_verify }; /// One stream's context keys and values a layer ([kv head][cap][head dim], as the target's attention reads them). pub const Lane = struct { @@ -35,6 +40,8 @@ pub const Lane = struct { staged: u32 = 0, held: [16]u32 = @splat(0), held_len: u32 = 0, + conf: [16]f32 = @splat(0), // each held draft's chance given the drafts before it landed (the confidence head) + conf_len: u32 = 0, // the drafts `conf` covers: the last settled hold's depth slot: ?u32 = null, // its block's place in the last signalled hold, whose picks still hold its drafts }; @@ -62,6 +69,7 @@ pub const Scratch = struct { best: mtl.Buffer, // the Markov scan's best score and token of each threadgroup, by stream at: mtl.Buffer, picks: mtl.Buffer, + conf: mtl.Buffer, // the confidence logits, step by step, a pass's blocks each }; pub const Drafter = struct { @@ -147,6 +155,7 @@ pub const Drafter = struct { .best = try d.buffer(r * block.scan_groups * 4), .at = try d.buffer(r * block.scan_groups * 4), .picks = try d.buffer(d.config.block * r * 8), + .conf = try d.buffer(d.config.block * r * 4), }; } @@ -156,7 +165,12 @@ pub const Drafter = struct { } pub fn drafter(d: *Drafter) dr.Drafter { - return .{ .ptr = d, .vtable = &.{ .taps = tapsFn, .facts = factsFn, .open = openFn, .absorb = absorbFn, .hold = holdFn, .held = heldFn, .release = releaseFn } }; + return .{ .ptr = d, .vtable = &.{ .taps = tapsFn, .facts = factsFn, .open = openFn, .absorb = absorbFn, .hold = holdFn, .held = heldFn, .release = releaseFn, .chances = chancesFn } }; + } + + /// Whether holds compute the confidence head (the scheduled verify, on a checkpoint that has the head). + pub fn confident(d: *const Drafter) bool { + return d.options.verify == .scheduled and d.w.confidence != null; } fn cast(ptr: *anyopaque) *Drafter { @@ -173,7 +187,46 @@ pub const Drafter = struct { fn factsFn(ptr: *anyopaque) dr.Facts { const d = cast(ptr); - return .{ .depth = d.config.block, .step_ms = d.step_ms, .prior = prior[0..d.config.block], .batched = true, .shared_guard = true, .stream_bytes = d.laneBytes(d.options.capacity) }; + var f: dr.Facts = .{ .depth = d.config.block, .step_ms = d.step_ms, .prior = prior[0..d.config.block], .batched = true, .shared_guard = true, .stream_bytes = d.laneBytes(d.options.capacity) }; + switch (d.options.verify) { + .depth => {}, + .fixed => { // every round the whole block, no plain stretches + f.chances = true; + f.confidence_floor = 0; + f.shared_guard = false; + }, + .scheduled => f.chances = d.confident(), + } + return f; + } + + /// A held chain's chances: the confidence head's running product, or for a fixed window, ones that fall by depth. + fn chancesFn(ptr: *anyopaque, s: *Stream, out: []f64) anyerror!bool { + const d = cast(ptr); + const l = d.lanes.get(s) orelse return false; + if (d.options.verify == .fixed) { // no wait: equal chances at a depth share a capped round's rows evenly + for (out, 1..) |*o, j| o.* = 1 - 1e-3 * @as(f64, @floatFromInt(j)); + return true; + } + if (!d.confident()) return false; + if (d.pending != null and l.slot != null and l.conf_len > 0) { + // the hold still runs: the lane's last settled confidences price it, so the host does not wait on the GPU + if (out.len > d.depths[l.slot.?]) return error.DraftHeldFewer; + var reach: f64 = 1; + for (out, 0..) |*o, j| { + reach *= l.conf[@min(j, l.conf_len - 1)]; + o.* = reach; + } + return true; + } + try d.settle(); + if (out.len > l.held_len) return error.DraftHeldFewer; + var reach: f64 = 1; + for (out, l.conf[0..out.len]) |*o, p| { + reach *= p; + o.* = reach; + } + return true; } pub fn laneBytes(d: *const Drafter, capacity: usize) usize { diff --git a/zig/src/families/lfm2_moe/dspark_block.zig b/zig/src/families/lfm2_moe/dspark_block.zig index d3f6cbabb8..0b712bf2dd 100644 --- a/zig/src/families/lfm2_moe/dspark_block.zig +++ b/zig/src/families/lfm2_moe/dspark_block.zig @@ -16,6 +16,7 @@ const QkArgs = extern struct { len: u32, cap: u32, eps: f32, log2_theta: f32 }; const AttnArgs = extern struct { len: u32, cap: u32, scale: f32, subs: u32 }; const Pick = round.contract.Pick; const MarkovArgs = extern struct { stride: u32, rank: u32, vocab: u32, prev_stride: u32 }; +const ConfidenceArgs = extern struct { n: u32, pass: u32, vocab: u32, rank: u32, hidden: u32 }; /// Tokens a Markov scan threadgroup takes (the kernel's MARKOV_SPAN), and threadgroups a stream's row takes. pub const scan_span = 512; pub const scan_groups = (c.vocab + scan_span - 1) / scan_span; @@ -233,6 +234,7 @@ pub fn hold(d: *Drafter, blocks: []const Block, signal: ?Signal) !?mtl.CommandBu x.norm(s.h, s.r, d.w.norm, true, s.x, tail); try x.project(fwd.matrix(emb), .{ .buf = s.x }, .{ .buf = s.logits }, tail); try markov(x, blocks.len, n, steps); + if (d.confident()) confidence(x, blocks.len, n, steps); if (signal) |sg| { // committed unwaited, the command buffer returned retained e.end(); cb.signalEvent(sg.event, sg.value); @@ -250,9 +252,18 @@ pub fn hold(d: *Drafter, blocks: []const Block, signal: ?Signal) !?mtl.CommandBu } b.lane.held_len = b.depth; } + for (blocks, 0..) |b, bi| readConfidence(d, b.lane, b.depth, bi); return null; } +/// Lane `l`'s drafts' chances from slot `bi` of the last hold's logits: sigmoid(logit + bias); none: left as they are. +fn readConfidence(d: *Drafter, l: *Lane, depth: u32, bi: usize) void { + if (!d.confident()) return; + const logits = d.s.conf.slice(f32, d.config.block * ds.pass_rows); + for (l.conf[0..depth], 0..) |*p, i| p.* = 1 / (1 + @exp(-(logits[i * ds.pass_rows + bi] + d.w.confidence_bias))); + l.conf_len = depth; +} + /// A finished signalled hold's picks into the lanes' held chains (a pick past the vocabulary reads as token 0). pub fn readPicks(d: *Drafter, lanes: []const *Lane, depths: []const u32) void { const picks = d.s.picks.slice(Pick, d.config.block * ds.pass_rows); @@ -262,9 +273,26 @@ pub fn readPicks(d: *Drafter, lanes: []const *Lane, depths: []const u32) void { h.* = if (p.token < c.vocab) p.token else 0; } l.held_len = depth; + readConfidence(d, l, depth, bi); } } +/// Each stream's draft i confidence logit (the final norm's row i, W1 of the token before draft i), for readConfidence. +fn confidence(x: Enc, streams: usize, n: u32, depth: u32) void { + const d = x.d; + const e = x.e; + x.pipe("lfm2_dspark_confidence"); + e.setBuffer(d.s.x, 0, 0); + e.setBuffer(d.w.w1, 0, 1); + e.setBuffer(d.s.ids, 0, 2); + e.setBuffer(d.s.picks, 0, 3); + e.setBuffer(d.w.confidence.?, 0, 4); + const args = ConfidenceArgs{ .n = n, .pass = ds.pass_rows, .vocab = c.vocab, .rank = d.config.rank, .hidden = c.hidden }; + e.setBytes(std.mem.asBytes(&args), 5); + e.setBuffer(d.s.conf, 0, 6); + e.dispatchGroups(mtl.Size.of(depth, streams, 1), mtl.Size.of(256, 1, 1)); +} + /// Step i of each stream: its row i's logits plus W2 W1[previous draft] (the pending token first), then the argmax. fn markov(x: Enc, streams: usize, n: u32, depth: u32) !void { const d = x.d; diff --git a/zig/src/families/lfm2_moe/dspark_test.zig b/zig/src/families/lfm2_moe/dspark_test.zig index 9ee31d868e..8fd1eb32e1 100644 --- a/zig/src/families/lfm2_moe/dspark_test.zig +++ b/zig/src/families/lfm2_moe/dspark_test.zig @@ -206,3 +206,8 @@ test "lfm2 dspark: lfm2_qk on half-split rows and gains gives interleaved RoPE's } } } + +test "lfm2 dspark: the drafter verifies by its confidence head unless asked otherwise" { + const o: @import("dspark.zig").Options = .{ .capacity = 64 }; + try std.testing.expectEqual(@import("dspark.zig").Verify.scheduled, o.verify); +} diff --git a/zig/src/families/lfm2_moe/dspark_weights.zig b/zig/src/families/lfm2_moe/dspark_weights.zig index 81dc727e56..bfd6d078e1 100644 --- a/zig/src/families/lfm2_moe/dspark_weights.zig +++ b/zig/src/families/lfm2_moe/dspark_weights.zig @@ -35,6 +35,9 @@ pub const Weights = struct { /// The Markov head: W1 [vocab, rank] bf16 as stored (rows gathered by token), W2 [vocab, rank] as 8-bit codes. w1: mtl.Buffer, w2: tile.Matrix, + /// The confidence head: its [hidden + rank] bf16 weights and fp32 bias (null: the checkpoint has none). + confidence: ?mtl.Buffer = null, + confidence_bias: f32 = 0, owned: std.ArrayList(mtl.Buffer) = .empty, pub fn deinit(w: *Weights, gpa: std.mem.Allocator) void { @@ -140,5 +143,14 @@ pub fn load(gpa: std.mem.Allocator, device: mtl.Device, io: std.Io, dir: []const w.w1 = try l.buffer(c.vocab * cfg.rank * 2); @memcpy(w.w1.slice(u16, c.vocab * cfg.rank), w1); w.w2 = try l.matrix("markov_head.markov_w2.weight", c.vocab, cfg.rank, false); + w.confidence = null; + w.confidence_bias = 0; + if (check.tensors.contains("confidence_head.proj.weight")) { + const width = c.hidden + cfg.rank; // the final norm's row, then W1 of the token before the draft + const proj = (try l.tensor("confidence_head.proj.weight", &.{ 1, width })).host(u16); + w.confidence = try l.buffer(width * 2); + @memcpy(w.confidence.?.slice(u16, width), proj); + w.confidence_bias = quant.toF32((try l.tensor("confidence_head.proj.bias", &.{1})).host(u16)[0]); + } return w; } diff --git a/zig/src/families/lfm2_moe/lfm2_moe.zig b/zig/src/families/lfm2_moe/lfm2_moe.zig index 8f861e31b4..75b40e5938 100644 --- a/zig/src/families/lfm2_moe/lfm2_moe.zig +++ b/zig/src/families/lfm2_moe/lfm2_moe.zig @@ -11,3 +11,4 @@ pub const dspark = @import("dspark.zig"); pub const dspark_config = @import("dspark_config.zig"); pub const dspark_block = @import("dspark_block.zig"); pub const dspark_quant = @import("dspark_quant.zig"); +pub const verify_costs = @import("verify_costs.zig"); diff --git a/zig/src/families/lfm2_moe/verify_costs.zig b/zig/src/families/lfm2_moe/verify_costs.zig new file mode 100644 index 0000000000..f80e4cdcef --- /dev/null +++ b/zig/src/families/lfm2_moe/verify_costs.zig @@ -0,0 +1,44 @@ +//! A verify cost table from tf-lfm2-verify-costs: shared forward ms by total rows, for the lane core's shared costs. +const std = @import("std"); +const lanes = @import("lanes"); +const st = @import("state.zig"); + +const Cost = lanes.config.Cost; + +/// Lines `streams rows_each total_rows ms` ('#' comments): each total's ms at the fewest streams that reach it. +pub fn parse(gpa: std.mem.Allocator, text: []const u8) ![]Cost { + var best: [st.batch_rows + 1]?struct { streams: u32, ms: f64 } = @splat(null); + var lines = std.mem.tokenizeAny(u8, text, "\r\n"); + while (lines.next()) |line| { + const body = std.mem.trim(u8, line, " \t"); + if (body.len == 0 or body[0] == '#') continue; + var fields = std.mem.tokenizeAny(u8, body, " \t"); + var v: [4][]const u8 = undefined; + for (&v) |*x| x.* = fields.next() orelse return error.BadVerifyCosts; + if (fields.next() != null) return error.BadVerifyCosts; + const streams = try std.fmt.parseInt(u32, v[0], 10); + const each = try std.fmt.parseInt(u32, v[1], 10); + const total = try std.fmt.parseInt(u32, v[2], 10); + const ms = try std.fmt.parseFloat(f64, v[3]); + if (streams == 0 or each == 0 or total != streams * each or total > st.batch_rows or !(ms > 0)) return error.BadVerifyCosts; + if (best[total]) |b| if (b.streams <= streams) continue; + best[total] = .{ .streams = streams, .ms = ms }; + } + var out: std.ArrayList(Cost) = .empty; + errdefer out.deinit(gpa); + for (best, 0..) |b, total| if (b) |x| try out.append(gpa, .{ .width = @intCast(total), .ms = x.ms }); + if (out.items.len == 0) return error.BadVerifyCosts; + return out.toOwnedSlice(gpa); +} + +test "lfm2 dspark: a verify cost table gives each total its fewest streams' ms and refuses malformed lines" { + const gpa = std.testing.allocator; + const t = try parse(gpa, "# streams rows_each total_rows ms\n1 1 1 3.7\n2 1 2 4.9\n1 2 2 4.6\n\n4 8 32 40.5\n"); + defer gpa.free(t); + try std.testing.expectEqual(@as(usize, 3), t.len); + try std.testing.expectEqual(@as(f64, 4.6), t[1].ms); // one stream of two rows, not two of one + try std.testing.expectEqual(@as(u32, 32), t[2].width); + try std.testing.expectError(error.BadVerifyCosts, parse(gpa, "2 2 5 4.0\n")); + try std.testing.expectError(error.BadVerifyCosts, parse(gpa, "8 8 64 9.0\n")); + try std.testing.expectError(error.BadVerifyCosts, parse(gpa, "# nothing\n")); +} diff --git a/zig/src/native/lfm2_moe.zig b/zig/src/native/lfm2_moe.zig index a78d67d97b..c8183fbac2 100644 --- a/zig/src/native/lfm2_moe.zig +++ b/zig/src/native/lfm2_moe.zig @@ -87,7 +87,8 @@ pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Ope release(h); return null; } - h.draft = q.dspark.Drafter.init(gpa, io, model, dir, .{ .capacity = window }) catch |err| { + const verify = verifyMode(o.draft_verify); + h.draft = q.dspark.Drafter.init(gpa, io, model, dir, .{ .capacity = window, .verify = verify }) catch |err| { problem.* = try std.fmt.allocPrint(a, "the native LFM2 drafter cannot load {s} ({s}): it takes LiquidAI/LFM2.5-24B-A2B-DSpark", .{ dir, @errorName(err) }); release(h); return null; @@ -109,7 +110,15 @@ pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Ope h.wrap.?.held_on_target = true; } errdefer if (h.wrap) |*w| w.deinit(); - const facts = if (h.wrap) |*w| w.facts(h.metal.facts()) else h.metal.facts(); + var facts = if (h.wrap) |*w| w.facts(h.metal.facts()) else h.metal.facts(); + const profiled = if (h.draft != null) verifyCosts(a, io, o.verify_costs) catch |err| { + problem.* = try std.fmt.allocPrint(a, "--verify-costs {s} is not a tf-lfm2-verify-costs table ({s})", .{ o.verify_costs.?, @errorName(err) }); + h.wrap.?.deinit(); + h.draft.?.deinit(); + release(h); + return null; + } else null; + if (profiled) |t| facts.shared_costs = t; h.config = try lanes.Config.init(gpa, facts, q.state.window_rows, q.state.window_rows - 1); errdefer h.config.deinit(gpa); h.clock = .{ .io = io }; @@ -150,10 +159,25 @@ pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Ope }; h.host.keepalive_target = .{ .ctx = &h.warm, .tick = mtl.keepalive.Target.tick }; try h.host.start(); - if (h.draft) |d| std.log.info("lfm2: DSpark drafts from --drafter, {d} a round at most, {d}-bit codes", .{ d.config.block, d.options.bits }); + if (h.draft) |d| std.log.info("lfm2: DSpark drafts from --drafter, {d} a round at most, {d}-bit codes, {t} verify", .{ d.config.block, d.options.bits, d.options.verify }); return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; } +/// --verify-costs: a tf-lfm2-verify-costs table's shared forward ms by total rows (null: none named; `a`-owned). +fn verifyCosts(a: std.mem.Allocator, io: std.Io, path: ?[]const u8) !?[]const lanes.config.Cost { + const file = path orelse return null; + const text = try std.Io.Dir.cwd().readFileAlloc(io, file, a, .limited(1 << 20)); + const t = try q.verify_costs.parse(a, text); + std.log.info("lfm2: verify costs from {s}, {d} totals of rows", .{ file, t.len }); + return t; +} + +/// --draft-verify as the drafter's verify (null: the default). +fn verifyMode(v: ?api.DraftVerify) q.dspark.Verify { + const asked = v orelse return q.dspark.default_verify; + return std.meta.stringToEnum(q.dspark.Verify, @tagName(asked)).?; +} + /// The prompt cache's copies of a stream: the target's state at a mark (backend.zig), and the drafter's context there. const Snaps = struct { const S = q.backend.Snaps; diff --git a/zig/src/native/metal.zig b/zig/src/native/metal.zig index 33533c3597..fa60fd2426 100644 --- a/zig/src/native/metal.zig +++ b/zig/src/native/metal.zig @@ -331,6 +331,7 @@ fn refusal(a: Allocator, o: api.Open, dense: bool) !?[]const u8 { const one_mac: ?[]const u8 = if (std.mem.eql(u8, t, "nemotron_h")) "Nemotron" else if (std.mem.eql(u8, t, "lfm2_moe")) "LFM2.5-24B-A2B" else if (isGemma4(t)) "Gemma 4" else if (std.mem.eql(u8, t, "step5_vl")) "Step-5" else if (std.mem.eql(u8, t, "spark2_5")) "Spark-X2.6" else if (!std.mem.eql(u8, t, "qwen3_5")) null else if (dense) "Qwen3.8-27B" else "Qwen3.5-2B"; if (o.speed_up != null) if (one_mac) |name| return try std.fmt.allocPrint(a, "--speed-up joins two Macs for Flash Next, GLM-5.3-Flash and GLM-5.3 only: serve {s} without it", .{name}); if (o.recall and std.mem.eql(u8, t, "qwen4_exp")) return recall_refusal; + if ((o.draft_verify != null or o.verify_costs != null) and (o.drafter == null or !std.mem.eql(u8, t, "lfm2_moe"))) return "--draft-verify and --verify-costs set the verify of LFM2.5-24B-A2B's DSpark drafter: serve it with --drafter"; if (spill.refusal(o)) |why| return why; if (o.learn == null) return null; // the 27B writes conversation states diff --git a/zig/src/server/cli.zig b/zig/src/server/cli.zig index d36a933aa7..20129f9d2f 100644 --- a/zig/src/server/cli.zig +++ b/zig/src/server/cli.zig @@ -101,6 +101,8 @@ pub const flags = [_]Flag{ .{ .name = "--slide-graph", .native = true }, .{ .name = "--drafter", .native = true }, .{ .name = "--drafter-bits", .choices = &.{ "0", "4" }, .native = true }, + .{ .name = "--draft-verify", .choices = &.{ "depth", "fixed", "scheduled" }, .native = builtin.os.tag == .macos }, + .{ .name = "--verify-costs", .native = builtin.os.tag == .macos }, .{ .name = "--record", .native = true }, .{ .name = "--kv-cache", .choices = &.{ "bf16", "fp8" }, .native = true }, .{ .name = "--prompt-precision", .choices = &.{ "bf16", "int8" }, .native = true }, @@ -211,6 +213,8 @@ pub const Args = struct { compact_memory: ?[]const u8 = null, drafter: ?[]const u8 = null, drafter_bits: u8 = 4, + draft_verify: ?api.DraftVerify = null, // the LFM2 DSpark drafter's verify rows (null: the engine's default) + verify_costs: ?[]const u8 = null, // a tf-lfm2-verify-costs table pricing them (null: costs timed at load) record: ?[]const u8 = null, // drafter training rows written under this folder kv_cache: []const u8 = "bf16", // fp8: e4m3 attention caches with a scale a row, about twice the context prompt_precision: []const u8 = "bf16", // int8: prompt matmuls quantize their inputs to int8 a 64-input group @@ -414,7 +418,7 @@ fn apply(a: Allocator, out: *Args, name: []const u8, value: ?[]const u8, u: *Usa out.learn_daily_gib = try gib(u, a, name, v); } else if (is(name, "--prompt-seed")) { out.prompt_seed = true; - } else if (is(name, "--max-tokens")) out.max_tokens = try int(u, a, name, v) else if (is(name, "--temperature")) out.temperature = try float(u, a, name, v) else if (is(name, "--top-p")) out.top_p = try float(u, a, name, v) else if (is(name, "--top-k")) out.top_k = try int(u, a, name, v) else if (is(name, "--min-p")) out.min_p = try float(u, a, name, v) else if (is(name, "--thinking")) out.thinking = true else if (is(name, "--no-thinking")) out.thinking = false else if (is(name, "--reasoning-effort")) out.reasoning_effort = v else if (is(name, "--thinking-budget")) out.thinking_budget = try int(u, a, name, v) else if (is(name, "--thinking-reserve")) out.thinking_reserve = try int(u, a, name, v) else if (is(name, "--loop-guard")) out.loop_guard = true else if (is(name, "--drafter")) out.drafter = v else if (is(name, "--drafter-bits")) out.drafter_bits = @intCast(try int(u, a, name, v)) else if (is(name, "--record")) out.record = v else if (is(name, "--kv-cache")) out.kv_cache = v else if (is(name, "--prompt-precision")) out.prompt_precision = v else if (is(name, "--no-drafts")) out.no_drafts = true else if (is(name, "--keep-warm")) out.keep_warm = try int(u, a, name, v) else if (is(name, "--idle-unload")) out.idle_unload = try int(u, a, name, v) else if (is(name, "--compact-at")) { + } else if (is(name, "--max-tokens")) out.max_tokens = try int(u, a, name, v) else if (is(name, "--temperature")) out.temperature = try float(u, a, name, v) else if (is(name, "--top-p")) out.top_p = try float(u, a, name, v) else if (is(name, "--top-k")) out.top_k = try int(u, a, name, v) else if (is(name, "--min-p")) out.min_p = try float(u, a, name, v) else if (is(name, "--thinking")) out.thinking = true else if (is(name, "--no-thinking")) out.thinking = false else if (is(name, "--reasoning-effort")) out.reasoning_effort = v else if (is(name, "--thinking-budget")) out.thinking_budget = try int(u, a, name, v) else if (is(name, "--thinking-reserve")) out.thinking_reserve = try int(u, a, name, v) else if (is(name, "--loop-guard")) out.loop_guard = true else if (is(name, "--drafter")) out.drafter = v else if (is(name, "--drafter-bits")) out.drafter_bits = @intCast(try int(u, a, name, v)) else if (is(name, "--draft-verify")) out.draft_verify = std.meta.stringToEnum(api.DraftVerify, v) else if (is(name, "--verify-costs")) out.verify_costs = v else if (is(name, "--record")) out.record = v else if (is(name, "--kv-cache")) out.kv_cache = v else if (is(name, "--prompt-precision")) out.prompt_precision = v else if (is(name, "--no-drafts")) out.no_drafts = true else if (is(name, "--keep-warm")) out.keep_warm = try int(u, a, name, v) else if (is(name, "--idle-unload")) out.idle_unload = try int(u, a, name, v) else if (is(name, "--compact-at")) { if (std.mem.eql(u8, v, "auto")) out.compact_auto = true else { const f = try float(u, a, name, v); if (!(f > 0 and f <= 1)) return fail(u, a, "argument --compact-at: expected auto or a fraction in (0, 1]: '{s}'", .{v}); @@ -541,6 +545,8 @@ pub fn request(a: Allocator, dir: []const u8, model_type: []const u8, args: Args .drafts = !args.no_drafts, .drafter = args.drafter, .drafter_bits = args.drafter_bits, + .draft_verify = args.draft_verify, + .verify_costs = args.verify_costs, .record = args.record, .kv8 = std.mem.eql(u8, args.kv_cache, "fp8"), .prompt_int8 = std.mem.eql(u8, args.prompt_precision, "int8"), diff --git a/zig/src/server/cli_test.zig b/zig/src/server/cli_test.zig index f85890a632..7c8668186d 100644 --- a/zig/src/server/cli_test.zig +++ b/zig/src/server/cli_test.zig @@ -165,6 +165,13 @@ test "parse and capabilities share the table" { try std.testing.expectEqualStrings("x", (try parse(a, &.{ "m", "--drafter", "x" }, &u)).drafter.?); try std.testing.expectEqual(@as(u8, 0), (try parse(a, &.{ "m", "--drafter-bits", "0" }, &u)).drafter_bits); try std.testing.expectError(error.Usage, parse(a, &.{ "m", "--drafter-bits", "8" }, &u)); + if (@import("builtin").os.tag == .macos) { + const verify = try parse(a, &.{ "m", "--draft-verify", "scheduled", "--verify-costs", "/t" }, &u); + try std.testing.expectEqual(api.DraftVerify.scheduled, verify.draft_verify.?); + try std.testing.expectEqual(api.DraftVerify.scheduled, (try request(a, "d", "t", verify)).draft_verify.?); + try std.testing.expectEqualStrings("/t", (try request(a, "d", "t", verify)).verify_costs.?); + try std.testing.expectError(error.Usage, parse(a, &.{ "m", "--draft-verify", "wide" }, &u)); + } try std.testing.expectEqualStrings("/r", (try parse(a, &.{ "m", "--record", "/r" }, &u)).record.?); try std.testing.expectEqualStrings("fp8", (try parse(a, &.{ "m", "--kv-cache", "fp8" }, &u)).kv_cache); try std.testing.expectError(error.Usage, parse(a, &.{ "m", "--kv-cache", "int4" }, &u)); diff --git a/zig/src/server/main.zig b/zig/src/server/main.zig index 5dab0c4568..dcb493118e 100644 --- a/zig/src/server/main.zig +++ b/zig/src/server/main.zig @@ -25,7 +25,7 @@ const usage_line = "usage: tensorfold serve [-h] [--host HOST] [--port PORT] [-- " [--thinking | --no-thinking] [--reasoning-effort {low,medium,high,xhigh}]" ++ " [--thinking-budget THINKING_BUDGET] [--thinking-reserve THINKING_RESERVE] [--loop-guard] [--no-drafts]" ++ " [--drafter DRAFTER]" ++ - " [--drafter-bits {0,4}] [--record RECORD] [--keep-warm SECONDS] [--idle-unload SECONDS]" ++ + " [--drafter-bits {0,4}] [--draft-verify {depth,fixed,scheduled}] [--verify-costs VERIFY_COSTS] [--record RECORD] [--keep-warm SECONDS] [--idle-unload SECONDS]" ++ " [--silence-floor-dbfs DBFS] [--silence-seconds SECONDS]" ++ " [--silence-onset-s SECONDS] [--silence-onset-voiced-s SECONDS]" ++ " [--silence-mid-s SECONDS] [--silence-mid-min-voiced-s SECONDS] [--silence-redraws N]" ++ diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 4a278138ac..d4c3dab915 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -87,6 +87,7 @@ test { _ = @import("families/lfm2_moe/dspark_config.zig"); _ = @import("families/lfm2_moe/dspark_quant.zig"); _ = @import("families/lfm2_moe/dspark_test.zig"); + _ = @import("families/lfm2_moe/verify_costs.zig"); } pub const lane_projection = @import("core/lane_projection.zig"); diff --git a/zig/tests/lfm2_dspark.zig b/zig/tests/lfm2_dspark.zig index f95ce615e5..53157c06a1 100644 --- a/zig/tests/lfm2_dspark.zig +++ b/zig/tests/lfm2_dspark.zig @@ -125,7 +125,7 @@ pub fn main(init: std.process.Init) !void { defer pool.pop(); const m = try q.Model.load(gpa, init.io, args[1]); defer m.deinit(); - const d = try q.dspark.Drafter.init(gpa, init.io, m, args[2], .{ .capacity = length + 64, .bits = bits }); + const d = try q.dspark.Drafter.init(gpa, init.io, m, args[2], .{ .capacity = length + 64, .bits = bits, .verify = .scheduled }); defer d.deinit(); const cfg = d.config; // the compared block's rows (a hold of depth d runs a d-row block) @@ -156,6 +156,8 @@ pub fn main(init: std.process.Init) !void { _ = try q.dspark_block.hold(d, &.{.{ .lane = lane, .pending = ids[ctx], .anchor = ctx, .depth = block }}, null); var held: [16]u32 = undefined; @memcpy(held[0..block], lane.held[0..block]); + var conf: [16]f32 = undefined; + @memcpy(conf[0..block], lane.conf[0..block]); const gpu = held[0..block]; for ([_]usize{ 1, 4, 16 }) |rows| { var best: f64 = std.math.floatMax(f64); @@ -269,6 +271,9 @@ pub fn main(init: std.process.Init) !void { const w1 = try r.get("markov_head.markov_w1.weight"); const w2 = try r.get("markov_head.markov_w2.weight"); const bias = try a.alloc(f32, c.vocab); + const proj = try r.get("confidence_head.proj.weight"); + const proj_bias = (try r.get("confidence_head.proj.bias"))[0]; + var conf_worst: f32 = 0; var prev = ids[ctx]; var same: usize = 0; var worst: f32 = 0; @@ -294,8 +299,16 @@ pub fn main(init: std.process.Init) !void { same += @intFromBool(ok); std.debug.print("step {d}: reference {d} (margin {d:.3}, second {d}), GPU {d} {s}\n", .{ i, best, margin, second, gpu[i], if (ok) "equal" else "DIFFERENT" }); if (!ok) worst = @max(worst, margin); + // the confidence head: [final norm row i, W1 of the token before draft i] . proj + bias, through a sigmoid + var logit: f32 = proj_bias; + for (h[i * c.hidden ..][0..c.hidden], proj[0..c.hidden]) |x0, w0| logit += x0 * w0; + for (w1[@as(usize, prev) * cfg.rank ..][0..cfg.rank], proj[c.hidden..]) |x0, w0| logit += x0 * w0; + const want = 1 / (1 + @exp(-logit)); + conf_worst = @max(conf_worst, @abs(want - conf[i])); + std.debug.print("step {d}: confidence reference {d:.4}, GPU {d:.4}\n", .{ i, want, conf[i] }); prev = gpu[i]; // teacher-forced: step i + 1 follows the GPU's pick, so each step is compared alone } std.debug.print("the reference's largest margin over a different GPU pick: {d:.3}\n", .{worst}); + std.debug.print("the confidences' worst difference from the fp32 reference: {d:.5}\n", .{conf_worst}); std.debug.print("{d} of {d} drafts equal the fp32 reference's picks on the same chain ({d}-token context)\n", .{ same, n, ctx }); } diff --git a/zig/tests/lfm2_verify_costs.zig b/zig/tests/lfm2_verify_costs.zig new file mode 100644 index 0000000000..eef8ad7be3 --- /dev/null +++ b/zig/tests/lfm2_verify_costs.zig @@ -0,0 +1,67 @@ +//! LFM2 verify forwards timed by stream count and rows a stream: the round costs a DSpark verify budget weighs. +const std = @import("std"); +const mtl = @import("metal"); +const q = @import("tensorfold").lfm2_moe; +const st = q.state; +const c = q.config; + +fn now() i128 { + var ts: std.c.timespec = undefined; + _ = std.c.clock_gettime(.MONOTONIC, &ts); + return @as(i128, ts.sec) * std.time.ns_per_s + ts.nsec; +} + +/// Stream counts profiled; each runs every rows-a-stream width whose total fits a shared forward. +const counts = [_]usize{ 1, 2, 3, 4, 6, 8, 12, 16 }; +const reps = 9; + +pub fn main(init: std.process.Init) !void { + const gpa = init.gpa; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 3 or args.len > 4) return error.ExpectedModelOutputHistory; + const history: usize = if (args.len == 4) try std.fmt.parseInt(usize, args[3], 10) else 1024; + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const m = try q.Model.load(gpa, init.io, args[1]); + defer m.deinit(); + var scratch = try st.Scratch.init(gpa, m.device, &m.config.attention, @max(history, st.batch_rows)); + defer scratch.deinit(); + var caches: [16]st.Cache = undefined; + var made: usize = 0; + defer for (caches[0..made]) |*k| k.deinit(); + const ids = try gpa.alloc(u32, history); + defer gpa.free(ids); + for (caches[0..], 0..) |*k, i| { + k.* = try st.Cache.init(gpa, m.device, &m.config.attention, history + st.window_rows + 1); + made += 1; + for (ids, 0..) |*id, j| id.* = @intCast((j * 7919 + 13 + i * 104729) % 120000); + try q.forward.run(m, &scratch, &.{.{ .cache = k, .rows = history }}, ids, false, .last); + } + var out: std.Io.Writer.Allocating = .init(gpa); + defer out.deinit(); + const w = &out.writer; + try w.print("# LFM2 verify forward ms (median of {d}) by streams and rows a stream, {d}-token histories\n", .{ reps, history }); + try w.print("# streams rows_each total_rows ms\n", .{}); + var segments: [st.batch_rows]q.forward.Segment = undefined; + var tokens: [st.batch_rows]u32 = undefined; + for (&tokens, 0..) |*t, i| t.* = @intCast((i * 31 + 7) % 120000); + for (counts) |streams| { + var each: usize = 1; + while (each <= st.window_rows and streams * each <= st.batch_rows) : (each += 1) { + for (segments[0..streams], caches[0..streams]) |*s, *k| s.* = .{ .cache = k, .rows = each }; + var times: [reps + 2]f64 = undefined; + for (×) |*t| { + const t0 = now(); + try q.forward.run(m, &scratch, segments[0..streams], tokens[0 .. streams * each], true, .all); + t.* = @as(f64, @floatFromInt(now() - t0)) / 1e6; + for (caches[0..streams]) |*k| k.len = history; + } + const kept = times[2..]; // the first two calls warm a new width's pipelines + std.mem.sort(f64, kept, {}, std.sort.asc(f64)); + const ms = kept[reps / 2]; + try w.print("{d} {d} {d} {d:.3}\n", .{ streams, each, streams * each, ms }); + std.debug.print("{d} streams x {d} rows ({d}): {d:.3} ms\n", .{ streams, each, streams * each, ms }); + } + } + try std.Io.Dir.cwd().writeFile(init.io, .{ .sub_path = args[2], .data = out.written() }); +} From 0f609b3241f04c70050659260bee91de4e070852 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:37:47 +0100 Subject: [PATCH 22/46] lfm2: a short prompt's tapped rows wait in the DSpark stage like verified rows Each prompt chunk's rows were absorbed at once in a command buffer of their own, waited on before the first token, and the prompt's last row in a second one at the first draft request. The prefill pass is waited on before its rows are handed over, so a chunk that fits the lane's stage (128 rows) can be copied there on the host like a round's verified rows. The next hold, or a kept prompt state, absorbs the stage first, as before. Longer chunks are absorbed at once. --- zig/src/families/lfm2_moe/dspark.zig | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/zig/src/families/lfm2_moe/dspark.zig b/zig/src/families/lfm2_moe/dspark.zig index 0263d03f64..404c23516c 100644 --- a/zig/src/families/lfm2_moe/dspark.zig +++ b/zig/src/families/lfm2_moe/dspark.zig @@ -348,7 +348,7 @@ pub const Drafter = struct { }; } - /// Rows from `start` replace the context there: verified rows wait on the host for the next hold, others run now. + /// Rows from `start` replace the context there: rows that fit the stage wait on the host for the next hold. fn absorbFn(ptr: *anyopaque, items: []const dr.Absorb) anyerror!void { const d = cast(ptr); try d.settle(); @@ -371,7 +371,7 @@ pub const Drafter = struct { // rows from `start` replace what was staged past it; rows before the absorbed end go to the GPU now const keep: u32 = if (it.start >= l.end) @intCast(it.start - l.end) else 0; l.staged = @min(l.staged, keep); - if (!it.prompt and it.start >= l.end and l.staged + f.rows <= stage_rows) { + if (it.start >= l.end and l.staged + f.rows <= stage_rows) { const from = (src.contents() + f.offset)[0 .. f.rows * width]; @memcpy((l.stage.contents() + l.staged * width)[0 .. f.rows * width], from); l.staged += f.rows; From df6743bfc687a013795a72cad3dc66f4eb91de9a Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:38:01 +0100 Subject: [PATCH 23/46] lfm2: DSpark prices its own hold pass for the lane core A DSpark hold is a pass of its own: the drafter runs its five layers and the target's head for every drafting stream before the verify. The lane core can now weigh that pass (Facts.copy_first, Facts.hold_costs, Model.own_copy_rates); the drafter gives it what it needs. - copy_first: a stream with a copy ahead verifies the copy, and the drafter is asked no hold for it. - hold_costs: after its warm-up the drafter times holds of 1, 2, 4, 8 and 16 blocks at one draft and at four (two at 16 blocks, one 32-row pass) on empty lanes, best of three. Each settled one-pass hold then moves the timed passes at its block counts a tenth of the way toward its own GPU ms, so the price follows the served contexts. - own_copy_rates: with a drafter, a stream's copy windows are priced by its own copies. Test: a 4-block hold at twice its price raises the 4-block passes by a tenth and leaves the others; one between 4 and 8 blocks moves both; a hold without GPU time teaches nothing. --- zig/src/families/lfm2_moe/dspark.zig | 56 +++++++++++++++++++++-- zig/src/families/lfm2_moe/dspark_test.zig | 22 +++++++++ zig/src/native/lfm2_moe.zig | 2 + 3 files changed, 75 insertions(+), 5 deletions(-) diff --git a/zig/src/families/lfm2_moe/dspark.zig b/zig/src/families/lfm2_moe/dspark.zig index 404c23516c..eb3fa4d90f 100644 --- a/zig/src/families/lfm2_moe/dspark.zig +++ b/zig/src/families/lfm2_moe/dspark.zig @@ -20,6 +20,8 @@ pub const pass_rows = 32; pub const absorb_rows = 2048; /// Verified rows a lane keeps on the host until its next hold absorbs them (a round's rows for no GPU work). pub const stage_rows = 128; +/// The blocks a hold is timed with at load, each at one draft and at four (two at 16 blocks: one 32-row pass). +const hold_grid = [_]u32{ 1, 2, 4, 8, 16 }; const kw = c.kv_heads * c.head_dim; /// Verify rows by the core's depth rule, the whole block (fixed), or the confidence head's chances (scheduled). @@ -85,6 +87,8 @@ pub const Drafter = struct { /// One stream's hold wall ms at full depth, and that over the block (the depth rule's cost a draft). hold_ms: f64 = 0, step_ms: f64 = 0, + holds: [hold_grid.len * 2]lanes.hold_price.HoldCost = undefined, // hold passes timed at load, then by settled holds + holds_len: usize = 0, event: mtl.Event = undefined, signalled: u64 = 0, pending: ?mtl.CommandBuffer = null, @@ -187,7 +191,7 @@ pub const Drafter = struct { fn factsFn(ptr: *anyopaque) dr.Facts { const d = cast(ptr); - var f: dr.Facts = .{ .depth = d.config.block, .step_ms = d.step_ms, .prior = prior[0..d.config.block], .batched = true, .shared_guard = true, .stream_bytes = d.laneBytes(d.options.capacity) }; + var f: dr.Facts = .{ .depth = d.config.block, .step_ms = d.step_ms, .prior = prior[0..d.config.block], .batched = true, .shared_guard = true, .stream_bytes = d.laneBytes(d.options.capacity), .copy_first = true }; switch (d.options.verify) { .depth => {}, .fixed => { // every round the whole block, no plain stretches @@ -195,7 +199,10 @@ pub const Drafter = struct { f.confidence_floor = 0; f.shared_guard = false; }, - .scheduled => f.chances = d.confident(), + .scheduled => { + f.chances = d.confident(); + f.hold_costs = d.holds[0..d.holds_len]; + }, } return f; } @@ -461,9 +468,29 @@ pub const Drafter = struct { std.log.err("lfm2 dspark: a hold failed on the GPU: {s}", .{msg}); return error.DraftGpuFailed; } + d.priceHold((cb.gpuEnd() - cb.gpuStart()) * 1000); block.readPicks(d, d.pass[0..d.pass_len], d.depths[0..d.pass_len]); } + /// A settled hold's GPU ms against its price: the timed passes at its block counts move a tenth of the way to it. + pub fn priceHold(d: *Drafter, ms: f64) void { + const timed = d.holds[0..d.holds_len]; + if (timed.len == 0 or d.pass_len == 0 or !(ms > 0)) return; + const deepest = std.mem.max(u32, d.depths[0..d.pass_len]); + const priced = lanes.hold_price.holdMs(timed, d.pass_len, deepest); + if (!(priced > 0)) return; + const step = 1 + 0.1 * (ms / priced - 1); + var lo: u32 = 0; + var hi: u32 = std.math.maxInt(u32); + for (timed) |h| { + if (h.streams <= d.pass_len) lo = @max(lo, h.streams); + if (h.streams >= d.pass_len) hi = @min(hi, h.streams); + } + for (timed) |*h| if (h.streams >= lo and h.streams <= hi) { + h.ms *= step; + }; + } + /// The signalled hold settled and its lanes' picks forgotten (a lane may be dropped, or a hold may overwrite). fn unpass(d: *Drafter) !void { defer { @@ -490,10 +517,13 @@ pub const Drafter = struct { return heldFn(ptr, &.{s}, &.{out}); } - /// Holds on an empty context: the pipelines compiled, then one-draft and full-depth holds timed (best of three). + /// Holds on empty contexts: the pipelines compiled, then holds timed by blocks and depth (best of three). fn warm(d: *Drafter) !void { - const l = try d.newLane(64); - defer d.dropLane(l); + var made: [hold_grid[hold_grid.len - 1]]*Lane = undefined; + var n: usize = 0; + defer for (made[0..n]) |l| d.dropLane(l); + while (n < made.len) : (n += 1) made[n] = try d.newLane(64); + const l = made[0]; var best: [2]f64 = @splat(std.math.floatMax(f64)); for (0..4) |i| for ([_]u32{ 1, d.config.block }, &best) |depth, *b| { const t0 = std.Io.Clock.awake.now(d.io); @@ -504,5 +534,21 @@ pub const Drafter = struct { d.hold_ms = best[1]; d.step_ms = best[1] / @as(f64, @floatFromInt(d.config.block)); std.log.info("lfm2 dspark: a 1-draft hold takes {d:.2} ms, a {d}-draft hold {d:.2} ms", .{ best[0], d.config.block, best[1] }); + var blocks: [hold_grid[hold_grid.len - 1]]block.Block = undefined; + for (made, &blocks) |lane, *b| b.* = .{ .lane = lane, .pending = 1, .anchor = 0, .depth = 1 }; + d.holds_len = 0; + for (hold_grid) |streams| for ([_]u32{ 1, @min(4, pass_rows / streams) }, 0..) |depth, j| { + if (j == 1 and depth == 1) continue; + for (blocks[0..streams]) |*b| b.depth = depth; + var ms_best = std.math.floatMax(f64); + for (0..4) |i| { + const t0 = std.Io.Clock.awake.now(d.io); + _ = try block.hold(d, blocks[0..streams], null); + const ms = @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(d.io)).toNanoseconds())) / 1e6; + if (i > 0) ms_best = @min(ms_best, ms); + } + d.holds[d.holds_len] = .{ .streams = streams, .depth = depth, .ms = ms_best }; + d.holds_len += 1; + }; } }; diff --git a/zig/src/families/lfm2_moe/dspark_test.zig b/zig/src/families/lfm2_moe/dspark_test.zig index 8fd1eb32e1..385866b291 100644 --- a/zig/src/families/lfm2_moe/dspark_test.zig +++ b/zig/src/families/lfm2_moe/dspark_test.zig @@ -4,6 +4,8 @@ const mtl = @import("metal"); const c = @import("config.zig"); const kernels = @import("kernels.zig"); const quant = @import("dspark_quant.zig"); +const ds = @import("dspark.zig"); +const lanes = @import("lanes"); const shared = mtl.ResourceOptions.shared; @@ -211,3 +213,23 @@ test "lfm2 dspark: the drafter verifies by its confidence head unless asked othe const o: @import("dspark.zig").Options = .{ .capacity = 64 }; try std.testing.expectEqual(@import("dspark.zig").Verify.scheduled, o.verify); } + +test "lfm2 dspark: a settled hold moves the timed passes at its block counts a tenth of the way to its GPU ms" { + var d: ds.Drafter = undefined; + const grid = [_]lanes.hold_price.HoldCost{ .{ .streams = 1, .depth = 1, .ms = 1 }, .{ .streams = 4, .depth = 1, .ms = 2 }, .{ .streams = 4, .depth = 4, .ms = 4 }, .{ .streams = 8, .depth = 1, .ms = 3 } }; + @memcpy(d.holds[0..grid.len], &grid); + d.holds_len = grid.len; + d.pass_len = 4; + d.depths[0..4].* = .{ 1, 1, 1, 1 }; + d.priceHold(4); // priced 2 ms: the 4-block passes rise by a tenth of the doubling + try std.testing.expectApproxEqAbs(@as(f64, 2.2), d.holds[1].ms, 1e-12); + try std.testing.expectApproxEqAbs(@as(f64, 4.4), d.holds[2].ms, 1e-12); + try std.testing.expectEqual(@as(f64, 1), d.holds[0].ms); + try std.testing.expectEqual(@as(f64, 3), d.holds[3].ms); + d.pass_len = 6; // between 4 and 8 blocks: both move + d.depths[0..6].* = .{ 1, 1, 1, 1, 1, 1 }; + d.priceHold(0); + try std.testing.expectEqual(@as(f64, 2.2), d.holds[1].ms); // no GPU time: nothing learned + d.priceHold(2.6 * 2); + try std.testing.expect(d.holds[1].ms > 2.2 and d.holds[3].ms > 3 and d.holds[0].ms == 1); +} diff --git a/zig/src/native/lfm2_moe.zig b/zig/src/native/lfm2_moe.zig index c8183fbac2..767cbbe975 100644 --- a/zig/src/native/lfm2_moe.zig +++ b/zig/src/native/lfm2_moe.zig @@ -119,6 +119,8 @@ pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Ope return null; } else null; if (profiled) |t| facts.shared_costs = t; + // a drafting stream prices its copy windows by its own copies: on distinct prompts they keep breaking + if (h.wrap != null) facts.own_copy_rates = true; h.config = try lanes.Config.init(gpa, facts, q.state.window_rows, q.state.window_rows - 1); errdefer h.config.deinit(gpa); h.clock = .{ .io = io }; From 262208f321dd230a3b76974c3184d6b174bbeeaa Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:38:12 +0100 Subject: [PATCH 24/46] lfm2: tf-lfm2-dspark times shared holds; tf-lfm2-verify-costs reads real rows tf-lfm2-dspark times a shared round's hold over 2 to 16 streams, by the wall clock and on the GPU, and absorbs of 128 rows and of a whole absorb pass, so a hold's price can be checked against a verify's. tf-lfm2-verify-costs reads token id files after HISTORY: each of the 16 streams takes its history and then its verify rows from its own sixteenth of the files' ids (the history shrinks to fit), so a shared round's rows route the experts their own text routes. Without id files the rows are synthetic as before; the table's header names their kind. --- zig/tests/lfm2_dspark.zig | 42 ++++++++++++++++++++++++++++++++- zig/tests/lfm2_verify_costs.zig | 30 +++++++++++++++++++---- 2 files changed, 66 insertions(+), 6 deletions(-) diff --git a/zig/tests/lfm2_dspark.zig b/zig/tests/lfm2_dspark.zig index 53157c06a1..5ee8140de2 100644 --- a/zig/tests/lfm2_dspark.zig +++ b/zig/tests/lfm2_dspark.zig @@ -109,6 +109,16 @@ fn compare(label: []const u8, gpu: []const u16, row: []const f32, bias: []const std.debug.print("{s}: cosine {d:.5}, worst difference {d:.3}\n", .{ label, dot / @sqrt(gg * rr), worst }); } +/// A hold's GPU ms: the blocks held in a signalled command buffer, waited on and read for its GPU times. +fn gpuHold(d: *q.dspark.Drafter, blocks: []const q.dspark_block.Block) !f64 { + d.signalled += 1; + const cb = (try q.dspark_block.hold(d, blocks, .{ .event = d.event, .value = d.signalled })).?; + defer cb.release(); + cb.wait(); + if (cb.failure() != null) return error.GpuFailed; + return (cb.gpuEnd() - cb.gpuStart()) * 1000; +} + fn absorb(d: *q.dspark.Drafter, piece: q.dspark_block.Piece) !void { const run = q.dspark_block.begin(d); try q.dspark_block.absorbPass(d, run.e, &.{piece}); @@ -159,7 +169,8 @@ pub fn main(init: std.process.Init) !void { var conf: [16]f32 = undefined; @memcpy(conf[0..block], lane.conf[0..block]); const gpu = held[0..block]; - for ([_]usize{ 1, 4, 16 }) |rows| { + for ([_]usize{ 1, 4, 16, 128, q.dspark.absorb_rows }) |rows| { + if (rows > ctx) continue; var best: f64 = std.math.floatMax(f64); for (0..5) |_| { const t0 = std.Io.Clock.awake.now(init.io); @@ -178,6 +189,35 @@ pub fn main(init: std.process.Init) !void { } std.debug.print("hold at depth {d}: {d:.2} ms (best of 5)\n", .{ depth, best }); } + // a shared round's hold: S streams' blocks of d rows in one pass (S x d at most pass_rows), each on this context + var lanes_: [16]*q.dspark.Lane = undefined; + lanes_[0] = lane; + var made: usize = 1; + defer for (lanes_[1..made]) |l| d.dropLane(l); + while (made < lanes_.len) : (made += 1) { + const l = try d.newLane(length + 32); + lanes_[made] = l; + var from: usize = 0; + while (from < ctx) : (from += q.dspark.absorb_rows) { + const rows = @min(q.dspark.absorb_rows, ctx - from); + try absorb(d, .{ .lane = l, .src = .{ .buf = taps.buf, .off = from * width * 2 }, .rows = rows, .p0 = from }); + } + l.end = ctx; + } + for ([_]usize{ 2, 4, 8, 16 }) |streams| for ([_]u32{ 1, 2, 3, 4, 8 }) |depth| { + if (streams * depth > q.dspark.pass_rows) continue; + var blocks: [16]q.dspark_block.Block = undefined; + for (blocks[0..streams], lanes_[0..streams]) |*b, l| b.* = .{ .lane = l, .pending = ids[ctx], .anchor = ctx, .depth = depth }; + var best: f64 = std.math.floatMax(f64); + var gpu_best: f64 = std.math.floatMax(f64); + for (0..5) |_| { + const t0 = std.Io.Clock.awake.now(init.io); + _ = try q.dspark_block.hold(d, blocks[0..streams], null); + best = @min(best, @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(init.io)).toNanoseconds())) / 1e6); + gpu_best = @min(gpu_best, try gpuHold(d, blocks[0..streams])); + } + std.debug.print("hold of {d} streams at depth {d}: {d:.2} ms, GPU {d:.2} ms (best of 5)\n", .{ streams, depth, best, gpu_best }); + }; _ = try q.dspark_block.hold(d, &.{.{ .lane = lane, .pending = ids[ctx], .anchor = ctx, .depth = block }}, null); // the host reference, fp32 throughout var r: Ref = .{ .gpa = gpa, .check = ck.Checkpoint.init(gpa), .cfg = cfg }; diff --git a/zig/tests/lfm2_verify_costs.zig b/zig/tests/lfm2_verify_costs.zig index eef8ad7be3..f3100d633e 100644 --- a/zig/tests/lfm2_verify_costs.zig +++ b/zig/tests/lfm2_verify_costs.zig @@ -15,11 +15,27 @@ fn now() i128 { const counts = [_]usize{ 1, 2, 3, 4, 6, 8, 12, 16 }; const reps = 9; +fn readIds(gpa: std.mem.Allocator, io: std.Io, path: []const u8, list: *std.ArrayList(u32)) !void { + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 22)); + defer gpa.free(bytes); + var it = std.mem.tokenizeAny(u8, bytes, " \n\r\t,"); + while (it.next()) |t| try list.append(gpa, try std.fmt.parseInt(u32, t, 10)); +} + pub fn main(init: std.process.Init) !void { const gpa = init.gpa; const args = try init.minimal.args.toSlice(init.arena.allocator()); - if (args.len < 3 or args.len > 4) return error.ExpectedModelOutputHistory; - const history: usize = if (args.len == 4) try std.fmt.parseInt(usize, args[3], 10) else 1024; + if (args.len < 3) return error.ExpectedModelOutputHistory; + var history: usize = if (args.len >= 4) try std.fmt.parseInt(usize, args[3], 10) else 1024; + // Token id files: every file's ids end to end; stream i reads its history and then its rows from its own sixteenth. + var real: std.ArrayList(u32) = .empty; + defer real.deinit(gpa); + for (args[@min(args.len, 4)..]) |path| try readIds(gpa, init.io, path, &real); + const part = real.items.len / 16; + if (real.items.len > 0) { + if (part < st.window_rows + 64) return error.IdFilesTooShort; + history = @min(history, part - st.window_rows); + } const pool = mtl.objc.Pool.push(); defer pool.pop(); const m = try q.Model.load(gpa, init.io, args[1]); @@ -34,13 +50,13 @@ pub fn main(init: std.process.Init) !void { for (caches[0..], 0..) |*k, i| { k.* = try st.Cache.init(gpa, m.device, &m.config.attention, history + st.window_rows + 1); made += 1; - for (ids, 0..) |*id, j| id.* = @intCast((j * 7919 + 13 + i * 104729) % 120000); + for (ids, 0..) |*id, j| id.* = if (part > 0) real.items[i * part + j] else @intCast((j * 7919 + 13 + i * 104729) % 120000); try q.forward.run(m, &scratch, &.{.{ .cache = k, .rows = history }}, ids, false, .last); } var out: std.Io.Writer.Allocating = .init(gpa); defer out.deinit(); const w = &out.writer; - try w.print("# LFM2 verify forward ms (median of {d}) by streams and rows a stream, {d}-token histories\n", .{ reps, history }); + try w.print("# LFM2 verify forward ms (median of {d}) by streams and rows a stream, {d}-token histories, {s} rows\n", .{ reps, history, if (part > 0) "real" else "synthetic" }); try w.print("# streams rows_each total_rows ms\n", .{}); var segments: [st.batch_rows]q.forward.Segment = undefined; var tokens: [st.batch_rows]u32 = undefined; @@ -48,7 +64,11 @@ pub fn main(init: std.process.Init) !void { for (counts) |streams| { var each: usize = 1; while (each <= st.window_rows and streams * each <= st.batch_rows) : (each += 1) { - for (segments[0..streams], caches[0..streams]) |*s, *k| s.* = .{ .cache = k, .rows = each }; + for (segments[0..streams], caches[0..streams], 0..) |*s, *k, i| { + s.* = .{ .cache = k, .rows = each }; + // real rows: the tokens after the stream's history, as a verify of drafts that all land + if (part > 0) @memcpy(tokens[i * each ..][0..each], real.items[i * part + history ..][0..each]); + } var times: [reps + 2]f64 = undefined; for (×) |*t| { const t0 = now(); From 0b76aeed3ab0301488d15390d2fd96a6fb1c03f1 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:40:54 +0100 Subject: [PATCH 25/46] lfm2: the DSpark drafter loads one matrix at a time and keeps a lane in one buffer Memory of the drafter, no change to its drafts: - The loader keeps the matrices in the safetensors file (Checkpoint.addFileSplit) and reads each into one reused stage, quantizes it and drops it. Load no longer holds the 655 MB bf16 file beside the 8-bit codes, only the largest matrix (W1 and W2). The codes are the same bytes. - A lane is one Metal buffer of laneBytes: its layers' keys and values, then its staged rows. Opening a stream allocated eleven buffers. Lane.kv gives a layer's offset; the rows and their layout are unchanged. - The absorb norms its context rows in place, so the rotated query heads go to the freed buffer and x, which only a hold reads, takes 32 rows instead of 2,048. - The waited passes end in one helper that commits, waits and reports a failed command buffer. - At load the drafter logs its weights, scratch and a stream's bytes. tf-lfm2-dspark counts host allocations: twenty rounds of hold, chances, held drafts and a staged absorb through the drafter's interface, after four warm ones, must take no allocation or resize of the drafter's allocator and one unwaited hold pass each. --- zig/src/families/lfm2_moe/dspark.zig | 62 +++++++++----------- zig/src/families/lfm2_moe/dspark_block.zig | 41 ++++++------- zig/src/families/lfm2_moe/dspark_weights.zig | 21 +++++-- zig/tests/lfm2_dspark.zig | 42 +++++++++++++ 4 files changed, 104 insertions(+), 62 deletions(-) diff --git a/zig/src/families/lfm2_moe/dspark.zig b/zig/src/families/lfm2_moe/dspark.zig index eb3fa4d90f..76d8369521 100644 --- a/zig/src/families/lfm2_moe/dspark.zig +++ b/zig/src/families/lfm2_moe/dspark.zig @@ -34,17 +34,21 @@ pub const Options = struct { capacity: usize, bits: u8 = 8, verify: Verify = def /// One stream's context keys and values a layer ([kv head][cap][head dim], as the target's attention reads them). pub const Lane = struct { - keys: [8]mtl.Buffer, - values: [8]mtl.Buffer, + mem: mtl.Buffer, // each layer's keys then values, then the staged rows (shared storage) cap: usize, // positions each holds: the context and one block past it + stage_at: usize, // the staged rows' byte offset in `mem` end: u64 = 0, // context rows absorbed: positions [0, end) - stage: mtl.Buffer, // tapped rows [end, end + staged), not absorbed yet (shared storage) - staged: u32 = 0, + staged: u32 = 0, // tapped rows [end, end + staged), not absorbed yet held: [16]u32 = @splat(0), held_len: u32 = 0, conf: [16]f32 = @splat(0), // each held draft's chance given the drafts before it landed (the confidence head) conf_len: u32 = 0, // the drafts `conf` covers: the last settled hold's depth slot: ?u32 = null, // its block's place in the last signalled hold, whose picks still hold its drafts + + /// Layer `layer`'s keys, or its values when `values`, in `mem`. + pub fn kv(l: *const Lane, layer: usize, values: bool) block.Ref { + return .{ .buf = l.mem, .off = (2 * layer + @intFromBool(values)) * l.cap * kw * 2 }; + } }; /// The passes' working buffers, one pass at a time on the target's queue. @@ -66,8 +70,8 @@ pub const Scratch = struct { act: mtl.Buffer, logits: mtl.Buffer, sums: mtl.Buffer, - fcx: mtl.Buffer, - hctx: mtl.Buffer, + fcx: mtl.Buffer, // an absorb pass's context rows: fc's output, normed in place + dump: mtl.Buffer, // the query heads lfm2_qk rotates from context rows, which no one reads best: mtl.Buffer, // the Markov scan's best score and token of each threadgroup, by stream at: mtl.Buffer, picks: mtl.Buffer, @@ -139,7 +143,7 @@ pub const Drafter = struct { d.s = .{ .ids = try d.buffer(r * 4), .h = try d.buffer(r * c.hidden * 2), - .x = try d.buffer(wide * c.hidden * 2), + .x = try d.buffer(r * c.hidden * 2), .r = try d.buffer(r * c.hidden * 2), .y = try d.buffer(r * c.hidden * 2), .q = try d.buffer(r * c.hidden * 2), @@ -155,7 +159,7 @@ pub const Drafter = struct { .logits = try d.buffer(r * c.vocab * 2), .sums = try d.buffer(sums * 4), .fcx = try d.buffer(absorb_rows * c.hidden * 2), - .hctx = try d.buffer(absorb_rows * c.hidden * 2), + .dump = try d.buffer(absorb_rows * c.hidden * 2), .best = try d.buffer(r * block.scan_groups * 4), .at = try d.buffer(r * block.scan_groups * 4), .picks = try d.buffer(d.config.block * r * 8), @@ -236,8 +240,9 @@ pub const Drafter = struct { return true; } + /// A stream's lane at `capacity` positions: its context keys and values a layer, and its staged rows. pub fn laneBytes(d: *const Drafter, capacity: usize) usize { - return 2 * d.config.layers * (capacity + d.config.block) * kw * 2 + stage_rows * d.tapBytes(); + return @intCast(d.keptBytes(capacity + d.config.block) + stage_rows * d.tapBytes()); } /// A tapped row's bytes: every tap's hidden row, bf16. @@ -248,7 +253,7 @@ pub const Drafter = struct { /// The lane's staged rows as an absorb piece (null: none staged). pub fn stagedPiece(l: *Lane) ?block.Piece { if (l.staged == 0) return null; - return .{ .lane = l, .src = .{ .buf = l.stage }, .rows = l.staged, .p0 = l.end }; + return .{ .lane = l, .src = .{ .buf = l.mem, .off = l.stage_at }, .rows = l.staged, .p0 = l.end }; } /// The lane's staged rows absorbed now, in a command buffer of their own. @@ -272,33 +277,17 @@ pub const Drafter = struct { try d.lanes.put(d.gpa, s, l); } + /// A lane in one buffer of laneBytes(capacity). pub fn newLane(d: *Drafter, capacity: usize) !*Lane { const l = try d.gpa.create(Lane); errdefer d.gpa.destroy(l); - l.* = .{ .keys = undefined, .values = undefined, .stage = try d.target.device.buffer(stage_rows * d.tapBytes(), shared), .cap = capacity + d.config.block }; - errdefer l.stage.deinit(); - var made: usize = 0; - errdefer for (0..made) |i| { - l.keys[i].deinit(); - l.values[i].deinit(); - }; - for (0..d.config.layers) |i| { - l.keys[i] = try d.target.device.buffer(l.cap * kw * 2, shared); - l.values[i] = d.target.device.buffer(l.cap * kw * 2, shared) catch |e| { - l.keys[i].deinit(); - return e; - }; - made += 1; - } + const at: usize = @intCast(d.keptBytes(capacity + d.config.block)); + l.* = .{ .mem = try d.target.device.buffer(d.laneBytes(capacity), shared), .cap = capacity + d.config.block, .stage_at = at }; return l; } pub fn dropLane(d: *Drafter, l: *Lane) void { - l.stage.deinit(); - for (0..d.config.layers) |i| { - l.keys[i].deinit(); - l.values[i].deinit(); - } + l.mem.deinit(); d.gpa.destroy(l); } @@ -348,8 +337,8 @@ pub const Drafter = struct { fn copyLane(d: *const Drafter, l: *const Lane, at: usize, bytes: []u8, out: bool) void { const n = at * c.head_dim * 2; var off: usize = 0; - for (0..d.config.layers) |i| for ([_]mtl.Buffer{ l.keys[i], l.values[i] }) |buf| for (0..c.kv_heads) |h| { - const live = (buf.contents() + h * l.cap * c.head_dim * 2)[0..n]; + for (0..d.config.layers) |i| for ([_]block.Ref{ l.kv(i, false), l.kv(i, true) }) |r| for (0..c.kv_heads) |h| { + const live = (r.buf.contents() + r.off + h * l.cap * c.head_dim * 2)[0..n]; if (out) @memcpy(bytes[off..][0..n], live) else @memcpy(live, bytes[off..][0..n]); off += n; }; @@ -380,7 +369,7 @@ pub const Drafter = struct { l.staged = @min(l.staged, keep); if (it.start >= l.end and l.staged + f.rows <= stage_rows) { const from = (src.contents() + f.offset)[0 .. f.rows * width]; - @memcpy((l.stage.contents() + l.staged * width)[0 .. f.rows * width], from); + @memcpy((l.mem.contents() + l.stage_at + l.staged * width)[0 .. f.rows * width], from); l.staged += f.rows; continue; } @@ -517,6 +506,12 @@ pub const Drafter = struct { return heldFn(ptr, &.{s}, &.{out}); } + fn mibOf(bufs: []const mtl.Buffer) f64 { + var n: usize = 0; + for (bufs) |b| n += b.length(); + return @as(f64, @floatFromInt(n)) / (1024.0 * 1024.0); + } + /// Holds on empty contexts: the pipelines compiled, then holds timed by blocks and depth (best of three). fn warm(d: *Drafter) !void { var made: [hold_grid[hold_grid.len - 1]]*Lane = undefined; @@ -534,6 +529,7 @@ pub const Drafter = struct { d.hold_ms = best[1]; d.step_ms = best[1] / @as(f64, @floatFromInt(d.config.block)); std.log.info("lfm2 dspark: a 1-draft hold takes {d:.2} ms, a {d}-draft hold {d:.2} ms", .{ best[0], d.config.block, best[1] }); + std.log.info("lfm2 dspark: weights {d:.1} MiB, scratch {d:.1} MiB, a stream {d} KiB a position and {d} KiB staged", .{ mibOf(d.w.buffers()), mibOf(d.owned.items), d.keptBytes(1) / 1024, stage_rows * d.tapBytes() / 1024 }); var blocks: [hold_grid[hold_grid.len - 1]]block.Block = undefined; for (made, &blocks) |lane, *b| b.* = .{ .lane = lane, .pending = 1, .anchor = 0, .depth = 1 }; d.holds_len = 0; diff --git a/zig/src/families/lfm2_moe/dspark_block.zig b/zig/src/families/lfm2_moe/dspark_block.zig index 0b712bf2dd..64e9dd38ce 100644 --- a/zig/src/families/lfm2_moe/dspark_block.zig +++ b/zig/src/families/lfm2_moe/dspark_block.zig @@ -73,8 +73,8 @@ const Enc = struct { const args = QkArgs{ .len = @intCast(len), .cap = @intCast(lane.cap), .eps = c.eps, .log2_theta = std.math.log2(x.d.config.theta) }; x.e.setBytes(std.mem.asBytes(&args), 5); x.e.setBuffer(queries.buf, queries.off, 6); - x.e.setBuffer(lane.keys[layer], 0, 7); - x.e.setBuffer(lane.values[layer], 0, 8); + x.e.setBuffer(lane.mem, lane.kv(layer, false).off, 7); + x.e.setBuffer(lane.mem, lane.kv(layer, true).off, 8); x.e.dispatchGroups(mtl.Size.of(c.query_heads + c.kv_heads, rows, 1), mtl.Size.of(32, 1, 1)); } @@ -85,8 +85,8 @@ const Enc = struct { const used = (keys + st.chunk_keys - 1) / st.chunk_keys * st.subs; x.pipe("lfm2_attn_block"); x.e.setBuffer(s.queries, row * c.hidden * 2, 0); - x.e.setBuffer(lane.keys[layer], 0, 1); - x.e.setBuffer(lane.values[layer], 0, 2); + x.e.setBuffer(lane.mem, lane.kv(layer, false).off, 1); + x.e.setBuffer(lane.mem, lane.kv(layer, true).off, 2); x.e.setBytes(std.mem.asBytes(&args), 3); x.e.setBuffer(s.pm, 0, 4); x.e.setBuffer(s.pl, 0, 5); @@ -103,18 +103,6 @@ const Enc = struct { } }; -fn failed(msg: [*:0]const u8) anyerror { - std.log.err("lfm2 dspark: command buffer failed: {s}", .{msg}); - return error.GpuFailed; -} - -fn finish(cb: mtl.CommandBuffer, e: mtl.ComputeEncoder) !void { - e.end(); - cb.commit(); - cb.wait(); - if (cb.failure()) |msg| return failed(msg); -} - /// One stream's tapped rows in an absorb pass: `rows` rows at `src`, the context from position `p0`. pub const Piece = struct { lane: *Lane, src: Ref, rows: usize, p0: u64 }; @@ -128,16 +116,15 @@ pub fn absorbPass(d: *Drafter, e: mtl.ComputeEncoder, pieces: []const Piece) !vo rows += pc.rows; } if (rows == 0 or rows > ds.absorb_rows) return error.DraftContextGap; - x.norm(s.fcx, s.fcx, d.w.hidden_norm, false, s.hctx, rows); + x.norm(s.fcx, s.fcx, d.w.hidden_norm, false, s.fcx, rows); // in place: a thread reads its values first for (0..d.config.layers) |i| { const ly = &d.w.layers[i]; - try x.project(ly.k, .{ .buf = s.hctx }, .{ .buf = s.k }, rows); - try x.project(ly.v, .{ .buf = s.hctx }, .{ .buf = s.v }, rows); + try x.project(ly.k, .{ .buf = s.fcx }, .{ .buf = s.k }, rows); + try x.project(ly.v, .{ .buf = s.fcx }, .{ .buf = s.v }, rows); var at: usize = 0; for (pieces) |pc| { - // the kernel rotates query heads too: they read the context rows and land in x, which no one reads - const q = Ref{ .buf = s.hctx, .off = at * c.hidden * 2 }; - x.qk(q, .{ .buf = s.k, .off = at * kw * 2 }, .{ .buf = s.v, .off = at * kw * 2 }, i, pc.lane, pc.p0, .{ .buf = s.x }, pc.rows); + const q = Ref{ .buf = s.fcx, .off = at * c.hidden * 2 }; + x.qk(q, .{ .buf = s.k, .off = at * kw * 2 }, .{ .buf = s.v, .off = at * kw * 2 }, i, pc.lane, pc.p0, .{ .buf = s.dump }, pc.rows); at += pc.rows; } } @@ -150,7 +137,13 @@ pub fn begin(d: *Drafter) struct { cb: mtl.CommandBuffer, e: mtl.ComputeEncoder } pub fn end(cb: mtl.CommandBuffer, e: mtl.ComputeEncoder) !void { - return finish(cb, e); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.log.err("lfm2 dspark: command buffer failed: {s}", .{msg}); + return error.GpuFailed; + } } /// A GPU event a hold signals instead of being waited on, and the value it sets. @@ -242,7 +235,7 @@ pub fn hold(d: *Drafter, blocks: []const Block, signal: ?Signal) !?mtl.CommandBu cb.commit(); return kept; } - try finish(cb, e); + try end(cb, e); const picks = s.picks.slice(Pick, d.config.block * ds.pass_rows); for (blocks, 0..) |b, bi| { for (b.lane.held[0..steps], 0..) |*h, i| { diff --git a/zig/src/families/lfm2_moe/dspark_weights.zig b/zig/src/families/lfm2_moe/dspark_weights.zig index bfd6d078e1..f0e081cb25 100644 --- a/zig/src/families/lfm2_moe/dspark_weights.zig +++ b/zig/src/families/lfm2_moe/dspark_weights.zig @@ -55,6 +55,7 @@ const Loader = struct { gpa: std.mem.Allocator, device: mtl.Device, check: *const ckpt.Checkpoint, + stage: *ckpt.Stage, // a matrix read from the file, one at a time, into one reused buffer owned: *std.ArrayList(mtl.Buffer), bits: u8, @@ -65,8 +66,11 @@ const Loader = struct { return b; } + /// A vector as loaded, or a matrix read into the stage (valid until the next matrix), of `shape` in bf16. fn tensor(l: Loader, name: []const u8, shape: []const usize) !ckpt.Tensor { - const t = try l.check.get(name); + l.stage.reset(); + if (shape.len == 2) try l.stage.reserve(l.check, &.{name}); + const t = if (shape.len == 2) try l.stage.read(l.check, name) else try l.check.get(name); if (t.dtype != .bf16 or t.rank != shape.len or !std.mem.eql(usize, t.shape[0..t.rank], shape)) { std.log.err("{s}: unexpected DSpark tensor dtype or shape", .{name}); return error.BadDraftWeight; @@ -108,18 +112,25 @@ fn part(buf: []u8, layer: usize, name: []const u8) ![]const u8 { return std.fmt.bufPrint(buf, "layers.{d}.{s}.weight", .{ layer, name }); } -/// The drafter's checkpoint read and quantized at `bits` (4 or 8); the checkpoint's own buffers are released after. +/// Matrices stay in the file until the loader reads them; vectors load with the header. +fn matrices(_: []const u8, rank: usize) bool { + return rank == 2; +} + +/// The drafter's checkpoint read and quantized at `bits` (4 or 8), a matrix at a time; its buffers are released after. pub fn load(gpa: std.mem.Allocator, device: mtl.Device, io: std.Io, dir: []const u8, cfg: *const dc.Config, bits: u8) !Weights { if (bits != 4 and bits != 8) return error.UnsupportedDraftBits; var check = ckpt.Checkpoint.init(gpa); defer check.deinit(); + var stage: ckpt.Stage = .{ .device = device }; + defer stage.deinit(); const files = try shards.shardFiles(gpa, io, dir); defer shards.freeShardFiles(gpa, files); - for (files) |path| try check.addFile(device, path, ""); + for (files) |path| try check.addFileSplit(device, path, "", null, matrices); var w: Weights = undefined; w.owned = .empty; errdefer w.deinit(gpa); - const l: Loader = .{ .gpa = gpa, .device = device, .check = &check, .owned = &w.owned, .bits = bits }; + const l: Loader = .{ .gpa = gpa, .device = device, .check = &check, .stage = &stage, .owned = &w.owned, .bits = bits }; const qw = c.query_heads * c.head_dim; const kw = c.kv_heads * c.head_dim; w.fc = try l.matrix("fc.weight", c.hidden, cfg.tap_count * c.hidden, false); @@ -145,7 +156,7 @@ pub fn load(gpa: std.mem.Allocator, device: mtl.Device, io: std.Io, dir: []const w.w2 = try l.matrix("markov_head.markov_w2.weight", c.vocab, cfg.rank, false); w.confidence = null; w.confidence_bias = 0; - if (check.tensors.contains("confidence_head.proj.weight")) { + if (check.has("confidence_head.proj.weight")) { const width = c.hidden + cfg.rank; // the final norm's row, then W1 of the token before the draft const proj = (try l.tensor("confidence_head.proj.weight", &.{ 1, width })).host(u16); w.confidence = try l.buffer(width * 2); diff --git a/zig/tests/lfm2_dspark.zig b/zig/tests/lfm2_dspark.zig index 5ee8140de2..5a9885b98e 100644 --- a/zig/tests/lfm2_dspark.zig +++ b/zig/tests/lfm2_dspark.zig @@ -125,6 +125,47 @@ fn absorb(d: *q.dspark.Drafter, piece: q.dspark_block.Piece) !void { try q.dspark_block.end(run.cb, run.e); } +/// Served rounds through the drafter's interface: after four warm rounds, no host allocation and one hold pass a round. +fn steady(gpa: std.mem.Allocator, d: *q.dspark.Drafter, taps: mtl.Buffer, ids: []const u32, ctx: usize) !void { + const lanes = tf.lanes; + var s = try lanes.Stream.init(gpa, .{ .id = "steady", .prompt = ids[0 .. ctx + 1], .max_new = 160 }); + defer s.deinit(gpa); + const dr = d.drafter(); + try dr.open(&s); + defer dr.release(&s); + const width: u32 = @intCast(d.tapBytes()); + const feats = struct { + fn at(buf: mtl.Buffer, row: usize, rows: usize, w: u32) lanes.backend.Features { + return .{ .buffer = @intFromPtr(buf.id), .offset = row * w, .rows = @intCast(rows), .row_bytes = w, .space = .device, .dtype = .bf16 }; + } + }; + try dr.absorb(&.{.{ .stream = &s, .features = feats.at(taps, 0, ctx, width), .start = 0, .follow = &.{} }}); + var counted = std.testing.FailingAllocator.init(gpa, .{}); + var anchor: u64 = ctx; + const rounds = 24; + const warm = 4; + var holds: u64 = 0; + for (0..rounds) |r| { + if (r == warm) { + d.gpa = counted.allocator(); + holds = d.signalled; + } + try dr.hold(&.{.{ .stream = &s, .pending = ids[@intCast(anchor % ctx)], .position = anchor + 1, .depth = 3 }}); + var chances: [3]f64 = undefined; + _ = try dr.chances(&s, &chances); + var drafts: [3]u32 = undefined; + try dr.held(&.{&s}, &.{&drafts}); + const rows = r % 3 + 1; + const at: usize = @intCast(anchor % (ctx - 3)); + try dr.absorb(&.{.{ .stream = &s, .features = feats.at(taps, at, rows, width), .start = anchor, .follow = &.{}, .prompt = false }}); + anchor += rows; + } + d.gpa = gpa; + const n = rounds - warm; + std.debug.print("steady rounds: {d} rounds of hold, chances, held and a verified absorb: {d} host allocations, {d} resizes, {d} unwaited holds\n", .{ n, counted.allocations, counted.resize_index, d.signalled - holds }); + if (counted.allocations != 0 or counted.resize_index != 0 or d.signalled - holds != n) return error.SteadyRoundsAllocate; +} + pub fn main(init: std.process.Init) !void { const gpa = init.gpa; const args = try init.minimal.args.toSlice(init.arena.allocator()); @@ -218,6 +259,7 @@ pub fn main(init: std.process.Init) !void { } std.debug.print("hold of {d} streams at depth {d}: {d:.2} ms, GPU {d:.2} ms (best of 5)\n", .{ streams, depth, best, gpu_best }); }; + try steady(gpa, d, taps.buf, ids, ctx); _ = try q.dspark_block.hold(d, &.{.{ .lane = lane, .pending = ids[ctx], .anchor = ctx, .depth = block }}, null); // the host reference, fp32 throughout var r: Ref = .{ .gpa = gpa, .check = ck.Checkpoint.init(gpa), .cfg = cfg }; From 7dd1a2ef800c35837eac2c57fb6006aa5b4a9ca5 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 10:06:31 +0100 Subject: [PATCH 26/46] docs: the LFM2 recipe's DSpark verify schedule and its flags The recipe names the confidence head, the three --draft-verify modes (scheduled by default), how the lane core prices a hold pass and a copy, the lone-stream guard, and --verify-costs with tf-lfm2-verify-costs. --- docs/recipes/lfm2.5-24b-a2b.md | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/docs/recipes/lfm2.5-24b-a2b.md b/docs/recipes/lfm2.5-24b-a2b.md index f4272b59a5..64e26c258b 100644 --- a/docs/recipes/lfm2.5-24b-a2b.md +++ b/docs/recipes/lfm2.5-24b-a2b.md @@ -82,3 +82,29 @@ drafted rounds measure faster than plain ones at that stream count. against an fp32 host reference of the released checkpoint (the same taps and the target's embedding and head), prints each draft against the reference's pick with the margin between the reference's two best tokens, and times the drafter's passes; a near tie can pick the other token. + +### Verify schedule + +The drafter also carries a confidence head: a linear map of a draft's final block row and `W1[previous token]` +(2,304 inputs) to one logit, through a sigmoid (`lfm2_dspark_confidence`). It does not change the drafts. +`--draft-verify` picks how they are verified: + +- `scheduled` (the default): a chain's chance at draft j is the product of its first j confidences. Each round the + lane core's window allocator takes the likeliest drafts over all streams while the expected tokens a ms rise, so + each stream gets its own verify length, and the next hold drafts as deep as the stream's grant. A hold still + running on the GPU is priced by its lane's last confidences, so the host does not wait for it. +- `depth`: the lane core's depth rule, from each stream's measured acceptance by depth. +- `fixed`: every round verifies the whole block, as LiquidAI's published benchmarks did; a shared round still fits + 32 rows. + +With `scheduled` a stream with a copy ahead verifies the copy and asks no hold, and a drafting stream prices its copy +windows by its own copies. The lane core also weighs the hold pass itself: the drafter times holds of 1 to 16 +blocks at load and corrects those prices toward each settled hold's GPU time. A lone stream holds drafts only while +its drafted rounds measure faster than plain ones (`zig/src/core/lanes/lone_guard.zig`); shared rounds keep the +shared-round guard. + +`--verify-costs FILE` replaces the costs timed at load with a table from `tf-lfm2-verify-costs MODEL FILE [HISTORY +[IDS...]]` (lines `streams rows_each total_rows ms`), measured on the serving Mac. With token id files each stream +reads its history and verify rows from its own part of the ids, so its rows route its own text's experts. Both flags +are refused unless an `lfm2_moe` checkpoint is served with `--drafter`. `tf-lfm2-dspark` also prints the confidence +head against its fp32 reference and times shared holds of 2 to 16 streams. From 9d35bdc5bc9a454031b1a2058e7e96ab1ed9a281 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 10:06:40 +0100 Subject: [PATCH 27/46] docs: the LFM2 recipe's DSpark loader and lane memory The recipe says how the drafter loads and keeps its lanes, what it logs at load and that tf-lfm2-dspark counts host allocations over served rounds. --- docs/recipes/lfm2.5-24b-a2b.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/docs/recipes/lfm2.5-24b-a2b.md b/docs/recipes/lfm2.5-24b-a2b.md index 64e26c258b..48b2666927 100644 --- a/docs/recipes/lfm2.5-24b-a2b.md +++ b/docs/recipes/lfm2.5-24b-a2b.md @@ -108,3 +108,8 @@ shared-round guard. reads its history and verify rows from its own part of the ids, so its rows route its own text's experts. Both flags are refused unless an `lfm2_moe` checkpoint is served with `--drafter`. `tf-lfm2-dspark` also prints the confidence head against its fp32 reference and times shared holds of 2 to 16 streams. + +The loader reads the checkpoint one matrix at a time and keeps only the 8-bit codes; a stream's drafter lane is one +Metal buffer (its layers' keys and values, then its staged rows), and its bytes join the streams the load limit +admits. At load the drafter logs its weights, scratch and a stream's bytes. `tf-lfm2-dspark` also counts host +allocations over twenty served drafter rounds, which must be none. From 4fa02928a7484ccf7bc1a529936618e9489eb24e Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:14:47 +0100 Subject: [PATCH 28/46] lanes: the shared guard is off when Config.timed is false, so a schedule does not depend on the clock The guard planned, observed and chose with the wall clock even when timed was false. Round planning, joining a shared round and observation now skip it in that mode. A test runs the same three streams on a fast and a slow fake clock and checks equal holds and replies. --- zig/src/core/lanes/drafted_guard_test.zig | 53 +++++++++++++++++++++++ zig/src/core/lanes/engine.zig | 6 +-- zig/src/core/lanes/open.zig | 2 +- 3 files changed, 57 insertions(+), 4 deletions(-) diff --git a/zig/src/core/lanes/drafted_guard_test.zig b/zig/src/core/lanes/drafted_guard_test.zig index ca6d6bc568..3715bec357 100644 --- a/zig/src/core/lanes/drafted_guard_test.zig +++ b/zig/src/core/lanes/drafted_guard_test.zig @@ -100,3 +100,56 @@ test "the guard changes only which shared rounds draft: replies equal one-token } try std.testing.expect(on.holds < off.holds); // its plain stretch holds no drafts } + +/// A clock that reads the same ms for every span: two of them with different ms stand for a fast and a slow machine. +const FlatClock = struct { + ms: f64, + + fn clock(c: *FlatClock) @import("backend.zig").Clock { + return .{ .ptr = c, .vtable = &.{ .start = start, .elapsed_ms = elapsed } }; + } + + fn start(_: *anyopaque) void {} + + fn elapsed(ptr: *anyopaque, _: @import("backend.zig").Mark) f64 { + const c: *FlatClock = @ptrCast(@alignCast(ptr)); + return c.ms; + } +}; + +/// Three streams, the guard asked for or not, timed or not, on a clock reading `ms` for every span. +fn runClock(timed: bool, guard: bool, ms: f64) !Run { + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var head: dt.FakeDrafter = .{ .facts_ = .{ .depth = 4, .step_ms = 0.5, .batched = true, .shared_guard = guard } }; + defer head.deinit(); + var wrapped = try Drafted.init(gpa, target.backend(), head.drafter()); + defer wrapped.deinit(); + const m: Model = .{ .exact_width = 16, .gpu_tokens = true, .hidden_rows = true, .batch_rows = 32, .max_streams = 4 }; + var cfg = try Config.init(gpa, wrapped.facts(m), 16, 15); + defer cfg.deinit(gpa); + cfg.timed = timed; + var clock: FlatClock = .{ .ms = ms }; + var engine = Engine.init(gpa, &cfg, wrapped.backend(), clock.clock()); + defer engine.deinit(); + var streams: [3]sm.Stream = undefined; + for (&streams, 0..) |*s, i| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = dt.long[0 .. 11 + i], .max_new = 48, .eos = &.{96}, .drafts = true }); + defer for (&streams) |*s| s.deinit(gpa); + for (&streams) |*s| try engine.addStream(s); + while (engine.activeCount() > 0) try engine.step(); + var out: Run = .{ .holds = head.holds }; + for (&streams, 0..) |*s, i| { + out.lens[i] = s.emitted().len; + @memcpy(out.tokens[i][0..out.lens[i]], s.emitted()); + } + return out; +} + +test "with timed off the shared guard keeps no plain stretch: the same schedule on a fast and a slow clock" { + const fast = try runClock(false, true, 0.5); + const slow = try runClock(false, true, 400.0); + try std.testing.expectEqual(fast.holds, slow.holds); + for (0..3) |i| try std.testing.expectEqualSlices(u32, fast.tokens[i][0..fast.lens[i]], slow.tokens[i][0..slow.lens[i]]); + const unasked = try runClock(false, false, 0.5); + try std.testing.expectEqual(unasked.holds, fast.holds); // as if the guard were not asked for +} diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index 5349b38904..664447ceaf 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -261,7 +261,7 @@ pub const Engine = struct { e.clock.start(); const a = e.arena.allocator(); const plans = try a.alloc(Plan, entries.len); - e.plain_round = e.cfg.shared_guard and !try e.guard.drafts(entries.len); + e.plain_round = e.cfg.shared_guard and e.cfg.timed and !try e.guard.drafts(entries.len); defer e.plain_round = false; for (entries, plans) |s, *p| p.* = try win.plan(e, s, null); try win.allocate(e, plans); @@ -293,7 +293,7 @@ pub const Engine = struct { const budgets = try a.alloc(i64, heads.items.len); const sampled = win.anySampled(plans); // the guard's plain stretches: no stream drafts, so the round costs what it would without a head - const guarded = e.cfg.shared_guard and !try e.guard.drafts(plans.len); + const guarded = e.cfg.shared_guard and e.cfg.timed and !try e.guard.drafts(plans.len); if (guarded) @memset(budgets, 0) else try e.rule.budgets(whos, win.probe(e), sampled, budgets); const ids = try a.alloc(ev.Value, heads.items.len); for (ids, heads.items) |*id, i| id.* = str(plans[i].stream.id); @@ -323,7 +323,7 @@ pub const Engine = struct { try trail.event(e, &.{ f("ev", str("overhead")), f("streams", int(plans.len)), f("rows", int(total)), f("ms", ev.bits(ms)) }); try e.rule.observeOverhead(@intCast(plans.len), total, ms, sampled, e.rule.headMs(budgets)); } - if (e.cfg.shared_guard) try e.guard.observe(plans.len, shared_guard.committed(outcomes), e.clock.elapsedMs(.overhead)); + if (e.cfg.shared_guard and e.cfg.timed) try e.guard.observe(plans.len, shared_guard.committed(outcomes), e.clock.elapsedMs(.overhead)); if (e.backend.round_observer != null) e.backend.observed(paths); } diff --git a/zig/src/core/lanes/open.zig b/zig/src/core/lanes/open.zig index 733d3bb2f8..153da72e53 100644 --- a/zig/src/core/lanes/open.zig +++ b/zig/src/core/lanes/open.zig @@ -137,7 +137,7 @@ pub fn firstDepth(e: *Engine, s: *Stream, reach: Reach) !u32 { var who = win.who(s); // a stream joining shared rounds the guard runs plain holds no drafts, as the guard measured it faster const others = e.activeCount(); - if (others > 0 and e.cfg.family_streams and e.cfg.shared_guard and !try e.guard.drafts(others + 1)) return 0; + if (others > 0 and e.cfg.family_streams and e.cfg.shared_guard and e.cfg.timed and !try e.guard.drafts(others + 1)) return 0; if (reach == .head and e.cfg.node_probabilities) return @intCast(try e.rule.headDepth(&who, null, win.probe(e))); return @intCast(try e.rule.depth(who)); } From db5130ba0d4531062bea055f9436345b7f02cb45 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 18:44:14 +0100 Subject: [PATCH 29/46] lfm2 dspark: a held pick that is nonfinite or past the vocabulary is refused, and hidden_act must be silu The unwaited hold read its picks with nonfinite ignored and a bad id as 0; the GPU copy into a verify window did the same. Both now refuse as the synchronous path does (DraftNonfinite): readPicks errors, and lfm2_draft_ids sets a flag the verify reads. Planted NaN, Inf and invalid-id tests cover the host and the kernel. The config refuses a hidden_act other than silu by name. --- zig/kernels/metal/lfm2/lfm2_dspark.metal | 10 +++-- zig/src/families/lfm2_moe/dspark.zig | 2 +- zig/src/families/lfm2_moe/dspark_block.zig | 21 +++++++--- zig/src/families/lfm2_moe/dspark_config.zig | 4 ++ zig/src/families/lfm2_moe/dspark_test.zig | 44 +++++++++++++++++++++ zig/src/families/lfm2_moe/forward.zig | 3 ++ zig/src/families/lfm2_moe/state.zig | 2 + zig/tests/contract/metal_lfm2_dspark.zig | 14 ++++++- 8 files changed, 88 insertions(+), 12 deletions(-) diff --git a/zig/kernels/metal/lfm2/lfm2_dspark.metal b/zig/kernels/metal/lfm2/lfm2_dspark.metal index 050fd5e8cc..34fbc3d704 100644 --- a/zig/kernels/metal/lfm2/lfm2_dspark.metal +++ b/zig/kernels/metal/lfm2/lfm2_dspark.metal @@ -98,10 +98,14 @@ constant constexpr uint MARKOV_SPAN = 512; } } -// A verify window's held drafts from the drafter's picks: picks[i * stride] into ids[i], a token past the vocabulary 0. +// A verify window's held drafts from the drafter's picks: picks[i * stride] into ids[i]. A pick flagged nonfinite +// (the word after it) or past the vocabulary becomes token 0 and sets bad, which the host reads after the verify. [[kernel]] void lfm2_draft_ids(const device uint* picks [[buffer(0)]], constant uint4& p [[buffer(1)]], - device uint* ids [[buffer(2)]], uint i [[thread_position_in_grid]]) { + device uint* ids [[buffer(2)]], device atomic_uint* bad [[buffer(3)]], + uint i [[thread_position_in_grid]]) { if (i >= p.y) return; const uint token = picks[size_t(i) * p.x]; - ids[i] = token < p.z ? token : 0u; + const bool ok = token < p.z && picks[size_t(i) * p.x + 1] == 0u; + ids[i] = ok ? token : 0u; + if (!ok) atomic_store_explicit(bad, 1u, memory_order_relaxed); } diff --git a/zig/src/families/lfm2_moe/dspark.zig b/zig/src/families/lfm2_moe/dspark.zig index 850c1ce34c..3effd0bf0c 100644 --- a/zig/src/families/lfm2_moe/dspark.zig +++ b/zig/src/families/lfm2_moe/dspark.zig @@ -408,7 +408,7 @@ pub const Drafter = struct { std.log.err("lfm2 dspark: a hold failed on the GPU: {s}", .{msg}); return error.DraftGpuFailed; } - block.readPicks(d, d.pass[0..d.pass_len], d.depths[0..d.pass_len]); + try block.readPicks(d, d.pass[0..d.pass_len], d.depths[0..d.pass_len]); } /// The signalled hold settled and its lanes' picks forgotten (a lane may be dropped, or a hold may overwrite). diff --git a/zig/src/families/lfm2_moe/dspark_block.zig b/zig/src/families/lfm2_moe/dspark_block.zig index d3f6cbabb8..fab4f7886e 100644 --- a/zig/src/families/lfm2_moe/dspark_block.zig +++ b/zig/src/families/lfm2_moe/dspark_block.zig @@ -253,15 +253,24 @@ pub fn hold(d: *Drafter, blocks: []const Block, signal: ?Signal) !?mtl.CommandBu return null; } -/// A finished signalled hold's picks into the lanes' held chains (a pick past the vocabulary reads as token 0). -pub fn readPicks(d: *Drafter, lanes: []const *Lane, depths: []const u32) void { +/// A finished signalled hold's picks into the lanes' held chains; a nonfinite or invalid pick is refused as the +/// synchronous path refuses it, and no lane keeps a chain then. +pub fn readPicks(d: *Drafter, lanes: []const *Lane, depths: []const u32) !void { const picks = d.s.picks.slice(Pick, d.config.block * ds.pass_rows); - for (lanes, depths, 0..) |l, depth, bi| { - for (l.held[0..depth], 0..) |*h, i| { + var chains: [ds.pass_rows][]u32 = undefined; + for (lanes, chains[0..lanes.len]) |l, *ch| ch.* = &l.held; + try copyPicks(picks, chains[0..lanes.len], depths); + for (lanes, depths) |l, depth| l.held_len = depth; +} + +/// Step i of block bi is picks[i * ds.pass_rows + bi]: each block's first `depth` picks into its chain. +pub fn copyPicks(picks: []const Pick, chains: []const []u32, depths: []const u32) !void { + for (chains, depths, 0..) |chain, depth, bi| { + for (chain[0..depth], 0..) |*h, i| { const p = picks[i * ds.pass_rows + bi]; - h.* = if (p.token < c.vocab) p.token else 0; + if (p.nonfinite != 0 or p.token >= c.vocab) return error.DraftNonfinite; + h.* = p.token; } - l.held_len = depth; } } diff --git a/zig/src/families/lfm2_moe/dspark_config.zig b/zig/src/families/lfm2_moe/dspark_config.zig index 7b3b628a02..79e518aa48 100644 --- a/zig/src/families/lfm2_moe/dspark_config.zig +++ b/zig/src/families/lfm2_moe/dspark_config.zig @@ -55,6 +55,9 @@ pub const Config = struct { } out.tap_count = @intCast(taps.array.items.len); try kinds(o, out.layers); + // the block's MLP is the fixed SiLU gate + const act = o.get("hidden_act") orelse return error.BadDraftConfig; + if (act != .string or !std.mem.eql(u8, act.string, "silu")) return error.UnsupportedDraftActivation; // the vendored defaults decide the block: non-causal attention, interleaved RoPE, a plain Markov head if (o.get("is_causal") != null) return error.UnsupportedDraftMask; const neox = o.get("rope_is_neox_style") orelse return error.BadDraftConfig; @@ -126,6 +129,7 @@ test "lfm2 dspark: other drafters and shapes are refused with their reason" { try std.testing.expectError(error.DraftForAnotherTarget, edited("\"num_target_layers\":40", "\"num_target_layers\":24")); try std.testing.expectError(error.DraftForAnotherTarget, edited("26,34]", "26,44]")); try std.testing.expectError(error.UnsupportedDraftShape, edited("\"head_dim\":64", "\"head_dim\":128")); + try std.testing.expectError(error.UnsupportedDraftActivation, edited("\"silu\"", "\"gelu\"")); try std.testing.expectError(error.UnsupportedDraftRope, edited("\"rope_is_neox_style\":false", "\"rope_is_neox_style\":true")); try std.testing.expectError(error.UnsupportedMarkovHead, edited("\"vanilla\"", "\"gated\"")); try std.testing.expectError(error.UnsupportedDraftMask, edited("\"full_attention\"]", "\"sliding_attention\"]")); diff --git a/zig/src/families/lfm2_moe/dspark_test.zig b/zig/src/families/lfm2_moe/dspark_test.zig index 9ee31d868e..f4d7dbf461 100644 --- a/zig/src/families/lfm2_moe/dspark_test.zig +++ b/zig/src/families/lfm2_moe/dspark_test.zig @@ -4,6 +4,9 @@ const mtl = @import("metal"); const c = @import("config.zig"); const kernels = @import("kernels.zig"); const quant = @import("dspark_quant.zig"); +const block = @import("dspark_block.zig"); +const ds = @import("dspark.zig"); +const Pick = @import("tree_round_gpu").contract.Pick; const shared = mtl.ResourceOptions.shared; @@ -206,3 +209,44 @@ test "lfm2 dspark: lfm2_qk on half-split rows and gains gives interleaved RoPE's } } } + +/// lfm2_draft_ids on one pick (token, nonfinite flag) at a stride of 4 words: the id it wrote and the bad flag. +fn draftIds(r: *Rig, pick: Pick) ![2]u32 { + const picks = try r.buffer(u32, 4); + defer picks.deinit(); + const ids = try r.buffer(u32, 1); + defer ids.deinit(); + const bad = try r.buffer(u32, 1); + defer bad.deinit(); + picks.slice(u32, 4)[0] = pick.token; + picks.slice(u32, 4)[1] = pick.nonfinite; + try r.run("lfm2_draft_ids", &.{ picks, ids, bad }, [4]u32{ 4, 1, c.vocab, 0 }, 1, .{ 1, 1 }); + return .{ ids.slice(u32, 1)[0], bad.slice(u32, 1)[0] }; +} + +test "lfm2 dspark: held ids carry a flagged or invalid pick as token 0 and the bad flag, a good pick untouched" { + var r = try Rig.init(); + defer r.deinit(); + try std.testing.expectEqual([2]u32{ 5, 0 }, try draftIds(&r, .{ .token = 5 })); + try std.testing.expectEqual([2]u32{ c.vocab - 1, 0 }, try draftIds(&r, .{ .token = c.vocab - 1 })); + for ([_]u32{ 1, 2, 4, 7 }) |kind| try std.testing.expectEqual([2]u32{ 0, 1 }, try draftIds(&r, .{ .token = 9, .nonfinite = kind })); + try std.testing.expectEqual([2]u32{ 0, 1 }, try draftIds(&r, .{ .token = std.math.maxInt(u32) })); + try std.testing.expectEqual([2]u32{ 0, 1 }, try draftIds(&r, .{ .token = c.vocab })); +} + +test "lfm2 dspark: a settled hold refuses a NaN, infinity or invalid pick as the synchronous path does" { + var picks: [4 * ds.pass_rows]Pick = @splat(.{ .token = 3 }); + var a: [4]u32 = undefined; + var b: [4]u32 = undefined; + const chains = [_][]u32{ &a, &b }; + try block.copyPicks(&picks, &chains, &.{ 2, 2 }); + try std.testing.expectEqualSlices(u32, &.{ 3, 3 }, a[0..2]); + const bad = [_]Pick{ .{ .token = 3, .nonfinite = 1 }, .{ .token = 3, .nonfinite = 2 }, .{ .token = 3, .nonfinite = 4 }, .{ .token = std.math.maxInt(u32) }, .{ .token = c.vocab } }; + for (bad) |p| { + picks[ds.pass_rows + 1] = p; // the second block's second step + try std.testing.expectError(error.DraftNonfinite, block.copyPicks(&picks, &chains, &.{ 2, 2 })); + picks[ds.pass_rows + 1] = .{ .token = 3 }; + } + picks[3 * ds.pass_rows] = .{ .nonfinite = 1 }; // a step past the block's depth is not read + try block.copyPicks(&picks, &chains, &.{ 2, 2 }); +} diff --git a/zig/src/families/lfm2_moe/forward.zig b/zig/src/families/lfm2_moe/forward.zig index 10040c0d77..51408e2fd4 100644 --- a/zig/src/families/lfm2_moe/forward.zig +++ b/zig/src/families/lfm2_moe/forward.zig @@ -146,11 +146,13 @@ pub fn runHeld(m: *Model, s: *st.Scratch, segments: []const Segment, tokens: []c 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 }; + if (held.len > 0) @memset(s.bad.contents()[0..4], 0); for (held) |h| { enc.pipe("lfm2_draft_ids"); e.setBuffer(h.src, h.off, 0); e.setValue([4]u32{ h.stride, h.count, c.vocab, 0 }, 1); e.setBuffer(s.ids, h.row * 4, 2); + e.setBuffer(s.bad, 0, 3); enc.threads(.{ h.count, 1, 1 }, .{ 32, 1, 1 }); } const embedding = m.weights.embedding; @@ -198,6 +200,7 @@ pub fn runHeld(m: *Model, s: *st.Scratch, segments: []const Segment, tokens: []c std.debug.print("LFM2 forward: {s}\n", .{msg}); return error.GpuFailed; } + if (held.len > 0 and std.mem.readInt(u32, s.bad.contents()[0..4], .little) != 0) return error.DraftNonfinite; var base: usize = 0; for (segments) |seg| { const logit_row = if (head == .all) base + seg.rows - 1 else 0; diff --git a/zig/src/families/lfm2_moe/state.zig b/zig/src/families/lfm2_moe/state.zig index 1bc749ae95..883f670fbe 100644 --- a/zig/src/families/lfm2_moe/state.zig +++ b/zig/src/families/lfm2_moe/state.zig @@ -119,6 +119,7 @@ pub const Scratch = struct { memory: Storage, rows: usize, ids: mtl.Buffer, + bad: mtl.Buffer, // set by lfm2_draft_ids when a held pick is flagged nonfinite or past the vocabulary h: mtl.Buffer, x: mtl.Buffer, r: mtl.Buffer, @@ -170,6 +171,7 @@ pub const Scratch = struct { const mem = &out.memory; const slots = rows * c.topk; out.ids = try mem.alloc(rows * 4); + out.bad = try mem.alloc(4); inline for (.{ "h", "x", "r", "y", "q" }) |name| @field(out, name) = try mem.alloc(rows * c.hidden * 2); out.queries = try mem.alloc((rows + attn_rows) * c.hidden * 2); // a prompt attention tile reads attn_rows from its first out.bcx = try mem.alloc(rows * 3 * c.hidden * 2); diff --git a/zig/tests/contract/metal_lfm2_dspark.zig b/zig/tests/contract/metal_lfm2_dspark.zig index 63d7bcf590..f119d00c66 100644 --- a/zig/tests/contract/metal_lfm2_dspark.zig +++ b/zig/tests/contract/metal_lfm2_dspark.zig @@ -63,6 +63,7 @@ pub const Probe = struct { ids: []u32 = &.{}, args: [8]rig_.Ref = undefined, y: rig_.Ref = undefined, + bad: rig_.Ref = undefined, // draft_ids: the flag word it sets for a flagged or invalid pick pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, w: Which) !Probe { const kernels: [2][:0]const u8 = switch (w) { @@ -109,10 +110,13 @@ pub const Probe = struct { 1 => vocab - 1, else => vocab + @as(u32, @intCast(i % 5)), }; + // the flag word after each pick: zero, but planted nonzero on the last of every four streams + for (0..s.rows) |r| p.ids[r * stride + 1] = if (r % 4 == 3) 1 else 0; p.args[0] = try p.rig.input(std.mem.sliceAsBytes(p.ids)); const at = [4]u32{ stride, @intCast(s.rows), vocab, 0 }; p.args[1] = try p.rig.input(std.mem.asBytes(&at)); p.y = try p.rig.output(p.count()); + p.bad = try p.rig.output(4); return true; } if (s.n != vocab or s.k != rank or s.rows == 0) return false; @@ -165,6 +169,7 @@ pub const Probe = struct { 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.y.buf, p.y.off, 2); + e.setBuffer(p.bad.buf, p.bad.off, 3); e.dispatchThreads(mtl.Size.of(p.rows, 1, 1), mtl.Size.of(32, 1, 1)); return p.rig.finish(pass); } @@ -194,11 +199,16 @@ pub const Probe = struct { /// Each stream's pick against float64 scores with the kernel's bf16 roundings: the best, but for near ties. pub fn check(p: *Probe, _: g.Gate) !?[]const u8 { const out = p.rig.bytes(p.y, p.count()); - if (p.which == .draft_ids) { // exact: each stream's first pick, or 0 for a token past the vocabulary + if (p.which == .draft_ids) { // exact: each stream's first pick, or 0 for a token past the vocabulary or a flag + var any = false; for (0..p.rows) |r| { const token = p.ids[r * stride]; - if (std.mem.readInt(u32, out[4 * r ..][0..4], .little) != (if (token < vocab) token else 0)) return "an id differs from its pick"; + const ok = token < vocab and p.ids[r * stride + 1] == 0; + any = any or !ok; + if (std.mem.readInt(u32, out[4 * r ..][0..4], .little) != (if (ok) token else 0)) return "an id differs from its pick"; } + const raised = std.mem.readInt(u32, p.rig.bytes(p.bad, 4)[0..4], .little); + if ((raised != 0) != any) return "the bad flag does not follow the flagged and invalid picks"; return null; } for (0..p.rows) |s| { From 0437d36a2981432b55af26ed97ab8a0c100ba311 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 18:44:43 +0100 Subject: [PATCH 30/46] tf-lfm2-dspark: the reference reads the target's embedding width from the checkpoint and joins its workers on every path --- zig/tests/lfm2_dspark.zig | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/zig/tests/lfm2_dspark.zig b/zig/tests/lfm2_dspark.zig index f95ce615e5..d47052c8d1 100644 --- a/zig/tests/lfm2_dspark.zig +++ b/zig/tests/lfm2_dspark.zig @@ -49,12 +49,12 @@ fn matmul(gpa: std.mem.Allocator, w: []const f32, x: []const f32, y: []f32, n: u const hs = try gpa.alloc(std.Thread, threads); defer gpa.free(hs); var made: usize = 0; + defer for (hs[0..made]) |h| h.join(); // a failed spawn still joins the workers already running var at: usize = 0; while (at < n) : (at += per) { hs[made] = try std.Thread.spawn(.{}, Job.run, .{ w, x, y, n, k, rows, at, @min(n, at + per) }); made += 1; } - for (hs[0..made]) |h| h.join(); } fn rms(x: []f32, gain: []const f32) void { @@ -76,19 +76,21 @@ fn rope(x: []f32, pos: u64, theta: f32) void { }; } -/// The target's embedding row `id`, dequantized from its 6-bit codes (MLX's scale * q + bias). +/// The target's embedding row `id`, dequantized from its codes at the checkpoint's own width (scale * q + bias). fn embedRow(m: *q.Model, id: u32, out: []f32) void { const e = m.weights.embedding; - const codes = (e.w.contents() + e.w_off)[@as(usize, id) * c.hidden * 6 / 8 ..][0 .. c.hidden * 6 / 8]; + const bits: usize = e.bits; + const row_bytes = c.hidden * bits / 8; + const codes = (e.w.contents() + e.w_off)[@as(usize, id) * row_bytes ..][0..row_bytes]; const sc = @as([*]const u16, @ptrCast(@alignCast(e.scales.contents() + e.s_off)))[@as(usize, id) * c.hidden / 64 ..][0 .. c.hidden / 64]; const bi = @as([*]const u16, @ptrCast(@alignCast(e.biases.contents() + e.b_off)))[@as(usize, id) * c.hidden / 64 ..][0 .. c.hidden / 64]; - for (0..c.hidden / 4) |u| { - const w = codes[u * 3 ..][0..3]; - const qs = [4]u32{ w[0] & 0x3f, (w[0] >> 6) | ((@as(u32, w[1]) & 0x0f) << 2), (w[1] >> 4) | ((@as(u32, w[2]) & 0x03) << 4), w[2] >> 2 }; - for (qs, 0..) |qv, j| { - const g = (u * 4 + j) / 64; - out[u * 4 + j] = quant.toF32(sc[g]) * @as(f32, @floatFromInt(qv)) + quant.toF32(bi[g]); - } + for (0..c.hidden) |j| { + // codes are a little-endian bit stream: code j takes bits [j * bits, (j + 1) * bits) + const at = j * bits; + var word: u32 = codes[at / 8]; + if (at / 8 + 1 < row_bytes) word |= @as(u32, codes[at / 8 + 1]) << 8; + const qv = (word >> @intCast(at % 8)) & ((@as(u32, 1) << @intCast(bits)) - 1); + out[j] = quant.toF32(sc[j / 64]) * @as(f32, @floatFromInt(qv)) + quant.toF32(bi[j / 64]); } } From f635e63ca521995cee4f0a8c7bf639c55a9ea583 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 18:46:01 +0100 Subject: [PATCH 31/46] lfm2 dspark: run the released bf16 matrices as stored (--drafter-bits 0, the default), no quantizing at load The drafter's linears and Markov W2 stay bf16: new rows and tile kernels for the linears and a bf16 scan. An explicit --drafter-bits other than 0 is refused by name. 8-bit codes stay behind Options.bits for tools. --- zig/kernels/metal/lfm2/lfm2_dspark.metal | 131 ++++++++++++++++-- zig/src/core/engine_api.zig | 2 + zig/src/families/lfm2_moe/dspark.zig | 4 +- .../families/lfm2_moe/dspark_bf16_test.zig | 121 ++++++++++++++++ zig/src/families/lfm2_moe/dspark_block.zig | 43 +++++- zig/src/families/lfm2_moe/dspark_test.zig | 10 +- zig/src/families/lfm2_moe/dspark_weights.zig | 43 +++--- zig/src/families/lfm2_moe/kernels.zig | 3 + zig/src/native/lfm2_moe.zig | 8 +- zig/src/server/cli.zig | 5 +- zig/src/server/cli_test.zig | 2 +- zig/src/tensorfold.zig | 1 + zig/tests/lfm2_dspark.zig | 2 +- 13 files changed, 327 insertions(+), 48 deletions(-) create mode 100644 zig/src/families/lfm2_moe/dspark_bf16_test.zig diff --git a/zig/kernels/metal/lfm2/lfm2_dspark.metal b/zig/kernels/metal/lfm2/lfm2_dspark.metal index 050fd5e8cc..e63d902867 100644 --- a/zig/kernels/metal/lfm2/lfm2_dspark.metal +++ b/zig/kernels/metal/lfm2/lfm2_dspark.metal @@ -1,18 +1,22 @@ // DSpark's Markov step, compiled after lfm2.metal's text (its includes and constants). -// DSpark's Markov scan: a threadgroup's best of 512 tokens of a row, bf16(base + bf16(W2 . W1[prev])), W2 8-bit codes. +// DSpark's Markov scan, W2 as 8-bit codes (BF16 false) or bf16 as released (true): a threadgroup's best of 512 tokens. struct lfm2_markov_args { uint stride, rank, vocab, prev_stride; }; constant constexpr uint MARKOV_SPAN = 512; -[[kernel]] void lfm2_markov_scan(const device bfloat* logits [[buffer(0)]], const device bfloat* w1 [[buffer(1)]], - const device uint8_t* w2 [[buffer(2)]], const device uint* prev [[buffer(3)]], - constant lfm2_markov_args& p [[buffer(4)]], device float* best [[buffer(5)]], - device uint* at [[buffer(6)]], const device bfloat* w2s [[buffer(7)]], - const device bfloat* w2b [[buffer(8)]], uint2 tg [[threadgroup_position_in_grid]], - uint2 grid [[threadgroups_per_grid]], uint sg [[simdgroup_index_in_threadgroup]], - uint lane [[thread_index_in_simdgroup]]) { - threadgroup float e[256]; - threadgroup float sv[8]; - threadgroup uint si[8]; + +// Eight bf16 in a uint4 as floats (a bf16 is the high half of its float). +inline void lfm2_unpack8(uint4 v, thread float* f) { + for (int j = 0; j < 4; ++j) { + f[2 * j] = as_type(v[j] << 16); + f[2 * j + 1] = as_type(v[j] & 0xffff0000u); + } +} +template +inline void lfm2_markov_scan_t(const device bfloat* logits, const device bfloat* w1, const device uint8_t* w2, + const device uint* prev, constant lfm2_markov_args& p, device float* best, + device uint* at, const device bfloat* w2s, const device bfloat* w2b, uint2 tg, + uint2 grid, uint sg, uint lane, threadgroup float* e, threadgroup float* sv, + threadgroup uint* si) { const uint s = tg.y, t = sg * 32 + lane, id = prev[size_t(s) * p.prev_stride]; for (uint i = t; i < p.rank; i += 256) e[i] = id < p.vocab ? float(w1[size_t(id) * p.rank + i]) : 0.0f; threadgroup_barrier(mem_flags::mem_threadgroup); @@ -24,6 +28,12 @@ constant constexpr uint MARKOV_SPAN = 512; for (uint v = first + sg; v < last; v += 8) { float d = 0.0f; for (uint i = 8 * lane; i < p.rank; i += 256) { + if (BF16) { + float f[8]; + lfm2_unpack8(*(const device uint4*)((const device bfloat*)w2 + size_t(v) * p.rank + i), f); + for (int j = 0; j < 8; ++j) d = fma(f[j], e[i + j], d); + continue; + } const size_t g = (size_t(v) * p.rank + i) / 64; const float sc = w2s[g], bi = w2b[g]; const uchar4 lo = *(const device uchar4*)(w2 + size_t(v) * p.rank + i); @@ -58,6 +68,32 @@ constant constexpr uint MARKOV_SPAN = 512; at[size_t(s) * grid.x + tg.x] = bi; } +[[kernel]] void lfm2_markov_scan(const device bfloat* logits [[buffer(0)]], const device bfloat* w1 [[buffer(1)]], + const device uint8_t* w2 [[buffer(2)]], const device uint* prev [[buffer(3)]], + constant lfm2_markov_args& p [[buffer(4)]], device float* best [[buffer(5)]], + device uint* at [[buffer(6)]], const device bfloat* w2s [[buffer(7)]], + const device bfloat* w2b [[buffer(8)]], uint2 tg [[threadgroup_position_in_grid]], + uint2 grid [[threadgroups_per_grid]], uint sg [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + threadgroup float e[256]; + threadgroup float sv[8]; + threadgroup uint si[8]; + lfm2_markov_scan_t(logits, w1, w2, prev, p, best, at, w2s, w2b, tg, grid, sg, lane, e, sv, si); +} + +// The same scan over the released bf16 W2 [vocab, rank] (buffers 7 and 8 unused). +[[kernel]] void lfm2_markov_scan_bf16(const device bfloat* logits [[buffer(0)]], const device bfloat* w1 [[buffer(1)]], + const device uint8_t* w2 [[buffer(2)]], const device uint* prev [[buffer(3)]], + constant lfm2_markov_args& p [[buffer(4)]], device float* best [[buffer(5)]], + device uint* at [[buffer(6)]], uint2 tg [[threadgroup_position_in_grid]], + uint2 grid [[threadgroups_per_grid]], uint sg [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + threadgroup float e[256]; + threadgroup float sv[8]; + threadgroup uint si[8]; + lfm2_markov_scan_t(logits, w1, w2, prev, p, best, at, w1, w1, tg, grid, sg, lane, e, sv, si); +} + // DSpark's Markov step pick: the scan's best over a stream's threadgroups (ties to the lowest token), as a Pick. [[kernel]] void lfm2_markov_pick(const device float* best [[buffer(0)]], const device uint* at [[buffer(1)]], constant uint& groups [[buffer(2)]], device uint* picks [[buffer(3)]], @@ -105,3 +141,76 @@ constant constexpr uint MARKOV_SPAN = 512; const uint token = picks[size_t(i) * p.x]; ids[i] = token < p.z ? token : 0u; } + +// DSpark's bf16 linears (the released matrices as stored): y [rows, n] = x [rows, k] W [n, k]^T, fp32 sums, bf16 out. +struct lfm2_bf16_args { uint rows, n, k; }; +constant constexpr uint BF16_ROWS = 16; + +// Few rows: a simdgroup a column of W over 16 rows, so W is read once a 16 rows. k a multiple of 8. +[[kernel]] void lfm2_bf16_rows(const device bfloat* x [[buffer(0)]], const device bfloat* w [[buffer(1)]], + constant lfm2_bf16_args& p [[buffer(2)]], device bfloat* y [[buffer(3)]], + uint2 tg [[threadgroup_position_in_grid]], uint sg [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + const uint col = tg.x * 8 + sg, r0 = tg.y * BF16_ROWS; + if (col >= p.n) return; + const uint rows = min(BF16_ROWS, p.rows - r0); + float acc[BF16_ROWS]; + for (uint r = 0; r < BF16_ROWS; ++r) acc[r] = 0.0f; + const device bfloat* wr = w + size_t(col) * p.k; + const device bfloat* xr = x + size_t(r0) * p.k; + for (uint i = 8 * lane; i < p.k; i += 256) { + float wv[8]; + lfm2_unpack8(*(const device uint4*)(wr + i), wv); + for (uint r = 0; r < BF16_ROWS; ++r) { + if (r >= rows) break; + float xv[8]; + lfm2_unpack8(*(const device uint4*)(xr + size_t(r) * p.k + i), xv); + for (int j = 0; j < 8; ++j) acc[r] = fma(wv[j], xv[j], acc[r]); + } + } + for (uint r = 0; r < BF16_ROWS; ++r) { + const float sum = simd_sum(acc[r]); + if (lane == 0 && r < rows) y[size_t(r0 + r) * p.n + col] = static_cast(sum); + } +} + +// Many rows: 64x64 tiles of y, four simdgroups of 32x32, k in steps of 32. k a multiple of 32. +[[kernel]] void lfm2_bf16_mma(const device bfloat* x [[buffer(0)]], const device bfloat* w [[buffer(1)]], + constant lfm2_bf16_args& p [[buffer(2)]], device bfloat* y [[buffer(3)]], + uint2 tg [[threadgroup_position_in_grid]], uint sg [[simdgroup_index_in_threadgroup]], + uint t [[thread_index_in_threadgroup]]) { + threadgroup bfloat xs[64 * 32]; + threadgroup bfloat ws[64 * 32]; + threadgroup float out[4][32 * 32]; + const uint c0 = tg.x * 64, r0 = tg.y * 64; + simdgroup_matrix acc[4][4]; + for (uint i = 0; i < 4; ++i) + for (uint j = 0; j < 4; ++j) acc[i][j] = simdgroup_matrix(0.0f); + const uint sr = (sg / 2) * 32, sc = (sg % 2) * 32; + for (uint k0 = 0; k0 < p.k; k0 += 32) { + for (uint i = t; i < 64 * 4; i += 128) { // 64 rows of 32 bf16 as 4 chunks of 8 + const uint r = i / 4, q = (i % 4) * 8; + for (uint j = 0; j < 8; ++j) { + xs[r * 32 + q + j] = r0 + r < p.rows ? x[size_t(r0 + r) * p.k + k0 + q + j] : bfloat(0.0f); + ws[r * 32 + q + j] = c0 + r < p.n ? w[size_t(c0 + r) * p.k + k0 + q + j] : bfloat(0.0f); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint kk = 0; kk < 32; kk += 8) { + simdgroup_matrix a[4], b[4]; + for (uint i = 0; i < 4; ++i) simdgroup_load(a[i], xs + (sr + 8 * i) * 32 + kk, 32); + for (uint j = 0; j < 4; ++j) simdgroup_load(b[j], ws + (sc + 8 * j) * 32 + kk, 32, ulong2(0, 0), true); + for (uint i = 0; i < 4; ++i) + for (uint j = 0; j < 4; ++j) simdgroup_multiply_accumulate(acc[i][j], a[i], b[j], acc[i][j]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + for (uint i = 0; i < 4; ++i) + for (uint j = 0; j < 4; ++j) simdgroup_store(acc[i][j], out[sg] + 8 * i * 32 + 8 * j, 32); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (uint i = t % 32; i < 32 * 32; i += 32) { + const uint r = i / 32, c = i % 32; + const uint row = r0 + sr + r, col = c0 + sc + c; + if (row < p.rows && col < p.n) y[size_t(row) * p.n + col] = static_cast(out[sg][i]); + } +} diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index a50d0de404..4838c4524e 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -277,6 +277,8 @@ pub const Open = struct { /// --drafter: a draft model's directory (Qwen3.8-27B's DFlash2; null: none), and --drafter-bits: 4, or 0 for bf16. drafter: ?[]const u8 = null, drafter_bits: u8 = 4, + /// Whether --drafter-bits was given: LFM2's drafter runs bf16 unless it is refused a width it has no mode for. + drafter_bits_given: bool = false, /// --record: each stream's committed rows and tapped features written under this folder, for drafter training. record: ?[]const u8 = null, speed_up: ?[]const u8 = null, diff --git a/zig/src/families/lfm2_moe/dspark.zig b/zig/src/families/lfm2_moe/dspark.zig index 850c1ce34c..2b4e9698b9 100644 --- a/zig/src/families/lfm2_moe/dspark.zig +++ b/zig/src/families/lfm2_moe/dspark.zig @@ -22,8 +22,8 @@ pub const absorb_rows = 2048; pub const stage_rows = 128; const kw = c.kv_heads * c.head_dim; -/// `bits`: the width the bf16 linears are quantized to (8 keeps the fp32 reference's drafts; 4 does not). -pub const Options = struct { capacity: usize, bits: u8 = 8 }; +/// `bits`: 0 runs the released bf16 matrices as stored; 8 quantizes the linears to 8-bit codes at load. +pub const Options = struct { capacity: usize, bits: u8 = 0 }; /// One stream's context keys and values a layer ([kv head][cap][head dim], as the target's attention reads them). pub const Lane = struct { diff --git a/zig/src/families/lfm2_moe/dspark_bf16_test.zig b/zig/src/families/lfm2_moe/dspark_bf16_test.zig new file mode 100644 index 0000000000..293549bde1 --- /dev/null +++ b/zig/src/families/lfm2_moe/dspark_bf16_test.zig @@ -0,0 +1,121 @@ +//! The DSpark drafter's bf16 kernels on the GPU against host references: the linears and the Markov scan over bf16 W2. +const std = @import("std"); +const mtl = @import("metal"); +const quant = @import("dspark_quant.zig"); +const base = @import("dspark_test.zig"); + +const Args = extern struct { rows: u32, n: u32, k: u32 }; + +/// y = x W^T through `key` on random bf16 data, each output within a bf16 step of the fp64 product. +fn linear(comptime key: []const u8, rows: usize, n: usize, k: usize, groups: [2]usize, threads: usize) !void { + var r = try base.Rig.init(); + defer r.deinit(); + const x = try r.buffer(u16, rows * k); + defer x.deinit(); + const w = try r.buffer(u16, n * k); + defer w.deinit(); + const y = try r.buffer(u16, rows * n); + defer y.deinit(); + base.noise(31, x.slice(u16, rows * k), 2); + base.noise(32, w.slice(u16, n * k), 0.5); + const cb = r.queue.commandBuffer(); + const e = cb.compute(.serial); + e.setPipeline(r.k.get(key)); + e.setBuffer(x, 0, 0); + e.setBuffer(w, 0, 1); + e.setValue(Args{ .rows = @intCast(rows), .n = @intCast(n), .k = @intCast(k) }, 2); + e.setBuffer(y, 0, 3); + e.dispatchGroups(mtl.Size.of(groups[0], groups[1], 1), mtl.Size.of(threads, 1, 1)); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |_| return error.GpuFailed; + const xs = x.slice(u16, rows * k); + const ws = w.slice(u16, n * k); + const ys = y.slice(u16, rows * n); + for (0..rows) |i| for (0..n) |j| { + var dot: f64 = 0; + for (0..k) |t| dot += @as(f64, quant.toF32(xs[i * k + t])) * @as(f64, quant.toF32(ws[j * k + t])); + const got: f64 = quant.toF32(ys[i * n + j]); + try std.testing.expect(@abs(got - dot) <= @abs(dot) * 0x1p-7 + 1e-3); + }; +} + +test "lfm2 dspark: the bf16 rows kernel is x W^T for one row, a full 16 and a ragged 19 (two row groups)" { + for ([_]usize{ 1, 16, 19 }) |rows| try linear("lfm2_bf16_rows", rows, 100, 288, .{ 13, (rows + 15) / 16 }, 256); +} + +test "lfm2 dspark: the bf16 tile kernel is x W^T for a ragged 70 rows and 100 columns, k a multiple of 32" { + try linear("lfm2_bf16_mma", 70, 100, 288, .{ 2, 2 }, 128); +} + +test "lfm2 dspark: the bf16 tile kernel at the drafter's widths" { + try linear("lfm2_bf16_mma", 33, 512, 2048, .{ 8, 1 }, 128); +} + +test "lfm2 dspark: a Markov step over bf16 W2 picks the argmax of bf16(base + bf16(W2 . W1[prev])), a NaN row flagged" { + var r = try base.Rig.init(); + defer r.deinit(); + const rank = 256; + const vocab = 1300; + const streams = 4; + const groups = (vocab + 511) / 512; + const w1 = try r.buffer(u16, vocab * rank); + defer w1.deinit(); + base.noise(11, w1.slice(u16, vocab * rank), 1); + const w2 = try r.buffer(u16, vocab * rank); + defer w2.deinit(); + base.noise(12, w2.slice(u16, vocab * rank), 0.5); + const ids = try r.buffer(u32, 2 * streams); + defer ids.deinit(); + @memcpy(ids.slice(u32, 2 * streams), &[_]u32{ 17, 9, vocab - 1, 9, vocab + 5, 0, 3, 3 }); + const logits = try r.buffer(u16, streams * 2 * vocab); + defer logits.deinit(); + base.noise(14, logits.slice(u16, streams * 2 * vocab), 4); + logits.slice(u16, streams * 2 * vocab)[3 * 2 * vocab + 700] = 0x7fc0; + const best = try r.buffer(f32, streams * groups); + defer best.deinit(); + const at = try r.buffer(u32, streams * groups); + defer at.deinit(); + const picks = try r.buffer(u32, 2 * streams); + defer picks.deinit(); + const cb = r.queue.commandBuffer(); + const e = cb.compute(.serial); + e.setPipeline(r.k.get("lfm2_markov_scan_bf16")); + for ([_]mtl.Buffer{ logits, w1, w2, ids }, 0..) |b, i| e.setBuffer(b, 0, i); + e.setValue([4]u32{ 2 * vocab, rank, vocab, 2 }, 4); + for ([_]mtl.Buffer{ best, at }, 5..) |b, i| e.setBuffer(b, 0, i); + e.dispatchGroups(mtl.Size.of(groups, streams, 1), mtl.Size.of(256, 1, 1)); + e.setPipeline(r.k.get("lfm2_markov_pick")); + e.setBuffer(best, 0, 0); + e.setBuffer(at, 0, 1); + e.setValue(@as(u32, groups), 2); + e.setBuffer(picks, 0, 3); + e.dispatchGroups(mtl.Size.of(streams, 1, 1), mtl.Size.of(256, 1, 1)); + e.end(); + cb.commit(); + cb.wait(); + const got = picks.slice(u32, 2 * streams); + const logit = logits.slice(u16, streams * 2 * vocab); + for (0..streams) |s| { + if (s == 3) { + try std.testing.expectEqual(@as(u32, 1), got[2 * s + 1]); + continue; + } + try std.testing.expectEqual(@as(u32, 0), got[2 * s + 1]); + const id = ids.slice(u32, 2 * streams)[2 * s]; + var top: f64 = -std.math.inf(f64); + var mine: f64 = 0; + for (0..vocab) |v| { + var dot: f64 = 0; + for (0..rank) |i| { + const e1: f64 = if (id < vocab) quant.toF32(w1.slice(u16, vocab * rank)[id * rank + i]) else 0; + dot += @as(f64, quant.toF32(w2.slice(u16, vocab * rank)[v * rank + i])) * e1; + } + const score: f64 = quant.toF32(quant.toBf16(quant.toF32(logit[s * 2 * vocab + v]) + quant.toF32(quant.toBf16(@floatCast(dot))))); + top = @max(top, score); + if (v == got[2 * s]) mine = score; + } + try std.testing.expect(got[2 * s] < vocab and mine >= top - @abs(top) * 0x1p-7); + } +} diff --git a/zig/src/families/lfm2_moe/dspark_block.zig b/zig/src/families/lfm2_moe/dspark_block.zig index d3f6cbabb8..f94aba037e 100644 --- a/zig/src/families/lfm2_moe/dspark_block.zig +++ b/zig/src/families/lfm2_moe/dspark_block.zig @@ -8,6 +8,7 @@ const tile = @import("../../core/tile_qmm.zig"); const am = @import("../../core/affine_mm.zig"); const round = @import("tree_round_gpu"); const ds = @import("dspark.zig"); +const dw = @import("dspark_weights.zig"); const Drafter = ds.Drafter; const Lane = ds.Lane; @@ -15,6 +16,7 @@ const kw = c.kv_heads * c.head_dim; const QkArgs = extern struct { len: u32, cap: u32, eps: f32, log2_theta: f32 }; const AttnArgs = extern struct { len: u32, cap: u32, scale: f32, subs: u32 }; const Pick = round.contract.Pick; +const Bf16Args = extern struct { rows: u32, n: u32, k: u32 }; const MarkovArgs = extern struct { stride: u32, rank: u32, vocab: u32, prev_stride: u32 }; /// Tokens a Markov scan threadgroup takes (the kernel's MARKOV_SPAN), and threadgroups a stream's row takes. pub const scan_span = 512; @@ -35,7 +37,11 @@ const Enc = struct { } /// y[rows] = x[rows] W^T: the twin up to its rows, past them the M5's tensor-unit kernels, else the tiles. - fn project(x: Enc, w: tile.Matrix, in: Ref, out: Ref, rows: usize) !void { + fn project(x: Enc, lin: dw.Linear, in: Ref, out: Ref, rows: usize) !void { + const w = switch (lin) { + .codes => |m| m, + .bf16 => |m| return x.bf16(m, in, out, rows), + }; const t = &x.d.target.kernels.tiles; if (rows > tile.twin_rows) if (x.d.target.kernels.naxPipes(w.bits)) |p| { const sums = Ref{ .buf = x.d.s.sums }; @@ -48,6 +54,20 @@ const Enc = struct { try tile.encode(x.e, t, w, in.buf, in.off, x.d.s.sums, 0, out.buf, out.off, rows); } + /// y[rows] = x[rows] W^T over the released bf16 matrix: 16 rows a pass of W up to 32 rows, else 64x64 tiles. + fn bf16(x: Enc, m: dw.Bf16, in: Ref, out: Ref, rows: usize) void { + const args = Bf16Args{ .rows = @intCast(rows), .n = @intCast(m.n), .k = @intCast(m.k) }; + const few = rows <= 32; + if (few) x.pipe("lfm2_bf16_rows") else x.pipe("lfm2_bf16_mma"); + x.e.setBuffer(in.buf, in.off, 0); + x.e.setBuffer(m.w, 0, 1); + x.e.setBytes(std.mem.asBytes(&args), 2); + x.e.setBuffer(out.buf, out.off, 3); + if (few) { + x.e.dispatchGroups(mtl.Size.of((m.n + 7) / 8, (rows + 15) / 16, 1), mtl.Size.of(256, 1, 1)); + } else x.e.dispatchGroups(mtl.Size.of((m.n + 63) / 64, (rows + 63) / 64, 1), mtl.Size.of(128, 1, 1)); + } + /// h += r when `add`, then out = norm(h) * gain: LFM2's own norm (the drafter's width and eps are the target's). fn norm(x: Enc, h: mtl.Buffer, r: mtl.Buffer, gain: mtl.Buffer, add: bool, out: mtl.Buffer, rows: usize) void { x.pipe("lfm2_norm"); @@ -231,7 +251,7 @@ pub fn hold(d: *Drafter, blocks: []const Block, signal: ?Signal) !?mtl.CommandBu try x.project(ly.down, .{ .buf = s.act }, .{ .buf = s.r }, out); } x.norm(s.h, s.r, d.w.norm, true, s.x, tail); - try x.project(fwd.matrix(emb), .{ .buf = s.x }, .{ .buf = s.logits }, tail); + try x.project(.{ .codes = fwd.matrix(emb) }, .{ .buf = s.x }, .{ .buf = s.logits }, tail); try markov(x, blocks.len, n, steps); if (signal) |sg| { // committed unwaited, the command buffer returned retained e.end(); @@ -272,17 +292,28 @@ fn markov(x: Enc, streams: usize, n: u32, depth: u32) !void { const e = x.e; for (0..depth) |i| { // step i's previous token: the pending one (each block's first id), then step i - 1's pick - x.pipe("lfm2_markov_scan"); + switch (d.w.w2) { + .codes => x.pipe("lfm2_markov_scan"), + .bf16 => x.pipe("lfm2_markov_scan_bf16"), + } e.setBuffer(s.logits, i * c.vocab * 2, 0); e.setBuffer(d.w.w1, 0, 1); - e.setBuffer(d.w.w2.w, d.w.w2.w_off, 2); + e.setBuffer(switch (d.w.w2) { + .codes => |m| m.w, + .bf16 => |m| m.w, + }, switch (d.w.w2) { + .codes => |m| m.w_off, + .bf16 => 0, + }, 2); if (i == 0) e.setBuffer(s.ids, 0, 3) else e.setBuffer(s.picks, (i - 1) * ds.pass_rows * @sizeOf(Pick), 3); const args = MarkovArgs{ .stride = n * c.vocab, .rank = d.config.rank, .vocab = c.vocab, .prev_stride = if (i == 0) n else @sizeOf(Pick) / 4 }; e.setBytes(std.mem.asBytes(&args), 4); e.setBuffer(s.best, 0, 5); e.setBuffer(s.at, 0, 6); - e.setBuffer(d.w.w2.scales, d.w.w2.s_off, 7); - e.setBuffer(d.w.w2.biases, d.w.w2.b_off, 8); + if (d.w.w2 == .codes) { + e.setBuffer(d.w.w2.codes.scales, d.w.w2.codes.s_off, 7); + e.setBuffer(d.w.w2.codes.biases, d.w.w2.codes.b_off, 8); + } e.dispatchGroups(mtl.Size.of(scan_groups, streams, 1), mtl.Size.of(256, 1, 1)); x.pipe("lfm2_markov_pick"); e.setBuffer(s.best, 0, 0); diff --git a/zig/src/families/lfm2_moe/dspark_test.zig b/zig/src/families/lfm2_moe/dspark_test.zig index 9ee31d868e..e9129ceabf 100644 --- a/zig/src/families/lfm2_moe/dspark_test.zig +++ b/zig/src/families/lfm2_moe/dspark_test.zig @@ -7,12 +7,12 @@ const quant = @import("dspark_quant.zig"); const shared = mtl.ResourceOptions.shared; -const Rig = struct { +pub const Rig = struct { device: mtl.Device, queue: mtl.Queue, k: kernels.Kernels, - fn init() !Rig { + pub fn init() !Rig { const device = mtl.Device.init() catch return error.SkipZigTest; errdefer device.deinit(); const queue = try device.queue(); @@ -20,13 +20,13 @@ const Rig = struct { return .{ .device = device, .queue = queue, .k = try kernels.load(device) }; } - fn deinit(r: *Rig) void { + pub 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 { + pub 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; @@ -46,7 +46,7 @@ const Rig = struct { } }; -fn noise(seed: u64, out: []u16, scale: f32) void { +pub fn noise(seed: u64, out: []u16, scale: f32) void { var prng = std.Random.DefaultPrng.init(seed); for (out) |*v| v.* = quant.toBf16((prng.random().float(f32) * 2 - 1) * scale); } diff --git a/zig/src/families/lfm2_moe/dspark_weights.zig b/zig/src/families/lfm2_moe/dspark_weights.zig index 81dc727e56..11249916c2 100644 --- a/zig/src/families/lfm2_moe/dspark_weights.zig +++ b/zig/src/families/lfm2_moe/dspark_weights.zig @@ -1,4 +1,4 @@ -//! The DSpark drafter's weights: bf16 linears quantized to MLX affine codes at load, norms and W1 kept as stored. +//! The DSpark drafter's weights: the released bf16 linears as stored, or on request quantized to 8-bit affine codes. const std = @import("std"); const mtl = @import("metal"); const c = @import("config.zig"); @@ -10,17 +10,23 @@ const tile = @import("../../core/tile_qmm.zig"); const shared = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; +/// An [n, k] bf16 matrix as released (rows only reordered for half-split RoPE where the checkpoint is read). +pub const Bf16 = struct { w: mtl.Buffer, n: usize, k: usize }; + +/// A linear: bf16 as stored (bits 0), or MLX affine codes at the loader's width. +pub const Linear = union(enum) { bf16: Bf16, codes: tile.Matrix }; + /// A bf16 vector in its own buffer (norm gains; RoPE-paired ones in the half-split order the kernels rotate). pub const Vector = mtl.Buffer; pub const Layer = struct { - q: tile.Matrix, - k: tile.Matrix, - v: tile.Matrix, - o: tile.Matrix, - gate: tile.Matrix, - up: tile.Matrix, - down: tile.Matrix, + q: Linear, + k: Linear, + v: Linear, + o: Linear, + gate: Linear, + up: Linear, + down: Linear, input_norm: Vector, post_norm: Vector, q_norm: Vector, @@ -28,13 +34,13 @@ pub const Layer = struct { }; pub const Weights = struct { - fc: tile.Matrix, + fc: Linear, hidden_norm: Vector, norm: Vector, layers: [8]Layer, - /// The Markov head: W1 [vocab, rank] bf16 as stored (rows gathered by token), W2 [vocab, rank] as 8-bit codes. + /// The Markov head: W1 [vocab, rank] bf16 as stored (rows gathered by token), W2 [vocab, rank] as the linears. w1: mtl.Buffer, - w2: tile.Matrix, + w2: Linear, owned: std.ArrayList(mtl.Buffer) = .empty, pub fn deinit(w: *Weights, gpa: std.mem.Allocator) void { @@ -80,8 +86,8 @@ const Loader = struct { return b; } - /// An [n, k] bf16 matrix as affine codes at the loader's width; rows reordered for half-split RoPE when `paired`. - fn matrix(l: Loader, name: []const u8, n: usize, k: usize, paired: bool) !tile.Matrix { + /// An [n, k] bf16 matrix as it is or as affine codes; rows reordered for half-split RoPE when `paired`. + fn matrix(l: Loader, name: []const u8, n: usize, k: usize, paired: bool) !Linear { const t = try l.tensor(name, &.{ n, k }); var rows = t.host(u16); var moved: ?[]u16 = null; @@ -92,12 +98,17 @@ const Loader = struct { for (0..n) |r| @memcpy(m[r * k ..][0..k], rows[quant.halfSplit(r, c.head_dim) * k ..][0..k]); rows = m; } + if (l.bits == 0) { + const b = try l.buffer(n * k * 2); + @memcpy(b.slice(u16, n * k), rows); + return .{ .bf16 = .{ .w = b, .n = n, .k = k } }; + } const codes = try l.buffer(n * k * l.bits / 8); const scales = try l.buffer(n * k / 64 * 2); const biases = try l.buffer(n * k / 64 * 2); const job: quant.Job = .{ .src = rows, .k = k, .bits = l.bits, .codes = codes.slice(u8, n * k * l.bits / 8), .scales = scales.slice(u16, n * k / 64), .biases = biases.slice(u16, n * k / 64) }; try job.all(l.gpa, n); - return .{ .w = codes, .scales = scales, .biases = biases, .n = n, .k = k, .bits = l.bits }; + return .{ .codes = .{ .w = codes, .scales = scales, .biases = biases, .n = n, .k = k, .bits = l.bits } }; } }; @@ -105,9 +116,9 @@ fn part(buf: []u8, layer: usize, name: []const u8) ![]const u8 { return std.fmt.bufPrint(buf, "layers.{d}.{s}.weight", .{ layer, name }); } -/// The drafter's checkpoint read and quantized at `bits` (4 or 8); the checkpoint's own buffers are released after. +/// The drafter's checkpoint read as released (`bits` 0) or quantized to 8-bit codes; its own buffers are freed after. pub fn load(gpa: std.mem.Allocator, device: mtl.Device, io: std.Io, dir: []const u8, cfg: *const dc.Config, bits: u8) !Weights { - if (bits != 4 and bits != 8) return error.UnsupportedDraftBits; + if (bits != 0 and bits != 8) return error.UnsupportedDraftBits; var check = ckpt.Checkpoint.init(gpa); defer check.deinit(); const files = try shards.shardFiles(gpa, io, dir); diff --git a/zig/src/families/lfm2_moe/kernels.zig b/zig/src/families/lfm2_moe/kernels.zig index 936ba33197..7d7658cce4 100644 --- a/zig/src/families/lfm2_moe/kernels.zig +++ b/zig/src/families/lfm2_moe/kernels.zig @@ -31,6 +31,9 @@ const names = [_][:0]const u8{ "lfm2_attn_merge_segs", "lfm2_tap", "lfm2_markov_scan", + "lfm2_markov_scan_bf16", + "lfm2_bf16_rows", + "lfm2_bf16_mma", "lfm2_markov_pick", "lfm2_draft_ids", }; diff --git a/zig/src/native/lfm2_moe.zig b/zig/src/native/lfm2_moe.zig index a78d67d97b..e12e0a189a 100644 --- a/zig/src/native/lfm2_moe.zig +++ b/zig/src/native/lfm2_moe.zig @@ -82,12 +82,12 @@ pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Ope h.draft = null; h.wrap = null; if (if (o.drafts) o.drafter else null) |dir| { - if (o.drafter_bits == 0) { - problem.* = "--drafter-bits 0: the LFM2 DSpark drafter runs its bf16 linears as 8-bit codes, quantized at load"; + if (o.drafter_bits_given and o.drafter_bits != 0) { + problem.* = "--drafter-bits: the LFM2 DSpark drafter runs the released bf16 matrices; it takes 0 and no other width"; release(h); return null; } - h.draft = q.dspark.Drafter.init(gpa, io, model, dir, .{ .capacity = window }) catch |err| { + h.draft = q.dspark.Drafter.init(gpa, io, model, dir, .{ .capacity = window, .bits = 0 }) catch |err| { problem.* = try std.fmt.allocPrint(a, "the native LFM2 drafter cannot load {s} ({s}): it takes LiquidAI/LFM2.5-24B-A2B-DSpark", .{ dir, @errorName(err) }); release(h); return null; @@ -150,7 +150,7 @@ pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Ope }; h.host.keepalive_target = .{ .ctx = &h.warm, .tick = mtl.keepalive.Target.tick }; try h.host.start(); - if (h.draft) |d| std.log.info("lfm2: DSpark drafts from --drafter, {d} a round at most, {d}-bit codes", .{ d.config.block, d.options.bits }); + if (h.draft) |d| std.log.info("lfm2: DSpark drafts from --drafter, {d} a round at most, bf16 matrices as released", .{d.config.block}); return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; } diff --git a/zig/src/server/cli.zig b/zig/src/server/cli.zig index 0e2da51172..7f5909ab2b 100644 --- a/zig/src/server/cli.zig +++ b/zig/src/server/cli.zig @@ -210,7 +210,7 @@ pub const Args = struct { compact_keep: ?u32 = null, compact_memory: ?[]const u8 = null, drafter: ?[]const u8 = null, - drafter_bits: u8 = 4, + drafter_bits: ?u8 = null, // null: not given (a family without its own default takes 4) record: ?[]const u8 = null, // drafter training rows written under this folder kv_cache: []const u8 = "bf16", // fp8: e4m3 attention caches with a scale a row, about twice the context prompt_precision: []const u8 = "bf16", // int8: prompt matmuls quantize their inputs to int8 a 64-input group @@ -540,7 +540,8 @@ pub fn request(a: Allocator, dir: []const u8, model_type: []const u8, args: Args .lanes_fixed = parallelFixed(args.parallel), .drafts = !args.no_drafts, .drafter = args.drafter, - .drafter_bits = args.drafter_bits, + .drafter_bits = args.drafter_bits orelse 4, + .drafter_bits_given = args.drafter_bits != null, .record = args.record, .kv8 = std.mem.eql(u8, args.kv_cache, "fp8"), .prompt_int8 = std.mem.eql(u8, args.prompt_precision, "int8"), diff --git a/zig/src/server/cli_test.zig b/zig/src/server/cli_test.zig index f85890a632..a2ea3c4db2 100644 --- a/zig/src/server/cli_test.zig +++ b/zig/src/server/cli_test.zig @@ -163,7 +163,7 @@ test "parse and capabilities share the table" { try std.testing.expectEqual(@as(usize, 2), args.api_key.len); try std.testing.expect(!args.thinking); try std.testing.expectEqualStrings("x", (try parse(a, &.{ "m", "--drafter", "x" }, &u)).drafter.?); - try std.testing.expectEqual(@as(u8, 0), (try parse(a, &.{ "m", "--drafter-bits", "0" }, &u)).drafter_bits); + try std.testing.expectEqual(@as(?u8, 0), (try parse(a, &.{ "m", "--drafter-bits", "0" }, &u)).drafter_bits); try std.testing.expectError(error.Usage, parse(a, &.{ "m", "--drafter-bits", "8" }, &u)); try std.testing.expectEqualStrings("/r", (try parse(a, &.{ "m", "--record", "/r" }, &u)).record.?); try std.testing.expectEqualStrings("fp8", (try parse(a, &.{ "m", "--kv-cache", "fp8" }, &u)).kv_cache); diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 4a278138ac..ed6764f781 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -87,6 +87,7 @@ test { _ = @import("families/lfm2_moe/dspark_config.zig"); _ = @import("families/lfm2_moe/dspark_quant.zig"); _ = @import("families/lfm2_moe/dspark_test.zig"); + _ = @import("families/lfm2_moe/dspark_bf16_test.zig"); } pub const lane_projection = @import("core/lane_projection.zig"); diff --git a/zig/tests/lfm2_dspark.zig b/zig/tests/lfm2_dspark.zig index f95ce615e5..cc1a80baef 100644 --- a/zig/tests/lfm2_dspark.zig +++ b/zig/tests/lfm2_dspark.zig @@ -119,7 +119,7 @@ pub fn main(init: std.process.Init) !void { const gpa = init.gpa; const args = try init.minimal.args.toSlice(init.arena.allocator()); if (args.len < 4 or args.len > 6) return error.ExpectedModelDrafterPromptLengthRowsBits; - const bits: u8 = if (args.len == 6) try std.fmt.parseInt(u8, args[5], 10) else 8; + const bits: u8 = if (args.len == 6) try std.fmt.parseInt(u8, args[5], 10) else 0; const length = try std.fmt.parseInt(usize, args[3], 10); const pool = mtl.objc.Pool.push(); defer pool.pop(); 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 32/46] 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 b71e5f1cc281119cea4e358a787c49a4ab7da4da Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 18:48:29 +0100 Subject: [PATCH 33/46] lfm2 dspark: one-line comments for the held-pick refusal --- zig/kernels/metal/lfm2/lfm2_dspark.metal | 3 +-- zig/src/families/lfm2_moe/dspark_block.zig | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/zig/kernels/metal/lfm2/lfm2_dspark.metal b/zig/kernels/metal/lfm2/lfm2_dspark.metal index 34fbc3d704..1c5b834a85 100644 --- a/zig/kernels/metal/lfm2/lfm2_dspark.metal +++ b/zig/kernels/metal/lfm2/lfm2_dspark.metal @@ -98,8 +98,7 @@ constant constexpr uint MARKOV_SPAN = 512; } } -// A verify window's held drafts from the drafter's picks: picks[i * stride] into ids[i]. A pick flagged nonfinite -// (the word after it) or past the vocabulary becomes token 0 and sets bad, which the host reads after the verify. +// Held drafts from picks[i * stride] into ids[i]; a nonfinite or past-vocabulary pick is token 0 and sets bad. [[kernel]] void lfm2_draft_ids(const device uint* picks [[buffer(0)]], constant uint4& p [[buffer(1)]], device uint* ids [[buffer(2)]], device atomic_uint* bad [[buffer(3)]], uint i [[thread_position_in_grid]]) { diff --git a/zig/src/families/lfm2_moe/dspark_block.zig b/zig/src/families/lfm2_moe/dspark_block.zig index fab4f7886e..08ec1ca7f8 100644 --- a/zig/src/families/lfm2_moe/dspark_block.zig +++ b/zig/src/families/lfm2_moe/dspark_block.zig @@ -253,8 +253,7 @@ pub fn hold(d: *Drafter, blocks: []const Block, signal: ?Signal) !?mtl.CommandBu return null; } -/// A finished signalled hold's picks into the lanes' held chains; a nonfinite or invalid pick is refused as the -/// synchronous path refuses it, and no lane keeps a chain then. +/// A finished hold's picks into the lanes' held chains; a nonfinite or invalid pick is refused, as when synchronous. pub fn readPicks(d: *Drafter, lanes: []const *Lane, depths: []const u32) !void { const picks = d.s.picks.slice(Pick, d.config.block * ds.pass_rows); var chains: [ds.pass_rows][]u32 = undefined; From 874856e454d48363cc4fb7bc8daf8217d83df3e6 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 18:49:03 +0100 Subject: [PATCH 34/46] lfm2 dspark: one-line comment for Options --- zig/src/families/lfm2_moe/dspark.zig | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/zig/src/families/lfm2_moe/dspark.zig b/zig/src/families/lfm2_moe/dspark.zig index 84ea325213..a8bbd89201 100644 --- a/zig/src/families/lfm2_moe/dspark.zig +++ b/zig/src/families/lfm2_moe/dspark.zig @@ -29,8 +29,7 @@ pub const Verify = enum { depth, fixed, scheduled }; /// The verify when none is asked: the confidence-scheduled one. pub const default_verify: Verify = .scheduled; -/// `bits`: the width the bf16 linears are quantized to (8 keeps the fp32 reference's drafts; 4 does not). -/// `verify_asked`: --draft-verify was given, so a scheduled verify without the confidence head is refused, not depth. +/// `bits`: the quantized width (8 keeps the fp32 reference's drafts; 4 does not); `verify_asked`: --draft-verify given. pub const Options = struct { capacity: usize, bits: u8 = 8, verify: Verify = default_verify, verify_asked: bool = false }; /// An asked scheduled verify needs the checkpoint's confidence head; the default one falls back to the depth rule. From 095c4d70fa9f014fe0a7bd106d425ed27ec188f7 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 18:52:03 +0100 Subject: [PATCH 35/46] kernel contract: the draft ids probe zeroes its bad flag each run --- zig/tests/contract/metal_lfm2_dspark.zig | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/zig/tests/contract/metal_lfm2_dspark.zig b/zig/tests/contract/metal_lfm2_dspark.zig index f119d00c66..bd1def8d1b 100644 --- a/zig/tests/contract/metal_lfm2_dspark.zig +++ b/zig/tests/contract/metal_lfm2_dspark.zig @@ -63,7 +63,7 @@ pub const Probe = struct { ids: []u32 = &.{}, args: [8]rig_.Ref = undefined, y: rig_.Ref = undefined, - bad: rig_.Ref = undefined, // draft_ids: the flag word it sets for a flagged or invalid pick + bad: rig_.Ref = undefined, // draft_ids: the flag word it sets for a flagged or invalid pick, zeroed each run pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, w: Which) !Probe { const kernels: [2][:0]const u8 = switch (w) { @@ -116,7 +116,6 @@ pub const Probe = struct { const at = [4]u32{ stride, @intCast(s.rows), vocab, 0 }; p.args[1] = try p.rig.input(std.mem.asBytes(&at)); p.y = try p.rig.output(p.count()); - p.bad = try p.rig.output(4); return true; } if (s.n != vocab or s.k != rank or s.rows == 0) return false; @@ -169,6 +168,8 @@ pub const Probe = struct { 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.y.buf, p.y.off, 2); + p.bad = try p.rig.scratch(4); + @memset(p.rig.bytes(p.bad, 4), 0); e.setBuffer(p.bad.buf, p.bad.off, 3); e.dispatchThreads(mtl.Size.of(p.rows, 1, 1), mtl.Size.of(32, 1, 1)); return p.rig.finish(pass); From 9beda28081f9d8ded7185e0afb2b135303fde46b Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 19:09:04 +0100 Subject: [PATCH 36/46] lfm2 dspark: the bf16 tile kernel loads its tiles 16 bytes at a time --- zig/kernels/metal/lfm2/lfm2_dspark.metal | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/zig/kernels/metal/lfm2/lfm2_dspark.metal b/zig/kernels/metal/lfm2/lfm2_dspark.metal index 2842817430..656d0aa2a1 100644 --- a/zig/kernels/metal/lfm2/lfm2_dspark.metal +++ b/zig/kernels/metal/lfm2/lfm2_dspark.metal @@ -191,12 +191,12 @@ constant constexpr uint BF16_ROWS = 16; for (uint j = 0; j < 4; ++j) acc[i][j] = simdgroup_matrix(0.0f); const uint sr = (sg / 2) * 32, sc = (sg % 2) * 32; for (uint k0 = 0; k0 < p.k; k0 += 32) { - for (uint i = t; i < 64 * 4; i += 128) { // 64 rows of 32 bf16 as 4 chunks of 8 + for (uint i = t; i < 64 * 4; i += 128) { // 64 rows of 32 bf16 as 4 chunks of 8, 16 bytes a load const uint r = i / 4, q = (i % 4) * 8; - for (uint j = 0; j < 8; ++j) { - xs[r * 32 + q + j] = r0 + r < p.rows ? x[size_t(r0 + r) * p.k + k0 + q + j] : bfloat(0.0f); - ws[r * 32 + q + j] = c0 + r < p.n ? w[size_t(c0 + r) * p.k + k0 + q + j] : bfloat(0.0f); - } + const uint4 xv = r0 + r < p.rows ? *(const device uint4*)(x + size_t(r0 + r) * p.k + k0 + q) : uint4(0); + const uint4 wv = c0 + r < p.n ? *(const device uint4*)(w + size_t(c0 + r) * p.k + k0 + q) : uint4(0); + *(threadgroup uint4*)(xs + r * 32 + q) = xv; + *(threadgroup uint4*)(ws + r * 32 + q) = wv; } threadgroup_barrier(mem_flags::mem_threadgroup); for (uint kk = 0; kk < 32; kk += 8) { From 1cf5175aa06dd75b7d3db4c50368b72751705d3e Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 19:18:44 +0100 Subject: [PATCH 37/46] lfm2 dspark: Options' comment fits 120 columns --- zig/src/families/lfm2_moe/dspark.zig | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zig/src/families/lfm2_moe/dspark.zig b/zig/src/families/lfm2_moe/dspark.zig index b73098f591..ab3c3f3b33 100644 --- a/zig/src/families/lfm2_moe/dspark.zig +++ b/zig/src/families/lfm2_moe/dspark.zig @@ -29,7 +29,7 @@ pub const Verify = enum { depth, fixed, scheduled }; /// The verify when none is asked: the confidence-scheduled one. pub const default_verify: Verify = .scheduled; -/// `bits`: 0 runs the released bf16 matrices as stored, 8 quantizes to 8-bit codes; `verify_asked`: --draft-verify given. +/// `bits`: 0 runs the bf16 matrices as stored, 8 quantizes them at load; `verify_asked`: --draft-verify given. pub const Options = struct { capacity: usize, bits: u8 = 0, verify: Verify = default_verify, verify_asked: bool = false }; /// An asked scheduled verify needs the checkpoint's confidence head; the default one falls back to the depth rule. From 1987a20090c4f95dad57ebdd71e47efc7529fa0f Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 19:20:32 +0100 Subject: [PATCH 38/46] lfm2 dspark: --drafter-bits 8 serves the 8-bit codes, named in --help, the README, the recipe and the startup log --drafter-bits takes 0, 4 and 8. LFM2's drafter runs bf16 as released when it is not given or 0, 8-bit codes quantized at load on 8, and refuses 4. The other families refuse 8 by name. --- README.md | 2 +- docs/recipes/lfm2.5-24b-a2b.md | 6 ++++-- zig/src/core/engine_api.zig | 4 ++-- zig/src/native/cuda_options.zig | 4 ++++ zig/src/native/lfm2_moe.zig | 12 ++++++++---- zig/src/server/cli.zig | 2 +- zig/src/server/cli_test.zig | 4 +++- zig/src/server/main.zig | 2 +- 8 files changed, 24 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index f013d4160e..049f2fb09f 100644 --- a/README.md +++ b/README.md @@ -287,7 +287,7 @@ The binary's `capabilities --json` response lists its supported flags and platfo | `--loop-guard` | Close reasoning that repeats itself verbatim (a short cycle, or a long passage) where the engine supports it. | | `--no-drafts` | Produce the plain reference through the same engine. | | `--drafter DRAFTER` | An external drafter, as a directory or a repo id in the Hugging Face cache, for models that take one: Qwen3.8-27B uses DFlash2 and Gemma 4 31B the drafter named in its recipe. Without one they decode plain. | -| `--drafter-bits 0\|4` | Drafter weight precision where the model offers a choice: 4 bits (the default), or 0 to keep bf16. Qwen3.8-27B on CUDA takes 4 only. | +| `--drafter-bits 0\|4\|8` | Drafter weight precision where the model offers a choice: 4 bits (the default), or 0 to keep bf16. Qwen3.8-27B on CUDA takes 4 only. LFM2.5-24B-A2B's DSpark drafter runs bf16 as released (0, the default when not given) or 8-bit codes quantized at load (8); it has no 4-bit mode. | | `--record DIR` | Write each stream's committed tokens, positions and tapped rows under DIR, for drafter training. CUDA Nemotron only; a single stream runs through the lane rounds, and recording stops below 25 GiB free disk. | | `--keep-warm SECONDS` | Keep the Metal GPU active while idle for this long after a request, default 900; zero disables it. | | `--cfg-gamma GAMMA` | Metal Drama speech: semantic guidance strength, finite 1–5; 0/default disables it. Only reference-conditioned requests are guided. Enabled replies report `X-Cfg-Gamma`; codec heads stay conditional. | diff --git a/docs/recipes/lfm2.5-24b-a2b.md b/docs/recipes/lfm2.5-24b-a2b.md index f4272b59a5..194f5c9fc2 100644 --- a/docs/recipes/lfm2.5-24b-a2b.md +++ b/docs/recipes/lfm2.5-24b-a2b.md @@ -69,8 +69,10 @@ The drafter has no embedding and no head of its own: drafter was trained on. - Markov head: row i's logits plus `W2 W1[previous token]` (rank 256); draft i is their argmax. -Its bf16 linears are quantized at load to 8-bit MLX affine codes in groups of 64 and run on the family's kernels -and `core/tile_qmm`; the q and k rows are reordered so that the target's half-split RoPE kernel gives the +The drafter runs the released bf16 matrices as they are (`--drafter-bits 0`, the default). `--drafter-bits 8` instead +quantizes them at load to 8-bit MLX affine codes in groups of 64, about 5% faster on one stream and 361 MiB of weights +against 625 MiB; the server logs which mode runs, and 4 bits is refused. Both run on the family's kernels +(bf16 on `lfm2_bf16_rows`, `lfm2_bf16_mma` and `lfm2_markov_scan_bf16`, 8-bit on `core/tile_qmm`); the q and k rows are reordered so that the target's half-split RoPE kernel gives the checkpoint's interleaved RoPE. Its own kernels are `lfm2_tap`, `lfm2_attn_block` and `lfm2_attn_block_merge` (non-causal block attention, the target's attention kernels with a template flag), `lfm2_markov_scan` and `lfm2_markov_pick` (a Markov step) and `lfm2_draft_ids` (a hold's drafts into the next verify's rows on the GPU, so diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index 4838c4524e..2b5f908856 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -274,10 +274,10 @@ pub const Open = struct { /// --parallel named a number: an engine whose memory fits fewer streams refuses instead of serving fewer. lanes_fixed: bool = false, drafts: bool = true, - /// --drafter: a draft model's directory (Qwen3.8-27B's DFlash2; null: none), and --drafter-bits: 4, or 0 for bf16. + /// --drafter: a draft model's directory (null: none); --drafter-bits: 4, 0 for bf16, or 8 (LFM2 only). drafter: ?[]const u8 = null, drafter_bits: u8 = 4, - /// Whether --drafter-bits was given: LFM2's drafter runs bf16 unless it is refused a width it has no mode for. + /// Whether --drafter-bits was given: LFM2's drafter runs bf16 as released unless it is asked for 8-bit codes. drafter_bits_given: bool = false, /// --record: each stream's committed rows and tapped features written under this folder, for drafter training. record: ?[]const u8 = null, diff --git a/zig/src/native/cuda_options.zig b/zig/src/native/cuda_options.zig index d6a6eefe28..0c07152bdb 100644 --- a/zig/src/native/cuda_options.zig +++ b/zig/src/native/cuda_options.zig @@ -36,6 +36,10 @@ pub fn set(comptime F: type, o: api.Open, opts: *F.Options, problem: *[]const u8 return false; }; } + if (o.drafter_bits_given and o.drafter_bits == 8) { + problem.* = "--drafter-bits 8: only LFM2.5-24B-A2B's Metal DSpark drafter takes it; this family takes 0 (bf16) or 4"; + return false; + } if (comptime @hasField(F.Options, "drafter")) opts.drafter = o.drafter; if (comptime @hasField(F.Options, "drafter_bits")) opts.drafter_bits = o.drafter_bits; if (@hasField(F.Options, "kv8")) opts.kv8 = o.kv8; diff --git a/zig/src/native/lfm2_moe.zig b/zig/src/native/lfm2_moe.zig index e12e0a189a..a946d4bb22 100644 --- a/zig/src/native/lfm2_moe.zig +++ b/zig/src/native/lfm2_moe.zig @@ -82,12 +82,13 @@ pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Ope h.draft = null; h.wrap = null; if (if (o.drafts) o.drafter else null) |dir| { - if (o.drafter_bits_given and o.drafter_bits != 0) { - problem.* = "--drafter-bits: the LFM2 DSpark drafter runs the released bf16 matrices; it takes 0 and no other width"; + const bits: u8 = if (o.drafter_bits_given) o.drafter_bits else 0; + if (bits != 0 and bits != 8) { + problem.* = "--drafter-bits: the LFM2 DSpark drafter takes 0 (bf16 as released, the default) or 8 (8-bit codes quantized at load)"; release(h); return null; } - h.draft = q.dspark.Drafter.init(gpa, io, model, dir, .{ .capacity = window, .bits = 0 }) catch |err| { + h.draft = q.dspark.Drafter.init(gpa, io, model, dir, .{ .capacity = window, .bits = bits }) catch |err| { problem.* = try std.fmt.allocPrint(a, "the native LFM2 drafter cannot load {s} ({s}): it takes LiquidAI/LFM2.5-24B-A2B-DSpark", .{ dir, @errorName(err) }); release(h); return null; @@ -150,7 +151,10 @@ pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Ope }; h.host.keepalive_target = .{ .ctx = &h.warm, .tick = mtl.keepalive.Target.tick }; try h.host.start(); - if (h.draft) |d| std.log.info("lfm2: DSpark drafts from --drafter, {d} a round at most, bf16 matrices as released", .{d.config.block}); + if (h.draft) |d| { + const mode = if (d.options.bits == 8) "8-bit codes quantized at load (--drafter-bits 8), not the released bf16" else "bf16 matrices as released"; + std.log.info("lfm2: DSpark drafts from --drafter, {d} a round at most, {s}", .{ d.config.block, mode }); + } return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; } diff --git a/zig/src/server/cli.zig b/zig/src/server/cli.zig index 7f5909ab2b..2eaf3739a4 100644 --- a/zig/src/server/cli.zig +++ b/zig/src/server/cli.zig @@ -100,7 +100,7 @@ pub const flags = [_]Flag{ .{ .name = "--slide", .kind = .store_true, .native = true }, .{ .name = "--slide-graph", .native = true }, .{ .name = "--drafter", .native = true }, - .{ .name = "--drafter-bits", .choices = &.{ "0", "4" }, .native = true }, + .{ .name = "--drafter-bits", .choices = &.{ "0", "4", "8" }, .native = true }, .{ .name = "--record", .native = true }, .{ .name = "--kv-cache", .choices = &.{ "bf16", "fp8" }, .native = true }, .{ .name = "--prompt-precision", .choices = &.{ "bf16", "int8" }, .native = true }, diff --git a/zig/src/server/cli_test.zig b/zig/src/server/cli_test.zig index a2ea3c4db2..f35a44e631 100644 --- a/zig/src/server/cli_test.zig +++ b/zig/src/server/cli_test.zig @@ -164,7 +164,9 @@ test "parse and capabilities share the table" { try std.testing.expect(!args.thinking); try std.testing.expectEqualStrings("x", (try parse(a, &.{ "m", "--drafter", "x" }, &u)).drafter.?); try std.testing.expectEqual(@as(?u8, 0), (try parse(a, &.{ "m", "--drafter-bits", "0" }, &u)).drafter_bits); - try std.testing.expectError(error.Usage, parse(a, &.{ "m", "--drafter-bits", "8" }, &u)); + try std.testing.expectEqual(@as(?u8, 8), (try parse(a, &.{ "m", "--drafter-bits", "8" }, &u)).drafter_bits); + try std.testing.expectEqual(@as(?u8, null), (try parse(a, &.{ "m", "--drafter", "x" }, &u)).drafter_bits); + try std.testing.expectError(error.Usage, parse(a, &.{ "m", "--drafter-bits", "3" }, &u)); try std.testing.expectEqualStrings("/r", (try parse(a, &.{ "m", "--record", "/r" }, &u)).record.?); try std.testing.expectEqualStrings("fp8", (try parse(a, &.{ "m", "--kv-cache", "fp8" }, &u)).kv_cache); try std.testing.expectError(error.Usage, parse(a, &.{ "m", "--kv-cache", "int4" }, &u)); diff --git a/zig/src/server/main.zig b/zig/src/server/main.zig index 5dab0c4568..12fbfd93f2 100644 --- a/zig/src/server/main.zig +++ b/zig/src/server/main.zig @@ -25,7 +25,7 @@ const usage_line = "usage: tensorfold serve [-h] [--host HOST] [--port PORT] [-- " [--thinking | --no-thinking] [--reasoning-effort {low,medium,high,xhigh}]" ++ " [--thinking-budget THINKING_BUDGET] [--thinking-reserve THINKING_RESERVE] [--loop-guard] [--no-drafts]" ++ " [--drafter DRAFTER]" ++ - " [--drafter-bits {0,4}] [--record RECORD] [--keep-warm SECONDS] [--idle-unload SECONDS]" ++ + " [--drafter-bits {0,4,8}] [--record RECORD] [--keep-warm SECONDS] [--idle-unload SECONDS]" ++ " [--silence-floor-dbfs DBFS] [--silence-seconds SECONDS]" ++ " [--silence-onset-s SECONDS] [--silence-onset-voiced-s SECONDS]" ++ " [--silence-mid-s SECONDS] [--silence-mid-min-voiced-s SECONDS] [--silence-redraws N]" ++ From 4aa86a42e0ee6ce4461697933f8619ad0f9702dd Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 19:21:22 +0100 Subject: [PATCH 39/46] lfm2 dspark: the weights loader's comment fits 120 columns --- zig/src/families/lfm2_moe/dspark_weights.zig | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zig/src/families/lfm2_moe/dspark_weights.zig b/zig/src/families/lfm2_moe/dspark_weights.zig index c97a647b8b..9d87dd6eed 100644 --- a/zig/src/families/lfm2_moe/dspark_weights.zig +++ b/zig/src/families/lfm2_moe/dspark_weights.zig @@ -128,7 +128,7 @@ fn matrices(_: []const u8, rank: usize) bool { return rank == 2; } -/// The drafter's checkpoint read as released (`bits` 0) or quantized to 8-bit codes, a matrix at a time; buffers freed after. +/// The checkpoint read as released (`bits` 0) or quantized to 8-bit codes, a matrix at a time; buffers freed after. pub fn load(gpa: std.mem.Allocator, device: mtl.Device, io: std.Io, dir: []const u8, cfg: *const dc.Config, bits: u8) !Weights { if (bits != 0 and bits != 8) return error.UnsupportedDraftBits; var check = ckpt.Checkpoint.init(gpa); From 9c94ead07e6618e1d09382ccdd094b94e0263a35 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 19:33:11 +0100 Subject: [PATCH 40/46] kernel contract: the drafter's bf16 linears and bf16 Markov scan carry gates 1-4, recorded M1/M2 limits and six planted controls The few-rows and 64x64 tile kernels and the scan over a bf16 W2 pass gates 1 to 4. Six controls edit the source: a lane skipping half of k, a group's last row, an untransposed W tile, a lost last column, a NaN score not flagged and every stream reading the first stream's previous token. Each goes red at its gate. The tile kernel compiles for 832 threads on M1 and M2. --- docs/kernel-contract.md | 1 + zig/src/core/contract/lfm2_records.zig | 3 + zig/tests/contract/metal_lfm2_bf16.zig | 133 +++++++++++++++++++++ zig/tests/contract/metal_lfm2_controls.zig | 20 +++- zig/tests/contract/metal_lfm2_dspark.zig | 26 ++-- zig/tests/contract/metal_main.zig | 24 ++++ 6 files changed, 197 insertions(+), 10 deletions(-) create mode 100644 zig/tests/contract/metal_lfm2_bf16.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 7c95665c69..03a85c50c7 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -135,6 +135,7 @@ LFM2.5-24B-A2B's shared rounds: the greedy pick (`lfm2_argmax`) and the conv and 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. The DSpark drafter's tap copy, Markov scan and pick, draft ids and block attention are there too, with six more controls. +Its bf16 linears (few rows, 64x64 tiles) and the Markov scan over a bf16 W2 have four and two more. Their launches are in `lfm2_records.zig` with the M1 and M2 limits from `--offline`; `lfm2_attn_segs` and `lfm2_attn_block` reserve their 128 threads in their source, as Metal cannot build them 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 diff --git a/zig/src/core/contract/lfm2_records.zig b/zig/src/core/contract/lfm2_records.zig index c138251702..2d0add5d11 100644 --- a/zig/src/core/contract/lfm2_records.zig +++ b/zig/src/core/contract/lfm2_records.zig @@ -9,6 +9,9 @@ pub const records = [_]M{ .{ .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 }, .{ .name = "metal lfm2 markov scan lfm2_markov_scan", .threads = 256, .tg_memory = 1088, .m1 = 1024, .m2 = 1024 }, + .{ .name = "metal lfm2 markov scan bf16 lfm2_markov_scan_bf16", .threads = 256, .tg_memory = 1088, .m1 = 1024, .m2 = 1024 }, + .{ .name = "metal lfm2 bf16 rows lfm2_bf16_rows", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "metal lfm2 bf16 tile lfm2_bf16_mma", .threads = 128, .tg_memory = 24576, .m1 = 832, .m2 = 832 }, .{ .name = "metal lfm2 markov pick lfm2_markov_pick", .threads = 256, .tg_memory = 3072, .m1 = 1024, .m2 = 1024 }, .{ .name = "metal lfm2 draft ids lfm2_draft_ids", .threads = 32, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, .{ .name = "metal lfm2 block attention lfm2_attn_block", .threads = 128, .tg_memory = 18688, .reserved = 128 }, diff --git a/zig/tests/contract/metal_lfm2_bf16.zig b/zig/tests/contract/metal_lfm2_bf16.zig new file mode 100644 index 0000000000..0a191a48fa --- /dev/null +++ b/zig/tests/contract/metal_lfm2_bf16.zig @@ -0,0 +1,133 @@ +//! LFM2's DSpark bf16 linears under gates 1-4: x W^T with bf16 inputs and fp32 sums, few rows or 64x64 tiles. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const controls = @import("metal_lfm2_controls.zig"); + +const c = contract.spec; +const g = contract.gates; +const Dense = contract.dense.Dense; + +pub const Which = enum { rows, tile }; +const names = [_][]const u8{ "metal lfm2 bf16 rows", "metal lfm2 bf16 tile" }; + +/// The drafter's matrices: q and o (2,048 from 2,048), k and v (512), gate and up (6,144), down (2,048 from 6,144). +const rows_served = [_]c.Shape{ + .{ .rows = 1, .n = 2048, .k = 2048 }, .{ .rows = 9, .n = 512, .k = 2048 }, .{ .rows = 16, .n = 6144, .k = 2048 }, + .{ .rows = 32, .n = 2048, .k = 6144 }, .{ .rows = 4, .n = 2048, .k = 10240 }, +}; +/// A cold prompt's absorb: the taps' projection (2,048 from 10,240) and k and v, up to 2,048 rows a pass. +const tile_served = [_]c.Shape{ + .{ .rows = 33, .n = 2048, .k = 2048 }, .{ .rows = 128, .n = 512, .k = 2048 }, + .{ .rows = 2048, .n = 2048, .k = 10240 }, .{ .rows = 70, .n = 6144, .k = 2048 }, +}; + +/// A lane's k / 32 fp32 fmas in turn, then the simdgroup's five-level sum. +fn rowsDepth(s: c.Shape) usize { + return s.k / 32 + 8; +} + +/// One 8-wide MMA step a k step of 8, accumulating in fp32. +fn tileDepth(s: c.Shape) usize { + return s.k / 8 + 8; +} + +pub fn spec(w: Which) c.Spec { + return switch (w) { + // rows in groups of 16, a column a simdgroup (any n), k in 256-wide lane steps of 8 + .rows => .{ .name = names[0], .accum = .f32, .depth = rowsDepth, .out = .bf16, .tiles = .{ .{ .size = 16 }, .{ .size = 8 }, .{ .size = 256, .step = 8, .min = 8 } }, .served = &rows_served }, + .tile => .{ .name = names[1], .accum = .f32, .depth = tileDepth, .out = .bf16, .tiles = .{ .{ .size = 64 }, .{ .size = 64 }, .{ .size = 32, .step = 32, .min = 32 } }, .served = &tile_served }, + }; +} + +const Args = extern struct { rows: u32, n: u32, k: u32 }; + +pub const Probe = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + which: Which, + mutation: controls.Mutation, + pipe: mtl.Pipeline, + d: ?Dense = null, + x: rig_.Ref = undefined, + w: rig_.Ref = undefined, + out: rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, w: Which, mutation: controls.Mutation) !Probe { + const text = try controls.source(a, mutation); + defer if (mutation != .none) a.free(text); + const key: []const u8 = if (w == .rows) "lfm2_bf16_rows" else "lfm2_bf16_mma"; + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .which = w, .mutation = mutation, .pipe = try rig_.pipeline(device, text, key) }; + } + + pub fn deinit(p: *Probe) void { + if (p.d) |d| d.deinit(); + p.rig.deinit(); + p.pipe.deinit(); + } + + const vt = g.hooks(Probe); + pub fn kernel(p: *Probe) g.Kernel { + var s = spec(p.which); + if (p.mutation != .none) s.name = p.mutation.name(); + return .{ .spec = s, .ctx = p, .vt = &vt, .control = p.mutation.gate() }; + } + + pub fn prepare(p: *Probe, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.k == 0) return false; + if (p.which == .rows and (shape.k % 8 != 0 or shape.rows > 32)) return false; + if (p.which == .tile and shape.k % 32 != 0) return false; + if (p.d) |d| d.deinit(); + p.d = null; + p.rig.reset(); + const d = try Dense.init(p.a, shape, 1); + p.d = d; + d.fillX(fill, seed, 2); + // W as bf16 bits in [-1, 1), or gate 4's mirrored integers whose halves cancel (partials under 2^24) + const bits = try p.a.alloc(u16, shape.n * shape.k); + defer p.a.free(bits); + var prng = std.Random.DefaultPrng.init(seed ^ 0x6d); + const rng = prng.random(); + const k = shape.k; + for (bits) |*b| b.* = c.toBf16(if (fill == .stress) 0 else rng.float(f32) * 2 - 1); + if (fill == .stress) for (0..shape.n) |j| for (0..k / 2) |i| { + const v: f32 = @floatFromInt(1 + rng.uintLessThan(u32, 7)); + bits[j * k + i] = c.toBf16(v); + bits[j * k + k / 2 + i] = c.toBf16(-v); + }; + for (d.w, bits) |*wv, b| wv.* = c.fromBf16(b); + p.x = try p.rig.input(std.mem.sliceAsBytes(d.x)); + p.w = try p.rig.input(std.mem.sliceAsBytes(bits)); + p.out = try p.rig.output(d.outputs() * c.size(spec(p.which).out)); + return true; + } + + pub fn run(p: *Probe) !void { + const d = p.d.?; + const pass = p.rig.begin(); + const e = pass.e; + e.setPipeline(p.pipe); + e.setBuffer(p.x.buf, p.x.off, 0); + e.setBuffer(p.w.buf, p.w.off, 1); + const args = Args{ .rows = @intCast(d.shape.rows), .n = @intCast(d.shape.n), .k = @intCast(d.shape.k) }; + e.setBytes(std.mem.asBytes(&args), 2); + e.setBuffer(p.out.buf, p.out.off, 3); + if (p.which == .rows) { + e.dispatchGroups(mtl.Size.of((d.shape.n + 7) / 8, (d.shape.rows + 15) / 16, 1), mtl.Size.of(256, 1, 1)); + } else e.dispatchGroups(mtl.Size.of((d.shape.n + 63) / 64, (d.shape.rows + 63) / 64, 1), mtl.Size.of(128, 1, 1)); + try p.rig.finish(pass); + } + + pub fn count(p: *Probe) usize { + return p.d.?.outputs(); + } + + pub fn reference(p: *Probe, want: []f64, mag: []f64) !void { + return p.d.?.reference(want, mag); + } + + pub fn simulate(p: *Probe, acc: c.Accum, order: g.Order, out: []f64) !void { + return p.d.?.simulate(acc, order, out); + } +}; diff --git a/zig/tests/contract/metal_lfm2_controls.zig b/zig/tests/contract/metal_lfm2_controls.zig index 2a36ae0615..d7a1f7cd3b 100644 --- a/zig/tests/contract/metal_lfm2_controls.zig +++ b/zig/tests/contract/metal_lfm2_controls.zig @@ -15,6 +15,12 @@ pub const Mutation = enum { markov_first_prev, draft_ids_unzeroed, block_attention_causal, + markov_bf16_nan_dropped, + markov_bf16_first_prev, + bf16_rows_short_k, + bf16_rows_last_row, + bf16_tile_untransposed, + bf16_tile_last_column, pub fn name(m: Mutation) []const u8 { return switch (m) { @@ -29,6 +35,12 @@ pub const Mutation = enum { .markov_first_prev => "control: metal lfm2 markov step, every stream reads the first stream's previous token", .draft_ids_unzeroed => "control: metal lfm2 draft ids, a pick past the vocabulary keeps its token", .block_attention_causal => "control: metal lfm2 block attention, a row misses the last key", + .markov_bf16_nan_dropped => "control: metal lfm2 markov step bf16, a NaN score is not flagged", + .markov_bf16_first_prev => "control: metal lfm2 markov step bf16, every stream reads the first stream's previous token", + .bf16_rows_short_k => "control: metal lfm2 bf16 rows, a lane skips every second step of k", + .bf16_rows_last_row => "control: metal lfm2 bf16 rows, a group's last row is not computed", + .bf16_tile_untransposed => "control: metal lfm2 bf16 tile, W's tile is read untransposed", + .bf16_tile_last_column => "control: metal lfm2 bf16 tile, the last column is not written", }; } @@ -36,7 +48,7 @@ pub const Mutation = enum { pub fn gate(m: Mutation) ?contract.gates.Gate { return switch (m) { .none => null, - .argmax_no_flag, .markov_nan_dropped => .nonfinite, + .argmax_no_flag, .markov_nan_dropped, .markov_bf16_nan_dropped => .nonfinite, else => .shapes, }; } @@ -55,6 +67,12 @@ pub const Mutation = enum { .markov_first_prev => .{ .from = "id = prev[size_t(s) * p.prev_stride];", .to = "id = prev[0];" }, .draft_ids_unzeroed => .{ .from = "ids[i] = ok ? token : 0u;", .to = "ids[i] = token;" }, .block_attention_causal => .{ .from = "count[b] = FULL ? p.len + 1 :", .to = "count[b] = FULL ? p.len :" }, + .markov_bf16_nan_dropped => .{ .from = "nan = nan || isnan(sc);", .to = "nan = false;" }, + .markov_bf16_first_prev => .{ .from = "id = prev[size_t(s) * p.prev_stride];", .to = "id = prev[0];" }, + .bf16_rows_short_k => .{ .from = "for (uint i = 8 * lane; i < p.k; i += 256) {", .to = "for (uint i = 8 * lane; i < p.k; i += 512) {" }, + .bf16_rows_last_row => .{ .from = "if (r >= rows) break;", .to = "if (r + 1 >= rows) break;" }, + .bf16_tile_untransposed => .{ .from = "32, ulong2(0, 0), true);", .to = "32, ulong2(0, 0), false);" }, + .bf16_tile_last_column => .{ .from = "if (row < p.rows && col < p.n) y[", .to = "if (row < p.rows && col + 1 < p.n) y[" }, .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_dspark.zig b/zig/tests/contract/metal_lfm2_dspark.zig index 180efd8c58..04b7fc489a 100644 --- a/zig/tests/contract/metal_lfm2_dspark.zig +++ b/zig/tests/contract/metal_lfm2_dspark.zig @@ -8,13 +8,13 @@ const g = contract.gates; const rig_ = @import("rig.zig"); const controls = @import("metal_lfm2_controls.zig"); -pub const Which = enum { markov_step, draft_ids }; +pub const Which = enum { markov_step, draft_ids, markov_step_bf16 }; const rank = 256; const vocab = 128000; const span = 512; // tokens a scan threadgroup takes (the kernel's MARKOV_SPAN) const groups = (vocab + span - 1) / span; const table = 1000; // W1's rows in the check: ids are inside it or past the vocabulary (a zero row) -const names = [_][]const u8{ "metal lfm2 markov step", "metal lfm2 draft ids" }; +const names = [_][]const u8{ "metal lfm2 markov step", "metal lfm2 draft ids", "metal lfm2 markov step bf16" }; const stride = 64; // words between a stream's picks: a hold pass's 32 blocks, a token and a flag each const ids_shapes = [_]c.Shape{ .{ .rows = 1, .n = 1, .k = 1 }, .{ .rows = 4, .n = 1, .k = 1 }, .{ .rows = 15, .n = 1, .k = 1 } }; @@ -31,7 +31,7 @@ fn depth(_: c.Shape) usize { pub fn spec(w: Which) c.Spec { return switch (w) { .draft_ids => .{ .name = names[1], .accum = .none, .depth = none, .out = .raw, .served = &ids_shapes, .sums = .{ .f16 = 0, .bf16 = 0, .why = "a copy moves tokens: no sums" } }, - .markov_step => .{ .name = names[0], .accum = .none, .depth = depth, .out = .raw, .served = &step_shapes, .nan = .{ .stated = "a row with a NaN score is flagged nonfinite" }, .sums = .{ .f16 = 0, .bf16 = 0, .why = "picks are tokens: the check bounds each pick's score by gamma(256), not sum bits" } }, + .markov_step, .markov_step_bf16 => .{ .name = if (w == .markov_step) names[0] else names[2], .accum = .none, .depth = depth, .out = .raw, .served = &step_shapes, .nan = .{ .stated = "a row with a NaN score is flagged nonfinite" }, .sums = .{ .f16 = 0, .bf16 = 0, .why = "picks are tokens: the check bounds each pick's score by gamma(256), not sum bits" } }, }; } @@ -62,6 +62,7 @@ pub const Probe = struct { codes: []u8 = &.{}, scales: []u16 = &.{}, biases: []u16 = &.{}, + w2: []u16 = &.{}, // the bf16 step's W2 ids: []u32 = &.{}, args: [8]rig_.Ref = undefined, y: rig_.Ref = undefined, @@ -72,6 +73,7 @@ pub const Probe = struct { defer if (mutation != .none) a.free(text); const kernels: [2][:0]const u8 = switch (w) { .markov_step => .{ "lfm2_markov_scan", "lfm2_markov_pick" }, + .markov_step_bf16 => .{ "lfm2_markov_scan_bf16", "lfm2_markov_pick" }, .draft_ids => .{ "lfm2_draft_ids", "lfm2_draft_ids" }, }; const first = try rig_.pipeline(device, text, kernels[0]); @@ -81,7 +83,7 @@ pub const Probe = struct { } fn drop(p: *Probe) void { - inline for (.{ "src", "w1", "codes", "scales", "biases", "ids" }) |f| { + inline for (.{ "src", "w1", "codes", "scales", "biases", "w2", "ids" }) |f| { p.a.free(@field(p, f)); @field(p, f) = &.{}; } @@ -138,6 +140,10 @@ pub const Probe = struct { sc.* = c.toBf16(@as(f32, @floatFromInt(i % 7 + 1)) / 4096); bi.* = c.toBf16(@as(f32, @floatFromInt(i % 5)) / 64 - 0.03); } + if (p.which == .markov_step_bf16) { + p.w2 = try p.a.alloc(u16, vocab * rank); + fill(p.w2, seed +% 9, .random); + } // the previous token every second word (the picks' layout): inside W1, its last row, or past the vocabulary p.ids = try p.a.alloc(u32, 2 * s.rows); for (p.ids, 0..) |*id, i| id.* = if (i % 2 == 1) 0xdead else switch ((i / 2) % 3) { @@ -147,7 +153,7 @@ pub const Probe = struct { }; p.args[0] = try p.rig.input(std.mem.sliceAsBytes(p.src)); p.args[1] = try p.rig.input(std.mem.sliceAsBytes(p.w1)); - p.args[2] = try p.rig.input(p.codes); + p.args[2] = if (p.which == .markov_step_bf16) try p.rig.input(std.mem.sliceAsBytes(p.w2)) else try p.rig.input(p.codes); p.args[3] = try p.rig.input(std.mem.sliceAsBytes(p.ids)); const args = [4]u32{ 2 * vocab, rank, vocab, 2 }; p.args[4] = try p.rig.input(std.mem.asBytes(&args)); @@ -161,7 +167,7 @@ pub const Probe = struct { pub fn count(p: *Probe) usize { return switch (p.which) { - .markov_step => p.rows * 8, + .markov_step, .markov_step_bf16 => p.rows * 8, .draft_ids => p.rows * 4, }; } @@ -186,8 +192,10 @@ pub const Probe = struct { for (p.args[0..5], 0..) |r, i| e.setBuffer(r.buf, r.off, i); e.setBuffer(best.buf, best.off, 5); e.setBuffer(at.buf, at.off, 6); - e.setBuffer(p.args[5].buf, p.args[5].off, 7); - e.setBuffer(p.args[6].buf, p.args[6].off, 8); + if (p.which == .markov_step) { + e.setBuffer(p.args[5].buf, p.args[5].off, 7); + e.setBuffer(p.args[6].buf, p.args[6].off, 8); + } e.dispatchGroups(mtl.Size.of(groups, p.rows, 1), mtl.Size.of(256, 1, 1)); e.setPipeline(p.pipes[1]); e.setBuffer(best.buf, best.off, 0); @@ -231,7 +239,7 @@ pub const Probe = struct { var size: f64 = 0; if (id < table) for (0..rank) |i| { const k = (v * rank + i) / 64; - const w = @as(f64, c.fromBf16(p.scales[k])) * @as(f64, @floatFromInt(p.codes[v * rank + i])) + c.fromBf16(p.biases[k]); + const w: f64 = if (p.which == .markov_step_bf16) c.fromBf16(p.w2[v * rank + i]) else @as(f64, c.fromBf16(p.scales[k])) * @as(f64, @floatFromInt(p.codes[v * rank + i])) + c.fromBf16(p.biases[k]); dot += w * c.fromBf16(p.w1[id * rank + i]); size += @abs(w * c.fromBf16(p.w1[id * rank + i])); }; diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 740d0e5672..d90aef8a24 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -27,6 +27,7 @@ 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_dspark = @import("metal_lfm2_dspark.zig"); +const lfm2_bf16 = @import("metal_lfm2_bf16.zig"); const lfm2_attn = @import("metal_lfm2_attn.zig"); const lfm2_controls = @import("metal_lfm2_controls.zig"); const g = contract.gates; @@ -211,6 +212,7 @@ pub fn main(init: std.process.Init) !u8 { } var drafter_list: std.ArrayList(*lfm2_dspark.Probe) = .empty; var block_list: std.ArrayList(*lfm2_attn.Probe) = .empty; + var bf16_list: std.ArrayList(*lfm2_bf16.Probe) = .empty; defer { for (drafter_list.items) |p| { p.deinit(); @@ -220,12 +222,18 @@ pub fn main(init: std.process.Init) !u8 { p.deinit(); a.destroy(p); } + for (bf16_list.items) |p| { + p.deinit(); + a.destroy(p); + } drafter_list.deinit(a); block_list.deinit(a); + bf16_list.deinit(a); } for (std.enums.values(lfm2_controls.Mutation)) |m| { const drafter_which: ?lfm2_dspark.Which = switch (m) { .markov_nan_dropped, .markov_first_prev => .markov_step, + .markov_bf16_nan_dropped, .markov_bf16_first_prev => .markov_step_bf16, .draft_ids_unzeroed => .draft_ids, else => null, }; @@ -240,6 +248,22 @@ pub fn main(init: std.process.Init) !u8 { try kernels.append(a, p.kernel()); } } + const bf16_which: ?lfm2_bf16.Which = switch (m) { + .bf16_rows_short_k, .bf16_rows_last_row => .rows, + .bf16_tile_untransposed, .bf16_tile_last_column => .tile, + else => null, + }; + if (m == .none or bf16_which != null) { + for (std.enums.values(lfm2_bf16.Which)) |w| { + if (m != .none and w != bf16_which.?) continue; + if (!wanted(o, if (m == .none) lfm2_bf16.spec(w).name else m.name())) continue; + const p = try a.create(lfm2_bf16.Probe); + errdefer a.destroy(p); + p.* = try lfm2_bf16.Probe.init(a, device, queue, w, m); + try bf16_list.append(a, p); + try kernels.append(a, p.kernel()); + } + } if (m == .none or m == .block_attention_causal) { if (!wanted(o, if (m == .none) lfm2_attn.name else m.name())) continue; const p = try a.create(lfm2_attn.Probe); 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 41/46] 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 18a23ed82e1c551d2675c670c47c13663a175616 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 19:39:12 +0100 Subject: [PATCH 42/46] lfm2: --drafter is refused by name on an M1 or M2 GPU, and an M3 is logged as untested with the drafter The check is its own, after the engine's: it uses the same generationServed helper, with a test beside it. The drafter kernels' records keep their M1 and M2 limits. --- zig/src/core/contract/lfm2_records.zig | 1 + zig/src/native/lfm2_moe.zig | 17 +++++++++++++++++ 2 files changed, 18 insertions(+) diff --git a/zig/src/core/contract/lfm2_records.zig b/zig/src/core/contract/lfm2_records.zig index 2d0add5d11..5227176040 100644 --- a/zig/src/core/contract/lfm2_records.zig +++ b/zig/src/core/contract/lfm2_records.zig @@ -1,6 +1,7 @@ //! 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; +// The drafter's kernels carry M1 and M2 limits too, but --drafter is refused on those chips (native/lfm2_moe.zig). 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 }, diff --git a/zig/src/native/lfm2_moe.zig b/zig/src/native/lfm2_moe.zig index e8faaab103..20f28a10e1 100644 --- a/zig/src/native/lfm2_moe.zig +++ b/zig/src/native/lfm2_moe.zig @@ -57,6 +57,17 @@ fn generationServed(generation: u32) bool { return generation == 0 or generation >= first_generation; } +/// Why --drafter is refused on a GPU of this generation, or null: it is served from the M3 on, as the engine is. +fn drafterRefusal(generation: u32) ?[]const u8 { + if (generationServed(generation)) return null; + return "--drafter: the LFM2 DSpark drafter needs an M3 or later GPU; M1 and M2 are not qualified for its kernels"; +} + +test "the drafter is refused on M1 and M2 by name, served from the M3 on" { + try std.testing.expect(drafterRefusal(13) != null and drafterRefusal(14) != null); + try std.testing.expect(drafterRefusal(15) == null and drafterRefusal(16) == null and drafterRefusal(0) == null); +} + 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)); @@ -100,6 +111,12 @@ pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Ope h.draft = null; h.wrap = null; if (if (o.drafts) o.drafter else null) |dir| { + if (drafterRefusal(model.device.generation())) |why| { + problem.* = why; + release(h); + return null; + } + if (model.device.generation() == first_generation) std.log.info("lfm2: an M3 is not yet tested with the DSpark drafter", .{}); const bits: u8 = if (o.drafter_bits_given) o.drafter_bits else 0; if (bits != 0 and bits != 8) { problem.* = "--drafter-bits: the LFM2 DSpark drafter takes 0 (bf16 as released, the default) or 8 (8-bit codes quantized at load)"; From 6539b72beacf5bf911d5a5141f50da88cfcab4c6 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 19:53:24 +0100 Subject: [PATCH 43/46] lanes: a lone stream that drafting loses on keeps its loss for the next stream, and probes less often A loss is counted on the stretch's own rate. After two losses in a row the plain wait grows until a probe costs at most 3% of the plain time. The band keeps its last losing verdict, so the next stream in that band starts plain. Any winning stretch clears it. --- zig/src/core/lanes/depth.zig | 3 +- zig/src/core/lanes/lone_guard.zig | 101 +++++++++++++++++++++++++++--- 2 files changed, 96 insertions(+), 8 deletions(-) diff --git a/zig/src/core/lanes/depth.zig b/zig/src/core/lanes/depth.zig index e02c1e6240..9e2d3a2f7a 100644 --- a/zig/src/core/lanes/depth.zig +++ b/zig/src/core/lanes/depth.zig @@ -239,7 +239,8 @@ pub const Rule = struct { /// A lone round's committed tokens and wall ms at the stream's context band, for the lone guard. pub fn observeLone(r: *Rule, who: Who, band: usize, tokens: usize, ms: f64) !void { if (!r.lone_guard) return; - (try r.state(who)).lone.observe(r.lone_k, r.lone_plain.at(who.sampled, band), tokens, ms); + const memo = r.lone_plain.memo(who.sampled, band); + (try r.state(who)).lone.observeBand(r.lone_k, r.lone_plain.at(who.sampled, band), memo, tokens, ms); } /// Python `_head_depth`: the next round's head drafts for a stream (0 when forced, 1 or none with a copy ahead). diff --git a/zig/src/core/lanes/lone_guard.zig b/zig/src/core/lanes/lone_guard.zig index d861174dc3..4657ffe829 100644 --- a/zig/src/core/lanes/lone_guard.zig +++ b/zig/src/core/lanes/lone_guard.zig @@ -7,6 +7,8 @@ pub const Constants = struct { wait: u32 = 16, // plain rounds before drafting is measured again; doubles each time plain wins again wait_most: u32 = 512, margin: f64 = 1.05, // drafted rounds must beat plain ones' tokens a ms by this much + probe_cost: f64 = 0.03, // after two losses a probe costs at most this share of plain time + confirm: u32 = 2, // losing verdicts in a row (a remembered one counts) before the wait follows probe_cost blend: f64 = 0.5, // the newest stretch's weight in a rate }; @@ -15,6 +17,13 @@ pub const Mode = enum { drafted, plain }; /// Plain rounds' tokens a ms by sampling mode and context band, kept from stream to stream (no draft changes them). pub const Plain = struct { rates: [2][bands]?f64 = @splat(@splat(null)), + verdicts: [2][bands]?Verdict = @splat(@splat(null)), // each band's last losing verdict, until drafting wins there + + pub const Verdict = struct { drafted: f64, wait: u32 }; + + pub fn memo(p: *Plain, sampled: bool, band: usize) *?Verdict { + return &p.verdicts[@intFromBool(sampled)][@min(band, bands - 1)]; + } pub const bands = 6; @@ -34,6 +43,9 @@ pub const Lone = struct { pending: bool = false, // a drafted verdict waits for the plain rate this stretch takes fresh: bool = false, // a drafted stretch's first hold: one draft and its slack, whatever the last grant drafted: ?f64 = null, + drafted_ms: f64 = 0, // the last drafted stretch's wall ms + drafted_now: f64 = 0, // the last drafted stretch's own tokens a ms (unblended) + losses: u32 = 0, // losing verdicts in a row wait: u32 = 0, /// Whether the next hold may draft. @@ -43,13 +55,26 @@ pub const Lone = struct { /// A lone round of held drafts or none (copies left out): its tokens and wall ms; `plain` the band's rate. pub fn observe(l: *Lone, k: Constants, plain: *?f64, tokens: usize, ms: f64) void { + l.observeBand(k, plain, null, tokens, ms); + } + + /// observe, with the band's last verdict: a stream starting where drafting lost starts plain for its wait. + pub fn observeBand(l: *Lone, k: Constants, plain: *?f64, memo: ?*?Plain.Verdict, tokens: usize, ms: f64) void { + var switched = false; if (!l.started) { l.started = true; l.left = k.drafted; l.wait = k.wait; + if (memo) |m| if (m.*) |v| if (plain.*) |p| if (v.drafted < p * k.margin) { + l.drafted = v.drafted; + l.losses = k.confirm - 1; // one more losing probe confirms the band's loss again + l.switchTo(.plain, v.wait); + l.wait = @min(v.wait *| 2, k.wait_most); + switched = true; + }; } if (l.skip) { - l.skip = false; + l.skip = switched; // the next round still verifies this round's drafted hold return; } l.tokens += @floatFromInt(tokens); @@ -57,35 +82,54 @@ pub const Lone = struct { l.left -|= 1; if (l.left > 0) return; const now = if (l.ms > 0) l.tokens / l.ms else 0; + const stretch_ms = l.ms; l.tokens = 0; l.ms = 0; switch (l.mode) { .drafted => { l.drafted = blend(k, l.drafted, now); + l.drafted_ms = stretch_ms; + l.drafted_now = now; if (plain.* == null) { l.pending = true; l.switchTo(.plain, k.plain); - } else l.verdict(k, plain.*.?); + } else l.verdict(k, plain.*.?, memo); }, .plain => { plain.* = blend(k, plain.*, now); if (l.pending) { l.pending = false; - l.verdict(k, plain.*.?); + l.verdict(k, plain.*.?, memo); } else l.switchTo(.drafted, k.drafted); }, } } - /// Drafting goes on while its rate beats plain's by the margin, else plain rounds for a wait that doubles. - fn verdict(l: *Lone, k: Constants, plain: f64) void { + /// Drafting goes on while it beats plain by the margin, else plain rounds for a wait that doubles. + fn verdict(l: *Lone, k: Constants, plain: f64, memo: ?*?Plain.Verdict) void { if (l.drafted.? >= plain * k.margin) { l.wait = k.wait; + l.losses = 0; l.switchTo(.drafted, k.drafted); + if (memo) |m| m.* = null; return; } - l.switchTo(.plain, l.wait); - l.wait = @min(l.wait *| 2, k.wait_most); + var wait = l.wait; + if (l.drafted_now >= plain * k.margin) { + // the blend lost but this stretch won: no loss counted, the band forgets its loss + l.losses = 0; + if (memo) |m| m.* = null; + } else { + l.losses += 1; + if (l.losses >= k.confirm) { + const lost_ms = l.drafted_ms * @max(0, 1 - l.drafted_now / plain); + const need: f64 = @ceil(lost_ms * plain / k.probe_cost); // plain rounds of about one token each + wait = @min(k.wait_most, @max(wait, @as(u32, @intFromFloat(@min(need, 1e6))))); + } + if (memo) |m| m.* = .{ .drafted = l.drafted.?, .wait = wait }; + } + l.switchTo(.plain, wait); + l.wait = @min(wait *| 2, k.wait_most); } fn switchTo(l: *Lone, mode: Mode, rounds: u32) void { @@ -153,3 +197,46 @@ test "a drafted stretch's first round is not measured and its first hold is fres try std.testing.expect(l.drafts() and l.fresh); // the probe after the wait try std.testing.expectApproxEqAbs(@as(f64, 0.2), rate.?, 1e-12); } + +test "a probe that loses a third waits long enough that probes cost at most probe_cost; a band's verdict carries over" { + const k: Constants = .{}; + var rate: ?f64 = null; + var memo: ?Plain.Verdict = null; + var l: Lone = .{}; + // 130k on the M5: plain 8.67 ms a token, drafted 15.7 ms for 1.2 tokens + var drafted: usize = 0; + for (0..200) |_| { + const d = l.drafts(); + drafted += @intFromBool(d); + l.observeBand(k, &rate, &memo, if (d) 12 else 10, if (d) 157 else 86.7); + } + try std.testing.expect(drafted <= 2 * (1 + k.drafted)); // two losing stretches, then the probe_cost wait + try std.testing.expect(memo != null and memo.?.wait > 200); + var next: Lone = .{}; + var later: usize = 0; + for (0..200) |_| { + const d = next.drafts(); + later += @intFromBool(d); + next.observeBand(k, &rate, &memo, if (d) 12 else 10, if (d) 157 else 86.7); + } + try std.testing.expect(later <= 1); +} + +test "a winning stretch the blend still calls a loss counts no loss: no probe_cost wait" { + const k: Constants = .{ .drafted = 2, .wait = 2 }; + var rate: ?f64 = 0.2; + var memo: ?Plain.Verdict = null; + var l: Lone = .{}; + l.observeBand(k, &rate, &memo, 9, 1); // skipped + l.observeBand(k, &rate, &memo, 1, 10); + l.observeBand(k, &rate, &memo, 1, 10); // 0.1: a loss, wait 2 + try std.testing.expect(!l.drafts() and l.losses == 1 and memo != null); + l.observeBand(k, &rate, &memo, 9, 1); // skipped + l.observeBand(k, &rate, &memo, 1, 5); + l.observeBand(k, &rate, &memo, 1, 5); + try std.testing.expect(l.drafts()); + l.observeBand(k, &rate, &memo, 9, 1); // skipped + l.observeBand(k, &rate, &memo, 2, 8); + l.observeBand(k, &rate, &memo, 2, 8); // 0.25 alone wins; blended 0.175 loses + try std.testing.expect(!l.drafts() and l.losses == 0 and memo == null and l.left == 4); +} From caf9c601d882e19336c2f6e8acec5175468be293 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 20:21:37 +0100 Subject: [PATCH 44/46] metal.zig back to 600 lines: its one comment line goes, the --draft-verify refusal stays --- zig/src/native/metal.zig | 1 - 1 file changed, 1 deletion(-) diff --git a/zig/src/native/metal.zig b/zig/src/native/metal.zig index daa08eb95e..a5eda7c55c 100644 --- a/zig/src/native/metal.zig +++ b/zig/src/native/metal.zig @@ -337,7 +337,6 @@ fn refusal(a: Allocator, o: api.Open, dense: bool) !?[]const u8 { if ((o.draft_verify != null or o.verify_costs != null) and (o.drafter == null or !std.mem.eql(u8, t, "lfm2_moe"))) return "--draft-verify and --verify-costs set the verify of LFM2.5-24B-A2B's DSpark drafter: serve it with --drafter"; if (spill.refusal(o)) |why| return why; if (o.learn == null) return null; - // the 27B writes conversation states if (o.learn_conversations and std.mem.eql(u8, t, "qwen3_5") and dense) return null; const name: []const u8 = if (std.mem.eql(u8, t, "glm_moe_dsa")) "GLM-5.3 (EXLR8)" else if (std.mem.eql(u8, t, "lfm2_moe")) "LFM2.5-24B-A2B" else if (isGemma4(t)) "Gemma 4" else if (std.mem.eql(u8, t, "step5_vl")) "Step-5" else if (std.mem.eql(u8, t, "spark2_5")) "Spark-X2.6" else if (!std.mem.eql(u8, t, "qwen3_5")) return null else if (dense) "Qwen3.8-27B" else "Qwen3.5-2B"; if (o.learn_conversations) return try std.fmt.allocPrint(a, "{s} keeps conversation states on disk for Qwen3.8-27B, Nemotron, Flash Next and GLM-5.3-Flash: serve {s} without it", .{ spill.flag(o), name }); 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 45/46] 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; From 1bb6e5db28e9bf8434b7dc33c80def5b7b54af8a Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 20:55:29 +0100 Subject: [PATCH 46/46] lfm2: the drafter's verify is not a choice: the confidence head's chances, else the depth rule --draft-verify and --verify-costs are gone with the fixed mode, the cost-table loader and tf-lfm2-verify-costs. A checkpoint with the confidence head verifies by its chances. One without it verifies by the depth rule. --- docs/recipes/lfm2.5-24b-a2b.md | 21 ++---- zig/build/metal_targets.zig | 1 - zig/src/core/engine_api.zig | 6 -- zig/src/families/lfm2_moe/dspark.zig | 37 ++------- zig/src/families/lfm2_moe/dspark_test.zig | 14 ---- zig/src/families/lfm2_moe/lfm2_moe.zig | 1 - zig/src/families/lfm2_moe/verify_costs.zig | 44 ----------- zig/src/native/lfm2_moe.zig | 28 +------ zig/src/native/metal.zig | 1 - zig/src/server/cli.zig | 8 +- zig/src/server/cli_test.zig | 7 -- zig/src/server/main.zig | 2 +- zig/src/tensorfold.zig | 1 - zig/tests/lfm2_dspark.zig | 2 +- zig/tests/lfm2_verify_costs.zig | 87 ---------------------- 15 files changed, 18 insertions(+), 242 deletions(-) delete mode 100644 zig/src/families/lfm2_moe/verify_costs.zig delete mode 100644 zig/tests/lfm2_verify_costs.zig diff --git a/docs/recipes/lfm2.5-24b-a2b.md b/docs/recipes/lfm2.5-24b-a2b.md index e4ff65165a..48a4065c34 100644 --- a/docs/recipes/lfm2.5-24b-a2b.md +++ b/docs/recipes/lfm2.5-24b-a2b.md @@ -89,24 +89,17 @@ the drafter's passes; a near tie can pick the other token. The drafter also carries a confidence head: a linear map of a draft's final block row and `W1[previous token]` (2,304 inputs) to one logit, through a sigmoid (`lfm2_dspark_confidence`). It does not change the drafts. -`--draft-verify` picks how they are verified: +The verify rows follow it: a chain's chance at draft j is the product of its first j confidences. Each round the +lane core's window allocator takes the likeliest drafts over all streams while the expected tokens a ms rise, so +each stream gets its own verify length, and the next hold drafts as deep as the stream's grant. A hold still +running on the GPU is priced by its lane's last confidences, so the host does not wait for it. A checkpoint +without the head verifies by the core's depth rule, from each stream's measured acceptance by depth. -- `scheduled` (the default): a chain's chance at draft j is the product of its first j confidences. Each round the - lane core's window allocator takes the likeliest drafts over all streams while the expected tokens a ms rise, so - each stream gets its own verify length, and the next hold drafts as deep as the stream's grant. A hold still - running on the GPU is priced by its lane's last confidences, so the host does not wait for it. -- `depth`: the lane core's depth rule, from each stream's measured acceptance by depth. -- `fixed`: every round verifies the whole block, as LiquidAI's published benchmarks did; a shared round still fits - 32 rows. - -With `scheduled` a stream with a copy ahead verifies the copy and asks no hold, and a drafting stream prices its copy +A stream with a copy ahead verifies the copy and asks no hold, and a drafting stream prices its copy windows by its own copies. The lane core also weighs the hold pass itself: the drafter times holds of 1 to 16 blocks at load and corrects those prices toward each settled hold's GPU time. A lone stream holds drafts only while its drafted rounds measure faster than plain ones (`zig/src/core/lanes/lone_guard.zig`); shared rounds keep the shared-round guard. -`--verify-costs FILE` replaces the costs timed at load with a table from `tf-lfm2-verify-costs MODEL FILE [HISTORY -[IDS...]]` (lines `streams rows_each total_rows ms`), measured on the serving Mac. With token id files each stream -reads its history and verify rows from its own part of the ids, so its rows route its own text's experts. Both flags -are refused unless an `lfm2_moe` checkpoint is served with `--drafter`. `tf-lfm2-dspark` also prints the confidence +`tf-lfm2-dspark` also prints the confidence head against its fp32 reference and times shared holds of 2 to 16 streams. diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index 9896a764f2..727daf7e2e 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -293,7 +293,6 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s .{ .name = "tf-lfm2-bench", .path = "zig/tests/lfm2_bench.zig", .about = "LFM2-MoE prompt-chunk and one-row decode timing" }, .{ .name = "tf-lfm2-exact", .path = "zig/tests/lfm2_exact.zig", .about = "LFM2-MoE prompt chunks, windows and committed-state exactness" }, .{ .name = "tf-lfm2-dspark", .path = "zig/tests/lfm2_dspark.zig", .about = "LFM2's DSpark drafter on the GPU against an fp32 host reference of the released drafter" }, - .{ .name = "tf-lfm2-verify-costs", .path = "zig/tests/lfm2_verify_costs.zig", .about = "LFM2 verify forward ms by stream count and rows a stream, written as a cost table" }, .{ .name = "tf-lfm2-fp64", .path = "zig/tests/lfm2_fp64.zig", .about = "LFM2 projection classes and attention on the tensor units, the tiles and the twin, each against float64" }, .{ .name = "tensorfold", .path = "zig/src/main.zig", .about = "The native engine's command line" }, .{ .name = "tf-nemotron-fixtures", .path = "zig/tests/nemotron_fixtures.zig", .about = "Nemotron kernels against the Python engine's captured ops" }, diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index 96c3781b25..7191dabbff 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -270,9 +270,6 @@ pub const Family = struct { model_type: []const u8, formats: []const []const u8, /// The math an NVFP4 checkpoint runs: its own (FP4 rows into NVFP4 layers, FP8 into FP8 ones), or bf16 rows throughout. pub const Precision = enum { checkpoint, full }; -/// An external drafter's verify: the depth rule's rows, its whole block, or its confidence head's survivals. -pub const DraftVerify = enum { depth, fixed, scheduled }; - /// What a server asks of the engine it opens: the checkpoint, and the serve flags an engine reads. pub const Open = struct { dir: []const u8, @@ -291,9 +288,6 @@ pub const Open = struct { /// --drafter: a draft model's directory (null: none); --drafter-bits: 4, 0 for bf16, or 8 (LFM2 only). drafter: ?[]const u8 = null, drafter_bits: u8 = 4, - /// --draft-verify and --verify-costs (a tf-lfm2-verify-costs table): the LFM2 DSpark drafter's verify rows. - draft_verify: ?DraftVerify = null, - verify_costs: ?[]const u8 = null, /// Whether --drafter-bits was given: LFM2's drafter runs bf16 as released unless it is asked for 8-bit codes. drafter_bits_given: bool = false, /// --record: each stream's committed rows and tapped features written under this folder, for drafter training. diff --git a/zig/src/families/lfm2_moe/dspark.zig b/zig/src/families/lfm2_moe/dspark.zig index ab3c3f3b33..93ea1f715d 100644 --- a/zig/src/families/lfm2_moe/dspark.zig +++ b/zig/src/families/lfm2_moe/dspark.zig @@ -24,18 +24,8 @@ pub const stage_rows = 128; const hold_grid = [_]u32{ 1, 2, 4, 8, 16 }; const kw = c.kv_heads * c.head_dim; -/// Verify rows by the core's depth rule, the whole block (fixed), or the confidence head's chances (scheduled). -pub const Verify = enum { depth, fixed, scheduled }; -/// The verify when none is asked: the confidence-scheduled one. -pub const default_verify: Verify = .scheduled; - -/// `bits`: 0 runs the bf16 matrices as stored, 8 quantizes them at load; `verify_asked`: --draft-verify given. -pub const Options = struct { capacity: usize, bits: u8 = 0, verify: Verify = default_verify, verify_asked: bool = false }; - -/// An asked scheduled verify needs the checkpoint's confidence head; the default one falls back to the depth rule. -pub fn admitVerify(o: Options, has_head: bool) !void { - if (o.verify == .scheduled and o.verify_asked and !has_head) return error.ScheduledVerifyNeedsConfidenceHead; -} +/// `bits`: 0 runs the bf16 matrices as stored, 8 quantizes them at load. +pub const Options = struct { capacity: usize, bits: u8 = 0 }; /// One stream's context keys and values a layer ([kv head][cap][head dim], as the target's attention reads them). pub const Lane = struct { @@ -111,7 +101,6 @@ pub const Drafter = struct { errdefer d.w.deinit(gpa); errdefer for (d.owned.items) |b| b.deinit(); errdefer d.owned.deinit(gpa); - try admitVerify(options, d.w.confidence != null); try d.scratch(); try d.warm(); return d; @@ -178,9 +167,9 @@ pub const Drafter = struct { return .{ .ptr = d, .vtable = &.{ .taps = tapsFn, .facts = factsFn, .open = openFn, .absorb = absorbFn, .hold = holdFn, .held = heldFn, .release = releaseFn, .chances = chancesFn } }; } - /// Whether holds compute the confidence head (the scheduled verify, on a checkpoint that has the head). + /// Whether holds compute the confidence head: the verify rows follow its chances, else the core's depth rule. pub fn confident(d: *const Drafter) bool { - return d.options.verify == .scheduled and d.w.confidence != null; + return d.w.confidence != null; } fn cast(ptr: *anyopaque) *Drafter { @@ -198,18 +187,8 @@ pub const Drafter = struct { fn factsFn(ptr: *anyopaque) dr.Facts { const d = cast(ptr); var f: dr.Facts = .{ .depth = d.config.block, .step_ms = d.step_ms, .prior = prior[0..d.config.block], .batched = true, .shared_guard = true, .stream_bytes = d.laneBytes(d.options.capacity), .copy_first = true }; - switch (d.options.verify) { - .depth => {}, - .fixed => { // every round the whole block, no plain stretches - f.chances = true; - f.confidence_floor = 0; - f.shared_guard = false; - }, - .scheduled => { - f.chances = d.confident(); - f.hold_costs = d.holds[0..d.holds_len]; - }, - } + f.chances = d.confident(); + f.hold_costs = d.holds[0..d.holds_len]; return f; } @@ -217,10 +196,6 @@ pub const Drafter = struct { fn chancesFn(ptr: *anyopaque, s: *Stream, out: []f64) anyerror!bool { const d = cast(ptr); const l = d.lanes.get(s) orelse return false; - if (d.options.verify == .fixed) { // no wait: equal chances at a depth share a capped round's rows evenly - for (out, 1..) |*o, j| o.* = 1 - 1e-3 * @as(f64, @floatFromInt(j)); - return true; - } if (!d.confident()) return false; if (d.pending != null and l.slot != null and l.conf_len > 0) { // the hold still runs: the lane's last settled confidences price it, so the host does not wait on the GPU diff --git a/zig/src/families/lfm2_moe/dspark_test.zig b/zig/src/families/lfm2_moe/dspark_test.zig index 55c341161c..043a590353 100644 --- a/zig/src/families/lfm2_moe/dspark_test.zig +++ b/zig/src/families/lfm2_moe/dspark_test.zig @@ -210,12 +210,6 @@ test "lfm2 dspark: lfm2_qk on half-split rows and gains gives interleaved RoPE's } } } - -test "lfm2 dspark: the drafter verifies by its confidence head unless asked otherwise" { - const o: @import("dspark.zig").Options = .{ .capacity = 64 }; - try std.testing.expectEqual(@import("dspark.zig").Verify.scheduled, o.verify); -} - test "lfm2 dspark: a settled hold moves the timed passes at its block counts a tenth of the way to its GPU ms" { var d: ds.Drafter = undefined; const grid = [_]lanes.hold_price.HoldCost{ .{ .streams = 1, .depth = 1, .ms = 1 }, .{ .streams = 4, .depth = 1, .ms = 2 }, .{ .streams = 4, .depth = 4, .ms = 4 }, .{ .streams = 8, .depth = 1, .ms = 3 } }; @@ -276,11 +270,3 @@ test "lfm2 dspark: a settled hold refuses a NaN, infinity or invalid pick as the picks[3 * ds.pass_rows] = .{ .nonfinite = 1 }; // a step past the block's depth is not read try block.copyPicks(&picks, &chains, &.{ 2, 2 }); } - -test "lfm2 dspark: an asked scheduled verify without the confidence head is refused by name, the default falls back" { - const asked: ds.Options = .{ .capacity = 64, .verify_asked = true }; - try std.testing.expectError(error.ScheduledVerifyNeedsConfidenceHead, ds.admitVerify(asked, false)); - try ds.admitVerify(asked, true); - try ds.admitVerify(.{ .capacity = 64 }, false); - try ds.admitVerify(.{ .capacity = 64, .verify = .depth, .verify_asked = true }, false); -} diff --git a/zig/src/families/lfm2_moe/lfm2_moe.zig b/zig/src/families/lfm2_moe/lfm2_moe.zig index 75b40e5938..8f861e31b4 100644 --- a/zig/src/families/lfm2_moe/lfm2_moe.zig +++ b/zig/src/families/lfm2_moe/lfm2_moe.zig @@ -11,4 +11,3 @@ pub const dspark = @import("dspark.zig"); pub const dspark_config = @import("dspark_config.zig"); pub const dspark_block = @import("dspark_block.zig"); pub const dspark_quant = @import("dspark_quant.zig"); -pub const verify_costs = @import("verify_costs.zig"); diff --git a/zig/src/families/lfm2_moe/verify_costs.zig b/zig/src/families/lfm2_moe/verify_costs.zig deleted file mode 100644 index f80e4cdcef..0000000000 --- a/zig/src/families/lfm2_moe/verify_costs.zig +++ /dev/null @@ -1,44 +0,0 @@ -//! A verify cost table from tf-lfm2-verify-costs: shared forward ms by total rows, for the lane core's shared costs. -const std = @import("std"); -const lanes = @import("lanes"); -const st = @import("state.zig"); - -const Cost = lanes.config.Cost; - -/// Lines `streams rows_each total_rows ms` ('#' comments): each total's ms at the fewest streams that reach it. -pub fn parse(gpa: std.mem.Allocator, text: []const u8) ![]Cost { - var best: [st.batch_rows + 1]?struct { streams: u32, ms: f64 } = @splat(null); - var lines = std.mem.tokenizeAny(u8, text, "\r\n"); - while (lines.next()) |line| { - const body = std.mem.trim(u8, line, " \t"); - if (body.len == 0 or body[0] == '#') continue; - var fields = std.mem.tokenizeAny(u8, body, " \t"); - var v: [4][]const u8 = undefined; - for (&v) |*x| x.* = fields.next() orelse return error.BadVerifyCosts; - if (fields.next() != null) return error.BadVerifyCosts; - const streams = try std.fmt.parseInt(u32, v[0], 10); - const each = try std.fmt.parseInt(u32, v[1], 10); - const total = try std.fmt.parseInt(u32, v[2], 10); - const ms = try std.fmt.parseFloat(f64, v[3]); - if (streams == 0 or each == 0 or total != streams * each or total > st.batch_rows or !(ms > 0)) return error.BadVerifyCosts; - if (best[total]) |b| if (b.streams <= streams) continue; - best[total] = .{ .streams = streams, .ms = ms }; - } - var out: std.ArrayList(Cost) = .empty; - errdefer out.deinit(gpa); - for (best, 0..) |b, total| if (b) |x| try out.append(gpa, .{ .width = @intCast(total), .ms = x.ms }); - if (out.items.len == 0) return error.BadVerifyCosts; - return out.toOwnedSlice(gpa); -} - -test "lfm2 dspark: a verify cost table gives each total its fewest streams' ms and refuses malformed lines" { - const gpa = std.testing.allocator; - const t = try parse(gpa, "# streams rows_each total_rows ms\n1 1 1 3.7\n2 1 2 4.9\n1 2 2 4.6\n\n4 8 32 40.5\n"); - defer gpa.free(t); - try std.testing.expectEqual(@as(usize, 3), t.len); - try std.testing.expectEqual(@as(f64, 4.6), t[1].ms); // one stream of two rows, not two of one - try std.testing.expectEqual(@as(u32, 32), t[2].width); - try std.testing.expectError(error.BadVerifyCosts, parse(gpa, "2 2 5 4.0\n")); - try std.testing.expectError(error.BadVerifyCosts, parse(gpa, "8 8 64 9.0\n")); - try std.testing.expectError(error.BadVerifyCosts, parse(gpa, "# nothing\n")); -} diff --git a/zig/src/native/lfm2_moe.zig b/zig/src/native/lfm2_moe.zig index dd222e9af0..d5114d2a0c 100644 --- a/zig/src/native/lfm2_moe.zig +++ b/zig/src/native/lfm2_moe.zig @@ -88,8 +88,7 @@ pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Ope release(h); return null; } - const verify = verifyMode(o.draft_verify); - h.draft = q.dspark.Drafter.init(gpa, io, model, dir, .{ .capacity = window, .bits = bits, .verify = verify, .verify_asked = o.draft_verify != null }) catch |err| { + h.draft = q.dspark.Drafter.init(gpa, io, model, dir, .{ .capacity = window, .bits = bits }) catch |err| { problem.* = try std.fmt.allocPrint(a, "the native LFM2 drafter cannot load {s} ({s}): it takes LiquidAI/LFM2.5-24B-A2B-DSpark", .{ dir, @errorName(err) }); release(h); return null; @@ -112,14 +111,6 @@ pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Ope } errdefer if (h.wrap) |*w| w.deinit(); var facts = if (h.wrap) |*w| w.facts(h.metal.facts()) else h.metal.facts(); - const profiled = if (h.draft != null) verifyCosts(a, io, o.verify_costs) catch |err| { - problem.* = try std.fmt.allocPrint(a, "--verify-costs {s} is not a tf-lfm2-verify-costs table ({s})", .{ o.verify_costs.?, @errorName(err) }); - h.wrap.?.deinit(); - h.draft.?.deinit(); - release(h); - return null; - } else null; - if (profiled) |t| facts.shared_costs = t; // a drafting stream prices its copy windows by its own copies: on distinct prompts they keep breaking if (h.wrap != null) facts.own_copy_rates = true; h.config = try lanes.Config.init(gpa, facts, q.state.window_rows, q.state.window_rows - 1); @@ -164,26 +155,11 @@ pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Ope try h.host.start(); if (h.draft) |d| { const mode = if (d.options.bits == 8) "8-bit codes quantized at load (--drafter-bits 8), not the released bf16" else "bf16 matrices as released"; - std.log.info("lfm2: DSpark drafts from --drafter, {d} a round at most, {s}, {t} verify", .{ d.config.block, mode, d.options.verify }); + std.log.info("lfm2: DSpark drafts from --drafter, {d} a round at most, {s}", .{ d.config.block, mode }); } return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; } -/// --verify-costs: a tf-lfm2-verify-costs table's shared forward ms by total rows (null: none named; `a`-owned). -fn verifyCosts(a: std.mem.Allocator, io: std.Io, path: ?[]const u8) !?[]const lanes.config.Cost { - const file = path orelse return null; - const text = try std.Io.Dir.cwd().readFileAlloc(io, file, a, .limited(1 << 20)); - const t = try q.verify_costs.parse(a, text); - std.log.info("lfm2: verify costs from {s}, {d} totals of rows", .{ file, t.len }); - return t; -} - -/// --draft-verify as the drafter's verify (null: the default). -fn verifyMode(v: ?api.DraftVerify) q.dspark.Verify { - const asked = v orelse return q.dspark.default_verify; - return std.meta.stringToEnum(q.dspark.Verify, @tagName(asked)).?; -} - /// The prompt cache's copies of a stream: the target's state at a mark (backend.zig), and the drafter's context there. const Snaps = struct { const S = q.backend.Snaps; diff --git a/zig/src/native/metal.zig b/zig/src/native/metal.zig index a5eda7c55c..41dc28577d 100644 --- a/zig/src/native/metal.zig +++ b/zig/src/native/metal.zig @@ -334,7 +334,6 @@ fn refusal(a: Allocator, o: api.Open, dense: bool) !?[]const u8 { const one_mac: ?[]const u8 = if (std.mem.eql(u8, t, "nemotron_h")) "Nemotron" else if (std.mem.eql(u8, t, "lfm2_moe")) "LFM2.5-24B-A2B" else if (isGemma4(t)) "Gemma 4" else if (std.mem.eql(u8, t, "step5_vl")) "Step-5" else if (std.mem.eql(u8, t, "spark2_5")) "Spark-X2.6" else if (!std.mem.eql(u8, t, "qwen3_5")) null else if (dense) "Qwen3.8-27B" else "Qwen3.5-2B"; if (o.speed_up != null) if (one_mac) |name| return try std.fmt.allocPrint(a, "--speed-up joins two Macs for Flash Next, GLM-5.3-Flash and GLM-5.3 only: serve {s} without it", .{name}); if (o.recall and std.mem.eql(u8, t, "qwen4_exp")) return recall_refusal; - if ((o.draft_verify != null or o.verify_costs != null) and (o.drafter == null or !std.mem.eql(u8, t, "lfm2_moe"))) return "--draft-verify and --verify-costs set the verify of LFM2.5-24B-A2B's DSpark drafter: serve it with --drafter"; if (spill.refusal(o)) |why| return why; if (o.learn == null) return null; if (o.learn_conversations and std.mem.eql(u8, t, "qwen3_5") and dense) return null; diff --git a/zig/src/server/cli.zig b/zig/src/server/cli.zig index 6e41813208..5a38d9fd3a 100644 --- a/zig/src/server/cli.zig +++ b/zig/src/server/cli.zig @@ -101,8 +101,6 @@ pub const flags = [_]Flag{ .{ .name = "--slide-graph", .native = true }, .{ .name = "--drafter", .native = true }, .{ .name = "--drafter-bits", .choices = &.{ "0", "4", "8" }, .native = true }, - .{ .name = "--draft-verify", .choices = &.{ "depth", "fixed", "scheduled" }, .native = builtin.os.tag == .macos }, - .{ .name = "--verify-costs", .native = builtin.os.tag == .macos }, .{ .name = "--record", .native = true }, .{ .name = "--kv-cache", .choices = &.{ "bf16", "fp8" }, .native = true }, .{ .name = "--prompt-precision", .choices = &.{ "bf16", "int8" }, .native = true }, @@ -213,8 +211,6 @@ pub const Args = struct { compact_memory: ?[]const u8 = null, drafter: ?[]const u8 = null, drafter_bits: ?u8 = null, // null: not given (a family without its own default takes 4) - draft_verify: ?api.DraftVerify = null, // the LFM2 DSpark drafter's verify rows (null: the engine's default) - verify_costs: ?[]const u8 = null, // a tf-lfm2-verify-costs table pricing them (null: costs timed at load) record: ?[]const u8 = null, // drafter training rows written under this folder kv_cache: []const u8 = "bf16", // fp8: e4m3 attention caches with a scale a row, about twice the context prompt_precision: []const u8 = "bf16", // int8: prompt matmuls quantize their inputs to int8 a 64-input group @@ -418,7 +414,7 @@ fn apply(a: Allocator, out: *Args, name: []const u8, value: ?[]const u8, u: *Usa out.learn_daily_gib = try gib(u, a, name, v); } else if (is(name, "--prompt-seed")) { out.prompt_seed = true; - } else if (is(name, "--max-tokens")) out.max_tokens = try int(u, a, name, v) else if (is(name, "--temperature")) out.temperature = try float(u, a, name, v) else if (is(name, "--top-p")) out.top_p = try float(u, a, name, v) else if (is(name, "--top-k")) out.top_k = try int(u, a, name, v) else if (is(name, "--min-p")) out.min_p = try float(u, a, name, v) else if (is(name, "--thinking")) out.thinking = true else if (is(name, "--no-thinking")) out.thinking = false else if (is(name, "--reasoning-effort")) out.reasoning_effort = v else if (is(name, "--thinking-budget")) out.thinking_budget = try int(u, a, name, v) else if (is(name, "--thinking-reserve")) out.thinking_reserve = try int(u, a, name, v) else if (is(name, "--loop-guard")) out.loop_guard = true else if (is(name, "--drafter")) out.drafter = v else if (is(name, "--drafter-bits")) out.drafter_bits = @intCast(try int(u, a, name, v)) else if (is(name, "--draft-verify")) out.draft_verify = std.meta.stringToEnum(api.DraftVerify, v) else if (is(name, "--verify-costs")) out.verify_costs = v else if (is(name, "--record")) out.record = v else if (is(name, "--kv-cache")) out.kv_cache = v else if (is(name, "--prompt-precision")) out.prompt_precision = v else if (is(name, "--no-drafts")) out.no_drafts = true else if (is(name, "--keep-warm")) out.keep_warm = try int(u, a, name, v) else if (is(name, "--idle-unload")) out.idle_unload = try int(u, a, name, v) else if (is(name, "--compact-at")) { + } else if (is(name, "--max-tokens")) out.max_tokens = try int(u, a, name, v) else if (is(name, "--temperature")) out.temperature = try float(u, a, name, v) else if (is(name, "--top-p")) out.top_p = try float(u, a, name, v) else if (is(name, "--top-k")) out.top_k = try int(u, a, name, v) else if (is(name, "--min-p")) out.min_p = try float(u, a, name, v) else if (is(name, "--thinking")) out.thinking = true else if (is(name, "--no-thinking")) out.thinking = false else if (is(name, "--reasoning-effort")) out.reasoning_effort = v else if (is(name, "--thinking-budget")) out.thinking_budget = try int(u, a, name, v) else if (is(name, "--thinking-reserve")) out.thinking_reserve = try int(u, a, name, v) else if (is(name, "--loop-guard")) out.loop_guard = true else if (is(name, "--drafter")) out.drafter = v else if (is(name, "--drafter-bits")) out.drafter_bits = @intCast(try int(u, a, name, v)) else if (is(name, "--record")) out.record = v else if (is(name, "--kv-cache")) out.kv_cache = v else if (is(name, "--prompt-precision")) out.prompt_precision = v else if (is(name, "--no-drafts")) out.no_drafts = true else if (is(name, "--keep-warm")) out.keep_warm = try int(u, a, name, v) else if (is(name, "--idle-unload")) out.idle_unload = try int(u, a, name, v) else if (is(name, "--compact-at")) { if (std.mem.eql(u8, v, "auto")) out.compact_auto = true else { const f = try float(u, a, name, v); if (!(f > 0 and f <= 1)) return fail(u, a, "argument --compact-at: expected auto or a fraction in (0, 1]: '{s}'", .{v}); @@ -546,8 +542,6 @@ pub fn request(a: Allocator, dir: []const u8, model_type: []const u8, args: Args .drafter = args.drafter, .drafter_bits = args.drafter_bits orelse 4, .drafter_bits_given = args.drafter_bits != null, - .draft_verify = args.draft_verify, - .verify_costs = args.verify_costs, .record = args.record, .kv8 = std.mem.eql(u8, args.kv_cache, "fp8"), .prompt_int8 = std.mem.eql(u8, args.prompt_precision, "int8"), diff --git a/zig/src/server/cli_test.zig b/zig/src/server/cli_test.zig index dbbf7ac152..f35a44e631 100644 --- a/zig/src/server/cli_test.zig +++ b/zig/src/server/cli_test.zig @@ -167,13 +167,6 @@ test "parse and capabilities share the table" { try std.testing.expectEqual(@as(?u8, 8), (try parse(a, &.{ "m", "--drafter-bits", "8" }, &u)).drafter_bits); try std.testing.expectEqual(@as(?u8, null), (try parse(a, &.{ "m", "--drafter", "x" }, &u)).drafter_bits); try std.testing.expectError(error.Usage, parse(a, &.{ "m", "--drafter-bits", "3" }, &u)); - if (@import("builtin").os.tag == .macos) { - const verify = try parse(a, &.{ "m", "--draft-verify", "scheduled", "--verify-costs", "/t" }, &u); - try std.testing.expectEqual(api.DraftVerify.scheduled, verify.draft_verify.?); - try std.testing.expectEqual(api.DraftVerify.scheduled, (try request(a, "d", "t", verify)).draft_verify.?); - try std.testing.expectEqualStrings("/t", (try request(a, "d", "t", verify)).verify_costs.?); - try std.testing.expectError(error.Usage, parse(a, &.{ "m", "--draft-verify", "wide" }, &u)); - } try std.testing.expectEqualStrings("/r", (try parse(a, &.{ "m", "--record", "/r" }, &u)).record.?); try std.testing.expectEqualStrings("fp8", (try parse(a, &.{ "m", "--kv-cache", "fp8" }, &u)).kv_cache); try std.testing.expectError(error.Usage, parse(a, &.{ "m", "--kv-cache", "int4" }, &u)); diff --git a/zig/src/server/main.zig b/zig/src/server/main.zig index 056ebe0eb7..12fbfd93f2 100644 --- a/zig/src/server/main.zig +++ b/zig/src/server/main.zig @@ -25,7 +25,7 @@ const usage_line = "usage: tensorfold serve [-h] [--host HOST] [--port PORT] [-- " [--thinking | --no-thinking] [--reasoning-effort {low,medium,high,xhigh}]" ++ " [--thinking-budget THINKING_BUDGET] [--thinking-reserve THINKING_RESERVE] [--loop-guard] [--no-drafts]" ++ " [--drafter DRAFTER]" ++ - " [--drafter-bits {0,4,8}] [--draft-verify {depth,fixed,scheduled}] [--verify-costs VERIFY_COSTS] [--record RECORD] [--keep-warm SECONDS] [--idle-unload SECONDS]" ++ + " [--drafter-bits {0,4,8}] [--record RECORD] [--keep-warm SECONDS] [--idle-unload SECONDS]" ++ " [--silence-floor-dbfs DBFS] [--silence-seconds SECONDS]" ++ " [--silence-onset-s SECONDS] [--silence-onset-voiced-s SECONDS]" ++ " [--silence-mid-s SECONDS] [--silence-mid-min-voiced-s SECONDS] [--silence-redraws N]" ++ diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 1b3822179f..fc49c99405 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -87,7 +87,6 @@ test { _ = @import("families/lfm2_moe/dspark_config.zig"); _ = @import("families/lfm2_moe/dspark_quant.zig"); _ = @import("families/lfm2_moe/dspark_test.zig"); - _ = @import("families/lfm2_moe/verify_costs.zig"); _ = @import("families/lfm2_moe/dspark_bf16_test.zig"); } diff --git a/zig/tests/lfm2_dspark.zig b/zig/tests/lfm2_dspark.zig index 7c74a66e09..07b72b093a 100644 --- a/zig/tests/lfm2_dspark.zig +++ b/zig/tests/lfm2_dspark.zig @@ -137,7 +137,7 @@ pub fn main(init: std.process.Init) !void { defer pool.pop(); const m = try q.Model.load(gpa, init.io, args[1]); defer m.deinit(); - const d = try q.dspark.Drafter.init(gpa, init.io, m, args[2], .{ .capacity = length + 64, .bits = bits, .verify = .scheduled }); + const d = try q.dspark.Drafter.init(gpa, init.io, m, args[2], .{ .capacity = length + 64, .bits = bits }); defer d.deinit(); const cfg = d.config; // the compared block's rows (a hold of depth d runs a d-row block) diff --git a/zig/tests/lfm2_verify_costs.zig b/zig/tests/lfm2_verify_costs.zig deleted file mode 100644 index f3100d633e..0000000000 --- a/zig/tests/lfm2_verify_costs.zig +++ /dev/null @@ -1,87 +0,0 @@ -//! LFM2 verify forwards timed by stream count and rows a stream: the round costs a DSpark verify budget weighs. -const std = @import("std"); -const mtl = @import("metal"); -const q = @import("tensorfold").lfm2_moe; -const st = q.state; -const c = q.config; - -fn now() i128 { - var ts: std.c.timespec = undefined; - _ = std.c.clock_gettime(.MONOTONIC, &ts); - return @as(i128, ts.sec) * std.time.ns_per_s + ts.nsec; -} - -/// Stream counts profiled; each runs every rows-a-stream width whose total fits a shared forward. -const counts = [_]usize{ 1, 2, 3, 4, 6, 8, 12, 16 }; -const reps = 9; - -fn readIds(gpa: std.mem.Allocator, io: std.Io, path: []const u8, list: *std.ArrayList(u32)) !void { - const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 22)); - defer gpa.free(bytes); - var it = std.mem.tokenizeAny(u8, bytes, " \n\r\t,"); - while (it.next()) |t| try list.append(gpa, try std.fmt.parseInt(u32, t, 10)); -} - -pub fn main(init: std.process.Init) !void { - const gpa = init.gpa; - const args = try init.minimal.args.toSlice(init.arena.allocator()); - if (args.len < 3) return error.ExpectedModelOutputHistory; - var history: usize = if (args.len >= 4) try std.fmt.parseInt(usize, args[3], 10) else 1024; - // Token id files: every file's ids end to end; stream i reads its history and then its rows from its own sixteenth. - var real: std.ArrayList(u32) = .empty; - defer real.deinit(gpa); - for (args[@min(args.len, 4)..]) |path| try readIds(gpa, init.io, path, &real); - const part = real.items.len / 16; - if (real.items.len > 0) { - if (part < st.window_rows + 64) return error.IdFilesTooShort; - history = @min(history, part - st.window_rows); - } - const pool = mtl.objc.Pool.push(); - defer pool.pop(); - const m = try q.Model.load(gpa, init.io, args[1]); - defer m.deinit(); - var scratch = try st.Scratch.init(gpa, m.device, &m.config.attention, @max(history, st.batch_rows)); - defer scratch.deinit(); - var caches: [16]st.Cache = undefined; - var made: usize = 0; - defer for (caches[0..made]) |*k| k.deinit(); - const ids = try gpa.alloc(u32, history); - defer gpa.free(ids); - for (caches[0..], 0..) |*k, i| { - k.* = try st.Cache.init(gpa, m.device, &m.config.attention, history + st.window_rows + 1); - made += 1; - for (ids, 0..) |*id, j| id.* = if (part > 0) real.items[i * part + j] else @intCast((j * 7919 + 13 + i * 104729) % 120000); - try q.forward.run(m, &scratch, &.{.{ .cache = k, .rows = history }}, ids, false, .last); - } - var out: std.Io.Writer.Allocating = .init(gpa); - defer out.deinit(); - const w = &out.writer; - try w.print("# LFM2 verify forward ms (median of {d}) by streams and rows a stream, {d}-token histories, {s} rows\n", .{ reps, history, if (part > 0) "real" else "synthetic" }); - try w.print("# streams rows_each total_rows ms\n", .{}); - var segments: [st.batch_rows]q.forward.Segment = undefined; - var tokens: [st.batch_rows]u32 = undefined; - for (&tokens, 0..) |*t, i| t.* = @intCast((i * 31 + 7) % 120000); - for (counts) |streams| { - var each: usize = 1; - while (each <= st.window_rows and streams * each <= st.batch_rows) : (each += 1) { - for (segments[0..streams], caches[0..streams], 0..) |*s, *k, i| { - s.* = .{ .cache = k, .rows = each }; - // real rows: the tokens after the stream's history, as a verify of drafts that all land - if (part > 0) @memcpy(tokens[i * each ..][0..each], real.items[i * part + history ..][0..each]); - } - var times: [reps + 2]f64 = undefined; - for (×) |*t| { - const t0 = now(); - try q.forward.run(m, &scratch, segments[0..streams], tokens[0 .. streams * each], true, .all); - t.* = @as(f64, @floatFromInt(now() - t0)) / 1e6; - for (caches[0..streams]) |*k| k.len = history; - } - const kept = times[2..]; // the first two calls warm a new width's pipelines - std.mem.sort(f64, kept, {}, std.sort.asc(f64)); - const ms = kept[reps / 2]; - try w.print("{d} {d} {d} {d:.3}\n", .{ streams, each, streams * each, ms }); - std.debug.print("{d} streams x {d} rows ({d}): {d:.3} ms\n", .{ streams, each, streams * each, ms }); - } - } - try std.Io.Dir.cwd().writeFile(init.io, .{ .sub_path = args[2], .data = out.written() }); -}