diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 1ed43ac146..7fa49b6419 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -20,7 +20,7 @@ The test of a kernel fills a `contract.spec.Spec` and hands the harness a few ho and the magnitude of the terms each output sums, and can replay the sums in another accumulator. - Its served envelope: the shapes the engine calls it at, and for each dimension its tile, the step a valid size moves by and the least valid size. -- Its accumulator (`f32` for every kernel here) and the depth of its longest chain of roundings at a shape, plus any +- Its accumulator (`f32` where a kernel sums) and the depth of its longest chain of roundings at a shape, plus any fp32 operations after the sums (`ops`) and intermediate roundings to the output type (`mids`). - Whether it may write in place. If it replaces another kernel, the kernel whose bits it must equal. - Its NaN rule: `propagate` (as float64 arithmetic on its formula does), or a stated rule the test checks itself @@ -28,7 +28,8 @@ The test of a kernel fills a `contract.spec.Spec` and hands the harness a few ho - For gate 5, its launch in `zig/src/core/contract/registry.zig`: threads a group and threadgroup bytes (the source's, and any its encoder sets) on Metal, or symbol, image, threads a block, dynamic shared bytes and cluster width on CUDA, with the limits its compiles recorded. HIP launches are in `zig/src/core/contract/hip_records.zig`: symbol, code - object, threads a workgroup and dynamic LDS bytes. + object, threads a workgroup and dynamic LDS bytes. The seam's records are in `metal_pair_records.zig`, which + `registry.zig` concatenates. ## The gates @@ -99,7 +100,8 @@ The program exits 1 on any failed gate, any control that stayed green, or any li the inputs for a shape and a gate's fill (usually through an input family), `run` launches once and waits, `count` and `reference` give the outputs and their float64 values, `simulate` replays the sums for gate 4. 2. Add it to the runner's list, with its served shapes and tiles. -3. Add its launches to `registry.zig`, and fill the records from `--offline` (Metal) or `kernel-contract-limits` +3. Add its launches to `registry.zig`, or to a records file it concatenates as the seam's are, and fill them from + `--offline` (Metal) or `kernel-contract-limits` (CUDA: nvcc 13.0's in `cuda`, and each later release's set where that release compiles it differently); a HIP kernel's go to `hip_records.zig`, filled from `hip-kernel-contract-limits`. 4. Post the `SUMMARY` lines from each backend you ran in the pull request, and say which chips you could not run. @@ -127,8 +129,8 @@ float64 output class, without an inherited exception. ## Coverage -Metal: softmax, RMS norm, bf16 top-k, the row projection (3, 4, 6 and 8 bits), affine_mm (4 and 8 bits), simd_qmm -(4 bits) and simd_qmm_bits (3, 6 and 8 bits), and the lane projection (6-bit groups of 32 on every chip; 4- and 8-bit +Metal: softmax, RMS norm, bf16 top-k, the row projection (3, 4, 6 and 8 bits), affine_mm (4 and 8 bits), simd_qmm (4 +bits) and simd_qmm_bits (3, 6 and 8 bits), and the lane projection (6-bit groups of 32 on every chip; 4- and 8-bit formats on the M5's tensor units, where its reg kernel must equal the cooperative one bit for bit), and tile_qmm's prompt tiles and their decode twin (4, 6 and 8 bits; the twin bit-equal to the tiles). Gate 5 also holds the rest of tile_qmm's pipelines and GLM-5.3-Flash's hc_expand_10. @@ -141,6 +143,13 @@ Metal also checks Nemotron's routed experts on the engine's route tables, one ro against float64 at tf_load16's bf16 x sums (up's relu2 as an interval over its two bf16 roundings), and the rows2, rows4 and six geometry pipelines of each projection bit-equal to them (two controls: member rows read from one row, and w3's x sums without their bf16 steps). +The row-split seam's four kernels (`pair_rows.metal`, launched by `pair_rows.zig`) move bf16 frames, never sums: +`pair_stage` copies this rank's half of a frame row into the slot the relay sends, and `pair_gather` writes both halves +of every row, its own natively and the peer's from the relayed frame. `pair_argmax` flags a row whose logits are not +finite (its id word carries bit 31, its value word a NaN) and takes the lower id on a tie; `pair_pick` turns a flagged +rank's word into the id sentinel, so both Macs refuse a non-finite round where one Mac does. Six controls edit the +source: a stage reading the peer's half, a gather writing its own half twice, and an argmax or pick tie to the higher +id must go red at `shapes`, and an argmax or pick that does not flag at `nonfinite`. CUDA: the RMS norm rows (16-byte form and the fp32-out tree), the bf16 residual's norm tails (MLX's rounded rows and the once-rounded ones, each checked on its residual and its output, with four controls each dropping a rounding point), Nemotron's router logits (and their 128-row tiles, bit-equal to them) and top-k, the latent top-k's decode form (bit-equal to `tf_super_topk`), the latent combine and diff --git a/zig/kernels/metal/core/pair_rows.metal b/zig/kernels/metal/core/pair_rows.metal new file mode 100644 index 0000000000..a505ba2cd8 --- /dev/null +++ b/zig/kernels/metal/core/pair_rows.metal @@ -0,0 +1,129 @@ +// Two Macs' tensor-parallel exchanges: the host relays parts through shared events, and these kernels move bf16 bits. +#include +using namespace metal; +// A rank compares one row's logits across ranks; no kernel sums across them. + +// Sync page words (shared with pair_rows.zig): FAILED = the host gave up on a seam; OBSERVED = this side's refusals. +constant constexpr uint FAILED = 16, OBSERVED = 20; +// An argmax id word's bit 31: the row had a logit that is not finite, so pair_pick sends the row to the id sentinel. +constant constexpr uint PAIR_FLAG = 0x80000000u; +constant constexpr uint PAIR_SENTINEL = 0xFFFFFFFFu; + +// Kernels' argument block, byte-identical to pair_rows.zig's PairRows: width, the frame row's stride, this rank, one. +struct PairRows { + uint width, stride, rank, x; +}; + +// bf16 bits as the exact fp32 value they name (argmax reads a bf16 logits stream this way). +inline float pair_bf16(uint bits) { + return as_type(bits << 16); +} + +// True once the host gave up: FAILED is set before the wait releases, so a later kernel poisons its output rows. +inline bool pair_link_broken(device atomic_uint* sync, uint t) { + if (atomic_load_explicit(&sync[FAILED], memory_order_relaxed) == 0u) return false; + if (t == 0) atomic_fetch_add_explicit(&sync[OBSERVED], 1u, memory_order_relaxed); + return true; +} + +// Row r: this rank's own half of the frame row into the exchange slot; a bf16 copy, so it carries the frame's bits. +kernel void pair_stage(device atomic_uint* sync [[buffer(0)]], constant PairRows& a [[buffer(1)]], + const device ushort* src [[buffer(2)]], device ushort* dst [[buffer(3)]], + uint r [[threadgroup_position_in_grid]], uint t [[thread_index_in_threadgroup]]) { + const ulong in = ulong(r) * a.stride; + const ulong out = ulong(r) * a.width; + if (pair_link_broken(sync, t)) { + for (uint i = t; i < a.width; i += 256) dst[out + i] = 0x7FC0u; // bf16 NaN: the link stays down + return; + } + for (uint i = t; i < a.width; i += 256) dst[out + i] = src[in + i]; +} + +// Row r: this rank's half at its own offset in the frame row and the peer's at the other; a bf16 copy, no arithmetic. +kernel void pair_gather(device atomic_uint* sync [[buffer(0)]], constant PairRows& a [[buffer(1)]], + const device ushort* mine [[buffer(2)]], device atomic_uint* peer [[buffer(3)]], + device ushort* out [[buffer(4)]], uint r [[threadgroup_position_in_grid]], + uint t [[thread_index_in_threadgroup]]) { + const ulong mrow = ulong(r) * a.width; + const ulong orow = ulong(r) * a.stride; + const ulong own = a.rank == 0 ? orow : orow + a.width; + const ulong got = a.rank == 0 ? orow + a.width : orow; + if (pair_link_broken(sync, t)) { + for (uint i = t; i < a.width; i += 256) { + out[own + i] = 0x7FC0u; + out[got + i] = 0x7FC0u; + } + return; + } + for (uint i = t; i < a.width; i += 256) { + // The relay wrote the peer's half: an untracked frame read, so one atomic word (two bf16) at a time. + const uint word = atomic_load_explicit(&peer[(mrow + i) >> 1], memory_order_relaxed); + out[own + i] = mine[mrow + i]; + out[got + i] = ushort(word >> ((i & 1) * 16)); + } +} + +// Row r's argmax over this rank's logits (ids lo..lo+n): (fp32 value, id), lower id on tie; dims.w 0 = fp32, 1 = bf16. +kernel void pair_argmax(const device float* logits [[buffer(0)]], constant uint4& dims [[buffer(1)]], + device uint* out [[buffer(2)]], uint r [[threadgroup_position_in_grid]], + uint t [[thread_index_in_threadgroup]], uint sg [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + const ulong row = ulong(r) * dims.x; + const device float* xf = logits + row; + const device ushort* xb = reinterpret_cast(logits) + row; + const bool wide = dims.w != 0u; + float best = -INFINITY; + uint at = PAIR_SENTINEL; + bool bad = false; + for (uint i = t; i < dims.z; i += 1024) { + const float v = wide ? pair_bf16(uint(xb[i])) : xf[i]; + if (!isfinite(v)) { bad = true; continue; } // the flagged id is the argmax over finite logits + if (v > best) best = v, at = dims.y + i; + } + for (ushort o = 16; o > 0; o >>= 1) { + const float ob = simd_shuffle_xor(best, o); + const uint oa = simd_shuffle_xor(at, o); + if (ob > best || (ob == best && oa < at)) best = ob, at = oa; + } + // The flag is the row's, not the thread's: one lane saw a non-finite logit, so every lane must know. + uint dirty = bad ? 1u : 0u; + for (ushort o = 16; o > 0; o >>= 1) dirty |= simd_shuffle_xor(dirty, o); + threadgroup float vb[32]; + threadgroup uint va[32]; + threadgroup uint vd[32]; + if (lane == 0) vb[sg] = best, va[sg] = at, vd[sg] = dirty; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg != 0) return; + best = vb[lane]; + at = va[lane]; + dirty = vd[lane]; + for (ushort o = 16; o > 0; o >>= 1) { + const float ob = simd_shuffle_xor(best, o); + const uint oa = simd_shuffle_xor(at, o); + if (ob > best || (ob == best && oa < at)) best = ob, at = oa; + dirty |= simd_shuffle_xor(dirty, o); + } + if (lane == 0) { + out[2 * r] = dirty != 0u ? as_type(NAN) : as_type(best); + out[2 * r + 1] = dirty != 0u ? (at | PAIR_FLAG) : at; + } +} + +// Each row's pick over both Macs' argmaxes: a flagged row draws no token; else the larger value, lower id on a tie. +kernel void pair_pick(device atomic_uint* sync [[buffer(0)]], constant uint2& rx [[buffer(1)]], + const device uint* mine [[buffer(2)]], device atomic_uint* peer [[buffer(3)]], + device uint* picks [[buffer(4)]], uint t [[thread_index_in_threadgroup]]) { + if (pair_link_broken(sync, t)) { + if (t == 0) for (uint r = 0; r < rx.x; ++r) picks[r] = 0xFFFFFFFFu; // no such token: visibly wrong, never a token + return; + } + if (t != 0) return; + for (uint r = 0; r < rx.x; ++r) { + const float vm = as_type(mine[2 * r]); + const uint am = mine[2 * r + 1]; + const float vp = as_type(atomic_load_explicit(&peer[2 * r], memory_order_relaxed)); + const uint ap = atomic_load_explicit(&peer[2 * r + 1], memory_order_relaxed); + // A non-finite row on either Mac draws no token: the same sentinel a broken link reads, never an id. + picks[r] = ((am | ap) & PAIR_FLAG) != 0u ? PAIR_SENTINEL : (vp > vm || (vp == vm && ap < am)) ? ap : am; + } +} diff --git a/zig/kernels/metal/qwen27/attention_mpp.metal b/zig/kernels/metal/qwen27/attention_mpp.metal index b81b68cba8..78f2f84d11 100644 --- a/zig/kernels/metal/qwen27/attention_mpp.metal +++ b/zig/kernels/metal/qwen27/attention_mpp.metal @@ -263,8 +263,71 @@ kernel void q27_attn_tail(device const bfloat* QB [[buffer(0)]], device const Q2 } } -// A row's chunk partials combined in key order, a block of loads ahead; `gated` applies the query gate (q27_gate). +// A row's chunk partials in key order, loads ahead: `gated` the query gate, `heads` the rank's (12 = one Mac's 24 / 2). kernel void q27_attn_merge(device const float* POA [[buffer(0)]], device const float* PMA [[buffer(1)]], device const float* PLA [[buffer(2)]], + device const float* POB [[buffer(3)]], device const float* PMB [[buffer(4)]], device const float* PLB [[buffer(5)]], + device const int* meta [[buffer(6)]], device const int* nodes [[buffer(7)]], device bfloat* OUT [[buffer(8)]], + device const bfloat* QG [[buffer(9)]], constant uint& gated [[buffer(10)]], constant uint& heads [[buffer(11)]], + uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]], uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + const uint lane = thread_index_in_simdgroup; + const uint hk = threadgroup_position_in_grid.x; + const uint r = threadgroup_position_in_grid.y; + const int NCB = meta[3], W = meta[4], RPA = meta[5], CA = meta[6]; + const int st = nodes[NW * (int(r) / G) + 1]; + const int NCBS = meta[MG + st * MS + 5]; + const int ra = meta[MG + st * MS + 3] * 16 + (int(r) / G - meta[MG + st * MS + 6]) * G + int(r) % G; + const int CT = nodes[NW * (int(r) / G) + 3]; // the row's own prefix chunks + constexpr int DP = D / 32, MB = 8; + const int node = r / G, g = r % G; + float m = -INFINITY, l = 0.0f, o[DP]; + for (int i = 0; i < DP; i++) o[i] = 0.0f; + for (int c0 = 0; c0 < CT; c0 += MB) { + float mcs[MB], lcs[MB]; + float4 lo[MB], hi[MB]; + for (int j = 0; j < MB; j++) { + const int64_t row = ((int64_t)hk * CA + min(c0 + j, CT - 1)) * RPA + ra; + mcs[j] = PMA[row]; + lcs[j] = PLA[row]; + lo[j] = *(device const float4*)(POA + row * D + lane * DP); + hi[j] = *(device const float4*)(POA + row * D + lane * DP + 4); + } + for (int j = 0; j < MB; j++) { + if (c0 + j >= CT) break; + const float mc = mcs[j]; + if (mc == -INFINITY) continue; + const float lc = lcs[j]; + const float nm = max(m, mc); + const float f1 = fast::exp(m - nm), f2 = fast::exp(mc - nm); + l = l * f1 + lc * f2; + for (int i = 0; i < 4; i++) o[i] = o[i] * f1 + lo[j][i] * f2; + for (int i = 0; i < 4; i++) o[4 + i] = o[4 + i] * f1 + hi[j][i] * f2; + m = nm; + } + } + for (int c = 0; c < NCBS; c++) { + const int64_t row = (((int64_t)hk * NCB + c) * W + node) * 16 + g; + const float mc = PMB[row]; + if (mc == -INFINITY) continue; + const float lc = PLB[row]; + const float nm = max(m, mc); + const float f1 = fast::exp(m - nm), f2 = fast::exp(mc - nm); + l = l * f1 + lc * f2; + for (int i = 0; i < DP; i++) o[i] = o[i] * f1 + POB[row * D + lane * DP + i] * f2; + m = nm; + } + const int h = hk * G + g; + for (int i = 0; i < DP; i++) { + const int64_t index = ((int64_t)node * (int64_t)heads + h) * D + lane * DP + i; + const bfloat v = static_cast(o[i] / l); + if (gated) { + const float gate = float(QG[((int64_t)node * (int64_t)heads + h) * 2 * D + D + lane * DP + i]); + OUT[index] = static_cast(float(v) * q27_round(1.0f / (1.0f + exp(-gate)), 0)); + } else OUT[index] = v; + } +} + +// A single Mac's 24 heads are known at compile time: main's kernel, so one Mac runs unpaired code. +kernel void q27_attn_merge_single(device const float* POA [[buffer(0)]], device const float* PMA [[buffer(1)]], device const float* PLA [[buffer(2)]], device const float* POB [[buffer(3)]], device const float* PMB [[buffer(4)]], device const float* PLB [[buffer(5)]], device const int* meta [[buffer(6)]], device const int* nodes [[buffer(7)]], device bfloat* OUT [[buffer(8)]], device const bfloat* QG [[buffer(9)]], constant uint& gated [[buffer(10)]], diff --git a/zig/kernels/metal/qwen27/attention_prompt.metal b/zig/kernels/metal/qwen27/attention_prompt.metal index 8ba07e0501..5cd595390e 100644 --- a/zig/kernels/metal/qwen27/attention_prompt.metal +++ b/zig/kernels/metal/qwen27/attention_prompt.metal @@ -4,7 +4,8 @@ using namespace metal; using namespace mpp::tensor_ops; constant constexpr int PG = 6, PD = 256, PH = 128, PCK = 512, PTK = 32; -struct Q27PromptArgs { int start, rows, chunks, rp, key_stride, value_stride; float scale2; }; +// kv_heads is the rank's, not the model's: a split attention writes local node rows. +struct Q27PromptArgs { int start, rows, chunks, rp, key_stride, value_stride; float scale2; int kv_heads; }; typedef vec q27f8; typedef vec q27b8; typedef vec q27h8; @@ -166,5 +167,29 @@ kernel void q27_attn_prompt_merge(device const float* PO [[buffer(0)]], device c m = nm; } const int node = r / PG, head = hk * PG + r % PG; + const int64_t at = ((int64_t)node * (PG * args.kv_heads) + head) * PD + lane * DP; + for (int i = 0; i < DP; i++) OUT[at + i] = static_cast(o[i] / l); +} + +// A single Mac's 24 query heads are known at compile time: main's kernel, so one Mac runs unpaired code. +kernel void q27_attn_prompt_merge_single(device const float* PO [[buffer(0)]], device const float* PM [[buffer(1)]], device const float* PL [[buffer(2)]], + constant Q27PromptArgs& args [[buffer(3)]], device bfloat* OUT [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], uint lane [[thread_index_in_simdgroup]]) { + const int hk = int(tg.x), r = int(tg.y); + constexpr int DP = PD / 32; + float m = -INFINITY, l = 0.0f, o[DP]; + for (int i = 0; i < DP; i++) o[i] = 0.0f; + for (int c = 0; c < args.chunks; c++) { + const int64_t row = ((int64_t)hk * args.chunks + c) * args.rp + r; + const float mc = PM[row]; + if (mc == -INFINITY) continue; + const float lc = PL[row]; + const float nm = max(m, mc); + const float f1 = fast::exp2(m - nm), f2 = fast::exp2(mc - nm); + l = l * f1 + lc * f2; + for (int i = 0; i < DP; i++) o[i] = o[i] * f1 + PO[row * PD + lane * DP + i] * f2; + m = nm; + } + const int node = r / PG, head = hk * PG + r % PG; for (int i = 0; i < DP; i++) OUT[((int64_t)node * (PG * 4) + head) * PD + lane * DP + i] = static_cast(o[i] / l); } diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index 2506fb4905..22cb484ef0 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -34,6 +34,8 @@ pub const core_exlr8_dense = @embedFile("core/exlr8_dense.metal"); pub const core_exlr8_prompt = @embedFile("core/exlr8_prompt.metal"); pub const core_exlr8_dense_prompt = @embedFile("core/exlr8_dense_prompt.metal"); pub const core_pair = @embedFile("core/pair.metal"); +/// The row-split seam's kernels (see `pair_rows.zig`): `pair.metal` above stays as it is for glm_dsa. +pub const core_pair_rows = @embedFile("core/pair_rows.metal"); /// GLM-5.3 (glm_moe_dsa): MLA with rope over a bf16 latent cache, the DSA indexer and its top-k, fp32 glue. pub const glm_dsa_mla = @embedFile("glm_dsa/mla.metal"); diff --git a/zig/src/core/contract.zig b/zig/src/core/contract.zig index 2c976103b2..a04a28ec3f 100644 --- a/zig/src/core/contract.zig +++ b/zig/src/core/contract.zig @@ -21,6 +21,7 @@ test { _ = @import("contract/glm_rows_records.zig"); _ = @import("contract/nemotron_experts_records.zig"); _ = @import("contract/moss_image_records.zig"); + _ = @import("contract/metal_pair_records.zig"); _ = spec; _ = gates; _ = rows; diff --git a/zig/src/core/contract/metal_pair_records.zig b/zig/src/core/contract/metal_pair_records.zig new file mode 100644 index 0000000000..5e9eebe17d --- /dev/null +++ b/zig/src/core/contract/metal_pair_records.zig @@ -0,0 +1,25 @@ +//! The pair seam's four launches (core/pair_rows.metal) with the M1/M2 limits the offline compile measures. +const std = @import("std"); +const limits = @import("limits.zig"); +const registry = @import("registry.zig"); + +pub const records = [_]limits.MetalLaunch{ + .{ .name = "metal pair stage pair_stage", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "metal pair gather pair_gather", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "metal pair argmax pair_argmax", .threads = 1024, .tg_memory = 384, .m1 = 1024, .m2 = 1024 }, + .{ .name = "metal pair pick pair_pick", .threads = 32, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, +}; + +test "the pair seam's four launches are in the live and offline collector and fit every family" { + for (records) |record| { + var found: usize = 0; + for (registry.all_metal) |active| if (std.mem.eql(u8, active.name, record.name)) { + try std.testing.expectEqualDeep(record, active); + found += 1; + }; + try std.testing.expectEqual(@as(usize, 1), found); + try std.testing.expectEqual(@as(usize, 0), try registry.checkMetal(std.testing.allocator, &.{record}, true)); + // A zero here would be a launch no M1/M2 compiler was asked about. + try std.testing.expect(record.m1 > 0 and record.m2 > 0); + } +} diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 7d83d71249..e3a4f9abd9 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -572,7 +572,7 @@ test { _ = paging; } /// The Metal collector keeps measured slots and the restored vision declarations. -pub const all_metal = metal ++ @import("delta_slots_records.zig").records ++ @import("flashnext_pending.zig").records ++ @import("simd_pairs_records.zig").records ++ @import("glm_rows_records.zig").records ++ @import("weight_only_records.zig").metal ++ @import("hc_activation_records.zig").metal ++ @import("qwen4_records.zig").metal ++ @import("lfm2_records.zig").records ++ @import("nemotron_experts_records.zig").records ++ +pub const all_metal = metal ++ @import("delta_slots_records.zig").records ++ @import("flashnext_pending.zig").records ++ @import("simd_pairs_records.zig").records ++ @import("metal_pair_records.zig").records ++ @import("glm_rows_records.zig").records ++ @import("weight_only_records.zig").metal ++ @import("hc_activation_records.zig").metal ++ @import("qwen4_records.zig").metal ++ @import("lfm2_records.zig").records ++ @import("nemotron_experts_records.zig").records ++ @import("moss_image_records.zig").records; const learning = @import("registry_learning.zig"); pub const gradient_pending = learning.gradient_pending; diff --git a/zig/src/core/pair_rows.zig b/zig/src/core/pair_rows.zig new file mode 100644 index 0000000000..38017cb531 --- /dev/null +++ b/zig/src/core/pair_rows.zig @@ -0,0 +1,463 @@ +//! Two Macs' exchanges: GPU part, posted(x), host MCDMA relay, landed(x); the event wait is a real scheduler hand-off. +const std = @import("std"); +const platform = @import("sys"); +const mtl = @import("metal"); +const fabric = @import("fabric"); +const sources = @import("kernel_sources"); +const settings = @import("../families/flashnext/tp_settings.zig"); +const Fence = @import("fence.zig").Fence; +const Owner = @import("../families/flashnext/owned.zig").Owner; +const pair_layout = @import("pair_rows_layout.zig"); +const pair_loopback = @import("pair_rows_loopback.zig"); + +pub const Settings = settings.Settings; +pub const Ref = @import("exlr8_projection.zig").Ref; + +pub const command_words = pair_layout.command_words; +pub const Slot = pair_layout.Slot; +pub const seam_failed = pair_layout.seam_failed; +pub const Seam = pair_layout.Seam; +pub const elems_per_word = pair_layout.elems_per_word; +pub const stageBytes = pair_layout.stageBytes; +pub const gatherBytes = pair_layout.gatherBytes; +pub const pickBytes = pair_layout.pickBytes; +pub const frameBytes = pair_layout.frameBytes; +pub const toFrame = pair_layout.toFrame; +pub const seedValue = pair_layout.seedValue; +pub const Loopback = pair_loopback.Loopback; + +const PAGE = pair_layout.PAGE; +const pages = pair_layout.pages; +const CMD_BYTES = pair_layout.CMD_BYTES; +const JOBS = pair_layout.JOBS; +const slotOf = pair_layout.slotOf; +const Layout = pair_layout.Layout; +const FLAG = pair_layout.FLAG; +const HELLO = pair_layout.HELLO; +const CMD_FLAG = pair_layout.CMD_FLAG; +const CMD_ACK = pair_layout.CMD_ACK; +const READY = pair_layout.READY; +const BEAT = pair_layout.BEAT; +const POSTED = pair_layout.POSTED; +const FAILED = pair_layout.FAILED; +const OBSERVED = pair_layout.OBSERVED; +const PairRows = pair_layout.PairRows; +const seam = pair_layout.seam; + +/// Mach's clock: 24 MHz, so a millisecond is 24,000 ticks -- the clock the probes' seam numbers are in. +const ticks_per_ms: u64 = 24_000; + +/// How long the host waits for the GPU to reach a published seam: a hung GPU must fail the request, not wedge it. +const post_timeout_ms: u64 = 5_000; +/// The peer's part wait: 250 ms, 300x the measured tail (p99 ~700 us), so a dead peer fails fast. +const peer_timeout_ms: u64 = 250; +/// The cold-link peer wait: the first request pays one-time Metal pipelines (1.07 s measured, 0.14 s warm). +const cold_peer_timeout_ms: u64 = 5_000; +/// How long after the link's first relay the cold bound applies: 8 s covers a longer first prompt (prefill chunks). +const cold_window_ms: u64 = 8_000; + +/// How long rank 1 waits for rank 0's acknowledgement of a replay: a stuck peer fails the request, it is not wedged. +const cmd_ack_timeout_ms: u64 = 60_000; + +/// How long either rank waits for the other's ready word: a cold 27B load on the far Mac, five minutes at most. +const ready_timeout_ms: u64 = 300_000; + +/// The close's last flush: microseconds when the peer is there, and never the library's whole timeout when it is gone. +const last_flush_ns: u64 = 2 * std.time.ns_per_s; + +threadlocal var raised = false; + +/// This thread at the user-interactive QoS once: busy CPUs must not delay the pair's sends and commands. +pub fn raise() void { + if (raised) return; + _ = platform.c.pthread_set_qos_class_self_np(.USER_INTERACTIVE, 0); + raised = true; +} + +/// Whether 32-bit counter a has reached b across the wrap. +fn reached(a: u32, b: u32) bool { + return @as(i32, @bitCast(a -% b)) >= 0; +} + +fn ticksNow() u64 { + return platform.c.mach_absolute_time(); +} + +/// A short sleep, for the host threads' waits that are not on the seam's critical path. +fn sleepUs(us: u64) void { + const ts: platform.c.timespec = .{ .sec = @intCast(us / 1_000_000), .nsec = @intCast(us % 1_000_000 * 1000) }; + _ = platform.c.nanosleep(&ts, null); +} + +pub const Pair = struct { + rank: u32, + peer: u32, + l: Layout, + ep: *fabric.mcdma.Endpoint, + rd: fabric.rdma.Rdma, + win: []u8, + wbuf: mtl.Buffer, + stage_pipe: mtl.Pipeline, + gather_pipe: mtl.Pipeline, + argmax_pipe: mtl.Pipeline, + pick_pipe: mtl.Pipeline, + /// Orders the window's untracked writes across the encoders a seam splits. + fence: Fence, + /// GPU -> host: seam x's part is written, relay it. Host -> GPU: the peer's part for x is in place, add it. + posted: mtl.SharedEvent, + landed: mtl.SharedEvent, + /// Exchanges encoded so far: the same count on both Macs. + x: u64 = 0, + bytes: []u32, + queued: std.atomic.Value(u32) = .init(0), + thread: ?std.Thread = null, + stop: std.atomic.Value(bool) = .init(false), + broken: std.atomic.Value(bool) = .init(false), + /// The exchange the service thread is on: the round number in a failure line. + seam_x: std.atomic.Value(u32) = .init(0), + /// The request this pair is on, for failure lines; `beginRound` bumps it -- a label only, failures are permanent. + round: std.atomic.Value(u32) = .init(0), + /// Commands sent (rank 0) or taken (rank 1) so far. + cmd: u64 = 0, + /// Beats this Mac's watchdog has signalled to the peer: its counter for the peer's silence check. + beats: u64 = 0, + /// The GPU objects this pair holds: released in `deinit`, the window's buffer before the memory behind it. + owned: Owner(mtl) = .{}, + /// The allocator `init` was given: `bytes`, the settings arena and this struct go back to it. + gpa: std.mem.Allocator, + /// The parsed settings file: the endpoint's config points into it, so it outlives `ep`. + arena: *std.heap.ArenaAllocator, + + /// Connect to the peer named in `settings_path` (speed-up mode's file), exchanges up to `bytes`, and start sending. + pub fn init(gpa: std.mem.Allocator, io: std.Io, device: mtl.Device, settings_path: []const u8, bytes: usize) !*Pair { + const l: Layout = .{ .slot = pages(bytes) }; + const text = try std.Io.Dir.cwd().readFileAlloc(io, settings_path, gpa, .limited(1 << 20)); + defer gpa.free(text); + const arena = try gpa.create(std.heap.ArenaAllocator); + errdefer gpa.destroy(arena); + arena.* = .init(gpa); + errdefer arena.deinit(); + // the endpoint's config keeps `s.links`, so the parse outlives `text` in the arena, not on the stack + const s = try settings.read(arena.allocator(), text); + if (s.rank > 1 or s.links.len != 1) return error.PairTwoRanksOneLink; + const lib = try gpa.dupeSentinel(u8, s.library, 0); + defer gpa.free(lib); + const ep = try fabric.mcdma.Endpoint.create(gpa, lib, .{ .rank = s.rank, .ranks = 2, .window_bytes = l.window(), .staging_bytes = 4 << 20, .links = s.links, .timeout_ns = 60 * std.time.ns_per_s, .connect_timeout_ns = 300 * std.time.ns_per_s }); + errdefer ep.deinit(); + var owned: Owner(mtl) = .{}; + errdefer owned.deinit(); + const rd = ep.rdma(); + const win = rd.window(); + const lib_m = try mtl.Library.fromSource(device, sources.core_pair_rows, mtl.CompileOptions.mlx()); + defer lib_m.deinit(); + const fence = try Fence.init(device); + errdefer fence.deinit(); + const posted = try device.sharedEvent(); + errdefer posted.deinit(); + const landed = try device.sharedEvent(); + errdefer landed.deinit(); + const p = try gpa.create(Pair); + errdefer gpa.destroy(p); + const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; + p.* = .{ + .rank = s.rank, + .peer = 1 - s.rank, + .l = l, + .ep = ep, + .rd = rd, + .win = win, + .wbuf = try owned.keep(gpa, try device.bufferNoCopy(win.ptr, win.len, opts)), + .stage_pipe = try owned.keep(gpa, try mtl.Pipeline.init(device, lib_m, "pair_stage", false)), + .gather_pipe = try owned.keep(gpa, try mtl.Pipeline.init(device, lib_m, "pair_gather", false)), + .argmax_pipe = try owned.keep(gpa, try mtl.Pipeline.init(device, lib_m, "pair_argmax", false)), + .pick_pipe = try owned.keep(gpa, try mtl.Pipeline.init(device, lib_m, "pair_pick", false)), + .fence = fence, + .posted = posted, + .landed = landed, + .bytes = try gpa.alloc(u32, JOBS), + .gpa = gpa, + .arena = arena, + .owned = owned, + }; + owned = .{}; // the pair owns them now: the error path must not release them twice + errdefer gpa.free(p.bytes); + // both Macs up before the first exchange: a word each way + try rd.signal(p.peer, l.flags() + HELLO, 1); + while (@atomicLoad(u64, p.word64(l.flags() + HELLO), .acquire) < 1) std.atomic.spinLoopHint(); + p.thread = try std.Thread.spawn(.{}, service, .{p}); + return p; + } + + /// Rank 0 tells rank 1 it is done (its `next` returns null); then the sending thread and the link close. + pub fn deinit(p: *Pair) void { + if (p.rank == 0 and !p.failed()) p.send(0, &.{}) catch {}; + p.stop.store(true, .release); + if (p.thread) |t| t.join(); + if (!p.failed()) flushLast(p); + p.posted.deinit(); + p.landed.deinit(); + p.fence.deinit(); + p.owned.deinit(); // the window's buffer before the endpoint frees the memory it aliases + p.ep.deinit(); + const gpa = p.gpa; + gpa.free(p.bytes); + p.arena.deinit(); + gpa.destroy(p.arena); + gpa.destroy(p); + } + + /// Hand the last seam over without the library's whole timeout: a peer that is gone leaves nothing to deliver. + fn flushLast(p: *Pair) void { + if (fabric.words.checkedNowNs()) |now| p.rd.setDeadline(now + last_flush_ns) else |_| return; + p.rd.flush() catch {}; + } + + fn word64(p: *const Pair, off: usize) *u64 { + return @ptrCast(@alignCast(p.win.ptr + off)); + } + fn word32(p: *const Pair, off: usize) *u32 { + return @ptrCast(@alignCast(p.win.ptr + off)); + } + + /// Seams the host gave up on: nonzero means the run is invalid. + pub fn gaveUp(p: *const Pair) u32 { + return @atomicLoad(u32, p.word32(p.l.sync() + FAILED * 4), .acquire); + } + + /// Seams the GPU itself refused after seeing FAILED: this side's own record of a failed relay. + pub fn observed(p: *const Pair) u32 { + return @atomicLoad(u32, p.word32(p.l.sync() + OBSERVED * 4), .acquire); + } + + /// The exchange the service thread is on (or the one it gave up on, if it has). + pub fn failedSeam(p: *const Pair) u32 { + return p.seam_x.load(.acquire); + } + + /// Labels a request; never a reset: after a seam failure the pair is dead (a half-written ring corrupts). + pub fn beginRound(p: *Pair) u32 { + return p.round.fetchAdd(1, .monotonic) + 1; + } + + /// Where this rank's part for exchange x goes, and where the peer's lands: the relay's two addresses. + pub fn sendOffset(p: *const Pair, x: u32) usize { + return p.l.send(x); + } + pub fn recvOffset(p: *const Pair, x: u32) usize { + return p.l.recv(x); + } + + /// The link is down: failed relay, dead peer, GPU never at the seam; check after every seam and command buffer. + pub fn failed(p: *const Pair) bool { + return p.broken.load(.acquire) or p.gaveUp() != 0; + } + + /// The peer's watchdog is alive: its last beat, a word its watchdog wrote here; silence means the peer is gone. + pub fn heartbeats(p: *const Pair) u64 { + return @atomicLoad(u64, p.word64(p.l.flags() + pair_layout.BEAT), .acquire); + } + + /// Tell the peer this rank's watchdog is alive; a failed signal is the peer gone, not nothing to do. + pub fn beat(p: *Pair) void { + p.beats +%= 1; + p.rd.signal(p.peer, p.l.flags() + pair_layout.BEAT, p.beats) catch p.broken.store(true, .release); + } + + /// The next exchange: take it before the producer writes, then pass the same x to stage/gather/argmax/pick. + pub fn nextX(p: *const Pair) u32 { + return @truncate(p.x + 1); + } + + /// The exchange slot for this seam's part: where this rank's rows go, and where the peer's land. + pub fn part(p: *const Pair) Seam { + const x: u32 = @truncate(p.x + 1); + return .{ .x = x, .ref = .{ .buf = p.wbuf, .off = p.l.send(x) } }; + } + + /// Publish x and its bytes to the service thread; x must be `part`'s, so an out-of-order one panics in release too. + fn queue(p: *Pair, x: u32, bytes: usize) void { + if (x != @as(u32, @truncate(p.x + 1))) @panic("pair: an exchange was encoded out of order (part() not called first)"); + if (bytes > p.l.slot) @panic("pair: this exchange's part does not fit the window's slot (speed-up's rows too large)"); + p.x += 1; + p.bytes[x % JOBS] = @intCast(bytes); + p.queued.store(x, .release); + } + + /// Before `seam`: this rank's `width` columns of `rows` rows from `src` into the slot, bf16 to bf16, no arithmetic. + pub fn stage(p: *Pair, e: *mtl.ComputeEncoder, x: u32, src: Ref, rows: u32, width: u32, stride: u32) void { + if (rows == 0 or width == 0 or width > stride) @panic("pair: stage() wants rows > 0 and a width inside the source row"); + e.setPipeline(p.stage_pipe); + e.setBuffer(p.wbuf, p.l.sync(), 0); + e.setValue(PairRows{ .width = width, .stride = stride, .rank = p.rank, .x = x }, 1); + e.setBuffer(src.buf, src.off, 2); + e.setBuffer(p.wbuf, p.l.send(x), 3); + e.dispatchGroups(.{ .width = rows }, .{ .width = 256 }); + } + + /// After `stage` (or the producer): relay, then pack this rank's `width` at `rank * width` and the peer's after. + pub fn gather(p: *Pair, e: *mtl.ComputeEncoder, x: u32, mine: Ref, out: Ref, rows: u32, width: u32, out_stride: u32) void { + if (rows == 0 or width == 0 or width % elems_per_word != 0) @panic("pair: gather() wants rows > 0 and a width that fills whole window words"); + if (out_stride < 2 * width) @panic("pair: gather()'s frame row is narrower than both halves"); + if (mine.buf.id != p.wbuf.id or mine.off != p.l.send(x)) @panic("pair: gather()'s part is not the region part() returned"); + p.queue(x, stageBytes(rows, width)); + seam(e, p.fence, p.posted, p.landed, x); + e.setPipeline(p.gather_pipe); + e.setBuffer(p.wbuf, p.l.sync(), 0); + e.setValue(PairRows{ .width = width, .stride = out_stride, .rank = p.rank, .x = x }, 1); + e.setBuffer(mine.buf, mine.off, 2); + e.setBuffer(p.wbuf, p.l.recv(x), 3); + e.setBuffer(out.buf, out.off, 4); + e.dispatchGroups(.{ .width = rows }, .{ .width = 256 }); + } + + /// Each row's argmax over this rank's n fp32 logits from id lo, row stride `stride`, then the pick over both Macs'. + pub fn argmax(p: *Pair, e: *mtl.ComputeEncoder, x: u32, logits: Ref, stride: u32, lo: u32, n: u32, rows: u32, picks: Ref) void { + p.argmaxWidth(e, x, logits, stride, lo, n, rows, picks, 0); + } + + /// The same seam over bf16 logits (Qwen3.8-27B's `frame.logits`): one kernel, `dims.w` selects the load width. + pub fn argmaxBf16(p: *Pair, e: *mtl.ComputeEncoder, x: u32, logits: Ref, stride: u32, lo: u32, n: u32, rows: u32, picks: Ref) void { + p.argmaxWidth(e, x, logits, stride, lo, n, rows, picks, 1); + } + + /// `dtype` 0 is fp32 logits, 1 is bf16; the exchange is (fp32 value, id) either way, so pair_pick is one path. + fn argmaxWidth(p: *Pair, e: *mtl.ComputeEncoder, x: u32, logits: Ref, stride: u32, lo: u32, n: u32, rows: u32, picks: Ref, dtype: u32) void { + if (rows == 0 or n == 0) @panic("pair: argmax() wants rows > 0 and n > 0"); + p.queue(x, pickBytes(rows)); + e.setPipeline(p.argmax_pipe); + e.setBuffer(logits.buf, logits.off, 0); + e.setValue([4]u32{ stride, lo, n, dtype }, 1); + e.setBuffer(p.wbuf, p.l.send(x), 2); + e.dispatchGroups(.{ .width = rows }, .{ .width = 1024 }); + seam(e, p.fence, p.posted, p.landed, x); // the encoder boundary orders the argmax before the part is relayed + e.setPipeline(p.pick_pipe); + e.setBuffer(p.wbuf, p.l.sync(), 0); + e.setValue([2]u32{ rows, x }, 1); + e.setBuffer(p.wbuf, p.l.send(x), 2); + e.setBuffer(p.wbuf, p.l.recv(x), 3); + e.setBuffer(picks.buf, picks.off, 4); + e.dispatchGroups(.{ .width = 1 }, .{ .width = 32 }); + } + + /// Rank 0: command `kind` (never 0, which ends rank 1's loop) and its words, once rank 1 took the one before. + pub fn send(p: *Pair, kind: u32, words: []const u32) !void { + if (p.rank != 0) return error.FollowsPeer; + if (words.len > command_words) return error.CommandTooLong; + raise(); + try p.reach(p.l.flags() + CMD_ACK, p.cmd, cmd_ack_timeout_ms); + p.cmd += 1; + const head: [2]u32 = .{ kind, @intCast(words.len) }; + const out = p.l.cmdOut(); + @memcpy(p.win[out..][0..8], std.mem.asBytes(&head)); + @memcpy(p.win[out + 8 ..][0 .. 4 * words.len], std.mem.sliceAsBytes(words)); + try p.ep.writeSignalFrom(p.peer, out, p.l.cmdIn(), 8 + 4 * words.len, p.l.flags() + CMD_FLAG, p.cmd); + } + + /// Rank 1: rank 0's next command's words into `out` and its kind; null once rank 0 is done or this Mac stops. + pub fn next(p: *Pair, gpa: std.mem.Allocator, out: *std.ArrayList(u32)) !?u32 { + raise(); + p.reach(p.l.flags() + CMD_FLAG, p.cmd + 1, 0) catch |err| return if (err == error.PairStopped) null else err; + p.cmd += 1; + const in = p.l.cmdIn(); + const head: *const [2]u32 = @ptrCast(@alignCast(p.win.ptr + in)); + if (head[1] > command_words) return error.CommandOutOfStep; + out.clearRetainingCapacity(); + try out.appendSlice(gpa, @as([*]const u32, @ptrCast(@alignCast(p.win.ptr + in + 8)))[0..head[1]]); + const kind = head[0]; + try p.rd.signal(p.peer, p.l.flags() + CMD_ACK, p.cmd); + return if (kind == 0) null else kind; + } + + /// Both Macs loaded before the first command: a word each way, each waited for with the load's bound. + pub fn ready(p: *Pair) !void { + try p.rd.signal(p.peer, p.l.flags() + READY, 1); + try p.reach(p.l.flags() + READY, 1, ready_timeout_ms); + } + + /// This Mac's command wait ends (rank 1 at close). + pub fn stopFollowing(p: *Pair) void { + p.stop.store(true, .release); + } + + /// Wait until this Mac's word at `off` is `value`: spins, then short sleeps, or fails after `limit_ms` (0: none). + fn reach(p: *Pair, off: usize, value: u64, limit_ms: u64) !void { + const t0 = ticksNow(); + var spins: usize = 0; + while (@atomicLoad(u64, p.word64(off), .acquire) < value) { + if (p.stop.load(.acquire)) return error.PairStopped; + if (p.failed()) return error.PairLinkFailed; + if (limit_ms != 0 and ticksNow() -% t0 > limit_ms * ticks_per_ms) return error.PairPeerSilent; + spins += 1; + if (spins > 200_000) sleepUs(20) else std.atomic.spinLoopHint(); + } + } + + /// Give up on x: release every event wait with the sentinel, so the blocked GPU resumes and the request aborts. + fn fail(p: *Pair, x: u32, comptime what: []const u8, peer_ms: u64) void { + p.broken.store(true, .release); + _ = @atomicRmw(u32, p.word32(p.l.sync() + FAILED * 4), .Add, 1, .monotonic); + std.log.err("pair rank {d}: request {d} seam {d} failed: {s}; the link is down and stays down (posted {d}, flag {d}, queued {d}, GPU-observed {d}; deadlines {d} ms post, {d} ms peer): restart both processes. The blocked GPU is released with the failure sentinel", .{ + p.rank, p.round.load(.acquire), x, what, p.posted.value(), @atomicLoad(u64, p.word64(p.l.flags() + FLAG), .acquire), p.queued.load(.acquire), p.observed(), post_timeout_ms, peer_ms, + }); + p.landed.set(seam_failed); + } + + /// One exchange at a time: part written, relay it, peer's lands, release the GPU; every wait has a deadline. + fn service(p: *Pair) void { + raise(); + var x: u32 = 1; + // Ticks after which this link is no longer cold: set at the first relay, when both ranks have just loaded. + var cold_until: u64 = 0; + while (!p.stop.load(.acquire)) { + p.seam_x.store(x, .release); + // 1. The GPU must reach the seam: sleep while idle, spin once published, fail if it never arrives. + var armed: u64 = 0; + while (p.posted.value() < @as(u64, x)) { + if (p.stop.load(.acquire)) return; + if (!reached(p.queued.load(.acquire), x)) { + sleepUs(20); + continue; + } + if (armed == 0) armed = ticksNow(); + if (ticksNow() -% armed > post_timeout_ms * ticks_per_ms) return p.fail(x, "this Mac's GPU never reached the seam", peer_timeout_ms); + std.atomic.spinLoopHint(); + } + // 2. Relay: part to the peer's recv slot, arrival flag in the same MCDMA write, data and count together. + const at = x % JOBS; + p.ep.writeSignalFrom(p.peer, p.l.send(x), p.l.recv(x), p.bytes[at], p.l.flags() + FLAG, x) catch |err| { + std.log.err("pair rank {d}: seam {d} relay write failed: {s}", .{ p.rank, x, @errorName(err) }); + return p.fail(x, "the relay write failed", peer_timeout_ms); + }; + // 3. The peer's relay puts its part in our slot; its flag is proof; the cold bound, then the steady one. + if (x == 1) cold_until = ticksNow() +% cold_window_ms * ticks_per_ms; + const cold = cold_until != 0 and ticksNow() < cold_until; + const budget_ms = if (cold) cold_peer_timeout_ms else peer_timeout_ms; + const t1 = ticksNow(); + while (@atomicLoad(u64, p.word64(p.l.flags() + FLAG), .acquire) < x) { + if (p.stop.load(.acquire)) return; + if (ticksNow() -% t1 > budget_ms * ticks_per_ms) return p.fail(x, "the peer's part never landed", budget_ms); + std.atomic.spinLoopHint(); + } + // 4. Release the GPU: both parts are in; the event's wait is the acquire that makes them visible. + @atomicStore(u32, p.word32(p.l.sync() + POSTED * 4), x, .release); + p.landed.set(x); + x +%= 1; + } + } +}; + +test "the window's regions do not overlap and its words sit on their pages" { + const l: Layout = .{ .slot = pages(stageBytes(256, 6144)) }; + try std.testing.expect(l.send(1) >= l.recv(1) + l.slot and l.send(0) >= 2 * l.slot + PAGE and l.send(0) + l.slot <= l.send(1)); + try std.testing.expect(l.cmdIn() >= l.send(1) + l.slot and l.cmdOut() >= l.cmdIn() + CMD_BYTES + 64); + try std.testing.expect(l.flags() >= l.cmdOut() + 8 + 4 * command_words and BEAT + 8 <= PAGE and l.window() % PAGE == 0); + try std.testing.expect(reached(1, 0xFFFFFFFF) and !reached(5, 6)); +} + +test "the parity ring alternates and one sentinel releases every exchange" { + try std.testing.expect(slotOf(1) == .b and slotOf(2) == .a and slotOf(3) == .b); + const l: Layout = .{ .slot = pages(stageBytes(16, 5120)) }; + try std.testing.expect(l.recv(1) == l.slot and l.recv(2) == 0 and l.send(2) == 2 * l.slot + PAGE); + try std.testing.expectEqual(pages(stageBytes(16, 5120)), l.slot); // the 27B's 16-row prefill part, page-rounded + try std.testing.expect(seam_failed > std.math.maxInt(u32)); // every waitFor(x) in a buffer is satisfied by it +} diff --git a/zig/src/core/pair_rows_layout.zig b/zig/src/core/pair_rows_layout.zig new file mode 100644 index 0000000000..a1a1ff14d6 --- /dev/null +++ b/zig/src/core/pair_rows_layout.zig @@ -0,0 +1,125 @@ +//! The seam window's regions and the bf16 frame ABI, shared by the link and the host-copy Loopback. +const std = @import("std"); +const mtl = @import("metal"); +const Fence = @import("fence.zig").Fence; + +pub const Ref = @import("exlr8_projection.zig").Ref; + +pub const PAGE = 16384; +pub fn pages(n: usize) usize { + return (n + PAGE - 1) / PAGE * PAGE; +} +/// A command's words at most: a 262,144-token prompt and its head. +pub const command_words = 262144 + 64; +pub const CMD_BYTES = pages(8 + 4 * command_words); +pub const JOBS = 4096; + +/// Which parity slot x uses: rank r cannot finish x without the peer's x, so the peer is at most one exchange ahead. +pub const Slot = enum(u1) { a = 0, b = 1 }; + +pub fn slotOf(x: u32) Slot { + return @fromBackingInt(@intCast(x % 2)); +} + +/// The window on both Macs: parts in/out by parity (a page before each out), commands, flags, host words. +pub const Layout = struct { + slot: usize, + + pub fn recv(l: Layout, x: u32) usize { + return @as(usize, @backingInt(slotOf(x))) * l.slot; + } + pub fn send(l: Layout, x: u32) usize { + return 2 * l.slot + @as(usize, @backingInt(slotOf(x))) * (PAGE + l.slot) + PAGE; + } + pub fn cmdIn(l: Layout) usize { + return 4 * l.slot + 2 * PAGE; + } + pub fn cmdOut(l: Layout) usize { + return l.cmdIn() + CMD_BYTES + PAGE; + } + pub fn flags(l: Layout) usize { + return l.cmdOut() + CMD_BYTES; + } + pub fn sync(l: Layout) usize { + return l.flags() + PAGE; + } + pub fn window(l: Layout) usize { + return l.sync() + PAGE; + } +}; +pub const FLAG = 0; // u64 in the flag page: the last exchange whose peer part has landed (written by the peer's relay) +pub const HELLO = 64; +pub const CMD_FLAG = 128; // u64, at rank 1: the last command rank 0 has written +pub const CMD_ACK = 192; // u64, at rank 0: the last command rank 1 has copied out +pub const READY = 256; // u64: the peer has loaded its model +pub const BEAT = 320; // u64: the peer's watchdog ran; a rank idle in no round still sees a dead peer through this + +// The sync page's words shared with pair_rows.metal: POSTED last released exchange, FAILED give-ups, OBSERVED refusals. +pub const POSTED = 0; +pub const FAILED = 16; +pub const OBSERVED = 20; + +pub const seam_failed: u64 = std.math.maxInt(u64); + +pub const Seam = struct { x: u32, ref: Ref }; + +/// Elements in one 32-bit window word: the relay and the link move whole words, and a frame element is a bf16. +pub const elems_per_word: usize = 2; + +/// `stage`'s outgoing copy: this rank's `cols` of `rows` rows, bf16 -- the frame seam's payload. +pub fn stageBytes(rows: u32, cols: u32) usize { + return @as(usize, rows) * cols * @sizeOf(u16); +} + +/// `gather`'s `out`: both halves of `rows` rows packed into the frame, so the frame row is `2 * cols` wide. +pub fn gatherBytes(rows: u32, cols: u32) usize { + return @as(usize, rows) * 2 * cols * @sizeOf(u16); +} + +/// The head seam's exchange: an (fp32 value, u32 id) a row; pair_pick compares logits, not words. +pub fn pickBytes(rows: u32) usize { + return @as(usize, rows) * @sizeOf(f32) * 2; +} + +/// A frame tensor's bytes: the bf16 grid `stage` reads and `gather` writes its halves into. +pub fn frameBytes(rows: u32, width: u32) usize { + return @as(usize, rows) * width * @sizeOf(u16); +} + +/// One fp32 value as the frame stores it (bf16, nearest even), the host side of the kernels' frame elements. +pub fn toFrame(v: f32) u16 { + return bf16Bits(v); +} + +/// One rounding to bf16, nearest-even: the host side of a frame element (pair_rows.metal copies the bits it is given). +fn bf16Bits(v: f32) u16 { + const u: u32 = @bitCast(v); + return @truncate((u +% 0x7FFF +% ((u >> 16) & 1)) >> 16); +} + +/// The seed pattern as a frame element: 1 + m/128, exact in both bf16 and fp32, so expected bytes hold either way. +pub fn seedValue(r: u32, i: u32, salt: u16) f32 { + const k = r *% 5120 +% i *% 37; + const m: u32 = (salt & 0x1F) + ((k *% 53) & 0x1F); + return 1.0 + @as(f32, @floatFromInt(m)) / 128.0; +} + +/// One frame element at `at` of a bf16 payload: `stage`'s and `gather`'s slot is the frame's bits, two bytes each. +pub fn writeHalf(out: []u8, at: usize, v: f32) void { + std.mem.writeInt(u16, out[at * @sizeOf(u16) ..][0..@sizeOf(u16)], toFrame(v), .little); +} + +/// The kernels' argument block: four u32, so its layout and pair_rows.metal's struct stay one struct. +pub const PairRows = extern struct { width: u32, stride: u32, rank: u32, x: u32 }; + +/// Between encoders only (core/segments.zig:59-66): fence, close, posted, wait, reopen, wait -- a real hand-off. +pub fn seam(enc: *mtl.ComputeEncoder, fence: Fence, posted: mtl.SharedEvent, landed: mtl.SharedEvent, x: u32) void { + const cb = enc.commandBuffer(); + const mode = enc.dispatchType(); + fence.update(enc.*); + enc.end(); + cb.signal(posted, @as(u64, x)); + cb.waitFor(landed, @as(u64, x)); + enc.* = cb.compute(mode); + fence.wait(enc.*); +} diff --git a/zig/src/core/pair_rows_loopback.zig b/zig/src/core/pair_rows_loopback.zig new file mode 100644 index 0000000000..5086fa1e3e --- /dev/null +++ b/zig/src/core/pair_rows_loopback.zig @@ -0,0 +1,249 @@ +//! The same seam with a host copy instead of the link: the numeric reference the real transport is checked against. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const Fence = @import("fence.zig").Fence; +const pair_layout = @import("pair_rows_layout.zig"); + +const FAILED = pair_layout.FAILED; +const JOBS = pair_layout.JOBS; +const Layout = pair_layout.Layout; +const PAGE = pair_layout.PAGE; +const PairRows = pair_layout.PairRows; +const Seam = pair_layout.Seam; +const elems_per_word = pair_layout.elems_per_word; +const pages = pair_layout.pages; +const stageBytes = pair_layout.stageBytes; +const pickBytes = pair_layout.pickBytes; +const seam = pair_layout.seam; +const seam_failed = pair_layout.seam_failed; +const writeHalf = pair_layout.writeHalf; +const Ref = pair_layout.Ref; + +/// The same seam with a host copy instead of the link: the numeric reference; bytes arrive via `arrive`/`seed`. +pub const Loopback = struct { + allocator: std.mem.Allocator, + rank: u32, + l: Layout, + win: []u8, + wbuf: mtl.Buffer, + stage_pipe: mtl.Pipeline, + gather_pipe: mtl.Pipeline, + argmax_pipe: mtl.Pipeline, + pick_pipe: mtl.Pipeline, + /// Per-exchange part bytes (stage/gather and argmax queue different sizes): the relay copies exactly this much. + bytes: []u32, + fence: Fence, + /// This rank's parts: the GPU signals it per exchange, and the relay reads `bytesFor(x)` after that. + posted: mtl.SharedEvent, + /// The peer's parts: the relay sets it once the copy has landed, and the seam waits for it. + landed: mtl.SharedEvent, + x: u64 = 0, + relay: ?std.Thread = null, + stop: std.atomic.Value(bool) = .init(false), + + pub fn init(device: mtl.Device, bytes: usize, rank: u32) !Loopback { + const l: Layout = .{ .slot = pages(bytes) }; + // the window is a no-copy Metal buffer: page-aligned, page-multiple allocation, never the gpa's good luck + const win = try std.heap.page_allocator.alloc(u8, l.window()); + errdefer std.heap.page_allocator.free(win); + std.debug.assert(@intFromPtr(win.ptr) % PAGE == 0 and win.len % PAGE == 0); + @memset(win, 0); + const lib = try mtl.Library.fromSource(device, sources.core_pair_rows, mtl.CompileOptions.mlx()); + defer lib.deinit(); + const wbuf = try device.bufferNoCopy(win.ptr, win.len, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + errdefer wbuf.deinit(); + const fence = try Fence.init(device); + errdefer fence.deinit(); + const posted = try device.sharedEvent(); + errdefer posted.deinit(); + const landed = try device.sharedEvent(); + errdefer landed.deinit(); + const stage_pipe = try mtl.Pipeline.init(device, lib, "pair_stage", false); + errdefer stage_pipe.deinit(); + const gather_pipe = try mtl.Pipeline.init(device, lib, "pair_gather", false); + errdefer gather_pipe.deinit(); + const argmax_pipe = try mtl.Pipeline.init(device, lib, "pair_argmax", false); + errdefer argmax_pipe.deinit(); + const pick_pipe = try mtl.Pipeline.init(device, lib, "pair_pick", false); + errdefer pick_pipe.deinit(); + const sizes = try std.heap.page_allocator.alloc(u32, JOBS); + errdefer std.heap.page_allocator.free(sizes); + @memset(sizes, 0); + return .{ .allocator = std.heap.page_allocator, .rank = rank, .l = l, .win = win, .wbuf = wbuf, .stage_pipe = stage_pipe, .gather_pipe = gather_pipe, .argmax_pipe = argmax_pipe, .pick_pipe = pick_pipe, .bytes = sizes, .fence = fence, .posted = posted, .landed = landed }; + } + + pub fn deinit(p: *Loopback) void { + p.stopRelay(); + p.stage_pipe.deinit(); + p.gather_pipe.deinit(); + p.argmax_pipe.deinit(); + p.pick_pipe.deinit(); + p.allocator.free(p.bytes); + p.wbuf.deinit(); + p.posted.deinit(); + p.landed.deinit(); + p.fence.deinit(); + p.allocator.free(p.win); + } + + fn word32(p: *const Loopback, off: usize) *u32 { + return @ptrCast(@alignCast(p.win.ptr + off)); + } + + /// Where this rank's part for exchange x goes (what the relay copies out, once `posted` passes x). + pub fn sendOffset(p: *const Loopback, x: u32) usize { + return p.l.send(x); + } + + /// Where the peer's part for exchange x goes (what the relay copies in, before setting `landed` to x). + pub fn recvOffset(p: *const Loopback, x: u32) usize { + return p.l.recv(x); + } + + /// Where the producer writes this rank's bf16 columns of the next exchange; the same shape as `Pair.part`. + pub fn nextX(p: *const Loopback) u32 { + return @truncate(p.x + 1); + } + + pub fn part(p: *const Loopback) Seam { + const x: u32 = @truncate(p.x + 1); + return .{ .x = x, .ref = .{ .buf = p.wbuf, .off = p.l.send(x) } }; + } + + /// The seed pattern as a frame element: 1 + m/128, exact in bf16, so expected bytes hold either way. + fn seedValue(r: u32, i: u32, salt: u16) f32 { + const k = r *% 5120 +% i *% 37; + const m: u32 = (salt & 0x1F) + ((k *% 53) & 0x1F); + return 1.0 + @as(f32, @floatFromInt(m)) / 128.0; + } + + /// Fill both slots with the seed and mark all arrived: arithmetic check only, timing must not use this. + pub fn seed(p: *Loopback, rows: u32, width: u32) void { + const bytes = stageBytes(rows, width); + if (bytes > p.l.slot) @panic("loopback: the seeded part does not fit the window's slot"); + for (0..2) |x| { + const recv = p.win[p.l.recv(@intCast(x))..][0..bytes]; + const send = p.win[p.l.send(@intCast(x))..][0..bytes]; + for (0..rows) |r| for (0..width) |i| { + const at = r * width + i; + writeHalf(recv, at, seedValue(@intCast(r), @intCast(i), 0x15)); + writeHalf(send, at, seedValue(@intCast(r), @intCast(i), 0x0A)); + }; + } + p.satisfy(); + } + + /// Release every wait with the sentinel, as a failed link does: arithmetic check only, timing runs `startRelay`. + pub fn satisfy(p: *Loopback) void { + p.landed.set(seam_failed); + } + + /// The in-process relay: wait for posted, copy in, release -- Pair's two events in order, memcpy for the wire. + pub fn startRelay(p: *Loopback) !void { + std.debug.assert(p.relay == null); + p.stop.store(false, .release); + p.relay = try std.Thread.spawn(.{}, relayLoop, .{p}); + } + + pub fn stopRelay(p: *Loopback) void { + p.stop.store(true, .release); + if (p.relay) |t| t.join(); + p.relay = null; + } + + fn relayLoop(p: *Loopback) void { + var x: u64 = 1; + while (!p.stop.load(.acquire)) { + if (p.posted.value() < x) continue; // spin: the host's side of the seam is part of what is measured + const at: u32 = @truncate(x); + const n: usize = p.bytesFor(at); + @memcpy(p.win[p.l.recv(at)..][0..n], p.win[p.l.send(at)..][0..n]); + p.landed.set(x); + x += 1; + } + } + + /// The part's bytes for exchange `x` as the last `gather`/`argmax` queued it (the relay's copy length). + pub fn bytesFor(p: *const Loopback, x: u32) u32 { + return @atomicLoad(u32, &p.bytes[x % JOBS], .acquire); + } + + /// Land a host-written peer part for exchange x (the relay's post-copy half): copy it in, then release the wait. + pub fn arrive(p: *Loopback, x: u32, bytes: []const u8) void { + std.debug.assert(bytes.len <= p.l.slot); + @memcpy(p.win[p.l.recv(x)..][0..bytes.len], bytes); + p.landed.set(x); + } + + /// The seam `Pair.stage` encodes (before its `gather`): this rank's `width` columns of the frame row into the slot. + pub fn stage(p: *Loopback, e: *mtl.ComputeEncoder, x: u32, src: Ref, rows: u32, width: u32, stride: u32) void { + if (rows == 0 or width == 0 or width > stride) @panic("loopback: stage() wants rows > 0 and a width inside the source row"); + if (stageBytes(rows, width) > p.l.slot) @panic("loopback: this exchange's part does not fit the window's slot"); + e.setPipeline(p.stage_pipe); + e.setBuffer(p.wbuf, p.l.sync(), 0); + e.setValue(PairRows{ .width = width, .stride = stride, .rank = p.rank, .x = x }, 1); + e.setBuffer(src.buf, src.off, 2); + e.setBuffer(p.wbuf, p.l.send(x), 3); + e.dispatchGroups(.{ .width = rows }, .{ .width = 256 }); + } + + /// The seam `Pair.gather` encodes: the relay, then the frame row as this rank's half and the peer's, both bf16. + pub fn gather(p: *Loopback, e: *mtl.ComputeEncoder, x: u32, mine: Ref, out: Ref, rows: u32, width: u32, out_stride: u32) void { + if (rows == 0 or width == 0 or width % elems_per_word != 0) @panic("loopback: gather() wants rows > 0 and a width that fills whole window words"); + if (out_stride < 2 * width) @panic("loopback: gather()'s frame row is narrower than both halves"); + if (x != @as(u32, @truncate(p.x + 1))) @panic("loopback: an exchange was encoded out of order (part() not called first)"); + if (mine.buf.id != p.wbuf.id or mine.off != p.l.send(x)) @panic("loopback: gather()'s part is not the region part() returned"); + if (stageBytes(rows, width) > p.l.slot) @panic("loopback: this exchange's part does not fit the window's slot"); + p.x += 1; + @atomicStore(u32, &p.bytes[x % JOBS], @intCast(stageBytes(rows, width)), .release); + seam(e, p.fence, p.posted, p.landed, x); + e.setPipeline(p.gather_pipe); + e.setBuffer(p.wbuf, p.l.sync(), 0); + e.setValue(PairRows{ .width = width, .stride = out_stride, .rank = p.rank, .x = x }, 1); + e.setBuffer(mine.buf, mine.off, 2); + e.setBuffer(p.wbuf, p.l.recv(x), 3); + e.setBuffer(out.buf, out.off, 4); + e.dispatchGroups(.{ .width = rows }, .{ .width = 256 }); + } + + /// The loopback's head seam: this rank's argmax over n bf16 logits, the relay copy, then the pick over both parts. + pub fn argmaxBf16(p: *Loopback, e: *mtl.ComputeEncoder, x: u32, logits: Ref, stride: u32, lo: u32, n: u32, rows: u32, picks: Ref) void { + p.argmaxWidth(e, x, logits, stride, lo, n, rows, picks, 1); + } + + /// The same seam over fp32 logits (GLM-5.3's frame): the loopback covers both widths the kernel reads. + pub fn argmax(p: *Loopback, e: *mtl.ComputeEncoder, x: u32, logits: Ref, stride: u32, lo: u32, n: u32, rows: u32, picks: Ref) void { + p.argmaxWidth(e, x, logits, stride, lo, n, rows, picks, 0); + } + + fn argmaxWidth(p: *Loopback, e: *mtl.ComputeEncoder, x: u32, logits: Ref, stride: u32, lo: u32, n: u32, rows: u32, picks: Ref, dtype: u32) void { + if (rows == 0 or n == 0) @panic("loopback: argmax() wants rows > 0 and n > 0"); + if (x != @as(u32, @truncate(p.x + 1))) @panic("loopback: an exchange was encoded out of order (part() not called first)"); + p.x += 1; + @atomicStore(u32, &p.bytes[x % JOBS], @intCast(pickBytes(rows)), .release); // two u32 a row: (value, id) + e.setPipeline(p.argmax_pipe); + e.setBuffer(logits.buf, logits.off, 0); + e.setValue([4]u32{ stride, lo, n, dtype }, 1); + e.setBuffer(p.wbuf, p.l.send(x), 2); + e.dispatchGroups(.{ .width = rows }, .{ .width = 1024 }); + seam(e, p.fence, p.posted, p.landed, x); // the encoder boundary orders the argmax before the part is relayed + e.setPipeline(p.pick_pipe); + e.setBuffer(p.wbuf, p.l.sync(), 0); + e.setValue([2]u32{ rows, x }, 1); + e.setBuffer(p.wbuf, p.l.send(x), 2); + e.setBuffer(p.wbuf, p.l.recv(x), 3); + e.setBuffer(picks.buf, picks.off, 4); + e.dispatchGroups(.{ .width = 1 }, .{ .width = 32 }); + } + + /// Seams the GPU refused (must be 0 for a run to be valid). + pub fn gaveUp(p: *const Loopback) u32 { + return @atomicLoad(u32, p.word32(p.l.sync() + FAILED * 4), .acquire); + } + + /// There is no link to fail here, so this is `gaveUp` under the seam's name, and it reads 0. + pub fn failed(p: *const Loopback) bool { + return p.gaveUp() != 0; + } +}; diff --git a/zig/src/core/simd_qmm.zig b/zig/src/core/simd_qmm.zig index 844c2099e5..e8e24feb5f 100644 --- a/zig/src/core/simd_qmm.zig +++ b/zig/src/core/simd_qmm.zig @@ -8,6 +8,11 @@ pub fn splits(n: usize) u32 { return if (n <= 64) 32 else if (n <= 6144) 16 else 8; } +/// Chunks for `n`; a split rank passes the whole width as `whole` so it keeps that width's tree (0 keeps `n`). +pub fn chunks(n: usize, whole: usize) u32 { + return splits(if (whole != 0) whole else n); +} + /// Output tiles a simdgroup within the 16 KB reduction buffer; tiles never change a row's bits. pub fn tiles(n: usize, rt: u32, s: u32) u32 { var nt: u32 = if (n % 32 == 0) 4 else if (n % 16 == 0) 2 else 1; @@ -72,10 +77,10 @@ pub const Pipelines = struct { pub const Call = struct { pipeline: mtl.Pipeline, args: Args, groups: mtl.Size, threads: usize }; /// y = x W^T, 4-bit groups of 64: the scalar twin at 1-2 rows if `scalar`, chunk pairs if `paired`, else MMA tiles. -pub fn call(p: *const Pipelines, n: usize, k: usize, rows: usize, scalar: bool, paired: bool) !Call { +pub fn call(p: *const Pipelines, n: usize, k: usize, rows: usize, scalar: bool, paired: bool, whole: usize) !Call { if (n == 0 or n % 8 != 0 or n > std.math.maxInt(i32) or k == 0 or k % 64 != 0 or k > std.math.maxInt(i32)) return error.UnsupportedSimdShape; if (rows == 0 or rows > std.math.maxInt(i32)) return error.NoRows; - const s = splits(n); + const s = chunks(n, whole); const si = std.mem.indexOfScalar(u32, &s_values, s).?; const args = Args{ .k = @intCast(k), .n = @intCast(n), .r = @intCast(rows) }; if (scalar and scalarBlock(rows, s) != 0) { @@ -127,7 +132,7 @@ pub fn pairedExact(p: *const Pipelines, device: mtl.Device, m: Matrix) !bool { defer queue.deinit(); const cb = queue.commandBuffer(); const e = cb.compute(.serial); - for (outs, [_]bool{ false, true }) |o, pair| encode(e, try call(p, m.n, m.k, rows, false, pair), m, x, 0, o, 0); + for (outs, [_]bool{ false, true }) |o, pair| encode(e, try call(p, m.n, m.k, rows, false, pair, 0), m, x, 0, o, 0); e.end(); cb.commit(); cb.wait(); @@ -157,9 +162,9 @@ pub fn twinExact(p: *const Pipelines, device: mtl.Device, m: Matrix) !bool { defer queue.deinit(); const cb = queue.commandBuffer(); const e = cb.compute(.serial); - encode(e, try call(p, m.n, m.k, rows, false, false), m, x, 0, tiled, 0); - for (0..rows) |r| encode(e, try call(p, m.n, m.k, 1, true, false), m, x, r * m.k * 2, twin, r * m.n * 2); - for ([_]usize{ 0, rows - 2 }, 0..) |r, i| encode(e, try call(p, m.n, m.k, 2, true, false), m, x, r * m.k * 2, twos, i * 2 * m.n * 2); + encode(e, try call(p, m.n, m.k, rows, false, false, 0), m, x, 0, tiled, 0); + for (0..rows) |r| encode(e, try call(p, m.n, m.k, 1, true, false, 0), m, x, r * m.k * 2, twin, r * m.n * 2); + for ([_]usize{ 0, rows - 2 }, 0..) |r, i| encode(e, try call(p, m.n, m.k, 2, true, false, 0), m, x, r * m.k * 2, twos, i * 2 * m.n * 2); e.end(); cb.commit(); cb.wait(); @@ -177,6 +182,14 @@ test "chunks follow output width only and tiles fit the reduction buffer" { try std.testing.expectEqual(@as(u32, 32), splits(48)); try std.testing.expectEqual(@as(u32, 16), splits(5120)); try std.testing.expectEqual(@as(u32, 8), splits(34816)); + try std.testing.expectEqual(splits(6240), chunks(3168, 6240)); + try std.testing.expectEqual(splits(6240), chunks(6240, 0)); + // a pair's q and q | k | v key their checkpoint width: 12,288 and 10,240 over 6,144 and 5,120 a rank's. + try std.testing.expectEqual(@as(u32, 8), chunks(6144, 12288)); + try std.testing.expectEqual(@as(u32, 8), chunks(5120, 10240)); + // the rank's window alone would take 16 chunks, the split this key keeps a pair from forking. + try std.testing.expectEqual(@as(u32, 16), chunks(6144, 0)); + try std.testing.expectEqual(@as(u32, 16), chunks(5120, 0)); for (s_values) |s| for ([_]u32{ 1, 2 }) |rt| { const nt = tiles(34816, rt, s); try std.testing.expect(s * rt * nt * 64 * 4 <= 16384); diff --git a/zig/src/core/simd_qmm_bits.zig b/zig/src/core/simd_qmm_bits.zig index 73643f7d24..908013cb92 100644 --- a/zig/src/core/simd_qmm_bits.zig +++ b/zig/src/core/simd_qmm_bits.zig @@ -61,11 +61,12 @@ pub const Pipelines = struct { pub const Call = struct { pipeline: mtl.Pipeline, args: Args, groups: mtl.Size, threads: usize }; /// y[m * ld + j] = x W^T, g64 codes of `bits`: the twin at 1-2 rows when `scalar` (checked equal), else tiles. -pub fn call(p: *Pipelines, bits: u8, n: usize, k: usize, rows: usize, ld: usize, scalar: bool) !Call { +pub fn call(p: *Pipelines, bits: u8, n: usize, k: usize, rows: usize, ld: usize, scalar: bool, whole: usize) !Call { + // `whole` is the split rank's full width, whose own tree its K chunks keep (0: `n`). if (n == 0 or n % 8 != 0 or n > std.math.maxInt(i32) or k == 0 or k % 64 != 0 or k > std.math.maxInt(i32)) return error.UnsupportedSimdShape; if (ld < n or ld > std.math.maxInt(i32)) return error.UnsupportedSimdStride; if (rows == 0 or rows > std.math.maxInt(i32)) return error.NoRows; - const s = simd.splits(n); + const s = simd.chunks(n, whole); const args = Args{ .k = @intCast(k), .n = @intCast(n), .r = @intCast(rows), .ld = @intCast(ld) }; if (scalar and simd.scalarBlock(rows, s) != 0) { const nr: usize = if (n > 2048) 2 else 1; @@ -99,8 +100,8 @@ pub fn encode(e: mtl.ComputeEncoder, c: Call, m: Matrix, x: mtl.Buffer, x_off: u e.dispatchGroups(c.groups, .{ .width = c.threads }); } -/// Whether the scalar twin's 1-2 row results equal the tiles' for an [n, k] matrix of `bits` codes. -pub fn twinExact(p: *Pipelines, device: mtl.Device, bits: u8, m: Matrix, n: usize, k: usize) !bool { +/// Whether the scalar twin's 1-2 row results equal the tiles' for an [n, k] `bits` matrix, chunked as `call` takes it. +pub fn twinExact(p: *Pipelines, device: mtl.Device, bits: u8, m: Matrix, n: usize, k: usize, whole: usize) !bool { const rows = 8; const pool = mtl.objc.Pool.push(); defer pool.pop(); @@ -118,9 +119,9 @@ pub fn twinExact(p: *Pipelines, device: mtl.Device, bits: u8, m: Matrix, n: usiz defer queue.deinit(); const cb = queue.commandBuffer(); const e = cb.compute(.serial); - encode(e, try call(p, bits, n, k, rows, n, false), m, x, 0, tiled, 0); - for (0..rows) |r| encode(e, try call(p, bits, n, k, 1, n, true), m, x, r * k * 2, twin, r * n * 2); - for ([_]usize{ 0, rows - 2 }, 0..) |r, i| encode(e, try call(p, bits, n, k, 2, n, true), m, x, r * k * 2, pairs, i * 2 * n * 2); + encode(e, try call(p, bits, n, k, rows, n, false, whole), m, x, 0, tiled, 0); + for (0..rows) |r| encode(e, try call(p, bits, n, k, 1, n, true, whole), m, x, r * k * 2, twin, r * n * 2); + for ([_]usize{ 0, rows - 2 }, 0..) |r, i| encode(e, try call(p, bits, n, k, 2, n, true, whole), m, x, r * k * 2, pairs, i * 2 * n * 2); e.end(); cb.commit(); cb.wait(); diff --git a/zig/src/families/gemma4/weights.zig b/zig/src/families/gemma4/weights.zig index e417cd02fc..483a109f67 100644 --- a/zig/src/families/gemma4/weights.zig +++ b/zig/src/families/gemma4/weights.zig @@ -304,7 +304,7 @@ pub fn load(gpa: std.mem.Allocator, device: mtl.Device, c: *const ckpt.Checkpoin pub fn projectRow(ps: *Projections, device: mtl.Device, e: mtl.ComputeEncoder, raw: rowp.Weights, l: Linear, x: lane.Ref, y: mtl.Buffer, rows: usize) !void { const p = try ps.simd(device); const mtx = rowp.simd_bits.Matrix{ .w = raw.w, .scales = raw.scales, .biases = raw.biases }; - const call = try rowp.simd_bits.call(p, raw.bits, raw.n, raw.k, rows, l.y_stride orelse l.n, false); + const call = try rowp.simd_bits.call(p, raw.bits, raw.n, raw.k, rows, l.y_stride orelse l.n, false, 0); rowp.simd_bits.encode(e, call, mtx, x.buf, x.off, y, l.col * 2); } diff --git a/zig/src/families/nemotron/forward.zig b/zig/src/families/nemotron/forward.zig index 41d6c1a285..4d43535fa8 100644 --- a/zig/src/families/nemotron/forward.zig +++ b/zig/src/families/nemotron/forward.zig @@ -135,7 +135,7 @@ pub const Forward = struct { var calls: [3]row.simd_bits.Call = undefined; for (parts, 0..) |m, i| { if (!tiled(m.rows)) return false; - calls[i] = row.simd_bits.call(p, m.rows.bits, m.rows.n, m.rows.k, rows, lin.n, false) catch return false; + calls[i] = row.simd_bits.call(p, m.rows.bits, m.rows.n, m.rows.k, rows, lin.n, false, 0) catch return false; } for (parts, calls[0..parts.len]) |m, c| { e.pipe(c.pipeline); diff --git a/zig/src/families/nemotron/model.zig b/zig/src/families/nemotron/model.zig index 69cf472cd3..cb3337e0fe 100644 --- a/zig/src/families/nemotron/model.zig +++ b/zig/src/families/nemotron/model.zig @@ -124,10 +124,10 @@ pub const Model = struct { var it = m.weights.linears.valueIterator(); while (it.next()) |lin| if (lin.tiles == null) { for (lin.stack[0..lin.stacked]) |s| if (fwd.tiled(s.rows)) { - _ = try row.simd_bits.call(p, s.rows.bits, s.rows.n, s.rows.k, prefill.lane_rows + 1, lin.n, false); + _ = try row.simd_bits.call(p, s.rows.bits, s.rows.n, s.rows.k, prefill.lane_rows + 1, lin.n, false, 0); }; if (lin.rows) |r| if (fwd.tiled(r)) { - _ = try row.simd_bits.call(p, r.bits, r.n, r.k, prefill.lane_rows + 1, r.n, false); + _ = try row.simd_bits.call(p, r.bits, r.n, r.k, prefill.lane_rows + 1, r.n, false, 0); }; }; var gathers: std.ArrayList(prefill_w.Gathers) = .empty; diff --git a/zig/src/families/qwen27/attention.zig b/zig/src/families/qwen27/attention.zig index 994ff77797..c57aedf21e 100644 --- a/zig/src/families/qwen27/attention.zig +++ b/zig/src/families/qwen27/attention.zig @@ -9,7 +9,7 @@ pub const Plan = planning.Plan; pub const Cache = struct { keys: Ref, values: Ref, capacity: u32, key_stride: u32 = 0, value_stride: u32 = 0 }; const CachePtrs = extern struct { keys: u64, values: u64 }; const PackArgs = extern struct { rows: u32, packed_rows: u32, heads: u32 = 24, kv_heads: u32 = 4, dim: u32 = 256, group: u32 = 6 }; -const PromptArgs = extern struct { start: i32, rows: i32, chunks: i32, rp: i32, key_stride: i32, value_stride: i32, scale2: f32 }; +const PromptArgs = extern struct { start: i32, rows: i32, chunks: i32, rp: i32, key_stride: i32, value_stride: i32, scale2: f32, kv_heads: i32 }; pub const Limits = struct { rows: u32 = 128, streams: u32 = 8, @@ -196,8 +196,8 @@ pub const Preprocess = struct { qg: Ref, key: Ref, q_gain: Ref, k_gain: Ref, nor pub const Attention = struct { profiler: ?*profile.Trace = null, prefix: [16]mtl.Pipeline, - io: [10]mtl.Pipeline, - prompt: ?[2]mtl.Pipeline, // a prompt chunk's kernel and merge: tensor ops on the M5, simdgroup matrices before it + io: [11]mtl.Pipeline, + prompt: ?[3]mtl.Pipeline, // a prompt chunk's kernel and merge: tensor ops on the M5, simdgroup matrices before it units: bool, pub fn init(device: mtl.Device, library: mtl.Library) !Attention { var p: [16]mtl.Pipeline = undefined; @@ -208,10 +208,10 @@ pub const Attention = struct { pipe.* = try mtl.Pipeline.init(device, library, try std.fmt.bufPrint(&name, "q27_attn_prefix_{d}", .{i}), false); pn += 1; } - var io: [10]mtl.Pipeline = undefined; + var io: [11]mtl.Pipeline = undefined; var n: usize = 0; errdefer for (io[0..n]) |pipe| pipe.deinit(); - for ([_][]const u8{ "q27_attn_pack_a", "q27_attn_pack_b", "q27_attn_append", "q27_attn_tail", "q27_attn_merge", "q27_attn_gather", "q27_attn_scatter", "q27_head_norm", "q27_rope", "q27_gate" }, &io) |name, *pipe| { + for ([_][]const u8{ "q27_attn_pack_a", "q27_attn_pack_b", "q27_attn_append", "q27_attn_tail", "q27_attn_merge", "q27_attn_gather", "q27_attn_scatter", "q27_head_norm", "q27_rope", "q27_gate", "q27_attn_merge_single" }, &io) |name, *pipe| { pipe.* = try mtl.Pipeline.init(device, library, name, false); n += 1; } @@ -222,8 +222,8 @@ pub const Attention = struct { for (a.prefix ++ a.io) |pipe| pipe.deinit(); if (a.prompt) |pipes| for (pipes) |pipe| pipe.deinit(); } - /// The prompt kernel and its merge; before M5 null without them or without `prompt_threads` a group. - fn promptPipelines(device: mtl.Device, library: mtl.Library, units: bool) !?[2]mtl.Pipeline { + /// The prompt kernel plus pair and single-Mac merges; before M5 null unless a group holds `prompt_threads`. + fn promptPipelines(device: mtl.Device, library: mtl.Library, units: bool) !?[3]mtl.Pipeline { const chunks = try mtl.Pipeline.init(device, library, if (units) "q27_attn_prompt_8" else "q27_attn_prompt_sg_8", false); if (!units and chunks.maxThreads() < prompt_threads) { std.log.info("qwen27: the prompt attention kernel runs {d} threads a group here, not {d}; prompt rows take the decode path", .{ chunks.maxThreads(), prompt_threads }); @@ -231,7 +231,9 @@ pub const Attention = struct { return null; } errdefer chunks.deinit(); - return .{ chunks, try mtl.Pipeline.init(device, library, "q27_attn_prompt_merge", false) }; + const merge = try mtl.Pipeline.init(device, library, "q27_attn_prompt_merge", false); + errdefer merge.deinit(); + return .{ chunks, merge, try mtl.Pipeline.init(device, library, "q27_attn_prompt_merge_single", false) }; } pub fn preprocess(a: *const Attention, input_encoder: mtl.ComputeEncoder, p: Preprocess) !void { try p.geometry.check(); @@ -282,8 +284,8 @@ pub const Attention = struct { e.barrier(); try timing.end(); } - /// `prompt`: a prompt chunk's rows (one stream, a chain) take the prompt kernel. - pub fn encode(a: *const Attention, input_encoder: mtl.ComputeEncoder, scratch: *const Scratch, binding: *const Binding, plan: Plan, q: Ref, k: Ref, v: Ref, out: Ref, caches: []const Cache, raw_gate: ?Ref, prompt: bool) !void { + /// `prompt` picks the prompt kernel; `split` reads the head count, constant when one Mac runs it. + pub fn encode(a: *const Attention, input_encoder: mtl.ComputeEncoder, scratch: *const Scratch, binding: *const Binding, plan: Plan, q: Ref, k: Ref, v: Ref, out: Ref, caches: []const Cache, raw_gate: ?Ref, prompt: bool, split: bool) !void { try plan.geometry.check(); const g = plan.geometry; const q_bytes = @as(usize, plan.rows) * g.heads * g.dim * 2; @@ -324,7 +326,7 @@ pub const Attention = struct { e.dispatchThreads(mtl.Size.of(g.dim, args.packed_rows, g.kv_heads), mtl.Size.of(32, 4, 1)); if (prompt and a.prompt != null and (plan.geometry.shared_prefix or !a.units) and plan.streams == 1) { e.barrier(); - try a.promptRows(e, scratch, binding, plan, out); + try a.promptRows(e, scratch, binding, plan, out, split); if (raw_gate) |gate| { e.setPipeline(a.io[9]); bind(e, .{ out, gate }); @@ -333,11 +335,11 @@ pub const Attention = struct { e.dispatchThreads(mtl.Size.of(g.dim, g.heads, plan.rows), mtl.Size.of(32, 1, 1)); e.barrier(); } - } else try a.treeRows(e, scratch, binding, plan, q, out, raw_gate, args); + } else try a.treeRows(e, scratch, binding, plan, q, out, raw_gate, args, split); try timing.end(); } /// A prompt chunk's rows over absolute 512-key chunks (attention_prompt.metal), merged in key order. - fn promptRows(a: *const Attention, e: mtl.ComputeEncoder, scratch: *const Scratch, binding: *const Binding, plan: Plan, out: Ref) !void { + fn promptRows(a: *const Attention, e: mtl.ComputeEncoder, scratch: *const Scratch, binding: *const Binding, plan: Plan, out: Ref, split: bool) !void { const g = plan.geometry; const start: i32 = plan.meta[planning.global_words + 4]; for (plan.positions, 0..) |position, i| if (position != start + @as(i32, @intCast(i))) return error.PromptRowsNotChain; // the kernel places row i at start + i @@ -346,7 +348,7 @@ pub const Attention = struct { // before M5 (no shared prefix) the partials take the tail's buffers, which hold every chunk's packed rows const parts = if (g.shared_prefix) [3]Field{ .poa, .pma, .pla } else [3]Field{ .pob, .pmb, .plb }; if (!g.shared_prefix and @as(usize, chunks) * plan.tiles * 16 > @as(usize, scratch.tail_chunks) * scratch.limits.rows * 16) return error.AttentionLimits; - const args = PromptArgs{ .start = start, .rows = @intCast(plan.rows), .chunks = @intCast(chunks), .rp = @intCast(plan.tiles * 16), .key_stride = @intCast(binding.strides[0][0]), .value_stride = @intCast(binding.strides[0][1]), .scale2 = 0.0625 * std.math.log2e }; + const args = PromptArgs{ .start = start, .rows = @intCast(plan.rows), .chunks = @intCast(chunks), .rp = @intCast(plan.tiles * 16), .key_stride = @intCast(binding.strides[0][0]), .value_stride = @intCast(binding.strides[0][1]), .scale2 = 0.0625 * std.math.log2e, .kv_heads = @intCast(g.kv_heads) }; const pairs = prompt_threads / 64; // simdgroup pairs, each halving the head dim of one 16-row tile e.setPipeline(a.prompt.?[0]); bind(e, .{ scratch.get(.qa), binding.get(.caches) }); @@ -356,7 +358,7 @@ pub const Attention = struct { e.setBuffer(scratch.get(parts[2]).buffer, 0, 5); e.dispatchGroups(mtl.Size.of(g.kv_heads, chunks, (plan.tiles + pairs - 1) / pairs), mtl.Size.of(prompt_threads, 1, 1)); e.barrier(); - e.setPipeline(a.prompt.?[1]); + e.setPipeline(a.prompt.?[if (split) 1 else 2]); bind(e, .{ scratch.get(parts[0]), scratch.get(parts[1]), scratch.get(parts[2]) }); e.setValue(args, 3); e.setBuffer(out.buffer, out.offset, 4); @@ -364,7 +366,7 @@ pub const Attention = struct { e.barrier(); } /// Trees and decoded rows: cached keys by chunk (prefix), the round's own by path (tail), merged and gated. - fn treeRows(a: *const Attention, e: mtl.ComputeEncoder, scratch: *const Scratch, binding: *const Binding, plan: Plan, q: Ref, out: Ref, raw_gate: ?Ref, args: PackArgs) !void { + fn treeRows(a: *const Attention, e: mtl.ComputeEncoder, scratch: *const Scratch, binding: *const Binding, plan: Plan, q: Ref, out: Ref, raw_gate: ?Ref, args: PackArgs, split: bool) !void { const g = plan.geometry; if (!plan.fused) { e.setPipeline(a.io[1]); @@ -400,9 +402,10 @@ pub const Attention = struct { e.dispatchGroups(mtl.Size.of(g.kv_heads, plan.tail_chunks, plan.rows), mtl.Size.of(128, 1, 1)); e.barrier(); } - e.setPipeline(a.io[4]); + e.setPipeline(a.io[if (split) 4 else 10]); bind(e, .{ scratch.get(.poa), scratch.get(.pma), scratch.get(.pla), scratch.get(.pob), scratch.get(.pmb), scratch.get(.plb), binding.get(.meta), scratch.get(.nodes), out, raw_gate orelse out }); e.setValue(@as(u32, @intFromBool(raw_gate != null)), 10); + if (split) e.setValue(g.heads, 11); e.dispatchGroups(mtl.Size.of(g.kv_heads, plan.rows * g.group(), 1), mtl.Size.of(32, 1, 1)); e.barrier(); } diff --git a/zig/src/families/qwen27/attention_plan.zig b/zig/src/families/qwen27/attention_plan.zig index 0308f18d6d..36afdf3696 100644 --- a/zig/src/families/qwen27/attention_plan.zig +++ b/zig/src/families/qwen27/attention_plan.zig @@ -46,7 +46,9 @@ pub const Geometry = struct { } pub fn check(g: Geometry) !void { - if (g.heads != 24 or g.kv_heads != 4 or g.dim != 256 or g.rotary != 64) return error.TargetAttentionShape; + // Shaders take the GQA group (6) as a constant, not head count: 12/2 on a rank equals 24/4 on one Mac. + if (g.dim != 256 or g.rotary != 64 or g.heads == 0 or g.kv_heads == 0 or g.heads % g.kv_heads != 0) return error.TargetAttentionShape; + if (g.heads / g.kv_heads != 6) return error.TargetAttentionShape; } pub fn promptRows(g: Geometry, requested: u32) u32 { return @min(requested, @as(u32, if (g.shared_prefix and g.fuse_chain) chain_rows else max_rows)); @@ -435,7 +437,13 @@ test "target attention reads every config dimension without widening kernel qual const g = try Geometry.fromConfig(c, false); try std.testing.expect(g.sameShape(.{})); try std.testing.expect(!g.shared_prefix); + c.heads = 12; // a rank of a two-way split: the same GQA6 over half the heads + c.kv_heads = 2; + const half = try Geometry.fromConfig(c, true); + try std.testing.expectEqual(@as(u32, 12), half.heads); + try std.testing.expectEqual(@as(u32, 6), half.group()); c.heads = 16; // the supplied 9B has GQA4; the unchanged shaders still require GQA6 + c.kv_heads = 4; try std.testing.expectError(error.TargetAttentionShape, Geometry.fromConfig(c, true)); c.heads = 24; c.kv_heads = 2; diff --git a/zig/src/families/qwen27/config.zig b/zig/src/families/qwen27/config.zig index e4122d631e..8305d3c56a 100644 --- a/zig/src/families/qwen27/config.zig +++ b/zig/src/families/qwen27/config.zig @@ -27,6 +27,9 @@ pub const Config = struct { eos_count: usize = 0, max_position: usize = 262144, mrope_interleaved: bool = true, + rank: u32 = 0, + ranks: u32 = 1, + split_attention: bool = false, pub fn kind(c: Config, layer: usize) Kind { return c.kinds[layer]; @@ -51,8 +54,153 @@ pub const Config = struct { pub fn isEos(c: Config, token: u32) bool { return std.mem.indexOfScalar(u32, c.eos[0..c.eos_count], token) != null; } + + // Rank/ranks: 0/1 on one Mac. The fused zba row holds this rank's z, then every rank's b, then every rank's a. + + /// This rank's first value head: where its b, a, A_log and dt_bias rows start; 0 on one Mac or a replicated shard. + pub fn valueHead0(c: Config) usize { + return if (c.split_attention) c.rank * c.v_heads else 0; + } + + /// The `b` block's width: every rank's, because the fused projection keeps it replicated. + fn zbaHeads(c: Config) usize { + return if (c.split_attention) c.ranks * c.v_heads else c.v_heads; + } + + /// The rank's `b` row inside the fused `zba` row. + pub fn zbaB(c: Config) usize { + return c.gdnValueDim() + c.valueHead0(); + } + + /// The rank's `a` row inside the fused `zba` row, after the `b` block. + pub fn zbaA(c: Config) usize { + return c.gdnValueDim() + c.zbaHeads() + c.valueHead0(); + } + + /// The fused `zba` row's width. A multiple of 32 keeps `projection.prepare` on its tiled kernel. + pub fn zbaDim(c: Config) usize { + return c.gdnValueDim() + 2 * c.zbaHeads(); + } +}; + +/// A rank's slice of one tensor: elements `[from, from+len)` of the checkpoint's rows, or of an input row. +pub const Span = struct { + from: usize = 0, + len: usize = 0, + + pub fn end(s: Span) usize { + return s.from + s.len; + } +}; + +/// A rank's share of the 27B: the checkpoint holds `global`'s shapes, `local` is what this rank runs. +pub const Shard = struct { + rank: u32, + ranks: u32, + global: Config, + local: Config, + + /// This rank's first output row in a module whose `local` rows each rank takes a block of. + pub fn row0(s: Shard, local: usize) usize { + return @as(usize, s.rank) * local; + } + + /// This rank's rows of a global-row module: local rows on a split dimension, the whole tensor on a replicated one. + fn rows(s: Shard, global: usize, local: usize) Span { + return if (local == global) .{ .from = 0, .len = global } else .{ .from = s.row0(local), .len = local }; + } + + /// Whether attention and linear attention split across ranks; an mlp shard leaves them whole, so one seam a layer. + pub fn splitsAttention(s: Shard) bool { + return s.local.split_attention; + } + + /// self_attn.q_proj rows: q and its gate interleave every head_dim, so slices land on a 2*head_dim boundary. + pub fn qRows(s: Shard) Span { + return s.rows(s.global.qDim(), s.local.qDim()); + } + + /// `self_attn.k_proj` and `v_proj` rows (stacked k|v in that order): whole KV heads, one GQA group each. + pub fn kvRows(s: Shard) Span { + return s.rows(s.global.kvDim(), s.local.kvDim()); + } + + /// linear_attn.in_proj_qkv rows and conv1d channels: one q | k | v tensor, each block split on its own head count. + pub fn qkvSpans(s: Shard) [3]Span { + const gqk = s.global.k_heads * s.global.dk; + const q = s.rows(gqk, s.local.k_heads * s.local.dk); + const k = s.rows(gqk, s.local.k_heads * s.local.dk); + const v = s.rows(s.global.gdnValueDim(), s.local.gdnValueDim()); + return .{ q, .{ .from = gqk + k.from, .len = k.len }, .{ .from = 2 * gqk + v.from, .len = v.len } }; + } + + /// linear_attn.conv1d channels in qkvSpans order: the same three blocks, so local conv state indexes like qkv rows. + pub fn convSpans(s: Shard) [3]Span { + return s.qkvSpans(); + } + + /// `linear_attn.in_proj_z` rows: this rank's `v_heads * dv` value channels, one contiguous block. + pub fn zRows(s: Shard) Span { + return s.rows(s.global.gdnValueDim(), s.local.gdnValueDim()); + } + + /// linear_attn.in_proj_a and in_proj_b: replicated whole into zba; a 24-row split would drop the tiled projection. + pub fn abRows(s: Shard) Span { + return .{ .from = 0, .len = s.global.v_heads }; + } + + /// `linear_attn.A_log` and `dt_bias`: one constant a value head, so a rank takes its own heads' slice. + pub fn valueRows(s: Shard) Span { + return s.rows(s.global.v_heads, s.local.v_heads); + } + + /// `mlp.gate_proj` and `up_proj` rows (stacked gate|up): this rank's share of the intermediate width. + pub fn intermediateRows(s: Shard) Span { + return s.rows(s.global.intermediate, s.local.intermediate); + } + + /// `lm_head` rows: this rank's share of the vocabulary, one contiguous block of tokens. + pub fn vocabRows(s: Shard) Span { + return s.rows(s.global.vocab, s.local.vocab); + } + + /// `o_proj`/`out_proj`/`down_proj` rows: this rank's block over the whole K, or the width when replicated. + pub fn outputRows(s: Shard) Span { + return if (s.splitsAttention()) s.rows(s.global.hidden, s.global.hidden / s.ranks) else .{ .from = 0, .len = s.global.hidden }; + } }; +/// Mode: full halves attention, linear attention and MLP; mlp splits only the MLP and head, so one seam a layer. +pub const Mode = enum { full, mlp }; + +/// Split the even dimensions: heads, GQA groups, GDN heads, MLP width and vocabulary; embeddings and norms stay whole. +pub fn split(c: *Config, rank: u32, ranks: u32, mode: Mode) !Shard { + if (ranks < 2 or ranks > 8 or rank >= ranks) return error.BadShard; + if (c.intermediate % ranks != 0 or c.vocab % ranks != 0) return error.BadShard; + if (mode == .full and (c.heads % ranks != 0 or c.kv_heads % ranks != 0 or c.k_heads % ranks != 0 or c.v_heads % ranks != 0 or c.hidden % ranks != 0)) return error.BadShard; + const global = c.*; + if (mode == .full) { + c.heads /= ranks; + c.kv_heads /= ranks; + c.k_heads /= ranks; + c.v_heads /= ranks; + } + c.intermediate /= ranks; + c.vocab /= ranks; + c.rank = rank; + c.ranks = ranks; + c.split_attention = mode == .full; + // the local shapes must keep the group structure the attention, GDN and projection kernels run + if (c.heads == 0 or c.kv_heads == 0 or c.k_heads == 0 or c.v_heads == 0) return error.BadShard; + if (c.heads % c.kv_heads != 0 or c.v_heads % c.k_heads != 0) return error.BadShard; + if (c.heads / c.kv_heads != global.heads / global.kv_heads) return error.BadShard; + if (c.qDim() % 4 != 0 or c.kvDim() % 4 != 0 or c.gdnQkvDim() % 4 != 0 or c.gdnValueDim() % 4 != 0) return error.BadShard; + if (c.vocab % 4 != 0 or c.intermediate % 4 != 0) return error.BadShard; + if (c.layers != global.layers) return error.BadShard; + try shapeProducts(c.*); + return .{ .rank = rank, .ranks = ranks, .global = global, .local = c.* }; +} + pub fn object(v: std.json.Value) !std.json.ObjectMap { return if (v == .object) v.object else error.BadConfig; } @@ -184,121 +332,12 @@ pub fn parse(gpa: std.mem.Allocator, json: []const u8, generation: ?[]const u8) return c; } -const fixture = - \\{"model_type":"qwen3_5","hidden_size":64,"intermediate_size":128,"num_hidden_layers":4, - \\"num_attention_heads":4,"num_key_value_heads":2,"head_dim":16,"vocab_size":97, - \\"linear_num_key_heads":2,"linear_num_value_heads":4,"linear_key_head_dim":16, - \\"linear_value_head_dim":16,"linear_conv_kernel_dim":4,"eos_token_id":[2,11]} -; - -fn edit(a: std.mem.Allocator, key: []const u8, value: std.json.Value) ![]u8 { - var p = try std.json.parseFromSlice(std.json.Value, a, fixture, .{}); - defer p.deinit(); - try p.value.object.put(p.arena.allocator(), key, value); - return std.json.Stringify.valueAlloc(a, p.value, .{}); -} - -test "flat config preserves dimensions and interval schedule" { - const c = try parse(std.testing.allocator, fixture, null); - try std.testing.expectEqual(@as(usize, 128), c.qDim()); - try std.testing.expectEqual(@as(usize, 32), c.kvDim()); - try std.testing.expectEqual(@as(usize, 128), c.gdnQkvDim()); - try std.testing.expectEqual(Kind.linear, c.kind(0)); - try std.testing.expectEqual(Kind.attention, c.kind(3)); -} - -test "nested config merges generation EOS in stable order" { - const a = std.testing.allocator; - const text = try std.fmt.allocPrint(a, "{{\"model_type\":\"qwen3_5\",\"eos_token_id\":[11,2],\"text_config\":{s}}}", .{fixture}); - defer a.free(text); - const c = try parse(a, text, "{\"eos_token_id\":[2,3,11]}"); - try std.testing.expectEqualSlices(u32, &.{ 11, 2, 3 }, c.eos[0..c.eos_count]); - try std.testing.expect(c.isEos(3) and !c.isEos(4)); -} - -test "explicit layer types preserve their order" { - const a = std.testing.allocator; - var p = try std.json.parseFromSlice(std.json.Value, a, "[\"full_attention\",\"linear_attention\",\"full_attention\",\"linear_attention\"]", .{}); - defer p.deinit(); - const json = try edit(a, "layer_types", p.value); - defer a.free(json); - const c = try parse(a, json, null); - try std.testing.expectEqualSlices(Kind, &.{ .attention, .linear, .attention, .linear }, c.kinds[0..4]); -} - -test "rope parameters preserve precision and section metadata" { - const a = std.testing.allocator; - const p = try std.json.parseFromSlice(std.json.Value, a, "{\"rope_theta\":1234567.25,\"partial_rotary_factor\":0.5,\"mrope_section\":[2,1,1]}", .{}); - defer p.deinit(); - const json = try edit(a, "rope_parameters", p.value); - defer a.free(json); - const c = try parse(a, json, null); - try std.testing.expectEqual(@as(f64, 1234567.25), c.rope_theta); - try std.testing.expectEqual(@as(usize, 8), c.rope_dims); - try std.testing.expectEqualSlices(usize, &.{ 2, 1, 1 }, &c.mrope_section); -} - -test "dense config refuses MoE and tied embedding heads" { - const a = std.testing.allocator; - const moe = try edit(a, "num_experts", .{ .integer = 8 }); - defer a.free(moe); - try std.testing.expectError(error.NotQwenDense, parse(a, moe, null)); - const tied = try edit(a, "tie_word_embeddings", .{ .bool = true }); - defer a.free(tied); - try std.testing.expectError(error.TiedEmbeddingHead, parse(a, tied, null)); -} - -test "invalid scalar shapes and types fail without a cast trap" { - const a = std.testing.allocator; - for ([_]std.json.Value{ .{ .integer = -1 }, .{ .integer = 0 }, .{ .bool = true }, .{ .integer = 4294967296 } }) |v| { - const json = try edit(a, "hidden_size", v); - defer a.free(json); - try std.testing.expectError(error.BadConfig, parse(a, json, null)); - } - const json = try edit(a, "num_hidden_layers", .{ .integer = 129 }); - defer a.free(json); - try std.testing.expectError(error.BadConfig, parse(a, json, null)); -} - -test "missing and out of vocabulary EOS fail" { - const a = std.testing.allocator; - const missing = try edit(a, "eos_token_id", .null); - defer a.free(missing); - try std.testing.expectError(error.MissingEos, parse(a, missing, null)); - const c = try parse(a, missing, "{\"eos_token_id\":7}"); - try std.testing.expect(c.isEos(7)); - const wrong = try edit(a, "eos_token_id", .{ .integer = 97 }); - defer a.free(wrong); - try std.testing.expectError(error.BadEos, parse(a, wrong, null)); -} - -test "malformed layer and rope metadata fail" { - const a = std.testing.allocator; - const layers = try edit(a, "layer_types", .{ .bool = true }); - defer a.free(layers); - try std.testing.expectError(error.BadConfig, parse(a, layers, null)); - const rope = try edit(a, "partial_rotary_factor", .{ .float = 2 }); - defer a.free(rope); - try std.testing.expectError(error.BadConfig, parse(a, rope, null)); - try std.testing.expectError(error.BadConfig, parse(a, "[]", null)); -} - -test "head grouping requires complete key and value groups" { - const a = std.testing.allocator; - const attention = try edit(a, "num_key_value_heads", .{ .integer = 3 }); - defer a.free(attention); - try std.testing.expectError(error.BadConfig, parse(a, attention, null)); - const gdn = try edit(a, "linear_num_value_heads", .{ .integer = 3 }); - defer a.free(gdn); - try std.testing.expectError(error.BadConfig, parse(a, gdn, null)); +/// The embedding table stays whole on every rank, so its gather bound is the table's rows, not a shard's width. +pub fn embedBound(c: *const Config, rows: u64) !u32 { + if (rows < c.vocab) return error.BadEmbedding; + return std.math.cast(u32, rows) orelse error.BadEmbedding; } -test "text execution refuses unsupported gates and activation contracts" { - const a = std.testing.allocator; - const gate = try edit(a, "attn_output_gate", .{ .bool = false }); - defer a.free(gate); - try std.testing.expectError(error.UnsupportedAttentionGate, parse(a, gate, null)); - const act = try edit(a, "output_gate_type", .{ .string = "sigmoid" }); - defer a.free(act); - try std.testing.expectError(error.UnsupportedActivation, parse(a, act, null)); +test { + _ = @import("config_test.zig"); } diff --git a/zig/src/families/qwen27/config_test.zig b/zig/src/families/qwen27/config_test.zig new file mode 100644 index 0000000000..d0882d41cd --- /dev/null +++ b/zig/src/families/qwen27/config_test.zig @@ -0,0 +1,280 @@ +const std = @import("std"); +const cfg = @import("config.zig"); +const parse = cfg.parse; +const Config = cfg.Config; +const Kind = cfg.Kind; +const Span = cfg.Span; +const Mode = cfg.Mode; +const split = cfg.split; +const object = cfg.object; + +const fixture = + \\{"model_type":"qwen3_5","hidden_size":64,"intermediate_size":128,"num_hidden_layers":4, + \\"num_attention_heads":4,"num_key_value_heads":2,"head_dim":16,"vocab_size":97, + \\"linear_num_key_heads":2,"linear_num_value_heads":4,"linear_key_head_dim":16, + \\"linear_value_head_dim":16,"linear_conv_kernel_dim":4,"eos_token_id":[2,11]} +; + +fn edit(a: std.mem.Allocator, key: []const u8, value: std.json.Value) ![]u8 { + var p = try std.json.parseFromSlice(std.json.Value, a, fixture, .{}); + defer p.deinit(); + try p.value.object.put(p.arena.allocator(), key, value); + return std.json.Stringify.valueAlloc(a, p.value, .{}); +} + +test "flat config preserves dimensions and interval schedule" { + const c = try parse(std.testing.allocator, fixture, null); + try std.testing.expectEqual(@as(usize, 128), c.qDim()); + try std.testing.expectEqual(@as(usize, 32), c.kvDim()); + try std.testing.expectEqual(@as(usize, 128), c.gdnQkvDim()); + try std.testing.expectEqual(Kind.linear, c.kind(0)); + try std.testing.expectEqual(Kind.attention, c.kind(3)); +} + +test "nested config merges generation EOS in stable order" { + const a = std.testing.allocator; + const text = try std.fmt.allocPrint(a, "{{\"model_type\":\"qwen3_5\",\"eos_token_id\":[11,2],\"text_config\":{s}}}", .{fixture}); + defer a.free(text); + const c = try parse(a, text, "{\"eos_token_id\":[2,3,11]}"); + try std.testing.expectEqualSlices(u32, &.{ 11, 2, 3 }, c.eos[0..c.eos_count]); + try std.testing.expect(c.isEos(3) and !c.isEos(4)); +} + +test "explicit layer types preserve their order" { + const a = std.testing.allocator; + var p = try std.json.parseFromSlice(std.json.Value, a, "[\"full_attention\",\"linear_attention\",\"full_attention\",\"linear_attention\"]", .{}); + defer p.deinit(); + const json = try edit(a, "layer_types", p.value); + defer a.free(json); + const c = try parse(a, json, null); + try std.testing.expectEqualSlices(Kind, &.{ .attention, .linear, .attention, .linear }, c.kinds[0..4]); +} + +test "rope parameters preserve precision and section metadata" { + const a = std.testing.allocator; + const p = try std.json.parseFromSlice(std.json.Value, a, "{\"rope_theta\":1234567.25,\"partial_rotary_factor\":0.5,\"mrope_section\":[2,1,1]}", .{}); + defer p.deinit(); + const json = try edit(a, "rope_parameters", p.value); + defer a.free(json); + const c = try parse(a, json, null); + try std.testing.expectEqual(@as(f64, 1234567.25), c.rope_theta); + try std.testing.expectEqual(@as(usize, 8), c.rope_dims); + try std.testing.expectEqualSlices(usize, &.{ 2, 1, 1 }, &c.mrope_section); +} + +test "dense config refuses MoE and tied embedding heads" { + const a = std.testing.allocator; + const moe = try edit(a, "num_experts", .{ .integer = 8 }); + defer a.free(moe); + try std.testing.expectError(error.NotQwenDense, parse(a, moe, null)); + const tied = try edit(a, "tie_word_embeddings", .{ .bool = true }); + defer a.free(tied); + try std.testing.expectError(error.TiedEmbeddingHead, parse(a, tied, null)); +} + +test "invalid scalar shapes and types fail without a cast trap" { + const a = std.testing.allocator; + for ([_]std.json.Value{ .{ .integer = -1 }, .{ .integer = 0 }, .{ .bool = true }, .{ .integer = 4294967296 } }) |v| { + const json = try edit(a, "hidden_size", v); + defer a.free(json); + try std.testing.expectError(error.BadConfig, parse(a, json, null)); + } + const json = try edit(a, "num_hidden_layers", .{ .integer = 129 }); + defer a.free(json); + try std.testing.expectError(error.BadConfig, parse(a, json, null)); +} + +test "missing and out of vocabulary EOS fail" { + const a = std.testing.allocator; + const missing = try edit(a, "eos_token_id", .null); + defer a.free(missing); + try std.testing.expectError(error.MissingEos, parse(a, missing, null)); + const c = try parse(a, missing, "{\"eos_token_id\":7}"); + try std.testing.expect(c.isEos(7)); + const wrong = try edit(a, "eos_token_id", .{ .integer = 97 }); + defer a.free(wrong); + try std.testing.expectError(error.BadEos, parse(a, wrong, null)); +} + +test "malformed layer and rope metadata fail" { + const a = std.testing.allocator; + const layers = try edit(a, "layer_types", .{ .bool = true }); + defer a.free(layers); + try std.testing.expectError(error.BadConfig, parse(a, layers, null)); + const rope = try edit(a, "partial_rotary_factor", .{ .float = 2 }); + defer a.free(rope); + try std.testing.expectError(error.BadConfig, parse(a, rope, null)); + try std.testing.expectError(error.BadConfig, parse(a, "[]", null)); +} + +test "head grouping requires complete key and value groups" { + const a = std.testing.allocator; + const attention = try edit(a, "num_key_value_heads", .{ .integer = 3 }); + defer a.free(attention); + try std.testing.expectError(error.BadConfig, parse(a, attention, null)); + const gdn = try edit(a, "linear_num_value_heads", .{ .integer = 3 }); + defer a.free(gdn); + try std.testing.expectError(error.BadConfig, parse(a, gdn, null)); +} + +test "text execution refuses unsupported gates and activation contracts" { + const a = std.testing.allocator; + const gate = try edit(a, "attn_output_gate", .{ .bool = false }); + defer a.free(gate); + try std.testing.expectError(error.UnsupportedAttentionGate, parse(a, gate, null)); + const act = try edit(a, "output_gate_type", .{ .string = "sigmoid" }); + defer a.free(act); + try std.testing.expectError(error.UnsupportedActivation, parse(a, act, null)); +} + +/// The real 27B's shapes, so the shard arithmetic is pinned against the checkpoint rather than a toy config. +fn model27() Config { + return .{ + .hidden = 5120, + .intermediate = 17408, + .layers = 64, + .vocab = 248320, + .heads = 24, + .kv_heads = 4, + .head_dim = 256, + .k_heads = 16, + .v_heads = 48, + .dk = 128, + .dv = 128, + .conv_kernel = 4, + .eps = 1e-6, + .rope_dims = 64, + .rope_theta = 1e7, + }; +} + +fn expectSpan(span: Span, from: usize, len: usize) !void { + try std.testing.expectEqual(from, span.from); + try std.testing.expectEqual(len, span.len); +} + +test "a full shard takes whole heads and boundary-aligned row blocks" { + var d = model27(); + const s = try split(&d, 1, 2, .full); + try std.testing.expectEqual(@as(usize, 64), s.local.layers); + try std.testing.expectEqual(@as(usize, 12), s.local.heads); + try std.testing.expectEqual(@as(usize, 2), s.local.kv_heads); + try std.testing.expectEqual(@as(usize, 8), s.local.k_heads); + try std.testing.expectEqual(@as(usize, 24), s.local.v_heads); + try std.testing.expectEqual(@as(usize, 8704), s.local.intermediate); + try std.testing.expectEqual(@as(usize, 124160), s.local.vocab); + try std.testing.expect(s.splitsAttention()); + // q and its gate interleave every head: rank 1 starts at head 12, a 2*head_dim (512) boundary + try expectSpan(s.qRows(), 6144, 6144); + try std.testing.expectEqual(@as(usize, 0), s.qRows().from % (2 * s.global.head_dim)); + try expectSpan(s.kvRows(), 512, 512); + // GDN q | k | v is one tensor of three blocks: q and k split 8 of 16 heads, v splits 24 of 48 + const blocks = s.qkvSpans(); + try expectSpan(blocks[0], 1024, 1024); + try expectSpan(blocks[1], 3072, 1024); + try expectSpan(blocks[2], 7168, 3072); + try std.testing.expectEqual(@as(usize, 5120), blocks[0].len + blocks[1].len + blocks[2].len); + const conv = s.convSpans(); + for (blocks, conv) |want, got| { + try expectSpan(got, want.from, want.len); + } + try expectSpan(s.zRows(), 3072, 3072); + try expectSpan(s.valueRows(), 24, 24); + try expectSpan(s.abRows(), 0, 48); + try expectSpan(s.intermediateRows(), 8704, 8704); + try expectSpan(s.vocabRows(), 124160, 124160); + // o_proj, out_proj and down_proj own rows of the hidden width, each over the whole input K + try expectSpan(s.outputRows(), 2560, 2560); +} + +test "the embedding bound is the whole table on a pair rank, not the shard's width" { + var d = model27(); + const s = try split(&d, 1, 2, .full); + // the table keeps 248320 rows on both ranks while the head's shard is 124160 wide + try std.testing.expectEqual(@as(usize, 124160), s.local.vocab); + try std.testing.expectEqual(@as(u32, 248320), try cfg.embedBound(&s.local, 248320)); + try std.testing.expectEqual(@as(u32, 248320), try cfg.embedBound(&s.global, 248320)); + // a bound below the shard's own width is a table that lost rows, so it stays an error + try std.testing.expectError(error.BadEmbedding, cfg.embedBound(&s.local, 124159)); +} + +test "rank 0's blocks start at the tensor's front" { + var d = model27(); + const s = try split(&d, 0, 2, .full); + try expectSpan(s.qRows(), 0, 6144); + try expectSpan(s.kvRows(), 0, 512); + const blocks = s.qkvSpans(); + try expectSpan(blocks[0], 0, 1024); + try expectSpan(blocks[1], 2048, 1024); + try expectSpan(blocks[2], 4096, 3072); + try expectSpan(s.zRows(), 0, 3072); + try expectSpan(s.valueRows(), 0, 24); + try expectSpan(s.outputRows(), 0, 2560); +} + +test "the GDN qkv blocks tile the global tensor exactly across ranks" { + var seen: [10240]bool = @splat(false); + for (0..2) |r| { + var d = model27(); + const s = try split(&d, @intCast(r), 2, .full); + for (s.qkvSpans()) |span| { + try std.testing.expect(span.end() <= seen.len); + for (span.from..span.end()) |i| { + try std.testing.expect(!seen[i]); + seen[i] = true; + } + } + } + for (seen) |hit| try std.testing.expect(hit); +} + +test "an mlp shard replicates attention and splits only the mlp and the head" { + var d = model27(); + const s = try split(&d, 1, 2, .mlp); + try std.testing.expect(!s.splitsAttention()); + try std.testing.expectEqual(@as(usize, 24), s.local.heads); + try std.testing.expectEqual(@as(usize, 4), s.local.kv_heads); + try std.testing.expectEqual(@as(usize, 16), s.local.k_heads); + try std.testing.expectEqual(@as(usize, 48), s.local.v_heads); + try std.testing.expectEqual(@as(usize, 8704), s.local.intermediate); + try std.testing.expectEqual(@as(usize, 124160), s.local.vocab); + try expectSpan(s.qRows(), 0, 12288); + try expectSpan(s.kvRows(), 0, 1024); + const blocks = s.qkvSpans(); + try expectSpan(blocks[0], 0, 2048); + try expectSpan(blocks[1], 2048, 2048); + try expectSpan(blocks[2], 4096, 6144); + try expectSpan(s.zRows(), 0, 6144); + try expectSpan(s.abRows(), 0, 48); + // attention is replicated, so every value head is this rank's: A_log and dt_bias stay whole too + try expectSpan(s.valueRows(), 0, 48); + try expectSpan(s.intermediateRows(), 8704, 8704); + // an mlp shard keeps the mixers and down_proj replicated, so their output rows are the whole width + try expectSpan(s.outputRows(), 0, 5120); +} + +test "the fused zba row keeps the tiled projection and each rank's a and b" { + for ([_]Mode{ .full, .mlp }) |mode| { + var d = model27(); + const s = try split(&d, 1, 2, mode); + const c = s.local; + try std.testing.expectEqual(@as(usize, 0), c.zbaDim() % 32); + const whole_attention = mode == .mlp; + const ranks: usize = c.ranks; + const b_width = c.v_heads * (if (whole_attention) 1 else ranks); + const head0 = if (whole_attention) 0 else c.v_heads; + try std.testing.expectEqual(c.gdnValueDim() + 2 * b_width, c.zbaDim()); + try std.testing.expectEqual(c.gdnValueDim() + head0, c.zbaB()); + try std.testing.expectEqual(c.gdnValueDim() + b_width + head0, c.zbaA()); + try std.testing.expectEqual(c.zbaDim(), c.zbaA() + c.v_heads); + var e = model27(); + const z = try split(&e, 0, 2, mode); + try std.testing.expectEqual(z.local.gdnValueDim(), z.local.zbaB()); + try std.testing.expectEqual(z.local.gdnValueDim() + 48, z.local.zbaA()); + } + // one Mac is today's layout: z | b | a, each value head once + const u = model27(); + try std.testing.expectEqual(@as(usize, 6144), u.zbaB()); + try std.testing.expectEqual(@as(usize, 6192), u.zbaA()); + try std.testing.expectEqual(@as(usize, 6240), u.zbaDim()); +} diff --git a/zig/src/families/qwen27/decode_round.zig b/zig/src/families/qwen27/decode_round.zig index d2b3a9eb2b..160b9b3a49 100644 --- a/zig/src/families/qwen27/decode_round.zig +++ b/zig/src/families/qwen27/decode_round.zig @@ -4,6 +4,7 @@ const mtl = @import("metal"); const Model = @import("model.zig").Model; const Frame = @import("gpu_frame.zig").Frame; const forward = @import("forward.zig"); +const Ref = @import("projection.zig").Ref; /// Layers a verify encodes before committing them, so the GPU starts while the CPU encodes the rest. const commit_every = 4; const wts = @import("gpu_weights.zig"); @@ -59,6 +60,10 @@ pub const Runner = struct { images: ?image_rows.Images = null, /// A kept state's blit still reading a slot, by slot: the slot's next write waits for it (settleSlot). reading: [attn.planning.max_streams]?mtl.CommandBuffer = @splat(null), + /// A tensor-parallel rank's all-reduce, threaded straight into the layer program; null on one Mac. + reduce: ?forward.Reduce = null, + /// The pair's head-seam picks (one u32 a row), required with `reduce`: unexchanged halves make the ranks diverge. + picks: ?Ref = null, pub fn init(a: std.mem.Allocator, m: *Model, slots: u32, capacity: u32) !Runner { return initReady(a, m, slots, capacity, slots); @@ -248,11 +253,16 @@ pub const Runner = struct { defer pool.pop(); const spanned = if (r.trace) |t| t.spanned else 0; // a profiled pass's spans commit buffers of their own const handed = if (r.model.ane) |share| share.gpu_seconds else 0; - var stages = forward.Stages{ .model = r.model, .rows = @intCast(round.ids.len), .hooks = .{ .ptr = r, .gdn = gdnForward, .attention = attentionForward, .tap = if (r.capture_taps) tap else null }, .head = head, .images = placed.images }; + var stages = forward.Stages{ .model = r.model, .rows = @intCast(round.ids.len), .hooks = .{ .ptr = r, .gdn = gdnForward, .attention = attentionForward, .tap = if (r.capture_taps) tap else null }, .head = head, .images = placed.images, .reduce = r.reduce, .picks = r.picks }; r.last_gpu_seconds = @import("core").segments.runCommitting(r.model.device, &.{r.model.queue}, r.model.config.layers, .concurrent, &stages, commit_every) catch |err| { r.failed = true; return err; }; + // A seam cannot fail while encoding: the transport reports it out of band, and a broken pair fails the round. + if (r.reduce) |red| if (red.broken()) { + r.failed = true; + return error.PairFailed; + }; if (r.model.ane) |share| share.drain() catch |err| { r.failed = true; return err; @@ -436,7 +446,7 @@ pub const Runner = struct { try r.model.glue.unstack(e, f.get(.mixed), f.get(.value), .{ .rows = rows, .width = kv, .stride = 2 * kv, .offset = kv }); e.barrier(); try r.attention.preprocess(e, .{ .qg = f.get(.q_gate), .key = f.get(.key), .q_gain = wts.ref(weights.q_norm), .k_gain = wts.ref(weights.k_norm), .norm_q = f.get(.norm_query), .norm_k = f.get(.norm_key), .query = f.get(.query), .rotated_key = f.get(.rotated_key), .positions = f.get(.positions), .rows = rows, .eps = r.model.config.eps, .theta = r.model.config.rope_theta, .h_end = r.current_placed.h_end, .w_end = r.current_placed.w_end, .plane = f.capacity, .geometry = r.geometry }); - try r.attention.encode(e, &r.scratch, &r.bindings[layer], plan, f.get(.query), f.get(.rotated_key), f.get(.value), f.get(.mixer_out), caches, f.get(.q_gate), r.model.kernels.prompt); + try r.attention.encode(e, &r.scratch, &r.bindings[layer], plan, f.get(.query), f.get(.rotated_key), f.get(.value), f.get(.mixer_out), caches, f.get(.q_gate), r.model.kernels.prompt, r.reduce != null); } fn tap(ptr: *anyopaque, e: mtl.ComputeEncoder, layer: usize, hidden: @import("projection.zig").Ref, rows: u32) !void { diff --git a/zig/src/families/qwen27/dflash/runtime_callbacks.zig b/zig/src/families/qwen27/dflash/runtime_callbacks.zig index 46d252448c..bfea3d8552 100644 --- a/zig/src/families/qwen27/dflash/runtime_callbacks.zig +++ b/zig/src/families/qwen27/dflash/runtime_callbacks.zig @@ -29,7 +29,7 @@ fn project(ptr: *anyopaque, epoch: u64, x: op.View, w: ck.Tensor, mode: op.Mode) const input = try b.frame.ref(x); const output = try b.frame.ref(y); const s = @import("core").row_projection.simd; - s.encode(e, try s.call(&b.simd.?, t.n, t.k, x.rows, false, false), t.matrix(), input.buf, input.off, output.buf, output.off); + s.encode(e, try s.call(&b.simd.?, t.n, t.k, x.rows, false, false, 0), t.matrix(), input.buf, input.off, output.buf, output.off); } else { const adapter = b.prepared.get(@intFromPtr(w.bytes.ptr)) orelse return error.DraftPreparationUnqualified; const input = try b.frame.ref(x); diff --git a/zig/src/families/qwen27/forward.zig b/zig/src/families/qwen27/forward.zig index 7cf8e564bc..c2e6d8964a 100644 --- a/zig/src/families/qwen27/forward.zig +++ b/zig/src/families/qwen27/forward.zig @@ -4,7 +4,8 @@ const mtl = @import("metal"); const Model = @import("model.zig").Model; const wts = @import("gpu_weights.zig"); const Config = @import("config.zig").Config; -const Ref = @import("projection.zig").Ref; +/// A tensor in a frame buffer: the `Reduce` callbacks pass refs, so the type the seam API exposes here. +pub const Ref = @import("projection.zig").Ref; const Frame = @import("gpu_frame.zig").Frame; const ane_mlp = @import("ane_mlp.zig"); @@ -15,6 +16,47 @@ pub const Mixer = struct { tap: ?*const fn (*anyopaque, mtl.ComputeEncoder, usize, Ref, u32) anyerror!void = null, }; +/// One exchange: `x` its index, `ref` this rank's `width` columns of the seam in the window's slot for x. +pub const Seam = struct { x: u32, ref: Ref }; + +/// A rank's per-layer all-gather behind one interface: transport is a model flag, not a fork in the layer program. +pub const Reduce = struct { + ptr: *anyopaque, + /// Pure, and taken BEFORE the producer kernel writes `ref`: this seam's index and where this rank's part goes. + part: *const fn (*anyopaque) Seam, + /// This rank's own columns of the seam's input rows into x's slot: a bf16 copy, so it carries the frame's bits. + stage: *const fn (*anyopaque, *mtl.ComputeEncoder, u32, Ref, u32, u32, u32) void, + /// Relays x, waits for the peer's columns, and reopens the encoder: the caller uses the reopened copy. + gather: *const fn (*anyopaque, *mtl.ComputeEncoder, u32, Ref, Ref, u32, u32, u32) void, + /// Out of band: a seam cannot fail while encoding, so `failed()` is checked after each seam and command buffer. + failed: *const fn (*anyopaque) bool, + /// Head seam: this rank's argmax over `lo..lo+n` (stride `stride`) meets the peer's in `picks`, one u32 a row. + argmax: ?*const fn (*anyopaque, *mtl.ComputeEncoder, u32, Ref, u32, u32, u32, u32, Ref) void = null, + pub fn next(r: Reduce) Seam { + return r.part(r.ptr); + } + /// The next exchange's index, for a seam whose part the kernel writes itself (`argmax` writes the picks). + pub fn nextX(r: Reduce) u32 { + return r.part(r.ptr).x; + } + /// A seam's input rows: stage this rank's columns of `src`, then take both halves of the frame into `out`. + pub fn rowsIn(r: Reduce, e: *mtl.ComputeEncoder, se: Seam, src: Ref, rows: u32, width: u32, stride: u32, out: Ref, out_stride: u32) void { + r.stage(r.ptr, e, se.x, src, rows, width, stride); + r.gather(r.ptr, e, se.x, se.ref, out, rows, width, out_stride); + } + /// A seam's output rows: the producer wrote this rank's rows to `se.ref`; both ranks' rows land in `out`. + pub fn rowsOut(r: Reduce, e: *mtl.ComputeEncoder, se: Seam, out: Ref, rows: u32, width: u32, out_stride: u32) void { + r.gather(r.ptr, e, se.x, se.ref, out, rows, width, out_stride); + } + pub fn pick(r: Reduce, e: *mtl.ComputeEncoder, x: u32, logits: Ref, stride: u32, lo: u32, n: u32, rows: u32, picks: Ref) void { + r.argmax.?(r.ptr, e, x, logits, stride, lo, n, rows, picks); + } + /// Whether the peer link failed out of band: after a seam, before the next seam's kernel is encoded. + pub fn broken(r: Reduce) bool { + return r.failed(r.ptr); + } +}; + /// Serial decode and prompt segments own separate frames and order shared caches at the mixer. pub const Head = enum { all, last, none }; pub const Stages = struct { @@ -27,6 +69,10 @@ pub const Stages = struct { head: Head = .all, /// An image prompt's rows: the frame's `takes` rows take theirs over their placeholders' embeddings. images: ?Ref = null, + /// A tensor-parallel rank's all-reduce: null on one Mac, and then every seam is compiled out of the program. + reduce: ?Reduce = null, + /// Head seam picks (one u32 a row), required with `Reduce.argmax`: the ranks must pick the same token. + picks: ?Ref = null, pub fn begin(s: *Stages, lane: *@import("core").segments.Lane) !void { const m = s.model; if (s.rows == 0 or s.rows > m.frame.capacity or lane.k != 0) return error.BadFrameRows; @@ -67,21 +113,43 @@ pub const Stages = struct { pub fn post(s: *Stages, lane: *@import("core").segments.Lane, index: usize) !void { const m = s.model; const f = &m.frame; - const e = lane.enc; + var e = lane.enc; const layer = m.weights.layers[index]; const h: u32 = @intCast(m.config.hidden); const out = switch (layer.mixer) { .linear => |g| g.out, .attention => |a| a.out, }; - try m.kernels.quant(e, out, f.get(.mixer_out), null, f.get(.dims), f.get(.output), s.rows); - e.barrier(); + // Seam A: a rank owns output rows, not a K half: gather the mixer's columns before o_proj, its rows after. + const split_attention = s.reduce != null and m.shard.splitsAttention(); + var out_in = f.get(.mixer_out); + if (split_attention) { + std.debug.assert(m.shard.ranks == 2); // two Macs: gather's other half is the peer's half + const cols: u32 = @intCast(@divExact(out.k, 2)); // this rank's mixer columns, rows contiguous in `.mixer_out` + const first = s.reduce.?.next(); + s.reduce.?.rowsIn(&e, first, out_in, s.rows, cols, cols, f.get(.mixed), out.k); + lane.enc = e; // rowsIn reopened the encoder; the lane carries it to the next layer + if (s.reduce.?.broken()) return error.PairFailed; + e.barrier(); + out_in = f.get(.mixed); + } + const product = if (split_attention) s.reduce.?.next() else null; + if (product) |sm| try m.kernels.quant(lane.enc, out, out_in, null, f.get(.dims), sm.ref, s.rows) else try m.kernels.quant(lane.enc, out, out_in, null, f.get(.dims), f.get(.output), s.rows); + lane.enc.barrier(); + if (product) |sm| { + s.reduce.?.rowsOut(&e, sm, f.get(.output), s.rows, out.n, h); + lane.enc = e; + if (s.reduce.?.broken()) return error.PairFailed; + e.barrier(); + } try m.kernels.norm(e, s.hidden, f.get(.output), wts.ref(layer.post_norm), f.get(.eps), f.get(.dims), s.residual, f.get(.input), f.get(.sums), h, s.rows); e.barrier(); std.mem.swap(Ref, &s.hidden, &s.residual); var up = f.get(.mixed); up.offset += @as(usize, m.config.intermediate) * 2; const width: u32 = @intCast(m.config.intermediate); + // No ANE share on a pair (a rank's MLP columns are the GPU's); the assert keeps the branches from disagreeing. + std.debug.assert(s.reduce == null or m.ane == null); const shared = if (m.ane) |sh| m.kernels.prompt and s.rows <= sh.capacity() else false; if (if (shared) m.ane else null) |share| { // ANE slices run the shared columns; the GPU adds its remaining columns. @@ -111,14 +179,36 @@ pub const Stages = struct { try m.kernels.mlp(e, f.get(.mixed), up, f.get(.dims), f.get(.activated), f.get(.act_sums), width, s.rows); } lane.enc.barrier(); - try m.kernels.quant(lane.enc, layer.down, f.get(.activated), sums(f, .act_sums, s.rows), f.get(.dims), f.get(.output), s.rows); + // Seam B: gate|up gives each rank its columns, so down_proj's full K needs the peer's; `.mlp` keeps it whole. + var down_in = f.get(.activated); + const acts = if (s.reduce) |r| r.next() else null; + if (acts) |sm| { + s.reduce.?.rowsIn(&e, sm, down_in, s.rows, width, width, f.get(.mixed), layer.down.k); + lane.enc = e; + if (s.reduce.?.broken()) return error.PairFailed; + e.barrier(); + down_in = f.get(.mixed); + } + // The MLP is column-split on every pair (`.mlp` mode splits it too), so its seam is unconditional there. + const split_down = layer.down.n != h; + const seam_b = if (s.reduce != null and split_down) s.reduce.?.next() else null; + // A gathered frame's group sums are the peer's to write, so the kernel derives them where there is no peer. + const down_sums: ?@import("quant_gpu.zig").Sums = if (s.reduce == null) sums(f, .act_sums, s.rows) else null; + if (seam_b) |sm| try m.kernels.quant(lane.enc, layer.down, down_in, down_sums, f.get(.dims), sm.ref, s.rows) else try m.kernels.quant(lane.enc, layer.down, down_in, down_sums, f.get(.dims), f.get(.output), s.rows); lane.enc.barrier(); + if (seam_b) |sm| { + s.reduce.?.rowsOut(&e, sm, f.get(.output), s.rows, layer.down.n, h); + lane.enc = e; + if (s.reduce.?.broken()) return error.PairFailed; + e.barrier(); + } s.pending = f.get(.output); } pub fn finish(s: *Stages, lane: *@import("core").segments.Lane) !void { const m = s.model; const f = &m.frame; - const e = lane.enc; + // The head seam reopens the encoder, so it goes back in the lane; rank 0's half-vocab argmax is not a token. + var e = lane.enc; try m.kernels.norm(e, s.hidden, s.pending, wts.ref(m.weights.norm), f.get(.eps), f.get(.dims), s.residual, f.get(.input), f.get(.sums), @intCast(m.config.hidden), s.rows); e.barrier(); std.mem.swap(Ref, &s.hidden, &s.residual); @@ -134,12 +224,35 @@ pub const Stages = struct { try m.kernels.quant(e, m.weights.head, input, head_sums, f.get(.dims), f.get(.logits), if (s.head == .last) 1 else s.rows); } e.barrier(); + if (s.reduce) |r| { + if (s.head == .none) return; + const picks = s.picks orelse return error.HeadSeamMissing; + const seam = r.argmax orelse return error.HeadSeamMissing; + const span = m.shard.vocabRows(); + const local_rows: u32 = if (s.head == .last) 1 else s.rows; + seam(r.ptr, &e, r.nextX(), f.get(.logits), @intCast(m.config.vocab), @intCast(span.from), @intCast(span.len), local_rows, picks); + lane.enc = e; // the head seam reopened the encoder; the caller commits this one + if (r.broken()) return error.PairFailed; + } } }; -pub fn encode(m: *Model, e: mtl.ComputeEncoder, rows: u32, mixer: Mixer) !void { +/// `e` is the caller's encoder, but seams reopen it: the final one is written back through the pointer, so commit it. +pub fn encode(m: *Model, e: *mtl.ComputeEncoder, rows: u32, mixer: Mixer) !void { var stages = Stages{ .model = m, .rows = rows, .hooks = mixer }; - var lane = @import("core").segments.Lane{ .k = 0, .cb = undefined, .enc = e, .fence = undefined }; + try run(m, e, &stages); +} + +/// The tensor-parallel entry point: the one-Mac program with each layer's all-reduces inserted behind `reduce`. +pub fn encodeTp(m: *Model, e: *mtl.ComputeEncoder, rows: u32, mixer: Mixer, reduce: Reduce) !void { + var stages = Stages{ .model = m, .rows = rows, .hooks = mixer, .reduce = reduce }; + try run(m, e, &stages); +} + +fn run(m: *Model, e: *mtl.ComputeEncoder, stages: *Stages) !void { + // `cb`/`fence` stay undefined: a seam is the only closer, and it orders its writes through the transport's fence. + var lane = @import("core").segments.Lane{ .k = 0, .cb = undefined, .enc = e.*, .fence = undefined }; + defer e.* = lane.enc; try stages.begin(&lane); for (0..m.config.layers) |index| { try stages.pre(&lane, index); diff --git a/zig/src/families/qwen27/gdn_contract.zig b/zig/src/families/qwen27/gdn_contract.zig index 3dce6d0892..62ad130390 100644 --- a/zig/src/families/qwen27/gdn_contract.zig +++ b/zig/src/families/qwen27/gdn_contract.zig @@ -37,7 +37,7 @@ pub const Params = extern struct { slots: u32, weight_flags: u32, eps: f32, - /// 0: z, a and b in their own rows; else one fused zba row of this many values (z, then b, then a). + /// 0: z, a, b in own rows; else one fused zba row (z, b, a) whose rank blocks sit at zbaB/zbaA and is zbaDim wide. zba_stride: u32 = 0, advance: u32 = 0, // 1: a shared round's one-row windows step their state in the scan; the replay skips them }; @@ -84,7 +84,8 @@ pub const Bytes = struct { pub fn init(s: Shape, p: Params, k: WeightKinds) !Bytes { var expected = try s.params(p.rows, p.slots, k); - if (p.zba_stride != 0 and p.zba_stride < s.value + 2 * s.c.v_heads) return error.BadGdnDispatch; + // The row must span this rank's b and a reads past the other rank's padding, so the bound is the config stride. + if (p.zba_stride != 0 and p.zba_stride < s.c.zbaDim()) return error.BadGdnDispatch; expected.zba_stride = p.zba_stride; if (!std.mem.eql(u8, std.mem.asBytes(&p), std.mem.asBytes(&expected))) return error.BadGdnDispatch; return .{ diff --git a/zig/src/families/qwen27/gdn_hook.zig b/zig/src/families/qwen27/gdn_hook.zig index f2d6ce2e0f..944fd7aa30 100644 --- a/zig/src/families/qwen27/gdn_hook.zig +++ b/zig/src/families/qwen27/gdn_hook.zig @@ -33,12 +33,11 @@ pub const Hook = struct { } }; -/// z, b and a read in place from the fused zba rows. +/// z, b and a read in place from fused rows: a rank's b starts rank*24 after z, a 48 later (zbaDim = width + 96). fn refs(weight: Weights.Gdn, frame: *Frame) !pool_gpu.Refs { - const width = frame.config.gdnValueDim(); - const heads = frame.config.v_heads; + const c = frame.config; const mixed = frame.get(.mixed); - return .{ .z = mixed, .b = .{ .buffer = mixed.buffer, .offset = mixed.offset + width * 2 }, .a = .{ .buffer = mixed.buffer, .offset = mixed.offset + (width + heads) * 2 }, .output = frame.get(.mixer_out), .scalars = try pool_gpu.Scalars.from(weight), .zba_stride = @intCast(width + 2 * heads) }; + return .{ .z = mixed, .b = .{ .buffer = mixed.buffer, .offset = mixed.offset + c.zbaB() * 2 }, .a = .{ .buffer = mixed.buffer, .offset = mixed.offset + c.zbaA() * 2 }, .output = frame.get(.mixer_out), .scalars = try pool_gpu.Scalars.from(weight), .zba_stride = @intCast(c.zbaDim()) }; } fn devicePlan(plan: Plan, frame: *Frame) pool_gpu.DevicePlan { diff --git a/zig/src/families/qwen27/glue_gpu.zig b/zig/src/families/qwen27/glue_gpu.zig index d746877919..5fd13d67b1 100644 --- a/zig/src/families/qwen27/glue_gpu.zig +++ b/zig/src/families/qwen27/glue_gpu.zig @@ -1,7 +1,8 @@ //! Small operators use borrowed buffers and caller-owned encoders; copies never reinterpret activations. const mtl = @import("metal"); const source = @import("qwen_runtime_sources"); -const Config = @import("config.zig").Config; +const config = @import("config.zig"); +const Config = config.Config; const abi = @import("glue.zig"); const projection = @import("projection.zig"); const profile = @import("core").gpu_profile; @@ -51,12 +52,14 @@ pub const Kernels = struct { e.dispatchThreads(mtl.Size.of(width, rows, 1), mtl.Size.of(256, 1, 1)); } + /// The table is whole on every rank; only the head is sliced, so the bound is its row count, not `c.vocab`. pub fn embedding(k: Kernels, input_encoder: mtl.ComputeEncoder, c: Config, p: projection.Triple, ids: Ref, out: Ref, rows: u32) !void { + const vocab = try config.embedBound(&c, p.weight.shape[0]); var span = profile.begin(k.profiler, input_encoder, .embedding); errdefer span.cancel(); const e = span.encoder; - const a = abi.Embed{ .width = @intCast(c.hidden), .rows = rows, .vocab = @intCast(c.vocab), .bits = p.bits, .group = p.group, .metadata = .bf16 }; + const a = abi.Embed{ .width = @intCast(c.hidden), .rows = rows, .vocab = @intCast(vocab), .bits = p.bits, .group = p.group, .metadata = .bf16 }; try abi.checkEmbed(a); e.setPipeline(k.embed); e.setBuffer(ids.buffer, ids.offset, 0); diff --git a/zig/src/families/qwen27/gpu_weights.zig b/zig/src/families/qwen27/gpu_weights.zig index bb0fba68e3..ac0895d5b3 100644 --- a/zig/src/families/qwen27/gpu_weights.zig +++ b/zig/src/families/qwen27/gpu_weights.zig @@ -1,8 +1,8 @@ -//! The Metal graph borrows native scalar tensors and owns only byte-preserving affine layouts. +//! Borrows native scalar tensors and owns only byte-preserving affine layouts; rows and columns follow config.Shard. const std = @import("std"); const mtl = @import("metal"); const cp = @import("core").checkpoint_metal; -const Config = @import("config.zig").Config; +const cfg = @import("config.zig"); const projection = @import("projection.zig"); const lane = @import("core").lane_projection; const widths = @import("widths.zig"); @@ -89,31 +89,91 @@ const Builder = struct { return .{ .weight = tensors[0], .scales = tensors[1], .biases = tensors[2], .bits = w.bits, .group = w.group }; } - /// `parts` (stacked in order, `n` rows of `k` each, at their widths) staged from the files and repacked. + /// One part's staged tensors, shape-checked and viewed at the rank's rows; the stage must outlive these views. + fn viewed(p: Part, k: usize, tensors: [3]Tensor) !projection.Triple { + const weight = try checked(tensors[0], .u32, &.{ p.n, k * p.w.bits / 32 }); + const scales = try checked(tensors[1], .bf16, &.{ p.n, k / p.w.group }); + const biases = try checked(tensors[2], .bf16, &.{ p.n, k / p.w.group }); + return .{ .weight = try rowView(weight, p.span), .scales = try rowView(scales, p.span), .biases = try rowView(biases, p.span), .bits = p.w.bits, .group = p.w.group, .gn = p.n }; + } + + /// Parts at their widths and K windows, all staged before any is viewed: the stage is one reused buffer. fn linear(b: Builder, parts: []const Part, k: usize) !Linear { var arena = std.heap.ArenaAllocator.init(b.allocator); defer arena.deinit(); - const names = try arena.allocator().alloc([]const u8, parts.len * 3); + if (parts.len > 3) return error.BadProjectionShape; + for (parts) |p| if (k * p.w.bits % 32 != 0 or k % p.w.group != 0) return error.BadNativeTensor; + const a = arena.allocator(); + const names = try a.alloc([]const u8, parts.len * 3); for (parts, 0..) |p, i| for ([_][]const u8{ ".weight", ".scales", ".biases" }, 0..) |suffix, j| { - names[i * 3 + j] = try std.mem.concat(arena.allocator(), u8, &.{ b.prefix, p.module, suffix }); + names[i * 3 + j] = try std.mem.concat(a, u8, &.{ b.prefix, p.module, suffix }); }; try b.stage.reserve(b.checkpoint, names); defer b.stage.reset(); var triples: [3]projection.Triple = undefined; - if (parts.len > triples.len) return error.BadProjectionShape; for (parts, triples[0..parts.len], 0..) |p, *t, i| { - if (k * p.w.bits % 32 != 0 or k % p.w.group != 0) return error.BadNativeTensor; - t.* = .{ - .weight = try checked(try b.stage.read(b.checkpoint, names[i * 3]), .u32, &.{ p.n, k * p.w.bits / 32 }), - .scales = try checked(try b.stage.read(b.checkpoint, names[i * 3 + 1]), .bf16, &.{ p.n, k / p.w.group }), - .biases = try checked(try b.stage.read(b.checkpoint, names[i * 3 + 2]), .bf16, &.{ p.n, k / p.w.group }), - .bits = p.w.bits, - .group = p.w.group, - }; + t.* = try viewed(p, k, .{ + try b.stage.read(b.checkpoint, names[i * 3]), + try b.stage.read(b.checkpoint, names[i * 3 + 1]), + try b.stage.read(b.checkpoint, names[i * 3 + 2]), + }); } return b.stack(triples[0..parts.len]); } + /// One tensor a rank takes in blocks (GDN q | k | v); `n` is the split's key, `rows` each block's checkpoint rows. + fn blocks(b: Builder, module: []const u8, n: usize, k: usize, w: Width, spans: []const cfg.Span, rows: []const usize) !Linear { + var arena = std.heap.ArenaAllocator.init(b.allocator); + defer arena.deinit(); + if (k * w.bits % 32 != 0 or k % w.group != 0) return error.BadNativeTensor; + const suffixes = [_][]const u8{ ".weight", ".scales", ".biases" }; + const names = try arena.allocator().alloc([]const u8, suffixes.len); + for (suffixes, 0..) |suffix, j| names[j] = try std.mem.concat(arena.allocator(), u8, &.{ b.prefix, module, suffix }); + try b.stage.reserve(b.checkpoint, names); + defer b.stage.reset(); + const weight = try checked(try b.stage.read(b.checkpoint, names[0]), .u32, &.{ n, k * w.bits / 32 }); + const scales = try checked(try b.stage.read(b.checkpoint, names[1]), .bf16, &.{ n, k / w.group }); + const biases = try checked(try b.stage.read(b.checkpoint, names[2]), .bf16, &.{ n, k / w.group }); + var triples: [3]projection.Triple = undefined; + if (spans.len > triples.len or rows.len != spans.len) return error.BadProjectionShape; + for (spans, rows, triples[0..spans.len]) |span, gn, *t| { + t.* = .{ .weight = try rowView(weight, span), .scales = try rowView(scales, span), .biases = try rowView(biases, span), .bits = w.bits, .group = w.group, .gn = gn, .whole = n }; + } + return b.stack(triples[0..spans.len]); + } + + /// `conv1d.weight`'s row blocks fused in the same q | k | v order, as one bf16 copy (the file's is not repacked). + fn conv(b: Builder, module: []const u8, kernel: usize, n: usize, spans: []const cfg.Span) !Tensor { + const t = try b.tensor(module, .bf16, &.{ n, kernel, 1 }); + var at: usize = 0; + for (spans) |span| { + if (span.len == 0 or span.from != at) break; + at = span.end(); + } + if (at == n) return t; // the blocks are the whole tensor in order: no copy + var rows: usize = 0; + for (spans) |span| rows += span.len; + const out = try b.device.buffer(rows * kernel * @sizeOf(u16), mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + b.owned.append(b.allocator, out) catch |err| { + out.deinit(); + return err; + }; + const from = t.host(u16); + const to = out.slice(u16, rows * kernel); + var base: usize = 0; + for (spans) |span| { + if (span.end() > n) return error.BadNativeTensor; + @memcpy(to[base * kernel ..][0 .. span.len * kernel], from[span.from * kernel ..][0 .. span.len * kernel]); + base += span.len; + } + var view = t; + view.buffer = out; + view.offset = 0; + view.bytes = rows * kernel * @sizeOf(u16); + view.shape[0] = rows; + return view; + } + /// Fused members of one width take one matrix as before; a mixed stack becomes one matrix a member. fn stack(b: Builder, parts: []const projection.Triple) !Linear { var ws: [3]Width = undefined; @@ -133,12 +193,14 @@ const Builder = struct { } var fused = runs[0].linear; fused.n = first; + fused.gn = 0; + for (runs) |run| fused.gn += run.linear.gn; fused.raw = null; fused.runs = runs; return fused; } - fn layer(b: Builder, c: Config, i: usize) !Layer { + fn layer(b: Builder, s: *const cfg.Shard, i: usize) !Layer { var arena = std.heap.ArenaAllocator.init(b.allocator); defer arena.deinit(); const a = arena.allocator(); @@ -148,41 +210,83 @@ const Builder = struct { return std.mem.concat(allocator, u8, &.{ p, leaf }); } }; + const c = s.local; + const g = s.global; const w = b.widths.layers[i]; const M = widths.Member; + const qkv = s.qkvSpans(); + const qkvRows = [3]usize{ g.k_heads * g.dk, g.k_heads * g.dk, g.gdnValueDim() }; + const convRows = s.convSpans(); const mixer: @FieldType(Layer, "mixer") = switch (c.kind(i)) { .linear => .{ .linear = .{ - .qkv = try b.linear(&.{.{ .module = try Name.full(a, prefix, "linear_attn.in_proj_qkv"), .n = c.gdnQkvDim(), .w = w[@backingInt(M.qkv)] }}, c.hidden), - .zba = try b.linear(&.{ .{ .module = try Name.full(a, prefix, "linear_attn.in_proj_z"), .n = c.gdnValueDim(), .w = w[@backingInt(M.z)] }, .{ .module = try Name.full(a, prefix, "linear_attn.in_proj_b"), .n = c.v_heads, .w = w[@backingInt(M.b)] }, .{ .module = try Name.full(a, prefix, "linear_attn.in_proj_a"), .n = c.v_heads, .w = w[@backingInt(M.a)] } }, c.hidden), - .out = try b.linear(&.{.{ .module = try Name.full(a, prefix, "linear_attn.out_proj"), .n = c.hidden, .w = w[@backingInt(M.out)] }}, c.gdnValueDim()), - .conv = try b.tensor(try Name.full(a, prefix, "linear_attn.conv1d.weight"), .bf16, &.{ c.gdnQkvDim(), c.conv_kernel, 1 }), - .a_log = try b.tensor(try Name.full(a, prefix, "linear_attn.A_log"), .bf16, &.{c.v_heads}), - .dt = try b.tensor(try Name.full(a, prefix, "linear_attn.dt_bias"), .bf16, &.{c.v_heads}), + .qkv = try b.blocks(try Name.full(a, prefix, "linear_attn.in_proj_qkv"), g.gdnQkvDim(), g.hidden, w[@backingInt(M.qkv)], &qkv, &qkvRows), + .zba = try b.linear(&.{ + .{ .module = try Name.full(a, prefix, "linear_attn.in_proj_z"), .n = g.gdnValueDim(), .w = w[@backingInt(M.z)], .span = s.zRows() }, + .{ .module = try Name.full(a, prefix, "linear_attn.in_proj_b"), .n = g.v_heads, .w = w[@backingInt(M.b)], .span = s.abRows() }, + .{ .module = try Name.full(a, prefix, "linear_attn.in_proj_a"), .n = g.v_heads, .w = w[@backingInt(M.a)], .span = s.abRows() }, + }, g.hidden), + .out = try b.linear(&.{.{ .module = try Name.full(a, prefix, "linear_attn.out_proj"), .n = g.hidden, .w = w[@backingInt(M.out)], .span = s.outputRows() }}, g.gdnValueDim()), + .conv = try b.conv(try Name.full(a, prefix, "linear_attn.conv1d.weight"), g.conv_kernel, g.gdnQkvDim(), &convRows), + .a_log = try rowView(try b.tensor(try Name.full(a, prefix, "linear_attn.A_log"), .bf16, &.{g.v_heads}), s.valueRows()), + .dt = try rowView(try b.tensor(try Name.full(a, prefix, "linear_attn.dt_bias"), .bf16, &.{g.v_heads}), s.valueRows()), .norm = try b.tensor(try Name.full(a, prefix, "linear_attn.norm.weight"), .bf16, &.{c.dv}), } }, .attention => .{ .attention = .{ - .q = try b.linear(&.{.{ .module = try Name.full(a, prefix, "self_attn.q_proj"), .n = c.qDim(), .w = w[@backingInt(M.q)] }}, c.hidden), - .kv = try b.linear(&.{ .{ .module = try Name.full(a, prefix, "self_attn.k_proj"), .n = c.kvDim(), .w = w[@backingInt(M.k)] }, .{ .module = try Name.full(a, prefix, "self_attn.v_proj"), .n = c.kvDim(), .w = w[@backingInt(M.v)] } }, c.hidden), - .out = try b.linear(&.{.{ .module = try Name.full(a, prefix, "self_attn.o_proj"), .n = c.hidden, .w = w[@backingInt(M.o)] }}, c.heads * c.head_dim), + .q = try b.linear(&.{.{ .module = try Name.full(a, prefix, "self_attn.q_proj"), .n = g.qDim(), .w = w[@backingInt(M.q)], .span = s.qRows() }}, g.hidden), + .kv = try b.linear(&.{ + .{ .module = try Name.full(a, prefix, "self_attn.k_proj"), .n = g.kvDim(), .w = w[@backingInt(M.k)], .span = s.kvRows() }, + .{ .module = try Name.full(a, prefix, "self_attn.v_proj"), .n = g.kvDim(), .w = w[@backingInt(M.v)], .span = s.kvRows() }, + }, g.hidden), + .out = try b.linear(&.{.{ .module = try Name.full(a, prefix, "self_attn.o_proj"), .n = g.hidden, .w = w[@backingInt(M.o)], .span = s.outputRows() }}, g.heads * g.head_dim), .q_norm = try b.tensor(try Name.full(a, prefix, "self_attn.q_norm.weight"), .bf16, &.{c.head_dim}), .k_norm = try b.tensor(try Name.full(a, prefix, "self_attn.k_norm.weight"), .bf16, &.{c.head_dim}), } }, }; - const gu = try b.linear(&.{ .{ .module = try Name.full(a, prefix, "mlp.gate_proj"), .n = c.intermediate, .w = w[@backingInt(M.gate)] }, .{ .module = try Name.full(a, prefix, "mlp.up_proj"), .n = c.intermediate, .w = w[@backingInt(M.up)] } }, c.hidden); - const down = try b.linear(&.{.{ .module = try Name.full(a, prefix, "mlp.down_proj"), .n = c.hidden, .w = w[@backingInt(M.down)] }}, c.intermediate); + const gu = try b.linear(&.{ + .{ .module = try Name.full(a, prefix, "mlp.gate_proj"), .n = g.intermediate, .w = w[@backingInt(M.gate)], .span = s.intermediateRows() }, + .{ .module = try Name.full(a, prefix, "mlp.up_proj"), .n = g.intermediate, .w = w[@backingInt(M.up)], .span = s.intermediateRows() }, + }, g.hidden); + const down = try b.linear(&.{.{ .module = try Name.full(a, prefix, "mlp.down_proj"), .n = g.hidden, .w = w[@backingInt(M.down)], .span = s.outputRows() }}, g.intermediate); return .{ .norm = try b.tensor(try Name.full(a, prefix, "input_layernorm.weight"), .bf16, &.{c.hidden}), .post_norm = try b.tensor(try Name.full(a, prefix, "post_attention_layernorm.weight"), .bf16, &.{c.hidden}), .gu = gu, .down = down, .mixer = mixer }; } }; -const Part = struct { module: []const u8, n: usize, w: Width }; +const Part = struct { + /// The module's checkpoint name, without the shard's namespace prefix. + module: []const u8, + /// The rows the checkpoint's tensor has, before this rank takes its span. + n: usize, + w: Width, + /// The rows this rank takes: the whole tensor unless the rank splits this matrix. + span: cfg.Span = .{}, +}; + +/// Rows `[span.from, end)` of `t`: rows are contiguous, so the same buffer at a later offset and a shorter shape. +fn rowView(t: Tensor, span: cfg.Span) !Tensor { + if (span.len == 0) return t; + if (span.from == 0 and span.len == t.shape[0]) return t; + if (span.end() > t.shape[0]) return error.BadNativeTensor; + var row: usize = 1; + for (t.shape[1..t.rank]) |d| row *= d; + const elem = t.dtype.size(); + const at = span.from * row * elem; + if (at % elem != 0) return error.BadNativeTensor; + var view = t; + view.offset = t.offset + at; + view.bytes = row * span.len * elem; + view.shape[0] = span.len; + return view; +} fn checked(t: Tensor, dtype: cp.DType, shape: []const usize) !Tensor { if (t.dtype != dtype or t.rank != shape.len or !std.mem.eql(usize, t.shape[0..t.rank], shape)) return error.BadNativeTensor; return t; } -/// The graph from `checkpoint`, whose `repacked` tensors stay in their files: each Linear is staged, repacked, dropped. -pub fn load(allocator: std.mem.Allocator, device: mtl.Device, checkpoint: *const cp.Checkpoint, c: Config, ws: *const widths.Widths) !Weights { +/// Graph from `checkpoint`; each Linear is staged, repacked, dropped. `s` is the rank's share, `local` from `global`. +pub fn load(allocator: std.mem.Allocator, device: mtl.Device, checkpoint: *const cp.Checkpoint, s: *const cfg.Shard, ws: *const widths.Widths) !Weights { + const c = s.local; + const g = s.global; const prefix: []const u8 = if (checkpoint.has("language_model.model.embed_tokens.weight")) "language_model." else if (checkpoint.has("model.embed_tokens.weight")) "" else return error.MissingTextNamespace; if (checkpoint.has("language_model.model.embed_tokens.weight") and checkpoint.has("model.embed_tokens.weight")) return error.AmbiguousTextNamespace; var owned: std.ArrayList(mtl.Buffer) = .empty; @@ -200,10 +304,10 @@ pub fn load(allocator: std.mem.Allocator, device: mtl.Device, checkpoint: *const var stage = cp.Stage{ .device = device }; defer stage.deinit(); const b = Builder{ .allocator = allocator, .device = device, .checkpoint = checkpoint, .prefix = prefix, .owned = &owned, .stage = &stage, .runs = &runs, .widths = ws }; - for (layers, 0..) |*l, i| l.* = try b.layer(c, i); - const embed = try b.quant("model.embed_tokens", c.vocab, c.hidden, ws.embed); + for (layers, 0..) |*l, i| l.* = try b.layer(s, i); + const embed = try b.quant("model.embed_tokens", g.vocab, g.hidden, ws.embed); const norm = try b.tensor("model.norm.weight", .bf16, &.{c.hidden}); - const head = try b.linear(&.{.{ .module = "lm_head", .n = c.vocab, .w = ws.head }}, c.hidden); + const head = try b.linear(&.{.{ .module = "lm_head", .n = g.vocab, .w = ws.head, .span = s.vocabRows() }}, g.hidden); if (device.tensorUnits()) try laneOrder(allocator, device, layers, head); return .{ .allocator = allocator, .layers = layers, .embed = embed, .norm = norm, .head = head, .owned = owned, .runs = runs }; } diff --git a/zig/src/families/qwen27/model.zig b/zig/src/families/qwen27/model.zig index 1e5aa91252..cb83ae0790 100644 --- a/zig/src/families/qwen27/model.zig +++ b/zig/src/families/qwen27/model.zig @@ -18,6 +18,9 @@ pub const Model = struct { allocator: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, + /// This rank's share of the family's weights; `local` shapes are what this graph builds, `global` the file's. + shard: cfg.Shard, + /// The local config, so downstream shapes are the rank's own. config: cfg.Config, checkpoint: core.checkpoint_metal.Checkpoint, weights: wts.Weights, @@ -41,6 +44,16 @@ pub const Model = struct { return loadWith(a, io, dir, rows, null, true); } + /// One rank of a `ranks`-way tensor-parallel run: the same checkpoint, sliced in the engine. + pub fn loadShard(a: std.mem.Allocator, io: std.Io, dir: []const u8, rows: u32, rank: u32, ranks: u32, mode: cfg.Mode, share_fits: ?*const fn (u64) bool) !*Model { + return loadShared(a, io, dir, rows, rank, ranks, mode, share_fits, false); + } + + /// One rank of a pair, its share decided after everything else is accounted for (the serving order). + pub fn loadShardDeferred(a: std.mem.Allocator, io: std.Io, dir: []const u8, rows: u32, rank: u32, ranks: u32, mode: cfg.Mode) !*Model { + return loadShared(a, io, dir, rows, rank, ranks, mode, null, true); + } + fn loadWith( a: std.mem.Allocator, io: std.Io, @@ -48,27 +61,44 @@ pub const Model = struct { rows: u32, share_fits: ?*const fn (u64) bool, later: bool, + ) !*Model { + return loadShared(a, io, dir, rows, 0, 1, .full, share_fits, later); + } + + fn loadShared( + a: std.mem.Allocator, + io: std.Io, + dir: []const u8, + rows: u32, + rank: u32, + ranks: u32, + mode: cfg.Mode, + share_fits: ?*const fn (u64) bool, + later: bool, ) !*Model { const path = try std.fs.path.join(a, &.{ dir, "config.json" }); defer a.free(path); const text = try std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(16 << 20)); defer a.free(text); const c = try cfg.parse(a, text, null); - try supported(c); + var local = c; + const shard = if (ranks == 1) cfg.Shard{ .rank = 0, .ranks = 1, .global = c, .local = local } else try cfg.split(&local, rank, ranks, mode); + try supported(shard.global); var formats = try affine.Formats.init(a, text); defer formats.deinit(); - var ws = try widths.Widths.resolve(a, &formats, c); + var ws = try widths.Widths.resolve(a, &formats, shard.local); defer ws.deinit(); const m = try a.create(Model); errdefer a.destroy(m); m.allocator = a; - m.config = c; + m.shard = shard; + m.config = shard.local; m.widths_tag = ws.tag(); m.device = try mtl.Device.init(); errdefer m.device.deinit(); if (!qualifiedGeneration(m.device.generation())) return error.QwenMetalQualificationRequired; var why: widths.Why = .{}; - widths.admit(&ws, c, .{ .tensor_units = m.device.tensorUnits() }, &why) catch |err| { + widths.admit(&ws, shard.local, .{ .tensor_units = m.device.tensorUnits() }, &why) catch |err| { std.log.err("qwen27: {s} (layer {?d}) is {d}-bit group {d}; this chip's projection kernels do not run it ({s})", .{ why.module, why.layer, why.width.bits, why.width.group, @errorName(why.reason) }); return err; }; @@ -80,14 +110,14 @@ pub const Model = struct { defer core.checkpoint_host.freeShardFiles(a, files); // projections stay in the files until load repacks them: the weights live once, in their lane layouts for (files) |file| try m.checkpoint.addFileSplit(m.device, file, "", "language_model.", wts.repacked); - m.weights = try wts.load(a, m.device, &m.checkpoint, c, &ws); + m.weights = try wts.load(a, m.device, &m.checkpoint, &m.shard, &ws); errdefer m.weights.deinit(); m.kernels = .{ .allocator = a, .device = m.device, .raw_tiles = !m.device.tensorUnits() }; errdefer m.kernels.deinit(); m.glue = try glue.Kernels.init(m.device); errdefer m.glue.deinit(); // Serving defers automatic widening until its context and streams are known; explicit widths stay exact. - m.frame = try Frame.init(m.device, c, @max(rows, promptRows(if (later) 0 else m.device.generation()))); + m.frame = try Frame.init(m.device, shard.local, @max(rows, promptRows(if (later) 0 else m.device.generation()))); errdefer m.frame.deinit(); m.ane = null; m.share_pending = true; @@ -106,7 +136,9 @@ pub const Model = struct { defer m.checkpoint.closeFiles(); const c = m.config; const columns = try ane_policy.columns(if (platform.c.getenv("TF_ANE_COLUMNS")) |v| std.mem.span(v) else null); - const share = ane_policy.requestedOnChip(columns, m.device.tensorUnits()) and mtl.ane.available() and m.weights.uniformMlp(); + // The share is a whole MLP's columns: a rank holds half of them, so tensor parallel keeps prompts on the GPU. + const share = m.shard.ranks == 1 and ane_policy.requestedOnChip(columns, m.device.tensorUnits()) and mtl.ane.available() and m.weights.uniformMlp(); + if (m.shard.ranks != 1) std.log.info("qwen27: tensor parallel rank {d}/{d} runs prompts on the GPU alone", .{ m.shard.rank, m.shard.ranks }); const bytes = ane_mlp.Share.bytes(c, columns); const slices = (m.frame.capacity + ane_mlp.width - 1) / ane_mlp.width; const required: ?u64 = if (share) ane_mlp.Share.totalBytes(c, columns, slices, ane_policy.build_runtime_overhead) catch |err| switch (err) { diff --git a/zig/src/families/qwen27/projection.zig b/zig/src/families/qwen27/projection.zig index a3603ee3f2..7d32bb6798 100644 --- a/zig/src/families/qwen27/projection.zig +++ b/zig/src/families/qwen27/projection.zig @@ -6,11 +6,12 @@ const row = @import("core").row_projection; const lane = @import("core").lane_projection; pub const Ref = struct { buffer: mtl.Buffer, offset: usize = 0 }; -/// A device matrix; `runs` non-empty means fused members of different widths, each its own product into these rows. -pub const Linear = struct { words: Ref, pairs: Ref, n: u32, k: u32, tile: u32, slices: u32, raw: ?row.Weights = null, bits: u8 = 4, group: u32 = 64, runs: []const Run = &.{} }; +/// A device matrix; `runs` non-empty means fused members of different widths. `gn`: the checkpoint's rows (0 = `n`). +pub const Linear = struct { words: Ref, pairs: Ref, n: u32, k: u32, tile: u32, slices: u32, raw: ?row.Weights = null, bits: u8 = 4, group: u32 = 64, runs: []const Run = &.{}, gn: usize = 0 }; /// One product inside a mixed fused matrix, writing output columns [first, first + linear.n). pub const Run = struct { linear: Linear, first: u32 }; -pub const Triple = struct { weight: checkpoint.Tensor, scales: checkpoint.Tensor, biases: checkpoint.Tensor, bits: u8 = 4, group: u32 = 64 }; +/// Staged tensors: `gn` the part's rows when `weight.shape[0]` is a slice, `whole` its tensor's rows as a block. +pub const Triple = struct { weight: checkpoint.Tensor, scales: checkpoint.Tensor, biases: checkpoint.Tensor, bits: u8 = 4, group: u32 = 64, gn: usize = 0, whole: usize = 0 }; pub fn split(n: usize, k: usize) !u32 { if (n == 0 or k == 0 or n % 4 != 0 or k % 64 != 0) return error.BadProjectionShape; @@ -33,24 +34,53 @@ pub fn promptSplit(n: usize, k: usize) ?u32 { return null; } +/// A part's rows in the checkpoint, which are the staged rows on one Mac. +fn globalRows(p: Triple) usize { + return if (p.gn != 0) p.gn else p.weight.shape[0]; +} + +/// The checkpoint tensor's rows, which the split keys on: a block of one tensor keys the whole tensor. +fn tensorRows(p: Triple) usize { + return if (p.whole != 0) p.whole else globalRows(p); +} + +/// The fused matrix's rows in the checkpoint: what the M5 and prompt tables key on. +fn splitRows(parts: []const Triple) usize { + var rows: usize = 0; + for (parts) |p| rows += globalRows(p); + return rows; +} + +/// The K split a staged matrix takes: the tensor's rows key the fp32 order, the fused checkpoint rows the M5 table. +const Split = struct { slices: u32, m5: ?u32 }; + +fn splitFor(parts: []const Triple, k: usize) !Split { + return .{ .slices = try split(tensorRows(parts[0]), k), .m5 = m5Split(splitRows(parts), k) }; +} + pub fn prepare(device: mtl.Device, allocator: std.mem.Allocator, owned: *std.ArrayList(mtl.Buffer), parts: []const Triple, tile: u32) !Linear { if (tile != 32 or parts.len == 0) return error.BadProjectionShape; const bits = parts[0].bits; const group = parts[0].group; if (bits == 0 or bits > 8 or group == 0) return error.BadAffineMetadata; + // the checkpoint's packed width, whole on every rank's member const k = parts[0].weight.shape[1] * 32 / bits; - var slices = try split(parts[0].weight.shape[0], k); + if (k == 0 or k % group != 0) return error.BadAffineMetadata; + // The split keys on the leading part's tensor; the fused checkpoint rows key the M5 table. + const gn = splitRows(parts); + const key = try splitFor(parts, k); + var slices = key.slices; var n: usize = 0; for (parts) |p| { const rows = p.weight.shape[0]; - if (p.bits != bits or p.group != group or k % group != 0) return error.BadAffineMetadata; - if (p.weight.rank != 2 or p.weight.dtype != .u32 or p.weight.shape[1] * 32 != k * bits or rows % 4 != 0 or try split(rows, k) != slices) return error.BadProjectionShape; + if (p.bits != bits or p.group != group) return error.BadAffineMetadata; + if (p.weight.rank != 2 or p.weight.dtype != .u32 or p.weight.shape[1] * 32 != k * bits or rows % 4 != 0) return error.BadProjectionShape; if (p.scales.dtype != .bf16 or p.biases.dtype != .bf16 or p.scales.rank != 2 or p.biases.rank != 2 or p.scales.shape[0] != rows or p.biases.shape[0] != rows or p.scales.shape[1] != k / group or p.biases.shape[1] != k / group) return error.BadAffineMetadata; n = try std.math.add(usize, n, rows); } - if (!device.tensorUnits() or n % tile != 0) return prepareRows(device, allocator, owned, parts, n, k, tile, slices); - if (bits != 4 or group != 64) return prepareLanes(device, allocator, owned, parts, n, k, tile, slices); - slices = m5Split(n, k) orelse slices; + if (!device.tensorUnits() or n % tile != 0) return prepareRows(device, allocator, owned, parts, n, k, tile, slices, gn); + if (bits != 4 or group != 64) return prepareLanes(device, allocator, owned, parts, n, k, tile, slices, gn); + slices = key.m5 orelse slices; const w = try device.buffer(n * k / 2, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); owned.append(allocator, w) catch |err| { w.deinit(); @@ -64,21 +94,22 @@ pub fn prepare(device: mtl.Device, allocator: std.mem.Allocator, owned: *std.Arr const dst = w.slice(u32, n * k / 8); const pair = sb.slice(u16, n * (k / 64) * 2); var base: usize = 0; + const groups = k / 64; for (parts) |p| { const rows = p.weight.shape[0]; const words = p.weight.host(u32); const scales = p.scales.host(u16); const biases = p.biases.host(u16); - for (0..rows) |r| for (0..k / 64) |g| { + for (0..rows) |r| for (0..groups) |g| { const src = r * (k / 8) + g * 8; - const to = (((base + r) / tile) * (k / 64) + g) * tile * 8 + ((base + r) % tile) * 8; + const to = (((base + r) / tile) * groups + g) * tile * 8 + ((base + r) % tile) * 8; @memcpy(dst[to..][0..8], words[src..][0..8]); - pair[(g * n + base + r) * 2] = scales[r * (k / 64) + g]; - pair[(g * n + base + r) * 2 + 1] = biases[r * (k / 64) + g]; + pair[(g * n + base + r) * 2] = scales[r * groups + g]; + pair[(g * n + base + r) * 2 + 1] = biases[r * groups + g]; }; base += rows; } - return .{ .words = .{ .buffer = w }, .pairs = .{ .buffer = sb }, .n = @intCast(n), .k = @intCast(k), .tile = tile, .slices = slices }; + return .{ .words = .{ .buffer = w }, .pairs = .{ .buffer = sb }, .n = @intCast(n), .k = @intCast(k), .tile = tile, .slices = slices, .gn = gn }; } test "stacked gate-up keeps the member K split rather than recomputing it from stacked width" { @@ -89,29 +120,103 @@ test "stacked gate-up keeps the member K split rather than recomputing it from s try std.testing.expectEqual(@as(u32, 8), try split(5120, 17408)); } -fn prepareRows(device: mtl.Device, allocator: std.mem.Allocator, owned: *std.ArrayList(mtl.Buffer), parts: []const Triple, n: usize, k: usize, tile: u32, slices: u32) !Linear { +/// A staged matrix of `rows` rows whose checkpoint matrix has `gn` (0: the rows are the whole matrix). +fn tri(rows: usize, gn: usize, blank: checkpoint.Tensor) Triple { + var t = blank; + t.shape[0] = rows; + return .{ .weight = t, .scales = t, .biases = t, .gn = gn }; +} + +test "a rank's row slice keeps the checkpoint matrix's K split" { + const blank: checkpoint.Tensor = .{ .buffer = undefined, .offset = 0, .bytes = 0, .dtype = .u32 }; + // gate | up: a rank stages 8,704 rows a member of the checkpoint's 17,408, so the leading member keys 2, not 4. + const gate = [2]Triple{ tri(8704, 17408, blank), tri(8704, 17408, blank) }; + try std.testing.expectEqual(@as(usize, 17408), globalRows(gate[0])); + try std.testing.expectEqual(@as(u32, 2), try split(globalRows(gate[0]), 5120)); + try std.testing.expectEqual(@as(u32, 4), try split(gate[0].weight.shape[0], 5120)); + // The fused matrix's checkpoint rows double, as the single Mac's stack: 34,816, the tables' key. + try std.testing.expectEqual(@as(usize, 34816), splitRows(&gate)); + // q: 12,288 checkpoint rows, 6,144 a rank's, 4 and 8 slices. + const q = tri(6144, 12288, blank); + try std.testing.expectEqual(@as(u32, 4), try split(globalRows(q), 5120)); + try std.testing.expectEqual(@as(u32, 8), try split(q.weight.shape[0], 5120)); + // down_proj: the M5 and prompt tables hold the checkpoint's 5,120 rows over 17,408, which a rank's 2,560 miss. + const down = tri(2560, 5120, blank); + try std.testing.expectEqual(@as(?u32, 4), m5Split(splitRows(&.{down}), 17408)); + try std.testing.expectEqual(@as(?u32, null), m5Split(down.weight.shape[0], 17408)); + try std.testing.expectEqual(@as(?u32, 2), promptSplit(splitRows(&.{down}), 17408)); + try std.testing.expectEqual(@as(?u32, null), promptSplit(down.weight.shape[0], 17408)); + // q | k | v: three blocks of one 10,240-row tensor, so the blocks' checkpoint rows sum to its width. + var qkv = [3]Triple{ tri(1024, 2048, blank), tri(1024, 2048, blank), tri(3072, 6144, blank) }; + for (&qkv) |*part| part.whole = 10240; + try std.testing.expectEqual(@as(usize, 10240), splitRows(&qkv)); + try std.testing.expectEqual(@as(?u32, 2), m5Split(splitRows(&qkv), 5120)); + try std.testing.expectEqual(@as(?u32, 1), promptSplit(splitRows(&qkv), 5120)); + // The tensor's 10,240 rows key 4, as main's single tensor did; keyed on the block's 2,048 it would take 8. + try std.testing.expectEqual(@as(u32, 4), try split(tensorRows(qkv[0]), 5120)); + try std.testing.expectEqual(@as(u32, 8), try split(globalRows(qkv[0]), 5120)); + // A rank's three local blocks are 5,120 rows and key neither table, so the fallback split would stand. + var local: usize = 0; + for (qkv) |part| local += part.weight.shape[0]; + try std.testing.expectEqual(@as(usize, 5120), local); + try std.testing.expectEqual(@as(?u32, null), m5Split(local, 5120)); + try std.testing.expectEqual(@as(?u32, null), promptSplit(local, 5120)); +} + +test "every one-Mac stack keeps main's K split" { + const blank: checkpoint.Tensor = .{ .buffer = undefined, .offset = 0, .bytes = 0, .dtype = .u32 }; + // One Mac stages every block whole: `gn` the checkpoint's rows, `whole` the tensor's when the parts block one. + const Table = struct { rows: [3]usize, len: usize, whole: usize, k: usize, slices: u32, m5: ?u32 }; + const table = [_]Table{ + .{ .rows = .{ 2048, 2048, 6144 }, .len = 3, .whole = 10240, .k = 5120, .slices = 4, .m5 = 2 }, + .{ .rows = .{ 12288, 0, 0 }, .len = 1, .whole = 0, .k = 5120, .slices = 4, .m5 = 2 }, + .{ .rows = .{ 1024, 1024, 0 }, .len = 2, .whole = 0, .k = 5120, .slices = 8, .m5 = 8 }, + .{ .rows = .{ 5120, 0, 0 }, .len = 1, .whole = 0, .k = 6144, .slices = 8, .m5 = 4 }, + .{ .rows = .{ 17408, 17408, 0 }, .len = 2, .whole = 0, .k = 5120, .slices = 2, .m5 = null }, + .{ .rows = .{ 5120, 0, 0 }, .len = 1, .whole = 0, .k = 17408, .slices = 8, .m5 = 4 }, + .{ .rows = .{ 6144, 48, 48 }, .len = 3, .whole = 0, .k = 5120, .slices = 8, .m5 = 4 }, + .{ .rows = .{ 248320, 0, 0 }, .len = 1, .whole = 0, .k = 5120, .slices = 1, .m5 = null }, + }; + for (table) |t| { + var parts: [3]Triple = undefined; + for (t.rows[0..t.len], parts[0..t.len]) |rows, *p| { + p.* = .{ .weight = blank, .scales = blank, .biases = blank, .gn = rows, .whole = t.whole }; + } + const key = try splitFor(parts[0..t.len], t.k); + try std.testing.expectEqual(t.slices, key.slices); + try std.testing.expectEqual(t.m5, key.m5); + } +} + +fn prepareRows(device: mtl.Device, allocator: std.mem.Allocator, owned: *std.ArrayList(mtl.Buffer), parts: []const Triple, n: usize, k: usize, tile: u32, slices: u32, gn: usize) !Linear { const options = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; const bits: usize = parts[0].bits; const ng = k / parts[0].group; const row_bytes = k * bits / 8; + const srow_words = k * bits / 32; const w = try own(device, allocator, owned, n * row_bytes, options); const scale = try own(device, allocator, owned, n * ng * 2, options); const bias = try own(device, allocator, owned, n * ng * 2, options); var base: usize = 0; for (parts) |p| { const rows = p.weight.shape[0]; - @memcpy(w.contents()[base * row_bytes .. (base + rows) * row_bytes], std.mem.sliceAsBytes(p.weight.host(u32))); - @memcpy(scale.contents()[base * ng * 2 .. (base + rows) * ng * 2], std.mem.sliceAsBytes(p.scales.host(u16))); - @memcpy(bias.contents()[base * ng * 2 .. (base + rows) * ng * 2], std.mem.sliceAsBytes(p.biases.host(u16))); + const words = p.weight.host(u32); + const scales = p.scales.host(u16); + const biases = p.biases.host(u16); + for (0..rows) |r| { + @memcpy(w.contents()[(base + r) * row_bytes ..][0..row_bytes], std.mem.sliceAsBytes(words[r * srow_words ..][0..srow_words])); + @memcpy(scale.contents()[(base + r) * ng * 2 ..][0 .. ng * 2], std.mem.sliceAsBytes(scales[r * ng ..][0..ng])); + @memcpy(bias.contents()[(base + r) * ng * 2 ..][0 .. ng * 2], std.mem.sliceAsBytes(biases[r * ng ..][0..ng])); + } base += rows; } const raw = row.Weights{ .w = w, .scales = scale, .biases = bias, .n = n, .k = k, .group = parts[0].group, .bits = parts[0].bits, .sum = .f32 }; try raw.validate(); - return .{ .words = .{ .buffer = w }, .pairs = .{ .buffer = scale }, .n = @intCast(n), .k = @intCast(k), .tile = tile, .slices = slices, .raw = raw, .bits = parts[0].bits, .group = parts[0].group }; + return .{ .words = .{ .buffer = w }, .pairs = .{ .buffer = scale }, .n = @intCast(n), .k = @intCast(k), .tile = tile, .slices = slices, .raw = raw, .bits = parts[0].bits, .group = parts[0].group, .gn = gn }; } /// Widths other than 4/64 on tensor units: lane_projection.weight_layout byte copies for the non-reg lane kernel. -fn prepareLanes(device: mtl.Device, allocator: std.mem.Allocator, owned: *std.ArrayList(mtl.Buffer), parts: []const Triple, n: usize, k: usize, tile: u32, slices: u32) !Linear { +fn prepareLanes(device: mtl.Device, allocator: std.mem.Allocator, owned: *std.ArrayList(mtl.Buffer), parts: []const Triple, n: usize, k: usize, tile: u32, slices: u32, gn: usize) !Linear { const options = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; const layout = lane.Layout{ .n = n, .k = k, .format = .{ .bits = parts[0].bits, .group = parts[0].group }, .sk = slices }; try layout.validate(); @@ -122,13 +227,13 @@ fn prepareLanes(device: mtl.Device, allocator: std.mem.Allocator, owned: *std.Ar packTiles(.{ .words = p.weight.host(u32), .scales = p.scales.host(u16), .biases = p.biases.host(u16), .rows = p.weight.shape[0] }, base, n, k, parts[0].bits, parts[0].group, w.slice(u32, layout.weightWords()), sb.slice(u16, layout.metadataElements())); base += p.weight.shape[0]; } - return .{ .words = .{ .buffer = w }, .pairs = .{ .buffer = sb }, .n = @intCast(n), .k = @intCast(k), .tile = tile, .slices = slices, .bits = parts[0].bits, .group = parts[0].group }; + return .{ .words = .{ .buffer = w }, .pairs = .{ .buffer = sb }, .n = @intCast(n), .k = @intCast(k), .tile = tile, .slices = slices, .bits = parts[0].bits, .group = parts[0].group, .gn = gn }; } /// One member as stored: words [rows][k*bits/32], scales and biases [rows][k/group]. pub const Member = struct { words: []const u32, scales: []const u16, biases: []const u16, rows: usize }; -/// Copies a member into lane_projection.weight_layout rows [base, base + rows) of an n-row stack, bytes unchanged. +/// Copies a member's rows into lane_projection.weight_layout rows [base, base + rows), bytes unchanged. pub fn packTiles(p: Member, base: usize, n: usize, k: usize, bits: usize, group: usize, out_w: []u32, out_sb: []u16) void { const ng = k / group; const wp = group * bits / 32; diff --git a/zig/src/families/qwen27/quant_gpu.zig b/zig/src/families/qwen27/quant_gpu.zig index e38b913e2a..4d8dd66175 100644 --- a/zig/src/families/qwen27/quant_gpu.zig +++ b/zig/src/families/qwen27/quant_gpu.zig @@ -26,7 +26,7 @@ pub const Kernels = struct { /// Per [n, k] shape: whether the scalar twin's 1-2 row bits equal the tiles' (checked once on this chip). twins: std.AutoHashMapUnmanaged([2]usize, bool) = .empty, /// The same for simd_qmm_bits, per [bits, n, k]. - bit_twins: std.AutoHashMapUnmanaged([3]usize, bool) = .empty, + bit_twins: std.AutoHashMapUnmanaged([4]usize, bool) = .empty, /// Per [n, k]: whether simd_qmm's chunk pairs give its tiles' bits (checked once on this chip). pair_twins: std.AutoHashMapUnmanaged([2]usize, bool) = .empty, /// Set around prompt chunks: prompt rows take projection.promptSplit, decoded rows the decode split. @@ -126,26 +126,29 @@ pub const Kernels = struct { if (skip > 0 and (p.raw != null or skip >= p.n / 64)) return error.BadCoreProjectionRows; if (p.raw) |raw| if (raw.bits != 4 and k.tiled(raw)) { if (rows == 0 or rows > 128 or p.tile != 32) return error.BadCoreProjectionRows; - return k.strided(input_encoder, raw, x, y, rows, @intCast(raw.n)); + return k.strided(input_encoder, raw, x, y, rows, @intCast(raw.n), p.gn); }; var span = profile.begin(k.profiler, input_encoder, .dense); - const sk = if (k.prompt) projection.promptSplit(p.n, p.k) orelse p.slices else p.slices; + // Prompt rows take the checkpoint matrix's split, as decoded rows already do through `p.slices`. + const gn = if (p.gn != 0) p.gn else p.n; + const sk = if (k.prompt) projection.promptSplit(gn, p.k) orelse p.slices else p.slices; span.shape = .{ .n = p.n, .k = p.k, .sk = sk, .bits = p.bits }; errdefer span.cancel(); const e = span.encoder; _ = dims; if (rows == 0 or rows > @import("gpu_frame.zig").max_rows or p.tile != 32) return error.BadCoreProjectionRows; - if (p.raw) |raw| if (raw.bits == 4 and raw.group == 64 and raw.sum == .f32 and raw.n % 8 == 0) { - const m = row.simd.Matrix{ .w = raw.w, .w_off = raw.w_off, .scales = raw.scales, .s_off = raw.s_off, .biases = raw.biases, .b_off = raw.b_off, .n = raw.n, .k = raw.k }; - const scalar = try k.twin(m); // loads the tiles before they're read - const call = try row.simd.call(&k.simd.?, raw.n, raw.k, rows, scalar, try k.pairTwin(m)); - row.simd.encode(e, call, m, x.buffer, x.offset, y.buffer, y.offset); - e.barrier(); - try span.end(); - return; - }; if (p.raw) |raw| { + if (raw.bits == 4 and raw.group == 64 and raw.sum == .f32 and raw.n % 8 == 0) { + const m = row.simd.Matrix{ .w = raw.w, .w_off = raw.w_off, .scales = raw.scales, .s_off = raw.s_off, .biases = raw.biases, .b_off = raw.b_off, .n = raw.n, .k = raw.k }; + const scalar = try k.twin(m); // loads the tiles before they're read + const paired = try k.pairTwin(m); + const call = try row.simd.call(&k.simd.?, raw.n, raw.k, rows, scalar, paired, gn); + row.simd.encode(e, call, m, x.buffer, x.offset, y.buffer, y.offset); + e.barrier(); + try span.end(); + return; + } if (k.row_pipelines == null) k.row_pipelines = try row.Pipelines.load(k.device); const call = try row.call(&k.row_pipelines.?, raw, rows, false); e.setPipeline(call.pipeline); @@ -189,7 +192,7 @@ pub const Kernels = struct { if (skip > 0 or cut[1] > 0 or rows == 0 or rows > 128) return error.BadCoreProjectionRows; for (p.runs) |run| { if (run.linear.raw) |raw| { - try k.strided(input_encoder, raw, x, .{ .buffer = y.buffer, .offset = y.offset + @as(usize, run.first) * 2 }, rows, p.n); + try k.strided(input_encoder, raw, x, .{ .buffer = y.buffer, .offset = y.offset + @as(usize, run.first) * 2 }, rows, p.n, run.linear.gn); continue; } for (0..rows) |r| { // lane kernels take one row a launch until they write at a stride @@ -202,7 +205,8 @@ pub const Kernels = struct { } /// simd_qmm's tiles at any width (4-bit: its bits), else the row kernel; output row m at y + m * stride values. - fn strided(k: *Kernels, input_encoder: mtl.ComputeEncoder, raw: row.Weights, x: Ref, y: Ref, rows: u32, stride: u32) !void { + fn strided(k: *Kernels, input_encoder: mtl.ComputeEncoder, raw: row.Weights, x: Ref, y: Ref, rows: u32, stride: u32, whole: usize) !void { + // `whole` is the checkpoint matrix's rows (0: the staged rows), so a rank splits K as one Mac does. var span = profile.begin(k.profiler, input_encoder, .dense); span.shape = .{ .n = @intCast(raw.n), .k = @intCast(raw.k), .sk = 1, .bits = raw.bits }; errdefer span.cancel(); @@ -210,8 +214,8 @@ pub const Kernels = struct { if (k.tiled(raw)) { if (k.simd_bits == null) k.simd_bits = try row.simd_bits.Pipelines.load(k.device); const m = row.simd_bits.Matrix{ .w = raw.w, .w_off = raw.w_off, .scales = raw.scales, .s_off = raw.s_off, .biases = raw.biases, .b_off = raw.b_off }; - const scalar = rows <= row.simd.max_scalar_rows and try k.bitTwin(m, raw); // loads the tiles before they're read - const call = try row.simd_bits.call(&k.simd_bits.?, raw.bits, raw.n, raw.k, rows, stride, scalar); + const scalar = rows <= row.simd.max_scalar_rows and try k.bitTwin(m, raw, whole); // loads the tiles before they're read + const call = try row.simd_bits.call(&k.simd_bits.?, raw.bits, raw.n, raw.k, rows, stride, scalar, whole); row.simd_bits.encode(e, call, m, x.buffer, x.offset, y.buffer, y.offset); e.barrier(); try span.end(); @@ -231,11 +235,11 @@ pub const Kernels = struct { try span.end(); } - /// Checks simd_qmm_bits' scalar twin against its tiles once per width and shape. - fn bitTwin(k: *Kernels, m: row.simd_bits.Matrix, raw: row.Weights) !bool { - const key = [3]usize{ raw.bits, raw.n, raw.k }; + /// Checks simd_qmm_bits' scalar twin against its tiles once per width, shape and K split key. + fn bitTwin(k: *Kernels, m: row.simd_bits.Matrix, raw: row.Weights, whole: usize) !bool { + const key = [4]usize{ raw.bits, raw.n, raw.k, whole }; if (k.bit_twins.get(key)) |ok| return ok; - const ok = try row.simd_bits.twinExact(&k.simd_bits.?, k.device, raw.bits, m, raw.n, raw.k); + const ok = try row.simd_bits.twinExact(&k.simd_bits.?, k.device, raw.bits, m, raw.n, raw.k, whole); if (!ok) std.log.info("qwen27: {d}-bit {d}x{d} rows take simdgroup tiles at every width (the scalar twin's bits differ on this chip)", .{ raw.bits, raw.n, raw.k }); try k.bit_twins.put(k.allocator, key, ok); return ok; diff --git a/zig/src/families/qwen27/session.zig b/zig/src/families/qwen27/session.zig index b0296cdf95..220c9b205d 100644 --- a/zig/src/families/qwen27/session.zig +++ b/zig/src/families/qwen27/session.zig @@ -34,17 +34,29 @@ pub const Session = struct { pub fn step(s: Session, token: u32) !void { try s.feed(&.{token}, .last); } + /// The last round's token: a pair reads the head seam's `picks`, whose `maxInt(u32)` sentinel marks a broken link. pub fn greedy(s: Session) !u32 { if (s.runner.failed or s.runner.active != null or s.runner.offsets[s.slot] == 0) return error.RoundNotReady; + if (s.runner.reduce) |reduce| { + const picks = s.runner.picks orelse return error.HeadSeamMissing; + const words: [*]const u32 = @ptrCast(@alignCast(picks.buffer.contents() + picks.offset)); + const token = words[0]; + if (token >= s.runner.model.shard.global.vocab) return if (reduce.broken()) error.PairFailed else error.NonfiniteLogits; + return token; + } const values = s.runner.model.frame.get(.logits).buffer.slice(u16, s.runner.model.config.vocab); return argmax(values); } - /// Hash committed recurrent bytes, visible KV prefixes and the last row's logits (not unused capacity). + /// Per-rank digest (same rank, never the peer): committed recurrent bytes, visible KV prefixes, the last logits. pub fn fingerprint(s: Session) ![32]u8 { const r = s.runner; if (r.active != null or r.failed or r.offsets[0] == 0) return error.RoundNotReady; var hash = std.crypto.hash.sha2.Sha256.init(.{}); hash.update(std.mem.asBytes(&r.offsets[0])); + if (r.reduce != null) { + const picks = r.picks orelse return error.HeadSeamMissing; + hash.update(picks.buffer.contents()[picks.offset..][0..@sizeOf(u32)]); + } for (0..r.gdn.budget.layers) |layer| { hash.update(try r.gdn.committed(layer, 0, true)); hash.update(try r.gdn.committed(layer, 0, false)); @@ -64,7 +76,8 @@ pub const Session = struct { } fn feed(s: Session, ids: []const u32, head: @import("forward.zig").Head) !void { - var round = try Round.init(s.runner.allocator, &.{.{ .slot = s.slot, .start = s.runner.offsets[s.slot], .capacity = s.runner.capacity, .ids = ids }}, @intCast(s.runner.model.config.conv_kernel), s.runner.slots, @intCast(s.runner.model.config.vocab)); + // Ids are global on every rank (only head rows split): validate against the whole vocab, not `config.vocab`. + var round = try Round.init(s.runner.allocator, &.{.{ .slot = s.slot, .start = s.runner.offsets[s.slot], .capacity = s.runner.capacity, .ids = ids }}, @intCast(s.runner.model.config.conv_kernel), s.runner.slots, @intCast(s.runner.model.shard.global.vocab)); defer round.deinit(); if (!s.reference_tree) { try s.runner.advance(&round, head); @@ -125,3 +138,35 @@ test "top-two margins retain first-ID ties and reject nonfinite scores" { try std.testing.expectEqualSlices(f32, &.{ 2, 2 }, &top.logits); try std.testing.expectError(error.NonfiniteLogits, topTwo(&.{ 0x3f80, 0x7fc0 })); } + +/// One rank's block of the vocabulary: `values` are its local logits rows, `lo` its first global token id. +pub const Block = struct { values: []const u16, lo: u32 }; + +/// Top two over the ranks' blocks, shifted by `lo`: a cross-block tie keeps the lower id, as `argmax`/`pair_pick` do. +pub fn topTwoBlocks(blocks: []const Block) !TopTwo { + if (blocks.len == 0) return error.BadLogits; + var result = TopTwo{ .ids = .{ 0, 1 }, .logits = .{ -std.math.inf(f32), -std.math.inf(f32) } }; + for (blocks) |block| { + const one = try topTwo(block.values); + for (one.ids, one.logits) |id, value| { + const global = block.lo + id; + if (value > result.logits[0]) { + result.logits[1] = result.logits[0]; + result.ids[1] = result.ids[0]; + result.logits[0] = value; + result.ids[0] = global; + } else if (value > result.logits[1]) { + result.logits[1] = value; + result.ids[1] = global; + } + } + } + return result; +} +test "a pair's merged top two shifts each block to global ids and a cross-block tie keeps the lower id" { + const top = try topTwoBlocks(&.{ .{ .values = &.{ 0x3f80, 0x4000, 0x4040 }, .lo = 0 }, .{ .values = &.{ 0x4040, 0xbf80 }, .lo = 4 } }); + try std.testing.expectEqualSlices(u32, &.{ 2, 4 }, &top.ids); + try std.testing.expectEqualSlices(f32, &.{ 3, 3 }, &top.logits); + try std.testing.expectError(error.NonfiniteLogits, topTwoBlocks(&.{.{ .values = &.{ 0x3f80, 0x7fc0 }, .lo = 0 }})); + try std.testing.expectError(error.BadLogits, topTwoBlocks(&.{})); +} diff --git a/zig/src/metal/command.zig b/zig/src/metal/command.zig index 49118d0b3d..45c2b29358 100644 --- a/zig/src/metal/command.zig +++ b/zig/src/metal/command.zig @@ -195,4 +195,14 @@ pub const ComputeEncoder = struct { pub fn end(self: ComputeEncoder) void { objc.msg(void, self.id, "endEncoding", .{}); } + + /// The command buffer this encoder was opened from, for the event ops a seam encodes BETWEEN encoders. + pub fn commandBuffer(self: ComputeEncoder) CommandBuffer { + return .{ .id = objc.msg(Id, self.id, "commandBuffer", .{}) }; + } + + /// Whether this encoder dispatches serially or concurrently: a seam reopens it in the same mode. + pub fn dispatchType(self: ComputeEncoder) types.DispatchType { + return @fromBackingInt(@intCast(objc.msg(usize, self.id, "dispatchType", .{}))); + } }; diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 4fa4fbe147..4cc2deba22 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -32,6 +32,8 @@ pub const exlr8 = @import("core/exlr8.zig"); pub const exlr8_projection = @import("core/exlr8_projection.zig"); pub const exlr8_dense_gpu = @import("core/exlr8_dense_gpu.zig"); pub const exlr8_prompt = @import("core/exlr8_prompt.zig"); +/// The two-Mac pair seam's row split: the transport, its layout and the one-Mac Loopback the stage-2 checks use. +pub const pair_rows = @import("core/pair_rows.zig"); pub const glm_dsa = @import("families/glm_dsa/glm_dsa.zig"); pub const nemotron = @import("families/nemotron/nemotron.zig"); pub const flashnext_vision_contract = @import("families/flashnext/vision_contract.zig"); diff --git a/zig/tests/affine_ref.zig b/zig/tests/affine_ref.zig index b367fa120a..691ecf11ce 100644 --- a/zig/tests/affine_ref.zig +++ b/zig/tests/affine_ref.zig @@ -272,7 +272,7 @@ fn affine(t: *Tally, device: mtl.Device, queue: mtl.Queue, s: *const Synth, ref: /// simd_qmm over rows [start, start + m), the scalar twin when `scalar` and it takes m rows. fn simdWindow(e: mtl.ComputeEncoder, p: *const sq.Pipelines, s: *const Synth, y: Out, scalar: bool, start: usize, m: usize) !void { const mat = sq.Matrix{ .w = s.w, .scales = s.s, .biases = s.b, .n = s.n, .k = s.k }; - sq.encode(e, try sq.call(p, s.n, s.k, m, scalar, true), mat, s.x, start * s.k * 2, y.buf, start * s.n * 2); + sq.encode(e, try sq.call(p, s.n, s.k, m, scalar, true, 0), mat, s.x, start * s.k * 2, y.buf, start * s.n * 2); } /// simd_qmm (4-bit, group 64): tiles at every window against float64 and the wide call; the twin at 1 and 2 rows. diff --git a/zig/tests/contract/metal_affine.zig b/zig/tests/contract/metal_affine.zig index 46f921e9f8..5d65c349db 100644 --- a/zig/tests/contract/metal_affine.zig +++ b/zig/tests/contract/metal_affine.zig @@ -198,11 +198,11 @@ pub const Proj = struct { e.dispatchGroups(mtl.Size.of(cl.groups, cl.slots, 1), mtl.Size.of(cl.threads, 1, 1)); }, .simd => { - const cl = try sq.call(&p.sq_p.?, shape.n, shape.k, shape.rows, false, true); + const cl = try sq.call(&p.sq_p.?, shape.n, shape.k, shape.rows, false, true, 0); sq.encode(e, cl, .{ .w = p.w.buf, .w_off = p.w.off, .scales = p.s.buf, .s_off = p.s.off, .biases = p.b.buf, .b_off = p.b.off, .n = shape.n, .k = shape.k }, p.x.buf, p.x.off, p.out.buf, p.out.off); }, .simd_bits => { - const cl = try sqb.call(&p.sqb_p.?, p.v.bits, shape.n, shape.k, shape.rows, shape.n, false); + const cl = try sqb.call(&p.sqb_p.?, p.v.bits, shape.n, shape.k, shape.rows, shape.n, false, 0); sqb.encode(e, cl, .{ .w = p.w.buf, .w_off = p.w.off, .scales = p.s.buf, .s_off = p.s.off, .biases = p.b.buf, .b_off = p.b.off }, p.x.buf, p.x.off, p.out.buf, p.out.off); }, .affine_mm => { diff --git a/zig/tests/contract/metal_limits.zig b/zig/tests/contract/metal_limits.zig index 53609bffa0..c0e267c0d0 100644 --- a/zig/tests/contract/metal_limits.zig +++ b/zig/tests/contract/metal_limits.zig @@ -100,6 +100,8 @@ fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch) !La text = sources.core_deltanet; } else if (std.mem.startsWith(u8, d.name, "qwen27 ")) { text = if (has(d.name, " copy ")) tf.qwen27.runtime_sources.copy else tf.qwen27.runtime_sources.glue; + } else if (std.mem.startsWith(u8, d.name, "metal pair ")) { + text = sources.core_pair_rows; } else if (has(d.name, "tile_qmm")) text = sources.core_tile_qmm else if (has(d.name, "softmax")) text = sources.ops_softmax else if (has(d.name, "rms")) text = sources.ops_embed_norm else if (has(d.name, "top-k")) text = topk_text else if (has(d.name, "row projection")) text = sources.core_row_projection else if (has(d.name, "simd_qmm_bits")) text = sources.core_simd_qmm_bits else if (has(d.name, "simd_qmm")) text = sources.core_simd_qmm else if (has(d.name, "affine_mm")) { const bits: u8 = if (has(d.name, "q8")) 8 else 4; const group: u16 = if (has(d.name, "g32")) 32 else 64; diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 46801e171c..9896a1cd97 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -17,6 +17,8 @@ const drama_combine = @import("metal_drama_combine.zig"); const drama_cfg = @import("metal_drama_cfg.zig"); const glm_eh_k = @import("metal_glm_eh.zig"); const glm_rows_k = @import("metal_glm_rows.zig"); +const pair_k = @import("metal_pair.zig"); +const pair_controls = @import("metal_pair_controls.zig"); const flash_pool = @import("metal_flashnext_pool.zig"); const flash = @import("metal_flashnext.zig"); @@ -202,6 +204,32 @@ pub fn main(init: std.process.Init) !u8 { try rows_list.append(a, r); try kernels.append(a, r.kernel()); } + var pair_list: std.ArrayList(*pair_k.Pair) = .empty; + defer { + for (pair_list.items) |p| { + p.deinit(); + a.destroy(p); + } + pair_list.deinit(a); + } + for (std.enums.values(pair_controls.Mutation)) |m| { + for (pair_k.all) |w| { + const control_for = switch (m) { + .none => true, + .stage_peer_half => w == .stage, + .gather_own_halves => w == .gather, + .argmax_high_tie, .argmax_no_flag => w == .argmax, + .pick_no_flag, .pick_high_tie => w == .pick, + }; + if (!control_for) continue; + if (!wanted(o, if (m == .none) pair_k.spec(w).name else m.name())) continue; + const p = try a.create(pair_k.Pair); + errdefer a.destroy(p); + p.* = try pair_k.Pair.init(a, device, queue, w, m); + try pair_list.append(a, p); + try kernels.append(a, p.kernel()); + } + } var flash_list: [8]flash.Probe = undefined; var flash_made: usize = 0; defer for (flash_list[0..flash_made]) |*probe| probe.deinit(); diff --git a/zig/tests/contract/metal_pair.zig b/zig/tests/contract/metal_pair.zig new file mode 100644 index 0000000000..9a15b020f8 --- /dev/null +++ b/zig/tests/contract/metal_pair.zig @@ -0,0 +1,485 @@ +//! The pair seam's four kernels (core/pair_rows.metal), and the float64 rules the gates check. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); +const rig_ = @import("rig.zig"); +const controls = @import("metal_pair_controls.zig"); + +const c = contract.spec; +const g = contract.gates; + +/// The seam's four kernels; the tail of each registry name is the MSL symbol launchOf hands the source. +pub const Which = enum { stage, gather, argmax, pick }; + +pub const all = [_]Which{ .stage, .gather, .argmax, .pick }; + +/// What pair_rows.zig dispatches per kernel, so `init` can compare it against Metal's own reflection. +pub const Launched = struct { threads: u32, tg_memory: u32 }; + +pub fn launch(w: Which) Launched { + return switch (w) { + .stage => .{ .threads = 256, .tg_memory = 0 }, + .gather => .{ .threads = 256, .tg_memory = 0 }, + .argmax => .{ .threads = 1024, .tg_memory = 384 }, + .pick => .{ .threads = 32, .tg_memory = 0 }, + }; +} + +/// The MSL symbol pair_rows.metal defines. +pub fn symbol(w: Which) []const u8 { + return switch (w) { + .stage => "pair_stage", + .gather => "pair_gather", + .argmax => "pair_argmax", + .pick => "pair_pick", + }; +} + +/// The spec's name: `metal `, so launchOf's last word is the symbol and its prefix picks core_pair_rows. +pub fn name(w: Which) []const u8 { + return switch (w) { + .stage => "metal pair stage pair_stage", + .gather => "metal pair gather pair_gather", + .argmax => "metal pair argmax pair_argmax", + .pick => "metal pair pick pair_pick", + }; +} + +/// Qwen3.8-27B's rank-local slice: the seam reduces ids [124160, 248320) on rank 1. +const logit_lo: u32 = 124160; + +/// An argmax id word's bit 31: the row had a logit that is not finite, so pair_pick sends the row to the id sentinel. +const flag: u32 = 0x80000000; + +/// The id a flagged row, and a broken link, read as: no token, visibly wrong. +const sentinel: u32 = 0xFFFFFFFF; + +/// A thread's chain over every 1024th logit, five shuffle levels, then the 32 simdgroup winners. +fn chainLogits(s: c.Shape) usize { + return (s.k + 1023) / 1024 + 5 + 5; +} + +/// Only moves and compares values: no sums to carry an accumulator's error. +fn noSums(_: c.Shape) usize { + return 0; +} + +const move_why = "a move writes each bf16 element as it read it: no sum, no rounding, so no accumulator's error can enter"; + +const argmax_rule = "a winner is (value bits, id): NaN loses every comparison, a row of only NaNs reports the -inf sentinel with id 0xFFFFFFFF, and a row with any non-finite logit reports NaN bits and id | 0x80000000"; +const pick_rule = "the pick copies one rank's (value, id) word pair; an id with bit 31 set, and the broken-link word, draw 0xFFFFFFFF: the same sentinel, so one downstream check covers both"; + +/// A kernel's spec. The raw kernels check themselves; the two moves' bound is exact, so gate 3 keeps every edge. +pub fn spec(w: Which) c.Spec { + return switch (w) { + // One spec for both moves; a gather's `n` is its frame row stride, and the 17,408 shape is seam B's. + .stage, .gather => .{ + .name = name(w), + .accum = .none, + .depth = noSums, + .out = .bf16, + .tiles = .{ .{}, .{}, .{ .size = 256, .step = 256, .min = 256 } }, + .served = if (w == .stage) &.{.{ .rows = 8, .n = 5120, .k = 5120 }} else &.{ + .{ .rows = 8, .n = 10240, .k = 5120 }, .{ .rows = 8, .n = 17408, .k = 8704 }, + .{ .rows = 8, .n = 10496, .k = 5120 }, + }, + .sums = .{ .f16 = 0, .bf16 = 0, .exact = false, .why = move_why }, + }, + .argmax => .{ + .name = name(w), + .accum = .none, + .depth = chainLogits, + .out = .raw, + .nan = .{ .stated = argmax_rule }, + .tiles = .{ .{}, .{}, .{ .size = 2048, .step = 64, .min = 64 } }, + .served = &.{.{ .rows = 8, .n = 1, .k = 124160 }}, + }, + .pick => .{ + .name = name(w), + .accum = .none, + .depth = noSums, + .out = .raw, + .nan = .{ .stated = pick_rule }, + .served = &.{ .{ .rows = 32, .n = 1, .k = 1 }, .{ .rows = 128, .n = 1, .k = 1 } }, + }, + }; +} + +/// `constant PairRows&` at buffer 1: `width`, the frame row `stride` and which rank's half this side writes. +const PairRows = extern struct { + width: u32, + stride: u32, + rank: u32, + x: u32, +}; + +fn rowsArgs(width: usize, stride: usize, rank: u32, x: u32) PairRows { + return .{ .width = @intCast(width), .stride = @intCast(stride), .rank = rank, .x = x }; +} + +fn bit(v: f32) u32 { + return @bitCast(v); +} + +fn floatOf(v: u32) f32 { + return @bitCast(v); +} + +/// One prepared call: the buffers the kernel is bound to, the host values its reference reads, and its own check. +pub const Pair = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + which: Which, + mutation: controls.Mutation, + pipe: mtl.Pipeline, + /// The hooks this probe exposes: the raw kernels check their own rule, the two computing ones use the reference. + v: g.VTable, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + rows: usize = 0, + width: usize = 0, + /// The rank this prepare writes first: the gather's `mine` half is its payload. + rank: u32 = 0, + /// The row stride in PairRows, in bf16 elements: the stage's source frame row, the gather's destination row. + stride: usize = 0, + + // Host inputs of the last prepare, as the reference and the check read them (freed by `drop`). + sh: ?[]u16 = null, // stage: the src payload, bf16 (rows, each a frame row of `stride`) + mh: ?[]u16 = null, // gather: this rank's half (mine), bf16 + nh: ?[]u16 = null, // gather: the peer's half (recv), bf16 + kh: ?[]u16 = null, // argmax: the bf16 logits + pw: ?[]u32 = null, // pick: this rank's (value bits, id) words + qw: ?[]u32 = null, // pick: the peer's words + + msg: [192]u8 = undefined, + sync: rig_.Ref = undefined, + mine: rig_.Ref = undefined, + peer: rig_.Ref = undefined, + logits: rig_.Ref = undefined, + out: rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, which: Which, mutation: controls.Mutation) !Pair { + const src = if (mutation == .none) sources.core_pair_rows else try controls.source(a, mutation); + defer if (mutation != .none) a.free(src); + const pipe = try rig_.pipeline(device, src, symbol(which)); + errdefer pipe.deinit(); + const want = launch(which); + const tg = rig_.tgMemory(pipe); + if (tg != want.tg_memory) { + std.debug.print("metal pair: {s} takes {d} bytes of threadgroup memory, its record says {d}\n", .{ symbol(which), tg, want.tg_memory }); + return error.ThreadgroupMemoryMismatch; + } + var v = g.hooks(Pair); + if (which != .argmax and which != .pick) v.check = null; + if (which != .stage and which != .gather) v.bounds = null; + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .which = which, .mutation = mutation, .pipe = pipe, .v = v }; + } + + pub fn deinit(self: *Pair) void { + self.drop(); + self.rig.deinit(); + self.pipe.deinit(); + } + + pub fn kernel(self: *Pair) g.Kernel { + var s = spec(self.which); + if (self.mutation != .none) { + s.name = self.mutation.name(); + return .{ .spec = s, .ctx = self, .vt = &self.v, .control = self.mutation.gate() }; + } + return .{ .spec = s, .ctx = self, .vt = &self.v }; + } + + fn drop(self: *Pair) void { + if (self.sh) |s| self.a.free(s); + if (self.mh) |s| self.a.free(s); + if (self.nh) |s| self.a.free(s); + if (self.kh) |s| self.a.free(s); + if (self.pw) |s| self.a.free(s); + if (self.qw) |s| self.a.free(s); + self.sh = null; + self.mh = null; + self.nh = null; + self.kh = null; + self.pw = null; + self.qw = null; + } + + /// Ordinary values: what the seam carries in production, symmetric about zero. + fn ordinary(rng: std.Random) f32 { + return rng.float(f32) * 2 - 1; + } + + /// A bf16 payload: what the two moves carry, with gate 1's non-finite plants where the reference looks for them. + fn payload(self: *Pair, len: usize, fill: g.Fill, rng: std.Random) ![]u16 { + const s = try self.a.alloc(u16, len); + for (s) |*v| v.* = c.toBf16(ordinary(rng)); + if (fill == .nonfinite) { + const nan = c.toBf16(std.math.nan(f32)); + const inf = c.toBf16(std.math.inf(f32)); + const ninf = c.toBf16(-std.math.inf(f32)); + for ([_]u16{ nan, inf, ninf, nan, ninf }, 0..) |plant, p| { + const at = (p + 1) * (len / 8 + 1) + p * 3; + if (at < len) s[at] = plant; + } + } + return s; + } + + /// Inputs for `shape` as `fill` asks, outputs poisoned and guarded; false when the kernel does not serve it. + pub fn prepare(self: *Pair, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.rows > 512) return false; + self.drop(); + self.rig.reset(); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + switch (self.which) { + .stage, .gather => { + if (shape.k == 0 or shape.k > 32768 or shape.k % 256 != 0) return false; + // The gather writes this rank's half first; alternating the rank with the seed covers both orders. + self.rank = @intCast(seed & 1); + self.sync = try self.rig.scratch(128); + if (self.which == .stage) { + // The source is a frame row: the stage reads this rank's `width` columns out of its `stride`. + self.stride = 2 * shape.k; + const len = shape.rows * self.stride; + const mine = try self.payload(len, fill, rng); + self.sh = mine; + self.mine = try self.rig.input(std.mem.sliceAsBytes(mine)); + self.out = try self.rig.output(shape.rows * shape.k * 2); + } else { + const n = shape.rows * shape.k; + // A frame row holds both halves, so its stride is at least 2*width: `n` sets a wider one. + self.stride = if (shape.n > 2 * shape.k) shape.n else 2 * shape.k; + const mine = try self.payload(n, fill, rng); + const peer = try self.payload(n, fill, rng); + self.mh = mine; + self.nh = peer; + self.mine = try self.rig.input(std.mem.sliceAsBytes(mine)); + self.peer = try self.rig.input(std.mem.sliceAsBytes(peer)); + self.out = try self.rig.output(2 * ((shape.rows - 1) * self.stride + 2 * shape.k)); + } + }, + .argmax => { + if (shape.k == 0 or shape.k > 262144) return false; + const n = shape.rows * shape.k; + const k = try self.a.alloc(u16, n); + self.kh = k; + for (k) |*v| v.* = c.toBf16(ordinary(rng)); + // Two in-range ids at +2.0: the lower id sits in the first half, so the tie rule's pick is fixed. + for (0..shape.rows) |r| { + const row = k[r * shape.k ..][0..shape.k]; + row[0] = c.toBf16(2.0); + row[shape.k / 2] = c.toBf16(2.0); + } + if (fill == .nonfinite) { + const nan = c.toBf16(std.math.nan(f32)); + const inf = c.toBf16(std.math.inf(f32)); + const ninf = c.toBf16(-std.math.inf(f32)); + for (0..shape.rows) |r| { + const row = k[r * shape.k ..][0..shape.k]; + row[shape.k / 3] = switch (r % 4) { + 0 => nan, + 1 => inf, + 2 => ninf, + else => row[shape.k / 3], + }; + if (r + 1 == shape.rows) for (row) |*v| { + v.* = nan; + }; + } + } + self.logits = try self.rig.input(std.mem.sliceAsBytes(k)); + self.out = try self.rig.output(shape.rows * 8); + }, + .pick => { + if (shape.n != 1 or shape.k != 1) return false; + self.pw = try self.a.alloc(u32, 2 * shape.rows); + self.qw = try self.a.alloc(u32, 2 * shape.rows); + for (0..shape.rows) |r| { + self.pw.?[2 * r] = bit(ordinary(rng)); + self.qw.?[2 * r] = bit(ordinary(rng)); + self.pw.?[2 * r + 1] = rng.uintLessThan(u32, 1 << 20); + self.qw.?[2 * r + 1] = rng.uintLessThan(u32, 1 << 20); + } + // Equal values on both ranks: the tie rule decides, and its direction alternates row by row. + for (0..shape.rows) |r| { + if (r % 4 != 2 and r % 4 != 3) continue; + self.qw.?[2 * r] = self.pw.?[2 * r]; + self.pw.?[2 * r + 1] = 300 + @as(u32, @intCast(r)); + self.qw.?[2 * r + 1] = if (r % 8 == 2) 700 + @as(u32, @intCast(r)) else 100 + @as(u32, @intCast(r)); + } + if (fill == .nonfinite) { + // Two plants prove the id flag is the rule: a flagged row draws the sentinel, a plain one picks. + for (0..shape.rows) |r| { + switch (r % 4) { + 0 => self.pw.?[2 * r + 1] |= flag, + 1 => self.qw.?[2 * r + 1] |= flag, + 2 => self.pw.?[2 * r] = bit(-std.math.inf(f32)), + else => self.qw.?[2 * r] = bit(std.math.inf(f32)), + } + } + } + self.sync = try self.rig.scratch(128); + self.mine = try self.rig.input(std.mem.sliceAsBytes(self.pw.?)); + self.peer = try self.rig.input(std.mem.sliceAsBytes(self.qw.?)); + self.out = try self.rig.output(shape.rows * 4); + }, + } + self.shape = shape; + self.rows = shape.rows; + self.width = shape.k; + return true; + } + + /// One launch on the prepared buffers, waited for. + pub fn run(self: *Pair) !void { + const p = self.rig.begin(); + const e = p.e; + const groups = mtl.Size.of(self.rows, 1, 1); + const threads = mtl.Size.of(launch(self.which).threads, 1, 1); + switch (self.which) { + .stage => { + @memset(self.rig.bytes(self.sync, 128), 0); + e.setPipeline(self.pipe); + e.setBuffer(self.sync.buf, self.sync.off, 0); + e.setValue(rowsArgs(self.width, self.stride, self.rank, 1), 1); + e.setBuffer(self.mine.buf, self.mine.off, 2); + e.setBuffer(self.out.buf, self.out.off, 3); + e.dispatchGroups(groups, threads); + }, + .gather => { + @memset(self.rig.bytes(self.sync, 128), 0); + const frame = self.rig.bytes(self.out, 2 * self.count()); + @memset(frame, 0); // the frame's gaps: no launch writes them, so their planted value is zero + e.setPipeline(self.pipe); + e.setBuffer(self.sync.buf, self.sync.off, 0); + e.setValue(rowsArgs(self.width, self.stride, self.rank, 1), 1); + e.setBuffer(self.mine.buf, self.mine.off, 2); + e.setBuffer(self.peer.buf, self.peer.off, 3); + e.setBuffer(self.out.buf, self.out.off, 4); + e.dispatchGroups(groups, threads); + }, + .argmax => { + e.setPipeline(self.pipe); + e.setBuffer(self.logits.buf, self.logits.off, 0); + const dims = [4]u32{ @intCast(self.width), logit_lo, @intCast(self.width), 1 }; + e.setValue(dims, 1); + e.setBuffer(self.out.buf, self.out.off, 2); + e.dispatchGroups(groups, threads); + }, + .pick => { + @memset(self.rig.bytes(self.sync, 128), 0); + e.setPipeline(self.pipe); + e.setBuffer(self.sync.buf, self.sync.off, 0); + e.setValue([2]u32{ @intCast(self.rows), 1 }, 1); + e.setBuffer(self.mine.buf, self.mine.off, 2); + e.setBuffer(self.peer.buf, self.peer.off, 3); + e.setBuffer(self.out.buf, self.out.off, 4); + e.dispatchGroups(mtl.Size.of(1, 1, 1), threads); + }, + } + try self.rig.finish(p); + } + + /// Output elements of the prepared call: bf16 elements for the two moves, the two raw kernels' bytes. + pub fn count(self: *Pair) usize { + return switch (self.which) { + .stage => self.rows * self.width, + .gather => (self.rows - 1) * self.stride + 2 * self.width, + .argmax => self.rows * 8, + .pick => self.rows * 4, + }; + } + + /// Float64 outputs of the prepared inputs and each one's sums' term magnitude (a move's term is the value itself). + pub fn reference(self: *Pair, want: []f64, mag: []f64) !void { + const width = self.width; + switch (self.which) { + .stage => { + const sh = self.sh orelse return error.NoReference; + for (0..self.rows) |r| { + const row = sh[r * self.stride ..][0..width]; + for (0..width) |i| { + const v = c.fromBf16(row[i]); + want[r * width + i] = v; + mag[r * width + i] = @abs(v); + } + } + }, + .gather => { + const mh = self.mh orelse return error.NoReference; + const nh = self.nh orelse return error.NoReference; + @memset(want, 0); // the frame's gaps: no launch writes them, so their reference value is zero + for (0..self.rows) |row| { + const at = row * self.stride; + for (0..width) |i| { + want[at + i] = c.fromBf16(if (self.rank == 0) mh[row * width + i] else nh[row * width + i]); + want[at + width + i] = c.fromBf16(if (self.rank == 0) nh[row * width + i] else mh[row * width + i]); + } + } + for (want, mag) |w, *m| m.* = @abs(w); + }, + else => return error.NoReference, + } + } + + /// The moves are exact: every output element must equal the reference's bf16 bit for bit. + pub fn bounds(self: *Pair, want: []const f64, mag: []const f64, allow: []f64) !void { + _ = want; + _ = mag; + if (self.which != .stage and self.which != .gather) return error.NotMove; + for (allow) |*l| l.* = 0; + } + + /// The seam's own checks: the raw kernels' rule at gate 1 (non-finite inputs) and gate 3 (every served edge). + pub fn check(self: *Pair, _: g.Gate) !?[]const u8 { + const out = try self.rig.outputBytes(); + switch (self.which) { + .argmax => { + const k = self.kh orelse return error.NoReference; + const n = self.width; + for (0..self.rows) |r| { + var best: f32 = -std.math.inf(f32); + var at: u32 = sentinel; + var bad = false; + for (0..n) |i| { + const v = c.fromBf16(k[r * n + i]); + if (!std.math.isFinite(v)) { + bad = true; + continue; + } + if (v > best) { + best = v; + at = logit_lo + @as(u32, @intCast(i)); + } + } + const got = std.mem.readInt(u32, out[r * 8 ..][0..4], .little); + const id = std.mem.readInt(u32, out[r * 8 + 4 ..][0..4], .little); + if (bad) { + if (!std.math.isNan(floatOf(got))) return try std.fmt.bufPrint(&self.msg, "row {d}: value {e} where a non-finite logit must write NaN", .{ r, floatOf(got) }); + if (id != at | flag) return try std.fmt.bufPrint(&self.msg, "row {d}: id {d} where flagging gives {d}", .{ r, id, at | flag }); + } else if (got != bit(best) or id != at) { + return try std.fmt.bufPrint(&self.msg, "row {d}: ({e} {d}) where its own rule gives ({e} {d})", .{ r, floatOf(got), id, best, at }); + } + } + }, + .pick => { + const pw = self.pw orelse return error.NoReference; + const qw = self.qw orelse return error.NoReference; + for (0..self.rows) |r| { + const vm = floatOf(pw[2 * r]); + const am = pw[2 * r + 1]; + const vp = floatOf(qw[2 * r]); + const ap = qw[2 * r + 1]; + const flagged = ((am | ap) & flag) != 0; + const want = if (flagged) sentinel else if (vp > vm or (vp == vm and ap < am)) ap else am; + const got = std.mem.readInt(u32, out[r * 4 ..][0..4], .little); + if (got != want) return try std.fmt.bufPrint(&self.msg, "row {d} picks {d} where its own rule gives {d}", .{ r, got, want }); + } + }, + else => return error.NotRaw, + } + return null; + } +}; diff --git a/zig/tests/contract/metal_pair_controls.zig b/zig/tests/contract/metal_pair_controls.zig new file mode 100644 index 0000000000..1e7f4edbab --- /dev/null +++ b/zig/tests/contract/metal_pair_controls.zig @@ -0,0 +1,59 @@ +//! Planted controls for the pair seam's four kernels: one fault in pair_rows.metal's text, and the gate to catch it. +const std = @import("std"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); + +pub const Mutation = enum { + none, + stage_peer_half, + gather_own_halves, + argmax_high_tie, + argmax_no_flag, + pick_no_flag, + pick_high_tie, + + pub fn name(m: Mutation) []const u8 { + return switch (m) { + .none => "", + .stage_peer_half => "control: metal pair stage, this rank stages the peer's half", + .gather_own_halves => "control: metal pair gather, this rank writes its own half on both sides", + .argmax_high_tie => "control: metal pair argmax, a tie goes to the highest id", + .argmax_no_flag => "control: metal pair argmax, a nonfinite logit is not flagged", + .pick_no_flag => "control: metal pair pick, a flagged row is not turned into the sentinel", + .pick_high_tie => "control: metal pair pick, a tie goes to the higher id", + }; + } + + /// The gate that must go red. + pub fn gate(m: Mutation) ?contract.gates.Gate { + return switch (m) { + .none => null, + .argmax_no_flag, .pick_no_flag => .nonfinite, + else => .shapes, + }; + } + + const Edit = struct { from: []const u8, to: []const u8 }; + + fn edit(m: Mutation) Edit { + return switch (m) { + .none => unreachable, + // The stage's source is a frame row of `stride` = 2*width: the peer's half is in bounds and different. + .stage_peer_half => .{ .from = "const ulong in = ulong(r) * a.stride;", .to = "const ulong in = ulong(r) * a.stride + a.width;" }, + .gather_own_halves => .{ .from = "out[got + i] = ushort(word >> ((i & 1) * 16));", .to = "out[got + i] = mine[mrow + i];" }, + .argmax_high_tie => .{ .from = "if (ob > best || (ob == best && oa < at)) best = ob, at = oa;", .to = "if (ob > best || (ob == best && oa > at)) best = ob, at = oa;" }, + .argmax_no_flag => .{ .from = "if (!isfinite(v)) { bad = true; continue; }", .to = "if (false) { bad = true; continue; }" }, + .pick_high_tie => .{ .from = "(vp > vm || (vp == vm && ap < am)) ? ap : am;", .to = "(vp > vm || (vp == vm && ap > am)) ? ap : am;" }, + // The flag is planted only at the nonfinite fill, so dropping the rule is invisible at shapes. + .pick_no_flag => .{ .from = "((am | ap) & PAIR_FLAG) != 0u", .to = "false" }, + }; + } +}; + +/// pair_rows.metal's text for the kernel under test; a control's one edit must change it, or it proves nothing. +pub fn source(a: std.mem.Allocator, m: Mutation) ![]const u8 { + if (m == .none) return sources.core_pair_rows; + const e = m.edit(); + if (std.mem.indexOf(u8, sources.core_pair_rows, e.from) == null) return error.ControlEditMissing; + return std.mem.replaceOwned(u8, a, sources.core_pair_rows, e.from, e.to); +} diff --git a/zig/tests/qwen27/attention.zig b/zig/tests/qwen27/attention.zig index 09b3f32eea..6b582e82d0 100644 --- a/zig/tests/qwen27/attention.zig +++ b/zig/tests/qwen27/attention.zig @@ -49,7 +49,7 @@ const Frame = struct { try op.preprocess(e, .{ .qg = f.get(.qg), .key = f.get(.raw_k), .q_gain = f.get(.q_gain), .k_gain = f.get(.k_gain), .norm_q = f.get(.norm_q), .norm_k = f.get(.norm_k), .query = f.get(.query), .rotated_key = f.get(.key), .positions = scratch.get(.positions), .rows = f.rows, .eps = 1e-6, .theta = 10000000 }); } fn encode(f: *const Frame, op: *const q.Attention, e: mtl.ComputeEncoder, scratch: *const q.Scratch, binding: *const q.Binding, plan: q.Plan, caches: []const Ref, gate: bool) !void { - try op.encode(e, scratch, binding, plan, f.get(.query), f.get(.key), f.get(.v), f.get(.out), caches, if (gate) f.get(.qg) else null, false); + try op.encode(e, scratch, binding, plan, f.get(.query), f.get(.key), f.get(.v), f.get(.out), caches, if (gate) f.get(.qg) else null, false, false); } }; const CacheOwner = struct { diff --git a/zig/tests/qwen27/attention_fused.zig b/zig/tests/qwen27/attention_fused.zig index c368e6d186..63a394a7a3 100644 --- a/zig/tests/qwen27/attention_fused.zig +++ b/zig/tests/qwen27/attention_fused.zig @@ -56,7 +56,7 @@ fn run(queue: mtl.Queue, op: *const attn.Attention, scratch: *attn.Scratch, bind try binding.upload(plan, caches); const cb = queue.commandBuffer(); const e = cb.compute(.concurrent); - try op.encode(e, scratch, binding, plan, query, key, value, out, caches, gate, false); + try op.encode(e, scratch, binding, plan, query, key, value, out, caches, gate, false, false); e.end(); cb.commit(); cb.wait(); @@ -98,7 +98,7 @@ fn bench(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, op: *const for (×) |*t| { const cb = queue.commandBuffer(); const e = cb.compute(.concurrent); - for (0..layers) |i| try op.encode(e, &scratch, &bindings[i], plan, .{ .buffer = query }, .{ .buffer = key }, .{ .buffer = value }, .{ .buffer = out }, caches[i .. i + 1], .{ .buffer = raw_gate }, false); + for (0..layers) |i| try op.encode(e, &scratch, &bindings[i], plan, .{ .buffer = query }, .{ .buffer = key }, .{ .buffer = value }, .{ .buffer = out }, caches[i .. i + 1], .{ .buffer = raw_gate }, false, false); e.end(); cb.commit(); cb.wait(); diff --git a/zig/tests/qwen27/attention_precision.zig b/zig/tests/qwen27/attention_precision.zig index 5621e9c7ba..4d703de3aa 100644 --- a/zig/tests/qwen27/attention_precision.zig +++ b/zig/tests/qwen27/attention_precision.zig @@ -157,7 +157,7 @@ pub fn main(init: std.process.Init) !void { try binding.upload(plan, &caches); const cb = queue.commandBuffer(); const e = cb.compute(.concurrent); - try op.encode(e, &scratch, &binding, plan, .{ .buffer = query }, .{ .buffer = key }, .{ .buffer = val }, .{ .buffer = out }, &caches, null, prompt and pass == 0); + try op.encode(e, &scratch, &binding, plan, .{ .buffer = query }, .{ .buffer = key }, .{ .buffer = val }, .{ .buffer = out }, &caches, null, prompt and pass == 0, false); e.end(); cb.commit(); cb.wait(); diff --git a/zig/tests/qwen27/forward_diff_runner.zig b/zig/tests/qwen27/forward_diff_runner.zig index 70c61eca42..84644a8d7a 100644 --- a/zig/tests/qwen27/forward_diff_runner.zig +++ b/zig/tests/qwen27/forward_diff_runner.zig @@ -29,7 +29,7 @@ const Hooks = struct { try r.model.glue.unstack(e, f.get(.mixed), f.get(.value), .{ .rows = rows, .width = 1024, .stride = 2048, .offset = 1024 }); e.barrier(); try r.attention.preprocess(e, .{ .qg = f.get(.q_gate), .key = f.get(.key), .q_gain = Weights.ref(weights.q_norm), .k_gain = Weights.ref(weights.k_norm), .norm_q = f.get(.norm_query), .norm_k = f.get(.norm_key), .query = f.get(.query), .rotated_key = f.get(.rotated_key), .positions = f.get(.positions), .rows = rows, .eps = r.model.config.eps, .theta = r.model.config.rope_theta }); - try r.attention.encode(e, &r.scratch, &r.bindings[layer], plan, f.get(.query), f.get(.rotated_key), f.get(.value), f.get(.mixer_out), &.{cache}, f.get(.q_gate), false); + try r.attention.encode(e, &r.scratch, &r.bindings[layer], plan, f.get(.query), f.get(.rotated_key), f.get(.value), f.get(.mixer_out), &.{cache}, f.get(.q_gate), false, false); } fn tap(ptr: *anyopaque, e: mtl.ComputeEncoder, layer: usize, hidden: Ref, rows: u32) !void { const h: *Hooks = @ptrCast(@alignCast(ptr)); diff --git a/zig/tests/qwen27/synthetic_graph.zig b/zig/tests/qwen27/synthetic_graph.zig index 3da574b7d1..f0e707634f 100644 --- a/zig/tests/qwen27/synthetic_graph.zig +++ b/zig/tests/qwen27/synthetic_graph.zig @@ -95,7 +95,8 @@ pub fn createLayers(a: std.mem.Allocator, layers: usize) !*q.model.Model { } var widths = try q.widths.Widths.uniformAll(a, c.layers); defer widths.deinit(); - var weights = try q.gpu_weights.load(a, device, &cp, c, &widths); + const shard = q.config.Shard{ .rank = 0, .ranks = 1, .global = c, .local = c }; + var weights = try q.gpu_weights.load(a, device, &cp, &shard, &widths); errdefer weights.deinit(); const glue = try q.glue_gpu.Kernels.init(device); errdefer glue.deinit(); @@ -105,6 +106,6 @@ pub fn createLayers(a: std.mem.Allocator, layers: usize) !*q.model.Model { f.deinit(); } const m = try a.create(q.model.Model); - m.* = .{ .allocator = a, .device = device, .queue = queue, .config = c, .checkpoint = cp, .weights = weights, .kernels = .{ .allocator = a, .device = device }, .glue = glue, .frame = frame }; + m.* = .{ .allocator = a, .device = device, .queue = queue, .shard = shard, .config = c, .checkpoint = cp, .weights = weights, .kernels = .{ .allocator = a, .device = device }, .glue = glue, .frame = frame }; return m; } diff --git a/zig/tests/qwen27/widths.zig b/zig/tests/qwen27/widths.zig index 9c8ff9b6e2..9e3ebf784e 100644 --- a/zig/tests/qwen27/widths.zig +++ b/zig/tests/qwen27/widths.zig @@ -306,7 +306,7 @@ const Gpu = struct { const cb = g.queue.commandBuffer(); const e = cb.compute(.serial); if (!std.mem.eql(u8, path(raw).name, "row")) { // simd_qmm_bits at every width, 4-bit too - const tiles = try row.simd_bits.call(g.bits, raw.bits, raw.n, raw.k, wide, raw.n, false); + const tiles = try row.simd_bits.call(g.bits, raw.bits, raw.n, raw.k, wide, raw.n, false, 0); row.simd_bits.encode(e, tiles, .{ .w = raw.w, .w_off = raw.w_off, .scales = raw.scales, .s_off = raw.s_off, .biases = raw.biases, .b_off = raw.b_off }, x, 0, y, 0); return commit(cb, e); }