diff --git a/build.zig b/build.zig index 1bf8bfa260..8fd7b6e7d6 100644 --- a/build.zig +++ b/build.zig @@ -3,6 +3,8 @@ 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 glm_checks = @import("zig/build/glm_checks.zig"); const hip_build = @import("zig/build/hip.zig"); const grammar_build = @import("zig/build/grammar.zig"); const counters_build = @import("zig/build/metal_counters.zig"); @@ -455,6 +457,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); @@ -474,6 +477,8 @@ fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b // The core lane, shared-context and Qwen3.8-27B programs and host tests. qwen27_build.targets(b, target, optimize, metal, sources, engine, lanes, mods.qwen27, test_step); + core_checks.targets(b, target, optimize, metal, engine); + glm_checks.targets(b, target, optimize, metal, engine, sources); _ = nativeServer(b, target, optimize, metal, engine, api, engines, build_options, test_step); diff --git a/zig/build/core_checks.zig b/zig/build/core_checks.zig new file mode 100644 index 0000000000..6749819464 --- /dev/null +++ b/zig/build/core_checks.zig @@ -0,0 +1,32 @@ +//! 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); +} + +/// 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/build/glm_checks.zig b/zig/build/glm_checks.zig new file mode 100644 index 0000000000..bf65167ddd --- /dev/null +++ b/zig/build/glm_checks.zig @@ -0,0 +1,25 @@ +//! GLM-5.3-Flash's device checks on synthetic weights: programs that run on the GPU, each under its own step. + +const std = @import("std"); + +/// tf-glm-rows-check: the products at every MLX width against float64, every row's bits in every window and launch. +pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, metal: *std.Build.Module, engine: *std.Build.Module, sources: *std.Build.Module) void { + const exe = b.addExecutable(.{ .name = "tf-glm-rows-check", .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/glm_rows_check.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "tensorfold", .module = engine }, .{ .name = "kernel_sources", .module = sources } }, + }) }); + b.installArtifact(exe); + const about = "GLM-5.3-Flash's products at every MLX width on synthetic weights against float64, each row's bits in every launch"; + b.step("tf-glm-rows-check", about).dependOn(&b.addInstallArtifact(exe, .{}).step); + const speed = b.addExecutable(.{ .name = "tf-glm-rows-speed", .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/glm_rows_speed.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "tensorfold", .module = engine }, .{ .name = "kernel_sources", .module = sources } }, + }) }); + b.step("tf-glm-rows-speed", "GLM-5.3-Flash one-row decode products at the checkpoints' shapes: GPU time a launch, GB/s, a token by class").dependOn(&b.addInstallArtifact(speed, .{}).step); +} diff --git a/zig/kernels/metal/core/affine_mm.metal b/zig/kernels/metal/core/affine_mm.metal index 16e54006d2..612c7e936a 100644 --- a/zig/kernels/metal/core/affine_mm.metal +++ b/zig/kernels/metal/core/affine_mm.metal @@ -4,7 +4,7 @@ using namespace metal; #include "../nax.h" using namespace tfp; -// Defines: TF_BITS (4 or 8), TF_GROUP (32 or 64), TF_OUT_T (bfloat, or float for partial sums another pass adds). +// Defines: TF_BITS (2-6 or 8), TF_GROUP (32 or 64), TF_OUT_T (bfloat, or float for partial sums another pass adds). #ifndef TF_OUT_T #define TF_OUT_T bfloat #endif @@ -66,7 +66,29 @@ 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; words at 3, 5 and 6 bits). +#if TF_BITS == 2 +struct Codes { + uint2 v; +}; + +inline Codes am_load(const device uchar* w) { + return Codes{*(const device uint2*)w}; +} +#elif TF_BITS % 4 != 0 +struct Codes { + uint v[TF_BITS]; +}; + +inline Codes am_load(const device uchar* w) { + Codes c; + TF_UNROLL + for (short i = 0; i < TF_BITS; i++) { + c.v[i] = ((const device uint*)w)[i]; + } + return c; +} +#else struct Codes { uint4 v[TF_BITS / 4]; }; @@ -79,11 +101,38 @@ 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 != 0 + TF_UNROLL + for (short i = 0; i < 8; i++) { // codes 4i .. 4i + 3 of the little-endian bit stream, as themselves + float f[4]; + TF_UNROLL + for (short j = 0; j < 4; j++) { + const short at = TF_BITS * (4 * i + j), word = at / 32, shift = at % 32; + uint v = c.v[word] >> shift; + if (shift + TF_BITS > 32) v |= c.v[word + 1] << (32 - shift); + f[j] = float(v & ((1u << TF_BITS) - 1u)); + } + o[i / 2][2 * (i % 2)] = as_type(ushort2(as_type(bfloat(f[0])), as_type(bfloat(f[1])))); + o[i / 2][2 * (i % 2) + 1] = as_type(ushort2(as_type(bfloat(f[2])), as_type(bfloat(f[3])))); + } +#elif TF_BITS == 4 TF_UNROLL for (short i = 0; i < 4; i++) { const uint w = c.v[0][i]; 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/kernels/metal/core/row_projection.metal b/zig/kernels/metal/core/row_projection.metal index dae395aeb0..d985ee1e7f 100644 --- a/zig/kernels/metal/core/row_projection.metal +++ b/zig/kernels/metal/core/row_projection.metal @@ -1,7 +1,10 @@ -// 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; +// 2^-e, exact: a code read in place at bit e times its input scaled by 2^-e is the same product. +inline float tf_rp_p2(int e) { return as_type(uint(127 - e) << 23); } + // 16 bf16 inputs as floats (scaled to the nibble positions at 4 bits) and their group sum: fp32 in order, or rounded to bf16 at each add (BF16SUM, the Nemotron expert convention). template inline float tf_rp_load16(const device bfloat* x, thread float* xt) { @@ -10,22 +13,22 @@ inline float tf_rp_load16(const device bfloat* x, thread float* xt) { const bfloat a = x[i], b = x[i + 1], c = x[i + 2], d = x[i + 3]; if (BF16SUM) sum += float(bfloat(float(bfloat(float(bfloat(float(a) + float(b))) + float(c))) + float(d))); else { sum += float(a); sum += float(b); sum += float(c); sum += float(d); } - xt[i] = float(a); - xt[i + 1] = BITS == 4 ? float(b) / 16.0f : float(b); - xt[i + 2] = BITS == 4 ? float(c) / 256.0f : float(c); - xt[i + 3] = BITS == 4 ? float(d) / 4096.0f : float(d); + xt[i] = BITS == 2 ? float(a) * tf_rp_p2(2 * i) : float(a); + xt[i + 1] = BITS == 4 ? float(b) / 16.0f : BITS == 2 ? float(b) * tf_rp_p2(2 * i + 2) : float(b); + xt[i + 2] = BITS == 4 ? float(c) / 256.0f : BITS == 2 ? float(c) * tf_rp_p2(2 * i + 4) : float(c); + xt[i + 3] = BITS == 4 ? float(d) / 4096.0f : BITS == 2 ? float(d) * tf_rp_p2(2 * i + 6) : float(d); } return sum; } -// 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) { float q[16]; if (BITS == 2) { const uint v = *(const device uint*)w; - for (int i = 0; i < 16; i++) q[i] = float((v >> (2 * i)) & 3u); + for (int i = 0; i < 16; i++) q[i] = float(v & (3u << (2 * i))); } else if (BITS == 4) { const device uint16_t* ws = (const device uint16_t*)w; for (int i = 0; i < 4; i++) { @@ -40,6 +43,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 +110,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 +121,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 +136,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 +212,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/kernels/metal/glm_rows.metal b/zig/kernels/metal/glm_rows.metal new file mode 100644 index 0000000000..9c87b9b592 --- /dev/null +++ b/zig/kernels/metal/glm_rows.metal @@ -0,0 +1,243 @@ +// GLM-5.3-Flash's products at any MLX width: core/row_projection.metal's row sums (compiled first), any window. + +struct GlmRowDims { + int k, n, group, rows; // a matrix [n, k] at `group`; x rows + int heads, ldx, hsx; // matrices back to back (one a head); x's row and head strides + int ldo, hso, hsw; // out's row and head strides; a head's matrix rows +}; + +struct GlmPickDims { + int k, n, group, slots; // slots: (row, pick) pairs, or sorted pairs + int repeat, experts, ldo; // a row's picks (x row = slot / repeat); experts in W; out's row stride + float lim; // the SwiGLU clip +}; + +// The MoE kernels' SwiGLU of a gate and up sum: glm_swiglu's, its sigmoid_fast put first (wide.zig glue_math). +inline bfloat glm_rows_glu(float gv, float uv, float limit) { + const float lim = float(bfloat(limit)); + const bfloat gt = bfloat(metal::min(float(bfloat(gv)), lim)); + const bfloat up = bfloat(metal::min(metal::max(float(bfloat(uv)), -lim), lim)); + return (gt * sigmoid_fast(gt)) * up; +} + +// out[m, h, c] = x[m, h] . W[h][c]: 4 columns a simdgroup, 2 simdgroups a threadgroup; grid (n / 8, heads). +template +[[kernel]] void glm_rows(const device bfloat* X [[buffer(0)]], const device uint32_t* W [[buffer(1)]], + const device bfloat* S [[buffer(2)]], const device bfloat* B [[buffer(3)]], + constant GlmRowDims& d [[buffer(4)]], device bfloat* OUT [[buffer(5)]], + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int row0 = (int(tg.x) * 2 + int(sg)) * 4; + if (row0 >= d.n) return; + const int h = int(tg.y); + const size_t at = size_t(h) * d.hsw + row0; + const device uint8_t* wr = (const device uint8_t*)W + at * (d.k * BITS / 8); + const device bfloat* sr = S + at * (d.k / d.group); + const device bfloat* br = B + at * (d.k / d.group); + const device bfloat* x0 = X + size_t(h) * d.hsx; + device bfloat* o0 = OUT + size_t(h) * d.hso + row0; + const int cols = metal::min(4, d.n - row0); + int m = 0; + if (d.ldx == d.k) { + #pragma clang loop unroll(disable) + for (; m + 2 <= d.rows; m += 2) { + float acc[2][4]; + tf_rp_rows<4, 2, BITS, false>(wr, sr, br, x0 + size_t(m) * d.k, d.k, d.group, lane, acc); + if (lane == 0) + for (int b = 0; b < 2; b++) + for (int j = 0; j < cols; j++) o0[size_t(m + b) * d.ldo + j] = bfloat(acc[b][j]); + } + } + #pragma clang loop unroll(disable) + for (; m < d.rows; m++) { + float acc[1][4]; + tf_rp_rows<4, 1, BITS, false>(wr, sr, br, x0 + size_t(m) * d.ldx, d.k, d.group, lane, acc); + if (lane == 0) + for (int j = 0; j < cols; j++) o0[size_t(m) * d.ldo + j] = bfloat(acc[0][j]); + } +} + +// Slot y through expert ids[y] on x row y / repeat (GLU: gate W and up W2, SwiGLU'd); grid (n / 8, slots). +template +[[kernel]] void glm_rows_ids(const device bfloat* X [[buffer(0)]], const device uint32_t* W [[buffer(1)]], + const device bfloat* S [[buffer(2)]], const device bfloat* B [[buffer(3)]], + constant GlmPickDims& d [[buffer(4)]], device bfloat* OUT [[buffer(5)]], + const device int* ids [[buffer(6)]], const device uint32_t* W2 [[buffer(7)]], + const device bfloat* S2 [[buffer(8)]], const device bfloat* B2 [[buffer(9)]], + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int row0 = (int(tg.x) * 2 + int(sg)) * 4; + if (row0 >= d.n) return; + const int slot = int(tg.y); + const int expert = ids[slot]; + device bfloat* out = OUT + size_t(slot) * d.ldo + row0; + const int cols = metal::min(4, d.n - row0); + if (expert < 0 || expert >= d.experts) { + if (int(lane) < cols) out[lane] = bfloat(0.0f); + return; + } + const size_t at = size_t(expert) * d.n + row0; + const size_t kb = size_t(d.k) * BITS / 8, kg = size_t(d.k / d.group); + const device bfloat* x = X + size_t(slot / d.repeat) * d.k; + float acc[1][4]; + tf_rp_rows<4, 1, BITS, false>((const device uint8_t*)W + at * kb, S + at * kg, B + at * kg, x, d.k, d.group, lane, acc); + if (GLU) { + float up[1][4]; + tf_rp_rows<4, 1, BITS, false>((const device uint8_t*)W2 + at * kb, S2 + at * kg, B2 + at * kg, x, d.k, d.group, lane, up); + if (lane == 0) + for (int j = 0; j < cols; j++) out[j] = glm_rows_glu(acc[0][j], up[0][j], d.lim); + } else if (lane == 0) { + for (int j = 0; j < cols; j++) out[j] = bfloat(acc[0][j]); + } +} + +// Slots as glm_rows_ids (BITS), then X2's rows through the shared expert (SBITS; GLU: gate then up rows) into OUT2. +template // RPS columns a simdgroup; KC: k compiled in +[[kernel]] void glm_moe_rows(const device bfloat* X [[buffer(0)]], const device uint32_t* W [[buffer(1)]], + const device bfloat* S [[buffer(2)]], const device bfloat* B [[buffer(3)]], + constant GlmPickDims& d [[buffer(4)]], device bfloat* OUT [[buffer(5)]], + const device int* ids [[buffer(6)]], const device uint32_t* W2 [[buffer(7)]], + const device bfloat* S2 [[buffer(8)]], const device bfloat* B2 [[buffer(9)]], + const device bfloat* X2 [[buffer(10)]], const device uint32_t* SW [[buffer(11)]], + const device bfloat* SS [[buffer(12)]], const device bfloat* SB [[buffer(13)]], + device bfloat* OUT2 [[buffer(14)]], + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int row0 = (int(tg.x) * 2 + int(sg)) * RPS; + if (row0 >= d.n) return; + const int y = int(tg.y); + const int cols = metal::min(RPS, d.n - row0); + const int K = KC > 0 ? KC : d.k; + const size_t kg = size_t(K / d.group); + const device uint8_t* gw; + const device bfloat *gs, *gb, *x; + device bfloat* out; + size_t up_at = 0; // the up rows' offset from the gate rows, in rows + if (y < d.slots) { + const int expert = ids[y]; + out = OUT + size_t(y) * d.ldo + row0; + if (expert < 0 || expert >= d.experts) { + if (int(lane) < cols) out[lane] = bfloat(0.0f); + return; + } + const size_t at = size_t(expert) * d.n + row0; + const size_t kb = size_t(K) * BITS / 8; + x = X + size_t(y / d.repeat) * K; + float acc[1][RPS]; + tf_rp_rows((const device uint8_t*)W + at * kb, S + at * kg, B + at * kg, x, K, d.group, lane, acc); + if (GLU) { + float up[1][RPS]; + tf_rp_rows((const device uint8_t*)W2 + at * kb, S2 + at * kg, B2 + at * kg, x, K, d.group, lane, up); + if (lane == 0) + for (int j = 0; j < cols; j++) out[j] = glm_rows_glu(acc[0][j], up[0][j], d.lim); + } else if (lane == 0) { + for (int j = 0; j < cols; j++) out[j] = bfloat(acc[0][j]); + } + return; + } + const int r = y - d.slots; + const size_t skb = size_t(K) * SBITS / 8; + x = X2 + size_t(r) * K; + out = OUT2 + size_t(r) * d.n + row0; + gw = (const device uint8_t*)SW + size_t(row0) * skb; + gs = SS + size_t(row0) * kg; + gb = SB + size_t(row0) * kg; + float acc[1][RPS]; + tf_rp_rows(gw, gs, gb, x, K, d.group, lane, acc); + if (GLU) { + up_at = size_t(d.n); + float up[1][RPS]; + tf_rp_rows(gw + up_at * skb, gs + up_at * kg, gb + up_at * kg, x, K, d.group, lane, up); + if (lane == 0) + for (int j = 0; j < cols; j++) out[j] = glm_rows_glu(acc[0][j], up[0][j], d.lim); + } else if (lane == 0) { + for (int j = 0; j < cols; j++) out[j] = bfloat(acc[0][j]); + } +} + +// Expert e's sorted pairs: GLU x row order[i] / repeat to out row i, else x row i to row order[i]; grid (n / 8, e). +template +[[kernel]] void glm_rows_sorted(const device bfloat* X [[buffer(0)]], const device uint32_t* W [[buffer(1)]], + const device bfloat* S [[buffer(2)]], const device bfloat* B [[buffer(3)]], + constant GlmPickDims& d [[buffer(4)]], device bfloat* OUT [[buffer(5)]], + const device int* offsets [[buffer(6)]], const device int* order [[buffer(7)]], + const device uint32_t* W2 [[buffer(8)]], const device bfloat* S2 [[buffer(9)]], + const device bfloat* B2 [[buffer(10)]], + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int row0 = (int(tg.x) * 2 + int(sg)) * 4; + if (row0 >= d.n) return; + const int e = int(tg.y); + const int lo = offsets[e], hi = offsets[e + 1]; + const size_t at = size_t(e) * d.n + row0; + const size_t kb = size_t(d.k) * BITS / 8, kg = size_t(d.k / d.group); + const device uint8_t* w = (const device uint8_t*)W + at * kb; + const int cols = metal::min(4, d.n - row0); + #pragma clang loop unroll(disable) + for (int i = lo; i < hi; i++) { + const int src = GLU ? order[i] / d.repeat : i; + const int dst = GLU ? i : order[i]; + float acc[1][4]; + tf_rp_rows<4, 1, BITS, false>(w, S + at * kg, B + at * kg, X + size_t(src) * d.k, d.k, d.group, lane, acc); + device bfloat* out = OUT + size_t(dst) * d.ldo + row0; + if (GLU) { + float up[1][4]; + tf_rp_rows<4, 1, BITS, false>((const device uint8_t*)W2 + at * kb, S2 + at * kg, B2 + at * kg, X + size_t(src) * d.k, d.k, d.group, lane, up); + if (lane == 0) + for (int j = 0; j < cols; j++) out[j] = glm_rows_glu(acc[0][j], up[0][j], d.lim); + } else if (lane == 0) { + for (int j = 0; j < cols; j++) out[j] = bfloat(acc[0][j]); + } + } +} + +// out[m, c] = x[m] . W[c] for bf16 W [n, k] (an unquantized eh_proj): fp32 sums, lanes over 16-input pieces in order. +[[kernel]] void glm_rows_bf16(const device bfloat* X [[buffer(0)]], const device bfloat* W [[buffer(1)]], + constant GlmRowDims& d [[buffer(4)]], device bfloat* OUT [[buffer(5)]], + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int row0 = (int(tg.x) * 2 + int(sg)) * 4; + if (row0 >= d.n) return; + const int cols = metal::min(4, d.n - row0); + #pragma clang loop unroll(disable) + for (int m = 0; m < d.rows; m++) { + const device bfloat* x = X + size_t(m) * d.ldx; + float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (int k0 = int(lane) * 16; k0 < d.k; k0 += 512) { + float xt[16]; + for (int i = 0; i < 16; i++) xt[i] = float(x[k0 + i]); + for (int j = 0; j < cols; j++) { + const device bfloat* w = W + size_t(row0 + j) * d.k + k0; + float s = 0.0f; + for (int i = 0; i < 16; i++) s += xt[i] * float(w[i]); + acc[j] += s; + } + } + for (int j = 0; j < 4; j++) acc[j] = simd_sum(acc[j]); + if (lane == 0) + for (int j = 0; j < cols; j++) OUT[size_t(m) * d.ldo + row0 + j] = bfloat(acc[j]); + } +} + +#define GLM_ROWS_AT(B) \ + template [[host_name("glm_rows_q" #B)]] [[kernel]] decltype(glm_rows) glm_rows; \ + template [[host_name("glm_rows_ids_q" #B)]] [[kernel]] decltype(glm_rows_ids) glm_rows_ids; \ + template [[host_name("glm_rows_glu_ids_q" #B)]] [[kernel]] decltype(glm_rows_ids) glm_rows_ids; \ + template [[host_name("glm_rows_sorted_q" #B)]] [[kernel]] decltype(glm_rows_sorted) glm_rows_sorted; \ + template [[host_name("glm_rows_glu_sorted_q" #B)]] [[kernel]] decltype(glm_rows_sorted) glm_rows_sorted; +GLM_ROWS_AT(2) +GLM_ROWS_AT(3) +GLM_ROWS_AT(4) +GLM_ROWS_AT(5) +GLM_ROWS_AT(6) +GLM_ROWS_AT(8) + +#define GLM_MOE_AT(B) \ + template [[host_name("glm_moe_rows_q" #B "_s8")]] [[kernel]] decltype(glm_moe_rows) glm_moe_rows; \ + template [[host_name("glm_moe_rows_glu_q" #B "_s8")]] [[kernel]] decltype(glm_moe_rows) glm_moe_rows; +GLM_MOE_AT(2) +GLM_MOE_AT(3) +GLM_MOE_AT(4) +GLM_MOE_AT(5) +GLM_MOE_AT(6) +GLM_MOE_AT(8) diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index 6313b656cd..a1113ada56 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -82,6 +82,8 @@ pub const flashnext_gdn = @embedFile("decode/fn_gdn.metal"); pub const glm = @import("glm/kernels.zig"); /// GLM-5.3-Flash's glue: MLX's arithmetic where the Python family calls MLX ops; selection, argmax, small steps. pub const glm_glue = @embedFile("glm_glue.metal"); +/// Its products at the widths the generated kernels do not read, after core_row_projection (glm/wide.zig). +pub const glm_rows = @embedFile("glm_rows.metal"); /// GLM-5.3-Flash's latent attention for rows that read every key: 64 heads as one matrix on the tensor units. pub const glm_attn = @embedFile("glm_attn.metal"); /// A prompt chunk's KDA layer in three passes, appended to the generated kda_rows source (its helpers). diff --git a/zig/src/core/affine_mm.zig b/zig/src/core/affine_mm.zig index dd25dce831..25386372a8 100644 --- a/zig/src/core/affine_mm.zig +++ b/zig/src/core/affine_mm.zig @@ -4,12 +4,12 @@ const mtl = @import("metal"); const ks = @import("kernel_sources"); const frags = @import("frags.zig"); -/// The weights' bits (4 or 8) and group (32 or 64), and the output's type. +/// The weights' bits (2, 3, 4, 5, 6 or 8) and group (32 or 64), and the output's type. pub const Format = struct { bits: u8 = 4, group: u16 = 64, out_f32: bool = false }; /// The kernel source for `f` (nax.h inlined for this GPU). pub fn source(device: mtl.Device, a: std.mem.Allocator, f: Format) ![]u8 { - if ((f.bits != 4 and f.bits != 8) or (f.group != 32 and f.group != 64)) return error.UnsupportedFormat; + if (f.bits < 2 or f.bits > 8 or f.bits == 7 or (f.group != 32 and f.group != 64)) return error.UnsupportedFormat; const body = try frags.source(device, a, ks.core_affine_mm); defer a.free(body); return std.fmt.allocPrint(a, "#define TF_BITS {d}\n#define TF_GROUP {d}\n#define TF_OUT_T {s}\n{s}", .{ f.bits, f.group, if (f.out_f32) "float" else "bfloat", body }); @@ -111,6 +111,15 @@ pub fn gatherGlu(e: mtl.ComputeEncoder, p: Pipes, x: anytype, sums: anytype, gat e.dispatchGroups(mtl.Size.of(gate.n / 32, tiles, 1), mtl.Size.of(32, 2, 2)); } +/// The code at bit `at` of a little-endian bit stream (MLX's packing at every width). +fn code(bytes: []const u8, at: usize, bits: usize) u32 { + var v: u32 = 0; + for (0..3) |i| if (at / 8 + i < bytes.len) { + v |= @as(u32, bytes[at / 8 + i]) << @intCast(8 * i); + }; + return (v >> @intCast(at % 8)) & ((@as(u32, 1) << @intCast(bits)) - 1); +} + fn bf16(v: f32) u16 { // round to nearest even const u: u32 = @bitCast(v); return @intCast((u + 0x7fff + ((u >> 16) & 1)) >> 16); @@ -135,7 +144,7 @@ test "dense and gathered, each format, against exact sums on the host" { const R = counts[0] + counts[1] + counts[2]; var prng = std.Random.DefaultPrng.init(7); const rnd = prng.random(); - for ([_][2]usize{ .{ 4, 64 }, .{ 4, 32 }, .{ 8, 64 }, .{ 8, 32 } }) |fmt| { + for ([_][2]usize{ .{ 4, 64 }, .{ 4, 32 }, .{ 8, 64 }, .{ 8, 32 }, .{ 2, 64 }, .{ 2, 32 }, .{ 3, 64 }, .{ 5, 64 }, .{ 6, 64 }, .{ 5, 32 }, .{ 6, 32 } }) |fmt| { const bits = fmt[0]; const group = fmt[1]; const w_bytes = E * N * K * bits / 8; @@ -207,7 +216,7 @@ test "dense and gathered, each format, against exact sums on the host" { const off: f64 = if (bits == 4) 128 else 0; for (0..group) |j| { const k = g * group + j; - const qv: u32 = if (bits == 4) (wb[(row * K + k) / 2] >> @intCast(4 * (k % 2))) & 15 else wb[row * K + k]; + const qv: u32 = code(wb, (row * K + k) * bits, bits); const xv: f64 = f32of(x[r * K + k]); dot += @as(f64, @floatFromInt(qv)) * xv; dmag += (@as(f64, @floatFromInt(qv)) + off) * @abs(xv); diff --git a/zig/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/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; +} 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 })); +} 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/src/families/glm/affine.zig b/zig/src/families/glm/affine.zig new file mode 100644 index 0000000000..cd80ddbb1c --- /dev/null +++ b/zig/src/families/glm/affine.zig @@ -0,0 +1,211 @@ +//! GLM-5.3-Flash's affine matrices at any MLX width: stored shapes, byte spans and the widths each device path reads. +const std = @import("std"); +const mtl = @import("metal"); +const quant = @import("quant.zig"); + +pub const Width = quant.Width; + +pub const Ref = struct { + buf: mtl.Buffer, + off: usize = 0, + + pub fn at(r: Ref, by: usize) Ref { + return .{ .buf = r.buf, .off = r.off + by }; + } + + pub fn addr(r: Ref) [*]u8 { + return r.buf.contents() + r.off; + } +}; + +/// The most runs a fused matrix splits into: its members. +pub const max_runs = 6; + +/// A run of a split stack: one member (0.6.6 calls each apart), rows [row, row + n) of the fused output, bytes its own. +pub const Run = struct { w: Ref, s: Ref, b: Ref, row: u32, n: u32, bits: u8, group: u16 }; + +/// An MLX affine matrix: w [n, k bits / 32] u32, s and b [n, k / group] bf16 (stacked: [e, n, ...]). +pub const Affine = struct { + w: Ref, + s: Ref, + b: Ref, + n: u32, + k: u32, + bits: u8, + group: u16, + more: u8 = 0, // runs past the first, whose bytes and width are the matrix's own: 0 for one width + runs: [max_runs - 1]Run = undefined, + + pub fn width(q: Affine) Width { + return .{ .bits = q.bits, .group = q.group }; + } + + /// Run `i` (the first: the matrix's own bytes and width). + pub fn run(q: Affine, i: usize) Run { + if (i > 0) return q.runs[i - 1]; + return .{ .w = q.w, .s = q.s, .b = q.b, .row = 0, .n = if (q.more == 0) q.n else q.runs[0].row, .bits = q.bits, .group = q.group }; + } +}; + +/// Bytes of [n, k]'s packed codes at `w`. +pub fn wBytes(n: usize, k: usize, w: Width) usize { + return n * k * w.bits / 8; +} + +/// Bytes of [n, k]'s bf16 scale (or bias) table at `w`. +pub fn sBytes(n: usize, k: usize, w: Width) usize { + return n * (k / w.group) * 2; +} + +/// Bytes of [n, k]'s codes and both tables at `w`. +pub fn bytes(n: usize, k: usize, w: Width) usize { + return wBytes(n, k, w) + 2 * sBytes(n, k, w); +} + +/// [n, k]'s stored shapes at `w`: u32 words [n, k bits / 32], bf16 tables [n, k / group]; null when k does not pack. +pub fn shapes(n: usize, k: usize, w: Width) ?[2][2]usize { + if (k == 0 or k % w.group != 0 or k * w.bits % 32 != 0) return null; + return .{ .{ n, k * w.bits / 32 }, .{ n, k / w.group } }; +} + +/// Columns [lo, lo + k) as (offset, length) in a row's words and in its tables; null unless whole groups and words. +pub fn cut(lo: usize, k: usize, w: Width) ?[2][2]usize { + if (lo % w.group != 0 or shapes(1, k, w) == null) return null; + return .{ .{ wBytes(1, lo, w), wBytes(1, k, w) }, .{ sBytes(1, lo, w), sBytes(1, k, w) } }; +} + +/// Matrices (each run of a split stack) by width, for a plan's summary. +pub const Tally = struct { + count: [9][3]u32 = @splat(@splat(0)), // by bits, then group 32, 64, 128 + + pub fn add(t: *Tally, v: anytype) void { + const T = @TypeOf(v); + if (T == Affine) { + for (0..@as(usize, v.more) + 1) |i| { + const r = v.run(i); + if (r.n > 0 and r.bits > 0) t.count[r.bits][std.math.log2_int(u16, r.group) - 5] += 1; + } + return; + } + switch (@typeInfo(T)) { + .@"struct" => |s| inline for (s.field_names) |name| t.add(@field(v, name)), + .@"union" => switch (v) { + inline else => |x| t.add(x), + }, + .optional => if (v) |x| t.add(x), + .array => for (v) |x| t.add(x), + else => {}, + } + } + + /// "4-bit/64 x12, 8-bit/64 x3" in `buf`. + pub fn summary(t: Tally, buf: []u8) []const u8 { + var n: usize = 0; + for (t.count, 0..) |g, bits| for (g, 0..) |x, gi| if (x > 0) { + const sep = if (n == 0) "" else ", "; + n += (std.fmt.bufPrint(buf[n..], "{s}{d}-bit/{d} x{d}", .{ sep, bits, @as(u32, 32) << @intCast(gi), x }) catch return buf[0..n]).len; + }; + return buf[0..n]; + } +}; + +/// The device paths that read affine matrices, by projection class. +pub const Path = enum { embed, kda_in, kda_gate, kda_out, mla_in, mla_q, kv_b, mla_out, dense_mlp, shared_expert, routed_experts, head, eh_proj }; + +/// The code widths the row kernels (glm_rows.metal, after core/row_projection.metal) unpack, in groups of 32, 64, 128. +pub const row_bits = [_]u8{ 2, 3, 4, 5, 6, 8 }; + +/// Whether a path's kernels read `w`: kv_b 4/64, the embedding 4/6/8 in groups of 64, the rest the row kernels' widths. +pub fn reads(path: Path, w: Width) bool { + return switch (path) { + .kv_b => std.meta.eql(w, quant.uniform), + .embed => w.group == 64 and (w.bits == 4 or w.bits == 6 or w.bits == 8), + else => std.mem.indexOfScalar(u8, &row_bits, w.bits) != null and (w.group == 32 or w.group == 64 or w.group == 128), + }; +} + +/// Refuses, by tensor name, a width the path's kernels do not read: a load fails before any weight is read. +pub fn admit(name: []const u8, path: Path, w: Width, rows: usize) error{ UnsupportedWidth, UnsupportedShape }!void { + if (!reads(path, w)) { + quant.say("glm: {s} is {d}-bit in groups of {d}, a width the {t} kernels do not read", .{ name, w.bits, w.group, path }); + return error.UnsupportedWidth; + } + if (rows % 4 == 0) return; + quant.say("glm: {s} has {d} rows; the row kernels read a matrix's rows in fours", .{ name, rows }); + return error.UnsupportedShape; +} + +/// Whether every matrix in `v` takes the generated kernels (one width, 4/64): no other width's pipelines needed. +pub fn allNative(v: anytype) bool { + const T = @TypeOf(v); + if (T == Affine) return v.n == 0 or (v.more == 0 and v.bits == 4 and v.group == 64); + return switch (@typeInfo(T)) { + .@"struct" => |s| inline for (s.field_names) |name| { + if (!allNative(@field(v, name))) break false; + } else true, + .@"union" => switch (v) { + inline else => |x| allNative(x), + }, + .optional => if (v) |x| allNative(x) else true, + .array => for (v) |x| { + if (!allNative(x)) break false; + } else true, + else => true, + }; +} + +const testing = std.testing; + +test "stored shapes and bytes follow the width; 4/64 keeps the 4-bit matrix's" { + try testing.expectEqual([2][2]usize{ .{ 4096, 512 }, .{ 4096, 64 } }, shapes(4096, 4096, quant.uniform).?); + try testing.expectEqual(@as(usize, 4096 * 4096 / 2), wBytes(4096, 4096, quant.uniform)); + try testing.expectEqual(@as(usize, 4096 * (4096 / 64) * 2), sBytes(4096, 4096, quant.uniform)); + const cases = [_]struct { Width, [2]usize }{ + .{ .{ .bits = 2, .group = 64 }, .{ 16, 4 } }, + .{ .{ .bits = 3, .group = 32 }, .{ 24, 8 } }, + .{ .{ .bits = 5, .group = 64 }, .{ 40, 4 } }, + .{ .{ .bits = 6, .group = 128 }, .{ 48, 2 } }, + .{ .{ .bits = 8, .group = 64 }, .{ 64, 4 } }, + }; + for (cases) |c| { + const s = shapes(7, 256, c[0]).?; + try testing.expectEqual(c[1][0], s[0][1]); + try testing.expectEqual(c[1][1], s[1][1]); + try testing.expectEqual(7 * 4 * s[0][1], wBytes(7, 256, c[0])); + } + try testing.expect(shapes(1, 96, .{ .bits = 4, .group = 64 }) == null); +} + +test "a column cut keeps whole groups and words at every width" { + try testing.expectEqual([2][2]usize{ .{ 2048, 2048 }, .{ 128, 128 } }, cut(4096, 4096, quant.uniform).?); + try testing.expectEqual([2][2]usize{ .{ 640, 640 }, .{ 32, 32 } }, cut(1024, 1024, .{ .bits = 5, .group = 64 }).?); + try testing.expect(cut(32, 1024, quant.uniform) == null); + try testing.expect(cut(64, 1024, .{ .bits = 6, .group = 128 }) == null); +} + +test "every device path admits 4/64; the oQ widths where their kernels read them, others refused by name" { + for (std.enums.values(Path)) |p| { + try admit("x", p, quant.uniform, 64); + try testing.expectError(error.UnsupportedWidth, admit("x", p, .{ .bits = 7, .group = 64 }, 64)); + try testing.expectError(error.UnsupportedWidth, admit("x", p, .{ .bits = 4, .group = 16 }, 64)); + try testing.expectError(error.UnsupportedShape, admit("x", p, quant.uniform, 66)); + } + for ([_]Width{ .{ .bits = 2, .group = 64 }, .{ .bits = 5, .group = 64 }, .{ .bits = 6, .group = 64 }, .{ .bits = 8, .group = 64 }, .{ .bits = 3, .group = 32 }, .{ .bits = 5, .group = 128 } }) |w| + for ([_]Path{ .kda_in, .kda_gate, .kda_out, .mla_in, .mla_q, .mla_out, .dense_mlp, .shared_expert, .routed_experts, .head, .eh_proj }) |p| try admit("x", p, w, 64); + try admit("x", .embed, .{ .bits = 8, .group = 64 }, 64); + try testing.expectError(error.UnsupportedWidth, admit("x", .embed, .{ .bits = 2, .group = 64 }, 64)); + try testing.expectError(error.UnsupportedWidth, admit("x", .kv_b, .{ .bits = 8, .group = 64 }, 64)); +} + +test "a matrix of another width or a split stack needs the other widths' pipelines; 4/64 and none do not" { + const r: Ref = .{ .buf = undefined }; + const q: Affine = .{ .w = r, .s = r, .b = r, .n = 64, .k = 128, .bits = 4, .group = 64 }; + const none: Affine = .{ .w = r, .s = r, .b = r, .n = 0, .k = 128, .bits = 0, .group = 0 }; + try testing.expect(allNative(.{ q, none, @as(?Affine, null), [2]Affine{ q, q } })); + var eight = q; + eight.bits = 8; + try testing.expect(!allNative(.{ q, [2]Affine{ q, eight } })); + var split = q; + split.more = 1; + try testing.expect(!allNative(@as(?Affine, split))); +} diff --git a/zig/src/families/glm/checks.zig b/zig/src/families/glm/checks.zig index 4b521379c1..98404ca22a 100644 --- a/zig/src/families/glm/checks.zig +++ b/zig/src/families/glm/checks.zig @@ -158,6 +158,7 @@ pub fn checkMatmul(e: *Engine, prompt: []const u32) !void { .kda => |*a| a, .mla => return error.NotKda, }; + if (!@import("wide.zig").native(a.in_proj)) return error.UnsupportedWidth; // the 4/64 product's comparison only const pool = mtl.objc.Pool.push(); defer pool.pop(); e.sync(); diff --git a/zig/src/families/glm/engine.zig b/zig/src/families/glm/engine.zig index 54095bbad2..e1ffe5ac41 100644 --- a/zig/src/families/glm/engine.zig +++ b/zig/src/families/glm/engine.zig @@ -205,6 +205,7 @@ pub const Engine = struct { e.w.deinit(); gpa.destroy(e.w); } + if (!wts.native(e.w, &e.c)) try e.k.widen(gpa, e.device); // a width, split stack or layout past the 4/64 kernels try e.prepare(); if (link) |s| { const rows = e.c.byRows(); diff --git a/zig/src/families/glm/forward.zig b/zig/src/families/glm/forward.zig index 018b99d3d3..330d9324e6 100644 --- a/zig/src/families/glm/forward.zig +++ b/zig/src/families/glm/forward.zig @@ -9,6 +9,7 @@ const Ep = @import("ep.zig").Ep; const moe_route = @import("../../core/moe_route.zig"); const hc_shape = @import("kernels.zig").hc_shape; const hc_core = @import("../../core/hc.zig"); +const wide = @import("wide.zig"); const Ref = wts.Ref; /// The MoE block (moe.zig). pub const moe = @import("moe.zig").moe; @@ -112,11 +113,11 @@ pub fn size(w: usize, h: usize, d: usize) mtl.Size { return mtl.Size.of(w, h, d); } -/// x [rows, K] through a 4-bit matrix: qmv_rows, MLX's one-row qmv_fast sums for every row. -pub fn qmv(x: *const Ctx, e: mtl.ComputeEncoder, pipe: mtl.Pipeline, in: Ref, q: wts.Q4, out: Ref, rows: u32) void { +/// x [rows, K] through a 4-bit/64 matrix (qmv_rows, MLX's one-row qmv_fast sums), else each member at its width (wide). +pub fn qmv(x: *const Ctx, e: mtl.ComputeEncoder, pipe: mtl.Pipeline, in: Ref, q: wts.Affine, out: Ref, rows: u32) void { + if (!wide.native(q)) return wide.dense(&x.k.wide.?, e, q, .{ .in = in, .ldx = q.k, .out = out, .ldo = q.n, .rows = rows }); e.setPipeline(pipe); bind(e, 0, .{ in, q.w, q.s, q.b, out }); - _ = x; e.dispatchThreads(size(32 * rows, q.n / 4, 1), size(32 * rows, 1, 1)); } @@ -149,11 +150,16 @@ pub fn embed(x: *Ctx, e: mtl.ComputeEncoder, ids: Ref, rows: u32) void { pub fn embedRows(x: *const Ctx, e: mtl.ComputeEncoder, ids: Ref, out: Ref, rows: u32) void { const q = x.w.embed; - e.setPipeline(x.k.embed); + const per: u32 = switch (q.bits) { // values a thread: two codes a byte at 4 bits, four in three bytes at 6, one at 8 + 4 => 2, + 6 => 4, + else => 1, + }; + e.setPipeline(if (q.bits == 4) x.k.embed else x.k.wide.?.embed[if (q.bits == 6) 0 else 1]); bind(e, 0, .{ ids, q.w, q.s, q.b }); shape(e, 4, .{x.c.hidden}); bind(e, 5, .{out}); - e.dispatchThreads(size(x.c.hidden / 2, rows, 1), size(256, 1, 1)); + e.dispatchThreads(size(x.c.hidden / per, rows, 1), size(256, 1, 1)); } /// A block boundary: the pending branch into the streams, then (with `hc`) the next block's mix, split and norm (core/hc.zig). @@ -202,10 +208,25 @@ pub fn kdaStep(x: *const Ctx, e: mtl.ComputeEncoder, ki: usize, w: *const wts.Kd const c = x.c; const L = &x.s.kda[ki]; const cur = L.cur; - e.setPipeline(if (c.tp > 1) x.k.kda_rows_tp else x.k.kda_rows); + const gate = wide.kdaGate(if (x.k.wide) |*p| p else null, w.f_b, w.g_b); + var fb = [3]Ref{ w.f_b.w, w.f_b.s, w.f_b.b }; + var gb = [3]Ref{ w.g_b.w, w.g_b.s, w.g_b.b }; + if (gate == .pre) { // f_b and g_b at their widths on f_a and g_a first, [rows, width] each (in gu: no MLP runs yet) + const W = c.kdaWidth(); + const f = x.sc.gu; + const g = f.at(@as(usize, st.max_rows) * W * 2); + wide.dense(&x.k.wide.?, e, w.f_b, .{ .in = proj.at(3 * W * 2), .ldx = c.kdaProj(), .out = f, .ldo = W, .rows = rows }); + wide.dense(&x.k.wide.?, e, w.g_b, .{ .in = proj.at((3 * W + c.kda_dim) * 2), .ldx = c.kdaProj(), .out = g, .ldo = W, .rows = rows }); + fb = .{ f, f, f }; + gb = .{ g, g, g }; + } + e.setPipeline(switch (gate) { + .generated => if (c.tp > 1) x.k.kda_rows_tp else x.k.kda_rows, + .pre, .fused => |p| p, + }); bind(e, 0, .{proj}); shape(e, 1, .{ rows, c.kdaProj() }); - bind(e, 2, .{ L.cs[cur], w.conv_w, w.f_b.w, w.f_b.s, w.f_b.b, w.g_b.w, w.g_b.s, w.g_b.b, w.a, w.dt_bias, L.st[cur], w.o_norm }); + bind(e, 2, .{ L.cs[cur], w.conv_w, fb[0], fb[1], fb[2], gb[0], gb[1], gb[2], w.a, w.dt_bias, L.st[cur], w.o_norm }); e.setValue(c.lower_bound, 14); e.setValue(c.eps, 15); bind(e, 16, .{ y, L.st[1 - cur], L.cs[1 - cur] }); @@ -297,9 +318,10 @@ pub fn mlaCache(x: *const Ctx, e: mtl.ComputeEncoder, mi: usize, w: *const wts.M } } -/// Every row's 64 heads of q_nope (in `qp`, a row qrProj apart) into the latent: MLX's qvm on kv_b's key half. +/// Every row's 64 heads of q_nope (`qp`, rows qrProj apart) into the latent: qvm on kv_b's key half, or embed_q's qmv. pub fn absorb(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Mla, qp: Ref, ql: Ref, rows: u32) void { const c = x.c; + if (w.absorbed) |pair| return wide.perHead(&x.k.wide.?, e, pair[0], c.mla_heads, .{ .in = qp, .ldx = c.qrProj(), .out = ql, .ldo = c.mla_heads * c.kv_lora, .rows = rows }, c.nope, c.kv_lora); e.setPipeline(x.k.absorb); bind(e, 0, .{ w.kv_b.w, w.kv_b.s, w.kv_b.b, qp, ql }); e.setValue(c.qrProj(), 5); @@ -307,9 +329,10 @@ pub fn absorb(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Mla, qp: Ref, e.dispatchGroups(size(1, c.kv_lora / 64, rows * c.mla_heads), size(64, 1, 1)); } -/// Every row's 64 heads' latent outputs to values [rows, heads * v]: MLX's batched qmv_fast on kv_b's value half. +/// Every row's 64 heads' latent outputs to values [rows, heads * v]: qmv_fast on kv_b's value half, or unembed_out's. pub fn unabsorb(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Mla, att: Ref, vals: Ref, rows: u32) void { const c = x.c; + if (w.absorbed) |pair| return wide.perHead(&x.k.wide.?, e, pair[1], c.mla_heads, .{ .in = att, .ldx = c.mla_heads * c.kv_lora, .out = vals, .ldo = c.mla_heads * c.v_dim, .rows = rows }, c.kv_lora, c.v_dim); e.setPipeline(x.k.unabsorb); bind(e, 0, .{ w.kv_b.w, w.kv_b.s, w.kv_b.b, att, vals }); e.setValue(c.mla_heads, 5); diff --git a/zig/src/families/glm/glm.zig b/zig/src/families/glm/glm.zig index c9eaf4f724..11a7502c8f 100644 --- a/zig/src/families/glm/glm.zig +++ b/zig/src/families/glm/glm.zig @@ -13,6 +13,8 @@ pub const pool = @import("pool.zig"); pub const backend = @import("backend.zig"); pub const mirror = @import("mirror.zig"); pub const timing = @import("timing.zig"); +pub const affine = @import("affine.zig"); +pub const wide = @import("wide.zig"); test { _ = config; @@ -31,4 +33,9 @@ test { _ = @import("kept.zig"); _ = mirror; _ = timing; + _ = @import("quant.zig"); + _ = @import("affine.zig"); + _ = @import("names.zig"); + _ = @import("weights_test.zig"); + _ = @import("wide.zig"); } diff --git a/zig/src/families/glm/kernels.zig b/zig/src/families/glm/kernels.zig index 5f0f3e7916..d3ee3a580b 100644 --- a/zig/src/families/glm/kernels.zig +++ b/zig/src/families/glm/kernels.zig @@ -6,6 +6,7 @@ const frags = @import("../../core/frags.zig"); const moe_route = @import("../../core/moe_route.zig"); const affine_mm = @import("../../core/affine_mm.zig"); const hc = @import("../../core/hc.zig"); +const wide = @import("wide.zig"); /// The route's shape (config.zig refuses other checkpoints). pub const route_shape: moe_route.Shape = .{ .hidden = 4096, .experts = 288, .topk = 8 }; @@ -100,13 +101,23 @@ pub const Kernels = struct { absorb_nax_tp: mtl.Pipeline, mm_bf16: affine_mm.Pipes, // prompt chunks' 4-bit g64 matmuls on the tensor units (core/affine_mm.zig): dense, gathers mm_f32: affine_mm.Pipes, // and with fp32 out: expert parallel by rows' down partials + wide: ?wide.Pipes = null, // the products at other widths (wide.zig), for a checkpoint that has one pub fn deinit(k: *Kernels) void { const info = @typeInfo(Kernels).@"struct"; inline for (info.field_names, info.field_types) |name, T| { - if (T == mtl.Pipeline) @field(k, name).deinit() else if (T != u64) for (&@field(k, name)) |*p| p.deinit(); + if (T == mtl.Pipeline) @field(k, name).deinit() else if (T == ?wide.Pipes) { + if (k.wide) |*w| w.deinit(); + } else if (T != u64) for (&@field(k, name)) |*p| p.deinit(); } } + + /// The other widths' pipelines, once: their sources join the hash a learned state's identity reads. + pub fn widen(k: *Kernels, gpa: std.mem.Allocator, device: mtl.Device) !void { + if (k.wide != null) return; + k.wide = try wide.load(gpa, device); + k.source_hash ^= k.wide.?.hash; + } }; /// The generated kernel `key` (tools/zig/gen_glm_kernels.py) and the field its first function fills. @@ -182,11 +193,21 @@ const Job = struct { }; /// Generated kernel `key`'s template at other parameters, under `name`. -fn instance(comptime key: []const u8, comptime params: []const u8, comptime name: []const u8) []const u8 { +pub fn instance(comptime key: []const u8, comptime params: []const u8, comptime name: []const u8) []const u8 { const f = comptime kernelOf(key).functions[0]; return "\ntemplate [[host_name(\"" ++ name ++ "\")]] [[kernel]] decltype(" ++ f ++ "<" ++ params ++ ">) " ++ f ++ "<" ++ params ++ ">;\n"; } +/// Generated kernel `key`'s source. +pub fn generatedSource(comptime key: []const u8) []const u8 { + return kernelOf(key).source; +} + +/// Generated kernel `key`'s first function. +pub fn generatedFunction(comptime key: []const u8) [:0]const u8 { + return kernelOf(key).functions[0]; +} + fn kernelOf(comptime key: []const u8) sources.glm.Kernel { @setEvalBranchQuota(20000); for (sources.glm.all) |k| if (comptime std.mem.eql(u8, k.key, key)) return k; @@ -197,6 +218,7 @@ fn kernelOf(comptime key: []const u8) sources.glm.Kernel { pub fn load(gpa: std.mem.Allocator, device: mtl.Device) !*Kernels { const k = try gpa.create(Kernels); errdefer gpa.destroy(k); + k.wide = null; var jobs: [generated.len + 22]Job = undefined; inline for (generated, 0..) |g, i| { const src = comptime kernelOf(g.key); diff --git a/zig/src/families/glm/load_plan.zig b/zig/src/families/glm/load_plan.zig index 13f032535c..53a0dddbcd 100644 --- a/zig/src/families/glm/load_plan.zig +++ b/zig/src/families/glm/load_plan.zig @@ -3,6 +3,7 @@ const std = @import("std"); const mtl = @import("metal"); const cfg = @import("config.zig"); const wts = @import("weights.zig"); +const affine = @import("affine.zig"); const st = @import("state.zig"); const prompt_mod = @import("prompt.zig"); const ep_mod = @import("ep.zig"); @@ -27,6 +28,70 @@ pub fn planBytes(gpa: std.mem.Allocator, device: mtl.Device, dir: []const u8, c: return plan.bytes; } +/// Bytes of layer i's dense buffer (everything but routed experts) at its matrices' widths, generously rounded. +pub fn layerBytes(l: anytype, i: usize) !usize { + const c = l.c; + const D: usize = c.hidden; + const a = wts.prefix ++ "layers.{d}.self_attn."; + const m = wts.prefix ++ "layers.{d}.mlp."; + var n: usize = 0; + const q = struct { + fn b(rows: usize, k: usize, w: wts.Width) usize { + return affine.bytes(rows, k, w) + 3 * 256; + } + }; + n += 2 * (24 * 4 * D * 2 + 1024) + 2 * (D * 2 + 256); + switch (c.kind(@intCast(i))) { + .kda => { + const W: usize = c.kdaWidth(); + n += try stackBytes(l, a, .{i}, &wts.kda_in, c.kdaProj(), D, 64) + q.b(W, c.kda_dim, try l.widthAt(a ++ "f_b_proj", .{i})); + n += q.b(W, c.kda_dim, try l.widthAt(a ++ "g_b_proj", .{i})) + q.b(D, W, try l.widthAt(a ++ "o_proj", .{i})) + 3 * W * c.conv * 4 + 8 * 1024 + W * 4; + }, + .mla => { + n += try stackBytes(l, a, .{i}, &wts.mla_x, c.xProj(), D, 64) + try stackBytes(l, a, .{i}, &wts.mla_qr, c.qrProj(), c.q_lora, 1); + n += q.b(D, c.mla_heads * c.v_dim, try l.widthAt(a ++ "o_proj", .{i})) + 16 * 1024 + c.i_dim * D * 2; + if (l.has(a ++ "embed_q.weight", .{i})) { + n += q.b(c.mla_heads * c.kv_lora, c.nope, try l.widthAt(a ++ "embed_q", .{i})) + q.b(c.mla_heads * c.v_dim, c.kv_lora, try l.widthAt(a ++ "unembed_out", .{i})); + } else n += q.b(c.mla_heads * (c.nope + c.v_dim), c.kv_lora, try l.widthAt(a ++ "kv_b_proj", .{i})); + }, + } + if (c.isMoe(@intCast(i))) { + n += c.experts * D * 2 + c.experts * 4 + try stackBytes(l, m ++ "shared_experts.", .{i}, &wts.gate_up, 2 * c.moe_inter, D, 1) + q.b(D, c.moe_inter, try l.widthAt(m ++ "shared_experts.down_proj", .{i})) + 2048; + } else n += try stackBytes(l, m, .{i}, &wts.gate_up, 2 * c.dense_inter, D, 1) + q.b(D, c.dense_inter, try l.widthAt(m ++ "down_proj", .{i})); + const eh = wts.prefix ++ "layers.{d}.eh_proj"; + if (i == c.layers) n += 3 * (D * 2 + 256) + if (l.has(eh ++ ".scales", .{i})) q.b(D, 2 * D, try l.widthAt(eh, .{i})) else D * 2 * D * 2 + 256; + return n; +} + +/// Bytes of layer i's routed experts on this Mac at the first one's widths (the rest must match). +pub fn expertBytes(l: anytype, i: usize) !usize { + const c = l.c; + const D: usize = c.hidden; + const every = i == c.layers; + const m: usize = if (every) c.moe_inter else c.inter[1] - c.inter[0]; + const lo: usize = if (every) 0 else c.own[0]; + const per = affine.bytes(m, D, try l.widths.expert(i, lo, "gate_proj")) + affine.bytes(m, D, try l.widths.expert(i, lo, "up_proj")) + affine.bytes(D, m, try l.widths.expert(i, lo, "down_proj")); + return (if (every) c.experts else c.own[1] - c.own[0]) * per + 9 * 256; +} + +/// A stack's bytes: its rows, padded to `pad` a run (a part each once widths differ), at the costliest bytes a row. +fn stackBytes(l: anytype, comptime fmt: []const u8, args: anytype, parts: []const []const u8, rows: usize, k: usize, pad: usize) !usize { + var name: [200]u8 = undefined; + const base = try std.fmt.bufPrint(&name, fmt, args); + var most: usize = 0; + var first: ?wts.Width = null; + var mixed = false; + for (parts) |p| { + var full: [240]u8 = undefined; + const w = try l.widths.of(try std.fmt.bufPrint(&full, "{s}{s}", .{ base, p })); + most = @max(most, affine.bytes(1, k, w)); + mixed = mixed or (first != null and !std.meta.eql(first.?, w)); + if (first == null) first = w; + } + const runs = if (mixed) parts.len else 1; + return (std.mem.alignForward(usize, rows, pad) + (runs - 1) * pad) * most + runs * 3 * 256; +} + /// The fewest bytes of caches and buffers a load of `cap` tokens takes (the shortest prompt chunks), counted, not allocated. pub fn leastArena(gpa: std.mem.Allocator, c: *const cfg.Config, cap: u32, chunked: bool) !usize { var dry: st.Arena = .{ .device = undefined, .gpa = gpa, .dry = true }; diff --git a/zig/src/families/glm/moe.zig b/zig/src/families/glm/moe.zig index ab3faf68fc..9604439859 100644 --- a/zig/src/families/glm/moe.zig +++ b/zig/src/families/glm/moe.zig @@ -6,6 +6,7 @@ const st = @import("state.zig"); const fwd = @import("forward.zig"); const moe_route = @import("../../core/moe_route.zig"); const route_shape = @import("kernels.zig").route_shape; +const wide = @import("wide.zig"); const Ref = wts.Ref; const Ctx = fwd.Ctx; const Class = fwd.Class; @@ -39,7 +40,7 @@ pub fn moe(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, r if (s & Class.exchange == 0) ep.send(e, sc.ye, rows, on(x, "x_pack") and on(x, "x_locpost"), on(x, "x_locpost")); if (s & Class.shared == 0) experts(x, e, w, x_in, rows, 1, .{ sc.none, sc.none, sc.none }); if (s & Class.exchange == 0 and on(x, "x_unpack")) ep.receive(e, sc.ye, rows); - } else if (s & (Class.shared | Class.routed) == 0 and on(x, "e_gateup") and on(x, "e_down") and on(x, "s_gateup") and on(x, "s_down")) { + } else if (native(w, .both) and s & (Class.shared | Class.routed) == 0 and on(x, "e_gateup") and on(x, "e_down") and on(x, "s_gateup") and on(x, "s_down")) { if (s & Class.route == 0) route(x, e, w, x_in, rows); both(x, e, w, x_in, rows); // the shared expert as the gathers' last slot: two launches fewer if (s & Class.combine != 0 or !on(x, "combine")) return; @@ -49,6 +50,9 @@ pub fn moe(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, r bind(e, 4, .{sc.branch}); e.dispatchThreads(size(rows * c.hidden, 1, 1), size(256, 1, 1)); return; + } else if (x.k.wide != null and wide.moeOne(parts(w)) and s & (Class.shared | Class.routed) == 0 and on(x, "e_gateup") and on(x, "e_down") and on(x, "s_gateup") and on(x, "s_down")) { + if (s & Class.route == 0) route(x, e, w, x_in, rows); + wide.moe(&x.k.wide.?, e, parts(w), c.experts, c.topk, c.swiglu_limit, .{ .x = x_in, .ids = sc.pick, .act = sc.act, .acts = sc.acts, .ye = sc.ye, .ys = sc.ys }, rows); } else { if (s & Class.shared == 0) experts(x, e, w, x_in, rows, 1, .{ sc.none, sc.none, sc.none }); if (s & Class.route == 0) route(x, e, w, x_in, rows); @@ -62,6 +66,10 @@ pub fn moe(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, r e.dispatchThreads(size(rows * c.hidden, 1, 1), size(256, 1, 1)); } +fn parts(w: *const wts.Moe) wide.Moe { + return .{ .gate = w.gate, .up = w.up, .down = w.down, .sh_gate_up = w.sh_gate_up, .sh_down = w.sh_down }; +} + /// The router's fp32 logits and the route: picks, weights and the window's unique experts with their members. fn route(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, rows: u32) void { const c = x.c; @@ -96,6 +104,17 @@ fn route(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, row e.dispatchThreads(size(512, 1, 1), size(512, 1, 1)); } +/// Whether the shared expert, the routed ones or both take the generated 4/64 kernels. +pub fn native(w: *const wts.Moe, which: enum { shared, routed, both }) bool { + const shared = wide.native(w.sh_gate_up) and wide.native(w.sh_down); + const routed = wide.native(w.gate) and wide.native(w.up) and wide.native(w.down); + return switch (which) { + .shared => shared, + .routed => routed, + .both => shared and routed, + }; +} + /// Part 1: the shared expert into `ys`; part 2: the routed experts of `group` (unique ids, members, count) into `ye`. fn experts(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, rows: u32, part: u32, group: [3]Ref) void { const c = x.c; @@ -103,6 +122,8 @@ fn experts(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, r const sc = x.sc; const D = c.hidden; const N = c.moe_inter; + if (part == 1 and !native(w, .shared)) return sharedRows(x, e, w, x_in, rows); + if (part == 2 and !native(w, .routed)) return routedRows(x, e, w, x_in, rows); const zs: u32 = if (part == 1) 1 else rows * c.topk; const slots: u32 = if (part == 1) 1 else c.topk; const act = if (part == 1) sc.acts else sc.act; @@ -176,3 +197,31 @@ fn gateUp(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, ro bind(e, 15, .{act}); e.dispatchThreads(size(32 * rows, if (part == 1) N else N / 4, zs), size(32 * rows, 1, 1)); } + +/// The shared expert at its widths apart from the routed launch (0.6.6 fuses only a 4-bit one): gate/up, SwiGLU, down. +fn sharedRows(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, rows: u32) void { + const c = x.c; + const sc = x.sc; + const N = c.moe_inter; + const p = &x.k.wide.?; + if (on(x, "s_gateup") and w.sh_gate_up.more == 0) { // gate and up one width: one launch, glm_swiglu's bits + wide.glu(p, e, w.sh_gate_up, N, x_in, sc.acts, rows, c.swiglu_limit); + } else if (on(x, "s_gateup")) { + wide.dense(p, e, w.sh_gate_up, .{ .in = x_in, .ldx = c.hidden, .out = sc.gu, .ldo = 2 * N, .rows = rows }); + e.setPipeline(x.k.swiglu); + bind(e, 0, .{ sc.gu, sc.acts }); + e.setValue(fwd.Rows{ .rows = @intCast(rows), .width = @intCast(N), .x_stride = @intCast(2 * N), .y_stride = @intCast(N), .eps = c.swiglu_limit }, 2); + e.dispatchThreads(size(N, rows, 1), size(256, 1, 1)); + } + if (on(x, "s_down")) wide.dense(p, e, w.sh_down, .{ .in = sc.acts, .ldx = N, .out = sc.ys, .ldo = c.hidden, .rows = rows }); +} + +/// The routed experts at their widths, each (row, pick) its own slot: gate and up SwiGLU'd into `act`, down into `ye`. +fn routedRows(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, rows: u32) void { + const c = x.c; + const sc = x.sc; + const p = &x.k.wide.?; + const slots = rows * c.topk; + if (on(x, "e_gateup")) wide.byIds(p, e, .{ .w = w.gate, .up = w.up, .experts = c.experts, .repeat = c.topk, .lim = c.swiglu_limit }, x_in, sc.pick, sc.act, slots); + if (on(x, "e_down")) wide.byIds(p, e, .{ .w = w.down, .experts = c.experts, .repeat = 1 }, sc.act, sc.pick, sc.ye, slots); +} diff --git a/zig/src/families/glm/mtp.zig b/zig/src/families/glm/mtp.zig index d683c2389f..2f1234c1fb 100644 --- a/zig/src/families/glm/mtp.zig +++ b/zig/src/families/glm/mtp.zig @@ -37,7 +37,9 @@ pub fn absorb(x: *fwd.Ctx, e: mtl.ComputeEncoder, h: Ref, next: Ref, rows: u32, fwd.embedRows(x, e, next, sc.m_emb, rows); fwd.rms(x, e, sc.m_emb, m.enorm, sc.m_eh, rows, D, D, 2 * D, c.eps); fwd.rms(x, e, h, m.hnorm, sc.m_eh.at(@as(usize, D) * 2), rows, D, D, 2 * D, c.eps); - fwd.qmv(x, e, x.k.qmv_kda_out, sc.m_eh, m.eh_proj, sc.m_x, rows); + if (m.eh_bf16) |eh| { // stored unquantized (mlx-lm's): fp32 sums of its bf16 rows + @import("wide.zig").bf16(&x.k.wide.?, e, eh, D, 2 * D, .{ .in = sc.m_eh, .ldx = 2 * D, .out = sc.m_x, .ldo = D, .rows = rows }); + } else fwd.qmv(x, e, x.k.qmv_kda_out, sc.m_eh, m.eh_proj, sc.m_x, rows); fwd.rms(x, e, sc.m_x, L.in_norm, sc.m_xn, rows, D, D, D, c.eps); fwd.mlaKeys(x, e, c.countKind(.mla), &L.attn.mla, sc.m_xn, rows, pos); } diff --git a/zig/src/families/glm/names.zig b/zig/src/families/glm/names.zig new file mode 100644 index 0000000000..abea94eabb --- /dev/null +++ b/zig/src/families/glm/names.zig @@ -0,0 +1,104 @@ +//! GLM-5.3-Flash's checkpoint index in either layout (the original or mlx-lm's), each tensor by its original name. +const std = @import("std"); +const mtl = @import("metal"); +const st = @import("../../core/safetensors.zig"); +const load_fix = @import("weights_fix.zig"); + +const Src = load_fix.Src; +const lm = "model.language_model."; + +/// A stored tensor's name in the original layout, or null for another tower's (the vision tower is never read). +pub fn original(buf: []u8, name: []const u8, mtp_layer: u32) error{NoSpaceLeft}!?[]const u8 { + if (std.mem.startsWith(u8, name, lm) or std.mem.startsWith(u8, name, "lm_head.")) return name; + const mlx = "language_model."; + if (!std.mem.startsWith(u8, name, mlx)) return null; + var short = name[mlx.len..]; + if (std.mem.startsWith(u8, short, "lm_head.")) return short; + var num: [16]u8 = undefined; + var head: []const u8 = ""; + if (std.mem.startsWith(u8, short, "model.")) { + short = short["model.".len..]; + } else if (std.mem.startsWith(u8, short, "mtp.0.")) { // the MTP block as the layer after the backbone's + short = short["mtp.0.".len..]; + if (std.mem.startsWith(u8, short, "block.")) short = short["block.".len..]; + if (std.mem.eql(u8, short, "norm.weight")) short = "shared_head.norm.weight"; + head = try std.fmt.bufPrint(&num, "layers.{d}.", .{mtp_layer}); + } else return null; + inline for (.{ "attn", "ffn" }) |site| inline for (.{ "fn", "base", "scale" }) |part| { + const tail = "." ++ site ++ "_hc." ++ part; + if (std.mem.endsWith(u8, short, tail)) return try std.fmt.bufPrint(buf, lm ++ "{s}{s}.hc_" ++ site ++ "_" ++ part, .{ head, short[0 .. short.len - tail.len] }); + }; + const gate = ".self_attn.forget_gate."; + if (std.mem.indexOf(u8, short, gate)) |at| return try std.fmt.bufPrint(buf, lm ++ "{s}{s}.self_attn.{s}", .{ head, short[0..at], short[at + gate.len ..] }); + return try std.fmt.bufPrint(buf, lm ++ "{s}{s}", .{ head, short }); +} + +/// The index and every shard's header: each language-model tensor's shard and byte range, by its original name. +pub fn read(gpa: std.mem.Allocator, arena: std.mem.Allocator, dir: []const u8, mtp_layer: u32, shards: *std.ArrayList([:0]const u8), names: *std.StringHashMapUnmanaged(Src)) !void { + const path = try std.fmt.allocPrintSentinel(arena, "{s}/model.safetensors.index.json", .{dir}, 0); + const f = try mtl.MappedFile.open(path); + defer f.deinit(); + const doc = try std.json.parseFromSliceLeaky(std.json.Value, arena, f.bytes[0..f.size], .{ .allocate = .alloc_always }); + var files: std.StringArrayHashMapUnmanaged(void) = .empty; + var it = doc.object.get("weight_map").?.object.iterator(); + while (it.next()) |kv| try files.put(arena, kv.value_ptr.string, {}); + for (files.keys(), 0..) |name, si| { + const full = try std.fmt.allocPrintSentinel(arena, "{s}/{s}", .{ dir, name }, 0); + try shards.append(gpa, full); + const fd = std.c.open(full, .{ .ACCMODE = .RDONLY }); + if (fd < 0) return error.OpenFailed; + defer _ = std.c.close(fd); + var head: [8]u8 = undefined; + try load_fix.readAll(fd, &head, 0); + const hlen = std.mem.readInt(u64, &head, .little); + const header = try arena.alloc(u8, hlen); + try load_fix.readAll(fd, header, 8); + const end = std.c.lseek(fd, 0, std.c.SEEK.END); + const data = 8 + hlen; + const doc2 = try std.json.parseFromSliceLeaky(std.json.Value, arena, header, .{}); + var eit = doc2.object.iterator(); + while (eit.next()) |kv| { + var buf: [240]u8 = undefined; + // the language model's tensors only: the vision tower's (5-D patch kernels among them) are never read + const tensor = (try original(&buf, kv.key_ptr.*, mtp_layer)) orelse continue; + const o = kv.value_ptr.object; + const dtype = st.DType.parse(o.get("dtype").?.string) orelse return error.UnsupportedDType; + const dims = o.get("shape").?.array.items; + if (dims.len > st.max_rank) return error.RankTooHigh; + var shape: [st.max_rank]usize = @splat(1); + for (dims, 0..) |d, j| shape[j] = @intCast(d.integer); + const offs = o.get("data_offsets").?.array.items; + const b: u64 = @intCast(offs[0].integer); + const e2: u64 = @intCast(offs[1].integer); + if (e2 < b or data + e2 > @as(u64, @intCast(end))) return error.BadSafetensors; + const got = try names.getOrPut(arena, try arena.dupe(u8, tensor)); + if (got.found_existing) { // one tensor stored under both layouts' names, or twice + @import("quant.zig").say("glm: {s} is stored twice in the checkpoint", .{tensor}); + return error.BadSafetensors; + } + got.value_ptr.* = .{ .shard = @intCast(si), .off = data + b, .len = e2 - b, .dtype = dtype, .shape = shape, .rank = @intCast(dims.len) }; + } + } +} + +test "mlx-lm's names map onto the original layout's, the MTP block onto layer 45" { + var buf: [240]u8 = undefined; + const cases = [_][2][]const u8{ + .{ "language_model.model.embed_tokens.weight", "model.language_model.embed_tokens.weight" }, + .{ "language_model.lm_head.scales", "lm_head.scales" }, + .{ "lm_head.weight", "lm_head.weight" }, + .{ "model.language_model.layers.3.self_attn.kv_b_proj.weight", "model.language_model.layers.3.self_attn.kv_b_proj.weight" }, + .{ "language_model.model.layers.7.attn_hc.fn", "model.language_model.layers.7.hc_attn_fn" }, + .{ "language_model.model.layers.7.ffn_hc.scale", "model.language_model.layers.7.hc_ffn_scale" }, + .{ "language_model.model.layers.3.mlp.switch_mlp.up_proj.biases", "model.language_model.layers.3.mlp.switch_mlp.up_proj.biases" }, + .{ "language_model.model.layers.3.self_attn.embed_q.weight", "model.language_model.layers.3.self_attn.embed_q.weight" }, + .{ "language_model.model.layers.0.self_attn.forget_gate.A_log", "model.language_model.layers.0.self_attn.A_log" }, + .{ "language_model.mtp.0.block.self_attn.o_proj.weight", "model.language_model.layers.45.self_attn.o_proj.weight" }, + .{ "language_model.mtp.0.eh_proj.weight", "model.language_model.layers.45.eh_proj.weight" }, + .{ "language_model.mtp.0.enorm.weight", "model.language_model.layers.45.enorm.weight" }, + .{ "language_model.mtp.0.norm.weight", "model.language_model.layers.45.shared_head.norm.weight" }, + }; + for (cases) |c| try std.testing.expectEqualStrings(c[1], (try original(&buf, c[0], 45)).?); + for ([_][]const u8{ "vision_tower.blocks.0.attn.qkv.weight", "language_model.mtp.1.eh_proj.weight", "model.visual.x" }) |skip| + try std.testing.expect((try original(&buf, skip, 45)) == null); +} diff --git a/zig/src/families/glm/prompt.zig b/zig/src/families/glm/prompt.zig index e4f94f0600..4fed98c967 100644 --- a/zig/src/families/glm/prompt.zig +++ b/zig/src/families/glm/prompt.zig @@ -10,6 +10,7 @@ const ep_mod = @import("ep.zig"); const affine_mm = @import("../../core/affine_mm.zig"); const hc = @import("../../core/hc.zig"); const kernels = @import("kernels.zig"); +const wide = @import("wide.zig"); const Kernels = kernels.Kernels; const moe_route = @import("../../core/moe_route.zig"); const Ref = wts.Ref; @@ -186,28 +187,32 @@ fn run(e: mtl.ComputeEncoder, grid: [3]usize, group: [3]usize) void { } /// y [M, N] = x [M, K] W^T for a 4-bit matrix (64x64 tiles; N rounded up to the weights' padded rows, y's row pitch). -fn qmm(p: *const Prompt, e: mtl.ComputeEncoder, x: Ref, q: wts.Q4, y: Ref, M: u32) void { +fn qmm(p: *const Prompt, e: mtl.ComputeEncoder, x: Ref, q: wts.Affine, y: Ref, M: u32) void { + if (!wide.native(q)) return qmmAt(p, e, x, q.k, q, y, M); affine_mm.rowSums(e, p.mm.mm_bf16, 64, x, p.sums, M, q.k); affine_mm.dense(e, p.mm.mm_bf16, x, p.sums, q, y, M); } /// `qmm` over x's rows `ld` elements apart (a slice of a wider row). -fn qmmAt(p: *const Prompt, e: mtl.ComputeEncoder, x: Ref, ld: u32, q: wts.Q4, y: Ref, M: u32) void { - affine_mm.rowSumsAt(e, p.mm.mm_bf16, 64, x, p.sums, M, q.k, ld); - affine_mm.denseBatch(e, p.mm.mm_bf16, x, p.sums, q, y, M, 1, .{ .x_row = @intCast(ld), .y_row = @intCast(std.mem.alignForward(u32, q.n, 64)) }); +fn qmmAt(p: *const Prompt, e: mtl.ComputeEncoder, x: Ref, ld: u32, q: wts.Affine, y: Ref, M: u32) void { + const pitch = std.mem.alignForward(u32, q.n, 64); + if (!wide.native(q)) return wide.denseMm(&p.mm.wide.?, e, q, .{ .in = x, .ldx = ld, .out = y, .ldo = pitch, .rows = M }, p.sums); + const mm = p.mm.mm_bf16; + affine_mm.rowSumsAt(e, mm, 64, x, p.sums, M, q.k, ld); + affine_mm.denseBatch(e, mm, x, p.sums, q, y, M, 1, .{ .x_row = @intCast(ld), .y_row = @intCast(pitch) }); } /// `qmm` with fp32 out (TP2: one Mac's partial of a row-split projection). -fn qmmF32(p: *const Prompt, e: mtl.ComputeEncoder, x: Ref, q: wts.Q4, y: Ref, M: u32) void { +fn qmmF32(p: *const Prompt, e: mtl.ComputeEncoder, x: Ref, q: wts.Affine, y: Ref, M: u32) void { affine_mm.rowSums(e, p.mm.mm_bf16, 64, x, p.sums, M, q.k); affine_mm.dense(e, p.mm.mm_f32, x, p.sums, q, y, M); } /// The routed experts' gate, up and SwiGLU on the M rows of `x_in` its pairs read (in expert order) into `act` ([n, gate.n]). -fn gateUp(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, M: u32, n: u32) void { - affine_mm.rowSums(e, p.mm.mm_bf16, 64, x_in, p.sums, M, w.gate.k); - affine_mm.gatherGlu(e, p.mm.mm_bf16, x_in, p.sums, w.gate, w.up, p.offsets, p.order, p.act, n, x.c.experts, x.c.topk, x.c.swiglu_limit); - affine_mm.rowSums(e, p.mm.mm_bf16, 64, p.act, p.sums, n, w.down.k); +fn gateUp(p: *const Prompt, mm: affine_mm.Pipes, x: *const fwd.Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, M: u32, n: u32) void { + affine_mm.rowSums(e, mm, 64, x_in, p.sums, M, w.gate.k); + affine_mm.gatherGlu(e, mm, x_in, p.sums, w.gate, w.up, p.offsets, p.order, p.act, n, x.c.experts, x.c.topk, x.c.swiglu_limit); + affine_mm.rowSums(e, mm, 64, p.act, p.sums, n, w.down.k); } fn swiglu(x: *const fwd.Ctx, e: mtl.ComputeEncoder, gu: Ref, act: Ref, rows: u32, width: u32) void { @@ -247,7 +252,7 @@ fn moe(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, w: *const wts if (s & Class.route == 0) route(p, x, e, w, x_in, M); if (ep) |t| { if (s & Class.routed == 0) { // down's fp32 partials put back in (row, slot) order for the decode's combine - gateUp(p, x, e, w, x_in, M, n); + gateUp(p, p.mm.mm_bf16, x, e, w, x_in, M, n); affine_mm.gatherTo(e, p.mm.mm_f32, p.act, p.sums, w.down, p.offsets, p.yf, p.order, n, E); } if (s & Class.exchange == 0) t.sendRows(e, p.yf, p.wts, M); @@ -255,9 +260,14 @@ fn moe(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, w: *const wts if (s & Class.combine == 0) t.receiveRows(e, p.ys, out, M); return; } - if (s & Class.routed == 0) { // down's rows put back in (row, slot) order for the combine - gateUp(p, x, e, w, x_in, M, n); - affine_mm.gatherTo(e, p.mm.mm_bf16, p.act, p.sums, w.down, p.offsets, p.yp, p.order, n, E); + const path = wide.routed(w.gate, w.up, w.down); + if (s & Class.routed == 0 and path == .rows) { // each expert's sorted pairs at its width + wide.bySorted(&p.mm.wide.?, e, .{ .w = w.gate, .up = w.up, .experts = E, .repeat = c.topk, .lim = c.swiglu_limit }, x_in, p.offsets, p.order, p.act, c.moe_inter); + wide.bySorted(&p.mm.wide.?, e, .{ .w = w.down, .experts = E, .repeat = 1 }, p.act, p.offsets, p.order, p.yp, D); + } else if (s & Class.routed == 0) { // down's rows put back in (row, slot) order for the combine + const mm = if (path == .mm4) p.mm.mm_bf16 else p.mm.wide.?.mm2; + gateUp(p, mm, x, e, w, x_in, M, n); + affine_mm.gatherTo(e, mm, p.act, p.sums, w.down, p.offsets, p.yp, p.order, n, E); } if (s & Class.combine != 0) return; e.setPipeline(k.moe_combine); @@ -355,6 +365,7 @@ fn unabsorb(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, w: *cons const H = c.mla_heads; const K = c.kv_lora; const per = c.nope + c.v_dim; // kv_b's rows a head: its key half, then its value half + if (w.absorbed) |pair| return wide.headsBatch(&p.mm.wide.?, e, pair[1], H, .{ .in = p.att, .ldx = H * K, .out = p.vals, .ldo = H * c.v_dim, .rows = M }, K, p.sums); var q = w.kv_b; q.w = q.w.at(@as(usize, c.nope) * K / 2); q.s = q.s.at(@as(usize, c.nope) * (K / 64) * 2); @@ -390,7 +401,9 @@ fn mla(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, mi: usize, w: qmm(p, e, p.qr, w.qr_proj, p.qp, M); } if (fwd.on(x, "mla_cache")) fwd.mlaCache(x, e, mi, w, x_in, p.xp, XP, p.iw, M, pos); - if (fwd.on(x, "mla_absorb")) { + if (fwd.on(x, "mla_absorb") and w.absorbed != null) { + wide.headsBatch(&p.mm.wide.?, e, w.absorbed.?[0], c.mla_heads, .{ .in = p.qp, .ldx = c.qrProj(), .out = p.ql, .ldo = c.mla_heads * c.kv_lora, .rows = M }, c.nope, p.sums); // mlx-lm's embed_q + } else if (fwd.on(x, "mla_absorb")) { e.setPipeline(if (c.tp > 1) x.k.absorb_nax_tp else x.k.absorb_nax); // the row kernel's qvm, weights in three bf16 parts bind(e, 0, .{ w.kv_b.w, w.kv_b.s, w.kv_b.b, p.qp, p.ql }); e.setValue([2]i32{ @intCast(M), @intCast(c.qrProj()) }, 5); @@ -499,7 +512,7 @@ pub fn mtp(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, h: Ref, n fwd.embedRows(x, e, next, p.m_emb, M); fwd.rms(x, e, p.m_emb, m.enorm, p.m_eh, M, D, D, 2 * D, c.eps); fwd.rms(x, e, h, m.hnorm, p.m_eh.at(@as(usize, D) * 2), M, D, D, 2 * D, c.eps); - qmm(p, e, p.m_eh, m.eh_proj, p.m_x, M); + if (m.eh_bf16) |eh| wide.bf16(&p.mm.wide.?, e, eh, D, 2 * D, .{ .in = p.m_eh, .ldx = 2 * D, .out = p.m_x, .ldo = D, .rows = M }) else qmm(p, e, p.m_eh, m.eh_proj, p.m_x, M); fwd.rms(x, e, p.m_x, L.in_norm, p.m_xn, M, D, D, D, c.eps); qmm(p, e, p.m_xn, L.attn.mla.x_proj, p.xp, M); fwd.mlaCache(x, e, c.countKind(.mla), &L.attn.mla, p.m_xn, p.xp, @intCast(std.mem.alignForward(usize, c.xProj(), 64)), p.iw, M, pos); diff --git a/zig/src/families/glm/quant.zig b/zig/src/families/glm/quant.zig new file mode 100644 index 0000000000..6b63242db6 --- /dev/null +++ b/zig/src/families/glm/quant.zig @@ -0,0 +1,174 @@ +//! GLM-5.3-Flash's MLX affine width per module: config.json's quantization table, read through core/quant. +const std = @import("std"); +const quant = @import("../../core/quant/quant.zig"); + +pub const Width = quant.mlx.Width; + +/// The width every GLM kernel is generated for: 4 bits in groups of 64. +pub const uniform: Width = .{ .bits = 4, .group = 64 }; + +/// The checkpoint's widths: the core table with its override keys canonical (`layers.N.x`, `embed_tokens`, `lm_head`). +pub const Widths = struct { + table: quant.mlx.Config, + mtp_layer: u32, + + /// The width of the module a tensor name names, in either layout; an entry the table refuses fails by name. + pub fn of(w: *const Widths, name: []const u8) !Width { + var buf: [256]u8 = undefined; + return w.table.width(try canonical(&buf, name, w.mtp_layer)) catch |err| { + say("glm: config.json quantizes {s} in a format the engine does not read (MLX affine, 2 to 8 bits)", .{name}); + return err; + }; + } + + /// Layer `i`'s routed expert `e`'s `proj`: its own entry, else its layer's stack's (`switch_mlp`), else the global. + pub fn expert(w: *const Widths, i: usize, e: usize, proj: []const u8) !Width { + var a: [128]u8 = undefined; + var b: [128]u8 = undefined; + const own = try std.fmt.bufPrint(&a, "layers.{d}.mlp.experts.{d}.{s}", .{ i, e, proj }); + const stack = try std.fmt.bufPrint(&b, "layers.{d}.mlp.switch_mlp.{s}", .{ i, proj }); + const x = w.table.overrides.get(own); + const y = w.table.overrides.get(stack); + if (x != null and y != null and !std.meta.eql(x.?, y.?)) return conflict(own, stack); + return w.of(if (x != null) own else stack); + } +}; + +/// A refusal's reason, logged as an error; tests check the error itself and keep the log quiet. +pub fn say(comptime fmt: []const u8, args: anytype) void { + if (@import("builtin").is_test) return; + std.log.err(fmt, args); +} + +fn conflict(a: []const u8, b: []const u8) error{ConflictingQuantization} { + say("glm: config.json gives {s} and {s}, one module, different widths", .{ a, b }); + return error.ConflictingQuantization; +} + +/// core/quant's module key, then GLM's own: `mtp.0.[block.]` as its layer, `self_attn.forget_gate.` as `self_attn.`. +pub fn canonical(buf: []u8, name: []const u8, mtp_layer: u32) error{NoSpaceLeft}![]const u8 { + var short = @import("../../core/quant/keys.zig").canonical(name); + var layer: []const u8 = ""; + var num: [12]u8 = undefined; + if (std.mem.startsWith(u8, short, "mtp.0.")) { + short = short["mtp.0.".len..]; + if (std.mem.startsWith(u8, short, "block.")) short = short["block.".len..]; + layer = try std.fmt.bufPrint(&num, "layers.{d}.", .{mtp_layer}); + } + const gate = ".self_attn.forget_gate."; + const at = std.mem.indexOf(u8, short, gate) orelse return std.fmt.bufPrint(buf, "{s}{s}", .{ layer, short }); + return std.fmt.bufPrint(buf, "{s}{s}.self_attn.{s}", .{ layer, short[0..at], short[at + gate.len ..] }); +} + +/// The quantization table of config.json's root, or of its text_config: `quantization`, else `quantization_config`. +fn block(root: std.json.ObjectMap) ?std.json.ObjectMap { + const text: ?std.json.ObjectMap = if (root.get("text_config")) |t| (if (t == .object) t.object else null) else null; + for ([_]?std.json.ObjectMap{ root, text }) |o| if (o) |obj| for ([_][]const u8{ "quantization", "quantization_config" }) |k| { + if (obj.get(k)) |v| if (v == .object) return v.object; + }; + return null; +} + +/// config.json's widths: MLX affine only, overrides keyed by canonical module, aliases that disagree refused. +pub fn parse(a: std.mem.Allocator, json: []const u8, mtp_layer: u32) !Widths { + const doc = try std.json.parseFromSliceLeaky(std.json.Value, a, json, .{ .allocate = .alloc_always }); + if (doc != .object) return error.BadConfig; + const q = block(doc.object) orelse { + say("glm: config.json has no quantization table; the engine reads MLX affine checkpoints", .{}); + return error.UnsupportedQuantization; + }; + const stated = quant.detect(a, .{ .root = doc.object, .quantization = q }) catch |err| { + say("glm: config.json's quantization is not MLX affine at 2, 3, 4, 5, 6 or 8 bits in groups of 32, 64 or 128, one width a module", .{}); + return err; + }; + const table = switch (stated) { + .mlx => |t| t, + else => return error.UnsupportedQuantization, + }; + var out: Widths = .{ .table = .{ .global = table.global }, .mtp_layer = mtp_layer }; + var keys: std.StringHashMapUnmanaged([]const u8) = .empty; // canonical -> the key that set it + var it = table.overrides.iterator(); + while (it.next()) |kv| { + var buf: [256]u8 = undefined; + const key = try a.dupe(u8, try canonical(&buf, kv.key_ptr.*, mtp_layer)); + const got = try out.table.overrides.getOrPut(a, key); + if (got.found_existing and !std.meta.eql(got.value_ptr.*, kv.value_ptr.*)) return conflict(keys.get(key).?, kv.key_ptr.*); + got.value_ptr.* = kv.value_ptr.*; + try keys.put(a, key, kv.key_ptr.*); + } + return out; +} + +const testing = std.testing; + +test "override keys resolve by module whatever their prefix, the MTP block as its layer" { + var buf: [256]u8 = undefined; + const cases = [_][2][]const u8{ + .{ "language_model.model.layers.3.self_attn.q_b_proj", "layers.3.self_attn.q_b_proj" }, + .{ "model.language_model.layers.3.self_attn.q_b_proj", "layers.3.self_attn.q_b_proj" }, + .{ "model.layers.3.self_attn.q_b_proj", "layers.3.self_attn.q_b_proj" }, + .{ "language_model.lm_head", "lm_head" }, + .{ "lm_head", "lm_head" }, + .{ "language_model.model.embed_tokens", "embed_tokens" }, + .{ "language_model.mtp.0.block.self_attn.o_proj", "layers.45.self_attn.o_proj" }, + .{ "language_model.mtp.0.eh_proj", "layers.45.eh_proj" }, + .{ "model.layers.7.self_attn.forget_gate.f_b_proj", "layers.7.self_attn.f_b_proj" }, + .{ "vision_model.blocks.0.attn.qkv", "vision_model.blocks.0.attn.qkv" }, + }; + for (cases) |c| try testing.expectEqualStrings(c[1], try canonical(&buf, c[0], 45)); +} + +/// An oQ-style table: global 4/64, overrides at 8, 5 and 6 bits under both prefixes, the MTP block's, a stack's. +const mixed = + \\{"text_config": {}, "quantization": {"group_size": 64, "bits": 4, "mode": "affine", + \\ "language_model.lm_head": {"group_size": 64, "bits": 8}, + \\ "language_model.model.embed_tokens": {"group_size": 64, "bits": 8}, + \\ "language_model.model.layers.0.self_attn.q_proj": {"group_size": 64, "bits": 5}, + \\ "model.layers.0.self_attn.o_proj": {"group_size": 32, "bits": 6}, + \\ "language_model.model.layers.3.mlp.shared_experts.gate_proj": {"group_size": 128, "bits": 3}, + \\ "language_model.model.layers.3.mlp.switch_mlp.down_proj": {"group_size": 64, "bits": 2}, + \\ "language_model.mtp.0.block.self_attn.kv_b_proj": {"group_size": 64, "bits": 6}, + \\ "language_model.model.layers.5.mlp.shared_experts.up_proj": {"group_size": 64, "bits": 7}}} +; + +test "each GLM tensor of a mixed oQ table resolves to its module's width, the rest to the global" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const w = try parse(arena.allocator(), mixed, 45); + const cases = [_]struct { []const u8, Width }{ + .{ "lm_head", .{ .bits = 8, .group = 64 } }, + .{ "model.language_model.embed_tokens", .{ .bits = 8, .group = 64 } }, + .{ "model.language_model.layers.0.self_attn.q_proj", .{ .bits = 5, .group = 64 } }, + .{ "model.language_model.layers.0.self_attn.k_proj", uniform }, + .{ "model.language_model.layers.0.self_attn.o_proj", .{ .bits = 6, .group = 32 } }, + .{ "model.language_model.layers.1.self_attn.o_proj", uniform }, + .{ "model.language_model.layers.3.mlp.shared_experts.gate_proj", .{ .bits = 3, .group = 128 } }, + .{ "model.language_model.layers.45.self_attn.kv_b_proj", .{ .bits = 6, .group = 64 } }, + .{ "model.language_model.layers.44.self_attn.kv_b_proj", uniform }, + }; + for (cases) |c| try testing.expectEqual(c[1], try w.of(c[0])); + try testing.expectEqual(Width{ .bits = 2, .group = 64 }, try w.expert(3, 17, "down_proj")); + try testing.expectEqual(uniform, try w.expert(3, 17, "gate_proj")); + try testing.expectEqual(uniform, try w.expert(4, 0, "down_proj")); + try testing.expectError(error.UnsupportedQuantization, w.of("model.language_model.layers.5.mlp.shared_experts.up_proj")); +} + +test "a uniform table gives every module 4/64; aliases that disagree, no table and other modes are refused" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const flat = try parse(a, "{\"quantization\": {\"group_size\": 64, \"bits\": 4, \"mode\": \"affine\"}}", 45); + try testing.expectEqual(uniform, try flat.of("model.language_model.layers.12.self_attn.q_proj")); + try testing.expectEqual(uniform, try flat.expert(12, 287, "up_proj")); + try testing.expectEqual(@as(u32, 0), flat.table.overrides.count()); + const agree = "{\"quantization_config\": {\"group_size\": 64, \"bits\": 4, \"model.layers.1.mlp.down_proj\": {\"group_size\": 64, \"bits\": 8}, \"language_model.model.layers.1.mlp.down_proj\": {\"group_size\": 64, \"bits\": 8}}}"; + try testing.expectEqual(Width{ .bits = 8, .group = 64 }, try (try parse(a, agree, 45)).of("model.language_model.layers.1.mlp.down_proj")); + const clash = "{\"quantization\": {\"group_size\": 64, \"bits\": 4, \"model.layers.1.mlp.down_proj\": {\"group_size\": 64, \"bits\": 8}, \"language_model.model.layers.1.mlp.down_proj\": {\"group_size\": 64, \"bits\": 6}}}"; + try testing.expectError(error.UnsupportedQuantization, parse(a, clash, 45)); // core/quant refuses the aliases first + const both = "{\"quantization\": {\"group_size\": 64, \"bits\": 4, \"model.layers.3.mlp.experts.2.up_proj\": {\"group_size\": 64, \"bits\": 8}, \"model.layers.3.mlp.switch_mlp.up_proj\": {\"group_size\": 64, \"bits\": 2}}}"; + const mtp = "{\"quantization\": {\"group_size\": 64, \"bits\": 4, \"language_model.mtp.0.block.self_attn.o_proj\": {\"group_size\": 64, \"bits\": 6}, \"model.layers.45.self_attn.o_proj\": {\"group_size\": 64, \"bits\": 8}}}"; + try testing.expectError(error.ConflictingQuantization, parse(a, mtp, 45)); // GLM's own alias: the MTP block is layer 45 + try testing.expectError(error.ConflictingQuantization, (try parse(a, both, 45)).expert(3, 2, "up_proj")); + try testing.expectError(error.UnsupportedQuantization, parse(a, "{\"text_config\": {}}", 45)); + try testing.expectError(error.UnsupportedQuantization, parse(a, "{\"quantization\": {\"group_size\": 64, \"bits\": 4, \"mode\": \"mxfp4\"}}", 45)); +} diff --git a/zig/src/families/glm/state.zig b/zig/src/families/glm/state.zig index ab01e184e9..2604c1ef2b 100644 --- a/zig/src/families/glm/state.zig +++ b/zig/src/families/glm/state.zig @@ -175,6 +175,7 @@ pub fn initState(arena: *Arena, c: *const cfg.Config, cap: u32, share: ?*const S pub fn init(arena: *Arena, c: *const cfg.Config, cap: u32) !struct { state: State, scratch: Scratch } { const R: usize = max_rows; const D: usize = c.hidden; + if (c.kdaWidth() > c.dense_inter) return error.ScratchTooSmall; // kdaStep puts f_b and g_b products made apart in gu const s = try initState(arena, c, cap, null); const H: usize = c.mla_heads; const V: usize = c.vocab; diff --git a/zig/src/families/glm/weights.zig b/zig/src/families/glm/weights.zig index 9fae40b3c9..5beb758f14 100644 --- a/zig/src/families/glm/weights.zig +++ b/zig/src/families/glm/weights.zig @@ -3,46 +3,43 @@ const std = @import("std"); const mtl = @import("metal"); const st = @import("../../core/safetensors.zig"); const cfg = @import("config.zig"); +const quant = @import("quant.zig"); +const affine = @import("affine.zig"); +const layout = @import("names.zig"); +const load_plan = @import("load_plan.zig"); -pub const Ref = struct { - buf: mtl.Buffer, - off: usize = 0, +pub const Width = quant.Width; - pub fn at(r: Ref, bytes: usize) Ref { - return .{ .buf = r.buf, .off = r.off + bytes }; - } - - pub fn addr(r: Ref) [*]u8 { - return r.buf.contents() + r.off; - } -}; +pub const Ref = affine.Ref; +pub const Affine = affine.Affine; -/// A 4-bit affine matrix in groups of 64: w [n, k/8] u32, s and b [n, k/64] bf16 (stacked: [e, n, ...]). -pub const Q4 = struct { - w: Ref, - s: Ref, - b: Ref, - n: u32, - k: u32, - - pub fn wBytes(n: usize, k: usize) usize { - return n * k / 2; - } - - pub fn sBytes(n: usize, k: usize) usize { - return n * (k / 64) * 2; - } -}; +/// The members of the fused matrices, in their rows' order. +pub const kda_in = [_][]const u8{ "q_proj", "k_proj", "v_proj", "f_a_proj", "g_a_proj", "b_proj" }; +pub const mla_x = [_][]const u8{ "q_a_proj", "kv_a_proj_with_mqa", "indexer.wk", "indexer.weights_proj" }; +pub const mla_qr = [_][]const u8{ "q_b_proj", "indexer.wq_b" }; +pub const gate_up = [_][]const u8{ "gate_proj", "up_proj" }; pub const Hc = struct { fnp: Ref, base: Ref, scale: Ref }; -pub const Kda = struct { in_proj: Q4, f_b: Q4, g_b: Q4, o_proj: Q4, conv_w: Ref, a_log: Ref, a: Ref, dt_bias: Ref, o_norm: Ref }; - -pub const Mla = struct { x_proj: Q4, qr_proj: Q4, kv_b: Q4, o_proj: Q4, q_norm: Ref, kv_norm: Ref, k_norm_w: Ref, k_norm_b: Ref, ape: Ref, igate: Ref }; // igate packed as a router +pub const Kda = struct { in_proj: Affine, f_b: Affine, g_b: Affine, o_proj: Affine, conv_w: Ref, a_log: Ref, a: Ref, dt_bias: Ref, o_norm: Ref }; + +pub const Mla = struct { + x_proj: Affine, + qr_proj: Affine, + kv_b: Affine, // none when the checkpoint stores mlx-lm's absorbed pair + o_proj: Affine, + q_norm: Ref, + kv_norm: Ref, + k_norm_w: Ref, + k_norm_b: Ref, + ape: Ref, + igate: Ref, // packed as a router + absorbed: ?[2]Affine = null, // mlx-lm's embed_q [H kv_lora, nope] and unembed_out [H v, kv_lora] +}; -pub const Dense = struct { gate_up: Q4, down: Q4 }; +pub const Dense = struct { gate_up: Affine, down: Affine }; -pub const Moe = struct { router: Ref, bias: Ref, gate: Q4, up: Q4, down: Q4, sh_gate_up: Q4, sh_down: Q4 }; +pub const Moe = struct { router: Ref, bias: Ref, gate: Affine, up: Affine, down: Affine, sh_gate_up: Affine, sh_down: Affine }; pub const Layer = struct { hc: ?[2]Hc, // attention's and the MLP's hyper-connections (none on the MTP layer) @@ -52,13 +49,13 @@ pub const Layer = struct { mlp: union(enum) { dense: Dense, moe: Moe }, }; -pub const Mtp = struct { eh_proj: Q4, enorm: Ref, hnorm: Ref, norm: Ref }; +pub const Mtp = struct { eh_proj: Affine, enorm: Ref, hnorm: Ref, norm: Ref, eh_bf16: ?Ref = null }; // stored bf16 pub const Weights = struct { layers: [cfg.max_layers]Layer = undefined, // the backbone's, then the MTP layer at index c.layers mtp: ?Mtp = null, - embed: Q4 = undefined, - head: Q4 = undefined, + embed: Affine = undefined, + head: Affine = undefined, norm: Ref = undefined, buffers: std.ArrayList(mtl.Buffer) = .empty, bytes: usize = 0, @@ -70,16 +67,13 @@ pub const Weights = struct { } }; -const prefix = "model.language_model."; +pub const prefix = "model.language_model."; const load_fix = @import("weights_fix.zig"); -const paced_read = @import("../../core/paced_read.zig"); const Src = load_fix.Src; const Fix = load_fix.Fix; const Transform = load_fix.Transform; -const readAll = load_fix.readAll; - -const Copy = struct { shard: u32, off: u64, len: u64, dst: [*]u8 }; +const Copy = load_fix.Copy; const Loader = struct { gpa: std.mem.Allocator, @@ -93,6 +87,7 @@ const Loader = struct { copies: std.ArrayList(Copy) = .empty, transforms: std.ArrayList(Transform) = .empty, dry: bool = false, // plan and check only: no buffers, no reads + widths: quant.Widths = undefined, // the buffer being filled and its next free byte cur: ?mtl.Buffer = null, cur_len: usize = 0, @@ -109,10 +104,24 @@ const Loader = struct { fn expect(s: Src, name: []const u8, dtype: st.DType, shape: []const usize) !void { if (s.dtype == dtype and s.rank == shape.len and std.mem.eql(usize, s.shape[0..s.rank], shape)) return; - std.log.err("glm: {s} is {t} {any}; the kernels read {t} {any}", .{ name, s.dtype, s.shape[0..s.rank], dtype, shape }); + quant.say("glm: {s} is {t} {any}; the kernels read {t} {any}", .{ name, s.dtype, s.shape[0..s.rank], dtype, shape }); return error.UnexpectedTensor; } + /// Refuses two members of one fused matrix stored at different widths. + fn mixed(a: []const u8, aw: Width, b: []const u8, bw: Width) error{MixedWidths} { + quant.say("glm: {s} is {d}-bit in groups of {d} and {s} {d}-bit in groups of {d}; one matrix holds both", .{ a, aw.bits, aw.group, b, bw.bits, bw.group }); + return error.MixedWidths; + } + + /// Input columns [lo, lo + k) of `module` at `w`, refused unless whole groups and words. + fn cutOf(module: []const u8, lo: usize, k: usize, w: Width) error{UnexpectedTensor}![2][2]usize { + return affine.cut(lo, k, w) orelse { + quant.say("glm: {s}: columns [{d}, {d}) split a group or word at {d} bits in groups of {d}", .{ module, lo, lo + k, w.bits, w.group }); + return error.UnexpectedTensor; + }; + } + /// A new buffer of `len` bytes that the next `take`s fill. fn begin(l: *Loader, len: usize) !void { const b: ?mtl.Buffer = if (l.dry) null else try l.device.buffer(@max(len, 16), mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); @@ -138,11 +147,13 @@ const Loader = struct { try l.copies.append(l.gpa, .{ .shard = s.shard, .off = s.off, .len = s.len, .dst = dst.addr() }); } - /// A tensor as stored (bf16 norms, fp32 biases), checked. + /// A tensor as stored (bf16 norms, fp32 biases), checked; fp32 stored as bf16 (mlx-lm's) is widened, exactly. fn plain(l: *Loader, comptime fmt: []const u8, args: anytype, dtype: st.DType, shape: []const usize) !Ref { const s = try l.src(fmt, args); var name: [200]u8 = undefined; - try expect(s, try std.fmt.bufPrint(&name, fmt, args), dtype, shape); + const widen = dtype == .f32 and s.dtype == .bf16; + try expect(s, try std.fmt.bufPrint(&name, fmt, args), if (widen) .bf16 else dtype, shape); + if (widen) return l.fixed(s, 2 * s.len, .to_f32); const r = l.take(s.len); try l.copyTo(s, r); return r; @@ -154,54 +165,144 @@ const Loader = struct { return r; } - /// 4-bit projections that read one input, stacked by rows in order; each part checked [n_i, k]. - fn q4(l: *Loader, k: usize, comptime fmt: []const u8, parts: []const []const u8, ns: []const usize, args: anytype) !Q4 { + pub fn has(l: *Loader, comptime fmt: []const u8, args: anytype) bool { + var name: [200]u8 = undefined; + return l.names.contains(std.fmt.bufPrint(&name, fmt, args) catch return false); + } + + /// No matrix: the place of one the checkpoint stores another way (refused on the device before it is read). + fn none(l: *Loader, k: usize) Affine { + const r = l.take(0); + return .{ .w = r, .s = r, .b = r, .n = 0, .k = @intCast(k), .bits = 0, .group = 0 }; + } + + /// Expert `e`'s three tensors cut from layer i's mlx-lm stack `switch_mlp.proj` [experts, n, ...], checked. + fn slab(l: *Loader, i: usize, e: usize, proj: []const u8, n: usize, k: usize, w: Width) ![3]Src { + var buf: [200]u8 = undefined; + const module = try std.fmt.bufPrint(&buf, prefix ++ "layers.{d}.mlp.switch_mlp.{s}", .{ i, proj }); + const sh = affine.shapes(n, k, w) orelse return error.UnexpectedTensor; + var out: [3]Src = undefined; + for ([_][]const u8{ "weight", "scales", "biases" }, &out, 0..) |comp, *o, j| { + var full: [240]u8 = undefined; + const nm = try std.fmt.bufPrint(&full, "{s}.{s}", .{ module, comp }); + const s = l.names.get(nm) orelse return error.MissingTensor; + try expect(s, nm, if (j == 0) .u32 else .bf16, &.{ l.c.experts, sh[@min(j, 1)][0], sh[@min(j, 1)][1] }); + const per = s.len / l.c.experts; + o.* = .{ .shard = s.shard, .off = s.off + e * per, .len = per, .dtype = s.dtype, .shape = .{ n, sh[@min(j, 1)][1], 1, 1, 1 }, .rank = 2 }; + } + return out; + } + + /// This Mac's heads of mlx-lm's per-head `module` stored [heads, r, k]: rows [h0 r, (h0 + H) r) as one matrix. + fn perHead(l: *Loader, comptime fmt: []const u8, args: anytype, heads: usize, h0: usize, H: usize, r: usize, k: usize) !Affine { + var buf: [200]u8 = undefined; + const module = try std.fmt.bufPrint(&buf, fmt, args); + const w = try l.widths.of(module); + const sh = affine.shapes(r, k, w) orelse return error.UnexpectedTensor; + const out: Affine = .{ .w = l.take(affine.wBytes(H * r, k, w)), .s = l.take(affine.sBytes(H * r, k, w)), .b = l.take(affine.sBytes(H * r, k, w)), .n = @intCast(H * r), .k = @intCast(k), .bits = w.bits, .group = w.group }; + for ([_][]const u8{ "weight", "scales", "biases" }, [_]Ref{ out.w, out.s, out.b }, 0..) |comp, dst, j| { + var full: [240]u8 = undefined; + const nm = try std.fmt.bufPrint(&full, "{s}.{s}", .{ module, comp }); + const s = l.names.get(nm) orelse return error.MissingTensor; + try expect(s, nm, if (j == 0) .u32 else .bf16, &.{ heads, r, sh[@min(j, 1)][1] }); + const row = if (j == 0) affine.wBytes(1, k, w) else affine.sBytes(1, k, w); + try l.copyTo(.{ .shard = s.shard, .off = s.off + h0 * r * row, .len = H * r * row, .dtype = s.dtype, .shape = s.shape, .rank = s.rank }, dst); + } + return out; + } + + /// The width config.json gives the module `fmt` names, for a buffer's budget. + pub fn widthAt(l: *Loader, comptime fmt: []const u8, args: anytype) !Width { + var name: [200]u8 = undefined; + return l.widths.of(try std.fmt.bufPrint(&name, fmt, args)); + } + + /// The weight, scales and biases of `module` stored [all, k] at `w`, checked. + fn stored(l: *Loader, module: []const u8, all: usize, k: usize, w: Width) ![3]Src { + const sh = affine.shapes(all, k, w) orelse { + quant.say("glm: {s} has {d} input columns, which do not pack at {d} bits in groups of {d}", .{ module, k, w.bits, w.group }); + return error.UnexpectedTensor; + }; + var out: [3]Src = undefined; + for ([_][]const u8{ "weight", "scales", "biases" }, &out, 0..) |comp, *o, j| { + var full: [240]u8 = undefined; + const nm = try std.fmt.bufPrint(&full, "{s}.{s}", .{ module, comp }); + o.* = l.names.get(nm) orelse { + quant.say("glm: the checkpoint has no tensor {s}", .{nm}); + return error.MissingTensor; + }; + try expect(o.*, nm, if (j == 0) .u32 else .bf16, &sh[@min(j, 1)]); + } + return out; + } + + /// Projections that read one input, stacked by rows in order; each part checked [n_i, k]. + fn mat(l: *Loader, path: affine.Path, k: usize, comptime fmt: []const u8, parts: []const []const u8, ns: []const usize, args: anytype) !Affine { const zeros: [8]usize = @splat(0); - return l.q4Rows(k, fmt, parts, zeros[0..parts.len], ns, ns, 1, args); + return l.stack(path, k, fmt, parts, zeros[0..parts.len], ns, ns, 1, args); } - /// 4-bit rows [lo, lo + n) of each part stored [all, k], stacked; room for `pad`-row tiles. - fn q4Rows(l: *Loader, k: usize, comptime fmt: []const u8, parts: []const []const u8, los: []const usize, ns: []const usize, alls: []const usize, pad: usize, args: anytype) !Q4 { - var n: usize = 0; - for (ns) |x| n += x; - const room = std.mem.alignForward(usize, n, pad); - const out: Q4 = .{ .w = l.take(Q4.wBytes(room, k)), .s = l.take(Q4.sBytes(room, k)), .b = l.take(Q4.sBytes(room, k)), .n = @intCast(n), .k = @intCast(k) }; + /// Rows [lo, lo + n) of each part stored [all, k], stacked, padded to `pad`-row tiles; mixed widths: a run a part. + fn stack(l: *Loader, path: affine.Path, k: usize, comptime fmt: []const u8, parts: []const []const u8, los: []const usize, ns: []const usize, alls: []const usize, pad: usize, args: anytype) !Affine { + var name: [200]u8 = undefined; + const base = try std.fmt.bufPrint(&name, fmt, args); + var srcs: [affine.max_runs][3]Src = undefined; + var ws: [affine.max_runs]Width = undefined; + var mixed_widths = false; + for (parts, alls, 0..) |p, all, i| { + var full: [240]u8 = undefined; + const module = try std.fmt.bufPrint(&full, "{s}{s}", .{ base, p }); + ws[i] = try l.widths.of(module); + srcs[i] = try l.stored(module, all, k, ws[i]); + mixed_widths = mixed_widths or !std.meta.eql(ws[0], ws[i]); + } + // 0.6.6's QSplit: once widths differ, every part is its own product, even beside one of its width + const runs = if (mixed_widths) parts.len else 1; + var starts: [affine.max_runs + 1]usize = undefined; // run r: parts [starts[r], starts[r + 1]) + for (0..runs) |r| starts[r] = r; + starts[runs] = parts.len; + for (parts, ws[0..parts.len], ns) |p, pw, rows| { + var full: [240]u8 = undefined; + try affine.admit(try std.fmt.bufPrint(&full, "{s}{s}", .{ base, p }), path, pw, rows); + } + var out: Affine = undefined; var row: usize = 0; - for (parts, los, ns, alls) |p, lo, rows, all| { - const comps = [_][]const u8{ "weight", "scales", "biases" }; - const dsts = [_]Ref{ out.w.at(Q4.wBytes(row, k)), out.s.at(Q4.sBytes(row, k)), out.b.at(Q4.sBytes(row, k)) }; - for (comps, dsts, 0..) |comp, dst, j| { - var name: [200]u8 = undefined; - const base = try std.fmt.bufPrint(&name, fmt, args); - var full: [240]u8 = undefined; - const nm = try std.fmt.bufPrint(&full, "{s}{s}.{s}", .{ base, p, comp }); - const s = l.names.get(nm) orelse { - std.log.err("glm: the checkpoint has no tensor {s}", .{nm}); - return error.MissingTensor; - }; - try expect(s, nm, if (j == 0) .u32 else .bf16, &.{ all, if (j == 0) k / 8 else k / 64 }); - const row_bytes: usize = if (j == 0) k / 2 else k / 64 * 2; - try l.copyTo(.{ .shard = s.shard, .off = s.off + lo * row_bytes, .len = rows * row_bytes, .dtype = s.dtype, .shape = s.shape, .rank = s.rank }, dst); + for (0..runs) |r| { + const w = ws[starts[r]]; + var n: usize = 0; + for (ns[starts[r]..starts[r + 1]]) |x| n += x; + const room = std.mem.alignForward(usize, n, pad); + const run: affine.Run = .{ .w = l.take(affine.wBytes(room, k, w)), .s = l.take(affine.sBytes(room, k, w)), .b = l.take(affine.sBytes(room, k, w)), .row = @intCast(row), .n = @intCast(n), .bits = w.bits, .group = w.group }; + if (r == 0) out = .{ .w = run.w, .s = run.s, .b = run.b, .n = 0, .k = @intCast(k), .bits = w.bits, .group = w.group } else out.runs[r - 1] = run; + var at: usize = 0; // the run's own rows + for (srcs[starts[r]..starts[r + 1]], los[starts[r]..starts[r + 1]], ns[starts[r]..starts[r + 1]]) |three, lo, x| { + const dsts = [_]Ref{ run.w.at(affine.wBytes(at, k, w)), run.s.at(affine.sBytes(at, k, w)), run.b.at(affine.sBytes(at, k, w)) }; + for (three, dsts, 0..) |s, dst, j| { + const row_bytes = if (j == 0) affine.wBytes(1, k, w) else affine.sBytes(1, k, w); + try l.copyTo(.{ .shard = s.shard, .off = s.off + lo * row_bytes, .len = x * row_bytes, .dtype = s.dtype, .shape = s.shape, .rank = s.rank }, dst); + } + at += x; } - row += rows; + row += n; } + out.n = @intCast(row); + out.more = @intCast(runs - 1); return out; } - /// One 4-bit matrix stored [n, k_all]: its input columns [lo, lo + k) (TP2: o_proj over one Mac's heads). - fn q4Cols(l: *Loader, comptime fmt: []const u8, args: anytype, n: usize, k_all: usize, lo: usize, k: usize) !Q4 { - const out: Q4 = .{ .w = l.take(Q4.wBytes(n, k)), .s = l.take(Q4.sBytes(n, k)), .b = l.take(Q4.sBytes(n, k)), .n = @intCast(n), .k = @intCast(k) }; - const comps = [_][]const u8{ "weight", "scales", "biases" }; - for (comps, [_]Ref{ out.w, out.s, out.b }, 0..) |comp, dst, j| { - var name: [200]u8 = undefined; - const base = try std.fmt.bufPrint(&name, fmt, args); - var full: [240]u8 = undefined; - const nm = try std.fmt.bufPrint(&full, "{s}.{s}", .{ base, comp }); - const s = l.names.get(nm) orelse return error.MissingTensor; - try expect(s, nm, if (j == 0) .u32 else .bf16, &.{ n, if (j == 0) k_all / 8 else k_all / 64 }); - const row_all: usize = if (j == 0) k_all / 2 else k_all / 64 * 2; - const span: [2]usize = if (j == 0) .{ lo / 2, k / 2 } else .{ lo / 64 * 2, k / 64 * 2 }; - if (!l.dry) try l.transforms.append(l.gpa, .{ .src = s, .dst = dst.addr(), .fix = .cols, .arg = .{ row_all, span[0], span[1] } }); + /// One matrix stored [n, k_all]: its input columns [lo, lo + k) (TP2: o_proj over one Mac's heads). + fn cols(l: *Loader, path: affine.Path, comptime fmt: []const u8, args: anytype, n: usize, k_all: usize, lo: usize, k: usize) !Affine { + var name: [200]u8 = undefined; + const module = try std.fmt.bufPrint(&name, fmt, args); + const w = try l.widths.of(module); + const three = try l.stored(module, n, k_all, w); + try affine.admit(module, path, w, n); + const span = try cutOf(module, lo, k, w); + const out: Affine = .{ .w = l.take(affine.wBytes(n, k, w)), .s = l.take(affine.sBytes(n, k, w)), .b = l.take(affine.sBytes(n, k, w)), .n = @intCast(n), .k = @intCast(k), .bits = w.bits, .group = w.group }; + for (three, [_]Ref{ out.w, out.s, out.b }, 0..) |s, dst, j| { + const row_all = if (j == 0) affine.wBytes(1, k_all, w) else affine.sBytes(1, k_all, w); + const at = span[@min(j, 1)]; + if (!l.dry) try l.transforms.append(l.gpa, .{ .src = s, .dst = dst.addr(), .fix = .cols, .arg = .{ row_all, at[0], at[1] } }); } return out; } @@ -210,14 +311,18 @@ const Loader = struct { fn plainRange(l: *Loader, comptime fmt: []const u8, args: anytype, dtype: st.DType, all: usize, lo: usize, n: usize, size: usize) !Ref { const s = try l.src(fmt, args); var name: [200]u8 = undefined; - try expect(s, try std.fmt.bufPrint(&name, fmt, args), dtype, &.{all}); + const widen = dtype == .f32 and s.dtype == .bf16; + try expect(s, try std.fmt.bufPrint(&name, fmt, args), if (widen) .bf16 else dtype, &.{all}); + const at = if (widen) size / 2 else size; // a stored element's bytes + const part: Src = .{ .shard = s.shard, .off = s.off + lo * at, .len = n * at, .dtype = s.dtype, .shape = s.shape, .rank = s.rank }; + if (widen) return l.fixed(part, n * size, .to_f32); const r = l.take(n * size); - try l.copyTo(.{ .shard = s.shard, .off = s.off + lo * size, .len = n * size, .dtype = s.dtype, .shape = s.shape, .rank = s.rank }, r); + try l.copyTo(part, r); return r; } - /// This Mac's routed experts' `proj` stacked [own, n, k]; by rows, gate/up keep rows `inter` and down keeps input columns `inter`. - fn experts(l: *Loader, i: usize, proj: []const u8, n_all: usize, k_all: usize, part: enum { whole, rows, cols }) !Q4 { + /// This Mac's routed experts' `proj` stacked [own, n, k] at one width; by rows, gate/up keep rows `inter`. + fn experts(l: *Loader, i: usize, proj: []const u8, n_all: usize, k_all: usize, part: enum { whole, rows, cols }) !Affine { const every = i == l.c.layers; // the MTP layer: every expert whole on every Mac (its drafts need no exchange) const lo = if (every) 0 else l.c.own[0]; const e = (if (every) l.c.experts else l.c.own[1]) - lo; @@ -225,19 +330,26 @@ const Loader = struct { const in_n: usize = l.c.inter[1] - l.c.inter[0]; const n = if (part == .rows) in_n else n_all; const k = if (part == .cols) in_n else k_all; - const out: Q4 = .{ .w = l.take(e * Q4.wBytes(n, k)), .s = l.take(e * Q4.sBytes(n, k)), .b = l.take(e * Q4.sBytes(n, k)), .n = @intCast(n), .k = @intCast(k) }; + var first: [200]u8 = undefined; + const name0 = try std.fmt.bufPrint(&first, prefix ++ "layers.{d}.mlp.experts.{d}.{s}", .{ i, lo, proj }); + const w = try l.widths.expert(i, lo, proj); + try affine.admit(name0, .routed_experts, w, n); + const span = if (part == .cols) try cutOf(name0, in_lo, in_n, w) else undefined; + const out: Affine = .{ .w = l.take(e * affine.wBytes(n, k, w)), .s = l.take(e * affine.sBytes(n, k, w)), .b = l.take(e * affine.sBytes(n, k, w)), .n = @intCast(n), .k = @intCast(k), .bits = w.bits, .group = w.group }; for (0..e) |x| { - const comps = [_][]const u8{ "weight", "scales", "biases" }; - const dsts = [_]Ref{ out.w.at(x * Q4.wBytes(n, k)), out.s.at(x * Q4.sBytes(n, k)), out.b.at(x * Q4.sBytes(n, k)) }; - for (comps, dsts, 0..) |comp, dst, j| { - const s = try l.src(prefix ++ "layers.{d}.mlp.experts.{d}.{s}.{s}", .{ i, lo + x, proj, comp }); - var name: [200]u8 = undefined; - try expect(s, try std.fmt.bufPrint(&name, "experts.{d}.{d}.{s}.{s}", .{ i, x, proj, comp }), if (j == 0) .u32 else .bf16, &.{ n_all, if (j == 0) k_all / 8 else k_all / 64 }); - const row_all: usize = if (j == 0) k_all / 2 else k_all / 64 * 2; // a stored row's bytes + var buf: [200]u8 = undefined; + const module = try std.fmt.bufPrint(&buf, prefix ++ "layers.{d}.mlp.experts.{d}.{s}", .{ i, lo + x, proj }); + const xw = try l.widths.expert(i, lo + x, proj); + if (!std.meta.eql(w, xw)) return mixed(name0, w, module, xw); + const own = l.has("{s}.weight", .{module}); + const three = if (own) try l.stored(module, n_all, k_all, w) else try l.slab(i, lo + x, proj, n_all, k_all, w); + const dsts = [_]Ref{ out.w.at(x * affine.wBytes(n, k, w)), out.s.at(x * affine.sBytes(n, k, w)), out.b.at(x * affine.sBytes(n, k, w)) }; + for (three, dsts, 0..) |s, dst, j| { + const row_all = if (j == 0) affine.wBytes(1, k_all, w) else affine.sBytes(1, k_all, w); switch (part) { .whole => try l.copyTo(s, dst), .rows => try l.copyTo(.{ .shard = s.shard, .off = s.off + in_lo * row_all, .len = in_n * row_all, .dtype = s.dtype, .shape = s.shape, .rank = s.rank }, dst), - .cols => if (!l.dry) try l.transforms.append(l.gpa, .{ .src = s, .dst = dst.addr(), .fix = .cols, .arg = .{ row_all, if (j == 0) in_lo / 2 else in_lo / 64 * 2, if (j == 0) in_n / 2 else in_n / 64 * 2 } }), + .cols => if (!l.dry) try l.transforms.append(l.gpa, .{ .src = s, .dst = dst.addr(), .fix = .cols, .arg = .{ row_all, span[@min(j, 1)][0], span[@min(j, 1)][1] } }), } } } @@ -247,14 +359,32 @@ const Loader = struct { fn hc(l: *Loader, i: usize, site: []const u8) !Hc { const D: usize = l.c.hidden; const f = try l.src(prefix ++ "layers.{d}.hc_{s}_fn", .{ i, site }); - try expect(f, "hc fn", .bf16, &.{ 24, 4 * D }); + try expect(f, "hc fn", if (f.dtype == .f32) .f32 else .bf16, &.{ 24, 4 * D }); // fp32 (mlx-lm's): bf16 if exact return .{ - .fnp = try l.fixed(f, f.len, .hc_pack), + .fnp = try l.fixed(f, 24 * 4 * D * 2, .hc_pack), .base = try l.plain(prefix ++ "layers.{d}.hc_{s}_base", .{ i, site }, .f32, &.{24}), .scale = try l.plain(prefix ++ "layers.{d}.hc_{s}_scale", .{ i, site }, .f32, &.{3}), }; } + /// Layer i's q, k, v conv taps [channels, 1, taps]: three tensors, or mlx-lm's one [3 channels, taps, 1] cut up. + fn convs(l: *Loader, i: usize, channels: usize) ![3]Src { + const a = prefix ++ "layers.{d}.self_attn."; + if (l.has(a ++ "conv1d.weight", .{i})) { + const s = try l.src(a ++ "conv1d.weight", .{i}); + try expect(s, "conv1d", .bf16, &.{ 3 * channels, l.c.conv, 1 }); + var out: [3]Src = undefined; + for (&out, 0..) |*o, p| o.* = .{ .shard = s.shard, .off = s.off + p * s.len / 3, .len = s.len / 3, .dtype = .bf16, .shape = .{ channels, 1, l.c.conv, 1, 1 }, .rank = 3 }; + return out; + } + var out: [3]Src = undefined; + inline for (.{ "q", "k", "v" }, 0..) |x, p| { + out[p] = try l.src(a ++ x ++ "_conv1d.weight", .{i}); + try expect(out[p], x ++ "_conv1d", .bf16, &.{ channels, 1, l.c.conv }); + } + return out; + } + fn kda(l: *Loader, i: usize) !Kda { const c = l.c; const D: usize = c.hidden; @@ -266,10 +396,9 @@ const Loader = struct { const Wa = W * c.tp; // every head's channels, as stored const Ha = H * c.tp; const a = prefix ++ "layers.{d}.self_attn."; - const conv_q = try l.src(a ++ "q_conv1d.weight", .{i}); - try expect(conv_q, "q_conv1d", .bf16, &.{ Wa, 1, c.conv }); - const conv = try l.fixed(conv_q, 3 * W * c.conv * 4, .conv); - const extra: [3]?Src = .{ try l.src(a ++ "k_conv1d.weight", .{i}), try l.src(a ++ "v_conv1d.weight", .{i}), null }; + const qkv = try l.convs(i, Wa); + const conv = try l.fixed(qkv[0], 3 * W * c.conv * 4, .conv); + const extra: [3]?Src = .{ qkv[1], qkv[2], null }; if (!l.dry) { const t = &l.transforms.items[l.transforms.items.len - 1]; t.extra = extra; @@ -278,10 +407,10 @@ const Loader = struct { const o_norm = try l.src(a ++ "o_norm.weight", .{i}); try expect(o_norm, "o_norm", .bf16, &.{dk}); return .{ - .in_proj = try l.q4Rows(D, a, &.{ "q_proj", "k_proj", "v_proj", "f_a_proj", "g_a_proj", "b_proj" }, &.{ ch0, ch0, ch0, 0, 0, h0 }, &.{ W, W, W, dk, dk, H }, &.{ Wa, Wa, Wa, dk, dk, Ha }, 64, .{i}), - .f_b = try l.q4Rows(dk, a, &.{"f_b_proj"}, &.{ch0}, &.{W}, &.{Wa}, 1, .{i}), - .g_b = try l.q4Rows(dk, a, &.{"g_b_proj"}, &.{ch0}, &.{W}, &.{Wa}, 1, .{i}), - .o_proj = if (c.tp == 1) try l.q4(W, a, &.{"o_proj"}, &.{D}, .{i}) else try l.q4Cols(a ++ "o_proj", .{i}, D, Wa, ch0, W), + .in_proj = try l.stack(.kda_in, D, a, &kda_in, &.{ ch0, ch0, ch0, 0, 0, h0 }, &.{ W, W, W, dk, dk, H }, &.{ Wa, Wa, Wa, dk, dk, Ha }, 64, .{i}), + .f_b = try l.stack(.kda_gate, dk, a, &.{"f_b_proj"}, &.{ch0}, &.{W}, &.{Wa}, 1, .{i}), + .g_b = try l.stack(.kda_gate, dk, a, &.{"g_b_proj"}, &.{ch0}, &.{W}, &.{Wa}, 1, .{i}), + .o_proj = if (c.tp == 1) try l.mat(.kda_out, W, a, &.{"o_proj"}, &.{D}, .{i}) else try l.cols(.kda_out, a ++ "o_proj", .{i}, D, Wa, ch0, W), .conv_w = conv, .a_log = try l.plainRange(a ++ "A_log", .{i}, .f32, Ha, h0, H, 4), .a = l.take(H * 4), @@ -299,11 +428,13 @@ const Loader = struct { const Ha = H * c.tp; // every head, as stored const gate = try l.src(a ++ "indexer.index_kpool_compress_gate", .{i}); try expect(gate, "index_kpool_compress_gate", .bf16, &.{ c.i_dim, D }); + const pair = l.has(a ++ "embed_q.weight", .{i}); return .{ - .x_proj = try l.q4Rows(D, a, &.{ "q_a_proj", "kv_a_proj_with_mqa", "indexer.wk", "indexer.weights_proj" }, &.{ 0, 0, 0, 0 }, &.{ c.q_lora, c.kv_lora, c.i_dim, c.i_heads }, &.{ c.q_lora, c.kv_lora, c.i_dim, c.i_heads }, 64, .{i}), - .qr_proj = try l.q4Rows(c.q_lora, a, &.{ "q_b_proj", "indexer.wq_b" }, &.{ h0 * c.nope, 0 }, &.{ H * c.nope, c.i_heads * c.i_dim }, &.{ Ha * c.nope, c.i_heads * c.i_dim }, 1, .{i}), - .kv_b = try l.q4Rows(c.kv_lora, a, &.{"kv_b_proj"}, &.{h0 * (c.nope + c.v_dim)}, &.{H * (c.nope + c.v_dim)}, &.{Ha * (c.nope + c.v_dim)}, 1, .{i}), - .o_proj = if (c.tp == 1) try l.q4(H * c.v_dim, a, &.{"o_proj"}, &.{D}, .{i}) else try l.q4Cols(a ++ "o_proj", .{i}, D, Ha * c.v_dim, h0 * c.v_dim, H * c.v_dim), + .x_proj = try l.stack(.mla_in, D, a, &mla_x, &.{ 0, 0, 0, 0 }, &.{ c.q_lora, c.kv_lora, c.i_dim, c.i_heads }, &.{ c.q_lora, c.kv_lora, c.i_dim, c.i_heads }, 64, .{i}), + .qr_proj = try l.stack(.mla_q, c.q_lora, a, &mla_qr, &.{ h0 * c.nope, 0 }, &.{ H * c.nope, c.i_heads * c.i_dim }, &.{ Ha * c.nope, c.i_heads * c.i_dim }, 1, .{i}), + .kv_b = if (pair) l.none(c.kv_lora) else try l.stack(.kv_b, c.kv_lora, a, &.{"kv_b_proj"}, &.{h0 * (c.nope + c.v_dim)}, &.{H * (c.nope + c.v_dim)}, &.{Ha * (c.nope + c.v_dim)}, 1, .{i}), + .absorbed = if (pair) .{ try l.perHead(a ++ "embed_q", .{i}, Ha, h0, H, c.kv_lora, c.nope), try l.perHead(a ++ "unembed_out", .{i}, Ha, h0, H, c.v_dim, c.kv_lora) } else null, + .o_proj = if (c.tp == 1) try l.mat(.mla_out, H * c.v_dim, a, &.{"o_proj"}, &.{D}, .{i}) else try l.cols(.mla_out, a ++ "o_proj", .{i}, D, Ha * c.v_dim, h0 * c.v_dim, H * c.v_dim), .q_norm = try l.plain(a ++ "q_a_layernorm.weight", .{i}, .bf16, &.{c.q_lora}), .kv_norm = try l.plain(a ++ "kv_a_layernorm.weight", .{i}, .bf16, &.{c.kv_lora}), .k_norm_w = try l.plain(a ++ "indexer.k_norm.weight", .{i}, .bf16, &.{c.i_dim}), @@ -313,42 +444,11 @@ const Loader = struct { }; } - /// Bytes of layer i's dense buffer (everything but routed experts), generously rounded. - fn layerBytes(l: *Loader, i: usize) usize { - const c = l.c; - const D: usize = c.hidden; - var n: usize = 0; - const q = struct { - fn b(rows: usize, k: usize) usize { - return Q4.wBytes(rows, k) + 2 * Q4.sBytes(rows, k) + 3 * 256; - } - }; - n += 2 * (24 * 4 * D * 2 + 1024) + 2 * (D * 2 + 256); - switch (c.kind(@intCast(i))) { - .kda => n += q.b(std.mem.alignForward(usize, c.kdaProj(), 64), D) + 2 * q.b(c.kdaWidth(), c.kda_dim) + q.b(D, c.kdaWidth()) + 3 * c.kdaWidth() * c.conv * 4 + 8 * 1024 + c.kdaWidth() * 4, - .mla => n += q.b(std.mem.alignForward(usize, c.xProj(), 64), D) + q.b(c.qrProj(), c.q_lora) + q.b(c.mla_heads * (c.nope + c.v_dim), c.kv_lora) + q.b(D, c.mla_heads * c.v_dim) + 16 * 1024 + c.i_dim * D * 2, - } - if (c.isMoe(@intCast(i))) { - n += c.experts * D * 2 + c.experts * 4 + q.b(2 * c.moe_inter, D) + q.b(D, c.moe_inter) + 2048; - } else n += q.b(2 * c.dense_inter, D) + q.b(D, c.dense_inter); - if (i == c.layers) n += q.b(D, 2 * D) + 3 * (D * 2 + 256); - return n; - } - - fn expertBytes(l: *Loader, i: usize) usize { - const c = l.c; - const D: usize = c.hidden; - const every = i == c.layers; - const m: usize = if (every) c.moe_inter else c.inter[1] - c.inter[0]; - const per = 2 * (Q4.wBytes(m, D) + 2 * Q4.sBytes(m, D)) + Q4.wBytes(D, m) + 2 * Q4.sBytes(D, m); - return (if (every) c.experts else c.own[1] - c.own[0]) * per + 9 * 256; - } - fn layer(l: *Loader, i: usize) !Layer { const c = l.c; const D: usize = c.hidden; const plain_layer = i == c.layers; // the MTP block: no hyper-connections - try l.begin(l.layerBytes(i)); + try l.begin(try load_plan.layerBytes(l, i)); var out: Layer = .{ .hc = if (plain_layer) null else .{ try l.hc(i, "attn"), try l.hc(i, "ffn") }, .in_norm = try l.plain(prefix ++ "layers.{d}.input_layernorm.weight", .{i}, .bf16, &.{D}), @@ -361,8 +461,10 @@ const Loader = struct { }; if (plain_layer) { // the MTP head's own tensors, in the layer's buffer (before the experts begin theirs) const p = prefix ++ "layers.{d}."; + const bf16 = !l.has(p ++ "eh_proj.scales", .{i}); l.w.mtp = .{ - .eh_proj = try l.q4(2 * D, p, &.{"eh_proj"}, &.{D}, .{i}), + .eh_proj = if (bf16) l.none(2 * D) else try l.mat(.eh_proj, 2 * D, p, &.{"eh_proj"}, &.{D}, .{i}), + .eh_bf16 = if (bf16) try l.plain(p ++ "eh_proj.weight", .{i}, .bf16, &.{ D, 2 * D }) else null, .enorm = try l.plain(p ++ "enorm.weight", .{i}, .bf16, &.{D}), .hnorm = try l.plain(p ++ "hnorm.weight", .{i}, .bf16, &.{D}), .norm = try l.plain(p ++ "shared_head.norm.weight", .{i}, .bf16, &.{D}), @@ -371,139 +473,39 @@ const Loader = struct { const m = prefix ++ "layers.{d}.mlp."; if (c.isMoe(@intCast(i))) { const router = try l.src(m ++ "gate.weight", .{i}); - try expect(router, "mlp.gate.weight", .bf16, &.{ c.experts, D }); + try expect(router, "mlp.gate.weight", if (router.dtype == .f32) .f32 else .bf16, &.{ c.experts, D }); var moe: Moe = .{ - .router = try l.fixed(router, router.len, .router_pack), + .router = try l.fixed(router, @as(usize, c.experts) * D * 2, .router_pack), .bias = try l.plain(m ++ "gate.e_score_correction_bias", .{i}, .f32, &.{c.experts}), - .sh_gate_up = try l.q4(D, m ++ "shared_experts.", &.{ "gate_proj", "up_proj" }, &.{ c.moe_inter, c.moe_inter }, .{i}), - .sh_down = try l.q4(c.moe_inter, m ++ "shared_experts.", &.{"down_proj"}, &.{D}, .{i}), + .sh_gate_up = try l.mat(.shared_expert, D, m ++ "shared_experts.", &gate_up, &.{ c.moe_inter, c.moe_inter }, .{i}), + .sh_down = try l.mat(.shared_expert, c.moe_inter, m ++ "shared_experts.", &.{"down_proj"}, &.{D}, .{i}), .gate = undefined, .up = undefined, .down = undefined, }; - try l.begin(l.expertBytes(i)); + try l.begin(try load_plan.expertBytes(l, i)); const by_rows = c.byRows() and i != c.layers; moe.gate = try l.experts(i, "gate_proj", c.moe_inter, D, if (by_rows) .rows else .whole); moe.up = try l.experts(i, "up_proj", c.moe_inter, D, if (by_rows) .rows else .whole); + if (!std.meta.eql(moe.gate.width(), moe.up.width())) { // one launch takes both, SwiGLU'd + var buf: [96]u8 = undefined; + return mixed(try std.fmt.bufPrint(&buf, "layer {d}'s routed gate_proj", .{i}), moe.gate.width(), "up_proj", moe.up.width()); + } moe.down = try l.experts(i, "down_proj", D, c.moe_inter, if (by_rows) .cols else .whole); out.mlp = .{ .moe = moe }; } else { const n: usize = c.dense_inter; // this Mac's rows; TP2: rank r's half [r n, (r + 1) n) const lo = c.tp_rank * n; out.mlp = .{ .dense = .{ - .gate_up = try l.q4Rows(D, m, &.{ "gate_proj", "up_proj" }, &.{ lo, lo }, &.{ n, n }, &.{ n * c.tp, n * c.tp }, 1, .{i}), - .down = if (c.tp == 1) try l.q4(n, m, &.{"down_proj"}, &.{D}, .{i}) else try l.q4Cols(m ++ "down_proj", .{i}, D, n * c.tp, lo, n), + .gate_up = try l.stack(.dense_mlp, D, m, &gate_up, &.{ lo, lo }, &.{ n, n }, &.{ n * c.tp, n * c.tp }, 1, .{i}), + .down = if (c.tp == 1) try l.mat(.dense_mlp, n, m, &.{"down_proj"}, &.{D}, .{i}) else try l.cols(.dense_mlp, m ++ "down_proj", .{i}, D, n * c.tp, lo, n), } }; } return out; } - - /// The index and every shard's header: each tensor's shard and byte range. - fn index(l: *Loader) !void { - const path = try std.fmt.allocPrintSentinel(l.arena, "{s}/model.safetensors.index.json", .{l.dir}, 0); - const f = try mtl.MappedFile.open(path); - defer f.deinit(); - const doc = try std.json.parseFromSliceLeaky(std.json.Value, l.arena, f.bytes[0..f.size], .{ .allocate = .alloc_always }); - var files: std.StringArrayHashMapUnmanaged(void) = .empty; - var it = doc.object.get("weight_map").?.object.iterator(); - while (it.next()) |kv| try files.put(l.arena, kv.value_ptr.string, {}); - for (files.keys(), 0..) |name, si| { - const full = try std.fmt.allocPrintSentinel(l.arena, "{s}/{s}", .{ l.dir, name }, 0); - try l.shards.append(l.gpa, full); - const fd = std.c.open(full, .{ .ACCMODE = .RDONLY }); - if (fd < 0) return error.OpenFailed; - defer _ = std.c.close(fd); - var head: [8]u8 = undefined; - try readAll(fd, &head, 0); - const hlen = std.mem.readInt(u64, &head, .little); - const header = try l.arena.alloc(u8, hlen); - try readAll(fd, header, 8); - const end = std.c.lseek(fd, 0, std.c.SEEK.END); - const data = 8 + hlen; - const doc2 = try std.json.parseFromSliceLeaky(std.json.Value, l.arena, header, .{}); - var eit = doc2.object.iterator(); - while (eit.next()) |kv| { - const tensor = kv.key_ptr.*; - // the language model's tensors only: the vision tower's (5-D patch kernels among them) are never read - if (!std.mem.startsWith(u8, tensor, prefix) and !std.mem.startsWith(u8, tensor, "lm_head.")) continue; - const o = kv.value_ptr.object; - const dtype = st.DType.parse(o.get("dtype").?.string) orelse return error.UnsupportedDType; - const dims = o.get("shape").?.array.items; - if (dims.len > st.max_rank) return error.RankTooHigh; - var shape: [st.max_rank]usize = @splat(1); - for (dims, 0..) |d, j| shape[j] = @intCast(d.integer); - const offs = o.get("data_offsets").?.array.items; - const b: u64 = @intCast(offs[0].integer); - const e2: u64 = @intCast(offs[1].integer); - if (e2 < b or data + e2 > @as(u64, @intCast(end))) return error.BadSafetensors; - try l.names.put(l.arena, tensor, .{ .shard = @intCast(si), .off = data + b, .len = e2 - b, .dtype = dtype, .shape = shape, .rank = @intCast(dims.len) }); - } - } - } -}; - -const Pool = struct { - fds: []std.c.fd_t, - jobs: []const Copy, - next: std.atomic.Value(usize) = .init(0), - failed: std.atomic.Value(bool) = .init(false), - - fn run(p: *Pool) void { - const buf = paced_read.stage(std.heap.page_allocator, 16 << 20) catch { - p.failed.store(true, .release); - return; - }; - defer std.heap.page_allocator.free(buf); - while (true) { - const i = p.next.fetchAdd(1, .monotonic); - if (i >= p.jobs.len) return; - const j = p.jobs[i]; - paced_read.read(p.fds[j.shard], j.dst[0..j.len], j.off, buf) catch p.failed.store(true, .release); - } - } }; -/// Every copy on `threads` threads, uncached and paced by free memory (a large load otherwise fills the page cache too). -fn runCopies(l: *Loader, threads: usize) !void { - const fds = try l.gpa.alloc(std.c.fd_t, l.shards.items.len); - defer l.gpa.free(fds); - for (l.shards.items, fds) |path, *fd| { - fd.* = std.c.open(path, .{ .ACCMODE = .RDONLY }); - if (fd.* < 0) return error.OpenFailed; - _ = std.c.fcntl(fd.*, 48, @as(c_int, 1)); // F_NOCACHE - } - defer for (fds) |fd| { - _ = std.c.close(fd); - }; - var pool: Pool = .{ .fds = fds, .jobs = l.copies.items }; - const workers = try l.gpa.alloc(?std.Thread, threads); - defer l.gpa.free(workers); - for (workers) |*t| t.* = std.Thread.spawn(.{}, Pool.run, .{&pool}) catch null; - pool.run(); - for (workers) |t| if (t) |th| th.join(); - if (pool.failed.load(.acquire)) return error.ShortRead; - // the transforms on the same threads (each reads its own tensors whole) - const Tx = struct { - l: *Loader, - fds: []std.c.fd_t, - next: std.atomic.Value(usize) = .init(0), - failed: std.atomic.Value(bool) = .init(false), - fn run(x: *@This()) void { - while (true) { - const i = x.next.fetchAdd(1, .monotonic); - if (i >= x.l.transforms.items.len) return; - load_fix.transform(x.l.gpa, x.fds, x.l.transforms.items[i]) catch x.failed.store(true, .release); - } - } - }; - var tx: Tx = .{ .l = l, .fds = fds }; - for (workers) |*t| t.* = std.Thread.spawn(.{}, Tx.run, .{&tx}) catch null; - tx.run(); - for (workers) |t| if (t) |th| th.join(); - if (tx.failed.load(.acquire)) return error.TransformFailed; -} - -/// The first `c.run` layers, the MTP layer (when stored) and the head; `dry`: check names, dtypes and shapes, read nothing. +/// The first `c.run` layers, the MTP layer (if stored), the head; `dry`: check names, dtypes and shapes, read nothing. pub fn load(gpa: std.mem.Allocator, device: mtl.Device, dir: []const u8, c: *const cfg.Config, threads: usize, dry: bool) !*Weights { var arena_state = std.heap.ArenaAllocator.init(gpa); defer arena_state.deinit(); @@ -517,21 +519,62 @@ pub fn load(gpa: std.mem.Allocator, device: mtl.Device, dir: []const u8, c: *con defer l.shards.deinit(gpa); defer l.copies.deinit(gpa); defer l.transforms.deinit(gpa); - try l.index(); + try layout.read(gpa, l.arena, dir, c.layers, &l.shards, &l.names); + l.widths = try readWidths(l.arena, dir, c); const D: usize = c.hidden; const V: usize = c.vocab; - try l.begin(2 * (Q4.wBytes(V, D) + 2 * Q4.sBytes(V, D)) + D * 2 + 8 * 256); - w.embed = try l.q4(D, prefix ++ "embed_tokens", &.{""}, &.{V}, .{}); + try l.begin(affine.bytes(V, D, try l.widths.of(prefix ++ "embed_tokens")) + affine.bytes(V, D, try l.widths.of("lm_head")) + D * 2 + 8 * 256); + w.embed = try l.mat(.embed, D, prefix ++ "embed_tokens", &.{""}, &.{V}, .{}); const vp = c.vocabPart(); // TP2: this Mac's half of the head's rows - w.head = try l.q4Rows(D, "lm_head", &.{""}, &.{vp[0]}, &.{vp[1] - vp[0]}, &.{V}, 1, .{}); + w.head = try l.stack(.head, D, "lm_head", &.{""}, &.{vp[0]}, &.{vp[1] - vp[0]}, &.{V}, 1, .{}); w.norm = try l.plain(prefix ++ "norm.weight", .{}, .bf16, &.{D}); for (0..c.run) |i| w.layers[i] = try l.layer(i); const has_mtp = c.mtp > 0 and l.names.contains(prefix ++ "layers.45.eh_proj.weight"); if (has_mtp) w.layers[c.layers] = try l.layer(c.layers); + if ((c.tp > 1 or c.byRows() or c.own[1] - c.own[0] != c.experts) and !native(w, c)) { + quant.say("glm: this checkpoint has widths other than 4/64; the pair's split kernels read only 4/64 (load it on one Mac)", .{}); + return error.UnsupportedWidth; + } if (dry) { - std.debug.print("glm plan: {d} of {d} layers{s}, {d:.2} GB in buffers, every tensor's name, dtype and shape checked\n", .{ c.run, c.layers, if (has_mtp) " and the MTP layer" else "", @as(f64, @floatFromInt(w.bytes)) / 1e9 }); + var buf: [256]u8 = undefined; + const widths = tally(w, c).summary(&buf); + std.debug.print("glm plan: {d} of {d} layers{s}, {d:.2} GB in buffers, every tensor's name, dtype and shape checked; {s}\n", .{ c.run, c.layers, if (has_mtp) " and the MTP layer" else "", @as(f64, @floatFromInt(w.bytes)) / 1e9, widths }); return w; } - try runCopies(&l, threads); + try load_fix.runCopies(gpa, l.shards.items, l.copies.items, l.transforms.items, threads); return w; } + +/// A plan's matrices by width: the head, the embedding, every loaded layer. +pub fn tally(w: *const Weights, c: *const cfg.Config) affine.Tally { + var t: affine.Tally = .{}; + t.add(w.embed); + t.add(w.head); + t.add(w.mtp); + for (w.layers[0..c.run]) |l| t.add(l); + if (w.mtp != null) t.add(w.layers[c.layers]); + return t; +} + +/// Whether every loaded matrix takes the generated 4/64 kernels (no other width's pipelines needed). +pub fn native(w: *const Weights, c: *const cfg.Config) bool { + if (!affine.allNative(.{ w.embed, w.head, w.mtp })) return false; + if (w.mtp) |m| if (m.eh_bf16 != null) return false; + for (w.layers[0..c.run]) |l| if (!layerNative(l)) return false; + return w.mtp == null or layerNative(w.layers[c.layers]); +} + +fn layerNative(l: Layer) bool { + return switch (l.attn) { + .mla => |m| m.absorbed == null and affine.allNative(l), + .kda => affine.allNative(l), + }; +} + +/// config.json's widths by module. +fn readWidths(arena: std.mem.Allocator, dir: []const u8, c: *const cfg.Config) !quant.Widths { + const path = try std.fmt.allocPrintSentinel(arena, "{s}/config.json", .{dir}, 0); + const f = try mtl.MappedFile.open(path); + defer f.deinit(); + return quant.parse(arena, f.bytes[0..f.size], c.layers); +} diff --git a/zig/src/families/glm/weights_fix.zig b/zig/src/families/glm/weights_fix.zig index b63e0a7c55..ec0861fe6e 100644 --- a/zig/src/families/glm/weights_fix.zig +++ b/zig/src/families/glm/weights_fix.zig @@ -1,6 +1,7 @@ -//! GLM-5.3-Flash weights rewritten as they load: packed mixes and router, conv taps, fp32 norms, column slices. +//! GLM-5.3-Flash weights as they load: parallel reads into place, and rewrites (mixes, router, conv, norms, columns). const std = @import("std"); const st = @import("../../core/safetensors.zig"); +const paced_read = @import("../../core/paced_read.zig"); /// One tensor's bytes on disk. pub const Src = struct { shard: u32, off: u64, len: u64, dtype: st.DType, shape: [st.max_rank]usize, rank: u8 }; @@ -8,6 +9,9 @@ pub const Fix = enum { hc_pack, router_pack, conv, to_f32, transpose_bf16, cols /// A tensor read whole into host memory and rewritten into its place (small: mixes, router, conv taps, norms). pub const Transform = struct { src: Src, dst: [*]u8, fix: Fix, extra: [3]?Src = .{ null, null, null }, arg: [3]usize = .{ 0, 0, 0 } }; +/// One tensor's bytes read straight into their place in a weights buffer. +pub const Copy = struct { shard: u32, off: u64, len: u64, dst: [*]u8 }; + pub fn readAll(fd: std.c.fd_t, dest: []u8, at: u64) !void { var done: usize = 0; while (done < dest.len) { @@ -21,7 +25,12 @@ pub fn transform(gpa: std.mem.Allocator, fds: []std.c.fd_t, t: Transform) !void const raw = try gpa.alloc(u8, t.src.len); defer gpa.free(raw); try readAll(fds[t.src.shard], raw, t.src.off); - const in16: []const u16 = @alignCast(std.mem.bytesAsSlice(u16, raw)); + const narrow = t.src.dtype == .f32 and (t.fix == .hc_pack or t.fix == .router_pack); // mlx-lm's fp32 mixes, router + if (narrow) bf16Exact(raw) catch |err| { + @import("quant.zig").say("glm: an fp32 {t} tensor (shard {d}, at byte {d}, shape {any}) is not exact in bf16", .{ t.fix, t.src.shard, t.src.off, t.src.shape[0..t.src.rank] }); + return err; + }; + const in16: []const u16 = @alignCast(std.mem.bytesAsSlice(u16, raw[0..if (narrow) raw.len / 2 else raw.len])); switch (t.fix) { .hc_pack => { // [og 6][tm 4][i 16][sgn 8][lane 32][tn 4] -> [og][sgn][lane][i][tm][tn] const out: [*]u16 = @ptrCast(@alignCast(t.dst)); @@ -79,3 +88,85 @@ pub fn transform(gpa: std.mem.Allocator, fds: []std.c.fd_t, t: Transform) !void }, } } + +/// fp32 values as bf16 in place (the first half of `raw`), refused unless each is a bf16 (0.6.6's bf16_if_exact). +fn bf16Exact(raw: []u8) error{InexactBf16}!void { + const n = raw.len / 4; + for (0..n) |i| { + const v = std.mem.readInt(u32, raw[4 * i ..][0..4], .little); + if (v & 0xffff != 0) return error.InexactBf16; + std.mem.writeInt(u16, raw[2 * i ..][0..2], @truncate(v >> 16), .little); + } +} + +test "fp32 values that are bf16 narrow in place; any other value is refused" { + var raw: [12]u8 = undefined; + for ([_]f32{ 1.5, -0.0078125, 3.0 }, 0..) |v, i| std.mem.writeInt(u32, raw[4 * i ..][0..4], @bitCast(v), .little); + try bf16Exact(&raw); + try std.testing.expectEqual(@as(u16, 0x3fc0), std.mem.readInt(u16, raw[0..2], .little)); + try std.testing.expectEqual(@as(u16, 0x4040), std.mem.readInt(u16, raw[4..6], .little)); + std.mem.writeInt(u32, raw[0..4], @bitCast(@as(f32, 1.0000001)), .little); + try std.testing.expectError(error.InexactBf16, bf16Exact(&raw)); +} + +const Pool = struct { + fds: []std.c.fd_t, + jobs: []const Copy, + next: std.atomic.Value(usize) = .init(0), + failed: std.atomic.Value(bool) = .init(false), + + fn run(p: *Pool) void { + const buf = paced_read.stage(std.heap.page_allocator, 16 << 20) catch { + p.failed.store(true, .release); + return; + }; + defer std.heap.page_allocator.free(buf); + while (true) { + const i = p.next.fetchAdd(1, .monotonic); + if (i >= p.jobs.len) return; + const j = p.jobs[i]; + paced_read.read(p.fds[j.shard], j.dst[0..j.len], j.off, buf) catch p.failed.store(true, .release); + } + } +}; + +/// Every copy on `threads` threads, uncached, paced by free memory (a large load otherwise fills the page cache too). +pub fn runCopies(gpa: std.mem.Allocator, shards: []const [:0]const u8, copies: []const Copy, transforms: []const Transform, threads: usize) !void { + const fds = try gpa.alloc(std.c.fd_t, shards.len); + defer gpa.free(fds); + for (shards, fds) |path, *fd| { + fd.* = std.c.open(path, .{ .ACCMODE = .RDONLY }); + if (fd.* < 0) return error.OpenFailed; + _ = std.c.fcntl(fd.*, 48, @as(c_int, 1)); // F_NOCACHE + } + defer for (fds) |fd| { + _ = std.c.close(fd); + }; + var pool: Pool = .{ .fds = fds, .jobs = copies }; + const workers = try gpa.alloc(?std.Thread, threads); + defer gpa.free(workers); + for (workers) |*t| t.* = std.Thread.spawn(.{}, Pool.run, .{&pool}) catch null; + pool.run(); + for (workers) |t| if (t) |th| th.join(); + if (pool.failed.load(.acquire)) return error.ShortRead; + // the transforms on the same threads (each reads its own tensors whole) + const Tx = struct { + gpa: std.mem.Allocator, + jobs: []const Transform, + fds: []std.c.fd_t, + next: std.atomic.Value(usize) = .init(0), + failed: std.atomic.Value(bool) = .init(false), + fn run(x: *@This()) void { + while (true) { + const i = x.next.fetchAdd(1, .monotonic); + if (i >= x.jobs.len) return; + transform(x.gpa, x.fds, x.jobs[i]) catch x.failed.store(true, .release); + } + } + }; + var tx: Tx = .{ .gpa = gpa, .jobs = transforms, .fds = fds }; + for (workers) |*t| t.* = std.Thread.spawn(.{}, Tx.run, .{&tx}) catch null; + tx.run(); + for (workers) |t| if (t) |th| th.join(); + if (tx.failed.load(.acquire)) return error.TransformFailed; +} diff --git a/zig/src/families/glm/weights_test.zig b/zig/src/families/glm/weights_test.zig new file mode 100644 index 0000000000..87e0c7fd01 --- /dev/null +++ b/zig/src/families/glm/weights_test.zig @@ -0,0 +1,468 @@ +//! GLM-5.3-Flash's loader on synthetic checkpoints: one sparse shard whose header names every tensor at its width. +const std = @import("std"); +const cfg = @import("config.zig"); +const wts = @import("weights.zig"); +const quant = @import("quant.zig"); + +const testing = std.testing; +const Width = quant.Width; + +const Entry = struct { name: []const u8, dtype: []const u8, shape: []const usize }; + +/// A module stored at a width other than 4/64. +const Pack = struct { []const u8, Width }; + +/// The tensors the loader reads for `c` (its first `c.run` layers and the MTP layer), stored at `stored`'s widths. +const Fixture = struct { + a: std.mem.Allocator, + c: *const cfg.Config, + stored: []const Pack = &.{}, // module -> its width on disk (else 4/64) + list: std.ArrayList(Entry) = .empty, + mlx: bool = false, // mlx-lm's layout: its names, stacked experts, the absorbed pair, bf16 everywhere unquantized + fused: bool = false, // with mlx: one conv1d tensor for q, k and v, and fp32 kept (mixes, router, biases) + + fn plain(f: *Fixture, comptime fmt: []const u8, args: anytype, dtype: []const u8, shape: []const usize) !void { + const t = if (f.mlx and !f.fused and std.mem.eql(u8, dtype, "F32")) "BF16" else dtype; + try f.list.append(f.a, .{ .name = try std.fmt.allocPrint(f.a, fmt, args), .dtype = t, .shape = try f.a.dupe(usize, shape) }); + } + + /// A stack of `lead` matrices [n, k] (mlx-lm's experts and per-head maps), at its module's width. + fn q3(f: *Fixture, comptime fmt: []const u8, args: anytype, lead: usize, n: usize, k: usize) !void { + const module = try std.fmt.allocPrint(f.a, fmt, args); + var w = quant.uniform; + for (f.stored) |p| if (std.mem.eql(u8, p[0], module)) { + w = p[1]; + }; + try f.plain("{s}.weight", .{module}, "U32", &.{ lead, n, k * w.bits / 32 }); + try f.plain("{s}.scales", .{module}, "BF16", &.{ lead, n, k / w.group }); + try f.plain("{s}.biases", .{module}, "BF16", &.{ lead, n, k / w.group }); + } + + fn q(f: *Fixture, comptime fmt: []const u8, args: anytype, n: usize, k: usize) !void { + const module = try std.fmt.allocPrint(f.a, fmt, args); + var w = quant.uniform; + for (f.stored) |p| if (std.mem.eql(u8, p[0], module)) { + w = p[1]; + }; + try f.plain("{s}.weight", .{module}, "U32", &.{ n, k * w.bits / 32 }); + try f.plain("{s}.scales", .{module}, "BF16", &.{ n, k / w.group }); + try f.plain("{s}.biases", .{module}, "BF16", &.{ n, k / w.group }); + } + + fn layer(f: *Fixture, i: usize) !void { + const c = f.c; + const D: usize = c.hidden; + const p = "model.language_model.layers.{d}."; + const a = p ++ "self_attn."; + if (i != c.layers) for ([_][]const u8{ "attn", "ffn" }) |site| { + try f.plain(p ++ "hc_{s}_fn", .{ i, site }, if (f.fused) "F32" else "BF16", &.{ 24, 4 * D }); + try f.plain(p ++ "hc_{s}_base", .{ i, site }, "F32", &.{24}); + try f.plain(p ++ "hc_{s}_scale", .{ i, site }, "F32", &.{3}); + }; + try f.plain(p ++ "input_layernorm.weight", .{i}, "BF16", &.{D}); + try f.plain(p ++ "post_attention_layernorm.weight", .{i}, "BF16", &.{D}); + switch (c.kind(@intCast(i))) { + .kda => { + const W: usize = c.kdaWidth(); + const dk: usize = c.kda_dim; + if (f.fused) try f.plain(a ++ "conv1d.weight", .{i}, "BF16", &.{ 3 * W, c.conv, 1 }) else for ([_][]const u8{ "q", "k", "v" }) |x| + try f.plain(a ++ "{s}_conv1d.weight", .{ i, x }, "BF16", &.{ W, 1, c.conv }); + try f.plain(a ++ "o_norm.weight", .{i}, "BF16", &.{dk}); + for ([_][]const u8{ "q_proj", "k_proj", "v_proj" }) |x| try f.q(a ++ "{s}", .{ i, x }, W, D); + try f.q(a ++ "f_a_proj", .{i}, dk, D); + try f.q(a ++ "g_a_proj", .{i}, dk, D); + try f.q(a ++ "b_proj", .{i}, c.kda_heads, D); + try f.q(a ++ "f_b_proj", .{i}, W, dk); + try f.q(a ++ "g_b_proj", .{i}, W, dk); + try f.q(a ++ "o_proj", .{i}, D, W); + try f.plain(a ++ "A_log", .{i}, "F32", &.{c.kda_heads}); + try f.plain(a ++ "dt_bias", .{i}, "F32", &.{W}); + }, + .mla => { + const H: usize = c.mla_heads; + try f.plain(a ++ "indexer.index_kpool_compress_gate", .{i}, "BF16", &.{ c.i_dim, D }); + try f.q(a ++ "q_a_proj", .{i}, c.q_lora, D); + try f.q(a ++ "kv_a_proj_with_mqa", .{i}, c.kv_lora, D); + try f.q(a ++ "indexer.wk", .{i}, c.i_dim, D); + try f.q(a ++ "indexer.weights_proj", .{i}, c.i_heads, D); + try f.q(a ++ "q_b_proj", .{i}, H * c.nope, c.q_lora); + try f.q(a ++ "indexer.wq_b", .{i}, c.i_heads * c.i_dim, c.q_lora); + if (f.mlx) { + try f.q3(a ++ "embed_q", .{i}, H, c.kv_lora, c.nope); + try f.q3(a ++ "unembed_out", .{i}, H, c.v_dim, c.kv_lora); + } else try f.q(a ++ "kv_b_proj", .{i}, H * (c.nope + c.v_dim), c.kv_lora); + try f.q(a ++ "o_proj", .{i}, D, H * c.v_dim); + try f.plain(a ++ "q_a_layernorm.weight", .{i}, "BF16", &.{c.q_lora}); + try f.plain(a ++ "kv_a_layernorm.weight", .{i}, "BF16", &.{c.kv_lora}); + try f.plain(a ++ "indexer.k_norm.weight", .{i}, "BF16", &.{c.i_dim}); + try f.plain(a ++ "indexer.k_norm.bias", .{i}, "BF16", &.{c.i_dim}); + try f.plain(a ++ "indexer.index_kpool_compress_ape", .{i}, "BF16", &.{ c.kpool, c.i_dim }); + }, + } + if (i == c.layers) { + if (f.mlx) try f.plain(p ++ "eh_proj.weight", .{i}, "BF16", &.{ D, 2 * D }) else try f.q(p ++ "eh_proj", .{i}, D, 2 * D); + for ([_][]const u8{ "enorm", "hnorm", "shared_head.norm" }) |x| try f.plain(p ++ "{s}.weight", .{ i, x }, "BF16", &.{D}); + } + const m = p ++ "mlp."; + if (c.isMoe(@intCast(i))) { + try f.plain(m ++ "gate.weight", .{i}, if (f.fused) "F32" else "BF16", &.{ c.experts, D }); + try f.plain(m ++ "gate.e_score_correction_bias", .{i}, "F32", &.{c.experts}); + try f.q(m ++ "shared_experts.gate_proj", .{i}, c.moe_inter, D); + try f.q(m ++ "shared_experts.up_proj", .{i}, c.moe_inter, D); + try f.q(m ++ "shared_experts.down_proj", .{i}, D, c.moe_inter); + if (f.mlx) for ([_][]const u8{ "gate_proj", "up_proj", "down_proj" }) |x| { + const down = x[0] == 'd'; + try f.q3(m ++ "switch_mlp.{s}", .{ i, x }, c.experts, if (down) D else c.moe_inter, if (down) c.moe_inter else D); + }; + if (!f.mlx) for (0..c.experts) |e| { + try f.q(m ++ "experts.{d}.gate_proj", .{ i, e }, c.moe_inter, D); + try f.q(m ++ "experts.{d}.up_proj", .{ i, e }, c.moe_inter, D); + try f.q(m ++ "experts.{d}.down_proj", .{ i, e }, D, c.moe_inter); + }; + } else { + try f.q(m ++ "gate_proj", .{i}, c.dense_inter, D); + try f.q(m ++ "up_proj", .{i}, c.dense_inter, D); + try f.q(m ++ "down_proj", .{i}, D, c.dense_inter); + } + } + + /// config.json, the index and one shard (header only, its data a hole) in `dir`. + fn write(f: *Fixture, tmp: std.testing.TmpDir, dir: []const u8, config: []const u8) !void { + const c = f.c; + try f.q("model.language_model.embed_tokens", .{}, c.vocab, c.hidden); + try f.q("lm_head", .{}, c.vocab, c.hidden); + try f.plain("model.language_model.norm.weight", .{}, "BF16", &.{c.hidden}); + for (0..c.run) |i| try f.layer(i); + try f.layer(c.layers); + var head: std.ArrayList(u8) = .empty; + var index: std.ArrayList(u8) = .empty; + try index.appendSlice(f.a, "{\"weight_map\": {"); + var at: usize = 0; + for (f.list.items, 0..) |*e, j| { + if (f.mlx) e.name = try mlxName(f.a, e.name); + var n: usize = if (std.mem.eql(u8, e.dtype, "BF16")) 2 else 4; + for (e.shape) |d| n *= d; + const sep = if (j == 0) "{" else ","; + try head.appendSlice(f.a, try std.fmt.allocPrint(f.a, "{s}\"{s}\":{{\"dtype\":\"{s}\",\"shape\":[{s}],\"data_offsets\":[{d},{d}]}}", .{ sep, e.name, e.dtype, try dims(f.a, e.shape), at, at + n })); + try index.appendSlice(f.a, try std.fmt.allocPrint(f.a, "{s}\"{s}\": \"model-00001.safetensors\"", .{ if (j == 0) "" else ",", e.name })); + at += n; + } + try head.appendSlice(f.a, "}"); + try index.appendSlice(f.a, "}}"); + const io = testing.io; + try tmp.dir.writeFile(io, .{ .sub_path = "config.json", .data = config }); + try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = index.items }); + var len: [8]u8 = undefined; + std.mem.writeInt(u64, &len, head.items.len, .little); + try tmp.dir.writeFile(io, .{ .sub_path = "model-00001.safetensors", .data = try std.mem.concat(f.a, u8, &.{ &len, head.items }) }); + const path = try std.fmt.allocPrintSentinel(f.a, "{s}/model-00001.safetensors", .{dir}, 0); + const fd = std.c.open(path, .{ .ACCMODE = .WRONLY }, @as(std.c.mode_t, 0)); + if (fd < 0) return error.TestFile; + defer _ = std.c.close(fd); + if (std.c.ftruncate(fd, @intCast(8 + head.items.len + at)) != 0) return error.TestFile; + } +}; + +/// An original-layout name as mlx-lm stores it. +fn mlxName(a: std.mem.Allocator, name: []const u8) ![]const u8 { + if (std.mem.startsWith(u8, name, "lm_head.")) return std.mem.concat(a, u8, &.{ "language_model.", name }); + const short = name["model.language_model.".len..]; + const mtp = "layers.45."; + if (std.mem.startsWith(u8, short, mtp)) { + const rest = short[mtp.len..]; + if (std.mem.eql(u8, rest, "shared_head.norm.weight")) return "language_model.mtp.0.norm.weight"; + const top = std.mem.startsWith(u8, rest, "eh_proj.") or std.mem.startsWith(u8, rest, "enorm.") or std.mem.startsWith(u8, rest, "hnorm."); + return std.mem.concat(a, u8, &.{ "language_model.mtp.0.", if (top) "" else "block.", rest }); + } + inline for (.{ "attn", "ffn" }) |site| inline for (.{ "fn", "base", "scale" }) |part| { + const old = ".hc_" ++ site ++ "_" ++ part; + if (std.mem.endsWith(u8, short, old)) return std.mem.concat(a, u8, &.{ "language_model.model.", short[0 .. short.len - old.len], "." ++ site ++ "_hc." ++ part }); + }; + return std.mem.concat(a, u8, &.{ "language_model.model.", short }); +} + +fn dims(a: std.mem.Allocator, shape: []const usize) ![]const u8 { + var out: std.ArrayList(u8) = .empty; + for (shape, 0..) |d, j| try out.appendSlice(a, try std.fmt.allocPrint(a, "{s}{d}", .{ if (j == 0) "" else ",", d })); + return out.items; +} + +/// A small GLM: four layers (three KDA and dense, one MLA and MoE) and the MTP layer, four routed experts. +fn small() cfg.Config { + var c: cfg.Config = .{}; + c.run = 4; + c.experts = 4; + c.own = .{ 0, 4 }; + return c; +} + +/// Every matrix's rows, columns and buffer offsets, and every other tensor's offset, hashed in load order. +fn digest(h: *std.hash.Wyhash, v: anytype) void { + const T = @TypeOf(v); + switch (@typeInfo(T)) { + .@"struct" => |s| { + if (@hasField(T, "buf") and @hasField(T, "off")) return h.update(std.mem.asBytes(&v.off)); + if (@hasField(T, "w") and @hasField(T, "n")) { + h.update(std.mem.asBytes(&v.n)); + h.update(std.mem.asBytes(&v.k)); + inline for (.{ v.w, v.s, v.b }) |r| digest(h, r); + return; + } + inline for (s.field_names) |name| digest(h, @field(v, name)); + }, + .@"union" => switch (v) { + inline else => |x| digest(h, x), + }, + .optional => if (v) |x| digest(h, x), + .array => for (v) |x| digest(h, x), + else => {}, + } +} + +fn plan(w: *const wts.Weights, c: *const cfg.Config) u64 { + var h = std.hash.Wyhash.init(0); + inline for (.{ w.embed, w.head, w.norm, w.mtp }) |x| digest(&h, x); + for (w.layers[0..c.run]) |l| digest(&h, l); + digest(&h, w.layers[c.layers]); + h.update(std.mem.asBytes(&w.bytes)); + return h.final(); +} + +const flat = "{\"quantization\": {\"group_size\": 64, \"bits\": 4, \"mode\": \"affine\"}}"; + +/// A dry load of a fixture stored at `stored` under `config`: the loader's plan, or its refusal. +fn dry(a: std.mem.Allocator, config: []const u8, stored: []const Pack) !*wts.Weights { + return dryAs(a, config, stored, false); +} + +fn dryAs(a: std.mem.Allocator, config: []const u8, stored: []const Pack, mlx: bool) !*wts.Weights { + return dryWith(a, config, stored, mlx, false); +} + +fn dryWith(a: std.mem.Allocator, config: []const u8, stored: []const Pack, mlx: bool, fused: bool) !*wts.Weights { + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + const dir = try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}", .{tmp.sub_path}); + const c = small(); + var f: Fixture = .{ .a = a, .c = &c, .stored = stored, .mlx = mlx, .fused = fused }; + try f.write(tmp, dir, config); + return wts.load(testing.allocator, undefined, dir, &c, 1, true); +} + +fn free(w: *wts.Weights) void { + w.deinit(); + testing.allocator.destroy(w); +} + +test "a uniform 4/64 checkpoint plans the same matrices, offsets and bytes as the 4-bit-only loader" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const c = small(); + const w = try dry(arena.allocator(), flat, &.{}); + defer free(w); + try testing.expectEqual(@as(usize, 1_512_774_688), w.bytes); + try testing.expectEqual(@as(u64, 0xc1e86c075603176a), plan(w, &c)); + try testing.expectEqual(@as(u8, 0), w.layers[0].attn.kda.in_proj.more); // one width: one run, the matrix itself +} + +/// An oQ-style table: o_proj of layer 0 at 8 bits under the short prefix, layer 3's shared up at 5 under the long. +const oq = + \\{"quantization": {"group_size": 64, "bits": 4, "mode": "affine", + \\ "model.layers.0.self_attn.o_proj": {"group_size": 64, "bits": 8}, + \\ "language_model.model.layers.3.mlp.shared_experts.down_proj": {"group_size": 32, "bits": 5}, + \\ "language_model.mtp.0.block.mlp.switch_mlp.gate_proj": {"group_size": 64, "bits": 2}}} +; +const oq_packed = [_]Pack{ + .{ "model.language_model.layers.0.self_attn.o_proj", .{ .bits = 8, .group = 64 } }, + .{ "model.language_model.layers.3.mlp.shared_experts.down_proj", .{ .bits = 5, .group = 32 } }, +} ++ blk: { + var out: [4]Pack = undefined; + for (&out, 0..) |*o, e| o.* = .{ std.fmt.comptimePrint("model.language_model.layers.45.mlp.experts.{d}.gate_proj", .{e}), .{ .bits = 2, .group = 64 } }; + break :blk out; +}; + +test "the device admits each tensor at its override's width, whatever the key's prefix; a gate and up at two are refused" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + try testing.expectError(error.MixedWidths, dry(a, oq, &oq_packed)); // the MTP layer's gate at 2, its up at 4 + const both = oq[0 .. oq.len - 2] ++ ",\n \"language_model.mtp.0.block.mlp.switch_mlp.up_proj\": {\"group_size\": 64, \"bits\": 2}}}"; + const ups = comptime oq_packed ++ blk: { + var out: [4]Pack = undefined; + for (&out, 0..) |*o, e| o.* = .{ std.fmt.comptimePrint("model.language_model.layers.45.mlp.experts.{d}.up_proj", .{e}), .{ .bits = 2, .group = 64 } }; + break :blk out; + }; + const c = small(); + const w = try dry(a, both, &ups); + defer free(w); + try testing.expectEqual(Width{ .bits = 8, .group = 64 }, w.layers[0].attn.kda.o_proj.width()); + try testing.expectEqual(Width{ .bits = 5, .group = 32 }, w.layers[3].mlp.moe.sh_down.width()); + try testing.expectEqual(Width{ .bits = 2, .group = 64 }, w.layers[c.layers].mlp.moe.gate.width()); + try testing.expectEqual(Width{ .bits = 2, .group = 64 }, w.layers[c.layers].mlp.moe.up.width()); + try testing.expectEqual(quant.uniform, w.layers[1].attn.kda.o_proj.width()); + try testing.expectEqual(quant.uniform, w.layers[3].mlp.moe.gate.width()); + try testing.expect(!wts.native(w, &c)); +} + +test "a tensor stored at another width than the config gives it is refused, never read at a guessed width" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const eight = [_]Pack{.{ "model.language_model.layers.0.self_attn.o_proj", .{ .bits = 8, .group = 64 } }}; + try testing.expectError(error.UnexpectedTensor, dry(a, flat, &eight)); // global 4/64, stored at 8 + try testing.expectError(error.UnexpectedTensor, dry(a, oq, &.{})); // override 8, stored at 4 +} + +/// oQ4's split stacks on this fixture: the KDA input (q, v at 5; k, f_a, g_a, b at 8), the MLA x_proj (6 and 5). +const oq_split = + \\{"quantization": {"group_size": 64, "bits": 4, "mode": "affine", + \\ "model.layers.0.self_attn.q_proj": {"group_size": 64, "bits": 5}, + \\ "model.layers.0.self_attn.k_proj": {"group_size": 64, "bits": 8}, + \\ "model.layers.0.self_attn.v_proj": {"group_size": 64, "bits": 5}, + \\ "model.layers.0.self_attn.f_a_proj": {"group_size": 64, "bits": 8}, + \\ "model.layers.0.self_attn.g_a_proj": {"group_size": 64, "bits": 8}, + \\ "model.layers.0.self_attn.b_proj": {"group_size": 64, "bits": 8}, + \\ "model.layers.3.self_attn.q_a_proj": {"group_size": 64, "bits": 6}, + \\ "model.layers.3.self_attn.kv_a_proj_with_mqa": {"group_size": 64, "bits": 6}, + \\ "model.layers.3.self_attn.indexer.wk": {"group_size": 64, "bits": 5}, + \\ "model.layers.3.self_attn.indexer.weights_proj": {"group_size": 64, "bits": 5}}} +; +const oq_split_stored = blk: { + const bits = [_]struct { []const u8, u8 }{ + .{ "0.self_attn.q_proj", 5 }, + .{ "0.self_attn.k_proj", 8 }, + .{ "0.self_attn.v_proj", 5 }, + .{ "0.self_attn.f_a_proj", 8 }, + .{ "0.self_attn.g_a_proj", 8 }, + .{ "0.self_attn.b_proj", 8 }, + .{ "3.self_attn.q_a_proj", 6 }, + .{ "3.self_attn.kv_a_proj_with_mqa", 6 }, + .{ "3.self_attn.indexer.wk", 5 }, + .{ "3.self_attn.indexer.weights_proj", 5 }, + }; + var out: [bits.len]Pack = undefined; + for (&out, bits) |*o, b| o.* = .{ "model.language_model.layers." ++ b[0], .{ .bits = b[1], .group = 64 } }; + break :blk out; +}; + +/// Run `i` of `q` is rows [row, row + n) at `bits`, its codes and tables its own, `pad`-row tiles apart. +fn expectRun(q: wts.Affine, i: usize, row: u32, n: u32, bits: u8, pad: usize) !void { + const r = q.run(i); + try testing.expectEqual(row, r.row); + try testing.expectEqual(n, r.n); + try testing.expectEqual(Width{ .bits = bits, .group = 64 }, Width{ .bits = r.bits, .group = r.group }); + const w: Width = .{ .bits = r.bits, .group = r.group }; + const room = std.mem.alignForward(usize, n, pad); + try testing.expectEqual(std.mem.alignForward(usize, room * q.k * bits / 8, 256), r.s.off - r.w.off); + try testing.expectEqual(std.mem.alignForward(usize, room * (q.k / w.group) * 2, 256), r.b.off - r.s.off); +} + +test "a stack of mixed widths plans one run a member, as 0.6.6's QSplit calls each part, at its rows of the output" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const c = small(); + const w = try dry(arena.allocator(), oq_split, &oq_split_stored); + defer free(w); + const in = w.layers[0].attn.kda.in_proj; + const W = c.kdaWidth(); + const dk = c.kda_dim; + try testing.expectEqual(@as(u8, 5), in.more); // f_a, g_a and b share 8 bits yet stay apart + try testing.expectEqual(c.kdaProj(), in.n); + try expectRun(in, 0, 0, W, 5, 64); + try expectRun(in, 1, W, W, 8, 64); + try expectRun(in, 2, 2 * W, W, 5, 64); + try expectRun(in, 3, 3 * W, dk, 8, 64); + try expectRun(in, 4, 3 * W + dk, dk, 8, 64); + try expectRun(in, 5, 3 * W + 2 * dk, c.kda_heads, 8, 64); + const x = w.layers[3].attn.mla.x_proj; + try testing.expectEqual(@as(u8, 3), x.more); + try expectRun(x, 0, 0, c.q_lora, 6, 64); + try expectRun(x, 1, c.q_lora, c.kv_lora, 6, 64); + try expectRun(x, 2, c.q_lora + c.kv_lora, c.i_dim, 5, 64); + try expectRun(x, 3, c.q_lora + c.kv_lora + c.i_dim, c.i_heads, 5, 64); + try testing.expectEqual(@as(u8, 0), w.layers[1].attn.kda.in_proj.more); + try testing.expectEqual(@as(u8, 0), w.layers[3].attn.mla.qr_proj.more); + const t = wts.tally(w, &c); + try testing.expectEqual(@as(u32, 4), t.count[5][1]); // q, v and the indexer's pair + try testing.expectEqual(@as(u32, 2), t.count[6][1]); + try testing.expectEqual(@as(u32, 4), t.count[8][1]); + var buf: [256]u8 = undefined; + try testing.expect(std.mem.indexOf(u8, t.summary(&buf), "5-bit/64 x4, 6-bit/64 x2, 8-bit/64 x4") != null); +} + +/// oQ's spelling in mlx-lm's layout: the embedding at 8, a stack at 2, the absorbed pair at 5 and 6. +const oq_mlx = + \\{"quantization": {"group_size": 64, "bits": 4, "mode": "affine", + \\ "language_model.model.embed_tokens": {"group_size": 64, "bits": 8}, + \\ "language_model.model.layers.3.mlp.switch_mlp.down_proj": {"group_size": 64, "bits": 2}, + \\ "language_model.model.layers.3.self_attn.embed_q": {"group_size": 64, "bits": 5}, + \\ "language_model.mtp.0.block.self_attn.unembed_out": {"group_size": 64, "bits": 6}, + \\ "language_model.mtp.0.block.mlp.switch_mlp.gate_proj": {"group_size": 64, "bits": 2}}} +; +const oq_mlx_stored = [_]Pack{ + .{ "model.language_model.embed_tokens", .{ .bits = 8, .group = 64 } }, + .{ "model.language_model.layers.3.mlp.switch_mlp.down_proj", .{ .bits = 2, .group = 64 } }, + .{ "model.language_model.layers.3.self_attn.embed_q", .{ .bits = 5, .group = 64 } }, + .{ "model.language_model.layers.45.self_attn.unembed_out", .{ .bits = 6, .group = 64 } }, + .{ "model.language_model.layers.45.mlp.switch_mlp.gate_proj", .{ .bits = 2, .group = 64 } }, +}; + +test "an mlx-lm checkpoint's every tensor is found, each at its own width" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const c = small(); + const both = oq_mlx[0 .. oq_mlx.len - 2] ++ ",\n \"language_model.mtp.0.block.mlp.switch_mlp.up_proj\": {\"group_size\": 64, \"bits\": 2}}}"; + const up: Pack = .{ "model.language_model.layers.45.mlp.switch_mlp.up_proj", .{ .bits = 2, .group = 64 } }; + const w = try dryAs(arena.allocator(), both, &(oq_mlx_stored ++ [_]Pack{up}), true); + defer free(w); + try testing.expectEqual(Width{ .bits = 8, .group = 64 }, w.embed.width()); + try testing.expectEqual(quant.uniform, w.head.width()); + try testing.expectEqual(quant.uniform, w.layers[0].attn.kda.o_proj.width()); + const moe = w.layers[3].mlp.moe; + try testing.expectEqual(Width{ .bits = 2, .group = 64 }, moe.down.width()); + try testing.expectEqual(quant.uniform, moe.gate.width()); + const mla = w.layers[3].attn.mla; + try testing.expectEqual(@as(u32, 0), mla.kv_b.n); + try testing.expectEqual(Width{ .bits = 5, .group = 64 }, mla.absorbed.?[0].width()); + try testing.expectEqual(@as(u32, c.mla_heads * c.kv_lora), mla.absorbed.?[0].n); + try testing.expectEqual(@as(u32, c.nope), mla.absorbed.?[0].k); + try testing.expectEqual(quant.uniform, mla.absorbed.?[1].width()); + try testing.expectEqual(Width{ .bits = 6, .group = 64 }, w.layers[c.layers].attn.mla.absorbed.?[1].width()); + try testing.expectEqual(Width{ .bits = 2, .group = 64 }, w.layers[c.layers].mlp.moe.gate.width()); + try testing.expectEqual(Width{ .bits = 2, .group = 64 }, w.layers[c.layers].mlp.moe.up.width()); + try testing.expect(w.mtp.?.eh_bf16 != null); +} + +test "the device loads mlx-lm's absorbed pair and bf16 eh_proj; a routed gate and up at two widths are refused" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + try testing.expectError(error.MixedWidths, dryAs(a, oq_mlx, &oq_mlx_stored, true)); // the MTP gate at 2, up at 4 + const c = small(); + const w = try dryAs(a, flat, &.{}, true); + defer free(w); + try testing.expectEqual(quant.uniform, w.layers[3].attn.mla.absorbed.?[0].width()); + try testing.expect(w.mtp.?.eh_bf16 != null); + try testing.expect(!wts.native(w, &c)); // the pair and the bf16 eh_proj take the row kernels +} + +test "a uniform checkpoint in the original layout needs no other width's pipelines" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const c = small(); + const w = try dry(arena.allocator(), flat, &.{}); + defer free(w); + try testing.expect(wts.native(w, &c)); +} + +test "mlx-lm's one conv1d tensor and fp32 mixes and router plan as the three conv tensors and bf16 ones do" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const c = small(); + const w = try dryWith(a, flat, &.{}, true, true); + defer free(w); + const v = try dryWith(a, flat, &.{}, true, false); + defer free(v); + try testing.expectEqual(plan(v, &c), plan(w, &c)); // the same buffers, offsets and bytes +} diff --git a/zig/src/families/glm/wide.zig b/zig/src/families/glm/wide.zig new file mode 100644 index 0000000000..4dbf759211 --- /dev/null +++ b/zig/src/families/glm/wide.zig @@ -0,0 +1,420 @@ +//! GLM-5.3-Flash's products at widths its generated kernels do not read (glm_rows.metal): a launch a stack member. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const affine = @import("affine.zig"); +const affine_mm = @import("../../core/affine_mm.zig"); +const kernels = @import("kernels.zig"); +const Affine = affine.Affine; +const Ref = affine.Ref; + +const bit_widths = affine.row_bits; +const nw = bit_widths.len; + +/// Whether `q` takes the generated kernels: one width, 4 bits in groups of 64 (the uniform checkpoints' launches). +pub fn native(q: Affine) bool { + return q.more == 0 and q.bits == 4 and q.group == 64; +} + +pub const RowDims = extern struct { k: i32, n: i32, group: i32, rows: i32, heads: i32 = 1, ldx: i32, hsx: i32 = 0, ldo: i32, hso: i32 = 0, hsw: i32 = 0 }; +pub const PickDims = extern struct { k: i32, n: i32, group: i32, slots: i32, repeat: i32, experts: i32, ldo: i32, lim: f32 }; + +pub const Pipes = struct { + rows: [nw]mtl.Pipeline, // dense, strided, by head + ids: [nw]mtl.Pipeline, // a slot through its expert + glu_ids: [nw]mtl.Pipeline, // gate and up through it, SwiGLU'd + sorted: [nw]mtl.Pipeline, // an expert's sorted pairs (prompt chunks) + glu_sorted: [nw]mtl.Pipeline, + moe: [nw]mtl.Pipeline, // routed experts by pick at each width and the 8-bit shared expert, one launch + moe_glu: [nw]mtl.Pipeline, // the same through gate and up, SwiGLU'd + bf16: mtl.Pipeline, // an unquantized matrix (the MTP head's eh_proj) + kda_pre: mtl.Pipeline, // kda_rows reading f_b's and g_b's products made apart (any width but 4/64) + kda_fg: [nw]mtl.Pipeline, // kda_rows making f_b's and g_b's dots of one width (g64) as glm_rows does: .pre's bits + embed: [2]mtl.Pipeline, // the embedding at 6 and 8 bits, groups of 64 + mm8: affine_mm.Pipes, // prompt chunks' products of one width at 8 bits in groups of 64 (core/affine_mm.zig) + mm2: affine_mm.Pipes, // and at 2 bits in groups of 64: oQ2's routed experts, gathered + mm5: affine_mm.Pipes, // and at 5 and 6 bits (oQ4's), split stacks' members too + mm6: affine_mm.Pipes, + zeros: mtl.Buffer, // expert ids all 0: a dense GLU through the by-ids kernel (the shared expert) + hash: u64, // every source compiled here + + pub fn deinit(p: *Pipes) void { + inline for (.{ &p.rows, &p.ids, &p.glu_ids, &p.sorted, &p.glu_sorted, &p.moe, &p.moe_glu, &p.embed, &p.kda_fg }) |set| for (set) |x| x.deinit(); + p.bf16.deinit(); + p.kda_pre.deinit(); + p.zeros.deinit(); + inline for (.{ p.mm8, p.mm2, p.mm5, p.mm6 }) |set| for (set) |x| x.deinit(); + } +}; + +/// kda_rows's f_b dot read from a product made apart (a row's sum in its first lane of four, zeros in the rest: exact). +const pre_f = .{ "result = quad_dot(x, wb, s, bb);", "result = qlid == 0u ? float(((device const bfloat*)FBW)[(size_t)r * W + row]) : 0.0f;" }; +const pre_g = .{ "result = quad_dot(x, wb, s, bb);", "result = qlid == 0u ? float(((device const bfloat*)GBW)[(size_t)r * W + row]) : 0.0f;" }; + +/// kda_rows's f_b and g_b dots by quads of lanes, cut (from its block's start to the barrier after it) for `fg_block`. +const fg_cut = .{ " {\n constexpr int PER = D / 4;", "\n threadgroup_barrier(mem_flags::mem_threadgroup);\n for (uint idx = tid; idx < 3u * (uint)D;" }; +/// The head's f_b and g_b rows as glm_rows sums them (core tf_rp_rows, 4 rows a simdgroup, then bf16): .pre's values. +const fg_block = + \\ for (uint t0 = simdgroup_index_in_threadgroup * 4u; t0 < 2u * (uint)D; t0 += (uint)TY * 4u) { + \\ const bool f = t0 < (uint)D; + \\ const uint d0 = f ? t0 : t0 - (uint)D; + \\ const size_t row = (size_t)h * D + d0; + \\ float acc[1][4]; + \\ if (f) tf_rp_rows<4, 1, FB, false>((const device uint8_t*)FBW + row * (D * FB / 8), FBS + row * (D / 64), FBB + row * (D / 64), prow + FA, D, 64, lane, acc); + \\ else tf_rp_rows<4, 1, GB, false>((const device uint8_t*)GBW + row * (D * GB / 8), GBS + row * (D / 64), GBB + row * (D / 64), prow + GA, D, 64, lane, acc); + \\ if (lane == 0u) + \\ for (int j = 0; j < 4; j++) (f ? sa : sgate)[d0 + j] = float(bfloat(acc[0][j])); + \\ } +; +/// core/row_projection.metal's row sums only (its kernels and instances cut). +const rp_sums = blk: { + @setEvalBranchQuota(100_000); + const r = sources.core_row_projection; + break :blk r[0..std.mem.indexOf(u8, r, "// A value out: bf16").?]; +}; + +/// glm_glue.metal's bfloat overloads and sigmoid_fast, the text glm_swiglu compiles: glm_rows's SwiGLU takes its bits. +pub const glue_math = blk: { + @setEvalBranchQuota(100_000); + const g = sources.glm_glue; + const a = std.mem.indexOf(u8, g, "// MLX's bfloat math overloads").?; + const b = a + std.mem.indexOf(u8, g[a..], "// The dense MLP's SwiGLU").?; + break :blk "typedef bfloat bfloat16_t;\n" ++ g[a..b]; +}; + +/// The pipelines built so far: a failed compile releases them all. +const Made = struct { + list: [nw * 8 + 4 + 4 * affine_mm.names.len]mtl.Pipeline = undefined, + n: usize = 0, + released: usize = 0, + + fn add(m: *Made, device: mtl.Device, l: mtl.Library, name: []const u8) !mtl.Pipeline { + m.list[m.n] = try mtl.Pipeline.init(device, l, name, false); + m.n += 1; + return m.list[m.n - 1]; + } + + fn drop(m: *Made) void { + for (m.list[0..m.n]) |p| p.deinit(); + m.released = m.n; + m.n = 0; + } +}; + +fn lib(m: *Made, device: mtl.Device, src: []const u8, names: []const [:0]const u8, out: []mtl.Pipeline) !void { + const l = try mtl.Library.fromSource(device, src, mtl.CompileOptions.mlx()); + defer l.deinit(); + for (names, out) |n, *p| p.* = try m.add(device, l, n); +} + +pub fn load(gpa: std.mem.Allocator, device: mtl.Device) !Pipes { + var made: Made = .{}; + return build(gpa, device, comptime kernels.generatedSource("kda_rows"), &made); +} + +fn build(gpa: std.mem.Allocator, device: mtl.Device, kda: []const u8, made: *Made) !Pipes { + errdefer made.drop(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + var p: Pipes = undefined; + var h = std.hash.Wyhash.init(0x77); + const rows_src = try std.mem.concat(gpa, u8, &.{ sources.core_row_projection, glue_math, sources.glm_rows }); + defer gpa.free(rows_src); + h.update(rows_src); + { + const l = try mtl.Library.fromSource(device, rows_src, mtl.CompileOptions.mlx()); + defer l.deinit(); + inline for (bit_widths, 0..) |b, i| { + const q = std.fmt.comptimePrint("_q{d}", .{b}); + p.rows[i] = try made.add(device, l, "glm_rows" ++ q); + p.ids[i] = try made.add(device, l, "glm_rows_ids" ++ q); + p.glu_ids[i] = try made.add(device, l, "glm_rows_glu_ids" ++ q); + p.sorted[i] = try made.add(device, l, "glm_rows_sorted" ++ q); + p.glu_sorted[i] = try made.add(device, l, "glm_rows_glu_sorted" ++ q); + p.moe[i] = try made.add(device, l, "glm_moe_rows" ++ q ++ "_s8"); + p.moe_glu[i] = try made.add(device, l, "glm_moe_rows_glu" ++ q ++ "_s8"); + } + p.bf16 = try made.add(device, l, "glm_rows_bf16"); + } + const pre_a = try std.mem.replaceOwned(u8, gpa, kda, pre_f[0], pre_f[1]); // its own instance, under its own name + defer gpa.free(pre_a); + const pre_src = try std.mem.replaceOwned(u8, gpa, pre_a, pre_g[0], pre_g[1]); + defer gpa.free(pre_src); + if (std.mem.indexOf(u8, pre_src, pre_f[0]) != null or std.mem.indexOf(u8, pre_src, pre_g[0]) != null or std.mem.indexOf(u8, pre_src, pre_f[1]) == null) return error.KernelCompile; + h.update(pre_src); + try lib(made, device, pre_src, &.{comptime kernels.generatedFunction("kda_rows")}, @as(*[1]mtl.Pipeline, &p.kda_pre)); + const a = std.mem.indexOf(u8, kda, fg_cut[0]) orelse return error.KernelCompile; + const z = a + (std.mem.indexOf(u8, kda[a..], fg_cut[1]) orelse return error.KernelCompile); + const own = z + (std.mem.lastIndexOf(u8, kda[z..], "\ntemplate [[host_name(") orelse return error.KernelCompile); + const fg_src = try std.mem.concat(gpa, u8, &.{ rp_sums, kda[0..a], fg_block, kda[z..own], fg_instances }); // not its own + defer gpa.free(fg_src); + h.update(fg_src); + try lib(made, device, fg_src, &fg_names, &p.kda_fg); + h.update(sources.ops_embed_norm); + try lib(made, device, sources.ops_embed_norm, &.{ "tf_embed_b6_g64", "tf_embed_b8_g64" }, &p.embed); + for ([_]u8{ 8, 2, 5, 6 }, [_]*affine_mm.Pipes{ &p.mm8, &p.mm2, &p.mm5, &p.mm6 }) |bits, out| { + const src = try affine_mm.source(device, gpa, .{ .bits = bits, .group = 64 }); + defer gpa.free(src); + h.update(src); + try lib(made, device, src, &affine_mm.names, out); + } + p.zeros = try device.buffer(64 * 4, mtl.ResourceOptions.shared); + @memset(p.zeros.slice(i32, 64), 0); + p.hash = h.final(); + return p; +} + +/// kda_rows with f_b and g_b at one width a pipeline (`fg_block`), under the names glm_kda_fg_q. +const fg_names = blk: { + var n: [nw][:0]const u8 = undefined; + for (bit_widths, 0..) |b, i| n[i] = std.fmt.comptimePrint("glm_kda_fg_q{d}", .{b}); + break :blk n; +}; +const fg_instances = blk: { + var s: []const u8 = ""; + for (bit_widths, fg_names) |b, name| s = s ++ kernels.instance("kda_rows", std.fmt.comptimePrint("64, 128, 4, 32, {d}, {d}", .{ b, b }), name); + break :blk s; +}; + +fn bind(e: mtl.ComputeEncoder, first: usize, refs: anytype) void { + inline for (refs, 0..) |r, i| e.setBuffer(r.buf, r.off, first + i); +} + +fn wBytes(n: u32, k: u32, q: Affine) usize { + return affine.wBytes(n, k, q.width()); +} + +fn sBytes(n: u32, k: u32, q: Affine) usize { + return affine.sBytes(n, k, q.width()); +} + +fn at(bits: u8) usize { + return std.mem.indexOfScalar(u8, &bit_widths, bits).?; +} + +/// Rows of x (`ldx` apart) through a matrix into out (`ldo` apart). +pub const Io = struct { in: Ref, ldx: u32, out: Ref, ldo: u32, rows: u32 }; + +fn launch(p: *const Pipes, e: mtl.ComputeEncoder, w: [3]Ref, bits: u8, d: RowDims, in: Ref, out: Ref) void { + e.setPipeline(p.rows[at(bits)]); + bind(e, 0, .{ in, w[0], w[1], w[2] }); + e.setValue(d, 4); + bind(e, 5, .{out}); + e.dispatchGroups(mtl.Size.of(@intCast(@divTrunc(d.n + 7, 8)), @intCast(d.heads), 1), mtl.Size.of(64, 1, 1)); +} + +/// x through every member of `q` at its own width, each into its own output rows. +pub fn dense(p: *const Pipes, e: mtl.ComputeEncoder, q: Affine, io: Io) void { + for (0..@as(usize, q.more) + 1) |i| { + const r = q.run(i); + launch(p, e, .{ r.w, r.s, r.b }, r.bits, .{ .k = @intCast(q.k), .n = @intCast(r.n), .group = r.group, .rows = @intCast(io.rows), .ldx = @intCast(io.ldx), .ldo = @intCast(io.ldo) }, io.in, io.out.at(@as(usize, r.row) * 2)); + } +} + +/// `heads` matrices stacked [heads, n / heads, k] (embed_q, unembed_out): head h reads x + h hsx, writes out + h hso. +pub fn perHead(p: *const Pipes, e: mtl.ComputeEncoder, q: Affine, heads: u32, io: Io, hsx: u32, hso: u32) void { + const r = q.n / heads; + launch(p, e, .{ q.w, q.s, q.b }, q.bits, .{ .k = @intCast(q.k), .n = @intCast(r), .group = q.group, .rows = @intCast(io.rows), .heads = @intCast(heads), .ldx = @intCast(io.ldx), .hsx = @intCast(hsx), .ldo = @intCast(io.ldo), .hso = @intCast(hso), .hsw = @intCast(r) }, io.in, io.out); +} + +/// A chunk's x slices (`kx` apart) through q's heads into out [rows, heads, n / heads]: one affine_mm batch, else rows. +pub fn headsBatch(p: *const Pipes, e: mtl.ComputeEncoder, q: Affine, h: u32, io: Io, kx: u32, sums: Ref) void { + const ny = q.n / h; + const pipes = mm(p, q) orelse return perHead(p, e, q, h, io, kx, ny); + var head = q; + head.n = ny; + affine_mm.rowSumsAt(e, pipes, 64, io.in, sums, io.rows, h * kx, io.ldx); // a head's sums one row apart + affine_mm.denseBatch(e, pipes, io.in, sums, head, io.out, io.rows, h, .{ .x_row = @intCast(io.ldx), .y_row = @intCast(io.ldo), .sums_row = @intCast(h), .x_batch = @intCast(kx), .y_batch = @intCast(ny), .sums_batch = 1, .w_batch = @intCast(ny) }); +} + +/// x [rows, k] through bf16 W [n, k]. +pub fn bf16(p: *const Pipes, e: mtl.ComputeEncoder, w: Ref, n: u32, k: u32, io: Io) void { + e.setPipeline(p.bf16); + bind(e, 0, .{ io.in, w }); + e.setValue(RowDims{ .k = @intCast(k), .n = @intCast(n), .group = 0, .rows = @intCast(io.rows), .ldx = @intCast(io.ldx), .ldo = @intCast(io.ldo) }, 4); + bind(e, 5, .{io.out}); + e.dispatchGroups(mtl.Size.of((n + 7) / 8, 1, 1), mtl.Size.of(64, 1, 1)); +} + +/// The routed experts `w` (stacked [experts, n, k]) by pick: with `up`, gate and up SwiGLU'd. +pub const Picks = struct { w: Affine, up: ?Affine = null, experts: u32, repeat: u32, lim: f32 = 0 }; + +fn pickDims(m: Picks, slots: u32, ldo: u32) PickDims { + return .{ .k = @intCast(m.w.k), .n = @intCast(m.w.n), .group = m.w.group, .slots = @intCast(slots), .repeat = @intCast(m.repeat), .experts = @intCast(m.experts), .ldo = @intCast(ldo), .lim = m.lim }; +} + +/// Slot y [0, slots) through expert ids[y] (i32) on x row y / repeat, into out row y. +pub fn byIds(p: *const Pipes, e: mtl.ComputeEncoder, m: Picks, x: Ref, ids: Ref, out: Ref, slots: u32) void { + const u = m.up orelse m.w; + e.setPipeline(if (m.up != null) p.glu_ids[at(m.w.bits)] else p.ids[at(m.w.bits)]); + bind(e, 0, .{ x, m.w.w, m.w.s, m.w.b }); + e.setValue(pickDims(m, slots, m.w.n), 4); + bind(e, 5, .{ out, ids, u.w, u.s, u.b }); + e.dispatchGroups(mtl.Size.of((m.w.n + 7) / 8, slots, 1), mtl.Size.of(64, 1, 1)); +} + +/// x through [gate; up] (one width, `n` rows each) SwiGLU'd into out: the by-ids GLU, every id 0 (glm_swiglu's bits). +pub fn glu(p: *const Pipes, e: mtl.ComputeEncoder, gu: Affine, n: u32, x: Ref, out: Ref, rows: u32, lim: f32) void { + std.debug.assert(gu.more == 0 and gu.n == 2 * n and rows <= 64); + var gate = gu; + gate.n = n; + var up = gate; + up.w = gu.w.at(wBytes(n, gu.k, gu)); + up.s = gu.s.at(sBytes(n, gu.k, gu)); + up.b = gu.b.at(sBytes(n, gu.k, gu)); + byIds(p, e, .{ .w = gate, .up = up, .experts = 1, .repeat = 1, .lim = lim }, x, .{ .buf = p.zeros }, out, rows); +} + +/// A MoE block's routed experts and shared expert (`sh`: [gate; up] one run, then down), all g64, the shared at 8 bits. +pub const Moe = struct { gate: Affine, up: Affine, down: Affine, sh_gate_up: Affine, sh_down: Affine }; + +/// Whether `m` takes `moe`: every matrix one run, g64, the routed of one width (not 4/64), the shared at 8 bits. +pub fn moeOne(m: Moe) bool { + if (native(m.gate) and native(m.up) and native(m.down)) return false; // routed 4/64 keeps the generated gathers + for ([_]Affine{ m.gate, m.up, m.down, m.sh_gate_up, m.sh_down }) |q| if (q.more != 0 or q.group != 64) return false; + return m.up.bits == m.gate.bits and m.down.bits == m.gate.bits and m.sh_gate_up.bits == 8 and m.sh_down.bits == 8; +} + +/// byIds's GLU and down for the slots, glu and dense for the shared expert's rows: a launch each step, the same bits. +pub fn moe(p: *const Pipes, e: mtl.ComputeEncoder, m: Moe, experts: u32, topk: u32, lim: f32, io: MoeIo, rows: u32) void { + const slots = rows * topk; + const g = m.sh_gate_up; + const steps = [2]struct { pipe: mtl.Pipeline, w: Affine, up: ?Affine, sh: Affine, x: Ref, x2: Ref, out: Ref, out2: Ref, repeat: u32 }{ + .{ .pipe = p.moe_glu[at(m.gate.bits)], .w = m.gate, .up = m.up, .sh = g, .x = io.x, .x2 = io.x, .out = io.act, .out2 = io.acts, .repeat = topk }, + .{ .pipe = p.moe[at(m.down.bits)], .w = m.down, .up = null, .sh = m.sh_down, .x = io.act, .x2 = io.acts, .out = io.ye, .out2 = io.ys, .repeat = 1 }, + }; + for (steps) |s| { + const u = s.up orelse s.w; + e.setPipeline(s.pipe); + bind(e, 0, .{ s.x, s.w.w, s.w.s, s.w.b }); + e.setValue(PickDims{ .k = @intCast(s.w.k), .n = @intCast(s.w.n), .group = 64, .slots = @intCast(slots), .repeat = @intCast(s.repeat), .experts = @intCast(experts), .ldo = @intCast(s.w.n), .lim = lim }, 4); + bind(e, 5, .{ s.out, io.ids, u.w, u.s, u.b, s.x2, s.sh.w, s.sh.s, s.sh.b, s.out2 }); + e.dispatchGroups(mtl.Size.of((s.w.n + 7) / 8, slots + rows, 1), mtl.Size.of(64, 1, 1)); + } +} + +/// `moe`'s buffers: x rows, the picks' expert ids, act and acts (gate|up out), ye and ys (down out). +pub const MoeIo = struct { x: Ref, ids: Ref, act: Ref, acts: Ref, ye: Ref, ys: Ref }; + +/// Each expert's sorted pairs (offsets, order): with `up` x row order / repeat to row i, else x row i to row order[i]. +pub fn bySorted(p: *const Pipes, e: mtl.ComputeEncoder, m: Picks, x: Ref, offsets: Ref, order: Ref, out: Ref, ldo: u32) void { + const u = m.up orelse m.w; + e.setPipeline(if (m.up != null) p.glu_sorted[at(m.w.bits)] else p.sorted[at(m.w.bits)]); + bind(e, 0, .{ x, m.w.w, m.w.s, m.w.b }); + e.setValue(pickDims(m, 0, ldo), 4); + bind(e, 5, .{ out, offsets, order, u.w, u.s, u.b }); + e.dispatchGroups(mtl.Size.of((m.w.n + 7) / 8, m.experts, 1), mtl.Size.of(64, 1, 1)); +} + +/// A prompt chunk's affine_mm pipelines for `q` (one width, a format wide.zig compiles), else null: the row kernels. +pub fn mm(p: *const Pipes, q: Affine) ?affine_mm.Pipes { + if (q.more != 0 or q.group != 64) return null; + return switch (q.bits) { + 8 => p.mm8, + 2 => p.mm2, + 5 => p.mm5, + 6 => p.mm6, + else => null, + }; +} + +/// A chunk's product into io.out: each member on affine_mm at its width if it fills 64-row tiles, else the row kernels. +pub fn denseMm(p: *const Pipes, e: mtl.ComputeEncoder, q: Affine, io: Io, sums: Ref) void { + var summed = false; + for (0..@as(usize, q.more) + 1) |i| { + const r = q.run(i); + const one: Affine = .{ .w = r.w, .s = r.s, .b = r.b, .n = r.n, .k = q.k, .bits = r.bits, .group = r.group }; + var part = io; + part.out = io.out.at(@as(usize, r.row) * 2); + const pipes = (if (q.more == 0 or r.n % 64 == 0) mm(p, one) else null) orelse { + dense(p, e, one, part); + continue; + }; + if (!summed) affine_mm.rowSumsAt(e, pipes, 64, io.in, sums, io.rows, q.k, io.ldx); // every member's: one k, g64 + summed = true; + affine_mm.denseBatch(e, pipes, io.in, sums, one, part.out, io.rows, 1, .{ .x_row = @intCast(io.ldx), .y_row = @intCast(io.ldo) }); + } +} + +/// A prompt chunk's routed experts: affine_mm gathers at 4/64 (`mm4`) or 2/64 (`mm2`, oQ2's), else sorted row products. +pub const Routed = enum { mm4, mm2, rows }; + +pub fn routed(gate: Affine, up: Affine, down: Affine) Routed { + if (native(gate) and native(up) and native(down)) return .mm4; + for ([_]Affine{ gate, up, down }) |q| if (q.more != 0 or q.bits != 2 or q.group != 64) return .rows; + return .mm2; +} + +/// f_b/g_b dots: generated (both 4/64), `.fused` (one width, g64, glm_rows's sums inside), else `.pre`: same bits. +pub const KdaGate = union(enum) { generated, fused: mtl.Pipeline, pre: mtl.Pipeline }; + +pub fn kdaGate(p: ?*const Pipes, f: Affine, g: Affine) KdaGate { + if (native(f) and native(g)) return .generated; + if (f.more == 0 and g.more == 0 and f.group == 64 and g.group == 64 and f.bits == g.bits) return .{ .fused = p.?.kda_fg[at(f.bits)] }; + return .{ .pre = p.?.kda_pre }; +} + +test "4/64 of one width alone takes the generated kernels" { + const r: Ref = .{ .buf = undefined }; + const q: Affine = .{ .w = r, .s = r, .b = r, .n = 64, .k = 128, .bits = 4, .group = 64 }; + try std.testing.expect(native(q)); + var m = q; + m.bits = 8; + try std.testing.expect(!native(m)); + m = q; + m.more = 1; + try std.testing.expect(!native(m)); + try std.testing.expect(kdaGate(null, q, q) == .generated); +} + +test "a failed compile releases every pipeline built before it; a load counts each one it builds" { + const gpa = std.testing.allocator; + const device = try mtl.Device.init(); + defer device.deinit(); + var made: Made = .{}; + try std.testing.expectError(error.KernelCompile, build(gpa, device, "", &made)); // after the row kernels + try std.testing.expectEqual(@as(usize, nw * 7 + 1), made.released); + try std.testing.expectEqual(@as(usize, 0), made.n); + var all: Made = .{}; + var p = try build(gpa, device, comptime kernels.generatedSource("kda_rows"), &all); + p.deinit(); + try std.testing.expectEqual(all.list.len, all.n); +} + +test "glm_rows's SwiGLU compiles glm_glue's sigmoid_fast and bfloat overloads, no copy of its own" { + try std.testing.expect(std.mem.indexOf(u8, glue_math, "inline U sigmoid_fast(U x)") != null); + try std.testing.expect(std.mem.indexOf(u8, glue_math, "METAL_FUNC bfloat16_t exp(bfloat16_t x)") != null); + try std.testing.expect(std.mem.indexOf(u8, glue_math, "kernel void") == null); // the helpers only + try std.testing.expect(std.mem.indexOf(u8, sources.glm_rows, "exp(") == null); + try std.testing.expect(std.mem.indexOf(u8, sources.glm_rows, "sigmoid_fast(gt)") != null); +} + +test "prompt chunks gather routed experts on affine_mm at 4/64 and 2/64 only; other widths take the sorted row products" { + const r: Ref = .{ .buf = undefined }; + const q: Affine = .{ .w = r, .s = r, .b = r, .n = 64, .k = 128, .bits = 2, .group = 64 }; + var four = q; + four.bits = 4; + try std.testing.expectEqual(Routed.mm4, routed(four, four, four)); + try std.testing.expectEqual(Routed.mm2, routed(q, q, q)); + var three = q; + three.bits = 3; + var g32 = q; + g32.group = 32; + try std.testing.expectEqual(Routed.rows, routed(q, q, three)); + try std.testing.expectEqual(Routed.rows, routed(q, q, four)); + try std.testing.expectEqual(Routed.rows, routed(g32, g32, g32)); +} + +test "f_b and g_b of one width (g64) take glm_rows's dots inside kda_rows, two widths the products made apart" { + const r: Ref = .{ .buf = undefined }; + const q: Affine = .{ .w = r, .s = r, .b = r, .n = 64, .k = 128, .bits = 8, .group = 64 }; + var four = q; + four.bits = 4; + var g32 = q; + g32.group = 32; + const p: Pipes = undefined; + for ([_][2]Affine{ .{ four, q }, .{ q, four }, .{ g32, g32 } }) |fg| try std.testing.expect(kdaGate(&p, fg[0], fg[1]) == .pre); + var five = q; + five.bits = 5; + for ([_]Affine{ q, five }) |x| try std.testing.expect(kdaGate(&p, x, x) == .fused); +} diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 1823eb56a9..d9d7c8a072 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -13,6 +13,9 @@ 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"); +/// 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/affine_ref.zig b/zig/tests/affine_ref.zig new file mode 100644 index 0000000000..4849616083 --- /dev/null +++ b/zig/tests/affine_ref.zig @@ -0,0 +1,368 @@ +//! 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); + const at = (j * self.k + i) * self.bits; // MLX's little-endian bit stream at every width + var v: u32 = 0; + for (0..3) |b| if (at / 8 + b < bytes.len) { + v |= @as(u32, bytes[at / 8 + b]) << @intCast(8 * b); + }; + return (v >> @intCast(at % 8)) & ((@as(u32, 1) << @intCast(self.bits)) - 1); + } + + /// y = x W^T in float64, and the sizes the fp32 sums carry when codes enter as q + off (affine_mm's 4-bit: 128). + 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); + }; + for ([_][2]usize{ .{ 192, 1024 }, .{ 256, 4096 } }) |shape| for ([_]u8{ 3, 5, 6 }) |bits| for ([_]usize{ 32, 64 }) |group| { // oQ4's 5 and 6 + const s = try Synth.init(device, shape[0], shape[1], bits, group, 0x3716 + 100 * @as(usize, bits) + shape[0] + group); + defer s.deinit(); + const ref = try s.reference(0); + defer ref.deinit(); + for ([_]bool{ false, true }) |f32_out| try affine(&t, device, queue, &s, ref, f32_out); + }; + var p = try sq.Pipelines.load(device); + defer p.deinit(); + for ([_][2]usize{ .{ 40, 512 }, .{ 256, 4096 }, .{ 6400, 256 } }) |shape| { // chunks of 32, 16 and 8 + 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}); +} 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, }); diff --git a/zig/tests/glm_rows_check.zig b/zig/tests/glm_rows_check.zig new file mode 100644 index 0000000000..b2866b7c9c --- /dev/null +++ b/zig/tests/glm_rows_check.zig @@ -0,0 +1,569 @@ +//! GLM-5.3-Flash's products at every MLX width (glm/wide.zig) against fp64, each row's bits in every window and launch. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const sources = @import("kernel_sources"); +const wide = tf.glm.wide; +const affine = tf.glm.affine; +const mlx = tf.quant.mlx; +const row = tf.row_projection; +const affine_mm = tf.affine_mm; + +const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; +const a = std.heap.page_allocator; +const rows = 16; +const pad = 32; // a strided x row's and out row's extra values +const Ref = affine.Ref; + +fn toBf(v: f32) u16 { + const u: u32 = @bitCast(v); + return @truncate((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); +} +fn fromBf(v: u16) f64 { + return @as(f32, @bitCast(@as(u32, v) << 16)); +} +pub fn buf(device: mtl.Device, bytes: usize) !mtl.Buffer { + return device.buffer(@max(bytes, 16), opts); +} +pub fn ref(b: mtl.Buffer) Ref { + return .{ .buf = b }; +} + +/// A synthetic MLX matrix [n, k], its fp64 dequantization and its codes alone; x: `rows` rows of inputs. +pub const Mat = struct { + q: affine.Affine, + bufs: [4]mtl.Buffer, + deq: []f64, + codes: []f64, + + pub fn init(device: mtl.Device, n: usize, k: usize, bits: u8, group: usize, seed: u64) !Mat { + var prng = std.Random.DefaultPrng.init(seed); + const r = prng.random(); + const words = n * k * bits / 32; + const tables = n * (k / group); + const bs = [4]mtl.Buffer{ try buf(device, words * 4), try buf(device, tables * 2), try buf(device, tables * 2), try buf(device, rows * k * 2) }; + for (bs[0].slice(u32, words)) |*v| v.* = r.int(u32); + for (bs[1].slice(u16, tables)) |*v| v.* = toBf(r.float(f32) * 0.02 + 0.001); + for (bs[2].slice(u16, tables)) |*v| v.* = toBf((r.float(f32) - 0.5) * 0.1); + for (bs[3].slice(u16, rows * k)) |*v| v.* = toBf((r.float(f32) - 0.5) * 2); + const m = Mat{ .q = .{ .w = ref(bs[0]), .s = ref(bs[1]), .b = ref(bs[2]), .n = @intCast(n), .k = @intCast(k), .bits = bits, .group = @intCast(group) }, .bufs = bs, .deq = try a.alloc(f64, n * k), .codes = try a.alloc(f64, n * k) }; + try mlx.reference(m.host(bs[1].slice(u16, tables), bs[2].slice(u16, tables)), m.deq); + const ones = try a.alloc(u16, tables); + defer a.free(ones); + @memset(ones, toBf(1)); + const zeros = try a.alloc(u16, tables); + defer a.free(zeros); + @memset(zeros, 0); + try mlx.reference(m.host(ones, zeros), m.codes); + return m; + } + + pub fn deinit(m: Mat) void { + for (m.bufs) |b| b.deinit(); + a.free(m.deq); + a.free(m.codes); + } + + pub fn x(m: Mat) mtl.Buffer { + return m.bufs[3]; + } + + fn host(m: Mat, scales: []const u16, biases: []const u16) mlx.Host { + const q = m.q; + const g = q.k / q.group; + return .{ + .words = .{ .dtype = .i32, .rank = 2, .shape = .{ q.n, q.k * q.bits / 32, 1, 1, 1 }, .bytes = m.bufs[0].slice(u8, q.n * q.k * q.bits / 8) }, + .scales = .{ .dtype = .bf16, .rank = 2, .shape = .{ q.n, g, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(scales) }, + .biases = .{ .dtype = .bf16, .rank = 2, .shape = .{ q.n, g, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(biases) }, + .bits = q.bits, + .group = q.group, + }; + } + + /// Row i against column c in fp64, and the fp32 sums' bound (output rounding plus k-deep accumulation). + fn reference(m: Mat, i: usize, c: usize) [2]f64 { + const k: usize = m.q.k; + const g_n = k / m.q.group; + const sc = m.bufs[1].slice(u16, m.q.n * g_n); + const bi = m.bufs[2].slice(u16, m.q.n * g_n); + const xs = m.x().slice(u16, rows * k); + var y: f64 = 0; + var terms: f64 = 0; + for (0..g_n) |g| { + var dot: f64 = 0; + var sum: f64 = 0; + for (0..m.q.group) |j| { + const kk = g * m.q.group + j; + const xv = fromBf(xs[i * k + kk]); + y += xv * m.deq[c * k + kk]; + dot += @abs(xv * m.codes[c * k + kk]); + sum += @abs(xv); + } + terms += @abs(fromBf(sc[c * g_n + g])) * dot + @abs(fromBf(bi[c * g_n + g])) * sum; + } + const gamma = @import("glm_rows_classes.zig").gamma(k); + return .{ y, std.math.ldexp(@abs(y) + gamma * terms, -8) + gamma * terms }; + } + + /// Expert (or head) e of `parts` stacked matrices: its rows, codes and tables. + fn part(m: Mat, e: usize, parts: usize) affine.Affine { + var q = m.q; + const n = q.n / parts; + const wb = affine.wBytes(n, q.k, q.width()); + const sb = affine.sBytes(n, q.k, q.width()); + q.w = q.w.at(e * wb); + q.s = q.s.at(e * sb); + q.b = q.b.at(e * sb); + q.n = @intCast(n); + return q; + } +}; + +pub const Gpu = struct { + device: mtl.Device, + queue: mtl.Queue, + p: wide.Pipes, + core: row.Pipelines, + swiglu: mtl.Pipeline, + + /// One command buffer around `f(e, args)`. + pub fn run(g: *const Gpu, comptime f: anytype, args: anytype) !void { + const cb = g.queue.commandBuffer(); + const e = cb.compute(.serial); + @call(.auto, f, .{e} ++ args); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.debug.print("GPU failed: {s}\n", .{msg}); + return error.GpuFailed; + } + } +}; + +pub fn dense(e: mtl.ComputeEncoder, p: *const wide.Pipes, q: affine.Affine, io: wide.Io) void { + wide.dense(p, e, q, io); +} + +const Tally = struct { values: usize = 0, worst: f64 = 0, pairs: usize = 0 }; + +pub fn eql(got: []const u16, want: []const u16) !void { + if (!std.mem.eql(u16, got, want)) { + for (got, want, 0..) |x, y, i| if (x != y) { + std.debug.print("first difference at {d}: {x} vs {x}\n", .{ i, x, y }); + break; + }; + return error.BitsDiffer; + } +} + +/// One matrix: fp64 per column, windows and strided rows equal to each row alone, the core row kernel's bits. +fn denseCheck(g: *const Gpu, m: Mat, t: *Tally) !void { + const n: usize = m.q.n; + const k: usize = m.q.k; + const alone = try buf(g.device, rows * n * 2); + defer alone.deinit(); + for (0..rows) |i| try g.run(dense, .{ &g.p, m.q, wide.Io{ .in = ref(m.x()).at(i * k * 2), .ldx = m.q.k, .out = ref(alone).at(i * n * 2), .ldo = m.q.n, .rows = 1 } }); + const got = alone.slice(u16, rows * n); + for (0..rows) |i| for (0..n) |c| { + const r = m.reference(i, c); + const err = @abs(fromBf(got[i * n + c]) - r[0]); + t.worst = @max(t.worst, err / r[1]); + if (err > r[1]) { + std.debug.print("row {d} col {d}: {d} vs fp64 {d} (bound {e})\n", .{ i, c, fromBf(got[i * n + c]), r[0], r[1] }); + return error.PastBound; + } + t.values += 1; + }; + if (n % 8 == 0) { // the core row kernel's sums (its tiles: 8 columns): the same bits for a row + const core = try buf(g.device, n * 2); + defer core.deinit(); + const call = try row.call(&g.core, .{ .w = m.bufs[0], .scales = m.bufs[1], .biases = m.bufs[2], .n = n, .k = k, .group = m.q.group, .bits = m.q.bits }, 1, false); + try g.run(coreRow, .{ call, m, core }); + try eql(core.slice(u16, n), got[0..n]); + t.pairs += n; + } + // windows, contiguous (two rows a weight read) and strided (one at a time, out rows apart with a lead offset) + const xs = try buf(g.device, rows * (k + pad) * 2); + defer xs.deinit(); + for (0..rows) |i| @memcpy(xs.slice(u16, rows * (k + pad))[i * (k + pad) ..][0..k], m.x().slice(u16, rows * k)[i * k ..][0..k]); + const out = try buf(g.device, (rows * (n + pad) + 4) * 2); + defer out.deinit(); + for ([_]usize{ 2, 3, 16 }) |w| for ([_]bool{ false, true }) |strided| { + const ld: usize = if (strided) n + pad else n; + const at: usize = if (strided) 4 else 0; + try g.run(dense, .{ &g.p, m.q, wide.Io{ .in = ref(if (strided) xs else m.x()), .ldx = @intCast(if (strided) k + pad else k), .out = ref(out).at(at * 2), .ldo = @intCast(ld), .rows = @intCast(w) } }); + for (0..w) |i| try eql(out.slice(u16, rows * (n + pad) + 4)[at + i * ld ..][0..n], got[i * n ..][0..n]); + t.pairs += w * n; + }; +} + +fn coreRow(e: mtl.ComputeEncoder, c: row.Call, m: Mat, out: mtl.Buffer) void { + e.setPipeline(c.pipeline); + e.setBuffer(m.x(), 0, 0); + e.setBuffer(m.bufs[0], 0, 1); + e.setBuffer(m.bufs[1], 0, 2); + e.setBuffer(m.bufs[2], 0, 3); + e.setValue(c.dims, 4); + e.setBuffer(out, 0, 5); + e.dispatchGroups(mtl.Size.of(c.groups, c.slots, 1), mtl.Size.of(32 * row.groups, 1, 1)); +} + +/// Four heads (or experts): per head, by pick (GLU and plain) and sorted (8, 5, 2, 1 pairs), each against dense. +fn partsCheck(g: *const Gpu, m: Mat, up: Mat, t: *Tally) !void { + const parts = 4; + const k: usize = m.q.k; + const r: usize = m.q.n / parts; + const n_rows = 3; // x rows a head + const repeat = 2; // picks an x row: 8 rows, so an expert's sorted pairs fill 4- and 8-pair passes and leave tails + const slots = 8 * repeat; + // per head: x [rows, heads, k], out [rows, heads, r] + const ph = try buf(g.device, n_rows * parts * r * 2); + defer ph.deinit(); + try g.run(perHead, .{ &g.p, m.q, @as(u32, parts), wide.Io{ .in = ref(m.x()), .ldx = @intCast(parts * k), .out = ref(ph), .ldo = @intCast(parts * r), .rows = n_rows }, @as(u32, @intCast(k)), @as(u32, @intCast(r)) }); + const one = try buf(g.device, n_rows * r * 2); + defer one.deinit(); + for (0..parts) |h| { + try g.run(dense, .{ &g.p, m.part(h, parts), wide.Io{ .in = ref(m.x()).at(h * k * 2), .ldx = @intCast(parts * k), .out = ref(one), .ldo = @intCast(r), .rows = n_rows } }); + for (0..n_rows) |i| try eql(ph.slice(u16, n_rows * parts * r)[(i * parts + h) * r ..][0..r], one.slice(u16, n_rows * r)[i * r ..][0..r]); + t.pairs += n_rows * r; + } + // by pick: slot y through expert ids[y] on x row y / repeat + const ids = try buf(g.device, slots * 4); + defer ids.deinit(); + const picks = [slots]i32{ 2, 0, 0, 2, 2, 1, 3, 2, 2, 0, 1, 2, 2, 0, 0, 2 }; + @memcpy(ids.slice(i32, slots), &picks); + var gate = m.q; + gate.n = @intCast(r); + var upq = up.q; + upq.n = @intCast(r); + const plain = try buf(g.device, slots * r * 2); + defer plain.deinit(); + const glu = try buf(g.device, slots * r * 2); + defer glu.deinit(); + const lim: f32 = 0.75; + try g.run(byIds, .{ &g.p, wide.Picks{ .w = gate, .experts = parts, .repeat = repeat }, ref(m.x()), ref(ids), ref(plain), @as(u32, slots) }); + try g.run(byIds, .{ &g.p, wide.Picks{ .w = gate, .up = upq, .experts = parts, .repeat = repeat, .lim = lim }, ref(m.x()), ref(ids), ref(glu), @as(u32, slots) }); + const gu = try buf(g.device, slots * 2 * r * 2); + defer gu.deinit(); + const act = try buf(g.device, slots * r * 2); + defer act.deinit(); + for (0..slots) |y| { + const e: usize = @intCast(picks[y]); + const x_at = ref(m.x()).at(y / repeat * k * 2); + try g.run(dense, .{ &g.p, m.part(e, parts), wide.Io{ .in = x_at, .ldx = @intCast(k), .out = ref(gu).at(y * 2 * r * 2), .ldo = @intCast(2 * r), .rows = 1 } }); + try g.run(dense, .{ &g.p, up.part(e, parts), wide.Io{ .in = x_at, .ldx = @intCast(k), .out = ref(gu).at((y * 2 * r + r) * 2), .ldo = @intCast(2 * r), .rows = 1 } }); + try eql(plain.slice(u16, slots * r)[y * r ..][0..r], gu.slice(u16, slots * 2 * r)[y * 2 * r ..][0..r]); + } + try g.run(swiglu, .{ g.swiglu, ref(gu), ref(act), @as(u32, slots), @as(u32, @intCast(r)), lim }); + try eql(glu.slice(u16, slots * r), act.slice(u16, slots * r)); + t.pairs += 2 * slots * r; + // sorted: each expert's pairs in pick order; GLU out row i, down out row order[i] from x row i + var offsets: [parts + 1]i32 = @splat(0); + var order: [slots]i32 = undefined; + var at: usize = 0; + for (0..parts) |e| { + offsets[e] = @intCast(at); + for (picks, 0..) |pk, y| if (pk == e) { + order[at] = @intCast(y); + at += 1; + }; + } + offsets[parts] = @intCast(at); + const ob = try buf(g.device, (parts + 1 + slots) * 4); + defer ob.deinit(); + @memcpy(ob.slice(i32, parts + 1), &offsets); + @memcpy(ob.slice(i32, parts + 1 + slots)[parts + 1 ..], &order); + const order_ref = ref(ob).at((parts + 1) * 4); + const sorted = try buf(g.device, slots * r * 2); + defer sorted.deinit(); + try g.run(bySorted, .{ &g.p, wide.Picks{ .w = gate, .up = upq, .experts = parts, .repeat = repeat, .lim = lim }, ref(m.x()), ref(ob), order_ref, ref(sorted), @as(u32, @intCast(r)) }); + for (0..slots) |i| try eql(sorted.slice(u16, slots * r)[i * r ..][0..r], glu.slice(u16, slots * r)[@as(usize, @intCast(order[i])) * r ..][0..r]); + // the down: x row i (the sorted pairs' inputs, here x rows by order) into out row order[i] + const xs = try buf(g.device, slots * k * 2); + defer xs.deinit(); + for (0..slots) |i| @memcpy(xs.slice(u16, slots * k)[i * k ..][0..k], m.x().slice(u16, rows * k)[@as(usize, @intCast(order[i])) * k ..][0..k]); + try g.run(bySorted, .{ &g.p, wide.Picks{ .w = gate, .experts = parts, .repeat = 1 }, ref(xs), ref(ob), order_ref, ref(sorted), @as(u32, @intCast(r)) }); + try g.run(byIds, .{ &g.p, wide.Picks{ .w = gate, .experts = parts, .repeat = 1 }, ref(m.x()), ref(ids), ref(plain), @as(u32, slots) }); + try eql(sorted.slice(u16, slots * r), plain.slice(u16, slots * r)); + t.pairs += 2 * slots * r; +} + +fn mmDense(e: mtl.ComputeEncoder, mm: affine_mm.Pipes, q: affine.Affine, x: Ref, sums: Ref, y: Ref, n_rows: u32) void { + affine_mm.rowSums(e, mm, 64, x, sums, n_rows, q.k); + affine_mm.dense(e, mm, x, sums, q, y, n_rows); +} + +/// dims: x rows, sorted pairs, experts, picks a row. +fn mmGlu(e: mtl.ComputeEncoder, mm: affine_mm.Pipes, gate: affine.Affine, up: affine.Affine, x: Ref, sums: Ref, sorted: [2]Ref, act: Ref, dims: [4]u32, lim: f32) void { + affine_mm.rowSums(e, mm, 64, x, sums, dims[0], gate.k); + affine_mm.gatherGlu(e, mm, x, sums, gate, up, sorted[0], sorted[1], act, dims[1], dims[2], dims[3], lim); +} + +fn mmDown(e: mtl.ComputeEncoder, mm: affine_mm.Pipes, q: affine.Affine, x: Ref, sums: Ref, sorted: [2]Ref, y: Ref, slots: u32, experts: u32) void { + affine_mm.rowSums(e, mm, 64, x, sums, slots, q.k); + affine_mm.gatherTo(e, mm, x, sums, q, sorted[0], y, sorted[1], slots, experts); +} + +/// oQ2's prompt experts, affine_mm at 2/64: gathered GLU == glm_swiglu(dense gate | up), gathered down == dense. +fn gatherCheck(g: *const Gpu, m: Mat, up: Mat, t: *Tally) !void { + const parts = 4; + const k: usize = m.q.k; + const r: usize = m.q.n / parts; + const repeat = 2; + const slots = 8 * repeat; + const picks = [slots]i32{ 2, 0, 0, 2, 2, 1, 3, 2, 2, 0, 1, 2, 2, 0, 0, 2 }; + const lim: f32 = 0.75; + const mm = g.p.mm2; + var order: [slots]i32 = undefined; + var offsets: [parts]i32 = undefined; + var at: usize = 0; + for (0..parts) |e| { + offsets[e] = @intCast(at); + for (picks, 0..) |pk, y| if (pk == e) { + order[at] = @intCast(y); + at += 1; + }; + } + const ob = try buf(g.device, (parts + slots) * 4); + defer ob.deinit(); + @memcpy(ob.slice(i32, parts), &offsets); + @memcpy(ob.slice(i32, parts + slots)[parts..], &order); + const sorted = [2]Ref{ ref(ob), ref(ob).at(parts * 4) }; + const bs = [_]mtl.Buffer{ try buf(g.device, rows * (k / 64) * 4), try buf(g.device, 2 * parts * rows * r * 2), try buf(g.device, slots * 2 * r * 2), try buf(g.device, 2 * slots * r * 2), try buf(g.device, slots * k * 2), try buf(g.device, (parts + 1) * slots * r * 2) }; + defer for (bs) |b| b.deinit(); + const sums = ref(bs[0]); + var gate = m.q; + gate.n = @intCast(r); + var upq = up.q; + upq.n = @intCast(r); + // every expert's gate and up, dense on every x row: [2][parts][rows, r] + for (0..parts) |e| for ([2]Mat{ m, up }, 0..) |w, j| { + try g.run(mmDense, .{ mm, w.part(e, parts), ref(m.x()), sums, ref(bs[1]).at(((j * parts + e) * rows * r) * 2), @as(u32, rows) }); + }; + const dg = bs[1].slice(u16, 2 * parts * rows * r); + const gu = bs[2].slice(u16, slots * 2 * r); + for (order, 0..) |o, i| { + const e: usize = @intCast(picks[@intCast(o)]); + const xr: usize = @as(usize, @intCast(o)) / repeat; + for (0..2) |j| @memcpy(gu[(2 * i + j) * r ..][0..r], dg[((j * parts + e) * rows + xr) * r ..][0..r]); + } + const act = ref(bs[3]); + try g.run(mmGlu, .{ mm, gate, upq, ref(m.x()), sums, sorted, act, [4]u32{ rows, slots, parts, repeat }, lim }); + try g.run(swiglu, .{ g.swiglu, ref(bs[2]), act.at(slots * r * 2), @as(u32, slots), @as(u32, @intCast(r)), lim }); + const both = bs[3].slice(u16, 2 * slots * r); + try eql(both[0 .. slots * r], both[slots * r ..]); + // the down: sorted pair i's input (x row order[i]) into row order[i], == its expert's dense product of that row + const xs = bs[4].slice(u16, slots * k); + for (order, 0..) |o, i| @memcpy(xs[i * k ..][0..k], m.x().slice(u16, rows * k)[@as(usize, @intCast(o)) * k ..][0..k]); + const y = bs[5].slice(u16, (parts + 1) * slots * r); + try g.run(mmDown, .{ mm, gate, ref(bs[4]), sums, sorted, ref(bs[5]), @as(u32, slots), @as(u32, parts) }); + for (0..parts) |e| try g.run(mmDense, .{ mm, m.part(e, parts), ref(bs[4]), sums, ref(bs[5]).at((1 + e) * slots * r * 2), @as(u32, slots) }); + for (order, 0..) |o, i| { + const e: usize = @intCast(picks[@intCast(o)]); + try eql(y[@as(usize, @intCast(o)) * r ..][0..r], y[((1 + e) * slots + i) * r ..][0..r]); + } + t.pairs += 2 * slots * r; +} + +/// A split stack (a run a member, each its width) in one call, at the decode's and the prompt's pitch == each alone. +fn splitCheck(g: *const Gpu, k: usize, ns: []const usize, bits: []const u8, t: *Tally) !void { + var ms: [affine.max_runs]Mat = undefined; + for (ns, bits, 0..) |n, b, i| ms[i] = try Mat.init(g.device, n, k, b, 64, 500 + 31 * i + n); + defer for (ms[0..ns.len]) |m| m.deinit(); + var q = ms[0].q; + var at: u32 = @intCast(ns[0]); + for (ms[1..ns.len], 0..) |m, i| { + q.runs[i] = .{ .w = m.q.w, .s = m.q.s, .b = m.q.b, .row = at, .n = m.q.n, .bits = m.q.bits, .group = m.q.group }; + at += m.q.n; + } + q.more = @intCast(ns.len - 1); + q.n = at; + const x = ref(ms[0].x()); + for ([_]usize{ at, std.mem.alignForward(usize, at, 64) }) |ld| { + const out = try buf(g.device, rows * ld * 2); + defer out.deinit(); + try g.run(dense, .{ &g.p, q, wide.Io{ .in = x, .ldx = @intCast(k), .out = ref(out), .ldo = @intCast(ld), .rows = rows } }); + for (ms[0..ns.len], 0..) |m, i| { + const r = q.run(i); + const one = try buf(g.device, rows * r.n * 2); + defer one.deinit(); + try g.run(dense, .{ &g.p, m.q, wide.Io{ .in = x, .ldx = @intCast(k), .out = ref(one), .ldo = r.n, .rows = rows } }); + for (0..rows) |y| try eql(out.slice(u16, rows * ld)[y * ld + r.row ..][0..r.n], one.slice(u16, rows * r.n)[y * r.n ..][0..r.n]); + t.pairs += rows * r.n; + } + } +} + +fn headsLaunch(e: mtl.ComputeEncoder, p: *const wide.Pipes, q: affine.Affine, h: u32, io: wide.Io, kx: u32, sums: Ref) void { + wide.headsBatch(p, e, q, h, io, kx, sums); +} + +/// A chunk's embed_q / unembed_out (wide.headsBatch): the heads in one affine_mm batch == each head's own product. +fn headsCheck(g: *const Gpu, mm: affine_mm.Pipes, bits: u8, n: usize, k: usize, t: *Tally) !void { + const nh = 4; + const m = 33; // two 16-row steps and one more + const ldx = nh * k + 96; // a row's slices, then columns no head reads (qp's indexer queries) + const w = try Mat.init(g.device, nh * n, k, bits, 64, 4100 + n + bits); + defer w.deinit(); + var prng = std.Random.DefaultPrng.init(4200 + k); + const r = prng.random(); + const bs = [_]mtl.Buffer{ try buf(g.device, m * ldx * 2), try buf(g.device, m * nh * (k / 64) * 4), try buf(g.device, m * nh * n * 2), try buf(g.device, m * k * 2), try buf(g.device, m * n * 2) }; + defer for (bs) |b| b.deinit(); + for (bs[0].slice(u16, m * ldx)) |*v| v.* = toBf((r.float(f32) - 0.5) * 2); + try g.run(headsLaunch, .{ &g.p, w.q, @as(u32, nh), wide.Io{ .in = ref(bs[0]), .ldx = @intCast(ldx), .out = ref(bs[2]), .ldo = @intCast(nh * n), .rows = m }, @as(u32, @intCast(k)), ref(bs[1]) }); + const all = bs[2].slice(u16, m * nh * n); + for (0..nh) |h| { + for (0..m) |i| @memcpy(bs[3].slice(u16, m * k)[i * k ..][0..k], bs[0].slice(u16, m * ldx)[i * ldx + h * k ..][0..k]); + try g.run(mmDense, .{ mm, w.part(h, nh), ref(bs[3]), ref(bs[1]), ref(bs[4]), @as(u32, m) }); + for (0..m) |i| try eql(all[(i * nh + h) * n ..][0..n], bs[4].slice(u16, m * n)[i * n ..][0..n]); + t.pairs += m * n; + } +} + +fn denseMm(e: mtl.ComputeEncoder, p: *const wide.Pipes, q: affine.Affine, io: wide.Io, sums: Ref) void { + wide.denseMm(p, e, q, io, sums); +} + +/// A chunk's split stack (wide.denseMm) at the 64-padded pitch == each member alone (affine_mm or the row kernels). +fn splitMmCheck(g: *const Gpu, k: usize, ns: []const usize, bits: []const u8, t: *Tally) !usize { + const m = 33; + var ms: [affine.max_runs]Mat = undefined; + for (ns, bits, 0..) |n, b, i| ms[i] = try Mat.init(g.device, n, k, b, 64, 900 + 37 * i + n); + defer for (ms[0..ns.len]) |x| x.deinit(); + var q = ms[0].q; + var at: u32 = @intCast(ns[0]); + for (ms[1..ns.len], 0..) |x, i| { + q.runs[i] = .{ .w = x.q.w, .s = x.q.s, .b = x.q.b, .row = at, .n = x.q.n, .bits = x.q.bits, .group = x.q.group }; + at += x.q.n; + } + q.more = @intCast(ns.len - 1); + q.n = at; + var prng = std.Random.DefaultPrng.init(9100 + k); + const r = prng.random(); + const ld = std.mem.alignForward(usize, at, 64); + const bs = [_]mtl.Buffer{ try buf(g.device, m * k * 2), try buf(g.device, m * (k / 64) * 4), try buf(g.device, m * ld * 2), try buf(g.device, m * ld * 2) }; + defer for (bs) |b| b.deinit(); + for (bs[0].slice(u16, m * k)) |*v| v.* = toBf((r.float(f32) - 0.5) * 2); + try g.run(denseMm, .{ &g.p, q, wide.Io{ .in = ref(bs[0]), .ldx = @intCast(k), .out = ref(bs[2]), .ldo = @intCast(ld), .rows = m }, ref(bs[1]) }); + var on_mm: usize = 0; + for (ms[0..ns.len], 0..) |x, i| { + const run = q.run(i); + if (wide.mm(&g.p, x.q) != null and run.n % 64 == 0) { + try g.run(mmDense, .{ wide.mm(&g.p, x.q).?, x.q, ref(bs[0]), ref(bs[1]), ref(bs[3]), @as(u32, m) }); + on_mm += 1; + } else try g.run(dense, .{ &g.p, x.q, wide.Io{ .in = ref(bs[0]), .ldx = @intCast(k), .out = ref(bs[3]), .ldo = run.n, .rows = m } }); + for (0..m) |y| try eql(bs[2].slice(u16, m * ld)[y * ld + run.row ..][0..run.n], bs[3].slice(u16, m * ld)[y * run.n ..][0..run.n]); + t.pairs += m * run.n; + } + return on_mm; +} + +pub fn gluLaunch(e: mtl.ComputeEncoder, p: *const wide.Pipes, gu: affine.Affine, n: u32, x: Ref, out: Ref, m: u32) void { + wide.glu(p, e, gu, n, x, out, m, 0.75); +} + +/// The shared expert's gate and up at one width (wide.glu, one launch) == glm_swiglu(dense [gate; up]) bit for bit. +fn sharedGluCheck(g: *const Gpu) !usize { + var pairs: usize = 0; + for ([_]u8{ 8, 5, 4, 2 }) |bits| { + const n = 2048; + const m = try Mat.init(g.device, 2 * n, 4096, bits, 64, 8800 + @as(u64, bits)); + defer m.deinit(); + const bs = [_]mtl.Buffer{ try buf(g.device, rows * 2 * n * 2), try buf(g.device, rows * n * 2), try buf(g.device, rows * n * 2) }; + defer for (bs) |b| b.deinit(); + try g.run(dense, .{ &g.p, m.q, wide.Io{ .in = ref(m.x()), .ldx = 4096, .out = ref(bs[0]), .ldo = 2 * n, .rows = rows } }); + try g.run(swiglu, .{ g.swiglu, ref(bs[0]), ref(bs[1]), @as(u32, rows), @as(u32, n), @as(f32, 0.75) }); + try g.run(gluLaunch, .{ &g.p, m.q, @as(u32, n), ref(m.x()), ref(bs[2]), @as(u32, rows) }); + try eql(bs[2].slice(u16, rows * n), bs[1].slice(u16, rows * n)); + pairs += rows * n; + } + std.debug.print("shared expert gate|up at one width in one launch == glm_swiglu(dense [gate; up]): {d} bit-equal values\n", .{pairs}); + return pairs; +} + +fn perHead(e: mtl.ComputeEncoder, p: *const wide.Pipes, q: affine.Affine, heads: u32, io: wide.Io, hsx: u32, hso: u32) void { + wide.perHead(p, e, q, heads, io, hsx, hso); +} +pub fn byIds(e: mtl.ComputeEncoder, p: *const wide.Pipes, m: wide.Picks, x: Ref, ids: Ref, out: Ref, slots: u32) void { + wide.byIds(p, e, m, x, ids, out, slots); +} +fn bySorted(e: mtl.ComputeEncoder, p: *const wide.Pipes, m: wide.Picks, x: Ref, offsets: Ref, order: Ref, out: Ref, ldo: u32) void { + wide.bySorted(p, e, m, x, offsets, order, out, ldo); +} +fn swiglu(e: mtl.ComputeEncoder, pipe: mtl.Pipeline, gu: Ref, act: Ref, n_rows: u32, width: u32, lim: f32) void { + e.setPipeline(pipe); + e.setBuffer(gu.buf, gu.off, 0); + e.setBuffer(act.buf, act.off, 1); + e.setValue(extern struct { rows: i32, width: i32, x_stride: i32, y_stride: i32, eps: f32 }{ .rows = @intCast(n_rows), .width = @intCast(width), .x_stride = @intCast(2 * width), .y_stride = @intCast(width), .eps = lim }, 2); + e.dispatchThreads(mtl.Size.of(width, n_rows, 1), mtl.Size.of(256, 1, 1)); +} + +pub fn main() !void { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + const glue = try mtl.Library.fromSource(device, sources.glm_glue, mtl.CompileOptions.mlx()); + defer glue.deinit(); + var g = Gpu{ .device = device, .queue = queue, .p = try wide.load(a, device), .core = try row.Pipelines.load(device), .swiglu = try mtl.Pipeline.init(device, glue, "glm_swiglu", false) }; + defer g.p.deinit(); + defer g.core.deinit(); + defer g.swiglu.deinit(); + std.debug.print("GPU {s}, tensor units: {}\n", .{ device.name(), device.tensorUnits() }); + var all = Tally{}; + const shapes = [_][2]usize{ .{ 132, 384 }, .{ 64, 2688 }, .{ 40, 1536 } }; + for (affine.row_bits) |bits| for ([_]usize{ 32, 64, 128 }) |group| { + var t = Tally{}; + for (shapes, 0..) |s, si| { + const m = try Mat.init(device, s[0], s[1], bits, group, 1000 * @as(u64, bits) + 10 * group + si); + defer m.deinit(); + try denseCheck(&g, m, &t); + if (s[0] == 64) { + const up = try Mat.init(device, s[0], s[1], bits, group, 7 + 1000 * @as(u64, bits) + 10 * group + si); + defer up.deinit(); + try partsCheck(&g, m, up, &t); + } + } + // oQ2's routed experts (and oQ4's at 4 bits): four experts [2048, 4096] (gate/up) and [4096, 2048] (down) + if (group == 64 and (bits == 2 or bits == 4)) for ([_][2]usize{ .{ 4 * 2048, 4096 }, .{ 4 * 4096, 2048 } }, 0..) |s, si| { + const m = try Mat.init(device, s[0], s[1], bits, group, 77 + 1000 * @as(u64, bits) + si); + defer m.deinit(); + const up = try Mat.init(device, s[0], s[1], bits, group, 78 + 1000 * @as(u64, bits) + si); + defer up.deinit(); + try partsCheck(&g, m, up, &t); + if (bits == 2) try gatherCheck(&g, m, up, &t); // the prompt's 2/64 experts on affine_mm + }; + std.debug.print("q{d} g{d}: {d} values vs fp64, worst err/bound {d:.4}; {d} bit-equal pairs (windows, strides, heads, picks, sorted, core)\n", .{ bits, group, t.values, t.worst, t.pairs }); + all.values += t.values; + all.pairs += t.pairs; + all.worst = @max(all.worst, t.worst); + }; + var split = Tally{}; + try splitCheck(&g, 256, &.{ 64, 64, 64, 16, 16, 8 }, &.{ 5, 8, 5, 8, 8, 8 }, &split); // oQ4 layer 6's KDA input + try splitCheck(&g, 512, &.{ 48, 16, 8, 8 }, &.{ 6, 6, 5, 5 }, &split); // oQ4 layer 7's and the MTP's MLA x_proj + std.debug.print("split stacks through one call == each member alone: {d} bit-equal values\n", .{split.pairs}); + var sm = Tally{}; + const kda_mm = try splitMmCheck(&g, 256, &.{ 128, 128, 128, 64, 64, 64 }, &.{ 5, 8, 5, 8, 8, 8 }, &sm); // oQ4 layer 6 + const x_mm = try splitMmCheck(&g, 512, &.{ 192, 64, 64, 32 }, &.{ 6, 6, 5, 5 }, &sm); // layer 7's x_proj: 32 on rows + std.debug.print("prompt split stacks == each member alone ({d} + {d} members on affine_mm): {d} bit-equal values\n", .{ kda_mm, x_mm, sm.pairs }); + all.pairs += sm.pairs; + var hd = Tally{}; + for ([_][2]usize{ .{ 512, 256 }, .{ 256, 512 } }) |s| try headsCheck(&g, g.p.mm8, 8, s[0], s[1], &hd); // embed_q, unembed_out + std.debug.print("prompt embed_q / unembed_out q8: the heads in one affine_mm batch == each head alone: {d} bit-equal values\n", .{hd.pairs}); + all.pairs += hd.pairs; + all.pairs += split.pairs; + all.pairs += try @import("glm_rows_kda.zig").check(&g.p, device, queue); + all.pairs += try @import("glm_rows_kda.zig").fusedCheck(&g.p, device, queue); + all.pairs += try sharedGluCheck(&g); + all.pairs += try @import("glm_rows_moe.zig").check(&g); + all.values += try @import("glm_rows_kda.zig").bf16Check(&g.p, device, queue, &all.worst); + all.values += try @import("glm_rows_kda.zig").embedCheck(&g.p, device, queue, &all.worst); + const cl = try @import("glm_rows_classes.zig").check(&g.p, device, queue); + all.values += cl.values; + all.pairs += cl.pairs; + all.worst = @max(all.worst, cl.worst); + std.debug.print("PASS {d} values within the fp32-sum bound (worst err/bound {d:.4}), {d} bit-equal pairs\n", .{ all.values, all.worst, all.pairs }); +} diff --git a/zig/tests/glm_rows_classes.zig b/zig/tests/glm_rows_classes.zig new file mode 100644 index 0000000000..4b9f99d4f6 --- /dev/null +++ b/zig/tests/glm_rows_classes.zig @@ -0,0 +1,181 @@ +//! glm_rows_check's projection classes: each GLM product at the oQ checkpoints' widths, decode and prompt, vs fp64. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const wide = tf.glm.wide; +const affine = tf.glm.affine; +const affine_mm = tf.affine_mm; + +const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; +const rows = 16; // a decode window +const prompt_rows = 33; // a prompt chunk: two 16-row steps and one more + +/// A class's shape [n, k] (a head's for embed_q / unembed_out) and the widths oQ4, oQ2 and the 4-bit build store it at. +const Class = struct { name: []const u8, k: u32, n: u32, bits: []const u8 }; +const classes = [_]Class{ + .{ .name = "kda_in q/k/v", .k = 4096, .n = 8192, .bits = &.{ 5, 8 } }, + .{ .name = "kda_in f_a/g_a", .k = 4096, .n = 128, .bits = &.{ 5, 8 } }, + .{ .name = "kda_in b", .k = 4096, .n = 64, .bits = &.{ 5, 8 } }, + .{ .name = "kda f_b/g_b", .k = 128, .n = 8192, .bits = &.{ 5, 8 } }, + .{ .name = "kda o_proj", .k = 8192, .n = 4096, .bits = &.{ 6, 8 } }, + .{ .name = "mla q_a", .k = 4096, .n = 1536, .bits = &.{ 6, 8 } }, + .{ .name = "mla kv_a", .k = 4096, .n = 512, .bits = &.{ 6, 8 } }, + .{ .name = "indexer wk", .k = 4096, .n = 128, .bits = &.{ 5, 8 } }, + .{ .name = "indexer weights_proj", .k = 4096, .n = 32, .bits = &.{ 5, 8 } }, + .{ .name = "mla x_proj (one width)", .k = 4096, .n = 2208, .bits = &.{8} }, + .{ .name = "mla q_b", .k = 1536, .n = 16384, .bits = &.{ 6, 8 } }, + .{ .name = "indexer wq_b", .k = 1536, .n = 4096, .bits = &.{ 5, 8 } }, + .{ .name = "embed_q (a head)", .k = 256, .n = 512, .bits = &.{ 5, 8 } }, + .{ .name = "unembed_out (a head)", .k = 512, .n = 256, .bits = &.{ 5, 8 } }, + .{ .name = "mla o_proj", .k = 16384, .n = 4096, .bits = &.{ 6, 8 } }, + .{ .name = "dense gate/up", .k = 4096, .n = 12288, .bits = &.{ 5, 8 } }, + .{ .name = "dense down", .k = 12288, .n = 4096, .bits = &.{ 6, 8 } }, + .{ .name = "shared gate/up", .k = 4096, .n = 2048, .bits = &.{8} }, + .{ .name = "shared down", .k = 2048, .n = 4096, .bits = &.{8} }, + .{ .name = "routed gate/up", .k = 4096, .n = 2048, .bits = &.{ 2, 4 } }, + .{ .name = "routed down", .k = 2048, .n = 4096, .bits = &.{ 2, 4 } }, + .{ .name = "lm_head", .k = 4096, .n = 154880, .bits = &.{8} }, +}; + +fn toBf(v: f32) u16 { + const u: u32 = @bitCast(v); + return @truncate((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); +} +fn fromBf(v: u16) f64 { + return @as(f32, @bitCast(@as(u32, v) << 16)); +} + +/// Code j of a row's little-endian bit stream at `bits` (MLX's packing at every width). +pub fn code(row: []const u8, j: usize, bits: u8) u32 { + const at = j * bits; + var v: u32 = 0; + for (0..4) |b| { + if (at / 8 + b < row.len) v |= @as(u32, row[at / 8 + b]) << @intCast(8 * b); + } + return (v >> @intCast(at % 8)) & ((@as(u32, 1) << @intCast(bits)) - 1); +} + +const Tally = struct { values: usize = 0, worst: f64 = 0, worst_col: usize = 0, pairs: usize = 0, err: f64 = 0, ymax: f64 = 0, caught: usize = 0 }; + +/// fp32 sums' error per unit of |terms|: under k / 8 + 64 additions deep (row kernels and affine_mm), 2^-24 each. +pub fn gamma(k: usize) f64 { + return @as(f64, @floatFromInt(k / 8 + 64)) * 0x1p-24; +} + +/// out row i vs fp64 on every column within the fp32 sums' bound; `scale` != 1: a control, columns past it counted. +fn against(q: affine.Affine, w: []const u8, s: []const u16, b: []const u16, x: []const u16, out: []const u16, pitch: usize, i: usize, t: *Tally, scale: f64) !void { + const k: usize = q.k; + const g_n = k / q.group; + const rb = k * q.bits / 8; + const gm = gamma(k); + for (0..q.n) |c| { + var y: f64 = 0; + var terms: f64 = 0; + for (0..g_n) |g| { + const sc = scale * fromBf(s[c * g_n + g]); + const bi = fromBf(b[c * g_n + g]); + var dot: f64 = 0; + var sum: f64 = 0; + for (0..q.group) |j| { + const kk = g * q.group + j; + const xv = fromBf(x[i * k + kk]); + const cv: f64 = @floatFromInt(code(w[c * rb ..][0..rb], kk, q.bits)); + y += xv * (sc * cv + bi); + dot += @abs(xv * cv); + sum += @abs(xv); + } + terms += @abs(sc) * dot + @abs(bi) * sum; + } + const bound = std.math.ldexp(@abs(y) + gm * terms, -8) + gm * terms; + const err = @abs(fromBf(out[i * pitch + c]) - y); + if (scale != 1) { + t.caught += @intFromBool(err > bound); + continue; + } + t.err = @max(t.err, err); + t.ymax = @max(t.ymax, @abs(y)); + if (err / bound > t.worst) { + t.worst = err / bound; + t.worst_col = c; + } + if (err > bound) { + std.debug.print("col {d} row {d}: {d} vs fp64 {d} (bound {e})\n", .{ c, i, fromBf(out[i * pitch + c]), y, bound }); + return error.PastBound; + } + t.values += 1; + } +} + +fn submit(queue: mtl.Queue, comptime f: anytype, args: anytype) !void { + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + @call(.auto, f, .{e} ++ args); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.debug.print("GPU failed: {s}\n", .{msg}); + return error.GpuFailed; + } +} + +fn rowsLaunch(e: mtl.ComputeEncoder, p: *const wide.Pipes, q: affine.Affine, x: mtl.Buffer, at: usize, out: mtl.Buffer, m: u32) void { + wide.dense(p, e, q, .{ .in = (affine.Ref{ .buf = x }).at(at), .ldx = q.k, .out = .{ .buf = out }, .ldo = q.n, .rows = m }); +} + +fn mmLaunch(e: mtl.ComputeEncoder, mm: affine_mm.Pipes, q: affine.Affine, x: mtl.Buffer, sums: mtl.Buffer, out: mtl.Buffer, m: u32) void { + const xr: affine.Ref = .{ .buf = x }; + const sr: affine.Ref = .{ .buf = sums }; + affine_mm.rowSums(e, mm, 64, xr, sr, m, q.k); + affine_mm.dense(e, mm, xr, sr, q, affine.Ref{ .buf = out }, m); +} + +/// Each class at its widths (g64): decode rows (alone == windowed) and the prompt's product vs fp64, every column. +pub fn check(p: *const wide.Pipes, device: mtl.Device, queue: mtl.Queue) !Tally { + var all = Tally{}; + var prng = std.Random.DefaultPrng.init(716); + const r = prng.random(); + for (classes) |cl| for (cl.bits) |bits| { + const k: usize = cl.k; + const n: usize = cl.n; + const np = std.mem.alignForward(usize, n, 64); // the loader's padded rows (affine_mm reads whole 64-row tiles) + const w_bytes = np * k * bits / 8; + const tables = np * (k / 64); + const bufs = [_]mtl.Buffer{ try device.buffer(w_bytes, opts), try device.buffer(tables * 2, opts), try device.buffer(tables * 2, opts), try device.buffer(prompt_rows * k * 2, opts), try device.buffer(prompt_rows * np * 2, opts), try device.buffer(prompt_rows * np * 2, opts), try device.buffer(prompt_rows * (k / 64) * 4, opts) }; + defer for (bufs) |x| x.deinit(); + r.bytes(bufs[0].slice(u8, w_bytes)); + for (bufs[1].slice(u16, tables)) |*v| v.* = toBf(r.float(f32) * 0.02 + 0.001); + for (bufs[2].slice(u16, tables)) |*v| v.* = toBf((r.float(f32) - 0.5) * 0.1); + for (bufs[3].slice(u16, prompt_rows * k)) |*v| v.* = toBf((r.float(f32) - 0.5) * 2); + const q: affine.Affine = .{ .w = .{ .buf = bufs[0] }, .s = .{ .buf = bufs[1] }, .b = .{ .buf = bufs[2] }, .n = cl.n, .k = cl.k, .bits = bits, .group = 64 }; + const w = bufs[0].slice(u8, w_bytes); + const s = bufs[1].slice(u16, tables); + const b = bufs[2].slice(u16, tables); + const x = bufs[3].slice(u16, prompt_rows * k); + var t = Tally{}; + // decode: a 16-row window, then rows 0 and 15 alone + try submit(queue, rowsLaunch, .{ p, q, bufs[3], 0, bufs[4], rows }); + for ([_]usize{ 0, rows - 1 }) |i| { + try submit(queue, rowsLaunch, .{ p, q, bufs[3], i * k * 2, bufs[5], 1 }); + if (!std.mem.eql(u16, bufs[5].slice(u16, n), bufs[4].slice(u16, rows * n)[i * n ..][0..n])) return error.WindowDiffersFromRowAlone; + t.pairs += n; + } + for ([_]usize{ 0, rows - 1 }) |i| try against(q, w, s, b, x, bufs[4].slice(u16, rows * n), n, i, &t, 1); + const dec = t.worst; + try against(q, w, s, b, x, bufs[4].slice(u16, rows * n), n, 0, &t, 1.01); // scales 1% off must be seen + if (t.caught == 0) return error.ControlNotCaught; + // prompt: the chunk's product at rows 0 and 32; on the row kernels every row's bits equal the decode's + const mm: ?affine_mm.Pipes = wide.mm(p, q); // prompt.zig's (2: the routed experts' gathers) + const on_mm = mm != null; + if (mm) |pipes| try submit(queue, mmLaunch, .{ pipes, q, bufs[3], bufs[6], bufs[5], prompt_rows }) else try submit(queue, rowsLaunch, .{ p, q, bufs[3], 0, bufs[5], prompt_rows }); + t.worst = 0; + const pitch = if (on_mm) np else n; + for ([_]usize{ 0, prompt_rows - 1 }) |i| try against(q, w, s, b, x, bufs[5].slice(u16, prompt_rows * pitch), pitch, i, &t, 1); + if (!on_mm and !std.mem.eql(u16, bufs[5].slice(u16, rows * n), bufs[4].slice(u16, rows * n))) return error.PromptDiffersFromDecode; + std.debug.print("{s} [{d}, {d}] q{d}: decode worst err/bound {d:.4}, prompt ({s}) {d:.4} at col {d}; max |err| {e:.3} = {e:.3} of max |y|; scales 1% off caught in {d} of {d} columns\n", .{ cl.name, n, k, bits, dec, if (on_mm) "affine_mm" else "row kernels", t.worst, t.worst_col, t.err, t.err / t.ymax, t.caught, n }); + all.values += t.values; + all.pairs += t.pairs; + all.worst = @max(all.worst, @max(dec, t.worst)); + }; + return all; +} diff --git a/zig/tests/glm_rows_kda.zig b/zig/tests/glm_rows_kda.zig new file mode 100644 index 0000000000..70294fe562 --- /dev/null +++ b/zig/tests/glm_rows_kda.zig @@ -0,0 +1,297 @@ +//! glm_rows_check's KDA step (f_b, g_b made apart == the generated 4-bit dots, exact) and bf16 product (fp64) parts. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const wide = tf.glm.wide; +const affine = tf.glm.affine; + +const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; +const H = 64; +const D = 128; +const W = H * D; +const PS = 3 * W + 2 * D + H; +const R = 3; +const TAPS = 4; + +fn toBf(v: f32) u16 { + const u: u32 = @bitCast(v); + return @truncate((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); +} +fn fromBf(v: u16) f64 { + return @as(f32, @bitCast(@as(u32, v) << 16)); +} + +fn submit(queue: mtl.Queue, comptime f: anytype, args: anytype) !void { + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + @call(.auto, f, .{e} ++ args); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.debug.print("GPU failed: {s}\n", .{msg}); + return error.GpuFailed; + } +} + +const Bufs = struct { p: mtl.Buffer, cs: mtl.Buffer, cw: mtl.Buffer, fb: [3]mtl.Buffer, gb: [3]mtl.Buffer, a: mtl.Buffer, dtb: mtl.Buffer, st: mtl.Buffer, onw: mtl.Buffer }; +const Outs = struct { y: mtl.Buffer, st: mtl.Buffer, cs: mtl.Buffer }; + +fn step(e: mtl.ComputeEncoder, pipe: mtl.Pipeline, b: Bufs, f: [3]mtl.Buffer, g: [3]mtl.Buffer, o: Outs) void { + e.setPipeline(pipe); + e.setBuffer(b.p, 0, 0); + e.setBytes(std.mem.asBytes(&[2]i32{ R, PS }), 1); + const ins = [_]mtl.Buffer{ b.cs, b.cw, f[0], f[1], f[2], g[0], g[1], g[2], b.a, b.dtb, b.st, b.onw }; + for (ins, 2..) |x, i| e.setBuffer(x, 0, i); + e.setValue(@as(f32, -5.0), 14); + e.setValue(@as(f32, 1e-6), 15); + e.setBuffer(o.y, 0, 16); + e.setBuffer(o.st, 0, 17); + e.setBuffer(o.cs, 0, 18); + e.dispatchThreads(mtl.Size.of(32, 32, H), mtl.Size.of(32, 32, 1)); +} + +fn products(e: mtl.ComputeEncoder, p: *const wide.Pipes, fq: affine.Affine, gq: affine.Affine, proj: mtl.Buffer, f: mtl.Buffer, g: mtl.Buffer) void { + const in: affine.Ref = .{ .buf = proj }; + wide.dense(p, e, fq, .{ .in = in.at(3 * W * 2), .ldx = PS, .out = .{ .buf = f }, .ldo = W, .rows = R }); + wide.dense(p, e, gq, .{ .in = in.at((3 * W + D) * 2), .ldx = PS, .out = .{ .buf = g }, .ldo = W, .rows = R }); +} + +/// kda_rows on f_b, g_b products made apart == the generated 4-bit dots: every code, dyadic scales and biases, exact. +pub fn check(p: *const wide.Pipes, device: mtl.Device, queue: mtl.Queue) !usize { + const kernels = tf.glm.kernels; + const gen_lib = try mtl.Library.fromSource(device, comptime kernels.generatedSource("kda_rows"), mtl.CompileOptions.mlx()); + defer gen_lib.deinit(); + const generated = try mtl.Pipeline.init(device, gen_lib, comptime kernels.generatedFunction("kda_rows"), false); + defer generated.deinit(); + var prng = std.Random.DefaultPrng.init(4242); + const r = prng.random(); + const mk = struct { + fn of(dev: mtl.Device, bytes: usize) !mtl.Buffer { + return dev.buffer(bytes, opts); + } + }.of; + var b: Bufs = undefined; + b.p = try mk(device, R * PS * 2); + b.cs = try mk(device, (TAPS - 1) * 3 * W * 2); + b.cw = try mk(device, TAPS * 3 * W * 4); + for (&b.fb, &b.gb, [_]usize{ W * D / 2, W * 2 * 2, W * 2 * 2 }) |*f, *g, n| { + f.* = try mk(device, n); + g.* = try mk(device, n); + } + b.a = try mk(device, H * 4); + b.dtb = try mk(device, W * 4); + b.st = try mk(device, H * D * D * 4); + b.onw = try mk(device, D * 4); + defer for ([_]mtl.Buffer{ b.p, b.cs, b.cw, b.fb[0], b.fb[1], b.fb[2], b.gb[0], b.gb[1], b.gb[2], b.a, b.dtb, b.st, b.onw }) |x| x.deinit(); + const pv = b.p.slice(u16, R * PS); + for (pv, 0..) |*v, i| { + const c = i % PS; + const small = (c >= 3 * W and c < 3 * W + 2 * D); // f_a and g_a: small integers, the dots exact in any order + v.* = toBf(if (small) @floatFromInt(r.intRangeAtMost(i32, -2, 2)) else (r.float(f32) - 0.5) * 2); + } + for (b.cs.slice(u16, (TAPS - 1) * 3 * W)) |*v| v.* = toBf((r.float(f32) - 0.5) * 2); + for (b.cw.slice(f32, TAPS * 3 * W)) |*v| v.* = (r.float(f32) - 0.5) * 0.5; + for ([_]mtl.Buffer{ b.fb[0], b.gb[0] }) |x| for (x.slice(u8, W * D / 2)) |*v| { + v.* = r.int(u8); // every 4-bit code + }; + for ([_]mtl.Buffer{ b.fb[1], b.gb[1] }) |x| for (x.slice(u16, W * 2)) |*v| { + v.* = toBf(std.math.ldexp(@as(f32, 1), -r.intRangeAtMost(i32, 5, 8))); // a scale a group + }; + for ([_]mtl.Buffer{ b.fb[2], b.gb[2] }) |x| for (x.slice(u16, W * 2)) |*v| { + const m = r.intRangeAtMost(i32, 1, 8); + v.* = toBf(@as(f32, @floatFromInt(if (r.boolean()) m else -m)) / 64); // a nonzero bias a group + }; + for (b.a.slice(f32, H)) |*v| v.* = -r.float(f32); + for (b.dtb.slice(f32, W)) |*v| v.* = (r.float(f32) - 0.5) * 0.1; + for (b.st.slice(f32, H * D * D)) |*v| v.* = (r.float(f32) - 0.5) * 0.1; + for (b.onw.slice(f32, D)) |*v| v.* = 1 + (r.float(f32) - 0.5) * 0.1; + var outs: [2]Outs = undefined; + for (&outs) |*o| o.* = .{ .y = try mk(device, R * W * 2), .st = try mk(device, H * D * D * 4), .cs = try mk(device, (TAPS - 1) * 3 * W * 2) }; + defer for (outs) |o| for ([_]mtl.Buffer{ o.y, o.st, o.cs }) |x| x.deinit(); + try submit(queue, step, .{ generated, b, b.fb, b.gb, outs[0] }); // the generated 4-bit dots + const fq: affine.Affine = .{ .w = .{ .buf = b.fb[0] }, .s = .{ .buf = b.fb[1] }, .b = .{ .buf = b.fb[2] }, .n = W, .k = D, .bits = 4, .group = 64 }; + var gq = fq; + gq.w = .{ .buf = b.gb[0] }; + gq.s = .{ .buf = b.gb[1] }; + gq.b = .{ .buf = b.gb[2] }; + const f = try mk(device, 16 * W * 2); + defer f.deinit(); + const g = try mk(device, 16 * W * 2); + defer g.deinit(); + try submit(queue, products, .{ p, fq, gq, b.p, f, g }); + try submit(queue, step, .{ p.kda_pre, b, [3]mtl.Buffer{ f, f, f }, [3]mtl.Buffer{ g, g, g }, outs[1] }); + var same: usize = 0; + for ([_][2][]const u8{ .{ std.mem.sliceAsBytes(outs[0].y.slice(u16, R * W)), std.mem.sliceAsBytes(outs[1].y.slice(u16, R * W)) }, .{ outs[0].st.slice(u8, H * D * D * 4), outs[1].st.slice(u8, H * D * D * 4) }, .{ outs[0].cs.slice(u8, (TAPS - 1) * 3 * W * 2), outs[1].cs.slice(u8, (TAPS - 1) * 3 * W * 2) } }) |pair| { + if (!std.mem.eql(u8, pair[0], pair[1])) return error.KdaPreDiffers; + same += pair[0].len; + } + var nonzero: usize = 0; + for (outs[0].y.slice(u16, R * W)) |v| nonzero += @intFromBool(v & 0x7fff != 0); + if (nonzero < R * W / 2) return error.KdaOutputsMostlyZero; + std.debug.print("kda_rows on f_b, g_b made apart == the generated 4-bit dots: {d} rows, outputs, state and conv window byte-equal ({d} bytes)\n", .{ R, same }); + return R * W; +} + +/// kda_rows with f_b, g_b made inside (`.fused`) == made apart (`.pre`) at every width, g64: Y, state, conv equal. +pub fn fusedCheck(p: *const wide.Pipes, device: mtl.Device, queue: mtl.Queue) !usize { + var prng = std.Random.DefaultPrng.init(4343); + const r = prng.random(); + const mk = struct { + fn of(dev: mtl.Device, bytes: usize) !mtl.Buffer { + return dev.buffer(bytes, opts); + } + }.of; + var b: Bufs = undefined; + b.p = try mk(device, R * PS * 2); + b.cs = try mk(device, (TAPS - 1) * 3 * W * 2); + b.cw = try mk(device, TAPS * 3 * W * 4); + b.a = try mk(device, H * 4); + b.dtb = try mk(device, W * 4); + b.st = try mk(device, H * D * D * 4); + b.onw = try mk(device, D * 4); + const f = try mk(device, 16 * W * 2); + const g = try mk(device, 16 * W * 2); + defer for ([_]mtl.Buffer{ b.p, b.cs, b.cw, b.a, b.dtb, b.st, b.onw, f, g }) |x| x.deinit(); + for (b.p.slice(u16, R * PS)) |*v| v.* = toBf((r.float(f32) - 0.5) * 2); + for (b.cs.slice(u16, (TAPS - 1) * 3 * W)) |*v| v.* = toBf((r.float(f32) - 0.5) * 2); + for (b.cw.slice(f32, TAPS * 3 * W)) |*v| v.* = (r.float(f32) - 0.5) * 0.5; + for (b.a.slice(f32, H)) |*v| v.* = -r.float(f32); + for (b.dtb.slice(f32, W)) |*v| v.* = (r.float(f32) - 0.5) * 0.1; + for (b.st.slice(f32, H * D * D)) |*v| v.* = (r.float(f32) - 0.5) * 0.1; + for (b.onw.slice(f32, D)) |*v| v.* = 1 + (r.float(f32) - 0.5) * 0.1; + var outs: [2]Outs = undefined; + for (&outs) |*o| o.* = .{ .y = try mk(device, R * W * 2), .st = try mk(device, H * D * D * 4), .cs = try mk(device, (TAPS - 1) * 3 * W * 2) }; + defer for (outs) |o| for ([_]mtl.Buffer{ o.y, o.st, o.cs }) |x| x.deinit(); + var same: usize = 0; + for (affine.row_bits, 0..) |bits, i| { + var wq: [2]affine.Affine = undefined; + var wb: [2][3]mtl.Buffer = undefined; + for (&wq, &wb) |*q, *x| { + x.* = .{ try mk(device, W * D * @as(usize, bits) / 8), try mk(device, W * 2 * 2), try mk(device, W * 2 * 2) }; + r.bytes(x[0].slice(u8, W * D * @as(usize, bits) / 8)); // every code + for (x[1].slice(u16, W * 2)) |*v| v.* = toBf(r.float(f32) * 0.02 + 0.001); + for (x[2].slice(u16, W * 2)) |*v| v.* = toBf((r.float(f32) - 0.5) * 0.1); + q.* = .{ .w = .{ .buf = x[0] }, .s = .{ .buf = x[1] }, .b = .{ .buf = x[2] }, .n = W, .k = D, .bits = bits, .group = 64 }; + } + defer for (wb) |x| for (x) |y| y.deinit(); + if (bits != 4 and wide.kdaGate(p, wq[0], wq[1]) != .fused) return error.NotFused; // 4/64 takes the generated dots + try submit(queue, products, .{ p, wq[0], wq[1], b.p, f, g }); + try submit(queue, step, .{ p.kda_pre, b, [3]mtl.Buffer{ f, f, f }, [3]mtl.Buffer{ g, g, g }, outs[0] }); + try submit(queue, step, .{ p.kda_fg[i], b, wb[0], wb[1], outs[1] }); + for ([_][2][]const u8{ .{ std.mem.sliceAsBytes(outs[0].y.slice(u16, R * W)), std.mem.sliceAsBytes(outs[1].y.slice(u16, R * W)) }, .{ outs[0].st.slice(u8, H * D * D * 4), outs[1].st.slice(u8, H * D * D * 4) }, .{ outs[0].cs.slice(u8, (TAPS - 1) * 3 * W * 2), outs[1].cs.slice(u8, (TAPS - 1) * 3 * W * 2) } }) |pair| { + if (!std.mem.eql(u8, pair[0], pair[1])) { + std.debug.print("kda_rows fused q{d} differs from .pre\n", .{bits}); + return error.KdaFusedDiffers; + } + same += pair[0].len; + } + var nonzero: usize = 0; + for (outs[1].y.slice(u16, R * W)) |v| nonzero += @intFromBool(v & 0x7fff != 0); + if (nonzero < R * W / 2) return error.KdaOutputsMostlyZero; + } + std.debug.print("kda_rows with f_b, g_b made inside == made apart (.pre) at q2..q8 g64: {d} rows, outputs, state and conv window byte-equal ({d} bytes)\n", .{ R, same }); + return same / 4; +} + +/// The bf16 product (an unquantized eh_proj) against fp64 per column, and a row alone == the row in a window. +pub fn bf16Check(p: *const wide.Pipes, device: mtl.Device, queue: mtl.Queue, worst: *f64) !usize { + const n = 40; + const k = 8192; + const rows = 5; + var prng = std.Random.DefaultPrng.init(99); + const r = prng.random(); + const w = try device.buffer(n * k * 2, opts); + defer w.deinit(); + const x = try device.buffer(rows * k * 2, opts); + defer x.deinit(); + for (w.slice(u16, n * k)) |*v| v.* = toBf((r.float(f32) - 0.5) * 0.05); + for (x.slice(u16, rows * k)) |*v| v.* = toBf((r.float(f32) - 0.5) * 2); + const all = try device.buffer(rows * n * 2, opts); + defer all.deinit(); + const one = try device.buffer(n * 2, opts); + defer one.deinit(); + const run = struct { + fn f(e: mtl.ComputeEncoder, pp: *const wide.Pipes, wb: mtl.Buffer, in: affine.Ref, out: mtl.Buffer, m: u32) void { + wide.bf16(pp, e, .{ .buf = wb }, n, k, .{ .in = in, .ldx = k, .out = .{ .buf = out }, .ldo = n, .rows = m }); + } + }.f; + try submit(queue, run, .{ p, w, affine.Ref{ .buf = x }, all, rows }); + var values: usize = 0; + for (0..rows) |i| { + try submit(queue, run, .{ p, w, (affine.Ref{ .buf = x }).at(i * k * 2), one, 1 }); + if (!std.mem.eql(u16, one.slice(u16, n), all.slice(u16, rows * n)[i * n ..][0..n])) return error.Bf16WindowDiffers; + for (0..n) |c| { + var y: f64 = 0; + var terms: f64 = 0; + for (0..k) |j| { + const t = fromBf(x.slice(u16, rows * k)[i * k + j]) * fromBf(w.slice(u16, n * k)[c * k + j]); + y += t; + terms += @abs(t); + } + const bound = std.math.ldexp(@abs(y), -8) + @as(f64, k + 32) * 0x1p-24 * terms; + const err = @abs(fromBf(one.slice(u16, n)[c]) - y); + worst.* = @max(worst.*, err / bound); + if (err > bound) return error.Bf16PastBound; + values += 1; + } + } + std.debug.print("bf16 [{d}, {d}]: {d} values vs fp64 within the fp32-sum bound; each row alone == the row in a {d}-row window\n", .{ n, k, values, rows }); + return values; +} + +/// The 6- and 8-bit embedding's gathers vs fp64 within two bf16 roundings; a neighbouring group's bias must be seen. +pub fn embedCheck(p: *const wide.Pipes, device: mtl.Device, queue: mtl.Queue, worst: *f64) !usize { + const V = 40; + const width = 4096; + const ids = [_]u32{ 7, 0, 39, 7, 21 }; + var prng = std.Random.DefaultPrng.init(808); + const r = prng.random(); + var values: usize = 0; + for ([_]u8{ 6, 8 }, 0..) |bits, pi| { + const rb: usize = width * @as(usize, bits) / 8; + const G = width / 64; + const t = try device.buffer(V * rb, opts); + defer t.deinit(); + const s = try device.buffer(V * G * 2, opts); + defer s.deinit(); + const b = try device.buffer(V * G * 2, opts); + defer b.deinit(); + const idb = try device.buffer(ids.len * 4, opts); + defer idb.deinit(); + const out = try device.buffer(ids.len * width * 2, opts); + defer out.deinit(); + r.bytes(t.slice(u8, V * rb)); + for (s.slice(u16, V * G)) |*v| v.* = toBf(r.float(f32) * 0.02 + 0.001); + for (b.slice(u16, V * G)) |*v| v.* = toBf((r.float(f32) - 0.5) * 0.2); + @memcpy(idb.slice(u32, ids.len), &ids); + const run = struct { + fn f(e: mtl.ComputeEncoder, pipe: mtl.Pipeline, bs: [5]mtl.Buffer, per: u32) void { + e.setPipeline(pipe); + for (bs[0..4], 0..) |x, i| e.setBuffer(x, 0, i); + e.setBytes(std.mem.asBytes(&[1]i32{width}), 4); + e.setBuffer(bs[4], 0, 5); + e.dispatchThreads(mtl.Size.of(width / per, ids.len, 1), mtl.Size.of(256, 1, 1)); + } + }.f; + try submit(queue, run, .{ p.embed[pi], [5]mtl.Buffer{ idb, t, s, b, out }, @as(u32, if (bits == 6) 4 else 1) }); + var caught: usize = 0; + for (ids, 0..) |id, i| for (0..width) |j| { + const g = j / 64; + const q: f64 = @floatFromInt(@import("glm_rows_classes.zig").code(t.slice(u8, V * rb)[id * rb ..][0..rb], j, bits)); + const sq = fromBf(s.slice(u16, V * G)[id * G + g]) * q; + const bias = fromBf(b.slice(u16, V * G)[id * G + g]); + const got = fromBf(out.slice(u16, ids.len * width)[i * width + j]); + const bound = std.math.ldexp(2 * @abs(sq) + @abs(bias), -8) + std.math.ldexp(@abs(sq), -16); + const err = @abs(got - (sq + bias)); + worst.* = @max(worst.*, err / bound); + if (err > bound) return error.EmbedPastBound; + const other = fromBf(b.slice(u16, V * G)[id * G + (g ^ 1)]); + caught += @intFromBool(@abs(got - (sq + other)) > bound); + values += 1; + }; + if (caught == 0) return error.ControlNotCaught; + std.debug.print("embedding q{d} g64: {d} gathered values vs fp64 within two bf16 roundings; a neighbour group's bias caught in {d}\n", .{ bits, ids.len * width, caught }); + } + return values; +} diff --git a/zig/tests/glm_rows_moe.zig b/zig/tests/glm_rows_moe.zig new file mode 100644 index 0000000000..41de7fbe8b --- /dev/null +++ b/zig/tests/glm_rows_moe.zig @@ -0,0 +1,51 @@ +//! glm_rows_check's MoE launches: the routed and shared experts' two-launch path (wide.moe) == its four launches. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const wide = tf.glm.wide; +const affine = tf.glm.affine; +const rc = @import("glm_rows_check.zig"); + +fn moeLaunch(e: mtl.ComputeEncoder, p: *const wide.Pipes, m: wide.Moe, io: wide.MoeIo, n_rows: u32) void { + wide.moe(p, e, m, 4, 8, 0.75, io, n_rows); +} + +/// wide.moe (routed at q2..q8, shared at 8, two launches) == byIds GLU/down + glu/dense bit for bit; id 5 of 4: zeros. +pub fn check(g: *const rc.Gpu) !usize { + const n = 256; + const k = 512; + const n_rows = 2; + const slots = 8 * n_rows; + var pairs: usize = 0; + const sh = [2]rc.Mat{ try rc.Mat.init(g.device, 2 * n, k, 8, 64, 7300), try rc.Mat.init(g.device, k, n, 8, 64, 7301) }; + defer for (sh) |m| m.deinit(); + for (affine.row_bits) |bits| { + const ms = [3]rc.Mat{ try rc.Mat.init(g.device, 4 * n, k, bits, 64, 7400 + @as(u64, bits)), try rc.Mat.init(g.device, 4 * n, k, bits, 64, 7500 + @as(u64, bits)), try rc.Mat.init(g.device, 4 * k, n, bits, 64, 7600 + @as(u64, bits)) }; + defer for (ms) |m| m.deinit(); + var q: [3]affine.Affine = undefined; + for (&q, ms, [_]u32{ n, n, k }) |*x, m, r| { + x.* = m.q; + x.n = r; + } + const bs = [_]mtl.Buffer{ try rc.buf(g.device, slots * 4), try rc.buf(g.device, 2 * slots * n * 2), try rc.buf(g.device, 2 * n_rows * n * 2), try rc.buf(g.device, 2 * slots * k * 2), try rc.buf(g.device, 2 * n_rows * k * 2) }; + defer for (bs) |b| b.deinit(); + const picks = [slots]i32{ 2, 0, 3, 1, 1, 2, 0, 3, 3, 3, 0, 5, 1, 2, 2, 0 }; + @memcpy(bs[0].slice(i32, slots), &picks); + const m: wide.Moe = .{ .gate = q[0], .up = q[1], .down = q[2], .sh_gate_up = sh[0].q, .sh_down = sh[1].q }; + if (!wide.moeOne(m) and bits != 4) return error.NotOneLaunch; // 4/64 routed keeps the generated gathers + const io: wide.MoeIo = .{ .x = rc.ref(ms[0].x()), .ids = rc.ref(bs[0]), .act = rc.ref(bs[1]), .acts = rc.ref(bs[2]), .ye = rc.ref(bs[3]), .ys = rc.ref(bs[4]) }; + try g.run(moeLaunch, .{ &g.p, m, io, @as(u32, n_rows) }); + const half = [5]usize{ 0, slots * n, n_rows * n, slots * k, n_rows * k }; // the reference in each buffer's second half + try g.run(rc.byIds, .{ &g.p, wide.Picks{ .w = q[0], .up = q[1], .experts = 4, .repeat = 8, .lim = 0.75 }, io.x, io.ids, io.act.at(half[1] * 2), @as(u32, slots) }); + try g.run(rc.gluLaunch, .{ &g.p, sh[0].q, @as(u32, n), io.x, io.acts.at(half[2] * 2), @as(u32, n_rows) }); + try g.run(rc.byIds, .{ &g.p, wide.Picks{ .w = q[2], .experts = 4, .repeat = 1 }, io.act, io.ids, io.ye.at(half[3] * 2), @as(u32, slots) }); + try g.run(rc.dense, .{ &g.p, sh[1].q, wide.Io{ .in = io.acts, .ldx = n, .out = io.ys.at(half[4] * 2), .ldo = k, .rows = n_rows } }); + for (bs[1..], half[1..]) |b, h| { + const v = b.slice(u16, 2 * h); + try rc.eql(v[0..h], v[h..]); + pairs += h; + } + } + std.debug.print("routed experts at q2..q8 and the 8-bit shared expert in two launches == byIds, glu and dense: {d} bit-equal values\n", .{pairs}); + return pairs; +} diff --git a/zig/tests/glm_rows_speed.metal b/zig/tests/glm_rows_speed.metal new file mode 100644 index 0000000000..561a5a3772 --- /dev/null +++ b/zig/tests/glm_rows_speed.metal @@ -0,0 +1,44 @@ +// tf-glm-rows-speed's candidates: glm_rows's one-row sums at RPS columns a simdgroup, SG a group, K compiled in. +template +[[kernel]] void glm_rows_x(const device bfloat* X [[buffer(0)]], const device uint32_t* W [[buffer(1)]], + const device bfloat* S [[buffer(2)]], const device bfloat* B [[buffer(3)]], + constant GlmRowDims& d [[buffer(4)]], device bfloat* OUT [[buffer(5)]], + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int row0 = (int(tg.x) * SG + int(sg)) * RPS; + if (row0 >= d.n) return; + const int k = K > 0 ? K : d.k; + const int h = int(tg.y); + const size_t at = size_t(h) * d.hsw + row0; + const device uint8_t* wr = (const device uint8_t*)W + at * (k * BITS / 8); + const device bfloat* sr = S + at * (k / 64); + const device bfloat* br = B + at * (k / 64); + const device bfloat* x0 = X + size_t(h) * d.hsx; + device bfloat* o0 = OUT + size_t(h) * d.hso + row0; + const int cols = metal::min(RPS, d.n - row0); + #pragma clang loop unroll(disable) + for (int m = 0; m < d.rows; m++) { + float acc[1][RPS]; + tf_rp_rows(wr, sr, br, x0 + size_t(m) * d.ldx, k, 64, lane, acc); + if (lane == 0) + for (int j = 0; j < cols; j++) o0[size_t(m) * d.ldo + j] = bfloat(acc[0][j]); + } +} + +#define GX(B, R, S, K) \ + template [[host_name("x_q" #B "_r" #R "_s" #S "_k" #K)]] [[kernel]] decltype(glm_rows_x) glm_rows_x; +#define GXK(B, K) GX(B, 4, 2, K) GX(B, 8, 2, K) GX(B, 4, 4, K) GX(B, 8, 4, K) GX(B, 2, 4, K) GX(B, 2, 8, K) GX(B, 4, 8, K) +#define GXB(B) GXK(B, 0) GXK(B, 4096) GXK(B, 8192) GXK(B, 16384) GXK(B, 12288) GXK(B, 1536) GXK(B, 2048) GXK(B, 128) GXK(B, 256) GXK(B, 512) +GXB(8) + + +// glm_moe_rows (glm_rows.metal) at other columns a simdgroup and with k compiled in: oQ2's routed q2 + shared q8. +#define GXM(R, K) \ + template [[host_name("x_moe_glu_r" #R "_k" #K)]] [[kernel]] decltype(glm_moe_rows<2, 8, true, R, K>) glm_moe_rows<2, 8, true, R, K>; \ + template [[host_name("x_moe_r" #R "_k" #K)]] [[kernel]] decltype(glm_moe_rows<2, 8, false, R, K>) glm_moe_rows<2, 8, false, R, K>; +GXM(8, 0) +GXM(2, 0) +template [[host_name("x_moe_glu_r4_k4096")]] [[kernel]] decltype(glm_moe_rows<2, 8, true, 4, 4096>) glm_moe_rows<2, 8, true, 4, 4096>; +template [[host_name("x_moe_glu_r8_k4096")]] [[kernel]] decltype(glm_moe_rows<2, 8, true, 8, 4096>) glm_moe_rows<2, 8, true, 8, 4096>; +template [[host_name("x_moe_r4_k2048")]] [[kernel]] decltype(glm_moe_rows<2, 8, false, 4, 2048>) glm_moe_rows<2, 8, false, 4, 2048>; +template [[host_name("x_moe_r8_k2048")]] [[kernel]] decltype(glm_moe_rows<2, 8, false, 8, 2048>) glm_moe_rows<2, 8, false, 8, 2048>; diff --git a/zig/tests/glm_rows_speed.zig b/zig/tests/glm_rows_speed.zig new file mode 100644 index 0000000000..788e3e1f50 --- /dev/null +++ b/zig/tests/glm_rows_speed.zig @@ -0,0 +1,598 @@ +//! tf-glm-rows-speed: GLM-5.3-Flash's decode products at the checkpoints' widths and shapes, GPU time and GB/s. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const sources = @import("kernel_sources"); +const wide = tf.glm.wide; +const affine = tf.glm.affine; +const kernels = tf.glm.kernels; +const Ref = affine.Ref; +const Affine = affine.Affine; + +const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; +const a = std.heap.page_allocator; +const hot = 768 << 20; // a class's matrices rotate over at least this many bytes: past the GPU's caches +const reps = 32; // launches a command buffer +const runs = 7; // command buffers a timing (median) + +const H = 64; // KDA and MLA heads +const KD = 128; +const W = H * KD; +const PS = 3 * W + 2 * KD + H; // KDA's stacked projection +const XP = 2208; // MLA x_proj rows +const QR = 20480; // qr_proj rows +const N = 2048; // an expert's intermediate rows +const D = 4096; +const E = 64; // experts in a pool +const TOPK = 8; + +fn buf(device: mtl.Device, bytes: usize) !mtl.Buffer { + return device.buffer(@max(bytes, 16), opts); +} + +fn fill(b: mtl.Buffer, bytes: usize, seed: u64) void { + var s = seed | 1; + for (b.slice(u64, bytes / 8)) |*v| { + s ^= s << 13; + s ^= s >> 7; + s ^= s << 17; + v.* = s; + } +} + +/// `copies` matrices [n, k] at bits/64 (experts > 1: each one stacked [experts, n, k]); codes random, tables small. +const Pool = struct { + q: []Affine, + bufs: []mtl.Buffer, + bytes: usize, // one copy's + + fn init(device: mtl.Device, n: usize, k: usize, bits: u8, experts: usize, min_copies: usize) !Pool { + const w: affine.Width = .{ .bits = bits, .group = 64 }; + const one = affine.bytes(n, k, w) * experts; + const copies = if (experts > 1) min_copies else @max(min_copies, @min(64, (hot + one - 1) / one)); // a pool's picks walk its experts + const p = Pool{ .q = try a.alloc(Affine, copies), .bufs = try a.alloc(mtl.Buffer, 3 * copies), .bytes = one }; + for (0..copies) |c| { + const wb = affine.wBytes(n, k, w) * experts; + const sb = affine.sBytes(n, k, w) * experts; + p.bufs[3 * c] = try buf(device, wb); + fill(p.bufs[3 * c], wb, 0x9e37 + c * 7919 + n + k); + for (1..3) |j| { + p.bufs[3 * c + j] = try buf(device, sb); + @memset(p.bufs[3 * c + j].slice(u16, sb / 2), if (j == 1) 0x3c00 else 0xbc00); // 2^-7, -2^-7 + } + p.q[c] = .{ .w = .{ .buf = p.bufs[3 * c] }, .s = .{ .buf = p.bufs[3 * c + 1] }, .b = .{ .buf = p.bufs[3 * c + 2] }, .n = @intCast(n), .k = @intCast(k), .bits = bits, .group = 64 }; + } + return p; + } + + fn deinit(p: Pool) void { + for (p.bufs) |b| b.deinit(); + a.free(p.q); + a.free(p.bufs); + } + + fn at(p: Pool, i: usize) Affine { + return p.q[i % p.q.len]; + } +}; + +/// Activations and outputs every launch shares (wide enough for any class). +const Io = struct { x: mtl.Buffer, y: mtl.Buffer, z: mtl.Buffer, ids: mtl.Buffer, uids: mtl.Buffer, umem: mtl.Buffer, ucount: mtl.Buffer }; + +const Bench = struct { + device: mtl.Device, + queue: mtl.Queue, + k: *kernels.Kernels, + io: Io, + reps: usize = reps, + + /// Median GPU microseconds a launch of `f(e, i, args)` over `reps` launches i in one command buffer. + fn time(b: *const Bench, comptime f: anytype, args: anytype) !f64 { + var t: [runs + 1]f64 = undefined; + for (&t) |*v| { + const cb = b.queue.commandBuffer(); + const e = cb.compute(.serial); + for (0..b.reps) |i| @call(.auto, f, .{ e, i } ++ args); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.debug.print("GPU failed: {s}\n", .{msg}); + return error.GpuFailed; + } + v.* = cb.gpuSeconds() * 1e6 / @as(f64, @floatFromInt(b.reps)); + } + const s = t[1..]; // the first warms the pipelines + std.mem.sort(f64, s, {}, std.sort.asc(f64)); + return s[runs / 2]; + } + + /// One launch of `f(e, 0, args)`, waited on. + fn run(b: *const Bench, comptime f: anytype, args: anytype) !void { + const cb = b.queue.commandBuffer(); + const e = cb.compute(.serial); + @call(.auto, f, .{ e, @as(usize, 0) } ++ args); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.debug.print("GPU failed: {s}\n", .{msg}); + return error.GpuFailed; + } + } +}; + +fn bind(e: mtl.ComputeEncoder, first: usize, refs: anytype) void { + inline for (refs, 0..) |r, i| e.setBuffer(r.buf, r.off, first + i); +} +fn shape(e: mtl.ComputeEncoder, index: usize, dims: anytype) void { + var v: [dims.len]i32 = undefined; + inline for (dims, 0..) |d, i| v[i] = @intCast(d); + e.setBytes(std.mem.asBytes(&v), index); +} +fn rf(b: mtl.Buffer) Ref { + return .{ .buf = b }; +} + +// Uniform 4/64: the generated kernels' launches (forward.zig, moe.zig). +fn qmv(e: mtl.ComputeEncoder, i: usize, pipe: mtl.Pipeline, p: Pool, io: Io) void { + const q = p.at(i); + e.setPipeline(pipe); + bind(e, 0, .{ rf(io.x), q.w, q.s, q.b, rf(io.y) }); + e.dispatchThreads(mtl.Size.of(32, q.n / 4, 1), mtl.Size.of(32, 1, 1)); +} +fn absorb(e: mtl.ComputeEncoder, i: usize, k: *const kernels.Kernels, p: Pool, io: Io) void { + const q = p.at(i); + e.setPipeline(k.absorb); + bind(e, 0, .{ q.w, q.s, q.b, rf(io.x), rf(io.y) }); + e.setValue(@as(u32, QR), 5); + e.setValue(@as(u32, H), 6); + e.dispatchGroups(mtl.Size.of(1, 512 / 64, H), mtl.Size.of(64, 1, 1)); +} +fn unabsorb(e: mtl.ComputeEncoder, i: usize, k: *const kernels.Kernels, p: Pool, io: Io) void { + const q = p.at(i); + e.setPipeline(k.unabsorb); + bind(e, 0, .{ q.w, q.s, q.b, rf(io.x), rf(io.y) }); + e.setValue(@as(u32, H), 5); + e.dispatchGroups(mtl.Size.of(1, 256 / 4, H), mtl.Size.of(32, 1, 1)); +} +/// The routed experts and the shared one in one gather (moe.zig `both`, one row: two outputs a simdgroup). +fn moeBoth(e: mtl.ComputeEncoder, i: usize, k: *const kernels.Kernels, gu: [2]Pool, down: Pool, sh: [2]Pool, io: Io) void { + const g = gu[0].at(i); + const u = gu[1].at(i); + const s = sh[0].at(i); + const ids = rf(io.uids).at((i % reps) * 128 * 4); + e.setPipeline(k.moe_gateup_0_one); + bind(e, 0, .{rf(io.x)}); + shape(e, 1, .{ 1, D }); + bind(e, 2, .{ g.w, g.s, g.b, u.w, u.s, u.b, s.w, s.s, s.b, ids, rf(io.umem), rf(io.ucount) }); + e.setValue(@as(f32, 10.0), 14); + bind(e, 15, .{rf(io.z)}); + e.dispatchThreads(mtl.Size.of(32, N / 2, TOPK + 1), mtl.Size.of(32, 1, 1)); + const d = down.at(i); + const sd = sh[1].at(i); + e.setPipeline(k.moe_down_0_one); + bind(e, 0, .{rf(io.z)}); + shape(e, 1, .{ 1, TOPK + 1, N }); + bind(e, 2, .{ d.w, d.s, d.b, sd.w, sd.s, sd.b, ids, rf(io.umem), rf(io.ucount), rf(io.y) }); + e.dispatchThreads(mtl.Size.of(32, D / 2, TOPK + 1), mtl.Size.of(32, 1, 1)); +} +/// oQ4's routed experts: the generated 4/64 gathers without the shared slot (moe.zig `experts` part 2). +fn moeRouted4(e: mtl.ComputeEncoder, i: usize, k: *const kernels.Kernels, gu: [2]Pool, down: Pool, io: Io) void { + const g = gu[0].at(i); + const u = gu[1].at(i); + const ids = rf(io.uids).at((i % reps) * 128 * 4); + e.setPipeline(k.moe_gateup_2); + bind(e, 0, .{rf(io.x)}); + shape(e, 1, .{ 1, D }); + bind(e, 2, .{ g.w, g.s, g.b, u.w, u.s, u.b, g.w, g.s, g.b, ids, rf(io.umem), rf(io.ucount) }); + e.setValue(@as(f32, 10.0), 14); + bind(e, 15, .{rf(io.z)}); + e.dispatchThreads(mtl.Size.of(32, N / 4, TOPK), mtl.Size.of(32, 1, 1)); + const d = down.at(i); + e.setPipeline(k.moe_down_2); + bind(e, 0, .{rf(io.z)}); + shape(e, 1, .{ 1, TOPK, N }); + bind(e, 2, .{ d.w, d.s, d.b, d.w, d.s, d.b, ids, rf(io.umem), rf(io.ucount), rf(io.y) }); + e.dispatchThreads(mtl.Size.of(32, D / 4, TOPK), mtl.Size.of(32, 1, 1)); +} + +// Other widths: wide.zig's launches as the engine makes them. +fn rows(e: mtl.ComputeEncoder, i: usize, p: *const wide.Pipes, pool: Pool, io: Io) void { + const q = pool.at(i); + wide.dense(p, e, q, .{ .in = rf(io.x), .ldx = q.k, .out = rf(io.y), .ldo = q.n, .rows = 1 }); +} +fn perHead(e: mtl.ComputeEncoder, i: usize, p: *const wide.Pipes, pool: Pool, io: Io, ldx: u32, hsx: u32) void { + const q = pool.at(i); + wide.perHead(p, e, q, H, .{ .in = rf(io.x), .ldx = ldx, .out = rf(io.y), .ldo = q.n, .rows = 1 }, hsx, q.n / H); +} +fn shared(e: mtl.ComputeEncoder, i: usize, p: *const wide.Pipes, k: *const kernels.Kernels, sh: [2]Pool, io: Io) void { + _ = k; + wide.glu(p, e, sh[0].at(i), N, rf(io.x), rf(io.z).at(2 * N * 2), 1, 10.0); // moe.zig sharedRows: one launch + const dq = sh[1].at(i); + wide.dense(p, e, dq, .{ .in = rf(io.z).at(2 * N * 2), .ldx = N, .out = rf(io.y), .ldo = D, .rows = 1 }); +} +/// A KDA layer's step state (16 copies: each launch its own, as each layer has). +const Kda = struct { + p: mtl.Buffer, + cs: [16]mtl.Buffer, + cw: mtl.Buffer, + a: mtl.Buffer, + dtb: mtl.Buffer, + st: [16]mtl.Buffer, + st_out: [16]mtl.Buffer, + onw: mtl.Buffer, + fg: mtl.Buffer, // f_b's and g_b's products made apart (.pre) + + fn init(device: mtl.Device) !Kda { + var k: Kda = undefined; + k.p = try buf(device, PS * 2); + @memset(k.p.slice(u16, PS), 0x3c00); + k.cw = try buf(device, 4 * 3 * W * 4); + @memset(k.cw.slice(f32, 4 * 3 * W), 0.25); + k.a = try buf(device, H * 4); + @memset(k.a.slice(f32, H), 0.5); + k.dtb = try buf(device, W * 4); + @memset(k.dtb.slice(f32, W), 0); + k.onw = try buf(device, KD * 4); + @memset(k.onw.slice(f32, KD), 1); + k.fg = try buf(device, 2 * 16 * W * 2); + for (&k.cs, &k.st, &k.st_out) |*c, *s, *o| { + c.* = try buf(device, 3 * 3 * W * 2); + @memset(c.slice(u16, 3 * 3 * W), 0); + s.* = try buf(device, H * KD * KD * 4); + @memset(s.slice(f32, H * KD * KD), 0); + o.* = try buf(device, H * KD * KD * 4); + } + return k; + } +}; + +/// kda_rows (generated dots, or `.pre`: f_b, g_b made apart first) on launch i's state and gates. +fn kdaStep(e: mtl.ComputeEncoder, i: usize, pipe: mtl.Pipeline, s: *const Kda, gates: [2]Pool, pre: ?*const wide.Pipes, io: Io) void { + const f = gates[0].at(i); + const g = gates[1].at(i); + var fb = [3]Ref{ f.w, f.s, f.b }; + var gb = [3]Ref{ g.w, g.s, g.b }; + if (pre) |p| { + const fo = rf(s.fg); + const go = fo.at(16 * W * 2); + wide.dense(p, e, f, .{ .in = rf(s.p).at(3 * W * 2), .ldx = PS, .out = fo, .ldo = W, .rows = 1 }); + wide.dense(p, e, g, .{ .in = rf(s.p).at((3 * W + KD) * 2), .ldx = PS, .out = go, .ldo = W, .rows = 1 }); + fb = .{ fo, fo, fo }; + gb = .{ go, go, go }; + } + const c = i % 16; + e.setPipeline(pipe); + bind(e, 0, .{rf(s.p)}); + shape(e, 1, .{ 1, PS }); + bind(e, 2, .{ rf(s.cs[c]), rf(s.cw), fb[0], fb[1], fb[2], gb[0], gb[1], gb[2], rf(s.a), rf(s.dtb), rf(s.st[c]), rf(s.onw) }); + e.setValue(@as(f32, -5.0), 14); + e.setValue(@as(f32, 1e-5), 15); + bind(e, 16, .{ rf(io.y), rf(s.st_out[c]), rf(s.cs[(c + 1) % 16]) }); + e.dispatchThreads(mtl.Size.of(32, 32, H), mtl.Size.of(32, 32, 1)); +} + +/// glm_rows.metal with glm_rows_speed.metal's candidate instances. +fn candidateLib(device: mtl.Device) !mtl.Library { + const src = try std.mem.concat(a, u8, &.{ sources.core_row_projection, wide.glue_math, sources.glm_rows, @embedFile("glm_rows_speed.metal") }); + defer a.free(src); + return mtl.Library.fromSource(device, src, mtl.CompileOptions.mlx()); +} + +const PM = 2048; // a prompt chunk's rows + +/// A prompt chunk's per-head product (embed_q: x qp's slices; unembed_out: the latent rows) on the row kernels. +fn headsRows(e: mtl.ComputeEncoder, i: usize, p: *const wide.Pipes, pool: Pool, io: Io, ldx: u32, kx: u32) void { + const q = pool.at(i); + wide.perHead(p, e, q, H, .{ .in = rf(io.x), .ldx = ldx, .out = rf(io.y), .ldo = q.n, .rows = PM }, kx, q.n / H); +} +/// The same as one affine_mm batch (prompt.zig headsMm). +fn headsMm(e: mtl.ComputeEncoder, i: usize, mm: tf.affine_mm.Pipes, pool: Pool, io: Io, ldx: u32, kx: u32) void { + var head = pool.at(i); + head.n /= H; + tf.affine_mm.rowSumsAt(e, mm, 64, rf(io.x), rf(io.z), PM, H * kx, ldx); + tf.affine_mm.denseBatch(e, mm, rf(io.x), rf(io.z), head, rf(io.y), PM, H, .{ .x_row = @intCast(ldx), .y_row = @intCast(H * head.n), .sums_row = H, .x_batch = @intCast(kx), .y_batch = @intCast(head.n), .sums_batch = 1, .w_batch = @intCast(head.n) }); +} +fn ehRows(e: mtl.ComputeEncoder, i: usize, p: *const wide.Pipes, eh: mtl.Buffer, io: Io) void { + _ = i; + wide.bf16(p, e, .{ .buf = eh }, D, 2 * D, .{ .in = rf(io.x), .ldx = 2 * D, .out = rf(io.y), .ldo = D, .rows = PM }); +} + +/// A chunk's dense product: row kernels, or affine_mm (`mm`). +fn chunkRows(e: mtl.ComputeEncoder, i: usize, p: *const wide.Pipes, pool: Pool, io: Io) void { + const q = pool.at(i); + wide.dense(p, e, q, .{ .in = rf(io.x), .ldx = q.k, .out = rf(io.y), .ldo = q.n, .rows = PM }); +} +fn chunkMm(e: mtl.ComputeEncoder, i: usize, mm: tf.affine_mm.Pipes, pool: Pool, io: Io) void { + const q = pool.at(i); + tf.affine_mm.rowSums(e, mm, 64, rf(io.x), rf(io.z), PM, q.k); + tf.affine_mm.dense(e, mm, rf(io.x), rf(io.z), q, rf(io.y), PM); +} + +/// A prompt chunk's products where oQ's paths differ from uniform's, ms a chunk of 2048 rows. +fn prompt(b: *Bench, k: *const kernels.Kernels, p: *const wide.Pipes) !void { + b.reps = 2; + defer b.reps = reps; + std.debug.print("prompt chunk ({d} rows), ms a launch:\n", .{PM}); + const eq = try Pool.init(b.device, H * 512, 256, 8, 1, 2); + defer eq.deinit(); + std.debug.print(" embed_q q8 per-head rows {d:8.2} affine_mm batch {d:8.2} (x11 a chunk)\n", .{ try b.time(headsRows, .{ p, eq, b.io, @as(u32, QR), @as(u32, 256) }) / 1e3, try b.time(headsMm, .{ p.mm8, eq, b.io, @as(u32, QR), @as(u32, 256) }) / 1e3 }); + const uo = try Pool.init(b.device, H * 256, 512, 8, 1, 2); + defer uo.deinit(); + std.debug.print(" unembed_out q8 per-head rows {d:8.2} affine_mm batch {d:8.2} (x11 a chunk)\n", .{ try b.time(headsRows, .{ p, uo, b.io, @as(u32, H * 512), @as(u32, 512) }) / 1e3, try b.time(headsMm, .{ p.mm8, uo, b.io, @as(u32, H * 512), @as(u32, 512) }) / 1e3 }); + { // the MTP head's bf16 eh_proj [4096, 8192] over the chunk (prompt.zig mtp) + const eh = try buf(b.device, D * 2 * D * 2); + defer eh.deinit(); + @memset(eh.slice(u16, 2 * D * D), 0x3c00); + std.debug.print(" eh_proj bf16 [4096, 8192] rows {d:8.2} (x1 a chunk with MTP)\n", .{try b.time(ehRows, .{ p, eh, b.io }) / 1e3}); + } + for ([_][3]usize{ .{ PS, D, 9 }, .{ D, W, 10 }, .{ 1536, D, 1 } }) |s| { + const m8 = try Pool.init(b.device, s[0], s[1], 8, 1, 2); + defer m8.deinit(); + const m4 = try Pool.init(b.device, s[0], s[1], 4, 1, 2); + defer m4.deinit(); + const t8 = try b.time(chunkMm, .{ p.mm8, m8, b.io }); + const t4 = try b.time(chunkMm, .{ k.mm_bf16, m4, b.io }); + for ([_]u8{ 5, 6 }) |bits| { + const m = try Pool.init(b.device, s[0], s[1], bits, 1, 2); + defer m.deinit(); + const on_mm = if (wide.mm(p, m.at(0))) |mm| try b.time(chunkMm, .{ mm, m, b.io }) else 0; + std.debug.print(" [{d}, {d}] q{d} rows {d:8.2} affine_mm {d:8.2} (affine_mm q8 {d:.2}, q4 {d:.2})\n", .{ s[0], s[1], bits, try b.time(chunkRows, .{ p, m, b.io }) / 1e3, on_mm / 1e3, t8 / 1e3, t4 / 1e3 }); + } + } +} + +fn cand(e: mtl.ComputeEncoder, i: usize, pipe: mtl.Pipeline, rps: u32, sg: u32, pool: Pool, io: Io) void { + const q = pool.at(i); + e.setPipeline(pipe); + bind(e, 0, .{ rf(io.x), q.w, q.s, q.b }); + e.setValue(wide.RowDims{ .k = @intCast(q.k), .n = @intCast(q.n), .group = 64, .rows = 1, .ldx = @intCast(q.k), .ldo = @intCast(q.n) }, 4); + bind(e, 5, .{rf(io.y)}); + e.dispatchGroups(mtl.Size.of((q.n + rps * sg - 1) / (rps * sg), 1, 1), mtl.Size.of(32 * sg, 1, 1)); +} + +/// glm_rows_speed.metal's launch shapes on `pool` at 8 bits: each one's bits == glm_rows's, then its time. +fn candidates(b: *const Bench, lib: mtl.Library, p: *const wide.Pipes, pool: Pool, name: []const u8) !void { + const q = pool.at(0); + const base = try b.time(rows, .{ p, pool, b.io }); + try b.run(rows, .{ p, pool, b.io }); + const want = try a.dupe(u16, b.io.y.slice(u16, q.n)); + defer a.free(want); + var best: f64 = base; + var best_name: [40]u8 = undefined; + var best_len: usize = 0; + for ([_][2]u32{ .{ 4, 2 }, .{ 8, 2 }, .{ 4, 4 }, .{ 8, 4 }, .{ 2, 4 }, .{ 2, 8 }, .{ 4, 8 } }) |rs| for ([_]u32{ 0, q.k }) |kk| { + var nb: [40]u8 = undefined; + const fname = try std.fmt.bufPrintSentinel(&nb, "x_q8_r{d}_s{d}_k{d}", .{ rs[0], rs[1], kk }, 0); + const pipe = try mtl.Pipeline.init(b.device, lib, fname, false); + defer pipe.deinit(); + @memset(b.io.y.slice(u16, q.n), 0); + try b.run(cand, .{ pipe, rs[0], rs[1], pool, b.io }); + if (!std.mem.eql(u16, b.io.y.slice(u16, q.n), want)) { + std.debug.print(" {s} {s}: BITS DIFFER\n", .{ name, fname }); + return error.BitsDiffer; + } + const us = try b.time(cand, .{ pipe, rs[0], rs[1], pool, b.io }); + if (us < best) { + best = us; + @memcpy(best_name[0..fname.len], fname); + best_len = fname.len; + } + std.debug.print(" {s:<16} {s:<24} {d:8.1} us ({d:5.3}x glm_rows)\n", .{ name, fname, us, base / us }); + }; + std.debug.print(" {s:<16} glm_rows {d:.1} us, best {s} {d:.1} us ({d:.0} GB/s)\n", .{ name, base, best_name[0..best_len], best, @as(f64, @floatFromInt(pool.bytes)) / best / 1e3 }); +} + +/// oQ2's MoE: routed by pick and the shared expert, a launch for gate|up and one for down (wide.moe). +fn moeTwo(e: mtl.ComputeEncoder, i: usize, p: *const wide.Pipes, gu: [2]Pool, down: Pool, sh: [2]Pool, io: Io) void { + const m: wide.Moe = .{ .gate = gu[0].at(i), .up = gu[1].at(i), .down = down.at(i), .sh_gate_up = sh[0].at(i), .sh_down = sh[1].at(i) }; + const z = rf(io.z); + wide.moe(p, e, m, E, TOPK, 10.0, .{ .x = rf(io.x), .ids = rf(io.ids).at((i % reps) * TOPK * 4), .act = z, .acts = z.at(TOPK * N * 2), .ye = rf(io.y), .ys = rf(io.y).at(TOPK * D * 2) }, 1); +} + +/// moeTwo through candidate pipelines (glu, down) at `rps` columns a simdgroup. +fn moeCand(e: mtl.ComputeEncoder, i: usize, pipes: [2]mtl.Pipeline, rps: [2]u32, gu: [2]Pool, down: Pool, sh: [2]Pool, io: Io) void { + const z = rf(io.z); + const ids = rf(io.ids).at((i % reps) * TOPK * 4); + const steps = [2]struct { w: Affine, up: Affine, sh: Affine, x: Ref, out: Ref, out2: Ref, repeat: u32 }{ + .{ .w = gu[0].at(i), .up = gu[1].at(i), .sh = sh[0].at(i), .x = rf(io.x), .out = z, .out2 = z.at(TOPK * N * 2), .repeat = TOPK }, + .{ .w = down.at(i), .up = down.at(i), .sh = sh[1].at(i), .x = z, .out = rf(io.y), .out2 = rf(io.y).at(TOPK * D * 2), .repeat = 1 }, + }; + for (steps, pipes, rps, 0..) |st, pipe, r, j| { + e.setPipeline(pipe); + bind(e, 0, .{ st.x, st.w.w, st.w.s, st.w.b }); + e.setValue(wide.PickDims{ .k = @intCast(st.w.k), .n = @intCast(st.w.n), .group = 64, .slots = TOPK, .repeat = @intCast(st.repeat), .experts = E, .ldo = @intCast(st.w.n), .lim = 10.0 }, 4); + bind(e, 5, .{ st.out, ids, st.up.w, st.up.s, st.up.b, if (j == 0) rf(io.x) else z.at(TOPK * N * 2), st.sh.w, st.sh.s, st.sh.b, st.out2 }); + e.dispatchGroups(mtl.Size.of((st.w.n + 2 * r - 1) / (2 * r), TOPK + 1, 1), mtl.Size.of(64, 1, 1)); + } +} + +/// glm_moe_rows candidates (glm_rows_speed.metal) for oQ2's MoE: bits == wide.moe's, then the time a layer. +fn moeCandidates(b: *const Bench, lib: mtl.Library, p: *const wide.Pipes, gu: [2]Pool, down: Pool, sh: [2]Pool) !void { + const base = try b.time(moeTwo, .{ p, gu, down, sh, b.io }); + try b.run(moeTwo, .{ p, gu, down, sh, b.io }); + const want = try a.dupe(u16, b.io.y.slice(u16, (TOPK + 1) * D)); + defer a.free(want); + for ([_][2][:0]const u8{ .{ "x_moe_glu_r8_k0", "x_moe_r8_k0" }, .{ "x_moe_glu_r2_k0", "x_moe_r2_k0" }, .{ "x_moe_glu_r4_k4096", "x_moe_r4_k2048" }, .{ "x_moe_glu_r8_k4096", "x_moe_r8_k2048" } }) |names| { + var pipes: [2]mtl.Pipeline = undefined; + for (&pipes, names) |*x, n| x.* = try mtl.Pipeline.init(b.device, lib, n, false); + defer for (pipes) |x| x.deinit(); + const rps = [2]u32{ names[0][11] - '0', names[1][7] - '0' }; + @memset(b.io.y.slice(u16, (TOPK + 1) * D), 0); + try b.run(moeCand, .{ pipes, rps, gu, down, sh, b.io }); + if (!std.mem.eql(u16, b.io.y.slice(u16, (TOPK + 1) * D), want)) return error.BitsDiffer; + const us = try b.time(moeCand, .{ pipes, rps, gu, down, sh, b.io }); + std.debug.print(" moe q2+q8 {s} / {s}: {d:8.1} us ({d:5.3}x wide.moe {d:.1} us)\n", .{ names[0], names[1], us, base / us, base }); + } +} + +fn routedRows(e: mtl.ComputeEncoder, i: usize, p: *const wide.Pipes, gu: [2]Pool, down: Pool, io: Io) void { + const ids = rf(io.ids).at((i % reps) * TOPK * 4); + wide.byIds(p, e, .{ .w = gu[0].at(i), .up = gu[1].at(i), .experts = E, .repeat = TOPK, .lim = 10.0 }, rf(io.x), ids, rf(io.z), TOPK); + wide.byIds(p, e, .{ .w = down.at(i), .experts = E, .repeat = 1 }, rf(io.z), ids, rf(io.y), TOPK); +} + +/// A token's GPU time by class: launches a token times the median launch. +const Row = struct { name: []const u8, us: f64, count: f64, bytes: f64 }; +const Sheet = struct { + rows: [24]Row = undefined, + n: usize = 0, + + fn add(s: *Sheet, name: []const u8, us: f64, count: f64, bytes: usize) void { + s.rows[s.n] = .{ .name = name, .us = us, .count = count, .bytes = @floatFromInt(bytes) }; + s.n += 1; + std.debug.print(" {s:<26} {d:9.1} us {d:7.1} MB {d:6.1} GB/s x{d:<5.0} {d:7.3} ms/token\n", .{ name, us, @as(f64, @floatFromInt(bytes)) / 1e6, @as(f64, @floatFromInt(bytes)) / us / 1e3, count, us * count / 1e3 }); + } + + fn total(s: *const Sheet, label: []const u8) void { + var ms: f64 = 0; + var gb: f64 = 0; + for (s.rows[0..s.n]) |x| { + ms += x.us * x.count / 1e3; + gb += x.bytes * x.count / 1e9; + } + std.debug.print(" {s}: {d:.3} ms a token over the products, {d:.3} GB read ({d:.0} GB/s)\n", .{ label, ms, gb, gb / ms * 1e3 }); + } +}; + +pub fn main() !void { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + const k = try kernels.load(a, device); + defer k.deinit(); + try k.widen(a, device); + const p = &k.wide.?; + const io = Io{ .x = try buf(device, 160 << 20), .y = try buf(device, 160 << 20), .z = try buf(device, 64 << 20), .ids = try buf(device, reps * TOPK * 4), .uids = try buf(device, reps * 128 * 4), .umem = try buf(device, 128 * 16 * 4), .ucount = try buf(device, 4) }; + @memset(io.x.slice(u16, 32 << 20), 0x3c00); + for (0..reps) |i| for (0..TOPK) |j| { // launch i's picks: distinct experts, the pool walked in order + io.ids.slice(i32, reps * TOPK)[i * TOPK + j] = @intCast((i * TOPK + j) % E); + io.uids.slice(i32, reps * 128)[i * 128 + j] = @intCast((i * TOPK + j) % E); + }; + @memset(io.umem.slice(i32, 128 * 16), -1); + for (0..TOPK) |u| io.umem.slice(i32, 128 * 16)[u * 16] = @intCast(u); // unique id u: row 0's pick u + io.ucount.slice(i32, 1)[0] = TOPK; + var b = Bench{ .device = device, .queue = queue, .k = k, .io = io }; + std.debug.print("GPU {s}, tensor units {}; one row, {d} launches a command buffer, median of {d}\n", .{ device.name(), device.tensorUnits(), reps, runs }); + const shapes = [_]struct { name: []const u8, n: usize, k: usize, uni: ?mtl.Pipeline, count: [3]f64 }{ // counts: uniform, oQ2, oQ4 + .{ .name = "kda in", .n = PS, .k = D, .uni = k.qmv_kda_in, .count = .{ 34, 34, 34 } }, + .{ .name = "kda f_b/g_b", .n = W, .k = KD, .uni = null, .count = .{ 0, 0, 0 } }, // inside the KDA step's time + .{ .name = "kda out", .n = D, .k = W, .uni = k.qmv_kda_out, .count = .{ 34, 34, 34 } }, + .{ .name = "mla x_proj", .n = XP, .k = D, .uni = k.qmv_x, .count = .{ 11, 11, 11 } }, + .{ .name = "mla qr_proj", .n = QR, .k = 1536, .uni = k.qmv_qr, .count = .{ 11, 11, 11 } }, + .{ .name = "mla out", .n = D, .k = 16384, .uni = k.qmv_mla_out, .count = .{ 11, 11, 11 } }, + .{ .name = "dense gate|up", .n = 24576, .k = D, .uni = k.qmv_dense_gu, .count = .{ 3, 3, 3 } }, + .{ .name = "dense down", .n = D, .k = 12288, .uni = k.qmv_dense_down, .count = .{ 3, 3, 3 } }, + .{ .name = "lm_head", .n = 154880, .k = D, .uni = k.qmv_head, .count = .{ 1, 1, 1 } }, + }; + var uni = Sheet{}; + var oq2 = Sheet{}; + var oq4 = Sheet{}; + std.debug.print("generated 4/64 (uniform) and glm_rows at 8, 6, 5 bits:\n", .{}); + for (shapes) |s| { + if (s.uni) |pipe| { + const m = try Pool.init(device, s.n, s.k, 4, 1, 1); + defer m.deinit(); + uni.add(s.name, try b.time(qmv, .{ pipe, m, io }), s.count[0], m.bytes); + } + for ([_]u8{ 8, 6, 5 }) |bits| { + const m = try Pool.init(device, s.n, s.k, bits, 1, 1); + defer m.deinit(); + const us = try b.time(rows, .{ p, m, io }); + var name_buf: [40]u8 = undefined; + const name = try std.fmt.bufPrint(&name_buf, "{s} q{d}", .{ s.name, bits }); + if (bits == 8) { + oq2.add(try a.dupe(u8, name), us, s.count[1], m.bytes); + oq4.add(try a.dupe(u8, name), us, s.count[2], m.bytes); // oQ4's 5/6-bit members priced at 8 (an upper bound) + } else std.debug.print(" {s:<26} {d:9.1} us {d:7.1} MB {d:6.1} GB/s\n", .{ name, us, @as(f64, @floatFromInt(m.bytes)) / 1e6, @as(f64, @floatFromInt(m.bytes)) / us / 1e3 }); + } + } + { // the KDA step: generated 4/4 dots (uniform), generated 8/8 dots, f_b and g_b made apart at 8 and 5 bits (.pre) + const st = try Kda.init(device); + const g4 = [2]Pool{ try Pool.init(device, W, KD, 4, 1, 16), try Pool.init(device, W, KD, 4, 1, 16) }; + const g8 = [2]Pool{ try Pool.init(device, W, KD, 8, 1, 16), try Pool.init(device, W, KD, 8, 1, 16) }; + const g5 = [2]Pool{ try Pool.init(device, W, KD, 5, 1, 16), try Pool.init(device, W, KD, 5, 1, 16) }; + const step_bytes = 2 * H * KD * KD * 4; // the state read and written + uni.add("kda step (dots 4/4)", try b.time(kdaStep, .{ k.kda_rows, &st, g4, null, io }), 34, step_bytes + g4[0].bytes * 2); + const us_pre = try b.time(kdaStep, .{ p.kda_pre, &st, g8, p, io }); + std.debug.print(" {s:<26} {d:9.1} us\n", .{ "kda step (.pre q8)", us_pre }); + const us_pre5 = try b.time(kdaStep, .{ p.kda_pre, &st, g5, p, io }); + std.debug.print(" {s:<26} {d:9.1} us\n", .{ "kda step (.pre q5)", us_pre5 }); + const f8 = wide.kdaGate(p, g8[0].at(0), g8[1].at(0)).fused; + const us_f8 = try b.time(kdaStep, .{ f8, &st, g8, null, io }); + oq2.add("kda step (fused q8)", us_f8, 34, step_bytes + g8[0].bytes * 2); + oq4.add("kda step (fused q8)", us_f8, 34, step_bytes + g8[0].bytes * 2); + const f5 = wide.kdaGate(p, g5[0].at(0), g5[1].at(0)).fused; + std.debug.print(" {s:<26} {d:9.1} us\n", .{ "kda step (fused q5)", try b.time(kdaStep, .{ f5, &st, g5, null, io }) }); + } + { // MLA absorb / unabsorb: kv_b's halves (uniform) or embed_q / unembed_out a head at 8 bits + const kv = try Pool.init(device, H * 512, 512, 4, 1, 1); + defer kv.deinit(); + uni.add("mla absorb", try b.time(absorb, .{ k, kv, io }), 11, kv.bytes / 2); + uni.add("mla unabsorb", try b.time(unabsorb, .{ k, kv, io }), 11, kv.bytes / 2); + const eq = try Pool.init(device, H * 512, 256, 8, 1, 1); + defer eq.deinit(); + const us_eq = try b.time(perHead, .{ p, eq, io, @as(u32, QR), @as(u32, 256) }); + oq2.add("mla embed_q q8", us_eq, 11, eq.bytes); + oq4.add("mla embed_q q8", us_eq, 11, eq.bytes); + const uo = try Pool.init(device, H * 256, 512, 8, 1, 1); + defer uo.deinit(); + const us_uo = try b.time(perHead, .{ p, uo, io, @as(u32, H * 512), @as(u32, 512) }); + oq2.add("mla unembed_out q8", us_uo, 11, uo.bytes); + oq4.add("mla unembed_out q8", us_uo, 11, uo.bytes); + } + { // MoE: the routed experts (8 picks) with the shared expert + const gu4 = [2]Pool{ try Pool.init(device, N, D, 4, E, 1), try Pool.init(device, N, D, 4, E, 1) }; + defer for (gu4) |x| x.deinit(); + const dn4 = try Pool.init(device, D, N, 4, E, 1); + defer dn4.deinit(); + const routed4 = 3 * affine.bytes(N, D, .{ .bits = 4, .group = 64 }) * TOPK; + const sh4 = [2]Pool{ try Pool.init(device, 2 * N, D, 4, 1, 2), try Pool.init(device, D, N, 4, 1, 2) }; + defer for (sh4) |x| x.deinit(); + const shared4 = sh4[0].bytes + sh4[1].bytes; + uni.add("moe routed+shared 4", try b.time(moeBoth, .{ k, gu4, dn4, sh4, io }), 42, routed4 + shared4); + oq4.add("moe routed 4", try b.time(moeRouted4, .{ k, gu4, dn4, io }), 42, routed4); + const sh8 = [2]Pool{ try Pool.init(device, 2 * N, D, 8, 1, 2), try Pool.init(device, D, N, 8, 1, 2) }; + defer for (sh8) |x| x.deinit(); + const us_sh = try b.time(shared, .{ p, k, sh8, io }); + std.debug.print(" {s:<26} {d:9.1} us (oQ2: in the routed launches below)\n", .{ "moe shared q8", us_sh }); + oq4.add("moe shared q8", us_sh, 42, sh8[0].bytes + sh8[1].bytes); + const alt = try b.time(moeTwo, .{ p, gu4, dn4, sh8, io }); // oQ4's routed 4/64 on glm_rows sums with the shared + std.debug.print(" {s:<26} {d:9.1} us (not taken: changes routed bits)\n", .{ "moe routed 4 + shared q8", alt }); + const gu2 = [2]Pool{ try Pool.init(device, N, D, 2, E, 1), try Pool.init(device, N, D, 2, E, 1) }; + defer for (gu2) |x| x.deinit(); + const dn2 = try Pool.init(device, D, N, 2, E, 1); + defer dn2.deinit(); + std.debug.print(" {s:<26} {d:9.1} us (apart)\n", .{ "moe routed q2 (ids)", try b.time(routedRows, .{ p, gu2, dn2, io }) }); + oq2.add("moe routed q2 + shared q8", try b.time(moeTwo, .{ p, gu2, dn2, sh8, io }), 42, 3 * affine.bytes(N, D, .{ .bits = 2, .group = 64 }) * TOPK + sh8[0].bytes + sh8[1].bytes); + if (std.c.getenv("GLM_SPEED_CANDIDATES") != null) { + const lib = try candidateLib(device); + defer lib.deinit(); + try moeCandidates(&b, lib, p, gu2, dn2, sh8); + } + } + if (std.c.getenv("GLM_SPEED_CANDIDATES") != null) { + const lib = try candidateLib(device); + defer lib.deinit(); + std.debug.print("candidates at 8 bits (bits checked == glm_rows):\n", .{}); + for (shapes) |s| { + const m = try Pool.init(device, s.n, s.k, 8, 1, 1); + defer m.deinit(); + try candidates(&b, lib, p, m, s.name); + } + } + if (std.c.getenv("GLM_SPEED_PROMPT") != null) try prompt(&b, k, p); + std.debug.print("uniform 4/64:\n", .{}); + for (uni.rows[0..uni.n]) |x| std.debug.print(" {s:<26} {d:7.3} ms/token\n", .{ x.name, x.us * x.count / 1e3 }); + uni.total("uniform 4/64"); + std.debug.print("oQ2:\n", .{}); + for (oq2.rows[0..oq2.n]) |x| std.debug.print(" {s:<26} {d:7.3} ms/token\n", .{ x.name, x.us * x.count / 1e3 }); + oq2.total("oQ2"); + std.debug.print("oQ4 (5/6-bit members priced at 8 bits):\n", .{}); + for (oq4.rows[0..oq4.n]) |x| std.debug.print(" {s:<26} {d:7.3} ms/token\n", .{ x.name, x.us * x.count / 1e3 }); + oq4.total("oQ4"); +} 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; +}