diff --git a/CHANGELOG.md b/CHANGELOG.md index 99b2565df1..f94e6b6540 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- HIP: the decode round's residual-add and norm tail issues each row's loads before its branches, and a MoE tail loads its slots' rows together, with output bytes unchanged. On a gfx1151, one-stream decode is 2.7% faster on Qwen3.6-35B-A3B and 1.1% on the 27B. `tf_tail` is now under the kernel contract (#2247). Thanks to @AbrahamCete. + - Docs: `docs/optimization-ledger.md` lists which speed and memory optimizations apply to which native family, and where each stands, with links to the PRs that measured them (#2232). Thanks to @jschmied. - Docs: the pull request pre-flight asks a PR that adds, finishes or drops a family's optimization to update that family's cell in the optimization ledger (#2243). Thanks to @jschmied. diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 5281aa5b8f..b19db41e0a 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -298,6 +298,9 @@ form no registration runs. A fixed plan keeps the linear. Gate 4 proves an fp16 stress input (all-ones states: one weight at every width). Six controls: another K split (3, 4 and 5 bits; `linear_wc`), swapped scales, a swapped suh. HIP: the RMS norm rows, a wave a row (fp32, fp16 and bf16 out) and a block a row (fp16 and bf16 out). +HIP also checks the decode tail (`tf_tail`): the residual add, of a row of the activation kind or of nine slots' fp32 +rows times their weights summed in slot order, and the RMS norm after it, against float64 at Qwen3.6-35B-A3B's and +Qwen3.8-27B's widths, decode rows and windows, and the widest row the launch takes. HIP also checks the 64-row prompt attention tile (`prefill.hip`): `tf_fa_paged_bf16` and `tf_fa_paged_f16` over a permuted page table at the 27B's, the 35B-A3B's and Nemotron's head counts, each output within attention_ref's float64 bound, and `tf_fa_wide_bf16` and `tf_fa_wide_f16`, which no model calls, bit-equal to them on flat keys. A build runs diff --git a/zig/kernels/hip/decode/decode.hip b/zig/kernels/hip/decode/decode.hip index 19e2eff375..4b5fee3880 100644 --- a/zig/kernels/hip/decode/decode.hip +++ b/zig/kernels/hip/decode/decode.hip @@ -152,32 +152,67 @@ __device__ __forceinline__ void store_kind(void* p, int kind, long long i, float constexpr int kTailThreads = 512; constexpr int kTailWidth = 8192; +constexpr int kTailAhead = 4; // a thread's elements loaded before any is used, so their latencies overlap +constexpr int kTailSlots = 8; // the slots whose rows are loaded before any is summed -// One row a block, one launch: x = round(x + t), t = y or the slot-ordered weighted sum; normed = rms(x) * weight. -__device__ void tail_body(void* x, const void* y, const float* wts, const float* weight, void* normed, int kind, - int slots, int width, float eps) { - __shared__ float row_v[kTailWidth]; - __shared__ float part[kTailThreads / 32]; +// A thread's elements of the row: x = round(x + t), t = y or (SLOTS) the slots' sum in slot order; its sum of v * v. +template +__device__ __forceinline__ float tail_add(void* x, const void* y, const float* wts, float* row_v, int kind, int slots, + int width) { const int row = blockIdx.x; const long long at = static_cast(row) * width; float sum = 0.f; - for (int i = threadIdx.x; i < width; i += kTailThreads) { - float t; - if (slots > 0) { - float acc = 0.f; - for (int s = 0; s < slots; ++s) { - acc += static_cast(y)[(static_cast(row) * slots + s) * width + i] * - wts[row * slots + s]; + for (int base = threadIdx.x; base < width; base += kTailAhead * kTailThreads) { + // an element past the row loads the row's last, so no load waits behind a branch + int col[kTailAhead]; + float xv[kTailAhead], t[kTailAhead]; +#pragma unroll + for (int k = 0; k < kTailAhead; ++k) { + col[k] = min(base + k * kTailThreads, width - 1); + xv[k] = load_kind(x, kind, at + col[k]); + t[k] = SLOTS ? 0.f : load_kind(y, kind, at + col[k]); + } + // each element's slots in slot order, a group of slots' loads issued before their sums + if constexpr (SLOTS) { + for (int s0 = 0; s0 < slots; s0 += kTailSlots) { + float yv[kTailSlots][kTailAhead]; +#pragma unroll + for (int j = 0; j < kTailSlots; ++j) { + if (s0 + j >= slots) break; + const long long first = (static_cast(row) * slots + s0 + j) * width; + const float* ys = static_cast(y) + first; +#pragma unroll + for (int k = 0; k < kTailAhead; ++k) yv[j][k] = ys[col[k]]; + } +#pragma unroll + for (int j = 0; j < kTailSlots; ++j) { + if (s0 + j >= slots) break; + const float w = wts[row * slots + s0 + j]; +#pragma unroll + for (int k = 0; k < kTailAhead; ++k) t[k] += yv[j][k] * w; + } } - t = round_kind(acc, kind); - } else { - t = load_kind(y, kind, at + i); } - const float v = round_kind(load_kind(x, kind, at + i) + t, kind); - store_kind(x, kind, at + i, v); - row_v[i] = v; - sum += v * v; +#pragma unroll + for (int k = 0; k < kTailAhead; ++k) { + const int i = base + k * kTailThreads; + if (i >= width) break; + const float v = round_kind(xv[k] + (SLOTS ? round_kind(t[k], kind) : t[k]), kind); + store_kind(x, kind, at + i, v); + row_v[i] = v; + sum += v * v; + } } + return sum; +} + +// tail_body at one activation kind, so a thread's loads and stores take no branch on it. +template +__device__ __forceinline__ void tail_kind(void* x, const void* y, const float* wts, const float* weight, void* normed, + float* row_v, float* part, int slots, int width, float eps) { + const long long at = static_cast(blockIdx.x) * width; + float sum = slots > 0 ? tail_add(x, y, wts, row_v, KIND, slots, width) + : tail_add(x, y, wts, row_v, KIND, slots, width); #pragma unroll for (int mask = 16; mask > 0; mask >>= 1) sum += __shfl_xor(sum, mask, 32); if ((threadIdx.x & 31) == 0) part[threadIdx.x >> 5] = sum; @@ -186,10 +221,32 @@ __device__ void tail_body(void* x, const void* y, const float* wts, const float* #pragma unroll for (int w = 0; w < kTailThreads / 32; ++w) total += part[w]; const float inv = rsqrtf(total / static_cast(width) + eps); - for (int i = threadIdx.x; i < width; i += kTailThreads) { - float value = row_v[i] * inv; - if (weight != nullptr) value *= weight[i]; - store_kind(normed, kind, at + i, value); + for (int base = threadIdx.x; base < width; base += kTailAhead * kTailThreads) { + float w[kTailAhead]; +#pragma unroll + for (int k = 0; k < kTailAhead; ++k) { + w[k] = weight != nullptr ? weight[min(base + k * kTailThreads, width - 1)] : 1.f; + } +#pragma unroll + for (int k = 0; k < kTailAhead; ++k) { + const int i = base + k * kTailThreads; + if (i >= width) break; + float value = row_v[i] * inv; + if (weight != nullptr) value *= w[k]; + store_kind(normed, KIND, at + i, value); + } + } +} + +// One row a block, one launch: x = round(x + t), t = y or the slot-ordered weighted sum; normed = rms(x) * weight. +__device__ void tail_body(void* x, const void* y, const float* wts, const float* weight, void* normed, int kind, + int slots, int width, float eps) { + __shared__ float row_v[kTailWidth]; + __shared__ float part[kTailThreads / 32]; + if (kind == 1) { + tail_kind<1>(x, y, wts, weight, normed, row_v, part, slots, width, eps); + } else { + tail_kind<2>(x, y, wts, weight, normed, row_v, part, slots, width, eps); } } diff --git a/zig/src/core/contract/hip_records.zig b/zig/src/core/contract/hip_records.zig index 966c2ebed1..e973ec5dbd 100644 --- a/zig/src/core/contract/hip_records.zig +++ b/zig/src/core/contract/hip_records.zig @@ -494,7 +494,7 @@ const own_records = [_]Record{ }, }, } ++ @import("hip_records_pair.zig").records ++ @import("hip_records_prompt.zig").records ++ - @import("hip_records_score.zig").records ++ @import("hip_records_embed.zig").records; + @import("hip_records_score.zig").records ++ @import("hip_records_embed.zig").records ++ @import("hip_records_tail.zig").records; pub const records = own_records ++ @import("hip_records_stream.zig").records ++ @import("hip_records_router.zig").records ++ @import("hip_records_item.zig").records ++ @import("hip_records_qwen3.zig").records ++ @import("hip_records_columns.zig").records ++ diff --git a/zig/src/core/contract/hip_records_families.zig b/zig/src/core/contract/hip_records_families.zig index 0bc4158947..a7bdc2c992 100644 --- a/zig/src/core/contract/hip_records_families.zig +++ b/zig/src/core/contract/hip_records_families.zig @@ -16,6 +16,6 @@ pub const all = join: { test "the HIP records stay main's; the family lists only regroup GLM-OCR's and the joint schema's" { const every = @import("hip_records.zig").records; - try std.testing.expectEqual(@as(usize, 192), every.len); // main's records, GLM-OCR's 12 and the joint schema's 10 among them - try std.testing.expectEqual(@as(usize, 170), every.len - all.len); + try std.testing.expectEqual(@as(usize, 193), every.len); // main's records, GLM-OCR's 12 and the joint schema's 10 among them + try std.testing.expectEqual(@as(usize, 171), every.len - all.len); } diff --git a/zig/src/core/contract/hip_records_tail.zig b/zig/src/core/contract/hip_records_tail.zig new file mode 100644 index 0000000000..95c30aa7ed --- /dev/null +++ b/zig/src/core/contract/hip_records_tail.zig @@ -0,0 +1,12 @@ +//! Gate 5's record of the decode tail (decode.hip: tf_tail): a row a block of 512 threads, the row held in LDS. +const Record = @import("hip_records.zig").Record; + +pub const records = [_]Record{.{ + .launch = .{ .name = "decode tail tf_tail", .symbol = "tf_tail", .image = "decode", .threads = 512 }, + .usage = .{ + .{ .vgprs = 125, .sgprs = 58, .lds = 32832, .private = 0, .bound = 512, .wave = 32 }, + .{ .vgprs = 120, .sgprs = 58, .lds = 32832, .private = 0, .bound = 512, .wave = 32 }, + .{ .vgprs = 120, .sgprs = 58, .lds = 32832, .private = 0, .bound = 512, .wave = 32 }, + .{ .vgprs = 120, .sgprs = 58, .lds = 32832, .private = 0, .bound = 512, .wave = 32 }, + }, +}}; diff --git a/zig/src/hip/kernel_tests.zig b/zig/src/hip/kernel_tests.zig index e71b693d85..a9bf5694cb 100644 --- a/zig/src/hip/kernel_tests.zig +++ b/zig/src/hip/kernel_tests.zig @@ -21,6 +21,7 @@ test { _ = @import("kernel_tests/gdn.zig"); _ = @import("kernel_tests/mamba.zig"); _ = @import("kernel_tests/decode.zig"); + _ = @import("kernel_tests/tail_pin.zig"); _ = @import("kernel_tests/chain.zig"); _ = @import("kernel_tests/attention.zig"); _ = @import("kernel_tests/attention_pin.zig"); diff --git a/zig/src/hip/kernel_tests/tail_pin.zig b/zig/src/hip/kernel_tests/tail_pin.zig new file mode 100644 index 0000000000..5b8f3fcf88 --- /dev/null +++ b/zig/src/hip/kernel_tests/tail_pin.zig @@ -0,0 +1,101 @@ +//! The decode tail's exact output bits (the new residual and its norm), pinned by a digest per architecture. +const std = @import("std"); +const rig = @import("rig.zig"); +const Rig = rig.Rig; +const gpa = rig.gpa; +const at = rig.at; +const kernels = @import("../kernels.zig"); + +/// The digest of every case's residual and norm, in both activation kinds, on an architecture whose bits were recorded. +const pins = [_]struct { arch: []const u8, digest: u64 }{ + .{ .arch = "gfx1151", .digest = 0x325db42cb918dd99 }, +}; + +/// `rows` rows of `width` plus a plain row (`slots` 0) or the slots' weighted sum; `cancel` makes the slot order show. +const Case = struct { rows: usize, width: usize, slots: usize, weight: bool = true, cancel: bool = false }; + +/// Cancelling slots (past a slot group too), one row and a window at the 35B's and 27B's widths, ragged, unweighted. +const cases = [_]Case{ + .{ .rows = 2, .width = 2048, .slots = 3, .cancel = true }, + .{ .rows = 2, .width = 2048, .slots = 9, .cancel = true }, + .{ .rows = 1, .width = 2048, .slots = 0 }, + .{ .rows = 1, .width = 2048, .slots = 9 }, + .{ .rows = 4, .width = 2048, .slots = 9 }, + .{ .rows = 1, .width = 5120, .slots = 0 }, + .{ .rows = 3, .width = 5120, .slots = 0 }, + .{ .rows = 2, .width = 1000, .slots = 3 }, + .{ .rows = 2, .width = 2048, .slots = 0, .weight = false }, + .{ .rows = 1, .width = 8192, .slots = 2 }, +}; + +/// FNV-1a over `bytes`, continuing from `h`. +fn fnv(h: u64, bytes: []const u8) u64 { + var x = h; + for (bytes) |b| { + x ^= b; + x *%= 0x100000001b3; + } + return x; +} + +/// One case's residual and norm through tf_tail at activation kind `kind` (1 fp16, 2 bf16), folded into `h`. +fn run(t: *Rig, rng: *rig.Rng, c: Case, kind: c_int, h: u64) !u64 { + const bits = if (kind == 1) &rig.f16Bits else &rig.bf16Bits; + const n = c.rows * c.width; + const hx = try gpa.alloc(u16, n); + defer gpa.free(hx); + for (hx) |*v| v.* = bits(rng.unit()); + const hy16 = try gpa.alloc(u16, if (c.slots > 0) 0 else n); + defer gpa.free(hy16); + for (hy16) |*v| v.* = bits(rng.unit() / 8.0); + const hy32 = try gpa.alloc(f32, c.rows * c.slots * c.width); + defer gpa.free(hy32); + for (hy32, 0..) |*v, i| { + const slot = (i / c.width) % @max(c.slots, 1); + const big: f32 = if (!c.cancel or slot > 1) 0 else if (slot == 0) 4096 else -4096; + v.* = big + rng.unit() / 8.0; + } + const hw = try gpa.alloc(f32, c.rows * @max(c.slots, 1)); + defer gpa.free(hw); + for (hw) |*v| v.* = if (c.cancel) 0.5 else (rng.unit() + 1.0) / @as(f32, @floatFromInt(@max(c.slots, 1))); + const hn = try gpa.alloc(f32, c.width); + defer gpa.free(hn); + for (hn) |*v| v.* = 1.0 + rng.unit() / 4.0; + var x = try t.upload(hx); + defer x.free(); + var y = if (c.slots > 0) try t.upload(hy32) else try t.upload(hy16); + defer y.free(); + var wts = try t.upload(hw); + defer wts.free(); + var weight = try t.upload(hn); + defer weight.free(); + var normed = try t.alloc(n * 2); + defer normed.free(); + const wp: ?[*]const f32 = if (c.slots > 0) @ptrFromInt(at(wts)) else null; + const np: ?[*]const f32 = if (c.weight) @ptrFromInt(at(weight)) else null; + try t.on.tf_tail(@ptrFromInt(at(x)), @ptrFromInt(at(y)), wp, np, @ptrFromInt(at(normed)), kind, @intCast(c.rows), @intCast(c.slots), @intCast(c.width), 1e-6, t.stream.handle); + try t.stream.synchronize(); + const gx = try rig.download(u16, x); + defer gpa.free(gx); + const gn = try rig.download(u16, normed); + defer gpa.free(gn); + return fnv(fnv(h, std.mem.sliceAsBytes(gx)), std.mem.sliceAsBytes(gn)); +} + +test "the decode tail writes its pinned residual and norm bits in fp16 and bf16, plain, slot-summed, ragged and wide" { + const t = try Rig.open(1 << 20); + defer t.close(); + var rng: rig.Rng = .{ .state = 0x9E3779B97F4A7C15 }; + var h: u64 = 0xcbf29ce484222325; + for ([_]c_int{ 2, 1 }) |kind| { + for (cases) |c| h = try run(t, &rng, c, kind, h); + } + const want = for (pins) |p| { + if (std.mem.eql(u8, p.arch, kernels.arch)) break p.digest; + } else { + std.debug.print("decode tail: no digest pinned for {s}; this run wrote 0x{x:0>16}\n", .{ kernels.arch, h }); + return error.SkipZigTest; + }; + if (h != want) std.debug.print("decode tail digest on {s}: 0x{x:0>16}\n", .{ kernels.arch, h }); + try std.testing.expectEqual(want, h); +} diff --git a/zig/tests/hip/contract_main.zig b/zig/tests/hip/contract_main.zig index a2252abf96..dd39d3135d 100644 --- a/zig/tests/hip/contract_main.zig +++ b/zig/tests/hip/contract_main.zig @@ -15,6 +15,7 @@ const moe_k = @import("contract_nemo_moe.zig"); const walk_k = @import("contract_walk.zig"); const prefill_k = @import("contract_prefill.zig"); const add_rms_k = @import("contract_add_rms.zig"); +const tail_k = @import("contract_tail.zig"); const window_k = @import("contract_window.zig"); const relu2_k = @import("contract_relu2.zig"); const kv_k = @import("contract_kv_write.zig"); @@ -191,6 +192,10 @@ fn run(a: std.mem.Allocator, io: std.Io) !usize { for (&add_rms, std.enums.values(add_rms_k.Mode)) |*m, w| m.* = .init(a, &r, s, &l, w); defer for (&add_rms) |*m| m.deinit(); for (&add_rms) |*m| try kernels.append(a, m.kernel()); + var tails: [2]tail_k.Tail = undefined; + for (&tails, std.enums.values(tail_k.Mode)) |*m, w| m.* = .init(a, &r, s, &l, w); + defer for (&tails) |*m| m.deinit(); + for (&tails) |*m| try kernels.append(a, m.kernel()); var convs: [2]window_k.ConvWindow = undefined; for (&convs, std.enums.values(window_k.Part)) |*m, p| m.init(a, &r, s, &l, p); defer for (&convs) |*m| m.deinit(); diff --git a/zig/tests/hip/contract_tail.zig b/zig/tests/hip/contract_tail.zig new file mode 100644 index 0000000000..dd738c1ca0 --- /dev/null +++ b/zig/tests/hip/contract_tail.zig @@ -0,0 +1,261 @@ +//! tf_tail under the contract on HIP: the decode tail's residual add (a row or the slots' sum) and the norm after it. +const std = @import("std"); +const hip = @import("hip"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +const kit = @import("contract_kit.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const rows = contract.rows; + +/// What the residual adds: a row of the kind (slots 0), or a MoE layer's slots' fp32 rows times their weights. +pub const Mode = enum { plain, moe }; + +/// Qwen3.6-35B-A3B's MoE tail: eight routed experts and the shared one. +const slots = 9; +const threads = 512; +const widest = 8192; +const eps: f32 = 1e-6; + +/// A term's square, its thread's strided chain, its wave's xor tree, then the block's 16 wave partials in order. +fn depthPlain(s: c.Shape) usize { + return 1 + (s.k + threads - 1) / threads + 5 + threads / 32; +} + +/// The slots' products and adds, in slot order, before the norm's sum. +fn depthMoe(s: c.Shape) usize { + return depthPlain(s) + 2 * slots; +} + +/// The 35B's 2,048 at a decode row and a window's rows, the 27B's 5,120, and the widest row the launch takes. +const served_plain = [_]c.Shape{ + .{ .rows = 1, .n = 2048, .k = 2048 }, + .{ .rows = 16, .n = 2048, .k = 2048 }, + .{ .rows = 1, .n = 5120, .k = 5120 }, + .{ .rows = 16, .n = 5120, .k = 5120 }, + .{ .rows = 3, .n = widest, .k = widest }, +}; + +/// The 35B's MoE tail at a decode row and a window's rows, and the widest row the launch takes. +const served_moe = [_]c.Shape{ + .{ .rows = 1, .n = 2048, .k = 2048 }, + .{ .rows = 16, .n = 2048, .k = 2048 }, + .{ .rows = 3, .n = widest, .k = widest }, +}; + +const why = "a 16-bit output follows the sum of squares as ss^-1/2: a bf16 sum's error is a few output ulps"; +const why_bf16 = why ++ "; fp16's overflowing sums give 0, which the bound's residual roundings put at 51x to 85x"; + +pub fn spec(mode: Mode, k: kit.Kind) c.Spec { + const names = [2][2][]const u8{ + .{ "hip tf_tail fp16 (y fp16)", "hip tf_tail bf16 (y bf16)" }, + .{ "hip tf_tail fp16 (9 fp32 slots)", "hip tf_tail bf16 (9 fp32 slots)" }, + }; + return .{ + .name = names[@backingInt(mode)][@intFromBool(k == .bf16)], + .accum = .f32, + .depth = if (mode == .moe) &depthMoe else &depthPlain, + // the residual's add, the mean, eps, rsqrtf (2 ulps), the products by inv and by the weight + .ops = 7, + // the residual rounded before the norm reads it, directly and through the sum; the slots' sum rounded too + .mids = if (mode == .moe) 4 else 2, + .out = kit.elem(k), + // a pass over the row is 4 elements a thread of 512, its last ragged + .tiles = .{ .{}, .{}, .{ .size = 4 * threads } }, + .served = if (mode == .moe) &served_moe else &served_plain, + .sums = if (k == .bf16) .{ .f16 = 40, .bf16 = 0, .why = why_bf16 } else .{ .bf16 = 0, .why = why }, + }; +} + +/// shape.rows rows of shape.k: x [rows][k] of the kind updated in place, normed [rows][k] of the kind. +pub const Tail = struct { + rig: rig_.Rig, + mode: Mode, + l: *const hip.Launcher, + k: kit.Kind, + arena: std.heap.ArenaAllocator, + shape: c.Shape = undefined, + x: []f64 = &.{}, + y: []f64 = &.{}, // [rows][k], or the slots' [rows][slots][k] + wts: []f64 = &.{}, // [rows][slots] + w: []f64 = &.{}, + x_bytes: []u8 = &.{}, + in: [3]u64 = @splat(0), // y, wts (slots only), the norm's weight + out: [2]u64 = @splat(0), // x, normed + + pub fn init(a: std.mem.Allocator, rt: *const hip.Runtime, s: hip.Stream, l: *const hip.Launcher, m: Mode) Tail { + return .{ .rig = .{ .a = a, .r = rt, .s = s }, .mode = m, .l = l, .k = kit.act(), .arena = .init(a) }; + } + + pub fn deinit(self: *Tail) void { + self.arena.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *Tail) g.Kernel { + return .{ .spec = spec(self.mode, self.k), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Tail); + + fn signed(rng: std.Random, lo: f64, hi: f64) f64 { + const v = lo + rng.float(f64) * (hi - lo); + return if (rng.boolean()) v else -v; + } + + /// x in +-[1, 2) and an addend of x's sign at about a quarter to a half of it, so the add never cancels. + pub fn prepare(self: *Tail, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.k == 0 or shape.k > widest) return false; + self.rig.reset(); + _ = self.arena.reset(.retain_capacity); + self.shape = shape; + const a = self.arena.allocator(); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const k = self.k; + const n = shape.rows * shape.k; + const moe = self.mode == .moe; + self.x = try a.alloc(f64, n); + self.y = try a.alloc(f64, if (moe) n * slots else n); + self.wts = try a.alloc(f64, if (moe) shape.rows * slots else 0); + self.w = try kit.floats(a, rng, shape.k, 0.75, 1.25); + for (self.x) |*v| v.* = kit.round(k, signed(rng, 1, 2)); + for (self.wts) |*v| v.* = @as(f32, @floatCast(0.05 + rng.float(f64) * 0.1)); // nine weights near 0.1 each + for (0..n) |i| { + const sign: f64 = if (self.x[i] < 0) -1 else 1; + if (!moe) { + self.y[i] = kit.round(k, sign * (0.25 + rng.float(f64) * 0.25) * @abs(self.x[i])); + continue; + } + const r = i / shape.k; + for (0..slots) |j| { + const part = sign * (0.25 + rng.float(f64) * 0.25) * @abs(self.x[i]) / (slots * 0.1); + self.y[(r * slots + j) * shape.k + i % shape.k] = @as(f32, @floatCast(part)); + } + } + if (fill == .stress) try self.stress(seed); + if (fill == .nonfinite) self.plant(); + self.x_bytes = try a.alloc(u8, n * kit.size(k)); + for (self.x, 0..) |v, i| kit.put(k, self.x_bytes, i, v); + const yb = if (moe) try kit.f32Bytes(a, self.y) else blk: { + const b = try a.alloc(u8, n * kit.size(k)); + for (self.y, 0..) |v, i| kit.put(k, b, i, v); + break :blk b; + }; + self.in[0] = try self.rig.input(yb); + self.in[1] = if (moe) try self.rig.input(try kit.f32Bytes(a, self.wts)) else 0; + self.in[2] = try self.rig.input(try kit.f32Bytes(a, self.w)); + for (&self.out) |*at| at.* = try self.rig.output(self.x_bytes.len); + return true; + } + + /// Gate 4: the norm's rows of integers (rows.zig's) and nothing added, so its sums are exact in fp32. + fn stress(self: *Tail, seed: u64) !void { + const s = self.shape; + var r = try rows.Rows.init(self.arena.allocator(), .rms, .{ .rows = s.rows, .n = s.k, .k = s.k }); + r.fill(.stress, seed); + for (self.x, r.x) |*v, b| v.* = c.fromBf16(b); + @memset(self.y, 0); + } + + /// NaN in row 0's x, +Inf in row 1's addend and -Inf in row 2's (a slot's row), an overflowing sum in row 3. + fn plant(self: *Tail) void { + const s = self.shape; + const big: f64 = if (self.k == .f16) 60000 else 3e38; + if (5 < self.x.len) self.x[5] = std.math.nan(f64); + const at = [_]usize{ s.k + 9, 2 * s.k + 1 }; + for ([_]f64{ std.math.inf(f64), -std.math.inf(f64) }, at) |v, i| if (i < self.x.len) switch (self.mode) { + .moe => self.y[(i / s.k * slots + 1) * s.k + i % s.k] = v, + .plain => self.y[i] = v, + }; + if (3 * s.k + 2 < self.x.len) self.x[3 * s.k + 2] = kit.round(self.k, big); + if (3 * s.k + 2 < self.x.len and self.mode == .plain) self.y[3 * s.k + 2] = kit.round(self.k, big); + } + + fn ptr(at: u64) ?*anyopaque { + return @ptrFromInt(at); + } + + /// x back to the prepared rows, then the launch as the engine makes it (no weights for a plain row). + pub fn run(self: *Tail) !void { + try self.rig.put(self.out[0], self.x_bytes); + const s = self.shape; + const wts: ?[*]const f32 = if (self.mode == .moe) @ptrFromInt(self.in[1]) else null; + try self.l.tf_tail( + ptr(self.out[0]), + ptr(self.in[0]), + wts, + @ptrFromInt(self.in[2]), + ptr(self.out[1]), + @backingInt(self.k), + @intCast(s.rows), + if (self.mode == .moe) slots else 0, + @intCast(s.k), + eps, + self.rig.s.handle, + ); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Tail) usize { + return 2 * self.x.len; + } + + /// Element i's addend in float64, and the magnitude of the terms the slots' sum carries into it. + fn addend(self: *Tail, i: usize) struct { f64, f64 } { + const s = self.shape; + if (self.mode == .plain) return .{ self.y[i], 0 }; + const r = i / s.k; + var acc: f64 = 0; + var abs: f64 = 0; + for (0..slots) |j| { + const t = self.y[(r * slots + j) * s.k + i % s.k] * self.wts[r * slots + j]; + acc += t; + abs += @abs(t); + } + return .{ acc, abs }; + } + + /// x + the addend, then rms(x + addend) * w, in float64; the norm's magnitude as rows.zig's. + pub fn reference(self: *Tail, want: []f64, mag: []f64) !void { + const s = self.shape; + const n = self.x.len; + for (0..n) |i| { + const b, const m = self.addend(i); + const v = self.x[i] + b; + want[i] = if (std.math.isInf(kit.round(self.k, v))) kit.round(self.k, v) else v; // past the kind's range + mag[i] = m; + } + for (0..s.rows) |r| { + const row = want[r * s.k ..][0..s.k]; + var ss: f64 = 0; + for (row) |v| ss += v * v; + const inv = 1 / @sqrt(ss / @as(f64, @floatFromInt(s.k)) + eps); + for (row, 0..) |v, j| { + const i = n + r * s.k + j; + want[i] = v * inv * self.w[j]; + mag[i] = @abs(want[i]) / 2 + mag[r * s.k + j] * inv * self.w[j]; + } + } + } + + /// Gate 4's rows, nothing added: the norm's sum of squares in `acc` and `order`, the rest exact. + pub fn simulate(self: *Tail, acc: c.Accum, order: g.Order, out: []f64) !void { + const s = self.shape; + const n = self.x.len; + const t = try self.rig.a.alloc(f64, s.k); + defer self.rig.a.free(t); + for (0..s.rows) |r| { + const row = self.x[r * s.k ..][0..s.k]; + for (t, row) |*v, x| v.* = x * x; + const ss = rows.sumIn(acc, order, t); + const inv = 1 / @sqrt(ss / @as(f64, @floatFromInt(s.k)) + eps); + for (row, 0..) |x, j| { + out[r * s.k + j] = x; + out[n + r * s.k + j] = x * inv * self.w[j]; + } + } + } +};