From e73425a4d5b861c57d92a47f5086357e9fc33434 Mon Sep 17 00:00:00 2001 From: mrv777 Date: Sun, 11 Oct 2026 16:25:47 -0500 Subject: [PATCH 1/2] GLM: fix sparse attention thread limit on M1 --- tools/zig/gen_glm_kernels.py | 6 ++- zig/kernels/metal/glm/sparse_attention.metal | 1 + .../metal/glm/sparse_attention_tp.metal | 1 + zig/src/families/glm/attn_check.zig | 44 ++++++++++++------- 4 files changed, 35 insertions(+), 17 deletions(-) diff --git a/tools/zig/gen_glm_kernels.py b/tools/zig/gen_glm_kernels.py index edc54913a..9485fdc2c 100644 --- a/tools/zig/gen_glm_kernels.py +++ b/tools/zig/gen_glm_kernels.py @@ -66,6 +66,7 @@ class Spec: header: str = "" template: list = field(default_factory=list) extra: list = field(default_factory=list) # more template instantiations of the same function + threads: int = 0 def _digest(text: str) -> str: @@ -173,7 +174,8 @@ def router(rr: int) -> list: out.append(Spec(key, f"tf_glm5_indexed_sparse_attention_{_digest(SA._SOURCE)}", SA._SOURCE, [Arg("queries", bf, big, 3), Arg("keys", bf, big, 2), Arg("indices", i32, big, 2), Arg("scale", f32, one), Arg("meta", i32, one)], - [Arg("out", bf, big, 3)], "", [("QK_DIM", RANK), ("TOPK", TOPK_KEYS), ("HEADS", heads)])) + [Arg("out", bf, big, 3)], "", [("QK_DIM", RANK), ("TOPK", TOPK_KEYS), ("HEADS", heads)], + threads=1024)) return out @@ -181,6 +183,8 @@ def rendered(spec: Spec) -> tuple[str, str]: """(function name, file text) for a spec; extra instantiations get their own host names.""" fname, text = kernel_text(spec.name, spec.inputs, spec.outputs, spec.source, spec.header, spec.template) + if spec.threads: + text = text.replace("[[kernel]] void", f"[[max_total_threads_per_threadgroup({spec.threads})]]\n[[kernel]] void", 1) for template in spec.extra: host = function_name(spec.name, spec.inputs, spec.outputs, template) args = fname + "<" + ", ".join(_template_value(v) for _, v in template) + ">" diff --git a/zig/kernels/metal/glm/sparse_attention.metal b/zig/kernels/metal/glm/sparse_attention.metal index 057ab321f..a0848a9a5 100644 --- a/zig/kernels/metal/glm/sparse_attention.metal +++ b/zig/kernels/metal/glm/sparse_attention.metal @@ -16,6 +16,7 @@ inline bfloat16_t exp(bfloat16_t v) { return bfloat16_t(metal::precise::exp(floa } } template +[[max_total_threads_per_threadgroup(1024)]] [[kernel]] void custom_kernel_tf_glm5_indexed_sparse_attention_ad9e57bce4__512_2051_64_bfloat16_t_bfloat16_t_int32_t_floatc_int32_tc_bfloat16_t( const device bfloat16_t* queries [[buffer(0)]], const device bfloat16_t* keys [[buffer(1)]], diff --git a/zig/kernels/metal/glm/sparse_attention_tp.metal b/zig/kernels/metal/glm/sparse_attention_tp.metal index 5c0044ec4..29d06f354 100644 --- a/zig/kernels/metal/glm/sparse_attention_tp.metal +++ b/zig/kernels/metal/glm/sparse_attention_tp.metal @@ -16,6 +16,7 @@ inline bfloat16_t exp(bfloat16_t v) { return bfloat16_t(metal::precise::exp(floa } } template +[[max_total_threads_per_threadgroup(1024)]] [[kernel]] void custom_kernel_tf_glm5_indexed_sparse_attention_ad9e57bce4__512_2051_32_bfloat16_t_bfloat16_t_int32_t_floatc_int32_tc_bfloat16_t( const device bfloat16_t* queries [[buffer(0)]], const device bfloat16_t* keys [[buffer(1)]], diff --git a/zig/src/families/glm/attn_check.zig b/zig/src/families/glm/attn_check.zig index 82fdc9f31..aea2d08bb 100644 --- a/zig/src/families/glm/attn_check.zig +++ b/zig/src/families/glm/attn_check.zig @@ -19,7 +19,7 @@ fn f32of(h: u16) f32 { return @bitCast(@as(u32, h) << 16); } -test "sparse attention on the tensor units: the row kernel's values within bf16 rounding" { +fn sparseCheck(comptime heads: usize) !void { const device = mtl.Device.init() catch return error.SkipZigTest; defer device.deinit(); const pool = mtl.objc.Pool.push(); @@ -28,13 +28,15 @@ test "sparse attention on the tensor units: the row kernel's values within bf16 defer queue.deinit(); const a = std.testing.allocator; const gen = for (sources.glm.all) |k| { - if (std.mem.eql(u8, k.key, "sparse_attention")) break k; + if (std.mem.eql(u8, k.key, if (heads == 64) "sparse_attention" else "sparse_attention_tp")) break k; } else return error.MissingKernel; const lib_row = try mtl.Library.fromSource(device, gen.source, mtl.CompileOptions.mlx()); defer lib_row.deinit(); const row_pipe = try mtl.Pipeline.init(device, lib_row, gen.functions[0], false); defer row_pipe.deinit(); - const nax_src = try frags.source(device, a, sources.glm_sparse_nax); + try std.testing.expect(row_pipe.maxThreads() >= 1024); + const raw = if (heads == 32) "#define GLM_HEADS 32\n" ++ sources.glm_sparse_nax else sources.glm_sparse_nax; + const nax_src = try frags.source(device, a, raw); defer a.free(nax_src); const lib_nax = try mtl.Library.fromSource(device, nax_src, mtl.CompileOptions.mlx()); defer lib_nax.deinit(); @@ -42,12 +44,12 @@ test "sparse attention on the tensor units: the row kernel's values within bf16 defer nax_pipe.deinit(); const opts = mtl.ResourceOptions.shared; - const bufs = [_]mtl.Buffer{ try device.buffer(ROWS * H * RANK * 2, opts), try device.buffer(CAP * RANK * 2, opts), try device.buffer(ROWS * WIDTH * 4, opts), try device.buffer(ROWS * H * RANK * 2, opts), try device.buffer(ROWS * H * RANK * 2, opts) }; + const bufs = [_]mtl.Buffer{ try device.buffer(ROWS * heads * RANK * 2, opts), try device.buffer(CAP * RANK * 2, opts), try device.buffer(ROWS * WIDTH * 4, opts), try device.buffer(ROWS * heads * RANK * 2, opts), try device.buffer(ROWS * heads * RANK * 2, opts) }; defer for (bufs) |b| b.deinit(); var prng = std.Random.DefaultPrng.init(11); const rnd = prng.random(); - const q = bufs[0].slice(u16, ROWS * H * RANK); - for (q, 0..) |*v, i| v.* = bf16((rnd.float(f32) * 4 - 2) * (if (i / (H * RANK) == 4) @as(f32, 6) else 1)); // row 4: peaked + const q = bufs[0].slice(u16, ROWS * heads * RANK); + for (q, 0..) |*v, i| v.* = bf16((rnd.float(f32) * 4 - 2) * (if (i / (heads * RANK) == 4) @as(f32, 6) else 1)); // row 4: peaked const keys = bufs[1].slice(u16, CAP * RANK); for (keys) |*v| v.* = bf16(rnd.float(f32) * 2 - 1); const lists = bufs[2].slice(i32, ROWS * WIDTH); @@ -71,29 +73,29 @@ test "sparse attention on the tensor units: the row kernel's values within bf16 enc.setValue(scale, 3); enc.setValue(@as(i32, CAP), 4); enc.setBuffer(bufs[3], 0, 5); - enc.dispatchThreads(mtl.Size.of(1024, ROWS * H, 1), mtl.Size.of(1024, 1, 1)); + enc.dispatchThreads(mtl.Size.of(1024, ROWS * heads, 1), mtl.Size.of(1024, 1, 1)); enc.setPipeline(nax_pipe); for (bufs[0..3], 0..) |b, i| enc.setBuffer(b, 0, i); enc.setValue(scale, 3); enc.setValue([4]i32{ WIDTH, CAP, 0, 0 }, 4); enc.setBuffer(bufs[4], 0, 5); - enc.dispatchGroups(mtl.Size.of(H / 16, ROWS, 1), mtl.Size.of(128, 1, 1)); + enc.dispatchGroups(mtl.Size.of(heads / 16, ROWS, 1), mtl.Size.of(128, 1, 1)); enc.end(); cb.commit(); cb.wait(); try std.testing.expect(cb.failure() == null); - const row_out = bufs[3].slice(u16, ROWS * H * RANK); - const nax_out = bufs[4].slice(u16, ROWS * H * RANK); + const row_out = bufs[3].slice(u16, ROWS * heads * RANK); + const nax_out = bufs[4].slice(u16, ROWS * heads * RANK); const scores = try a.alloc(f64, WIDTH); defer a.free(scores); var same: usize = 0; var worst: f64 = 0; // the largest difference from the row kernel's output, in bf16 steps of the larger value const outs = try a.alloc(f64, RANK); defer a.free(outs); - for (0..ROWS) |r| for (0..H) |h| { + for (0..ROWS) |r| for (0..heads) |h| { // the host's fp64 softmax over the valid listed keys - const qr = q[(r * H + h) * RANK ..][0..RANK]; + const qr = q[(r * heads + h) * RANK ..][0..RANK]; var top: f64 = -std.math.inf(f64); for (lists[r * WIDTH ..][0..WIDTH], scores) |k, *sc| { sc.* = -std.math.inf(f64); @@ -118,20 +120,30 @@ test "sparse attention on the tensor units: the row kernel's values within bf16 big = @max(big, @abs(outs[d])); } for (0..RANK) |d| { - const i = (r * H + h) * RANK + d; + const i = (r * heads + h) * RANK + d; const got: f64 = f32of(nax_out[i]); const theirs: f64 = f32of(row_out[i]); // a bf16 step of the value, but no finer than fp32 sums of the head's values can tell (2^-16 of its largest) const floor = big / 65536.0; - try std.testing.expect(@abs(got - outs[d]) <= @max(@abs(outs[d]) / 128.0, floor)); + const bound = @max(@abs(outs[d]) / 128.0, floor); + try std.testing.expect(@abs(got - outs[d]) <= bound); + try std.testing.expect(@abs(theirs - outs[d]) <= bound); if (nax_out[i] == row_out[i]) same += 1 else worst = @max(worst, @abs(got - theirs) / @max(@max(@abs(got), @abs(theirs)) / 128.0, floor)); } }; - const share = @as(f64, @floatFromInt(same)) / @as(f64, @floatFromInt(ROWS * H * RANK)); - std.debug.print("sparse attention on the tensor units: {d:.2}% of outputs equal the row kernel's, at most {d:.2} bf16 steps apart\n", .{ share * 100, worst }); + const share = @as(f64, @floatFromInt(same)) / @as(f64, @floatFromInt(ROWS * heads * RANK)); + std.debug.print("sparse attention ({d} heads): {d:.2}% of outputs equal the row kernel's, at most {d:.2} bf16 steps apart\n", .{ heads, share * 100, worst }); try std.testing.expect(share > 0.99 and worst <= 2); } +test "sparse attention on the tensor units: the row kernel's values within bf16 rounding" { + try sparseCheck(64); +} + +test "TP2 sparse attention: 32 heads against the row kernel and fp64" { + try sparseCheck(32); +} + test "the MLA absorb on the tensor units: the row kernel's values within bf16 rounding" { const device = mtl.Device.init() catch return error.SkipZigTest; defer device.deinit(); From 434a62fc46a60d1458b7966ad970d05f44050ff9 Mon Sep 17 00:00:00 2001 From: mrv777 Date: Sun, 11 Oct 2026 17:38:42 -0500 Subject: [PATCH 2/2] GLM: fix KDA and router thread limits on M1 --- tools/zig/gen_glm_kernels.py | 9 +- zig/kernels/metal/glm/kda_rows.metal | 1 + zig/kernels/metal/glm/kda_rows_tp.metal | 1 + zig/kernels/metal/glm/router.metal | 1 + zig/src/families/glm/glm.zig | 1 + zig/src/families/glm/launch_check.zig | 350 ++++++++++++++++++++++++ 6 files changed, 359 insertions(+), 4 deletions(-) create mode 100644 zig/src/families/glm/launch_check.zig diff --git a/tools/zig/gen_glm_kernels.py b/tools/zig/gen_glm_kernels.py index 9485fdc2c..21a8fe8e0 100644 --- a/tools/zig/gen_glm_kernels.py +++ b/tools/zig/gen_glm_kernels.py @@ -88,8 +88,9 @@ def specs() -> list[Spec]: out: list[Spec] = [] def plain(key: str, base: str, body: str, header: str, inputs: list[Arg], outputs: list[Arg], template: list, - extra: list | None = None) -> None: - out.append(Spec(key, f"{base}_{_digest(header + body)}", body, inputs, outputs, header, template, extra or [])) + extra: list | None = None, threads: int = 0) -> None: + out.append(Spec(key, f"{base}_{_digest(header + body)}", body, inputs, outputs, + header, template, extra or [], threads=threads)) qmv_in = [Arg("X", bf, big, 2), Arg("W", u32, big, 2), Arg("S", bf, big, 2), Arg("B", bf, big, 2)] for k, n in PROJECTIONS: @@ -122,13 +123,13 @@ def plain(key: str, base: str, body: str, header: str, inputs: list[Arg], output Arg("GBB", bf, big, 2), Arg("A", f32, big), Arg("DTB", f32, big), Arg("ST", f32, big, 4), Arg("ONW", f32, big), Arg("LB", f32, one), Arg("EPS", f32, one)], [Arg("Y", bf, big, 2), Arg("ST_OUT", f32, big, 4), Arg("CS_OUT", bf, big, 2)], - [("H", heads), ("D", KDA_DIM), ("TAPS", TAPS), ("TY", KD.TY), ("FB", 4), ("GB", 4)]) + [("H", heads), ("D", KDA_DIM), ("TAPS", TAPS), ("TY", KD.TY), ("FB", 4), ("GB", 4)], threads=1024) def router(rr: int) -> list: return [("K", HIDDEN), ("NE", EXPERTS), ("RR", rr), ("C", 8 if rr <= 4 else 4), ("NT", M.ROUTER_TG)] plain("router", "tf_glm5_fused_router_tg", M._ROUTER_TG, F._HEADER, [Arg("X", f32, big, 2), Arg("RP", bf, big, 5)], - [Arg("OUT", f32, big, 2)], router(1), [router(rr) for rr in range(2, MAX_ROWS + 1)]) + [Arg("OUT", f32, big, 2)], router(1), [router(rr) for rr in range(2, MAX_ROWS + 1)], threads=M.ROUTER_TG) plain("moe_route", "tf_glm5_fused_moe_route", M._MOE_ROUTE, F._HEADER, [Arg("LOGITS", f32, big, 2), Arg("BIAS", f32, big), Arg("SCALE", f32, one)], [Arg("PICK", i32, big, 2), Arg("WTS", f32, big, 2), Arg("UIDS", i32, big), Arg("UMEM", i32, big, 2), diff --git a/zig/kernels/metal/glm/kda_rows.metal b/zig/kernels/metal/glm/kda_rows.metal index 9a6b0e658..e8652269f 100644 --- a/zig/kernels/metal/glm/kda_rows.metal +++ b/zig/kernels/metal/glm/kda_rows.metal @@ -92,6 +92,7 @@ inline float sq_acc(float acc, float v) { } #pragma clang fp contract(on) template +[[max_total_threads_per_threadgroup(1024)]] [[kernel]] void custom_kernel_tf_glm5_kda_rows_918a341e71__64_128_4_32_4_4_bfloat16_t_bfloat16_t_float_uint32_t_bfloat16_t_bfloat16_t_uint32_t_bfloat16_t_bfloat16_t_float_float_float_float_floatc_floatc_bfloat16_t_float_bfloat16_t( const device bfloat16_t* P [[buffer(0)]], const constant int* P_shape [[buffer(1)]], diff --git a/zig/kernels/metal/glm/kda_rows_tp.metal b/zig/kernels/metal/glm/kda_rows_tp.metal index abfb2043a..be91d2cad 100644 --- a/zig/kernels/metal/glm/kda_rows_tp.metal +++ b/zig/kernels/metal/glm/kda_rows_tp.metal @@ -92,6 +92,7 @@ inline float sq_acc(float acc, float v) { } #pragma clang fp contract(on) template +[[max_total_threads_per_threadgroup(1024)]] [[kernel]] void custom_kernel_tf_glm5_kda_rows_918a341e71__32_128_4_32_4_4_bfloat16_t_bfloat16_t_float_uint32_t_bfloat16_t_bfloat16_t_uint32_t_bfloat16_t_bfloat16_t_float_float_float_float_floatc_floatc_bfloat16_t_float_bfloat16_t( const device bfloat16_t* P [[buffer(0)]], const constant int* P_shape [[buffer(1)]], diff --git a/zig/kernels/metal/glm/router.metal b/zig/kernels/metal/glm/router.metal index f288e81b0..e67478dca 100644 --- a/zig/kernels/metal/glm/router.metal +++ b/zig/kernels/metal/glm/router.metal @@ -172,6 +172,7 @@ inline float mul_add(float acc, float a, float b) { return a * b + acc; } inline float add_nc(float a, float b) { return a + b; } #pragma clang fp contract(on) template +[[max_total_threads_per_threadgroup(1024)]] [[kernel]] void custom_kernel_tf_glm5_fused_router_tg_331f012361__4096_288_1_8_1024_float_bfloat16_t_float( const device float* X [[buffer(0)]], const device bfloat16_t* RP [[buffer(1)]], diff --git a/zig/src/families/glm/glm.zig b/zig/src/families/glm/glm.zig index 34bf628cc..25bbe6b5c 100644 --- a/zig/src/families/glm/glm.zig +++ b/zig/src/families/glm/glm.zig @@ -25,6 +25,7 @@ test { _ = @import("../../core/paced_read.zig"); _ = @import("../../core/copy_index.zig"); _ = @import("attn_check.zig"); + _ = @import("launch_check.zig"); _ = @import("load_plan.zig"); _ = @import("page_cache.zig"); _ = pool; diff --git a/zig/src/families/glm/launch_check.zig b/zig/src/families/glm/launch_check.zig new file mode 100644 index 000000000..76a17d300 --- /dev/null +++ b/zig/src/families/glm/launch_check.zig @@ -0,0 +1,350 @@ +//! GLM's 1,024-thread router and KDA launches: capacity, host values, full coverage and recurrent seams. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); + +fn kernel(comptime key: []const u8) sources.glm.Kernel { + for (sources.glm.all) |k| if (std.mem.eql(u8, key, k.key)) return k; + unreachable; +} + +fn bf(v: f32) u16 { + const u: u32 = @bitCast(v); + return @intCast((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +fn value(v: u16) f32 { + return @bitCast(@as(u32, v) << 16); +} + +fn dyadic(i: usize) f32 { + return @as(f32, @floatFromInt(@as(i32, @intCast(i % 17)) - 8)) / 32; +} + +fn routerWeight(e: usize, k: usize) f32 { + return dyadic(e * 11 + k * 7 + k / 32); +} + +test "GLM router: every row count supports 1024 threads and equals FP64" { + const d = mtl.Device.init() catch return error.SkipZigTest; + defer d.deinit(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const q = try d.queue(); + defer q.deinit(); + const gen = comptime kernel("router"); + const lib = try mtl.Library.fromSource(d, gen.source, mtl.CompileOptions.mlx()); + defer lib.deinit(); + const K = 4096; + const NE = 288; + const R = 16; + const x = try d.buffer(R * K * 4, mtl.ResourceOptions.shared); + defer x.deinit(); + const w = try d.buffer(NE * K * 2, mtl.ResourceOptions.shared); + defer w.deinit(); + const y = try d.buffer((R * NE + 32) * 4, mtl.ResourceOptions.shared); + defer y.deinit(); + const xs = x.slice(f32, R * K); + for (xs, 0..) |*v, i| v.* = dyadic((i / K) * 7 + (i % K) * 5); + const weights = w.slice(u16, NE * K); + for (0..NE / 4) |g| for (0..8) |m| for (0..K / 32) |it| for (0..4) |tm| for (0..4) |tn| { + const at = ((((g * 8 + m) * (K / 32) + it) * 4 + tm) * 4 + tn); + weights[at] = bf(routerWeight(g * 4 + tn, it * 32 + m * 4 + tm)); + }; + const refs = try std.testing.allocator.alloc(f64, R * NE); + defer std.testing.allocator.free(refs); + for (refs, 0..) |*v, i| { + v.* = 0; + for (0..K) |k| v.* += @as(f64, xs[(i / NE) * K + k]) * routerWeight(i % NE, k); + } + for (gen.functions, 1..) |name, rows| { + const p = try mtl.Pipeline.init(d, lib, name, false); + defer p.deinit(); + try std.testing.expect(p.maxThreads() >= 1024); + @memset(y.slice(u32, R * NE + 32), 0x7fc00000); + const cb = q.commandBuffer(); + const e = cb.compute(.serial); + e.setPipeline(p); + e.setBuffer(x, 0, 0); + e.setBuffer(w, 0, 1); + e.setBuffer(y, 0, 2); + e.dispatchThreads(mtl.Size.of(1024 * NE / 16, 1, 1), mtl.Size.of(1024, 1, 1)); + e.end(); + cb.commit(); + cb.wait(); + try std.testing.expect(cb.failure() == null); + for (y.slice(f32, R * NE + 32)[0 .. rows * NE], refs[0 .. rows * NE]) |actual, expected| { + try std.testing.expect(std.math.isFinite(actual)); + try std.testing.expectEqual(expected, @as(f64, actual)); + } + for (y.slice(u32, R * NE + 32)[rows * NE ..]) |bits| try std.testing.expectEqual(@as(u32, 0x7fc00000), bits); + } +} + +const Kda = struct { + d: mtl.Device, + q: mtl.Queue, + p: mtl.Pipeline, + b: [19]mtl.Buffer, + h: usize, + w: usize, + pitch: usize, + + fn init(d: mtl.Device, q: mtl.Queue, p: mtl.Pipeline, h: usize) !Kda { + const w = h * 128; + const pitch = 3 * w + 256 + h; + const sizes = [_]usize{ 16 * pitch * 2, 8, 3 * 3 * w * 2, 4 * 3 * w * 4, w * 16 * 4, w * 2 * 2, w * 2 * 2, w * 16 * 4, w * 2 * 2, w * 2 * 2, h * 4, w * 4, w * 128 * 4, 128 * 4, 4, 4, 16 * w * 2, w * 128 * 4, 3 * 3 * w * 2 }; + var k = Kda{ .d = d, .q = q, .p = p, .b = undefined, .h = h, .w = w, .pitch = pitch }; + var made: usize = 0; + errdefer for (k.b[0..made]) |b| b.deinit(); + for (sizes) |n| { + k.b[made] = try d.buffer(n, mtl.ResourceOptions.shared); + @memset(k.b[made].contents()[0..n], 0); + made += 1; + } + k.scalarInputs(); + return k; + } + + fn deinit(k: Kda) void { + for (k.b) |b| b.deinit(); + } + + fn scalarInputs(k: Kda) void { + const ps = k.b[0].slice(u16, 16 * k.pitch); + for (0..16) |r| { + for (0..3 * k.w) |i| ps[r * k.pitch + i] = bf(activation(@intCast(r), i / k.w, (i % k.w) / 128, i % 128)); + @memset(ps[r * k.pitch + 3 * k.w ..][0..256], bf(1.0 / 16.0)); + for (0..k.h) |h| ps[r * k.pitch + 3 * k.w + 256 + h] = bf(betaInput(r, h)); + } + const cs = k.b[2].slice(u16, 9 * k.w); + for (cs, 0..) |*v, i| v.* = bf(activation(@as(i32, @intCast(i / (3 * k.w))) - 3, (i % (3 * k.w)) / k.w, (i % k.w) / 128, i % 128)); + @memset(k.b[3].slice(f32, 12 * k.w), 0.25); + // These packed projections give constant decay input 0.1875 and output gate 0.4375. + @memset(k.b[4].slice(u32, k.w * 16), 0x11111111); + @memset(k.b[5].slice(u16, k.w * 2), bf(1.0 / 64.0)); + @memset(k.b[6].slice(u16, k.w * 2), bf(1.0 / 128.0)); + @memset(k.b[7].slice(u32, k.w * 16), 0x22222222); + @memset(k.b[8].slice(u16, k.w * 2), bf(1.0 / 32.0)); + @memset(k.b[9].slice(u16, k.w * 2), bf(-1.0 / 128.0)); + @memset(k.b[10].slice(f32, k.h), 1); + for (k.b[11].slice(f32, k.w), 0..) |*v, i| v.* = dtBias(i / 128); + k.seed(); + for (k.b[13].slice(f32, 128), 0..) |*v, i| v.* = gain(i); + k.b[14].slice(f32, 1)[0] = -5; + k.b[15].slice(f32, 1)[0] = 1e-5; + } + + fn seed(k: Kda) void { + for (k.b[12].slice(f32, k.w * 128), 0..) |*v, i| v.* = initial(i / (128 * 128), (i / 128) % 128); + } + + fn run(k: Kda, p: mtl.Pipeline, first: usize, rows: usize) !void { + try k.runGroup(p, first, rows, 32); + } + + fn runGroup(k: Kda, p: mtl.Pipeline, first: usize, rows: usize, ty: usize) !void { + @memcpy(k.b[1].slice(i32, 2), &[_]i32{ @intCast(rows), @intCast(k.pitch) }); + @memset(k.b[16].slice(u16, 16 * k.w), 0x7fc0); + @memset(k.b[17].slice(u32, k.w * 128), 0x7fc00000); + @memset(k.b[18].slice(u16, 9 * k.w), 0x7fc0); + const cb = k.q.commandBuffer(); + const e = cb.compute(.serial); + e.setPipeline(p); + for (k.b, 0..) |b, i| e.setBuffer(b, if (i == 0) first * k.pitch * 2 else 0, i); + e.dispatchThreads(mtl.Size.of(32, ty, k.h), mtl.Size.of(32, ty, 1)); + e.end(); + cb.commit(); + cb.wait(); + try std.testing.expect(cb.failure() == null); + for (k.b[16].slice(u16, 16 * k.w)[0 .. rows * k.w]) |v| try std.testing.expect(std.math.isFinite(value(v))); + for (k.b[16].slice(u16, 16 * k.w)[rows * k.w ..]) |v| try std.testing.expectEqual(@as(u16, 0x7fc0), v); + for (k.b[17].slice(f32, k.w * 128)) |v| try std.testing.expect(std.math.isFinite(v)); + for (k.b[18].slice(u16, 9 * k.w)) |v| try std.testing.expect(std.math.isFinite(value(v))); + } + + fn reference(k: Kda, rows: usize) !void { + var worst_y: f64 = 0; + var worst_state: f64 = 0; + for (0..k.h) |h| { + var state: [128]f64 = undefined; + for (&state, 0..) |*v, dv| v.* = initial(h, dv); + const decay = @exp(-5.0 / (1.0 + @exp(-(@as(f64, dtBias(h)) + 0.1875)))); + for (0..rows) |r| { + var y: [128]f64 = undefined; + var norm: f64 = 0; + const qv = sl(@as(f64, activation(0, 0, h, 0))); + const kv = sl(@as(f64, activation(0, 1, h, 0))); + const qb: f64 = value(bf(@floatCast(qv / @sqrt(128 * qv * qv + 1e-6) / @sqrt(@as(f64, 128))))); + const kb: f64 = value(bf(@floatCast(kv / @sqrt(128 * kv * kv + 1e-6)))); + const beta = sigmoidBf(betaInput(r, h)); + for (&state, &y, 0..) |*st, *out, dv| { + var conv: f64 = 0; + for (0..4) |tap| conv += @as(f64, activation(@as(i32, @intCast(r + tap)) - 3, 2, h, dv)) * 0.25; + const vv = sl(value(bf(@floatCast(conv)))); + st.* *= decay; + st.* += kb * (vv - 128 * st.* * kb) * beta; + out.* = value(bf(@floatCast(128 * st.* * qb))); + norm += out.* * out.*; + } + const inv = 1.0 / @sqrt(norm / 128 + @as(f64, k.b[15].slice(f32, 1)[0])); + for (y, 0..) |out, dv| { + const expected = out * inv * gain(dv) / (1 + @exp(-@as(f64, 0.4375))); + const actual = value(k.b[16].slice(u16, rows * k.w)[r * k.w + h * 128 + dv]); + const err = @abs(@as(f64, actual) - expected); + const unit = @as(u32, bf(@floatCast(@abs(expected))) & 0x7f80) << 16; + const step: f64 = @as(f32, @bitCast(unit)) / 128; + worst_y = @max(worst_y, err / @max(step, 1e-8)); + try std.testing.expect(err <= 2 * step + 1e-6); + } + } + for (state, 0..) |expected, dv| for (0..128) |dk| { + const actual = k.b[17].slice(f32, k.w * 128)[(h * 128 + dv) * 128 + dk]; + worst_state = @max(worst_state, @abs(@as(f64, actual) - expected)); + try std.testing.expectApproxEqAbs(expected, @as(f64, actual), 1e-5); + }; + } + for (k.b[18].slice(u16, 9 * k.w), 0..) |v, i| { + const row = @as(i32, @intCast(rows + i / (3 * k.w))) - 3; + const expected = bf(activation(row, (i % (3 * k.w)) / k.w, (i % k.w) / 128, i % 128)); + try std.testing.expectEqual(expected, v); + } + std.debug.print("KDA heads={d} rows={d}: worst {d:.4} bf16 steps, state {d:.9}\n", .{ k.h, rows, worst_y, worst_state }); + } + + fn seams(k: Kda) !void { + try k.run(k.p, 0, 16); + const a = std.testing.allocator; + const y = try a.dupe(u16, k.b[16].slice(u16, 16 * k.w)); + defer a.free(y); + const st = try a.dupe(u8, k.b[17].contents()[0 .. k.w * 128 * 4]); + defer a.free(st); + const cs = try a.dupe(u8, k.b[18].contents()[0 .. 9 * k.w * 2]); + defer a.free(cs); + var first: usize = 0; + for ([_]usize{ 1, 2, 5, 8 }) |rows| { + try k.run(k.p, first, rows); + try std.testing.expectEqualSlices(u16, y[first * k.w ..][0 .. rows * k.w], k.b[16].slice(u16, rows * k.w)); + @memcpy(k.b[12].contents()[0..st.len], k.b[17].contents()[0..st.len]); + @memcpy(k.b[2].contents()[0..cs.len], k.b[18].contents()[0..cs.len]); + first += rows; + } + try std.testing.expectEqualSlices(u8, st, k.b[17].contents()[0..st.len]); + try std.testing.expectEqualSlices(u8, cs, k.b[18].contents()[0..cs.len]); + } + + fn randomControl(k: Kda, source: []const u8) !void { + const a = std.testing.allocator; + const from = try std.fmt.allocPrint(a, "<{d}, 128, 4, 32, 4, 4>", .{k.h}); + defer a.free(from); + const to = try std.fmt.allocPrint(a, "<{d}, 128, 4, 16, 4, 4>", .{k.h}); + defer a.free(to); + try std.testing.expectEqual(@as(usize, 2), std.mem.count(u8, source, from)); + const reduced = try std.mem.replaceOwned(u8, a, source, from, to); + defer a.free(reduced); + const control = try std.mem.replaceOwned(u8, a, reduced, "max_total_threads_per_threadgroup(1024)", "max_total_threads_per_threadgroup(512)"); + defer a.free(control); + const lib = try mtl.Library.fromSource(k.d, control, mtl.CompileOptions.mlx()); + defer lib.deinit(); + const gen = if (k.h == 64) comptime kernel("kda_rows") else comptime kernel("kda_rows_tp"); + const p = try mtl.Pipeline.init(k.d, lib, gen.functions[0], false); + defer p.deinit(); + try std.testing.expect(p.maxThreads() >= 512); + var prng = std.Random.DefaultPrng.init(679); + const rnd = prng.random(); + for ([_]usize{ 0, 2 }) |idx| { + for (k.b[idx].slice(u16, k.b[idx].length() / 2)) |*v| v.* = bf(dyadic(rnd.int(u16))); + } + for ([_]usize{ 4, 7 }) |idx| { + for (k.b[idx].slice(u32, k.w * 16)) |*v| v.* = rnd.int(u32); + } + for ([_]usize{ 5, 8 }) |idx| { + for (k.b[idx].slice(u16, k.w * 2)) |*v| v.* = bf(@as(f32, @floatFromInt(rnd.int(u8) % 3 + 1)) / 64); + } + for ([_]usize{ 6, 9 }) |idx| { + for (k.b[idx].slice(u16, k.w * 2)) |*v| v.* = bf(dyadic(rnd.int(u16)) / 4); + } + for (k.b[3].slice(f32, 12 * k.w)) |*v| v.* = dyadic(rnd.int(u16)); + for (k.b[12].slice(f32, k.w * 128)) |*v| v.* = dyadic(rnd.int(u16)) / 8; + // TY=16 covers all 128 state rows with the same per-value reductions at a legal 512-thread launch. + for ([_]usize{ 1, 2, 8, 16 }) |rows| { + try k.run(k.p, 0, rows); + var saved: [3][]u8 = undefined; + var made: usize = 0; + defer for (saved[0..made]) |bytes| a.free(bytes); + for ([_]usize{ 16, 17, 18 }, 0..) |idx, i| { + saved[i] = try a.dupe(u8, k.b[idx].contents()[0..k.b[idx].length()]); + made += 1; + } + try k.runGroup(p, 0, rows, 16); + for ([_]usize{ 16, 17, 18 }, saved) |idx, bytes| { + try std.testing.expectEqualSlices(u8, bytes, k.b[idx].contents()[0..bytes.len]); + } + } + } +}; + +fn activation(r: i32, part: usize, h: usize, dv: usize) f32 { + return switch (part) { + 0 => @as(f32, @floatFromInt(h % 4 + 1)) / 8, + 1 => @as(f32, @floatFromInt(h % 5 + 1)) / 8, + else => @as(f32, @floatFromInt(r + 4 + @as(i32, @intCast((h + dv) % 7)))) / 16, + }; +} + +fn betaInput(r: usize, h: usize) f32 { + return @as(f32, @floatFromInt(@as(i32, @intCast((r + h) % 3)) - 1)) / 4; +} + +fn dtBias(h: usize) f32 { + return @as(f32, @floatFromInt(@as(i32, @intCast(h % 5)) - 2)) / 8; +} + +fn gain(dv: usize) f32 { + return 1 + @as(f32, @floatFromInt(dv % 7)) / 16; +} + +fn initial(h: usize, dv: usize) f32 { + return @as(f32, @floatFromInt(@as(i32, @intCast((h + dv) % 9)) - 4)) / 256; +} + +fn sigmoidBf(x: f32) f64 { + const e = value(bf(@floatCast(@exp(@abs(@as(f64, x)))))); + const y = value(bf(1 / value(bf(1 + e)))); + return if (x < 0) y else value(bf(1 - y)); +} + +fn sl(x: f64) f64 { + return value(bf(@floatCast(x * sigmoidBf(@floatCast(x))))); +} + +fn kdaCheck(comptime heads: usize) !void { + const d = mtl.Device.init() catch return error.SkipZigTest; + defer d.deinit(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const q = try d.queue(); + defer q.deinit(); + const gen = comptime kernel(if (heads == 64) "kda_rows" else "kda_rows_tp"); + const lib = try mtl.Library.fromSource(d, gen.source, mtl.CompileOptions.mlx()); + defer lib.deinit(); + const p = try mtl.Pipeline.init(d, lib, gen.functions[0], false); + defer p.deinit(); + try std.testing.expect(p.maxThreads() >= 1024); + const k = try Kda.init(d, q, p, heads); + defer k.deinit(); + for ([_]usize{ 1, 2, 8, 16 }) |rows| { + try k.run(p, 0, rows); + try k.reference(rows); + } + try k.randomControl(gen.source); + try k.seams(); +} + +test "GLM KDA 64 heads: 1024 threads, FP64 values and recurrent seams" { + try kdaCheck(64); +} + +test "GLM KDA TP2 32 heads: 1024 threads, FP64 values and recurrent seams" { + try kdaCheck(32); +}