diff --git a/build.zig b/build.zig index 1bf8bfa260..273e5b9486 100644 --- a/build.zig +++ b/build.zig @@ -3,6 +3,8 @@ const builtin = @import("builtin"); const cuda_build = @import("zig/build/cuda.zig"); const dist_build = @import("zig/build/dist.zig"); const qwen27_build = @import("zig/build/qwen27.zig"); +const core_checks = @import("zig/build/core_checks.zig"); +const flashnext_checks = @import("zig/build/flashnext_checks.zig"); const hip_build = @import("zig/build/hip.zig"); const grammar_build = @import("zig/build/grammar.zig"); const counters_build = @import("zig/build/metal_counters.zig"); @@ -455,6 +457,7 @@ fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "row_projection", .module = row_mod } }, }) }); b.step("tf-row-check", "The core row projection on synthetic 4-bit matrices: widths bit-equal, a CPU reference, indexed experts").dependOn(&b.addInstallArtifact(row_check, .{}).step); + core_checks.rowWidths(b, target, optimize, metal, row_mod); test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = row_mod })).step); exl3_build.targets(b, target, optimize, metal, test_step); exlr8_build.metalTargets(b, target, optimize, metal, mods.sources, test_step); @@ -474,6 +477,8 @@ fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b // The core lane, shared-context and Qwen3.8-27B programs and host tests. qwen27_build.targets(b, target, optimize, metal, sources, engine, lanes, mods.qwen27, test_step); + core_checks.targets(b, target, optimize, metal, engine); + flashnext_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/flashnext_checks.zig b/zig/build/flashnext_checks.zig new file mode 100644 index 0000000000..2b0e7a8386 --- /dev/null +++ b/zig/build/flashnext_checks.zig @@ -0,0 +1,16 @@ +//! Flash Next's device checks on synthetic weights: programs that run on the GPU, each under its own step. + +const std = @import("std"); + +/// tf-flashnext-widths: every product class at every oQ width 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-flashnext-widths", .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/flashnext_widths.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "tensorfold", .module = engine } }, + }) }); + const about = "Flash Next's products at every oQ width against float64 (worst column), every row's bits in every window"; + b.step("tf-flashnext-widths", about).dependOn(&b.addInstallArtifact(exe, .{}).step); +} diff --git a/zig/flashnext_host.zig b/zig/flashnext_host.zig index 06782189f6..01600db8c1 100644 --- a/zig/flashnext_host.zig +++ b/zig/flashnext_host.zig @@ -2,4 +2,7 @@ test { _ = @import("src/families/flashnext/host_test.zig"); _ = @import("src/families/flashnext/pack.zig"); + _ = @import("src/families/flashnext/lane_w.zig"); + _ = @import("src/families/flashnext/engine_widths.zig"); + _ = @import("src/families/flashnext/prompt_w.zig"); } 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/kernels/metal/prefill/qmmw_nax.metal b/zig/kernels/metal/prefill/qmmw_nax.metal new file mode 100644 index 0000000000..0f58a3899e --- /dev/null +++ b/zig/kernels/metal/prefill/qmmw_nax.metal @@ -0,0 +1,340 @@ +// Prompt projections at any MLX width (WB bits, groups of WG, defined first): qmm6_nax's tiles and fp32 sums. +#include +using namespace metal; +typedef bfloat bfloat16_t; +#include "../nax.h" +using namespace tfp; + +namespace tfqw { + +// One thread's 32 codes (WB words, little-endian bit stream) as bf16(s q + b) in out[0..32). +template +inline void dequantw(thread const uint (&w)[WB + 1], T scale, T bias, threadgroup O* out) { + const float s = float(scale), b = float(bias); + TF_UNROLL + for (int i = 0; i < 32; i++) { + const int bit = WB * i, word = bit >> 5, sh = bit & 31; + uint v = w[word] >> sh; + if (sh + WB > 32) v |= w[word + 1] << (32 - sh); + out[i] = O(static_cast(s * float(v & ((1u << WB) - 1u)) + b)); + } +} + +inline void loadw(thread uint (&w)[WB + 1], const device uint* src) { + TF_UNROLL + for (int i = 0; i < WB; i++) w[i] = src[i]; + w[WB] = 0; +} + +// x times a [64 rows, K] block, 64 deep a step: thread t dequantizes row t / 2's codes k0 + k + 32 (t % 2) onward. +template +inline void k_loopw(thread frag (&acc)[TM][2], const device T* x, int K, int ldx, int live, bool inside, + const device uint* wq, const device T* srow, const device T* brow, int k0, threadgroup T* tile, + int tn, uint t, short2 home) { + constexpr int PAD = 64 + 16 / sizeof(T); + threadgroup T* mine = tile + (t / 2) * PAD + 32 * (t % 2); + TF_UNROLL + for (short i = 0; i < TM; i++) { + acc[i][0] = frag(0); + acc[i][1] = frag(0); + } + const int hk = k0 + 32 * int(t % 2); + uint w[WB + 1]; + loadw(w, wq); + T sc = srow[hk / WG], bi = brow[hk / WG]; + for (int k = 0; k < K; k += 64) { + threadgroup_barrier(mem_flags::mem_threadgroup); + dequantw(w, sc, bi, mine); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (k + 64 < K) { + wq += 2 * WB; + loadw(w, wq); + sc = srow[(hk + k + 64) / WG]; + bi = brow[(hk + k + 64) / WG]; + } +#pragma clang loop unroll(disable) + for (int kk = 0; kk < 64; kk += 32) { + if (live > 0) { + frag a[TM][2], b[2][2]; + TF_UNROLL + for (short i = 0; i < 2; i++) { + TF_UNROLL + for (short j = 0; j < 2; j++) { + frag_get_t(b[j][i], (const threadgroup T*)tile, PAD, tn + 16 * i, kk + 16 * j, home); + } + } + TF_UNROLL + for (short i = 0; i < TM; i++) { + TF_UNROLL + for (short j = 0; j < 2; j++) { + if (inside) { + frag_get(a[i][j], x, ldx, 16 * i, kk + 16 * j, home); + } else { + frag_get_in(a[i][j], x, ldx, 16 * i, kk + 16 * j, home, live, kk + 32); + } + } + } + TF_UNROLL + for (short m = 0; m < TM; m++) { + TF_UNROLL + for (short j = 0; j < 2; j++) { + mma_16x32(acc[m][0], acc[m][1], a[m][j], b[j][0], b[j][1]); + } + } + } + } + x += 64; + } + threadgroup_barrier(mem_flags::mem_threadgroup); +} + +// Gate and up at once (k_loopw's sums, twice): both tiles dequantized side by side, each x fragment feeds both. +template +inline void k_loopwx2(thread frag (&ag)[TM][2], thread frag (&au)[TM][2], const device T* x, int K, + int live, bool inside, const device uint* wg, const device T* sg, const device T* bg, + const device uint* wu, const device T* su, const device T* bu, threadgroup T* tg_, + threadgroup T* tu, int tn, uint t, short2 home) { + constexpr int PAD = 64 + 16 / sizeof(T); + threadgroup T* mg = tg_ + (t / 2) * PAD + 32 * (t % 2); + threadgroup T* mu = tu + (t / 2) * PAD + 32 * (t % 2); + TF_UNROLL + for (short i = 0; i < TM; i++) { + ag[i][0] = frag(0); + ag[i][1] = frag(0); + au[i][0] = frag(0); + au[i][1] = frag(0); + } + const int hk = 32 * int(t % 2); + uint gw[WB + 1], uw[WB + 1]; + loadw(gw, wg); + loadw(uw, wu); + T gs = sg[hk / WG], gbi = bg[hk / WG], us = su[hk / WG], ubi = bu[hk / WG]; + for (int k = 0; k < K; k += 64) { + threadgroup_barrier(mem_flags::mem_threadgroup); + dequantw(gw, gs, gbi, mg); + dequantw(uw, us, ubi, mu); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (k + 64 < K) { + wg += 2 * WB; + wu += 2 * WB; + loadw(gw, wg); + loadw(uw, wu); + const int g = (hk + k + 64) / WG; + gs = sg[g]; + gbi = bg[g]; + us = su[g]; + ubi = bu[g]; + } +#pragma clang loop unroll(disable) + for (int kk = 0; kk < 64; kk += 32) { + if (live > 0) { + frag a[TM][2], b[2][2], c[2][2]; + TF_UNROLL + for (short i = 0; i < 2; i++) { + TF_UNROLL + for (short j = 0; j < 2; j++) { + frag_get_t(b[j][i], (const threadgroup T*)tg_, PAD, tn + 16 * i, kk + 16 * j, home); + frag_get_t(c[j][i], (const threadgroup T*)tu, PAD, tn + 16 * i, kk + 16 * j, home); + } + } + TF_UNROLL + for (short i = 0; i < TM; i++) { + TF_UNROLL + for (short j = 0; j < 2; j++) { + if (inside) { + frag_get(a[i][j], x, K, 16 * i, kk + 16 * j, home); + } else { + frag_get_in(a[i][j], x, K, 16 * i, kk + 16 * j, home, live, kk + 32); + } + } + } + TF_UNROLL + for (short m = 0; m < TM; m++) { + TF_UNROLL + for (short j = 0; j < 2; j++) { + mma_16x32(ag[m][0], ag[m][1], a[m][j], b[j][0], b[j][1]); + mma_16x32(au[m][0], au[m][1], a[m][j], c[j][0], c[j][1]); + } + } + } + } + x += 64; + } + threadgroup_barrier(mem_flags::mem_threadgroup); +} + +// The expert activation from the two products, as qmm6_nax's act6 rounds it. +inline float actw(float g, float u) { + const float gb = float(bfloat16_t(g)), ub = float(bfloat16_t(u)); + return float(bfloat16_t(float(bfloat16_t(gb / (1.0f + metal::exp(-gb)))) * ub)); +} + +template +inline void store_act(thread const frag (&ag)[TM][2], thread const frag (&au)[TM][2], device bfloat16_t* y, + int ld, int live, int nc, short2 home) { + TF_UNROLL + for (short i = 0; i < TM; i++) { + TF_UNROLL + for (short j = 0; j < 2; j++) { + frag f; + TF_UNROLL + for (short e = 0; e < 8; e++) f[e] = actw(ag[i][j][e], au[i][j][e]); + if (live == 16 * TM && nc >= 32) { + frag_put(f, y, ld, 16 * i, 16 * j, home); + } else { + frag_put_in(f, y, ld, 16 * i, 16 * j, home, live, nc); + } + } + } +} + +template +inline void store(thread const frag (&acc)[TM][2], device T* y, int ld, int live, int nc, short2 home) { + TF_UNROLL + for (short i = 0; i < TM; i++) { + TF_UNROLL + for (short j = 0; j < 2; j++) { + if (live == 16 * TM && nc >= 32) { + frag_put(acc[i][j], y, ld, 16 * i, 16 * j, home); + } else { + frag_put_in(acc[i][j], y, ld, 16 * i, 16 * j, home, live, nc); + } + } + } +} + +// The tile'th BM-row tile over experts in order: its expert, first row and row count (qmm6_nax's expert_tile). +template +inline bool expert_tile(const device int32_t* offsets, int experts, int total, int tile, uint lane, thread int& expert, + thread int& row, thread int& rows) { + int before = 0; + for (int base = 0; base < experts; base += 32) { + const int e = base + int(lane); + const int first = e < experts ? offsets[e] : total, last = e + 1 < experts ? offsets[e + 1] : total; + const int count = (last - first + BM - 1) / BM; + const int upto = before + simd_prefix_inclusive_sum(count); + const int owner = simd_sum(int(upto <= tile)); + if (owner < 32) { + expert = base + owner; + row = simd_shuffle(first, ushort(owner)) + (tile - simd_shuffle(upto - count, ushort(owner))) * BM; + rows = min(BM, simd_shuffle(last, ushort(owner)) - row); + return true; + } + before = simd_shuffle(upto, ushort(31)); + } + return false; +} + +} // namespace tfqw + +// y = x W^T for W [N, K] at WB/WG: 64 x 64 tiles, 4 simdgroups of 32 x 32. P: K N M, y's row stride (0: N). +[[kernel]] void tf_qmmw_t_nax(const device uint* W [[buffer(0)]], const device bfloat16_t* S [[buffer(1)]], + const device bfloat16_t* B [[buffer(2)]], const device bfloat16_t* X [[buffer(3)]], const device int* P [[buffer(4)]], + device bfloat16_t* Y [[buffer(5)]], uint sg [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint3 tg [[threadgroup_position_in_grid]]) { + threadgroup bfloat16_t tile[64 * (64 + 16 / sizeof(bfloat16_t))]; + const int K = P[0], N = P[1], M = P[2], LD = P[3] > 0 ? P[3] : N; + const int row = int(tg.y) * 64, col = int(tg.x) * 64, t = int(sg) * 32 + int(lane); + const int tm = 32 * int(sg / 2), tn = 32 * int(sg % 2), live = min(32, M - (row + tm)); + const long wrow = long(min(col + t / 2, N - 1)); + const int WPR = K * WB / 32, KG = K / WG; + frag acc[2][2]; + tfqw::k_loopw(acc, X + long(row + tm) * K, K, K, live, live == 32, W + wrow * WPR + WB * (t % 2), + S + wrow * KG, B + wrow * KG, 0, tile, tn, uint(t), frag_home(ushort(lane))); + if (col + tn < N && live > 0) + tfqw::store(acc, Y + long(row + tm) * LD + col + tn, LD, live, N - (col + tn), frag_home(ushort(lane))); +} + +// One K part of x W^T: fp32 partials PART[z][M][N] for part z of P[3] parts. P: K N M parts. +[[kernel]] void tf_qmmw_splitk_nax(const device uint* W [[buffer(0)]], const device bfloat16_t* S [[buffer(1)]], + const device bfloat16_t* B [[buffer(2)]], const device bfloat16_t* X [[buffer(3)]], const device int* P [[buffer(4)]], + device float* PART [[buffer(5)]], uint sg [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint3 tg [[threadgroup_position_in_grid]]) { + threadgroup bfloat16_t tile[64 * (64 + 16 / sizeof(bfloat16_t))]; + const int K = P[0], N = P[1], M = P[2], parts = P[3], KP = K / parts, z = int(tg.z); + const int row = int(tg.y) * 64, col = int(tg.x) * 64, t = int(sg) * 32 + int(lane); + const int tm = 32 * int(sg / 2), tn = 32 * int(sg % 2), live = min(32, M - (row + tm)); + const long wrow = long(min(col + t / 2, N - 1)); + const int WPR = K * WB / 32, KG = K / WG, k0 = z * KP; + frag acc[2][2]; + tfqw::k_loopw(acc, X + long(row + tm) * K + k0, KP, K, live, live == 32, + W + wrow * WPR + (k0 / 32) * WB + WB * (t % 2), S + wrow * KG, B + wrow * KG, k0, tile, tn, + uint(t), frag_home(ushort(lane))); + if (col + tn < N && live > 0) + tfqw::store(acc, PART + (long(z) * M + row + tm) * N + col + tn, N, live, N - (col + tn), frag_home(ushort(lane))); +} + +// y[r] = x[r] W_e^T for rows sorted by expert, W [E, N, K]: BM-row tiles within an expert. P: M N K experts. +template +inline void gatherw(const device bfloat16_t* X, const device uint* W, const device bfloat16_t* S, + const device bfloat16_t* B, const device int32_t* O, device bfloat16_t* Y, const device int* P, + threadgroup bfloat16_t* tile, uint3 tg, uint sg, uint lane) { + constexpr int SM = BM / 2; + const int M = P[0], N = P[1], K = P[2]; + int expert, row, rows; + if (!tfqw::expert_tile(O, P[3], M, int(tg.y), lane, expert, row, rows)) { + return; + } + const int col = int(tg.x) * 64, t = int(sg) * 32 + int(lane); + const int tm = SM * int(sg / 2), tn = 32 * int(sg % 2), live = clamp(rows - tm, 0, SM); + const long wrow = long(expert) * N + min(col + t / 2, N - 1); + const int WPR = K * WB / 32, KG = K / WG; + frag acc[SM / 16][2]; + tfqw::k_loopw(acc, X + long(row + tm) * K, K, K, live, row + tm + SM <= M, + W + wrow * WPR + WB * (t % 2), S + wrow * KG, B + wrow * KG, 0, tile, tn, uint(t), + frag_home(ushort(lane))); + if (col + tn < N && live > 0) + tfqw::store(acc, Y + long(row + tm) * N + col + tn, N, live, N - (col + tn), frag_home(ushort(lane))); +} + +// Experts' gate and up products and activation in one pass, for rows sorted by expert. P: M N K experts. +template +inline void gather_guw(const device bfloat16_t* X, const device uint* WG_, const device bfloat16_t* SG, + const device bfloat16_t* BG, const device uint* WU, const device bfloat16_t* SU, + const device bfloat16_t* BU, const device int32_t* O, device bfloat16_t* Y, const device int* P, + threadgroup bfloat16_t* tile_g, threadgroup bfloat16_t* tile_u, uint3 tg, uint sg, uint lane) { + constexpr int SM = BM / 2; + const int M = P[0], N = P[1], K = P[2]; + int expert, row, rows; + if (!tfqw::expert_tile(O, P[3], M, int(tg.y), lane, expert, row, rows)) { + return; + } + const int col = int(tg.x) * 64, t = int(sg) * 32 + int(lane); + const int tm = SM * int(sg / 2), tn = 32 * int(sg % 2), live = clamp(rows - tm, 0, SM); + const long wrow = long(expert) * N + min(col + t / 2, N - 1); + const int WPR = K * WB / 32, KG = K / WG; + frag ag[SM / 16][2], au[SM / 16][2]; + tfqw::k_loopwx2(ag, au, X + long(row + tm) * K, K, live, row + tm + SM <= M, + WG_ + wrow * WPR + WB * (t % 2), SG + wrow * KG, BG + wrow * KG, + WU + wrow * WPR + WB * (t % 2), SU + wrow * KG, BU + wrow * KG, tile_g, tile_u, tn, + uint(t), frag_home(ushort(lane))); + if (col + tn < N && live > 0) + tfqw::store_act(ag, au, Y + long(row + tm) * N + col + tn, N, live, N - (col + tn), frag_home(ushort(lane))); +} + +#define TF_GATHERW(BM) \ + [[kernel]] void tf_gather_qmmw_nax_##BM(const device bfloat16_t* X [[buffer(0)]], const device uint* W [[buffer(1)]], \ + const device bfloat16_t* S [[buffer(2)]], const device bfloat16_t* B [[buffer(3)]], \ + const device int32_t* O [[buffer(4)]], const device int32_t* P [[buffer(5)]], device bfloat16_t* Y [[buffer(6)]], \ + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], \ + uint3 tg [[threadgroup_position_in_grid]]) { \ + threadgroup bfloat16_t tile[64 * (64 + 16 / sizeof(bfloat16_t))]; \ + gatherw(X, W, S, B, O, Y, P, tile, tg, sg, lane); \ + } +TF_GATHERW(32) +TF_GATHERW(64) +TF_GATHERW(128) + +#define TF_GATHER_GUW(BM) \ + [[kernel]] void tf_gather_guw_nax_##BM(const device bfloat16_t* X [[buffer(0)]], const device uint* WG_ [[buffer(1)]], \ + const device bfloat16_t* SG [[buffer(2)]], const device bfloat16_t* BG [[buffer(3)]], \ + const device uint* WU [[buffer(4)]], const device bfloat16_t* SU [[buffer(5)]], \ + const device bfloat16_t* BU [[buffer(6)]], const device int32_t* O [[buffer(7)]], \ + const device int32_t* P [[buffer(8)]], device bfloat16_t* Y [[buffer(9)]], \ + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], \ + uint3 tg [[threadgroup_position_in_grid]]) { \ + threadgroup bfloat16_t tile_g[64 * (64 + 16 / sizeof(bfloat16_t))], tile_u[64 * (64 + 16 / sizeof(bfloat16_t))]; \ + gather_guw(X, WG_, SG, BG, WU, SU, BU, O, Y, P, tile_g, tile_u, tg, sg, lane); \ + } +TF_GATHER_GUW(32) +TF_GATHER_GUW(64) diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index 6313b656cd..e916b97a2c 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -49,6 +49,8 @@ pub const nax = @embedFile("nax.h"); /// Flash Next 6-bit (group 32) prompt projections on the tensor units: dense and sorted-expert gather. pub const flashnext_qmm6 = @embedFile("prefill/qmm6_nax.metal") ++ @embedFile("prefill/qmm6_nax_b.metal"); +/// Flash Next's prompt products at any MLX width (WB and WG defined before it; nax.h inlined at load). +pub const flashnext_qmmw = @embedFile("prefill/qmmw_nax.metal"); /// Flash Next prompt-chunk glue: hyper-connection pieces, router rows, top-k, the expert sort, gathers and scatters. pub const flashnext_prompt = @embedFile("prefill/fn_prompt.metal"); /// Flash Next block selection in GPU-side rounds: per-row metadata from the arena and pooling at absolute blocks. 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/flashnext/affine.zig b/zig/src/families/flashnext/affine.zig index 6152735775..40aa21b977 100644 --- a/zig/src/families/flashnext/affine.zig +++ b/zig/src/families/flashnext/affine.zig @@ -18,10 +18,6 @@ pub const Spec = struct { if (count % 32 != 0) return error.InvalidAffineWidth; return count / 32; } - - pub fn checkFlashKernel(s: Spec) !void { - if (s.bits != 6 or s.group != 32) return error.UnsupportedFlashAffineKernel; - } }; pub const Shape = struct { experts: usize, n: usize, k: usize, words: usize, groups: usize, dtype: st.DType }; diff --git a/zig/src/families/flashnext/config.zig b/zig/src/families/flashnext/config.zig index 7241ee0b54..df990164e1 100644 --- a/zig/src/families/flashnext/config.zig +++ b/zig/src/families/flashnext/config.zig @@ -54,6 +54,8 @@ pub const Config = struct { eos: [16]u32 = @splat(0), eos_count: usize = 0, global_affine: affine.Spec, + /// Set by a tensor reader: each lookup then goes through core/quant, which fills in no missing bits or group. + widths: ?*const @import("widths.zig").Widths = null, pub fn deinit(c: *Config) void { c.parsed.deinit(); @@ -94,6 +96,10 @@ pub const Config = struct { } pub fn quantization(c: *const Config, path: []const u8) !?affine.Spec { + if (c.widths) |w| { + const width = (try w.of(path)) orelse return null; + return try affine.Spec.init(width.bits, width.group); + } const wanted = canonical(path); var found = false; var answer: ?affine.Spec = c.global_affine; @@ -113,7 +119,7 @@ pub const Config = struct { } }; -fn canonical(value: []const u8) []const u8 { +pub fn canonical(value: []const u8) []const u8 { var path = if (std.mem.endsWith(u8, value, ".weight")) value[0 .. value.len - 7] else value; for ([_][]const u8{ "model.language_model.", "language_model.", "text_model.", "model." }) |prefix| { if (std.mem.startsWith(u8, path, prefix)) path = path[prefix.len..]; diff --git a/zig/src/families/flashnext/dense.zig b/zig/src/families/flashnext/dense.zig index 1dda0a2e10..0ca09e6068 100644 --- a/zig/src/families/flashnext/dense.zig +++ b/zig/src/families/flashnext/dense.zig @@ -61,7 +61,7 @@ pub fn recorded(r: *Run, role: []const u8) !struct { n: usize, k: usize, sk: usi } /// y = x W over the layout's tiles, for the rows `mdims` holds. -pub fn project(r: *Run, l: Layout, x: Buf, w: Lane, mdims: Buf, y: Buf) !void { +pub fn project(r: *Run, l: Layout, x: Buf, w: Lane.Six, mdims: Buf, y: Buf) !void { const pipe = try pipeline(r, l); r.enc.setPipeline(pipe); for ([_]Buf{ x, w.wq, w.sbt, mdims, y }, 0..) |b, j| r.enc.setBuffer(b.b, b.off, j); @@ -70,7 +70,7 @@ pub fn project(r: *Run, l: Layout, x: Buf, w: Lane, mdims: Buf, y: Buf) !void { } /// A target lane projection on the core kernel at the recorded kernel's shape and K slices. -pub fn lane(r: *Run, role: []const u8, x: Buf, w: Lane, mdims: Buf, y: Buf) !void { +pub fn lane(r: *Run, role: []const u8, x: Buf, w: Lane.Six, mdims: Buf, y: Buf) !void { if (r.skip & Run.class(role) != 0) return; const s = try recorded(r, role); try project(r, .{ .n = s.n, .k = s.k, .sk = s.sk, .pf = r.lane_pf }, x, w, mdims, y); diff --git a/zig/src/families/flashnext/drafts.zig b/zig/src/families/flashnext/drafts.zig new file mode 100644 index 0000000000..34f5d72c63 --- /dev/null +++ b/zig/src/families/flashnext/drafts.zig @@ -0,0 +1,59 @@ +//! The decode's draft widths: the depth rule over recent landing, and copy drafts from the token history. +const std = @import("std"); + +/// Round width from recent landing: 3 (prose), 6 while drafts land (code), the widest while nearly all land. +pub const DepthRule = struct { + rate: f64 = 0.6, // moving average of the head's drafts landed over drafts offered + depth: usize = 3, + pair: bool = false, // speed-up mode: even windows only (an odd window's last row costs a Mac a whole expert row) + + pub fn pick(self: *const DepthRule) usize { + return self.depth; + } + + /// A head round's drafts and how many landed: 3, 6 or 8 on one Mac, 3 or 7 on the pair. + pub fn update(self: *DepthRule, depth: usize, landed: usize) void { + self.rate = 0.7 * self.rate + 0.3 * @as(f64, @floatFromInt(landed)) / @as(f64, @floatFromInt(depth)); + if (self.pair) { + if (self.depth == 3 and self.rate > 0.95) self.depth = 7 else if (self.depth == 7 and self.rate < 0.8) self.depth = 3; + return; + } + if (self.depth == 3 and self.rate > 0.8) { + self.depth = 6; + } else if (self.depth == 6 and self.rate > 0.74) { + self.depth = 8; + } else if (self.depth == 8 and self.rate < 0.62) { + self.depth = 6; + } else if (self.depth == 6 and self.rate < 0.65) self.depth = 3; + } +}; + +test "the pair's depth rule keeps every window even" { + var pair: DepthRule = .{ .pair = true }; + for (0..20) |_| pair.update(pair.pick(), pair.pick()); + try std.testing.expectEqual(@as(usize, 7), pair.pick()); + for (0..20) |_| pair.update(pair.pick(), 1); + try std.testing.expectEqual(@as(usize, 3), pair.pick()); + var one: DepthRule = .{}; + for (0..20) |_| one.update(one.pick(), one.pick()); + try std.testing.expectEqual(@as(usize, 8), one.pick()); +} + +/// Copies into `out` what followed the latest earlier match of `hist`'s longest repeated suffix (`min`..8), or 0. +pub fn copyDrafts(hist: []const u32, min: usize, out: []u32) usize { + const n = hist.len; + if (n < min + 1) return 0; + var len: usize = @min(8, n - 1); + while (len >= min) : (len -= 1) { + const suffix = hist[n - len ..]; + var end: usize = n - 1; + while (end >= len) : (end -= 1) { + if (std.mem.eql(u32, hist[end - len .. end], suffix)) { + const k = @min(out.len, n - end); + @memcpy(out[0..k], hist[end .. end + k]); + return k; + } + } + } + return 0; +} diff --git a/zig/src/families/flashnext/engine.zig b/zig/src/families/flashnext/engine.zig index d588ec1b32..a2969b5774 100644 --- a/zig/src/families/flashnext/engine.zig +++ b/zig/src/families/flashnext/engine.zig @@ -49,9 +49,11 @@ const f32Buf = fz.f32Buf; const jsonInt = fz.jsonInt; /// Rounds the token ring holds (fz_accept writes round & 511). -const RING = 512; +pub const RING = 512; /// Positions a reply keeps free past its last token: two rounds in flight. pub const MARGIN = 2 * MAXR; +/// An MoE layer's expert projections in the engine's order: gate, up, shared gate, shared up, down, shared down. +const expert_projs = [_][]const u8{ "switch_mlp.gate_proj", "switch_mlp.up_proj", "shared_expert.gate_proj", "shared_expert.up_proj", "switch_mlp.down_proj", "shared_expert.down_proj" }; /// Bytes a conversation's next kept state usually adds (about 4k tokens of reply and tool output): the pool readies that much more. pub const NEXT_TURN = 4096 * 13 * 2304; @@ -199,7 +201,9 @@ pub const Engine = struct { } const index_mapped = try mtl.MappedFile.open(try std.fmt.allocPrintSentinel(arena, "{s}/model.safetensors.index.json", .{model_dir}, 0)); try maps.append(arena, index_mapped); - const id = try pack_io.identityFromMapped(arena, index_mapped.bytes[0..index_mapped.size], mapped_shards.items); + const config_mapped = try mtl.MappedFile.open(try std.fmt.allocPrintSentinel(arena, "{s}/config.json", .{model_dir}, 0)); + try maps.append(arena, config_mapped); + const id = try pack_io.identityFromMapped(arena, index_mapped.bytes[0..index_mapped.size], config_mapped.bytes[0..config_mapped.size], mapped_shards.items); for (maps.items) |*m| m.deinit(); maps.clearRetainingCapacity(); break :blk id; @@ -223,6 +227,7 @@ pub const Engine = struct { // the reference values: the dump's ref.json, or derived from the checkpoint's config (config.pleRef) var ple_values: config.PleRef = undefined; var attention_scale: f64 = undefined; + var cfg = try config.Config.read(arena, io, model_dir); // read once: the reference values and the widths if (dump_dir) |d| { const ref_file = try mtl.MappedFile.open(try std.fmt.allocPrintSentinel(arena, "{s}/ref.json", .{d}, 0)); const ref = try std.json.parseFromSliceLeaky(std.json.Value, arena, ref_file.bytes[0..ref_file.size], .{}); @@ -243,7 +248,6 @@ pub const Engine = struct { } attention_scale = ref.object.get("attention_scale").?.float; } else { - var cfg = try config.Config.read(arena, io, model_dir); ple_values = try config.pleRef(&cfg); attention_scale = 1.0 / std.math.sqrt(@as(f64, @floatFromInt(cfg.head_dim))); } @@ -271,8 +275,7 @@ pub const Engine = struct { L.pool = try r.loadf("L{d}.att.pool", .{i}); L.pooled = .{ .b = try r.buffer(CAP / 4 * 128 * 2) }; } - const projs = [_][]const u8{ "switch_mlp.gate_proj", "switch_mlp.up_proj", "shared_expert.gate_proj", "shared_expert.up_proj", "switch_mlp.down_proj", "shared_expert.down_proj" }; - for (projs, 0..) |proj, j| { + for (expert_projs, 0..) |proj, j| { for ([_][]const u8{ "weight", "scales", "biases" }, 0..) |suffix, k| { L.ex[j * 3 + k] = try r.loadf("language_model.model.layers.{d}.mlp.{s}.{s}", .{ i, proj, suffix }); } @@ -309,6 +312,7 @@ pub const Engine = struct { m.ple.tables[3 * g + 1] = try r.group(16 * g, 16, "scales"); m.ple.tables[3 * g + 2] = try r.group(16 * g, 16, "biases"); } + m.ple.scale = try fz.ple_scale.of(r, index.object.get("weight_map").?.object); const B = struct { fn of(rr: *Run, n: usize) !Buf { return .{ .b = try rr.buffer(n) }; @@ -360,11 +364,11 @@ pub const Engine = struct { { const ids_n = try fz.draftCount((try r.entry("mtp.draft_ids")).len); const ids = try r.load("mtp.draft_ids"); + const mtp_root = @import("index.zig").mtpRoot(index.object.get("weight_map").?.object); var ex: [18]Buf = undefined; - const projs = [_][]const u8{ "switch_mlp.gate_proj", "switch_mlp.up_proj", "shared_expert.gate_proj", "shared_expert.up_proj", "switch_mlp.down_proj", "shared_expert.down_proj" }; - for (projs, 0..) |proj, j| { + for (expert_projs, 0..) |proj, j| { for ([_][]const u8{ "weight", "scales", "biases" }, 0..) |suffix, k| { - ex[j * 3 + k] = try r.loadf("language_model.mtp.layers.0.mlp.{s}.{s}", .{ proj, suffix }); + ex[j * 3 + k] = try r.loadf("{s}.layers.0.mlp.{s}.{s}", .{ mtp_root, proj, suffix }); } } if (r.xpack) try r.repack(&ex); @@ -422,6 +426,7 @@ pub const Engine = struct { e.m = m; e.pr = try arena.create(Prompt); e.pr.* = try Prompt.init(r, pack_dir, r.xnew_header); + const uniform = try @import("engine_widths.zig").apply(arena, &cfg, index.object.get("weight_map").?.object, e.pr, m); e.pr2 = try arena.create(Prompt); e.pr2.* = try e.pr.sibling(); try e.rounds(); @@ -429,55 +434,14 @@ pub const Engine = struct { e.msel = try GSelect.init(r, 0); // speed-up mode: TP=2 with the peer the settings name (tp.zig); both ranks then run the same replies in lockstep const tp_path: ?[]const u8 = speed_up orelse if (std.c.getenv("TF_FLASHNEXT_TP")) |p| std.mem.span(p) else null; + if (tp_path != null and !uniform) return error.SpeedUpNeeds6Bit32; // TP splits the 6/32 kernels only if (tp_path) |path| r.tp = try fz.Tp2.init(arena, r.device, path); return e; } - /// The GPU-side rounds' buffers: DeltaNet states (kept, window rows), the PLE tail pair, the arena, the token ring, - /// per-width row counts, and the MTP slots reading their positions from the arena. + /// The GPU-side rounds' buffers (engine_rounds.zig). fn rounds(e: *Engine) !void { - const r = e.r; - const m = e.m; - const n_lin: usize = 36; - e.g_cs = try r.buffer(n_lin * CS_ROW); - e.g_so = try r.buffer(n_lin * SO_ROW); - e.o_cs = try r.buffer(n_lin * MAXR * CS_ROW); - e.o_so = try r.buffer(n_lin * MAXR * SO_ROW); - var gi: usize = 0; - for (&m.layers) |*L| if (L.linear) { - L.cs[0] = .{ .b = e.g_cs, .off = gi * CS_ROW }; - L.cs[1] = .{ .b = e.o_cs, .off = gi * MAXR * CS_ROW }; - L.so[0] = .{ .b = e.g_so, .off = gi * SO_ROW }; - L.so[1] = .{ .b = e.o_so, .off = gi * MAXR * SO_ROW }; - gi += 1; - }; - if (r.gdn_kept != null) m.recs = .{ .{ .b = try r.buffer(n_lin * fz.gdn_step.RECORD) }, .{ .b = try r.buffer(n_lin * fz.gdn_step.RECORD) } }; - m.marks = try fz.Marks.init(r, n_lin); - e.cins = .{ m.ple.cin, .{ .b = try r.buffer((PLE_TAIL + MAXR) * WIDE * 2) } }; - r.ar = .{ .b = try r.buffer(4 * fz.AR_WORDS) }; - e.ring = try r.buffer(fz.RING_WORDS * 4 * RING); - e.hist = try r.buffer((CAP + 64) * 4); - e.wids = .{ .b = try r.buffer(64) }; - for (1..MAXR + 1) |n| { - e.rows_w[n] = try i32Buf(r, &.{@intCast(n)}); - e.mdims_w[n] = try i32Buf(r, &.{ @intCast(n), 16, 0, 0, 0, 0, 0, 0 }); - } - e.host = .{ .rows = m.t.rows, .mdims = m.t.mdims, .pos8 = m.t.pos8, .nk8 = m.t.nk8, .kvmeta = m.t.kvmeta, .slots = m.mtp.slots }; - e.gpu_slots = m.mtp.slots; - for (1..MAXR + 1) |n| { // the head absorbing a window of n rows: slot MAXR - 1 + n - const sl = &e.gpu_slots[MAXR - 1 + n]; - try m.mtpMeta(sl, n); - sl.pos8 = .{ .b = r.ar.b, .off = fz.AR_ABS_POS * 4 }; - sl.nk8 = .{ .b = r.ar.b, .off = fz.AR_ABS_NK * 4 }; - sl.kvmeta = .{ .b = r.ar.b, .off = fz.AR_ABS_KV * 4 }; - } - for (1..MAXR - 1) |j| { // chained draft j: slot j - try m.mtpMeta(&e.gpu_slots[j], 1); - const b = (fz.AR_CHAIN + (j - 1) * fz.AR_CHAIN_STRIDE) * 4; - e.gpu_slots[j].pos8 = .{ .b = r.ar.b, .off = b }; - e.gpu_slots[j].nk8 = .{ .b = r.ar.b, .off = b + 16 * 4 }; - e.gpu_slots[j].kvmeta = .{ .b = r.ar.b, .off = b + 32 * 4 }; - } + return @import("engine_rounds.zig").rounds(e); } /// Host-driven calls (prompt chunks, the head's first drafts) read the host's metadata buffers. diff --git a/zig/src/families/flashnext/engine_rounds.zig b/zig/src/families/flashnext/engine_rounds.zig new file mode 100644 index 0000000000..a016f629e8 --- /dev/null +++ b/zig/src/families/flashnext/engine_rounds.zig @@ -0,0 +1,57 @@ +//! The GPU-side rounds' buffers an engine sets up once at load. +const fz = @import("replay.zig"); +const engine = @import("engine.zig"); +const Engine = engine.Engine; +const MAXR = fz.MAXR; +const CS_ROW = fz.CS_ROW; +const SO_ROW = fz.SO_ROW; +const PLE_TAIL = fz.PLE_TAIL; +const WIDE = fz.WIDE; +const CAP = fz.CAP; +const i32Buf = fz.i32Buf; + +/// DeltaNet states (kept, window rows), PLE tail pair, arena, token ring, per-width row counts, MTP slots. +pub fn rounds(e: *Engine) !void { + const r = e.r; + const m = e.m; + const n_lin: usize = 36; + e.g_cs = try r.buffer(n_lin * CS_ROW); + e.g_so = try r.buffer(n_lin * SO_ROW); + e.o_cs = try r.buffer(n_lin * MAXR * CS_ROW); + e.o_so = try r.buffer(n_lin * MAXR * SO_ROW); + var gi: usize = 0; + for (&m.layers) |*L| if (L.linear) { + L.cs[0] = .{ .b = e.g_cs, .off = gi * CS_ROW }; + L.cs[1] = .{ .b = e.o_cs, .off = gi * MAXR * CS_ROW }; + L.so[0] = .{ .b = e.g_so, .off = gi * SO_ROW }; + L.so[1] = .{ .b = e.o_so, .off = gi * MAXR * SO_ROW }; + gi += 1; + }; + if (r.gdn_kept != null) m.recs = .{ .{ .b = try r.buffer(n_lin * fz.gdn_step.RECORD) }, .{ .b = try r.buffer(n_lin * fz.gdn_step.RECORD) } }; + m.marks = try fz.Marks.init(r, n_lin); + e.cins = .{ m.ple.cin, .{ .b = try r.buffer((PLE_TAIL + MAXR) * WIDE * 2) } }; + r.ar = .{ .b = try r.buffer(4 * fz.AR_WORDS) }; + e.ring = try r.buffer(fz.RING_WORDS * 4 * engine.RING); + e.hist = try r.buffer((CAP + 64) * 4); + e.wids = .{ .b = try r.buffer(64) }; + for (1..MAXR + 1) |n| { + e.rows_w[n] = try i32Buf(r, &.{@intCast(n)}); + e.mdims_w[n] = try i32Buf(r, &.{ @intCast(n), 16, 0, 0, 0, 0, 0, 0 }); + } + e.host = .{ .rows = m.t.rows, .mdims = m.t.mdims, .pos8 = m.t.pos8, .nk8 = m.t.nk8, .kvmeta = m.t.kvmeta, .slots = m.mtp.slots }; + e.gpu_slots = m.mtp.slots; + for (1..MAXR + 1) |n| { // the head absorbing a window of n rows: slot MAXR - 1 + n + const sl = &e.gpu_slots[MAXR - 1 + n]; + try m.mtpMeta(sl, n); + sl.pos8 = .{ .b = r.ar.b, .off = fz.AR_ABS_POS * 4 }; + sl.nk8 = .{ .b = r.ar.b, .off = fz.AR_ABS_NK * 4 }; + sl.kvmeta = .{ .b = r.ar.b, .off = fz.AR_ABS_KV * 4 }; + } + for (1..MAXR - 1) |j| { // chained draft j: slot j + try m.mtpMeta(&e.gpu_slots[j], 1); + const b = (fz.AR_CHAIN + (j - 1) * fz.AR_CHAIN_STRIDE) * 4; + e.gpu_slots[j].pos8 = .{ .b = r.ar.b, .off = b }; + e.gpu_slots[j].nk8 = .{ .b = r.ar.b, .off = b + 16 * 4 }; + e.gpu_slots[j].kvmeta = .{ .b = r.ar.b, .off = b + 32 * 4 }; + } +} diff --git a/zig/src/families/flashnext/engine_warm.zig b/zig/src/families/flashnext/engine_warm.zig new file mode 100644 index 0000000000..cf7342a562 --- /dev/null +++ b/zig/src/families/flashnext/engine_warm.zig @@ -0,0 +1,67 @@ +//! Every kernel a checkpoint's widths need, compiled at load: a compile failure refuses the load, not a request. +const std = @import("std"); +const replay = @import("replay.zig"); +const variants = @import("variants.zig"); +const lane_w = @import("lane_w.zig"); +const Run = replay.Run; +const Model = replay.Model; + +/// `name` at every recorded row count on its kernel at `consts` (none: the recorded width, compiled already). +fn role(r: *Run, name: []const u8, consts: []const variants.Const) !void { + if (consts.len == 0) return; + var buf: [96]u8 = undefined; + for (1..replay.MAXR + 1) |n| if (r.roles.get(try std.fmt.bufPrint(&buf, "{s}|{d}", .{ name, n }))) |s| { + _ = try variants.patched(r, s.v, consts); + }; +} + +fn hc(r: *Run, h: replay.Hc, comptime at: []const u8) !void { + const w = h.w orelse return; + var c: [2][2]variants.Const = undefined; + try role(r, "qa_hc_down@" ++ at, variants.width(&c[0], "BITS", "GS", w[0])); + try role(r, "qa_hc_up@" ++ at, variants.width(&c[1], "BITS", "GS", w[1])); +} + +fn mtpHc(r: *Run, h: replay.Hc, comptime at: []const u8) !void { + const w = h.w orelse return; + var c: [2][2]variants.Const = undefined; + try role(r, "mtp:qa_hc_down@" ++ at, variants.width(&c[0], "BITS", "GS", w[0])); + try role(r, "mtp:qa_hc_up@" ++ at, variants.width(&c[1], "BITS", "GS", w[1])); +} + +fn experts(r: *Run, xw: ?variants.Experts, gu: []const u8, down: []const u8) !void { + const w = xw orelse return; + var cg: [4]variants.Const = undefined; + var cd: [4]variants.Const = undefined; + try role(r, gu, variants.four(&cg, w.gu)); + try role(r, down, variants.four(&cd, w.down)); +} + +/// A lane's runs at both row tilings (16 and 32 rows a tile). +fn lane(r: *Run, l: replay.Lane) !void { + const w = l.w orelse return; + for (w.runs) |p| for ([_]usize{ 1, 2 }) |tmr| { + _ = try lane_w.pipeline(r.arena, r.device, &r.wide_pipes, .{ .n = p.n, .k = p.k, .bits = p.bits, .group = p.group, .tmr = tmr, .ld = w.n, .mapped = p.map != null, .tiled = p.tiled }); + }; +} + +/// Every decode kernel `m`'s widths launch; none on a 6/32 checkpoint. +pub fn all(r: *Run, m: *const Model) !void { + for (&m.layers) |*L| { + try hc(r, L.ahc, "ahc"); + try hc(r, L.mhc, "mhc"); + try experts(r, L.xw, "qa_expert_gateup@moe.gate", "qa_expert_down_y@moe.down"); + try lane(r, L.proj); + try lane(r, L.out); + } + try hc(r, m.mix, "mix"); + try mtpHc(r, m.mtp.ahc, "mtp.ahc"); + try mtpHc(r, m.mtp.mhc, "mtp.mhc"); + try mtpHc(r, m.mtp.mix, "mtp.mix"); + try experts(r, m.mtp.xw, "mtp:qa_expert_gateup@mtp.moe.gate", "mtp:qa_expert_down_y@mtp.moe.down"); + for ([_]replay.Lane{ m.head, m.ple.kv, m.mtp.proj, m.mtp.out, m.mtp.fce, m.mtp.fch, m.mtp.draft }) |l| try lane(r, l); + var c: [2]variants.Const = undefined; + try role(r, "qa_embed_rows@embed", variants.width(&c, "BITS", "GS", m.embed_w)); + try role(r, "mtp:qa_embed_rows@embed", variants.width(&c, "BITS", "GS", m.embed_w)); + try role(r, "qa_ple_lookup@ple", variants.width(&c, "BITS", "GS", m.ple.w)); +} diff --git a/zig/src/families/flashnext/engine_widths.zig b/zig/src/families/flashnext/engine_widths.zig new file mode 100644 index 0000000000..8162786bf5 --- /dev/null +++ b/zig/src/families/flashnext/engine_widths.zig @@ -0,0 +1,96 @@ +//! The widths of the tensors the engine reads straight from the checkpoint: experts, embedding, n-gram tables. +const std = @import("std"); +const config = @import("config.zig"); +const widths = @import("widths.zig"); +const variants = @import("variants.zig"); +const replay = @import("replay.zig"); + +/// Null wherever a tensor is the recorded 6/32: those keep their recorded kernels. +pub const Set = struct { + layers: [config.max_layers]?variants.Experts = @splat(null), + mtp: ?variants.Experts = null, + embed: ?[2]usize = null, + ple: ?[2]usize = null, + /// Every quantized matrix 6/32: the uniform format, its packs pinned and its fast paths on. + uniform: bool = true, +}; + +fn width(w: *const widths.Widths, path: []const u8) ![2]usize { + const got = (try w.of(path)) orelse return error.UnquantizedAffineTensor; + return .{ got.bits, got.group }; +} + +fn same(x: [2]usize, y: [2]usize) bool { + return x[0] == y[0] and x[1] == y[1]; +} + +/// A layer's experts: gate and up share one width (one kernel reads both), routed and shared each their own. +pub fn experts(w: *const widths.Widths, stem: []const u8) !?variants.Experts { + var buf: [192]u8 = undefined; + var v: [6][2]usize = undefined; + const projs = [_][]const u8{ "switch_mlp.gate_proj", "switch_mlp.up_proj", "shared_expert.gate_proj", "shared_expert.up_proj", "switch_mlp.down_proj", "shared_expert.down_proj" }; + for (projs, &v) |p, *out| out.* = try width(w, try std.fmt.bufPrint(&buf, "{s}.mlp.{s}", .{ stem, p })); + if (!same(v[0], v[1]) or !same(v[2], v[3])) { + std.log.warn("{s}: gate and up of one expert differ in width; one kernel reads both", .{stem}); + return error.ExpertGateUpWidths; + } + const all_six = for (v) |x| { + if (!same(x, .{ 6, 32 })) break false; + } else true; + if (all_six) return null; + return .{ .gu = .{ v[0][0], v[0][1], v[2][0], v[2][1] }, .down = .{ v[4][0], v[4][1], v[5][0], v[5][1] } }; +} + +fn orNull(x: [2]usize) ?[2]usize { + return if (same(x, .{ 6, 32 })) null else x; +} + +fn sixUniform(w: *const widths.Widths) bool { + return w.uniform() and w.q.global.bits == 6 and w.q.global.group == 32; +} + +/// Whether a checkpoint folder is the uniform 6/32 format, which check programs that read only it require. +pub fn uniformDir(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !bool { + var c = try config.Config.read(gpa, io, dir); + defer c.deinit(); + var w = try widths.Widths.init(gpa, &c); + defer w.deinit(); + return sixUniform(&w); +} + +/// Every directly read tensor's width from the checkpoint's config; `spelling` names the n-gram shards. +pub fn read(c: *const config.Config, w: *const widths.Widths, spelling: []const u8) !Set { + var s: Set = .{}; + var buf: [192]u8 = undefined; + for (0..c.layers) |i| s.layers[i] = try experts(w, try std.fmt.bufPrint(&buf, "language_model.model.layers.{d}", .{i})); + if (c.mtp.layers > 0) s.mtp = try experts(w, "language_model.mtp.layers.0"); + s.embed = orNull(try width(w, "language_model.model.embed_tokens")); + for (0..c.layers) |i| if (c.ple[i]) { + const first = try width(w, try std.fmt.bufPrint(&buf, "language_model.model.layers.{d}.ple.ple_embedding.ngram_embedding.{s}0", .{ i, spelling })); + for (1..c.ngram_shards) |k| { + const other = try width(w, try std.fmt.bufPrint(&buf, "language_model.model.layers.{d}.ple.ple_embedding.ngram_embedding.{s}{d}", .{ i, spelling, k })); + if (!same(first, other)) return error.NgramTableWidths; // one lookup kernel reads every shard + } + s.ple = orNull(first); + }; + s.uniform = sixUniform(w); + return s; +} + +test { + _ = @import("engine_widths_test.zig"); +} + +/// Widths from the engine's own config and index into `m` and the prompt, kernels compiled; true when all 6/32. +pub fn apply(a: std.mem.Allocator, c: *const config.Config, map: std.json.ObjectMap, p: *replay.Prompt, m: *replay.Model) !bool { + var w = try widths.Widths.init(a, c); + defer w.deinit(); + const ws = try read(c, &w, @import("index.zig").ngramSpelling(map)); + for (&m.layers, 0..) |*L, i| L.xw = ws.layers[i]; + m.mtp.xw = ws.mtp; + m.embed_w = ws.embed; + m.ple.w = ws.ple; + try replay.prompt_w.attach(p, m); + try @import("engine_warm.zig").all(p.r, m); + return ws.uniform; +} diff --git a/zig/src/families/flashnext/engine_widths_test.zig b/zig/src/families/flashnext/engine_widths_test.zig new file mode 100644 index 0000000000..eb80b6bffb --- /dev/null +++ b/zig/src/families/flashnext/engine_widths_test.zig @@ -0,0 +1,92 @@ +//! The engine's directly read tensors take their declared widths: experts by layer, the embedding, the n-gram tables. +const std = @import("std"); +const config = @import("config.zig"); +const widths = @import("widths.zig"); +const ew = @import("engine_widths.zig"); +const a = std.testing.allocator; +const checked_config = @embedFile("fixtures/config.json"); + +/// The checked config at global `bits`/`group` (both quantization blocks) with `entries` as overrides. +fn with(bits: []const u8, group: []const u8, entries: []const u8) ![]u8 { + const declaration = try std.fmt.allocPrint(a, "\"mode\": \"affine\", {s}", .{entries}); + defer a.free(declaration); + const once = try std.mem.replaceOwned(u8, a, checked_config, "\"mode\": \"affine\"", declaration); + defer a.free(once); + const twice = try std.mem.replaceOwned(u8, a, once, "\"bits\": 6,", bits); + defer a.free(twice); + return std.mem.replaceOwned(u8, a, twice, "\"group_size\": 32,", group); +} + +fn readSet(bytes: []const u8) !ew.Set { + var c = try config.parse(a, bytes); + defer c.deinit(); + var w = try widths.Widths.init(a, &c); + defer w.deinit(); + return ew.read(&c, &w, "shard_"); +} + +test "oQ2's experts, embedding and n-gram tables take their own widths, layer by layer" { + const bytes = try with("\"bits\": 2,", "\"group_size\": 32,", + \\"language_model.model.embed_tokens": {"bits": 8, "group_size": 64}, + \\"language_model.model.layers.3.mlp.shared_expert.gate_proj": {"bits": 8, "group_size": 128}, + \\"language_model.model.layers.3.mlp.shared_expert.up_proj": {"bits": 8, "group_size": 128}, + \\"language_model.model.layers.3.mlp.shared_expert.down_proj": {"bits": 8, "group_size": 128}, + \\"language_model.mtp.layers.0.mlp.switch_mlp.gate_proj": {"bits": 4, "group_size": 32}, + \\"language_model.mtp.layers.0.mlp.switch_mlp.up_proj": {"bits": 4, "group_size": 32} + ); + defer a.free(bytes); + const s = try readSet(bytes); + try std.testing.expect(!s.uniform); + try std.testing.expectEqual([4]usize{ 2, 32, 8, 128 }, s.layers[3].?.gu); + try std.testing.expectEqual([4]usize{ 2, 32, 8, 128 }, s.layers[3].?.down); + try std.testing.expectEqual([4]usize{ 2, 32, 2, 32 }, s.layers[0].?.gu); + try std.testing.expectEqual([4]usize{ 4, 32, 2, 32 }, s.mtp.?.gu); + try std.testing.expectEqual([4]usize{ 2, 32, 2, 32 }, s.mtp.?.down); + try std.testing.expectEqual([2]usize{ 8, 64 }, s.embed.?); + try std.testing.expectEqual([2]usize{ 2, 32 }, s.ple.?); +} + +test "the uniform 6/32 checkpoint keeps every recorded kernel" { + const s = try readSet(checked_config); + try std.testing.expect(s.uniform); + try std.testing.expect(s.embed == null and s.ple == null and s.mtp == null); + for (s.layers) |l| try std.testing.expect(l == null); +} + +test "an expert whose gate and up differ in width is refused: one kernel reads both" { + const bytes = try with("\"bits\": 6,", "\"group_size\": 32,", + \\"language_model.model.layers.0.mlp.switch_mlp.up_proj": {"bits": 4, "group_size": 64} + ); + defer a.free(bytes); + try std.testing.expectError(error.ExpertGateUpWidths, readSet(bytes)); +} + +test "n-gram shards of different widths are refused: one lookup kernel reads every shard" { + const bytes = try with("\"bits\": 6,", "\"group_size\": 32,", + \\"language_model.model.layers.1.ple.ple_embedding.ngram_embedding.shard_1": {"bits": 4, "group_size": 32} + ); + defer a.free(bytes); + try std.testing.expectError(error.NgramTableWidths, readSet(bytes)); +} + +test "a check program that reads only 6/32 refuses a 4/32 checkpoint and one 6/32 checkpoint with an override" { + const io = std.testing.io; + const cases = [_]struct { bits: []const u8, entries: []const u8, six: bool }{ + .{ .bits = "\"bits\": 6,", .entries = "", .six = true }, + .{ .bits = "\"bits\": 4,", .entries = "", .six = false }, + .{ .bits = "\"bits\": 6,", .entries = "\"language_model.model.embed_tokens\": {\"bits\": 8, \"group_size\": 64}", .six = false }, + }; + for (cases) |case| { + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const bytes = if (case.entries.len == 0) + try std.mem.replaceOwned(u8, a, checked_config, "\"bits\": 6,", case.bits) + else + try with(case.bits, "\"group_size\": 32,", case.entries); + defer a.free(bytes); + try tmp.dir.writeFile(io, .{ .sub_path = "config.json", .data = bytes }); + const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(dir); + try std.testing.expectEqual(case.six, try ew.uniformDir(a, io, dir)); + } +} diff --git a/zig/src/families/flashnext/host_test.zig b/zig/src/families/flashnext/host_test.zig index 22edefa19a..bc198ce4ad 100644 --- a/zig/src/families/flashnext/host_test.zig +++ b/zig/src/families/flashnext/host_test.zig @@ -119,8 +119,10 @@ test "module quantization preserves false true and dictionary-default override s try std.testing.expectEqual(@as(usize, 64), override.group); try std.testing.expectEqual(@as(u8, 6), (try c.quantization("layers.0.inherit")).?.bits); try std.testing.expectEqual(@as(usize, 32), (try c.quantization("unknown.projection")).?.group); - try std.testing.expectError(error.UnsupportedFlashAffineKernel, override.checkFlashKernel()); - try c.global_affine.checkFlashKernel(); + const widths = @import("widths.zig"); + try std.testing.expect(widths.launched(.{ .bits = override.bits, .group = @intCast(override.group) })); + try std.testing.expect(widths.launched(.{ .bits = c.global_affine.bits, .group = @intCast(c.global_affine.group) })); + try std.testing.expect(!widths.launched(.{ .bits = 7, .group = 64 }) and !widths.launched(.{ .bits = 4, .group = 16 })); } test "malformed FlashNext configuration fields are refused" { diff --git a/zig/src/families/flashnext/index.zig b/zig/src/families/flashnext/index.zig index b3210aed15..44c3c6d012 100644 --- a/zig/src/families/flashnext/index.zig +++ b/zig/src/families/flashnext/index.zig @@ -1,11 +1,30 @@ //! Public index admission proves required tensor names and safe shard basenames before payload headers are available. const std = @import("std"); const Config = @import("config.zig").Config; +const widths = @import("widths.zig"); + +/// The first required matrix the device kernels cannot read yet, by tensor stem, with its declared width. +pub const Refusal = struct { + buf: [256]u8 = undefined, + len: usize = 0, + width: ?widths.Width = null, + + pub fn name(r: *const Refusal) []const u8 { + return r.buf[0..r.len]; + } + + fn set(r: *Refusal, stem: []const u8, width: ?widths.Width) void { + r.len = @min(stem.len, r.buf.len); + @memcpy(r.buf[0..r.len], stem[0..r.len]); + r.width = width; + } +}; const Check = struct { map: std.json.ObjectMap, names: usize = 0, - runtime: ?*const Config = null, + runtime: ?*const widths.Widths = null, + refusal: ?*Refusal = null, fn take(c: *Check, name: []const u8) !void { if (c.map.get(name) == null) return error.MissingFlashTensor; @@ -24,9 +43,12 @@ const Check = struct { const biases = c.map.get(try std.fmt.bufPrint(&buf, "{s}.biases", .{stem})) != null; if (scales != biases) return error.IncompleteFlashAffine; if (scales) { - if (c.runtime) |cfg| { - const spec = (try cfg.quantization(stem)) orelse return error.UnsupportedFlashAffineKernel; - try spec.checkFlashKernel(); + if (c.runtime) |w| { + const width = try w.of(stem); + if (width == null or !widths.launched(width.?)) { + if (c.refusal) |r| r.set(stem, width); + return error.UnsupportedFlashAffineKernel; + } } try c.part(stem, "scales"); try c.part(stem, "biases"); @@ -57,15 +79,33 @@ const Check = struct { } } - fn ple(c: *Check, stem: []const u8, shards: usize, spelling: []const u8) !void { + /// An n-gram table shard: one lookup reads a row of `row` values a group at a time, so its group divides the row. + fn table(c: *Check, stem: []const u8, row: usize) !void { + try c.matrix(stem); + const w = c.runtime orelse return; + const width = (try w.of(stem)) orelse return; + if (row % width.group == 0) return; + if (c.refusal) |r| r.set(stem, width); + return error.UnsupportedFlashAffineKernel; + } + + fn ple(c: *Check, stem: []const u8, shards: usize, spelling: []const u8, row: usize) !void { for ([_][]const u8{ "conv1d.weight", "norm_conv.weight", "norm_key.weight", "norm_query.weight", "ple_embedding.layer_multipliers", "ple_embedding.ngram_heads_offsets", "ple_embedding.ngram_heads_vocab_sizes" }) |suffix| try c.part(stem, suffix); try c.proj(stem, "key_proj"); try c.proj(stem, "value_proj"); var buf: [256]u8 = undefined; - for (0..shards) |i| try c.proj(stem, try std.fmt.bufPrint(&buf, "ple_embedding.ngram_embedding.{s}{d}", .{ spelling, i })); + var name: [256]u8 = undefined; + for (0..shards) |i| try c.table(try std.fmt.bufPrint(&name, "{s}.{s}", .{ stem, try std.fmt.bufPrint(&buf, "ple_embedding.ngram_embedding.{s}{d}", .{ spelling, i }) }), row); } }; +/// Where the index keeps the MTP head: `language_model.mtp` (mlx-lm), else oMLX's top-level `mtp`. +pub fn mtpRoot(map: std.json.ObjectMap) []const u8 { + for (map.keys()) |name| if (std.mem.startsWith(u8, name, "language_model.mtp.")) return "language_model.mtp"; + for (map.keys()) |name| if (std.mem.startsWith(u8, name, "mtp.")) return "mtp"; + return "language_model.mtp"; +} + /// The index's spelling of its n-gram table shards, `shard_N` when a `shard_0` is listed, else oMLX's `shards.N`. pub fn ngramSpelling(map: std.json.ObjectMap) []const u8 { for (map.keys()) |name| if (std.mem.endsWith(u8, name, ".ngram_embedding.shard_0.weight")) return "shard_"; @@ -75,16 +115,15 @@ pub fn ngramSpelling(map: std.json.ObjectMap) []const u8 { pub const Inventory = struct { tensor_names: usize, required_names: usize, shards: usize, headers_verified: bool = false }; pub fn admit(gpa: std.mem.Allocator, bytes: []const u8, config: *const Config) !Inventory { - return admitWith(gpa, bytes, config, false); + return admitWith(gpa, bytes, config, null, null); } -/// Runtime kernels consume only six-bit/group-32 text matrices, unlike the generic pack metadata reader. -pub fn admitRuntime(gpa: std.mem.Allocator, bytes: []const u8, config: *const Config) !Inventory { - try config.global_affine.checkFlashKernel(); - return admitWith(gpa, bytes, config, true); +/// Each required matrix at its own module width, refused by name when no device kernel reads that width yet. +pub fn admitRuntime(gpa: std.mem.Allocator, bytes: []const u8, config: *const Config, w: *const widths.Widths, refusal: *Refusal) !Inventory { + return admitWith(gpa, bytes, config, w, refusal); } -fn admitWith(gpa: std.mem.Allocator, bytes: []const u8, config: *const Config, runtime: bool) !Inventory { +fn admitWith(gpa: std.mem.Allocator, bytes: []const u8, config: *const Config, runtime: ?*const widths.Widths, refusal: ?*Refusal) !Inventory { const p = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); defer p.deinit(); if (p.value != .object) return error.InvalidFlashIndex; @@ -98,7 +137,7 @@ fn admitWith(gpa: std.mem.Allocator, bytes: []const u8, config: *const Config, r if (file != .string or file.string.len == 0 or file.string[0] == '.' or std.mem.indexOfAny(u8, file.string, "/\\") != null or !std.mem.endsWith(u8, file.string, ".safetensors")) return error.UnsafeFlashShard; try files.put(gpa, file.string, {}); } - var c = Check{ .map = map.object, .runtime = if (runtime) config else null }; + var c = Check{ .map = map.object, .runtime = runtime, .refusal = refusal }; const spelling = ngramSpelling(map.object); // every shard is checked under this one, so an index mixing the two is refused try c.matrix("language_model.model.embed_tokens"); try c.matrix("language_model.lm_head"); @@ -109,14 +148,16 @@ fn admitWith(gpa: std.mem.Allocator, bytes: []const u8, config: *const Config, r try c.layer(layer, try config.kind(i)); if (config.ple[i]) { var buf: [160]u8 = undefined; - try c.ple(try std.fmt.bufPrint(&buf, "{s}.ple", .{layer}), config.ngram_shards, spelling); + try c.ple(try std.fmt.bufPrint(&buf, "{s}.ple", .{layer}), config.ngram_shards, spelling, config.hidden / try config.ngramHeadCount()); } } if (config.mtp.layers > 0) { - for ([_][]const u8{ "fc_embedding", "fc_hidden" }) |suffix| try c.proj("language_model.mtp", suffix); - for ([_][]const u8{ "pre_fc_norm_embedding.weight", "pre_fc_norm_hidden.weight" }) |suffix| try c.part("language_model.mtp", suffix); - try c.hc("language_model.mtp.hyper_connection_mixer", false); - for (0..config.mtp.layers) |i| try c.layer(try std.fmt.bufPrint(&stem, "language_model.mtp.layers.{d}", .{i}), config.mtp.kinds[i]); + const root = mtpRoot(map.object); + for ([_][]const u8{ "fc_embedding", "fc_hidden" }) |suffix| try c.proj(root, suffix); + for ([_][]const u8{ "pre_fc_norm_embedding.weight", "pre_fc_norm_hidden.weight" }) |suffix| try c.part(root, suffix); + var mixer: [64]u8 = undefined; + try c.hc(try std.fmt.bufPrint(&mixer, "{s}.hyper_connection_mixer", .{root}), false); + for (0..config.mtp.layers) |i| try c.layer(try std.fmt.bufPrint(&stem, "{s}.layers.{d}", .{ root, i }), config.mtp.kinds[i]); } return .{ .tensor_names = map.object.count(), .required_names = c.names, .shards = files.count() }; } diff --git a/zig/src/families/flashnext/lane_w.zig b/zig/src/families/flashnext/lane_w.zig new file mode 100644 index 0000000000..020cfbeb38 --- /dev/null +++ b/zig/src/families/flashnext/lane_w.zig @@ -0,0 +1,201 @@ +//! Lane products at any MLX width on the recorded lane kernel's text: one launch per bit-width run, mapped columns. +const std = @import("std"); +const mtl = @import("metal"); +const ks = @import("kernel_sources"); +const replay = @import("replay.zig"); +const Run = replay.Run; +const Buf = replay.Buf; +const pack_lane = @import("pack_lane.zig"); +const variants = @import("variants.zig"); + +/// One run's product: n columns over k inputs at one width, and where its columns land in the fused output. +pub const Product = struct { wq: Buf, sbt: Buf, n: usize, k: usize, bits: u8, group: u16, map: ?Buf = null, tiled: bool = true }; + +/// A lane projection read at its own widths: its runs and the fused output's width. +pub const Wide = struct { runs: []const Product, n: usize }; + +/// A pack lane: the recorded 6/32 `.wq` and `.sbt`, or its runs at other widths (`w`); never both. +pub const Lane = struct { + six: ?Six = null, + w: ?*const Wide = null, + + pub const Six = struct { wq: Buf, sbt: Buf }; + + /// The buffers a 6/32 kernel reads; a lane at other widths has none, so no 6-bit kernel reads one of its runs. + pub fn recorded(l: Lane) error{LaneNot6Bit32}!Six { + return l.six orelse error.LaneNot6Bit32; + } +}; + +/// One compiled variant: shape, width, row tiles, output stride, column map, words tiled (else MLX's, decode.py nt 0). +pub const Spec = struct { n: usize, k: usize, bits: u8, group: u16, tmr: usize, ld: usize, mapped: bool, tiled: bool = true }; + +/// Compiled variants by their spec. +pub const Cache = std.AutoHashMapUnmanaged(Spec, mtl.Pipeline); + +/// K slices for an n by k product as lane_qmm.split_k: fixed by the shape, never by the row count. +pub fn splitK(n: usize, k: usize) usize { + const tiles = (n + 31) / 32; + var sk: usize = 1; + while (sk < 8 and tiles * sk < 1024 and (k / 64) / (sk * 2) >= 8) sk *= 2; + return sk; +} + +/// The recorded lane kernel's text for this GPU: the -lanes twin without tensor units. +pub fn template(tensor_units: bool) ![]const u8 { + const name = if (tensor_units) "lane_qmm_bytes_grouped_68a0697c6f38c26d.metal" else "lane_qmm_bytes_grouped_68a0697c6f38c26d-lanes.metal"; + for (ks.flashnext_gen.sources) |s| if (std.mem.eql(u8, s.name, name)) return s.text; + return error.NoSource; +} + +fn swap(a: std.mem.Allocator, text: []const u8, from: []const u8, to: []const u8) ![]const u8 { + if (std.mem.count(u8, text, from) != 1) return error.LanePatch; + return std.mem.replaceOwned(u8, a, text, from, to); +} + +/// The lane kernel at `s` (any width staged as bytes, sums over GS, stores at stride `ld`): the recorded arithmetic. +pub fn source(a: std.mem.Allocator, text_in: []const u8, s: Spec) ![]const u8 { + var text = text_in; + for ([_][]const u8{ "TMR", "N", "K", "SK", "BITS", "GS", "TILED" }, [_]usize{ s.tmr, s.n, s.k, splitK(s.n, s.k), s.bits, s.group, @intFromBool(s.tiled) }) |name, v| + text = try variants.constant(a, text, name, v); + text = try swap(a, text, "static_assert(BITS == 5 || BITS == 6 || BITS == 8, \"bytes for 5-, 6- and 8-bit weights\");", "static_assert(BITS >= 2 && BITS <= 8, \"codes of 2 to 8 bits staged as bytes\");"); + text = try swap(a, text, "for (int i = 0; i < 32; i++) acc += float(X[m * K + g * 32 + i]);", "for (int i = 0; i < FZ_GS; i++) acc += float(X[m * K + g * FZ_GS + i]);"); + const lanes = std.mem.indexOf(u8, text, "Y[m * N + nn + (j & 1) + 8 * (j >> 1)]") != null; + text = if (lanes) + try swap(a, text, "Y[m * N + nn + (j & 1) + 8 * (j >> 1)]", "Y[m * FZ_LD + FZ_COL(nn + (j & 1) + 8 * (j >> 1))]") + else + try swap(a, text, "Y[m * N + nn + j]", "Y[m * FZ_LD + FZ_COL(nn + j)]"); + if (s.mapped) text = try swap(a, text, " device bfloat16_t* Y [[buffer(5)]],\n", " device bfloat16_t* Y [[buffer(5)]],\n const device uint32_t* FZ_MAP [[buffer(6)]],\n"); + const col = if (s.mapped) "(FZ_MAP[c])" else "(c)"; + return std.fmt.allocPrint(a, "#define FZ_GS {d}\n#define FZ_LD {d}\n#define FZ_COL(c) {s}\n{s}", .{ s.group, s.ld, col, text }); +} + +/// The kernel function a source defines. +pub fn functionName(a: std.mem.Allocator, text: []const u8) ![:0]const u8 { + const at = (std.mem.indexOf(u8, text, "[[kernel]] void ") orelse return error.LanePatch) + "[[kernel]] void ".len; + const end = at + (std.mem.indexOfScalar(u8, text[at..], '(') orelse return error.LanePatch); + return std.fmt.allocPrintSentinel(a, "{s}", .{text[at..end]}, 0); +} + +/// The pipeline for `s` on `device`, compiled once and kept in `cache`. +pub fn pipeline(a: std.mem.Allocator, device: mtl.Device, cache: *Cache, s: Spec) !mtl.Pipeline { + if (s.n == 0 or s.n % 4 != 0 or s.tiled != (s.n % 32 == 0)) return error.LaneShape; + if (cache.get(s)) |p| return p; + const text = try source(a, try template(device.tensorUnits()), s); + const lib = try mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); + const pipe = try mtl.Pipeline.init(device, lib, try functionName(a, text), false); + try cache.put(a, s, pipe); + return pipe; +} + +/// Rows tile 16 at a time up to 32 (lane_qmm's ROW_BLOCK): a row's sums never depend on its window. +pub fn rowBlock(rows: usize) usize { + return @min((rows + 15) / 16 * 16, 32); +} + +/// y[rows, w.n] = x[rows, k] W^T, every run at its own width; `mdims` holds rows and rows padded to 16. +pub fn encode(enc: mtl.ComputeEncoder, a: std.mem.Allocator, device: mtl.Device, cache: *Cache, w: *const Wide, x: Buf, mdims: Buf, rows: usize, y: Buf) !void { + const block = rowBlock(rows); + const mp = (rows + 15) / 16 * 16; + for (w.runs) |p| { + const s: Spec = .{ .n = p.n, .k = p.k, .bits = p.bits, .group = p.group, .tmr = block / 16, .ld = w.n, .mapped = p.map != null, .tiled = p.tiled }; + enc.setPipeline(try pipeline(a, device, cache, s)); + for ([_]Buf{ x, x, p.wq, p.sbt, mdims, y }, 0..) |b, j| enc.setBuffer(b.b, b.off, j); + if (p.map) |m| enc.setBuffer(m.b, m.off, 6); + const sk = splitK(p.n, p.k); + enc.dispatchThreads(mtl.Size.of((p.n + 31) / 32 * 32 * sk, (mp + block - 1) / block, 1), mtl.Size.of(32 * sk, 1, 1)); + } +} + +/// `w` on the run's encoder at `rows` rows, unless the role is knocked out. +pub fn launch(r: *Run, w: *const Wide, role: []const u8, x: Buf, mdims: Buf, rows: usize, y: Buf) !void { + if (r.skip & Run.class(role) != 0) return; + try encode(r.enc, r.arena, r.device, &r.wide_pipes, w, x, mdims, rows, y); + if (!r.serial) r.enc.barrier(); +} + +/// Prefetches a lane's weights: the recorded pair, or every run's. +pub fn touch(r: *Run, l: Lane) void { + if (l.six) |s| for ([_]Buf{ s.wq, s.sbt }) |b| r.touch(b); + if (l.w) |w| for (w.runs) |p| for ([_]Buf{ p.wq, p.sbt }) |b| r.touch(b); +} + +/// Run columns to fused columns from a `.members` table of (run, row in run, row in fused output) per member. +pub fn columnMaps(a: std.mem.Allocator, members: []const u32, run_rows: []const usize) ![][]u32 { + const maps = try a.alloc([]u32, run_rows.len); + for (maps, run_rows) |*m, n| { + m.* = try a.alloc(u32, n); + @memset(m.*, std.math.maxInt(u32)); + } + const count = members.len / 3; + for (0..count) |i| { + const run = members[i * 3]; + if (run >= run_rows.len) return error.BadLaneMembers; + const start = members[i * 3 + 1]; + var end: usize = run_rows[run]; // a member runs to the next start in its run + for (0..count) |j| if (members[j * 3] == run and members[j * 3 + 1] > start) { + end = @min(end, members[j * 3 + 1]); + }; + if (start >= end) return error.BadLaneMembers; + for (start..end) |c| maps[run][c] = members[i * 3 + 2] + @as(u32, @intCast(c - start)); + } + var total: usize = 0; + for (run_rows) |n| total += n; + const seen = try a.alloc(bool, total); // every fused column written exactly once + @memset(seen, false); + for (maps) |m| for (m) |c| { + if (c >= total or seen[c]) return error.BadLaneMembers; + seen[c] = true; + }; + return maps; +} + +/// A pack lane's widths: null for a 6/32 lane (no `.runs` table), else its runs, each `.rN` when there are several. +pub fn of(r: *Run, stem: []const u8) !?*const Wide { + const a = r.arena; + const runs_name = try std.fmt.allocPrint(a, "{s}.runs", .{stem}); + if (!r.index.contains(runs_name)) return null; + const t = try variants.table(r, runs_name); + const n_runs = t.len / 3; + if (n_runs == 0 or t.len % 3 != 0) return error.BadLaneRuns; + const run_rows = try a.alloc(usize, n_runs); + for (run_rows, 0..) |*n, i| n.* = t[i * 3 + 2]; + var maps: ?[][]u32 = null; + if (n_runs > 1) { + const members_name = try std.fmt.allocPrint(a, "{s}.members", .{stem}); + maps = try columnMaps(a, try variants.table(r, members_name), run_rows); + } + const runs = try a.alloc(Product, n_runs); + var total: usize = 0; + for (runs, 0..) |*p, i| { + const base = if (n_runs == 1) stem else try std.fmt.allocPrint(a, "{s}.r{d}", .{ stem, i }); + const wq_name = try std.fmt.allocPrint(a, "{s}.wq", .{base}); + const bits: u8 = @intCast(t[i * 3]); + const n = run_rows[i]; + if (n == 0 or bits == 0) return error.BadLaneRuns; + const k = (try r.entry(wq_name)).len * 8 / (n * bits); + p.* = .{ .wq = try r.load(wq_name), .sbt = try r.loadf("{s}.sbt", .{base}), .n = n, .k = k, .bits = bits, .group = @intCast(t[i * 3 + 1]), .tiled = pack_lane.tiled(n) }; + if (maps) |m| { + const b = try r.buffer(n * 4); + @memcpy(b.slice(u32, n), m[i]); + p.map = .{ .b = b }; + } + total += n; + } + const w = try a.create(Wide); + w.* = .{ .runs = runs, .n = total }; + return w; +} + +/// A pack lane: its runs when not 6/32 (`of`), else the recorded `.wq` and `.sbt`. +pub fn laneOf(r: *Run, comptime fmt: []const u8, args: anytype) !Lane { + var name: [96]u8 = undefined; + const stem = try std.fmt.bufPrint(&name, fmt, args); + var full: [128]u8 = undefined; + if (try of(r, stem)) |w| return .{ .w = w }; + return .{ .six = .{ .wq = try r.load(try std.fmt.bufPrint(&full, "{s}.wq", .{stem})), .sbt = try r.load(try std.fmt.bufPrint(&full, "{s}.sbt", .{stem})) } }; +} + +test { + _ = @import("lane_w_test.zig"); +} diff --git a/zig/src/families/flashnext/lane_w_test.zig b/zig/src/families/flashnext/lane_w_test.zig new file mode 100644 index 0000000000..6190a20ee5 --- /dev/null +++ b/zig/src/families/flashnext/lane_w_test.zig @@ -0,0 +1,108 @@ +//! The lane products' widths on the host: K slices as recorded, every recorded text patched, columns put back. +const std = @import("std"); +const ks = @import("kernel_sources"); +const lane_w = @import("lane_w.zig"); +const mtl = @import("metal"); +const replay = @import("replay.zig"); +const variants = @import("variants.zig"); +const pio = @import("pack_io.zig"); + +test "split_k gives every recorded lane kernel's SK from its shape alone" { + const recorded = [_][3]usize{ .{ 13952, 2560, 4 }, .{ 16480, 2560, 2 }, .{ 2560, 6144, 8 }, .{ 248320, 2560, 1 }, .{ 12800, 2560, 4 }, .{ 2560, 2560, 4 }, .{ 79616, 2560, 1 } }; + for (recorded) |c| try std.testing.expectEqual(c[2], lane_w.splitK(c[0], c[1])); +} + +test "every recorded lane text takes any width, a row stride and a column map, and keeps its arithmetic" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var seen: usize = 0; + for (ks.flashnext_gen.sources) |src| { + if (std.mem.indexOf(u8, src.name, "lane_qmm_bytes_grouped") == null) continue; + seen += 1; + for ([_]bool{ false, true }) |mapped| { + const out = try lane_w.source(a, src.text, .{ .n = 512, .k = 6144, .bits = 3, .group = 64, .tmr = 2, .ld = 13952, .mapped = mapped }); + try std.testing.expect(std.mem.indexOf(u8, out, "constexpr int BITS = 3;") != null); + try std.testing.expect(std.mem.indexOf(u8, out, "constexpr int SK = 8;") != null); + try std.testing.expect(std.mem.indexOf(u8, out, "constexpr int TMR = 2;") != null); + try std.testing.expect(std.mem.indexOf(u8, out, "#define FZ_GS 64\n#define FZ_LD 13952\n") != null); + try std.testing.expectEqual(@intFromBool(mapped), std.mem.count(u8, out, "FZ_MAP [[buffer(6)]]")); + try std.testing.expectEqual(@as(usize, 0), std.mem.count(u8, out, "g * 32 + i")); + // the sums: one fma per scale and bias, unchanged + try std.testing.expectEqual(std.mem.count(u8, src.text, "fma("), std.mem.count(u8, out, "fma(")); + } + } + try std.testing.expectEqual(@as(usize, 16), seen); +} + +test "a run's columns land at their members' fused columns, adjacent or not" { + const a = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(a); + defer arena.deinit(); + // q 4 rows and idx 2 rows in run 0, v 3 rows in run 1: fused q | k(=run 0 too, 1 row) | v | idx + const members = [_]u32{ 0, 0, 0, 0, 4, 4, 1, 0, 5, 0, 5, 8 }; + const maps = try lane_w.columnMaps(arena.allocator(), &members, &.{ 7, 3 }); + try std.testing.expectEqualSlices(u32, &.{ 0, 1, 2, 3, 4, 8, 9 }, maps[0]); + try std.testing.expectEqualSlices(u32, &.{ 5, 6, 7 }, maps[1]); + try std.testing.expectError(error.BadLaneMembers, lane_w.columnMaps(arena.allocator(), &.{ 0, 0, 0, 2, 0, 7 }, &.{ 7, 3 })); +} + +/// A pack of a lane at two widths (8/64, 3/64), a 6/32 lane, and a hyper-connection with a widths table. +fn widePack(gpa: std.mem.Allocator, a: std.mem.Allocator, io: std.Io, path: []const u8) !void { + const synthetic = @import("pack_widen_test.zig").synthetic; + var out = pio.Out.init(gpa); + defer out.deinit(); + try @import("pack_lane.zig").putLane(&out, gpa, "L0.att.proj", &.{ try synthetic(a, 1, 64, 256, 8, 64), try synthetic(a, 2, 32, 256, 3, 64) }); + try @import("pack_lane.zig").putLane(&out, gpa, "L0.att.o", &.{try synthetic(a, 3, 32, 256, 6, 32)}); + for ([_][]const u8{ "scale", "down.w", "down.s", "down.b", "up.w", "up.s", "up.b" }) |part| { + const bytes = try gpa.alloc(u8, 64); + @memset(bytes, 0); + try out.put(try std.fmt.allocPrint(a, "L0.ahc.{s}", .{part}), "U8", &.{64}, bytes); + } + const w = try gpa.alloc(u32, 4); + w[0..4].* = .{ 4, 64, 4, 64 }; + try out.put("L0.ahc.widths", "U32", &.{4}, w); + try out.write(gpa, io, path, "test"); +} + +test "a lane's and a hyper-connection's width tables are read on the host: only weights reach the device" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const path = try std.fmt.allocPrintSentinel(a, ".zig-cache/tmp/{s}/pack.safetensors", .{tmp.sub_path}, 0); + try widePack(std.testing.allocator, a, std.testing.io, path); + var r = replay.Run{ .arena = a, .device = device, .queue = try device.queue() }; + try r.indexFile(path); + var weights: usize = 0; + for ([_][]const u8{ "L0.att.proj.r0.wq", "L0.att.proj.r0.sbt", "L0.att.proj.r1.wq", "L0.att.proj.r1.sbt" }) |n| weights += (try r.entry(n)).len; + const w = (try lane_w.of(&r, "L0.att.proj")).?; + try std.testing.expectEqual(@as(usize, 96), w.n); + try std.testing.expectEqual(weights, r.loaded); + const hc = try variants.hcOf(&r, "L0.ahc", .{}); + try std.testing.expectEqual(weights + 7 * 64, r.loaded); + try std.testing.expectEqual([2][2]usize{ .{ 4, 64 }, .{ 4, 64 } }, hc.w.?); +} + +test "a lane at other widths gives a 6-bit kernel nothing to read; a 6/32 lane gives its recorded pair" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const path = try std.fmt.allocPrintSentinel(a, ".zig-cache/tmp/{s}/pack.safetensors", .{tmp.sub_path}, 0); + try widePack(std.testing.allocator, a, std.testing.io, path); + var r = replay.Run{ .arena = a, .device = device, .queue = try device.queue() }; + try r.indexFile(path); + const wide = try lane_w.laneOf(&r, "L{d}.att.proj", .{0}); + try std.testing.expectEqual(@as(usize, 2), wide.w.?.runs.len); + try std.testing.expectError(error.LaneNot6Bit32, wide.recorded()); + const six = try lane_w.laneOf(&r, "L{d}.att.o", .{0}); + try std.testing.expect(six.w == null); + try std.testing.expectEqual((try r.entry("L0.att.o.wq")).len, (try six.recorded()).wq.b.length()); +} diff --git a/zig/src/families/flashnext/load_admission.zig b/zig/src/families/flashnext/load_admission.zig index 6970de32c5..ca3cda481e 100644 --- a/zig/src/families/flashnext/load_admission.zig +++ b/zig/src/families/flashnext/load_admission.zig @@ -2,27 +2,49 @@ const std = @import("std"); const config = @import("config.zig"); const index = @import("index.zig"); +const widths = @import("widths.zig"); -pub const problem = "Flash Next runtime kernels require the supported 6-bit affine, group_size 32 format; use a matching checkpoint"; +pub const problem = "Flash Next runtime kernels read MLX affine 2- to 8-bit weights in groups of 32, 64 or 128 (n-gram tables in groups that divide a row); the log names the first other one"; + +pub const stored_problem = "the checkpoint's stored weights contradict the widths its config declares, or a layout one kernel reads at one width differs; the log names the first one"; +pub const speed_up_problem = "speed-up mode (two Macs) splits the 6-bit group-32 kernels only; serve this checkpoint on one Mac"; +pub const scale_problem = "an n-gram table scale must be one finite, nonzero value; the log names it"; pub fn problemFor(err: anyerror) ?[]const u8 { return switch (err) { - error.UnsupportedFlashAffineKernel, error.UnsupportedAffineBits, error.UnsupportedAffineGroup, error.UnsupportedAffineMode => problem, + error.UnsupportedFlashAffineKernel, error.UnsupportedAffineBits, error.UnsupportedAffineGroup, error.UnsupportedAffineMode, error.UnsupportedQuantization => problem, + error.InvalidAffineShape, error.InvalidAffineRank, error.InvalidAffineStorage, error.InvalidAffineWidth, error.InvalidAffineByteCount, error.UnquantizedAffineTensor, error.ExpertGateUpWidths, error.NgramTableWidths, error.LaneRowsNotFours, error.TooManyLaneMembers => stored_problem, + error.SpeedUpNeeds6Bit32 => speed_up_problem, + error.NgramTableScale => scale_problem, else => null, }; } pub fn check(gpa: std.mem.Allocator, io: std.Io, directory: []const u8) !void { + var refusal: index.Refusal = .{}; + checkWith(gpa, io, directory, &refusal) catch |err| { + if (refusal.width) |w| std.log.warn("{s}: {d}-bit group {d} has no Flash Next device kernel", .{ refusal.name(), w.bits, w.group }); + if (refusal.len > 0 and refusal.width == null) std.log.warn("{s}: stored affine but declared unquantized", .{refusal.name()}); + return err; + }; +} + +/// Every required matrix resolved at its own module width; the first one no device kernel reads lands in `refusal`. +pub fn checkWith(gpa: std.mem.Allocator, io: std.Io, directory: []const u8, refusal: *index.Refusal) !void { var c = try config.Config.read(gpa, io, directory); defer c.deinit(); - try c.global_affine.checkFlashKernel(); + var w = try widths.Widths.init(gpa, &c); + defer w.deinit(); + // refused before the index + if (w.uniform() and !widths.launched(w.q.global)) return error.UnsupportedFlashAffineKernel; const path = try std.fs.path.join(gpa, &.{ directory, "model.safetensors.index.json" }); defer gpa.free(path); const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 26)); defer gpa.free(bytes); - _ = try index.admitRuntime(gpa, bytes, &c); + _ = try index.admitRuntime(gpa, bytes, &c, &w, refusal); } test { _ = @import("load_admission_test.zig"); + _ = @import("widths.zig"); } diff --git a/zig/src/families/flashnext/load_admission_test.zig b/zig/src/families/flashnext/load_admission_test.zig index 4a51579762..cde234ec9c 100644 --- a/zig/src/families/flashnext/load_admission_test.zig +++ b/zig/src/families/flashnext/load_admission_test.zig @@ -19,7 +19,7 @@ fn overridden(path: []const u8, value: []const u8) ![]u8 { return std.mem.replaceOwned(u8, a, checked_config, "\"mode\": \"affine\"", declaration); } -test "Flash Next load refuses a four-bit group-32 checkpoint before touching any weights or cache" { +test "Flash Next load admits a four-bit group-32 checkpoint without touching any weights or cache" { var tmp = std.testing.tmpDir(.{}); defer tmp.cleanup(); const bytes = try std.mem.replaceOwned(u8, a, checked_config, "\"bits\": 6", "\"bits\": 4"); @@ -28,11 +28,10 @@ test "Flash Next load refuses a four-bit group-32 checkpoint before touching any defer metadata.deinit(); try std.testing.expectEqual(@as(u8, 4), metadata.global_affine.bits); try std.testing.expectEqual(@as(usize, 32), metadata.global_affine.group); - try tmp.dir.writeFile(io, .{ .sub_path = "config.json", .data = bytes }); - const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + const dir = try fixture(tmp, bytes, checked_index); defer a.free(dir); - try std.testing.expectError(error.UnsupportedFlashAffineKernel, admission.check(a, io, dir)); - try std.testing.expectError(error.FileNotFound, tmp.dir.openFile(io, "model.safetensors.index.json", .{})); + try admission.check(a, io, dir); + try std.testing.expectError(error.FileNotFound, tmp.dir.openFile(io, "model-00001-of-00030.safetensors", .{})); try std.testing.expectError(error.FileNotFound, tmp.dir.openDir(io, "zig-pack", .{})); } @@ -46,7 +45,7 @@ test "load admission accepts the supported text and MTP matrices without reading try std.testing.expectError(error.FileNotFound, tmp.dir.openDir(io, "zig-pack", .{})); } -test "global unsupported affine groups refuse before the index is opened" { +test "a global group of 64 or 128 is admitted for every matrix, and the n-gram tables' 160-wide rows refuse it" { for ([_][]const u8{ "64", "128" }) |group| { var tmp = std.testing.tmpDir(.{}); defer tmp.cleanup(); @@ -54,18 +53,29 @@ test "global unsupported affine groups refuse before the index is opened" { defer a.free(declaration); const bytes = try std.mem.replaceOwned(u8, a, checked_config, "\"group_size\": 32", declaration); defer a.free(bytes); - try tmp.dir.writeFile(io, .{ .sub_path = "config.json", .data = bytes }); - const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + const dir = try fixture(tmp, bytes, checked_index); defer a.free(dir); - try std.testing.expectError(error.UnsupportedFlashAffineKernel, admission.check(a, io, dir)); + var refusal: @import("index.zig").Refusal = .{}; + try std.testing.expectError(error.UnsupportedFlashAffineKernel, admission.checkWith(a, io, dir, &refusal)); + try std.testing.expect(std.mem.indexOf(u8, refusal.name(), ".ple.ple_embedding.ngram_embedding.shard_0") != null); } } -test "required text and MTP affine overrides cannot bypass the runtime format" { +test "text and MTP overrides at other widths are admitted, and a stored matrix declared plain is refused" { for ([_]struct { path: []const u8, value: []const u8 }{ .{ .path = "language_model.model.layers.0.linear_attn.in_proj_qkv", .value = "{\"bits\":4,\"group_size\":32}" }, .{ .path = "model.layers.0.linear_attn.in_proj_qkv.weight", .value = "{\"bits\":6,\"group_size\":64}" }, - .{ .path = "language_model.mtp.layers.0.self_attn.q_proj", .value = "{\"bits\":4,\"group_size\":32}" }, + .{ .path = "language_model.mtp.layers.0.self_attn.q_proj", .value = "{\"bits\":3,\"group_size\":128}" }, + }) |override| { + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const bytes = try overridden(override.path, override.value); + defer a.free(bytes); + const dir = try fixture(tmp, bytes, checked_index); + defer a.free(dir); + try admission.check(a, io, dir); + } + for ([_]struct { path: []const u8, value: []const u8 }{ .{ .path = "language_model.lm_head", .value = "{}" }, .{ .path = "language_model.model.embed_tokens", .value = "false" }, }) |override| { @@ -121,6 +131,45 @@ test "normal and dump load paths admit before any engine device or pack work and const mapping = std.mem.indexOfPos(u8, native, open, "tf.flashnext_engine.quantizationProblem(e)"); try std.testing.expect(mapping != null); const problem = admission.problemFor(error.UnsupportedFlashAffineKernel).?; - try std.testing.expect(std.mem.indexOf(u8, problem, "supported 6-bit affine, group_size 32") != null); + try std.testing.expect(std.mem.indexOf(u8, problem, "2- to 8-bit weights in groups of 32, 64 or 128") != null); try std.testing.expect(admission.problemFor(error.OutOfMemory) == null); } + +test "a module the device kernels cannot read is refused under its checkpoint tensor name, whatever its key's prefix" { + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const bytes = try overridden("model.layers.1.ple.ple_embedding.ngram_embedding.shard_3.weight", "{\"bits\":8,\"group_size\":64}"); + defer a.free(bytes); + const dir = try fixture(tmp, bytes, checked_index); + defer a.free(dir); + var refusal: @import("index.zig").Refusal = .{}; + try std.testing.expectError(error.UnsupportedFlashAffineKernel, admission.checkWith(a, io, dir, &refusal)); + try std.testing.expectEqualStrings("language_model.model.layers.1.ple.ple_embedding.ngram_embedding.shard_3", refusal.name()); + try std.testing.expectEqual(@as(u8, 8), refusal.width.?.bits); + try std.testing.expectEqual(@as(u16, 64), refusal.width.?.group); +} + +test "an index keeping the MTP head at top-level mtp (oMLX) is admitted as one keeping it under language_model" { + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const index = try std.mem.replaceOwned(u8, a, checked_index, "\"language_model.mtp.", "\"mtp."); + defer a.free(index); + try std.testing.expect(std.mem.indexOf(u8, index, "\"mtp.fc_embedding.weight\"") != null); + const dir = try fixture(tmp, checked_config, index); + defer a.free(dir); + try admission.check(a, io, dir); + const missing = try std.mem.replaceOwned(u8, a, index, "\"mtp.fc_hidden.weight\"", "\"mtp.fc_hidden.other\""); + defer a.free(missing); + var tmp2 = std.testing.tmpDir(.{}); + defer tmp2.cleanup(); + const dir2 = try fixture(tmp2, checked_config, missing); + defer a.free(dir2); + try std.testing.expectError(error.MissingFlashTensor, admission.check(a, io, dir2)); +} + +test "the width and layout refusals past admission map to text a user can act on" { + for ([_]anyerror{ error.InvalidAffineShape, error.ExpertGateUpWidths, error.NgramTableWidths, error.LaneRowsNotFours, error.TooManyLaneMembers }) |e| + try std.testing.expectEqualStrings(admission.stored_problem, admission.problemFor(e).?); + try std.testing.expectEqualStrings(admission.speed_up_problem, admission.problemFor(error.SpeedUpNeeds6Bit32).?); + try std.testing.expectEqualStrings(admission.scale_problem, admission.problemFor(error.NgramTableScale).?); +} diff --git a/zig/src/families/flashnext/pack.zig b/zig/src/families/flashnext/pack.zig index 7b1105e9d2..2834475232 100644 --- a/zig/src/families/flashnext/pack.zig +++ b/zig/src/families/flashnext/pack.zig @@ -7,6 +7,9 @@ const affine = @import("affine.zig"); const config_mod = @import("config.zig"); const index_mod = @import("index.zig"); const pio = @import("pack_io.zig"); +const quant_mlx = @import("../../core/quant/mlx.zig"); +const quant_widen = @import("../../core/quant/widen.zig"); +const widen = @import("pack_widen.zig"); pub const Report = struct { decode_tensors: usize = 0, @@ -16,31 +19,24 @@ pub const Report = struct { draft_ids: usize = 0, }; -const Out = pio.Out; +pub const Out = pio.Out; -/// Which of a quantized linear's three tensors. -const Part = enum { weight, scales, biases }; - -/// A checkpoint quantized linear's three tensors with its verified affine spec. -pub const Q = struct { - w: st.Tensor, - s: st.Tensor, - b: st.Tensor, - rows: usize, - k: usize, - kw: usize, - spec: affine.Spec, - - fn part(r: Q, comptime which: Part) []const u8 { - return switch (which) { - .weight => r.w.bytes, - .scales => r.s.bytes, - .biases => r.b.bytes, - }; - } -}; +pub const lane = @import("pack_lane.zig"); +pub const Q = lane.Q; +pub const tileWeight = lane.tileWeight; +pub const packScales = lane.packScales; +const concatPart = lane.concatPart; +const putLane = lane.putLane; +/// A quantized linear read only at its own declared width, refused under its name otherwise. fn qlinear(ck: *ckpt.Checkpoint, cfg: *const config_mod.Config, path: []const u8) !Q { + return qlinearAt(ck, cfg, path) catch |err| { + std.log.warn("{s}: not read at its declared width: {s}", .{ path, @errorName(err) }); + return err; + }; +} + +fn qlinearAt(ck: *ckpt.Checkpoint, cfg: *const config_mod.Config, path: []const u8) !Q { const spec = (try cfg.quantization(path)) orelse return error.NotQuantized; var buf: [320]u8 = undefined; const w = try ck.get(try std.fmt.bufPrint(&buf, "{s}.weight", .{path})); @@ -48,7 +44,6 @@ fn qlinear(ck: *ckpt.Checkpoint, cfg: *const config_mod.Config, path: []const u8 const b = try ck.get(try std.fmt.bufPrint(&buf, "{s}.biases", .{path})); if (w.dtype != .u32 or w.rank != 2) return error.UnexpectedTensor; if (s.dtype != .bf16 or s.rank != 2 or b.dtype != .bf16 or b.rank != 2) return error.UnexpectedTensor; - if (spec.group == 128) return error.Group128NotPacked; // the pack never carries group 128; the halve-to-64 widen is unimplemented const rows = w.dim(0); const kw = w.dim(1); if (rows != s.dim(0) or rows != b.dim(0)) return error.UnexpectedTensor; @@ -58,79 +53,6 @@ fn qlinear(ck: *ckpt.Checkpoint, cfg: *const config_mod.Config, path: []const u8 return .{ .w = w, .s = s, .b = b, .rows = rows, .k = k, .kw = kw, .spec = spec }; } -/// The lane pipeline's word permutation: rows tiled nt at a time, a group's words innermost (lane_qmm.tile_weight). -pub fn tileWeight(gpa: std.mem.Allocator, flat: []const u8, rows: usize, kw: usize, nt: usize, w: usize) ![]u8 { - const words = std.mem.bytesAsSlice(u32, flat); - const out = try gpa.alloc(u8, flat.len); - errdefer gpa.free(out); - const out_words = std.mem.bytesAsSlice(u32, out); - const groups = kw / w; - for (0..rows / nt) |t| { - for (0..groups) |g| { - for (0..nt) |j| { - for (0..w) |c| { - out_words[((t * groups + g) * nt + j) * w + c] = words[(t * nt + j) * kw + g * w + c]; - } - } - } - } - return out; -} - -/// (K/GS, N, 2) bf16 (scale, bias) pairs, group-major, over a row-concatenated stack (lane_qmm.pack_scales). -pub fn packScales(gpa: std.mem.Allocator, members: []const Q) ![]u8 { - var rows: usize = 0; - for (members) |m| rows += m.rows; - const kg = members[0].s.dim(1); - const out = try gpa.alloc(u8, kg * rows * 4); - errdefer gpa.free(out); - const out16 = std.mem.bytesAsSlice(u16, out); - var at: usize = 0; - for (members) |m| { - const s16 = std.mem.bytesAsSlice(u16, m.part(.scales)); - const b16 = std.mem.bytesAsSlice(u16, m.part(.biases)); - for (0..kg) |g| { - for (0..m.rows) |j| { - out16[(g * rows + at + j) * 2] = s16[j * kg + g]; - out16[(g * rows + at + j) * 2 + 1] = b16[j * kg + g]; - } - } - at += m.rows; - } - return out; -} - -/// Row-concatenated members in the checkpoint's own MLX layout: a pure byte move per member. -fn concatPart(gpa: std.mem.Allocator, members: []const Q, comptime part: Part) ![]u8 { - var n: usize = 0; - for (members) |m| n += m.part(part).len; - const out = try gpa.alloc(u8, n); - errdefer gpa.free(out); - var at: usize = 0; - for (members) |m| { - @memcpy(out[at..][0..m.part(part).len], m.part(part)); - at += m.part(part).len; - } - return out; -} - -/// One lane projection's two pack tensors: `.wq` tiled, `.sbt` group-major pairs (decode.py's lane pipeline). -fn putLane(out: *Out, gpa: std.mem.Allocator, name: []const u8, members: []const Q) !void { - var rows: usize = 0; - for (members) |m| rows += m.rows; - const spec = members[0].spec; - for (members) |m| if (m.spec.bits != spec.bits or m.spec.group != spec.group) return error.MixedAffine; - const nt: usize = if (spec.bits == 4 and rows % 64 == 0) 64 else if (rows % 32 == 0) 32 else return error.UnsupportedLaneWidth; - const w = spec.group * spec.bits / 32; - const flat = try concatPart(gpa, members, .weight); - defer gpa.free(flat); - const wq = try tileWeight(gpa, flat, rows, members[0].kw, nt, w); - const sbt = try packScales(gpa, members); - const a = out.arena.allocator(); - try out.put(try std.fmt.allocPrint(a, "{s}.wq", .{name}), "U32", &.{ rows, members[0].kw }, wq); - try out.put(try std.fmt.allocPrint(a, "{s}.sbt", .{name}), "BF16", &.{ members[0].s.dim(1), rows, 2 }, sbt); -} - fn bf16ToF32(bits: u16) f32 { return @bitCast(@as(u32, bits) << 16); } @@ -195,22 +117,29 @@ fn putHc(out: *Out, gpa: std.mem.Allocator, a: std.mem.Allocator, ck: *ckpt.Chec } var storage: [2]Q = undefined; for (down_paths[0..down_n], 0..) |p, i| storage[i] = try qlinear(ck, cfg, p); - const down = storage[0..down_n]; - for (down) |m| if (m.spec.bits != down[0].spec.bits or m.spec.group != down[0].spec.group) return error.MixedAffine; - var rows: usize = 0; - for (down) |m| rows += m.rows; - const scale = try centeredScale(gpa, ck, try std.fmt.allocPrint(a, "{s}.hc_norm.weight", .{stem}), around_one); - try out.put(try std.fmt.allocPrint(a, "{s}.scale", .{name}), "F32", &.{wide}, scale); - const down_w = try concatPart(gpa, down, .weight); - const down_s = try concatPart(gpa, down, .scales); - const down_b = try concatPart(gpa, down, .biases); - try out.put(try std.fmt.allocPrint(a, "{s}.down.w", .{name}), "U32", &.{ rows, down[0].kw }, down_w); - try out.put(try std.fmt.allocPrint(a, "{s}.down.s", .{name}), "BF16", &.{ rows, down[0].s.dim(1) }, down_s); - try out.put(try std.fmt.allocPrint(a, "{s}.down.b", .{name}), "BF16", &.{ rows, down[0].b.dim(1) }, down_b); + const down = try hcDown(a, storage[0..down_n]); // a pair of two widths widens to the wider (QWeights.of) const up = try qlinear(ck, cfg, try std.fmt.allocPrint(a, "{s}.input_mix_weight_up", .{stem})); + const scale_name = try std.fmt.allocPrint(a, "{s}.scale", .{name}); + try out.put(scale_name, "F32", &.{wide}, try centeredScale(gpa, ck, try std.fmt.allocPrint(a, "{s}.hc_norm.weight", .{stem}), around_one)); + try out.put(try std.fmt.allocPrint(a, "{s}.down.w", .{name}), "U32", &.{ down.rows, down.kw }, try gpa.dupe(u8, down.part(.weight))); + try out.put(try std.fmt.allocPrint(a, "{s}.down.s", .{name}), "BF16", &.{ down.rows, down.s.dim(1) }, try gpa.dupe(u8, down.part(.scales))); + try out.put(try std.fmt.allocPrint(a, "{s}.down.b", .{name}), "BF16", &.{ down.rows, down.b.dim(1) }, try gpa.dupe(u8, down.part(.biases))); try out.put(try std.fmt.allocPrint(a, "{s}.up.w", .{name}), "U32", &.{ up.rows, up.kw }, try gpa.dupe(u8, up.part(.weight))); try out.put(try std.fmt.allocPrint(a, "{s}.up.s", .{name}), "BF16", &.{ up.rows, up.s.dim(1) }, try gpa.dupe(u8, up.part(.scales))); try out.put(try std.fmt.allocPrint(a, "{s}.up.b", .{name}), "BF16", &.{ up.rows, up.b.dim(1) }, try gpa.dupe(u8, up.part(.biases))); + // the uniform packs' pinned bytes + if (down.spec.bits == 6 and down.spec.group == 32 and up.spec.bits == 6 and up.spec.group == 32) return; + const widths_name = try std.fmt.allocPrint(a, "{s}.widths", .{name}); + const w = try gpa.alloc(u32, 4); + w[0..4].* = .{ down.spec.bits, @intCast(down.spec.group), up.spec.bits, @intCast(up.spec.group) }; + try out.put(widths_name, "U32", &.{4}, w); +} + +/// The down projection and the inject rows as one matrix: one width as stored, two widened exactly to the wider. +pub fn hcDown(a: std.mem.Allocator, members: []const Q) !Q { + var hosts: [2]quant_mlx.Host = undefined; + for (members, 0..) |m, i| hosts[i] = widen.host(m); + return widen.fromHost(try quant_widen.fuse(a, hosts[0..members.len])); } /// A gate's rows are stored bf16, or an affine linear dequantized on the host as code * scale + bias per group. @@ -332,6 +261,20 @@ fn pleStarts(gpa: std.mem.Allocator, ck: *ckpt.Checkpoint, cfg: *const config_mo return out; } +/// An n-gram table's weight_scale must be one finite, nonzero bf16 or f32 value: the lookups multiply rows by it. +fn tableScale(ck: *ckpt.Checkpoint, map: std.json.ObjectMap, name: []const u8) !void { + if (!map.contains(name)) return; + const t = try ck.get(name); + const v: ?f32 = switch (t.dtype) { + .bf16 => if (t.bytes.len == 2) @bitCast(@as(u32, std.mem.readInt(u16, t.bytes[0..2], .little)) << 16) else null, + .f32 => if (t.bytes.len == 4) @bitCast(std.mem.readInt(u32, t.bytes[0..4], .little)) else null, + else => null, + }; + if (v) |x| if (std.math.isFinite(x) and x != 0) return; + std.log.warn("{s}: an n-gram table scale must be one finite, nonzero value", .{name}); + return error.NgramTableScale; +} + /// Draft vocabulary ids are sorted and padded to a multiple of 64, and an id past the vocabulary is refused. pub fn draftIdList(gpa: std.mem.Allocator, io: Io, path: []const u8, vocab: usize) ![]u32 { const text = try Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 26)); @@ -371,9 +314,12 @@ fn putDraftHead(out: *Out, gpa: std.mem.Allocator, a: std.mem.Allocator, name: [ const rows = ids.len; const kw = lm_head.kw; const kg = lm_head.s.dim(1); - const wq = try gpa.alloc(u8, rows * kw * 4); + const wq = try gpa.alloc(u8, rows * kw * 4); // putLane copies the gathered rows, so they free here + defer gpa.free(wq); const s = try gpa.alloc(u8, rows * kg * 2); + defer gpa.free(s); const b = try gpa.alloc(u8, rows * kg * 2); + defer gpa.free(b); const w32 = std.mem.bytesAsSlice(u32, wq); const src32 = std.mem.bytesAsSlice(u32, lm_head.part(.weight)); const s16 = std.mem.bytesAsSlice(u16, s); @@ -387,10 +333,7 @@ fn putDraftHead(out: *Out, gpa: std.mem.Allocator, a: std.mem.Allocator, name: [ @memcpy(b16[r * kg ..][0..kg], src_b16[id * kg ..][0..kg]); } const gathered = [_]Q{.{ .w = .{ .dtype = .u32, .rank = 2, .shape = .{ rows, kw, 0, 0, 0 }, .bytes = wq }, .s = .{ .dtype = .bf16, .rank = 2, .shape = .{ rows, kg, 0, 0, 0 }, .bytes = s }, .b = .{ .dtype = .bf16, .rank = 2, .shape = .{ rows, kg, 0, 0, 0 }, .bytes = b }, .rows = rows, .k = lm_head.k, .kw = kw, .spec = lm_head.spec }}; - try putLane(out, gpa, name, gathered[0..]); // putLane copies the bytes, so the gathered buffers free here - gpa.free(wq); - gpa.free(s); - gpa.free(b); + try putLane(out, gpa, name, gathered[0..]); _ = a; } @@ -401,6 +344,9 @@ pub fn build(gpa: std.mem.Allocator, io: Io, model_dir: []const u8, out_dir: []c const a = arena_state.allocator(); var cfg = try config_mod.Config.read(gpa, io, model_dir); defer cfg.deinit(); + var strict = try @import("widths.zig").Widths.init(gpa, &cfg); + defer strict.deinit(); + cfg.widths = &strict; // the vision tower's tensors stay unread: a rank-five entry must not stop the no-dump build var ck = try ckpt.Checkpoint.openModelPrefix(gpa, io, model_dir, "language_model."); defer ck.close(); @@ -416,6 +362,11 @@ pub fn build(gpa: std.mem.Allocator, io: Io, model_dir: []const u8, out_dir: []c defer parsed_index.deinit(); const weight_map = parsed_index.value.object.get("weight_map").?.object; const spelling = index_mod.ngramSpelling(weight_map); + const mtp_root = index_mod.mtpRoot(weight_map); // oMLX keeps the head outside language_model. + var ck_mtp = if (std.mem.eql(u8, mtp_root, "mtp")) try ckpt.Checkpoint.openModelPrefix(gpa, io, model_dir, "mtp.") else null; + defer if (ck_mtp) |*m| m.close(); + const ckm = if (ck_mtp) |*m| m else &ck; + try @import("pack_raw.zig").checkAll(&ck, ckm, &strict, &cfg, mtp_root, spelling); // the directly read tensors const identity = try pio.sourceIdentity(gpa, io, model_dir); defer gpa.free(identity); @@ -491,13 +442,14 @@ pub fn build(gpa: std.mem.Allocator, io: Io, model_dir: []const u8, out_dir: []c try decode.put(try std.fmt.allocPrint(a, "ple.{s}", .{pack_name}), "F32", &.{wide}, scale); } try putConvSlice(&decode, gpa, &ck, "ple.conv", try std.fmt.allocPrint(a, "{s}.ple.conv1d.weight", .{stem}), true); + try tableScale(&ck, weight_map, try std.fmt.allocPrint(a, "{s}.ple.ple_embedding.ngram_embedding.weight_scale", .{stem})); const starts = try pleStarts(gpa, &ck, &cfg, pl, spelling); try decode.put("ple.starts", "U32", &.{8}, starts); } if (draft_vocab) |vocab_path| { const ids = try draftIdList(gpa, io, vocab_path, cfg.vocab); report.draft_ids = ids.len; - try putMtp(&decode, gpa, a, &ck, &cfg, ids, &report, wide); + try putMtp(&decode, gpa, a, ckm, &ck, mtp_root, &cfg, ids, &report, wide); try decode.put("mtp.draft_ids", "U32", &.{ids.len}, ids); } report.decode_tensors = decode.names.items.len; @@ -527,34 +479,52 @@ pub fn build(gpa: std.mem.Allocator, io: Io, model_dir: []const u8, out_dir: []c if (cfg.mtp.layers > 0) { var mtp_mlx = Out.init(gpa); defer mtp_mlx.deinit(); - const mtp_stem = "language_model.mtp.layers.0"; - try putMlxProjection(&mtp_mlx, gpa, a, &ck, &cfg, "mtp.att.proj", &.{ "self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.indexer.index_qk_proj" }, mtp_stem); - try putMlxProjection(&mtp_mlx, gpa, a, &ck, &cfg, "mtp.fce", &.{"fc_embedding"}, "language_model.mtp"); - try putMlxProjection(&mtp_mlx, gpa, a, &ck, &cfg, "mtp.fch", &.{"fc_hidden"}, "language_model.mtp"); + const mtp_stem = try std.fmt.allocPrint(a, "{s}.layers.0", .{mtp_root}); + try putMlxProjection(&mtp_mlx, gpa, a, ckm, &cfg, "mtp.att.proj", &.{ "self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.indexer.index_qk_proj" }, mtp_stem); + try putMlxProjection(&mtp_mlx, gpa, a, ckm, &cfg, "mtp.fce", &.{"fc_embedding"}, mtp_root); + try putMlxProjection(&mtp_mlx, gpa, a, ckm, &cfg, "mtp.fch", &.{"fc_hidden"}, mtp_root); report.mtp_mlx_tensors = mtp_mlx.names.items.len; try mtp_mlx.write(gpa, io, try std.fs.path.join(a, &.{ out_dir, "pack_mtp_mlx.safetensors" }), identity); } return report; } -/// One export_mlx_proj-style projection: `.mw` the checkpoint's packed words, `.ms`/`.mb` its scales and biases. +/// One export_mlx_proj-style projection (`.mw` words, `.ms`/`.mb` tables); `.parts` gives other widths per member. fn putMlxProjection(out: *Out, gpa: std.mem.Allocator, a: std.mem.Allocator, ck: *ckpt.Checkpoint, cfg: *const config_mod.Config, name: []const u8, members: []const []const u8, stem: []const u8) !void { var storage: [4]Q = undefined; for (members, 0..) |n, i| storage[i] = try qlinear(ck, cfg, try std.fmt.allocPrint(a, "{s}.{s}", .{ stem, n })); - const list = storage[0..members.len]; - const mw = try concatPart(gpa, list, .weight); - const ms = try concatPart(gpa, list, .scales); - const mb = try concatPart(gpa, list, .biases); + return mlxProjection(out, gpa, a, name, storage[0..members.len]); +} + +/// The members' MLX bytes in member order; 6/32 alone keeps its 2-D shapes, other widths go flat with `.parts`. +pub fn mlxProjection(out: *Out, gpa: std.mem.Allocator, a: std.mem.Allocator, name: []const u8, list: []const Q) !void { var rows: usize = 0; for (list) |m| rows += m.rows; - try out.put(try std.fmt.allocPrint(a, "{s}.mw", .{name}), "U32", &.{ rows, list[0].kw }, mw); - try out.put(try std.fmt.allocPrint(a, "{s}.ms", .{name}), "BF16", &.{ rows, list[0].s.dim(1) }, ms); - try out.put(try std.fmt.allocPrint(a, "{s}.mb", .{name}), "BF16", &.{ rows, list[0].b.dim(1) }, mb); + const six = for (list) |m| { + if (m.spec.bits != 6 or m.spec.group != 32) break false; + } else true; + // the uniform packs keep their bytes and shapes; members of other widths go flat, in member order + try mlxPart(out, gpa, a, name, "mw", "U32", list, .weight, if (six) .{ rows, list[0].kw } else null); + try mlxPart(out, gpa, a, name, "ms", "BF16", list, .scales, if (six) .{ rows, list[0].s.dim(1) } else null); + try mlxPart(out, gpa, a, name, "mb", "BF16", list, .biases, if (six) .{ rows, list[0].b.dim(1) } else null); + if (six) return; + const parts_name = try std.fmt.allocPrint(a, "{s}.parts", .{name}); // each member's rows, bits and group + const parts = try gpa.alloc(u32, list.len * 3); + for (list, 0..) |m, i| parts[i * 3 ..][0..3].* = .{ @intCast(m.rows), m.spec.bits, @intCast(m.spec.group) }; + try out.put(parts_name, "U32", &.{ list.len, 3 }, parts); +} + +/// One of a projection's three tensors, members concatenated: at `shape` when given, else flat. +fn mlxPart(out: *Out, gpa: std.mem.Allocator, a: std.mem.Allocator, name: []const u8, suffix: []const u8, dtype: []const u8, list: []const Q, comptime part: lane.Part, shape: ?[2]usize) !void { + const full = try std.fmt.allocPrint(a, "{s}.{s}", .{ name, suffix }); + const bytes = try concatPart(gpa, list, part); + if (shape) |sh| return out.put(full, dtype, &sh, bytes); + return out.put(full, dtype, &.{bytes.len / @as(usize, if (part == .weight) 4 else 2)}, bytes); } /// The mtp.* decode tensors: the MTP layer's hc/att/moe classes, its norms, and the draft head over the id list. -fn putMtp(out: *Out, gpa: std.mem.Allocator, a: std.mem.Allocator, ck: *ckpt.Checkpoint, cfg: *const config_mod.Config, ids: []const u32, report: *Report, wide: usize) !void { - const layer_stem = "language_model.mtp.layers.0"; +fn putMtp(out: *Out, gpa: std.mem.Allocator, a: std.mem.Allocator, ck: *ckpt.Checkpoint, text: *ckpt.Checkpoint, root: []const u8, cfg: *const config_mod.Config, ids: []const u32, report: *Report, wide: usize) !void { + const layer_stem = try std.fmt.allocPrint(a, "{s}.layers.0", .{root}); try putHc(out, gpa, a, ck, cfg, "mtp.ahc", try std.fmt.allocPrint(a, "{s}.attn_hyper_connection", .{layer_stem}), true, report.norms_around_one, wide); try putHc(out, gpa, a, ck, cfg, "mtp.mhc", try std.fmt.allocPrint(a, "{s}.mlp_hyper_connection", .{layer_stem}), true, report.norms_around_one, wide); try putRouter(out, gpa, a, ck, cfg, "mtp.moe.router", try std.fmt.allocPrint(a, "{s}.mlp.gate", .{layer_stem}), try std.fmt.allocPrint(a, "{s}.mlp.shared_expert_gate", .{layer_stem})); @@ -573,18 +543,18 @@ fn putMtp(out: *Out, gpa: std.mem.Allocator, a: std.mem.Allocator, ck: *ckpt.Che const scale = try centeredScale(gpa, ck, try std.fmt.allocPrint(a, "{s}.self_attn.indexer.{s}", .{ layer_stem, n }), report.norms_around_one); try out.put(try std.fmt.allocPrint(a, "mtp.att.{s}", .{pack_name}), "F32", &.{cfg.index_dim}, scale); } - try putHc(out, gpa, a, ck, cfg, "mtp.mix", "language_model.mtp.hyper_connection_mixer", false, report.norms_around_one, wide); - const fce = try qlinear(ck, cfg, "language_model.mtp.fc_embedding"); + try putHc(out, gpa, a, ck, cfg, "mtp.mix", try std.fmt.allocPrint(a, "{s}.hyper_connection_mixer", .{root}), false, report.norms_around_one, wide); + const fce = try qlinear(ck, cfg, try std.fmt.allocPrint(a, "{s}.fc_embedding", .{root})); const fce_members = [_]Q{fce}; try putLane(out, gpa, "mtp.fce", fce_members[0..]); - const fch = try qlinear(ck, cfg, "language_model.mtp.fc_hidden"); + const fch = try qlinear(ck, cfg, try std.fmt.allocPrint(a, "{s}.fc_hidden", .{root})); const fch_members = [_]Q{fch}; try putLane(out, gpa, "mtp.fch", fch_members[0..]); - const enorm = try centeredScale(gpa, ck, "language_model.mtp.pre_fc_norm_embedding.weight", report.norms_around_one); + const enorm = try centeredScale(gpa, ck, try std.fmt.allocPrint(a, "{s}.pre_fc_norm_embedding.weight", .{root}), report.norms_around_one); try out.put("mtp.enorm.scale", "F32", &.{cfg.hidden}, enorm); - const hnorm = try centeredScale(gpa, ck, "language_model.mtp.pre_fc_norm_hidden.weight", report.norms_around_one); + const hnorm = try centeredScale(gpa, ck, try std.fmt.allocPrint(a, "{s}.pre_fc_norm_hidden.weight", .{root}), report.norms_around_one); try out.put("mtp.hnorm.scale", "F32", &.{wide}, hnorm); - const head = try qlinear(ck, cfg, "language_model.lm_head"); + const head = try qlinear(text, cfg, "language_model.lm_head"); // the head sits with the text model try putDraftHead(out, gpa, a, "mtp.draft", head, ids); } @@ -592,4 +562,5 @@ pub const compareFile = pio.compareFile; test { _ = @import("pack_test.zig"); + _ = @import("pack_hc_test.zig"); } diff --git a/zig/src/families/flashnext/pack_cache_test.zig b/zig/src/families/flashnext/pack_cache_test.zig index 2855ff8a1f..f291237ea7 100644 --- a/zig/src/families/flashnext/pack_cache_test.zig +++ b/zig/src/families/flashnext/pack_cache_test.zig @@ -295,3 +295,28 @@ test "cache readiness reads bounded metadata rather than copying the weight body var bounded: std.heap.FixedBufferAllocator = .init(&storage); try std.testing.expect(try pio.packReady(bounded.allocator(), io, path, "checkpoint")); } + +test "packs cached or left in the model folder under the format-less identity are rebuilt, never reused" { + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const paths = try Paths.init(tmp); + defer paths.deinit(); + const h = pio.hashBytes("index"); + const shards = [_]pio.ShardId{.{ .name = "model-00001-of-00001.safetensors", .size = 4, .header_sha256 = h }}; + const current = try pio.identityString(a, h, h, &shards); + defer a.free(current); + const old = current[std.mem.indexOfScalar(u8, current, '\n').? + 1 ..]; // the 6/32-only packer's identity text + try std.testing.expect(std.mem.startsWith(u8, old, "index ")); + var stale: Builder = .{ .identity = old }; + const old_dir = try cache.ensure(a, io, paths.model, paths.root, old, &stale); + defer a.free(old_dir); + try tmp.dir.createDir(io, "model/zig-pack", .default_dir); + const legacy = try std.fs.path.join(a, &.{ paths.model, "zig-pack" }); + defer a.free(legacy); + try stale.build(a, io, paths.model, legacy, ""); + var builder: Builder = .{ .identity = current }; + const dir = try cache.ensure(a, io, paths.model, paths.root, current, &builder); + defer a.free(dir); + try std.testing.expectEqual(@as(usize, 1), builder.calls); + try std.testing.expect(!std.mem.eql(u8, old_dir, dir) and !std.mem.eql(u8, legacy, dir)); +} diff --git a/zig/src/families/flashnext/pack_hash_test.zig b/zig/src/families/flashnext/pack_hash_test.zig new file mode 100644 index 0000000000..af41511c25 --- /dev/null +++ b/zig/src/families/flashnext/pack_hash_test.zig @@ -0,0 +1,47 @@ +//! The uniform 6-bit/group-32 packs keep their bytes: one digest over every tensor's name, dtype, shape and payload. +const std = @import("std"); +const pack = @import("pack.zig"); +const fixture = @import("pack_test.zig"); +const st = @import("../../core/safetensors.zig"); + +const gpa = std.testing.allocator; +const io = std.testing.io; + +/// The digest of the synthetic 6/32 checkpoint's three packs' tensors: a uniform pack keeps every byte. +const uniform_digest = "21fa29e16708022181ed842263cc03478cf76de2e213426f94feb05b58d8112e"; + +fn absorb(h: *std.crypto.hash.sha2.Sha256, a: std.mem.Allocator, path: []const u8) !void { + var file = try st.File.open(a, io, path); + defer file.close(io); + const names = try a.dupe([]const u8, file.names.keys()); + std.mem.sort([]const u8, names, {}, struct { + fn less(_: void, x: []const u8, y: []const u8) bool { + return std.mem.lessThan(u8, x, y); + } + }.less); + for (names) |name| { + const e = file.names.get(name).?; + h.update(name); + h.update(&.{ @backingInt(e.dtype), e.rank }); + for (e.shape[0..e.rank]) |d| h.update(std.mem.asBytes(&@as(u64, d))); + h.update(file.map.memory[file.data + e.begin .. file.data + e.end]); + } +} + +test "the uniform 6-bit group-32 packs hash as they did before mixed widths" { + var arena_state: std.heap.ArenaAllocator = .init(gpa); + defer arena_state.deinit(); + const a = arena_state.allocator(); + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try fixture.writeCheckpoint(tmp, a); + try tmp.dir.createDirPath(io, "out"); + _ = try pack.build(a, io, try fixture.tmpPath(a, tmp, "."), try fixture.tmpPath(a, tmp, "out"), null); + var h = std.crypto.hash.sha2.Sha256.init(.{}); + for ([_][]const u8{ "out/pack.safetensors", "out/pack_mlx.safetensors", "out/pack_mtp_mlx.safetensors" }) |name| { + try absorb(&h, a, try std.fs.path.join(a, &.{ try fixture.tmpPath(a, tmp, "."), name })); + } + var digest: [32]u8 = undefined; + h.final(&digest); + try std.testing.expectEqualStrings(uniform_digest, &std.fmt.bytesToHex(digest, .lower)); +} diff --git a/zig/src/families/flashnext/pack_hc_test.zig b/zig/src/families/flashnext/pack_hc_test.zig new file mode 100644 index 0000000000..637a4d3206 --- /dev/null +++ b/zig/src/families/flashnext/pack_hc_test.zig @@ -0,0 +1,72 @@ +//! Hyper-connection packs: a down and inject pair of two widths widens exactly to the wider, as QWeights.of. +const std = @import("std"); +const pack = @import("pack.zig"); +const mlx = @import("../../core/quant/mlx.zig"); +const st = @import("../../core/safetensors.zig"); +const affine = @import("affine.zig"); +const synthetic = @import("pack_widen_test.zig").synthetic; + +fn values(a: std.mem.Allocator, q: pack.Q) ![]f64 { + const out = try a.alloc(f64, q.rows * q.k); + try mlx.reference(.{ .words = q.w, .scales = q.s, .biases = q.b, .bits = q.spec.bits, .group = @intCast(q.spec.group) }, out); + return out; +} + +test "oQ4e's 4/64 down with a 5/64 inject packs as one 5/64 matrix holding the same values" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const down = try synthetic(a, 11, 64, 1024, 4, 64); + const inject = try synthetic(a, 12, 32, 1024, 5, 64); + const both = try pack.hcDown(a, &.{ down, inject }); + try std.testing.expectEqual(@as(u8, 5), both.spec.bits); + try std.testing.expectEqual(@as(usize, 64), both.spec.group); + try std.testing.expectEqual(@as(usize, 96), both.rows); + const got = try values(a, both); + try std.testing.expectEqualSlices(f64, try values(a, down), got[0 .. 64 * 1024]); + try std.testing.expectEqualSlices(f64, try values(a, inject), got[64 * 1024 ..]); + // one width keeps its bytes: the uniform pair is a plain row concatenation + const same = try pack.hcDown(a, &.{ down, try synthetic(a, 13, 32, 1024, 4, 64) }); + try std.testing.expectEqualSlices(u8, down.w.bytes, same.w.bytes[0..down.w.bytes.len]); +} + +test "the up projection's 320 inputs can't be stored in groups of 128, so it packs at its stored width" { + const e = struct { + fn of(dtype: st.DType, rows: usize, cols: usize) st.Entry { + return .{ .dtype = dtype, .rank = 2, .shape = .{ rows, cols, 0, 0, 0 }, .begin = 0, .end = rows * cols * dtype.size() }; + } + }.of; + // 8-bit codes over 320 inputs are 80 words a row; three groups of 128 would need 96 + const at128 = affine.matrix(e(.u32, 64, 80), e(.bf16, 64, 3), e(.bf16, 64, 3), try affine.Spec.init(8, 128)); + try std.testing.expectError(error.InvalidAffineShape, at128); + try std.testing.expectEqual(@as(usize, 320), (try affine.matrix(e(.u32, 64, 80), e(.bf16, 64, 5), e(.bf16, 64, 5), try affine.Spec.init(8, 64))).k); +} + +test "a prompt projection of several widths keeps each member's bytes in order and tells their widths" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var out = pack.Out.init(std.testing.allocator); + defer out.deinit(); + const q = try synthetic(a, 21, 64, 256, 3, 64); + const v = try synthetic(a, 22, 32, 256, 8, 128); + try pack.mlxProjection(&out, std.testing.allocator, a, "L3.att.proj", &.{ q, v }); + const names = out.names.items; + const at = struct { + fn find(n: []const []const u8, name: []const u8) usize { + for (n, 0..) |x, i| if (std.mem.eql(u8, x, name)) return i; + unreachable; + } + }.find; + const mw = out.blobs.items[at(names, "L3.att.proj.mw")]; + try std.testing.expectEqualSlices(u8, q.w.bytes, mw[0..q.w.bytes.len]); + try std.testing.expectEqualSlices(u8, v.w.bytes, mw[q.w.bytes.len..]); + try std.testing.expectEqualSlices(usize, &.{mw.len / 4}, out.shapes.items[at(names, "L3.att.proj.mw")]); + const parts = out.blobs.items[at(names, "L3.att.proj.parts")]; + try std.testing.expectEqualSlices(u8, std.mem.sliceAsBytes(&[_]u32{ 64, 3, 64, 32, 8, 128 }), parts); + // 6/32 members write the pinned 2-D shapes and no table + const six = [_]pack.Q{ try synthetic(a, 23, 32, 256, 6, 32), try synthetic(a, 24, 32, 256, 6, 32) }; + try pack.mlxProjection(&out, std.testing.allocator, a, "L7.att.proj", &six); + try std.testing.expectEqualSlices(usize, &.{ 64, 48 }, out.shapes.items[at(out.names.items, "L7.att.proj.mw")]); + for (out.names.items) |n| try std.testing.expect(!std.mem.eql(u8, n, "L7.att.proj.parts")); +} diff --git a/zig/src/families/flashnext/pack_io.zig b/zig/src/families/flashnext/pack_io.zig index 942bfac500..6dbe6760ed 100644 --- a/zig/src/families/flashnext/pack_io.zig +++ b/zig/src/families/flashnext/pack_io.zig @@ -28,15 +28,22 @@ pub const Out = struct { out.* = undefined; } + /// Takes `bytes` (from `gpa`) even when it fails: then they are freed and every list keeps its length. pub fn put(out: *Out, name: []const u8, dtype: []const u8, shape: []const usize, bytes: anytype) !void { + errdefer out.gpa.free(bytes); const a = out.arena.allocator(); - try out.names.append(out.gpa, try a.dupe(u8, name)); - try out.dtypes.append(out.gpa, dtype); - const shape_copy = try a.alloc(usize, shape.len); - @memcpy(shape_copy, shape); - try out.shapes.append(out.gpa, shape_copy); - try out.blobs.append(out.gpa, std.mem.sliceAsBytes(bytes)); - try out.blob_aligns.append(out.gpa, comptime std.mem.Alignment.fromByteUnits(@alignOf(std.meta.Child(@TypeOf(bytes))))); + const name_copy = try a.dupe(u8, name); + const shape_copy = try a.dupe(usize, shape); + try out.names.ensureUnusedCapacity(out.gpa, 1); + try out.dtypes.ensureUnusedCapacity(out.gpa, 1); + try out.shapes.ensureUnusedCapacity(out.gpa, 1); + try out.blobs.ensureUnusedCapacity(out.gpa, 1); + try out.blob_aligns.ensureUnusedCapacity(out.gpa, 1); + out.names.appendAssumeCapacity(name_copy); + out.dtypes.appendAssumeCapacity(dtype); + out.shapes.appendAssumeCapacity(shape_copy); + out.blobs.appendAssumeCapacity(std.mem.sliceAsBytes(bytes)); + out.blob_aligns.appendAssumeCapacity(comptime std.mem.Alignment.fromByteUnits(@alignOf(std.meta.Child(@TypeOf(bytes))))); } /// Writes the file with `source` recorded under __metadata__ so a reader can refuse a pack from another checkpoint. @@ -102,8 +109,11 @@ pub fn hashBytes(bytes: []const u8) [32]u8 { return h; } -/// The identity's canonical text: the index file's hash, then each shard's name, size and header hash by name. -pub fn identityString(gpa: std.mem.Allocator, index_sha256: ?[32]u8, shards: []const ShardId) ![]u8 { +/// The pack layout's version, first in the identity: 3 has width tables, ragged runs untiled, the config's hash. +pub const pack_format = 3; + +/// The identity's text: pack format, index and config hashes (the config names the widths), each shard's id. +pub fn identityString(gpa: std.mem.Allocator, index_sha256: ?[32]u8, config_sha256: ?[32]u8, shards: []const ShardId) ![]u8 { const sorted = try gpa.dupe(ShardId, shards); defer gpa.free(sorted); std.mem.sort(ShardId, sorted, {}, struct { @@ -113,7 +123,9 @@ pub fn identityString(gpa: std.mem.Allocator, index_sha256: ?[32]u8, shards: []c }.lt); var w: std.Io.Writer.Allocating = .init(gpa); errdefer w.deinit(); + try w.writer.print("flashnext pack {d}\n", .{pack_format}); if (index_sha256) |h| try w.writer.print("index {s}\n", .{&std.fmt.bytesToHex(h, .lower)}) else try w.writer.writeAll("index absent\n"); + if (config_sha256) |h| try w.writer.print("config {s}\n", .{&std.fmt.bytesToHex(h, .lower)}) else try w.writer.writeAll("config absent\n"); for (sorted) |s| try w.writer.print("shard {s} {d} {s}\n", .{ s.name, s.size, &std.fmt.bytesToHex(s.header_sha256, .lower) }); return w.toOwnedSlice(); } @@ -134,8 +146,8 @@ fn headerLen(memory: []const u8) !usize { return hl; } -/// Canonical identity text comes from the index bytes and each shard's mapped bytes, shared by builder and engine. -pub fn identityFromMapped(gpa: std.mem.Allocator, index_bytes: ?[]const u8, shards: []const MappedShard) ![]u8 { +/// The identity from the index and config bytes and each shard's mapped bytes, shared by builder and engine. +pub fn identityFromMapped(gpa: std.mem.Allocator, index_bytes: ?[]const u8, config_bytes: ?[]const u8, shards: []const MappedShard) ![]u8 { var ids: std.ArrayList(ShardId) = .empty; defer ids.deinit(gpa); for (shards) |s| { @@ -145,11 +157,18 @@ pub fn identityFromMapped(gpa: std.mem.Allocator, index_bytes: ?[]const u8, shar const hl = try headerLen(s.bytes[0..size]); try ids.append(gpa, .{ .name = s.name, .size = s.size, .header_sha256 = hashBytes(s.bytes[8..][0..hl]) }); } - return identityString(gpa, if (index_bytes) |b| hashBytes(b) else null, ids.items); + return identityString(gpa, if (index_bytes) |b| hashBytes(b) else null, if (config_bytes) |b| hashBytes(b) else null, ids.items); } -/// A checkpoint identity is the index hash plus each shard size and header hash, and mappings then unmap. +/// A checkpoint identity is the index and config hashes plus each shard size and header hash; mappings then unmap. pub fn sourceIdentity(gpa: std.mem.Allocator, io: Io, model_dir: []const u8) ![]u8 { + const config_path = try std.fs.path.join(gpa, &.{ model_dir, "config.json" }); + defer gpa.free(config_path); + const config = Io.Dir.cwd().readFileAlloc(io, config_path, gpa, .limited(1 << 24)) catch |e| switch (e) { + error.FileNotFound => null, + else => return e, + }; + defer if (config) |c| gpa.free(c); const Open = struct { file: Io.File, map: Io.File.MemoryMap }; var shards: std.ArrayList(MappedShard) = .empty; defer shards.deinit(gpa); @@ -162,14 +181,14 @@ pub fn sourceIdentity(gpa: std.mem.Allocator, io: Io, model_dir: []const u8) ![] const index_path = try std.fs.path.join(gpa, &.{ model_dir, "model.safetensors.index.json" }); defer gpa.free(index_path); var mapped = mapFile(io, index_path) catch |e| switch (e) { - error.FileNotFound => return identityFromMapped(gpa, null, shards.items), + error.FileNotFound => return identityFromMapped(gpa, null, config, shards.items), else => return e, }; defer mapped.map.destroy(io); defer mapped.file.close(io); var parsed = try std.json.parseFromSlice(std.json.Value, gpa, mapped.map.memory, .{}); defer parsed.deinit(); - const weight_map = parsed.value.object.get("weight_map") orelse return identityFromMapped(gpa, null, shards.items); + const weight_map = parsed.value.object.get("weight_map") orelse return identityFromMapped(gpa, null, config, shards.items); var names: std.ArrayList([]const u8) = .empty; defer { for (names.items) |n| gpa.free(n); @@ -191,7 +210,7 @@ pub fn sourceIdentity(gpa: std.mem.Allocator, io: Io, model_dir: []const u8) ![] try maps.append(gpa, .{ .file = shard.file, .map = shard.map }); try shards.append(gpa, .{ .name = name, .size = shard.len, .bytes = shard.map.memory }); } - const out = try identityFromMapped(gpa, mapped.map.memory, shards.items); + const out = try identityFromMapped(gpa, mapped.map.memory, config, shards.items); for (maps.items) |*m| { m.map.destroy(io); m.file.close(io); diff --git a/zig/src/families/flashnext/pack_lane.zig b/zig/src/families/flashnext/pack_lane.zig new file mode 100644 index 0000000000..24203adfc7 --- /dev/null +++ b/zig/src/families/flashnext/pack_lane.zig @@ -0,0 +1,188 @@ +//! Lane packs: words tiled, scale pairs group-major, and one run per bit width when a fused stack mixes widths. +const std = @import("std"); +const st = @import("../../core/safetensors.zig"); +const affine = @import("affine.zig"); +const pio = @import("pack_io.zig"); +const widen = @import("pack_widen.zig"); + +/// Which of a quantized linear's three tensors. +pub const Part = enum { weight, scales, biases }; + +/// A checkpoint quantized linear's three tensors with its verified affine spec. +pub const Q = struct { + w: st.Tensor, + s: st.Tensor, + b: st.Tensor, + rows: usize, + k: usize, + kw: usize, + spec: affine.Spec, + + pub fn part(r: Q, comptime which: Part) []const u8 { + return switch (which) { + .weight => r.w.bytes, + .scales => r.s.bytes, + .biases => r.b.bytes, + }; + } +}; + +/// The lane pipeline's word permutation: rows tiled nt at a time, a group's words innermost (lane_qmm.tile_weight). +pub fn tileWeight(gpa: std.mem.Allocator, flat: []const u8, rows: usize, kw: usize, nt: usize, w: usize) ![]u8 { + const words = std.mem.bytesAsSlice(u32, flat); + const out = try gpa.alloc(u8, flat.len); + errdefer gpa.free(out); + const out_words = std.mem.bytesAsSlice(u32, out); + const groups = kw / w; + for (0..rows / nt) |t| { + for (0..groups) |g| { + for (0..nt) |j| { + for (0..w) |c| { + out_words[((t * groups + g) * nt + j) * w + c] = words[(t * nt + j) * kw + g * w + c]; + } + } + } + } + return out; +} + +/// (K/GS, N, 2) bf16 (scale, bias) pairs, group-major, over a row-concatenated stack (lane_qmm.pack_scales). +pub fn packScales(gpa: std.mem.Allocator, members: []const Q) ![]u8 { + var rows: usize = 0; + for (members) |m| rows += m.rows; + const kg = members[0].s.dim(1); + const out = try gpa.alloc(u8, kg * rows * 4); + errdefer gpa.free(out); + const out16 = std.mem.bytesAsSlice(u16, out); + var at: usize = 0; + for (members) |m| { + const s16 = std.mem.bytesAsSlice(u16, m.part(.scales)); + const b16 = std.mem.bytesAsSlice(u16, m.part(.biases)); + for (0..kg) |g| { + for (0..m.rows) |j| { + out16[(g * rows + at + j) * 2] = s16[j * kg + g]; + out16[(g * rows + at + j) * 2 + 1] = b16[j * kg + g]; + } + } + at += m.rows; + } + return out; +} + +/// Row-concatenated members in the checkpoint's own MLX layout: a pure byte move per member. +pub fn concatPart(gpa: std.mem.Allocator, members: []const Q, comptime part: Part) ![]u8 { + var n: usize = 0; + for (members) |m| n += m.part(part).len; + const out = try gpa.alloc(u8, n); + errdefer gpa.free(out); + var at: usize = 0; + for (members) |m| { + @memcpy(out[at..][0..m.part(part).len], m.part(part)); + at += m.part(part).len; + } + return out; +} + +/// One bit width's members, in member order, at their smallest group: one product, as decode.py's `_stacked`. +pub const Run = struct { bits: u8, group: usize, members: u16, rows: usize }; + +/// Runs by bit width in first-seen order, adjacent or not, so split_k(N, K) sees 0.6.6's N; group 128 reads as 64. +pub fn runsOf(specs: []const affine.Spec, rows: []const usize, out: []Run) ![]Run { + if (specs.len == 0 or specs.len != rows.len) return error.EmptyLaneStack; + if (specs.len > 16) return error.TooManyLaneMembers; + var n: usize = 0; + for (specs, rows, 0..) |s, r, i| { + const bit = @as(u16, 1) << @intCast(i); + for (out[0..n]) |*run| { + if (run.bits != s.bits) continue; + run.members |= bit; + run.rows += r; + run.group = @min(run.group, s.group); + break; + } else { + if (n == out.len) return error.TooManyLaneRuns; + out[n] = .{ .bits = s.bits, .group = s.group, .members = bit, .rows = r }; + n += 1; + } + } + for (out[0..n]) |*run| run.group = @min(run.group, 64); + return out[0..n]; +} + +/// Whether a run of `rows` columns is tiled: 32 at a time when they divide, else MLX's layout (decode.py's `nt`). +pub fn tiled(rows: usize) bool { + return rows % 32 == 0; +} + +/// One width's `.wq` (tiled, or MLX's layout when ragged) and `.sbt` group-major pairs (decode.py's lane pipeline). +fn putRun(out: *pio.Out, gpa: std.mem.Allocator, name: []const u8, members: []const Q) !void { + var rows: usize = 0; + for (members) |m| rows += m.rows; + const spec = members[0].spec; + if (rows % 4 != 0) return error.LaneRowsNotFours; // lane_matmul's column pairs and scale pairs + const w = spec.group * spec.bits / 32; + const a = out.arena.allocator(); + const wq_name = try std.fmt.allocPrint(a, "{s}.wq", .{name}); + const sbt_name = try std.fmt.allocPrint(a, "{s}.sbt", .{name}); + const flat = try concatPart(gpa, members, .weight); + defer gpa.free(flat); + // put owns each blob from the call on, failing or not + try out.put(wq_name, "U32", &.{ rows, members[0].kw }, if (tiled(rows)) try tileWeight(gpa, flat, rows, members[0].kw, 32, w) else try gpa.dupe(u8, flat)); + try out.put(sbt_name, "BF16", &.{ members[0].s.dim(1), rows, 2 }, try packScales(gpa, members)); +} + +/// A lane projection; any width but 6/32 adds a `.runs` table, and several widths `.rN` per run and `.members`. +pub fn putLane(out: *pio.Out, gpa: std.mem.Allocator, name: []const u8, given: []const Q) !void { + putStack(out, gpa, name, given) catch |e| { + std.log.warn("lane {s}: {s}", .{ name, @errorName(e) }); // the pack refusal names its lane + return e; + }; +} + +fn putStack(out: *pio.Out, gpa: std.mem.Allocator, name: []const u8, given: []const Q) !void { + var specs: [16]affine.Spec = undefined; + var rows: [16]usize = undefined; + if (given.len > specs.len) return error.TooManyLaneRuns; + for (given, 0..) |m, i| { + specs[i] = m.spec; + rows[i] = m.rows; + } + var buf: [8]Run = undefined; + const runs = try runsOf(specs[0..given.len], rows[0..given.len], &buf); + const a = out.arena.allocator(); + var place: [16][3]u32 = undefined; // run, row within the run, row within the fused output + var fused: u32 = 0; + for (runs, 0..) |r, ri| { + var storage: [16]Q = undefined; + var count: usize = 0; + var at: u32 = 0; + for (given, 0..) |m, i| { + if ((r.members >> @intCast(i)) & 1 == 0) continue; + storage[count] = try widen.regroup(a, m, r.group); + count += 1; + place[i] = .{ @intCast(ri), at, 0 }; + at += @intCast(m.rows); + } + const run_name = if (runs.len == 1) name else try std.fmt.allocPrint(a, "{s}.r{d}", .{ name, ri }); + try putRun(out, gpa, run_name, storage[0..count]); + } + // the uniform packs keep their bytes + if (runs.len == 1 and runs[0].bits == 6 and runs[0].group == 32 and tiled(runs[0].rows)) return; + const runs_name = try std.fmt.allocPrint(a, "{s}.runs", .{name}); + const table = try gpa.alloc(u32, runs.len * 3); + for (runs, 0..) |r, i| table[i * 3 ..][0..3].* = .{ r.bits, @intCast(r.group), @intCast(r.rows) }; + try out.put(runs_name, "U32", &.{ runs.len, 3 }, table); + if (runs.len == 1) return; + for (given, 0..) |m, i| { + place[i][2] = fused; + fused += @intCast(m.rows); + } + const members_name = try std.fmt.allocPrint(a, "{s}.members", .{name}); + const members = try gpa.alloc(u32, given.len * 3); + for (place[0..given.len], 0..) |p, i| members[i * 3 ..][0..3].* = p; + try out.put(members_name, "U32", &.{ given.len, 3 }, members); +} + +test { + _ = @import("pack_lane_test.zig"); +} diff --git a/zig/src/families/flashnext/pack_lane_test.zig b/zig/src/families/flashnext/pack_lane_test.zig new file mode 100644 index 0000000000..42a14b6001 --- /dev/null +++ b/zig/src/families/flashnext/pack_lane_test.zig @@ -0,0 +1,136 @@ +//! A fused stack packs one product per bit width, members of equal bits together wherever they sit (`_stacked`). +const std = @import("std"); +const affine = @import("affine.zig"); +const lane = @import("pack_lane.zig"); +const pio = @import("pack_io.zig"); +const synthetic = @import("pack_widen_test.zig").synthetic; + +fn spec(bits: u8, group: usize) affine.Spec { + return affine.Spec.init(bits, group) catch unreachable; +} + +const att = [_]usize{ 12288, 512, 512, 640 }; // row counts of q, k, v and index_qk +const gdn = [_]usize{ 10240, 6144, 48, 48 }; // row counts of in_proj_qkv, z, b and a + +fn expectRuns(specs: []const affine.Spec, rows: []const usize, want: []const lane.Run) !void { + var buf: [8]lane.Run = undefined; + const runs = try lane.runsOf(specs, rows[0..specs.len], &buf); + try std.testing.expectEqual(want.len, runs.len); + for (want, runs) |w, r| try std.testing.expectEqual(w, r); +} + +test "members of equal bits form one run at their smallest group, in first-seen order, adjacent or not" { + // oQ2/oQ3/oQ4e attention: an 8-bit v between 6-bit q, k and index_qk; 0.6.6 runs q, k and index_qk as one product + try expectRuns(&.{ spec(6, 64), spec(6, 64), spec(8, 64), spec(6, 64) }, &att, &.{ + .{ .bits = 6, .group = 64, .members = 0b1011, .rows = 12288 + 512 + 640 }, + .{ .bits = 8, .group = 64, .members = 0b0100, .rows = 512 }, + }); + // a 6/128 v among 6/64 members is regrouped into the one 6/64 product + try expectRuns(&.{ spec(6, 64), spec(6, 64), spec(6, 128), spec(6, 64) }, &att, &.{ + .{ .bits = 6, .group = 64, .members = 0b1111, .rows = 13952 }, + }); + // oQ4e layer 21's GDN inputs: qkv 4/64, z and b 5/128, a 6/64; the 5-bit run is read in groups of 64 on the lane + try expectRuns(&.{ spec(4, 64), spec(5, 128), spec(5, 128), spec(6, 64) }, &gdn, &.{ + .{ .bits = 4, .group = 64, .members = 0b0001, .rows = 10240 }, + .{ .bits = 5, .group = 64, .members = 0b0110, .rows = 6192 }, + .{ .bits = 6, .group = 64, .members = 0b1000, .rows = 48 }, + }); + // equal bits in groups of 32 and 64 share the group of 32, so the run's group sums are 0.6.6's + try expectRuns(&.{ spec(4, 64), spec(4, 32) }, &att, &.{.{ .bits = 4, .group = 32, .members = 0b11, .rows = 12800 }}); + try expectRuns(&.{ spec(6, 32), spec(6, 32), spec(6, 32), spec(6, 32) }, &att, &.{ + .{ .bits = 6, .group = 32, .members = 0b1111, .rows = 13952 }, + }); +} + +fn find(out: *const pio.Out, name: []const u8) ?usize { + for (out.names.items, 0..) |n, i| if (std.mem.eql(u8, n, name)) return i; + return null; +} + +test "a stack of several widths writes a pack per run with its rows' places, and a regrouped one keeps the plain names" { + const gpa = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var out = pio.Out.init(gpa); + defer out.deinit(); + const q = try synthetic(a, 1, 32, 256, 6, 64); + const v = try synthetic(a, 2, 32, 256, 8, 64); + const idx = try synthetic(a, 3, 32, 256, 6, 64); + try lane.putLane(&out, gpa, "L1.att.proj", &.{ q, v, idx }); + try std.testing.expect(find(&out, "L1.att.proj.wq") == null); + try std.testing.expectEqualSlices(usize, &.{ 64, 48 }, out.shapes.items[find(&out, "L1.att.proj.r0.wq").?]); + try std.testing.expectEqualSlices(usize, &.{ 32, 64 }, out.shapes.items[find(&out, "L1.att.proj.r1.wq").?]); + // index_qk's first word follows q's 32 rows in the 6-bit run: the second tile's first word + const r0 = out.blobs.items[find(&out, "L1.att.proj.r0.wq").?]; + try std.testing.expectEqualSlices(u8, idx.w.bytes[0..4], r0[32 * 48 * 4 ..][0..4]); + const runs = out.blobs.items[find(&out, "L1.att.proj.runs").?]; // (bits, group, rows) per run + try std.testing.expectEqualSlices(u8, std.mem.sliceAsBytes(&[_]u32{ 6, 64, 64, 8, 64, 32 }), runs); + const places = out.blobs.items[find(&out, "L1.att.proj.members").?]; // (run, row in run, row in fused output) + try std.testing.expectEqualSlices(u8, std.mem.sliceAsBytes(&[_]u32{ 0, 0, 0, 1, 0, 32, 0, 32, 64 }), places); + const one = [_]lane.Q{ try synthetic(a, 4, 32, 256, 6, 64), try synthetic(a, 5, 32, 256, 6, 128) }; + try lane.putLane(&out, gpa, "L3.att.proj", &one); + try std.testing.expectEqualSlices(usize, &.{ 4, 64, 2 }, out.shapes.items[find(&out, "L3.att.proj.sbt").?]); + const regrouped = out.blobs.items[find(&out, "L3.att.proj.runs").?]; // one run, its width told + try std.testing.expectEqualSlices(u8, std.mem.sliceAsBytes(&[_]u32{ 6, 64, 64 }), regrouped); + try std.testing.expect(find(&out, "L3.att.proj.members") == null); + const plain = [_]lane.Q{ try synthetic(a, 6, 32, 256, 6, 32), try synthetic(a, 7, 32, 256, 6, 32) }; + try lane.putLane(&out, gpa, "L7.att.proj", &plain); + try std.testing.expect(find(&out, "L7.att.proj.runs") == null); // 6/32 writes no table: its bytes are pinned + const four = [_]lane.Q{try synthetic(a, 8, 64, 256, 4, 64)}; + try lane.putLane(&out, gpa, "L9.att.o", &four); + try std.testing.expectEqualSlices(usize, &.{ 64, 32 }, out.shapes.items[find(&out, "L9.att.o.wq").?]); + // 64 rows of 4 bits tile 32 columns at a time, as the lane launches read them: the second tile follows the first + const wq = out.blobs.items[find(&out, "L9.att.o.wq").?]; + try std.testing.expectEqualSlices(u8, four[0].w.bytes[32 * 32 * 4 ..][0..4], wq[4 * 32 * 8 * 4 ..][0..4]); +} + +test "runs whose columns aren't a multiple of 32 keep MLX's layout untiled, as decode.py's nt 0, at the GDN rows" { + const gpa = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var out = pio.Out.init(gpa); + defer out.deinit(); + const widths = [_][2]usize{ .{ 4, 64 }, .{ 5, 128 }, .{ 5, 128 }, .{ 6, 64 } }; + var ms: [4]lane.Q = undefined; + for (&ms, gdn, widths, 0..) |*m, rows, w, i| m.* = try synthetic(a, 20 + i, rows, 256, @intCast(w[0]), w[1]); + try lane.putLane(&out, gpa, "L21.gdn.in", &ms); + const runs = out.blobs.items[find(&out, "L21.gdn.in.runs").?]; + try std.testing.expectEqualSlices(u8, std.mem.sliceAsBytes(&[_]u32{ 4, 64, 10240, 5, 64, 6192, 6, 64, 48 }), runs); + try std.testing.expect(lane.tiled(10240) and !lane.tiled(6192) and !lane.tiled(48)); + // the ragged 5-bit run is z then b row by row, and the 48-row a is its own bytes + const r1 = out.blobs.items[find(&out, "L21.gdn.in.r1.wq").?]; + try std.testing.expectEqualSlices(u8, ms[1].w.bytes, r1[0..ms[1].w.bytes.len]); + try std.testing.expectEqualSlices(u8, ms[2].w.bytes, r1[ms[1].w.bytes.len..]); + try std.testing.expectEqualSlices(u8, ms[3].w.bytes, out.blobs.items[find(&out, "L21.gdn.in.r2.wq").?]); + const places = out.blobs.items[find(&out, "L21.gdn.in.members").?]; + try std.testing.expectEqualSlices(u8, std.mem.sliceAsBytes(&[_]u32{ 0, 0, 0, 1, 0, 10240, 1, 6144, 16384, 2, 0, 16432 }), places); + // a 6/32 lane whose rows don't tile carries its table, so it is never read as a tiled uniform lane + const odd = [_]lane.Q{try synthetic(a, 30, 48, 256, 6, 32)}; + try lane.putLane(&out, gpa, "L0.gdn.out", &odd); + try std.testing.expect(find(&out, "L0.gdn.out.runs") != null); + // columns not in fours have no lane layout: refused by name + const bad = [_]lane.Q{try synthetic(a, 31, 50, 256, 6, 32)}; + try std.testing.expectError(error.LaneRowsNotFours, lane.putLane(&out, gpa, "L0.att.o", &bad)); +} + +fn packMixed(gpa: std.mem.Allocator, stack: []const lane.Q, proj: []const lane.Q) !void { + var out = pio.Out.init(gpa); + defer out.deinit(); + try lane.putLane(&out, gpa, "L21.gdn.in", stack); + try @import("pack.zig").mlxProjection(&out, gpa, out.arena.allocator(), "L21.gdn.in", proj); +} + +test "a lane stack and a prompt projection of several widths free every buffer when any allocation fails" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + // oQ4e's layer-21 GDN input in miniature: 4/64, 5/128 twice, 6/64, with a ragged 48-row 6-bit run + const l21 = [_]lane.Q{ try synthetic(a, 31, 64, 256, 4, 64), try synthetic(a, 32, 64, 256, 5, 128), try synthetic(a, 33, 48, 256, 5, 128), try synthetic(a, 34, 48, 256, 6, 64) }; + // a fixed buffer under the failing allocator: every run grows and resizes the same way at every optimize mode + const backing = try std.testing.allocator.alloc(u8, 64 << 20); + defer std.testing.allocator.free(backing); + var fixed: std.heap.FixedBufferAllocator = .init(backing); + try std.testing.checkAllAllocationFailures(fixed.allocator(), packMixed, .{ &l21, &l21 }); +} diff --git a/zig/src/families/flashnext/pack_raw.zig b/zig/src/families/flashnext/pack_raw.zig new file mode 100644 index 0000000000..d2328f0868 --- /dev/null +++ b/zig/src/families/flashnext/pack_raw.zig @@ -0,0 +1,42 @@ +//! Tensors the engine reads raw (experts, embedding, n-gram tables) checked at pack build against declared widths. +const std = @import("std"); +const st = @import("../../core/safetensors.zig"); +const ckpt = @import("../../core/checkpoint.zig"); +const config_mod = @import("config.zig"); +const Widths = @import("widths.zig").Widths; + +fn entry(t: st.Tensor) st.Entry { + return .{ .dtype = t.dtype, .rank = t.rank, .shape = t.shape, .begin = 0, .end = t.bytes.len }; +} + +/// Module `name`'s stored affine triple against the width declared for `key` (unquantized: admission's to judge). +pub fn check(ck: *ckpt.Checkpoint, w: *const Widths, key: []const u8, name: []const u8) !void { + var buf: [3][256]u8 = undefined; + var e: [3]st.Entry = undefined; + for ([_][]const u8{ "weight", "scales", "biases" }) |p| if (!ck.has(try std.fmt.bufPrint(&buf[0], "{s}.{s}", .{ name, p }))) return; + for ([_][]const u8{ "weight", "scales", "biases" }, 0..) |p, i| e[i] = entry(try ck.get(try std.fmt.bufPrint(&buf[i], "{s}.{s}", .{ name, p }))); + _ = w.shaped(key, e[0], e[1], e[2]) catch |err| { + std.log.warn("{s}: not stored at its declared width: {s}", .{ name, @errorName(err) }); + return err; + }; +} + +const projs = [_][]const u8{ "switch_mlp.gate_proj", "switch_mlp.up_proj", "shared_expert.gate_proj", "shared_expert.up_proj", "switch_mlp.down_proj", "shared_expert.down_proj" }; + +/// Every raw-read module: each layer's and the MTP head's experts, the embedding and the n-gram shards. +pub fn checkAll(ck: *ckpt.Checkpoint, ckm: *ckpt.Checkpoint, w: *const Widths, cfg: *const config_mod.Config, mtp_root: []const u8, spelling: []const u8) !void { + var buf: [256]u8 = undefined; + var key: [256]u8 = undefined; + for (0..cfg.layers) |i| for (projs) |p| { + const name = try std.fmt.bufPrint(&buf, "language_model.model.layers.{d}.mlp.{s}", .{ i, p }); + try check(ck, w, name, name); + }; + if (cfg.mtp.layers > 0) for (projs) |p| { + try check(ckm, w, try std.fmt.bufPrint(&key, "language_model.mtp.layers.0.mlp.{s}", .{p}), try std.fmt.bufPrint(&buf, "{s}.layers.0.mlp.{s}", .{ mtp_root, p })); + }; + try check(ck, w, "language_model.model.embed_tokens", "language_model.model.embed_tokens"); + for (0..cfg.layers) |i| if (cfg.ple[i]) for (0..cfg.ngram_shards) |k| { + const name = try std.fmt.bufPrint(&buf, "language_model.model.layers.{d}.ple.ple_embedding.ngram_embedding.{s}{d}", .{ i, spelling, k }); + try check(ck, w, name, name); + }; +} diff --git a/zig/src/families/flashnext/pack_source_test.zig b/zig/src/families/flashnext/pack_source_test.zig index e184b47e87..642e0ff667 100644 --- a/zig/src/families/flashnext/pack_source_test.zig +++ b/zig/src/families/flashnext/pack_source_test.zig @@ -109,9 +109,10 @@ test "the io and mapped identity paths produce the same text" { const bytes = try tmp.dir.readFileAlloc(io, kv.value_ptr.string, a, .limited(1 << 30)); try shards.append(a, .{ .name = kv.value_ptr.string, .size = bytes.len, .bytes = bytes }); } - const through_maps = try pio.identityFromMapped(a, index_bytes, shards.items); + const config_bytes = try tmp.dir.readFileAlloc(io, "config.json", a, .limited(1 << 24)); + const through_maps = try pio.identityFromMapped(a, index_bytes, config_bytes, shards.items); try std.testing.expectEqualStrings(through_io, through_maps); - try std.testing.expect(std.mem.startsWith(u8, through_maps, "index ")); + try std.testing.expect(std.mem.startsWith(u8, through_maps, "flashnext pack 3\nindex ")); } test "the identity string does not depend on the shard list's order" { @@ -123,9 +124,9 @@ test "the identity string does not depend on the shard list's order" { .{ .name = "model-00002.safetensors", .size = 20, .header_sha256 = h2 }, }; const two = [_]pio.ShardId{ one[1], one[0] }; - const s1 = try pio.identityString(a, h1, one[0..]); + const s1 = try pio.identityString(a, h1, h2, one[0..]); defer a.free(s1); - const s2 = try pio.identityString(a, h1, two[0..]); + const s2 = try pio.identityString(a, h1, h2, two[0..]); defer a.free(s2); try std.testing.expectEqualStrings(s1, s2); try std.testing.expect(std.mem.indexOf(u8, s1, "shard model-00001.safetensors 10 ") != null); @@ -165,3 +166,19 @@ test "the no-dump pack cache is ready only complete and matching, resolved from const changed = try pio.sourceIdentity(a, io, model_dir); try std.testing.expect(!try pio.packsReady(a, io, cache_dir, changed)); } + +test "a config-only change (the widths the pack is cut at) gives a new identity, so a stale pack is never reused" { + var arena_state: std.heap.ArenaAllocator = .init(gpa); + defer arena_state.deinit(); + const a = arena_state.allocator(); + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try fixture.writeCheckpoint(tmp, a); + const model_dir = try fixture.tmpPath(a, tmp, "."); + const before = try pio.sourceIdentity(a, io, model_dir); + try std.testing.expect(std.mem.indexOf(u8, before, "\nconfig ") != null); + const config = try tmp.dir.readFileAlloc(io, "config.json", a, .limited(1 << 24)); + try tmp.dir.writeFile(io, .{ .sub_path = "config.json", .data = try std.mem.concat(a, u8, &.{ config, " " }) }); + const after = try pio.sourceIdentity(a, io, model_dir); + try std.testing.expect(!std.mem.eql(u8, before, after)); +} diff --git a/zig/src/families/flashnext/pack_test.zig b/zig/src/families/flashnext/pack_test.zig index 2cb19d3baf..66db631c95 100644 --- a/zig/src/families/flashnext/pack_test.zig +++ b/zig/src/families/flashnext/pack_test.zig @@ -162,6 +162,7 @@ pub fn writeCheckpoint(tmp: std.testing.TmpDir, a: std.mem.Allocator) !void { // the vision tower's rank-five entry: the builder never reads it, and it must not stop the load try list.append(a, .{ .name = "model.visual.patch_embed.proj.weight", .dtype = "BF16", .shape = &.{ 1, 1, 1, 1, 1 }, .bytes = &.{ 0, 0 } }); try qspecs(&list, a, "language_model.lm_head", 128, 500_000, 0x3C00, 0x3B00); + try qspecs(&list, a, "language_model.model.embed_tokens", 128, 600_000, 0x3C00, 0x3B00); // read raw by the engine const mtp = "language_model.mtp"; try hcSpecs(&list, a, try std.fmt.allocPrint(a, "{s}.layers.0.attn_hyper_connection", .{mtp}), 13_000_000, true, 0x3C00, 0x3400); try hcSpecs(&list, a, try std.fmt.allocPrint(a, "{s}.layers.0.mlp_hyper_connection", .{mtp}), 13_500_000, true, 0x3C00, 0x3400); @@ -523,3 +524,35 @@ test "the draft vocabulary builds the mtp decode tensors and gathers the head's try expectTensor(&file, "mtp.hnorm.scale", .f32, &.{32}); try expectTensor(&file, "mtp.fce.wq", .u32, &.{ 32, 6 }); } + +test { + _ = @import("pack_hash_test.zig"); +} + +test "a pack never reads a tensor at a width the config leaves half-declared, though its shapes would fit" { + var arena_state: std.heap.ArenaAllocator = .init(gpa); + defer arena_state.deinit(); + const a = arena_state.allocator(); + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try writeCheckpoint(tmp, a); + // stored 6/32; a 3-bit group-64 reading of the same words and tables fits every shape check + const declared = "\"quantization\": {\"mtp.fc_embedding\": {\"bits\": 3}, "; + try tmp.dir.writeFile(io, .{ .sub_path = "config.json", .data = try std.mem.replaceOwned(u8, a, config_json, "\"quantization\": {", declared) }); + try tmp.dir.createDirPath(io, "out"); + try std.testing.expectError(error.UnsupportedQuantization, pack.build(a, io, try tmpPath(a, tmp, "."), try tmpPath(a, tmp, "out"), null)); +} + +test "a directly read tensor (expert, embedding, n-gram table) whose shapes contradict its declared width is refused" { + var arena_state: std.heap.ArenaAllocator = .init(gpa); + defer arena_state.deinit(); + const a = arena_state.allocator(); + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try writeCheckpoint(tmp, a); + // stored 6/32; the embedding declared 8/64 would be read by the lookup kernel at the wrong width + const declared = "\"quantization\": {\"language_model.model.embed_tokens\": {\"bits\": 8, \"group_size\": 64}, "; + try tmp.dir.writeFile(io, .{ .sub_path = "config.json", .data = try std.mem.replaceOwned(u8, a, config_json, "\"quantization\": {", declared) }); + try tmp.dir.createDirPath(io, "out"); + try std.testing.expectError(error.InvalidAffineShape, pack.build(a, io, try tmpPath(a, tmp, "."), try tmpPath(a, tmp, "out"), null)); +} diff --git a/zig/src/families/flashnext/pack_widen.zig b/zig/src/families/flashnext/pack_widen.zig new file mode 100644 index 0000000000..f9a6ddfcc6 --- /dev/null +++ b/zig/src/families/flashnext/pack_widen.zig @@ -0,0 +1,28 @@ +//! Exact regrouping for lane packs: core/quant widen at the same bits, a group read as smaller ones. +const std = @import("std"); +const quant_mlx = @import("../../core/quant/mlx.zig"); +const quant_widen = @import("../../core/quant/widen.zig"); +const affine = @import("affine.zig"); +const Q = @import("pack_lane.zig").Q; + +/// A pack linear as core/quant sees it. +pub fn host(q: Q) quant_mlx.Host { + return .{ .words = q.w, .scales = q.s, .biases = q.b, .bits = q.spec.bits, .group = @intCast(q.spec.group) }; +} + +/// A core/quant linear back as a pack linear, its K from the tables. +pub fn fromHost(h: quant_mlx.Host) !Q { + const k = h.scales.dim(1) * h.group; + return .{ .w = h.words, .s = h.scales, .b = h.biases, .rows = h.words.dim(0), .k = k, .kw = h.words.dim(1), .spec = try affine.Spec.init(h.bits, h.group) }; +} + +/// `q` in groups of `group`, which divides its own: codes kept, each scale and bias repeated (`QWeights.widened`). +pub fn regroup(a: std.mem.Allocator, q: Q, group: usize) !Q { + if (group == q.spec.group) return q; + if (group > q.spec.group or q.spec.group % group != 0) return error.InexactRegroup; + return fromHost(try quant_widen.widen(a, host(q), .{ .bits = q.spec.bits, .group = @intCast(group) })); +} + +test { + _ = @import("pack_widen_test.zig"); +} diff --git a/zig/src/families/flashnext/pack_widen_test.zig b/zig/src/families/flashnext/pack_widen_test.zig new file mode 100644 index 0000000000..ca8689450e --- /dev/null +++ b/zig/src/families/flashnext/pack_widen_test.zig @@ -0,0 +1,67 @@ +//! A group read as smaller ones dequantizes to the same float64 bits at every width, and the codes keep their bytes. +const std = @import("std"); +const st = @import("../../core/safetensors.zig"); +const mlx = @import("../../core/quant/mlx.zig"); +const affine = @import("affine.zig"); +const lane = @import("pack_lane.zig"); +const widen = @import("pack_widen.zig"); + +/// A [rows, k] matrix at `bits`/`group` with seeded codes and finite bf16 scales and biases, owned by `a`. +pub fn synthetic(a: std.mem.Allocator, seed: u64, rows: usize, k: usize, bits: u8, group: usize) !lane.Q { + var prng = std.Random.DefaultPrng.init(seed); + const rand = prng.random(); + const kw = k * bits / 32; + const w = try a.alloc(u8, rows * kw * 4); + rand.bytes(w); + const tables = [2][]u8{ try a.alloc(u8, rows * (k / group) * 2), try a.alloc(u8, rows * (k / group) * 2) }; + for (tables) |t| for (std.mem.bytesAsSlice(u16, t)) |*v| { + const f: f32 = (rand.float(f32) - 0.5) * 4; + v.* = @intCast(@as(u32, @bitCast(f)) >> 16); + }; + var shape_w: [st.max_rank]usize = @splat(0); + shape_w[0] = rows; + shape_w[1] = kw; + var shape_s = shape_w; + shape_s[1] = k / group; + return .{ + .w = .{ .dtype = .u32, .rank = 2, .shape = shape_w, .bytes = w }, + .s = .{ .dtype = .bf16, .rank = 2, .shape = shape_s, .bytes = tables[0] }, + .b = .{ .dtype = .bf16, .rank = 2, .shape = shape_s, .bytes = tables[1] }, + .rows = rows, + .k = k, + .kw = kw, + .spec = try affine.Spec.init(bits, group), + }; +} + +fn reference(a: std.mem.Allocator, q: lane.Q) ![]f64 { + const out = try a.alloc(f64, q.rows * q.k); + try mlx.reference(.{ .words = q.w, .scales = q.s, .biases = q.b, .bits = q.spec.bits, .group = @intCast(q.spec.group) }, out); + return out; +} + +test "a group read as smaller ones is an exact identity in float64 at every width" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + for ([_][2]usize{ .{ 128, 64 }, .{ 128, 32 }, .{ 64, 32 } }) |pair| for ([_]u8{ 2, 3, 4, 5, 6, 8 }, 0..) |bits, i| { + const q = try synthetic(a, 0x5eed + i + pair[1], 4, 512, bits, pair[0]); + const h = try widen.regroup(a, q, pair[1]); + try std.testing.expectEqual(pair[1], h.spec.group); + try std.testing.expectEqual(bits, h.spec.bits); + try std.testing.expectEqual(512 / pair[1], h.s.dim(1)); + try std.testing.expectEqualSlices(u8, q.w.bytes, h.w.bytes); + const want = try reference(a, q); + const got = try reference(a, h); + try std.testing.expectEqualSlices(u64, std.mem.bytesAsSlice(u64, std.mem.sliceAsBytes(want)), std.mem.bytesAsSlice(u64, std.mem.sliceAsBytes(got))); + }; +} + +test "a group is never read as a larger one, and its own group passes through untouched" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const q = try synthetic(a, 7, 2, 128, 6, 32); + try std.testing.expectEqual(q.s.bytes.ptr, (try widen.regroup(a, q, 32)).s.bytes.ptr); + try std.testing.expectError(error.InexactRegroup, widen.regroup(a, q, 64)); +} diff --git a/zig/src/families/flashnext/ple_scale.zig b/zig/src/families/flashnext/ple_scale.zig new file mode 100644 index 0000000000..3182c23fcd --- /dev/null +++ b/zig/src/families/flashnext/ple_scale.zig @@ -0,0 +1,91 @@ +//! The n-gram table's weight_scale on the looked-up rows, as oMLX's runtime applies it: bf16(row * scale). +const std = @import("std"); +const mtl = @import("metal"); +const replay = @import("replay.zig"); + +const Buf = replay.Buf; + +const source = + \\#include + \\using namespace metal; + \\kernel void tf_ple_scale(device bfloat* x [[buffer(0)]], constant float& s [[buffer(1)]], + \\ constant uint& n [[buffer(2)]], uint i [[thread_position_in_grid]]) { + \\ if (i < n) x[i] = bfloat(float(x[i]) * s); + \\} +; + +/// A scale other than 1 and its kernel, compiled at load. +pub const Scale = struct { + pipe: mtl.Pipeline, + value: f32, + + /// Multiplies `n` bf16 values at `x` in place, then a barrier unless the encoder is serial. + pub fn apply(s: Scale, enc: mtl.ComputeEncoder, x: Buf, n: usize, serial: bool) void { + const count: u32 = @intCast(n); + enc.setPipeline(s.pipe); + enc.setBuffer(x.b, x.off, 0); + enc.setBytes(std.mem.asBytes(&s.value), 1); + enc.setBytes(std.mem.asBytes(&count), 2); + enc.dispatchThreads(mtl.Size.of(n, 1, 1), mtl.Size.of(256, 1, 1)); + if (!serial) enc.barrier(); + } +}; + +/// The checkpoint's n-gram weight_scale as one bf16 or f32 value; null when there is none or it is 1. +pub fn of(r: *replay.Run, map: std.json.ObjectMap) !?Scale { + const name = for (map.keys()) |k| { + if (std.mem.endsWith(u8, k, ".ple.ple_embedding.ngram_embedding.weight_scale")) break k; + } else return null; + const e = try r.entry(name); + const b = try r.load(name); + defer b.b.deinit(); + const bytes = b.b.contents()[0..e.len]; + const value: f32 = switch (e.len) { + 2 => @bitCast(@as(u32, std.mem.readInt(u16, bytes[0..2], .little)) << 16), + 4 => @bitCast(std.mem.readInt(u32, bytes[0..4], .little)), + else => return error.NgramTableScale, + }; + if (!std.math.isFinite(value) or value == 0) return error.NgramTableScale; + if (value == 1) return null; + return try compile(r.device, value); +} + +/// The kernel for `value`; a compile failure refuses the load. +pub fn compile(device: mtl.Device, value: f32) !Scale { + const lib = try mtl.Library.fromSource(device, source, mtl.CompileOptions.mlx()); + return .{ .pipe = try mtl.Pipeline.init(device, lib, "tf_ple_scale", false), .value = value }; +} + +test "the scale rounds each row value once to bf16 after the product, as bf16 row * bf16 scale does" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + var r = replay.Run{ .arena = arena.allocator(), .device = device, .queue = try device.queue() }; + const scale: f32 = @bitCast(@as(u32, 0x3951) << 16); // oQ4e's weight_scale + const s = try compile(device, scale); + const n = 4099; + const buf: Buf = .{ .b = try r.buffer(n * 2) }; + const x = std.mem.bytesAsSlice(u16, buf.b.contents()[0 .. n * 2]); + var rng = std.Random.DefaultPrng.init(7); + for (x) |*v| { // normal bf16 values of both signs whose products stay normal + const exp: u16 = 32 + rng.random().uintLessThan(u16, 96); + v.* = (rng.random().int(u16) & 0x807f) | exp << 7; + } + const want = try std.testing.allocator.alloc(u16, n); + defer std.testing.allocator.free(want); + for (x, want) |v, *w| w.* = bf16(@as(f32, @bitCast(@as(u32, v) << 16)) * scale); + const cb = r.queue.commandBuffer(); + const enc = cb.compute(.serial); + s.apply(enc, buf, n, true); + enc.end(); + cb.commit(); + cb.wait(); + try std.testing.expectEqualSlices(u16, want, x); +} + +/// f32 to bf16, round to nearest even. +fn bf16(f: f32) u16 { + const u: u32 = @bitCast(f); + return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} diff --git a/zig/src/families/flashnext/prompt_w.zig b/zig/src/families/flashnext/prompt_w.zig new file mode 100644 index 0000000000..d23540bc5e --- /dev/null +++ b/zig/src/families/flashnext/prompt_w.zig @@ -0,0 +1,189 @@ +//! Prompt products at any MLX width: qmmw_nax compiled per (bits, group), and which weights take which width. +const std = @import("std"); +const mtl = @import("metal"); +pub const ks = @import("kernel_sources"); +const frags = @import("../../core/frags.zig"); +const replay = @import("replay.zig"); +const Run = replay.Run; +const Buf = replay.Buf; +const Prompt = replay.Prompt; +const Model = replay.Model; +const variants = @import("variants.zig"); + +/// One width's kernels: dense, split-K, sorted-expert gathers (32, 64, 128 rows) and fused gate/up (32, 64 rows). +pub const Pipes = struct { qmm: mtl.Pipeline, splitk: mtl.Pipeline, gather: [3]mtl.Pipeline, gu: [2]mtl.Pipeline }; + +/// A fused member: its rows and width, in member order. +pub const Part = struct { rows: usize, bits: u8, group: u16 }; + +/// A weight at a width other than 6/32: one width, or members of their own widths. +pub const Weight = struct { bits: u8 = 0, group: u16 = 0, parts: []const Part = &.{} }; + +/// A weight's device buffer and offset. +const Key = struct { id: usize, off: usize }; + +/// Every prompt weight by buffer and offset (6/32 too: a missing one is refused, never read at 6 bits) and kernels. +pub const Wide = struct { + weights: std.AutoHashMapUnmanaged(Key, Weight) = .empty, + pipes: std.AutoHashMapUnmanaged(u32, Pipes) = .empty, + missed: bool = false, // a launch met a weight or width the load didn't register: the prompt pass fails + + /// The weight's width when not 6/32 (null: the recorded 6/32 launch); an unregistered weight marks `missed`. + pub fn wide(w: *Wide, b: Buf) ?Weight { + const got = w.weights.get(key(b)) orelse { + w.missed = true; + return null; + }; + return if (got.parts.len == 0 and got.bits == 6 and got.group == 32) null else got; + } + + /// The kernels at a width compiled at load; one never compiled marks `missed` (any width's kernels run meanwhile). + pub fn of(w: *Wide, bits: u8, group: u16) Pipes { + return w.pipes.get(@as(u32, bits) << 16 | group) orelse { + w.missed = true; + var any = w.pipes.valueIterator(); + return any.next().?.*; + }; + } + + /// The pass's refusal when a launch met a weight or width the load didn't register. + pub fn check(w: *const Wide) !void { + if (w.missed) return error.PromptWeightWidthUnknown; + } +}; + +fn key(b: Buf) Key { + return .{ .id = @intFromPtr(b.b.id), .off = b.off }; +} + +/// qmmw_nax at `bits`/`group` with nax.h inlined for this GPU. +pub fn compile(r: *Run, bits: u8, group: u16) !Pipes { + const head = try std.fmt.allocPrint(r.arena, "#define WB {d}\n#define WG {d}\n", .{ bits, group }); + const text = try frags.source(r.device, r.arena, try std.mem.concat(r.arena, u8, &.{ head, ks.flashnext_qmmw })); + const lib = try mtl.Library.fromSource(r.device, text, mtl.CompileOptions.mlx()); + const p = struct { + fn of(rr: *Run, l: mtl.Library, name: [:0]const u8) !mtl.Pipeline { + return mtl.Pipeline.init(rr.device, l, name, false); + } + }; + return .{ + .qmm = try p.of(r, lib, "tf_qmmw_t_nax"), + .splitk = try p.of(r, lib, "tf_qmmw_splitk_nax"), + .gather = .{ try p.of(r, lib, "tf_gather_qmmw_nax_32"), try p.of(r, lib, "tf_gather_qmmw_nax_64"), try p.of(r, lib, "tf_gather_qmmw_nax_128") }, + .gu = .{ try p.of(r, lib, "tf_gather_guw_nax_32"), try p.of(r, lib, "tf_gather_guw_nax_64") }, + }; +} + +fn need(w: *Wide, r: *Run, bits: u8, group: u16) !void { + const k = @as(u32, bits) << 16 | group; + if (!w.pipes.contains(k)) try w.pipes.put(r.arena, k, try compile(r, bits, group)); +} + +fn put(w: *Wide, r: *Run, b: Buf, weight: Weight) !void { + try w.weights.put(r.arena, key(b), weight); + if (weight.parts.len == 0 and weight.bits == 6 and weight.group == 32) return; + if (weight.parts.len == 0) return need(w, r, weight.bits, weight.group); + for (weight.parts) |p| try need(w, r, p.bits, p.group); +} + +fn width(v: [2]usize) Weight { + return .{ .bits = @intCast(v[0]), .group = @intCast(v[1]) }; +} + +/// A pack_mlx projection's `.parts` table, when it has one. +fn parts(r: *Run, stem: []const u8) !?[]const Part { + const name = try std.fmt.allocPrint(r.arena, "{s}.parts", .{stem}); + if (!r.index.contains(name)) return null; + const t = try variants.table(r, name); + const n = t.len / 3; + const out = try r.arena.alloc(Part, n); + for (out, 0..) |*p, i| p.* = .{ .rows = t[i * 3], .bits = @intCast(t[i * 3 + 1]), .group = @intCast(t[i * 3 + 2]) }; + return out; +} + +const six: Weight = .{ .bits = 6, .group = 32 }; + +fn hc(w: *Wide, r: *Run, h: replay.Hc) !void { + try put(w, r, h.uw, if (h.w) |hw| width(hw[1]) else six); + if (h.w) |hw| try need(w, r, @intCast(hw[0][0]), @intCast(hw[0][1])); // the down projection's split-K +} + +/// The prompt's weights at their widths: null on a 6/32 checkpoint, which keeps every recorded launch. +pub fn attach(p: *Prompt, m: *const Model) !void { + const r = p.r; + const w = try r.arena.create(Wide); + w.* = .{}; + for (&m.layers, 0..) |*L, i| { + try hc(w, r, L.ahc); + try hc(w, r, L.mhc); + const x = L.xw orelse variants.Experts{ .gu = .{ 6, 32, 6, 32 }, .down = .{ 6, 32, 6, 32 } }; + for ([_]usize{ 0, 3, 6, 9, 12, 15 }, [_][2]usize{ .{ x.gu[0], x.gu[1] }, .{ x.gu[0], x.gu[1] }, .{ x.gu[2], x.gu[3] }, .{ x.gu[2], x.gu[3] }, .{ x.down[0], x.down[1] }, .{ x.down[2], x.down[3] } }) |j, v| + try put(w, r, L.ex[j], width(v)); // routed gate and up (one fused pass or two), shared, then the downs + const kind = if (L.linear) "gdn" else "att"; + for ([_][]const u8{ if (L.linear) "in" else "proj", if (L.linear) "out" else "o" }, [_]Buf{ p.proj[i][0], p.out[i][0] }) |n, b| { + try put(w, r, b, if (try parts(r, try std.fmt.allocPrint(r.arena, "L{d}.{s}.{s}", .{ i, kind, n }))) |ps| uniformOr(ps) else six); + } + } + try hc(w, r, m.mix); + try hc(w, r, m.mtp.ahc); + try put(w, r, p.ple_kv[0], if (try parts(r, "ple.kv")) |ps| uniformOr(ps) else six); + if (p.mtp_w) |mw| for ([_][]const u8{ "mtp.att.proj", "mtp.fce", "mtp.fch" }, 0..) |n, j| { + try put(w, r, mw[j][0], if (try parts(r, n)) |ps| uniformOr(ps) else six); + }; + p.wide = if (w.pipes.count() == 0) null else w; // all 6/32: every recorded launch as before +} + +/// Members of one width are one product at that width; several widths stay members. +fn uniformOr(ps: []const Part) Weight { + for (ps[1..]) |x| if (x.bits != ps[0].bits or x.group != ps[0].group) return .{ .parts = ps }; + return .{ .bits = ps[0].bits, .group = ps[0].group }; +} + +/// y[rows, n] (row stride ldy, 0: n) = x[rows, k] W^T at the weight's width(s): one launch per member when several. +pub fn qmm(p: *Prompt, wd: *Wide, w: Weight, x: Buf, ws: [3]Buf, k: usize, n: usize, rows: usize, y: Buf, ldy: usize) void { + const ld = if (ldy == 0) n else ldy; + if (w.parts.len == 0) return launch(p, wd.of(w.bits, w.group).qmm, ws, x, k, n, rows, y, ld); + var at = [3]usize{ ws[0].off, ws[1].off, ws[2].off }; + var col: usize = 0; + for (w.parts) |part| { + const m: [3]Buf = .{ .{ .b = ws[0].b, .off = at[0] }, .{ .b = ws[1].b, .off = at[1] }, .{ .b = ws[2].b, .off = at[2] } }; + launch(p, wd.of(part.bits, part.group).qmm, m, x, k, part.rows, rows, y.at(col * 2), ld); + at[0] += part.rows * k * part.bits / 8; + at[1] += part.rows * (k / part.group) * 2; + at[2] += part.rows * (k / part.group) * 2; + col += part.rows; + } +} + +fn launch(p: *Prompt, pipe: mtl.Pipeline, w: [3]Buf, x: Buf, k: usize, n: usize, rows: usize, y: Buf, ld: usize) void { + p.bind(pipe, &.{ w[0], w[1], w[2], x }); + const prm = [4]i32{ @intCast(k), @intCast(n), @intCast(rows), @intCast(ld) }; + p.r.enc.setBytes(std.mem.asBytes(&prm), 4); + p.r.enc.setBuffer(y.b, y.off, 5); + p.r.enc.dispatchThreads(mtl.Size.of(((n + 63) / 64) * 128, (rows + 63) / 64, 1), mtl.Size.of(128, 1, 1)); + p.barrier(); +} + +/// The sorted-expert gather at the weight's width: tiles of 32, 64 or 128 rows as Prompt.gather picks them. +pub fn gather(p: *Prompt, pipes: Pipes, bm: usize, x: Buf, w: []const Buf, k: usize, n: usize, pairs: usize, y: Buf) void { + p.bind(pipes.gather[if (bm == 128) 2 else if (bm == 64) 1 else 0], &.{ x, w[0], w[1], w[2], p.b.off }); + const prm = [4]i32{ @intCast(pairs), @intCast(n), @intCast(k), 512 }; + p.r.enc.setBytes(std.mem.asBytes(&prm), 5); + p.r.enc.setBuffer(y.b, y.off, 6); + p.r.enc.dispatchThreads(mtl.Size.of(((n + 63) / 64) * 128, pairs / bm + 512, 1), mtl.Size.of(128, 1, 1)); + p.barrier(); +} + +/// The fused gate/up gather at the weights' width. +pub fn gatherGU(p: *Prompt, pipes: Pipes, bm: usize, x: Buf, wg: []const Buf, wu: []const Buf, k: usize, n: usize, pairs: usize, y: Buf) void { + p.bind(pipes.gu[if (bm == 64) 1 else 0], &.{ x, wg[0], wg[1], wg[2], wu[0], wu[1], wu[2], p.b.off }); + const prm = [4]i32{ @intCast(pairs), @intCast(n), @intCast(k), 512 }; + p.r.enc.setBytes(std.mem.asBytes(&prm), 8); + p.r.enc.setBuffer(y.b, y.off, 9); + p.r.enc.dispatchThreads(mtl.Size.of(((n + 63) / 64) * 128, pairs / bm + 512, 1), mtl.Size.of(128, 1, 1)); + p.barrier(); +} + +test { + _ = @import("prompt_w_test.zig"); +} diff --git a/zig/src/families/flashnext/prompt_w_test.zig b/zig/src/families/flashnext/prompt_w_test.zig new file mode 100644 index 0000000000..f58a484eed --- /dev/null +++ b/zig/src/families/flashnext/prompt_w_test.zig @@ -0,0 +1,12 @@ +//! The prompt kernels at any width compile from one text, and fused members' offsets follow their own widths. +const std = @import("std"); +const ks = @import("kernel_sources"); + +test "the any-width prompt text names its kernels and declares no width of its own" { + const text = ks.flashnext_qmmw; + for ([_][]const u8{ "tf_qmmw_t_nax", "tf_qmmw_splitk_nax", "tf_gather_guw_nax_##BM", "tf_gather_qmmw_nax_##BM" }) |name| + try std.testing.expect(std.mem.indexOf(u8, text, name) != null); + try std.testing.expect(std.mem.indexOf(u8, text, "#define WB") == null); + try std.testing.expect(std.mem.indexOf(u8, text, "#define WG") == null); + try std.testing.expectEqual(@as(usize, 1), std.mem.count(u8, text, "#include \"../nax.h\"")); +} diff --git a/zig/src/families/flashnext/replay.zig b/zig/src/families/flashnext/replay.zig index f78fa08c28..cf54294663 100644 --- a/zig/src/families/flashnext/replay.zig +++ b/zig/src/families/flashnext/replay.zig @@ -10,6 +10,11 @@ const mark_slots = @import("marks.zig"); const verdict_kernels = @import("verdict.zig"); pub const split = @import("split.zig"); pub const dense = @import("dense.zig"); +pub const lane_w = @import("lane_w.zig"); +pub const ple_scale = @import("ple_scale.zig"); +pub const variants = @import("variants.zig"); +pub const prompt_w = @import("prompt_w.zig"); +pub const engine_widths = @import("engine_widths.zig"); pub const gdn_step = @import("gdn.zig"); pub const Tp2 = tpm.Tp2; pub const frags = @import("../../core/frags.zig"); @@ -887,6 +892,7 @@ pub const Run = struct { lane_new: bool = false, // the target's lane projections on core lane (dense.zig): the recorded bits, weights read ahead lane_pf: usize = 1, // groups a simdgroup reads ahead in core lane lane_pipes: std.AutoHashMapUnmanaged(u64, mtl.Pipeline) = .empty, + wide_pipes: lane_w.Cache = .empty, // lane products at other widths than 6/32 (lane_w.zig) lane_shape: std.AutoHashMapUnmanaged(*Variant, [3]usize) = .empty, // recorded lane kernels' N, K, SK gdn_pipe: ?mtl.Pipeline = null, // the DeltaNet window step on fz_gdn (gdn.zig) when set gdn_kept: ?mtl.Pipeline = null, // GPU-side rounds: fz_gdn keeping one state a layer, kept rows replayed (gdn.zig) @@ -1118,9 +1124,13 @@ pub const Run = struct { /// A hyper-connection's down then up projection at `as_rows` from normed streams `hn` (sums of squares in `ssp`): /// the mixed rows and inject gates; with `hc_up` the up projection of 3-8 rows on fz_hc_up (the same bits). pub fn hcProject(r: *Run, down: []const u8, up: []const u8, as_rows: usize, hn: Buf, ssp: Buf, hc: Hc, eps: Buf, rows: Buf, part: Buf, mixed: Buf, inj: Buf) !void { - try r.callAs(down, as_rows, &.{ hn, ssp, hc.scale, hc.dw, hc.ds, hc.db, eps, rows }, &.{part}); + var cd: [2]variants.Const = undefined; + var cu: [2]variants.Const = undefined; + const kd = variants.width(&cd, "BITS", "GS", if (hc.w) |w| w[0] else null); + const ku = variants.width(&cu, "BITS", "GS", if (hc.w) |w| w[1] else null); + try variants.callAs(r, down, as_rows, &.{ hn, ssp, hc.scale, hc.dw, hc.ds, hc.db, eps, rows }, &.{part}, kd); const ins = [_]Buf{ hn, ssp, hc.scale, part, hc.uw, hc.us, hc.ub, eps, rows }; - if (!r.hc_up or as_rows < 3 or as_rows > 8) return r.callAs(up, as_rows, &ins, &.{ mixed, inj }); + if (!r.hc_up or hc.w != null or as_rows < 3 or as_rows > 8) return variants.callAs(r, up, as_rows, &ins, &.{ mixed, inj }, ku); if (r.skip & class(up) != 0 or r.hcskip & 4 != 0) return; const nd: u32 = @intCast(hc.dw.b.length() / (WIDE * 6 / 8)); r.enc.setPipeline(r.hc_up_pipe); @@ -1195,7 +1205,7 @@ pub const Run = struct { } /// y = x W for `rows` rows of K inputs through fz_dense (blocks of 8 rows; 16 K-slices for narrow outputs). /// One row's dense projection over output tiles [t0, t0 + tn) only (32 outputs a tile), in the full layout. - pub fn denseTiles(r: *Run, x: Buf, k: usize, l: Lane, y: Buf, t0: usize, tn: usize) void { + pub fn denseTiles(r: *Run, x: Buf, k: usize, l: Lane.Six, y: Buf, t0: usize, tn: usize) void { if (tn == 0) return; const n = l.wq.b.length() * 4 / (3 * k); const narrow = n <= 4096; @@ -1211,7 +1221,7 @@ pub const Run = struct { if (!r.serial) r.enc.barrier(); } - pub fn denseRows(r: *Run, x: Buf, k: usize, l: Lane, rows: usize, y: Buf) void { + pub fn denseRows(r: *Run, x: Buf, k: usize, l: Lane.Six, rows: usize, y: Buf) void { const n = l.wq.b.length() * 4 / (3 * k); const narrow = n <= 4096; var at: usize = 0; @@ -1414,6 +1424,7 @@ pub fn catchSlot(words: []i32, k: usize, first: usize) usize { test { _ = tpm; // tp.zig's tests: the window layout and the host service's wraps + _ = @import("drafts.zig"); // the depth rule's tests } test "catch-up slots: every layer's dispatch reads its own first block after the later layers' writes" { @@ -1596,8 +1607,9 @@ pub const GSelect = struct { } }; -pub const Hc = struct { scale: Buf, dw: Buf, ds: Buf, db: Buf, uw: Buf, us: Buf, ub: Buf }; -pub const Lane = struct { wq: Buf, sbt: Buf }; +/// A pack hyper-connection; `w`: the down and up widths when not 6/32. +pub const Hc = struct { scale: Buf, dw: Buf, ds: Buf, db: Buf, uw: Buf, us: Buf, ub: Buf, w: ?[2][2]usize = null }; +pub const Lane = lane_w.Lane; // the recorded 6/32 pair, or runs at other widths pub const Layer = struct { ahc: Hc, mhc: Hc, @@ -1613,6 +1625,7 @@ pub const Layer = struct { iqn: Buf = undefined, router: Buf, ex: [18]Buf, // gate, up, shared gate, shared up (weight, scales, biases), then down, shared down + xw: ?variants.Experts = null, // the experts' widths when not 6/32 (engine_widths.zig) cs: [2]Buf = undefined, so: [2]Buf = undefined, keys: Buf = undefined, @@ -1623,22 +1636,8 @@ pub const Layer = struct { pooled_n: usize = 0, }; -pub fn hcOf(r: *Run, comptime fmt: []const u8, args: anytype) !Hc { - var name: [96]u8 = undefined; - const stem = try std.fmt.bufPrint(&name, fmt, args); - var full: [128]u8 = undefined; - const parts = [_][]const u8{ "scale", "down.w", "down.s", "down.b", "up.w", "up.s", "up.b" }; - var out: [7]Buf = undefined; - for (parts, 0..) |p, i| out[i] = try r.load(try std.fmt.bufPrint(&full, "{s}.{s}", .{ stem, p })); - return .{ .scale = out[0], .dw = out[1], .ds = out[2], .db = out[3], .uw = out[4], .us = out[5], .ub = out[6] }; -} - -pub fn laneOf(r: *Run, comptime fmt: []const u8, args: anytype) !Lane { - var name: [96]u8 = undefined; - const stem = try std.fmt.bufPrint(&name, fmt, args); - var full: [128]u8 = undefined; - return .{ .wq = try r.load(try std.fmt.bufPrint(&full, "{s}.wq", .{stem})), .sbt = try r.load(try std.fmt.bufPrint(&full, "{s}.sbt", .{stem})) }; -} +pub const hcOf = variants.hcOf; +pub const laneOf = lane_w.laneOf; pub const Ple = struct { kv: Lane, @@ -1654,6 +1653,8 @@ pub const Ple = struct { mult: [3]i64, sizes: [16]i64, offsets: [16]i64, + w: ?[2]usize = null, // the n-gram tables' width when not 6/32 + scale: ?ple_scale.Scale = null, // the n-gram tables' weight_scale when not 1 }; /// Every intermediate a window of up to MAXR rows writes, and the per-window values the kernels read. @@ -1750,6 +1751,7 @@ pub const Mtp = struct { ids: Buf, ids_n: usize, ex: [18]Buf, + xw: ?variants.Experts = null, keys: Buf, vals: Buf, raw: Buf, @@ -1784,6 +1786,7 @@ pub const Model = struct { mix: Hc, head: Lane, embed: [3]Buf, + embed_w: ?[2]usize = null, // the embedding's width when not 6/32 ple: Ple, t: Tmp, pos: usize = 0, @@ -1817,10 +1820,12 @@ pub const Model = struct { } pub fn lane(m: *Model, x: Buf, k: usize, l: Lane, role: []const u8, y: Buf) !void { - if (m.r.dense_target) return m.r.denseRows(x, k, l, m.r.rows, y); - if (m.r.lane_new) return dense.lane(m.r, role, x, l, m.t.mdims, y); + if (l.w) |w| return lane_w.launch(m.r, w, role, x, m.t.mdims, m.r.rows, y); + const six = try l.recorded(); + if (m.r.dense_target) return m.r.denseRows(x, k, six, m.r.rows, y); + if (m.r.lane_new) return dense.lane(m.r, role, x, six, m.t.mdims, y); if (!m.r.fused_xsum) try m.r.call(if (k == D) "lane_qmm_xsum#[2560]" else "lane_qmm_xsum#[6144]", &.{ x, m.t.mdims }, &.{m.t.xs}); - try m.r.call(role, &.{ x, m.t.xs, l.wq, l.sbt, m.t.mdims }, &.{y}); + try m.r.call(role, &.{ x, m.t.xs, six.wq, six.sbt, m.t.mdims }, &.{y}); } pub fn hcProject(m: *Model, hn: Buf, hc: Hc, down: []const u8, up: []const u8, inj: Buf) !void { @@ -1945,7 +1950,8 @@ pub const Model = struct { if (!r.serial) r.enc.barrier(); } var cur: usize = 0; - try r.call("qa_embed_rows@embed", &.{ ids, m.embed[0], m.embed[1], m.embed[2] }, &.{t.h[0]}); + var ce: [2]variants.Const = undefined; + try variants.callAs(r, "qa_embed_rows@embed", rows, &.{ ids, m.embed[0], m.embed[1], m.embed[2] }, &.{t.h[0]}, variants.width(&ce, "BITS", "GS", m.embed_w)); var pending: enum { none, grouped } = .none; for (0..LAYERS) |i| { const L = &m.layers[i]; @@ -1958,7 +1964,9 @@ pub const Model = struct { tabs[0] = t.ple_ids; tabs[1] = p.starts; for (0..3 * GROUPS) |j| tabs[2 + j] = p.tables[j]; - try r.call("qa_ple_lookup@ple", &tabs, &.{t.emb}); + var cp: [2]variants.Const = undefined; + try variants.callAs(r, "qa_ple_lookup@ple", rows, &tabs, &.{t.emb}, variants.width(&cp, "BITS", "GS", p.w)); + if (p.scale) |sc| sc.apply(r.enc, t.emb, rows * D, r.serial); try m.lane(t.emb, D, p.kv, "lane_qmm_bytes_grouped@ple.kv", t.kvp); try r.call("q4_ple_gate@ple", &.{ t.kvp, t.h[cur], p.ks, p.qs, p.cs, t.eps }, &.{ t.gated, .{ .b = p.cin.b, .off = PLE_TAIL * WIDE * 2 } }); try r.call("q4_ple_conv@ple", &.{ p.cin, p.conv, t.gated, t.h[cur] }, &.{t.hout}); @@ -1977,7 +1985,8 @@ pub const Model = struct { if (!r.fused_xsum) try r.call("lane_qmm_xsum#[2560]", &.{ t.mixed, t.mdims }, &.{t.xs}); // [q 64 tiles | k 64 | v 192 | z 192 | b, a 3]: this Mac's q, k, v and z heads, b and a whole, 8 K slices const cols = [_][2]usize{ .{ 32 * k0, 32 }, .{ 64 + 32 * k0, 32 }, .{ 128 + 96 * k0, 96 }, .{ 320 + 96 * k0, 96 }, .{ 512, 3 } }; - try split.laneTiles(r, "lane_qmm_bytes_grouped@gdn.in", &.{ t.mixed, t.xs, L.proj.wq, L.proj.sbt, t.mdims }, t.p, &cols, 8, null); + const pj = try L.proj.recorded(); + try split.laneTiles(r, "lane_qmm_bytes_grouped@gdn.in", &.{ t.mixed, t.xs, pj.wq, pj.sbt, t.mdims }, t.p, &cols, 8, null); const a = m.state; const cs_in: Buf = .{ .b = L.cs[a].b, .off = L.cs[a].off + m.state_row * CS_ROW }; const so_in: Buf = .{ .b = L.so[a].b, .off = L.so[a].off + m.state_row * SO_ROW }; @@ -1987,7 +1996,8 @@ pub const Model = struct { } else try split.gdnHeads(r, "q4_gdn@gdn", if (r.gdn_step and rows > 1) 8 else rows, &.{ t.p, cs_in, so_in, L.conv, L.alog, L.dt, L.norm, t.eps, t.rows }, &.{ t.gout, L.cs[1 - a], L.so[1 - a] }, tp.rank); const pn = tp.partNext(); const part: Buf = .{ .b = pn.b, .off = pn.off }; - try split.laneTiles(r, "lane_qmm_bytes_grouped@gdn.out", &.{ t.gout, t.xs, L.out.wq, L.out.sbt, t.mdims }, part, &.{.{ 0, 80 }}, 8, .{ 96 * k0, 96 }); // one kernel at every width: drafted == plain + const po = try L.out.recorded(); + try split.laneTiles(r, "lane_qmm_bytes_grouped@gdn.out", &.{ t.gout, t.xs, po.wq, po.sbt, t.mdims }, part, &.{.{ 0, 80 }}, 8, .{ 96 * k0, 96 }); // one kernel at every width: drafted == plain if (r.skip & TP_CLASS == 0) { // the partials' sum and the stream update in the exchange's launch tp.plain(r.enc, t.h[cur], t.inj_a, t.h[1 - cur], t.ssp, t.rows, rows); plain = true; @@ -1997,8 +2007,7 @@ pub const Model = struct { const a = m.state; const cs_in: Buf = .{ .b = L.cs[a].b, .off = L.cs[a].off + m.state_row * CS_ROW }; const so_in: Buf = .{ .b = L.so[a].b, .off = L.so[a].off + m.state_row * SO_ROW }; - r.touch(L.out.wq); - r.touch(L.out.sbt); + lane_w.touch(r, L.out); const gins = [_]Buf{ t.p, cs_in, so_in, L.conv, L.alog, L.dt, L.norm, t.eps, t.rows }; const gouts = [_]Buf{ t.gout, L.cs[1 - a], L.so[1 - a] }; if (m.keptState()) |recs| { @@ -2008,8 +2017,7 @@ pub const Model = struct { try m.lane(t.gout, 6144, L.out, "lane_qmm_bytes_grouped@gdn.out", t.branch); } else { try m.lane(t.mixed, D, L.proj, "lane_qmm_bytes_grouped@att.proj", t.p); - r.touch(L.out.wq); - r.touch(L.out.sbt); + lane_w.touch(r, L.out); try r.call("q4_attn_prep@att", &.{ t.p, t.pos8, L.qn, L.kn, L.iqn, t.eps, t.log2base }, &.{ t.q, t.kout, t.iq }); if (r.skip & (1 << 5) == 0) { r.enc.setPipeline(r.kv_pipe); @@ -2054,7 +2062,7 @@ pub const Model = struct { try m.hcProject(t.h[cur], L.mhc, "qa_hc_down@mhc", "qa_hc_up@mhc", t.inj_m); try r.router("q4_router_float@moe", t.mixed, L.router, t.rows, t.lg, rows); r.tp_layer = true; - try r.experts("qa_expert_gateup@moe.gate", "qa_expert_down_y@moe.down", t.mixed, t.lg, L.ex, t.act, t.pick, t.wts, t.rows, t.ydown); + if (L.xw) |w| try variants.experts(r, "qa_expert_gateup@moe.gate", "qa_expert_down_y@moe.down", t.mixed, t.lg, L.ex, t.act, t.pick, t.wts, t.rows, t.ydown, w) else try r.experts("qa_expert_gateup@moe.gate", "qa_expert_down_y@moe.down", t.mixed, t.lg, L.ex, t.act, t.pick, t.wts, t.rows, t.ydown); r.tp_layer = false; if (r.tp) |tp| if (r.skip & TP_CLASS == 0) tp.exchange(r.enc, t.ydown, t.wts, t.lg, rows); if (r.probe) |pb| r.copyKept(t.pick, .{ .b = pb.b, .off = pb.off + i * MAXR * 10 * 4 }, rows * 10, 0, 0, 0, 1, -1); @@ -2068,7 +2076,8 @@ pub const Model = struct { if (r.tp) |tp| { // TP: this Mac's half of the vocabulary (the same kernel, half its tiles), then one swap const half = HEAD_TILES / 2; if (!r.fused_xsum) try r.call("lane_qmm_xsum#[2560]", &.{ t.mixed, t.mdims }, &.{t.xs}); - try split.laneTiles(r, "lane_qmm_bytes_grouped@head", &.{ t.mixed, t.xs, m.head.wq, m.head.sbt, t.mdims }, t.logits, &.{.{ half * tp.rank, half }}, 1, null); + const hd = try m.head.recorded(); + try split.laneTiles(r, "lane_qmm_bytes_grouped@head", &.{ t.mixed, t.xs, hd.wq, hd.sbt, t.mdims }, t.logits, &.{.{ half * tp.rank, half }}, 1, null); if (r.skip & TP_CLASS == 0) tp.argmax(r.enc, t.logits, HEAD_TILES * 32, half * 32 * tp.rank, half * 32, t.picks, t.rows, rows); return; } @@ -2132,13 +2141,14 @@ pub const Model = struct { const t = &m.t; const h = &m.mtp; r.rows = rows; - try r.call("mtp:qa_embed_rows@embed", &.{ ids, m.embed[0], m.embed[1], m.embed[2] }, &.{h.emb}); + var ce: [2]variants.Const = undefined; + try variants.callAs(r, "mtp:qa_embed_rows@embed", rows, &.{ ids, m.embed[0], m.embed[1], m.embed[2] }, &.{h.emb}, variants.width(&ce, "BITS", "GS", m.embed_w)); try r.call("mtp:q4_rms_rows@mtp.enorm", &.{ h.emb, h.enorm, t.eps }, &.{h.en}); if (!r.fused_xsum) try r.call("mtp:lane_qmm_xsum#[2560]", &.{ h.en, sl.md }, &.{t.xs}); - if (r.dense) r.denseRows(h.en, D, h.fce, rows, h.e) else try r.call("mtp:lane_qmm_bytes_grouped@mtp.fce", &.{ h.en, t.xs, h.fce.wq, h.fce.sbt, sl.md }, &.{h.e}); + try m.mtpLane(h.fce, "mtp:lane_qmm_bytes_grouped@mtp.fce", h.en, sl.md, rows, D, h.e); try r.call("mtp:q4_rms_rows@mtp.hnorm", &.{ streams, h.hnorm, t.eps }, &.{h.hn}); if (!r.fused_xsum) try r.call("mtp:lane_qmm_xsum#[4R, 2560]", &.{ h.hn, sl.md4 }, &.{t.xs}); - if (r.dense) r.denseRows(h.hn, D, h.fch, 4 * rows, h.hs) else try r.call("mtp:lane_qmm_bytes_grouped@mtp.fch", &.{ h.hn, t.xs, h.fch.wq, h.fch.sbt, sl.md4 }, &.{h.hs}); + try m.mtpLane(h.fch, "mtp:lane_qmm_bytes_grouped@mtp.fch", h.hn, sl.md4, 4 * rows, D, h.hs); r.enc.setPipeline(r.add_pipe); for ([_]Buf{ h.e, h.hs, h.h[0], sl.n_add }, 0..) |b, j| r.enc.setBuffer(b.b, b.off, j); r.enc.dispatchThreads(mtl.Size.of(rows * WIDE, 1, 1), mtl.Size.of(256, 1, 1)); @@ -2148,7 +2158,7 @@ pub const Model = struct { const up = [_][]const u8{ "mtp:qa_hc_up@mtp.ahc", "mtp:qa_hc_up@mtp.mhc", "mtp:qa_hc_up@mtp.mix" }; try m.mtpProject(h.h[1], h.ahc, down[0], up[0], t.inj_a, sl.rows); if (!r.fused_xsum) try r.call("mtp:lane_qmm_xsum#[2560]", &.{ t.mixed, sl.md }, &.{t.xs}); - if (r.dense) r.denseRows(t.mixed, D, h.proj, rows, t.p) else try r.call("mtp:lane_qmm_bytes_grouped@mtp.att.proj", &.{ t.mixed, t.xs, h.proj.wq, h.proj.sbt, sl.md }, &.{t.p}); + try m.mtpLane(h.proj, "mtp:lane_qmm_bytes_grouped@mtp.att.proj", t.mixed, sl.md, rows, D, t.p); try r.call("mtp:q4_attn_prep@mtp.att", &.{ t.p, sl.pos8, h.qn, h.kn, h.iqn, t.eps, t.log2base }, &.{ t.q, t.kout, t.iq }); r.enc.setPipeline(r.kv_pipe); for ([_]Buf{ t.kout, t.p, h.keys, h.vals, h.raw, sl.kvmeta }, 0..) |b, j| r.enc.setBuffer(b.b, b.off, j); @@ -2201,41 +2211,53 @@ pub const Model = struct { } r.rows = 1; if (!r.fused_xsum) try r.call("mtp:lane_qmm_xsum#[6144]", &.{ t.aout, h.md1 }, &.{t.xs}); - if (r.dense) r.denseRows(t.aout, 6144, h.out, 1, t.branch) else try r.call("mtp:lane_qmm_bytes_grouped@mtp.att.o", &.{ t.aout, t.xs, h.out.wq, h.out.sbt, h.md1 }, &.{t.branch}); + try m.mtpLane(h.out, "mtp:lane_qmm_bytes_grouped@mtp.att.o", t.aout, h.md1, 1, 6144, t.branch); try r.call("mtp:q4_hc_norm_plain#[10240]", &.{ h.h[1], t.inj_a, t.branch }, &.{ h.h[0], t.ssp }); try m.mtpProject(h.h[0], h.mhc, down[1], up[1], t.inj_m, h.md1); try r.router("mtp:q4_router_float@mtp.moe", t.mixed, h.router, h.md1, t.lg, 1); - try r.experts("mtp:qa_expert_gateup@mtp.moe.gate", "mtp:qa_expert_down_y@mtp.moe.down", t.mixed, t.lg, h.ex, t.act, t.pick, t.wts, h.md1, t.ydown); + if (h.xw) |w| try variants.experts(r, "mtp:qa_expert_gateup@mtp.moe.gate", "mtp:qa_expert_down_y@mtp.moe.down", t.mixed, t.lg, h.ex, t.act, t.pick, t.wts, h.md1, t.ydown, w) else try r.experts("mtp:qa_expert_gateup@mtp.moe.gate", "mtp:qa_expert_down_y@mtp.moe.down", t.mixed, t.lg, h.ex, t.act, t.pick, t.wts, h.md1, t.ydown); try r.call("mtp:q4_hc_norm_grouped#[10240]", &.{ h.h[0], t.inj_m, t.ydown, t.wts, t.lg }, &.{ h.h[1], t.ssp }); h.last = h.h[1]; try m.mtpProject(h.h[1], h.mix, down[2], up[2], t.inj_a, h.md1); const x = t.mixed; if (!r.fused_xsum) try r.call("mtp:lane_qmm_xsum#[2560]", &.{ x, h.md1 }, &.{t.xs}); - if (r.tp != null and r.dense) { // TP: this Mac's half of the draft vocabulary's tiles, then one exact argmax swap + // TP: this Mac's half of the draft vocabulary's tiles, then one exact argmax swap + if (r.tp != null and r.dense and h.draft.w == null) { const tp = r.tp.?; - const n = h.draft.wq.b.length() * 4 / (3 * D); + const six = try h.draft.recorded(); + const n = six.wq.b.length() * 4 / (3 * D); const tiles = n / 32; const half = (tiles + 1) / 2; const t0: usize = if (tp.rank == 0) 0 else half; const tn: usize = if (tp.rank == 0) half else tiles - half; - r.denseTiles(x, D, h.draft, h.logits, t0, tn); + r.denseTiles(x, D, six, h.logits, t0, tn); const lo = @min(t0 * 32, h.ids_n); const hi = @min((t0 + tn) * 32, h.ids_n); tp.argmax(r.enc, h.logits, n, lo, hi - lo, Buf{ .b = tp.pick_tmp }, Buf{ .b = tp.one }, 1); tp.mapIds(r.enc, h.ids, out); return; } - if (r.dense) r.denseRows(x, D, h.draft, 1, h.logits) else try r.call("mtp:lane_qmm_bytes_grouped@mtp.draft", &.{ x, t.xs, h.draft.wq, h.draft.sbt, h.md1 }, &.{h.logits}); + try m.mtpLane(h.draft, "mtp:lane_qmm_bytes_grouped@mtp.draft", x, h.md1, 1, D, h.logits); r.enc.setPipeline(r.argids_pipe); for ([_]Buf{ h.logits, out, h.n_ids, h.ids }, 0..) |b, j| r.enc.setBuffer(b.b, b.off, j); r.enc.dispatchThreads(mtl.Size.of(1024, 1, 1), mtl.Size.of(1024, 1, 1)); if (!r.serial) r.enc.barrier(); } + /// An MTP head lane: at its widths' runs, as dense rows (r.dense), or the recorded launch. + fn mtpLane(m: *Model, l: Lane, role: []const u8, x: Buf, md: Buf, rows: usize, k: usize, y: Buf) !void { + if (l.w) |w| return lane_w.launch(m.r, w, role, x, md, rows, y); + const six = try l.recorded(); + if (m.r.dense) return m.r.denseRows(x, k, six, rows, y); + try m.r.call(role, &.{ x, m.t.xs, six.wq, six.sbt, md }, &.{y}); + } + pub fn mtpProject(m: *Model, hn: Buf, hc: Hc, down: []const u8, up: []const u8, inj: Buf, rows: Buf) !void { const t = &m.t; - try m.r.call(down, &.{ hn, t.ssp, hc.scale, hc.dw, hc.ds, hc.db, t.eps, rows }, &.{t.part}); - try m.r.call(up, &.{ hn, t.ssp, hc.scale, t.part, hc.uw, hc.us, hc.ub, t.eps, rows }, &.{ t.mixed, inj }); + var c: [2][2]variants.Const = undefined; + const k = [2][]const variants.Const{ variants.width(&c[0], "BITS", "GS", if (hc.w) |w| w[0] else null), variants.width(&c[1], "BITS", "GS", if (hc.w) |w| w[1] else null) }; + try variants.callAs(m.r, down, m.r.rows, &.{ hn, t.ssp, hc.scale, hc.dw, hc.ds, hc.db, t.eps, rows }, &.{t.part}, k[0]); + try variants.callAs(m.r, up, m.r.rows, &.{ hn, t.ssp, hc.scale, t.part, hc.uw, hc.us, hc.ub, t.eps, rows }, &.{ t.mixed, inj }, k[1]); } /// Absorb `nexts.len` rows (chunks of up to MAXR) from `streams`; returns the draft after the last row. @@ -2274,68 +2296,11 @@ pub const Model = struct { } }; -/// Drafts a round from recent landing: 3 (prose), 6 while drafts land (code-like text), the widest window while -/// nearly all land (copied text). /// FZ_PROFILE's knock-out bit for speed-up mode's exchanges: the bit after `Run.class`'s eight. pub const TP_CLASS: u32 = 1 << 8; -pub const DepthRule = struct { - rate: f64 = 0.6, // moving average of the head's drafts landed over drafts offered - depth: usize = 3, - pair: bool = false, // speed-up mode: even windows only (an odd window's last row costs a Mac a whole expert row) - - pub fn pick(self: *const DepthRule) usize { - return self.depth; - } - - /// A head round's drafts and how many landed: 3, 6 or 8 drafts on one Mac, 3 or 7 on the pair (copied rounds set - /// their own width). - pub fn update(self: *DepthRule, depth: usize, landed: usize) void { - self.rate = 0.7 * self.rate + 0.3 * @as(f64, @floatFromInt(landed)) / @as(f64, @floatFromInt(depth)); - if (self.pair) { - if (self.depth == 3 and self.rate > 0.95) self.depth = 7 else if (self.depth == 7 and self.rate < 0.8) self.depth = 3; - return; - } - if (self.depth == 3 and self.rate > 0.8) { - self.depth = 6; - } else if (self.depth == 6 and self.rate > 0.74) { - self.depth = 8; - } else if (self.depth == 8 and self.rate < 0.62) { - self.depth = 6; - } else if (self.depth == 6 and self.rate < 0.65) self.depth = 3; - } -}; - -test "the pair's depth rule keeps every window even" { - var pair: DepthRule = .{ .pair = true }; - for (0..20) |_| pair.update(pair.pick(), pair.pick()); - try std.testing.expectEqual(@as(usize, 7), pair.pick()); - for (0..20) |_| pair.update(pair.pick(), 1); - try std.testing.expectEqual(@as(usize, 3), pair.pick()); - var one: DepthRule = .{}; - for (0..20) |_| one.update(one.pick(), one.pick()); - try std.testing.expectEqual(@as(usize, 8), one.pick()); -} - -/// Copy lanes: the longest suffix of `hist` (`min`..8 tokens) seen earlier; the tokens after its latest earlier -/// occurrence go into `out`. Returns how many (0 when nothing matches). -pub fn copyDrafts(hist: []const u32, min: usize, out: []u32) usize { - const n = hist.len; - if (n < min + 1) return 0; - var len: usize = @min(8, n - 1); - while (len >= min) : (len -= 1) { - const suffix = hist[n - len ..]; - var end: usize = n - 1; - while (end >= len) : (end -= 1) { - if (std.mem.eql(u32, hist[end - len .. end], suffix)) { - const k = @min(out.len, n - end); - @memcpy(out[0..k], hist[end .. end + k]); - return k; - } - } - } - return 0; -} +pub const DepthRule = @import("drafts.zig").DepthRule; // drafts a round from recent landing +pub const copyDrafts = @import("drafts.zig").copyDrafts; /// Prompt chunks of up to PMAX rows: projections on the 6-bit tensor-unit kernels, experts sorted by expert and /// gathered, the decode's row kernels at the chunk's rows, and DeltaNet storing only its last row's state. @@ -2351,6 +2316,7 @@ pub const Prompt = struct { seg_queue: ?mtl.Queue = null, // a later staggered segment's queue (made at first use) last: Buf = undefined, // the last chunk's streams before the final mixer (each segment's own) gu32: mtl.Pipeline, + wide: ?*prompt_w.Wide = null, // weights at other widths than 6/32 (prompt_w.zig) fused_gu: bool = true, // the experts' gate and up products and activation in one pass (same bits) expert_bm: usize = 0, // expert tile rows (0: by the average rows an expert) mtp_w: ?[3][3]Buf = null, // the MTP head's attention projection, fce, fch in MLX layout (pack_mtp_mlx.safetensors) @@ -2379,6 +2345,7 @@ pub const Prompt = struct { pub fn init(r: *Run, dir: []const u8, header: []const u8) !Prompt { var p: Prompt = undefined; // every field is set below: `undefined` skips the declared defaults p.r = r; + p.wide = null; p.skip = 0; p.sel = null; p.step = PMAX; @@ -2542,6 +2509,7 @@ pub const Prompt = struct { /// y[rows, n] (row stride ldy, 0: n) = x[rows, k] W^T, W 6-bit g32 in MLX's layout. pub fn qmm(p: *Prompt, x: Buf, w: [3]Buf, k: usize, n: usize, rows: usize, y: Buf, ldy: usize) void { if (p.skip & 8 != 0) return; + if (p.wide) |wd| if (wd.wide(w[0])) |ww| return prompt_w.qmm(p, wd, ww, x, w, k, n, rows, y, ldy); const tall = p.tall_tiles and rows >= 512; p.bind(if (tall) p.qmm6_128 else p.qmm6, &.{ w[0], w[1], w[2], x }); const prm = [4]i32{ @intCast(k), @intCast(n), @intCast(rows), @intCast(ldy) }; @@ -2557,6 +2525,7 @@ pub const Prompt = struct { if (p.skip & 1 != 0) return; // tiles of 32, 64 or 128 rows as experts average 32 and 128 rows (128: FZ prompt chunks of thousands of rows) const bm: usize = if (p.expert_bm != 0) p.expert_bm else if (p.tall_tiles and pairs >= 512 * 128) 128 else if (pairs >= 512 * 32) 64 else 32; + if (p.wide) |wd| if (wd.wide(w[0])) |ww| return prompt_w.gather(p, wd.of(ww.bits, ww.group), bm, x, w, k, n, pairs, y); p.bind(if (bm == 128) p.gather128 else if (bm == 64) p.gather64 else p.gather32, &.{ x, w[0], w[1], w[2], p.b.off }); const prm = [4]i32{ @intCast(pairs), @intCast(n), @intCast(k), 512 }; p.r.enc.setBytes(std.mem.asBytes(&prm), 5); @@ -2569,6 +2538,7 @@ pub const Prompt = struct { pub fn gatherGU(p: *Prompt, x: Buf, wg: []const Buf, wu: []const Buf, k: usize, n: usize, pairs: usize, y: Buf) void { if (p.skip & 1 != 0) return; const bm: usize = if (p.expert_bm == 32 or p.expert_bm == 64) p.expert_bm else if (pairs >= 512 * 32) 64 else 32; + if (p.wide) |wd| if (wd.wide(wg[0])) |ww| return prompt_w.gatherGU(p, wd.of(ww.bits, ww.group), bm, x, wg, wu, k, n, pairs, y); p.bind(if (bm == 64) p.gu64 else p.gu32, &.{ x, wg[0], wg[1], wg[2], wu[0], wu[1], wu[2], p.b.off }); const prm = [4]i32{ @intCast(pairs), @intCast(n), @intCast(k), 512 }; p.r.enc.setBytes(std.mem.asBytes(&prm), 8); @@ -2581,13 +2551,13 @@ pub const Prompt = struct { pub fn hc(p: *Prompt, m: *Model, hn: Buf, w: Hc, rows: usize, inj: Buf) void { if (p.skip & 16 != 0) return; const b = &p.b; - const nd = w.dw.b.length() / (WIDE * 6 / 32 * 4); + const nd = w.dw.b.length() / (WIDE * (if (w.w) |hw| hw[0][0] else 6) / 8); // the down rows (+ inject rows) p.bind(p.pl[0], &.{ hn, b.ssp, w.scale, m.t.eps, b.normed }); p.r.enc.dispatchThreads(mtl.Size.of(WIDE, rows, 1), mtl.Size.of(256, 1, 1)); p.barrier(); if (p.skip & 8 == 0) { // the down projection's 10,240-deep sum in 8 parts, added in order const parts: usize = 8; - p.bind(p.splitk, &.{ w.dw, w.ds, w.db, b.normed }); + p.bind(if (w.w) |hw| p.wide.?.of(@intCast(hw[0][0]), @intCast(hw[0][1])).splitk else p.splitk, &.{ w.dw, w.ds, w.db, b.normed }); const prm = [4]i32{ WIDE, @intCast(nd), @intCast(rows), @intCast(parts) }; p.r.enc.setBytes(std.mem.asBytes(&prm), 4); p.r.enc.setBuffer(b.part.b, b.part.off, 5); @@ -2876,7 +2846,9 @@ pub const Prompt = struct { tabs[0] = b.pids; tabs[1] = pp.starts; for (0..3 * GROUPS) |j| tabs[2 + j] = pp.tables[j]; - try r.callRows("qa_ple_lookup@ple", rows, &tabs, &.{b.emb}, null); + var cp: [2]variants.Const = undefined; + try r.callRows("qa_ple_lookup@ple", rows, &tabs, &.{b.emb}, try variants.oneRow(r, "qa_ple_lookup@ple", variants.width(&cp, "BITS", "GS", pp.w))); + if (pp.scale) |sc| sc.apply(r.enc, b.emb, rows * D, r.serial); p.qmm(b.emb, p.ple_kv, D, 12800, rows, b.kvp, 0); try r.callRows("q4_ple_gate@ple", rows, &.{ b.kvp, b.h[s.cur], pp.ks, pp.qs, pp.cs, t.eps }, &.{ b.gated, .{ .b = b.cin.b, .off = PLE_TAIL * WIDE * 2 } }, null); try r.callRows("q4_ple_conv@ple", rows, &.{ b.cin, pp.conv, b.gated, b.h[s.cur] }, &.{b.hout}, null); @@ -3060,7 +3032,9 @@ pub const Prompt = struct { const s = h.at(l); const b = &s.p.b; const m = h.m; - try s.p.r.callRows("qa_embed_rows@embed", s.rows, &.{ b.ids, m.embed[0], m.embed[1], m.embed[2] }, &.{b.h[0]}, null); + var ce: [2]variants.Const = undefined; + const ev = try variants.oneRow(s.p.r, "qa_embed_rows@embed", variants.width(&ce, "BITS", "GS", m.embed_w)); + try s.p.r.callRows("qa_embed_rows@embed", s.rows, &.{ b.ids, m.embed[0], m.embed[1], m.embed[2] }, &.{b.h[0]}, ev); } pub fn wait(_: *Hooks, i: usize) segments.Wait { return if (i == 1) .pre else .mixer; @@ -3151,6 +3125,7 @@ pub const Prompt = struct { @memcpy(ps[0].b.cin.b.contents()[0 .. PLE_TAIL * WIDE * 2], cin_old[0 .. PLE_TAIL * WIDE * 2]); var hooks: Hooks = .{ .m = m, .segs = segs[0..N] }; m.gpu_seconds += try segments.run(r.device, queues[0..N], LAYERS, if (r.serial) .serial else .concurrent, &hooks); + if (ps[0].wide) |wd| try wd.check(); // a weight the load left out of the width table fails the pass const sl = &segs[N - 1]; m.state = if (N % 2 == 0) a else 1 - a; // each segment flipped it once m.state_row = 0; @@ -3294,21 +3269,23 @@ pub const Prompt = struct { b.n_add.b.slice(u32, 1)[0] = @intCast(n * WIDE); const cb = r.queue.commandBuffer(); r.enc = cb.compute(if (r.serial) .serial else .concurrent); - try r.callRows("mtp:qa_embed_rows@embed", n, &.{ b.mids, m.embed[0], m.embed[1], m.embed[2] }, &.{b.emb}, null); + var ce: [2]variants.Const = undefined; + try r.callRows("mtp:qa_embed_rows@embed", n, &.{ b.mids, m.embed[0], m.embed[1], m.embed[2] }, &.{b.emb}, try variants.oneRow(r, "mtp:qa_embed_rows@embed", variants.width(&ce, "BITS", "GS", m.embed_w))); try r.callRows("mtp:q4_rms_rows@mtp.enorm", n, &.{ b.emb, h.enorm, t.eps }, &.{b.branch}, null); - if (p.mtp_w) |w| p.qmm(b.branch, w[1], D, D, n, b.aout, 0) else r.denseRows(b.branch, D, h.fce, n, b.aout); + if (p.mtp_w) |w| p.qmm(b.branch, w[1], D, D, n, b.aout, 0) else r.denseRows(b.branch, D, try h.fce.recorded(), n, b.aout); try r.callRows("mtp:q4_rms_rows@mtp.hnorm", n, &.{ streams, h.hnorm, t.eps }, &.{b.normed}, null); - if (p.mtp_w) |w| p.qmm(b.normed, w[2], D, D, 4 * n, b.gated, 0) else r.denseRows(b.normed, D, h.fch, 4 * n, b.gated); + if (p.mtp_w) |w| p.qmm(b.normed, w[2], D, D, 4 * n, b.gated, 0) else r.denseRows(b.normed, D, try h.fch.recorded(), 4 * n, b.gated); p.bind(r.add_pipe, &.{ b.aout, b.gated, b.xs, b.n_add }); r.enc.dispatchThreads(mtl.Size.of(n * WIDE, 1, 1), mtl.Size.of(256, 1, 1)); p.barrier(); try r.callRows("mtp:q4_hc_norm_none#[10240]", n, &.{b.xs}, &.{ b.hout, b.ssp }, null); p.hc(m, b.hout, h.ahc, n, b.inj_a); - if (p.mtp_w) |w| p.qmm(b.mixed, w[0], D, 13952, n, b.p, 0) else r.denseRows(b.mixed, D, h.proj, n, b.p); + if (p.mtp_w) |w| p.qmm(b.mixed, w[0], D, 13952, n, b.p, 0) else r.denseRows(b.mixed, D, try h.proj.recorded(), n, b.p); try r.callRows("mtp:q4_attn_prep@mtp.att", n, &.{ b.p, b.pos, h.qn, h.kn, h.iqn, t.eps, t.log2base }, &.{ b.q, b.kout, b.iq }, null); p.bind(r.kv_pipe, &.{ b.kout, b.p, h.keys, h.vals, h.raw, b.kvmeta }); r.enc.dispatchThreads(mtl.Size.of(512 * n, 1, 1), mtl.Size.of(256, 1, 1)); try m.finish(cb); + if (p.wide) |wd| try wd.check(); } }; diff --git a/zig/src/families/flashnext/variants.zig b/zig/src/families/flashnext/variants.zig new file mode 100644 index 0000000000..5d0390f5c8 --- /dev/null +++ b/zig/src/families/flashnext/variants.zig @@ -0,0 +1,110 @@ +//! Recorded any-width kernels (0.6.6's qa_ texts) at another width: the same text with its constants swapped. +const std = @import("std"); +const mtl = @import("metal"); +const replay = @import("replay.zig"); +const Run = replay.Run; +const Buf = replay.Buf; +const Variant = replay.Variant; +const Hc = replay.Hc; + +/// A recorded `constexpr int` and the value it takes. +pub const Const = struct { name: []const u8, value: usize }; + +/// `text` with `constexpr int name = ...;` set to `value`; the declaration must appear exactly once. +pub fn constant(a: std.mem.Allocator, text: []const u8, name: []const u8, value: usize) ![]const u8 { + const key = try std.fmt.allocPrint(a, "constexpr int {s} = ", .{name}); + if (std.mem.count(u8, text, key) != 1) return error.KernelPatch; + const at = std.mem.indexOf(u8, text, key).? + key.len; + const end = at + (std.mem.indexOfScalar(u8, text[at..], ';') orelse return error.KernelPatch); + return std.fmt.allocPrint(a, "{s}{d}{s}", .{ text[0..at], value, text[end..] }); +} + +/// A pack table of u32s read into host memory: read on the host only, so never a device buffer. +pub fn table(r: *Run, name: []const u8) ![]const u32 { + const e = try r.entry(name); + if (e.len % 4 != 0) return error.BadPackTable; + const out = try r.arena.alloc(u32, e.len / 4); + try replay.readAll(e.fd, std.mem.sliceAsBytes(out), e.at); + return out; +} + +/// The recorded variant `v` at `consts`, compiled once and kept beside the recorded variants under its own key. +pub fn patched(r: *Run, v: *Variant, consts: []const Const) !*Variant { + var key: std.ArrayList(u8) = .empty; + try key.appendSlice(r.arena, v.name); + for (consts) |c| try key.print(r.arena, "|{s}={d}", .{ c.name, c.value }); + if (r.variants.get(key.items)) |hit| return hit; + var text = try Run.variantText(r.arena, v); + for (consts) |c| text = try constant(r.arena, text, c.name, c.value); + const lib = try mtl.Library.fromSource(r.device, text, mtl.CompileOptions.mlx()); + const nv = try r.arena.create(Variant); + nv.* = v.*; + nv.pipe = try mtl.Pipeline.init(r.device, lib, try std.fmt.allocPrintSentinel(r.arena, "{s}", .{v.name}, 0), false); + nv.text = text; + try r.variants.put(r.arena, key.items, nv); + return nv; +} + +/// `role` launched as recorded at `as_rows` rows on its kernel at `consts` (Run.callAs when there are none). +pub fn callAs(r: *Run, role: []const u8, as_rows: usize, ins: []const Buf, outs: []const Buf, consts: []const Const) !void { + if (consts.len == 0) return r.callAs(role, as_rows, ins, outs); + if (r.skip & Run.class(role) != 0) return; + if (r.hcskip != 0) for ([_][]const u8{ "hc_norm", "hc_down", "hc_up" }, 0..) |n, i| { + if (r.hcskip & (@as(u32, 1) << @intCast(i)) != 0 and std.mem.indexOf(u8, role, n) != null) return; + }; + var name: [96]u8 = undefined; + const s = r.roles.get(try std.fmt.bufPrint(&name, "{s}|{d}", .{ role, as_rows })) orelse return error.NoSite; + try r.bindV(try patched(r, s.v, consts), ins, outs); + r.enc.dispatchThreads(s.grid, s.tg); + if (!r.serial) r.enc.barrier(); +} + +/// The one-row variant of `role` at `consts`, as callRows takes an override; null for the recorded width. +pub fn oneRow(r: *Run, role: []const u8, consts: []const Const) !?*Variant { + if (consts.len == 0) return null; + var name: [96]u8 = undefined; + const s = r.roles.get(try std.fmt.bufPrint(&name, "{s}|1", .{role})) orelse return error.NoSite; + return try patched(r, s.v, consts); +} + +/// A width's two constants under the names a kernel declares them; none for the recorded 6/32 (null). +pub fn width(buf: *[2]Const, bits: []const u8, group: []const u8, w: ?[2]usize) []const Const { + const v = w orelse return &.{}; + buf.* = .{ .{ .name = bits, .value = v[0] }, .{ .name = group, .value = v[1] } }; + return buf; +} + +/// Expert widths as the expert kernels name them: routed (WB, WG) and shared (SWB, SWG), gate/up and down apart. +pub const Experts = struct { gu: [4]usize, down: [4]usize }; + +pub fn four(buf: *[4]Const, w: [4]usize) []const Const { + buf.* = .{ .{ .name = "WB", .value = w[0] }, .{ .name = "WG", .value = w[1] }, .{ .name = "SWB", .value = w[2] }, .{ .name = "SWG", .value = w[3] } }; + return buf; +} + +/// The routed and shared experts on the recorded any-width kernels at `w` (Run.experts' recorded path). +pub fn experts(r: *Run, gu_role: []const u8, down_role: []const u8, x: Buf, lg: Buf, e: [18]Buf, act: Buf, pick: Buf, wts: Buf, rows_buf: Buf, y: Buf, w: Experts) !void { + if (r.skip & (1 << 2) != 0) return; + var cg: [4]Const = undefined; + var cd: [4]Const = undefined; + try callAs(r, gu_role, r.rows, &.{ x, lg, e[0], e[1], e[2], e[3], e[4], e[5], e[6], e[7], e[8], e[9], e[10], e[11] }, &.{ act, pick, wts }, four(&cg, w.gu)); + try callAs(r, down_role, r.rows, &.{ act, pick, e[12], e[13], e[14], e[15], e[16], e[17], rows_buf }, &.{y}, four(&cd, w.down)); +} + +/// A pack hyper-connection: its seven tensors, and its `.widths` table (down, up) when not 6/32. +pub fn hcOf(r: *Run, comptime fmt: []const u8, args: anytype) !Hc { + var name: [96]u8 = undefined; + const stem = try std.fmt.bufPrint(&name, fmt, args); + var full: [128]u8 = undefined; + const parts = [_][]const u8{ "scale", "down.w", "down.s", "down.b", "up.w", "up.s", "up.b" }; + var out: [7]Buf = undefined; + for (parts, 0..) |p, i| out[i] = try r.load(try std.fmt.bufPrint(&full, "{s}.{s}", .{ stem, p })); + var hc: Hc = .{ .scale = out[0], .dw = out[1], .ds = out[2], .db = out[3], .uw = out[4], .us = out[5], .ub = out[6] }; + const widths = try std.fmt.bufPrint(&full, "{s}.widths", .{stem}); // any width but 6/32 (pack.putHc) + if (r.index.contains(widths)) { + const v = try table(r, widths); + if (v.len != 4) return error.BadPackTable; + hc.w = .{ .{ v[0], v[1] }, .{ v[2], v[3] } }; + } + return hc; +} diff --git a/zig/src/families/flashnext/widths.zig b/zig/src/families/flashnext/widths.zig new file mode 100644 index 0000000000..f74f69a4e6 --- /dev/null +++ b/zig/src/families/flashnext/widths.zig @@ -0,0 +1,76 @@ +//! Per-module MLX widths through core/quant, which indexes overrides by module; `false` and `{}` mark plain modules. +const std = @import("std"); +const mlx = @import("../../core/quant/mlx.zig"); +const st = @import("../../core/safetensors.zig"); +const affine = @import("affine.zig"); +const config = @import("config.zig"); + +pub const Width = mlx.Width; + +/// The widths the device kernels read: MLX affine 2, 3, 4, 5, 6 or 8 bits in groups of 32, 64 or 128. +pub fn launched(w: Width) bool { + const bits_ok = w.bits >= 2 and w.bits <= 8 and w.bits != 7; + return bits_ok and (w.group == 32 or w.group == 64 or w.group == 128); +} + +pub const Widths = struct { + arena: std.heap.ArenaAllocator, + /// Overrides by config key and by canonical module (`modules`); null is a width core/quant refuses (no guessing). + q: mlx.Config, + /// Modules the config declares unquantized (`false` or `{}`). + plain: std.StringHashMapUnmanaged(void) = .empty, + + pub fn init(gpa: std.mem.Allocator, c: *const config.Config) !Widths { + var w: Widths = .{ .arena = .init(gpa), .q = undefined }; + errdefer w.arena.deinit(); + const a = w.arena.allocator(); + // core/quant indexes overrides by module and refuses spellings that disagree + w.q = (mlx.detect(a, .{ .root = c.parsed.value.object, .quantization = c.quant }) catch |err| return switch (err) { + error.UnsupportedQuantization => error.ConflictingAffineAlias, // the only table core/quant refuses + else => err, + }) orelse return error.UnsupportedAffineMode; + var it = c.quant.iterator(); + while (it.next()) |item| { + const plain = switch (item.value_ptr.*) { + .bool => |on| !on, + .object => |o| o.count() == 0, + else => continue, + }; + const key = config.canonical(item.key_ptr.*); + if (!plain) { + if (w.plain.contains(key)) return error.ConflictingAffineAlias; + continue; + } + if (w.q.modules.get(key)) |declared| if (declared != null) return error.ConflictingAffineAlias; + try w.plain.put(a, key, {}); + } + return w; + } + + pub fn deinit(w: *Widths) void { + w.arena.deinit(); + w.* = undefined; + } + + /// No override at all: every quantized matrix has the global width. + pub fn uniform(w: *const Widths) bool { + return w.q.overrides.count() == 0 and w.plain.count() == 0; + } + + /// The module's width under any spelling of its key, or null when the config declares it unquantized. + pub fn of(w: *const Widths, path: []const u8) !?Width { + const key = config.canonical(path); + if (w.plain.contains(key)) return null; + return try w.q.width(key); + } + + /// The module's stored matrix checked against its own width, never one inferred from the shapes. + pub fn shaped(w: *const Widths, path: []const u8, weight: st.Entry, scales: st.Entry, biases: st.Entry) !affine.Shape { + const width = (try w.of(path)) orelse return error.UnquantizedAffineTensor; + return affine.matrix(weight, scales, biases, try affine.Spec.init(width.bits, width.group)); + } +}; + +test { + _ = @import("widths_test.zig"); +} diff --git a/zig/src/families/flashnext/widths_test.zig b/zig/src/families/flashnext/widths_test.zig new file mode 100644 index 0000000000..fa32e5004c --- /dev/null +++ b/zig/src/families/flashnext/widths_test.zig @@ -0,0 +1,85 @@ +//! Per-module widths: one module under any key spelling, its own width or a refusal, never a guess from shapes. +const std = @import("std"); +const st = @import("../../core/safetensors.zig"); +const config = @import("config.zig"); +const widths = @import("widths.zig"); +const a = std.testing.allocator; +const checked_config = @embedFile("fixtures/config.json"); + +fn overridden(entries: []const u8) ![]u8 { + const declaration = try std.fmt.allocPrint(a, "\"mode\": \"affine\", {s}", .{entries}); + defer a.free(declaration); + return std.mem.replaceOwned(u8, a, checked_config, "\"mode\": \"affine\"", declaration); +} + +const mixed = + \\"language_model.model.layers.0.self_attn.v_proj": {"bits": 8, "group_size": 64}, + \\"model.layers.1.linear_attn.in_proj_z.weight": {"bits": 5, "group_size": 128}, + \\"language_model.model.layers.0.mlp.gate": false, + \\"model.layers.2.self_attn.o_proj": {"bits": 4} +; + +fn expectWidth(w: *const widths.Widths, path: []const u8, bits: u8, group: u16) !void { + const got = (try w.of(path)).?; + try std.testing.expectEqual(bits, got.bits); + try std.testing.expectEqual(group, got.group); +} + +test "a mixed-override config resolves each module under either prefix and never fills a missing field" { + const bytes = try overridden(mixed); + defer a.free(bytes); + var c = try config.parse(a, bytes); + defer c.deinit(); + var w = try widths.Widths.init(a, &c); + defer w.deinit(); + try std.testing.expect(!w.uniform()); + try expectWidth(&w, "model.layers.0.self_attn.v_proj", 8, 64); + try expectWidth(&w, "language_model.model.layers.0.self_attn.v_proj.weight", 8, 64); + try expectWidth(&w, "language_model.model.layers.1.linear_attn.in_proj_z", 5, 128); + try expectWidth(&w, "language_model.model.layers.1.linear_attn.in_proj_qkv", 6, 32); + try std.testing.expect((try w.of("model.layers.0.mlp.gate")) == null); + // the family's own reader fills 4-bit group 64 here; core/quant refuses the half-declared width + try std.testing.expectError(error.UnsupportedQuantization, w.of("language_model.model.layers.2.self_attn.o_proj")); +} + +test "two spellings of one module that disagree are refused" { + const bytes = try overridden( + \\"language_model.model.layers.0.self_attn.v_proj": {"bits": 8, "group_size": 64}, + \\"model.layers.0.self_attn.v_proj.weight": {"bits": 4, "group_size": 64} + ); + defer a.free(bytes); + var c = try config.parse(a, bytes); + defer c.deinit(); + try std.testing.expectError(error.ConflictingAffineAlias, widths.Widths.init(a, &c)); +} + +fn entry(dtype: st.DType, rows: usize, cols: usize) st.Entry { + var shape: [st.max_rank]usize = @splat(0); + shape[0] = rows; + shape[1] = cols; + return .{ .dtype = dtype, .rank = 2, .shape = shape, .begin = 0, .end = rows * cols * dtype.size() }; +} + +/// A stored [rows, k] matrix's three entries at `bits`/`group`. +fn stored(rows: usize, k: usize, bits: usize, group: usize) [3]st.Entry { + return .{ entry(.u32, rows, k * bits / 32), entry(.bf16, rows, k / group), entry(.bf16, rows, k / group) }; +} + +test "a tensor packed at its override width is admitted and one at a guessed width is refused" { + const bytes = try overridden(mixed); + defer a.free(bytes); + var c = try config.parse(a, bytes); + defer c.deinit(); + var w = try widths.Widths.init(a, &c); + defer w.deinit(); + const v = "model.layers.0.self_attn.v_proj"; // its override is spelled under language_model. + const at_override = stored(8, 128, 8, 64); + const shape = try w.shaped(v, at_override[0], at_override[1], at_override[2]); + try std.testing.expectEqual(@as(usize, 128), shape.k); + try std.testing.expectEqual(@as(usize, 2), shape.groups); + const at_global = stored(8, 128, 6, 32); // the global width, which a prefix-blind lookup would guess + try std.testing.expectError(error.InvalidAffineShape, w.shaped(v, at_global[0], at_global[1], at_global[2])); + // a module with no override is read at the global width only, whatever width its shapes would fit + try std.testing.expectError(error.InvalidAffineShape, w.shaped("model.layers.0.self_attn.k_proj", at_override[0], at_override[1], at_override[2])); + try std.testing.expectError(error.UnquantizedAffineTensor, w.shaped("model.layers.0.mlp.gate", at_global[0], at_global[1], at_global[2])); +} diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 1823eb56a9..f34f41a4ff 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -13,6 +13,7 @@ pub const ids_json = @import("core/ids_json.zig"); pub const lanes = @import("lanes"); pub const segments = @import("core/segments.zig"); pub const row_projection = @import("core/row_projection.zig"); +pub const affine_mm = @import("core/affine_mm.zig"); pub const exl3_format = @import("core/exl3_format.zig"); pub const exl3_projection = @import("core/exl3_projection.zig"); pub const exlr8 = @import("core/exlr8.zig"); diff --git a/zig/tests/affine_ref.zig b/zig/tests/affine_ref.zig new file mode 100644 index 0000000000..5f32bd864a --- /dev/null +++ b/zig/tests/affine_ref.zig @@ -0,0 +1,358 @@ +//! The shared prompt products against float64 on synthetic weights, and every row's bits in every window width. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const mm = tf.affine_mm; +const sq = tf.row_projection.simd; + +const opts = mtl.ResourceOptions.shared; +const a = std.heap.page_allocator; +const unit: f64 = 0x1p-24; // fp32's unit roundoff +const wide = 129; // rows of the reference call: two 64-row tiles and one more +const windows = [_]usize{ 1, 2, 3, 8, 16, 17, 33, 64, 65 }; // 17, 33 and 65 cross the 16-, 32- and 64-row tiles + +fn toBf(v: f32) u16 { + const u: u32 = @bitCast(v); + return @truncate((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); +} + +fn fromBf(v: u16) f64 { + return @as(f32, @bitCast(@as(u32, v) << 16)); +} + +/// Synthetic MLX affine weights [n, k] (codes little-endian in bytes) and `wide` rows of x; `deq` is s q + b exactly. +const Synth = struct { + n: usize, + k: usize, + bits: u8, + group: usize, + x: mtl.Buffer, + w: mtl.Buffer, + s: mtl.Buffer, + b: mtl.Buffer, + deq: []f64, + + fn init(device: mtl.Device, n: usize, k: usize, bits: u8, group: usize, seed: u64) !Synth { + var prng = std.Random.DefaultPrng.init(seed); + const r = prng.random(); + const sn = n * (k / group); + const self = Synth{ .n = n, .k = k, .bits = bits, .group = group, .x = try device.buffer(wide * k * 2, opts), .w = try device.buffer(n * k * bits / 8, opts), .s = try device.buffer(sn * 2, opts), .b = try device.buffer(sn * 2, opts), .deq = try a.alloc(f64, n * k) }; + for (self.x.slice(u16, wide * k)) |*v| v.* = toBf(r.float(f32) * 2 - 1); + r.bytes(self.w.slice(u8, n * k * bits / 8)); + for (self.s.slice(u16, sn), self.b.slice(u16, sn)) |*s, *b| { + s.* = toBf((r.float(f32) + 0.5) / 64); + b.* = toBf(-(r.float(f32) + 0.5) * 0.15); + } + for (0..n) |j| for (0..k) |i| { + const g = j * (k / group) + i / group; + self.deq[j * k + i] = fromBf(self.s.slice(u16, sn)[g]) * @as(f64, @floatFromInt(self.code(j, i))) + fromBf(self.b.slice(u16, sn)[g]); + }; + return self; + } + + fn deinit(self: Synth) void { + for ([_]mtl.Buffer{ self.x, self.w, self.s, self.b }) |x| x.deinit(); + a.free(self.deq); + } + + fn code(self: Synth, j: usize, i: usize) u32 { + const bytes = self.w.slice(u8, self.n * self.k * self.bits / 8); + if (self.bits == 8) return bytes[j * self.k + i]; + if (self.bits == 2) return (bytes[(j * self.k + i) / 4] >> @intCast(2 * (i % 4))) & 3; + return (bytes[(j * self.k + i) / 2] >> @intCast(4 * (i % 2))) & 15; + } + + /// y = x W^T in float64, and the sizes the fp32 sums carry when codes enter as q + off (affine_mm's 4-bit: 128). + fn reference(self: Synth, off: f64) !Ref { + const n = self.n; + const k = self.k; + const groups = k / self.group; + const ref = Ref{ .want = try a.alloc(f64, wide * n), .mag = try a.alloc(f64, wide * n), .scaled = try a.alloc(f64, wide * n) }; + const x = self.x.slice(u16, wide * k); + const sc = self.s.slice(u16, n * groups); + const bi = self.b.slice(u16, n * groups); + const xd = try a.alloc(f64, k); + defer a.free(xd); + const xg = try a.alloc(f64, groups); // each group's sum of |x| + defer a.free(xg); + for (0..wide) |r| { + @memset(xg, 0); + for (xd, 0..) |*v, i| { + v.* = fromBf(x[r * k + i]); + xg[i / self.group] += @abs(v.*); + } + for (0..n) |j| { + var sum: f64 = 0; + var mag: f64 = 0; + var scaled: f64 = 0; + for (0..groups) |g| { + var qx: f64 = 0; + var qxs: f64 = 0; + for (g * self.group..(g + 1) * self.group) |i| { + const q: f64 = @floatFromInt(self.code(j, i)); + sum += xd[i] * self.deq[j * k + i]; + qx += (q + off) * @abs(xd[i]); + qxs += q * xd[i]; + } + const s = fromBf(sc[j * groups + g]); + mag += s * qx + @abs(fromBf(bi[j * groups + g]) - off * s) * xg[g]; + scaled += s * qxs; + } + ref.want[r * n + j] = sum; + ref.mag[r * n + j] = mag; + ref.scaled[r * n + j] = scaled; + } + } + return ref; + } +}; + +/// The float64 product, its terms' fp32 magnitudes, and its scales' share (a control: those scales 1% off must fail). +const Ref = struct { + want: []f64, + mag: []f64, + scaled: []f64, + fn deinit(r: Ref) void { + for ([_][]f64{ r.want, r.mag, r.scaled }) |v| a.free(v); + } +}; + +/// Every check's outcome: failures, and the largest difference as a fraction of its bound. +const Tally = struct { failed: usize = 0, worst: f64 = 0 }; + +/// One product's output: `rows` rows of n bf16 or fp32 values. +const Out = struct { + buf: mtl.Buffer, + f32_out: bool, + fn esize(o: Out) usize { + return if (o.f32_out) 4 else 2; + } + fn get(o: Out, i: usize) f64 { + return if (o.f32_out) o.buf.slice(f32, i + 1)[i] else fromBf(o.buf.slice(u16, i + 1)[i]); + } +}; + +/// The bound: fp32 sums of a group's products, its row sum and a chain of groups, plus one bf16 ulp of the value. +fn bound(want: f64, mag: f64, s: *const Synth, f32_out: bool, extra: usize) f64 { + const terms: f64 = @floatFromInt(s.group + s.k / s.group + 8 + extra); + return terms * unit * mag + (if (f32_out) 0 else @abs(want) * 0x1p-7) + 1e-30; +} + +/// `rows` rows of `out` against the float64 reference: the worst column's max abs and relative error, and the bound. +fn compare(t: *Tally, name: []const u8, out: Out, rows: usize, ref: Ref, s: *const Synth, extra: usize) void { + const n = s.n; + var col: usize = 0; + var col_abs: f64 = -1; + var col_rel: f64 = 0; + var frac: f64 = 0; + var ctl: f64 = 0; + var at: usize = 0; + for (0..n) |j| { + var abs: f64 = 0; + var top: f64 = 0; + for (0..rows) |r| { + const i = r * n + j; + const d = @abs(out.get(i) - ref.want[i]); + const tol = bound(ref.want[i], ref.mag[i], s, out.f32_out, extra); + const q = d / tol; + ctl = @max(ctl, @abs(out.get(i) - ref.want[i] - 0.01 * ref.scaled[i]) / tol); + if (!(d <= abs)) abs = d; + top = @max(top, @abs(ref.want[i])); + if (!(q <= frac)) { + frac = q; + at = i; + } + } + if (!(abs <= col_abs)) { + col_abs = abs; + col = j; + col_rel = abs / top; + } + } + t.worst = @max(t.worst, frac); + const ok = frac <= 1 and ctl > 1; + if (!ok) t.failed += 1; + std.debug.print("RESULT {s}: worst column {d}: max abs {e:.3}, relative {e:.3}; {d:.4} of its bound (row {d} col {d}), " ++ + "scales 1% off {d:.1}{s}\n", .{ name, col, col_abs, col_rel, frac, at / n, at % n, ctl, if (ok) "" else " FAIL" }); +} + +/// Rows [0, rows) of `got` equal `want` byte for byte; the first difference printed. +fn sameRows(t: *Tally, name: []const u8, m: usize, want: Out, got: Out, rows: usize, n: usize) bool { + const bytes = rows * n * want.esize(); + const w = want.buf.slice(u8, bytes); + const g = got.buf.slice(u8, bytes); + if (std.mem.eql(u8, w, g)) return true; + const i = std.mem.indexOfDiff(u8, w, g).? / want.esize(); + var differ: usize = 0; + var most: f64 = 0; + for (0..rows) |r| { + const row = r * n * want.esize(); + if (!std.mem.eql(u8, w[row..][0 .. n * want.esize()], g[row..][0 .. n * want.esize()])) differ += 1; + for (r * n..(r + 1) * n) |e| most = @max(most, @abs(got.get(e) - want.get(e))); + } + std.debug.print("RESULT {s}: {d}-row windows DIFFER in {d} of {d} rows, by up to {e:.3}; first row {d} col {d}: {d} against {d} in the {d}-row call\n", .{ name, m, differ, rows, most, i / n, i % n, got.get(i), want.get(i), wide }); + t.failed += 1; + return false; +} + +fn newOut(device: mtl.Device, s: *const Synth, f32_out: bool) !Out { + const o = Out{ .buf = try device.buffer(wide * s.n * 4, opts), .f32_out = f32_out }; + @memset(o.buf.slice(u8, wide * s.n * 4), 0xff); // NaN: a launch that writes nothing fails + return o; +} + +fn commit(cb: anytype, e: mtl.ComputeEncoder) !void { + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |why| { + std.debug.print("GPU failed: {s}\n", .{why}); + return error.GpuFailed; + } +} + +const Kernel = enum { dense, gather }; + +/// affine_mm over rows [start, start + m): its row sums, then the dense product or a gather with one expert. +fn affineWindow(e: mtl.ComputeEncoder, p: mm.Pipes, s: *const Synth, sums: mtl.Buffer, offs: mtl.Buffer, y: Out, kernel: Kernel, start: usize, m: usize) void { + const x = .{ .buf = s.x, .off = start * s.k * 2 }; + const xs = .{ .buf = sums, .off = start * (s.k / s.group) * 4 }; + const yo = .{ .buf = y.buf, .off = start * s.n * y.esize() }; + const z: usize = 0; + const q = .{ .w = .{ .buf = s.w, .off = z }, .s = .{ .buf = s.s, .off = z }, .b = .{ .buf = s.b, .off = z }, .n = @as(u32, @intCast(s.n)), .k = @as(u32, @intCast(s.k)) }; + mm.rowSums(e, p, @intCast(s.group), x, xs, @intCast(m), @intCast(s.k)); + switch (kernel) { + .dense => mm.dense(e, p, x, xs, q, yo, @intCast(m)), + .gather => mm.gather(e, p, x, xs, q, .{ .buf = offs, .off = z }, yo, @intCast(m), 1), + } +} + +/// affine_mm at one format: the wide call against float64, then every window's rows against the wide call's. +fn affine(t: *Tally, device: mtl.Device, queue: mtl.Queue, s: *const Synth, ref: Ref, f32_out: bool) !void { + const src = try mm.source(device, a, .{ .bits = s.bits, .group = @intCast(s.group), .out_f32 = f32_out }); + defer a.free(src); + const lib = try mtl.Library.fromSource(device, src, mtl.CompileOptions.mlx()); + defer lib.deinit(); + var p: mm.Pipes = undefined; + for (mm.names, 0..) |name, i| p[i] = try mtl.Pipeline.init(device, lib, name, false); + defer for (p) |x| x.deinit(); + const sums = try device.buffer(wide * (s.k / s.group) * 4, opts); + defer sums.deinit(); + const offs = try device.buffer(4, opts); + defer offs.deinit(); + offs.slice(i32, 1)[0] = 0; + var name: [96]u8 = undefined; + for ([_]Kernel{ .dense, .gather }) |kernel| { + const label = try std.fmt.bufPrint(&name, "affine_mm {s} q{d} g{d} {s} [{d}, {d}]", .{ @tagName(kernel), s.bits, s.group, if (f32_out) "fp32" else "bf16", s.n, s.k }); + const all = try newOut(device, s, f32_out); + defer all.buf.deinit(); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + affineWindow(e, p, s, sums, offs, all, kernel, 0, wide); + try commit(cb, e); + compare(t, label, all, wide, ref, s, 0); + var same = true; + for (windows) |m| { + const got = try newOut(device, s, f32_out); + defer got.buf.deinit(); + const wb = queue.commandBuffer(); + const we = wb.compute(.serial); + var start: usize = 0; + while (start + m <= wide) : (start += m) affineWindow(we, p, s, sums, offs, got, kernel, start, m); + try commit(wb, we); + same = sameRows(t, label, m, all, got, wide / m * m, s.n) and same; + } + if (same) std.debug.print("RESULT {s}: every row's bits equal in windows of 1, 2, 3, 8, 16, 17, 33, 64 and 65 rows\n", .{label}); + } +} + +/// simd_qmm over rows [start, start + m), the scalar twin when `scalar` and it takes m rows. +fn simdWindow(e: mtl.ComputeEncoder, p: *const sq.Pipelines, s: *const Synth, y: Out, scalar: bool, start: usize, m: usize) !void { + const mat = sq.Matrix{ .w = s.w, .scales = s.s, .biases = s.b, .n = s.n, .k = s.k }; + sq.encode(e, try sq.call(p, s.n, s.k, m, scalar), mat, s.x, start * s.k * 2, y.buf, start * s.n * 2); +} + +/// simd_qmm (4-bit, group 64): tiles at every window against float64 and the wide call; the twin at 1 and 2 rows. +fn simd(t: *Tally, device: mtl.Device, queue: mtl.Queue, p: *const sq.Pipelines, s: *const Synth, ref: Ref) !void { + var name: [96]u8 = undefined; + const label = try std.fmt.bufPrint(&name, "simd_qmm tiles q4 g64 bf16 [{d}, {d}]", .{ s.n, s.k }); + const all = try newOut(device, s, false); + defer all.buf.deinit(); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + try simdWindow(e, p, s, all, false, 0, wide); + try commit(cb, e); + const extra = std.math.log2_int(u32, sq.splits(s.n)) + 1; // the chunks' fixed tree + compare(t, label, all, wide, ref, s, extra); + var same = true; + for (windows) |m| { + const got = try newOut(device, s, false); + defer got.buf.deinit(); + const wb = queue.commandBuffer(); + const we = wb.compute(.serial); + var start: usize = 0; + while (start + m <= wide) : (start += m) try simdWindow(we, p, s, got, false, start, m); + try commit(wb, we); + same = sameRows(t, label, m, all, got, wide / m * m, s.n) and same; + } + if (same) std.debug.print("RESULT {s}: every row's bits equal in windows of 1, 2, 3, 8, 16, 17, 33, 64 and 65 rows\n", .{label}); + // the twin: callers take it only where twinExact says it equals the tiles, so a difference fails only then + const exact = try sq.twinExact(p, device, .{ .w = s.w, .scales = s.s, .biases = s.b, .n = s.n, .k = s.k }); + for ([_]usize{ 1, 2 }) |m| { + const twin = try std.fmt.bufPrint(&name, "simd_qmm twin q4 g64 bf16 [{d}, {d}] at {d} rows", .{ s.n, s.k, m }); + const got = try newOut(device, s, false); + defer got.buf.deinit(); + const wb = queue.commandBuffer(); + const we = wb.compute(.serial); + var start: usize = 0; + while (start + m <= wide) : (start += m) try simdWindow(we, p, s, got, true, start, m); + try commit(wb, we); + compare(t, twin, got, wide / m * m, ref, s, 0); + const rows = wide / m * m; + const eq = std.mem.eql(u8, all.buf.slice(u8, rows * s.n * 2), got.buf.slice(u8, rows * s.n * 2)); + std.debug.print("RESULT {s}: {s} the tiles' bits; twinExact says {s}{s}\n", .{ twin, if (eq) "equals" else "differs from", if (exact) "equal" else "differ", if (exact and !eq) " FAIL" else "" }); + if (exact and !eq) t.failed += 1; + } +} + +pub fn main(_: std.process.Init) !void { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + std.debug.print("device {s}, tensor units {}\n", .{ device.name(), device.tensorUnits() }); + var t = Tally{}; + var seed: u64 = 0x716; + for ([_][2]usize{ .{ 192, 1024 }, .{ 256, 4096 } }) |shape| for ([_]u8{ 4, 8 }) |bits| for ([_]usize{ 32, 64 }) |group| { + seed += 1; + const s = try Synth.init(device, shape[0], shape[1], bits, group, seed); + defer s.deinit(); + const ref = try s.reference(if (bits == 4) 128 else 0); + defer ref.deinit(); + for ([_]bool{ false, true }) |f32_out| try affine(&t, device, queue, &s, ref, f32_out); + }; + for ([_][2]usize{ .{ 192, 1024 }, .{ 256, 4096 } }) |shape| for ([_]usize{ 32, 64 }) |group| { // 2 bits: oQ2's experts + const s = try Synth.init(device, shape[0], shape[1], 2, group, 0x2716 + shape[0] + group); + defer s.deinit(); + const ref = try s.reference(0); + defer ref.deinit(); + for ([_]bool{ false, true }) |f32_out| try affine(&t, device, queue, &s, ref, f32_out); + }; + var p = try sq.Pipelines.load(device); + defer p.deinit(); + for ([_][2]usize{ .{ 40, 512 }, .{ 256, 4096 }, .{ 6400, 256 } }) |shape| { // chunks of 32, 16 and 8 + seed += 1; + const s = try Synth.init(device, shape[0], shape[1], 4, 64, seed); + defer s.deinit(); + const ref = try s.reference(0); + defer ref.deinit(); + try simd(&t, device, queue, &p, &s, ref); + } + if (t.failed > 0) { + std.debug.print("FAIL {d} checks; the largest difference is {d:.4} of its float64 bound\n", .{ t.failed, t.worst }); + return error.CheckFailed; + } + std.debug.print("PASS every product within {d:.4} of its float64 bound, every row's bits equal at every window\n", .{t.worst}); +} 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/flashnext_bench.zig b/zig/tests/flashnext_bench.zig index a4d2b87894..7a0c73d528 100644 --- a/zig/tests/flashnext_bench.zig +++ b/zig/tests/flashnext_bench.zig @@ -30,6 +30,11 @@ const Mode = enum { one, tp, stream, new, tpnew }; /// One projection class: its role, input width, input and output buffers, and every layer's weights. const Pass = struct { name: []const u8, role: []const u8, k: usize, x: Buf, y: Buf, lanes: []const Lane }; +/// A lane's recorded 6/32 pair: the harness loads 6/32 packs only (flashnext_run refuses others). +fn six(l: Lane) Lane.Six { + return l.six.?; +} + const tp_in = [_][2]usize{ .{ 0, 32 }, .{ 64, 32 }, .{ 128, 96 }, .{ 320, 96 }, .{ 512, 3 } }; const Bench = struct { @@ -52,7 +57,7 @@ const Bench = struct { /// Bytes a pass reads (speed-up mode: rank 0's tiles or K half). fn bytes(p: Pass, mode: Mode) f64 { var n: usize = 0; - for (p.lanes) |l| n += (l.wq.b.length() - l.wq.off) + (l.sbt.b.length() - l.sbt.off); + for (p.lanes) |l| n += (six(l).wq.b.length() - six(l).wq.off) + (six(l).sbt.b.length() - six(l).sbt.off); const all: f64 = @floatFromInt(n); if (mode != .tp and mode != .tpnew) return all; if (std.mem.endsWith(u8, p.role, "gdn.in")) return all * 259.0 / 515.0; @@ -62,13 +67,13 @@ const Bench = struct { /// Where a pass writes (the speed-up out-projection's fp32 partial, or the class's output), and its bytes at `rows`. fn out(b: *Bench, p: Pass, mode: Mode, rows: usize) struct { buf: Buf, len: usize } { - const n = (p.lanes[0].wq.b.length() - p.lanes[0].wq.off) * 4 / (3 * p.k); + const n = (six(p.lanes[0]).wq.b.length() - six(p.lanes[0]).wq.off) * 4 / (3 * p.k); if ((mode == .tp or mode == .tpnew) and std.mem.endsWith(u8, p.role, "gdn.out")) return .{ .buf = b.part, .len = rows * fz.D * 4 }; return .{ .buf = p.y, .len = rows * n * 2 }; } fn layout(b: *Bench, p: Pass, l: Lane, tp: bool) !dense.Layout { - const n = (l.wq.b.length() - l.wq.off) * 4 / (3 * p.k); + const n = (six(l).wq.b.length() - six(l).wq.off) * 4 / (3 * p.k); const sk = (try dense.recorded(b.r, p.role)).sk; const pf = b.r.lane_pf; if (!tp) return .{ .n = n, .k = p.k, .sk = sk, .pf = pf }; @@ -83,18 +88,18 @@ const Bench = struct { const t = &b.m.t; switch (mode) { .one => try b.m.lane(p.x, p.k, l, p.role, p.y), - .new, .tpnew => try dense.project(r, try b.layout(p, l, mode == .tpnew), p.x, l, t.mdims, b.out(p, mode, 1).buf), + .new, .tpnew => try dense.project(r, try b.layout(p, l, mode == .tpnew), p.x, six(l), t.mdims, b.out(p, mode, 1).buf), .stream => { r.enc.setPipeline(b.stream); - r.enc.setBuffer(l.wq.b, l.wq.off, 0); - r.enc.setBuffer(l.sbt.b, l.sbt.off, 1); + r.enc.setBuffer(six(l).wq.b, six(l).wq.off, 0); + r.enc.setBuffer(six(l).sbt.b, six(l).sbt.off, 1); r.enc.setBuffer(b.sink.b, 0, 2); - const n = [2]u32{ @intCast((l.wq.b.length() - l.wq.off) / 16), @intCast((l.sbt.b.length() - l.sbt.off) / 16) }; + const n = [2]u32{ @intCast((six(l).wq.b.length() - six(l).wq.off) / 16), @intCast((six(l).sbt.b.length() - six(l).sbt.off) / 16) }; r.enc.setBytes(std.mem.asBytes(&n), 3); r.enc.dispatchThreads(mtl.Size.of(160 * 1024, 1, 1), mtl.Size.of(1024, 1, 1)); }, .tp => { - const ins = [_]Buf{ p.x, t.xs, l.wq, l.sbt, t.mdims }; + const ins = [_]Buf{ p.x, t.xs, six(l).wq, six(l).sbt, t.mdims }; if (std.mem.endsWith(u8, p.role, "gdn.in")) { try split.laneTiles(r, p.role, &ins, p.y, &tp_in, 8, null); } else if (std.mem.endsWith(u8, p.role, "gdn.out")) { @@ -136,7 +141,7 @@ const Bench = struct { for (p.lanes) |l| { const lay = try b.layout(p, l, tp); r.enc.setPipeline(pipe); - for ([_]Buf{ p.x, l.wq, l.sbt, b.m.t.mdims, b.out(p, if (tp) .tpnew else .new, 1).buf }, 0..) |bf, j| r.enc.setBuffer(bf.b, bf.off, j); + for ([_]Buf{ p.x, six(l).wq, six(l).sbt, b.m.t.mdims, b.out(p, if (tp) .tpnew else .new, 1).buf }, 0..) |bf, j| r.enc.setBuffer(bf.b, bf.off, j); r.enc.dispatchThreads(mtl.Size.of(dense.tiles(lay) * 32 * lay.sk, 1, 1), mtl.Size.of(32 * lay.sk, 1, 1)); } } diff --git a/zig/tests/flashnext_run.zig b/zig/tests/flashnext_run.zig index eca3fc964a..87efe89c96 100644 --- a/zig/tests/flashnext_run.zig +++ b/zig/tests/flashnext_run.zig @@ -29,7 +29,6 @@ const KW = fz.KW; const Select = fz.Select; const GSelect = fz.GSelect; const Hc = fz.Hc; -const Lane = fz.Lane; const Layer = fz.Layer; const hcOf = fz.hcOf; const laneOf = fz.laneOf; @@ -37,7 +36,6 @@ const Ple = fz.Ple; const Tmp = fz.Tmp; const f32Buf = fz.f32Buf; const i32Buf = fz.i32Buf; -const Slot = fz.Slot; const Mtp = fz.Mtp; const Model = fz.Model; const DepthRule = fz.DepthRule; @@ -286,6 +284,8 @@ pub fn main(init: std.process.Init) !void { } return; } + // the replay below reads every expert, embedding and table at 6/32: other widths are the engine's (Engine.load) + if (!try fz.engine_widths.uniformDir(arena, init.io, args[1])) return error.HarnessNeeds6Bit32; var r = Run{ .arena = arena, .device = device, .queue = try device.queue() }; r.fused_xsum = std.c.getenv("FZ_FUSED_XSUM") != null; r.serial = std.c.getenv("FZ_SERIAL") != null; diff --git a/zig/tests/flashnext_widths.zig b/zig/tests/flashnext_widths.zig new file mode 100644 index 0000000000..b7cdc2905c --- /dev/null +++ b/zig/tests/flashnext_widths.zig @@ -0,0 +1,36 @@ +//! tf-flashnext-widths: Flash Next's products at every oQ width against float64, each row's bits in every window. +const std = @import("std"); +const mtl = @import("metal"); +const lib = @import("flashnext_widths_lib.zig"); +const lane = @import("flashnext_widths_lane.zig"); +const hc = @import("flashnext_widths_hc.zig"); +const moe = @import("flashnext_widths_moe.zig"); +const prompt = @import("flashnext_widths_prompt.zig"); +const fz = @import("tensorfold").flashnext_replay; + +pub fn main(init: 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(); + const args = try init.minimal.args.toSlice(lib.a); + const only: ?[]const u8 = if (args.len > 1) args[1] else null; // a class name runs that class alone + std.debug.print("device {s}, tensor units {}\n", .{ device.name(), device.tensorUnits() }); + var t = lib.Tally{}; + if (only == null or std.mem.eql(u8, only.?, "lane")) try lane.run(&t, device, queue); + if (only == null or !std.mem.eql(u8, only.?, "lane")) { // the recorded kernels the any-width ones come from + const r = try lib.a.create(fz.Run); + r.* = .{ .arena = lib.a, .device = device, .queue = queue, .serial = true }; + try r.compileChecked(); + if (only == null or std.mem.eql(u8, only.?, "hc")) try hc.run(&t, r); + if (only == null or std.mem.eql(u8, only.?, "moe")) try moe.run(&t, r); + if (only == null or std.mem.eql(u8, only.?, "prompt")) try prompt.run(&t, r); + } + if (t.failed > 0) { + std.debug.print("FAIL {d} of {d} checks; the largest difference is {d:.4} of its float64 bound\n", .{ t.failed, t.checks, t.worst }); + return error.CheckFailed; + } + std.debug.print("PASS {d} checks: every product within {d:.4} of its float64 bound, every row's bits equal in every window\n", .{ t.checks, t.worst }); +} diff --git a/zig/tests/flashnext_widths_hc.zig b/zig/tests/flashnext_widths_hc.zig new file mode 100644 index 0000000000..7af69e1e43 --- /dev/null +++ b/zig/tests/flashnext_widths_hc.zig @@ -0,0 +1,238 @@ +//! tf-flashnext-widths, hyper-connections: the recorded any-width kernels at each oQ width, decode rows 1 to 16. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const lib = @import("flashnext_widths_lib.zig"); +const fz = tf.flashnext_replay; +const variants = fz.variants; +const Buf = fz.Buf; +const Run = fz.Run; +const a = lib.a; + +const W = 4 * 2560; // the streams' width: the down projection's K +const LOW = 320; +const counts = [_]usize{ 1, 2, 3, 4, 8, 9, 16 }; // the row kernels at 1-2 rows, the MMA kernels past them +const R = 16; + +/// The streams and their sums of squares set so every stream's norm is exactly one: x = bf16(hn) as stored. +const Streams = struct { hn: mtl.Buffer, hd: []f64, ssp: mtl.Buffer, nw: mtl.Buffer, eps: mtl.Buffer }; + +fn streams(device: mtl.Device) !Streams { + const in = try lib.Inputs.init(device, R, W, 0x4c); + const ssp = try device.buffer(R * 10 * 4 * 4, lib.opts); + @memset(ssp.slice(f32, R * 40), 256); // ten partial sums of 256 a stream: mean square 1, rsqrt 1 + const nw = try device.buffer(W * 4, lib.opts); + @memset(nw.slice(f32, W), 1); + const eps = try device.buffer(16, lib.opts); + eps.slice(f32, 1)[0] = 0; + return .{ .hn = in.buf, .hd = in.x, .ssp = ssp, .nw = nw, .eps = eps }; +} + +const Weights = struct { dw: Buf, ds: Buf, db: Buf, uw: Buf, us: Buf, ub: Buf }; + +fn upload(device: mtl.Device, down: lib.Mat, up: lib.Mat) !Weights { + return .{ + .dw = .{ .b = try lib.upload(device, down.words) }, + .ds = .{ .b = try lib.upload(device, std.mem.sliceAsBytes(down.scales)) }, + .db = .{ .b = try lib.upload(device, std.mem.sliceAsBytes(down.biases)) }, + .uw = .{ .b = try lib.upload(device, up.words) }, + .us = .{ .b = try lib.upload(device, std.mem.sliceAsBytes(up.scales)) }, + .ub = .{ .b = try lib.upload(device, std.mem.sliceAsBytes(up.biases)) }, + }; +} + +/// One run's outputs: the down partials and the up's mixed rows and inject gates. +const Out = struct { part: mtl.Buffer, mixed: mtl.Buffer, inj: mtl.Buffer }; + +fn outputs(device: mtl.Device, nd: usize) !Out { + return .{ .part = try lib.poisoned(device, 10 * R * nd * 4), .mixed = try lib.poisoned(device, R * 2560 * 2), .inj = try lib.poisoned(device, R * 4 * 2) }; +} + +/// Rows [start, start + rows) through the down and up roles at `consts` (none: the recorded 6/32 kernels). +fn project(r: *Run, s: Streams, w: Weights, which: []const u8, start: usize, rows: usize, nd: usize, o: Out, kd: []const variants.Const, ku: []const variants.Const) !void { + var role: [64]u8 = undefined; + const rb = try r.buffer(16); + rb.slice(i32, 1)[0] = @intCast(rows); + const rows_buf: Buf = .{ .b = rb }; + const hn: Buf = .{ .b = s.hn, .off = start * W * 2 }; + const ssp: Buf = .{ .b = s.ssp, .off = start * 40 * 4 }; + const part: Buf = .{ .b = o.part, .off = start * 10 * nd * 4 }; + const down = try std.fmt.bufPrint(&role, "qa_hc_down@{s}", .{which}); + try variants.callAs(r, down, rows, &.{ hn, ssp, .{ .b = s.nw }, w.dw, w.ds, w.db, .{ .b = s.eps }, rows_buf }, &.{part}, kd); + const up = try std.fmt.bufPrint(&role, "qa_hc_up@{s}", .{which}); + const ins = [_]Buf{ hn, ssp, .{ .b = s.nw }, part, w.uw, w.us, w.ub, .{ .b = s.eps }, rows_buf }; + try variants.callAs(r, up, rows, &ins, &.{ .{ .b = o.mixed, .off = start * 2560 * 2 }, .{ .b = o.inj, .off = start * 4 * 2 } }, ku); +} + +fn runAll(r: *Run, s: Streams, w: Weights, which: []const u8, rows: usize, nd: usize, o: Out, kd: []const variants.Const, ku: []const variants.Const) !void { + const cb = r.queue.commandBuffer(); + r.enc = cb.compute(.serial); + var start: usize = 0; + while (start + rows <= R) : (start += rows) try project(r, s, w, which, start, rows, nd, o, kd, ku); + try lib.commit(cb, r.enc); +} + +/// One HC at a down and up width: recorded 6/32 bits, float64, row bits; `full`: codes past 63, no 6/32 twin. +fn check(t: *lib.Tally, r: *Run, s: Streams, which: []const u8, nd: usize, dw: [2]usize, uw: [2]usize, seed: u64, full: bool) !void { + var label_buf: [128]u8 = undefined; + const label = try std.fmt.bufPrint(&label_buf, "hc {s} down q{d}/g{d} up q{d}/g{d}{s}", .{ which, dw[0], dw[1], uw[0], uw[1], if (full) " full-range codes" else "" }); + const top: u32 = if (full) 256 else 64; // below 64: codes the recorded 6/32 kernels read too + const down = try lib.Mat.initBelow(nd, W, @intCast(dw[0]), dw[1], seed, @min(top, @as(u32, 1) << @intCast(dw[0]))); + const up = try lib.Mat.initBelow(W, LOW, @intCast(uw[0]), uw[1], seed + 1, @min(top, @as(u32, 1) << @intCast(uw[0]))); + const w = try upload(r.device, down, up); + const w6 = if (full) w else try upload(r.device, try down.recode(6, 32), try up.recode(6, 32)); + var cd: [2]variants.Const = undefined; + var cu: [2]variants.Const = undefined; + const kd = variants.width(&cd, "BITS", "GS", dw); + const ku = variants.width(&cu, "BITS", "GS", uw); + const one = try outputs(r.device, nd); + try runAll(r, s, w, which, 1, nd, one, kd, ku); + var same = true; + for (counts) |n| { + const got = try outputs(r.device, nd); + try runAll(r, s, w, which, n, nd, got, kd, ku); + const rows = R / n * n; + if (!full) { + const rec = try outputs(r.device, nd); + try runAll(r, s, w6, which, n, nd, rec, &.{}, &.{}); + same = lib.sameRows(t, label, n, rec.mixed.slice(u8, R * 5120), got.mixed.slice(u8, R * 5120), rows, 5120) and same; + same = lib.sameRows(t, label, n, rec.inj.slice(u8, R * 8), got.inj.slice(u8, R * 8), rows, 8) and same; + } + same = lib.sameRows(t, label, n, one.mixed.slice(u8, R * 5120), got.mixed.slice(u8, R * 5120), rows, 5120) and same; + same = lib.sameRows(t, label, n, one.inj.slice(u8, R * 8), got.inj.slice(u8, R * 8), rows, 8) and same; + } + if (same) std.debug.print("RESULT {s}: {s}every row's bits at 1-16 rows\n", .{ label, if (full) "" else "the recorded 6/32 kernels' bits on the same values, and " }); + // the down partials of the one-row launches against float64: the ten K parts summed + const terms = try down.dense(); + defer for (terms) |v| a.free(v); + const wrong = try down.skewed(1.01).dense(); + defer for (wrong) |v| a.free(v); + var worst = lib.Worst{}; + var ctrl = lib.Worst{}; + const part = one.part.slice(f32, 10 * R * nd); + for (0..R) |row| for (0..nd) |o| { + var sum: f64 = 0; + for (0..10) |k| sum += part[(row * 10 + k) * nd + o]; + const ref = lib.dot(s.hd[row * W ..][0..W], terms, o); + worst.add(o, sum, ref[0], lib.ratio(sum, ref[0], ref[1], 32 + W / dw[1] + 24, false)); + const cref = lib.dot(s.hd[row * W ..][0..W], wrong, o); + ctrl.add(o, sum, cref[0], lib.ratio(sum, cref[0], cref[1], 32 + W / dw[1] + 24, false)); + }; + var down_label: [160]u8 = undefined; + const dl = try std.fmt.bufPrint(&down_label, "{s} down", .{label}); + worst.report(t, dl); + ctrl.control(t, dl); + try upCheck(t, r, s, which, nd, down, terms, up, kd, ku, label); +} + +/// One bf16 step where fp32's sum and exp could land either side of a bf16 tie (2^-18 of p: a few fp32 steps). +fn tieStep(p: f64, got: f64) f64 { + return if (@abs(@abs(p - got) - lib.bfStep(p) / 2) <= @abs(p) * 0x1p-18) lib.bfStep(p) else 0; +} + +/// The up product against float64 through sigmoid(up) * stream 0 (other streams zero), centred so 1% scales show. +fn upCheck(t: *lib.Tally, r: *Run, s: Streams, which: []const u8, nd: usize, down: lib.Mat, dterms: [2][]f64, up0: lib.Mat, kd: []const variants.Const, ku: []const variants.Const, label: []const u8) !void { + var up = up0; + const gl = LOW / up.group; + up.biases = try a.alloc(u16, up0.biases.len); + for (0..W) |o| for (0..gl) |g| { + var codes: f64 = 0; + for (0..up.group) |i| codes += @floatFromInt(up.code(o, g * up.group + i)); + up.biases[o * gl + g] = lib.toBf(@floatCast(-lib.fromBf(up.scales[o * gl + g]) * codes / @as(f64, @floatFromInt(up.group)))); + }; + // stream 0's scale: a power of two, so bf16(scale x) is exact; picked on row 0 from float64 estimates + const hd = try a.alloc(f64, R * W); + @memset(hd, 0); + var dsum: [LOW]f64 = undefined; + for (0..LOW) |c| { + var v: f64 = 0; + for (0..2560) |i| v += dterms[0][c * W + i] * s.hd[i]; + dsum[c] = v; + } + var scale: f64 = 1; + var best: f64 = std.math.inf(f64); + var e: i32 = -14; + while (e <= 0) : (e += 1) { + const f = std.math.ldexp(@as(f64, 1), e); + var accs: [256]f64 = undefined; + for (&accs, 0..) |*acc, o| { + acc.* = 0; + for (0..LOW) |i| { + const v = f * dsum[i] / 4; + acc.* += up.deq(o * 10, i) * v / (1 + @exp(-v)); + } + acc.* = @abs(acc.*); + } + std.mem.sort(f64, &accs, {}, std.sort.asc(f64)); + const miss = @abs(@log2(accs[128] + 1e-30)); + if (miss < best) { + best = miss; + scale = f; + } + } + const hn = try r.device.buffer(R * W * 2, lib.opts); + const hv = hn.slice(u16, R * W); + @memset(hv, 0); + for (0..R) |row| for (0..2560) |i| { + hd[row * W + i] = scale * s.hd[row * W + i]; + hv[row * W + i] = lib.toBf(@floatCast(hd[row * W + i])); + }; + const s0: Streams = .{ .hn = hn, .hd = hd, .ssp = s.ssp, .nw = s.nw, .eps = s.eps }; + const o = try outputs(r.device, nd); + try runAll(r, s0, try upload(r.device, down, up), which, 1, nd, o, kd, ku); + const part = o.part.slice(f32, 10 * R * nd); + const mixed = o.mixed.slice(u16, R * 2560); + var worst = lib.Worst{}; + var ctrl = lib.Worst{}; + var act: [LOW]f64 = undefined; + var slack: [LOW]f64 = undefined; + const steps: f64 = @floatFromInt(32 + LOW / up.group + 16); + for (0..R) |row| { + for (0..LOW) |c| { // the kernel's prologue: the parts summed in fp32, bf16, a quarter, bf16(silu) + var v: f32 = 0; + for (0..10) |k| v += part[(row * 10 + k) * nd + c]; + const v4 = lib.bfRound(v) / 4; + const p = v4 / (1 + @exp(-v4)); + act[c] = lib.bfRound(p); + slack[c] = tieStep(p, act[c]); + } + for (0..2560) |d| for ([_]f64{ 1, 1.01 }) |skew| { + var acc: [3]f64 = .{ 0, 0, 0 }; // the product, its magnitude, the activations' tie slack + for (0..LOW) |i| { + const g = d * gl + i / up.group; + const sq = skew * lib.fromBf(up.scales[g]) * @as(f64, @floatFromInt(up.code(d, i))); + const b = lib.fromBf(up.biases[g]); + acc[0] += (sq + b) * act[i]; + acc[1] += (sq + @abs(b)) * @abs(act[i]); + acc[2] += @abs(sq + b) * slack[i]; + } + const ea = steps * lib.unit * acc[1] + acc[2]; + const sig = 1 / (1 + @exp(-acc[0])); + const normed = hd[row * W + d]; + const want = sig * normed; + // bf16(up), the sigmoid's slope at most 1/4, bf16(sigmoid), bf16(sigmoid times the stream) + const tol = @abs(normed) * (0.25 * (ea + 0x1p-8 * (@abs(acc[0]) + ea)) + 0x1p-8 * sig + 1e-6) + 0x1p-8 * 1.01 * @abs(want) + 1e-30; + const got = 4 * lib.fromBf(mixed[row * 2560 + d]); + (if (skew == 1) &worst else &ctrl).add(d, got, want, @abs(got - want) / tol); + }; + } + var up_label: [160]u8 = undefined; + const ul = try std.fmt.bufPrint(&up_label, "{s} up (stream scale 2^{d})", .{ label, std.math.log2(scale) }); + worst.report(t, ul); + ctrl.control(t, ul); +} + +pub fn run(t: *lib.Tally, r: *Run) !void { + const s = try streams(r.device); + var seed: u64 = 0x4c00; + for ([_]u8{ 2, 3, 4, 5, 6, 8 }) |bits| for ([_]usize{ 32, 64, 128 }) |group| { + seed += 4; + const up_group = @min(group, 64); // 320 inputs: the pack reads an up projection's group 128 as 64 + try check(t, r, s, "ahc", 324, .{ bits, group }, .{ bits, up_group }, seed, false); + if (bits == 8) try check(t, r, s, "ahc", 324, .{ bits, group }, .{ bits, up_group }, seed + 2, true); + }; + // the mixer (no inject rows) at the oQ2 and oQ4e widths, and a down and up of different widths + try check(t, r, s, "mix", 320, .{ 2, 32 }, .{ 2, 32 }, 0x4d01, false); + try check(t, r, s, "mix", 320, .{ 4, 32 }, .{ 4, 32 }, 0x4d02, false); + try check(t, r, s, "mhc", 324, .{ 5, 64 }, .{ 4, 64 }, 0x4d03, false); +} diff --git a/zig/tests/flashnext_widths_lane.zig b/zig/tests/flashnext_widths_lane.zig new file mode 100644 index 0000000000..9052ac0a2e --- /dev/null +++ b/zig/tests/flashnext_widths_lane.zig @@ -0,0 +1,188 @@ +//! tf-flashnext-widths, decode lanes: fused stacks packed as the engine packs them, every run at its own width. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const lib = @import("flashnext_widths_lib.zig"); +const lane_w = tf.flashnext_replay.lane_w; +const pack = tf.flashnext_pack; +const st = tf.safetensors; +const Buf = tf.flashnext_replay.Buf; +const a = lib.a; + +const wide = 64; // rows of the reference call: two 32-row blocks (the MTP head's 4-row windows reach 64) +const windows = [_]usize{ 1, 2, 3, 8, 16, 17, 33 }; // 17 and 33 cross the 16- and 32-row blocks + +/// A stack member: rows, bits and group as the checkpoint stores it. +const Member = struct { rows: usize, bits: u8, group: usize }; + +fn tensor(dtype: st.DType, rows: usize, cols: usize, bytes: []const u8) st.Tensor { + var shape: [st.max_rank]usize = @splat(0); + shape[0] = rows; + shape[1] = cols; + return .{ .dtype = dtype, .rank = 2, .shape = shape, .bytes = bytes }; +} + +fn q(m: lib.Mat) !pack.Q { + const kw = m.k * m.bits / 32; + const tables = std.mem.sliceAsBytes(m.scales); + return .{ .w = tensor(.u32, m.n, kw, m.words), .s = tensor(.bf16, m.n, m.k / m.group, tables), .b = tensor(.bf16, m.n, m.k / m.group, std.mem.sliceAsBytes(m.biases)), .rows = m.n, .k = m.k, .kw = kw, .spec = try @FieldType(pack.Q, "spec").init(m.bits, m.group) }; +} + +fn blob(out: *const pack.Out, name: []const u8) ?[]const u8 { + for (out.names.items, 0..) |n, i| if (std.mem.eql(u8, n, name)) return out.blobs.items[i]; + return null; +} + +/// The device lane the engine would load from these pack tensors (lane_w.of's reading, from memory). +fn wideOf(device: mtl.Device, out: *const pack.Out, k: usize, rows: usize) !lane_w.Wide { + const uniform = [_]u32{ 6, 32, @intCast(rows) }; // a 6/32 lane writes no table + const table = if (blob(out, "t.runs")) |b| std.mem.bytesAsSlice(u32, @as([]align(1) const u8, b)) else &uniform; + const n_runs = table.len / 3; + const run_rows = try a.alloc(usize, n_runs); + for (run_rows, 0..) |*n, i| n.* = table[i * 3 + 2]; + const maps = if (n_runs > 1) blk: { + const m = blob(out, "t.members").?; + const members = try a.alloc(u32, m.len / 4); + @memcpy(std.mem.sliceAsBytes(members), m); + break :blk try lane_w.columnMaps(a, members, run_rows); + } else null; + const runs = try a.alloc(lane_w.Product, n_runs); + var total: usize = 0; + for (runs, 0..) |*p, i| { + const base = if (n_runs == 1) "t" else try std.fmt.allocPrint(a, "t.r{d}", .{i}); + const wq = blob(out, try std.fmt.allocPrint(a, "{s}.wq", .{base})).?; + const sbt = blob(out, try std.fmt.allocPrint(a, "{s}.sbt", .{base})).?; + p.* = .{ .wq = .{ .b = try lib.upload(device, wq) }, .sbt = .{ .b = try lib.upload(device, sbt) }, .n = run_rows[i], .k = k, .bits = @intCast(table[i * 3]), .group = @intCast(table[i * 3 + 1]), .tiled = pack.lane.tiled(run_rows[i]) }; + if (maps) |m| p.map = .{ .b = try lib.upload(device, std.mem.sliceAsBytes(m[i])) }; + total += run_rows[i]; + } + return .{ .runs = runs, .n = total }; +} + +fn mdims(device: mtl.Device, rows: usize) !Buf { + const b = try device.buffer(32, lib.opts); + b.slice(i32, 2)[0] = @intCast(rows); + b.slice(i32, 2)[1] = @intCast((rows + 15) / 16 * 16); + return .{ .b = b }; +} + +/// Rows [start, start + m) of x through the lane at its widths, in windows of m rows. +fn launchWindows(device: mtl.Device, queue: mtl.Queue, cache: *lane_w.Cache, w: *const lane_w.Wide, x: lib.Inputs, m: usize, y: mtl.Buffer) !void { + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + var start: usize = 0; + while (start + m <= wide) : (start += m) { + try lane_w.encode(e, a, device, cache, w, .{ .b = x.buf, .off = start * x.k * 2 }, try mdims(device, m), m, .{ .b = y, .off = start * w.n * 2 }); + } + try lib.commit(cb, e); +} + +/// One stack: packed by the engine's packer, run at 64 rows against float64, then every window's rows bit for bit. +fn stack(t: *lib.Tally, device: mtl.Device, queue: mtl.Queue, cache: *lane_w.Cache, members: []const Member, k: usize, seed: u64) !void { + var mats = try a.alloc(lib.Mat, members.len); + var qs = try a.alloc(pack.Q, members.len); + var label: std.ArrayList(u8) = .empty; + try label.print(a, "lane K {d}:", .{k}); + for (members, 0..) |m, i| { + mats[i] = try lib.Mat.init(m.rows, k, m.bits, m.group, seed + i); + qs[i] = try q(mats[i]); + try label.print(a, " {d}x q{d}/g{d}", .{ m.rows, m.bits, m.group }); + } + var out = pack.Out.init(a); + try pack.lane.putLane(&out, a, "t", qs); + var rows: usize = 0; + for (members) |m| rows += m.rows; + const w = try wideOf(device, &out, k, rows); + try label.print(a, " (SK {d})", .{lane_w.splitK(w.runs[0].n, k)}); + const x = try lib.Inputs.init(device, wide, k, seed ^ 0x5eed); + const all = try lib.poisoned(device, wide * w.n * 2); + try launchWindows(device, queue, cache, &w, x, wide, all); + var worst = lib.Worst{}; + var ctrl = lib.Worst{}; + const got = all.slice(u16, wide * w.n); + var col: usize = 0; + for (mats) |m| { + const terms = try m.dense(); + const wrong = try m.skewed(1.01).dense(); + defer for (terms ++ wrong) |v| a.free(v); + for (0..m.n) |j| { + for (0..wide) |r| { + const g = lib.fromBf(got[r * w.n + col + j]); + const ref = lib.dot(x.row(r), terms, j); + worst.add(col + j, g, ref[0], lib.ratio(g, ref[0], ref[1], m.group + k / m.group + 16, true)); + const cref = lib.dot(x.row(r), wrong, j); + ctrl.add(col + j, g, cref[0], lib.ratio(g, cref[0], cref[1], m.group + k / m.group + 16, true)); + } + } + col += m.n; + } + worst.report(t, label.items); + ctrl.control(t, label.items); + var same = true; + for (windows) |m| { + const y = try lib.poisoned(device, wide * w.n * 2); + try launchWindows(device, queue, cache, &w, x, m, y); + same = lib.sameRows(t, label.items, m, all.slice(u8, wide * w.n * 2), y.slice(u8, wide * w.n * 2), wide / m * m, w.n * 2) and same; + } + if (same) std.debug.print("RESULT {s}: every row's bits equal in windows of 1, 2, 3, 8, 16, 17 and 33 rows ({d} runs)\n", .{ label.items, w.runs.len }); +} + +/// The generic source at 6/32 against the recorded kernel at a recorded shape: the same bits. +fn recordedTwin(t: *lib.Tally, device: mtl.Device, queue: mtl.Queue, cache: *lane_w.Cache) !void { + const n = 2560; + const k = 6144; + const m = try lib.Mat.init(n, k, 6, 32, 0x632); + var out = pack.Out.init(a); + try pack.lane.putLane(&out, a, "t", &.{try q(m)}); + const wq: Buf = .{ .b = try lib.upload(device, blob(&out, "t.wq").?) }; + const sbt: Buf = .{ .b = try lib.upload(device, blob(&out, "t.sbt").?) }; + const x = try lib.Inputs.init(device, 16, k, 0x633); + const md = try mdims(device, 16); + const text = try lane_w.template(device.tensorUnits()); + const rec = try mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); + const rec_pipe = try mtl.Pipeline.init(device, rec, try lane_w.functionName(a, text), false); + const y_rec = try lib.poisoned(device, 16 * n * 2); + const y_gen = try lib.poisoned(device, 16 * n * 2); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + e.setPipeline(rec_pipe); + for ([_]Buf{ .{ .b = x.buf }, .{ .b = x.buf }, wq, sbt, md, .{ .b = y_rec } }, 0..) |b, j| e.setBuffer(b.b, b.off, j); + e.dispatchThreads(mtl.Size.of(n / 32 * 32 * 8, 1, 1), mtl.Size.of(32 * 8, 1, 1)); + const w = lane_w.Wide{ .runs = &.{.{ .wq = wq, .sbt = sbt, .n = n, .k = k, .bits = 6, .group = 32 }}, .n = n }; + try lane_w.encode(e, a, device, cache, &w, .{ .b = x.buf }, md, 16, .{ .b = y_gen }); + try lib.commit(cb, e); + const same = lib.sameRows(t, "lane recorded 6/32 [2560, 6144] against the generic source", 16, y_rec.slice(u8, 16 * n * 2), y_gen.slice(u8, 16 * n * 2), 16, n * 2); + if (same) std.debug.print("RESULT lane recorded 6/32 [2560, 6144]: the generic source's bits equal the recorded kernel's\n", .{}); +} + +pub fn run(t: *lib.Tally, device: mtl.Device, queue: mtl.Queue) !void { + var cache: lane_w.Cache = .empty; + try recordedTwin(t, device, queue, &cache); + var seed: u64 = 0x1a9e; + for ([_]u8{ 2, 3, 4, 5, 6, 8 }) |bits| for ([_]usize{ 32, 64, 128 }) |group| { + for ([_][2]usize{ .{ 256, 2560 }, .{ 64, 6144 } }) |shape| { + seed += 16; + try stack(t, device, queue, &cache, &.{.{ .rows = shape[0], .bits = bits, .group = group }}, shape[1], seed); + } + }; + // the public checkpoints' mixed stacks (q, k, v, index_qk and the DeltaNet inputs), rows cut down + const stacks = [_][]const Member{ + &.{ .{ .rows = 256, .bits = 3, .group = 64 }, .{ .rows = 32, .bits = 3, .group = 64 }, .{ .rows = 32, .bits = 8, .group = 64 }, .{ .rows = 64, .bits = 3, .group = 64 } }, + &.{ .{ .rows = 256, .bits = 6, .group = 64 }, .{ .rows = 32, .bits = 6, .group = 64 }, .{ .rows = 32, .bits = 6, .group = 128 }, .{ .rows = 64, .bits = 6, .group = 64 } }, + &.{ .{ .rows = 256, .bits = 4, .group = 64 }, .{ .rows = 64, .bits = 5, .group = 128 }, .{ .rows = 32, .bits = 5, .group = 128 }, .{ .rows = 32, .bits = 6, .group = 64 } }, + &.{ .{ .rows = 256, .bits = 3, .group = 64 }, .{ .rows = 64, .bits = 6, .group = 64 }, .{ .rows = 32, .bits = 6, .group = 64 }, .{ .rows = 32, .bits = 6, .group = 64 } }, + &.{ .{ .rows = 128, .bits = 4, .group = 32 }, .{ .rows = 64, .bits = 4, .group = 64 }, .{ .rows = 32, .bits = 2, .group = 32 } }, + // oQ4e layer 21's DeltaNet inputs: runs of 112 (5-bit z and b) and 48 (6-bit a) columns, untiled + &.{ .{ .rows = 256, .bits = 4, .group = 64 }, .{ .rows = 64, .bits = 5, .group = 128 }, .{ .rows = 48, .bits = 5, .group = 128 }, .{ .rows = 48, .bits = 6, .group = 64 } }, + &.{.{ .rows = 48, .bits = 8, .group = 64 }}, + }; + for (stacks) |s| { + seed += 16; + try stack(t, device, queue, &cache, s, 2560, seed); + } + // K slices 1 and 2 (the vocabulary head and draft at 8/64, gdn.in), at 8 bits and a low width + for ([_]usize{ 32768, 16384 }) |n| for ([_][2]usize{ .{ 8, 64 }, .{ 3, 64 } }) |wd| { + seed += 16; + try stack(t, device, queue, &cache, &.{.{ .rows = n, .bits = @intCast(wd[0]), .group = wd[1] }}, 2560, seed); + }; +} diff --git a/zig/tests/flashnext_widths_lib.zig b/zig/tests/flashnext_widths_lib.zig new file mode 100644 index 0000000000..8c87f5af50 --- /dev/null +++ b/zig/tests/flashnext_widths_lib.zig @@ -0,0 +1,272 @@ +//! Shared parts of tf-flashnext-widths: synthetic MLX matrices, float64 references and the row checks. +const std = @import("std"); +const mtl = @import("metal"); + +pub const opts = mtl.ResourceOptions.shared; +pub const a = std.heap.page_allocator; +pub const unit: f64 = 0x1p-24; // fp32's unit roundoff + +pub fn toBf(v: f32) u16 { + const u: u32 = @bitCast(v); + return @truncate((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); +} + +pub fn fromBf(v: u16) f64 { + return @as(f32, @bitCast(@as(u32, v) << 16)); +} + +/// One synthetic MLX affine matrix [n, k]: codes kept beside their packed words, bf16 tables. +pub const Mat = struct { + n: usize, + k: usize, + bits: u8, + group: usize, + codes: []u8, + words: []u8, + scales: []u16, + biases: []u16, + skew: f64 = 1, // the references' scale factor: 1.01 makes the negative controls' wrong references + + /// The same matrix with its scales `f` times their values in every reference (the kernels' bytes unchanged). + pub fn skewed(m: Mat, f: f64) Mat { + var out = m; + out.skew = f; + return out; + } + + pub fn init(n: usize, k: usize, bits: u8, group: usize, seed: u64) !Mat { + return initBelow(n, k, bits, group, seed, @as(u32, 1) << @intCast(bits)); + } + + /// Codes below `top` only (an 8-bit matrix whose codes fit 6 bits, so it also exists at 6 bits). + pub fn initBelow(n: usize, k: usize, bits: u8, group: usize, seed: u64, top: u32) !Mat { + var prng = std.Random.DefaultPrng.init(seed); + const r = prng.random(); + const tables = n * (k / group); + const codes = try a.alloc(u8, n * k); + for (codes) |*c| c.* = @intCast(r.uintLessThan(u32, top)); + const m = Mat{ .n = n, .k = k, .bits = bits, .group = group, .codes = codes, .words = try packCodes(codes, bits), .scales = try a.alloc(u16, tables), .biases = try a.alloc(u16, tables) }; + for (m.scales, m.biases) |*s, *b| { + s.* = toBf((r.float(f32) + 0.5) / 64); + b.* = toBf(-(r.float(f32) + 0.5) * 0.15); + } + return m; + } + + /// The same codes stored at `bits` in groups of `group` (which divides this group): the same values exactly. + pub fn recode(m: Mat, bits: u8, group: usize) !Mat { + if (m.group % group != 0) return error.GroupDoesNotDivide; + for (m.codes) |c| if (c >> @intCast(bits) != 0) return error.CodeTooWide; + const reps = m.group / group; + const out = Mat{ .n = m.n, .k = m.k, .bits = bits, .group = group, .codes = m.codes, .words = try packCodes(m.codes, bits), .scales = try a.alloc(u16, m.scales.len * reps), .biases = try a.alloc(u16, m.biases.len * reps) }; + for (m.scales, m.biases, 0..) |sv, bv, i| for (0..reps) |j| { + out.scales[i * reps + j] = sv; + out.biases[i * reps + j] = bv; + }; + return out; + } + + pub fn deinit(m: Mat) void { + a.free(m.words); + a.free(m.scales); + a.free(m.biases); + } + + pub fn code(m: Mat, j: usize, i: usize) u32 { + return m.codes[j * m.k + i]; + } + + pub fn deq(m: Mat, j: usize, i: usize) f64 { + const g = j * (m.k / m.group) + i / m.group; + return m.skew * fromBf(m.scales[g]) * @as(f64, @floatFromInt(m.code(j, i))) + fromBf(m.biases[g]); + } + + /// s q + b and s q + |b| for every weight, row-major [n, k]: the reference's terms and their magnitudes. + pub fn dense(m: Mat) ![2][]f64 { + const out = [2][]f64{ try a.alloc(f64, m.n * m.k), try a.alloc(f64, m.n * m.k) }; + for (0..m.n) |j| for (0..m.k) |i| { + const g = j * (m.k / m.group) + i / m.group; + const sq = m.skew * fromBf(m.scales[g]) * @as(f64, @floatFromInt(m.code(j, i))); + out[0][j * m.k + i] = sq + fromBf(m.biases[g]); + out[1][j * m.k + i] = sq + @abs(fromBf(m.biases[g])); + }; + return out; + } + + /// Weights as qmmw's dequantw rounds them, bf16(s q + b), their magnitudes, and a bf16 step where fp32 nears a tie. + pub fn denseBf(m: Mat) ![3][]f64 { + const out = [3][]f64{ try a.alloc(f64, m.n * m.k), try a.alloc(f64, m.n * m.k), try a.alloc(f64, m.n * m.k) }; + for (0..m.n) |j| for (0..m.k) |i| { + const v = m.deq(j, i); + const w = bfRound(v); + const tie = @abs(@abs(v - w) - bfStep(v) / 2) <= @abs(v) * 0x1p-23; // fp32's error, twice over + out[0][j * m.k + i] = w; + out[1][j * m.k + i] = @abs(w); + out[2][j * m.k + i] = if (tie) bfStep(v) else 0; + }; + return out; + } +}; + +/// v rounded to bf16 (nearest, ties to even), as float64. +pub fn bfRound(v: f64) f64 { + return fromBf(toBf(@floatCast(v))); +} + +/// The bf16 step at v's binade. +pub fn bfStep(v: f64) f64 { + if (v == 0) return 0x1p-133; + return std.math.ldexp(@as(f64, 1), std.math.frexp(v).exponent - 8); +} + +/// The float64 product of one row with one weight row, and its magnitude: sum of (s q + |b|) |x|. +pub fn dot(x: []const f64, terms: [2][]f64, j: usize) [2]f64 { + const k = x.len; + var want: f64 = 0; + var mag: f64 = 0; + for (x, terms[0][j * k ..][0..k], terms[1][j * k ..][0..k]) |xv, d, m| { + want += xv * d; + mag += @abs(xv) * m; + } + return .{ want, mag }; +} + +/// dot over denseBf's weights: the product, its magnitude sum |x| |w|, and the tie slack sum |x| step. +pub fn dot3(x: []const f64, terms: [3][]f64, j: usize) [3]f64 { + const k = x.len; + var out: [3]f64 = .{ 0, 0, 0 }; + for (x, terms[0][j * k ..][0..k], terms[1][j * k ..][0..k], terms[2][j * k ..][0..k]) |xv, w, m, e| { + out[0] += xv * w; + out[1] += @abs(xv) * m; + out[2] += @abs(xv) * e; + } + return out; +} + +/// ratio for a dot3 reference: fp32 sums of `terms` steps over the magnitude, the tie slack, one bf16 rounding. +pub fn ratio3(got: f64, ref: [3]f64, terms: usize, bf16_out: bool) f64 { + const tol = @as(f64, @floatFromInt(terms)) * unit * ref[1] + ref[2] + (if (bf16_out) @abs(ref[0]) * 0x1p-8 else 0) + 1e-30; + return @abs(got - ref[0]) / tol; +} + +/// Codes as an MLX little-endian bit stream of `bits`-wide values. +pub fn packCodes(codes: []const u8, bits: u8) ![]u8 { + const out = try a.alloc(u8, (codes.len * bits + 31) / 32 * 4); + @memset(out, 0); + for (codes, 0..) |c, i| { + const bit = i * bits; + const v = @as(u32, c) << @intCast(bit % 8); + for (0..3) |b| { + if (bit / 8 + b < out.len) out[bit / 8 + b] |= @truncate(v >> @intCast(8 * b)); + } + } + return out; +} + +/// Every check's outcome: failures, the largest difference as a fraction of its bound, and the negative controls. +pub const Tally = struct { failed: usize = 0, worst: f64 = 0, checks: usize = 0, controls: usize = 0, least: f64 = std.math.inf(f64) }; + +/// `rows` rows of x [rows, k] as bf16 (in a device buffer) and float64. +pub const Inputs = struct { + buf: mtl.Buffer, + x: []f64, + rows: usize, + k: usize, + + pub fn init(device: mtl.Device, rows: usize, k: usize, seed: u64) !Inputs { + var prng = std.Random.DefaultPrng.init(seed); + const r = prng.random(); + const in = Inputs{ .buf = try device.buffer(rows * k * 2, opts), .x = try a.alloc(f64, rows * k), .rows = rows, .k = k }; + for (in.buf.slice(u16, rows * k), in.x) |*v, *d| { + v.* = toBf(r.float(f32) * 2 - 1); + d.* = fromBf(v.*); + } + return in; + } + + pub fn row(in: Inputs, r: usize) []const f64 { + return in.x[r * in.k ..][0..in.k]; + } +}; + +/// A column's result against float64: the error over its bound (fp32 sums of `terms` steps, one bf16 rounding). +pub fn ratio(got: f64, want: f64, mag: f64, terms: usize, bf16_out: bool) f64 { + const tol = @as(f64, @floatFromInt(terms)) * unit * mag + (if (bf16_out) @abs(want) * 0x1p-8 else 0) + 1e-30; + return @abs(got - want) / tol; +} + +/// The worst column of one product, printed with its max abs and relative error and its share of the bound. +pub const Worst = struct { + frac: f64 = 0, + col: usize = 0, + abs: f64 = 0, + rel: f64 = 0, + nan: bool = false, + + pub fn add(w: *Worst, col: usize, got: f64, want: f64, frac: f64) void { + if (std.math.isNan(got)) w.nan = true; + const d = @abs(got - want); + if (frac > w.frac or std.math.isNan(got)) { + w.frac = if (std.math.isNan(got)) std.math.inf(f64) else frac; + w.col = col; + w.abs = d; + w.rel = d / (@abs(want) + 1e-30); + } + } + + pub fn report(w: Worst, t: *Tally, label: []const u8) void { + t.checks += 1; + t.worst = @max(t.worst, w.frac); + const ok = w.frac <= 1 and !w.nan; + 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{s}\n", .{ label, w.col, w.abs, w.rel, w.frac, if (ok) "" else " FAIL" }); + } + + /// A negative control: the same results against a reference with scales 1% off must exceed the bound somewhere. + pub fn control(w: Worst, t: *Tally, label: []const u8) void { + t.checks += 1; + t.controls += 1; + t.least = @min(t.least, w.frac); + const ok = w.frac > 1; + if (!ok) t.failed += 1; + std.debug.print("RESULT {s}: negative control, scales 1% off: {d:.3} of its bound{s}\n", .{ label, w.frac, if (ok) "" else " FAIL (the bound would pass a wrong scale)" }); + } +}; + +/// Rows [0, rows) of two outputs (row stride `ld` bytes) equal byte for byte; the first difference printed. +pub fn sameRows(t: *Tally, label: []const u8, window: usize, want: []const u8, got: []const u8, rows: usize, ld: usize) bool { + t.checks += 1; + if (std.mem.eql(u8, want[0 .. rows * ld], got[0 .. rows * ld])) return true; + const i = std.mem.indexOfDiff(u8, want, got).?; + var differ: usize = 0; + for (0..rows) |r| { + if (!std.mem.eql(u8, want[r * ld ..][0..ld], got[r * ld ..][0..ld])) differ += 1; + } + std.debug.print("RESULT {s}: {d}-row windows DIFFER in {d} of {d} rows; first at row {d} byte {d} FAIL\n", .{ label, window, differ, rows, i / ld, i % ld }); + t.failed += 1; + return false; +} + +pub 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; + } +} + +/// A device buffer of `n` bytes filled with 0xff (NaN in bf16 and fp32): a launch that writes nothing fails. +pub fn poisoned(device: mtl.Device, n: usize) !mtl.Buffer { + const b = try device.buffer(@max(n, 16), opts); + @memset(b.slice(u8, @max(n, 16)), 0xff); + return b; +} + +/// A device buffer holding `bytes`. +pub fn upload(device: mtl.Device, bytes: []const u8) !mtl.Buffer { + const b = try device.buffer(@max(bytes.len, 16), opts); + @memcpy(b.slice(u8, bytes.len), bytes); + return b; +} diff --git a/zig/tests/flashnext_widths_moe.zig b/zig/tests/flashnext_widths_moe.zig new file mode 100644 index 0000000000..e8a216b7f5 --- /dev/null +++ b/zig/tests/flashnext_widths_moe.zig @@ -0,0 +1,294 @@ +//! tf-flashnext-widths, experts, embedding and n-gram tables: the recorded any-width kernels at each oQ width. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const lib = @import("flashnext_widths_lib.zig"); +const fz = tf.flashnext_replay; +const variants = fz.variants; +const Buf = fz.Buf; +const Run = fz.Run; +const a = lib.a; + +const D = 2560; +const N = 640; // an expert's intermediate width +const E = 16; // experts held: the logits pick among these only +const R = 16; +const counts = [_]usize{ 1, 2, 3, 8, 16 }; + +fn up(device: mtl.Device, m: lib.Mat) ![3]Buf { + return .{ .{ .b = try lib.upload(device, m.words) }, .{ .b = try lib.upload(device, std.mem.sliceAsBytes(m.scales)) }, .{ .b = try lib.upload(device, std.mem.sliceAsBytes(m.biases)) } }; +} + +/// Codes below 64, so every matrix also exists at 6/32 with the same values; `full`: codes over their whole range. +fn mat(n: usize, k: usize, w: [2]usize, seed: u64, full: bool) !lib.Mat { + return lib.Mat.initBelow(n, k, @intCast(w[0]), w[1], seed, @min(@as(u32, if (full) 256 else 64), @as(u32, 1) << @intCast(w[0]))); +} + +/// Gate, up, shared gate, shared up, down and shared down at the routed and shared widths, and their 6/32 twins. +const Set = struct { + e: [18]Buf, + e6: [18]Buf, + m: [6]lib.Mat, // gate, up, shared gate, shared up, down, shared down + + fn deinit(s: Set) void { + for (s.e) |b| b.b.deinit(); + if (s.e6[0].b.id != s.e[0].b.id) for (s.e6) |b| b.b.deinit(); + for (s.m) |m| { + m.deinit(); + a.free(m.codes); + } + } +}; + +fn set(device: mtl.Device, routed: [2]usize, shared: [2]usize, seed: u64, full: bool) !Set { + const shapes = [_][3]usize{ .{ E * N, D, 0 }, .{ E * N, D, 0 }, .{ N, D, 1 }, .{ N, D, 1 }, .{ E * D, N, 0 }, .{ D, N, 1 } }; + var s: Set = undefined; + for (shapes, 0..) |sh, j| { + const m = try mat(sh[0], sh[1], if (sh[2] == 1) shared else routed, seed + j, full); + const b = try up(device, m); + const m6 = if (full) m else try m.recode(6, 32); // full-range 8-bit codes have no 6/32 twin + defer if (!full) m6.deinit(); + const b6 = if (full) b else try up(device, m6); + for (0..3) |k| { + s.e[j * 3 + k] = b[k]; + s.e6[j * 3 + k] = b6[k]; + } + s.m[j] = m; + } + return s; +} + +const Io = struct { x: lib.Inputs, lg: mtl.Buffer }; + +fn inputs(device: mtl.Device) !Io { + const lg = try device.buffer(R * 513 * 4, lib.opts); + const v = lg.slice(f32, R * 513); + for (0..R) |r| for (0..513) |e| { + v[r * 513 + e] = if (e < E) 4.0 - @as(f32, @floatFromInt((e * 7 + r * 3) % E)) * 0.25 else if (e == 512) 0.5 else -100; + }; + return .{ .x = try lib.Inputs.init(device, R, D, 0x3e), .lg = lg }; +} + +const Out = struct { act: mtl.Buffer, pick: mtl.Buffer, wts: mtl.Buffer, y: mtl.Buffer }; + +fn outputs(device: mtl.Device) !Out { + return .{ .act = try lib.poisoned(device, R * 11 * N * 2), .pick = try lib.poisoned(device, R * 10 * 4), .wts = try lib.poisoned(device, R * 10 * 4), .y = try lib.poisoned(device, R * 11 * D * 2) }; +} + +/// Every window of `rows` rows through gate/up then down, at `w` (null: the recorded 6/32 kernels on `e`). +fn experts(r: *Run, io: Io, e: [18]Buf, rows: usize, o: Out, w: ?variants.Experts) !void { + const cb = r.queue.commandBuffer(); + r.enc = cb.compute(.serial); + var start: usize = 0; + while (start + rows <= R) : (start += rows) { + const rb = try r.buffer(16); + rb.slice(i32, 1)[0] = @intCast(rows); + r.rows = rows; + const x: Buf = .{ .b = io.x.buf, .off = start * D * 2 }; + const lg: Buf = .{ .b = io.lg, .off = start * 513 * 4 }; + const act: Buf = .{ .b = o.act, .off = start * 11 * N * 2 }; + const pick: Buf = .{ .b = o.pick, .off = start * 40 }; + const wts: Buf = .{ .b = o.wts, .off = start * 40 }; + const y: Buf = .{ .b = o.y, .off = start * 11 * D * 2 }; + if (w) |ww| try variants.experts(r, "qa_expert_gateup@moe.gate", "qa_expert_down_y@moe.down", x, lg, e, act, pick, wts, .{ .b = rb }, y, ww) else { + try r.call("qa_expert_gateup@moe.gate", &.{ x, lg, e[0], e[1], e[2], e[3], e[4], e[5], e[6], e[7], e[8], e[9], e[10], e[11] }, &.{ act, pick, wts }); + try r.call("qa_expert_down_y@moe.down", &.{ act, pick, e[12], e[13], e[14], e[15], e[16], e[17], .{ .b = rb } }, &.{y}); + } + } + try lib.commit(cb, r.enc); +} + +fn same4(t: *lib.Tally, label: []const u8, n: usize, want: Out, got: Out) bool { + const rows = R / n * n; + var ok = lib.sameRows(t, label, n, want.act.slice(u8, R * 11 * N * 2), got.act.slice(u8, R * 11 * N * 2), rows, 11 * N * 2); + ok = lib.sameRows(t, label, n, want.pick.slice(u8, R * 40), got.pick.slice(u8, R * 40), rows, 40) and ok; + ok = lib.sameRows(t, label, n, want.wts.slice(u8, R * 40), got.wts.slice(u8, R * 40), rows, 40) and ok; + return lib.sameRows(t, label, n, want.y.slice(u8, R * 11 * D * 2), got.y.slice(u8, R * 11 * D * 2), rows, 11 * D * 2) and ok; +} + +fn expertCheck(t: *lib.Tally, r: *Run, io: Io, routed: [2]usize, shared: [2]usize, seed: u64, full: bool) !void { + var label_buf: [128]u8 = undefined; + const label = try std.fmt.bufPrint(&label_buf, "experts routed q{d}/g{d} shared q{d}/g{d}{s}", .{ routed[0], routed[1], shared[0], shared[1], if (full) " full-range codes" else "" }); + const s = try set(r.device, routed, shared, seed, full); + defer s.deinit(); + const w: variants.Experts = .{ .gu = .{ routed[0], routed[1], shared[0], shared[1] }, .down = .{ routed[0], routed[1], shared[0], shared[1] } }; + const one = try outputs(r.device); + try experts(r, io, s.e, 1, one, w); + // 6 and 8 bits read 8 codes a lane step, other widths 16 (0.6.6's lane_values): only the former sum as 6/32 does + const as_recorded = !full and (routed[0] == 6 or routed[0] == 8) and (shared[0] == 6 or shared[0] == 8); + var same = true; + for (counts) |n| { + const got = try outputs(r.device); + try experts(r, io, s.e, n, got, w); + if (as_recorded) { + const rec = try outputs(r.device); + try experts(r, io, s.e6, n, rec, null); + same = same4(t, label, n, rec, got) and same; + } + same = same4(t, label, n, one, got) and same; + } + if (same) std.debug.print("RESULT {s}: every row's bits at 1-16 rows{s}\n", .{ label, if (as_recorded) ", and the recorded 6/32 kernels' bits on the same values" else "" }); + try gateUp(t, label, s, io, one); + // the down products against float64 from the activations the gate/up launch wrote + var worst = [2]lib.Worst{ .{}, .{} }; // against the right scales, and the negative control's 1% off + for ([_]f64{ 1, 1.01 }, &worst) |skew, *wo| { + const routed_terms = try s.m[4].skewed(skew).dense(); + const shared_terms = try s.m[5].skewed(skew).dense(); + defer for (routed_terms ++ shared_terms) |v| a.free(v); + const act = one.act.slice(u16, R * 11 * N); + const pick = one.pick.slice(u32, R * 10); + const y = one.y.slice(u16, R * 11 * D); + const xa = try a.alloc(f64, N); + defer a.free(xa); + for (0..R) |row| for (0..11) |slot| { + for (xa, 0..) |*v, i| v.* = lib.fromBf(act[(row * 11 + slot) * N + i]); + for (0..D) |d| { + const ref = if (slot < 10) lib.dot(xa, routed_terms, pick[row * 10 + slot] * D + d) else lib.dot(xa, shared_terms, d); + const got = lib.fromBf(y[(row * 11 + slot) * D + d]); + const g = if (slot < 10) routed[1] else shared[1]; + wo.add(d, got, ref[0], lib.ratio(got, ref[0], ref[1], 32 + N / g + 32, true)); + } + }; + } + var down_label: [160]u8 = undefined; + const dl = try std.fmt.bufPrint(&down_label, "{s} down", .{label}); + worst[0].report(t, dl); + worst[1].control(t, dl); +} + +/// Four rows' activations silu(x Wg) * (x Wu) against float64: sums' bounds through the activation and roundings. +fn gateUp(t: *lib.Tally, label: []const u8, s: Set, io: Io, one: Out) !void { + var worst = [2]lib.Worst{ .{}, .{} }; // against the right scales, and the negative control's 1% off + const act = one.act.slice(u16, R * 11 * N); + const pick = one.pick.slice(u32, R * 10); + for ([_]f64{ 1, 1.01 }, &worst) |skew, *w| for (0..4) |row| for (0..11) |slot| for (0..N) |n| { + const x = io.x.row(row); + const base: usize = if (slot < 10) pick[row * 10 + slot] * N else 0; + const gm = (if (slot < 10) s.m[0] else s.m[2]).skewed(skew); + const um = (if (slot < 10) s.m[1] else s.m[3]).skewed(skew); + var gv: [2]f64 = .{ 0, 0 }; // the product and its magnitude sum |x| (s q + |b|) + var uv: [2]f64 = .{ 0, 0 }; + for (0..D) |i| { + const gq = gm.skew * lib.fromBf(gm.scales[(base + n) * (D / gm.group) + i / gm.group]) * @as(f64, @floatFromInt(gm.code(base + n, i))); + const uq = um.skew * lib.fromBf(um.scales[(base + n) * (D / um.group) + i / um.group]) * @as(f64, @floatFromInt(um.code(base + n, i))); + gv[0] += x[i] * gm.deq(base + n, i); + gv[1] += @abs(x[i]) * (gq + @abs(lib.fromBf(gm.biases[(base + n) * (D / gm.group) + i / gm.group]))); + uv[0] += x[i] * um.deq(base + n, i); + uv[1] += @abs(x[i]) * (uq + @abs(lib.fromBf(um.biases[(base + n) * (D / um.group) + i / um.group]))); + } + const sig = 1 / (1 + @exp(-gv[0])); + const silu = gv[0] * sig; + const want = silu * uv[0]; + const steps: f64 = @floatFromInt(32 + D / gm.group + 16); + const dg = steps * lib.unit * gv[1] + 0x1p-8 * @abs(gv[0]); + const du = steps * lib.unit * uv[1] + 0x1p-8 * @abs(uv[0]); + const tol = 1.1 * @abs(uv[0]) * dg + @abs(silu) * du + 0x1p-7 * @abs(want) + 1e-6; + const got = lib.fromBf(act[(row * 11 + slot) * N + n]); + w.add(n, got, want, @abs(got - want) / tol); + }; + var gu_label: [160]u8 = undefined; + const gl = try std.fmt.bufPrint(&gu_label, "{s} gate/up", .{label}); + worst[0].report(t, gl); + worst[1].control(t, gl); +} + +/// The embedding rows (4 stream copies, and the MTP head's one) and the n-gram lookups at width `w`. +fn lookups(t: *lib.Tally, r: *Run, w: [2]usize, seed: u64, full: bool) !void { + var label_buf: [96]u8 = undefined; + const label = try std.fmt.bufPrint(&label_buf, "embed and n-gram tables q{d}/g{d}{s}", .{ w[0], w[1], if (full) " full-range codes" else "" }); + var c: [2]variants.Const = undefined; + const k = variants.width(&c, "BITS", "GS", w); + const emb = try mat(64, D, w, seed, full); + const e = try up(r.device, emb); + const e6 = if (full) e else try up(r.device, try emb.recode(6, 32)); + const ids = try r.buffer(R * 16 * 4); + for (ids.slice(u32, R * 16), 0..) |*v, i| v.* = @intCast((i * 37 + 5) % 64); + var tabs: [26]Buf = undefined; + var tabs6: [26]Buf = undefined; + tabs[0] = .{ .b = ids }; + const starts = try r.buffer(32); + for (starts.slice(u32, 8), 0..) |*v, g| v.* = @intCast(8 * g); + tabs[1] = .{ .b = starts }; + var ple_mats: [8]lib.Mat = undefined; + const tables: usize = if (160 % w[1] == 0) 3 else 2; // an n-gram row is 160 wide: groups of 32 only + for (0..8) |g| { + ple_mats[g] = try mat(8, 160, .{ w[0], if (tables == 3) w[1] else 32 }, seed + 1 + g, full); + const b = try up(r.device, ple_mats[g]); + const b6 = if (full) b else try up(r.device, try ple_mats[g].recode(6, 32)); + for (0..3) |j| { + tabs[2 + 3 * g + j] = b[j]; + tabs6[2 + 3 * g + j] = b6[j]; + } + } + tabs6[0] = tabs[0]; + tabs6[1] = tabs[1]; + defer { + for (e) |b| b.b.deinit(); + if (!full) for (e6) |b| b.b.deinit(); + } + const roles = [_][]const u8{ "qa_embed_rows@embed", "mtp:qa_embed_rows@embed", "qa_ple_lookup@ple" }; + const sizes = [_]usize{ 4 * D * 2, D * 2, 16 * 160 * 2 }; + var same = true; + var worst = lib.Worst{}; + var ctrl = lib.Worst{}; + var ones: [3]mtl.Buffer = undefined; // each role's one-row outputs: a full-range pass's rows compare with them + for (roles[0..tables], sizes[0..tables], 0..) |role, size, which| for (counts) |n| { + const got = try lib.poisoned(r.device, R * size); + const rec = try lib.poisoned(r.device, R * size); + const cb = r.queue.commandBuffer(); + r.enc = cb.compute(.serial); + var start: usize = 0; + while (start + n <= R) : (start += n) { + r.rows = n; + const id: Buf = .{ .b = ids, .off = start * @as(usize, if (which == 2) 16 * 4 else 4) }; + const out: Buf = .{ .b = got, .off = start * size }; + const out6: Buf = .{ .b = rec, .off = start * size }; + if (which < 2) { + try variants.callAs(r, role, n, &.{ id, e[0], e[1], e[2] }, &.{out}, k); + try r.callAs(role, n, &.{ id, e6[0], e6[1], e6[2] }, &.{out6}); + } else { + var tw = tabs; + var tw6 = tabs6; + tw[0] = id; + tw6[0] = id; + try variants.callAs(r, role, n, &tw, &.{out}, k); + try r.callAs(role, n, &tw6, &.{out6}); + } + } + try lib.commit(cb, r.enc); + if (n == 1) ones[which] = got; + const want_rows = if (full) ones[which] else rec; + same = lib.sameRows(t, label, n, want_rows.slice(u8, R * size), got.slice(u8, R * size), R / n * n, size) and same; + if (n != 1) continue; + // the values against s q + b in float64: two bf16 roundings at most + const v = got.slice(u16, R * size / 2); + for (0..R) |row| for (0..size / 2) |i| for ([_]f64{ 1, 1.01 }) |skew| { + const want = if (which == 2) blk: { + const head = i / 160; + const id = ids.slice(u32, R * 16)[row * 16 + head]; + break :blk ple_mats[id / 8].skewed(skew).deq(id % 8, i % 160); + } else emb.skewed(skew).deq(ids.slice(u32, R * 16)[row], i % D); + const g = lib.fromBf(v[row * size / 2 + i]); + (if (skew == 1) &worst else &ctrl).add(i, g, want, @abs(g - want) / (@abs(want) * 0x1p-7 + 1e-6)); + }; + }; + if (same) std.debug.print("RESULT {s}: {s} at 1-16 rows\n", .{ label, if (full) "every row's bits" else "the recorded 6/32 kernels' bits on the same values" }); + worst.report(t, label); + ctrl.control(t, label); +} + +pub fn run(t: *lib.Tally, r: *Run) !void { + const io = try inputs(r.device); + var seed: u64 = 0xe000; + for ([_]u8{ 2, 3, 4, 5, 6, 8 }) |bits| for ([_]usize{ 32, 64, 128 }) |group| { + seed += 8; + try expertCheck(t, r, io, .{ bits, group }, .{ 8, 128 }, seed, false); + try lookups(t, r, .{ bits, group }, seed + 4, false); + }; + try expertCheck(t, r, io, .{ 4, 32 }, .{ 4, 32 }, 0xe801, false); // the MTP head's experts in oQ2 + // 8-bit codes past 63 (the embedding and the shared expert are 8-bit in every oQ checkpoint): float64 and rows + for ([_]usize{ 32, 64, 128 }, 0..) |group, i| { + try expertCheck(t, r, io, .{ 8, group }, .{ 8, 128 }, 0xe900 + 8 * i, true); + try lookups(t, r, .{ 8, group }, 0xe904 + 8 * i, true); + } +} diff --git a/zig/tests/flashnext_widths_prompt.zig b/zig/tests/flashnext_widths_prompt.zig new file mode 100644 index 0000000000..43f85c1751 --- /dev/null +++ b/zig/tests/flashnext_widths_prompt.zig @@ -0,0 +1,249 @@ +//! tf-flashnext-widths, prompt products: qmmw_nax at each oQ width, dense, split-K, members and expert gathers. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const lib = @import("flashnext_widths_lib.zig"); +const fz = tf.flashnext_replay; +const pw = fz.prompt_w; +const Buf = fz.Buf; +const Run = fz.Run; +const a = lib.a; + +const wide = 129; // two 64-row tiles and one row more +const windows = [_]usize{ 1, 2, 3, 8, 16, 17, 33, 64, 65 }; + +/// A prompt with only what the any-width launches read: the run and the expert offsets. +fn prompt(r: *Run, offsets: Buf) fz.Prompt { + var p: fz.Prompt = undefined; + p.r = r; + p.skip = 0; + p.b.off = offsets; + return p; +} + +fn weights(device: mtl.Device, m: lib.Mat) ![3]Buf { + return .{ .{ .b = try lib.upload(device, m.words) }, .{ .b = try lib.upload(device, std.mem.sliceAsBytes(m.scales)) }, .{ .b = try lib.upload(device, std.mem.sliceAsBytes(m.biases)) } }; +} + +/// Rows [start, start + m) of x through the dense product at `w`'s member widths, every window of m rows. +fn dense(p: *fz.Prompt, wd: *pw.Wide, w: pw.Weight, ws: [3]Buf, x: lib.Inputs, n: usize, m: usize, y: mtl.Buffer) !void { + const cb = p.r.queue.commandBuffer(); + p.r.enc = cb.compute(.serial); + var start: usize = 0; + while (start + m <= wide) : (start += m) pw.qmm(p, wd, w, .{ .b = x.buf, .off = start * x.k * 2 }, ws, x.k, n, m, .{ .b = y, .off = start * n * 2 }, 0); + try lib.commit(cb, p.r.enc); +} + +/// A fused projection of `members` (rows, bits, group): float64 per member, then every window's rows bit for bit. +fn denseCheck(t: *lib.Tally, p: *fz.Prompt, wd: *pw.Wide, members: []const [3]usize, k: usize, seed: u64) !void { + var label: std.ArrayList(u8) = .empty; + try label.print(a, "prompt dense K {d}:", .{k}); + var mats = try a.alloc(lib.Mat, members.len); + var parts = try a.alloc(pw.Part, members.len); + var words: std.ArrayList(u8) = .empty; + var scales: std.ArrayList(u8) = .empty; + var biases: std.ArrayList(u8) = .empty; + var n: usize = 0; + for (members, 0..) |mb, i| { + mats[i] = try lib.Mat.init(mb[0], k, @intCast(mb[1]), mb[2], seed + i); + parts[i] = .{ .rows = mb[0], .bits = @intCast(mb[1]), .group = @intCast(mb[2]) }; + try words.appendSlice(a, mats[i].words); + try scales.appendSlice(a, std.mem.sliceAsBytes(mats[i].scales)); + try biases.appendSlice(a, std.mem.sliceAsBytes(mats[i].biases)); + try label.print(a, " {d}x q{d}/g{d}", .{ mb[0], mb[1], mb[2] }); + n += mb[0]; + const kk = @as(u32, @intCast(mb[1])) << 16 | @as(u32, @intCast(mb[2])); + if (!wd.pipes.contains(kk)) try wd.pipes.put(a, kk, try pw.compile(p.r, @intCast(mb[1]), @intCast(mb[2]))); + } + const d = p.r.device; + const ws: [3]Buf = .{ .{ .b = try lib.upload(d, words.items) }, .{ .b = try lib.upload(d, scales.items) }, .{ .b = try lib.upload(d, biases.items) } }; + const w: pw.Weight = if (members.len == 1) .{ .bits = parts[0].bits, .group = parts[0].group } else .{ .parts = parts }; + const x = try lib.Inputs.init(d, wide, k, seed ^ 0x9e); + const all = try lib.poisoned(d, wide * n * 2); + try dense(p, wd, w, ws, x, n, wide, all); + var worst = lib.Worst{}; + var ctrl = lib.Worst{}; + const got = all.slice(u16, wide * n); + var col: usize = 0; + for (mats) |m| { + // each weight rounds to bf16 before its fp32 product (qmmw's dequantw): the reference rounds it the same way + const terms = try m.denseBf(); + defer for (terms) |v| a.free(v); + const wrong = try m.skewed(1.01).denseBf(); + defer for (wrong) |v| a.free(v); + for (0..m.n) |j| for (0..wide) |r| { + const g = lib.fromBf(got[r * n + col + j]); + const ref = lib.dot3(x.row(r), terms, j); + worst.add(col + j, g, ref[0], lib.ratio3(g, ref, k + 64, true)); + const cref = lib.dot3(x.row(r), wrong, j); + ctrl.add(col + j, g, cref[0], lib.ratio3(g, cref, k + 64, true)); + }; + col += m.n; + } + worst.report(t, label.items); + ctrl.control(t, label.items); + var same = true; + for (windows) |m| { + const y = try lib.poisoned(d, wide * n * 2); + try dense(p, wd, w, ws, x, n, m, y); + same = lib.sameRows(t, label.items, m, all.slice(u8, wide * n * 2), y.slice(u8, wide * n * 2), wide / m * m, n * 2) 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.items}); +} + +/// The generic text at 6/32 against the recorded tf_qmm6_t_nax on the same weights: the same bits. +fn recordedTwin(t: *lib.Tally, p: *fz.Prompt, wd: *pw.Wide) !void { + const d = p.r.device; + const m = try lib.Mat.init(192, 2560, 6, 32, 0x632); + const ws = try weights(d, m); + const x = try lib.Inputs.init(d, wide, 2560, 0x633); + try wd.pipes.put(a, 6 << 16 | 32, try pw.compile(p.r, 6, 32)); + const gen = try lib.poisoned(d, wide * 192 * 2); + try dense(p, wd, .{ .bits = 6, .group = 32 }, ws, x, 192, wide, gen); + const text = try tf.flashnext_replay.frags.source(d, a, pw.ks.flashnext_qmm6); + const rlib = try mtl.Library.fromSource(d, text, mtl.CompileOptions.mlx()); + const pipe = try mtl.Pipeline.init(d, rlib, "tf_qmm6_t_nax", false); + const rec = try lib.poisoned(d, wide * 192 * 2); + const cb = p.r.queue.commandBuffer(); + p.r.enc = cb.compute(.serial); + p.bind(pipe, &.{ ws[0], ws[1], ws[2], .{ .b = x.buf } }); + const prm = [4]i32{ 2560, 192, wide, 0 }; + p.r.enc.setBytes(std.mem.asBytes(&prm), 4); + p.r.enc.setBuffer(rec, 0, 5); + p.r.enc.dispatchThreads(mtl.Size.of(3 * 128, 3, 1), mtl.Size.of(128, 1, 1)); + try lib.commit(cb, p.r.enc); + if (lib.sameRows(t, "prompt recorded qmm6 [192, 2560] against the generic text at 6/32", wide, rec.slice(u8, wide * 384), gen.slice(u8, wide * 384), wide, 384)) + std.debug.print("RESULT prompt qmm6 [192, 2560]: the generic text's bits equal the recorded kernel's\n", .{}); +} + +/// The hyper-connection's down projection as the prompt splits it: 8 K parts of fp32 partials, summed against float64. +fn splitCheck(t: *lib.Tally, p: *fz.Prompt, wd: *pw.Wide, bits: u8, group: usize, seed: u64) !void { + var label_buf: [96]u8 = undefined; + const label = try std.fmt.bufPrint(&label_buf, "prompt split-K [324, 10240] q{d}/g{d}", .{ bits, group }); + const n = 324; + const k = 10240; + const d = p.r.device; + const m = try lib.Mat.init(n, k, bits, group, seed); + defer m.deinit(); + const ws = try weights(d, m); + const x = try lib.Inputs.init(d, wide, k, seed + 1); + const kk = @as(u32, bits) << 16 | @as(u32, @intCast(group)); + if (!wd.pipes.contains(kk)) try wd.pipes.put(a, kk, try pw.compile(p.r, bits, @intCast(group))); + const part = try lib.poisoned(d, 8 * wide * n * 4); + const cb = p.r.queue.commandBuffer(); + p.r.enc = cb.compute(.serial); + p.bind(wd.of(bits, @intCast(group)).splitk, &.{ ws[0], ws[1], ws[2], .{ .b = x.buf } }); + const prm = [4]i32{ k, n, wide, 8 }; + p.r.enc.setBytes(std.mem.asBytes(&prm), 4); + p.r.enc.setBuffer(part, 0, 5); + p.r.enc.dispatchThreads(mtl.Size.of(((n + 63) / 64) * 128, (wide + 63) / 64, 8), mtl.Size.of(128, 1, 1)); + try lib.commit(cb, p.r.enc); + const terms = try m.denseBf(); + defer for (terms) |v| a.free(v); + const wrong = try m.skewed(1.01).denseBf(); + defer for (wrong) |v| a.free(v); + const got = part.slice(f32, 8 * wide * n); + var worst = lib.Worst{}; + var ctrl = lib.Worst{}; + for (0..wide) |r| for (0..n) |j| { + var sum: f64 = 0; + for (0..8) |z| sum += got[(z * wide + r) * n + j]; + // each part is an fp32 sum of K / 8 products (any order: at most that many steps), the parts added here exactly + worst.add(j, sum, lib.dot3(x.row(r), terms, j)[0], lib.ratio3(sum, lib.dot3(x.row(r), terms, j), k / 8 + 64, false)); + ctrl.add(j, sum, lib.dot3(x.row(r), wrong, j)[0], lib.ratio3(sum, lib.dot3(x.row(r), wrong, j), k / 8 + 64, false)); + }; + worst.report(t, label); + ctrl.control(t, label); +} + +/// Sorted-expert gathers (rows grouped by expert, 8 experts used): the product and the fused gate/up activation. +fn gatherCheck(t: *lib.Tally, p: *fz.Prompt, wd: *pw.Wide, bits: u8, group: usize, seed: u64) !void { + var label_buf: [96]u8 = undefined; + const label = try std.fmt.bufPrint(&label_buf, "prompt expert gathers [8 x 128, 2560] q{d}/g{d}", .{ bits, group }); + const e_n = 8; + const n = 128; + const k = 2560; + const d = p.r.device; + const gm = try lib.Mat.init(e_n * n, k, bits, group, seed); + const um = try lib.Mat.init(e_n * n, k, bits, group, seed + 1); + defer gm.deinit(); + defer um.deinit(); + const gw = try weights(d, gm); + const uw = try weights(d, um); + const x = try lib.Inputs.init(d, wide, k, seed + 2); + const off = p.b.off.b.slice(i32, 512); + var expert: [wide]usize = undefined; + for (&expert, 0..) |*e, i| e.* = i * e_n / wide; // rows sorted by expert + for (off, 0..) |*o, e| { + o.* = wide; + for (expert, 0..) |x_e, i| if (x_e >= e) { + o.* = @intCast(i); + break; + }; + } + const kk = @as(u32, bits) << 16 | @as(u32, @intCast(group)); + if (!wd.pipes.contains(kk)) try wd.pipes.put(a, kk, try pw.compile(p.r, bits, @intCast(group))); + const pipes = wd.of(bits, @intCast(group)); + const y = try lib.poisoned(d, wide * n * 2); + const act = try lib.poisoned(d, wide * n * 2); + const cb = p.r.queue.commandBuffer(); + p.r.enc = cb.compute(.serial); + pw.gather(p, pipes, 32, .{ .b = x.buf }, &gw, k, n, wide, .{ .b = y }); + pw.gatherGU(p, pipes, 32, .{ .b = x.buf }, &gw, &uw, k, n, wide, .{ .b = act }); + try lib.commit(cb, p.r.enc); + var worst = [2]lib.Worst{ .{}, .{} }; // product, gate/up activation + var ctrl = [2]lib.Worst{ .{}, .{} }; + for ([_]f64{ 1, 1.01 }, 0..) |skew, which| { + const gt = try gm.skewed(skew).denseBf(); + defer for (gt) |v| a.free(v); + const ut = try um.skewed(skew).denseBf(); + defer for (ut) |v| a.free(v); + const w = if (which == 0) &worst else &ctrl; + for (0..wide) |r| for (0..n) |j| { + const g = lib.dot3(x.row(r), gt, expert[r] * n + j); + const u = lib.dot3(x.row(r), ut, expert[r] * n + j); + const got = lib.fromBf(y.slice(u16, wide * n)[r * n + j]); + w[0].add(j, got, g[0], lib.ratio3(got, g, k + 64, true)); + const sig = 1 / (1 + @exp(-g[0])); + const want = g[0] * sig * u[0]; + // the sums' fp32 and tie bounds, then the gate's bf16 rounding; silu's slope is below 1.1 + const dg = @as(f64, k + 64) * lib.unit * g[1] + g[2] + 0x1p-8 * @abs(g[0]); + const du = @as(f64, k + 64) * lib.unit * u[1] + u[2] + 0x1p-8 * @abs(u[0]); + const tol = 1.1 * @abs(u[0]) * dg + @abs(g[0] * sig) * du + 0x1p-7 * @abs(want) + 1e-6; + const ga = lib.fromBf(act.slice(u16, wide * n)[r * n + j]); + w[1].add(j, ga, want, @abs(ga - want) / tol); + }; + } + var act_label: [128]u8 = undefined; + const gu = try std.fmt.bufPrint(&act_label, "{s} gate/up", .{label}); + worst[0].report(t, label); + ctrl[0].control(t, label); + worst[1].report(t, gu); + ctrl[1].control(t, gu); +} + +pub fn run(t: *lib.Tally, r: *Run) !void { + const offsets = try r.buffer(512 * 4); + var p = prompt(r, .{ .b = offsets }); + const wd = try a.create(pw.Wide); + wd.* = .{}; + try recordedTwin(t, &p, wd); + var seed: u64 = 0x9000; + for ([_]usize{ 2, 3, 4, 5, 6, 8 }) |bits| for ([_]usize{ 32, 64, 128 }) |group| { + seed += 8; + try denseCheck(t, &p, wd, &.{.{ 192, bits, group }}, 2560, seed); + try denseCheck(t, &p, wd, &.{.{ 64, bits, group }}, 640, seed + 4); + try splitCheck(t, &p, wd, @intCast(bits), group, seed + 5); + try gatherCheck(t, &p, wd, @intCast(bits), group, seed + 6); + }; + // the oQ fused stacks, members at their own widths (q, k, v, index_qk and the DeltaNet inputs, rows cut down) + try denseCheck(t, &p, wd, &.{ .{ 256, 3, 64 }, .{ 32, 3, 64 }, .{ 32, 8, 64 }, .{ 64, 3, 64 } }, 2560, 0x9a01); + try denseCheck(t, &p, wd, &.{ .{ 256, 4, 64 }, .{ 64, 5, 128 }, .{ 32, 5, 128 }, .{ 32, 6, 64 } }, 2560, 0x9a02); + try denseCheck(t, &p, wd, &.{ .{ 128, 6, 64 }, .{ 32, 6, 64 }, .{ 32, 6, 128 }, .{ 64, 6, 64 } }, 2560, 0x9a03); + // oQ4e layer 21's DeltaNet inputs: the 5-bit run's b and the 6-bit a keep their 48 rows + try denseCheck(t, &p, wd, &.{ .{ 256, 4, 64 }, .{ 64, 5, 128 }, .{ 48, 5, 128 }, .{ 48, 6, 64 } }, 2560, 0x9a04); + // the prompt's other shapes: HC up (K 320), N 10240, and att.o and gdn.out (K 6144) at their oQ widths + for ([_]u8{ 2, 4, 8 }, 0..) |bits, i| try denseCheck(t, &p, wd, &.{.{ 256, bits, 64 }}, 320, 0x9b00 + i); + for ([_][2]usize{ .{ 3, 64 }, .{ 4, 64 }, .{ 5, 128 }, .{ 8, 64 } }, 0..) |w, i| try denseCheck(t, &p, wd, &.{.{ 64, w[0], w[1] }}, 6144, 0x9b10 + i); + try denseCheck(t, &p, wd, &.{.{ 10240, 3, 64 }}, 320, 0x9b20); +} 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; +}