core: add the pair seam and register its Metal kernels with the kernel contract

docs/kernel-contract.md

@@ -127,10 +127,14 @@ 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
-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.
+(4 bits, and its fp32 store: the same tiles compiled with `TF_OUT` set to float, so a rank can write its exact part
+into a seam) 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), and the TP seam's
+four exchange kernels (`pair_add_rms`, `pair_sum`, `pair_argmax` and `pair_pick`; the input is one rank's own part
+plus the peer's bytes, so an exchange is a function of those bytes alone, `pair_sum`'s fp32 parts meet in the
+frame's one bf16 rounding, and `pair_argmax` and `pair_pick` state their own NaN and tie rules). Gate 5 also holds
+the rest of tile_qmm's pipelines and GLM-5.3-Flash's hc_expand_10.
 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.metal

@@ -1,42 +1,75 @@
-// Two Macs' tensor-parallel exchanges inside a command buffer: post, wait for the peer's half, add in rank order.
+// Two Macs' tensor-parallel exchanges: the host relays parts through shared events; these kernels only do arithmetic.
 #include <metal_stdlib>
 using namespace metal;
 
-constant constexpr uint POSTED = 0, GAVE_UP = 16;
+// Sync page words (shared with pair.zig): FAILED = the host gave up on a seam; OBSERVED = this side's refusals.
+constant constexpr uint FAILED = 16, OBSERVED = 20;
 
-// Thread 0 posts exchange x (the host then sends this rank's part) and polls until the peer's part has landed.
-inline void pair_post_wait(device atomic_uint* sync, device atomic_uint* flag, uint x, uint t) {
-  if (t != 0) return;
-  atomic_store_explicit(&sync[POSTED], x, memory_order_relaxed);
-  uint polls = 0;
-  while (int(atomic_load_explicit(flag, memory_order_relaxed) - x) < 0) {
-    if (++polls > 400000000u) {
-      atomic_fetch_add_explicit(&sync[GAVE_UP], 1u, memory_order_relaxed);
-      return;
-    }
-  }
+// 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;
 }
 
 inline float pair_lanes(float v) {
   for (ushort d = 16; d > 0; d >>= 1) v += simd_shuffle_xor(v, d);
   return v;
 }
 
+// Kernels' argument block, byte-identical to pair.zig's PairRows: four u32 then four f32; no payload element.
 struct PairRows {
   uint width, stride, rank, x;
   float eps, pad0, pad1, pad2;
 };
 
-// Row r: h += (rank 0's part + rank 1's part), then y = w * (h / rms(h)); the peer's words read as atomics.
-kernel void pair_add_rms(device atomic_uint* sync [[buffer(0)]], device atomic_uint* flag [[buffer(1)]],
-                         constant PairRows& a [[buffer(2)]], const device float* mine [[buffer(3)]],
-                         device atomic_uint* peer [[buffer(4)]], device float* h [[buffer(5)]],
-                         const device bfloat* w [[buffer(6)]], device float* y [[buffer(7)]],
-                         uint r [[threadgroup_position_in_grid]], uint t [[thread_index_in_threadgroup]],
-                         uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]]) {
+// SEAM_ELEM_BF16: 0 = fp32 payload (shipping), 1 = bf16; producer must match. Rationale: docs/speed-up-mode.md
+#define SEAM_ELEM_BF16 0
+
+// bf16 bits as the exact fp32 value they name, and fp32 as bf16 bits rounded to nearest even (one rounding).
+inline float pair_bf16(uint bits) {
+  return as_type<float>(bits << 16);
+}
+inline ushort pair_bf16_bits(float f) {
+  const uint u = as_type<uint>(f);
+  return ushort((u + 0x7FFFu + ((u >> 16) & 1u)) >> 16);
+}
+
+// One payload element of this rank's part (`mine` or `out`), indexed by element, whichever type the payload holds.
+inline float seam_get(const device ushort* p, ulong i) {
+#if SEAM_ELEM_BF16
+  return pair_bf16(uint(p[i]));
+#else
+  return reinterpret_cast<const device float*>(p)[i];
+#endif
+}
+// Frame, not payload: pair_sum's out is bf16 under either payload, so an f32 payload buys an exact half-sum only.
+inline void frame_put(device ushort* p, ulong i, float v) {
+  p[i] = pair_bf16_bits(v);
+}
+
+// The peer's element i: atomic word load; bf16 packs two a word (an even width keeps rows whole), fp32 is one a word.
+inline float seam_peer(const device atomic_uint* words, ulong i) {
+#if SEAM_ELEM_BF16
+  const uint w = atomic_load_explicit(&words[i >> 1], memory_order_relaxed);
+  return pair_bf16(w >> ((i & 1) * 16));
+#else
+  return as_type<float>(atomic_load_explicit(&words[i], memory_order_relaxed));
+#endif
+}
+
+// Row r: h += rank 0's part + rank 1's part, then y = w * (h / rms(h)); payload is fp32 (streamBytes), not SEAM_ELEM.
+kernel void pair_add_rms(device atomic_uint* sync [[buffer(0)]], constant PairRows& a [[buffer(1)]],
+                         const device float* mine [[buffer(2)]], device atomic_uint* peer [[buffer(3)]],
+                         device float* h [[buffer(4)]], const device bfloat* w [[buffer(5)]],
+                         device float* y [[buffer(6)]], uint r [[threadgroup_position_in_grid]],
+                         uint t [[thread_index_in_threadgroup]], uint sg [[simdgroup_index_in_threadgroup]],
+                         uint lane [[thread_index_in_simdgroup]]) {
   threadgroup float part[8];
-  pair_post_wait(sync, flag, a.x, t);
-  threadgroup_barrier(mem_flags::mem_device);
+  if (pair_link_broken(sync, t)) {
+    for (uint i = t; i < a.width; i += 256) y[ulong(r) * a.stride + i] = NAN;
+    return;
+  }
   device float* x = h + ulong(r) * a.stride;
   float s = 0.0f;
   for (uint i = t; i < a.width; i += 256) {
@@ -56,16 +89,37 @@ kernel void pair_add_rms(device atomic_uint* sync [[buffer(0)]], device atomic_u
   for (uint i = t; i < a.width; i += 256) y[ulong(r) * a.stride + i] = float(w[i]) * (x[i] * rs);
 }
 
-// Row r's argmax over this rank's logits (n of them, from vocabulary id lo): (value, id), the lower id on a tie.
+// Row r: out = frame(float(mine) + float(peer)), one rounding; payload is a.width, frame a.stride, equal here.
+kernel void pair_sum(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;
+  if (pair_link_broken(sync, t)) {
+    for (uint i = t; i < a.width; i += 256) frame_put(out, orow + i, NAN);
+    return;
+  }
+  for (uint i = t; i < a.width; i += 256) {
+    const float m = seam_get(mine, mrow + i);
+    const float p = seam_peer(peer, mrow + i);
+    frame_put(out, orow + i, a.rank == 0 ? m + p : p + m);
+  }
+}
+
+// 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 device float* x = logits + ulong(r) * dims.x;
+  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 = 0xFFFFFFFFu;
   for (uint i = t; i < dims.z; i += 1024) {
-    const float v = x[i];
+    const float v = wide ? pair_bf16(uint(xb[i])) : xf[i];
     if (v > best) best = v, at = dims.y + i;
   }
   for (ushort o = 16; o > 0; o >>= 1) {
@@ -88,12 +142,14 @@ kernel void pair_argmax(const device float* logits [[buffer(0)]], constant uint4
   if (lane == 0) out[2 * r] = as_type<uint>(best), out[2 * r + 1] = at;
 }
 
-// Post, wait for the peer's halves, then each row's pick: the larger value, the lower id on a tie (one Mac's argmax).
-kernel void pair_pick(device atomic_uint* sync [[buffer(0)]], device atomic_uint* flag [[buffer(1)]],
-                      constant uint2& rx [[buffer(2)]], const device uint* mine [[buffer(3)]],
-                      device atomic_uint* peer [[buffer(4)]], device uint* picks [[buffer(5)]],
-                      uint t [[thread_index_in_threadgroup]]) {
-  pair_post_wait(sync, flag, rx.y, t);
+// Each row's pick over both Macs' argmaxes: the larger value, the lower id on a tie (one Mac's argmax).
+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]);

zig/kernels/metal/core/simd_qmm.metal

@@ -2,6 +2,11 @@
 #include <metal_stdlib>
 using namespace metal;
 
+// OUT type: bf16 for the frame, fp32 for the seam producer (exact partial); TF_OUT=float selects the fp32 build.
+#ifndef TF_OUT
+#define TF_OUT bfloat
+#endif
+
 // K inputs, N outputs, R rows; one is 1.0 read at run time so the compiler can't fold or reorder the fma sums.
 struct TfSqArgs { int K, N, R; float one; };
 
@@ -28,7 +33,7 @@ template <int S, int NT, int RT>
                           const device bfloat* SC [[buffer(2)]],
                           const device bfloat* BI [[buffer(3)]],
                           constant TfSqArgs& a [[buffer(4)]],
-                          device bfloat* OUT [[buffer(5)]],
+                          device TF_OUT* OUT [[buffer(5)]],
                           uint sgi [[simdgroup_index_in_threadgroup]],
                           uint sgs [[simdgroups_per_threadgroup]],
                           uint lane [[thread_index_in_simdgroup]],
@@ -100,7 +105,7 @@ template <int S, int NT, int RT>
         for (int t = 0; t < NT; t++)
           for (int e = 0; e < 2; e++) {
             const int row = rb + 8 * rt + fn + e, n = nb + 8 * t + fm;
-            if (row < R && n < N) OUT[size_t(row) * N + n] = bfloat(acc[rt][t][e]);
+            if (row < R && n < N) OUT[size_t(row) * N + n] = TF_OUT(acc[rt][t][e]);
           }
       return;
     }
@@ -117,7 +122,7 @@ template <int S, int NT, int RT>
     const int rt = idx / (NT * 64), t = (idx / 64) % NT, l = (idx % 64) / 2, e = idx % 2;
     const int lq = l / 4;
     const int row = rb + 8 * rt + (lq & 2) * 2 + (l % 2) * 2 + e, n = nb + 8 * t + (lq & 4) + ((l / 2) % 4);
-    if (row < R && n < N) OUT[size_t(row) * N + n] = bfloat(v[0]);
+    if (row < R && n < N) OUT[size_t(row) * N + n] = TF_OUT(v[0]);
   }
 }
 
