From b27726a4dfd2d5865edfa7e9c3112f39e662d222 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:19:36 -1000 Subject: [PATCH 01/18] MLX quantization: an override answers its module in any spelling oQ checkpoints write per-module widths as language_model.model.layers.N..., while loaders ask for model.layers.N... or the bare path, so an exact-key lookup missed every override and read the tensor at the global width. The table keeps each entry under the config's own key, as before, and also indexes it under its canonical module key (no .weight suffix, no known prefix, the rule the 27B and Flash Next families use); `width` takes the key as given, then its module. Two spellings of one module that declare different widths refuse the table as unsupported, the error `detect` already returns, so its callers' error sets are unchanged; the same width twice is fine. Test that fails before: the fixture reads a 5-bit down_proj declared under the oQ prefix by the loader spelling; on main it fails the packed shape check (read at the global width). Also: every prefix form and a trailing .weight name one module; aliases must agree; an override in any spelling answers every spelling and stays under its own key. Gate on main 0efbd447 (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test: `get` logs an error on the DeferredTensor path the test expects). lean_check: 539 problems (main 539), no new findings. --- zig/src/core/quant/keys.zig | 44 +++++++++ zig/src/core/quant/mlx.zig | 174 ++++++++++++++++++++++++++++++++++-- 2 files changed, 212 insertions(+), 6 deletions(-) create mode 100644 zig/src/core/quant/keys.zig diff --git a/zig/src/core/quant/keys.zig b/zig/src/core/quant/keys.zig new file mode 100644 index 0000000000..fe20231e4f --- /dev/null +++ b/zig/src/core/quant/keys.zig @@ -0,0 +1,44 @@ +//! Override keys by module: every spelling checkpoints and loaders give one projection reduces to one canonical key. + +const std = @import("std"); + +/// The prefixes a module path may carry, each stripped at most once and in this order. +const prefixes = [_][]const u8{ "model.language_model.", "language_model.", "text_model.", "model." }; + +/// The module `key` names, without a trailing `.weight` or the known prefixes; a slice of `key`. +pub fn canonical(key: []const u8) []const u8 { + var k = key; + if (std.mem.endsWith(u8, k, ".weight")) k = k[0 .. k.len - ".weight".len]; + for (prefixes) |prefix| { + if (std.mem.startsWith(u8, k, prefix)) k = k[prefix.len..]; + } + return k; +} + +test "every prefix form and a trailing weight name one module" { + for ([_][]const u8{ + "layers.0.mlp.down_proj", + "model.layers.0.mlp.down_proj", + "language_model.model.layers.0.mlp.down_proj", + "language_model.layers.0.mlp.down_proj", + "model.language_model.layers.0.mlp.down_proj", + "model.language_model.model.layers.0.mlp.down_proj", + "text_model.layers.0.mlp.down_proj", + "text_model.model.layers.0.mlp.down_proj", + "layers.0.mlp.down_proj.weight", + "model.layers.0.mlp.down_proj.weight", + "language_model.model.layers.0.mlp.down_proj.weight", + "model.language_model.layers.0.mlp.down_proj.weight", + }) |key| try std.testing.expectEqualStrings("layers.0.mlp.down_proj", canonical(key)); +} + +test "multi-token prediction and vision modules keep their own names" { + for ([_][]const u8{ "mtp.layers.0.mlp.gate_proj", "model.mtp.layers.0.mlp.gate_proj", "language_model.mtp.layers.0.mlp.gate_proj", "language_model.model.mtp.layers.0.mlp.gate_proj.weight" }) |key| + try std.testing.expectEqualStrings("mtp.layers.0.mlp.gate_proj", canonical(key)); + try std.testing.expectEqualStrings("vision_tower.blocks.0.attn.qkv", canonical("vision_tower.blocks.0.attn.qkv.weight")); + try std.testing.expectEqualStrings("visual.blocks.0.attn.qkv", canonical("model.visual.blocks.0.attn.qkv")); + try std.testing.expectEqualStrings("lm_head", canonical("language_model.lm_head")); + // a prefix only counts at the start and a suffix only at the end + try std.testing.expectEqualStrings("layers.0.weight_proj", canonical("layers.0.weight_proj")); + try std.testing.expectEqualStrings("vision_tower.model.x", canonical("vision_tower.model.x")); +} diff --git a/zig/src/core/quant/mlx.zig b/zig/src/core/quant/mlx.zig index f2334eb01a..4b8a289821 100644 --- a/zig/src/core/quant/mlx.zig +++ b/zig/src/core/quant/mlx.zig @@ -4,6 +4,7 @@ const std = @import("std"); const types = @import("types.zig"); const slice = @import("slice.zig"); const convert = @import("convert.zig"); +const keys = @import("keys.zig"); const Tensor = types.Tensor; const Allocator = std.mem.Allocator; @@ -18,19 +19,36 @@ pub const bit_widths = [_]u8{ 2, 3, 4, 5, 6, 8 }; /// One affine width: groups of `group` weights share a scale and a bias, `bits` per weight. pub const Width = struct { bits: u8, group: u16 }; -/// The config's `quantization` table: the global width and the per-tensor overrides (mixed-width conversions). +/// The config's `quantization` table: the global width and the per-module overrides (mixed-width conversions). pub const Config = struct { global: Width, - /// A per-tensor entry; `null` marks one this format refuses, which fails when the tensor is looked up. + /// A per-tensor entry under the config's own key; `null` marks one this format refuses, which fails on lookup. overrides: std.StringHashMapUnmanaged(?Width) = .empty, + /// The same entries under their module's canonical key, so every spelling of a module finds its width. + modules: std.StringHashMapUnmanaged(?Width) = .empty, - /// The tensor's own width when the table names it as an object, else the global one. + /// The width the table names `key` with, else its module under any spelling, else the global one. pub fn width(q: Config, key: []const u8) error{UnsupportedQuantization}!Width { - if (q.overrides.get(key)) |w| return w orelse error.UnsupportedQuantization; + const named = q.overrides.get(key) orelse q.modules.get(keys.canonical(key)); + if (named) |w| return w orelse error.UnsupportedQuantization; return q.global; } + + /// One entry under its own key and its module's; an alias already there must declare the same width. + fn add(q: *Config, a: Allocator, key: []const u8, w: ?Width) DetectError!void { + try q.overrides.put(a, key, w); + const slot = try q.modules.getOrPut(a, keys.canonical(key)); + if (slot.found_existing and !std.meta.eql(slot.value_ptr.*, w)) { + std.log.warn("{s}: another spelling of module {s} declares a different width", .{ key, slot.key_ptr.* }); + return error.UnsupportedQuantization; + } + slot.value_ptr.* = w; + } }; +/// Why a table is refused: allocation, or two spellings of one module that declare different widths (unsupported). +pub const DetectError = Allocator.Error || error{UnsupportedQuantization}; + /// The affine width an object of the config names, or null when it is not one the kernels run. fn affineWidth(o: std.json.ObjectMap) ?Width { const bits = intOrNull(o.get("bits")) orelse return null; @@ -51,11 +69,11 @@ fn intOrNull(v: ?std.json.Value) ?i64 { } /// The checkpoint's `quantization` object when it is MLX affine (`mode` absent or "affine"), else null. -pub fn detect(a: Allocator, src: types.Sources) Allocator.Error!?Config { +pub fn detect(a: Allocator, src: types.Sources) DetectError!?Config { const obj = src.quantization; var q: Config = .{ .global = affineWidth(obj) orelse return null }; var it = obj.iterator(); - while (it.next()) |kv| if (kv.value_ptr.* == .object) try q.overrides.put(a, kv.key_ptr.*, affineWidth(kv.value_ptr.object)); + while (it.next()) |kv| if (kv.value_ptr.* == .object) try q.add(a, kv.key_ptr.*, affineWidth(kv.value_ptr.object)); return q; } @@ -328,3 +346,147 @@ test "the config's per-tensor widths override the global one" { const gptq = (try std.json.parseFromSliceLeaky(std.json.Value, a, "{\"group_size\": 64, \"bits\": 4, \"mode\": \"gptq\"}", .{})).object; try std.testing.expect((try detect(a, .{ .root = gptq, .quantization = gptq })) == null); } + +/// `codes` packed `bits` wide, low code first, into little-endian words (a code may cross into the next word). +fn packCodes(bits: u8, codes: []const u8, words: []u32) void { + @memset(words, 0); + for (codes, 0..) |code, i| { + const bit = i * bits; + const v: u64 = @as(u64, code) << @intCast(bit & 31); + words[bit >> 5] |= @truncate(v); + if ((bit & 31) + bits > 32) words[(bit >> 5) + 1] |= @truncate(v >> 32); + } +} + +/// A checkpoint as `read` sees it: tensors by name, and the config's widths by the key the loader asks for. +const Fixture = struct { + q: Config, + names: []const []const u8, + tensors: []const Tensor, + + pub fn has(f: Fixture, name: []const u8) bool { + return f.find(name) != null; + } + + pub fn get(f: Fixture, name: []const u8) error{MissingTensor}!Tensor { + return f.tensors[f.find(name) orelse return error.MissingTensor]; + } + + pub fn width(f: Fixture, key: []const u8) error{UnsupportedQuantization}!Width { + return f.q.width(key); + } + + fn find(f: Fixture, name: []const u8) ?usize { + for (f.names, 0..) |n, i| if (std.mem.eql(u8, n, name)) return i; + return null; + } +}; + +test "an override under the oQ prefix reads a 5-bit tensor the loader names its own way" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const text = + \\{"group_size": 64, "bits": 4, "mode": "affine", + \\ "language_model.model.layers.0.mlp.down_proj": {"group_size": 64, "bits": 5, "mode": "affine"}} + ; + const obj = (try std.json.parseFromSliceLeaky(std.json.Value, a, text, .{})).object; + const q = (try detect(a, .{ .root = obj, .quantization = obj })).?; + // one row of 64 codes 5 bits wide: ten words and one group; the global 4 bits would need eight words + var codes: [64]u8 = undefined; + for (&codes, 0..) |*c, i| c.* = @intCast(i % 32); + var words: [10]u32 = undefined; + packCodes(5, &codes, &words); + var scale = [_]u16{0x3F80}; // bf16 1.0 + var bias = [_]u16{0x0000}; + const fx: Fixture = .{ + .q = q, + .names = &.{ "model.layers.0.mlp.down_proj.weight", "model.layers.0.mlp.down_proj.scales", "model.layers.0.mlp.down_proj.biases" }, + .tensors = &.{ + .{ .dtype = .u32, .rank = 2, .shape = .{ 1, 10, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(&words) }, + .{ .dtype = .bf16, .rank = 2, .shape = .{ 1, 1, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(&scale) }, + .{ .dtype = .bf16, .rank = 2, .shape = .{ 1, 1, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(&bias) }, + }, + }; + const h = try read(a, fx, "model.layers.0.mlp.down_proj"); + try std.testing.expectEqual(Width{ .bits = 5, .group = 64 }, Width{ .bits = h.bits, .group = h.group }); + var out: [64]f64 = undefined; + try reference(h, &out); + for (out, codes) |v, c| try std.testing.expectEqual(@as(f64, @floatFromInt(c)), v); +} + +/// The table `text` declares, parsed into `a`. +fn tableOf(a: Allocator, text: []const u8) DetectError!?Config { + const obj = (std.json.parseFromSliceLeaky(std.json.Value, a, text, .{}) catch unreachable).object; + return detect(a, .{ .root = obj, .quantization = obj }); +} + +test "an override declared in any spelling answers every spelling of its module" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const spellings = [_][]const u8{ + "layers.0.mlp.down_proj", + "model.layers.0.mlp.down_proj", + "language_model.model.layers.0.mlp.down_proj", + "model.language_model.layers.0.mlp.down_proj", + "model.language_model.model.layers.0.mlp.down_proj", + "text_model.layers.0.mlp.down_proj", + "model.layers.0.mlp.down_proj.weight", + "language_model.model.layers.0.mlp.down_proj.weight", + }; + for (spellings) |declared| { + const text = try std.fmt.allocPrint(a, "{{\"group_size\": 64, \"bits\": 4, \"{s}\": {{\"group_size\": 32, \"bits\": 3}}}}", .{declared}); + const q = (try tableOf(a, text)).?; + for (spellings) |asked| try std.testing.expectEqual(Width{ .bits = 3, .group = 32 }, try q.width(asked)); + // the config's own key stays in `overrides`, for loaders that read the map by it + try std.testing.expectEqual(@as(?Width, .{ .bits = 3, .group = 32 }), q.overrides.get(declared).?); + try std.testing.expectEqual(Width{ .bits = 4, .group = 64 }, try q.width("model.layers.1.mlp.down_proj")); + } + const mtp = (try tableOf(a, "{\"group_size\": 64, \"bits\": 4, \"language_model.mtp.fc\": {\"group_size\": 64, \"bits\": 8}}")).?; + try std.testing.expectEqual(Width{ .bits = 8, .group = 64 }, try mtp.width("mtp.fc.weight")); + try std.testing.expectEqual(Width{ .bits = 4, .group = 64 }, try mtp.width("fc")); +} + +test "aliases of one module must agree on its width" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + try std.testing.expectError(error.UnsupportedQuantization, tableOf(a, + \\{"group_size": 64, "bits": 4, "language_model.model.layers.0.mlp.up_proj": {"group_size": 64, "bits": 5}, + \\ "model.layers.0.mlp.up_proj.weight": {"group_size": 64, "bits": 6}} + )); + try std.testing.expectError(error.UnsupportedQuantization, tableOf(a, + \\{"group_size": 64, "bits": 4, "language_model.model.layers.0.mlp.up_proj": {"group_size": 64, "bits": 5}, + \\ "layers.0.mlp.up_proj": {"group_size": 32, "bits": 5}} + )); + // an unsupported entry and a width for the same module disagree too + try std.testing.expectError(error.UnsupportedQuantization, tableOf(a, + \\{"group_size": 64, "bits": 4, "language_model.model.layers.0.mlp.up_proj": {"group_size": 64, "bits": 5}, + \\ "model.layers.0.mlp.up_proj": {"group_size": 64, "bits": 5, "mode": "mxfp4"}} + )); + const same = (try tableOf(a, + \\{"group_size": 64, "bits": 4, "language_model.model.layers.0.mlp.up_proj": {"group_size": 64, "bits": 5}, + \\ "model.layers.0.mlp.up_proj": {"group_size": 64, "bits": 5, "mode": "affine"}} + )).?; + try std.testing.expectEqual(Width{ .bits = 5, .group = 64 }, try same.width("layers.0.mlp.up_proj")); +} + +test "an entry the kernels do not run refuses every spelling of its module" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const q = (try tableOf(a, + \\{"group_size": 64, "bits": 4, "language_model.model.layers.0.mlp.gate_proj": {"group_size": 64, "bits": 4, "mode": "mxfp4"}, + \\ "model.layers.0.mlp.gate_proj": {"group_size": 16, "bits": 4}, + \\ "language_model.model.layers.1.mlp.gate_proj": {"group_size": 64, "bits": 7}} + )).?; + for ([_][]const u8{ "layers.0.mlp.gate_proj", "model.layers.0.mlp.gate_proj.weight", "model.language_model.layers.0.mlp.gate_proj" }) |key| + try std.testing.expectError(error.UnsupportedQuantization, q.width(key)); + try std.testing.expectError(error.UnsupportedQuantization, q.width("model.layers.1.mlp.gate_proj")); + try std.testing.expectEqual(Width{ .bits = 4, .group = 64 }, try q.width("model.layers.2.mlp.gate_proj")); +} + +test { + _ = keys; +} From 7f1c3eed66da0aa5eb7b7ab88a36cc9ee320d02e Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:18:51 -1000 Subject: [PATCH 02/18] MLX affine weights widen exactly: codes re-packed at more bits, group tables repeated for a smaller group A projection of mixed-width members (the Flash Next hyper-connection down and inject pair, where 0.6.6's QWeights.of widens) reads as one matrix at their common width, the most bits and the smallest group. Narrowing and a group that does not divide the source group are refused. Used only where 0.6.6 widens; no caller on this commit. Tests (widen.zig, new): every source width and every wider target dequantize to the same fp64 values, bit for bit; a stack of experts widens row for row; the same width comes back as it is; narrowing bits, growing groups and widths the kernels do not run are refused; the common width takes the most bits and the smallest group; fused members keep each one's values in order. Gate on main 0efbd447 (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test: `get` logs an error on the DeferredTensor path the test expects). lean_check: 539 problems (main 539), no new findings. --- zig/src/core/quant/quant.zig | 2 + zig/src/core/quant/widen.zig | 321 +++++++++++++++++++++++++++++++++++ 2 files changed, 323 insertions(+) create mode 100644 zig/src/core/quant/widen.zig diff --git a/zig/src/core/quant/quant.zig b/zig/src/core/quant/quant.zig index 944bca7a3d..756fd87b90 100644 --- a/zig/src/core/quant/quant.zig +++ b/zig/src/core/quant/quant.zig @@ -10,6 +10,7 @@ pub const dense = @import("dense.zig"); /// FP8 and NVFP4 values and fp64 references; host only (no Format: their kernels' layouts are each backend's). pub const fp8 = @import("fp8.zig"); pub const nvfp4 = @import("nvfp4.zig"); +pub const widen = @import("widen.zig"); const Allocator = std.mem.Allocator; @@ -202,4 +203,5 @@ test { _ = dense; _ = fp8; _ = nvfp4; + _ = widen; } diff --git a/zig/src/core/quant/widen.zig b/zig/src/core/quant/widen.zig new file mode 100644 index 0000000000..1a26a56330 --- /dev/null +++ b/zig/src/core/quant/widen.zig @@ -0,0 +1,321 @@ +//! Exact widening of MLX affine projections: codes re-packed at more bits, group tables repeated for smaller groups. + +const std = @import("std"); +const mlx = @import("mlx.zig"); +const types = @import("types.zig"); + +const Allocator = std.mem.Allocator; +const Host = mlx.Host; +const Width = mlx.Width; +const Tensor = types.Tensor; + +pub const Error = error{ + /// The target has fewer bits than the source: codes would not fit. + NarrowsBits, + /// The target group does not divide the source group: one scale could not cover a target group. + GroupDoesNotDivide, + /// The target is not an affine width the kernels run (bits 2, 3, 4, 5, 6, 8; groups 32, 64, 128). + UnsupportedQuantization, + /// The words and tables disagree on shape, or fused members on K or table type. + UnexpectedTensor, +} || Allocator.Error; + +/// The width a set of members widens to exactly: the most bits and the smallest group among them. +pub fn common(widths: []const Width) Width { + std.debug.assert(widths.len > 0); + var w = widths[0]; + for (widths[1..]) |x| w = .{ .bits = @max(w.bits, x.bits), .group = @min(w.group, x.group) }; + return w; +} + +fn supported(w: Width) bool { + const bits_ok = std.mem.indexOfScalar(u8, &mlx.bit_widths, w.bits) != null; + return bits_ok and (w.group == 32 or w.group == 64 or w.group == 128); +} + +/// Rows of a host with its leading dimensions folded (a stack's experts times N), after checking the three agree. +fn foldedRows(h: Host) error{UnexpectedTensor}!usize { + const r = h.words.rank; + if (r < 2 or h.scales.rank != r or h.biases.rank != r) return error.UnexpectedTensor; + if (h.scales.dtype != h.biases.dtype) return error.UnexpectedTensor; + var n: usize = 1; + for (0..r - 1) |i| { + if (h.scales.shape[i] != h.words.shape[i] or h.biases.shape[i] != h.words.shape[i]) return error.UnexpectedTensor; + n *= h.words.shape[i]; + } + if (h.scales.shape[r - 1] != h.biases.shape[r - 1]) return error.UnexpectedTensor; + if (h.words.shape[r - 1] * 32 != h.scales.shape[r - 1] * h.group * h.bits) return error.UnexpectedTensor; + return n; +} + +/// The same dequantized values stored at width `to`: never fewer bits, and a group that divides the source's. +pub fn widen(a: Allocator, h: Host, to: Width) Error!Host { + if (!supported(to)) return error.UnsupportedQuantization; + const n = try foldedRows(h); + if (to.bits == h.bits and to.group == h.group) return h; + if (to.bits < h.bits) return error.NarrowsBits; + if (h.group % to.group != 0) return error.GroupDoesNotDivide; + const reps = h.group / to.group; + return .{ + .words = if (to.bits == h.bits) h.words else try repack(a, h.words, n, h.bits, to.bits), + .scales = if (reps == 1) h.scales else try repeat(a, h.scales, n, reps), + .biases = if (reps == 1) h.biases else try repeat(a, h.biases, n, reps), + .bits = to.bits, + .group = to.group, + }; +} + +/// Words of `bits`-wide codes re-packed at `wider` bits, code for code, low code first (`n` folded rows). +fn repack(a: Allocator, words: Tensor, n: usize, bits: u8, wider: u8) Allocator.Error!Tensor { + const r = words.rank; + const from_row = words.shape[r - 1]; + const k = from_row * 32 / bits; + const to_row = k * wider / 32; + const out = try a.alloc(u8, n * to_row * 4); + @memset(out, 0); + const mask: u64 = (@as(u64, 1) << @intCast(bits)) - 1; + for (0..n) |row| for (0..k) |c| { + const code = codeAt(words.bytes, row * from_row, c * bits) & mask; + const bit = c * wider; + const at = row * to_row + (bit >> 5); + const v = code << @intCast(bit & 31); + orWord(out, at, @truncate(v)); + if ((bit & 31) + wider > 32) orWord(out, at + 1, @truncate(v >> 32)); + }; + var t = words; + t.shape[r - 1] = to_row; + t.bytes = out; + return t; +} + +/// The (unmasked) code starting `bit` bits into the row of words beginning at word `base`. +fn codeAt(bytes: []const u8, base: usize, bit: usize) u64 { + const at = base + (bit >> 5); + const lo: u64 = std.mem.readInt(u32, bytes[4 * at ..][0..4], .little); + const hi: u64 = if ((bit & 31) != 0 and 4 * (at + 2) <= bytes.len) std.mem.readInt(u32, bytes[4 * (at + 1) ..][0..4], .little) else 0; + return (lo | hi << 32) >> @intCast(bit & 31); +} + +fn orWord(bytes: []u8, at: usize, v: u32) void { + const w = bytes[4 * at ..][0..4]; + std.mem.writeInt(u32, w, std.mem.readInt(u32, w, .little) | v, .little); +} + +/// A group table with every entry repeated `reps` times along its last dimension (`n` folded rows). +fn repeat(a: Allocator, t: Tensor, n: usize, reps: usize) Allocator.Error!Tensor { + const size = t.dtype.size(); + const groups = t.shape[t.rank - 1]; + const out = try a.alloc(u8, n * groups * reps * size); + for (0..n * groups) |g| { + const src = t.bytes[g * size ..][0..size]; + for (0..reps) |j| @memcpy(out[(g * reps + j) * size ..][0..size], src); + } + var o = t; + o.shape[t.rank - 1] = groups * reps; + o.bytes = out; + return o; +} + +/// One rank-2 projection of `hosts`' rows in order, all widened to their `common` width; K and tables must agree. +pub fn fuse(a: Allocator, hosts: []const Host) Error!Host { + if (hosts.len == 0) return error.UnexpectedTensor; + var to: Width = .{ .bits = hosts[0].bits, .group = hosts[0].group }; + for (hosts[1..]) |h| to = common(&.{ to, .{ .bits = h.bits, .group = h.group } }); + const parts = try a.alloc(Host, hosts.len); + var n: usize = 0; + for (hosts, parts) |h, *p| { + if (h.words.rank != 2 or h.scales.dtype != hosts[0].scales.dtype) return error.UnexpectedTensor; + if (h.scales.shape[1] * h.group != hosts[0].scales.shape[1] * hosts[0].group) return error.UnexpectedTensor; + p.* = try widen(a, h, to); + n += p.words.shape[0]; + } + if (hosts.len == 1) return parts[0]; + var out = parts[0]; + out.words = try concatRows(a, parts, .words, n); + out.scales = try concatRows(a, parts, .scales, n); + out.biases = try concatRows(a, parts, .biases, n); + return out; +} + +fn concatRows(a: Allocator, parts: []const Host, comptime field: std.meta.FieldEnum(Host), n: usize) Allocator.Error!Tensor { + var total: usize = 0; + for (parts) |p| total += @field(p, @tagName(field)).bytes.len; + const out = try a.alloc(u8, total); + var at: usize = 0; + for (parts) |p| { + const b = @field(p, @tagName(field)).bytes; + @memcpy(out[at..][0..b.len], b); + at += b.len; + } + var t = @field(parts[0], @tagName(field)); + t.shape[0] = n; + t.bytes = out; + return t; +} + +// ---- tests ---- + +const testing = std.testing; +const groups_all = [_]u16{ 32, 64, 128 }; + +fn tensor(dtype: types.DType, shape: []const usize, bytes: []const u8) Tensor { + var t: Tensor = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .bytes = bytes }; + @memcpy(t.shape[0..shape.len], shape); + return t; +} + +/// A host of `lead` x `k` codes of width `w` with bf16 tables: row 0 all zeros, row 1 all max, the rest random. +fn sample(a: Allocator, rng: std.Random, lead: []const usize, k: usize, w: Width) !Host { + var n: usize = 1; + for (lead) |d| n *= d; + const row = k * w.bits / 32; + const words = try a.alloc(u8, n * row * 4); + @memset(words, 0); + const max: u64 = (@as(u64, 1) << @intCast(w.bits)) - 1; + for (0..n) |r| for (0..k) |c| { + const code: u64 = switch (r % 4) { + 0 => 0, + 1 => max, + else => rng.int(u64) & max, + }; + const bit = c * w.bits; + const v = code << @intCast(bit & 31); + orWord(words, r * row + (bit >> 5), @truncate(v)); + if ((bit & 31) + w.bits > 32) orWord(words, r * row + (bit >> 5) + 1, @truncate(v >> 32)); + }; + const g = k / w.group; + const scales = try a.alloc(u8, n * g * 2); + const biases = try a.alloc(u8, n * g * 2); + for (0..n * g) |i| { + std.mem.writeInt(u16, scales[2 * i ..][0..2], bf16Of(rng.floatNorm(f32) * 0.02), .little); + std.mem.writeInt(u16, biases[2 * i ..][0..2], bf16Of(rng.floatNorm(f32) * 0.1), .little); + } + var ws: [5]usize = undefined; + var ts: [5]usize = undefined; + @memcpy(ws[0..lead.len], lead); + @memcpy(ts[0..lead.len], lead); + ws[lead.len] = row; + ts[lead.len] = g; + const r = lead.len + 1; + return .{ + .words = tensor(.i32, ws[0..r], words), + .scales = tensor(.bf16, ts[0..r], scales), + .biases = tensor(.bf16, ts[0..r], biases), + .bits = w.bits, + .group = w.group, + }; +} + +fn bf16Of(x: f32) u16 { + return @truncate(@as(u32, @bitCast(x)) >> 16); +} + +fn dequantized(a: Allocator, h: Host) ![]f64 { + const out = try a.alloc(f64, mlx.elements(h)); + try mlx.reference(h, out); + return out; +} + +fn expectBits(want: []const f64, got: []const f64) !void { + try testing.expectEqual(want.len, got.len); + for (want, got) |x, y| try testing.expectEqual(@as(u64, @bitCast(x)), @as(u64, @bitCast(y))); +} + +test "every width widens to every wider one with the same fp64 values, bit for bit" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var prng = std.Random.DefaultPrng.init(0x51de); + const rng = prng.random(); + var checked: usize = 0; + for (mlx.bit_widths) |bits| for (groups_all) |group| { + const from: Width = .{ .bits = bits, .group = group }; + const h = try sample(a, rng, &.{6}, 256, from); + const want = try dequantized(a, h); + for (mlx.bit_widths) |to_bits| for (groups_all) |to_group| { + if (to_bits < bits or group % to_group != 0) continue; + const to: Width = .{ .bits = to_bits, .group = to_group }; + const w = try widen(a, h, to); + try testing.expectEqual(to_bits, w.bits); + try testing.expectEqual(to_group, w.group); + try testing.expectEqual(@as(usize, 256) * to_bits / 32, w.words.shape[1]); + try testing.expectEqual(@as(usize, 256) / to_group, w.scales.shape[1]); + try expectBits(want, try dequantized(a, w)); + checked += 1; + }; + }; + // 18 sources; each reaches the wider-or-equal bits times the groups dividing its own + try testing.expectEqual(@as(usize, 126), checked); +} + +test "a stack of experts widens row for row" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var prng = std.Random.DefaultPrng.init(0xe0e1); + const rng = prng.random(); + const h = try sample(a, rng, &.{ 3, 4 }, 384, .{ .bits = 3, .group = 128 }); + const w = try widen(a, h, .{ .bits = 5, .group = 32 }); + try testing.expectEqual(@as(u8, 3), w.words.rank); + try testing.expectEqualSlices(usize, &.{ 3, 4, 384 * 5 / 32 }, w.words.shape[0..3]); + try testing.expectEqualSlices(usize, &.{ 3, 4, 12 }, w.biases.shape[0..3]); + try expectBits(try dequantized(a, h), try dequantized(a, w)); +} + +test "the same width is returned as it is" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var prng = std.Random.DefaultPrng.init(1); + const h = try sample(a, prng.random(), &.{2}, 128, .{ .bits = 4, .group = 64 }); + const w = try widen(a, h, .{ .bits = 4, .group = 64 }); + try testing.expectEqual(h.words.bytes.ptr, w.words.bytes.ptr); + try testing.expectEqual(h.scales.bytes.ptr, w.scales.bytes.ptr); +} + +test "narrowing bits, growing groups and widths the kernels do not run are refused" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var prng = std.Random.DefaultPrng.init(2); + const rng = prng.random(); + const h = try sample(a, rng, &.{2}, 256, .{ .bits = 5, .group = 64 }); + try testing.expectError(error.NarrowsBits, widen(a, h, .{ .bits = 4, .group = 64 })); + try testing.expectError(error.NarrowsBits, widen(a, h, .{ .bits = 2, .group = 32 })); + try testing.expectError(error.GroupDoesNotDivide, widen(a, h, .{ .bits = 5, .group = 128 })); + try testing.expectError(error.GroupDoesNotDivide, widen(a, h, .{ .bits = 8, .group = 128 })); + try testing.expectError(error.UnsupportedQuantization, widen(a, h, .{ .bits = 7, .group = 64 })); + try testing.expectError(error.UnsupportedQuantization, widen(a, h, .{ .bits = 8, .group = 16 })); + try testing.expectError(error.UnsupportedQuantization, widen(a, h, .{ .bits = 8, .group = 48 })); + var bad = h; + bad.bits = 4; + try testing.expectError(error.UnexpectedTensor, widen(a, bad, .{ .bits = 8, .group = 32 })); +} + +test "the common width takes the most bits and the smallest group" { + try testing.expectEqual(Width{ .bits = 5, .group = 32 }, common(&.{ .{ .bits = 4, .group = 64 }, .{ .bits = 5, .group = 128 }, .{ .bits = 2, .group = 32 } })); + try testing.expectEqual(Width{ .bits = 6, .group = 32 }, common(&.{.{ .bits = 6, .group = 32 }})); +} + +test "fused members keep each one's values in order at the common width" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var prng = std.Random.DefaultPrng.init(0xf05e); + const rng = prng.random(); + const k = 512; + const members = [_]Host{ + try sample(a, rng, &.{4}, k, .{ .bits = 4, .group = 64 }), + try sample(a, rng, &.{3}, k, .{ .bits = 5, .group = 128 }), + try sample(a, rng, &.{5}, k, .{ .bits = 2, .group = 32 }), + }; + const f = try fuse(a, &members); + try testing.expectEqual(@as(u8, 5), f.bits); + try testing.expectEqual(@as(u16, 32), f.group); + try testing.expectEqual(@as(usize, 12), f.words.shape[0]); + var want: std.ArrayList(f64) = .empty; + for (members) |m| try want.appendSlice(a, try dequantized(a, m)); + try expectBits(want.items, try dequantized(a, f)); + const other = try sample(a, rng, &.{2}, 256, .{ .bits = 4, .group = 64 }); + try testing.expectError(error.UnexpectedTensor, fuse(a, &.{ members[0], other })); +} From dcebc462b94a9c563c1d8164b2494d45411ac829 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 04:08:50 -1000 Subject: [PATCH 03/18] The row projection reads 3- and 5-bit MLX codes and can write into a wider output Chips without tensor units now decode every MLX affine width (2, 3, 4, 5, 6, 8 bits) in the row kernel: 3- and 5-bit rows unpack eight codes from three or five bytes, low code first, as MLX packs them. The 2-, 4-, 6- and 8-bit paths are unchanged. A 3- or 5-bit row must end on a word, as the checkpoint stores it. callStrided launches the row kernel with an output row stride, so one member of a fused stack (one run of equal width) writes its columns straight into the stack output. The default call, its dims and its kernels are unchanged. tf-row-widths-check (zig/build/core_checks.zig) runs it on the GPU: each width at groups of 32, 64 and 128 against an fp64 product from the MLX reference dequantization, within the fp32-sum bound; one row alone against the same row in 2-, 3-, 8- and 16-row windows, relu2, indexed slots and a strided output, bit for bit; a fused stack of 48-column members at mixed widths (4-bit among them, group 64) written through the stride at windows up to 48 rows, each member as its own call; and, given the unchanged kernel source, the widths it already served byte for byte against it. Tests that fail before: "shapes the row kernels serve" now validates 3 and 5 bits (main refuses them), "3- and 5-bit rows take whole words of codes", "every width has its own pipeline slot" and "a strided call writes rows at least a matrix wide" (no callStrided on main). tf-row-widths-check on an M3 Ultra (no tensor units), given main's row_projection.metal: PASS 64512 values within the fp32-sum bound (worst err/bound 0.9400), 408192 bit-equal pairs, 301056 outputs byte-equal to main's kernels at the widths main serves (3 and 5 bits have no main kernel to compare). Gate on main 0efbd447 (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test: `get` logs an error on the DeferredTensor path the test expects). lean_check: 536 problems (main 539), no new findings. --- build.zig | 2 + zig/build/core_checks.zig | 18 ++ zig/kernels/metal/core/row_projection.metal | 56 +++- zig/src/core/row_projection.zig | 50 ++- zig/tests/row_widths_check.zig | 324 ++++++++++++++++++++ 5 files changed, 439 insertions(+), 11 deletions(-) create mode 100644 zig/build/core_checks.zig create mode 100644 zig/tests/row_widths_check.zig diff --git a/build.zig b/build.zig index 1bf8bfa260..2ba6052374 100644 --- a/build.zig +++ b/build.zig @@ -3,6 +3,7 @@ const builtin = @import("builtin"); const cuda_build = @import("zig/build/cuda.zig"); const dist_build = @import("zig/build/dist.zig"); const qwen27_build = @import("zig/build/qwen27.zig"); +const core_checks = @import("zig/build/core_checks.zig"); const hip_build = @import("zig/build/hip.zig"); const grammar_build = @import("zig/build/grammar.zig"); const counters_build = @import("zig/build/metal_counters.zig"); @@ -455,6 +456,7 @@ fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "row_projection", .module = row_mod } }, }) }); b.step("tf-row-check", "The core row projection on synthetic 4-bit matrices: widths bit-equal, a CPU reference, indexed experts").dependOn(&b.addInstallArtifact(row_check, .{}).step); + core_checks.rowWidths(b, target, optimize, metal, row_mod); test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = row_mod })).step); exl3_build.targets(b, target, optimize, metal, test_step); exlr8_build.metalTargets(b, target, optimize, metal, mods.sources, test_step); diff --git a/zig/build/core_checks.zig b/zig/build/core_checks.zig new file mode 100644 index 0000000000..75f39c00c2 --- /dev/null +++ b/zig/build/core_checks.zig @@ -0,0 +1,18 @@ +//! The core's device checks on synthetic weights: programs that run on the GPU, each under its own step. + +const std = @import("std"); + +/// tf-row-widths-check: the row projection at every MLX width and group against float64, windows, strides, stock bytes. +pub fn rowWidths(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, metal: *std.Build.Module, row_mod: *std.Build.Module) void { + const tokenizer = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = optimize, .link_libc = true }); + const core = b.createModule(.{ .root_source_file = b.path("zig/src/core/root.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "tokenizer", .module = tokenizer }} }); + const exe = b.addExecutable(.{ .name = "tf-row-widths-check", .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/row_widths_check.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "row_projection", .module = row_mod }, .{ .name = "core", .module = core } }, + }) }); + const about = "The row projection at every MLX width and group: fp64 reference, windows, strides, stock bytes"; + b.step("tf-row-widths-check", about).dependOn(&b.addInstallArtifact(exe, .{}).step); +} diff --git a/zig/kernels/metal/core/row_projection.metal b/zig/kernels/metal/core/row_projection.metal index dae395aeb0..5a503e3acc 100644 --- a/zig/kernels/metal/core/row_projection.metal +++ b/zig/kernels/metal/core/row_projection.metal @@ -1,4 +1,4 @@ -// Affine 2-, 4-, 6- and 8-bit row projections for chips without tensor units: RPS columns a simdgroup, MB rows a weight read, each row its own fp32 sums. +// Affine 2- to 8-bit row projections without tensor units: RPS columns a simdgroup, MB rows a weight read. #include using namespace metal; @@ -18,7 +18,7 @@ inline float tf_rp_load16(const device bfloat* x, thread float* xt) { return sum; } -// MB rows' dots with one weight word group (16 codes; 6-bit: four values in three bytes), the codes unpacked to floats once; each row's sums in one order. +// MB rows' dots with 16 codes (3, 5, 6 bits: 8, 8, 4 codes in 3, 5, 3 bytes), unpacked once; sums in one order. template inline void tf_rp_qdot16(const device uint8_t* w, const thread float (*xt)[16], float scale, float bias, const thread float* sum, thread float* out) { @@ -40,6 +40,20 @@ inline void tf_rp_qdot16(const device uint8_t* w, const thread float (*xt)[16], q[4 * i + 2] = float(((p[1] >> 4) & 0x0f) | ((p[2] & 0x03) << 4)); q[4 * i + 3] = float((p[2] >> 2) & 0x3f); } + } else if (BITS == 3) { + for (int i = 0; i < 2; i++) { + const device uint8_t* p = w + 3 * i; + const uint v = uint(p[0]) | (uint(p[1]) << 8) | (uint(p[2]) << 16); + for (int j = 0; j < 8; j++) q[8 * i + j] = float((v >> (3 * j)) & 7u); + } + } else if (BITS == 5) { + for (int i = 0; i < 2; i++) { + const device uint8_t* p = w + 5 * i; + const uint lo = uint(p[0]) | (uint(p[1]) << 8) | (uint(p[2]) << 16) | (uint(p[3]) << 24), hi = p[4]; + for (int j = 0; j < 6; j++) q[8 * i + j] = float((lo >> (5 * j)) & 31u); + q[8 * i + 6] = float(((lo >> 30) | (hi << 2)) & 31u); + q[8 * i + 7] = float(hi >> 3); + } } else { for (int i = 0; i < 16; i++) q[i] = float(w[i]); } @@ -93,8 +107,8 @@ inline bfloat tf_rp_out(float v) { return bfloat(v); } -// dims: K, N, GS, rows. X [rows, K] bf16; W [N, K*BITS/32] words; S, B [N, K/GS] bf16; OUT [rows, N] bf16. -template +// dims: K, N, GS, rows, OUT's row stride when STRIDED (else N). X [rows, K]; W [N, K*BITS/32] words; S, B [N, K/GS]. +template [[kernel]] void tf_row_projection(const device bfloat* X [[buffer(0)]], const device uint32_t* W [[buffer(1)]], const device bfloat* S [[buffer(2)]], @@ -104,7 +118,7 @@ template uint simdgroup_index_in_threadgroup [[simdgroup_index_in_threadgroup]], uint thread_index_in_simdgroup [[thread_index_in_simdgroup]], uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]]) { - const int K = dims[0], N = dims[1], GS = dims[2], rows = dims[3]; + const int K = dims[0], N = dims[1], GS = dims[2], rows = dims[3], LD = STRIDED ? dims[4] : N; const uint lane = thread_index_in_simdgroup; const int row0 = (int(threadgroup_position_in_grid.x) * SG + int(simdgroup_index_in_threadgroup)) * RPS; if (row0 >= N) return; @@ -119,14 +133,14 @@ template tf_rp_rows(wr, sr, br, X + size_t(m) * K, K, GS, lane, acc); if (lane == 0) for (int b = 0; b < MB; b++) - for (int j = 0; j < RPS; j++) OUT[size_t(m + b) * N + row0 + j] = tf_rp_out(acc[b][j]); + for (int j = 0; j < RPS; j++) OUT[size_t(m + b) * LD + row0 + j] = tf_rp_out(acc[b][j]); } #pragma clang loop unroll(disable) for (; m < rows; m++) { float acc[1][RPS]; tf_rp_rows(wr, sr, br, X + size_t(m) * K, K, GS, lane, acc); if (lane == 0) - for (int j = 0; j < RPS; j++) OUT[size_t(m) * N + row0 + j] = tf_rp_out(acc[0][j]); + for (int j = 0; j < RPS; j++) OUT[size_t(m) * LD + row0 + j] = tf_rp_out(acc[0][j]); } } @@ -195,3 +209,31 @@ template [[host_name("tf_row_projection_q8_bf16")]] [[kernel]] decltype(tf_row_p template [[host_name("tf_row_projection_indexed_q8_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 8, true>) tf_row_projection_indexed<4, 2, false, 8, true>; template [[host_name("tf_row_projection_relu2_q8_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 8, true>) tf_row_projection<4, 2, 2, true, 8, true>; template [[host_name("tf_row_projection_indexed_relu2_q8_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 8, true>) tf_row_projection_indexed<4, 2, true, 8, true>; +template [[host_name("tf_row_projection_q3_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 3, false>) tf_row_projection<4, 2, 2, false, 3, false>; +template [[host_name("tf_row_projection_indexed_q3_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 3, false>) tf_row_projection_indexed<4, 2, false, 3, false>; +template [[host_name("tf_row_projection_relu2_q3_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 3, false>) tf_row_projection<4, 2, 2, true, 3, false>; +template [[host_name("tf_row_projection_indexed_relu2_q3_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 3, false>) tf_row_projection_indexed<4, 2, true, 3, false>; +template [[host_name("tf_row_projection_q3_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 3, true>) tf_row_projection<4, 2, 2, false, 3, true>; +template [[host_name("tf_row_projection_indexed_q3_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 3, true>) tf_row_projection_indexed<4, 2, false, 3, true>; +template [[host_name("tf_row_projection_relu2_q3_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 3, true>) tf_row_projection<4, 2, 2, true, 3, true>; +template [[host_name("tf_row_projection_indexed_relu2_q3_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 3, true>) tf_row_projection_indexed<4, 2, true, 3, true>; +template [[host_name("tf_row_projection_q5_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 5, false>) tf_row_projection<4, 2, 2, false, 5, false>; +template [[host_name("tf_row_projection_indexed_q5_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 5, false>) tf_row_projection_indexed<4, 2, false, 5, false>; +template [[host_name("tf_row_projection_relu2_q5_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 5, false>) tf_row_projection<4, 2, 2, true, 5, false>; +template [[host_name("tf_row_projection_indexed_relu2_q5_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 5, false>) tf_row_projection_indexed<4, 2, true, 5, false>; +template [[host_name("tf_row_projection_q5_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 5, true>) tf_row_projection<4, 2, 2, false, 5, true>; +template [[host_name("tf_row_projection_indexed_q5_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 5, true>) tf_row_projection_indexed<4, 2, false, 5, true>; +template [[host_name("tf_row_projection_relu2_q5_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 5, true>) tf_row_projection<4, 2, 2, true, 5, true>; +template [[host_name("tf_row_projection_indexed_relu2_q5_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 5, true>) tf_row_projection_indexed<4, 2, true, 5, true>; +template [[host_name("tf_row_projection_strided_q2_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 2, false, true>) tf_row_projection<4, 2, 2, false, 2, false, true>; +template [[host_name("tf_row_projection_strided_q2_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 2, true, true>) tf_row_projection<4, 2, 2, false, 2, true, true>; +template [[host_name("tf_row_projection_strided_q3_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 3, false, true>) tf_row_projection<4, 2, 2, false, 3, false, true>; +template [[host_name("tf_row_projection_strided_q3_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 3, true, true>) tf_row_projection<4, 2, 2, false, 3, true, true>; +template [[host_name("tf_row_projection_strided_q4_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 4, false, true>) tf_row_projection<4, 2, 2, false, 4, false, true>; +template [[host_name("tf_row_projection_strided_q4_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 4, true, true>) tf_row_projection<4, 2, 2, false, 4, true, true>; +template [[host_name("tf_row_projection_strided_q5_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 5, false, true>) tf_row_projection<4, 2, 2, false, 5, false, true>; +template [[host_name("tf_row_projection_strided_q5_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 5, true, true>) tf_row_projection<4, 2, 2, false, 5, true, true>; +template [[host_name("tf_row_projection_strided_q6_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 6, false, true>) tf_row_projection<4, 2, 2, false, 6, false, true>; +template [[host_name("tf_row_projection_strided_q6_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 6, true, true>) tf_row_projection<4, 2, 2, false, 6, true, true>; +template [[host_name("tf_row_projection_strided_q8_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 8, false, true>) tf_row_projection<4, 2, 2, false, 8, false, true>; +template [[host_name("tf_row_projection_strided_q8_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 8, true, true>) tf_row_projection<4, 2, 2, false, 8, true, true>; diff --git a/zig/src/core/row_projection.zig b/zig/src/core/row_projection.zig index 2bd1cba34a..d28555029c 100644 --- a/zig/src/core/row_projection.zig +++ b/zig/src/core/row_projection.zig @@ -1,4 +1,4 @@ -//! The lane projection's path for chips without tensor units: row-layout 2-, 4- and 8-bit weights, one read a window, each row its own fp32 sums. +//! The lane projection's path for chips without tensor units: row-layout 2- to 8-bit weights, one read a window. const std = @import("std"); const mtl = @import("metal"); const sources = @import("kernel_sources"); @@ -14,7 +14,7 @@ pub const columns = 4; pub const groups = 2; /// The code widths the kernels unpack. -pub const bit_widths = [_]u8{ 2, 4, 6, 8 }; +pub const bit_widths = [_]u8{ 2, 3, 4, 5, 6, 8 }; /// A group's sum of inputs for the bias term: fp32 in order (the tensor path's xsum), or rounded to bf16 at each add (Nemotron's expert kernels). pub const Sum = enum { f32, bf16 }; @@ -34,7 +34,7 @@ pub const Weights = struct { sum: Sum = .f32, pub fn validate(x: Weights) !void { - if (x.n == 0 or x.n % (columns * groups) != 0 or x.k == 0 or x.k % 16 != 0 or x.group < 16 or x.group % 16 != 0 or x.k % x.group != 0 or 512 % x.group != 0) return error.UnsupportedRowShape; + if (x.n == 0 or x.n % (columns * groups) != 0 or x.k == 0 or x.k % 16 != 0 or x.group < 16 or x.group % 16 != 0 or x.k % x.group != 0 or 512 % x.group != 0 or x.k * x.bits % 32 != 0) return error.UnsupportedRowShape; if (std.mem.indexOfScalar(u8, &bit_widths, x.bits) == null) return error.UnsupportedRowBits; } @@ -63,6 +63,7 @@ pub const Pipelines = struct { relu2: [variants]mtl.Pipeline, indexed: [variants]mtl.Pipeline, indexed_relu2: [variants]mtl.Pipeline, + strided: [variants]mtl.Pipeline, pub fn load(device: mtl.Device) !Pipelines { const pool = mtl.objc.Pool.push(); @@ -76,6 +77,7 @@ pub const Pipelines = struct { p.relu2[i] = try mtl.Pipeline.init(device, lib, std.fmt.comptimePrint("tf_row_projection_relu2_q{d}_{s}", .{ bits, sum }), false); p.indexed[i] = try mtl.Pipeline.init(device, lib, std.fmt.comptimePrint("tf_row_projection_indexed_q{d}_{s}", .{ bits, sum }), false); p.indexed_relu2[i] = try mtl.Pipeline.init(device, lib, std.fmt.comptimePrint("tf_row_projection_indexed_relu2_q{d}_{s}", .{ bits, sum }), false); + p.strided[i] = try mtl.Pipeline.init(device, lib, std.fmt.comptimePrint("tf_row_projection_strided_q{d}_{s}", .{ bits, sum }), false); }; return p; } @@ -85,6 +87,7 @@ pub const Pipelines = struct { for (p.relu2) |x| x.deinit(); for (p.indexed) |x| x.deinit(); for (p.indexed_relu2) |x| x.deinit(); + for (p.strided) |x| x.deinit(); } }; @@ -99,6 +102,14 @@ pub fn call(p: *const Pipelines, w: Weights, rows: usize, relu2: bool) !Call { return .{ .pipeline = if (relu2) p.relu2[i] else p.plain[i], .dims = .{ @intCast(w.k), @intCast(w.n), @intCast(w.group), @intCast(rows), 1, 1 }, .groups = w.n / (columns * groups) }; } +/// `call` with output row m at y + m * `stride` values: a fused stack member writes its columns of the rows. +pub fn callStrided(p: *const Pipelines, w: Weights, rows: usize, stride: usize) !Call { + try w.validate(); + if (rows == 0) return error.NoRows; + if (stride < w.n or stride > std.math.maxInt(i32)) return error.UnsupportedRowStride; + return .{ .pipeline = p.strided[w.variant()], .dims = .{ @intCast(w.k), @intCast(w.n), @intCast(w.group), @intCast(rows), @intCast(stride), 1 }, .groups = w.n / (columns * groups) }; +} + /// y[slots, n] = x[slot / repeat] W[ids[slot]]^T; the grid's y is the slot. pub fn callIndexed(p: *const Pipelines, x: Experts, slots: usize, relu2: bool) !Call { try x.validate(); @@ -127,5 +138,36 @@ test "shapes the row kernels serve" { try (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 256, .k = 2688, .bits = 8 }).validate(); try (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 256, .k = 2688, .bits = 2 }).validate(); try (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 256, .k = 2688, .bits = 6 }).validate(); - try std.testing.expectError(error.UnsupportedRowBits, (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 256, .k = 2688, .bits = 3 }).validate()); + try (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 256, .k = 2688, .bits = 3 }).validate(); + try (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 256, .k = 2688, .bits = 5 }).validate(); + try std.testing.expectError(error.UnsupportedRowBits, (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 256, .k = 2688, .bits = 7 }).validate()); +} + +test "3- and 5-bit rows take whole words of codes" { + for ([_]u8{ 3, 5 }) |bits| for ([_]usize{ 32, 64, 128 }) |group| { + try (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 256, .k = 2688, .group = group, .bits = bits }).validate(); + try (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 5120, .k = 17408, .group = group, .bits = bits }).validate(); + }; + // 48 codes of 3 or 5 bits end inside a word: the checkpoint cannot hold that row + try std.testing.expectError(error.UnsupportedRowShape, (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 256, .k = 48, .group = 16, .bits = 3 }).validate()); + try std.testing.expectError(error.UnsupportedRowShape, (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 256, .k = 48, .group = 16, .bits = 5 }).validate()); + try (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 256, .k = 48, .group = 16, .bits = 6 }).validate(); +} + +test "every width has its own pipeline slot" { + var seen: [variants]bool = @splat(false); + for (bit_widths) |bits| for ([_]Sum{ .f32, .bf16 }) |sum| { + const i = (Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 8, .k = 64, .bits = bits, .sum = sum }).variant(); + try std.testing.expect(!seen[i]); + seen[i] = true; + }; +} + +test "a strided call writes rows at least a matrix wide" { + const p: Pipelines = undefined; + const w = Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = 64, .k = 128, .bits = 5 }; + try std.testing.expectError(error.UnsupportedRowStride, callStrided(&p, w, 2, 63)); + try std.testing.expectError(error.NoRows, callStrided(&p, w, 0, 64)); + const c = try callStrided(&p, w, 2, 192); + try std.testing.expectEqual([6]i32{ 128, 64, 64, 2, 192, 1 }, c.dims); } diff --git a/zig/tests/row_widths_check.zig b/zig/tests/row_widths_check.zig new file mode 100644 index 0000000000..327f7bc3d1 --- /dev/null +++ b/zig/tests/row_widths_check.zig @@ -0,0 +1,324 @@ +//! The core row projection at every MLX width and group: an fp64 reference, windows, strides and stock bytes. +const std = @import("std"); +const mtl = @import("metal"); +const row = @import("row_projection"); +const mlx = @import("core").quant.mlx; + +const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; +const a = std.heap.page_allocator; +const usage = "usage: tf-row-widths-check [STOCK.metal] (STOCK: the unchanged row_projection.metal for the byte comparison)"; + +const groups_tested = [_]usize{ 32, 64, 128 }; +const Shape = struct { n: usize, k: usize }; +// all tail (K < 512), full passes and a tail, full passes only +const shapes = [_]Shape{ .{ .n = 64, .k = 384 }, .{ .n = 128, .k = 2688 }, .{ .n = 32, .k = 1536 } }; +const rows = 16; +// a fused stack's windows reach this many rows +const x_rows = 48; +const sentinel: u16 = 0xffff; + +fn toBf(v: f32) u16 { + const u: u32 = @bitCast(v); + return @truncate((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); +} +fn fromBf(v: u16) f32 { + return @bitCast(@as(u32, v) << 16); +} + +/// One synthetic MLX matrix in the row layout, its inputs and its fp64 dequantization (and codes alone). +const Mat = struct { + w: row.Weights, + x: mtl.Buffer, + deq: []f64, + codes: []f64, + + fn init(device: mtl.Device, s: Shape, bits: u8, group: usize, sum: row.Sum, seed: u64) !Mat { + var prng = std.Random.DefaultPrng.init(seed); + const r = prng.random(); + const words = s.n * s.k * bits / 32; + const tables = s.n * (s.k / group); + const m = Mat{ + .w = .{ .w = try device.buffer(words * 4, opts), .scales = try device.buffer(tables * 2, opts), .biases = try device.buffer(tables * 2, opts), .n = s.n, .k = s.k, .group = group, .bits = bits, .sum = sum }, + .x = try device.buffer(x_rows * s.k * 2, opts), + .deq = try a.alloc(f64, s.n * s.k), + .codes = try a.alloc(f64, s.n * s.k), + }; + for (m.w.w.slice(u32, words)) |*v| v.* = r.int(u32); + for (m.w.scales.slice(u16, tables)) |*v| v.* = toBf(r.float(f32) * 0.02 + 0.001); + for (m.w.biases.slice(u16, tables)) |*v| v.* = toBf((r.float(f32) - 0.5) * 0.1); + for (m.x.slice(u16, x_rows * s.k)) |*v| v.* = toBf((r.float(f32) - 0.5) * 2); + try mlx.reference(host(m.w, m.w.scales.slice(u16, tables), m.w.biases.slice(u16, tables)), m.deq); + const ones = try a.alloc(u16, tables); + defer a.free(ones); + const zeros = try a.alloc(u16, tables); + defer a.free(zeros); + @memset(ones, toBf(1)); + @memset(zeros, 0); + try mlx.reference(host(m.w, ones, zeros), m.codes); + return m; + } + + fn deinit(m: Mat) void { + for ([_]mtl.Buffer{ m.w.w, m.w.scales, m.w.biases, m.x }) |b| b.deinit(); + a.free(m.deq); + a.free(m.codes); + } + + /// The checkpoint's view of the same words and tables, as `mlx.reference` reads it. + fn host(w: row.Weights, scales: []const u16, biases: []const u16) mlx.Host { + const groups = w.k / w.group; + return .{ + .words = .{ .dtype = .i32, .rank = 2, .shape = .{ w.n, w.k * w.bits / 32, 1, 1, 1 }, .bytes = w.w.slice(u8, w.n * w.k * w.bits / 8) }, + .scales = .{ .dtype = .bf16, .rank = 2, .shape = .{ w.n, groups, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(scales) }, + .biases = .{ .dtype = .bf16, .rank = 2, .shape = .{ w.n, groups, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(biases) }, + .bits = w.bits, + .group = @intCast(w.group), + }; + } + + /// Row m against column c in fp64, and the fp32 sums' error bound (output rounding plus k-deep accumulation). + fn reference(m: Mat, i: usize, c: usize) [2]f64 { + const k = m.w.k; + const g_n = k / m.w.group; + const sc = m.w.scales.slice(u16, m.w.n * g_n); + const bi = m.w.biases.slice(u16, m.w.n * g_n); + const x = m.x.slice(u16, rows * k); + var y: f64 = 0; + var terms: f64 = 0; + for (0..g_n) |g| { + var dot: f64 = 0; + var xs: f64 = 0; + for (0..m.w.group) |j| { + const kk = g * m.w.group + j; + const xv: f64 = fromBf(x[i * k + kk]); + y += xv * m.deq[c * k + kk]; + dot += @abs(xv * m.codes[c * k + kk]); + xs += @abs(xv); + } + terms += @abs(@as(f64, fromBf(sc[c * g_n + g]))) * dot + @abs(@as(f64, fromBf(bi[c * g_n + g]))) * xs; + } + const u32_round = std.math.ldexp(@as(f64, 1), -24); + const gamma = @as(f64, @floatFromInt(k + 64)) * u32_round; + return .{ y, std.math.ldexp(@abs(y) + gamma * terms, -8) + gamma * terms }; + } +}; + +/// One launch: buffers x, W, S, B, dims, out at 0-5 and ids at 6 when given; grid (groups, slots). +const Launch = struct { pipeline: mtl.Pipeline, dims: [6]i32, groups: usize, slots: usize = 1, x_off: usize = 0, out_off: usize = 0, ids: ?mtl.Buffer = null }; + +fn run(queue: mtl.Queue, l: Launch, m: Mat, out: mtl.Buffer) !void { + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + e.setPipeline(l.pipeline); + e.setBuffer(m.x, l.x_off, 0); + e.setBuffer(m.w.w, 0, 1); + e.setBuffer(m.w.scales, 0, 2); + e.setBuffer(m.w.biases, 0, 3); + e.setValue(l.dims, 4); + e.setBuffer(out, l.out_off, 5); + if (l.ids) |ids| e.setBuffer(ids, 0, 6); + e.dispatchGroups(mtl.Size.of(l.groups, l.slots, 1), mtl.Size.of(32 * row.groups, 1, 1)); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.debug.print("GPU failed: {s}\n", .{msg}); + return error.GpuFailed; + } +} + +fn of(c: row.Call) Launch { + return .{ .pipeline = c.pipeline, .dims = c.dims, .groups = c.groups, .slots = c.slots }; +} + +/// The stock library's pipelines for the widths it already served, by name. +const Stock = struct { + lib: mtl.Library, + fn get(s: Stock, device: mtl.Device, kind: []const u8, bits: u8, sum: row.Sum) !mtl.Pipeline { + var buf: [96]u8 = undefined; + const name = try std.fmt.bufPrint(&buf, "tf_row_projection{s}_q{d}_{s}", .{ kind, bits, @tagName(sum) }); + return mtl.Pipeline.init(device, s.lib, name, false); + } +}; + +const Tally = struct { worst_err: f64 = 0, worst_ratio: f64 = 0, worst_col: usize = 0, values: usize = 0, exact: usize = 0, pairs: usize = 0, stock: usize = 0 }; + +pub fn main(init: std.process.Init) !void { + const arena = init.arena.allocator(); + const args = try init.minimal.args.toSlice(arena); + if (args.len > 2) { + std.debug.print("{s}\n", .{usage}); + std.process.exit(2); + } + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try row.Pipelines.load(device); + defer p.deinit(); + var stock: ?Stock = null; + if (args.len == 2) { + const f = try mtl.MappedFile.open(try arena.dupeSentinel(u8, args[1], 0)); + defer f.deinit(); + stock = .{ .lib = try mtl.Library.fromSource(device, f.bytes[0..f.size], mtl.CompileOptions.mlx()) }; + } + defer if (stock) |s| s.lib.deinit(); + const out = try device.buffer(rows * 512 * 2 * 4, opts); + defer out.deinit(); + const one = try device.buffer(rows * 512 * 2 * 4, opts); + defer one.deinit(); + const ids = try device.buffer(rows * 4, opts); + defer ids.deinit(); + var all = Tally{}; + for (row.bit_widths) |bits| for (groups_tested) |group| { + var t = Tally{}; + for (shapes, 0..) |s, si| for ([_]row.Sum{ .f32, .bf16 }) |sum| { + const m = try Mat.init(device, s, bits, group, sum, 1000 * @as(u64, bits) + 10 * group + si); + defer m.deinit(); + try check(device, queue, &p, stock, m, out, one, ids, &t); + }; + std.debug.print("q{d} g{d}: {d} values vs fp64, worst |err| {e:.3} (col {d}), worst err/bound {d:.4}, bf16(ref) exact {d}/{d}; {d} bit-equal pairs; {d} stock bytes equal\n", .{ bits, group, t.values, t.worst_err, t.worst_col, t.worst_ratio, t.exact, t.values, t.pairs, t.stock }); + all.values += t.values; + all.pairs += t.pairs; + all.stock += t.stock; + all.worst_ratio = @max(all.worst_ratio, t.worst_ratio); + }; + all.pairs += try stack(device, queue, &p); + if (stock == null) std.debug.print("no stock source given: the byte comparison was not run\n", .{}); + std.debug.print("PASS {d} values within the fp32-sum bound (worst err/bound {d:.4}), {d} bit-equal pairs, {d} stock outputs byte-equal\n", .{ all.values, all.worst_ratio, all.pairs, all.stock }); +} + +/// One matrix: the reference (f32 sums), windows, relu2, indexed and strided outputs against one row alone, then stock. +fn check(device: mtl.Device, queue: mtl.Queue, p: *const row.Pipelines, stock: ?Stock, m: Mat, out: mtl.Buffer, one: mtl.Buffer, ids: mtl.Buffer, t: *Tally) !void { + const n = m.w.n; + const k = m.w.k; + // each row alone: the reference everything else is compared with + const alone = try a.alloc(u16, rows * n); + defer a.free(alone); + for (0..rows) |i| { + try run(queue, launchAt(try row.call(p, m.w, 1, false), i * k * 2, 0), m, one); + @memcpy(alone[i * n ..][0..n], one.slice(u16, n)); + } + if (m.w.sum == .f32) for (0..rows) |i| for (0..n) |c| { + const ref = m.reference(i, c); + const got: f64 = fromBf(alone[i * n + c]); + const err = @abs(got - ref[0]); + if (err > ref[1]) { + std.debug.print("q{d} g{d} n {d} k {d} row {d} col {d}: GPU {d} fp64 {d} bound {e}\n", .{ m.w.bits, m.w.group, n, k, i, c, got, ref[0], ref[1] }); + return error.ReferenceMismatch; + } + if (err > t.worst_err) { + t.worst_err = err; + t.worst_col = c; + } + t.worst_ratio = @max(t.worst_ratio, err / ref[1]); + if (alone[i * n + c] == toBf(@floatCast(ref[0]))) t.exact += 1; + t.values += 1; + }; + // windows of 2, 3, 8 and 16 rows (ending at the last row): each row as it was alone + for ([_]usize{ 2, 3, 8, 16 }) |w| { + const start = rows - w; + try run(queue, launchAt(try row.call(p, m.w, w, false), start * k * 2, 0), m, out); + for (0..w) |i| { + if (!std.mem.eql(u16, out.slice(u16, w * n)[i * n ..][0..n], alone[(start + i) * n ..][0..n])) return error.WindowChangesBits; + t.pairs += n; + } + } + // relu2 is relu squared of the plain output + try run(queue, of(try row.call(p, m.w, rows, true)), m, out); + for (out.slice(u16, rows * n), alone) |r2, v| { + const h = @max(fromBf(v), 0); + if (r2 != toBf(h * h)) return error.Relu2DiffersFromPlain; + } + // indexed slots of one expert: each slot as its row alone; an id past the count is a zero row + for (ids.slice(u32, rows), 0..) |*id, i| id.* = if (i == rows - 1) 1 else 0; + var li = of(try row.callIndexed(p, .{ .one = m.w, .experts = 1 }, rows, false)); + li.ids = ids; + try run(queue, li, m, out); + for (0..rows) |i| { + const got = out.slice(u16, rows * n)[i * n ..][0..n]; + if (i == rows - 1) { + for (got) |v| if (v != 0) return error.MissingExpertNotZero; + } else if (!std.mem.eql(u16, got, alone[i * n ..][0..n])) return error.IndexedDiffersFromDense; + } + // strided: the matrix's rows at column 16 of a wider buffer; its columns as alone, every other value untouched + const stride = 2 * n + 24; + @memset(out.slice(u16, rows * stride), sentinel); + try run(queue, launchAt(try row.callStrided(p, m.w, rows, stride), 0, 16 * 2), m, out); + for (0..rows) |i| for (0..stride) |c| { + const v = out.slice(u16, rows * stride)[i * stride + c]; + if (c >= 16 and c < 16 + n) { + if (v != alone[i * n + c - 16]) return error.StridedDiffersFromAlone; + t.pairs += 1; + } else if (v != sentinel) return error.StridedWritesOutside; + }; + const s = stock orelse return; + if (std.mem.indexOfScalar(u8, &.{ 2, 4, 6, 8 }, m.w.bits) == null) return; + // the stock kernels: plain and relu2 windows of 1, 3 and 16 rows, indexed slots, byte for byte + for ([_]usize{ 1, 3, 16 }) |w| for ([_]bool{ false, true }) |relu2| { + const c = try row.call(p, m.w, w, relu2); + try run(queue, of(c), m, out); + const pipe = try s.get(device, if (relu2) "_relu2" else "", m.w.bits, m.w.sum); + defer pipe.deinit(); + var l = of(c); + l.pipeline = pipe; + try run(queue, l, m, one); + if (!std.mem.eql(u16, out.slice(u16, w * n), one.slice(u16, w * n))) return error.StockBytesDiffer; + t.stock += w * n; + }; + const pipe = try s.get(device, "_indexed", m.w.bits, m.w.sum); + defer pipe.deinit(); + try run(queue, li, m, out); + li.pipeline = pipe; + try run(queue, li, m, one); + if (!std.mem.eql(u16, out.slice(u16, rows * n), one.slice(u16, rows * n))) return error.StockBytesDiffer; + t.stock += rows * n; +} + +fn launchAt(c: row.Call, x_off: usize, out_off: usize) Launch { + var l = of(c); + l.x_off = x_off; + l.out_off = out_off; + return l; +} + +/// A fused stack of 48-column members at mixed widths, group 64: each member's strided columns as its own call. +fn stack(device: mtl.Device, queue: mtl.Queue, p: *const row.Pipelines) !usize { + const widths = [_]u8{ 4, 5, 4, 6, 3, 8, 2, 4 }; + const n = 48; + const stride = widths.len * n; + var members: [widths.len]Mat = undefined; + for (&members, widths, 0..) |*m, bits, i| m.* = try Mat.init(device, .{ .n = n, .k = 2688 }, bits, 64, .f32, 77 + i); + defer for (members) |m| m.deinit(); + // every member reads the same inputs + for (members[1..]) |m| @memcpy(m.x.slice(u16, x_rows * 2688), members[0].x.slice(u16, x_rows * 2688)); + const fused = try device.buffer(x_rows * stride * 2, opts); + defer fused.deinit(); + const own = try device.buffer(x_rows * n * 2, opts); + defer own.deinit(); + const alone = try device.buffer(n * 2, opts); + defer alone.deinit(); + var pairs: usize = 0; + for ([_]usize{ 1, 2, 3, 16, 48 }) |w| { + @memset(fused.slice(u16, w * stride), sentinel); + for (members, 0..) |m, j| try run(queue, launchAt(try row.callStrided(p, m.w, w, stride), 0, j * n * 2), m, fused); + for (members, 0..) |m, j| { + try run(queue, of(try row.call(p, m.w, w, false)), m, own); + for (0..w) |i| { + const got = fused.slice(u16, w * stride)[i * stride + j * n ..][0..n]; + if (!std.mem.eql(u16, got, own.slice(u16, w * n)[i * n ..][0..n])) return error.StackDiffersFromMember; + if (w == 48) { + try run(queue, launchAt(try row.call(p, m.w, 1, false), i * 2688 * 2, 0), m, alone); + if (!std.mem.eql(u16, got, alone.slice(u16, n))) return error.StackDiffersFromRowAlone; + pairs += n; + } + pairs += n; + } + } + for (fused.slice(u16, w * stride)) |v| if (v == sentinel) return error.StackLeftAGap; + } + std.debug.print("stack of {d} members (bits 4 5 4 6 3 8 2 4, n {d}, group 64) at windows 1, 2, 3, 16, 48: each member as its own call\n", .{ widths.len, n }); + return pairs; +} From 26681b8459e2b36f295ebc9a55f9860103976370 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:35:27 -1000 Subject: [PATCH 04/18] Lane projections take every MLX affine width, an output row stride and 16-row members Format accepts 2-, 3- and 5-bit codes beside 4, 6 and 8, in groups of 32, 64 or 128. The non-4-bit path widens 2- and 3-bit codes to nibbles for the bf16 x uint4 op and 5-, 6- and 8-bit codes to bytes for the uint8 op, the staging the Python lane matmul uses, so each width keeps that arithmetic. 3- and 5-bit groups of 32 are an odd number of words a lane; they are fetched a word at a time instead of in pairs, which would drop the last word. Layout.y_stride (default N) sets the elements between output rows in every kernel (tf_lane, tf_lane_pair, tf_lane_reg), so a member of a fused stack writes its columns of the fused rows; the output bound checks follow it. The store of the simdgroup-matrix layout (Macs without tensor units, #679) takes the same stride and keeps its column guard. N may be a multiple of 16 on the plain kernel: the last 32-column tile is packed with zero codes past N, its metadata reads and stores are already guarded by N, and weightWords counts the padded tiles (sourceWords is MLX's W). Cooperative and reg layouts still need N % 32 == 0. encodeSums checks its sums buffer by the group count, as #698 sizes the dispatch. Without tensor units the admitted layout is unchanged: 6-bit group 32 on the plain kernel (#679); every other width still answers TensorUnitsRequired there, now checked at 2, 3 and 5 bits too. With the defaults the 4-, 6- and 8-bit kernels preprocess (metal -E) to the same text as main's for every group, schedule, output, prepared-sums/reg layout and both fragment layouts tried (148 variants), so the uniform paths keep their kernels. Host tests that fail before: 2, 3 and 5 bits validate; y_stride reaches the source and below N is refused; a 48-row member validates (cooperative and reg refuse it); every width x group x {64, 48} rows packs MLX codes bit for bit with zero padding. The device program covers every width, 48-row members and strided output on tensor units; without them (M3 Ultra): PASS 14 layouts, 1904 bit-equal row/window pairs, 208 windows equal to fn_lane's recorded bytes, 7424 CPU values (max roundoff-bound fraction 0.0144), 10 refusals, and a 48-row member and two strided layouts at 6/32 equal the unstrided rows with the gaps untouched. fn_lane's recorded windows are compared only on the layouts they were recorded at. Gate on main 0efbd447 (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test: `get` logs an error on the DeferredTensor path the test expects). lean_check: 536 problems (main 539), no new findings. --- zig/kernels/metal/core/lane_projection.metal | 62 ++++++++- zig/src/core/lane_projection.zig | 134 +++++++++++++------ zig/tests/core_lane_projection.zig | 77 +++++++++-- 3 files changed, 219 insertions(+), 54 deletions(-) diff --git a/zig/kernels/metal/core/lane_projection.metal b/zig/kernels/metal/core/lane_projection.metal index 0c3cf622c1..81cfd5d9d4 100644 --- a/zig/kernels/metal/core/lane_projection.metal +++ b/zig/kernels/metal/core/lane_projection.metal @@ -12,18 +12,29 @@ using namespace mpp::tensor_ops; #ifndef TF_PRECOMPUTE_SUMS #define TF_PRECOMPUTE_SUMS 0 #endif +#ifndef TF_YS +#define TF_YS TF_N +#endif +// Output rows lie TF_YS elements apart, so a member of a fused stack writes its columns of the fused rows. +#define TF_ODD_WORDS (TF_GROUP * TF_BITS / 32 % 2) constant constexpr int NT = 32, GS = TF_GROUP, KG = TF_K / GS, NF = NT / 16; constant constexpr int WP = GS * TF_BITS / 32; +#if TF_ODD_WORDS +struct tf_raw { uint w[WP]; uint4 s[NF]; }; // 3- and 5-bit groups of 32 are an odd word count, loaded a word at a time +#else struct tf_raw { uint2 w[WP / 2]; uint4 s[NF]; }; // a lane's packed code words and its scale/bias quads for one group +#endif inline tf_raw tf_lane_fetch(const device uint* W, const device uint4* SB, int tile, int g, int g_end, int n0, int lane, short fn) { tf_raw raw; const bool ok = g < g_end; const device uint2* src = (const device uint2*)(W + (((size_t)tile * KG + (ok ? g : 0)) * NT + lane) * WP); -#if TF_BITS != 4 +#if TF_ODD_WORDS + for (int i = 0; i < WP; i++) raw.w[i] = ok && n0 + lane < TF_N ? ((const device uint*)src)[i] : 0u; +#elif TF_BITS != 4 for (int i = 0; i < WP / 2; i++) raw.w[i] = ok && n0 + lane < TF_N ? src[i] : uint2(0); #else for (int i = 0; i < WP / 2; i++) raw.w[i] = uint2(0); @@ -53,6 +64,7 @@ inline float tf_lane_xsum(const device bfloat* X, int m, int g, int M) { return acc; } +// Codes: 4-bit straight to the uint4 op; 2/3-bit widened to nibbles for it, 5/6/8-bit to bytes for the uint8 op. #if TF_BITS == 4 #define TF_STAGE_CODES(raw) #define TF_GROUP_PRODUCT(g) \ @@ -62,9 +74,42 @@ inline float tf_lane_xsum(const device bfloat* X, int m, int g, int M) { auto P = op.template get_destination_cooperative_tensor(); \ op.run(a, b, P); #else +#if TF_BITS < 4 +inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, int lane) { + uint words[WP + 1]; +#if TF_ODD_WORDS + for (int i = 0; i < WP; i++) words[i] = raw.w[i]; +#else + for (int i = 0; i < WP / 2; i++) { words[2*i] = raw.w[i].x; words[2*i+1] = raw.w[i].y; } +#endif + words[WP] = 0u; +#if TF_BITS == 3 + for (int c = 0; c < GS / 8; c++) { + const int bit = 24 * c, i = bit >> 5, sh = bit & 31; + uint pack = words[i] >> sh; + if (sh > 8) pack |= words[i + 1] << (32 - sh); + uint nib = 0u; + for (int j = 0; j < 8; j++) nib |= ((pack >> (3 * j)) & 7u) << (4 * j); + stage[lane * (GS / 8) + c] = nib; + } +#else + for (int c = 0; c < GS / 8; c++) { // word c/2's half c%2: 8 codes of 2 bits to 8 nibbles + uint v = (words[c >> 1] >> (16 * (c & 1))) & 0xFFFFu; + v = (v | (v << 8)) & 0x00FF00FFu; + v = (v | (v << 4)) & 0x0F0F0F0Fu; + v = (v | (v << 2)) & 0x33333333u; + stage[lane * (GS / 8) + c] = v; + } +#endif +} +#else inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, int lane) { uint words[WP + 1]; +#if TF_ODD_WORDS + for (int i = 0; i < WP; i++) words[i] = raw.w[i]; +#else for (int i = 0; i < WP / 2; i++) { words[2*i] = raw.w[i].x; words[2*i+1] = raw.w[i].y; } +#endif words[WP] = 0u; for (int c = 0; c < GS / 4; c++) { uint word = 0u; @@ -77,6 +122,7 @@ inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, in stage[lane * (GS / 4) + c] = word; } } +#endif #define TF_STAGE_CODES(raw) tf_stage_codes(stage, raw, lane); #define TF_GROUP_PRODUCT(g) \ simdgroup_barrier(mem_flags::mem_threadgroup); \ @@ -136,7 +182,11 @@ inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, in constexpr auto desc = matmul2d_descriptor(16, NT, GS, false, true, false, matmul2d_descriptor::mode::multiply); matmul2d op; tensor, tensor_inline> tA((device bfloat*)X, dextents(TF_K, M)); -#if TF_BITS != 4 +#if TF_BITS < 4 + threadgroup uint stage_all[TF_SK * NT * (GS / 8)]; + threadgroup uint* stage = stage_all + sg * NT * (GS / 8); + tensor, tensor_inline> b((threadgroup uchar*)stage, dextents(GS, NT)); +#elif TF_BITS != 4 threadgroup uint stage_all[TF_SK * NT * (GS / 4)]; threadgroup uint* stage = stage_all + sg * NT * (GS / 4); tensor, tensor_inline> b((threadgroup uint8_t*)stage, dextents(GS, NT)); @@ -169,7 +219,7 @@ inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, in if (m < M && nn < TF_N) for (int j = 0; j < 4; j++) { const int col = nn + (TF_SIMD_FRAGS ? (j & 1) + 8 * (j >> 1) : j); - if (col < TF_N) Y[(size_t)m * TF_N + col] = static_cast(C[f * 8 + r * 4 + j]); + if (col < TF_N) Y[(size_t)m * TF_YS + col] = static_cast(C[f * 8 + r * 4 + j]); } } } @@ -219,7 +269,7 @@ uint group_id [[simdgroup_index_in_threadgroup]],uint lane [[thread_index_in_sim threadgroup_barrier(mem_flags::mem_threadgroup); if(sg==0) for(int part=1;part> 3) & 1), mb = ((lane >> 1) & 3) + 4 * ((lane >> 4) & 1); @@ -288,7 +338,7 @@ inline tf_reg_load tf_reg_fetch(const device uint4* wl, const device uint4* sbv, if (ks == 0) for (int i = 0; i < 16; i++) { const int m = mb + 8 * ((i >> 2) & 1), n = n0 + nb + 16 * (i >> 3) + (i & 3); - if (m < M && n < TF_N) Y[(size_t)m * TF_N + n] = static_cast(C[i]); + if (m < M && n < TF_N) Y[(size_t)m * TF_YS + n] = static_cast(C[i]); } } #endif diff --git a/zig/src/core/lane_projection.zig b/zig/src/core/lane_projection.zig index 892a8a541b..001b684892 100644 --- a/zig/src/core/lane_projection.zig +++ b/zig/src/core/lane_projection.zig @@ -4,7 +4,7 @@ const mtl = @import("metal"); const ks = @import("kernel_sources"); const frags = @import("frags.zig"); -pub const weight_layout = "u32[N/32][K/group][32][group*bits/32], low-bit-first codes"; +pub const weight_layout = "u32[ceil(N/32)][K/group][32][group*bits/32], low-bit-first codes, columns past N zero"; pub const reg_layout = "u32[N/32][K/64][2][32 lanes][4] 4-bit codes, regBlock order"; pub const metadata_layout = "bf16[K/group][N][scale,bias]"; pub const Output = enum { bf16, f32 }; @@ -14,7 +14,7 @@ pub const Format = struct { group: usize, pub fn validate(f: Format) !void { - if (f.bits != 4 and f.bits != 6 and f.bits != 8) return error.UnsupportedBits; + if (f.bits < 2 or f.bits > 8 or f.bits == 7) return error.UnsupportedBits; if (f.group != 32 and f.group != 64 and f.group != 128) return error.UnsupportedGroup; } }; @@ -34,12 +34,16 @@ pub const Layout = struct { cooperative: bool = false, /// 4-bit g64 codes in reg_layout, unpacked into tensor-op registers; every word equals the cooperative kernel's. reg: bool = false, + /// Elements between output rows (default N): a member of a fused stack writes its columns of the fused rows. + y_stride: ?usize = null, pub fn validate(l: Layout) !void { try l.format.validate(); if (l.cooperative and l.format.bits != 4) return error.UnsupportedCooperativeFormat; if (l.reg and (l.cooperative or l.format.bits != 4 or l.format.group != 64 or !l.precompute_sums)) return error.UnsupportedRegFormat; - if (l.n == 0 or l.n % 32 != 0 or l.k == 0 or l.k % l.format.group != 0 or + if ((l.cooperative or l.reg) and l.n % 32 != 0) return error.InvalidShape; + if (l.y_stride) |s| if (s < l.n or s > std.math.maxInt(i32) / 16) return error.InvalidOutputStride; + if (l.n == 0 or l.n % 16 != 0 or l.k == 0 or l.k % l.format.group != 0 or l.n > std.math.maxInt(i32) / 16 or l.k > std.math.maxInt(i32) / 16) return error.InvalidShape; if (l.sk == 0 or l.sk > 32 or (l.pf != 1 and l.pf != 2)) return error.InvalidSchedule; const g = l.groupCut(); @@ -49,7 +53,7 @@ pub const Layout = struct { if (l.threadgroupBytes() > 32768) return error.ThreadgroupMemoryLimit; var end: usize = 0; for (l.ranges) |r| { - if (r[1] == 0 or r[0] < end or r[0] >= l.n / 32 or r[1] > l.n / 32 - r[0]) return error.InvalidTileRange; + if (r[1] == 0 or r[0] < end or r[0] >= l.tileCount() or r[1] > l.tileCount() - r[0]) return error.InvalidTileRange; end = r[0] + r[1]; } _ = try std.math.mul(usize, l.n, l.k); @@ -59,8 +63,13 @@ pub const Layout = struct { return l.groups orelse .{ 0, l.k / l.format.group }; } + /// 32-column tiles over N; a last tile of 16 columns is padded with zero codes and never stored. + pub fn tileCount(l: Layout) usize { + return (l.n + 31) / 32; + } + pub fn tiles(l: Layout) usize { - if (l.ranges.len == 0) return l.n / 32; + if (l.ranges.len == 0) return l.tileCount(); var count: usize = 0; for (l.ranges) |r| count += r[1]; return count; @@ -71,6 +80,11 @@ pub const Layout = struct { } pub fn weightWords(l: Layout) usize { + return l.tileCount() * 32 * (l.k / 32) * l.format.bits; + } + + /// MLX's packed W for this layout: N rows of K*bits/32 words. + pub fn sourceWords(l: Layout) usize { return l.n * (l.k / 32) * l.format.bits; } @@ -100,6 +114,7 @@ pub fn source(a: std.mem.Allocator, l: Layout) ![]u8 { try text.appendSlice(a, if (l.ranges.len == 0) " return j;\n}\n" else " return 0;\n}\n"); try text.print(a, "#define TF_PRECOMPUTE_SUMS {d}\n", .{@intFromBool(l.precompute_sums)}); if (l.reg) try text.appendSlice(a, "#define TF_REG 1\n"); + if (l.y_stride) |s| try text.print(a, "#define TF_YS {d}\n", .{s}); try text.appendSlice(a, ks.core_lane_projection); return text.toOwnedSlice(a); } @@ -110,8 +125,8 @@ pub fn pack(l: Layout, w: []const u32, scales: []const u16, biases: []const u16, const ng = l.k / l.format.group; const wp = l.format.group * l.format.bits / 32; const row_words = l.k / 32 * l.format.bits; - if (w.len != l.weightWords() or scales.len != l.n * ng or biases.len != l.n * ng or - out_w.len != w.len or out_sb.len != l.metadataElements()) return error.InvalidStorage; + if (w.len != l.sourceWords() or scales.len != l.n * ng or biases.len != l.n * ng or + out_w.len != l.weightWords() or out_sb.len != l.metadataElements()) return error.InvalidStorage; if (overlap(std.mem.sliceAsBytes(w), std.mem.sliceAsBytes(out_w)) or overlap(std.mem.sliceAsBytes(w), std.mem.sliceAsBytes(out_sb)) or overlap(std.mem.sliceAsBytes(scales), std.mem.sliceAsBytes(out_w)) or @@ -119,11 +134,14 @@ pub fn pack(l: Layout, w: []const u32, scales: []const u16, biases: []const u16, overlap(std.mem.sliceAsBytes(scales), std.mem.sliceAsBytes(out_sb)) or overlap(std.mem.sliceAsBytes(biases), std.mem.sliceAsBytes(out_sb)) or overlap(std.mem.sliceAsBytes(out_w), std.mem.sliceAsBytes(out_sb))) return error.OverlappingStorage; - for (0..l.n / 32) |tile| for (0..ng) |g| for (0..32) |col| { + for (0..l.tileCount()) |tile| for (0..ng) |g| for (0..32) |col| { const n = tile * 32 + col; - const src = n * row_words + g * wp; const dst = ((tile * ng + g) * 32 + col) * wp; - @memcpy(out_w[dst..][0..wp], w[src..][0..wp]); + if (n >= l.n) { + @memset(out_w[dst..][0..wp], 0); + continue; + } + @memcpy(out_w[dst..][0..wp], w[n * row_words + g * wp ..][0..wp]); const sb = (g * l.n + n) * 2; out_sb[sb] = scales[n * ng + g]; out_sb[sb + 1] = biases[n * ng + g]; @@ -207,6 +225,7 @@ pub const Projection = struct { weight_bytes: usize, metadata_bytes: usize, output_bytes: usize, + y_stride: usize, tile_count: usize, sk: usize, sums_pipe: ?mtl.Pipeline = null, @@ -233,7 +252,7 @@ pub const Projection = struct { errdefer sums_pipe.?.deinit(); sums_buffer = try device.buffer(@as(usize, if (l.reg) reg_rows else 16) * (l.k / l.format.group) * 4, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); } - return .{ .pipe = p, .n = l.n, .k = l.k, .group_count = l.k / l.format.group, .weight_bytes = l.weightWords() * 4, .metadata_bytes = l.metadataElements() * 2, .output_bytes = if (l.output == .f32) 4 else 2, .tile_count = l.tiles(), .sk = l.sk, .sums_pipe = sums_pipe, .sums_buffer = sums_buffer, .threads_per_slice = if (l.cooperative) 64 else 32, .reg = l.reg }; + return .{ .pipe = p, .n = l.n, .k = l.k, .group_count = l.k / l.format.group, .weight_bytes = l.weightWords() * 4, .metadata_bytes = l.metadataElements() * 2, .output_bytes = if (l.output == .f32) 4 else 2, .y_stride = l.y_stride orelse l.n, .tile_count = l.tiles(), .sk = l.sk, .sums_pipe = sums_pipe, .sums_buffer = sums_buffer, .threads_per_slice = if (l.cooperative) 64 else 32, .reg = l.reg }; } pub fn deinit(p: Projection) void { @@ -251,7 +270,7 @@ pub const Projection = struct { pub fn encodeSums(p: Projection, e: mtl.ComputeEncoder, x: Ref, w: Ref, sb: Ref, sums: Sums, y: Ref, rows: u32) !void { if (!p.reg) return error.SumsNeedReg; if (rows == 0 or rows > reg_rows or sums.stride < rows) return error.InvalidRows; - try checkRef(sums.ref, ((p.k / 64 - 1) * @as(usize, sums.stride) + rows) * 4, 4); + try checkRef(sums.ref, ((p.group_count - 1) * @as(usize, sums.stride) + rows) * 4, 4); return p.encodeImpl(e, x, w, sb, sums, y, rows); } @@ -262,7 +281,7 @@ pub const Projection = struct { try checkRef(x, rows * p.k * 2, 2); try checkRef(w, p.weight_bytes, if (p.reg) 16 else 8); try checkRef(sb, p.metadata_bytes, 16); - try checkRef(y, rows * p.n * p.output_bytes, p.output_bytes); + try checkRef(y, ((rows - 1) * p.y_stride + p.n) * p.output_bytes, p.output_bytes); if (given == null) if (p.sums_pipe) |pipe| { e.setPipeline(pipe); e.setBuffer(x.buf, x.off, 0); @@ -302,9 +321,35 @@ test "layouts reject unsupported formats, schedules, cuts and overlapping tile r const good = Layout{ .n = 64, .k = 256, .format = .{ .bits = 4, .group = 64 } }; const s = try source(arena.allocator(), good); try std.testing.expect(std.mem.indexOf(u8, s, "#define TF_BITS 4\n#define TF_GROUP 64\n") != null); + try std.testing.expect(std.mem.indexOf(u8, s[0 .. s.len - ks.core_lane_projection.len], "TF_YS") == null); var bad = good; - bad.format.bits = 5; - try std.testing.expectError(error.UnsupportedBits, bad.validate()); + for ([_]u8{ 2, 3, 5, 6, 8 }) |bits| { + bad.format.bits = bits; + try bad.validate(); + } + for ([_]u8{ 0, 1, 7, 9 }) |bits| { + bad.format.bits = bits; + try std.testing.expectError(error.UnsupportedBits, bad.validate()); + } + bad = good; + bad.y_stride = 96; + try std.testing.expect(std.mem.indexOf(u8, try source(arena.allocator(), bad), "#define TF_YS 96\n") != null); + bad.y_stride = 63; + try std.testing.expectError(error.InvalidOutputStride, bad.validate()); + bad = good; + bad.n = 48; + bad.ranges = &.{.{ 1, 1 }}; + try bad.validate(); + try std.testing.expectEqual(2, bad.tileCount()); + bad.cooperative = true; + try std.testing.expectError(error.InvalidShape, bad.validate()); + bad.cooperative = false; + bad.reg = true; + bad.precompute_sums = true; + try std.testing.expectError(error.InvalidShape, bad.validate()); + bad = good; + bad.ranges = &.{.{ 2, 1 }}; + try std.testing.expectError(error.InvalidTileRange, bad.validate()); bad = good; bad.format.group = 16; try std.testing.expectError(error.UnsupportedGroup, bad.validate()); @@ -334,34 +379,45 @@ test "layouts reject unsupported formats, schedules, cuts and overlapping tile r try std.testing.expectError(error.InvalidTileRange, bad.validate()); } -test "every affine layout packs raw words and metadata without changing their bits" { - for ([_]u8{ 4, 6, 8 }) |bits| for ([_]usize{ 32, 64, 128 }) |group| { - const l = Layout{ .n = 64, .k = 2 * group, .format = .{ .bits = bits, .group = group } }; - const a = std.testing.allocator; - const w = try a.alloc(u32, l.weightWords()); - defer a.free(w); - const s = try a.alloc(u16, 128); - defer a.free(s); - const b = try a.alloc(u16, 128); - defer a.free(b); - const pw = try a.alloc(u32, w.len); - defer a.free(pw); - const sb = try a.alloc(u16, 256); - defer a.free(sb); - for (w, 0..) |*v, i| v.* = @as(u32, @intCast(i)) *% 0x9e3779b9; - for (s, 0..) |*v, i| v.* = @intCast(i); - for (b, 0..) |*v, i| v.* = @intCast(65535 - i); +test "every width, group and 16-row member packs MLX codes, metadata and zero padding without changing a bit" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var prng: std.Random.DefaultPrng = .init(0x1a9e); + for ([_]u8{ 2, 3, 4, 5, 6, 8 }) |bits| for ([_]usize{ 32, 64, 128 }) |group| for ([_]usize{ 64, 48 }) |rows| { + const l = Layout{ .n = rows, .k = 2 * group, .format = .{ .bits = bits, .group = group } }; + const codes = try a.alloc(u8, l.n * l.k); + const w = try a.alloc(u32, l.sourceWords()); + @memset(w, 0); + for (codes, 0..) |*v, i| { + v.* = @intCast(prng.random().intRangeAtMost(u16, 0, (@as(u16, 1) << @intCast(bits)) - 1)); + // MLX's packing written bit by bit: row n is one low-bit-first stream of K codes. + for (0..bits) |bit| w[(i * bits + bit) / 32] |= @as(u32, (v.* >> @intCast(bit)) & 1) << @intCast((i * bits + bit) % 32); + } + const s = try a.alloc(u16, l.n * 2); + const b = try a.alloc(u16, l.n * 2); + for (s, b, 0..) |*sv, *bv, i| { + sv.* = @intCast(i); + bv.* = @intCast(65535 - i); + } + const pw = try a.alloc(u32, l.weightWords()); + @memset(pw, 0xffff_ffff); + const sb = try a.alloc(u16, l.metadataElements()); try pack(l, w, s, b, pw, sb); const wp = group * bits / 32; - for (0..64) |n| for (0..2) |g| { - const src = (n * 2 + g) * wp; - const dst = (((n / 32) * 2 + g) * 32 + n % 32) * wp; - try std.testing.expectEqualSlices(u32, w[src..][0..wp], pw[dst..][0..wp]); - try std.testing.expectEqual(s[n * 2 + g], sb[(g * 64 + n) * 2]); - try std.testing.expectEqual(b[n * 2 + g], sb[(g * 64 + n) * 2 + 1]); + for (0..l.tileCount() * 32) |n| for (0..l.k) |k| { + const at = ((((n / 32) * 2 + k / group) * 32 + n % 32) * wp) * 32 + (k % group) * bits; + var code: u32 = 0; + for (0..bits) |bit| code |= ((pw[(at + bit) / 32] >> @intCast((at + bit) % 32)) & 1) << @intCast(bit); + try std.testing.expectEqual(@as(u32, if (n < l.n) codes[n * l.k + k] else 0), code); + }; + for (0..l.n) |n| for (0..2) |g| { + try std.testing.expectEqual(s[n * 2 + g], sb[(g * l.n + n) * 2]); + try std.testing.expectEqual(b[n * 2 + g], sb[(g * l.n + n) * 2 + 1]); }; try std.testing.expectError(error.InvalidStorage, pack(l, w[1..], s, b, pw, sb)); - try std.testing.expectError(error.OverlappingStorage, pack(l, w, s, b, w, sb)); + try std.testing.expectError(error.InvalidStorage, pack(l, w, s, b, pw[1..], sb)); + try std.testing.expectError(error.OverlappingStorage, pack(l, pw[0..w.len], s, b, pw, sb)); }; } diff --git a/zig/tests/core_lane_projection.zig b/zig/tests/core_lane_projection.zig index bce8c3299a..bdd3e01b88 100644 --- a/zig/tests/core_lane_projection.zig +++ b/zig/tests/core_lane_projection.zig @@ -24,7 +24,7 @@ const Data = struct { biases: []u16, fn init(l: core.Layout) !Data { - const d = Data{ .x = try a.alloc(u16, rows_max * l.k), .codes = try a.alloc(u8, l.n * l.k), .words = try a.alloc(u32, l.weightWords()), .scales = try a.alloc(u16, l.n * l.k / l.format.group), .biases = try a.alloc(u16, l.n * l.k / l.format.group) }; + const d = Data{ .x = try a.alloc(u16, rows_max * l.k), .codes = try a.alloc(u8, l.n * l.k), .words = try a.alloc(u32, l.sourceWords()), .scales = try a.alloc(u16, l.n * l.k / l.format.group), .biases = try a.alloc(u16, l.n * l.k / l.format.group) }; @memset(d.words, 0); var prng: std.Random.DefaultPrng = .init(0x1a9e + @as(u64, l.format.bits) * 1000 + l.format.group); const rnd = prng.random(); @@ -169,7 +169,8 @@ fn check(device: mtl.Device, queue: mtl.Queue, l: core.Layout, counts: *Counts) counts.cpu_values += 1; }; } - if (l.format.bits == 6 and l.format.group == 32) { + const pinned = if (frozen.bytes(l)) |_| true else |_| false; // fn_lane's recorded 6/32 layouts only + if (pinned) { for (1..rows_max + 1) |m| { const count = m * l.n * elem; const got = windows.contents()[guard + (m - 1) * window_bytes ..][0..count]; @@ -180,20 +181,67 @@ fn check(device: mtl.Device, queue: mtl.Queue, l: core.Layout, counts: *Counts) for ([_]mtl.Buffer{ x, w, sb, one, windows }) |b| try checkGuard(b); counts.layouts += 1; std.debug.print("PASS b{d}g{d} sk{d} pf{d} {s} groups{any} tiles{d}: widths1-16, guards{s}\n", .{ - l.format.bits, l.format.group, l.sk, l.pf, @tagName(l.output), l.groupCut(), l.tiles(), - if (l.format.bits == 6 and l.format.group == 32) ", frozen fn_lane byte-equal" else "", + l.format.bits, l.format.group, l.sk, l.pf, @tagName(l.output), l.groupCut(), l.tiles(), + if (pinned) ", frozen fn_lane byte-equal" else "", + }); +} + +/// A member writing into fused rows: each row equals the unstrided encode's and nothing around it is touched. +fn strided(device: mtl.Device, queue: mtl.Queue, l: core.Layout) !void { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + var fused = l; + fused.y_stride = l.n + 48; + const p = try core.Projection.init(a, device, l); + defer p.deinit(); + const q = try core.Projection.init(a, device, fused); + defer q.deinit(); + const d = try Data.init(l); + defer d.deinit(); + const x = try buffer(device, d.x.len * 2); + defer x.deinit(); + const w = try buffer(device, l.weightWords() * 4); + defer w.deinit(); + const sb = try buffer(device, l.metadataElements() * 2); + defer sb.deinit(); + const elem: usize = if (l.output == .f32) 4 else 2; + const ys = fused.y_stride.? * elem; + const plain = try buffer(device, rows_max * l.n * elem); + defer plain.deinit(); + const stack = try buffer(device, rows_max * ys); + defer stack.deinit(); + @memcpy(x.contents()[guard..][0 .. d.x.len * 2], std.mem.sliceAsBytes(d.x)); + const packed_w = @as([*]u32, @ptrCast(@alignCast(w.contents() + guard)))[0..l.weightWords()]; + const packed_sb = @as([*]u16, @ptrCast(@alignCast(sb.contents() + guard)))[0..l.metadataElements()]; + try core.pack(l, d.words, d.scales, d.biases, packed_w, packed_sb); + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try p.encode(enc, ref(x, 0), ref(w, 0), ref(sb, 0), ref(plain, 0), rows_max); + try q.encode(enc, ref(x, 0), ref(w, 0), ref(sb, 0), ref(stack, 16 * elem), rows_max); + enc.end(); + try finish(cb); + for (0..rows_max) |r| { + const row = stack.contents()[guard + r * ys ..][0..ys]; + const want = plain.contents()[guard + r * l.n * elem ..][0 .. l.n * elem]; + if (!std.mem.eql(u8, want, row[16 * elem ..][0 .. l.n * elem])) return error.StridedRowMismatch; + for (row[0 .. 16 * elem]) |v| if (v != 0xa5) return error.StridedGapOverwrite; + for (row[(16 + l.n) * elem ..]) |v| if (v != 0xa5) return error.StridedGapOverwrite; + } + for ([_]mtl.Buffer{ x, w, sb, plain, stack }) |b| try checkGuard(b); + std.debug.print("PASS b{d}g{d} n{d} {s}: output rows {d} apart equal the unstrided rows, gaps untouched\n", .{ + l.format.bits, l.format.group, l.n, @tagName(l.output), fused.y_stride.?, }); } fn packBench(io: std.Io, l: core.Layout) !void { try l.validate(); - const w = try a.alloc(u32, l.weightWords()); + const w = try a.alloc(u32, l.sourceWords()); defer a.free(w); const s = try a.alloc(u16, l.n * l.k / l.format.group); defer a.free(s); const b = try a.alloc(u16, s.len); defer a.free(b); - const pw = try a.alloc(u32, w.len); + const pw = try a.alloc(u32, l.weightWords()); defer a.free(pw); const sb = try a.alloc(u16, 2 * s.len); defer a.free(sb); @@ -261,7 +309,7 @@ pub fn main(init: std.process.Init) !void { try check(device, queue, .{ .n = 96, .k = 224, .format = .{ .bits = 6, .group = 32 }, .sk = 4, .pf = 2, .ranges = &.{.{ 1, 1 }}, .groups = .{ 1, 3 }, .output = .f32 }, &counts); if (!device.tensorUnits()) { var unsupported = core.Layout{ .n = 64, .k = 224, .format = .{ .bits = 6, .group = 32 } }; - for ([_]u8{ 4, 8 }) |bits| { + for ([_]u8{ 2, 3, 4, 5, 8 }) |bits| { unsupported.format.bits = bits; try std.testing.expectError(error.TensorUnitsRequired, core.Projection.init(a, device, unsupported)); } @@ -278,12 +326,16 @@ pub fn main(init: std.process.Init) !void { unsupported.reg = true; try std.testing.expectError(error.TensorUnitsRequired, core.Projection.init(a, device, unsupported)); } + // 48 rows (a zero-padded last tile) and a member writing its columns of fused rows, on the same path + try check(device, queue, .{ .n = 48, .k = 160, .format = .{ .bits = 6, .group = 32 }, .sk = 2, .pf = 2, .output = .f32 }, &counts); + try strided(device, queue, .{ .n = 48, .k = 96, .format = .{ .bits = 6, .group = 32 } }); + try strided(device, queue, .{ .n = 64, .k = 96, .format = .{ .bits = 6, .group = 32 }, .sk = 2, .output = .f32 }); std.debug.print("PASS Flash Next: {d} layouts, {d} bit-equal row/window pairs, {d} frozen windows, {d} CPU values, {d} refusals; max fp32 error {d:.9}, max roundoff-bound fraction {d:.6}\n", .{ - counts.layouts, counts.row_pairs, counts.golden_windows, counts.cpu_values, @as(usize, if (device.tensorUnits()) 0 else 7), counts.max_error, counts.max_fraction, + counts.layouts, counts.row_pairs, counts.golden_windows, counts.cpu_values, @as(usize, if (device.tensorUnits()) 0 else 10), counts.max_error, counts.max_fraction, }); return; } - for ([_]u8{ 4, 6, 8 }) |bits| for ([_]usize{ 32, 64, 128 }) |group| { + for ([_]u8{ 2, 3, 4, 5, 6, 8 }) |bits| for ([_]usize{ 32, 64, 128 }) |group| { for ([_]usize{ 1, 2, 4 }) |sk| for ([_]usize{ 1, 2 }) |pf| { const l = core.Layout{ .n = 64, .k = 7 * group, .format = .{ .bits = bits, .group = group }, .sk = sk, .pf = pf, .output = .f32 }; try check(device, queue, l, &counts); @@ -309,6 +361,13 @@ pub fn main(init: std.process.Init) !void { const group: usize = if (bits == 4) 64 else if (bits == 6) 32 else 128; try check(device, queue, .{ .n = 96, .k = 7 * group, .format = .{ .bits = bits, .group = group }, .sk = 4, .pf = 2, .ranges = &.{.{ 1, 1 }}, .groups = .{ 1, 3 }, .output = .f32 }, &counts); } + for ([_]u8{ 2, 3, 4, 5, 6, 8 }) |bits| for ([_]usize{ 32, 64, 128 }) |group| { + // 48 rows: the last tile's 16 padded columns are read as zero codes and never stored. + try check(device, queue, .{ .n = 48, .k = 5 * group, .format = .{ .bits = bits, .group = group }, .sk = 2, .pf = 2, .output = .f32 }, &counts); + try strided(device, queue, .{ .n = 48, .k = 3 * group, .format = .{ .bits = bits, .group = group } }); + try strided(device, queue, .{ .n = 64, .k = 3 * group, .format = .{ .bits = bits, .group = group }, .sk = 2, .output = .f32 }); + }; + try strided(device, queue, .{ .n = 96, .k = 9 * 64, .format = .{ .bits = 4, .group = 64 }, .sk = 2, .precompute_sums = true, .cooperative = true }); std.debug.print("PASS {d} layouts, {d} bit-equal row/window pairs, {d} frozen fn_lane byte-equal windows, {d} CPU-reference values; max fp32 error {d:.9}, max roundoff-bound fraction {d:.6}\n", .{ counts.layouts, counts.row_pairs, counts.golden_windows, counts.cpu_values, counts.max_error, counts.max_fraction, }); From 0649371e605d865451b9915a5e2bd89cb4790d3b Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 04:10:26 -1000 Subject: [PATCH 05/18] The prompt product affine_mm reads 2-bit MLX codes, exactly affine_mm took 4- and 8-bit weights only, so a 2-bit matrix (oQ2's routed experts) had no tiled prompt product. At 2 bits a thread loads its 32 codes of a step as 8 bytes and writes each code into the bf16 tile as itself, with no offset (as at 8 bits; 0..3 is exact in bf16), so the group sums are the 4- and 8-bit path's. A first try with the 4-bit "+128" offset stayed within the float64 bound, but a reference with scales 1% off no longer exceeded it at K = 4096, so the codes enter exactly. The 4- and 8-bit sources preprocess as before. Test that fails before: "dense and gathered, each format, against exact sums on the host" adds 2/64 and 2/32, dense and gathered; source() refused 2 bits. The 4- and 8-bit sources preprocess (metal -E) to main's text at both groups, both outputs and both fragment layouts (16 variants). tf-affine-ref (next commit) holds the 2-bit products to float64. Gate on main 0efbd447 (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test: `get` logs an error on the DeferredTensor path the test expects). lean_check: 536 problems (main 539), no new findings. --- zig/kernels/metal/core/affine_mm.metal | 28 +++++++++++++++++++++++--- zig/src/core/affine_mm.zig | 12 +++++++---- 2 files changed, 33 insertions(+), 7 deletions(-) diff --git a/zig/kernels/metal/core/affine_mm.metal b/zig/kernels/metal/core/affine_mm.metal index 16e54006d2..464e459fe4 100644 --- a/zig/kernels/metal/core/affine_mm.metal +++ b/zig/kernels/metal/core/affine_mm.metal @@ -4,7 +4,7 @@ using namespace metal; #include "../nax.h" using namespace tfp; -// Defines: TF_BITS (4 or 8), TF_GROUP (32 or 64), TF_OUT_T (bfloat, or float for partial sums another pass adds). +// Defines: TF_BITS (2, 4 or 8), TF_GROUP (32 or 64), TF_OUT_T (bfloat, or float for partial sums another pass adds). #ifndef TF_OUT_T #define TF_OUT_T bfloat #endif @@ -66,7 +66,16 @@ inline void am_first(thread frag& lo, thread frag& hi, thread cons #endif } -// A thread's 32 codes of one weight row in a step (16-byte loads). +// A thread's 32 codes of one weight row in a step (16-byte loads; 8 bytes at 2 bits). +#if TF_BITS == 2 +struct Codes { + uint2 v; +}; + +inline Codes am_load(const device uchar* w) { + return Codes{*(const device uint2*)w}; +} +#else struct Codes { uint4 v[TF_BITS / 4]; }; @@ -79,11 +88,24 @@ inline Codes am_load(const device uchar* w) { } return c; } +#endif // The codes plus OFF, exactly, as bf16 into the tile row (16-byte stores). inline void am_codes(thread const Codes& c, threadgroup bfloat* out) { threadgroup uint4* o = (threadgroup uint4*)out; -#if TF_BITS == 4 +#if TF_BITS == 2 + TF_UNROLL + for (short i = 0; i < 2; i++) { // a word's 16 codes, as themselves (no offset): byte j holds codes 4j..4j+3 + const uint w = c.v[i]; + const float4 f0 = float4(as_type(w & 0x03030303u)), f1 = float4(as_type((w >> 2) & 0x03030303u)); + const float4 f2 = float4(as_type((w >> 4) & 0x03030303u)), f3 = float4(as_type((w >> 6) & 0x03030303u)); + TF_UNROLL + for (short j = 0; j < 4; j++) { + o[2 * i + j / 2][2 * (j % 2)] = as_type(ushort2(as_type(bfloat(f0[j])), as_type(bfloat(f1[j])))); + o[2 * i + j / 2][2 * (j % 2) + 1] = as_type(ushort2(as_type(bfloat(f2[j])), as_type(bfloat(f3[j])))); + } + } +#elif TF_BITS == 4 TF_UNROLL for (short i = 0; i < 4; i++) { const uint w = c.v[0][i]; diff --git a/zig/src/core/affine_mm.zig b/zig/src/core/affine_mm.zig index dd25dce831..4b5575a84a 100644 --- a/zig/src/core/affine_mm.zig +++ b/zig/src/core/affine_mm.zig @@ -4,12 +4,12 @@ const mtl = @import("metal"); const ks = @import("kernel_sources"); const frags = @import("frags.zig"); -/// The weights' bits (4 or 8) and group (32 or 64), and the output's type. +/// The weights' bits (2, 4 or 8) and group (32 or 64), and the output's type. pub const Format = struct { bits: u8 = 4, group: u16 = 64, out_f32: bool = false }; /// The kernel source for `f` (nax.h inlined for this GPU). pub fn source(device: mtl.Device, a: std.mem.Allocator, f: Format) ![]u8 { - if ((f.bits != 4 and f.bits != 8) or (f.group != 32 and f.group != 64)) return error.UnsupportedFormat; + if ((f.bits != 2 and f.bits != 4 and f.bits != 8) or (f.group != 32 and f.group != 64)) return error.UnsupportedFormat; const body = try frags.source(device, a, ks.core_affine_mm); defer a.free(body); return std.fmt.allocPrint(a, "#define TF_BITS {d}\n#define TF_GROUP {d}\n#define TF_OUT_T {s}\n{s}", .{ f.bits, f.group, if (f.out_f32) "float" else "bfloat", body }); @@ -135,7 +135,7 @@ test "dense and gathered, each format, against exact sums on the host" { const R = counts[0] + counts[1] + counts[2]; var prng = std.Random.DefaultPrng.init(7); const rnd = prng.random(); - for ([_][2]usize{ .{ 4, 64 }, .{ 4, 32 }, .{ 8, 64 }, .{ 8, 32 } }) |fmt| { + for ([_][2]usize{ .{ 4, 64 }, .{ 4, 32 }, .{ 8, 64 }, .{ 8, 32 }, .{ 2, 64 }, .{ 2, 32 } }) |fmt| { const bits = fmt[0]; const group = fmt[1]; const w_bytes = E * N * K * bits / 8; @@ -207,7 +207,11 @@ test "dense and gathered, each format, against exact sums on the host" { const off: f64 = if (bits == 4) 128 else 0; for (0..group) |j| { const k = g * group + j; - const qv: u32 = if (bits == 4) (wb[(row * K + k) / 2] >> @intCast(4 * (k % 2))) & 15 else wb[row * K + k]; + const qv: u32 = switch (bits) { + 2 => (wb[(row * K + k) / 4] >> @intCast(2 * (k % 4))) & 3, + 4 => (wb[(row * K + k) / 2] >> @intCast(4 * (k % 2))) & 15, + else => wb[row * K + k], + }; const xv: f64 = f32of(x[r * K + k]); dot += @as(f64, @floatFromInt(qv)) * xv; dmag += (@as(f64, @floatFromInt(qv)) + off) * @abs(xv); From 6d2a62aa796947d1cbe41c2464603e74b556b321 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 04:11:00 -1000 Subject: [PATCH 06/18] tests: affine-ref, the prompt products against float64 on Metal, every row's bits in every window The device side of #633's quant-ref for the shared prompt products, item 5 of #716. tf-affine-ref takes no arguments and needs no model: seeded synthetic codes, scales, biases and inputs, packed as MLX affine. The engine module now names core/affine_mm so a check program can reach it. - affine_mm at 2, 4 and 8 bits, groups 32 and 64, bf16 and fp32 out, dense and as a one-expert gather, at [192, 1024] and [256, 4096]. - simd_qmm (4-bit, group 64) as tiles and as the scalar twin at 1 and 2 rows, at [40, 512], [256, 4096] and [6400, 256] (K chunks of 32, 16 and 8). The reference dequantizes the same codes in float64 and multiplies in float64. Each product prints its worst column (max abs error, and that error over the column's largest |y|) and its largest error as a fraction of the bound: (group + K / group + 8) fp32 roundings of the terms' magnitudes as the kernels sum them (4-bit affine_mm codes plus 128), plus one bf16 ulp for bf16 outputs. A reference with its scales 1% off must exceed the bound. Then every row of a 129-row call must equal, byte for byte, the same row computed in windows of 1, 2, 3, 8, 16, 17, 33, 64 and 65 rows (every full window over the 129 rows, so each row is seen at every position), which crosses the 16-row fragments, the 32- and 64-row tiles and the gather's switch between them. The twin must equal the tiles wherever twinExact says it does, since only then do callers take it. Receipts, M3 Ultra (no tensor units: nax.h's simdgroup-matrix layout), on main 0efbd447: - every product within 0.4935 of its bound; every row's bits equal at every window, both kernels, every format; the twin equals the tiles at 1 and 2 rows on all three shapes, as twinExact says. - worst columns, bf16 out (dense and gather equal), at [192, 1024] and [256, 4096]: affine_mm q2 g32 2.64e-2 (rel 3.23e-3) and 6.22e-2 (3.78e-3), q2 g64 2.87e-2 (3.51e-3) and 3.12e-2 (2.47e-3), q4 g32 2.45e-2 (2.65e-3) and 3.12e-2 (3.01e-3), q4 g64 1.56e-2 (2.89e-3) and 3.12e-2 (2.53e-3), q8 g32 0.496 (3.41e-3) and 0.968 (3.63e-3), q8 g64 0.453 (2.87e-3) and 0.956 (3.38e-3); simd_qmm 1.19e-2 (2.57e-3), 3.13e-2 (2.87e-3) and 1.54e-2 (3.74e-3). All under 2^-8 of the column's largest |y|; 0.20 to 0.49 of the bound. - fp32 out: at most 0.0081 of the bound (q2 1.5e-6 to 7.7e-6, rel 1.9e-7 to 6.0e-7). - scales 1% off: 2.3 to 29 times the bound for bf16 outputs, 2.7 to 391 for fp32, 44 to 273 for simd_qmm. Not run: an M5, where affine_mm takes the tensor units; `zig build tf-affine-ref && ./zig-out/bin/tf-affine-ref`. The fp32 bound is the worst case for any summation order, so a changed order still passes; the measured errors sit two orders below it, and the bit checks are what pin the order. Gate on main 0efbd447 (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test: `get` logs an error on the DeferredTensor path the test expects). lean_check: 536 problems (main 539), no new findings. --- build.zig | 1 + zig/build/core_checks.zig | 14 ++ zig/src/tensorfold.zig | 1 + zig/tests/affine_ref.zig | 358 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 374 insertions(+) create mode 100644 zig/tests/affine_ref.zig diff --git a/build.zig b/build.zig index 2ba6052374..d3b8a0a760 100644 --- a/build.zig +++ b/build.zig @@ -476,6 +476,7 @@ fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b // The core lane, shared-context and Qwen3.8-27B programs and host tests. qwen27_build.targets(b, target, optimize, metal, sources, engine, lanes, mods.qwen27, test_step); + core_checks.targets(b, target, optimize, metal, engine); _ = nativeServer(b, target, optimize, metal, engine, api, engines, build_options, test_step); diff --git a/zig/build/core_checks.zig b/zig/build/core_checks.zig index 75f39c00c2..6749819464 100644 --- a/zig/build/core_checks.zig +++ b/zig/build/core_checks.zig @@ -16,3 +16,17 @@ pub fn rowWidths(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std. const about = "The row projection at every MLX width and group: fp64 reference, windows, strides, stock bytes"; b.step("tf-row-widths-check", about).dependOn(&b.addInstallArtifact(exe, .{}).step); } + +/// tf-affine-ref: the prompt products (affine_mm, simd_qmm) against float64, every row's bits in every window. +pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, metal: *std.Build.Module, engine: *std.Build.Module) void { + const exe = b.addExecutable(.{ .name = "tf-affine-ref", .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/affine_ref.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "tensorfold", .module = engine } }, + }) }); + b.installArtifact(exe); + const about = "The prompt products (affine_mm, simd_qmm) on synthetic weights against float64, every row's bits in every window"; + b.step("tf-affine-ref", about).dependOn(&b.addInstallArtifact(exe, .{}).step); +} diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 1823eb56a9..f34f41a4ff 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -13,6 +13,7 @@ pub const ids_json = @import("core/ids_json.zig"); pub const lanes = @import("lanes"); pub const segments = @import("core/segments.zig"); pub const row_projection = @import("core/row_projection.zig"); +pub const affine_mm = @import("core/affine_mm.zig"); pub const exl3_format = @import("core/exl3_format.zig"); pub const exl3_projection = @import("core/exl3_projection.zig"); pub const exlr8 = @import("core/exlr8.zig"); diff --git a/zig/tests/affine_ref.zig b/zig/tests/affine_ref.zig new file mode 100644 index 0000000000..5f32bd864a --- /dev/null +++ b/zig/tests/affine_ref.zig @@ -0,0 +1,358 @@ +//! The shared prompt products against float64 on synthetic weights, and every row's bits in every window width. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const mm = tf.affine_mm; +const sq = tf.row_projection.simd; + +const opts = mtl.ResourceOptions.shared; +const a = std.heap.page_allocator; +const unit: f64 = 0x1p-24; // fp32's unit roundoff +const wide = 129; // rows of the reference call: two 64-row tiles and one more +const windows = [_]usize{ 1, 2, 3, 8, 16, 17, 33, 64, 65 }; // 17, 33 and 65 cross the 16-, 32- and 64-row tiles + +fn toBf(v: f32) u16 { + const u: u32 = @bitCast(v); + return @truncate((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); +} + +fn fromBf(v: u16) f64 { + return @as(f32, @bitCast(@as(u32, v) << 16)); +} + +/// Synthetic MLX affine weights [n, k] (codes little-endian in bytes) and `wide` rows of x; `deq` is s q + b exactly. +const Synth = struct { + n: usize, + k: usize, + bits: u8, + group: usize, + x: mtl.Buffer, + w: mtl.Buffer, + s: mtl.Buffer, + b: mtl.Buffer, + deq: []f64, + + fn init(device: mtl.Device, n: usize, k: usize, bits: u8, group: usize, seed: u64) !Synth { + var prng = std.Random.DefaultPrng.init(seed); + const r = prng.random(); + const sn = n * (k / group); + const self = Synth{ .n = n, .k = k, .bits = bits, .group = group, .x = try device.buffer(wide * k * 2, opts), .w = try device.buffer(n * k * bits / 8, opts), .s = try device.buffer(sn * 2, opts), .b = try device.buffer(sn * 2, opts), .deq = try a.alloc(f64, n * k) }; + for (self.x.slice(u16, wide * k)) |*v| v.* = toBf(r.float(f32) * 2 - 1); + r.bytes(self.w.slice(u8, n * k * bits / 8)); + for (self.s.slice(u16, sn), self.b.slice(u16, sn)) |*s, *b| { + s.* = toBf((r.float(f32) + 0.5) / 64); + b.* = toBf(-(r.float(f32) + 0.5) * 0.15); + } + for (0..n) |j| for (0..k) |i| { + const g = j * (k / group) + i / group; + self.deq[j * k + i] = fromBf(self.s.slice(u16, sn)[g]) * @as(f64, @floatFromInt(self.code(j, i))) + fromBf(self.b.slice(u16, sn)[g]); + }; + return self; + } + + fn deinit(self: Synth) void { + for ([_]mtl.Buffer{ self.x, self.w, self.s, self.b }) |x| x.deinit(); + a.free(self.deq); + } + + fn code(self: Synth, j: usize, i: usize) u32 { + const bytes = self.w.slice(u8, self.n * self.k * self.bits / 8); + if (self.bits == 8) return bytes[j * self.k + i]; + if (self.bits == 2) return (bytes[(j * self.k + i) / 4] >> @intCast(2 * (i % 4))) & 3; + return (bytes[(j * self.k + i) / 2] >> @intCast(4 * (i % 2))) & 15; + } + + /// y = x W^T in float64, and the sizes the fp32 sums carry when codes enter as q + off (affine_mm's 4-bit: 128). + fn reference(self: Synth, off: f64) !Ref { + const n = self.n; + const k = self.k; + const groups = k / self.group; + const ref = Ref{ .want = try a.alloc(f64, wide * n), .mag = try a.alloc(f64, wide * n), .scaled = try a.alloc(f64, wide * n) }; + const x = self.x.slice(u16, wide * k); + const sc = self.s.slice(u16, n * groups); + const bi = self.b.slice(u16, n * groups); + const xd = try a.alloc(f64, k); + defer a.free(xd); + const xg = try a.alloc(f64, groups); // each group's sum of |x| + defer a.free(xg); + for (0..wide) |r| { + @memset(xg, 0); + for (xd, 0..) |*v, i| { + v.* = fromBf(x[r * k + i]); + xg[i / self.group] += @abs(v.*); + } + for (0..n) |j| { + var sum: f64 = 0; + var mag: f64 = 0; + var scaled: f64 = 0; + for (0..groups) |g| { + var qx: f64 = 0; + var qxs: f64 = 0; + for (g * self.group..(g + 1) * self.group) |i| { + const q: f64 = @floatFromInt(self.code(j, i)); + sum += xd[i] * self.deq[j * k + i]; + qx += (q + off) * @abs(xd[i]); + qxs += q * xd[i]; + } + const s = fromBf(sc[j * groups + g]); + mag += s * qx + @abs(fromBf(bi[j * groups + g]) - off * s) * xg[g]; + scaled += s * qxs; + } + ref.want[r * n + j] = sum; + ref.mag[r * n + j] = mag; + ref.scaled[r * n + j] = scaled; + } + } + return ref; + } +}; + +/// The float64 product, its terms' fp32 magnitudes, and its scales' share (a control: those scales 1% off must fail). +const Ref = struct { + want: []f64, + mag: []f64, + scaled: []f64, + fn deinit(r: Ref) void { + for ([_][]f64{ r.want, r.mag, r.scaled }) |v| a.free(v); + } +}; + +/// Every check's outcome: failures, and the largest difference as a fraction of its bound. +const Tally = struct { failed: usize = 0, worst: f64 = 0 }; + +/// One product's output: `rows` rows of n bf16 or fp32 values. +const Out = struct { + buf: mtl.Buffer, + f32_out: bool, + fn esize(o: Out) usize { + return if (o.f32_out) 4 else 2; + } + fn get(o: Out, i: usize) f64 { + return if (o.f32_out) o.buf.slice(f32, i + 1)[i] else fromBf(o.buf.slice(u16, i + 1)[i]); + } +}; + +/// The bound: fp32 sums of a group's products, its row sum and a chain of groups, plus one bf16 ulp of the value. +fn bound(want: f64, mag: f64, s: *const Synth, f32_out: bool, extra: usize) f64 { + const terms: f64 = @floatFromInt(s.group + s.k / s.group + 8 + extra); + return terms * unit * mag + (if (f32_out) 0 else @abs(want) * 0x1p-7) + 1e-30; +} + +/// `rows` rows of `out` against the float64 reference: the worst column's max abs and relative error, and the bound. +fn compare(t: *Tally, name: []const u8, out: Out, rows: usize, ref: Ref, s: *const Synth, extra: usize) void { + const n = s.n; + var col: usize = 0; + var col_abs: f64 = -1; + var col_rel: f64 = 0; + var frac: f64 = 0; + var ctl: f64 = 0; + var at: usize = 0; + for (0..n) |j| { + var abs: f64 = 0; + var top: f64 = 0; + for (0..rows) |r| { + const i = r * n + j; + const d = @abs(out.get(i) - ref.want[i]); + const tol = bound(ref.want[i], ref.mag[i], s, out.f32_out, extra); + const q = d / tol; + ctl = @max(ctl, @abs(out.get(i) - ref.want[i] - 0.01 * ref.scaled[i]) / tol); + if (!(d <= abs)) abs = d; + top = @max(top, @abs(ref.want[i])); + if (!(q <= frac)) { + frac = q; + at = i; + } + } + if (!(abs <= col_abs)) { + col_abs = abs; + col = j; + col_rel = abs / top; + } + } + t.worst = @max(t.worst, frac); + const ok = frac <= 1 and ctl > 1; + if (!ok) t.failed += 1; + std.debug.print("RESULT {s}: worst column {d}: max abs {e:.3}, relative {e:.3}; {d:.4} of its bound (row {d} col {d}), " ++ + "scales 1% off {d:.1}{s}\n", .{ name, col, col_abs, col_rel, frac, at / n, at % n, ctl, if (ok) "" else " FAIL" }); +} + +/// Rows [0, rows) of `got` equal `want` byte for byte; the first difference printed. +fn sameRows(t: *Tally, name: []const u8, m: usize, want: Out, got: Out, rows: usize, n: usize) bool { + const bytes = rows * n * want.esize(); + const w = want.buf.slice(u8, bytes); + const g = got.buf.slice(u8, bytes); + if (std.mem.eql(u8, w, g)) return true; + const i = std.mem.indexOfDiff(u8, w, g).? / want.esize(); + var differ: usize = 0; + var most: f64 = 0; + for (0..rows) |r| { + const row = r * n * want.esize(); + if (!std.mem.eql(u8, w[row..][0 .. n * want.esize()], g[row..][0 .. n * want.esize()])) differ += 1; + for (r * n..(r + 1) * n) |e| most = @max(most, @abs(got.get(e) - want.get(e))); + } + std.debug.print("RESULT {s}: {d}-row windows DIFFER in {d} of {d} rows, by up to {e:.3}; first row {d} col {d}: {d} against {d} in the {d}-row call\n", .{ name, m, differ, rows, most, i / n, i % n, got.get(i), want.get(i), wide }); + t.failed += 1; + return false; +} + +fn newOut(device: mtl.Device, s: *const Synth, f32_out: bool) !Out { + const o = Out{ .buf = try device.buffer(wide * s.n * 4, opts), .f32_out = f32_out }; + @memset(o.buf.slice(u8, wide * s.n * 4), 0xff); // NaN: a launch that writes nothing fails + return o; +} + +fn commit(cb: anytype, e: mtl.ComputeEncoder) !void { + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |why| { + std.debug.print("GPU failed: {s}\n", .{why}); + return error.GpuFailed; + } +} + +const Kernel = enum { dense, gather }; + +/// affine_mm over rows [start, start + m): its row sums, then the dense product or a gather with one expert. +fn affineWindow(e: mtl.ComputeEncoder, p: mm.Pipes, s: *const Synth, sums: mtl.Buffer, offs: mtl.Buffer, y: Out, kernel: Kernel, start: usize, m: usize) void { + const x = .{ .buf = s.x, .off = start * s.k * 2 }; + const xs = .{ .buf = sums, .off = start * (s.k / s.group) * 4 }; + const yo = .{ .buf = y.buf, .off = start * s.n * y.esize() }; + const z: usize = 0; + const q = .{ .w = .{ .buf = s.w, .off = z }, .s = .{ .buf = s.s, .off = z }, .b = .{ .buf = s.b, .off = z }, .n = @as(u32, @intCast(s.n)), .k = @as(u32, @intCast(s.k)) }; + mm.rowSums(e, p, @intCast(s.group), x, xs, @intCast(m), @intCast(s.k)); + switch (kernel) { + .dense => mm.dense(e, p, x, xs, q, yo, @intCast(m)), + .gather => mm.gather(e, p, x, xs, q, .{ .buf = offs, .off = z }, yo, @intCast(m), 1), + } +} + +/// affine_mm at one format: the wide call against float64, then every window's rows against the wide call's. +fn affine(t: *Tally, device: mtl.Device, queue: mtl.Queue, s: *const Synth, ref: Ref, f32_out: bool) !void { + const src = try mm.source(device, a, .{ .bits = s.bits, .group = @intCast(s.group), .out_f32 = f32_out }); + defer a.free(src); + const lib = try mtl.Library.fromSource(device, src, mtl.CompileOptions.mlx()); + defer lib.deinit(); + var p: mm.Pipes = undefined; + for (mm.names, 0..) |name, i| p[i] = try mtl.Pipeline.init(device, lib, name, false); + defer for (p) |x| x.deinit(); + const sums = try device.buffer(wide * (s.k / s.group) * 4, opts); + defer sums.deinit(); + const offs = try device.buffer(4, opts); + defer offs.deinit(); + offs.slice(i32, 1)[0] = 0; + var name: [96]u8 = undefined; + for ([_]Kernel{ .dense, .gather }) |kernel| { + const label = try std.fmt.bufPrint(&name, "affine_mm {s} q{d} g{d} {s} [{d}, {d}]", .{ @tagName(kernel), s.bits, s.group, if (f32_out) "fp32" else "bf16", s.n, s.k }); + const all = try newOut(device, s, f32_out); + defer all.buf.deinit(); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + affineWindow(e, p, s, sums, offs, all, kernel, 0, wide); + try commit(cb, e); + compare(t, label, all, wide, ref, s, 0); + var same = true; + for (windows) |m| { + const got = try newOut(device, s, f32_out); + defer got.buf.deinit(); + const wb = queue.commandBuffer(); + const we = wb.compute(.serial); + var start: usize = 0; + while (start + m <= wide) : (start += m) affineWindow(we, p, s, sums, offs, got, kernel, start, m); + try commit(wb, we); + same = sameRows(t, label, m, all, got, wide / m * m, s.n) and same; + } + if (same) std.debug.print("RESULT {s}: every row's bits equal in windows of 1, 2, 3, 8, 16, 17, 33, 64 and 65 rows\n", .{label}); + } +} + +/// simd_qmm over rows [start, start + m), the scalar twin when `scalar` and it takes m rows. +fn simdWindow(e: mtl.ComputeEncoder, p: *const sq.Pipelines, s: *const Synth, y: Out, scalar: bool, start: usize, m: usize) !void { + const mat = sq.Matrix{ .w = s.w, .scales = s.s, .biases = s.b, .n = s.n, .k = s.k }; + sq.encode(e, try sq.call(p, s.n, s.k, m, scalar), mat, s.x, start * s.k * 2, y.buf, start * s.n * 2); +} + +/// simd_qmm (4-bit, group 64): tiles at every window against float64 and the wide call; the twin at 1 and 2 rows. +fn simd(t: *Tally, device: mtl.Device, queue: mtl.Queue, p: *const sq.Pipelines, s: *const Synth, ref: Ref) !void { + var name: [96]u8 = undefined; + const label = try std.fmt.bufPrint(&name, "simd_qmm tiles q4 g64 bf16 [{d}, {d}]", .{ s.n, s.k }); + const all = try newOut(device, s, false); + defer all.buf.deinit(); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + try simdWindow(e, p, s, all, false, 0, wide); + try commit(cb, e); + const extra = std.math.log2_int(u32, sq.splits(s.n)) + 1; // the chunks' fixed tree + compare(t, label, all, wide, ref, s, extra); + var same = true; + for (windows) |m| { + const got = try newOut(device, s, false); + defer got.buf.deinit(); + const wb = queue.commandBuffer(); + const we = wb.compute(.serial); + var start: usize = 0; + while (start + m <= wide) : (start += m) try simdWindow(we, p, s, got, false, start, m); + try commit(wb, we); + same = sameRows(t, label, m, all, got, wide / m * m, s.n) and same; + } + if (same) std.debug.print("RESULT {s}: every row's bits equal in windows of 1, 2, 3, 8, 16, 17, 33, 64 and 65 rows\n", .{label}); + // the twin: callers take it only where twinExact says it equals the tiles, so a difference fails only then + const exact = try sq.twinExact(p, device, .{ .w = s.w, .scales = s.s, .biases = s.b, .n = s.n, .k = s.k }); + for ([_]usize{ 1, 2 }) |m| { + const twin = try std.fmt.bufPrint(&name, "simd_qmm twin q4 g64 bf16 [{d}, {d}] at {d} rows", .{ s.n, s.k, m }); + const got = try newOut(device, s, false); + defer got.buf.deinit(); + const wb = queue.commandBuffer(); + const we = wb.compute(.serial); + var start: usize = 0; + while (start + m <= wide) : (start += m) try simdWindow(we, p, s, got, true, start, m); + try commit(wb, we); + compare(t, twin, got, wide / m * m, ref, s, 0); + const rows = wide / m * m; + const eq = std.mem.eql(u8, all.buf.slice(u8, rows * s.n * 2), got.buf.slice(u8, rows * s.n * 2)); + std.debug.print("RESULT {s}: {s} the tiles' bits; twinExact says {s}{s}\n", .{ twin, if (eq) "equals" else "differs from", if (exact) "equal" else "differ", if (exact and !eq) " FAIL" else "" }); + if (exact and !eq) t.failed += 1; + } +} + +pub fn main(_: std.process.Init) !void { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + std.debug.print("device {s}, tensor units {}\n", .{ device.name(), device.tensorUnits() }); + var t = Tally{}; + var seed: u64 = 0x716; + for ([_][2]usize{ .{ 192, 1024 }, .{ 256, 4096 } }) |shape| for ([_]u8{ 4, 8 }) |bits| for ([_]usize{ 32, 64 }) |group| { + seed += 1; + const s = try Synth.init(device, shape[0], shape[1], bits, group, seed); + defer s.deinit(); + const ref = try s.reference(if (bits == 4) 128 else 0); + defer ref.deinit(); + for ([_]bool{ false, true }) |f32_out| try affine(&t, device, queue, &s, ref, f32_out); + }; + for ([_][2]usize{ .{ 192, 1024 }, .{ 256, 4096 } }) |shape| for ([_]usize{ 32, 64 }) |group| { // 2 bits: oQ2's experts + const s = try Synth.init(device, shape[0], shape[1], 2, group, 0x2716 + shape[0] + group); + defer s.deinit(); + const ref = try s.reference(0); + defer ref.deinit(); + for ([_]bool{ false, true }) |f32_out| try affine(&t, device, queue, &s, ref, f32_out); + }; + var p = try sq.Pipelines.load(device); + defer p.deinit(); + for ([_][2]usize{ .{ 40, 512 }, .{ 256, 4096 }, .{ 6400, 256 } }) |shape| { // chunks of 32, 16 and 8 + seed += 1; + const s = try Synth.init(device, shape[0], shape[1], 4, 64, seed); + defer s.deinit(); + const ref = try s.reference(0); + defer ref.deinit(); + try simd(&t, device, queue, &p, &s, ref); + } + if (t.failed > 0) { + std.debug.print("FAIL {d} checks; the largest difference is {d:.4} of its float64 bound\n", .{ t.failed, t.worst }); + return error.CheckFailed; + } + std.debug.print("PASS every product within {d:.4} of its float64 bound, every row's bits equal at every window\n", .{t.worst}); +} From 9c4d7e3005bb5fc01d30f61d1a5414c10d213e2d Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 04:28:08 -1000 Subject: [PATCH 07/18] GLM-5.3-Flash reads each module's MLX affine width and mlx-lm's layout; a width-carrying matrix replaces Q4 Every GLM projection was typed 4-bit/group-64 (weights.zig's Q4), so an oQ checkpoint could not even be planned. - Widths (quant.zig): each module's MLX affine width from config.json through core/quant (keys.canonical, then GLM's own spellings: the MTP block is layer 45, forget_gate folds into the KDA gate). Aliases that disagree, a missing table and other modes are refused; a tensor stored at another width than declared is refused by name, never read at a guessed width. - Affine{bits, group} in place of Q4 (affine.zig): stored shapes, byte spans and column cuts follow each tensor's width; a fused matrix whose members differ in format plans one run a member at its output rows, as 0.6.6's QSplit calls each part (glm5_next/linear.py); a one-format stack stays one matrix. The load plan's layer and expert bytes are counted at those widths. - mlx-lm's layout (names.zig), which the public oQ checkpoints use: the MTP block, hyper-connections, stacked experts, the absorbed embed_q/unembed_out pair and a bf16 eh_proj; a fused conv1d.weight is cut into q/k/v; fp32 hyper-connection mixes and router narrow to bf16 only when exact (else refused naming fix, shard, offset, shape). - The dry-run plan prints its matrices by width. A real load still refuses any matrix the generated 4/64 kernels cannot read (the next commit gives the device every width); TP2 and expert parallel refuse any width but 4/64. The uniform 4/64 checkpoint plans the same matrices, offsets and bytes as the 4-bit-only loader (a digest pinned before the change). Tests that fail before (no width reading, no mlx-lm names on main): weights_test.zig, oQ-shaped fake checkpoints in both layouts, each tensor found at its own width; overrides under either prefix; a tensor stored at another width, two spellings or the MTP block and layer 45 at two widths, an odd member and a routed gate/up pair at two widths refused by name; oQ4's KDA input and MLA x_proj planned as one run a member; mlx-lm's conv1d and fp32 mixes planned as the original layout's tensors. quant.zig, affine.zig, names.zig and weights_fix.zig have their own unit tests. Gate on main 0efbd447 + the core commits (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test logs an error on the path it expects). lean_check: 533 problems (main 539), no new findings. --- zig/src/families/glm/affine.zig | 211 ++++++++++ zig/src/families/glm/forward.zig | 2 +- zig/src/families/glm/glm.zig | 5 + zig/src/families/glm/load_plan.zig | 65 +++ zig/src/families/glm/names.zig | 104 +++++ zig/src/families/glm/prompt.zig | 6 +- zig/src/families/glm/quant.zig | 174 ++++++++ zig/src/families/glm/weights.zig | 568 ++++++++++++++------------ zig/src/families/glm/weights_fix.zig | 95 ++++- zig/src/families/glm/weights_test.zig | 468 +++++++++++++++++++++ 10 files changed, 1432 insertions(+), 266 deletions(-) create mode 100644 zig/src/families/glm/affine.zig create mode 100644 zig/src/families/glm/names.zig create mode 100644 zig/src/families/glm/quant.zig create mode 100644 zig/src/families/glm/weights_test.zig diff --git a/zig/src/families/glm/affine.zig b/zig/src/families/glm/affine.zig new file mode 100644 index 0000000000..cd80ddbb1c --- /dev/null +++ b/zig/src/families/glm/affine.zig @@ -0,0 +1,211 @@ +//! GLM-5.3-Flash's affine matrices at any MLX width: stored shapes, byte spans and the widths each device path reads. +const std = @import("std"); +const mtl = @import("metal"); +const quant = @import("quant.zig"); + +pub const Width = quant.Width; + +pub const Ref = struct { + buf: mtl.Buffer, + off: usize = 0, + + pub fn at(r: Ref, by: usize) Ref { + return .{ .buf = r.buf, .off = r.off + by }; + } + + pub fn addr(r: Ref) [*]u8 { + return r.buf.contents() + r.off; + } +}; + +/// The most runs a fused matrix splits into: its members. +pub const max_runs = 6; + +/// A run of a split stack: one member (0.6.6 calls each apart), rows [row, row + n) of the fused output, bytes its own. +pub const Run = struct { w: Ref, s: Ref, b: Ref, row: u32, n: u32, bits: u8, group: u16 }; + +/// An MLX affine matrix: w [n, k bits / 32] u32, s and b [n, k / group] bf16 (stacked: [e, n, ...]). +pub const Affine = struct { + w: Ref, + s: Ref, + b: Ref, + n: u32, + k: u32, + bits: u8, + group: u16, + more: u8 = 0, // runs past the first, whose bytes and width are the matrix's own: 0 for one width + runs: [max_runs - 1]Run = undefined, + + pub fn width(q: Affine) Width { + return .{ .bits = q.bits, .group = q.group }; + } + + /// Run `i` (the first: the matrix's own bytes and width). + pub fn run(q: Affine, i: usize) Run { + if (i > 0) return q.runs[i - 1]; + return .{ .w = q.w, .s = q.s, .b = q.b, .row = 0, .n = if (q.more == 0) q.n else q.runs[0].row, .bits = q.bits, .group = q.group }; + } +}; + +/// Bytes of [n, k]'s packed codes at `w`. +pub fn wBytes(n: usize, k: usize, w: Width) usize { + return n * k * w.bits / 8; +} + +/// Bytes of [n, k]'s bf16 scale (or bias) table at `w`. +pub fn sBytes(n: usize, k: usize, w: Width) usize { + return n * (k / w.group) * 2; +} + +/// Bytes of [n, k]'s codes and both tables at `w`. +pub fn bytes(n: usize, k: usize, w: Width) usize { + return wBytes(n, k, w) + 2 * sBytes(n, k, w); +} + +/// [n, k]'s stored shapes at `w`: u32 words [n, k bits / 32], bf16 tables [n, k / group]; null when k does not pack. +pub fn shapes(n: usize, k: usize, w: Width) ?[2][2]usize { + if (k == 0 or k % w.group != 0 or k * w.bits % 32 != 0) return null; + return .{ .{ n, k * w.bits / 32 }, .{ n, k / w.group } }; +} + +/// Columns [lo, lo + k) as (offset, length) in a row's words and in its tables; null unless whole groups and words. +pub fn cut(lo: usize, k: usize, w: Width) ?[2][2]usize { + if (lo % w.group != 0 or shapes(1, k, w) == null) return null; + return .{ .{ wBytes(1, lo, w), wBytes(1, k, w) }, .{ sBytes(1, lo, w), sBytes(1, k, w) } }; +} + +/// Matrices (each run of a split stack) by width, for a plan's summary. +pub const Tally = struct { + count: [9][3]u32 = @splat(@splat(0)), // by bits, then group 32, 64, 128 + + pub fn add(t: *Tally, v: anytype) void { + const T = @TypeOf(v); + if (T == Affine) { + for (0..@as(usize, v.more) + 1) |i| { + const r = v.run(i); + if (r.n > 0 and r.bits > 0) t.count[r.bits][std.math.log2_int(u16, r.group) - 5] += 1; + } + return; + } + switch (@typeInfo(T)) { + .@"struct" => |s| inline for (s.field_names) |name| t.add(@field(v, name)), + .@"union" => switch (v) { + inline else => |x| t.add(x), + }, + .optional => if (v) |x| t.add(x), + .array => for (v) |x| t.add(x), + else => {}, + } + } + + /// "4-bit/64 x12, 8-bit/64 x3" in `buf`. + pub fn summary(t: Tally, buf: []u8) []const u8 { + var n: usize = 0; + for (t.count, 0..) |g, bits| for (g, 0..) |x, gi| if (x > 0) { + const sep = if (n == 0) "" else ", "; + n += (std.fmt.bufPrint(buf[n..], "{s}{d}-bit/{d} x{d}", .{ sep, bits, @as(u32, 32) << @intCast(gi), x }) catch return buf[0..n]).len; + }; + return buf[0..n]; + } +}; + +/// The device paths that read affine matrices, by projection class. +pub const Path = enum { embed, kda_in, kda_gate, kda_out, mla_in, mla_q, kv_b, mla_out, dense_mlp, shared_expert, routed_experts, head, eh_proj }; + +/// The code widths the row kernels (glm_rows.metal, after core/row_projection.metal) unpack, in groups of 32, 64, 128. +pub const row_bits = [_]u8{ 2, 3, 4, 5, 6, 8 }; + +/// Whether a path's kernels read `w`: kv_b 4/64, the embedding 4/6/8 in groups of 64, the rest the row kernels' widths. +pub fn reads(path: Path, w: Width) bool { + return switch (path) { + .kv_b => std.meta.eql(w, quant.uniform), + .embed => w.group == 64 and (w.bits == 4 or w.bits == 6 or w.bits == 8), + else => std.mem.indexOfScalar(u8, &row_bits, w.bits) != null and (w.group == 32 or w.group == 64 or w.group == 128), + }; +} + +/// Refuses, by tensor name, a width the path's kernels do not read: a load fails before any weight is read. +pub fn admit(name: []const u8, path: Path, w: Width, rows: usize) error{ UnsupportedWidth, UnsupportedShape }!void { + if (!reads(path, w)) { + quant.say("glm: {s} is {d}-bit in groups of {d}, a width the {t} kernels do not read", .{ name, w.bits, w.group, path }); + return error.UnsupportedWidth; + } + if (rows % 4 == 0) return; + quant.say("glm: {s} has {d} rows; the row kernels read a matrix's rows in fours", .{ name, rows }); + return error.UnsupportedShape; +} + +/// Whether every matrix in `v` takes the generated kernels (one width, 4/64): no other width's pipelines needed. +pub fn allNative(v: anytype) bool { + const T = @TypeOf(v); + if (T == Affine) return v.n == 0 or (v.more == 0 and v.bits == 4 and v.group == 64); + return switch (@typeInfo(T)) { + .@"struct" => |s| inline for (s.field_names) |name| { + if (!allNative(@field(v, name))) break false; + } else true, + .@"union" => switch (v) { + inline else => |x| allNative(x), + }, + .optional => if (v) |x| allNative(x) else true, + .array => for (v) |x| { + if (!allNative(x)) break false; + } else true, + else => true, + }; +} + +const testing = std.testing; + +test "stored shapes and bytes follow the width; 4/64 keeps the 4-bit matrix's" { + try testing.expectEqual([2][2]usize{ .{ 4096, 512 }, .{ 4096, 64 } }, shapes(4096, 4096, quant.uniform).?); + try testing.expectEqual(@as(usize, 4096 * 4096 / 2), wBytes(4096, 4096, quant.uniform)); + try testing.expectEqual(@as(usize, 4096 * (4096 / 64) * 2), sBytes(4096, 4096, quant.uniform)); + const cases = [_]struct { Width, [2]usize }{ + .{ .{ .bits = 2, .group = 64 }, .{ 16, 4 } }, + .{ .{ .bits = 3, .group = 32 }, .{ 24, 8 } }, + .{ .{ .bits = 5, .group = 64 }, .{ 40, 4 } }, + .{ .{ .bits = 6, .group = 128 }, .{ 48, 2 } }, + .{ .{ .bits = 8, .group = 64 }, .{ 64, 4 } }, + }; + for (cases) |c| { + const s = shapes(7, 256, c[0]).?; + try testing.expectEqual(c[1][0], s[0][1]); + try testing.expectEqual(c[1][1], s[1][1]); + try testing.expectEqual(7 * 4 * s[0][1], wBytes(7, 256, c[0])); + } + try testing.expect(shapes(1, 96, .{ .bits = 4, .group = 64 }) == null); +} + +test "a column cut keeps whole groups and words at every width" { + try testing.expectEqual([2][2]usize{ .{ 2048, 2048 }, .{ 128, 128 } }, cut(4096, 4096, quant.uniform).?); + try testing.expectEqual([2][2]usize{ .{ 640, 640 }, .{ 32, 32 } }, cut(1024, 1024, .{ .bits = 5, .group = 64 }).?); + try testing.expect(cut(32, 1024, quant.uniform) == null); + try testing.expect(cut(64, 1024, .{ .bits = 6, .group = 128 }) == null); +} + +test "every device path admits 4/64; the oQ widths where their kernels read them, others refused by name" { + for (std.enums.values(Path)) |p| { + try admit("x", p, quant.uniform, 64); + try testing.expectError(error.UnsupportedWidth, admit("x", p, .{ .bits = 7, .group = 64 }, 64)); + try testing.expectError(error.UnsupportedWidth, admit("x", p, .{ .bits = 4, .group = 16 }, 64)); + try testing.expectError(error.UnsupportedShape, admit("x", p, quant.uniform, 66)); + } + for ([_]Width{ .{ .bits = 2, .group = 64 }, .{ .bits = 5, .group = 64 }, .{ .bits = 6, .group = 64 }, .{ .bits = 8, .group = 64 }, .{ .bits = 3, .group = 32 }, .{ .bits = 5, .group = 128 } }) |w| + for ([_]Path{ .kda_in, .kda_gate, .kda_out, .mla_in, .mla_q, .mla_out, .dense_mlp, .shared_expert, .routed_experts, .head, .eh_proj }) |p| try admit("x", p, w, 64); + try admit("x", .embed, .{ .bits = 8, .group = 64 }, 64); + try testing.expectError(error.UnsupportedWidth, admit("x", .embed, .{ .bits = 2, .group = 64 }, 64)); + try testing.expectError(error.UnsupportedWidth, admit("x", .kv_b, .{ .bits = 8, .group = 64 }, 64)); +} + +test "a matrix of another width or a split stack needs the other widths' pipelines; 4/64 and none do not" { + const r: Ref = .{ .buf = undefined }; + const q: Affine = .{ .w = r, .s = r, .b = r, .n = 64, .k = 128, .bits = 4, .group = 64 }; + const none: Affine = .{ .w = r, .s = r, .b = r, .n = 0, .k = 128, .bits = 0, .group = 0 }; + try testing.expect(allNative(.{ q, none, @as(?Affine, null), [2]Affine{ q, q } })); + var eight = q; + eight.bits = 8; + try testing.expect(!allNative(.{ q, [2]Affine{ q, eight } })); + var split = q; + split.more = 1; + try testing.expect(!allNative(@as(?Affine, split))); +} diff --git a/zig/src/families/glm/forward.zig b/zig/src/families/glm/forward.zig index 018b99d3d3..de761d70fd 100644 --- a/zig/src/families/glm/forward.zig +++ b/zig/src/families/glm/forward.zig @@ -113,7 +113,7 @@ pub fn size(w: usize, h: usize, d: usize) mtl.Size { } /// x [rows, K] through a 4-bit matrix: qmv_rows, MLX's one-row qmv_fast sums for every row. -pub fn qmv(x: *const Ctx, e: mtl.ComputeEncoder, pipe: mtl.Pipeline, in: Ref, q: wts.Q4, out: Ref, rows: u32) void { +pub fn qmv(x: *const Ctx, e: mtl.ComputeEncoder, pipe: mtl.Pipeline, in: Ref, q: wts.Affine, out: Ref, rows: u32) void { e.setPipeline(pipe); bind(e, 0, .{ in, q.w, q.s, q.b, out }); _ = x; diff --git a/zig/src/families/glm/glm.zig b/zig/src/families/glm/glm.zig index c9eaf4f724..1e116ccb3c 100644 --- a/zig/src/families/glm/glm.zig +++ b/zig/src/families/glm/glm.zig @@ -13,6 +13,7 @@ pub const pool = @import("pool.zig"); pub const backend = @import("backend.zig"); pub const mirror = @import("mirror.zig"); pub const timing = @import("timing.zig"); +pub const affine = @import("affine.zig"); test { _ = config; @@ -31,4 +32,8 @@ test { _ = @import("kept.zig"); _ = mirror; _ = timing; + _ = @import("quant.zig"); + _ = @import("affine.zig"); + _ = @import("names.zig"); + _ = @import("weights_test.zig"); } diff --git a/zig/src/families/glm/load_plan.zig b/zig/src/families/glm/load_plan.zig index 13f032535c..53a0dddbcd 100644 --- a/zig/src/families/glm/load_plan.zig +++ b/zig/src/families/glm/load_plan.zig @@ -3,6 +3,7 @@ const std = @import("std"); const mtl = @import("metal"); const cfg = @import("config.zig"); const wts = @import("weights.zig"); +const affine = @import("affine.zig"); const st = @import("state.zig"); const prompt_mod = @import("prompt.zig"); const ep_mod = @import("ep.zig"); @@ -27,6 +28,70 @@ pub fn planBytes(gpa: std.mem.Allocator, device: mtl.Device, dir: []const u8, c: return plan.bytes; } +/// Bytes of layer i's dense buffer (everything but routed experts) at its matrices' widths, generously rounded. +pub fn layerBytes(l: anytype, i: usize) !usize { + const c = l.c; + const D: usize = c.hidden; + const a = wts.prefix ++ "layers.{d}.self_attn."; + const m = wts.prefix ++ "layers.{d}.mlp."; + var n: usize = 0; + const q = struct { + fn b(rows: usize, k: usize, w: wts.Width) usize { + return affine.bytes(rows, k, w) + 3 * 256; + } + }; + n += 2 * (24 * 4 * D * 2 + 1024) + 2 * (D * 2 + 256); + switch (c.kind(@intCast(i))) { + .kda => { + const W: usize = c.kdaWidth(); + n += try stackBytes(l, a, .{i}, &wts.kda_in, c.kdaProj(), D, 64) + q.b(W, c.kda_dim, try l.widthAt(a ++ "f_b_proj", .{i})); + n += q.b(W, c.kda_dim, try l.widthAt(a ++ "g_b_proj", .{i})) + q.b(D, W, try l.widthAt(a ++ "o_proj", .{i})) + 3 * W * c.conv * 4 + 8 * 1024 + W * 4; + }, + .mla => { + n += try stackBytes(l, a, .{i}, &wts.mla_x, c.xProj(), D, 64) + try stackBytes(l, a, .{i}, &wts.mla_qr, c.qrProj(), c.q_lora, 1); + n += q.b(D, c.mla_heads * c.v_dim, try l.widthAt(a ++ "o_proj", .{i})) + 16 * 1024 + c.i_dim * D * 2; + if (l.has(a ++ "embed_q.weight", .{i})) { + n += q.b(c.mla_heads * c.kv_lora, c.nope, try l.widthAt(a ++ "embed_q", .{i})) + q.b(c.mla_heads * c.v_dim, c.kv_lora, try l.widthAt(a ++ "unembed_out", .{i})); + } else n += q.b(c.mla_heads * (c.nope + c.v_dim), c.kv_lora, try l.widthAt(a ++ "kv_b_proj", .{i})); + }, + } + if (c.isMoe(@intCast(i))) { + n += c.experts * D * 2 + c.experts * 4 + try stackBytes(l, m ++ "shared_experts.", .{i}, &wts.gate_up, 2 * c.moe_inter, D, 1) + q.b(D, c.moe_inter, try l.widthAt(m ++ "shared_experts.down_proj", .{i})) + 2048; + } else n += try stackBytes(l, m, .{i}, &wts.gate_up, 2 * c.dense_inter, D, 1) + q.b(D, c.dense_inter, try l.widthAt(m ++ "down_proj", .{i})); + const eh = wts.prefix ++ "layers.{d}.eh_proj"; + if (i == c.layers) n += 3 * (D * 2 + 256) + if (l.has(eh ++ ".scales", .{i})) q.b(D, 2 * D, try l.widthAt(eh, .{i})) else D * 2 * D * 2 + 256; + return n; +} + +/// Bytes of layer i's routed experts on this Mac at the first one's widths (the rest must match). +pub fn expertBytes(l: anytype, i: usize) !usize { + const c = l.c; + const D: usize = c.hidden; + const every = i == c.layers; + const m: usize = if (every) c.moe_inter else c.inter[1] - c.inter[0]; + const lo: usize = if (every) 0 else c.own[0]; + const per = affine.bytes(m, D, try l.widths.expert(i, lo, "gate_proj")) + affine.bytes(m, D, try l.widths.expert(i, lo, "up_proj")) + affine.bytes(D, m, try l.widths.expert(i, lo, "down_proj")); + return (if (every) c.experts else c.own[1] - c.own[0]) * per + 9 * 256; +} + +/// A stack's bytes: its rows, padded to `pad` a run (a part each once widths differ), at the costliest bytes a row. +fn stackBytes(l: anytype, comptime fmt: []const u8, args: anytype, parts: []const []const u8, rows: usize, k: usize, pad: usize) !usize { + var name: [200]u8 = undefined; + const base = try std.fmt.bufPrint(&name, fmt, args); + var most: usize = 0; + var first: ?wts.Width = null; + var mixed = false; + for (parts) |p| { + var full: [240]u8 = undefined; + const w = try l.widths.of(try std.fmt.bufPrint(&full, "{s}{s}", .{ base, p })); + most = @max(most, affine.bytes(1, k, w)); + mixed = mixed or (first != null and !std.meta.eql(first.?, w)); + if (first == null) first = w; + } + const runs = if (mixed) parts.len else 1; + return (std.mem.alignForward(usize, rows, pad) + (runs - 1) * pad) * most + runs * 3 * 256; +} + /// The fewest bytes of caches and buffers a load of `cap` tokens takes (the shortest prompt chunks), counted, not allocated. pub fn leastArena(gpa: std.mem.Allocator, c: *const cfg.Config, cap: u32, chunked: bool) !usize { var dry: st.Arena = .{ .device = undefined, .gpa = gpa, .dry = true }; diff --git a/zig/src/families/glm/names.zig b/zig/src/families/glm/names.zig new file mode 100644 index 0000000000..abea94eabb --- /dev/null +++ b/zig/src/families/glm/names.zig @@ -0,0 +1,104 @@ +//! GLM-5.3-Flash's checkpoint index in either layout (the original or mlx-lm's), each tensor by its original name. +const std = @import("std"); +const mtl = @import("metal"); +const st = @import("../../core/safetensors.zig"); +const load_fix = @import("weights_fix.zig"); + +const Src = load_fix.Src; +const lm = "model.language_model."; + +/// A stored tensor's name in the original layout, or null for another tower's (the vision tower is never read). +pub fn original(buf: []u8, name: []const u8, mtp_layer: u32) error{NoSpaceLeft}!?[]const u8 { + if (std.mem.startsWith(u8, name, lm) or std.mem.startsWith(u8, name, "lm_head.")) return name; + const mlx = "language_model."; + if (!std.mem.startsWith(u8, name, mlx)) return null; + var short = name[mlx.len..]; + if (std.mem.startsWith(u8, short, "lm_head.")) return short; + var num: [16]u8 = undefined; + var head: []const u8 = ""; + if (std.mem.startsWith(u8, short, "model.")) { + short = short["model.".len..]; + } else if (std.mem.startsWith(u8, short, "mtp.0.")) { // the MTP block as the layer after the backbone's + short = short["mtp.0.".len..]; + if (std.mem.startsWith(u8, short, "block.")) short = short["block.".len..]; + if (std.mem.eql(u8, short, "norm.weight")) short = "shared_head.norm.weight"; + head = try std.fmt.bufPrint(&num, "layers.{d}.", .{mtp_layer}); + } else return null; + inline for (.{ "attn", "ffn" }) |site| inline for (.{ "fn", "base", "scale" }) |part| { + const tail = "." ++ site ++ "_hc." ++ part; + if (std.mem.endsWith(u8, short, tail)) return try std.fmt.bufPrint(buf, lm ++ "{s}{s}.hc_" ++ site ++ "_" ++ part, .{ head, short[0 .. short.len - tail.len] }); + }; + const gate = ".self_attn.forget_gate."; + if (std.mem.indexOf(u8, short, gate)) |at| return try std.fmt.bufPrint(buf, lm ++ "{s}{s}.self_attn.{s}", .{ head, short[0..at], short[at + gate.len ..] }); + return try std.fmt.bufPrint(buf, lm ++ "{s}{s}", .{ head, short }); +} + +/// The index and every shard's header: each language-model tensor's shard and byte range, by its original name. +pub fn read(gpa: std.mem.Allocator, arena: std.mem.Allocator, dir: []const u8, mtp_layer: u32, shards: *std.ArrayList([:0]const u8), names: *std.StringHashMapUnmanaged(Src)) !void { + const path = try std.fmt.allocPrintSentinel(arena, "{s}/model.safetensors.index.json", .{dir}, 0); + const f = try mtl.MappedFile.open(path); + defer f.deinit(); + const doc = try std.json.parseFromSliceLeaky(std.json.Value, arena, f.bytes[0..f.size], .{ .allocate = .alloc_always }); + var files: std.StringArrayHashMapUnmanaged(void) = .empty; + var it = doc.object.get("weight_map").?.object.iterator(); + while (it.next()) |kv| try files.put(arena, kv.value_ptr.string, {}); + for (files.keys(), 0..) |name, si| { + const full = try std.fmt.allocPrintSentinel(arena, "{s}/{s}", .{ dir, name }, 0); + try shards.append(gpa, full); + const fd = std.c.open(full, .{ .ACCMODE = .RDONLY }); + if (fd < 0) return error.OpenFailed; + defer _ = std.c.close(fd); + var head: [8]u8 = undefined; + try load_fix.readAll(fd, &head, 0); + const hlen = std.mem.readInt(u64, &head, .little); + const header = try arena.alloc(u8, hlen); + try load_fix.readAll(fd, header, 8); + const end = std.c.lseek(fd, 0, std.c.SEEK.END); + const data = 8 + hlen; + const doc2 = try std.json.parseFromSliceLeaky(std.json.Value, arena, header, .{}); + var eit = doc2.object.iterator(); + while (eit.next()) |kv| { + var buf: [240]u8 = undefined; + // the language model's tensors only: the vision tower's (5-D patch kernels among them) are never read + const tensor = (try original(&buf, kv.key_ptr.*, mtp_layer)) orelse continue; + const o = kv.value_ptr.object; + const dtype = st.DType.parse(o.get("dtype").?.string) orelse return error.UnsupportedDType; + const dims = o.get("shape").?.array.items; + if (dims.len > st.max_rank) return error.RankTooHigh; + var shape: [st.max_rank]usize = @splat(1); + for (dims, 0..) |d, j| shape[j] = @intCast(d.integer); + const offs = o.get("data_offsets").?.array.items; + const b: u64 = @intCast(offs[0].integer); + const e2: u64 = @intCast(offs[1].integer); + if (e2 < b or data + e2 > @as(u64, @intCast(end))) return error.BadSafetensors; + const got = try names.getOrPut(arena, try arena.dupe(u8, tensor)); + if (got.found_existing) { // one tensor stored under both layouts' names, or twice + @import("quant.zig").say("glm: {s} is stored twice in the checkpoint", .{tensor}); + return error.BadSafetensors; + } + got.value_ptr.* = .{ .shard = @intCast(si), .off = data + b, .len = e2 - b, .dtype = dtype, .shape = shape, .rank = @intCast(dims.len) }; + } + } +} + +test "mlx-lm's names map onto the original layout's, the MTP block onto layer 45" { + var buf: [240]u8 = undefined; + const cases = [_][2][]const u8{ + .{ "language_model.model.embed_tokens.weight", "model.language_model.embed_tokens.weight" }, + .{ "language_model.lm_head.scales", "lm_head.scales" }, + .{ "lm_head.weight", "lm_head.weight" }, + .{ "model.language_model.layers.3.self_attn.kv_b_proj.weight", "model.language_model.layers.3.self_attn.kv_b_proj.weight" }, + .{ "language_model.model.layers.7.attn_hc.fn", "model.language_model.layers.7.hc_attn_fn" }, + .{ "language_model.model.layers.7.ffn_hc.scale", "model.language_model.layers.7.hc_ffn_scale" }, + .{ "language_model.model.layers.3.mlp.switch_mlp.up_proj.biases", "model.language_model.layers.3.mlp.switch_mlp.up_proj.biases" }, + .{ "language_model.model.layers.3.self_attn.embed_q.weight", "model.language_model.layers.3.self_attn.embed_q.weight" }, + .{ "language_model.model.layers.0.self_attn.forget_gate.A_log", "model.language_model.layers.0.self_attn.A_log" }, + .{ "language_model.mtp.0.block.self_attn.o_proj.weight", "model.language_model.layers.45.self_attn.o_proj.weight" }, + .{ "language_model.mtp.0.eh_proj.weight", "model.language_model.layers.45.eh_proj.weight" }, + .{ "language_model.mtp.0.enorm.weight", "model.language_model.layers.45.enorm.weight" }, + .{ "language_model.mtp.0.norm.weight", "model.language_model.layers.45.shared_head.norm.weight" }, + }; + for (cases) |c| try std.testing.expectEqualStrings(c[1], (try original(&buf, c[0], 45)).?); + for ([_][]const u8{ "vision_tower.blocks.0.attn.qkv.weight", "language_model.mtp.1.eh_proj.weight", "model.visual.x" }) |skip| + try std.testing.expect((try original(&buf, skip, 45)) == null); +} diff --git a/zig/src/families/glm/prompt.zig b/zig/src/families/glm/prompt.zig index e4f94f0600..8981004096 100644 --- a/zig/src/families/glm/prompt.zig +++ b/zig/src/families/glm/prompt.zig @@ -186,19 +186,19 @@ fn run(e: mtl.ComputeEncoder, grid: [3]usize, group: [3]usize) void { } /// y [M, N] = x [M, K] W^T for a 4-bit matrix (64x64 tiles; N rounded up to the weights' padded rows, y's row pitch). -fn qmm(p: *const Prompt, e: mtl.ComputeEncoder, x: Ref, q: wts.Q4, y: Ref, M: u32) void { +fn qmm(p: *const Prompt, e: mtl.ComputeEncoder, x: Ref, q: wts.Affine, y: Ref, M: u32) void { affine_mm.rowSums(e, p.mm.mm_bf16, 64, x, p.sums, M, q.k); affine_mm.dense(e, p.mm.mm_bf16, x, p.sums, q, y, M); } /// `qmm` over x's rows `ld` elements apart (a slice of a wider row). -fn qmmAt(p: *const Prompt, e: mtl.ComputeEncoder, x: Ref, ld: u32, q: wts.Q4, y: Ref, M: u32) void { +fn qmmAt(p: *const Prompt, e: mtl.ComputeEncoder, x: Ref, ld: u32, q: wts.Affine, y: Ref, M: u32) void { affine_mm.rowSumsAt(e, p.mm.mm_bf16, 64, x, p.sums, M, q.k, ld); affine_mm.denseBatch(e, p.mm.mm_bf16, x, p.sums, q, y, M, 1, .{ .x_row = @intCast(ld), .y_row = @intCast(std.mem.alignForward(u32, q.n, 64)) }); } /// `qmm` with fp32 out (TP2: one Mac's partial of a row-split projection). -fn qmmF32(p: *const Prompt, e: mtl.ComputeEncoder, x: Ref, q: wts.Q4, y: Ref, M: u32) void { +fn qmmF32(p: *const Prompt, e: mtl.ComputeEncoder, x: Ref, q: wts.Affine, y: Ref, M: u32) void { affine_mm.rowSums(e, p.mm.mm_bf16, 64, x, p.sums, M, q.k); affine_mm.dense(e, p.mm.mm_f32, x, p.sums, q, y, M); } diff --git a/zig/src/families/glm/quant.zig b/zig/src/families/glm/quant.zig new file mode 100644 index 0000000000..6b63242db6 --- /dev/null +++ b/zig/src/families/glm/quant.zig @@ -0,0 +1,174 @@ +//! GLM-5.3-Flash's MLX affine width per module: config.json's quantization table, read through core/quant. +const std = @import("std"); +const quant = @import("../../core/quant/quant.zig"); + +pub const Width = quant.mlx.Width; + +/// The width every GLM kernel is generated for: 4 bits in groups of 64. +pub const uniform: Width = .{ .bits = 4, .group = 64 }; + +/// The checkpoint's widths: the core table with its override keys canonical (`layers.N.x`, `embed_tokens`, `lm_head`). +pub const Widths = struct { + table: quant.mlx.Config, + mtp_layer: u32, + + /// The width of the module a tensor name names, in either layout; an entry the table refuses fails by name. + pub fn of(w: *const Widths, name: []const u8) !Width { + var buf: [256]u8 = undefined; + return w.table.width(try canonical(&buf, name, w.mtp_layer)) catch |err| { + say("glm: config.json quantizes {s} in a format the engine does not read (MLX affine, 2 to 8 bits)", .{name}); + return err; + }; + } + + /// Layer `i`'s routed expert `e`'s `proj`: its own entry, else its layer's stack's (`switch_mlp`), else the global. + pub fn expert(w: *const Widths, i: usize, e: usize, proj: []const u8) !Width { + var a: [128]u8 = undefined; + var b: [128]u8 = undefined; + const own = try std.fmt.bufPrint(&a, "layers.{d}.mlp.experts.{d}.{s}", .{ i, e, proj }); + const stack = try std.fmt.bufPrint(&b, "layers.{d}.mlp.switch_mlp.{s}", .{ i, proj }); + const x = w.table.overrides.get(own); + const y = w.table.overrides.get(stack); + if (x != null and y != null and !std.meta.eql(x.?, y.?)) return conflict(own, stack); + return w.of(if (x != null) own else stack); + } +}; + +/// A refusal's reason, logged as an error; tests check the error itself and keep the log quiet. +pub fn say(comptime fmt: []const u8, args: anytype) void { + if (@import("builtin").is_test) return; + std.log.err(fmt, args); +} + +fn conflict(a: []const u8, b: []const u8) error{ConflictingQuantization} { + say("glm: config.json gives {s} and {s}, one module, different widths", .{ a, b }); + return error.ConflictingQuantization; +} + +/// core/quant's module key, then GLM's own: `mtp.0.[block.]` as its layer, `self_attn.forget_gate.` as `self_attn.`. +pub fn canonical(buf: []u8, name: []const u8, mtp_layer: u32) error{NoSpaceLeft}![]const u8 { + var short = @import("../../core/quant/keys.zig").canonical(name); + var layer: []const u8 = ""; + var num: [12]u8 = undefined; + if (std.mem.startsWith(u8, short, "mtp.0.")) { + short = short["mtp.0.".len..]; + if (std.mem.startsWith(u8, short, "block.")) short = short["block.".len..]; + layer = try std.fmt.bufPrint(&num, "layers.{d}.", .{mtp_layer}); + } + const gate = ".self_attn.forget_gate."; + const at = std.mem.indexOf(u8, short, gate) orelse return std.fmt.bufPrint(buf, "{s}{s}", .{ layer, short }); + return std.fmt.bufPrint(buf, "{s}{s}.self_attn.{s}", .{ layer, short[0..at], short[at + gate.len ..] }); +} + +/// The quantization table of config.json's root, or of its text_config: `quantization`, else `quantization_config`. +fn block(root: std.json.ObjectMap) ?std.json.ObjectMap { + const text: ?std.json.ObjectMap = if (root.get("text_config")) |t| (if (t == .object) t.object else null) else null; + for ([_]?std.json.ObjectMap{ root, text }) |o| if (o) |obj| for ([_][]const u8{ "quantization", "quantization_config" }) |k| { + if (obj.get(k)) |v| if (v == .object) return v.object; + }; + return null; +} + +/// config.json's widths: MLX affine only, overrides keyed by canonical module, aliases that disagree refused. +pub fn parse(a: std.mem.Allocator, json: []const u8, mtp_layer: u32) !Widths { + const doc = try std.json.parseFromSliceLeaky(std.json.Value, a, json, .{ .allocate = .alloc_always }); + if (doc != .object) return error.BadConfig; + const q = block(doc.object) orelse { + say("glm: config.json has no quantization table; the engine reads MLX affine checkpoints", .{}); + return error.UnsupportedQuantization; + }; + const stated = quant.detect(a, .{ .root = doc.object, .quantization = q }) catch |err| { + say("glm: config.json's quantization is not MLX affine at 2, 3, 4, 5, 6 or 8 bits in groups of 32, 64 or 128, one width a module", .{}); + return err; + }; + const table = switch (stated) { + .mlx => |t| t, + else => return error.UnsupportedQuantization, + }; + var out: Widths = .{ .table = .{ .global = table.global }, .mtp_layer = mtp_layer }; + var keys: std.StringHashMapUnmanaged([]const u8) = .empty; // canonical -> the key that set it + var it = table.overrides.iterator(); + while (it.next()) |kv| { + var buf: [256]u8 = undefined; + const key = try a.dupe(u8, try canonical(&buf, kv.key_ptr.*, mtp_layer)); + const got = try out.table.overrides.getOrPut(a, key); + if (got.found_existing and !std.meta.eql(got.value_ptr.*, kv.value_ptr.*)) return conflict(keys.get(key).?, kv.key_ptr.*); + got.value_ptr.* = kv.value_ptr.*; + try keys.put(a, key, kv.key_ptr.*); + } + return out; +} + +const testing = std.testing; + +test "override keys resolve by module whatever their prefix, the MTP block as its layer" { + var buf: [256]u8 = undefined; + const cases = [_][2][]const u8{ + .{ "language_model.model.layers.3.self_attn.q_b_proj", "layers.3.self_attn.q_b_proj" }, + .{ "model.language_model.layers.3.self_attn.q_b_proj", "layers.3.self_attn.q_b_proj" }, + .{ "model.layers.3.self_attn.q_b_proj", "layers.3.self_attn.q_b_proj" }, + .{ "language_model.lm_head", "lm_head" }, + .{ "lm_head", "lm_head" }, + .{ "language_model.model.embed_tokens", "embed_tokens" }, + .{ "language_model.mtp.0.block.self_attn.o_proj", "layers.45.self_attn.o_proj" }, + .{ "language_model.mtp.0.eh_proj", "layers.45.eh_proj" }, + .{ "model.layers.7.self_attn.forget_gate.f_b_proj", "layers.7.self_attn.f_b_proj" }, + .{ "vision_model.blocks.0.attn.qkv", "vision_model.blocks.0.attn.qkv" }, + }; + for (cases) |c| try testing.expectEqualStrings(c[1], try canonical(&buf, c[0], 45)); +} + +/// An oQ-style table: global 4/64, overrides at 8, 5 and 6 bits under both prefixes, the MTP block's, a stack's. +const mixed = + \\{"text_config": {}, "quantization": {"group_size": 64, "bits": 4, "mode": "affine", + \\ "language_model.lm_head": {"group_size": 64, "bits": 8}, + \\ "language_model.model.embed_tokens": {"group_size": 64, "bits": 8}, + \\ "language_model.model.layers.0.self_attn.q_proj": {"group_size": 64, "bits": 5}, + \\ "model.layers.0.self_attn.o_proj": {"group_size": 32, "bits": 6}, + \\ "language_model.model.layers.3.mlp.shared_experts.gate_proj": {"group_size": 128, "bits": 3}, + \\ "language_model.model.layers.3.mlp.switch_mlp.down_proj": {"group_size": 64, "bits": 2}, + \\ "language_model.mtp.0.block.self_attn.kv_b_proj": {"group_size": 64, "bits": 6}, + \\ "language_model.model.layers.5.mlp.shared_experts.up_proj": {"group_size": 64, "bits": 7}}} +; + +test "each GLM tensor of a mixed oQ table resolves to its module's width, the rest to the global" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const w = try parse(arena.allocator(), mixed, 45); + const cases = [_]struct { []const u8, Width }{ + .{ "lm_head", .{ .bits = 8, .group = 64 } }, + .{ "model.language_model.embed_tokens", .{ .bits = 8, .group = 64 } }, + .{ "model.language_model.layers.0.self_attn.q_proj", .{ .bits = 5, .group = 64 } }, + .{ "model.language_model.layers.0.self_attn.k_proj", uniform }, + .{ "model.language_model.layers.0.self_attn.o_proj", .{ .bits = 6, .group = 32 } }, + .{ "model.language_model.layers.1.self_attn.o_proj", uniform }, + .{ "model.language_model.layers.3.mlp.shared_experts.gate_proj", .{ .bits = 3, .group = 128 } }, + .{ "model.language_model.layers.45.self_attn.kv_b_proj", .{ .bits = 6, .group = 64 } }, + .{ "model.language_model.layers.44.self_attn.kv_b_proj", uniform }, + }; + for (cases) |c| try testing.expectEqual(c[1], try w.of(c[0])); + try testing.expectEqual(Width{ .bits = 2, .group = 64 }, try w.expert(3, 17, "down_proj")); + try testing.expectEqual(uniform, try w.expert(3, 17, "gate_proj")); + try testing.expectEqual(uniform, try w.expert(4, 0, "down_proj")); + try testing.expectError(error.UnsupportedQuantization, w.of("model.language_model.layers.5.mlp.shared_experts.up_proj")); +} + +test "a uniform table gives every module 4/64; aliases that disagree, no table and other modes are refused" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const flat = try parse(a, "{\"quantization\": {\"group_size\": 64, \"bits\": 4, \"mode\": \"affine\"}}", 45); + try testing.expectEqual(uniform, try flat.of("model.language_model.layers.12.self_attn.q_proj")); + try testing.expectEqual(uniform, try flat.expert(12, 287, "up_proj")); + try testing.expectEqual(@as(u32, 0), flat.table.overrides.count()); + const agree = "{\"quantization_config\": {\"group_size\": 64, \"bits\": 4, \"model.layers.1.mlp.down_proj\": {\"group_size\": 64, \"bits\": 8}, \"language_model.model.layers.1.mlp.down_proj\": {\"group_size\": 64, \"bits\": 8}}}"; + try testing.expectEqual(Width{ .bits = 8, .group = 64 }, try (try parse(a, agree, 45)).of("model.language_model.layers.1.mlp.down_proj")); + const clash = "{\"quantization\": {\"group_size\": 64, \"bits\": 4, \"model.layers.1.mlp.down_proj\": {\"group_size\": 64, \"bits\": 8}, \"language_model.model.layers.1.mlp.down_proj\": {\"group_size\": 64, \"bits\": 6}}}"; + try testing.expectError(error.UnsupportedQuantization, parse(a, clash, 45)); // core/quant refuses the aliases first + const both = "{\"quantization\": {\"group_size\": 64, \"bits\": 4, \"model.layers.3.mlp.experts.2.up_proj\": {\"group_size\": 64, \"bits\": 8}, \"model.layers.3.mlp.switch_mlp.up_proj\": {\"group_size\": 64, \"bits\": 2}}}"; + const mtp = "{\"quantization\": {\"group_size\": 64, \"bits\": 4, \"language_model.mtp.0.block.self_attn.o_proj\": {\"group_size\": 64, \"bits\": 6}, \"model.layers.45.self_attn.o_proj\": {\"group_size\": 64, \"bits\": 8}}}"; + try testing.expectError(error.ConflictingQuantization, parse(a, mtp, 45)); // GLM's own alias: the MTP block is layer 45 + try testing.expectError(error.ConflictingQuantization, (try parse(a, both, 45)).expert(3, 2, "up_proj")); + try testing.expectError(error.UnsupportedQuantization, parse(a, "{\"text_config\": {}}", 45)); + try testing.expectError(error.UnsupportedQuantization, parse(a, "{\"quantization\": {\"group_size\": 64, \"bits\": 4, \"mode\": \"mxfp4\"}}", 45)); +} diff --git a/zig/src/families/glm/weights.zig b/zig/src/families/glm/weights.zig index 9fae40b3c9..9e7d19b77d 100644 --- a/zig/src/families/glm/weights.zig +++ b/zig/src/families/glm/weights.zig @@ -3,46 +3,43 @@ const std = @import("std"); const mtl = @import("metal"); const st = @import("../../core/safetensors.zig"); const cfg = @import("config.zig"); +const quant = @import("quant.zig"); +const affine = @import("affine.zig"); +const layout = @import("names.zig"); +const load_plan = @import("load_plan.zig"); -pub const Ref = struct { - buf: mtl.Buffer, - off: usize = 0, +pub const Width = quant.Width; - pub fn at(r: Ref, bytes: usize) Ref { - return .{ .buf = r.buf, .off = r.off + bytes }; - } - - pub fn addr(r: Ref) [*]u8 { - return r.buf.contents() + r.off; - } -}; +pub const Ref = affine.Ref; +pub const Affine = affine.Affine; -/// A 4-bit affine matrix in groups of 64: w [n, k/8] u32, s and b [n, k/64] bf16 (stacked: [e, n, ...]). -pub const Q4 = struct { - w: Ref, - s: Ref, - b: Ref, - n: u32, - k: u32, - - pub fn wBytes(n: usize, k: usize) usize { - return n * k / 2; - } - - pub fn sBytes(n: usize, k: usize) usize { - return n * (k / 64) * 2; - } -}; +/// The members of the fused matrices, in their rows' order. +pub const kda_in = [_][]const u8{ "q_proj", "k_proj", "v_proj", "f_a_proj", "g_a_proj", "b_proj" }; +pub const mla_x = [_][]const u8{ "q_a_proj", "kv_a_proj_with_mqa", "indexer.wk", "indexer.weights_proj" }; +pub const mla_qr = [_][]const u8{ "q_b_proj", "indexer.wq_b" }; +pub const gate_up = [_][]const u8{ "gate_proj", "up_proj" }; pub const Hc = struct { fnp: Ref, base: Ref, scale: Ref }; -pub const Kda = struct { in_proj: Q4, f_b: Q4, g_b: Q4, o_proj: Q4, conv_w: Ref, a_log: Ref, a: Ref, dt_bias: Ref, o_norm: Ref }; - -pub const Mla = struct { x_proj: Q4, qr_proj: Q4, kv_b: Q4, o_proj: Q4, q_norm: Ref, kv_norm: Ref, k_norm_w: Ref, k_norm_b: Ref, ape: Ref, igate: Ref }; // igate packed as a router +pub const Kda = struct { in_proj: Affine, f_b: Affine, g_b: Affine, o_proj: Affine, conv_w: Ref, a_log: Ref, a: Ref, dt_bias: Ref, o_norm: Ref }; + +pub const Mla = struct { + x_proj: Affine, + qr_proj: Affine, + kv_b: Affine, // none when the checkpoint stores mlx-lm's absorbed pair + o_proj: Affine, + q_norm: Ref, + kv_norm: Ref, + k_norm_w: Ref, + k_norm_b: Ref, + ape: Ref, + igate: Ref, // packed as a router + absorbed: ?[2]Affine = null, // mlx-lm's embed_q [H kv_lora, nope] and unembed_out [H v, kv_lora] +}; -pub const Dense = struct { gate_up: Q4, down: Q4 }; +pub const Dense = struct { gate_up: Affine, down: Affine }; -pub const Moe = struct { router: Ref, bias: Ref, gate: Q4, up: Q4, down: Q4, sh_gate_up: Q4, sh_down: Q4 }; +pub const Moe = struct { router: Ref, bias: Ref, gate: Affine, up: Affine, down: Affine, sh_gate_up: Affine, sh_down: Affine }; pub const Layer = struct { hc: ?[2]Hc, // attention's and the MLP's hyper-connections (none on the MTP layer) @@ -52,13 +49,13 @@ pub const Layer = struct { mlp: union(enum) { dense: Dense, moe: Moe }, }; -pub const Mtp = struct { eh_proj: Q4, enorm: Ref, hnorm: Ref, norm: Ref }; +pub const Mtp = struct { eh_proj: Affine, enorm: Ref, hnorm: Ref, norm: Ref, eh_bf16: ?Ref = null }; // stored bf16 pub const Weights = struct { layers: [cfg.max_layers]Layer = undefined, // the backbone's, then the MTP layer at index c.layers mtp: ?Mtp = null, - embed: Q4 = undefined, - head: Q4 = undefined, + embed: Affine = undefined, + head: Affine = undefined, norm: Ref = undefined, buffers: std.ArrayList(mtl.Buffer) = .empty, bytes: usize = 0, @@ -70,16 +67,13 @@ pub const Weights = struct { } }; -const prefix = "model.language_model."; +pub const prefix = "model.language_model."; const load_fix = @import("weights_fix.zig"); -const paced_read = @import("../../core/paced_read.zig"); const Src = load_fix.Src; const Fix = load_fix.Fix; const Transform = load_fix.Transform; -const readAll = load_fix.readAll; - -const Copy = struct { shard: u32, off: u64, len: u64, dst: [*]u8 }; +const Copy = load_fix.Copy; const Loader = struct { gpa: std.mem.Allocator, @@ -93,6 +87,7 @@ const Loader = struct { copies: std.ArrayList(Copy) = .empty, transforms: std.ArrayList(Transform) = .empty, dry: bool = false, // plan and check only: no buffers, no reads + widths: quant.Widths = undefined, // the buffer being filled and its next free byte cur: ?mtl.Buffer = null, cur_len: usize = 0, @@ -109,10 +104,24 @@ const Loader = struct { fn expect(s: Src, name: []const u8, dtype: st.DType, shape: []const usize) !void { if (s.dtype == dtype and s.rank == shape.len and std.mem.eql(usize, s.shape[0..s.rank], shape)) return; - std.log.err("glm: {s} is {t} {any}; the kernels read {t} {any}", .{ name, s.dtype, s.shape[0..s.rank], dtype, shape }); + quant.say("glm: {s} is {t} {any}; the kernels read {t} {any}", .{ name, s.dtype, s.shape[0..s.rank], dtype, shape }); return error.UnexpectedTensor; } + /// Refuses two members of one fused matrix stored at different widths. + fn mixed(a: []const u8, aw: Width, b: []const u8, bw: Width) error{MixedWidths} { + quant.say("glm: {s} is {d}-bit in groups of {d} and {s} {d}-bit in groups of {d}; one matrix holds both", .{ a, aw.bits, aw.group, b, bw.bits, bw.group }); + return error.MixedWidths; + } + + /// Input columns [lo, lo + k) of `module` at `w`, refused unless whole groups and words. + fn cutOf(module: []const u8, lo: usize, k: usize, w: Width) error{UnexpectedTensor}![2][2]usize { + return affine.cut(lo, k, w) orelse { + quant.say("glm: {s}: columns [{d}, {d}) split a group or word at {d} bits in groups of {d}", .{ module, lo, lo + k, w.bits, w.group }); + return error.UnexpectedTensor; + }; + } + /// A new buffer of `len` bytes that the next `take`s fill. fn begin(l: *Loader, len: usize) !void { const b: ?mtl.Buffer = if (l.dry) null else try l.device.buffer(@max(len, 16), mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); @@ -138,11 +147,13 @@ const Loader = struct { try l.copies.append(l.gpa, .{ .shard = s.shard, .off = s.off, .len = s.len, .dst = dst.addr() }); } - /// A tensor as stored (bf16 norms, fp32 biases), checked. + /// A tensor as stored (bf16 norms, fp32 biases), checked; fp32 stored as bf16 (mlx-lm's) is widened, exactly. fn plain(l: *Loader, comptime fmt: []const u8, args: anytype, dtype: st.DType, shape: []const usize) !Ref { const s = try l.src(fmt, args); var name: [200]u8 = undefined; - try expect(s, try std.fmt.bufPrint(&name, fmt, args), dtype, shape); + const widen = dtype == .f32 and s.dtype == .bf16; + try expect(s, try std.fmt.bufPrint(&name, fmt, args), if (widen) .bf16 else dtype, shape); + if (widen) return l.fixed(s, 2 * s.len, .to_f32); const r = l.take(s.len); try l.copyTo(s, r); return r; @@ -154,54 +165,144 @@ const Loader = struct { return r; } - /// 4-bit projections that read one input, stacked by rows in order; each part checked [n_i, k]. - fn q4(l: *Loader, k: usize, comptime fmt: []const u8, parts: []const []const u8, ns: []const usize, args: anytype) !Q4 { + pub fn has(l: *Loader, comptime fmt: []const u8, args: anytype) bool { + var name: [200]u8 = undefined; + return l.names.contains(std.fmt.bufPrint(&name, fmt, args) catch return false); + } + + /// No matrix: the place of one the checkpoint stores another way (refused on the device before it is read). + fn none(l: *Loader, k: usize) Affine { + const r = l.take(0); + return .{ .w = r, .s = r, .b = r, .n = 0, .k = @intCast(k), .bits = 0, .group = 0 }; + } + + /// Expert `e`'s three tensors cut from layer i's mlx-lm stack `switch_mlp.proj` [experts, n, ...], checked. + fn slab(l: *Loader, i: usize, e: usize, proj: []const u8, n: usize, k: usize, w: Width) ![3]Src { + var buf: [200]u8 = undefined; + const module = try std.fmt.bufPrint(&buf, prefix ++ "layers.{d}.mlp.switch_mlp.{s}", .{ i, proj }); + const sh = affine.shapes(n, k, w) orelse return error.UnexpectedTensor; + var out: [3]Src = undefined; + for ([_][]const u8{ "weight", "scales", "biases" }, &out, 0..) |comp, *o, j| { + var full: [240]u8 = undefined; + const nm = try std.fmt.bufPrint(&full, "{s}.{s}", .{ module, comp }); + const s = l.names.get(nm) orelse return error.MissingTensor; + try expect(s, nm, if (j == 0) .u32 else .bf16, &.{ l.c.experts, sh[@min(j, 1)][0], sh[@min(j, 1)][1] }); + const per = s.len / l.c.experts; + o.* = .{ .shard = s.shard, .off = s.off + e * per, .len = per, .dtype = s.dtype, .shape = .{ n, sh[@min(j, 1)][1], 1, 1, 1 }, .rank = 2 }; + } + return out; + } + + /// This Mac's heads of mlx-lm's per-head `module` stored [heads, r, k]: rows [h0 r, (h0 + H) r) as one matrix. + fn perHead(l: *Loader, comptime fmt: []const u8, args: anytype, heads: usize, h0: usize, H: usize, r: usize, k: usize) !Affine { + var buf: [200]u8 = undefined; + const module = try std.fmt.bufPrint(&buf, fmt, args); + const w = try l.widths.of(module); + const sh = affine.shapes(r, k, w) orelse return error.UnexpectedTensor; + const out: Affine = .{ .w = l.take(affine.wBytes(H * r, k, w)), .s = l.take(affine.sBytes(H * r, k, w)), .b = l.take(affine.sBytes(H * r, k, w)), .n = @intCast(H * r), .k = @intCast(k), .bits = w.bits, .group = w.group }; + for ([_][]const u8{ "weight", "scales", "biases" }, [_]Ref{ out.w, out.s, out.b }, 0..) |comp, dst, j| { + var full: [240]u8 = undefined; + const nm = try std.fmt.bufPrint(&full, "{s}.{s}", .{ module, comp }); + const s = l.names.get(nm) orelse return error.MissingTensor; + try expect(s, nm, if (j == 0) .u32 else .bf16, &.{ heads, r, sh[@min(j, 1)][1] }); + const row = if (j == 0) affine.wBytes(1, k, w) else affine.sBytes(1, k, w); + try l.copyTo(.{ .shard = s.shard, .off = s.off + h0 * r * row, .len = H * r * row, .dtype = s.dtype, .shape = s.shape, .rank = s.rank }, dst); + } + return out; + } + + /// The width config.json gives the module `fmt` names, for a buffer's budget. + pub fn widthAt(l: *Loader, comptime fmt: []const u8, args: anytype) !Width { + var name: [200]u8 = undefined; + return l.widths.of(try std.fmt.bufPrint(&name, fmt, args)); + } + + /// The weight, scales and biases of `module` stored [all, k] at `w`, checked. + fn stored(l: *Loader, module: []const u8, all: usize, k: usize, w: Width) ![3]Src { + const sh = affine.shapes(all, k, w) orelse { + quant.say("glm: {s} has {d} input columns, which do not pack at {d} bits in groups of {d}", .{ module, k, w.bits, w.group }); + return error.UnexpectedTensor; + }; + var out: [3]Src = undefined; + for ([_][]const u8{ "weight", "scales", "biases" }, &out, 0..) |comp, *o, j| { + var full: [240]u8 = undefined; + const nm = try std.fmt.bufPrint(&full, "{s}.{s}", .{ module, comp }); + o.* = l.names.get(nm) orelse { + quant.say("glm: the checkpoint has no tensor {s}", .{nm}); + return error.MissingTensor; + }; + try expect(o.*, nm, if (j == 0) .u32 else .bf16, &sh[@min(j, 1)]); + } + return out; + } + + /// Projections that read one input, stacked by rows in order; each part checked [n_i, k]. + fn mat(l: *Loader, path: affine.Path, k: usize, comptime fmt: []const u8, parts: []const []const u8, ns: []const usize, args: anytype) !Affine { const zeros: [8]usize = @splat(0); - return l.q4Rows(k, fmt, parts, zeros[0..parts.len], ns, ns, 1, args); + return l.stack(path, k, fmt, parts, zeros[0..parts.len], ns, ns, 1, args); } - /// 4-bit rows [lo, lo + n) of each part stored [all, k], stacked; room for `pad`-row tiles. - fn q4Rows(l: *Loader, k: usize, comptime fmt: []const u8, parts: []const []const u8, los: []const usize, ns: []const usize, alls: []const usize, pad: usize, args: anytype) !Q4 { - var n: usize = 0; - for (ns) |x| n += x; - const room = std.mem.alignForward(usize, n, pad); - const out: Q4 = .{ .w = l.take(Q4.wBytes(room, k)), .s = l.take(Q4.sBytes(room, k)), .b = l.take(Q4.sBytes(room, k)), .n = @intCast(n), .k = @intCast(k) }; + /// Rows [lo, lo + n) of each part stored [all, k], stacked, padded to `pad`-row tiles; mixed widths: a run a part. + fn stack(l: *Loader, path: affine.Path, k: usize, comptime fmt: []const u8, parts: []const []const u8, los: []const usize, ns: []const usize, alls: []const usize, pad: usize, args: anytype) !Affine { + var name: [200]u8 = undefined; + const base = try std.fmt.bufPrint(&name, fmt, args); + var srcs: [affine.max_runs][3]Src = undefined; + var ws: [affine.max_runs]Width = undefined; + var mixed_widths = false; + for (parts, alls, 0..) |p, all, i| { + var full: [240]u8 = undefined; + const module = try std.fmt.bufPrint(&full, "{s}{s}", .{ base, p }); + ws[i] = try l.widths.of(module); + srcs[i] = try l.stored(module, all, k, ws[i]); + mixed_widths = mixed_widths or !std.meta.eql(ws[0], ws[i]); + } + // 0.6.6's QSplit: once widths differ, every part is its own product, even beside one of its width + const runs = if (mixed_widths) parts.len else 1; + var starts: [affine.max_runs + 1]usize = undefined; // run r: parts [starts[r], starts[r + 1]) + for (0..runs) |r| starts[r] = r; + starts[runs] = parts.len; + for (parts, ws[0..parts.len], ns) |p, pw, rows| { + var full: [240]u8 = undefined; + try affine.admit(try std.fmt.bufPrint(&full, "{s}{s}", .{ base, p }), path, pw, rows); + } + var out: Affine = undefined; var row: usize = 0; - for (parts, los, ns, alls) |p, lo, rows, all| { - const comps = [_][]const u8{ "weight", "scales", "biases" }; - const dsts = [_]Ref{ out.w.at(Q4.wBytes(row, k)), out.s.at(Q4.sBytes(row, k)), out.b.at(Q4.sBytes(row, k)) }; - for (comps, dsts, 0..) |comp, dst, j| { - var name: [200]u8 = undefined; - const base = try std.fmt.bufPrint(&name, fmt, args); - var full: [240]u8 = undefined; - const nm = try std.fmt.bufPrint(&full, "{s}{s}.{s}", .{ base, p, comp }); - const s = l.names.get(nm) orelse { - std.log.err("glm: the checkpoint has no tensor {s}", .{nm}); - return error.MissingTensor; - }; - try expect(s, nm, if (j == 0) .u32 else .bf16, &.{ all, if (j == 0) k / 8 else k / 64 }); - const row_bytes: usize = if (j == 0) k / 2 else k / 64 * 2; - try l.copyTo(.{ .shard = s.shard, .off = s.off + lo * row_bytes, .len = rows * row_bytes, .dtype = s.dtype, .shape = s.shape, .rank = s.rank }, dst); + for (0..runs) |r| { + const w = ws[starts[r]]; + var n: usize = 0; + for (ns[starts[r]..starts[r + 1]]) |x| n += x; + const room = std.mem.alignForward(usize, n, pad); + const run: affine.Run = .{ .w = l.take(affine.wBytes(room, k, w)), .s = l.take(affine.sBytes(room, k, w)), .b = l.take(affine.sBytes(room, k, w)), .row = @intCast(row), .n = @intCast(n), .bits = w.bits, .group = w.group }; + if (r == 0) out = .{ .w = run.w, .s = run.s, .b = run.b, .n = 0, .k = @intCast(k), .bits = w.bits, .group = w.group } else out.runs[r - 1] = run; + var at: usize = 0; // the run's own rows + for (srcs[starts[r]..starts[r + 1]], los[starts[r]..starts[r + 1]], ns[starts[r]..starts[r + 1]]) |three, lo, x| { + const dsts = [_]Ref{ run.w.at(affine.wBytes(at, k, w)), run.s.at(affine.sBytes(at, k, w)), run.b.at(affine.sBytes(at, k, w)) }; + for (three, dsts, 0..) |s, dst, j| { + const row_bytes = if (j == 0) affine.wBytes(1, k, w) else affine.sBytes(1, k, w); + try l.copyTo(.{ .shard = s.shard, .off = s.off + lo * row_bytes, .len = x * row_bytes, .dtype = s.dtype, .shape = s.shape, .rank = s.rank }, dst); + } + at += x; } - row += rows; + row += n; } + out.n = @intCast(row); + out.more = @intCast(runs - 1); return out; } - /// One 4-bit matrix stored [n, k_all]: its input columns [lo, lo + k) (TP2: o_proj over one Mac's heads). - fn q4Cols(l: *Loader, comptime fmt: []const u8, args: anytype, n: usize, k_all: usize, lo: usize, k: usize) !Q4 { - const out: Q4 = .{ .w = l.take(Q4.wBytes(n, k)), .s = l.take(Q4.sBytes(n, k)), .b = l.take(Q4.sBytes(n, k)), .n = @intCast(n), .k = @intCast(k) }; - const comps = [_][]const u8{ "weight", "scales", "biases" }; - for (comps, [_]Ref{ out.w, out.s, out.b }, 0..) |comp, dst, j| { - var name: [200]u8 = undefined; - const base = try std.fmt.bufPrint(&name, fmt, args); - var full: [240]u8 = undefined; - const nm = try std.fmt.bufPrint(&full, "{s}.{s}", .{ base, comp }); - const s = l.names.get(nm) orelse return error.MissingTensor; - try expect(s, nm, if (j == 0) .u32 else .bf16, &.{ n, if (j == 0) k_all / 8 else k_all / 64 }); - const row_all: usize = if (j == 0) k_all / 2 else k_all / 64 * 2; - const span: [2]usize = if (j == 0) .{ lo / 2, k / 2 } else .{ lo / 64 * 2, k / 64 * 2 }; - if (!l.dry) try l.transforms.append(l.gpa, .{ .src = s, .dst = dst.addr(), .fix = .cols, .arg = .{ row_all, span[0], span[1] } }); + /// One matrix stored [n, k_all]: its input columns [lo, lo + k) (TP2: o_proj over one Mac's heads). + fn cols(l: *Loader, path: affine.Path, comptime fmt: []const u8, args: anytype, n: usize, k_all: usize, lo: usize, k: usize) !Affine { + var name: [200]u8 = undefined; + const module = try std.fmt.bufPrint(&name, fmt, args); + const w = try l.widths.of(module); + const three = try l.stored(module, n, k_all, w); + try affine.admit(module, path, w, n); + const span = try cutOf(module, lo, k, w); + const out: Affine = .{ .w = l.take(affine.wBytes(n, k, w)), .s = l.take(affine.sBytes(n, k, w)), .b = l.take(affine.sBytes(n, k, w)), .n = @intCast(n), .k = @intCast(k), .bits = w.bits, .group = w.group }; + for (three, [_]Ref{ out.w, out.s, out.b }, 0..) |s, dst, j| { + const row_all = if (j == 0) affine.wBytes(1, k_all, w) else affine.sBytes(1, k_all, w); + const at = span[@min(j, 1)]; + if (!l.dry) try l.transforms.append(l.gpa, .{ .src = s, .dst = dst.addr(), .fix = .cols, .arg = .{ row_all, at[0], at[1] } }); } return out; } @@ -210,14 +311,18 @@ const Loader = struct { fn plainRange(l: *Loader, comptime fmt: []const u8, args: anytype, dtype: st.DType, all: usize, lo: usize, n: usize, size: usize) !Ref { const s = try l.src(fmt, args); var name: [200]u8 = undefined; - try expect(s, try std.fmt.bufPrint(&name, fmt, args), dtype, &.{all}); + const widen = dtype == .f32 and s.dtype == .bf16; + try expect(s, try std.fmt.bufPrint(&name, fmt, args), if (widen) .bf16 else dtype, &.{all}); + const at = if (widen) size / 2 else size; // a stored element's bytes + const part: Src = .{ .shard = s.shard, .off = s.off + lo * at, .len = n * at, .dtype = s.dtype, .shape = s.shape, .rank = s.rank }; + if (widen) return l.fixed(part, n * size, .to_f32); const r = l.take(n * size); - try l.copyTo(.{ .shard = s.shard, .off = s.off + lo * size, .len = n * size, .dtype = s.dtype, .shape = s.shape, .rank = s.rank }, r); + try l.copyTo(part, r); return r; } - /// This Mac's routed experts' `proj` stacked [own, n, k]; by rows, gate/up keep rows `inter` and down keeps input columns `inter`. - fn experts(l: *Loader, i: usize, proj: []const u8, n_all: usize, k_all: usize, part: enum { whole, rows, cols }) !Q4 { + /// This Mac's routed experts' `proj` stacked [own, n, k] at one width; by rows, gate/up keep rows `inter`. + fn experts(l: *Loader, i: usize, proj: []const u8, n_all: usize, k_all: usize, part: enum { whole, rows, cols }) !Affine { const every = i == l.c.layers; // the MTP layer: every expert whole on every Mac (its drafts need no exchange) const lo = if (every) 0 else l.c.own[0]; const e = (if (every) l.c.experts else l.c.own[1]) - lo; @@ -225,19 +330,26 @@ const Loader = struct { const in_n: usize = l.c.inter[1] - l.c.inter[0]; const n = if (part == .rows) in_n else n_all; const k = if (part == .cols) in_n else k_all; - const out: Q4 = .{ .w = l.take(e * Q4.wBytes(n, k)), .s = l.take(e * Q4.sBytes(n, k)), .b = l.take(e * Q4.sBytes(n, k)), .n = @intCast(n), .k = @intCast(k) }; + var first: [200]u8 = undefined; + const name0 = try std.fmt.bufPrint(&first, prefix ++ "layers.{d}.mlp.experts.{d}.{s}", .{ i, lo, proj }); + const w = try l.widths.expert(i, lo, proj); + try affine.admit(name0, .routed_experts, w, n); + const span = if (part == .cols) try cutOf(name0, in_lo, in_n, w) else undefined; + const out: Affine = .{ .w = l.take(e * affine.wBytes(n, k, w)), .s = l.take(e * affine.sBytes(n, k, w)), .b = l.take(e * affine.sBytes(n, k, w)), .n = @intCast(n), .k = @intCast(k), .bits = w.bits, .group = w.group }; for (0..e) |x| { - const comps = [_][]const u8{ "weight", "scales", "biases" }; - const dsts = [_]Ref{ out.w.at(x * Q4.wBytes(n, k)), out.s.at(x * Q4.sBytes(n, k)), out.b.at(x * Q4.sBytes(n, k)) }; - for (comps, dsts, 0..) |comp, dst, j| { - const s = try l.src(prefix ++ "layers.{d}.mlp.experts.{d}.{s}.{s}", .{ i, lo + x, proj, comp }); - var name: [200]u8 = undefined; - try expect(s, try std.fmt.bufPrint(&name, "experts.{d}.{d}.{s}.{s}", .{ i, x, proj, comp }), if (j == 0) .u32 else .bf16, &.{ n_all, if (j == 0) k_all / 8 else k_all / 64 }); - const row_all: usize = if (j == 0) k_all / 2 else k_all / 64 * 2; // a stored row's bytes + var buf: [200]u8 = undefined; + const module = try std.fmt.bufPrint(&buf, prefix ++ "layers.{d}.mlp.experts.{d}.{s}", .{ i, lo + x, proj }); + const xw = try l.widths.expert(i, lo + x, proj); + if (!std.meta.eql(w, xw)) return mixed(name0, w, module, xw); + const own = l.has("{s}.weight", .{module}); + const three = if (own) try l.stored(module, n_all, k_all, w) else try l.slab(i, lo + x, proj, n_all, k_all, w); + const dsts = [_]Ref{ out.w.at(x * affine.wBytes(n, k, w)), out.s.at(x * affine.sBytes(n, k, w)), out.b.at(x * affine.sBytes(n, k, w)) }; + for (three, dsts, 0..) |s, dst, j| { + const row_all = if (j == 0) affine.wBytes(1, k_all, w) else affine.sBytes(1, k_all, w); switch (part) { .whole => try l.copyTo(s, dst), .rows => try l.copyTo(.{ .shard = s.shard, .off = s.off + in_lo * row_all, .len = in_n * row_all, .dtype = s.dtype, .shape = s.shape, .rank = s.rank }, dst), - .cols => if (!l.dry) try l.transforms.append(l.gpa, .{ .src = s, .dst = dst.addr(), .fix = .cols, .arg = .{ row_all, if (j == 0) in_lo / 2 else in_lo / 64 * 2, if (j == 0) in_n / 2 else in_n / 64 * 2 } }), + .cols => if (!l.dry) try l.transforms.append(l.gpa, .{ .src = s, .dst = dst.addr(), .fix = .cols, .arg = .{ row_all, span[@min(j, 1)][0], span[@min(j, 1)][1] } }), } } } @@ -247,14 +359,32 @@ const Loader = struct { fn hc(l: *Loader, i: usize, site: []const u8) !Hc { const D: usize = l.c.hidden; const f = try l.src(prefix ++ "layers.{d}.hc_{s}_fn", .{ i, site }); - try expect(f, "hc fn", .bf16, &.{ 24, 4 * D }); + try expect(f, "hc fn", if (f.dtype == .f32) .f32 else .bf16, &.{ 24, 4 * D }); // fp32 (mlx-lm's): bf16 if exact return .{ - .fnp = try l.fixed(f, f.len, .hc_pack), + .fnp = try l.fixed(f, 24 * 4 * D * 2, .hc_pack), .base = try l.plain(prefix ++ "layers.{d}.hc_{s}_base", .{ i, site }, .f32, &.{24}), .scale = try l.plain(prefix ++ "layers.{d}.hc_{s}_scale", .{ i, site }, .f32, &.{3}), }; } + /// Layer i's q, k, v conv taps [channels, 1, taps]: three tensors, or mlx-lm's one [3 channels, taps, 1] cut up. + fn convs(l: *Loader, i: usize, channels: usize) ![3]Src { + const a = prefix ++ "layers.{d}.self_attn."; + if (l.has(a ++ "conv1d.weight", .{i})) { + const s = try l.src(a ++ "conv1d.weight", .{i}); + try expect(s, "conv1d", .bf16, &.{ 3 * channels, l.c.conv, 1 }); + var out: [3]Src = undefined; + for (&out, 0..) |*o, p| o.* = .{ .shard = s.shard, .off = s.off + p * s.len / 3, .len = s.len / 3, .dtype = .bf16, .shape = .{ channels, 1, l.c.conv, 1, 1 }, .rank = 3 }; + return out; + } + var out: [3]Src = undefined; + inline for (.{ "q", "k", "v" }, 0..) |x, p| { + out[p] = try l.src(a ++ x ++ "_conv1d.weight", .{i}); + try expect(out[p], x ++ "_conv1d", .bf16, &.{ channels, 1, l.c.conv }); + } + return out; + } + fn kda(l: *Loader, i: usize) !Kda { const c = l.c; const D: usize = c.hidden; @@ -266,10 +396,9 @@ const Loader = struct { const Wa = W * c.tp; // every head's channels, as stored const Ha = H * c.tp; const a = prefix ++ "layers.{d}.self_attn."; - const conv_q = try l.src(a ++ "q_conv1d.weight", .{i}); - try expect(conv_q, "q_conv1d", .bf16, &.{ Wa, 1, c.conv }); - const conv = try l.fixed(conv_q, 3 * W * c.conv * 4, .conv); - const extra: [3]?Src = .{ try l.src(a ++ "k_conv1d.weight", .{i}), try l.src(a ++ "v_conv1d.weight", .{i}), null }; + const qkv = try l.convs(i, Wa); + const conv = try l.fixed(qkv[0], 3 * W * c.conv * 4, .conv); + const extra: [3]?Src = .{ qkv[1], qkv[2], null }; if (!l.dry) { const t = &l.transforms.items[l.transforms.items.len - 1]; t.extra = extra; @@ -278,10 +407,10 @@ const Loader = struct { const o_norm = try l.src(a ++ "o_norm.weight", .{i}); try expect(o_norm, "o_norm", .bf16, &.{dk}); return .{ - .in_proj = try l.q4Rows(D, a, &.{ "q_proj", "k_proj", "v_proj", "f_a_proj", "g_a_proj", "b_proj" }, &.{ ch0, ch0, ch0, 0, 0, h0 }, &.{ W, W, W, dk, dk, H }, &.{ Wa, Wa, Wa, dk, dk, Ha }, 64, .{i}), - .f_b = try l.q4Rows(dk, a, &.{"f_b_proj"}, &.{ch0}, &.{W}, &.{Wa}, 1, .{i}), - .g_b = try l.q4Rows(dk, a, &.{"g_b_proj"}, &.{ch0}, &.{W}, &.{Wa}, 1, .{i}), - .o_proj = if (c.tp == 1) try l.q4(W, a, &.{"o_proj"}, &.{D}, .{i}) else try l.q4Cols(a ++ "o_proj", .{i}, D, Wa, ch0, W), + .in_proj = try l.stack(.kda_in, D, a, &kda_in, &.{ ch0, ch0, ch0, 0, 0, h0 }, &.{ W, W, W, dk, dk, H }, &.{ Wa, Wa, Wa, dk, dk, Ha }, 64, .{i}), + .f_b = try l.stack(.kda_gate, dk, a, &.{"f_b_proj"}, &.{ch0}, &.{W}, &.{Wa}, 1, .{i}), + .g_b = try l.stack(.kda_gate, dk, a, &.{"g_b_proj"}, &.{ch0}, &.{W}, &.{Wa}, 1, .{i}), + .o_proj = if (c.tp == 1) try l.mat(.kda_out, W, a, &.{"o_proj"}, &.{D}, .{i}) else try l.cols(.kda_out, a ++ "o_proj", .{i}, D, Wa, ch0, W), .conv_w = conv, .a_log = try l.plainRange(a ++ "A_log", .{i}, .f32, Ha, h0, H, 4), .a = l.take(H * 4), @@ -299,11 +428,13 @@ const Loader = struct { const Ha = H * c.tp; // every head, as stored const gate = try l.src(a ++ "indexer.index_kpool_compress_gate", .{i}); try expect(gate, "index_kpool_compress_gate", .bf16, &.{ c.i_dim, D }); + const pair = l.has(a ++ "embed_q.weight", .{i}); return .{ - .x_proj = try l.q4Rows(D, a, &.{ "q_a_proj", "kv_a_proj_with_mqa", "indexer.wk", "indexer.weights_proj" }, &.{ 0, 0, 0, 0 }, &.{ c.q_lora, c.kv_lora, c.i_dim, c.i_heads }, &.{ c.q_lora, c.kv_lora, c.i_dim, c.i_heads }, 64, .{i}), - .qr_proj = try l.q4Rows(c.q_lora, a, &.{ "q_b_proj", "indexer.wq_b" }, &.{ h0 * c.nope, 0 }, &.{ H * c.nope, c.i_heads * c.i_dim }, &.{ Ha * c.nope, c.i_heads * c.i_dim }, 1, .{i}), - .kv_b = try l.q4Rows(c.kv_lora, a, &.{"kv_b_proj"}, &.{h0 * (c.nope + c.v_dim)}, &.{H * (c.nope + c.v_dim)}, &.{Ha * (c.nope + c.v_dim)}, 1, .{i}), - .o_proj = if (c.tp == 1) try l.q4(H * c.v_dim, a, &.{"o_proj"}, &.{D}, .{i}) else try l.q4Cols(a ++ "o_proj", .{i}, D, Ha * c.v_dim, h0 * c.v_dim, H * c.v_dim), + .x_proj = try l.stack(.mla_in, D, a, &mla_x, &.{ 0, 0, 0, 0 }, &.{ c.q_lora, c.kv_lora, c.i_dim, c.i_heads }, &.{ c.q_lora, c.kv_lora, c.i_dim, c.i_heads }, 64, .{i}), + .qr_proj = try l.stack(.mla_q, c.q_lora, a, &mla_qr, &.{ h0 * c.nope, 0 }, &.{ H * c.nope, c.i_heads * c.i_dim }, &.{ Ha * c.nope, c.i_heads * c.i_dim }, 1, .{i}), + .kv_b = if (pair) l.none(c.kv_lora) else try l.stack(.kv_b, c.kv_lora, a, &.{"kv_b_proj"}, &.{h0 * (c.nope + c.v_dim)}, &.{H * (c.nope + c.v_dim)}, &.{Ha * (c.nope + c.v_dim)}, 1, .{i}), + .absorbed = if (pair) .{ try l.perHead(a ++ "embed_q", .{i}, Ha, h0, H, c.kv_lora, c.nope), try l.perHead(a ++ "unembed_out", .{i}, Ha, h0, H, c.v_dim, c.kv_lora) } else null, + .o_proj = if (c.tp == 1) try l.mat(.mla_out, H * c.v_dim, a, &.{"o_proj"}, &.{D}, .{i}) else try l.cols(.mla_out, a ++ "o_proj", .{i}, D, Ha * c.v_dim, h0 * c.v_dim, H * c.v_dim), .q_norm = try l.plain(a ++ "q_a_layernorm.weight", .{i}, .bf16, &.{c.q_lora}), .kv_norm = try l.plain(a ++ "kv_a_layernorm.weight", .{i}, .bf16, &.{c.kv_lora}), .k_norm_w = try l.plain(a ++ "indexer.k_norm.weight", .{i}, .bf16, &.{c.i_dim}), @@ -313,42 +444,11 @@ const Loader = struct { }; } - /// Bytes of layer i's dense buffer (everything but routed experts), generously rounded. - fn layerBytes(l: *Loader, i: usize) usize { - const c = l.c; - const D: usize = c.hidden; - var n: usize = 0; - const q = struct { - fn b(rows: usize, k: usize) usize { - return Q4.wBytes(rows, k) + 2 * Q4.sBytes(rows, k) + 3 * 256; - } - }; - n += 2 * (24 * 4 * D * 2 + 1024) + 2 * (D * 2 + 256); - switch (c.kind(@intCast(i))) { - .kda => n += q.b(std.mem.alignForward(usize, c.kdaProj(), 64), D) + 2 * q.b(c.kdaWidth(), c.kda_dim) + q.b(D, c.kdaWidth()) + 3 * c.kdaWidth() * c.conv * 4 + 8 * 1024 + c.kdaWidth() * 4, - .mla => n += q.b(std.mem.alignForward(usize, c.xProj(), 64), D) + q.b(c.qrProj(), c.q_lora) + q.b(c.mla_heads * (c.nope + c.v_dim), c.kv_lora) + q.b(D, c.mla_heads * c.v_dim) + 16 * 1024 + c.i_dim * D * 2, - } - if (c.isMoe(@intCast(i))) { - n += c.experts * D * 2 + c.experts * 4 + q.b(2 * c.moe_inter, D) + q.b(D, c.moe_inter) + 2048; - } else n += q.b(2 * c.dense_inter, D) + q.b(D, c.dense_inter); - if (i == c.layers) n += q.b(D, 2 * D) + 3 * (D * 2 + 256); - return n; - } - - fn expertBytes(l: *Loader, i: usize) usize { - const c = l.c; - const D: usize = c.hidden; - const every = i == c.layers; - const m: usize = if (every) c.moe_inter else c.inter[1] - c.inter[0]; - const per = 2 * (Q4.wBytes(m, D) + 2 * Q4.sBytes(m, D)) + Q4.wBytes(D, m) + 2 * Q4.sBytes(D, m); - return (if (every) c.experts else c.own[1] - c.own[0]) * per + 9 * 256; - } - fn layer(l: *Loader, i: usize) !Layer { const c = l.c; const D: usize = c.hidden; const plain_layer = i == c.layers; // the MTP block: no hyper-connections - try l.begin(l.layerBytes(i)); + try l.begin(try load_plan.layerBytes(l, i)); var out: Layer = .{ .hc = if (plain_layer) null else .{ try l.hc(i, "attn"), try l.hc(i, "ffn") }, .in_norm = try l.plain(prefix ++ "layers.{d}.input_layernorm.weight", .{i}, .bf16, &.{D}), @@ -361,8 +461,10 @@ const Loader = struct { }; if (plain_layer) { // the MTP head's own tensors, in the layer's buffer (before the experts begin theirs) const p = prefix ++ "layers.{d}."; + const bf16 = !l.has(p ++ "eh_proj.scales", .{i}); l.w.mtp = .{ - .eh_proj = try l.q4(2 * D, p, &.{"eh_proj"}, &.{D}, .{i}), + .eh_proj = if (bf16) l.none(2 * D) else try l.mat(.eh_proj, 2 * D, p, &.{"eh_proj"}, &.{D}, .{i}), + .eh_bf16 = if (bf16) try l.plain(p ++ "eh_proj.weight", .{i}, .bf16, &.{ D, 2 * D }) else null, .enorm = try l.plain(p ++ "enorm.weight", .{i}, .bf16, &.{D}), .hnorm = try l.plain(p ++ "hnorm.weight", .{i}, .bf16, &.{D}), .norm = try l.plain(p ++ "shared_head.norm.weight", .{i}, .bf16, &.{D}), @@ -371,139 +473,39 @@ const Loader = struct { const m = prefix ++ "layers.{d}.mlp."; if (c.isMoe(@intCast(i))) { const router = try l.src(m ++ "gate.weight", .{i}); - try expect(router, "mlp.gate.weight", .bf16, &.{ c.experts, D }); + try expect(router, "mlp.gate.weight", if (router.dtype == .f32) .f32 else .bf16, &.{ c.experts, D }); var moe: Moe = .{ - .router = try l.fixed(router, router.len, .router_pack), + .router = try l.fixed(router, @as(usize, c.experts) * D * 2, .router_pack), .bias = try l.plain(m ++ "gate.e_score_correction_bias", .{i}, .f32, &.{c.experts}), - .sh_gate_up = try l.q4(D, m ++ "shared_experts.", &.{ "gate_proj", "up_proj" }, &.{ c.moe_inter, c.moe_inter }, .{i}), - .sh_down = try l.q4(c.moe_inter, m ++ "shared_experts.", &.{"down_proj"}, &.{D}, .{i}), + .sh_gate_up = try l.mat(.shared_expert, D, m ++ "shared_experts.", &gate_up, &.{ c.moe_inter, c.moe_inter }, .{i}), + .sh_down = try l.mat(.shared_expert, c.moe_inter, m ++ "shared_experts.", &.{"down_proj"}, &.{D}, .{i}), .gate = undefined, .up = undefined, .down = undefined, }; - try l.begin(l.expertBytes(i)); + try l.begin(try load_plan.expertBytes(l, i)); const by_rows = c.byRows() and i != c.layers; moe.gate = try l.experts(i, "gate_proj", c.moe_inter, D, if (by_rows) .rows else .whole); moe.up = try l.experts(i, "up_proj", c.moe_inter, D, if (by_rows) .rows else .whole); + if (!std.meta.eql(moe.gate.width(), moe.up.width())) { // one launch takes both, SwiGLU'd + var buf: [96]u8 = undefined; + return mixed(try std.fmt.bufPrint(&buf, "layer {d}'s routed gate_proj", .{i}), moe.gate.width(), "up_proj", moe.up.width()); + } moe.down = try l.experts(i, "down_proj", D, c.moe_inter, if (by_rows) .cols else .whole); out.mlp = .{ .moe = moe }; } else { const n: usize = c.dense_inter; // this Mac's rows; TP2: rank r's half [r n, (r + 1) n) const lo = c.tp_rank * n; out.mlp = .{ .dense = .{ - .gate_up = try l.q4Rows(D, m, &.{ "gate_proj", "up_proj" }, &.{ lo, lo }, &.{ n, n }, &.{ n * c.tp, n * c.tp }, 1, .{i}), - .down = if (c.tp == 1) try l.q4(n, m, &.{"down_proj"}, &.{D}, .{i}) else try l.q4Cols(m ++ "down_proj", .{i}, D, n * c.tp, lo, n), + .gate_up = try l.stack(.dense_mlp, D, m, &gate_up, &.{ lo, lo }, &.{ n, n }, &.{ n * c.tp, n * c.tp }, 1, .{i}), + .down = if (c.tp == 1) try l.mat(.dense_mlp, n, m, &.{"down_proj"}, &.{D}, .{i}) else try l.cols(.dense_mlp, m ++ "down_proj", .{i}, D, n * c.tp, lo, n), } }; } return out; } - - /// The index and every shard's header: each tensor's shard and byte range. - fn index(l: *Loader) !void { - const path = try std.fmt.allocPrintSentinel(l.arena, "{s}/model.safetensors.index.json", .{l.dir}, 0); - const f = try mtl.MappedFile.open(path); - defer f.deinit(); - const doc = try std.json.parseFromSliceLeaky(std.json.Value, l.arena, f.bytes[0..f.size], .{ .allocate = .alloc_always }); - var files: std.StringArrayHashMapUnmanaged(void) = .empty; - var it = doc.object.get("weight_map").?.object.iterator(); - while (it.next()) |kv| try files.put(l.arena, kv.value_ptr.string, {}); - for (files.keys(), 0..) |name, si| { - const full = try std.fmt.allocPrintSentinel(l.arena, "{s}/{s}", .{ l.dir, name }, 0); - try l.shards.append(l.gpa, full); - const fd = std.c.open(full, .{ .ACCMODE = .RDONLY }); - if (fd < 0) return error.OpenFailed; - defer _ = std.c.close(fd); - var head: [8]u8 = undefined; - try readAll(fd, &head, 0); - const hlen = std.mem.readInt(u64, &head, .little); - const header = try l.arena.alloc(u8, hlen); - try readAll(fd, header, 8); - const end = std.c.lseek(fd, 0, std.c.SEEK.END); - const data = 8 + hlen; - const doc2 = try std.json.parseFromSliceLeaky(std.json.Value, l.arena, header, .{}); - var eit = doc2.object.iterator(); - while (eit.next()) |kv| { - const tensor = kv.key_ptr.*; - // the language model's tensors only: the vision tower's (5-D patch kernels among them) are never read - if (!std.mem.startsWith(u8, tensor, prefix) and !std.mem.startsWith(u8, tensor, "lm_head.")) continue; - const o = kv.value_ptr.object; - const dtype = st.DType.parse(o.get("dtype").?.string) orelse return error.UnsupportedDType; - const dims = o.get("shape").?.array.items; - if (dims.len > st.max_rank) return error.RankTooHigh; - var shape: [st.max_rank]usize = @splat(1); - for (dims, 0..) |d, j| shape[j] = @intCast(d.integer); - const offs = o.get("data_offsets").?.array.items; - const b: u64 = @intCast(offs[0].integer); - const e2: u64 = @intCast(offs[1].integer); - if (e2 < b or data + e2 > @as(u64, @intCast(end))) return error.BadSafetensors; - try l.names.put(l.arena, tensor, .{ .shard = @intCast(si), .off = data + b, .len = e2 - b, .dtype = dtype, .shape = shape, .rank = @intCast(dims.len) }); - } - } - } -}; - -const Pool = struct { - fds: []std.c.fd_t, - jobs: []const Copy, - next: std.atomic.Value(usize) = .init(0), - failed: std.atomic.Value(bool) = .init(false), - - fn run(p: *Pool) void { - const buf = paced_read.stage(std.heap.page_allocator, 16 << 20) catch { - p.failed.store(true, .release); - return; - }; - defer std.heap.page_allocator.free(buf); - while (true) { - const i = p.next.fetchAdd(1, .monotonic); - if (i >= p.jobs.len) return; - const j = p.jobs[i]; - paced_read.read(p.fds[j.shard], j.dst[0..j.len], j.off, buf) catch p.failed.store(true, .release); - } - } }; -/// Every copy on `threads` threads, uncached and paced by free memory (a large load otherwise fills the page cache too). -fn runCopies(l: *Loader, threads: usize) !void { - const fds = try l.gpa.alloc(std.c.fd_t, l.shards.items.len); - defer l.gpa.free(fds); - for (l.shards.items, fds) |path, *fd| { - fd.* = std.c.open(path, .{ .ACCMODE = .RDONLY }); - if (fd.* < 0) return error.OpenFailed; - _ = std.c.fcntl(fd.*, 48, @as(c_int, 1)); // F_NOCACHE - } - defer for (fds) |fd| { - _ = std.c.close(fd); - }; - var pool: Pool = .{ .fds = fds, .jobs = l.copies.items }; - const workers = try l.gpa.alloc(?std.Thread, threads); - defer l.gpa.free(workers); - for (workers) |*t| t.* = std.Thread.spawn(.{}, Pool.run, .{&pool}) catch null; - pool.run(); - for (workers) |t| if (t) |th| th.join(); - if (pool.failed.load(.acquire)) return error.ShortRead; - // the transforms on the same threads (each reads its own tensors whole) - const Tx = struct { - l: *Loader, - fds: []std.c.fd_t, - next: std.atomic.Value(usize) = .init(0), - failed: std.atomic.Value(bool) = .init(false), - fn run(x: *@This()) void { - while (true) { - const i = x.next.fetchAdd(1, .monotonic); - if (i >= x.l.transforms.items.len) return; - load_fix.transform(x.l.gpa, x.fds, x.l.transforms.items[i]) catch x.failed.store(true, .release); - } - } - }; - var tx: Tx = .{ .l = l, .fds = fds }; - for (workers) |*t| t.* = std.Thread.spawn(.{}, Tx.run, .{&tx}) catch null; - tx.run(); - for (workers) |t| if (t) |th| th.join(); - if (tx.failed.load(.acquire)) return error.TransformFailed; -} - -/// The first `c.run` layers, the MTP layer (when stored) and the head; `dry`: check names, dtypes and shapes, read nothing. +/// The first `c.run` layers, the MTP layer (if stored), the head; `dry`: check names, dtypes and shapes, read nothing. pub fn load(gpa: std.mem.Allocator, device: mtl.Device, dir: []const u8, c: *const cfg.Config, threads: usize, dry: bool) !*Weights { var arena_state = std.heap.ArenaAllocator.init(gpa); defer arena_state.deinit(); @@ -517,21 +519,67 @@ pub fn load(gpa: std.mem.Allocator, device: mtl.Device, dir: []const u8, c: *con defer l.shards.deinit(gpa); defer l.copies.deinit(gpa); defer l.transforms.deinit(gpa); - try l.index(); + try layout.read(gpa, l.arena, dir, c.layers, &l.shards, &l.names); + l.widths = try readWidths(l.arena, dir, c); const D: usize = c.hidden; const V: usize = c.vocab; - try l.begin(2 * (Q4.wBytes(V, D) + 2 * Q4.sBytes(V, D)) + D * 2 + 8 * 256); - w.embed = try l.q4(D, prefix ++ "embed_tokens", &.{""}, &.{V}, .{}); + try l.begin(affine.bytes(V, D, try l.widths.of(prefix ++ "embed_tokens")) + affine.bytes(V, D, try l.widths.of("lm_head")) + D * 2 + 8 * 256); + w.embed = try l.mat(.embed, D, prefix ++ "embed_tokens", &.{""}, &.{V}, .{}); const vp = c.vocabPart(); // TP2: this Mac's half of the head's rows - w.head = try l.q4Rows(D, "lm_head", &.{""}, &.{vp[0]}, &.{vp[1] - vp[0]}, &.{V}, 1, .{}); + w.head = try l.stack(.head, D, "lm_head", &.{""}, &.{vp[0]}, &.{vp[1] - vp[0]}, &.{V}, 1, .{}); w.norm = try l.plain(prefix ++ "norm.weight", .{}, .bf16, &.{D}); for (0..c.run) |i| w.layers[i] = try l.layer(i); const has_mtp = c.mtp > 0 and l.names.contains(prefix ++ "layers.45.eh_proj.weight"); if (has_mtp) w.layers[c.layers] = try l.layer(c.layers); + if ((c.tp > 1 or c.byRows() or c.own[1] - c.own[0] != c.experts) and !native(w, c)) { + quant.say("glm: this checkpoint has widths other than 4/64; the pair's split kernels read only 4/64 (load it on one Mac)", .{}); + return error.UnsupportedWidth; + } + if (!dry and !native(w, c)) { + var buf: [256]u8 = undefined; + quant.say("glm: this checkpoint's matrices are {s}; the device kernels read 4/64 only", .{tally(w, c).summary(&buf)}); + return error.UnsupportedWidth; + } if (dry) { - std.debug.print("glm plan: {d} of {d} layers{s}, {d:.2} GB in buffers, every tensor's name, dtype and shape checked\n", .{ c.run, c.layers, if (has_mtp) " and the MTP layer" else "", @as(f64, @floatFromInt(w.bytes)) / 1e9 }); + var buf: [256]u8 = undefined; + const widths = tally(w, c).summary(&buf); + std.debug.print("glm plan: {d} of {d} layers{s}, {d:.2} GB in buffers, every tensor's name, dtype and shape checked; {s}\n", .{ c.run, c.layers, if (has_mtp) " and the MTP layer" else "", @as(f64, @floatFromInt(w.bytes)) / 1e9, widths }); return w; } - try runCopies(&l, threads); + try load_fix.runCopies(gpa, l.shards.items, l.copies.items, l.transforms.items, threads); return w; } + +/// A plan's matrices by width: the head, the embedding, every loaded layer. +pub fn tally(w: *const Weights, c: *const cfg.Config) affine.Tally { + var t: affine.Tally = .{}; + t.add(w.embed); + t.add(w.head); + t.add(w.mtp); + for (w.layers[0..c.run]) |l| t.add(l); + if (w.mtp != null) t.add(w.layers[c.layers]); + return t; +} + +/// Whether every loaded matrix takes the generated 4/64 kernels (no other width's pipelines needed). +pub fn native(w: *const Weights, c: *const cfg.Config) bool { + if (!affine.allNative(.{ w.embed, w.head, w.mtp })) return false; + if (w.mtp) |m| if (m.eh_bf16 != null) return false; + for (w.layers[0..c.run]) |l| if (!layerNative(l)) return false; + return w.mtp == null or layerNative(w.layers[c.layers]); +} + +fn layerNative(l: Layer) bool { + return switch (l.attn) { + .mla => |m| m.absorbed == null and affine.allNative(l), + .kda => affine.allNative(l), + }; +} + +/// config.json's widths by module. +fn readWidths(arena: std.mem.Allocator, dir: []const u8, c: *const cfg.Config) !quant.Widths { + const path = try std.fmt.allocPrintSentinel(arena, "{s}/config.json", .{dir}, 0); + const f = try mtl.MappedFile.open(path); + defer f.deinit(); + return quant.parse(arena, f.bytes[0..f.size], c.layers); +} diff --git a/zig/src/families/glm/weights_fix.zig b/zig/src/families/glm/weights_fix.zig index b63e0a7c55..ec0861fe6e 100644 --- a/zig/src/families/glm/weights_fix.zig +++ b/zig/src/families/glm/weights_fix.zig @@ -1,6 +1,7 @@ -//! GLM-5.3-Flash weights rewritten as they load: packed mixes and router, conv taps, fp32 norms, column slices. +//! GLM-5.3-Flash weights as they load: parallel reads into place, and rewrites (mixes, router, conv, norms, columns). const std = @import("std"); const st = @import("../../core/safetensors.zig"); +const paced_read = @import("../../core/paced_read.zig"); /// One tensor's bytes on disk. pub const Src = struct { shard: u32, off: u64, len: u64, dtype: st.DType, shape: [st.max_rank]usize, rank: u8 }; @@ -8,6 +9,9 @@ pub const Fix = enum { hc_pack, router_pack, conv, to_f32, transpose_bf16, cols /// A tensor read whole into host memory and rewritten into its place (small: mixes, router, conv taps, norms). pub const Transform = struct { src: Src, dst: [*]u8, fix: Fix, extra: [3]?Src = .{ null, null, null }, arg: [3]usize = .{ 0, 0, 0 } }; +/// One tensor's bytes read straight into their place in a weights buffer. +pub const Copy = struct { shard: u32, off: u64, len: u64, dst: [*]u8 }; + pub fn readAll(fd: std.c.fd_t, dest: []u8, at: u64) !void { var done: usize = 0; while (done < dest.len) { @@ -21,7 +25,12 @@ pub fn transform(gpa: std.mem.Allocator, fds: []std.c.fd_t, t: Transform) !void const raw = try gpa.alloc(u8, t.src.len); defer gpa.free(raw); try readAll(fds[t.src.shard], raw, t.src.off); - const in16: []const u16 = @alignCast(std.mem.bytesAsSlice(u16, raw)); + const narrow = t.src.dtype == .f32 and (t.fix == .hc_pack or t.fix == .router_pack); // mlx-lm's fp32 mixes, router + if (narrow) bf16Exact(raw) catch |err| { + @import("quant.zig").say("glm: an fp32 {t} tensor (shard {d}, at byte {d}, shape {any}) is not exact in bf16", .{ t.fix, t.src.shard, t.src.off, t.src.shape[0..t.src.rank] }); + return err; + }; + const in16: []const u16 = @alignCast(std.mem.bytesAsSlice(u16, raw[0..if (narrow) raw.len / 2 else raw.len])); switch (t.fix) { .hc_pack => { // [og 6][tm 4][i 16][sgn 8][lane 32][tn 4] -> [og][sgn][lane][i][tm][tn] const out: [*]u16 = @ptrCast(@alignCast(t.dst)); @@ -79,3 +88,85 @@ pub fn transform(gpa: std.mem.Allocator, fds: []std.c.fd_t, t: Transform) !void }, } } + +/// fp32 values as bf16 in place (the first half of `raw`), refused unless each is a bf16 (0.6.6's bf16_if_exact). +fn bf16Exact(raw: []u8) error{InexactBf16}!void { + const n = raw.len / 4; + for (0..n) |i| { + const v = std.mem.readInt(u32, raw[4 * i ..][0..4], .little); + if (v & 0xffff != 0) return error.InexactBf16; + std.mem.writeInt(u16, raw[2 * i ..][0..2], @truncate(v >> 16), .little); + } +} + +test "fp32 values that are bf16 narrow in place; any other value is refused" { + var raw: [12]u8 = undefined; + for ([_]f32{ 1.5, -0.0078125, 3.0 }, 0..) |v, i| std.mem.writeInt(u32, raw[4 * i ..][0..4], @bitCast(v), .little); + try bf16Exact(&raw); + try std.testing.expectEqual(@as(u16, 0x3fc0), std.mem.readInt(u16, raw[0..2], .little)); + try std.testing.expectEqual(@as(u16, 0x4040), std.mem.readInt(u16, raw[4..6], .little)); + std.mem.writeInt(u32, raw[0..4], @bitCast(@as(f32, 1.0000001)), .little); + try std.testing.expectError(error.InexactBf16, bf16Exact(&raw)); +} + +const Pool = struct { + fds: []std.c.fd_t, + jobs: []const Copy, + next: std.atomic.Value(usize) = .init(0), + failed: std.atomic.Value(bool) = .init(false), + + fn run(p: *Pool) void { + const buf = paced_read.stage(std.heap.page_allocator, 16 << 20) catch { + p.failed.store(true, .release); + return; + }; + defer std.heap.page_allocator.free(buf); + while (true) { + const i = p.next.fetchAdd(1, .monotonic); + if (i >= p.jobs.len) return; + const j = p.jobs[i]; + paced_read.read(p.fds[j.shard], j.dst[0..j.len], j.off, buf) catch p.failed.store(true, .release); + } + } +}; + +/// Every copy on `threads` threads, uncached, paced by free memory (a large load otherwise fills the page cache too). +pub fn runCopies(gpa: std.mem.Allocator, shards: []const [:0]const u8, copies: []const Copy, transforms: []const Transform, threads: usize) !void { + const fds = try gpa.alloc(std.c.fd_t, shards.len); + defer gpa.free(fds); + for (shards, fds) |path, *fd| { + fd.* = std.c.open(path, .{ .ACCMODE = .RDONLY }); + if (fd.* < 0) return error.OpenFailed; + _ = std.c.fcntl(fd.*, 48, @as(c_int, 1)); // F_NOCACHE + } + defer for (fds) |fd| { + _ = std.c.close(fd); + }; + var pool: Pool = .{ .fds = fds, .jobs = copies }; + const workers = try gpa.alloc(?std.Thread, threads); + defer gpa.free(workers); + for (workers) |*t| t.* = std.Thread.spawn(.{}, Pool.run, .{&pool}) catch null; + pool.run(); + for (workers) |t| if (t) |th| th.join(); + if (pool.failed.load(.acquire)) return error.ShortRead; + // the transforms on the same threads (each reads its own tensors whole) + const Tx = struct { + gpa: std.mem.Allocator, + jobs: []const Transform, + fds: []std.c.fd_t, + next: std.atomic.Value(usize) = .init(0), + failed: std.atomic.Value(bool) = .init(false), + fn run(x: *@This()) void { + while (true) { + const i = x.next.fetchAdd(1, .monotonic); + if (i >= x.jobs.len) return; + transform(x.gpa, x.fds, x.jobs[i]) catch x.failed.store(true, .release); + } + } + }; + var tx: Tx = .{ .gpa = gpa, .jobs = transforms, .fds = fds }; + for (workers) |*t| t.* = std.Thread.spawn(.{}, Tx.run, .{&tx}) catch null; + tx.run(); + for (workers) |t| if (t) |th| th.join(); + if (tx.failed.load(.acquire)) return error.TransformFailed; +} diff --git a/zig/src/families/glm/weights_test.zig b/zig/src/families/glm/weights_test.zig new file mode 100644 index 0000000000..87e0c7fd01 --- /dev/null +++ b/zig/src/families/glm/weights_test.zig @@ -0,0 +1,468 @@ +//! GLM-5.3-Flash's loader on synthetic checkpoints: one sparse shard whose header names every tensor at its width. +const std = @import("std"); +const cfg = @import("config.zig"); +const wts = @import("weights.zig"); +const quant = @import("quant.zig"); + +const testing = std.testing; +const Width = quant.Width; + +const Entry = struct { name: []const u8, dtype: []const u8, shape: []const usize }; + +/// A module stored at a width other than 4/64. +const Pack = struct { []const u8, Width }; + +/// The tensors the loader reads for `c` (its first `c.run` layers and the MTP layer), stored at `stored`'s widths. +const Fixture = struct { + a: std.mem.Allocator, + c: *const cfg.Config, + stored: []const Pack = &.{}, // module -> its width on disk (else 4/64) + list: std.ArrayList(Entry) = .empty, + mlx: bool = false, // mlx-lm's layout: its names, stacked experts, the absorbed pair, bf16 everywhere unquantized + fused: bool = false, // with mlx: one conv1d tensor for q, k and v, and fp32 kept (mixes, router, biases) + + fn plain(f: *Fixture, comptime fmt: []const u8, args: anytype, dtype: []const u8, shape: []const usize) !void { + const t = if (f.mlx and !f.fused and std.mem.eql(u8, dtype, "F32")) "BF16" else dtype; + try f.list.append(f.a, .{ .name = try std.fmt.allocPrint(f.a, fmt, args), .dtype = t, .shape = try f.a.dupe(usize, shape) }); + } + + /// A stack of `lead` matrices [n, k] (mlx-lm's experts and per-head maps), at its module's width. + fn q3(f: *Fixture, comptime fmt: []const u8, args: anytype, lead: usize, n: usize, k: usize) !void { + const module = try std.fmt.allocPrint(f.a, fmt, args); + var w = quant.uniform; + for (f.stored) |p| if (std.mem.eql(u8, p[0], module)) { + w = p[1]; + }; + try f.plain("{s}.weight", .{module}, "U32", &.{ lead, n, k * w.bits / 32 }); + try f.plain("{s}.scales", .{module}, "BF16", &.{ lead, n, k / w.group }); + try f.plain("{s}.biases", .{module}, "BF16", &.{ lead, n, k / w.group }); + } + + fn q(f: *Fixture, comptime fmt: []const u8, args: anytype, n: usize, k: usize) !void { + const module = try std.fmt.allocPrint(f.a, fmt, args); + var w = quant.uniform; + for (f.stored) |p| if (std.mem.eql(u8, p[0], module)) { + w = p[1]; + }; + try f.plain("{s}.weight", .{module}, "U32", &.{ n, k * w.bits / 32 }); + try f.plain("{s}.scales", .{module}, "BF16", &.{ n, k / w.group }); + try f.plain("{s}.biases", .{module}, "BF16", &.{ n, k / w.group }); + } + + fn layer(f: *Fixture, i: usize) !void { + const c = f.c; + const D: usize = c.hidden; + const p = "model.language_model.layers.{d}."; + const a = p ++ "self_attn."; + if (i != c.layers) for ([_][]const u8{ "attn", "ffn" }) |site| { + try f.plain(p ++ "hc_{s}_fn", .{ i, site }, if (f.fused) "F32" else "BF16", &.{ 24, 4 * D }); + try f.plain(p ++ "hc_{s}_base", .{ i, site }, "F32", &.{24}); + try f.plain(p ++ "hc_{s}_scale", .{ i, site }, "F32", &.{3}); + }; + try f.plain(p ++ "input_layernorm.weight", .{i}, "BF16", &.{D}); + try f.plain(p ++ "post_attention_layernorm.weight", .{i}, "BF16", &.{D}); + switch (c.kind(@intCast(i))) { + .kda => { + const W: usize = c.kdaWidth(); + const dk: usize = c.kda_dim; + if (f.fused) try f.plain(a ++ "conv1d.weight", .{i}, "BF16", &.{ 3 * W, c.conv, 1 }) else for ([_][]const u8{ "q", "k", "v" }) |x| + try f.plain(a ++ "{s}_conv1d.weight", .{ i, x }, "BF16", &.{ W, 1, c.conv }); + try f.plain(a ++ "o_norm.weight", .{i}, "BF16", &.{dk}); + for ([_][]const u8{ "q_proj", "k_proj", "v_proj" }) |x| try f.q(a ++ "{s}", .{ i, x }, W, D); + try f.q(a ++ "f_a_proj", .{i}, dk, D); + try f.q(a ++ "g_a_proj", .{i}, dk, D); + try f.q(a ++ "b_proj", .{i}, c.kda_heads, D); + try f.q(a ++ "f_b_proj", .{i}, W, dk); + try f.q(a ++ "g_b_proj", .{i}, W, dk); + try f.q(a ++ "o_proj", .{i}, D, W); + try f.plain(a ++ "A_log", .{i}, "F32", &.{c.kda_heads}); + try f.plain(a ++ "dt_bias", .{i}, "F32", &.{W}); + }, + .mla => { + const H: usize = c.mla_heads; + try f.plain(a ++ "indexer.index_kpool_compress_gate", .{i}, "BF16", &.{ c.i_dim, D }); + try f.q(a ++ "q_a_proj", .{i}, c.q_lora, D); + try f.q(a ++ "kv_a_proj_with_mqa", .{i}, c.kv_lora, D); + try f.q(a ++ "indexer.wk", .{i}, c.i_dim, D); + try f.q(a ++ "indexer.weights_proj", .{i}, c.i_heads, D); + try f.q(a ++ "q_b_proj", .{i}, H * c.nope, c.q_lora); + try f.q(a ++ "indexer.wq_b", .{i}, c.i_heads * c.i_dim, c.q_lora); + if (f.mlx) { + try f.q3(a ++ "embed_q", .{i}, H, c.kv_lora, c.nope); + try f.q3(a ++ "unembed_out", .{i}, H, c.v_dim, c.kv_lora); + } else try f.q(a ++ "kv_b_proj", .{i}, H * (c.nope + c.v_dim), c.kv_lora); + try f.q(a ++ "o_proj", .{i}, D, H * c.v_dim); + try f.plain(a ++ "q_a_layernorm.weight", .{i}, "BF16", &.{c.q_lora}); + try f.plain(a ++ "kv_a_layernorm.weight", .{i}, "BF16", &.{c.kv_lora}); + try f.plain(a ++ "indexer.k_norm.weight", .{i}, "BF16", &.{c.i_dim}); + try f.plain(a ++ "indexer.k_norm.bias", .{i}, "BF16", &.{c.i_dim}); + try f.plain(a ++ "indexer.index_kpool_compress_ape", .{i}, "BF16", &.{ c.kpool, c.i_dim }); + }, + } + if (i == c.layers) { + if (f.mlx) try f.plain(p ++ "eh_proj.weight", .{i}, "BF16", &.{ D, 2 * D }) else try f.q(p ++ "eh_proj", .{i}, D, 2 * D); + for ([_][]const u8{ "enorm", "hnorm", "shared_head.norm" }) |x| try f.plain(p ++ "{s}.weight", .{ i, x }, "BF16", &.{D}); + } + const m = p ++ "mlp."; + if (c.isMoe(@intCast(i))) { + try f.plain(m ++ "gate.weight", .{i}, if (f.fused) "F32" else "BF16", &.{ c.experts, D }); + try f.plain(m ++ "gate.e_score_correction_bias", .{i}, "F32", &.{c.experts}); + try f.q(m ++ "shared_experts.gate_proj", .{i}, c.moe_inter, D); + try f.q(m ++ "shared_experts.up_proj", .{i}, c.moe_inter, D); + try f.q(m ++ "shared_experts.down_proj", .{i}, D, c.moe_inter); + if (f.mlx) for ([_][]const u8{ "gate_proj", "up_proj", "down_proj" }) |x| { + const down = x[0] == 'd'; + try f.q3(m ++ "switch_mlp.{s}", .{ i, x }, c.experts, if (down) D else c.moe_inter, if (down) c.moe_inter else D); + }; + if (!f.mlx) for (0..c.experts) |e| { + try f.q(m ++ "experts.{d}.gate_proj", .{ i, e }, c.moe_inter, D); + try f.q(m ++ "experts.{d}.up_proj", .{ i, e }, c.moe_inter, D); + try f.q(m ++ "experts.{d}.down_proj", .{ i, e }, D, c.moe_inter); + }; + } else { + try f.q(m ++ "gate_proj", .{i}, c.dense_inter, D); + try f.q(m ++ "up_proj", .{i}, c.dense_inter, D); + try f.q(m ++ "down_proj", .{i}, D, c.dense_inter); + } + } + + /// config.json, the index and one shard (header only, its data a hole) in `dir`. + fn write(f: *Fixture, tmp: std.testing.TmpDir, dir: []const u8, config: []const u8) !void { + const c = f.c; + try f.q("model.language_model.embed_tokens", .{}, c.vocab, c.hidden); + try f.q("lm_head", .{}, c.vocab, c.hidden); + try f.plain("model.language_model.norm.weight", .{}, "BF16", &.{c.hidden}); + for (0..c.run) |i| try f.layer(i); + try f.layer(c.layers); + var head: std.ArrayList(u8) = .empty; + var index: std.ArrayList(u8) = .empty; + try index.appendSlice(f.a, "{\"weight_map\": {"); + var at: usize = 0; + for (f.list.items, 0..) |*e, j| { + if (f.mlx) e.name = try mlxName(f.a, e.name); + var n: usize = if (std.mem.eql(u8, e.dtype, "BF16")) 2 else 4; + for (e.shape) |d| n *= d; + const sep = if (j == 0) "{" else ","; + try head.appendSlice(f.a, try std.fmt.allocPrint(f.a, "{s}\"{s}\":{{\"dtype\":\"{s}\",\"shape\":[{s}],\"data_offsets\":[{d},{d}]}}", .{ sep, e.name, e.dtype, try dims(f.a, e.shape), at, at + n })); + try index.appendSlice(f.a, try std.fmt.allocPrint(f.a, "{s}\"{s}\": \"model-00001.safetensors\"", .{ if (j == 0) "" else ",", e.name })); + at += n; + } + try head.appendSlice(f.a, "}"); + try index.appendSlice(f.a, "}}"); + const io = testing.io; + try tmp.dir.writeFile(io, .{ .sub_path = "config.json", .data = config }); + try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = index.items }); + var len: [8]u8 = undefined; + std.mem.writeInt(u64, &len, head.items.len, .little); + try tmp.dir.writeFile(io, .{ .sub_path = "model-00001.safetensors", .data = try std.mem.concat(f.a, u8, &.{ &len, head.items }) }); + const path = try std.fmt.allocPrintSentinel(f.a, "{s}/model-00001.safetensors", .{dir}, 0); + const fd = std.c.open(path, .{ .ACCMODE = .WRONLY }, @as(std.c.mode_t, 0)); + if (fd < 0) return error.TestFile; + defer _ = std.c.close(fd); + if (std.c.ftruncate(fd, @intCast(8 + head.items.len + at)) != 0) return error.TestFile; + } +}; + +/// An original-layout name as mlx-lm stores it. +fn mlxName(a: std.mem.Allocator, name: []const u8) ![]const u8 { + if (std.mem.startsWith(u8, name, "lm_head.")) return std.mem.concat(a, u8, &.{ "language_model.", name }); + const short = name["model.language_model.".len..]; + const mtp = "layers.45."; + if (std.mem.startsWith(u8, short, mtp)) { + const rest = short[mtp.len..]; + if (std.mem.eql(u8, rest, "shared_head.norm.weight")) return "language_model.mtp.0.norm.weight"; + const top = std.mem.startsWith(u8, rest, "eh_proj.") or std.mem.startsWith(u8, rest, "enorm.") or std.mem.startsWith(u8, rest, "hnorm."); + return std.mem.concat(a, u8, &.{ "language_model.mtp.0.", if (top) "" else "block.", rest }); + } + inline for (.{ "attn", "ffn" }) |site| inline for (.{ "fn", "base", "scale" }) |part| { + const old = ".hc_" ++ site ++ "_" ++ part; + if (std.mem.endsWith(u8, short, old)) return std.mem.concat(a, u8, &.{ "language_model.model.", short[0 .. short.len - old.len], "." ++ site ++ "_hc." ++ part }); + }; + return std.mem.concat(a, u8, &.{ "language_model.model.", short }); +} + +fn dims(a: std.mem.Allocator, shape: []const usize) ![]const u8 { + var out: std.ArrayList(u8) = .empty; + for (shape, 0..) |d, j| try out.appendSlice(a, try std.fmt.allocPrint(a, "{s}{d}", .{ if (j == 0) "" else ",", d })); + return out.items; +} + +/// A small GLM: four layers (three KDA and dense, one MLA and MoE) and the MTP layer, four routed experts. +fn small() cfg.Config { + var c: cfg.Config = .{}; + c.run = 4; + c.experts = 4; + c.own = .{ 0, 4 }; + return c; +} + +/// Every matrix's rows, columns and buffer offsets, and every other tensor's offset, hashed in load order. +fn digest(h: *std.hash.Wyhash, v: anytype) void { + const T = @TypeOf(v); + switch (@typeInfo(T)) { + .@"struct" => |s| { + if (@hasField(T, "buf") and @hasField(T, "off")) return h.update(std.mem.asBytes(&v.off)); + if (@hasField(T, "w") and @hasField(T, "n")) { + h.update(std.mem.asBytes(&v.n)); + h.update(std.mem.asBytes(&v.k)); + inline for (.{ v.w, v.s, v.b }) |r| digest(h, r); + return; + } + inline for (s.field_names) |name| digest(h, @field(v, name)); + }, + .@"union" => switch (v) { + inline else => |x| digest(h, x), + }, + .optional => if (v) |x| digest(h, x), + .array => for (v) |x| digest(h, x), + else => {}, + } +} + +fn plan(w: *const wts.Weights, c: *const cfg.Config) u64 { + var h = std.hash.Wyhash.init(0); + inline for (.{ w.embed, w.head, w.norm, w.mtp }) |x| digest(&h, x); + for (w.layers[0..c.run]) |l| digest(&h, l); + digest(&h, w.layers[c.layers]); + h.update(std.mem.asBytes(&w.bytes)); + return h.final(); +} + +const flat = "{\"quantization\": {\"group_size\": 64, \"bits\": 4, \"mode\": \"affine\"}}"; + +/// A dry load of a fixture stored at `stored` under `config`: the loader's plan, or its refusal. +fn dry(a: std.mem.Allocator, config: []const u8, stored: []const Pack) !*wts.Weights { + return dryAs(a, config, stored, false); +} + +fn dryAs(a: std.mem.Allocator, config: []const u8, stored: []const Pack, mlx: bool) !*wts.Weights { + return dryWith(a, config, stored, mlx, false); +} + +fn dryWith(a: std.mem.Allocator, config: []const u8, stored: []const Pack, mlx: bool, fused: bool) !*wts.Weights { + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + const dir = try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}", .{tmp.sub_path}); + const c = small(); + var f: Fixture = .{ .a = a, .c = &c, .stored = stored, .mlx = mlx, .fused = fused }; + try f.write(tmp, dir, config); + return wts.load(testing.allocator, undefined, dir, &c, 1, true); +} + +fn free(w: *wts.Weights) void { + w.deinit(); + testing.allocator.destroy(w); +} + +test "a uniform 4/64 checkpoint plans the same matrices, offsets and bytes as the 4-bit-only loader" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const c = small(); + const w = try dry(arena.allocator(), flat, &.{}); + defer free(w); + try testing.expectEqual(@as(usize, 1_512_774_688), w.bytes); + try testing.expectEqual(@as(u64, 0xc1e86c075603176a), plan(w, &c)); + try testing.expectEqual(@as(u8, 0), w.layers[0].attn.kda.in_proj.more); // one width: one run, the matrix itself +} + +/// An oQ-style table: o_proj of layer 0 at 8 bits under the short prefix, layer 3's shared up at 5 under the long. +const oq = + \\{"quantization": {"group_size": 64, "bits": 4, "mode": "affine", + \\ "model.layers.0.self_attn.o_proj": {"group_size": 64, "bits": 8}, + \\ "language_model.model.layers.3.mlp.shared_experts.down_proj": {"group_size": 32, "bits": 5}, + \\ "language_model.mtp.0.block.mlp.switch_mlp.gate_proj": {"group_size": 64, "bits": 2}}} +; +const oq_packed = [_]Pack{ + .{ "model.language_model.layers.0.self_attn.o_proj", .{ .bits = 8, .group = 64 } }, + .{ "model.language_model.layers.3.mlp.shared_experts.down_proj", .{ .bits = 5, .group = 32 } }, +} ++ blk: { + var out: [4]Pack = undefined; + for (&out, 0..) |*o, e| o.* = .{ std.fmt.comptimePrint("model.language_model.layers.45.mlp.experts.{d}.gate_proj", .{e}), .{ .bits = 2, .group = 64 } }; + break :blk out; +}; + +test "the device admits each tensor at its override's width, whatever the key's prefix; a gate and up at two are refused" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + try testing.expectError(error.MixedWidths, dry(a, oq, &oq_packed)); // the MTP layer's gate at 2, its up at 4 + const both = oq[0 .. oq.len - 2] ++ ",\n \"language_model.mtp.0.block.mlp.switch_mlp.up_proj\": {\"group_size\": 64, \"bits\": 2}}}"; + const ups = comptime oq_packed ++ blk: { + var out: [4]Pack = undefined; + for (&out, 0..) |*o, e| o.* = .{ std.fmt.comptimePrint("model.language_model.layers.45.mlp.experts.{d}.up_proj", .{e}), .{ .bits = 2, .group = 64 } }; + break :blk out; + }; + const c = small(); + const w = try dry(a, both, &ups); + defer free(w); + try testing.expectEqual(Width{ .bits = 8, .group = 64 }, w.layers[0].attn.kda.o_proj.width()); + try testing.expectEqual(Width{ .bits = 5, .group = 32 }, w.layers[3].mlp.moe.sh_down.width()); + try testing.expectEqual(Width{ .bits = 2, .group = 64 }, w.layers[c.layers].mlp.moe.gate.width()); + try testing.expectEqual(Width{ .bits = 2, .group = 64 }, w.layers[c.layers].mlp.moe.up.width()); + try testing.expectEqual(quant.uniform, w.layers[1].attn.kda.o_proj.width()); + try testing.expectEqual(quant.uniform, w.layers[3].mlp.moe.gate.width()); + try testing.expect(!wts.native(w, &c)); +} + +test "a tensor stored at another width than the config gives it is refused, never read at a guessed width" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const eight = [_]Pack{.{ "model.language_model.layers.0.self_attn.o_proj", .{ .bits = 8, .group = 64 } }}; + try testing.expectError(error.UnexpectedTensor, dry(a, flat, &eight)); // global 4/64, stored at 8 + try testing.expectError(error.UnexpectedTensor, dry(a, oq, &.{})); // override 8, stored at 4 +} + +/// oQ4's split stacks on this fixture: the KDA input (q, v at 5; k, f_a, g_a, b at 8), the MLA x_proj (6 and 5). +const oq_split = + \\{"quantization": {"group_size": 64, "bits": 4, "mode": "affine", + \\ "model.layers.0.self_attn.q_proj": {"group_size": 64, "bits": 5}, + \\ "model.layers.0.self_attn.k_proj": {"group_size": 64, "bits": 8}, + \\ "model.layers.0.self_attn.v_proj": {"group_size": 64, "bits": 5}, + \\ "model.layers.0.self_attn.f_a_proj": {"group_size": 64, "bits": 8}, + \\ "model.layers.0.self_attn.g_a_proj": {"group_size": 64, "bits": 8}, + \\ "model.layers.0.self_attn.b_proj": {"group_size": 64, "bits": 8}, + \\ "model.layers.3.self_attn.q_a_proj": {"group_size": 64, "bits": 6}, + \\ "model.layers.3.self_attn.kv_a_proj_with_mqa": {"group_size": 64, "bits": 6}, + \\ "model.layers.3.self_attn.indexer.wk": {"group_size": 64, "bits": 5}, + \\ "model.layers.3.self_attn.indexer.weights_proj": {"group_size": 64, "bits": 5}}} +; +const oq_split_stored = blk: { + const bits = [_]struct { []const u8, u8 }{ + .{ "0.self_attn.q_proj", 5 }, + .{ "0.self_attn.k_proj", 8 }, + .{ "0.self_attn.v_proj", 5 }, + .{ "0.self_attn.f_a_proj", 8 }, + .{ "0.self_attn.g_a_proj", 8 }, + .{ "0.self_attn.b_proj", 8 }, + .{ "3.self_attn.q_a_proj", 6 }, + .{ "3.self_attn.kv_a_proj_with_mqa", 6 }, + .{ "3.self_attn.indexer.wk", 5 }, + .{ "3.self_attn.indexer.weights_proj", 5 }, + }; + var out: [bits.len]Pack = undefined; + for (&out, bits) |*o, b| o.* = .{ "model.language_model.layers." ++ b[0], .{ .bits = b[1], .group = 64 } }; + break :blk out; +}; + +/// Run `i` of `q` is rows [row, row + n) at `bits`, its codes and tables its own, `pad`-row tiles apart. +fn expectRun(q: wts.Affine, i: usize, row: u32, n: u32, bits: u8, pad: usize) !void { + const r = q.run(i); + try testing.expectEqual(row, r.row); + try testing.expectEqual(n, r.n); + try testing.expectEqual(Width{ .bits = bits, .group = 64 }, Width{ .bits = r.bits, .group = r.group }); + const w: Width = .{ .bits = r.bits, .group = r.group }; + const room = std.mem.alignForward(usize, n, pad); + try testing.expectEqual(std.mem.alignForward(usize, room * q.k * bits / 8, 256), r.s.off - r.w.off); + try testing.expectEqual(std.mem.alignForward(usize, room * (q.k / w.group) * 2, 256), r.b.off - r.s.off); +} + +test "a stack of mixed widths plans one run a member, as 0.6.6's QSplit calls each part, at its rows of the output" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const c = small(); + const w = try dry(arena.allocator(), oq_split, &oq_split_stored); + defer free(w); + const in = w.layers[0].attn.kda.in_proj; + const W = c.kdaWidth(); + const dk = c.kda_dim; + try testing.expectEqual(@as(u8, 5), in.more); // f_a, g_a and b share 8 bits yet stay apart + try testing.expectEqual(c.kdaProj(), in.n); + try expectRun(in, 0, 0, W, 5, 64); + try expectRun(in, 1, W, W, 8, 64); + try expectRun(in, 2, 2 * W, W, 5, 64); + try expectRun(in, 3, 3 * W, dk, 8, 64); + try expectRun(in, 4, 3 * W + dk, dk, 8, 64); + try expectRun(in, 5, 3 * W + 2 * dk, c.kda_heads, 8, 64); + const x = w.layers[3].attn.mla.x_proj; + try testing.expectEqual(@as(u8, 3), x.more); + try expectRun(x, 0, 0, c.q_lora, 6, 64); + try expectRun(x, 1, c.q_lora, c.kv_lora, 6, 64); + try expectRun(x, 2, c.q_lora + c.kv_lora, c.i_dim, 5, 64); + try expectRun(x, 3, c.q_lora + c.kv_lora + c.i_dim, c.i_heads, 5, 64); + try testing.expectEqual(@as(u8, 0), w.layers[1].attn.kda.in_proj.more); + try testing.expectEqual(@as(u8, 0), w.layers[3].attn.mla.qr_proj.more); + const t = wts.tally(w, &c); + try testing.expectEqual(@as(u32, 4), t.count[5][1]); // q, v and the indexer's pair + try testing.expectEqual(@as(u32, 2), t.count[6][1]); + try testing.expectEqual(@as(u32, 4), t.count[8][1]); + var buf: [256]u8 = undefined; + try testing.expect(std.mem.indexOf(u8, t.summary(&buf), "5-bit/64 x4, 6-bit/64 x2, 8-bit/64 x4") != null); +} + +/// oQ's spelling in mlx-lm's layout: the embedding at 8, a stack at 2, the absorbed pair at 5 and 6. +const oq_mlx = + \\{"quantization": {"group_size": 64, "bits": 4, "mode": "affine", + \\ "language_model.model.embed_tokens": {"group_size": 64, "bits": 8}, + \\ "language_model.model.layers.3.mlp.switch_mlp.down_proj": {"group_size": 64, "bits": 2}, + \\ "language_model.model.layers.3.self_attn.embed_q": {"group_size": 64, "bits": 5}, + \\ "language_model.mtp.0.block.self_attn.unembed_out": {"group_size": 64, "bits": 6}, + \\ "language_model.mtp.0.block.mlp.switch_mlp.gate_proj": {"group_size": 64, "bits": 2}}} +; +const oq_mlx_stored = [_]Pack{ + .{ "model.language_model.embed_tokens", .{ .bits = 8, .group = 64 } }, + .{ "model.language_model.layers.3.mlp.switch_mlp.down_proj", .{ .bits = 2, .group = 64 } }, + .{ "model.language_model.layers.3.self_attn.embed_q", .{ .bits = 5, .group = 64 } }, + .{ "model.language_model.layers.45.self_attn.unembed_out", .{ .bits = 6, .group = 64 } }, + .{ "model.language_model.layers.45.mlp.switch_mlp.gate_proj", .{ .bits = 2, .group = 64 } }, +}; + +test "an mlx-lm checkpoint's every tensor is found, each at its own width" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const c = small(); + const both = oq_mlx[0 .. oq_mlx.len - 2] ++ ",\n \"language_model.mtp.0.block.mlp.switch_mlp.up_proj\": {\"group_size\": 64, \"bits\": 2}}}"; + const up: Pack = .{ "model.language_model.layers.45.mlp.switch_mlp.up_proj", .{ .bits = 2, .group = 64 } }; + const w = try dryAs(arena.allocator(), both, &(oq_mlx_stored ++ [_]Pack{up}), true); + defer free(w); + try testing.expectEqual(Width{ .bits = 8, .group = 64 }, w.embed.width()); + try testing.expectEqual(quant.uniform, w.head.width()); + try testing.expectEqual(quant.uniform, w.layers[0].attn.kda.o_proj.width()); + const moe = w.layers[3].mlp.moe; + try testing.expectEqual(Width{ .bits = 2, .group = 64 }, moe.down.width()); + try testing.expectEqual(quant.uniform, moe.gate.width()); + const mla = w.layers[3].attn.mla; + try testing.expectEqual(@as(u32, 0), mla.kv_b.n); + try testing.expectEqual(Width{ .bits = 5, .group = 64 }, mla.absorbed.?[0].width()); + try testing.expectEqual(@as(u32, c.mla_heads * c.kv_lora), mla.absorbed.?[0].n); + try testing.expectEqual(@as(u32, c.nope), mla.absorbed.?[0].k); + try testing.expectEqual(quant.uniform, mla.absorbed.?[1].width()); + try testing.expectEqual(Width{ .bits = 6, .group = 64 }, w.layers[c.layers].attn.mla.absorbed.?[1].width()); + try testing.expectEqual(Width{ .bits = 2, .group = 64 }, w.layers[c.layers].mlp.moe.gate.width()); + try testing.expectEqual(Width{ .bits = 2, .group = 64 }, w.layers[c.layers].mlp.moe.up.width()); + try testing.expect(w.mtp.?.eh_bf16 != null); +} + +test "the device loads mlx-lm's absorbed pair and bf16 eh_proj; a routed gate and up at two widths are refused" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + try testing.expectError(error.MixedWidths, dryAs(a, oq_mlx, &oq_mlx_stored, true)); // the MTP gate at 2, up at 4 + const c = small(); + const w = try dryAs(a, flat, &.{}, true); + defer free(w); + try testing.expectEqual(quant.uniform, w.layers[3].attn.mla.absorbed.?[0].width()); + try testing.expect(w.mtp.?.eh_bf16 != null); + try testing.expect(!wts.native(w, &c)); // the pair and the bf16 eh_proj take the row kernels +} + +test "a uniform checkpoint in the original layout needs no other width's pipelines" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const c = small(); + const w = try dry(arena.allocator(), flat, &.{}); + defer free(w); + try testing.expect(wts.native(w, &c)); +} + +test "mlx-lm's one conv1d tensor and fp32 mixes and router plan as the three conv tensors and bf16 ones do" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const c = small(); + const w = try dryWith(a, flat, &.{}, true, true); + defer free(w); + const v = try dryWith(a, flat, &.{}, true, false); + defer free(v); + try testing.expectEqual(plan(v, &c), plan(w, &c)); // the same buffers, offsets and bytes +} From 72f5a10b731cc4671ce8b1c160c7183a5d8a19f6 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 04:28:08 -1000 Subject: [PATCH 08/18] GLM-5.3-Flash runs every MLX width on GPUs without tensor units, split stacks a product a member GLM has no tensor-unit-only kernel (its *_nax kernels take the simdgroup-matrix fragments without tensor units), so the pre-M5 path only lacked products at other widths. glm_rows.metal (wide.zig) runs them over core row_projection's fp32 group sums, 2/3/4/5/6/8 bits in groups of 32/64/128, every row its own sums (one row alone == the same row in any window): - dense with input and output row strides and per-head batching: each split-stack member at its own output rows (oQ4's KDA input q, v at 5 bits and k, f_a, g_a, b at 8 is six products), mlx-lm's embed_q and unembed_out a head; - by expert id (decode's routed experts, gate and up fused with SwiGLU, glm_glue's sigmoid_fast and bfloat overloads compiled in, not copied), and by sorted pairs (prompt's routed experts); then the existing combine; - bf16 (the MTP head's unquantized eh_proj); the embedding at 6 and 8 bits. KDA f_b and g_b at 4/64 keep the generated kernel's fused dots; at any other width they are row products first and one instance of the generated kda_rows reads them in place of its dots (0.6.6's unfused shape); the products borrow the dense MLP scratch, so state refuses a config whose KDA width outgrows it. A shared expert not at 4/64 runs its own products (no shared-in-routed launch); routed experts not at 4/64 run by id and by sorted pair. Prompt chunks: a one-width 8/64 matrix runs on core affine_mm, split stacks and other widths on the row kernels. The new pipelines load only when a checkpoint has a matrix the 4/64 kernels cannot read, and their source hash joins the kernels' identity; a failed compile releases every pipeline built before it. The uniform 4/64 checkpoint compiles and launches what it did. The previous commit's refusals stay: a member whose rows are not a multiple of 4 (the row kernels read four rows), routed gate and up at two widths, any width but 4/64 with TP2 or expert parallel. Tests that fail before (no wide.zig on the previous commit): 4/64 of one width alone takes the generated kernels; f_b and g_b at any other width (8/64 too) take the products made apart; a failed compile releases every pipeline built before it and a load counts each one (GPU); glm_rows's SwiGLU compiles glm_glue's helpers, no copy of its own. The previous commit's load-time refusal of every other width is gone. Gate on main 0efbd447 + the core commits (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test logs an error on the path it expects). lean_check: 533 problems (main 539), no new findings. --- zig/kernels/metal/glm_rows.metal | 169 +++++++++++++++++++++ zig/kernels/metal/sources.zig | 2 + zig/src/families/glm/checks.zig | 1 + zig/src/families/glm/engine.zig | 1 + zig/src/families/glm/forward.zig | 39 ++++- zig/src/families/glm/glm.zig | 2 + zig/src/families/glm/kernels.zig | 26 +++- zig/src/families/glm/moe.zig | 42 +++++- zig/src/families/glm/mtp.zig | 4 +- zig/src/families/glm/prompt.zig | 22 ++- zig/src/families/glm/state.zig | 1 + zig/src/families/glm/weights.zig | 5 - zig/src/families/glm/wide.zig | 243 +++++++++++++++++++++++++++++++ 13 files changed, 535 insertions(+), 22 deletions(-) create mode 100644 zig/kernels/metal/glm_rows.metal create mode 100644 zig/src/families/glm/wide.zig diff --git a/zig/kernels/metal/glm_rows.metal b/zig/kernels/metal/glm_rows.metal new file mode 100644 index 0000000000..25073b268b --- /dev/null +++ b/zig/kernels/metal/glm_rows.metal @@ -0,0 +1,169 @@ +// GLM-5.3-Flash's products at any MLX width: core/row_projection.metal's row sums (compiled first), any window. + +struct GlmRowDims { + int k, n, group, rows; // a matrix [n, k] at `group`; x rows + int heads, ldx, hsx; // matrices back to back (one a head); x's row and head strides + int ldo, hso, hsw; // out's row and head strides; a head's matrix rows +}; + +struct GlmPickDims { + int k, n, group, slots; // slots: (row, pick) pairs, or sorted pairs + int repeat, experts, ldo; // a row's picks (x row = slot / repeat); experts in W; out's row stride + float lim; // the SwiGLU clip +}; + +// The MoE kernels' SwiGLU of a gate and up sum: glm_swiglu's, its sigmoid_fast put first (wide.zig glue_math). +inline bfloat glm_rows_glu(float gv, float uv, float limit) { + const float lim = float(bfloat(limit)); + const bfloat gt = bfloat(metal::min(float(bfloat(gv)), lim)); + const bfloat up = bfloat(metal::min(metal::max(float(bfloat(uv)), -lim), lim)); + return (gt * sigmoid_fast(gt)) * up; +} + +// out[m, h, c] = x[m, h] . W[h][c]: 4 columns a simdgroup, 2 simdgroups a threadgroup; grid (n / 8, heads). +template +[[kernel]] void glm_rows(const device bfloat* X [[buffer(0)]], const device uint32_t* W [[buffer(1)]], + const device bfloat* S [[buffer(2)]], const device bfloat* B [[buffer(3)]], + constant GlmRowDims& d [[buffer(4)]], device bfloat* OUT [[buffer(5)]], + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int row0 = (int(tg.x) * 2 + int(sg)) * 4; + if (row0 >= d.n) return; + const int h = int(tg.y); + const size_t at = size_t(h) * d.hsw + row0; + const device uint8_t* wr = (const device uint8_t*)W + at * (d.k * BITS / 8); + const device bfloat* sr = S + at * (d.k / d.group); + const device bfloat* br = B + at * (d.k / d.group); + const device bfloat* x0 = X + size_t(h) * d.hsx; + device bfloat* o0 = OUT + size_t(h) * d.hso + row0; + const int cols = metal::min(4, d.n - row0); + int m = 0; + if (d.ldx == d.k) { + #pragma clang loop unroll(disable) + for (; m + 2 <= d.rows; m += 2) { + float acc[2][4]; + tf_rp_rows<4, 2, BITS, false>(wr, sr, br, x0 + size_t(m) * d.k, d.k, d.group, lane, acc); + if (lane == 0) + for (int b = 0; b < 2; b++) + for (int j = 0; j < cols; j++) o0[size_t(m + b) * d.ldo + j] = bfloat(acc[b][j]); + } + } + #pragma clang loop unroll(disable) + for (; m < d.rows; m++) { + float acc[1][4]; + tf_rp_rows<4, 1, BITS, false>(wr, sr, br, x0 + size_t(m) * d.ldx, d.k, d.group, lane, acc); + if (lane == 0) + for (int j = 0; j < cols; j++) o0[size_t(m) * d.ldo + j] = bfloat(acc[0][j]); + } +} + +// Slot y through expert ids[y] on x row y / repeat (GLU: gate W and up W2, SwiGLU'd); grid (n / 8, slots). +template +[[kernel]] void glm_rows_ids(const device bfloat* X [[buffer(0)]], const device uint32_t* W [[buffer(1)]], + const device bfloat* S [[buffer(2)]], const device bfloat* B [[buffer(3)]], + constant GlmPickDims& d [[buffer(4)]], device bfloat* OUT [[buffer(5)]], + const device int* ids [[buffer(6)]], const device uint32_t* W2 [[buffer(7)]], + const device bfloat* S2 [[buffer(8)]], const device bfloat* B2 [[buffer(9)]], + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int row0 = (int(tg.x) * 2 + int(sg)) * 4; + if (row0 >= d.n) return; + const int slot = int(tg.y); + const int expert = ids[slot]; + device bfloat* out = OUT + size_t(slot) * d.ldo + row0; + const int cols = metal::min(4, d.n - row0); + if (expert < 0 || expert >= d.experts) { + if (int(lane) < cols) out[lane] = bfloat(0.0f); + return; + } + const size_t at = size_t(expert) * d.n + row0; + const size_t kb = size_t(d.k) * BITS / 8, kg = size_t(d.k / d.group); + const device bfloat* x = X + size_t(slot / d.repeat) * d.k; + float acc[1][4]; + tf_rp_rows<4, 1, BITS, false>((const device uint8_t*)W + at * kb, S + at * kg, B + at * kg, x, d.k, d.group, lane, acc); + if (GLU) { + float up[1][4]; + tf_rp_rows<4, 1, BITS, false>((const device uint8_t*)W2 + at * kb, S2 + at * kg, B2 + at * kg, x, d.k, d.group, lane, up); + if (lane == 0) + for (int j = 0; j < cols; j++) out[j] = glm_rows_glu(acc[0][j], up[0][j], d.lim); + } else if (lane == 0) { + for (int j = 0; j < cols; j++) out[j] = bfloat(acc[0][j]); + } +} + +// Expert e's sorted pairs: GLU x row order[i] / repeat to out row i, else x row i to row order[i]; grid (n / 8, e). +template +[[kernel]] void glm_rows_sorted(const device bfloat* X [[buffer(0)]], const device uint32_t* W [[buffer(1)]], + const device bfloat* S [[buffer(2)]], const device bfloat* B [[buffer(3)]], + constant GlmPickDims& d [[buffer(4)]], device bfloat* OUT [[buffer(5)]], + const device int* offsets [[buffer(6)]], const device int* order [[buffer(7)]], + const device uint32_t* W2 [[buffer(8)]], const device bfloat* S2 [[buffer(9)]], + const device bfloat* B2 [[buffer(10)]], + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int row0 = (int(tg.x) * 2 + int(sg)) * 4; + if (row0 >= d.n) return; + const int e = int(tg.y); + const int lo = offsets[e], hi = offsets[e + 1]; + const size_t at = size_t(e) * d.n + row0; + const size_t kb = size_t(d.k) * BITS / 8, kg = size_t(d.k / d.group); + const device uint8_t* w = (const device uint8_t*)W + at * kb; + const int cols = metal::min(4, d.n - row0); + #pragma clang loop unroll(disable) + for (int i = lo; i < hi; i++) { + const int src = GLU ? order[i] / d.repeat : i; + const int dst = GLU ? i : order[i]; + float acc[1][4]; + tf_rp_rows<4, 1, BITS, false>(w, S + at * kg, B + at * kg, X + size_t(src) * d.k, d.k, d.group, lane, acc); + device bfloat* out = OUT + size_t(dst) * d.ldo + row0; + if (GLU) { + float up[1][4]; + tf_rp_rows<4, 1, BITS, false>((const device uint8_t*)W2 + at * kb, S2 + at * kg, B2 + at * kg, X + size_t(src) * d.k, d.k, d.group, lane, up); + if (lane == 0) + for (int j = 0; j < cols; j++) out[j] = glm_rows_glu(acc[0][j], up[0][j], d.lim); + } else if (lane == 0) { + for (int j = 0; j < cols; j++) out[j] = bfloat(acc[0][j]); + } + } +} + +// out[m, c] = x[m] . W[c] for bf16 W [n, k] (an unquantized eh_proj): fp32 sums, lanes over 16-input pieces in order. +[[kernel]] void glm_rows_bf16(const device bfloat* X [[buffer(0)]], const device bfloat* W [[buffer(1)]], + constant GlmRowDims& d [[buffer(4)]], device bfloat* OUT [[buffer(5)]], + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int row0 = (int(tg.x) * 2 + int(sg)) * 4; + if (row0 >= d.n) return; + const int cols = metal::min(4, d.n - row0); + #pragma clang loop unroll(disable) + for (int m = 0; m < d.rows; m++) { + const device bfloat* x = X + size_t(m) * d.ldx; + float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (int k0 = int(lane) * 16; k0 < d.k; k0 += 512) { + float xt[16]; + for (int i = 0; i < 16; i++) xt[i] = float(x[k0 + i]); + for (int j = 0; j < cols; j++) { + const device bfloat* w = W + size_t(row0 + j) * d.k + k0; + float s = 0.0f; + for (int i = 0; i < 16; i++) s += xt[i] * float(w[i]); + acc[j] += s; + } + } + for (int j = 0; j < 4; j++) acc[j] = simd_sum(acc[j]); + if (lane == 0) + for (int j = 0; j < cols; j++) OUT[size_t(m) * d.ldo + row0 + j] = bfloat(acc[j]); + } +} + +#define GLM_ROWS_AT(B) \ + template [[host_name("glm_rows_q" #B)]] [[kernel]] decltype(glm_rows) glm_rows; \ + template [[host_name("glm_rows_ids_q" #B)]] [[kernel]] decltype(glm_rows_ids) glm_rows_ids; \ + template [[host_name("glm_rows_glu_ids_q" #B)]] [[kernel]] decltype(glm_rows_ids) glm_rows_ids; \ + template [[host_name("glm_rows_sorted_q" #B)]] [[kernel]] decltype(glm_rows_sorted) glm_rows_sorted; \ + template [[host_name("glm_rows_glu_sorted_q" #B)]] [[kernel]] decltype(glm_rows_sorted) glm_rows_sorted; +GLM_ROWS_AT(2) +GLM_ROWS_AT(3) +GLM_ROWS_AT(4) +GLM_ROWS_AT(5) +GLM_ROWS_AT(6) +GLM_ROWS_AT(8) diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index 6313b656cd..a1113ada56 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -82,6 +82,8 @@ pub const flashnext_gdn = @embedFile("decode/fn_gdn.metal"); pub const glm = @import("glm/kernels.zig"); /// GLM-5.3-Flash's glue: MLX's arithmetic where the Python family calls MLX ops; selection, argmax, small steps. pub const glm_glue = @embedFile("glm_glue.metal"); +/// Its products at the widths the generated kernels do not read, after core_row_projection (glm/wide.zig). +pub const glm_rows = @embedFile("glm_rows.metal"); /// GLM-5.3-Flash's latent attention for rows that read every key: 64 heads as one matrix on the tensor units. pub const glm_attn = @embedFile("glm_attn.metal"); /// A prompt chunk's KDA layer in three passes, appended to the generated kda_rows source (its helpers). diff --git a/zig/src/families/glm/checks.zig b/zig/src/families/glm/checks.zig index 4b521379c1..98404ca22a 100644 --- a/zig/src/families/glm/checks.zig +++ b/zig/src/families/glm/checks.zig @@ -158,6 +158,7 @@ pub fn checkMatmul(e: *Engine, prompt: []const u32) !void { .kda => |*a| a, .mla => return error.NotKda, }; + if (!@import("wide.zig").native(a.in_proj)) return error.UnsupportedWidth; // the 4/64 product's comparison only const pool = mtl.objc.Pool.push(); defer pool.pop(); e.sync(); diff --git a/zig/src/families/glm/engine.zig b/zig/src/families/glm/engine.zig index 54095bbad2..e1ffe5ac41 100644 --- a/zig/src/families/glm/engine.zig +++ b/zig/src/families/glm/engine.zig @@ -205,6 +205,7 @@ pub const Engine = struct { e.w.deinit(); gpa.destroy(e.w); } + if (!wts.native(e.w, &e.c)) try e.k.widen(gpa, e.device); // a width, split stack or layout past the 4/64 kernels try e.prepare(); if (link) |s| { const rows = e.c.byRows(); diff --git a/zig/src/families/glm/forward.zig b/zig/src/families/glm/forward.zig index de761d70fd..a9104bfd4e 100644 --- a/zig/src/families/glm/forward.zig +++ b/zig/src/families/glm/forward.zig @@ -9,6 +9,7 @@ const Ep = @import("ep.zig").Ep; const moe_route = @import("../../core/moe_route.zig"); const hc_shape = @import("kernels.zig").hc_shape; const hc_core = @import("../../core/hc.zig"); +const wide = @import("wide.zig"); const Ref = wts.Ref; /// The MoE block (moe.zig). pub const moe = @import("moe.zig").moe; @@ -112,11 +113,11 @@ pub fn size(w: usize, h: usize, d: usize) mtl.Size { return mtl.Size.of(w, h, d); } -/// x [rows, K] through a 4-bit matrix: qmv_rows, MLX's one-row qmv_fast sums for every row. +/// x [rows, K] through a 4-bit/64 matrix (qmv_rows, MLX's one-row qmv_fast sums), else each member at its width (wide). pub fn qmv(x: *const Ctx, e: mtl.ComputeEncoder, pipe: mtl.Pipeline, in: Ref, q: wts.Affine, out: Ref, rows: u32) void { + if (!wide.native(q)) return wide.dense(&x.k.wide.?, e, q, .{ .in = in, .ldx = q.k, .out = out, .ldo = q.n, .rows = rows }); e.setPipeline(pipe); bind(e, 0, .{ in, q.w, q.s, q.b, out }); - _ = x; e.dispatchThreads(size(32 * rows, q.n / 4, 1), size(32 * rows, 1, 1)); } @@ -149,11 +150,16 @@ pub fn embed(x: *Ctx, e: mtl.ComputeEncoder, ids: Ref, rows: u32) void { pub fn embedRows(x: *const Ctx, e: mtl.ComputeEncoder, ids: Ref, out: Ref, rows: u32) void { const q = x.w.embed; - e.setPipeline(x.k.embed); + const per: u32 = switch (q.bits) { // values a thread: two codes a byte at 4 bits, four in three bytes at 6, one at 8 + 4 => 2, + 6 => 4, + else => 1, + }; + e.setPipeline(if (q.bits == 4) x.k.embed else x.k.wide.?.embed[if (q.bits == 6) 0 else 1]); bind(e, 0, .{ ids, q.w, q.s, q.b }); shape(e, 4, .{x.c.hidden}); bind(e, 5, .{out}); - e.dispatchThreads(size(x.c.hidden / 2, rows, 1), size(256, 1, 1)); + e.dispatchThreads(size(x.c.hidden / per, rows, 1), size(256, 1, 1)); } /// A block boundary: the pending branch into the streams, then (with `hc`) the next block's mix, split and norm (core/hc.zig). @@ -202,10 +208,25 @@ pub fn kdaStep(x: *const Ctx, e: mtl.ComputeEncoder, ki: usize, w: *const wts.Kd const c = x.c; const L = &x.s.kda[ki]; const cur = L.cur; - e.setPipeline(if (c.tp > 1) x.k.kda_rows_tp else x.k.kda_rows); + const gate = wide.kdaGate(if (x.k.wide) |*p| p else null, w.f_b, w.g_b); + var fb = [3]Ref{ w.f_b.w, w.f_b.s, w.f_b.b }; + var gb = [3]Ref{ w.g_b.w, w.g_b.s, w.g_b.b }; + if (gate == .pre) { // f_b and g_b at their widths on f_a and g_a first, [rows, width] each (in gu: no MLP runs yet) + const W = c.kdaWidth(); + const f = x.sc.gu; + const g = f.at(@as(usize, st.max_rows) * W * 2); + wide.dense(&x.k.wide.?, e, w.f_b, .{ .in = proj.at(3 * W * 2), .ldx = c.kdaProj(), .out = f, .ldo = W, .rows = rows }); + wide.dense(&x.k.wide.?, e, w.g_b, .{ .in = proj.at((3 * W + c.kda_dim) * 2), .ldx = c.kdaProj(), .out = g, .ldo = W, .rows = rows }); + fb = .{ f, f, f }; + gb = .{ g, g, g }; + } + e.setPipeline(switch (gate) { + .generated => if (c.tp > 1) x.k.kda_rows_tp else x.k.kda_rows, + .pre => |p| p, + }); bind(e, 0, .{proj}); shape(e, 1, .{ rows, c.kdaProj() }); - bind(e, 2, .{ L.cs[cur], w.conv_w, w.f_b.w, w.f_b.s, w.f_b.b, w.g_b.w, w.g_b.s, w.g_b.b, w.a, w.dt_bias, L.st[cur], w.o_norm }); + bind(e, 2, .{ L.cs[cur], w.conv_w, fb[0], fb[1], fb[2], gb[0], gb[1], gb[2], w.a, w.dt_bias, L.st[cur], w.o_norm }); e.setValue(c.lower_bound, 14); e.setValue(c.eps, 15); bind(e, 16, .{ y, L.st[1 - cur], L.cs[1 - cur] }); @@ -297,9 +318,10 @@ pub fn mlaCache(x: *const Ctx, e: mtl.ComputeEncoder, mi: usize, w: *const wts.M } } -/// Every row's 64 heads of q_nope (in `qp`, a row qrProj apart) into the latent: MLX's qvm on kv_b's key half. +/// Every row's 64 heads of q_nope (`qp`, rows qrProj apart) into the latent: qvm on kv_b's key half, or embed_q's qmv. pub fn absorb(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Mla, qp: Ref, ql: Ref, rows: u32) void { const c = x.c; + if (w.absorbed) |pair| return wide.perHead(&x.k.wide.?, e, pair[0], c.mla_heads, .{ .in = qp, .ldx = c.qrProj(), .out = ql, .ldo = c.mla_heads * c.kv_lora, .rows = rows }, c.nope, c.kv_lora); e.setPipeline(x.k.absorb); bind(e, 0, .{ w.kv_b.w, w.kv_b.s, w.kv_b.b, qp, ql }); e.setValue(c.qrProj(), 5); @@ -307,9 +329,10 @@ pub fn absorb(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Mla, qp: Ref, e.dispatchGroups(size(1, c.kv_lora / 64, rows * c.mla_heads), size(64, 1, 1)); } -/// Every row's 64 heads' latent outputs to values [rows, heads * v]: MLX's batched qmv_fast on kv_b's value half. +/// Every row's 64 heads' latent outputs to values [rows, heads * v]: qmv_fast on kv_b's value half, or unembed_out's. pub fn unabsorb(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Mla, att: Ref, vals: Ref, rows: u32) void { const c = x.c; + if (w.absorbed) |pair| return wide.perHead(&x.k.wide.?, e, pair[1], c.mla_heads, .{ .in = att, .ldx = c.mla_heads * c.kv_lora, .out = vals, .ldo = c.mla_heads * c.v_dim, .rows = rows }, c.kv_lora, c.v_dim); e.setPipeline(x.k.unabsorb); bind(e, 0, .{ w.kv_b.w, w.kv_b.s, w.kv_b.b, att, vals }); e.setValue(c.mla_heads, 5); diff --git a/zig/src/families/glm/glm.zig b/zig/src/families/glm/glm.zig index 1e116ccb3c..11a7502c8f 100644 --- a/zig/src/families/glm/glm.zig +++ b/zig/src/families/glm/glm.zig @@ -14,6 +14,7 @@ pub const backend = @import("backend.zig"); pub const mirror = @import("mirror.zig"); pub const timing = @import("timing.zig"); pub const affine = @import("affine.zig"); +pub const wide = @import("wide.zig"); test { _ = config; @@ -36,4 +37,5 @@ test { _ = @import("affine.zig"); _ = @import("names.zig"); _ = @import("weights_test.zig"); + _ = @import("wide.zig"); } diff --git a/zig/src/families/glm/kernels.zig b/zig/src/families/glm/kernels.zig index 5f0f3e7916..d3ee3a580b 100644 --- a/zig/src/families/glm/kernels.zig +++ b/zig/src/families/glm/kernels.zig @@ -6,6 +6,7 @@ const frags = @import("../../core/frags.zig"); const moe_route = @import("../../core/moe_route.zig"); const affine_mm = @import("../../core/affine_mm.zig"); const hc = @import("../../core/hc.zig"); +const wide = @import("wide.zig"); /// The route's shape (config.zig refuses other checkpoints). pub const route_shape: moe_route.Shape = .{ .hidden = 4096, .experts = 288, .topk = 8 }; @@ -100,13 +101,23 @@ pub const Kernels = struct { absorb_nax_tp: mtl.Pipeline, mm_bf16: affine_mm.Pipes, // prompt chunks' 4-bit g64 matmuls on the tensor units (core/affine_mm.zig): dense, gathers mm_f32: affine_mm.Pipes, // and with fp32 out: expert parallel by rows' down partials + wide: ?wide.Pipes = null, // the products at other widths (wide.zig), for a checkpoint that has one pub fn deinit(k: *Kernels) void { const info = @typeInfo(Kernels).@"struct"; inline for (info.field_names, info.field_types) |name, T| { - if (T == mtl.Pipeline) @field(k, name).deinit() else if (T != u64) for (&@field(k, name)) |*p| p.deinit(); + if (T == mtl.Pipeline) @field(k, name).deinit() else if (T == ?wide.Pipes) { + if (k.wide) |*w| w.deinit(); + } else if (T != u64) for (&@field(k, name)) |*p| p.deinit(); } } + + /// The other widths' pipelines, once: their sources join the hash a learned state's identity reads. + pub fn widen(k: *Kernels, gpa: std.mem.Allocator, device: mtl.Device) !void { + if (k.wide != null) return; + k.wide = try wide.load(gpa, device); + k.source_hash ^= k.wide.?.hash; + } }; /// The generated kernel `key` (tools/zig/gen_glm_kernels.py) and the field its first function fills. @@ -182,11 +193,21 @@ const Job = struct { }; /// Generated kernel `key`'s template at other parameters, under `name`. -fn instance(comptime key: []const u8, comptime params: []const u8, comptime name: []const u8) []const u8 { +pub fn instance(comptime key: []const u8, comptime params: []const u8, comptime name: []const u8) []const u8 { const f = comptime kernelOf(key).functions[0]; return "\ntemplate [[host_name(\"" ++ name ++ "\")]] [[kernel]] decltype(" ++ f ++ "<" ++ params ++ ">) " ++ f ++ "<" ++ params ++ ">;\n"; } +/// Generated kernel `key`'s source. +pub fn generatedSource(comptime key: []const u8) []const u8 { + return kernelOf(key).source; +} + +/// Generated kernel `key`'s first function. +pub fn generatedFunction(comptime key: []const u8) [:0]const u8 { + return kernelOf(key).functions[0]; +} + fn kernelOf(comptime key: []const u8) sources.glm.Kernel { @setEvalBranchQuota(20000); for (sources.glm.all) |k| if (comptime std.mem.eql(u8, k.key, key)) return k; @@ -197,6 +218,7 @@ fn kernelOf(comptime key: []const u8) sources.glm.Kernel { pub fn load(gpa: std.mem.Allocator, device: mtl.Device) !*Kernels { const k = try gpa.create(Kernels); errdefer gpa.destroy(k); + k.wide = null; var jobs: [generated.len + 22]Job = undefined; inline for (generated, 0..) |g, i| { const src = comptime kernelOf(g.key); diff --git a/zig/src/families/glm/moe.zig b/zig/src/families/glm/moe.zig index ab3faf68fc..458a2ae444 100644 --- a/zig/src/families/glm/moe.zig +++ b/zig/src/families/glm/moe.zig @@ -6,6 +6,7 @@ const st = @import("state.zig"); const fwd = @import("forward.zig"); const moe_route = @import("../../core/moe_route.zig"); const route_shape = @import("kernels.zig").route_shape; +const wide = @import("wide.zig"); const Ref = wts.Ref; const Ctx = fwd.Ctx; const Class = fwd.Class; @@ -39,7 +40,7 @@ pub fn moe(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, r if (s & Class.exchange == 0) ep.send(e, sc.ye, rows, on(x, "x_pack") and on(x, "x_locpost"), on(x, "x_locpost")); if (s & Class.shared == 0) experts(x, e, w, x_in, rows, 1, .{ sc.none, sc.none, sc.none }); if (s & Class.exchange == 0 and on(x, "x_unpack")) ep.receive(e, sc.ye, rows); - } else if (s & (Class.shared | Class.routed) == 0 and on(x, "e_gateup") and on(x, "e_down") and on(x, "s_gateup") and on(x, "s_down")) { + } else if (native(w, .both) and s & (Class.shared | Class.routed) == 0 and on(x, "e_gateup") and on(x, "e_down") and on(x, "s_gateup") and on(x, "s_down")) { if (s & Class.route == 0) route(x, e, w, x_in, rows); both(x, e, w, x_in, rows); // the shared expert as the gathers' last slot: two launches fewer if (s & Class.combine != 0 or !on(x, "combine")) return; @@ -96,6 +97,17 @@ fn route(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, row e.dispatchThreads(size(512, 1, 1), size(512, 1, 1)); } +/// Whether the shared expert, the routed ones or both take the generated 4/64 kernels. +pub fn native(w: *const wts.Moe, which: enum { shared, routed, both }) bool { + const shared = wide.native(w.sh_gate_up) and wide.native(w.sh_down); + const routed = wide.native(w.gate) and wide.native(w.up) and wide.native(w.down); + return switch (which) { + .shared => shared, + .routed => routed, + .both => shared and routed, + }; +} + /// Part 1: the shared expert into `ys`; part 2: the routed experts of `group` (unique ids, members, count) into `ye`. fn experts(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, rows: u32, part: u32, group: [3]Ref) void { const c = x.c; @@ -103,6 +115,8 @@ fn experts(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, r const sc = x.sc; const D = c.hidden; const N = c.moe_inter; + if (part == 1 and !native(w, .shared)) return sharedRows(x, e, w, x_in, rows); + if (part == 2 and !native(w, .routed)) return routedRows(x, e, w, x_in, rows); const zs: u32 = if (part == 1) 1 else rows * c.topk; const slots: u32 = if (part == 1) 1 else c.topk; const act = if (part == 1) sc.acts else sc.act; @@ -176,3 +190,29 @@ fn gateUp(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, ro bind(e, 15, .{act}); e.dispatchThreads(size(32 * rows, if (part == 1) N else N / 4, zs), size(32 * rows, 1, 1)); } + +/// The shared expert at its widths apart from the routed launch (0.6.6 fuses only a 4-bit one): gate/up, SwiGLU, down. +fn sharedRows(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, rows: u32) void { + const c = x.c; + const sc = x.sc; + const N = c.moe_inter; + const p = &x.k.wide.?; + if (on(x, "s_gateup")) { + wide.dense(p, e, w.sh_gate_up, .{ .in = x_in, .ldx = c.hidden, .out = sc.gu, .ldo = 2 * N, .rows = rows }); + e.setPipeline(x.k.swiglu); + bind(e, 0, .{ sc.gu, sc.acts }); + e.setValue(fwd.Rows{ .rows = @intCast(rows), .width = @intCast(N), .x_stride = @intCast(2 * N), .y_stride = @intCast(N), .eps = c.swiglu_limit }, 2); + e.dispatchThreads(size(N, rows, 1), size(256, 1, 1)); + } + if (on(x, "s_down")) wide.dense(p, e, w.sh_down, .{ .in = sc.acts, .ldx = N, .out = sc.ys, .ldo = c.hidden, .rows = rows }); +} + +/// The routed experts at their widths, each (row, pick) its own slot: gate and up SwiGLU'd into `act`, down into `ye`. +fn routedRows(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, rows: u32) void { + const c = x.c; + const sc = x.sc; + const p = &x.k.wide.?; + const slots = rows * c.topk; + if (on(x, "e_gateup")) wide.byIds(p, e, .{ .w = w.gate, .up = w.up, .experts = c.experts, .repeat = c.topk, .lim = c.swiglu_limit }, x_in, sc.pick, sc.act, slots); + if (on(x, "e_down")) wide.byIds(p, e, .{ .w = w.down, .experts = c.experts, .repeat = 1 }, sc.act, sc.pick, sc.ye, slots); +} diff --git a/zig/src/families/glm/mtp.zig b/zig/src/families/glm/mtp.zig index d683c2389f..2f1234c1fb 100644 --- a/zig/src/families/glm/mtp.zig +++ b/zig/src/families/glm/mtp.zig @@ -37,7 +37,9 @@ pub fn absorb(x: *fwd.Ctx, e: mtl.ComputeEncoder, h: Ref, next: Ref, rows: u32, fwd.embedRows(x, e, next, sc.m_emb, rows); fwd.rms(x, e, sc.m_emb, m.enorm, sc.m_eh, rows, D, D, 2 * D, c.eps); fwd.rms(x, e, h, m.hnorm, sc.m_eh.at(@as(usize, D) * 2), rows, D, D, 2 * D, c.eps); - fwd.qmv(x, e, x.k.qmv_kda_out, sc.m_eh, m.eh_proj, sc.m_x, rows); + if (m.eh_bf16) |eh| { // stored unquantized (mlx-lm's): fp32 sums of its bf16 rows + @import("wide.zig").bf16(&x.k.wide.?, e, eh, D, 2 * D, .{ .in = sc.m_eh, .ldx = 2 * D, .out = sc.m_x, .ldo = D, .rows = rows }); + } else fwd.qmv(x, e, x.k.qmv_kda_out, sc.m_eh, m.eh_proj, sc.m_x, rows); fwd.rms(x, e, sc.m_x, L.in_norm, sc.m_xn, rows, D, D, D, c.eps); fwd.mlaKeys(x, e, c.countKind(.mla), &L.attn.mla, sc.m_xn, rows, pos); } diff --git a/zig/src/families/glm/prompt.zig b/zig/src/families/glm/prompt.zig index 8981004096..f163d11619 100644 --- a/zig/src/families/glm/prompt.zig +++ b/zig/src/families/glm/prompt.zig @@ -10,6 +10,7 @@ const ep_mod = @import("ep.zig"); const affine_mm = @import("../../core/affine_mm.zig"); const hc = @import("../../core/hc.zig"); const kernels = @import("kernels.zig"); +const wide = @import("wide.zig"); const Kernels = kernels.Kernels; const moe_route = @import("../../core/moe_route.zig"); const Ref = wts.Ref; @@ -187,14 +188,19 @@ fn run(e: mtl.ComputeEncoder, grid: [3]usize, group: [3]usize) void { /// y [M, N] = x [M, K] W^T for a 4-bit matrix (64x64 tiles; N rounded up to the weights' padded rows, y's row pitch). fn qmm(p: *const Prompt, e: mtl.ComputeEncoder, x: Ref, q: wts.Affine, y: Ref, M: u32) void { + if (!wide.native(q)) return qmmAt(p, e, x, q.k, q, y, M); affine_mm.rowSums(e, p.mm.mm_bf16, 64, x, p.sums, M, q.k); affine_mm.dense(e, p.mm.mm_bf16, x, p.sums, q, y, M); } /// `qmm` over x's rows `ld` elements apart (a slice of a wider row). fn qmmAt(p: *const Prompt, e: mtl.ComputeEncoder, x: Ref, ld: u32, q: wts.Affine, y: Ref, M: u32) void { - affine_mm.rowSumsAt(e, p.mm.mm_bf16, 64, x, p.sums, M, q.k, ld); - affine_mm.denseBatch(e, p.mm.mm_bf16, x, p.sums, q, y, M, 1, .{ .x_row = @intCast(ld), .y_row = @intCast(std.mem.alignForward(u32, q.n, 64)) }); + const pitch = std.mem.alignForward(u32, q.n, 64); + const mm = if (wide.native(q)) p.mm.mm_bf16 else if (q.more == 0 and q.bits == 8 and q.group == 64) p.mm.wide.?.mm8 else { + return wide.dense(&p.mm.wide.?, e, q, .{ .in = x, .ldx = ld, .out = y, .ldo = pitch, .rows = M }); // a split stack, other widths + }; + affine_mm.rowSumsAt(e, mm, 64, x, p.sums, M, q.k, ld); + affine_mm.denseBatch(e, mm, x, p.sums, q, y, M, 1, .{ .x_row = @intCast(ld), .y_row = @intCast(pitch) }); } /// `qmm` with fp32 out (TP2: one Mac's partial of a row-split projection). @@ -255,7 +261,10 @@ fn moe(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, w: *const wts if (s & Class.combine == 0) t.receiveRows(e, p.ys, out, M); return; } - if (s & Class.routed == 0) { // down's rows put back in (row, slot) order for the combine + if (s & Class.routed == 0 and !@import("moe.zig").native(w, .routed)) { // each expert's sorted pairs at its width + wide.bySorted(&p.mm.wide.?, e, .{ .w = w.gate, .up = w.up, .experts = E, .repeat = c.topk, .lim = c.swiglu_limit }, x_in, p.offsets, p.order, p.act, c.moe_inter); + wide.bySorted(&p.mm.wide.?, e, .{ .w = w.down, .experts = E, .repeat = 1 }, p.act, p.offsets, p.order, p.yp, D); + } else if (s & Class.routed == 0) { // down's rows put back in (row, slot) order for the combine gateUp(p, x, e, w, x_in, M, n); affine_mm.gatherTo(e, p.mm.mm_bf16, p.act, p.sums, w.down, p.offsets, p.yp, p.order, n, E); } @@ -355,6 +364,7 @@ fn unabsorb(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, w: *cons const H = c.mla_heads; const K = c.kv_lora; const per = c.nope + c.v_dim; // kv_b's rows a head: its key half, then its value half + if (w.absorbed) |pair| return wide.perHead(&p.mm.wide.?, e, pair[1], H, .{ .in = p.att, .ldx = H * K, .out = p.vals, .ldo = H * c.v_dim, .rows = M }, K, c.v_dim); var q = w.kv_b; q.w = q.w.at(@as(usize, c.nope) * K / 2); q.s = q.s.at(@as(usize, c.nope) * (K / 64) * 2); @@ -390,7 +400,9 @@ fn mla(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, mi: usize, w: qmm(p, e, p.qr, w.qr_proj, p.qp, M); } if (fwd.on(x, "mla_cache")) fwd.mlaCache(x, e, mi, w, x_in, p.xp, XP, p.iw, M, pos); - if (fwd.on(x, "mla_absorb")) { + if (fwd.on(x, "mla_absorb") and w.absorbed != null) { + fwd.absorb(x, e, w, p.qp, p.ql, M); // mlx-lm's embed_q a head at its width, the decode's launch over the chunk + } else if (fwd.on(x, "mla_absorb")) { e.setPipeline(if (c.tp > 1) x.k.absorb_nax_tp else x.k.absorb_nax); // the row kernel's qvm, weights in three bf16 parts bind(e, 0, .{ w.kv_b.w, w.kv_b.s, w.kv_b.b, p.qp, p.ql }); e.setValue([2]i32{ @intCast(M), @intCast(c.qrProj()) }, 5); @@ -499,7 +511,7 @@ pub fn mtp(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, h: Ref, n fwd.embedRows(x, e, next, p.m_emb, M); fwd.rms(x, e, p.m_emb, m.enorm, p.m_eh, M, D, D, 2 * D, c.eps); fwd.rms(x, e, h, m.hnorm, p.m_eh.at(@as(usize, D) * 2), M, D, D, 2 * D, c.eps); - qmm(p, e, p.m_eh, m.eh_proj, p.m_x, M); + if (m.eh_bf16) |eh| wide.bf16(&p.mm.wide.?, e, eh, D, 2 * D, .{ .in = p.m_eh, .ldx = 2 * D, .out = p.m_x, .ldo = D, .rows = M }) else qmm(p, e, p.m_eh, m.eh_proj, p.m_x, M); fwd.rms(x, e, p.m_x, L.in_norm, p.m_xn, M, D, D, D, c.eps); qmm(p, e, p.m_xn, L.attn.mla.x_proj, p.xp, M); fwd.mlaCache(x, e, c.countKind(.mla), &L.attn.mla, p.m_xn, p.xp, @intCast(std.mem.alignForward(usize, c.xProj(), 64)), p.iw, M, pos); diff --git a/zig/src/families/glm/state.zig b/zig/src/families/glm/state.zig index ab01e184e9..2604c1ef2b 100644 --- a/zig/src/families/glm/state.zig +++ b/zig/src/families/glm/state.zig @@ -175,6 +175,7 @@ pub fn initState(arena: *Arena, c: *const cfg.Config, cap: u32, share: ?*const S pub fn init(arena: *Arena, c: *const cfg.Config, cap: u32) !struct { state: State, scratch: Scratch } { const R: usize = max_rows; const D: usize = c.hidden; + if (c.kdaWidth() > c.dense_inter) return error.ScratchTooSmall; // kdaStep puts f_b and g_b products made apart in gu const s = try initState(arena, c, cap, null); const H: usize = c.mla_heads; const V: usize = c.vocab; diff --git a/zig/src/families/glm/weights.zig b/zig/src/families/glm/weights.zig index 9e7d19b77d..5beb758f14 100644 --- a/zig/src/families/glm/weights.zig +++ b/zig/src/families/glm/weights.zig @@ -535,11 +535,6 @@ pub fn load(gpa: std.mem.Allocator, device: mtl.Device, dir: []const u8, c: *con quant.say("glm: this checkpoint has widths other than 4/64; the pair's split kernels read only 4/64 (load it on one Mac)", .{}); return error.UnsupportedWidth; } - if (!dry and !native(w, c)) { - var buf: [256]u8 = undefined; - quant.say("glm: this checkpoint's matrices are {s}; the device kernels read 4/64 only", .{tally(w, c).summary(&buf)}); - return error.UnsupportedWidth; - } if (dry) { var buf: [256]u8 = undefined; const widths = tally(w, c).summary(&buf); diff --git a/zig/src/families/glm/wide.zig b/zig/src/families/glm/wide.zig new file mode 100644 index 0000000000..860b6ba032 --- /dev/null +++ b/zig/src/families/glm/wide.zig @@ -0,0 +1,243 @@ +//! GLM-5.3-Flash's products at widths its generated kernels do not read (glm_rows.metal): a launch a stack member. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const affine = @import("affine.zig"); +const affine_mm = @import("../../core/affine_mm.zig"); +const kernels = @import("kernels.zig"); +const Affine = affine.Affine; +const Ref = affine.Ref; + +const bit_widths = affine.row_bits; +const nw = bit_widths.len; + +/// Whether `q` takes the generated kernels: one width, 4 bits in groups of 64 (the uniform checkpoints' launches). +pub fn native(q: Affine) bool { + return q.more == 0 and q.bits == 4 and q.group == 64; +} + +pub const RowDims = extern struct { k: i32, n: i32, group: i32, rows: i32, heads: i32 = 1, ldx: i32, hsx: i32 = 0, ldo: i32, hso: i32 = 0, hsw: i32 = 0 }; +pub const PickDims = extern struct { k: i32, n: i32, group: i32, slots: i32, repeat: i32, experts: i32, ldo: i32, lim: f32 }; + +pub const Pipes = struct { + rows: [nw]mtl.Pipeline, // dense, strided, by head + ids: [nw]mtl.Pipeline, // a slot through its expert + glu_ids: [nw]mtl.Pipeline, // gate and up through it, SwiGLU'd + sorted: [nw]mtl.Pipeline, // an expert's sorted pairs (prompt chunks) + glu_sorted: [nw]mtl.Pipeline, + bf16: mtl.Pipeline, // an unquantized matrix (the MTP head's eh_proj) + kda_pre: mtl.Pipeline, // kda_rows reading f_b's and g_b's products made apart (any width but 4/64) + embed: [2]mtl.Pipeline, // the embedding at 6 and 8 bits, groups of 64 + mm8: affine_mm.Pipes, // prompt chunks' products of one width at 8 bits in groups of 64 (core/affine_mm.zig) + hash: u64, // every source compiled here + + pub fn deinit(p: *Pipes) void { + inline for (.{ &p.rows, &p.ids, &p.glu_ids, &p.sorted, &p.glu_sorted, &p.embed }) |set| for (set) |x| x.deinit(); + p.bf16.deinit(); + p.kda_pre.deinit(); + for (p.mm8) |x| x.deinit(); + } +}; + +/// kda_rows's f_b dot read from a product made apart (a row's sum in its first lane of four, zeros in the rest: exact). +const pre_f = .{ "result = quad_dot(x, wb, s, bb);", "result = qlid == 0u ? float(((device const bfloat*)FBW)[(size_t)r * W + row]) : 0.0f;" }; +const pre_g = .{ "result = quad_dot(x, wb, s, bb);", "result = qlid == 0u ? float(((device const bfloat*)GBW)[(size_t)r * W + row]) : 0.0f;" }; + +/// glm_glue.metal's bfloat overloads and sigmoid_fast, the text glm_swiglu compiles: glm_rows's SwiGLU takes its bits. +pub const glue_math = blk: { + @setEvalBranchQuota(100_000); + const g = sources.glm_glue; + const a = std.mem.indexOf(u8, g, "// MLX's bfloat math overloads").?; + const b = a + std.mem.indexOf(u8, g[a..], "// The dense MLP's SwiGLU").?; + break :blk "typedef bfloat bfloat16_t;\n" ++ g[a..b]; +}; + +/// The pipelines built so far: a failed compile releases them all. +const Made = struct { + list: [nw * 5 + 4 + affine_mm.names.len]mtl.Pipeline = undefined, + n: usize = 0, + released: usize = 0, + + fn add(m: *Made, device: mtl.Device, l: mtl.Library, name: []const u8) !mtl.Pipeline { + m.list[m.n] = try mtl.Pipeline.init(device, l, name, false); + m.n += 1; + return m.list[m.n - 1]; + } + + fn drop(m: *Made) void { + for (m.list[0..m.n]) |p| p.deinit(); + m.released = m.n; + m.n = 0; + } +}; + +fn lib(m: *Made, device: mtl.Device, src: []const u8, names: []const [:0]const u8, out: []mtl.Pipeline) !void { + const l = try mtl.Library.fromSource(device, src, mtl.CompileOptions.mlx()); + defer l.deinit(); + for (names, out) |n, *p| p.* = try m.add(device, l, n); +} + +pub fn load(gpa: std.mem.Allocator, device: mtl.Device) !Pipes { + var made: Made = .{}; + return build(gpa, device, comptime kernels.generatedSource("kda_rows"), &made); +} + +fn build(gpa: std.mem.Allocator, device: mtl.Device, kda: []const u8, made: *Made) !Pipes { + errdefer made.drop(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + var p: Pipes = undefined; + var h = std.hash.Wyhash.init(0x77); + const rows_src = try std.mem.concat(gpa, u8, &.{ sources.core_row_projection, glue_math, sources.glm_rows }); + defer gpa.free(rows_src); + h.update(rows_src); + { + const l = try mtl.Library.fromSource(device, rows_src, mtl.CompileOptions.mlx()); + defer l.deinit(); + inline for (bit_widths, 0..) |b, i| { + const q = std.fmt.comptimePrint("_q{d}", .{b}); + p.rows[i] = try made.add(device, l, "glm_rows" ++ q); + p.ids[i] = try made.add(device, l, "glm_rows_ids" ++ q); + p.glu_ids[i] = try made.add(device, l, "glm_rows_glu_ids" ++ q); + p.sorted[i] = try made.add(device, l, "glm_rows_sorted" ++ q); + p.glu_sorted[i] = try made.add(device, l, "glm_rows_glu_sorted" ++ q); + } + p.bf16 = try made.add(device, l, "glm_rows_bf16"); + } + const pre_a = try std.mem.replaceOwned(u8, gpa, kda, pre_f[0], pre_f[1]); // its own instance, under its own name + defer gpa.free(pre_a); + const pre_src = try std.mem.replaceOwned(u8, gpa, pre_a, pre_g[0], pre_g[1]); + defer gpa.free(pre_src); + if (std.mem.indexOf(u8, pre_src, pre_f[0]) != null or std.mem.indexOf(u8, pre_src, pre_g[0]) != null or std.mem.indexOf(u8, pre_src, pre_f[1]) == null) return error.KernelCompile; + h.update(pre_src); + try lib(made, device, pre_src, &.{comptime kernels.generatedFunction("kda_rows")}, @as(*[1]mtl.Pipeline, &p.kda_pre)); + h.update(sources.ops_embed_norm); + try lib(made, device, sources.ops_embed_norm, &.{ "tf_embed_b6_g64", "tf_embed_b8_g64" }, &p.embed); + const mm8 = try affine_mm.source(device, gpa, .{ .bits = 8, .group = 64 }); + defer gpa.free(mm8); + h.update(mm8); + try lib(made, device, mm8, &affine_mm.names, &p.mm8); + p.hash = h.final(); + return p; +} + +fn bind(e: mtl.ComputeEncoder, first: usize, refs: anytype) void { + inline for (refs, 0..) |r, i| e.setBuffer(r.buf, r.off, first + i); +} + +fn at(bits: u8) usize { + return std.mem.indexOfScalar(u8, &bit_widths, bits).?; +} + +/// Rows of x (`ldx` apart) through a matrix into out (`ldo` apart). +pub const Io = struct { in: Ref, ldx: u32, out: Ref, ldo: u32, rows: u32 }; + +fn launch(p: *const Pipes, e: mtl.ComputeEncoder, w: [3]Ref, bits: u8, d: RowDims, in: Ref, out: Ref) void { + e.setPipeline(p.rows[at(bits)]); + bind(e, 0, .{ in, w[0], w[1], w[2] }); + e.setValue(d, 4); + bind(e, 5, .{out}); + e.dispatchGroups(mtl.Size.of(@intCast(@divTrunc(d.n + 7, 8)), @intCast(d.heads), 1), mtl.Size.of(64, 1, 1)); +} + +/// x through every member of `q` at its own width, each into its own output rows. +pub fn dense(p: *const Pipes, e: mtl.ComputeEncoder, q: Affine, io: Io) void { + for (0..@as(usize, q.more) + 1) |i| { + const r = q.run(i); + launch(p, e, .{ r.w, r.s, r.b }, r.bits, .{ .k = @intCast(q.k), .n = @intCast(r.n), .group = r.group, .rows = @intCast(io.rows), .ldx = @intCast(io.ldx), .ldo = @intCast(io.ldo) }, io.in, io.out.at(@as(usize, r.row) * 2)); + } +} + +/// `heads` matrices stacked [heads, n / heads, k] (embed_q, unembed_out): head h reads x + h hsx, writes out + h hso. +pub fn perHead(p: *const Pipes, e: mtl.ComputeEncoder, q: Affine, heads: u32, io: Io, hsx: u32, hso: u32) void { + const r = q.n / heads; + launch(p, e, .{ q.w, q.s, q.b }, q.bits, .{ .k = @intCast(q.k), .n = @intCast(r), .group = q.group, .rows = @intCast(io.rows), .heads = @intCast(heads), .ldx = @intCast(io.ldx), .hsx = @intCast(hsx), .ldo = @intCast(io.ldo), .hso = @intCast(hso), .hsw = @intCast(r) }, io.in, io.out); +} + +/// x [rows, k] through bf16 W [n, k]. +pub fn bf16(p: *const Pipes, e: mtl.ComputeEncoder, w: Ref, n: u32, k: u32, io: Io) void { + e.setPipeline(p.bf16); + bind(e, 0, .{ io.in, w }); + e.setValue(RowDims{ .k = @intCast(k), .n = @intCast(n), .group = 0, .rows = @intCast(io.rows), .ldx = @intCast(io.ldx), .ldo = @intCast(io.ldo) }, 4); + bind(e, 5, .{io.out}); + e.dispatchGroups(mtl.Size.of((n + 7) / 8, 1, 1), mtl.Size.of(64, 1, 1)); +} + +/// The routed experts `w` (stacked [experts, n, k]) by pick: with `up`, gate and up SwiGLU'd. +pub const Picks = struct { w: Affine, up: ?Affine = null, experts: u32, repeat: u32, lim: f32 = 0 }; + +fn pickDims(m: Picks, slots: u32, ldo: u32) PickDims { + return .{ .k = @intCast(m.w.k), .n = @intCast(m.w.n), .group = m.w.group, .slots = @intCast(slots), .repeat = @intCast(m.repeat), .experts = @intCast(m.experts), .ldo = @intCast(ldo), .lim = m.lim }; +} + +/// Slot y [0, slots) through expert ids[y] (i32) on x row y / repeat, into out row y. +pub fn byIds(p: *const Pipes, e: mtl.ComputeEncoder, m: Picks, x: Ref, ids: Ref, out: Ref, slots: u32) void { + const u = m.up orelse m.w; + e.setPipeline(if (m.up != null) p.glu_ids[at(m.w.bits)] else p.ids[at(m.w.bits)]); + bind(e, 0, .{ x, m.w.w, m.w.s, m.w.b }); + e.setValue(pickDims(m, slots, m.w.n), 4); + bind(e, 5, .{ out, ids, u.w, u.s, u.b }); + e.dispatchGroups(mtl.Size.of((m.w.n + 7) / 8, slots, 1), mtl.Size.of(64, 1, 1)); +} + +/// Each expert's sorted pairs (offsets, order): with `up` x row order / repeat to row i, else x row i to row order[i]. +pub fn bySorted(p: *const Pipes, e: mtl.ComputeEncoder, m: Picks, x: Ref, offsets: Ref, order: Ref, out: Ref, ldo: u32) void { + const u = m.up orelse m.w; + e.setPipeline(if (m.up != null) p.glu_sorted[at(m.w.bits)] else p.sorted[at(m.w.bits)]); + bind(e, 0, .{ x, m.w.w, m.w.s, m.w.b }); + e.setValue(pickDims(m, 0, ldo), 4); + bind(e, 5, .{ out, offsets, order, u.w, u.s, u.b }); + e.dispatchGroups(mtl.Size.of((m.w.n + 7) / 8, m.experts, 1), mtl.Size.of(64, 1, 1)); +} + +/// kda_rows's f_b/g_b dots: generated (both 4/64), else (`.pre`) glm_rows products made apart, the fp64-checked ones. +pub const KdaGate = union(enum) { generated, pre: mtl.Pipeline }; + +pub fn kdaGate(p: ?*const Pipes, f: Affine, g: Affine) KdaGate { + if (native(f) and native(g)) return .generated; + return .{ .pre = p.?.kda_pre }; +} + +test "4/64 of one width alone takes the generated kernels" { + const r: Ref = .{ .buf = undefined }; + const q: Affine = .{ .w = r, .s = r, .b = r, .n = 64, .k = 128, .bits = 4, .group = 64 }; + try std.testing.expect(native(q)); + var m = q; + m.bits = 8; + try std.testing.expect(!native(m)); + m = q; + m.more = 1; + try std.testing.expect(!native(m)); + try std.testing.expect(kdaGate(null, q, q) == .generated); +} + +test "a failed compile releases every pipeline built before it; a load counts each one it builds" { + const gpa = std.testing.allocator; + const device = try mtl.Device.init(); + defer device.deinit(); + var made: Made = .{}; + try std.testing.expectError(error.KernelCompile, build(gpa, device, "", &made)); // after the row kernels + try std.testing.expectEqual(@as(usize, nw * 5 + 1), made.released); + try std.testing.expectEqual(@as(usize, 0), made.n); + var all: Made = .{}; + var p = try build(gpa, device, comptime kernels.generatedSource("kda_rows"), &all); + p.deinit(); + try std.testing.expectEqual(all.list.len, all.n); +} + +test "glm_rows's SwiGLU compiles glm_glue's sigmoid_fast and bfloat overloads, no copy of its own" { + try std.testing.expect(std.mem.indexOf(u8, glue_math, "inline U sigmoid_fast(U x)") != null); + try std.testing.expect(std.mem.indexOf(u8, glue_math, "METAL_FUNC bfloat16_t exp(bfloat16_t x)") != null); + try std.testing.expect(std.mem.indexOf(u8, glue_math, "kernel void") == null); // the helpers only + try std.testing.expect(std.mem.indexOf(u8, sources.glm_rows, "exp(") == null); + try std.testing.expect(std.mem.indexOf(u8, sources.glm_rows, "sigmoid_fast(gt)") != null); +} + +test "f_b and g_b at any other width (8/64 too) take the products made apart, not the template's own dots" { + const r: Ref = .{ .buf = undefined }; + const q: Affine = .{ .w = r, .s = r, .b = r, .n = 64, .k = 128, .bits = 8, .group = 64 }; + var four = q; + four.bits = 4; + const p: Pipes = undefined; + for ([_][2]Affine{ .{ q, q }, .{ four, q }, .{ q, four } }) |fg| try std.testing.expect(kdaGate(&p, fg[0], fg[1]) == .pre); +} From 883c1e20eab5a12e5027662b74beeef2bc263781 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 04:28:08 -1000 Subject: [PATCH 09/18] tf-glm-rows-check: GLM-5.3-Flash products at every MLX width against float64, each row bits in every launch No model files: seeded synthetic codes, scales, biases and inputs, packed as MLX affine. - Every bits {2,3,4,5,6,8} x group {32,64,128} on three shapes, and every GLM projection class at the public oQ checkpoints' widths at real shapes (KDA in, f_b/g_b, KDA o, MLA x/qr/o, embed_q/unembed_out, dense, shared, routed, lm_head), decode rows and the prompt's product separately, against float64 on every column of two rows. The bound is the fp32 sums' depth (k / 8 + 64 additions) plus a bf16 rounding; each class prints its worst column (max |err| and its share of max |y|). A reference with scales 1% off must exceed the bound in every class and width. - Bit-equal pairs: one row alone == the same row in 2-, 3- and 16-row windows; strided == contiguous; per head == each head dense; by id == dense; GLU by id == glm_swiglu(dense gate | up); sorted == by id (experts of 8, 5, 2 and 1 pairs); the row kernels == core row_projection; a split stack in one call == each member alone (oQ4's KDA input 5/8/5/8/8/8 and x_proj 6/6/5/5, at the decode pitch and the 64-padded prompt pitch); experts [2048, 4096] and [4096, 2048] at 2 and 4 bits by pick and sorted == dense. - KDA: the separate f_b/g_b products then kda_rows == the generated 4/64 kda_rows byte for byte (output, state, conv window) on every 4-bit code, random per-group power-of-two scales and nonzero biases. - The 6- and 8-bit embedding against float64 within two bf16 roundings; a neighbouring group's bias is caught. The engine module exports core/quant (`tensorfold.quant`) for the check's MLX reader. tf-glm-rows-check (M3 Ultra, no tensor units): PASS 1,328,072 values within the fp32-sum bound of float64 (worst 0.9839 of it, bf16 output rounding; worst-column |err| / max|y| 1.7e-3 to 3.6e-3 a class), 1,780,096 bit-equal pairs; scales 1% off exceed the bound in every class and width. Gate on main 0efbd447 + the core commits (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test logs an error on the path it expects). lean_check: 533 problems (main 539), no new findings. --- build.zig | 2 + zig/build/glm_checks.zig | 17 ++ zig/src/tensorfold.zig | 2 + zig/tests/glm_rows_check.zig | 391 +++++++++++++++++++++++++++++++++ zig/tests/glm_rows_classes.zig | 180 +++++++++++++++ zig/tests/glm_rows_kda.zig | 236 ++++++++++++++++++++ 6 files changed, 828 insertions(+) create mode 100644 zig/build/glm_checks.zig create mode 100644 zig/tests/glm_rows_check.zig create mode 100644 zig/tests/glm_rows_classes.zig create mode 100644 zig/tests/glm_rows_kda.zig diff --git a/build.zig b/build.zig index d3b8a0a760..8fd7b6e7d6 100644 --- a/build.zig +++ b/build.zig @@ -4,6 +4,7 @@ const cuda_build = @import("zig/build/cuda.zig"); const dist_build = @import("zig/build/dist.zig"); const qwen27_build = @import("zig/build/qwen27.zig"); const core_checks = @import("zig/build/core_checks.zig"); +const glm_checks = @import("zig/build/glm_checks.zig"); const hip_build = @import("zig/build/hip.zig"); const grammar_build = @import("zig/build/grammar.zig"); const counters_build = @import("zig/build/metal_counters.zig"); @@ -477,6 +478,7 @@ fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b // The core lane, shared-context and Qwen3.8-27B programs and host tests. qwen27_build.targets(b, target, optimize, metal, sources, engine, lanes, mods.qwen27, test_step); core_checks.targets(b, target, optimize, metal, engine); + glm_checks.targets(b, target, optimize, metal, engine, sources); _ = nativeServer(b, target, optimize, metal, engine, api, engines, build_options, test_step); diff --git a/zig/build/glm_checks.zig b/zig/build/glm_checks.zig new file mode 100644 index 0000000000..2088a10fbc --- /dev/null +++ b/zig/build/glm_checks.zig @@ -0,0 +1,17 @@ +//! GLM-5.3-Flash's device checks on synthetic weights: programs that run on the GPU, each under its own step. + +const std = @import("std"); + +/// tf-glm-rows-check: the products at every MLX width against float64, every row's bits in every window and launch. +pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, metal: *std.Build.Module, engine: *std.Build.Module, sources: *std.Build.Module) void { + const exe = b.addExecutable(.{ .name = "tf-glm-rows-check", .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/glm_rows_check.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "tensorfold", .module = engine }, .{ .name = "kernel_sources", .module = sources } }, + }) }); + b.installArtifact(exe); + const about = "GLM-5.3-Flash's products at every MLX width on synthetic weights against float64, each row's bits in every launch"; + b.step("tf-glm-rows-check", about).dependOn(&b.addInstallArtifact(exe, .{}).step); +} diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index f34f41a4ff..d9d7c8a072 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -14,6 +14,8 @@ 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"); +/// The checkpoint formats (MLX affine and its fp64 reader for the device checks, exact widen). +pub const quant = @import("core/quant/quant.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/glm_rows_check.zig b/zig/tests/glm_rows_check.zig new file mode 100644 index 0000000000..cc525805c4 --- /dev/null +++ b/zig/tests/glm_rows_check.zig @@ -0,0 +1,391 @@ +//! GLM-5.3-Flash's products at every MLX width (glm/wide.zig) against fp64, each row's bits in every window and launch. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const sources = @import("kernel_sources"); +const wide = tf.glm.wide; +const affine = tf.glm.affine; +const mlx = tf.quant.mlx; +const row = tf.row_projection; + +const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; +const a = std.heap.page_allocator; +const rows = 16; +const pad = 32; // a strided x row's and out row's extra values +const Ref = affine.Ref; + +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)); +} +fn buf(device: mtl.Device, bytes: usize) !mtl.Buffer { + return device.buffer(@max(bytes, 16), opts); +} +fn ref(b: mtl.Buffer) Ref { + return .{ .buf = b }; +} + +/// A synthetic MLX matrix [n, k], its fp64 dequantization and its codes alone; x: `rows` rows of inputs. +const Mat = struct { + q: affine.Affine, + bufs: [4]mtl.Buffer, + deq: []f64, + codes: []f64, + + fn init(device: mtl.Device, n: usize, k: usize, bits: u8, group: usize, seed: u64) !Mat { + var prng = std.Random.DefaultPrng.init(seed); + const r = prng.random(); + const words = n * k * bits / 32; + const tables = n * (k / group); + const bs = [4]mtl.Buffer{ try buf(device, words * 4), try buf(device, tables * 2), try buf(device, tables * 2), try buf(device, rows * k * 2) }; + for (bs[0].slice(u32, words)) |*v| v.* = r.int(u32); + for (bs[1].slice(u16, tables)) |*v| v.* = toBf(r.float(f32) * 0.02 + 0.001); + for (bs[2].slice(u16, tables)) |*v| v.* = toBf((r.float(f32) - 0.5) * 0.1); + for (bs[3].slice(u16, rows * k)) |*v| v.* = toBf((r.float(f32) - 0.5) * 2); + const m = Mat{ .q = .{ .w = ref(bs[0]), .s = ref(bs[1]), .b = ref(bs[2]), .n = @intCast(n), .k = @intCast(k), .bits = bits, .group = @intCast(group) }, .bufs = bs, .deq = try a.alloc(f64, n * k), .codes = try a.alloc(f64, n * k) }; + try mlx.reference(m.host(bs[1].slice(u16, tables), bs[2].slice(u16, tables)), m.deq); + const ones = try a.alloc(u16, tables); + defer a.free(ones); + @memset(ones, toBf(1)); + const zeros = try a.alloc(u16, tables); + defer a.free(zeros); + @memset(zeros, 0); + try mlx.reference(m.host(ones, zeros), m.codes); + return m; + } + + fn deinit(m: Mat) void { + for (m.bufs) |b| b.deinit(); + a.free(m.deq); + a.free(m.codes); + } + + fn x(m: Mat) mtl.Buffer { + return m.bufs[3]; + } + + fn host(m: Mat, scales: []const u16, biases: []const u16) mlx.Host { + const q = m.q; + const g = q.k / q.group; + return .{ + .words = .{ .dtype = .i32, .rank = 2, .shape = .{ q.n, q.k * q.bits / 32, 1, 1, 1 }, .bytes = m.bufs[0].slice(u8, q.n * q.k * q.bits / 8) }, + .scales = .{ .dtype = .bf16, .rank = 2, .shape = .{ q.n, g, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(scales) }, + .biases = .{ .dtype = .bf16, .rank = 2, .shape = .{ q.n, g, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(biases) }, + .bits = q.bits, + .group = q.group, + }; + } + + /// Row i against column c in fp64, and the fp32 sums' bound (output rounding plus k-deep accumulation). + fn reference(m: Mat, i: usize, c: usize) [2]f64 { + const k: usize = m.q.k; + const g_n = k / m.q.group; + const sc = m.bufs[1].slice(u16, m.q.n * g_n); + const bi = m.bufs[2].slice(u16, m.q.n * g_n); + const xs = m.x().slice(u16, rows * k); + var y: f64 = 0; + var terms: f64 = 0; + for (0..g_n) |g| { + var dot: f64 = 0; + var sum: f64 = 0; + for (0..m.q.group) |j| { + const kk = g * m.q.group + j; + const xv = fromBf(xs[i * k + kk]); + y += xv * m.deq[c * k + kk]; + dot += @abs(xv * m.codes[c * k + kk]); + sum += @abs(xv); + } + terms += @abs(fromBf(sc[c * g_n + g])) * dot + @abs(fromBf(bi[c * g_n + g])) * sum; + } + const gamma = @import("glm_rows_classes.zig").gamma(k); + return .{ y, std.math.ldexp(@abs(y) + gamma * terms, -8) + gamma * terms }; + } + + /// Expert (or head) e of `parts` stacked matrices: its rows, codes and tables. + fn part(m: Mat, e: usize, parts: usize) affine.Affine { + var q = m.q; + const n = q.n / parts; + const wb = affine.wBytes(n, q.k, q.width()); + const sb = affine.sBytes(n, q.k, q.width()); + q.w = q.w.at(e * wb); + q.s = q.s.at(e * sb); + q.b = q.b.at(e * sb); + q.n = @intCast(n); + return q; + } +}; + +const Gpu = struct { + device: mtl.Device, + queue: mtl.Queue, + p: wide.Pipes, + core: row.Pipelines, + swiglu: mtl.Pipeline, + + /// One command buffer around `f(e, args)`. + fn run(g: *const Gpu, comptime f: anytype, args: anytype) !void { + const cb = g.queue.commandBuffer(); + const e = cb.compute(.serial); + @call(.auto, f, .{e} ++ args); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.debug.print("GPU failed: {s}\n", .{msg}); + return error.GpuFailed; + } + } +}; + +fn dense(e: mtl.ComputeEncoder, p: *const wide.Pipes, q: affine.Affine, io: wide.Io) void { + wide.dense(p, e, q, io); +} + +const Tally = struct { values: usize = 0, worst: f64 = 0, pairs: usize = 0 }; + +fn eql(got: []const u16, want: []const u16) !void { + if (!std.mem.eql(u16, got, want)) { + for (got, want, 0..) |x, y, i| if (x != y) { + std.debug.print("first difference at {d}: {x} vs {x}\n", .{ i, x, y }); + break; + }; + return error.BitsDiffer; + } +} + +/// One matrix: fp64 per column, windows and strided rows equal to each row alone, the core row kernel's bits. +fn denseCheck(g: *const Gpu, m: Mat, t: *Tally) !void { + const n: usize = m.q.n; + const k: usize = m.q.k; + const alone = try buf(g.device, rows * n * 2); + defer alone.deinit(); + for (0..rows) |i| try g.run(dense, .{ &g.p, m.q, wide.Io{ .in = ref(m.x()).at(i * k * 2), .ldx = m.q.k, .out = ref(alone).at(i * n * 2), .ldo = m.q.n, .rows = 1 } }); + const got = alone.slice(u16, rows * n); + for (0..rows) |i| for (0..n) |c| { + const r = m.reference(i, c); + const err = @abs(fromBf(got[i * n + c]) - r[0]); + t.worst = @max(t.worst, err / r[1]); + if (err > r[1]) { + std.debug.print("row {d} col {d}: {d} vs fp64 {d} (bound {e})\n", .{ i, c, fromBf(got[i * n + c]), r[0], r[1] }); + return error.PastBound; + } + t.values += 1; + }; + if (n % 8 == 0) { // the core row kernel's sums (its tiles: 8 columns): the same bits for a row + const core = try buf(g.device, n * 2); + defer core.deinit(); + const call = try row.call(&g.core, .{ .w = m.bufs[0], .scales = m.bufs[1], .biases = m.bufs[2], .n = n, .k = k, .group = m.q.group, .bits = m.q.bits }, 1, false); + try g.run(coreRow, .{ call, m, core }); + try eql(core.slice(u16, n), got[0..n]); + t.pairs += n; + } + // windows, contiguous (two rows a weight read) and strided (one at a time, out rows apart with a lead offset) + const xs = try buf(g.device, rows * (k + pad) * 2); + defer xs.deinit(); + for (0..rows) |i| @memcpy(xs.slice(u16, rows * (k + pad))[i * (k + pad) ..][0..k], m.x().slice(u16, rows * k)[i * k ..][0..k]); + const out = try buf(g.device, (rows * (n + pad) + 4) * 2); + defer out.deinit(); + for ([_]usize{ 2, 3, 16 }) |w| for ([_]bool{ false, true }) |strided| { + const ld: usize = if (strided) n + pad else n; + const at: usize = if (strided) 4 else 0; + try g.run(dense, .{ &g.p, m.q, wide.Io{ .in = ref(if (strided) xs else m.x()), .ldx = @intCast(if (strided) k + pad else k), .out = ref(out).at(at * 2), .ldo = @intCast(ld), .rows = @intCast(w) } }); + for (0..w) |i| try eql(out.slice(u16, rows * (n + pad) + 4)[at + i * ld ..][0..n], got[i * n ..][0..n]); + t.pairs += w * n; + }; +} + +fn coreRow(e: mtl.ComputeEncoder, c: row.Call, m: Mat, out: mtl.Buffer) void { + e.setPipeline(c.pipeline); + e.setBuffer(m.x(), 0, 0); + e.setBuffer(m.bufs[0], 0, 1); + e.setBuffer(m.bufs[1], 0, 2); + e.setBuffer(m.bufs[2], 0, 3); + e.setValue(c.dims, 4); + e.setBuffer(out, 0, 5); + e.dispatchGroups(mtl.Size.of(c.groups, c.slots, 1), mtl.Size.of(32 * row.groups, 1, 1)); +} + +/// Four heads (or experts): per head, by pick (GLU and plain) and sorted (8, 5, 2, 1 pairs), each against dense. +fn partsCheck(g: *const Gpu, m: Mat, up: Mat, t: *Tally) !void { + const parts = 4; + const k: usize = m.q.k; + const r: usize = m.q.n / parts; + const n_rows = 3; // x rows a head + const repeat = 2; // picks an x row: 8 rows, so an expert's sorted pairs fill 4- and 8-pair passes and leave tails + const slots = 8 * repeat; + // per head: x [rows, heads, k], out [rows, heads, r] + const ph = try buf(g.device, n_rows * parts * r * 2); + defer ph.deinit(); + try g.run(perHead, .{ &g.p, m.q, @as(u32, parts), wide.Io{ .in = ref(m.x()), .ldx = @intCast(parts * k), .out = ref(ph), .ldo = @intCast(parts * r), .rows = n_rows }, @as(u32, @intCast(k)), @as(u32, @intCast(r)) }); + const one = try buf(g.device, n_rows * r * 2); + defer one.deinit(); + for (0..parts) |h| { + try g.run(dense, .{ &g.p, m.part(h, parts), wide.Io{ .in = ref(m.x()).at(h * k * 2), .ldx = @intCast(parts * k), .out = ref(one), .ldo = @intCast(r), .rows = n_rows } }); + for (0..n_rows) |i| try eql(ph.slice(u16, n_rows * parts * r)[(i * parts + h) * r ..][0..r], one.slice(u16, n_rows * r)[i * r ..][0..r]); + t.pairs += n_rows * r; + } + // by pick: slot y through expert ids[y] on x row y / repeat + const ids = try buf(g.device, slots * 4); + defer ids.deinit(); + const picks = [slots]i32{ 2, 0, 0, 2, 2, 1, 3, 2, 2, 0, 1, 2, 2, 0, 0, 2 }; + @memcpy(ids.slice(i32, slots), &picks); + var gate = m.q; + gate.n = @intCast(r); + var upq = up.q; + upq.n = @intCast(r); + const plain = try buf(g.device, slots * r * 2); + defer plain.deinit(); + const glu = try buf(g.device, slots * r * 2); + defer glu.deinit(); + const lim: f32 = 0.75; + try g.run(byIds, .{ &g.p, wide.Picks{ .w = gate, .experts = parts, .repeat = repeat }, ref(m.x()), ref(ids), ref(plain), @as(u32, slots) }); + try g.run(byIds, .{ &g.p, wide.Picks{ .w = gate, .up = upq, .experts = parts, .repeat = repeat, .lim = lim }, ref(m.x()), ref(ids), ref(glu), @as(u32, slots) }); + const gu = try buf(g.device, slots * 2 * r * 2); + defer gu.deinit(); + const act = try buf(g.device, slots * r * 2); + defer act.deinit(); + for (0..slots) |y| { + const e: usize = @intCast(picks[y]); + const x_at = ref(m.x()).at(y / repeat * k * 2); + try g.run(dense, .{ &g.p, m.part(e, parts), wide.Io{ .in = x_at, .ldx = @intCast(k), .out = ref(gu).at(y * 2 * r * 2), .ldo = @intCast(2 * r), .rows = 1 } }); + try g.run(dense, .{ &g.p, up.part(e, parts), wide.Io{ .in = x_at, .ldx = @intCast(k), .out = ref(gu).at((y * 2 * r + r) * 2), .ldo = @intCast(2 * r), .rows = 1 } }); + try eql(plain.slice(u16, slots * r)[y * r ..][0..r], gu.slice(u16, slots * 2 * r)[y * 2 * r ..][0..r]); + } + try g.run(swiglu, .{ g.swiglu, ref(gu), ref(act), @as(u32, slots), @as(u32, @intCast(r)), lim }); + try eql(glu.slice(u16, slots * r), act.slice(u16, slots * r)); + t.pairs += 2 * slots * r; + // sorted: each expert's pairs in pick order; GLU out row i, down out row order[i] from x row i + var offsets: [parts + 1]i32 = @splat(0); + var order: [slots]i32 = undefined; + var at: usize = 0; + for (0..parts) |e| { + offsets[e] = @intCast(at); + for (picks, 0..) |pk, y| if (pk == e) { + order[at] = @intCast(y); + at += 1; + }; + } + offsets[parts] = @intCast(at); + const ob = try buf(g.device, (parts + 1 + slots) * 4); + defer ob.deinit(); + @memcpy(ob.slice(i32, parts + 1), &offsets); + @memcpy(ob.slice(i32, parts + 1 + slots)[parts + 1 ..], &order); + const order_ref = ref(ob).at((parts + 1) * 4); + const sorted = try buf(g.device, slots * r * 2); + defer sorted.deinit(); + try g.run(bySorted, .{ &g.p, wide.Picks{ .w = gate, .up = upq, .experts = parts, .repeat = repeat, .lim = lim }, ref(m.x()), ref(ob), order_ref, ref(sorted), @as(u32, @intCast(r)) }); + for (0..slots) |i| try eql(sorted.slice(u16, slots * r)[i * r ..][0..r], glu.slice(u16, slots * r)[@as(usize, @intCast(order[i])) * r ..][0..r]); + // the down: x row i (the sorted pairs' inputs, here x rows by order) into out row order[i] + const xs = try buf(g.device, slots * k * 2); + defer xs.deinit(); + for (0..slots) |i| @memcpy(xs.slice(u16, slots * k)[i * k ..][0..k], m.x().slice(u16, rows * k)[@as(usize, @intCast(order[i])) * k ..][0..k]); + try g.run(bySorted, .{ &g.p, wide.Picks{ .w = gate, .experts = parts, .repeat = 1 }, ref(xs), ref(ob), order_ref, ref(sorted), @as(u32, @intCast(r)) }); + try g.run(byIds, .{ &g.p, wide.Picks{ .w = gate, .experts = parts, .repeat = 1 }, ref(m.x()), ref(ids), ref(plain), @as(u32, slots) }); + try eql(sorted.slice(u16, slots * r), plain.slice(u16, slots * r)); + t.pairs += 2 * slots * r; +} + +/// A split stack (a run a member, each its width) in one call, at the decode's and the prompt's pitch == each alone. +fn splitCheck(g: *const Gpu, k: usize, ns: []const usize, bits: []const u8, t: *Tally) !void { + var ms: [affine.max_runs]Mat = undefined; + for (ns, bits, 0..) |n, b, i| ms[i] = try Mat.init(g.device, n, k, b, 64, 500 + 31 * i + n); + defer for (ms[0..ns.len]) |m| m.deinit(); + var q = ms[0].q; + var at: u32 = @intCast(ns[0]); + for (ms[1..ns.len], 0..) |m, i| { + q.runs[i] = .{ .w = m.q.w, .s = m.q.s, .b = m.q.b, .row = at, .n = m.q.n, .bits = m.q.bits, .group = m.q.group }; + at += m.q.n; + } + q.more = @intCast(ns.len - 1); + q.n = at; + const x = ref(ms[0].x()); + for ([_]usize{ at, std.mem.alignForward(usize, at, 64) }) |ld| { + const out = try buf(g.device, rows * ld * 2); + defer out.deinit(); + try g.run(dense, .{ &g.p, q, wide.Io{ .in = x, .ldx = @intCast(k), .out = ref(out), .ldo = @intCast(ld), .rows = rows } }); + for (ms[0..ns.len], 0..) |m, i| { + const r = q.run(i); + const one = try buf(g.device, rows * r.n * 2); + defer one.deinit(); + try g.run(dense, .{ &g.p, m.q, wide.Io{ .in = x, .ldx = @intCast(k), .out = ref(one), .ldo = r.n, .rows = rows } }); + for (0..rows) |y| try eql(out.slice(u16, rows * ld)[y * ld + r.row ..][0..r.n], one.slice(u16, rows * r.n)[y * r.n ..][0..r.n]); + t.pairs += rows * r.n; + } + } +} + +fn perHead(e: mtl.ComputeEncoder, p: *const wide.Pipes, q: affine.Affine, heads: u32, io: wide.Io, hsx: u32, hso: u32) void { + wide.perHead(p, e, q, heads, io, hsx, hso); +} +fn byIds(e: mtl.ComputeEncoder, p: *const wide.Pipes, m: wide.Picks, x: Ref, ids: Ref, out: Ref, slots: u32) void { + wide.byIds(p, e, m, x, ids, out, slots); +} +fn bySorted(e: mtl.ComputeEncoder, p: *const wide.Pipes, m: wide.Picks, x: Ref, offsets: Ref, order: Ref, out: Ref, ldo: u32) void { + wide.bySorted(p, e, m, x, offsets, order, out, ldo); +} +fn swiglu(e: mtl.ComputeEncoder, pipe: mtl.Pipeline, gu: Ref, act: Ref, n_rows: u32, width: u32, lim: f32) void { + e.setPipeline(pipe); + e.setBuffer(gu.buf, gu.off, 0); + e.setBuffer(act.buf, act.off, 1); + e.setValue(extern struct { rows: i32, width: i32, x_stride: i32, y_stride: i32, eps: f32 }{ .rows = @intCast(n_rows), .width = @intCast(width), .x_stride = @intCast(2 * width), .y_stride = @intCast(width), .eps = lim }, 2); + e.dispatchThreads(mtl.Size.of(width, n_rows, 1), mtl.Size.of(256, 1, 1)); +} + +pub fn main() !void { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + const glue = try mtl.Library.fromSource(device, sources.glm_glue, mtl.CompileOptions.mlx()); + defer glue.deinit(); + var g = Gpu{ .device = device, .queue = queue, .p = try wide.load(a, device), .core = try row.Pipelines.load(device), .swiglu = try mtl.Pipeline.init(device, glue, "glm_swiglu", false) }; + defer g.p.deinit(); + defer g.core.deinit(); + defer g.swiglu.deinit(); + std.debug.print("GPU {s}, tensor units: {}\n", .{ device.name(), device.tensorUnits() }); + var all = Tally{}; + const shapes = [_][2]usize{ .{ 132, 384 }, .{ 64, 2688 }, .{ 40, 1536 } }; + for (affine.row_bits) |bits| for ([_]usize{ 32, 64, 128 }) |group| { + var t = Tally{}; + for (shapes, 0..) |s, si| { + const m = try Mat.init(device, s[0], s[1], bits, group, 1000 * @as(u64, bits) + 10 * group + si); + defer m.deinit(); + try denseCheck(&g, m, &t); + if (s[0] == 64) { + const up = try Mat.init(device, s[0], s[1], bits, group, 7 + 1000 * @as(u64, bits) + 10 * group + si); + defer up.deinit(); + try partsCheck(&g, m, up, &t); + } + } + // oQ2's routed experts (and oQ4's at 4 bits): four experts [2048, 4096] (gate/up) and [4096, 2048] (down) + if (group == 64 and (bits == 2 or bits == 4)) for ([_][2]usize{ .{ 4 * 2048, 4096 }, .{ 4 * 4096, 2048 } }, 0..) |s, si| { + const m = try Mat.init(device, s[0], s[1], bits, group, 77 + 1000 * @as(u64, bits) + si); + defer m.deinit(); + const up = try Mat.init(device, s[0], s[1], bits, group, 78 + 1000 * @as(u64, bits) + si); + defer up.deinit(); + try partsCheck(&g, m, up, &t); + }; + std.debug.print("q{d} g{d}: {d} values vs fp64, worst err/bound {d:.4}; {d} bit-equal pairs (windows, strides, heads, picks, sorted, core)\n", .{ bits, group, t.values, t.worst, t.pairs }); + all.values += t.values; + all.pairs += t.pairs; + all.worst = @max(all.worst, t.worst); + }; + var split = Tally{}; + try splitCheck(&g, 256, &.{ 64, 64, 64, 16, 16, 8 }, &.{ 5, 8, 5, 8, 8, 8 }, &split); // oQ4 layer 6's KDA input + try splitCheck(&g, 512, &.{ 48, 16, 8, 8 }, &.{ 6, 6, 5, 5 }, &split); // oQ4 layer 7's and the MTP's MLA x_proj + std.debug.print("split stacks through one call == each member alone: {d} bit-equal values\n", .{split.pairs}); + all.pairs += split.pairs; + all.pairs += try @import("glm_rows_kda.zig").check(&g.p, device, queue); + all.values += try @import("glm_rows_kda.zig").bf16Check(&g.p, device, queue, &all.worst); + all.values += try @import("glm_rows_kda.zig").embedCheck(&g.p, device, queue, &all.worst); + const cl = try @import("glm_rows_classes.zig").check(&g.p, device, queue); + all.values += cl.values; + all.pairs += cl.pairs; + all.worst = @max(all.worst, cl.worst); + std.debug.print("PASS {d} values within the fp32-sum bound (worst err/bound {d:.4}), {d} bit-equal pairs\n", .{ all.values, all.worst, all.pairs }); +} diff --git a/zig/tests/glm_rows_classes.zig b/zig/tests/glm_rows_classes.zig new file mode 100644 index 0000000000..75f5027a45 --- /dev/null +++ b/zig/tests/glm_rows_classes.zig @@ -0,0 +1,180 @@ +//! glm_rows_check's projection classes: each GLM product at the oQ checkpoints' widths, decode and prompt, vs fp64. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const wide = tf.glm.wide; +const affine = tf.glm.affine; +const affine_mm = tf.affine_mm; + +const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; +const rows = 16; // a decode window +const prompt_rows = 33; // a prompt chunk: two 16-row steps and one more + +/// A class's shape [n, k] (a head's for embed_q / unembed_out) and the widths oQ4, oQ2 and the 4-bit build store it at. +const Class = struct { name: []const u8, k: u32, n: u32, bits: []const u8 }; +const classes = [_]Class{ + .{ .name = "kda_in q/k/v", .k = 4096, .n = 8192, .bits = &.{ 5, 8 } }, + .{ .name = "kda_in f_a/g_a", .k = 4096, .n = 128, .bits = &.{ 5, 8 } }, + .{ .name = "kda_in b", .k = 4096, .n = 64, .bits = &.{ 5, 8 } }, + .{ .name = "kda f_b/g_b", .k = 128, .n = 8192, .bits = &.{ 5, 8 } }, + .{ .name = "kda o_proj", .k = 8192, .n = 4096, .bits = &.{ 6, 8 } }, + .{ .name = "mla q_a", .k = 4096, .n = 1536, .bits = &.{ 6, 8 } }, + .{ .name = "mla kv_a", .k = 4096, .n = 512, .bits = &.{ 6, 8 } }, + .{ .name = "indexer wk", .k = 4096, .n = 128, .bits = &.{ 5, 8 } }, + .{ .name = "indexer weights_proj", .k = 4096, .n = 32, .bits = &.{ 5, 8 } }, + .{ .name = "mla x_proj (one width)", .k = 4096, .n = 2208, .bits = &.{8} }, + .{ .name = "mla q_b", .k = 1536, .n = 16384, .bits = &.{ 6, 8 } }, + .{ .name = "indexer wq_b", .k = 1536, .n = 4096, .bits = &.{ 5, 8 } }, + .{ .name = "embed_q (a head)", .k = 256, .n = 512, .bits = &.{ 5, 8 } }, + .{ .name = "unembed_out (a head)", .k = 512, .n = 256, .bits = &.{ 5, 8 } }, + .{ .name = "mla o_proj", .k = 16384, .n = 4096, .bits = &.{ 6, 8 } }, + .{ .name = "dense gate/up", .k = 4096, .n = 12288, .bits = &.{ 5, 8 } }, + .{ .name = "dense down", .k = 12288, .n = 4096, .bits = &.{ 6, 8 } }, + .{ .name = "shared gate/up", .k = 4096, .n = 2048, .bits = &.{8} }, + .{ .name = "shared down", .k = 2048, .n = 4096, .bits = &.{8} }, + .{ .name = "routed gate/up", .k = 4096, .n = 2048, .bits = &.{ 2, 4 } }, + .{ .name = "routed down", .k = 2048, .n = 4096, .bits = &.{ 2, 4 } }, + .{ .name = "lm_head", .k = 4096, .n = 154880, .bits = &.{8} }, +}; + +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)); +} + +/// Code j of a row's little-endian bit stream at `bits` (MLX's packing at every width). +pub fn code(row: []const u8, j: usize, bits: u8) u32 { + const at = j * bits; + var v: u32 = 0; + for (0..4) |b| { + if (at / 8 + b < row.len) v |= @as(u32, row[at / 8 + b]) << @intCast(8 * b); + } + return (v >> @intCast(at % 8)) & ((@as(u32, 1) << @intCast(bits)) - 1); +} + +const Tally = struct { values: usize = 0, worst: f64 = 0, worst_col: usize = 0, pairs: usize = 0, err: f64 = 0, ymax: f64 = 0, caught: usize = 0 }; + +/// fp32 sums' error per unit of |terms|: under k / 8 + 64 additions deep (row kernels and affine_mm), 2^-24 each. +pub fn gamma(k: usize) f64 { + return @as(f64, @floatFromInt(k / 8 + 64)) * 0x1p-24; +} + +/// out row i vs fp64 on every column within the fp32 sums' bound; `scale` != 1: a control, columns past it counted. +fn against(q: affine.Affine, w: []const u8, s: []const u16, b: []const u16, x: []const u16, out: []const u16, pitch: usize, i: usize, t: *Tally, scale: f64) !void { + const k: usize = q.k; + const g_n = k / q.group; + const rb = k * q.bits / 8; + const gm = gamma(k); + for (0..q.n) |c| { + var y: f64 = 0; + var terms: f64 = 0; + for (0..g_n) |g| { + const sc = scale * fromBf(s[c * g_n + g]); + const bi = fromBf(b[c * g_n + g]); + var dot: f64 = 0; + var sum: f64 = 0; + for (0..q.group) |j| { + const kk = g * q.group + j; + const xv = fromBf(x[i * k + kk]); + const cv: f64 = @floatFromInt(code(w[c * rb ..][0..rb], kk, q.bits)); + y += xv * (sc * cv + bi); + dot += @abs(xv * cv); + sum += @abs(xv); + } + terms += @abs(sc) * dot + @abs(bi) * sum; + } + const bound = std.math.ldexp(@abs(y) + gm * terms, -8) + gm * terms; + const err = @abs(fromBf(out[i * pitch + c]) - y); + if (scale != 1) { + t.caught += @intFromBool(err > bound); + continue; + } + t.err = @max(t.err, err); + t.ymax = @max(t.ymax, @abs(y)); + if (err / bound > t.worst) { + t.worst = err / bound; + t.worst_col = c; + } + if (err > bound) { + std.debug.print("col {d} row {d}: {d} vs fp64 {d} (bound {e})\n", .{ c, i, fromBf(out[i * pitch + c]), y, bound }); + return error.PastBound; + } + t.values += 1; + } +} + +fn submit(queue: mtl.Queue, comptime f: anytype, args: anytype) !void { + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + @call(.auto, f, .{e} ++ args); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.debug.print("GPU failed: {s}\n", .{msg}); + return error.GpuFailed; + } +} + +fn rowsLaunch(e: mtl.ComputeEncoder, p: *const wide.Pipes, q: affine.Affine, x: mtl.Buffer, at: usize, out: mtl.Buffer, m: u32) void { + wide.dense(p, e, q, .{ .in = (affine.Ref{ .buf = x }).at(at), .ldx = q.k, .out = .{ .buf = out }, .ldo = q.n, .rows = m }); +} + +fn mm8(e: mtl.ComputeEncoder, p: *const wide.Pipes, q: affine.Affine, x: mtl.Buffer, sums: mtl.Buffer, out: mtl.Buffer, m: u32) void { + const xr: affine.Ref = .{ .buf = x }; + const sr: affine.Ref = .{ .buf = sums }; + affine_mm.rowSums(e, p.mm8, 64, xr, sr, m, q.k); + affine_mm.dense(e, p.mm8, xr, sr, q, affine.Ref{ .buf = out }, m); +} + +/// Each class at its widths (g64): decode rows (alone == windowed) and the prompt's product vs fp64, every column. +pub fn check(p: *const wide.Pipes, device: mtl.Device, queue: mtl.Queue) !Tally { + var all = Tally{}; + var prng = std.Random.DefaultPrng.init(716); + const r = prng.random(); + for (classes) |cl| for (cl.bits) |bits| { + const k: usize = cl.k; + const n: usize = cl.n; + const np = std.mem.alignForward(usize, n, 64); // the loader's padded rows (affine_mm reads whole 64-row tiles) + const w_bytes = np * k * bits / 8; + const tables = np * (k / 64); + const bufs = [_]mtl.Buffer{ try device.buffer(w_bytes, opts), try device.buffer(tables * 2, opts), try device.buffer(tables * 2, opts), try device.buffer(prompt_rows * k * 2, opts), try device.buffer(prompt_rows * np * 2, opts), try device.buffer(prompt_rows * np * 2, opts), try device.buffer(prompt_rows * (k / 64) * 4, opts) }; + defer for (bufs) |x| x.deinit(); + r.bytes(bufs[0].slice(u8, w_bytes)); + for (bufs[1].slice(u16, tables)) |*v| v.* = toBf(r.float(f32) * 0.02 + 0.001); + for (bufs[2].slice(u16, tables)) |*v| v.* = toBf((r.float(f32) - 0.5) * 0.1); + for (bufs[3].slice(u16, prompt_rows * k)) |*v| v.* = toBf((r.float(f32) - 0.5) * 2); + const q: affine.Affine = .{ .w = .{ .buf = bufs[0] }, .s = .{ .buf = bufs[1] }, .b = .{ .buf = bufs[2] }, .n = cl.n, .k = cl.k, .bits = bits, .group = 64 }; + const w = bufs[0].slice(u8, w_bytes); + const s = bufs[1].slice(u16, tables); + const b = bufs[2].slice(u16, tables); + const x = bufs[3].slice(u16, prompt_rows * k); + var t = Tally{}; + // decode: a 16-row window, then rows 0 and 15 alone + try submit(queue, rowsLaunch, .{ p, q, bufs[3], 0, bufs[4], rows }); + for ([_]usize{ 0, rows - 1 }) |i| { + try submit(queue, rowsLaunch, .{ p, q, bufs[3], i * k * 2, bufs[5], 1 }); + if (!std.mem.eql(u16, bufs[5].slice(u16, n), bufs[4].slice(u16, rows * n)[i * n ..][0..n])) return error.WindowDiffersFromRowAlone; + t.pairs += n; + } + for ([_]usize{ 0, rows - 1 }) |i| try against(q, w, s, b, x, bufs[4].slice(u16, rows * n), n, i, &t, 1); + const dec = t.worst; + try against(q, w, s, b, x, bufs[4].slice(u16, rows * n), n, 0, &t, 1.01); // scales 1% off must be seen + if (t.caught == 0) return error.ControlNotCaught; + // prompt: the chunk's product at rows 0 and 32; on the row kernels every row's bits equal the decode's + const eight = bits == 8; + if (eight) try submit(queue, mm8, .{ p, q, bufs[3], bufs[6], bufs[5], prompt_rows }) else try submit(queue, rowsLaunch, .{ p, q, bufs[3], 0, bufs[5], prompt_rows }); + t.worst = 0; + const pitch = if (eight) np else n; + for ([_]usize{ 0, prompt_rows - 1 }) |i| try against(q, w, s, b, x, bufs[5].slice(u16, prompt_rows * pitch), pitch, i, &t, 1); + if (!eight and !std.mem.eql(u16, bufs[5].slice(u16, rows * n), bufs[4].slice(u16, rows * n))) return error.PromptDiffersFromDecode; + std.debug.print("{s} [{d}, {d}] q{d}: decode worst err/bound {d:.4}, prompt ({s}) {d:.4} at col {d}; max |err| {e:.3} = {e:.3} of max |y|; scales 1% off caught in {d} of {d} columns\n", .{ cl.name, n, k, bits, dec, if (eight) "affine_mm" else "row kernels", t.worst, t.worst_col, t.err, t.err / t.ymax, t.caught, n }); + all.values += t.values; + all.pairs += t.pairs; + all.worst = @max(all.worst, @max(dec, t.worst)); + }; + return all; +} diff --git a/zig/tests/glm_rows_kda.zig b/zig/tests/glm_rows_kda.zig new file mode 100644 index 0000000000..d3d307a81d --- /dev/null +++ b/zig/tests/glm_rows_kda.zig @@ -0,0 +1,236 @@ +//! glm_rows_check's KDA step (f_b, g_b made apart == the generated 4-bit dots, exact) and bf16 product (fp64) parts. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const wide = tf.glm.wide; +const affine = tf.glm.affine; + +const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; +const H = 64; +const D = 128; +const W = H * D; +const PS = 3 * W + 2 * D + H; +const R = 3; +const TAPS = 4; + +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)); +} + +fn submit(queue: mtl.Queue, comptime f: anytype, args: anytype) !void { + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + @call(.auto, f, .{e} ++ args); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.debug.print("GPU failed: {s}\n", .{msg}); + return error.GpuFailed; + } +} + +const Bufs = struct { p: mtl.Buffer, cs: mtl.Buffer, cw: mtl.Buffer, fb: [3]mtl.Buffer, gb: [3]mtl.Buffer, a: mtl.Buffer, dtb: mtl.Buffer, st: mtl.Buffer, onw: mtl.Buffer }; +const Outs = struct { y: mtl.Buffer, st: mtl.Buffer, cs: mtl.Buffer }; + +fn step(e: mtl.ComputeEncoder, pipe: mtl.Pipeline, b: Bufs, f: [3]mtl.Buffer, g: [3]mtl.Buffer, o: Outs) void { + e.setPipeline(pipe); + e.setBuffer(b.p, 0, 0); + e.setBytes(std.mem.asBytes(&[2]i32{ R, PS }), 1); + const ins = [_]mtl.Buffer{ b.cs, b.cw, f[0], f[1], f[2], g[0], g[1], g[2], b.a, b.dtb, b.st, b.onw }; + for (ins, 2..) |x, i| e.setBuffer(x, 0, i); + e.setValue(@as(f32, -5.0), 14); + e.setValue(@as(f32, 1e-6), 15); + e.setBuffer(o.y, 0, 16); + e.setBuffer(o.st, 0, 17); + e.setBuffer(o.cs, 0, 18); + e.dispatchThreads(mtl.Size.of(32, 32, H), mtl.Size.of(32, 32, 1)); +} + +fn products(e: mtl.ComputeEncoder, p: *const wide.Pipes, fq: affine.Affine, gq: affine.Affine, proj: mtl.Buffer, f: mtl.Buffer, g: mtl.Buffer) void { + const in: affine.Ref = .{ .buf = proj }; + wide.dense(p, e, fq, .{ .in = in.at(3 * W * 2), .ldx = PS, .out = .{ .buf = f }, .ldo = W, .rows = R }); + wide.dense(p, e, gq, .{ .in = in.at((3 * W + D) * 2), .ldx = PS, .out = .{ .buf = g }, .ldo = W, .rows = R }); +} + +/// kda_rows on f_b, g_b products made apart == the generated 4-bit dots: every code, dyadic scales and biases, exact. +pub fn check(p: *const wide.Pipes, device: mtl.Device, queue: mtl.Queue) !usize { + const kernels = tf.glm.kernels; + const gen_lib = try mtl.Library.fromSource(device, comptime kernels.generatedSource("kda_rows"), mtl.CompileOptions.mlx()); + defer gen_lib.deinit(); + const generated = try mtl.Pipeline.init(device, gen_lib, comptime kernels.generatedFunction("kda_rows"), false); + defer generated.deinit(); + var prng = std.Random.DefaultPrng.init(4242); + const r = prng.random(); + const mk = struct { + fn of(dev: mtl.Device, bytes: usize) !mtl.Buffer { + return dev.buffer(bytes, opts); + } + }.of; + var b: Bufs = undefined; + b.p = try mk(device, R * PS * 2); + b.cs = try mk(device, (TAPS - 1) * 3 * W * 2); + b.cw = try mk(device, TAPS * 3 * W * 4); + for (&b.fb, &b.gb, [_]usize{ W * D / 2, W * 2 * 2, W * 2 * 2 }) |*f, *g, n| { + f.* = try mk(device, n); + g.* = try mk(device, n); + } + b.a = try mk(device, H * 4); + b.dtb = try mk(device, W * 4); + b.st = try mk(device, H * D * D * 4); + b.onw = try mk(device, D * 4); + defer for ([_]mtl.Buffer{ b.p, b.cs, b.cw, b.fb[0], b.fb[1], b.fb[2], b.gb[0], b.gb[1], b.gb[2], b.a, b.dtb, b.st, b.onw }) |x| x.deinit(); + const pv = b.p.slice(u16, R * PS); + for (pv, 0..) |*v, i| { + const c = i % PS; + const small = (c >= 3 * W and c < 3 * W + 2 * D); // f_a and g_a: small integers, the dots exact in any order + v.* = toBf(if (small) @floatFromInt(r.intRangeAtMost(i32, -2, 2)) else (r.float(f32) - 0.5) * 2); + } + for (b.cs.slice(u16, (TAPS - 1) * 3 * W)) |*v| v.* = toBf((r.float(f32) - 0.5) * 2); + for (b.cw.slice(f32, TAPS * 3 * W)) |*v| v.* = (r.float(f32) - 0.5) * 0.5; + for ([_]mtl.Buffer{ b.fb[0], b.gb[0] }) |x| for (x.slice(u8, W * D / 2)) |*v| { + v.* = r.int(u8); // every 4-bit code + }; + for ([_]mtl.Buffer{ b.fb[1], b.gb[1] }) |x| for (x.slice(u16, W * 2)) |*v| { + v.* = toBf(std.math.ldexp(@as(f32, 1), -r.intRangeAtMost(i32, 5, 8))); // a scale a group + }; + for ([_]mtl.Buffer{ b.fb[2], b.gb[2] }) |x| for (x.slice(u16, W * 2)) |*v| { + const m = r.intRangeAtMost(i32, 1, 8); + v.* = toBf(@as(f32, @floatFromInt(if (r.boolean()) m else -m)) / 64); // a nonzero bias a group + }; + for (b.a.slice(f32, H)) |*v| v.* = -r.float(f32); + for (b.dtb.slice(f32, W)) |*v| v.* = (r.float(f32) - 0.5) * 0.1; + for (b.st.slice(f32, H * D * D)) |*v| v.* = (r.float(f32) - 0.5) * 0.1; + for (b.onw.slice(f32, D)) |*v| v.* = 1 + (r.float(f32) - 0.5) * 0.1; + var outs: [2]Outs = undefined; + for (&outs) |*o| o.* = .{ .y = try mk(device, R * W * 2), .st = try mk(device, H * D * D * 4), .cs = try mk(device, (TAPS - 1) * 3 * W * 2) }; + defer for (outs) |o| for ([_]mtl.Buffer{ o.y, o.st, o.cs }) |x| x.deinit(); + try submit(queue, step, .{ generated, b, b.fb, b.gb, outs[0] }); // the generated 4-bit dots + const fq: affine.Affine = .{ .w = .{ .buf = b.fb[0] }, .s = .{ .buf = b.fb[1] }, .b = .{ .buf = b.fb[2] }, .n = W, .k = D, .bits = 4, .group = 64 }; + var gq = fq; + gq.w = .{ .buf = b.gb[0] }; + gq.s = .{ .buf = b.gb[1] }; + gq.b = .{ .buf = b.gb[2] }; + const f = try mk(device, 16 * W * 2); + defer f.deinit(); + const g = try mk(device, 16 * W * 2); + defer g.deinit(); + try submit(queue, products, .{ p, fq, gq, b.p, f, g }); + try submit(queue, step, .{ p.kda_pre, b, [3]mtl.Buffer{ f, f, f }, [3]mtl.Buffer{ g, g, g }, outs[1] }); + var same: usize = 0; + for ([_][2][]const u8{ .{ std.mem.sliceAsBytes(outs[0].y.slice(u16, R * W)), std.mem.sliceAsBytes(outs[1].y.slice(u16, R * W)) }, .{ outs[0].st.slice(u8, H * D * D * 4), outs[1].st.slice(u8, H * D * D * 4) }, .{ outs[0].cs.slice(u8, (TAPS - 1) * 3 * W * 2), outs[1].cs.slice(u8, (TAPS - 1) * 3 * W * 2) } }) |pair| { + if (!std.mem.eql(u8, pair[0], pair[1])) return error.KdaPreDiffers; + same += pair[0].len; + } + var nonzero: usize = 0; + for (outs[0].y.slice(u16, R * W)) |v| nonzero += @intFromBool(v & 0x7fff != 0); + if (nonzero < R * W / 2) return error.KdaOutputsMostlyZero; + std.debug.print("kda_rows on f_b, g_b made apart == the generated 4-bit dots: {d} rows, outputs, state and conv window byte-equal ({d} bytes)\n", .{ R, same }); + return R * W; +} + +/// The bf16 product (an unquantized eh_proj) against fp64 per column, and a row alone == the row in a window. +pub fn bf16Check(p: *const wide.Pipes, device: mtl.Device, queue: mtl.Queue, worst: *f64) !usize { + const n = 40; + const k = 8192; + const rows = 5; + var prng = std.Random.DefaultPrng.init(99); + const r = prng.random(); + const w = try device.buffer(n * k * 2, opts); + defer w.deinit(); + const x = try device.buffer(rows * k * 2, opts); + defer x.deinit(); + for (w.slice(u16, n * k)) |*v| v.* = toBf((r.float(f32) - 0.5) * 0.05); + for (x.slice(u16, rows * k)) |*v| v.* = toBf((r.float(f32) - 0.5) * 2); + const all = try device.buffer(rows * n * 2, opts); + defer all.deinit(); + const one = try device.buffer(n * 2, opts); + defer one.deinit(); + const run = struct { + fn f(e: mtl.ComputeEncoder, pp: *const wide.Pipes, wb: mtl.Buffer, in: affine.Ref, out: mtl.Buffer, m: u32) void { + wide.bf16(pp, e, .{ .buf = wb }, n, k, .{ .in = in, .ldx = k, .out = .{ .buf = out }, .ldo = n, .rows = m }); + } + }.f; + try submit(queue, run, .{ p, w, affine.Ref{ .buf = x }, all, rows }); + var values: usize = 0; + for (0..rows) |i| { + try submit(queue, run, .{ p, w, (affine.Ref{ .buf = x }).at(i * k * 2), one, 1 }); + if (!std.mem.eql(u16, one.slice(u16, n), all.slice(u16, rows * n)[i * n ..][0..n])) return error.Bf16WindowDiffers; + for (0..n) |c| { + var y: f64 = 0; + var terms: f64 = 0; + for (0..k) |j| { + const t = fromBf(x.slice(u16, rows * k)[i * k + j]) * fromBf(w.slice(u16, n * k)[c * k + j]); + y += t; + terms += @abs(t); + } + const bound = std.math.ldexp(@abs(y), -8) + @as(f64, k + 32) * 0x1p-24 * terms; + const err = @abs(fromBf(one.slice(u16, n)[c]) - y); + worst.* = @max(worst.*, err / bound); + if (err > bound) return error.Bf16PastBound; + values += 1; + } + } + std.debug.print("bf16 [{d}, {d}]: {d} values vs fp64 within the fp32-sum bound; each row alone == the row in a {d}-row window\n", .{ n, k, values, rows }); + return values; +} + +/// The 6- and 8-bit embedding's gathers vs fp64 within two bf16 roundings; a neighbouring group's bias must be seen. +pub fn embedCheck(p: *const wide.Pipes, device: mtl.Device, queue: mtl.Queue, worst: *f64) !usize { + const V = 40; + const width = 4096; + const ids = [_]u32{ 7, 0, 39, 7, 21 }; + var prng = std.Random.DefaultPrng.init(808); + const r = prng.random(); + var values: usize = 0; + for ([_]u8{ 6, 8 }, 0..) |bits, pi| { + const rb: usize = width * @as(usize, bits) / 8; + const G = width / 64; + const t = try device.buffer(V * rb, opts); + defer t.deinit(); + const s = try device.buffer(V * G * 2, opts); + defer s.deinit(); + const b = try device.buffer(V * G * 2, opts); + defer b.deinit(); + const idb = try device.buffer(ids.len * 4, opts); + defer idb.deinit(); + const out = try device.buffer(ids.len * width * 2, opts); + defer out.deinit(); + r.bytes(t.slice(u8, V * rb)); + for (s.slice(u16, V * G)) |*v| v.* = toBf(r.float(f32) * 0.02 + 0.001); + for (b.slice(u16, V * G)) |*v| v.* = toBf((r.float(f32) - 0.5) * 0.2); + @memcpy(idb.slice(u32, ids.len), &ids); + const run = struct { + fn f(e: mtl.ComputeEncoder, pipe: mtl.Pipeline, bs: [5]mtl.Buffer, per: u32) void { + e.setPipeline(pipe); + for (bs[0..4], 0..) |x, i| e.setBuffer(x, 0, i); + e.setBytes(std.mem.asBytes(&[1]i32{width}), 4); + e.setBuffer(bs[4], 0, 5); + e.dispatchThreads(mtl.Size.of(width / per, ids.len, 1), mtl.Size.of(256, 1, 1)); + } + }.f; + try submit(queue, run, .{ p.embed[pi], [5]mtl.Buffer{ idb, t, s, b, out }, @as(u32, if (bits == 6) 4 else 1) }); + var caught: usize = 0; + for (ids, 0..) |id, i| for (0..width) |j| { + const g = j / 64; + const q: f64 = @floatFromInt(@import("glm_rows_classes.zig").code(t.slice(u8, V * rb)[id * rb ..][0..rb], j, bits)); + const sq = fromBf(s.slice(u16, V * G)[id * G + g]) * q; + const bias = fromBf(b.slice(u16, V * G)[id * G + g]); + const got = fromBf(out.slice(u16, ids.len * width)[i * width + j]); + const bound = std.math.ldexp(2 * @abs(sq) + @abs(bias), -8) + std.math.ldexp(@abs(sq), -16); + const err = @abs(got - (sq + bias)); + worst.* = @max(worst.*, err / bound); + if (err > bound) return error.EmbedPastBound; + const other = fromBf(b.slice(u16, V * G)[id * G + (g ^ 1)]); + caught += @intFromBool(@abs(got - (sq + other)) > bound); + values += 1; + }; + if (caught == 0) return error.ControlNotCaught; + std.debug.print("embedding q{d} g64: {d} gathered values vs fp64 within two bf16 roundings; a neighbour group's bias caught in {d}\n", .{ bits, ids.len * width, caught }); + } + return values; +} From 5fe7e7890f0bcf654a91bbb346eb8b19928fa281 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 04:28:08 -1000 Subject: [PATCH 10/18] GLM-5.3-Flash prompt chunks gather 2-bit routed experts on core affine_mm oQ2's routed experts (2-bit, group 64) ran a prompt chunk on the sorted row kernels: 81% of a 2048-row chunk's GPU time. With core affine_mm reading 2-bit codes, wide.routed(gate, up, down) picks affine_mm's 4/64 pipelines, its 2/64 ones (wide.Pipes.mm2) or the sorted row kernels, and 2/64 experts take gatherGlu/gatherTo as 4/64 ones do; other widths stay on the row kernels. Decode is unchanged. Test that fails before: "prompt chunks gather routed experts on affine_mm at 4/64 and 2/64 only; other widths take the sorted row products" (no wide.routed). tf-glm-rows-check adds gatherCheck: 2/64 experts gathered on affine_mm == the dense product (GLU == glm_swiglu(dense gate | up), down == dense) at expert shapes, and the routed q2 prompt classes now run on affine_mm against float64. tf-glm-rows-check (M3 Ultra, no tensor units): PASS 1,328,072 values within the fp32-sum bound of float64 (worst 0.9839 of it, bf16 output rounding; worst-column |err| / max|y| 1.7e-3 to 3.6e-3 a class), 1,976,704 bit-equal pairs; scales 1% off exceed the bound in every class and width. TensorFold/GLM-5.3-Flash-MLX-oQ2-MTP @dc55627 on an M3 Ultra (plan 109.01 GB: 2-bit/64 x126, 4-bit/64 x3, 8-bit/64 x290; loads in 21-24 s): drafted == `"draft": false` 4/4 at 160 tokens (MTP acceptance 74-88%; plain 39.5-40.1 tok/s, drafted 50.9-54.4); resumed == fresh (a 2427-token prompt at 0 and 2425 cached tokens; the second turn resumed at 2425 == a server with `--prompt-cache-gib 0`); concurrent == solo, 2 streams. Cache off: prefill 2427 tokens 5.09-5.11 s (475-477 tok/s; 163 tok/s with the experts on the sorted row kernels), decode 256 tokens 40.8-41.0 tok/s. Gate on main 0efbd447 + the core commits (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test logs an error on the path it expects). lean_check: 533 problems (main 539), no new findings. --- zig/src/families/glm/prompt.zig | 18 ++++---- zig/src/families/glm/wide.zig | 33 +++++++++++++- zig/tests/glm_rows_check.zig | 79 +++++++++++++++++++++++++++++++++ zig/tests/glm_rows_classes.zig | 17 +++---- 4 files changed, 130 insertions(+), 17 deletions(-) diff --git a/zig/src/families/glm/prompt.zig b/zig/src/families/glm/prompt.zig index f163d11619..df293dd0f4 100644 --- a/zig/src/families/glm/prompt.zig +++ b/zig/src/families/glm/prompt.zig @@ -210,10 +210,10 @@ fn qmmF32(p: *const Prompt, e: mtl.ComputeEncoder, x: Ref, q: wts.Affine, y: Ref } /// The routed experts' gate, up and SwiGLU on the M rows of `x_in` its pairs read (in expert order) into `act` ([n, gate.n]). -fn gateUp(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, M: u32, n: u32) void { - affine_mm.rowSums(e, p.mm.mm_bf16, 64, x_in, p.sums, M, w.gate.k); - affine_mm.gatherGlu(e, p.mm.mm_bf16, x_in, p.sums, w.gate, w.up, p.offsets, p.order, p.act, n, x.c.experts, x.c.topk, x.c.swiglu_limit); - affine_mm.rowSums(e, p.mm.mm_bf16, 64, p.act, p.sums, n, w.down.k); +fn gateUp(p: *const Prompt, mm: affine_mm.Pipes, x: *const fwd.Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, M: u32, n: u32) void { + affine_mm.rowSums(e, mm, 64, x_in, p.sums, M, w.gate.k); + affine_mm.gatherGlu(e, mm, x_in, p.sums, w.gate, w.up, p.offsets, p.order, p.act, n, x.c.experts, x.c.topk, x.c.swiglu_limit); + affine_mm.rowSums(e, mm, 64, p.act, p.sums, n, w.down.k); } fn swiglu(x: *const fwd.Ctx, e: mtl.ComputeEncoder, gu: Ref, act: Ref, rows: u32, width: u32) void { @@ -253,7 +253,7 @@ fn moe(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, w: *const wts if (s & Class.route == 0) route(p, x, e, w, x_in, M); if (ep) |t| { if (s & Class.routed == 0) { // down's fp32 partials put back in (row, slot) order for the decode's combine - gateUp(p, x, e, w, x_in, M, n); + gateUp(p, p.mm.mm_bf16, x, e, w, x_in, M, n); affine_mm.gatherTo(e, p.mm.mm_f32, p.act, p.sums, w.down, p.offsets, p.yf, p.order, n, E); } if (s & Class.exchange == 0) t.sendRows(e, p.yf, p.wts, M); @@ -261,12 +261,14 @@ fn moe(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, w: *const wts if (s & Class.combine == 0) t.receiveRows(e, p.ys, out, M); return; } - if (s & Class.routed == 0 and !@import("moe.zig").native(w, .routed)) { // each expert's sorted pairs at its width + const path = wide.routed(w.gate, w.up, w.down); + if (s & Class.routed == 0 and path == .rows) { // each expert's sorted pairs at its width wide.bySorted(&p.mm.wide.?, e, .{ .w = w.gate, .up = w.up, .experts = E, .repeat = c.topk, .lim = c.swiglu_limit }, x_in, p.offsets, p.order, p.act, c.moe_inter); wide.bySorted(&p.mm.wide.?, e, .{ .w = w.down, .experts = E, .repeat = 1 }, p.act, p.offsets, p.order, p.yp, D); } else if (s & Class.routed == 0) { // down's rows put back in (row, slot) order for the combine - gateUp(p, x, e, w, x_in, M, n); - affine_mm.gatherTo(e, p.mm.mm_bf16, p.act, p.sums, w.down, p.offsets, p.yp, p.order, n, E); + const mm = if (path == .mm4) p.mm.mm_bf16 else p.mm.wide.?.mm2; + gateUp(p, mm, x, e, w, x_in, M, n); + affine_mm.gatherTo(e, mm, p.act, p.sums, w.down, p.offsets, p.yp, p.order, n, E); } if (s & Class.combine != 0) return; e.setPipeline(k.moe_combine); diff --git a/zig/src/families/glm/wide.zig b/zig/src/families/glm/wide.zig index 860b6ba032..9c9fc9dafe 100644 --- a/zig/src/families/glm/wide.zig +++ b/zig/src/families/glm/wide.zig @@ -29,6 +29,7 @@ pub const Pipes = struct { kda_pre: mtl.Pipeline, // kda_rows reading f_b's and g_b's products made apart (any width but 4/64) embed: [2]mtl.Pipeline, // the embedding at 6 and 8 bits, groups of 64 mm8: affine_mm.Pipes, // prompt chunks' products of one width at 8 bits in groups of 64 (core/affine_mm.zig) + mm2: affine_mm.Pipes, // and at 2 bits in groups of 64: oQ2's routed experts, gathered hash: u64, // every source compiled here pub fn deinit(p: *Pipes) void { @@ -36,6 +37,7 @@ pub const Pipes = struct { p.bf16.deinit(); p.kda_pre.deinit(); for (p.mm8) |x| x.deinit(); + for (p.mm2) |x| x.deinit(); } }; @@ -54,7 +56,7 @@ pub const glue_math = blk: { /// The pipelines built so far: a failed compile releases them all. const Made = struct { - list: [nw * 5 + 4 + affine_mm.names.len]mtl.Pipeline = undefined, + list: [nw * 5 + 4 + 2 * affine_mm.names.len]mtl.Pipeline = undefined, n: usize = 0, released: usize = 0, @@ -117,6 +119,10 @@ fn build(gpa: std.mem.Allocator, device: mtl.Device, kda: []const u8, made: *Mad defer gpa.free(mm8); h.update(mm8); try lib(made, device, mm8, &affine_mm.names, &p.mm8); + const mm2 = try affine_mm.source(device, gpa, .{ .bits = 2, .group = 64 }); + defer gpa.free(mm2); + h.update(mm2); + try lib(made, device, mm2, &affine_mm.names, &p.mm2); p.hash = h.final(); return p; } @@ -190,6 +196,15 @@ pub fn bySorted(p: *const Pipes, e: mtl.ComputeEncoder, m: Picks, x: Ref, offset e.dispatchGroups(mtl.Size.of((m.w.n + 7) / 8, m.experts, 1), mtl.Size.of(64, 1, 1)); } +/// A prompt chunk's routed experts: affine_mm gathers at 4/64 (`mm4`) or 2/64 (`mm2`, oQ2's), else sorted row products. +pub const Routed = enum { mm4, mm2, rows }; + +pub fn routed(gate: Affine, up: Affine, down: Affine) Routed { + if (native(gate) and native(up) and native(down)) return .mm4; + for ([_]Affine{ gate, up, down }) |q| if (q.more != 0 or q.bits != 2 or q.group != 64) return .rows; + return .mm2; +} + /// kda_rows's f_b/g_b dots: generated (both 4/64), else (`.pre`) glm_rows products made apart, the fp64-checked ones. pub const KdaGate = union(enum) { generated, pre: mtl.Pipeline }; @@ -233,6 +248,22 @@ test "glm_rows's SwiGLU compiles glm_glue's sigmoid_fast and bfloat overloads, n try std.testing.expect(std.mem.indexOf(u8, sources.glm_rows, "sigmoid_fast(gt)") != null); } +test "prompt chunks gather routed experts on affine_mm at 4/64 and 2/64 only; other widths take the sorted row products" { + const r: Ref = .{ .buf = undefined }; + const q: Affine = .{ .w = r, .s = r, .b = r, .n = 64, .k = 128, .bits = 2, .group = 64 }; + var four = q; + four.bits = 4; + try std.testing.expectEqual(Routed.mm4, routed(four, four, four)); + try std.testing.expectEqual(Routed.mm2, routed(q, q, q)); + var three = q; + three.bits = 3; + var g32 = q; + g32.group = 32; + try std.testing.expectEqual(Routed.rows, routed(q, q, three)); + try std.testing.expectEqual(Routed.rows, routed(q, q, four)); + try std.testing.expectEqual(Routed.rows, routed(g32, g32, g32)); +} + test "f_b and g_b at any other width (8/64 too) take the products made apart, not the template's own dots" { const r: Ref = .{ .buf = undefined }; const q: Affine = .{ .w = r, .s = r, .b = r, .n = 64, .k = 128, .bits = 8, .group = 64 }; diff --git a/zig/tests/glm_rows_check.zig b/zig/tests/glm_rows_check.zig index cc525805c4..1f1858ca32 100644 --- a/zig/tests/glm_rows_check.zig +++ b/zig/tests/glm_rows_check.zig @@ -7,6 +7,7 @@ const wide = tf.glm.wide; const affine = tf.glm.affine; const mlx = tf.quant.mlx; const row = tf.row_projection; +const affine_mm = tf.affine_mm; const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; const a = std.heap.page_allocator; @@ -288,6 +289,83 @@ fn partsCheck(g: *const Gpu, m: Mat, up: Mat, t: *Tally) !void { t.pairs += 2 * slots * r; } +fn mmDense(e: mtl.ComputeEncoder, mm: affine_mm.Pipes, q: affine.Affine, x: Ref, sums: Ref, y: Ref, n_rows: u32) void { + affine_mm.rowSums(e, mm, 64, x, sums, n_rows, q.k); + affine_mm.dense(e, mm, x, sums, q, y, n_rows); +} + +/// dims: x rows, sorted pairs, experts, picks a row. +fn mmGlu(e: mtl.ComputeEncoder, mm: affine_mm.Pipes, gate: affine.Affine, up: affine.Affine, x: Ref, sums: Ref, sorted: [2]Ref, act: Ref, dims: [4]u32, lim: f32) void { + affine_mm.rowSums(e, mm, 64, x, sums, dims[0], gate.k); + affine_mm.gatherGlu(e, mm, x, sums, gate, up, sorted[0], sorted[1], act, dims[1], dims[2], dims[3], lim); +} + +fn mmDown(e: mtl.ComputeEncoder, mm: affine_mm.Pipes, q: affine.Affine, x: Ref, sums: Ref, sorted: [2]Ref, y: Ref, slots: u32, experts: u32) void { + affine_mm.rowSums(e, mm, 64, x, sums, slots, q.k); + affine_mm.gatherTo(e, mm, x, sums, q, sorted[0], y, sorted[1], slots, experts); +} + +/// oQ2's prompt experts, affine_mm at 2/64: gathered GLU == glm_swiglu(dense gate | up), gathered down == dense. +fn gatherCheck(g: *const Gpu, m: Mat, up: Mat, t: *Tally) !void { + const parts = 4; + const k: usize = m.q.k; + const r: usize = m.q.n / parts; + const repeat = 2; + const slots = 8 * repeat; + const picks = [slots]i32{ 2, 0, 0, 2, 2, 1, 3, 2, 2, 0, 1, 2, 2, 0, 0, 2 }; + const lim: f32 = 0.75; + const mm = g.p.mm2; + var order: [slots]i32 = undefined; + var offsets: [parts]i32 = undefined; + var at: usize = 0; + for (0..parts) |e| { + offsets[e] = @intCast(at); + for (picks, 0..) |pk, y| if (pk == e) { + order[at] = @intCast(y); + at += 1; + }; + } + const ob = try buf(g.device, (parts + slots) * 4); + defer ob.deinit(); + @memcpy(ob.slice(i32, parts), &offsets); + @memcpy(ob.slice(i32, parts + slots)[parts..], &order); + const sorted = [2]Ref{ ref(ob), ref(ob).at(parts * 4) }; + const bs = [_]mtl.Buffer{ try buf(g.device, rows * (k / 64) * 4), try buf(g.device, 2 * parts * rows * r * 2), try buf(g.device, slots * 2 * r * 2), try buf(g.device, 2 * slots * r * 2), try buf(g.device, slots * k * 2), try buf(g.device, (parts + 1) * slots * r * 2) }; + defer for (bs) |b| b.deinit(); + const sums = ref(bs[0]); + var gate = m.q; + gate.n = @intCast(r); + var upq = up.q; + upq.n = @intCast(r); + // every expert's gate and up, dense on every x row: [2][parts][rows, r] + for (0..parts) |e| for ([2]Mat{ m, up }, 0..) |w, j| { + try g.run(mmDense, .{ mm, w.part(e, parts), ref(m.x()), sums, ref(bs[1]).at(((j * parts + e) * rows * r) * 2), @as(u32, rows) }); + }; + const dg = bs[1].slice(u16, 2 * parts * rows * r); + const gu = bs[2].slice(u16, slots * 2 * r); + for (order, 0..) |o, i| { + const e: usize = @intCast(picks[@intCast(o)]); + const xr: usize = @as(usize, @intCast(o)) / repeat; + for (0..2) |j| @memcpy(gu[(2 * i + j) * r ..][0..r], dg[((j * parts + e) * rows + xr) * r ..][0..r]); + } + const act = ref(bs[3]); + try g.run(mmGlu, .{ mm, gate, upq, ref(m.x()), sums, sorted, act, [4]u32{ rows, slots, parts, repeat }, lim }); + try g.run(swiglu, .{ g.swiglu, ref(bs[2]), act.at(slots * r * 2), @as(u32, slots), @as(u32, @intCast(r)), lim }); + const both = bs[3].slice(u16, 2 * slots * r); + try eql(both[0 .. slots * r], both[slots * r ..]); + // the down: sorted pair i's input (x row order[i]) into row order[i], == its expert's dense product of that row + const xs = bs[4].slice(u16, slots * k); + for (order, 0..) |o, i| @memcpy(xs[i * k ..][0..k], m.x().slice(u16, rows * k)[@as(usize, @intCast(o)) * k ..][0..k]); + const y = bs[5].slice(u16, (parts + 1) * slots * r); + try g.run(mmDown, .{ mm, gate, ref(bs[4]), sums, sorted, ref(bs[5]), @as(u32, slots), @as(u32, parts) }); + for (0..parts) |e| try g.run(mmDense, .{ mm, m.part(e, parts), ref(bs[4]), sums, ref(bs[5]).at((1 + e) * slots * r * 2), @as(u32, slots) }); + for (order, 0..) |o, i| { + const e: usize = @intCast(picks[@intCast(o)]); + try eql(y[@as(usize, @intCast(o)) * r ..][0..r], y[((1 + e) * slots + i) * r ..][0..r]); + } + t.pairs += 2 * slots * r; +} + /// A split stack (a run a member, each its width) in one call, at the decode's and the prompt's pitch == each alone. fn splitCheck(g: *const Gpu, k: usize, ns: []const usize, bits: []const u8, t: *Tally) !void { var ms: [affine.max_runs]Mat = undefined; @@ -369,6 +447,7 @@ pub fn main() !void { const up = try Mat.init(device, s[0], s[1], bits, group, 78 + 1000 * @as(u64, bits) + si); defer up.deinit(); try partsCheck(&g, m, up, &t); + if (bits == 2) try gatherCheck(&g, m, up, &t); // the prompt's 2/64 experts on affine_mm }; std.debug.print("q{d} g{d}: {d} values vs fp64, worst err/bound {d:.4}; {d} bit-equal pairs (windows, strides, heads, picks, sorted, core)\n", .{ bits, group, t.values, t.worst, t.pairs }); all.values += t.values; diff --git a/zig/tests/glm_rows_classes.zig b/zig/tests/glm_rows_classes.zig index 75f5027a45..966df1d6c4 100644 --- a/zig/tests/glm_rows_classes.zig +++ b/zig/tests/glm_rows_classes.zig @@ -123,11 +123,11 @@ fn rowsLaunch(e: mtl.ComputeEncoder, p: *const wide.Pipes, q: affine.Affine, x: wide.dense(p, e, q, .{ .in = (affine.Ref{ .buf = x }).at(at), .ldx = q.k, .out = .{ .buf = out }, .ldo = q.n, .rows = m }); } -fn mm8(e: mtl.ComputeEncoder, p: *const wide.Pipes, q: affine.Affine, x: mtl.Buffer, sums: mtl.Buffer, out: mtl.Buffer, m: u32) void { +fn mmLaunch(e: mtl.ComputeEncoder, mm: affine_mm.Pipes, q: affine.Affine, x: mtl.Buffer, sums: mtl.Buffer, out: mtl.Buffer, m: u32) void { const xr: affine.Ref = .{ .buf = x }; const sr: affine.Ref = .{ .buf = sums }; - affine_mm.rowSums(e, p.mm8, 64, xr, sr, m, q.k); - affine_mm.dense(e, p.mm8, xr, sr, q, affine.Ref{ .buf = out }, m); + affine_mm.rowSums(e, mm, 64, xr, sr, m, q.k); + affine_mm.dense(e, mm, xr, sr, q, affine.Ref{ .buf = out }, m); } /// Each class at its widths (g64): decode rows (alone == windowed) and the prompt's product vs fp64, every column. @@ -165,13 +165,14 @@ pub fn check(p: *const wide.Pipes, device: mtl.Device, queue: mtl.Queue) !Tally try against(q, w, s, b, x, bufs[4].slice(u16, rows * n), n, 0, &t, 1.01); // scales 1% off must be seen if (t.caught == 0) return error.ControlNotCaught; // prompt: the chunk's product at rows 0 and 32; on the row kernels every row's bits equal the decode's - const eight = bits == 8; - if (eight) try submit(queue, mm8, .{ p, q, bufs[3], bufs[6], bufs[5], prompt_rows }) else try submit(queue, rowsLaunch, .{ p, q, bufs[3], 0, bufs[5], prompt_rows }); + const mm: ?affine_mm.Pipes = if (bits == 8) p.mm8 else if (bits == 2) p.mm2 else null; // prompt.zig's (2: routed) + const on_mm = mm != null; + if (mm) |pipes| try submit(queue, mmLaunch, .{ pipes, q, bufs[3], bufs[6], bufs[5], prompt_rows }) else try submit(queue, rowsLaunch, .{ p, q, bufs[3], 0, bufs[5], prompt_rows }); t.worst = 0; - const pitch = if (eight) np else n; + const pitch = if (on_mm) np else n; for ([_]usize{ 0, prompt_rows - 1 }) |i| try against(q, w, s, b, x, bufs[5].slice(u16, prompt_rows * pitch), pitch, i, &t, 1); - if (!eight and !std.mem.eql(u16, bufs[5].slice(u16, rows * n), bufs[4].slice(u16, rows * n))) return error.PromptDiffersFromDecode; - std.debug.print("{s} [{d}, {d}] q{d}: decode worst err/bound {d:.4}, prompt ({s}) {d:.4} at col {d}; max |err| {e:.3} = {e:.3} of max |y|; scales 1% off caught in {d} of {d} columns\n", .{ cl.name, n, k, bits, dec, if (eight) "affine_mm" else "row kernels", t.worst, t.worst_col, t.err, t.err / t.ymax, t.caught, n }); + if (!on_mm and !std.mem.eql(u16, bufs[5].slice(u16, rows * n), bufs[4].slice(u16, rows * n))) return error.PromptDiffersFromDecode; + std.debug.print("{s} [{d}, {d}] q{d}: decode worst err/bound {d:.4}, prompt ({s}) {d:.4} at col {d}; max |err| {e:.3} = {e:.3} of max |y|; scales 1% off caught in {d} of {d} columns\n", .{ cl.name, n, k, bits, dec, if (on_mm) "affine_mm" else "row kernels", t.worst, t.worst_col, t.err, t.err / t.ymax, t.caught, n }); all.values += t.values; all.pairs += t.pairs; all.worst = @max(all.worst, @max(dec, t.worst)); From 92d3af7621f6d62d02f8dfc5541af9d9c35dbfb7 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 09:39:30 -1000 Subject: [PATCH 11/18] tf-glm-rows-speed: GLM-5.3-Flash decode products at the checkpoints' shapes, GPU time and GB/s by class Times every decode product class on synthetic weights at the shapes and widths the uniform 4/64, oQ2 and oQ4 checkpoints read: the generated 4/64 kernels, glm_rows at 8, 6 and 5 bits, the KDA step (generated dots or .pre), the MoE launches (uniform's fused gather, oQ4's routed 4/64, oQ2's 2-bit picks, the 8-bit shared expert). Each class rotates over at least 768 MB of weights; a token's time is summed by class for each checkpoint. GLM_SPEED_CANDIDATES=1 also times glm_rows launch shapes that keep its sums (each checked bit-equal first); GLM_SPEED_PROMPT=1 a 2048-row chunk's per-head and 5/6-bit products. --- zig/build/glm_checks.zig | 8 + zig/tests/glm_rows_speed.metal | 33 ++ zig/tests/glm_rows_speed.zig | 539 +++++++++++++++++++++++++++++++++ 3 files changed, 580 insertions(+) create mode 100644 zig/tests/glm_rows_speed.metal create mode 100644 zig/tests/glm_rows_speed.zig diff --git a/zig/build/glm_checks.zig b/zig/build/glm_checks.zig index 2088a10fbc..bf65167ddd 100644 --- a/zig/build/glm_checks.zig +++ b/zig/build/glm_checks.zig @@ -14,4 +14,12 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu b.installArtifact(exe); const about = "GLM-5.3-Flash's products at every MLX width on synthetic weights against float64, each row's bits in every launch"; b.step("tf-glm-rows-check", about).dependOn(&b.addInstallArtifact(exe, .{}).step); + const speed = b.addExecutable(.{ .name = "tf-glm-rows-speed", .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/glm_rows_speed.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "tensorfold", .module = engine }, .{ .name = "kernel_sources", .module = sources } }, + }) }); + b.step("tf-glm-rows-speed", "GLM-5.3-Flash one-row decode products at the checkpoints' shapes: GPU time a launch, GB/s, a token by class").dependOn(&b.addInstallArtifact(speed, .{}).step); } diff --git a/zig/tests/glm_rows_speed.metal b/zig/tests/glm_rows_speed.metal new file mode 100644 index 0000000000..5b700bf300 --- /dev/null +++ b/zig/tests/glm_rows_speed.metal @@ -0,0 +1,33 @@ +// tf-glm-rows-speed's candidates: glm_rows's one-row sums at RPS columns a simdgroup, SG a group, K compiled in. +template +[[kernel]] void glm_rows_x(const device bfloat* X [[buffer(0)]], const device uint32_t* W [[buffer(1)]], + const device bfloat* S [[buffer(2)]], const device bfloat* B [[buffer(3)]], + constant GlmRowDims& d [[buffer(4)]], device bfloat* OUT [[buffer(5)]], + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int row0 = (int(tg.x) * SG + int(sg)) * RPS; + if (row0 >= d.n) return; + const int k = K > 0 ? K : d.k; + const int h = int(tg.y); + const size_t at = size_t(h) * d.hsw + row0; + const device uint8_t* wr = (const device uint8_t*)W + at * (k * BITS / 8); + const device bfloat* sr = S + at * (k / 64); + const device bfloat* br = B + at * (k / 64); + const device bfloat* x0 = X + size_t(h) * d.hsx; + device bfloat* o0 = OUT + size_t(h) * d.hso + row0; + const int cols = metal::min(RPS, d.n - row0); + #pragma clang loop unroll(disable) + for (int m = 0; m < d.rows; m++) { + float acc[1][RPS]; + tf_rp_rows(wr, sr, br, x0 + size_t(m) * d.ldx, k, 64, lane, acc); + if (lane == 0) + for (int j = 0; j < cols; j++) o0[size_t(m) * d.ldo + j] = bfloat(acc[0][j]); + } +} + +#define GX(B, R, S, K) \ + template [[host_name("x_q" #B "_r" #R "_s" #S "_k" #K)]] [[kernel]] decltype(glm_rows_x) glm_rows_x; +#define GXK(B, K) GX(B, 4, 2, K) GX(B, 8, 2, K) GX(B, 4, 4, K) GX(B, 8, 4, K) GX(B, 2, 4, K) GX(B, 2, 8, K) GX(B, 4, 8, K) +#define GXB(B) GXK(B, 0) GXK(B, 4096) GXK(B, 8192) GXK(B, 16384) GXK(B, 12288) GXK(B, 1536) GXK(B, 2048) GXK(B, 128) GXK(B, 256) GXK(B, 512) +GXB(8) + diff --git a/zig/tests/glm_rows_speed.zig b/zig/tests/glm_rows_speed.zig new file mode 100644 index 0000000000..ab25dc27c8 --- /dev/null +++ b/zig/tests/glm_rows_speed.zig @@ -0,0 +1,539 @@ +//! tf-glm-rows-speed: GLM-5.3-Flash's decode products at the checkpoints' widths and shapes, GPU time and GB/s. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const sources = @import("kernel_sources"); +const wide = tf.glm.wide; +const affine = tf.glm.affine; +const kernels = tf.glm.kernels; +const Ref = affine.Ref; +const Affine = affine.Affine; + +const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; +const a = std.heap.page_allocator; +const hot = 768 << 20; // a class's matrices rotate over at least this many bytes: past the GPU's caches +const reps = 32; // launches a command buffer +const runs = 7; // command buffers a timing (median) + +const H = 64; // KDA and MLA heads +const KD = 128; +const W = H * KD; +const PS = 3 * W + 2 * KD + H; // KDA's stacked projection +const XP = 2208; // MLA x_proj rows +const QR = 20480; // qr_proj rows +const N = 2048; // an expert's intermediate rows +const D = 4096; +const E = 64; // experts in a pool +const TOPK = 8; + +fn buf(device: mtl.Device, bytes: usize) !mtl.Buffer { + return device.buffer(@max(bytes, 16), opts); +} + +fn fill(b: mtl.Buffer, bytes: usize, seed: u64) void { + var s = seed | 1; + for (b.slice(u64, bytes / 8)) |*v| { + s ^= s << 13; + s ^= s >> 7; + s ^= s << 17; + v.* = s; + } +} + +/// `copies` matrices [n, k] at bits/64 (experts > 1: each one stacked [experts, n, k]); codes random, tables small. +const Pool = struct { + q: []Affine, + bufs: []mtl.Buffer, + bytes: usize, // one copy's + + fn init(device: mtl.Device, n: usize, k: usize, bits: u8, experts: usize, min_copies: usize) !Pool { + const w: affine.Width = .{ .bits = bits, .group = 64 }; + const one = affine.bytes(n, k, w) * experts; + const copies = if (experts > 1) min_copies else @max(min_copies, @min(64, (hot + one - 1) / one)); // a pool's picks walk its experts + const p = Pool{ .q = try a.alloc(Affine, copies), .bufs = try a.alloc(mtl.Buffer, 3 * copies), .bytes = one }; + for (0..copies) |c| { + const wb = affine.wBytes(n, k, w) * experts; + const sb = affine.sBytes(n, k, w) * experts; + p.bufs[3 * c] = try buf(device, wb); + fill(p.bufs[3 * c], wb, 0x9e37 + c * 7919 + n + k); + for (1..3) |j| { + p.bufs[3 * c + j] = try buf(device, sb); + @memset(p.bufs[3 * c + j].slice(u16, sb / 2), if (j == 1) 0x3c00 else 0xbc00); // 2^-7, -2^-7 + } + p.q[c] = .{ .w = .{ .buf = p.bufs[3 * c] }, .s = .{ .buf = p.bufs[3 * c + 1] }, .b = .{ .buf = p.bufs[3 * c + 2] }, .n = @intCast(n), .k = @intCast(k), .bits = bits, .group = 64 }; + } + return p; + } + + fn deinit(p: Pool) void { + for (p.bufs) |b| b.deinit(); + a.free(p.q); + a.free(p.bufs); + } + + fn at(p: Pool, i: usize) Affine { + return p.q[i % p.q.len]; + } +}; + +/// Activations and outputs every launch shares (wide enough for any class). +const Io = struct { x: mtl.Buffer, y: mtl.Buffer, z: mtl.Buffer, ids: mtl.Buffer, uids: mtl.Buffer, umem: mtl.Buffer, ucount: mtl.Buffer }; + +const Bench = struct { + device: mtl.Device, + queue: mtl.Queue, + k: *kernels.Kernels, + io: Io, + reps: usize = reps, + + /// Median GPU microseconds a launch of `f(e, i, args)` over `reps` launches i in one command buffer. + fn time(b: *const Bench, comptime f: anytype, args: anytype) !f64 { + var t: [runs + 1]f64 = undefined; + for (&t) |*v| { + const cb = b.queue.commandBuffer(); + const e = cb.compute(.serial); + for (0..b.reps) |i| @call(.auto, f, .{ e, i } ++ args); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.debug.print("GPU failed: {s}\n", .{msg}); + return error.GpuFailed; + } + v.* = cb.gpuSeconds() * 1e6 / @as(f64, @floatFromInt(b.reps)); + } + const s = t[1..]; // the first warms the pipelines + std.mem.sort(f64, s, {}, std.sort.asc(f64)); + return s[runs / 2]; + } + + /// One launch of `f(e, 0, args)`, waited on. + fn run(b: *const Bench, comptime f: anytype, args: anytype) !void { + const cb = b.queue.commandBuffer(); + const e = cb.compute(.serial); + @call(.auto, f, .{ e, @as(usize, 0) } ++ args); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.debug.print("GPU failed: {s}\n", .{msg}); + return error.GpuFailed; + } + } +}; + +fn bind(e: mtl.ComputeEncoder, first: usize, refs: anytype) void { + inline for (refs, 0..) |r, i| e.setBuffer(r.buf, r.off, first + i); +} +fn shape(e: mtl.ComputeEncoder, index: usize, dims: anytype) void { + var v: [dims.len]i32 = undefined; + inline for (dims, 0..) |d, i| v[i] = @intCast(d); + e.setBytes(std.mem.asBytes(&v), index); +} +fn rf(b: mtl.Buffer) Ref { + return .{ .buf = b }; +} + +// Uniform 4/64: the generated kernels' launches (forward.zig, moe.zig). +fn qmv(e: mtl.ComputeEncoder, i: usize, pipe: mtl.Pipeline, p: Pool, io: Io) void { + const q = p.at(i); + e.setPipeline(pipe); + bind(e, 0, .{ rf(io.x), q.w, q.s, q.b, rf(io.y) }); + e.dispatchThreads(mtl.Size.of(32, q.n / 4, 1), mtl.Size.of(32, 1, 1)); +} +fn absorb(e: mtl.ComputeEncoder, i: usize, k: *const kernels.Kernels, p: Pool, io: Io) void { + const q = p.at(i); + e.setPipeline(k.absorb); + bind(e, 0, .{ q.w, q.s, q.b, rf(io.x), rf(io.y) }); + e.setValue(@as(u32, QR), 5); + e.setValue(@as(u32, H), 6); + e.dispatchGroups(mtl.Size.of(1, 512 / 64, H), mtl.Size.of(64, 1, 1)); +} +fn unabsorb(e: mtl.ComputeEncoder, i: usize, k: *const kernels.Kernels, p: Pool, io: Io) void { + const q = p.at(i); + e.setPipeline(k.unabsorb); + bind(e, 0, .{ q.w, q.s, q.b, rf(io.x), rf(io.y) }); + e.setValue(@as(u32, H), 5); + e.dispatchGroups(mtl.Size.of(1, 256 / 4, H), mtl.Size.of(32, 1, 1)); +} +/// The routed experts and the shared one in one gather (moe.zig `both`, one row: two outputs a simdgroup). +fn moeBoth(e: mtl.ComputeEncoder, i: usize, k: *const kernels.Kernels, gu: [2]Pool, down: Pool, sh: [2]Pool, io: Io) void { + const g = gu[0].at(i); + const u = gu[1].at(i); + const s = sh[0].at(i); + const ids = rf(io.uids).at((i % reps) * 128 * 4); + e.setPipeline(k.moe_gateup_0_one); + bind(e, 0, .{rf(io.x)}); + shape(e, 1, .{ 1, D }); + bind(e, 2, .{ g.w, g.s, g.b, u.w, u.s, u.b, s.w, s.s, s.b, ids, rf(io.umem), rf(io.ucount) }); + e.setValue(@as(f32, 10.0), 14); + bind(e, 15, .{rf(io.z)}); + e.dispatchThreads(mtl.Size.of(32, N / 2, TOPK + 1), mtl.Size.of(32, 1, 1)); + const d = down.at(i); + const sd = sh[1].at(i); + e.setPipeline(k.moe_down_0_one); + bind(e, 0, .{rf(io.z)}); + shape(e, 1, .{ 1, TOPK + 1, N }); + bind(e, 2, .{ d.w, d.s, d.b, sd.w, sd.s, sd.b, ids, rf(io.umem), rf(io.ucount), rf(io.y) }); + e.dispatchThreads(mtl.Size.of(32, D / 2, TOPK + 1), mtl.Size.of(32, 1, 1)); +} +/// oQ4's routed experts: the generated 4/64 gathers without the shared slot (moe.zig `experts` part 2). +fn moeRouted4(e: mtl.ComputeEncoder, i: usize, k: *const kernels.Kernels, gu: [2]Pool, down: Pool, io: Io) void { + const g = gu[0].at(i); + const u = gu[1].at(i); + const ids = rf(io.uids).at((i % reps) * 128 * 4); + e.setPipeline(k.moe_gateup_2); + bind(e, 0, .{rf(io.x)}); + shape(e, 1, .{ 1, D }); + bind(e, 2, .{ g.w, g.s, g.b, u.w, u.s, u.b, g.w, g.s, g.b, ids, rf(io.umem), rf(io.ucount) }); + e.setValue(@as(f32, 10.0), 14); + bind(e, 15, .{rf(io.z)}); + e.dispatchThreads(mtl.Size.of(32, N / 4, TOPK), mtl.Size.of(32, 1, 1)); + const d = down.at(i); + e.setPipeline(k.moe_down_2); + bind(e, 0, .{rf(io.z)}); + shape(e, 1, .{ 1, TOPK, N }); + bind(e, 2, .{ d.w, d.s, d.b, d.w, d.s, d.b, ids, rf(io.umem), rf(io.ucount), rf(io.y) }); + e.dispatchThreads(mtl.Size.of(32, D / 4, TOPK), mtl.Size.of(32, 1, 1)); +} + +// Other widths: wide.zig's launches as the engine makes them. +fn rows(e: mtl.ComputeEncoder, i: usize, p: *const wide.Pipes, pool: Pool, io: Io) void { + const q = pool.at(i); + wide.dense(p, e, q, .{ .in = rf(io.x), .ldx = q.k, .out = rf(io.y), .ldo = q.n, .rows = 1 }); +} +fn perHead(e: mtl.ComputeEncoder, i: usize, p: *const wide.Pipes, pool: Pool, io: Io, ldx: u32, hsx: u32) void { + const q = pool.at(i); + wide.perHead(p, e, q, H, .{ .in = rf(io.x), .ldx = ldx, .out = rf(io.y), .ldo = q.n, .rows = 1 }, hsx, q.n / H); +} +fn shared(e: mtl.ComputeEncoder, i: usize, p: *const wide.Pipes, k: *const kernels.Kernels, sh: [2]Pool, io: Io) void { + const gu = sh[0].at(i); + wide.dense(p, e, gu, .{ .in = rf(io.x), .ldx = D, .out = rf(io.z), .ldo = 2 * N, .rows = 1 }); + e.setPipeline(k.swiglu); + bind(e, 0, .{ rf(io.z), rf(io.z).at(2 * N * 2) }); + e.setValue(extern struct { rows: i32, width: i32, xs: i32, ys: i32, lim: f32 }{ .rows = 1, .width = N, .xs = 2 * N, .ys = N, .lim = 10.0 }, 2); + e.dispatchThreads(mtl.Size.of(N, 1, 1), mtl.Size.of(256, 1, 1)); + const dq = sh[1].at(i); + wide.dense(p, e, dq, .{ .in = rf(io.z).at(2 * N * 2), .ldx = N, .out = rf(io.y), .ldo = D, .rows = 1 }); +} +/// A KDA layer's step state (16 copies: each launch its own, as each layer has). +const Kda = struct { + p: mtl.Buffer, + cs: [16]mtl.Buffer, + cw: mtl.Buffer, + a: mtl.Buffer, + dtb: mtl.Buffer, + st: [16]mtl.Buffer, + st_out: [16]mtl.Buffer, + onw: mtl.Buffer, + fg: mtl.Buffer, // f_b's and g_b's products made apart (.pre) + + fn init(device: mtl.Device) !Kda { + var k: Kda = undefined; + k.p = try buf(device, PS * 2); + @memset(k.p.slice(u16, PS), 0x3c00); + k.cw = try buf(device, 4 * 3 * W * 4); + @memset(k.cw.slice(f32, 4 * 3 * W), 0.25); + k.a = try buf(device, H * 4); + @memset(k.a.slice(f32, H), 0.5); + k.dtb = try buf(device, W * 4); + @memset(k.dtb.slice(f32, W), 0); + k.onw = try buf(device, KD * 4); + @memset(k.onw.slice(f32, KD), 1); + k.fg = try buf(device, 2 * 16 * W * 2); + for (&k.cs, &k.st, &k.st_out) |*c, *s, *o| { + c.* = try buf(device, 3 * 3 * W * 2); + @memset(c.slice(u16, 3 * 3 * W), 0); + s.* = try buf(device, H * KD * KD * 4); + @memset(s.slice(f32, H * KD * KD), 0); + o.* = try buf(device, H * KD * KD * 4); + } + return k; + } +}; + +/// kda_rows (generated dots, or `.pre`: f_b, g_b made apart first) on launch i's state and gates. +fn kdaStep(e: mtl.ComputeEncoder, i: usize, pipe: mtl.Pipeline, s: *const Kda, gates: [2]Pool, pre: ?*const wide.Pipes, io: Io) void { + const f = gates[0].at(i); + const g = gates[1].at(i); + var fb = [3]Ref{ f.w, f.s, f.b }; + var gb = [3]Ref{ g.w, g.s, g.b }; + if (pre) |p| { + const fo = rf(s.fg); + const go = fo.at(16 * W * 2); + wide.dense(p, e, f, .{ .in = rf(s.p).at(3 * W * 2), .ldx = PS, .out = fo, .ldo = W, .rows = 1 }); + wide.dense(p, e, g, .{ .in = rf(s.p).at((3 * W + KD) * 2), .ldx = PS, .out = go, .ldo = W, .rows = 1 }); + fb = .{ fo, fo, fo }; + gb = .{ go, go, go }; + } + const c = i % 16; + e.setPipeline(pipe); + bind(e, 0, .{rf(s.p)}); + shape(e, 1, .{ 1, PS }); + bind(e, 2, .{ rf(s.cs[c]), rf(s.cw), fb[0], fb[1], fb[2], gb[0], gb[1], gb[2], rf(s.a), rf(s.dtb), rf(s.st[c]), rf(s.onw) }); + e.setValue(@as(f32, -5.0), 14); + e.setValue(@as(f32, 1e-5), 15); + bind(e, 16, .{ rf(io.y), rf(s.st_out[c]), rf(s.cs[(c + 1) % 16]) }); + e.dispatchThreads(mtl.Size.of(32, 32, H), mtl.Size.of(32, 32, 1)); +} + +const PM = 2048; // a prompt chunk's rows + +/// A prompt chunk's per-head product (embed_q: x qp's slices; unembed_out: the latent rows) on the row kernels. +fn headsRows(e: mtl.ComputeEncoder, i: usize, p: *const wide.Pipes, pool: Pool, io: Io, ldx: u32, kx: u32) void { + const q = pool.at(i); + wide.perHead(p, e, q, H, .{ .in = rf(io.x), .ldx = ldx, .out = rf(io.y), .ldo = q.n, .rows = PM }, kx, q.n / H); +} +/// The same as one affine_mm batch (prompt.zig headsMm). +fn headsMm(e: mtl.ComputeEncoder, i: usize, mm: tf.affine_mm.Pipes, pool: Pool, io: Io, ldx: u32, kx: u32) void { + var head = pool.at(i); + head.n /= H; + tf.affine_mm.rowSumsAt(e, mm, 64, rf(io.x), rf(io.z), PM, H * kx, ldx); + tf.affine_mm.denseBatch(e, mm, rf(io.x), rf(io.z), head, rf(io.y), PM, H, .{ .x_row = @intCast(ldx), .y_row = @intCast(H * head.n), .sums_row = H, .x_batch = @intCast(kx), .y_batch = @intCast(head.n), .sums_batch = 1, .w_batch = @intCast(head.n) }); +} +fn ehRows(e: mtl.ComputeEncoder, i: usize, p: *const wide.Pipes, eh: mtl.Buffer, io: Io) void { + _ = i; + wide.bf16(p, e, .{ .buf = eh }, D, 2 * D, .{ .in = rf(io.x), .ldx = 2 * D, .out = rf(io.y), .ldo = D, .rows = PM }); +} + +/// A chunk's dense product: row kernels, or affine_mm (`mm`). +fn chunkRows(e: mtl.ComputeEncoder, i: usize, p: *const wide.Pipes, pool: Pool, io: Io) void { + const q = pool.at(i); + wide.dense(p, e, q, .{ .in = rf(io.x), .ldx = q.k, .out = rf(io.y), .ldo = q.n, .rows = PM }); +} +fn chunkMm(e: mtl.ComputeEncoder, i: usize, mm: tf.affine_mm.Pipes, pool: Pool, io: Io) void { + const q = pool.at(i); + tf.affine_mm.rowSums(e, mm, 64, rf(io.x), rf(io.z), PM, q.k); + tf.affine_mm.dense(e, mm, rf(io.x), rf(io.z), q, rf(io.y), PM); +} + +/// A prompt chunk's products where oQ's paths differ from uniform's, ms a chunk of 2048 rows. +fn prompt(b: *Bench, k: *const kernels.Kernels, p: *const wide.Pipes) !void { + b.reps = 2; + defer b.reps = reps; + std.debug.print("prompt chunk ({d} rows), ms a launch:\n", .{PM}); + const eq = try Pool.init(b.device, H * 512, 256, 8, 1, 2); + defer eq.deinit(); + std.debug.print(" embed_q q8 per-head rows {d:8.2} affine_mm batch {d:8.2} (x11 a chunk)\n", .{ try b.time(headsRows, .{ p, eq, b.io, @as(u32, QR), @as(u32, 256) }) / 1e3, try b.time(headsMm, .{ p.mm8, eq, b.io, @as(u32, QR), @as(u32, 256) }) / 1e3 }); + const uo = try Pool.init(b.device, H * 256, 512, 8, 1, 2); + defer uo.deinit(); + std.debug.print(" unembed_out q8 per-head rows {d:8.2} affine_mm batch {d:8.2} (x11 a chunk)\n", .{ try b.time(headsRows, .{ p, uo, b.io, @as(u32, H * 512), @as(u32, 512) }) / 1e3, try b.time(headsMm, .{ p.mm8, uo, b.io, @as(u32, H * 512), @as(u32, 512) }) / 1e3 }); + { // the MTP head's bf16 eh_proj [4096, 8192] over the chunk (prompt.zig mtp) + const eh = try buf(b.device, D * 2 * D * 2); + defer eh.deinit(); + @memset(eh.slice(u16, 2 * D * D), 0x3c00); + std.debug.print(" eh_proj bf16 [4096, 8192] rows {d:8.2} (x1 a chunk with MTP)\n", .{try b.time(ehRows, .{ p, eh, b.io }) / 1e3}); + } + for ([_][3]usize{ .{ PS, D, 9 }, .{ D, W, 10 }, .{ 1536, D, 1 } }) |s| { + const m8 = try Pool.init(b.device, s[0], s[1], 8, 1, 2); + defer m8.deinit(); + const m4 = try Pool.init(b.device, s[0], s[1], 4, 1, 2); + defer m4.deinit(); + const t8 = try b.time(chunkMm, .{ p.mm8, m8, b.io }); + const t4 = try b.time(chunkMm, .{ k.mm_bf16, m4, b.io }); + for ([_]u8{ 5, 6 }) |bits| { + const m = try Pool.init(b.device, s[0], s[1], bits, 1, 2); + defer m.deinit(); + std.debug.print(" [{d}, {d}] q{d} rows {d:8.2} (affine_mm q8 {d:.2}, q4 {d:.2})\n", .{ s[0], s[1], bits, try b.time(chunkRows, .{ p, m, b.io }) / 1e3, t8 / 1e3, t4 / 1e3 }); + } + } +} + +fn cand(e: mtl.ComputeEncoder, i: usize, pipe: mtl.Pipeline, rps: u32, sg: u32, pool: Pool, io: Io) void { + const q = pool.at(i); + e.setPipeline(pipe); + bind(e, 0, .{ rf(io.x), q.w, q.s, q.b }); + e.setValue(wide.RowDims{ .k = @intCast(q.k), .n = @intCast(q.n), .group = 64, .rows = 1, .ldx = @intCast(q.k), .ldo = @intCast(q.n) }, 4); + bind(e, 5, .{rf(io.y)}); + e.dispatchGroups(mtl.Size.of((q.n + rps * sg - 1) / (rps * sg), 1, 1), mtl.Size.of(32 * sg, 1, 1)); +} + +/// glm_rows_speed.metal's launch shapes on `pool` at 8 bits: each one's bits == glm_rows's, then its time. +fn candidates(b: *const Bench, lib: mtl.Library, p: *const wide.Pipes, pool: Pool, name: []const u8) !void { + const q = pool.at(0); + const base = try b.time(rows, .{ p, pool, b.io }); + try b.run(rows, .{ p, pool, b.io }); + const want = try a.dupe(u16, b.io.y.slice(u16, q.n)); + defer a.free(want); + var best: f64 = base; + var best_name: [40]u8 = undefined; + var best_len: usize = 0; + for ([_][2]u32{ .{ 4, 2 }, .{ 8, 2 }, .{ 4, 4 }, .{ 8, 4 }, .{ 2, 4 }, .{ 2, 8 }, .{ 4, 8 } }) |rs| for ([_]u32{ 0, q.k }) |kk| { + var nb: [40]u8 = undefined; + const fname = try std.fmt.bufPrintSentinel(&nb, "x_q8_r{d}_s{d}_k{d}", .{ rs[0], rs[1], kk }, 0); + const pipe = try mtl.Pipeline.init(b.device, lib, fname, false); + defer pipe.deinit(); + @memset(b.io.y.slice(u16, q.n), 0); + try b.run(cand, .{ pipe, rs[0], rs[1], pool, b.io }); + if (!std.mem.eql(u16, b.io.y.slice(u16, q.n), want)) { + std.debug.print(" {s} {s}: BITS DIFFER\n", .{ name, fname }); + return error.BitsDiffer; + } + const us = try b.time(cand, .{ pipe, rs[0], rs[1], pool, b.io }); + if (us < best) { + best = us; + @memcpy(best_name[0..fname.len], fname); + best_len = fname.len; + } + std.debug.print(" {s:<16} {s:<24} {d:8.1} us ({d:5.3}x glm_rows)\n", .{ name, fname, us, base / us }); + }; + std.debug.print(" {s:<16} glm_rows {d:.1} us, best {s} {d:.1} us ({d:.0} GB/s)\n", .{ name, base, best_name[0..best_len], best, @as(f64, @floatFromInt(pool.bytes)) / best / 1e3 }); +} + +fn routedRows(e: mtl.ComputeEncoder, i: usize, p: *const wide.Pipes, gu: [2]Pool, down: Pool, io: Io) void { + const ids = rf(io.ids).at((i % reps) * TOPK * 4); + wide.byIds(p, e, .{ .w = gu[0].at(i), .up = gu[1].at(i), .experts = E, .repeat = TOPK, .lim = 10.0 }, rf(io.x), ids, rf(io.z), TOPK); + wide.byIds(p, e, .{ .w = down.at(i), .experts = E, .repeat = 1 }, rf(io.z), ids, rf(io.y), TOPK); +} + +/// A token's GPU time by class: launches a token times the median launch. +const Row = struct { name: []const u8, us: f64, count: f64, bytes: f64 }; +const Sheet = struct { + rows: [24]Row = undefined, + n: usize = 0, + + fn add(s: *Sheet, name: []const u8, us: f64, count: f64, bytes: usize) void { + s.rows[s.n] = .{ .name = name, .us = us, .count = count, .bytes = @floatFromInt(bytes) }; + s.n += 1; + std.debug.print(" {s:<26} {d:9.1} us {d:7.1} MB {d:6.1} GB/s x{d:<5.0} {d:7.3} ms/token\n", .{ name, us, @as(f64, @floatFromInt(bytes)) / 1e6, @as(f64, @floatFromInt(bytes)) / us / 1e3, count, us * count / 1e3 }); + } + + fn total(s: *const Sheet, label: []const u8) void { + var ms: f64 = 0; + var gb: f64 = 0; + for (s.rows[0..s.n]) |x| { + ms += x.us * x.count / 1e3; + gb += x.bytes * x.count / 1e9; + } + std.debug.print(" {s}: {d:.3} ms a token over the products, {d:.3} GB read ({d:.0} GB/s)\n", .{ label, ms, gb, gb / ms * 1e3 }); + } +}; + +pub fn main() !void { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + const k = try kernels.load(a, device); + defer k.deinit(); + try k.widen(a, device); + const p = &k.wide.?; + const io = Io{ .x = try buf(device, 160 << 20), .y = try buf(device, 160 << 20), .z = try buf(device, 64 << 20), .ids = try buf(device, reps * TOPK * 4), .uids = try buf(device, reps * 128 * 4), .umem = try buf(device, 128 * 16 * 4), .ucount = try buf(device, 4) }; + @memset(io.x.slice(u16, 32 << 20), 0x3c00); + for (0..reps) |i| for (0..TOPK) |j| { // launch i's picks: distinct experts, the pool walked in order + io.ids.slice(i32, reps * TOPK)[i * TOPK + j] = @intCast((i * TOPK + j) % E); + io.uids.slice(i32, reps * 128)[i * 128 + j] = @intCast((i * TOPK + j) % E); + }; + @memset(io.umem.slice(i32, 128 * 16), -1); + for (0..TOPK) |u| io.umem.slice(i32, 128 * 16)[u * 16] = @intCast(u); // unique id u: row 0's pick u + io.ucount.slice(i32, 1)[0] = TOPK; + var b = Bench{ .device = device, .queue = queue, .k = k, .io = io }; + std.debug.print("GPU {s}, tensor units {}; one row, {d} launches a command buffer, median of {d}\n", .{ device.name(), device.tensorUnits(), reps, runs }); + const shapes = [_]struct { name: []const u8, n: usize, k: usize, uni: ?mtl.Pipeline, count: [3]f64 }{ // counts: uniform, oQ2, oQ4 + .{ .name = "kda in", .n = PS, .k = D, .uni = k.qmv_kda_in, .count = .{ 34, 34, 34 } }, + .{ .name = "kda f_b/g_b", .n = W, .k = KD, .uni = null, .count = .{ 0, 0, 0 } }, // inside the KDA step's time + .{ .name = "kda out", .n = D, .k = W, .uni = k.qmv_kda_out, .count = .{ 34, 34, 34 } }, + .{ .name = "mla x_proj", .n = XP, .k = D, .uni = k.qmv_x, .count = .{ 11, 11, 11 } }, + .{ .name = "mla qr_proj", .n = QR, .k = 1536, .uni = k.qmv_qr, .count = .{ 11, 11, 11 } }, + .{ .name = "mla out", .n = D, .k = 16384, .uni = k.qmv_mla_out, .count = .{ 11, 11, 11 } }, + .{ .name = "dense gate|up", .n = 24576, .k = D, .uni = k.qmv_dense_gu, .count = .{ 3, 3, 3 } }, + .{ .name = "dense down", .n = D, .k = 12288, .uni = k.qmv_dense_down, .count = .{ 3, 3, 3 } }, + .{ .name = "lm_head", .n = 154880, .k = D, .uni = k.qmv_head, .count = .{ 1, 1, 1 } }, + }; + var uni = Sheet{}; + var oq2 = Sheet{}; + var oq4 = Sheet{}; + std.debug.print("generated 4/64 (uniform) and glm_rows at 8, 6, 5 bits:\n", .{}); + for (shapes) |s| { + if (s.uni) |pipe| { + const m = try Pool.init(device, s.n, s.k, 4, 1, 1); + defer m.deinit(); + uni.add(s.name, try b.time(qmv, .{ pipe, m, io }), s.count[0], m.bytes); + } + for ([_]u8{ 8, 6, 5 }) |bits| { + const m = try Pool.init(device, s.n, s.k, bits, 1, 1); + defer m.deinit(); + const us = try b.time(rows, .{ p, m, io }); + var name_buf: [40]u8 = undefined; + const name = try std.fmt.bufPrint(&name_buf, "{s} q{d}", .{ s.name, bits }); + if (bits == 8) { + oq2.add(try a.dupe(u8, name), us, s.count[1], m.bytes); + oq4.add(try a.dupe(u8, name), us, s.count[2], m.bytes); // oQ4's 5/6-bit members priced at 8 (an upper bound) + } else std.debug.print(" {s:<26} {d:9.1} us {d:7.1} MB {d:6.1} GB/s\n", .{ name, us, @as(f64, @floatFromInt(m.bytes)) / 1e6, @as(f64, @floatFromInt(m.bytes)) / us / 1e3 }); + } + } + { // the KDA step: generated 4/4 dots (uniform), generated 8/8 dots, f_b and g_b made apart at 8 and 5 bits (.pre) + const st = try Kda.init(device); + const g4 = [2]Pool{ try Pool.init(device, W, KD, 4, 1, 16), try Pool.init(device, W, KD, 4, 1, 16) }; + const g8 = [2]Pool{ try Pool.init(device, W, KD, 8, 1, 16), try Pool.init(device, W, KD, 8, 1, 16) }; + const g5 = [2]Pool{ try Pool.init(device, W, KD, 5, 1, 16), try Pool.init(device, W, KD, 5, 1, 16) }; + const step_bytes = 2 * H * KD * KD * 4; // the state read and written + uni.add("kda step (dots 4/4)", try b.time(kdaStep, .{ k.kda_rows, &st, g4, null, io }), 34, step_bytes + g4[0].bytes * 2); + const us_pre = try b.time(kdaStep, .{ p.kda_pre, &st, g8, p, io }); + oq2.add("kda step (.pre q8)", us_pre, 34, step_bytes + g8[0].bytes * 2); + oq4.add("kda step (.pre q8)", us_pre, 34, step_bytes + g8[0].bytes * 2); + const us_pre5 = try b.time(kdaStep, .{ p.kda_pre, &st, g5, p, io }); + std.debug.print(" {s:<26} {d:9.1} us\n", .{ "kda step (.pre q5)", us_pre5 }); + } + { // MLA absorb / unabsorb: kv_b's halves (uniform) or embed_q / unembed_out a head at 8 bits + const kv = try Pool.init(device, H * 512, 512, 4, 1, 1); + defer kv.deinit(); + uni.add("mla absorb", try b.time(absorb, .{ k, kv, io }), 11, kv.bytes / 2); + uni.add("mla unabsorb", try b.time(unabsorb, .{ k, kv, io }), 11, kv.bytes / 2); + const eq = try Pool.init(device, H * 512, 256, 8, 1, 1); + defer eq.deinit(); + const us_eq = try b.time(perHead, .{ p, eq, io, @as(u32, QR), @as(u32, 256) }); + oq2.add("mla embed_q q8", us_eq, 11, eq.bytes); + oq4.add("mla embed_q q8", us_eq, 11, eq.bytes); + const uo = try Pool.init(device, H * 256, 512, 8, 1, 1); + defer uo.deinit(); + const us_uo = try b.time(perHead, .{ p, uo, io, @as(u32, H * 512), @as(u32, 512) }); + oq2.add("mla unembed_out q8", us_uo, 11, uo.bytes); + oq4.add("mla unembed_out q8", us_uo, 11, uo.bytes); + } + { // MoE: the routed experts (8 picks) with the shared expert + const gu4 = [2]Pool{ try Pool.init(device, N, D, 4, E, 1), try Pool.init(device, N, D, 4, E, 1) }; + defer for (gu4) |x| x.deinit(); + const dn4 = try Pool.init(device, D, N, 4, E, 1); + defer dn4.deinit(); + const routed4 = 3 * affine.bytes(N, D, .{ .bits = 4, .group = 64 }) * TOPK; + const sh4 = [2]Pool{ try Pool.init(device, 2 * N, D, 4, 1, 2), try Pool.init(device, D, N, 4, 1, 2) }; + defer for (sh4) |x| x.deinit(); + const shared4 = sh4[0].bytes + sh4[1].bytes; + uni.add("moe routed+shared 4", try b.time(moeBoth, .{ k, gu4, dn4, sh4, io }), 42, routed4 + shared4); + oq4.add("moe routed 4", try b.time(moeRouted4, .{ k, gu4, dn4, io }), 42, routed4); + const sh8 = [2]Pool{ try Pool.init(device, 2 * N, D, 8, 1, 2), try Pool.init(device, D, N, 8, 1, 2) }; + defer for (sh8) |x| x.deinit(); + const us_sh = try b.time(shared, .{ p, k, sh8, io }); + oq2.add("moe shared q8", us_sh, 42, sh8[0].bytes + sh8[1].bytes); + oq4.add("moe shared q8", us_sh, 42, sh8[0].bytes + sh8[1].bytes); + const gu2 = [2]Pool{ try Pool.init(device, N, D, 2, E, 1), try Pool.init(device, N, D, 2, E, 1) }; + defer for (gu2) |x| x.deinit(); + const dn2 = try Pool.init(device, D, N, 2, E, 1); + defer dn2.deinit(); + oq2.add("moe routed q2 (ids)", try b.time(routedRows, .{ p, gu2, dn2, io }), 42, 3 * affine.bytes(N, D, .{ .bits = 2, .group = 64 }) * TOPK); + } + if (std.c.getenv("GLM_SPEED_CANDIDATES") != null) { + const src = try std.mem.concat(a, u8, &.{ sources.core_row_projection, wide.glue_math, sources.glm_rows, @embedFile("glm_rows_speed.metal") }); + const lib = try mtl.Library.fromSource(device, src, mtl.CompileOptions.mlx()); + defer lib.deinit(); + std.debug.print("candidates at 8 bits (bits checked == glm_rows):\n", .{}); + for (shapes) |s| { + const m = try Pool.init(device, s.n, s.k, 8, 1, 1); + defer m.deinit(); + try candidates(&b, lib, p, m, s.name); + } + } + if (std.c.getenv("GLM_SPEED_PROMPT") != null) try prompt(&b, k, p); + std.debug.print("uniform 4/64:\n", .{}); + for (uni.rows[0..uni.n]) |x| std.debug.print(" {s:<26} {d:7.3} ms/token\n", .{ x.name, x.us * x.count / 1e3 }); + uni.total("uniform 4/64"); + std.debug.print("oQ2:\n", .{}); + for (oq2.rows[0..oq2.n]) |x| std.debug.print(" {s:<26} {d:7.3} ms/token\n", .{ x.name, x.us * x.count / 1e3 }); + oq2.total("oQ2"); + std.debug.print("oQ4 (5/6-bit members priced at 8 bits):\n", .{}); + for (oq4.rows[0..oq4.n]) |x| std.debug.print(" {s:<26} {d:7.3} ms/token\n", .{ x.name, x.us * x.count / 1e3 }); + oq4.total("oQ4"); +} From 8715b25661a9ccd7308c017cd483ffb6f0fd5484 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 09:39:56 -1000 Subject: [PATCH 12/18] GLM-5.3-Flash prompt chunks run mlx-lm's embed_q and unembed_out on affine_mm, every head in one batch A prompt chunk took the decode's per-head row launch over every row for the absorbed MLA pair (each row a weight read, 64 heads x 2048 rows). At a width affine_mm reads (8/64 in oQ4 and oQ2), wide.headsBatch now runs the heads as one affine_mm batch over the strided rows (x's slices a head apart, sums one row apart), as the 4-bit path does with kv_b's value half. Other widths keep the row launch. Uniform 4/64 checkpoints never reach it. Test: tf-glm-rows-check's headsCheck: embed_q [4 x 512, 256] and unembed_out [4 x 256, 512] at 8 bits, 33 rows with a strided x, the batch == each head's own affine_mm product bit for bit (classes hold that product against fp64). --- zig/src/families/glm/prompt.zig | 6 +++--- zig/src/families/glm/wide.zig | 20 ++++++++++++++++++++ zig/tests/glm_rows_check.zig | 30 ++++++++++++++++++++++++++++++ 3 files changed, 53 insertions(+), 3 deletions(-) diff --git a/zig/src/families/glm/prompt.zig b/zig/src/families/glm/prompt.zig index df293dd0f4..e825a48dac 100644 --- a/zig/src/families/glm/prompt.zig +++ b/zig/src/families/glm/prompt.zig @@ -196,7 +196,7 @@ fn qmm(p: *const Prompt, e: mtl.ComputeEncoder, x: Ref, q: wts.Affine, y: Ref, M /// `qmm` over x's rows `ld` elements apart (a slice of a wider row). fn qmmAt(p: *const Prompt, e: mtl.ComputeEncoder, x: Ref, ld: u32, q: wts.Affine, y: Ref, M: u32) void { const pitch = std.mem.alignForward(u32, q.n, 64); - const mm = if (wide.native(q)) p.mm.mm_bf16 else if (q.more == 0 and q.bits == 8 and q.group == 64) p.mm.wide.?.mm8 else { + const mm = if (wide.native(q)) p.mm.mm_bf16 else wide.mm(&p.mm.wide.?, q) orelse { return wide.dense(&p.mm.wide.?, e, q, .{ .in = x, .ldx = ld, .out = y, .ldo = pitch, .rows = M }); // a split stack, other widths }; affine_mm.rowSumsAt(e, mm, 64, x, p.sums, M, q.k, ld); @@ -366,7 +366,7 @@ fn unabsorb(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, w: *cons const H = c.mla_heads; const K = c.kv_lora; const per = c.nope + c.v_dim; // kv_b's rows a head: its key half, then its value half - if (w.absorbed) |pair| return wide.perHead(&p.mm.wide.?, e, pair[1], H, .{ .in = p.att, .ldx = H * K, .out = p.vals, .ldo = H * c.v_dim, .rows = M }, K, c.v_dim); + if (w.absorbed) |pair| return wide.headsBatch(&p.mm.wide.?, e, pair[1], H, .{ .in = p.att, .ldx = H * K, .out = p.vals, .ldo = H * c.v_dim, .rows = M }, K, p.sums); var q = w.kv_b; q.w = q.w.at(@as(usize, c.nope) * K / 2); q.s = q.s.at(@as(usize, c.nope) * (K / 64) * 2); @@ -403,7 +403,7 @@ fn mla(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, mi: usize, w: } if (fwd.on(x, "mla_cache")) fwd.mlaCache(x, e, mi, w, x_in, p.xp, XP, p.iw, M, pos); if (fwd.on(x, "mla_absorb") and w.absorbed != null) { - fwd.absorb(x, e, w, p.qp, p.ql, M); // mlx-lm's embed_q a head at its width, the decode's launch over the chunk + wide.headsBatch(&p.mm.wide.?, e, w.absorbed.?[0], c.mla_heads, .{ .in = p.qp, .ldx = c.qrProj(), .out = p.ql, .ldo = c.mla_heads * c.kv_lora, .rows = M }, c.nope, p.sums); // mlx-lm's embed_q } else if (fwd.on(x, "mla_absorb")) { e.setPipeline(if (c.tp > 1) x.k.absorb_nax_tp else x.k.absorb_nax); // the row kernel's qvm, weights in three bf16 parts bind(e, 0, .{ w.kv_b.w, w.kv_b.s, w.kv_b.b, p.qp, p.ql }); diff --git a/zig/src/families/glm/wide.zig b/zig/src/families/glm/wide.zig index 9c9fc9dafe..9cda79085b 100644 --- a/zig/src/families/glm/wide.zig +++ b/zig/src/families/glm/wide.zig @@ -160,6 +160,16 @@ pub fn perHead(p: *const Pipes, e: mtl.ComputeEncoder, q: Affine, heads: u32, io launch(p, e, .{ q.w, q.s, q.b }, q.bits, .{ .k = @intCast(q.k), .n = @intCast(r), .group = q.group, .rows = @intCast(io.rows), .heads = @intCast(heads), .ldx = @intCast(io.ldx), .hsx = @intCast(hsx), .ldo = @intCast(io.ldo), .hso = @intCast(hso), .hsw = @intCast(r) }, io.in, io.out); } +/// A chunk's x slices (`kx` apart) through q's heads into out [rows, heads, n / heads]: one affine_mm batch, else rows. +pub fn headsBatch(p: *const Pipes, e: mtl.ComputeEncoder, q: Affine, h: u32, io: Io, kx: u32, sums: Ref) void { + const ny = q.n / h; + const pipes = mm(p, q) orelse return perHead(p, e, q, h, io, kx, ny); + var head = q; + head.n = ny; + affine_mm.rowSumsAt(e, pipes, 64, io.in, sums, io.rows, h * kx, io.ldx); // a head's sums one row apart + affine_mm.denseBatch(e, pipes, io.in, sums, head, io.out, io.rows, h, .{ .x_row = @intCast(io.ldx), .y_row = @intCast(io.ldo), .sums_row = @intCast(h), .x_batch = @intCast(kx), .y_batch = @intCast(ny), .sums_batch = 1, .w_batch = @intCast(ny) }); +} + /// x [rows, k] through bf16 W [n, k]. pub fn bf16(p: *const Pipes, e: mtl.ComputeEncoder, w: Ref, n: u32, k: u32, io: Io) void { e.setPipeline(p.bf16); @@ -196,6 +206,16 @@ pub fn bySorted(p: *const Pipes, e: mtl.ComputeEncoder, m: Picks, x: Ref, offset e.dispatchGroups(mtl.Size.of((m.w.n + 7) / 8, m.experts, 1), mtl.Size.of(64, 1, 1)); } +/// A prompt chunk's affine_mm pipelines for `q` (one width, a format wide.zig compiles), else null: the row kernels. +pub fn mm(p: *const Pipes, q: Affine) ?affine_mm.Pipes { + if (q.more != 0 or q.group != 64) return null; + return switch (q.bits) { + 8 => p.mm8, + 2 => p.mm2, + else => null, + }; +} + /// A prompt chunk's routed experts: affine_mm gathers at 4/64 (`mm4`) or 2/64 (`mm2`, oQ2's), else sorted row products. pub const Routed = enum { mm4, mm2, rows }; diff --git a/zig/tests/glm_rows_check.zig b/zig/tests/glm_rows_check.zig index 1f1858ca32..faf07b4da4 100644 --- a/zig/tests/glm_rows_check.zig +++ b/zig/tests/glm_rows_check.zig @@ -395,6 +395,32 @@ fn splitCheck(g: *const Gpu, k: usize, ns: []const usize, bits: []const u8, t: * } } +fn headsLaunch(e: mtl.ComputeEncoder, p: *const wide.Pipes, q: affine.Affine, h: u32, io: wide.Io, kx: u32, sums: Ref) void { + wide.headsBatch(p, e, q, h, io, kx, sums); +} + +/// A chunk's embed_q / unembed_out (wide.headsBatch): the heads in one affine_mm batch == each head's own product. +fn headsCheck(g: *const Gpu, mm: affine_mm.Pipes, bits: u8, n: usize, k: usize, t: *Tally) !void { + const nh = 4; + const m = 33; // two 16-row steps and one more + const ldx = nh * k + 96; // a row's slices, then columns no head reads (qp's indexer queries) + const w = try Mat.init(g.device, nh * n, k, bits, 64, 4100 + n + bits); + defer w.deinit(); + var prng = std.Random.DefaultPrng.init(4200 + k); + const r = prng.random(); + const bs = [_]mtl.Buffer{ try buf(g.device, m * ldx * 2), try buf(g.device, m * nh * (k / 64) * 4), try buf(g.device, m * nh * n * 2), try buf(g.device, m * k * 2), try buf(g.device, m * n * 2) }; + defer for (bs) |b| b.deinit(); + for (bs[0].slice(u16, m * ldx)) |*v| v.* = toBf((r.float(f32) - 0.5) * 2); + try g.run(headsLaunch, .{ &g.p, w.q, @as(u32, nh), wide.Io{ .in = ref(bs[0]), .ldx = @intCast(ldx), .out = ref(bs[2]), .ldo = @intCast(nh * n), .rows = m }, @as(u32, @intCast(k)), ref(bs[1]) }); + const all = bs[2].slice(u16, m * nh * n); + for (0..nh) |h| { + for (0..m) |i| @memcpy(bs[3].slice(u16, m * k)[i * k ..][0..k], bs[0].slice(u16, m * ldx)[i * ldx + h * k ..][0..k]); + try g.run(mmDense, .{ mm, w.part(h, nh), ref(bs[3]), ref(bs[1]), ref(bs[4]), @as(u32, m) }); + for (0..m) |i| try eql(all[(i * nh + h) * n ..][0..n], bs[4].slice(u16, m * n)[i * n ..][0..n]); + t.pairs += m * n; + } +} + fn perHead(e: mtl.ComputeEncoder, p: *const wide.Pipes, q: affine.Affine, heads: u32, io: wide.Io, hsx: u32, hso: u32) void { wide.perHead(p, e, q, heads, io, hsx, hso); } @@ -458,6 +484,10 @@ pub fn main() !void { try splitCheck(&g, 256, &.{ 64, 64, 64, 16, 16, 8 }, &.{ 5, 8, 5, 8, 8, 8 }, &split); // oQ4 layer 6's KDA input try splitCheck(&g, 512, &.{ 48, 16, 8, 8 }, &.{ 6, 6, 5, 5 }, &split); // oQ4 layer 7's and the MTP's MLA x_proj std.debug.print("split stacks through one call == each member alone: {d} bit-equal values\n", .{split.pairs}); + var hd = Tally{}; + for ([_][2]usize{ .{ 512, 256 }, .{ 256, 512 } }) |s| try headsCheck(&g, g.p.mm8, 8, s[0], s[1], &hd); // embed_q, unembed_out + std.debug.print("prompt embed_q / unembed_out q8: the heads in one affine_mm batch == each head alone: {d} bit-equal values\n", .{hd.pairs}); + all.pairs += hd.pairs; all.pairs += split.pairs; all.pairs += try @import("glm_rows_kda.zig").check(&g.p, device, queue); all.values += try @import("glm_rows_kda.zig").bf16Check(&g.p, device, queue, &all.worst); From 433704053e64b019be42ce223216b03c373d4693 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 09:40:36 -1000 Subject: [PATCH 13/18] The prompt product affine_mm reads 3-, 5- and 6-bit MLX codes, exactly A thread's half step (32 codes, 4 x bits bytes) loads as words and each code is cut from the little-endian bit stream and enters the tile as itself (offset 0, as at 2 and 8 bits; codes up to 63 are exact in bf16). source() admits every width 2-8 but 7. The 2-, 4- and 8-bit sources preprocess byte-identically to before (metal -E, both fragment layouts' defines, both output types), so their compiled kernels are unchanged. Tests: core/affine_mm.zig's dense-and-gathered test adds 3/64, 5/64, 6/64, 5/32 and 6/32 (codes read from the bit stream on the host; fails before: source() refused them); tf-affine-ref adds q3, q5, q6 at g32/g64 on [192, 1024] and [256, 4096] (own seeds, the earlier formats' receipts unchanged). --- zig/kernels/metal/core/affine_mm.metal | 31 ++++++++++++++++++++++++-- zig/src/core/affine_mm.zig | 21 ++++++++++------- zig/tests/affine_ref.zig | 16 ++++++++++--- 3 files changed, 55 insertions(+), 13 deletions(-) diff --git a/zig/kernels/metal/core/affine_mm.metal b/zig/kernels/metal/core/affine_mm.metal index 464e459fe4..612c7e936a 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 (2, 4 or 8), TF_GROUP (32 or 64), TF_OUT_T (bfloat, or float for partial sums another pass adds). +// Defines: TF_BITS (2-6 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,7 @@ 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; 8 bytes at 2 bits). +// A thread's 32 codes of one weight row in a step (16-byte loads; 8 bytes at 2 bits; words at 3, 5 and 6 bits). #if TF_BITS == 2 struct Codes { uint2 v; @@ -75,6 +75,19 @@ struct Codes { inline Codes am_load(const device uchar* w) { return Codes{*(const device uint2*)w}; } +#elif TF_BITS % 4 != 0 +struct Codes { + uint v[TF_BITS]; +}; + +inline Codes am_load(const device uchar* w) { + Codes c; + TF_UNROLL + for (short i = 0; i < TF_BITS; i++) { + c.v[i] = ((const device uint*)w)[i]; + } + return c; +} #else struct Codes { uint4 v[TF_BITS / 4]; @@ -105,6 +118,20 @@ inline void am_codes(thread const Codes& c, threadgroup bfloat* out) { 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 != 0 + TF_UNROLL + for (short i = 0; i < 8; i++) { // codes 4i .. 4i + 3 of the little-endian bit stream, as themselves + float f[4]; + TF_UNROLL + for (short j = 0; j < 4; j++) { + const short at = TF_BITS * (4 * i + j), word = at / 32, shift = at % 32; + uint v = c.v[word] >> shift; + if (shift + TF_BITS > 32) v |= c.v[word + 1] << (32 - shift); + f[j] = float(v & ((1u << TF_BITS) - 1u)); + } + o[i / 2][2 * (i % 2)] = as_type(ushort2(as_type(bfloat(f[0])), as_type(bfloat(f[1])))); + o[i / 2][2 * (i % 2) + 1] = as_type(ushort2(as_type(bfloat(f[2])), as_type(bfloat(f[3])))); + } #elif TF_BITS == 4 TF_UNROLL for (short i = 0; i < 4; i++) { diff --git a/zig/src/core/affine_mm.zig b/zig/src/core/affine_mm.zig index 4b5575a84a..25386372a8 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 (2, 4 or 8) and group (32 or 64), and the output's type. +/// The weights' bits (2, 3, 4, 5, 6 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 != 2 and f.bits != 4 and f.bits != 8) or (f.group != 32 and f.group != 64)) return error.UnsupportedFormat; + if (f.bits < 2 or f.bits > 8 or f.bits == 7 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 }); @@ -111,6 +111,15 @@ pub fn gatherGlu(e: mtl.ComputeEncoder, p: Pipes, x: anytype, sums: anytype, gat e.dispatchGroups(mtl.Size.of(gate.n / 32, tiles, 1), mtl.Size.of(32, 2, 2)); } +/// The code at bit `at` of a little-endian bit stream (MLX's packing at every width). +fn code(bytes: []const u8, at: usize, bits: usize) u32 { + var v: u32 = 0; + for (0..3) |i| if (at / 8 + i < bytes.len) { + v |= @as(u32, bytes[at / 8 + i]) << @intCast(8 * i); + }; + return (v >> @intCast(at % 8)) & ((@as(u32, 1) << @intCast(bits)) - 1); +} + fn bf16(v: f32) u16 { // round to nearest even const u: u32 = @bitCast(v); return @intCast((u + 0x7fff + ((u >> 16) & 1)) >> 16); @@ -135,7 +144,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 }, .{ 2, 64 }, .{ 2, 32 } }) |fmt| { + for ([_][2]usize{ .{ 4, 64 }, .{ 4, 32 }, .{ 8, 64 }, .{ 8, 32 }, .{ 2, 64 }, .{ 2, 32 }, .{ 3, 64 }, .{ 5, 64 }, .{ 6, 64 }, .{ 5, 32 }, .{ 6, 32 } }) |fmt| { const bits = fmt[0]; const group = fmt[1]; const w_bytes = E * N * K * bits / 8; @@ -207,11 +216,7 @@ 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 = 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 qv: u32 = code(wb, (row * K + k) * bits, bits); const xv: f64 = f32of(x[r * K + k]); dot += @as(f64, @floatFromInt(qv)) * xv; dmag += (@as(f64, @floatFromInt(qv)) + off) * @abs(xv); diff --git a/zig/tests/affine_ref.zig b/zig/tests/affine_ref.zig index 5f32bd864a..4849616083 100644 --- a/zig/tests/affine_ref.zig +++ b/zig/tests/affine_ref.zig @@ -57,9 +57,12 @@ const Synth = struct { 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; + const at = (j * self.k + i) * self.bits; // MLX's little-endian bit stream at every width + var v: u32 = 0; + for (0..3) |b| if (at / 8 + b < bytes.len) { + v |= @as(u32, bytes[at / 8 + b]) << @intCast(8 * b); + }; + return (v >> @intCast(at % 8)) & ((@as(u32, 1) << @intCast(self.bits)) - 1); } /// 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). @@ -340,6 +343,13 @@ pub fn main(_: std.process.Init) !void { 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 ([_]u8{ 3, 5, 6 }) |bits| for ([_]usize{ 32, 64 }) |group| { // oQ4's 5 and 6 + const s = try Synth.init(device, shape[0], shape[1], bits, group, 0x3716 + 100 * @as(usize, bits) + 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 From 2f5ffff06c2d1f085bab94ac64dfd7d258a5c010 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 09:45:06 -1000 Subject: [PATCH 14/18] GLM-5.3-Flash prompt chunks run 5- and 6-bit products and split-stack members on affine_mm oQ4's 5- and 6-bit matrices (KDA in, f_b/g_b, o_proj, MLA, the MTP layer) and its split stacks (layer 6's KDA input, layer 7's and the MTP's x_proj and qr_proj) took the row kernels in prompt chunks, each row a weight read. wide.zig now compiles affine_mm at 5/64 and 6/64 (mm5, mm6) and wide.denseMm runs each member of a matrix on affine_mm at its width where its rows fill 64-row tiles (a single run: its padded rows, as at 8 bits), into its rows of the chunk's output, the x group sums made once; members of other widths or rows (the 32-row indexer weights_proj) stay on the row kernels. Uniform 4/64 matrices keep their launches. Tests: tf-glm-rows-check's classes now take prompt.zig's choice (wide.mm), so every 5- and 6-bit class's prompt product is held against fp64 on affine_mm; splitMmCheck runs layer 6's KDA input (5/8/5/8/8/8, every member on affine_mm) and layer 7's x_proj (6/6/5 on affine_mm, the 32-row 5-bit member on rows) through wide.denseMm at the 64-padded pitch, each member == its own product bit for bit. --- zig/src/families/glm/prompt.zig | 5 ++-- zig/src/families/glm/wide.zig | 41 ++++++++++++++++++++++--------- zig/tests/glm_rows_check.zig | 43 +++++++++++++++++++++++++++++++++ zig/tests/glm_rows_classes.zig | 2 +- zig/tests/glm_rows_speed.zig | 3 ++- 5 files changed, 78 insertions(+), 16 deletions(-) diff --git a/zig/src/families/glm/prompt.zig b/zig/src/families/glm/prompt.zig index e825a48dac..4fed98c967 100644 --- a/zig/src/families/glm/prompt.zig +++ b/zig/src/families/glm/prompt.zig @@ -196,9 +196,8 @@ fn qmm(p: *const Prompt, e: mtl.ComputeEncoder, x: Ref, q: wts.Affine, y: Ref, M /// `qmm` over x's rows `ld` elements apart (a slice of a wider row). fn qmmAt(p: *const Prompt, e: mtl.ComputeEncoder, x: Ref, ld: u32, q: wts.Affine, y: Ref, M: u32) void { const pitch = std.mem.alignForward(u32, q.n, 64); - const mm = if (wide.native(q)) p.mm.mm_bf16 else wide.mm(&p.mm.wide.?, q) orelse { - return wide.dense(&p.mm.wide.?, e, q, .{ .in = x, .ldx = ld, .out = y, .ldo = pitch, .rows = M }); // a split stack, other widths - }; + if (!wide.native(q)) return wide.denseMm(&p.mm.wide.?, e, q, .{ .in = x, .ldx = ld, .out = y, .ldo = pitch, .rows = M }, p.sums); + const mm = p.mm.mm_bf16; affine_mm.rowSumsAt(e, mm, 64, x, p.sums, M, q.k, ld); affine_mm.denseBatch(e, mm, x, p.sums, q, y, M, 1, .{ .x_row = @intCast(ld), .y_row = @intCast(pitch) }); } diff --git a/zig/src/families/glm/wide.zig b/zig/src/families/glm/wide.zig index 9cda79085b..d56336af34 100644 --- a/zig/src/families/glm/wide.zig +++ b/zig/src/families/glm/wide.zig @@ -30,14 +30,15 @@ pub const Pipes = struct { embed: [2]mtl.Pipeline, // the embedding at 6 and 8 bits, groups of 64 mm8: affine_mm.Pipes, // prompt chunks' products of one width at 8 bits in groups of 64 (core/affine_mm.zig) mm2: affine_mm.Pipes, // and at 2 bits in groups of 64: oQ2's routed experts, gathered + mm5: affine_mm.Pipes, // and at 5 and 6 bits (oQ4's), split stacks' members too + mm6: affine_mm.Pipes, hash: u64, // every source compiled here pub fn deinit(p: *Pipes) void { inline for (.{ &p.rows, &p.ids, &p.glu_ids, &p.sorted, &p.glu_sorted, &p.embed }) |set| for (set) |x| x.deinit(); p.bf16.deinit(); p.kda_pre.deinit(); - for (p.mm8) |x| x.deinit(); - for (p.mm2) |x| x.deinit(); + inline for (.{ p.mm8, p.mm2, p.mm5, p.mm6 }) |set| for (set) |x| x.deinit(); } }; @@ -56,7 +57,7 @@ pub const glue_math = blk: { /// The pipelines built so far: a failed compile releases them all. const Made = struct { - list: [nw * 5 + 4 + 2 * affine_mm.names.len]mtl.Pipeline = undefined, + list: [nw * 5 + 4 + 4 * affine_mm.names.len]mtl.Pipeline = undefined, n: usize = 0, released: usize = 0, @@ -115,14 +116,12 @@ fn build(gpa: std.mem.Allocator, device: mtl.Device, kda: []const u8, made: *Mad try lib(made, device, pre_src, &.{comptime kernels.generatedFunction("kda_rows")}, @as(*[1]mtl.Pipeline, &p.kda_pre)); h.update(sources.ops_embed_norm); try lib(made, device, sources.ops_embed_norm, &.{ "tf_embed_b6_g64", "tf_embed_b8_g64" }, &p.embed); - const mm8 = try affine_mm.source(device, gpa, .{ .bits = 8, .group = 64 }); - defer gpa.free(mm8); - h.update(mm8); - try lib(made, device, mm8, &affine_mm.names, &p.mm8); - const mm2 = try affine_mm.source(device, gpa, .{ .bits = 2, .group = 64 }); - defer gpa.free(mm2); - h.update(mm2); - try lib(made, device, mm2, &affine_mm.names, &p.mm2); + for ([_]u8{ 8, 2, 5, 6 }, [_]*affine_mm.Pipes{ &p.mm8, &p.mm2, &p.mm5, &p.mm6 }) |bits, out| { + const src = try affine_mm.source(device, gpa, .{ .bits = bits, .group = 64 }); + defer gpa.free(src); + h.update(src); + try lib(made, device, src, &affine_mm.names, out); + } p.hash = h.final(); return p; } @@ -212,10 +211,30 @@ pub fn mm(p: *const Pipes, q: Affine) ?affine_mm.Pipes { return switch (q.bits) { 8 => p.mm8, 2 => p.mm2, + 5 => p.mm5, + 6 => p.mm6, else => null, }; } +/// A chunk's product into io.out: each member on affine_mm at its width if it fills 64-row tiles, else the row kernels. +pub fn denseMm(p: *const Pipes, e: mtl.ComputeEncoder, q: Affine, io: Io, sums: Ref) void { + var summed = false; + for (0..@as(usize, q.more) + 1) |i| { + const r = q.run(i); + const one: Affine = .{ .w = r.w, .s = r.s, .b = r.b, .n = r.n, .k = q.k, .bits = r.bits, .group = r.group }; + var part = io; + part.out = io.out.at(@as(usize, r.row) * 2); + const pipes = (if (q.more == 0 or r.n % 64 == 0) mm(p, one) else null) orelse { + dense(p, e, one, part); + continue; + }; + if (!summed) affine_mm.rowSumsAt(e, pipes, 64, io.in, sums, io.rows, q.k, io.ldx); // every member's: one k, g64 + summed = true; + affine_mm.denseBatch(e, pipes, io.in, sums, one, part.out, io.rows, 1, .{ .x_row = @intCast(io.ldx), .y_row = @intCast(io.ldo) }); + } +} + /// A prompt chunk's routed experts: affine_mm gathers at 4/64 (`mm4`) or 2/64 (`mm2`, oQ2's), else sorted row products. pub const Routed = enum { mm4, mm2, rows }; diff --git a/zig/tests/glm_rows_check.zig b/zig/tests/glm_rows_check.zig index faf07b4da4..27f4b28173 100644 --- a/zig/tests/glm_rows_check.zig +++ b/zig/tests/glm_rows_check.zig @@ -421,6 +421,44 @@ fn headsCheck(g: *const Gpu, mm: affine_mm.Pipes, bits: u8, n: usize, k: usize, } } +fn denseMm(e: mtl.ComputeEncoder, p: *const wide.Pipes, q: affine.Affine, io: wide.Io, sums: Ref) void { + wide.denseMm(p, e, q, io, sums); +} + +/// A chunk's split stack (wide.denseMm) at the 64-padded pitch == each member alone (affine_mm or the row kernels). +fn splitMmCheck(g: *const Gpu, k: usize, ns: []const usize, bits: []const u8, t: *Tally) !usize { + const m = 33; + var ms: [affine.max_runs]Mat = undefined; + for (ns, bits, 0..) |n, b, i| ms[i] = try Mat.init(g.device, n, k, b, 64, 900 + 37 * i + n); + defer for (ms[0..ns.len]) |x| x.deinit(); + var q = ms[0].q; + var at: u32 = @intCast(ns[0]); + for (ms[1..ns.len], 0..) |x, i| { + q.runs[i] = .{ .w = x.q.w, .s = x.q.s, .b = x.q.b, .row = at, .n = x.q.n, .bits = x.q.bits, .group = x.q.group }; + at += x.q.n; + } + q.more = @intCast(ns.len - 1); + q.n = at; + var prng = std.Random.DefaultPrng.init(9100 + k); + const r = prng.random(); + const ld = std.mem.alignForward(usize, at, 64); + const bs = [_]mtl.Buffer{ try buf(g.device, m * k * 2), try buf(g.device, m * (k / 64) * 4), try buf(g.device, m * ld * 2), try buf(g.device, m * ld * 2) }; + defer for (bs) |b| b.deinit(); + for (bs[0].slice(u16, m * k)) |*v| v.* = toBf((r.float(f32) - 0.5) * 2); + try g.run(denseMm, .{ &g.p, q, wide.Io{ .in = ref(bs[0]), .ldx = @intCast(k), .out = ref(bs[2]), .ldo = @intCast(ld), .rows = m }, ref(bs[1]) }); + var on_mm: usize = 0; + for (ms[0..ns.len], 0..) |x, i| { + const run = q.run(i); + if (wide.mm(&g.p, x.q) != null and run.n % 64 == 0) { + try g.run(mmDense, .{ wide.mm(&g.p, x.q).?, x.q, ref(bs[0]), ref(bs[1]), ref(bs[3]), @as(u32, m) }); + on_mm += 1; + } else try g.run(dense, .{ &g.p, x.q, wide.Io{ .in = ref(bs[0]), .ldx = @intCast(k), .out = ref(bs[3]), .ldo = run.n, .rows = m } }); + for (0..m) |y| try eql(bs[2].slice(u16, m * ld)[y * ld + run.row ..][0..run.n], bs[3].slice(u16, m * ld)[y * run.n ..][0..run.n]); + t.pairs += m * run.n; + } + return on_mm; +} + fn perHead(e: mtl.ComputeEncoder, p: *const wide.Pipes, q: affine.Affine, heads: u32, io: wide.Io, hsx: u32, hso: u32) void { wide.perHead(p, e, q, heads, io, hsx, hso); } @@ -484,6 +522,11 @@ pub fn main() !void { try splitCheck(&g, 256, &.{ 64, 64, 64, 16, 16, 8 }, &.{ 5, 8, 5, 8, 8, 8 }, &split); // oQ4 layer 6's KDA input try splitCheck(&g, 512, &.{ 48, 16, 8, 8 }, &.{ 6, 6, 5, 5 }, &split); // oQ4 layer 7's and the MTP's MLA x_proj std.debug.print("split stacks through one call == each member alone: {d} bit-equal values\n", .{split.pairs}); + var sm = Tally{}; + const kda_mm = try splitMmCheck(&g, 256, &.{ 128, 128, 128, 64, 64, 64 }, &.{ 5, 8, 5, 8, 8, 8 }, &sm); // oQ4 layer 6 + const x_mm = try splitMmCheck(&g, 512, &.{ 192, 64, 64, 32 }, &.{ 6, 6, 5, 5 }, &sm); // layer 7's x_proj: 32 on rows + std.debug.print("prompt split stacks == each member alone ({d} + {d} members on affine_mm): {d} bit-equal values\n", .{ kda_mm, x_mm, sm.pairs }); + all.pairs += sm.pairs; var hd = Tally{}; for ([_][2]usize{ .{ 512, 256 }, .{ 256, 512 } }) |s| try headsCheck(&g, g.p.mm8, 8, s[0], s[1], &hd); // embed_q, unembed_out std.debug.print("prompt embed_q / unembed_out q8: the heads in one affine_mm batch == each head alone: {d} bit-equal values\n", .{hd.pairs}); diff --git a/zig/tests/glm_rows_classes.zig b/zig/tests/glm_rows_classes.zig index 966df1d6c4..4b9f99d4f6 100644 --- a/zig/tests/glm_rows_classes.zig +++ b/zig/tests/glm_rows_classes.zig @@ -165,7 +165,7 @@ pub fn check(p: *const wide.Pipes, device: mtl.Device, queue: mtl.Queue) !Tally try against(q, w, s, b, x, bufs[4].slice(u16, rows * n), n, 0, &t, 1.01); // scales 1% off must be seen if (t.caught == 0) return error.ControlNotCaught; // prompt: the chunk's product at rows 0 and 32; on the row kernels every row's bits equal the decode's - const mm: ?affine_mm.Pipes = if (bits == 8) p.mm8 else if (bits == 2) p.mm2 else null; // prompt.zig's (2: routed) + const mm: ?affine_mm.Pipes = wide.mm(p, q); // prompt.zig's (2: the routed experts' gathers) const on_mm = mm != null; if (mm) |pipes| try submit(queue, mmLaunch, .{ pipes, q, bufs[3], bufs[6], bufs[5], prompt_rows }) else try submit(queue, rowsLaunch, .{ p, q, bufs[3], 0, bufs[5], prompt_rows }); t.worst = 0; diff --git a/zig/tests/glm_rows_speed.zig b/zig/tests/glm_rows_speed.zig index ab25dc27c8..be91017c21 100644 --- a/zig/tests/glm_rows_speed.zig +++ b/zig/tests/glm_rows_speed.zig @@ -334,7 +334,8 @@ fn prompt(b: *Bench, k: *const kernels.Kernels, p: *const wide.Pipes) !void { for ([_]u8{ 5, 6 }) |bits| { const m = try Pool.init(b.device, s[0], s[1], bits, 1, 2); defer m.deinit(); - std.debug.print(" [{d}, {d}] q{d} rows {d:8.2} (affine_mm q8 {d:.2}, q4 {d:.2})\n", .{ s[0], s[1], bits, try b.time(chunkRows, .{ p, m, b.io }) / 1e3, t8 / 1e3, t4 / 1e3 }); + const on_mm = if (wide.mm(p, m.at(0))) |mm| try b.time(chunkMm, .{ mm, m, b.io }) else 0; + std.debug.print(" [{d}, {d}] q{d} rows {d:8.2} affine_mm {d:8.2} (affine_mm q8 {d:.2}, q4 {d:.2})\n", .{ s[0], s[1], bits, try b.time(chunkRows, .{ p, m, b.io }) / 1e3, on_mm / 1e3, t8 / 1e3, t4 / 1e3 }); } } } From 3d8bb7eda0385849cc9ec6e1af0e00ab33e9c012 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 09:52:10 -1000 Subject: [PATCH 15/18] GLM-5.3-Flash decode makes f_b's and g_b's dots inside kda_rows at one width, the same bits as made apart f_b and g_b at a width other than 4/64 (8/64 in oQ2 and 26 of oQ4's KDA layers, 5/64 in 8) were two glm_rows launches a KDA layer, written to scratch and read back by kda_rows (`.pre`). When both share one width in groups of 64, a kda_rows instance now makes them in its own threadgroup: each simdgroup sums four of the head's f_b or g_b rows with core tf_rp_rows (glm_rows's sums, lanes and simd_sum) and rounds them to bf16 as glm_rows's output does, so f and g are .pre's values bit for bit and so is everything after. Two launches and a round trip a KDA layer fewer (68 a decode token). Other widths keep .pre; uniform 4/64 keeps the generated dots. Test: tf-glm-rows-check's fusedCheck: at 2, 3, 4, 5, 6 and 8 bits in groups of 64, every code, random scales, biases, inputs, state and conv window, 3 rows: outputs, state and conv window byte-equal to .pre's (fails before: no fused pipeline). wide.zig's unit test: one width g64 -> .fused; 4|8, 8|4 and g32 -> .pre. --- zig/src/families/glm/forward.zig | 2 +- zig/src/families/glm/wide.zig | 60 +++++++++++++++++++++++++++---- zig/tests/glm_rows_check.zig | 1 + zig/tests/glm_rows_kda.zig | 61 ++++++++++++++++++++++++++++++++ zig/tests/glm_rows_speed.zig | 9 +++-- 5 files changed, 124 insertions(+), 9 deletions(-) diff --git a/zig/src/families/glm/forward.zig b/zig/src/families/glm/forward.zig index a9104bfd4e..330d9324e6 100644 --- a/zig/src/families/glm/forward.zig +++ b/zig/src/families/glm/forward.zig @@ -222,7 +222,7 @@ pub fn kdaStep(x: *const Ctx, e: mtl.ComputeEncoder, ki: usize, w: *const wts.Kd } e.setPipeline(switch (gate) { .generated => if (c.tp > 1) x.k.kda_rows_tp else x.k.kda_rows, - .pre => |p| p, + .pre, .fused => |p| p, }); bind(e, 0, .{proj}); shape(e, 1, .{ rows, c.kdaProj() }); diff --git a/zig/src/families/glm/wide.zig b/zig/src/families/glm/wide.zig index d56336af34..560a168375 100644 --- a/zig/src/families/glm/wide.zig +++ b/zig/src/families/glm/wide.zig @@ -27,6 +27,7 @@ pub const Pipes = struct { glu_sorted: [nw]mtl.Pipeline, bf16: mtl.Pipeline, // an unquantized matrix (the MTP head's eh_proj) kda_pre: mtl.Pipeline, // kda_rows reading f_b's and g_b's products made apart (any width but 4/64) + kda_fg: [nw]mtl.Pipeline, // kda_rows making f_b's and g_b's dots of one width (g64) as glm_rows does: .pre's bits embed: [2]mtl.Pipeline, // the embedding at 6 and 8 bits, groups of 64 mm8: affine_mm.Pipes, // prompt chunks' products of one width at 8 bits in groups of 64 (core/affine_mm.zig) mm2: affine_mm.Pipes, // and at 2 bits in groups of 64: oQ2's routed experts, gathered @@ -35,7 +36,7 @@ pub const Pipes = struct { hash: u64, // every source compiled here pub fn deinit(p: *Pipes) void { - inline for (.{ &p.rows, &p.ids, &p.glu_ids, &p.sorted, &p.glu_sorted, &p.embed }) |set| for (set) |x| x.deinit(); + inline for (.{ &p.rows, &p.ids, &p.glu_ids, &p.sorted, &p.glu_sorted, &p.embed, &p.kda_fg }) |set| for (set) |x| x.deinit(); p.bf16.deinit(); p.kda_pre.deinit(); inline for (.{ p.mm8, p.mm2, p.mm5, p.mm6 }) |set| for (set) |x| x.deinit(); @@ -46,6 +47,28 @@ pub const Pipes = struct { const pre_f = .{ "result = quad_dot(x, wb, s, bb);", "result = qlid == 0u ? float(((device const bfloat*)FBW)[(size_t)r * W + row]) : 0.0f;" }; const pre_g = .{ "result = quad_dot(x, wb, s, bb);", "result = qlid == 0u ? float(((device const bfloat*)GBW)[(size_t)r * W + row]) : 0.0f;" }; +/// kda_rows's f_b and g_b dots by quads of lanes, cut (from its block's start to the barrier after it) for `fg_block`. +const fg_cut = .{ " {\n constexpr int PER = D / 4;", "\n threadgroup_barrier(mem_flags::mem_threadgroup);\n for (uint idx = tid; idx < 3u * (uint)D;" }; +/// The head's f_b and g_b rows as glm_rows sums them (core tf_rp_rows, 4 rows a simdgroup, then bf16): .pre's values. +const fg_block = + \\ for (uint t0 = simdgroup_index_in_threadgroup * 4u; t0 < 2u * (uint)D; t0 += (uint)TY * 4u) { + \\ const bool f = t0 < (uint)D; + \\ const uint d0 = f ? t0 : t0 - (uint)D; + \\ const size_t row = (size_t)h * D + d0; + \\ float acc[1][4]; + \\ if (f) tf_rp_rows<4, 1, FB, false>((const device uint8_t*)FBW + row * (D * FB / 8), FBS + row * (D / 64), FBB + row * (D / 64), prow + FA, D, 64, lane, acc); + \\ else tf_rp_rows<4, 1, GB, false>((const device uint8_t*)GBW + row * (D * GB / 8), GBS + row * (D / 64), GBB + row * (D / 64), prow + GA, D, 64, lane, acc); + \\ if (lane == 0u) + \\ for (int j = 0; j < 4; j++) (f ? sa : sgate)[d0 + j] = float(bfloat(acc[0][j])); + \\ } +; +/// core/row_projection.metal's row sums only (its kernels and instances cut). +const rp_sums = blk: { + @setEvalBranchQuota(100_000); + const r = sources.core_row_projection; + break :blk r[0..std.mem.indexOf(u8, r, "// A value out: bf16").?]; +}; + /// glm_glue.metal's bfloat overloads and sigmoid_fast, the text glm_swiglu compiles: glm_rows's SwiGLU takes its bits. pub const glue_math = blk: { @setEvalBranchQuota(100_000); @@ -57,7 +80,7 @@ pub const glue_math = blk: { /// The pipelines built so far: a failed compile releases them all. const Made = struct { - list: [nw * 5 + 4 + 4 * affine_mm.names.len]mtl.Pipeline = undefined, + list: [nw * 6 + 4 + 4 * affine_mm.names.len]mtl.Pipeline = undefined, n: usize = 0, released: usize = 0, @@ -114,6 +137,13 @@ fn build(gpa: std.mem.Allocator, device: mtl.Device, kda: []const u8, made: *Mad if (std.mem.indexOf(u8, pre_src, pre_f[0]) != null or std.mem.indexOf(u8, pre_src, pre_g[0]) != null or std.mem.indexOf(u8, pre_src, pre_f[1]) == null) return error.KernelCompile; h.update(pre_src); try lib(made, device, pre_src, &.{comptime kernels.generatedFunction("kda_rows")}, @as(*[1]mtl.Pipeline, &p.kda_pre)); + const a = std.mem.indexOf(u8, kda, fg_cut[0]) orelse return error.KernelCompile; + const z = a + (std.mem.indexOf(u8, kda[a..], fg_cut[1]) orelse return error.KernelCompile); + const own = z + (std.mem.lastIndexOf(u8, kda[z..], "\ntemplate [[host_name(") orelse return error.KernelCompile); + const fg_src = try std.mem.concat(gpa, u8, &.{ rp_sums, kda[0..a], fg_block, kda[z..own], fg_instances }); // not its own + defer gpa.free(fg_src); + h.update(fg_src); + try lib(made, device, fg_src, &fg_names, &p.kda_fg); h.update(sources.ops_embed_norm); try lib(made, device, sources.ops_embed_norm, &.{ "tf_embed_b6_g64", "tf_embed_b8_g64" }, &p.embed); for ([_]u8{ 8, 2, 5, 6 }, [_]*affine_mm.Pipes{ &p.mm8, &p.mm2, &p.mm5, &p.mm6 }) |bits, out| { @@ -126,6 +156,18 @@ fn build(gpa: std.mem.Allocator, device: mtl.Device, kda: []const u8, made: *Mad return p; } +/// kda_rows with f_b and g_b at one width a pipeline (`fg_block`), under the names glm_kda_fg_q. +const fg_names = blk: { + var n: [nw][:0]const u8 = undefined; + for (bit_widths, 0..) |b, i| n[i] = std.fmt.comptimePrint("glm_kda_fg_q{d}", .{b}); + break :blk n; +}; +const fg_instances = blk: { + var s: []const u8 = ""; + for (bit_widths, fg_names) |b, name| s = s ++ kernels.instance("kda_rows", std.fmt.comptimePrint("64, 128, 4, 32, {d}, {d}", .{ b, b }), name); + break :blk s; +}; + fn bind(e: mtl.ComputeEncoder, first: usize, refs: anytype) void { inline for (refs, 0..) |r, i| e.setBuffer(r.buf, r.off, first + i); } @@ -244,11 +286,12 @@ pub fn routed(gate: Affine, up: Affine, down: Affine) Routed { return .mm2; } -/// kda_rows's f_b/g_b dots: generated (both 4/64), else (`.pre`) glm_rows products made apart, the fp64-checked ones. -pub const KdaGate = union(enum) { generated, pre: mtl.Pipeline }; +/// f_b/g_b dots: generated (both 4/64), `.fused` (one width, g64, glm_rows's sums inside), else `.pre`: same bits. +pub const KdaGate = union(enum) { generated, fused: mtl.Pipeline, pre: mtl.Pipeline }; pub fn kdaGate(p: ?*const Pipes, f: Affine, g: Affine) KdaGate { if (native(f) and native(g)) return .generated; + if (f.more == 0 and g.more == 0 and f.group == 64 and g.group == 64 and f.bits == g.bits) return .{ .fused = p.?.kda_fg[at(f.bits)] }; return .{ .pre = p.?.kda_pre }; } @@ -303,11 +346,16 @@ test "prompt chunks gather routed experts on affine_mm at 4/64 and 2/64 only; ot try std.testing.expectEqual(Routed.rows, routed(g32, g32, g32)); } -test "f_b and g_b at any other width (8/64 too) take the products made apart, not the template's own dots" { +test "f_b and g_b of one width (g64) take glm_rows's dots inside kda_rows, two widths the products made apart" { const r: Ref = .{ .buf = undefined }; const q: Affine = .{ .w = r, .s = r, .b = r, .n = 64, .k = 128, .bits = 8, .group = 64 }; var four = q; four.bits = 4; + var g32 = q; + g32.group = 32; const p: Pipes = undefined; - for ([_][2]Affine{ .{ q, q }, .{ four, q }, .{ q, four } }) |fg| try std.testing.expect(kdaGate(&p, fg[0], fg[1]) == .pre); + for ([_][2]Affine{ .{ four, q }, .{ q, four }, .{ g32, g32 } }) |fg| try std.testing.expect(kdaGate(&p, fg[0], fg[1]) == .pre); + var five = q; + five.bits = 5; + for ([_]Affine{ q, five }) |x| try std.testing.expect(kdaGate(&p, x, x) == .fused); } diff --git a/zig/tests/glm_rows_check.zig b/zig/tests/glm_rows_check.zig index 27f4b28173..f5d4ae89c2 100644 --- a/zig/tests/glm_rows_check.zig +++ b/zig/tests/glm_rows_check.zig @@ -533,6 +533,7 @@ pub fn main() !void { all.pairs += hd.pairs; all.pairs += split.pairs; all.pairs += try @import("glm_rows_kda.zig").check(&g.p, device, queue); + all.pairs += try @import("glm_rows_kda.zig").fusedCheck(&g.p, device, queue); all.values += try @import("glm_rows_kda.zig").bf16Check(&g.p, device, queue, &all.worst); all.values += try @import("glm_rows_kda.zig").embedCheck(&g.p, device, queue, &all.worst); const cl = try @import("glm_rows_classes.zig").check(&g.p, device, queue); diff --git a/zig/tests/glm_rows_kda.zig b/zig/tests/glm_rows_kda.zig index d3d307a81d..70294fe562 100644 --- a/zig/tests/glm_rows_kda.zig +++ b/zig/tests/glm_rows_kda.zig @@ -133,6 +133,67 @@ pub fn check(p: *const wide.Pipes, device: mtl.Device, queue: mtl.Queue) !usize return R * W; } +/// kda_rows with f_b, g_b made inside (`.fused`) == made apart (`.pre`) at every width, g64: Y, state, conv equal. +pub fn fusedCheck(p: *const wide.Pipes, device: mtl.Device, queue: mtl.Queue) !usize { + var prng = std.Random.DefaultPrng.init(4343); + const r = prng.random(); + const mk = struct { + fn of(dev: mtl.Device, bytes: usize) !mtl.Buffer { + return dev.buffer(bytes, opts); + } + }.of; + var b: Bufs = undefined; + b.p = try mk(device, R * PS * 2); + b.cs = try mk(device, (TAPS - 1) * 3 * W * 2); + b.cw = try mk(device, TAPS * 3 * W * 4); + b.a = try mk(device, H * 4); + b.dtb = try mk(device, W * 4); + b.st = try mk(device, H * D * D * 4); + b.onw = try mk(device, D * 4); + const f = try mk(device, 16 * W * 2); + const g = try mk(device, 16 * W * 2); + defer for ([_]mtl.Buffer{ b.p, b.cs, b.cw, b.a, b.dtb, b.st, b.onw, f, g }) |x| x.deinit(); + for (b.p.slice(u16, R * PS)) |*v| v.* = toBf((r.float(f32) - 0.5) * 2); + for (b.cs.slice(u16, (TAPS - 1) * 3 * W)) |*v| v.* = toBf((r.float(f32) - 0.5) * 2); + for (b.cw.slice(f32, TAPS * 3 * W)) |*v| v.* = (r.float(f32) - 0.5) * 0.5; + for (b.a.slice(f32, H)) |*v| v.* = -r.float(f32); + for (b.dtb.slice(f32, W)) |*v| v.* = (r.float(f32) - 0.5) * 0.1; + for (b.st.slice(f32, H * D * D)) |*v| v.* = (r.float(f32) - 0.5) * 0.1; + for (b.onw.slice(f32, D)) |*v| v.* = 1 + (r.float(f32) - 0.5) * 0.1; + var outs: [2]Outs = undefined; + for (&outs) |*o| o.* = .{ .y = try mk(device, R * W * 2), .st = try mk(device, H * D * D * 4), .cs = try mk(device, (TAPS - 1) * 3 * W * 2) }; + defer for (outs) |o| for ([_]mtl.Buffer{ o.y, o.st, o.cs }) |x| x.deinit(); + var same: usize = 0; + for (affine.row_bits, 0..) |bits, i| { + var wq: [2]affine.Affine = undefined; + var wb: [2][3]mtl.Buffer = undefined; + for (&wq, &wb) |*q, *x| { + x.* = .{ try mk(device, W * D * @as(usize, bits) / 8), try mk(device, W * 2 * 2), try mk(device, W * 2 * 2) }; + r.bytes(x[0].slice(u8, W * D * @as(usize, bits) / 8)); // every code + for (x[1].slice(u16, W * 2)) |*v| v.* = toBf(r.float(f32) * 0.02 + 0.001); + for (x[2].slice(u16, W * 2)) |*v| v.* = toBf((r.float(f32) - 0.5) * 0.1); + q.* = .{ .w = .{ .buf = x[0] }, .s = .{ .buf = x[1] }, .b = .{ .buf = x[2] }, .n = W, .k = D, .bits = bits, .group = 64 }; + } + defer for (wb) |x| for (x) |y| y.deinit(); + if (bits != 4 and wide.kdaGate(p, wq[0], wq[1]) != .fused) return error.NotFused; // 4/64 takes the generated dots + try submit(queue, products, .{ p, wq[0], wq[1], b.p, f, g }); + try submit(queue, step, .{ p.kda_pre, b, [3]mtl.Buffer{ f, f, f }, [3]mtl.Buffer{ g, g, g }, outs[0] }); + try submit(queue, step, .{ p.kda_fg[i], b, wb[0], wb[1], outs[1] }); + for ([_][2][]const u8{ .{ std.mem.sliceAsBytes(outs[0].y.slice(u16, R * W)), std.mem.sliceAsBytes(outs[1].y.slice(u16, R * W)) }, .{ outs[0].st.slice(u8, H * D * D * 4), outs[1].st.slice(u8, H * D * D * 4) }, .{ outs[0].cs.slice(u8, (TAPS - 1) * 3 * W * 2), outs[1].cs.slice(u8, (TAPS - 1) * 3 * W * 2) } }) |pair| { + if (!std.mem.eql(u8, pair[0], pair[1])) { + std.debug.print("kda_rows fused q{d} differs from .pre\n", .{bits}); + return error.KdaFusedDiffers; + } + same += pair[0].len; + } + var nonzero: usize = 0; + for (outs[1].y.slice(u16, R * W)) |v| nonzero += @intFromBool(v & 0x7fff != 0); + if (nonzero < R * W / 2) return error.KdaOutputsMostlyZero; + } + std.debug.print("kda_rows with f_b, g_b made inside == made apart (.pre) at q2..q8 g64: {d} rows, outputs, state and conv window byte-equal ({d} bytes)\n", .{ R, same }); + return same / 4; +} + /// The bf16 product (an unquantized eh_proj) against fp64 per column, and a row alone == the row in a window. pub fn bf16Check(p: *const wide.Pipes, device: mtl.Device, queue: mtl.Queue, worst: *f64) !usize { const n = 40; diff --git a/zig/tests/glm_rows_speed.zig b/zig/tests/glm_rows_speed.zig index be91017c21..75b1a53eea 100644 --- a/zig/tests/glm_rows_speed.zig +++ b/zig/tests/glm_rows_speed.zig @@ -473,10 +473,15 @@ pub fn main() !void { const step_bytes = 2 * H * KD * KD * 4; // the state read and written uni.add("kda step (dots 4/4)", try b.time(kdaStep, .{ k.kda_rows, &st, g4, null, io }), 34, step_bytes + g4[0].bytes * 2); const us_pre = try b.time(kdaStep, .{ p.kda_pre, &st, g8, p, io }); - oq2.add("kda step (.pre q8)", us_pre, 34, step_bytes + g8[0].bytes * 2); - oq4.add("kda step (.pre q8)", us_pre, 34, step_bytes + g8[0].bytes * 2); + std.debug.print(" {s:<26} {d:9.1} us\n", .{ "kda step (.pre q8)", us_pre }); const us_pre5 = try b.time(kdaStep, .{ p.kda_pre, &st, g5, p, io }); std.debug.print(" {s:<26} {d:9.1} us\n", .{ "kda step (.pre q5)", us_pre5 }); + const f8 = wide.kdaGate(p, g8[0].at(0), g8[1].at(0)).fused; + const us_f8 = try b.time(kdaStep, .{ f8, &st, g8, null, io }); + oq2.add("kda step (fused q8)", us_f8, 34, step_bytes + g8[0].bytes * 2); + oq4.add("kda step (fused q8)", us_f8, 34, step_bytes + g8[0].bytes * 2); + const f5 = wide.kdaGate(p, g5[0].at(0), g5[1].at(0)).fused; + std.debug.print(" {s:<26} {d:9.1} us\n", .{ "kda step (fused q5)", try b.time(kdaStep, .{ f5, &st, g5, null, io }) }); } { // MLA absorb / unabsorb: kv_b's halves (uniform) or embed_q / unembed_out a head at 8 bits const kv = try Pool.init(device, H * 512, 512, 4, 1, 1); From b6e156750a670cf60c4d025614177b381380d5e3 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 09:53:50 -1000 Subject: [PATCH 16/18] GLM-5.3-Flash decode runs the shared expert's gate, up and SwiGLU in one launch at its width, the same bits The shared expert at a width other than 4/64 (8/64 in oQ4 and oQ2) took a glm_rows launch for [gate; up] and a glm_swiglu launch every MoE layer. Gate and up of one width are one run of the stack, so wide.glu now runs them through the routed experts' by-ids GLU kernel with every id 0 (expert 0 of one, up N rows on): its sums are glm_rows's and its SwiGLU glm_swiglu's, so acts are the same bits. One launch a MoE layer fewer (42 a decode token). A split [gate; up] keeps the two launches; uniform 4/64 keeps the shared slot of the routed gather. Test: tf-glm-rows-check's sharedGluCheck: [2 x 2048, 4096] at 8, 5, 4 and 2 bits, 16 rows: wide.glu == glm_swiglu(dense [gate; up]) bit for bit (fails before: no wide.glu). --- zig/src/families/glm/moe.zig | 4 +++- zig/src/families/glm/wide.zig | 24 ++++++++++++++++++++++++ zig/tests/glm_rows_check.zig | 24 ++++++++++++++++++++++++ zig/tests/glm_rows_speed.zig | 8 ++------ 4 files changed, 53 insertions(+), 7 deletions(-) diff --git a/zig/src/families/glm/moe.zig b/zig/src/families/glm/moe.zig index 458a2ae444..0e3d84b293 100644 --- a/zig/src/families/glm/moe.zig +++ b/zig/src/families/glm/moe.zig @@ -197,7 +197,9 @@ fn sharedRows(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref const sc = x.sc; const N = c.moe_inter; const p = &x.k.wide.?; - if (on(x, "s_gateup")) { + if (on(x, "s_gateup") and w.sh_gate_up.more == 0) { // gate and up one width: one launch, glm_swiglu's bits + wide.glu(p, e, w.sh_gate_up, N, x_in, sc.acts, rows, c.swiglu_limit); + } else if (on(x, "s_gateup")) { wide.dense(p, e, w.sh_gate_up, .{ .in = x_in, .ldx = c.hidden, .out = sc.gu, .ldo = 2 * N, .rows = rows }); e.setPipeline(x.k.swiglu); bind(e, 0, .{ sc.gu, sc.acts }); diff --git a/zig/src/families/glm/wide.zig b/zig/src/families/glm/wide.zig index 560a168375..7e40bf3da0 100644 --- a/zig/src/families/glm/wide.zig +++ b/zig/src/families/glm/wide.zig @@ -33,12 +33,14 @@ pub const Pipes = struct { mm2: affine_mm.Pipes, // and at 2 bits in groups of 64: oQ2's routed experts, gathered mm5: affine_mm.Pipes, // and at 5 and 6 bits (oQ4's), split stacks' members too mm6: affine_mm.Pipes, + zeros: mtl.Buffer, // expert ids all 0: a dense GLU through the by-ids kernel (the shared expert) hash: u64, // every source compiled here pub fn deinit(p: *Pipes) void { inline for (.{ &p.rows, &p.ids, &p.glu_ids, &p.sorted, &p.glu_sorted, &p.embed, &p.kda_fg }) |set| for (set) |x| x.deinit(); p.bf16.deinit(); p.kda_pre.deinit(); + p.zeros.deinit(); inline for (.{ p.mm8, p.mm2, p.mm5, p.mm6 }) |set| for (set) |x| x.deinit(); } }; @@ -152,6 +154,8 @@ fn build(gpa: std.mem.Allocator, device: mtl.Device, kda: []const u8, made: *Mad h.update(src); try lib(made, device, src, &affine_mm.names, out); } + p.zeros = try device.buffer(64 * 4, mtl.ResourceOptions.shared); + @memset(p.zeros.slice(i32, 64), 0); p.hash = h.final(); return p; } @@ -172,6 +176,14 @@ fn bind(e: mtl.ComputeEncoder, first: usize, refs: anytype) void { inline for (refs, 0..) |r, i| e.setBuffer(r.buf, r.off, first + i); } +fn wBytes(n: u32, k: u32, q: Affine) usize { + return affine.wBytes(n, k, q.width()); +} + +fn sBytes(n: u32, k: u32, q: Affine) usize { + return affine.sBytes(n, k, q.width()); +} + fn at(bits: u8) usize { return std.mem.indexOfScalar(u8, &bit_widths, bits).?; } @@ -237,6 +249,18 @@ pub fn byIds(p: *const Pipes, e: mtl.ComputeEncoder, m: Picks, x: Ref, ids: Ref, e.dispatchGroups(mtl.Size.of((m.w.n + 7) / 8, slots, 1), mtl.Size.of(64, 1, 1)); } +/// x through [gate; up] (one width, `n` rows each) SwiGLU'd into out: the by-ids GLU, every id 0 (glm_swiglu's bits). +pub fn glu(p: *const Pipes, e: mtl.ComputeEncoder, gu: Affine, n: u32, x: Ref, out: Ref, rows: u32, lim: f32) void { + std.debug.assert(gu.more == 0 and gu.n == 2 * n and rows <= 64); + var gate = gu; + gate.n = n; + var up = gate; + up.w = gu.w.at(wBytes(n, gu.k, gu)); + up.s = gu.s.at(sBytes(n, gu.k, gu)); + up.b = gu.b.at(sBytes(n, gu.k, gu)); + byIds(p, e, .{ .w = gate, .up = up, .experts = 1, .repeat = 1, .lim = lim }, x, .{ .buf = p.zeros }, out, rows); +} + /// Each expert's sorted pairs (offsets, order): with `up` x row order / repeat to row i, else x row i to row order[i]. pub fn bySorted(p: *const Pipes, e: mtl.ComputeEncoder, m: Picks, x: Ref, offsets: Ref, order: Ref, out: Ref, ldo: u32) void { const u = m.up orelse m.w; diff --git a/zig/tests/glm_rows_check.zig b/zig/tests/glm_rows_check.zig index f5d4ae89c2..cce14c3698 100644 --- a/zig/tests/glm_rows_check.zig +++ b/zig/tests/glm_rows_check.zig @@ -459,6 +459,29 @@ fn splitMmCheck(g: *const Gpu, k: usize, ns: []const usize, bits: []const u8, t: return on_mm; } +fn gluLaunch(e: mtl.ComputeEncoder, p: *const wide.Pipes, gu: affine.Affine, n: u32, x: Ref, out: Ref, m: u32) void { + wide.glu(p, e, gu, n, x, out, m, 0.75); +} + +/// The shared expert's gate and up at one width (wide.glu, one launch) == glm_swiglu(dense [gate; up]) bit for bit. +fn sharedGluCheck(g: *const Gpu) !usize { + var pairs: usize = 0; + for ([_]u8{ 8, 5, 4, 2 }) |bits| { + const n = 2048; + const m = try Mat.init(g.device, 2 * n, 4096, bits, 64, 8800 + @as(u64, bits)); + defer m.deinit(); + const bs = [_]mtl.Buffer{ try buf(g.device, rows * 2 * n * 2), try buf(g.device, rows * n * 2), try buf(g.device, rows * n * 2) }; + defer for (bs) |b| b.deinit(); + try g.run(dense, .{ &g.p, m.q, wide.Io{ .in = ref(m.x()), .ldx = 4096, .out = ref(bs[0]), .ldo = 2 * n, .rows = rows } }); + try g.run(swiglu, .{ g.swiglu, ref(bs[0]), ref(bs[1]), @as(u32, rows), @as(u32, n), @as(f32, 0.75) }); + try g.run(gluLaunch, .{ &g.p, m.q, @as(u32, n), ref(m.x()), ref(bs[2]), @as(u32, rows) }); + try eql(bs[2].slice(u16, rows * n), bs[1].slice(u16, rows * n)); + pairs += rows * n; + } + std.debug.print("shared expert gate|up at one width in one launch == glm_swiglu(dense [gate; up]): {d} bit-equal values\n", .{pairs}); + return pairs; +} + fn perHead(e: mtl.ComputeEncoder, p: *const wide.Pipes, q: affine.Affine, heads: u32, io: wide.Io, hsx: u32, hso: u32) void { wide.perHead(p, e, q, heads, io, hsx, hso); } @@ -534,6 +557,7 @@ pub fn main() !void { all.pairs += split.pairs; all.pairs += try @import("glm_rows_kda.zig").check(&g.p, device, queue); all.pairs += try @import("glm_rows_kda.zig").fusedCheck(&g.p, device, queue); + all.pairs += try sharedGluCheck(&g); all.values += try @import("glm_rows_kda.zig").bf16Check(&g.p, device, queue, &all.worst); all.values += try @import("glm_rows_kda.zig").embedCheck(&g.p, device, queue, &all.worst); const cl = try @import("glm_rows_classes.zig").check(&g.p, device, queue); diff --git a/zig/tests/glm_rows_speed.zig b/zig/tests/glm_rows_speed.zig index 75b1a53eea..a1eee3c1e4 100644 --- a/zig/tests/glm_rows_speed.zig +++ b/zig/tests/glm_rows_speed.zig @@ -207,12 +207,8 @@ fn perHead(e: mtl.ComputeEncoder, i: usize, p: *const wide.Pipes, pool: Pool, io wide.perHead(p, e, q, H, .{ .in = rf(io.x), .ldx = ldx, .out = rf(io.y), .ldo = q.n, .rows = 1 }, hsx, q.n / H); } fn shared(e: mtl.ComputeEncoder, i: usize, p: *const wide.Pipes, k: *const kernels.Kernels, sh: [2]Pool, io: Io) void { - const gu = sh[0].at(i); - wide.dense(p, e, gu, .{ .in = rf(io.x), .ldx = D, .out = rf(io.z), .ldo = 2 * N, .rows = 1 }); - e.setPipeline(k.swiglu); - bind(e, 0, .{ rf(io.z), rf(io.z).at(2 * N * 2) }); - e.setValue(extern struct { rows: i32, width: i32, xs: i32, ys: i32, lim: f32 }{ .rows = 1, .width = N, .xs = 2 * N, .ys = N, .lim = 10.0 }, 2); - e.dispatchThreads(mtl.Size.of(N, 1, 1), mtl.Size.of(256, 1, 1)); + _ = k; + wide.glu(p, e, sh[0].at(i), N, rf(io.x), rf(io.z).at(2 * N * 2), 1, 10.0); // moe.zig sharedRows: one launch const dq = sh[1].at(i); wide.dense(p, e, dq, .{ .in = rf(io.z).at(2 * N * 2), .ldx = N, .out = rf(io.y), .ldo = D, .rows = 1 }); } From e696d4e6cab904eb658d4cb9257e97ecb43eba3c Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 10:12:30 -1000 Subject: [PATCH 17/18] GLM-5.3-Flash decode runs routed experts by pick and the 8-bit shared expert in two launches, the same bits oQ2's MoE layers (routed 2/64, shared 8/64) took four launches: the routed GLU and down by pick, the shared expert's GLU and dense down. glm_moe_rows now runs both in one grid (slots, then the rows through the shared expert), each output with the same tf_rp_rows sums and SwiGLU as the kernel it replaces, so act, acts, ye and ys are the same bits. Two launches a MoE layer fewer (84 a decode token) and the shared expert's small products share the routed grid. Taken when every expert matrix is one run in groups of 64, the routed of one width, the shared at 8 bits; uniform 4/64 keeps the generated gathers and oQ4 (routed 4/64) its generated routed kernels. Test: tf-glm-rows-check's glm_rows_moe.zig: at q2..q8 routed with an 8-bit shared expert, 2 rows of 8 picks over 4 experts (one id past them), wide.moe == byIds GLU and down, wide.glu and dense, bit for bit (fails before: no wide.moe). glm_rows_check's helpers are pub for it (the file stays under 600 lines). --- zig/kernels/metal/glm_rows.metal | 74 ++++++++++++++++++++++++++++++++ zig/src/families/glm/moe.zig | 7 +++ zig/src/families/glm/wide.zig | 41 ++++++++++++++++-- zig/tests/glm_rows_check.zig | 25 +++++------ zig/tests/glm_rows_moe.zig | 51 ++++++++++++++++++++++ zig/tests/glm_rows_speed.metal | 11 +++++ zig/tests/glm_rows_speed.zig | 65 ++++++++++++++++++++++++++-- 7 files changed, 255 insertions(+), 19 deletions(-) create mode 100644 zig/tests/glm_rows_moe.zig diff --git a/zig/kernels/metal/glm_rows.metal b/zig/kernels/metal/glm_rows.metal index 25073b268b..9c87b9b592 100644 --- a/zig/kernels/metal/glm_rows.metal +++ b/zig/kernels/metal/glm_rows.metal @@ -91,6 +91,70 @@ template } } +// Slots as glm_rows_ids (BITS), then X2's rows through the shared expert (SBITS; GLU: gate then up rows) into OUT2. +template // RPS columns a simdgroup; KC: k compiled in +[[kernel]] void glm_moe_rows(const device bfloat* X [[buffer(0)]], const device uint32_t* W [[buffer(1)]], + const device bfloat* S [[buffer(2)]], const device bfloat* B [[buffer(3)]], + constant GlmPickDims& d [[buffer(4)]], device bfloat* OUT [[buffer(5)]], + const device int* ids [[buffer(6)]], const device uint32_t* W2 [[buffer(7)]], + const device bfloat* S2 [[buffer(8)]], const device bfloat* B2 [[buffer(9)]], + const device bfloat* X2 [[buffer(10)]], const device uint32_t* SW [[buffer(11)]], + const device bfloat* SS [[buffer(12)]], const device bfloat* SB [[buffer(13)]], + device bfloat* OUT2 [[buffer(14)]], + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int row0 = (int(tg.x) * 2 + int(sg)) * RPS; + if (row0 >= d.n) return; + const int y = int(tg.y); + const int cols = metal::min(RPS, d.n - row0); + const int K = KC > 0 ? KC : d.k; + const size_t kg = size_t(K / d.group); + const device uint8_t* gw; + const device bfloat *gs, *gb, *x; + device bfloat* out; + size_t up_at = 0; // the up rows' offset from the gate rows, in rows + if (y < d.slots) { + const int expert = ids[y]; + out = OUT + size_t(y) * d.ldo + row0; + if (expert < 0 || expert >= d.experts) { + if (int(lane) < cols) out[lane] = bfloat(0.0f); + return; + } + const size_t at = size_t(expert) * d.n + row0; + const size_t kb = size_t(K) * BITS / 8; + x = X + size_t(y / d.repeat) * K; + float acc[1][RPS]; + tf_rp_rows((const device uint8_t*)W + at * kb, S + at * kg, B + at * kg, x, K, d.group, lane, acc); + if (GLU) { + float up[1][RPS]; + tf_rp_rows((const device uint8_t*)W2 + at * kb, S2 + at * kg, B2 + at * kg, x, K, d.group, lane, up); + if (lane == 0) + for (int j = 0; j < cols; j++) out[j] = glm_rows_glu(acc[0][j], up[0][j], d.lim); + } else if (lane == 0) { + for (int j = 0; j < cols; j++) out[j] = bfloat(acc[0][j]); + } + return; + } + const int r = y - d.slots; + const size_t skb = size_t(K) * SBITS / 8; + x = X2 + size_t(r) * K; + out = OUT2 + size_t(r) * d.n + row0; + gw = (const device uint8_t*)SW + size_t(row0) * skb; + gs = SS + size_t(row0) * kg; + gb = SB + size_t(row0) * kg; + float acc[1][RPS]; + tf_rp_rows(gw, gs, gb, x, K, d.group, lane, acc); + if (GLU) { + up_at = size_t(d.n); + float up[1][RPS]; + tf_rp_rows(gw + up_at * skb, gs + up_at * kg, gb + up_at * kg, x, K, d.group, lane, up); + if (lane == 0) + for (int j = 0; j < cols; j++) out[j] = glm_rows_glu(acc[0][j], up[0][j], d.lim); + } else if (lane == 0) { + for (int j = 0; j < cols; j++) out[j] = bfloat(acc[0][j]); + } +} + // Expert e's sorted pairs: GLU x row order[i] / repeat to out row i, else x row i to row order[i]; grid (n / 8, e). template [[kernel]] void glm_rows_sorted(const device bfloat* X [[buffer(0)]], const device uint32_t* W [[buffer(1)]], @@ -167,3 +231,13 @@ GLM_ROWS_AT(4) GLM_ROWS_AT(5) GLM_ROWS_AT(6) GLM_ROWS_AT(8) + +#define GLM_MOE_AT(B) \ + template [[host_name("glm_moe_rows_q" #B "_s8")]] [[kernel]] decltype(glm_moe_rows) glm_moe_rows; \ + template [[host_name("glm_moe_rows_glu_q" #B "_s8")]] [[kernel]] decltype(glm_moe_rows) glm_moe_rows; +GLM_MOE_AT(2) +GLM_MOE_AT(3) +GLM_MOE_AT(4) +GLM_MOE_AT(5) +GLM_MOE_AT(6) +GLM_MOE_AT(8) diff --git a/zig/src/families/glm/moe.zig b/zig/src/families/glm/moe.zig index 0e3d84b293..9604439859 100644 --- a/zig/src/families/glm/moe.zig +++ b/zig/src/families/glm/moe.zig @@ -50,6 +50,9 @@ pub fn moe(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, r bind(e, 4, .{sc.branch}); e.dispatchThreads(size(rows * c.hidden, 1, 1), size(256, 1, 1)); return; + } else if (x.k.wide != null and wide.moeOne(parts(w)) and s & (Class.shared | Class.routed) == 0 and on(x, "e_gateup") and on(x, "e_down") and on(x, "s_gateup") and on(x, "s_down")) { + if (s & Class.route == 0) route(x, e, w, x_in, rows); + wide.moe(&x.k.wide.?, e, parts(w), c.experts, c.topk, c.swiglu_limit, .{ .x = x_in, .ids = sc.pick, .act = sc.act, .acts = sc.acts, .ye = sc.ye, .ys = sc.ys }, rows); } else { if (s & Class.shared == 0) experts(x, e, w, x_in, rows, 1, .{ sc.none, sc.none, sc.none }); if (s & Class.route == 0) route(x, e, w, x_in, rows); @@ -63,6 +66,10 @@ pub fn moe(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, r e.dispatchThreads(size(rows * c.hidden, 1, 1), size(256, 1, 1)); } +fn parts(w: *const wts.Moe) wide.Moe { + return .{ .gate = w.gate, .up = w.up, .down = w.down, .sh_gate_up = w.sh_gate_up, .sh_down = w.sh_down }; +} + /// The router's fp32 logits and the route: picks, weights and the window's unique experts with their members. fn route(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, rows: u32) void { const c = x.c; diff --git a/zig/src/families/glm/wide.zig b/zig/src/families/glm/wide.zig index 7e40bf3da0..4dbf759211 100644 --- a/zig/src/families/glm/wide.zig +++ b/zig/src/families/glm/wide.zig @@ -25,6 +25,8 @@ pub const Pipes = struct { glu_ids: [nw]mtl.Pipeline, // gate and up through it, SwiGLU'd sorted: [nw]mtl.Pipeline, // an expert's sorted pairs (prompt chunks) glu_sorted: [nw]mtl.Pipeline, + moe: [nw]mtl.Pipeline, // routed experts by pick at each width and the 8-bit shared expert, one launch + moe_glu: [nw]mtl.Pipeline, // the same through gate and up, SwiGLU'd bf16: mtl.Pipeline, // an unquantized matrix (the MTP head's eh_proj) kda_pre: mtl.Pipeline, // kda_rows reading f_b's and g_b's products made apart (any width but 4/64) kda_fg: [nw]mtl.Pipeline, // kda_rows making f_b's and g_b's dots of one width (g64) as glm_rows does: .pre's bits @@ -37,7 +39,7 @@ pub const Pipes = struct { hash: u64, // every source compiled here pub fn deinit(p: *Pipes) void { - inline for (.{ &p.rows, &p.ids, &p.glu_ids, &p.sorted, &p.glu_sorted, &p.embed, &p.kda_fg }) |set| for (set) |x| x.deinit(); + inline for (.{ &p.rows, &p.ids, &p.glu_ids, &p.sorted, &p.glu_sorted, &p.moe, &p.moe_glu, &p.embed, &p.kda_fg }) |set| for (set) |x| x.deinit(); p.bf16.deinit(); p.kda_pre.deinit(); p.zeros.deinit(); @@ -82,7 +84,7 @@ pub const glue_math = blk: { /// The pipelines built so far: a failed compile releases them all. const Made = struct { - list: [nw * 6 + 4 + 4 * affine_mm.names.len]mtl.Pipeline = undefined, + list: [nw * 8 + 4 + 4 * affine_mm.names.len]mtl.Pipeline = undefined, n: usize = 0, released: usize = 0, @@ -129,6 +131,8 @@ fn build(gpa: std.mem.Allocator, device: mtl.Device, kda: []const u8, made: *Mad p.glu_ids[i] = try made.add(device, l, "glm_rows_glu_ids" ++ q); p.sorted[i] = try made.add(device, l, "glm_rows_sorted" ++ q); p.glu_sorted[i] = try made.add(device, l, "glm_rows_glu_sorted" ++ q); + p.moe[i] = try made.add(device, l, "glm_moe_rows" ++ q ++ "_s8"); + p.moe_glu[i] = try made.add(device, l, "glm_moe_rows_glu" ++ q ++ "_s8"); } p.bf16 = try made.add(device, l, "glm_rows_bf16"); } @@ -261,6 +265,37 @@ pub fn glu(p: *const Pipes, e: mtl.ComputeEncoder, gu: Affine, n: u32, x: Ref, o byIds(p, e, .{ .w = gate, .up = up, .experts = 1, .repeat = 1, .lim = lim }, x, .{ .buf = p.zeros }, out, rows); } +/// A MoE block's routed experts and shared expert (`sh`: [gate; up] one run, then down), all g64, the shared at 8 bits. +pub const Moe = struct { gate: Affine, up: Affine, down: Affine, sh_gate_up: Affine, sh_down: Affine }; + +/// Whether `m` takes `moe`: every matrix one run, g64, the routed of one width (not 4/64), the shared at 8 bits. +pub fn moeOne(m: Moe) bool { + if (native(m.gate) and native(m.up) and native(m.down)) return false; // routed 4/64 keeps the generated gathers + for ([_]Affine{ m.gate, m.up, m.down, m.sh_gate_up, m.sh_down }) |q| if (q.more != 0 or q.group != 64) return false; + return m.up.bits == m.gate.bits and m.down.bits == m.gate.bits and m.sh_gate_up.bits == 8 and m.sh_down.bits == 8; +} + +/// byIds's GLU and down for the slots, glu and dense for the shared expert's rows: a launch each step, the same bits. +pub fn moe(p: *const Pipes, e: mtl.ComputeEncoder, m: Moe, experts: u32, topk: u32, lim: f32, io: MoeIo, rows: u32) void { + const slots = rows * topk; + const g = m.sh_gate_up; + const steps = [2]struct { pipe: mtl.Pipeline, w: Affine, up: ?Affine, sh: Affine, x: Ref, x2: Ref, out: Ref, out2: Ref, repeat: u32 }{ + .{ .pipe = p.moe_glu[at(m.gate.bits)], .w = m.gate, .up = m.up, .sh = g, .x = io.x, .x2 = io.x, .out = io.act, .out2 = io.acts, .repeat = topk }, + .{ .pipe = p.moe[at(m.down.bits)], .w = m.down, .up = null, .sh = m.sh_down, .x = io.act, .x2 = io.acts, .out = io.ye, .out2 = io.ys, .repeat = 1 }, + }; + for (steps) |s| { + const u = s.up orelse s.w; + e.setPipeline(s.pipe); + bind(e, 0, .{ s.x, s.w.w, s.w.s, s.w.b }); + e.setValue(PickDims{ .k = @intCast(s.w.k), .n = @intCast(s.w.n), .group = 64, .slots = @intCast(slots), .repeat = @intCast(s.repeat), .experts = @intCast(experts), .ldo = @intCast(s.w.n), .lim = lim }, 4); + bind(e, 5, .{ s.out, io.ids, u.w, u.s, u.b, s.x2, s.sh.w, s.sh.s, s.sh.b, s.out2 }); + e.dispatchGroups(mtl.Size.of((s.w.n + 7) / 8, slots + rows, 1), mtl.Size.of(64, 1, 1)); + } +} + +/// `moe`'s buffers: x rows, the picks' expert ids, act and acts (gate|up out), ye and ys (down out). +pub const MoeIo = struct { x: Ref, ids: Ref, act: Ref, acts: Ref, ye: Ref, ys: Ref }; + /// Each expert's sorted pairs (offsets, order): with `up` x row order / repeat to row i, else x row i to row order[i]. pub fn bySorted(p: *const Pipes, e: mtl.ComputeEncoder, m: Picks, x: Ref, offsets: Ref, order: Ref, out: Ref, ldo: u32) void { const u = m.up orelse m.w; @@ -338,7 +373,7 @@ test "a failed compile releases every pipeline built before it; a load counts ea defer device.deinit(); var made: Made = .{}; try std.testing.expectError(error.KernelCompile, build(gpa, device, "", &made)); // after the row kernels - try std.testing.expectEqual(@as(usize, nw * 5 + 1), made.released); + try std.testing.expectEqual(@as(usize, nw * 7 + 1), made.released); try std.testing.expectEqual(@as(usize, 0), made.n); var all: Made = .{}; var p = try build(gpa, device, comptime kernels.generatedSource("kda_rows"), &all); diff --git a/zig/tests/glm_rows_check.zig b/zig/tests/glm_rows_check.zig index cce14c3698..b2866b7c9c 100644 --- a/zig/tests/glm_rows_check.zig +++ b/zig/tests/glm_rows_check.zig @@ -22,21 +22,21 @@ fn toBf(v: f32) u16 { fn fromBf(v: u16) f64 { return @as(f32, @bitCast(@as(u32, v) << 16)); } -fn buf(device: mtl.Device, bytes: usize) !mtl.Buffer { +pub fn buf(device: mtl.Device, bytes: usize) !mtl.Buffer { return device.buffer(@max(bytes, 16), opts); } -fn ref(b: mtl.Buffer) Ref { +pub fn ref(b: mtl.Buffer) Ref { return .{ .buf = b }; } /// A synthetic MLX matrix [n, k], its fp64 dequantization and its codes alone; x: `rows` rows of inputs. -const Mat = struct { +pub const Mat = struct { q: affine.Affine, bufs: [4]mtl.Buffer, deq: []f64, codes: []f64, - fn init(device: mtl.Device, n: usize, k: usize, bits: u8, group: usize, seed: u64) !Mat { + pub fn init(device: mtl.Device, n: usize, k: usize, bits: u8, group: usize, seed: u64) !Mat { var prng = std.Random.DefaultPrng.init(seed); const r = prng.random(); const words = n * k * bits / 32; @@ -58,13 +58,13 @@ const Mat = struct { return m; } - fn deinit(m: Mat) void { + pub fn deinit(m: Mat) void { for (m.bufs) |b| b.deinit(); a.free(m.deq); a.free(m.codes); } - fn x(m: Mat) mtl.Buffer { + pub fn x(m: Mat) mtl.Buffer { return m.bufs[3]; } @@ -119,7 +119,7 @@ const Mat = struct { } }; -const Gpu = struct { +pub const Gpu = struct { device: mtl.Device, queue: mtl.Queue, p: wide.Pipes, @@ -127,7 +127,7 @@ const Gpu = struct { swiglu: mtl.Pipeline, /// One command buffer around `f(e, args)`. - fn run(g: *const Gpu, comptime f: anytype, args: anytype) !void { + pub fn run(g: *const Gpu, comptime f: anytype, args: anytype) !void { const cb = g.queue.commandBuffer(); const e = cb.compute(.serial); @call(.auto, f, .{e} ++ args); @@ -141,13 +141,13 @@ const Gpu = struct { } }; -fn dense(e: mtl.ComputeEncoder, p: *const wide.Pipes, q: affine.Affine, io: wide.Io) void { +pub fn dense(e: mtl.ComputeEncoder, p: *const wide.Pipes, q: affine.Affine, io: wide.Io) void { wide.dense(p, e, q, io); } const Tally = struct { values: usize = 0, worst: f64 = 0, pairs: usize = 0 }; -fn eql(got: []const u16, want: []const u16) !void { +pub fn eql(got: []const u16, want: []const u16) !void { if (!std.mem.eql(u16, got, want)) { for (got, want, 0..) |x, y, i| if (x != y) { std.debug.print("first difference at {d}: {x} vs {x}\n", .{ i, x, y }); @@ -459,7 +459,7 @@ fn splitMmCheck(g: *const Gpu, k: usize, ns: []const usize, bits: []const u8, t: return on_mm; } -fn gluLaunch(e: mtl.ComputeEncoder, p: *const wide.Pipes, gu: affine.Affine, n: u32, x: Ref, out: Ref, m: u32) void { +pub fn gluLaunch(e: mtl.ComputeEncoder, p: *const wide.Pipes, gu: affine.Affine, n: u32, x: Ref, out: Ref, m: u32) void { wide.glu(p, e, gu, n, x, out, m, 0.75); } @@ -485,7 +485,7 @@ fn sharedGluCheck(g: *const Gpu) !usize { fn perHead(e: mtl.ComputeEncoder, p: *const wide.Pipes, q: affine.Affine, heads: u32, io: wide.Io, hsx: u32, hso: u32) void { wide.perHead(p, e, q, heads, io, hsx, hso); } -fn byIds(e: mtl.ComputeEncoder, p: *const wide.Pipes, m: wide.Picks, x: Ref, ids: Ref, out: Ref, slots: u32) void { +pub fn byIds(e: mtl.ComputeEncoder, p: *const wide.Pipes, m: wide.Picks, x: Ref, ids: Ref, out: Ref, slots: u32) void { wide.byIds(p, e, m, x, ids, out, slots); } fn bySorted(e: mtl.ComputeEncoder, p: *const wide.Pipes, m: wide.Picks, x: Ref, offsets: Ref, order: Ref, out: Ref, ldo: u32) void { @@ -558,6 +558,7 @@ pub fn main() !void { all.pairs += try @import("glm_rows_kda.zig").check(&g.p, device, queue); all.pairs += try @import("glm_rows_kda.zig").fusedCheck(&g.p, device, queue); all.pairs += try sharedGluCheck(&g); + all.pairs += try @import("glm_rows_moe.zig").check(&g); all.values += try @import("glm_rows_kda.zig").bf16Check(&g.p, device, queue, &all.worst); all.values += try @import("glm_rows_kda.zig").embedCheck(&g.p, device, queue, &all.worst); const cl = try @import("glm_rows_classes.zig").check(&g.p, device, queue); diff --git a/zig/tests/glm_rows_moe.zig b/zig/tests/glm_rows_moe.zig new file mode 100644 index 0000000000..41de7fbe8b --- /dev/null +++ b/zig/tests/glm_rows_moe.zig @@ -0,0 +1,51 @@ +//! glm_rows_check's MoE launches: the routed and shared experts' two-launch path (wide.moe) == its four launches. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const wide = tf.glm.wide; +const affine = tf.glm.affine; +const rc = @import("glm_rows_check.zig"); + +fn moeLaunch(e: mtl.ComputeEncoder, p: *const wide.Pipes, m: wide.Moe, io: wide.MoeIo, n_rows: u32) void { + wide.moe(p, e, m, 4, 8, 0.75, io, n_rows); +} + +/// wide.moe (routed at q2..q8, shared at 8, two launches) == byIds GLU/down + glu/dense bit for bit; id 5 of 4: zeros. +pub fn check(g: *const rc.Gpu) !usize { + const n = 256; + const k = 512; + const n_rows = 2; + const slots = 8 * n_rows; + var pairs: usize = 0; + const sh = [2]rc.Mat{ try rc.Mat.init(g.device, 2 * n, k, 8, 64, 7300), try rc.Mat.init(g.device, k, n, 8, 64, 7301) }; + defer for (sh) |m| m.deinit(); + for (affine.row_bits) |bits| { + const ms = [3]rc.Mat{ try rc.Mat.init(g.device, 4 * n, k, bits, 64, 7400 + @as(u64, bits)), try rc.Mat.init(g.device, 4 * n, k, bits, 64, 7500 + @as(u64, bits)), try rc.Mat.init(g.device, 4 * k, n, bits, 64, 7600 + @as(u64, bits)) }; + defer for (ms) |m| m.deinit(); + var q: [3]affine.Affine = undefined; + for (&q, ms, [_]u32{ n, n, k }) |*x, m, r| { + x.* = m.q; + x.n = r; + } + const bs = [_]mtl.Buffer{ try rc.buf(g.device, slots * 4), try rc.buf(g.device, 2 * slots * n * 2), try rc.buf(g.device, 2 * n_rows * n * 2), try rc.buf(g.device, 2 * slots * k * 2), try rc.buf(g.device, 2 * n_rows * k * 2) }; + defer for (bs) |b| b.deinit(); + const picks = [slots]i32{ 2, 0, 3, 1, 1, 2, 0, 3, 3, 3, 0, 5, 1, 2, 2, 0 }; + @memcpy(bs[0].slice(i32, slots), &picks); + const m: wide.Moe = .{ .gate = q[0], .up = q[1], .down = q[2], .sh_gate_up = sh[0].q, .sh_down = sh[1].q }; + if (!wide.moeOne(m) and bits != 4) return error.NotOneLaunch; // 4/64 routed keeps the generated gathers + const io: wide.MoeIo = .{ .x = rc.ref(ms[0].x()), .ids = rc.ref(bs[0]), .act = rc.ref(bs[1]), .acts = rc.ref(bs[2]), .ye = rc.ref(bs[3]), .ys = rc.ref(bs[4]) }; + try g.run(moeLaunch, .{ &g.p, m, io, @as(u32, n_rows) }); + const half = [5]usize{ 0, slots * n, n_rows * n, slots * k, n_rows * k }; // the reference in each buffer's second half + try g.run(rc.byIds, .{ &g.p, wide.Picks{ .w = q[0], .up = q[1], .experts = 4, .repeat = 8, .lim = 0.75 }, io.x, io.ids, io.act.at(half[1] * 2), @as(u32, slots) }); + try g.run(rc.gluLaunch, .{ &g.p, sh[0].q, @as(u32, n), io.x, io.acts.at(half[2] * 2), @as(u32, n_rows) }); + try g.run(rc.byIds, .{ &g.p, wide.Picks{ .w = q[2], .experts = 4, .repeat = 1 }, io.act, io.ids, io.ye.at(half[3] * 2), @as(u32, slots) }); + try g.run(rc.dense, .{ &g.p, sh[1].q, wide.Io{ .in = io.acts, .ldx = n, .out = io.ys.at(half[4] * 2), .ldo = k, .rows = n_rows } }); + for (bs[1..], half[1..]) |b, h| { + const v = b.slice(u16, 2 * h); + try rc.eql(v[0..h], v[h..]); + pairs += h; + } + } + std.debug.print("routed experts at q2..q8 and the 8-bit shared expert in two launches == byIds, glu and dense: {d} bit-equal values\n", .{pairs}); + return pairs; +} diff --git a/zig/tests/glm_rows_speed.metal b/zig/tests/glm_rows_speed.metal index 5b700bf300..561a5a3772 100644 --- a/zig/tests/glm_rows_speed.metal +++ b/zig/tests/glm_rows_speed.metal @@ -31,3 +31,14 @@ template #define GXB(B) GXK(B, 0) GXK(B, 4096) GXK(B, 8192) GXK(B, 16384) GXK(B, 12288) GXK(B, 1536) GXK(B, 2048) GXK(B, 128) GXK(B, 256) GXK(B, 512) GXB(8) + +// glm_moe_rows (glm_rows.metal) at other columns a simdgroup and with k compiled in: oQ2's routed q2 + shared q8. +#define GXM(R, K) \ + template [[host_name("x_moe_glu_r" #R "_k" #K)]] [[kernel]] decltype(glm_moe_rows<2, 8, true, R, K>) glm_moe_rows<2, 8, true, R, K>; \ + template [[host_name("x_moe_r" #R "_k" #K)]] [[kernel]] decltype(glm_moe_rows<2, 8, false, R, K>) glm_moe_rows<2, 8, false, R, K>; +GXM(8, 0) +GXM(2, 0) +template [[host_name("x_moe_glu_r4_k4096")]] [[kernel]] decltype(glm_moe_rows<2, 8, true, 4, 4096>) glm_moe_rows<2, 8, true, 4, 4096>; +template [[host_name("x_moe_glu_r8_k4096")]] [[kernel]] decltype(glm_moe_rows<2, 8, true, 8, 4096>) glm_moe_rows<2, 8, true, 8, 4096>; +template [[host_name("x_moe_r4_k2048")]] [[kernel]] decltype(glm_moe_rows<2, 8, false, 4, 2048>) glm_moe_rows<2, 8, false, 4, 2048>; +template [[host_name("x_moe_r8_k2048")]] [[kernel]] decltype(glm_moe_rows<2, 8, false, 8, 2048>) glm_moe_rows<2, 8, false, 8, 2048>; diff --git a/zig/tests/glm_rows_speed.zig b/zig/tests/glm_rows_speed.zig index a1eee3c1e4..788e3e1f50 100644 --- a/zig/tests/glm_rows_speed.zig +++ b/zig/tests/glm_rows_speed.zig @@ -273,6 +273,13 @@ fn kdaStep(e: mtl.ComputeEncoder, i: usize, pipe: mtl.Pipeline, s: *const Kda, g e.dispatchThreads(mtl.Size.of(32, 32, H), mtl.Size.of(32, 32, 1)); } +/// glm_rows.metal with glm_rows_speed.metal's candidate instances. +fn candidateLib(device: mtl.Device) !mtl.Library { + const src = try std.mem.concat(a, u8, &.{ sources.core_row_projection, wide.glue_math, sources.glm_rows, @embedFile("glm_rows_speed.metal") }); + defer a.free(src); + return mtl.Library.fromSource(device, src, mtl.CompileOptions.mlx()); +} + const PM = 2048; // a prompt chunk's rows /// A prompt chunk's per-head product (embed_q: x qp's slices; unembed_out: the latent rows) on the row kernels. @@ -377,6 +384,49 @@ fn candidates(b: *const Bench, lib: mtl.Library, p: *const wide.Pipes, pool: Poo std.debug.print(" {s:<16} glm_rows {d:.1} us, best {s} {d:.1} us ({d:.0} GB/s)\n", .{ name, base, best_name[0..best_len], best, @as(f64, @floatFromInt(pool.bytes)) / best / 1e3 }); } +/// oQ2's MoE: routed by pick and the shared expert, a launch for gate|up and one for down (wide.moe). +fn moeTwo(e: mtl.ComputeEncoder, i: usize, p: *const wide.Pipes, gu: [2]Pool, down: Pool, sh: [2]Pool, io: Io) void { + const m: wide.Moe = .{ .gate = gu[0].at(i), .up = gu[1].at(i), .down = down.at(i), .sh_gate_up = sh[0].at(i), .sh_down = sh[1].at(i) }; + const z = rf(io.z); + wide.moe(p, e, m, E, TOPK, 10.0, .{ .x = rf(io.x), .ids = rf(io.ids).at((i % reps) * TOPK * 4), .act = z, .acts = z.at(TOPK * N * 2), .ye = rf(io.y), .ys = rf(io.y).at(TOPK * D * 2) }, 1); +} + +/// moeTwo through candidate pipelines (glu, down) at `rps` columns a simdgroup. +fn moeCand(e: mtl.ComputeEncoder, i: usize, pipes: [2]mtl.Pipeline, rps: [2]u32, gu: [2]Pool, down: Pool, sh: [2]Pool, io: Io) void { + const z = rf(io.z); + const ids = rf(io.ids).at((i % reps) * TOPK * 4); + const steps = [2]struct { w: Affine, up: Affine, sh: Affine, x: Ref, out: Ref, out2: Ref, repeat: u32 }{ + .{ .w = gu[0].at(i), .up = gu[1].at(i), .sh = sh[0].at(i), .x = rf(io.x), .out = z, .out2 = z.at(TOPK * N * 2), .repeat = TOPK }, + .{ .w = down.at(i), .up = down.at(i), .sh = sh[1].at(i), .x = z, .out = rf(io.y), .out2 = rf(io.y).at(TOPK * D * 2), .repeat = 1 }, + }; + for (steps, pipes, rps, 0..) |st, pipe, r, j| { + e.setPipeline(pipe); + bind(e, 0, .{ st.x, st.w.w, st.w.s, st.w.b }); + e.setValue(wide.PickDims{ .k = @intCast(st.w.k), .n = @intCast(st.w.n), .group = 64, .slots = TOPK, .repeat = @intCast(st.repeat), .experts = E, .ldo = @intCast(st.w.n), .lim = 10.0 }, 4); + bind(e, 5, .{ st.out, ids, st.up.w, st.up.s, st.up.b, if (j == 0) rf(io.x) else z.at(TOPK * N * 2), st.sh.w, st.sh.s, st.sh.b, st.out2 }); + e.dispatchGroups(mtl.Size.of((st.w.n + 2 * r - 1) / (2 * r), TOPK + 1, 1), mtl.Size.of(64, 1, 1)); + } +} + +/// glm_moe_rows candidates (glm_rows_speed.metal) for oQ2's MoE: bits == wide.moe's, then the time a layer. +fn moeCandidates(b: *const Bench, lib: mtl.Library, p: *const wide.Pipes, gu: [2]Pool, down: Pool, sh: [2]Pool) !void { + const base = try b.time(moeTwo, .{ p, gu, down, sh, b.io }); + try b.run(moeTwo, .{ p, gu, down, sh, b.io }); + const want = try a.dupe(u16, b.io.y.slice(u16, (TOPK + 1) * D)); + defer a.free(want); + for ([_][2][:0]const u8{ .{ "x_moe_glu_r8_k0", "x_moe_r8_k0" }, .{ "x_moe_glu_r2_k0", "x_moe_r2_k0" }, .{ "x_moe_glu_r4_k4096", "x_moe_r4_k2048" }, .{ "x_moe_glu_r8_k4096", "x_moe_r8_k2048" } }) |names| { + var pipes: [2]mtl.Pipeline = undefined; + for (&pipes, names) |*x, n| x.* = try mtl.Pipeline.init(b.device, lib, n, false); + defer for (pipes) |x| x.deinit(); + const rps = [2]u32{ names[0][11] - '0', names[1][7] - '0' }; + @memset(b.io.y.slice(u16, (TOPK + 1) * D), 0); + try b.run(moeCand, .{ pipes, rps, gu, down, sh, b.io }); + if (!std.mem.eql(u16, b.io.y.slice(u16, (TOPK + 1) * D), want)) return error.BitsDiffer; + const us = try b.time(moeCand, .{ pipes, rps, gu, down, sh, b.io }); + std.debug.print(" moe q2+q8 {s} / {s}: {d:8.1} us ({d:5.3}x wide.moe {d:.1} us)\n", .{ names[0], names[1], us, base / us, base }); + } +} + fn routedRows(e: mtl.ComputeEncoder, i: usize, p: *const wide.Pipes, gu: [2]Pool, down: Pool, io: Io) void { const ids = rf(io.ids).at((i % reps) * TOPK * 4); wide.byIds(p, e, .{ .w = gu[0].at(i), .up = gu[1].at(i), .experts = E, .repeat = TOPK, .lim = 10.0 }, rf(io.x), ids, rf(io.z), TOPK); @@ -509,17 +559,24 @@ pub fn main() !void { const sh8 = [2]Pool{ try Pool.init(device, 2 * N, D, 8, 1, 2), try Pool.init(device, D, N, 8, 1, 2) }; defer for (sh8) |x| x.deinit(); const us_sh = try b.time(shared, .{ p, k, sh8, io }); - oq2.add("moe shared q8", us_sh, 42, sh8[0].bytes + sh8[1].bytes); + std.debug.print(" {s:<26} {d:9.1} us (oQ2: in the routed launches below)\n", .{ "moe shared q8", us_sh }); oq4.add("moe shared q8", us_sh, 42, sh8[0].bytes + sh8[1].bytes); + const alt = try b.time(moeTwo, .{ p, gu4, dn4, sh8, io }); // oQ4's routed 4/64 on glm_rows sums with the shared + std.debug.print(" {s:<26} {d:9.1} us (not taken: changes routed bits)\n", .{ "moe routed 4 + shared q8", alt }); const gu2 = [2]Pool{ try Pool.init(device, N, D, 2, E, 1), try Pool.init(device, N, D, 2, E, 1) }; defer for (gu2) |x| x.deinit(); const dn2 = try Pool.init(device, D, N, 2, E, 1); defer dn2.deinit(); - oq2.add("moe routed q2 (ids)", try b.time(routedRows, .{ p, gu2, dn2, io }), 42, 3 * affine.bytes(N, D, .{ .bits = 2, .group = 64 }) * TOPK); + std.debug.print(" {s:<26} {d:9.1} us (apart)\n", .{ "moe routed q2 (ids)", try b.time(routedRows, .{ p, gu2, dn2, io }) }); + oq2.add("moe routed q2 + shared q8", try b.time(moeTwo, .{ p, gu2, dn2, sh8, io }), 42, 3 * affine.bytes(N, D, .{ .bits = 2, .group = 64 }) * TOPK + sh8[0].bytes + sh8[1].bytes); + if (std.c.getenv("GLM_SPEED_CANDIDATES") != null) { + const lib = try candidateLib(device); + defer lib.deinit(); + try moeCandidates(&b, lib, p, gu2, dn2, sh8); + } } if (std.c.getenv("GLM_SPEED_CANDIDATES") != null) { - const src = try std.mem.concat(a, u8, &.{ sources.core_row_projection, wide.glue_math, sources.glm_rows, @embedFile("glm_rows_speed.metal") }); - const lib = try mtl.Library.fromSource(device, src, mtl.CompileOptions.mlx()); + const lib = try candidateLib(device); defer lib.deinit(); std.debug.print("candidates at 8 bits (bits checked == glm_rows):\n", .{}); for (shapes) |s| { From 67ab196ee24bd4c30f984caddb8c41862d6bee2a Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Tue, 13 Oct 2026 00:05:51 -1000 Subject: [PATCH 18/18] Core row projections read 2-bit codes in place, as they read 4-bit nibbles: the same products tf_rp_qdot16 read each 2-bit code with a shift and a mask; it now masks the code in place (code * 4^i) and tf_rp_load16 scales the code's input by 4^-i, a power of two, so every product is the same float and every sum the same order. 4-bit codes already work this way. Every 2-bit row product takes it, GLM-5.3-Flash's routed experts by pick and in the fused MoE launch among them. tf-glm-rows-speed, oQ2's MoE layer on M3 Ultra (routed q2 + shared q8, both launches): 158.7 -> 132.4 us (1.20x), every output equal to the shift-and-mask kernel's. --- zig/kernels/metal/core/row_projection.metal | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/zig/kernels/metal/core/row_projection.metal b/zig/kernels/metal/core/row_projection.metal index 5a503e3acc..d985ee1e7f 100644 --- a/zig/kernels/metal/core/row_projection.metal +++ b/zig/kernels/metal/core/row_projection.metal @@ -2,6 +2,9 @@ #include using namespace metal; +// 2^-e, exact: a code read in place at bit e times its input scaled by 2^-e is the same product. +inline float tf_rp_p2(int e) { return as_type(uint(127 - e) << 23); } + // 16 bf16 inputs as floats (scaled to the nibble positions at 4 bits) and their group sum: fp32 in order, or rounded to bf16 at each add (BF16SUM, the Nemotron expert convention). template inline float tf_rp_load16(const device bfloat* x, thread float* xt) { @@ -10,10 +13,10 @@ inline float tf_rp_load16(const device bfloat* x, thread float* xt) { const bfloat a = x[i], b = x[i + 1], c = x[i + 2], d = x[i + 3]; if (BF16SUM) sum += float(bfloat(float(bfloat(float(bfloat(float(a) + float(b))) + float(c))) + float(d))); else { sum += float(a); sum += float(b); sum += float(c); sum += float(d); } - xt[i] = float(a); - xt[i + 1] = BITS == 4 ? float(b) / 16.0f : float(b); - xt[i + 2] = BITS == 4 ? float(c) / 256.0f : float(c); - xt[i + 3] = BITS == 4 ? float(d) / 4096.0f : float(d); + xt[i] = BITS == 2 ? float(a) * tf_rp_p2(2 * i) : float(a); + xt[i + 1] = BITS == 4 ? float(b) / 16.0f : BITS == 2 ? float(b) * tf_rp_p2(2 * i + 2) : float(b); + xt[i + 2] = BITS == 4 ? float(c) / 256.0f : BITS == 2 ? float(c) * tf_rp_p2(2 * i + 4) : float(c); + xt[i + 3] = BITS == 4 ? float(d) / 4096.0f : BITS == 2 ? float(d) * tf_rp_p2(2 * i + 6) : float(d); } return sum; } @@ -25,7 +28,7 @@ inline void tf_rp_qdot16(const device uint8_t* w, const thread float (*xt)[16], float q[16]; if (BITS == 2) { const uint v = *(const device uint*)w; - for (int i = 0; i < 16; i++) q[i] = float((v >> (2 * i)) & 3u); + for (int i = 0; i < 16; i++) q[i] = float(v & (3u << (2 * i))); } else if (BITS == 4) { const device uint16_t* ws = (const device uint16_t*)w; for (int i = 0; i < 4; i++) {