qwen27: rank-slice the 27B family and gather the whole embedding table

docs/kernel-contract.md

@@ -20,15 +20,16 @@ The test of a kernel fills a `contract.spec.Spec` and hands the harness a few ho
   and the magnitude of the terms each output sums, and can replay the sums in another accumulator.
 - Its served envelope: the shapes the engine calls it at, and for each dimension its tile, the step a valid size
   moves by and the least valid size.
-- Its accumulator (`f32` for every kernel here) and the depth of its longest chain of roundings at a shape, plus any
+- Its accumulator (`f32` where a kernel sums) and the depth of its longest chain of roundings at a shape, plus any
   fp32 operations after the sums (`ops`) and intermediate roundings to the output type (`mids`).
 - Whether it may write in place. If it replaces another kernel, the kernel whose bits it must equal.
 - Its NaN rule: `propagate` (as float64 arithmetic on its formula does), or a stated rule the test checks itself
   (top-k, argmax, sampling and routing kernels state theirs).
 - For gate 5, its launch in `zig/src/core/contract/registry.zig`: threads a group and threadgroup bytes (the source's,
   and any its encoder sets) on Metal, or symbol, image, threads a block, dynamic shared bytes and cluster width on CUDA,
   with the limits its compiles recorded. HIP launches are in `zig/src/core/contract/hip_records.zig`: symbol, code
-  object, threads a workgroup and dynamic LDS bytes.
+  object, threads a workgroup and dynamic LDS bytes. The seam's records are in `metal_pair_records.zig`, which
+  `registry.zig` concatenates.
 
 ## The gates
 
@@ -99,7 +100,8 @@ The program exits 1 on any failed gate, any control that stayed green, or any li
    the inputs for a shape and a gate's fill (usually through an input family), `run` launches once and waits, `count`
    and `reference` give the outputs and their float64 values, `simulate` replays the sums for gate 4.
 2. Add it to the runner's list, with its served shapes and tiles.
-3. Add its launches to `registry.zig`, and fill the records from `--offline` (Metal) or `kernel-contract-limits`
+3. Add its launches to `registry.zig`, or to a records file it concatenates as the seam's are, and fill them from
+   `--offline` (Metal) or `kernel-contract-limits`
    (CUDA: nvcc 13.0's in `cuda`, and each later release's set where that release compiles it differently); a HIP
    kernel's go to `hip_records.zig`, filled from `hip-kernel-contract-limits`.
 4. Post the `SUMMARY` lines from each backend you ran in the pull request, and say which chips you could not run.
@@ -127,8 +129,8 @@ float64 output class, without an inherited exception.
 
 ## Coverage
 
-Metal: softmax, RMS norm, bf16 top-k, the row projection (3, 4, 6 and 8 bits), affine_mm (4 and 8 bits), simd_qmm
-(4 bits) and simd_qmm_bits (3, 6 and 8 bits), and the lane projection (6-bit groups of 32 on every chip; 4- and 8-bit
+Metal: softmax, RMS norm, bf16 top-k, the row projection (3, 4, 6 and 8 bits), affine_mm (4 and 8 bits), simd_qmm (4
+bits) and simd_qmm_bits (3, 6 and 8 bits), and the lane projection (6-bit groups of 32 on every chip; 4- and 8-bit
 formats on the M5's tensor units, where its reg kernel must equal the cooperative one bit for bit), and tile_qmm's
 prompt tiles and their decode twin (4, 6 and 8 bits; the twin bit-equal to the tiles). Gate 5 also holds the rest of
 tile_qmm's pipelines and GLM-5.3-Flash's hc_expand_10.
@@ -141,6 +143,13 @@ Metal also checks Nemotron's routed experts on the engine's route tables, one ro
 against float64 at tf_load16's bf16 x sums (up's relu2 as an interval over its two bf16 roundings), and the rows2, rows4
 and six geometry pipelines of each projection bit-equal to them (two controls: member rows read from one row, and
 w3's x sums without their bf16 steps).
+The row-split seam's four kernels (`pair_rows.metal`, launched by `pair_rows.zig`) move bf16 frames, never sums:
+`pair_stage` copies this rank's half of a frame row into the slot the relay sends, and `pair_gather` writes both halves
+of every row, its own natively and the peer's from the relayed frame. `pair_argmax` flags a row whose logits are not
+finite (its id word carries bit 31, its value word a NaN) and takes the lower id on a tie; `pair_pick` turns a flagged
+rank's word into the id sentinel, so both Macs refuse a non-finite round where one Mac does. Six controls edit the
+source: a stage reading the peer's half, a gather writing its own half twice, and an argmax or pick tie to the higher
+id must go red at `shapes`, and an argmax or pick that does not flag at `nonfinite`.
 CUDA: the RMS norm rows (16-byte form and the fp32-out tree), the bf16 residual's norm tails (MLX's rounded rows and
 the once-rounded ones, each checked on its residual and its output, with four controls each dropping a rounding point), Nemotron's router logits (and their 128-row tiles,
 bit-equal to them) and top-k, the latent top-k's decode form (bit-equal to `tf_super_topk`), the latent combine and