@@ -128,7 +133,7 @@ template <int S, int NT, int RT, int P>
                           const device bfloat* SC [[buffer(2)]],
                           const device bfloat* BI [[buffer(3)]],
                           constant TfSqArgs& a [[buffer(4)]],
-                          device bfloat* OUT [[buffer(5)]],
+                          device TF_OUT* OUT [[buffer(5)]],
                           uint sgi [[simdgroup_index_in_threadgroup]],
                           uint lane [[thread_index_in_simdgroup]],
                           uint3 tg [[threadgroup_position_in_grid]]) {
@@ -209,7 +214,7 @@ template <int S, int NT, int RT, int P>
       for (int t = 0; t < NT; t++)
         for (int e = 0; e < 2; e++) {
           const int row = rb + 8 * rt + fn + e, n = nb + 8 * t + fm;
-          if (row < R && n < N) OUT[size_t(row) * N + n] = bfloat(acc[0][rt][t][e]);
+          if (row < R && n < N) OUT[size_t(row) * N + n] = TF_OUT(acc[0][rt][t][e]);
         }
     return;
   }
@@ -225,7 +230,7 @@ template <int S, int NT, int RT, int P>
     const int rt = idx / (NT * 64), t = (idx / 64) % NT, l = (idx % 64) / 2, e = idx % 2;
     const int lq = l / 4;
     const int row = rb + 8 * rt + (lq & 2) * 2 + (l % 2) * 2 + e, n = nb + 8 * t + (lq & 4) + ((l / 2) % 4);
