From dcccfe7bf13a6acc380537dfe7bddeef308bf71b Mon Sep 17 00:00:00 2001 From: Andrey Kolesnikov <15747200+akol1@users.noreply.github.com> Date: Tue, 27 Oct 2026 18:47:48 -0700 Subject: [PATCH 1/3] core: add the row-split pair seam and split simd_qmm over K MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Step 1 of 6 of the two-Mac 27B series; stacked on main (1b9a987e) alone, and steps 2-6 stack on this commit. The seam exchanges whole output rows instead of K halves: rank r owns 2,560 of the 5,120 output columns (8,704 of 17,408 for the down projection), reduces over the full K on its own rank, and `pair_gather` copies whole rows in bf16 without doing any arithmetic, so no exchange rounds or sums a partial product across ranks. A broken link surfaces as bf16 NaN `0x7FC0`. The seam's four kernels (`pair_stage`, `pair_gather`, `pair_argmax`, `pair_pick`) move bf16 frames and compare values only, so they carry no sums and no accumulator's error. One Mac fails a row holding any NaN or ±Inf, so `pair_argmax` reports such a row as NaN bits with bit 31 set in its id word, and `pair_pick` turns that flag into the id sentinel: both Macs then refuse a non-finite round where one Mac would, through the same downstream check as a broken link. Five planted controls (a stage on the peer's half, a gather writing its own half twice, an argmax and a pick tie to the higher id, an argmax that does not flag) each go red on the gate named in `docs/kernel-contract.md`; the next commit adds the sixth, for `pair_pick`'s flag rule itself. That transport is a new module, `pair_rows.zig` and `pair_rows.metal`, over `pair_rows_layout.zig`'s slots and the relay protocol unchanged. `pair.metal` and `pair.zig` stay for GLM-5.3-Flash, whose transport this series does not touch: a shipped path cannot change without an A/B against main. `simd_qmm` gains a K-chunk width: `chunks(n, whole)` lets a rank that still holds the whole output width keep the split one Mac derives from that width, instead of one derived from the rank's narrower slice. --- docs/kernel-contract.md | 19 +- zig/kernels/metal/core/pair_rows.metal | 129 +++++ zig/kernels/metal/sources.zig | 2 + zig/src/core/contract.zig | 1 + zig/src/core/contract/metal_pair_records.zig | 25 + zig/src/core/contract/registry.zig | 2 +- zig/src/core/pair_rows.zig | 463 +++++++++++++++++ zig/src/core/pair_rows_layout.zig | 125 +++++ zig/src/core/pair_rows_loopback.zig | 249 +++++++++ zig/src/core/simd_qmm.zig | 19 +- .../qwen27/dflash/runtime_callbacks.zig | 2 +- zig/src/families/qwen27/quant_gpu.zig | 2 +- zig/src/metal/command.zig | 10 + zig/src/tensorfold.zig | 2 + zig/tests/affine_ref.zig | 2 +- zig/tests/contract/metal_affine.zig | 2 +- zig/tests/contract/metal_limits.zig | 2 + zig/tests/contract/metal_main.zig | 28 + zig/tests/contract/metal_pair.zig | 485 ++++++++++++++++++ zig/tests/contract/metal_pair_controls.zig | 55 ++ 20 files changed, 1608 insertions(+), 16 deletions(-) create mode 100644 zig/kernels/metal/core/pair_rows.metal create mode 100644 zig/src/core/contract/metal_pair_records.zig create mode 100644 zig/src/core/pair_rows.zig create mode 100644 zig/src/core/pair_rows_layout.zig create mode 100644 zig/src/core/pair_rows_loopback.zig create mode 100644 zig/tests/contract/metal_pair.zig create mode 100644 zig/tests/contract/metal_pair_controls.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 1ed43ac146..42ba822a3f 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -20,7 +20,7 @@ The test of a kernel fills a `contract.spec.Spec` and hands the harness a few ho and the magnitude of the terms each output sums, and can replay the sums in another accumulator. - Its served envelope: the shapes the engine calls it at, and for each dimension its tile, the step a valid size moves by and the least valid size. -- Its accumulator (`f32` for every kernel here) and the depth of its longest chain of roundings at a shape, plus any +- Its accumulator (`f32` where a kernel sums) and the depth of its longest chain of roundings at a shape, plus any fp32 operations after the sums (`ops`) and intermediate roundings to the output type (`mids`). - Whether it may write in place. If it replaces another kernel, the kernel whose bits it must equal. - Its NaN rule: `propagate` (as float64 arithmetic on its formula does), or a stated rule the test checks itself @@ -28,7 +28,8 @@ The test of a kernel fills a `contract.spec.Spec` and hands the harness a few ho - For gate 5, its launch in `zig/src/core/contract/registry.zig`: threads a group and threadgroup bytes (the source's, and any its encoder sets) on Metal, or symbol, image, threads a block, dynamic shared bytes and cluster width on CUDA, with the limits its compiles recorded. HIP launches are in `zig/src/core/contract/hip_records.zig`: symbol, code - object, threads a workgroup and dynamic LDS bytes. + object, threads a workgroup and dynamic LDS bytes. The seam's records are in `metal_pair_records.zig`, which + `registry.zig` concatenates. ## The gates @@ -99,7 +100,8 @@ The program exits 1 on any failed gate, any control that stayed green, or any li the inputs for a shape and a gate's fill (usually through an input family), `run` launches once and waits, `count` and `reference` give the outputs and their float64 values, `simulate` replays the sums for gate 4. 2. Add it to the runner's list, with its served shapes and tiles. -3. Add its launches to `registry.zig`, and fill the records from `--offline` (Metal) or `kernel-contract-limits` +3. Add its launches to `registry.zig`, or to a records file it concatenates as the seam's are, and fill them from + `--offline` (Metal) or `kernel-contract-limits` (CUDA: nvcc 13.0's in `cuda`, and each later release's set where that release compiles it differently); a HIP kernel's go to `hip_records.zig`, filled from `hip-kernel-contract-limits`. 4. Post the `SUMMARY` lines from each backend you ran in the pull request, and say which chips you could not run. @@ -127,8 +129,8 @@ float64 output class, without an inherited exception. ## Coverage -Metal: softmax, RMS norm, bf16 top-k, the row projection (3, 4, 6 and 8 bits), affine_mm (4 and 8 bits), simd_qmm -(4 bits) and simd_qmm_bits (3, 6 and 8 bits), and the lane projection (6-bit groups of 32 on every chip; 4- and 8-bit +Metal: softmax, RMS norm, bf16 top-k, the row projection (3, 4, 6 and 8 bits), affine_mm (4 and 8 bits), simd_qmm (4 +bits) and simd_qmm_bits (3, 6 and 8 bits), and the lane projection (6-bit groups of 32 on every chip; 4- and 8-bit formats on the M5's tensor units, where its reg kernel must equal the cooperative one bit for bit), and tile_qmm's prompt tiles and their decode twin (4, 6 and 8 bits; the twin bit-equal to the tiles). Gate 5 also holds the rest of tile_qmm's pipelines and GLM-5.3-Flash's hc_expand_10. @@ -141,6 +143,13 @@ Metal also checks Nemotron's routed experts on the engine's route tables, one ro against float64 at tf_load16's bf16 x sums (up's relu2 as an interval over its two bf16 roundings), and the rows2, rows4 and six geometry pipelines of each projection bit-equal to them (two controls: member rows read from one row, and w3's x sums without their bf16 steps). +The row-split seam's four kernels (`pair_rows.metal`, launched by `pair_rows.zig`) move bf16 frames, never sums: +`pair_stage` copies this rank's half of a frame row into the slot the relay sends, and `pair_gather` writes both halves +of every row, its own natively and the peer's from the relayed frame. `pair_argmax` flags a row whose logits are not +finite (its id word carries bit 31, its value word a NaN) and takes the lower id on a tie; `pair_pick` turns a flagged +rank's word into the id sentinel, so both Macs refuse a non-finite round where one Mac does. Five controls edit the +source: a stage reading the peer's half, a gather writing its own half twice, and an argmax or pick tie to the higher +id must go red at `shapes`, and an argmax that does not flag at `nonfinite`. CUDA: the RMS norm rows (16-byte form and the fp32-out tree), the bf16 residual's norm tails (MLX's rounded rows and the once-rounded ones, each checked on its residual and its output, with four controls each dropping a rounding point), Nemotron's router logits (and their 128-row tiles, bit-equal to them) and top-k, the latent top-k's decode form (bit-equal to `tf_super_topk`), the latent combine and diff --git a/zig/kernels/metal/core/pair_rows.metal b/zig/kernels/metal/core/pair_rows.metal new file mode 100644 index 0000000000..a505ba2cd8 --- /dev/null +++ b/zig/kernels/metal/core/pair_rows.metal @@ -0,0 +1,129 @@ +// Two Macs' tensor-parallel exchanges: the host relays parts through shared events, and these kernels move bf16 bits. +#include +using namespace metal; +// A rank compares one row's logits across ranks; no kernel sums across them. + +// Sync page words (shared with pair_rows.zig): FAILED = the host gave up on a seam; OBSERVED = this side's refusals. +constant constexpr uint FAILED = 16, OBSERVED = 20; +// An argmax id word's bit 31: the row had a logit that is not finite, so pair_pick sends the row to the id sentinel. +constant constexpr uint PAIR_FLAG = 0x80000000u; +constant constexpr uint PAIR_SENTINEL = 0xFFFFFFFFu; + +// Kernels' argument block, byte-identical to pair_rows.zig's PairRows: width, the frame row's stride, this rank, one. +struct PairRows { + uint width, stride, rank, x; +}; + +// bf16 bits as the exact fp32 value they name (argmax reads a bf16 logits stream this way). +inline float pair_bf16(uint bits) { + return as_type(bits << 16); +} + +// True once the host gave up: FAILED is set before the wait releases, so a later kernel poisons its output rows. +inline bool pair_link_broken(device atomic_uint* sync, uint t) { + if (atomic_load_explicit(&sync[FAILED], memory_order_relaxed) == 0u) return false; + if (t == 0) atomic_fetch_add_explicit(&sync[OBSERVED], 1u, memory_order_relaxed); + return true; +} + +// Row r: this rank's own half of the frame row into the exchange slot; a bf16 copy, so it carries the frame's bits. +kernel void pair_stage(device atomic_uint* sync [[buffer(0)]], constant PairRows& a [[buffer(1)]], + const device ushort* src [[buffer(2)]], device ushort* dst [[buffer(3)]], + uint r [[threadgroup_position_in_grid]], uint t [[thread_index_in_threadgroup]]) { + const ulong in = ulong(r) * a.stride; + const ulong out = ulong(r) * a.width; + if (pair_link_broken(sync, t)) { + for (uint i = t; i < a.width; i += 256) dst[out + i] = 0x7FC0u; // bf16 NaN: the link stays down + return; + } + for (uint i = t; i < a.width; i += 256) dst[out + i] = src[in + i]; +} + +// Row r: this rank's half at its own offset in the frame row and the peer's at the other; a bf16 copy, no arithmetic. +kernel void pair_gather(device atomic_uint* sync [[buffer(0)]], constant PairRows& a [[buffer(1)]], + const device ushort* mine [[buffer(2)]], device atomic_uint* peer [[buffer(3)]], + device ushort* out [[buffer(4)]], uint r [[threadgroup_position_in_grid]], + uint t [[thread_index_in_threadgroup]]) { + const ulong mrow = ulong(r) * a.width; + const ulong orow = ulong(r) * a.stride; + const ulong own = a.rank == 0 ? orow : orow + a.width; + const ulong got = a.rank == 0 ? orow + a.width : orow; + if (pair_link_broken(sync, t)) { + for (uint i = t; i < a.width; i += 256) { + out[own + i] = 0x7FC0u; + out[got + i] = 0x7FC0u; + } + return; + } + for (uint i = t; i < a.width; i += 256) { + // The relay wrote the peer's half: an untracked frame read, so one atomic word (two bf16) at a time. + const uint word = atomic_load_explicit(&peer[(mrow + i) >> 1], memory_order_relaxed); + out[own + i] = mine[mrow + i]; + out[got + i] = ushort(word >> ((i & 1) * 16)); + } +} + +// Row r's argmax over this rank's logits (ids lo..lo+n): (fp32 value, id), lower id on tie; dims.w 0 = fp32, 1 = bf16. +kernel void pair_argmax(const device float* logits [[buffer(0)]], constant uint4& dims [[buffer(1)]], + device uint* out [[buffer(2)]], uint r [[threadgroup_position_in_grid]], + uint t [[thread_index_in_threadgroup]], uint sg [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + const ulong row = ulong(r) * dims.x; + const device float* xf = logits + row; + const device ushort* xb = reinterpret_cast(logits) + row; + const bool wide = dims.w != 0u; + float best = -INFINITY; + uint at = PAIR_SENTINEL; + bool bad = false; + for (uint i = t; i < dims.z; i += 1024) { + const float v = wide ? pair_bf16(uint(xb[i])) : xf[i]; + if (!isfinite(v)) { bad = true; continue; } // the flagged id is the argmax over finite logits + if (v > best) best = v, at = dims.y + i; + } + for (ushort o = 16; o > 0; o >>= 1) { + const float ob = simd_shuffle_xor(best, o); + const uint oa = simd_shuffle_xor(at, o); + if (ob > best || (ob == best && oa < at)) best = ob, at = oa; + } + // The flag is the row's, not the thread's: one lane saw a non-finite logit, so every lane must know. + uint dirty = bad ? 1u : 0u; + for (ushort o = 16; o > 0; o >>= 1) dirty |= simd_shuffle_xor(dirty, o); + threadgroup float vb[32]; + threadgroup uint va[32]; + threadgroup uint vd[32]; + if (lane == 0) vb[sg] = best, va[sg] = at, vd[sg] = dirty; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg != 0) return; + best = vb[lane]; + at = va[lane]; + dirty = vd[lane]; + for (ushort o = 16; o > 0; o >>= 1) { + const float ob = simd_shuffle_xor(best, o); + const uint oa = simd_shuffle_xor(at, o); + if (ob > best || (ob == best && oa < at)) best = ob, at = oa; + dirty |= simd_shuffle_xor(dirty, o); + } + if (lane == 0) { + out[2 * r] = dirty != 0u ? as_type(NAN) : as_type(best); + out[2 * r + 1] = dirty != 0u ? (at | PAIR_FLAG) : at; + } +} + +// Each row's pick over both Macs' argmaxes: a flagged row draws no token; else the larger value, lower id on a tie. +kernel void pair_pick(device atomic_uint* sync [[buffer(0)]], constant uint2& rx [[buffer(1)]], + const device uint* mine [[buffer(2)]], device atomic_uint* peer [[buffer(3)]], + device uint* picks [[buffer(4)]], uint t [[thread_index_in_threadgroup]]) { + if (pair_link_broken(sync, t)) { + if (t == 0) for (uint r = 0; r < rx.x; ++r) picks[r] = 0xFFFFFFFFu; // no such token: visibly wrong, never a token + return; + } + if (t != 0) return; + for (uint r = 0; r < rx.x; ++r) { + const float vm = as_type(mine[2 * r]); + const uint am = mine[2 * r + 1]; + const float vp = as_type(atomic_load_explicit(&peer[2 * r], memory_order_relaxed)); + const uint ap = atomic_load_explicit(&peer[2 * r + 1], memory_order_relaxed); + // A non-finite row on either Mac draws no token: the same sentinel a broken link reads, never an id. + picks[r] = ((am | ap) & PAIR_FLAG) != 0u ? PAIR_SENTINEL : (vp > vm || (vp == vm && ap < am)) ? ap : am; + } +} diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index 2506fb4905..22cb484ef0 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -34,6 +34,8 @@ pub const core_exlr8_dense = @embedFile("core/exlr8_dense.metal"); pub const core_exlr8_prompt = @embedFile("core/exlr8_prompt.metal"); pub const core_exlr8_dense_prompt = @embedFile("core/exlr8_dense_prompt.metal"); pub const core_pair = @embedFile("core/pair.metal"); +/// The row-split seam's kernels (see `pair_rows.zig`): `pair.metal` above stays as it is for glm_dsa. +pub const core_pair_rows = @embedFile("core/pair_rows.metal"); /// GLM-5.3 (glm_moe_dsa): MLA with rope over a bf16 latent cache, the DSA indexer and its top-k, fp32 glue. pub const glm_dsa_mla = @embedFile("glm_dsa/mla.metal"); diff --git a/zig/src/core/contract.zig b/zig/src/core/contract.zig index 2c976103b2..a04a28ec3f 100644 --- a/zig/src/core/contract.zig +++ b/zig/src/core/contract.zig @@ -21,6 +21,7 @@ test { _ = @import("contract/glm_rows_records.zig"); _ = @import("contract/nemotron_experts_records.zig"); _ = @import("contract/moss_image_records.zig"); + _ = @import("contract/metal_pair_records.zig"); _ = spec; _ = gates; _ = rows; diff --git a/zig/src/core/contract/metal_pair_records.zig b/zig/src/core/contract/metal_pair_records.zig new file mode 100644 index 0000000000..5e9eebe17d --- /dev/null +++ b/zig/src/core/contract/metal_pair_records.zig @@ -0,0 +1,25 @@ +//! The pair seam's four launches (core/pair_rows.metal) with the M1/M2 limits the offline compile measures. +const std = @import("std"); +const limits = @import("limits.zig"); +const registry = @import("registry.zig"); + +pub const records = [_]limits.MetalLaunch{ + .{ .name = "metal pair stage pair_stage", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "metal pair gather pair_gather", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "metal pair argmax pair_argmax", .threads = 1024, .tg_memory = 384, .m1 = 1024, .m2 = 1024 }, + .{ .name = "metal pair pick pair_pick", .threads = 32, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, +}; + +test "the pair seam's four launches are in the live and offline collector and fit every family" { + for (records) |record| { + var found: usize = 0; + for (registry.all_metal) |active| if (std.mem.eql(u8, active.name, record.name)) { + try std.testing.expectEqualDeep(record, active); + found += 1; + }; + try std.testing.expectEqual(@as(usize, 1), found); + try std.testing.expectEqual(@as(usize, 0), try registry.checkMetal(std.testing.allocator, &.{record}, true)); + // A zero here would be a launch no M1/M2 compiler was asked about. + try std.testing.expect(record.m1 > 0 and record.m2 > 0); + } +} diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 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..1856304072 100644 --- a/zig/src/core/simd_qmm.zig +++ b/zig/src/core/simd_qmm.zig @@ -8,6 +8,11 @@ pub fn splits(n: usize) u32 { return if (n <= 64) 32 else if (n <= 6144) 16 else 8; } +/// Chunks for `n`; a split rank passes the whole width as `whole` so it keeps that width's tree (0 keeps `n`). +pub fn chunks(n: usize, whole: usize) u32 { + return splits(if (whole != 0) whole else n); +} + /// Output tiles a simdgroup within the 16 KB reduction buffer; tiles never change a row's bits. pub fn tiles(n: usize, rt: u32, s: u32) u32 { var nt: u32 = if (n % 32 == 0) 4 else if (n % 16 == 0) 2 else 1; @@ -72,10 +77,10 @@ pub const Pipelines = struct { pub const Call = struct { pipeline: mtl.Pipeline, args: Args, groups: mtl.Size, threads: usize }; /// y = x W^T, 4-bit groups of 64: the scalar twin at 1-2 rows if `scalar`, chunk pairs if `paired`, else MMA tiles. -pub fn call(p: *const Pipelines, n: usize, k: usize, rows: usize, scalar: bool, paired: bool) !Call { +pub fn call(p: *const Pipelines, n: usize, k: usize, rows: usize, scalar: bool, paired: bool, whole: usize) !Call { if (n == 0 or n % 8 != 0 or n > std.math.maxInt(i32) or k == 0 or k % 64 != 0 or k > std.math.maxInt(i32)) return error.UnsupportedSimdShape; if (rows == 0 or rows > std.math.maxInt(i32)) return error.NoRows; - const s = splits(n); + const s = chunks(n, whole); const si = std.mem.indexOfScalar(u32, &s_values, s).?; const args = Args{ .k = @intCast(k), .n = @intCast(n), .r = @intCast(rows) }; if (scalar and scalarBlock(rows, s) != 0) { @@ -127,7 +132,7 @@ pub fn pairedExact(p: *const Pipelines, device: mtl.Device, m: Matrix) !bool { defer queue.deinit(); const cb = queue.commandBuffer(); const e = cb.compute(.serial); - for (outs, [_]bool{ false, true }) |o, pair| encode(e, try call(p, m.n, m.k, rows, false, pair), m, x, 0, o, 0); + for (outs, [_]bool{ false, true }) |o, pair| encode(e, try call(p, m.n, m.k, rows, false, pair, 0), m, x, 0, o, 0); e.end(); cb.commit(); cb.wait(); @@ -157,9 +162,9 @@ pub fn twinExact(p: *const Pipelines, device: mtl.Device, m: Matrix) !bool { defer queue.deinit(); const cb = queue.commandBuffer(); const e = cb.compute(.serial); - encode(e, try call(p, m.n, m.k, rows, false, false), m, x, 0, tiled, 0); - for (0..rows) |r| encode(e, try call(p, m.n, m.k, 1, true, false), m, x, r * m.k * 2, twin, r * m.n * 2); - for ([_]usize{ 0, rows - 2 }, 0..) |r, i| encode(e, try call(p, m.n, m.k, 2, true, false), m, x, r * m.k * 2, twos, i * 2 * m.n * 2); + encode(e, try call(p, m.n, m.k, rows, false, false, 0), m, x, 0, tiled, 0); + for (0..rows) |r| encode(e, try call(p, m.n, m.k, 1, true, false, 0), m, x, r * m.k * 2, twin, r * m.n * 2); + for ([_]usize{ 0, rows - 2 }, 0..) |r, i| encode(e, try call(p, m.n, m.k, 2, true, false, 0), m, x, r * m.k * 2, twos, i * 2 * m.n * 2); e.end(); cb.commit(); cb.wait(); @@ -177,6 +182,8 @@ test "chunks follow output width only and tiles fit the reduction buffer" { try std.testing.expectEqual(@as(u32, 32), splits(48)); try std.testing.expectEqual(@as(u32, 16), splits(5120)); try std.testing.expectEqual(@as(u32, 8), splits(34816)); + try std.testing.expectEqual(splits(6240), chunks(3168, 6240)); + try std.testing.expectEqual(splits(6240), chunks(6240, 0)); for (s_values) |s| for ([_]u32{ 1, 2 }) |rt| { const nt = tiles(34816, rt, s); try std.testing.expect(s * rt * nt * 64 * 4 <= 16384); diff --git a/zig/src/families/qwen27/dflash/runtime_callbacks.zig b/zig/src/families/qwen27/dflash/runtime_callbacks.zig index 46d252448c..bfea3d8552 100644 --- a/zig/src/families/qwen27/dflash/runtime_callbacks.zig +++ b/zig/src/families/qwen27/dflash/runtime_callbacks.zig @@ -29,7 +29,7 @@ fn project(ptr: *anyopaque, epoch: u64, x: op.View, w: ck.Tensor, mode: op.Mode) const input = try b.frame.ref(x); const output = try b.frame.ref(y); const s = @import("core").row_projection.simd; - s.encode(e, try s.call(&b.simd.?, t.n, t.k, x.rows, false, false), t.matrix(), input.buf, input.off, output.buf, output.off); + s.encode(e, try s.call(&b.simd.?, t.n, t.k, x.rows, false, false, 0), t.matrix(), input.buf, input.off, output.buf, output.off); } else { const adapter = b.prepared.get(@intFromPtr(w.bytes.ptr)) orelse return error.DraftPreparationUnqualified; const input = try b.frame.ref(x); diff --git a/zig/src/families/qwen27/quant_gpu.zig b/zig/src/families/qwen27/quant_gpu.zig index e38b913e2a..728cc3955f 100644 --- a/zig/src/families/qwen27/quant_gpu.zig +++ b/zig/src/families/qwen27/quant_gpu.zig @@ -139,7 +139,7 @@ pub const Kernels = struct { if (p.raw) |raw| if (raw.bits == 4 and raw.group == 64 and raw.sum == .f32 and raw.n % 8 == 0) { const m = row.simd.Matrix{ .w = raw.w, .w_off = raw.w_off, .scales = raw.scales, .s_off = raw.s_off, .biases = raw.biases, .b_off = raw.b_off, .n = raw.n, .k = raw.k }; const scalar = try k.twin(m); // loads the tiles before they're read - const call = try row.simd.call(&k.simd.?, raw.n, raw.k, rows, scalar, try k.pairTwin(m)); + const call = try row.simd.call(&k.simd.?, raw.n, raw.k, rows, scalar, try k.pairTwin(m), 0); row.simd.encode(e, call, m, x.buffer, x.offset, y.buffer, y.offset); e.barrier(); try span.end(); diff --git a/zig/src/metal/command.zig b/zig/src/metal/command.zig index 49118d0b3d..45c2b29358 100644 --- a/zig/src/metal/command.zig +++ b/zig/src/metal/command.zig @@ -195,4 +195,14 @@ pub const ComputeEncoder = struct { pub fn end(self: ComputeEncoder) void { objc.msg(void, self.id, "endEncoding", .{}); } + + /// The command buffer this encoder was opened from, for the event ops a seam encodes BETWEEN encoders. + pub fn commandBuffer(self: ComputeEncoder) CommandBuffer { + return .{ .id = objc.msg(Id, self.id, "commandBuffer", .{}) }; + } + + /// Whether this encoder dispatches serially or concurrently: a seam reopens it in the same mode. + pub fn dispatchType(self: ComputeEncoder) types.DispatchType { + return @fromBackingInt(@intCast(objc.msg(usize, self.id, "dispatchType", .{}))); + } }; diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 4fa4fbe147..4cc2deba22 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -32,6 +32,8 @@ pub const exlr8 = @import("core/exlr8.zig"); pub const exlr8_projection = @import("core/exlr8_projection.zig"); pub const exlr8_dense_gpu = @import("core/exlr8_dense_gpu.zig"); pub const exlr8_prompt = @import("core/exlr8_prompt.zig"); +/// The two-Mac pair seam's row split: the transport, its layout and the one-Mac Loopback the stage-2 checks use. +pub const pair_rows = @import("core/pair_rows.zig"); pub const glm_dsa = @import("families/glm_dsa/glm_dsa.zig"); pub const nemotron = @import("families/nemotron/nemotron.zig"); pub const flashnext_vision_contract = @import("families/flashnext/vision_contract.zig"); diff --git a/zig/tests/affine_ref.zig b/zig/tests/affine_ref.zig index b367fa120a..691ecf11ce 100644 --- a/zig/tests/affine_ref.zig +++ b/zig/tests/affine_ref.zig @@ -272,7 +272,7 @@ fn affine(t: *Tally, device: mtl.Device, queue: mtl.Queue, s: *const Synth, ref: /// simd_qmm over rows [start, start + m), the scalar twin when `scalar` and it takes m rows. fn simdWindow(e: mtl.ComputeEncoder, p: *const sq.Pipelines, s: *const Synth, y: Out, scalar: bool, start: usize, m: usize) !void { const mat = sq.Matrix{ .w = s.w, .scales = s.s, .biases = s.b, .n = s.n, .k = s.k }; - sq.encode(e, try sq.call(p, s.n, s.k, m, scalar, true), mat, s.x, start * s.k * 2, y.buf, start * s.n * 2); + sq.encode(e, try sq.call(p, s.n, s.k, m, scalar, true, 0), mat, s.x, start * s.k * 2, y.buf, start * s.n * 2); } /// simd_qmm (4-bit, group 64): tiles at every window against float64 and the wide call; the twin at 1 and 2 rows. diff --git a/zig/tests/contract/metal_affine.zig b/zig/tests/contract/metal_affine.zig index 46f921e9f8..3aa37d5de0 100644 --- a/zig/tests/contract/metal_affine.zig +++ b/zig/tests/contract/metal_affine.zig @@ -198,7 +198,7 @@ pub const Proj = struct { e.dispatchGroups(mtl.Size.of(cl.groups, cl.slots, 1), mtl.Size.of(cl.threads, 1, 1)); }, .simd => { - const cl = try sq.call(&p.sq_p.?, shape.n, shape.k, shape.rows, false, true); + const cl = try sq.call(&p.sq_p.?, shape.n, shape.k, shape.rows, false, true, 0); sq.encode(e, cl, .{ .w = p.w.buf, .w_off = p.w.off, .scales = p.s.buf, .s_off = p.s.off, .biases = p.b.buf, .b_off = p.b.off, .n = shape.n, .k = shape.k }, p.x.buf, p.x.off, p.out.buf, p.out.off); }, .simd_bits => { diff --git a/zig/tests/contract/metal_limits.zig b/zig/tests/contract/metal_limits.zig index 53609bffa0..c0e267c0d0 100644 --- a/zig/tests/contract/metal_limits.zig +++ b/zig/tests/contract/metal_limits.zig @@ -100,6 +100,8 @@ fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch) !La text = sources.core_deltanet; } else if (std.mem.startsWith(u8, d.name, "qwen27 ")) { text = if (has(d.name, " copy ")) tf.qwen27.runtime_sources.copy else tf.qwen27.runtime_sources.glue; + } else if (std.mem.startsWith(u8, d.name, "metal pair ")) { + text = sources.core_pair_rows; } else if (has(d.name, "tile_qmm")) text = sources.core_tile_qmm else if (has(d.name, "softmax")) text = sources.ops_softmax else if (has(d.name, "rms")) text = sources.ops_embed_norm else if (has(d.name, "top-k")) text = topk_text else if (has(d.name, "row projection")) text = sources.core_row_projection else if (has(d.name, "simd_qmm_bits")) text = sources.core_simd_qmm_bits else if (has(d.name, "simd_qmm")) text = sources.core_simd_qmm else if (has(d.name, "affine_mm")) { const bits: u8 = if (has(d.name, "q8")) 8 else 4; const group: u16 = if (has(d.name, "g32")) 32 else 64; diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 46801e171c..7fd93c2ff8 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -17,6 +17,8 @@ const drama_combine = @import("metal_drama_combine.zig"); const drama_cfg = @import("metal_drama_cfg.zig"); const glm_eh_k = @import("metal_glm_eh.zig"); const glm_rows_k = @import("metal_glm_rows.zig"); +const pair_k = @import("metal_pair.zig"); +const pair_controls = @import("metal_pair_controls.zig"); const flash_pool = @import("metal_flashnext_pool.zig"); const flash = @import("metal_flashnext.zig"); @@ -202,6 +204,32 @@ pub fn main(init: std.process.Init) !u8 { try rows_list.append(a, r); try kernels.append(a, r.kernel()); } + var pair_list: std.ArrayList(*pair_k.Pair) = .empty; + defer { + for (pair_list.items) |p| { + p.deinit(); + a.destroy(p); + } + pair_list.deinit(a); + } + for (std.enums.values(pair_controls.Mutation)) |m| { + for (pair_k.all) |w| { + const control_for = switch (m) { + .none => true, + .stage_peer_half => w == .stage, + .gather_own_halves => w == .gather, + .argmax_high_tie, .argmax_no_flag => w == .argmax, + .pick_high_tie => w == .pick, + }; + if (!control_for) continue; + if (!wanted(o, if (m == .none) pair_k.spec(w).name else m.name())) continue; + const p = try a.create(pair_k.Pair); + errdefer a.destroy(p); + p.* = try pair_k.Pair.init(a, device, queue, w, m); + try pair_list.append(a, p); + try kernels.append(a, p.kernel()); + } + } var flash_list: [8]flash.Probe = undefined; var flash_made: usize = 0; defer for (flash_list[0..flash_made]) |*probe| probe.deinit(); diff --git a/zig/tests/contract/metal_pair.zig b/zig/tests/contract/metal_pair.zig new file mode 100644 index 0000000000..9a15b020f8 --- /dev/null +++ b/zig/tests/contract/metal_pair.zig @@ -0,0 +1,485 @@ +//! The pair seam's four kernels (core/pair_rows.metal), and the float64 rules the gates check. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); +const rig_ = @import("rig.zig"); +const controls = @import("metal_pair_controls.zig"); + +const c = contract.spec; +const g = contract.gates; + +/// The seam's four kernels; the tail of each registry name is the MSL symbol launchOf hands the source. +pub const Which = enum { stage, gather, argmax, pick }; + +pub const all = [_]Which{ .stage, .gather, .argmax, .pick }; + +/// What pair_rows.zig dispatches per kernel, so `init` can compare it against Metal's own reflection. +pub const Launched = struct { threads: u32, tg_memory: u32 }; + +pub fn launch(w: Which) Launched { + return switch (w) { + .stage => .{ .threads = 256, .tg_memory = 0 }, + .gather => .{ .threads = 256, .tg_memory = 0 }, + .argmax => .{ .threads = 1024, .tg_memory = 384 }, + .pick => .{ .threads = 32, .tg_memory = 0 }, + }; +} + +/// The MSL symbol pair_rows.metal defines. +pub fn symbol(w: Which) []const u8 { + return switch (w) { + .stage => "pair_stage", + .gather => "pair_gather", + .argmax => "pair_argmax", + .pick => "pair_pick", + }; +} + +/// The spec's name: `metal `, so launchOf's last word is the symbol and its prefix picks core_pair_rows. +pub fn name(w: Which) []const u8 { + return switch (w) { + .stage => "metal pair stage pair_stage", + .gather => "metal pair gather pair_gather", + .argmax => "metal pair argmax pair_argmax", + .pick => "metal pair pick pair_pick", + }; +} + +/// Qwen3.8-27B's rank-local slice: the seam reduces ids [124160, 248320) on rank 1. +const logit_lo: u32 = 124160; + +/// An argmax id word's bit 31: the row had a logit that is not finite, so pair_pick sends the row to the id sentinel. +const flag: u32 = 0x80000000; + +/// The id a flagged row, and a broken link, read as: no token, visibly wrong. +const sentinel: u32 = 0xFFFFFFFF; + +/// A thread's chain over every 1024th logit, five shuffle levels, then the 32 simdgroup winners. +fn chainLogits(s: c.Shape) usize { + return (s.k + 1023) / 1024 + 5 + 5; +} + +/// Only moves and compares values: no sums to carry an accumulator's error. +fn noSums(_: c.Shape) usize { + return 0; +} + +const move_why = "a move writes each bf16 element as it read it: no sum, no rounding, so no accumulator's error can enter"; + +const argmax_rule = "a winner is (value bits, id): NaN loses every comparison, a row of only NaNs reports the -inf sentinel with id 0xFFFFFFFF, and a row with any non-finite logit reports NaN bits and id | 0x80000000"; +const pick_rule = "the pick copies one rank's (value, id) word pair; an id with bit 31 set, and the broken-link word, draw 0xFFFFFFFF: the same sentinel, so one downstream check covers both"; + +/// A kernel's spec. The raw kernels check themselves; the two moves' bound is exact, so gate 3 keeps every edge. +pub fn spec(w: Which) c.Spec { + return switch (w) { + // One spec for both moves; a gather's `n` is its frame row stride, and the 17,408 shape is seam B's. + .stage, .gather => .{ + .name = name(w), + .accum = .none, + .depth = noSums, + .out = .bf16, + .tiles = .{ .{}, .{}, .{ .size = 256, .step = 256, .min = 256 } }, + .served = if (w == .stage) &.{.{ .rows = 8, .n = 5120, .k = 5120 }} else &.{ + .{ .rows = 8, .n = 10240, .k = 5120 }, .{ .rows = 8, .n = 17408, .k = 8704 }, + .{ .rows = 8, .n = 10496, .k = 5120 }, + }, + .sums = .{ .f16 = 0, .bf16 = 0, .exact = false, .why = move_why }, + }, + .argmax => .{ + .name = name(w), + .accum = .none, + .depth = chainLogits, + .out = .raw, + .nan = .{ .stated = argmax_rule }, + .tiles = .{ .{}, .{}, .{ .size = 2048, .step = 64, .min = 64 } }, + .served = &.{.{ .rows = 8, .n = 1, .k = 124160 }}, + }, + .pick => .{ + .name = name(w), + .accum = .none, + .depth = noSums, + .out = .raw, + .nan = .{ .stated = pick_rule }, + .served = &.{ .{ .rows = 32, .n = 1, .k = 1 }, .{ .rows = 128, .n = 1, .k = 1 } }, + }, + }; +} + +/// `constant PairRows&` at buffer 1: `width`, the frame row `stride` and which rank's half this side writes. +const PairRows = extern struct { + width: u32, + stride: u32, + rank: u32, + x: u32, +}; + +fn rowsArgs(width: usize, stride: usize, rank: u32, x: u32) PairRows { + return .{ .width = @intCast(width), .stride = @intCast(stride), .rank = rank, .x = x }; +} + +fn bit(v: f32) u32 { + return @bitCast(v); +} + +fn floatOf(v: u32) f32 { + return @bitCast(v); +} + +/// One prepared call: the buffers the kernel is bound to, the host values its reference reads, and its own check. +pub const Pair = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + which: Which, + mutation: controls.Mutation, + pipe: mtl.Pipeline, + /// The hooks this probe exposes: the raw kernels check their own rule, the two computing ones use the reference. + v: g.VTable, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + rows: usize = 0, + width: usize = 0, + /// The rank this prepare writes first: the gather's `mine` half is its payload. + rank: u32 = 0, + /// The row stride in PairRows, in bf16 elements: the stage's source frame row, the gather's destination row. + stride: usize = 0, + + // Host inputs of the last prepare, as the reference and the check read them (freed by `drop`). + sh: ?[]u16 = null, // stage: the src payload, bf16 (rows, each a frame row of `stride`) + mh: ?[]u16 = null, // gather: this rank's half (mine), bf16 + nh: ?[]u16 = null, // gather: the peer's half (recv), bf16 + kh: ?[]u16 = null, // argmax: the bf16 logits + pw: ?[]u32 = null, // pick: this rank's (value bits, id) words + qw: ?[]u32 = null, // pick: the peer's words + + msg: [192]u8 = undefined, + sync: rig_.Ref = undefined, + mine: rig_.Ref = undefined, + peer: rig_.Ref = undefined, + logits: rig_.Ref = undefined, + out: rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, which: Which, mutation: controls.Mutation) !Pair { + const src = if (mutation == .none) sources.core_pair_rows else try controls.source(a, mutation); + defer if (mutation != .none) a.free(src); + const pipe = try rig_.pipeline(device, src, symbol(which)); + errdefer pipe.deinit(); + const want = launch(which); + const tg = rig_.tgMemory(pipe); + if (tg != want.tg_memory) { + std.debug.print("metal pair: {s} takes {d} bytes of threadgroup memory, its record says {d}\n", .{ symbol(which), tg, want.tg_memory }); + return error.ThreadgroupMemoryMismatch; + } + var v = g.hooks(Pair); + if (which != .argmax and which != .pick) v.check = null; + if (which != .stage and which != .gather) v.bounds = null; + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .which = which, .mutation = mutation, .pipe = pipe, .v = v }; + } + + pub fn deinit(self: *Pair) void { + self.drop(); + self.rig.deinit(); + self.pipe.deinit(); + } + + pub fn kernel(self: *Pair) g.Kernel { + var s = spec(self.which); + if (self.mutation != .none) { + s.name = self.mutation.name(); + return .{ .spec = s, .ctx = self, .vt = &self.v, .control = self.mutation.gate() }; + } + return .{ .spec = s, .ctx = self, .vt = &self.v }; + } + + fn drop(self: *Pair) void { + if (self.sh) |s| self.a.free(s); + if (self.mh) |s| self.a.free(s); + if (self.nh) |s| self.a.free(s); + if (self.kh) |s| self.a.free(s); + if (self.pw) |s| self.a.free(s); + if (self.qw) |s| self.a.free(s); + self.sh = null; + self.mh = null; + self.nh = null; + self.kh = null; + self.pw = null; + self.qw = null; + } + + /// Ordinary values: what the seam carries in production, symmetric about zero. + fn ordinary(rng: std.Random) f32 { + return rng.float(f32) * 2 - 1; + } + + /// A bf16 payload: what the two moves carry, with gate 1's non-finite plants where the reference looks for them. + fn payload(self: *Pair, len: usize, fill: g.Fill, rng: std.Random) ![]u16 { + const s = try self.a.alloc(u16, len); + for (s) |*v| v.* = c.toBf16(ordinary(rng)); + if (fill == .nonfinite) { + const nan = c.toBf16(std.math.nan(f32)); + const inf = c.toBf16(std.math.inf(f32)); + const ninf = c.toBf16(-std.math.inf(f32)); + for ([_]u16{ nan, inf, ninf, nan, ninf }, 0..) |plant, p| { + const at = (p + 1) * (len / 8 + 1) + p * 3; + if (at < len) s[at] = plant; + } + } + return s; + } + + /// Inputs for `shape` as `fill` asks, outputs poisoned and guarded; false when the kernel does not serve it. + pub fn prepare(self: *Pair, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.rows > 512) return false; + self.drop(); + self.rig.reset(); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + switch (self.which) { + .stage, .gather => { + if (shape.k == 0 or shape.k > 32768 or shape.k % 256 != 0) return false; + // The gather writes this rank's half first; alternating the rank with the seed covers both orders. + self.rank = @intCast(seed & 1); + self.sync = try self.rig.scratch(128); + if (self.which == .stage) { + // The source is a frame row: the stage reads this rank's `width` columns out of its `stride`. + self.stride = 2 * shape.k; + const len = shape.rows * self.stride; + const mine = try self.payload(len, fill, rng); + self.sh = mine; + self.mine = try self.rig.input(std.mem.sliceAsBytes(mine)); + self.out = try self.rig.output(shape.rows * shape.k * 2); + } else { + const n = shape.rows * shape.k; + // A frame row holds both halves, so its stride is at least 2*width: `n` sets a wider one. + self.stride = if (shape.n > 2 * shape.k) shape.n else 2 * shape.k; + const mine = try self.payload(n, fill, rng); + const peer = try self.payload(n, fill, rng); + self.mh = mine; + self.nh = peer; + self.mine = try self.rig.input(std.mem.sliceAsBytes(mine)); + self.peer = try self.rig.input(std.mem.sliceAsBytes(peer)); + self.out = try self.rig.output(2 * ((shape.rows - 1) * self.stride + 2 * shape.k)); + } + }, + .argmax => { + if (shape.k == 0 or shape.k > 262144) return false; + const n = shape.rows * shape.k; + const k = try self.a.alloc(u16, n); + self.kh = k; + for (k) |*v| v.* = c.toBf16(ordinary(rng)); + // Two in-range ids at +2.0: the lower id sits in the first half, so the tie rule's pick is fixed. + for (0..shape.rows) |r| { + const row = k[r * shape.k ..][0..shape.k]; + row[0] = c.toBf16(2.0); + row[shape.k / 2] = c.toBf16(2.0); + } + if (fill == .nonfinite) { + const nan = c.toBf16(std.math.nan(f32)); + const inf = c.toBf16(std.math.inf(f32)); + const ninf = c.toBf16(-std.math.inf(f32)); + for (0..shape.rows) |r| { + const row = k[r * shape.k ..][0..shape.k]; + row[shape.k / 3] = switch (r % 4) { + 0 => nan, + 1 => inf, + 2 => ninf, + else => row[shape.k / 3], + }; + if (r + 1 == shape.rows) for (row) |*v| { + v.* = nan; + }; + } + } + self.logits = try self.rig.input(std.mem.sliceAsBytes(k)); + self.out = try self.rig.output(shape.rows * 8); + }, + .pick => { + if (shape.n != 1 or shape.k != 1) return false; + self.pw = try self.a.alloc(u32, 2 * shape.rows); + self.qw = try self.a.alloc(u32, 2 * shape.rows); + for (0..shape.rows) |r| { + self.pw.?[2 * r] = bit(ordinary(rng)); + self.qw.?[2 * r] = bit(ordinary(rng)); + self.pw.?[2 * r + 1] = rng.uintLessThan(u32, 1 << 20); + self.qw.?[2 * r + 1] = rng.uintLessThan(u32, 1 << 20); + } + // Equal values on both ranks: the tie rule decides, and its direction alternates row by row. + for (0..shape.rows) |r| { + if (r % 4 != 2 and r % 4 != 3) continue; + self.qw.?[2 * r] = self.pw.?[2 * r]; + self.pw.?[2 * r + 1] = 300 + @as(u32, @intCast(r)); + self.qw.?[2 * r + 1] = if (r % 8 == 2) 700 + @as(u32, @intCast(r)) else 100 + @as(u32, @intCast(r)); + } + if (fill == .nonfinite) { + // Two plants prove the id flag is the rule: a flagged row draws the sentinel, a plain one picks. + for (0..shape.rows) |r| { + switch (r % 4) { + 0 => self.pw.?[2 * r + 1] |= flag, + 1 => self.qw.?[2 * r + 1] |= flag, + 2 => self.pw.?[2 * r] = bit(-std.math.inf(f32)), + else => self.qw.?[2 * r] = bit(std.math.inf(f32)), + } + } + } + self.sync = try self.rig.scratch(128); + self.mine = try self.rig.input(std.mem.sliceAsBytes(self.pw.?)); + self.peer = try self.rig.input(std.mem.sliceAsBytes(self.qw.?)); + self.out = try self.rig.output(shape.rows * 4); + }, + } + self.shape = shape; + self.rows = shape.rows; + self.width = shape.k; + return true; + } + + /// One launch on the prepared buffers, waited for. + pub fn run(self: *Pair) !void { + const p = self.rig.begin(); + const e = p.e; + const groups = mtl.Size.of(self.rows, 1, 1); + const threads = mtl.Size.of(launch(self.which).threads, 1, 1); + switch (self.which) { + .stage => { + @memset(self.rig.bytes(self.sync, 128), 0); + e.setPipeline(self.pipe); + e.setBuffer(self.sync.buf, self.sync.off, 0); + e.setValue(rowsArgs(self.width, self.stride, self.rank, 1), 1); + e.setBuffer(self.mine.buf, self.mine.off, 2); + e.setBuffer(self.out.buf, self.out.off, 3); + e.dispatchGroups(groups, threads); + }, + .gather => { + @memset(self.rig.bytes(self.sync, 128), 0); + const frame = self.rig.bytes(self.out, 2 * self.count()); + @memset(frame, 0); // the frame's gaps: no launch writes them, so their planted value is zero + e.setPipeline(self.pipe); + e.setBuffer(self.sync.buf, self.sync.off, 0); + e.setValue(rowsArgs(self.width, self.stride, self.rank, 1), 1); + e.setBuffer(self.mine.buf, self.mine.off, 2); + e.setBuffer(self.peer.buf, self.peer.off, 3); + e.setBuffer(self.out.buf, self.out.off, 4); + e.dispatchGroups(groups, threads); + }, + .argmax => { + e.setPipeline(self.pipe); + e.setBuffer(self.logits.buf, self.logits.off, 0); + const dims = [4]u32{ @intCast(self.width), logit_lo, @intCast(self.width), 1 }; + e.setValue(dims, 1); + e.setBuffer(self.out.buf, self.out.off, 2); + e.dispatchGroups(groups, threads); + }, + .pick => { + @memset(self.rig.bytes(self.sync, 128), 0); + e.setPipeline(self.pipe); + e.setBuffer(self.sync.buf, self.sync.off, 0); + e.setValue([2]u32{ @intCast(self.rows), 1 }, 1); + e.setBuffer(self.mine.buf, self.mine.off, 2); + e.setBuffer(self.peer.buf, self.peer.off, 3); + e.setBuffer(self.out.buf, self.out.off, 4); + e.dispatchGroups(mtl.Size.of(1, 1, 1), threads); + }, + } + try self.rig.finish(p); + } + + /// Output elements of the prepared call: bf16 elements for the two moves, the two raw kernels' bytes. + pub fn count(self: *Pair) usize { + return switch (self.which) { + .stage => self.rows * self.width, + .gather => (self.rows - 1) * self.stride + 2 * self.width, + .argmax => self.rows * 8, + .pick => self.rows * 4, + }; + } + + /// Float64 outputs of the prepared inputs and each one's sums' term magnitude (a move's term is the value itself). + pub fn reference(self: *Pair, want: []f64, mag: []f64) !void { + const width = self.width; + switch (self.which) { + .stage => { + const sh = self.sh orelse return error.NoReference; + for (0..self.rows) |r| { + const row = sh[r * self.stride ..][0..width]; + for (0..width) |i| { + const v = c.fromBf16(row[i]); + want[r * width + i] = v; + mag[r * width + i] = @abs(v); + } + } + }, + .gather => { + const mh = self.mh orelse return error.NoReference; + const nh = self.nh orelse return error.NoReference; + @memset(want, 0); // the frame's gaps: no launch writes them, so their reference value is zero + for (0..self.rows) |row| { + const at = row * self.stride; + for (0..width) |i| { + want[at + i] = c.fromBf16(if (self.rank == 0) mh[row * width + i] else nh[row * width + i]); + want[at + width + i] = c.fromBf16(if (self.rank == 0) nh[row * width + i] else mh[row * width + i]); + } + } + for (want, mag) |w, *m| m.* = @abs(w); + }, + else => return error.NoReference, + } + } + + /// The moves are exact: every output element must equal the reference's bf16 bit for bit. + pub fn bounds(self: *Pair, want: []const f64, mag: []const f64, allow: []f64) !void { + _ = want; + _ = mag; + if (self.which != .stage and self.which != .gather) return error.NotMove; + for (allow) |*l| l.* = 0; + } + + /// The seam's own checks: the raw kernels' rule at gate 1 (non-finite inputs) and gate 3 (every served edge). + pub fn check(self: *Pair, _: g.Gate) !?[]const u8 { + const out = try self.rig.outputBytes(); + switch (self.which) { + .argmax => { + const k = self.kh orelse return error.NoReference; + const n = self.width; + for (0..self.rows) |r| { + var best: f32 = -std.math.inf(f32); + var at: u32 = sentinel; + var bad = false; + for (0..n) |i| { + const v = c.fromBf16(k[r * n + i]); + if (!std.math.isFinite(v)) { + bad = true; + continue; + } + if (v > best) { + best = v; + at = logit_lo + @as(u32, @intCast(i)); + } + } + const got = std.mem.readInt(u32, out[r * 8 ..][0..4], .little); + const id = std.mem.readInt(u32, out[r * 8 + 4 ..][0..4], .little); + if (bad) { + if (!std.math.isNan(floatOf(got))) return try std.fmt.bufPrint(&self.msg, "row {d}: value {e} where a non-finite logit must write NaN", .{ r, floatOf(got) }); + if (id != at | flag) return try std.fmt.bufPrint(&self.msg, "row {d}: id {d} where flagging gives {d}", .{ r, id, at | flag }); + } else if (got != bit(best) or id != at) { + return try std.fmt.bufPrint(&self.msg, "row {d}: ({e} {d}) where its own rule gives ({e} {d})", .{ r, floatOf(got), id, best, at }); + } + } + }, + .pick => { + const pw = self.pw orelse return error.NoReference; + const qw = self.qw orelse return error.NoReference; + for (0..self.rows) |r| { + const vm = floatOf(pw[2 * r]); + const am = pw[2 * r + 1]; + const vp = floatOf(qw[2 * r]); + const ap = qw[2 * r + 1]; + const flagged = ((am | ap) & flag) != 0; + const want = if (flagged) sentinel else if (vp > vm or (vp == vm and ap < am)) ap else am; + const got = std.mem.readInt(u32, out[r * 4 ..][0..4], .little); + if (got != want) return try std.fmt.bufPrint(&self.msg, "row {d} picks {d} where its own rule gives {d}", .{ r, got, want }); + } + }, + else => return error.NotRaw, + } + return null; + } +}; diff --git a/zig/tests/contract/metal_pair_controls.zig b/zig/tests/contract/metal_pair_controls.zig new file mode 100644 index 0000000000..64c2bb91a8 --- /dev/null +++ b/zig/tests/contract/metal_pair_controls.zig @@ -0,0 +1,55 @@ +//! Planted controls for the pair seam's four kernels: one fault in pair_rows.metal's text, and the gate to catch it. +const std = @import("std"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); + +pub const Mutation = enum { + none, + stage_peer_half, + gather_own_halves, + argmax_high_tie, + argmax_no_flag, + pick_high_tie, + + pub fn name(m: Mutation) []const u8 { + return switch (m) { + .none => "", + .stage_peer_half => "control: metal pair stage, this rank stages the peer's half", + .gather_own_halves => "control: metal pair gather, this rank writes its own half on both sides", + .argmax_high_tie => "control: metal pair argmax, a tie goes to the highest id", + .argmax_no_flag => "control: metal pair argmax, a nonfinite logit is not flagged", + .pick_high_tie => "control: metal pair pick, a tie goes to the higher id", + }; + } + + /// The gate that must go red. + pub fn gate(m: Mutation) ?contract.gates.Gate { + return switch (m) { + .none => null, + .argmax_no_flag => .nonfinite, + else => .shapes, + }; + } + + const Edit = struct { from: []const u8, to: []const u8 }; + + fn edit(m: Mutation) Edit { + return switch (m) { + .none => unreachable, + // The stage's source is a frame row of `stride` = 2*width: the peer's half is in bounds and different. + .stage_peer_half => .{ .from = "const ulong in = ulong(r) * a.stride;", .to = "const ulong in = ulong(r) * a.stride + a.width;" }, + .gather_own_halves => .{ .from = "out[got + i] = ushort(word >> ((i & 1) * 16));", .to = "out[got + i] = mine[mrow + i];" }, + .argmax_high_tie => .{ .from = "if (ob > best || (ob == best && oa < at)) best = ob, at = oa;", .to = "if (ob > best || (ob == best && oa > at)) best = ob, at = oa;" }, + .argmax_no_flag => .{ .from = "if (!isfinite(v)) { bad = true; continue; }", .to = "if (false) { bad = true; continue; }" }, + .pick_high_tie => .{ .from = "(vp > vm || (vp == vm && ap < am)) ? ap : am;", .to = "(vp > vm || (vp == vm && ap > am)) ? ap : am;" }, + }; + } +}; + +/// pair_rows.metal's text for the kernel under test; a control's one edit must change it, or it proves nothing. +pub fn source(a: std.mem.Allocator, m: Mutation) ![]const u8 { + if (m == .none) return sources.core_pair_rows; + const e = m.edit(); + if (std.mem.indexOf(u8, sources.core_pair_rows, e.from) == null) return error.ControlEditMissing; + return std.mem.replaceOwned(u8, a, sources.core_pair_rows, e.from, e.to); +} From 52b11e46defdfe8fc82db1266fd21838f177750d Mon Sep 17 00:00:00 2001 From: Andrey Kolesnikov <15747200+akol1@users.noreply.github.com> Date: Wed, 28 Oct 2026 07:10:01 -0700 Subject: [PATCH 2/3] qwen27: give pair_pick's flag rule a control that goes red The pick's sentinel rule was the one pair rule without a planted control. The argmax control proves the row's non-finite flag is written, but nothing proved `pair_pick` acts on it: a pick that dropped the `((am | ap) & PAIR_FLAG) != 0u` test would still pass every gate. The sixth control replaces that test with `false`, planted only at the nonfinite fill -- at shapes both ranks' words are finite, so dropping the rule is invisible there and the control has to be red at the gate that plants it. The control is red as planted, and the pair docs now count six. --- docs/kernel-contract.md | 4 ++-- zig/tests/contract/metal_main.zig | 2 +- zig/tests/contract/metal_pair_controls.zig | 6 +++++- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 42ba822a3f..7fa49b6419 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -147,9 +147,9 @@ The row-split seam's four kernels (`pair_rows.metal`, launched by `pair_rows.zig `pair_stage` copies this rank's half of a frame row into the slot the relay sends, and `pair_gather` writes both halves of every row, its own natively and the peer's from the relayed frame. `pair_argmax` flags a row whose logits are not finite (its id word carries bit 31, its value word a NaN) and takes the lower id on a tie; `pair_pick` turns a flagged -rank's word into the id sentinel, so both Macs refuse a non-finite round where one Mac does. Five controls edit the +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 that does not flag at `nonfinite`. +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/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 7fd93c2ff8..9896a1cd97 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -219,7 +219,7 @@ pub fn main(init: std.process.Init) !u8 { .stage_peer_half => w == .stage, .gather_own_halves => w == .gather, .argmax_high_tie, .argmax_no_flag => w == .argmax, - .pick_high_tie => w == .pick, + .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; diff --git a/zig/tests/contract/metal_pair_controls.zig b/zig/tests/contract/metal_pair_controls.zig index 64c2bb91a8..1e7f4edbab 100644 --- a/zig/tests/contract/metal_pair_controls.zig +++ b/zig/tests/contract/metal_pair_controls.zig @@ -9,6 +9,7 @@ pub const Mutation = enum { gather_own_halves, argmax_high_tie, argmax_no_flag, + pick_no_flag, pick_high_tie, pub fn name(m: Mutation) []const u8 { @@ -18,6 +19,7 @@ pub const Mutation = enum { .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", }; } @@ -26,7 +28,7 @@ pub const Mutation = enum { pub fn gate(m: Mutation) ?contract.gates.Gate { return switch (m) { .none => null, - .argmax_no_flag => .nonfinite, + .argmax_no_flag, .pick_no_flag => .nonfinite, else => .shapes, }; } @@ -42,6 +44,8 @@ pub const Mutation = enum { .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" }, }; } }; From 55865b017d4b4818290bea895247c3a63883f47c Mon Sep 17 00:00:00 2001 From: Andrey Kolesnikov <15747200+akol1@users.noreply.github.com> Date: Tue, 27 Oct 2026 18:48:36 -0700 Subject: [PATCH 3/3] qwen27: rank-slice the family and hand the pair its seams MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Step 2 of 6, stacked on step 1 (`qwen27/pair-seam`) and main (1b9a987e); step 3 stacks on this commit. A rank's layer program now carries one `Reduce` interface instead of assuming the whole matrix. Attention's o_proj and the MLP's down_proj seam this rank's output columns, gather the whole frame, and let the producer kernel reduce over full K on its own rank; the head seams the last row's logits and takes the pair's argmax over the whole vocab, so the picked token is the single Mac's token, not rank 0's half-vocab pick. On one Mac `reduce` is null and every seam is compiled out of the program: the single-Mac program, its digests and its tests are the ones that ran before. The rank-local K-split is fixed while the slice is. A projection's split is keyed on the checkpoint tensor's rows, not the view's rows: `Triple` carries `whole` (the matrix's row count) and `prepare` derives `tensorRows` from it, while the M5 and prompt tables keep the checkpoint rows they always keyed. `in_proj_qkv` loads as three blocks of one tensor (10,240 rows), so its gate/up-style stack keeps main's 4 K slices on 6,144 rows instead of reading the q block's 2,048 rows and splitting into 8; every width main served on one Mac keeps main's split, and host tests "a rank's row slice keeps the checkpoint matrix's K split" and "every one-Mac stack keeps main's K split" cover the eight stacks (qkv 10,240 rows: 4, q 6,144: 4, against the 2,048-row view's 8). `Part` drops its unused `k0`, `Triple` its `k`/`k0`, and `Shard`/the family's private `cols`/`attentionCols`/`gdnCols`/`downCols` and the column-window code `prepare`/`packTiles` used them for are gone, with `config_test.zig`'s column-span assertions deleted and the row fields left as they were. `simd_qmm.call` takes the rank's `whole` width and `chunks(n, whole)` picks the chunk fence from it, so a rank's K tree is its own full matrix's while a one-Mac stack passes 0 and keeps exactly main's fence (the core test pins `chunks` on 6,144/12,288 and 5,120/10,240 at 8, and on 6,144 and 5,120 alone at 16, which is the pair's own-width split the `gn` forwarding is there to avoid). `quant_gpu.strided` forwards the `gn` of the linear it serves. `q27_attn_merge` takes the rank's head count (buffer 11) and strides a row by that many heads instead of 24 (`G * 4`), so the non-prompt multi-row rounds a served chat feeds every earlier reply through merge the rank's own rows; `attention.zig` binds `g.heads`. The buffered form is only what a pair runs. One Mac has its 24 heads at compile time, and main's kernel used them there (`G * 4`, and `PG * 4` in the prompt merge), so each site now has a second instance: `q27_attn_merge_single` and `q27_attn_prompt_merge_single` carry main's body verbatim — same reduction, same float arithmetic, main's indexing, no `heads` or `kv_heads` buffer — and the encoder takes a trailing `split` and picks `io[4]`/`prompt[1]` for a pair and `io[10]`/`prompt[2]` for one Mac (`decode_round.zig` passes `r.reduce != null`). The two instances compute the same values, so only the emitted code differs. A function constant would need `MTLFunctionConstantValues`, which the runtime library never creates (`PipelineBuild` would fail the contract's first gate), so the second compiled instance is the route that needs no new Metal machinery. The config gains its shard parse (`config_test.zig` covers it), the embedding table stays whole on every rank — only the head is sliced — with `embedBound` reading the table's bound rather than the family's `vocab_size`, and the session keeps the pair's pick as the round's token and surfaces `error.PairFailed` from the head seam's sentinel rather than returning `maxInt(u32)`. The seam's `Reduce` is threaded through `forward.encodeTp`/`session`, and the prompt attention kernel's signature follows the seam's reopened encoder. --- zig/kernels/metal/qwen27/attention_mpp.metal | 65 +++- .../metal/qwen27/attention_prompt.metal | 27 +- zig/src/core/simd_qmm.zig | 6 + zig/src/core/simd_qmm_bits.zig | 15 +- zig/src/families/gemma4/weights.zig | 2 +- zig/src/families/nemotron/forward.zig | 2 +- zig/src/families/nemotron/model.zig | 4 +- zig/src/families/qwen27/attention.zig | 37 +-- zig/src/families/qwen27/attention_plan.zig | 10 +- zig/src/families/qwen27/config.zig | 269 ++++++++++------- zig/src/families/qwen27/config_test.zig | 280 ++++++++++++++++++ zig/src/families/qwen27/decode_round.zig | 14 +- zig/src/families/qwen27/forward.zig | 129 +++++++- zig/src/families/qwen27/gdn_contract.zig | 5 +- zig/src/families/qwen27/gdn_hook.zig | 7 +- zig/src/families/qwen27/glue_gpu.zig | 7 +- zig/src/families/qwen27/gpu_weights.zig | 168 +++++++++-- zig/src/families/qwen27/model.zig | 46 ++- zig/src/families/qwen27/projection.zig | 149 ++++++++-- zig/src/families/qwen27/quant_gpu.zig | 44 +-- zig/src/families/qwen27/session.zig | 49 ++- zig/tests/contract/metal_affine.zig | 2 +- zig/tests/qwen27/attention.zig | 2 +- zig/tests/qwen27/attention_fused.zig | 4 +- zig/tests/qwen27/attention_precision.zig | 2 +- zig/tests/qwen27/forward_diff_runner.zig | 2 +- zig/tests/qwen27/synthetic_graph.zig | 5 +- zig/tests/qwen27/widths.zig | 2 +- 28 files changed, 1098 insertions(+), 256 deletions(-) create mode 100644 zig/src/families/qwen27/config_test.zig 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/src/core/simd_qmm.zig b/zig/src/core/simd_qmm.zig index 1856304072..e8e24feb5f 100644 --- a/zig/src/core/simd_qmm.zig +++ b/zig/src/core/simd_qmm.zig @@ -184,6 +184,12 @@ test "chunks follow output width only and tiles fit the reduction buffer" { 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/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 728cc3955f..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), 0); - 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/tests/contract/metal_affine.zig b/zig/tests/contract/metal_affine.zig index 3aa37d5de0..5d65c349db 100644 --- a/zig/tests/contract/metal_affine.zig +++ b/zig/tests/contract/metal_affine.zig @@ -202,7 +202,7 @@ pub const Proj = struct { 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/qwen27/attention.zig b/zig/tests/qwen27/attention.zig index 09b3f32eea..6b582e82d0 100644 --- a/zig/tests/qwen27/attention.zig +++ b/zig/tests/qwen27/attention.zig @@ -49,7 +49,7 @@ const Frame = struct { try op.preprocess(e, .{ .qg = f.get(.qg), .key = f.get(.raw_k), .q_gain = f.get(.q_gain), .k_gain = f.get(.k_gain), .norm_q = f.get(.norm_q), .norm_k = f.get(.norm_k), .query = f.get(.query), .rotated_key = f.get(.key), .positions = scratch.get(.positions), .rows = f.rows, .eps = 1e-6, .theta = 10000000 }); } fn encode(f: *const Frame, op: *const q.Attention, e: mtl.ComputeEncoder, scratch: *const q.Scratch, binding: *const q.Binding, plan: q.Plan, caches: []const Ref, gate: bool) !void { - try op.encode(e, scratch, binding, plan, f.get(.query), f.get(.key), f.get(.v), f.get(.out), caches, if (gate) f.get(.qg) else null, false); + try op.encode(e, scratch, binding, plan, f.get(.query), f.get(.key), f.get(.v), f.get(.out), caches, if (gate) f.get(.qg) else null, false, false); } }; const CacheOwner = struct { diff --git a/zig/tests/qwen27/attention_fused.zig b/zig/tests/qwen27/attention_fused.zig index c368e6d186..63a394a7a3 100644 --- a/zig/tests/qwen27/attention_fused.zig +++ b/zig/tests/qwen27/attention_fused.zig @@ -56,7 +56,7 @@ fn run(queue: mtl.Queue, op: *const attn.Attention, scratch: *attn.Scratch, bind try binding.upload(plan, caches); const cb = queue.commandBuffer(); const e = cb.compute(.concurrent); - try op.encode(e, scratch, binding, plan, query, key, value, out, caches, gate, false); + try op.encode(e, scratch, binding, plan, query, key, value, out, caches, gate, false, false); e.end(); cb.commit(); cb.wait(); @@ -98,7 +98,7 @@ fn bench(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, op: *const for (×) |*t| { const cb = queue.commandBuffer(); const e = cb.compute(.concurrent); - for (0..layers) |i| try op.encode(e, &scratch, &bindings[i], plan, .{ .buffer = query }, .{ .buffer = key }, .{ .buffer = value }, .{ .buffer = out }, caches[i .. i + 1], .{ .buffer = raw_gate }, false); + for (0..layers) |i| try op.encode(e, &scratch, &bindings[i], plan, .{ .buffer = query }, .{ .buffer = key }, .{ .buffer = value }, .{ .buffer = out }, caches[i .. i + 1], .{ .buffer = raw_gate }, false, false); e.end(); cb.commit(); cb.wait(); diff --git a/zig/tests/qwen27/attention_precision.zig b/zig/tests/qwen27/attention_precision.zig index 5621e9c7ba..4d703de3aa 100644 --- a/zig/tests/qwen27/attention_precision.zig +++ b/zig/tests/qwen27/attention_precision.zig @@ -157,7 +157,7 @@ pub fn main(init: std.process.Init) !void { try binding.upload(plan, &caches); const cb = queue.commandBuffer(); const e = cb.compute(.concurrent); - try op.encode(e, &scratch, &binding, plan, .{ .buffer = query }, .{ .buffer = key }, .{ .buffer = val }, .{ .buffer = out }, &caches, null, prompt and pass == 0); + try op.encode(e, &scratch, &binding, plan, .{ .buffer = query }, .{ .buffer = key }, .{ .buffer = val }, .{ .buffer = out }, &caches, null, prompt and pass == 0, false); e.end(); cb.commit(); cb.wait(); diff --git a/zig/tests/qwen27/forward_diff_runner.zig b/zig/tests/qwen27/forward_diff_runner.zig index 70c61eca42..84644a8d7a 100644 --- a/zig/tests/qwen27/forward_diff_runner.zig +++ b/zig/tests/qwen27/forward_diff_runner.zig @@ -29,7 +29,7 @@ const Hooks = struct { try r.model.glue.unstack(e, f.get(.mixed), f.get(.value), .{ .rows = rows, .width = 1024, .stride = 2048, .offset = 1024 }); e.barrier(); try r.attention.preprocess(e, .{ .qg = f.get(.q_gate), .key = f.get(.key), .q_gain = Weights.ref(weights.q_norm), .k_gain = Weights.ref(weights.k_norm), .norm_q = f.get(.norm_query), .norm_k = f.get(.norm_key), .query = f.get(.query), .rotated_key = f.get(.rotated_key), .positions = f.get(.positions), .rows = rows, .eps = r.model.config.eps, .theta = r.model.config.rope_theta }); - try r.attention.encode(e, &r.scratch, &r.bindings[layer], plan, f.get(.query), f.get(.rotated_key), f.get(.value), f.get(.mixer_out), &.{cache}, f.get(.q_gate), false); + try r.attention.encode(e, &r.scratch, &r.bindings[layer], plan, f.get(.query), f.get(.rotated_key), f.get(.value), f.get(.mixer_out), &.{cache}, f.get(.q_gate), false, false); } fn tap(ptr: *anyopaque, e: mtl.ComputeEncoder, layer: usize, hidden: Ref, rows: u32) !void { const h: *Hooks = @ptrCast(@alignCast(ptr)); diff --git a/zig/tests/qwen27/synthetic_graph.zig b/zig/tests/qwen27/synthetic_graph.zig index 3da574b7d1..f0e707634f 100644 --- a/zig/tests/qwen27/synthetic_graph.zig +++ b/zig/tests/qwen27/synthetic_graph.zig @@ -95,7 +95,8 @@ pub fn createLayers(a: std.mem.Allocator, layers: usize) !*q.model.Model { } var widths = try q.widths.Widths.uniformAll(a, c.layers); defer widths.deinit(); - var weights = try q.gpu_weights.load(a, device, &cp, c, &widths); + const shard = q.config.Shard{ .rank = 0, .ranks = 1, .global = c, .local = c }; + var weights = try q.gpu_weights.load(a, device, &cp, &shard, &widths); errdefer weights.deinit(); const glue = try q.glue_gpu.Kernels.init(device); errdefer glue.deinit(); @@ -105,6 +106,6 @@ pub fn createLayers(a: std.mem.Allocator, layers: usize) !*q.model.Model { f.deinit(); } const m = try a.create(q.model.Model); - m.* = .{ .allocator = a, .device = device, .queue = queue, .config = c, .checkpoint = cp, .weights = weights, .kernels = .{ .allocator = a, .device = device }, .glue = glue, .frame = frame }; + m.* = .{ .allocator = a, .device = device, .queue = queue, .shard = shard, .config = c, .checkpoint = cp, .weights = weights, .kernels = .{ .allocator = a, .device = device }, .glue = glue, .frame = frame }; return m; } diff --git a/zig/tests/qwen27/widths.zig b/zig/tests/qwen27/widths.zig index 9c8ff9b6e2..9e3ebf784e 100644 --- a/zig/tests/qwen27/widths.zig +++ b/zig/tests/qwen27/widths.zig @@ -306,7 +306,7 @@ const Gpu = struct { const cb = g.queue.commandBuffer(); const e = cb.compute(.serial); if (!std.mem.eql(u8, path(raw).name, "row")) { // simd_qmm_bits at every width, 4-bit too - const tiles = try row.simd_bits.call(g.bits, raw.bits, raw.n, raw.k, wide, raw.n, false); + const tiles = try row.simd_bits.call(g.bits, raw.bits, raw.n, raw.k, wide, raw.n, false, 0); row.simd_bits.encode(e, tiles, .{ .w = raw.w, .w_off = raw.w_off, .scales = raw.scales, .s_off = raw.s_off, .biases = raw.biases, .b_off = raw.b_off }, x, 0, y, 0); return commit(cb, e); }