zig/kernels/metal/core/pair_rows.metal

@@ -0,0 +1,129 @@
+// Two Macs' tensor-parallel exchanges: the host relays parts through shared events, and these kernels move bf16 bits.
+#include <metal_stdlib>
+using namespace metal;
+// A rank compares one row's logits across ranks; no kernel sums across them.
+
+// Sync page words (shared with pair_rows.zig): FAILED = the host gave up on a seam; OBSERVED = this side's refusals.
+constant constexpr uint FAILED = 16, OBSERVED = 20;
+// An argmax id word's bit 31: the row had a logit that is not finite, so pair_pick sends the row to the id sentinel.
+constant constexpr uint PAIR_FLAG = 0x80000000u;
+constant constexpr uint PAIR_SENTINEL = 0xFFFFFFFFu;
+
+// Kernels' argument block, byte-identical to pair_rows.zig's PairRows: width, the frame row's stride, this rank, one.
+struct PairRows {
+  uint width, stride, rank, x;
+};
+
+// bf16 bits as the exact fp32 value they name (argmax reads a bf16 logits stream this way).
+inline float pair_bf16(uint bits) {
+  return as_type<float>(bits << 16);
+}
+
+// True once the host gave up: FAILED is set before the wait releases, so a later kernel poisons its output rows.
+inline bool pair_link_broken(device atomic_uint* sync, uint t) {
+  if (atomic_load_explicit(&sync[FAILED], memory_order_relaxed) == 0u) return false;
+  if (t == 0) atomic_fetch_add_explicit(&sync[OBSERVED], 1u, memory_order_relaxed);
+  return true;
+}
+
+// Row r: this rank's own half of the frame row into the exchange slot; a bf16 copy, so it carries the frame's bits.
+kernel void pair_stage(device atomic_uint* sync [[buffer(0)]], constant PairRows& a [[buffer(1)]],
+                       const device ushort* src [[buffer(2)]], device ushort* dst [[buffer(3)]],
+                       uint r [[threadgroup_position_in_grid]], uint t [[thread_index_in_threadgroup]]) {
+  const ulong in = ulong(r) * a.stride;
+  const ulong out = ulong(r) * a.width;
+  if (pair_link_broken(sync, t)) {
+    for (uint i = t; i < a.width; i += 256) dst[out + i] = 0x7FC0u; // bf16 NaN: the link stays down
+    return;
+  }
+  for (uint i = t; i < a.width; i += 256) dst[out + i] = src[in + i];
+}
+
+// Row r: this rank's half at its own offset in the frame row and the peer's at the other; a bf16 copy, no arithmetic.
+kernel void pair_gather(device atomic_uint* sync [[buffer(0)]], constant PairRows& a [[buffer(1)]],
+                        const device ushort* mine [[buffer(2)]], device atomic_uint* peer [[buffer(3)]],
+                        device ushort* out [[buffer(4)]], uint r [[threadgroup_position_in_grid]],
+                        uint t [[thread_index_in_threadgroup]]) {
+  const ulong mrow = ulong(r) * a.width;
+  const ulong orow = ulong(r) * a.stride;
+  const ulong own = a.rank == 0 ? orow : orow + a.width;
+  const ulong got = a.rank == 0 ? orow + a.width : orow;
+  if (pair_link_broken(sync, t)) {
+    for (uint i = t; i < a.width; i += 256) {
+      out[own + i] = 0x7FC0u;
+      out[got + i] = 0x7FC0u;
+    }
+    return;
+  }
+  for (uint i = t; i < a.width; i += 256) {
+    // The relay wrote the peer's half: an untracked frame read, so one atomic word (two bf16) at a time.
+    const uint word = atomic_load_explicit(&peer[(mrow + i) >> 1], memory_order_relaxed);
+    out[own + i] = mine[mrow + i];
+    out[got + i] = ushort(word >> ((i & 1) * 16));
+  }
+}
+
+// Row r's argmax over this rank's logits (ids lo..lo+n): (fp32 value, id), lower id on tie; dims.w 0 = fp32, 1 = bf16.
+kernel void pair_argmax(const device float* logits [[buffer(0)]], constant uint4& dims [[buffer(1)]],
+                        device uint* out [[buffer(2)]], uint r [[threadgroup_position_in_grid]],
+                        uint t [[thread_index_in_threadgroup]], uint sg [[simdgroup_index_in_threadgroup]],
+                        uint lane [[thread_index_in_simdgroup]]) {
+  const ulong row = ulong(r) * dims.x;
+  const device float* xf = logits + row;
+  const device ushort* xb = reinterpret_cast<const device ushort*>(logits) + row;
+  const bool wide = dims.w != 0u;
+  float best = -INFINITY;
+  uint at = PAIR_SENTINEL;
+  bool bad = false;
+  for (uint i = t; i < dims.z; i += 1024) {
+    const float v = wide ? pair_bf16(uint(xb[i])) : xf[i];
+    if (!isfinite(v)) { bad = true; continue; } // the flagged id is the argmax over finite logits
+    if (v > best) best = v, at = dims.y + i;
+  }
+  for (ushort o = 16; o > 0; o >>= 1) {
+    const float ob = simd_shuffle_xor(best, o);
+    const uint oa = simd_shuffle_xor(at, o);
+    if (ob > best || (ob == best && oa < at)) best = ob, at = oa;
+  }
+  // The flag is the row's, not the thread's: one lane saw a non-finite logit, so every lane must know.
+  uint dirty = bad ? 1u : 0u;
+  for (ushort o = 16; o > 0; o >>= 1) dirty |= simd_shuffle_xor(dirty, o);
+  threadgroup float vb[32];
+  threadgroup uint va[32];
+  threadgroup uint vd[32];
+  if (lane == 0) vb[sg] = best, va[sg] = at, vd[sg] = dirty;
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (sg != 0) return;
+  best = vb[lane];
+  at = va[lane];
+  dirty = vd[lane];
+  for (ushort o = 16; o > 0; o >>= 1) {
+    const float ob = simd_shuffle_xor(best, o);
+    const uint oa = simd_shuffle_xor(at, o);
+    if (ob > best || (ob == best && oa < at)) best = ob, at = oa;
+    dirty |= simd_shuffle_xor(dirty, o);
+  }
+  if (lane == 0) {
+    out[2 * r] = dirty != 0u ? as_type<uint>(NAN) : as_type<uint>(best);
+    out[2 * r + 1] = dirty != 0u ? (at | PAIR_FLAG) : at;
+  }
+}
+
+// Each row's pick over both Macs' argmaxes: a flagged row draws no token; else the larger value, lower id on a tie.
+kernel void pair_pick(device atomic_uint* sync [[buffer(0)]], constant uint2& rx [[buffer(1)]],
+                      const device uint* mine [[buffer(2)]], device atomic_uint* peer [[buffer(3)]],
+                      device uint* picks [[buffer(4)]], uint t [[thread_index_in_threadgroup]]) {
+  if (pair_link_broken(sync, t)) {
+    if (t == 0) for (uint r = 0; r < rx.x; ++r) picks[r] = 0xFFFFFFFFu;  // no such token: visibly wrong, never a token
+    return;
+  }
+  if (t != 0) return;
+  for (uint r = 0; r < rx.x; ++r) {
+    const float vm = as_type<float>(mine[2 * r]);
+    const uint am = mine[2 * r + 1];
+    const float vp = as_type<float>(atomic_load_explicit(&peer[2 * r], memory_order_relaxed));
+    const uint ap = atomic_load_explicit(&peer[2 * r + 1], memory_order_relaxed);
+    // A non-finite row on either Mac draws no token: the same sentinel a broken link reads, never an id.
+    picks[r] = ((am | ap) & PAIR_FLAG) != 0u ? PAIR_SENTINEL : (vp > vm || (vp == vm && ap < am)) ? ap : am;
+  }
+}

