GLM: fix KDA and router thread limits on M1

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:
@@ -87,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:
@@ -121,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),
@@ -173,14 +175,17 @@ 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
 
 
 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) + ">"

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 <int H, int D, int TAPS, int TY, int FB, int GB>
+[[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)]],

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 <int H, int D, int TAPS, int TY, int FB, int GB>
+[[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)]],

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 <int K, int NE, int RR, int C, int NT>
+[[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)]],

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 <int QK_DIM, int TOPK, int HEADS>
+[[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)]],

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 <int QK_DIM, int TOPK, int HEADS>
+[[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)]],

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,26 +28,28 @@ 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();
     const nax_pipe = try mtl.Pipeline.init(device, lib_nax, "glm_sparse_nax", false);
     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();

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;