diff --git a/zig/kernels/metal/glm_ep.metal b/zig/kernels/metal/glm_ep.metal index 20826ce23..996f36d94 100644 --- a/zig/kernels/metal/glm_ep.metal +++ b/zig/kernels/metal/glm_ep.metal @@ -29,28 +29,14 @@ kernel void ep_localize(const device int* PICK [[buffer(0)]], const device int* if (theirs) THEIRS[o2 + b2] = i; if (t == 0) { LCOUNT[0] = n0; COUNTS[0] = n1; COUNTS[1] = n2; atomic_store_explicit(CNT, uint(n1), memory_order_relaxed); } } -// Exchange x writes its send slot once the peer's x - 1 landed (posted after our x - 2 reached it); give-ups counted -inline bool ep_slot_free(device atomic_uint* flag, uint x, device atomic_uint* gave_up, uint t, threadgroup uint* ok) { - if (t == 0) { - uint polls = 0; - *ok = 1u; - while (int(atomic_load_explicit(flag, memory_order_relaxed) - (x - 1u)) < 0) { - if (++polls > 400000000u) { atomic_fetch_add_explicit(gave_up, 1u, memory_order_relaxed); *ok = 0u; break; } - } - if (atomic_load_explicit(gave_up, memory_order_relaxed) != 0u) *ok = 0u; - } - threadgroup_barrier(mem_flags::mem_device | mem_flags::mem_threadgroup); - return *ok != 0u; -} // This Mac's picks' outputs (Y [picks, 4096] bf16) packed in MINE order for the peer; a threadgroup a pick kernel void ep_pack(const device uint* Y [[buffer(0)]], const device int* MINE [[buffer(1)]], - const device int* COUNTS [[buffer(2)]], device uint* OUT [[buffer(3)]], device atomic_uint* flag [[buffer(4)]], - constant uint& x [[buffer(5)]], device atomic_uint* gave_up [[buffer(6)]], uint3 tg [[threadgroup_position_in_grid]], + const device int* COUNTS [[buffer(2)]], device uint* OUT [[buffer(3)]], + const device ulong* previous [[buffer(4)]], constant ulong& expected [[buffer(5)]], uint3 tg [[threadgroup_position_in_grid]], uint3 tpos [[thread_position_in_threadgroup]]) { + if (previous[0] != expected) return; const int i = int(tg.y); const uint t = tpos.x; - threadgroup uint ok; - if (!ep_slot_free(flag, x, gave_up, t, &ok)) return; if (i >= COUNTS[0]) return; const device uint* src = Y + size_t(MINE[i]) * WORDS; device uint* dst = OUT + size_t(i) * WORDS; @@ -62,10 +48,9 @@ inline float ep_mul_add(float acc, float a, float b) { return a * b + acc; } #pragma clang fp contract(on) kernel void ep_rcombine(const device float* YP [[buffer(0)]], const device float* WTS [[buffer(1)]], constant int& rows [[buffer(2)]], device float* OUT [[buffer(3)]], device atomic_uint* CNT [[buffer(4)]], - device atomic_uint* flag [[buffer(5)]], constant uint& x [[buffer(6)]], device atomic_uint* gave_up [[buffer(7)]], - uint gid [[thread_position_in_grid]], uint t [[thread_index_in_threadgroup]]) { - threadgroup uint ok; - if (!ep_slot_free(flag, x, gave_up, t, &ok)) return; + const device ulong* previous [[buffer(5)]], constant ulong& expected [[buffer(6)]], + uint gid [[thread_position_in_grid]]) { + if (previous[0] != expected) return; const int r = int(gid) / 4096, d = int(gid) % 4096; if (gid == 0) atomic_store_explicit(CNT, uint(rows * 2), memory_order_relaxed); if (r >= rows) return; @@ -74,84 +59,56 @@ kernel void ep_rcombine(const device float* YP [[buffer(0)]], const device float for (int k = 1; k < TOPK; k++) acc = ep_mul_add(acc, WTS[r * TOPK + k], y[size_t(k) * 4096]); OUT[size_t(r) * 4096 + d] = acc; } -// By rows: once the peer's sums land, each branch is both Macs' sums in rank order, rounded, plus the shared expert -kernel void ep_rfinal(device atomic_uint* flag [[buffer(0)]], constant uint& x [[buffer(1)]], - device atomic_uint* gave_up [[buffer(2)]], const device float* MINE [[buffer(3)]], device atomic_uint* PEER [[buffer(4)]], +// By rows: both sums in rank order, rounded, plus the shared expert. A missing epoch keeps the receive slot unread. +kernel void ep_rfinal(const device ulong* ready [[buffer(0)]], constant ulong& x [[buffer(1)]], + const device float* MINE [[buffer(3)]], const device uint* PEER [[buffer(4)]], constant uint& rank [[buffer(5)]], const device bfloat* YS [[buffer(6)]], device bfloat* OUT [[buffer(7)]], constant int& rows [[buffer(8)]], uint3 tg [[threadgroup_position_in_grid]], uint3 tpos [[thread_position_in_threadgroup]]) { const uint t = tpos.x; - if (t == 0) { - uint polls = 0; - while (int(atomic_load_explicit(flag, memory_order_relaxed) - x) < 0) { - if (++polls > 400000000u) { atomic_fetch_add_explicit(gave_up, 1u, memory_order_relaxed); break; } - } - } - threadgroup_barrier(mem_flags::mem_device); const int i = int(tg.x) * 256 + int(t); if (i >= rows * 4096) return; - const float mine = MINE[i], peer = as_type(atomic_load_explicit(&PEER[i], memory_order_relaxed)); + const float mine = MINE[i], peer = ready[0] == x ? as_type(PEER[i]) : 0.0f; const float total = rank == 0u ? mine + peer : peer + mine; OUT[i] = bfloat(total) + YS[i]; } -// TP: this Mac's fp32 partial rows into the send slot once it is free, and the entries the host sends (two a row) +// TP: this Mac's fp32 partial rows into the send slot, and the entries the host sends (two a row) kernel void ep_rsend(const device float* PART [[buffer(0)]], constant int& rows [[buffer(1)]], - device float* OUT [[buffer(2)]], device atomic_uint* CNT [[buffer(3)]], device atomic_uint* flag [[buffer(4)]], - constant uint& x [[buffer(5)]], device atomic_uint* gave_up [[buffer(6)]], - uint gid [[thread_position_in_grid]], uint t [[thread_index_in_threadgroup]]) { - threadgroup uint ok; - if (!ep_slot_free(flag, x, gave_up, t, &ok)) return; + device float* OUT [[buffer(2)]], device atomic_uint* CNT [[buffer(3)]], + const device ulong* previous [[buffer(4)]], constant ulong& expected [[buffer(5)]], uint gid [[thread_position_in_grid]]) { + if (previous[0] != expected) return; if (gid == 0) atomic_store_explicit(CNT, uint(rows * 2), memory_order_relaxed); if (int(gid) < rows * 4096) OUT[gid] = PART[gid]; } -// TP: wait for the peer's partials of exchange x; each value is both Macs' partials added in rank order, rounded -kernel void ep_rsum(device atomic_uint* flag [[buffer(0)]], constant uint& x [[buffer(1)]], - device atomic_uint* gave_up [[buffer(2)]], const device float* MINE [[buffer(3)]], device atomic_uint* PEER [[buffer(4)]], +// TP: both Macs' partials added in rank order and rounded +kernel void ep_rsum(const device ulong* ready [[buffer(0)]], constant ulong& x [[buffer(1)]], + const device float* MINE [[buffer(3)]], const device uint* PEER [[buffer(4)]], constant uint& rank [[buffer(5)]], device bfloat* OUT [[buffer(6)]], constant int& rows [[buffer(7)]], uint3 tg [[threadgroup_position_in_grid]], uint3 tpos [[thread_position_in_threadgroup]]) { const uint t = tpos.x; - if (t == 0) { - uint polls = 0; - while (int(atomic_load_explicit(flag, memory_order_relaxed) - x) < 0) { - if (++polls > 400000000u) { atomic_fetch_add_explicit(gave_up, 1u, memory_order_relaxed); break; } - } - } - threadgroup_barrier(mem_flags::mem_device); const int i = int(tg.x) * 256 + int(t); if (i >= rows * 4096) return; - const float mine = MINE[i], peer = as_type(atomic_load_explicit(&PEER[i], memory_order_relaxed)); + const float mine = MINE[i], peer = ready[0] == x ? as_type(PEER[i]) : 0.0f; OUT[i] = bfloat(rank == 0u ? mine + peer : peer + mine); } -// Exchange x packed: the host sends it once the GPU gets here -kernel void ep_post(device atomic_uint* posted [[buffer(0)]], constant uint& x [[buffer(1)]], - device atomic_uint* gave_up [[buffer(2)]]) { - if (atomic_load_explicit(gave_up, memory_order_relaxed) == 0u) atomic_store_explicit(posted, x, memory_order_relaxed); -} -// Once the peer's exchange x lands (its flag follows its bytes), its entries (atomic reads) into their picks' rows of Y -kernel void ep_unpack(device atomic_uint* flag [[buffer(0)]], constant uint& x [[buffer(1)]], - device atomic_uint* gave_up [[buffer(2)]], device atomic_uint* IN [[buffer(3)]], const device int* THEIRS [[buffer(4)]], +// Copy the peer's entries into their picks' rows of Y +kernel void ep_unpack(const device ulong* ready [[buffer(0)]], constant ulong& x [[buffer(1)]], + const device uint* IN [[buffer(3)]], const device int* THEIRS [[buffer(4)]], const device int* COUNTS [[buffer(5)]], device uint* Y [[buffer(6)]], uint3 tg [[threadgroup_position_in_grid]], uint3 tpos [[thread_position_in_threadgroup]]) { const uint t = tpos.x; - if (t == 0) { - uint polls = 0; - while (int(atomic_load_explicit(flag, memory_order_relaxed) - x) < 0) { - if (++polls > 400000000u) { atomic_fetch_add_explicit(gave_up, 1u, memory_order_relaxed); break; } - } - } - threadgroup_barrier(mem_flags::mem_device); const int i = int(tg.y); if (i >= COUNTS[1]) return; - device atomic_uint* src = IN + size_t(i) * WORDS; + const device uint* src = IN + size_t(i) * WORDS; device uint* dst = Y + size_t(THEIRS[i]) * WORDS; - for (uint j = t; j < uint(WORDS); j += 256) dst[j] = atomic_load_explicit(&src[j], memory_order_relaxed); + const bool landed = ready[0] == x; + for (uint j = t; j < uint(WORDS); j += 256) dst[j] = landed ? src[j] : 0u; } // TP: each row's best logit of this Mac's vocabulary rows [lo, lo + n) and its token id to send (n = 0: none) kernel void ep_amax_send(const device bfloat* LOGITS [[buffer(0)]], const device uint* PICKS [[buffer(1)]], constant uint4& arg [[buffer(2)]], device uint* OUT [[buffer(3)]], device atomic_uint* CNT [[buffer(4)]], - device atomic_uint* flag [[buffer(5)]], constant uint& x [[buffer(6)]], device atomic_uint* gave_up [[buffer(7)]], + const device ulong* previous [[buffer(5)]], constant ulong& expected [[buffer(6)]], uint t [[thread_index_in_threadgroup]]) { - threadgroup uint ok; - if (!ep_slot_free(flag, x, gave_up, t, &ok)) return; + if (previous[0] != expected) return; const uint rows = arg.x, n = arg.y, lo = arg.z; if (t == 0) atomic_store_explicit(CNT, 1u, memory_order_relaxed); if (t >= rows) return; @@ -161,18 +118,15 @@ kernel void ep_amax_send(const device bfloat* LOGITS [[buffer(0)]], const device OUT[2 * t + 1] = has ? lo + p : 0xffffffffu; } // TP: once the peer's candidates land, each row's pick is the larger logit, the lower id on a tie (one Mac's rule) -kernel void ep_amax_merge(device atomic_uint* flag [[buffer(0)]], constant uint& x [[buffer(1)]], - device atomic_uint* gave_up [[buffer(2)]], const device uint* MINE [[buffer(3)]], device atomic_uint* PEER [[buffer(4)]], +kernel void ep_amax_merge(const device ulong* ready [[buffer(0)]], constant ulong& x [[buffer(1)]], + const device uint* MINE [[buffer(3)]], const device uint* PEER [[buffer(4)]], constant uint& rows [[buffer(5)]], device uint* PICKS [[buffer(6)]], uint t [[thread_index_in_threadgroup]]) { - if (t == 0) { - uint polls = 0; - while (int(atomic_load_explicit(flag, memory_order_relaxed) - x) < 0) { - if (++polls > 400000000u) { atomic_fetch_add_explicit(gave_up, 1u, memory_order_relaxed); break; } - } - } - threadgroup_barrier(mem_flags::mem_device); if (t >= rows) return; - const float a = as_type(MINE[2 * t]), b = as_type(atomic_load_explicit(&PEER[2 * t], memory_order_relaxed)); - const uint ia = MINE[2 * t + 1], ib = atomic_load_explicit(&PEER[2 * t + 1], memory_order_relaxed); + if (ready[0] != x) { + PICKS[t] = 0u; + return; + } + const float a = as_type(MINE[2 * t]), b = as_type(PEER[2 * t]); + const uint ia = MINE[2 * t + 1], ib = PEER[2 * t + 1]; PICKS[t] = (b > a || (b == a && ib < ia)) ? ib : ia; } diff --git a/zig/src/core/contract/glm_ep_records.zig b/zig/src/core/contract/glm_ep_records.zig new file mode 100644 index 000000000..43a353162 --- /dev/null +++ b/zig/src/core/contract/glm_ep_records.zig @@ -0,0 +1,13 @@ +//! GLM's Metal handoff launches with no threadgroup scratch. +const M = @import("limits.zig").MetalLaunch; + +pub const records = [_]M{ + .{ .name = "metal glm ep ep_pack", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "metal glm ep ep_unpack", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "metal glm ep ep_rcombine", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "metal glm ep ep_rfinal", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "metal glm ep ep_rsend", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "metal glm ep ep_rsum", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "metal glm ep ep_amax_send", .threads = 64, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "metal glm ep ep_amax_merge", .threads = 64, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, +}; diff --git a/zig/src/core/contract/registry_joins.zig b/zig/src/core/contract/registry_joins.zig index 1050affbb..9b8f68e73 100644 --- a/zig/src/core/contract/registry_joins.zig +++ b/zig/src/core/contract/registry_joins.zig @@ -57,6 +57,7 @@ pub const all_metal = @import("lane_q8_records.zig").records ++ @import("flashnext_pending.zig").records ++ @import("simd_pairs_records.zig").records ++ @import("simd_half_records.zig").records ++ + @import("glm_ep_records.zig").records ++ @import("metal_pair_records.zig").records ++ @import("glm_rows_records.zig").records ++ @import("weight_only_records.zig").metal ++ diff --git a/zig/src/fabric/mcdma.zig b/zig/src/fabric/mcdma.zig index ecc12db4a..43e615a4e 100644 --- a/zig/src/fabric/mcdma.zig +++ b/zig/src/fabric/mcdma.zig @@ -249,12 +249,19 @@ pub const Endpoint = struct { /// Send `len` bytes at `local` to peer `offset`, then `value` at `flag`; head goes before `local`; keep till acked. pub fn writeSignalFrom(self: *Endpoint, peer: u32, local: usize, offset: usize, len: usize, flag: usize, value: u64) rma.Error!void { - if (len == 0) return signal(self, peer, flag, value); // the library takes a zero-length write-and-signal as a dead peer - if (local < abi.ws_room or !self.in(local - abi.ws_room, len + abi.ws_room) or !self.in(offset, len) or !self.in(flag, 8)) return error.OutOfBounds; + return self.writeSignalFromUntil(peer, local, offset, len, flag, value, std.math.maxInt(u64)); + } + + /// The call deadline bounds admission to the native write. The library owns the timeout once inside that call. + pub fn writeSignalFromUntil(self: *Endpoint, peer: u32, local: usize, offset: usize, len: usize, flag: usize, value: u64, until: u64) rma.Error!void { + if (len > 0 and (local < abi.ws_room or !self.in(local - abi.ws_room, len + abi.ws_room) or !self.in(offset, len) or !self.in(flag, 8))) return error.OutOfBounds; + if (len == 0) return self.signalUntil(peer, flag, value, until); if (flag % 8 != 0) return error.Unaligned; const port = try self.portFor(peer); - try port.mutex.lock(try self.end()); + const end_at = @min(until, try self.end()); + try port.mutex.lock(end_at); defer port.mutex.unlock(); + if (try words.checkedNowNs() >= end_at) return error.Timeout; if (port.link == roce) { // ConnectX has no fused message: the write, then the signal, on the one queue pair try self.checkPort(port, self.library.api.write(port.peer.?, local, offset, len)); port.pending = true; @@ -355,8 +362,14 @@ pub const Endpoint = struct { } fn flushPort(self: *Endpoint, port: *Port) rma.Error!void { + try self.flushPortUntil(port, try self.end()); + } + + fn flushPortUntil(self: *Endpoint, port: *Port, end_at: u64) rma.Error!void { if (!port.pending) return; - const timeout = try self.left(); + const now = try words.checkedNowNs(); + if (now >= end_at) return error.Timeout; + const timeout = end_at - now; try self.checkPort(port, self.library.api.flush(port.peer.?, timeout)); port.used = 0; port.pending = false; @@ -446,7 +459,10 @@ pub const Endpoint = struct { } fn signal(ptr: *anyopaque, peer: u32, offset: usize, value: u64) rma.Error!void { - const self = of(ptr); + return of(ptr).signalUntil(peer, offset, value, std.math.maxInt(u64)); + } + + fn signalUntil(self: *Endpoint, peer: u32, offset: usize, value: u64, until: u64) rma.Error!void { if (offset % 8 != 0) return error.Unaligned; if (!self.in(offset, 8)) return error.OutOfBounds; if (peer == self.config.rank) { @@ -454,8 +470,10 @@ pub const Endpoint = struct { return; } const port = try self.portFor(peer); - try port.mutex.lock(try self.end()); + const end_at = @min(until, try self.end()); + try port.mutex.lock(end_at); defer port.mutex.unlock(); + if (try words.checkedNowNs() >= end_at) return error.Timeout; try self.checkPort(port, self.library.api.signal(port.peer.?, offset, value)); port.pending = true; } @@ -480,6 +498,15 @@ pub const Endpoint = struct { } } + /// Retire a peer's posted writes without changing deadlines on concurrent control or benchmark calls. + pub fn flushPeerUntil(self: *Endpoint, peer: u32, until: u64) rma.Error!void { + const port = try self.portFor(peer); + const end_at = @min(until, try self.end()); + try port.mutex.lock(end_at); + defer port.mutex.unlock(); + try self.flushPortUntil(port, end_at); + } + fn kind(ptr: *anyopaque, peer: u32) rma.LinkKind { const self = of(ptr); if (peer == self.config.rank) return .local; diff --git a/zig/src/fabric/mcdma_test.zig b/zig/src/fabric/mcdma_test.zig index a1fde50d5..a3723537c 100644 --- a/zig/src/fabric/mcdma_test.zig +++ b/zig/src/fabric/mcdma_test.zig @@ -347,6 +347,49 @@ test "write2Signal refuses a bad flag before posting anything, over ConnectX as } } +test "an expired send deadline refuses payload and signal-only posts before calling the library" { + for ([_]c_int{ roce, abi.thunderbolt }) |link| { + Fake.reset(link); + const ep = try Fake.endpoint(Fake.cx); + defer ep.deinit(); + const expired = try words.checkedNowNs(); + try std.testing.expectError(error.Timeout, ep.writeSignalFromUntil(1, 1024, 2048, 64, 24, 3, expired)); + try std.testing.expectError(error.Timeout, ep.writeSignalFromUntil(1, 0, 0, 0, 32, 4, expired)); + try std.testing.expectEqual(@as(usize, 0), Fake.n); + try std.testing.expect(!ep.ports[0].pending); + } +} + +test "a bounded peer flush retires a signal-only post only after native success" { + Fake.reset(abi.thunderbolt); + const ep = try Fake.endpoint(Fake.cx); + defer ep.deinit(); + const budget = 500 * std.time.ns_per_ms; + + var deadline = (try words.checkedNowNs()) + budget; + try ep.writeSignalFromUntil(1, 0, 0, 0, 32, 4, deadline); + try std.testing.expectEqualDeep(Fake.Call{ .signal = .{ 32, 4 } }, Fake.calls[0]); + try std.testing.expect(ep.ports[0].pending); + try std.testing.expectError(error.Timeout, ep.flushPeerUntil(1, try words.checkedNowNs())); + try std.testing.expectEqual(@as(usize, 1), Fake.n); + try std.testing.expect(ep.ports[0].pending); + + try ep.flushPeerUntil(1, deadline); + try std.testing.expectEqual(@as(usize, 2), Fake.n); + try std.testing.expect(Fake.calls[1] == .flush); + try std.testing.expect(Fake.calls[1].flush > 0 and Fake.calls[1].flush <= budget); + try std.testing.expect(!ep.ports[0].pending); + + deadline = (try words.checkedNowNs()) + budget; + try ep.writeSignalFromUntil(1, 1024, 2048, 64, 24, 5, deadline); + Fake.flush_status = 6; + try std.testing.expectError(error.Timeout, ep.flushPeerUntil(1, deadline)); + try std.testing.expectEqual(@as(usize, 4), Fake.n); + try std.testing.expect(Fake.calls[3] == .flush); + try std.testing.expect(Fake.calls[3].flush > 0 and Fake.calls[3].flush <= budget); + try std.testing.expect(ep.ports[0].pending); +} + test "one endpoint with a Thunderbolt peer and a ConnectX peer picks each port's path and kind" { Fake.reset(abi.thunderbolt); Fake.mixed = true; diff --git a/zig/src/families/glm/checks.zig b/zig/src/families/glm/checks.zig index 00df4b892..5767a30be 100644 --- a/zig/src/families/glm/checks.zig +++ b/zig/src/families/glm/checks.zig @@ -23,9 +23,9 @@ pub fn capture(e: *Engine, prompt: []const u32, path: []const u8) !void { const dump = try e.arena.buffer(bytes); var x = e.ctx(); x.dump = dump; - const b = e.begin(); - fwd.backbone(&x, b.enc, e.prompt_ids, n, 0); - fwd.head(&x, b.enc, e.sc.hidden, dump.at(x.dump_at), e.sc.picks, n); + var b = e.begin(); + fwd.backbone(&x, &b, e.prompt_ids, n, 0); + fwd.head(&x, &b, e.sc.hidden, dump.at(x.dump_at), e.sc.picks, n); try e.finish(b.cb, b.enc); fwd.flipKda(&x); const file = platform.c.fopen(try std.fmt.allocPrintSentinel(e.gpa, "{s}", .{path}, 0), "wb") orelse return error.OpenFailed; @@ -57,9 +57,9 @@ pub fn trace(e: *Engine, prompt: []const u32, reply: []const u32, steps: usize, const n: u32 = if (at < P) @min(st.max_rows, P - at) else 1; x.dump = dump; x.dump_at = 0; - const b = e.begin(); - fwd.backbone(&x, b.enc, e.prompt_ids.at(@as(usize, at) * 4), n, at); - fwd.head(&x, b.enc, e.sc.hidden, dump.at(x.dump_at), e.sc.picks, n); + var b = e.begin(); + fwd.backbone(&x, &b, e.prompt_ids.at(@as(usize, at) * 4), n, at); + fwd.head(&x, &b, e.sc.hidden, dump.at(x.dump_at), e.sc.picks, n); try e.finish(b.cb, b.enc); fwd.flipKda(&x); const bytes = x.dump_at + @as(usize, n) * c.vocab * 2; @@ -88,9 +88,9 @@ pub fn forced(e: *Engine, prompt: []const u32, reply: []const u32) !struct { sam var first: ?usize = null; while (at < total) { const n: u32 = if (at < P) @min(st.max_rows, P - at) else 1; - const b = e.begin(); - fwd.backbone(&x, b.enc, e.prompt_ids.at(@as(usize, at) * 4), n, at); - fwd.head(&x, b.enc, e.sc.hidden.at(@as(usize, n - 1) * c.hidden * 2), e.sc.logits, e.sc.picks, 1); + var b = e.begin(); + fwd.backbone(&x, &b, e.prompt_ids.at(@as(usize, at) * 4), n, at); + fwd.head(&x, &b, e.sc.hidden.at(@as(usize, n - 1) * c.hidden * 2), e.sc.logits, e.sc.picks, 1); try e.finish(b.cb, b.enc); fwd.flipKda(&x); at += n; @@ -131,13 +131,13 @@ pub fn profile(e: *Engine, depth: u32, reps: usize, only: bool, parts: bool) !vo const m = if (mask == 0xffff_ffff) all else if (only and mask != 0) all & ~mask else mask; if (parts) x.pskip = m else x.skip = m; for (0..reps + 1) |rep| { - const b = e.begin(); + var b = e.begin(); if (d > 0 and x.skip & fwd.Class.mtp == 0) { - mtp.run(&x, b.enc, e.sc.hidden, e.sc.picks, R, e.s.mtp_pos, e.sc.ids.at(4)); - for (1..d) |j| mtp.chain(&x, b.enc, if (j == 1) e.sc.m_x.at(@as(usize, R - 1) * D * 2) else e.sc.m_x, e.sc.ids.at(j * 4), e.s.mtp_pos + R + @as(u32, @intCast(j)) - 1, e.sc.ids.at((j + 1) * 4)); + mtp.run(&x, &b, e.sc.hidden, e.sc.picks, R, e.s.mtp_pos, e.sc.ids.at(4)); + for (1..d) |j| mtp.chain(&x, &b, if (j == 1) e.sc.m_x.at(@as(usize, R - 1) * D * 2) else e.sc.m_x, e.sc.ids.at(j * 4), e.s.mtp_pos + R + @as(u32, @intCast(j)) - 1, e.sc.ids.at((j + 1) * 4)); } - fwd.backbone(&x, b.enc, e.sc.ids, R, e.s.pos); - fwd.head(&x, b.enc, e.sc.hidden, e.sc.logits, e.sc.picks, R); + fwd.backbone(&x, &b, e.sc.ids, R, e.s.pos); + fwd.head(&x, &b, e.sc.hidden, e.sc.logits, e.sc.picks, R); try e.finish(b.cb, b.enc); if (rep > 0) times[rep - 1] = (e.gpu[1] - e.gpu[0]) * 1e3; // the first run warms the class's state } @@ -213,8 +213,8 @@ pub fn profilePrompt(e: *Engine, prompt: []const u32, reps: usize, only: bool, p if (parts) x.pskip = m else x.skip = m; for (0..reps + 1) |rep| { e.s.reset(); - const b = e.begin(); - prompt_mod.backbone(pr, &x, b.enc, e.prompt_ids, rows, 0); + var b = e.begin(); + prompt_mod.backbone(pr, &x, &b, e.prompt_ids, rows, 0); try e.finish(b.cb, b.enc); if (rep > 0) times[rep - 1] = (e.gpu[1] - e.gpu[0]) * 1e3; } diff --git a/zig/src/families/glm/engine.zig b/zig/src/families/glm/engine.zig index 0a5a7f9fe..5f47f3b13 100644 --- a/zig/src/families/glm/engine.zig +++ b/zig/src/families/glm/engine.zig @@ -13,6 +13,7 @@ const prompt_mod = @import("prompt.zig"); const Fence = @import("../../core/fence.zig").Fence; const kernels = @import("kernels.zig"); const ep_mod = @import("ep.zig"); +const Pass = @import("pass.zig").Pass; const CopyIndex = @import("../../core/copy_index.zig").CopyIndex; const checks = @import("checks.zig"); const load_plan = @import("load_plan.zig"); @@ -235,6 +236,7 @@ pub const Engine = struct { me.rest[2] = @intFromBool(learn); // --learn on both Macs or neither: each holds its half of a learned state e.ep = try ep_mod.Ep.init(gpa, io, e.device, s, me); } + errdefer if (e.ep) |ep| ep.deinit(gpa); if (vision) |lim| try vision_mod.attach(e, gpa, dir, f.bytes[0..f.size], lim); // --vision: the tower and its rows errdefer if (e.vision) |v| v.deinit(); // a process's first prompt pays a one-time cost: every rank of a pair pays it here at the same step @@ -316,7 +318,7 @@ pub const Engine = struct { } /// A command buffer ordered after the last one this engine committed. - pub fn begin(e: *Engine) struct { cb: mtl.CommandBuffer, enc: mtl.ComputeEncoder } { + pub fn begin(e: *Engine) Pass { const cb = e.queue.commandBuffer(); if (e.ev > 0) cb.waitFor(e.event, e.ev); return .{ .cb = cb, .enc = cb.compute(.serial) }; @@ -333,6 +335,7 @@ pub const Engine = struct { e.ev += 1; cb.signal(e.event, e.ev); cb.commit(); + if (e.ep) |ep| ep.submit(); e.committed = platform.c.mach_absolute_time(); if (spin) { while (@backingInt(cb.status()) < @backingInt(mtl.types.CommandBufferStatus.completed)) std.atomic.spinLoopHint(); @@ -341,10 +344,14 @@ pub const Engine = struct { if (cb.failure()) |msg| { e.event.set(e.ev); // a failed buffer may never signal: the next one must not wait on it std.log.err("glm: command buffer failed: {s}", .{msg}); - if (e.ep) |ep| ep.failed.store(true, .release); // its exchanges are out of step with the peer's + if (e.ep) |ep| ep.gpuFailed(); return error.GpuFailed; } - if (e.ep) |ep| if (ep.failed.load(.acquire) or ep.gaveUp() > 0) return error.EpLinkFailed; + try e.checkPair(); + } + + pub fn checkPair(e: *const Engine) !void { + if (e.ep) |ep| if (ep.failed.load(.acquire)) return error.EpLinkFailed; } /// Wait until every command buffer this engine committed has completed (before the host touches shared state). @@ -375,6 +382,7 @@ pub const Engine = struct { /// One reply, greedy or sampled by `e.sampling`; `depth` drafts a round (drafted equals plain, depth 0). pub fn generate(e: *Engine, prompt: []const u32, max_tokens: usize, eos: []const u32, depth: usize, out: Out) !Result { + try e.checkPair(); const c = &e.c; const D = c.hidden; const P: u32 = @intCast(prompt.len); @@ -401,7 +409,7 @@ pub const Engine = struct { const n = @min(@as(u32, if (chunk) e.chunk_rows else st.max_rows), end - at); const last = at + n == P; const absorb = if (last) n - 1 else n; - const b = e.begin(); + var b = e.begin(); if (last and e.trace_last != null) { x.dump = e.trace_last; x.dump_at = 0; @@ -413,20 +421,20 @@ pub const Engine = struct { const pr = &e.pr.?; var px = x; px.sc = &pr.streams; - prompt_mod.backbone(pr, &px, b.enc, e.prompt_ids.at(@as(usize, at) * 4), n, at); + prompt_mod.backbone(pr, &px, &b, e.prompt_ids.at(@as(usize, at) * 4), n, at); fwd.flipKda(&x); const hidden = pr.streams.hidden; if (last) { // the last row into the decode scratch: the head's first token, the MTP head's first row copyRow(&x, b.enc, hidden.at(@as(usize, n - 1) * D * 2), e.sc.hidden, D); - fwd.head(&x, b.enc, e.sc.hidden, e.sc.logits, e.sc.picks, 1); + fwd.head(&x, &b, e.sc.hidden, e.sc.logits, e.sc.picks, 1); } if (d > 0 and absorb > 0) prompt_mod.mtp(pr, &px, b.enc, hidden, e.prompt_ids.at(@as(usize, at + 1) * 4), absorb, at); last_n = 1; } else { - fwd.backbone(&x, b.enc, e.prompt_ids.at(@as(usize, at) * 4), n, at); + fwd.backbone(&x, &b, e.prompt_ids.at(@as(usize, at) * 4), n, at); fwd.flipKda(&x); - if (last) fwd.head(&x, b.enc, e.sc.hidden.at(@as(usize, n - 1) * D * 2), e.sc.logits, e.sc.picks, 1); - if (d > 0 and absorb > 0) mtp.run(&x, b.enc, e.sc.hidden, e.prompt_ids.at(@as(usize, at + 1) * 4), absorb, at, e.sc.m_picks); + if (last) fwd.head(&x, &b, e.sc.hidden.at(@as(usize, n - 1) * D * 2), e.sc.logits, e.sc.picks, 1); + if (d > 0 and absorb > 0) mtp.run(&x, &b, e.sc.hidden, e.prompt_ids.at(@as(usize, at + 1) * 4), absorb, at, e.sc.m_picks); last_n = n; } if (d > 0 and absorb > 0) e.s.mtp_pos = at + absorb; // both prompt paths set it after their MTP absorb @@ -471,7 +479,7 @@ pub const Engine = struct { var round_draws: fwd.Draws = .{}; // the current decode round's rules, rebuilt each round while (true) { const t_enc = platform.c.mach_absolute_time(); - const b = e.begin(); + var b = e.begin(); if (window > 0) { // the last window's kept rows (the prompt's windows kept all theirs) fwd.keepKda(&x, b.enc, window, keep); fwd.flipKda(&x); @@ -505,24 +513,24 @@ pub const Engine = struct { } if (d > 0) { const next = if (window == 0) e.sc.next else e.sc.picks; - mtp.run(&x, b.enc, e.sc.hidden.at(@as(usize, h0) * D * 2), next, keep, e.s.mtp_pos, if (copied > 0) null else e.sc.ids.at(4)); + mtp.run(&x, &b, e.sc.hidden.at(@as(usize, h0) * D * 2), next, keep, e.s.mtp_pos, if (copied > 0) null else e.sc.ids.at(4)); const base = e.s.mtp_pos + keep; if (copied == 0) for (1..d) |j| { const h = if (j == 1) e.sc.m_x.at(@as(usize, keep - 1) * D * 2) else e.sc.m_x; if (e.rank_log) x.m_row = @intCast(j); - mtp.chain(&x, b.enc, h, e.sc.ids.at(j * 4), base + @as(u32, @intCast(j)) - 1, e.sc.ids.at((j + 1) * 4)); + mtp.chain(&x, &b, h, e.sc.ids.at(j * 4), base + @as(u32, @intCast(j)) - 1, e.sc.ids.at((j + 1) * 4)); }; x.m_row = 0; e.s.mtp_pos = base; } const R = if (copied > 0) copied + 1 else d + 1; - fwd.backbone(&x, b.enc, e.sc.ids, R, e.s.pos); + fwd.backbone(&x, &b, e.sc.ids, R, e.s.pos); // The decode rows draw keyed at their context positions; the MTP head's drafts drew greedily. x.draws = if (e.sampling) |p| if (p.temperature > 0) blk: { round_draws = fwd.rulesFor(p, e.s.pos + 1, c.vocab, R); break :blk &round_draws; } else null else null; - fwd.head(&x, b.enc, e.sc.hidden, e.sc.logits, e.sc.picks, R); + fwd.head(&x, &b, e.sc.hidden, e.sc.logits, e.sc.picks, R); try e.finish(b.cb, b.enc); res.encode_seconds += @as(f64, @floatFromInt(e.committed - t_enc)) / 24e6; res.gpu_seconds += e.gpu[1] - e.gpu[0]; diff --git a/zig/src/families/glm/ep.zig b/zig/src/families/glm/ep.zig index bfa64edd8..8498b4b76 100644 --- a/zig/src/families/glm/ep.zig +++ b/zig/src/families/glm/ep.zig @@ -3,6 +3,8 @@ const std = @import("std"); const platform = @import("sys"); const mtl = @import("metal"); const fabric = @import("fabric"); +const Fence = @import("../../core/fence.zig").Fence; +const Pass = @import("pass.zig").Pass; const settings = @import("../flashnext/tp_settings.zig"); const control = @import("ep_control.zig"); const Ref = @import("weights.zig").Ref; @@ -28,11 +30,13 @@ const SEND = RECV + 2 * SLOT; const FLAGS = SEND + 2 * (PAGE + SLOT); const FLAG = FLAGS; // u64: the last exchange whose peer entries have landed const SYNC = FLAGS + PAGE; -const POSTED = SYNC; // u32: the last exchange the GPU has packed (its low 32 bits) const COUNT = SYNC + 64; // u32 by parity: each exchange's packed entries -const GAVE_UP = SYNC + 128; // u32: GPU waits that gave up (nonzero: the run is invalid) const CONTROL = SYNC + PAGE; // the control protocol's region (ep_control.zig) pub const WINDOW = CONTROL + control.BYTES; +const READY_STRIDE = 64; // the most recent valid receive for each parity, in a separate Metal buffer +const post_timeout = 30 * 24_000_000; +const transfer_budget_ns = std.time.ns_per_s; +const failed_event = std.math.maxInt(u64); fn recvAt(x: u64) usize { return RECV + @as(usize, @intCast(x % 2)) * SLOT; @@ -42,11 +46,6 @@ fn sendAt(x: u64) usize { return SEND + @as(usize, @intCast(x % 2)) * (PAGE + SLOT) + PAGE; } -/// Whether the GPU's posted word (an exchange's low 32 bits) has reached exchange x, across the 32-bit wrap. -fn reached(posted: u32, x: u64) bool { - return @as(i32, @bitCast(posted -% @as(u32, @truncate(x)))) >= 0; -} - // Lists one exchange keeps (i32): this Mac's unique experts with local ids and members, the picks each Mac computes. const L_IDS = 0; const L_MEM = L_IDS + MAXP * 4; @@ -68,9 +67,12 @@ pub const Ep = struct { win: []u8, wbuf: mtl.Buffer, lists: mtl.Buffer, + ready: mtl.Buffer, + posted: mtl.SharedEvent, + landed: mtl.SharedEvent, + fence: Fence, localize_pipe: mtl.Pipeline, pack_pipe: mtl.Pipeline, - post_pipe: mtl.Pipeline, unpack_pipe: mtl.Pipeline, rcombine_pipe: mtl.Pipeline, rfinal_pipe: mtl.Pipeline, @@ -78,12 +80,17 @@ pub const Ep = struct { rsum_pipe: mtl.Pipeline, amax_send_pipe: mtl.Pipeline, amax_merge_pipe: mtl.Pipeline, - x: u64 = 0, // exchanges encoded so far, the same count on both Macs (the GPU sees the low 32 bits) + x: u64 = 0, // exchanges encoded so far, the same count on both Macs + submitted: std.atomic.Value(u64) = .init(0), ctl: control.Control, // identities, requests and stop decisions thread: ?std.Thread = null, + monitor: ?std.Thread = null, + publication: std.atomic.Mutex = .unlocked, + transfer: ?struct { exchange: u64, deadline: u64 } = null, // under publication stop: std.atomic.Value(bool) = .init(false), failed: std.atomic.Value(bool) = .init(false), trace: bool = false, // GLM_EP_TRACE: log every exchange the host sends + by_experts: bool, delay_ns: u64 = 0, // GLM_EP_DELAY_US: hold every prompt-chunk send (over 1 MiB) this long: a slow link, for stress runs sent: u64 = 0, // exchanges the host has sent held_ticks: u64 = 0, // their time from the GPU's post to the send returning @@ -109,16 +116,49 @@ pub const Ep = struct { errdefer wbuf.deinit(); const lists = try device.buffer(LISTS, opts); errdefer lists.deinit(); + const ready = try device.buffer(2 * READY_STRIDE, opts); + errdefer ready.deinit(); + @memset(ready.contents()[0 .. 2 * READY_STRIDE], 0); + const posted = try device.sharedEvent(); + errdefer posted.deinit(); + const landed = try device.sharedEvent(); + errdefer landed.deinit(); + const fence = try Fence.init(device); + errdefer fence.deinit(); const lib_m = try mtl.Library.fromSource(device, source, mtl.CompileOptions.mlx()); defer lib_m.deinit(); - var pipes: [10]mtl.Pipeline = undefined; + var pipes: [9]mtl.Pipeline = undefined; var made: usize = 0; errdefer for (pipes[0..made]) |pp| pp.deinit(); - for ([_][]const u8{ "ep_localize", "ep_pack", "ep_post", "ep_unpack", "ep_rcombine", "ep_rfinal", "ep_rsend", "ep_rsum", "ep_amax_send", "ep_amax_merge" }) |name| { + for ([_][]const u8{ "ep_localize", "ep_pack", "ep_unpack", "ep_rcombine", "ep_rfinal", "ep_rsend", "ep_rsum", "ep_amax_send", "ep_amax_merge" }) |name| { pipes[made] = try mtl.Pipeline.init(device, lib_m, name, false); made += 1; } - t.* = .{ .rank = s.rank, .peer = 1 - s.rank, .own = .{ me.own_lo, me.own_hi }, .link = link, .rd = rd, .win = win, .wbuf = wbuf, .lists = lists, .localize_pipe = pipes[0], .pack_pipe = pipes[1], .post_pipe = pipes[2], .unpack_pipe = pipes[3], .rcombine_pipe = pipes[4], .rfinal_pipe = pipes[5], .rsend_pipe = pipes[6], .rsum_pipe = pipes[7], .amax_send_pipe = pipes[8], .amax_merge_pipe = pipes[9], .ctl = undefined }; + t.* = .{ + .rank = s.rank, + .peer = 1 - s.rank, + .own = .{ me.own_lo, me.own_hi }, + .link = link, + .rd = rd, + .win = win, + .wbuf = wbuf, + .lists = lists, + .ready = ready, + .posted = posted, + .landed = landed, + .fence = fence, + .localize_pipe = pipes[0], + .pack_pipe = pipes[1], + .unpack_pipe = pipes[2], + .rcombine_pipe = pipes[3], + .rfinal_pipe = pipes[4], + .rsend_pipe = pipes[5], + .rsum_pipe = pipes[6], + .amax_send_pipe = pipes[7], + .amax_merge_pipe = pipes[8], + .ctl = undefined, + .by_experts = me.rest[0] == 0, + }; t.ctl = try control.Control.init(gpa, rd, CONTROL, &t.failed); errdefer t.ctl.deinit(gpa); t.ctl.inventory = true; @@ -127,6 +167,11 @@ pub const Ep = struct { errdefer if (t.benchmark_group) |b| gpa.free(b); t.trace = platform.c.getenv("GLM_EP_TRACE") != null; if (platform.c.getenv("GLM_EP_DELAY_US")) |v| t.delay_ns = 1000 * (std.fmt.parseInt(u64, std.mem.span(v), 10) catch 0); + t.monitor = try std.Thread.spawn(.{}, watchTransfers, .{t}); + errdefer { + t.stop.store(true, .release); + t.monitor.?.join(); + } t.thread = try std.Thread.spawn(.{}, service, .{t}); std.log.info("expert parallel: rank {d} of 2 holds experts {d}-{d}, connected", .{ t.rank, me.own_lo, me.own_hi - 1 }); return t; @@ -134,18 +179,26 @@ pub const Ep = struct { /// Say goodbye (rank 0), stop the sending thread and close the link (no GPU work may still use the window). pub fn deinit(t: *Ep, gpa: std.mem.Allocator) void { - t.ctl.bye(); - t.rd.flush() catch {}; t.stop.store(true, .release); if (t.thread) |th| th.join(); + if (t.monitor) |th| th.join(); + if (!t.failed.load(.acquire)) if (fabric.words.checkedNowNs()) |now| { + t.rd.setDeadline(now + std.time.ns_per_s); + t.ctl.bye(); + t.rd.flush() catch {}; + } else |_| {}; if (t.sent > 0) { const n: f64 = @floatFromInt(t.sent); const late: f64 = @floatFromInt(@max(t.st.late, 1)); if (t.st.big_n > 0) std.log.info("expert parallel rank {d}: {d} prompt-chunk exchanges, {d:.0} us a send, the peer's sums landed {d:.0} us after our post on average", .{ t.rank, t.st.big_n, @as(f64, @floatFromInt(t.big_held_ticks)) / @as(f64, @floatFromInt(@max(t.big_sent, 1))) / 24.0, @as(f64, @floatFromInt(t.st.big_ticks)) / @as(f64, @floatFromInt(t.st.big_n)) / 24.0 }); - std.log.info("expert parallel rank {d}: {d} exchanges, {d:.1} us a send; the peer's entries landed {d:.1} us after our post on average ({d} times, {d} before it); entries a exchange: ours {d:.2}, theirs {d:.2}, |difference| {d:.2}; {d} GPU waits gave up", .{ t.rank, t.sent, @as(f64, @floatFromInt(t.held_ticks)) / n / 24.0, @as(f64, @floatFromInt(t.st.late_ticks)) / late / 24.0, t.st.late, t.st.early, @as(f64, @floatFromInt(t.st.mine)) / n, @as(f64, @floatFromInt(t.st.theirs)) / n, @as(f64, @floatFromInt(t.st.imbalance)) / n, t.gaveUp() }); + std.log.info("expert parallel rank {d}: {d} exchanges, {d:.1} us a send; the peer's entries landed {d:.1} us after our post on average ({d} times, {d} before it); entries a exchange: ours {d:.2}, theirs {d:.2}, |difference| {d:.2}", .{ t.rank, t.sent, @as(f64, @floatFromInt(t.held_ticks)) / n / 24.0, @as(f64, @floatFromInt(t.st.late_ticks)) / late / 24.0, t.st.late, t.st.early, @as(f64, @floatFromInt(t.st.mine)) / n, @as(f64, @floatFromInt(t.st.theirs)) / n, @as(f64, @floatFromInt(t.st.imbalance)) / n }); } t.ctl.deinit(gpa); - for ([_]mtl.Pipeline{ t.localize_pipe, t.pack_pipe, t.post_pipe, t.unpack_pipe, t.rcombine_pipe, t.rfinal_pipe, t.rsend_pipe, t.rsum_pipe, t.amax_send_pipe, t.amax_merge_pipe }) |pp| pp.deinit(); + for ([_]mtl.Pipeline{ t.localize_pipe, t.pack_pipe, t.unpack_pipe, t.rcombine_pipe, t.rfinal_pipe, t.rsend_pipe, t.rsum_pipe, t.amax_send_pipe, t.amax_merge_pipe }) |pp| pp.deinit(); + t.fence.deinit(); + t.landed.deinit(); + t.posted.deinit(); + t.ready.deinit(); t.lists.deinit(); t.wbuf.deinit(); t.link.deinit(); @@ -153,9 +206,9 @@ pub const Ep = struct { gpa.destroy(t); } - /// GPU waits that gave up (a peer that never answered); nonzero means the replies since are invalid. - pub fn gaveUp(t: *const Ep) u32 { - return @atomicLoad(u32, t.word32(GAVE_UP), .acquire); + /// Arm the relay's post deadline only after Metal has a committed command buffer. + pub fn submit(t: *Ep) void { + t.submitted.store(t.x, .release); } fn list(t: *const Ep, off: usize) Ref { @@ -188,26 +241,23 @@ pub const Ep = struct { } /// After this Mac's routed outputs are in `ye` [rows * TOPK, D]: pack them and post; the host sends them. - pub fn send(t: *Ep, enc: mtl.ComputeEncoder, ye: Ref, rows: u32, pack: bool, post: bool) void { + pub fn send(t: *Ep, pass: *Pass, ye: Ref, rows: u32, pack: bool) void { + const enc = pass.enc; if (pack) { enc.setPipeline(t.pack_pipe); enc.setBuffer(ye.buf, ye.off, 0); enc.setBuffer(t.lists, MINE, 1); enc.setBuffer(t.lists, COUNTS, 2); enc.setBuffer(t.wbuf, sendAt(t.x), 3); - t.slotGuard(enc, 4); + t.bindPrevious(enc, 4); enc.dispatchGroups(mtl.Size.of(1, rows * TOPK, 1), mtl.Size.of(256, 1, 1)); } - if (!post) return; - enc.setPipeline(t.post_pipe); - enc.setBuffer(t.wbuf, POSTED, 0); - enc.setValue(@as(u32, @truncate(t.x)), 1); - enc.setBuffer(t.wbuf, GAVE_UP, 2); // after a give-up nothing more is sent: the run is already invalid - enc.dispatchThreads(mtl.Size.of(1, 1, 1), mtl.Size.of(1, 1, 1)); + t.post(pass); } /// By rows: our routed sums (`yp` partials weighted by `wts`) to the send slot, then the post (window or chunk). - pub fn sendRows(t: *Ep, enc: mtl.ComputeEncoder, yp: Ref, wts: Ref, rows: u32) void { + pub fn sendRows(t: *Ep, pass: *Pass, yp: Ref, wts: Ref, rows: u32) void { + const enc = pass.enc; std.debug.assert(rows <= PROMPT_ROWS); enc.setPipeline(t.rcombine_pipe); enc.setBuffer(yp.buf, yp.off, 0); @@ -215,24 +265,49 @@ pub const Ep = struct { enc.setValue(@as(i32, @intCast(rows)), 2); enc.setBuffer(t.wbuf, sendAt(t.x), 3); enc.setBuffer(t.wbuf, COUNT + 4 * @as(usize, @intCast(t.x % 2)), 4); - t.slotGuard(enc, 5); + t.bindPrevious(enc, 5); enc.dispatchThreads(mtl.Size.of(rows * 4096, 1, 1), mtl.Size.of(256, 1, 1)); - t.send(enc, yp, rows, false, true); + t.post(pass); + } + + /// GPU → host: the count and send slot are complete at this signal. + fn post(t: *const Ep, pass: *Pass) void { + t.handoff(pass, true); + } + + /// Host → GPU: wait before consuming, including before any later exchange can reuse a parity slot. + pub fn wait(t: *const Ep, pass: *Pass) void { + t.handoff(pass, false); + } + + fn handoff(t: *const Ep, pass: *Pass, sending: bool) void { + const mode = pass.enc.dispatchType(); + t.fence.update(pass.enc); + pass.enc.end(); + if (sending) pass.cb.signal(t.posted, t.x) else pass.cb.waitFor(t.landed, t.x); + pass.enc = pass.cb.compute(mode); + t.fence.wait(pass.enc); + } + + fn readyAt(x: u64) usize { + return @as(usize, @intCast(x % 2)) * READY_STRIDE; } - /// The send slot guard's operands (ep_slot_free) from buffer `first`: the flag, this exchange, the give-up count. - fn slotGuard(t: *const Ep, enc: mtl.ComputeEncoder, first: usize) void { - enc.setBuffer(t.wbuf, FLAG, first); - enc.setValue(@as(u32, @truncate(t.x)), first + 1); - enc.setBuffer(t.wbuf, GAVE_UP, first + 2); + fn bindReady(t: *const Ep, enc: mtl.ComputeEncoder) void { + enc.setBuffer(t.ready, readyAt(t.x), 0); + enc.setValue(t.x, 1); + } + + /// A terminal wake must not let a later producer overwrite a slot still in use by the transport. + fn bindPrevious(t: *const Ep, enc: mtl.ComputeEncoder, first: usize) void { + enc.setBuffer(t.ready, readyAt(t.x - 1), first); + enc.setValue(t.x - 1, first + 1); } /// By rows: once the peer's sums land, each row's branch: both sums in rank order, rounded, plus the shared `ys`. pub fn receiveRows(t: *Ep, enc: mtl.ComputeEncoder, ys: Ref, out: Ref, rows: u32) void { enc.setPipeline(t.rfinal_pipe); - enc.setBuffer(t.wbuf, FLAG, 0); - enc.setValue(@as(u32, @truncate(t.x)), 1); - enc.setBuffer(t.wbuf, GAVE_UP, 2); + t.bindReady(enc); enc.setBuffer(t.wbuf, sendAt(t.x), 3); enc.setBuffer(t.wbuf, recvAt(t.x), 4); enc.setValue(t.rank, 5); @@ -243,21 +318,20 @@ pub const Ep = struct { } /// TP: this Mac's fp32 partials `part` [rows, 4096] and the peer's, added in rank order and rounded into `out`. - pub fn reduce(t: *Ep, enc: mtl.ComputeEncoder, part: Ref, out: Ref, rows: u32) void { + pub fn reduce(t: *Ep, pass: *Pass, part: Ref, out: Ref, rows: u32) void { std.debug.assert(rows <= PROMPT_ROWS); t.begin(); - enc.setPipeline(t.rsend_pipe); - enc.setBuffer(part.buf, part.off, 0); - enc.setValue(@as(i32, @intCast(rows)), 1); - enc.setBuffer(t.wbuf, sendAt(t.x), 2); - enc.setBuffer(t.wbuf, COUNT + 4 * @as(usize, @intCast(t.x % 2)), 3); - t.slotGuard(enc, 4); - enc.dispatchThreads(mtl.Size.of(rows * 4096, 1, 1), mtl.Size.of(256, 1, 1)); - t.send(enc, part, rows, false, true); + pass.enc.setPipeline(t.rsend_pipe); + pass.enc.setBuffer(part.buf, part.off, 0); + pass.enc.setValue(@as(i32, @intCast(rows)), 1); + pass.enc.setBuffer(t.wbuf, sendAt(t.x), 2); + pass.enc.setBuffer(t.wbuf, COUNT + 4 * @as(usize, @intCast(t.x % 2)), 3); + t.bindPrevious(pass.enc, 4); + pass.enc.dispatchThreads(mtl.Size.of(rows * 4096, 1, 1), mtl.Size.of(256, 1, 1)); + t.exchange(pass); + const enc = pass.enc; enc.setPipeline(t.rsum_pipe); - enc.setBuffer(t.wbuf, FLAG, 0); - enc.setValue(@as(u32, @truncate(t.x)), 1); - enc.setBuffer(t.wbuf, GAVE_UP, 2); + t.bindReady(enc); enc.setBuffer(t.wbuf, sendAt(t.x), 3); enc.setBuffer(t.wbuf, recvAt(t.x), 4); enc.setValue(t.rank, 5); @@ -267,22 +341,21 @@ pub const Ep = struct { } /// TP: each row's pick from this Mac's vocabulary rows [lo, lo + n) (local picks in `picks`) and the peer's. - pub fn argmax(t: *Ep, enc: mtl.ComputeEncoder, logits: Ref, picks: Ref, n: u32, lo: u32, rows: u32) void { + pub fn argmax(t: *Ep, pass: *Pass, logits: Ref, picks: Ref, n: u32, lo: u32, rows: u32) void { std.debug.assert(rows <= 64); t.begin(); - enc.setPipeline(t.amax_send_pipe); - enc.setBuffer(logits.buf, logits.off, 0); - enc.setBuffer(picks.buf, picks.off, 1); - enc.setValue([4]u32{ rows, n, lo, 0 }, 2); - enc.setBuffer(t.wbuf, sendAt(t.x), 3); - enc.setBuffer(t.wbuf, COUNT + 4 * @as(usize, @intCast(t.x % 2)), 4); - t.slotGuard(enc, 5); - enc.dispatchThreads(mtl.Size.of(64, 1, 1), mtl.Size.of(64, 1, 1)); - t.send(enc, picks, rows, false, true); + pass.enc.setPipeline(t.amax_send_pipe); + pass.enc.setBuffer(logits.buf, logits.off, 0); + pass.enc.setBuffer(picks.buf, picks.off, 1); + pass.enc.setValue([4]u32{ rows, n, lo, 0 }, 2); + pass.enc.setBuffer(t.wbuf, sendAt(t.x), 3); + pass.enc.setBuffer(t.wbuf, COUNT + 4 * @as(usize, @intCast(t.x % 2)), 4); + t.bindPrevious(pass.enc, 5); + pass.enc.dispatchThreads(mtl.Size.of(64, 1, 1), mtl.Size.of(64, 1, 1)); + t.exchange(pass); + const enc = pass.enc; enc.setPipeline(t.amax_merge_pipe); - enc.setBuffer(t.wbuf, FLAG, 0); - enc.setValue(@as(u32, @truncate(t.x)), 1); - enc.setBuffer(t.wbuf, GAVE_UP, 2); + t.bindReady(enc); enc.setBuffer(t.wbuf, sendAt(t.x), 3); enc.setBuffer(t.wbuf, recvAt(t.x), 4); enc.setValue(rows, 5); @@ -290,12 +363,10 @@ pub const Ep = struct { enc.dispatchThreads(mtl.Size.of(64, 1, 1), mtl.Size.of(64, 1, 1)); } - /// Wait for the peer's outputs of this exchange and put them in their picks' rows of `ye`. + /// After the host's arrival event, put the peer's outputs in their picks' rows of `ye`. pub fn receive(t: *Ep, enc: mtl.ComputeEncoder, ye: Ref, rows: u32) void { enc.setPipeline(t.unpack_pipe); - enc.setBuffer(t.wbuf, FLAG, 0); - enc.setValue(@as(u32, @truncate(t.x)), 1); - enc.setBuffer(t.wbuf, GAVE_UP, 2); + t.bindReady(enc); enc.setBuffer(t.wbuf, recvAt(t.x), 3); enc.setBuffer(t.lists, THEIRS, 4); enc.setBuffer(t.lists, COUNTS, 5); @@ -303,6 +374,17 @@ pub const Ep = struct { enc.dispatchGroups(mtl.Size.of(1, rows * TOPK, 1), mtl.Size.of(256, 1, 1)); } + /// For reductions without independent GPU work between the post and the wait, use a single encoder boundary. + fn exchange(t: *const Ep, pass: *Pass) void { + const mode = pass.enc.dispatchType(); + t.fence.update(pass.enc); + pass.enc.end(); + pass.cb.signal(t.posted, t.x); + pass.cb.waitFor(t.landed, t.x); + pass.enc = pass.cb.compute(mode); + t.fence.wait(pass.enc); + } + fn word32(t: *const Ep, off: usize) *u32 { return @ptrCast(@alignCast(t.win.ptr + off)); } @@ -311,66 +393,127 @@ pub const Ep = struct { return @ptrCast(@alignCast(t.win.ptr + off)); } - /// Each posted exchange's entries to the peer, then its flag; 10 s with no answer fails the link, lands every flag. + /// GPU post, write to peer, CPU arrival, then GPU release. One exchange at a time protects both parity slots. fn service(t: *Ep) void { - _ = platform.c.pthread_set_qos_class_self_np(.USER_INTERACTIVE, 0); // busy CPUs must not hold back a send - const posted = t.word32(POSTED); + _ = platform.c.pthread_set_qos_class_self_np(.USER_INTERACTIVE, 0); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); const flag = t.word64(FLAG); var x: u64 = 1; - var waiting = false; // an exchange sent whose peer entries have not landed yet - var want: u64 = 0; - var want_at: u64 = 0; - while (true) { - var spins: usize = 0; - while (!reached(@atomicLoad(u32, posted, .acquire), x)) { + while (!t.stop.load(.acquire)) : (x += 1) { + var began: ?u64 = null; + while (t.posted.value() < x) { t.st.watch(@atomicLoad(u64, flag, .acquire), x - 1); - if (t.stop.load(.acquire)) return; - if (waiting and @atomicLoad(u64, flag, .acquire) >= want) waiting = false; - if (waiting and !t.failed.load(.acquire) and platform.c.mach_absolute_time() - want_at > 240_000_000) { - t.fail(want, error.PeerSilent); - t.land(want); - waiting = false; + if (t.stop.load(.acquire) or t.failed.load(.acquire)) return; + if (t.submitted.load(.acquire) < x) { + sleep(20_000); + } else { + const now = platform.c.mach_absolute_time(); + if (began == null) began = now; + if (now -% began.? > post_timeout) return t.fail(x, error.GpuPostTimeout); + std.atomic.spinLoopHint(); } - spins += 1; - if (spins > 1_000_000) { // idle: back off - const ts: platform.c.timespec = .{ .sec = 0, .nsec = 20_000 }; - _ = platform.c.nanosleep(&ts, null); - } else std.atomic.spinLoopHint(); } + if (t.failed.load(.acquire)) return; + if (t.posted.value() != x) return t.fail(x, error.GpuPostOutOfStep); + const deadline = t.armTransfer(x) orelse return; const seen = platform.c.mach_absolute_time(); const n: usize = @atomicLoad(u32, t.word32(COUNT + 4 * @as(usize, @intCast(x % 2))), .acquire); - const theirs: u32 = @bitCast(@as([*]const i32, @ptrCast(@alignCast(t.lists.contents() + COUNTS)))[1]); + const bytes = n * ENTRY; + if (bytes > SLOT) return t.fail(x, error.BadCount); + const counts: [*]const u32 = @ptrCast(@alignCast(t.lists.contents() + COUNTS)); + const theirs = if (t.by_experts) counts[1] else @as(u32, @intCast(n)); t.st.posted(x, seen, @intCast(n), theirs); t.st.watch(@atomicLoad(u64, flag, .acquire), x); if (t.trace) std.debug.print("EP rank{d} exchange {d}: {d} entries\n", .{ t.rank, x, n }); - if (t.delay_ns > 0 and n * ENTRY > 1 << 20) { - const ts: platform.c.timespec = .{ .sec = @intCast(t.delay_ns / std.time.ns_per_s), .nsec = @intCast(t.delay_ns % std.time.ns_per_s) }; - _ = platform.c.nanosleep(&ts, null); - } - if (t.failed.load(.acquire) or n * ENTRY > SLOT) { - if (n * ENTRY > SLOT) t.fail(x, error.BadCount); - t.land(x); - } else t.link.writeSignalFrom(t.peer, sendAt(x), recvAt(x), n * ENTRY, FLAG, x) catch |err| { - t.fail(x, err); - t.land(x); - }; - want = x; - want_at = platform.c.mach_absolute_time(); - waiting = true; + if (t.delay_ns > 0 and bytes > 1 << 20) t.delay(); + if (t.stop.load(.acquire) or t.failed.load(.acquire)) return; + t.link.writeSignalFromUntil(t.peer, sendAt(x), recvAt(x), bytes, FLAG, x, deadline) catch |err| return t.fail(x, err); + t.link.flushPeerUntil(t.peer, deadline) catch |err| return t.fail(x, err); + const sent_at = platform.c.mach_absolute_time(); t.sent += 1; - t.held_ticks += want_at - seen; - if (n * ENTRY > 1 << 20) { + t.held_ticks += sent_at -% seen; + if (bytes > 1 << 20) { t.big_sent += 1; - t.big_held_ticks += want_at - seen; + t.big_held_ticks += sent_at -% seen; } - x += 1; + while (true) { + const arrived = @atomicLoad(u64, flag, .acquire); + t.st.watch(arrived, x); + if (t.stop.load(.acquire) or t.failed.load(.acquire)) return; + // A fast peer can already have posted x + 1. It needs our x + 1 before it can reuse x's slot. + if (arrived > x + 1) return t.fail(x, error.EpOutOfStep); + if (arrived >= x) break; + std.atomic.spinLoopHint(); + } + if (!t.publishArrival(x)) return; + if (t.trace) std.debug.print("EP rank{d} exchange {d}: peer payload landed\n", .{ t.rank, x }); } } - /// The flag the GPU waits on, set here as if the peer's message had landed (its bytes stale). - fn land(t: *Ep, x: u64) void { - const flag = t.word64(FLAG); - if (@atomicLoad(u64, flag, .acquire) < x) @atomicStore(u64, flag, x, .release); + fn lockPublication(t: *Ep) void { + while (!t.publication.tryLock()) std.atomic.spinLoopHint(); + } + + fn armTransfer(t: *Ep, x: u64) ?u64 { + const now = fabric.words.checkedNowNs() catch |err| { + t.fail(x, err); + return null; + }; + t.lockPublication(); + defer t.publication.unlock(); + if (t.failed.load(.acquire)) return null; + const deadline = now +| transfer_budget_ns; + t.transfer = .{ .exchange = x, .deadline = deadline }; + return deadline; + } + + fn publishArrival(t: *Ep, x: u64) bool { + t.lockPublication(); + defer t.publication.unlock(); + if (t.failed.load(.acquire)) return false; + const now = fabric.words.checkedNowNs() catch |err| { + t.failLocked(x, err); + return false; + }; + if (now >= t.transfer.?.deadline) { + t.failLocked(x, error.EpTransferTimeout); + return false; + } + const ready: *u64 = @ptrCast(@alignCast(t.ready.contents() + readyAt(x))); + ready.* = x; + t.landed.set(x); + t.transfer = null; + return true; + } + + /// This thread can release GPU waits even while the relay is inside a native transport call. + fn watchTransfers(t: *Ep) void { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + while (!t.stop.load(.acquire) and !t.failed.load(.acquire)) { + sleep(std.time.ns_per_ms); + t.lockPublication(); + defer t.publication.unlock(); + if (t.transfer) |pending| { + const now = fabric.words.checkedNowNs() catch |err| return t.failLocked(pending.exchange, err); + if (now >= pending.deadline) return t.failLocked(pending.exchange, error.EpTransferTimeout); + } + } + } + + fn delay(t: *const Ep) void { + var left = t.delay_ns; + while (left > 0 and !t.stop.load(.acquire) and !t.failed.load(.acquire)) { + const step = @min(left, 10 * std.time.ns_per_ms); + sleep(step); + left -= step; + } + } + + fn sleep(ns: u64) void { + const ts: platform.c.timespec = .{ .sec = @intCast(ns / std.time.ns_per_s), .nsec = @intCast(ns % std.time.ns_per_s) }; + _ = platform.c.nanosleep(&ts, null); } /// The host's view of each exchange: when the GPU posted it and when the peer's entries landed (64 in flight). @@ -422,8 +565,21 @@ pub const Ep = struct { }; fn fail(t: *Ep, x: u64, err: anyerror) void { - std.log.err("expert parallel rank {d}: exchange {d} failed: {s}", .{ t.rank, x, @errorName(err) }); - t.failed.store(true, .release); + t.lockPublication(); + defer t.publication.unlock(); + t.failLocked(x, err); + } + + fn failLocked(t: *Ep, x: u64, err: anyerror) void { + if (!t.failed.swap(true, .acq_rel)) { + t.transfer = null; + t.landed.set(failed_event); + std.log.err("expert parallel rank {d}: exchange {d} failed: {s}", .{ t.rank, x, @errorName(err) }); + } + } + + pub fn gpuFailed(t: *Ep) void { + t.fail(t.posted.value() + 1, error.GpuFailed); } }; @@ -432,15 +588,7 @@ test "the window's regions tile without overlap, slots by parity, a page of room try std.testing.expect(recvAt(1) + SLOT <= SEND and recvAt(2) == RECV); try std.testing.expect(sendAt(1) - PAGE >= sendAt(2) + SLOT and sendAt(2) == SEND + PAGE); try std.testing.expect(sendAt(1) + SLOT <= FLAGS and FLAG + 8 <= SYNC); - try std.testing.expect(GAVE_UP + 4 <= CONTROL and COUNT + 8 <= GAVE_UP and CONTROL % PAGE == 0); + try std.testing.expect(COUNT + 8 <= CONTROL and CONTROL % PAGE == 0); try std.testing.expectEqual(@as(usize, 4096), D); try std.testing.expect(COUNTS + 8 <= LISTS); } - -test "the posted word reaches an exchange across the 32-bit wrap" { - try std.testing.expect(reached(5, 5) and reached(6, 5) and !reached(4, 5)); - const top: u64 = 0xffff_fffe; - try std.testing.expect(reached(0xffff_fffe, top) and !reached(0xffff_fffd, top)); - try std.testing.expect(reached(0, top + 2) and reached(1, top + 2) and !reached(0xffff_ffff, top + 2)); // 2^32 posts as 0 - try std.testing.expect(reached(3, (1 << 33) + 3) and !reached(2, (1 << 33) + 3)); -} diff --git a/zig/src/families/glm/forward.zig b/zig/src/families/glm/forward.zig index a8f304852..d2ba487f6 100644 --- a/zig/src/families/glm/forward.zig +++ b/zig/src/families/glm/forward.zig @@ -1,6 +1,7 @@ -//! A 1-16 row window through GLM-5.3-Flash on one serial encoder, each op in the Python family's decode arithmetic. +//! A 1-16 row window through GLM-5.3-Flash, each op in the Python family's decode arithmetic. const std = @import("std"); const mtl = @import("metal"); +const Pass = @import("pass.zig").Pass; const cfg = @import("config.zig"); const wts = @import("weights.zig"); const st = @import("state.zig"); @@ -212,7 +213,8 @@ pub fn boundary(x: *Ctx, e: mtl.ComputeEncoder, rows: u32, pending: bool, hc: ?w } /// KDA layer `ki` on `normed`: the stacked projection (kept for replay), each stream's fused step, the out-projection. -fn kda(x: *Ctx, e: mtl.ComputeEncoder, ki: usize, w: *const wts.Kda, rows: u32) void { +fn kda(x: *Ctx, pass: *Pass, ki: usize, w: *const wts.Kda, rows: u32) void { + const e = pass.enc; const sc = x.sc; const tp = x.c.tp > 1; // TP2: this Mac's heads, their out-projection partial summed with the peer's const proj = x.s.kda[ki].proj; // every stream's state shares it: the window's rows @@ -228,7 +230,7 @@ fn kda(x: *Ctx, e: mtl.ComputeEncoder, ki: usize, w: *const wts.Kda, rows: u32) if (!on(x, "kda_out")) return; if (!tp) return qmv(x, e, x.k.qmv_kda_out, sc.y, w.o_proj, sc.branch, rows); qmv(x, e, x.k.qmvp_kda_out, sc.y, w.o_proj, sc.yp, rows); - x.ep.?.reduce(e, sc.yp, sc.branch, rows); + x.ep.?.reduce(pass, sc.yp, sc.branch, rows); } /// The fused KDA step over `rows` of stacked projections `proj`: state and conv window from slot cur to the other. @@ -263,9 +265,9 @@ pub fn kdaStep(x: *const Ctx, e: mtl.ComputeEncoder, ki: usize, w: *const wts.Kd } /// MLA layer `mi` (its index among MLA caches) for rows at positions pos.. on `x_in`, into `branch`. -pub fn mla(x: *Ctx, e: mtl.ComputeEncoder, mi: usize, w: *const wts.Mla, x_in: Ref, rows: u32, pos: u32) void { - mlaKeys(x, e, mi, w, x_in, rows, pos); - mlaAttend(x, e, mi, w, rows, pos); +pub fn mla(x: *Ctx, pass: *Pass, mi: usize, w: *const wts.Mla, x_in: Ref, rows: u32, pos: u32) void { + mlaKeys(x, pass.enc, mi, w, x_in, rows, pos); + mlaAttend(x, pass, mi, w, rows, pos); } /// The rows' x_proj into `xp` and each segment's key, indexer and pool writes into its cache (and `iw`). @@ -281,7 +283,8 @@ pub fn mlaKeys(x: *Ctx, e: mtl.ComputeEncoder, mi: usize, w: *const wts.Mla, x_i } /// The rows' queries (from `xp`), attention over each segment's cache as written, unabsorb and o_proj into `branch`. -pub fn mlaAttend(x: *Ctx, e: mtl.ComputeEncoder, mi: usize, w: *const wts.Mla, rows: u32, pos: u32) void { +pub fn mlaAttend(x: *Ctx, pass: *Pass, mi: usize, w: *const wts.Mla, rows: u32, pos: u32) void { + const e = pass.enc; const c = x.c; const k = x.k; const sc = x.sc; @@ -296,7 +299,7 @@ pub fn mlaAttend(x: *Ctx, e: mtl.ComputeEncoder, mi: usize, w: *const wts.Mla, r if (!on(x, "mla_out")) return; if (c.tp == 1) return qmv(x, e, k.qmv_mla_out, sc.vals, w.o_proj, sc.branch, rows); qmv(x, e, k.qmvp_mla_out, sc.vals, w.o_proj, sc.yp, rows); // TP2: this Mac's heads' partial, summed with the peer's - x.ep.?.reduce(e, sc.yp, sc.branch, rows); + x.ep.?.reduce(pass, sc.yp, sc.branch, rows); } /// Rows' indexer weights from their x_proj outputs (a row `xp_stride` apart), into `iw`. @@ -448,7 +451,8 @@ pub fn attendIndexed(x: *const Ctx, e: mtl.ComputeEncoder, mi: usize, ql: Ref, i e.dispatchThreads(size(1024, n * c.mla_heads, 1), size(1024, 1, 1)); } -fn denseMlp(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Dense, x_in: Ref, rows: u32) void { +fn denseMlp(x: *const Ctx, pass: *Pass, w: *const wts.Dense, x_in: Ref, rows: u32) void { + const e = pass.enc; if (!on(x, "dense")) return; const c = x.c; const sc = x.sc; @@ -459,63 +463,64 @@ fn denseMlp(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Dense, x_in: Ref e.dispatchThreads(size(c.dense_inter, rows, 1), size(256, 1, 1)); if (c.tp == 1) return qmv(x, e, x.k.qmv_dense_down, sc.actd, w.down, sc.branch, rows); qmv(x, e, x.k.qmvp_dense_down, sc.actd, w.down, sc.yp, rows); // TP2: this Mac's half, summed with the peer's - x.ep.?.reduce(e, sc.yp, sc.branch, rows); + x.ep.?.reduce(pass, sc.yp, sc.branch, rows); } /// The backbone over the window (its tokens in `ids`) at positions pos..: final-normed rows into `hidden`. -pub fn backbone(x: *Ctx, e: mtl.ComputeEncoder, ids: Ref, rows: u32, pos: u32) void { +pub fn backbone(x: *Ctx, pass: *Pass, ids: Ref, rows: u32, pos: u32) void { const c = x.c; const sc = x.sc; const k = x.skip; - if (k & Class.ends == 0) embed(x, e, ids, rows, pos); + if (k & Class.ends == 0) embed(x, pass.enc, ids, rows, pos); const plane = @as(usize, rows) * c.hidden * 2; - snap(x, e, sc.h, plane); + snap(x, pass.enc, sc.h, plane); var pending = false; var ki: usize = 0; var mi: usize = 0; for (0..c.run) |li| { const L = &x.w.layers[li]; const hcs = L.hc.?; - if (k & Class.hc == 0) boundary(x, e, rows, pending, hcs[0], L.in_norm); - snap(x, e, sc.normed, plane); + if (k & Class.hc == 0) boundary(x, pass.enc, rows, pending, hcs[0], L.in_norm); + snap(x, pass.enc, sc.normed, plane); switch (L.attn) { .kda => |*a| { - if (k & Class.kda == 0) kda(x, e, ki, a, rows); + if (k & Class.kda == 0) kda(x, pass, ki, a, rows); ki += 1; }, .mla => |*a| { - if (k & Class.mla == 0) mla(x, e, mi, a, sc.normed, rows, pos); + if (k & Class.mla == 0) mla(x, pass, mi, a, sc.normed, rows, pos); mi += 1; }, } - snap(x, e, sc.branch, plane); - if (k & Class.hc == 0) boundary(x, e, rows, true, hcs[1], L.post_norm); - snap(x, e, sc.normed, plane); + snap(x, pass.enc, sc.branch, plane); + if (k & Class.hc == 0) boundary(x, pass.enc, rows, true, hcs[1], L.post_norm); + snap(x, pass.enc, sc.normed, plane); switch (L.mlp) { - .dense => |*d| if (k & Class.dense == 0) denseMlp(x, e, d, sc.normed, rows), - .moe => |*m| moe(x, e, m, sc.normed, rows), + .dense => |*d| if (k & Class.dense == 0) denseMlp(x, pass, d, sc.normed, rows), + .moe => |*m| moe(x, pass, m, sc.normed, rows), } - snap(x, e, sc.branch, plane); + snap(x, pass.enc, sc.branch, plane); pending = true; } - if (k & Class.hc == 0) boundary(x, e, rows, true, null, null); + if (k & Class.hc == 0) boundary(x, pass.enc, rows, true, null, null); if (k & Class.ends == 0) { - e.setPipeline(x.k.stream_mean); - bind(e, 0, .{ sc.x[x.xi], sc.raw }); - e.setValue([2]u32{ c.hidden, rows }, 2); - e.dispatchThreads(size(c.hidden, rows, 1), size(256, 1, 1)); - rms(x, e, sc.raw, x.w.norm, sc.hidden, rows, c.hidden, c.hidden, c.hidden, c.eps); + pass.enc.setPipeline(x.k.stream_mean); + bind(pass.enc, 0, .{ sc.x[x.xi], sc.raw }); + pass.enc.setValue([2]u32{ c.hidden, rows }, 2); + pass.enc.dispatchThreads(size(c.hidden, rows, 1), size(256, 1, 1)); + rms(x, pass.enc, sc.raw, x.w.norm, sc.hidden, rows, c.hidden, c.hidden, c.hidden, c.eps); } - snap(x, e, sc.hidden, plane); + snap(x, pass.enc, sc.hidden, plane); } /// LM head logits and draws for `rows` rows of `in` into `logits` and `picks` (u32): each row its rule's draw. -pub fn head(x: *const Ctx, e: mtl.ComputeEncoder, in: Ref, logits: Ref, picks: Ref, rows: u32) void { - headOver(x, e, in, logits, picks, rows, x.c.vocab, false); +pub fn head(x: *const Ctx, pass: *Pass, in: Ref, logits: Ref, picks: Ref, rows: u32) void { + headOver(x, pass, in, logits, picks, rows, x.c.vocab, false); } /// `head` over the first `vocab` tokens; `greedy` forces the MTP draft head's argmax draws. -pub fn headOver(x: *const Ctx, e: mtl.ComputeEncoder, in: Ref, logits: Ref, picks: Ref, rows: u32, vocab: u32, greedy: bool) void { +pub fn headOver(x: *const Ctx, pass: *Pass, in: Ref, logits: Ref, picks: Ref, rows: u32, vocab: u32, greedy: bool) void { + const e = pass.enc; if (x.skip & Class.head != 0) return; const part = x.c.vocabPart(); // TP2: this Mac's half, its picks merged with the peer's const n: u32 = if (vocab > part[0]) @min(vocab, part[1]) - part[0] else 0; @@ -527,7 +532,7 @@ pub fn headOver(x: *const Ctx, e: mtl.ComputeEncoder, in: Ref, logits: Ref, pick // draw_dispatch splits the rows into runs of one mode: sampled tf_sample_full, else argmax. @import("draw_dispatch.zig").encode(e, x.k.sample, x.k.argmax, logits.buf, logits.off, picks.buf, picks.off, &payload, n); } - if (x.c.tp > 1) x.ep.?.argmax(e, logits, picks, n, part[0], rows); + if (x.c.tp > 1) x.ep.?.argmax(pass, logits, picks, n, part[0], rows); } /// A judged window's first `keep` of `rows` in each KDA layer, replayed from the round's entry state on a rejection. diff --git a/zig/src/families/glm/mirror.zig b/zig/src/families/glm/mirror.zig index 02d9d3be7..99e624ab9 100644 --- a/zig/src/families/glm/mirror.zig +++ b/zig/src/families/glm/mirror.zig @@ -16,6 +16,7 @@ pub const Kind = enum(u32) { begin = 1, chunk, window, keep, draft, release, sav pub fn send(e: *Engine, kind: Kind, words: []const u32) !void { const ep = e.ep orelse return; if (ep.rank != 0) return error.FollowsPeer; + try e.checkPair(); try ep.ctl.sendCommand(@backingInt(kind), words); } diff --git a/zig/src/families/glm/moe.zig b/zig/src/families/glm/moe.zig index 04c074864..81d3fc756 100644 --- a/zig/src/families/glm/moe.zig +++ b/zig/src/families/glm/moe.zig @@ -1,6 +1,7 @@ //! GLM-5.3-Flash's MoE block on a decode window: shared expert, route, routed experts, combine. const std = @import("std"); const mtl = @import("metal"); +const Pass = @import("pass.zig").Pass; const wts = @import("weights.zig"); const st = @import("state.zig"); const fwd = @import("forward.zig"); @@ -16,47 +17,54 @@ const shape = fwd.shape; const size = fwd.size; /// MoE block (shared expert, route, routed experts, combine); expert parallel: our experts, sent, the peer's received. -pub fn moe(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, rows: u32) void { +pub fn moe(x: *const Ctx, pass: *Pass, w: *const wts.Moe, x_in: Ref, rows: u32) void { std.debug.assert(rows <= st.max_rows); const c = x.c; const k = x.k; const sc = x.sc; const top = c.topk; const s = x.skip; + const exchange = s & Class.exchange == 0 and on(x, "x_locpost"); if (x.ep) |ep| if (c.byRows()) { // every expert's half on each Mac: both compute every pick, then swap their sums - if (s & Class.exchange == 0 and on(x, "x_locpost")) ep.begin(); - if (s & Class.route == 0) route(x, e, w, x_in, rows); - if (s & Class.routed == 0) halfExperts(x, e, w, x_in, rows); - if (s & Class.exchange == 0 and on(x, "x_locpost")) ep.sendRows(e, sc.yp, sc.wts, rows); // with its begin - if (s & Class.shared == 0) experts(x, e, w, x_in, rows, 1, .{ sc.none, sc.none, sc.none }); - if (s & Class.combine == 0) ep.receiveRows(e, sc.ys, sc.branch, rows); + if (exchange) ep.begin(); + if (s & Class.route == 0) route(x, pass.enc, w, x_in, rows); + if (s & Class.routed == 0) halfExperts(x, pass.enc, w, x_in, rows); + if (exchange) ep.sendRows(pass, sc.yp, sc.wts, rows); + if (s & Class.shared == 0) experts(x, pass.enc, w, x_in, rows, 1, .{ sc.none, sc.none, sc.none }); + if (exchange) { + ep.wait(pass); + if (s & Class.combine == 0) ep.receiveRows(pass.enc, sc.ys, sc.branch, rows); + } return; }; if (x.ep) |ep| { - if (s & Class.exchange == 0 and on(x, "x_locpost")) ep.begin(); - if (s & Class.route == 0) route(x, e, w, x_in, rows); - if (!x.fused_route and s & Class.exchange == 0 and on(x, "x_locpost")) ep.localize(e, sc.pick, sc.uids, sc.umem, sc.ucount, rows); - if (s & Class.routed == 0) experts(x, e, w, x_in, rows, 2, ep.group()); - if (s & Class.exchange == 0) ep.send(e, sc.ye, rows, on(x, "x_pack") and on(x, "x_locpost"), on(x, "x_locpost")); - if (s & Class.shared == 0) experts(x, e, w, x_in, rows, 1, .{ sc.none, sc.none, sc.none }); - if (s & Class.exchange == 0 and on(x, "x_unpack")) ep.receive(e, sc.ye, rows); + if (exchange) ep.begin(); + if (s & Class.route == 0) route(x, pass.enc, w, x_in, rows); + if (!x.fused_route and exchange) ep.localize(pass.enc, sc.pick, sc.uids, sc.umem, sc.ucount, rows); + if (s & Class.routed == 0) experts(x, pass.enc, w, x_in, rows, 2, ep.group()); + if (exchange) ep.send(pass, sc.ye, rows, on(x, "x_pack")); + if (s & Class.shared == 0) experts(x, pass.enc, w, x_in, rows, 1, .{ sc.none, sc.none, sc.none }); + if (exchange) { + ep.wait(pass); + if (on(x, "x_unpack")) ep.receive(pass.enc, sc.ye, rows); + } } else if (native(w, .both) and s & (Class.shared | Class.routed) == 0 and on(x, "e_gateup") and on(x, "e_down") and on(x, "s_gateup") and on(x, "s_down")) { - if (s & Class.route == 0) route(x, e, w, x_in, rows); - both(x, e, w, x_in, rows); // the shared expert as the gathers' last slot: two launches fewer + if (s & Class.route == 0) route(x, pass.enc, w, x_in, rows); + both(x, pass.enc, w, x_in, rows); // the shared expert as the gathers' last slot: two launches fewer if (s & Class.combine != 0 or !on(x, "combine")) return; - e.setPipeline(k.combine_0); - bind(e, 0, .{ sc.ye, sc.ye, sc.wts }); - shape(e, 3, .{ rows, top }); - bind(e, 4, .{sc.branch}); - e.dispatchThreads(size(rows * c.hidden, 1, 1), size(256, 1, 1)); + pass.enc.setPipeline(k.combine_0); + bind(pass.enc, 0, .{ sc.ye, sc.ye, sc.wts }); + shape(pass.enc, 3, .{ rows, top }); + bind(pass.enc, 4, .{sc.branch}); + pass.enc.dispatchThreads(size(rows * c.hidden, 1, 1), size(256, 1, 1)); return; } else if (x.k.wide != null and wide.moeOne(parts(w)) and s & (Class.shared | Class.routed) == 0 and on(x, "e_gateup") and on(x, "e_down") and on(x, "s_gateup") and on(x, "s_down")) { - if (s & Class.route == 0) route(x, e, w, x_in, rows); + if (s & Class.route == 0) route(x, pass.enc, w, x_in, rows); wide.moe( &x.k.wide.?, - e, + pass.enc, parts(w), c.experts, c.topk, @@ -65,16 +73,16 @@ pub fn moe(x: *const Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, r rows, ); } else { - if (s & Class.shared == 0) experts(x, e, w, x_in, rows, 1, .{ sc.none, sc.none, sc.none }); - if (s & Class.route == 0) route(x, e, w, x_in, rows); - if (s & Class.routed == 0) experts(x, e, w, x_in, rows, 2, .{ sc.uids, sc.umem, sc.ucount }); + if (s & Class.shared == 0) experts(x, pass.enc, w, x_in, rows, 1, .{ sc.none, sc.none, sc.none }); + if (s & Class.route == 0) route(x, pass.enc, w, x_in, rows); + if (s & Class.routed == 0) experts(x, pass.enc, w, x_in, rows, 2, .{ sc.uids, sc.umem, sc.ucount }); } if (s & Class.combine != 0 or !on(x, "combine")) return; - e.setPipeline(k.moe_combine); - bind(e, 0, .{ sc.ys, sc.ye, sc.wts }); - shape(e, 3, .{ rows, top }); - bind(e, 4, .{sc.branch}); - e.dispatchThreads(size(rows * c.hidden, 1, 1), size(256, 1, 1)); + pass.enc.setPipeline(k.moe_combine); + bind(pass.enc, 0, .{ sc.ys, sc.ye, sc.wts }); + shape(pass.enc, 3, .{ rows, top }); + bind(pass.enc, 4, .{sc.branch}); + pass.enc.dispatchThreads(size(rows * c.hidden, 1, 1), size(256, 1, 1)); } fn parts(w: *const wts.Moe) wide.Moe { diff --git a/zig/src/families/glm/mtp.zig b/zig/src/families/glm/mtp.zig index f7b9fbe3f..ffd0f64ea 100644 --- a/zig/src/families/glm/mtp.zig +++ b/zig/src/families/glm/mtp.zig @@ -1,6 +1,7 @@ //! MTP draft head (layer 45): enorm(embed(next)) | hnorm(h), eh_proj, a plain MLA + MoE block, its norm, the LM head. const std = @import("std"); const mtl = @import("metal"); +const Pass = @import("pass.zig").Pass; const fwd = @import("forward.zig"); const wts = @import("weights.zig"); const Ref = wts.Ref; @@ -15,16 +16,16 @@ fn add(x: *const fwd.Ctx, e: mtl.ComputeEncoder, a: Ref, b: Ref, out: Ref, n: u3 } /// `rows` rows of h with next tokens at head positions pos..: output rows into m_x (pre-norm), drafts into `picks`. -pub fn run(x: *fwd.Ctx, e: mtl.ComputeEncoder, h: Ref, next: Ref, rows: u32, pos: u32, picks: ?Ref) void { - layer(x, e, h, next, rows, pos); +pub fn run(x: *fwd.Ctx, pass: *Pass, h: Ref, next: Ref, rows: u32, pos: u32, picks: ?Ref) void { + layer(x, pass, h, next, rows, pos); const out = picks orelse return; // the rows absorbed into the head's cache only, no draft - drafts(x, e, x.sc.m_x.at(@as(usize, rows - 1) * x.c.hidden * 2), 1, out); + drafts(x, pass, x.sc.m_x.at(@as(usize, rows - 1) * x.c.hidden * 2), 1, out); } /// The head's block over `rows` rows of h and their next tokens (positions pos.., or each segment's): m_x, pre-norm. -pub fn layer(x: *fwd.Ctx, e: mtl.ComputeEncoder, h: Ref, next: Ref, rows: u32, pos: u32) void { - absorb(x, e, h, next, rows, pos); - finish(x, e, x.sc.m_x, x.sc.m_xn, rows, pos, false); +pub fn layer(x: *fwd.Ctx, pass: *Pass, h: Ref, next: Ref, rows: u32, pos: u32) void { + absorb(x, pass.enc, h, next, rows, pos); + finish(x, pass, x.sc.m_x, x.sc.m_xn, rows, pos, false); } /// The rows' entries in the head's cache: embed, norms and eh_proj into m_x, its norm into m_xn, then the key writes. @@ -46,34 +47,34 @@ pub fn absorb(x: *fwd.Ctx, e: mtl.ComputeEncoder, h: Ref, next: Ref, rows: u32, } /// The rest of the block for absorbed rows (`res`: m_x's, `normed`: m_xn's; `gathered`: their x_proj again): m_x. -pub fn finish(x: *fwd.Ctx, e: mtl.ComputeEncoder, res: Ref, normed: Ref, rows: u32, pos: u32, gathered: bool) void { +pub fn finish(x: *fwd.Ctx, pass: *Pass, res: Ref, normed: Ref, rows: u32, pos: u32, gathered: bool) void { const c = x.c; const sc = x.sc; const D = c.hidden; const L = &x.w.layers[c.layers]; const mi = c.countKind(.mla); if (gathered) { // the rows' queries and indexer weights, as the absorb computed them where the rows were - fwd.qmv(x, e, x.k.qmv_x, normed, L.attn.mla.x_proj, sc.xp, rows); - fwd.indexWeights(x, e, sc.xp, c.xProj(), sc.iw, rows); + fwd.qmv(x, pass.enc, x.k.qmv_x, normed, L.attn.mla.x_proj, sc.xp, rows); + fwd.indexWeights(x, pass.enc, sc.xp, c.xProj(), sc.iw, rows); } - fwd.mlaAttend(x, e, mi, &L.attn.mla, rows, pos); - add(x, e, res, sc.branch, sc.m_out, rows * D); - fwd.rms(x, e, sc.m_out, L.post_norm, sc.m_xn, rows, D, D, D, c.eps); + fwd.mlaAttend(x, pass, mi, &L.attn.mla, rows, pos); + add(x, pass.enc, res, sc.branch, sc.m_out, rows * D); + fwd.rms(x, pass.enc, sc.m_out, L.post_norm, sc.m_xn, rows, D, D, D, c.eps); var local = x.*; // every Mac holds the MTP layer's experts whole: its MoE is one Mac's, with no exchange local.ep = null; - fwd.moe(&local, e, &L.mlp.moe, sc.m_xn, rows); - add(x, e, sc.m_out, sc.branch, sc.m_x, rows * D); + fwd.moe(&local, pass, &L.mlp.moe, sc.m_xn, rows); + add(x, pass.enc, sc.m_out, sc.branch, sc.m_x, rows * D); } /// Drafts for `rows` pre-norm head outputs at `out`: the head's norm, then the LM head over the draft vocabulary. -pub fn drafts(x: *const fwd.Ctx, e: mtl.ComputeEncoder, out: Ref, rows: u32, picks: Ref) void { +pub fn drafts(x: *const fwd.Ctx, pass: *Pass, out: Ref, rows: u32, picks: Ref) void { const c = x.c; const D = c.hidden; - fwd.rms(x, e, out, x.w.mtp.?.norm, x.sc.m_hn, rows, D, D, D, c.eps); - fwd.headOver(x, e, x.sc.m_hn, x.sc.m_logits.at(@as(usize, x.m_row) * c.vocab * 2), picks, rows, x.draft_vocab, true); + fwd.rms(x, pass.enc, out, x.w.mtp.?.norm, x.sc.m_hn, rows, D, D, D, c.eps); + fwd.headOver(x, pass, x.sc.m_hn, x.sc.m_logits.at(@as(usize, x.m_row) * c.vocab * 2), picks, rows, x.draft_vocab, true); } /// One chained draft from `h` (a row of the head's previous output, m_x) and the draft before it (`token`, u32). -pub fn chain(x: *fwd.Ctx, e: mtl.ComputeEncoder, h: Ref, token: Ref, pos: u32, picks: Ref) void { - run(x, e, h, token, 1, pos, picks); +pub fn chain(x: *fwd.Ctx, pass: *Pass, h: Ref, token: Ref, pos: u32, picks: Ref) void { + run(x, pass, h, token, 1, pos, picks); } diff --git a/zig/src/families/glm/ordered_prompt.zig b/zig/src/families/glm/ordered_prompt.zig index 21d68c28c..c9d93b945 100644 --- a/zig/src/families/glm/ordered_prompt.zig +++ b/zig/src/families/glm/ordered_prompt.zig @@ -1,6 +1,7 @@ //! Chunk projections retain their tiles; MLA and routed experts keep the window operators. const std = @import("std"); const mtl = @import("metal"); +const Pass = @import("pass.zig").Pass; const fwd = @import("forward.zig"); const st = @import("state.zig"); const wts = @import("weights.zig"); @@ -35,7 +36,7 @@ pub fn mla(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, mi: usize if (fwd.on(x, "mla_out")) _ = ordered.encode(x.k, e, w.o_proj, .{ .in = p.vals, .ldx = c.mla_heads * c.v_dim, .out = p.streams.branch, .ldo = c.hidden, .rows = M }); } -pub fn moe(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, in: Ref, M: u32) void { +pub fn moe(p: *const Prompt, x: *const fwd.Ctx, pass: *Pass, w: *const wts.Moe, in: Ref, M: u32) void { var r: u32 = 0; while (r < M) : (r += st.max_rows) { const n = @min(st.max_rows, M - r); @@ -43,6 +44,6 @@ pub fn moe(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, w: *const sc.branch = p.streams.branch.at(@as(usize, r) * x.c.hidden * 2); var y = x.*; y.sc = ≻ - fwd.moe(&y, e, w, in.at(@as(usize, r) * x.c.hidden * 2), n); + fwd.moe(&y, pass, w, in.at(@as(usize, r) * x.c.hidden * 2), n); } } diff --git a/zig/src/families/glm/pass.zig b/zig/src/families/glm/pass.zig new file mode 100644 index 000000000..3ecfde289 --- /dev/null +++ b/zig/src/families/glm/pass.zig @@ -0,0 +1,7 @@ +//! A GLM command buffer and its current compute encoder, which changes at an exchange. +const mtl = @import("metal"); + +pub const Pass = struct { + cb: mtl.CommandBuffer, + enc: mtl.ComputeEncoder, +}; diff --git a/zig/src/families/glm/plain_rounds.zig b/zig/src/families/glm/plain_rounds.zig index 91a074a44..3dd269e50 100644 --- a/zig/src/families/glm/plain_rounds.zig +++ b/zig/src/families/glm/plain_rounds.zig @@ -2,6 +2,7 @@ const std = @import("std"); const platform = @import("sys"); const mtl = @import("metal"); +const Pass = @import("pass.zig").Pass; const fwd = @import("forward.zig"); const Fence = @import("../../core/fence.zig").Fence; const Ref = @import("weights.zig").Ref; @@ -17,7 +18,7 @@ fn plainRound(e: *Engine, x: *fwd.Ctx, from: ?Ref, to: Ref, pos: u32) Pending { const fence = e.round_fence.?; const cb = e.queue.commandBuffer(); if (from == null and e.ev > 0) cb.waitFor(e.event, e.ev); - const b = .{ .cb = cb, .enc = cb.compute(.serial) }; + var b: Pass = .{ .cb = cb, .enc = cb.compute(.serial) }; if (from != null) fence.wait(b.enc); if (from) |f| { b.enc.setPipeline(x.k.copy_u32); @@ -25,12 +26,12 @@ fn plainRound(e: *Engine, x: *fwd.Ctx, from: ?Ref, to: Ref, pos: u32) Pending { b.enc.setValue(@as(u32, 1), 2); b.enc.dispatchThreads(fwd.size(1, 1, 1), fwd.size(1, 1, 1)); } - fwd.backbone(x, b.enc, e.sc.ids, 1, pos); + fwd.backbone(x, &b, e.sc.ids, 1, pos); const draws = @import("draw_rule.zig").promptDraws(e.sampling, pos + 1, e.c.vocab); const prior = x.draws; x.draws = &draws; defer x.draws = prior; - fwd.head(x, b.enc, e.sc.hidden, e.sc.logits, to, 1); + fwd.head(x, &b, e.sc.hidden, e.sc.logits, to, 1); fence.update(b.enc); b.enc.end(); e.ev += 1; diff --git a/zig/src/families/glm/prompt.zig b/zig/src/families/glm/prompt.zig index 6db9b46de..1f11ed0a9 100644 --- a/zig/src/families/glm/prompt.zig +++ b/zig/src/families/glm/prompt.zig @@ -2,6 +2,7 @@ const std = @import("std"); const platform = @import("sys"); const mtl = @import("metal"); +const Pass = @import("pass.zig").Pass; const cfg = @import("config.zig"); const wts = @import("weights.zig"); const st = @import("state.zig"); @@ -246,8 +247,9 @@ fn shared(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, w: *const } /// Shared expert, route, experts gathered by expert; by rows, this Mac's halves summed in slot order and swapped. -fn moe(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, M: u32, ep: ?*ep_mod.Ep) void { - if (x.k.ordered != null) return @import("ordered_prompt.zig").moe(p, x, e, w, x_in, M); +fn moe(p: *const Prompt, x: *const fwd.Ctx, pass: *Pass, w: *const wts.Moe, x_in: Ref, M: u32, ep: ?*ep_mod.Ep) void { + const e = pass.enc; + if (x.k.ordered != null) return @import("ordered_prompt.zig").moe(p, x, pass, w, x_in, M); const c = x.c; const k = x.k; const D = c.hidden; @@ -256,8 +258,9 @@ fn moe(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, w: *const wts const out = p.streams.branch; const s = x.skip; // classes a profile leaves out const Class = fwd.Class; + const exchange = s & Class.exchange == 0; if (ep) |t| { - if (s & Class.exchange == 0) t.begin(); + if (exchange) t.begin(); } else if (s & Class.shared == 0) shared(p, x, e, w, x_in, M); if (s & Class.route == 0) route(p, x, e, w, x_in, M); if (ep) |t| { @@ -265,9 +268,12 @@ fn moe(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, w: *const wts gateUp(p, p.mm.mm_bf16, x, e, w, x_in, M, n); affine_mm.gatherTo(e, p.mm.mm_f32, p.act, p.sums, w.down, p.offsets, p.yf, p.order, n, E); } - if (s & Class.exchange == 0) t.sendRows(e, p.yf, p.wts, M); - if (s & Class.shared == 0) shared(p, x, e, w, x_in, M); - if (s & Class.combine == 0) t.receiveRows(e, p.ys, out, M); + if (exchange) t.sendRows(pass, p.yf, p.wts, M); + if (s & Class.shared == 0) shared(p, x, pass.enc, w, x_in, M); + if (exchange) { + t.wait(pass); + if (s & Class.combine == 0) t.receiveRows(pass.enc, p.ys, out, M); + } return; } const path = wide.routed(w.gate, w.up, w.down); @@ -434,7 +440,8 @@ fn selectKeys(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, mi: us } /// MLA layer `mi` at pos..: tensor-unit projections, decode cache writes, each row's key list in the sparse kernel. -fn mla(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, mi: usize, w: *const wts.Mla, x_in: Ref, M: u32, pos: u32) void { +fn mla(p: *const Prompt, x: *const fwd.Ctx, pass: *Pass, mi: usize, w: *const wts.Mla, x_in: Ref, M: u32, pos: u32) void { + const e = pass.enc; if (x.k.ordered != null) return @import("ordered_prompt.zig").mla(p, x, e, mi, w, x_in, M, pos); const c = x.c; const XP: u32 = @intCast(std.mem.alignForward(usize, c.xProj(), 64)); @@ -490,69 +497,69 @@ fn mla(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, mi: usize, w: if (!fwd.on(x, "mla_out")) return; if (c.tp == 1) return qmm(p, e, p.vals, w.o_proj, p.streams.branch, M); qmmF32(p, e, p.vals, w.o_proj, p.part, M); // TP2: this Mac's heads' partial, summed with the peer's - x.ep.?.reduce(e, p.part, p.streams.branch, M); + x.ep.?.reduce(pass, p.part, p.streams.branch, M); } /// The backbone over a chunk (tokens in `ids`) at positions pos..: final-normed rows into `streams.hidden`. -pub fn backbone(p: *const Prompt, x: *fwd.Ctx, e: mtl.ComputeEncoder, ids: Ref, M: u32, pos: u32) void { +pub fn backbone(p: *const Prompt, x: *fwd.Ctx, pass: *Pass, ids: Ref, M: u32, pos: u32) void { const c = x.c; const ss = &p.streams; std.debug.assert(x.sc == ss and M <= p.rows); const s = x.skip; // classes a profile leaves out const Class = fwd.Class; - if (s & Class.ends == 0) fwd.embed(x, e, ids, M, pos); + if (s & Class.ends == 0) fwd.embed(x, pass.enc, ids, M, pos); const plane = @as(usize, M) * c.hidden * 2; // a trace's capture points: the decode backbone's - fwd.snap(x, e, ss.h, plane); + fwd.snap(x, pass.enc, ss.h, plane); var pending = false; var ki: usize = 0; var mi: usize = 0; for (0..c.run) |li| { const L = &x.w.layers[li]; const hcs = L.hc.?; - if (s & Class.hc == 0) fwd.boundary(x, e, M, pending, hcs[0], L.in_norm); - fwd.snap(x, e, ss.normed, plane); + if (s & Class.hc == 0) fwd.boundary(x, pass.enc, M, pending, hcs[0], L.in_norm); + fwd.snap(x, pass.enc, ss.normed, plane); switch (L.attn) { .kda => |*a| { if (s & Class.kda == 0) { - if (fwd.on(x, "kda_in")) qmm(p, e, ss.normed, a.in_proj, p.proj, M); - kdaChunk(p, x, e, ki, a, M); + if (fwd.on(x, "kda_in")) qmm(p, pass.enc, ss.normed, a.in_proj, p.proj, M); + kdaChunk(p, x, pass.enc, ki, a, M); if (c.tp > 1) { // TP2: this Mac's heads' out-projection partial, summed with the peer's - qmmF32(p, e, p.y, a.o_proj, p.part, M); - x.ep.?.reduce(e, p.part, ss.branch, M); - } else if (fwd.on(x, "kda_out")) qmm(p, e, p.y, a.o_proj, ss.branch, M); + qmmF32(p, pass.enc, p.y, a.o_proj, p.part, M); + x.ep.?.reduce(pass, p.part, ss.branch, M); + } else if (fwd.on(x, "kda_out")) qmm(p, pass.enc, p.y, a.o_proj, ss.branch, M); } ki += 1; }, .mla => |*a| { - if (s & Class.mla == 0) mla(p, x, e, mi, a, ss.normed, M, pos); + if (s & Class.mla == 0) mla(p, x, pass, mi, a, ss.normed, M, pos); mi += 1; }, } - fwd.snap(x, e, ss.branch, plane); - if (s & Class.hc == 0) fwd.boundary(x, e, M, true, hcs[1], L.post_norm); - fwd.snap(x, e, ss.normed, plane); + fwd.snap(x, pass.enc, ss.branch, plane); + if (s & Class.hc == 0) fwd.boundary(x, pass.enc, M, true, hcs[1], L.post_norm); + fwd.snap(x, pass.enc, ss.normed, plane); switch (L.mlp) { .dense => |*d| if (s & Class.dense == 0) { - qmm(p, e, ss.normed, d.gate_up, p.gu, M); - swiglu(x, e, p.gu, p.actd, M, c.dense_inter); + qmm(p, pass.enc, ss.normed, d.gate_up, p.gu, M); + swiglu(x, pass.enc, p.gu, p.actd, M, c.dense_inter); if (c.tp > 1) { // TP2: this Mac's half, summed with the peer's - qmmF32(p, e, p.actd, d.down, p.part, M); - x.ep.?.reduce(e, p.part, ss.branch, M); - } else qmm(p, e, p.actd, d.down, ss.branch, M); + qmmF32(p, pass.enc, p.actd, d.down, p.part, M); + x.ep.?.reduce(pass, p.part, ss.branch, M); + } else qmm(p, pass.enc, p.actd, d.down, ss.branch, M); }, - .moe => |*m| moe(p, x, e, m, ss.normed, M, if (c.byRows()) x.ep else null), + .moe => |*m| moe(p, x, pass, m, ss.normed, M, if (c.byRows()) x.ep else null), } - fwd.snap(x, e, ss.branch, plane); + fwd.snap(x, pass.enc, ss.branch, plane); pending = true; } - if (s & Class.hc == 0) fwd.boundary(x, e, M, true, null, null); + if (s & Class.hc == 0) fwd.boundary(x, pass.enc, M, true, null, null); if (s & Class.ends != 0) return; - e.setPipeline(x.k.stream_mean); - bind(e, 0, .{ ss.x[x.xi], ss.raw }); - e.setValue([2]u32{ c.hidden, M }, 2); - e.dispatchThreads(size(c.hidden, M, 1), size(256, 1, 1)); - fwd.rms(x, e, ss.raw, x.w.norm, ss.hidden, M, c.hidden, c.hidden, c.hidden, c.eps); - fwd.snap(x, e, ss.hidden, plane); + pass.enc.setPipeline(x.k.stream_mean); + bind(pass.enc, 0, .{ ss.x[x.xi], ss.raw }); + pass.enc.setValue([2]u32{ c.hidden, M }, 2); + pass.enc.dispatchThreads(size(c.hidden, M, 1), size(256, 1, 1)); + fwd.rms(x, pass.enc, ss.raw, x.w.norm, ss.hidden, M, c.hidden, c.hidden, c.hidden, c.eps); + fwd.snap(x, pass.enc, ss.hidden, plane); } /// MTP head over M prompt rows (`h`, next tokens) at head positions pos..: cache entries only (the output is unread). diff --git a/zig/src/families/glm/slots.zig b/zig/src/families/glm/slots.zig index 2d62ebec4..d55d6dba2 100644 --- a/zig/src/families/glm/slots.zig +++ b/zig/src/families/glm/slots.zig @@ -9,6 +9,7 @@ const prompt_mod = @import("prompt.zig"); const snapshot = @import("snapshot.zig"); const pool_mod = @import("pool.zig"); const kept_mod = @import("kept.zig"); +const Pass = @import("pass.zig").Pass; const Engine = @import("engine.zig").Engine; const Ref = @import("weights.zig").Ref; const Fence = @import("../../core/fence.zig").Fence; @@ -69,7 +70,7 @@ pub const Slots = struct { fence: ?Fence = null, // orders queued rounds on the GPU pool: ?pool_mod.Pool = null, // one Mac: the streams share the engine's MLA caches (GLM_POOL=0: a window each) - const Open = struct { cb: mtl.CommandBuffer, enc: mtl.ComputeEncoder, pool: mtl.objc.Pool }; + const Open = struct { pass: Pass, pool: mtl.objc.Pool }; /// `n` slots: the engine's caches first (no generate while they serve), the rest sharing its projections. pub fn init(gpa: std.mem.Allocator, e: *Engine, n: u32) !Slots { @@ -122,14 +123,14 @@ pub const Slots = struct { } /// The open command buffer's encoder; its autorelease pool stays pushed until `flush` (one thread drives them). - fn encoder(sl: *Slots) mtl.ComputeEncoder { + fn encoder(sl: *Slots) *Pass { if (sl.open == null) { const pool = mtl.objc.Pool.push(); const b = sl.e.begin(); - if (sl.launched != null) sl.fence.?.wait(b.enc); // after the keeps and drafts committed ahead - sl.open = .{ .cb = b.cb, .enc = b.enc, .pool = pool }; + if (sl.launched != null) sl.fence.?.wait(b.enc); + sl.open = .{ .pass = b, .pool = pool }; } - return sl.open.?.enc; + return &sl.open.?.pass; } /// Commit the open command buffer and wait for it (and for work committed ahead of it). @@ -142,7 +143,13 @@ pub const Slots = struct { const o = sl.open orelse return sl.land(); sl.open = null; defer o.pool.pop(); - const done = sl.e.finishBy(o.cb, o.enc, spin); + if (sl.e.ep != null) { + // Encoding overlaps the drafts, but a failed draft must never be a GPU dependency of this buffer. + errdefer o.pass.enc.end(); + try sl.land(); + try sl.e.checkPair(); + } + const done = sl.e.finishBy(o.pass.cb, o.pass.enc, spin); try sl.land(); return done; } @@ -154,10 +161,11 @@ pub const Slots = struct { if (sl.fence == null) sl.fence = try Fence.init(sl.e.device); sl.open = null; defer o.pool.pop(); - sl.fence.?.update(o.enc); - o.enc.end(); - o.cb.commit(); - sl.launched = o.cb.retain(); + sl.fence.?.update(o.pass.enc); + o.pass.enc.end(); + o.pass.cb.commit(); + if (sl.e.ep) |ep| ep.submit(); + sl.launched = o.pass.cb.retain(); } /// Work committed ahead, finished: its GPU span kept for the window log, its failure returned. @@ -169,8 +177,10 @@ pub const Slots = struct { sl.ahead = .{ cb.gpuStart(), cb.gpuEnd() }; if (cb.failure()) |msg| { std.log.err("glm: keeps and drafts committed ahead failed: {s}", .{msg}); + if (sl.e.ep) |ep| ep.gpuFailed(); return error.GpuFailed; } + try sl.e.checkPair(); } /// Slot `i`'s cache length once its last window settles (a window no keep has settled keeps every row). @@ -230,25 +240,25 @@ pub const Slots = struct { var x = sl.ctx(slot); x.inject = if (e.vision) |v| v.injected(i) else null; // an image prompt's rows over its placeholders (slot_images.zig) if (draws) |d| x.draws = d; // the head's one output row reads rule row 0, keyed at the prompt's length - const b = e.begin(); + var b = e.begin(); const ids = e.prompt_ids.at(@as(usize, at) * 4); const next = e.prompt_ids.at(@as(usize, at + 1) * 4); if (e.pr != null and n > st.max_rows) { const pr = &e.pr.?; var px = x; px.sc = &pr.streams; - prompt_mod.backbone(pr, &px, b.enc, ids, n, at); + prompt_mod.backbone(pr, &px, &b, ids, n, at); fwd.flipKda(&x); if (last) Engine.copyRow(&x, b.enc, pr.streams.hidden.at(@as(usize, n - 1) * D * 2), slot.last, D); if (slot.mtp and absorb > 0) prompt_mod.mtp(pr, &px, b.enc, pr.streams.hidden, next, absorb, at); } else { - fwd.backbone(&x, b.enc, ids, n, at); + fwd.backbone(&x, &b, ids, n, at); fwd.flipKda(&x); if (last) Engine.copyRow(&x, b.enc, e.sc.hidden.at(@as(usize, n - 1) * D * 2), slot.last, D); if (slot.mtp and absorb > 0) mtp.absorb(&x, b.enc, e.sc.hidden, next, absorb, at); } if (slot.mtp and absorb > 0) slot.s.mtp_pos = at + absorb; - if (last) fwd.head(&x, b.enc, slot.last, e.sc.logits, e.sc.picks, 1); + if (last) fwd.head(&x, &b, slot.last, e.sc.logits, e.sc.picks, 1); try e.finish(b.cb, b.enc); slot.s.pos = at + n; // the pass stands here: a snapshot of it is this prefix's state if (last) { @@ -281,7 +291,7 @@ pub const Slots = struct { const ids = Engine.u32s(e.sc.ids, st.max_rows); var x = e.ctx(); if (draws) |d| x.draws = d; // each row's rule is its stream's, keyed at its own position - const enc = sl.encoder(); + const pass = sl.encoder(); var total: u32 = 0; for (wins, segs[0..wins.len], 0..) |w, *g, wi| { const slot = try sl.slotAt(w.slot); @@ -290,7 +300,7 @@ pub const Slots = struct { if (w.held > slot.held_n or total + rows > st.max_rows or slot.s.pos + rows + 1 > slot.s.cap) return error.WindowOutOfStep; ids[total] = w.pending; @memcpy(ids[total + 1 + w.held ..][0..w.tokens.len], w.tokens); - if (w.held > 0) copyWords(&x, enc, slot.held, e.sc.ids.at(@as(usize, total + 1) * 4), w.held); + if (w.held > 0) copyWords(&x, pass.enc, slot.held, e.sc.ids.at(@as(usize, total + 1) * 4), w.held); g.* = .{ .s = &slot.s, .row0 = total, .rows = rows, .pos = slot.s.pos }; slot.row0 = total; slot.rows = rows; @@ -299,8 +309,8 @@ pub const Slots = struct { } x.s = segs[0].s; x.segs = segs[0..wins.len]; - fwd.backbone(&x, enc, e.sc.ids, total, segs[0].pos); - fwd.head(&x, enc, e.sc.hidden, e.sc.logits, e.sc.picks, total); + fwd.backbone(&x, pass, e.sc.ids, total, segs[0].pos); + fwd.head(&x, pass, e.sc.hidden, e.sc.logits, e.sc.picks, total); sl.ahead = .{ 0, 0 }; try sl.flushBy(true); const ret = platform.c.mach_absolute_time(); @@ -334,9 +344,11 @@ pub const Slots = struct { if (cb.failure()) |msg| { sl.e.event.set(sl.e.ev); // a failed buffer may never signal: the next one must not wait on it std.log.err("glm: a queued round failed: {s}", .{msg}); + if (sl.e.ep) |ep| ep.gpuFailed(); return error.GpuFailed; } } + try sl.e.checkPair(); return Engine.u32s(sl.ring, ring_len)[h]; } @@ -365,19 +377,20 @@ pub const Slots = struct { const pool = mtl.objc.Pool.push(); defer pool.pop(); const cb = e.queue.commandBuffer(); - const enc = cb.compute(.serial); - fence.wait(enc); // after the round queued before it; synchronous work has finished before this is encoded + var pass: Pass = .{ .cb = cb, .enc = cb.compute(.serial) }; + fence.wait(pass.enc); // after the round queued before it; synchronous work has finished before this is encoded var x = sl.ctx(slot); const draws = @import("draw_rule.zig").promptDraws(sampling, slot.s.pos + 1, e.c.vocab); x.draws = &draws; - copyWords(&x, enc, sl.ring.at(@as(usize, from) * 4), e.sc.ids, 1); - fwd.backbone(&x, enc, e.sc.ids, 1, slot.s.pos); - fwd.head(&x, enc, e.sc.hidden, e.sc.logits, sl.ring.at(@as(usize, to) * 4), 1); - fence.update(enc); - enc.end(); + copyWords(&x, pass.enc, sl.ring.at(@as(usize, from) * 4), e.sc.ids, 1); + fwd.backbone(&x, &pass, e.sc.ids, 1, slot.s.pos); + fwd.head(&x, &pass, e.sc.hidden, e.sc.logits, sl.ring.at(@as(usize, to) * 4), 1); + fence.update(pass.enc); + pass.enc.end(); e.ev += 1; cb.signal(e.event, e.ev); // a synchronous window after it waits for this cb.commit(); + if (e.ep) |ep| ep.submit(); sl.ring_cb[to] = .{ .id = mtl.objc.retain(cb.id) }; slot.row0 = 0; slot.rows = 1; @@ -397,7 +410,7 @@ pub const Slots = struct { /// A slot's last window settled at its first `kept` rows: rejected rows replayed out of its KDA states, flipped. fn settle(sl: *Slots, slot: *Slot, kept: u32) void { var x = sl.ctx(slot); - if (kept < slot.rows) fwd.keepKdaAt(&x, sl.encoder(), slot.row0, slot.rows, kept); + if (kept < slot.rows) fwd.keepKdaAt(&x, sl.encoder().enc, slot.row0, slot.rows, kept); fwd.flipKda(&x); slot.s.pos += kept; slot.rows = 0; @@ -437,21 +450,21 @@ pub const Slots = struct { const ids = Engine.u32s(sl.ids, st.max_rows); var segs: [st.max_rows]fwd.Seg = undefined; var x = e.ctx(); - const enc = sl.encoder(); + const pass = sl.encoder(); var row: u32 = 0; for (order[0..reqs.len], segs[0..reqs.len]) |k, *g| { const r = reqs[k]; const slot = &sl.slots[r.slot]; const n: u32 = @intCast(r.follow.len); const h = if (r.prompt) slot.last else e.sc.hidden.at(@as(usize, slot.row0) * D * 2); - copyWords(&x, enc, h, sl.rows.at(row * D * 2), @intCast(n * D / 2)); + copyWords(&x, pass.enc, h, sl.rows.at(row * D * 2), @intCast(n * D / 2)); @memcpy(ids[row..][0..n], r.follow); g.* = .{ .s = &slot.s, .row0 = row, .rows = n, .pos = slot.s.mtp_pos }; row += n; } x.s = segs[0].s; x.segs = segs[0..reqs.len]; - mtp.absorb(&x, enc, sl.rows, sl.ids, row, 0); // every row's cache entries; only last rows go on to draft + mtp.absorb(&x, pass.enc, sl.rows, sl.ids, row, 0); // every row's cache entries; only last rows go on to draft var drafting: u32 = 0; var lasts: [st.max_rows]fwd.Seg = undefined; for (order[0..reqs.len], segs[0..reqs.len], 0..) |k, g, i| { @@ -460,18 +473,18 @@ pub const Slots = struct { slot.held_n = reqs[k].depth; if (reqs[k].depth == 0) continue; const at = (g.row0 + g.rows - 1) * D * 2; - copyWords(&x, enc, e.sc.m_x.at(at), sl.lasts.at(i * D * 2), @intCast(D / 2)); - copyWords(&x, enc, e.sc.m_xn.at(at), sl.rows.at(i * D * 2), @intCast(D / 2)); + copyWords(&x, pass.enc, e.sc.m_x.at(at), sl.lasts.at(i * D * 2), @intCast(D / 2)); + copyWords(&x, pass.enc, e.sc.m_xn.at(at), sl.rows.at(i * D * 2), @intCast(D / 2)); lasts[i] = .{ .s = &slot.s, .row0 = @intCast(i), .rows = 1, .pos = slot.s.mtp_pos - 1 }; drafting += 1; } if (drafting == 0) return; x.s = lasts[0].s; x.segs = lasts[0..drafting]; - mtp.finish(&x, enc, sl.lasts, sl.rows, drafting, 0, true); + mtp.finish(&x, pass, sl.lasts, sl.rows, drafting, 0, true); const y = e.ctx(); - mtp.drafts(&y, enc, e.sc.m_x, drafting, sl.picks); - sl.hold(enc, order[0..drafting], reqs, 0); + mtp.drafts(&y, pass, e.sc.m_x, drafting, sl.picks); + sl.hold(pass.enc, order[0..drafting], reqs, 0); var level: u32 = 1; while (true) : (level += 1) { var n: u32 = 0; // the streams drafting past this level: a prefix of `order` @@ -482,9 +495,9 @@ pub const Slots = struct { if (n == 0) break; x.s = segs[0].s; x.segs = segs[0..n]; - mtp.layer(&x, enc, e.sc.m_x, sl.picks, n, 0); // the level before's outputs, a prefix - mtp.drafts(&y, enc, e.sc.m_x, n, sl.picks); - sl.hold(enc, order[0..n], reqs, level); + mtp.layer(&x, pass, e.sc.m_x, sl.picks, n, 0); // the level before's outputs, a prefix + mtp.drafts(&y, pass, e.sc.m_x, n, sl.picks); + sl.hold(pass.enc, order[0..n], reqs, level); } } diff --git a/zig/tests/contract/metal_glm_ep_copy.zig b/zig/tests/contract/metal_glm_ep_copy.zig new file mode 100644 index 000000000..95086dfa2 --- /dev/null +++ b/zig/tests/contract/metal_glm_ep_copy.zig @@ -0,0 +1,424 @@ +//! GLM exchange copies and token selection under both committed and stale receive epochs. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); + +const c = contract.spec; +const g = contract.gates; +const source = @import("kernel_sources").glm_ep; +const width = 4096; +const words = width / 2; +const topk = 8; +const vocab_part = 154880 / 2; +const epoch = 0x1_0000_0003; +const stale_epoch = epoch - (1 << 32); +const poison: u32 = 0xd3d3d3d3; +const absent = std.math.maxInt(u32); + +const Which = enum { + ep_pack, + ep_unpack, + ep_rsend, + ep_amax_send, + ep_amax_merge, + + fn name(w: Which) []const u8 { + return switch (w) { + inline else => |kind| "metal glm ep " ++ @tagName(kind), + }; + } + + fn threads(w: Which) u32 { + return switch (w) { + .ep_pack, .ep_unpack, .ep_rsend => 256, + .ep_amax_send, .ep_amax_merge => 64, + }; + } + + fn spec(w: Which) c.Spec { + return .{ + .name = w.name(), + .accum = .none, + .depth = noSums, + .out = .raw, + .tiles = .{ .{ .size = w.maxRows() }, .{}, .{} }, + .nan = .{ .stated = switch (w) { + .ep_pack, .ep_rsend => "the matching epoch copies every bit; a stale epoch leaves the send slot untouched", + .ep_unpack => "mapped picks copy every bit when ready or become zero when stale; other picks are untouched", + .ep_amax_send => "a selected BF16 value widens to FP32; a missing pick sends minus infinity and no id", + .ep_amax_merge => "the peer wins a larger value or a lower id on a tie; a stale epoch writes token zero", + } }, + .served = switch (w) { + .ep_pack, .ep_unpack => &.{ + .{ .rows = 16, .n = words, .k = topk }, .{ .rows = 1, .n = 0, .k = topk }, + }, + .ep_rsend => &.{ + .{ .rows = 4096, .n = width, .k = 1 }, .{ .rows = 16, .n = width, .k = 1 }, + }, + .ep_amax_send => &.{ + .{ .rows = 64, .n = vocab_part, .k = 1 }, .{ .rows = 16, .n = vocab_part - 64, .k = 1 }, + .{ .rows = 64, .n = 0, .k = 1 }, + }, + .ep_amax_merge => &.{.{ .rows = 64, .n = 2, .k = 1 }}, + }, + }; + } + + fn maxRows(w: Which) usize { + return switch (w) { + .ep_pack, .ep_unpack => 16, + .ep_rsend => 4096, + .ep_amax_send, .ep_amax_merge => 64, + }; + } + + fn countsEntries(w: Which) bool { + return w == .ep_rsend or w == .ep_amax_send; + } + + fn planted(w: Which, a: std.mem.Allocator) ![]u8 { + const edit = switch (w) { + .ep_pack, .ep_rsend, .ep_amax_send => .{ "if (previous[0] != expected) return;", "if (uint(previous[0]) != uint(expected)) return;" }, + .ep_unpack => .{ "ready[0] == x", "uint(ready[0]) == uint(x)" }, + .ep_amax_merge => .{ "ready[0] != x", "uint(ready[0]) != uint(x)" }, + }; + const begin = std.mem.indexOf(u8, source, @tagName(w)) orelse return error.ControlEditMissing; + const end = std.mem.indexOf(u8, source[begin..], "\n}") orelse return error.ControlEditMissing; + const at = std.mem.indexOf(u8, source[begin..][0..end], edit[0]) orelse return error.ControlEditMissing; + return std.mem.concat(a, u8, &.{ source[0 .. begin + at], edit[1], source[begin + at + edit[0].len ..] }); + } +}; + +fn noSums(_: c.Shape) usize { + return 0; +} + +fn wordsAt(rig: *const rig_.Rig, ref: rig_.Ref, n: usize) []u32 { + return @ptrCast(@alignCast(rig.bytes(ref, n * @sizeOf(u32)))); +} + +fn halvesAt(rig: *const rig_.Rig, ref: rig_.Ref, n: usize) []u16 { + return @ptrCast(@alignCast(rig.bytes(ref, n * @sizeOf(u16)))); +} + +const Slot = struct { + stamp: rig_.Ref, + out: rig_.Ref, + count: ?rig_.Ref = null, +}; + +const Probe = struct { + rig: rig_.Rig, + which: Which, + pipe: mtl.Pipeline, + shape: c.Shape = undefined, + control: bool, + values: rig_.Ref = undefined, + picks: rig_.Ref = undefined, + counts: rig_.Ref = undefined, + slots: [2]Slot = undefined, + entries: usize = 0, + message: [120]u8 = undefined, + + const vt = g.hooks(Probe); + + fn kernel(p: *Probe) g.Kernel { + var desc = p.which.spec(); + if (p.control) { + desc.name = switch (p.which) { + inline else => |kind| "control: metal glm ep " ++ @tagName(kind) ++ " ignores high epoch bits", + }; + } + return .{ .spec = desc, .ctx = p, .vt = &vt, .control = if (p.control) .shapes else null }; + } + + pub fn count(p: *Probe) usize { + const shape = p.shape; + return switch (p.which) { + .ep_pack, .ep_unpack => shape.rows * topk * words, + .ep_rsend => shape.rows * width, + .ep_amax_send => 2 * 64, + .ep_amax_merge => 64, + }; + } + + pub fn reference(_: *Probe, _: []f64, _: []f64) !void { + return error.CheckedByItsOwnRule; + } + + fn route(p: *Probe, fill: g.Fill, rng: std.Random) !void { + const picks = p.shape.rows * topk; + p.entries = if (p.shape.n == 0) 0 else if (p.shape.rows == 16) picks else picks * 2 / 3; + const order = try p.rig.a.alloc(u32, picks); + defer p.rig.a.free(order); + @memset(order, absent); + for (order[0..p.entries], 0..) |*at, i| at.* = @intCast(i * picks / p.entries); + p.picks = try p.rig.input(std.mem.sliceAsBytes(order)); + const owned: i32 = @intCast(p.entries); + const other: i32 = @intCast(picks - p.entries); + const counts: [2]i32 = if (p.which == .ep_pack) .{ owned, other } else .{ other, owned }; + p.counts = try p.rig.input(std.mem.asBytes(&counts)); + try p.payload(fill, rng, true); + } + + fn payload(p: *Probe, fill: g.Fill, rng: std.Random, packed_words: bool) !void { + const a = p.rig.a; + const data = try a.alloc(u32, p.count()); + defer a.free(data); + for (data) |*v| v.* = if (packed_words) rng.int(u32) & 0xbfff_bfff else rng.int(u32) & 0xbfff_ffff; + if (fill == .nonfinite) { + data[0..3].* = if (packed_words) .{ 0x7fc1_ff80, 0x7f80_8000, 0xff80_7fc2 } else .{ 0x7fc0_0081, 0xff80_0000, 0x7f80_0000 }; + } + p.values = try p.rig.input(std.mem.sliceAsBytes(data)); + } + + fn logits(p: *Probe, fill: g.Fill, rng: std.Random) !void { + const a = p.rig.a; + const n = p.shape.n; + const buf = try a.alloc(u16, @max(1, p.shape.rows * n)); + defer a.free(buf); + @memset(buf, 0x3f80); + var chosen: [64]u32 = @splat(absent); + if (n != 0) for (chosen[0..p.shape.rows], 0..) |*id, r| { + if (r % 8 == 7) continue; + id.* = if (r == 0) @intCast(n - 1) else rng.uintLessThan(u32, @intCast(n)); + buf[r * n + id.*] = if (fill == .nonfinite) switch (r) { + 0 => 0x7fc1, + 1 => 0x7f80, + 2 => 0xff80, + 3 => 0x8000, + else => 0x4000, + } else 0x4000 | @as(u16, rng.int(u7)); + }; + p.values = try p.rig.input(std.mem.sliceAsBytes(buf)); + p.picks = try p.rig.input(std.mem.asBytes(&chosen)); + } + + fn candidates(p: *Probe, fill: g.Fill, rng: std.Random) !void { + var mine: [128]u32 = undefined; + var peer: [128]u32 = undefined; + for (0..64) |r| { + const ma = &mine[2 * r]; + const mb = &peer[2 * r]; + const a: f32 = rng.float(f32) + 1; + const b: f32 = rng.float(f32) - 1; + ma.* = @bitCast(if (r % 4 == 0) b else a); + mb.* = @bitCast(if (r % 4 == 0) a else if (r % 4 == 1) b else a); + mine[2 * r + 1] = if (r % 4 == 3) 19 else 41; + peer[2 * r + 1] = if (r % 4 == 3) 41 else 19; + } + if (fill == .nonfinite) { + mine[0..10].* = .{ 0x7fc0_0018, 41, 0x7f80_0000, 10, 0xff80_0000, absent, 0x8000_0000, 41, 0x3f80_0000, 41 }; + peer[0..10].* = .{ 0x7f80_0000, 19, 0x7f80_0000, 9, 0xff80_0000, absent, 0, 19, 0x7fc0_0081, 19 }; + } + p.values = try p.rig.input(std.mem.asBytes(&mine)); + p.picks = try p.rig.input(std.mem.asBytes(&peer)); + } + + pub fn prepare(p: *Probe, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.rows > p.which.maxRows()) return false; + p.rig.reset(); + p.shape = shape; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + switch (p.which) { + .ep_pack, .ep_unpack => try p.route(fill, rng), + .ep_rsend => try p.payload(fill, rng, false), + .ep_amax_send => try p.logits(fill, rng), + .ep_amax_merge => try p.candidates(fill, rng), + } + for (&p.slots, [_]u64{ epoch, stale_epoch }) |*slot, seen_epoch| { + slot.* = .{ + .stamp = try p.rig.input(std.mem.asBytes(&seen_epoch)), + .out = try p.rig.output(p.count() * 4), + .count = if (p.which.countsEntries()) try p.rig.output(4) else null, + }; + @memset(wordsAt(&p.rig, slot.out, p.count()), poison); + if (slot.count) |at| wordsAt(&p.rig, at, 1)[0] = poison; + } + return true; + } + + fn bind(e: mtl.ComputeEncoder, at: usize, refs: []const rig_.Ref) void { + for (refs, at..) |r, i| e.setBuffer(r.buf, r.off, i); + } + + fn stamp(e: mtl.ComputeEncoder, at: usize, slot: Slot) void { + bind(e, at, &.{slot.stamp}); + e.setValue(@as(u64, epoch), at + 1); + } + + fn offset(p: *const Probe) u32 { + return if (p.shape.rows % 2 == 0) vocab_part else 0; + } + + pub fn run(p: *Probe) !void { + const pass = p.rig.begin(); + const e = pass.e; + const rows: u32 = @intCast(p.shape.rows); + const group = mtl.Size.of(p.which.threads(), 1, 1); + e.setPipeline(p.pipe); + for (p.slots) |slot| { + switch (p.which) { + .ep_pack => { + bind(e, 0, &.{ p.values, p.picks, p.counts, slot.out }); + stamp(e, 4, slot); + e.dispatchGroups(mtl.Size.of(1, rows * topk, 1), group); + }, + .ep_unpack => { + stamp(e, 0, slot); + bind(e, 3, &.{ p.values, p.picks, p.counts, slot.out }); + e.dispatchGroups(mtl.Size.of(1, rows * topk, 1), group); + }, + .ep_rsend => { + bind(e, 0, &.{p.values}); + e.setValue(@as(i32, @intCast(rows)), 1); + bind(e, 2, &.{ slot.out, slot.count.? }); + stamp(e, 4, slot); + e.dispatchThreads(mtl.Size.of(rows * width, 1, 1), group); + }, + .ep_amax_send => { + bind(e, 0, &.{ p.values, p.picks }); + e.setValue([4]u32{ rows, @intCast(p.shape.n), p.offset(), 0 }, 2); + bind(e, 3, &.{ slot.out, slot.count.? }); + stamp(e, 5, slot); + e.dispatchThreads(group, group); + }, + .ep_amax_merge => { + stamp(e, 0, slot); + bind(e, 3, &.{ p.values, p.picks }); + e.setValue(rows, 5); + bind(e, 6, &.{slot.out}); + e.dispatchThreads(group, group); + }, + } + } + try p.rig.finish(pass); + } + + fn checkRoute(p: *Probe, out: []const u32, ready: bool) ?usize { + const picked = wordsAt(&p.rig, p.picks, p.shape.rows * topk); + const data = wordsAt(&p.rig, p.values, p.count()); + var unpacked: usize = 0; + for (0..picked.len) |row| { + if (p.which == .ep_pack) { + const source_row = if (ready and row < p.entries) picked[row] else absent; + if (checkWords(out, data, row, source_row, poison)) |at| return at; + } else { + const mapped = unpacked < p.entries and picked[unpacked] == row; + const source_row = if (mapped and ready) @as(u32, @intCast(unpacked)) else absent; + if (checkWords(out, data, row, source_row, if (mapped) 0 else poison)) |at| return at; + if (mapped) unpacked += 1; + } + } + return null; + } + + fn checkWords(out: []const u32, data: []const u32, row: usize, source_row: u32, empty: u32) ?usize { + const start = row * words; + for (out[start..][0..words], 0..) |got, col| { + const want = if (source_row == absent) empty else data[@as(usize, source_row) * words + col]; + if (got != want) return start + col; + } + return null; + } + + fn checkPartial(p: *Probe, out: []const u32, ready: bool) ?usize { + const data = wordsAt(&p.rig, p.values, p.count()); + for (out, data, 0..) |got, bits, i| { + if (got != if (ready) bits else poison) return i; + } + return null; + } + + fn checkLogits(p: *Probe, out: []const u32, ready: bool) ?usize { + const n = p.shape.n; + const data = halvesAt(&p.rig, p.values, @max(1, p.shape.rows * n)); + const ids = wordsAt(&p.rig, p.picks, 64); + for (ids, 0..) |id, row| { + const wrote = ready and row < p.shape.rows; + const has = wrote and id < n; + const want = if (!wrote) poison else if (has) @as(u32, data[row * n + id]) << 16 else 0xff80_0000; + const got = out[2 * row]; + const nan = std.math.isNan(@as(f32, @bitCast(want))); + if (if (nan) !std.math.isNan(@as(f32, @bitCast(got))) else got != want) return 2 * row; + const want_id = if (!wrote) poison else if (has) p.offset() + id else absent; + if (out[2 * row + 1] != want_id) return 2 * row + 1; + } + return null; + } + + fn checkPick(p: *Probe, out: []const u32, ready: bool) ?usize { + const mine = wordsAt(&p.rig, p.values, 128); + const peer = wordsAt(&p.rig, p.picks, 128); + for (out, 0..) |got, r| { + const local_value: f64 = @as(f32, @bitCast(mine[2 * r])); + const peer_value: f64 = @as(f32, @bitCast(peer[2 * r])); + const peer_id = peer[2 * r + 1]; + const local_id = mine[2 * r + 1]; + const winner = if (peer_value > local_value or (peer_value == local_value and peer_id < local_id)) peer_id else local_id; + if (got != if (r >= p.shape.rows) poison else if (!ready) 0 else winner) return r; + } + return null; + } + + pub fn check(p: *Probe, _: g.Gate) !?[]const u8 { + if (!try p.rig.intact()) return "a launch wrote outside its buffers"; + for (p.slots, [_]bool{ true, false }) |slot, ready| { + const out = wordsAt(&p.rig, slot.out, p.count()); + const mismatch = switch (p.which) { + .ep_pack, .ep_unpack => p.checkRoute(out, ready), + .ep_rsend => p.checkPartial(out, ready), + .ep_amax_send => p.checkLogits(out, ready), + .ep_amax_merge => p.checkPick(out, ready), + }; + const when = if (ready) "matching" else "stale"; + if (mismatch) |at| return try std.fmt.bufPrint(&p.message, "{s} epoch differs at output word {d}", .{ when, at }); + if (slot.count) |ref| { + const want: u32 = if (!ready) poison else if (p.which == .ep_rsend) @intCast(p.shape.rows * 2) else 1; + if (wordsAt(&p.rig, ref, 1)[0] != want) return try std.fmt.bufPrint(&p.message, "{s} epoch wrote an incorrect entry count", .{when}); + } + } + return null; + } +}; + +pub const Suite = struct { + a: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + kernels: *std.ArrayList(g.Kernel), + probes: std.ArrayList(*Probe) = .empty, + + pub fn collect(s: *Suite, only: ?[]const u8) !void { + for (std.enums.values(Which)) |which| { + if (only) |needle| if (std.mem.indexOf(u8, which.name(), needle) == null) continue; + try s.add(which, false); + try s.add(which, true); + } + } + + fn add(s: *Suite, which: Which, control: bool) !void { + try s.probes.ensureUnusedCapacity(s.a, 1); + try s.kernels.ensureUnusedCapacity(s.a, 1); + const text = if (control) try which.planted(s.a) else source; + defer if (control) s.a.free(text); + const pipe = try rig_.pipeline(s.device, text, @tagName(which)); + errdefer pipe.deinit(); + if (pipe.maxThreads() < which.threads() or rig_.tgMemory(pipe) != 0) return error.UnexpectedLaunchLimits; + const p = try s.a.create(Probe); + p.* = .{ .rig = .{ .a = s.a, .device = s.device, .queue = s.queue }, .which = which, .pipe = pipe, .control = control }; + s.probes.appendAssumeCapacity(p); + s.kernels.appendAssumeCapacity(p.kernel()); + } + + pub fn deinit(s: *Suite) void { + for (s.probes.items) |p| { + p.rig.deinit(); + p.pipe.deinit(); + s.a.destroy(p); + } + s.probes.deinit(s.a); + } +}; diff --git a/zig/tests/contract/metal_glm_ep_reduce.zig b/zig/tests/contract/metal_glm_ep_reduce.zig new file mode 100644 index 000000000..99db46690 --- /dev/null +++ b/zig/tests/contract/metal_glm_ep_reduce.zig @@ -0,0 +1,344 @@ +//! GLM expert-parallel reductions: guarded send counts, rank sums, and full-width receive epochs. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const rows = contract.rows; +const c = contract.spec; +const g = contract.gates; +const rig = @import("rig.zig"); +const source = @import("kernel_sources").glm_ep; +const max_rows = @import("tensorfold").glm.ep.PROMPT_ROWS; + +const width = 4096; +const picks = 8; +const threads = 256; +const epoch: u64 = 0x1_0000_0003; +const stale: u64 = @as(u32, @truncate(epoch)); +const Kind = enum { + ep_rcombine, + ep_rfinal, + ep_rsum, + + fn combine(k: Kind) bool { + return k == .ep_rcombine; + } + + fn elem(k: Kind) c.Elem { + return if (k.combine()) .f32 else .bf16; + } + + fn spec(k: Kind, label: []const u8) c.Spec { + return .{ + .name = label, + .accum = .f32, + .depth = if (k.combine()) combineDepth else receiveDepth, + .out = k.elem(), + .nan = .{ .stated = "NaN and infinity classes match the reference; a stale epoch blocks transfer data" }, + .served = if (k.combine()) &combine_shapes else &receive_shapes, + .tiles = .{ .{ .size = 16 }, .{}, .{} }, + }; + } +}; + +fn shapes(comptime k: usize) [7]c.Shape { + var set: [7]c.Shape = undefined; + for (&set, [_]usize{ 5, 2, 32, 63, 64, max_rows - 1, max_rows }) |*s, n| + s.* = .{ .rows = n, .n = width, .k = k }; + return set; +} +const combine_shapes = shapes(picks); +const receive_shapes = shapes(2); +fn combineDepth(_: c.Shape) usize { + return picks * 2 - 1; +} +fn receiveDepth(_: c.Shape) usize { + return 1; +} + +fn combineInput(i: usize, salt: u64) f32 { + const h = std.hash.int(i +% salt); + return @as(f32, @floatFromInt(@as(i32, @as(u16, @truncate(h))) - 32768)) / 8192; +} + +const Probe = struct { + const vt = g.hooks(Probe); + + rig: rig.Rig, + pipe: mtl.Pipeline, + kind: Kind, + shape: c.Shape = undefined, + fill: g.Fill = .random, + seed: u64 = 0, + data: rig.Ref = undefined, + extra: rig.Ref = undefined, + shared: rig.Ref = undefined, + ready: rig.Ref = undefined, + expired: rig.Ref = undefined, + out: rig.Ref = undefined, + count_ref: rig.Ref = undefined, + + fn rank(p: *const Probe) u32 { + const parity: u32 = @intCast(p.shape.rows % 2); + return parity ^ @intFromBool(p.kind == .ep_rfinal); + } + + fn deinit(p: *Probe) void { + p.rig.deinit(); + p.pipe.deinit(); + } + + fn part(p: *const Probe, slot: usize, i: usize) f32 { + if (p.fill == .stress) { + const stress = [_]f32{ 0x1p20, 1 - 0x1p20, 1, -1, 2, -2, 3, -3 }; + return stress[slot] * if (p.kind.combine()) @as(f32, picks) else 1; + } + if (p.fill == .nonfinite) { + const inf = std.math.inf(f32); + if (slot == 0 and i < 4) return ([_]f32{ std.math.nan(f32), inf, -inf, inf })[i]; + if (slot == 1 and i == 3) return -inf; + } + if (p.kind == .ep_rfinal and i == 4) return if (slot == 0) 1 else 0x1p-8; + return combineInput(i, p.seed +% (slot + 1) * 0x9e3779b9); + } + + fn weight(p: *const Probe, row: usize, slot: usize) f32 { + if (p.fill == .stress) return 1.0 / @as(f32, picks); + return (4 + combineInput(row * picks + slot, p.seed)) / 32 + 0x1p-8; + } + + fn sharedAt(p: *const Probe, i: usize) u16 { + if (p.fill == .stress) return 0; + if (p.fill == .nonfinite and i == 5) return c.toBf16(std.math.inf(f32)); + if (i == 4) return c.toBf16(-1); + return c.toBf16(combineInput(i, p.seed) / 4); + } + + pub fn prepare(p: *Probe, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.rows > max_rows or shape.n != width) return false; + p.rig.reset(); + p.shape = shape; + p.fill = fill; + p.seed = seed; + if (p.kind.combine()) { + const data = try p.rig.a.alloc(f32, p.count() * picks); + defer p.rig.a.free(data); + for (data, 0..) |*v, i| v.* = p.part(i / width % picks, i / (width * picks) * width + i % width); + p.data = try p.rig.input(std.mem.sliceAsBytes(data)); + const wts = try p.rig.a.alloc(f32, shape.rows * picks); + defer p.rig.a.free(wts); + for (wts, 0..) |*v, i| v.* = p.weight(i / picks, i % picks); + p.extra = try p.rig.input(std.mem.sliceAsBytes(wts)); + p.count_ref = try p.rig.scratch(4); + } else { + const data = try p.rig.a.alloc(f32, p.count()); + defer p.rig.a.free(data); + const own = p.rank(); + for (data, 0..) |*v, i| v.* = p.part(own, i); + p.data = try p.rig.input(std.mem.sliceAsBytes(data)); + for (data, 0..) |*v, i| v.* = p.part(1 - own, i); + p.extra = try p.rig.input(std.mem.sliceAsBytes(data)); + if (p.kind == .ep_rfinal) { + const shared = try p.rig.a.alloc(u16, p.count()); + defer p.rig.a.free(shared); + for (shared, 0..) |*v, i| v.* = p.sharedAt(i); + p.shared = try p.rig.input(std.mem.sliceAsBytes(shared)); + } + } + p.ready = try p.rig.input(std.mem.asBytes(&epoch)); + p.expired = try p.rig.input(std.mem.asBytes(&stale)); + p.out = try p.rig.output(p.outLen()); + return true; + } + + pub fn count(p: *const Probe) usize { + return p.shape.rows * width; + } + + fn outLen(p: *const Probe) usize { + return p.count() * c.size(p.kind.elem()); + } + + fn dispatch(p: *Probe, ready: rig.Ref) !void { + @memset(p.rig.bytes(p.out, p.outLen()), 0xff); + const pass = p.rig.begin(); + const e = pass.e; + const n: i32 = @intCast(p.shape.rows); + e.setPipeline(p.pipe); + if (p.kind.combine()) { + @memset(p.rig.bytes(p.count_ref, 4), 0xff); + e.setBuffer(p.data.buf, p.data.off, 0); + e.setBuffer(p.extra.buf, p.extra.off, 1); + e.setValue(n, 2); + e.setBuffer(p.out.buf, p.out.off, 3); + e.setBuffer(p.count_ref.buf, p.count_ref.off, 4); + e.setBuffer(ready.buf, ready.off, 5); + e.setValue(epoch, 6); + e.dispatchThreads(mtl.Size.of(p.count(), 1, 1), mtl.Size.of(threads, 1, 1)); + } else { + e.setBuffer(ready.buf, ready.off, 0); + e.setValue(epoch, 1); + e.setBuffer(p.data.buf, p.data.off, 3); + e.setBuffer(p.extra.buf, p.extra.off, 4); + e.setValue(p.rank(), 5); + const out_binding: u32 = if (p.kind == .ep_rfinal) 7 else 6; + if (p.kind == .ep_rfinal) e.setBuffer(p.shared.buf, p.shared.off, 6); + e.setBuffer(p.out.buf, p.out.off, out_binding); + e.setValue(n, out_binding + 1); + e.dispatchGroups(mtl.Size.of(p.count() / threads, 1, 1), mtl.Size.of(threads, 1, 1)); + } + try p.rig.finish(pass); + } + + pub fn run(p: *Probe) !void { + try p.dispatch(p.ready); + if (p.kind.combine() and p.sent() != p.shape.rows * 2) return error.WrongExchangeCount; + } + + fn sent(p: *const Probe) u32 { + return std.mem.readInt(u32, p.rig.bytes(p.count_ref, 4)[0..4], .little); + } + + fn terms(p: *const Probe, i: usize, out: *[picks]f64) []const f64 { + if (p.kind.combine()) { + for (out, 0..) |*v, slot| v.* = @as(f64, p.weight(i / width, slot)) * p.part(slot, i); + return out; + } + out[0] = p.part(0, i); + out[1] = p.part(1, i); + return out[0..2]; + } + + fn finish(p: *const Probe, i: usize, value: f64) f64 { + if (p.kind.combine()) return value; + const rounded = contract.bf16_round.bf16(value); + if (p.kind == .ep_rfinal) return contract.bf16_round.bf16(rounded + c.fromBf16(p.sharedAt(i))); + return rounded; + } + + fn ordered(p: *const Probe, i: usize, ready: bool) f32 { + if (!ready) { + const local = p.part(p.rank(), i) + 0.0; + return @floatCast(p.finish(i, local)); + } + var values: [picks]f64 = undefined; + const v = p.terms(i, &values); + return @floatCast(p.finish(i, rows.sumIn(.f32, .sequential, v))); + } + + pub fn reference(p: *Probe, want: []f64, mag: []f64) !void { + for (want, mag, 0..) |*w, *m, i| { + var values: [picks]f64 = undefined; + const terms_ = p.terms(i, &values); + w.* = p.finish(i, rows.sumIn(.none, .sequential, terms_)); + m.* = 0; + for (terms_) |v| m.* += @abs(v); + } + } + + pub fn simulate(p: *Probe, acc: c.Accum, order: g.Order, out: []f64) !void { + for (out, 0..) |*v, i| { + var values: [picks]f64 = undefined; + v.* = p.finish(i, rows.sumIn(acc, order, p.terms(i, &values))); + } + } + + fn compare(p: *Probe, gate: g.Gate, ready: bool) !?[]const u8 { + if (!try p.rig.intact()) return "a launch wrote outside its buffers"; + const out = p.rig.bytes(p.out, p.outLen()); + if (p.kind.combine() and !ready) { + if (p.sent() != std.math.maxInt(u32) or !std.mem.allEqual(u8, out, 0xff)) return "stale previous epoch changed the send payload or count"; + return null; + } + var planted: usize = 0; + for (0..p.count()) |i| { + const got = c.read(p.kind.elem(), out, i); + if (gate == .nonfinite) { + var values: [picks]f64 = undefined; + const want = if (ready) p.finish(i, rows.sumIn(.none, .sequential, p.terms(i, &values))) else p.ordered(i, false); + if (c.classOf(want) != .finite) planted += 1; + if (c.classOf(got) != c.classOf(want)) { + p.report(ready, i, got, want); + return "non-finite classes differ from float64"; + } + } else { + const want = p.ordered(i, ready); + if (@as(u32, @bitCast(@as(f32, @floatCast(got)))) != @as(u32, @bitCast(want))) { + p.report(ready, i, got, want); + return "output bits differ from the ordered fp32 and bf16 rounding reference"; + } + } + } + if (gate == .nonfinite and ready and (planted == 0 or planted == p.count())) return "the fixture needs both finite and non-finite outputs"; + return null; + } + + fn report(p: *const Probe, ready: bool, i: usize, got: f64, want: f64) void { + std.debug.print("{t} rank {d} ready {} at row {d} column {d}: got {e}, want {e}\n", .{ + p.kind, p.rank(), ready, i / width, i % width, got, want, + }); + } + + pub fn check(p: *Probe, gate: g.Gate) !?[]const u8 { + if (try p.compare(gate, true)) |why| return why; + try p.dispatch(p.expired); + return p.compare(gate, false); + } +}; + +fn name(comptime k: Kind, comptime planted: bool) []const u8 { + return "metal glm ep " ++ @tagName(k) ++ (if (planted) " low-word epoch control" else ""); +} + +pub const Suite = struct { + a: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + kernels: *std.ArrayList(g.Kernel), + probes: std.ArrayList(*Probe) = .empty, + + pub fn collect(s: *Suite, only: ?[]const u8) !void { + inline for (.{ false, true }) |planted| { + var library: ?mtl.Library = null; + defer if (library) |lib| lib.deinit(); + inline for (std.enums.values(Kind)) |k| { + const label = name(k, planted); + if (only == null or std.mem.indexOf(u8, label, only.?) != null) { + if (library == null) library = try s.makeLibrary(planted); + try s.add(library.?, k, planted, label); + } + } + } + } + + pub fn deinit(s: *Suite) void { + for (s.probes.items) |p| { + p.deinit(); + s.a.destroy(p); + } + s.probes.deinit(s.a); + } + + fn makeLibrary(s: *Suite, planted: bool) !mtl.Library { + if (!planted) return mtl.Library.fromSource(s.device, source, mtl.CompileOptions.mlx()); + const sends = try std.mem.replaceOwned(u8, s.a, source, "previous[0] != expected", "uint(previous[0]) != uint(expected)"); + defer s.a.free(sends); + const receives = try std.mem.replaceOwned(u8, s.a, sends, "ready[0] == x", "uint(ready[0]) == uint(x)"); + defer s.a.free(receives); + return mtl.Library.fromSource(s.device, receives, mtl.CompileOptions.mlx()); + } + + fn add(s: *Suite, library: mtl.Library, comptime k: Kind, comptime planted: bool, label: []const u8) !void { + const p = try s.a.create(Probe); + errdefer s.a.destroy(p); + p.* = .{ + .rig = .{ .a = s.a, .device = s.device, .queue = s.queue }, + .pipe = try mtl.Pipeline.init(s.device, library, @tagName(k), false), + .kind = k, + }; + errdefer p.deinit(); + try s.probes.ensureUnusedCapacity(s.a, 1); + try s.kernels.ensureUnusedCapacity(s.a, 1); + s.probes.appendAssumeCapacity(p); + s.kernels.appendAssumeCapacity(.{ .spec = k.spec(label), .ctx = p, .vt = &Probe.vt, .control = if (planted) .shapes else null }); + } +}; diff --git a/zig/tests/contract/metal_limits.zig b/zig/tests/contract/metal_limits.zig index e33553113..008e68286 100644 --- a/zig/tests/contract/metal_limits.zig +++ b/zig/tests/contract/metal_limits.zig @@ -58,6 +58,10 @@ fn resolve(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch, comp return std.mem.indexOf(u8, s, p) != null; } }.f; + if (std.mem.startsWith(u8, d.name, "metal glm ep ")) { + if (!assemble) return undefined; + return .{ .decl = d, .function = fun, .here = sources.glm_ep, .pre_m5 = sources.glm_ep }; + } if (std.mem.eql(u8, d.name, "contract masked values tf_masked_control")) { if (!assemble) return undefined; const control_text = @import("masked_value_probe.zig").control_source; diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 24bd0ea4f..8a47ca3d0 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -421,6 +421,12 @@ pub fn main(init: std.process.Init) !u8 { var glm_windows: @import("metal_glm_rows.zig").Suite = .{ .a = a, .device = device, .queue = queue, .kernels = &kernels }; defer glm_windows.deinit(); try glm_windows.collect(o.only); + var glm_ep_copy: @import("metal_glm_ep_copy.zig").Suite = .{ .a = a, .device = device, .queue = queue, .kernels = &kernels }; + defer glm_ep_copy.deinit(); + try glm_ep_copy.collect(o.only); + var glm_ep_reduce: @import("metal_glm_ep_reduce.zig").Suite = .{ .a = a, .device = device, .queue = queue, .kernels = &kernels }; + defer glm_ep_reduce.deinit(); + try glm_ep_reduce.collect(o.only); var flint: @import("metal_flint.zig").Suite = .{ .a = a, .device = device, .queue = queue, .kernels = &kernels }; defer flint.deinit(); try flint.collect(o.only);