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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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 696c31286e8c36cc792a06f7d265bd0d6641cc75 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:33:31 -1000 Subject: [PATCH 07/16] The 27B resolves each loaded module's affine width and admits it against the kernels this chip runs Only the modules the engine loads are resolved, so draft-head and vision overrides never refuse a checkpoint. A width or group the projection kernels cannot run yet refuses naming the module, layer and width; fused members of different widths are admitted as runs of one width. --- zig/src/families/qwen27/qwen27.zig | 2 + zig/src/families/qwen27/widths.zig | 255 +++++++++++++++++++++++++++++ 2 files changed, 257 insertions(+) create mode 100644 zig/src/families/qwen27/widths.zig diff --git a/zig/src/families/qwen27/qwen27.zig b/zig/src/families/qwen27/qwen27.zig index a1ba407bbd..423e1b7fd7 100644 --- a/zig/src/families/qwen27/qwen27.zig +++ b/zig/src/families/qwen27/qwen27.zig @@ -2,6 +2,7 @@ pub const dflash = @import("dflash/root.zig"); pub const config = @import("config.zig"); pub const affine = @import("affine.zig"); +pub const widths = @import("widths.zig"); pub const checkpoint = @import("checkpoint.zig"); pub const weights = @import("weights.zig"); pub const attention = @import("attention.zig"); @@ -49,6 +50,7 @@ test { _ = prompt_state; _ = config; _ = affine; + _ = widths; _ = checkpoint; _ = weights; _ = @import("weights_test.zig"); diff --git a/zig/src/families/qwen27/widths.zig b/zig/src/families/qwen27/widths.zig new file mode 100644 index 0000000000..69a786ba2c --- /dev/null +++ b/zig/src/families/qwen27/widths.zig @@ -0,0 +1,255 @@ +//! Affine widths of the modules the engine loads, admitted against the projection kernels this chip runs. +const std = @import("std"); +const core = @import("core"); +const cfg = @import("config.zig"); +const affine = @import("affine.zig"); +const glue = @import("glue.zig"); +const row = core.row_projection; +const lane = core.lane_projection; + +/// One affine width, as core.quant.mlx.Width: `bits` per weight, a scale and bias a `group` of K. +pub const Width = struct { bits: u8, group: u16 }; +/// The qualified format: today's kernels and launches, unchanged. +pub const uniform = Width{ .bits = 4, .group = 64 }; + +pub fn isUniform(w: Width) bool { + return w.bits == uniform.bits and w.group == uniform.group; +} + +/// A layer's projections; a layer stores entries for both mixers and uses those of its kind. +pub const Member = enum { gate, up, down, qkv, z, b, a, out, q, k, v, o }; +const members = @backingInt(Member.o) + 1; + +pub fn leaf(m: Member) []const u8 { + return switch (m) { + .gate => "mlp.gate_proj", + .up => "mlp.up_proj", + .down => "mlp.down_proj", + .qkv => "linear_attn.in_proj_qkv", + .z => "linear_attn.in_proj_z", + .b => "linear_attn.in_proj_b", + .a => "linear_attn.in_proj_a", + .out => "linear_attn.out_proj", + .q => "self_attn.q_proj", + .k => "self_attn.k_proj", + .v => "self_attn.v_proj", + .o => "self_attn.o_proj", + }; +} + +/// Output rows and input width of a member. +pub fn shape(c: cfg.Config, m: Member) [2]usize { + return switch (m) { + .gate, .up => .{ c.intermediate, c.hidden }, + .down => .{ c.hidden, c.intermediate }, + .qkv => .{ c.gdnQkvDim(), c.hidden }, + .z => .{ c.gdnValueDim(), c.hidden }, + .b, .a => .{ c.v_heads, c.hidden }, + .out => .{ c.hidden, c.gdnValueDim() }, + .q => .{ c.qDim(), c.hidden }, + .k, .v => .{ c.kvDim(), c.hidden }, + .o => .{ c.hidden, c.heads * c.head_dim }, + }; +} + +/// The device matrices of a layer: members fused into one output, in output-row order. +pub fn stacks(kind: cfg.Kind) []const []const Member { + return switch (kind) { + .linear => &.{ &.{ .gate, .up }, &.{.down}, &.{.qkv}, &.{ .z, .b, .a }, &.{.out} }, + .attention => &.{ &.{ .gate, .up }, &.{.down}, &.{.q}, &.{ .k, .v }, &.{.o} }, + }; +} + +/// Consecutive members of one width as [first, count]; a stack of one width is one run. +pub fn runs(ws: []const Width, out: [][2]usize) [][2]usize { + var n: usize = 0; + for (ws, 0..) |w, i| { + if (n > 0 and std.meta.eql(ws[out[n - 1][0]], w)) { + out[n - 1][1] += 1; + } else { + out[n] = .{ i, 1 }; + n += 1; + } + } + return out[0..n]; +} + +pub const Widths = struct { + allocator: std.mem.Allocator, + layers: [][members]Width, + embed: Width, + head: Width, + + /// Resolves only loaded modules, so draft-head (MTP) and vision overrides never decide admission. + pub fn resolve(a: std.mem.Allocator, f: *const affine.Formats, c: cfg.Config) !Widths { + if ((try f.resolve(null)) == null) return error.MissingAffineFormat; + const layers = try a.alloc([members]Width, c.layers); + errdefer a.free(layers); + var name: [96]u8 = undefined; + for (layers, 0..) |*l, i| { + l.* = @splat(uniform); + for (stacks(c.kind(i))) |stack| for (stack) |m| { + l[@backingInt(m)] = try one(f, try std.fmt.bufPrint(&name, "model.layers.{d}.{s}", .{ i, leaf(m) })); + }; + } + return .{ .allocator = a, .layers = layers, .embed = try one(f, "model.embed_tokens"), .head = try one(f, "lm_head") }; + } + + pub fn deinit(w: *Widths) void { + w.allocator.free(w.layers); + } + + pub fn of(w: *const Widths, layer: usize, m: Member) Width { + return w.layers[layer][@backingInt(m)]; + } +}; + +fn one(f: *const affine.Formats, path: []const u8) !Width { + const s = (try f.resolve(path)) orelse return error.UnqualifiedQwenFormat; + return .{ .bits = s.bits, .group = @intCast(s.group_size) }; +} + +/// What the chip's projection kernels run: the lane kernels on tensor units, the row kernels before them. +pub const Chip = struct { tensor_units: bool, row_bits: []const u8 = &row.bit_widths }; + +/// The first module admission refused, for the load's error message. +pub const Why = struct { layer: ?usize = null, module: []const u8 = "", width: Width = uniform, reason: anyerror = error.None }; + +pub fn admit(w: *const Widths, c: cfg.Config, chip: Chip, why: *Why) !void { + var buf: [3][2]usize = undefined; + for (0..c.layers) |i| for (stacks(c.kind(i))) |stack| { + var ws: [3]Width = undefined; + for (stack, 0..) |m, j| ws[j] = w.of(i, m); + const list = runs(ws[0..stack.len], &buf); + if (list.len == 1 and isUniform(ws[0])) continue; + for (list) |r| { + var n: usize = 0; + for (stack[r[0]..][0..r[1]]) |m| n += shape(c, m)[0]; + check(chip, n, shape(c, stack[r[0]])[1], ws[r[0]]) catch |err| { + why.* = .{ .layer = i, .module = leaf(stack[r[0]]), .width = ws[r[0]], .reason = err }; + return error.UnqualifiedQwenFormat; + }; + } + }; + glue.checkEmbed(.{ .width = @intCast(c.hidden), .rows = 1, .vocab = @intCast(c.vocab), .bits = w.embed.bits, .group = w.embed.group, .metadata = .bf16 }) catch |err| { + why.* = .{ .module = "embed_tokens", .width = w.embed, .reason = err }; + return error.UnqualifiedQwenFormat; + }; + if (!isUniform(w.head)) check(chip, c.vocab, c.hidden, w.head) catch |err| { + why.* = .{ .module = "lm_head", .width = w.head, .reason = err }; + return error.UnqualifiedQwenFormat; + }; +} + +/// One run's kernel: 4/64 keeps today's choice; other widths take the non-reg lane kernel or the row kernel. +fn check(chip: Chip, n: usize, k: usize, w: Width) !void { + _ = try @import("projection.zig").split(n, k); + // The lane sums and odd word counts assume group 64 off the 4-bit kernel. + if (chip.tensor_units and w.bits != 4 and w.group != 64) return error.UnsupportedLaneGroup; + if (chip.tensor_units) return (lane.Layout{ .n = n, .k = k, .format = .{ .bits = w.bits, .group = w.group }, .precompute_sums = isUniform(w), .reg = isUniform(w) }).validate(); + if (isUniform(w) and n % 8 == 0) return; + if (std.mem.indexOfScalar(u8, chip.row_bits, w.bits) == null) return error.UnsupportedRowBits; + return (row.Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = n, .k = k, .group = w.group }).validate(); +} + +const all_rows = [_]u8{ 2, 3, 4, 5, 6, 8 }; + +fn testConfig(a: std.mem.Allocator, quantization: []const u8) ![]u8 { + return std.fmt.allocPrint(a, + \\{{"model_type":"qwen3_5","text_config":{{"model_type":"qwen3_5_text","hidden_size":5120,"intermediate_size":17408, + \\"num_hidden_layers":64,"num_attention_heads":24,"num_key_value_heads":4,"head_dim":256,"vocab_size":248320, + \\"linear_num_key_heads":16,"linear_num_value_heads":48,"linear_key_head_dim":128,"linear_value_head_dim":128, + \\"linear_conv_kernel_dim":4,"partial_rotary_factor":0.25,"eos_token_id":248044}},"quantization":{s}}} + , .{quantization}); +} + +fn testAdmit(quantization: []const u8, chip: Chip, why: *Why) !void { + const a = std.testing.allocator; + const text = try testConfig(a, quantization); + defer a.free(text); + const c = try cfg.parse(a, text, null); + var f = try affine.Formats.init(a, text); + defer f.deinit(); + var w = try Widths.resolve(a, &f, c); + defer w.deinit(); + return admit(&w, c, chip, why); +} + +/// TensorFold/Qwen3.8-27B-oQ4 override shapes: 4-bit g64, 5-bit GDN, MLP and attention members, one 6-bit v. +const oq4 = + \\{"group_size":64,"bits":4,"mode":"affine", + \\"language_model.model.layers.37.linear_attn.in_proj_a":{"group_size":64,"bits":5}, + \\"language_model.model.layers.36.linear_attn.in_proj_b":{"group_size":64,"bits":5}, + \\"language_model.model.layers.26.linear_attn.in_proj_z":{"group_size":64,"bits":5}, + \\"language_model.model.layers.62.linear_attn.out_proj":{"group_size":64,"bits":5}, + \\"language_model.model.layers.25.mlp.down_proj":{"group_size":64,"bits":5}, + \\"language_model.model.layers.27.self_attn.k_proj":{"group_size":64,"bits":5}, + \\"language_model.model.layers.27.self_attn.o_proj":{"group_size":64,"bits":5}, + \\"language_model.model.layers.27.self_attn.v_proj":{"group_size":64,"bits":6}} +; + +/// The override shapes of TensorFold/Qwen3.8-27B-oQ2: 2-bit g64 with an 8-bit embedding and 3/5/6-bit members. +const oq2 = + \\{"group_size":64,"bits":2,"mode":"affine", + \\"language_model.model.embed_tokens":{"group_size":64,"bits":8}, + \\"language_model.model.layers.50.linear_attn.in_proj_a":{"group_size":64,"bits":3}, + \\"language_model.model.layers.44.linear_attn.in_proj_a":{"group_size":64,"bits":5}, + \\"language_model.model.layers.42.linear_attn.in_proj_b":{"group_size":64,"bits":5}, + \\"language_model.model.layers.12.linear_attn.in_proj_z":{"group_size":64,"bits":5}, + \\"language_model.model.layers.10.mlp.down_proj":{"group_size":64,"bits":5}, + \\"language_model.model.layers.35.self_attn.k_proj":{"group_size":64,"bits":5}, + \\"language_model.model.layers.35.self_attn.v_proj":{"group_size":64,"bits":6}} +; + +test "uniform 4-bit g64 admits on both chip kinds and ignores draft-head and vision overrides" { + var why: Why = .{}; + const q = "{\"bits\":4,\"group_size\":64,\"mtp.layers.0.mlp.gate_proj\":{\"bits\":8,\"group_size\":64},\"vision_tower.blocks.0.attn.qkv\":false}"; + try testAdmit(q, .{ .tensor_units = false }, &why); + try testAdmit(q, .{ .tensor_units = true }, &why); +} + +test "oQ4 and oQ2 formats admit once the row kernels unpack every width they declare" { + var why: Why = .{}; + try testAdmit(oq4, .{ .tensor_units = false, .row_bits = &all_rows }, &why); + try testAdmit(oq2, .{ .tensor_units = false, .row_bits = &all_rows }, &why); +} + +test "a width the row kernels do not unpack refuses naming the module, never guessing four bits" { + var why: Why = .{}; + try std.testing.expectError(error.UnqualifiedQwenFormat, testAdmit(oq4, .{ .tensor_units = false, .row_bits = &.{ 2, 4, 6, 8 } }, &why)); + try std.testing.expectEqual(@as(?usize, 25), why.layer); + try std.testing.expectEqualStrings("mlp.down_proj", why.module); + try std.testing.expectEqual(Width{ .bits = 5, .group = 64 }, why.width); + try std.testing.expectError(error.UnqualifiedQwenFormat, testAdmit(oq2, .{ .tensor_units = false, .row_bits = &.{ 2, 4, 6, 8 } }, &why)); + try std.testing.expectEqual(@as(u8, 5), why.width.bits); +} + +test "a dense projection and an unsupported global group refuse" { + var why: Why = .{}; + try std.testing.expectError(error.UnqualifiedQwenFormat, testAdmit("{\"bits\":4,\"group_size\":64,\"model.layers.3.self_attn.q_proj\":false}", .{ .tensor_units = false }, &why)); + try std.testing.expectError(error.BadAffine, testAdmit("{\"bits\":4,\"group_size\":16}", .{ .tensor_units = false }, &why)); +} + +test "fused members split into runs of one width" { + var buf: [3][2]usize = undefined; + const four = uniform; + const five = Width{ .bits = 5, .group = 64 }; + try std.testing.expectEqualSlices([2]usize, &.{.{ 0, 3 }}, runs(&.{ four, four, four }, &buf)); + try std.testing.expectEqualSlices([2]usize, &.{ .{ 0, 1 }, .{ 1, 2 } }, runs(&.{ five, four, four }, &buf)); + try std.testing.expectEqualSlices([2]usize, &.{ .{ 0, 1 }, .{ 1, 1 }, .{ 2, 1 } }, runs(&.{ four, five, four }, &buf)); +} + +test "resolved widths follow overrides under any checkpoint prefix" { + const a = std.testing.allocator; + const text = try testConfig(a, oq4); + defer a.free(text); + const c = try cfg.parse(a, text, null); + var f = try affine.Formats.init(a, text); + defer f.deinit(); + var w = try Widths.resolve(a, &f, c); + defer w.deinit(); + try std.testing.expectEqual(Width{ .bits = 6, .group = 64 }, w.of(27, .v)); + try std.testing.expectEqual(Width{ .bits = 5, .group = 64 }, w.of(27, .k)); + try std.testing.expectEqual(uniform, w.of(27, .q)); + try std.testing.expectEqual(uniform, w.head); +} From ca919c0e81f76d0d682bfee127b91bc91ab70219 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:33:31 -1000 Subject: [PATCH 08/16] The 27B loads mixed-width oQ checkpoints: each projection carries its width, mixed fused groups run per width Uniform 4-bit g64 keeps its bytes, kernels and launches. Other widths take the row kernel before tensor units and the non-reg lane kernel on them; fused members of different widths become runs, each its own product into the fused rows, one row a launch until the core output stride lands. The embedding takes the checkpoint's width. The Neural Engine share declines unless every MLP matrix is 4-bit g64, and the DFlash2 draft head refuses a non-4-bit embedding or head for now. Synthetic graphs load at uniform 4-bit g64. --- .../families/qwen27/dflash/runtime_head.zig | 2 +- .../families/qwen27/dflash/runtime_model.zig | 1 + zig/src/families/qwen27/glue_gpu.zig | 2 +- zig/src/families/qwen27/gpu_weights.zig | 115 ++++++++++++---- zig/src/families/qwen27/model.zig | 31 +++-- zig/src/families/qwen27/projection.zig | 123 ++++++++++++++---- zig/src/families/qwen27/quant_gpu.zig | 27 +++- zig/src/families/qwen27/widths.zig | 7 + zig/tests/qwen27/synthetic_graph.zig | 4 +- 9 files changed, 238 insertions(+), 74 deletions(-) diff --git a/zig/src/families/qwen27/dflash/runtime_head.zig b/zig/src/families/qwen27/dflash/runtime_head.zig index f758fa5c65..0601e3cc2e 100644 --- a/zig/src/families/qwen27/dflash/runtime_head.zig +++ b/zig/src/families/qwen27/dflash/runtime_head.zig @@ -15,7 +15,7 @@ pub const Head = struct { errdefer head.deinit(); if (source.n != 248320) return head; if (source.tile != 32 or source.k % 64 != 0) return error.BadDraftHead; - if (source.raw) |raw| if (raw.bits != 4 or raw.group != 64) return error.BadDraftHead; + if (source.bits != 4 or source.group != 64 or source.runs.len > 0) return error.BadDraftHead; const k: usize = source.k; const input_words = if (source.raw) |raw| raw.w else source.words.buffer; const input_offset = if (source.raw) |raw| raw.w_off else source.words.offset; diff --git a/zig/src/families/qwen27/dflash/runtime_model.zig b/zig/src/families/qwen27/dflash/runtime_model.zig index 8bbdfab534..b75378458f 100644 --- a/zig/src/families/qwen27/dflash/runtime_model.zig +++ b/zig/src/families/qwen27/dflash/runtime_model.zig @@ -132,6 +132,7 @@ pub const Model = struct { }; pub fn loadGraph(a: std.mem.Allocator, b: *Backend) !weights.Graph { const target = b.target; + if (!target.weights.uniformEnds()) return error.DraftNeedsUniformHead; // the borrowed embedding and head are 4-bit g64 const prefix: []const u8 = if (target.checkpoint.has("language_model.model.embed_tokens.weight")) "language_model." else ""; var arena = std.heap.ArenaAllocator.init(a); defer arena.deinit(); diff --git a/zig/src/families/qwen27/glue_gpu.zig b/zig/src/families/qwen27/glue_gpu.zig index 65a4b3b890..857a18944e 100644 --- a/zig/src/families/qwen27/glue_gpu.zig +++ b/zig/src/families/qwen27/glue_gpu.zig @@ -45,7 +45,7 @@ pub const Kernels = struct { errdefer span.cancel(); const e = span.encoder; - const a = abi.Embed{ .width = @intCast(c.hidden), .rows = rows, .vocab = @intCast(c.vocab), .bits = 4, .group = 64, .metadata = .bf16 }; + const a = abi.Embed{ .width = @intCast(c.hidden), .rows = rows, .vocab = @intCast(c.vocab), .bits = p.bits, .group = p.group, .metadata = .bf16 }; try abi.checkEmbed(a); e.setPipeline(k.embed); e.setBuffer(ids.buffer, ids.offset, 0); diff --git a/zig/src/families/qwen27/gpu_weights.zig b/zig/src/families/qwen27/gpu_weights.zig index b5642b8ea9..dd77c66eac 100644 --- a/zig/src/families/qwen27/gpu_weights.zig +++ b/zig/src/families/qwen27/gpu_weights.zig @@ -5,6 +5,8 @@ const cp = @import("core").checkpoint_metal; const Config = @import("config.zig").Config; const projection = @import("projection.zig"); const lane = @import("core").lane_projection; +const widths = @import("widths.zig"); +const Width = widths.Width; const Ref = projection.Ref; const Linear = projection.Linear; pub const Tensor = cp.Tensor; @@ -20,14 +22,32 @@ pub const Weights = struct { norm: Tensor, head: Linear, owned: std.ArrayList(mtl.Buffer) = .empty, + runs: std.ArrayList([]projection.Run) = .empty, pub fn deinit(w: *Weights) void { for (w.owned.items) |b| b.deinit(); w.owned.deinit(w.allocator); + for (w.runs.items) |r| w.allocator.free(r); + w.runs.deinit(w.allocator); w.allocator.free(w.layers); } + + /// Whether every MLP projection is 4-bit g64 in one matrix (the Neural Engine share's format). + pub fn uniformMlp(w: *const Weights) bool { + for (w.layers) |l| for ([_]Linear{ l.gu, l.down }) |p| if (!uniform(p)) return false; + return true; + } + + /// Whether the embedding and head are 4-bit g64 (the format the draft head borrows). + pub fn uniformEnds(w: *const Weights) bool { + return w.embed.bits == 4 and w.embed.group == 64 and uniform(w.head); + } }; +fn uniform(p: Linear) bool { + return p.bits == 4 and p.group == 64 and p.runs.len == 0; +} + /// Modules the graph reads only through repacked copies: their tensors stay in the files until load stages them. const repacked_modules = [_][]const u8{ ".mlp.gate_proj", ".mlp.up_proj", ".mlp.down_proj", ".linear_attn.in_proj_qkv", ".linear_attn.in_proj_z", ".linear_attn.in_proj_b", ".linear_attn.in_proj_a", ".linear_attn.out_proj", ".self_attn.q_proj", ".self_attn.k_proj", ".self_attn.v_proj", ".self_attn.o_proj" }; @@ -49,6 +69,8 @@ const Builder = struct { prefix: []const u8, owned: *std.ArrayList(mtl.Buffer), stage: *cp.Stage, + runs: *std.ArrayList([]projection.Run), + widths: *const widths.Widths, fn get(b: Builder, suffix: []const u8) !Tensor { const name = try std.mem.concat(b.allocator, u8, &.{ b.prefix, suffix }); @@ -60,18 +82,19 @@ const Builder = struct { return checked(try b.get(suffix), dtype, shape); } - fn quant(b: Builder, module: []const u8, n: usize, k: usize) !projection.Triple { + fn quant(b: Builder, module: []const u8, n: usize, k: usize, w: Width) !projection.Triple { + if (k * w.bits % 32 != 0 or k % w.group != 0) return error.BadNativeTensor; const suffixes = [_][]const u8{ ".weight", ".scales", ".biases" }; var tensors: [3]Tensor = undefined; for (suffixes, &tensors, 0..) |suffix, *t, i| { const name = try std.mem.concat(b.allocator, u8, &.{ module, suffix }); defer b.allocator.free(name); - t.* = try b.tensor(name, if (i == 0) .u32 else .bf16, &.{ n, if (i == 0) k / 8 else k / 64 }); + t.* = try b.tensor(name, if (i == 0) .u32 else .bf16, &.{ n, if (i == 0) k * w.bits / 32 else k / w.group }); } - return .{ .weight = tensors[0], .scales = tensors[1], .biases = tensors[2] }; + return .{ .weight = tensors[0], .scales = tensors[1], .biases = tensors[2], .bits = w.bits, .group = w.group }; } - /// `parts` (stacked in order, each `n` rows of `k`) staged from the files and repacked into one owned Linear. + /// `parts` (stacked in order, `n` rows of `k` each, at their widths) staged from the files and repacked. fn linear(b: Builder, parts: []const Part, k: usize) !Linear { var arena = std.heap.ArenaAllocator.init(b.allocator); defer arena.deinit(); @@ -83,12 +106,41 @@ const Builder = struct { defer b.stage.reset(); var triples: [3]projection.Triple = undefined; if (parts.len > triples.len) return error.BadProjectionShape; - for (parts, triples[0..parts.len], 0..) |p, *t, i| t.* = .{ - .weight = try checked(try b.stage.read(b.checkpoint, names[i * 3]), .u32, &.{ p.n, k / 8 }), - .scales = try checked(try b.stage.read(b.checkpoint, names[i * 3 + 1]), .bf16, &.{ p.n, k / 64 }), - .biases = try checked(try b.stage.read(b.checkpoint, names[i * 3 + 2]), .bf16, &.{ p.n, k / 64 }), + for (parts, triples[0..parts.len], 0..) |p, *t, i| { + if (k * p.w.bits % 32 != 0 or k % p.w.group != 0) return error.BadNativeTensor; + t.* = .{ + .weight = try checked(try b.stage.read(b.checkpoint, names[i * 3]), .u32, &.{ p.n, k * p.w.bits / 32 }), + .scales = try checked(try b.stage.read(b.checkpoint, names[i * 3 + 1]), .bf16, &.{ p.n, k / p.w.group }), + .biases = try checked(try b.stage.read(b.checkpoint, names[i * 3 + 2]), .bf16, &.{ p.n, k / p.w.group }), + .bits = p.w.bits, + .group = p.w.group, + }; + } + return b.stack(triples[0..parts.len]); + } + + /// Fused members of one width take one matrix as before; mixed widths become one matrix per run of a width. + fn stack(b: Builder, parts: []const projection.Triple) !Linear { + var ws: [3]Width = undefined; + for (parts, 0..) |p, i| ws[i] = .{ .bits = p.bits, .group = @intCast(p.group) }; + var buf: [3][2]usize = undefined; + const list = widths.runs(ws[0..parts.len], &buf); + if (list.len == 1) return projection.prepare(b.device, b.allocator, b.owned, parts, 32); + const runs = try b.allocator.alloc(projection.Run, list.len); + b.runs.append(b.allocator, runs) catch |err| { + b.allocator.free(runs); + return err; }; - return projection.prepare(b.device, b.allocator, b.owned, triples[0..parts.len], 32); + var first: u32 = 0; + for (list, runs) |r, *run| { + run.* = .{ .linear = try projection.prepare(b.device, b.allocator, b.owned, parts[r[0]..][0..r[1]], 32), .first = first }; + first += run.linear.n; + } + var fused = runs[0].linear; + fused.n = first; + fused.raw = null; + fused.runs = runs; + return fused; } fn layer(b: Builder, c: Config, i: usize) !Layer { @@ -101,31 +153,33 @@ const Builder = struct { return std.mem.concat(allocator, u8, &.{ p, leaf }); } }; + const w = b.widths.layers[i]; + const M = widths.Member; const mixer: @FieldType(Layer, "mixer") = switch (c.kind(i)) { .linear => .{ .linear = .{ - .qkv = try b.linear(&.{.{ .module = try Name.full(a, prefix, "linear_attn.in_proj_qkv"), .n = c.gdnQkvDim() }}, c.hidden), - .zba = try b.linear(&.{ .{ .module = try Name.full(a, prefix, "linear_attn.in_proj_z"), .n = c.gdnValueDim() }, .{ .module = try Name.full(a, prefix, "linear_attn.in_proj_b"), .n = c.v_heads }, .{ .module = try Name.full(a, prefix, "linear_attn.in_proj_a"), .n = c.v_heads } }, c.hidden), - .out = try b.linear(&.{.{ .module = try Name.full(a, prefix, "linear_attn.out_proj"), .n = c.hidden }}, c.gdnValueDim()), + .qkv = try b.linear(&.{.{ .module = try Name.full(a, prefix, "linear_attn.in_proj_qkv"), .n = c.gdnQkvDim(), .w = w[@backingInt(M.qkv)] }}, c.hidden), + .zba = try b.linear(&.{ .{ .module = try Name.full(a, prefix, "linear_attn.in_proj_z"), .n = c.gdnValueDim(), .w = w[@backingInt(M.z)] }, .{ .module = try Name.full(a, prefix, "linear_attn.in_proj_b"), .n = c.v_heads, .w = w[@backingInt(M.b)] }, .{ .module = try Name.full(a, prefix, "linear_attn.in_proj_a"), .n = c.v_heads, .w = w[@backingInt(M.a)] } }, c.hidden), + .out = try b.linear(&.{.{ .module = try Name.full(a, prefix, "linear_attn.out_proj"), .n = c.hidden, .w = w[@backingInt(M.out)] }}, c.gdnValueDim()), .conv = try b.tensor(try Name.full(a, prefix, "linear_attn.conv1d.weight"), .bf16, &.{ c.gdnQkvDim(), c.conv_kernel, 1 }), .a_log = try b.tensor(try Name.full(a, prefix, "linear_attn.A_log"), .bf16, &.{c.v_heads}), .dt = try b.tensor(try Name.full(a, prefix, "linear_attn.dt_bias"), .bf16, &.{c.v_heads}), .norm = try b.tensor(try Name.full(a, prefix, "linear_attn.norm.weight"), .bf16, &.{c.dv}), } }, .attention => .{ .attention = .{ - .q = try b.linear(&.{.{ .module = try Name.full(a, prefix, "self_attn.q_proj"), .n = c.qDim() }}, c.hidden), - .kv = try b.linear(&.{ .{ .module = try Name.full(a, prefix, "self_attn.k_proj"), .n = c.kvDim() }, .{ .module = try Name.full(a, prefix, "self_attn.v_proj"), .n = c.kvDim() } }, c.hidden), - .out = try b.linear(&.{.{ .module = try Name.full(a, prefix, "self_attn.o_proj"), .n = c.hidden }}, c.heads * c.head_dim), + .q = try b.linear(&.{.{ .module = try Name.full(a, prefix, "self_attn.q_proj"), .n = c.qDim(), .w = w[@backingInt(M.q)] }}, c.hidden), + .kv = try b.linear(&.{ .{ .module = try Name.full(a, prefix, "self_attn.k_proj"), .n = c.kvDim(), .w = w[@backingInt(M.k)] }, .{ .module = try Name.full(a, prefix, "self_attn.v_proj"), .n = c.kvDim(), .w = w[@backingInt(M.v)] } }, c.hidden), + .out = try b.linear(&.{.{ .module = try Name.full(a, prefix, "self_attn.o_proj"), .n = c.hidden, .w = w[@backingInt(M.o)] }}, c.heads * c.head_dim), .q_norm = try b.tensor(try Name.full(a, prefix, "self_attn.q_norm.weight"), .bf16, &.{c.head_dim}), .k_norm = try b.tensor(try Name.full(a, prefix, "self_attn.k_norm.weight"), .bf16, &.{c.head_dim}), } }, }; - const gu = try b.linear(&.{ .{ .module = try Name.full(a, prefix, "mlp.gate_proj"), .n = c.intermediate }, .{ .module = try Name.full(a, prefix, "mlp.up_proj"), .n = c.intermediate } }, c.hidden); - const down = try b.linear(&.{.{ .module = try Name.full(a, prefix, "mlp.down_proj"), .n = c.hidden }}, c.intermediate); + const gu = try b.linear(&.{ .{ .module = try Name.full(a, prefix, "mlp.gate_proj"), .n = c.intermediate, .w = w[@backingInt(M.gate)] }, .{ .module = try Name.full(a, prefix, "mlp.up_proj"), .n = c.intermediate, .w = w[@backingInt(M.up)] } }, c.hidden); + const down = try b.linear(&.{.{ .module = try Name.full(a, prefix, "mlp.down_proj"), .n = c.hidden, .w = w[@backingInt(M.down)] }}, c.intermediate); return .{ .norm = try b.tensor(try Name.full(a, prefix, "input_layernorm.weight"), .bf16, &.{c.hidden}), .post_norm = try b.tensor(try Name.full(a, prefix, "post_attention_layernorm.weight"), .bf16, &.{c.hidden}), .gu = gu, .down = down, .mixer = mixer }; } }; -const Part = struct { module: []const u8, n: usize }; +const Part = struct { module: []const u8, n: usize, w: Width }; fn checked(t: Tensor, dtype: cp.DType, shape: []const usize) !Tensor { if (t.dtype != dtype or t.rank != shape.len or !std.mem.eql(usize, t.shape[0..t.rank], shape)) return error.BadNativeTensor; @@ -133,7 +187,7 @@ fn checked(t: Tensor, dtype: cp.DType, shape: []const usize) !Tensor { } /// The graph from `checkpoint`, whose `repacked` tensors stay in their files: each Linear is staged, repacked, dropped. -pub fn load(allocator: std.mem.Allocator, device: mtl.Device, checkpoint: *const cp.Checkpoint, c: Config) !Weights { +pub fn load(allocator: std.mem.Allocator, device: mtl.Device, checkpoint: *const cp.Checkpoint, c: Config, ws: *const widths.Widths) !Weights { const prefix: []const u8 = if (checkpoint.has("language_model.model.embed_tokens.weight")) "language_model." else if (checkpoint.has("model.embed_tokens.weight")) "" else return error.MissingTextNamespace; if (checkpoint.has("language_model.model.embed_tokens.weight") and checkpoint.has("model.embed_tokens.weight")) return error.AmbiguousTextNamespace; var owned: std.ArrayList(mtl.Buffer) = .empty; @@ -141,17 +195,22 @@ pub fn load(allocator: std.mem.Allocator, device: mtl.Device, checkpoint: *const for (owned.items) |b| b.deinit(); owned.deinit(allocator); } + var runs: std.ArrayList([]projection.Run) = .empty; + errdefer { + for (runs.items) |r| allocator.free(r); + runs.deinit(allocator); + } const layers = try allocator.alloc(Layer, c.layers); errdefer allocator.free(layers); var stage = cp.Stage{ .device = device }; defer stage.deinit(); - const b = Builder{ .allocator = allocator, .device = device, .checkpoint = checkpoint, .prefix = prefix, .owned = &owned, .stage = &stage }; + const b = Builder{ .allocator = allocator, .device = device, .checkpoint = checkpoint, .prefix = prefix, .owned = &owned, .stage = &stage, .runs = &runs, .widths = ws }; for (layers, 0..) |*l, i| l.* = try b.layer(c, i); - const embed = try b.quant("model.embed_tokens", c.vocab, c.hidden); + const embed = try b.quant("model.embed_tokens", c.vocab, c.hidden, ws.embed); const norm = try b.tensor("model.norm.weight", .bf16, &.{c.hidden}); - const head = try b.linear(&.{.{ .module = "lm_head", .n = c.vocab }}, c.hidden); + const head = try b.linear(&.{.{ .module = "lm_head", .n = c.vocab, .w = ws.head }}, c.hidden); if (device.tensorUnits()) try laneOrder(allocator, device, layers, head); - return .{ .allocator = allocator, .layers = layers, .embed = embed, .norm = norm, .head = head, .owned = owned }; + return .{ .allocator = allocator, .layers = layers, .embed = embed, .norm = norm, .head = head, .owned = owned, .runs = runs }; } /// Tensor-unit projections leave load in lane_projection.reg_layout, the order quant_gpu's reg kernel reads. @@ -163,12 +222,18 @@ fn laneOrder(a: std.mem.Allocator, device: mtl.Device, layers: []const Layer, he .linear => |g| [_]Linear{ g.qkv, g.zba, g.out, l.gu, l.down }, .attention => |t| [_]Linear{ t.q, t.kv, t.out, l.gu, l.down }, }; - for (list) |p| try jobs.append(a, .{ .w = .{ .buf = p.words.buffer, .off = p.words.offset }, .n = p.n, .k = p.k }); + for (list) |p| try regJobs(a, &jobs, p); } - try jobs.append(a, .{ .w = .{ .buf = head.words.buffer, .off = head.words.offset }, .n = head.n, .k = head.k }); + try regJobs(a, &jobs, head); try lane.regOrder(device, jobs.items); } +/// Only 4-bit g64 matrices take the reg kernel; other widths keep lane_projection.weight_layout. +fn regJobs(a: std.mem.Allocator, jobs: *std.ArrayList(lane.Order), p: Linear) std.mem.Allocator.Error!void { + for (p.runs) |r| try regJobs(a, jobs, r.linear); + if (p.runs.len == 0 and p.bits == 4 and p.group == 64) try jobs.append(a, .{ .w = .{ .buf = p.words.buffer, .off = p.words.offset }, .n = p.n, .k = p.k }); +} + pub fn ref(t: Tensor) Ref { return .{ .buffer = t.buffer, .offset = t.offset }; } diff --git a/zig/src/families/qwen27/model.zig b/zig/src/families/qwen27/model.zig index 3a45318ef2..ea9f06a4de 100644 --- a/zig/src/families/qwen27/model.zig +++ b/zig/src/families/qwen27/model.zig @@ -4,6 +4,7 @@ const mtl = @import("metal"); const core = @import("core"); const cfg = @import("config.zig"); const affine = @import("affine.zig"); +const widths = @import("widths.zig"); const wts = @import("gpu_weights.zig"); const quant = @import("quant_gpu.zig"); const glue = @import("glue_gpu.zig"); @@ -33,7 +34,11 @@ pub const Model = struct { const text = try std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(16 << 20)); defer a.free(text); const c = try cfg.parse(a, text, null); - try supported(a, c, text); + try supported(c); + var formats = try affine.Formats.init(a, text); + defer formats.deinit(); + var ws = try widths.Widths.resolve(a, &formats, c); + defer ws.deinit(); const m = try a.create(Model); errdefer a.destroy(m); m.allocator = a; @@ -41,6 +46,11 @@ pub const Model = struct { m.device = try mtl.Device.init(); errdefer m.device.deinit(); if (!qualifiedGeneration(m.device.generation())) return error.QwenMetalQualificationRequired; + var why: widths.Why = .{}; + widths.admit(&ws, c, .{ .tensor_units = m.device.tensorUnits() }, &why) catch |err| { + std.log.err("qwen27: {s} (layer {?d}) is {d}-bit group {d}; this chip's projection kernels do not run it ({s})", .{ why.module, why.layer, why.width.bits, why.width.group, @errorName(why.reason) }); + return err; + }; m.queue = try m.device.queue(); errdefer m.queue.deinit(); m.checkpoint = core.checkpoint_metal.Checkpoint.init(a); @@ -49,7 +59,7 @@ pub const Model = struct { defer core.checkpoint_host.freeShardFiles(a, files); // projections stay in the files until load repacks them: the weights live once, in their lane layouts for (files) |file| try m.checkpoint.addFileSplit(m.device, file, "", "language_model.", wts.repacked); - m.weights = try wts.load(a, m.device, &m.checkpoint, c); + m.weights = try wts.load(a, m.device, &m.checkpoint, c, &ws); errdefer m.weights.deinit(); m.kernels = .{ .allocator = a, .device = m.device }; errdefer m.kernels.deinit(); @@ -59,7 +69,7 @@ pub const Model = struct { errdefer m.frame.deinit(); m.ane = null; const columns = if (std.c.getenv("TF_ANE_COLUMNS")) |v| try std.fmt.parseInt(u32, std.mem.span(v), 10) else ane_columns; - const share = columns > 0 and m.device.tensorUnits() and mtl.ane.available(); + const share = columns > 0 and m.device.tensorUnits() and mtl.ane.available() and m.weights.uniformMlp(); const fits = if (share_fits) |f| !share or f(ane_mlp.Share.bytes(c, columns)) else true; if (!fits) std.log.info("qwen27: the Neural Engine's prompt share ({d} MiB of fp16 weights) would take this Mac past its load limit; prompts run on the GPU alone", .{ane_mlp.Share.bytes(c, columns) >> 20}); if (share and fits) m.ane = ane_mlp.Share.init(a, m.device, m.queue, &m.checkpoint, c, columns) catch |err| blk: { @@ -83,21 +93,10 @@ pub const Model = struct { } }; -pub fn supported(a: std.mem.Allocator, c: cfg.Config, text: []const u8) !void { +/// The qualified shape and layer schedule; per-module widths are admitted by widths.admit once the chip is known. +pub fn supported(c: cfg.Config) !void { if (c.hidden != 5120 or c.intermediate != 17408 or c.layers != 64 or c.vocab != 248320 or c.heads != 24 or c.kv_heads != 4 or c.head_dim != 256 or c.k_heads != 16 or c.v_heads != 48 or c.dk != 128 or c.dv != 128 or c.conv_kernel != 4 or c.rope_dims != 64) return error.UnqualifiedQwenShape; for (c.kinds[0..c.layers], 0..) |kind, i| if (kind != (if ((i + 1) % 4 == 0) cfg.Kind.attention else cfg.Kind.linear)) return error.UnqualifiedLayerSchedule; - var formats = try affine.Formats.init(a, text); - defer formats.deinit(); - const declared = (try formats.resolve(null)) orelse return error.MissingAffineFormat; - if (declared.bits != 4 or declared.group_size != 64) return error.UnqualifiedQwenFormat; - if (formats.quant) |q| { - var it = q.iterator(); - while (it.next()) |item| { - if (std.mem.indexOfScalar(u8, item.key_ptr.*, '.') == null) continue; - const own = (try formats.resolve(item.key_ptr.*)) orelse return error.UnqualifiedQwenFormat; - if (own.bits != 4 or own.group_size != 64) return error.UnqualifiedQwenFormat; - } - } } fn qualifiedGeneration(generation: u32) bool { diff --git a/zig/src/families/qwen27/projection.zig b/zig/src/families/qwen27/projection.zig index a59eacd43f..cc52ed1d57 100644 --- a/zig/src/families/qwen27/projection.zig +++ b/zig/src/families/qwen27/projection.zig @@ -3,10 +3,14 @@ const std = @import("std"); const mtl = @import("metal"); const checkpoint = @import("core").checkpoint_metal; const row = @import("core").row_projection; +const lane = @import("core").lane_projection; pub const Ref = struct { buffer: mtl.Buffer, offset: usize = 0 }; -pub const Linear = struct { words: Ref, pairs: Ref, n: u32, k: u32, tile: u32, slices: u32, raw: ?row.Weights = null }; -pub const Triple = struct { weight: checkpoint.Tensor, scales: checkpoint.Tensor, biases: checkpoint.Tensor }; +/// A device matrix; `runs` non-empty means fused members of different widths, each its own product into these rows. +pub const Linear = struct { words: Ref, pairs: Ref, n: u32, k: u32, tile: u32, slices: u32, raw: ?row.Weights = null, bits: u8 = 4, group: u32 = 64, runs: []const Run = &.{} }; +/// Members of one width inside a fused matrix, writing output columns [first, first + linear.n). +pub const Run = struct { linear: Linear, first: u32 }; +pub const Triple = struct { weight: checkpoint.Tensor, scales: checkpoint.Tensor, biases: checkpoint.Tensor, bits: u8 = 4, group: u32 = 64 }; pub fn split(n: usize, k: usize) !u32 { if (n == 0 or k == 0 or n % 4 != 0 or k % 64 != 0) return error.BadProjectionShape; @@ -31,17 +35,22 @@ pub fn promptSplit(n: usize, k: usize) ?u32 { pub fn prepare(device: mtl.Device, allocator: std.mem.Allocator, owned: *std.ArrayList(mtl.Buffer), parts: []const Triple, tile: u32) !Linear { if (tile != 32 or parts.len == 0) return error.BadProjectionShape; - const k = parts[0].weight.shape[1] * 8; + const bits = parts[0].bits; + const group = parts[0].group; + if (bits == 0 or bits > 8 or group == 0) return error.BadAffineMetadata; + const k = parts[0].weight.shape[1] * 32 / bits; var slices = try split(parts[0].weight.shape[0], k); var n: usize = 0; for (parts) |p| { const rows = p.weight.shape[0]; - if (p.weight.rank != 2 or p.weight.dtype != .u32 or p.weight.shape[1] * 8 != k or rows % 4 != 0 or try split(rows, k) != slices) return error.BadProjectionShape; - if (p.scales.dtype != .bf16 or p.biases.dtype != .bf16 or p.scales.rank != 2 or p.biases.rank != 2 or p.scales.shape[0] != rows or p.biases.shape[0] != rows or p.scales.shape[1] != k / 64 or p.biases.shape[1] != k / 64) return error.BadAffineMetadata; + if (p.bits != bits or p.group != group or k % group != 0) return error.BadAffineMetadata; + if (p.weight.rank != 2 or p.weight.dtype != .u32 or p.weight.shape[1] * 32 != k * bits or rows % 4 != 0 or try split(rows, k) != slices) return error.BadProjectionShape; + if (p.scales.dtype != .bf16 or p.biases.dtype != .bf16 or p.scales.rank != 2 or p.biases.rank != 2 or p.scales.shape[0] != rows or p.biases.shape[0] != rows or p.scales.shape[1] != k / group or p.biases.shape[1] != k / group) return error.BadAffineMetadata; n = try std.math.add(usize, n, rows); } - if (n % tile != 0) return error.BadProjectionShape; if (!device.tensorUnits()) return prepareRows(device, allocator, owned, parts, n, k, tile, slices); + if (n % tile != 0) return error.BadProjectionShape; + if (bits != 4 or group != 64) return prepareLanes(device, allocator, owned, parts, n, k, tile, slices); slices = m5Split(n, k) orelse slices; const w = try device.buffer(n * k / 2, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); owned.append(allocator, w) catch |err| { @@ -83,30 +92,92 @@ test "stacked gate-up keeps the member K split rather than recomputing it from s fn prepareRows(device: mtl.Device, allocator: std.mem.Allocator, owned: *std.ArrayList(mtl.Buffer), parts: []const Triple, n: usize, k: usize, tile: u32, slices: u32) !Linear { const options = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; - const w = try device.buffer(n * k / 2, options); - owned.append(allocator, w) catch |err| { - w.deinit(); - return err; - }; - const scale = try device.buffer(n * (k / 64) * 2, options); - owned.append(allocator, scale) catch |err| { - scale.deinit(); - return err; - }; - const bias = try device.buffer(n * (k / 64) * 2, options); - owned.append(allocator, bias) catch |err| { - bias.deinit(); - return err; - }; + const bits: usize = parts[0].bits; + const ng = k / parts[0].group; + const row_bytes = k * bits / 8; + const w = try own(device, allocator, owned, n * row_bytes, options); + const scale = try own(device, allocator, owned, n * ng * 2, options); + const bias = try own(device, allocator, owned, n * ng * 2, options); var base: usize = 0; for (parts) |p| { const rows = p.weight.shape[0]; - @memcpy(w.contents()[base * k / 2 .. (base + rows) * k / 2], std.mem.sliceAsBytes(p.weight.host(u32))); - @memcpy(scale.contents()[base * (k / 64) * 2 .. (base + rows) * (k / 64) * 2], std.mem.sliceAsBytes(p.scales.host(u16))); - @memcpy(bias.contents()[base * (k / 64) * 2 .. (base + rows) * (k / 64) * 2], std.mem.sliceAsBytes(p.biases.host(u16))); + @memcpy(w.contents()[base * row_bytes .. (base + rows) * row_bytes], std.mem.sliceAsBytes(p.weight.host(u32))); + @memcpy(scale.contents()[base * ng * 2 .. (base + rows) * ng * 2], std.mem.sliceAsBytes(p.scales.host(u16))); + @memcpy(bias.contents()[base * ng * 2 .. (base + rows) * ng * 2], std.mem.sliceAsBytes(p.biases.host(u16))); base += rows; } - const raw = row.Weights{ .w = w, .scales = scale, .biases = bias, .n = n, .k = k, .group = 64, .bits = 4, .sum = .f32 }; + const raw = row.Weights{ .w = w, .scales = scale, .biases = bias, .n = n, .k = k, .group = parts[0].group, .bits = parts[0].bits, .sum = .f32 }; try raw.validate(); - return .{ .words = .{ .buffer = w }, .pairs = .{ .buffer = scale }, .n = @intCast(n), .k = @intCast(k), .tile = tile, .slices = slices, .raw = raw }; + return .{ .words = .{ .buffer = w }, .pairs = .{ .buffer = scale }, .n = @intCast(n), .k = @intCast(k), .tile = tile, .slices = slices, .raw = raw, .bits = parts[0].bits, .group = parts[0].group }; +} + +/// Widths other than 4/64 on tensor units: lane_projection.weight_layout byte copies for the non-reg lane kernel. +fn prepareLanes(device: mtl.Device, allocator: std.mem.Allocator, owned: *std.ArrayList(mtl.Buffer), parts: []const Triple, n: usize, k: usize, tile: u32, slices: u32) !Linear { + const options = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; + const layout = lane.Layout{ .n = n, .k = k, .format = .{ .bits = parts[0].bits, .group = parts[0].group }, .sk = slices }; + try layout.validate(); + const w = try own(device, allocator, owned, layout.weightWords() * 4, options); + const sb = try own(device, allocator, owned, layout.metadataElements() * 2, options); + var base: usize = 0; + for (parts) |p| { + packTiles(.{ .words = p.weight.host(u32), .scales = p.scales.host(u16), .biases = p.biases.host(u16), .rows = p.weight.shape[0] }, base, n, k, parts[0].bits, parts[0].group, w.slice(u32, layout.weightWords()), sb.slice(u16, layout.metadataElements())); + base += p.weight.shape[0]; + } + return .{ .words = .{ .buffer = w }, .pairs = .{ .buffer = sb }, .n = @intCast(n), .k = @intCast(k), .tile = tile, .slices = slices, .bits = parts[0].bits, .group = parts[0].group }; +} + +/// One member as stored: words [rows][k*bits/32], scales and biases [rows][k/group]. +pub const Member = struct { words: []const u32, scales: []const u16, biases: []const u16, rows: usize }; + +/// Copies a member into lane_projection.weight_layout rows [base, base + rows) of an n-row stack, bytes unchanged. +pub fn packTiles(p: Member, base: usize, n: usize, k: usize, bits: usize, group: usize, out_w: []u32, out_sb: []u16) void { + const ng = k / group; + const wp = group * bits / 32; + for (0..p.rows) |r| for (0..ng) |g| { + const at = base + r; + @memcpy(out_w[((at / 32 * ng + g) * 32 + at % 32) * wp ..][0..wp], p.words[(r * ng + g) * wp ..][0..wp]); + out_sb[(g * n + at) * 2] = p.scales[r * ng + g]; + out_sb[(g * n + at) * 2 + 1] = p.biases[r * ng + g]; + }; +} + +test "a stack packed member by member equals the core lane pack of the whole matrix" { + const a = std.testing.allocator; + const n = 64; + const k = 256; + var rng = std.Random.DefaultPrng.init(7); + for ([_]u8{ 4, 6, 8 }) |bits| for ([_]usize{ 32, 64, 128 }) |group| { + const l = lane.Layout{ .n = n, .k = k, .format = .{ .bits = bits, .group = group } }; + const w = try a.alloc(u32, l.weightWords()); + defer a.free(w); + const sb = try a.alloc(u16, n * (k / group) * 2); + defer a.free(sb); + rng.random().bytes(std.mem.sliceAsBytes(w)); + rng.random().bytes(std.mem.sliceAsBytes(sb)); + const scales = sb[0 .. n * (k / group)]; + const biases = sb[n * (k / group) ..]; + const want_w = try a.alloc(u32, w.len); + defer a.free(want_w); + const want_sb = try a.alloc(u16, l.metadataElements()); + defer a.free(want_sb); + try lane.pack(l, w, scales, biases, want_w, want_sb); + const got_w = try a.alloc(u32, w.len); + defer a.free(got_w); + const got_sb = try a.alloc(u16, want_sb.len); + defer a.free(got_sb); + const row_words = w.len / n; + const ng = k / group; + for ([_][2]usize{ .{ 0, 32 }, .{ 32, 16 }, .{ 48, 16 } }) |part| packTiles(.{ .words = w[part[0] * row_words ..][0 .. part[1] * row_words], .scales = scales[part[0] * ng ..][0 .. part[1] * ng], .biases = biases[part[0] * ng ..][0 .. part[1] * ng], .rows = part[1] }, part[0], n, k, bits, group, got_w, got_sb); + try std.testing.expectEqualSlices(u32, want_w, got_w); + try std.testing.expectEqualSlices(u16, want_sb, got_sb); + }; +} + +fn own(device: mtl.Device, allocator: std.mem.Allocator, owned: *std.ArrayList(mtl.Buffer), bytes: usize, options: usize) !mtl.Buffer { + const b = try device.buffer(bytes, options); + owned.append(allocator, b) catch |err| { + b.deinit(); + return err; + }; + return b; } diff --git a/zig/src/families/qwen27/quant_gpu.zig b/zig/src/families/qwen27/quant_gpu.zig index 7632937dd2..824a87eb95 100644 --- a/zig/src/families/qwen27/quant_gpu.zig +++ b/zig/src/families/qwen27/quant_gpu.zig @@ -10,7 +10,7 @@ pub const Ref = projection.Ref; /// Group sums the input's producer wrote: XS[group * stride + row], in the lane projection's order. pub const Sums = struct { ref: Ref, stride: u32 }; -const Key = struct { kind: Kind, n: u32 = 0, k: u32, slices: u32 = 0, tile: u32 = 0, tmr: u32 = 0, edge: bool = false, skip: u32 = 0, cut: [2]u32 = .{ 0, 0 } }; +const Key = struct { kind: Kind, n: u32 = 0, k: u32, slices: u32 = 0, tile: u32 = 0, tmr: u32 = 0, edge: bool = false, skip: u32 = 0, cut: [2]u32 = .{ 0, 0 }, bits: u8 = 4, group: u32 = 64 }; pub const Kind = enum { projection, xsum, norm, input_norm, mlp, ane_input, mlp_ane, add_ane }; pub const Kernels = struct { @@ -92,6 +92,11 @@ pub const Kernels = struct { } fn quantImpl(k: *Kernels, input_encoder: mtl.ComputeEncoder, p: projection.Linear, x: Ref, xs: ?Sums, dims: Ref, y: Ref, rows: u32, skip: u32, cut: [2]u32) !void { + if (p.runs.len > 0) return k.quantRuns(input_encoder, p, x, xs, dims, y, rows, skip, cut); + return k.quantOne(input_encoder, p, x, xs, dims, y, rows, skip, cut); + } + + fn quantOne(k: *Kernels, input_encoder: mtl.ComputeEncoder, p: projection.Linear, x: Ref, xs: ?Sums, dims: Ref, y: Ref, rows: u32, skip: u32, cut: [2]u32) !void { if ((skip > 0 or cut[1] > 0) and p.raw != null) return error.BadCoreProjectionRows; if (skip > 0 and (p.raw != null or skip >= p.n / 64)) return error.BadCoreProjectionRows; var span = profile.begin(k.profiler, input_encoder, .dense); @@ -126,13 +131,16 @@ pub const Kernels = struct { try span.end(); return; } - const key = Key{ .kind = .projection, .n = p.n, .k = p.k, .slices = sk, .skip = skip, .cut = cut }; + const reg = p.bits == 4 and p.group == 64; + if (!reg and (skip > 0 or cut[1] > 0)) return error.BadCoreProjectionRows; + const key = Key{ .kind = .projection, .n = p.n, .k = p.k, .slices = if (reg) sk else p.slices, .skip = skip, .cut = cut, .bits = p.bits, .group = p.group }; + const sums: ?Sums = if (reg) xs else null; // the non-reg kernel sums its own groups const half = p.n / 64; const ranges = [_][2]usize{ .{ skip, half - skip }, .{ half + skip, half - skip } }; const entry = try k.lanes.getOrPut(k.allocator, key); if (!entry.found_existing) { errdefer _ = k.lanes.remove(key); - entry.value_ptr.* = try lane.Projection.init(k.allocator, k.device, .{ .n = p.n, .k = p.k, .format = .{ .bits = 4, .group = 64 }, .sk = sk, .ranges = if (skip > 0) &ranges else &.{}, .groups = if (cut[1] > 0) .{ cut[0], cut[1] } else null, .output = if (cut[1] > 0) .f32 else .bf16, .precompute_sums = true, .reg = true }); + entry.value_ptr.* = try lane.Projection.init(k.allocator, k.device, .{ .n = p.n, .k = p.k, .format = .{ .bits = p.bits, .group = p.group }, .sk = key.slices, .ranges = if (skip > 0) &ranges else &.{}, .groups = if (cut[1] > 0) .{ cut[0], cut[1] } else null, .output = if (cut[1] > 0) .f32 else .bf16, .precompute_sums = reg, .reg = reg }); } const step: u32 = if (entry.value_ptr.reg) lane.reg_rows else 16; var first: u32 = 0; @@ -141,11 +149,22 @@ pub const Kernels = struct { const w = lane.Ref{ .buf = p.words.buffer, .off = p.words.offset }; const sb = lane.Ref{ .buf = p.pairs.buffer, .off = p.pairs.offset }; const yr = lane.Ref{ .buf = y.buffer, .off = y.offset + @as(usize, first) * p.n * @as(usize, if (cut[1] > 0) 4 else 2) }; - if (xs) |s| try entry.value_ptr.encodeSums(e, xr, w, sb, .{ .ref = .{ .buf = s.ref.buffer, .off = s.ref.offset + @as(usize, first) * 4 }, .stride = s.stride }, yr, @min(rows - first, step)) else try entry.value_ptr.encode(e, xr, w, sb, yr, @min(rows - first, step)); + if (sums) |s| try entry.value_ptr.encodeSums(e, xr, w, sb, .{ .ref = .{ .buf = s.ref.buffer, .off = s.ref.offset + @as(usize, first) * 4 }, .stride = s.stride }, yr, @min(rows - first, step)) else try entry.value_ptr.encode(e, xr, w, sb, yr, @min(rows - first, step)); } try span.end(); } + /// Mixed-width members: each run its own product into its columns, a row a launch so every row count matches. + fn quantRuns(k: *Kernels, input_encoder: mtl.ComputeEncoder, p: projection.Linear, x: Ref, xs: ?Sums, dims: Ref, y: Ref, rows: u32, skip: u32, cut: [2]u32) !void { + if (skip > 0 or cut[1] > 0 or rows == 0 or rows > 128) return error.BadCoreProjectionRows; + for (p.runs) |run| for (0..rows) |r| { + const xr = Ref{ .buffer = x.buffer, .offset = x.offset + r * p.k * 2 }; + const sr: ?Sums = if (xs) |s| .{ .ref = .{ .buffer = s.ref.buffer, .offset = s.ref.offset + r * 4 }, .stride = s.stride } else null; + const yr = Ref{ .buffer = y.buffer, .offset = y.offset + (r * p.n + run.first) * 2 }; + try k.quantOne(input_encoder, run.linear, xr, sr, dims, yr, 1, 0, .{ 0, 0 }); + }; + } + /// Loads the tiles on first use and checks the scalar twin against them once per shape. fn twin(k: *Kernels, m: row.simd.Matrix) !bool { if (k.simd == null) k.simd = try row.simd.Pipelines.load(k.device); diff --git a/zig/src/families/qwen27/widths.zig b/zig/src/families/qwen27/widths.zig index 69a786ba2c..ffb082842f 100644 --- a/zig/src/families/qwen27/widths.zig +++ b/zig/src/families/qwen27/widths.zig @@ -95,6 +95,13 @@ pub const Widths = struct { return .{ .allocator = a, .layers = layers, .embed = try one(f, "model.embed_tokens"), .head = try one(f, "lm_head") }; } + /// Every module 4-bit g64, for synthetic graphs. + pub fn uniformAll(a: std.mem.Allocator, layers: usize) !Widths { + const all = try a.alloc([members]Width, layers); + for (all) |*l| l.* = @splat(uniform); + return .{ .allocator = a, .layers = all, .embed = uniform, .head = uniform }; + } + pub fn deinit(w: *Widths) void { w.allocator.free(w.layers); } diff --git a/zig/tests/qwen27/synthetic_graph.zig b/zig/tests/qwen27/synthetic_graph.zig index a88aa39654..3da574b7d1 100644 --- a/zig/tests/qwen27/synthetic_graph.zig +++ b/zig/tests/qwen27/synthetic_graph.zig @@ -93,7 +93,9 @@ pub fn createLayers(a: std.mem.Allocator, layers: usize) !*q.model.Model { inline for (.{ "self_attn.q_norm.weight", "self_attn.k_norm.weight" }) |leaf| try builder.tensor(try builder.layerName(i, leaf), &.{256}, .bf16, 1); } } - var weights = try q.gpu_weights.load(a, device, &cp, c); + var widths = try q.widths.Widths.uniformAll(a, c.layers); + defer widths.deinit(); + var weights = try q.gpu_weights.load(a, device, &cp, c, &widths); errdefer weights.deinit(); const glue = try q.glue_gpu.Kernels.init(device); errdefer glue.deinit(); From 5dbbbb7fa1ed0184a01b4dd7b34e2ed8b737a36f Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Sun, 11 Oct 2026 19:05:01 -1000 Subject: [PATCH 09/16] The DFlash2 drafter borrows the target embedding and head at their checkpoint widths The draft vocabulary copies whole rows (or 32-row tiles) of k * bits / 8 bytes and k / group metadata, and the borrowed embedding is validated at its resolved width. The draft-head check covers 2-, 4- and 8-bit heads. --- zig/src/families/qwen27/dflash/runtime_head.zig | 12 ++++++------ zig/src/families/qwen27/dflash/runtime_model.zig | 7 +++---- zig/src/families/qwen27/gpu_weights.zig | 5 ----- zig/tests/qwen27/dflash_head.zig | 8 ++++---- 4 files changed, 13 insertions(+), 19 deletions(-) diff --git a/zig/src/families/qwen27/dflash/runtime_head.zig b/zig/src/families/qwen27/dflash/runtime_head.zig index 0601e3cc2e..be6581a9b7 100644 --- a/zig/src/families/qwen27/dflash/runtime_head.zig +++ b/zig/src/families/qwen27/dflash/runtime_head.zig @@ -14,18 +14,18 @@ pub const Head = struct { var head = Head{ .allocator = a, .linear = source }; errdefer head.deinit(); if (source.n != 248320) return head; - if (source.tile != 32 or source.k % 64 != 0) return error.BadDraftHead; - if (source.bits != 4 or source.group != 64 or source.runs.len > 0) return error.BadDraftHead; + if (source.tile != 32 or source.k % source.group != 0 or source.runs.len > 0) return error.BadDraftHead; const k: usize = source.k; + const row = k * source.bits / 8; // whole 32-row tiles in the lane layout, rows in the raw one const input_words = if (source.raw) |raw| raw.w else source.words.buffer; const input_offset = if (source.raw) |raw| raw.w_off else source.words.offset; - const groups = k / 64; - const words = try head.allocate(device, width * k / 2); + const groups = k / source.group; + const words = try head.allocate(device, width * row); var destination: usize = 0; for (spans) |span| { const rows = span[1] - span[0]; - const first = input_offset + @as(usize, span[0]) * k / 2; - @memcpy(words.contents()[destination * k / 2 ..][0 .. @as(usize, rows) * k / 2], input_words.contents()[first..][0 .. @as(usize, rows) * k / 2]); + const first = input_offset + @as(usize, span[0]) * row; + @memcpy(words.contents()[destination * row ..][0 .. @as(usize, rows) * row], input_words.contents()[first..][0 .. @as(usize, rows) * row]); destination += rows; } head.linear.n = width; diff --git a/zig/src/families/qwen27/dflash/runtime_model.zig b/zig/src/families/qwen27/dflash/runtime_model.zig index b75378458f..dc8bd4db3a 100644 --- a/zig/src/families/qwen27/dflash/runtime_model.zig +++ b/zig/src/families/qwen27/dflash/runtime_model.zig @@ -132,12 +132,11 @@ pub const Model = struct { }; pub fn loadGraph(a: std.mem.Allocator, b: *Backend) !weights.Graph { const target = b.target; - if (!target.weights.uniformEnds()) return error.DraftNeedsUniformHead; // the borrowed embedding and head are 4-bit g64 const prefix: []const u8 = if (target.checkpoint.has("language_model.model.embed_tokens.weight")) "language_model." else ""; var arena = std.heap.ArenaAllocator.init(a); defer arena.deinit(); const local = arena.allocator(); - const embed = try binding(local, target, prefix, "model.embed_tokens"); + const embed = try binding(local, target, prefix, "model.embed_tokens", .{ .bits = target.weights.embed.bits, .group_size = target.weights.embed.group }); var names: std.ArrayList([]const u8) = .empty; var iterator = b.checkpoint.tensors.keyIterator(); while (iterator.next()) |name| try names.append(local, name.*); @@ -145,10 +144,10 @@ pub fn loadGraph(a: std.mem.Allocator, b: *Backend) !weights.Graph { while (later.next()) |name| try names.append(local, name.*); return weights.load(a, .{ .source = b.source(), .names = names.items }, b.config, embed); } -fn binding(a: std.mem.Allocator, target: *Target, prefix: []const u8, module: []const u8) !ck.Linear { +fn binding(a: std.mem.Allocator, target: *Target, prefix: []const u8, module: []const u8, format: @import("../affine.zig").Spec) !ck.Linear { const stem = try std.mem.concat(a, u8, &.{ prefix, module }); const w = Backend.cpu(try target.checkpoint.get(try std.mem.concat(a, u8, &.{ stem, ".weight" }))); const s = Backend.cpu(try target.checkpoint.get(try std.mem.concat(a, u8, &.{ stem, ".scales" }))); const b = Backend.cpu(try target.checkpoint.get(try std.mem.concat(a, u8, &.{ stem, ".biases" }))); - return ck.validate(w, s, b, .{ .bits = 4, .group_size = 64 }, target.config.vocab, target.config.hidden); + return ck.validate(w, s, b, format, target.config.vocab, target.config.hidden); } diff --git a/zig/src/families/qwen27/gpu_weights.zig b/zig/src/families/qwen27/gpu_weights.zig index dd77c66eac..acd6207aa3 100644 --- a/zig/src/families/qwen27/gpu_weights.zig +++ b/zig/src/families/qwen27/gpu_weights.zig @@ -37,11 +37,6 @@ pub const Weights = struct { for (w.layers) |l| for ([_]Linear{ l.gu, l.down }) |p| if (!uniform(p)) return false; return true; } - - /// Whether the embedding and head are 4-bit g64 (the format the draft head borrows). - pub fn uniformEnds(w: *const Weights) bool { - return w.embed.bits == 4 and w.embed.group == 64 and uniform(w.head); - } }; fn uniform(p: Linear) bool { diff --git a/zig/tests/qwen27/dflash_head.zig b/zig/tests/qwen27/dflash_head.zig index 451473b63a..eaa4fb0e04 100644 --- a/zig/tests/qwen27/dflash_head.zig +++ b/zig/tests/qwen27/dflash_head.zig @@ -6,7 +6,7 @@ const Head = @import("tensorfold").qwen27.dflash.runtime_backend.Head; pub fn run(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue) !void { const n = 248320; const k = 128; - const sizes = [_]usize{ n * k / 2, n * 2 * 4, n * 2 * 2, n * 2 * 2, 7 * k * 2, 7 * n * 2, 7 * 98592 * 2, 64, 64 }; + const sizes = [_]usize{ n * k, n * 2 * 4, n * 2 * 2, n * 2 * 2, 7 * k * 2, 7 * n * 2, 7 * 98592 * 2, 64, 64 }; var buffers: [sizes.len]mtl.Buffer = undefined; var made: usize = 0; defer for (buffers[0..made]) |b| b.deinit(); @@ -28,9 +28,9 @@ pub fn run(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue) !void { var kernels = q.quant_gpu.Kernels{ .allocator = a, .device = device }; defer kernels.deinit(); var cases: usize = 0; - for ([_]bool{ false, true }) |raw| { + for ([_]u8{ 2, 4, 8 }) |bits| for ([_]bool{ false, true }) |raw| { if (!raw and !device.tensorUnits()) continue; // the lane layout needs tensor units; pre-M5 heads take raw rows - const linear = q.projection.Linear{ .words = .{ .buffer = buffers[0], .offset = 32 }, .pairs = .{ .buffer = buffers[1], .offset = 32 }, .n = n, .k = k, .tile = 32, .slices = 1, .raw = if (raw) .{ .w = buffers[0], .scales = buffers[2], .biases = buffers[3], .w_off = 32, .s_off = 32, .b_off = 32, .n = n, .k = k, .group = 64, .bits = 4, .sum = .f32 } else null }; + const linear = q.projection.Linear{ .words = .{ .buffer = buffers[0], .offset = 32 }, .pairs = .{ .buffer = buffers[1], .offset = 32 }, .n = n, .k = k, .tile = 32, .slices = 1, .bits = bits, .raw = if (raw) .{ .w = buffers[0], .scales = buffers[2], .biases = buffers[3], .w_off = 32, .s_off = 32, .b_off = 32, .n = n, .k = k, .group = 64, .bits = bits, .sum = .f32 } else null }; var selected = try Head.init(a, device, linear); defer selected.deinit(); const ids = selected.ids.?.slice(u32, 98592); @@ -53,6 +53,6 @@ pub fn run(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue) !void { }; cases += 1; } - } + }; std.debug.print("native draft vocabulary: {d} lane/raw width cases,98592 selected columns preserve full-head BF16 bytes and token IDs\n", .{cases}); } From 64c7d27caf631bfc1bd2c416be1b2211abe6a46b Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Sun, 11 Oct 2026 19:20:16 -1000 Subject: [PATCH 10/16] A mixed 27B stack runs every member as its own product, as the reference does The reference stacks fused members only when all share one width and group; otherwise each member is its own call at its own K split. Members of one width beside a member of another no longer share a product. --- zig/src/families/qwen27/gpu_weights.zig | 2 +- zig/src/families/qwen27/projection.zig | 2 +- zig/src/families/qwen27/quant_gpu.zig | 2 +- zig/src/families/qwen27/widths.zig | 22 +++++++++------------- 4 files changed, 12 insertions(+), 16 deletions(-) diff --git a/zig/src/families/qwen27/gpu_weights.zig b/zig/src/families/qwen27/gpu_weights.zig index acd6207aa3..5ab7edb8da 100644 --- a/zig/src/families/qwen27/gpu_weights.zig +++ b/zig/src/families/qwen27/gpu_weights.zig @@ -114,7 +114,7 @@ const Builder = struct { return b.stack(triples[0..parts.len]); } - /// Fused members of one width take one matrix as before; mixed widths become one matrix per run of a width. + /// Fused members of one width take one matrix as before; a mixed stack becomes one matrix a member. fn stack(b: Builder, parts: []const projection.Triple) !Linear { var ws: [3]Width = undefined; for (parts, 0..) |p, i| ws[i] = .{ .bits = p.bits, .group = @intCast(p.group) }; diff --git a/zig/src/families/qwen27/projection.zig b/zig/src/families/qwen27/projection.zig index cc52ed1d57..dbbd13ae8a 100644 --- a/zig/src/families/qwen27/projection.zig +++ b/zig/src/families/qwen27/projection.zig @@ -8,7 +8,7 @@ const lane = @import("core").lane_projection; pub const Ref = struct { buffer: mtl.Buffer, offset: usize = 0 }; /// A device matrix; `runs` non-empty means fused members of different widths, each its own product into these rows. pub const Linear = struct { words: Ref, pairs: Ref, n: u32, k: u32, tile: u32, slices: u32, raw: ?row.Weights = null, bits: u8 = 4, group: u32 = 64, runs: []const Run = &.{} }; -/// Members of one width inside a fused matrix, writing output columns [first, first + linear.n). +/// One product inside a mixed fused matrix, writing output columns [first, first + linear.n). pub const Run = struct { linear: Linear, first: u32 }; pub const Triple = struct { weight: checkpoint.Tensor, scales: checkpoint.Tensor, biases: checkpoint.Tensor, bits: u8 = 4, group: u32 = 64 }; diff --git a/zig/src/families/qwen27/quant_gpu.zig b/zig/src/families/qwen27/quant_gpu.zig index 824a87eb95..8fab227d2c 100644 --- a/zig/src/families/qwen27/quant_gpu.zig +++ b/zig/src/families/qwen27/quant_gpu.zig @@ -154,7 +154,7 @@ pub const Kernels = struct { try span.end(); } - /// Mixed-width members: each run its own product into its columns, a row a launch so every row count matches. + /// A mixed stack: each run its own product into its columns, a row a launch so every row count matches. fn quantRuns(k: *Kernels, input_encoder: mtl.ComputeEncoder, p: projection.Linear, x: Ref, xs: ?Sums, dims: Ref, y: Ref, rows: u32, skip: u32, cut: [2]u32) !void { if (skip > 0 or cut[1] > 0 or rows == 0 or rows > 128) return error.BadCoreProjectionRows; for (p.runs) |run| for (0..rows) |r| { diff --git a/zig/src/families/qwen27/widths.zig b/zig/src/families/qwen27/widths.zig index ffb082842f..ada7ce0f64 100644 --- a/zig/src/families/qwen27/widths.zig +++ b/zig/src/families/qwen27/widths.zig @@ -60,18 +60,14 @@ pub fn stacks(kind: cfg.Kind) []const []const Member { }; } -/// Consecutive members of one width as [first, count]; a stack of one width is one run. +/// A stack's products as [first, count]: one for a stack of one width, else one a member, as the reference runs them. pub fn runs(ws: []const Width, out: [][2]usize) [][2]usize { - var n: usize = 0; - for (ws, 0..) |w, i| { - if (n > 0 and std.meta.eql(ws[out[n - 1][0]], w)) { - out[n - 1][1] += 1; - } else { - out[n] = .{ i, 1 }; - n += 1; - } - } - return out[0..n]; + for (ws[1..]) |w| if (!std.meta.eql(ws[0], w)) { + for (out[0..ws.len], 0..) |*o, i| o.* = .{ i, 1 }; + return out[0..ws.len]; + }; + out[0] = .{ 0, ws.len }; + return out[0..1]; } pub const Widths = struct { @@ -237,12 +233,12 @@ test "a dense projection and an unsupported global group refuse" { try std.testing.expectError(error.BadAffine, testAdmit("{\"bits\":4,\"group_size\":16}", .{ .tensor_units = false }, &why)); } -test "fused members split into runs of one width" { +test "a mixed stack runs every member alone, even beside a member of its own width" { var buf: [3][2]usize = undefined; const four = uniform; const five = Width{ .bits = 5, .group = 64 }; try std.testing.expectEqualSlices([2]usize, &.{.{ 0, 3 }}, runs(&.{ four, four, four }, &buf)); - try std.testing.expectEqualSlices([2]usize, &.{ .{ 0, 1 }, .{ 1, 2 } }, runs(&.{ five, four, four }, &buf)); + try std.testing.expectEqualSlices([2]usize, &.{ .{ 0, 1 }, .{ 1, 1 }, .{ 2, 1 } }, runs(&.{ five, four, four }, &buf)); try std.testing.expectEqualSlices([2]usize, &.{ .{ 0, 1 }, .{ 1, 1 }, .{ 2, 1 } }, runs(&.{ four, five, four }, &buf)); } From f70e5cb5ecf62781cbde6252854894ee58902c08 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Sun, 11 Oct 2026 19:34:31 -1000 Subject: [PATCH 11/16] A mixed 27B stack writes each member into the fused rows at a stride before tensor units Each member is still its own product, now one row-kernel launch over all rows writing its columns at the stack width, in place of one launch a row. The row kernel gives a row the same bits in any window, so every row count still matches. --- zig/src/families/qwen27/quant_gpu.zig | 40 ++++++++++++++++++++++----- 1 file changed, 33 insertions(+), 7 deletions(-) diff --git a/zig/src/families/qwen27/quant_gpu.zig b/zig/src/families/qwen27/quant_gpu.zig index 8fab227d2c..63c0a8974c 100644 --- a/zig/src/families/qwen27/quant_gpu.zig +++ b/zig/src/families/qwen27/quant_gpu.zig @@ -154,15 +154,41 @@ pub const Kernels = struct { try span.end(); } - /// A mixed stack: each run its own product into its columns, a row a launch so every row count matches. + /// A mixed stack: each member its own product into its columns; row kernels write the fused rows at a stride. fn quantRuns(k: *Kernels, input_encoder: mtl.ComputeEncoder, p: projection.Linear, x: Ref, xs: ?Sums, dims: Ref, y: Ref, rows: u32, skip: u32, cut: [2]u32) !void { if (skip > 0 or cut[1] > 0 or rows == 0 or rows > 128) return error.BadCoreProjectionRows; - for (p.runs) |run| for (0..rows) |r| { - const xr = Ref{ .buffer = x.buffer, .offset = x.offset + r * p.k * 2 }; - const sr: ?Sums = if (xs) |s| .{ .ref = .{ .buffer = s.ref.buffer, .offset = s.ref.offset + r * 4 }, .stride = s.stride } else null; - const yr = Ref{ .buffer = y.buffer, .offset = y.offset + (r * p.n + run.first) * 2 }; - try k.quantOne(input_encoder, run.linear, xr, sr, dims, yr, 1, 0, .{ 0, 0 }); - }; + for (p.runs) |run| { + if (run.linear.raw) |raw| { + try k.rowStrided(input_encoder, raw, x, .{ .buffer = y.buffer, .offset = y.offset + @as(usize, run.first) * 2 }, rows, p.n); + continue; + } + for (0..rows) |r| { // lane kernels take one row a launch until they write at a stride + const xr = Ref{ .buffer = x.buffer, .offset = x.offset + r * p.k * 2 }; + const sr: ?Sums = if (xs) |s| .{ .ref = .{ .buffer = s.ref.buffer, .offset = s.ref.offset + r * 4 }, .stride = s.stride } else null; + const yr = Ref{ .buffer = y.buffer, .offset = y.offset + (r * p.n + run.first) * 2 }; + try k.quantOne(input_encoder, run.linear, xr, sr, dims, yr, 1, 0, .{ 0, 0 }); + } + } + } + + /// The row kernel at any width, output row m at y + m * stride values; a row's bits do not depend on the window. + fn rowStrided(k: *Kernels, input_encoder: mtl.ComputeEncoder, raw: row.Weights, x: Ref, y: Ref, rows: u32, stride: u32) !void { + var span = profile.begin(k.profiler, input_encoder, .dense); + span.shape = .{ .n = @intCast(raw.n), .k = @intCast(raw.k), .sk = 1 }; + errdefer span.cancel(); + const e = span.encoder; + if (k.row_pipelines == null) k.row_pipelines = try row.Pipelines.load(k.device); + const call = try row.callStrided(&k.row_pipelines.?, raw, rows, stride); + e.setPipeline(call.pipeline); + e.setBuffer(x.buffer, x.offset, 0); + e.setBuffer(raw.w, raw.w_off, 1); + e.setBuffer(raw.scales, raw.s_off, 2); + e.setBuffer(raw.biases, raw.b_off, 3); + e.setValue(call.dims, 4); + e.setBuffer(y.buffer, y.offset, 5); + e.dispatchGroups(.{ .width = call.groups }, .{ .width = call.threads }); + e.barrier(); + try span.end(); } /// Loads the tiles on first use and checks the scalar twin against them once per shape. From cef59343874d1b1878ed548644b47b0993f279d0 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 07:32:53 -1000 Subject: [PATCH 12/16] tests: qwen27-widths, every 27B projection class at every oQ width against float64 without tensor units tf-qwen27-widths runs the family's own products (quant_gpu.Kernels.quant on projection.prepare's matrices: the row kernels, or simd_qmm at 4/64, as a chip without tensor units picks them) on synthetic MLX affine weights at the shape of every projection class (GDN qkv, z, b, a and out; attention q, k, v and o; MLP gate_up and down; lm_head) at 2, 3, 4, 5, 6 and 8 bits, group 64, which covers every width oQ4 and oQ2 use: - a prompt call (128 rows) and decoded rows (one row a call) against core.quant.mlx.reference in float64, the worst column within the fp32-sum bound; scales 1% off exceed the bound in every class and width; - every row's bits equal in calls of 1, 2, 3, 8, 16, 17 and 128 rows; - the mixed stacks [z4 | b5 a5], [z4 b4 | a5], [z2 b2 | a3] and [k5 | v6], built as the loader builds them: every row's bits equal in every call, each member's columns within the float64 bound and equal bit for bit to its own product (the row kernel at its width, unstrided), and to the member loaded as a matrix of its own, except a 4-bit z, which alone takes simd_qmm (no output stride). The engine module exports core.quant for the reference. The program refuses a chip with tensor units, whose lane kernels it does not check. M3 Ultra (no tensor units): PASS 10,575,872 values within 0.482 of the float64 bound (bf16 output rounding; worst-column |err| / max|y| 2.5e-3 to 3.9e-3), scales 1% off at least 2.1 times the bound in every class; 1,515,200,224 bit-equal pairs; 31 s. A 4-bit z inside a mixed stack differs from the same z as a uniform layer's simd_qmm product in 51 and 72 of 128 rows, by up to 3.1e-2; both are within the bound. zig build test test-golden -Dcpu=apple_m1: 158/160 steps, 1566/1574 tests (8 skipped), goldens 309/319 with 0 unexpected differences; the one failure is main's core.checkpoint_metal split-checkpoint test. tools/zig/lean_check.py: 536 findings, 0 new against main. --- zig/build/qwen27.zig | 1 + zig/src/tensorfold.zig | 2 + zig/tests/qwen27/widths.zig | 495 ++++++++++++++++++++++++++++++++++++ 3 files changed, 498 insertions(+) create mode 100644 zig/tests/qwen27/widths.zig diff --git a/zig/build/qwen27.zig b/zig/build/qwen27.zig index fc5245b612..73a84a639a 100644 --- a/zig/build/qwen27.zig +++ b/zig/build/qwen27.zig @@ -97,6 +97,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu const engine_only: []const std.Build.Module.Import = &.{ .{ .name = "tensorfold", .module = engine }, .{ .name = "metal", .module = metal } }; _ = program(b, "tf-qwen27-synthetic", "Build bounded synthetic native hybrid chunk checks, no models", "zig/tests/qwen27/synthetic.zig", target, false, engine_only); _ = program(b, "tf-qwen27-dflash-synthetic", "Build model-free native draft context/forward/tree check", "zig/tests/qwen27/dflash_runtime.zig", target, false, engine_only); + _ = program(b, "tf-qwen27-widths", "Build the projections at every oQ width before tensor units: float64, row windows, mixed stacks", "zig/tests/qwen27/widths.zig", target, false, engine_only); _ = program(b, "tf-qwen27-dflash-refine", "Build the drafter refinement probe (top-1 by block position, masked versus given)", "zig/tests/qwen27/dflash_refine.zig", target, false, engine_only); const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} }); 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/qwen27/widths.zig b/zig/tests/qwen27/widths.zig new file mode 100644 index 0000000000..a1d78aaa7e --- /dev/null +++ b/zig/tests/qwen27/widths.zig @@ -0,0 +1,495 @@ +//! The 27B's projections without tensor units at every oQ width against float64, row bits in every call, mixed stacks. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const mlx = tf.quant.mlx; +const q27 = tf.qwen27; +const projection = q27.projection; +const row = tf.row_projection; +const cp = tf.checkpoint; + +const opts = mtl.ResourceOptions.shared; +const a = std.heap.page_allocator; +const unit: f64 = 0x1p-24; // fp32's unit roundoff +const group = 64; // every module of both oQ checkpoints +const wide = 128; // a prompt chunk's rows +const windows = [_]usize{ 1, 2, 3, 8, 16, 17, 128 }; +const bit_widths = [_]u8{ 2, 3, 4, 5, 6, 8 }; // every class at every width oQ4 and oQ2 use anywhere + +/// A projection class: output rows n of input width k (gate and up share a width, so one matrix as loaded). +const Class = struct { name: []const u8, n: usize, k: usize }; +const classes = [_]Class{ + .{ .name = "gdn in_proj_qkv", .n = 10240, .k = 5120 }, + .{ .name = "gdn in_proj_z", .n = 6144, .k = 5120 }, + .{ .name = "gdn in_proj_b", .n = 48, .k = 5120 }, + .{ .name = "gdn in_proj_a", .n = 48, .k = 5120 }, + .{ .name = "gdn out_proj", .n = 5120, .k = 6144 }, + .{ .name = "attn q_proj", .n = 12288, .k = 5120 }, + .{ .name = "attn k_proj", .n = 1024, .k = 5120 }, + .{ .name = "attn v_proj", .n = 1024, .k = 5120 }, + .{ .name = "attn o_proj", .n = 5120, .k = 6144 }, + .{ .name = "mlp gate_up", .n = 34816, .k = 5120 }, + .{ .name = "mlp down_proj", .n = 5120, .k = 17408 }, + .{ .name = "lm_head", .n = 248320, .k = 5120 }, +}; + +/// A fused member: name, rows, bits. +const Member = struct { name: []const u8, n: usize, bits: u8 }; +/// The mixed stacks of both checkpoints (zba and kv), k = 5120. +const stacks = [_][]const Member{ + &.{ .{ .name = "z4", .n = 6144, .bits = 4 }, .{ .name = "b5", .n = 48, .bits = 5 }, .{ .name = "a5", .n = 48, .bits = 5 } }, + &.{ .{ .name = "z4", .n = 6144, .bits = 4 }, .{ .name = "b4", .n = 48, .bits = 4 }, .{ .name = "a5", .n = 48, .bits = 5 } }, + &.{ .{ .name = "z2", .n = 6144, .bits = 2 }, .{ .name = "b2", .n = 48, .bits = 2 }, .{ .name = "a3", .n = 48, .bits = 3 } }, + &.{ .{ .name = "k5", .n = 1024, .bits = 5 }, .{ .name = "v6", .n = 1024, .bits = 6 } }, +}; + +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 shape2(r: usize, c: usize) [5]usize { + return .{ r, c, 1, 1, 1 }; +} + +/// Synthetic MLX affine weights [n, k] as the checkpoint stores them: words of random codes, bf16 scales and biases. +const Synth = struct { + n: usize, + k: usize, + bits: u8, + w: mtl.Buffer, + s: mtl.Buffer, + b: mtl.Buffer, + + fn init(device: mtl.Device, n: usize, k: usize, bits: u8, seed: u64) !Synth { + var prng = std.Random.DefaultPrng.init(seed); + const r = prng.random(); + const sn = n * (k / group); + const m = Synth{ .n = n, .k = k, .bits = bits, .w = try device.buffer(n * k * bits / 8, opts), .s = try device.buffer(sn * 2, opts), .b = try device.buffer(sn * 2, opts) }; + r.bytes(m.w.slice(u8, n * k * bits / 8)); + for (m.s.slice(u16, sn), m.b.slice(u16, sn)) |*s, *b| { + s.* = toBf((r.float(f32) + 0.5) / 64 / @as(f32, @floatFromInt(@as(u32, 1) << @intCast(bits - 2)))); + b.* = toBf(-(r.float(f32) + 0.5) * 0.15); + } + return m; + } + + fn deinit(m: Synth) void { + for ([_]mtl.Buffer{ m.w, m.s, m.b }) |x| x.deinit(); + } + + fn tensor(buffer: mtl.Buffer, dtype: cp.DType, r: usize, c: usize, size: usize) cp.Tensor { + return .{ .buffer = buffer, .offset = 0, .bytes = r * c * size, .dtype = dtype, .shape = shape2(r, c), .rank = 2 }; + } + + /// The loader's view: what gpu_weights hands projection.prepare. + fn triple(m: Synth) projection.Triple { + const g = m.k / group; + return .{ .weight = tensor(m.w, .u32, m.n, m.k * m.bits / 32, 4), .scales = tensor(m.s, .bf16, m.n, g, 2), .biases = tensor(m.b, .bf16, m.n, g, 2), .bits = m.bits, .group = group }; + } + + /// Row j alone as core.quant.mlx reads it. + fn host(m: Synth, j: usize) mlx.Host { + const g = m.k / group; + const wb = m.k * m.bits / 8; + return .{ + .words = .{ .dtype = .u32, .rank = 2, .shape = shape2(1, wb / 4), .bytes = m.w.slice(u8, m.n * wb)[j * wb ..][0..wb] }, + .scales = .{ .dtype = .bf16, .rank = 2, .shape = shape2(1, g), .bytes = m.s.slice(u8, m.n * g * 2)[j * g * 2 ..][0 .. g * 2] }, + .biases = .{ .dtype = .bf16, .rank = 2, .shape = shape2(1, g), .bytes = m.b.slice(u8, m.n * g * 2)[j * g * 2 ..][0 .. g * 2] }, + .bits = m.bits, + .group = group, + }; + } +}; + +/// `wide` rows of bf16 inputs in [-1, 1). +fn inputs(device: mtl.Device, k: usize, seed: u64) !mtl.Buffer { + var prng = std.Random.DefaultPrng.init(seed); + const x = try device.buffer(wide * k * 2, opts); + for (x.slice(u16, wide * k)) |*v| v.* = toBf(prng.random().float(f32) * 2 - 1); + return x; +} + +/// Every column of a matrix up to 1024 wide, else 512: the first and last 128 and 256 at an odd stride between. +fn pick(n: usize) ![]usize { + if (n <= 1024) { + const cols = try a.alloc(usize, n); + for (cols, 0..) |*c, i| c.* = i; + return cols; + } + const cols = try a.alloc(usize, 512); + const step = ((n - 256) / 256 - 1) | 1; + for (0..128) |i| { + cols[i] = i; + cols[511 - i] = n - 1 - i; + } + for (0..256) |i| cols[128 + i] = 128 + i * step; + return cols; +} + +/// y = x W^T in float64 on `cols`, the sizes the fp32 sums carry, and the scales' share (scales 1% off add 1% of it). +const Ref = struct { + cols: []usize, + want: []f64, + mag: []f64, + scaled: []f64, + + fn init(m: Synth, x: mtl.Buffer) !Ref { + const k = m.k; + const g_n = k / group; + const cols = try pick(m.n); + const c_n = cols.len; + const r = Ref{ .cols = cols, .want = try a.alloc(f64, wide * c_n), .mag = try a.alloc(f64, wide * c_n), .scaled = try a.alloc(f64, wide * c_n) }; + const xd = try a.alloc(f64, wide * k); + defer a.free(xd); + const xg = try a.alloc(f64, wide * g_n); // each group's sum of |x| + defer a.free(xg); + @memset(xg, 0); + for (x.slice(u16, wide * k), xd, 0..) |v, *d, i| { + d.* = fromBf(v); + xg[i / group] += @abs(d.*); + } + const deq = try a.alloc(f64, k); + defer a.free(deq); + const bias = m.b.slice(u16, m.n * g_n); + for (cols, 0..) |j, c| { + try mlx.reference(m.host(j), deq); + for (0..wide) |i| { + var sum: f64 = 0; + var mag: f64 = 0; + var scaled: f64 = 0; + for (0..g_n) |g| { + const b = fromBf(bias[j * g_n + g]); + var qx: f64 = 0; + var qxs: f64 = 0; + for (g * group..(g + 1) * group) |e| { + const xv = xd[i * k + e]; + const sq = deq[e] - b; // s q, exact + sum += xv * deq[e]; + qx += sq * @abs(xv); + qxs += sq * xv; + } + mag += qx + @abs(b) * xg[i * g_n + g]; + scaled += qxs; + } + r.want[i * c_n + c] = sum; + r.mag[i * c_n + c] = mag; + r.scaled[i * c_n + c] = scaled; + } + } + return r; + } + + fn deinit(r: Ref) void { + for ([_][]f64{ r.want, r.mag, r.scaled }) |v| a.free(v); + a.free(r.cols); + } +}; + +/// Failures, values against float64, bit-equal pairs, and the largest difference as a fraction of its bound. +const Tally = struct { failed: usize = 0, values: usize = 0, pairs: usize = 0, worst: f64 = 0, ctl_min: f64 = std.math.inf(f64) }; + +/// `rows` bf16 rows `pitch` apart; a member's columns start at `first`. +const View = struct { + vals: []const u16, + pitch: usize, + first: usize = 0, + fn get(v: View, r: usize, j: usize) f64 { + return fromBf(v.vals[r * v.pitch + v.first + j]); + } +}; + +/// fp32 sums of a group's products, its row sum and a chain of groups (plus the kernel's tree), and one bf16 ulp. +fn bound(want: f64, mag: f64, k: usize, extra: usize) f64 { + const terms: f64 = @floatFromInt(group + k / group + 8 + extra); + return terms * unit * mag + @abs(want) * 0x1p-7 + 1e-30; +} + +/// `rows` rows of `out` against the float64 reference: the worst column, its bound, and the 1%-off control. +fn compare(t: *Tally, name: []const u8, out: View, rows: usize, ref: Ref, k: usize, extra: usize) void { + const c_n = ref.cols.len; + var col: usize = 0; + var col_abs: f64 = -1; + var col_rel: f64 = 0; + var frac: f64 = 0; + var ctl: f64 = 0; + var at: [2]usize = .{ 0, 0 }; + for (ref.cols, 0..) |j, c| { + var abs: f64 = 0; + var top: f64 = 0; + for (0..rows) |r| { + const i = r * c_n + c; + const got = out.get(r, j); + const d = @abs(got - ref.want[i]); + const tol = bound(ref.want[i], ref.mag[i], k, extra); + ctl = @max(ctl, @abs(got - ref.want[i] - 0.01 * ref.scaled[i]) / tol); + if (!(d <= abs)) abs = d; + top = @max(top, @abs(ref.want[i])); + if (!(d / tol <= frac)) { + frac = d / tol; + at = .{ r, j }; + } + } + if (!(abs <= col_abs)) { + col_abs = abs; + col = j; + col_rel = abs / top; + } + } + t.values += rows * c_n; + t.worst = @max(t.worst, frac); + t.ctl_min = @min(t.ctl_min, ctl); + const ok = frac <= 1 and ctl > 1; + if (!ok) t.failed += 1; + std.debug.print("RESULT {s}: {d} columns, 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, c_n, col, col_abs, col_rel, frac, at[0], at[1], ctl, if (ok) "" else " FAIL" }); +} + +/// Rows [0, rows) of `got` equal `want` bit for bit; the first difference printed. +fn same(t: *Tally, name: []const u8, what: []const u8, want: View, got: View, rows: usize, n: usize) bool { + var differ: usize = 0; + var first: ?[2]usize = null; + var most: f64 = 0; + for (0..rows) |r| { + const w = want.vals[r * want.pitch + want.first ..][0..n]; + const g = got.vals[r * got.pitch + got.first ..][0..n]; + if (std.mem.eql(u16, w, g)) continue; + differ += 1; + if (first == null) first = .{ r, std.mem.indexOfDiff(u16, w, g).? }; + for (0..n) |j| most = @max(most, @abs(want.get(r, j) - got.get(r, j))); + } + if (first) |f| { + std.debug.print(" {s}: {s} differ in {d} of {d} rows, by up to {e:.3}; first row {d} col {d}: {d} against {d}\n", .{ name, what, differ, rows, most, f[0], f[1], got.get(f[0], f[1]), want.get(f[0], f[1]) }); + return false; + } + t.pairs += rows * n; + return true; +} + +const Gpu = struct { + device: mtl.Device, + queue: mtl.Queue, + kern: *q27.quant_gpu.Kernels, + pipes: *const row.Pipelines, + + fn commit(cb: mtl.CommandBuffer, 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; + } + } + + /// The family's product over `wide` rows in windows of `m` (a prompt chunk when m == wide), into a NaN-filled y. + fn quant(g: Gpu, lin: projection.Linear, x: mtl.Buffer, y: mtl.Buffer, m: usize) !void { + const n: usize = lin.n; + @memset(y.slice(u8, wide * n * 2), 0xff); + g.kern.prompt = m == wide; + defer g.kern.prompt = false; + const cb = g.queue.commandBuffer(); + const e = cb.compute(.serial); + var start: usize = 0; + while (start + m <= wide) : (start += m) try g.kern.quant(e, lin, .{ .buffer = x, .offset = start * lin.k * 2 }, null, .{ .buffer = x }, .{ .buffer = y, .offset = start * n * 2 }, @intCast(m)); + try commit(cb, e); + } + + /// The row kernel alone at the matrix's width over `wide` rows, unstrided: a member's own product. + fn rows(g: Gpu, raw: row.Weights, x: mtl.Buffer, y: mtl.Buffer) !void { + @memset(y.slice(u8, wide * raw.n * 2), 0xff); + const cb = g.queue.commandBuffer(); + const e = cb.compute(.serial); + const call = try row.call(g.pipes, raw, wide, false); + e.setPipeline(call.pipeline); + e.setBuffer(x, 0, 0); + e.setBuffer(raw.w, raw.w_off, 1); + e.setBuffer(raw.scales, raw.s_off, 2); + e.setBuffer(raw.biases, raw.b_off, 3); + e.setValue(call.dims, 4); + e.setBuffer(y, 0, 5); + e.dispatchGroups(.{ .width = call.groups }, .{ .width = call.threads }); + try commit(cb, e); + } +}; + +/// The kernel the family picks for a matrix before tensor units, and its tree's extra depth for the bound. +fn path(lin: projection.Linear) struct { name: []const u8, extra: usize } { + const raw = lin.raw.?; + if (raw.bits == 4 and raw.group == 64 and raw.n % 8 == 0) return .{ .name = "simd_qmm", .extra = std.math.log2_int(u32, row.simd.splits(raw.n)) + 1 }; + return .{ .name = "row", .extra = 0 }; +} + +fn release(owned: *std.ArrayList(mtl.Buffer)) void { + for (owned.items) |b| b.deinit(); + owned.deinit(a); +} + +/// One class at one width: prompt and decoded rows against float64, every window's rows against the prompt call's. +fn class(t: *Tally, g: Gpu, c: Class, bits: u8, seed: u64) !void { + const m = try Synth.init(g.device, c.n, c.k, bits, seed); + defer m.deinit(); + const x = try inputs(g.device, c.k, seed ^ 0x27); + defer x.deinit(); + var owned: std.ArrayList(mtl.Buffer) = .empty; + defer release(&owned); + const lin = try projection.prepare(g.device, a, &owned, &.{m.triple()}, 32); + const p = path(lin); + const ref = try Ref.init(m, x); + defer ref.deinit(); + var name: [128]u8 = undefined; + const label = try std.fmt.bufPrint(&name, "{s} [{d}, {d}] q{d} g64 on {s}", .{ c.name, c.n, c.k, bits, p.name }); + var outs: [windows.len]mtl.Buffer = undefined; + for (&outs, windows) |*y, w| { + y.* = try g.device.buffer(wide * c.n * 2, opts); + try g.quant(lin, x, y.*, w); + } + defer for (outs) |y| y.deinit(); + const view = struct { + fn of(y: mtl.Buffer, n: usize) View { + return .{ .vals = y.slice(u16, wide * n), .pitch = n }; + } + }.of; + var tag: [160]u8 = undefined; + compare(t, try std.fmt.bufPrint(&tag, "{s}, prompt ({d} rows)", .{ label, wide }), view(outs[windows.len - 1], c.n), wide, ref, c.k, p.extra); + compare(t, try std.fmt.bufPrint(&tag, "{s}, decode (1 row a call)", .{label}), view(outs[0], c.n), wide, ref, c.k, p.extra); + var ok = true; + for (windows[0 .. windows.len - 1], outs[0 .. windows.len - 1]) |w, y| { + var what: [48]u8 = undefined; + ok = same(t, label, try std.fmt.bufPrint(&what, "{d}-row windows", .{w}), view(outs[windows.len - 1], c.n), view(y, c.n), wide / w * w, c.n) and ok; + } + if (!ok) t.failed += 1; + std.debug.print("RESULT {s}: {s}\n", .{ label, if (ok) "every row's bits equal in calls of 1, 2, 3, 8, 16, 17 and 128 rows" else "a row's bits depend on the call FAIL" }); +} + +/// A mixed stack as gpu_weights.Builder.stack builds it: one product a member, each into its columns. +fn mixed(g: Gpu, owned: *std.ArrayList(mtl.Buffer), parts: []const projection.Triple, runs: []projection.Run) !projection.Linear { + var ws: [3]q27.widths.Width = undefined; + for (parts, 0..) |p, i| ws[i] = .{ .bits = p.bits, .group = @intCast(p.group) }; + var buf: [3][2]usize = undefined; + const list = q27.widths.runs(ws[0..parts.len], &buf); + if (list.len != parts.len) return error.NotMixed; + var first: u32 = 0; + for (list, runs[0..list.len]) |r, *run| { + run.* = .{ .linear = try projection.prepare(g.device, a, owned, parts[r[0]..][0..r[1]], 32), .first = first }; + first += run.linear.n; + } + var fused = runs[0].linear; + fused.n = first; + fused.raw = null; + fused.runs = runs[0..list.len]; + return fused; +} + +/// A mixed stack: rows equal in every call; each member's columns within float64's bound and equal to its own product. +fn stack(t: *Tally, g: Gpu, members: []const Member, seed: u64) !void { + const k = 5120; + var name: [64]u8 = undefined; + var len: usize = 0; + for (members, 0..) |m, i| len += (try std.fmt.bufPrint(name[len..], "{s}{s}", .{ if (i == 0) "[" else " ", m.name })).len; + const label = try std.fmt.bufPrint(name[len..], "]", .{}); + const stack_name = name[0 .. len + label.len]; + var synth: [3]Synth = undefined; + var parts: [3]projection.Triple = undefined; + var total: usize = 0; + for (members, 0..) |m, i| { + synth[i] = try Synth.init(g.device, m.n, k, m.bits, seed + i); + parts[i] = synth[i].triple(); + total += m.n; + } + defer for (synth[0..members.len]) |s| s.deinit(); + const x = try inputs(g.device, k, seed ^ 0x27); + defer x.deinit(); + var owned: std.ArrayList(mtl.Buffer) = .empty; + defer release(&owned); + var runs: [3]projection.Run = undefined; + const fused = try mixed(g, &owned, parts[0..members.len], &runs); + var outs: [windows.len]mtl.Buffer = undefined; + for (&outs, windows) |*y, w| { + y.* = try g.device.buffer(wide * total * 2, opts); + try g.quant(fused, x, y.*, w); + } + defer for (outs) |y| y.deinit(); + const prompt = View{ .vals = outs[windows.len - 1].slice(u16, wide * total), .pitch = total }; + var ok = true; + for (windows[0 .. windows.len - 1], outs[0 .. windows.len - 1]) |w, y| { + var what: [48]u8 = undefined; + ok = same(t, stack_name, try std.fmt.bufPrint(&what, "{d}-row windows", .{w}), prompt, .{ .vals = y.slice(u16, wide * total), .pitch = total }, wide / w * w, total) and ok; + } + if (!ok) t.failed += 1; + std.debug.print("RESULT stack {s}: {s}\n", .{ stack_name, if (ok) "every row's bits equal in calls of 1, 2, 3, 8, 16, 17 and 128 rows" else "a row's bits depend on the call FAIL" }); + const alone = try g.device.buffer(wide * synth[0].n * 2, opts); + defer alone.deinit(); + const matrix = try g.device.buffer(wide * synth[0].n * 2, opts); + defer matrix.deinit(); + for (members, runs[0..members.len], 0..) |m, run, i| { + var tag: [128]u8 = undefined; + const member = try std.fmt.bufPrint(&tag, "stack {s} member {s} [{d}, {d}] q{d} g64", .{ stack_name, m.name, m.n, k, m.bits }); + const col = View{ .vals = prompt.vals, .pitch = total, .first = run.first }; + const ref = try Ref.init(synth[i], x); + defer ref.deinit(); + var line: [192]u8 = undefined; + compare(t, try std.fmt.bufPrint(&line, "{s}, prompt ({d} rows)", .{ member, wide }), col, wide, ref, k, 0); + compare(t, try std.fmt.bufPrint(&line, "{s}, decode (1 row a call)", .{member}), .{ .vals = outs[0].slice(u16, wide * total), .pitch = total, .first = run.first }, wide, ref, k, 0); + try g.rows(run.linear.raw.?, x, alone); + const own = View{ .vals = alone.slice(u16, wide * m.n), .pitch = m.n }; + if (same(t, member, "the stack's columns and the member's own row product", own, col, wide, m.n)) { + std.debug.print("RESULT {s}: its columns equal its own row product bit for bit\n", .{member}); + } else { + std.debug.print("RESULT {s}: its columns differ from its own row product FAIL\n", .{member}); + t.failed += 1; + } + // the member as a matrix of its own through the family: the row kernel, or simd_qmm at 4/64 (no output stride) + var solo: std.ArrayList(mtl.Buffer) = .empty; + defer release(&solo); + const lin = try projection.prepare(g.device, a, &solo, parts[i .. i + 1], 32); + try g.quant(lin, x, matrix, wide); + const via = path(lin).name; + if (same(t, member, "the stack's columns and the member as its own matrix", .{ .vals = matrix.slice(u16, wide * m.n), .pitch = m.n }, col, wide, m.n)) { + std.debug.print("RESULT {s}: equals the member as its own matrix ({s}) bit for bit\n", .{ member, via }); + } else if (std.mem.eql(u8, via, "simd_qmm")) { + std.debug.print("NOTE {s}: as its own matrix it takes simd_qmm (no output stride), whose bits differ\n", .{member}); + } else { + std.debug.print("RESULT {s}: differs from the member as its own matrix ({s}) FAIL\n", .{ member, via }); + 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() }); + if (device.tensorUnits()) { + std.debug.print("this program checks the row path; tensor units take the lane kernels\n", .{}); + return error.TensorUnits; + } + var kern = q27.quant_gpu.Kernels{ .allocator = a, .device = device }; + defer kern.deinit(); + var pipes = try row.Pipelines.load(device); + defer pipes.deinit(); + const g = Gpu{ .device = device, .queue = queue, .kern = &kern, .pipes = &pipes }; + var t = Tally{}; + var seed: u64 = 0x2727; + for (classes) |c| for (bit_widths) |bits| { + seed += 1; + try class(&t, g, c, bits, seed); + }; + for (stacks) |s| { + seed += 16; + try stack(&t, g, s, seed); + } + 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 {d} values within {d:.4} of the float64 bound (scales 1% off: at least {d:.1} of it), {d} bit-equal pairs\n", .{ t.values, t.worst, t.ctl_min, t.pairs }); +} From b2eb33a98e5df33ba508eed67dc03bbca6293c8b Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 07:38:24 -1000 Subject: [PATCH 13/16] tests: qwen27-widths gathers the embedding at 8 and 3 bits against float64 tf-qwen27-widths now also runs the family's embedding gather (glue_gpu.Kernels.embedding, q27_embed) on a synthetic MLX affine table at the 27B's shape [248320, 5120], group 64, at 8 bits (oQ2's embedding) and 3 bits, for 64 tokens including both ends of the vocabulary and a repeat: - every gathered value equals core.quant.mlx.reference in float64 rounded once to bf16, bit for bit; - one-row calls give the same bits as one 64-row call; - the reference with scales 1% off differs in a large share of values, so the comparison sees a wrong decode. M3 Ultra (no tensor units): 8-bit 0 of 327,680 values differ, 3-bit 0 of 327,680; one-row calls equal; scales 1% off differ in 160,115 and 142,528 values. The projection checks are unchanged (PASS 10,575,872 values within 0.482 of the bound). zig build test test-golden -Dcpu=apple_m1: 158/160 steps, goldens 309/319 with 0 unexpected differences; the one failure is main's core.checkpoint_metal split-checkpoint test. tools/zig/lean_check.py: 536 findings, 0 new. --- zig/build/qwen27.zig | 2 +- zig/tests/qwen27/embed.zig | 154 ++++++++++++++++++++++++++++++++++++ zig/tests/qwen27/widths.zig | 3 +- 3 files changed, 157 insertions(+), 2 deletions(-) create mode 100644 zig/tests/qwen27/embed.zig diff --git a/zig/build/qwen27.zig b/zig/build/qwen27.zig index 73a84a639a..73b38e72f1 100644 --- a/zig/build/qwen27.zig +++ b/zig/build/qwen27.zig @@ -97,7 +97,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu const engine_only: []const std.Build.Module.Import = &.{ .{ .name = "tensorfold", .module = engine }, .{ .name = "metal", .module = metal } }; _ = program(b, "tf-qwen27-synthetic", "Build bounded synthetic native hybrid chunk checks, no models", "zig/tests/qwen27/synthetic.zig", target, false, engine_only); _ = program(b, "tf-qwen27-dflash-synthetic", "Build model-free native draft context/forward/tree check", "zig/tests/qwen27/dflash_runtime.zig", target, false, engine_only); - _ = program(b, "tf-qwen27-widths", "Build the projections at every oQ width before tensor units: float64, row windows, mixed stacks", "zig/tests/qwen27/widths.zig", target, false, engine_only); + _ = program(b, "tf-qwen27-widths", "Build the projections at every oQ width before tensor units: float64, row windows, mixed stacks, embedding", "zig/tests/qwen27/widths.zig", target, false, engine_only); _ = program(b, "tf-qwen27-dflash-refine", "Build the drafter refinement probe (top-1 by block position, masked versus given)", "zig/tests/qwen27/dflash_refine.zig", target, false, engine_only); const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} }); diff --git a/zig/tests/qwen27/embed.zig b/zig/tests/qwen27/embed.zig new file mode 100644 index 0000000000..a031fe8ccd --- /dev/null +++ b/zig/tests/qwen27/embed.zig @@ -0,0 +1,154 @@ +//! The 27B's embedding gather at oQ2's 8-bit and a 3-bit width: every column the float64 decode rounded once to bf16. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const mlx = tf.quant.mlx; +const q27 = tf.qwen27; + +const opts = mtl.ResourceOptions.shared; +const a = std.heap.page_allocator; +const group = 64; +const hidden: usize = 5120; +const vocab = 248320; +const rows = 64; + +fn toBf(v: f32) u16 { + const u: u32 = @bitCast(v); + return @truncate((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); +} + +/// One rounding of a float64 to bf16, ties to even (normal range; the gathered values are). +fn bf64(v: f64) u16 { + const u: u64 = @bitCast(v); + const r: f64 = @bitCast((u +% 0xfff_ffff_ffff +% ((u >> 45) & 1)) & ~@as(u64, 0x1fff_ffff_ffff)); + return @truncate(@as(u32, @bitCast(@as(f32, @floatCast(r)))) >> 16); +} + +/// Token ids: both ends of the vocabulary, a repeat, and random rows. +fn ids(seed: u64) [rows]u32 { + var prng = std.Random.DefaultPrng.init(seed); + var out: [rows]u32 = undefined; + for (&out) |*v| v.* = prng.random().uintLessThan(u32, vocab); + out[0] = 0; + out[1] = vocab - 1; + out[2] = 1; + out[3] = vocab - 2; + out[4] = out[7]; + return out; +} + +/// A [vocab, hidden] table at `bits`, the gathered rows filled with random codes, bf16 scales and biases. +const Table = struct { + w: mtl.Buffer, + s: mtl.Buffer, + b: mtl.Buffer, + bits: u8, + + fn init(device: mtl.Device, bits: u8, tokens: []const u32, seed: u64) !Table { + const wb = hidden * bits / 8; + const g = hidden / group; + const t = Table{ .w = try device.buffer(vocab * wb, opts), .s = try device.buffer(vocab * g * 2, opts), .b = try device.buffer(vocab * g * 2, opts), .bits = bits }; + var prng = std.Random.DefaultPrng.init(seed); + const r = prng.random(); + for (tokens) |id| { + r.bytes(t.w.slice(u8, vocab * wb)[id * wb ..][0..wb]); + for (t.s.slice(u16, vocab * g)[id * g ..][0..g], t.b.slice(u16, vocab * g)[id * g ..][0..g]) |*s, *b| { + s.* = toBf((r.float(f32) + 0.5) / 64 / @as(f32, @floatFromInt(@as(u32, 1) << @intCast(bits - 2)))); + b.* = toBf(-(r.float(f32) + 0.5) * 0.15); + } + } + return t; + } + + fn deinit(t: Table) void { + for ([_]mtl.Buffer{ t.w, t.s, t.b }) |x| x.deinit(); + } + + fn tensor(buffer: mtl.Buffer, dtype: tf.checkpoint.DType, c: usize, size: usize) tf.checkpoint.Tensor { + return .{ .buffer = buffer, .offset = 0, .bytes = vocab * c * size, .dtype = dtype, .shape = .{ vocab, c, 1, 1, 1 }, .rank = 2 }; + } + + /// What gpu_weights hands the embedding. + fn triple(t: Table) q27.projection.Triple { + const g = hidden / group; + return .{ .weight = tensor(t.w, .u32, hidden * t.bits / 32, 4), .scales = tensor(t.s, .bf16, g, 2), .biases = tensor(t.b, .bf16, g, 2), .bits = t.bits, .group = group }; + } + + /// Row `id` decoded in float64 by core.quant.mlx, its scales times `scale` (1 for the reference). + fn reference(t: Table, id: usize, scale: f64, out: []f64) !void { + const g = hidden / group; + const wb = hidden * t.bits / 8; + var s: [g * 4]u8 = undefined; // the control's scales, fp32 so 1% off is exactly that + for (0..g) |i| { + const v: f64 = @as(f32, @bitCast(@as(u32, t.s.slice(u16, vocab * g)[id * g + i]) << 16)); + std.mem.writeInt(u32, s[i * 4 ..][0..4], @bitCast(@as(f32, @floatCast(v * scale))), .little); + } + const shape: [5]usize = .{ 1, g, 1, 1, 1 }; + try mlx.reference(.{ + .words = .{ .dtype = .u32, .rank = 2, .shape = .{ 1, wb / 4, 1, 1, 1 }, .bytes = t.w.slice(u8, vocab * wb)[id * wb ..][0..wb] }, + .scales = if (scale == 1) .{ .dtype = .bf16, .rank = 2, .shape = shape, .bytes = t.s.slice(u8, vocab * g * 2)[id * g * 2 ..][0 .. g * 2] } else .{ .dtype = .f32, .rank = 2, .shape = shape, .bytes = &s }, + .biases = .{ .dtype = .bf16, .rank = 2, .shape = shape, .bytes = t.b.slice(u8, vocab * g * 2)[id * g * 2 ..][0 .. g * 2] }, + .bits = t.bits, + .group = group, + }, out); + } +}; + +fn commit(cb: mtl.CommandBuffer, 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; + } +} + +/// The family's gather of `rows` tokens in calls of `m` rows, into a NaN-filled bf16 output. +fn gather(device: mtl.Device, queue: mtl.Queue, k: q27.glue_gpu.Kernels, p: q27.projection.Triple, id: mtl.Buffer, m: usize) !mtl.Buffer { + const c = q27.config.Config{ .hidden = hidden, .intermediate = 17408, .layers = 64, .vocab = vocab, .heads = 24, .kv_heads = 4, .head_dim = 256, .k_heads = 16, .v_heads = 48, .dk = 128, .dv = 128, .conv_kernel = 4, .eps = 1e-6, .rope_dims = 64, .rope_theta = 1e7 }; + const y = try device.buffer(rows * hidden * 2, opts); + @memset(y.slice(u8, rows * hidden * 2), 0xff); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + var start: usize = 0; + while (start < rows) : (start += m) try k.embedding(e, c, p, .{ .buffer = id, .offset = start * 4 }, .{ .buffer = y, .offset = start * hidden * 2 }, @intCast(m)); + try commit(cb, e); + return y; +} + +/// Each width: gathered bits against float64 rounded once, one-row calls against one call, scales 1% off as control. +pub fn check(device: mtl.Device, queue: mtl.Queue) !usize { + const k = try q27.glue_gpu.Kernels.init(device); + defer k.deinit(); + var failed: usize = 0; + const deq = try a.alloc(f64, hidden); + defer a.free(deq); + for ([_]u8{ 8, 3 }, 0..) |bits, at| { + const tokens = ids(0xe27 + at); + const t = try Table.init(device, bits, &tokens, 0x27e + at); + defer t.deinit(); + const id = try device.buffer(rows * 4, opts); + defer id.deinit(); + @memcpy(id.slice(u32, rows), &tokens); + const all = try gather(device, queue, k, t.triple(), id, rows); + defer all.deinit(); + const one = try gather(device, queue, k, t.triple(), id, 1); + defer one.deinit(); + const got = all.slice(u16, rows * hidden); + var differ: usize = 0; + var ctl: usize = 0; + for (tokens, 0..) |token, r| { + try t.reference(token, 1, deq); + for (deq, got[r * hidden ..][0..hidden]) |d, g| differ += @intFromBool(g != bf64(d)); + try t.reference(token, 1.01, deq); + for (deq, got[r * hidden ..][0..hidden]) |d, g| ctl += @intFromBool(g != bf64(d)); + } + const same = std.mem.eql(u16, got, one.slice(u16, rows * hidden)); + const ok = differ == 0 and same and ctl * 10 > rows * hidden; + if (!ok) failed += 1; + std.debug.print("RESULT embedding [{d}, {d}] q{d} g64, {d} tokens: {d} of {d} values differ from float64 rounded once to bf16; " ++ + "one-row calls {s}; scales 1% off: {d} differ{s}\n", .{ vocab, hidden, bits, rows, differ, rows * hidden, if (same) "equal bit for bit" else "DIFFER", ctl, if (ok) "" else " FAIL" }); + } + return failed; +} diff --git a/zig/tests/qwen27/widths.zig b/zig/tests/qwen27/widths.zig index a1d78aaa7e..51d839fb3a 100644 --- a/zig/tests/qwen27/widths.zig +++ b/zig/tests/qwen27/widths.zig @@ -1,4 +1,4 @@ -//! The 27B's projections without tensor units at every oQ width against float64, row bits in every call, mixed stacks. +//! The 27B's projections, mixed stacks and embedding without tensor units at every oQ width against float64. const std = @import("std"); const mtl = @import("metal"); const tf = @import("tensorfold"); @@ -487,6 +487,7 @@ pub fn main(_: std.process.Init) !void { seed += 16; try stack(&t, g, s, seed); } + t.failed += try @import("embed.zig").check(device, queue); 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; From 4c7c799acf7bfb5064126e729dfe8bf875ae3c29 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 10:00:16 -1000 Subject: [PATCH 14/16] tests: qwen27-run --seams compares state, KV and logits bytes at prompt chunk seams and after decode --compare-chunk compared Session.fingerprint() (every GDN layer's recurrent and conv state, the visible KV rows of every attention layer, the last row's bf16 logits) at the prompt end only, and decode by tokens. With --compare-chunk C it now also compares the fingerprint after the decoded reply in both arms (the prefill to decode seam and every decode step's state), failing with DecodeStateOrLogitsDiffer. --seams N adds, for every N-th prompt row, a prefill of the prompt up to that row at --chunk and at C, fingerprints compared byte for byte; the receipt lists each seam and seams_byte_equal, and any difference fails with SeamStateOrLogitsDiffer. Test program only; nothing in the engine. M3 Ultra (no tensor units), the exact 2048-token public prefill array, chunk 128 against 16 and against 37, seams every 256 rows, 64 decoded tokens: every seam, the prompt end and the decoded state byte-equal on the uniform 4-bit checkpoint and on oQ4 and oQ2. --- zig/tests/qwen27/run.zig | 28 +++++++++++++++++++++++++--- 1 file changed, 25 insertions(+), 3 deletions(-) diff --git a/zig/tests/qwen27/run.zig b/zig/tests/qwen27/run.zig index 250f6ee2ab..bb1dfbdf80 100644 --- a/zig/tests/qwen27/run.zig +++ b/zig/tests/qwen27/run.zig @@ -10,12 +10,13 @@ const core_lane = @import("core_lane"); const help = \\tf-qwen27-run --run --model DIR --tokens IDS.json [--max-tokens 256] \\ [--chunk 128] [--compare-chunk 16] [--reference GENERATED.json] [--force-length] - \\ [--no-warmup] [--profile-steps N] [--teacher] + \\ [--no-warmup] [--profile-steps N] [--teacher] [--seams N] \\IDs and reference files are JSON arrays of unsigned token IDs, already chat-rendered. \\Greedy only, drafts off; JSON receipt on stdout, errors on stderr. \\--teacher feeds every reference token and records native top2 at all positions; margin timing is diagnostic. \\--max-tokens 0 reports prompt tok/s including the final single-row head. \\--compare-chunk reruns the whole prompt and decode and checks prompt state/KV/logits bytes. + \\--seams N (with --compare-chunk): state/KV/logits bytes every N prompt rows under both chunks and after decode. \\CPU only: --help or --emit-sources EXISTING_DIR, no Metal device or model load. ; @@ -85,6 +86,7 @@ pub fn main(init: std.process.Init) !void { var count: usize = 256; var chunk: usize = 128; var compare: usize = 0; + var seam_every: usize = 0; var run = false; var force = false; var warmup = true; @@ -111,9 +113,9 @@ pub fn main(init: std.process.Init) !void { if (i + 1 >= args.len) return error.MissingValue; i += 1; const value = args[i]; - if (std.mem.eql(u8, arg, "--model")) model = value else if (std.mem.eql(u8, arg, "--tokens")) file = value else if (std.mem.eql(u8, arg, "--reference")) reference = value else if (std.mem.eql(u8, arg, "--max-tokens")) count = try std.fmt.parseInt(usize, value, 10) else if (std.mem.eql(u8, arg, "--chunk")) chunk = try std.fmt.parseInt(usize, value, 10) else if (std.mem.eql(u8, arg, "--profile-steps")) profile_steps = try std.fmt.parseInt(usize, value, 10) else if (std.mem.eql(u8, arg, "--compare-chunk")) compare = try std.fmt.parseInt(usize, value, 10) else return error.UnknownOption; + if (std.mem.eql(u8, arg, "--model")) model = value else if (std.mem.eql(u8, arg, "--tokens")) file = value else if (std.mem.eql(u8, arg, "--reference")) reference = value else if (std.mem.eql(u8, arg, "--max-tokens")) count = try std.fmt.parseInt(usize, value, 10) else if (std.mem.eql(u8, arg, "--chunk")) chunk = try std.fmt.parseInt(usize, value, 10) else if (std.mem.eql(u8, arg, "--profile-steps")) profile_steps = try std.fmt.parseInt(usize, value, 10) else if (std.mem.eql(u8, arg, "--compare-chunk")) compare = try std.fmt.parseInt(usize, value, 10) else if (std.mem.eql(u8, arg, "--seams")) seam_every = try std.fmt.parseInt(usize, value, 10) else return error.UnknownOption; } - if (!run or model == null or file == null or chunk == 0 or chunk > 128 or compare > 128 or count > 65536) return error.BadOptions; + if (!run or model == null or file == null or chunk == 0 or chunk > 128 or compare > 128 or count > 65536 or (seam_every > 0 and compare == 0)) return error.BadOptions; const ids = try tokens(a, io, file.?); if (ids.len == 0 or ids.len + count > 262144) return error.BadPrompt; const expected = if (reference) |path| try tokens(a, io, path) else null; @@ -144,6 +146,7 @@ pub fn main(init: std.process.Init) !void { const generation = try generated(a, s, count, force, expected, if (profile_steps > 0) &trace else null, profile_steps, teacher); const out = generation.tokens; const decode_ns = decode_start.durationTo(std.Io.Clock.awake.now(io)).toNanoseconds(); + const decoded = if (compare > 0) try s.fingerprint() else digest; const reference_equal = if (expected) |wanted| std.mem.eql(u32, out, wanted) else false; if (compare > 0) { try runner.reset(0); @@ -151,6 +154,19 @@ pub fn main(init: std.process.Init) !void { if (!std.mem.eql(u8, &digest, &try s.fingerprint())) return error.PromptStateOrLogitsDiffer; const second = try generated(a, s, count, force, if (teacher) expected else null, null, 0, teacher); if (!std.mem.eql(u32, out, second.tokens)) return error.ChunkedTokensDiffer; + if (!std.mem.eql(u8, &decoded, &try s.fingerprint())) return error.DecodeStateOrLogitsDiffer; + } + const Seam = struct { at: usize, equal: bool }; + var seams: std.ArrayList(Seam) = .empty; + var at = seam_every; + while (seam_every > 0 and at < ids.len) : (at += seam_every) { + var sides: [2][32]u8 = undefined; + for ([_]usize{ chunk, compare }, &sides) |c, *side| { + try runner.reset(0); + try s.prefill(ids[0..at], c); + side.* = try s.fingerprint(); + } + try seams.append(a, .{ .at = at, .equal = std.mem.eql(u8, &sides[0], &sides[1]) }); } const digest_hex = std.fmt.bytesToHex(digest, .lower); const ProfileRecord = struct { class: []const u8, calls: u64, gpu_seconds: f64 }; @@ -168,6 +184,11 @@ pub fn main(init: std.process.Init) !void { .chunk = chunk, .compare_chunk = compare, .prompt_state_byte_equal = compare > 0, + .decode_state_byte_equal = compare > 0, + .seams = seams.items, + .seams_byte_equal = for (seams.items) |x| { + if (!x.equal) break false; + } else true, .reference_token_equal = reference_equal, .first_difference = generation.first_difference, .logit_dtype = "bf16", @@ -188,5 +209,6 @@ pub fn main(init: std.process.Init) !void { }, .{}); try std.Io.File.stdout().writeStreamingAll(io, receipt); try std.Io.File.stdout().writeStreamingAll(io, "\n"); + for (seams.items) |x| if (!x.equal) return error.SeamStateOrLogitsDiffer; if (!teacher and reference != null and !reference_equal) return error.ReferenceTokensDiffer; } From c090c1e1aa2e4815456758dd2cf20f581cb706a7 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Mon, 12 Oct 2026 09:18:22 -1000 Subject: [PATCH 15/16] The 27B keeps timed window costs per checkpoint widths, so an oQ checkpoint never drafts on another's costs Window costs are kept per build, chip and model shape, and every 27B checkpoint has one shape: whichever checkpoint a build served first wrote the table every later one read. An oQ2 server then chose its draft depth from the uniform 4-bit costs (wide windows nearly free) and drafted deeper than its row kernels pay for: same tokens, 28% more drafted rows, 15% slower code decode. A mixed-width checkpoint now adds a tag of its widths to the shape; a uniform 4-bit g64 checkpoint keeps its shape string, so its key and costs are unchanged. --- zig/src/families/qwen27/model.zig | 3 +++ zig/src/families/qwen27/widths.zig | 42 ++++++++++++++++++++++++++++++ zig/src/native/qwen27_lanes.zig | 3 +-- 3 files changed, 46 insertions(+), 2 deletions(-) diff --git a/zig/src/families/qwen27/model.zig b/zig/src/families/qwen27/model.zig index ea9f06a4de..62ddf72cfc 100644 --- a/zig/src/families/qwen27/model.zig +++ b/zig/src/families/qwen27/model.zig @@ -26,6 +26,8 @@ pub const Model = struct { frame: Frame, /// The Neural Engine's share of prompt-chunk MLPs, when it loaded. ane: ?*ane_mlp.Share = null, + /// widths.Widths.tag: 0 for a uniform 4-bit g64 checkpoint. + widths_tag: u64 = 0, /// `share_fits(bytes)`: whether the Neural Engine's share fits beside what is loaded (null: wherever it can load). pub fn load(a: std.mem.Allocator, io: std.Io, dir: []const u8, rows: u32, share_fits: ?*const fn (u64) bool) !*Model { @@ -43,6 +45,7 @@ pub const Model = struct { errdefer a.destroy(m); m.allocator = a; m.config = c; + m.widths_tag = ws.tag(); m.device = try mtl.Device.init(); errdefer m.device.deinit(); if (!qualifiedGeneration(m.device.generation())) return error.QwenMetalQualificationRequired; diff --git a/zig/src/families/qwen27/widths.zig b/zig/src/families/qwen27/widths.zig index ada7ce0f64..d435014c4a 100644 --- a/zig/src/families/qwen27/widths.zig +++ b/zig/src/families/qwen27/widths.zig @@ -105,6 +105,21 @@ pub const Widths = struct { pub fn of(w: *const Widths, layer: usize, m: Member) Width { return w.layers[layer][@backingInt(m)]; } + + /// 0 when every module is 4-bit g64, else a nonzero hash of every module's width. + pub fn tag(w: *const Widths) u64 { + var h = std.hash.Wyhash.init(0x27); + var mixed = false; + for (w.layers) |l| for (l) |x| { + mixed = mixed or !isUniform(x); + h.update(&.{ x.bits, @as(u8, @truncate(x.group)), @as(u8, @truncate(x.group >> 8)) }); + }; + for ([_]Width{ w.embed, w.head }) |x| { + mixed = mixed or !isUniform(x); + h.update(&.{ x.bits, @as(u8, @truncate(x.group)), @as(u8, @truncate(x.group >> 8)) }); + } + return if (mixed) h.final() | 1 else 0; + } }; fn one(f: *const affine.Formats, path: []const u8) !Width { @@ -112,6 +127,13 @@ fn one(f: *const affine.Formats, path: []const u8) !Width { return .{ .bits = s.bits, .group = @intCast(s.group_size) }; } +/// The timed window costs' model shape; mixed widths add their tag, so checkpoints of one shape keep their own costs. +pub fn costShape(c: cfg.Config, tag: u64, buf: []u8, rows: usize, drafted: bool) ![]const u8 { + const s = try std.fmt.bufPrint(buf, "{d}/{d}/{d}/{d}/{d}/{d}", .{ c.layers, c.hidden, c.intermediate, c.vocab, rows, @intFromBool(drafted) }); + if (tag == 0) return s; + return buf[0 .. s.len + (try std.fmt.bufPrint(buf[s.len..], "/w{x:0>16}", .{tag})).len]; +} + /// What the chip's projection kernels run: the lane kernels on tensor units, the row kernels before them. pub const Chip = struct { tensor_units: bool, row_bits: []const u8 = &row.bit_widths }; @@ -242,6 +264,26 @@ test "a mixed stack runs every member alone, even beside a member of its own wid try std.testing.expectEqualSlices([2]usize, &.{ .{ 0, 1 }, .{ 1, 1 }, .{ 2, 1 } }, runs(&.{ four, five, four }, &buf)); } +test "timed costs keep a uniform checkpoint's shape and give each mixed one its own" { + const a = std.testing.allocator; + var tags: [3]u64 = undefined; + for ([_][]const u8{ "{\"bits\":4,\"group_size\":64}", oq4, oq2 }, &tags) |quantization, *t| { + const text = try testConfig(a, quantization); + defer a.free(text); + const c = try cfg.parse(a, text, null); + var f = try affine.Formats.init(a, text); + defer f.deinit(); + var w = try Widths.resolve(a, &f, c); + defer w.deinit(); + t.* = w.tag(); + var buf: [96]u8 = undefined; + const got = try costShape(c, t.*, &buf, 16, true); + try std.testing.expectEqual(t.* == 0, std.mem.eql(u8, got, "64/5120/17408/248320/16/1")); + } + try std.testing.expectEqual(@as(u64, 0), tags[0]); + try std.testing.expect(tags[1] != 0 and tags[2] != 0 and tags[1] != tags[2]); +} + test "resolved widths follow overrides under any checkpoint prefix" { const a = std.testing.allocator; const text = try testConfig(a, oq4); diff --git a/zig/src/native/qwen27_lanes.zig b/zig/src/native/qwen27_lanes.zig index c177ef95fa..e09a508a8c 100644 --- a/zig/src/native/qwen27_lanes.zig +++ b/zig/src/native/qwen27_lanes.zig @@ -148,9 +148,8 @@ pub const Metal = struct { pub fn measure(b: *Metal, io: std.Io) !void { const r = b.runner; if (b.owners[0] != null or r.active != null) return error.StreamBusy; - const c = r.model.config; var shape: [128]u8 = undefined; - const parts = [_][]const u8{ "qwen27-lanes", std.mem.span(r.model.device.name()), try std.fmt.bufPrint(&shape, "{d}/{d}/{d}/{d}/{d}/{d}", .{ c.layers, c.hidden, c.intermediate, c.vocab, max_rows, @intFromBool(b.draft != null) }) }; + const parts = [_][]const u8{ "qwen27-lanes", std.mem.span(r.model.device.name()), try q.widths.costShape(r.model.config, r.model.widths_tag, &shape, max_rows, b.draft != null) }; const k = lanes.cost_cache.key(b.gpa, io, &parts) catch null; if (k) |key| if (lanes.cost_cache.load([max_rows]lanes.config.Cost, b.gpa, io, key)) |kept| { b.costs = kept; From a2420e759f434abb6771e0dd5bf9b0bcad3d005b Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Tue, 13 Oct 2026 09:13:51 -1000 Subject: [PATCH 16/16] The 27B runs a run whose rows a tensor-unit tile does not divide on the row kernels, on every chip With tensor units, projection.prepare refused any run with n % tile != 0 (BadProjectionShape) and widths.check gave every 4/64 run the reg lane layout, which needs n % 32 == 0. A lone 48-row GDN in_proj_a or in_proj_b is such a run: TensorFold/Qwen3.8-27B-oQ4's layer 37 (in_proj_a at 5/64, in_proj_b a 48-row 4/64 run) was refused at admission and -oQ2 failed at load on M5. Those runs now take the row kernels, as they do before tensor units; mixed stacks already send a row run to rowStrided. A 4/64 run off the 32-row lanes stays out of regOrder. Before tensor units nothing changes (those runs already took prepareRows), and a uniform 4/64 checkpoint's matrices are multiples of 32 rows. Test: widths.zig: a 48-row 4/64 run admits on tensor units through the row kernels, a 64-row one through the lanes (fails before: the lane layout refused n % 32 != 0). Reported with the fix by Seth For Privacy (@sethforprivacy) on an M5 Ultra (#756), with oQ4 and oQ2 served exactly there. --- zig/src/families/qwen27/gpu_weights.zig | 2 +- zig/src/families/qwen27/projection.zig | 3 +-- zig/src/families/qwen27/widths.zig | 7 ++++++- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/zig/src/families/qwen27/gpu_weights.zig b/zig/src/families/qwen27/gpu_weights.zig index 5ab7edb8da..bb0fba68e3 100644 --- a/zig/src/families/qwen27/gpu_weights.zig +++ b/zig/src/families/qwen27/gpu_weights.zig @@ -226,7 +226,7 @@ fn laneOrder(a: std.mem.Allocator, device: mtl.Device, layers: []const Layer, he /// Only 4-bit g64 matrices take the reg kernel; other widths keep lane_projection.weight_layout. fn regJobs(a: std.mem.Allocator, jobs: *std.ArrayList(lane.Order), p: Linear) std.mem.Allocator.Error!void { for (p.runs) |r| try regJobs(a, jobs, r.linear); - if (p.runs.len == 0 and p.bits == 4 and p.group == 64) try jobs.append(a, .{ .w = .{ .buf = p.words.buffer, .off = p.words.offset }, .n = p.n, .k = p.k }); + if (p.runs.len == 0 and p.bits == 4 and p.group == 64 and p.n % 32 == 0) try jobs.append(a, .{ .w = .{ .buf = p.words.buffer, .off = p.words.offset }, .n = p.n, .k = p.k }); } pub fn ref(t: Tensor) Ref { diff --git a/zig/src/families/qwen27/projection.zig b/zig/src/families/qwen27/projection.zig index dbbd13ae8a..cd3834e88d 100644 --- a/zig/src/families/qwen27/projection.zig +++ b/zig/src/families/qwen27/projection.zig @@ -48,8 +48,7 @@ pub fn prepare(device: mtl.Device, allocator: std.mem.Allocator, owned: *std.Arr if (p.scales.dtype != .bf16 or p.biases.dtype != .bf16 or p.scales.rank != 2 or p.biases.rank != 2 or p.scales.shape[0] != rows or p.biases.shape[0] != rows or p.scales.shape[1] != k / group or p.biases.shape[1] != k / group) return error.BadAffineMetadata; n = try std.math.add(usize, n, rows); } - if (!device.tensorUnits()) return prepareRows(device, allocator, owned, parts, n, k, tile, slices); - if (n % tile != 0) return error.BadProjectionShape; + if (!device.tensorUnits() or n % tile != 0) return prepareRows(device, allocator, owned, parts, n, k, tile, slices); if (bits != 4 or group != 64) return prepareLanes(device, allocator, owned, parts, n, k, tile, slices); slices = m5Split(n, k) orelse slices; const w = try device.buffer(n * k / 2, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); diff --git a/zig/src/families/qwen27/widths.zig b/zig/src/families/qwen27/widths.zig index d435014c4a..6de398e809 100644 --- a/zig/src/families/qwen27/widths.zig +++ b/zig/src/families/qwen27/widths.zig @@ -171,7 +171,7 @@ fn check(chip: Chip, n: usize, k: usize, w: Width) !void { _ = try @import("projection.zig").split(n, k); // The lane sums and odd word counts assume group 64 off the 4-bit kernel. if (chip.tensor_units and w.bits != 4 and w.group != 64) return error.UnsupportedLaneGroup; - if (chip.tensor_units) return (lane.Layout{ .n = n, .k = k, .format = .{ .bits = w.bits, .group = w.group }, .precompute_sums = isUniform(w), .reg = isUniform(w) }).validate(); + if (chip.tensor_units and n % 32 == 0) return (lane.Layout{ .n = n, .k = k, .format = .{ .bits = w.bits, .group = w.group }, .precompute_sums = isUniform(w), .reg = isUniform(w) }).validate(); if (isUniform(w) and n % 8 == 0) return; if (std.mem.indexOfScalar(u8, chip.row_bits, w.bits) == null) return error.UnsupportedRowBits; return (row.Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = n, .k = k, .group = w.group }).validate(); @@ -255,6 +255,11 @@ test "a dense projection and an unsupported global group refuse" { try std.testing.expectError(error.BadAffine, testAdmit("{\"bits\":4,\"group_size\":16}", .{ .tensor_units = false }, &why)); } +test "a 4/64 run off the 32-row lanes, such as a lone 48-row in_proj_b, takes the row kernels on tensor units" { + try check(.{ .tensor_units = true }, 48, 5120, uniform); + try check(.{ .tensor_units = true }, 64, 5120, uniform); +} + test "a mixed stack runs every member alone, even beside a member of its own width" { var buf: [3][2]usize = undefined; const four = uniform;