From b27726a4dfd2d5865edfa7e9c3112f39e662d222 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:19:36 -1000 Subject: [PATCH 01/10] MLX quantization: an override answers its module in any spelling oQ checkpoints write per-module widths as language_model.model.layers.N..., while loaders ask for model.layers.N... or the bare path, so an exact-key lookup missed every override and read the tensor at the global width. The table keeps each entry under the config's own key, as before, and also indexes it under its canonical module key (no .weight suffix, no known prefix, the rule the 27B and Flash Next families use); `width` takes the key as given, then its module. Two spellings of one module that declare different widths refuse the table as unsupported, the error `detect` already returns, so its callers' error sets are unchanged; the same width twice is fine. Test that fails before: the fixture reads a 5-bit down_proj declared under the oQ prefix by the loader spelling; on main it fails the packed shape check (read at the global width). Also: every prefix form and a trailing .weight name one module; aliases must agree; an override in any spelling answers every spelling and stays under its own key. Gate on main 0efbd447 (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test: `get` logs an error on the DeferredTensor path the test expects). lean_check: 539 problems (main 539), no new findings. --- zig/src/core/quant/keys.zig | 44 +++++++++ zig/src/core/quant/mlx.zig | 174 ++++++++++++++++++++++++++++++++++-- 2 files changed, 212 insertions(+), 6 deletions(-) create mode 100644 zig/src/core/quant/keys.zig diff --git a/zig/src/core/quant/keys.zig b/zig/src/core/quant/keys.zig new file mode 100644 index 0000000000..fe20231e4f --- /dev/null +++ b/zig/src/core/quant/keys.zig @@ -0,0 +1,44 @@ +//! Override keys by module: every spelling checkpoints and loaders give one projection reduces to one canonical key. + +const std = @import("std"); + +/// The prefixes a module path may carry, each stripped at most once and in this order. +const prefixes = [_][]const u8{ "model.language_model.", "language_model.", "text_model.", "model." }; + +/// The module `key` names, without a trailing `.weight` or the known prefixes; a slice of `key`. +pub fn canonical(key: []const u8) []const u8 { + var k = key; + if (std.mem.endsWith(u8, k, ".weight")) k = k[0 .. k.len - ".weight".len]; + for (prefixes) |prefix| { + if (std.mem.startsWith(u8, k, prefix)) k = k[prefix.len..]; + } + return k; +} + +test "every prefix form and a trailing weight name one module" { + for ([_][]const u8{ + "layers.0.mlp.down_proj", + "model.layers.0.mlp.down_proj", + "language_model.model.layers.0.mlp.down_proj", + "language_model.layers.0.mlp.down_proj", + "model.language_model.layers.0.mlp.down_proj", + "model.language_model.model.layers.0.mlp.down_proj", + "text_model.layers.0.mlp.down_proj", + "text_model.model.layers.0.mlp.down_proj", + "layers.0.mlp.down_proj.weight", + "model.layers.0.mlp.down_proj.weight", + "language_model.model.layers.0.mlp.down_proj.weight", + "model.language_model.layers.0.mlp.down_proj.weight", + }) |key| try std.testing.expectEqualStrings("layers.0.mlp.down_proj", canonical(key)); +} + +test "multi-token prediction and vision modules keep their own names" { + for ([_][]const u8{ "mtp.layers.0.mlp.gate_proj", "model.mtp.layers.0.mlp.gate_proj", "language_model.mtp.layers.0.mlp.gate_proj", "language_model.model.mtp.layers.0.mlp.gate_proj.weight" }) |key| + try std.testing.expectEqualStrings("mtp.layers.0.mlp.gate_proj", canonical(key)); + try std.testing.expectEqualStrings("vision_tower.blocks.0.attn.qkv", canonical("vision_tower.blocks.0.attn.qkv.weight")); + try std.testing.expectEqualStrings("visual.blocks.0.attn.qkv", canonical("model.visual.blocks.0.attn.qkv")); + try std.testing.expectEqualStrings("lm_head", canonical("language_model.lm_head")); + // a prefix only counts at the start and a suffix only at the end + try std.testing.expectEqualStrings("layers.0.weight_proj", canonical("layers.0.weight_proj")); + try std.testing.expectEqualStrings("vision_tower.model.x", canonical("vision_tower.model.x")); +} diff --git a/zig/src/core/quant/mlx.zig b/zig/src/core/quant/mlx.zig index f2334eb01a..4b8a289821 100644 --- a/zig/src/core/quant/mlx.zig +++ b/zig/src/core/quant/mlx.zig @@ -4,6 +4,7 @@ const std = @import("std"); const types = @import("types.zig"); const slice = @import("slice.zig"); const convert = @import("convert.zig"); +const keys = @import("keys.zig"); const Tensor = types.Tensor; const Allocator = std.mem.Allocator; @@ -18,19 +19,36 @@ pub const bit_widths = [_]u8{ 2, 3, 4, 5, 6, 8 }; /// One affine width: groups of `group` weights share a scale and a bias, `bits` per weight. pub const Width = struct { bits: u8, group: u16 }; -/// The config's `quantization` table: the global width and the per-tensor overrides (mixed-width conversions). +/// The config's `quantization` table: the global width and the per-module overrides (mixed-width conversions). pub const Config = struct { global: Width, - /// A per-tensor entry; `null` marks one this format refuses, which fails when the tensor is looked up. + /// A per-tensor entry under the config's own key; `null` marks one this format refuses, which fails on lookup. overrides: std.StringHashMapUnmanaged(?Width) = .empty, + /// The same entries under their module's canonical key, so every spelling of a module finds its width. + modules: std.StringHashMapUnmanaged(?Width) = .empty, - /// The tensor's own width when the table names it as an object, else the global one. + /// The width the table names `key` with, else its module under any spelling, else the global one. pub fn width(q: Config, key: []const u8) error{UnsupportedQuantization}!Width { - if (q.overrides.get(key)) |w| return w orelse error.UnsupportedQuantization; + const named = q.overrides.get(key) orelse q.modules.get(keys.canonical(key)); + if (named) |w| return w orelse error.UnsupportedQuantization; return q.global; } + + /// One entry under its own key and its module's; an alias already there must declare the same width. + fn add(q: *Config, a: Allocator, key: []const u8, w: ?Width) DetectError!void { + try q.overrides.put(a, key, w); + const slot = try q.modules.getOrPut(a, keys.canonical(key)); + if (slot.found_existing and !std.meta.eql(slot.value_ptr.*, w)) { + std.log.warn("{s}: another spelling of module {s} declares a different width", .{ key, slot.key_ptr.* }); + return error.UnsupportedQuantization; + } + slot.value_ptr.* = w; + } }; +/// Why a table is refused: allocation, or two spellings of one module that declare different widths (unsupported). +pub const DetectError = Allocator.Error || error{UnsupportedQuantization}; + /// The affine width an object of the config names, or null when it is not one the kernels run. fn affineWidth(o: std.json.ObjectMap) ?Width { const bits = intOrNull(o.get("bits")) orelse return null; @@ -51,11 +69,11 @@ fn intOrNull(v: ?std.json.Value) ?i64 { } /// The checkpoint's `quantization` object when it is MLX affine (`mode` absent or "affine"), else null. -pub fn detect(a: Allocator, src: types.Sources) Allocator.Error!?Config { +pub fn detect(a: Allocator, src: types.Sources) DetectError!?Config { const obj = src.quantization; var q: Config = .{ .global = affineWidth(obj) orelse return null }; var it = obj.iterator(); - while (it.next()) |kv| if (kv.value_ptr.* == .object) try q.overrides.put(a, kv.key_ptr.*, affineWidth(kv.value_ptr.object)); + while (it.next()) |kv| if (kv.value_ptr.* == .object) try q.add(a, kv.key_ptr.*, affineWidth(kv.value_ptr.object)); return q; } @@ -328,3 +346,147 @@ test "the config's per-tensor widths override the global one" { const gptq = (try std.json.parseFromSliceLeaky(std.json.Value, a, "{\"group_size\": 64, \"bits\": 4, \"mode\": \"gptq\"}", .{})).object; try std.testing.expect((try detect(a, .{ .root = gptq, .quantization = gptq })) == null); } + +/// `codes` packed `bits` wide, low code first, into little-endian words (a code may cross into the next word). +fn packCodes(bits: u8, codes: []const u8, words: []u32) void { + @memset(words, 0); + for (codes, 0..) |code, i| { + const bit = i * bits; + const v: u64 = @as(u64, code) << @intCast(bit & 31); + words[bit >> 5] |= @truncate(v); + if ((bit & 31) + bits > 32) words[(bit >> 5) + 1] |= @truncate(v >> 32); + } +} + +/// A checkpoint as `read` sees it: tensors by name, and the config's widths by the key the loader asks for. +const Fixture = struct { + q: Config, + names: []const []const u8, + tensors: []const Tensor, + + pub fn has(f: Fixture, name: []const u8) bool { + return f.find(name) != null; + } + + pub fn get(f: Fixture, name: []const u8) error{MissingTensor}!Tensor { + return f.tensors[f.find(name) orelse return error.MissingTensor]; + } + + pub fn width(f: Fixture, key: []const u8) error{UnsupportedQuantization}!Width { + return f.q.width(key); + } + + fn find(f: Fixture, name: []const u8) ?usize { + for (f.names, 0..) |n, i| if (std.mem.eql(u8, n, name)) return i; + return null; + } +}; + +test "an override under the oQ prefix reads a 5-bit tensor the loader names its own way" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const text = + \\{"group_size": 64, "bits": 4, "mode": "affine", + \\ "language_model.model.layers.0.mlp.down_proj": {"group_size": 64, "bits": 5, "mode": "affine"}} + ; + const obj = (try std.json.parseFromSliceLeaky(std.json.Value, a, text, .{})).object; + const q = (try detect(a, .{ .root = obj, .quantization = obj })).?; + // one row of 64 codes 5 bits wide: ten words and one group; the global 4 bits would need eight words + var codes: [64]u8 = undefined; + for (&codes, 0..) |*c, i| c.* = @intCast(i % 32); + var words: [10]u32 = undefined; + packCodes(5, &codes, &words); + var scale = [_]u16{0x3F80}; // bf16 1.0 + var bias = [_]u16{0x0000}; + const fx: Fixture = .{ + .q = q, + .names = &.{ "model.layers.0.mlp.down_proj.weight", "model.layers.0.mlp.down_proj.scales", "model.layers.0.mlp.down_proj.biases" }, + .tensors = &.{ + .{ .dtype = .u32, .rank = 2, .shape = .{ 1, 10, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(&words) }, + .{ .dtype = .bf16, .rank = 2, .shape = .{ 1, 1, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(&scale) }, + .{ .dtype = .bf16, .rank = 2, .shape = .{ 1, 1, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(&bias) }, + }, + }; + const h = try read(a, fx, "model.layers.0.mlp.down_proj"); + try std.testing.expectEqual(Width{ .bits = 5, .group = 64 }, Width{ .bits = h.bits, .group = h.group }); + var out: [64]f64 = undefined; + try reference(h, &out); + for (out, codes) |v, c| try std.testing.expectEqual(@as(f64, @floatFromInt(c)), v); +} + +/// The table `text` declares, parsed into `a`. +fn tableOf(a: Allocator, text: []const u8) DetectError!?Config { + const obj = (std.json.parseFromSliceLeaky(std.json.Value, a, text, .{}) catch unreachable).object; + return detect(a, .{ .root = obj, .quantization = obj }); +} + +test "an override declared in any spelling answers every spelling of its module" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const spellings = [_][]const u8{ + "layers.0.mlp.down_proj", + "model.layers.0.mlp.down_proj", + "language_model.model.layers.0.mlp.down_proj", + "model.language_model.layers.0.mlp.down_proj", + "model.language_model.model.layers.0.mlp.down_proj", + "text_model.layers.0.mlp.down_proj", + "model.layers.0.mlp.down_proj.weight", + "language_model.model.layers.0.mlp.down_proj.weight", + }; + for (spellings) |declared| { + const text = try std.fmt.allocPrint(a, "{{\"group_size\": 64, \"bits\": 4, \"{s}\": {{\"group_size\": 32, \"bits\": 3}}}}", .{declared}); + const q = (try tableOf(a, text)).?; + for (spellings) |asked| try std.testing.expectEqual(Width{ .bits = 3, .group = 32 }, try q.width(asked)); + // the config's own key stays in `overrides`, for loaders that read the map by it + try std.testing.expectEqual(@as(?Width, .{ .bits = 3, .group = 32 }), q.overrides.get(declared).?); + try std.testing.expectEqual(Width{ .bits = 4, .group = 64 }, try q.width("model.layers.1.mlp.down_proj")); + } + const mtp = (try tableOf(a, "{\"group_size\": 64, \"bits\": 4, \"language_model.mtp.fc\": {\"group_size\": 64, \"bits\": 8}}")).?; + try std.testing.expectEqual(Width{ .bits = 8, .group = 64 }, try mtp.width("mtp.fc.weight")); + try std.testing.expectEqual(Width{ .bits = 4, .group = 64 }, try mtp.width("fc")); +} + +test "aliases of one module must agree on its width" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + try std.testing.expectError(error.UnsupportedQuantization, tableOf(a, + \\{"group_size": 64, "bits": 4, "language_model.model.layers.0.mlp.up_proj": {"group_size": 64, "bits": 5}, + \\ "model.layers.0.mlp.up_proj.weight": {"group_size": 64, "bits": 6}} + )); + try std.testing.expectError(error.UnsupportedQuantization, tableOf(a, + \\{"group_size": 64, "bits": 4, "language_model.model.layers.0.mlp.up_proj": {"group_size": 64, "bits": 5}, + \\ "layers.0.mlp.up_proj": {"group_size": 32, "bits": 5}} + )); + // an unsupported entry and a width for the same module disagree too + try std.testing.expectError(error.UnsupportedQuantization, tableOf(a, + \\{"group_size": 64, "bits": 4, "language_model.model.layers.0.mlp.up_proj": {"group_size": 64, "bits": 5}, + \\ "model.layers.0.mlp.up_proj": {"group_size": 64, "bits": 5, "mode": "mxfp4"}} + )); + const same = (try tableOf(a, + \\{"group_size": 64, "bits": 4, "language_model.model.layers.0.mlp.up_proj": {"group_size": 64, "bits": 5}, + \\ "model.layers.0.mlp.up_proj": {"group_size": 64, "bits": 5, "mode": "affine"}} + )).?; + try std.testing.expectEqual(Width{ .bits = 5, .group = 64 }, try same.width("layers.0.mlp.up_proj")); +} + +test "an entry the kernels do not run refuses every spelling of its module" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const q = (try tableOf(a, + \\{"group_size": 64, "bits": 4, "language_model.model.layers.0.mlp.gate_proj": {"group_size": 64, "bits": 4, "mode": "mxfp4"}, + \\ "model.layers.0.mlp.gate_proj": {"group_size": 16, "bits": 4}, + \\ "language_model.model.layers.1.mlp.gate_proj": {"group_size": 64, "bits": 7}} + )).?; + for ([_][]const u8{ "layers.0.mlp.gate_proj", "model.layers.0.mlp.gate_proj.weight", "model.language_model.layers.0.mlp.gate_proj" }) |key| + try std.testing.expectError(error.UnsupportedQuantization, q.width(key)); + try std.testing.expectError(error.UnsupportedQuantization, q.width("model.layers.1.mlp.gate_proj")); + try std.testing.expectEqual(Width{ .bits = 4, .group = 64 }, try q.width("model.layers.2.mlp.gate_proj")); +} + +test { + _ = keys; +} From 7f1c3eed66da0aa5eb7b7ab88a36cc9ee320d02e Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:18:51 -1000 Subject: [PATCH 02/10] MLX affine weights widen exactly: codes re-packed at more bits, group tables repeated for a smaller group A projection of mixed-width members (the Flash Next hyper-connection down and inject pair, where 0.6.6's QWeights.of widens) reads as one matrix at their common width, the most bits and the smallest group. Narrowing and a group that does not divide the source group are refused. Used only where 0.6.6 widens; no caller on this commit. Tests (widen.zig, new): every source width and every wider target dequantize to the same fp64 values, bit for bit; a stack of experts widens row for row; the same width comes back as it is; narrowing bits, growing groups and widths the kernels do not run are refused; the common width takes the most bits and the smallest group; fused members keep each one's values in order. Gate on main 0efbd447 (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test: `get` logs an error on the DeferredTensor path the test expects). lean_check: 539 problems (main 539), no new findings. --- zig/src/core/quant/quant.zig | 2 + zig/src/core/quant/widen.zig | 321 +++++++++++++++++++++++++++++++++++ 2 files changed, 323 insertions(+) create mode 100644 zig/src/core/quant/widen.zig diff --git a/zig/src/core/quant/quant.zig b/zig/src/core/quant/quant.zig index 944bca7a3d..756fd87b90 100644 --- a/zig/src/core/quant/quant.zig +++ b/zig/src/core/quant/quant.zig @@ -10,6 +10,7 @@ pub const dense = @import("dense.zig"); /// FP8 and NVFP4 values and fp64 references; host only (no Format: their kernels' layouts are each backend's). pub const fp8 = @import("fp8.zig"); pub const nvfp4 = @import("nvfp4.zig"); +pub const widen = @import("widen.zig"); const Allocator = std.mem.Allocator; @@ -202,4 +203,5 @@ test { _ = dense; _ = fp8; _ = nvfp4; + _ = widen; } diff --git a/zig/src/core/quant/widen.zig b/zig/src/core/quant/widen.zig new file mode 100644 index 0000000000..1a26a56330 --- /dev/null +++ b/zig/src/core/quant/widen.zig @@ -0,0 +1,321 @@ +//! Exact widening of MLX affine projections: codes re-packed at more bits, group tables repeated for smaller groups. + +const std = @import("std"); +const mlx = @import("mlx.zig"); +const types = @import("types.zig"); + +const Allocator = std.mem.Allocator; +const Host = mlx.Host; +const Width = mlx.Width; +const Tensor = types.Tensor; + +pub const Error = error{ + /// The target has fewer bits than the source: codes would not fit. + NarrowsBits, + /// The target group does not divide the source group: one scale could not cover a target group. + GroupDoesNotDivide, + /// The target is not an affine width the kernels run (bits 2, 3, 4, 5, 6, 8; groups 32, 64, 128). + UnsupportedQuantization, + /// The words and tables disagree on shape, or fused members on K or table type. + UnexpectedTensor, +} || Allocator.Error; + +/// The width a set of members widens to exactly: the most bits and the smallest group among them. +pub fn common(widths: []const Width) Width { + std.debug.assert(widths.len > 0); + var w = widths[0]; + for (widths[1..]) |x| w = .{ .bits = @max(w.bits, x.bits), .group = @min(w.group, x.group) }; + return w; +} + +fn supported(w: Width) bool { + const bits_ok = std.mem.indexOfScalar(u8, &mlx.bit_widths, w.bits) != null; + return bits_ok and (w.group == 32 or w.group == 64 or w.group == 128); +} + +/// Rows of a host with its leading dimensions folded (a stack's experts times N), after checking the three agree. +fn foldedRows(h: Host) error{UnexpectedTensor}!usize { + const r = h.words.rank; + if (r < 2 or h.scales.rank != r or h.biases.rank != r) return error.UnexpectedTensor; + if (h.scales.dtype != h.biases.dtype) return error.UnexpectedTensor; + var n: usize = 1; + for (0..r - 1) |i| { + if (h.scales.shape[i] != h.words.shape[i] or h.biases.shape[i] != h.words.shape[i]) return error.UnexpectedTensor; + n *= h.words.shape[i]; + } + if (h.scales.shape[r - 1] != h.biases.shape[r - 1]) return error.UnexpectedTensor; + if (h.words.shape[r - 1] * 32 != h.scales.shape[r - 1] * h.group * h.bits) return error.UnexpectedTensor; + return n; +} + +/// The same dequantized values stored at width `to`: never fewer bits, and a group that divides the source's. +pub fn widen(a: Allocator, h: Host, to: Width) Error!Host { + if (!supported(to)) return error.UnsupportedQuantization; + const n = try foldedRows(h); + if (to.bits == h.bits and to.group == h.group) return h; + if (to.bits < h.bits) return error.NarrowsBits; + if (h.group % to.group != 0) return error.GroupDoesNotDivide; + const reps = h.group / to.group; + return .{ + .words = if (to.bits == h.bits) h.words else try repack(a, h.words, n, h.bits, to.bits), + .scales = if (reps == 1) h.scales else try repeat(a, h.scales, n, reps), + .biases = if (reps == 1) h.biases else try repeat(a, h.biases, n, reps), + .bits = to.bits, + .group = to.group, + }; +} + +/// Words of `bits`-wide codes re-packed at `wider` bits, code for code, low code first (`n` folded rows). +fn repack(a: Allocator, words: Tensor, n: usize, bits: u8, wider: u8) Allocator.Error!Tensor { + const r = words.rank; + const from_row = words.shape[r - 1]; + const k = from_row * 32 / bits; + const to_row = k * wider / 32; + const out = try a.alloc(u8, n * to_row * 4); + @memset(out, 0); + const mask: u64 = (@as(u64, 1) << @intCast(bits)) - 1; + for (0..n) |row| for (0..k) |c| { + const code = codeAt(words.bytes, row * from_row, c * bits) & mask; + const bit = c * wider; + const at = row * to_row + (bit >> 5); + const v = code << @intCast(bit & 31); + orWord(out, at, @truncate(v)); + if ((bit & 31) + wider > 32) orWord(out, at + 1, @truncate(v >> 32)); + }; + var t = words; + t.shape[r - 1] = to_row; + t.bytes = out; + return t; +} + +/// The (unmasked) code starting `bit` bits into the row of words beginning at word `base`. +fn codeAt(bytes: []const u8, base: usize, bit: usize) u64 { + const at = base + (bit >> 5); + const lo: u64 = std.mem.readInt(u32, bytes[4 * at ..][0..4], .little); + const hi: u64 = if ((bit & 31) != 0 and 4 * (at + 2) <= bytes.len) std.mem.readInt(u32, bytes[4 * (at + 1) ..][0..4], .little) else 0; + return (lo | hi << 32) >> @intCast(bit & 31); +} + +fn orWord(bytes: []u8, at: usize, v: u32) void { + const w = bytes[4 * at ..][0..4]; + std.mem.writeInt(u32, w, std.mem.readInt(u32, w, .little) | v, .little); +} + +/// A group table with every entry repeated `reps` times along its last dimension (`n` folded rows). +fn repeat(a: Allocator, t: Tensor, n: usize, reps: usize) Allocator.Error!Tensor { + const size = t.dtype.size(); + const groups = t.shape[t.rank - 1]; + const out = try a.alloc(u8, n * groups * reps * size); + for (0..n * groups) |g| { + const src = t.bytes[g * size ..][0..size]; + for (0..reps) |j| @memcpy(out[(g * reps + j) * size ..][0..size], src); + } + var o = t; + o.shape[t.rank - 1] = groups * reps; + o.bytes = out; + return o; +} + +/// One rank-2 projection of `hosts`' rows in order, all widened to their `common` width; K and tables must agree. +pub fn fuse(a: Allocator, hosts: []const Host) Error!Host { + if (hosts.len == 0) return error.UnexpectedTensor; + var to: Width = .{ .bits = hosts[0].bits, .group = hosts[0].group }; + for (hosts[1..]) |h| to = common(&.{ to, .{ .bits = h.bits, .group = h.group } }); + const parts = try a.alloc(Host, hosts.len); + var n: usize = 0; + for (hosts, parts) |h, *p| { + if (h.words.rank != 2 or h.scales.dtype != hosts[0].scales.dtype) return error.UnexpectedTensor; + if (h.scales.shape[1] * h.group != hosts[0].scales.shape[1] * hosts[0].group) return error.UnexpectedTensor; + p.* = try widen(a, h, to); + n += p.words.shape[0]; + } + if (hosts.len == 1) return parts[0]; + var out = parts[0]; + out.words = try concatRows(a, parts, .words, n); + out.scales = try concatRows(a, parts, .scales, n); + out.biases = try concatRows(a, parts, .biases, n); + return out; +} + +fn concatRows(a: Allocator, parts: []const Host, comptime field: std.meta.FieldEnum(Host), n: usize) Allocator.Error!Tensor { + var total: usize = 0; + for (parts) |p| total += @field(p, @tagName(field)).bytes.len; + const out = try a.alloc(u8, total); + var at: usize = 0; + for (parts) |p| { + const b = @field(p, @tagName(field)).bytes; + @memcpy(out[at..][0..b.len], b); + at += b.len; + } + var t = @field(parts[0], @tagName(field)); + t.shape[0] = n; + t.bytes = out; + return t; +} + +// ---- tests ---- + +const testing = std.testing; +const groups_all = [_]u16{ 32, 64, 128 }; + +fn tensor(dtype: types.DType, shape: []const usize, bytes: []const u8) Tensor { + var t: Tensor = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .bytes = bytes }; + @memcpy(t.shape[0..shape.len], shape); + return t; +} + +/// A host of `lead` x `k` codes of width `w` with bf16 tables: row 0 all zeros, row 1 all max, the rest random. +fn sample(a: Allocator, rng: std.Random, lead: []const usize, k: usize, w: Width) !Host { + var n: usize = 1; + for (lead) |d| n *= d; + const row = k * w.bits / 32; + const words = try a.alloc(u8, n * row * 4); + @memset(words, 0); + const max: u64 = (@as(u64, 1) << @intCast(w.bits)) - 1; + for (0..n) |r| for (0..k) |c| { + const code: u64 = switch (r % 4) { + 0 => 0, + 1 => max, + else => rng.int(u64) & max, + }; + const bit = c * w.bits; + const v = code << @intCast(bit & 31); + orWord(words, r * row + (bit >> 5), @truncate(v)); + if ((bit & 31) + w.bits > 32) orWord(words, r * row + (bit >> 5) + 1, @truncate(v >> 32)); + }; + const g = k / w.group; + const scales = try a.alloc(u8, n * g * 2); + const biases = try a.alloc(u8, n * g * 2); + for (0..n * g) |i| { + std.mem.writeInt(u16, scales[2 * i ..][0..2], bf16Of(rng.floatNorm(f32) * 0.02), .little); + std.mem.writeInt(u16, biases[2 * i ..][0..2], bf16Of(rng.floatNorm(f32) * 0.1), .little); + } + var ws: [5]usize = undefined; + var ts: [5]usize = undefined; + @memcpy(ws[0..lead.len], lead); + @memcpy(ts[0..lead.len], lead); + ws[lead.len] = row; + ts[lead.len] = g; + const r = lead.len + 1; + return .{ + .words = tensor(.i32, ws[0..r], words), + .scales = tensor(.bf16, ts[0..r], scales), + .biases = tensor(.bf16, ts[0..r], biases), + .bits = w.bits, + .group = w.group, + }; +} + +fn bf16Of(x: f32) u16 { + return @truncate(@as(u32, @bitCast(x)) >> 16); +} + +fn dequantized(a: Allocator, h: Host) ![]f64 { + const out = try a.alloc(f64, mlx.elements(h)); + try mlx.reference(h, out); + return out; +} + +fn expectBits(want: []const f64, got: []const f64) !void { + try testing.expectEqual(want.len, got.len); + for (want, got) |x, y| try testing.expectEqual(@as(u64, @bitCast(x)), @as(u64, @bitCast(y))); +} + +test "every width widens to every wider one with the same fp64 values, bit for bit" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var prng = std.Random.DefaultPrng.init(0x51de); + const rng = prng.random(); + var checked: usize = 0; + for (mlx.bit_widths) |bits| for (groups_all) |group| { + const from: Width = .{ .bits = bits, .group = group }; + const h = try sample(a, rng, &.{6}, 256, from); + const want = try dequantized(a, h); + for (mlx.bit_widths) |to_bits| for (groups_all) |to_group| { + if (to_bits < bits or group % to_group != 0) continue; + const to: Width = .{ .bits = to_bits, .group = to_group }; + const w = try widen(a, h, to); + try testing.expectEqual(to_bits, w.bits); + try testing.expectEqual(to_group, w.group); + try testing.expectEqual(@as(usize, 256) * to_bits / 32, w.words.shape[1]); + try testing.expectEqual(@as(usize, 256) / to_group, w.scales.shape[1]); + try expectBits(want, try dequantized(a, w)); + checked += 1; + }; + }; + // 18 sources; each reaches the wider-or-equal bits times the groups dividing its own + try testing.expectEqual(@as(usize, 126), checked); +} + +test "a stack of experts widens row for row" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var prng = std.Random.DefaultPrng.init(0xe0e1); + const rng = prng.random(); + const h = try sample(a, rng, &.{ 3, 4 }, 384, .{ .bits = 3, .group = 128 }); + const w = try widen(a, h, .{ .bits = 5, .group = 32 }); + try testing.expectEqual(@as(u8, 3), w.words.rank); + try testing.expectEqualSlices(usize, &.{ 3, 4, 384 * 5 / 32 }, w.words.shape[0..3]); + try testing.expectEqualSlices(usize, &.{ 3, 4, 12 }, w.biases.shape[0..3]); + try expectBits(try dequantized(a, h), try dequantized(a, w)); +} + +test "the same width is returned as it is" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var prng = std.Random.DefaultPrng.init(1); + const h = try sample(a, prng.random(), &.{2}, 128, .{ .bits = 4, .group = 64 }); + const w = try widen(a, h, .{ .bits = 4, .group = 64 }); + try testing.expectEqual(h.words.bytes.ptr, w.words.bytes.ptr); + try testing.expectEqual(h.scales.bytes.ptr, w.scales.bytes.ptr); +} + +test "narrowing bits, growing groups and widths the kernels do not run are refused" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var prng = std.Random.DefaultPrng.init(2); + const rng = prng.random(); + const h = try sample(a, rng, &.{2}, 256, .{ .bits = 5, .group = 64 }); + try testing.expectError(error.NarrowsBits, widen(a, h, .{ .bits = 4, .group = 64 })); + try testing.expectError(error.NarrowsBits, widen(a, h, .{ .bits = 2, .group = 32 })); + try testing.expectError(error.GroupDoesNotDivide, widen(a, h, .{ .bits = 5, .group = 128 })); + try testing.expectError(error.GroupDoesNotDivide, widen(a, h, .{ .bits = 8, .group = 128 })); + try testing.expectError(error.UnsupportedQuantization, widen(a, h, .{ .bits = 7, .group = 64 })); + try testing.expectError(error.UnsupportedQuantization, widen(a, h, .{ .bits = 8, .group = 16 })); + try testing.expectError(error.UnsupportedQuantization, widen(a, h, .{ .bits = 8, .group = 48 })); + var bad = h; + bad.bits = 4; + try testing.expectError(error.UnexpectedTensor, widen(a, bad, .{ .bits = 8, .group = 32 })); +} + +test "the common width takes the most bits and the smallest group" { + try testing.expectEqual(Width{ .bits = 5, .group = 32 }, common(&.{ .{ .bits = 4, .group = 64 }, .{ .bits = 5, .group = 128 }, .{ .bits = 2, .group = 32 } })); + try testing.expectEqual(Width{ .bits = 6, .group = 32 }, common(&.{.{ .bits = 6, .group = 32 }})); +} + +test "fused members keep each one's values in order at the common width" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var prng = std.Random.DefaultPrng.init(0xf05e); + const rng = prng.random(); + const k = 512; + const members = [_]Host{ + try sample(a, rng, &.{4}, k, .{ .bits = 4, .group = 64 }), + try sample(a, rng, &.{3}, k, .{ .bits = 5, .group = 128 }), + try sample(a, rng, &.{5}, k, .{ .bits = 2, .group = 32 }), + }; + const f = try fuse(a, &members); + try testing.expectEqual(@as(u8, 5), f.bits); + try testing.expectEqual(@as(u16, 32), f.group); + try testing.expectEqual(@as(usize, 12), f.words.shape[0]); + var want: std.ArrayList(f64) = .empty; + for (members) |m| try want.appendSlice(a, try dequantized(a, m)); + try expectBits(want.items, try dequantized(a, f)); + const other = try sample(a, rng, &.{2}, 256, .{ .bits = 4, .group = 64 }); + try testing.expectError(error.UnexpectedTensor, fuse(a, &.{ members[0], other })); +} From dcebc462b94a9c563c1d8164b2494d45411ac829 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 04:08:50 -1000 Subject: [PATCH 03/10] The row projection reads 3- and 5-bit MLX codes and can write into a wider output Chips without tensor units now decode every MLX affine width (2, 3, 4, 5, 6, 8 bits) in the row kernel: 3- and 5-bit rows unpack eight codes from three or five bytes, low code first, as MLX packs them. The 2-, 4-, 6- and 8-bit paths are unchanged. A 3- or 5-bit row must end on a word, as the checkpoint stores it. callStrided launches the row kernel with an output row stride, so one member of a fused stack (one run of equal width) writes its columns straight into the stack output. The default call, its dims and its kernels are unchanged. tf-row-widths-check (zig/build/core_checks.zig) runs it on the GPU: each width at groups of 32, 64 and 128 against an fp64 product from the MLX reference dequantization, within the fp32-sum bound; one row alone against the same row in 2-, 3-, 8- and 16-row windows, relu2, indexed slots and a strided output, bit for bit; a fused stack of 48-column members at mixed widths (4-bit among them, group 64) written through the stride at windows up to 48 rows, each member as its own call; and, given the unchanged kernel source, the widths it already served byte for byte against it. Tests that fail before: "shapes the row kernels serve" now validates 3 and 5 bits (main refuses them), "3- and 5-bit rows take whole words of codes", "every width has its own pipeline slot" and "a strided call writes rows at least a matrix wide" (no callStrided on main). tf-row-widths-check on an M3 Ultra (no tensor units), given main's row_projection.metal: PASS 64512 values within the fp32-sum bound (worst err/bound 0.9400), 408192 bit-equal pairs, 301056 outputs byte-equal to main's kernels at the widths main serves (3 and 5 bits have no main kernel to compare). Gate on main 0efbd447 (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test: `get` logs an error on the DeferredTensor path the test expects). lean_check: 536 problems (main 539), no new findings. --- build.zig | 2 + zig/build/core_checks.zig | 18 ++ zig/kernels/metal/core/row_projection.metal | 56 +++- zig/src/core/row_projection.zig | 50 ++- zig/tests/row_widths_check.zig | 324 ++++++++++++++++++++ 5 files changed, 439 insertions(+), 11 deletions(-) create mode 100644 zig/build/core_checks.zig create mode 100644 zig/tests/row_widths_check.zig diff --git a/build.zig b/build.zig index 1bf8bfa260..2ba6052374 100644 --- a/build.zig +++ b/build.zig @@ -3,6 +3,7 @@ const builtin = @import("builtin"); const cuda_build = @import("zig/build/cuda.zig"); const dist_build = @import("zig/build/dist.zig"); const qwen27_build = @import("zig/build/qwen27.zig"); +const core_checks = @import("zig/build/core_checks.zig"); const hip_build = @import("zig/build/hip.zig"); const grammar_build = @import("zig/build/grammar.zig"); const counters_build = @import("zig/build/metal_counters.zig"); @@ -455,6 +456,7 @@ fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "row_projection", .module = row_mod } }, }) }); b.step("tf-row-check", "The core row projection on synthetic 4-bit matrices: widths bit-equal, a CPU reference, indexed experts").dependOn(&b.addInstallArtifact(row_check, .{}).step); + core_checks.rowWidths(b, target, optimize, metal, row_mod); test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = row_mod })).step); exl3_build.targets(b, target, optimize, metal, test_step); exlr8_build.metalTargets(b, target, optimize, metal, mods.sources, test_step); diff --git a/zig/build/core_checks.zig b/zig/build/core_checks.zig new file mode 100644 index 0000000000..75f39c00c2 --- /dev/null +++ b/zig/build/core_checks.zig @@ -0,0 +1,18 @@ +//! The core's device checks on synthetic weights: programs that run on the GPU, each under its own step. + +const std = @import("std"); + +/// tf-row-widths-check: the row projection at every MLX width and group against float64, windows, strides, stock bytes. +pub fn rowWidths(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, metal: *std.Build.Module, row_mod: *std.Build.Module) void { + const tokenizer = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = optimize, .link_libc = true }); + const core = b.createModule(.{ .root_source_file = b.path("zig/src/core/root.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "tokenizer", .module = tokenizer }} }); + const exe = b.addExecutable(.{ .name = "tf-row-widths-check", .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/row_widths_check.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "row_projection", .module = row_mod }, .{ .name = "core", .module = core } }, + }) }); + const about = "The row projection at every MLX width and group: fp64 reference, windows, strides, stock bytes"; + b.step("tf-row-widths-check", about).dependOn(&b.addInstallArtifact(exe, .{}).step); +} diff --git a/zig/kernels/metal/core/row_projection.metal b/zig/kernels/metal/core/row_projection.metal index dae395aeb0..5a503e3acc 100644 --- a/zig/kernels/metal/core/row_projection.metal +++ b/zig/kernels/metal/core/row_projection.metal @@ -1,4 +1,4 @@ -// Affine 2-, 4-, 6- and 8-bit row projections for chips without tensor units: RPS columns a simdgroup, MB rows a weight read, each row its own fp32 sums. +// Affine 2- to 8-bit row projections without tensor units: RPS columns a simdgroup, MB rows a weight read. #include using namespace metal; @@ -18,7 +18,7 @@ inline float tf_rp_load16(const device bfloat* x, thread float* xt) { return sum; } -// MB rows' dots with one weight word group (16 codes; 6-bit: four values in three bytes), the codes unpacked to floats once; each row's sums in one order. +// MB rows' dots with 16 codes (3, 5, 6 bits: 8, 8, 4 codes in 3, 5, 3 bytes), unpacked once; sums in one order. template inline void tf_rp_qdot16(const device uint8_t* w, const thread float (*xt)[16], float scale, float bias, const thread float* sum, thread float* out) { @@ -40,6 +40,20 @@ inline void tf_rp_qdot16(const device uint8_t* w, const thread float (*xt)[16], q[4 * i + 2] = float(((p[1] >> 4) & 0x0f) | ((p[2] & 0x03) << 4)); q[4 * i + 3] = float((p[2] >> 2) & 0x3f); } + } else if (BITS == 3) { + for (int i = 0; i < 2; i++) { + const device uint8_t* p = w + 3 * i; + const uint v = uint(p[0]) | (uint(p[1]) << 8) | (uint(p[2]) << 16); + for (int j = 0; j < 8; j++) q[8 * i + j] = float((v >> (3 * j)) & 7u); + } + } else if (BITS == 5) { + for (int i = 0; i < 2; i++) { + const device uint8_t* p = w + 5 * i; + const uint lo = uint(p[0]) | (uint(p[1]) << 8) | (uint(p[2]) << 16) | (uint(p[3]) << 24), hi = p[4]; + for (int j = 0; j < 6; j++) q[8 * i + j] = float((lo >> (5 * j)) & 31u); + q[8 * i + 6] = float(((lo >> 30) | (hi << 2)) & 31u); + q[8 * i + 7] = float(hi >> 3); + } } else { for (int i = 0; i < 16; i++) q[i] = float(w[i]); } @@ -93,8 +107,8 @@ inline bfloat tf_rp_out(float v) { return bfloat(v); } -// dims: K, N, GS, rows. X [rows, K] bf16; W [N, K*BITS/32] words; S, B [N, K/GS] bf16; OUT [rows, N] bf16. -template +// dims: K, N, GS, rows, OUT's row stride when STRIDED (else N). X [rows, K]; W [N, K*BITS/32] words; S, B [N, K/GS]. +template [[kernel]] void tf_row_projection(const device bfloat* X [[buffer(0)]], const device uint32_t* W [[buffer(1)]], const device bfloat* S [[buffer(2)]], @@ -104,7 +118,7 @@ template uint simdgroup_index_in_threadgroup [[simdgroup_index_in_threadgroup]], uint thread_index_in_simdgroup [[thread_index_in_simdgroup]], uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]]) { - const int K = dims[0], N = dims[1], GS = dims[2], rows = dims[3]; + const int K = dims[0], N = dims[1], GS = dims[2], rows = dims[3], LD = STRIDED ? dims[4] : N; const uint lane = thread_index_in_simdgroup; const int row0 = (int(threadgroup_position_in_grid.x) * SG + int(simdgroup_index_in_threadgroup)) * RPS; if (row0 >= N) return; @@ -119,14 +133,14 @@ template tf_rp_rows(wr, sr, br, X + size_t(m) * K, K, GS, lane, acc); if (lane == 0) for (int b = 0; b < MB; b++) - for (int j = 0; j < RPS; j++) OUT[size_t(m + b) * N + row0 + j] = tf_rp_out(acc[b][j]); + for (int j = 0; j < RPS; j++) OUT[size_t(m + b) * LD + row0 + j] = tf_rp_out(acc[b][j]); } #pragma clang loop unroll(disable) for (; m < rows; m++) { float acc[1][RPS]; tf_rp_rows(wr, sr, br, X + size_t(m) * K, K, GS, lane, acc); if (lane == 0) - for (int j = 0; j < RPS; j++) OUT[size_t(m) * N + row0 + j] = tf_rp_out(acc[0][j]); + for (int j = 0; j < RPS; j++) OUT[size_t(m) * LD + row0 + j] = tf_rp_out(acc[0][j]); } } @@ -195,3 +209,31 @@ template [[host_name("tf_row_projection_q8_bf16")]] [[kernel]] decltype(tf_row_p template [[host_name("tf_row_projection_indexed_q8_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 8, true>) tf_row_projection_indexed<4, 2, false, 8, true>; template [[host_name("tf_row_projection_relu2_q8_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 8, true>) tf_row_projection<4, 2, 2, true, 8, true>; template [[host_name("tf_row_projection_indexed_relu2_q8_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 8, true>) tf_row_projection_indexed<4, 2, true, 8, true>; +template [[host_name("tf_row_projection_q3_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 3, false>) tf_row_projection<4, 2, 2, false, 3, false>; +template [[host_name("tf_row_projection_indexed_q3_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 3, false>) tf_row_projection_indexed<4, 2, false, 3, false>; +template [[host_name("tf_row_projection_relu2_q3_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 3, false>) tf_row_projection<4, 2, 2, true, 3, false>; +template [[host_name("tf_row_projection_indexed_relu2_q3_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 3, false>) tf_row_projection_indexed<4, 2, true, 3, false>; +template [[host_name("tf_row_projection_q3_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 3, true>) tf_row_projection<4, 2, 2, false, 3, true>; +template [[host_name("tf_row_projection_indexed_q3_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 3, true>) tf_row_projection_indexed<4, 2, false, 3, true>; +template [[host_name("tf_row_projection_relu2_q3_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 3, true>) tf_row_projection<4, 2, 2, true, 3, true>; +template [[host_name("tf_row_projection_indexed_relu2_q3_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 3, true>) tf_row_projection_indexed<4, 2, true, 3, true>; +template [[host_name("tf_row_projection_q5_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 5, false>) tf_row_projection<4, 2, 2, false, 5, false>; +template [[host_name("tf_row_projection_indexed_q5_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 5, false>) tf_row_projection_indexed<4, 2, false, 5, false>; +template [[host_name("tf_row_projection_relu2_q5_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 5, false>) tf_row_projection<4, 2, 2, true, 5, false>; +template [[host_name("tf_row_projection_indexed_relu2_q5_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 5, false>) tf_row_projection_indexed<4, 2, true, 5, false>; +template [[host_name("tf_row_projection_q5_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 5, true>) tf_row_projection<4, 2, 2, false, 5, true>; +template [[host_name("tf_row_projection_indexed_q5_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 5, true>) tf_row_projection_indexed<4, 2, false, 5, true>; +template [[host_name("tf_row_projection_relu2_q5_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 5, true>) tf_row_projection<4, 2, 2, true, 5, true>; +template [[host_name("tf_row_projection_indexed_relu2_q5_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 5, true>) tf_row_projection_indexed<4, 2, true, 5, true>; +template [[host_name("tf_row_projection_strided_q2_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 2, false, true>) tf_row_projection<4, 2, 2, false, 2, false, true>; +template [[host_name("tf_row_projection_strided_q2_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 2, true, true>) tf_row_projection<4, 2, 2, false, 2, true, true>; +template [[host_name("tf_row_projection_strided_q3_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 3, false, true>) tf_row_projection<4, 2, 2, false, 3, false, true>; +template [[host_name("tf_row_projection_strided_q3_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 3, true, true>) tf_row_projection<4, 2, 2, false, 3, true, true>; +template [[host_name("tf_row_projection_strided_q4_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 4, false, true>) tf_row_projection<4, 2, 2, false, 4, false, true>; +template [[host_name("tf_row_projection_strided_q4_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 4, true, true>) tf_row_projection<4, 2, 2, false, 4, true, true>; +template [[host_name("tf_row_projection_strided_q5_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 5, false, true>) tf_row_projection<4, 2, 2, false, 5, false, true>; +template [[host_name("tf_row_projection_strided_q5_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 5, true, true>) tf_row_projection<4, 2, 2, false, 5, true, true>; +template [[host_name("tf_row_projection_strided_q6_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 6, false, true>) tf_row_projection<4, 2, 2, false, 6, false, true>; +template [[host_name("tf_row_projection_strided_q6_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 6, true, true>) tf_row_projection<4, 2, 2, false, 6, true, true>; +template [[host_name("tf_row_projection_strided_q8_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 8, false, true>) tf_row_projection<4, 2, 2, false, 8, false, true>; +template [[host_name("tf_row_projection_strided_q8_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 8, true, true>) tf_row_projection<4, 2, 2, false, 8, true, true>; diff --git a/zig/src/core/row_projection.zig b/zig/src/core/row_projection.zig index 2bd1cba34a..d28555029c 100644 --- a/zig/src/core/row_projection.zig +++ b/zig/src/core/row_projection.zig @@ -1,4 +1,4 @@ -//! The lane projection's path for chips without tensor units: row-layout 2-, 4- and 8-bit weights, one read a window, each row its own fp32 sums. +//! The lane projection's path for chips without tensor units: row-layout 2- to 8-bit weights, one read a window. const std = @import("std"); const mtl = @import("metal"); const sources = @import("kernel_sources"); @@ -14,7 +14,7 @@ pub const columns = 4; pub const groups = 2; /// The code widths the kernels unpack. -pub const bit_widths = [_]u8{ 2, 4, 6, 8 }; +pub const bit_widths = [_]u8{ 2, 3, 4, 5, 6, 8 }; /// A group's sum of inputs for the bias term: fp32 in order (the tensor path's xsum), or rounded to bf16 at each add (Nemotron's expert kernels). pub const Sum = enum { f32, bf16 }; @@ -34,7 +34,7 @@ pub const Weights = struct { sum: Sum = .f32, pub fn validate(x: Weights) !void { - if (x.n == 0 or x.n % (columns * groups) != 0 or x.k == 0 or x.k % 16 != 0 or x.group < 16 or x.group % 16 != 0 or x.k % x.group != 0 or 512 % x.group != 0) return error.UnsupportedRowShape; + if (x.n == 0 or x.n % (columns * groups) != 0 or x.k == 0 or x.k % 16 != 0 or x.group < 16 or x.group % 16 != 0 or x.k % x.group != 0 or 512 % x.group != 0 or x.k * x.bits % 32 != 0) return error.UnsupportedRowShape; if (std.mem.indexOfScalar(u8, &bit_widths, x.bits) == null) return error.UnsupportedRowBits; } @@ -63,6 +63,7 @@ pub const Pipelines = struct { relu2: [variants]mtl.Pipeline, indexed: [variants]mtl.Pipeline, indexed_relu2: [variants]mtl.Pipeline, + strided: [variants]mtl.Pipeline, pub fn load(device: mtl.Device) !Pipelines { const pool = mtl.objc.Pool.push(); @@ -76,6 +77,7 @@ pub const Pipelines = struct { p.relu2[i] = try mtl.Pipeline.init(device, lib, std.fmt.comptimePrint("tf_row_projection_relu2_q{d}_{s}", .{ bits, sum }), false); p.indexed[i] = try mtl.Pipeline.init(device, lib, std.fmt.comptimePrint("tf_row_projection_indexed_q{d}_{s}", .{ bits, sum }), false); p.indexed_relu2[i] = try mtl.Pipeline.init(device, lib, std.fmt.comptimePrint("tf_row_projection_indexed_relu2_q{d}_{s}", .{ bits, sum }), false); + p.strided[i] = try mtl.Pipeline.init(device, lib, std.fmt.comptimePrint("tf_row_projection_strided_q{d}_{s}", .{ bits, sum }), false); }; return p; } @@ -85,6 +87,7 @@ pub const Pipelines = struct { for (p.relu2) |x| x.deinit(); for (p.indexed) |x| x.deinit(); for (p.indexed_relu2) |x| x.deinit(); + for (p.strided) |x| x.deinit(); } }; @@ -99,6 +102,14 @@ pub fn call(p: *const Pipelines, w: Weights, rows: usize, relu2: bool) !Call { return .{ .pipeline = if (relu2) p.relu2[i] else p.plain[i], .dims = .{ @intCast(w.k), @intCast(w.n), @intCast(w.group), @intCast(rows), 1, 1 }, .groups = w.n / (columns * groups) }; } +/// `call` with output row m at y + m * `stride` values: a fused stack member writes its columns of the rows. +pub fn callStrided(p: *const Pipelines, w: Weights, rows: usize, stride: usize) !Call { + try w.validate(); + if (rows == 0) return error.NoRows; + if (stride < w.n or stride > std.math.maxInt(i32)) return error.UnsupportedRowStride; + return .{ .pipeline = p.strided[w.variant()], .dims = .{ @intCast(w.k), @intCast(w.n), @intCast(w.group), @intCast(rows), @intCast(stride), 1 }, .groups = w.n / (columns * groups) }; +} + /// y[slots, n] = x[slot / repeat] W[ids[slot]]^T; the grid's y is the slot. pub fn callIndexed(p: *const Pipelines, x: Experts, slots: usize, relu2: bool) !Call { try x.validate(); @@ -127,5 +138,36 @@ test "shapes the row kernels serve" { try (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 256, .k = 2688, .bits = 8 }).validate(); try (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 256, .k = 2688, .bits = 2 }).validate(); try (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 256, .k = 2688, .bits = 6 }).validate(); - try std.testing.expectError(error.UnsupportedRowBits, (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 256, .k = 2688, .bits = 3 }).validate()); + try (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 256, .k = 2688, .bits = 3 }).validate(); + try (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 256, .k = 2688, .bits = 5 }).validate(); + try std.testing.expectError(error.UnsupportedRowBits, (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 256, .k = 2688, .bits = 7 }).validate()); +} + +test "3- and 5-bit rows take whole words of codes" { + for ([_]u8{ 3, 5 }) |bits| for ([_]usize{ 32, 64, 128 }) |group| { + try (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 256, .k = 2688, .group = group, .bits = bits }).validate(); + try (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 5120, .k = 17408, .group = group, .bits = bits }).validate(); + }; + // 48 codes of 3 or 5 bits end inside a word: the checkpoint cannot hold that row + try std.testing.expectError(error.UnsupportedRowShape, (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 256, .k = 48, .group = 16, .bits = 3 }).validate()); + try std.testing.expectError(error.UnsupportedRowShape, (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 256, .k = 48, .group = 16, .bits = 5 }).validate()); + try (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 256, .k = 48, .group = 16, .bits = 6 }).validate(); +} + +test "every width has its own pipeline slot" { + var seen: [variants]bool = @splat(false); + for (bit_widths) |bits| for ([_]Sum{ .f32, .bf16 }) |sum| { + const i = (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 8, .k = 64, .bits = bits, .sum = sum }).variant(); + try std.testing.expect(!seen[i]); + seen[i] = true; + }; +} + +test "a strided call writes rows at least a matrix wide" { + const p: Pipelines = undefined; + const w = Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 64, .k = 128, .bits = 5 }; + try std.testing.expectError(error.UnsupportedRowStride, callStrided(&p, w, 2, 63)); + try std.testing.expectError(error.NoRows, callStrided(&p, w, 0, 64)); + const c = try callStrided(&p, w, 2, 192); + try std.testing.expectEqual([6]i32{ 128, 64, 64, 2, 192, 1 }, c.dims); } diff --git a/zig/tests/row_widths_check.zig b/zig/tests/row_widths_check.zig new file mode 100644 index 0000000000..327f7bc3d1 --- /dev/null +++ b/zig/tests/row_widths_check.zig @@ -0,0 +1,324 @@ +//! The core row projection at every MLX width and group: an fp64 reference, windows, strides and stock bytes. +const std = @import("std"); +const mtl = @import("metal"); +const row = @import("row_projection"); +const mlx = @import("core").quant.mlx; + +const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; +const a = std.heap.page_allocator; +const usage = "usage: tf-row-widths-check [STOCK.metal] (STOCK: the unchanged row_projection.metal for the byte comparison)"; + +const groups_tested = [_]usize{ 32, 64, 128 }; +const Shape = struct { n: usize, k: usize }; +// all tail (K < 512), full passes and a tail, full passes only +const shapes = [_]Shape{ .{ .n = 64, .k = 384 }, .{ .n = 128, .k = 2688 }, .{ .n = 32, .k = 1536 } }; +const rows = 16; +// a fused stack's windows reach this many rows +const x_rows = 48; +const sentinel: u16 = 0xffff; + +fn toBf(v: f32) u16 { + const u: u32 = @bitCast(v); + return @truncate((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); +} +fn fromBf(v: u16) f32 { + return @bitCast(@as(u32, v) << 16); +} + +/// One synthetic MLX matrix in the row layout, its inputs and its fp64 dequantization (and codes alone). +const Mat = struct { + w: row.Weights, + x: mtl.Buffer, + deq: []f64, + codes: []f64, + + fn init(device: mtl.Device, s: Shape, bits: u8, group: usize, sum: row.Sum, seed: u64) !Mat { + var prng = std.Random.DefaultPrng.init(seed); + const r = prng.random(); + const words = s.n * s.k * bits / 32; + const tables = s.n * (s.k / group); + const m = Mat{ + .w = .{ .w = try device.buffer(words * 4, opts), .scales = try device.buffer(tables * 2, opts), .biases = try device.buffer(tables * 2, opts), .n = s.n, .k = s.k, .group = group, .bits = bits, .sum = sum }, + .x = try device.buffer(x_rows * s.k * 2, opts), + .deq = try a.alloc(f64, s.n * s.k), + .codes = try a.alloc(f64, s.n * s.k), + }; + for (m.w.w.slice(u32, words)) |*v| v.* = r.int(u32); + for (m.w.scales.slice(u16, tables)) |*v| v.* = toBf(r.float(f32) * 0.02 + 0.001); + for (m.w.biases.slice(u16, tables)) |*v| v.* = toBf((r.float(f32) - 0.5) * 0.1); + for (m.x.slice(u16, x_rows * s.k)) |*v| v.* = toBf((r.float(f32) - 0.5) * 2); + try mlx.reference(host(m.w, m.w.scales.slice(u16, tables), m.w.biases.slice(u16, tables)), m.deq); + const ones = try a.alloc(u16, tables); + defer a.free(ones); + const zeros = try a.alloc(u16, tables); + defer a.free(zeros); + @memset(ones, toBf(1)); + @memset(zeros, 0); + try mlx.reference(host(m.w, ones, zeros), m.codes); + return m; + } + + fn deinit(m: Mat) void { + for ([_]mtl.Buffer{ m.w.w, m.w.scales, m.w.biases, m.x }) |b| b.deinit(); + a.free(m.deq); + a.free(m.codes); + } + + /// The checkpoint's view of the same words and tables, as `mlx.reference` reads it. + fn host(w: row.Weights, scales: []const u16, biases: []const u16) mlx.Host { + const groups = w.k / w.group; + return .{ + .words = .{ .dtype = .i32, .rank = 2, .shape = .{ w.n, w.k * w.bits / 32, 1, 1, 1 }, .bytes = w.w.slice(u8, w.n * w.k * w.bits / 8) }, + .scales = .{ .dtype = .bf16, .rank = 2, .shape = .{ w.n, groups, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(scales) }, + .biases = .{ .dtype = .bf16, .rank = 2, .shape = .{ w.n, groups, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(biases) }, + .bits = w.bits, + .group = @intCast(w.group), + }; + } + + /// Row m against column c in fp64, and the fp32 sums' error bound (output rounding plus k-deep accumulation). + fn reference(m: Mat, i: usize, c: usize) [2]f64 { + const k = m.w.k; + const g_n = k / m.w.group; + const sc = m.w.scales.slice(u16, m.w.n * g_n); + const bi = m.w.biases.slice(u16, m.w.n * g_n); + const x = m.x.slice(u16, rows * k); + var y: f64 = 0; + var terms: f64 = 0; + for (0..g_n) |g| { + var dot: f64 = 0; + var xs: f64 = 0; + for (0..m.w.group) |j| { + const kk = g * m.w.group + j; + const xv: f64 = fromBf(x[i * k + kk]); + y += xv * m.deq[c * k + kk]; + dot += @abs(xv * m.codes[c * k + kk]); + xs += @abs(xv); + } + terms += @abs(@as(f64, fromBf(sc[c * g_n + g]))) * dot + @abs(@as(f64, fromBf(bi[c * g_n + g]))) * xs; + } + const u32_round = std.math.ldexp(@as(f64, 1), -24); + const gamma = @as(f64, @floatFromInt(k + 64)) * u32_round; + return .{ y, std.math.ldexp(@abs(y) + gamma * terms, -8) + gamma * terms }; + } +}; + +/// One launch: buffers x, W, S, B, dims, out at 0-5 and ids at 6 when given; grid (groups, slots). +const Launch = struct { pipeline: mtl.Pipeline, dims: [6]i32, groups: usize, slots: usize = 1, x_off: usize = 0, out_off: usize = 0, ids: ?mtl.Buffer = null }; + +fn run(queue: mtl.Queue, l: Launch, m: Mat, out: mtl.Buffer) !void { + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + e.setPipeline(l.pipeline); + e.setBuffer(m.x, l.x_off, 0); + e.setBuffer(m.w.w, 0, 1); + e.setBuffer(m.w.scales, 0, 2); + e.setBuffer(m.w.biases, 0, 3); + e.setValue(l.dims, 4); + e.setBuffer(out, l.out_off, 5); + if (l.ids) |ids| e.setBuffer(ids, 0, 6); + e.dispatchGroups(mtl.Size.of(l.groups, l.slots, 1), mtl.Size.of(32 * row.groups, 1, 1)); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.debug.print("GPU failed: {s}\n", .{msg}); + return error.GpuFailed; + } +} + +fn of(c: row.Call) Launch { + return .{ .pipeline = c.pipeline, .dims = c.dims, .groups = c.groups, .slots = c.slots }; +} + +/// The stock library's pipelines for the widths it already served, by name. +const Stock = struct { + lib: mtl.Library, + fn get(s: Stock, device: mtl.Device, kind: []const u8, bits: u8, sum: row.Sum) !mtl.Pipeline { + var buf: [96]u8 = undefined; + const name = try std.fmt.bufPrint(&buf, "tf_row_projection{s}_q{d}_{s}", .{ kind, bits, @tagName(sum) }); + return mtl.Pipeline.init(device, s.lib, name, false); + } +}; + +const Tally = struct { worst_err: f64 = 0, worst_ratio: f64 = 0, worst_col: usize = 0, values: usize = 0, exact: usize = 0, pairs: usize = 0, stock: usize = 0 }; + +pub fn main(init: std.process.Init) !void { + const arena = init.arena.allocator(); + const args = try init.minimal.args.toSlice(arena); + if (args.len > 2) { + std.debug.print("{s}\n", .{usage}); + std.process.exit(2); + } + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try row.Pipelines.load(device); + defer p.deinit(); + var stock: ?Stock = null; + if (args.len == 2) { + const f = try mtl.MappedFile.open(try arena.dupeSentinel(u8, args[1], 0)); + defer f.deinit(); + stock = .{ .lib = try mtl.Library.fromSource(device, f.bytes[0..f.size], mtl.CompileOptions.mlx()) }; + } + defer if (stock) |s| s.lib.deinit(); + const out = try device.buffer(rows * 512 * 2 * 4, opts); + defer out.deinit(); + const one = try device.buffer(rows * 512 * 2 * 4, opts); + defer one.deinit(); + const ids = try device.buffer(rows * 4, opts); + defer ids.deinit(); + var all = Tally{}; + for (row.bit_widths) |bits| for (groups_tested) |group| { + var t = Tally{}; + for (shapes, 0..) |s, si| for ([_]row.Sum{ .f32, .bf16 }) |sum| { + const m = try Mat.init(device, s, bits, group, sum, 1000 * @as(u64, bits) + 10 * group + si); + defer m.deinit(); + try check(device, queue, &p, stock, m, out, one, ids, &t); + }; + std.debug.print("q{d} g{d}: {d} values vs fp64, worst |err| {e:.3} (col {d}), worst err/bound {d:.4}, bf16(ref) exact {d}/{d}; {d} bit-equal pairs; {d} stock bytes equal\n", .{ bits, group, t.values, t.worst_err, t.worst_col, t.worst_ratio, t.exact, t.values, t.pairs, t.stock }); + all.values += t.values; + all.pairs += t.pairs; + all.stock += t.stock; + all.worst_ratio = @max(all.worst_ratio, t.worst_ratio); + }; + all.pairs += try stack(device, queue, &p); + if (stock == null) std.debug.print("no stock source given: the byte comparison was not run\n", .{}); + std.debug.print("PASS {d} values within the fp32-sum bound (worst err/bound {d:.4}), {d} bit-equal pairs, {d} stock outputs byte-equal\n", .{ all.values, all.worst_ratio, all.pairs, all.stock }); +} + +/// One matrix: the reference (f32 sums), windows, relu2, indexed and strided outputs against one row alone, then stock. +fn check(device: mtl.Device, queue: mtl.Queue, p: *const row.Pipelines, stock: ?Stock, m: Mat, out: mtl.Buffer, one: mtl.Buffer, ids: mtl.Buffer, t: *Tally) !void { + const n = m.w.n; + const k = m.w.k; + // each row alone: the reference everything else is compared with + const alone = try a.alloc(u16, rows * n); + defer a.free(alone); + for (0..rows) |i| { + try run(queue, launchAt(try row.call(p, m.w, 1, false), i * k * 2, 0), m, one); + @memcpy(alone[i * n ..][0..n], one.slice(u16, n)); + } + if (m.w.sum == .f32) for (0..rows) |i| for (0..n) |c| { + const ref = m.reference(i, c); + const got: f64 = fromBf(alone[i * n + c]); + const err = @abs(got - ref[0]); + if (err > ref[1]) { + std.debug.print("q{d} g{d} n {d} k {d} row {d} col {d}: GPU {d} fp64 {d} bound {e}\n", .{ m.w.bits, m.w.group, n, k, i, c, got, ref[0], ref[1] }); + return error.ReferenceMismatch; + } + if (err > t.worst_err) { + t.worst_err = err; + t.worst_col = c; + } + t.worst_ratio = @max(t.worst_ratio, err / ref[1]); + if (alone[i * n + c] == toBf(@floatCast(ref[0]))) t.exact += 1; + t.values += 1; + }; + // windows of 2, 3, 8 and 16 rows (ending at the last row): each row as it was alone + for ([_]usize{ 2, 3, 8, 16 }) |w| { + const start = rows - w; + try run(queue, launchAt(try row.call(p, m.w, w, false), start * k * 2, 0), m, out); + for (0..w) |i| { + if (!std.mem.eql(u16, out.slice(u16, w * n)[i * n ..][0..n], alone[(start + i) * n ..][0..n])) return error.WindowChangesBits; + t.pairs += n; + } + } + // relu2 is relu squared of the plain output + try run(queue, of(try row.call(p, m.w, rows, true)), m, out); + for (out.slice(u16, rows * n), alone) |r2, v| { + const h = @max(fromBf(v), 0); + if (r2 != toBf(h * h)) return error.Relu2DiffersFromPlain; + } + // indexed slots of one expert: each slot as its row alone; an id past the count is a zero row + for (ids.slice(u32, rows), 0..) |*id, i| id.* = if (i == rows - 1) 1 else 0; + var li = of(try row.callIndexed(p, .{ .one = m.w, .experts = 1 }, rows, false)); + li.ids = ids; + try run(queue, li, m, out); + for (0..rows) |i| { + const got = out.slice(u16, rows * n)[i * n ..][0..n]; + if (i == rows - 1) { + for (got) |v| if (v != 0) return error.MissingExpertNotZero; + } else if (!std.mem.eql(u16, got, alone[i * n ..][0..n])) return error.IndexedDiffersFromDense; + } + // strided: the matrix's rows at column 16 of a wider buffer; its columns as alone, every other value untouched + const stride = 2 * n + 24; + @memset(out.slice(u16, rows * stride), sentinel); + try run(queue, launchAt(try row.callStrided(p, m.w, rows, stride), 0, 16 * 2), m, out); + for (0..rows) |i| for (0..stride) |c| { + const v = out.slice(u16, rows * stride)[i * stride + c]; + if (c >= 16 and c < 16 + n) { + if (v != alone[i * n + c - 16]) return error.StridedDiffersFromAlone; + t.pairs += 1; + } else if (v != sentinel) return error.StridedWritesOutside; + }; + const s = stock orelse return; + if (std.mem.indexOfScalar(u8, &.{ 2, 4, 6, 8 }, m.w.bits) == null) return; + // the stock kernels: plain and relu2 windows of 1, 3 and 16 rows, indexed slots, byte for byte + for ([_]usize{ 1, 3, 16 }) |w| for ([_]bool{ false, true }) |relu2| { + const c = try row.call(p, m.w, w, relu2); + try run(queue, of(c), m, out); + const pipe = try s.get(device, if (relu2) "_relu2" else "", m.w.bits, m.w.sum); + defer pipe.deinit(); + var l = of(c); + l.pipeline = pipe; + try run(queue, l, m, one); + if (!std.mem.eql(u16, out.slice(u16, w * n), one.slice(u16, w * n))) return error.StockBytesDiffer; + t.stock += w * n; + }; + const pipe = try s.get(device, "_indexed", m.w.bits, m.w.sum); + defer pipe.deinit(); + try run(queue, li, m, out); + li.pipeline = pipe; + try run(queue, li, m, one); + if (!std.mem.eql(u16, out.slice(u16, rows * n), one.slice(u16, rows * n))) return error.StockBytesDiffer; + t.stock += rows * n; +} + +fn launchAt(c: row.Call, x_off: usize, out_off: usize) Launch { + var l = of(c); + l.x_off = x_off; + l.out_off = out_off; + return l; +} + +/// A fused stack of 48-column members at mixed widths, group 64: each member's strided columns as its own call. +fn stack(device: mtl.Device, queue: mtl.Queue, p: *const row.Pipelines) !usize { + const widths = [_]u8{ 4, 5, 4, 6, 3, 8, 2, 4 }; + const n = 48; + const stride = widths.len * n; + var members: [widths.len]Mat = undefined; + for (&members, widths, 0..) |*m, bits, i| m.* = try Mat.init(device, .{ .n = n, .k = 2688 }, bits, 64, .f32, 77 + i); + defer for (members) |m| m.deinit(); + // every member reads the same inputs + for (members[1..]) |m| @memcpy(m.x.slice(u16, x_rows * 2688), members[0].x.slice(u16, x_rows * 2688)); + const fused = try device.buffer(x_rows * stride * 2, opts); + defer fused.deinit(); + const own = try device.buffer(x_rows * n * 2, opts); + defer own.deinit(); + const alone = try device.buffer(n * 2, opts); + defer alone.deinit(); + var pairs: usize = 0; + for ([_]usize{ 1, 2, 3, 16, 48 }) |w| { + @memset(fused.slice(u16, w * stride), sentinel); + for (members, 0..) |m, j| try run(queue, launchAt(try row.callStrided(p, m.w, w, stride), 0, j * n * 2), m, fused); + for (members, 0..) |m, j| { + try run(queue, of(try row.call(p, m.w, w, false)), m, own); + for (0..w) |i| { + const got = fused.slice(u16, w * stride)[i * stride + j * n ..][0..n]; + if (!std.mem.eql(u16, got, own.slice(u16, w * n)[i * n ..][0..n])) return error.StackDiffersFromMember; + if (w == 48) { + try run(queue, launchAt(try row.call(p, m.w, 1, false), i * 2688 * 2, 0), m, alone); + if (!std.mem.eql(u16, got, alone.slice(u16, n))) return error.StackDiffersFromRowAlone; + pairs += n; + } + pairs += n; + } + } + for (fused.slice(u16, w * stride)) |v| if (v == sentinel) return error.StackLeftAGap; + } + std.debug.print("stack of {d} members (bits 4 5 4 6 3 8 2 4, n {d}, group 64) at windows 1, 2, 3, 16, 48: each member as its own call\n", .{ widths.len, n }); + return pairs; +} From 26681b8459e2b36f295ebc9a55f9860103976370 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:35:27 -1000 Subject: [PATCH 04/10] Lane projections take every MLX affine width, an output row stride and 16-row members Format accepts 2-, 3- and 5-bit codes beside 4, 6 and 8, in groups of 32, 64 or 128. The non-4-bit path widens 2- and 3-bit codes to nibbles for the bf16 x uint4 op and 5-, 6- and 8-bit codes to bytes for the uint8 op, the staging the Python lane matmul uses, so each width keeps that arithmetic. 3- and 5-bit groups of 32 are an odd number of words a lane; they are fetched a word at a time instead of in pairs, which would drop the last word. Layout.y_stride (default N) sets the elements between output rows in every kernel (tf_lane, tf_lane_pair, tf_lane_reg), so a member of a fused stack writes its columns of the fused rows; the output bound checks follow it. The store of the simdgroup-matrix layout (Macs without tensor units, #679) takes the same stride and keeps its column guard. N may be a multiple of 16 on the plain kernel: the last 32-column tile is packed with zero codes past N, its metadata reads and stores are already guarded by N, and weightWords counts the padded tiles (sourceWords is MLX's W). Cooperative and reg layouts still need N % 32 == 0. encodeSums checks its sums buffer by the group count, as #698 sizes the dispatch. Without tensor units the admitted layout is unchanged: 6-bit group 32 on the plain kernel (#679); every other width still answers TensorUnitsRequired there, now checked at 2, 3 and 5 bits too. With the defaults the 4-, 6- and 8-bit kernels preprocess (metal -E) to the same text as main's for every group, schedule, output, prepared-sums/reg layout and both fragment layouts tried (148 variants), so the uniform paths keep their kernels. Host tests that fail before: 2, 3 and 5 bits validate; y_stride reaches the source and below N is refused; a 48-row member validates (cooperative and reg refuse it); every width x group x {64, 48} rows packs MLX codes bit for bit with zero padding. The device program covers every width, 48-row members and strided output on tensor units; without them (M3 Ultra): PASS 14 layouts, 1904 bit-equal row/window pairs, 208 windows equal to fn_lane's recorded bytes, 7424 CPU values (max roundoff-bound fraction 0.0144), 10 refusals, and a 48-row member and two strided layouts at 6/32 equal the unstrided rows with the gaps untouched. fn_lane's recorded windows are compared only on the layouts they were recorded at. Gate on main 0efbd447 (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test: `get` logs an error on the DeferredTensor path the test expects). lean_check: 536 problems (main 539), no new findings. --- zig/kernels/metal/core/lane_projection.metal | 62 ++++++++- zig/src/core/lane_projection.zig | 134 +++++++++++++------ zig/tests/core_lane_projection.zig | 77 +++++++++-- 3 files changed, 219 insertions(+), 54 deletions(-) diff --git a/zig/kernels/metal/core/lane_projection.metal b/zig/kernels/metal/core/lane_projection.metal index 0c3cf622c1..81cfd5d9d4 100644 --- a/zig/kernels/metal/core/lane_projection.metal +++ b/zig/kernels/metal/core/lane_projection.metal @@ -12,18 +12,29 @@ using namespace mpp::tensor_ops; #ifndef TF_PRECOMPUTE_SUMS #define TF_PRECOMPUTE_SUMS 0 #endif +#ifndef TF_YS +#define TF_YS TF_N +#endif +// Output rows lie TF_YS elements apart, so a member of a fused stack writes its columns of the fused rows. +#define TF_ODD_WORDS (TF_GROUP * TF_BITS / 32 % 2) constant constexpr int NT = 32, GS = TF_GROUP, KG = TF_K / GS, NF = NT / 16; constant constexpr int WP = GS * TF_BITS / 32; +#if TF_ODD_WORDS +struct tf_raw { uint w[WP]; uint4 s[NF]; }; // 3- and 5-bit groups of 32 are an odd word count, loaded a word at a time +#else struct tf_raw { uint2 w[WP / 2]; uint4 s[NF]; }; // a lane's packed code words and its scale/bias quads for one group +#endif inline tf_raw tf_lane_fetch(const device uint* W, const device uint4* SB, int tile, int g, int g_end, int n0, int lane, short fn) { tf_raw raw; const bool ok = g < g_end; const device uint2* src = (const device uint2*)(W + (((size_t)tile * KG + (ok ? g : 0)) * NT + lane) * WP); -#if TF_BITS != 4 +#if TF_ODD_WORDS + for (int i = 0; i < WP; i++) raw.w[i] = ok && n0 + lane < TF_N ? ((const device uint*)src)[i] : 0u; +#elif TF_BITS != 4 for (int i = 0; i < WP / 2; i++) raw.w[i] = ok && n0 + lane < TF_N ? src[i] : uint2(0); #else for (int i = 0; i < WP / 2; i++) raw.w[i] = uint2(0); @@ -53,6 +64,7 @@ inline float tf_lane_xsum(const device bfloat* X, int m, int g, int M) { return acc; } +// Codes: 4-bit straight to the uint4 op; 2/3-bit widened to nibbles for it, 5/6/8-bit to bytes for the uint8 op. #if TF_BITS == 4 #define TF_STAGE_CODES(raw) #define TF_GROUP_PRODUCT(g) \ @@ -62,9 +74,42 @@ inline float tf_lane_xsum(const device bfloat* X, int m, int g, int M) { auto P = op.template get_destination_cooperative_tensor(); \ op.run(a, b, P); #else +#if TF_BITS < 4 +inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, int lane) { + uint words[WP + 1]; +#if TF_ODD_WORDS + for (int i = 0; i < WP; i++) words[i] = raw.w[i]; +#else + for (int i = 0; i < WP / 2; i++) { words[2*i] = raw.w[i].x; words[2*i+1] = raw.w[i].y; } +#endif + words[WP] = 0u; +#if TF_BITS == 3 + for (int c = 0; c < GS / 8; c++) { + const int bit = 24 * c, i = bit >> 5, sh = bit & 31; + uint pack = words[i] >> sh; + if (sh > 8) pack |= words[i + 1] << (32 - sh); + uint nib = 0u; + for (int j = 0; j < 8; j++) nib |= ((pack >> (3 * j)) & 7u) << (4 * j); + stage[lane * (GS / 8) + c] = nib; + } +#else + for (int c = 0; c < GS / 8; c++) { // word c/2's half c%2: 8 codes of 2 bits to 8 nibbles + uint v = (words[c >> 1] >> (16 * (c & 1))) & 0xFFFFu; + v = (v | (v << 8)) & 0x00FF00FFu; + v = (v | (v << 4)) & 0x0F0F0F0Fu; + v = (v | (v << 2)) & 0x33333333u; + stage[lane * (GS / 8) + c] = v; + } +#endif +} +#else inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, int lane) { uint words[WP + 1]; +#if TF_ODD_WORDS + for (int i = 0; i < WP; i++) words[i] = raw.w[i]; +#else for (int i = 0; i < WP / 2; i++) { words[2*i] = raw.w[i].x; words[2*i+1] = raw.w[i].y; } +#endif words[WP] = 0u; for (int c = 0; c < GS / 4; c++) { uint word = 0u; @@ -77,6 +122,7 @@ inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, in stage[lane * (GS / 4) + c] = word; } } +#endif #define TF_STAGE_CODES(raw) tf_stage_codes(stage, raw, lane); #define TF_GROUP_PRODUCT(g) \ simdgroup_barrier(mem_flags::mem_threadgroup); \ @@ -136,7 +182,11 @@ inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, in constexpr auto desc = matmul2d_descriptor(16, NT, GS, false, true, false, matmul2d_descriptor::mode::multiply); matmul2d op; tensor, tensor_inline> tA((device bfloat*)X, dextents(TF_K, M)); -#if TF_BITS != 4 +#if TF_BITS < 4 + threadgroup uint stage_all[TF_SK * NT * (GS / 8)]; + threadgroup uint* stage = stage_all + sg * NT * (GS / 8); + tensor, tensor_inline> b((threadgroup uchar*)stage, dextents(GS, NT)); +#elif TF_BITS != 4 threadgroup uint stage_all[TF_SK * NT * (GS / 4)]; threadgroup uint* stage = stage_all + sg * NT * (GS / 4); tensor, tensor_inline> b((threadgroup uint8_t*)stage, dextents(GS, NT)); @@ -169,7 +219,7 @@ inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, in if (m < M && nn < TF_N) for (int j = 0; j < 4; j++) { const int col = nn + (TF_SIMD_FRAGS ? (j & 1) + 8 * (j >> 1) : j); - if (col < TF_N) Y[(size_t)m * TF_N + col] = static_cast(C[f * 8 + r * 4 + j]); + if (col < TF_N) Y[(size_t)m * TF_YS + col] = static_cast(C[f * 8 + r * 4 + j]); } } } @@ -219,7 +269,7 @@ uint group_id [[simdgroup_index_in_threadgroup]],uint lane [[thread_index_in_sim threadgroup_barrier(mem_flags::mem_threadgroup); if(sg==0) for(int part=1;part> 3) & 1), mb = ((lane >> 1) & 3) + 4 * ((lane >> 4) & 1); @@ -288,7 +338,7 @@ inline tf_reg_load tf_reg_fetch(const device uint4* wl, const device uint4* sbv, if (ks == 0) for (int i = 0; i < 16; i++) { const int m = mb + 8 * ((i >> 2) & 1), n = n0 + nb + 16 * (i >> 3) + (i & 3); - if (m < M && n < TF_N) Y[(size_t)m * TF_N + n] = static_cast(C[i]); + if (m < M && n < TF_N) Y[(size_t)m * TF_YS + n] = static_cast(C[i]); } } #endif diff --git a/zig/src/core/lane_projection.zig b/zig/src/core/lane_projection.zig index 892a8a541b..001b684892 100644 --- a/zig/src/core/lane_projection.zig +++ b/zig/src/core/lane_projection.zig @@ -4,7 +4,7 @@ const mtl = @import("metal"); const ks = @import("kernel_sources"); const frags = @import("frags.zig"); -pub const weight_layout = "u32[N/32][K/group][32][group*bits/32], low-bit-first codes"; +pub const weight_layout = "u32[ceil(N/32)][K/group][32][group*bits/32], low-bit-first codes, columns past N zero"; pub const reg_layout = "u32[N/32][K/64][2][32 lanes][4] 4-bit codes, regBlock order"; pub const metadata_layout = "bf16[K/group][N][scale,bias]"; pub const Output = enum { bf16, f32 }; @@ -14,7 +14,7 @@ pub const Format = struct { group: usize, pub fn validate(f: Format) !void { - if (f.bits != 4 and f.bits != 6 and f.bits != 8) return error.UnsupportedBits; + if (f.bits < 2 or f.bits > 8 or f.bits == 7) return error.UnsupportedBits; if (f.group != 32 and f.group != 64 and f.group != 128) return error.UnsupportedGroup; } }; @@ -34,12 +34,16 @@ pub const Layout = struct { cooperative: bool = false, /// 4-bit g64 codes in reg_layout, unpacked into tensor-op registers; every word equals the cooperative kernel's. reg: bool = false, + /// Elements between output rows (default N): a member of a fused stack writes its columns of the fused rows. + y_stride: ?usize = null, pub fn validate(l: Layout) !void { try l.format.validate(); if (l.cooperative and l.format.bits != 4) return error.UnsupportedCooperativeFormat; if (l.reg and (l.cooperative or l.format.bits != 4 or l.format.group != 64 or !l.precompute_sums)) return error.UnsupportedRegFormat; - if (l.n == 0 or l.n % 32 != 0 or l.k == 0 or l.k % l.format.group != 0 or + if ((l.cooperative or l.reg) and l.n % 32 != 0) return error.InvalidShape; + if (l.y_stride) |s| if (s < l.n or s > std.math.maxInt(i32) / 16) return error.InvalidOutputStride; + if (l.n == 0 or l.n % 16 != 0 or l.k == 0 or l.k % l.format.group != 0 or l.n > std.math.maxInt(i32) / 16 or l.k > std.math.maxInt(i32) / 16) return error.InvalidShape; if (l.sk == 0 or l.sk > 32 or (l.pf != 1 and l.pf != 2)) return error.InvalidSchedule; const g = l.groupCut(); @@ -49,7 +53,7 @@ pub const Layout = struct { if (l.threadgroupBytes() > 32768) return error.ThreadgroupMemoryLimit; var end: usize = 0; for (l.ranges) |r| { - if (r[1] == 0 or r[0] < end or r[0] >= l.n / 32 or r[1] > l.n / 32 - r[0]) return error.InvalidTileRange; + if (r[1] == 0 or r[0] < end or r[0] >= l.tileCount() or r[1] > l.tileCount() - r[0]) return error.InvalidTileRange; end = r[0] + r[1]; } _ = try std.math.mul(usize, l.n, l.k); @@ -59,8 +63,13 @@ pub const Layout = struct { return l.groups orelse .{ 0, l.k / l.format.group }; } + /// 32-column tiles over N; a last tile of 16 columns is padded with zero codes and never stored. + pub fn tileCount(l: Layout) usize { + return (l.n + 31) / 32; + } + pub fn tiles(l: Layout) usize { - if (l.ranges.len == 0) return l.n / 32; + if (l.ranges.len == 0) return l.tileCount(); var count: usize = 0; for (l.ranges) |r| count += r[1]; return count; @@ -71,6 +80,11 @@ pub const Layout = struct { } pub fn weightWords(l: Layout) usize { + return l.tileCount() * 32 * (l.k / 32) * l.format.bits; + } + + /// MLX's packed W for this layout: N rows of K*bits/32 words. + pub fn sourceWords(l: Layout) usize { return l.n * (l.k / 32) * l.format.bits; } @@ -100,6 +114,7 @@ pub fn source(a: std.mem.Allocator, l: Layout) ![]u8 { try text.appendSlice(a, if (l.ranges.len == 0) " return j;\n}\n" else " return 0;\n}\n"); try text.print(a, "#define TF_PRECOMPUTE_SUMS {d}\n", .{@intFromBool(l.precompute_sums)}); if (l.reg) try text.appendSlice(a, "#define TF_REG 1\n"); + if (l.y_stride) |s| try text.print(a, "#define TF_YS {d}\n", .{s}); try text.appendSlice(a, ks.core_lane_projection); return text.toOwnedSlice(a); } @@ -110,8 +125,8 @@ pub fn pack(l: Layout, w: []const u32, scales: []const u16, biases: []const u16, const ng = l.k / l.format.group; const wp = l.format.group * l.format.bits / 32; const row_words = l.k / 32 * l.format.bits; - if (w.len != l.weightWords() or scales.len != l.n * ng or biases.len != l.n * ng or - out_w.len != w.len or out_sb.len != l.metadataElements()) return error.InvalidStorage; + if (w.len != l.sourceWords() or scales.len != l.n * ng or biases.len != l.n * ng or + out_w.len != l.weightWords() or out_sb.len != l.metadataElements()) return error.InvalidStorage; if (overlap(std.mem.sliceAsBytes(w), std.mem.sliceAsBytes(out_w)) or overlap(std.mem.sliceAsBytes(w), std.mem.sliceAsBytes(out_sb)) or overlap(std.mem.sliceAsBytes(scales), std.mem.sliceAsBytes(out_w)) or @@ -119,11 +134,14 @@ pub fn pack(l: Layout, w: []const u32, scales: []const u16, biases: []const u16, overlap(std.mem.sliceAsBytes(scales), std.mem.sliceAsBytes(out_sb)) or overlap(std.mem.sliceAsBytes(biases), std.mem.sliceAsBytes(out_sb)) or overlap(std.mem.sliceAsBytes(out_w), std.mem.sliceAsBytes(out_sb))) return error.OverlappingStorage; - for (0..l.n / 32) |tile| for (0..ng) |g| for (0..32) |col| { + for (0..l.tileCount()) |tile| for (0..ng) |g| for (0..32) |col| { const n = tile * 32 + col; - const src = n * row_words + g * wp; const dst = ((tile * ng + g) * 32 + col) * wp; - @memcpy(out_w[dst..][0..wp], w[src..][0..wp]); + if (n >= l.n) { + @memset(out_w[dst..][0..wp], 0); + continue; + } + @memcpy(out_w[dst..][0..wp], w[n * row_words + g * wp ..][0..wp]); const sb = (g * l.n + n) * 2; out_sb[sb] = scales[n * ng + g]; out_sb[sb + 1] = biases[n * ng + g]; @@ -207,6 +225,7 @@ pub const Projection = struct { weight_bytes: usize, metadata_bytes: usize, output_bytes: usize, + y_stride: usize, tile_count: usize, sk: usize, sums_pipe: ?mtl.Pipeline = null, @@ -233,7 +252,7 @@ pub const Projection = struct { errdefer sums_pipe.?.deinit(); sums_buffer = try device.buffer(@as(usize, if (l.reg) reg_rows else 16) * (l.k / l.format.group) * 4, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); } - return .{ .pipe = p, .n = l.n, .k = l.k, .group_count = l.k / l.format.group, .weight_bytes = l.weightWords() * 4, .metadata_bytes = l.metadataElements() * 2, .output_bytes = if (l.output == .f32) 4 else 2, .tile_count = l.tiles(), .sk = l.sk, .sums_pipe = sums_pipe, .sums_buffer = sums_buffer, .threads_per_slice = if (l.cooperative) 64 else 32, .reg = l.reg }; + return .{ .pipe = p, .n = l.n, .k = l.k, .group_count = l.k / l.format.group, .weight_bytes = l.weightWords() * 4, .metadata_bytes = l.metadataElements() * 2, .output_bytes = if (l.output == .f32) 4 else 2, .y_stride = l.y_stride orelse l.n, .tile_count = l.tiles(), .sk = l.sk, .sums_pipe = sums_pipe, .sums_buffer = sums_buffer, .threads_per_slice = if (l.cooperative) 64 else 32, .reg = l.reg }; } pub fn deinit(p: Projection) void { @@ -251,7 +270,7 @@ pub const Projection = struct { pub fn encodeSums(p: Projection, e: mtl.ComputeEncoder, x: Ref, w: Ref, sb: Ref, sums: Sums, y: Ref, rows: u32) !void { if (!p.reg) return error.SumsNeedReg; if (rows == 0 or rows > reg_rows or sums.stride < rows) return error.InvalidRows; - try checkRef(sums.ref, ((p.k / 64 - 1) * @as(usize, sums.stride) + rows) * 4, 4); + try checkRef(sums.ref, ((p.group_count - 1) * @as(usize, sums.stride) + rows) * 4, 4); return p.encodeImpl(e, x, w, sb, sums, y, rows); } @@ -262,7 +281,7 @@ pub const Projection = struct { try checkRef(x, rows * p.k * 2, 2); try checkRef(w, p.weight_bytes, if (p.reg) 16 else 8); try checkRef(sb, p.metadata_bytes, 16); - try checkRef(y, rows * p.n * p.output_bytes, p.output_bytes); + try checkRef(y, ((rows - 1) * p.y_stride + p.n) * p.output_bytes, p.output_bytes); if (given == null) if (p.sums_pipe) |pipe| { e.setPipeline(pipe); e.setBuffer(x.buf, x.off, 0); @@ -302,9 +321,35 @@ test "layouts reject unsupported formats, schedules, cuts and overlapping tile r const good = Layout{ .n = 64, .k = 256, .format = .{ .bits = 4, .group = 64 } }; const s = try source(arena.allocator(), good); try std.testing.expect(std.mem.indexOf(u8, s, "#define TF_BITS 4\n#define TF_GROUP 64\n") != null); + try std.testing.expect(std.mem.indexOf(u8, s[0 .. s.len - ks.core_lane_projection.len], "TF_YS") == null); var bad = good; - bad.format.bits = 5; - try std.testing.expectError(error.UnsupportedBits, bad.validate()); + for ([_]u8{ 2, 3, 5, 6, 8 }) |bits| { + bad.format.bits = bits; + try bad.validate(); + } + for ([_]u8{ 0, 1, 7, 9 }) |bits| { + bad.format.bits = bits; + try std.testing.expectError(error.UnsupportedBits, bad.validate()); + } + bad = good; + bad.y_stride = 96; + try std.testing.expect(std.mem.indexOf(u8, try source(arena.allocator(), bad), "#define TF_YS 96\n") != null); + bad.y_stride = 63; + try std.testing.expectError(error.InvalidOutputStride, bad.validate()); + bad = good; + bad.n = 48; + bad.ranges = &.{.{ 1, 1 }}; + try bad.validate(); + try std.testing.expectEqual(2, bad.tileCount()); + bad.cooperative = true; + try std.testing.expectError(error.InvalidShape, bad.validate()); + bad.cooperative = false; + bad.reg = true; + bad.precompute_sums = true; + try std.testing.expectError(error.InvalidShape, bad.validate()); + bad = good; + bad.ranges = &.{.{ 2, 1 }}; + try std.testing.expectError(error.InvalidTileRange, bad.validate()); bad = good; bad.format.group = 16; try std.testing.expectError(error.UnsupportedGroup, bad.validate()); @@ -334,34 +379,45 @@ test "layouts reject unsupported formats, schedules, cuts and overlapping tile r try std.testing.expectError(error.InvalidTileRange, bad.validate()); } -test "every affine layout packs raw words and metadata without changing their bits" { - for ([_]u8{ 4, 6, 8 }) |bits| for ([_]usize{ 32, 64, 128 }) |group| { - const l = Layout{ .n = 64, .k = 2 * group, .format = .{ .bits = bits, .group = group } }; - const a = std.testing.allocator; - const w = try a.alloc(u32, l.weightWords()); - defer a.free(w); - const s = try a.alloc(u16, 128); - defer a.free(s); - const b = try a.alloc(u16, 128); - defer a.free(b); - const pw = try a.alloc(u32, w.len); - defer a.free(pw); - const sb = try a.alloc(u16, 256); - defer a.free(sb); - for (w, 0..) |*v, i| v.* = @as(u32, @intCast(i)) *% 0x9e3779b9; - for (s, 0..) |*v, i| v.* = @intCast(i); - for (b, 0..) |*v, i| v.* = @intCast(65535 - i); +test "every width, group and 16-row member packs MLX codes, metadata and zero padding without changing a bit" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var prng: std.Random.DefaultPrng = .init(0x1a9e); + for ([_]u8{ 2, 3, 4, 5, 6, 8 }) |bits| for ([_]usize{ 32, 64, 128 }) |group| for ([_]usize{ 64, 48 }) |rows| { + const l = Layout{ .n = rows, .k = 2 * group, .format = .{ .bits = bits, .group = group } }; + const codes = try a.alloc(u8, l.n * l.k); + const w = try a.alloc(u32, l.sourceWords()); + @memset(w, 0); + for (codes, 0..) |*v, i| { + v.* = @intCast(prng.random().intRangeAtMost(u16, 0, (@as(u16, 1) << @intCast(bits)) - 1)); + // MLX's packing written bit by bit: row n is one low-bit-first stream of K codes. + for (0..bits) |bit| w[(i * bits + bit) / 32] |= @as(u32, (v.* >> @intCast(bit)) & 1) << @intCast((i * bits + bit) % 32); + } + const s = try a.alloc(u16, l.n * 2); + const b = try a.alloc(u16, l.n * 2); + for (s, b, 0..) |*sv, *bv, i| { + sv.* = @intCast(i); + bv.* = @intCast(65535 - i); + } + const pw = try a.alloc(u32, l.weightWords()); + @memset(pw, 0xffff_ffff); + const sb = try a.alloc(u16, l.metadataElements()); try pack(l, w, s, b, pw, sb); const wp = group * bits / 32; - for (0..64) |n| for (0..2) |g| { - const src = (n * 2 + g) * wp; - const dst = (((n / 32) * 2 + g) * 32 + n % 32) * wp; - try std.testing.expectEqualSlices(u32, w[src..][0..wp], pw[dst..][0..wp]); - try std.testing.expectEqual(s[n * 2 + g], sb[(g * 64 + n) * 2]); - try std.testing.expectEqual(b[n * 2 + g], sb[(g * 64 + n) * 2 + 1]); + for (0..l.tileCount() * 32) |n| for (0..l.k) |k| { + const at = ((((n / 32) * 2 + k / group) * 32 + n % 32) * wp) * 32 + (k % group) * bits; + var code: u32 = 0; + for (0..bits) |bit| code |= ((pw[(at + bit) / 32] >> @intCast((at + bit) % 32)) & 1) << @intCast(bit); + try std.testing.expectEqual(@as(u32, if (n < l.n) codes[n * l.k + k] else 0), code); + }; + for (0..l.n) |n| for (0..2) |g| { + try std.testing.expectEqual(s[n * 2 + g], sb[(g * l.n + n) * 2]); + try std.testing.expectEqual(b[n * 2 + g], sb[(g * l.n + n) * 2 + 1]); }; try std.testing.expectError(error.InvalidStorage, pack(l, w[1..], s, b, pw, sb)); - try std.testing.expectError(error.OverlappingStorage, pack(l, w, s, b, w, sb)); + try std.testing.expectError(error.InvalidStorage, pack(l, w, s, b, pw[1..], sb)); + try std.testing.expectError(error.OverlappingStorage, pack(l, pw[0..w.len], s, b, pw, sb)); }; } diff --git a/zig/tests/core_lane_projection.zig b/zig/tests/core_lane_projection.zig index bce8c3299a..bdd3e01b88 100644 --- a/zig/tests/core_lane_projection.zig +++ b/zig/tests/core_lane_projection.zig @@ -24,7 +24,7 @@ const Data = struct { biases: []u16, fn init(l: core.Layout) !Data { - const d = Data{ .x = try a.alloc(u16, rows_max * l.k), .codes = try a.alloc(u8, l.n * l.k), .words = try a.alloc(u32, l.weightWords()), .scales = try a.alloc(u16, l.n * l.k / l.format.group), .biases = try a.alloc(u16, l.n * l.k / l.format.group) }; + const d = Data{ .x = try a.alloc(u16, rows_max * l.k), .codes = try a.alloc(u8, l.n * l.k), .words = try a.alloc(u32, l.sourceWords()), .scales = try a.alloc(u16, l.n * l.k / l.format.group), .biases = try a.alloc(u16, l.n * l.k / l.format.group) }; @memset(d.words, 0); var prng: std.Random.DefaultPrng = .init(0x1a9e + @as(u64, l.format.bits) * 1000 + l.format.group); const rnd = prng.random(); @@ -169,7 +169,8 @@ fn check(device: mtl.Device, queue: mtl.Queue, l: core.Layout, counts: *Counts) counts.cpu_values += 1; }; } - if (l.format.bits == 6 and l.format.group == 32) { + const pinned = if (frozen.bytes(l)) |_| true else |_| false; // fn_lane's recorded 6/32 layouts only + if (pinned) { for (1..rows_max + 1) |m| { const count = m * l.n * elem; const got = windows.contents()[guard + (m - 1) * window_bytes ..][0..count]; @@ -180,20 +181,67 @@ fn check(device: mtl.Device, queue: mtl.Queue, l: core.Layout, counts: *Counts) for ([_]mtl.Buffer{ x, w, sb, one, windows }) |b| try checkGuard(b); counts.layouts += 1; std.debug.print("PASS b{d}g{d} sk{d} pf{d} {s} groups{any} tiles{d}: widths1-16, guards{s}\n", .{ - l.format.bits, l.format.group, l.sk, l.pf, @tagName(l.output), l.groupCut(), l.tiles(), - if (l.format.bits == 6 and l.format.group == 32) ", frozen fn_lane byte-equal" else "", + l.format.bits, l.format.group, l.sk, l.pf, @tagName(l.output), l.groupCut(), l.tiles(), + if (pinned) ", frozen fn_lane byte-equal" else "", + }); +} + +/// A member writing into fused rows: each row equals the unstrided encode's and nothing around it is touched. +fn strided(device: mtl.Device, queue: mtl.Queue, l: core.Layout) !void { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + var fused = l; + fused.y_stride = l.n + 48; + const p = try core.Projection.init(a, device, l); + defer p.deinit(); + const q = try core.Projection.init(a, device, fused); + defer q.deinit(); + const d = try Data.init(l); + defer d.deinit(); + const x = try buffer(device, d.x.len * 2); + defer x.deinit(); + const w = try buffer(device, l.weightWords() * 4); + defer w.deinit(); + const sb = try buffer(device, l.metadataElements() * 2); + defer sb.deinit(); + const elem: usize = if (l.output == .f32) 4 else 2; + const ys = fused.y_stride.? * elem; + const plain = try buffer(device, rows_max * l.n * elem); + defer plain.deinit(); + const stack = try buffer(device, rows_max * ys); + defer stack.deinit(); + @memcpy(x.contents()[guard..][0 .. d.x.len * 2], std.mem.sliceAsBytes(d.x)); + const packed_w = @as([*]u32, @ptrCast(@alignCast(w.contents() + guard)))[0..l.weightWords()]; + const packed_sb = @as([*]u16, @ptrCast(@alignCast(sb.contents() + guard)))[0..l.metadataElements()]; + try core.pack(l, d.words, d.scales, d.biases, packed_w, packed_sb); + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try p.encode(enc, ref(x, 0), ref(w, 0), ref(sb, 0), ref(plain, 0), rows_max); + try q.encode(enc, ref(x, 0), ref(w, 0), ref(sb, 0), ref(stack, 16 * elem), rows_max); + enc.end(); + try finish(cb); + for (0..rows_max) |r| { + const row = stack.contents()[guard + r * ys ..][0..ys]; + const want = plain.contents()[guard + r * l.n * elem ..][0 .. l.n * elem]; + if (!std.mem.eql(u8, want, row[16 * elem ..][0 .. l.n * elem])) return error.StridedRowMismatch; + for (row[0 .. 16 * elem]) |v| if (v != 0xa5) return error.StridedGapOverwrite; + for (row[(16 + l.n) * elem ..]) |v| if (v != 0xa5) return error.StridedGapOverwrite; + } + for ([_]mtl.Buffer{ x, w, sb, plain, stack }) |b| try checkGuard(b); + std.debug.print("PASS b{d}g{d} n{d} {s}: output rows {d} apart equal the unstrided rows, gaps untouched\n", .{ + l.format.bits, l.format.group, l.n, @tagName(l.output), fused.y_stride.?, }); } fn packBench(io: std.Io, l: core.Layout) !void { try l.validate(); - const w = try a.alloc(u32, l.weightWords()); + const w = try a.alloc(u32, l.sourceWords()); defer a.free(w); const s = try a.alloc(u16, l.n * l.k / l.format.group); defer a.free(s); const b = try a.alloc(u16, s.len); defer a.free(b); - const pw = try a.alloc(u32, w.len); + const pw = try a.alloc(u32, l.weightWords()); defer a.free(pw); const sb = try a.alloc(u16, 2 * s.len); defer a.free(sb); @@ -261,7 +309,7 @@ pub fn main(init: std.process.Init) !void { try check(device, queue, .{ .n = 96, .k = 224, .format = .{ .bits = 6, .group = 32 }, .sk = 4, .pf = 2, .ranges = &.{.{ 1, 1 }}, .groups = .{ 1, 3 }, .output = .f32 }, &counts); if (!device.tensorUnits()) { var unsupported = core.Layout{ .n = 64, .k = 224, .format = .{ .bits = 6, .group = 32 } }; - for ([_]u8{ 4, 8 }) |bits| { + for ([_]u8{ 2, 3, 4, 5, 8 }) |bits| { unsupported.format.bits = bits; try std.testing.expectError(error.TensorUnitsRequired, core.Projection.init(a, device, unsupported)); } @@ -278,12 +326,16 @@ pub fn main(init: std.process.Init) !void { unsupported.reg = true; try std.testing.expectError(error.TensorUnitsRequired, core.Projection.init(a, device, unsupported)); } + // 48 rows (a zero-padded last tile) and a member writing its columns of fused rows, on the same path + try check(device, queue, .{ .n = 48, .k = 160, .format = .{ .bits = 6, .group = 32 }, .sk = 2, .pf = 2, .output = .f32 }, &counts); + try strided(device, queue, .{ .n = 48, .k = 96, .format = .{ .bits = 6, .group = 32 } }); + try strided(device, queue, .{ .n = 64, .k = 96, .format = .{ .bits = 6, .group = 32 }, .sk = 2, .output = .f32 }); std.debug.print("PASS Flash Next: {d} layouts, {d} bit-equal row/window pairs, {d} frozen windows, {d} CPU values, {d} refusals; max fp32 error {d:.9}, max roundoff-bound fraction {d:.6}\n", .{ - counts.layouts, counts.row_pairs, counts.golden_windows, counts.cpu_values, @as(usize, if (device.tensorUnits()) 0 else 7), counts.max_error, counts.max_fraction, + counts.layouts, counts.row_pairs, counts.golden_windows, counts.cpu_values, @as(usize, if (device.tensorUnits()) 0 else 10), counts.max_error, counts.max_fraction, }); return; } - for ([_]u8{ 4, 6, 8 }) |bits| for ([_]usize{ 32, 64, 128 }) |group| { + for ([_]u8{ 2, 3, 4, 5, 6, 8 }) |bits| for ([_]usize{ 32, 64, 128 }) |group| { for ([_]usize{ 1, 2, 4 }) |sk| for ([_]usize{ 1, 2 }) |pf| { const l = core.Layout{ .n = 64, .k = 7 * group, .format = .{ .bits = bits, .group = group }, .sk = sk, .pf = pf, .output = .f32 }; try check(device, queue, l, &counts); @@ -309,6 +361,13 @@ pub fn main(init: std.process.Init) !void { const group: usize = if (bits == 4) 64 else if (bits == 6) 32 else 128; try check(device, queue, .{ .n = 96, .k = 7 * group, .format = .{ .bits = bits, .group = group }, .sk = 4, .pf = 2, .ranges = &.{.{ 1, 1 }}, .groups = .{ 1, 3 }, .output = .f32 }, &counts); } + for ([_]u8{ 2, 3, 4, 5, 6, 8 }) |bits| for ([_]usize{ 32, 64, 128 }) |group| { + // 48 rows: the last tile's 16 padded columns are read as zero codes and never stored. + try check(device, queue, .{ .n = 48, .k = 5 * group, .format = .{ .bits = bits, .group = group }, .sk = 2, .pf = 2, .output = .f32 }, &counts); + try strided(device, queue, .{ .n = 48, .k = 3 * group, .format = .{ .bits = bits, .group = group } }); + try strided(device, queue, .{ .n = 64, .k = 3 * group, .format = .{ .bits = bits, .group = group }, .sk = 2, .output = .f32 }); + }; + try strided(device, queue, .{ .n = 96, .k = 9 * 64, .format = .{ .bits = 4, .group = 64 }, .sk = 2, .precompute_sums = true, .cooperative = true }); std.debug.print("PASS {d} layouts, {d} bit-equal row/window pairs, {d} frozen fn_lane byte-equal windows, {d} CPU-reference values; max fp32 error {d:.9}, max roundoff-bound fraction {d:.6}\n", .{ counts.layouts, counts.row_pairs, counts.golden_windows, counts.cpu_values, counts.max_error, counts.max_fraction, }); From 0649371e605d865451b9915a5e2bd89cb4790d3b Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 04:10:26 -1000 Subject: [PATCH 05/10] The prompt product affine_mm reads 2-bit MLX codes, exactly affine_mm took 4- and 8-bit weights only, so a 2-bit matrix (oQ2's routed experts) had no tiled prompt product. At 2 bits a thread loads its 32 codes of a step as 8 bytes and writes each code into the bf16 tile as itself, with no offset (as at 8 bits; 0..3 is exact in bf16), so the group sums are the 4- and 8-bit path's. A first try with the 4-bit "+128" offset stayed within the float64 bound, but a reference with scales 1% off no longer exceeded it at K = 4096, so the codes enter exactly. The 4- and 8-bit sources preprocess as before. Test that fails before: "dense and gathered, each format, against exact sums on the host" adds 2/64 and 2/32, dense and gathered; source() refused 2 bits. The 4- and 8-bit sources preprocess (metal -E) to main's text at both groups, both outputs and both fragment layouts (16 variants). tf-affine-ref (next commit) holds the 2-bit products to float64. Gate on main 0efbd447 (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test: `get` logs an error on the DeferredTensor path the test expects). lean_check: 536 problems (main 539), no new findings. --- zig/kernels/metal/core/affine_mm.metal | 28 +++++++++++++++++++++++--- zig/src/core/affine_mm.zig | 12 +++++++---- 2 files changed, 33 insertions(+), 7 deletions(-) diff --git a/zig/kernels/metal/core/affine_mm.metal b/zig/kernels/metal/core/affine_mm.metal index 16e54006d2..464e459fe4 100644 --- a/zig/kernels/metal/core/affine_mm.metal +++ b/zig/kernels/metal/core/affine_mm.metal @@ -4,7 +4,7 @@ using namespace metal; #include "../nax.h" using namespace tfp; -// Defines: TF_BITS (4 or 8), TF_GROUP (32 or 64), TF_OUT_T (bfloat, or float for partial sums another pass adds). +// Defines: TF_BITS (2, 4 or 8), TF_GROUP (32 or 64), TF_OUT_T (bfloat, or float for partial sums another pass adds). #ifndef TF_OUT_T #define TF_OUT_T bfloat #endif @@ -66,7 +66,16 @@ inline void am_first(thread frag& lo, thread frag& hi, thread cons #endif } -// A thread's 32 codes of one weight row in a step (16-byte loads). +// A thread's 32 codes of one weight row in a step (16-byte loads; 8 bytes at 2 bits). +#if TF_BITS == 2 +struct Codes { + uint2 v; +}; + +inline Codes am_load(const device uchar* w) { + return Codes{*(const device uint2*)w}; +} +#else struct Codes { uint4 v[TF_BITS / 4]; }; @@ -79,11 +88,24 @@ inline Codes am_load(const device uchar* w) { } return c; } +#endif // The codes plus OFF, exactly, as bf16 into the tile row (16-byte stores). inline void am_codes(thread const Codes& c, threadgroup bfloat* out) { threadgroup uint4* o = (threadgroup uint4*)out; -#if TF_BITS == 4 +#if TF_BITS == 2 + TF_UNROLL + for (short i = 0; i < 2; i++) { // a word's 16 codes, as themselves (no offset): byte j holds codes 4j..4j+3 + const uint w = c.v[i]; + const float4 f0 = float4(as_type(w & 0x03030303u)), f1 = float4(as_type((w >> 2) & 0x03030303u)); + const float4 f2 = float4(as_type((w >> 4) & 0x03030303u)), f3 = float4(as_type((w >> 6) & 0x03030303u)); + TF_UNROLL + for (short j = 0; j < 4; j++) { + o[2 * i + j / 2][2 * (j % 2)] = as_type(ushort2(as_type(bfloat(f0[j])), as_type(bfloat(f1[j])))); + o[2 * i + j / 2][2 * (j % 2) + 1] = as_type(ushort2(as_type(bfloat(f2[j])), as_type(bfloat(f3[j])))); + } + } +#elif TF_BITS == 4 TF_UNROLL for (short i = 0; i < 4; i++) { const uint w = c.v[0][i]; diff --git a/zig/src/core/affine_mm.zig b/zig/src/core/affine_mm.zig index dd25dce831..4b5575a84a 100644 --- a/zig/src/core/affine_mm.zig +++ b/zig/src/core/affine_mm.zig @@ -4,12 +4,12 @@ const mtl = @import("metal"); const ks = @import("kernel_sources"); const frags = @import("frags.zig"); -/// The weights' bits (4 or 8) and group (32 or 64), and the output's type. +/// The weights' bits (2, 4 or 8) and group (32 or 64), and the output's type. pub const Format = struct { bits: u8 = 4, group: u16 = 64, out_f32: bool = false }; /// The kernel source for `f` (nax.h inlined for this GPU). pub fn source(device: mtl.Device, a: std.mem.Allocator, f: Format) ![]u8 { - if ((f.bits != 4 and f.bits != 8) or (f.group != 32 and f.group != 64)) return error.UnsupportedFormat; + if ((f.bits != 2 and f.bits != 4 and f.bits != 8) or (f.group != 32 and f.group != 64)) return error.UnsupportedFormat; const body = try frags.source(device, a, ks.core_affine_mm); defer a.free(body); return std.fmt.allocPrint(a, "#define TF_BITS {d}\n#define TF_GROUP {d}\n#define TF_OUT_T {s}\n{s}", .{ f.bits, f.group, if (f.out_f32) "float" else "bfloat", body }); @@ -135,7 +135,7 @@ test "dense and gathered, each format, against exact sums on the host" { const R = counts[0] + counts[1] + counts[2]; var prng = std.Random.DefaultPrng.init(7); const rnd = prng.random(); - for ([_][2]usize{ .{ 4, 64 }, .{ 4, 32 }, .{ 8, 64 }, .{ 8, 32 } }) |fmt| { + for ([_][2]usize{ .{ 4, 64 }, .{ 4, 32 }, .{ 8, 64 }, .{ 8, 32 }, .{ 2, 64 }, .{ 2, 32 } }) |fmt| { const bits = fmt[0]; const group = fmt[1]; const w_bytes = E * N * K * bits / 8; @@ -207,7 +207,11 @@ test "dense and gathered, each format, against exact sums on the host" { const off: f64 = if (bits == 4) 128 else 0; for (0..group) |j| { const k = g * group + j; - const qv: u32 = if (bits == 4) (wb[(row * K + k) / 2] >> @intCast(4 * (k % 2))) & 15 else wb[row * K + k]; + const qv: u32 = switch (bits) { + 2 => (wb[(row * K + k) / 4] >> @intCast(2 * (k % 4))) & 3, + 4 => (wb[(row * K + k) / 2] >> @intCast(4 * (k % 2))) & 15, + else => wb[row * K + k], + }; const xv: f64 = f32of(x[r * K + k]); dot += @as(f64, @floatFromInt(qv)) * xv; dmag += (@as(f64, @floatFromInt(qv)) + off) * @abs(xv); From 6d2a62aa796947d1cbe41c2464603e74b556b321 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 04:11:00 -1000 Subject: [PATCH 06/10] tests: affine-ref, the prompt products against float64 on Metal, every row's bits in every window The device side of #633's quant-ref for the shared prompt products, item 5 of #716. tf-affine-ref takes no arguments and needs no model: seeded synthetic codes, scales, biases and inputs, packed as MLX affine. The engine module now names core/affine_mm so a check program can reach it. - affine_mm at 2, 4 and 8 bits, groups 32 and 64, bf16 and fp32 out, dense and as a one-expert gather, at [192, 1024] and [256, 4096]. - simd_qmm (4-bit, group 64) as tiles and as the scalar twin at 1 and 2 rows, at [40, 512], [256, 4096] and [6400, 256] (K chunks of 32, 16 and 8). The reference dequantizes the same codes in float64 and multiplies in float64. Each product prints its worst column (max abs error, and that error over the column's largest |y|) and its largest error as a fraction of the bound: (group + K / group + 8) fp32 roundings of the terms' magnitudes as the kernels sum them (4-bit affine_mm codes plus 128), plus one bf16 ulp for bf16 outputs. A reference with its scales 1% off must exceed the bound. Then every row of a 129-row call must equal, byte for byte, the same row computed in windows of 1, 2, 3, 8, 16, 17, 33, 64 and 65 rows (every full window over the 129 rows, so each row is seen at every position), which crosses the 16-row fragments, the 32- and 64-row tiles and the gather's switch between them. The twin must equal the tiles wherever twinExact says it does, since only then do callers take it. Receipts, M3 Ultra (no tensor units: nax.h's simdgroup-matrix layout), on main 0efbd447: - every product within 0.4935 of its bound; every row's bits equal at every window, both kernels, every format; the twin equals the tiles at 1 and 2 rows on all three shapes, as twinExact says. - worst columns, bf16 out (dense and gather equal), at [192, 1024] and [256, 4096]: affine_mm q2 g32 2.64e-2 (rel 3.23e-3) and 6.22e-2 (3.78e-3), q2 g64 2.87e-2 (3.51e-3) and 3.12e-2 (2.47e-3), q4 g32 2.45e-2 (2.65e-3) and 3.12e-2 (3.01e-3), q4 g64 1.56e-2 (2.89e-3) and 3.12e-2 (2.53e-3), q8 g32 0.496 (3.41e-3) and 0.968 (3.63e-3), q8 g64 0.453 (2.87e-3) and 0.956 (3.38e-3); simd_qmm 1.19e-2 (2.57e-3), 3.13e-2 (2.87e-3) and 1.54e-2 (3.74e-3). All under 2^-8 of the column's largest |y|; 0.20 to 0.49 of the bound. - fp32 out: at most 0.0081 of the bound (q2 1.5e-6 to 7.7e-6, rel 1.9e-7 to 6.0e-7). - scales 1% off: 2.3 to 29 times the bound for bf16 outputs, 2.7 to 391 for fp32, 44 to 273 for simd_qmm. Not run: an M5, where affine_mm takes the tensor units; `zig build tf-affine-ref && ./zig-out/bin/tf-affine-ref`. The fp32 bound is the worst case for any summation order, so a changed order still passes; the measured errors sit two orders below it, and the bit checks are what pin the order. Gate on main 0efbd447 (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test: `get` logs an error on the DeferredTensor path the test expects). lean_check: 536 problems (main 539), no new findings. --- build.zig | 1 + zig/build/core_checks.zig | 14 ++ zig/src/tensorfold.zig | 1 + zig/tests/affine_ref.zig | 358 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 374 insertions(+) create mode 100644 zig/tests/affine_ref.zig diff --git a/build.zig b/build.zig index 2ba6052374..d3b8a0a760 100644 --- a/build.zig +++ b/build.zig @@ -476,6 +476,7 @@ fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b // The core lane, shared-context and Qwen3.8-27B programs and host tests. qwen27_build.targets(b, target, optimize, metal, sources, engine, lanes, mods.qwen27, test_step); + core_checks.targets(b, target, optimize, metal, engine); _ = nativeServer(b, target, optimize, metal, engine, api, engines, build_options, test_step); diff --git a/zig/build/core_checks.zig b/zig/build/core_checks.zig index 75f39c00c2..6749819464 100644 --- a/zig/build/core_checks.zig +++ b/zig/build/core_checks.zig @@ -16,3 +16,17 @@ pub fn rowWidths(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std. const about = "The row projection at every MLX width and group: fp64 reference, windows, strides, stock bytes"; b.step("tf-row-widths-check", about).dependOn(&b.addInstallArtifact(exe, .{}).step); } + +/// tf-affine-ref: the prompt products (affine_mm, simd_qmm) against float64, every row's bits in every window. +pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, metal: *std.Build.Module, engine: *std.Build.Module) void { + const exe = b.addExecutable(.{ .name = "tf-affine-ref", .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/affine_ref.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "tensorfold", .module = engine } }, + }) }); + b.installArtifact(exe); + const about = "The prompt products (affine_mm, simd_qmm) on synthetic weights against float64, every row's bits in every window"; + b.step("tf-affine-ref", about).dependOn(&b.addInstallArtifact(exe, .{}).step); +} diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 1823eb56a9..f34f41a4ff 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -13,6 +13,7 @@ pub const ids_json = @import("core/ids_json.zig"); pub const lanes = @import("lanes"); pub const segments = @import("core/segments.zig"); pub const row_projection = @import("core/row_projection.zig"); +pub const affine_mm = @import("core/affine_mm.zig"); pub const exl3_format = @import("core/exl3_format.zig"); pub const exl3_projection = @import("core/exl3_projection.zig"); pub const exlr8 = @import("core/exlr8.zig"); diff --git a/zig/tests/affine_ref.zig b/zig/tests/affine_ref.zig new file mode 100644 index 0000000000..5f32bd864a --- /dev/null +++ b/zig/tests/affine_ref.zig @@ -0,0 +1,358 @@ +//! The shared prompt products against float64 on synthetic weights, and every row's bits in every window width. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const mm = tf.affine_mm; +const sq = tf.row_projection.simd; + +const opts = mtl.ResourceOptions.shared; +const a = std.heap.page_allocator; +const unit: f64 = 0x1p-24; // fp32's unit roundoff +const wide = 129; // rows of the reference call: two 64-row tiles and one more +const windows = [_]usize{ 1, 2, 3, 8, 16, 17, 33, 64, 65 }; // 17, 33 and 65 cross the 16-, 32- and 64-row tiles + +fn toBf(v: f32) u16 { + const u: u32 = @bitCast(v); + return @truncate((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); +} + +fn fromBf(v: u16) f64 { + return @as(f32, @bitCast(@as(u32, v) << 16)); +} + +/// Synthetic MLX affine weights [n, k] (codes little-endian in bytes) and `wide` rows of x; `deq` is s q + b exactly. +const Synth = struct { + n: usize, + k: usize, + bits: u8, + group: usize, + x: mtl.Buffer, + w: mtl.Buffer, + s: mtl.Buffer, + b: mtl.Buffer, + deq: []f64, + + fn init(device: mtl.Device, n: usize, k: usize, bits: u8, group: usize, seed: u64) !Synth { + var prng = std.Random.DefaultPrng.init(seed); + const r = prng.random(); + const sn = n * (k / group); + const self = Synth{ .n = n, .k = k, .bits = bits, .group = group, .x = try device.buffer(wide * k * 2, opts), .w = try device.buffer(n * k * bits / 8, opts), .s = try device.buffer(sn * 2, opts), .b = try device.buffer(sn * 2, opts), .deq = try a.alloc(f64, n * k) }; + for (self.x.slice(u16, wide * k)) |*v| v.* = toBf(r.float(f32) * 2 - 1); + r.bytes(self.w.slice(u8, n * k * bits / 8)); + for (self.s.slice(u16, sn), self.b.slice(u16, sn)) |*s, *b| { + s.* = toBf((r.float(f32) + 0.5) / 64); + b.* = toBf(-(r.float(f32) + 0.5) * 0.15); + } + for (0..n) |j| for (0..k) |i| { + const g = j * (k / group) + i / group; + self.deq[j * k + i] = fromBf(self.s.slice(u16, sn)[g]) * @as(f64, @floatFromInt(self.code(j, i))) + fromBf(self.b.slice(u16, sn)[g]); + }; + return self; + } + + fn deinit(self: Synth) void { + for ([_]mtl.Buffer{ self.x, self.w, self.s, self.b }) |x| x.deinit(); + a.free(self.deq); + } + + fn code(self: Synth, j: usize, i: usize) u32 { + const bytes = self.w.slice(u8, self.n * self.k * self.bits / 8); + if (self.bits == 8) return bytes[j * self.k + i]; + if (self.bits == 2) return (bytes[(j * self.k + i) / 4] >> @intCast(2 * (i % 4))) & 3; + return (bytes[(j * self.k + i) / 2] >> @intCast(4 * (i % 2))) & 15; + } + + /// y = x W^T in float64, and the sizes the fp32 sums carry when codes enter as q + off (affine_mm's 4-bit: 128). + fn reference(self: Synth, off: f64) !Ref { + const n = self.n; + const k = self.k; + const groups = k / self.group; + const ref = Ref{ .want = try a.alloc(f64, wide * n), .mag = try a.alloc(f64, wide * n), .scaled = try a.alloc(f64, wide * n) }; + const x = self.x.slice(u16, wide * k); + const sc = self.s.slice(u16, n * groups); + const bi = self.b.slice(u16, n * groups); + const xd = try a.alloc(f64, k); + defer a.free(xd); + const xg = try a.alloc(f64, groups); // each group's sum of |x| + defer a.free(xg); + for (0..wide) |r| { + @memset(xg, 0); + for (xd, 0..) |*v, i| { + v.* = fromBf(x[r * k + i]); + xg[i / self.group] += @abs(v.*); + } + for (0..n) |j| { + var sum: f64 = 0; + var mag: f64 = 0; + var scaled: f64 = 0; + for (0..groups) |g| { + var qx: f64 = 0; + var qxs: f64 = 0; + for (g * self.group..(g + 1) * self.group) |i| { + const q: f64 = @floatFromInt(self.code(j, i)); + sum += xd[i] * self.deq[j * k + i]; + qx += (q + off) * @abs(xd[i]); + qxs += q * xd[i]; + } + const s = fromBf(sc[j * groups + g]); + mag += s * qx + @abs(fromBf(bi[j * groups + g]) - off * s) * xg[g]; + scaled += s * qxs; + } + ref.want[r * n + j] = sum; + ref.mag[r * n + j] = mag; + ref.scaled[r * n + j] = scaled; + } + } + return ref; + } +}; + +/// The float64 product, its terms' fp32 magnitudes, and its scales' share (a control: those scales 1% off must fail). +const Ref = struct { + want: []f64, + mag: []f64, + scaled: []f64, + fn deinit(r: Ref) void { + for ([_][]f64{ r.want, r.mag, r.scaled }) |v| a.free(v); + } +}; + +/// Every check's outcome: failures, and the largest difference as a fraction of its bound. +const Tally = struct { failed: usize = 0, worst: f64 = 0 }; + +/// One product's output: `rows` rows of n bf16 or fp32 values. +const Out = struct { + buf: mtl.Buffer, + f32_out: bool, + fn esize(o: Out) usize { + return if (o.f32_out) 4 else 2; + } + fn get(o: Out, i: usize) f64 { + return if (o.f32_out) o.buf.slice(f32, i + 1)[i] else fromBf(o.buf.slice(u16, i + 1)[i]); + } +}; + +/// The bound: fp32 sums of a group's products, its row sum and a chain of groups, plus one bf16 ulp of the value. +fn bound(want: f64, mag: f64, s: *const Synth, f32_out: bool, extra: usize) f64 { + const terms: f64 = @floatFromInt(s.group + s.k / s.group + 8 + extra); + return terms * unit * mag + (if (f32_out) 0 else @abs(want) * 0x1p-7) + 1e-30; +} + +/// `rows` rows of `out` against the float64 reference: the worst column's max abs and relative error, and the bound. +fn compare(t: *Tally, name: []const u8, out: Out, rows: usize, ref: Ref, s: *const Synth, extra: usize) void { + const n = s.n; + var col: usize = 0; + var col_abs: f64 = -1; + var col_rel: f64 = 0; + var frac: f64 = 0; + var ctl: f64 = 0; + var at: usize = 0; + for (0..n) |j| { + var abs: f64 = 0; + var top: f64 = 0; + for (0..rows) |r| { + const i = r * n + j; + const d = @abs(out.get(i) - ref.want[i]); + const tol = bound(ref.want[i], ref.mag[i], s, out.f32_out, extra); + const q = d / tol; + ctl = @max(ctl, @abs(out.get(i) - ref.want[i] - 0.01 * ref.scaled[i]) / tol); + if (!(d <= abs)) abs = d; + top = @max(top, @abs(ref.want[i])); + if (!(q <= frac)) { + frac = q; + at = i; + } + } + if (!(abs <= col_abs)) { + col_abs = abs; + col = j; + col_rel = abs / top; + } + } + t.worst = @max(t.worst, frac); + const ok = frac <= 1 and ctl > 1; + if (!ok) t.failed += 1; + std.debug.print("RESULT {s}: worst column {d}: max abs {e:.3}, relative {e:.3}; {d:.4} of its bound (row {d} col {d}), " ++ + "scales 1% off {d:.1}{s}\n", .{ name, col, col_abs, col_rel, frac, at / n, at % n, ctl, if (ok) "" else " FAIL" }); +} + +/// Rows [0, rows) of `got` equal `want` byte for byte; the first difference printed. +fn sameRows(t: *Tally, name: []const u8, m: usize, want: Out, got: Out, rows: usize, n: usize) bool { + const bytes = rows * n * want.esize(); + const w = want.buf.slice(u8, bytes); + const g = got.buf.slice(u8, bytes); + if (std.mem.eql(u8, w, g)) return true; + const i = std.mem.indexOfDiff(u8, w, g).? / want.esize(); + var differ: usize = 0; + var most: f64 = 0; + for (0..rows) |r| { + const row = r * n * want.esize(); + if (!std.mem.eql(u8, w[row..][0 .. n * want.esize()], g[row..][0 .. n * want.esize()])) differ += 1; + for (r * n..(r + 1) * n) |e| most = @max(most, @abs(got.get(e) - want.get(e))); + } + std.debug.print("RESULT {s}: {d}-row windows DIFFER in {d} of {d} rows, by up to {e:.3}; first row {d} col {d}: {d} against {d} in the {d}-row call\n", .{ name, m, differ, rows, most, i / n, i % n, got.get(i), want.get(i), wide }); + t.failed += 1; + return false; +} + +fn newOut(device: mtl.Device, s: *const Synth, f32_out: bool) !Out { + const o = Out{ .buf = try device.buffer(wide * s.n * 4, opts), .f32_out = f32_out }; + @memset(o.buf.slice(u8, wide * s.n * 4), 0xff); // NaN: a launch that writes nothing fails + return o; +} + +fn commit(cb: anytype, e: mtl.ComputeEncoder) !void { + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |why| { + std.debug.print("GPU failed: {s}\n", .{why}); + return error.GpuFailed; + } +} + +const Kernel = enum { dense, gather }; + +/// affine_mm over rows [start, start + m): its row sums, then the dense product or a gather with one expert. +fn affineWindow(e: mtl.ComputeEncoder, p: mm.Pipes, s: *const Synth, sums: mtl.Buffer, offs: mtl.Buffer, y: Out, kernel: Kernel, start: usize, m: usize) void { + const x = .{ .buf = s.x, .off = start * s.k * 2 }; + const xs = .{ .buf = sums, .off = start * (s.k / s.group) * 4 }; + const yo = .{ .buf = y.buf, .off = start * s.n * y.esize() }; + const z: usize = 0; + const q = .{ .w = .{ .buf = s.w, .off = z }, .s = .{ .buf = s.s, .off = z }, .b = .{ .buf = s.b, .off = z }, .n = @as(u32, @intCast(s.n)), .k = @as(u32, @intCast(s.k)) }; + mm.rowSums(e, p, @intCast(s.group), x, xs, @intCast(m), @intCast(s.k)); + switch (kernel) { + .dense => mm.dense(e, p, x, xs, q, yo, @intCast(m)), + .gather => mm.gather(e, p, x, xs, q, .{ .buf = offs, .off = z }, yo, @intCast(m), 1), + } +} + +/// affine_mm at one format: the wide call against float64, then every window's rows against the wide call's. +fn affine(t: *Tally, device: mtl.Device, queue: mtl.Queue, s: *const Synth, ref: Ref, f32_out: bool) !void { + const src = try mm.source(device, a, .{ .bits = s.bits, .group = @intCast(s.group), .out_f32 = f32_out }); + defer a.free(src); + const lib = try mtl.Library.fromSource(device, src, mtl.CompileOptions.mlx()); + defer lib.deinit(); + var p: mm.Pipes = undefined; + for (mm.names, 0..) |name, i| p[i] = try mtl.Pipeline.init(device, lib, name, false); + defer for (p) |x| x.deinit(); + const sums = try device.buffer(wide * (s.k / s.group) * 4, opts); + defer sums.deinit(); + const offs = try device.buffer(4, opts); + defer offs.deinit(); + offs.slice(i32, 1)[0] = 0; + var name: [96]u8 = undefined; + for ([_]Kernel{ .dense, .gather }) |kernel| { + const label = try std.fmt.bufPrint(&name, "affine_mm {s} q{d} g{d} {s} [{d}, {d}]", .{ @tagName(kernel), s.bits, s.group, if (f32_out) "fp32" else "bf16", s.n, s.k }); + const all = try newOut(device, s, f32_out); + defer all.buf.deinit(); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + affineWindow(e, p, s, sums, offs, all, kernel, 0, wide); + try commit(cb, e); + compare(t, label, all, wide, ref, s, 0); + var same = true; + for (windows) |m| { + const got = try newOut(device, s, f32_out); + defer got.buf.deinit(); + const wb = queue.commandBuffer(); + const we = wb.compute(.serial); + var start: usize = 0; + while (start + m <= wide) : (start += m) affineWindow(we, p, s, sums, offs, got, kernel, start, m); + try commit(wb, we); + same = sameRows(t, label, m, all, got, wide / m * m, s.n) and same; + } + if (same) std.debug.print("RESULT {s}: every row's bits equal in windows of 1, 2, 3, 8, 16, 17, 33, 64 and 65 rows\n", .{label}); + } +} + +/// simd_qmm over rows [start, start + m), the scalar twin when `scalar` and it takes m rows. +fn simdWindow(e: mtl.ComputeEncoder, p: *const sq.Pipelines, s: *const Synth, y: Out, scalar: bool, start: usize, m: usize) !void { + const mat = sq.Matrix{ .w = s.w, .scales = s.s, .biases = s.b, .n = s.n, .k = s.k }; + sq.encode(e, try sq.call(p, s.n, s.k, m, scalar), mat, s.x, start * s.k * 2, y.buf, start * s.n * 2); +} + +/// simd_qmm (4-bit, group 64): tiles at every window against float64 and the wide call; the twin at 1 and 2 rows. +fn simd(t: *Tally, device: mtl.Device, queue: mtl.Queue, p: *const sq.Pipelines, s: *const Synth, ref: Ref) !void { + var name: [96]u8 = undefined; + const label = try std.fmt.bufPrint(&name, "simd_qmm tiles q4 g64 bf16 [{d}, {d}]", .{ s.n, s.k }); + const all = try newOut(device, s, false); + defer all.buf.deinit(); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + try simdWindow(e, p, s, all, false, 0, wide); + try commit(cb, e); + const extra = std.math.log2_int(u32, sq.splits(s.n)) + 1; // the chunks' fixed tree + compare(t, label, all, wide, ref, s, extra); + var same = true; + for (windows) |m| { + const got = try newOut(device, s, false); + defer got.buf.deinit(); + const wb = queue.commandBuffer(); + const we = wb.compute(.serial); + var start: usize = 0; + while (start + m <= wide) : (start += m) try simdWindow(we, p, s, got, false, start, m); + try commit(wb, we); + same = sameRows(t, label, m, all, got, wide / m * m, s.n) and same; + } + if (same) std.debug.print("RESULT {s}: every row's bits equal in windows of 1, 2, 3, 8, 16, 17, 33, 64 and 65 rows\n", .{label}); + // the twin: callers take it only where twinExact says it equals the tiles, so a difference fails only then + const exact = try sq.twinExact(p, device, .{ .w = s.w, .scales = s.s, .biases = s.b, .n = s.n, .k = s.k }); + for ([_]usize{ 1, 2 }) |m| { + const twin = try std.fmt.bufPrint(&name, "simd_qmm twin q4 g64 bf16 [{d}, {d}] at {d} rows", .{ s.n, s.k, m }); + const got = try newOut(device, s, false); + defer got.buf.deinit(); + const wb = queue.commandBuffer(); + const we = wb.compute(.serial); + var start: usize = 0; + while (start + m <= wide) : (start += m) try simdWindow(we, p, s, got, true, start, m); + try commit(wb, we); + compare(t, twin, got, wide / m * m, ref, s, 0); + const rows = wide / m * m; + const eq = std.mem.eql(u8, all.buf.slice(u8, rows * s.n * 2), got.buf.slice(u8, rows * s.n * 2)); + std.debug.print("RESULT {s}: {s} the tiles' bits; twinExact says {s}{s}\n", .{ twin, if (eq) "equals" else "differs from", if (exact) "equal" else "differ", if (exact and !eq) " FAIL" else "" }); + if (exact and !eq) t.failed += 1; + } +} + +pub fn main(_: std.process.Init) !void { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + std.debug.print("device {s}, tensor units {}\n", .{ device.name(), device.tensorUnits() }); + var t = Tally{}; + var seed: u64 = 0x716; + for ([_][2]usize{ .{ 192, 1024 }, .{ 256, 4096 } }) |shape| for ([_]u8{ 4, 8 }) |bits| for ([_]usize{ 32, 64 }) |group| { + seed += 1; + const s = try Synth.init(device, shape[0], shape[1], bits, group, seed); + defer s.deinit(); + const ref = try s.reference(if (bits == 4) 128 else 0); + defer ref.deinit(); + for ([_]bool{ false, true }) |f32_out| try affine(&t, device, queue, &s, ref, f32_out); + }; + for ([_][2]usize{ .{ 192, 1024 }, .{ 256, 4096 } }) |shape| for ([_]usize{ 32, 64 }) |group| { // 2 bits: oQ2's experts + const s = try Synth.init(device, shape[0], shape[1], 2, group, 0x2716 + shape[0] + group); + defer s.deinit(); + const ref = try s.reference(0); + defer ref.deinit(); + for ([_]bool{ false, true }) |f32_out| try affine(&t, device, queue, &s, ref, f32_out); + }; + var p = try sq.Pipelines.load(device); + defer p.deinit(); + for ([_][2]usize{ .{ 40, 512 }, .{ 256, 4096 }, .{ 6400, 256 } }) |shape| { // chunks of 32, 16 and 8 + seed += 1; + const s = try Synth.init(device, shape[0], shape[1], 4, 64, seed); + defer s.deinit(); + const ref = try s.reference(0); + defer ref.deinit(); + try simd(&t, device, queue, &p, &s, ref); + } + if (t.failed > 0) { + std.debug.print("FAIL {d} checks; the largest difference is {d:.4} of its float64 bound\n", .{ t.failed, t.worst }); + return error.CheckFailed; + } + std.debug.print("PASS every product within {d:.4} of its float64 bound, every row's bits equal at every window\n", .{t.worst}); +} From d6d6a6136970d85590f2515ebba737a1d6a91158 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 04:41:38 -1000 Subject: [PATCH 07/10] Flash Next: the depth rule and copy drafts move to drafts.zig, the rounds buffers to engine_rounds.zig Moved, to make room under replay.zig's and engine.zig's lengths for the widths that follow: DepthRule, its test and copyDrafts (re-exported from replay.zig, the test reached from replay's test block), and the body of Engine.rounds (Engine keeps a one-line rounds, so the load test that reads the source still finds `try e.rounds();`; RING is now pub for it). Each moved doc comment is one line now; no behaviour changes. replay.zig 3321 -> 3265 lines, engine.zig 899 -> 857. Gate on main 0efbd447 + the core commits (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test logs an error on the path it expects); native and tf-flashnext-run (with flashnext_bench), -reuse, -pack and -think build. lean_check: 531 problems (main 539); the only findings not on main are the two long files' shorter line counts. --- zig/src/families/flashnext/drafts.zig | 59 +++++++++++++++++++ zig/src/families/flashnext/engine.zig | 48 +-------------- zig/src/families/flashnext/engine_rounds.zig | 57 ++++++++++++++++++ zig/src/families/flashnext/replay.zig | 62 +------------------- 4 files changed, 122 insertions(+), 104 deletions(-) create mode 100644 zig/src/families/flashnext/drafts.zig create mode 100644 zig/src/families/flashnext/engine_rounds.zig diff --git a/zig/src/families/flashnext/drafts.zig b/zig/src/families/flashnext/drafts.zig new file mode 100644 index 0000000000..34f5d72c63 --- /dev/null +++ b/zig/src/families/flashnext/drafts.zig @@ -0,0 +1,59 @@ +//! The decode's draft widths: the depth rule over recent landing, and copy drafts from the token history. +const std = @import("std"); + +/// Round width from recent landing: 3 (prose), 6 while drafts land (code), the widest while nearly all land. +pub const DepthRule = struct { + rate: f64 = 0.6, // moving average of the head's drafts landed over drafts offered + depth: usize = 3, + pair: bool = false, // speed-up mode: even windows only (an odd window's last row costs a Mac a whole expert row) + + pub fn pick(self: *const DepthRule) usize { + return self.depth; + } + + /// A head round's drafts and how many landed: 3, 6 or 8 on one Mac, 3 or 7 on the pair. + pub fn update(self: *DepthRule, depth: usize, landed: usize) void { + self.rate = 0.7 * self.rate + 0.3 * @as(f64, @floatFromInt(landed)) / @as(f64, @floatFromInt(depth)); + if (self.pair) { + if (self.depth == 3 and self.rate > 0.95) self.depth = 7 else if (self.depth == 7 and self.rate < 0.8) self.depth = 3; + return; + } + if (self.depth == 3 and self.rate > 0.8) { + self.depth = 6; + } else if (self.depth == 6 and self.rate > 0.74) { + self.depth = 8; + } else if (self.depth == 8 and self.rate < 0.62) { + self.depth = 6; + } else if (self.depth == 6 and self.rate < 0.65) self.depth = 3; + } +}; + +test "the pair's depth rule keeps every window even" { + var pair: DepthRule = .{ .pair = true }; + for (0..20) |_| pair.update(pair.pick(), pair.pick()); + try std.testing.expectEqual(@as(usize, 7), pair.pick()); + for (0..20) |_| pair.update(pair.pick(), 1); + try std.testing.expectEqual(@as(usize, 3), pair.pick()); + var one: DepthRule = .{}; + for (0..20) |_| one.update(one.pick(), one.pick()); + try std.testing.expectEqual(@as(usize, 8), one.pick()); +} + +/// Copies into `out` what followed the latest earlier match of `hist`'s longest repeated suffix (`min`..8), or 0. +pub fn copyDrafts(hist: []const u32, min: usize, out: []u32) usize { + const n = hist.len; + if (n < min + 1) return 0; + var len: usize = @min(8, n - 1); + while (len >= min) : (len -= 1) { + const suffix = hist[n - len ..]; + var end: usize = n - 1; + while (end >= len) : (end -= 1) { + if (std.mem.eql(u32, hist[end - len .. end], suffix)) { + const k = @min(out.len, n - end); + @memcpy(out[0..k], hist[end .. end + k]); + return k; + } + } + } + return 0; +} diff --git a/zig/src/families/flashnext/engine.zig b/zig/src/families/flashnext/engine.zig index d588ec1b32..16b20695dd 100644 --- a/zig/src/families/flashnext/engine.zig +++ b/zig/src/families/flashnext/engine.zig @@ -49,7 +49,7 @@ const f32Buf = fz.f32Buf; const jsonInt = fz.jsonInt; /// Rounds the token ring holds (fz_accept writes round & 511). -const RING = 512; +pub const RING = 512; /// Positions a reply keeps free past its last token: two rounds in flight. pub const MARGIN = 2 * MAXR; /// Bytes a conversation's next kept state usually adds (about 4k tokens of reply and tool output): the pool readies that much more. @@ -433,51 +433,9 @@ pub const Engine = struct { return e; } - /// The GPU-side rounds' buffers: DeltaNet states (kept, window rows), the PLE tail pair, the arena, the token ring, - /// per-width row counts, and the MTP slots reading their positions from the arena. + /// The GPU-side rounds' buffers (engine_rounds.zig). fn rounds(e: *Engine) !void { - const r = e.r; - const m = e.m; - const n_lin: usize = 36; - e.g_cs = try r.buffer(n_lin * CS_ROW); - e.g_so = try r.buffer(n_lin * SO_ROW); - e.o_cs = try r.buffer(n_lin * MAXR * CS_ROW); - e.o_so = try r.buffer(n_lin * MAXR * SO_ROW); - var gi: usize = 0; - for (&m.layers) |*L| if (L.linear) { - L.cs[0] = .{ .b = e.g_cs, .off = gi * CS_ROW }; - L.cs[1] = .{ .b = e.o_cs, .off = gi * MAXR * CS_ROW }; - L.so[0] = .{ .b = e.g_so, .off = gi * SO_ROW }; - L.so[1] = .{ .b = e.o_so, .off = gi * MAXR * SO_ROW }; - gi += 1; - }; - if (r.gdn_kept != null) m.recs = .{ .{ .b = try r.buffer(n_lin * fz.gdn_step.RECORD) }, .{ .b = try r.buffer(n_lin * fz.gdn_step.RECORD) } }; - m.marks = try fz.Marks.init(r, n_lin); - e.cins = .{ m.ple.cin, .{ .b = try r.buffer((PLE_TAIL + MAXR) * WIDE * 2) } }; - r.ar = .{ .b = try r.buffer(4 * fz.AR_WORDS) }; - e.ring = try r.buffer(fz.RING_WORDS * 4 * RING); - e.hist = try r.buffer((CAP + 64) * 4); - e.wids = .{ .b = try r.buffer(64) }; - for (1..MAXR + 1) |n| { - e.rows_w[n] = try i32Buf(r, &.{@intCast(n)}); - e.mdims_w[n] = try i32Buf(r, &.{ @intCast(n), 16, 0, 0, 0, 0, 0, 0 }); - } - e.host = .{ .rows = m.t.rows, .mdims = m.t.mdims, .pos8 = m.t.pos8, .nk8 = m.t.nk8, .kvmeta = m.t.kvmeta, .slots = m.mtp.slots }; - e.gpu_slots = m.mtp.slots; - for (1..MAXR + 1) |n| { // the head absorbing a window of n rows: slot MAXR - 1 + n - const sl = &e.gpu_slots[MAXR - 1 + n]; - try m.mtpMeta(sl, n); - sl.pos8 = .{ .b = r.ar.b, .off = fz.AR_ABS_POS * 4 }; - sl.nk8 = .{ .b = r.ar.b, .off = fz.AR_ABS_NK * 4 }; - sl.kvmeta = .{ .b = r.ar.b, .off = fz.AR_ABS_KV * 4 }; - } - for (1..MAXR - 1) |j| { // chained draft j: slot j - try m.mtpMeta(&e.gpu_slots[j], 1); - const b = (fz.AR_CHAIN + (j - 1) * fz.AR_CHAIN_STRIDE) * 4; - e.gpu_slots[j].pos8 = .{ .b = r.ar.b, .off = b }; - e.gpu_slots[j].nk8 = .{ .b = r.ar.b, .off = b + 16 * 4 }; - e.gpu_slots[j].kvmeta = .{ .b = r.ar.b, .off = b + 32 * 4 }; - } + return @import("engine_rounds.zig").rounds(e); } /// Host-driven calls (prompt chunks, the head's first drafts) read the host's metadata buffers. diff --git a/zig/src/families/flashnext/engine_rounds.zig b/zig/src/families/flashnext/engine_rounds.zig new file mode 100644 index 0000000000..a016f629e8 --- /dev/null +++ b/zig/src/families/flashnext/engine_rounds.zig @@ -0,0 +1,57 @@ +//! The GPU-side rounds' buffers an engine sets up once at load. +const fz = @import("replay.zig"); +const engine = @import("engine.zig"); +const Engine = engine.Engine; +const MAXR = fz.MAXR; +const CS_ROW = fz.CS_ROW; +const SO_ROW = fz.SO_ROW; +const PLE_TAIL = fz.PLE_TAIL; +const WIDE = fz.WIDE; +const CAP = fz.CAP; +const i32Buf = fz.i32Buf; + +/// DeltaNet states (kept, window rows), PLE tail pair, arena, token ring, per-width row counts, MTP slots. +pub fn rounds(e: *Engine) !void { + const r = e.r; + const m = e.m; + const n_lin: usize = 36; + e.g_cs = try r.buffer(n_lin * CS_ROW); + e.g_so = try r.buffer(n_lin * SO_ROW); + e.o_cs = try r.buffer(n_lin * MAXR * CS_ROW); + e.o_so = try r.buffer(n_lin * MAXR * SO_ROW); + var gi: usize = 0; + for (&m.layers) |*L| if (L.linear) { + L.cs[0] = .{ .b = e.g_cs, .off = gi * CS_ROW }; + L.cs[1] = .{ .b = e.o_cs, .off = gi * MAXR * CS_ROW }; + L.so[0] = .{ .b = e.g_so, .off = gi * SO_ROW }; + L.so[1] = .{ .b = e.o_so, .off = gi * MAXR * SO_ROW }; + gi += 1; + }; + if (r.gdn_kept != null) m.recs = .{ .{ .b = try r.buffer(n_lin * fz.gdn_step.RECORD) }, .{ .b = try r.buffer(n_lin * fz.gdn_step.RECORD) } }; + m.marks = try fz.Marks.init(r, n_lin); + e.cins = .{ m.ple.cin, .{ .b = try r.buffer((PLE_TAIL + MAXR) * WIDE * 2) } }; + r.ar = .{ .b = try r.buffer(4 * fz.AR_WORDS) }; + e.ring = try r.buffer(fz.RING_WORDS * 4 * engine.RING); + e.hist = try r.buffer((CAP + 64) * 4); + e.wids = .{ .b = try r.buffer(64) }; + for (1..MAXR + 1) |n| { + e.rows_w[n] = try i32Buf(r, &.{@intCast(n)}); + e.mdims_w[n] = try i32Buf(r, &.{ @intCast(n), 16, 0, 0, 0, 0, 0, 0 }); + } + e.host = .{ .rows = m.t.rows, .mdims = m.t.mdims, .pos8 = m.t.pos8, .nk8 = m.t.nk8, .kvmeta = m.t.kvmeta, .slots = m.mtp.slots }; + e.gpu_slots = m.mtp.slots; + for (1..MAXR + 1) |n| { // the head absorbing a window of n rows: slot MAXR - 1 + n + const sl = &e.gpu_slots[MAXR - 1 + n]; + try m.mtpMeta(sl, n); + sl.pos8 = .{ .b = r.ar.b, .off = fz.AR_ABS_POS * 4 }; + sl.nk8 = .{ .b = r.ar.b, .off = fz.AR_ABS_NK * 4 }; + sl.kvmeta = .{ .b = r.ar.b, .off = fz.AR_ABS_KV * 4 }; + } + for (1..MAXR - 1) |j| { // chained draft j: slot j + try m.mtpMeta(&e.gpu_slots[j], 1); + const b = (fz.AR_CHAIN + (j - 1) * fz.AR_CHAIN_STRIDE) * 4; + e.gpu_slots[j].pos8 = .{ .b = r.ar.b, .off = b }; + e.gpu_slots[j].nk8 = .{ .b = r.ar.b, .off = b + 16 * 4 }; + e.gpu_slots[j].kvmeta = .{ .b = r.ar.b, .off = b + 32 * 4 }; + } +} diff --git a/zig/src/families/flashnext/replay.zig b/zig/src/families/flashnext/replay.zig index f78fa08c28..ad35200ccc 100644 --- a/zig/src/families/flashnext/replay.zig +++ b/zig/src/families/flashnext/replay.zig @@ -1414,6 +1414,7 @@ pub fn catchSlot(words: []i32, k: usize, first: usize) usize { test { _ = tpm; // tp.zig's tests: the window layout and the host service's wraps + _ = @import("drafts.zig"); // the depth rule's tests } test "catch-up slots: every layer's dispatch reads its own first block after the later layers' writes" { @@ -2274,68 +2275,11 @@ pub const Model = struct { } }; -/// Drafts a round from recent landing: 3 (prose), 6 while drafts land (code-like text), the widest window while -/// nearly all land (copied text). /// FZ_PROFILE's knock-out bit for speed-up mode's exchanges: the bit after `Run.class`'s eight. pub const TP_CLASS: u32 = 1 << 8; -pub const DepthRule = struct { - rate: f64 = 0.6, // moving average of the head's drafts landed over drafts offered - depth: usize = 3, - pair: bool = false, // speed-up mode: even windows only (an odd window's last row costs a Mac a whole expert row) - - pub fn pick(self: *const DepthRule) usize { - return self.depth; - } - - /// A head round's drafts and how many landed: 3, 6 or 8 drafts on one Mac, 3 or 7 on the pair (copied rounds set - /// their own width). - pub fn update(self: *DepthRule, depth: usize, landed: usize) void { - self.rate = 0.7 * self.rate + 0.3 * @as(f64, @floatFromInt(landed)) / @as(f64, @floatFromInt(depth)); - if (self.pair) { - if (self.depth == 3 and self.rate > 0.95) self.depth = 7 else if (self.depth == 7 and self.rate < 0.8) self.depth = 3; - return; - } - if (self.depth == 3 and self.rate > 0.8) { - self.depth = 6; - } else if (self.depth == 6 and self.rate > 0.74) { - self.depth = 8; - } else if (self.depth == 8 and self.rate < 0.62) { - self.depth = 6; - } else if (self.depth == 6 and self.rate < 0.65) self.depth = 3; - } -}; - -test "the pair's depth rule keeps every window even" { - var pair: DepthRule = .{ .pair = true }; - for (0..20) |_| pair.update(pair.pick(), pair.pick()); - try std.testing.expectEqual(@as(usize, 7), pair.pick()); - for (0..20) |_| pair.update(pair.pick(), 1); - try std.testing.expectEqual(@as(usize, 3), pair.pick()); - var one: DepthRule = .{}; - for (0..20) |_| one.update(one.pick(), one.pick()); - try std.testing.expectEqual(@as(usize, 8), one.pick()); -} - -/// Copy lanes: the longest suffix of `hist` (`min`..8 tokens) seen earlier; the tokens after its latest earlier -/// occurrence go into `out`. Returns how many (0 when nothing matches). -pub fn copyDrafts(hist: []const u32, min: usize, out: []u32) usize { - const n = hist.len; - if (n < min + 1) return 0; - var len: usize = @min(8, n - 1); - while (len >= min) : (len -= 1) { - const suffix = hist[n - len ..]; - var end: usize = n - 1; - while (end >= len) : (end -= 1) { - if (std.mem.eql(u32, hist[end - len .. end], suffix)) { - const k = @min(out.len, n - end); - @memcpy(out[0..k], hist[end .. end + k]); - return k; - } - } - } - return 0; -} +pub const DepthRule = @import("drafts.zig").DepthRule; // drafts a round from recent landing +pub const copyDrafts = @import("drafts.zig").copyDrafts; /// Prompt chunks of up to PMAX rows: projections on the 6-bit tensor-unit kernels, experts sorted by expert and /// gathered, the decode's row kernels at the chunk's rows, and DeltaNet storing only its last row's state. From bfbea4178d2b0605871cb1563d19a7ac2862050f Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 04:42:39 -1000 Subject: [PATCH 08/10] Flash Next admits and packs each matrix at its own MLX width; the engine still runs 6-bit group 32 only Admission took one global 6-bit/group-32 width, and the packs refused group 128 and fused members of mixed widths. - Widths (widths.zig): each module's width through core/quant (module keys under any prefix; disagreeing spellings and half-declared overrides refused; `false` and `{}` mark plain modules). Admission reads MLX affine 2, 3, 4, 5, 6 or 8 bits in groups of 32, 64 or 128 for every text and MTP matrix, refuses the first other one under its checkpoint tensor name, and an n-gram table whose group does not divide its 160-wide rows. - Lane packs follow 0.6.6's _stacked: a fused projection's members grouped by bits in first-seen order, adjacent or not, each group one run at its smallest group, regrouped exactly (core/quant widen; group 128 read as 64), with .runs/.members tables. A run whose columns are not a multiple of 32 keeps MLX's layout untiled (decode.py's nt 0; oQ4e's layer-21 GDN input). A projection of one width keeps its single run. - Hyper-connections: a down and inject pair of two widths widens exactly to the wider (QWeights.of); the up projection packs at its stored width. Prompt (pack_mlx) projections of several widths keep each member's bytes in order with a .parts table. Experts, embedding and n-gram tables, which the engine reads straight from the checkpoint, are shape-checked at their declared widths when the pack is built. - oMLX's layout: an MTP head kept at top-level mtp.* is read where the index keeps it; an n-gram table's weight_scale (oQ conversions of the FP8 checkpoint keep its scale, as oQ4e does) must be one finite, nonzero value, else the pack is refused by name; the next commit applies it. - Pack format 3, whose identity also hashes config.json, so a config-only change rebuilds the pack; packs under the format-less identity are rebuilt, never reused. Pack building frees every buffer when an allocation fails. The uniform 6/32 packs keep every tensor byte for byte (a digest pinned before the change). The engine refuses any other width at load until the next commit gives the device its kernels. Tests that fail before: widths_test.zig; load_admission_test.zig's overrides at other widths, global groups 64/128, refusal under the tensor's name, oMLX's mtp.*; pack_lane_test.zig (runs by bits, a pack per run, ragged runs untiled at oQ4e's real rows 10240/6144/48/48, allocation failures free everything); pack_hc_test.zig (oQ4e's 4/64 down with a 5/64 inject as one 5/64 matrix; the up projection at its stored width; mixed prompt members); pack_widen_test.zig; pack_test.zig's half-declared and contradicting raw tensors; pack_source_test.zig's config-only change; pack_cache_test.zig's format-less identities. Gate on main 0efbd447 + the core commits (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test logs an error on the path it expects); native and tf-flashnext-run (with flashnext_bench), -reuse, -pack and -think build. lean_check: 531 problems (main 539); the only findings not on main are the two long files' shorter line counts. --- zig/src/families/flashnext/affine.zig | 4 - zig/src/families/flashnext/config.zig | 8 +- zig/src/families/flashnext/engine.zig | 22 +- zig/src/families/flashnext/host_test.zig | 6 +- zig/src/families/flashnext/index.zig | 77 ++++-- zig/src/families/flashnext/load_admission.zig | 30 ++- .../flashnext/load_admission_test.zig | 73 +++++- zig/src/families/flashnext/pack.zig | 239 ++++++++---------- .../families/flashnext/pack_cache_test.zig | 25 ++ zig/src/families/flashnext/pack_hash_test.zig | 47 ++++ zig/src/families/flashnext/pack_hc_test.zig | 72 ++++++ zig/src/families/flashnext/pack_io.zig | 51 ++-- zig/src/families/flashnext/pack_lane.zig | 188 ++++++++++++++ zig/src/families/flashnext/pack_lane_test.zig | 136 ++++++++++ zig/src/families/flashnext/pack_raw.zig | 42 +++ .../families/flashnext/pack_source_test.zig | 25 +- zig/src/families/flashnext/pack_test.zig | 33 +++ zig/src/families/flashnext/pack_widen.zig | 28 ++ .../families/flashnext/pack_widen_test.zig | 67 +++++ zig/src/families/flashnext/widths.zig | 76 ++++++ zig/src/families/flashnext/widths_test.zig | 85 +++++++ 21 files changed, 1132 insertions(+), 202 deletions(-) create mode 100644 zig/src/families/flashnext/pack_hash_test.zig create mode 100644 zig/src/families/flashnext/pack_hc_test.zig create mode 100644 zig/src/families/flashnext/pack_lane.zig create mode 100644 zig/src/families/flashnext/pack_lane_test.zig create mode 100644 zig/src/families/flashnext/pack_raw.zig create mode 100644 zig/src/families/flashnext/pack_widen.zig create mode 100644 zig/src/families/flashnext/pack_widen_test.zig create mode 100644 zig/src/families/flashnext/widths.zig create mode 100644 zig/src/families/flashnext/widths_test.zig diff --git a/zig/src/families/flashnext/affine.zig b/zig/src/families/flashnext/affine.zig index 6152735775..40aa21b977 100644 --- a/zig/src/families/flashnext/affine.zig +++ b/zig/src/families/flashnext/affine.zig @@ -18,10 +18,6 @@ pub const Spec = struct { if (count % 32 != 0) return error.InvalidAffineWidth; return count / 32; } - - pub fn checkFlashKernel(s: Spec) !void { - if (s.bits != 6 or s.group != 32) return error.UnsupportedFlashAffineKernel; - } }; pub const Shape = struct { experts: usize, n: usize, k: usize, words: usize, groups: usize, dtype: st.DType }; diff --git a/zig/src/families/flashnext/config.zig b/zig/src/families/flashnext/config.zig index 7241ee0b54..df990164e1 100644 --- a/zig/src/families/flashnext/config.zig +++ b/zig/src/families/flashnext/config.zig @@ -54,6 +54,8 @@ pub const Config = struct { eos: [16]u32 = @splat(0), eos_count: usize = 0, global_affine: affine.Spec, + /// Set by a tensor reader: each lookup then goes through core/quant, which fills in no missing bits or group. + widths: ?*const @import("widths.zig").Widths = null, pub fn deinit(c: *Config) void { c.parsed.deinit(); @@ -94,6 +96,10 @@ pub const Config = struct { } pub fn quantization(c: *const Config, path: []const u8) !?affine.Spec { + if (c.widths) |w| { + const width = (try w.of(path)) orelse return null; + return try affine.Spec.init(width.bits, width.group); + } const wanted = canonical(path); var found = false; var answer: ?affine.Spec = c.global_affine; @@ -113,7 +119,7 @@ pub const Config = struct { } }; -fn canonical(value: []const u8) []const u8 { +pub fn canonical(value: []const u8) []const u8 { var path = if (std.mem.endsWith(u8, value, ".weight")) value[0 .. value.len - 7] else value; for ([_][]const u8{ "model.language_model.", "language_model.", "text_model.", "model." }) |prefix| { if (std.mem.startsWith(u8, path, prefix)) path = path[prefix.len..]; diff --git a/zig/src/families/flashnext/engine.zig b/zig/src/families/flashnext/engine.zig index 16b20695dd..e8f86569f1 100644 --- a/zig/src/families/flashnext/engine.zig +++ b/zig/src/families/flashnext/engine.zig @@ -52,6 +52,8 @@ const jsonInt = fz.jsonInt; pub const RING = 512; /// Positions a reply keeps free past its last token: two rounds in flight. pub const MARGIN = 2 * MAXR; +/// An MoE layer's expert projections in the engine's order: gate, up, shared gate, shared up, down, shared down. +const expert_projs = [_][]const u8{ "switch_mlp.gate_proj", "switch_mlp.up_proj", "shared_expert.gate_proj", "shared_expert.up_proj", "switch_mlp.down_proj", "shared_expert.down_proj" }; /// Bytes a conversation's next kept state usually adds (about 4k tokens of reply and tool output): the pool readies that much more. pub const NEXT_TURN = 4096 * 13 * 2304; @@ -199,7 +201,9 @@ pub const Engine = struct { } const index_mapped = try mtl.MappedFile.open(try std.fmt.allocPrintSentinel(arena, "{s}/model.safetensors.index.json", .{model_dir}, 0)); try maps.append(arena, index_mapped); - const id = try pack_io.identityFromMapped(arena, index_mapped.bytes[0..index_mapped.size], mapped_shards.items); + const config_mapped = try mtl.MappedFile.open(try std.fmt.allocPrintSentinel(arena, "{s}/config.json", .{model_dir}, 0)); + try maps.append(arena, config_mapped); + const id = try pack_io.identityFromMapped(arena, index_mapped.bytes[0..index_mapped.size], config_mapped.bytes[0..config_mapped.size], mapped_shards.items); for (maps.items) |*m| m.deinit(); maps.clearRetainingCapacity(); break :blk id; @@ -223,6 +227,12 @@ pub const Engine = struct { // the reference values: the dump's ref.json, or derived from the checkpoint's config (config.pleRef) var ple_values: config.PleRef = undefined; var attention_scale: f64 = undefined; + var cfg = try config.Config.read(arena, io, model_dir); // read once: the reference values and the widths + var cfg_widths = try @import("widths.zig").Widths.init(arena, &cfg); + if (!cfg_widths.uniform() or cfg_widths.q.global.bits != 6 or cfg_widths.q.global.group != 32) { + std.log.warn("flashnext: widths other than 6-bit group 32; the device kernels read 6/32 only", .{}); + return error.UnsupportedFlashAffineKernel; + } if (dump_dir) |d| { const ref_file = try mtl.MappedFile.open(try std.fmt.allocPrintSentinel(arena, "{s}/ref.json", .{d}, 0)); const ref = try std.json.parseFromSliceLeaky(std.json.Value, arena, ref_file.bytes[0..ref_file.size], .{}); @@ -243,7 +253,6 @@ pub const Engine = struct { } attention_scale = ref.object.get("attention_scale").?.float; } else { - var cfg = try config.Config.read(arena, io, model_dir); ple_values = try config.pleRef(&cfg); attention_scale = 1.0 / std.math.sqrt(@as(f64, @floatFromInt(cfg.head_dim))); } @@ -271,8 +280,7 @@ pub const Engine = struct { L.pool = try r.loadf("L{d}.att.pool", .{i}); L.pooled = .{ .b = try r.buffer(CAP / 4 * 128 * 2) }; } - const projs = [_][]const u8{ "switch_mlp.gate_proj", "switch_mlp.up_proj", "shared_expert.gate_proj", "shared_expert.up_proj", "switch_mlp.down_proj", "shared_expert.down_proj" }; - for (projs, 0..) |proj, j| { + for (expert_projs, 0..) |proj, j| { for ([_][]const u8{ "weight", "scales", "biases" }, 0..) |suffix, k| { L.ex[j * 3 + k] = try r.loadf("language_model.model.layers.{d}.mlp.{s}.{s}", .{ i, proj, suffix }); } @@ -360,11 +368,11 @@ pub const Engine = struct { { const ids_n = try fz.draftCount((try r.entry("mtp.draft_ids")).len); const ids = try r.load("mtp.draft_ids"); + const mtp_root = @import("index.zig").mtpRoot(index.object.get("weight_map").?.object); var ex: [18]Buf = undefined; - const projs = [_][]const u8{ "switch_mlp.gate_proj", "switch_mlp.up_proj", "shared_expert.gate_proj", "shared_expert.up_proj", "switch_mlp.down_proj", "shared_expert.down_proj" }; - for (projs, 0..) |proj, j| { + for (expert_projs, 0..) |proj, j| { for ([_][]const u8{ "weight", "scales", "biases" }, 0..) |suffix, k| { - ex[j * 3 + k] = try r.loadf("language_model.mtp.layers.0.mlp.{s}.{s}", .{ proj, suffix }); + ex[j * 3 + k] = try r.loadf("{s}.layers.0.mlp.{s}.{s}", .{ mtp_root, proj, suffix }); } } if (r.xpack) try r.repack(&ex); diff --git a/zig/src/families/flashnext/host_test.zig b/zig/src/families/flashnext/host_test.zig index 22edefa19a..bc198ce4ad 100644 --- a/zig/src/families/flashnext/host_test.zig +++ b/zig/src/families/flashnext/host_test.zig @@ -119,8 +119,10 @@ test "module quantization preserves false true and dictionary-default override s try std.testing.expectEqual(@as(usize, 64), override.group); try std.testing.expectEqual(@as(u8, 6), (try c.quantization("layers.0.inherit")).?.bits); try std.testing.expectEqual(@as(usize, 32), (try c.quantization("unknown.projection")).?.group); - try std.testing.expectError(error.UnsupportedFlashAffineKernel, override.checkFlashKernel()); - try c.global_affine.checkFlashKernel(); + const widths = @import("widths.zig"); + try std.testing.expect(widths.launched(.{ .bits = override.bits, .group = @intCast(override.group) })); + try std.testing.expect(widths.launched(.{ .bits = c.global_affine.bits, .group = @intCast(c.global_affine.group) })); + try std.testing.expect(!widths.launched(.{ .bits = 7, .group = 64 }) and !widths.launched(.{ .bits = 4, .group = 16 })); } test "malformed FlashNext configuration fields are refused" { diff --git a/zig/src/families/flashnext/index.zig b/zig/src/families/flashnext/index.zig index b3210aed15..44c3c6d012 100644 --- a/zig/src/families/flashnext/index.zig +++ b/zig/src/families/flashnext/index.zig @@ -1,11 +1,30 @@ //! Public index admission proves required tensor names and safe shard basenames before payload headers are available. const std = @import("std"); const Config = @import("config.zig").Config; +const widths = @import("widths.zig"); + +/// The first required matrix the device kernels cannot read yet, by tensor stem, with its declared width. +pub const Refusal = struct { + buf: [256]u8 = undefined, + len: usize = 0, + width: ?widths.Width = null, + + pub fn name(r: *const Refusal) []const u8 { + return r.buf[0..r.len]; + } + + fn set(r: *Refusal, stem: []const u8, width: ?widths.Width) void { + r.len = @min(stem.len, r.buf.len); + @memcpy(r.buf[0..r.len], stem[0..r.len]); + r.width = width; + } +}; const Check = struct { map: std.json.ObjectMap, names: usize = 0, - runtime: ?*const Config = null, + runtime: ?*const widths.Widths = null, + refusal: ?*Refusal = null, fn take(c: *Check, name: []const u8) !void { if (c.map.get(name) == null) return error.MissingFlashTensor; @@ -24,9 +43,12 @@ const Check = struct { const biases = c.map.get(try std.fmt.bufPrint(&buf, "{s}.biases", .{stem})) != null; if (scales != biases) return error.IncompleteFlashAffine; if (scales) { - if (c.runtime) |cfg| { - const spec = (try cfg.quantization(stem)) orelse return error.UnsupportedFlashAffineKernel; - try spec.checkFlashKernel(); + if (c.runtime) |w| { + const width = try w.of(stem); + if (width == null or !widths.launched(width.?)) { + if (c.refusal) |r| r.set(stem, width); + return error.UnsupportedFlashAffineKernel; + } } try c.part(stem, "scales"); try c.part(stem, "biases"); @@ -57,15 +79,33 @@ const Check = struct { } } - fn ple(c: *Check, stem: []const u8, shards: usize, spelling: []const u8) !void { + /// An n-gram table shard: one lookup reads a row of `row` values a group at a time, so its group divides the row. + fn table(c: *Check, stem: []const u8, row: usize) !void { + try c.matrix(stem); + const w = c.runtime orelse return; + const width = (try w.of(stem)) orelse return; + if (row % width.group == 0) return; + if (c.refusal) |r| r.set(stem, width); + return error.UnsupportedFlashAffineKernel; + } + + fn ple(c: *Check, stem: []const u8, shards: usize, spelling: []const u8, row: usize) !void { for ([_][]const u8{ "conv1d.weight", "norm_conv.weight", "norm_key.weight", "norm_query.weight", "ple_embedding.layer_multipliers", "ple_embedding.ngram_heads_offsets", "ple_embedding.ngram_heads_vocab_sizes" }) |suffix| try c.part(stem, suffix); try c.proj(stem, "key_proj"); try c.proj(stem, "value_proj"); var buf: [256]u8 = undefined; - for (0..shards) |i| try c.proj(stem, try std.fmt.bufPrint(&buf, "ple_embedding.ngram_embedding.{s}{d}", .{ spelling, i })); + var name: [256]u8 = undefined; + for (0..shards) |i| try c.table(try std.fmt.bufPrint(&name, "{s}.{s}", .{ stem, try std.fmt.bufPrint(&buf, "ple_embedding.ngram_embedding.{s}{d}", .{ spelling, i }) }), row); } }; +/// Where the index keeps the MTP head: `language_model.mtp` (mlx-lm), else oMLX's top-level `mtp`. +pub fn mtpRoot(map: std.json.ObjectMap) []const u8 { + for (map.keys()) |name| if (std.mem.startsWith(u8, name, "language_model.mtp.")) return "language_model.mtp"; + for (map.keys()) |name| if (std.mem.startsWith(u8, name, "mtp.")) return "mtp"; + return "language_model.mtp"; +} + /// The index's spelling of its n-gram table shards, `shard_N` when a `shard_0` is listed, else oMLX's `shards.N`. pub fn ngramSpelling(map: std.json.ObjectMap) []const u8 { for (map.keys()) |name| if (std.mem.endsWith(u8, name, ".ngram_embedding.shard_0.weight")) return "shard_"; @@ -75,16 +115,15 @@ pub fn ngramSpelling(map: std.json.ObjectMap) []const u8 { pub const Inventory = struct { tensor_names: usize, required_names: usize, shards: usize, headers_verified: bool = false }; pub fn admit(gpa: std.mem.Allocator, bytes: []const u8, config: *const Config) !Inventory { - return admitWith(gpa, bytes, config, false); + return admitWith(gpa, bytes, config, null, null); } -/// Runtime kernels consume only six-bit/group-32 text matrices, unlike the generic pack metadata reader. -pub fn admitRuntime(gpa: std.mem.Allocator, bytes: []const u8, config: *const Config) !Inventory { - try config.global_affine.checkFlashKernel(); - return admitWith(gpa, bytes, config, true); +/// Each required matrix at its own module width, refused by name when no device kernel reads that width yet. +pub fn admitRuntime(gpa: std.mem.Allocator, bytes: []const u8, config: *const Config, w: *const widths.Widths, refusal: *Refusal) !Inventory { + return admitWith(gpa, bytes, config, w, refusal); } -fn admitWith(gpa: std.mem.Allocator, bytes: []const u8, config: *const Config, runtime: bool) !Inventory { +fn admitWith(gpa: std.mem.Allocator, bytes: []const u8, config: *const Config, runtime: ?*const widths.Widths, refusal: ?*Refusal) !Inventory { const p = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); defer p.deinit(); if (p.value != .object) return error.InvalidFlashIndex; @@ -98,7 +137,7 @@ fn admitWith(gpa: std.mem.Allocator, bytes: []const u8, config: *const Config, r if (file != .string or file.string.len == 0 or file.string[0] == '.' or std.mem.indexOfAny(u8, file.string, "/\\") != null or !std.mem.endsWith(u8, file.string, ".safetensors")) return error.UnsafeFlashShard; try files.put(gpa, file.string, {}); } - var c = Check{ .map = map.object, .runtime = if (runtime) config else null }; + var c = Check{ .map = map.object, .runtime = runtime, .refusal = refusal }; const spelling = ngramSpelling(map.object); // every shard is checked under this one, so an index mixing the two is refused try c.matrix("language_model.model.embed_tokens"); try c.matrix("language_model.lm_head"); @@ -109,14 +148,16 @@ fn admitWith(gpa: std.mem.Allocator, bytes: []const u8, config: *const Config, r try c.layer(layer, try config.kind(i)); if (config.ple[i]) { var buf: [160]u8 = undefined; - try c.ple(try std.fmt.bufPrint(&buf, "{s}.ple", .{layer}), config.ngram_shards, spelling); + try c.ple(try std.fmt.bufPrint(&buf, "{s}.ple", .{layer}), config.ngram_shards, spelling, config.hidden / try config.ngramHeadCount()); } } if (config.mtp.layers > 0) { - for ([_][]const u8{ "fc_embedding", "fc_hidden" }) |suffix| try c.proj("language_model.mtp", suffix); - for ([_][]const u8{ "pre_fc_norm_embedding.weight", "pre_fc_norm_hidden.weight" }) |suffix| try c.part("language_model.mtp", suffix); - try c.hc("language_model.mtp.hyper_connection_mixer", false); - for (0..config.mtp.layers) |i| try c.layer(try std.fmt.bufPrint(&stem, "language_model.mtp.layers.{d}", .{i}), config.mtp.kinds[i]); + const root = mtpRoot(map.object); + for ([_][]const u8{ "fc_embedding", "fc_hidden" }) |suffix| try c.proj(root, suffix); + for ([_][]const u8{ "pre_fc_norm_embedding.weight", "pre_fc_norm_hidden.weight" }) |suffix| try c.part(root, suffix); + var mixer: [64]u8 = undefined; + try c.hc(try std.fmt.bufPrint(&mixer, "{s}.hyper_connection_mixer", .{root}), false); + for (0..config.mtp.layers) |i| try c.layer(try std.fmt.bufPrint(&stem, "{s}.layers.{d}", .{ root, i }), config.mtp.kinds[i]); } return .{ .tensor_names = map.object.count(), .required_names = c.names, .shards = files.count() }; } diff --git a/zig/src/families/flashnext/load_admission.zig b/zig/src/families/flashnext/load_admission.zig index 6970de32c5..ca3cda481e 100644 --- a/zig/src/families/flashnext/load_admission.zig +++ b/zig/src/families/flashnext/load_admission.zig @@ -2,27 +2,49 @@ const std = @import("std"); const config = @import("config.zig"); const index = @import("index.zig"); +const widths = @import("widths.zig"); -pub const problem = "Flash Next runtime kernels require the supported 6-bit affine, group_size 32 format; use a matching checkpoint"; +pub const problem = "Flash Next runtime kernels read MLX affine 2- to 8-bit weights in groups of 32, 64 or 128 (n-gram tables in groups that divide a row); the log names the first other one"; + +pub const stored_problem = "the checkpoint's stored weights contradict the widths its config declares, or a layout one kernel reads at one width differs; the log names the first one"; +pub const speed_up_problem = "speed-up mode (two Macs) splits the 6-bit group-32 kernels only; serve this checkpoint on one Mac"; +pub const scale_problem = "an n-gram table scale must be one finite, nonzero value; the log names it"; pub fn problemFor(err: anyerror) ?[]const u8 { return switch (err) { - error.UnsupportedFlashAffineKernel, error.UnsupportedAffineBits, error.UnsupportedAffineGroup, error.UnsupportedAffineMode => problem, + error.UnsupportedFlashAffineKernel, error.UnsupportedAffineBits, error.UnsupportedAffineGroup, error.UnsupportedAffineMode, error.UnsupportedQuantization => problem, + error.InvalidAffineShape, error.InvalidAffineRank, error.InvalidAffineStorage, error.InvalidAffineWidth, error.InvalidAffineByteCount, error.UnquantizedAffineTensor, error.ExpertGateUpWidths, error.NgramTableWidths, error.LaneRowsNotFours, error.TooManyLaneMembers => stored_problem, + error.SpeedUpNeeds6Bit32 => speed_up_problem, + error.NgramTableScale => scale_problem, else => null, }; } pub fn check(gpa: std.mem.Allocator, io: std.Io, directory: []const u8) !void { + var refusal: index.Refusal = .{}; + checkWith(gpa, io, directory, &refusal) catch |err| { + if (refusal.width) |w| std.log.warn("{s}: {d}-bit group {d} has no Flash Next device kernel", .{ refusal.name(), w.bits, w.group }); + if (refusal.len > 0 and refusal.width == null) std.log.warn("{s}: stored affine but declared unquantized", .{refusal.name()}); + return err; + }; +} + +/// Every required matrix resolved at its own module width; the first one no device kernel reads lands in `refusal`. +pub fn checkWith(gpa: std.mem.Allocator, io: std.Io, directory: []const u8, refusal: *index.Refusal) !void { var c = try config.Config.read(gpa, io, directory); defer c.deinit(); - try c.global_affine.checkFlashKernel(); + var w = try widths.Widths.init(gpa, &c); + defer w.deinit(); + // refused before the index + if (w.uniform() and !widths.launched(w.q.global)) return error.UnsupportedFlashAffineKernel; const path = try std.fs.path.join(gpa, &.{ directory, "model.safetensors.index.json" }); defer gpa.free(path); const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 26)); defer gpa.free(bytes); - _ = try index.admitRuntime(gpa, bytes, &c); + _ = try index.admitRuntime(gpa, bytes, &c, &w, refusal); } test { _ = @import("load_admission_test.zig"); + _ = @import("widths.zig"); } diff --git a/zig/src/families/flashnext/load_admission_test.zig b/zig/src/families/flashnext/load_admission_test.zig index 4a51579762..cde234ec9c 100644 --- a/zig/src/families/flashnext/load_admission_test.zig +++ b/zig/src/families/flashnext/load_admission_test.zig @@ -19,7 +19,7 @@ fn overridden(path: []const u8, value: []const u8) ![]u8 { return std.mem.replaceOwned(u8, a, checked_config, "\"mode\": \"affine\"", declaration); } -test "Flash Next load refuses a four-bit group-32 checkpoint before touching any weights or cache" { +test "Flash Next load admits a four-bit group-32 checkpoint without touching any weights or cache" { var tmp = std.testing.tmpDir(.{}); defer tmp.cleanup(); const bytes = try std.mem.replaceOwned(u8, a, checked_config, "\"bits\": 6", "\"bits\": 4"); @@ -28,11 +28,10 @@ test "Flash Next load refuses a four-bit group-32 checkpoint before touching any defer metadata.deinit(); try std.testing.expectEqual(@as(u8, 4), metadata.global_affine.bits); try std.testing.expectEqual(@as(usize, 32), metadata.global_affine.group); - try tmp.dir.writeFile(io, .{ .sub_path = "config.json", .data = bytes }); - const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + const dir = try fixture(tmp, bytes, checked_index); defer a.free(dir); - try std.testing.expectError(error.UnsupportedFlashAffineKernel, admission.check(a, io, dir)); - try std.testing.expectError(error.FileNotFound, tmp.dir.openFile(io, "model.safetensors.index.json", .{})); + try admission.check(a, io, dir); + try std.testing.expectError(error.FileNotFound, tmp.dir.openFile(io, "model-00001-of-00030.safetensors", .{})); try std.testing.expectError(error.FileNotFound, tmp.dir.openDir(io, "zig-pack", .{})); } @@ -46,7 +45,7 @@ test "load admission accepts the supported text and MTP matrices without reading try std.testing.expectError(error.FileNotFound, tmp.dir.openDir(io, "zig-pack", .{})); } -test "global unsupported affine groups refuse before the index is opened" { +test "a global group of 64 or 128 is admitted for every matrix, and the n-gram tables' 160-wide rows refuse it" { for ([_][]const u8{ "64", "128" }) |group| { var tmp = std.testing.tmpDir(.{}); defer tmp.cleanup(); @@ -54,18 +53,29 @@ test "global unsupported affine groups refuse before the index is opened" { defer a.free(declaration); const bytes = try std.mem.replaceOwned(u8, a, checked_config, "\"group_size\": 32", declaration); defer a.free(bytes); - try tmp.dir.writeFile(io, .{ .sub_path = "config.json", .data = bytes }); - const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + const dir = try fixture(tmp, bytes, checked_index); defer a.free(dir); - try std.testing.expectError(error.UnsupportedFlashAffineKernel, admission.check(a, io, dir)); + var refusal: @import("index.zig").Refusal = .{}; + try std.testing.expectError(error.UnsupportedFlashAffineKernel, admission.checkWith(a, io, dir, &refusal)); + try std.testing.expect(std.mem.indexOf(u8, refusal.name(), ".ple.ple_embedding.ngram_embedding.shard_0") != null); } } -test "required text and MTP affine overrides cannot bypass the runtime format" { +test "text and MTP overrides at other widths are admitted, and a stored matrix declared plain is refused" { for ([_]struct { path: []const u8, value: []const u8 }{ .{ .path = "language_model.model.layers.0.linear_attn.in_proj_qkv", .value = "{\"bits\":4,\"group_size\":32}" }, .{ .path = "model.layers.0.linear_attn.in_proj_qkv.weight", .value = "{\"bits\":6,\"group_size\":64}" }, - .{ .path = "language_model.mtp.layers.0.self_attn.q_proj", .value = "{\"bits\":4,\"group_size\":32}" }, + .{ .path = "language_model.mtp.layers.0.self_attn.q_proj", .value = "{\"bits\":3,\"group_size\":128}" }, + }) |override| { + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const bytes = try overridden(override.path, override.value); + defer a.free(bytes); + const dir = try fixture(tmp, bytes, checked_index); + defer a.free(dir); + try admission.check(a, io, dir); + } + for ([_]struct { path: []const u8, value: []const u8 }{ .{ .path = "language_model.lm_head", .value = "{}" }, .{ .path = "language_model.model.embed_tokens", .value = "false" }, }) |override| { @@ -121,6 +131,45 @@ test "normal and dump load paths admit before any engine device or pack work and const mapping = std.mem.indexOfPos(u8, native, open, "tf.flashnext_engine.quantizationProblem(e)"); try std.testing.expect(mapping != null); const problem = admission.problemFor(error.UnsupportedFlashAffineKernel).?; - try std.testing.expect(std.mem.indexOf(u8, problem, "supported 6-bit affine, group_size 32") != null); + try std.testing.expect(std.mem.indexOf(u8, problem, "2- to 8-bit weights in groups of 32, 64 or 128") != null); try std.testing.expect(admission.problemFor(error.OutOfMemory) == null); } + +test "a module the device kernels cannot read is refused under its checkpoint tensor name, whatever its key's prefix" { + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const bytes = try overridden("model.layers.1.ple.ple_embedding.ngram_embedding.shard_3.weight", "{\"bits\":8,\"group_size\":64}"); + defer a.free(bytes); + const dir = try fixture(tmp, bytes, checked_index); + defer a.free(dir); + var refusal: @import("index.zig").Refusal = .{}; + try std.testing.expectError(error.UnsupportedFlashAffineKernel, admission.checkWith(a, io, dir, &refusal)); + try std.testing.expectEqualStrings("language_model.model.layers.1.ple.ple_embedding.ngram_embedding.shard_3", refusal.name()); + try std.testing.expectEqual(@as(u8, 8), refusal.width.?.bits); + try std.testing.expectEqual(@as(u16, 64), refusal.width.?.group); +} + +test "an index keeping the MTP head at top-level mtp (oMLX) is admitted as one keeping it under language_model" { + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const index = try std.mem.replaceOwned(u8, a, checked_index, "\"language_model.mtp.", "\"mtp."); + defer a.free(index); + try std.testing.expect(std.mem.indexOf(u8, index, "\"mtp.fc_embedding.weight\"") != null); + const dir = try fixture(tmp, checked_config, index); + defer a.free(dir); + try admission.check(a, io, dir); + const missing = try std.mem.replaceOwned(u8, a, index, "\"mtp.fc_hidden.weight\"", "\"mtp.fc_hidden.other\""); + defer a.free(missing); + var tmp2 = std.testing.tmpDir(.{}); + defer tmp2.cleanup(); + const dir2 = try fixture(tmp2, checked_config, missing); + defer a.free(dir2); + try std.testing.expectError(error.MissingFlashTensor, admission.check(a, io, dir2)); +} + +test "the width and layout refusals past admission map to text a user can act on" { + for ([_]anyerror{ error.InvalidAffineShape, error.ExpertGateUpWidths, error.NgramTableWidths, error.LaneRowsNotFours, error.TooManyLaneMembers }) |e| + try std.testing.expectEqualStrings(admission.stored_problem, admission.problemFor(e).?); + try std.testing.expectEqualStrings(admission.speed_up_problem, admission.problemFor(error.SpeedUpNeeds6Bit32).?); + try std.testing.expectEqualStrings(admission.scale_problem, admission.problemFor(error.NgramTableScale).?); +} diff --git a/zig/src/families/flashnext/pack.zig b/zig/src/families/flashnext/pack.zig index 7b1105e9d2..2834475232 100644 --- a/zig/src/families/flashnext/pack.zig +++ b/zig/src/families/flashnext/pack.zig @@ -7,6 +7,9 @@ const affine = @import("affine.zig"); const config_mod = @import("config.zig"); const index_mod = @import("index.zig"); const pio = @import("pack_io.zig"); +const quant_mlx = @import("../../core/quant/mlx.zig"); +const quant_widen = @import("../../core/quant/widen.zig"); +const widen = @import("pack_widen.zig"); pub const Report = struct { decode_tensors: usize = 0, @@ -16,31 +19,24 @@ pub const Report = struct { draft_ids: usize = 0, }; -const Out = pio.Out; +pub const Out = pio.Out; -/// Which of a quantized linear's three tensors. -const Part = enum { weight, scales, biases }; - -/// A checkpoint quantized linear's three tensors with its verified affine spec. -pub const Q = struct { - w: st.Tensor, - s: st.Tensor, - b: st.Tensor, - rows: usize, - k: usize, - kw: usize, - spec: affine.Spec, - - fn part(r: Q, comptime which: Part) []const u8 { - return switch (which) { - .weight => r.w.bytes, - .scales => r.s.bytes, - .biases => r.b.bytes, - }; - } -}; +pub const lane = @import("pack_lane.zig"); +pub const Q = lane.Q; +pub const tileWeight = lane.tileWeight; +pub const packScales = lane.packScales; +const concatPart = lane.concatPart; +const putLane = lane.putLane; +/// A quantized linear read only at its own declared width, refused under its name otherwise. fn qlinear(ck: *ckpt.Checkpoint, cfg: *const config_mod.Config, path: []const u8) !Q { + return qlinearAt(ck, cfg, path) catch |err| { + std.log.warn("{s}: not read at its declared width: {s}", .{ path, @errorName(err) }); + return err; + }; +} + +fn qlinearAt(ck: *ckpt.Checkpoint, cfg: *const config_mod.Config, path: []const u8) !Q { const spec = (try cfg.quantization(path)) orelse return error.NotQuantized; var buf: [320]u8 = undefined; const w = try ck.get(try std.fmt.bufPrint(&buf, "{s}.weight", .{path})); @@ -48,7 +44,6 @@ fn qlinear(ck: *ckpt.Checkpoint, cfg: *const config_mod.Config, path: []const u8 const b = try ck.get(try std.fmt.bufPrint(&buf, "{s}.biases", .{path})); if (w.dtype != .u32 or w.rank != 2) return error.UnexpectedTensor; if (s.dtype != .bf16 or s.rank != 2 or b.dtype != .bf16 or b.rank != 2) return error.UnexpectedTensor; - if (spec.group == 128) return error.Group128NotPacked; // the pack never carries group 128; the halve-to-64 widen is unimplemented const rows = w.dim(0); const kw = w.dim(1); if (rows != s.dim(0) or rows != b.dim(0)) return error.UnexpectedTensor; @@ -58,79 +53,6 @@ fn qlinear(ck: *ckpt.Checkpoint, cfg: *const config_mod.Config, path: []const u8 return .{ .w = w, .s = s, .b = b, .rows = rows, .k = k, .kw = kw, .spec = spec }; } -/// The lane pipeline's word permutation: rows tiled nt at a time, a group's words innermost (lane_qmm.tile_weight). -pub fn tileWeight(gpa: std.mem.Allocator, flat: []const u8, rows: usize, kw: usize, nt: usize, w: usize) ![]u8 { - const words = std.mem.bytesAsSlice(u32, flat); - const out = try gpa.alloc(u8, flat.len); - errdefer gpa.free(out); - const out_words = std.mem.bytesAsSlice(u32, out); - const groups = kw / w; - for (0..rows / nt) |t| { - for (0..groups) |g| { - for (0..nt) |j| { - for (0..w) |c| { - out_words[((t * groups + g) * nt + j) * w + c] = words[(t * nt + j) * kw + g * w + c]; - } - } - } - } - return out; -} - -/// (K/GS, N, 2) bf16 (scale, bias) pairs, group-major, over a row-concatenated stack (lane_qmm.pack_scales). -pub fn packScales(gpa: std.mem.Allocator, members: []const Q) ![]u8 { - var rows: usize = 0; - for (members) |m| rows += m.rows; - const kg = members[0].s.dim(1); - const out = try gpa.alloc(u8, kg * rows * 4); - errdefer gpa.free(out); - const out16 = std.mem.bytesAsSlice(u16, out); - var at: usize = 0; - for (members) |m| { - const s16 = std.mem.bytesAsSlice(u16, m.part(.scales)); - const b16 = std.mem.bytesAsSlice(u16, m.part(.biases)); - for (0..kg) |g| { - for (0..m.rows) |j| { - out16[(g * rows + at + j) * 2] = s16[j * kg + g]; - out16[(g * rows + at + j) * 2 + 1] = b16[j * kg + g]; - } - } - at += m.rows; - } - return out; -} - -/// Row-concatenated members in the checkpoint's own MLX layout: a pure byte move per member. -fn concatPart(gpa: std.mem.Allocator, members: []const Q, comptime part: Part) ![]u8 { - var n: usize = 0; - for (members) |m| n += m.part(part).len; - const out = try gpa.alloc(u8, n); - errdefer gpa.free(out); - var at: usize = 0; - for (members) |m| { - @memcpy(out[at..][0..m.part(part).len], m.part(part)); - at += m.part(part).len; - } - return out; -} - -/// One lane projection's two pack tensors: `.wq` tiled, `.sbt` group-major pairs (decode.py's lane pipeline). -fn putLane(out: *Out, gpa: std.mem.Allocator, name: []const u8, members: []const Q) !void { - var rows: usize = 0; - for (members) |m| rows += m.rows; - const spec = members[0].spec; - for (members) |m| if (m.spec.bits != spec.bits or m.spec.group != spec.group) return error.MixedAffine; - const nt: usize = if (spec.bits == 4 and rows % 64 == 0) 64 else if (rows % 32 == 0) 32 else return error.UnsupportedLaneWidth; - const w = spec.group * spec.bits / 32; - const flat = try concatPart(gpa, members, .weight); - defer gpa.free(flat); - const wq = try tileWeight(gpa, flat, rows, members[0].kw, nt, w); - const sbt = try packScales(gpa, members); - const a = out.arena.allocator(); - try out.put(try std.fmt.allocPrint(a, "{s}.wq", .{name}), "U32", &.{ rows, members[0].kw }, wq); - try out.put(try std.fmt.allocPrint(a, "{s}.sbt", .{name}), "BF16", &.{ members[0].s.dim(1), rows, 2 }, sbt); -} - fn bf16ToF32(bits: u16) f32 { return @bitCast(@as(u32, bits) << 16); } @@ -195,22 +117,29 @@ fn putHc(out: *Out, gpa: std.mem.Allocator, a: std.mem.Allocator, ck: *ckpt.Chec } var storage: [2]Q = undefined; for (down_paths[0..down_n], 0..) |p, i| storage[i] = try qlinear(ck, cfg, p); - const down = storage[0..down_n]; - for (down) |m| if (m.spec.bits != down[0].spec.bits or m.spec.group != down[0].spec.group) return error.MixedAffine; - var rows: usize = 0; - for (down) |m| rows += m.rows; - const scale = try centeredScale(gpa, ck, try std.fmt.allocPrint(a, "{s}.hc_norm.weight", .{stem}), around_one); - try out.put(try std.fmt.allocPrint(a, "{s}.scale", .{name}), "F32", &.{wide}, scale); - const down_w = try concatPart(gpa, down, .weight); - const down_s = try concatPart(gpa, down, .scales); - const down_b = try concatPart(gpa, down, .biases); - try out.put(try std.fmt.allocPrint(a, "{s}.down.w", .{name}), "U32", &.{ rows, down[0].kw }, down_w); - try out.put(try std.fmt.allocPrint(a, "{s}.down.s", .{name}), "BF16", &.{ rows, down[0].s.dim(1) }, down_s); - try out.put(try std.fmt.allocPrint(a, "{s}.down.b", .{name}), "BF16", &.{ rows, down[0].b.dim(1) }, down_b); + const down = try hcDown(a, storage[0..down_n]); // a pair of two widths widens to the wider (QWeights.of) const up = try qlinear(ck, cfg, try std.fmt.allocPrint(a, "{s}.input_mix_weight_up", .{stem})); + const scale_name = try std.fmt.allocPrint(a, "{s}.scale", .{name}); + try out.put(scale_name, "F32", &.{wide}, try centeredScale(gpa, ck, try std.fmt.allocPrint(a, "{s}.hc_norm.weight", .{stem}), around_one)); + try out.put(try std.fmt.allocPrint(a, "{s}.down.w", .{name}), "U32", &.{ down.rows, down.kw }, try gpa.dupe(u8, down.part(.weight))); + try out.put(try std.fmt.allocPrint(a, "{s}.down.s", .{name}), "BF16", &.{ down.rows, down.s.dim(1) }, try gpa.dupe(u8, down.part(.scales))); + try out.put(try std.fmt.allocPrint(a, "{s}.down.b", .{name}), "BF16", &.{ down.rows, down.b.dim(1) }, try gpa.dupe(u8, down.part(.biases))); try out.put(try std.fmt.allocPrint(a, "{s}.up.w", .{name}), "U32", &.{ up.rows, up.kw }, try gpa.dupe(u8, up.part(.weight))); try out.put(try std.fmt.allocPrint(a, "{s}.up.s", .{name}), "BF16", &.{ up.rows, up.s.dim(1) }, try gpa.dupe(u8, up.part(.scales))); try out.put(try std.fmt.allocPrint(a, "{s}.up.b", .{name}), "BF16", &.{ up.rows, up.b.dim(1) }, try gpa.dupe(u8, up.part(.biases))); + // the uniform packs' pinned bytes + if (down.spec.bits == 6 and down.spec.group == 32 and up.spec.bits == 6 and up.spec.group == 32) return; + const widths_name = try std.fmt.allocPrint(a, "{s}.widths", .{name}); + const w = try gpa.alloc(u32, 4); + w[0..4].* = .{ down.spec.bits, @intCast(down.spec.group), up.spec.bits, @intCast(up.spec.group) }; + try out.put(widths_name, "U32", &.{4}, w); +} + +/// The down projection and the inject rows as one matrix: one width as stored, two widened exactly to the wider. +pub fn hcDown(a: std.mem.Allocator, members: []const Q) !Q { + var hosts: [2]quant_mlx.Host = undefined; + for (members, 0..) |m, i| hosts[i] = widen.host(m); + return widen.fromHost(try quant_widen.fuse(a, hosts[0..members.len])); } /// A gate's rows are stored bf16, or an affine linear dequantized on the host as code * scale + bias per group. @@ -332,6 +261,20 @@ fn pleStarts(gpa: std.mem.Allocator, ck: *ckpt.Checkpoint, cfg: *const config_mo return out; } +/// An n-gram table's weight_scale must be one finite, nonzero bf16 or f32 value: the lookups multiply rows by it. +fn tableScale(ck: *ckpt.Checkpoint, map: std.json.ObjectMap, name: []const u8) !void { + if (!map.contains(name)) return; + const t = try ck.get(name); + const v: ?f32 = switch (t.dtype) { + .bf16 => if (t.bytes.len == 2) @bitCast(@as(u32, std.mem.readInt(u16, t.bytes[0..2], .little)) << 16) else null, + .f32 => if (t.bytes.len == 4) @bitCast(std.mem.readInt(u32, t.bytes[0..4], .little)) else null, + else => null, + }; + if (v) |x| if (std.math.isFinite(x) and x != 0) return; + std.log.warn("{s}: an n-gram table scale must be one finite, nonzero value", .{name}); + return error.NgramTableScale; +} + /// Draft vocabulary ids are sorted and padded to a multiple of 64, and an id past the vocabulary is refused. pub fn draftIdList(gpa: std.mem.Allocator, io: Io, path: []const u8, vocab: usize) ![]u32 { const text = try Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 26)); @@ -371,9 +314,12 @@ fn putDraftHead(out: *Out, gpa: std.mem.Allocator, a: std.mem.Allocator, name: [ const rows = ids.len; const kw = lm_head.kw; const kg = lm_head.s.dim(1); - const wq = try gpa.alloc(u8, rows * kw * 4); + const wq = try gpa.alloc(u8, rows * kw * 4); // putLane copies the gathered rows, so they free here + defer gpa.free(wq); const s = try gpa.alloc(u8, rows * kg * 2); + defer gpa.free(s); const b = try gpa.alloc(u8, rows * kg * 2); + defer gpa.free(b); const w32 = std.mem.bytesAsSlice(u32, wq); const src32 = std.mem.bytesAsSlice(u32, lm_head.part(.weight)); const s16 = std.mem.bytesAsSlice(u16, s); @@ -387,10 +333,7 @@ fn putDraftHead(out: *Out, gpa: std.mem.Allocator, a: std.mem.Allocator, name: [ @memcpy(b16[r * kg ..][0..kg], src_b16[id * kg ..][0..kg]); } const gathered = [_]Q{.{ .w = .{ .dtype = .u32, .rank = 2, .shape = .{ rows, kw, 0, 0, 0 }, .bytes = wq }, .s = .{ .dtype = .bf16, .rank = 2, .shape = .{ rows, kg, 0, 0, 0 }, .bytes = s }, .b = .{ .dtype = .bf16, .rank = 2, .shape = .{ rows, kg, 0, 0, 0 }, .bytes = b }, .rows = rows, .k = lm_head.k, .kw = kw, .spec = lm_head.spec }}; - try putLane(out, gpa, name, gathered[0..]); // putLane copies the bytes, so the gathered buffers free here - gpa.free(wq); - gpa.free(s); - gpa.free(b); + try putLane(out, gpa, name, gathered[0..]); _ = a; } @@ -401,6 +344,9 @@ pub fn build(gpa: std.mem.Allocator, io: Io, model_dir: []const u8, out_dir: []c const a = arena_state.allocator(); var cfg = try config_mod.Config.read(gpa, io, model_dir); defer cfg.deinit(); + var strict = try @import("widths.zig").Widths.init(gpa, &cfg); + defer strict.deinit(); + cfg.widths = &strict; // the vision tower's tensors stay unread: a rank-five entry must not stop the no-dump build var ck = try ckpt.Checkpoint.openModelPrefix(gpa, io, model_dir, "language_model."); defer ck.close(); @@ -416,6 +362,11 @@ pub fn build(gpa: std.mem.Allocator, io: Io, model_dir: []const u8, out_dir: []c defer parsed_index.deinit(); const weight_map = parsed_index.value.object.get("weight_map").?.object; const spelling = index_mod.ngramSpelling(weight_map); + const mtp_root = index_mod.mtpRoot(weight_map); // oMLX keeps the head outside language_model. + var ck_mtp = if (std.mem.eql(u8, mtp_root, "mtp")) try ckpt.Checkpoint.openModelPrefix(gpa, io, model_dir, "mtp.") else null; + defer if (ck_mtp) |*m| m.close(); + const ckm = if (ck_mtp) |*m| m else &ck; + try @import("pack_raw.zig").checkAll(&ck, ckm, &strict, &cfg, mtp_root, spelling); // the directly read tensors const identity = try pio.sourceIdentity(gpa, io, model_dir); defer gpa.free(identity); @@ -491,13 +442,14 @@ pub fn build(gpa: std.mem.Allocator, io: Io, model_dir: []const u8, out_dir: []c try decode.put(try std.fmt.allocPrint(a, "ple.{s}", .{pack_name}), "F32", &.{wide}, scale); } try putConvSlice(&decode, gpa, &ck, "ple.conv", try std.fmt.allocPrint(a, "{s}.ple.conv1d.weight", .{stem}), true); + try tableScale(&ck, weight_map, try std.fmt.allocPrint(a, "{s}.ple.ple_embedding.ngram_embedding.weight_scale", .{stem})); const starts = try pleStarts(gpa, &ck, &cfg, pl, spelling); try decode.put("ple.starts", "U32", &.{8}, starts); } if (draft_vocab) |vocab_path| { const ids = try draftIdList(gpa, io, vocab_path, cfg.vocab); report.draft_ids = ids.len; - try putMtp(&decode, gpa, a, &ck, &cfg, ids, &report, wide); + try putMtp(&decode, gpa, a, ckm, &ck, mtp_root, &cfg, ids, &report, wide); try decode.put("mtp.draft_ids", "U32", &.{ids.len}, ids); } report.decode_tensors = decode.names.items.len; @@ -527,34 +479,52 @@ pub fn build(gpa: std.mem.Allocator, io: Io, model_dir: []const u8, out_dir: []c if (cfg.mtp.layers > 0) { var mtp_mlx = Out.init(gpa); defer mtp_mlx.deinit(); - const mtp_stem = "language_model.mtp.layers.0"; - try putMlxProjection(&mtp_mlx, gpa, a, &ck, &cfg, "mtp.att.proj", &.{ "self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.indexer.index_qk_proj" }, mtp_stem); - try putMlxProjection(&mtp_mlx, gpa, a, &ck, &cfg, "mtp.fce", &.{"fc_embedding"}, "language_model.mtp"); - try putMlxProjection(&mtp_mlx, gpa, a, &ck, &cfg, "mtp.fch", &.{"fc_hidden"}, "language_model.mtp"); + const mtp_stem = try std.fmt.allocPrint(a, "{s}.layers.0", .{mtp_root}); + try putMlxProjection(&mtp_mlx, gpa, a, ckm, &cfg, "mtp.att.proj", &.{ "self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.indexer.index_qk_proj" }, mtp_stem); + try putMlxProjection(&mtp_mlx, gpa, a, ckm, &cfg, "mtp.fce", &.{"fc_embedding"}, mtp_root); + try putMlxProjection(&mtp_mlx, gpa, a, ckm, &cfg, "mtp.fch", &.{"fc_hidden"}, mtp_root); report.mtp_mlx_tensors = mtp_mlx.names.items.len; try mtp_mlx.write(gpa, io, try std.fs.path.join(a, &.{ out_dir, "pack_mtp_mlx.safetensors" }), identity); } return report; } -/// One export_mlx_proj-style projection: `.mw` the checkpoint's packed words, `.ms`/`.mb` its scales and biases. +/// One export_mlx_proj-style projection (`.mw` words, `.ms`/`.mb` tables); `.parts` gives other widths per member. fn putMlxProjection(out: *Out, gpa: std.mem.Allocator, a: std.mem.Allocator, ck: *ckpt.Checkpoint, cfg: *const config_mod.Config, name: []const u8, members: []const []const u8, stem: []const u8) !void { var storage: [4]Q = undefined; for (members, 0..) |n, i| storage[i] = try qlinear(ck, cfg, try std.fmt.allocPrint(a, "{s}.{s}", .{ stem, n })); - const list = storage[0..members.len]; - const mw = try concatPart(gpa, list, .weight); - const ms = try concatPart(gpa, list, .scales); - const mb = try concatPart(gpa, list, .biases); + return mlxProjection(out, gpa, a, name, storage[0..members.len]); +} + +/// The members' MLX bytes in member order; 6/32 alone keeps its 2-D shapes, other widths go flat with `.parts`. +pub fn mlxProjection(out: *Out, gpa: std.mem.Allocator, a: std.mem.Allocator, name: []const u8, list: []const Q) !void { var rows: usize = 0; for (list) |m| rows += m.rows; - try out.put(try std.fmt.allocPrint(a, "{s}.mw", .{name}), "U32", &.{ rows, list[0].kw }, mw); - try out.put(try std.fmt.allocPrint(a, "{s}.ms", .{name}), "BF16", &.{ rows, list[0].s.dim(1) }, ms); - try out.put(try std.fmt.allocPrint(a, "{s}.mb", .{name}), "BF16", &.{ rows, list[0].b.dim(1) }, mb); + const six = for (list) |m| { + if (m.spec.bits != 6 or m.spec.group != 32) break false; + } else true; + // the uniform packs keep their bytes and shapes; members of other widths go flat, in member order + try mlxPart(out, gpa, a, name, "mw", "U32", list, .weight, if (six) .{ rows, list[0].kw } else null); + try mlxPart(out, gpa, a, name, "ms", "BF16", list, .scales, if (six) .{ rows, list[0].s.dim(1) } else null); + try mlxPart(out, gpa, a, name, "mb", "BF16", list, .biases, if (six) .{ rows, list[0].b.dim(1) } else null); + if (six) return; + const parts_name = try std.fmt.allocPrint(a, "{s}.parts", .{name}); // each member's rows, bits and group + const parts = try gpa.alloc(u32, list.len * 3); + for (list, 0..) |m, i| parts[i * 3 ..][0..3].* = .{ @intCast(m.rows), m.spec.bits, @intCast(m.spec.group) }; + try out.put(parts_name, "U32", &.{ list.len, 3 }, parts); +} + +/// One of a projection's three tensors, members concatenated: at `shape` when given, else flat. +fn mlxPart(out: *Out, gpa: std.mem.Allocator, a: std.mem.Allocator, name: []const u8, suffix: []const u8, dtype: []const u8, list: []const Q, comptime part: lane.Part, shape: ?[2]usize) !void { + const full = try std.fmt.allocPrint(a, "{s}.{s}", .{ name, suffix }); + const bytes = try concatPart(gpa, list, part); + if (shape) |sh| return out.put(full, dtype, &sh, bytes); + return out.put(full, dtype, &.{bytes.len / @as(usize, if (part == .weight) 4 else 2)}, bytes); } /// The mtp.* decode tensors: the MTP layer's hc/att/moe classes, its norms, and the draft head over the id list. -fn putMtp(out: *Out, gpa: std.mem.Allocator, a: std.mem.Allocator, ck: *ckpt.Checkpoint, cfg: *const config_mod.Config, ids: []const u32, report: *Report, wide: usize) !void { - const layer_stem = "language_model.mtp.layers.0"; +fn putMtp(out: *Out, gpa: std.mem.Allocator, a: std.mem.Allocator, ck: *ckpt.Checkpoint, text: *ckpt.Checkpoint, root: []const u8, cfg: *const config_mod.Config, ids: []const u32, report: *Report, wide: usize) !void { + const layer_stem = try std.fmt.allocPrint(a, "{s}.layers.0", .{root}); try putHc(out, gpa, a, ck, cfg, "mtp.ahc", try std.fmt.allocPrint(a, "{s}.attn_hyper_connection", .{layer_stem}), true, report.norms_around_one, wide); try putHc(out, gpa, a, ck, cfg, "mtp.mhc", try std.fmt.allocPrint(a, "{s}.mlp_hyper_connection", .{layer_stem}), true, report.norms_around_one, wide); try putRouter(out, gpa, a, ck, cfg, "mtp.moe.router", try std.fmt.allocPrint(a, "{s}.mlp.gate", .{layer_stem}), try std.fmt.allocPrint(a, "{s}.mlp.shared_expert_gate", .{layer_stem})); @@ -573,18 +543,18 @@ fn putMtp(out: *Out, gpa: std.mem.Allocator, a: std.mem.Allocator, ck: *ckpt.Che const scale = try centeredScale(gpa, ck, try std.fmt.allocPrint(a, "{s}.self_attn.indexer.{s}", .{ layer_stem, n }), report.norms_around_one); try out.put(try std.fmt.allocPrint(a, "mtp.att.{s}", .{pack_name}), "F32", &.{cfg.index_dim}, scale); } - try putHc(out, gpa, a, ck, cfg, "mtp.mix", "language_model.mtp.hyper_connection_mixer", false, report.norms_around_one, wide); - const fce = try qlinear(ck, cfg, "language_model.mtp.fc_embedding"); + try putHc(out, gpa, a, ck, cfg, "mtp.mix", try std.fmt.allocPrint(a, "{s}.hyper_connection_mixer", .{root}), false, report.norms_around_one, wide); + const fce = try qlinear(ck, cfg, try std.fmt.allocPrint(a, "{s}.fc_embedding", .{root})); const fce_members = [_]Q{fce}; try putLane(out, gpa, "mtp.fce", fce_members[0..]); - const fch = try qlinear(ck, cfg, "language_model.mtp.fc_hidden"); + const fch = try qlinear(ck, cfg, try std.fmt.allocPrint(a, "{s}.fc_hidden", .{root})); const fch_members = [_]Q{fch}; try putLane(out, gpa, "mtp.fch", fch_members[0..]); - const enorm = try centeredScale(gpa, ck, "language_model.mtp.pre_fc_norm_embedding.weight", report.norms_around_one); + const enorm = try centeredScale(gpa, ck, try std.fmt.allocPrint(a, "{s}.pre_fc_norm_embedding.weight", .{root}), report.norms_around_one); try out.put("mtp.enorm.scale", "F32", &.{cfg.hidden}, enorm); - const hnorm = try centeredScale(gpa, ck, "language_model.mtp.pre_fc_norm_hidden.weight", report.norms_around_one); + const hnorm = try centeredScale(gpa, ck, try std.fmt.allocPrint(a, "{s}.pre_fc_norm_hidden.weight", .{root}), report.norms_around_one); try out.put("mtp.hnorm.scale", "F32", &.{wide}, hnorm); - const head = try qlinear(ck, cfg, "language_model.lm_head"); + const head = try qlinear(text, cfg, "language_model.lm_head"); // the head sits with the text model try putDraftHead(out, gpa, a, "mtp.draft", head, ids); } @@ -592,4 +562,5 @@ pub const compareFile = pio.compareFile; test { _ = @import("pack_test.zig"); + _ = @import("pack_hc_test.zig"); } diff --git a/zig/src/families/flashnext/pack_cache_test.zig b/zig/src/families/flashnext/pack_cache_test.zig index 2855ff8a1f..f291237ea7 100644 --- a/zig/src/families/flashnext/pack_cache_test.zig +++ b/zig/src/families/flashnext/pack_cache_test.zig @@ -295,3 +295,28 @@ test "cache readiness reads bounded metadata rather than copying the weight body var bounded: std.heap.FixedBufferAllocator = .init(&storage); try std.testing.expect(try pio.packReady(bounded.allocator(), io, path, "checkpoint")); } + +test "packs cached or left in the model folder under the format-less identity are rebuilt, never reused" { + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const paths = try Paths.init(tmp); + defer paths.deinit(); + const h = pio.hashBytes("index"); + const shards = [_]pio.ShardId{.{ .name = "model-00001-of-00001.safetensors", .size = 4, .header_sha256 = h }}; + const current = try pio.identityString(a, h, h, &shards); + defer a.free(current); + const old = current[std.mem.indexOfScalar(u8, current, '\n').? + 1 ..]; // the 6/32-only packer's identity text + try std.testing.expect(std.mem.startsWith(u8, old, "index ")); + var stale: Builder = .{ .identity = old }; + const old_dir = try cache.ensure(a, io, paths.model, paths.root, old, &stale); + defer a.free(old_dir); + try tmp.dir.createDir(io, "model/zig-pack", .default_dir); + const legacy = try std.fs.path.join(a, &.{ paths.model, "zig-pack" }); + defer a.free(legacy); + try stale.build(a, io, paths.model, legacy, ""); + var builder: Builder = .{ .identity = current }; + const dir = try cache.ensure(a, io, paths.model, paths.root, current, &builder); + defer a.free(dir); + try std.testing.expectEqual(@as(usize, 1), builder.calls); + try std.testing.expect(!std.mem.eql(u8, old_dir, dir) and !std.mem.eql(u8, legacy, dir)); +} diff --git a/zig/src/families/flashnext/pack_hash_test.zig b/zig/src/families/flashnext/pack_hash_test.zig new file mode 100644 index 0000000000..af41511c25 --- /dev/null +++ b/zig/src/families/flashnext/pack_hash_test.zig @@ -0,0 +1,47 @@ +//! The uniform 6-bit/group-32 packs keep their bytes: one digest over every tensor's name, dtype, shape and payload. +const std = @import("std"); +const pack = @import("pack.zig"); +const fixture = @import("pack_test.zig"); +const st = @import("../../core/safetensors.zig"); + +const gpa = std.testing.allocator; +const io = std.testing.io; + +/// The digest of the synthetic 6/32 checkpoint's three packs' tensors: a uniform pack keeps every byte. +const uniform_digest = "21fa29e16708022181ed842263cc03478cf76de2e213426f94feb05b58d8112e"; + +fn absorb(h: *std.crypto.hash.sha2.Sha256, a: std.mem.Allocator, path: []const u8) !void { + var file = try st.File.open(a, io, path); + defer file.close(io); + const names = try a.dupe([]const u8, file.names.keys()); + std.mem.sort([]const u8, names, {}, struct { + fn less(_: void, x: []const u8, y: []const u8) bool { + return std.mem.lessThan(u8, x, y); + } + }.less); + for (names) |name| { + const e = file.names.get(name).?; + h.update(name); + h.update(&.{ @backingInt(e.dtype), e.rank }); + for (e.shape[0..e.rank]) |d| h.update(std.mem.asBytes(&@as(u64, d))); + h.update(file.map.memory[file.data + e.begin .. file.data + e.end]); + } +} + +test "the uniform 6-bit group-32 packs hash as they did before mixed widths" { + var arena_state: std.heap.ArenaAllocator = .init(gpa); + defer arena_state.deinit(); + const a = arena_state.allocator(); + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try fixture.writeCheckpoint(tmp, a); + try tmp.dir.createDirPath(io, "out"); + _ = try pack.build(a, io, try fixture.tmpPath(a, tmp, "."), try fixture.tmpPath(a, tmp, "out"), null); + var h = std.crypto.hash.sha2.Sha256.init(.{}); + for ([_][]const u8{ "out/pack.safetensors", "out/pack_mlx.safetensors", "out/pack_mtp_mlx.safetensors" }) |name| { + try absorb(&h, a, try std.fs.path.join(a, &.{ try fixture.tmpPath(a, tmp, "."), name })); + } + var digest: [32]u8 = undefined; + h.final(&digest); + try std.testing.expectEqualStrings(uniform_digest, &std.fmt.bytesToHex(digest, .lower)); +} diff --git a/zig/src/families/flashnext/pack_hc_test.zig b/zig/src/families/flashnext/pack_hc_test.zig new file mode 100644 index 0000000000..637a4d3206 --- /dev/null +++ b/zig/src/families/flashnext/pack_hc_test.zig @@ -0,0 +1,72 @@ +//! Hyper-connection packs: a down and inject pair of two widths widens exactly to the wider, as QWeights.of. +const std = @import("std"); +const pack = @import("pack.zig"); +const mlx = @import("../../core/quant/mlx.zig"); +const st = @import("../../core/safetensors.zig"); +const affine = @import("affine.zig"); +const synthetic = @import("pack_widen_test.zig").synthetic; + +fn values(a: std.mem.Allocator, q: pack.Q) ![]f64 { + const out = try a.alloc(f64, q.rows * q.k); + try mlx.reference(.{ .words = q.w, .scales = q.s, .biases = q.b, .bits = q.spec.bits, .group = @intCast(q.spec.group) }, out); + return out; +} + +test "oQ4e's 4/64 down with a 5/64 inject packs as one 5/64 matrix holding the same values" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const down = try synthetic(a, 11, 64, 1024, 4, 64); + const inject = try synthetic(a, 12, 32, 1024, 5, 64); + const both = try pack.hcDown(a, &.{ down, inject }); + try std.testing.expectEqual(@as(u8, 5), both.spec.bits); + try std.testing.expectEqual(@as(usize, 64), both.spec.group); + try std.testing.expectEqual(@as(usize, 96), both.rows); + const got = try values(a, both); + try std.testing.expectEqualSlices(f64, try values(a, down), got[0 .. 64 * 1024]); + try std.testing.expectEqualSlices(f64, try values(a, inject), got[64 * 1024 ..]); + // one width keeps its bytes: the uniform pair is a plain row concatenation + const same = try pack.hcDown(a, &.{ down, try synthetic(a, 13, 32, 1024, 4, 64) }); + try std.testing.expectEqualSlices(u8, down.w.bytes, same.w.bytes[0..down.w.bytes.len]); +} + +test "the up projection's 320 inputs can't be stored in groups of 128, so it packs at its stored width" { + const e = struct { + fn of(dtype: st.DType, rows: usize, cols: usize) st.Entry { + return .{ .dtype = dtype, .rank = 2, .shape = .{ rows, cols, 0, 0, 0 }, .begin = 0, .end = rows * cols * dtype.size() }; + } + }.of; + // 8-bit codes over 320 inputs are 80 words a row; three groups of 128 would need 96 + const at128 = affine.matrix(e(.u32, 64, 80), e(.bf16, 64, 3), e(.bf16, 64, 3), try affine.Spec.init(8, 128)); + try std.testing.expectError(error.InvalidAffineShape, at128); + try std.testing.expectEqual(@as(usize, 320), (try affine.matrix(e(.u32, 64, 80), e(.bf16, 64, 5), e(.bf16, 64, 5), try affine.Spec.init(8, 64))).k); +} + +test "a prompt projection of several widths keeps each member's bytes in order and tells their widths" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var out = pack.Out.init(std.testing.allocator); + defer out.deinit(); + const q = try synthetic(a, 21, 64, 256, 3, 64); + const v = try synthetic(a, 22, 32, 256, 8, 128); + try pack.mlxProjection(&out, std.testing.allocator, a, "L3.att.proj", &.{ q, v }); + const names = out.names.items; + const at = struct { + fn find(n: []const []const u8, name: []const u8) usize { + for (n, 0..) |x, i| if (std.mem.eql(u8, x, name)) return i; + unreachable; + } + }.find; + const mw = out.blobs.items[at(names, "L3.att.proj.mw")]; + try std.testing.expectEqualSlices(u8, q.w.bytes, mw[0..q.w.bytes.len]); + try std.testing.expectEqualSlices(u8, v.w.bytes, mw[q.w.bytes.len..]); + try std.testing.expectEqualSlices(usize, &.{mw.len / 4}, out.shapes.items[at(names, "L3.att.proj.mw")]); + const parts = out.blobs.items[at(names, "L3.att.proj.parts")]; + try std.testing.expectEqualSlices(u8, std.mem.sliceAsBytes(&[_]u32{ 64, 3, 64, 32, 8, 128 }), parts); + // 6/32 members write the pinned 2-D shapes and no table + const six = [_]pack.Q{ try synthetic(a, 23, 32, 256, 6, 32), try synthetic(a, 24, 32, 256, 6, 32) }; + try pack.mlxProjection(&out, std.testing.allocator, a, "L7.att.proj", &six); + try std.testing.expectEqualSlices(usize, &.{ 64, 48 }, out.shapes.items[at(out.names.items, "L7.att.proj.mw")]); + for (out.names.items) |n| try std.testing.expect(!std.mem.eql(u8, n, "L7.att.proj.parts")); +} diff --git a/zig/src/families/flashnext/pack_io.zig b/zig/src/families/flashnext/pack_io.zig index 942bfac500..6dbe6760ed 100644 --- a/zig/src/families/flashnext/pack_io.zig +++ b/zig/src/families/flashnext/pack_io.zig @@ -28,15 +28,22 @@ pub const Out = struct { out.* = undefined; } + /// Takes `bytes` (from `gpa`) even when it fails: then they are freed and every list keeps its length. pub fn put(out: *Out, name: []const u8, dtype: []const u8, shape: []const usize, bytes: anytype) !void { + errdefer out.gpa.free(bytes); const a = out.arena.allocator(); - try out.names.append(out.gpa, try a.dupe(u8, name)); - try out.dtypes.append(out.gpa, dtype); - const shape_copy = try a.alloc(usize, shape.len); - @memcpy(shape_copy, shape); - try out.shapes.append(out.gpa, shape_copy); - try out.blobs.append(out.gpa, std.mem.sliceAsBytes(bytes)); - try out.blob_aligns.append(out.gpa, comptime std.mem.Alignment.fromByteUnits(@alignOf(std.meta.Child(@TypeOf(bytes))))); + const name_copy = try a.dupe(u8, name); + const shape_copy = try a.dupe(usize, shape); + try out.names.ensureUnusedCapacity(out.gpa, 1); + try out.dtypes.ensureUnusedCapacity(out.gpa, 1); + try out.shapes.ensureUnusedCapacity(out.gpa, 1); + try out.blobs.ensureUnusedCapacity(out.gpa, 1); + try out.blob_aligns.ensureUnusedCapacity(out.gpa, 1); + out.names.appendAssumeCapacity(name_copy); + out.dtypes.appendAssumeCapacity(dtype); + out.shapes.appendAssumeCapacity(shape_copy); + out.blobs.appendAssumeCapacity(std.mem.sliceAsBytes(bytes)); + out.blob_aligns.appendAssumeCapacity(comptime std.mem.Alignment.fromByteUnits(@alignOf(std.meta.Child(@TypeOf(bytes))))); } /// Writes the file with `source` recorded under __metadata__ so a reader can refuse a pack from another checkpoint. @@ -102,8 +109,11 @@ pub fn hashBytes(bytes: []const u8) [32]u8 { return h; } -/// The identity's canonical text: the index file's hash, then each shard's name, size and header hash by name. -pub fn identityString(gpa: std.mem.Allocator, index_sha256: ?[32]u8, shards: []const ShardId) ![]u8 { +/// The pack layout's version, first in the identity: 3 has width tables, ragged runs untiled, the config's hash. +pub const pack_format = 3; + +/// The identity's text: pack format, index and config hashes (the config names the widths), each shard's id. +pub fn identityString(gpa: std.mem.Allocator, index_sha256: ?[32]u8, config_sha256: ?[32]u8, shards: []const ShardId) ![]u8 { const sorted = try gpa.dupe(ShardId, shards); defer gpa.free(sorted); std.mem.sort(ShardId, sorted, {}, struct { @@ -113,7 +123,9 @@ pub fn identityString(gpa: std.mem.Allocator, index_sha256: ?[32]u8, shards: []c }.lt); var w: std.Io.Writer.Allocating = .init(gpa); errdefer w.deinit(); + try w.writer.print("flashnext pack {d}\n", .{pack_format}); if (index_sha256) |h| try w.writer.print("index {s}\n", .{&std.fmt.bytesToHex(h, .lower)}) else try w.writer.writeAll("index absent\n"); + if (config_sha256) |h| try w.writer.print("config {s}\n", .{&std.fmt.bytesToHex(h, .lower)}) else try w.writer.writeAll("config absent\n"); for (sorted) |s| try w.writer.print("shard {s} {d} {s}\n", .{ s.name, s.size, &std.fmt.bytesToHex(s.header_sha256, .lower) }); return w.toOwnedSlice(); } @@ -134,8 +146,8 @@ fn headerLen(memory: []const u8) !usize { return hl; } -/// Canonical identity text comes from the index bytes and each shard's mapped bytes, shared by builder and engine. -pub fn identityFromMapped(gpa: std.mem.Allocator, index_bytes: ?[]const u8, shards: []const MappedShard) ![]u8 { +/// The identity from the index and config bytes and each shard's mapped bytes, shared by builder and engine. +pub fn identityFromMapped(gpa: std.mem.Allocator, index_bytes: ?[]const u8, config_bytes: ?[]const u8, shards: []const MappedShard) ![]u8 { var ids: std.ArrayList(ShardId) = .empty; defer ids.deinit(gpa); for (shards) |s| { @@ -145,11 +157,18 @@ pub fn identityFromMapped(gpa: std.mem.Allocator, index_bytes: ?[]const u8, shar const hl = try headerLen(s.bytes[0..size]); try ids.append(gpa, .{ .name = s.name, .size = s.size, .header_sha256 = hashBytes(s.bytes[8..][0..hl]) }); } - return identityString(gpa, if (index_bytes) |b| hashBytes(b) else null, ids.items); + return identityString(gpa, if (index_bytes) |b| hashBytes(b) else null, if (config_bytes) |b| hashBytes(b) else null, ids.items); } -/// A checkpoint identity is the index hash plus each shard size and header hash, and mappings then unmap. +/// A checkpoint identity is the index and config hashes plus each shard size and header hash; mappings then unmap. pub fn sourceIdentity(gpa: std.mem.Allocator, io: Io, model_dir: []const u8) ![]u8 { + const config_path = try std.fs.path.join(gpa, &.{ model_dir, "config.json" }); + defer gpa.free(config_path); + const config = Io.Dir.cwd().readFileAlloc(io, config_path, gpa, .limited(1 << 24)) catch |e| switch (e) { + error.FileNotFound => null, + else => return e, + }; + defer if (config) |c| gpa.free(c); const Open = struct { file: Io.File, map: Io.File.MemoryMap }; var shards: std.ArrayList(MappedShard) = .empty; defer shards.deinit(gpa); @@ -162,14 +181,14 @@ pub fn sourceIdentity(gpa: std.mem.Allocator, io: Io, model_dir: []const u8) ![] const index_path = try std.fs.path.join(gpa, &.{ model_dir, "model.safetensors.index.json" }); defer gpa.free(index_path); var mapped = mapFile(io, index_path) catch |e| switch (e) { - error.FileNotFound => return identityFromMapped(gpa, null, shards.items), + error.FileNotFound => return identityFromMapped(gpa, null, config, shards.items), else => return e, }; defer mapped.map.destroy(io); defer mapped.file.close(io); var parsed = try std.json.parseFromSlice(std.json.Value, gpa, mapped.map.memory, .{}); defer parsed.deinit(); - const weight_map = parsed.value.object.get("weight_map") orelse return identityFromMapped(gpa, null, shards.items); + const weight_map = parsed.value.object.get("weight_map") orelse return identityFromMapped(gpa, null, config, shards.items); var names: std.ArrayList([]const u8) = .empty; defer { for (names.items) |n| gpa.free(n); @@ -191,7 +210,7 @@ pub fn sourceIdentity(gpa: std.mem.Allocator, io: Io, model_dir: []const u8) ![] try maps.append(gpa, .{ .file = shard.file, .map = shard.map }); try shards.append(gpa, .{ .name = name, .size = shard.len, .bytes = shard.map.memory }); } - const out = try identityFromMapped(gpa, mapped.map.memory, shards.items); + const out = try identityFromMapped(gpa, mapped.map.memory, config, shards.items); for (maps.items) |*m| { m.map.destroy(io); m.file.close(io); diff --git a/zig/src/families/flashnext/pack_lane.zig b/zig/src/families/flashnext/pack_lane.zig new file mode 100644 index 0000000000..24203adfc7 --- /dev/null +++ b/zig/src/families/flashnext/pack_lane.zig @@ -0,0 +1,188 @@ +//! Lane packs: words tiled, scale pairs group-major, and one run per bit width when a fused stack mixes widths. +const std = @import("std"); +const st = @import("../../core/safetensors.zig"); +const affine = @import("affine.zig"); +const pio = @import("pack_io.zig"); +const widen = @import("pack_widen.zig"); + +/// Which of a quantized linear's three tensors. +pub const Part = enum { weight, scales, biases }; + +/// A checkpoint quantized linear's three tensors with its verified affine spec. +pub const Q = struct { + w: st.Tensor, + s: st.Tensor, + b: st.Tensor, + rows: usize, + k: usize, + kw: usize, + spec: affine.Spec, + + pub fn part(r: Q, comptime which: Part) []const u8 { + return switch (which) { + .weight => r.w.bytes, + .scales => r.s.bytes, + .biases => r.b.bytes, + }; + } +}; + +/// The lane pipeline's word permutation: rows tiled nt at a time, a group's words innermost (lane_qmm.tile_weight). +pub fn tileWeight(gpa: std.mem.Allocator, flat: []const u8, rows: usize, kw: usize, nt: usize, w: usize) ![]u8 { + const words = std.mem.bytesAsSlice(u32, flat); + const out = try gpa.alloc(u8, flat.len); + errdefer gpa.free(out); + const out_words = std.mem.bytesAsSlice(u32, out); + const groups = kw / w; + for (0..rows / nt) |t| { + for (0..groups) |g| { + for (0..nt) |j| { + for (0..w) |c| { + out_words[((t * groups + g) * nt + j) * w + c] = words[(t * nt + j) * kw + g * w + c]; + } + } + } + } + return out; +} + +/// (K/GS, N, 2) bf16 (scale, bias) pairs, group-major, over a row-concatenated stack (lane_qmm.pack_scales). +pub fn packScales(gpa: std.mem.Allocator, members: []const Q) ![]u8 { + var rows: usize = 0; + for (members) |m| rows += m.rows; + const kg = members[0].s.dim(1); + const out = try gpa.alloc(u8, kg * rows * 4); + errdefer gpa.free(out); + const out16 = std.mem.bytesAsSlice(u16, out); + var at: usize = 0; + for (members) |m| { + const s16 = std.mem.bytesAsSlice(u16, m.part(.scales)); + const b16 = std.mem.bytesAsSlice(u16, m.part(.biases)); + for (0..kg) |g| { + for (0..m.rows) |j| { + out16[(g * rows + at + j) * 2] = s16[j * kg + g]; + out16[(g * rows + at + j) * 2 + 1] = b16[j * kg + g]; + } + } + at += m.rows; + } + return out; +} + +/// Row-concatenated members in the checkpoint's own MLX layout: a pure byte move per member. +pub fn concatPart(gpa: std.mem.Allocator, members: []const Q, comptime part: Part) ![]u8 { + var n: usize = 0; + for (members) |m| n += m.part(part).len; + const out = try gpa.alloc(u8, n); + errdefer gpa.free(out); + var at: usize = 0; + for (members) |m| { + @memcpy(out[at..][0..m.part(part).len], m.part(part)); + at += m.part(part).len; + } + return out; +} + +/// One bit width's members, in member order, at their smallest group: one product, as decode.py's `_stacked`. +pub const Run = struct { bits: u8, group: usize, members: u16, rows: usize }; + +/// Runs by bit width in first-seen order, adjacent or not, so split_k(N, K) sees 0.6.6's N; group 128 reads as 64. +pub fn runsOf(specs: []const affine.Spec, rows: []const usize, out: []Run) ![]Run { + if (specs.len == 0 or specs.len != rows.len) return error.EmptyLaneStack; + if (specs.len > 16) return error.TooManyLaneMembers; + var n: usize = 0; + for (specs, rows, 0..) |s, r, i| { + const bit = @as(u16, 1) << @intCast(i); + for (out[0..n]) |*run| { + if (run.bits != s.bits) continue; + run.members |= bit; + run.rows += r; + run.group = @min(run.group, s.group); + break; + } else { + if (n == out.len) return error.TooManyLaneRuns; + out[n] = .{ .bits = s.bits, .group = s.group, .members = bit, .rows = r }; + n += 1; + } + } + for (out[0..n]) |*run| run.group = @min(run.group, 64); + return out[0..n]; +} + +/// Whether a run of `rows` columns is tiled: 32 at a time when they divide, else MLX's layout (decode.py's `nt`). +pub fn tiled(rows: usize) bool { + return rows % 32 == 0; +} + +/// One width's `.wq` (tiled, or MLX's layout when ragged) and `.sbt` group-major pairs (decode.py's lane pipeline). +fn putRun(out: *pio.Out, gpa: std.mem.Allocator, name: []const u8, members: []const Q) !void { + var rows: usize = 0; + for (members) |m| rows += m.rows; + const spec = members[0].spec; + if (rows % 4 != 0) return error.LaneRowsNotFours; // lane_matmul's column pairs and scale pairs + const w = spec.group * spec.bits / 32; + const a = out.arena.allocator(); + const wq_name = try std.fmt.allocPrint(a, "{s}.wq", .{name}); + const sbt_name = try std.fmt.allocPrint(a, "{s}.sbt", .{name}); + const flat = try concatPart(gpa, members, .weight); + defer gpa.free(flat); + // put owns each blob from the call on, failing or not + try out.put(wq_name, "U32", &.{ rows, members[0].kw }, if (tiled(rows)) try tileWeight(gpa, flat, rows, members[0].kw, 32, w) else try gpa.dupe(u8, flat)); + try out.put(sbt_name, "BF16", &.{ members[0].s.dim(1), rows, 2 }, try packScales(gpa, members)); +} + +/// A lane projection; any width but 6/32 adds a `.runs` table, and several widths `.rN` per run and `.members`. +pub fn putLane(out: *pio.Out, gpa: std.mem.Allocator, name: []const u8, given: []const Q) !void { + putStack(out, gpa, name, given) catch |e| { + std.log.warn("lane {s}: {s}", .{ name, @errorName(e) }); // the pack refusal names its lane + return e; + }; +} + +fn putStack(out: *pio.Out, gpa: std.mem.Allocator, name: []const u8, given: []const Q) !void { + var specs: [16]affine.Spec = undefined; + var rows: [16]usize = undefined; + if (given.len > specs.len) return error.TooManyLaneRuns; + for (given, 0..) |m, i| { + specs[i] = m.spec; + rows[i] = m.rows; + } + var buf: [8]Run = undefined; + const runs = try runsOf(specs[0..given.len], rows[0..given.len], &buf); + const a = out.arena.allocator(); + var place: [16][3]u32 = undefined; // run, row within the run, row within the fused output + var fused: u32 = 0; + for (runs, 0..) |r, ri| { + var storage: [16]Q = undefined; + var count: usize = 0; + var at: u32 = 0; + for (given, 0..) |m, i| { + if ((r.members >> @intCast(i)) & 1 == 0) continue; + storage[count] = try widen.regroup(a, m, r.group); + count += 1; + place[i] = .{ @intCast(ri), at, 0 }; + at += @intCast(m.rows); + } + const run_name = if (runs.len == 1) name else try std.fmt.allocPrint(a, "{s}.r{d}", .{ name, ri }); + try putRun(out, gpa, run_name, storage[0..count]); + } + // the uniform packs keep their bytes + if (runs.len == 1 and runs[0].bits == 6 and runs[0].group == 32 and tiled(runs[0].rows)) return; + const runs_name = try std.fmt.allocPrint(a, "{s}.runs", .{name}); + const table = try gpa.alloc(u32, runs.len * 3); + for (runs, 0..) |r, i| table[i * 3 ..][0..3].* = .{ r.bits, @intCast(r.group), @intCast(r.rows) }; + try out.put(runs_name, "U32", &.{ runs.len, 3 }, table); + if (runs.len == 1) return; + for (given, 0..) |m, i| { + place[i][2] = fused; + fused += @intCast(m.rows); + } + const members_name = try std.fmt.allocPrint(a, "{s}.members", .{name}); + const members = try gpa.alloc(u32, given.len * 3); + for (place[0..given.len], 0..) |p, i| members[i * 3 ..][0..3].* = p; + try out.put(members_name, "U32", &.{ given.len, 3 }, members); +} + +test { + _ = @import("pack_lane_test.zig"); +} diff --git a/zig/src/families/flashnext/pack_lane_test.zig b/zig/src/families/flashnext/pack_lane_test.zig new file mode 100644 index 0000000000..42a14b6001 --- /dev/null +++ b/zig/src/families/flashnext/pack_lane_test.zig @@ -0,0 +1,136 @@ +//! A fused stack packs one product per bit width, members of equal bits together wherever they sit (`_stacked`). +const std = @import("std"); +const affine = @import("affine.zig"); +const lane = @import("pack_lane.zig"); +const pio = @import("pack_io.zig"); +const synthetic = @import("pack_widen_test.zig").synthetic; + +fn spec(bits: u8, group: usize) affine.Spec { + return affine.Spec.init(bits, group) catch unreachable; +} + +const att = [_]usize{ 12288, 512, 512, 640 }; // row counts of q, k, v and index_qk +const gdn = [_]usize{ 10240, 6144, 48, 48 }; // row counts of in_proj_qkv, z, b and a + +fn expectRuns(specs: []const affine.Spec, rows: []const usize, want: []const lane.Run) !void { + var buf: [8]lane.Run = undefined; + const runs = try lane.runsOf(specs, rows[0..specs.len], &buf); + try std.testing.expectEqual(want.len, runs.len); + for (want, runs) |w, r| try std.testing.expectEqual(w, r); +} + +test "members of equal bits form one run at their smallest group, in first-seen order, adjacent or not" { + // oQ2/oQ3/oQ4e attention: an 8-bit v between 6-bit q, k and index_qk; 0.6.6 runs q, k and index_qk as one product + try expectRuns(&.{ spec(6, 64), spec(6, 64), spec(8, 64), spec(6, 64) }, &att, &.{ + .{ .bits = 6, .group = 64, .members = 0b1011, .rows = 12288 + 512 + 640 }, + .{ .bits = 8, .group = 64, .members = 0b0100, .rows = 512 }, + }); + // a 6/128 v among 6/64 members is regrouped into the one 6/64 product + try expectRuns(&.{ spec(6, 64), spec(6, 64), spec(6, 128), spec(6, 64) }, &att, &.{ + .{ .bits = 6, .group = 64, .members = 0b1111, .rows = 13952 }, + }); + // oQ4e layer 21's GDN inputs: qkv 4/64, z and b 5/128, a 6/64; the 5-bit run is read in groups of 64 on the lane + try expectRuns(&.{ spec(4, 64), spec(5, 128), spec(5, 128), spec(6, 64) }, &gdn, &.{ + .{ .bits = 4, .group = 64, .members = 0b0001, .rows = 10240 }, + .{ .bits = 5, .group = 64, .members = 0b0110, .rows = 6192 }, + .{ .bits = 6, .group = 64, .members = 0b1000, .rows = 48 }, + }); + // equal bits in groups of 32 and 64 share the group of 32, so the run's group sums are 0.6.6's + try expectRuns(&.{ spec(4, 64), spec(4, 32) }, &att, &.{.{ .bits = 4, .group = 32, .members = 0b11, .rows = 12800 }}); + try expectRuns(&.{ spec(6, 32), spec(6, 32), spec(6, 32), spec(6, 32) }, &att, &.{ + .{ .bits = 6, .group = 32, .members = 0b1111, .rows = 13952 }, + }); +} + +fn find(out: *const pio.Out, name: []const u8) ?usize { + for (out.names.items, 0..) |n, i| if (std.mem.eql(u8, n, name)) return i; + return null; +} + +test "a stack of several widths writes a pack per run with its rows' places, and a regrouped one keeps the plain names" { + const gpa = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var out = pio.Out.init(gpa); + defer out.deinit(); + const q = try synthetic(a, 1, 32, 256, 6, 64); + const v = try synthetic(a, 2, 32, 256, 8, 64); + const idx = try synthetic(a, 3, 32, 256, 6, 64); + try lane.putLane(&out, gpa, "L1.att.proj", &.{ q, v, idx }); + try std.testing.expect(find(&out, "L1.att.proj.wq") == null); + try std.testing.expectEqualSlices(usize, &.{ 64, 48 }, out.shapes.items[find(&out, "L1.att.proj.r0.wq").?]); + try std.testing.expectEqualSlices(usize, &.{ 32, 64 }, out.shapes.items[find(&out, "L1.att.proj.r1.wq").?]); + // index_qk's first word follows q's 32 rows in the 6-bit run: the second tile's first word + const r0 = out.blobs.items[find(&out, "L1.att.proj.r0.wq").?]; + try std.testing.expectEqualSlices(u8, idx.w.bytes[0..4], r0[32 * 48 * 4 ..][0..4]); + const runs = out.blobs.items[find(&out, "L1.att.proj.runs").?]; // (bits, group, rows) per run + try std.testing.expectEqualSlices(u8, std.mem.sliceAsBytes(&[_]u32{ 6, 64, 64, 8, 64, 32 }), runs); + const places = out.blobs.items[find(&out, "L1.att.proj.members").?]; // (run, row in run, row in fused output) + try std.testing.expectEqualSlices(u8, std.mem.sliceAsBytes(&[_]u32{ 0, 0, 0, 1, 0, 32, 0, 32, 64 }), places); + const one = [_]lane.Q{ try synthetic(a, 4, 32, 256, 6, 64), try synthetic(a, 5, 32, 256, 6, 128) }; + try lane.putLane(&out, gpa, "L3.att.proj", &one); + try std.testing.expectEqualSlices(usize, &.{ 4, 64, 2 }, out.shapes.items[find(&out, "L3.att.proj.sbt").?]); + const regrouped = out.blobs.items[find(&out, "L3.att.proj.runs").?]; // one run, its width told + try std.testing.expectEqualSlices(u8, std.mem.sliceAsBytes(&[_]u32{ 6, 64, 64 }), regrouped); + try std.testing.expect(find(&out, "L3.att.proj.members") == null); + const plain = [_]lane.Q{ try synthetic(a, 6, 32, 256, 6, 32), try synthetic(a, 7, 32, 256, 6, 32) }; + try lane.putLane(&out, gpa, "L7.att.proj", &plain); + try std.testing.expect(find(&out, "L7.att.proj.runs") == null); // 6/32 writes no table: its bytes are pinned + const four = [_]lane.Q{try synthetic(a, 8, 64, 256, 4, 64)}; + try lane.putLane(&out, gpa, "L9.att.o", &four); + try std.testing.expectEqualSlices(usize, &.{ 64, 32 }, out.shapes.items[find(&out, "L9.att.o.wq").?]); + // 64 rows of 4 bits tile 32 columns at a time, as the lane launches read them: the second tile follows the first + const wq = out.blobs.items[find(&out, "L9.att.o.wq").?]; + try std.testing.expectEqualSlices(u8, four[0].w.bytes[32 * 32 * 4 ..][0..4], wq[4 * 32 * 8 * 4 ..][0..4]); +} + +test "runs whose columns aren't a multiple of 32 keep MLX's layout untiled, as decode.py's nt 0, at the GDN rows" { + const gpa = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var out = pio.Out.init(gpa); + defer out.deinit(); + const widths = [_][2]usize{ .{ 4, 64 }, .{ 5, 128 }, .{ 5, 128 }, .{ 6, 64 } }; + var ms: [4]lane.Q = undefined; + for (&ms, gdn, widths, 0..) |*m, rows, w, i| m.* = try synthetic(a, 20 + i, rows, 256, @intCast(w[0]), w[1]); + try lane.putLane(&out, gpa, "L21.gdn.in", &ms); + const runs = out.blobs.items[find(&out, "L21.gdn.in.runs").?]; + try std.testing.expectEqualSlices(u8, std.mem.sliceAsBytes(&[_]u32{ 4, 64, 10240, 5, 64, 6192, 6, 64, 48 }), runs); + try std.testing.expect(lane.tiled(10240) and !lane.tiled(6192) and !lane.tiled(48)); + // the ragged 5-bit run is z then b row by row, and the 48-row a is its own bytes + const r1 = out.blobs.items[find(&out, "L21.gdn.in.r1.wq").?]; + try std.testing.expectEqualSlices(u8, ms[1].w.bytes, r1[0..ms[1].w.bytes.len]); + try std.testing.expectEqualSlices(u8, ms[2].w.bytes, r1[ms[1].w.bytes.len..]); + try std.testing.expectEqualSlices(u8, ms[3].w.bytes, out.blobs.items[find(&out, "L21.gdn.in.r2.wq").?]); + const places = out.blobs.items[find(&out, "L21.gdn.in.members").?]; + try std.testing.expectEqualSlices(u8, std.mem.sliceAsBytes(&[_]u32{ 0, 0, 0, 1, 0, 10240, 1, 6144, 16384, 2, 0, 16432 }), places); + // a 6/32 lane whose rows don't tile carries its table, so it is never read as a tiled uniform lane + const odd = [_]lane.Q{try synthetic(a, 30, 48, 256, 6, 32)}; + try lane.putLane(&out, gpa, "L0.gdn.out", &odd); + try std.testing.expect(find(&out, "L0.gdn.out.runs") != null); + // columns not in fours have no lane layout: refused by name + const bad = [_]lane.Q{try synthetic(a, 31, 50, 256, 6, 32)}; + try std.testing.expectError(error.LaneRowsNotFours, lane.putLane(&out, gpa, "L0.att.o", &bad)); +} + +fn packMixed(gpa: std.mem.Allocator, stack: []const lane.Q, proj: []const lane.Q) !void { + var out = pio.Out.init(gpa); + defer out.deinit(); + try lane.putLane(&out, gpa, "L21.gdn.in", stack); + try @import("pack.zig").mlxProjection(&out, gpa, out.arena.allocator(), "L21.gdn.in", proj); +} + +test "a lane stack and a prompt projection of several widths free every buffer when any allocation fails" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + // oQ4e's layer-21 GDN input in miniature: 4/64, 5/128 twice, 6/64, with a ragged 48-row 6-bit run + const l21 = [_]lane.Q{ try synthetic(a, 31, 64, 256, 4, 64), try synthetic(a, 32, 64, 256, 5, 128), try synthetic(a, 33, 48, 256, 5, 128), try synthetic(a, 34, 48, 256, 6, 64) }; + // a fixed buffer under the failing allocator: every run grows and resizes the same way at every optimize mode + const backing = try std.testing.allocator.alloc(u8, 64 << 20); + defer std.testing.allocator.free(backing); + var fixed: std.heap.FixedBufferAllocator = .init(backing); + try std.testing.checkAllAllocationFailures(fixed.allocator(), packMixed, .{ &l21, &l21 }); +} diff --git a/zig/src/families/flashnext/pack_raw.zig b/zig/src/families/flashnext/pack_raw.zig new file mode 100644 index 0000000000..d2328f0868 --- /dev/null +++ b/zig/src/families/flashnext/pack_raw.zig @@ -0,0 +1,42 @@ +//! Tensors the engine reads raw (experts, embedding, n-gram tables) checked at pack build against declared widths. +const std = @import("std"); +const st = @import("../../core/safetensors.zig"); +const ckpt = @import("../../core/checkpoint.zig"); +const config_mod = @import("config.zig"); +const Widths = @import("widths.zig").Widths; + +fn entry(t: st.Tensor) st.Entry { + return .{ .dtype = t.dtype, .rank = t.rank, .shape = t.shape, .begin = 0, .end = t.bytes.len }; +} + +/// Module `name`'s stored affine triple against the width declared for `key` (unquantized: admission's to judge). +pub fn check(ck: *ckpt.Checkpoint, w: *const Widths, key: []const u8, name: []const u8) !void { + var buf: [3][256]u8 = undefined; + var e: [3]st.Entry = undefined; + for ([_][]const u8{ "weight", "scales", "biases" }) |p| if (!ck.has(try std.fmt.bufPrint(&buf[0], "{s}.{s}", .{ name, p }))) return; + for ([_][]const u8{ "weight", "scales", "biases" }, 0..) |p, i| e[i] = entry(try ck.get(try std.fmt.bufPrint(&buf[i], "{s}.{s}", .{ name, p }))); + _ = w.shaped(key, e[0], e[1], e[2]) catch |err| { + std.log.warn("{s}: not stored at its declared width: {s}", .{ name, @errorName(err) }); + return err; + }; +} + +const projs = [_][]const u8{ "switch_mlp.gate_proj", "switch_mlp.up_proj", "shared_expert.gate_proj", "shared_expert.up_proj", "switch_mlp.down_proj", "shared_expert.down_proj" }; + +/// Every raw-read module: each layer's and the MTP head's experts, the embedding and the n-gram shards. +pub fn checkAll(ck: *ckpt.Checkpoint, ckm: *ckpt.Checkpoint, w: *const Widths, cfg: *const config_mod.Config, mtp_root: []const u8, spelling: []const u8) !void { + var buf: [256]u8 = undefined; + var key: [256]u8 = undefined; + for (0..cfg.layers) |i| for (projs) |p| { + const name = try std.fmt.bufPrint(&buf, "language_model.model.layers.{d}.mlp.{s}", .{ i, p }); + try check(ck, w, name, name); + }; + if (cfg.mtp.layers > 0) for (projs) |p| { + try check(ckm, w, try std.fmt.bufPrint(&key, "language_model.mtp.layers.0.mlp.{s}", .{p}), try std.fmt.bufPrint(&buf, "{s}.layers.0.mlp.{s}", .{ mtp_root, p })); + }; + try check(ck, w, "language_model.model.embed_tokens", "language_model.model.embed_tokens"); + for (0..cfg.layers) |i| if (cfg.ple[i]) for (0..cfg.ngram_shards) |k| { + const name = try std.fmt.bufPrint(&buf, "language_model.model.layers.{d}.ple.ple_embedding.ngram_embedding.{s}{d}", .{ i, spelling, k }); + try check(ck, w, name, name); + }; +} diff --git a/zig/src/families/flashnext/pack_source_test.zig b/zig/src/families/flashnext/pack_source_test.zig index e184b47e87..642e0ff667 100644 --- a/zig/src/families/flashnext/pack_source_test.zig +++ b/zig/src/families/flashnext/pack_source_test.zig @@ -109,9 +109,10 @@ test "the io and mapped identity paths produce the same text" { const bytes = try tmp.dir.readFileAlloc(io, kv.value_ptr.string, a, .limited(1 << 30)); try shards.append(a, .{ .name = kv.value_ptr.string, .size = bytes.len, .bytes = bytes }); } - const through_maps = try pio.identityFromMapped(a, index_bytes, shards.items); + const config_bytes = try tmp.dir.readFileAlloc(io, "config.json", a, .limited(1 << 24)); + const through_maps = try pio.identityFromMapped(a, index_bytes, config_bytes, shards.items); try std.testing.expectEqualStrings(through_io, through_maps); - try std.testing.expect(std.mem.startsWith(u8, through_maps, "index ")); + try std.testing.expect(std.mem.startsWith(u8, through_maps, "flashnext pack 3\nindex ")); } test "the identity string does not depend on the shard list's order" { @@ -123,9 +124,9 @@ test "the identity string does not depend on the shard list's order" { .{ .name = "model-00002.safetensors", .size = 20, .header_sha256 = h2 }, }; const two = [_]pio.ShardId{ one[1], one[0] }; - const s1 = try pio.identityString(a, h1, one[0..]); + const s1 = try pio.identityString(a, h1, h2, one[0..]); defer a.free(s1); - const s2 = try pio.identityString(a, h1, two[0..]); + const s2 = try pio.identityString(a, h1, h2, two[0..]); defer a.free(s2); try std.testing.expectEqualStrings(s1, s2); try std.testing.expect(std.mem.indexOf(u8, s1, "shard model-00001.safetensors 10 ") != null); @@ -165,3 +166,19 @@ test "the no-dump pack cache is ready only complete and matching, resolved from const changed = try pio.sourceIdentity(a, io, model_dir); try std.testing.expect(!try pio.packsReady(a, io, cache_dir, changed)); } + +test "a config-only change (the widths the pack is cut at) gives a new identity, so a stale pack is never reused" { + var arena_state: std.heap.ArenaAllocator = .init(gpa); + defer arena_state.deinit(); + const a = arena_state.allocator(); + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try fixture.writeCheckpoint(tmp, a); + const model_dir = try fixture.tmpPath(a, tmp, "."); + const before = try pio.sourceIdentity(a, io, model_dir); + try std.testing.expect(std.mem.indexOf(u8, before, "\nconfig ") != null); + const config = try tmp.dir.readFileAlloc(io, "config.json", a, .limited(1 << 24)); + try tmp.dir.writeFile(io, .{ .sub_path = "config.json", .data = try std.mem.concat(a, u8, &.{ config, " " }) }); + const after = try pio.sourceIdentity(a, io, model_dir); + try std.testing.expect(!std.mem.eql(u8, before, after)); +} diff --git a/zig/src/families/flashnext/pack_test.zig b/zig/src/families/flashnext/pack_test.zig index 2cb19d3baf..66db631c95 100644 --- a/zig/src/families/flashnext/pack_test.zig +++ b/zig/src/families/flashnext/pack_test.zig @@ -162,6 +162,7 @@ pub fn writeCheckpoint(tmp: std.testing.TmpDir, a: std.mem.Allocator) !void { // the vision tower's rank-five entry: the builder never reads it, and it must not stop the load try list.append(a, .{ .name = "model.visual.patch_embed.proj.weight", .dtype = "BF16", .shape = &.{ 1, 1, 1, 1, 1 }, .bytes = &.{ 0, 0 } }); try qspecs(&list, a, "language_model.lm_head", 128, 500_000, 0x3C00, 0x3B00); + try qspecs(&list, a, "language_model.model.embed_tokens", 128, 600_000, 0x3C00, 0x3B00); // read raw by the engine const mtp = "language_model.mtp"; try hcSpecs(&list, a, try std.fmt.allocPrint(a, "{s}.layers.0.attn_hyper_connection", .{mtp}), 13_000_000, true, 0x3C00, 0x3400); try hcSpecs(&list, a, try std.fmt.allocPrint(a, "{s}.layers.0.mlp_hyper_connection", .{mtp}), 13_500_000, true, 0x3C00, 0x3400); @@ -523,3 +524,35 @@ test "the draft vocabulary builds the mtp decode tensors and gathers the head's try expectTensor(&file, "mtp.hnorm.scale", .f32, &.{32}); try expectTensor(&file, "mtp.fce.wq", .u32, &.{ 32, 6 }); } + +test { + _ = @import("pack_hash_test.zig"); +} + +test "a pack never reads a tensor at a width the config leaves half-declared, though its shapes would fit" { + var arena_state: std.heap.ArenaAllocator = .init(gpa); + defer arena_state.deinit(); + const a = arena_state.allocator(); + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try writeCheckpoint(tmp, a); + // stored 6/32; a 3-bit group-64 reading of the same words and tables fits every shape check + const declared = "\"quantization\": {\"mtp.fc_embedding\": {\"bits\": 3}, "; + try tmp.dir.writeFile(io, .{ .sub_path = "config.json", .data = try std.mem.replaceOwned(u8, a, config_json, "\"quantization\": {", declared) }); + try tmp.dir.createDirPath(io, "out"); + try std.testing.expectError(error.UnsupportedQuantization, pack.build(a, io, try tmpPath(a, tmp, "."), try tmpPath(a, tmp, "out"), null)); +} + +test "a directly read tensor (expert, embedding, n-gram table) whose shapes contradict its declared width is refused" { + var arena_state: std.heap.ArenaAllocator = .init(gpa); + defer arena_state.deinit(); + const a = arena_state.allocator(); + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try writeCheckpoint(tmp, a); + // stored 6/32; the embedding declared 8/64 would be read by the lookup kernel at the wrong width + const declared = "\"quantization\": {\"language_model.model.embed_tokens\": {\"bits\": 8, \"group_size\": 64}, "; + try tmp.dir.writeFile(io, .{ .sub_path = "config.json", .data = try std.mem.replaceOwned(u8, a, config_json, "\"quantization\": {", declared) }); + try tmp.dir.createDirPath(io, "out"); + try std.testing.expectError(error.InvalidAffineShape, pack.build(a, io, try tmpPath(a, tmp, "."), try tmpPath(a, tmp, "out"), null)); +} diff --git a/zig/src/families/flashnext/pack_widen.zig b/zig/src/families/flashnext/pack_widen.zig new file mode 100644 index 0000000000..f9a6ddfcc6 --- /dev/null +++ b/zig/src/families/flashnext/pack_widen.zig @@ -0,0 +1,28 @@ +//! Exact regrouping for lane packs: core/quant widen at the same bits, a group read as smaller ones. +const std = @import("std"); +const quant_mlx = @import("../../core/quant/mlx.zig"); +const quant_widen = @import("../../core/quant/widen.zig"); +const affine = @import("affine.zig"); +const Q = @import("pack_lane.zig").Q; + +/// A pack linear as core/quant sees it. +pub fn host(q: Q) quant_mlx.Host { + return .{ .words = q.w, .scales = q.s, .biases = q.b, .bits = q.spec.bits, .group = @intCast(q.spec.group) }; +} + +/// A core/quant linear back as a pack linear, its K from the tables. +pub fn fromHost(h: quant_mlx.Host) !Q { + const k = h.scales.dim(1) * h.group; + return .{ .w = h.words, .s = h.scales, .b = h.biases, .rows = h.words.dim(0), .k = k, .kw = h.words.dim(1), .spec = try affine.Spec.init(h.bits, h.group) }; +} + +/// `q` in groups of `group`, which divides its own: codes kept, each scale and bias repeated (`QWeights.widened`). +pub fn regroup(a: std.mem.Allocator, q: Q, group: usize) !Q { + if (group == q.spec.group) return q; + if (group > q.spec.group or q.spec.group % group != 0) return error.InexactRegroup; + return fromHost(try quant_widen.widen(a, host(q), .{ .bits = q.spec.bits, .group = @intCast(group) })); +} + +test { + _ = @import("pack_widen_test.zig"); +} diff --git a/zig/src/families/flashnext/pack_widen_test.zig b/zig/src/families/flashnext/pack_widen_test.zig new file mode 100644 index 0000000000..ca8689450e --- /dev/null +++ b/zig/src/families/flashnext/pack_widen_test.zig @@ -0,0 +1,67 @@ +//! A group read as smaller ones dequantizes to the same float64 bits at every width, and the codes keep their bytes. +const std = @import("std"); +const st = @import("../../core/safetensors.zig"); +const mlx = @import("../../core/quant/mlx.zig"); +const affine = @import("affine.zig"); +const lane = @import("pack_lane.zig"); +const widen = @import("pack_widen.zig"); + +/// A [rows, k] matrix at `bits`/`group` with seeded codes and finite bf16 scales and biases, owned by `a`. +pub fn synthetic(a: std.mem.Allocator, seed: u64, rows: usize, k: usize, bits: u8, group: usize) !lane.Q { + var prng = std.Random.DefaultPrng.init(seed); + const rand = prng.random(); + const kw = k * bits / 32; + const w = try a.alloc(u8, rows * kw * 4); + rand.bytes(w); + const tables = [2][]u8{ try a.alloc(u8, rows * (k / group) * 2), try a.alloc(u8, rows * (k / group) * 2) }; + for (tables) |t| for (std.mem.bytesAsSlice(u16, t)) |*v| { + const f: f32 = (rand.float(f32) - 0.5) * 4; + v.* = @intCast(@as(u32, @bitCast(f)) >> 16); + }; + var shape_w: [st.max_rank]usize = @splat(0); + shape_w[0] = rows; + shape_w[1] = kw; + var shape_s = shape_w; + shape_s[1] = k / group; + return .{ + .w = .{ .dtype = .u32, .rank = 2, .shape = shape_w, .bytes = w }, + .s = .{ .dtype = .bf16, .rank = 2, .shape = shape_s, .bytes = tables[0] }, + .b = .{ .dtype = .bf16, .rank = 2, .shape = shape_s, .bytes = tables[1] }, + .rows = rows, + .k = k, + .kw = kw, + .spec = try affine.Spec.init(bits, group), + }; +} + +fn reference(a: std.mem.Allocator, q: lane.Q) ![]f64 { + const out = try a.alloc(f64, q.rows * q.k); + try mlx.reference(.{ .words = q.w, .scales = q.s, .biases = q.b, .bits = q.spec.bits, .group = @intCast(q.spec.group) }, out); + return out; +} + +test "a group read as smaller ones is an exact identity in float64 at every width" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + for ([_][2]usize{ .{ 128, 64 }, .{ 128, 32 }, .{ 64, 32 } }) |pair| for ([_]u8{ 2, 3, 4, 5, 6, 8 }, 0..) |bits, i| { + const q = try synthetic(a, 0x5eed + i + pair[1], 4, 512, bits, pair[0]); + const h = try widen.regroup(a, q, pair[1]); + try std.testing.expectEqual(pair[1], h.spec.group); + try std.testing.expectEqual(bits, h.spec.bits); + try std.testing.expectEqual(512 / pair[1], h.s.dim(1)); + try std.testing.expectEqualSlices(u8, q.w.bytes, h.w.bytes); + const want = try reference(a, q); + const got = try reference(a, h); + try std.testing.expectEqualSlices(u64, std.mem.bytesAsSlice(u64, std.mem.sliceAsBytes(want)), std.mem.bytesAsSlice(u64, std.mem.sliceAsBytes(got))); + }; +} + +test "a group is never read as a larger one, and its own group passes through untouched" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const q = try synthetic(a, 7, 2, 128, 6, 32); + try std.testing.expectEqual(q.s.bytes.ptr, (try widen.regroup(a, q, 32)).s.bytes.ptr); + try std.testing.expectError(error.InexactRegroup, widen.regroup(a, q, 64)); +} diff --git a/zig/src/families/flashnext/widths.zig b/zig/src/families/flashnext/widths.zig new file mode 100644 index 0000000000..f74f69a4e6 --- /dev/null +++ b/zig/src/families/flashnext/widths.zig @@ -0,0 +1,76 @@ +//! Per-module MLX widths through core/quant, which indexes overrides by module; `false` and `{}` mark plain modules. +const std = @import("std"); +const mlx = @import("../../core/quant/mlx.zig"); +const st = @import("../../core/safetensors.zig"); +const affine = @import("affine.zig"); +const config = @import("config.zig"); + +pub const Width = mlx.Width; + +/// The widths the device kernels read: MLX affine 2, 3, 4, 5, 6 or 8 bits in groups of 32, 64 or 128. +pub fn launched(w: Width) bool { + const bits_ok = w.bits >= 2 and w.bits <= 8 and w.bits != 7; + return bits_ok and (w.group == 32 or w.group == 64 or w.group == 128); +} + +pub const Widths = struct { + arena: std.heap.ArenaAllocator, + /// Overrides by config key and by canonical module (`modules`); null is a width core/quant refuses (no guessing). + q: mlx.Config, + /// Modules the config declares unquantized (`false` or `{}`). + plain: std.StringHashMapUnmanaged(void) = .empty, + + pub fn init(gpa: std.mem.Allocator, c: *const config.Config) !Widths { + var w: Widths = .{ .arena = .init(gpa), .q = undefined }; + errdefer w.arena.deinit(); + const a = w.arena.allocator(); + // core/quant indexes overrides by module and refuses spellings that disagree + w.q = (mlx.detect(a, .{ .root = c.parsed.value.object, .quantization = c.quant }) catch |err| return switch (err) { + error.UnsupportedQuantization => error.ConflictingAffineAlias, // the only table core/quant refuses + else => err, + }) orelse return error.UnsupportedAffineMode; + var it = c.quant.iterator(); + while (it.next()) |item| { + const plain = switch (item.value_ptr.*) { + .bool => |on| !on, + .object => |o| o.count() == 0, + else => continue, + }; + const key = config.canonical(item.key_ptr.*); + if (!plain) { + if (w.plain.contains(key)) return error.ConflictingAffineAlias; + continue; + } + if (w.q.modules.get(key)) |declared| if (declared != null) return error.ConflictingAffineAlias; + try w.plain.put(a, key, {}); + } + return w; + } + + pub fn deinit(w: *Widths) void { + w.arena.deinit(); + w.* = undefined; + } + + /// No override at all: every quantized matrix has the global width. + pub fn uniform(w: *const Widths) bool { + return w.q.overrides.count() == 0 and w.plain.count() == 0; + } + + /// The module's width under any spelling of its key, or null when the config declares it unquantized. + pub fn of(w: *const Widths, path: []const u8) !?Width { + const key = config.canonical(path); + if (w.plain.contains(key)) return null; + return try w.q.width(key); + } + + /// The module's stored matrix checked against its own width, never one inferred from the shapes. + pub fn shaped(w: *const Widths, path: []const u8, weight: st.Entry, scales: st.Entry, biases: st.Entry) !affine.Shape { + const width = (try w.of(path)) orelse return error.UnquantizedAffineTensor; + return affine.matrix(weight, scales, biases, try affine.Spec.init(width.bits, width.group)); + } +}; + +test { + _ = @import("widths_test.zig"); +} diff --git a/zig/src/families/flashnext/widths_test.zig b/zig/src/families/flashnext/widths_test.zig new file mode 100644 index 0000000000..fa32e5004c --- /dev/null +++ b/zig/src/families/flashnext/widths_test.zig @@ -0,0 +1,85 @@ +//! Per-module widths: one module under any key spelling, its own width or a refusal, never a guess from shapes. +const std = @import("std"); +const st = @import("../../core/safetensors.zig"); +const config = @import("config.zig"); +const widths = @import("widths.zig"); +const a = std.testing.allocator; +const checked_config = @embedFile("fixtures/config.json"); + +fn overridden(entries: []const u8) ![]u8 { + const declaration = try std.fmt.allocPrint(a, "\"mode\": \"affine\", {s}", .{entries}); + defer a.free(declaration); + return std.mem.replaceOwned(u8, a, checked_config, "\"mode\": \"affine\"", declaration); +} + +const mixed = + \\"language_model.model.layers.0.self_attn.v_proj": {"bits": 8, "group_size": 64}, + \\"model.layers.1.linear_attn.in_proj_z.weight": {"bits": 5, "group_size": 128}, + \\"language_model.model.layers.0.mlp.gate": false, + \\"model.layers.2.self_attn.o_proj": {"bits": 4} +; + +fn expectWidth(w: *const widths.Widths, path: []const u8, bits: u8, group: u16) !void { + const got = (try w.of(path)).?; + try std.testing.expectEqual(bits, got.bits); + try std.testing.expectEqual(group, got.group); +} + +test "a mixed-override config resolves each module under either prefix and never fills a missing field" { + const bytes = try overridden(mixed); + defer a.free(bytes); + var c = try config.parse(a, bytes); + defer c.deinit(); + var w = try widths.Widths.init(a, &c); + defer w.deinit(); + try std.testing.expect(!w.uniform()); + try expectWidth(&w, "model.layers.0.self_attn.v_proj", 8, 64); + try expectWidth(&w, "language_model.model.layers.0.self_attn.v_proj.weight", 8, 64); + try expectWidth(&w, "language_model.model.layers.1.linear_attn.in_proj_z", 5, 128); + try expectWidth(&w, "language_model.model.layers.1.linear_attn.in_proj_qkv", 6, 32); + try std.testing.expect((try w.of("model.layers.0.mlp.gate")) == null); + // the family's own reader fills 4-bit group 64 here; core/quant refuses the half-declared width + try std.testing.expectError(error.UnsupportedQuantization, w.of("language_model.model.layers.2.self_attn.o_proj")); +} + +test "two spellings of one module that disagree are refused" { + const bytes = try overridden( + \\"language_model.model.layers.0.self_attn.v_proj": {"bits": 8, "group_size": 64}, + \\"model.layers.0.self_attn.v_proj.weight": {"bits": 4, "group_size": 64} + ); + defer a.free(bytes); + var c = try config.parse(a, bytes); + defer c.deinit(); + try std.testing.expectError(error.ConflictingAffineAlias, widths.Widths.init(a, &c)); +} + +fn entry(dtype: st.DType, rows: usize, cols: usize) st.Entry { + var shape: [st.max_rank]usize = @splat(0); + shape[0] = rows; + shape[1] = cols; + return .{ .dtype = dtype, .rank = 2, .shape = shape, .begin = 0, .end = rows * cols * dtype.size() }; +} + +/// A stored [rows, k] matrix's three entries at `bits`/`group`. +fn stored(rows: usize, k: usize, bits: usize, group: usize) [3]st.Entry { + return .{ entry(.u32, rows, k * bits / 32), entry(.bf16, rows, k / group), entry(.bf16, rows, k / group) }; +} + +test "a tensor packed at its override width is admitted and one at a guessed width is refused" { + const bytes = try overridden(mixed); + defer a.free(bytes); + var c = try config.parse(a, bytes); + defer c.deinit(); + var w = try widths.Widths.init(a, &c); + defer w.deinit(); + const v = "model.layers.0.self_attn.v_proj"; // its override is spelled under language_model. + const at_override = stored(8, 128, 8, 64); + const shape = try w.shaped(v, at_override[0], at_override[1], at_override[2]); + try std.testing.expectEqual(@as(usize, 128), shape.k); + try std.testing.expectEqual(@as(usize, 2), shape.groups); + const at_global = stored(8, 128, 6, 32); // the global width, which a prefix-blind lookup would guess + try std.testing.expectError(error.InvalidAffineShape, w.shaped(v, at_global[0], at_global[1], at_global[2])); + // a module with no override is read at the global width only, whatever width its shapes would fit + try std.testing.expectError(error.InvalidAffineShape, w.shaped("model.layers.0.self_attn.k_proj", at_override[0], at_override[1], at_override[2])); + try std.testing.expectError(error.UnquantizedAffineTensor, w.shaped("model.layers.0.mlp.gate", at_global[0], at_global[1], at_global[2])); +} From e369e939d48b24f4ebf6a41bfe0adf8e937ee159 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 04:42:39 -1000 Subject: [PATCH 09/10] Flash Next runs every MLX width on the GPU, with and without tensor units, in decode and in prompt - Lanes (decode, the MTP head, lane rounds; lane_w.zig): the recorded lane kernel's text compiled per (N, K, bits, group): codes staged as bytes, group sums over the group, K slices as lane_qmm.split_k; one launch a run, columns placed through the .members map (FZ_MAP in the kernel's store), untiled runs as MLX lays them out. A 6/32 lane keeps its recorded kernels and launches; a lane at other widths holds no 6/32 pair, so a 6-bit kernel can't read one of its runs (LaneNot6Bit32). - Hyper-connections, experts, embedding and n-gram tables (variants.zig): 0.6.6's any-width qa_ kernels, recorded at 6/32, launched at the tensor's BITS/GS (WB/WG, SWB/SWG for the experts). The 6-bit fast paths (fz_xgu/fz_xdown, fz_dense, fz_hc_up) run only on 6/32 tensors; an expert whose gate and up differ in width, and n-gram shards of different widths, are refused (one kernel reads both). - Prompt products (prompt_w.zig, prefill/qmmw_nax.metal): qmm6_nax's tiles and fp32 sums with codes read at WB bits in groups of WG: dense, split-K (the hyper-connection down), sorted-expert gathers and the fused gate/up; a fused projection of several widths runs one product a member into its columns. Every prompt weight is registered with its width at load, and a launch on an unregistered one fails the pass by name (PromptWeightWidthUnknown). - Fail closed: every width kernel compiles at load (engine_warm.zig), so a compile failure refuses the load, not a request. Width tables are read on the host, never through device buffers. Speed-up mode (two Macs) splits the 6/32 kernels only and refuses any other checkpoint by name (SpeedUpNeeds6Bit32). The replay harness refuses other widths (HarnessNeeds6Bit32). The config and index are read once at load. - The n-gram table's weight_scale (ple_scale.zig): oQ conversions of the FP8 checkpoint store the unscaled rows and oMLX's runtime multiplies each looked-up row by the scale (bf16 row times bf16 scale, one rounding). One in-place kernel after the lookup, in decode and in prompt, compiled at load only when the scale is not 1. The uniform 6/32 checkpoint keeps every recorded kernel and launch. Tests that fail before (none of these files exist on the previous commit): lane_w_test.zig (split_k gives every recorded kernel's SK; every recorded lane text takes any width, a row stride and a column map and keeps its arithmetic; runs land at their members' columns; width tables read on the host; a lane at other widths gives a 6-bit kernel nothing to read), engine_widths_test.zig (oQ2's experts, embedding and n-gram tables at their own widths; the uniform checkpoint keeps every recorded kernel; mixed gate/up and n-gram shards refused; the 6/32-only harness refuses others), prompt_w_test.zig, ple_scale.zig's scale against the bf16 product on the host. Gate on main 0efbd447 + the core commits (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test logs an error on the path it expects); native and tf-flashnext-run (with flashnext_bench), -reuse, -pack and -think build. lean_check: 531 problems (main 539); the only findings not on main are the two long files' shorter line counts. --- zig/flashnext_host.zig | 3 + zig/kernels/metal/prefill/qmmw_nax.metal | 340 ++++++++++++++++++ zig/kernels/metal/sources.zig | 2 + zig/src/families/flashnext/dense.zig | 4 +- zig/src/families/flashnext/engine.zig | 8 +- zig/src/families/flashnext/engine_warm.zig | 67 ++++ zig/src/families/flashnext/engine_widths.zig | 96 +++++ .../families/flashnext/engine_widths_test.zig | 92 +++++ zig/src/families/flashnext/lane_w.zig | 201 +++++++++++ zig/src/families/flashnext/lane_w_test.zig | 108 ++++++ zig/src/families/flashnext/ple_scale.zig | 91 +++++ zig/src/families/flashnext/prompt_w.zig | 189 ++++++++++ zig/src/families/flashnext/prompt_w_test.zig | 12 + zig/src/families/flashnext/replay.zig | 143 +++++--- zig/src/families/flashnext/variants.zig | 110 ++++++ zig/tests/flashnext_bench.zig | 23 +- zig/tests/flashnext_run.zig | 4 +- 17 files changed, 1420 insertions(+), 73 deletions(-) create mode 100644 zig/kernels/metal/prefill/qmmw_nax.metal create mode 100644 zig/src/families/flashnext/engine_warm.zig create mode 100644 zig/src/families/flashnext/engine_widths.zig create mode 100644 zig/src/families/flashnext/engine_widths_test.zig create mode 100644 zig/src/families/flashnext/lane_w.zig create mode 100644 zig/src/families/flashnext/lane_w_test.zig create mode 100644 zig/src/families/flashnext/ple_scale.zig create mode 100644 zig/src/families/flashnext/prompt_w.zig create mode 100644 zig/src/families/flashnext/prompt_w_test.zig create mode 100644 zig/src/families/flashnext/variants.zig diff --git a/zig/flashnext_host.zig b/zig/flashnext_host.zig index 06782189f6..01600db8c1 100644 --- a/zig/flashnext_host.zig +++ b/zig/flashnext_host.zig @@ -2,4 +2,7 @@ test { _ = @import("src/families/flashnext/host_test.zig"); _ = @import("src/families/flashnext/pack.zig"); + _ = @import("src/families/flashnext/lane_w.zig"); + _ = @import("src/families/flashnext/engine_widths.zig"); + _ = @import("src/families/flashnext/prompt_w.zig"); } diff --git a/zig/kernels/metal/prefill/qmmw_nax.metal b/zig/kernels/metal/prefill/qmmw_nax.metal new file mode 100644 index 0000000000..0f58a3899e --- /dev/null +++ b/zig/kernels/metal/prefill/qmmw_nax.metal @@ -0,0 +1,340 @@ +// Prompt projections at any MLX width (WB bits, groups of WG, defined first): qmm6_nax's tiles and fp32 sums. +#include +using namespace metal; +typedef bfloat bfloat16_t; +#include "../nax.h" +using namespace tfp; + +namespace tfqw { + +// One thread's 32 codes (WB words, little-endian bit stream) as bf16(s q + b) in out[0..32). +template +inline void dequantw(thread const uint (&w)[WB + 1], T scale, T bias, threadgroup O* out) { + const float s = float(scale), b = float(bias); + TF_UNROLL + for (int i = 0; i < 32; i++) { + const int bit = WB * i, word = bit >> 5, sh = bit & 31; + uint v = w[word] >> sh; + if (sh + WB > 32) v |= w[word + 1] << (32 - sh); + out[i] = O(static_cast(s * float(v & ((1u << WB) - 1u)) + b)); + } +} + +inline void loadw(thread uint (&w)[WB + 1], const device uint* src) { + TF_UNROLL + for (int i = 0; i < WB; i++) w[i] = src[i]; + w[WB] = 0; +} + +// x times a [64 rows, K] block, 64 deep a step: thread t dequantizes row t / 2's codes k0 + k + 32 (t % 2) onward. +template +inline void k_loopw(thread frag (&acc)[TM][2], const device T* x, int K, int ldx, int live, bool inside, + const device uint* wq, const device T* srow, const device T* brow, int k0, threadgroup T* tile, + int tn, uint t, short2 home) { + constexpr int PAD = 64 + 16 / sizeof(T); + threadgroup T* mine = tile + (t / 2) * PAD + 32 * (t % 2); + TF_UNROLL + for (short i = 0; i < TM; i++) { + acc[i][0] = frag(0); + acc[i][1] = frag(0); + } + const int hk = k0 + 32 * int(t % 2); + uint w[WB + 1]; + loadw(w, wq); + T sc = srow[hk / WG], bi = brow[hk / WG]; + for (int k = 0; k < K; k += 64) { + threadgroup_barrier(mem_flags::mem_threadgroup); + dequantw(w, sc, bi, mine); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (k + 64 < K) { + wq += 2 * WB; + loadw(w, wq); + sc = srow[(hk + k + 64) / WG]; + bi = brow[(hk + k + 64) / WG]; + } +#pragma clang loop unroll(disable) + for (int kk = 0; kk < 64; kk += 32) { + if (live > 0) { + frag a[TM][2], b[2][2]; + TF_UNROLL + for (short i = 0; i < 2; i++) { + TF_UNROLL + for (short j = 0; j < 2; j++) { + frag_get_t(b[j][i], (const threadgroup T*)tile, PAD, tn + 16 * i, kk + 16 * j, home); + } + } + TF_UNROLL + for (short i = 0; i < TM; i++) { + TF_UNROLL + for (short j = 0; j < 2; j++) { + if (inside) { + frag_get(a[i][j], x, ldx, 16 * i, kk + 16 * j, home); + } else { + frag_get_in(a[i][j], x, ldx, 16 * i, kk + 16 * j, home, live, kk + 32); + } + } + } + TF_UNROLL + for (short m = 0; m < TM; m++) { + TF_UNROLL + for (short j = 0; j < 2; j++) { + mma_16x32(acc[m][0], acc[m][1], a[m][j], b[j][0], b[j][1]); + } + } + } + } + x += 64; + } + threadgroup_barrier(mem_flags::mem_threadgroup); +} + +// Gate and up at once (k_loopw's sums, twice): both tiles dequantized side by side, each x fragment feeds both. +template +inline void k_loopwx2(thread frag (&ag)[TM][2], thread frag (&au)[TM][2], const device T* x, int K, + int live, bool inside, const device uint* wg, const device T* sg, const device T* bg, + const device uint* wu, const device T* su, const device T* bu, threadgroup T* tg_, + threadgroup T* tu, int tn, uint t, short2 home) { + constexpr int PAD = 64 + 16 / sizeof(T); + threadgroup T* mg = tg_ + (t / 2) * PAD + 32 * (t % 2); + threadgroup T* mu = tu + (t / 2) * PAD + 32 * (t % 2); + TF_UNROLL + for (short i = 0; i < TM; i++) { + ag[i][0] = frag(0); + ag[i][1] = frag(0); + au[i][0] = frag(0); + au[i][1] = frag(0); + } + const int hk = 32 * int(t % 2); + uint gw[WB + 1], uw[WB + 1]; + loadw(gw, wg); + loadw(uw, wu); + T gs = sg[hk / WG], gbi = bg[hk / WG], us = su[hk / WG], ubi = bu[hk / WG]; + for (int k = 0; k < K; k += 64) { + threadgroup_barrier(mem_flags::mem_threadgroup); + dequantw(gw, gs, gbi, mg); + dequantw(uw, us, ubi, mu); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (k + 64 < K) { + wg += 2 * WB; + wu += 2 * WB; + loadw(gw, wg); + loadw(uw, wu); + const int g = (hk + k + 64) / WG; + gs = sg[g]; + gbi = bg[g]; + us = su[g]; + ubi = bu[g]; + } +#pragma clang loop unroll(disable) + for (int kk = 0; kk < 64; kk += 32) { + if (live > 0) { + frag a[TM][2], b[2][2], c[2][2]; + TF_UNROLL + for (short i = 0; i < 2; i++) { + TF_UNROLL + for (short j = 0; j < 2; j++) { + frag_get_t(b[j][i], (const threadgroup T*)tg_, PAD, tn + 16 * i, kk + 16 * j, home); + frag_get_t(c[j][i], (const threadgroup T*)tu, PAD, tn + 16 * i, kk + 16 * j, home); + } + } + TF_UNROLL + for (short i = 0; i < TM; i++) { + TF_UNROLL + for (short j = 0; j < 2; j++) { + if (inside) { + frag_get(a[i][j], x, K, 16 * i, kk + 16 * j, home); + } else { + frag_get_in(a[i][j], x, K, 16 * i, kk + 16 * j, home, live, kk + 32); + } + } + } + TF_UNROLL + for (short m = 0; m < TM; m++) { + TF_UNROLL + for (short j = 0; j < 2; j++) { + mma_16x32(ag[m][0], ag[m][1], a[m][j], b[j][0], b[j][1]); + mma_16x32(au[m][0], au[m][1], a[m][j], c[j][0], c[j][1]); + } + } + } + } + x += 64; + } + threadgroup_barrier(mem_flags::mem_threadgroup); +} + +// The expert activation from the two products, as qmm6_nax's act6 rounds it. +inline float actw(float g, float u) { + const float gb = float(bfloat16_t(g)), ub = float(bfloat16_t(u)); + return float(bfloat16_t(float(bfloat16_t(gb / (1.0f + metal::exp(-gb)))) * ub)); +} + +template +inline void store_act(thread const frag (&ag)[TM][2], thread const frag (&au)[TM][2], device bfloat16_t* y, + int ld, int live, int nc, short2 home) { + TF_UNROLL + for (short i = 0; i < TM; i++) { + TF_UNROLL + for (short j = 0; j < 2; j++) { + frag f; + TF_UNROLL + for (short e = 0; e < 8; e++) f[e] = actw(ag[i][j][e], au[i][j][e]); + if (live == 16 * TM && nc >= 32) { + frag_put(f, y, ld, 16 * i, 16 * j, home); + } else { + frag_put_in(f, y, ld, 16 * i, 16 * j, home, live, nc); + } + } + } +} + +template +inline void store(thread const frag (&acc)[TM][2], device T* y, int ld, int live, int nc, short2 home) { + TF_UNROLL + for (short i = 0; i < TM; i++) { + TF_UNROLL + for (short j = 0; j < 2; j++) { + if (live == 16 * TM && nc >= 32) { + frag_put(acc[i][j], y, ld, 16 * i, 16 * j, home); + } else { + frag_put_in(acc[i][j], y, ld, 16 * i, 16 * j, home, live, nc); + } + } + } +} + +// The tile'th BM-row tile over experts in order: its expert, first row and row count (qmm6_nax's expert_tile). +template +inline bool expert_tile(const device int32_t* offsets, int experts, int total, int tile, uint lane, thread int& expert, + thread int& row, thread int& rows) { + int before = 0; + for (int base = 0; base < experts; base += 32) { + const int e = base + int(lane); + const int first = e < experts ? offsets[e] : total, last = e + 1 < experts ? offsets[e + 1] : total; + const int count = (last - first + BM - 1) / BM; + const int upto = before + simd_prefix_inclusive_sum(count); + const int owner = simd_sum(int(upto <= tile)); + if (owner < 32) { + expert = base + owner; + row = simd_shuffle(first, ushort(owner)) + (tile - simd_shuffle(upto - count, ushort(owner))) * BM; + rows = min(BM, simd_shuffle(last, ushort(owner)) - row); + return true; + } + before = simd_shuffle(upto, ushort(31)); + } + return false; +} + +} // namespace tfqw + +// y = x W^T for W [N, K] at WB/WG: 64 x 64 tiles, 4 simdgroups of 32 x 32. P: K N M, y's row stride (0: N). +[[kernel]] void tf_qmmw_t_nax(const device uint* W [[buffer(0)]], const device bfloat16_t* S [[buffer(1)]], + const device bfloat16_t* B [[buffer(2)]], const device bfloat16_t* X [[buffer(3)]], const device int* P [[buffer(4)]], + device bfloat16_t* Y [[buffer(5)]], uint sg [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint3 tg [[threadgroup_position_in_grid]]) { + threadgroup bfloat16_t tile[64 * (64 + 16 / sizeof(bfloat16_t))]; + const int K = P[0], N = P[1], M = P[2], LD = P[3] > 0 ? P[3] : N; + const int row = int(tg.y) * 64, col = int(tg.x) * 64, t = int(sg) * 32 + int(lane); + const int tm = 32 * int(sg / 2), tn = 32 * int(sg % 2), live = min(32, M - (row + tm)); + const long wrow = long(min(col + t / 2, N - 1)); + const int WPR = K * WB / 32, KG = K / WG; + frag acc[2][2]; + tfqw::k_loopw(acc, X + long(row + tm) * K, K, K, live, live == 32, W + wrow * WPR + WB * (t % 2), + S + wrow * KG, B + wrow * KG, 0, tile, tn, uint(t), frag_home(ushort(lane))); + if (col + tn < N && live > 0) + tfqw::store(acc, Y + long(row + tm) * LD + col + tn, LD, live, N - (col + tn), frag_home(ushort(lane))); +} + +// One K part of x W^T: fp32 partials PART[z][M][N] for part z of P[3] parts. P: K N M parts. +[[kernel]] void tf_qmmw_splitk_nax(const device uint* W [[buffer(0)]], const device bfloat16_t* S [[buffer(1)]], + const device bfloat16_t* B [[buffer(2)]], const device bfloat16_t* X [[buffer(3)]], const device int* P [[buffer(4)]], + device float* PART [[buffer(5)]], uint sg [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint3 tg [[threadgroup_position_in_grid]]) { + threadgroup bfloat16_t tile[64 * (64 + 16 / sizeof(bfloat16_t))]; + const int K = P[0], N = P[1], M = P[2], parts = P[3], KP = K / parts, z = int(tg.z); + const int row = int(tg.y) * 64, col = int(tg.x) * 64, t = int(sg) * 32 + int(lane); + const int tm = 32 * int(sg / 2), tn = 32 * int(sg % 2), live = min(32, M - (row + tm)); + const long wrow = long(min(col + t / 2, N - 1)); + const int WPR = K * WB / 32, KG = K / WG, k0 = z * KP; + frag acc[2][2]; + tfqw::k_loopw(acc, X + long(row + tm) * K + k0, KP, K, live, live == 32, + W + wrow * WPR + (k0 / 32) * WB + WB * (t % 2), S + wrow * KG, B + wrow * KG, k0, tile, tn, + uint(t), frag_home(ushort(lane))); + if (col + tn < N && live > 0) + tfqw::store(acc, PART + (long(z) * M + row + tm) * N + col + tn, N, live, N - (col + tn), frag_home(ushort(lane))); +} + +// y[r] = x[r] W_e^T for rows sorted by expert, W [E, N, K]: BM-row tiles within an expert. P: M N K experts. +template +inline void gatherw(const device bfloat16_t* X, const device uint* W, const device bfloat16_t* S, + const device bfloat16_t* B, const device int32_t* O, device bfloat16_t* Y, const device int* P, + threadgroup bfloat16_t* tile, uint3 tg, uint sg, uint lane) { + constexpr int SM = BM / 2; + const int M = P[0], N = P[1], K = P[2]; + int expert, row, rows; + if (!tfqw::expert_tile(O, P[3], M, int(tg.y), lane, expert, row, rows)) { + return; + } + const int col = int(tg.x) * 64, t = int(sg) * 32 + int(lane); + const int tm = SM * int(sg / 2), tn = 32 * int(sg % 2), live = clamp(rows - tm, 0, SM); + const long wrow = long(expert) * N + min(col + t / 2, N - 1); + const int WPR = K * WB / 32, KG = K / WG; + frag acc[SM / 16][2]; + tfqw::k_loopw(acc, X + long(row + tm) * K, K, K, live, row + tm + SM <= M, + W + wrow * WPR + WB * (t % 2), S + wrow * KG, B + wrow * KG, 0, tile, tn, uint(t), + frag_home(ushort(lane))); + if (col + tn < N && live > 0) + tfqw::store(acc, Y + long(row + tm) * N + col + tn, N, live, N - (col + tn), frag_home(ushort(lane))); +} + +// Experts' gate and up products and activation in one pass, for rows sorted by expert. P: M N K experts. +template +inline void gather_guw(const device bfloat16_t* X, const device uint* WG_, const device bfloat16_t* SG, + const device bfloat16_t* BG, const device uint* WU, const device bfloat16_t* SU, + const device bfloat16_t* BU, const device int32_t* O, device bfloat16_t* Y, const device int* P, + threadgroup bfloat16_t* tile_g, threadgroup bfloat16_t* tile_u, uint3 tg, uint sg, uint lane) { + constexpr int SM = BM / 2; + const int M = P[0], N = P[1], K = P[2]; + int expert, row, rows; + if (!tfqw::expert_tile(O, P[3], M, int(tg.y), lane, expert, row, rows)) { + return; + } + const int col = int(tg.x) * 64, t = int(sg) * 32 + int(lane); + const int tm = SM * int(sg / 2), tn = 32 * int(sg % 2), live = clamp(rows - tm, 0, SM); + const long wrow = long(expert) * N + min(col + t / 2, N - 1); + const int WPR = K * WB / 32, KG = K / WG; + frag ag[SM / 16][2], au[SM / 16][2]; + tfqw::k_loopwx2(ag, au, X + long(row + tm) * K, K, live, row + tm + SM <= M, + WG_ + wrow * WPR + WB * (t % 2), SG + wrow * KG, BG + wrow * KG, + WU + wrow * WPR + WB * (t % 2), SU + wrow * KG, BU + wrow * KG, tile_g, tile_u, tn, + uint(t), frag_home(ushort(lane))); + if (col + tn < N && live > 0) + tfqw::store_act(ag, au, Y + long(row + tm) * N + col + tn, N, live, N - (col + tn), frag_home(ushort(lane))); +} + +#define TF_GATHERW(BM) \ + [[kernel]] void tf_gather_qmmw_nax_##BM(const device bfloat16_t* X [[buffer(0)]], const device uint* W [[buffer(1)]], \ + const device bfloat16_t* S [[buffer(2)]], const device bfloat16_t* B [[buffer(3)]], \ + const device int32_t* O [[buffer(4)]], const device int32_t* P [[buffer(5)]], device bfloat16_t* Y [[buffer(6)]], \ + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], \ + uint3 tg [[threadgroup_position_in_grid]]) { \ + threadgroup bfloat16_t tile[64 * (64 + 16 / sizeof(bfloat16_t))]; \ + gatherw(X, W, S, B, O, Y, P, tile, tg, sg, lane); \ + } +TF_GATHERW(32) +TF_GATHERW(64) +TF_GATHERW(128) + +#define TF_GATHER_GUW(BM) \ + [[kernel]] void tf_gather_guw_nax_##BM(const device bfloat16_t* X [[buffer(0)]], const device uint* WG_ [[buffer(1)]], \ + const device bfloat16_t* SG [[buffer(2)]], const device bfloat16_t* BG [[buffer(3)]], \ + const device uint* WU [[buffer(4)]], const device bfloat16_t* SU [[buffer(5)]], \ + const device bfloat16_t* BU [[buffer(6)]], const device int32_t* O [[buffer(7)]], \ + const device int32_t* P [[buffer(8)]], device bfloat16_t* Y [[buffer(9)]], \ + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], \ + uint3 tg [[threadgroup_position_in_grid]]) { \ + threadgroup bfloat16_t tile_g[64 * (64 + 16 / sizeof(bfloat16_t))], tile_u[64 * (64 + 16 / sizeof(bfloat16_t))]; \ + gather_guw(X, WG_, SG, BG, WU, SU, BU, O, Y, P, tile_g, tile_u, tg, sg, lane); \ + } +TF_GATHER_GUW(32) +TF_GATHER_GUW(64) diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index 6313b656cd..e916b97a2c 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -49,6 +49,8 @@ pub const nax = @embedFile("nax.h"); /// Flash Next 6-bit (group 32) prompt projections on the tensor units: dense and sorted-expert gather. pub const flashnext_qmm6 = @embedFile("prefill/qmm6_nax.metal") ++ @embedFile("prefill/qmm6_nax_b.metal"); +/// Flash Next's prompt products at any MLX width (WB and WG defined before it; nax.h inlined at load). +pub const flashnext_qmmw = @embedFile("prefill/qmmw_nax.metal"); /// Flash Next prompt-chunk glue: hyper-connection pieces, router rows, top-k, the expert sort, gathers and scatters. pub const flashnext_prompt = @embedFile("prefill/fn_prompt.metal"); /// Flash Next block selection in GPU-side rounds: per-row metadata from the arena and pooling at absolute blocks. diff --git a/zig/src/families/flashnext/dense.zig b/zig/src/families/flashnext/dense.zig index 1dda0a2e10..0ca09e6068 100644 --- a/zig/src/families/flashnext/dense.zig +++ b/zig/src/families/flashnext/dense.zig @@ -61,7 +61,7 @@ pub fn recorded(r: *Run, role: []const u8) !struct { n: usize, k: usize, sk: usi } /// y = x W over the layout's tiles, for the rows `mdims` holds. -pub fn project(r: *Run, l: Layout, x: Buf, w: Lane, mdims: Buf, y: Buf) !void { +pub fn project(r: *Run, l: Layout, x: Buf, w: Lane.Six, mdims: Buf, y: Buf) !void { const pipe = try pipeline(r, l); r.enc.setPipeline(pipe); for ([_]Buf{ x, w.wq, w.sbt, mdims, y }, 0..) |b, j| r.enc.setBuffer(b.b, b.off, j); @@ -70,7 +70,7 @@ pub fn project(r: *Run, l: Layout, x: Buf, w: Lane, mdims: Buf, y: Buf) !void { } /// A target lane projection on the core kernel at the recorded kernel's shape and K slices. -pub fn lane(r: *Run, role: []const u8, x: Buf, w: Lane, mdims: Buf, y: Buf) !void { +pub fn lane(r: *Run, role: []const u8, x: Buf, w: Lane.Six, mdims: Buf, y: Buf) !void { if (r.skip & Run.class(role) != 0) return; const s = try recorded(r, role); try project(r, .{ .n = s.n, .k = s.k, .sk = s.sk, .pf = r.lane_pf }, x, w, mdims, y); diff --git a/zig/src/families/flashnext/engine.zig b/zig/src/families/flashnext/engine.zig index e8f86569f1..a2969b5774 100644 --- a/zig/src/families/flashnext/engine.zig +++ b/zig/src/families/flashnext/engine.zig @@ -228,11 +228,6 @@ pub const Engine = struct { var ple_values: config.PleRef = undefined; var attention_scale: f64 = undefined; var cfg = try config.Config.read(arena, io, model_dir); // read once: the reference values and the widths - var cfg_widths = try @import("widths.zig").Widths.init(arena, &cfg); - if (!cfg_widths.uniform() or cfg_widths.q.global.bits != 6 or cfg_widths.q.global.group != 32) { - std.log.warn("flashnext: widths other than 6-bit group 32; the device kernels read 6/32 only", .{}); - return error.UnsupportedFlashAffineKernel; - } if (dump_dir) |d| { const ref_file = try mtl.MappedFile.open(try std.fmt.allocPrintSentinel(arena, "{s}/ref.json", .{d}, 0)); const ref = try std.json.parseFromSliceLeaky(std.json.Value, arena, ref_file.bytes[0..ref_file.size], .{}); @@ -317,6 +312,7 @@ pub const Engine = struct { m.ple.tables[3 * g + 1] = try r.group(16 * g, 16, "scales"); m.ple.tables[3 * g + 2] = try r.group(16 * g, 16, "biases"); } + m.ple.scale = try fz.ple_scale.of(r, index.object.get("weight_map").?.object); const B = struct { fn of(rr: *Run, n: usize) !Buf { return .{ .b = try rr.buffer(n) }; @@ -430,6 +426,7 @@ pub const Engine = struct { e.m = m; e.pr = try arena.create(Prompt); e.pr.* = try Prompt.init(r, pack_dir, r.xnew_header); + const uniform = try @import("engine_widths.zig").apply(arena, &cfg, index.object.get("weight_map").?.object, e.pr, m); e.pr2 = try arena.create(Prompt); e.pr2.* = try e.pr.sibling(); try e.rounds(); @@ -437,6 +434,7 @@ pub const Engine = struct { e.msel = try GSelect.init(r, 0); // speed-up mode: TP=2 with the peer the settings name (tp.zig); both ranks then run the same replies in lockstep const tp_path: ?[]const u8 = speed_up orelse if (std.c.getenv("TF_FLASHNEXT_TP")) |p| std.mem.span(p) else null; + if (tp_path != null and !uniform) return error.SpeedUpNeeds6Bit32; // TP splits the 6/32 kernels only if (tp_path) |path| r.tp = try fz.Tp2.init(arena, r.device, path); return e; } diff --git a/zig/src/families/flashnext/engine_warm.zig b/zig/src/families/flashnext/engine_warm.zig new file mode 100644 index 0000000000..cf7342a562 --- /dev/null +++ b/zig/src/families/flashnext/engine_warm.zig @@ -0,0 +1,67 @@ +//! Every kernel a checkpoint's widths need, compiled at load: a compile failure refuses the load, not a request. +const std = @import("std"); +const replay = @import("replay.zig"); +const variants = @import("variants.zig"); +const lane_w = @import("lane_w.zig"); +const Run = replay.Run; +const Model = replay.Model; + +/// `name` at every recorded row count on its kernel at `consts` (none: the recorded width, compiled already). +fn role(r: *Run, name: []const u8, consts: []const variants.Const) !void { + if (consts.len == 0) return; + var buf: [96]u8 = undefined; + for (1..replay.MAXR + 1) |n| if (r.roles.get(try std.fmt.bufPrint(&buf, "{s}|{d}", .{ name, n }))) |s| { + _ = try variants.patched(r, s.v, consts); + }; +} + +fn hc(r: *Run, h: replay.Hc, comptime at: []const u8) !void { + const w = h.w orelse return; + var c: [2][2]variants.Const = undefined; + try role(r, "qa_hc_down@" ++ at, variants.width(&c[0], "BITS", "GS", w[0])); + try role(r, "qa_hc_up@" ++ at, variants.width(&c[1], "BITS", "GS", w[1])); +} + +fn mtpHc(r: *Run, h: replay.Hc, comptime at: []const u8) !void { + const w = h.w orelse return; + var c: [2][2]variants.Const = undefined; + try role(r, "mtp:qa_hc_down@" ++ at, variants.width(&c[0], "BITS", "GS", w[0])); + try role(r, "mtp:qa_hc_up@" ++ at, variants.width(&c[1], "BITS", "GS", w[1])); +} + +fn experts(r: *Run, xw: ?variants.Experts, gu: []const u8, down: []const u8) !void { + const w = xw orelse return; + var cg: [4]variants.Const = undefined; + var cd: [4]variants.Const = undefined; + try role(r, gu, variants.four(&cg, w.gu)); + try role(r, down, variants.four(&cd, w.down)); +} + +/// A lane's runs at both row tilings (16 and 32 rows a tile). +fn lane(r: *Run, l: replay.Lane) !void { + const w = l.w orelse return; + for (w.runs) |p| for ([_]usize{ 1, 2 }) |tmr| { + _ = try lane_w.pipeline(r.arena, r.device, &r.wide_pipes, .{ .n = p.n, .k = p.k, .bits = p.bits, .group = p.group, .tmr = tmr, .ld = w.n, .mapped = p.map != null, .tiled = p.tiled }); + }; +} + +/// Every decode kernel `m`'s widths launch; none on a 6/32 checkpoint. +pub fn all(r: *Run, m: *const Model) !void { + for (&m.layers) |*L| { + try hc(r, L.ahc, "ahc"); + try hc(r, L.mhc, "mhc"); + try experts(r, L.xw, "qa_expert_gateup@moe.gate", "qa_expert_down_y@moe.down"); + try lane(r, L.proj); + try lane(r, L.out); + } + try hc(r, m.mix, "mix"); + try mtpHc(r, m.mtp.ahc, "mtp.ahc"); + try mtpHc(r, m.mtp.mhc, "mtp.mhc"); + try mtpHc(r, m.mtp.mix, "mtp.mix"); + try experts(r, m.mtp.xw, "mtp:qa_expert_gateup@mtp.moe.gate", "mtp:qa_expert_down_y@mtp.moe.down"); + for ([_]replay.Lane{ m.head, m.ple.kv, m.mtp.proj, m.mtp.out, m.mtp.fce, m.mtp.fch, m.mtp.draft }) |l| try lane(r, l); + var c: [2]variants.Const = undefined; + try role(r, "qa_embed_rows@embed", variants.width(&c, "BITS", "GS", m.embed_w)); + try role(r, "mtp:qa_embed_rows@embed", variants.width(&c, "BITS", "GS", m.embed_w)); + try role(r, "qa_ple_lookup@ple", variants.width(&c, "BITS", "GS", m.ple.w)); +} diff --git a/zig/src/families/flashnext/engine_widths.zig b/zig/src/families/flashnext/engine_widths.zig new file mode 100644 index 0000000000..8162786bf5 --- /dev/null +++ b/zig/src/families/flashnext/engine_widths.zig @@ -0,0 +1,96 @@ +//! The widths of the tensors the engine reads straight from the checkpoint: experts, embedding, n-gram tables. +const std = @import("std"); +const config = @import("config.zig"); +const widths = @import("widths.zig"); +const variants = @import("variants.zig"); +const replay = @import("replay.zig"); + +/// Null wherever a tensor is the recorded 6/32: those keep their recorded kernels. +pub const Set = struct { + layers: [config.max_layers]?variants.Experts = @splat(null), + mtp: ?variants.Experts = null, + embed: ?[2]usize = null, + ple: ?[2]usize = null, + /// Every quantized matrix 6/32: the uniform format, its packs pinned and its fast paths on. + uniform: bool = true, +}; + +fn width(w: *const widths.Widths, path: []const u8) ![2]usize { + const got = (try w.of(path)) orelse return error.UnquantizedAffineTensor; + return .{ got.bits, got.group }; +} + +fn same(x: [2]usize, y: [2]usize) bool { + return x[0] == y[0] and x[1] == y[1]; +} + +/// A layer's experts: gate and up share one width (one kernel reads both), routed and shared each their own. +pub fn experts(w: *const widths.Widths, stem: []const u8) !?variants.Experts { + var buf: [192]u8 = undefined; + var v: [6][2]usize = undefined; + const projs = [_][]const u8{ "switch_mlp.gate_proj", "switch_mlp.up_proj", "shared_expert.gate_proj", "shared_expert.up_proj", "switch_mlp.down_proj", "shared_expert.down_proj" }; + for (projs, &v) |p, *out| out.* = try width(w, try std.fmt.bufPrint(&buf, "{s}.mlp.{s}", .{ stem, p })); + if (!same(v[0], v[1]) or !same(v[2], v[3])) { + std.log.warn("{s}: gate and up of one expert differ in width; one kernel reads both", .{stem}); + return error.ExpertGateUpWidths; + } + const all_six = for (v) |x| { + if (!same(x, .{ 6, 32 })) break false; + } else true; + if (all_six) return null; + return .{ .gu = .{ v[0][0], v[0][1], v[2][0], v[2][1] }, .down = .{ v[4][0], v[4][1], v[5][0], v[5][1] } }; +} + +fn orNull(x: [2]usize) ?[2]usize { + return if (same(x, .{ 6, 32 })) null else x; +} + +fn sixUniform(w: *const widths.Widths) bool { + return w.uniform() and w.q.global.bits == 6 and w.q.global.group == 32; +} + +/// Whether a checkpoint folder is the uniform 6/32 format, which check programs that read only it require. +pub fn uniformDir(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !bool { + var c = try config.Config.read(gpa, io, dir); + defer c.deinit(); + var w = try widths.Widths.init(gpa, &c); + defer w.deinit(); + return sixUniform(&w); +} + +/// Every directly read tensor's width from the checkpoint's config; `spelling` names the n-gram shards. +pub fn read(c: *const config.Config, w: *const widths.Widths, spelling: []const u8) !Set { + var s: Set = .{}; + var buf: [192]u8 = undefined; + for (0..c.layers) |i| s.layers[i] = try experts(w, try std.fmt.bufPrint(&buf, "language_model.model.layers.{d}", .{i})); + if (c.mtp.layers > 0) s.mtp = try experts(w, "language_model.mtp.layers.0"); + s.embed = orNull(try width(w, "language_model.model.embed_tokens")); + for (0..c.layers) |i| if (c.ple[i]) { + const first = try width(w, try std.fmt.bufPrint(&buf, "language_model.model.layers.{d}.ple.ple_embedding.ngram_embedding.{s}0", .{ i, spelling })); + for (1..c.ngram_shards) |k| { + const other = try width(w, try std.fmt.bufPrint(&buf, "language_model.model.layers.{d}.ple.ple_embedding.ngram_embedding.{s}{d}", .{ i, spelling, k })); + if (!same(first, other)) return error.NgramTableWidths; // one lookup kernel reads every shard + } + s.ple = orNull(first); + }; + s.uniform = sixUniform(w); + return s; +} + +test { + _ = @import("engine_widths_test.zig"); +} + +/// Widths from the engine's own config and index into `m` and the prompt, kernels compiled; true when all 6/32. +pub fn apply(a: std.mem.Allocator, c: *const config.Config, map: std.json.ObjectMap, p: *replay.Prompt, m: *replay.Model) !bool { + var w = try widths.Widths.init(a, c); + defer w.deinit(); + const ws = try read(c, &w, @import("index.zig").ngramSpelling(map)); + for (&m.layers, 0..) |*L, i| L.xw = ws.layers[i]; + m.mtp.xw = ws.mtp; + m.embed_w = ws.embed; + m.ple.w = ws.ple; + try replay.prompt_w.attach(p, m); + try @import("engine_warm.zig").all(p.r, m); + return ws.uniform; +} diff --git a/zig/src/families/flashnext/engine_widths_test.zig b/zig/src/families/flashnext/engine_widths_test.zig new file mode 100644 index 0000000000..eb80b6bffb --- /dev/null +++ b/zig/src/families/flashnext/engine_widths_test.zig @@ -0,0 +1,92 @@ +//! The engine's directly read tensors take their declared widths: experts by layer, the embedding, the n-gram tables. +const std = @import("std"); +const config = @import("config.zig"); +const widths = @import("widths.zig"); +const ew = @import("engine_widths.zig"); +const a = std.testing.allocator; +const checked_config = @embedFile("fixtures/config.json"); + +/// The checked config at global `bits`/`group` (both quantization blocks) with `entries` as overrides. +fn with(bits: []const u8, group: []const u8, entries: []const u8) ![]u8 { + const declaration = try std.fmt.allocPrint(a, "\"mode\": \"affine\", {s}", .{entries}); + defer a.free(declaration); + const once = try std.mem.replaceOwned(u8, a, checked_config, "\"mode\": \"affine\"", declaration); + defer a.free(once); + const twice = try std.mem.replaceOwned(u8, a, once, "\"bits\": 6,", bits); + defer a.free(twice); + return std.mem.replaceOwned(u8, a, twice, "\"group_size\": 32,", group); +} + +fn readSet(bytes: []const u8) !ew.Set { + var c = try config.parse(a, bytes); + defer c.deinit(); + var w = try widths.Widths.init(a, &c); + defer w.deinit(); + return ew.read(&c, &w, "shard_"); +} + +test "oQ2's experts, embedding and n-gram tables take their own widths, layer by layer" { + const bytes = try with("\"bits\": 2,", "\"group_size\": 32,", + \\"language_model.model.embed_tokens": {"bits": 8, "group_size": 64}, + \\"language_model.model.layers.3.mlp.shared_expert.gate_proj": {"bits": 8, "group_size": 128}, + \\"language_model.model.layers.3.mlp.shared_expert.up_proj": {"bits": 8, "group_size": 128}, + \\"language_model.model.layers.3.mlp.shared_expert.down_proj": {"bits": 8, "group_size": 128}, + \\"language_model.mtp.layers.0.mlp.switch_mlp.gate_proj": {"bits": 4, "group_size": 32}, + \\"language_model.mtp.layers.0.mlp.switch_mlp.up_proj": {"bits": 4, "group_size": 32} + ); + defer a.free(bytes); + const s = try readSet(bytes); + try std.testing.expect(!s.uniform); + try std.testing.expectEqual([4]usize{ 2, 32, 8, 128 }, s.layers[3].?.gu); + try std.testing.expectEqual([4]usize{ 2, 32, 8, 128 }, s.layers[3].?.down); + try std.testing.expectEqual([4]usize{ 2, 32, 2, 32 }, s.layers[0].?.gu); + try std.testing.expectEqual([4]usize{ 4, 32, 2, 32 }, s.mtp.?.gu); + try std.testing.expectEqual([4]usize{ 2, 32, 2, 32 }, s.mtp.?.down); + try std.testing.expectEqual([2]usize{ 8, 64 }, s.embed.?); + try std.testing.expectEqual([2]usize{ 2, 32 }, s.ple.?); +} + +test "the uniform 6/32 checkpoint keeps every recorded kernel" { + const s = try readSet(checked_config); + try std.testing.expect(s.uniform); + try std.testing.expect(s.embed == null and s.ple == null and s.mtp == null); + for (s.layers) |l| try std.testing.expect(l == null); +} + +test "an expert whose gate and up differ in width is refused: one kernel reads both" { + const bytes = try with("\"bits\": 6,", "\"group_size\": 32,", + \\"language_model.model.layers.0.mlp.switch_mlp.up_proj": {"bits": 4, "group_size": 64} + ); + defer a.free(bytes); + try std.testing.expectError(error.ExpertGateUpWidths, readSet(bytes)); +} + +test "n-gram shards of different widths are refused: one lookup kernel reads every shard" { + const bytes = try with("\"bits\": 6,", "\"group_size\": 32,", + \\"language_model.model.layers.1.ple.ple_embedding.ngram_embedding.shard_1": {"bits": 4, "group_size": 32} + ); + defer a.free(bytes); + try std.testing.expectError(error.NgramTableWidths, readSet(bytes)); +} + +test "a check program that reads only 6/32 refuses a 4/32 checkpoint and one 6/32 checkpoint with an override" { + const io = std.testing.io; + const cases = [_]struct { bits: []const u8, entries: []const u8, six: bool }{ + .{ .bits = "\"bits\": 6,", .entries = "", .six = true }, + .{ .bits = "\"bits\": 4,", .entries = "", .six = false }, + .{ .bits = "\"bits\": 6,", .entries = "\"language_model.model.embed_tokens\": {\"bits\": 8, \"group_size\": 64}", .six = false }, + }; + for (cases) |case| { + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const bytes = if (case.entries.len == 0) + try std.mem.replaceOwned(u8, a, checked_config, "\"bits\": 6,", case.bits) + else + try with(case.bits, "\"group_size\": 32,", case.entries); + defer a.free(bytes); + try tmp.dir.writeFile(io, .{ .sub_path = "config.json", .data = bytes }); + const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(dir); + try std.testing.expectEqual(case.six, try ew.uniformDir(a, io, dir)); + } +} diff --git a/zig/src/families/flashnext/lane_w.zig b/zig/src/families/flashnext/lane_w.zig new file mode 100644 index 0000000000..020cfbeb38 --- /dev/null +++ b/zig/src/families/flashnext/lane_w.zig @@ -0,0 +1,201 @@ +//! Lane products at any MLX width on the recorded lane kernel's text: one launch per bit-width run, mapped columns. +const std = @import("std"); +const mtl = @import("metal"); +const ks = @import("kernel_sources"); +const replay = @import("replay.zig"); +const Run = replay.Run; +const Buf = replay.Buf; +const pack_lane = @import("pack_lane.zig"); +const variants = @import("variants.zig"); + +/// One run's product: n columns over k inputs at one width, and where its columns land in the fused output. +pub const Product = struct { wq: Buf, sbt: Buf, n: usize, k: usize, bits: u8, group: u16, map: ?Buf = null, tiled: bool = true }; + +/// A lane projection read at its own widths: its runs and the fused output's width. +pub const Wide = struct { runs: []const Product, n: usize }; + +/// A pack lane: the recorded 6/32 `.wq` and `.sbt`, or its runs at other widths (`w`); never both. +pub const Lane = struct { + six: ?Six = null, + w: ?*const Wide = null, + + pub const Six = struct { wq: Buf, sbt: Buf }; + + /// The buffers a 6/32 kernel reads; a lane at other widths has none, so no 6-bit kernel reads one of its runs. + pub fn recorded(l: Lane) error{LaneNot6Bit32}!Six { + return l.six orelse error.LaneNot6Bit32; + } +}; + +/// One compiled variant: shape, width, row tiles, output stride, column map, words tiled (else MLX's, decode.py nt 0). +pub const Spec = struct { n: usize, k: usize, bits: u8, group: u16, tmr: usize, ld: usize, mapped: bool, tiled: bool = true }; + +/// Compiled variants by their spec. +pub const Cache = std.AutoHashMapUnmanaged(Spec, mtl.Pipeline); + +/// K slices for an n by k product as lane_qmm.split_k: fixed by the shape, never by the row count. +pub fn splitK(n: usize, k: usize) usize { + const tiles = (n + 31) / 32; + var sk: usize = 1; + while (sk < 8 and tiles * sk < 1024 and (k / 64) / (sk * 2) >= 8) sk *= 2; + return sk; +} + +/// The recorded lane kernel's text for this GPU: the -lanes twin without tensor units. +pub fn template(tensor_units: bool) ![]const u8 { + const name = if (tensor_units) "lane_qmm_bytes_grouped_68a0697c6f38c26d.metal" else "lane_qmm_bytes_grouped_68a0697c6f38c26d-lanes.metal"; + for (ks.flashnext_gen.sources) |s| if (std.mem.eql(u8, s.name, name)) return s.text; + return error.NoSource; +} + +fn swap(a: std.mem.Allocator, text: []const u8, from: []const u8, to: []const u8) ![]const u8 { + if (std.mem.count(u8, text, from) != 1) return error.LanePatch; + return std.mem.replaceOwned(u8, a, text, from, to); +} + +/// The lane kernel at `s` (any width staged as bytes, sums over GS, stores at stride `ld`): the recorded arithmetic. +pub fn source(a: std.mem.Allocator, text_in: []const u8, s: Spec) ![]const u8 { + var text = text_in; + for ([_][]const u8{ "TMR", "N", "K", "SK", "BITS", "GS", "TILED" }, [_]usize{ s.tmr, s.n, s.k, splitK(s.n, s.k), s.bits, s.group, @intFromBool(s.tiled) }) |name, v| + text = try variants.constant(a, text, name, v); + text = try swap(a, text, "static_assert(BITS == 5 || BITS == 6 || BITS == 8, \"bytes for 5-, 6- and 8-bit weights\");", "static_assert(BITS >= 2 && BITS <= 8, \"codes of 2 to 8 bits staged as bytes\");"); + text = try swap(a, text, "for (int i = 0; i < 32; i++) acc += float(X[m * K + g * 32 + i]);", "for (int i = 0; i < FZ_GS; i++) acc += float(X[m * K + g * FZ_GS + i]);"); + const lanes = std.mem.indexOf(u8, text, "Y[m * N + nn + (j & 1) + 8 * (j >> 1)]") != null; + text = if (lanes) + try swap(a, text, "Y[m * N + nn + (j & 1) + 8 * (j >> 1)]", "Y[m * FZ_LD + FZ_COL(nn + (j & 1) + 8 * (j >> 1))]") + else + try swap(a, text, "Y[m * N + nn + j]", "Y[m * FZ_LD + FZ_COL(nn + j)]"); + if (s.mapped) text = try swap(a, text, " device bfloat16_t* Y [[buffer(5)]],\n", " device bfloat16_t* Y [[buffer(5)]],\n const device uint32_t* FZ_MAP [[buffer(6)]],\n"); + const col = if (s.mapped) "(FZ_MAP[c])" else "(c)"; + return std.fmt.allocPrint(a, "#define FZ_GS {d}\n#define FZ_LD {d}\n#define FZ_COL(c) {s}\n{s}", .{ s.group, s.ld, col, text }); +} + +/// The kernel function a source defines. +pub fn functionName(a: std.mem.Allocator, text: []const u8) ![:0]const u8 { + const at = (std.mem.indexOf(u8, text, "[[kernel]] void ") orelse return error.LanePatch) + "[[kernel]] void ".len; + const end = at + (std.mem.indexOfScalar(u8, text[at..], '(') orelse return error.LanePatch); + return std.fmt.allocPrintSentinel(a, "{s}", .{text[at..end]}, 0); +} + +/// The pipeline for `s` on `device`, compiled once and kept in `cache`. +pub fn pipeline(a: std.mem.Allocator, device: mtl.Device, cache: *Cache, s: Spec) !mtl.Pipeline { + if (s.n == 0 or s.n % 4 != 0 or s.tiled != (s.n % 32 == 0)) return error.LaneShape; + if (cache.get(s)) |p| return p; + const text = try source(a, try template(device.tensorUnits()), s); + const lib = try mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); + const pipe = try mtl.Pipeline.init(device, lib, try functionName(a, text), false); + try cache.put(a, s, pipe); + return pipe; +} + +/// Rows tile 16 at a time up to 32 (lane_qmm's ROW_BLOCK): a row's sums never depend on its window. +pub fn rowBlock(rows: usize) usize { + return @min((rows + 15) / 16 * 16, 32); +} + +/// y[rows, w.n] = x[rows, k] W^T, every run at its own width; `mdims` holds rows and rows padded to 16. +pub fn encode(enc: mtl.ComputeEncoder, a: std.mem.Allocator, device: mtl.Device, cache: *Cache, w: *const Wide, x: Buf, mdims: Buf, rows: usize, y: Buf) !void { + const block = rowBlock(rows); + const mp = (rows + 15) / 16 * 16; + for (w.runs) |p| { + const s: Spec = .{ .n = p.n, .k = p.k, .bits = p.bits, .group = p.group, .tmr = block / 16, .ld = w.n, .mapped = p.map != null, .tiled = p.tiled }; + enc.setPipeline(try pipeline(a, device, cache, s)); + for ([_]Buf{ x, x, p.wq, p.sbt, mdims, y }, 0..) |b, j| enc.setBuffer(b.b, b.off, j); + if (p.map) |m| enc.setBuffer(m.b, m.off, 6); + const sk = splitK(p.n, p.k); + enc.dispatchThreads(mtl.Size.of((p.n + 31) / 32 * 32 * sk, (mp + block - 1) / block, 1), mtl.Size.of(32 * sk, 1, 1)); + } +} + +/// `w` on the run's encoder at `rows` rows, unless the role is knocked out. +pub fn launch(r: *Run, w: *const Wide, role: []const u8, x: Buf, mdims: Buf, rows: usize, y: Buf) !void { + if (r.skip & Run.class(role) != 0) return; + try encode(r.enc, r.arena, r.device, &r.wide_pipes, w, x, mdims, rows, y); + if (!r.serial) r.enc.barrier(); +} + +/// Prefetches a lane's weights: the recorded pair, or every run's. +pub fn touch(r: *Run, l: Lane) void { + if (l.six) |s| for ([_]Buf{ s.wq, s.sbt }) |b| r.touch(b); + if (l.w) |w| for (w.runs) |p| for ([_]Buf{ p.wq, p.sbt }) |b| r.touch(b); +} + +/// Run columns to fused columns from a `.members` table of (run, row in run, row in fused output) per member. +pub fn columnMaps(a: std.mem.Allocator, members: []const u32, run_rows: []const usize) ![][]u32 { + const maps = try a.alloc([]u32, run_rows.len); + for (maps, run_rows) |*m, n| { + m.* = try a.alloc(u32, n); + @memset(m.*, std.math.maxInt(u32)); + } + const count = members.len / 3; + for (0..count) |i| { + const run = members[i * 3]; + if (run >= run_rows.len) return error.BadLaneMembers; + const start = members[i * 3 + 1]; + var end: usize = run_rows[run]; // a member runs to the next start in its run + for (0..count) |j| if (members[j * 3] == run and members[j * 3 + 1] > start) { + end = @min(end, members[j * 3 + 1]); + }; + if (start >= end) return error.BadLaneMembers; + for (start..end) |c| maps[run][c] = members[i * 3 + 2] + @as(u32, @intCast(c - start)); + } + var total: usize = 0; + for (run_rows) |n| total += n; + const seen = try a.alloc(bool, total); // every fused column written exactly once + @memset(seen, false); + for (maps) |m| for (m) |c| { + if (c >= total or seen[c]) return error.BadLaneMembers; + seen[c] = true; + }; + return maps; +} + +/// A pack lane's widths: null for a 6/32 lane (no `.runs` table), else its runs, each `.rN` when there are several. +pub fn of(r: *Run, stem: []const u8) !?*const Wide { + const a = r.arena; + const runs_name = try std.fmt.allocPrint(a, "{s}.runs", .{stem}); + if (!r.index.contains(runs_name)) return null; + const t = try variants.table(r, runs_name); + const n_runs = t.len / 3; + if (n_runs == 0 or t.len % 3 != 0) return error.BadLaneRuns; + const run_rows = try a.alloc(usize, n_runs); + for (run_rows, 0..) |*n, i| n.* = t[i * 3 + 2]; + var maps: ?[][]u32 = null; + if (n_runs > 1) { + const members_name = try std.fmt.allocPrint(a, "{s}.members", .{stem}); + maps = try columnMaps(a, try variants.table(r, members_name), run_rows); + } + const runs = try a.alloc(Product, n_runs); + var total: usize = 0; + for (runs, 0..) |*p, i| { + const base = if (n_runs == 1) stem else try std.fmt.allocPrint(a, "{s}.r{d}", .{ stem, i }); + const wq_name = try std.fmt.allocPrint(a, "{s}.wq", .{base}); + const bits: u8 = @intCast(t[i * 3]); + const n = run_rows[i]; + if (n == 0 or bits == 0) return error.BadLaneRuns; + const k = (try r.entry(wq_name)).len * 8 / (n * bits); + p.* = .{ .wq = try r.load(wq_name), .sbt = try r.loadf("{s}.sbt", .{base}), .n = n, .k = k, .bits = bits, .group = @intCast(t[i * 3 + 1]), .tiled = pack_lane.tiled(n) }; + if (maps) |m| { + const b = try r.buffer(n * 4); + @memcpy(b.slice(u32, n), m[i]); + p.map = .{ .b = b }; + } + total += n; + } + const w = try a.create(Wide); + w.* = .{ .runs = runs, .n = total }; + return w; +} + +/// A pack lane: its runs when not 6/32 (`of`), else the recorded `.wq` and `.sbt`. +pub fn laneOf(r: *Run, comptime fmt: []const u8, args: anytype) !Lane { + var name: [96]u8 = undefined; + const stem = try std.fmt.bufPrint(&name, fmt, args); + var full: [128]u8 = undefined; + if (try of(r, stem)) |w| return .{ .w = w }; + return .{ .six = .{ .wq = try r.load(try std.fmt.bufPrint(&full, "{s}.wq", .{stem})), .sbt = try r.load(try std.fmt.bufPrint(&full, "{s}.sbt", .{stem})) } }; +} + +test { + _ = @import("lane_w_test.zig"); +} diff --git a/zig/src/families/flashnext/lane_w_test.zig b/zig/src/families/flashnext/lane_w_test.zig new file mode 100644 index 0000000000..6190a20ee5 --- /dev/null +++ b/zig/src/families/flashnext/lane_w_test.zig @@ -0,0 +1,108 @@ +//! The lane products' widths on the host: K slices as recorded, every recorded text patched, columns put back. +const std = @import("std"); +const ks = @import("kernel_sources"); +const lane_w = @import("lane_w.zig"); +const mtl = @import("metal"); +const replay = @import("replay.zig"); +const variants = @import("variants.zig"); +const pio = @import("pack_io.zig"); + +test "split_k gives every recorded lane kernel's SK from its shape alone" { + const recorded = [_][3]usize{ .{ 13952, 2560, 4 }, .{ 16480, 2560, 2 }, .{ 2560, 6144, 8 }, .{ 248320, 2560, 1 }, .{ 12800, 2560, 4 }, .{ 2560, 2560, 4 }, .{ 79616, 2560, 1 } }; + for (recorded) |c| try std.testing.expectEqual(c[2], lane_w.splitK(c[0], c[1])); +} + +test "every recorded lane text takes any width, a row stride and a column map, and keeps its arithmetic" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var seen: usize = 0; + for (ks.flashnext_gen.sources) |src| { + if (std.mem.indexOf(u8, src.name, "lane_qmm_bytes_grouped") == null) continue; + seen += 1; + for ([_]bool{ false, true }) |mapped| { + const out = try lane_w.source(a, src.text, .{ .n = 512, .k = 6144, .bits = 3, .group = 64, .tmr = 2, .ld = 13952, .mapped = mapped }); + try std.testing.expect(std.mem.indexOf(u8, out, "constexpr int BITS = 3;") != null); + try std.testing.expect(std.mem.indexOf(u8, out, "constexpr int SK = 8;") != null); + try std.testing.expect(std.mem.indexOf(u8, out, "constexpr int TMR = 2;") != null); + try std.testing.expect(std.mem.indexOf(u8, out, "#define FZ_GS 64\n#define FZ_LD 13952\n") != null); + try std.testing.expectEqual(@intFromBool(mapped), std.mem.count(u8, out, "FZ_MAP [[buffer(6)]]")); + try std.testing.expectEqual(@as(usize, 0), std.mem.count(u8, out, "g * 32 + i")); + // the sums: one fma per scale and bias, unchanged + try std.testing.expectEqual(std.mem.count(u8, src.text, "fma("), std.mem.count(u8, out, "fma(")); + } + } + try std.testing.expectEqual(@as(usize, 16), seen); +} + +test "a run's columns land at their members' fused columns, adjacent or not" { + const a = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(a); + defer arena.deinit(); + // q 4 rows and idx 2 rows in run 0, v 3 rows in run 1: fused q | k(=run 0 too, 1 row) | v | idx + const members = [_]u32{ 0, 0, 0, 0, 4, 4, 1, 0, 5, 0, 5, 8 }; + const maps = try lane_w.columnMaps(arena.allocator(), &members, &.{ 7, 3 }); + try std.testing.expectEqualSlices(u32, &.{ 0, 1, 2, 3, 4, 8, 9 }, maps[0]); + try std.testing.expectEqualSlices(u32, &.{ 5, 6, 7 }, maps[1]); + try std.testing.expectError(error.BadLaneMembers, lane_w.columnMaps(arena.allocator(), &.{ 0, 0, 0, 2, 0, 7 }, &.{ 7, 3 })); +} + +/// A pack of a lane at two widths (8/64, 3/64), a 6/32 lane, and a hyper-connection with a widths table. +fn widePack(gpa: std.mem.Allocator, a: std.mem.Allocator, io: std.Io, path: []const u8) !void { + const synthetic = @import("pack_widen_test.zig").synthetic; + var out = pio.Out.init(gpa); + defer out.deinit(); + try @import("pack_lane.zig").putLane(&out, gpa, "L0.att.proj", &.{ try synthetic(a, 1, 64, 256, 8, 64), try synthetic(a, 2, 32, 256, 3, 64) }); + try @import("pack_lane.zig").putLane(&out, gpa, "L0.att.o", &.{try synthetic(a, 3, 32, 256, 6, 32)}); + for ([_][]const u8{ "scale", "down.w", "down.s", "down.b", "up.w", "up.s", "up.b" }) |part| { + const bytes = try gpa.alloc(u8, 64); + @memset(bytes, 0); + try out.put(try std.fmt.allocPrint(a, "L0.ahc.{s}", .{part}), "U8", &.{64}, bytes); + } + const w = try gpa.alloc(u32, 4); + w[0..4].* = .{ 4, 64, 4, 64 }; + try out.put("L0.ahc.widths", "U32", &.{4}, w); + try out.write(gpa, io, path, "test"); +} + +test "a lane's and a hyper-connection's width tables are read on the host: only weights reach the device" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const path = try std.fmt.allocPrintSentinel(a, ".zig-cache/tmp/{s}/pack.safetensors", .{tmp.sub_path}, 0); + try widePack(std.testing.allocator, a, std.testing.io, path); + var r = replay.Run{ .arena = a, .device = device, .queue = try device.queue() }; + try r.indexFile(path); + var weights: usize = 0; + for ([_][]const u8{ "L0.att.proj.r0.wq", "L0.att.proj.r0.sbt", "L0.att.proj.r1.wq", "L0.att.proj.r1.sbt" }) |n| weights += (try r.entry(n)).len; + const w = (try lane_w.of(&r, "L0.att.proj")).?; + try std.testing.expectEqual(@as(usize, 96), w.n); + try std.testing.expectEqual(weights, r.loaded); + const hc = try variants.hcOf(&r, "L0.ahc", .{}); + try std.testing.expectEqual(weights + 7 * 64, r.loaded); + try std.testing.expectEqual([2][2]usize{ .{ 4, 64 }, .{ 4, 64 } }, hc.w.?); +} + +test "a lane at other widths gives a 6-bit kernel nothing to read; a 6/32 lane gives its recorded pair" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const path = try std.fmt.allocPrintSentinel(a, ".zig-cache/tmp/{s}/pack.safetensors", .{tmp.sub_path}, 0); + try widePack(std.testing.allocator, a, std.testing.io, path); + var r = replay.Run{ .arena = a, .device = device, .queue = try device.queue() }; + try r.indexFile(path); + const wide = try lane_w.laneOf(&r, "L{d}.att.proj", .{0}); + try std.testing.expectEqual(@as(usize, 2), wide.w.?.runs.len); + try std.testing.expectError(error.LaneNot6Bit32, wide.recorded()); + const six = try lane_w.laneOf(&r, "L{d}.att.o", .{0}); + try std.testing.expect(six.w == null); + try std.testing.expectEqual((try r.entry("L0.att.o.wq")).len, (try six.recorded()).wq.b.length()); +} diff --git a/zig/src/families/flashnext/ple_scale.zig b/zig/src/families/flashnext/ple_scale.zig new file mode 100644 index 0000000000..3182c23fcd --- /dev/null +++ b/zig/src/families/flashnext/ple_scale.zig @@ -0,0 +1,91 @@ +//! The n-gram table's weight_scale on the looked-up rows, as oMLX's runtime applies it: bf16(row * scale). +const std = @import("std"); +const mtl = @import("metal"); +const replay = @import("replay.zig"); + +const Buf = replay.Buf; + +const source = + \\#include + \\using namespace metal; + \\kernel void tf_ple_scale(device bfloat* x [[buffer(0)]], constant float& s [[buffer(1)]], + \\ constant uint& n [[buffer(2)]], uint i [[thread_position_in_grid]]) { + \\ if (i < n) x[i] = bfloat(float(x[i]) * s); + \\} +; + +/// A scale other than 1 and its kernel, compiled at load. +pub const Scale = struct { + pipe: mtl.Pipeline, + value: f32, + + /// Multiplies `n` bf16 values at `x` in place, then a barrier unless the encoder is serial. + pub fn apply(s: Scale, enc: mtl.ComputeEncoder, x: Buf, n: usize, serial: bool) void { + const count: u32 = @intCast(n); + enc.setPipeline(s.pipe); + enc.setBuffer(x.b, x.off, 0); + enc.setBytes(std.mem.asBytes(&s.value), 1); + enc.setBytes(std.mem.asBytes(&count), 2); + enc.dispatchThreads(mtl.Size.of(n, 1, 1), mtl.Size.of(256, 1, 1)); + if (!serial) enc.barrier(); + } +}; + +/// The checkpoint's n-gram weight_scale as one bf16 or f32 value; null when there is none or it is 1. +pub fn of(r: *replay.Run, map: std.json.ObjectMap) !?Scale { + const name = for (map.keys()) |k| { + if (std.mem.endsWith(u8, k, ".ple.ple_embedding.ngram_embedding.weight_scale")) break k; + } else return null; + const e = try r.entry(name); + const b = try r.load(name); + defer b.b.deinit(); + const bytes = b.b.contents()[0..e.len]; + const value: f32 = switch (e.len) { + 2 => @bitCast(@as(u32, std.mem.readInt(u16, bytes[0..2], .little)) << 16), + 4 => @bitCast(std.mem.readInt(u32, bytes[0..4], .little)), + else => return error.NgramTableScale, + }; + if (!std.math.isFinite(value) or value == 0) return error.NgramTableScale; + if (value == 1) return null; + return try compile(r.device, value); +} + +/// The kernel for `value`; a compile failure refuses the load. +pub fn compile(device: mtl.Device, value: f32) !Scale { + const lib = try mtl.Library.fromSource(device, source, mtl.CompileOptions.mlx()); + return .{ .pipe = try mtl.Pipeline.init(device, lib, "tf_ple_scale", false), .value = value }; +} + +test "the scale rounds each row value once to bf16 after the product, as bf16 row * bf16 scale does" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + var r = replay.Run{ .arena = arena.allocator(), .device = device, .queue = try device.queue() }; + const scale: f32 = @bitCast(@as(u32, 0x3951) << 16); // oQ4e's weight_scale + const s = try compile(device, scale); + const n = 4099; + const buf: Buf = .{ .b = try r.buffer(n * 2) }; + const x = std.mem.bytesAsSlice(u16, buf.b.contents()[0 .. n * 2]); + var rng = std.Random.DefaultPrng.init(7); + for (x) |*v| { // normal bf16 values of both signs whose products stay normal + const exp: u16 = 32 + rng.random().uintLessThan(u16, 96); + v.* = (rng.random().int(u16) & 0x807f) | exp << 7; + } + const want = try std.testing.allocator.alloc(u16, n); + defer std.testing.allocator.free(want); + for (x, want) |v, *w| w.* = bf16(@as(f32, @bitCast(@as(u32, v) << 16)) * scale); + const cb = r.queue.commandBuffer(); + const enc = cb.compute(.serial); + s.apply(enc, buf, n, true); + enc.end(); + cb.commit(); + cb.wait(); + try std.testing.expectEqualSlices(u16, want, x); +} + +/// f32 to bf16, round to nearest even. +fn bf16(f: f32) u16 { + const u: u32 = @bitCast(f); + return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} diff --git a/zig/src/families/flashnext/prompt_w.zig b/zig/src/families/flashnext/prompt_w.zig new file mode 100644 index 0000000000..d23540bc5e --- /dev/null +++ b/zig/src/families/flashnext/prompt_w.zig @@ -0,0 +1,189 @@ +//! Prompt products at any MLX width: qmmw_nax compiled per (bits, group), and which weights take which width. +const std = @import("std"); +const mtl = @import("metal"); +pub const ks = @import("kernel_sources"); +const frags = @import("../../core/frags.zig"); +const replay = @import("replay.zig"); +const Run = replay.Run; +const Buf = replay.Buf; +const Prompt = replay.Prompt; +const Model = replay.Model; +const variants = @import("variants.zig"); + +/// One width's kernels: dense, split-K, sorted-expert gathers (32, 64, 128 rows) and fused gate/up (32, 64 rows). +pub const Pipes = struct { qmm: mtl.Pipeline, splitk: mtl.Pipeline, gather: [3]mtl.Pipeline, gu: [2]mtl.Pipeline }; + +/// A fused member: its rows and width, in member order. +pub const Part = struct { rows: usize, bits: u8, group: u16 }; + +/// A weight at a width other than 6/32: one width, or members of their own widths. +pub const Weight = struct { bits: u8 = 0, group: u16 = 0, parts: []const Part = &.{} }; + +/// A weight's device buffer and offset. +const Key = struct { id: usize, off: usize }; + +/// Every prompt weight by buffer and offset (6/32 too: a missing one is refused, never read at 6 bits) and kernels. +pub const Wide = struct { + weights: std.AutoHashMapUnmanaged(Key, Weight) = .empty, + pipes: std.AutoHashMapUnmanaged(u32, Pipes) = .empty, + missed: bool = false, // a launch met a weight or width the load didn't register: the prompt pass fails + + /// The weight's width when not 6/32 (null: the recorded 6/32 launch); an unregistered weight marks `missed`. + pub fn wide(w: *Wide, b: Buf) ?Weight { + const got = w.weights.get(key(b)) orelse { + w.missed = true; + return null; + }; + return if (got.parts.len == 0 and got.bits == 6 and got.group == 32) null else got; + } + + /// The kernels at a width compiled at load; one never compiled marks `missed` (any width's kernels run meanwhile). + pub fn of(w: *Wide, bits: u8, group: u16) Pipes { + return w.pipes.get(@as(u32, bits) << 16 | group) orelse { + w.missed = true; + var any = w.pipes.valueIterator(); + return any.next().?.*; + }; + } + + /// The pass's refusal when a launch met a weight or width the load didn't register. + pub fn check(w: *const Wide) !void { + if (w.missed) return error.PromptWeightWidthUnknown; + } +}; + +fn key(b: Buf) Key { + return .{ .id = @intFromPtr(b.b.id), .off = b.off }; +} + +/// qmmw_nax at `bits`/`group` with nax.h inlined for this GPU. +pub fn compile(r: *Run, bits: u8, group: u16) !Pipes { + const head = try std.fmt.allocPrint(r.arena, "#define WB {d}\n#define WG {d}\n", .{ bits, group }); + const text = try frags.source(r.device, r.arena, try std.mem.concat(r.arena, u8, &.{ head, ks.flashnext_qmmw })); + const lib = try mtl.Library.fromSource(r.device, text, mtl.CompileOptions.mlx()); + const p = struct { + fn of(rr: *Run, l: mtl.Library, name: [:0]const u8) !mtl.Pipeline { + return mtl.Pipeline.init(rr.device, l, name, false); + } + }; + return .{ + .qmm = try p.of(r, lib, "tf_qmmw_t_nax"), + .splitk = try p.of(r, lib, "tf_qmmw_splitk_nax"), + .gather = .{ try p.of(r, lib, "tf_gather_qmmw_nax_32"), try p.of(r, lib, "tf_gather_qmmw_nax_64"), try p.of(r, lib, "tf_gather_qmmw_nax_128") }, + .gu = .{ try p.of(r, lib, "tf_gather_guw_nax_32"), try p.of(r, lib, "tf_gather_guw_nax_64") }, + }; +} + +fn need(w: *Wide, r: *Run, bits: u8, group: u16) !void { + const k = @as(u32, bits) << 16 | group; + if (!w.pipes.contains(k)) try w.pipes.put(r.arena, k, try compile(r, bits, group)); +} + +fn put(w: *Wide, r: *Run, b: Buf, weight: Weight) !void { + try w.weights.put(r.arena, key(b), weight); + if (weight.parts.len == 0 and weight.bits == 6 and weight.group == 32) return; + if (weight.parts.len == 0) return need(w, r, weight.bits, weight.group); + for (weight.parts) |p| try need(w, r, p.bits, p.group); +} + +fn width(v: [2]usize) Weight { + return .{ .bits = @intCast(v[0]), .group = @intCast(v[1]) }; +} + +/// A pack_mlx projection's `.parts` table, when it has one. +fn parts(r: *Run, stem: []const u8) !?[]const Part { + const name = try std.fmt.allocPrint(r.arena, "{s}.parts", .{stem}); + if (!r.index.contains(name)) return null; + const t = try variants.table(r, name); + const n = t.len / 3; + const out = try r.arena.alloc(Part, n); + for (out, 0..) |*p, i| p.* = .{ .rows = t[i * 3], .bits = @intCast(t[i * 3 + 1]), .group = @intCast(t[i * 3 + 2]) }; + return out; +} + +const six: Weight = .{ .bits = 6, .group = 32 }; + +fn hc(w: *Wide, r: *Run, h: replay.Hc) !void { + try put(w, r, h.uw, if (h.w) |hw| width(hw[1]) else six); + if (h.w) |hw| try need(w, r, @intCast(hw[0][0]), @intCast(hw[0][1])); // the down projection's split-K +} + +/// The prompt's weights at their widths: null on a 6/32 checkpoint, which keeps every recorded launch. +pub fn attach(p: *Prompt, m: *const Model) !void { + const r = p.r; + const w = try r.arena.create(Wide); + w.* = .{}; + for (&m.layers, 0..) |*L, i| { + try hc(w, r, L.ahc); + try hc(w, r, L.mhc); + const x = L.xw orelse variants.Experts{ .gu = .{ 6, 32, 6, 32 }, .down = .{ 6, 32, 6, 32 } }; + for ([_]usize{ 0, 3, 6, 9, 12, 15 }, [_][2]usize{ .{ x.gu[0], x.gu[1] }, .{ x.gu[0], x.gu[1] }, .{ x.gu[2], x.gu[3] }, .{ x.gu[2], x.gu[3] }, .{ x.down[0], x.down[1] }, .{ x.down[2], x.down[3] } }) |j, v| + try put(w, r, L.ex[j], width(v)); // routed gate and up (one fused pass or two), shared, then the downs + const kind = if (L.linear) "gdn" else "att"; + for ([_][]const u8{ if (L.linear) "in" else "proj", if (L.linear) "out" else "o" }, [_]Buf{ p.proj[i][0], p.out[i][0] }) |n, b| { + try put(w, r, b, if (try parts(r, try std.fmt.allocPrint(r.arena, "L{d}.{s}.{s}", .{ i, kind, n }))) |ps| uniformOr(ps) else six); + } + } + try hc(w, r, m.mix); + try hc(w, r, m.mtp.ahc); + try put(w, r, p.ple_kv[0], if (try parts(r, "ple.kv")) |ps| uniformOr(ps) else six); + if (p.mtp_w) |mw| for ([_][]const u8{ "mtp.att.proj", "mtp.fce", "mtp.fch" }, 0..) |n, j| { + try put(w, r, mw[j][0], if (try parts(r, n)) |ps| uniformOr(ps) else six); + }; + p.wide = if (w.pipes.count() == 0) null else w; // all 6/32: every recorded launch as before +} + +/// Members of one width are one product at that width; several widths stay members. +fn uniformOr(ps: []const Part) Weight { + for (ps[1..]) |x| if (x.bits != ps[0].bits or x.group != ps[0].group) return .{ .parts = ps }; + return .{ .bits = ps[0].bits, .group = ps[0].group }; +} + +/// y[rows, n] (row stride ldy, 0: n) = x[rows, k] W^T at the weight's width(s): one launch per member when several. +pub fn qmm(p: *Prompt, wd: *Wide, w: Weight, x: Buf, ws: [3]Buf, k: usize, n: usize, rows: usize, y: Buf, ldy: usize) void { + const ld = if (ldy == 0) n else ldy; + if (w.parts.len == 0) return launch(p, wd.of(w.bits, w.group).qmm, ws, x, k, n, rows, y, ld); + var at = [3]usize{ ws[0].off, ws[1].off, ws[2].off }; + var col: usize = 0; + for (w.parts) |part| { + const m: [3]Buf = .{ .{ .b = ws[0].b, .off = at[0] }, .{ .b = ws[1].b, .off = at[1] }, .{ .b = ws[2].b, .off = at[2] } }; + launch(p, wd.of(part.bits, part.group).qmm, m, x, k, part.rows, rows, y.at(col * 2), ld); + at[0] += part.rows * k * part.bits / 8; + at[1] += part.rows * (k / part.group) * 2; + at[2] += part.rows * (k / part.group) * 2; + col += part.rows; + } +} + +fn launch(p: *Prompt, pipe: mtl.Pipeline, w: [3]Buf, x: Buf, k: usize, n: usize, rows: usize, y: Buf, ld: usize) void { + p.bind(pipe, &.{ w[0], w[1], w[2], x }); + const prm = [4]i32{ @intCast(k), @intCast(n), @intCast(rows), @intCast(ld) }; + p.r.enc.setBytes(std.mem.asBytes(&prm), 4); + p.r.enc.setBuffer(y.b, y.off, 5); + p.r.enc.dispatchThreads(mtl.Size.of(((n + 63) / 64) * 128, (rows + 63) / 64, 1), mtl.Size.of(128, 1, 1)); + p.barrier(); +} + +/// The sorted-expert gather at the weight's width: tiles of 32, 64 or 128 rows as Prompt.gather picks them. +pub fn gather(p: *Prompt, pipes: Pipes, bm: usize, x: Buf, w: []const Buf, k: usize, n: usize, pairs: usize, y: Buf) void { + p.bind(pipes.gather[if (bm == 128) 2 else if (bm == 64) 1 else 0], &.{ x, w[0], w[1], w[2], p.b.off }); + const prm = [4]i32{ @intCast(pairs), @intCast(n), @intCast(k), 512 }; + p.r.enc.setBytes(std.mem.asBytes(&prm), 5); + p.r.enc.setBuffer(y.b, y.off, 6); + p.r.enc.dispatchThreads(mtl.Size.of(((n + 63) / 64) * 128, pairs / bm + 512, 1), mtl.Size.of(128, 1, 1)); + p.barrier(); +} + +/// The fused gate/up gather at the weights' width. +pub fn gatherGU(p: *Prompt, pipes: Pipes, bm: usize, x: Buf, wg: []const Buf, wu: []const Buf, k: usize, n: usize, pairs: usize, y: Buf) void { + p.bind(pipes.gu[if (bm == 64) 1 else 0], &.{ x, wg[0], wg[1], wg[2], wu[0], wu[1], wu[2], p.b.off }); + const prm = [4]i32{ @intCast(pairs), @intCast(n), @intCast(k), 512 }; + p.r.enc.setBytes(std.mem.asBytes(&prm), 8); + p.r.enc.setBuffer(y.b, y.off, 9); + p.r.enc.dispatchThreads(mtl.Size.of(((n + 63) / 64) * 128, pairs / bm + 512, 1), mtl.Size.of(128, 1, 1)); + p.barrier(); +} + +test { + _ = @import("prompt_w_test.zig"); +} diff --git a/zig/src/families/flashnext/prompt_w_test.zig b/zig/src/families/flashnext/prompt_w_test.zig new file mode 100644 index 0000000000..f58a484eed --- /dev/null +++ b/zig/src/families/flashnext/prompt_w_test.zig @@ -0,0 +1,12 @@ +//! The prompt kernels at any width compile from one text, and fused members' offsets follow their own widths. +const std = @import("std"); +const ks = @import("kernel_sources"); + +test "the any-width prompt text names its kernels and declares no width of its own" { + const text = ks.flashnext_qmmw; + for ([_][]const u8{ "tf_qmmw_t_nax", "tf_qmmw_splitk_nax", "tf_gather_guw_nax_##BM", "tf_gather_qmmw_nax_##BM" }) |name| + try std.testing.expect(std.mem.indexOf(u8, text, name) != null); + try std.testing.expect(std.mem.indexOf(u8, text, "#define WB") == null); + try std.testing.expect(std.mem.indexOf(u8, text, "#define WG") == null); + try std.testing.expectEqual(@as(usize, 1), std.mem.count(u8, text, "#include \"../nax.h\"")); +} diff --git a/zig/src/families/flashnext/replay.zig b/zig/src/families/flashnext/replay.zig index ad35200ccc..cf54294663 100644 --- a/zig/src/families/flashnext/replay.zig +++ b/zig/src/families/flashnext/replay.zig @@ -10,6 +10,11 @@ const mark_slots = @import("marks.zig"); const verdict_kernels = @import("verdict.zig"); pub const split = @import("split.zig"); pub const dense = @import("dense.zig"); +pub const lane_w = @import("lane_w.zig"); +pub const ple_scale = @import("ple_scale.zig"); +pub const variants = @import("variants.zig"); +pub const prompt_w = @import("prompt_w.zig"); +pub const engine_widths = @import("engine_widths.zig"); pub const gdn_step = @import("gdn.zig"); pub const Tp2 = tpm.Tp2; pub const frags = @import("../../core/frags.zig"); @@ -887,6 +892,7 @@ pub const Run = struct { lane_new: bool = false, // the target's lane projections on core lane (dense.zig): the recorded bits, weights read ahead lane_pf: usize = 1, // groups a simdgroup reads ahead in core lane lane_pipes: std.AutoHashMapUnmanaged(u64, mtl.Pipeline) = .empty, + wide_pipes: lane_w.Cache = .empty, // lane products at other widths than 6/32 (lane_w.zig) lane_shape: std.AutoHashMapUnmanaged(*Variant, [3]usize) = .empty, // recorded lane kernels' N, K, SK gdn_pipe: ?mtl.Pipeline = null, // the DeltaNet window step on fz_gdn (gdn.zig) when set gdn_kept: ?mtl.Pipeline = null, // GPU-side rounds: fz_gdn keeping one state a layer, kept rows replayed (gdn.zig) @@ -1118,9 +1124,13 @@ pub const Run = struct { /// A hyper-connection's down then up projection at `as_rows` from normed streams `hn` (sums of squares in `ssp`): /// the mixed rows and inject gates; with `hc_up` the up projection of 3-8 rows on fz_hc_up (the same bits). pub fn hcProject(r: *Run, down: []const u8, up: []const u8, as_rows: usize, hn: Buf, ssp: Buf, hc: Hc, eps: Buf, rows: Buf, part: Buf, mixed: Buf, inj: Buf) !void { - try r.callAs(down, as_rows, &.{ hn, ssp, hc.scale, hc.dw, hc.ds, hc.db, eps, rows }, &.{part}); + var cd: [2]variants.Const = undefined; + var cu: [2]variants.Const = undefined; + const kd = variants.width(&cd, "BITS", "GS", if (hc.w) |w| w[0] else null); + const ku = variants.width(&cu, "BITS", "GS", if (hc.w) |w| w[1] else null); + try variants.callAs(r, down, as_rows, &.{ hn, ssp, hc.scale, hc.dw, hc.ds, hc.db, eps, rows }, &.{part}, kd); const ins = [_]Buf{ hn, ssp, hc.scale, part, hc.uw, hc.us, hc.ub, eps, rows }; - if (!r.hc_up or as_rows < 3 or as_rows > 8) return r.callAs(up, as_rows, &ins, &.{ mixed, inj }); + if (!r.hc_up or hc.w != null or as_rows < 3 or as_rows > 8) return variants.callAs(r, up, as_rows, &ins, &.{ mixed, inj }, ku); if (r.skip & class(up) != 0 or r.hcskip & 4 != 0) return; const nd: u32 = @intCast(hc.dw.b.length() / (WIDE * 6 / 8)); r.enc.setPipeline(r.hc_up_pipe); @@ -1195,7 +1205,7 @@ pub const Run = struct { } /// y = x W for `rows` rows of K inputs through fz_dense (blocks of 8 rows; 16 K-slices for narrow outputs). /// One row's dense projection over output tiles [t0, t0 + tn) only (32 outputs a tile), in the full layout. - pub fn denseTiles(r: *Run, x: Buf, k: usize, l: Lane, y: Buf, t0: usize, tn: usize) void { + pub fn denseTiles(r: *Run, x: Buf, k: usize, l: Lane.Six, y: Buf, t0: usize, tn: usize) void { if (tn == 0) return; const n = l.wq.b.length() * 4 / (3 * k); const narrow = n <= 4096; @@ -1211,7 +1221,7 @@ pub const Run = struct { if (!r.serial) r.enc.barrier(); } - pub fn denseRows(r: *Run, x: Buf, k: usize, l: Lane, rows: usize, y: Buf) void { + pub fn denseRows(r: *Run, x: Buf, k: usize, l: Lane.Six, rows: usize, y: Buf) void { const n = l.wq.b.length() * 4 / (3 * k); const narrow = n <= 4096; var at: usize = 0; @@ -1597,8 +1607,9 @@ pub const GSelect = struct { } }; -pub const Hc = struct { scale: Buf, dw: Buf, ds: Buf, db: Buf, uw: Buf, us: Buf, ub: Buf }; -pub const Lane = struct { wq: Buf, sbt: Buf }; +/// A pack hyper-connection; `w`: the down and up widths when not 6/32. +pub const Hc = struct { scale: Buf, dw: Buf, ds: Buf, db: Buf, uw: Buf, us: Buf, ub: Buf, w: ?[2][2]usize = null }; +pub const Lane = lane_w.Lane; // the recorded 6/32 pair, or runs at other widths pub const Layer = struct { ahc: Hc, mhc: Hc, @@ -1614,6 +1625,7 @@ pub const Layer = struct { iqn: Buf = undefined, router: Buf, ex: [18]Buf, // gate, up, shared gate, shared up (weight, scales, biases), then down, shared down + xw: ?variants.Experts = null, // the experts' widths when not 6/32 (engine_widths.zig) cs: [2]Buf = undefined, so: [2]Buf = undefined, keys: Buf = undefined, @@ -1624,22 +1636,8 @@ pub const Layer = struct { pooled_n: usize = 0, }; -pub fn hcOf(r: *Run, comptime fmt: []const u8, args: anytype) !Hc { - var name: [96]u8 = undefined; - const stem = try std.fmt.bufPrint(&name, fmt, args); - var full: [128]u8 = undefined; - const parts = [_][]const u8{ "scale", "down.w", "down.s", "down.b", "up.w", "up.s", "up.b" }; - var out: [7]Buf = undefined; - for (parts, 0..) |p, i| out[i] = try r.load(try std.fmt.bufPrint(&full, "{s}.{s}", .{ stem, p })); - return .{ .scale = out[0], .dw = out[1], .ds = out[2], .db = out[3], .uw = out[4], .us = out[5], .ub = out[6] }; -} - -pub fn laneOf(r: *Run, comptime fmt: []const u8, args: anytype) !Lane { - var name: [96]u8 = undefined; - const stem = try std.fmt.bufPrint(&name, fmt, args); - var full: [128]u8 = undefined; - return .{ .wq = try r.load(try std.fmt.bufPrint(&full, "{s}.wq", .{stem})), .sbt = try r.load(try std.fmt.bufPrint(&full, "{s}.sbt", .{stem})) }; -} +pub const hcOf = variants.hcOf; +pub const laneOf = lane_w.laneOf; pub const Ple = struct { kv: Lane, @@ -1655,6 +1653,8 @@ pub const Ple = struct { mult: [3]i64, sizes: [16]i64, offsets: [16]i64, + w: ?[2]usize = null, // the n-gram tables' width when not 6/32 + scale: ?ple_scale.Scale = null, // the n-gram tables' weight_scale when not 1 }; /// Every intermediate a window of up to MAXR rows writes, and the per-window values the kernels read. @@ -1751,6 +1751,7 @@ pub const Mtp = struct { ids: Buf, ids_n: usize, ex: [18]Buf, + xw: ?variants.Experts = null, keys: Buf, vals: Buf, raw: Buf, @@ -1785,6 +1786,7 @@ pub const Model = struct { mix: Hc, head: Lane, embed: [3]Buf, + embed_w: ?[2]usize = null, // the embedding's width when not 6/32 ple: Ple, t: Tmp, pos: usize = 0, @@ -1818,10 +1820,12 @@ pub const Model = struct { } pub fn lane(m: *Model, x: Buf, k: usize, l: Lane, role: []const u8, y: Buf) !void { - if (m.r.dense_target) return m.r.denseRows(x, k, l, m.r.rows, y); - if (m.r.lane_new) return dense.lane(m.r, role, x, l, m.t.mdims, y); + if (l.w) |w| return lane_w.launch(m.r, w, role, x, m.t.mdims, m.r.rows, y); + const six = try l.recorded(); + if (m.r.dense_target) return m.r.denseRows(x, k, six, m.r.rows, y); + if (m.r.lane_new) return dense.lane(m.r, role, x, six, m.t.mdims, y); if (!m.r.fused_xsum) try m.r.call(if (k == D) "lane_qmm_xsum#[2560]" else "lane_qmm_xsum#[6144]", &.{ x, m.t.mdims }, &.{m.t.xs}); - try m.r.call(role, &.{ x, m.t.xs, l.wq, l.sbt, m.t.mdims }, &.{y}); + try m.r.call(role, &.{ x, m.t.xs, six.wq, six.sbt, m.t.mdims }, &.{y}); } pub fn hcProject(m: *Model, hn: Buf, hc: Hc, down: []const u8, up: []const u8, inj: Buf) !void { @@ -1946,7 +1950,8 @@ pub const Model = struct { if (!r.serial) r.enc.barrier(); } var cur: usize = 0; - try r.call("qa_embed_rows@embed", &.{ ids, m.embed[0], m.embed[1], m.embed[2] }, &.{t.h[0]}); + var ce: [2]variants.Const = undefined; + try variants.callAs(r, "qa_embed_rows@embed", rows, &.{ ids, m.embed[0], m.embed[1], m.embed[2] }, &.{t.h[0]}, variants.width(&ce, "BITS", "GS", m.embed_w)); var pending: enum { none, grouped } = .none; for (0..LAYERS) |i| { const L = &m.layers[i]; @@ -1959,7 +1964,9 @@ pub const Model = struct { tabs[0] = t.ple_ids; tabs[1] = p.starts; for (0..3 * GROUPS) |j| tabs[2 + j] = p.tables[j]; - try r.call("qa_ple_lookup@ple", &tabs, &.{t.emb}); + var cp: [2]variants.Const = undefined; + try variants.callAs(r, "qa_ple_lookup@ple", rows, &tabs, &.{t.emb}, variants.width(&cp, "BITS", "GS", p.w)); + if (p.scale) |sc| sc.apply(r.enc, t.emb, rows * D, r.serial); try m.lane(t.emb, D, p.kv, "lane_qmm_bytes_grouped@ple.kv", t.kvp); try r.call("q4_ple_gate@ple", &.{ t.kvp, t.h[cur], p.ks, p.qs, p.cs, t.eps }, &.{ t.gated, .{ .b = p.cin.b, .off = PLE_TAIL * WIDE * 2 } }); try r.call("q4_ple_conv@ple", &.{ p.cin, p.conv, t.gated, t.h[cur] }, &.{t.hout}); @@ -1978,7 +1985,8 @@ pub const Model = struct { if (!r.fused_xsum) try r.call("lane_qmm_xsum#[2560]", &.{ t.mixed, t.mdims }, &.{t.xs}); // [q 64 tiles | k 64 | v 192 | z 192 | b, a 3]: this Mac's q, k, v and z heads, b and a whole, 8 K slices const cols = [_][2]usize{ .{ 32 * k0, 32 }, .{ 64 + 32 * k0, 32 }, .{ 128 + 96 * k0, 96 }, .{ 320 + 96 * k0, 96 }, .{ 512, 3 } }; - try split.laneTiles(r, "lane_qmm_bytes_grouped@gdn.in", &.{ t.mixed, t.xs, L.proj.wq, L.proj.sbt, t.mdims }, t.p, &cols, 8, null); + const pj = try L.proj.recorded(); + try split.laneTiles(r, "lane_qmm_bytes_grouped@gdn.in", &.{ t.mixed, t.xs, pj.wq, pj.sbt, t.mdims }, t.p, &cols, 8, null); const a = m.state; const cs_in: Buf = .{ .b = L.cs[a].b, .off = L.cs[a].off + m.state_row * CS_ROW }; const so_in: Buf = .{ .b = L.so[a].b, .off = L.so[a].off + m.state_row * SO_ROW }; @@ -1988,7 +1996,8 @@ pub const Model = struct { } else try split.gdnHeads(r, "q4_gdn@gdn", if (r.gdn_step and rows > 1) 8 else rows, &.{ t.p, cs_in, so_in, L.conv, L.alog, L.dt, L.norm, t.eps, t.rows }, &.{ t.gout, L.cs[1 - a], L.so[1 - a] }, tp.rank); const pn = tp.partNext(); const part: Buf = .{ .b = pn.b, .off = pn.off }; - try split.laneTiles(r, "lane_qmm_bytes_grouped@gdn.out", &.{ t.gout, t.xs, L.out.wq, L.out.sbt, t.mdims }, part, &.{.{ 0, 80 }}, 8, .{ 96 * k0, 96 }); // one kernel at every width: drafted == plain + const po = try L.out.recorded(); + try split.laneTiles(r, "lane_qmm_bytes_grouped@gdn.out", &.{ t.gout, t.xs, po.wq, po.sbt, t.mdims }, part, &.{.{ 0, 80 }}, 8, .{ 96 * k0, 96 }); // one kernel at every width: drafted == plain if (r.skip & TP_CLASS == 0) { // the partials' sum and the stream update in the exchange's launch tp.plain(r.enc, t.h[cur], t.inj_a, t.h[1 - cur], t.ssp, t.rows, rows); plain = true; @@ -1998,8 +2007,7 @@ pub const Model = struct { const a = m.state; const cs_in: Buf = .{ .b = L.cs[a].b, .off = L.cs[a].off + m.state_row * CS_ROW }; const so_in: Buf = .{ .b = L.so[a].b, .off = L.so[a].off + m.state_row * SO_ROW }; - r.touch(L.out.wq); - r.touch(L.out.sbt); + lane_w.touch(r, L.out); const gins = [_]Buf{ t.p, cs_in, so_in, L.conv, L.alog, L.dt, L.norm, t.eps, t.rows }; const gouts = [_]Buf{ t.gout, L.cs[1 - a], L.so[1 - a] }; if (m.keptState()) |recs| { @@ -2009,8 +2017,7 @@ pub const Model = struct { try m.lane(t.gout, 6144, L.out, "lane_qmm_bytes_grouped@gdn.out", t.branch); } else { try m.lane(t.mixed, D, L.proj, "lane_qmm_bytes_grouped@att.proj", t.p); - r.touch(L.out.wq); - r.touch(L.out.sbt); + lane_w.touch(r, L.out); try r.call("q4_attn_prep@att", &.{ t.p, t.pos8, L.qn, L.kn, L.iqn, t.eps, t.log2base }, &.{ t.q, t.kout, t.iq }); if (r.skip & (1 << 5) == 0) { r.enc.setPipeline(r.kv_pipe); @@ -2055,7 +2062,7 @@ pub const Model = struct { try m.hcProject(t.h[cur], L.mhc, "qa_hc_down@mhc", "qa_hc_up@mhc", t.inj_m); try r.router("q4_router_float@moe", t.mixed, L.router, t.rows, t.lg, rows); r.tp_layer = true; - try r.experts("qa_expert_gateup@moe.gate", "qa_expert_down_y@moe.down", t.mixed, t.lg, L.ex, t.act, t.pick, t.wts, t.rows, t.ydown); + if (L.xw) |w| try variants.experts(r, "qa_expert_gateup@moe.gate", "qa_expert_down_y@moe.down", t.mixed, t.lg, L.ex, t.act, t.pick, t.wts, t.rows, t.ydown, w) else try r.experts("qa_expert_gateup@moe.gate", "qa_expert_down_y@moe.down", t.mixed, t.lg, L.ex, t.act, t.pick, t.wts, t.rows, t.ydown); r.tp_layer = false; if (r.tp) |tp| if (r.skip & TP_CLASS == 0) tp.exchange(r.enc, t.ydown, t.wts, t.lg, rows); if (r.probe) |pb| r.copyKept(t.pick, .{ .b = pb.b, .off = pb.off + i * MAXR * 10 * 4 }, rows * 10, 0, 0, 0, 1, -1); @@ -2069,7 +2076,8 @@ pub const Model = struct { if (r.tp) |tp| { // TP: this Mac's half of the vocabulary (the same kernel, half its tiles), then one swap const half = HEAD_TILES / 2; if (!r.fused_xsum) try r.call("lane_qmm_xsum#[2560]", &.{ t.mixed, t.mdims }, &.{t.xs}); - try split.laneTiles(r, "lane_qmm_bytes_grouped@head", &.{ t.mixed, t.xs, m.head.wq, m.head.sbt, t.mdims }, t.logits, &.{.{ half * tp.rank, half }}, 1, null); + const hd = try m.head.recorded(); + try split.laneTiles(r, "lane_qmm_bytes_grouped@head", &.{ t.mixed, t.xs, hd.wq, hd.sbt, t.mdims }, t.logits, &.{.{ half * tp.rank, half }}, 1, null); if (r.skip & TP_CLASS == 0) tp.argmax(r.enc, t.logits, HEAD_TILES * 32, half * 32 * tp.rank, half * 32, t.picks, t.rows, rows); return; } @@ -2133,13 +2141,14 @@ pub const Model = struct { const t = &m.t; const h = &m.mtp; r.rows = rows; - try r.call("mtp:qa_embed_rows@embed", &.{ ids, m.embed[0], m.embed[1], m.embed[2] }, &.{h.emb}); + var ce: [2]variants.Const = undefined; + try variants.callAs(r, "mtp:qa_embed_rows@embed", rows, &.{ ids, m.embed[0], m.embed[1], m.embed[2] }, &.{h.emb}, variants.width(&ce, "BITS", "GS", m.embed_w)); try r.call("mtp:q4_rms_rows@mtp.enorm", &.{ h.emb, h.enorm, t.eps }, &.{h.en}); if (!r.fused_xsum) try r.call("mtp:lane_qmm_xsum#[2560]", &.{ h.en, sl.md }, &.{t.xs}); - if (r.dense) r.denseRows(h.en, D, h.fce, rows, h.e) else try r.call("mtp:lane_qmm_bytes_grouped@mtp.fce", &.{ h.en, t.xs, h.fce.wq, h.fce.sbt, sl.md }, &.{h.e}); + try m.mtpLane(h.fce, "mtp:lane_qmm_bytes_grouped@mtp.fce", h.en, sl.md, rows, D, h.e); try r.call("mtp:q4_rms_rows@mtp.hnorm", &.{ streams, h.hnorm, t.eps }, &.{h.hn}); if (!r.fused_xsum) try r.call("mtp:lane_qmm_xsum#[4R, 2560]", &.{ h.hn, sl.md4 }, &.{t.xs}); - if (r.dense) r.denseRows(h.hn, D, h.fch, 4 * rows, h.hs) else try r.call("mtp:lane_qmm_bytes_grouped@mtp.fch", &.{ h.hn, t.xs, h.fch.wq, h.fch.sbt, sl.md4 }, &.{h.hs}); + try m.mtpLane(h.fch, "mtp:lane_qmm_bytes_grouped@mtp.fch", h.hn, sl.md4, 4 * rows, D, h.hs); r.enc.setPipeline(r.add_pipe); for ([_]Buf{ h.e, h.hs, h.h[0], sl.n_add }, 0..) |b, j| r.enc.setBuffer(b.b, b.off, j); r.enc.dispatchThreads(mtl.Size.of(rows * WIDE, 1, 1), mtl.Size.of(256, 1, 1)); @@ -2149,7 +2158,7 @@ pub const Model = struct { const up = [_][]const u8{ "mtp:qa_hc_up@mtp.ahc", "mtp:qa_hc_up@mtp.mhc", "mtp:qa_hc_up@mtp.mix" }; try m.mtpProject(h.h[1], h.ahc, down[0], up[0], t.inj_a, sl.rows); if (!r.fused_xsum) try r.call("mtp:lane_qmm_xsum#[2560]", &.{ t.mixed, sl.md }, &.{t.xs}); - if (r.dense) r.denseRows(t.mixed, D, h.proj, rows, t.p) else try r.call("mtp:lane_qmm_bytes_grouped@mtp.att.proj", &.{ t.mixed, t.xs, h.proj.wq, h.proj.sbt, sl.md }, &.{t.p}); + try m.mtpLane(h.proj, "mtp:lane_qmm_bytes_grouped@mtp.att.proj", t.mixed, sl.md, rows, D, t.p); try r.call("mtp:q4_attn_prep@mtp.att", &.{ t.p, sl.pos8, h.qn, h.kn, h.iqn, t.eps, t.log2base }, &.{ t.q, t.kout, t.iq }); r.enc.setPipeline(r.kv_pipe); for ([_]Buf{ t.kout, t.p, h.keys, h.vals, h.raw, sl.kvmeta }, 0..) |b, j| r.enc.setBuffer(b.b, b.off, j); @@ -2202,41 +2211,53 @@ pub const Model = struct { } r.rows = 1; if (!r.fused_xsum) try r.call("mtp:lane_qmm_xsum#[6144]", &.{ t.aout, h.md1 }, &.{t.xs}); - if (r.dense) r.denseRows(t.aout, 6144, h.out, 1, t.branch) else try r.call("mtp:lane_qmm_bytes_grouped@mtp.att.o", &.{ t.aout, t.xs, h.out.wq, h.out.sbt, h.md1 }, &.{t.branch}); + try m.mtpLane(h.out, "mtp:lane_qmm_bytes_grouped@mtp.att.o", t.aout, h.md1, 1, 6144, t.branch); try r.call("mtp:q4_hc_norm_plain#[10240]", &.{ h.h[1], t.inj_a, t.branch }, &.{ h.h[0], t.ssp }); try m.mtpProject(h.h[0], h.mhc, down[1], up[1], t.inj_m, h.md1); try r.router("mtp:q4_router_float@mtp.moe", t.mixed, h.router, h.md1, t.lg, 1); - try r.experts("mtp:qa_expert_gateup@mtp.moe.gate", "mtp:qa_expert_down_y@mtp.moe.down", t.mixed, t.lg, h.ex, t.act, t.pick, t.wts, h.md1, t.ydown); + if (h.xw) |w| try variants.experts(r, "mtp:qa_expert_gateup@mtp.moe.gate", "mtp:qa_expert_down_y@mtp.moe.down", t.mixed, t.lg, h.ex, t.act, t.pick, t.wts, h.md1, t.ydown, w) else try r.experts("mtp:qa_expert_gateup@mtp.moe.gate", "mtp:qa_expert_down_y@mtp.moe.down", t.mixed, t.lg, h.ex, t.act, t.pick, t.wts, h.md1, t.ydown); try r.call("mtp:q4_hc_norm_grouped#[10240]", &.{ h.h[0], t.inj_m, t.ydown, t.wts, t.lg }, &.{ h.h[1], t.ssp }); h.last = h.h[1]; try m.mtpProject(h.h[1], h.mix, down[2], up[2], t.inj_a, h.md1); const x = t.mixed; if (!r.fused_xsum) try r.call("mtp:lane_qmm_xsum#[2560]", &.{ x, h.md1 }, &.{t.xs}); - if (r.tp != null and r.dense) { // TP: this Mac's half of the draft vocabulary's tiles, then one exact argmax swap + // TP: this Mac's half of the draft vocabulary's tiles, then one exact argmax swap + if (r.tp != null and r.dense and h.draft.w == null) { const tp = r.tp.?; - const n = h.draft.wq.b.length() * 4 / (3 * D); + const six = try h.draft.recorded(); + const n = six.wq.b.length() * 4 / (3 * D); const tiles = n / 32; const half = (tiles + 1) / 2; const t0: usize = if (tp.rank == 0) 0 else half; const tn: usize = if (tp.rank == 0) half else tiles - half; - r.denseTiles(x, D, h.draft, h.logits, t0, tn); + r.denseTiles(x, D, six, h.logits, t0, tn); const lo = @min(t0 * 32, h.ids_n); const hi = @min((t0 + tn) * 32, h.ids_n); tp.argmax(r.enc, h.logits, n, lo, hi - lo, Buf{ .b = tp.pick_tmp }, Buf{ .b = tp.one }, 1); tp.mapIds(r.enc, h.ids, out); return; } - if (r.dense) r.denseRows(x, D, h.draft, 1, h.logits) else try r.call("mtp:lane_qmm_bytes_grouped@mtp.draft", &.{ x, t.xs, h.draft.wq, h.draft.sbt, h.md1 }, &.{h.logits}); + try m.mtpLane(h.draft, "mtp:lane_qmm_bytes_grouped@mtp.draft", x, h.md1, 1, D, h.logits); r.enc.setPipeline(r.argids_pipe); for ([_]Buf{ h.logits, out, h.n_ids, h.ids }, 0..) |b, j| r.enc.setBuffer(b.b, b.off, j); r.enc.dispatchThreads(mtl.Size.of(1024, 1, 1), mtl.Size.of(1024, 1, 1)); if (!r.serial) r.enc.barrier(); } + /// An MTP head lane: at its widths' runs, as dense rows (r.dense), or the recorded launch. + fn mtpLane(m: *Model, l: Lane, role: []const u8, x: Buf, md: Buf, rows: usize, k: usize, y: Buf) !void { + if (l.w) |w| return lane_w.launch(m.r, w, role, x, md, rows, y); + const six = try l.recorded(); + if (m.r.dense) return m.r.denseRows(x, k, six, rows, y); + try m.r.call(role, &.{ x, m.t.xs, six.wq, six.sbt, md }, &.{y}); + } + pub fn mtpProject(m: *Model, hn: Buf, hc: Hc, down: []const u8, up: []const u8, inj: Buf, rows: Buf) !void { const t = &m.t; - try m.r.call(down, &.{ hn, t.ssp, hc.scale, hc.dw, hc.ds, hc.db, t.eps, rows }, &.{t.part}); - try m.r.call(up, &.{ hn, t.ssp, hc.scale, t.part, hc.uw, hc.us, hc.ub, t.eps, rows }, &.{ t.mixed, inj }); + var c: [2][2]variants.Const = undefined; + const k = [2][]const variants.Const{ variants.width(&c[0], "BITS", "GS", if (hc.w) |w| w[0] else null), variants.width(&c[1], "BITS", "GS", if (hc.w) |w| w[1] else null) }; + try variants.callAs(m.r, down, m.r.rows, &.{ hn, t.ssp, hc.scale, hc.dw, hc.ds, hc.db, t.eps, rows }, &.{t.part}, k[0]); + try variants.callAs(m.r, up, m.r.rows, &.{ hn, t.ssp, hc.scale, t.part, hc.uw, hc.us, hc.ub, t.eps, rows }, &.{ t.mixed, inj }, k[1]); } /// Absorb `nexts.len` rows (chunks of up to MAXR) from `streams`; returns the draft after the last row. @@ -2295,6 +2316,7 @@ pub const Prompt = struct { seg_queue: ?mtl.Queue = null, // a later staggered segment's queue (made at first use) last: Buf = undefined, // the last chunk's streams before the final mixer (each segment's own) gu32: mtl.Pipeline, + wide: ?*prompt_w.Wide = null, // weights at other widths than 6/32 (prompt_w.zig) fused_gu: bool = true, // the experts' gate and up products and activation in one pass (same bits) expert_bm: usize = 0, // expert tile rows (0: by the average rows an expert) mtp_w: ?[3][3]Buf = null, // the MTP head's attention projection, fce, fch in MLX layout (pack_mtp_mlx.safetensors) @@ -2323,6 +2345,7 @@ pub const Prompt = struct { pub fn init(r: *Run, dir: []const u8, header: []const u8) !Prompt { var p: Prompt = undefined; // every field is set below: `undefined` skips the declared defaults p.r = r; + p.wide = null; p.skip = 0; p.sel = null; p.step = PMAX; @@ -2486,6 +2509,7 @@ pub const Prompt = struct { /// y[rows, n] (row stride ldy, 0: n) = x[rows, k] W^T, W 6-bit g32 in MLX's layout. pub fn qmm(p: *Prompt, x: Buf, w: [3]Buf, k: usize, n: usize, rows: usize, y: Buf, ldy: usize) void { if (p.skip & 8 != 0) return; + if (p.wide) |wd| if (wd.wide(w[0])) |ww| return prompt_w.qmm(p, wd, ww, x, w, k, n, rows, y, ldy); const tall = p.tall_tiles and rows >= 512; p.bind(if (tall) p.qmm6_128 else p.qmm6, &.{ w[0], w[1], w[2], x }); const prm = [4]i32{ @intCast(k), @intCast(n), @intCast(rows), @intCast(ldy) }; @@ -2501,6 +2525,7 @@ pub const Prompt = struct { if (p.skip & 1 != 0) return; // tiles of 32, 64 or 128 rows as experts average 32 and 128 rows (128: FZ prompt chunks of thousands of rows) const bm: usize = if (p.expert_bm != 0) p.expert_bm else if (p.tall_tiles and pairs >= 512 * 128) 128 else if (pairs >= 512 * 32) 64 else 32; + if (p.wide) |wd| if (wd.wide(w[0])) |ww| return prompt_w.gather(p, wd.of(ww.bits, ww.group), bm, x, w, k, n, pairs, y); p.bind(if (bm == 128) p.gather128 else if (bm == 64) p.gather64 else p.gather32, &.{ x, w[0], w[1], w[2], p.b.off }); const prm = [4]i32{ @intCast(pairs), @intCast(n), @intCast(k), 512 }; p.r.enc.setBytes(std.mem.asBytes(&prm), 5); @@ -2513,6 +2538,7 @@ pub const Prompt = struct { pub fn gatherGU(p: *Prompt, x: Buf, wg: []const Buf, wu: []const Buf, k: usize, n: usize, pairs: usize, y: Buf) void { if (p.skip & 1 != 0) return; const bm: usize = if (p.expert_bm == 32 or p.expert_bm == 64) p.expert_bm else if (pairs >= 512 * 32) 64 else 32; + if (p.wide) |wd| if (wd.wide(wg[0])) |ww| return prompt_w.gatherGU(p, wd.of(ww.bits, ww.group), bm, x, wg, wu, k, n, pairs, y); p.bind(if (bm == 64) p.gu64 else p.gu32, &.{ x, wg[0], wg[1], wg[2], wu[0], wu[1], wu[2], p.b.off }); const prm = [4]i32{ @intCast(pairs), @intCast(n), @intCast(k), 512 }; p.r.enc.setBytes(std.mem.asBytes(&prm), 8); @@ -2525,13 +2551,13 @@ pub const Prompt = struct { pub fn hc(p: *Prompt, m: *Model, hn: Buf, w: Hc, rows: usize, inj: Buf) void { if (p.skip & 16 != 0) return; const b = &p.b; - const nd = w.dw.b.length() / (WIDE * 6 / 32 * 4); + const nd = w.dw.b.length() / (WIDE * (if (w.w) |hw| hw[0][0] else 6) / 8); // the down rows (+ inject rows) p.bind(p.pl[0], &.{ hn, b.ssp, w.scale, m.t.eps, b.normed }); p.r.enc.dispatchThreads(mtl.Size.of(WIDE, rows, 1), mtl.Size.of(256, 1, 1)); p.barrier(); if (p.skip & 8 == 0) { // the down projection's 10,240-deep sum in 8 parts, added in order const parts: usize = 8; - p.bind(p.splitk, &.{ w.dw, w.ds, w.db, b.normed }); + p.bind(if (w.w) |hw| p.wide.?.of(@intCast(hw[0][0]), @intCast(hw[0][1])).splitk else p.splitk, &.{ w.dw, w.ds, w.db, b.normed }); const prm = [4]i32{ WIDE, @intCast(nd), @intCast(rows), @intCast(parts) }; p.r.enc.setBytes(std.mem.asBytes(&prm), 4); p.r.enc.setBuffer(b.part.b, b.part.off, 5); @@ -2820,7 +2846,9 @@ pub const Prompt = struct { tabs[0] = b.pids; tabs[1] = pp.starts; for (0..3 * GROUPS) |j| tabs[2 + j] = pp.tables[j]; - try r.callRows("qa_ple_lookup@ple", rows, &tabs, &.{b.emb}, null); + var cp: [2]variants.Const = undefined; + try r.callRows("qa_ple_lookup@ple", rows, &tabs, &.{b.emb}, try variants.oneRow(r, "qa_ple_lookup@ple", variants.width(&cp, "BITS", "GS", pp.w))); + if (pp.scale) |sc| sc.apply(r.enc, b.emb, rows * D, r.serial); p.qmm(b.emb, p.ple_kv, D, 12800, rows, b.kvp, 0); try r.callRows("q4_ple_gate@ple", rows, &.{ b.kvp, b.h[s.cur], pp.ks, pp.qs, pp.cs, t.eps }, &.{ b.gated, .{ .b = b.cin.b, .off = PLE_TAIL * WIDE * 2 } }, null); try r.callRows("q4_ple_conv@ple", rows, &.{ b.cin, pp.conv, b.gated, b.h[s.cur] }, &.{b.hout}, null); @@ -3004,7 +3032,9 @@ pub const Prompt = struct { const s = h.at(l); const b = &s.p.b; const m = h.m; - try s.p.r.callRows("qa_embed_rows@embed", s.rows, &.{ b.ids, m.embed[0], m.embed[1], m.embed[2] }, &.{b.h[0]}, null); + var ce: [2]variants.Const = undefined; + const ev = try variants.oneRow(s.p.r, "qa_embed_rows@embed", variants.width(&ce, "BITS", "GS", m.embed_w)); + try s.p.r.callRows("qa_embed_rows@embed", s.rows, &.{ b.ids, m.embed[0], m.embed[1], m.embed[2] }, &.{b.h[0]}, ev); } pub fn wait(_: *Hooks, i: usize) segments.Wait { return if (i == 1) .pre else .mixer; @@ -3095,6 +3125,7 @@ pub const Prompt = struct { @memcpy(ps[0].b.cin.b.contents()[0 .. PLE_TAIL * WIDE * 2], cin_old[0 .. PLE_TAIL * WIDE * 2]); var hooks: Hooks = .{ .m = m, .segs = segs[0..N] }; m.gpu_seconds += try segments.run(r.device, queues[0..N], LAYERS, if (r.serial) .serial else .concurrent, &hooks); + if (ps[0].wide) |wd| try wd.check(); // a weight the load left out of the width table fails the pass const sl = &segs[N - 1]; m.state = if (N % 2 == 0) a else 1 - a; // each segment flipped it once m.state_row = 0; @@ -3238,21 +3269,23 @@ pub const Prompt = struct { b.n_add.b.slice(u32, 1)[0] = @intCast(n * WIDE); const cb = r.queue.commandBuffer(); r.enc = cb.compute(if (r.serial) .serial else .concurrent); - try r.callRows("mtp:qa_embed_rows@embed", n, &.{ b.mids, m.embed[0], m.embed[1], m.embed[2] }, &.{b.emb}, null); + var ce: [2]variants.Const = undefined; + try r.callRows("mtp:qa_embed_rows@embed", n, &.{ b.mids, m.embed[0], m.embed[1], m.embed[2] }, &.{b.emb}, try variants.oneRow(r, "mtp:qa_embed_rows@embed", variants.width(&ce, "BITS", "GS", m.embed_w))); try r.callRows("mtp:q4_rms_rows@mtp.enorm", n, &.{ b.emb, h.enorm, t.eps }, &.{b.branch}, null); - if (p.mtp_w) |w| p.qmm(b.branch, w[1], D, D, n, b.aout, 0) else r.denseRows(b.branch, D, h.fce, n, b.aout); + if (p.mtp_w) |w| p.qmm(b.branch, w[1], D, D, n, b.aout, 0) else r.denseRows(b.branch, D, try h.fce.recorded(), n, b.aout); try r.callRows("mtp:q4_rms_rows@mtp.hnorm", n, &.{ streams, h.hnorm, t.eps }, &.{b.normed}, null); - if (p.mtp_w) |w| p.qmm(b.normed, w[2], D, D, 4 * n, b.gated, 0) else r.denseRows(b.normed, D, h.fch, 4 * n, b.gated); + if (p.mtp_w) |w| p.qmm(b.normed, w[2], D, D, 4 * n, b.gated, 0) else r.denseRows(b.normed, D, try h.fch.recorded(), 4 * n, b.gated); p.bind(r.add_pipe, &.{ b.aout, b.gated, b.xs, b.n_add }); r.enc.dispatchThreads(mtl.Size.of(n * WIDE, 1, 1), mtl.Size.of(256, 1, 1)); p.barrier(); try r.callRows("mtp:q4_hc_norm_none#[10240]", n, &.{b.xs}, &.{ b.hout, b.ssp }, null); p.hc(m, b.hout, h.ahc, n, b.inj_a); - if (p.mtp_w) |w| p.qmm(b.mixed, w[0], D, 13952, n, b.p, 0) else r.denseRows(b.mixed, D, h.proj, n, b.p); + if (p.mtp_w) |w| p.qmm(b.mixed, w[0], D, 13952, n, b.p, 0) else r.denseRows(b.mixed, D, try h.proj.recorded(), n, b.p); try r.callRows("mtp:q4_attn_prep@mtp.att", n, &.{ b.p, b.pos, h.qn, h.kn, h.iqn, t.eps, t.log2base }, &.{ b.q, b.kout, b.iq }, null); p.bind(r.kv_pipe, &.{ b.kout, b.p, h.keys, h.vals, h.raw, b.kvmeta }); r.enc.dispatchThreads(mtl.Size.of(512 * n, 1, 1), mtl.Size.of(256, 1, 1)); try m.finish(cb); + if (p.wide) |wd| try wd.check(); } }; diff --git a/zig/src/families/flashnext/variants.zig b/zig/src/families/flashnext/variants.zig new file mode 100644 index 0000000000..5d0390f5c8 --- /dev/null +++ b/zig/src/families/flashnext/variants.zig @@ -0,0 +1,110 @@ +//! Recorded any-width kernels (0.6.6's qa_ texts) at another width: the same text with its constants swapped. +const std = @import("std"); +const mtl = @import("metal"); +const replay = @import("replay.zig"); +const Run = replay.Run; +const Buf = replay.Buf; +const Variant = replay.Variant; +const Hc = replay.Hc; + +/// A recorded `constexpr int` and the value it takes. +pub const Const = struct { name: []const u8, value: usize }; + +/// `text` with `constexpr int name = ...;` set to `value`; the declaration must appear exactly once. +pub fn constant(a: std.mem.Allocator, text: []const u8, name: []const u8, value: usize) ![]const u8 { + const key = try std.fmt.allocPrint(a, "constexpr int {s} = ", .{name}); + if (std.mem.count(u8, text, key) != 1) return error.KernelPatch; + const at = std.mem.indexOf(u8, text, key).? + key.len; + const end = at + (std.mem.indexOfScalar(u8, text[at..], ';') orelse return error.KernelPatch); + return std.fmt.allocPrint(a, "{s}{d}{s}", .{ text[0..at], value, text[end..] }); +} + +/// A pack table of u32s read into host memory: read on the host only, so never a device buffer. +pub fn table(r: *Run, name: []const u8) ![]const u32 { + const e = try r.entry(name); + if (e.len % 4 != 0) return error.BadPackTable; + const out = try r.arena.alloc(u32, e.len / 4); + try replay.readAll(e.fd, std.mem.sliceAsBytes(out), e.at); + return out; +} + +/// The recorded variant `v` at `consts`, compiled once and kept beside the recorded variants under its own key. +pub fn patched(r: *Run, v: *Variant, consts: []const Const) !*Variant { + var key: std.ArrayList(u8) = .empty; + try key.appendSlice(r.arena, v.name); + for (consts) |c| try key.print(r.arena, "|{s}={d}", .{ c.name, c.value }); + if (r.variants.get(key.items)) |hit| return hit; + var text = try Run.variantText(r.arena, v); + for (consts) |c| text = try constant(r.arena, text, c.name, c.value); + const lib = try mtl.Library.fromSource(r.device, text, mtl.CompileOptions.mlx()); + const nv = try r.arena.create(Variant); + nv.* = v.*; + nv.pipe = try mtl.Pipeline.init(r.device, lib, try std.fmt.allocPrintSentinel(r.arena, "{s}", .{v.name}, 0), false); + nv.text = text; + try r.variants.put(r.arena, key.items, nv); + return nv; +} + +/// `role` launched as recorded at `as_rows` rows on its kernel at `consts` (Run.callAs when there are none). +pub fn callAs(r: *Run, role: []const u8, as_rows: usize, ins: []const Buf, outs: []const Buf, consts: []const Const) !void { + if (consts.len == 0) return r.callAs(role, as_rows, ins, outs); + if (r.skip & Run.class(role) != 0) return; + if (r.hcskip != 0) for ([_][]const u8{ "hc_norm", "hc_down", "hc_up" }, 0..) |n, i| { + if (r.hcskip & (@as(u32, 1) << @intCast(i)) != 0 and std.mem.indexOf(u8, role, n) != null) return; + }; + var name: [96]u8 = undefined; + const s = r.roles.get(try std.fmt.bufPrint(&name, "{s}|{d}", .{ role, as_rows })) orelse return error.NoSite; + try r.bindV(try patched(r, s.v, consts), ins, outs); + r.enc.dispatchThreads(s.grid, s.tg); + if (!r.serial) r.enc.barrier(); +} + +/// The one-row variant of `role` at `consts`, as callRows takes an override; null for the recorded width. +pub fn oneRow(r: *Run, role: []const u8, consts: []const Const) !?*Variant { + if (consts.len == 0) return null; + var name: [96]u8 = undefined; + const s = r.roles.get(try std.fmt.bufPrint(&name, "{s}|1", .{role})) orelse return error.NoSite; + return try patched(r, s.v, consts); +} + +/// A width's two constants under the names a kernel declares them; none for the recorded 6/32 (null). +pub fn width(buf: *[2]Const, bits: []const u8, group: []const u8, w: ?[2]usize) []const Const { + const v = w orelse return &.{}; + buf.* = .{ .{ .name = bits, .value = v[0] }, .{ .name = group, .value = v[1] } }; + return buf; +} + +/// Expert widths as the expert kernels name them: routed (WB, WG) and shared (SWB, SWG), gate/up and down apart. +pub const Experts = struct { gu: [4]usize, down: [4]usize }; + +pub fn four(buf: *[4]Const, w: [4]usize) []const Const { + buf.* = .{ .{ .name = "WB", .value = w[0] }, .{ .name = "WG", .value = w[1] }, .{ .name = "SWB", .value = w[2] }, .{ .name = "SWG", .value = w[3] } }; + return buf; +} + +/// The routed and shared experts on the recorded any-width kernels at `w` (Run.experts' recorded path). +pub fn experts(r: *Run, gu_role: []const u8, down_role: []const u8, x: Buf, lg: Buf, e: [18]Buf, act: Buf, pick: Buf, wts: Buf, rows_buf: Buf, y: Buf, w: Experts) !void { + if (r.skip & (1 << 2) != 0) return; + var cg: [4]Const = undefined; + var cd: [4]Const = undefined; + try callAs(r, gu_role, r.rows, &.{ x, lg, e[0], e[1], e[2], e[3], e[4], e[5], e[6], e[7], e[8], e[9], e[10], e[11] }, &.{ act, pick, wts }, four(&cg, w.gu)); + try callAs(r, down_role, r.rows, &.{ act, pick, e[12], e[13], e[14], e[15], e[16], e[17], rows_buf }, &.{y}, four(&cd, w.down)); +} + +/// A pack hyper-connection: its seven tensors, and its `.widths` table (down, up) when not 6/32. +pub fn hcOf(r: *Run, comptime fmt: []const u8, args: anytype) !Hc { + var name: [96]u8 = undefined; + const stem = try std.fmt.bufPrint(&name, fmt, args); + var full: [128]u8 = undefined; + const parts = [_][]const u8{ "scale", "down.w", "down.s", "down.b", "up.w", "up.s", "up.b" }; + var out: [7]Buf = undefined; + for (parts, 0..) |p, i| out[i] = try r.load(try std.fmt.bufPrint(&full, "{s}.{s}", .{ stem, p })); + var hc: Hc = .{ .scale = out[0], .dw = out[1], .ds = out[2], .db = out[3], .uw = out[4], .us = out[5], .ub = out[6] }; + const widths = try std.fmt.bufPrint(&full, "{s}.widths", .{stem}); // any width but 6/32 (pack.putHc) + if (r.index.contains(widths)) { + const v = try table(r, widths); + if (v.len != 4) return error.BadPackTable; + hc.w = .{ .{ v[0], v[1] }, .{ v[2], v[3] } }; + } + return hc; +} diff --git a/zig/tests/flashnext_bench.zig b/zig/tests/flashnext_bench.zig index a4d2b87894..7a0c73d528 100644 --- a/zig/tests/flashnext_bench.zig +++ b/zig/tests/flashnext_bench.zig @@ -30,6 +30,11 @@ const Mode = enum { one, tp, stream, new, tpnew }; /// One projection class: its role, input width, input and output buffers, and every layer's weights. const Pass = struct { name: []const u8, role: []const u8, k: usize, x: Buf, y: Buf, lanes: []const Lane }; +/// A lane's recorded 6/32 pair: the harness loads 6/32 packs only (flashnext_run refuses others). +fn six(l: Lane) Lane.Six { + return l.six.?; +} + const tp_in = [_][2]usize{ .{ 0, 32 }, .{ 64, 32 }, .{ 128, 96 }, .{ 320, 96 }, .{ 512, 3 } }; const Bench = struct { @@ -52,7 +57,7 @@ const Bench = struct { /// Bytes a pass reads (speed-up mode: rank 0's tiles or K half). fn bytes(p: Pass, mode: Mode) f64 { var n: usize = 0; - for (p.lanes) |l| n += (l.wq.b.length() - l.wq.off) + (l.sbt.b.length() - l.sbt.off); + for (p.lanes) |l| n += (six(l).wq.b.length() - six(l).wq.off) + (six(l).sbt.b.length() - six(l).sbt.off); const all: f64 = @floatFromInt(n); if (mode != .tp and mode != .tpnew) return all; if (std.mem.endsWith(u8, p.role, "gdn.in")) return all * 259.0 / 515.0; @@ -62,13 +67,13 @@ const Bench = struct { /// Where a pass writes (the speed-up out-projection's fp32 partial, or the class's output), and its bytes at `rows`. fn out(b: *Bench, p: Pass, mode: Mode, rows: usize) struct { buf: Buf, len: usize } { - const n = (p.lanes[0].wq.b.length() - p.lanes[0].wq.off) * 4 / (3 * p.k); + const n = (six(p.lanes[0]).wq.b.length() - six(p.lanes[0]).wq.off) * 4 / (3 * p.k); if ((mode == .tp or mode == .tpnew) and std.mem.endsWith(u8, p.role, "gdn.out")) return .{ .buf = b.part, .len = rows * fz.D * 4 }; return .{ .buf = p.y, .len = rows * n * 2 }; } fn layout(b: *Bench, p: Pass, l: Lane, tp: bool) !dense.Layout { - const n = (l.wq.b.length() - l.wq.off) * 4 / (3 * p.k); + const n = (six(l).wq.b.length() - six(l).wq.off) * 4 / (3 * p.k); const sk = (try dense.recorded(b.r, p.role)).sk; const pf = b.r.lane_pf; if (!tp) return .{ .n = n, .k = p.k, .sk = sk, .pf = pf }; @@ -83,18 +88,18 @@ const Bench = struct { const t = &b.m.t; switch (mode) { .one => try b.m.lane(p.x, p.k, l, p.role, p.y), - .new, .tpnew => try dense.project(r, try b.layout(p, l, mode == .tpnew), p.x, l, t.mdims, b.out(p, mode, 1).buf), + .new, .tpnew => try dense.project(r, try b.layout(p, l, mode == .tpnew), p.x, six(l), t.mdims, b.out(p, mode, 1).buf), .stream => { r.enc.setPipeline(b.stream); - r.enc.setBuffer(l.wq.b, l.wq.off, 0); - r.enc.setBuffer(l.sbt.b, l.sbt.off, 1); + r.enc.setBuffer(six(l).wq.b, six(l).wq.off, 0); + r.enc.setBuffer(six(l).sbt.b, six(l).sbt.off, 1); r.enc.setBuffer(b.sink.b, 0, 2); - const n = [2]u32{ @intCast((l.wq.b.length() - l.wq.off) / 16), @intCast((l.sbt.b.length() - l.sbt.off) / 16) }; + const n = [2]u32{ @intCast((six(l).wq.b.length() - six(l).wq.off) / 16), @intCast((six(l).sbt.b.length() - six(l).sbt.off) / 16) }; r.enc.setBytes(std.mem.asBytes(&n), 3); r.enc.dispatchThreads(mtl.Size.of(160 * 1024, 1, 1), mtl.Size.of(1024, 1, 1)); }, .tp => { - const ins = [_]Buf{ p.x, t.xs, l.wq, l.sbt, t.mdims }; + const ins = [_]Buf{ p.x, t.xs, six(l).wq, six(l).sbt, t.mdims }; if (std.mem.endsWith(u8, p.role, "gdn.in")) { try split.laneTiles(r, p.role, &ins, p.y, &tp_in, 8, null); } else if (std.mem.endsWith(u8, p.role, "gdn.out")) { @@ -136,7 +141,7 @@ const Bench = struct { for (p.lanes) |l| { const lay = try b.layout(p, l, tp); r.enc.setPipeline(pipe); - for ([_]Buf{ p.x, l.wq, l.sbt, b.m.t.mdims, b.out(p, if (tp) .tpnew else .new, 1).buf }, 0..) |bf, j| r.enc.setBuffer(bf.b, bf.off, j); + for ([_]Buf{ p.x, six(l).wq, six(l).sbt, b.m.t.mdims, b.out(p, if (tp) .tpnew else .new, 1).buf }, 0..) |bf, j| r.enc.setBuffer(bf.b, bf.off, j); r.enc.dispatchThreads(mtl.Size.of(dense.tiles(lay) * 32 * lay.sk, 1, 1), mtl.Size.of(32 * lay.sk, 1, 1)); } } diff --git a/zig/tests/flashnext_run.zig b/zig/tests/flashnext_run.zig index eca3fc964a..87efe89c96 100644 --- a/zig/tests/flashnext_run.zig +++ b/zig/tests/flashnext_run.zig @@ -29,7 +29,6 @@ const KW = fz.KW; const Select = fz.Select; const GSelect = fz.GSelect; const Hc = fz.Hc; -const Lane = fz.Lane; const Layer = fz.Layer; const hcOf = fz.hcOf; const laneOf = fz.laneOf; @@ -37,7 +36,6 @@ const Ple = fz.Ple; const Tmp = fz.Tmp; const f32Buf = fz.f32Buf; const i32Buf = fz.i32Buf; -const Slot = fz.Slot; const Mtp = fz.Mtp; const Model = fz.Model; const DepthRule = fz.DepthRule; @@ -286,6 +284,8 @@ pub fn main(init: std.process.Init) !void { } return; } + // the replay below reads every expert, embedding and table at 6/32: other widths are the engine's (Engine.load) + if (!try fz.engine_widths.uniformDir(arena, init.io, args[1])) return error.HarnessNeeds6Bit32; var r = Run{ .arena = arena, .device = device, .queue = try device.queue() }; r.fused_xsum = std.c.getenv("FZ_FUSED_XSUM") != null; r.serial = std.c.getenv("FZ_SERIAL") != null; From a42b541377c041c6d21e0627569d0f35b1b47adc Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 04:42:39 -1000 Subject: [PATCH 10/10] tf-flashnext-widths: Flash Next products at every MLX width against float64, each row's bits in every window No model files: seeded synthetic codes, scales, biases and inputs, packed as the engine packs them. Every projection class at bits {2,3,4,5,6,8} x group {32,64,128}, decode and prompt separately, plus each public checkpoint's mixed stacks and oQ4e's ragged one: lanes at K slices 1, 2, 4 and 8; hyper-connections (down partials and up); experts (gate/up through silu, down); the embedding and the n-gram tables (8-bit codes over their full range); prompt dense, split-K, expert gathers and the fused gate/up. Each check reports its worst column against a float64 bound, and a reference with scales 1% off must exceed it (negative control). Every row's bits must be equal alone and in every window (1-16 rows in decode, 1-65 in prompt), and the generic texts at 6/32 must give the recorded kernels' bits. `tf-flashnext-widths lane|hc|moe|prompt` runs one class. tf-flashnext-widths (M3 Ultra, no tensor units): PASS 2722 checks, worst column 0.9947 of its float64 bound (bf16 output rounding); all 262 negative controls exceed the bound, the least by 1.012x; every row's bits equal alone and in every window. The core lane check without tensor units passes as on the core commits. Qwen3.8-Flash-Next-MLX-oQ2-MTP @0a1bf6f on an M3 Ultra (its cached pack; loads in 21.1 s from a cold file cache, 9.0 s warm): drafted == `"draft": false` 4/4 (MTP acceptance 55-59%; plain 45.0-54.1 tok/s, drafted 84.3-91.2); resumed == fresh (a 5.3k-token prompt at 0 and 5262 cached tokens; the second turn resumed at 5262 == a server with `--prompt-cache-gib 0`); concurrent == solo, 2 streams; greedy text sane (RYB and RGB primaries, a doctested function, "5:45 PM"). Gate on main 0efbd447 + the core commits (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test logs an error on the path it expects); native and tf-flashnext-run (with flashnext_bench), -reuse, -pack and -think build. lean_check: 531 problems (main 539); the only findings not on main are the two long files' shorter line counts. --- build.zig | 2 + zig/build/flashnext_checks.zig | 16 ++ zig/tests/flashnext_widths.zig | 36 ++++ zig/tests/flashnext_widths_hc.zig | 238 +++++++++++++++++++++ zig/tests/flashnext_widths_lane.zig | 188 ++++++++++++++++ zig/tests/flashnext_widths_lib.zig | 272 ++++++++++++++++++++++++ zig/tests/flashnext_widths_moe.zig | 294 ++++++++++++++++++++++++++ zig/tests/flashnext_widths_prompt.zig | 249 ++++++++++++++++++++++ 8 files changed, 1295 insertions(+) create mode 100644 zig/build/flashnext_checks.zig create mode 100644 zig/tests/flashnext_widths.zig create mode 100644 zig/tests/flashnext_widths_hc.zig create mode 100644 zig/tests/flashnext_widths_lane.zig create mode 100644 zig/tests/flashnext_widths_lib.zig create mode 100644 zig/tests/flashnext_widths_moe.zig create mode 100644 zig/tests/flashnext_widths_prompt.zig diff --git a/build.zig b/build.zig index d3b8a0a760..273e5b9486 100644 --- a/build.zig +++ b/build.zig @@ -4,6 +4,7 @@ const cuda_build = @import("zig/build/cuda.zig"); const dist_build = @import("zig/build/dist.zig"); const qwen27_build = @import("zig/build/qwen27.zig"); const core_checks = @import("zig/build/core_checks.zig"); +const flashnext_checks = @import("zig/build/flashnext_checks.zig"); const hip_build = @import("zig/build/hip.zig"); const grammar_build = @import("zig/build/grammar.zig"); const counters_build = @import("zig/build/metal_counters.zig"); @@ -477,6 +478,7 @@ fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b // The core lane, shared-context and Qwen3.8-27B programs and host tests. qwen27_build.targets(b, target, optimize, metal, sources, engine, lanes, mods.qwen27, test_step); core_checks.targets(b, target, optimize, metal, engine); + flashnext_checks.targets(b, target, optimize, metal, engine); _ = nativeServer(b, target, optimize, metal, engine, api, engines, build_options, test_step); diff --git a/zig/build/flashnext_checks.zig b/zig/build/flashnext_checks.zig new file mode 100644 index 0000000000..2b0e7a8386 --- /dev/null +++ b/zig/build/flashnext_checks.zig @@ -0,0 +1,16 @@ +//! Flash Next's device checks on synthetic weights: programs that run on the GPU, each under its own step. + +const std = @import("std"); + +/// tf-flashnext-widths: every product class at every oQ width against float64, every row's bits in every window. +pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, metal: *std.Build.Module, engine: *std.Build.Module) void { + const exe = b.addExecutable(.{ .name = "tf-flashnext-widths", .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/flashnext_widths.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "tensorfold", .module = engine } }, + }) }); + const about = "Flash Next's products at every oQ width against float64 (worst column), every row's bits in every window"; + b.step("tf-flashnext-widths", about).dependOn(&b.addInstallArtifact(exe, .{}).step); +} diff --git a/zig/tests/flashnext_widths.zig b/zig/tests/flashnext_widths.zig new file mode 100644 index 0000000000..b7cdc2905c --- /dev/null +++ b/zig/tests/flashnext_widths.zig @@ -0,0 +1,36 @@ +//! tf-flashnext-widths: Flash Next's products at every oQ width against float64, each row's bits in every window. +const std = @import("std"); +const mtl = @import("metal"); +const lib = @import("flashnext_widths_lib.zig"); +const lane = @import("flashnext_widths_lane.zig"); +const hc = @import("flashnext_widths_hc.zig"); +const moe = @import("flashnext_widths_moe.zig"); +const prompt = @import("flashnext_widths_prompt.zig"); +const fz = @import("tensorfold").flashnext_replay; + +pub fn main(init: std.process.Init) !void { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + const args = try init.minimal.args.toSlice(lib.a); + const only: ?[]const u8 = if (args.len > 1) args[1] else null; // a class name runs that class alone + std.debug.print("device {s}, tensor units {}\n", .{ device.name(), device.tensorUnits() }); + var t = lib.Tally{}; + if (only == null or std.mem.eql(u8, only.?, "lane")) try lane.run(&t, device, queue); + if (only == null or !std.mem.eql(u8, only.?, "lane")) { // the recorded kernels the any-width ones come from + const r = try lib.a.create(fz.Run); + r.* = .{ .arena = lib.a, .device = device, .queue = queue, .serial = true }; + try r.compileChecked(); + if (only == null or std.mem.eql(u8, only.?, "hc")) try hc.run(&t, r); + if (only == null or std.mem.eql(u8, only.?, "moe")) try moe.run(&t, r); + if (only == null or std.mem.eql(u8, only.?, "prompt")) try prompt.run(&t, r); + } + if (t.failed > 0) { + std.debug.print("FAIL {d} of {d} checks; the largest difference is {d:.4} of its float64 bound\n", .{ t.failed, t.checks, t.worst }); + return error.CheckFailed; + } + std.debug.print("PASS {d} checks: every product within {d:.4} of its float64 bound, every row's bits equal in every window\n", .{ t.checks, t.worst }); +} diff --git a/zig/tests/flashnext_widths_hc.zig b/zig/tests/flashnext_widths_hc.zig new file mode 100644 index 0000000000..7af69e1e43 --- /dev/null +++ b/zig/tests/flashnext_widths_hc.zig @@ -0,0 +1,238 @@ +//! tf-flashnext-widths, hyper-connections: the recorded any-width kernels at each oQ width, decode rows 1 to 16. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const lib = @import("flashnext_widths_lib.zig"); +const fz = tf.flashnext_replay; +const variants = fz.variants; +const Buf = fz.Buf; +const Run = fz.Run; +const a = lib.a; + +const W = 4 * 2560; // the streams' width: the down projection's K +const LOW = 320; +const counts = [_]usize{ 1, 2, 3, 4, 8, 9, 16 }; // the row kernels at 1-2 rows, the MMA kernels past them +const R = 16; + +/// The streams and their sums of squares set so every stream's norm is exactly one: x = bf16(hn) as stored. +const Streams = struct { hn: mtl.Buffer, hd: []f64, ssp: mtl.Buffer, nw: mtl.Buffer, eps: mtl.Buffer }; + +fn streams(device: mtl.Device) !Streams { + const in = try lib.Inputs.init(device, R, W, 0x4c); + const ssp = try device.buffer(R * 10 * 4 * 4, lib.opts); + @memset(ssp.slice(f32, R * 40), 256); // ten partial sums of 256 a stream: mean square 1, rsqrt 1 + const nw = try device.buffer(W * 4, lib.opts); + @memset(nw.slice(f32, W), 1); + const eps = try device.buffer(16, lib.opts); + eps.slice(f32, 1)[0] = 0; + return .{ .hn = in.buf, .hd = in.x, .ssp = ssp, .nw = nw, .eps = eps }; +} + +const Weights = struct { dw: Buf, ds: Buf, db: Buf, uw: Buf, us: Buf, ub: Buf }; + +fn upload(device: mtl.Device, down: lib.Mat, up: lib.Mat) !Weights { + return .{ + .dw = .{ .b = try lib.upload(device, down.words) }, + .ds = .{ .b = try lib.upload(device, std.mem.sliceAsBytes(down.scales)) }, + .db = .{ .b = try lib.upload(device, std.mem.sliceAsBytes(down.biases)) }, + .uw = .{ .b = try lib.upload(device, up.words) }, + .us = .{ .b = try lib.upload(device, std.mem.sliceAsBytes(up.scales)) }, + .ub = .{ .b = try lib.upload(device, std.mem.sliceAsBytes(up.biases)) }, + }; +} + +/// One run's outputs: the down partials and the up's mixed rows and inject gates. +const Out = struct { part: mtl.Buffer, mixed: mtl.Buffer, inj: mtl.Buffer }; + +fn outputs(device: mtl.Device, nd: usize) !Out { + return .{ .part = try lib.poisoned(device, 10 * R * nd * 4), .mixed = try lib.poisoned(device, R * 2560 * 2), .inj = try lib.poisoned(device, R * 4 * 2) }; +} + +/// Rows [start, start + rows) through the down and up roles at `consts` (none: the recorded 6/32 kernels). +fn project(r: *Run, s: Streams, w: Weights, which: []const u8, start: usize, rows: usize, nd: usize, o: Out, kd: []const variants.Const, ku: []const variants.Const) !void { + var role: [64]u8 = undefined; + const rb = try r.buffer(16); + rb.slice(i32, 1)[0] = @intCast(rows); + const rows_buf: Buf = .{ .b = rb }; + const hn: Buf = .{ .b = s.hn, .off = start * W * 2 }; + const ssp: Buf = .{ .b = s.ssp, .off = start * 40 * 4 }; + const part: Buf = .{ .b = o.part, .off = start * 10 * nd * 4 }; + const down = try std.fmt.bufPrint(&role, "qa_hc_down@{s}", .{which}); + try variants.callAs(r, down, rows, &.{ hn, ssp, .{ .b = s.nw }, w.dw, w.ds, w.db, .{ .b = s.eps }, rows_buf }, &.{part}, kd); + const up = try std.fmt.bufPrint(&role, "qa_hc_up@{s}", .{which}); + const ins = [_]Buf{ hn, ssp, .{ .b = s.nw }, part, w.uw, w.us, w.ub, .{ .b = s.eps }, rows_buf }; + try variants.callAs(r, up, rows, &ins, &.{ .{ .b = o.mixed, .off = start * 2560 * 2 }, .{ .b = o.inj, .off = start * 4 * 2 } }, ku); +} + +fn runAll(r: *Run, s: Streams, w: Weights, which: []const u8, rows: usize, nd: usize, o: Out, kd: []const variants.Const, ku: []const variants.Const) !void { + const cb = r.queue.commandBuffer(); + r.enc = cb.compute(.serial); + var start: usize = 0; + while (start + rows <= R) : (start += rows) try project(r, s, w, which, start, rows, nd, o, kd, ku); + try lib.commit(cb, r.enc); +} + +/// One HC at a down and up width: recorded 6/32 bits, float64, row bits; `full`: codes past 63, no 6/32 twin. +fn check(t: *lib.Tally, r: *Run, s: Streams, which: []const u8, nd: usize, dw: [2]usize, uw: [2]usize, seed: u64, full: bool) !void { + var label_buf: [128]u8 = undefined; + const label = try std.fmt.bufPrint(&label_buf, "hc {s} down q{d}/g{d} up q{d}/g{d}{s}", .{ which, dw[0], dw[1], uw[0], uw[1], if (full) " full-range codes" else "" }); + const top: u32 = if (full) 256 else 64; // below 64: codes the recorded 6/32 kernels read too + const down = try lib.Mat.initBelow(nd, W, @intCast(dw[0]), dw[1], seed, @min(top, @as(u32, 1) << @intCast(dw[0]))); + const up = try lib.Mat.initBelow(W, LOW, @intCast(uw[0]), uw[1], seed + 1, @min(top, @as(u32, 1) << @intCast(uw[0]))); + const w = try upload(r.device, down, up); + const w6 = if (full) w else try upload(r.device, try down.recode(6, 32), try up.recode(6, 32)); + var cd: [2]variants.Const = undefined; + var cu: [2]variants.Const = undefined; + const kd = variants.width(&cd, "BITS", "GS", dw); + const ku = variants.width(&cu, "BITS", "GS", uw); + const one = try outputs(r.device, nd); + try runAll(r, s, w, which, 1, nd, one, kd, ku); + var same = true; + for (counts) |n| { + const got = try outputs(r.device, nd); + try runAll(r, s, w, which, n, nd, got, kd, ku); + const rows = R / n * n; + if (!full) { + const rec = try outputs(r.device, nd); + try runAll(r, s, w6, which, n, nd, rec, &.{}, &.{}); + same = lib.sameRows(t, label, n, rec.mixed.slice(u8, R * 5120), got.mixed.slice(u8, R * 5120), rows, 5120) and same; + same = lib.sameRows(t, label, n, rec.inj.slice(u8, R * 8), got.inj.slice(u8, R * 8), rows, 8) and same; + } + same = lib.sameRows(t, label, n, one.mixed.slice(u8, R * 5120), got.mixed.slice(u8, R * 5120), rows, 5120) and same; + same = lib.sameRows(t, label, n, one.inj.slice(u8, R * 8), got.inj.slice(u8, R * 8), rows, 8) and same; + } + if (same) std.debug.print("RESULT {s}: {s}every row's bits at 1-16 rows\n", .{ label, if (full) "" else "the recorded 6/32 kernels' bits on the same values, and " }); + // the down partials of the one-row launches against float64: the ten K parts summed + const terms = try down.dense(); + defer for (terms) |v| a.free(v); + const wrong = try down.skewed(1.01).dense(); + defer for (wrong) |v| a.free(v); + var worst = lib.Worst{}; + var ctrl = lib.Worst{}; + const part = one.part.slice(f32, 10 * R * nd); + for (0..R) |row| for (0..nd) |o| { + var sum: f64 = 0; + for (0..10) |k| sum += part[(row * 10 + k) * nd + o]; + const ref = lib.dot(s.hd[row * W ..][0..W], terms, o); + worst.add(o, sum, ref[0], lib.ratio(sum, ref[0], ref[1], 32 + W / dw[1] + 24, false)); + const cref = lib.dot(s.hd[row * W ..][0..W], wrong, o); + ctrl.add(o, sum, cref[0], lib.ratio(sum, cref[0], cref[1], 32 + W / dw[1] + 24, false)); + }; + var down_label: [160]u8 = undefined; + const dl = try std.fmt.bufPrint(&down_label, "{s} down", .{label}); + worst.report(t, dl); + ctrl.control(t, dl); + try upCheck(t, r, s, which, nd, down, terms, up, kd, ku, label); +} + +/// One bf16 step where fp32's sum and exp could land either side of a bf16 tie (2^-18 of p: a few fp32 steps). +fn tieStep(p: f64, got: f64) f64 { + return if (@abs(@abs(p - got) - lib.bfStep(p) / 2) <= @abs(p) * 0x1p-18) lib.bfStep(p) else 0; +} + +/// The up product against float64 through sigmoid(up) * stream 0 (other streams zero), centred so 1% scales show. +fn upCheck(t: *lib.Tally, r: *Run, s: Streams, which: []const u8, nd: usize, down: lib.Mat, dterms: [2][]f64, up0: lib.Mat, kd: []const variants.Const, ku: []const variants.Const, label: []const u8) !void { + var up = up0; + const gl = LOW / up.group; + up.biases = try a.alloc(u16, up0.biases.len); + for (0..W) |o| for (0..gl) |g| { + var codes: f64 = 0; + for (0..up.group) |i| codes += @floatFromInt(up.code(o, g * up.group + i)); + up.biases[o * gl + g] = lib.toBf(@floatCast(-lib.fromBf(up.scales[o * gl + g]) * codes / @as(f64, @floatFromInt(up.group)))); + }; + // stream 0's scale: a power of two, so bf16(scale x) is exact; picked on row 0 from float64 estimates + const hd = try a.alloc(f64, R * W); + @memset(hd, 0); + var dsum: [LOW]f64 = undefined; + for (0..LOW) |c| { + var v: f64 = 0; + for (0..2560) |i| v += dterms[0][c * W + i] * s.hd[i]; + dsum[c] = v; + } + var scale: f64 = 1; + var best: f64 = std.math.inf(f64); + var e: i32 = -14; + while (e <= 0) : (e += 1) { + const f = std.math.ldexp(@as(f64, 1), e); + var accs: [256]f64 = undefined; + for (&accs, 0..) |*acc, o| { + acc.* = 0; + for (0..LOW) |i| { + const v = f * dsum[i] / 4; + acc.* += up.deq(o * 10, i) * v / (1 + @exp(-v)); + } + acc.* = @abs(acc.*); + } + std.mem.sort(f64, &accs, {}, std.sort.asc(f64)); + const miss = @abs(@log2(accs[128] + 1e-30)); + if (miss < best) { + best = miss; + scale = f; + } + } + const hn = try r.device.buffer(R * W * 2, lib.opts); + const hv = hn.slice(u16, R * W); + @memset(hv, 0); + for (0..R) |row| for (0..2560) |i| { + hd[row * W + i] = scale * s.hd[row * W + i]; + hv[row * W + i] = lib.toBf(@floatCast(hd[row * W + i])); + }; + const s0: Streams = .{ .hn = hn, .hd = hd, .ssp = s.ssp, .nw = s.nw, .eps = s.eps }; + const o = try outputs(r.device, nd); + try runAll(r, s0, try upload(r.device, down, up), which, 1, nd, o, kd, ku); + const part = o.part.slice(f32, 10 * R * nd); + const mixed = o.mixed.slice(u16, R * 2560); + var worst = lib.Worst{}; + var ctrl = lib.Worst{}; + var act: [LOW]f64 = undefined; + var slack: [LOW]f64 = undefined; + const steps: f64 = @floatFromInt(32 + LOW / up.group + 16); + for (0..R) |row| { + for (0..LOW) |c| { // the kernel's prologue: the parts summed in fp32, bf16, a quarter, bf16(silu) + var v: f32 = 0; + for (0..10) |k| v += part[(row * 10 + k) * nd + c]; + const v4 = lib.bfRound(v) / 4; + const p = v4 / (1 + @exp(-v4)); + act[c] = lib.bfRound(p); + slack[c] = tieStep(p, act[c]); + } + for (0..2560) |d| for ([_]f64{ 1, 1.01 }) |skew| { + var acc: [3]f64 = .{ 0, 0, 0 }; // the product, its magnitude, the activations' tie slack + for (0..LOW) |i| { + const g = d * gl + i / up.group; + const sq = skew * lib.fromBf(up.scales[g]) * @as(f64, @floatFromInt(up.code(d, i))); + const b = lib.fromBf(up.biases[g]); + acc[0] += (sq + b) * act[i]; + acc[1] += (sq + @abs(b)) * @abs(act[i]); + acc[2] += @abs(sq + b) * slack[i]; + } + const ea = steps * lib.unit * acc[1] + acc[2]; + const sig = 1 / (1 + @exp(-acc[0])); + const normed = hd[row * W + d]; + const want = sig * normed; + // bf16(up), the sigmoid's slope at most 1/4, bf16(sigmoid), bf16(sigmoid times the stream) + const tol = @abs(normed) * (0.25 * (ea + 0x1p-8 * (@abs(acc[0]) + ea)) + 0x1p-8 * sig + 1e-6) + 0x1p-8 * 1.01 * @abs(want) + 1e-30; + const got = 4 * lib.fromBf(mixed[row * 2560 + d]); + (if (skew == 1) &worst else &ctrl).add(d, got, want, @abs(got - want) / tol); + }; + } + var up_label: [160]u8 = undefined; + const ul = try std.fmt.bufPrint(&up_label, "{s} up (stream scale 2^{d})", .{ label, std.math.log2(scale) }); + worst.report(t, ul); + ctrl.control(t, ul); +} + +pub fn run(t: *lib.Tally, r: *Run) !void { + const s = try streams(r.device); + var seed: u64 = 0x4c00; + for ([_]u8{ 2, 3, 4, 5, 6, 8 }) |bits| for ([_]usize{ 32, 64, 128 }) |group| { + seed += 4; + const up_group = @min(group, 64); // 320 inputs: the pack reads an up projection's group 128 as 64 + try check(t, r, s, "ahc", 324, .{ bits, group }, .{ bits, up_group }, seed, false); + if (bits == 8) try check(t, r, s, "ahc", 324, .{ bits, group }, .{ bits, up_group }, seed + 2, true); + }; + // the mixer (no inject rows) at the oQ2 and oQ4e widths, and a down and up of different widths + try check(t, r, s, "mix", 320, .{ 2, 32 }, .{ 2, 32 }, 0x4d01, false); + try check(t, r, s, "mix", 320, .{ 4, 32 }, .{ 4, 32 }, 0x4d02, false); + try check(t, r, s, "mhc", 324, .{ 5, 64 }, .{ 4, 64 }, 0x4d03, false); +} diff --git a/zig/tests/flashnext_widths_lane.zig b/zig/tests/flashnext_widths_lane.zig new file mode 100644 index 0000000000..9052ac0a2e --- /dev/null +++ b/zig/tests/flashnext_widths_lane.zig @@ -0,0 +1,188 @@ +//! tf-flashnext-widths, decode lanes: fused stacks packed as the engine packs them, every run at its own width. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const lib = @import("flashnext_widths_lib.zig"); +const lane_w = tf.flashnext_replay.lane_w; +const pack = tf.flashnext_pack; +const st = tf.safetensors; +const Buf = tf.flashnext_replay.Buf; +const a = lib.a; + +const wide = 64; // rows of the reference call: two 32-row blocks (the MTP head's 4-row windows reach 64) +const windows = [_]usize{ 1, 2, 3, 8, 16, 17, 33 }; // 17 and 33 cross the 16- and 32-row blocks + +/// A stack member: rows, bits and group as the checkpoint stores it. +const Member = struct { rows: usize, bits: u8, group: usize }; + +fn tensor(dtype: st.DType, rows: usize, cols: usize, bytes: []const u8) st.Tensor { + var shape: [st.max_rank]usize = @splat(0); + shape[0] = rows; + shape[1] = cols; + return .{ .dtype = dtype, .rank = 2, .shape = shape, .bytes = bytes }; +} + +fn q(m: lib.Mat) !pack.Q { + const kw = m.k * m.bits / 32; + const tables = std.mem.sliceAsBytes(m.scales); + return .{ .w = tensor(.u32, m.n, kw, m.words), .s = tensor(.bf16, m.n, m.k / m.group, tables), .b = tensor(.bf16, m.n, m.k / m.group, std.mem.sliceAsBytes(m.biases)), .rows = m.n, .k = m.k, .kw = kw, .spec = try @FieldType(pack.Q, "spec").init(m.bits, m.group) }; +} + +fn blob(out: *const pack.Out, name: []const u8) ?[]const u8 { + for (out.names.items, 0..) |n, i| if (std.mem.eql(u8, n, name)) return out.blobs.items[i]; + return null; +} + +/// The device lane the engine would load from these pack tensors (lane_w.of's reading, from memory). +fn wideOf(device: mtl.Device, out: *const pack.Out, k: usize, rows: usize) !lane_w.Wide { + const uniform = [_]u32{ 6, 32, @intCast(rows) }; // a 6/32 lane writes no table + const table = if (blob(out, "t.runs")) |b| std.mem.bytesAsSlice(u32, @as([]align(1) const u8, b)) else &uniform; + const n_runs = table.len / 3; + const run_rows = try a.alloc(usize, n_runs); + for (run_rows, 0..) |*n, i| n.* = table[i * 3 + 2]; + const maps = if (n_runs > 1) blk: { + const m = blob(out, "t.members").?; + const members = try a.alloc(u32, m.len / 4); + @memcpy(std.mem.sliceAsBytes(members), m); + break :blk try lane_w.columnMaps(a, members, run_rows); + } else null; + const runs = try a.alloc(lane_w.Product, n_runs); + var total: usize = 0; + for (runs, 0..) |*p, i| { + const base = if (n_runs == 1) "t" else try std.fmt.allocPrint(a, "t.r{d}", .{i}); + const wq = blob(out, try std.fmt.allocPrint(a, "{s}.wq", .{base})).?; + const sbt = blob(out, try std.fmt.allocPrint(a, "{s}.sbt", .{base})).?; + p.* = .{ .wq = .{ .b = try lib.upload(device, wq) }, .sbt = .{ .b = try lib.upload(device, sbt) }, .n = run_rows[i], .k = k, .bits = @intCast(table[i * 3]), .group = @intCast(table[i * 3 + 1]), .tiled = pack.lane.tiled(run_rows[i]) }; + if (maps) |m| p.map = .{ .b = try lib.upload(device, std.mem.sliceAsBytes(m[i])) }; + total += run_rows[i]; + } + return .{ .runs = runs, .n = total }; +} + +fn mdims(device: mtl.Device, rows: usize) !Buf { + const b = try device.buffer(32, lib.opts); + b.slice(i32, 2)[0] = @intCast(rows); + b.slice(i32, 2)[1] = @intCast((rows + 15) / 16 * 16); + return .{ .b = b }; +} + +/// Rows [start, start + m) of x through the lane at its widths, in windows of m rows. +fn launchWindows(device: mtl.Device, queue: mtl.Queue, cache: *lane_w.Cache, w: *const lane_w.Wide, x: lib.Inputs, m: usize, y: mtl.Buffer) !void { + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + var start: usize = 0; + while (start + m <= wide) : (start += m) { + try lane_w.encode(e, a, device, cache, w, .{ .b = x.buf, .off = start * x.k * 2 }, try mdims(device, m), m, .{ .b = y, .off = start * w.n * 2 }); + } + try lib.commit(cb, e); +} + +/// One stack: packed by the engine's packer, run at 64 rows against float64, then every window's rows bit for bit. +fn stack(t: *lib.Tally, device: mtl.Device, queue: mtl.Queue, cache: *lane_w.Cache, members: []const Member, k: usize, seed: u64) !void { + var mats = try a.alloc(lib.Mat, members.len); + var qs = try a.alloc(pack.Q, members.len); + var label: std.ArrayList(u8) = .empty; + try label.print(a, "lane K {d}:", .{k}); + for (members, 0..) |m, i| { + mats[i] = try lib.Mat.init(m.rows, k, m.bits, m.group, seed + i); + qs[i] = try q(mats[i]); + try label.print(a, " {d}x q{d}/g{d}", .{ m.rows, m.bits, m.group }); + } + var out = pack.Out.init(a); + try pack.lane.putLane(&out, a, "t", qs); + var rows: usize = 0; + for (members) |m| rows += m.rows; + const w = try wideOf(device, &out, k, rows); + try label.print(a, " (SK {d})", .{lane_w.splitK(w.runs[0].n, k)}); + const x = try lib.Inputs.init(device, wide, k, seed ^ 0x5eed); + const all = try lib.poisoned(device, wide * w.n * 2); + try launchWindows(device, queue, cache, &w, x, wide, all); + var worst = lib.Worst{}; + var ctrl = lib.Worst{}; + const got = all.slice(u16, wide * w.n); + var col: usize = 0; + for (mats) |m| { + const terms = try m.dense(); + const wrong = try m.skewed(1.01).dense(); + defer for (terms ++ wrong) |v| a.free(v); + for (0..m.n) |j| { + for (0..wide) |r| { + const g = lib.fromBf(got[r * w.n + col + j]); + const ref = lib.dot(x.row(r), terms, j); + worst.add(col + j, g, ref[0], lib.ratio(g, ref[0], ref[1], m.group + k / m.group + 16, true)); + const cref = lib.dot(x.row(r), wrong, j); + ctrl.add(col + j, g, cref[0], lib.ratio(g, cref[0], cref[1], m.group + k / m.group + 16, true)); + } + } + col += m.n; + } + worst.report(t, label.items); + ctrl.control(t, label.items); + var same = true; + for (windows) |m| { + const y = try lib.poisoned(device, wide * w.n * 2); + try launchWindows(device, queue, cache, &w, x, m, y); + same = lib.sameRows(t, label.items, m, all.slice(u8, wide * w.n * 2), y.slice(u8, wide * w.n * 2), wide / m * m, w.n * 2) and same; + } + if (same) std.debug.print("RESULT {s}: every row's bits equal in windows of 1, 2, 3, 8, 16, 17 and 33 rows ({d} runs)\n", .{ label.items, w.runs.len }); +} + +/// The generic source at 6/32 against the recorded kernel at a recorded shape: the same bits. +fn recordedTwin(t: *lib.Tally, device: mtl.Device, queue: mtl.Queue, cache: *lane_w.Cache) !void { + const n = 2560; + const k = 6144; + const m = try lib.Mat.init(n, k, 6, 32, 0x632); + var out = pack.Out.init(a); + try pack.lane.putLane(&out, a, "t", &.{try q(m)}); + const wq: Buf = .{ .b = try lib.upload(device, blob(&out, "t.wq").?) }; + const sbt: Buf = .{ .b = try lib.upload(device, blob(&out, "t.sbt").?) }; + const x = try lib.Inputs.init(device, 16, k, 0x633); + const md = try mdims(device, 16); + const text = try lane_w.template(device.tensorUnits()); + const rec = try mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); + const rec_pipe = try mtl.Pipeline.init(device, rec, try lane_w.functionName(a, text), false); + const y_rec = try lib.poisoned(device, 16 * n * 2); + const y_gen = try lib.poisoned(device, 16 * n * 2); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + e.setPipeline(rec_pipe); + for ([_]Buf{ .{ .b = x.buf }, .{ .b = x.buf }, wq, sbt, md, .{ .b = y_rec } }, 0..) |b, j| e.setBuffer(b.b, b.off, j); + e.dispatchThreads(mtl.Size.of(n / 32 * 32 * 8, 1, 1), mtl.Size.of(32 * 8, 1, 1)); + const w = lane_w.Wide{ .runs = &.{.{ .wq = wq, .sbt = sbt, .n = n, .k = k, .bits = 6, .group = 32 }}, .n = n }; + try lane_w.encode(e, a, device, cache, &w, .{ .b = x.buf }, md, 16, .{ .b = y_gen }); + try lib.commit(cb, e); + const same = lib.sameRows(t, "lane recorded 6/32 [2560, 6144] against the generic source", 16, y_rec.slice(u8, 16 * n * 2), y_gen.slice(u8, 16 * n * 2), 16, n * 2); + if (same) std.debug.print("RESULT lane recorded 6/32 [2560, 6144]: the generic source's bits equal the recorded kernel's\n", .{}); +} + +pub fn run(t: *lib.Tally, device: mtl.Device, queue: mtl.Queue) !void { + var cache: lane_w.Cache = .empty; + try recordedTwin(t, device, queue, &cache); + var seed: u64 = 0x1a9e; + for ([_]u8{ 2, 3, 4, 5, 6, 8 }) |bits| for ([_]usize{ 32, 64, 128 }) |group| { + for ([_][2]usize{ .{ 256, 2560 }, .{ 64, 6144 } }) |shape| { + seed += 16; + try stack(t, device, queue, &cache, &.{.{ .rows = shape[0], .bits = bits, .group = group }}, shape[1], seed); + } + }; + // the public checkpoints' mixed stacks (q, k, v, index_qk and the DeltaNet inputs), rows cut down + const stacks = [_][]const Member{ + &.{ .{ .rows = 256, .bits = 3, .group = 64 }, .{ .rows = 32, .bits = 3, .group = 64 }, .{ .rows = 32, .bits = 8, .group = 64 }, .{ .rows = 64, .bits = 3, .group = 64 } }, + &.{ .{ .rows = 256, .bits = 6, .group = 64 }, .{ .rows = 32, .bits = 6, .group = 64 }, .{ .rows = 32, .bits = 6, .group = 128 }, .{ .rows = 64, .bits = 6, .group = 64 } }, + &.{ .{ .rows = 256, .bits = 4, .group = 64 }, .{ .rows = 64, .bits = 5, .group = 128 }, .{ .rows = 32, .bits = 5, .group = 128 }, .{ .rows = 32, .bits = 6, .group = 64 } }, + &.{ .{ .rows = 256, .bits = 3, .group = 64 }, .{ .rows = 64, .bits = 6, .group = 64 }, .{ .rows = 32, .bits = 6, .group = 64 }, .{ .rows = 32, .bits = 6, .group = 64 } }, + &.{ .{ .rows = 128, .bits = 4, .group = 32 }, .{ .rows = 64, .bits = 4, .group = 64 }, .{ .rows = 32, .bits = 2, .group = 32 } }, + // oQ4e layer 21's DeltaNet inputs: runs of 112 (5-bit z and b) and 48 (6-bit a) columns, untiled + &.{ .{ .rows = 256, .bits = 4, .group = 64 }, .{ .rows = 64, .bits = 5, .group = 128 }, .{ .rows = 48, .bits = 5, .group = 128 }, .{ .rows = 48, .bits = 6, .group = 64 } }, + &.{.{ .rows = 48, .bits = 8, .group = 64 }}, + }; + for (stacks) |s| { + seed += 16; + try stack(t, device, queue, &cache, s, 2560, seed); + } + // K slices 1 and 2 (the vocabulary head and draft at 8/64, gdn.in), at 8 bits and a low width + for ([_]usize{ 32768, 16384 }) |n| for ([_][2]usize{ .{ 8, 64 }, .{ 3, 64 } }) |wd| { + seed += 16; + try stack(t, device, queue, &cache, &.{.{ .rows = n, .bits = @intCast(wd[0]), .group = wd[1] }}, 2560, seed); + }; +} diff --git a/zig/tests/flashnext_widths_lib.zig b/zig/tests/flashnext_widths_lib.zig new file mode 100644 index 0000000000..8c87f5af50 --- /dev/null +++ b/zig/tests/flashnext_widths_lib.zig @@ -0,0 +1,272 @@ +//! Shared parts of tf-flashnext-widths: synthetic MLX matrices, float64 references and the row checks. +const std = @import("std"); +const mtl = @import("metal"); + +pub const opts = mtl.ResourceOptions.shared; +pub const a = std.heap.page_allocator; +pub const unit: f64 = 0x1p-24; // fp32's unit roundoff + +pub fn toBf(v: f32) u16 { + const u: u32 = @bitCast(v); + return @truncate((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); +} + +pub fn fromBf(v: u16) f64 { + return @as(f32, @bitCast(@as(u32, v) << 16)); +} + +/// One synthetic MLX affine matrix [n, k]: codes kept beside their packed words, bf16 tables. +pub const Mat = struct { + n: usize, + k: usize, + bits: u8, + group: usize, + codes: []u8, + words: []u8, + scales: []u16, + biases: []u16, + skew: f64 = 1, // the references' scale factor: 1.01 makes the negative controls' wrong references + + /// The same matrix with its scales `f` times their values in every reference (the kernels' bytes unchanged). + pub fn skewed(m: Mat, f: f64) Mat { + var out = m; + out.skew = f; + return out; + } + + pub fn init(n: usize, k: usize, bits: u8, group: usize, seed: u64) !Mat { + return initBelow(n, k, bits, group, seed, @as(u32, 1) << @intCast(bits)); + } + + /// Codes below `top` only (an 8-bit matrix whose codes fit 6 bits, so it also exists at 6 bits). + pub fn initBelow(n: usize, k: usize, bits: u8, group: usize, seed: u64, top: u32) !Mat { + var prng = std.Random.DefaultPrng.init(seed); + const r = prng.random(); + const tables = n * (k / group); + const codes = try a.alloc(u8, n * k); + for (codes) |*c| c.* = @intCast(r.uintLessThan(u32, top)); + const m = Mat{ .n = n, .k = k, .bits = bits, .group = group, .codes = codes, .words = try packCodes(codes, bits), .scales = try a.alloc(u16, tables), .biases = try a.alloc(u16, tables) }; + for (m.scales, m.biases) |*s, *b| { + s.* = toBf((r.float(f32) + 0.5) / 64); + b.* = toBf(-(r.float(f32) + 0.5) * 0.15); + } + return m; + } + + /// The same codes stored at `bits` in groups of `group` (which divides this group): the same values exactly. + pub fn recode(m: Mat, bits: u8, group: usize) !Mat { + if (m.group % group != 0) return error.GroupDoesNotDivide; + for (m.codes) |c| if (c >> @intCast(bits) != 0) return error.CodeTooWide; + const reps = m.group / group; + const out = Mat{ .n = m.n, .k = m.k, .bits = bits, .group = group, .codes = m.codes, .words = try packCodes(m.codes, bits), .scales = try a.alloc(u16, m.scales.len * reps), .biases = try a.alloc(u16, m.biases.len * reps) }; + for (m.scales, m.biases, 0..) |sv, bv, i| for (0..reps) |j| { + out.scales[i * reps + j] = sv; + out.biases[i * reps + j] = bv; + }; + return out; + } + + pub fn deinit(m: Mat) void { + a.free(m.words); + a.free(m.scales); + a.free(m.biases); + } + + pub fn code(m: Mat, j: usize, i: usize) u32 { + return m.codes[j * m.k + i]; + } + + pub fn deq(m: Mat, j: usize, i: usize) f64 { + const g = j * (m.k / m.group) + i / m.group; + return m.skew * fromBf(m.scales[g]) * @as(f64, @floatFromInt(m.code(j, i))) + fromBf(m.biases[g]); + } + + /// s q + b and s q + |b| for every weight, row-major [n, k]: the reference's terms and their magnitudes. + pub fn dense(m: Mat) ![2][]f64 { + const out = [2][]f64{ try a.alloc(f64, m.n * m.k), try a.alloc(f64, m.n * m.k) }; + for (0..m.n) |j| for (0..m.k) |i| { + const g = j * (m.k / m.group) + i / m.group; + const sq = m.skew * fromBf(m.scales[g]) * @as(f64, @floatFromInt(m.code(j, i))); + out[0][j * m.k + i] = sq + fromBf(m.biases[g]); + out[1][j * m.k + i] = sq + @abs(fromBf(m.biases[g])); + }; + return out; + } + + /// Weights as qmmw's dequantw rounds them, bf16(s q + b), their magnitudes, and a bf16 step where fp32 nears a tie. + pub fn denseBf(m: Mat) ![3][]f64 { + const out = [3][]f64{ try a.alloc(f64, m.n * m.k), try a.alloc(f64, m.n * m.k), try a.alloc(f64, m.n * m.k) }; + for (0..m.n) |j| for (0..m.k) |i| { + const v = m.deq(j, i); + const w = bfRound(v); + const tie = @abs(@abs(v - w) - bfStep(v) / 2) <= @abs(v) * 0x1p-23; // fp32's error, twice over + out[0][j * m.k + i] = w; + out[1][j * m.k + i] = @abs(w); + out[2][j * m.k + i] = if (tie) bfStep(v) else 0; + }; + return out; + } +}; + +/// v rounded to bf16 (nearest, ties to even), as float64. +pub fn bfRound(v: f64) f64 { + return fromBf(toBf(@floatCast(v))); +} + +/// The bf16 step at v's binade. +pub fn bfStep(v: f64) f64 { + if (v == 0) return 0x1p-133; + return std.math.ldexp(@as(f64, 1), std.math.frexp(v).exponent - 8); +} + +/// The float64 product of one row with one weight row, and its magnitude: sum of (s q + |b|) |x|. +pub fn dot(x: []const f64, terms: [2][]f64, j: usize) [2]f64 { + const k = x.len; + var want: f64 = 0; + var mag: f64 = 0; + for (x, terms[0][j * k ..][0..k], terms[1][j * k ..][0..k]) |xv, d, m| { + want += xv * d; + mag += @abs(xv) * m; + } + return .{ want, mag }; +} + +/// dot over denseBf's weights: the product, its magnitude sum |x| |w|, and the tie slack sum |x| step. +pub fn dot3(x: []const f64, terms: [3][]f64, j: usize) [3]f64 { + const k = x.len; + var out: [3]f64 = .{ 0, 0, 0 }; + for (x, terms[0][j * k ..][0..k], terms[1][j * k ..][0..k], terms[2][j * k ..][0..k]) |xv, w, m, e| { + out[0] += xv * w; + out[1] += @abs(xv) * m; + out[2] += @abs(xv) * e; + } + return out; +} + +/// ratio for a dot3 reference: fp32 sums of `terms` steps over the magnitude, the tie slack, one bf16 rounding. +pub fn ratio3(got: f64, ref: [3]f64, terms: usize, bf16_out: bool) f64 { + const tol = @as(f64, @floatFromInt(terms)) * unit * ref[1] + ref[2] + (if (bf16_out) @abs(ref[0]) * 0x1p-8 else 0) + 1e-30; + return @abs(got - ref[0]) / tol; +} + +/// Codes as an MLX little-endian bit stream of `bits`-wide values. +pub fn packCodes(codes: []const u8, bits: u8) ![]u8 { + const out = try a.alloc(u8, (codes.len * bits + 31) / 32 * 4); + @memset(out, 0); + for (codes, 0..) |c, i| { + const bit = i * bits; + const v = @as(u32, c) << @intCast(bit % 8); + for (0..3) |b| { + if (bit / 8 + b < out.len) out[bit / 8 + b] |= @truncate(v >> @intCast(8 * b)); + } + } + return out; +} + +/// Every check's outcome: failures, the largest difference as a fraction of its bound, and the negative controls. +pub const Tally = struct { failed: usize = 0, worst: f64 = 0, checks: usize = 0, controls: usize = 0, least: f64 = std.math.inf(f64) }; + +/// `rows` rows of x [rows, k] as bf16 (in a device buffer) and float64. +pub const Inputs = struct { + buf: mtl.Buffer, + x: []f64, + rows: usize, + k: usize, + + pub fn init(device: mtl.Device, rows: usize, k: usize, seed: u64) !Inputs { + var prng = std.Random.DefaultPrng.init(seed); + const r = prng.random(); + const in = Inputs{ .buf = try device.buffer(rows * k * 2, opts), .x = try a.alloc(f64, rows * k), .rows = rows, .k = k }; + for (in.buf.slice(u16, rows * k), in.x) |*v, *d| { + v.* = toBf(r.float(f32) * 2 - 1); + d.* = fromBf(v.*); + } + return in; + } + + pub fn row(in: Inputs, r: usize) []const f64 { + return in.x[r * in.k ..][0..in.k]; + } +}; + +/// A column's result against float64: the error over its bound (fp32 sums of `terms` steps, one bf16 rounding). +pub fn ratio(got: f64, want: f64, mag: f64, terms: usize, bf16_out: bool) f64 { + const tol = @as(f64, @floatFromInt(terms)) * unit * mag + (if (bf16_out) @abs(want) * 0x1p-8 else 0) + 1e-30; + return @abs(got - want) / tol; +} + +/// The worst column of one product, printed with its max abs and relative error and its share of the bound. +pub const Worst = struct { + frac: f64 = 0, + col: usize = 0, + abs: f64 = 0, + rel: f64 = 0, + nan: bool = false, + + pub fn add(w: *Worst, col: usize, got: f64, want: f64, frac: f64) void { + if (std.math.isNan(got)) w.nan = true; + const d = @abs(got - want); + if (frac > w.frac or std.math.isNan(got)) { + w.frac = if (std.math.isNan(got)) std.math.inf(f64) else frac; + w.col = col; + w.abs = d; + w.rel = d / (@abs(want) + 1e-30); + } + } + + pub fn report(w: Worst, t: *Tally, label: []const u8) void { + t.checks += 1; + t.worst = @max(t.worst, w.frac); + const ok = w.frac <= 1 and !w.nan; + if (!ok) t.failed += 1; + std.debug.print("RESULT {s}: worst column {d}: max abs {e:.3}, relative {e:.3}; {d:.4} of its bound{s}\n", .{ label, w.col, w.abs, w.rel, w.frac, if (ok) "" else " FAIL" }); + } + + /// A negative control: the same results against a reference with scales 1% off must exceed the bound somewhere. + pub fn control(w: Worst, t: *Tally, label: []const u8) void { + t.checks += 1; + t.controls += 1; + t.least = @min(t.least, w.frac); + const ok = w.frac > 1; + if (!ok) t.failed += 1; + std.debug.print("RESULT {s}: negative control, scales 1% off: {d:.3} of its bound{s}\n", .{ label, w.frac, if (ok) "" else " FAIL (the bound would pass a wrong scale)" }); + } +}; + +/// Rows [0, rows) of two outputs (row stride `ld` bytes) equal byte for byte; the first difference printed. +pub fn sameRows(t: *Tally, label: []const u8, window: usize, want: []const u8, got: []const u8, rows: usize, ld: usize) bool { + t.checks += 1; + if (std.mem.eql(u8, want[0 .. rows * ld], got[0 .. rows * ld])) return true; + const i = std.mem.indexOfDiff(u8, want, got).?; + var differ: usize = 0; + for (0..rows) |r| { + if (!std.mem.eql(u8, want[r * ld ..][0..ld], got[r * ld ..][0..ld])) differ += 1; + } + std.debug.print("RESULT {s}: {d}-row windows DIFFER in {d} of {d} rows; first at row {d} byte {d} FAIL\n", .{ label, window, differ, rows, i / ld, i % ld }); + t.failed += 1; + return false; +} + +pub fn commit(cb: anytype, e: mtl.ComputeEncoder) !void { + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |why| { + std.debug.print("GPU failed: {s}\n", .{why}); + return error.GpuFailed; + } +} + +/// A device buffer of `n` bytes filled with 0xff (NaN in bf16 and fp32): a launch that writes nothing fails. +pub fn poisoned(device: mtl.Device, n: usize) !mtl.Buffer { + const b = try device.buffer(@max(n, 16), opts); + @memset(b.slice(u8, @max(n, 16)), 0xff); + return b; +} + +/// A device buffer holding `bytes`. +pub fn upload(device: mtl.Device, bytes: []const u8) !mtl.Buffer { + const b = try device.buffer(@max(bytes.len, 16), opts); + @memcpy(b.slice(u8, bytes.len), bytes); + return b; +} diff --git a/zig/tests/flashnext_widths_moe.zig b/zig/tests/flashnext_widths_moe.zig new file mode 100644 index 0000000000..e8a216b7f5 --- /dev/null +++ b/zig/tests/flashnext_widths_moe.zig @@ -0,0 +1,294 @@ +//! tf-flashnext-widths, experts, embedding and n-gram tables: the recorded any-width kernels at each oQ width. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const lib = @import("flashnext_widths_lib.zig"); +const fz = tf.flashnext_replay; +const variants = fz.variants; +const Buf = fz.Buf; +const Run = fz.Run; +const a = lib.a; + +const D = 2560; +const N = 640; // an expert's intermediate width +const E = 16; // experts held: the logits pick among these only +const R = 16; +const counts = [_]usize{ 1, 2, 3, 8, 16 }; + +fn up(device: mtl.Device, m: lib.Mat) ![3]Buf { + return .{ .{ .b = try lib.upload(device, m.words) }, .{ .b = try lib.upload(device, std.mem.sliceAsBytes(m.scales)) }, .{ .b = try lib.upload(device, std.mem.sliceAsBytes(m.biases)) } }; +} + +/// Codes below 64, so every matrix also exists at 6/32 with the same values; `full`: codes over their whole range. +fn mat(n: usize, k: usize, w: [2]usize, seed: u64, full: bool) !lib.Mat { + return lib.Mat.initBelow(n, k, @intCast(w[0]), w[1], seed, @min(@as(u32, if (full) 256 else 64), @as(u32, 1) << @intCast(w[0]))); +} + +/// Gate, up, shared gate, shared up, down and shared down at the routed and shared widths, and their 6/32 twins. +const Set = struct { + e: [18]Buf, + e6: [18]Buf, + m: [6]lib.Mat, // gate, up, shared gate, shared up, down, shared down + + fn deinit(s: Set) void { + for (s.e) |b| b.b.deinit(); + if (s.e6[0].b.id != s.e[0].b.id) for (s.e6) |b| b.b.deinit(); + for (s.m) |m| { + m.deinit(); + a.free(m.codes); + } + } +}; + +fn set(device: mtl.Device, routed: [2]usize, shared: [2]usize, seed: u64, full: bool) !Set { + const shapes = [_][3]usize{ .{ E * N, D, 0 }, .{ E * N, D, 0 }, .{ N, D, 1 }, .{ N, D, 1 }, .{ E * D, N, 0 }, .{ D, N, 1 } }; + var s: Set = undefined; + for (shapes, 0..) |sh, j| { + const m = try mat(sh[0], sh[1], if (sh[2] == 1) shared else routed, seed + j, full); + const b = try up(device, m); + const m6 = if (full) m else try m.recode(6, 32); // full-range 8-bit codes have no 6/32 twin + defer if (!full) m6.deinit(); + const b6 = if (full) b else try up(device, m6); + for (0..3) |k| { + s.e[j * 3 + k] = b[k]; + s.e6[j * 3 + k] = b6[k]; + } + s.m[j] = m; + } + return s; +} + +const Io = struct { x: lib.Inputs, lg: mtl.Buffer }; + +fn inputs(device: mtl.Device) !Io { + const lg = try device.buffer(R * 513 * 4, lib.opts); + const v = lg.slice(f32, R * 513); + for (0..R) |r| for (0..513) |e| { + v[r * 513 + e] = if (e < E) 4.0 - @as(f32, @floatFromInt((e * 7 + r * 3) % E)) * 0.25 else if (e == 512) 0.5 else -100; + }; + return .{ .x = try lib.Inputs.init(device, R, D, 0x3e), .lg = lg }; +} + +const Out = struct { act: mtl.Buffer, pick: mtl.Buffer, wts: mtl.Buffer, y: mtl.Buffer }; + +fn outputs(device: mtl.Device) !Out { + return .{ .act = try lib.poisoned(device, R * 11 * N * 2), .pick = try lib.poisoned(device, R * 10 * 4), .wts = try lib.poisoned(device, R * 10 * 4), .y = try lib.poisoned(device, R * 11 * D * 2) }; +} + +/// Every window of `rows` rows through gate/up then down, at `w` (null: the recorded 6/32 kernels on `e`). +fn experts(r: *Run, io: Io, e: [18]Buf, rows: usize, o: Out, w: ?variants.Experts) !void { + const cb = r.queue.commandBuffer(); + r.enc = cb.compute(.serial); + var start: usize = 0; + while (start + rows <= R) : (start += rows) { + const rb = try r.buffer(16); + rb.slice(i32, 1)[0] = @intCast(rows); + r.rows = rows; + const x: Buf = .{ .b = io.x.buf, .off = start * D * 2 }; + const lg: Buf = .{ .b = io.lg, .off = start * 513 * 4 }; + const act: Buf = .{ .b = o.act, .off = start * 11 * N * 2 }; + const pick: Buf = .{ .b = o.pick, .off = start * 40 }; + const wts: Buf = .{ .b = o.wts, .off = start * 40 }; + const y: Buf = .{ .b = o.y, .off = start * 11 * D * 2 }; + if (w) |ww| try variants.experts(r, "qa_expert_gateup@moe.gate", "qa_expert_down_y@moe.down", x, lg, e, act, pick, wts, .{ .b = rb }, y, ww) else { + try r.call("qa_expert_gateup@moe.gate", &.{ x, lg, e[0], e[1], e[2], e[3], e[4], e[5], e[6], e[7], e[8], e[9], e[10], e[11] }, &.{ act, pick, wts }); + try r.call("qa_expert_down_y@moe.down", &.{ act, pick, e[12], e[13], e[14], e[15], e[16], e[17], .{ .b = rb } }, &.{y}); + } + } + try lib.commit(cb, r.enc); +} + +fn same4(t: *lib.Tally, label: []const u8, n: usize, want: Out, got: Out) bool { + const rows = R / n * n; + var ok = lib.sameRows(t, label, n, want.act.slice(u8, R * 11 * N * 2), got.act.slice(u8, R * 11 * N * 2), rows, 11 * N * 2); + ok = lib.sameRows(t, label, n, want.pick.slice(u8, R * 40), got.pick.slice(u8, R * 40), rows, 40) and ok; + ok = lib.sameRows(t, label, n, want.wts.slice(u8, R * 40), got.wts.slice(u8, R * 40), rows, 40) and ok; + return lib.sameRows(t, label, n, want.y.slice(u8, R * 11 * D * 2), got.y.slice(u8, R * 11 * D * 2), rows, 11 * D * 2) and ok; +} + +fn expertCheck(t: *lib.Tally, r: *Run, io: Io, routed: [2]usize, shared: [2]usize, seed: u64, full: bool) !void { + var label_buf: [128]u8 = undefined; + const label = try std.fmt.bufPrint(&label_buf, "experts routed q{d}/g{d} shared q{d}/g{d}{s}", .{ routed[0], routed[1], shared[0], shared[1], if (full) " full-range codes" else "" }); + const s = try set(r.device, routed, shared, seed, full); + defer s.deinit(); + const w: variants.Experts = .{ .gu = .{ routed[0], routed[1], shared[0], shared[1] }, .down = .{ routed[0], routed[1], shared[0], shared[1] } }; + const one = try outputs(r.device); + try experts(r, io, s.e, 1, one, w); + // 6 and 8 bits read 8 codes a lane step, other widths 16 (0.6.6's lane_values): only the former sum as 6/32 does + const as_recorded = !full and (routed[0] == 6 or routed[0] == 8) and (shared[0] == 6 or shared[0] == 8); + var same = true; + for (counts) |n| { + const got = try outputs(r.device); + try experts(r, io, s.e, n, got, w); + if (as_recorded) { + const rec = try outputs(r.device); + try experts(r, io, s.e6, n, rec, null); + same = same4(t, label, n, rec, got) and same; + } + same = same4(t, label, n, one, got) and same; + } + if (same) std.debug.print("RESULT {s}: every row's bits at 1-16 rows{s}\n", .{ label, if (as_recorded) ", and the recorded 6/32 kernels' bits on the same values" else "" }); + try gateUp(t, label, s, io, one); + // the down products against float64 from the activations the gate/up launch wrote + var worst = [2]lib.Worst{ .{}, .{} }; // against the right scales, and the negative control's 1% off + for ([_]f64{ 1, 1.01 }, &worst) |skew, *wo| { + const routed_terms = try s.m[4].skewed(skew).dense(); + const shared_terms = try s.m[5].skewed(skew).dense(); + defer for (routed_terms ++ shared_terms) |v| a.free(v); + const act = one.act.slice(u16, R * 11 * N); + const pick = one.pick.slice(u32, R * 10); + const y = one.y.slice(u16, R * 11 * D); + const xa = try a.alloc(f64, N); + defer a.free(xa); + for (0..R) |row| for (0..11) |slot| { + for (xa, 0..) |*v, i| v.* = lib.fromBf(act[(row * 11 + slot) * N + i]); + for (0..D) |d| { + const ref = if (slot < 10) lib.dot(xa, routed_terms, pick[row * 10 + slot] * D + d) else lib.dot(xa, shared_terms, d); + const got = lib.fromBf(y[(row * 11 + slot) * D + d]); + const g = if (slot < 10) routed[1] else shared[1]; + wo.add(d, got, ref[0], lib.ratio(got, ref[0], ref[1], 32 + N / g + 32, true)); + } + }; + } + var down_label: [160]u8 = undefined; + const dl = try std.fmt.bufPrint(&down_label, "{s} down", .{label}); + worst[0].report(t, dl); + worst[1].control(t, dl); +} + +/// Four rows' activations silu(x Wg) * (x Wu) against float64: sums' bounds through the activation and roundings. +fn gateUp(t: *lib.Tally, label: []const u8, s: Set, io: Io, one: Out) !void { + var worst = [2]lib.Worst{ .{}, .{} }; // against the right scales, and the negative control's 1% off + const act = one.act.slice(u16, R * 11 * N); + const pick = one.pick.slice(u32, R * 10); + for ([_]f64{ 1, 1.01 }, &worst) |skew, *w| for (0..4) |row| for (0..11) |slot| for (0..N) |n| { + const x = io.x.row(row); + const base: usize = if (slot < 10) pick[row * 10 + slot] * N else 0; + const gm = (if (slot < 10) s.m[0] else s.m[2]).skewed(skew); + const um = (if (slot < 10) s.m[1] else s.m[3]).skewed(skew); + var gv: [2]f64 = .{ 0, 0 }; // the product and its magnitude sum |x| (s q + |b|) + var uv: [2]f64 = .{ 0, 0 }; + for (0..D) |i| { + const gq = gm.skew * lib.fromBf(gm.scales[(base + n) * (D / gm.group) + i / gm.group]) * @as(f64, @floatFromInt(gm.code(base + n, i))); + const uq = um.skew * lib.fromBf(um.scales[(base + n) * (D / um.group) + i / um.group]) * @as(f64, @floatFromInt(um.code(base + n, i))); + gv[0] += x[i] * gm.deq(base + n, i); + gv[1] += @abs(x[i]) * (gq + @abs(lib.fromBf(gm.biases[(base + n) * (D / gm.group) + i / gm.group]))); + uv[0] += x[i] * um.deq(base + n, i); + uv[1] += @abs(x[i]) * (uq + @abs(lib.fromBf(um.biases[(base + n) * (D / um.group) + i / um.group]))); + } + const sig = 1 / (1 + @exp(-gv[0])); + const silu = gv[0] * sig; + const want = silu * uv[0]; + const steps: f64 = @floatFromInt(32 + D / gm.group + 16); + const dg = steps * lib.unit * gv[1] + 0x1p-8 * @abs(gv[0]); + const du = steps * lib.unit * uv[1] + 0x1p-8 * @abs(uv[0]); + const tol = 1.1 * @abs(uv[0]) * dg + @abs(silu) * du + 0x1p-7 * @abs(want) + 1e-6; + const got = lib.fromBf(act[(row * 11 + slot) * N + n]); + w.add(n, got, want, @abs(got - want) / tol); + }; + var gu_label: [160]u8 = undefined; + const gl = try std.fmt.bufPrint(&gu_label, "{s} gate/up", .{label}); + worst[0].report(t, gl); + worst[1].control(t, gl); +} + +/// The embedding rows (4 stream copies, and the MTP head's one) and the n-gram lookups at width `w`. +fn lookups(t: *lib.Tally, r: *Run, w: [2]usize, seed: u64, full: bool) !void { + var label_buf: [96]u8 = undefined; + const label = try std.fmt.bufPrint(&label_buf, "embed and n-gram tables q{d}/g{d}{s}", .{ w[0], w[1], if (full) " full-range codes" else "" }); + var c: [2]variants.Const = undefined; + const k = variants.width(&c, "BITS", "GS", w); + const emb = try mat(64, D, w, seed, full); + const e = try up(r.device, emb); + const e6 = if (full) e else try up(r.device, try emb.recode(6, 32)); + const ids = try r.buffer(R * 16 * 4); + for (ids.slice(u32, R * 16), 0..) |*v, i| v.* = @intCast((i * 37 + 5) % 64); + var tabs: [26]Buf = undefined; + var tabs6: [26]Buf = undefined; + tabs[0] = .{ .b = ids }; + const starts = try r.buffer(32); + for (starts.slice(u32, 8), 0..) |*v, g| v.* = @intCast(8 * g); + tabs[1] = .{ .b = starts }; + var ple_mats: [8]lib.Mat = undefined; + const tables: usize = if (160 % w[1] == 0) 3 else 2; // an n-gram row is 160 wide: groups of 32 only + for (0..8) |g| { + ple_mats[g] = try mat(8, 160, .{ w[0], if (tables == 3) w[1] else 32 }, seed + 1 + g, full); + const b = try up(r.device, ple_mats[g]); + const b6 = if (full) b else try up(r.device, try ple_mats[g].recode(6, 32)); + for (0..3) |j| { + tabs[2 + 3 * g + j] = b[j]; + tabs6[2 + 3 * g + j] = b6[j]; + } + } + tabs6[0] = tabs[0]; + tabs6[1] = tabs[1]; + defer { + for (e) |b| b.b.deinit(); + if (!full) for (e6) |b| b.b.deinit(); + } + const roles = [_][]const u8{ "qa_embed_rows@embed", "mtp:qa_embed_rows@embed", "qa_ple_lookup@ple" }; + const sizes = [_]usize{ 4 * D * 2, D * 2, 16 * 160 * 2 }; + var same = true; + var worst = lib.Worst{}; + var ctrl = lib.Worst{}; + var ones: [3]mtl.Buffer = undefined; // each role's one-row outputs: a full-range pass's rows compare with them + for (roles[0..tables], sizes[0..tables], 0..) |role, size, which| for (counts) |n| { + const got = try lib.poisoned(r.device, R * size); + const rec = try lib.poisoned(r.device, R * size); + const cb = r.queue.commandBuffer(); + r.enc = cb.compute(.serial); + var start: usize = 0; + while (start + n <= R) : (start += n) { + r.rows = n; + const id: Buf = .{ .b = ids, .off = start * @as(usize, if (which == 2) 16 * 4 else 4) }; + const out: Buf = .{ .b = got, .off = start * size }; + const out6: Buf = .{ .b = rec, .off = start * size }; + if (which < 2) { + try variants.callAs(r, role, n, &.{ id, e[0], e[1], e[2] }, &.{out}, k); + try r.callAs(role, n, &.{ id, e6[0], e6[1], e6[2] }, &.{out6}); + } else { + var tw = tabs; + var tw6 = tabs6; + tw[0] = id; + tw6[0] = id; + try variants.callAs(r, role, n, &tw, &.{out}, k); + try r.callAs(role, n, &tw6, &.{out6}); + } + } + try lib.commit(cb, r.enc); + if (n == 1) ones[which] = got; + const want_rows = if (full) ones[which] else rec; + same = lib.sameRows(t, label, n, want_rows.slice(u8, R * size), got.slice(u8, R * size), R / n * n, size) and same; + if (n != 1) continue; + // the values against s q + b in float64: two bf16 roundings at most + const v = got.slice(u16, R * size / 2); + for (0..R) |row| for (0..size / 2) |i| for ([_]f64{ 1, 1.01 }) |skew| { + const want = if (which == 2) blk: { + const head = i / 160; + const id = ids.slice(u32, R * 16)[row * 16 + head]; + break :blk ple_mats[id / 8].skewed(skew).deq(id % 8, i % 160); + } else emb.skewed(skew).deq(ids.slice(u32, R * 16)[row], i % D); + const g = lib.fromBf(v[row * size / 2 + i]); + (if (skew == 1) &worst else &ctrl).add(i, g, want, @abs(g - want) / (@abs(want) * 0x1p-7 + 1e-6)); + }; + }; + if (same) std.debug.print("RESULT {s}: {s} at 1-16 rows\n", .{ label, if (full) "every row's bits" else "the recorded 6/32 kernels' bits on the same values" }); + worst.report(t, label); + ctrl.control(t, label); +} + +pub fn run(t: *lib.Tally, r: *Run) !void { + const io = try inputs(r.device); + var seed: u64 = 0xe000; + for ([_]u8{ 2, 3, 4, 5, 6, 8 }) |bits| for ([_]usize{ 32, 64, 128 }) |group| { + seed += 8; + try expertCheck(t, r, io, .{ bits, group }, .{ 8, 128 }, seed, false); + try lookups(t, r, .{ bits, group }, seed + 4, false); + }; + try expertCheck(t, r, io, .{ 4, 32 }, .{ 4, 32 }, 0xe801, false); // the MTP head's experts in oQ2 + // 8-bit codes past 63 (the embedding and the shared expert are 8-bit in every oQ checkpoint): float64 and rows + for ([_]usize{ 32, 64, 128 }, 0..) |group, i| { + try expertCheck(t, r, io, .{ 8, group }, .{ 8, 128 }, 0xe900 + 8 * i, true); + try lookups(t, r, .{ 8, group }, 0xe904 + 8 * i, true); + } +} diff --git a/zig/tests/flashnext_widths_prompt.zig b/zig/tests/flashnext_widths_prompt.zig new file mode 100644 index 0000000000..43f85c1751 --- /dev/null +++ b/zig/tests/flashnext_widths_prompt.zig @@ -0,0 +1,249 @@ +//! tf-flashnext-widths, prompt products: qmmw_nax at each oQ width, dense, split-K, members and expert gathers. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const lib = @import("flashnext_widths_lib.zig"); +const fz = tf.flashnext_replay; +const pw = fz.prompt_w; +const Buf = fz.Buf; +const Run = fz.Run; +const a = lib.a; + +const wide = 129; // two 64-row tiles and one row more +const windows = [_]usize{ 1, 2, 3, 8, 16, 17, 33, 64, 65 }; + +/// A prompt with only what the any-width launches read: the run and the expert offsets. +fn prompt(r: *Run, offsets: Buf) fz.Prompt { + var p: fz.Prompt = undefined; + p.r = r; + p.skip = 0; + p.b.off = offsets; + return p; +} + +fn weights(device: mtl.Device, m: lib.Mat) ![3]Buf { + return .{ .{ .b = try lib.upload(device, m.words) }, .{ .b = try lib.upload(device, std.mem.sliceAsBytes(m.scales)) }, .{ .b = try lib.upload(device, std.mem.sliceAsBytes(m.biases)) } }; +} + +/// Rows [start, start + m) of x through the dense product at `w`'s member widths, every window of m rows. +fn dense(p: *fz.Prompt, wd: *pw.Wide, w: pw.Weight, ws: [3]Buf, x: lib.Inputs, n: usize, m: usize, y: mtl.Buffer) !void { + const cb = p.r.queue.commandBuffer(); + p.r.enc = cb.compute(.serial); + var start: usize = 0; + while (start + m <= wide) : (start += m) pw.qmm(p, wd, w, .{ .b = x.buf, .off = start * x.k * 2 }, ws, x.k, n, m, .{ .b = y, .off = start * n * 2 }, 0); + try lib.commit(cb, p.r.enc); +} + +/// A fused projection of `members` (rows, bits, group): float64 per member, then every window's rows bit for bit. +fn denseCheck(t: *lib.Tally, p: *fz.Prompt, wd: *pw.Wide, members: []const [3]usize, k: usize, seed: u64) !void { + var label: std.ArrayList(u8) = .empty; + try label.print(a, "prompt dense K {d}:", .{k}); + var mats = try a.alloc(lib.Mat, members.len); + var parts = try a.alloc(pw.Part, members.len); + var words: std.ArrayList(u8) = .empty; + var scales: std.ArrayList(u8) = .empty; + var biases: std.ArrayList(u8) = .empty; + var n: usize = 0; + for (members, 0..) |mb, i| { + mats[i] = try lib.Mat.init(mb[0], k, @intCast(mb[1]), mb[2], seed + i); + parts[i] = .{ .rows = mb[0], .bits = @intCast(mb[1]), .group = @intCast(mb[2]) }; + try words.appendSlice(a, mats[i].words); + try scales.appendSlice(a, std.mem.sliceAsBytes(mats[i].scales)); + try biases.appendSlice(a, std.mem.sliceAsBytes(mats[i].biases)); + try label.print(a, " {d}x q{d}/g{d}", .{ mb[0], mb[1], mb[2] }); + n += mb[0]; + const kk = @as(u32, @intCast(mb[1])) << 16 | @as(u32, @intCast(mb[2])); + if (!wd.pipes.contains(kk)) try wd.pipes.put(a, kk, try pw.compile(p.r, @intCast(mb[1]), @intCast(mb[2]))); + } + const d = p.r.device; + const ws: [3]Buf = .{ .{ .b = try lib.upload(d, words.items) }, .{ .b = try lib.upload(d, scales.items) }, .{ .b = try lib.upload(d, biases.items) } }; + const w: pw.Weight = if (members.len == 1) .{ .bits = parts[0].bits, .group = parts[0].group } else .{ .parts = parts }; + const x = try lib.Inputs.init(d, wide, k, seed ^ 0x9e); + const all = try lib.poisoned(d, wide * n * 2); + try dense(p, wd, w, ws, x, n, wide, all); + var worst = lib.Worst{}; + var ctrl = lib.Worst{}; + const got = all.slice(u16, wide * n); + var col: usize = 0; + for (mats) |m| { + // each weight rounds to bf16 before its fp32 product (qmmw's dequantw): the reference rounds it the same way + const terms = try m.denseBf(); + defer for (terms) |v| a.free(v); + const wrong = try m.skewed(1.01).denseBf(); + defer for (wrong) |v| a.free(v); + for (0..m.n) |j| for (0..wide) |r| { + const g = lib.fromBf(got[r * n + col + j]); + const ref = lib.dot3(x.row(r), terms, j); + worst.add(col + j, g, ref[0], lib.ratio3(g, ref, k + 64, true)); + const cref = lib.dot3(x.row(r), wrong, j); + ctrl.add(col + j, g, cref[0], lib.ratio3(g, cref, k + 64, true)); + }; + col += m.n; + } + worst.report(t, label.items); + ctrl.control(t, label.items); + var same = true; + for (windows) |m| { + const y = try lib.poisoned(d, wide * n * 2); + try dense(p, wd, w, ws, x, n, m, y); + same = lib.sameRows(t, label.items, m, all.slice(u8, wide * n * 2), y.slice(u8, wide * n * 2), wide / m * m, n * 2) and same; + } + if (same) std.debug.print("RESULT {s}: every row's bits equal in windows of 1, 2, 3, 8, 16, 17, 33, 64 and 65 rows\n", .{label.items}); +} + +/// The generic text at 6/32 against the recorded tf_qmm6_t_nax on the same weights: the same bits. +fn recordedTwin(t: *lib.Tally, p: *fz.Prompt, wd: *pw.Wide) !void { + const d = p.r.device; + const m = try lib.Mat.init(192, 2560, 6, 32, 0x632); + const ws = try weights(d, m); + const x = try lib.Inputs.init(d, wide, 2560, 0x633); + try wd.pipes.put(a, 6 << 16 | 32, try pw.compile(p.r, 6, 32)); + const gen = try lib.poisoned(d, wide * 192 * 2); + try dense(p, wd, .{ .bits = 6, .group = 32 }, ws, x, 192, wide, gen); + const text = try tf.flashnext_replay.frags.source(d, a, pw.ks.flashnext_qmm6); + const rlib = try mtl.Library.fromSource(d, text, mtl.CompileOptions.mlx()); + const pipe = try mtl.Pipeline.init(d, rlib, "tf_qmm6_t_nax", false); + const rec = try lib.poisoned(d, wide * 192 * 2); + const cb = p.r.queue.commandBuffer(); + p.r.enc = cb.compute(.serial); + p.bind(pipe, &.{ ws[0], ws[1], ws[2], .{ .b = x.buf } }); + const prm = [4]i32{ 2560, 192, wide, 0 }; + p.r.enc.setBytes(std.mem.asBytes(&prm), 4); + p.r.enc.setBuffer(rec, 0, 5); + p.r.enc.dispatchThreads(mtl.Size.of(3 * 128, 3, 1), mtl.Size.of(128, 1, 1)); + try lib.commit(cb, p.r.enc); + if (lib.sameRows(t, "prompt recorded qmm6 [192, 2560] against the generic text at 6/32", wide, rec.slice(u8, wide * 384), gen.slice(u8, wide * 384), wide, 384)) + std.debug.print("RESULT prompt qmm6 [192, 2560]: the generic text's bits equal the recorded kernel's\n", .{}); +} + +/// The hyper-connection's down projection as the prompt splits it: 8 K parts of fp32 partials, summed against float64. +fn splitCheck(t: *lib.Tally, p: *fz.Prompt, wd: *pw.Wide, bits: u8, group: usize, seed: u64) !void { + var label_buf: [96]u8 = undefined; + const label = try std.fmt.bufPrint(&label_buf, "prompt split-K [324, 10240] q{d}/g{d}", .{ bits, group }); + const n = 324; + const k = 10240; + const d = p.r.device; + const m = try lib.Mat.init(n, k, bits, group, seed); + defer m.deinit(); + const ws = try weights(d, m); + const x = try lib.Inputs.init(d, wide, k, seed + 1); + const kk = @as(u32, bits) << 16 | @as(u32, @intCast(group)); + if (!wd.pipes.contains(kk)) try wd.pipes.put(a, kk, try pw.compile(p.r, bits, @intCast(group))); + const part = try lib.poisoned(d, 8 * wide * n * 4); + const cb = p.r.queue.commandBuffer(); + p.r.enc = cb.compute(.serial); + p.bind(wd.of(bits, @intCast(group)).splitk, &.{ ws[0], ws[1], ws[2], .{ .b = x.buf } }); + const prm = [4]i32{ k, n, wide, 8 }; + p.r.enc.setBytes(std.mem.asBytes(&prm), 4); + p.r.enc.setBuffer(part, 0, 5); + p.r.enc.dispatchThreads(mtl.Size.of(((n + 63) / 64) * 128, (wide + 63) / 64, 8), mtl.Size.of(128, 1, 1)); + try lib.commit(cb, p.r.enc); + const terms = try m.denseBf(); + defer for (terms) |v| a.free(v); + const wrong = try m.skewed(1.01).denseBf(); + defer for (wrong) |v| a.free(v); + const got = part.slice(f32, 8 * wide * n); + var worst = lib.Worst{}; + var ctrl = lib.Worst{}; + for (0..wide) |r| for (0..n) |j| { + var sum: f64 = 0; + for (0..8) |z| sum += got[(z * wide + r) * n + j]; + // each part is an fp32 sum of K / 8 products (any order: at most that many steps), the parts added here exactly + worst.add(j, sum, lib.dot3(x.row(r), terms, j)[0], lib.ratio3(sum, lib.dot3(x.row(r), terms, j), k / 8 + 64, false)); + ctrl.add(j, sum, lib.dot3(x.row(r), wrong, j)[0], lib.ratio3(sum, lib.dot3(x.row(r), wrong, j), k / 8 + 64, false)); + }; + worst.report(t, label); + ctrl.control(t, label); +} + +/// Sorted-expert gathers (rows grouped by expert, 8 experts used): the product and the fused gate/up activation. +fn gatherCheck(t: *lib.Tally, p: *fz.Prompt, wd: *pw.Wide, bits: u8, group: usize, seed: u64) !void { + var label_buf: [96]u8 = undefined; + const label = try std.fmt.bufPrint(&label_buf, "prompt expert gathers [8 x 128, 2560] q{d}/g{d}", .{ bits, group }); + const e_n = 8; + const n = 128; + const k = 2560; + const d = p.r.device; + const gm = try lib.Mat.init(e_n * n, k, bits, group, seed); + const um = try lib.Mat.init(e_n * n, k, bits, group, seed + 1); + defer gm.deinit(); + defer um.deinit(); + const gw = try weights(d, gm); + const uw = try weights(d, um); + const x = try lib.Inputs.init(d, wide, k, seed + 2); + const off = p.b.off.b.slice(i32, 512); + var expert: [wide]usize = undefined; + for (&expert, 0..) |*e, i| e.* = i * e_n / wide; // rows sorted by expert + for (off, 0..) |*o, e| { + o.* = wide; + for (expert, 0..) |x_e, i| if (x_e >= e) { + o.* = @intCast(i); + break; + }; + } + const kk = @as(u32, bits) << 16 | @as(u32, @intCast(group)); + if (!wd.pipes.contains(kk)) try wd.pipes.put(a, kk, try pw.compile(p.r, bits, @intCast(group))); + const pipes = wd.of(bits, @intCast(group)); + const y = try lib.poisoned(d, wide * n * 2); + const act = try lib.poisoned(d, wide * n * 2); + const cb = p.r.queue.commandBuffer(); + p.r.enc = cb.compute(.serial); + pw.gather(p, pipes, 32, .{ .b = x.buf }, &gw, k, n, wide, .{ .b = y }); + pw.gatherGU(p, pipes, 32, .{ .b = x.buf }, &gw, &uw, k, n, wide, .{ .b = act }); + try lib.commit(cb, p.r.enc); + var worst = [2]lib.Worst{ .{}, .{} }; // product, gate/up activation + var ctrl = [2]lib.Worst{ .{}, .{} }; + for ([_]f64{ 1, 1.01 }, 0..) |skew, which| { + const gt = try gm.skewed(skew).denseBf(); + defer for (gt) |v| a.free(v); + const ut = try um.skewed(skew).denseBf(); + defer for (ut) |v| a.free(v); + const w = if (which == 0) &worst else &ctrl; + for (0..wide) |r| for (0..n) |j| { + const g = lib.dot3(x.row(r), gt, expert[r] * n + j); + const u = lib.dot3(x.row(r), ut, expert[r] * n + j); + const got = lib.fromBf(y.slice(u16, wide * n)[r * n + j]); + w[0].add(j, got, g[0], lib.ratio3(got, g, k + 64, true)); + const sig = 1 / (1 + @exp(-g[0])); + const want = g[0] * sig * u[0]; + // the sums' fp32 and tie bounds, then the gate's bf16 rounding; silu's slope is below 1.1 + const dg = @as(f64, k + 64) * lib.unit * g[1] + g[2] + 0x1p-8 * @abs(g[0]); + const du = @as(f64, k + 64) * lib.unit * u[1] + u[2] + 0x1p-8 * @abs(u[0]); + const tol = 1.1 * @abs(u[0]) * dg + @abs(g[0] * sig) * du + 0x1p-7 * @abs(want) + 1e-6; + const ga = lib.fromBf(act.slice(u16, wide * n)[r * n + j]); + w[1].add(j, ga, want, @abs(ga - want) / tol); + }; + } + var act_label: [128]u8 = undefined; + const gu = try std.fmt.bufPrint(&act_label, "{s} gate/up", .{label}); + worst[0].report(t, label); + ctrl[0].control(t, label); + worst[1].report(t, gu); + ctrl[1].control(t, gu); +} + +pub fn run(t: *lib.Tally, r: *Run) !void { + const offsets = try r.buffer(512 * 4); + var p = prompt(r, .{ .b = offsets }); + const wd = try a.create(pw.Wide); + wd.* = .{}; + try recordedTwin(t, &p, wd); + var seed: u64 = 0x9000; + for ([_]usize{ 2, 3, 4, 5, 6, 8 }) |bits| for ([_]usize{ 32, 64, 128 }) |group| { + seed += 8; + try denseCheck(t, &p, wd, &.{.{ 192, bits, group }}, 2560, seed); + try denseCheck(t, &p, wd, &.{.{ 64, bits, group }}, 640, seed + 4); + try splitCheck(t, &p, wd, @intCast(bits), group, seed + 5); + try gatherCheck(t, &p, wd, @intCast(bits), group, seed + 6); + }; + // the oQ fused stacks, members at their own widths (q, k, v, index_qk and the DeltaNet inputs, rows cut down) + try denseCheck(t, &p, wd, &.{ .{ 256, 3, 64 }, .{ 32, 3, 64 }, .{ 32, 8, 64 }, .{ 64, 3, 64 } }, 2560, 0x9a01); + try denseCheck(t, &p, wd, &.{ .{ 256, 4, 64 }, .{ 64, 5, 128 }, .{ 32, 5, 128 }, .{ 32, 6, 64 } }, 2560, 0x9a02); + try denseCheck(t, &p, wd, &.{ .{ 128, 6, 64 }, .{ 32, 6, 64 }, .{ 32, 6, 128 }, .{ 64, 6, 64 } }, 2560, 0x9a03); + // oQ4e layer 21's DeltaNet inputs: the 5-bit run's b and the 6-bit a keep their 48 rows + try denseCheck(t, &p, wd, &.{ .{ 256, 4, 64 }, .{ 64, 5, 128 }, .{ 48, 5, 128 }, .{ 48, 6, 64 } }, 2560, 0x9a04); + // the prompt's other shapes: HC up (K 320), N 10240, and att.o and gdn.out (K 6144) at their oQ widths + for ([_]u8{ 2, 4, 8 }, 0..) |bits, i| try denseCheck(t, &p, wd, &.{.{ 256, bits, 64 }}, 320, 0x9b00 + i); + for ([_][2]usize{ .{ 3, 64 }, .{ 4, 64 }, .{ 5, 128 }, .{ 8, 64 } }, 0..) |w, i| try denseCheck(t, &p, wd, &.{.{ 64, w[0], w[1] }}, 6144, 0x9b10 + i); + try denseCheck(t, &p, wd, &.{.{ 10240, 3, 64 }}, 320, 0x9b20); +}