diff --git a/build.zig b/build.zig index 7f0a9d2403..c426401fe9 100644 --- a/build.zig +++ b/build.zig @@ -155,7 +155,7 @@ pub fn build(b: *std.Build) void { const gemma4_host_mods = metalEngineModules(b, b.graph.host, .Debug, draft_ids); gemma4_host_mods.sources.addImport("nemotron_packed_metallib", embedded(b, b.addWriteFiles().add("gemma4-compile-empty.metallib", ""), "nemotron_packed.metallib")); @import("zig/build/gemma4_host.zig").compileHost(b, gemma4_host_mods.metal, gemma4_host_mods.engine, gemma4_host_mods.lanes, gemma4_host_mods.sources); - @import("zig/build/qwen27_host.zig").compileHost(b, gemma4_host_mods.metal, gemma4_host_mods.engine, gemma4_host_mods.lanes); + test_step.dependOn(@import("zig/build/qwen27_host.zig").compileHost(b, gemma4_host_mods.metal, gemma4_host_mods.engine, gemma4_host_mods.lanes)); } @import("zig/build/qwen27_host.zig").tests(b); const xgrammar = grammar_build.option(b); 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/build/qwen27_host.zig b/zig/build/qwen27_host.zig index 1709ab50b5..4424ee3808 100644 --- a/zig/build/qwen27_host.zig +++ b/zig/build/qwen27_host.zig @@ -46,9 +46,14 @@ pub fn tests(b: *std.Build) void { b.step("test-qwen27-equality-host", "Shared lanes drafted, resumed and concurrent equality on the fake target").dependOn(&equal.step); } -pub fn compileHost(b: *std.Build, metal: *std.Build.Module, engine: *std.Build.Module, lanes: *std.Build.Module) void { +pub fn compileHost(b: *std.Build, metal: *std.Build.Module, engine: *std.Build.Module, lanes: *std.Build.Module) *std.Build.Step { const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = b.graph.host, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} }); const host = b.createModule(.{ .root_source_file = b.path("zig/src/native/qwen27_served.zig"), .target = b.graph.host, .link_libc = true, .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "tensorfold", .module = engine }, .{ .name = "engine_api", .module = api } } }); + // pair refusals are decided before a model loads: the host's own tests run on CPU alone + const served = b.step("test-qwen27-served-host", "Rank 1's named refusal and the pair's --vision refusal, no shaders or GPU"); + const served_run = b.addRunArtifact(b.addTest(.{ .root_module = host })); + served_run.has_side_effects = true; // a rerun must re-execute the tests, not replay a cached pass + served.dependOn(&served_run.step); const boundary = b.addTest(.{ .root_module = b.createModule(.{ .root_source_file = b.path("zig/qwen27_bank_boundary_host.zig"), .target = b.graph.host, .optimize = .debug, .link_libc = true, .imports = &.{.{ .name = "tensorfold", .module = engine }} }), .filters = &.{"qwen27 bank boundary:"} }); b.step("test-qwen27-bank-boundary-host", "Runner pending-keep guards with no device or shader execution").dependOn(&b.addRunArtifact(boundary).step); const object = b.addObject(.{ .name = "qwen27-speed-compile", .root_module = b.createModule(.{ @@ -60,6 +65,7 @@ pub fn compileHost(b: *std.Build, metal: *std.Build.Module, engine: *std.Build.M }) }); _ = object.getEmittedBin(); b.step("check-qwen27-speed-host", "Compile the actual Qwen27 startup and encoding without shaders").dependOn(&object.step); + return served; } fn tryNames(b: *std.Build, names: *std.ArrayList([]const u8), text: []const u8) void { diff --git a/zig/cost_rule_cases.zig b/zig/cost_rule_cases.zig index 91cd4fb9d4..42b3f715ba 100644 --- a/zig/cost_rule_cases.zig +++ b/zig/cost_rule_cases.zig @@ -16,7 +16,9 @@ test "Nemotron initial and drift passes share interleaved head sampling; the 27B try std.testing.expect(std.mem.indexOf(u8, body, "if (b.head != null) &head[0] else null") != null); try std.testing.expect(std.mem.indexOf(u8, body, "drifted(&head, &r.head)") != null); try std.testing.expect(std.mem.indexOf(u8, body, "fastest(heads") == null); + const warm = @embedFile("src/native/qwen27_warm.zig"); // the window timing moved here from the lane backend + try std.testing.expect(std.mem.indexOf(u8, warm, "cost_rule.fastest(timer, i)") != null); + try std.testing.expect(std.mem.indexOf(u8, warm, "cost_rule.measure(") == null); const qwen = @embedFile("src/native/qwen27_lanes.zig"); - try std.testing.expect(std.mem.indexOf(u8, qwen, "cost_rule.fastest(timer, i)") != null); try std.testing.expect(std.mem.indexOf(u8, qwen, "cost_rule.measure(") == null); } 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/native/metal.zig b/zig/src/native/metal.zig index 1391c64f86..a604e9ff46 100644 --- a/zig/src/native/metal.zig +++ b/zig/src/native/metal.zig @@ -335,7 +335,7 @@ fn refusal(a: Allocator, o: api.Open, dense: bool) !?[]const u8 { if (o.prompt_int8) return prompt_int8_refusal; if (o.slide and !std.mem.eql(u8, t, "nemotron_h")) return try std.fmt.allocPrint(a, "--slide: {s} checkpoints do not learn on Macs yet", .{t}); const one_mac: ?[]const u8 = if (std.mem.eql(u8, t, "nemotron_h")) "Nemotron" else if (std.mem.eql(u8, t, "lfm2_moe")) "LFM2.5-24B-A2B" else if (isGemma4(t)) "Gemma 4" else if (std.mem.eql(u8, t, "step5_vl")) "Step-5" else if (std.mem.eql(u8, t, "spark2_5")) "Spark-X2.6" else if (!std.mem.eql(u8, t, "qwen3_5")) null else if (dense) "Qwen3.8-27B" else "Qwen3.5-2B"; - if (o.speed_up != null) if (one_mac) |name| return try std.fmt.allocPrint(a, "--speed-up joins two Macs for Flash Next, GLM-5.3-Flash and GLM-5.3 only: serve {s} without it", .{name}); + if (o.speed_up != null and !(std.mem.eql(u8, t, "qwen3_5") and dense)) if (one_mac) |name| return try std.fmt.allocPrint(a, "--speed-up joins two Macs for Qwen3.8-27B, Flash Next, GLM-5.3-Flash and GLM-5.3 only: serve {s} without it", .{name}); if (o.recall and std.mem.eql(u8, t, "qwen4_exp")) return recall_refusal; if (spill.refusal(o)) |why| return why; if (o.learn == null) return null; @@ -361,7 +361,7 @@ fn qwen27Engine(a: Allocator, io: std.Io, dir: []const u8) bool { return true; } -/// Qwen3.8-27B on the lane host: `--parallel` streams in shared rounds, DFlash2 drafts from --drafter for one of them. +/// 27B: --parallel shared rounds, DFlash2 from --drafter, --speed-up opens the two-Mac seam before the weights load. fn openQwen27(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { const context = o.context orelse 32768; if (context <= 0 or context > 262144) { @@ -375,7 +375,7 @@ fn openQwen27(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[ var why: []const u8 = ""; const drafter = if (o.drafts) o.drafter else null; const learn: ?@import("qwen27_disk.zig").Learn = if (o.learn_conversations) if (o.learn) |root| .{ .root = root, .dir = o.dir, .drafter = drafter, .cap = gibBytes(o.learn_gib), .floor = gibBytes(o.learn_min_free_gib), .daily = gibBytes(o.learn_daily_gib), .spill = o.spill, .slots = spill.slots(o) } else null else null; - const h = qwen27.open(gpa, io, o.dir, .{ .context = @intCast(context), .streams = o.lanes, .streams_fixed = o.lanes_fixed, .drafter = drafter, .drafter_name = o.drafter_name, .drafter_bits = o.drafter_bits, .cache_gib = o.prompt_cache_gib, .cache_over_cap = o.prompt_cache_over_cap, .learn = learn, .vision = o.vision }, a, &why) catch |err| { + const h = qwen27.open(gpa, io, o.dir, .{ .pair = o.speed_up, .context = @intCast(context), .streams = o.lanes, .streams_fixed = o.lanes_fixed, .drafter = drafter, .drafter_name = o.drafter_name, .drafter_bits = o.drafter_bits, .cache_gib = o.prompt_cache_gib, .cache_over_cap = o.prompt_cache_over_cap, .learn = learn, .vision = o.vision }, a, &why) catch |err| { problem.* = if (why.len > 0) why else try std.fmt.allocPrint(a, "the native Qwen3.8-27B engine cannot load {s} ({s})", .{ o.dir, @errorName(err) }); return null; }; @@ -553,13 +553,13 @@ test "--spill-gib with --checkpoint-slots opens the four families with a disk st try std.testing.expectStringStartsWith((try refusal(a, .{ .dir = "d", .model_type = "glm5_next", .keep = 2 }, false)) orelse "", "--checkpoint-slots"); } -test "--speed-up is refused by Nemotron, the 27B and the 2B before any load, and not by the two-Mac families" { +test "--speed-up is refused by Nemotron and the 2B before any load, and opens the pair seam on the 27B" { var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); const a = arena.allocator(); const two: api.Open = .{ .dir = "d", .model_type = "qwen3_5", .speed_up = "pair.json" }; const tail = " only: serve {s} without it"; - try std.testing.expectStringEndsWith((try refusal(a, two, true)) orelse "", try std.fmt.allocPrint(a, tail, .{"Qwen3.8-27B"})); + try std.testing.expect((try refusal(a, two, true)) == null); // the 27B pairs: the pair seam is its two-Mac transport try std.testing.expectStringEndsWith((try refusal(a, two, false)) orelse "", try std.fmt.allocPrint(a, tail, .{"Qwen3.5-2B"})); try std.testing.expectStringEndsWith((try refusal(a, .{ .dir = "d", .model_type = "nemotron_h", .speed_up = "pair.json" }, false)) orelse "", try std.fmt.allocPrint(a, tail, .{"Nemotron"})); for ([_][]const u8{ "qwen4_exp", "glm5_next", "glm_moe_dsa" }) |t| try std.testing.expect((try refusal(a, .{ .dir = "d", .model_type = t, .speed_up = "pair.json" }, false)) == null); diff --git a/zig/src/native/qwen27_lanes.zig b/zig/src/native/qwen27_lanes.zig index ed44a0c58e..97e3893380 100644 --- a/zig/src/native/qwen27_lanes.zig +++ b/zig/src/native/qwen27_lanes.zig @@ -16,11 +16,21 @@ const Ring = q.tap_ring.Ring; const Context = q.dflash.runtime_model.Context; const Landing = @import("qwen27_landing.zig").Landing; const contract = tf.tree_round; +const pair_mod = @import("qwen27_pair.zig"); + +/// The pair transport's types and seam ABI, in their own unit: the lane backend and the link's adapter share them. +pub const Pair = pair_mod.Pair; +pub const Kind = pair_mod.Kind; +pub const Link = pair_mod.Link; +pub const batch_rows = pair_mod.batch_rows; +pub const pairReduce = pair_mod.pairReduce; +pub const pairBytes = pair_mod.pairBytes; +pub const pairChunk = pair_mod.chunk; +pub const seamBytes = pair_mod.seamBytes; /// A window's rows at most: the pending row and a DFlash2 tree. pub const max_rows = 16; -/// A shared round's rows at most: the frame's. -pub const batch_rows = 128; + /// Streams a runner holds at most (the attention plan's). pub const max_streams = q.attention_plan.max_streams; const ring = 64; @@ -73,6 +83,7 @@ pub const Metal = struct { fills: std.AutoHashMapUnmanaged(*lanes.Stream, slots.Fill) = .empty, // prompt passes the lane host runs by chunks image_fill: ?*const lanes.Stream = null, // the pass whose images the vision rows hold, until it closes promote_due: bool = false, // slot 0 freed: a stream moves in at the next backend call (slots.promoteDue) + link: ?*Link = null, // the pair: rank 0 sends each round here before it runs, rank 1 replays it pub fn init(gpa: std.mem.Allocator, runner: *q.decode_round.Runner, draft: ?*Draft) !*Metal { if (runner.slots == 0 or runner.slots > max_streams or runner.model.frame.capacity < batch_rows) return error.TargetBinding; @@ -122,7 +133,7 @@ pub const Metal = struct { const ringChunk = slots.ringChunk; const unbind = slots.unbind; const ready = slots.ready; - const refill = slots.refill; + pub const refill = slots.refill; pub fn backend(b: *Metal) be.Backend { return .{ .ptr = b, .vtable = &.{ .prefill = prefillFn, .prefill_step = slots.prefillStep, .cancel_prefill = slots.cancelPrefill, .first = firstFn, .first_masked = firstMaskedFn, .queue = queueFn, .read = readFn, .verify = verifyFn, .keep = keepFn, .draft = draft_fns.draftFn, .probabilities = draft_fns.probabilitiesFn, .tree = draft_fns.treeFn, .release = releaseFn, .reports_fill = true } }; @@ -161,66 +172,9 @@ pub const Metal = struct { return slot < b.draft_slots and live <= b.draft_slots and x.caches.len > 0; } - /// Each window width's ms (verify, draws, keep) on a scratch context in slot 0; resets it. Kept per build. + /// Each window width's ms on a scratch slot 0, kept per build: warm.measure, for the A/B tool. pub fn measure(b: *Metal, io: std.Io) !void { - const r = b.runner; - if (b.owners[0] != null or r.active != null) return error.StreamBusy; - var shape: [128]u8 = undefined; - const parts = [_][]const u8{ "qwen27-lanes", std.mem.span(r.model.device.name()), try q.widths.costShape(r.model.config, r.model.widths_tag, &shape, max_rows, b.draft != null) }; - const k = lanes.cost_cache.key(b.gpa, io, &parts) catch null; - if (k) |key| if (lanes.cost_cache.load([max_rows]lanes.config.Cost, b.gpa, io, key)) |kept| { - b.costs = kept; - b.timed = max_rows; - return; - }; - if (r.capacity <= max_rows) return error.ContextFull; - defer { - if (b.draft) |d| d.reset() catch {} else r.reset(0) catch {}; - } - var ids: [192]u32 = undefined; - for (&ids, 0..) |*t, i| t.* = @intCast((i * 7919 + 13) % r.model.config.vocab); - const prompt = ids[0..@min(ids.len, r.capacity - max_rows)]; // a small --context still holds a window past it - try b.refill(prompt); - var s = try lanes.Stream.init(b.gpa, .{ .id = "timing", .prompt = prompt, .max_new = 1 }); - defer s.deinit(b.gpa); - const timer: WindowTimer = .{ .b = b, .s = &s, .prompt = prompt }; - for (0..5) |_| _ = try timer.time(max_rows - 1); // the widest window first, while the GPU's clocks ramp up - var ms: [max_rows]f64 = undefined; - for (&ms, 0..) |*m, i| m.* = try lanes.cost_rule.fastest(timer, i); - try lanes.cost_rule.smooth(timer, &ms); - const ref = lanes.cost_cache.referenceKey(&parts); - if (lanes.cost_cache.load([max_rows]f64, b.gpa, io, ref)) |reference| if (lanes.cost_rule.drifted(&ms, &reference)) try lanes.cost_rule.again(timer, &ms); - for (&b.costs, ms, 0..) |*cost, m, i| cost.* = .{ .width = @intCast(i + 1), .ms = m }; - b.timed = max_rows; - if (k) |key| lanes.cost_cache.save([max_rows]lanes.config.Cost, b.gpa, io, key, b.costs); - lanes.cost_cache.save([max_rows]f64, b.gpa, io, ref, ms); - } - - /// Window entry `i` is `i + 1` rows, timed once; each keeps a row, so a full slot is refilled first. - const WindowTimer = struct { - b: *Metal, - s: *lanes.Stream, - prompt: []const u32, - pub fn time(t: WindowTimer, i: usize) !f64 { - if (t.b.runner.offsets[0] + i + 1 > t.b.runner.capacity) try t.b.refill(t.prompt); - return t.b.timeWindow(t.s, i + 1); - } - }; - - fn timeWindow(b: *Metal, s: *lanes.Stream, rows: usize) !f64 { - const t0 = mtl.clock.seconds(); - const positions: [max_rows]u64 = @splat(0); - const drafts: [max_rows]u32 = @splat(1); - var sampled: [max_rows]u32 = undefined; - var echoed: [max_rows]u32 = undefined; - b.owners[0] = s; - defer b.owners[0] = null; - const w = be.Window{ .stream = s, .pending = 0, .held = 0, .tokens = drafts[0 .. rows - 1], .parents = null, .positions = positions[0..rows] }; - var out = [_]be.Verified{.{ .sampled = sampled[0..rows], .drafts = echoed[0 .. rows - 1] }}; - try verifyFn(b, &.{w}, &out); - var path = [_]u32{0}; - try b.keepPaths(&.{&path}); - return (mtl.clock.seconds() - t0) * 1e3; + return @import("qwen27_warm.zig").measure(b, io); } pub fn self(ptr: *anyopaque) *Metal { @@ -268,6 +222,8 @@ pub const Metal = struct { /// Logits rows [first, first + rows) drawn at their positions: argmax on the GPU, or the stream's keyed sampler. fn draw(b: *Metal, s: *const lanes.Stream, first: usize, rows: usize, positions: []const u64, out: []u32, masks: []const u32) !void { + // a pair's token comes from the round's own head seam: a second local argmax would decide differently on rank 1 + if (b.runner.reduce != null) return error.PairGreedyOnly; const logits = b.runner.model.frame.get(.logits); const vocab: u32 = @intCast(b.runner.model.config.vocab); const width = (@as(usize, vocab) + 31) / 32; @@ -306,6 +262,16 @@ pub const Metal = struct { /// The tokens argmaxRows picked for rows [first, first + rows). fn picked(b: *const Metal, first: usize, rows: usize, out: []u32) !void { + if (b.runner.reduce != null) { + const seam_picks = b.runner.picks orelse return error.HeadSeamMissing; + const vocab: u32 = @intCast(b.runner.model.shard.global.vocab); + const ids: [*]const u32 = @ptrCast(@alignCast(seam_picks.buffer.contents() + seam_picks.offset)); + for (out[0..rows], ids[first..][0..rows]) |*t, id| { + if (id >= vocab) return error.NonfiniteLogits; + t.* = id; + } + return; + } const picks: [*]const contract.Pick = @ptrCast(@alignCast(b.picks.contents())); for (out[0..rows], picks[first..][0..rows]) |*t, p| { if (p.token == contract.invalid_token or p.nonfinite != 0) return error.NonfiniteLogits; @@ -314,7 +280,7 @@ pub const Metal = struct { } /// Keep the last verify's rows on `paths`: slot 0 alone keeps on the GPU with its taps, else the runner and rings. - fn keepPaths(b: *Metal, paths: []const []const u32) !void { + pub fn keepPaths(b: *Metal, paths: []const []const u32) !void { defer b.dropRound(); if (paths.len != b.count) return error.NothingToKeep; if (b.advanced) { // the verify kept each window's one row at its position; its logits already sit in row 0 @@ -417,8 +383,22 @@ pub const Metal = struct { // -- the vtable --------------------------------------------------------------------------------------------- + /// Rank 1 replays rank 0's rounds: refuse at the head of `prefill`, its first backend call (GLM's `beginPrompt`). + pub fn refuseFollower(m: *Metal) !void { + if (m.link) |l| if (l.pair.rank != 0) return error.FollowsPeer; + } + + /// A pair's token is the head seam's pick, so a sampled or grammar-constrained stream has no pair answer. + pub fn refusePairSampling(m: *Metal, s: *const lanes.Stream) !void { + // refuse it at `prefill`'s head too, before either rank settles, drops or binds a round (no second argmax). + if (m.link == null) return; + if (s.sampling != null or s.grammar != null) return error.PairGreedyOnly; + } + fn prefillFn(ptr: *anyopaque, s: *lanes.Stream) anyerror!void { const b = self(ptr); + try refusePairSampling(b, s); // a sampled or grammar-constrained row has no pair answer: refuse before any backend call + try refuseFollower(b); // before settle, dropRound, promoteDue and bind: rank 1 touched nothing yet b.settle() catch {}; b.dropRound(); slots.promoteDue(b); @@ -433,6 +413,11 @@ pub const Metal = struct { const b = self(ptr); _ = try b.slotFor(s); var token: [1]u32 = undefined; + if (b.runner.reduce != null) { + // the prompt's last chunk (head = .last) picked on both ranks: no second round, no unannounced exchange + try b.picked(0, 1, &token); + return b.push(token[0]); + } try b.draw(s, 0, 1, &.{position}, &token, &.{}); return b.push(token[0]); } @@ -467,7 +452,7 @@ pub const Metal = struct { } /// Every stream's window (pending token, then its drafts as a chain or tree) in one forward, every row drawn. - fn verifyFn(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { + pub fn verifyFn(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { const b = self(ptr); if (windows.len == 0 or windows.len > b.runner.slots or windows.len != out.len) return error.BadWindows; try b.settle(); @@ -503,10 +488,16 @@ pub const Metal = struct { if (windows.len > 1) b.dropTree(); const pool = mtl.objc.Pool.push(); defer pool.pop(); - b.round = try Round.init(b.gpa, inputs[0..windows.len], @intCast(r.model.config.conv_kernel), r.slots, @intCast(r.model.config.vocab)); + // Ids are global on every rank (only the head's rows are split), so the round validates the whole vocabulary. + b.round = try Round.init(b.gpa, inputs[0..windows.len], @intCast(r.model.config.conv_kernel), r.slots, @intCast(r.model.shard.global.vocab)); b.count = windows.len; errdefer b.dropRound(); const before = r.gpu_seconds; + if (b.runner.reduce != null) { + // landing 1 is one greedy stream of one-row windows: the words are the slot and the pending token. + if (windows.len != 1 or total != 1) return error.CommandOutOfStep; + if (b.link) |l| try l.step(inputs[0].slot, b.ids[0]); + } if (plain) try r.advance(&b.round.?, .all) else try r.verifyHead(&b.round.?, .all); for (windows) |w| slots.charge(b, w.stream, before); // a shared round counts whole in each reply b.advanced = plain; @@ -514,7 +505,7 @@ pub const Metal = struct { // every greedy window's rows from one argmax over the round; a sampled window draws on the host var greedy = false; for (windows) |w| greedy = greedy or (w.stream.sampling == null and w.masks.len == 0); - if (greedy) try b.argmaxRows(0, total); + if (greedy and b.runner.reduce == null) try b.argmaxRows(0, total); for (windows, out, b.round.?.firsts) |w, o, first| { const rows = w.rows(); if (w.stream.sampling != null or w.masks.len > 0) try b.draw(w.stream, first, rows, w.positions, o.sampled, w.masks) else try b.picked(first, rows, o.sampled); @@ -583,8 +574,10 @@ pub const Metal = struct { } fn releaseFn(ptr: *anyopaque, s: *lanes.Stream) void { - slots.endFill(self(ptr), s); - self(ptr).unbind(s); + const b = self(ptr); + slots.endFill(b, s); + if (b.link) |l| if (b.slotOf(s)) |slot| l.end(slot); + b.unbind(s); } }; diff --git a/zig/src/native/qwen27_link.zig b/zig/src/native/qwen27_link.zig new file mode 100644 index 0000000000..82afc00bd1 --- /dev/null +++ b/zig/src/native/qwen27_link.zig @@ -0,0 +1,76 @@ +//! The pair link's liveness: rank 1 replays rank 0's rounds, and both ranks stop when the link dies. +const std = @import("std"); +const platform = @import("sys"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const q = tf.qwen27; +const lb = @import("qwen27_lanes.zig"); + +/// Rank 1's follower: replay rank 0's rounds in order; a rank that cannot replay stops to be restarted. +pub fn follow(b: *lb.Metal) void { + const l = b.link orelse return; + while (true) { + const seen = l.pair.next(l.gpa, &l.words) catch |err| { + if (err != error.PairStopped) failFollow("a round", @errorName(err)); + return; + }; + const k = seen orelse return; + const kind = std.enums.fromInt(lb.Kind, k) orelse failFollow("an unknown command", k); + const w = l.words.items; + switch (kind) { + .begin => { + if (w.len != 1) failFollow("begin", w.len); + b.dropRound(); + b.runner.reset(w[0]) catch |err| return failFollow("begin", @errorName(err)); + }, + .chunk => { + if (w.len < 4 or w.len != 4 + w[1]) failFollow("chunk", w.len); + const session = q.session.Session{ .runner = b.runner, .slot = w[0] }; + const chunk = w[4..]; + if (w[3] != 0) { + session.decodeChunk(chunk, w[2] != 0) catch |err| return failFollow("decode chunk", @errorName(err)); + } else { + session.promptChunk(chunk, w[2] != 0) catch |err| return failFollow("prompt chunk", @errorName(err)); + } + }, + .step => { + if (w.len != 2) failFollow("step", w.len); + (q.session.Session{ .runner = b.runner, .slot = w[0] }).step(w[1]) catch |err| return failFollow("step", @errorName(err)); + }, + .end => {}, + } + } +} + +/// This rank can no longer hold rank 0's state: it stops rather than serve on a stale half, to be restarted. +fn failFollow(what: []const u8, detail: anytype) noreturn { + std.log.err("qwen27 rank 1: cannot replay rank 0's {s} ({any}): stopping so both ranks are restarted", .{ what, detail }); + std.process.exit(1); +} + +/// Both ranks stop when the link dies: a dead peer must not leave HTTP up on a half, so a supervisor restarts both. +pub fn watch(p: *lb.Pair, quit: *const std.atomic.Value(bool)) void { + var seen = p.heartbeats(); + var since = mtl.clock.seconds(); + while (!quit.load(.acquire)) { + if (p.failed()) down(p, "the pair link is down"); + p.beat(); + const beats = p.heartbeats(); + const now = mtl.clock.seconds(); + if (beats != seen) { + seen = beats; + since = now; + } else if (now - since > link_silence_s) down(p, "the peer stopped signalling"); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = @intCast(250 * std.time.ns_per_ms) }; + _ = platform.c.nanosleep(&ts, null); + } +} + +/// Beats the watchdog may miss before the peer counts as gone: twelve 250 ms intervals, longer than any seam. +const link_silence_s: f64 = 3.0; + +/// The link is not a serving state: name the half that is gone, then end the process the supervisor watches. +fn down(p: *lb.Pair, what: []const u8) noreturn { + std.log.err("qwen27 rank {d}: {s}: exiting so both ranks are restarted", .{ p.rank, what }); + std.process.exit(1); +} diff --git a/zig/src/native/qwen27_pair.zig b/zig/src/native/qwen27_pair.zig new file mode 100644 index 0000000000..dcd3fd4a12 --- /dev/null +++ b/zig/src/native/qwen27_pair.zig @@ -0,0 +1,113 @@ +//! The pair transport's types, the seam ABI the engine calls, and the window one exchange may carry. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const q = tf.qwen27; + +/// The pair's transport type: the row-split seam the engine exports, so this file names no other family's type. +pub const Pair = tf.pair_rows.Pair; + +/// The command set rank 1 replays: the four GPU operations a pair round is made of (glm_dsa's `follow.zig` shape). +pub const Kind = enum(u32) { begin = 1, chunk, step, end }; + +/// Rank 0 puts each round on the wire before running it; rank 1 replays them in order here. Owned by the host. +pub const Link = struct { + pair: *Pair, + gpa: std.mem.Allocator, + words: std.ArrayList(u32) = .empty, + + pub fn deinit(l: *Link) void { + l.words.deinit(l.gpa); + } + + fn put(l: *Link, kind: Kind) !void { + try l.pair.send(@backingInt(kind), l.words.items); + l.words.clearRetainingCapacity(); + } + + /// A prompt's chunk boundaries are rank 0's memory policy's, so the send carries them: rank 1 cannot infer them. + pub fn begin(l: *Link, slot: u32) !void { + l.words.clearRetainingCapacity(); + try l.words.append(l.gpa, slot); + try l.put(.begin); + } + + pub fn chunk(l: *Link, slot: u32, ids: []const u32, last: bool, decoded: bool) !void { + l.words.clearRetainingCapacity(); + try l.words.appendSlice(l.gpa, &.{ slot, @intCast(ids.len), @intFromBool(last), @intFromBool(decoded) }); + try l.words.appendSlice(l.gpa, ids); + try l.put(.chunk); + } + + pub fn step(l: *Link, slot: u32, token: u32) !void { + l.words.clearRetainingCapacity(); + try l.words.appendSlice(l.gpa, &.{ slot, token }); + try l.put(.step); + } + + pub fn end(l: *Link, slot: u32) void { + l.words.clearRetainingCapacity(); + l.words.append(l.gpa, slot) catch return; + l.put(.end) catch {}; + } +}; + +/// The pair's seam ends in the engine's shape: the transport's `Ref` is the frame's, so the casts are field copies. +pub fn pairReduce(p: *Pair) q.forward.Reduce { + return .{ .ptr = p, .part = pairPart, .stage = pairStage, .gather = pairGather, .failed = pairFailed, .argmax = pairArgmax }; +} + +fn pairPart(ptr: *anyopaque) q.forward.Seam { + const p: *Pair = @ptrCast(@alignCast(ptr)); + const s = p.part(); + return .{ .x = s.x, .ref = .{ .buffer = s.ref.buf, .offset = s.ref.off } }; +} + +fn pairStage(ptr: *anyopaque, e: *mtl.ComputeEncoder, x: u32, src: q.forward.Ref, rows: u32, width: u32, stride: u32) void { + const p: *Pair = @ptrCast(@alignCast(ptr)); + p.stage(e, x, .{ .buf = src.buffer, .off = src.offset }, rows, width, stride); +} + +fn pairGather(ptr: *anyopaque, e: *mtl.ComputeEncoder, x: u32, mine: q.forward.Ref, out: q.forward.Ref, rows: u32, width: u32, out_stride: u32) void { + const p: *Pair = @ptrCast(@alignCast(ptr)); + p.gather(e, x, .{ .buf = mine.buffer, .off = mine.offset }, .{ .buf = out.buffer, .off = out.offset }, rows, width, out_stride); +} + +fn pairFailed(ptr: *anyopaque) bool { + const p: *Pair = @ptrCast(@alignCast(ptr)); + return p.failed(); +} + +fn pairArgmax(ptr: *anyopaque, e: *mtl.ComputeEncoder, x: u32, logits: q.forward.Ref, stride: u32, lo: u32, n: u32, rows: u32, picks: q.forward.Ref) void { + const p: *Pair = @ptrCast(@alignCast(ptr)); + p.argmaxBf16(e, x, .{ .buf = logits.buffer, .off = logits.offset }, stride, lo, n, rows, .{ .buf = picks.buffer, .off = picks.offset }); +} + +/// The window an exchange may carry: a round's rows of the widest seam payload -- this rank's `intermediate / 2` bf16. +pub fn pairBytes() usize { + return seamBytes(batch_rows); +} + +/// The bytes a round of `rows` stages at the widest seam: this rank's `intermediate / 2` bf16. +pub fn seamBytes(rows: u32) usize { + return @as(usize, rows) * @divExact(17408, 2) * @sizeOf(u16); +} + +/// A pair's prompt pass at most: a chunk of the window's rows, so its widest seam fits `pairBytes()`. +pub fn chunk(requested: u32) u32 { + return @min(requested, batch_rows); +} + +/// A shared round's rows at most: the frame's. +pub const batch_rows = 128; + +test "the served pair's chunk times the widest seam fits the window" { + // a served pass asks for 512 rows here, 2048 on generation 17 (`model.promptRows`); the pair caps it at the window + const requested = q.model.promptRows(0); + const served = chunk(requested); + try std.testing.expect(served <= batch_rows); + try std.testing.expect(seamBytes(served) <= pairBytes()); + // the uncapped pass stages a seam wider than the window, so the cap is what keeps the exchange legal + try std.testing.expectEqual(@as(usize, batch_rows * @divExact(17408, 2) * @sizeOf(u16)), pairBytes()); + try std.testing.expect(seamBytes(512) > pairBytes()); +} diff --git a/zig/src/native/qwen27_pair_tests.zig b/zig/src/native/qwen27_pair_tests.zig new file mode 100644 index 0000000000..0aad7b4ff4 --- /dev/null +++ b/zig/src/native/qwen27_pair_tests.zig @@ -0,0 +1,62 @@ +//! The pair's refusals and its seam-sized chunk: the words, the guards, and the window the served pass may ask for. +const std = @import("std"); +const tf = @import("tensorfold"); +const q = tf.qwen27; +const lanes = tf.lanes; +const lb = @import("qwen27_lanes.zig"); +const served = @import("qwen27_served.zig"); + +test "rank 1 refuses a request by name before its backend moves round state" { + try std.testing.expectEqualStrings("speed-up mode: this Mac runs rank 0's requests; send requests to rank 0", served.words(null, error.FollowsPeer).?); + // `link` is the only `Metal` field the guard reads: the rest staying undefined proves no other step ran + var pair: lb.Pair = undefined; + pair.rank = 1; + var l: lb.Link = undefined; + l.pair = &pair; + var m: lb.Metal = undefined; + m.link = &l; + try std.testing.expectError(error.FollowsPeer, lb.Metal.refuseFollower(&m)); + pair.rank = 0; + try lb.Metal.refuseFollower(&m); // rank 0's own requests run on it + m.link = null; + try lb.Metal.refuseFollower(&m); // no pair: nothing to refuse for +} + +test "a pair refuses a sampled or grammar request by name at the head of prefill" { + try std.testing.expectEqualStrings("speed-up mode: a pair answers greedy only; send the request with sampling and grammar off", served.words(null, error.PairGreedyOnly).?); + // `link` is the only `Metal` field and the stream's two: the rest staying undefined proves no other step ran + var pair: lb.Pair = undefined; + pair.rank = 0; + var l: lb.Link = undefined; + l.pair = &pair; + var m: lb.Metal = undefined; + m.link = &l; + var s: lanes.Stream = undefined; + s.sampling = null; + s.grammar = null; + try lb.Metal.refusePairSampling(&m, &s); // a greedy stream runs on a pair + s.sampling = std.mem.zeroes(lanes.Sampling); + try std.testing.expectError(error.PairGreedyOnly, lb.Metal.refusePairSampling(&m, &s)); + s.sampling = null; + var grammar: lanes.grammar.Constraint = undefined; + s.grammar = &grammar; + try std.testing.expectError(error.PairGreedyOnly, lb.Metal.refusePairSampling(&m, &s)); + m.link = null; + try lb.Metal.refusePairSampling(&m, &s); // a single Mac samples and constrains freely +} + +test "the served pair's chunk times the widest seam fits the window" { + // the served pass asks for the generation's rows (`model.promptRows(generation)`); a pair caps them at the window + const chunk = lb.pairChunk(q.model.promptRows(0)); + try std.testing.expect(chunk <= lb.batch_rows); + try std.testing.expect(lb.seamBytes(chunk) <= lb.pairBytes()); + // the uncapped pass would stage a wider seam than the window, so the cap is what keeps the exchange legal + try std.testing.expect(lb.seamBytes(512) > lb.pairBytes()); +} + +test "--vision is refused on a pair, by name" { + var why: []const u8 = ""; + try std.testing.expectError(error.PairNoVision, served.pairFlags(.{ .context = 4096, .streams = 1, .vision = .{} }, false, &why)); + try std.testing.expect(std.mem.startsWith(u8, why, "--vision does not run on a pair")); + try served.pairFlags(.{ .context = 4096, .streams = 1 }, false, &why); // a single Mac reads images +} diff --git a/zig/src/native/qwen27_served.zig b/zig/src/native/qwen27_served.zig index 8e5bb8e37a..cd8f1391d8 100644 --- a/zig/src/native/qwen27_served.zig +++ b/zig/src/native/qwen27_served.zig @@ -12,6 +12,8 @@ const cache_fit = @import("cache_fit.zig"); const share_fit = @import("qwen27_share_fit.zig"); const Resident = @import("resident.zig").Resident; const lb = @import("qwen27_lanes.zig"); +const link = @import("qwen27_link.zig"); +const warm = @import("qwen27_warm.zig"); pub const lanes_backend = lb; const disk = @import("qwen27_disk.zig"); const images = @import("qwen27_images.zig"); @@ -29,6 +31,7 @@ pub const Options = struct { learn: ?disk.Learn = null, // --learn-conversations: kept states written behind to disk and read back vision: ?api.VisionOpen = null, // --vision: open loads the image tower, and requests may hold images images: ?*images.Vision = null, // the loaded tower attach serves (and the host then owns) + pair: ?[]const u8 = null, // --speed-up: this rank's settings file; `Pair.init` reads the rank out of it }; pub const Host = struct { @@ -49,6 +52,13 @@ pub const Host = struct { disk_stage: ?mtl.Buffer = null, // copies out of private kept states for the disk (qwen27_disk.zig) vision: ?*images.Vision = null, // --vision's tower and image rows spares: q.snapshot.Spares = .{}, // dropped kept states' buffers, for the next saves + pair: ?*lb.Pair = null, // --speed-up: the seam this rank owns; rank 1 follows through it + link: ?*lb.Link = null, // rank 0's round sender and rank 1's follower target; `Metal.link` points here + hw: ?mtl.Device = null, // the pair's device, opened before the model's so the seam exists for the load + picks: ?mtl.Buffer = null, // a round's head-seam picks (one u32 a row), read by the lane backend + follower: ?std.Thread = null, // rank 1's command loop + quit: std.atomic.Value(bool) = .init(false), // `close` sets it so the link watchdog never reads a freed pair + watcher: ?std.Thread = null, // the link watchdog both ranks run while the pair is up pub fn engine(h: *Host) api.Engine { return h.host.engine(); @@ -108,9 +118,13 @@ fn readMemory(_: ?*anyopaque, reset_peak: bool) ?api.Memory { } /// The words for a request the engine refuses. -fn words(_: ?*anyopaque, err: anyerror) ?[]const u8 { +pub fn words(_: ?*anyopaque, err: anyerror) ?[]const u8 { return switch (err) { error.PromptTooLong => "the prompt and max_tokens do not fit this server's --context", + // rank 1 replays rank 0's rounds and owns no request (qwen27_lanes.prefillFn refuses first) + error.FollowsPeer => "speed-up mode: this Mac runs rank 0's requests; send requests to rank 0", + // a pair's token is the head seam's pick: a sampled or grammar-constrained row is refused at prefill + error.PairGreedyOnly => "speed-up mode: a pair answers greedy only; send the request with sampling and grammar off", else => null, }; } @@ -271,19 +285,37 @@ fn laterBytes(h: *const Host) u64 { pub fn open(gpa: Allocator, io: std.Io, dir: []const u8, o: Options, a: Allocator, why: *[]const u8) !*Host { const pool = mtl.objc.Pool.push(); defer pool.pop(); - const model = try q.model.Model.loadDeferred(gpa, io, dir, 128); + // glm_dsa's order: the seam exists before the weights load and ready() comes once every GPU allocation is in. + var hw: ?mtl.Device = null; + var pair: ?*lb.Pair = null; + if (o.pair) |path| { + hw = try mtl.Device.init(); + pair = lb.Pair.init(gpa, io, hw.?, path, lb.pairBytes()) catch |err| { + why.* = try std.fmt.allocPrint(a, "the speed-up seam {s} would not open ({s})", .{ path, @errorName(err) }); + return err; + }; + } + errdefer if (pair) |p| p.deinit(); + errdefer if (hw) |d| d.deinit(); + const model = if (pair) |p| + try q.model.Model.loadShardDeferred(gpa, io, dir, 128, p.rank, 2, .full) + else + try q.model.Model.loadDeferred(gpa, io, dir, 128); errdefer model.deinit(); var with = o; // the tower before the slots: their room and the cache's budget count it as taken if (o.vision) |v| with.images = try images.open(gpa, io, model, dir, v, a, why); errdefer if (with.images) |v| v.deinit(); - const h = try attach(gpa, io, model, null, with, a, why); + const h = try attach(gpa, io, model, null, with, a, why, pair); h.owns_model = true; + h.pair = pair; + h.hw = hw; return h; } /// A loaded model served as `open` serves one, `drafter` bound in place of `o.drafter`; both stay the caller's. -pub fn attach(gpa: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q.dflash.runtime_model.Model, o: Options, a: Allocator, why: *[]const u8) !*Host { +pub fn attach(gpa: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q.dflash.runtime_model.Model, o: Options, a: Allocator, why: *[]const u8, pair: ?*lb.Pair) !*Host { + if (pair != null) try pairFlags(o, drafter != null, why); const pool = mtl.objc.Pool.push(); defer pool.pop(); const capacity = if (o.context == 0) @min(model.config.max_position, 32768) else o.context; @@ -320,10 +352,33 @@ pub fn attach(gpa: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q.df h.cache = null; h.resident = null; h.disk_stage = null; + // create() leaves unnamed Host fields undefined: a debug-filled optional is non-null, so name h.link first. + h.link = null; + h.hw = null; + h.picks = null; + h.follower = null; + h.quit = .init(false); + h.watcher = null; h.vision = o.images; // closed with the host; a failed attach leaves it to the caller h.trace = null; h.runner = try q.decode_round.Runner.initReady(gpa, model, slots, capacity + lb.max_rows, 1); errdefer h.runner.deinit(); + h.pair = pair; + if (pair) |p| { + if (slots != 1) { + why.* = try std.fmt.allocPrint(a, "--speed-up serves one greedy stream: pass --parallel 1", .{}); + return error.PairOneStream; + } + // The picks live where the head seam writes them, so the lane backend reads them rather than a local argmax. + const picks = try model.device.buffer(lb.batch_rows * @sizeOf(u32), mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + errdefer picks.deinit(); + h.picks = picks; + h.link = try gpa.create(lb.Link); + errdefer gpa.destroy(h.link.?); + h.link.?.* = .{ .pair = p, .gpa = gpa }; + h.runner.reduce = lb.pairReduce(p); + h.runner.picks = .{ .buffer = picks, .offset = 0 }; + } if (profileEvery()) |every| { h.trace = .{ .queue = model.queue, .report_every = every }; // enabled once measure() has priced the windows h.runner.profile(&h.trace.?); @@ -333,10 +388,13 @@ pub fn attach(gpa: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q.df errdefer if (h.draft) |d| d.deinit(); h.back = try lb.Metal.init(gpa, &h.runner, h.draft); errdefer h.back.deinit(); + h.back.link = h.link; h.back.chunk = h.runner.promptRows(@min(rows, row_policy.requested(model.device.generation(), value))); + // a pair stages each chunk's widest seam, so its pass may carry only the window's rows (qwen27_pair.chunk) + if (pair != null) h.back.chunk = lb.pairChunk(@intCast(h.back.chunk)); h.back.grow = .{ .ctx = h, .room = Host.roomFor, .added = Host.added }; h.back.vision = o.images; - if (h.draft != null) try h.back.measure(io); + if (h.draft != null) try warm.measure(h.back, io); h.cfg = try lanes.Config.init(gpa, h.back.facts(), lb.max_rows, lb.max_rows - 1); errdefer h.cfg.deinit(gpa); h.clock = .{ .io = io }; @@ -364,13 +422,36 @@ pub fn attach(gpa: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q.df }; if (model.ane) |share| share.footprint_growth = platform.memory_policy.growth(before, cache_fit.footprint()); } - try enableCache(h, o.cache_gib, o.cache_over_cap, a, why); + // a pair's prompt states are two Macs' and its slots are one: no reuse in landing 1, and no disk learning + if (pair == null) try enableCache(h, o.cache_gib, o.cache_over_cap, a, why); errdefer if (h.cache) |*store| store.deinit(); - if (o.learn) |l| try disk.attach(h, io, a, l); + if (pair == null) if (o.learn) |l| try disk.attach(h, io, a, l); errdefer if (h.cache) |*store| if (store.disk) |r| r.deinit(); h.host.keepalive_target = .{ .ctx = &h.warm, .tick = mtl.keepalive.Target.tick }; if (h.trace) |*t| t.enabled = true; // serving rounds count from here: measure() timed whole, unsplit buffers + // both Macs are loaded and allocated before either serves: rank 1 may now replay rank 0's rounds + if (pair) |p| p.ready() catch |err| { + why.* = if (err == error.PairPeerSilent) + "the peer never signalled ready in five minutes: start rank 1 first, with the same build on both Macs and each rank's own settings file" + else + try std.fmt.allocPrint(a, "the pair link was not ready ({s}): restart both ranks", .{@errorName(err)}); + return err; + }; + // a dead link is not a serving state: exit so a supervisor restarts both ranks, not serve seam errors + if (pair) |p| { + // fail closed: a rank whose watchdog is missing would keep serving HTTP on a half once the peer died + h.watcher = std.Thread.spawn(.{}, link.watch, .{ p, &h.quit }) catch |err| { + why.* = try std.fmt.allocPrint(a, "the pair link's watchdog would not start ({s}): a dead peer would leave this rank up, so this rank stops", .{@errorName(err)}); + return error.PairNoWatchdog; + }; + } try h.host.start(); + if (pair) |p| if (p.rank == 1) { + h.follower = std.Thread.spawn(.{}, link.follow, .{h.back}) catch |err| { + why.* = try std.fmt.allocPrint(a, "rank 1's follower would not start ({s})", .{@errorName(err)}); + return err; + }; + }; std.log.info("Qwen3.8-27B: {d} streams of {d} tokens{s}{s}", .{ slots, capacity, if (h.draft != null) ", DFlash2 drafts for one of them" else "", if (o.images != null) ", images read" else "" }); return h; } @@ -384,6 +465,26 @@ fn profileEvery() ?u64 { }; } +/// The flags a pair refuses, named: quietly serving without one of them would misreport what the reply ran on. +pub fn pairFlags(o: Options, drafter: bool, why: *[]const u8) !void { + if (drafter) { + why.* = "--drafter does not run on a pair: rank 1 has no drafter checkpoint, so both ranks decode plain"; + return error.PairNoDrafter; + } + if (o.cache_gib != null) { + why.* = "--prompt-cache-gib does not run on a pair: each Mac keeps its own halves of the slots' states"; + return error.PairNoCache; + } + if (o.learn != null) { + why.* = "--learn does not run on a pair: each Mac learns its own halves of the slots' states"; + return error.PairNoLearn; + } + if (o.vision != null or o.images != null) { + why.* = "--vision does not run on a pair: rank 1 replays rank 0's token ids, so an image prompt would run different rows on the two ranks"; + return error.PairNoVision; + } +} + /// Reserve later buffers and the loaded share's priced bytes that its construction left outside the footprint. fn laterFootprint(h: *const Host) !u64 { const tower: u64 = if (h.back.vision) |v| v.firstUse() else 0; @@ -415,7 +516,18 @@ const holdResident = q.residency.hold; pub fn close(ptr: *anyopaque) void { const h = Host.of(ptr); + h.quit.store(true, .release); // the watchdog exits before the pair it reads is torn down h.host.stop(); + // the pair goes after the host stops (no round is mid-flight) and before the model frees the device + if (h.pair) |p| p.stopFollowing(); + if (h.follower) |t| t.join(); + if (h.watcher) |t| t.join(); + if (h.link) |l| { + l.deinit(); + h.gpa.destroy(l); + } + if (h.pair) |p| p.deinit(); + if (h.hw) |d| d.deinit(); const rows = if (h.cache) |*store| store.disk else null; if (h.cache) |*store| store.deinit(); // its states leave the set as they go h.spares.deinit(); @@ -428,6 +540,7 @@ pub fn close(ptr: *anyopaque) void { if (h.draft) |d| d.deinit(); if (h.trace) |*t| t.reportRest(); h.runner.deinit(); + if (h.picks) |b| b.deinit(); if (h.vision) |v| v.deinit(); if (h.owns_model) h.model.deinit(); h.gpa.destroy(h); @@ -435,6 +548,7 @@ pub fn close(ptr: *anyopaque) void { test { _ = @import("qwen27_share_fit.zig"); + _ = @import("qwen27_pair_tests.zig"); } test "Qwen27 interior state uses a guaranteed cold request cut" { diff --git a/zig/src/native/qwen27_slots.zig b/zig/src/native/qwen27_slots.zig index b69664dc88..1eaa14a2d7 100644 --- a/zig/src/native/qwen27_slots.zig +++ b/zig/src/native/qwen27_slots.zig @@ -28,6 +28,8 @@ pub fn opening(b: *Metal, s: *lanes.Stream, slot: u32) !Fill { if (ids.len == 0 or ids.len + s.max_new + max_rows > b.runner.capacity) return error.PromptTooLong; if (s.yields()) return error.Yielded; // a background pass steps aside before it touches the slot if (b.vision != null and s.images.len > 0 and b.image_fill != null and b.image_fill != s) return error.LaneMemory; + // Announced before the first chunk: rank 1 opens where rank 0 does (glm_dsa order). + if (b.link) |l| try l.begin(slot); try recover(b, s); b.runner.rope[slot] = try images.layout(b.vision, s); const d = b.drafter(slot); @@ -123,6 +125,8 @@ pub fn chunk(b: *Metal, s: *lanes.Stream, f: *Fill) !bool { const part = ids[at..end]; const last = end == ids.len; const before = b.runner.gpu_seconds; + // Rank 0 sends the chunk before running it here (glm_dsa order): the send carries this pass's boundary. + if (b.link) |l| try l.chunk(slot, part, last, decoded); if (d) |x| { if (decoded) try x.decodeChunk(part, last) else try x.promptChunk(part, last); } else { diff --git a/zig/src/native/qwen27_warm.zig b/zig/src/native/qwen27_warm.zig new file mode 100644 index 0000000000..085edeb4a8 --- /dev/null +++ b/zig/src/native/qwen27_warm.zig @@ -0,0 +1,73 @@ +//! The lane backend's window timing: each width's ms on a scratch slot 0, kept in the per-build cost cache. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const q = tf.qwen27; +const lanes = tf.lanes; +const be = lanes.backend; + +const nl = @import("qwen27_lanes.zig"); // the native backend, not the core lane engine `lanes` names +const Metal = nl.Metal; +const max_rows = nl.max_rows; + +/// Each window width's ms (verify, draws, keep) on a scratch context in slot 0; resets it. Kept per build. +pub fn measure(b: *Metal, io: std.Io) !void { + const r = b.runner; + if (b.owners[0] != null or r.active != null) return error.StreamBusy; + var shape: [128]u8 = undefined; + const parts = [_][]const u8{ "qwen27-lanes", std.mem.span(r.model.device.name()), try q.widths.costShape(r.model.config, r.model.widths_tag, &shape, max_rows, b.draft != null) }; + const k = lanes.cost_cache.key(b.gpa, io, &parts) catch null; + if (k) |key| if (lanes.cost_cache.load([max_rows]lanes.config.Cost, b.gpa, io, key)) |kept| { + b.costs = kept; + b.timed = max_rows; + return; + }; + if (r.capacity <= max_rows) return error.ContextFull; + defer { + if (b.draft) |d| d.reset() catch {} else r.reset(0) catch {}; + } + var ids: [192]u32 = undefined; + for (&ids, 0..) |*t, i| t.* = @intCast((i * 7919 + 13) % r.model.config.vocab); + const prompt = ids[0..@min(ids.len, r.capacity - max_rows)]; // a small --context still holds a window past it + try b.refill(prompt); + var s = try lanes.Stream.init(b.gpa, .{ .id = "timing", .prompt = prompt, .max_new = 1 }); + defer s.deinit(b.gpa); + const timer: WindowTimer = .{ .b = b, .s = &s, .prompt = prompt }; + for (0..5) |_| _ = try timer.time(max_rows - 1); // the widest window first, while the GPU's clocks ramp up + var ms: [max_rows]f64 = undefined; + for (&ms, 0..) |*m, i| m.* = try lanes.cost_rule.fastest(timer, i); + try lanes.cost_rule.smooth(timer, &ms); + const ref = lanes.cost_cache.referenceKey(&parts); + if (lanes.cost_cache.load([max_rows]f64, b.gpa, io, ref)) |reference| if (lanes.cost_rule.drifted(&ms, &reference)) try lanes.cost_rule.again(timer, &ms); + for (&b.costs, ms, 0..) |*cost, m, i| cost.* = .{ .width = @intCast(i + 1), .ms = m }; + b.timed = max_rows; + if (k) |key| lanes.cost_cache.save([max_rows]lanes.config.Cost, b.gpa, io, key, b.costs); + lanes.cost_cache.save([max_rows]f64, b.gpa, io, ref, ms); +} + +/// Window entry `i` is `i + 1` rows, timed once; each keeps a row, so a full slot is refilled first. +const WindowTimer = struct { + b: *Metal, + s: *lanes.Stream, + prompt: []const u32, + pub fn time(t: WindowTimer, i: usize) !f64 { + if (t.b.runner.offsets[0] + i + 1 > t.b.runner.capacity) try t.b.refill(t.prompt); + return timeWindow(t.b, t.s, i + 1); + } +}; + +pub fn timeWindow(b: *Metal, s: *lanes.Stream, rows: usize) !f64 { + const t0 = mtl.clock.seconds(); + const positions: [max_rows]u64 = @splat(0); + const drafts: [max_rows]u32 = @splat(1); + var sampled: [max_rows]u32 = undefined; + var echoed: [max_rows]u32 = undefined; + b.owners[0] = s; + defer b.owners[0] = null; + const w = be.Window{ .stream = s, .pending = 0, .held = 0, .tokens = drafts[0 .. rows - 1], .parents = null, .positions = positions[0..rows] }; + var out = [_]be.Verified{.{ .sampled = sampled[0..rows], .drafts = echoed[0 .. rows - 1] }}; + try nl.Metal.verifyFn(b, &.{w}, &out); + var path = [_]u32{0}; + try b.keepPaths(&.{&path}); + return (mtl.clock.seconds() - t0) * 1e3; +} 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/dflash_served.zig b/zig/tests/qwen27/dflash_served.zig index c813567db6..f472c3389e 100644 --- a/zig/tests/qwen27/dflash_served.zig +++ b/zig/tests/qwen27/dflash_served.zig @@ -128,7 +128,7 @@ pub fn main(init: std.process.Init) !void { } var why: []const u8 = ""; // a context the window timing outgrows: a new build's first start refills the timing's slot - const served = host.attach(a, init.io, model, draft, .{ .context = 256, .streams = 1, .cache_gib = 0 }, a, &why) catch |err| { + const served = host.attach(a, init.io, model, draft, .{ .context = 256, .streams = 1, .cache_gib = 0 }, a, &why, null) catch |err| { std.debug.print("served host: {s}\n", .{why}); return err; }; 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/served.zig b/zig/tests/qwen27/served.zig index 858dabf5a3..d3714aebe7 100644 --- a/zig/tests/qwen27/served.zig +++ b/zig/tests/qwen27/served.zig @@ -58,7 +58,7 @@ pub fn main(init: std.process.Init) !void { } var box = Box{ .a = a, .io = init.io }; var why: []const u8 = ""; - const served = host.attach(a, init.io, model, null, .{ .context = 256, .streams = 1, .cache_gib = 0 }, a, &why) catch |err| { + const served = host.attach(a, init.io, model, null, .{ .context = 256, .streams = 1, .cache_gib = 0 }, a, &why, null) catch |err| { std.debug.print("served host: {s}\n", .{why}); return err; }; 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); }