diff --git a/build.zig b/build.zig index 1bf8bfa260..d3b8a0a760 100644 --- a/build.zig +++ b/build.zig @@ -3,6 +3,7 @@ const builtin = @import("builtin"); const cuda_build = @import("zig/build/cuda.zig"); const dist_build = @import("zig/build/dist.zig"); const qwen27_build = @import("zig/build/qwen27.zig"); +const core_checks = @import("zig/build/core_checks.zig"); const hip_build = @import("zig/build/hip.zig"); const grammar_build = @import("zig/build/grammar.zig"); const counters_build = @import("zig/build/metal_counters.zig"); @@ -455,6 +456,7 @@ fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "row_projection", .module = row_mod } }, }) }); b.step("tf-row-check", "The core row projection on synthetic 4-bit matrices: widths bit-equal, a CPU reference, indexed experts").dependOn(&b.addInstallArtifact(row_check, .{}).step); + core_checks.rowWidths(b, target, optimize, metal, row_mod); test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = row_mod })).step); exl3_build.targets(b, target, optimize, metal, test_step); exlr8_build.metalTargets(b, target, optimize, metal, mods.sources, test_step); @@ -474,6 +476,7 @@ fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b // The core lane, shared-context and Qwen3.8-27B programs and host tests. qwen27_build.targets(b, target, optimize, metal, sources, engine, lanes, mods.qwen27, test_step); + core_checks.targets(b, target, optimize, metal, engine); _ = nativeServer(b, target, optimize, metal, engine, api, engines, build_options, test_step); diff --git a/zig/build/core_checks.zig b/zig/build/core_checks.zig 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/qwen27.zig b/zig/build/qwen27.zig index fc5245b612..73b38e72f1 100644 --- a/zig/build/qwen27.zig +++ b/zig/build/qwen27.zig @@ -97,6 +97,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu const engine_only: []const std.Build.Module.Import = &.{ .{ .name = "tensorfold", .module = engine }, .{ .name = "metal", .module = metal } }; _ = program(b, "tf-qwen27-synthetic", "Build bounded synthetic native hybrid chunk checks, no models", "zig/tests/qwen27/synthetic.zig", target, false, engine_only); _ = program(b, "tf-qwen27-dflash-synthetic", "Build model-free native draft context/forward/tree check", "zig/tests/qwen27/dflash_runtime.zig", target, false, engine_only); + _ = program(b, "tf-qwen27-widths", "Build the projections at every oQ width before tensor units: float64, row windows, mixed stacks, embedding", "zig/tests/qwen27/widths.zig", target, false, engine_only); _ = program(b, "tf-qwen27-dflash-refine", "Build the drafter refinement probe (top-1 by block position, masked versus given)", "zig/tests/qwen27/dflash_refine.zig", target, false, engine_only); const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} }); diff --git a/zig/kernels/metal/core/affine_mm.metal b/zig/kernels/metal/core/affine_mm.metal index 16e54006d2..464e459fe4 100644 --- a/zig/kernels/metal/core/affine_mm.metal +++ b/zig/kernels/metal/core/affine_mm.metal @@ -4,7 +4,7 @@ using namespace metal; #include "../nax.h" using namespace tfp; -// Defines: TF_BITS (4 or 8), TF_GROUP (32 or 64), TF_OUT_T (bfloat, or float for partial sums another pass adds). +// Defines: TF_BITS (2, 4 or 8), TF_GROUP (32 or 64), TF_OUT_T (bfloat, or float for partial sums another pass adds). #ifndef TF_OUT_T #define TF_OUT_T bfloat #endif @@ -66,7 +66,16 @@ inline void am_first(thread frag& lo, thread frag& hi, thread cons #endif } -// A thread's 32 codes of one weight row in a step (16-byte loads). +// A thread's 32 codes of one weight row in a step (16-byte loads; 8 bytes at 2 bits). +#if TF_BITS == 2 +struct Codes { + uint2 v; +}; + +inline Codes am_load(const device uchar* w) { + return Codes{*(const device uint2*)w}; +} +#else struct Codes { uint4 v[TF_BITS / 4]; }; @@ -79,11 +88,24 @@ inline Codes am_load(const device uchar* w) { } return c; } +#endif // The codes plus OFF, exactly, as bf16 into the tile row (16-byte stores). inline void am_codes(thread const Codes& c, threadgroup bfloat* out) { threadgroup uint4* o = (threadgroup uint4*)out; -#if TF_BITS == 4 +#if TF_BITS == 2 + TF_UNROLL + for (short i = 0; i < 2; i++) { // a word's 16 codes, as themselves (no offset): byte j holds codes 4j..4j+3 + const uint w = c.v[i]; + const float4 f0 = float4(as_type(w & 0x03030303u)), f1 = float4(as_type((w >> 2) & 0x03030303u)); + const float4 f2 = float4(as_type((w >> 4) & 0x03030303u)), f3 = float4(as_type((w >> 6) & 0x03030303u)); + TF_UNROLL + for (short j = 0; j < 4; j++) { + o[2 * i + j / 2][2 * (j % 2)] = as_type(ushort2(as_type(bfloat(f0[j])), as_type(bfloat(f1[j])))); + o[2 * i + j / 2][2 * (j % 2) + 1] = as_type(ushort2(as_type(bfloat(f2[j])), as_type(bfloat(f3[j])))); + } + } +#elif TF_BITS == 4 TF_UNROLL for (short i = 0; i < 4; i++) { const uint w = c.v[0][i]; diff --git a/zig/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..5a503e3acc 100644 --- a/zig/kernels/metal/core/row_projection.metal +++ b/zig/kernels/metal/core/row_projection.metal @@ -1,4 +1,4 @@ -// Affine 2-, 4-, 6- and 8-bit row projections for chips without tensor units: RPS columns a simdgroup, MB rows a weight read, each row its own fp32 sums. +// Affine 2- to 8-bit row projections without tensor units: RPS columns a simdgroup, MB rows a weight read. #include using namespace metal; @@ -18,7 +18,7 @@ inline float tf_rp_load16(const device bfloat* x, thread float* xt) { return sum; } -// MB rows' dots with one weight word group (16 codes; 6-bit: four values in three bytes), the codes unpacked to floats once; each row's sums in one order. +// MB rows' dots with 16 codes (3, 5, 6 bits: 8, 8, 4 codes in 3, 5, 3 bytes), unpacked once; sums in one order. template inline void tf_rp_qdot16(const device uint8_t* w, const thread float (*xt)[16], float scale, float bias, const thread float* sum, thread float* out) { @@ -40,6 +40,20 @@ inline void tf_rp_qdot16(const device uint8_t* w, const thread float (*xt)[16], q[4 * i + 2] = float(((p[1] >> 4) & 0x0f) | ((p[2] & 0x03) << 4)); q[4 * i + 3] = float((p[2] >> 2) & 0x3f); } + } else if (BITS == 3) { + for (int i = 0; i < 2; i++) { + const device uint8_t* p = w + 3 * i; + const uint v = uint(p[0]) | (uint(p[1]) << 8) | (uint(p[2]) << 16); + for (int j = 0; j < 8; j++) q[8 * i + j] = float((v >> (3 * j)) & 7u); + } + } else if (BITS == 5) { + for (int i = 0; i < 2; i++) { + const device uint8_t* p = w + 5 * i; + const uint lo = uint(p[0]) | (uint(p[1]) << 8) | (uint(p[2]) << 16) | (uint(p[3]) << 24), hi = p[4]; + for (int j = 0; j < 6; j++) q[8 * i + j] = float((lo >> (5 * j)) & 31u); + q[8 * i + 6] = float(((lo >> 30) | (hi << 2)) & 31u); + q[8 * i + 7] = float(hi >> 3); + } } else { for (int i = 0; i < 16; i++) q[i] = float(w[i]); } @@ -93,8 +107,8 @@ inline bfloat tf_rp_out(float v) { return bfloat(v); } -// dims: K, N, GS, rows. X [rows, K] bf16; W [N, K*BITS/32] words; S, B [N, K/GS] bf16; OUT [rows, N] bf16. -template +// dims: K, N, GS, rows, OUT's row stride when STRIDED (else N). X [rows, K]; W [N, K*BITS/32] words; S, B [N, K/GS]. +template [[kernel]] void tf_row_projection(const device bfloat* X [[buffer(0)]], const device uint32_t* W [[buffer(1)]], const device bfloat* S [[buffer(2)]], @@ -104,7 +118,7 @@ template uint simdgroup_index_in_threadgroup [[simdgroup_index_in_threadgroup]], uint thread_index_in_simdgroup [[thread_index_in_simdgroup]], uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]]) { - const int K = dims[0], N = dims[1], GS = dims[2], rows = dims[3]; + const int K = dims[0], N = dims[1], GS = dims[2], rows = dims[3], LD = STRIDED ? dims[4] : N; const uint lane = thread_index_in_simdgroup; const int row0 = (int(threadgroup_position_in_grid.x) * SG + int(simdgroup_index_in_threadgroup)) * RPS; if (row0 >= N) return; @@ -119,14 +133,14 @@ template tf_rp_rows(wr, sr, br, X + size_t(m) * K, K, GS, lane, acc); if (lane == 0) for (int b = 0; b < MB; b++) - for (int j = 0; j < RPS; j++) OUT[size_t(m + b) * N + row0 + j] = tf_rp_out(acc[b][j]); + for (int j = 0; j < RPS; j++) OUT[size_t(m + b) * LD + row0 + j] = tf_rp_out(acc[b][j]); } #pragma clang loop unroll(disable) for (; m < rows; m++) { float acc[1][RPS]; tf_rp_rows(wr, sr, br, X + size_t(m) * K, K, GS, lane, acc); if (lane == 0) - for (int j = 0; j < RPS; j++) OUT[size_t(m) * N + row0 + j] = tf_rp_out(acc[0][j]); + for (int j = 0; j < RPS; j++) OUT[size_t(m) * LD + row0 + j] = tf_rp_out(acc[0][j]); } } @@ -195,3 +209,31 @@ template [[host_name("tf_row_projection_q8_bf16")]] [[kernel]] decltype(tf_row_p template [[host_name("tf_row_projection_indexed_q8_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 8, true>) tf_row_projection_indexed<4, 2, false, 8, true>; template [[host_name("tf_row_projection_relu2_q8_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 8, true>) tf_row_projection<4, 2, 2, true, 8, true>; template [[host_name("tf_row_projection_indexed_relu2_q8_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 8, true>) tf_row_projection_indexed<4, 2, true, 8, true>; +template [[host_name("tf_row_projection_q3_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 3, false>) tf_row_projection<4, 2, 2, false, 3, false>; +template [[host_name("tf_row_projection_indexed_q3_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 3, false>) tf_row_projection_indexed<4, 2, false, 3, false>; +template [[host_name("tf_row_projection_relu2_q3_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 3, false>) tf_row_projection<4, 2, 2, true, 3, false>; +template [[host_name("tf_row_projection_indexed_relu2_q3_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 3, false>) tf_row_projection_indexed<4, 2, true, 3, false>; +template [[host_name("tf_row_projection_q3_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 3, true>) tf_row_projection<4, 2, 2, false, 3, true>; +template [[host_name("tf_row_projection_indexed_q3_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 3, true>) tf_row_projection_indexed<4, 2, false, 3, true>; +template [[host_name("tf_row_projection_relu2_q3_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 3, true>) tf_row_projection<4, 2, 2, true, 3, true>; +template [[host_name("tf_row_projection_indexed_relu2_q3_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 3, true>) tf_row_projection_indexed<4, 2, true, 3, true>; +template [[host_name("tf_row_projection_q5_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 5, false>) tf_row_projection<4, 2, 2, false, 5, false>; +template [[host_name("tf_row_projection_indexed_q5_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 5, false>) tf_row_projection_indexed<4, 2, false, 5, false>; +template [[host_name("tf_row_projection_relu2_q5_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 5, false>) tf_row_projection<4, 2, 2, true, 5, false>; +template [[host_name("tf_row_projection_indexed_relu2_q5_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 5, false>) tf_row_projection_indexed<4, 2, true, 5, false>; +template [[host_name("tf_row_projection_q5_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 5, true>) tf_row_projection<4, 2, 2, false, 5, true>; +template [[host_name("tf_row_projection_indexed_q5_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 5, true>) tf_row_projection_indexed<4, 2, false, 5, true>; +template [[host_name("tf_row_projection_relu2_q5_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 5, true>) tf_row_projection<4, 2, 2, true, 5, true>; +template [[host_name("tf_row_projection_indexed_relu2_q5_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 5, true>) tf_row_projection_indexed<4, 2, true, 5, true>; +template [[host_name("tf_row_projection_strided_q2_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 2, false, true>) tf_row_projection<4, 2, 2, false, 2, false, true>; +template [[host_name("tf_row_projection_strided_q2_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 2, true, true>) tf_row_projection<4, 2, 2, false, 2, true, true>; +template [[host_name("tf_row_projection_strided_q3_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 3, false, true>) tf_row_projection<4, 2, 2, false, 3, false, true>; +template [[host_name("tf_row_projection_strided_q3_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 3, true, true>) tf_row_projection<4, 2, 2, false, 3, true, true>; +template [[host_name("tf_row_projection_strided_q4_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 4, false, true>) tf_row_projection<4, 2, 2, false, 4, false, true>; +template [[host_name("tf_row_projection_strided_q4_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 4, true, true>) tf_row_projection<4, 2, 2, false, 4, true, true>; +template [[host_name("tf_row_projection_strided_q5_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 5, false, true>) tf_row_projection<4, 2, 2, false, 5, false, true>; +template [[host_name("tf_row_projection_strided_q5_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 5, true, true>) tf_row_projection<4, 2, 2, false, 5, true, true>; +template [[host_name("tf_row_projection_strided_q6_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 6, false, true>) tf_row_projection<4, 2, 2, false, 6, false, true>; +template [[host_name("tf_row_projection_strided_q6_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 6, true, true>) tf_row_projection<4, 2, 2, false, 6, true, true>; +template [[host_name("tf_row_projection_strided_q8_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 8, false, true>) tf_row_projection<4, 2, 2, false, 8, false, true>; +template [[host_name("tf_row_projection_strided_q8_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 8, true, true>) tf_row_projection<4, 2, 2, false, 8, true, true>; diff --git a/zig/src/core/affine_mm.zig b/zig/src/core/affine_mm.zig index dd25dce831..4b5575a84a 100644 --- a/zig/src/core/affine_mm.zig +++ b/zig/src/core/affine_mm.zig @@ -4,12 +4,12 @@ const mtl = @import("metal"); const ks = @import("kernel_sources"); const frags = @import("frags.zig"); -/// The weights' bits (4 or 8) and group (32 or 64), and the output's type. +/// The weights' bits (2, 4 or 8) and group (32 or 64), and the output's type. pub const Format = struct { bits: u8 = 4, group: u16 = 64, out_f32: bool = false }; /// The kernel source for `f` (nax.h inlined for this GPU). pub fn source(device: mtl.Device, a: std.mem.Allocator, f: Format) ![]u8 { - if ((f.bits != 4 and f.bits != 8) or (f.group != 32 and f.group != 64)) return error.UnsupportedFormat; + if ((f.bits != 2 and f.bits != 4 and f.bits != 8) or (f.group != 32 and f.group != 64)) return error.UnsupportedFormat; const body = try frags.source(device, a, ks.core_affine_mm); defer a.free(body); return std.fmt.allocPrint(a, "#define TF_BITS {d}\n#define TF_GROUP {d}\n#define TF_OUT_T {s}\n{s}", .{ f.bits, f.group, if (f.out_f32) "float" else "bfloat", body }); @@ -135,7 +135,7 @@ test "dense and gathered, each format, against exact sums on the host" { const R = counts[0] + counts[1] + counts[2]; var prng = std.Random.DefaultPrng.init(7); const rnd = prng.random(); - for ([_][2]usize{ .{ 4, 64 }, .{ 4, 32 }, .{ 8, 64 }, .{ 8, 32 } }) |fmt| { + for ([_][2]usize{ .{ 4, 64 }, .{ 4, 32 }, .{ 8, 64 }, .{ 8, 32 }, .{ 2, 64 }, .{ 2, 32 } }) |fmt| { const bits = fmt[0]; const group = fmt[1]; const w_bytes = E * N * K * bits / 8; @@ -207,7 +207,11 @@ test "dense and gathered, each format, against exact sums on the host" { const off: f64 = if (bits == 4) 128 else 0; for (0..group) |j| { const k = g * group + j; - const qv: u32 = if (bits == 4) (wb[(row * K + k) / 2] >> @intCast(4 * (k % 2))) & 15 else wb[row * K + k]; + const qv: u32 = switch (bits) { + 2 => (wb[(row * K + k) / 4] >> @intCast(2 * (k % 4))) & 3, + 4 => (wb[(row * K + k) / 2] >> @intCast(4 * (k % 2))) & 15, + else => wb[row * K + k], + }; const xv: f64 = f32of(x[r * K + k]); dot += @as(f64, @floatFromInt(qv)) * xv; dmag += (@as(f64, @floatFromInt(qv)) + off) * @abs(xv); 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/qwen27/dflash/runtime_head.zig b/zig/src/families/qwen27/dflash/runtime_head.zig index f758fa5c65..be6581a9b7 100644 --- a/zig/src/families/qwen27/dflash/runtime_head.zig +++ b/zig/src/families/qwen27/dflash/runtime_head.zig @@ -14,18 +14,18 @@ pub const Head = struct { var head = Head{ .allocator = a, .linear = source }; errdefer head.deinit(); if (source.n != 248320) return head; - if (source.tile != 32 or source.k % 64 != 0) return error.BadDraftHead; - if (source.raw) |raw| if (raw.bits != 4 or raw.group != 64) return error.BadDraftHead; + if (source.tile != 32 or source.k % source.group != 0 or source.runs.len > 0) return error.BadDraftHead; const k: usize = source.k; + const row = k * source.bits / 8; // whole 32-row tiles in the lane layout, rows in the raw one const input_words = if (source.raw) |raw| raw.w else source.words.buffer; const input_offset = if (source.raw) |raw| raw.w_off else source.words.offset; - const groups = k / 64; - const words = try head.allocate(device, width * k / 2); + const groups = k / source.group; + const words = try head.allocate(device, width * row); var destination: usize = 0; for (spans) |span| { const rows = span[1] - span[0]; - const first = input_offset + @as(usize, span[0]) * k / 2; - @memcpy(words.contents()[destination * k / 2 ..][0 .. @as(usize, rows) * k / 2], input_words.contents()[first..][0 .. @as(usize, rows) * k / 2]); + const first = input_offset + @as(usize, span[0]) * row; + @memcpy(words.contents()[destination * row ..][0 .. @as(usize, rows) * row], input_words.contents()[first..][0 .. @as(usize, rows) * row]); destination += rows; } head.linear.n = width; diff --git a/zig/src/families/qwen27/dflash/runtime_model.zig b/zig/src/families/qwen27/dflash/runtime_model.zig index 8bbdfab534..dc8bd4db3a 100644 --- a/zig/src/families/qwen27/dflash/runtime_model.zig +++ b/zig/src/families/qwen27/dflash/runtime_model.zig @@ -136,7 +136,7 @@ pub fn loadGraph(a: std.mem.Allocator, b: *Backend) !weights.Graph { var arena = std.heap.ArenaAllocator.init(a); defer arena.deinit(); const local = arena.allocator(); - const embed = try binding(local, target, prefix, "model.embed_tokens"); + const embed = try binding(local, target, prefix, "model.embed_tokens", .{ .bits = target.weights.embed.bits, .group_size = target.weights.embed.group }); var names: std.ArrayList([]const u8) = .empty; var iterator = b.checkpoint.tensors.keyIterator(); while (iterator.next()) |name| try names.append(local, name.*); @@ -144,10 +144,10 @@ pub fn loadGraph(a: std.mem.Allocator, b: *Backend) !weights.Graph { while (later.next()) |name| try names.append(local, name.*); return weights.load(a, .{ .source = b.source(), .names = names.items }, b.config, embed); } -fn binding(a: std.mem.Allocator, target: *Target, prefix: []const u8, module: []const u8) !ck.Linear { +fn binding(a: std.mem.Allocator, target: *Target, prefix: []const u8, module: []const u8, format: @import("../affine.zig").Spec) !ck.Linear { const stem = try std.mem.concat(a, u8, &.{ prefix, module }); const w = Backend.cpu(try target.checkpoint.get(try std.mem.concat(a, u8, &.{ stem, ".weight" }))); const s = Backend.cpu(try target.checkpoint.get(try std.mem.concat(a, u8, &.{ stem, ".scales" }))); const b = Backend.cpu(try target.checkpoint.get(try std.mem.concat(a, u8, &.{ stem, ".biases" }))); - return ck.validate(w, s, b, .{ .bits = 4, .group_size = 64 }, target.config.vocab, target.config.hidden); + return ck.validate(w, s, b, format, target.config.vocab, target.config.hidden); } diff --git a/zig/src/families/qwen27/glue_gpu.zig b/zig/src/families/qwen27/glue_gpu.zig index 65a4b3b890..857a18944e 100644 --- a/zig/src/families/qwen27/glue_gpu.zig +++ b/zig/src/families/qwen27/glue_gpu.zig @@ -45,7 +45,7 @@ pub const Kernels = struct { errdefer span.cancel(); const e = span.encoder; - const a = abi.Embed{ .width = @intCast(c.hidden), .rows = rows, .vocab = @intCast(c.vocab), .bits = 4, .group = 64, .metadata = .bf16 }; + const a = abi.Embed{ .width = @intCast(c.hidden), .rows = rows, .vocab = @intCast(c.vocab), .bits = p.bits, .group = p.group, .metadata = .bf16 }; try abi.checkEmbed(a); e.setPipeline(k.embed); e.setBuffer(ids.buffer, ids.offset, 0); diff --git a/zig/src/families/qwen27/gpu_weights.zig b/zig/src/families/qwen27/gpu_weights.zig index b5642b8ea9..bb0fba68e3 100644 --- a/zig/src/families/qwen27/gpu_weights.zig +++ b/zig/src/families/qwen27/gpu_weights.zig @@ -5,6 +5,8 @@ const cp = @import("core").checkpoint_metal; const Config = @import("config.zig").Config; const projection = @import("projection.zig"); const lane = @import("core").lane_projection; +const widths = @import("widths.zig"); +const Width = widths.Width; const Ref = projection.Ref; const Linear = projection.Linear; pub const Tensor = cp.Tensor; @@ -20,14 +22,27 @@ pub const Weights = struct { norm: Tensor, head: Linear, owned: std.ArrayList(mtl.Buffer) = .empty, + runs: std.ArrayList([]projection.Run) = .empty, pub fn deinit(w: *Weights) void { for (w.owned.items) |b| b.deinit(); w.owned.deinit(w.allocator); + for (w.runs.items) |r| w.allocator.free(r); + w.runs.deinit(w.allocator); w.allocator.free(w.layers); } + + /// Whether every MLP projection is 4-bit g64 in one matrix (the Neural Engine share's format). + pub fn uniformMlp(w: *const Weights) bool { + for (w.layers) |l| for ([_]Linear{ l.gu, l.down }) |p| if (!uniform(p)) return false; + return true; + } }; +fn uniform(p: Linear) bool { + return p.bits == 4 and p.group == 64 and p.runs.len == 0; +} + /// Modules the graph reads only through repacked copies: their tensors stay in the files until load stages them. const repacked_modules = [_][]const u8{ ".mlp.gate_proj", ".mlp.up_proj", ".mlp.down_proj", ".linear_attn.in_proj_qkv", ".linear_attn.in_proj_z", ".linear_attn.in_proj_b", ".linear_attn.in_proj_a", ".linear_attn.out_proj", ".self_attn.q_proj", ".self_attn.k_proj", ".self_attn.v_proj", ".self_attn.o_proj" }; @@ -49,6 +64,8 @@ const Builder = struct { prefix: []const u8, owned: *std.ArrayList(mtl.Buffer), stage: *cp.Stage, + runs: *std.ArrayList([]projection.Run), + widths: *const widths.Widths, fn get(b: Builder, suffix: []const u8) !Tensor { const name = try std.mem.concat(b.allocator, u8, &.{ b.prefix, suffix }); @@ -60,18 +77,19 @@ const Builder = struct { return checked(try b.get(suffix), dtype, shape); } - fn quant(b: Builder, module: []const u8, n: usize, k: usize) !projection.Triple { + fn quant(b: Builder, module: []const u8, n: usize, k: usize, w: Width) !projection.Triple { + if (k * w.bits % 32 != 0 or k % w.group != 0) return error.BadNativeTensor; const suffixes = [_][]const u8{ ".weight", ".scales", ".biases" }; var tensors: [3]Tensor = undefined; for (suffixes, &tensors, 0..) |suffix, *t, i| { const name = try std.mem.concat(b.allocator, u8, &.{ module, suffix }); defer b.allocator.free(name); - t.* = try b.tensor(name, if (i == 0) .u32 else .bf16, &.{ n, if (i == 0) k / 8 else k / 64 }); + t.* = try b.tensor(name, if (i == 0) .u32 else .bf16, &.{ n, if (i == 0) k * w.bits / 32 else k / w.group }); } - return .{ .weight = tensors[0], .scales = tensors[1], .biases = tensors[2] }; + return .{ .weight = tensors[0], .scales = tensors[1], .biases = tensors[2], .bits = w.bits, .group = w.group }; } - /// `parts` (stacked in order, each `n` rows of `k`) staged from the files and repacked into one owned Linear. + /// `parts` (stacked in order, `n` rows of `k` each, at their widths) staged from the files and repacked. fn linear(b: Builder, parts: []const Part, k: usize) !Linear { var arena = std.heap.ArenaAllocator.init(b.allocator); defer arena.deinit(); @@ -83,12 +101,41 @@ const Builder = struct { defer b.stage.reset(); var triples: [3]projection.Triple = undefined; if (parts.len > triples.len) return error.BadProjectionShape; - for (parts, triples[0..parts.len], 0..) |p, *t, i| t.* = .{ - .weight = try checked(try b.stage.read(b.checkpoint, names[i * 3]), .u32, &.{ p.n, k / 8 }), - .scales = try checked(try b.stage.read(b.checkpoint, names[i * 3 + 1]), .bf16, &.{ p.n, k / 64 }), - .biases = try checked(try b.stage.read(b.checkpoint, names[i * 3 + 2]), .bf16, &.{ p.n, k / 64 }), + for (parts, triples[0..parts.len], 0..) |p, *t, i| { + if (k * p.w.bits % 32 != 0 or k % p.w.group != 0) return error.BadNativeTensor; + t.* = .{ + .weight = try checked(try b.stage.read(b.checkpoint, names[i * 3]), .u32, &.{ p.n, k * p.w.bits / 32 }), + .scales = try checked(try b.stage.read(b.checkpoint, names[i * 3 + 1]), .bf16, &.{ p.n, k / p.w.group }), + .biases = try checked(try b.stage.read(b.checkpoint, names[i * 3 + 2]), .bf16, &.{ p.n, k / p.w.group }), + .bits = p.w.bits, + .group = p.w.group, + }; + } + return b.stack(triples[0..parts.len]); + } + + /// Fused members of one width take one matrix as before; a mixed stack becomes one matrix a member. + fn stack(b: Builder, parts: []const projection.Triple) !Linear { + var ws: [3]Width = undefined; + for (parts, 0..) |p, i| ws[i] = .{ .bits = p.bits, .group = @intCast(p.group) }; + var buf: [3][2]usize = undefined; + const list = widths.runs(ws[0..parts.len], &buf); + if (list.len == 1) return projection.prepare(b.device, b.allocator, b.owned, parts, 32); + const runs = try b.allocator.alloc(projection.Run, list.len); + b.runs.append(b.allocator, runs) catch |err| { + b.allocator.free(runs); + return err; }; - return projection.prepare(b.device, b.allocator, b.owned, triples[0..parts.len], 32); + var first: u32 = 0; + for (list, runs) |r, *run| { + run.* = .{ .linear = try projection.prepare(b.device, b.allocator, b.owned, parts[r[0]..][0..r[1]], 32), .first = first }; + first += run.linear.n; + } + var fused = runs[0].linear; + fused.n = first; + fused.raw = null; + fused.runs = runs; + return fused; } fn layer(b: Builder, c: Config, i: usize) !Layer { @@ -101,31 +148,33 @@ const Builder = struct { return std.mem.concat(allocator, u8, &.{ p, leaf }); } }; + const w = b.widths.layers[i]; + const M = widths.Member; const mixer: @FieldType(Layer, "mixer") = switch (c.kind(i)) { .linear => .{ .linear = .{ - .qkv = try b.linear(&.{.{ .module = try Name.full(a, prefix, "linear_attn.in_proj_qkv"), .n = c.gdnQkvDim() }}, c.hidden), - .zba = try b.linear(&.{ .{ .module = try Name.full(a, prefix, "linear_attn.in_proj_z"), .n = c.gdnValueDim() }, .{ .module = try Name.full(a, prefix, "linear_attn.in_proj_b"), .n = c.v_heads }, .{ .module = try Name.full(a, prefix, "linear_attn.in_proj_a"), .n = c.v_heads } }, c.hidden), - .out = try b.linear(&.{.{ .module = try Name.full(a, prefix, "linear_attn.out_proj"), .n = c.hidden }}, c.gdnValueDim()), + .qkv = try b.linear(&.{.{ .module = try Name.full(a, prefix, "linear_attn.in_proj_qkv"), .n = c.gdnQkvDim(), .w = w[@backingInt(M.qkv)] }}, c.hidden), + .zba = try b.linear(&.{ .{ .module = try Name.full(a, prefix, "linear_attn.in_proj_z"), .n = c.gdnValueDim(), .w = w[@backingInt(M.z)] }, .{ .module = try Name.full(a, prefix, "linear_attn.in_proj_b"), .n = c.v_heads, .w = w[@backingInt(M.b)] }, .{ .module = try Name.full(a, prefix, "linear_attn.in_proj_a"), .n = c.v_heads, .w = w[@backingInt(M.a)] } }, c.hidden), + .out = try b.linear(&.{.{ .module = try Name.full(a, prefix, "linear_attn.out_proj"), .n = c.hidden, .w = w[@backingInt(M.out)] }}, c.gdnValueDim()), .conv = try b.tensor(try Name.full(a, prefix, "linear_attn.conv1d.weight"), .bf16, &.{ c.gdnQkvDim(), c.conv_kernel, 1 }), .a_log = try b.tensor(try Name.full(a, prefix, "linear_attn.A_log"), .bf16, &.{c.v_heads}), .dt = try b.tensor(try Name.full(a, prefix, "linear_attn.dt_bias"), .bf16, &.{c.v_heads}), .norm = try b.tensor(try Name.full(a, prefix, "linear_attn.norm.weight"), .bf16, &.{c.dv}), } }, .attention => .{ .attention = .{ - .q = try b.linear(&.{.{ .module = try Name.full(a, prefix, "self_attn.q_proj"), .n = c.qDim() }}, c.hidden), - .kv = try b.linear(&.{ .{ .module = try Name.full(a, prefix, "self_attn.k_proj"), .n = c.kvDim() }, .{ .module = try Name.full(a, prefix, "self_attn.v_proj"), .n = c.kvDim() } }, c.hidden), - .out = try b.linear(&.{.{ .module = try Name.full(a, prefix, "self_attn.o_proj"), .n = c.hidden }}, c.heads * c.head_dim), + .q = try b.linear(&.{.{ .module = try Name.full(a, prefix, "self_attn.q_proj"), .n = c.qDim(), .w = w[@backingInt(M.q)] }}, c.hidden), + .kv = try b.linear(&.{ .{ .module = try Name.full(a, prefix, "self_attn.k_proj"), .n = c.kvDim(), .w = w[@backingInt(M.k)] }, .{ .module = try Name.full(a, prefix, "self_attn.v_proj"), .n = c.kvDim(), .w = w[@backingInt(M.v)] } }, c.hidden), + .out = try b.linear(&.{.{ .module = try Name.full(a, prefix, "self_attn.o_proj"), .n = c.hidden, .w = w[@backingInt(M.o)] }}, c.heads * c.head_dim), .q_norm = try b.tensor(try Name.full(a, prefix, "self_attn.q_norm.weight"), .bf16, &.{c.head_dim}), .k_norm = try b.tensor(try Name.full(a, prefix, "self_attn.k_norm.weight"), .bf16, &.{c.head_dim}), } }, }; - const gu = try b.linear(&.{ .{ .module = try Name.full(a, prefix, "mlp.gate_proj"), .n = c.intermediate }, .{ .module = try Name.full(a, prefix, "mlp.up_proj"), .n = c.intermediate } }, c.hidden); - const down = try b.linear(&.{.{ .module = try Name.full(a, prefix, "mlp.down_proj"), .n = c.hidden }}, c.intermediate); + const gu = try b.linear(&.{ .{ .module = try Name.full(a, prefix, "mlp.gate_proj"), .n = c.intermediate, .w = w[@backingInt(M.gate)] }, .{ .module = try Name.full(a, prefix, "mlp.up_proj"), .n = c.intermediate, .w = w[@backingInt(M.up)] } }, c.hidden); + const down = try b.linear(&.{.{ .module = try Name.full(a, prefix, "mlp.down_proj"), .n = c.hidden, .w = w[@backingInt(M.down)] }}, c.intermediate); return .{ .norm = try b.tensor(try Name.full(a, prefix, "input_layernorm.weight"), .bf16, &.{c.hidden}), .post_norm = try b.tensor(try Name.full(a, prefix, "post_attention_layernorm.weight"), .bf16, &.{c.hidden}), .gu = gu, .down = down, .mixer = mixer }; } }; -const Part = struct { module: []const u8, n: usize }; +const Part = struct { module: []const u8, n: usize, w: Width }; fn checked(t: Tensor, dtype: cp.DType, shape: []const usize) !Tensor { if (t.dtype != dtype or t.rank != shape.len or !std.mem.eql(usize, t.shape[0..t.rank], shape)) return error.BadNativeTensor; @@ -133,7 +182,7 @@ fn checked(t: Tensor, dtype: cp.DType, shape: []const usize) !Tensor { } /// The graph from `checkpoint`, whose `repacked` tensors stay in their files: each Linear is staged, repacked, dropped. -pub fn load(allocator: std.mem.Allocator, device: mtl.Device, checkpoint: *const cp.Checkpoint, c: Config) !Weights { +pub fn load(allocator: std.mem.Allocator, device: mtl.Device, checkpoint: *const cp.Checkpoint, c: Config, ws: *const widths.Widths) !Weights { const prefix: []const u8 = if (checkpoint.has("language_model.model.embed_tokens.weight")) "language_model." else if (checkpoint.has("model.embed_tokens.weight")) "" else return error.MissingTextNamespace; if (checkpoint.has("language_model.model.embed_tokens.weight") and checkpoint.has("model.embed_tokens.weight")) return error.AmbiguousTextNamespace; var owned: std.ArrayList(mtl.Buffer) = .empty; @@ -141,17 +190,22 @@ pub fn load(allocator: std.mem.Allocator, device: mtl.Device, checkpoint: *const for (owned.items) |b| b.deinit(); owned.deinit(allocator); } + var runs: std.ArrayList([]projection.Run) = .empty; + errdefer { + for (runs.items) |r| allocator.free(r); + runs.deinit(allocator); + } const layers = try allocator.alloc(Layer, c.layers); errdefer allocator.free(layers); var stage = cp.Stage{ .device = device }; defer stage.deinit(); - const b = Builder{ .allocator = allocator, .device = device, .checkpoint = checkpoint, .prefix = prefix, .owned = &owned, .stage = &stage }; + const b = Builder{ .allocator = allocator, .device = device, .checkpoint = checkpoint, .prefix = prefix, .owned = &owned, .stage = &stage, .runs = &runs, .widths = ws }; for (layers, 0..) |*l, i| l.* = try b.layer(c, i); - const embed = try b.quant("model.embed_tokens", c.vocab, c.hidden); + const embed = try b.quant("model.embed_tokens", c.vocab, c.hidden, ws.embed); const norm = try b.tensor("model.norm.weight", .bf16, &.{c.hidden}); - const head = try b.linear(&.{.{ .module = "lm_head", .n = c.vocab }}, c.hidden); + const head = try b.linear(&.{.{ .module = "lm_head", .n = c.vocab, .w = ws.head }}, c.hidden); if (device.tensorUnits()) try laneOrder(allocator, device, layers, head); - return .{ .allocator = allocator, .layers = layers, .embed = embed, .norm = norm, .head = head, .owned = owned }; + return .{ .allocator = allocator, .layers = layers, .embed = embed, .norm = norm, .head = head, .owned = owned, .runs = runs }; } /// Tensor-unit projections leave load in lane_projection.reg_layout, the order quant_gpu's reg kernel reads. @@ -163,12 +217,18 @@ fn laneOrder(a: std.mem.Allocator, device: mtl.Device, layers: []const Layer, he .linear => |g| [_]Linear{ g.qkv, g.zba, g.out, l.gu, l.down }, .attention => |t| [_]Linear{ t.q, t.kv, t.out, l.gu, l.down }, }; - for (list) |p| try jobs.append(a, .{ .w = .{ .buf = p.words.buffer, .off = p.words.offset }, .n = p.n, .k = p.k }); + for (list) |p| try regJobs(a, &jobs, p); } - try jobs.append(a, .{ .w = .{ .buf = head.words.buffer, .off = head.words.offset }, .n = head.n, .k = head.k }); + try regJobs(a, &jobs, head); try lane.regOrder(device, jobs.items); } +/// Only 4-bit g64 matrices take the reg kernel; other widths keep lane_projection.weight_layout. +fn regJobs(a: std.mem.Allocator, jobs: *std.ArrayList(lane.Order), p: Linear) std.mem.Allocator.Error!void { + for (p.runs) |r| try regJobs(a, jobs, r.linear); + if (p.runs.len == 0 and p.bits == 4 and p.group == 64 and p.n % 32 == 0) try jobs.append(a, .{ .w = .{ .buf = p.words.buffer, .off = p.words.offset }, .n = p.n, .k = p.k }); +} + pub fn ref(t: Tensor) Ref { return .{ .buffer = t.buffer, .offset = t.offset }; } diff --git a/zig/src/families/qwen27/model.zig b/zig/src/families/qwen27/model.zig index 3a45318ef2..62ddf72cfc 100644 --- a/zig/src/families/qwen27/model.zig +++ b/zig/src/families/qwen27/model.zig @@ -4,6 +4,7 @@ const mtl = @import("metal"); const core = @import("core"); const cfg = @import("config.zig"); const affine = @import("affine.zig"); +const widths = @import("widths.zig"); const wts = @import("gpu_weights.zig"); const quant = @import("quant_gpu.zig"); const glue = @import("glue_gpu.zig"); @@ -25,6 +26,8 @@ pub const Model = struct { frame: Frame, /// The Neural Engine's share of prompt-chunk MLPs, when it loaded. ane: ?*ane_mlp.Share = null, + /// widths.Widths.tag: 0 for a uniform 4-bit g64 checkpoint. + widths_tag: u64 = 0, /// `share_fits(bytes)`: whether the Neural Engine's share fits beside what is loaded (null: wherever it can load). pub fn load(a: std.mem.Allocator, io: std.Io, dir: []const u8, rows: u32, share_fits: ?*const fn (u64) bool) !*Model { @@ -33,14 +36,24 @@ pub const Model = struct { const text = try std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(16 << 20)); defer a.free(text); const c = try cfg.parse(a, text, null); - try supported(a, c, text); + try supported(c); + var formats = try affine.Formats.init(a, text); + defer formats.deinit(); + var ws = try widths.Widths.resolve(a, &formats, c); + defer ws.deinit(); const m = try a.create(Model); errdefer a.destroy(m); m.allocator = a; m.config = c; + m.widths_tag = ws.tag(); m.device = try mtl.Device.init(); errdefer m.device.deinit(); if (!qualifiedGeneration(m.device.generation())) return error.QwenMetalQualificationRequired; + var why: widths.Why = .{}; + widths.admit(&ws, c, .{ .tensor_units = m.device.tensorUnits() }, &why) catch |err| { + std.log.err("qwen27: {s} (layer {?d}) is {d}-bit group {d}; this chip's projection kernels do not run it ({s})", .{ why.module, why.layer, why.width.bits, why.width.group, @errorName(why.reason) }); + return err; + }; m.queue = try m.device.queue(); errdefer m.queue.deinit(); m.checkpoint = core.checkpoint_metal.Checkpoint.init(a); @@ -49,7 +62,7 @@ pub const Model = struct { defer core.checkpoint_host.freeShardFiles(a, files); // projections stay in the files until load repacks them: the weights live once, in their lane layouts for (files) |file| try m.checkpoint.addFileSplit(m.device, file, "", "language_model.", wts.repacked); - m.weights = try wts.load(a, m.device, &m.checkpoint, c); + m.weights = try wts.load(a, m.device, &m.checkpoint, c, &ws); errdefer m.weights.deinit(); m.kernels = .{ .allocator = a, .device = m.device }; errdefer m.kernels.deinit(); @@ -59,7 +72,7 @@ pub const Model = struct { errdefer m.frame.deinit(); m.ane = null; const columns = if (std.c.getenv("TF_ANE_COLUMNS")) |v| try std.fmt.parseInt(u32, std.mem.span(v), 10) else ane_columns; - const share = columns > 0 and m.device.tensorUnits() and mtl.ane.available(); + const share = columns > 0 and m.device.tensorUnits() and mtl.ane.available() and m.weights.uniformMlp(); const fits = if (share_fits) |f| !share or f(ane_mlp.Share.bytes(c, columns)) else true; if (!fits) std.log.info("qwen27: the Neural Engine's prompt share ({d} MiB of fp16 weights) would take this Mac past its load limit; prompts run on the GPU alone", .{ane_mlp.Share.bytes(c, columns) >> 20}); if (share and fits) m.ane = ane_mlp.Share.init(a, m.device, m.queue, &m.checkpoint, c, columns) catch |err| blk: { @@ -83,21 +96,10 @@ pub const Model = struct { } }; -pub fn supported(a: std.mem.Allocator, c: cfg.Config, text: []const u8) !void { +/// The qualified shape and layer schedule; per-module widths are admitted by widths.admit once the chip is known. +pub fn supported(c: cfg.Config) !void { if (c.hidden != 5120 or c.intermediate != 17408 or c.layers != 64 or c.vocab != 248320 or c.heads != 24 or c.kv_heads != 4 or c.head_dim != 256 or c.k_heads != 16 or c.v_heads != 48 or c.dk != 128 or c.dv != 128 or c.conv_kernel != 4 or c.rope_dims != 64) return error.UnqualifiedQwenShape; for (c.kinds[0..c.layers], 0..) |kind, i| if (kind != (if ((i + 1) % 4 == 0) cfg.Kind.attention else cfg.Kind.linear)) return error.UnqualifiedLayerSchedule; - var formats = try affine.Formats.init(a, text); - defer formats.deinit(); - const declared = (try formats.resolve(null)) orelse return error.MissingAffineFormat; - if (declared.bits != 4 or declared.group_size != 64) return error.UnqualifiedQwenFormat; - if (formats.quant) |q| { - var it = q.iterator(); - while (it.next()) |item| { - if (std.mem.indexOfScalar(u8, item.key_ptr.*, '.') == null) continue; - const own = (try formats.resolve(item.key_ptr.*)) orelse return error.UnqualifiedQwenFormat; - if (own.bits != 4 or own.group_size != 64) return error.UnqualifiedQwenFormat; - } - } } fn qualifiedGeneration(generation: u32) bool { diff --git a/zig/src/families/qwen27/projection.zig b/zig/src/families/qwen27/projection.zig index a59eacd43f..cd3834e88d 100644 --- a/zig/src/families/qwen27/projection.zig +++ b/zig/src/families/qwen27/projection.zig @@ -3,10 +3,14 @@ const std = @import("std"); const mtl = @import("metal"); const checkpoint = @import("core").checkpoint_metal; const row = @import("core").row_projection; +const lane = @import("core").lane_projection; pub const Ref = struct { buffer: mtl.Buffer, offset: usize = 0 }; -pub const Linear = struct { words: Ref, pairs: Ref, n: u32, k: u32, tile: u32, slices: u32, raw: ?row.Weights = null }; -pub const Triple = struct { weight: checkpoint.Tensor, scales: checkpoint.Tensor, biases: checkpoint.Tensor }; +/// A device matrix; `runs` non-empty means fused members of different widths, each its own product into these rows. +pub const Linear = struct { words: Ref, pairs: Ref, n: u32, k: u32, tile: u32, slices: u32, raw: ?row.Weights = null, bits: u8 = 4, group: u32 = 64, runs: []const Run = &.{} }; +/// One product inside a mixed fused matrix, writing output columns [first, first + linear.n). +pub const Run = struct { linear: Linear, first: u32 }; +pub const Triple = struct { weight: checkpoint.Tensor, scales: checkpoint.Tensor, biases: checkpoint.Tensor, bits: u8 = 4, group: u32 = 64 }; pub fn split(n: usize, k: usize) !u32 { if (n == 0 or k == 0 or n % 4 != 0 or k % 64 != 0) return error.BadProjectionShape; @@ -31,17 +35,21 @@ pub fn promptSplit(n: usize, k: usize) ?u32 { pub fn prepare(device: mtl.Device, allocator: std.mem.Allocator, owned: *std.ArrayList(mtl.Buffer), parts: []const Triple, tile: u32) !Linear { if (tile != 32 or parts.len == 0) return error.BadProjectionShape; - const k = parts[0].weight.shape[1] * 8; + const bits = parts[0].bits; + const group = parts[0].group; + if (bits == 0 or bits > 8 or group == 0) return error.BadAffineMetadata; + const k = parts[0].weight.shape[1] * 32 / bits; var slices = try split(parts[0].weight.shape[0], k); var n: usize = 0; for (parts) |p| { const rows = p.weight.shape[0]; - if (p.weight.rank != 2 or p.weight.dtype != .u32 or p.weight.shape[1] * 8 != k or rows % 4 != 0 or try split(rows, k) != slices) return error.BadProjectionShape; - if (p.scales.dtype != .bf16 or p.biases.dtype != .bf16 or p.scales.rank != 2 or p.biases.rank != 2 or p.scales.shape[0] != rows or p.biases.shape[0] != rows or p.scales.shape[1] != k / 64 or p.biases.shape[1] != k / 64) return error.BadAffineMetadata; + if (p.bits != bits or p.group != group or k % group != 0) return error.BadAffineMetadata; + if (p.weight.rank != 2 or p.weight.dtype != .u32 or p.weight.shape[1] * 32 != k * bits or rows % 4 != 0 or try split(rows, k) != slices) return error.BadProjectionShape; + if (p.scales.dtype != .bf16 or p.biases.dtype != .bf16 or p.scales.rank != 2 or p.biases.rank != 2 or p.scales.shape[0] != rows or p.biases.shape[0] != rows or p.scales.shape[1] != k / group or p.biases.shape[1] != k / group) return error.BadAffineMetadata; n = try std.math.add(usize, n, rows); } - if (n % tile != 0) return error.BadProjectionShape; - if (!device.tensorUnits()) return prepareRows(device, allocator, owned, parts, n, k, tile, slices); + if (!device.tensorUnits() or n % tile != 0) return prepareRows(device, allocator, owned, parts, n, k, tile, slices); + if (bits != 4 or group != 64) return prepareLanes(device, allocator, owned, parts, n, k, tile, slices); slices = m5Split(n, k) orelse slices; const w = try device.buffer(n * k / 2, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); owned.append(allocator, w) catch |err| { @@ -83,30 +91,92 @@ test "stacked gate-up keeps the member K split rather than recomputing it from s fn prepareRows(device: mtl.Device, allocator: std.mem.Allocator, owned: *std.ArrayList(mtl.Buffer), parts: []const Triple, n: usize, k: usize, tile: u32, slices: u32) !Linear { const options = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; - const w = try device.buffer(n * k / 2, options); - owned.append(allocator, w) catch |err| { - w.deinit(); - return err; - }; - const scale = try device.buffer(n * (k / 64) * 2, options); - owned.append(allocator, scale) catch |err| { - scale.deinit(); - return err; - }; - const bias = try device.buffer(n * (k / 64) * 2, options); - owned.append(allocator, bias) catch |err| { - bias.deinit(); - return err; - }; + const bits: usize = parts[0].bits; + const ng = k / parts[0].group; + const row_bytes = k * bits / 8; + const w = try own(device, allocator, owned, n * row_bytes, options); + const scale = try own(device, allocator, owned, n * ng * 2, options); + const bias = try own(device, allocator, owned, n * ng * 2, options); var base: usize = 0; for (parts) |p| { const rows = p.weight.shape[0]; - @memcpy(w.contents()[base * k / 2 .. (base + rows) * k / 2], std.mem.sliceAsBytes(p.weight.host(u32))); - @memcpy(scale.contents()[base * (k / 64) * 2 .. (base + rows) * (k / 64) * 2], std.mem.sliceAsBytes(p.scales.host(u16))); - @memcpy(bias.contents()[base * (k / 64) * 2 .. (base + rows) * (k / 64) * 2], std.mem.sliceAsBytes(p.biases.host(u16))); + @memcpy(w.contents()[base * row_bytes .. (base + rows) * row_bytes], std.mem.sliceAsBytes(p.weight.host(u32))); + @memcpy(scale.contents()[base * ng * 2 .. (base + rows) * ng * 2], std.mem.sliceAsBytes(p.scales.host(u16))); + @memcpy(bias.contents()[base * ng * 2 .. (base + rows) * ng * 2], std.mem.sliceAsBytes(p.biases.host(u16))); base += rows; } - const raw = row.Weights{ .w = w, .scales = scale, .biases = bias, .n = n, .k = k, .group = 64, .bits = 4, .sum = .f32 }; + const raw = row.Weights{ .w = w, .scales = scale, .biases = bias, .n = n, .k = k, .group = parts[0].group, .bits = parts[0].bits, .sum = .f32 }; try raw.validate(); - return .{ .words = .{ .buffer = w }, .pairs = .{ .buffer = scale }, .n = @intCast(n), .k = @intCast(k), .tile = tile, .slices = slices, .raw = raw }; + return .{ .words = .{ .buffer = w }, .pairs = .{ .buffer = scale }, .n = @intCast(n), .k = @intCast(k), .tile = tile, .slices = slices, .raw = raw, .bits = parts[0].bits, .group = parts[0].group }; +} + +/// Widths other than 4/64 on tensor units: lane_projection.weight_layout byte copies for the non-reg lane kernel. +fn prepareLanes(device: mtl.Device, allocator: std.mem.Allocator, owned: *std.ArrayList(mtl.Buffer), parts: []const Triple, n: usize, k: usize, tile: u32, slices: u32) !Linear { + const options = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; + const layout = lane.Layout{ .n = n, .k = k, .format = .{ .bits = parts[0].bits, .group = parts[0].group }, .sk = slices }; + try layout.validate(); + const w = try own(device, allocator, owned, layout.weightWords() * 4, options); + const sb = try own(device, allocator, owned, layout.metadataElements() * 2, options); + var base: usize = 0; + for (parts) |p| { + packTiles(.{ .words = p.weight.host(u32), .scales = p.scales.host(u16), .biases = p.biases.host(u16), .rows = p.weight.shape[0] }, base, n, k, parts[0].bits, parts[0].group, w.slice(u32, layout.weightWords()), sb.slice(u16, layout.metadataElements())); + base += p.weight.shape[0]; + } + return .{ .words = .{ .buffer = w }, .pairs = .{ .buffer = sb }, .n = @intCast(n), .k = @intCast(k), .tile = tile, .slices = slices, .bits = parts[0].bits, .group = parts[0].group }; +} + +/// One member as stored: words [rows][k*bits/32], scales and biases [rows][k/group]. +pub const Member = struct { words: []const u32, scales: []const u16, biases: []const u16, rows: usize }; + +/// Copies a member into lane_projection.weight_layout rows [base, base + rows) of an n-row stack, bytes unchanged. +pub fn packTiles(p: Member, base: usize, n: usize, k: usize, bits: usize, group: usize, out_w: []u32, out_sb: []u16) void { + const ng = k / group; + const wp = group * bits / 32; + for (0..p.rows) |r| for (0..ng) |g| { + const at = base + r; + @memcpy(out_w[((at / 32 * ng + g) * 32 + at % 32) * wp ..][0..wp], p.words[(r * ng + g) * wp ..][0..wp]); + out_sb[(g * n + at) * 2] = p.scales[r * ng + g]; + out_sb[(g * n + at) * 2 + 1] = p.biases[r * ng + g]; + }; +} + +test "a stack packed member by member equals the core lane pack of the whole matrix" { + const a = std.testing.allocator; + const n = 64; + const k = 256; + var rng = std.Random.DefaultPrng.init(7); + for ([_]u8{ 4, 6, 8 }) |bits| for ([_]usize{ 32, 64, 128 }) |group| { + const l = lane.Layout{ .n = n, .k = k, .format = .{ .bits = bits, .group = group } }; + const w = try a.alloc(u32, l.weightWords()); + defer a.free(w); + const sb = try a.alloc(u16, n * (k / group) * 2); + defer a.free(sb); + rng.random().bytes(std.mem.sliceAsBytes(w)); + rng.random().bytes(std.mem.sliceAsBytes(sb)); + const scales = sb[0 .. n * (k / group)]; + const biases = sb[n * (k / group) ..]; + const want_w = try a.alloc(u32, w.len); + defer a.free(want_w); + const want_sb = try a.alloc(u16, l.metadataElements()); + defer a.free(want_sb); + try lane.pack(l, w, scales, biases, want_w, want_sb); + const got_w = try a.alloc(u32, w.len); + defer a.free(got_w); + const got_sb = try a.alloc(u16, want_sb.len); + defer a.free(got_sb); + const row_words = w.len / n; + const ng = k / group; + for ([_][2]usize{ .{ 0, 32 }, .{ 32, 16 }, .{ 48, 16 } }) |part| packTiles(.{ .words = w[part[0] * row_words ..][0 .. part[1] * row_words], .scales = scales[part[0] * ng ..][0 .. part[1] * ng], .biases = biases[part[0] * ng ..][0 .. part[1] * ng], .rows = part[1] }, part[0], n, k, bits, group, got_w, got_sb); + try std.testing.expectEqualSlices(u32, want_w, got_w); + try std.testing.expectEqualSlices(u16, want_sb, got_sb); + }; +} + +fn own(device: mtl.Device, allocator: std.mem.Allocator, owned: *std.ArrayList(mtl.Buffer), bytes: usize, options: usize) !mtl.Buffer { + const b = try device.buffer(bytes, options); + owned.append(allocator, b) catch |err| { + b.deinit(); + return err; + }; + return b; } diff --git a/zig/src/families/qwen27/quant_gpu.zig b/zig/src/families/qwen27/quant_gpu.zig index 7632937dd2..63c0a8974c 100644 --- a/zig/src/families/qwen27/quant_gpu.zig +++ b/zig/src/families/qwen27/quant_gpu.zig @@ -10,7 +10,7 @@ pub const Ref = projection.Ref; /// Group sums the input's producer wrote: XS[group * stride + row], in the lane projection's order. pub const Sums = struct { ref: Ref, stride: u32 }; -const Key = struct { kind: Kind, n: u32 = 0, k: u32, slices: u32 = 0, tile: u32 = 0, tmr: u32 = 0, edge: bool = false, skip: u32 = 0, cut: [2]u32 = .{ 0, 0 } }; +const Key = struct { kind: Kind, n: u32 = 0, k: u32, slices: u32 = 0, tile: u32 = 0, tmr: u32 = 0, edge: bool = false, skip: u32 = 0, cut: [2]u32 = .{ 0, 0 }, bits: u8 = 4, group: u32 = 64 }; pub const Kind = enum { projection, xsum, norm, input_norm, mlp, ane_input, mlp_ane, add_ane }; pub const Kernels = struct { @@ -92,6 +92,11 @@ pub const Kernels = struct { } fn quantImpl(k: *Kernels, input_encoder: mtl.ComputeEncoder, p: projection.Linear, x: Ref, xs: ?Sums, dims: Ref, y: Ref, rows: u32, skip: u32, cut: [2]u32) !void { + if (p.runs.len > 0) return k.quantRuns(input_encoder, p, x, xs, dims, y, rows, skip, cut); + return k.quantOne(input_encoder, p, x, xs, dims, y, rows, skip, cut); + } + + fn quantOne(k: *Kernels, input_encoder: mtl.ComputeEncoder, p: projection.Linear, x: Ref, xs: ?Sums, dims: Ref, y: Ref, rows: u32, skip: u32, cut: [2]u32) !void { if ((skip > 0 or cut[1] > 0) and p.raw != null) return error.BadCoreProjectionRows; if (skip > 0 and (p.raw != null or skip >= p.n / 64)) return error.BadCoreProjectionRows; var span = profile.begin(k.profiler, input_encoder, .dense); @@ -126,13 +131,16 @@ pub const Kernels = struct { try span.end(); return; } - const key = Key{ .kind = .projection, .n = p.n, .k = p.k, .slices = sk, .skip = skip, .cut = cut }; + const reg = p.bits == 4 and p.group == 64; + if (!reg and (skip > 0 or cut[1] > 0)) return error.BadCoreProjectionRows; + const key = Key{ .kind = .projection, .n = p.n, .k = p.k, .slices = if (reg) sk else p.slices, .skip = skip, .cut = cut, .bits = p.bits, .group = p.group }; + const sums: ?Sums = if (reg) xs else null; // the non-reg kernel sums its own groups const half = p.n / 64; const ranges = [_][2]usize{ .{ skip, half - skip }, .{ half + skip, half - skip } }; const entry = try k.lanes.getOrPut(k.allocator, key); if (!entry.found_existing) { errdefer _ = k.lanes.remove(key); - entry.value_ptr.* = try lane.Projection.init(k.allocator, k.device, .{ .n = p.n, .k = p.k, .format = .{ .bits = 4, .group = 64 }, .sk = sk, .ranges = if (skip > 0) &ranges else &.{}, .groups = if (cut[1] > 0) .{ cut[0], cut[1] } else null, .output = if (cut[1] > 0) .f32 else .bf16, .precompute_sums = true, .reg = true }); + entry.value_ptr.* = try lane.Projection.init(k.allocator, k.device, .{ .n = p.n, .k = p.k, .format = .{ .bits = p.bits, .group = p.group }, .sk = key.slices, .ranges = if (skip > 0) &ranges else &.{}, .groups = if (cut[1] > 0) .{ cut[0], cut[1] } else null, .output = if (cut[1] > 0) .f32 else .bf16, .precompute_sums = reg, .reg = reg }); } const step: u32 = if (entry.value_ptr.reg) lane.reg_rows else 16; var first: u32 = 0; @@ -141,8 +149,45 @@ pub const Kernels = struct { const w = lane.Ref{ .buf = p.words.buffer, .off = p.words.offset }; const sb = lane.Ref{ .buf = p.pairs.buffer, .off = p.pairs.offset }; const yr = lane.Ref{ .buf = y.buffer, .off = y.offset + @as(usize, first) * p.n * @as(usize, if (cut[1] > 0) 4 else 2) }; - if (xs) |s| try entry.value_ptr.encodeSums(e, xr, w, sb, .{ .ref = .{ .buf = s.ref.buffer, .off = s.ref.offset + @as(usize, first) * 4 }, .stride = s.stride }, yr, @min(rows - first, step)) else try entry.value_ptr.encode(e, xr, w, sb, yr, @min(rows - first, step)); + if (sums) |s| try entry.value_ptr.encodeSums(e, xr, w, sb, .{ .ref = .{ .buf = s.ref.buffer, .off = s.ref.offset + @as(usize, first) * 4 }, .stride = s.stride }, yr, @min(rows - first, step)) else try entry.value_ptr.encode(e, xr, w, sb, yr, @min(rows - first, step)); + } + try span.end(); + } + + /// A mixed stack: each member its own product into its columns; row kernels write the fused rows at a stride. + fn quantRuns(k: *Kernels, input_encoder: mtl.ComputeEncoder, p: projection.Linear, x: Ref, xs: ?Sums, dims: Ref, y: Ref, rows: u32, skip: u32, cut: [2]u32) !void { + if (skip > 0 or cut[1] > 0 or rows == 0 or rows > 128) return error.BadCoreProjectionRows; + for (p.runs) |run| { + if (run.linear.raw) |raw| { + try k.rowStrided(input_encoder, raw, x, .{ .buffer = y.buffer, .offset = y.offset + @as(usize, run.first) * 2 }, rows, p.n); + continue; + } + for (0..rows) |r| { // lane kernels take one row a launch until they write at a stride + const xr = Ref{ .buffer = x.buffer, .offset = x.offset + r * p.k * 2 }; + const sr: ?Sums = if (xs) |s| .{ .ref = .{ .buffer = s.ref.buffer, .offset = s.ref.offset + r * 4 }, .stride = s.stride } else null; + const yr = Ref{ .buffer = y.buffer, .offset = y.offset + (r * p.n + run.first) * 2 }; + try k.quantOne(input_encoder, run.linear, xr, sr, dims, yr, 1, 0, .{ 0, 0 }); + } } + } + + /// The row kernel at any width, output row m at y + m * stride values; a row's bits do not depend on the window. + fn rowStrided(k: *Kernels, input_encoder: mtl.ComputeEncoder, raw: row.Weights, x: Ref, y: Ref, rows: u32, stride: u32) !void { + var span = profile.begin(k.profiler, input_encoder, .dense); + span.shape = .{ .n = @intCast(raw.n), .k = @intCast(raw.k), .sk = 1 }; + errdefer span.cancel(); + const e = span.encoder; + if (k.row_pipelines == null) k.row_pipelines = try row.Pipelines.load(k.device); + const call = try row.callStrided(&k.row_pipelines.?, raw, rows, stride); + e.setPipeline(call.pipeline); + e.setBuffer(x.buffer, x.offset, 0); + e.setBuffer(raw.w, raw.w_off, 1); + e.setBuffer(raw.scales, raw.s_off, 2); + e.setBuffer(raw.biases, raw.b_off, 3); + e.setValue(call.dims, 4); + e.setBuffer(y.buffer, y.offset, 5); + e.dispatchGroups(.{ .width = call.groups }, .{ .width = call.threads }); + e.barrier(); try span.end(); } diff --git a/zig/src/families/qwen27/qwen27.zig b/zig/src/families/qwen27/qwen27.zig index a1ba407bbd..423e1b7fd7 100644 --- a/zig/src/families/qwen27/qwen27.zig +++ b/zig/src/families/qwen27/qwen27.zig @@ -2,6 +2,7 @@ pub const dflash = @import("dflash/root.zig"); pub const config = @import("config.zig"); pub const affine = @import("affine.zig"); +pub const widths = @import("widths.zig"); pub const checkpoint = @import("checkpoint.zig"); pub const weights = @import("weights.zig"); pub const attention = @import("attention.zig"); @@ -49,6 +50,7 @@ test { _ = prompt_state; _ = config; _ = affine; + _ = widths; _ = checkpoint; _ = weights; _ = @import("weights_test.zig"); diff --git a/zig/src/families/qwen27/widths.zig b/zig/src/families/qwen27/widths.zig new file mode 100644 index 0000000000..6de398e809 --- /dev/null +++ b/zig/src/families/qwen27/widths.zig @@ -0,0 +1,305 @@ +//! Affine widths of the modules the engine loads, admitted against the projection kernels this chip runs. +const std = @import("std"); +const core = @import("core"); +const cfg = @import("config.zig"); +const affine = @import("affine.zig"); +const glue = @import("glue.zig"); +const row = core.row_projection; +const lane = core.lane_projection; + +/// One affine width, as core.quant.mlx.Width: `bits` per weight, a scale and bias a `group` of K. +pub const Width = struct { bits: u8, group: u16 }; +/// The qualified format: today's kernels and launches, unchanged. +pub const uniform = Width{ .bits = 4, .group = 64 }; + +pub fn isUniform(w: Width) bool { + return w.bits == uniform.bits and w.group == uniform.group; +} + +/// A layer's projections; a layer stores entries for both mixers and uses those of its kind. +pub const Member = enum { gate, up, down, qkv, z, b, a, out, q, k, v, o }; +const members = @backingInt(Member.o) + 1; + +pub fn leaf(m: Member) []const u8 { + return switch (m) { + .gate => "mlp.gate_proj", + .up => "mlp.up_proj", + .down => "mlp.down_proj", + .qkv => "linear_attn.in_proj_qkv", + .z => "linear_attn.in_proj_z", + .b => "linear_attn.in_proj_b", + .a => "linear_attn.in_proj_a", + .out => "linear_attn.out_proj", + .q => "self_attn.q_proj", + .k => "self_attn.k_proj", + .v => "self_attn.v_proj", + .o => "self_attn.o_proj", + }; +} + +/// Output rows and input width of a member. +pub fn shape(c: cfg.Config, m: Member) [2]usize { + return switch (m) { + .gate, .up => .{ c.intermediate, c.hidden }, + .down => .{ c.hidden, c.intermediate }, + .qkv => .{ c.gdnQkvDim(), c.hidden }, + .z => .{ c.gdnValueDim(), c.hidden }, + .b, .a => .{ c.v_heads, c.hidden }, + .out => .{ c.hidden, c.gdnValueDim() }, + .q => .{ c.qDim(), c.hidden }, + .k, .v => .{ c.kvDim(), c.hidden }, + .o => .{ c.hidden, c.heads * c.head_dim }, + }; +} + +/// The device matrices of a layer: members fused into one output, in output-row order. +pub fn stacks(kind: cfg.Kind) []const []const Member { + return switch (kind) { + .linear => &.{ &.{ .gate, .up }, &.{.down}, &.{.qkv}, &.{ .z, .b, .a }, &.{.out} }, + .attention => &.{ &.{ .gate, .up }, &.{.down}, &.{.q}, &.{ .k, .v }, &.{.o} }, + }; +} + +/// A stack's products as [first, count]: one for a stack of one width, else one a member, as the reference runs them. +pub fn runs(ws: []const Width, out: [][2]usize) [][2]usize { + for (ws[1..]) |w| if (!std.meta.eql(ws[0], w)) { + for (out[0..ws.len], 0..) |*o, i| o.* = .{ i, 1 }; + return out[0..ws.len]; + }; + out[0] = .{ 0, ws.len }; + return out[0..1]; +} + +pub const Widths = struct { + allocator: std.mem.Allocator, + layers: [][members]Width, + embed: Width, + head: Width, + + /// Resolves only loaded modules, so draft-head (MTP) and vision overrides never decide admission. + pub fn resolve(a: std.mem.Allocator, f: *const affine.Formats, c: cfg.Config) !Widths { + if ((try f.resolve(null)) == null) return error.MissingAffineFormat; + const layers = try a.alloc([members]Width, c.layers); + errdefer a.free(layers); + var name: [96]u8 = undefined; + for (layers, 0..) |*l, i| { + l.* = @splat(uniform); + for (stacks(c.kind(i))) |stack| for (stack) |m| { + l[@backingInt(m)] = try one(f, try std.fmt.bufPrint(&name, "model.layers.{d}.{s}", .{ i, leaf(m) })); + }; + } + return .{ .allocator = a, .layers = layers, .embed = try one(f, "model.embed_tokens"), .head = try one(f, "lm_head") }; + } + + /// Every module 4-bit g64, for synthetic graphs. + pub fn uniformAll(a: std.mem.Allocator, layers: usize) !Widths { + const all = try a.alloc([members]Width, layers); + for (all) |*l| l.* = @splat(uniform); + return .{ .allocator = a, .layers = all, .embed = uniform, .head = uniform }; + } + + pub fn deinit(w: *Widths) void { + w.allocator.free(w.layers); + } + + pub fn of(w: *const Widths, layer: usize, m: Member) Width { + return w.layers[layer][@backingInt(m)]; + } + + /// 0 when every module is 4-bit g64, else a nonzero hash of every module's width. + pub fn tag(w: *const Widths) u64 { + var h = std.hash.Wyhash.init(0x27); + var mixed = false; + for (w.layers) |l| for (l) |x| { + mixed = mixed or !isUniform(x); + h.update(&.{ x.bits, @as(u8, @truncate(x.group)), @as(u8, @truncate(x.group >> 8)) }); + }; + for ([_]Width{ w.embed, w.head }) |x| { + mixed = mixed or !isUniform(x); + h.update(&.{ x.bits, @as(u8, @truncate(x.group)), @as(u8, @truncate(x.group >> 8)) }); + } + return if (mixed) h.final() | 1 else 0; + } +}; + +fn one(f: *const affine.Formats, path: []const u8) !Width { + const s = (try f.resolve(path)) orelse return error.UnqualifiedQwenFormat; + return .{ .bits = s.bits, .group = @intCast(s.group_size) }; +} + +/// The timed window costs' model shape; mixed widths add their tag, so checkpoints of one shape keep their own costs. +pub fn costShape(c: cfg.Config, tag: u64, buf: []u8, rows: usize, drafted: bool) ![]const u8 { + const s = try std.fmt.bufPrint(buf, "{d}/{d}/{d}/{d}/{d}/{d}", .{ c.layers, c.hidden, c.intermediate, c.vocab, rows, @intFromBool(drafted) }); + if (tag == 0) return s; + return buf[0 .. s.len + (try std.fmt.bufPrint(buf[s.len..], "/w{x:0>16}", .{tag})).len]; +} + +/// What the chip's projection kernels run: the lane kernels on tensor units, the row kernels before them. +pub const Chip = struct { tensor_units: bool, row_bits: []const u8 = &row.bit_widths }; + +/// The first module admission refused, for the load's error message. +pub const Why = struct { layer: ?usize = null, module: []const u8 = "", width: Width = uniform, reason: anyerror = error.None }; + +pub fn admit(w: *const Widths, c: cfg.Config, chip: Chip, why: *Why) !void { + var buf: [3][2]usize = undefined; + for (0..c.layers) |i| for (stacks(c.kind(i))) |stack| { + var ws: [3]Width = undefined; + for (stack, 0..) |m, j| ws[j] = w.of(i, m); + const list = runs(ws[0..stack.len], &buf); + if (list.len == 1 and isUniform(ws[0])) continue; + for (list) |r| { + var n: usize = 0; + for (stack[r[0]..][0..r[1]]) |m| n += shape(c, m)[0]; + check(chip, n, shape(c, stack[r[0]])[1], ws[r[0]]) catch |err| { + why.* = .{ .layer = i, .module = leaf(stack[r[0]]), .width = ws[r[0]], .reason = err }; + return error.UnqualifiedQwenFormat; + }; + } + }; + glue.checkEmbed(.{ .width = @intCast(c.hidden), .rows = 1, .vocab = @intCast(c.vocab), .bits = w.embed.bits, .group = w.embed.group, .metadata = .bf16 }) catch |err| { + why.* = .{ .module = "embed_tokens", .width = w.embed, .reason = err }; + return error.UnqualifiedQwenFormat; + }; + if (!isUniform(w.head)) check(chip, c.vocab, c.hidden, w.head) catch |err| { + why.* = .{ .module = "lm_head", .width = w.head, .reason = err }; + return error.UnqualifiedQwenFormat; + }; +} + +/// One run's kernel: 4/64 keeps today's choice; other widths take the non-reg lane kernel or the row kernel. +fn check(chip: Chip, n: usize, k: usize, w: Width) !void { + _ = try @import("projection.zig").split(n, k); + // The lane sums and odd word counts assume group 64 off the 4-bit kernel. + if (chip.tensor_units and w.bits != 4 and w.group != 64) return error.UnsupportedLaneGroup; + if (chip.tensor_units and n % 32 == 0) return (lane.Layout{ .n = n, .k = k, .format = .{ .bits = w.bits, .group = w.group }, .precompute_sums = isUniform(w), .reg = isUniform(w) }).validate(); + if (isUniform(w) and n % 8 == 0) return; + if (std.mem.indexOfScalar(u8, chip.row_bits, w.bits) == null) return error.UnsupportedRowBits; + return (row.Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = n, .k = k, .group = w.group }).validate(); +} + +const all_rows = [_]u8{ 2, 3, 4, 5, 6, 8 }; + +fn testConfig(a: std.mem.Allocator, quantization: []const u8) ![]u8 { + return std.fmt.allocPrint(a, + \\{{"model_type":"qwen3_5","text_config":{{"model_type":"qwen3_5_text","hidden_size":5120,"intermediate_size":17408, + \\"num_hidden_layers":64,"num_attention_heads":24,"num_key_value_heads":4,"head_dim":256,"vocab_size":248320, + \\"linear_num_key_heads":16,"linear_num_value_heads":48,"linear_key_head_dim":128,"linear_value_head_dim":128, + \\"linear_conv_kernel_dim":4,"partial_rotary_factor":0.25,"eos_token_id":248044}},"quantization":{s}}} + , .{quantization}); +} + +fn testAdmit(quantization: []const u8, chip: Chip, why: *Why) !void { + const a = std.testing.allocator; + const text = try testConfig(a, quantization); + defer a.free(text); + const c = try cfg.parse(a, text, null); + var f = try affine.Formats.init(a, text); + defer f.deinit(); + var w = try Widths.resolve(a, &f, c); + defer w.deinit(); + return admit(&w, c, chip, why); +} + +/// TensorFold/Qwen3.8-27B-oQ4 override shapes: 4-bit g64, 5-bit GDN, MLP and attention members, one 6-bit v. +const oq4 = + \\{"group_size":64,"bits":4,"mode":"affine", + \\"language_model.model.layers.37.linear_attn.in_proj_a":{"group_size":64,"bits":5}, + \\"language_model.model.layers.36.linear_attn.in_proj_b":{"group_size":64,"bits":5}, + \\"language_model.model.layers.26.linear_attn.in_proj_z":{"group_size":64,"bits":5}, + \\"language_model.model.layers.62.linear_attn.out_proj":{"group_size":64,"bits":5}, + \\"language_model.model.layers.25.mlp.down_proj":{"group_size":64,"bits":5}, + \\"language_model.model.layers.27.self_attn.k_proj":{"group_size":64,"bits":5}, + \\"language_model.model.layers.27.self_attn.o_proj":{"group_size":64,"bits":5}, + \\"language_model.model.layers.27.self_attn.v_proj":{"group_size":64,"bits":6}} +; + +/// The override shapes of TensorFold/Qwen3.8-27B-oQ2: 2-bit g64 with an 8-bit embedding and 3/5/6-bit members. +const oq2 = + \\{"group_size":64,"bits":2,"mode":"affine", + \\"language_model.model.embed_tokens":{"group_size":64,"bits":8}, + \\"language_model.model.layers.50.linear_attn.in_proj_a":{"group_size":64,"bits":3}, + \\"language_model.model.layers.44.linear_attn.in_proj_a":{"group_size":64,"bits":5}, + \\"language_model.model.layers.42.linear_attn.in_proj_b":{"group_size":64,"bits":5}, + \\"language_model.model.layers.12.linear_attn.in_proj_z":{"group_size":64,"bits":5}, + \\"language_model.model.layers.10.mlp.down_proj":{"group_size":64,"bits":5}, + \\"language_model.model.layers.35.self_attn.k_proj":{"group_size":64,"bits":5}, + \\"language_model.model.layers.35.self_attn.v_proj":{"group_size":64,"bits":6}} +; + +test "uniform 4-bit g64 admits on both chip kinds and ignores draft-head and vision overrides" { + var why: Why = .{}; + const q = "{\"bits\":4,\"group_size\":64,\"mtp.layers.0.mlp.gate_proj\":{\"bits\":8,\"group_size\":64},\"vision_tower.blocks.0.attn.qkv\":false}"; + try testAdmit(q, .{ .tensor_units = false }, &why); + try testAdmit(q, .{ .tensor_units = true }, &why); +} + +test "oQ4 and oQ2 formats admit once the row kernels unpack every width they declare" { + var why: Why = .{}; + try testAdmit(oq4, .{ .tensor_units = false, .row_bits = &all_rows }, &why); + try testAdmit(oq2, .{ .tensor_units = false, .row_bits = &all_rows }, &why); +} + +test "a width the row kernels do not unpack refuses naming the module, never guessing four bits" { + var why: Why = .{}; + try std.testing.expectError(error.UnqualifiedQwenFormat, testAdmit(oq4, .{ .tensor_units = false, .row_bits = &.{ 2, 4, 6, 8 } }, &why)); + try std.testing.expectEqual(@as(?usize, 25), why.layer); + try std.testing.expectEqualStrings("mlp.down_proj", why.module); + try std.testing.expectEqual(Width{ .bits = 5, .group = 64 }, why.width); + try std.testing.expectError(error.UnqualifiedQwenFormat, testAdmit(oq2, .{ .tensor_units = false, .row_bits = &.{ 2, 4, 6, 8 } }, &why)); + try std.testing.expectEqual(@as(u8, 5), why.width.bits); +} + +test "a dense projection and an unsupported global group refuse" { + var why: Why = .{}; + try std.testing.expectError(error.UnqualifiedQwenFormat, testAdmit("{\"bits\":4,\"group_size\":64,\"model.layers.3.self_attn.q_proj\":false}", .{ .tensor_units = false }, &why)); + try std.testing.expectError(error.BadAffine, testAdmit("{\"bits\":4,\"group_size\":16}", .{ .tensor_units = false }, &why)); +} + +test "a 4/64 run off the 32-row lanes, such as a lone 48-row in_proj_b, takes the row kernels on tensor units" { + try check(.{ .tensor_units = true }, 48, 5120, uniform); + try check(.{ .tensor_units = true }, 64, 5120, uniform); +} + +test "a mixed stack runs every member alone, even beside a member of its own width" { + var buf: [3][2]usize = undefined; + const four = uniform; + const five = Width{ .bits = 5, .group = 64 }; + try std.testing.expectEqualSlices([2]usize, &.{.{ 0, 3 }}, runs(&.{ four, four, four }, &buf)); + try std.testing.expectEqualSlices([2]usize, &.{ .{ 0, 1 }, .{ 1, 1 }, .{ 2, 1 } }, runs(&.{ five, four, four }, &buf)); + try std.testing.expectEqualSlices([2]usize, &.{ .{ 0, 1 }, .{ 1, 1 }, .{ 2, 1 } }, runs(&.{ four, five, four }, &buf)); +} + +test "timed costs keep a uniform checkpoint's shape and give each mixed one its own" { + const a = std.testing.allocator; + var tags: [3]u64 = undefined; + for ([_][]const u8{ "{\"bits\":4,\"group_size\":64}", oq4, oq2 }, &tags) |quantization, *t| { + const text = try testConfig(a, quantization); + defer a.free(text); + const c = try cfg.parse(a, text, null); + var f = try affine.Formats.init(a, text); + defer f.deinit(); + var w = try Widths.resolve(a, &f, c); + defer w.deinit(); + t.* = w.tag(); + var buf: [96]u8 = undefined; + const got = try costShape(c, t.*, &buf, 16, true); + try std.testing.expectEqual(t.* == 0, std.mem.eql(u8, got, "64/5120/17408/248320/16/1")); + } + try std.testing.expectEqual(@as(u64, 0), tags[0]); + try std.testing.expect(tags[1] != 0 and tags[2] != 0 and tags[1] != tags[2]); +} + +test "resolved widths follow overrides under any checkpoint prefix" { + const a = std.testing.allocator; + const text = try testConfig(a, oq4); + defer a.free(text); + const c = try cfg.parse(a, text, null); + var f = try affine.Formats.init(a, text); + defer f.deinit(); + var w = try Widths.resolve(a, &f, c); + defer w.deinit(); + try std.testing.expectEqual(Width{ .bits = 6, .group = 64 }, w.of(27, .v)); + try std.testing.expectEqual(Width{ .bits = 5, .group = 64 }, w.of(27, .k)); + try std.testing.expectEqual(uniform, w.of(27, .q)); + try std.testing.expectEqual(uniform, w.head); +} diff --git a/zig/src/native/qwen27_lanes.zig b/zig/src/native/qwen27_lanes.zig index c177ef95fa..e09a508a8c 100644 --- a/zig/src/native/qwen27_lanes.zig +++ b/zig/src/native/qwen27_lanes.zig @@ -148,9 +148,8 @@ pub const Metal = struct { pub fn measure(b: *Metal, io: std.Io) !void { const r = b.runner; if (b.owners[0] != null or r.active != null) return error.StreamBusy; - const c = r.model.config; var shape: [128]u8 = undefined; - const parts = [_][]const u8{ "qwen27-lanes", std.mem.span(r.model.device.name()), try std.fmt.bufPrint(&shape, "{d}/{d}/{d}/{d}/{d}/{d}", .{ c.layers, c.hidden, c.intermediate, c.vocab, max_rows, @intFromBool(b.draft != null) }) }; + const parts = [_][]const u8{ "qwen27-lanes", std.mem.span(r.model.device.name()), try q.widths.costShape(r.model.config, r.model.widths_tag, &shape, max_rows, b.draft != null) }; const k = lanes.cost_cache.key(b.gpa, io, &parts) catch null; if (k) |key| if (lanes.cost_cache.load([max_rows]lanes.config.Cost, b.gpa, io, key)) |kept| { b.costs = kept; 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..5f32bd864a --- /dev/null +++ b/zig/tests/affine_ref.zig @@ -0,0 +1,358 @@ +//! The shared prompt products against float64 on synthetic weights, and every row's bits in every window width. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const mm = tf.affine_mm; +const sq = tf.row_projection.simd; + +const opts = mtl.ResourceOptions.shared; +const a = std.heap.page_allocator; +const unit: f64 = 0x1p-24; // fp32's unit roundoff +const wide = 129; // rows of the reference call: two 64-row tiles and one more +const windows = [_]usize{ 1, 2, 3, 8, 16, 17, 33, 64, 65 }; // 17, 33 and 65 cross the 16-, 32- and 64-row tiles + +fn toBf(v: f32) u16 { + const u: u32 = @bitCast(v); + return @truncate((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); +} + +fn fromBf(v: u16) f64 { + return @as(f32, @bitCast(@as(u32, v) << 16)); +} + +/// Synthetic MLX affine weights [n, k] (codes little-endian in bytes) and `wide` rows of x; `deq` is s q + b exactly. +const Synth = struct { + n: usize, + k: usize, + bits: u8, + group: usize, + x: mtl.Buffer, + w: mtl.Buffer, + s: mtl.Buffer, + b: mtl.Buffer, + deq: []f64, + + fn init(device: mtl.Device, n: usize, k: usize, bits: u8, group: usize, seed: u64) !Synth { + var prng = std.Random.DefaultPrng.init(seed); + const r = prng.random(); + const sn = n * (k / group); + const self = Synth{ .n = n, .k = k, .bits = bits, .group = group, .x = try device.buffer(wide * k * 2, opts), .w = try device.buffer(n * k * bits / 8, opts), .s = try device.buffer(sn * 2, opts), .b = try device.buffer(sn * 2, opts), .deq = try a.alloc(f64, n * k) }; + for (self.x.slice(u16, wide * k)) |*v| v.* = toBf(r.float(f32) * 2 - 1); + r.bytes(self.w.slice(u8, n * k * bits / 8)); + for (self.s.slice(u16, sn), self.b.slice(u16, sn)) |*s, *b| { + s.* = toBf((r.float(f32) + 0.5) / 64); + b.* = toBf(-(r.float(f32) + 0.5) * 0.15); + } + for (0..n) |j| for (0..k) |i| { + const g = j * (k / group) + i / group; + self.deq[j * k + i] = fromBf(self.s.slice(u16, sn)[g]) * @as(f64, @floatFromInt(self.code(j, i))) + fromBf(self.b.slice(u16, sn)[g]); + }; + return self; + } + + fn deinit(self: Synth) void { + for ([_]mtl.Buffer{ self.x, self.w, self.s, self.b }) |x| x.deinit(); + a.free(self.deq); + } + + fn code(self: Synth, j: usize, i: usize) u32 { + const bytes = self.w.slice(u8, self.n * self.k * self.bits / 8); + if (self.bits == 8) return bytes[j * self.k + i]; + if (self.bits == 2) return (bytes[(j * self.k + i) / 4] >> @intCast(2 * (i % 4))) & 3; + return (bytes[(j * self.k + i) / 2] >> @intCast(4 * (i % 2))) & 15; + } + + /// y = x W^T in float64, and the sizes the fp32 sums carry when codes enter as q + off (affine_mm's 4-bit: 128). + fn reference(self: Synth, off: f64) !Ref { + const n = self.n; + const k = self.k; + const groups = k / self.group; + const ref = Ref{ .want = try a.alloc(f64, wide * n), .mag = try a.alloc(f64, wide * n), .scaled = try a.alloc(f64, wide * n) }; + const x = self.x.slice(u16, wide * k); + const sc = self.s.slice(u16, n * groups); + const bi = self.b.slice(u16, n * groups); + const xd = try a.alloc(f64, k); + defer a.free(xd); + const xg = try a.alloc(f64, groups); // each group's sum of |x| + defer a.free(xg); + for (0..wide) |r| { + @memset(xg, 0); + for (xd, 0..) |*v, i| { + v.* = fromBf(x[r * k + i]); + xg[i / self.group] += @abs(v.*); + } + for (0..n) |j| { + var sum: f64 = 0; + var mag: f64 = 0; + var scaled: f64 = 0; + for (0..groups) |g| { + var qx: f64 = 0; + var qxs: f64 = 0; + for (g * self.group..(g + 1) * self.group) |i| { + const q: f64 = @floatFromInt(self.code(j, i)); + sum += xd[i] * self.deq[j * k + i]; + qx += (q + off) * @abs(xd[i]); + qxs += q * xd[i]; + } + const s = fromBf(sc[j * groups + g]); + mag += s * qx + @abs(fromBf(bi[j * groups + g]) - off * s) * xg[g]; + scaled += s * qxs; + } + ref.want[r * n + j] = sum; + ref.mag[r * n + j] = mag; + ref.scaled[r * n + j] = scaled; + } + } + return ref; + } +}; + +/// The float64 product, its terms' fp32 magnitudes, and its scales' share (a control: those scales 1% off must fail). +const Ref = struct { + want: []f64, + mag: []f64, + scaled: []f64, + fn deinit(r: Ref) void { + for ([_][]f64{ r.want, r.mag, r.scaled }) |v| a.free(v); + } +}; + +/// Every check's outcome: failures, and the largest difference as a fraction of its bound. +const Tally = struct { failed: usize = 0, worst: f64 = 0 }; + +/// One product's output: `rows` rows of n bf16 or fp32 values. +const Out = struct { + buf: mtl.Buffer, + f32_out: bool, + fn esize(o: Out) usize { + return if (o.f32_out) 4 else 2; + } + fn get(o: Out, i: usize) f64 { + return if (o.f32_out) o.buf.slice(f32, i + 1)[i] else fromBf(o.buf.slice(u16, i + 1)[i]); + } +}; + +/// The bound: fp32 sums of a group's products, its row sum and a chain of groups, plus one bf16 ulp of the value. +fn bound(want: f64, mag: f64, s: *const Synth, f32_out: bool, extra: usize) f64 { + const terms: f64 = @floatFromInt(s.group + s.k / s.group + 8 + extra); + return terms * unit * mag + (if (f32_out) 0 else @abs(want) * 0x1p-7) + 1e-30; +} + +/// `rows` rows of `out` against the float64 reference: the worst column's max abs and relative error, and the bound. +fn compare(t: *Tally, name: []const u8, out: Out, rows: usize, ref: Ref, s: *const Synth, extra: usize) void { + const n = s.n; + var col: usize = 0; + var col_abs: f64 = -1; + var col_rel: f64 = 0; + var frac: f64 = 0; + var ctl: f64 = 0; + var at: usize = 0; + for (0..n) |j| { + var abs: f64 = 0; + var top: f64 = 0; + for (0..rows) |r| { + const i = r * n + j; + const d = @abs(out.get(i) - ref.want[i]); + const tol = bound(ref.want[i], ref.mag[i], s, out.f32_out, extra); + const q = d / tol; + ctl = @max(ctl, @abs(out.get(i) - ref.want[i] - 0.01 * ref.scaled[i]) / tol); + if (!(d <= abs)) abs = d; + top = @max(top, @abs(ref.want[i])); + if (!(q <= frac)) { + frac = q; + at = i; + } + } + if (!(abs <= col_abs)) { + col_abs = abs; + col = j; + col_rel = abs / top; + } + } + t.worst = @max(t.worst, frac); + const ok = frac <= 1 and ctl > 1; + if (!ok) t.failed += 1; + std.debug.print("RESULT {s}: worst column {d}: max abs {e:.3}, relative {e:.3}; {d:.4} of its bound (row {d} col {d}), " ++ + "scales 1% off {d:.1}{s}\n", .{ name, col, col_abs, col_rel, frac, at / n, at % n, ctl, if (ok) "" else " FAIL" }); +} + +/// Rows [0, rows) of `got` equal `want` byte for byte; the first difference printed. +fn sameRows(t: *Tally, name: []const u8, m: usize, want: Out, got: Out, rows: usize, n: usize) bool { + const bytes = rows * n * want.esize(); + const w = want.buf.slice(u8, bytes); + const g = got.buf.slice(u8, bytes); + if (std.mem.eql(u8, w, g)) return true; + const i = std.mem.indexOfDiff(u8, w, g).? / want.esize(); + var differ: usize = 0; + var most: f64 = 0; + for (0..rows) |r| { + const row = r * n * want.esize(); + if (!std.mem.eql(u8, w[row..][0 .. n * want.esize()], g[row..][0 .. n * want.esize()])) differ += 1; + for (r * n..(r + 1) * n) |e| most = @max(most, @abs(got.get(e) - want.get(e))); + } + std.debug.print("RESULT {s}: {d}-row windows DIFFER in {d} of {d} rows, by up to {e:.3}; first row {d} col {d}: {d} against {d} in the {d}-row call\n", .{ name, m, differ, rows, most, i / n, i % n, got.get(i), want.get(i), wide }); + t.failed += 1; + return false; +} + +fn newOut(device: mtl.Device, s: *const Synth, f32_out: bool) !Out { + const o = Out{ .buf = try device.buffer(wide * s.n * 4, opts), .f32_out = f32_out }; + @memset(o.buf.slice(u8, wide * s.n * 4), 0xff); // NaN: a launch that writes nothing fails + return o; +} + +fn commit(cb: anytype, e: mtl.ComputeEncoder) !void { + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |why| { + std.debug.print("GPU failed: {s}\n", .{why}); + return error.GpuFailed; + } +} + +const Kernel = enum { dense, gather }; + +/// affine_mm over rows [start, start + m): its row sums, then the dense product or a gather with one expert. +fn affineWindow(e: mtl.ComputeEncoder, p: mm.Pipes, s: *const Synth, sums: mtl.Buffer, offs: mtl.Buffer, y: Out, kernel: Kernel, start: usize, m: usize) void { + const x = .{ .buf = s.x, .off = start * s.k * 2 }; + const xs = .{ .buf = sums, .off = start * (s.k / s.group) * 4 }; + const yo = .{ .buf = y.buf, .off = start * s.n * y.esize() }; + const z: usize = 0; + const q = .{ .w = .{ .buf = s.w, .off = z }, .s = .{ .buf = s.s, .off = z }, .b = .{ .buf = s.b, .off = z }, .n = @as(u32, @intCast(s.n)), .k = @as(u32, @intCast(s.k)) }; + mm.rowSums(e, p, @intCast(s.group), x, xs, @intCast(m), @intCast(s.k)); + switch (kernel) { + .dense => mm.dense(e, p, x, xs, q, yo, @intCast(m)), + .gather => mm.gather(e, p, x, xs, q, .{ .buf = offs, .off = z }, yo, @intCast(m), 1), + } +} + +/// affine_mm at one format: the wide call against float64, then every window's rows against the wide call's. +fn affine(t: *Tally, device: mtl.Device, queue: mtl.Queue, s: *const Synth, ref: Ref, f32_out: bool) !void { + const src = try mm.source(device, a, .{ .bits = s.bits, .group = @intCast(s.group), .out_f32 = f32_out }); + defer a.free(src); + const lib = try mtl.Library.fromSource(device, src, mtl.CompileOptions.mlx()); + defer lib.deinit(); + var p: mm.Pipes = undefined; + for (mm.names, 0..) |name, i| p[i] = try mtl.Pipeline.init(device, lib, name, false); + defer for (p) |x| x.deinit(); + const sums = try device.buffer(wide * (s.k / s.group) * 4, opts); + defer sums.deinit(); + const offs = try device.buffer(4, opts); + defer offs.deinit(); + offs.slice(i32, 1)[0] = 0; + var name: [96]u8 = undefined; + for ([_]Kernel{ .dense, .gather }) |kernel| { + const label = try std.fmt.bufPrint(&name, "affine_mm {s} q{d} g{d} {s} [{d}, {d}]", .{ @tagName(kernel), s.bits, s.group, if (f32_out) "fp32" else "bf16", s.n, s.k }); + const all = try newOut(device, s, f32_out); + defer all.buf.deinit(); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + affineWindow(e, p, s, sums, offs, all, kernel, 0, wide); + try commit(cb, e); + compare(t, label, all, wide, ref, s, 0); + var same = true; + for (windows) |m| { + const got = try newOut(device, s, f32_out); + defer got.buf.deinit(); + const wb = queue.commandBuffer(); + const we = wb.compute(.serial); + var start: usize = 0; + while (start + m <= wide) : (start += m) affineWindow(we, p, s, sums, offs, got, kernel, start, m); + try commit(wb, we); + same = sameRows(t, label, m, all, got, wide / m * m, s.n) and same; + } + if (same) std.debug.print("RESULT {s}: every row's bits equal in windows of 1, 2, 3, 8, 16, 17, 33, 64 and 65 rows\n", .{label}); + } +} + +/// simd_qmm over rows [start, start + m), the scalar twin when `scalar` and it takes m rows. +fn simdWindow(e: mtl.ComputeEncoder, p: *const sq.Pipelines, s: *const Synth, y: Out, scalar: bool, start: usize, m: usize) !void { + const mat = sq.Matrix{ .w = s.w, .scales = s.s, .biases = s.b, .n = s.n, .k = s.k }; + sq.encode(e, try sq.call(p, s.n, s.k, m, scalar), mat, s.x, start * s.k * 2, y.buf, start * s.n * 2); +} + +/// simd_qmm (4-bit, group 64): tiles at every window against float64 and the wide call; the twin at 1 and 2 rows. +fn simd(t: *Tally, device: mtl.Device, queue: mtl.Queue, p: *const sq.Pipelines, s: *const Synth, ref: Ref) !void { + var name: [96]u8 = undefined; + const label = try std.fmt.bufPrint(&name, "simd_qmm tiles q4 g64 bf16 [{d}, {d}]", .{ s.n, s.k }); + const all = try newOut(device, s, false); + defer all.buf.deinit(); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + try simdWindow(e, p, s, all, false, 0, wide); + try commit(cb, e); + const extra = std.math.log2_int(u32, sq.splits(s.n)) + 1; // the chunks' fixed tree + compare(t, label, all, wide, ref, s, extra); + var same = true; + for (windows) |m| { + const got = try newOut(device, s, false); + defer got.buf.deinit(); + const wb = queue.commandBuffer(); + const we = wb.compute(.serial); + var start: usize = 0; + while (start + m <= wide) : (start += m) try simdWindow(we, p, s, got, false, start, m); + try commit(wb, we); + same = sameRows(t, label, m, all, got, wide / m * m, s.n) and same; + } + if (same) std.debug.print("RESULT {s}: every row's bits equal in windows of 1, 2, 3, 8, 16, 17, 33, 64 and 65 rows\n", .{label}); + // the twin: callers take it only where twinExact says it equals the tiles, so a difference fails only then + const exact = try sq.twinExact(p, device, .{ .w = s.w, .scales = s.s, .biases = s.b, .n = s.n, .k = s.k }); + for ([_]usize{ 1, 2 }) |m| { + const twin = try std.fmt.bufPrint(&name, "simd_qmm twin q4 g64 bf16 [{d}, {d}] at {d} rows", .{ s.n, s.k, m }); + const got = try newOut(device, s, false); + defer got.buf.deinit(); + const wb = queue.commandBuffer(); + const we = wb.compute(.serial); + var start: usize = 0; + while (start + m <= wide) : (start += m) try simdWindow(we, p, s, got, true, start, m); + try commit(wb, we); + compare(t, twin, got, wide / m * m, ref, s, 0); + const rows = wide / m * m; + const eq = std.mem.eql(u8, all.buf.slice(u8, rows * s.n * 2), got.buf.slice(u8, rows * s.n * 2)); + std.debug.print("RESULT {s}: {s} the tiles' bits; twinExact says {s}{s}\n", .{ twin, if (eq) "equals" else "differs from", if (exact) "equal" else "differ", if (exact and !eq) " FAIL" else "" }); + if (exact and !eq) t.failed += 1; + } +} + +pub fn main(_: std.process.Init) !void { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + std.debug.print("device {s}, tensor units {}\n", .{ device.name(), device.tensorUnits() }); + var t = Tally{}; + var seed: u64 = 0x716; + for ([_][2]usize{ .{ 192, 1024 }, .{ 256, 4096 } }) |shape| for ([_]u8{ 4, 8 }) |bits| for ([_]usize{ 32, 64 }) |group| { + seed += 1; + const s = try Synth.init(device, shape[0], shape[1], bits, group, seed); + defer s.deinit(); + const ref = try s.reference(if (bits == 4) 128 else 0); + defer ref.deinit(); + for ([_]bool{ false, true }) |f32_out| try affine(&t, device, queue, &s, ref, f32_out); + }; + for ([_][2]usize{ .{ 192, 1024 }, .{ 256, 4096 } }) |shape| for ([_]usize{ 32, 64 }) |group| { // 2 bits: oQ2's experts + const s = try Synth.init(device, shape[0], shape[1], 2, group, 0x2716 + shape[0] + group); + defer s.deinit(); + const ref = try s.reference(0); + defer ref.deinit(); + for ([_]bool{ false, true }) |f32_out| try affine(&t, device, queue, &s, ref, f32_out); + }; + var p = try sq.Pipelines.load(device); + defer p.deinit(); + for ([_][2]usize{ .{ 40, 512 }, .{ 256, 4096 }, .{ 6400, 256 } }) |shape| { // chunks of 32, 16 and 8 + seed += 1; + const s = try Synth.init(device, shape[0], shape[1], 4, 64, seed); + defer s.deinit(); + const ref = try s.reference(0); + defer ref.deinit(); + try simd(&t, device, queue, &p, &s, ref); + } + if (t.failed > 0) { + std.debug.print("FAIL {d} checks; the largest difference is {d:.4} of its float64 bound\n", .{ t.failed, t.worst }); + return error.CheckFailed; + } + std.debug.print("PASS every product within {d:.4} of its float64 bound, every row's bits equal at every window\n", .{t.worst}); +} 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/qwen27/dflash_head.zig b/zig/tests/qwen27/dflash_head.zig index 451473b63a..eaa4fb0e04 100644 --- a/zig/tests/qwen27/dflash_head.zig +++ b/zig/tests/qwen27/dflash_head.zig @@ -6,7 +6,7 @@ const Head = @import("tensorfold").qwen27.dflash.runtime_backend.Head; pub fn run(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue) !void { const n = 248320; const k = 128; - const sizes = [_]usize{ n * k / 2, n * 2 * 4, n * 2 * 2, n * 2 * 2, 7 * k * 2, 7 * n * 2, 7 * 98592 * 2, 64, 64 }; + const sizes = [_]usize{ n * k, n * 2 * 4, n * 2 * 2, n * 2 * 2, 7 * k * 2, 7 * n * 2, 7 * 98592 * 2, 64, 64 }; var buffers: [sizes.len]mtl.Buffer = undefined; var made: usize = 0; defer for (buffers[0..made]) |b| b.deinit(); @@ -28,9 +28,9 @@ pub fn run(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue) !void { var kernels = q.quant_gpu.Kernels{ .allocator = a, .device = device }; defer kernels.deinit(); var cases: usize = 0; - for ([_]bool{ false, true }) |raw| { + for ([_]u8{ 2, 4, 8 }) |bits| for ([_]bool{ false, true }) |raw| { if (!raw and !device.tensorUnits()) continue; // the lane layout needs tensor units; pre-M5 heads take raw rows - const linear = q.projection.Linear{ .words = .{ .buffer = buffers[0], .offset = 32 }, .pairs = .{ .buffer = buffers[1], .offset = 32 }, .n = n, .k = k, .tile = 32, .slices = 1, .raw = if (raw) .{ .w = buffers[0], .scales = buffers[2], .biases = buffers[3], .w_off = 32, .s_off = 32, .b_off = 32, .n = n, .k = k, .group = 64, .bits = 4, .sum = .f32 } else null }; + const linear = q.projection.Linear{ .words = .{ .buffer = buffers[0], .offset = 32 }, .pairs = .{ .buffer = buffers[1], .offset = 32 }, .n = n, .k = k, .tile = 32, .slices = 1, .bits = bits, .raw = if (raw) .{ .w = buffers[0], .scales = buffers[2], .biases = buffers[3], .w_off = 32, .s_off = 32, .b_off = 32, .n = n, .k = k, .group = 64, .bits = bits, .sum = .f32 } else null }; var selected = try Head.init(a, device, linear); defer selected.deinit(); const ids = selected.ids.?.slice(u32, 98592); @@ -53,6 +53,6 @@ pub fn run(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue) !void { }; cases += 1; } - } + }; std.debug.print("native draft vocabulary: {d} lane/raw width cases,98592 selected columns preserve full-head BF16 bytes and token IDs\n", .{cases}); } diff --git a/zig/tests/qwen27/embed.zig b/zig/tests/qwen27/embed.zig new file mode 100644 index 0000000000..a031fe8ccd --- /dev/null +++ b/zig/tests/qwen27/embed.zig @@ -0,0 +1,154 @@ +//! The 27B's embedding gather at oQ2's 8-bit and a 3-bit width: every column the float64 decode rounded once to bf16. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const mlx = tf.quant.mlx; +const q27 = tf.qwen27; + +const opts = mtl.ResourceOptions.shared; +const a = std.heap.page_allocator; +const group = 64; +const hidden: usize = 5120; +const vocab = 248320; +const rows = 64; + +fn toBf(v: f32) u16 { + const u: u32 = @bitCast(v); + return @truncate((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); +} + +/// One rounding of a float64 to bf16, ties to even (normal range; the gathered values are). +fn bf64(v: f64) u16 { + const u: u64 = @bitCast(v); + const r: f64 = @bitCast((u +% 0xfff_ffff_ffff +% ((u >> 45) & 1)) & ~@as(u64, 0x1fff_ffff_ffff)); + return @truncate(@as(u32, @bitCast(@as(f32, @floatCast(r)))) >> 16); +} + +/// Token ids: both ends of the vocabulary, a repeat, and random rows. +fn ids(seed: u64) [rows]u32 { + var prng = std.Random.DefaultPrng.init(seed); + var out: [rows]u32 = undefined; + for (&out) |*v| v.* = prng.random().uintLessThan(u32, vocab); + out[0] = 0; + out[1] = vocab - 1; + out[2] = 1; + out[3] = vocab - 2; + out[4] = out[7]; + return out; +} + +/// A [vocab, hidden] table at `bits`, the gathered rows filled with random codes, bf16 scales and biases. +const Table = struct { + w: mtl.Buffer, + s: mtl.Buffer, + b: mtl.Buffer, + bits: u8, + + fn init(device: mtl.Device, bits: u8, tokens: []const u32, seed: u64) !Table { + const wb = hidden * bits / 8; + const g = hidden / group; + const t = Table{ .w = try device.buffer(vocab * wb, opts), .s = try device.buffer(vocab * g * 2, opts), .b = try device.buffer(vocab * g * 2, opts), .bits = bits }; + var prng = std.Random.DefaultPrng.init(seed); + const r = prng.random(); + for (tokens) |id| { + r.bytes(t.w.slice(u8, vocab * wb)[id * wb ..][0..wb]); + for (t.s.slice(u16, vocab * g)[id * g ..][0..g], t.b.slice(u16, vocab * g)[id * g ..][0..g]) |*s, *b| { + s.* = toBf((r.float(f32) + 0.5) / 64 / @as(f32, @floatFromInt(@as(u32, 1) << @intCast(bits - 2)))); + b.* = toBf(-(r.float(f32) + 0.5) * 0.15); + } + } + return t; + } + + fn deinit(t: Table) void { + for ([_]mtl.Buffer{ t.w, t.s, t.b }) |x| x.deinit(); + } + + fn tensor(buffer: mtl.Buffer, dtype: tf.checkpoint.DType, c: usize, size: usize) tf.checkpoint.Tensor { + return .{ .buffer = buffer, .offset = 0, .bytes = vocab * c * size, .dtype = dtype, .shape = .{ vocab, c, 1, 1, 1 }, .rank = 2 }; + } + + /// What gpu_weights hands the embedding. + fn triple(t: Table) q27.projection.Triple { + const g = hidden / group; + return .{ .weight = tensor(t.w, .u32, hidden * t.bits / 32, 4), .scales = tensor(t.s, .bf16, g, 2), .biases = tensor(t.b, .bf16, g, 2), .bits = t.bits, .group = group }; + } + + /// Row `id` decoded in float64 by core.quant.mlx, its scales times `scale` (1 for the reference). + fn reference(t: Table, id: usize, scale: f64, out: []f64) !void { + const g = hidden / group; + const wb = hidden * t.bits / 8; + var s: [g * 4]u8 = undefined; // the control's scales, fp32 so 1% off is exactly that + for (0..g) |i| { + const v: f64 = @as(f32, @bitCast(@as(u32, t.s.slice(u16, vocab * g)[id * g + i]) << 16)); + std.mem.writeInt(u32, s[i * 4 ..][0..4], @bitCast(@as(f32, @floatCast(v * scale))), .little); + } + const shape: [5]usize = .{ 1, g, 1, 1, 1 }; + try mlx.reference(.{ + .words = .{ .dtype = .u32, .rank = 2, .shape = .{ 1, wb / 4, 1, 1, 1 }, .bytes = t.w.slice(u8, vocab * wb)[id * wb ..][0..wb] }, + .scales = if (scale == 1) .{ .dtype = .bf16, .rank = 2, .shape = shape, .bytes = t.s.slice(u8, vocab * g * 2)[id * g * 2 ..][0 .. g * 2] } else .{ .dtype = .f32, .rank = 2, .shape = shape, .bytes = &s }, + .biases = .{ .dtype = .bf16, .rank = 2, .shape = shape, .bytes = t.b.slice(u8, vocab * g * 2)[id * g * 2 ..][0 .. g * 2] }, + .bits = t.bits, + .group = group, + }, out); + } +}; + +fn commit(cb: mtl.CommandBuffer, e: mtl.ComputeEncoder) !void { + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |why| { + std.debug.print("GPU failed: {s}\n", .{why}); + return error.GpuFailed; + } +} + +/// The family's gather of `rows` tokens in calls of `m` rows, into a NaN-filled bf16 output. +fn gather(device: mtl.Device, queue: mtl.Queue, k: q27.glue_gpu.Kernels, p: q27.projection.Triple, id: mtl.Buffer, m: usize) !mtl.Buffer { + const c = q27.config.Config{ .hidden = hidden, .intermediate = 17408, .layers = 64, .vocab = vocab, .heads = 24, .kv_heads = 4, .head_dim = 256, .k_heads = 16, .v_heads = 48, .dk = 128, .dv = 128, .conv_kernel = 4, .eps = 1e-6, .rope_dims = 64, .rope_theta = 1e7 }; + const y = try device.buffer(rows * hidden * 2, opts); + @memset(y.slice(u8, rows * hidden * 2), 0xff); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + var start: usize = 0; + while (start < rows) : (start += m) try k.embedding(e, c, p, .{ .buffer = id, .offset = start * 4 }, .{ .buffer = y, .offset = start * hidden * 2 }, @intCast(m)); + try commit(cb, e); + return y; +} + +/// Each width: gathered bits against float64 rounded once, one-row calls against one call, scales 1% off as control. +pub fn check(device: mtl.Device, queue: mtl.Queue) !usize { + const k = try q27.glue_gpu.Kernels.init(device); + defer k.deinit(); + var failed: usize = 0; + const deq = try a.alloc(f64, hidden); + defer a.free(deq); + for ([_]u8{ 8, 3 }, 0..) |bits, at| { + const tokens = ids(0xe27 + at); + const t = try Table.init(device, bits, &tokens, 0x27e + at); + defer t.deinit(); + const id = try device.buffer(rows * 4, opts); + defer id.deinit(); + @memcpy(id.slice(u32, rows), &tokens); + const all = try gather(device, queue, k, t.triple(), id, rows); + defer all.deinit(); + const one = try gather(device, queue, k, t.triple(), id, 1); + defer one.deinit(); + const got = all.slice(u16, rows * hidden); + var differ: usize = 0; + var ctl: usize = 0; + for (tokens, 0..) |token, r| { + try t.reference(token, 1, deq); + for (deq, got[r * hidden ..][0..hidden]) |d, g| differ += @intFromBool(g != bf64(d)); + try t.reference(token, 1.01, deq); + for (deq, got[r * hidden ..][0..hidden]) |d, g| ctl += @intFromBool(g != bf64(d)); + } + const same = std.mem.eql(u16, got, one.slice(u16, rows * hidden)); + const ok = differ == 0 and same and ctl * 10 > rows * hidden; + if (!ok) failed += 1; + std.debug.print("RESULT embedding [{d}, {d}] q{d} g64, {d} tokens: {d} of {d} values differ from float64 rounded once to bf16; " ++ + "one-row calls {s}; scales 1% off: {d} differ{s}\n", .{ vocab, hidden, bits, rows, differ, rows * hidden, if (same) "equal bit for bit" else "DIFFER", ctl, if (ok) "" else " FAIL" }); + } + return failed; +} diff --git a/zig/tests/qwen27/run.zig b/zig/tests/qwen27/run.zig index 250f6ee2ab..bb1dfbdf80 100644 --- a/zig/tests/qwen27/run.zig +++ b/zig/tests/qwen27/run.zig @@ -10,12 +10,13 @@ const core_lane = @import("core_lane"); const help = \\tf-qwen27-run --run --model DIR --tokens IDS.json [--max-tokens 256] \\ [--chunk 128] [--compare-chunk 16] [--reference GENERATED.json] [--force-length] - \\ [--no-warmup] [--profile-steps N] [--teacher] + \\ [--no-warmup] [--profile-steps N] [--teacher] [--seams N] \\IDs and reference files are JSON arrays of unsigned token IDs, already chat-rendered. \\Greedy only, drafts off; JSON receipt on stdout, errors on stderr. \\--teacher feeds every reference token and records native top2 at all positions; margin timing is diagnostic. \\--max-tokens 0 reports prompt tok/s including the final single-row head. \\--compare-chunk reruns the whole prompt and decode and checks prompt state/KV/logits bytes. + \\--seams N (with --compare-chunk): state/KV/logits bytes every N prompt rows under both chunks and after decode. \\CPU only: --help or --emit-sources EXISTING_DIR, no Metal device or model load. ; @@ -85,6 +86,7 @@ pub fn main(init: std.process.Init) !void { var count: usize = 256; var chunk: usize = 128; var compare: usize = 0; + var seam_every: usize = 0; var run = false; var force = false; var warmup = true; @@ -111,9 +113,9 @@ pub fn main(init: std.process.Init) !void { if (i + 1 >= args.len) return error.MissingValue; i += 1; const value = args[i]; - if (std.mem.eql(u8, arg, "--model")) model = value else if (std.mem.eql(u8, arg, "--tokens")) file = value else if (std.mem.eql(u8, arg, "--reference")) reference = value else if (std.mem.eql(u8, arg, "--max-tokens")) count = try std.fmt.parseInt(usize, value, 10) else if (std.mem.eql(u8, arg, "--chunk")) chunk = try std.fmt.parseInt(usize, value, 10) else if (std.mem.eql(u8, arg, "--profile-steps")) profile_steps = try std.fmt.parseInt(usize, value, 10) else if (std.mem.eql(u8, arg, "--compare-chunk")) compare = try std.fmt.parseInt(usize, value, 10) else return error.UnknownOption; + if (std.mem.eql(u8, arg, "--model")) model = value else if (std.mem.eql(u8, arg, "--tokens")) file = value else if (std.mem.eql(u8, arg, "--reference")) reference = value else if (std.mem.eql(u8, arg, "--max-tokens")) count = try std.fmt.parseInt(usize, value, 10) else if (std.mem.eql(u8, arg, "--chunk")) chunk = try std.fmt.parseInt(usize, value, 10) else if (std.mem.eql(u8, arg, "--profile-steps")) profile_steps = try std.fmt.parseInt(usize, value, 10) else if (std.mem.eql(u8, arg, "--compare-chunk")) compare = try std.fmt.parseInt(usize, value, 10) else if (std.mem.eql(u8, arg, "--seams")) seam_every = try std.fmt.parseInt(usize, value, 10) else return error.UnknownOption; } - if (!run or model == null or file == null or chunk == 0 or chunk > 128 or compare > 128 or count > 65536) return error.BadOptions; + if (!run or model == null or file == null or chunk == 0 or chunk > 128 or compare > 128 or count > 65536 or (seam_every > 0 and compare == 0)) return error.BadOptions; const ids = try tokens(a, io, file.?); if (ids.len == 0 or ids.len + count > 262144) return error.BadPrompt; const expected = if (reference) |path| try tokens(a, io, path) else null; @@ -144,6 +146,7 @@ pub fn main(init: std.process.Init) !void { const generation = try generated(a, s, count, force, expected, if (profile_steps > 0) &trace else null, profile_steps, teacher); const out = generation.tokens; const decode_ns = decode_start.durationTo(std.Io.Clock.awake.now(io)).toNanoseconds(); + const decoded = if (compare > 0) try s.fingerprint() else digest; const reference_equal = if (expected) |wanted| std.mem.eql(u32, out, wanted) else false; if (compare > 0) { try runner.reset(0); @@ -151,6 +154,19 @@ pub fn main(init: std.process.Init) !void { if (!std.mem.eql(u8, &digest, &try s.fingerprint())) return error.PromptStateOrLogitsDiffer; const second = try generated(a, s, count, force, if (teacher) expected else null, null, 0, teacher); if (!std.mem.eql(u32, out, second.tokens)) return error.ChunkedTokensDiffer; + if (!std.mem.eql(u8, &decoded, &try s.fingerprint())) return error.DecodeStateOrLogitsDiffer; + } + const Seam = struct { at: usize, equal: bool }; + var seams: std.ArrayList(Seam) = .empty; + var at = seam_every; + while (seam_every > 0 and at < ids.len) : (at += seam_every) { + var sides: [2][32]u8 = undefined; + for ([_]usize{ chunk, compare }, &sides) |c, *side| { + try runner.reset(0); + try s.prefill(ids[0..at], c); + side.* = try s.fingerprint(); + } + try seams.append(a, .{ .at = at, .equal = std.mem.eql(u8, &sides[0], &sides[1]) }); } const digest_hex = std.fmt.bytesToHex(digest, .lower); const ProfileRecord = struct { class: []const u8, calls: u64, gpu_seconds: f64 }; @@ -168,6 +184,11 @@ pub fn main(init: std.process.Init) !void { .chunk = chunk, .compare_chunk = compare, .prompt_state_byte_equal = compare > 0, + .decode_state_byte_equal = compare > 0, + .seams = seams.items, + .seams_byte_equal = for (seams.items) |x| { + if (!x.equal) break false; + } else true, .reference_token_equal = reference_equal, .first_difference = generation.first_difference, .logit_dtype = "bf16", @@ -188,5 +209,6 @@ pub fn main(init: std.process.Init) !void { }, .{}); try std.Io.File.stdout().writeStreamingAll(io, receipt); try std.Io.File.stdout().writeStreamingAll(io, "\n"); + for (seams.items) |x| if (!x.equal) return error.SeamStateOrLogitsDiffer; if (!teacher and reference != null and !reference_equal) return error.ReferenceTokensDiffer; } diff --git a/zig/tests/qwen27/synthetic_graph.zig b/zig/tests/qwen27/synthetic_graph.zig index a88aa39654..3da574b7d1 100644 --- a/zig/tests/qwen27/synthetic_graph.zig +++ b/zig/tests/qwen27/synthetic_graph.zig @@ -93,7 +93,9 @@ pub fn createLayers(a: std.mem.Allocator, layers: usize) !*q.model.Model { inline for (.{ "self_attn.q_norm.weight", "self_attn.k_norm.weight" }) |leaf| try builder.tensor(try builder.layerName(i, leaf), &.{256}, .bf16, 1); } } - var weights = try q.gpu_weights.load(a, device, &cp, c); + var widths = try q.widths.Widths.uniformAll(a, c.layers); + defer widths.deinit(); + var weights = try q.gpu_weights.load(a, device, &cp, c, &widths); errdefer weights.deinit(); const glue = try q.glue_gpu.Kernels.init(device); errdefer glue.deinit(); diff --git a/zig/tests/qwen27/widths.zig b/zig/tests/qwen27/widths.zig new file mode 100644 index 0000000000..51d839fb3a --- /dev/null +++ b/zig/tests/qwen27/widths.zig @@ -0,0 +1,496 @@ +//! The 27B's projections, mixed stacks and embedding without tensor units at every oQ width against float64. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const mlx = tf.quant.mlx; +const q27 = tf.qwen27; +const projection = q27.projection; +const row = tf.row_projection; +const cp = tf.checkpoint; + +const opts = mtl.ResourceOptions.shared; +const a = std.heap.page_allocator; +const unit: f64 = 0x1p-24; // fp32's unit roundoff +const group = 64; // every module of both oQ checkpoints +const wide = 128; // a prompt chunk's rows +const windows = [_]usize{ 1, 2, 3, 8, 16, 17, 128 }; +const bit_widths = [_]u8{ 2, 3, 4, 5, 6, 8 }; // every class at every width oQ4 and oQ2 use anywhere + +/// A projection class: output rows n of input width k (gate and up share a width, so one matrix as loaded). +const Class = struct { name: []const u8, n: usize, k: usize }; +const classes = [_]Class{ + .{ .name = "gdn in_proj_qkv", .n = 10240, .k = 5120 }, + .{ .name = "gdn in_proj_z", .n = 6144, .k = 5120 }, + .{ .name = "gdn in_proj_b", .n = 48, .k = 5120 }, + .{ .name = "gdn in_proj_a", .n = 48, .k = 5120 }, + .{ .name = "gdn out_proj", .n = 5120, .k = 6144 }, + .{ .name = "attn q_proj", .n = 12288, .k = 5120 }, + .{ .name = "attn k_proj", .n = 1024, .k = 5120 }, + .{ .name = "attn v_proj", .n = 1024, .k = 5120 }, + .{ .name = "attn o_proj", .n = 5120, .k = 6144 }, + .{ .name = "mlp gate_up", .n = 34816, .k = 5120 }, + .{ .name = "mlp down_proj", .n = 5120, .k = 17408 }, + .{ .name = "lm_head", .n = 248320, .k = 5120 }, +}; + +/// A fused member: name, rows, bits. +const Member = struct { name: []const u8, n: usize, bits: u8 }; +/// The mixed stacks of both checkpoints (zba and kv), k = 5120. +const stacks = [_][]const Member{ + &.{ .{ .name = "z4", .n = 6144, .bits = 4 }, .{ .name = "b5", .n = 48, .bits = 5 }, .{ .name = "a5", .n = 48, .bits = 5 } }, + &.{ .{ .name = "z4", .n = 6144, .bits = 4 }, .{ .name = "b4", .n = 48, .bits = 4 }, .{ .name = "a5", .n = 48, .bits = 5 } }, + &.{ .{ .name = "z2", .n = 6144, .bits = 2 }, .{ .name = "b2", .n = 48, .bits = 2 }, .{ .name = "a3", .n = 48, .bits = 3 } }, + &.{ .{ .name = "k5", .n = 1024, .bits = 5 }, .{ .name = "v6", .n = 1024, .bits = 6 } }, +}; + +fn toBf(v: f32) u16 { + const u: u32 = @bitCast(v); + return @truncate((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); +} + +fn fromBf(v: u16) f64 { + return @as(f32, @bitCast(@as(u32, v) << 16)); +} + +fn shape2(r: usize, c: usize) [5]usize { + return .{ r, c, 1, 1, 1 }; +} + +/// Synthetic MLX affine weights [n, k] as the checkpoint stores them: words of random codes, bf16 scales and biases. +const Synth = struct { + n: usize, + k: usize, + bits: u8, + w: mtl.Buffer, + s: mtl.Buffer, + b: mtl.Buffer, + + fn init(device: mtl.Device, n: usize, k: usize, bits: u8, seed: u64) !Synth { + var prng = std.Random.DefaultPrng.init(seed); + const r = prng.random(); + const sn = n * (k / group); + const m = Synth{ .n = n, .k = k, .bits = bits, .w = try device.buffer(n * k * bits / 8, opts), .s = try device.buffer(sn * 2, opts), .b = try device.buffer(sn * 2, opts) }; + r.bytes(m.w.slice(u8, n * k * bits / 8)); + for (m.s.slice(u16, sn), m.b.slice(u16, sn)) |*s, *b| { + s.* = toBf((r.float(f32) + 0.5) / 64 / @as(f32, @floatFromInt(@as(u32, 1) << @intCast(bits - 2)))); + b.* = toBf(-(r.float(f32) + 0.5) * 0.15); + } + return m; + } + + fn deinit(m: Synth) void { + for ([_]mtl.Buffer{ m.w, m.s, m.b }) |x| x.deinit(); + } + + fn tensor(buffer: mtl.Buffer, dtype: cp.DType, r: usize, c: usize, size: usize) cp.Tensor { + return .{ .buffer = buffer, .offset = 0, .bytes = r * c * size, .dtype = dtype, .shape = shape2(r, c), .rank = 2 }; + } + + /// The loader's view: what gpu_weights hands projection.prepare. + fn triple(m: Synth) projection.Triple { + const g = m.k / group; + return .{ .weight = tensor(m.w, .u32, m.n, m.k * m.bits / 32, 4), .scales = tensor(m.s, .bf16, m.n, g, 2), .biases = tensor(m.b, .bf16, m.n, g, 2), .bits = m.bits, .group = group }; + } + + /// Row j alone as core.quant.mlx reads it. + fn host(m: Synth, j: usize) mlx.Host { + const g = m.k / group; + const wb = m.k * m.bits / 8; + return .{ + .words = .{ .dtype = .u32, .rank = 2, .shape = shape2(1, wb / 4), .bytes = m.w.slice(u8, m.n * wb)[j * wb ..][0..wb] }, + .scales = .{ .dtype = .bf16, .rank = 2, .shape = shape2(1, g), .bytes = m.s.slice(u8, m.n * g * 2)[j * g * 2 ..][0 .. g * 2] }, + .biases = .{ .dtype = .bf16, .rank = 2, .shape = shape2(1, g), .bytes = m.b.slice(u8, m.n * g * 2)[j * g * 2 ..][0 .. g * 2] }, + .bits = m.bits, + .group = group, + }; + } +}; + +/// `wide` rows of bf16 inputs in [-1, 1). +fn inputs(device: mtl.Device, k: usize, seed: u64) !mtl.Buffer { + var prng = std.Random.DefaultPrng.init(seed); + const x = try device.buffer(wide * k * 2, opts); + for (x.slice(u16, wide * k)) |*v| v.* = toBf(prng.random().float(f32) * 2 - 1); + return x; +} + +/// Every column of a matrix up to 1024 wide, else 512: the first and last 128 and 256 at an odd stride between. +fn pick(n: usize) ![]usize { + if (n <= 1024) { + const cols = try a.alloc(usize, n); + for (cols, 0..) |*c, i| c.* = i; + return cols; + } + const cols = try a.alloc(usize, 512); + const step = ((n - 256) / 256 - 1) | 1; + for (0..128) |i| { + cols[i] = i; + cols[511 - i] = n - 1 - i; + } + for (0..256) |i| cols[128 + i] = 128 + i * step; + return cols; +} + +/// y = x W^T in float64 on `cols`, the sizes the fp32 sums carry, and the scales' share (scales 1% off add 1% of it). +const Ref = struct { + cols: []usize, + want: []f64, + mag: []f64, + scaled: []f64, + + fn init(m: Synth, x: mtl.Buffer) !Ref { + const k = m.k; + const g_n = k / group; + const cols = try pick(m.n); + const c_n = cols.len; + const r = Ref{ .cols = cols, .want = try a.alloc(f64, wide * c_n), .mag = try a.alloc(f64, wide * c_n), .scaled = try a.alloc(f64, wide * c_n) }; + const xd = try a.alloc(f64, wide * k); + defer a.free(xd); + const xg = try a.alloc(f64, wide * g_n); // each group's sum of |x| + defer a.free(xg); + @memset(xg, 0); + for (x.slice(u16, wide * k), xd, 0..) |v, *d, i| { + d.* = fromBf(v); + xg[i / group] += @abs(d.*); + } + const deq = try a.alloc(f64, k); + defer a.free(deq); + const bias = m.b.slice(u16, m.n * g_n); + for (cols, 0..) |j, c| { + try mlx.reference(m.host(j), deq); + for (0..wide) |i| { + var sum: f64 = 0; + var mag: f64 = 0; + var scaled: f64 = 0; + for (0..g_n) |g| { + const b = fromBf(bias[j * g_n + g]); + var qx: f64 = 0; + var qxs: f64 = 0; + for (g * group..(g + 1) * group) |e| { + const xv = xd[i * k + e]; + const sq = deq[e] - b; // s q, exact + sum += xv * deq[e]; + qx += sq * @abs(xv); + qxs += sq * xv; + } + mag += qx + @abs(b) * xg[i * g_n + g]; + scaled += qxs; + } + r.want[i * c_n + c] = sum; + r.mag[i * c_n + c] = mag; + r.scaled[i * c_n + c] = scaled; + } + } + return r; + } + + fn deinit(r: Ref) void { + for ([_][]f64{ r.want, r.mag, r.scaled }) |v| a.free(v); + a.free(r.cols); + } +}; + +/// Failures, values against float64, bit-equal pairs, and the largest difference as a fraction of its bound. +const Tally = struct { failed: usize = 0, values: usize = 0, pairs: usize = 0, worst: f64 = 0, ctl_min: f64 = std.math.inf(f64) }; + +/// `rows` bf16 rows `pitch` apart; a member's columns start at `first`. +const View = struct { + vals: []const u16, + pitch: usize, + first: usize = 0, + fn get(v: View, r: usize, j: usize) f64 { + return fromBf(v.vals[r * v.pitch + v.first + j]); + } +}; + +/// fp32 sums of a group's products, its row sum and a chain of groups (plus the kernel's tree), and one bf16 ulp. +fn bound(want: f64, mag: f64, k: usize, extra: usize) f64 { + const terms: f64 = @floatFromInt(group + k / group + 8 + extra); + return terms * unit * mag + @abs(want) * 0x1p-7 + 1e-30; +} + +/// `rows` rows of `out` against the float64 reference: the worst column, its bound, and the 1%-off control. +fn compare(t: *Tally, name: []const u8, out: View, rows: usize, ref: Ref, k: usize, extra: usize) void { + const c_n = ref.cols.len; + var col: usize = 0; + var col_abs: f64 = -1; + var col_rel: f64 = 0; + var frac: f64 = 0; + var ctl: f64 = 0; + var at: [2]usize = .{ 0, 0 }; + for (ref.cols, 0..) |j, c| { + var abs: f64 = 0; + var top: f64 = 0; + for (0..rows) |r| { + const i = r * c_n + c; + const got = out.get(r, j); + const d = @abs(got - ref.want[i]); + const tol = bound(ref.want[i], ref.mag[i], k, extra); + ctl = @max(ctl, @abs(got - ref.want[i] - 0.01 * ref.scaled[i]) / tol); + if (!(d <= abs)) abs = d; + top = @max(top, @abs(ref.want[i])); + if (!(d / tol <= frac)) { + frac = d / tol; + at = .{ r, j }; + } + } + if (!(abs <= col_abs)) { + col_abs = abs; + col = j; + col_rel = abs / top; + } + } + t.values += rows * c_n; + t.worst = @max(t.worst, frac); + t.ctl_min = @min(t.ctl_min, ctl); + const ok = frac <= 1 and ctl > 1; + if (!ok) t.failed += 1; + std.debug.print("RESULT {s}: {d} columns, worst column {d}: max abs {e:.3}, relative {e:.3}; {d:.4} of its bound (row {d} col {d}), " ++ + "scales 1% off {d:.1}{s}\n", .{ name, c_n, col, col_abs, col_rel, frac, at[0], at[1], ctl, if (ok) "" else " FAIL" }); +} + +/// Rows [0, rows) of `got` equal `want` bit for bit; the first difference printed. +fn same(t: *Tally, name: []const u8, what: []const u8, want: View, got: View, rows: usize, n: usize) bool { + var differ: usize = 0; + var first: ?[2]usize = null; + var most: f64 = 0; + for (0..rows) |r| { + const w = want.vals[r * want.pitch + want.first ..][0..n]; + const g = got.vals[r * got.pitch + got.first ..][0..n]; + if (std.mem.eql(u16, w, g)) continue; + differ += 1; + if (first == null) first = .{ r, std.mem.indexOfDiff(u16, w, g).? }; + for (0..n) |j| most = @max(most, @abs(want.get(r, j) - got.get(r, j))); + } + if (first) |f| { + std.debug.print(" {s}: {s} differ in {d} of {d} rows, by up to {e:.3}; first row {d} col {d}: {d} against {d}\n", .{ name, what, differ, rows, most, f[0], f[1], got.get(f[0], f[1]), want.get(f[0], f[1]) }); + return false; + } + t.pairs += rows * n; + return true; +} + +const Gpu = struct { + device: mtl.Device, + queue: mtl.Queue, + kern: *q27.quant_gpu.Kernels, + pipes: *const row.Pipelines, + + fn commit(cb: mtl.CommandBuffer, e: mtl.ComputeEncoder) !void { + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |why| { + std.debug.print("GPU failed: {s}\n", .{why}); + return error.GpuFailed; + } + } + + /// The family's product over `wide` rows in windows of `m` (a prompt chunk when m == wide), into a NaN-filled y. + fn quant(g: Gpu, lin: projection.Linear, x: mtl.Buffer, y: mtl.Buffer, m: usize) !void { + const n: usize = lin.n; + @memset(y.slice(u8, wide * n * 2), 0xff); + g.kern.prompt = m == wide; + defer g.kern.prompt = false; + const cb = g.queue.commandBuffer(); + const e = cb.compute(.serial); + var start: usize = 0; + while (start + m <= wide) : (start += m) try g.kern.quant(e, lin, .{ .buffer = x, .offset = start * lin.k * 2 }, null, .{ .buffer = x }, .{ .buffer = y, .offset = start * n * 2 }, @intCast(m)); + try commit(cb, e); + } + + /// The row kernel alone at the matrix's width over `wide` rows, unstrided: a member's own product. + fn rows(g: Gpu, raw: row.Weights, x: mtl.Buffer, y: mtl.Buffer) !void { + @memset(y.slice(u8, wide * raw.n * 2), 0xff); + const cb = g.queue.commandBuffer(); + const e = cb.compute(.serial); + const call = try row.call(g.pipes, raw, wide, false); + e.setPipeline(call.pipeline); + e.setBuffer(x, 0, 0); + e.setBuffer(raw.w, raw.w_off, 1); + e.setBuffer(raw.scales, raw.s_off, 2); + e.setBuffer(raw.biases, raw.b_off, 3); + e.setValue(call.dims, 4); + e.setBuffer(y, 0, 5); + e.dispatchGroups(.{ .width = call.groups }, .{ .width = call.threads }); + try commit(cb, e); + } +}; + +/// The kernel the family picks for a matrix before tensor units, and its tree's extra depth for the bound. +fn path(lin: projection.Linear) struct { name: []const u8, extra: usize } { + const raw = lin.raw.?; + if (raw.bits == 4 and raw.group == 64 and raw.n % 8 == 0) return .{ .name = "simd_qmm", .extra = std.math.log2_int(u32, row.simd.splits(raw.n)) + 1 }; + return .{ .name = "row", .extra = 0 }; +} + +fn release(owned: *std.ArrayList(mtl.Buffer)) void { + for (owned.items) |b| b.deinit(); + owned.deinit(a); +} + +/// One class at one width: prompt and decoded rows against float64, every window's rows against the prompt call's. +fn class(t: *Tally, g: Gpu, c: Class, bits: u8, seed: u64) !void { + const m = try Synth.init(g.device, c.n, c.k, bits, seed); + defer m.deinit(); + const x = try inputs(g.device, c.k, seed ^ 0x27); + defer x.deinit(); + var owned: std.ArrayList(mtl.Buffer) = .empty; + defer release(&owned); + const lin = try projection.prepare(g.device, a, &owned, &.{m.triple()}, 32); + const p = path(lin); + const ref = try Ref.init(m, x); + defer ref.deinit(); + var name: [128]u8 = undefined; + const label = try std.fmt.bufPrint(&name, "{s} [{d}, {d}] q{d} g64 on {s}", .{ c.name, c.n, c.k, bits, p.name }); + var outs: [windows.len]mtl.Buffer = undefined; + for (&outs, windows) |*y, w| { + y.* = try g.device.buffer(wide * c.n * 2, opts); + try g.quant(lin, x, y.*, w); + } + defer for (outs) |y| y.deinit(); + const view = struct { + fn of(y: mtl.Buffer, n: usize) View { + return .{ .vals = y.slice(u16, wide * n), .pitch = n }; + } + }.of; + var tag: [160]u8 = undefined; + compare(t, try std.fmt.bufPrint(&tag, "{s}, prompt ({d} rows)", .{ label, wide }), view(outs[windows.len - 1], c.n), wide, ref, c.k, p.extra); + compare(t, try std.fmt.bufPrint(&tag, "{s}, decode (1 row a call)", .{label}), view(outs[0], c.n), wide, ref, c.k, p.extra); + var ok = true; + for (windows[0 .. windows.len - 1], outs[0 .. windows.len - 1]) |w, y| { + var what: [48]u8 = undefined; + ok = same(t, label, try std.fmt.bufPrint(&what, "{d}-row windows", .{w}), view(outs[windows.len - 1], c.n), view(y, c.n), wide / w * w, c.n) and ok; + } + if (!ok) t.failed += 1; + std.debug.print("RESULT {s}: {s}\n", .{ label, if (ok) "every row's bits equal in calls of 1, 2, 3, 8, 16, 17 and 128 rows" else "a row's bits depend on the call FAIL" }); +} + +/// A mixed stack as gpu_weights.Builder.stack builds it: one product a member, each into its columns. +fn mixed(g: Gpu, owned: *std.ArrayList(mtl.Buffer), parts: []const projection.Triple, runs: []projection.Run) !projection.Linear { + var ws: [3]q27.widths.Width = undefined; + for (parts, 0..) |p, i| ws[i] = .{ .bits = p.bits, .group = @intCast(p.group) }; + var buf: [3][2]usize = undefined; + const list = q27.widths.runs(ws[0..parts.len], &buf); + if (list.len != parts.len) return error.NotMixed; + var first: u32 = 0; + for (list, runs[0..list.len]) |r, *run| { + run.* = .{ .linear = try projection.prepare(g.device, a, owned, parts[r[0]..][0..r[1]], 32), .first = first }; + first += run.linear.n; + } + var fused = runs[0].linear; + fused.n = first; + fused.raw = null; + fused.runs = runs[0..list.len]; + return fused; +} + +/// A mixed stack: rows equal in every call; each member's columns within float64's bound and equal to its own product. +fn stack(t: *Tally, g: Gpu, members: []const Member, seed: u64) !void { + const k = 5120; + var name: [64]u8 = undefined; + var len: usize = 0; + for (members, 0..) |m, i| len += (try std.fmt.bufPrint(name[len..], "{s}{s}", .{ if (i == 0) "[" else " ", m.name })).len; + const label = try std.fmt.bufPrint(name[len..], "]", .{}); + const stack_name = name[0 .. len + label.len]; + var synth: [3]Synth = undefined; + var parts: [3]projection.Triple = undefined; + var total: usize = 0; + for (members, 0..) |m, i| { + synth[i] = try Synth.init(g.device, m.n, k, m.bits, seed + i); + parts[i] = synth[i].triple(); + total += m.n; + } + defer for (synth[0..members.len]) |s| s.deinit(); + const x = try inputs(g.device, k, seed ^ 0x27); + defer x.deinit(); + var owned: std.ArrayList(mtl.Buffer) = .empty; + defer release(&owned); + var runs: [3]projection.Run = undefined; + const fused = try mixed(g, &owned, parts[0..members.len], &runs); + var outs: [windows.len]mtl.Buffer = undefined; + for (&outs, windows) |*y, w| { + y.* = try g.device.buffer(wide * total * 2, opts); + try g.quant(fused, x, y.*, w); + } + defer for (outs) |y| y.deinit(); + const prompt = View{ .vals = outs[windows.len - 1].slice(u16, wide * total), .pitch = total }; + var ok = true; + for (windows[0 .. windows.len - 1], outs[0 .. windows.len - 1]) |w, y| { + var what: [48]u8 = undefined; + ok = same(t, stack_name, try std.fmt.bufPrint(&what, "{d}-row windows", .{w}), prompt, .{ .vals = y.slice(u16, wide * total), .pitch = total }, wide / w * w, total) and ok; + } + if (!ok) t.failed += 1; + std.debug.print("RESULT stack {s}: {s}\n", .{ stack_name, if (ok) "every row's bits equal in calls of 1, 2, 3, 8, 16, 17 and 128 rows" else "a row's bits depend on the call FAIL" }); + const alone = try g.device.buffer(wide * synth[0].n * 2, opts); + defer alone.deinit(); + const matrix = try g.device.buffer(wide * synth[0].n * 2, opts); + defer matrix.deinit(); + for (members, runs[0..members.len], 0..) |m, run, i| { + var tag: [128]u8 = undefined; + const member = try std.fmt.bufPrint(&tag, "stack {s} member {s} [{d}, {d}] q{d} g64", .{ stack_name, m.name, m.n, k, m.bits }); + const col = View{ .vals = prompt.vals, .pitch = total, .first = run.first }; + const ref = try Ref.init(synth[i], x); + defer ref.deinit(); + var line: [192]u8 = undefined; + compare(t, try std.fmt.bufPrint(&line, "{s}, prompt ({d} rows)", .{ member, wide }), col, wide, ref, k, 0); + compare(t, try std.fmt.bufPrint(&line, "{s}, decode (1 row a call)", .{member}), .{ .vals = outs[0].slice(u16, wide * total), .pitch = total, .first = run.first }, wide, ref, k, 0); + try g.rows(run.linear.raw.?, x, alone); + const own = View{ .vals = alone.slice(u16, wide * m.n), .pitch = m.n }; + if (same(t, member, "the stack's columns and the member's own row product", own, col, wide, m.n)) { + std.debug.print("RESULT {s}: its columns equal its own row product bit for bit\n", .{member}); + } else { + std.debug.print("RESULT {s}: its columns differ from its own row product FAIL\n", .{member}); + t.failed += 1; + } + // the member as a matrix of its own through the family: the row kernel, or simd_qmm at 4/64 (no output stride) + var solo: std.ArrayList(mtl.Buffer) = .empty; + defer release(&solo); + const lin = try projection.prepare(g.device, a, &solo, parts[i .. i + 1], 32); + try g.quant(lin, x, matrix, wide); + const via = path(lin).name; + if (same(t, member, "the stack's columns and the member as its own matrix", .{ .vals = matrix.slice(u16, wide * m.n), .pitch = m.n }, col, wide, m.n)) { + std.debug.print("RESULT {s}: equals the member as its own matrix ({s}) bit for bit\n", .{ member, via }); + } else if (std.mem.eql(u8, via, "simd_qmm")) { + std.debug.print("NOTE {s}: as its own matrix it takes simd_qmm (no output stride), whose bits differ\n", .{member}); + } else { + std.debug.print("RESULT {s}: differs from the member as its own matrix ({s}) FAIL\n", .{ member, via }); + t.failed += 1; + } + } +} + +pub fn main(_: std.process.Init) !void { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + std.debug.print("device {s}, tensor units {}\n", .{ device.name(), device.tensorUnits() }); + if (device.tensorUnits()) { + std.debug.print("this program checks the row path; tensor units take the lane kernels\n", .{}); + return error.TensorUnits; + } + var kern = q27.quant_gpu.Kernels{ .allocator = a, .device = device }; + defer kern.deinit(); + var pipes = try row.Pipelines.load(device); + defer pipes.deinit(); + const g = Gpu{ .device = device, .queue = queue, .kern = &kern, .pipes = &pipes }; + var t = Tally{}; + var seed: u64 = 0x2727; + for (classes) |c| for (bit_widths) |bits| { + seed += 1; + try class(&t, g, c, bits, seed); + }; + for (stacks) |s| { + seed += 16; + try stack(&t, g, s, seed); + } + t.failed += try @import("embed.zig").check(device, queue); + if (t.failed > 0) { + std.debug.print("FAIL {d} checks; the largest difference is {d:.4} of its float64 bound\n", .{ t.failed, t.worst }); + return error.CheckFailed; + } + std.debug.print("PASS {d} values within {d:.4} of the float64 bound (scales 1% off: at least {d:.1} of it), {d} bit-equal pairs\n", .{ t.values, t.worst, t.ctl_min, t.pairs }); +} 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; +}