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 1/6] 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 2/6] 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 3/6] 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 4/6] 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 5/6] 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 6/6] 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}); +}