zig/kernels/metal/qwen27/attention_mpp.metal

@@ -263,8 +263,71 @@ kernel void q27_attn_tail(device const bfloat* QB [[buffer(0)]], device const Q2
   }
 }
 
-// A row's chunk partials combined in key order, a block of loads ahead; `gated` applies the query gate (q27_gate).
+// A row's chunk partials in key order, loads ahead: `gated` the query gate, `heads` the rank's (12 = one Mac's 24 / 2).
 kernel void q27_attn_merge(device const float* POA [[buffer(0)]], device const float* PMA [[buffer(1)]], device const float* PLA [[buffer(2)]],
+    device const float* POB [[buffer(3)]], device const float* PMB [[buffer(4)]], device const float* PLB [[buffer(5)]],
+    device const int* meta [[buffer(6)]], device const int* nodes [[buffer(7)]], device bfloat* OUT [[buffer(8)]],
+    device const bfloat* QG [[buffer(9)]], constant uint& gated [[buffer(10)]], constant uint& heads [[buffer(11)]],
+    uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]], uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) {
+  const uint lane = thread_index_in_simdgroup;
+  const uint hk = threadgroup_position_in_grid.x;
+  const uint r = threadgroup_position_in_grid.y;
+  const int NCB = meta[3], W = meta[4], RPA = meta[5], CA = meta[6];
+  const int st = nodes[NW * (int(r) / G) + 1];
+  const int NCBS = meta[MG + st * MS + 5];
+  const int ra = meta[MG + st * MS + 3] * 16 + (int(r) / G - meta[MG + st * MS + 6]) * G + int(r) % G;
+  const int CT = nodes[NW * (int(r) / G) + 3]; // the row's own prefix chunks
+  constexpr int DP = D / 32, MB = 8;
+  const int node = r / G, g = r % G;
+  float m = -INFINITY, l = 0.0f, o[DP];
+  for (int i = 0; i < DP; i++) o[i] = 0.0f;
+  for (int c0 = 0; c0 < CT; c0 += MB) {
+    float mcs[MB], lcs[MB];
+    float4 lo[MB], hi[MB];
+    for (int j = 0; j < MB; j++) {
+      const int64_t row = ((int64_t)hk * CA + min(c0 + j, CT - 1)) * RPA + ra;
+      mcs[j] = PMA[row];
+      lcs[j] = PLA[row];
+      lo[j] = *(device const float4*)(POA + row * D + lane * DP);
+      hi[j] = *(device const float4*)(POA + row * D + lane * DP + 4);
+    }
+    for (int j = 0; j < MB; j++) {
+      if (c0 + j >= CT) break;
+      const float mc = mcs[j];
+      if (mc == -INFINITY) continue;
+      const float lc = lcs[j];
+      const float nm = max(m, mc);
+      const float f1 = fast::exp(m - nm), f2 = fast::exp(mc - nm);
+      l = l * f1 + lc * f2;
+      for (int i = 0; i < 4; i++) o[i] = o[i] * f1 + lo[j][i] * f2;
+      for (int i = 0; i < 4; i++) o[4 + i] = o[4 + i] * f1 + hi[j][i] * f2;
+      m = nm;
+    }
+  }
+  for (int c = 0; c < NCBS; c++) {
+    const int64_t row = (((int64_t)hk * NCB + c) * W + node) * 16 + g;
+    const float mc = PMB[row];
+    if (mc == -INFINITY) continue;
+    const float lc = PLB[row];
+    const float nm = max(m, mc);
+    const float f1 = fast::exp(m - nm), f2 = fast::exp(mc - nm);
+    l = l * f1 + lc * f2;
+    for (int i = 0; i < DP; i++) o[i] = o[i] * f1 + POB[row * D + lane * DP + i] * f2;
+    m = nm;
+  }
+  const int h = hk * G + g;
+  for (int i = 0; i < DP; i++) {
+    const int64_t index = ((int64_t)node * (int64_t)heads + h) * D + lane * DP + i;
+    const bfloat v = static_cast<bfloat>(o[i] / l);
+    if (gated) {
+      const float gate = float(QG[((int64_t)node * (int64_t)heads + 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;
+  }
+}
+
+// A single Mac's 24 heads are known at compile time: main's kernel, so one Mac runs unpaired code.
+kernel void q27_attn_merge_single(device const float* POA [[buffer(0)]], device const float* PMA [[buffer(1)]], device const float* PLA [[buffer(2)]],
     device const float* POB [[buffer(3)]], device const float* PMB [[buffer(4)]], device const float* PLB [[buffer(5)]],
     device const int* meta [[buffer(6)]], device const int* nodes [[buffer(7)]], device bfloat* OUT [[buffer(8)]],
     device const bfloat* QG [[buffer(9)]], constant uint& gated [[buffer(10)]],

zig/kernels/metal/qwen27/attention_prompt.metal

@@ -4,7 +4,8 @@
 using namespace metal;
 using namespace mpp::tensor_ops;
 constant constexpr int PG = 6, PD = 256, PH = 128, PCK = 512, PTK = 32;
-struct Q27PromptArgs { int start, rows, chunks, rp, key_stride, value_stride; float scale2; };
+// kv_heads is the rank's, not the model's: a split attention writes local node rows.
+struct Q27PromptArgs { int start, rows, chunks, rp, key_stride, value_stride; float scale2; int kv_heads; };
 typedef vec<float, 8> q27f8;
 typedef vec<bfloat, 8> q27b8;
 typedef vec<half, 8> q27h8;
@@ -166,5 +167,29 @@ 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;
+  const int64_t at = ((int64_t)node * (PG * args.kv_heads) + head) * PD + lane * DP;
+  for (int i = 0; i < DP; i++) OUT[at + i] = static_cast<bfloat>(o[i] / l);
+}
+
+// A single Mac's 24 query heads are known at compile time: main's kernel, so one Mac runs unpaired code.
+kernel void q27_attn_prompt_merge_single(device const float* PO [[buffer(0)]], device const float* PM [[buffer(1)]], device const float* PL [[buffer(2)]],
+    constant Q27PromptArgs& args [[buffer(3)]], device bfloat* OUT [[buffer(4)]],
+    uint3 tg [[threadgroup_position_in_grid]], uint lane [[thread_index_in_simdgroup]]) {
+  const int hk = int(tg.x), r = int(tg.y);
+  constexpr int DP = PD / 32;
+  float m = -INFINITY, l = 0.0f, o[DP];
+  for (int i = 0; i < DP; i++) o[i] = 0.0f;
+  for (int c = 0; c < args.chunks; c++) {
+    const int64_t row = ((int64_t)hk * args.chunks + c) * args.rp + r;
+    const float mc = PM[row];
+    if (mc == -INFINITY) continue;
+    const float lc = PL[row];
+    const float nm = max(m, mc);
+    const float f1 = fast::exp2(m - nm), f2 = fast::exp2(mc - nm);
+    l = l * f1 + lc * f2;
+    for (int i = 0; i < DP; i++) o[i] = o[i] * f1 + PO[row * PD + lane * DP + i] * f2;
+    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);
 }

zig/kernels/metal/sources.zig

@@ -34,6 +34,8 @@ pub const core_exlr8_dense = @embedFile("core/exlr8_dense.metal");
 pub const core_exlr8_prompt = @embedFile("core/exlr8_prompt.metal");
 pub const core_exlr8_dense_prompt = @embedFile("core/exlr8_dense_prompt.metal");
 pub const core_pair = @embedFile("core/pair.metal");
+/// The row-split seam's kernels (see `pair_rows.zig`): `pair.metal` above stays as it is for glm_dsa.
+pub const core_pair_rows = @embedFile("core/pair_rows.metal");
 
 /// GLM-5.3 (glm_moe_dsa): MLA with rope over a bf16 latent cache, the DSA indexer and its top-k, fp32 glue.
 pub const glm_dsa_mla = @embedFile("glm_dsa/mla.metal");

zig/src/core/contract.zig

@@ -21,6 +21,7 @@ test {
     _ = @import("contract/glm_rows_records.zig");
     _ = @import("contract/nemotron_experts_records.zig");
     _ = @import("contract/moss_image_records.zig");
+    _ = @import("contract/metal_pair_records.zig");
     _ = spec;
     _ = gates;
     _ = rows;

zig/src/core/contract/metal_pair_records.zig

@@ -0,0 +1,25 @@
+//! The pair seam's four launches (core/pair_rows.metal) with the M1/M2 limits the offline compile measures.
+const std = @import("std");
+const limits = @import("limits.zig");
+const registry = @import("registry.zig");
+
+pub const records = [_]limits.MetalLaunch{
+    .{ .name = "metal pair stage pair_stage", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 },
+    .{ .name = "metal pair gather pair_gather", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 },
+    .{ .name = "metal pair argmax pair_argmax", .threads = 1024, .tg_memory = 384, .m1 = 1024, .m2 = 1024 },
+    .{ .name = "metal pair pick pair_pick", .threads = 32, .tg_memory = 0, .m1 = 1024, .m2 = 1024 },
+};
+
+test "the pair seam's four launches are in the live and offline collector and fit every family" {
+    for (records) |record| {
+        var found: usize = 0;
+        for (registry.all_metal) |active| if (std.mem.eql(u8, active.name, record.name)) {
+            try std.testing.expectEqualDeep(record, active);
+            found += 1;
+        };
+        try std.testing.expectEqual(@as(usize, 1), found);
+        try std.testing.expectEqual(@as(usize, 0), try registry.checkMetal(std.testing.allocator, &.{record}, true));
+        // A zero here would be a launch no M1/M2 compiler was asked about.
+        try std.testing.expect(record.m1 > 0 and record.m2 > 0);
+    }
+}