qwen27: score and decide passes queued onto the lane host's idle engine thread

zig/build/qwen27.zig

@@ -16,7 +16,7 @@ pub const Shared = struct {
 };
 
 /// The 27B's Metal sources, embedded by name (zig/kernels/metal/qwen27).
-const parts = [_][]const u8{ "xsum", "norm_xs", "norm_input", "mlp_xs", "mlp_ane", "glue", "copy", "attention_mpp", "attention_io", "attention_prompt", "prompt_glue", "prompt_state" };
+const parts = [_][]const u8{ "xsum", "norm_xs", "norm_input", "mlp_xs", "mlp_ane", "glue", "copy", "moe", "attention_mpp", "attention_io", "attention_prompt", "prompt_glue", "prompt_state" };
 
 /// The shared modules, imported into `engine` (whose `core` is itself) for the server and the release binary.
 pub fn engineModules(b: *std.Build, target: std.Build.ResolvedTarget, metal: *std.Build.Module, sources: *std.Build.Module, engine: *std.Build.Module) Shared {
@@ -99,7 +99,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu
     const see = program(b, "tf-lanes-see-check", "Build synthetic shared-context checks, no GPU execution", "zig/tests/lanes_see.zig", target, false, &.{ .{ .name = "metal", .module = metal }, .{ .name = "shared_attention", .module = attention }, .{ .name = "shared_kv", .module = kv_layout } });
     b.step("test-lanes-see", "Run small shared-context, mask, position and KV append device checks").dependOn(&b.addRunArtifact(see).step);
 
-    const qwen_core = b.createModule(.{ .root_source_file = b.path("zig/qwen27_core.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true, .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "tokenizer", .module = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true }) } } });
+    const qwen_core = b.createModule(.{ .root_source_file = b.path("zig/qwen27_core.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true, .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "kernel_sources", .module = sources }, .{ .name = "tokenizer", .module = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true }) } } });
     const qwen27 = b.createModule(.{ .root_source_file = b.path("zig/src/families/qwen27/qwen27.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true, .imports = &.{ .{ .name = "core", .module = qwen_core }, .{ .name = "metal", .module = metal } } });
     for ([_]struct { []const u8, *std.Build.Module }{ .{ "core", qwen_core }, .{ "core_lane", core_lane }, .{ "core_row", core_row }, .{ "bf16_topk", s.bf16_topk }, .{ "bf16_topk_gpu", s.bf16_topk_gpu }, .{ "tree_round", s.tree_round }, .{ "tree_round_gpu", s.tree_round_gpu }, .{ "tree_commit_gpu", s.tree_commit }, .{ "draft_ops", s.draft_ops }, .{ "kernel_sources", sources } }) |m| qwen_core.addImport(m[0], m[1]);
     for ([_]struct { []const u8, *std.Build.Module }{ .{ "core_lane", core_lane }, .{ "qwen_runtime_sources", s.runtime_sources }, .{ "core_delta", s.delta }, .{ "qwen_gdn_source", s.gdn_source } }) |m| qwen27.addImport(m[0], m[1]);

zig/kernels/metal/core/row_projection.metal

@@ -190,53 +190,102 @@ template <int RPS, int SG, bool RELU2, int BITS, bool BF16SUM>
   if (lane == 0)
     for (int j = 0; j < RPS; j++) out[j] = tf_rp_out<RELU2>(acc[0][j]);
 }
+// tf_row_projection_indexed with an extra matrix at id `experts` (buffers 7-9): a shared expert among the routed ones.
+template <int RPS, int SG, bool RELU2, int BITS, bool BF16SUM>
+[[kernel]] void tf_row_projection_indexed_extra(const device bfloat* X [[buffer(0)]],
+                                          const device uint32_t* W [[buffer(1)]],
+                                          const device bfloat* S [[buffer(2)]],
+                                          const device bfloat* B [[buffer(3)]],
+                                          constant int* dims [[buffer(4)]],
+                                          device bfloat* OUT [[buffer(5)]],
+                                          const device uint* ids [[buffer(6)]],
+                                          const device uint32_t* W2 [[buffer(7)]],
+                                          const device bfloat* S2 [[buffer(8)]],
+                                          const device bfloat* B2 [[buffer(9)]],
+                                          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], experts = dims[4], repeat = dims[5];
+  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;
+  const uint slot = threadgroup_position_in_grid.y;
+  const uint expert = ids[slot];
+  device bfloat* out = OUT + size_t(slot) * N + row0;
+  if (expert > uint(experts)) {
+    if (lane < uint(RPS)) out[lane] = bfloat(0.0f);
+    return;
+  }
+  // id `experts` is the extra matrix (W2, S2, B2: one more [N, K] of the same width), the shared expert's
+  const bool extra = expert == uint(experts);
+  const size_t at = (extra ? 0 : size_t(expert) * N) + size_t(row0);
+  const device uint8_t* wr = (const device uint8_t*)(extra ? W2 : W) + at * (K * BITS / 8);
+  const device bfloat* sr = (extra ? S2 : S) + at * (K / GS);
+  const device bfloat* br = (extra ? B2 : B) + at * (K / GS);
+  float acc[1][RPS];
+  tf_rp_rows<RPS, 1, BITS, BF16SUM>(wr, sr, br, X + size_t(slot / repeat) * K, K, GS, lane, acc);
+  if (lane == 0)
+    for (int j = 0; j < RPS; j++) out[j] = tf_rp_out<RELU2>(acc[0][j]);
+}
 
 template [[host_name("tf_row_projection_q2_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 2, false>) tf_row_projection<4, 2, 2, false, 2, false>;
 template [[host_name("tf_row_projection_indexed_q2_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 2, false>) tf_row_projection_indexed<4, 2, false, 2, false>;
+template [[host_name("tf_row_projection_indexed_extra_q2_f32")]] [[kernel]] decltype(tf_row_projection_indexed_extra<4, 2, false, 2, false>) tf_row_projection_indexed_extra<4, 2, false, 2, false>;
 template [[host_name("tf_row_projection_relu2_q2_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 2, false>) tf_row_projection<4, 2, 2, true, 2, false>;
 template [[host_name("tf_row_projection_indexed_relu2_q2_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 2, false>) tf_row_projection_indexed<4, 2, true, 2, false>;
 template [[host_name("tf_row_projection_q2_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 2, true>) tf_row_projection<4, 2, 2, false, 2, true>;
 template [[host_name("tf_row_projection_indexed_q2_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 2, true>) tf_row_projection_indexed<4, 2, false, 2, true>;
+template [[host_name("tf_row_projection_indexed_extra_q2_bf16")]] [[kernel]] decltype(tf_row_projection_indexed_extra<4, 2, false, 2, true>) tf_row_projection_indexed_extra<4, 2, false, 2, true>;
 template [[host_name("tf_row_projection_relu2_q2_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 2, true>) tf_row_projection<4, 2, 2, true, 2, true>;
 template [[host_name("tf_row_projection_indexed_relu2_q2_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 2, true>) tf_row_projection_indexed<4, 2, true, 2, true>;
 template [[host_name("tf_row_projection_q4_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 4, false>) tf_row_projection<4, 2, 2, false, 4, false>;
 template [[host_name("tf_row_projection_indexed_q4_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 4, false>) tf_row_projection_indexed<4, 2, false, 4, false>;
+template [[host_name("tf_row_projection_indexed_extra_q4_f32")]] [[kernel]] decltype(tf_row_projection_indexed_extra<4, 2, false, 4, false>) tf_row_projection_indexed_extra<4, 2, false, 4, false>;
 template [[host_name("tf_row_projection_relu2_q4_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 4, false>) tf_row_projection<4, 2, 2, true, 4, false>;
 template [[host_name("tf_row_projection_indexed_relu2_q4_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 4, false>) tf_row_projection_indexed<4, 2, true, 4, false>;
 template [[host_name("tf_row_projection_q4_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 4, true>) tf_row_projection<4, 2, 2, false, 4, true>;
 template [[host_name("tf_row_projection_indexed_q4_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 4, true>) tf_row_projection_indexed<4, 2, false, 4, true>;
+template [[host_name("tf_row_projection_indexed_extra_q4_bf16")]] [[kernel]] decltype(tf_row_projection_indexed_extra<4, 2, false, 4, true>) tf_row_projection_indexed_extra<4, 2, false, 4, true>;
 template [[host_name("tf_row_projection_relu2_q4_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 4, true>) tf_row_projection<4, 2, 2, true, 4, true>;
 template [[host_name("tf_row_projection_indexed_relu2_q4_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 4, true>) tf_row_projection_indexed<4, 2, true, 4, true>;
 template [[host_name("tf_row_projection_q6_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 6, false>) tf_row_projection<4, 2, 2, false, 6, false>;
 template [[host_name("tf_row_projection_indexed_q6_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 6, false>) tf_row_projection_indexed<4, 2, false, 6, false>;
+template [[host_name("tf_row_projection_indexed_extra_q6_f32")]] [[kernel]] decltype(tf_row_projection_indexed_extra<4, 2, false, 6, false>) tf_row_projection_indexed_extra<4, 2, false, 6, false>;
 template [[host_name("tf_row_projection_relu2_q6_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 6, false>) tf_row_projection<4, 2, 2, true, 6, false>;
 template [[host_name("tf_row_projection_indexed_relu2_q6_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 6, false>) tf_row_projection_indexed<4, 2, true, 6, false>;
 template [[host_name("tf_row_projection_q6_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 6, true>) tf_row_projection<4, 2, 2, false, 6, true>;
 template [[host_name("tf_row_projection_indexed_q6_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 6, true>) tf_row_projection_indexed<4, 2, false, 6, true>;
+template [[host_name("tf_row_projection_indexed_extra_q6_bf16")]] [[kernel]] decltype(tf_row_projection_indexed_extra<4, 2, false, 6, true>) tf_row_projection_indexed_extra<4, 2, false, 6, true>;
 template [[host_name("tf_row_projection_relu2_q6_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 6, true>) tf_row_projection<4, 2, 2, true, 6, true>;
 template [[host_name("tf_row_projection_indexed_relu2_q6_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 6, true>) tf_row_projection_indexed<4, 2, true, 6, true>;
 template [[host_name("tf_row_projection_q8_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 8, false>) tf_row_projection<4, 2, 2, false, 8, false>;
 template [[host_name("tf_row_projection_indexed_q8_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 8, false>) tf_row_projection_indexed<4, 2, false, 8, false>;
+template [[host_name("tf_row_projection_indexed_extra_q8_f32")]] [[kernel]] decltype(tf_row_projection_indexed_extra<4, 2, false, 8, false>) tf_row_projection_indexed_extra<4, 2, false, 8, false>;
 template [[host_name("tf_row_projection_relu2_q8_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 8, false>) tf_row_projection<4, 2, 2, true, 8, false>;
 template [[host_name("tf_row_projection_indexed_relu2_q8_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 8, false>) tf_row_projection_indexed<4, 2, true, 8, false>;
 template [[host_name("tf_row_projection_q8_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 8, true>) tf_row_projection<4, 2, 2, false, 8, true>;
 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_indexed_extra_q8_bf16")]] [[kernel]] decltype(tf_row_projection_indexed_extra<4, 2, false, 8, true>) tf_row_projection_indexed_extra<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_indexed_extra_q3_f32")]] [[kernel]] decltype(tf_row_projection_indexed_extra<4, 2, false, 3, false>) tf_row_projection_indexed_extra<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_indexed_extra_q3_bf16")]] [[kernel]] decltype(tf_row_projection_indexed_extra<4, 2, false, 3, true>) tf_row_projection_indexed_extra<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_indexed_extra_q5_f32")]] [[kernel]] decltype(tf_row_projection_indexed_extra<4, 2, false, 5, false>) tf_row_projection_indexed_extra<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_indexed_extra_q5_bf16")]] [[kernel]] decltype(tf_row_projection_indexed_extra<4, 2, false, 5, true>) tf_row_projection_indexed_extra<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>;

zig/kernels/metal/qwen27/attention_mpp.metal

@@ -3,7 +3,11 @@
 #include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
 using namespace metal;
 using namespace mpp::tensor_ops;
-constant constexpr int G = 6, D = 256, TK = 64, MAXD = 128, MG = 8, MS = 12, NW = 4; // a row's key chunk is its node word 2
+#ifndef Q27_G
+#define Q27_G 6 // query heads a KV head (the 27B's); the model's host prepends its own
+#define Q27_KVH 4
+#endif
+constant constexpr int G = Q27_G, KVH = Q27_KVH, D = 256, TK = 64, MAXD = 128, MG = 8, MS = 12, NW = 4; // a row's key chunk is its node word 2
 struct Q27AttnReadPtrs { device const bfloat* keys; device const bfloat* values; };
 
 // SG tiles a threadgroup, a simdgroup pair each: both run scores and softmax, one owns each half of the head dim.
@@ -317,10 +321,10 @@ kernel void q27_attn_merge(device const float* POA [[buffer(0)]], device const f
   }
   const int h = hk * G + g;
   for (int i = 0; i < DP; i++) {
-    const int64_t index = ((int64_t)node * (G * 4) + h) * D + lane * DP + i;
+    const int64_t index = ((int64_t)node * (G * KVH) + h) * D + lane * DP + i;
     const bfloat v = static_cast<bfloat>(o[i] / l);
     if (gated) {
-      const float gate = float(QG[((int64_t)node * (G * 4) + h) * 2 * D + D + lane * DP + i]);
+      const float gate = float(QG[((int64_t)node * (G * KVH) + h) * 2 * D + D + lane * DP + i]);
       OUT[index] = static_cast<bfloat>(float(v) * q27_round(1.0f / (1.0f + exp(-gate)), 0));
     } else OUT[index] = v;
   }

zig/kernels/metal/qwen27/attention_prompt.metal

@@ -3,7 +3,7 @@
 #include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
 using namespace metal;
 using namespace mpp::tensor_ops;
-constant constexpr int PG = 6, PD = 256, PH = 128, PCK = 512, PTK = 32;
+constant constexpr int PG = Q27_G, PKVH = Q27_KVH, PD = 256, PH = 128, PCK = 512, PTK = 32; // after attention_mpp.metal's G
 struct Q27PromptArgs { int start, rows, chunks, rp, key_stride, value_stride; float scale2; };
 typedef vec<float, 8> q27f8;
 typedef vec<bfloat, 8> q27b8;
@@ -166,5 +166,5 @@ kernel void q27_attn_prompt_merge(device const float* PO [[buffer(0)]], device c
     m = nm;
   }
   const int node = r / PG, head = hk * PG + r % PG;
-  for (int i = 0; i < DP; i++) OUT[((int64_t)node * (PG * 4) + head) * PD + lane * DP + i] = static_cast<bfloat>(o[i] / l);
+  for (int i = 0; i < DP; i++) OUT[((int64_t)node * (PG * PKVH) + head) * PD + lane * DP + i] = static_cast<bfloat>(o[i] / l);
 }