-    if (row < R && n < N) OUT[size_t(row) * N + n] = bfloat(v[0]);
+    if (row < R && n < N) OUT[size_t(row) * N + n] = TF_OUT(v[0]);
   }
 }
 
@@ -236,7 +241,7 @@ template <int S, int NR, int XB, int RS>
                              const device bfloat* SC [[buffer(2)]],
                              const device bfloat* BI [[buffer(3)]],
                              constant TfSqArgs& a [[buffer(4)]],
-                             device bfloat* OUT [[buffer(5)]],
+                             device TF_OUT* OUT [[buffer(5)]],
                              uint sgi [[simdgroup_index_in_threadgroup]],
                              uint sgs [[simdgroups_per_threadgroup]],
                              uint lane [[thread_index_in_simdgroup]],
@@ -320,7 +325,7 @@ template <int S, int NR, int XB, int RS>
       float v = acc[u][r];
       for (int m = 1; m < S; m <<= 1) v = fma(simd_shuffle_xor(v, ushort(m)), one, v);
       const int n = n0 + SLOTS * u;
-      if (n < N && c == 0) OUT[size_t(r) * N + n] = bfloat(v);
+      if (n < N && c == 0) OUT[size_t(r) * N + n] = TF_OUT(v);
     }
 }
 

zig/src/core/contract.zig

