diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index af412e6887..59e0dd9123 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -20,7 +20,7 @@ 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 @@ -28,7 +28,8 @@ The test of a kernel fills a `contract.spec.Spec` and hands the harness a few ho - 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. Five 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 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 diff --git a/zig/kernels/metal/core/pair_rows.metal b/zig/kernels/metal/core/pair_rows.metal new file mode 100644 index 0000000000..a505ba2cd8 --- /dev/null +++ b/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 +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(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(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(NAN) : as_type(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(mine[2 * r]); + const uint am = mine[2 * r + 1]; + const float vp = as_type(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; + } +} diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index 2506fb4905..22cb484ef0 100644 --- a/zig/kernels/metal/sources.zig +++ b/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"); diff --git a/zig/src/core/contract.zig b/zig/src/core/contract.zig index 2c976103b2..a04a28ec3f 100644 --- a/zig/src/core/contract.zig +++ b/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; diff --git a/zig/src/core/contract/metal_pair_records.zig b/zig/src/core/contract/metal_pair_records.zig new file mode 100644 index 0000000000..5e9eebe17d --- /dev/null +++ b/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); + } +} diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 61a83c50a4..b03f65e89a 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -570,7 +570,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 ++ @import("lfm2_records.zig").records ++ @import("nemotron_experts_records.zig").records ++ +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 ++ @import("lfm2_records.zig").records ++ @import("nemotron_experts_records.zig").records ++ @import("moss_image_records.zig").records; const learning = @import("registry_learning.zig"); pub const gradient_pending = learning.gradient_pending; diff --git a/zig/src/core/pair_rows.zig b/zig/src/core/pair_rows.zig new file mode 100644 index 0000000000..38017cb531 --- /dev/null +++ b/zig/src/core/pair_rows.zig @@ -0,0 +1,463 @@ +//! Two Macs' exchanges: GPU part, posted(x), host MCDMA relay, landed(x); the event wait is a real scheduler hand-off. +const std = @import("std"); +const platform = @import("sys"); +const mtl = @import("metal"); +const fabric = @import("fabric"); +const sources = @import("kernel_sources"); +const settings = @import("../families/flashnext/tp_settings.zig"); +const Fence = @import("fence.zig").Fence; +const Owner = @import("../families/flashnext/owned.zig").Owner; +const pair_layout = @import("pair_rows_layout.zig"); +const pair_loopback = @import("pair_rows_loopback.zig"); + +pub const Settings = settings.Settings; +pub const Ref = @import("exlr8_projection.zig").Ref; + +pub const command_words = pair_layout.command_words; +pub const Slot = pair_layout.Slot; +pub const seam_failed = pair_layout.seam_failed; +pub const Seam = pair_layout.Seam; +pub const elems_per_word = pair_layout.elems_per_word; +pub const stageBytes = pair_layout.stageBytes; +pub const gatherBytes = pair_layout.gatherBytes; +pub const pickBytes = pair_layout.pickBytes; +pub const frameBytes = pair_layout.frameBytes; +pub const toFrame = pair_layout.toFrame; +pub const seedValue = pair_layout.seedValue; +pub const Loopback = pair_loopback.Loopback; + +const PAGE = pair_layout.PAGE; +const pages = pair_layout.pages; +const CMD_BYTES = pair_layout.CMD_BYTES; +const JOBS = pair_layout.JOBS; +const slotOf = pair_layout.slotOf; +const Layout = pair_layout.Layout; +const FLAG = pair_layout.FLAG; +const HELLO = pair_layout.HELLO; +const CMD_FLAG = pair_layout.CMD_FLAG; +const CMD_ACK = pair_layout.CMD_ACK; +const READY = pair_layout.READY; +const BEAT = pair_layout.BEAT; +const POSTED = pair_layout.POSTED; +const FAILED = pair_layout.FAILED; +const OBSERVED = pair_layout.OBSERVED; +const PairRows = pair_layout.PairRows; +const seam = pair_layout.seam; + +/// Mach's clock: 24 MHz, so a millisecond is 24,000 ticks -- the clock the probes' seam numbers are in. +const ticks_per_ms: u64 = 24_000; + +/// How long the host waits for the GPU to reach a published seam: a hung GPU must fail the request, not wedge it. +const post_timeout_ms: u64 = 5_000; +/// The peer's part wait: 250 ms, 300x the measured tail (p99 ~700 us), so a dead peer fails fast. +const peer_timeout_ms: u64 = 250; +/// The cold-link peer wait: the first request pays one-time Metal pipelines (1.07 s measured, 0.14 s warm). +const cold_peer_timeout_ms: u64 = 5_000; +/// How long after the link's first relay the cold bound applies: 8 s covers a longer first prompt (prefill chunks). +const cold_window_ms: u64 = 8_000; + +/// How long rank 1 waits for rank 0's acknowledgement of a replay: a stuck peer fails the request, it is not wedged. +const cmd_ack_timeout_ms: u64 = 60_000; + +/// How long either rank waits for the other's ready word: a cold 27B load on the far Mac, five minutes at most. +const ready_timeout_ms: u64 = 300_000; + +/// The close's last flush: microseconds when the peer is there, and never the library's whole timeout when it is gone. +const last_flush_ns: u64 = 2 * std.time.ns_per_s; + +threadlocal var raised = false; + +/// This thread at the user-interactive QoS once: busy CPUs must not delay the pair's sends and commands. +pub fn raise() void { + if (raised) return; + _ = platform.c.pthread_set_qos_class_self_np(.USER_INTERACTIVE, 0); + raised = true; +} + +/// Whether 32-bit counter a has reached b across the wrap. +fn reached(a: u32, b: u32) bool { + return @as(i32, @bitCast(a -% b)) >= 0; +} + +fn ticksNow() u64 { + return platform.c.mach_absolute_time(); +} + +/// A short sleep, for the host threads' waits that are not on the seam's critical path. +fn sleepUs(us: u64) void { + const ts: platform.c.timespec = .{ .sec = @intCast(us / 1_000_000), .nsec = @intCast(us % 1_000_000 * 1000) }; + _ = platform.c.nanosleep(&ts, null); +} + +pub const Pair = struct { + rank: u32, + peer: u32, + l: Layout, + ep: *fabric.mcdma.Endpoint, + rd: fabric.rdma.Rdma, + win: []u8, + wbuf: mtl.Buffer, + stage_pipe: mtl.Pipeline, + gather_pipe: mtl.Pipeline, + argmax_pipe: mtl.Pipeline, + pick_pipe: mtl.Pipeline, + /// Orders the window's untracked writes across the encoders a seam splits. + fence: Fence, + /// GPU -> host: seam x's part is written, relay it. Host -> GPU: the peer's part for x is in place, add it. + posted: mtl.SharedEvent, + landed: mtl.SharedEvent, + /// Exchanges encoded so far: the same count on both Macs. + x: u64 = 0, + bytes: []u32, + queued: std.atomic.Value(u32) = .init(0), + thread: ?std.Thread = null, + stop: std.atomic.Value(bool) = .init(false), + broken: std.atomic.Value(bool) = .init(false), + /// The exchange the service thread is on: the round number in a failure line. + seam_x: std.atomic.Value(u32) = .init(0), + /// The request this pair is on, for failure lines; `beginRound` bumps it -- a label only, failures are permanent. + round: std.atomic.Value(u32) = .init(0), + /// Commands sent (rank 0) or taken (rank 1) so far. + cmd: u64 = 0, + /// Beats this Mac's watchdog has signalled to the peer: its counter for the peer's silence check. + beats: u64 = 0, + /// The GPU objects this pair holds: released in `deinit`, the window's buffer before the memory behind it. + owned: Owner(mtl) = .{}, + /// The allocator `init` was given: `bytes`, the settings arena and this struct go back to it. + gpa: std.mem.Allocator, + /// The parsed settings file: the endpoint's config points into it, so it outlives `ep`. + arena: *std.heap.ArenaAllocator, + + /// Connect to the peer named in `settings_path` (speed-up mode's file), exchanges up to `bytes`, and start sending. + pub fn init(gpa: std.mem.Allocator, io: std.Io, device: mtl.Device, settings_path: []const u8, bytes: usize) !*Pair { + const l: Layout = .{ .slot = pages(bytes) }; + const text = try std.Io.Dir.cwd().readFileAlloc(io, settings_path, gpa, .limited(1 << 20)); + defer gpa.free(text); + const arena = try gpa.create(std.heap.ArenaAllocator); + errdefer gpa.destroy(arena); + arena.* = .init(gpa); + errdefer arena.deinit(); + // the endpoint's config keeps `s.links`, so the parse outlives `text` in the arena, not on the stack + const s = try settings.read(arena.allocator(), text); + if (s.rank > 1 or s.links.len != 1) return error.PairTwoRanksOneLink; + const lib = try gpa.dupeSentinel(u8, s.library, 0); + defer gpa.free(lib); + const ep = try fabric.mcdma.Endpoint.create(gpa, lib, .{ .rank = s.rank, .ranks = 2, .window_bytes = l.window(), .staging_bytes = 4 << 20, .links = s.links, .timeout_ns = 60 * std.time.ns_per_s, .connect_timeout_ns = 300 * std.time.ns_per_s }); + errdefer ep.deinit(); + var owned: Owner(mtl) = .{}; + errdefer owned.deinit(); + const rd = ep.rdma(); + const win = rd.window(); + const lib_m = try mtl.Library.fromSource(device, sources.core_pair_rows, mtl.CompileOptions.mlx()); + defer lib_m.deinit(); + const fence = try Fence.init(device); + errdefer fence.deinit(); + const posted = try device.sharedEvent(); + errdefer posted.deinit(); + const landed = try device.sharedEvent(); + errdefer landed.deinit(); + const p = try gpa.create(Pair); + errdefer gpa.destroy(p); + const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; + p.* = .{ + .rank = s.rank, + .peer = 1 - s.rank, + .l = l, + .ep = ep, + .rd = rd, + .win = win, + .wbuf = try owned.keep(gpa, try device.bufferNoCopy(win.ptr, win.len, opts)), + .stage_pipe = try owned.keep(gpa, try mtl.Pipeline.init(device, lib_m, "pair_stage", false)), + .gather_pipe = try owned.keep(gpa, try mtl.Pipeline.init(device, lib_m, "pair_gather", false)), + .argmax_pipe = try owned.keep(gpa, try mtl.Pipeline.init(device, lib_m, "pair_argmax", false)), + .pick_pipe = try owned.keep(gpa, try mtl.Pipeline.init(device, lib_m, "pair_pick", false)), + .fence = fence, + .posted = posted, + .landed = landed, + .bytes = try gpa.alloc(u32, JOBS), + .gpa = gpa, + .arena = arena, + .owned = owned, + }; + owned = .{}; // the pair owns them now: the error path must not release them twice + errdefer gpa.free(p.bytes); + // both Macs up before the first exchange: a word each way + try rd.signal(p.peer, l.flags() + HELLO, 1); + while (@atomicLoad(u64, p.word64(l.flags() + HELLO), .acquire) < 1) std.atomic.spinLoopHint(); + p.thread = try std.Thread.spawn(.{}, service, .{p}); + return p; + } + + /// Rank 0 tells rank 1 it is done (its `next` returns null); then the sending thread and the link close. + pub fn deinit(p: *Pair) void { + if (p.rank == 0 and !p.failed()) p.send(0, &.{}) catch {}; + p.stop.store(true, .release); + if (p.thread) |t| t.join(); + if (!p.failed()) flushLast(p); + p.posted.deinit(); + p.landed.deinit(); + p.fence.deinit(); + p.owned.deinit(); // the window's buffer before the endpoint frees the memory it aliases + p.ep.deinit(); + const gpa = p.gpa; + gpa.free(p.bytes); + p.arena.deinit(); + gpa.destroy(p.arena); + gpa.destroy(p); + } + + /// Hand the last seam over without the library's whole timeout: a peer that is gone leaves nothing to deliver. + fn flushLast(p: *Pair) void { + if (fabric.words.checkedNowNs()) |now| p.rd.setDeadline(now + last_flush_ns) else |_| return; + p.rd.flush() catch {}; + } + + fn word64(p: *const Pair, off: usize) *u64 { + return @ptrCast(@alignCast(p.win.ptr + off)); + } + fn word32(p: *const Pair, off: usize) *u32 { + return @ptrCast(@alignCast(p.win.ptr + off)); + } + + /// Seams the host gave up on: nonzero means the run is invalid. + pub fn gaveUp(p: *const Pair) u32 { + return @atomicLoad(u32, p.word32(p.l.sync() + FAILED * 4), .acquire); + } + + /// Seams the GPU itself refused after seeing FAILED: this side's own record of a failed relay. + pub fn observed(p: *const Pair) u32 { + return @atomicLoad(u32, p.word32(p.l.sync() + OBSERVED * 4), .acquire); + } + + /// The exchange the service thread is on (or the one it gave up on, if it has). + pub fn failedSeam(p: *const Pair) u32 { + return p.seam_x.load(.acquire); + } + + /// Labels a request; never a reset: after a seam failure the pair is dead (a half-written ring corrupts). + pub fn beginRound(p: *Pair) u32 { + return p.round.fetchAdd(1, .monotonic) + 1; + } + + /// Where this rank's part for exchange x goes, and where the peer's lands: the relay's two addresses. + pub fn sendOffset(p: *const Pair, x: u32) usize { + return p.l.send(x); + } + pub fn recvOffset(p: *const Pair, x: u32) usize { + return p.l.recv(x); + } + + /// The link is down: failed relay, dead peer, GPU never at the seam; check after every seam and command buffer. + pub fn failed(p: *const Pair) bool { + return p.broken.load(.acquire) or p.gaveUp() != 0; + } + + /// The peer's watchdog is alive: its last beat, a word its watchdog wrote here; silence means the peer is gone. + pub fn heartbeats(p: *const Pair) u64 { + return @atomicLoad(u64, p.word64(p.l.flags() + pair_layout.BEAT), .acquire); + } + + /// Tell the peer this rank's watchdog is alive; a failed signal is the peer gone, not nothing to do. + pub fn beat(p: *Pair) void { + p.beats +%= 1; + p.rd.signal(p.peer, p.l.flags() + pair_layout.BEAT, p.beats) catch p.broken.store(true, .release); + } + + /// The next exchange: take it before the producer writes, then pass the same x to stage/gather/argmax/pick. + pub fn nextX(p: *const Pair) u32 { + return @truncate(p.x + 1); + } + + /// The exchange slot for this seam's part: where this rank's rows go, and where the peer's land. + pub fn part(p: *const Pair) Seam { + const x: u32 = @truncate(p.x + 1); + return .{ .x = x, .ref = .{ .buf = p.wbuf, .off = p.l.send(x) } }; + } + + /// Publish x and its bytes to the service thread; x must be `part`'s, so an out-of-order one panics in release too. + fn queue(p: *Pair, x: u32, bytes: usize) void { + if (x != @as(u32, @truncate(p.x + 1))) @panic("pair: an exchange was encoded out of order (part() not called first)"); + if (bytes > p.l.slot) @panic("pair: this exchange's part does not fit the window's slot (speed-up's rows too large)"); + p.x += 1; + p.bytes[x % JOBS] = @intCast(bytes); + p.queued.store(x, .release); + } + + /// Before `seam`: this rank's `width` columns of `rows` rows from `src` into the slot, bf16 to bf16, no arithmetic. + pub fn stage(p: *Pair, e: *mtl.ComputeEncoder, x: u32, src: Ref, rows: u32, width: u32, stride: u32) void { + if (rows == 0 or width == 0 or width > stride) @panic("pair: stage() wants rows > 0 and a width inside the source row"); + e.setPipeline(p.stage_pipe); + e.setBuffer(p.wbuf, p.l.sync(), 0); + e.setValue(PairRows{ .width = width, .stride = stride, .rank = p.rank, .x = x }, 1); + e.setBuffer(src.buf, src.off, 2); + e.setBuffer(p.wbuf, p.l.send(x), 3); + e.dispatchGroups(.{ .width = rows }, .{ .width = 256 }); + } + + /// After `stage` (or the producer): relay, then pack this rank's `width` at `rank * width` and the peer's after. + pub fn gather(p: *Pair, e: *mtl.ComputeEncoder, x: u32, mine: Ref, out: Ref, rows: u32, width: u32, out_stride: u32) void { + if (rows == 0 or width == 0 or width % elems_per_word != 0) @panic("pair: gather() wants rows > 0 and a width that fills whole window words"); + if (out_stride < 2 * width) @panic("pair: gather()'s frame row is narrower than both halves"); + if (mine.buf.id != p.wbuf.id or mine.off != p.l.send(x)) @panic("pair: gather()'s part is not the region part() returned"); + p.queue(x, stageBytes(rows, width)); + seam(e, p.fence, p.posted, p.landed, x); + e.setPipeline(p.gather_pipe); + e.setBuffer(p.wbuf, p.l.sync(), 0); + e.setValue(PairRows{ .width = width, .stride = out_stride, .rank = p.rank, .x = x }, 1); + e.setBuffer(mine.buf, mine.off, 2); + e.setBuffer(p.wbuf, p.l.recv(x), 3); + e.setBuffer(out.buf, out.off, 4); + e.dispatchGroups(.{ .width = rows }, .{ .width = 256 }); + } + + /// Each row's argmax over this rank's n fp32 logits from id lo, row stride `stride`, then the pick over both Macs'. + pub fn argmax(p: *Pair, e: *mtl.ComputeEncoder, x: u32, logits: Ref, stride: u32, lo: u32, n: u32, rows: u32, picks: Ref) void { + p.argmaxWidth(e, x, logits, stride, lo, n, rows, picks, 0); + } + + /// The same seam over bf16 logits (Qwen3.8-27B's `frame.logits`): one kernel, `dims.w` selects the load width. + pub fn argmaxBf16(p: *Pair, e: *mtl.ComputeEncoder, x: u32, logits: Ref, stride: u32, lo: u32, n: u32, rows: u32, picks: Ref) void { + p.argmaxWidth(e, x, logits, stride, lo, n, rows, picks, 1); + } + + /// `dtype` 0 is fp32 logits, 1 is bf16; the exchange is (fp32 value, id) either way, so pair_pick is one path. + fn argmaxWidth(p: *Pair, e: *mtl.ComputeEncoder, x: u32, logits: Ref, stride: u32, lo: u32, n: u32, rows: u32, picks: Ref, dtype: u32) void { + if (rows == 0 or n == 0) @panic("pair: argmax() wants rows > 0 and n > 0"); + p.queue(x, pickBytes(rows)); + e.setPipeline(p.argmax_pipe); + e.setBuffer(logits.buf, logits.off, 0); + e.setValue([4]u32{ stride, lo, n, dtype }, 1); + e.setBuffer(p.wbuf, p.l.send(x), 2); + e.dispatchGroups(.{ .width = rows }, .{ .width = 1024 }); + seam(e, p.fence, p.posted, p.landed, x); // the encoder boundary orders the argmax before the part is relayed + e.setPipeline(p.pick_pipe); + e.setBuffer(p.wbuf, p.l.sync(), 0); + e.setValue([2]u32{ rows, x }, 1); + e.setBuffer(p.wbuf, p.l.send(x), 2); + e.setBuffer(p.wbuf, p.l.recv(x), 3); + e.setBuffer(picks.buf, picks.off, 4); + e.dispatchGroups(.{ .width = 1 }, .{ .width = 32 }); + } + + /// Rank 0: command `kind` (never 0, which ends rank 1's loop) and its words, once rank 1 took the one before. + pub fn send(p: *Pair, kind: u32, words: []const u32) !void { + if (p.rank != 0) return error.FollowsPeer; + if (words.len > command_words) return error.CommandTooLong; + raise(); + try p.reach(p.l.flags() + CMD_ACK, p.cmd, cmd_ack_timeout_ms); + p.cmd += 1; + const head: [2]u32 = .{ kind, @intCast(words.len) }; + const out = p.l.cmdOut(); + @memcpy(p.win[out..][0..8], std.mem.asBytes(&head)); + @memcpy(p.win[out + 8 ..][0 .. 4 * words.len], std.mem.sliceAsBytes(words)); + try p.ep.writeSignalFrom(p.peer, out, p.l.cmdIn(), 8 + 4 * words.len, p.l.flags() + CMD_FLAG, p.cmd); + } + + /// Rank 1: rank 0's next command's words into `out` and its kind; null once rank 0 is done or this Mac stops. + pub fn next(p: *Pair, gpa: std.mem.Allocator, out: *std.ArrayList(u32)) !?u32 { + raise(); + p.reach(p.l.flags() + CMD_FLAG, p.cmd + 1, 0) catch |err| return if (err == error.PairStopped) null else err; + p.cmd += 1; + const in = p.l.cmdIn(); + const head: *const [2]u32 = @ptrCast(@alignCast(p.win.ptr + in)); + if (head[1] > command_words) return error.CommandOutOfStep; + out.clearRetainingCapacity(); + try out.appendSlice(gpa, @as([*]const u32, @ptrCast(@alignCast(p.win.ptr + in + 8)))[0..head[1]]); + const kind = head[0]; + try p.rd.signal(p.peer, p.l.flags() + CMD_ACK, p.cmd); + return if (kind == 0) null else kind; + } + + /// Both Macs loaded before the first command: a word each way, each waited for with the load's bound. + pub fn ready(p: *Pair) !void { + try p.rd.signal(p.peer, p.l.flags() + READY, 1); + try p.reach(p.l.flags() + READY, 1, ready_timeout_ms); + } + + /// This Mac's command wait ends (rank 1 at close). + pub fn stopFollowing(p: *Pair) void { + p.stop.store(true, .release); + } + + /// Wait until this Mac's word at `off` is `value`: spins, then short sleeps, or fails after `limit_ms` (0: none). + fn reach(p: *Pair, off: usize, value: u64, limit_ms: u64) !void { + const t0 = ticksNow(); + var spins: usize = 0; + while (@atomicLoad(u64, p.word64(off), .acquire) < value) { + if (p.stop.load(.acquire)) return error.PairStopped; + if (p.failed()) return error.PairLinkFailed; + if (limit_ms != 0 and ticksNow() -% t0 > limit_ms * ticks_per_ms) return error.PairPeerSilent; + spins += 1; + if (spins > 200_000) sleepUs(20) else std.atomic.spinLoopHint(); + } + } + + /// Give up on x: release every event wait with the sentinel, so the blocked GPU resumes and the request aborts. + fn fail(p: *Pair, x: u32, comptime what: []const u8, peer_ms: u64) void { + p.broken.store(true, .release); + _ = @atomicRmw(u32, p.word32(p.l.sync() + FAILED * 4), .Add, 1, .monotonic); + std.log.err("pair rank {d}: request {d} seam {d} failed: {s}; the link is down and stays down (posted {d}, flag {d}, queued {d}, GPU-observed {d}; deadlines {d} ms post, {d} ms peer): restart both processes. The blocked GPU is released with the failure sentinel", .{ + p.rank, p.round.load(.acquire), x, what, p.posted.value(), @atomicLoad(u64, p.word64(p.l.flags() + FLAG), .acquire), p.queued.load(.acquire), p.observed(), post_timeout_ms, peer_ms, + }); + p.landed.set(seam_failed); + } + + /// One exchange at a time: part written, relay it, peer's lands, release the GPU; every wait has a deadline. + fn service(p: *Pair) void { + raise(); + var x: u32 = 1; + // Ticks after which this link is no longer cold: set at the first relay, when both ranks have just loaded. + var cold_until: u64 = 0; + while (!p.stop.load(.acquire)) { + p.seam_x.store(x, .release); + // 1. The GPU must reach the seam: sleep while idle, spin once published, fail if it never arrives. + var armed: u64 = 0; + while (p.posted.value() < @as(u64, x)) { + if (p.stop.load(.acquire)) return; + if (!reached(p.queued.load(.acquire), x)) { + sleepUs(20); + continue; + } + if (armed == 0) armed = ticksNow(); + if (ticksNow() -% armed > post_timeout_ms * ticks_per_ms) return p.fail(x, "this Mac's GPU never reached the seam", peer_timeout_ms); + std.atomic.spinLoopHint(); + } + // 2. Relay: part to the peer's recv slot, arrival flag in the same MCDMA write, data and count together. + const at = x % JOBS; + p.ep.writeSignalFrom(p.peer, p.l.send(x), p.l.recv(x), p.bytes[at], p.l.flags() + FLAG, x) catch |err| { + std.log.err("pair rank {d}: seam {d} relay write failed: {s}", .{ p.rank, x, @errorName(err) }); + return p.fail(x, "the relay write failed", peer_timeout_ms); + }; + // 3. The peer's relay puts its part in our slot; its flag is proof; the cold bound, then the steady one. + if (x == 1) cold_until = ticksNow() +% cold_window_ms * ticks_per_ms; + const cold = cold_until != 0 and ticksNow() < cold_until; + const budget_ms = if (cold) cold_peer_timeout_ms else peer_timeout_ms; + const t1 = ticksNow(); + while (@atomicLoad(u64, p.word64(p.l.flags() + FLAG), .acquire) < x) { + if (p.stop.load(.acquire)) return; + if (ticksNow() -% t1 > budget_ms * ticks_per_ms) return p.fail(x, "the peer's part never landed", budget_ms); + std.atomic.spinLoopHint(); + } + // 4. Release the GPU: both parts are in; the event's wait is the acquire that makes them visible. + @atomicStore(u32, p.word32(p.l.sync() + POSTED * 4), x, .release); + p.landed.set(x); + x +%= 1; + } + } +}; + +test "the window's regions do not overlap and its words sit on their pages" { + const l: Layout = .{ .slot = pages(stageBytes(256, 6144)) }; + try std.testing.expect(l.send(1) >= l.recv(1) + l.slot and l.send(0) >= 2 * l.slot + PAGE and l.send(0) + l.slot <= l.send(1)); + try std.testing.expect(l.cmdIn() >= l.send(1) + l.slot and l.cmdOut() >= l.cmdIn() + CMD_BYTES + 64); + try std.testing.expect(l.flags() >= l.cmdOut() + 8 + 4 * command_words and BEAT + 8 <= PAGE and l.window() % PAGE == 0); + try std.testing.expect(reached(1, 0xFFFFFFFF) and !reached(5, 6)); +} + +test "the parity ring alternates and one sentinel releases every exchange" { + try std.testing.expect(slotOf(1) == .b and slotOf(2) == .a and slotOf(3) == .b); + const l: Layout = .{ .slot = pages(stageBytes(16, 5120)) }; + try std.testing.expect(l.recv(1) == l.slot and l.recv(2) == 0 and l.send(2) == 2 * l.slot + PAGE); + try std.testing.expectEqual(pages(stageBytes(16, 5120)), l.slot); // the 27B's 16-row prefill part, page-rounded + try std.testing.expect(seam_failed > std.math.maxInt(u32)); // every waitFor(x) in a buffer is satisfied by it +} diff --git a/zig/src/core/pair_rows_layout.zig b/zig/src/core/pair_rows_layout.zig new file mode 100644 index 0000000000..a1a1ff14d6 --- /dev/null +++ b/zig/src/core/pair_rows_layout.zig @@ -0,0 +1,125 @@ +//! The seam window's regions and the bf16 frame ABI, shared by the link and the host-copy Loopback. +const std = @import("std"); +const mtl = @import("metal"); +const Fence = @import("fence.zig").Fence; + +pub const Ref = @import("exlr8_projection.zig").Ref; + +pub const PAGE = 16384; +pub fn pages(n: usize) usize { + return (n + PAGE - 1) / PAGE * PAGE; +} +/// A command's words at most: a 262,144-token prompt and its head. +pub const command_words = 262144 + 64; +pub const CMD_BYTES = pages(8 + 4 * command_words); +pub const JOBS = 4096; + +/// Which parity slot x uses: rank r cannot finish x without the peer's x, so the peer is at most one exchange ahead. +pub const Slot = enum(u1) { a = 0, b = 1 }; + +pub fn slotOf(x: u32) Slot { + return @fromBackingInt(@intCast(x % 2)); +} + +/// The window on both Macs: parts in/out by parity (a page before each out), commands, flags, host words. +pub const Layout = struct { + slot: usize, + + pub fn recv(l: Layout, x: u32) usize { + return @as(usize, @backingInt(slotOf(x))) * l.slot; + } + pub fn send(l: Layout, x: u32) usize { + return 2 * l.slot + @as(usize, @backingInt(slotOf(x))) * (PAGE + l.slot) + PAGE; + } + pub fn cmdIn(l: Layout) usize { + return 4 * l.slot + 2 * PAGE; + } + pub fn cmdOut(l: Layout) usize { + return l.cmdIn() + CMD_BYTES + PAGE; + } + pub fn flags(l: Layout) usize { + return l.cmdOut() + CMD_BYTES; + } + pub fn sync(l: Layout) usize { + return l.flags() + PAGE; + } + pub fn window(l: Layout) usize { + return l.sync() + PAGE; + } +}; +pub const FLAG = 0; // u64 in the flag page: the last exchange whose peer part has landed (written by the peer's relay) +pub const HELLO = 64; +pub const CMD_FLAG = 128; // u64, at rank 1: the last command rank 0 has written +pub const CMD_ACK = 192; // u64, at rank 0: the last command rank 1 has copied out +pub const READY = 256; // u64: the peer has loaded its model +pub const BEAT = 320; // u64: the peer's watchdog ran; a rank idle in no round still sees a dead peer through this + +// The sync page's words shared with pair_rows.metal: POSTED last released exchange, FAILED give-ups, OBSERVED refusals. +pub const POSTED = 0; +pub const FAILED = 16; +pub const OBSERVED = 20; + +pub const seam_failed: u64 = std.math.maxInt(u64); + +pub const Seam = struct { x: u32, ref: Ref }; + +/// Elements in one 32-bit window word: the relay and the link move whole words, and a frame element is a bf16. +pub const elems_per_word: usize = 2; + +/// `stage`'s outgoing copy: this rank's `cols` of `rows` rows, bf16 -- the frame seam's payload. +pub fn stageBytes(rows: u32, cols: u32) usize { + return @as(usize, rows) * cols * @sizeOf(u16); +} + +/// `gather`'s `out`: both halves of `rows` rows packed into the frame, so the frame row is `2 * cols` wide. +pub fn gatherBytes(rows: u32, cols: u32) usize { + return @as(usize, rows) * 2 * cols * @sizeOf(u16); +} + +/// The head seam's exchange: an (fp32 value, u32 id) a row; pair_pick compares logits, not words. +pub fn pickBytes(rows: u32) usize { + return @as(usize, rows) * @sizeOf(f32) * 2; +} + +/// A frame tensor's bytes: the bf16 grid `stage` reads and `gather` writes its halves into. +pub fn frameBytes(rows: u32, width: u32) usize { + return @as(usize, rows) * width * @sizeOf(u16); +} + +/// One fp32 value as the frame stores it (bf16, nearest even), the host side of the kernels' frame elements. +pub fn toFrame(v: f32) u16 { + return bf16Bits(v); +} + +/// One rounding to bf16, nearest-even: the host side of a frame element (pair_rows.metal copies the bits it is given). +fn bf16Bits(v: f32) u16 { + const u: u32 = @bitCast(v); + return @truncate((u +% 0x7FFF +% ((u >> 16) & 1)) >> 16); +} + +/// The seed pattern as a frame element: 1 + m/128, exact in both bf16 and fp32, so expected bytes hold either way. +pub fn seedValue(r: u32, i: u32, salt: u16) f32 { + const k = r *% 5120 +% i *% 37; + const m: u32 = (salt & 0x1F) + ((k *% 53) & 0x1F); + return 1.0 + @as(f32, @floatFromInt(m)) / 128.0; +} + +/// One frame element at `at` of a bf16 payload: `stage`'s and `gather`'s slot is the frame's bits, two bytes each. +pub fn writeHalf(out: []u8, at: usize, v: f32) void { + std.mem.writeInt(u16, out[at * @sizeOf(u16) ..][0..@sizeOf(u16)], toFrame(v), .little); +} + +/// The kernels' argument block: four u32, so its layout and pair_rows.metal's struct stay one struct. +pub const PairRows = extern struct { width: u32, stride: u32, rank: u32, x: u32 }; + +/// Between encoders only (core/segments.zig:59-66): fence, close, posted, wait, reopen, wait -- a real hand-off. +pub fn seam(enc: *mtl.ComputeEncoder, fence: Fence, posted: mtl.SharedEvent, landed: mtl.SharedEvent, x: u32) void { + const cb = enc.commandBuffer(); + const mode = enc.dispatchType(); + fence.update(enc.*); + enc.end(); + cb.signal(posted, @as(u64, x)); + cb.waitFor(landed, @as(u64, x)); + enc.* = cb.compute(mode); + fence.wait(enc.*); +} diff --git a/zig/src/core/pair_rows_loopback.zig b/zig/src/core/pair_rows_loopback.zig new file mode 100644 index 0000000000..5086fa1e3e --- /dev/null +++ b/zig/src/core/pair_rows_loopback.zig @@ -0,0 +1,249 @@ +//! The same seam with a host copy instead of the link: the numeric reference the real transport is checked against. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const Fence = @import("fence.zig").Fence; +const pair_layout = @import("pair_rows_layout.zig"); + +const FAILED = pair_layout.FAILED; +const JOBS = pair_layout.JOBS; +const Layout = pair_layout.Layout; +const PAGE = pair_layout.PAGE; +const PairRows = pair_layout.PairRows; +const Seam = pair_layout.Seam; +const elems_per_word = pair_layout.elems_per_word; +const pages = pair_layout.pages; +const stageBytes = pair_layout.stageBytes; +const pickBytes = pair_layout.pickBytes; +const seam = pair_layout.seam; +const seam_failed = pair_layout.seam_failed; +const writeHalf = pair_layout.writeHalf; +const Ref = pair_layout.Ref; + +/// The same seam with a host copy instead of the link: the numeric reference; bytes arrive via `arrive`/`seed`. +pub const Loopback = struct { + allocator: std.mem.Allocator, + rank: u32, + l: Layout, + win: []u8, + wbuf: mtl.Buffer, + stage_pipe: mtl.Pipeline, + gather_pipe: mtl.Pipeline, + argmax_pipe: mtl.Pipeline, + pick_pipe: mtl.Pipeline, + /// Per-exchange part bytes (stage/gather and argmax queue different sizes): the relay copies exactly this much. + bytes: []u32, + fence: Fence, + /// This rank's parts: the GPU signals it per exchange, and the relay reads `bytesFor(x)` after that. + posted: mtl.SharedEvent, + /// The peer's parts: the relay sets it once the copy has landed, and the seam waits for it. + landed: mtl.SharedEvent, + x: u64 = 0, + relay: ?std.Thread = null, + stop: std.atomic.Value(bool) = .init(false), + + pub fn init(device: mtl.Device, bytes: usize, rank: u32) !Loopback { + const l: Layout = .{ .slot = pages(bytes) }; + // the window is a no-copy Metal buffer: page-aligned, page-multiple allocation, never the gpa's good luck + const win = try std.heap.page_allocator.alloc(u8, l.window()); + errdefer std.heap.page_allocator.free(win); + std.debug.assert(@intFromPtr(win.ptr) % PAGE == 0 and win.len % PAGE == 0); + @memset(win, 0); + const lib = try mtl.Library.fromSource(device, sources.core_pair_rows, mtl.CompileOptions.mlx()); + defer lib.deinit(); + const wbuf = try device.bufferNoCopy(win.ptr, win.len, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + errdefer wbuf.deinit(); + const fence = try Fence.init(device); + errdefer fence.deinit(); + const posted = try device.sharedEvent(); + errdefer posted.deinit(); + const landed = try device.sharedEvent(); + errdefer landed.deinit(); + const stage_pipe = try mtl.Pipeline.init(device, lib, "pair_stage", false); + errdefer stage_pipe.deinit(); + const gather_pipe = try mtl.Pipeline.init(device, lib, "pair_gather", false); + errdefer gather_pipe.deinit(); + const argmax_pipe = try mtl.Pipeline.init(device, lib, "pair_argmax", false); + errdefer argmax_pipe.deinit(); + const pick_pipe = try mtl.Pipeline.init(device, lib, "pair_pick", false); + errdefer pick_pipe.deinit(); + const sizes = try std.heap.page_allocator.alloc(u32, JOBS); + errdefer std.heap.page_allocator.free(sizes); + @memset(sizes, 0); + return .{ .allocator = std.heap.page_allocator, .rank = rank, .l = l, .win = win, .wbuf = wbuf, .stage_pipe = stage_pipe, .gather_pipe = gather_pipe, .argmax_pipe = argmax_pipe, .pick_pipe = pick_pipe, .bytes = sizes, .fence = fence, .posted = posted, .landed = landed }; + } + + pub fn deinit(p: *Loopback) void { + p.stopRelay(); + p.stage_pipe.deinit(); + p.gather_pipe.deinit(); + p.argmax_pipe.deinit(); + p.pick_pipe.deinit(); + p.allocator.free(p.bytes); + p.wbuf.deinit(); + p.posted.deinit(); + p.landed.deinit(); + p.fence.deinit(); + p.allocator.free(p.win); + } + + fn word32(p: *const Loopback, off: usize) *u32 { + return @ptrCast(@alignCast(p.win.ptr + off)); + } + + /// Where this rank's part for exchange x goes (what the relay copies out, once `posted` passes x). + pub fn sendOffset(p: *const Loopback, x: u32) usize { + return p.l.send(x); + } + + /// Where the peer's part for exchange x goes (what the relay copies in, before setting `landed` to x). + pub fn recvOffset(p: *const Loopback, x: u32) usize { + return p.l.recv(x); + } + + /// Where the producer writes this rank's bf16 columns of the next exchange; the same shape as `Pair.part`. + pub fn nextX(p: *const Loopback) u32 { + return @truncate(p.x + 1); + } + + pub fn part(p: *const Loopback) Seam { + const x: u32 = @truncate(p.x + 1); + return .{ .x = x, .ref = .{ .buf = p.wbuf, .off = p.l.send(x) } }; + } + + /// The seed pattern as a frame element: 1 + m/128, exact in bf16, so expected bytes hold either way. + fn seedValue(r: u32, i: u32, salt: u16) f32 { + const k = r *% 5120 +% i *% 37; + const m: u32 = (salt & 0x1F) + ((k *% 53) & 0x1F); + return 1.0 + @as(f32, @floatFromInt(m)) / 128.0; + } + + /// Fill both slots with the seed and mark all arrived: arithmetic check only, timing must not use this. + pub fn seed(p: *Loopback, rows: u32, width: u32) void { + const bytes = stageBytes(rows, width); + if (bytes > p.l.slot) @panic("loopback: the seeded part does not fit the window's slot"); + for (0..2) |x| { + const recv = p.win[p.l.recv(@intCast(x))..][0..bytes]; + const send = p.win[p.l.send(@intCast(x))..][0..bytes]; + for (0..rows) |r| for (0..width) |i| { + const at = r * width + i; + writeHalf(recv, at, seedValue(@intCast(r), @intCast(i), 0x15)); + writeHalf(send, at, seedValue(@intCast(r), @intCast(i), 0x0A)); + }; + } + p.satisfy(); + } + + /// Release every wait with the sentinel, as a failed link does: arithmetic check only, timing runs `startRelay`. + pub fn satisfy(p: *Loopback) void { + p.landed.set(seam_failed); + } + + /// The in-process relay: wait for posted, copy in, release -- Pair's two events in order, memcpy for the wire. + pub fn startRelay(p: *Loopback) !void { + std.debug.assert(p.relay == null); + p.stop.store(false, .release); + p.relay = try std.Thread.spawn(.{}, relayLoop, .{p}); + } + + pub fn stopRelay(p: *Loopback) void { + p.stop.store(true, .release); + if (p.relay) |t| t.join(); + p.relay = null; + } + + fn relayLoop(p: *Loopback) void { + var x: u64 = 1; + while (!p.stop.load(.acquire)) { + if (p.posted.value() < x) continue; // spin: the host's side of the seam is part of what is measured + const at: u32 = @truncate(x); + const n: usize = p.bytesFor(at); + @memcpy(p.win[p.l.recv(at)..][0..n], p.win[p.l.send(at)..][0..n]); + p.landed.set(x); + x += 1; + } + } + + /// The part's bytes for exchange `x` as the last `gather`/`argmax` queued it (the relay's copy length). + pub fn bytesFor(p: *const Loopback, x: u32) u32 { + return @atomicLoad(u32, &p.bytes[x % JOBS], .acquire); + } + + /// Land a host-written peer part for exchange x (the relay's post-copy half): copy it in, then release the wait. + pub fn arrive(p: *Loopback, x: u32, bytes: []const u8) void { + std.debug.assert(bytes.len <= p.l.slot); + @memcpy(p.win[p.l.recv(x)..][0..bytes.len], bytes); + p.landed.set(x); + } + + /// The seam `Pair.stage` encodes (before its `gather`): this rank's `width` columns of the frame row into the slot. + pub fn stage(p: *Loopback, e: *mtl.ComputeEncoder, x: u32, src: Ref, rows: u32, width: u32, stride: u32) void { + if (rows == 0 or width == 0 or width > stride) @panic("loopback: stage() wants rows > 0 and a width inside the source row"); + if (stageBytes(rows, width) > p.l.slot) @panic("loopback: this exchange's part does not fit the window's slot"); + e.setPipeline(p.stage_pipe); + e.setBuffer(p.wbuf, p.l.sync(), 0); + e.setValue(PairRows{ .width = width, .stride = stride, .rank = p.rank, .x = x }, 1); + e.setBuffer(src.buf, src.off, 2); + e.setBuffer(p.wbuf, p.l.send(x), 3); + e.dispatchGroups(.{ .width = rows }, .{ .width = 256 }); + } + + /// The seam `Pair.gather` encodes: the relay, then the frame row as this rank's half and the peer's, both bf16. + pub fn gather(p: *Loopback, e: *mtl.ComputeEncoder, x: u32, mine: Ref, out: Ref, rows: u32, width: u32, out_stride: u32) void { + if (rows == 0 or width == 0 or width % elems_per_word != 0) @panic("loopback: gather() wants rows > 0 and a width that fills whole window words"); + if (out_stride < 2 * width) @panic("loopback: gather()'s frame row is narrower than both halves"); + if (x != @as(u32, @truncate(p.x + 1))) @panic("loopback: an exchange was encoded out of order (part() not called first)"); + if (mine.buf.id != p.wbuf.id or mine.off != p.l.send(x)) @panic("loopback: gather()'s part is not the region part() returned"); + if (stageBytes(rows, width) > p.l.slot) @panic("loopback: this exchange's part does not fit the window's slot"); + p.x += 1; + @atomicStore(u32, &p.bytes[x % JOBS], @intCast(stageBytes(rows, width)), .release); + seam(e, p.fence, p.posted, p.landed, x); + e.setPipeline(p.gather_pipe); + e.setBuffer(p.wbuf, p.l.sync(), 0); + e.setValue(PairRows{ .width = width, .stride = out_stride, .rank = p.rank, .x = x }, 1); + e.setBuffer(mine.buf, mine.off, 2); + e.setBuffer(p.wbuf, p.l.recv(x), 3); + e.setBuffer(out.buf, out.off, 4); + e.dispatchGroups(.{ .width = rows }, .{ .width = 256 }); + } + + /// The loopback's head seam: this rank's argmax over n bf16 logits, the relay copy, then the pick over both parts. + pub fn argmaxBf16(p: *Loopback, e: *mtl.ComputeEncoder, x: u32, logits: Ref, stride: u32, lo: u32, n: u32, rows: u32, picks: Ref) void { + p.argmaxWidth(e, x, logits, stride, lo, n, rows, picks, 1); + } + + /// The same seam over fp32 logits (GLM-5.3's frame): the loopback covers both widths the kernel reads. + pub fn argmax(p: *Loopback, e: *mtl.ComputeEncoder, x: u32, logits: Ref, stride: u32, lo: u32, n: u32, rows: u32, picks: Ref) void { + p.argmaxWidth(e, x, logits, stride, lo, n, rows, picks, 0); + } + + fn argmaxWidth(p: *Loopback, e: *mtl.ComputeEncoder, x: u32, logits: Ref, stride: u32, lo: u32, n: u32, rows: u32, picks: Ref, dtype: u32) void { + if (rows == 0 or n == 0) @panic("loopback: argmax() wants rows > 0 and n > 0"); + if (x != @as(u32, @truncate(p.x + 1))) @panic("loopback: an exchange was encoded out of order (part() not called first)"); + p.x += 1; + @atomicStore(u32, &p.bytes[x % JOBS], @intCast(pickBytes(rows)), .release); // two u32 a row: (value, id) + e.setPipeline(p.argmax_pipe); + e.setBuffer(logits.buf, logits.off, 0); + e.setValue([4]u32{ stride, lo, n, dtype }, 1); + e.setBuffer(p.wbuf, p.l.send(x), 2); + e.dispatchGroups(.{ .width = rows }, .{ .width = 1024 }); + seam(e, p.fence, p.posted, p.landed, x); // the encoder boundary orders the argmax before the part is relayed + e.setPipeline(p.pick_pipe); + e.setBuffer(p.wbuf, p.l.sync(), 0); + e.setValue([2]u32{ rows, x }, 1); + e.setBuffer(p.wbuf, p.l.send(x), 2); + e.setBuffer(p.wbuf, p.l.recv(x), 3); + e.setBuffer(picks.buf, picks.off, 4); + e.dispatchGroups(.{ .width = 1 }, .{ .width = 32 }); + } + + /// Seams the GPU refused (must be 0 for a run to be valid). + pub fn gaveUp(p: *const Loopback) u32 { + return @atomicLoad(u32, p.word32(p.l.sync() + FAILED * 4), .acquire); + } + + /// There is no link to fail here, so this is `gaveUp` under the seam's name, and it reads 0. + pub fn failed(p: *const Loopback) bool { + return p.gaveUp() != 0; + } +}; diff --git a/zig/src/core/simd_qmm.zig b/zig/src/core/simd_qmm.zig index 844c2099e5..1856304072 100644 --- a/zig/src/core/simd_qmm.zig +++ b/zig/src/core/simd_qmm.zig @@ -8,6 +8,11 @@ pub fn splits(n: usize) u32 { return if (n <= 64) 32 else if (n <= 6144) 16 else 8; } +/// Chunks for `n`; a split rank passes the whole width as `whole` so it keeps that width's tree (0 keeps `n`). +pub fn chunks(n: usize, whole: usize) u32 { + return splits(if (whole != 0) whole else n); +} + /// Output tiles a simdgroup within the 16 KB reduction buffer; tiles never change a row's bits. pub fn tiles(n: usize, rt: u32, s: u32) u32 { var nt: u32 = if (n % 32 == 0) 4 else if (n % 16 == 0) 2 else 1; @@ -72,10 +77,10 @@ pub const Pipelines = struct { pub const Call = struct { pipeline: mtl.Pipeline, args: Args, groups: mtl.Size, threads: usize }; /// y = x W^T, 4-bit groups of 64: the scalar twin at 1-2 rows if `scalar`, chunk pairs if `paired`, else MMA tiles. -pub fn call(p: *const Pipelines, n: usize, k: usize, rows: usize, scalar: bool, paired: bool) !Call { +pub fn call(p: *const Pipelines, n: usize, k: usize, rows: usize, scalar: bool, paired: bool, whole: usize) !Call { if (n == 0 or n % 8 != 0 or n > std.math.maxInt(i32) or k == 0 or k % 64 != 0 or k > std.math.maxInt(i32)) return error.UnsupportedSimdShape; if (rows == 0 or rows > std.math.maxInt(i32)) return error.NoRows; - const s = splits(n); + const s = chunks(n, whole); const si = std.mem.indexOfScalar(u32, &s_values, s).?; const args = Args{ .k = @intCast(k), .n = @intCast(n), .r = @intCast(rows) }; if (scalar and scalarBlock(rows, s) != 0) { @@ -127,7 +132,7 @@ pub fn pairedExact(p: *const Pipelines, device: mtl.Device, m: Matrix) !bool { defer queue.deinit(); const cb = queue.commandBuffer(); const e = cb.compute(.serial); - for (outs, [_]bool{ false, true }) |o, pair| encode(e, try call(p, m.n, m.k, rows, false, pair), m, x, 0, o, 0); + for (outs, [_]bool{ false, true }) |o, pair| encode(e, try call(p, m.n, m.k, rows, false, pair, 0), m, x, 0, o, 0); e.end(); cb.commit(); cb.wait(); @@ -157,9 +162,9 @@ pub fn twinExact(p: *const Pipelines, device: mtl.Device, m: Matrix) !bool { defer queue.deinit(); const cb = queue.commandBuffer(); const e = cb.compute(.serial); - encode(e, try call(p, m.n, m.k, rows, false, false), m, x, 0, tiled, 0); - for (0..rows) |r| encode(e, try call(p, m.n, m.k, 1, true, false), m, x, r * m.k * 2, twin, r * m.n * 2); - for ([_]usize{ 0, rows - 2 }, 0..) |r, i| encode(e, try call(p, m.n, m.k, 2, true, false), m, x, r * m.k * 2, twos, i * 2 * m.n * 2); + encode(e, try call(p, m.n, m.k, rows, false, false, 0), m, x, 0, tiled, 0); + for (0..rows) |r| encode(e, try call(p, m.n, m.k, 1, true, false, 0), m, x, r * m.k * 2, twin, r * m.n * 2); + for ([_]usize{ 0, rows - 2 }, 0..) |r, i| encode(e, try call(p, m.n, m.k, 2, true, false, 0), m, x, r * m.k * 2, twos, i * 2 * m.n * 2); e.end(); cb.commit(); cb.wait(); @@ -177,6 +182,8 @@ test "chunks follow output width only and tiles fit the reduction buffer" { try std.testing.expectEqual(@as(u32, 32), splits(48)); try std.testing.expectEqual(@as(u32, 16), splits(5120)); try std.testing.expectEqual(@as(u32, 8), splits(34816)); + try std.testing.expectEqual(splits(6240), chunks(3168, 6240)); + try std.testing.expectEqual(splits(6240), chunks(6240, 0)); for (s_values) |s| for ([_]u32{ 1, 2 }) |rt| { const nt = tiles(34816, rt, s); try std.testing.expect(s * rt * nt * 64 * 4 <= 16384); diff --git a/zig/src/families/qwen27/dflash/runtime_callbacks.zig b/zig/src/families/qwen27/dflash/runtime_callbacks.zig index 46d252448c..bfea3d8552 100644 --- a/zig/src/families/qwen27/dflash/runtime_callbacks.zig +++ b/zig/src/families/qwen27/dflash/runtime_callbacks.zig @@ -29,7 +29,7 @@ fn project(ptr: *anyopaque, epoch: u64, x: op.View, w: ck.Tensor, mode: op.Mode) const input = try b.frame.ref(x); const output = try b.frame.ref(y); const s = @import("core").row_projection.simd; - s.encode(e, try s.call(&b.simd.?, t.n, t.k, x.rows, false, false), t.matrix(), input.buf, input.off, output.buf, output.off); + s.encode(e, try s.call(&b.simd.?, t.n, t.k, x.rows, false, false, 0), t.matrix(), input.buf, input.off, output.buf, output.off); } else { const adapter = b.prepared.get(@intFromPtr(w.bytes.ptr)) orelse return error.DraftPreparationUnqualified; const input = try b.frame.ref(x); diff --git a/zig/src/families/qwen27/quant_gpu.zig b/zig/src/families/qwen27/quant_gpu.zig index e38b913e2a..728cc3955f 100644 --- a/zig/src/families/qwen27/quant_gpu.zig +++ b/zig/src/families/qwen27/quant_gpu.zig @@ -139,7 +139,7 @@ pub const Kernels = struct { if (p.raw) |raw| if (raw.bits == 4 and raw.group == 64 and raw.sum == .f32 and raw.n % 8 == 0) { const m = row.simd.Matrix{ .w = raw.w, .w_off = raw.w_off, .scales = raw.scales, .s_off = raw.s_off, .biases = raw.biases, .b_off = raw.b_off, .n = raw.n, .k = raw.k }; const scalar = try k.twin(m); // loads the tiles before they're read - const call = try row.simd.call(&k.simd.?, raw.n, raw.k, rows, scalar, try k.pairTwin(m)); + const call = try row.simd.call(&k.simd.?, raw.n, raw.k, rows, scalar, try k.pairTwin(m), 0); row.simd.encode(e, call, m, x.buffer, x.offset, y.buffer, y.offset); e.barrier(); try span.end(); diff --git a/zig/src/metal/command.zig b/zig/src/metal/command.zig index 49118d0b3d..45c2b29358 100644 --- a/zig/src/metal/command.zig +++ b/zig/src/metal/command.zig @@ -195,4 +195,14 @@ pub const ComputeEncoder = struct { pub fn end(self: ComputeEncoder) void { objc.msg(void, self.id, "endEncoding", .{}); } + + /// The command buffer this encoder was opened from, for the event ops a seam encodes BETWEEN encoders. + pub fn commandBuffer(self: ComputeEncoder) CommandBuffer { + return .{ .id = objc.msg(Id, self.id, "commandBuffer", .{}) }; + } + + /// Whether this encoder dispatches serially or concurrently: a seam reopens it in the same mode. + pub fn dispatchType(self: ComputeEncoder) types.DispatchType { + return @fromBackingInt(@intCast(objc.msg(usize, self.id, "dispatchType", .{}))); + } }; diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 4fa4fbe147..4cc2deba22 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -32,6 +32,8 @@ pub const exlr8 = @import("core/exlr8.zig"); pub const exlr8_projection = @import("core/exlr8_projection.zig"); pub const exlr8_dense_gpu = @import("core/exlr8_dense_gpu.zig"); pub const exlr8_prompt = @import("core/exlr8_prompt.zig"); +/// The two-Mac pair seam's row split: the transport, its layout and the one-Mac Loopback the stage-2 checks use. +pub const pair_rows = @import("core/pair_rows.zig"); pub const glm_dsa = @import("families/glm_dsa/glm_dsa.zig"); pub const nemotron = @import("families/nemotron/nemotron.zig"); pub const flashnext_vision_contract = @import("families/flashnext/vision_contract.zig"); diff --git a/zig/tests/affine_ref.zig b/zig/tests/affine_ref.zig index b367fa120a..691ecf11ce 100644 --- a/zig/tests/affine_ref.zig +++ b/zig/tests/affine_ref.zig @@ -272,7 +272,7 @@ fn affine(t: *Tally, device: mtl.Device, queue: mtl.Queue, s: *const Synth, ref: /// simd_qmm over rows [start, start + m), the scalar twin when `scalar` and it takes m rows. fn simdWindow(e: mtl.ComputeEncoder, p: *const sq.Pipelines, s: *const Synth, y: Out, scalar: bool, start: usize, m: usize) !void { const mat = sq.Matrix{ .w = s.w, .scales = s.s, .biases = s.b, .n = s.n, .k = s.k }; - sq.encode(e, try sq.call(p, s.n, s.k, m, scalar, true), mat, s.x, start * s.k * 2, y.buf, start * s.n * 2); + sq.encode(e, try sq.call(p, s.n, s.k, m, scalar, true, 0), mat, s.x, start * s.k * 2, y.buf, start * s.n * 2); } /// simd_qmm (4-bit, group 64): tiles at every window against float64 and the wide call; the twin at 1 and 2 rows. diff --git a/zig/tests/contract/metal_affine.zig b/zig/tests/contract/metal_affine.zig index 46f921e9f8..3aa37d5de0 100644 --- a/zig/tests/contract/metal_affine.zig +++ b/zig/tests/contract/metal_affine.zig @@ -198,7 +198,7 @@ pub const Proj = struct { e.dispatchGroups(mtl.Size.of(cl.groups, cl.slots, 1), mtl.Size.of(cl.threads, 1, 1)); }, .simd => { - const cl = try sq.call(&p.sq_p.?, shape.n, shape.k, shape.rows, false, true); + const cl = try sq.call(&p.sq_p.?, shape.n, shape.k, shape.rows, false, true, 0); sq.encode(e, cl, .{ .w = p.w.buf, .w_off = p.w.off, .scales = p.s.buf, .s_off = p.s.off, .biases = p.b.buf, .b_off = p.b.off, .n = shape.n, .k = shape.k }, p.x.buf, p.x.off, p.out.buf, p.out.off); }, .simd_bits => { diff --git a/zig/tests/contract/metal_limits.zig b/zig/tests/contract/metal_limits.zig index 53609bffa0..c0e267c0d0 100644 --- a/zig/tests/contract/metal_limits.zig +++ b/zig/tests/contract/metal_limits.zig @@ -100,6 +100,8 @@ fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch) !La text = sources.core_deltanet; } else if (std.mem.startsWith(u8, d.name, "qwen27 ")) { text = if (has(d.name, " copy ")) tf.qwen27.runtime_sources.copy else tf.qwen27.runtime_sources.glue; + } else if (std.mem.startsWith(u8, d.name, "metal pair ")) { + text = sources.core_pair_rows; } else if (has(d.name, "tile_qmm")) text = sources.core_tile_qmm else if (has(d.name, "softmax")) text = sources.ops_softmax else if (has(d.name, "rms")) text = sources.ops_embed_norm else if (has(d.name, "top-k")) text = topk_text else if (has(d.name, "row projection")) text = sources.core_row_projection else if (has(d.name, "simd_qmm_bits")) text = sources.core_simd_qmm_bits else if (has(d.name, "simd_qmm")) text = sources.core_simd_qmm else if (has(d.name, "affine_mm")) { const bits: u8 = if (has(d.name, "q8")) 8 else 4; const group: u16 = if (has(d.name, "g32")) 32 else 64; diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 46801e171c..7fd93c2ff8 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -17,6 +17,8 @@ const drama_combine = @import("metal_drama_combine.zig"); const drama_cfg = @import("metal_drama_cfg.zig"); const glm_eh_k = @import("metal_glm_eh.zig"); const glm_rows_k = @import("metal_glm_rows.zig"); +const pair_k = @import("metal_pair.zig"); +const pair_controls = @import("metal_pair_controls.zig"); const flash_pool = @import("metal_flashnext_pool.zig"); const flash = @import("metal_flashnext.zig"); @@ -202,6 +204,32 @@ pub fn main(init: std.process.Init) !u8 { try rows_list.append(a, r); try kernels.append(a, r.kernel()); } + var pair_list: std.ArrayList(*pair_k.Pair) = .empty; + defer { + for (pair_list.items) |p| { + p.deinit(); + a.destroy(p); + } + pair_list.deinit(a); + } + for (std.enums.values(pair_controls.Mutation)) |m| { + for (pair_k.all) |w| { + const control_for = switch (m) { + .none => true, + .stage_peer_half => w == .stage, + .gather_own_halves => w == .gather, + .argmax_high_tie, .argmax_no_flag => w == .argmax, + .pick_high_tie => w == .pick, + }; + if (!control_for) continue; + if (!wanted(o, if (m == .none) pair_k.spec(w).name else m.name())) continue; + const p = try a.create(pair_k.Pair); + errdefer a.destroy(p); + p.* = try pair_k.Pair.init(a, device, queue, w, m); + try pair_list.append(a, p); + try kernels.append(a, p.kernel()); + } + } var flash_list: [8]flash.Probe = undefined; var flash_made: usize = 0; defer for (flash_list[0..flash_made]) |*probe| probe.deinit(); diff --git a/zig/tests/contract/metal_pair.zig b/zig/tests/contract/metal_pair.zig new file mode 100644 index 0000000000..9a15b020f8 --- /dev/null +++ b/zig/tests/contract/metal_pair.zig @@ -0,0 +1,485 @@ +//! The pair seam's four kernels (core/pair_rows.metal), and the float64 rules the gates check. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); +const rig_ = @import("rig.zig"); +const controls = @import("metal_pair_controls.zig"); + +const c = contract.spec; +const g = contract.gates; + +/// The seam's four kernels; the tail of each registry name is the MSL symbol launchOf hands the source. +pub const Which = enum { stage, gather, argmax, pick }; + +pub const all = [_]Which{ .stage, .gather, .argmax, .pick }; + +/// What pair_rows.zig dispatches per kernel, so `init` can compare it against Metal's own reflection. +pub const Launched = struct { threads: u32, tg_memory: u32 }; + +pub fn launch(w: Which) Launched { + return switch (w) { + .stage => .{ .threads = 256, .tg_memory = 0 }, + .gather => .{ .threads = 256, .tg_memory = 0 }, + .argmax => .{ .threads = 1024, .tg_memory = 384 }, + .pick => .{ .threads = 32, .tg_memory = 0 }, + }; +} + +/// The MSL symbol pair_rows.metal defines. +pub fn symbol(w: Which) []const u8 { + return switch (w) { + .stage => "pair_stage", + .gather => "pair_gather", + .argmax => "pair_argmax", + .pick => "pair_pick", + }; +} + +/// The spec's name: `metal `, so launchOf's last word is the symbol and its prefix picks core_pair_rows. +pub fn name(w: Which) []const u8 { + return switch (w) { + .stage => "metal pair stage pair_stage", + .gather => "metal pair gather pair_gather", + .argmax => "metal pair argmax pair_argmax", + .pick => "metal pair pick pair_pick", + }; +} + +/// Qwen3.8-27B's rank-local slice: the seam reduces ids [124160, 248320) on rank 1. +const logit_lo: u32 = 124160; + +/// 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. +const flag: u32 = 0x80000000; + +/// The id a flagged row, and a broken link, read as: no token, visibly wrong. +const sentinel: u32 = 0xFFFFFFFF; + +/// A thread's chain over every 1024th logit, five shuffle levels, then the 32 simdgroup winners. +fn chainLogits(s: c.Shape) usize { + return (s.k + 1023) / 1024 + 5 + 5; +} + +/// Only moves and compares values: no sums to carry an accumulator's error. +fn noSums(_: c.Shape) usize { + return 0; +} + +const move_why = "a move writes each bf16 element as it read it: no sum, no rounding, so no accumulator's error can enter"; + +const argmax_rule = "a winner is (value bits, id): NaN loses every comparison, a row of only NaNs reports the -inf sentinel with id 0xFFFFFFFF, and a row with any non-finite logit reports NaN bits and id | 0x80000000"; +const pick_rule = "the pick copies one rank's (value, id) word pair; an id with bit 31 set, and the broken-link word, draw 0xFFFFFFFF: the same sentinel, so one downstream check covers both"; + +/// A kernel's spec. The raw kernels check themselves; the two moves' bound is exact, so gate 3 keeps every edge. +pub fn spec(w: Which) c.Spec { + return switch (w) { + // One spec for both moves; a gather's `n` is its frame row stride, and the 17,408 shape is seam B's. + .stage, .gather => .{ + .name = name(w), + .accum = .none, + .depth = noSums, + .out = .bf16, + .tiles = .{ .{}, .{}, .{ .size = 256, .step = 256, .min = 256 } }, + .served = if (w == .stage) &.{.{ .rows = 8, .n = 5120, .k = 5120 }} else &.{ + .{ .rows = 8, .n = 10240, .k = 5120 }, .{ .rows = 8, .n = 17408, .k = 8704 }, + .{ .rows = 8, .n = 10496, .k = 5120 }, + }, + .sums = .{ .f16 = 0, .bf16 = 0, .exact = false, .why = move_why }, + }, + .argmax => .{ + .name = name(w), + .accum = .none, + .depth = chainLogits, + .out = .raw, + .nan = .{ .stated = argmax_rule }, + .tiles = .{ .{}, .{}, .{ .size = 2048, .step = 64, .min = 64 } }, + .served = &.{.{ .rows = 8, .n = 1, .k = 124160 }}, + }, + .pick => .{ + .name = name(w), + .accum = .none, + .depth = noSums, + .out = .raw, + .nan = .{ .stated = pick_rule }, + .served = &.{ .{ .rows = 32, .n = 1, .k = 1 }, .{ .rows = 128, .n = 1, .k = 1 } }, + }, + }; +} + +/// `constant PairRows&` at buffer 1: `width`, the frame row `stride` and which rank's half this side writes. +const PairRows = extern struct { + width: u32, + stride: u32, + rank: u32, + x: u32, +}; + +fn rowsArgs(width: usize, stride: usize, rank: u32, x: u32) PairRows { + return .{ .width = @intCast(width), .stride = @intCast(stride), .rank = rank, .x = x }; +} + +fn bit(v: f32) u32 { + return @bitCast(v); +} + +fn floatOf(v: u32) f32 { + return @bitCast(v); +} + +/// One prepared call: the buffers the kernel is bound to, the host values its reference reads, and its own check. +pub const Pair = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + which: Which, + mutation: controls.Mutation, + pipe: mtl.Pipeline, + /// The hooks this probe exposes: the raw kernels check their own rule, the two computing ones use the reference. + v: g.VTable, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + rows: usize = 0, + width: usize = 0, + /// The rank this prepare writes first: the gather's `mine` half is its payload. + rank: u32 = 0, + /// The row stride in PairRows, in bf16 elements: the stage's source frame row, the gather's destination row. + stride: usize = 0, + + // Host inputs of the last prepare, as the reference and the check read them (freed by `drop`). + sh: ?[]u16 = null, // stage: the src payload, bf16 (rows, each a frame row of `stride`) + mh: ?[]u16 = null, // gather: this rank's half (mine), bf16 + nh: ?[]u16 = null, // gather: the peer's half (recv), bf16 + kh: ?[]u16 = null, // argmax: the bf16 logits + pw: ?[]u32 = null, // pick: this rank's (value bits, id) words + qw: ?[]u32 = null, // pick: the peer's words + + msg: [192]u8 = undefined, + sync: rig_.Ref = undefined, + mine: rig_.Ref = undefined, + peer: rig_.Ref = undefined, + logits: rig_.Ref = undefined, + out: rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, which: Which, mutation: controls.Mutation) !Pair { + const src = if (mutation == .none) sources.core_pair_rows else try controls.source(a, mutation); + defer if (mutation != .none) a.free(src); + const pipe = try rig_.pipeline(device, src, symbol(which)); + errdefer pipe.deinit(); + const want = launch(which); + const tg = rig_.tgMemory(pipe); + if (tg != want.tg_memory) { + std.debug.print("metal pair: {s} takes {d} bytes of threadgroup memory, its record says {d}\n", .{ symbol(which), tg, want.tg_memory }); + return error.ThreadgroupMemoryMismatch; + } + var v = g.hooks(Pair); + if (which != .argmax and which != .pick) v.check = null; + if (which != .stage and which != .gather) v.bounds = null; + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .which = which, .mutation = mutation, .pipe = pipe, .v = v }; + } + + pub fn deinit(self: *Pair) void { + self.drop(); + self.rig.deinit(); + self.pipe.deinit(); + } + + pub fn kernel(self: *Pair) g.Kernel { + var s = spec(self.which); + if (self.mutation != .none) { + s.name = self.mutation.name(); + return .{ .spec = s, .ctx = self, .vt = &self.v, .control = self.mutation.gate() }; + } + return .{ .spec = s, .ctx = self, .vt = &self.v }; + } + + fn drop(self: *Pair) void { + if (self.sh) |s| self.a.free(s); + if (self.mh) |s| self.a.free(s); + if (self.nh) |s| self.a.free(s); + if (self.kh) |s| self.a.free(s); + if (self.pw) |s| self.a.free(s); + if (self.qw) |s| self.a.free(s); + self.sh = null; + self.mh = null; + self.nh = null; + self.kh = null; + self.pw = null; + self.qw = null; + } + + /// Ordinary values: what the seam carries in production, symmetric about zero. + fn ordinary(rng: std.Random) f32 { + return rng.float(f32) * 2 - 1; + } + + /// A bf16 payload: what the two moves carry, with gate 1's non-finite plants where the reference looks for them. + fn payload(self: *Pair, len: usize, fill: g.Fill, rng: std.Random) ![]u16 { + const s = try self.a.alloc(u16, len); + for (s) |*v| v.* = c.toBf16(ordinary(rng)); + if (fill == .nonfinite) { + const nan = c.toBf16(std.math.nan(f32)); + const inf = c.toBf16(std.math.inf(f32)); + const ninf = c.toBf16(-std.math.inf(f32)); + for ([_]u16{ nan, inf, ninf, nan, ninf }, 0..) |plant, p| { + const at = (p + 1) * (len / 8 + 1) + p * 3; + if (at < len) s[at] = plant; + } + } + return s; + } + + /// Inputs for `shape` as `fill` asks, outputs poisoned and guarded; false when the kernel does not serve it. + pub fn prepare(self: *Pair, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.rows > 512) return false; + self.drop(); + self.rig.reset(); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + switch (self.which) { + .stage, .gather => { + if (shape.k == 0 or shape.k > 32768 or shape.k % 256 != 0) return false; + // The gather writes this rank's half first; alternating the rank with the seed covers both orders. + self.rank = @intCast(seed & 1); + self.sync = try self.rig.scratch(128); + if (self.which == .stage) { + // The source is a frame row: the stage reads this rank's `width` columns out of its `stride`. + self.stride = 2 * shape.k; + const len = shape.rows * self.stride; + const mine = try self.payload(len, fill, rng); + self.sh = mine; + self.mine = try self.rig.input(std.mem.sliceAsBytes(mine)); + self.out = try self.rig.output(shape.rows * shape.k * 2); + } else { + const n = shape.rows * shape.k; + // A frame row holds both halves, so its stride is at least 2*width: `n` sets a wider one. + self.stride = if (shape.n > 2 * shape.k) shape.n else 2 * shape.k; + const mine = try self.payload(n, fill, rng); + const peer = try self.payload(n, fill, rng); + self.mh = mine; + self.nh = peer; + self.mine = try self.rig.input(std.mem.sliceAsBytes(mine)); + self.peer = try self.rig.input(std.mem.sliceAsBytes(peer)); + self.out = try self.rig.output(2 * ((shape.rows - 1) * self.stride + 2 * shape.k)); + } + }, + .argmax => { + if (shape.k == 0 or shape.k > 262144) return false; + const n = shape.rows * shape.k; + const k = try self.a.alloc(u16, n); + self.kh = k; + for (k) |*v| v.* = c.toBf16(ordinary(rng)); + // Two in-range ids at +2.0: the lower id sits in the first half, so the tie rule's pick is fixed. + for (0..shape.rows) |r| { + const row = k[r * shape.k ..][0..shape.k]; + row[0] = c.toBf16(2.0); + row[shape.k / 2] = c.toBf16(2.0); + } + if (fill == .nonfinite) { + const nan = c.toBf16(std.math.nan(f32)); + const inf = c.toBf16(std.math.inf(f32)); + const ninf = c.toBf16(-std.math.inf(f32)); + for (0..shape.rows) |r| { + const row = k[r * shape.k ..][0..shape.k]; + row[shape.k / 3] = switch (r % 4) { + 0 => nan, + 1 => inf, + 2 => ninf, + else => row[shape.k / 3], + }; + if (r + 1 == shape.rows) for (row) |*v| { + v.* = nan; + }; + } + } + self.logits = try self.rig.input(std.mem.sliceAsBytes(k)); + self.out = try self.rig.output(shape.rows * 8); + }, + .pick => { + if (shape.n != 1 or shape.k != 1) return false; + self.pw = try self.a.alloc(u32, 2 * shape.rows); + self.qw = try self.a.alloc(u32, 2 * shape.rows); + for (0..shape.rows) |r| { + self.pw.?[2 * r] = bit(ordinary(rng)); + self.qw.?[2 * r] = bit(ordinary(rng)); + self.pw.?[2 * r + 1] = rng.uintLessThan(u32, 1 << 20); + self.qw.?[2 * r + 1] = rng.uintLessThan(u32, 1 << 20); + } + // Equal values on both ranks: the tie rule decides, and its direction alternates row by row. + for (0..shape.rows) |r| { + if (r % 4 != 2 and r % 4 != 3) continue; + self.qw.?[2 * r] = self.pw.?[2 * r]; + self.pw.?[2 * r + 1] = 300 + @as(u32, @intCast(r)); + self.qw.?[2 * r + 1] = if (r % 8 == 2) 700 + @as(u32, @intCast(r)) else 100 + @as(u32, @intCast(r)); + } + if (fill == .nonfinite) { + // Two plants prove the id flag is the rule: a flagged row draws the sentinel, a plain one picks. + for (0..shape.rows) |r| { + switch (r % 4) { + 0 => self.pw.?[2 * r + 1] |= flag, + 1 => self.qw.?[2 * r + 1] |= flag, + 2 => self.pw.?[2 * r] = bit(-std.math.inf(f32)), + else => self.qw.?[2 * r] = bit(std.math.inf(f32)), + } + } + } + self.sync = try self.rig.scratch(128); + self.mine = try self.rig.input(std.mem.sliceAsBytes(self.pw.?)); + self.peer = try self.rig.input(std.mem.sliceAsBytes(self.qw.?)); + self.out = try self.rig.output(shape.rows * 4); + }, + } + self.shape = shape; + self.rows = shape.rows; + self.width = shape.k; + return true; + } + + /// One launch on the prepared buffers, waited for. + pub fn run(self: *Pair) !void { + const p = self.rig.begin(); + const e = p.e; + const groups = mtl.Size.of(self.rows, 1, 1); + const threads = mtl.Size.of(launch(self.which).threads, 1, 1); + switch (self.which) { + .stage => { + @memset(self.rig.bytes(self.sync, 128), 0); + e.setPipeline(self.pipe); + e.setBuffer(self.sync.buf, self.sync.off, 0); + e.setValue(rowsArgs(self.width, self.stride, self.rank, 1), 1); + e.setBuffer(self.mine.buf, self.mine.off, 2); + e.setBuffer(self.out.buf, self.out.off, 3); + e.dispatchGroups(groups, threads); + }, + .gather => { + @memset(self.rig.bytes(self.sync, 128), 0); + const frame = self.rig.bytes(self.out, 2 * self.count()); + @memset(frame, 0); // the frame's gaps: no launch writes them, so their planted value is zero + e.setPipeline(self.pipe); + e.setBuffer(self.sync.buf, self.sync.off, 0); + e.setValue(rowsArgs(self.width, self.stride, self.rank, 1), 1); + e.setBuffer(self.mine.buf, self.mine.off, 2); + e.setBuffer(self.peer.buf, self.peer.off, 3); + e.setBuffer(self.out.buf, self.out.off, 4); + e.dispatchGroups(groups, threads); + }, + .argmax => { + e.setPipeline(self.pipe); + e.setBuffer(self.logits.buf, self.logits.off, 0); + const dims = [4]u32{ @intCast(self.width), logit_lo, @intCast(self.width), 1 }; + e.setValue(dims, 1); + e.setBuffer(self.out.buf, self.out.off, 2); + e.dispatchGroups(groups, threads); + }, + .pick => { + @memset(self.rig.bytes(self.sync, 128), 0); + e.setPipeline(self.pipe); + e.setBuffer(self.sync.buf, self.sync.off, 0); + e.setValue([2]u32{ @intCast(self.rows), 1 }, 1); + e.setBuffer(self.mine.buf, self.mine.off, 2); + e.setBuffer(self.peer.buf, self.peer.off, 3); + e.setBuffer(self.out.buf, self.out.off, 4); + e.dispatchGroups(mtl.Size.of(1, 1, 1), threads); + }, + } + try self.rig.finish(p); + } + + /// Output elements of the prepared call: bf16 elements for the two moves, the two raw kernels' bytes. + pub fn count(self: *Pair) usize { + return switch (self.which) { + .stage => self.rows * self.width, + .gather => (self.rows - 1) * self.stride + 2 * self.width, + .argmax => self.rows * 8, + .pick => self.rows * 4, + }; + } + + /// Float64 outputs of the prepared inputs and each one's sums' term magnitude (a move's term is the value itself). + pub fn reference(self: *Pair, want: []f64, mag: []f64) !void { + const width = self.width; + switch (self.which) { + .stage => { + const sh = self.sh orelse return error.NoReference; + for (0..self.rows) |r| { + const row = sh[r * self.stride ..][0..width]; + for (0..width) |i| { + const v = c.fromBf16(row[i]); + want[r * width + i] = v; + mag[r * width + i] = @abs(v); + } + } + }, + .gather => { + const mh = self.mh orelse return error.NoReference; + const nh = self.nh orelse return error.NoReference; + @memset(want, 0); // the frame's gaps: no launch writes them, so their reference value is zero + for (0..self.rows) |row| { + const at = row * self.stride; + for (0..width) |i| { + want[at + i] = c.fromBf16(if (self.rank == 0) mh[row * width + i] else nh[row * width + i]); + want[at + width + i] = c.fromBf16(if (self.rank == 0) nh[row * width + i] else mh[row * width + i]); + } + } + for (want, mag) |w, *m| m.* = @abs(w); + }, + else => return error.NoReference, + } + } + + /// The moves are exact: every output element must equal the reference's bf16 bit for bit. + pub fn bounds(self: *Pair, want: []const f64, mag: []const f64, allow: []f64) !void { + _ = want; + _ = mag; + if (self.which != .stage and self.which != .gather) return error.NotMove; + for (allow) |*l| l.* = 0; + } + + /// The seam's own checks: the raw kernels' rule at gate 1 (non-finite inputs) and gate 3 (every served edge). + pub fn check(self: *Pair, _: g.Gate) !?[]const u8 { + const out = try self.rig.outputBytes(); + switch (self.which) { + .argmax => { + const k = self.kh orelse return error.NoReference; + const n = self.width; + for (0..self.rows) |r| { + var best: f32 = -std.math.inf(f32); + var at: u32 = sentinel; + var bad = false; + for (0..n) |i| { + const v = c.fromBf16(k[r * n + i]); + if (!std.math.isFinite(v)) { + bad = true; + continue; + } + if (v > best) { + best = v; + at = logit_lo + @as(u32, @intCast(i)); + } + } + const got = std.mem.readInt(u32, out[r * 8 ..][0..4], .little); + const id = std.mem.readInt(u32, out[r * 8 + 4 ..][0..4], .little); + if (bad) { + if (!std.math.isNan(floatOf(got))) return try std.fmt.bufPrint(&self.msg, "row {d}: value {e} where a non-finite logit must write NaN", .{ r, floatOf(got) }); + if (id != at | flag) return try std.fmt.bufPrint(&self.msg, "row {d}: id {d} where flagging gives {d}", .{ r, id, at | flag }); + } else if (got != bit(best) or id != at) { + return try std.fmt.bufPrint(&self.msg, "row {d}: ({e} {d}) where its own rule gives ({e} {d})", .{ r, floatOf(got), id, best, at }); + } + } + }, + .pick => { + const pw = self.pw orelse return error.NoReference; + const qw = self.qw orelse return error.NoReference; + for (0..self.rows) |r| { + const vm = floatOf(pw[2 * r]); + const am = pw[2 * r + 1]; + const vp = floatOf(qw[2 * r]); + const ap = qw[2 * r + 1]; + const flagged = ((am | ap) & flag) != 0; + const want = if (flagged) sentinel else if (vp > vm or (vp == vm and ap < am)) ap else am; + const got = std.mem.readInt(u32, out[r * 4 ..][0..4], .little); + if (got != want) return try std.fmt.bufPrint(&self.msg, "row {d} picks {d} where its own rule gives {d}", .{ r, got, want }); + } + }, + else => return error.NotRaw, + } + return null; + } +}; diff --git a/zig/tests/contract/metal_pair_controls.zig b/zig/tests/contract/metal_pair_controls.zig new file mode 100644 index 0000000000..64c2bb91a8 --- /dev/null +++ b/zig/tests/contract/metal_pair_controls.zig @@ -0,0 +1,55 @@ +//! Planted controls for the pair seam's four kernels: one fault in pair_rows.metal's text, and the gate to catch it. +const std = @import("std"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); + +pub const Mutation = enum { + none, + stage_peer_half, + gather_own_halves, + argmax_high_tie, + argmax_no_flag, + pick_high_tie, + + pub fn name(m: Mutation) []const u8 { + return switch (m) { + .none => "", + .stage_peer_half => "control: metal pair stage, this rank stages the peer's half", + .gather_own_halves => "control: metal pair gather, this rank writes its own half on both sides", + .argmax_high_tie => "control: metal pair argmax, a tie goes to the highest id", + .argmax_no_flag => "control: metal pair argmax, a nonfinite logit is not flagged", + .pick_high_tie => "control: metal pair pick, a tie goes to the higher id", + }; + } + + /// The gate that must go red. + pub fn gate(m: Mutation) ?contract.gates.Gate { + return switch (m) { + .none => null, + .argmax_no_flag => .nonfinite, + else => .shapes, + }; + } + + const Edit = struct { from: []const u8, to: []const u8 }; + + fn edit(m: Mutation) Edit { + return switch (m) { + .none => unreachable, + // The stage's source is a frame row of `stride` = 2*width: the peer's half is in bounds and different. + .stage_peer_half => .{ .from = "const ulong in = ulong(r) * a.stride;", .to = "const ulong in = ulong(r) * a.stride + a.width;" }, + .gather_own_halves => .{ .from = "out[got + i] = ushort(word >> ((i & 1) * 16));", .to = "out[got + i] = mine[mrow + i];" }, + .argmax_high_tie => .{ .from = "if (ob > best || (ob == best && oa < at)) best = ob, at = oa;", .to = "if (ob > best || (ob == best && oa > at)) best = ob, at = oa;" }, + .argmax_no_flag => .{ .from = "if (!isfinite(v)) { bad = true; continue; }", .to = "if (false) { bad = true; continue; }" }, + .pick_high_tie => .{ .from = "(vp > vm || (vp == vm && ap < am)) ? ap : am;", .to = "(vp > vm || (vp == vm && ap > am)) ? ap : am;" }, + }; + } +}; + +/// pair_rows.metal's text for the kernel under test; a control's one edit must change it, or it proves nothing. +pub fn source(a: std.mem.Allocator, m: Mutation) ![]const u8 { + if (m == .none) return sources.core_pair_rows; + const e = m.edit(); + if (std.mem.indexOf(u8, sources.core_pair_rows, e.from) == null) return error.ControlEditMissing; + return std.mem.replaceOwned(u8, a, sources.core_pair_rows, e.from, e.to); +}