@@ -19,6 +19,7 @@ test {
     _ = @import("contract/vit_records.zig");
     _ = @import("contract/simd_pairs_records.zig");
     _ = @import("contract/glm_rows_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.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 add_rms pair_add_rms", .threads = 256, .tg_memory = 32, .m1 = 1024, .m2 = 1024 },
+    .{ .name = "metal pair sum pair_sum", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 },
+    .{ .name = "metal pair argmax pair_argmax", .threads = 1024, .tg_memory = 256, .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);
+    }
+}

zig/src/core/contract/registry.zig

@@ -393,6 +393,7 @@ pub const metal = [_]M{
     .{ .name = "affine_mm q4 g64 tf_affine_row_sums", .threads = 64, .tg_memory = 0, .m1 = 1024, .m2 = 1024 },
     .{ .name = "affine_mm q4 g64 tf_affine_mm", .threads = 128, .tg_memory = 9728, .offline = false },
     .{ .name = "affine_mm q8 g32 tf_affine_mm", .threads = 128, .tg_memory = 10240, .offline = false },
+    // The fp32 store is this same source with TF_OUT defined to float, so these launches cover both builds.
     .{ .name = "simd_qmm tf_sq_mma_16_4_1", .threads = 512, .tg_memory = 16384, .m1 = 768, .m2 = 768, .adapts = true },
     .{ .name = "simd_qmm tf_sq_mma_16_2_2", .threads = 512, .tg_memory = 16384, .m1 = 704, .m2 = 704, .adapts = true },
     .{ .name = "simd_qmm tf_sq_mma_8_4_1", .threads = 256, .tg_memory = 8192, .m1 = 768, .m2 = 768, .adapts = true },
@@ -566,7 +567,7 @@ test {
     _ = paging;
 }
 /// The Metal collector keeps measured slots and the restored vision declarations.
-pub const all_metal = metal ++ @import("delta_slots_records.zig").records ++ @import("flashnext_pending.zig").records ++ @import("simd_pairs_records.zig").records ++ @import("glm_rows_records.zig").records ++ @import("weight_only_records.zig").metal ++ @import("hc_activation_records.zig").metal ++ @import("qwen4_records.zig").metal;
+pub const all_metal = metal ++ @import("delta_slots_records.zig").records ++ @import("flashnext_pending.zig").records ++ @import("simd_pairs_records.zig").records ++ @import("metal_pair_records.zig").records ++ @import("glm_rows_records.zig").records ++ @import("weight_only_records.zig").metal ++ @import("hc_activation_records.zig").metal ++ @import("qwen4_records.zig").metal;
 const learning = @import("registry_learning.zig");
 pub const gradient_pending = learning.gradient_pending;
 pub const qwen_learning_pending = learning.qwen_learning_pending;