diff --git a/CHANGELOG.md b/CHANGELOG.md index 91bcfefaec..35cc9b70be 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 +- Kernel contract (Metal): the M5's tensor units keep 10 mantissa bits of an fp32 `mma_16x32` operand, cut toward zero. The contract now states that cut for every kernel with an fp32 operand: Gemma 4's, Spark's, Flash Next's, Nemotron's and MiniMax's attention and Nemotron's SSD products. Each has a control that fails without it, and no output changes (#1748, part of #1546). Thanks to @BobClawblaw. + - GLM-5.3-Flash on CUDA: `TF_GLM_CHUNK_SOLO=4096` (opt-in, on every rank) lets a prompt pass take 4,096 rows while no other stream is decoding, with the same bytes. On GB10s a cold prompt's first token comes 6% to 9% sooner from 8K to 128K tokens. It holds 1.6 to 1.9 GiB more for prompt buffers, and a request that arrives during a long pass waits for it to end. Unset, nothing changes (#1823, option C of #1554). Thanks to @Bizuayeu. - Qwen3.8-27B on M5: checkpoints that mix bit widths in a fused projection (such as the oQ4 and oQ2 quantizations) answer prompts longer than 128 tokens, instead of failing with `BadCoreProjectionRows`; each 128-row piece is the call main already made, so replies keep main's tokens (#1822). Thanks to @grantoverton. diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index b30ed4c065..09b68a97bc 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -121,6 +121,16 @@ HIP's prompt attention tile (`tf_fa_paged_*`, `tf_fa_wide_*`) writes 0 for a row float64 gives NaN, as its running max is `fmaxf` and a row at max -inf takes p = 0; and an infinite value at a key whose fp32 weight is 0 (past the row's causal edge in its 64-key tile, or under fp32's least) gives NaN where float64 gives inf or a finite value, as its P x V takes every key of the tile. +Qwen3.8 Flash Next's `tf_sattn_nax` and `tf_attn256_nax` score masked keys -FLT_MAX, the value each row's running +max starts from, so a row whose every score is -Inf weighs its masked keys 1 and gives a finite average where float64 +gives NaN. `tf_sattn_nax`'s padding lanes also read key 0 at p = 0, so a non-finite value in key 0's row reaches every +row whose key count is not a multiple of 32 (0 x Inf is NaN where float64 gives the infinity). +MiniMax-M3.1-Flash's `mm_prompt_attention`, Spark-X2.6-4B's `sp_prompt_attention`, Gemma 4's `gemma_prompt_attention` +(31B and 12B, sliding and global), Nemotron's `tf_attention_nax` and Qwen3.8 Flash Next's `tf_attn256_nax` give NaN +for a row whose tile or block reads a non-finite value at a key masked for that row, and Gemma 4's +`gemma_decode_attention` for a decode row whose 32-key step reads one: P = 0 times Inf inside the P x V mma, where +float64 leaves the masked key out. These were seen on the M5, where their cases run. Each kernel has an entry of its own +that plants such a value and carries the finding, so its other entries keep a clean gate 1. HIP’s MLX stream tile keeps its finite dot2 sums unchanged. A non-finite result with an actual non-finite operand is recomputed by full quantization groups using the factored float64 formula. The plain tile is @@ -138,6 +148,43 @@ bits) and simd_qmm_bits (3, 6 and 8 bits), and the lane projection (6-bit groups 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. +On the M5's tensor units, Metal also checks what nax.h's `mma_16x32` keeps of an fp32 operand +(`metal_mma_operand.zig`): the left, the right and both (attention's fp32 P, `gemm_nax`'s f32 kernels), each output a +single product bit-equal to float64 on operands with their low 13 mantissa bits cleared. That is 10 bits kept, rounded +toward zero, fp32's exponents and subnormals kept, NaN and the infinities passed. Four controls state the operand as +fp32 or as rounded to nearest at 10 bits, on each side, and must go red at `shapes`. The cut was measured on the tensor +units only, so the cases run only there; what M1-M4 keep of an fp32 operand is not measured. +`core/attention_range.zig` names that cut (`operand_cut`, 2^-10) as a form's `p_cut`: one-sided, since each p only +shrinks, so it moves an output down by up to 2^-10 of its positive values' share and up by its negative values'. On the +M5, Nemotron's `tf_attention_nax` states its fp32 P with it, and MiniMax-M3.1-Flash's `mm_prompt_attention` does on +every chip (before the M5 that only widens its range); each bf16 output must lie between its range's ends rounded to +bf16. Nemotron's adds a case with P spread over about +[2^-6, 1] and values in [0, 1), where the cut's bias adds up; its control, the same range without the cut, goes red at +`shapes`. Before the M5, `tf_attention_nax` keeps attention_ref's bound. +Qwen3.8 Flash Next's prompt attention on the M5 does too: `tf_sattn_nax` (the indexer's dense and selected rows) and +`tf_attn256_nax` (the causal chunk without the indexer), head dim 256, 24 heads over 2. Each bf16 output, +bf16(bf16(o) * bsig(gate)), must lie between its range's ends rounded to bf16 and through the gate's rounded sigmoid. +Each adds a P-spread case whose cutless range is a planted control, and a source edit gate 3 must catch: `tf_sattn_nax` +counting its padding lane (key 0 again), `tf_attn256_nax` seeing one key past the row. Before the M5 neither is +registered. +On the M5, Nemotron's SSD steps (`prefill_mixers.zig`'s Mamba mixer) run `gemm_nax`'s five f32 kernels with both +operands fp32, so each product enters cut, smaller by up to 1 - (1 - 2^-10)^2. Their entries +(`metal_nemotron_ssd_gemm.zig`, y = sur x dtx and the state dd^T x bh^T a head, at the step's served row counts) state +it as a one-sided range: float64's value less that share of the positive terms and plus that of the negative, beside +the fp32 sums' `mma_step` roundings for every k16 step. Gate 4 runs 11-bit integers times 5-bit ones, which the cut +keeps. Two cases take operands in (0, 1], where the bias adds up; their controls state the operands as fp32 and go red +at `shapes`, as does a source edit dropping the ragged K tail's last column. Before the M5 they are not registered; +their launches are in `nemotron_ssd_records.zig`. +Spark-X2.6-4B's `sp_prompt_attention` registers both pipelines the engine builds (sliding layers over a 512-key window +in a 2,560-slot ring, and full layers), its fp32 P stated with the cut, each with a P spread case +and its cutless control, and one source edit each (a key past the window, a key past the row); before the M5 they are +not registered. +On the M5, Gemma 4's `gemma_prompt_attention` (the 31B's and 12B's sliding and global layers) and +`gemma_decode_attention` with `gemma_attention_merge` (their global layers; the 12B's assistant launches the 12B's) +state their fp32 P with the cut too, each form's counts taken from the kernel's own steps, at the sliding ring, image +rows and decode chunks the engine serves. A P-spread case on the 31B's global prompt and decode has a cutless control +each, red at `shapes`, and a score edit (2 log2e) in each kernel is a control. Their launches are in +`gemma_attention_records.zig`. LFM2.5-24B-A2B's shared rounds: the greedy pick (`lfm2_argmax`) and the conv and attention kernels that run every stream's rows in one dispatch, each bit-equal to the per-stream kernels it replaces. Four controls edit the source: a tie to the highest id, a nonfinite logit not flagged, and every stream reading or writing the first stream's state or cache. diff --git a/zig/src/core/attention_range.zig b/zig/src/core/attention_range.zig index d8cb6bc57f..dcd8b6f1c0 100644 --- a/zig/src/core/attention_range.zig +++ b/zig/src/core/attention_range.zig @@ -18,6 +18,8 @@ pub const Form = struct { chunk_exp: bool, /// P's error into P x V, relative: bf16 2^-8, two bf16 halves 2^-16, fp32 0. p_rel: f64, + /// An fp32 P cut toward zero as it enters P x V, relative and one-sided (the M5's tensor units: operand_cut). + p_cut: f64 = 0, /// The chains of o's terms p_j v_j and of l's terms p_j, as roundings of their terms' magnitude. num: f64, den: f64, @@ -33,6 +35,9 @@ pub const Form = struct { /// An mma k16 step's error in 2^-24 of |acc| + sum |a b|: 17 terms cut at the largest's ulp, then rounded down. pub const mma_step = 36; +/// What the M5's tensor units keep of an fp32 operand: 10 mantissa bits, cut toward zero (metal_mma_operand.zig). +pub const operand_cut = 0x1p-10; + /// window_rows.cu's decode at head dim d over n keys: lanes' chains of d / 32 then a 5-level tree, 256-key chunks. pub fn decode(d: usize, n: usize) Form { const chunks: f64 = @floatFromInt(n / 256 + 2); @@ -105,6 +110,47 @@ pub fn head(q: []const f64, k: []const f64, v: []const f64, n: usize, d: usize, } } +/// P's cut after `head` (w, tau as it left them): each p only shrinks, so positive values move o down, negative up. +pub fn cutSides(f: Form, w: []const f64, tau: []const f64, v: []const f64, n: usize, d: usize, down: []f64, up: []f64) void { + var l: f64 = 0; + var tmax: f64 = 0; + for (w[0..n], tau[0..n]) |p, t| { + l += p; + if (std.math.isFinite(t)) tmax = @max(tmax, t); + } + const kappa = (1 + tmax) / (1 - tmax); + for (0..d) |e| { + var pos: f64 = 0; + var neg: f64 = 0; + for (0..n) |j| { + const x = v[j * d + e]; + if (x > 0) pos += w[j] * x else neg -= w[j] * x; + } + const finite = std.math.isFinite(pos) and std.math.isFinite(neg) and tmax < 0.5; + down[e] = if (finite) f.p_cut * pos / l * kappa * (1 + 0x1p-20) else 0; + up[e] = if (finite) f.p_cut * neg / l * kappa * (1 + 0x1p-20) else 0; + } +} + +/// A bf16 output's ends, point - half - down to point + half + up, each rounded to nearest as the output is. +pub fn ends(point: f64, half: f64, down: f64, up: f64) [2]f64 { + return .{ bfr.bf16(point - half - down), bfr.bf16(point + half + up) }; +} + +/// A bf16 output's allowance: the range's half width and the cut's wider side, plus the output's rounding. +pub fn allowBf16(point: f64, half: f64, down: f64, up: f64) f64 { + if (!std.math.isFinite(point)) return 0; + if (!std.math.isFinite(half)) return std.math.floatMax(f64); // scores too far apart for the first order: any value + const h = half + @max(down, up); + return h + 0x1p-8 * (@abs(point) + h) + 0x1p-133; +} + +/// A bf16 output's signed interval: its range's rounded ends, open where float64 or the range is not finite. +pub fn intervalBf16(point: f64, half: f64, down: f64, up: f64) [2]f64 { + if (!std.math.isFinite(point) or !std.math.isFinite(half)) return .{ -std.math.floatMax(f64), std.math.floatMax(f64) }; + return ends(point, half, down, up); +} + /// What a check saw: outputs whose range holds one bf16 value, several, and outputs non-finite in float64. pub const Stats = struct { decided: usize = 0, ranged: usize = 0, nonfinite: usize = 0 }; @@ -175,3 +221,21 @@ test "a weight's spread moves an output by about its share of |v - o|" { head(&.{1}, &.{ 30, 30 }, &.{ 0, 1 }, 2, 1, 1, prompt(256, 2, true), &w, &t, &pt, &hf); try std.testing.expect(hf[0] > small and hf[0] < 0x1p-9); } + +test "the cut widens a range on one side: down by positive values' share, up by negative values'" { + var w: [2]f64 = undefined; + var t: [2]f64 = undefined; + var pt: [2]f64 = undefined; + var hf: [2]f64 = undefined; + var down: [2]f64 = undefined; + var up: [2]f64 = undefined; + var f = prompt(256, 2, false); + f.p_rel = 0; + f.p_cut = operand_cut; + // equal scores, values 2 and 6 in one column, -2 and -6 in the other: o = 4 and -4 + head(&.{ 1, 0 }, &.{ 1, 0, 1, 0 }, &.{ 2, -2, 6, -6 }, 2, 2, 1, f, &w, &t, &pt, &hf); + cutSides(f, &w, &t, &.{ 2, -2, 6, -6 }, 2, 2, &down, &up); + try std.testing.expectEqualSlices(f64, &.{ 4, -4 }, &pt); + try std.testing.expect(down[0] >= 4 * operand_cut and down[0] < 4.01 * operand_cut and up[0] == 0); + try std.testing.expect(up[1] >= 4 * operand_cut and up[1] < 4.01 * operand_cut and down[1] == 0); +} diff --git a/zig/src/core/contract/flashnext_attention_records.zig b/zig/src/core/contract/flashnext_attention_records.zig new file mode 100644 index 0000000000..34eeb1ec53 --- /dev/null +++ b/zig/src/core/contract/flashnext_attention_records.zig @@ -0,0 +1,6 @@ +//! Flash Next's prompt attention on the tensor units, as replay_prompt_segments.zig launches it: 128 threads a group. +const M = @import("limits.zig").MetalLaunch; +pub const records = [_]M{ + .{ .name = "flashnext prompt attention tf_sattn_nax", .threads = 128, .tg_memory = 4 * 16 * 32 * 4, .offline = false }, + .{ .name = "flashnext prompt attention tf_attn256_nax", .threads = 128, .tg_memory = 0, .offline = false }, +}; diff --git a/zig/src/core/contract/gemma_attention_records.zig b/zig/src/core/contract/gemma_attention_records.zig new file mode 100644 index 0000000000..28ea501253 --- /dev/null +++ b/zig/src/core/contract/gemma_attention_records.zig @@ -0,0 +1,11 @@ +//! Gemma 4's attention on the tensor units (metal_gemma4_attention.zig): each source reserves its threads. +const M = @import("limits.zig").MetalLaunch; +// tensor ops before the M5 too (nax.h's simdgroup fragments still multiply through mpp): no offline build +pub const records = [_]M{ + .{ .name = "gemma4 31B attention sliding gemma_prompt_attention", .threads = 128, .tg_memory = 17920, .reserved = 256, .offline = false }, + .{ .name = "gemma4 31B attention global gemma_prompt_attention", .threads = 256, .tg_memory = 18432, .reserved = 256, .offline = false }, + .{ .name = "gemma4 31B attention global gemma_decode_attention", .threads = 128, .tg_memory = 8192, .reserved = 128, .offline = false }, + .{ .name = "gemma4 12B attention sliding gemma_prompt_attention", .threads = 128, .tg_memory = 17920, .reserved = 256, .offline = false }, + .{ .name = "gemma4 12B attention global gemma_prompt_attention", .threads = 256, .tg_memory = 18432, .reserved = 256, .offline = false }, + .{ .name = "gemma4 12B attention global gemma_decode_attention", .threads = 128, .tg_memory = 8192, .reserved = 128, .offline = false }, +}; diff --git a/zig/src/core/contract/nemotron_ssd_records.zig b/zig/src/core/contract/nemotron_ssd_records.zig new file mode 100644 index 0000000000..b0fbc08799 --- /dev/null +++ b/zig/src/core/contract/nemotron_ssd_records.zig @@ -0,0 +1,15 @@ +//! Nemotron's SSD steps on gemm_nax's f32 kernels, as prefill_mixers.mamba launches them through prefill_launch.gemm. +const std = @import("std"); +const M = @import("limits.zig").MetalLaunch; +const gemm_fn = "custom_kernel_tf_gemm_nax_f32_{s}_64_128_256_2_4_float_float_int32_t_float"; + +/// y = sur x dtx and the state dd^T x bh^T by whole blocks; tensor-op kernels whose source reserves 256 threads. +pub const records = blk: { + var out: [5]M = undefined; + for (&out, [_][]const u8{ "n_n_n_n_n", "n_n_t_n_n", "n_n_t_n_t", "t_t_t_t_n", "t_t_t_t_t" }) |*r, tags| r.* = .{ .name = "nemotron ssd gemm " ++ std.fmt.comptimePrint(gemm_fn, .{tags}), .threads = 256, .tg_memory = 0, .reserved = 256 }; + break :blk out; +}; + +test "Nemotron's SSD gemm launches fit every Apple family" { + try std.testing.expectEqual(@as(usize, 0), try @import("registry.zig").checkMetal(std.testing.allocator, &records, true)); +} diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index ac87ea207c..8e8e10c965 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -566,10 +566,11 @@ const paging = @import("paging_pending.zig"); test { _ = paging; _ = @import("minimax_records.zig"); + _ = @import("nemotron_ssd_records.zig"); } /// 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("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("nemotron_pages_records.zig").records ++ @import("nemotron_attention_records.zig").records ++ - @import("moss_image_records.zig").records ++ @import("dsv41_metal_records.zig").records ++ @import("minimax_records.zig").records; + @import("nemotron_ssd_records.zig").records ++ @import("flashnext_attention_records.zig").records ++ @import("spark25_records.zig").records ++ @import("gemma_attention_records.zig").records ++ @import("moss_image_records.zig").records ++ @import("dsv41_metal_records.zig").records ++ @import("minimax_records.zig").records; const learning = @import("registry_learning.zig"); pub const gradient_pending = learning.gradient_pending; pub const qwen_learning_pending = learning.qwen_learning_pending; diff --git a/zig/src/core/contract/spark25_records.zig b/zig/src/core/contract/spark25_records.zig new file mode 100644 index 0000000000..0acb4ae800 --- /dev/null +++ b/zig/src/core/contract/spark25_records.zig @@ -0,0 +1,7 @@ +//! Spark-X2.6-4B's prompt attention launches (metal_spark25_attention.zig): DH / 64 simdgroups, 256 reserved. +const M = @import("limits.zig").MetalLaunch; +pub const records = [_]M{ + // tensor ops before the M5 too (nax.h's simdgroup fragments still multiply through mpp): no offline build + .{ .name = "spark25 attention sliding sp_prompt_attention", .threads = 128, .tg_memory = 17920, .reserved = 256, .offline = false }, + .{ .name = "spark25 attention full sp_prompt_attention", .threads = 128, .tg_memory = 17920, .reserved = 256, .offline = false }, +}; diff --git a/zig/tests/contract/metal_flashnext_attn.zig b/zig/tests/contract/metal_flashnext_attn.zig new file mode 100644 index 0000000000..26e05b33af --- /dev/null +++ b/zig/tests/contract/metal_flashnext_attn.zig @@ -0,0 +1,473 @@ +//! Flash Next's prompt attention on the M5's tensor units (tf_sattn_nax, tf_attn256_nax), as replay_prompt.zig builds. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); +const rig_ = @import("rig.zig"); +const c = contract.spec; +const g = contract.gates; +const ar = contract.attention_range; + +const D = 256; +const H = 24; +const GQA = 12; +const KVH = H / GQA; +const PW = 13952; // the attention projection's row: each head's 256 queries, then its 256 gates +const BK = 32; // both kernels' key blocks +const BQ = 64; // tf_attn256_nax's query rows a threadgroup +const TOP = tf.flashnext_replay.TOP; +const KW = tf.flashnext_replay.KW; +const scale: f64 = 0.0625; // engine_load.zig's attention_scale: 1 / sqrt(head_dim) +const nan16 = 0x7fc0; +/// Gate 4's keys: the first 1,024 score high (their probabilities exactly 1), every other key scores far below. +const stress_on = 1024; + +/// fn_attn.metal's sparse rows (replay_prompt_segments.zig, while the indexer runs) or qmm6_nax_b.metal's causal chunk. +pub const Which = enum { sattn, attn256 }; +/// P cut as the tensor units take it; P spread over [0, 1]; that range without the cut; an edit; the known plants. +pub const Case = enum { cut, spread, cutless, edit, masked }; + +fn name(w: Which, k: Case) []const u8 { + return switch (w) { + inline else => |x| switch (k) { + inline else => |y| comptime label(x, y), + }, + }; +} + +fn label(comptime w: Which, comptime k: Case) []const u8 { + const base = "metal flashnext prompt attention " ++ (if (w == .sattn) "tf_sattn_nax" else "tf_attn256_nax"); + return switch (k) { + .cut => base ++ " 24 heads over 2, head dim 256", + .spread => base ++ ", P spread over [0, 1]", + .cutless => "control: " ++ base ++ ", P spread, its range without the operand cut", + .edit => "control: " ++ base ++ (if (w == .sattn) ", its padding lane counted (key 0 again)" else ", one key past the row seen"), + .masked => base ++ ", non-finite values its masked keys reach", + }; +} + +/// Each kernel's planted control: one source edit gate 3 must catch. +fn edit(w: Which) struct { from: []const u8, to: []const u8 } { + return switch (w) { + .sattn => .{ .from = "s[f][e] = j < n ? v * scale2 : masked;", .to = "s[f][e] = j <= n ? v * scale2 : masked;" }, + .attn256 => .{ .from = "(col >= kL || col > row) ? a_masked", .to = "(col >= kL || col > row + 1) ? a_masked" }, + }; +} + +/// tf_sattn_nax: four simdgroups' 4 mma steps over 64 dims each, then 3 adds; a 32-key block: rescale, 2 P x V steps. +fn sattnForm(n: usize, cut: f64) ar.Form { + const blocks: f64 = @floatFromInt(n / BK + 2); + // den: total times keep and one add a block, plus row_fold's 5-level tree of 32 weights; div: fast 1 / total + return .{ .score = 4 * ar.mma_step + 3, .fast_exp = true, .chunk_exp = false, .p_rel = 0, .p_cut = cut, .num = (2 * ar.mma_step + 1) * blocks, .den = 2 * blocks + 5, .score_round = 4, .div = 6 }; +} + +/// tf_attn256_nax: 16 mma steps over 256 dims; a 32-key block: rescale, 2 P x V steps, two 4-level folds into total. +fn attnForm(n: usize, cut: f64) ar.Form { + const blocks: f64 = @floatFromInt(n / BK + 2); + return .{ .score = 16 * ar.mma_step, .fast_exp = true, .chunk_exp = false, .p_rel = 0, .p_cut = cut, .num = (2 * ar.mma_step + 1) * blocks, .den = 3 * blocks + 4, .score_round = 4, .div = 6 }; +} + +/// P x V's chain at the most keys a row reads (the range's own counts replace it through `bounds`). +fn depth(s: c.Shape) usize { + return (2 * ar.mma_step + 1) * ((s.k + s.rows) / BK + 2); +} + +// rows: a prompt chunk's rows; n: the head dim; k: the chunk's first position +const sattn_served = [_]c.Shape{ .{ .rows = 64, .n = D, .k = 0 }, .{ .rows = 24, .n = D, .k = 2036 }, .{ .rows = 6, .n = D, .k = 30000 } }; +const attn_served = [_]c.Shape{ .{ .rows = 96, .n = D, .k = 0 }, .{ .rows = 40, .n = D, .k = 700 }, .{ .rows = 8, .n = D, .k = 1100 } }; + +pub fn spec(w: Which, k: Case) c.Spec { + const keys: c.Tile = .{ .size = BK, .step = 1, .min = 0 }; + return .{ + .name = name(w, k), + .accum = .f32, + .depth = depth, + .out = .bf16, + // sattn: a threadgroup a row, so rows have no edge; attn256: 64-row query tiles + .tiles = if (w == .sattn) .{ .{}, .{}, keys } else .{ .{ .size = BQ }, .{}, keys }, + // sattn: a chunk from the start, one crossing into selected blocks, one deep in the context; attn256 dense + .served = if (w == .sattn) &sattn_served else &attn_served, + .sums = .{ .exact = true }, + .known = if (k == .masked) .{ .gate = .nonfinite, .finding = finding(w) } else null, + }; +} + +/// Each kernel's gate-1 finding: masked keys score -FLT_MAX, where its running max starts (docs: Known findings). +fn finding(w: Which) []const u8 { + return switch (w) { + .sattn => "a non-finite value in key 0's row reaches every row through its padding lanes (key 0 at p = 0: 0 x Inf = NaN), " ++ + "and a row whose every score is -Inf weighs its padding lanes 1 (its max starts at -FLT_MAX), so it gives key 0's values where float64 gives NaN", + .attn256 => "a non-finite value at a key masked for a row but read in its 32-key block reaches that row as NaN (P = 0 times Inf inside the P x V mma), " ++ + "and a row whose every score is -Inf weighs its masked keys 1 (its max starts at -FLT_MAX), so it averages them where float64 gives NaN", + }; +} + +fn bf(v: f64) f64 { + return c.fromBf16(c.toBf16(@floatCast(v))); +} + +/// bsig's range: Metal's fast exp (attention_range's 40 + 2 |x| ulps), the add, fast division (2.5 ulps), bf16. +fn sigRange(x: f64) [2]f64 { + const e = @exp(-x); + const ee = (40 + 2 * @abs(x)) * 0x1p-23; + const lo = (1 + e * (1 - ee)) * (1 - 0x1p-24); + const hi = (1 + e * (1 + ee)) * (1 + 0x1p-24); + return .{ bf(1 / hi * (1 - 2.5 * 0x1p-23)), bf(1 / lo * (1 + 2.5 * 0x1p-23)) }; +} + +/// bf16(o) times bsig, rounded: the four corners of the two ranges; bf16 times bf16 is exact in fp32 but below normals. +fn gated(o: [2]f64, s: [2]f64) [2]f64 { + const p = [4]f64{ o[0] * s[0], o[0] * s[1], o[1] * s[0], o[1] * s[1] }; + const lo = @min(@min(p[0], p[1]), @min(p[2], p[3])); + const hi = @max(@max(p[0], p[1]), @max(p[2], p[3])); + return .{ bf(lo - 0x1p-149), bf(hi + 0x1p-149) }; +} + +fn fillRandom(rng: std.Random, out: []u16, r: f32) void { + for (out) |*v| v.* = c.toBf16((rng.float(f32) * 2 - 1) * r); +} + +pub const Probe = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + which: Which, + case: Case, + rows: usize = 0, + pos: usize = 0, + cap: usize = 0, + fill: g.Fill = .random, + q: []u16 = &.{}, + k: []u16 = &.{}, + v: []u16 = &.{}, + gp: []u16 = &.{}, + ids: []i32 = &.{}, + nk: []i32 = &.{}, + sparse: []i32 = &.{}, + lo: []f64 = &.{}, + hi: []f64 = &.{}, + refs: [8]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + /// The library as replay_prompt.zig builds it here (frags.source, MLX's options), the control's edit if asked. + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, w: Which, k: Case) !Probe { + const src = try tf.frags.source(device, a, if (w == .sattn) sources.flashnext_attn else sources.flashnext_qmm6); + defer a.free(src); + var text = src; + if (k == .edit) { + const e = edit(w); + if (std.mem.count(u8, src, e.from) != 1) return error.ControlEditMissing; + text = try std.mem.replaceOwned(u8, a, src, e.from, e.to); + } + defer if (k == .edit) a.free(text); + const pipe = try rig_.pipeline(device, text, if (w == .sattn) "tf_sattn_nax" else "tf_attn256_nax"); + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .pipe = pipe, .which = w, .case = k }; + } + + fn free(p: *Probe) void { + for ([_]*[]u16{ &p.q, &p.k, &p.v, &p.gp }) |x| { + p.a.free(x.*); + x.* = &.{}; + } + for ([_]*[]i32{ &p.ids, &p.nk, &p.sparse }) |x| { + p.a.free(x.*); + x.* = &.{}; + } + for ([_]*[]f64{ &p.lo, &p.hi }) |x| { + p.a.free(x.*); + x.* = &.{}; + } + } + + pub fn deinit(p: *Probe) void { + p.free(); + p.rig.deinit(); + p.pipe.deinit(); + } + + fn at(p: *const Probe, kv: usize, key: usize) usize { + return (kv * p.cap + key) * D; + } + + /// Select.meta's rows: dense while a row's complete 4-key blocks are at most TOP, else TOP blocks and the tail. + fn select(p: *Probe, rng: std.Random) !void { + const blocks = try p.a.alloc(bool, (p.pos + p.rows) / 4 + 1); + defer p.a.free(blocks); + for (0..p.rows) |r| { + const e = p.pos + r + 1; + const cb = e / 4; + const ids = p.ids[r * KW ..][0..KW]; + @memset(ids, @intCast(p.pos + p.rows)); // a dense row never reads its ids: they name a NaN cache row here + p.sparse[r] = @intFromBool(cb > TOP); + if (cb <= TOP) { + p.nk[r] = @intCast(e); + continue; + } + @memset(blocks, false); + if (p.fill == .stress) @memset(blocks[0..TOP], true) else { + var kept: usize = 0; + while (kept < TOP) { + const b = rng.uintLessThan(usize, cb); + if (blocks[b]) continue; + blocks[b] = true; + kept += 1; + } + } + // the chosen complete blocks in ascending order, then the tail's keys + var n: usize = 0; + for (blocks[0..cb], 0..) |on, b| if (on) for (0..4) |i| { + ids[n] = @intCast(4 * b + i); + n += 1; + }; + for (4 * cb..e) |key| { + ids[n] = @intCast(key); + n += 1; + } + p.nk[r] = @intCast(n); + } + } + + /// Scores q . k = q0 k0 alone, so P = e^(s - max) spreads over about [2^-4, 1]; values in [0, 1) add the cut up. + fn spread(p: *Probe, rng: std.Random) void { + for (0..p.rows * H) |r| { + @memset(p.q[r * D ..][0..D], 0); + p.q[r * D] = c.toBf16(0.5 + rng.float(f32)); + } + for (0..KVH) |kv| for (0..p.pos + p.rows) |key| { + p.k[p.at(kv, key)] = c.toBf16(-32 * rng.float(f32)); + for (p.v[p.at(kv, key)..][0..D]) |*x| x.* = c.toBf16(rng.float(f32)); + }; + @memset(p.gp, c.toBf16(20)); // bsig(20) is 1: the outputs show bf16(o) itself + } + + /// Plants both agree on: NaN at key 0; infinities at a key no row reads unseen (sattn: key 1); NaN queries and key. + fn plants(p: *Probe) void { + const inf_key: usize = if (p.which == .sattn) 1 else 0; + for (0..KVH) |kv| { + p.v[p.at(kv, 0)] = nan16; + p.v[p.at(kv, inf_key) + 1] = 0x7f80; + p.v[p.at(kv, inf_key) + 2] = 0xff80; + } + for (1..4) |r| if (r < p.rows) { + p.q[(r * H + 5 * r) * D + r] = nan16; + }; + p.k[p.at(1, 2) + 7] = nan16; + } + + /// The findings' plants: infinities in key 0's values and at key 5 (masked for rows 0 to 4), a row of -Inf scores. + fn knownPlants(p: *Probe) void { + p.plants(); + for (0..KVH) |kv| { + p.v[p.at(kv, 0) + 3] = 0x7f80; + p.v[p.at(kv, 5) + 4] = 0x7f80; + } + // row 4's head 0 scores -Inf at every key it sees: q +Inf against keys 0 to 4, each -1 in that dim + p.q[(4 * H + 0) * D + 9] = 0x7f80; + for (0..5) |key| p.k[p.at(0, key) + 9] = c.toBf16(-1); + } + + pub fn prepare(p: *Probe, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n != D) return false; + if (fill == .stress and shape.k < stress_on) return false; + p.rig.reset(); + p.free(); + p.rows = shape.rows; + p.pos = shape.k; + p.fill = fill; + const keys = p.pos + p.rows; + // sattn reads only its rows' keys (one NaN row past them); attn256 reads whole 32-key blocks + p.cap = if (p.which == .sattn) keys + 1 else std.mem.alignForward(usize, keys, BK); + const pad = if (p.which == .sattn) p.rows else std.mem.alignForward(usize, p.rows, BQ); + p.q = try p.a.alloc(u16, pad * H * D); + p.k = try p.a.alloc(u16, KVH * p.cap * D); + p.v = try p.a.alloc(u16, KVH * p.cap * D); + p.gp = try p.a.alloc(u16, p.rows * PW); + p.ids = try p.a.alloc(i32, p.rows * KW); + p.nk = try p.a.alloc(i32, p.rows); + p.sparse = try p.a.alloc(i32, p.rows); + p.lo = try p.a.alloc(f64, p.count()); + p.hi = try p.a.alloc(f64, p.count()); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + fillRandom(rng, p.q, 2); + fillRandom(rng, p.k, 2); + fillRandom(rng, p.v, 2); + fillRandom(rng, p.gp, 4); + switch (fill) { + .random => if (p.case == .spread or p.case == .cutless) p.spread(rng), + .nonfinite => if (p.case == .masked) p.knownPlants() else p.plants(), + .stress => { + // keys of 3s or -3s: the first 1,024 score high, the rest probability 0; values cancel but for 8 nudged + @memset(p.q, c.toBf16(3)); + for (0..KVH) |kv| for (0..keys) |key| @memset(p.k[p.at(kv, key)..][0..D], c.toBf16(if (key < stress_on) 3 else -3)); + for (0..KVH) |kv| for (0..stress_on / 2) |key| for (0..D) |e| { + const x: f32 = @floatFromInt(1 + rng.uintLessThan(u32, 255)); + p.v[p.at(kv, key) + e] = c.toBf16(x); + p.v[p.at(kv, key + stress_on / 2) + e] = c.toBf16(-x + @as(f32, @floatFromInt(@intFromBool(key < 8)))); + }; + @memset(p.gp, c.toBf16(20)); + }, + } + if (p.which == .sattn) { + // cache rows past the chunk hold anything: NaN here, which tf_sattn_nax never reads + for (0..KVH) |kv| for ([_][]u16{ p.k, p.v }) |x| @memset(x[p.at(kv, keys)..][0..((p.cap - keys) * D)], nan16); + try p.select(rng); + } else { + // query rows past the chunk are read into tiles whose outputs are never written: NaN here + @memset(p.q[p.rows * H * D ..], nan16); + } + try p.bind(); + return true; + } + + fn bind(p: *Probe) !void { + const f = [1]f32{@floatCast(scale)}; + const list = if (p.which == .sattn) + [_][]const u8{ std.mem.sliceAsBytes(p.q), std.mem.sliceAsBytes(p.k), std.mem.sliceAsBytes(p.v), std.mem.sliceAsBytes(p.ids), std.mem.sliceAsBytes(p.nk), std.mem.sliceAsBytes(p.sparse), std.mem.sliceAsBytes(p.gp), std.mem.asBytes(&f) } + else + [_][]const u8{ std.mem.sliceAsBytes(p.q), std.mem.sliceAsBytes(p.k), std.mem.sliceAsBytes(p.v), std.mem.sliceAsBytes(p.gp), std.mem.asBytes(&f), "", "", "" }; + for (list, &p.refs) |bytes, *r| if (bytes.len > 0) { + r.* = try p.rig.input(bytes); + }; + p.out = try p.rig.output(p.count() * 2); + } + + /// One launch as replay_prompt_segments.zig encodes it. + pub fn run(p: *Probe) !void { + const ps = p.rig.begin(); + ps.e.setPipeline(p.pipe); + if (p.which == .sattn) { + for (p.refs[0..8], 0..) |r, i| ps.e.setBuffer(r.buf, r.off, i); + const cap: i32 = @intCast(p.cap); + ps.e.setBytes(std.mem.asBytes(&cap), 8); + ps.e.setBuffer(p.out.buf, p.out.off, 9); + ps.e.dispatchThreads(mtl.Size.of(p.rows * 128, KVH, 1), mtl.Size.of(128, 1, 1)); + } else { + for (p.refs[0..4], 0..) |r, i| ps.e.setBuffer(r.buf, r.off, i); + const ap = [4]i32{ @intCast(p.rows), @intCast(p.pos + p.rows), @intCast(p.pos), @intCast(p.cap) }; + ps.e.setBytes(std.mem.asBytes(&ap), 4); + ps.e.setBuffer(p.refs[4].buf, p.refs[4].off, 5); + ps.e.setBuffer(p.out.buf, p.out.off, 6); + ps.e.dispatchThreads(mtl.Size.of(((p.rows + BQ - 1) / BQ) * 128, H, 1), mtl.Size.of(128, 1, 1)); + } + try p.rig.finish(ps); + } + + pub fn count(p: *const Probe) usize { + return p.rows * H * D; + } + + /// Row r's keys: a sparse row's ids, else every position up to the row's own. + fn rowKeys(p: *const Probe, r: usize, out: []usize) []usize { + const n: usize = if (p.which == .sattn) @intCast(p.nk[r]) else p.pos + r + 1; + for (out[0..n], 0..) |*x, j| x.* = if (p.which == .sattn and p.sparse[r] != 0) @intCast(p.ids[r * KW + j]) else j; + return out[0..n]; + } + + /// Each head's float64 output at bf16(bf16(o) * bsig(gate)), and its range: o's with the cut, then the gate's. + pub fn reference(p: *Probe, want: []f64, mag: []f64) !void { + const most = p.pos + p.rows; + const buf = try p.a.alloc(f64, 2 * most * D + 2 * most + 5 * D); + defer p.a.free(buf); + const keys = try p.a.alloc(usize, most); + defer p.a.free(keys); + const kf = buf[0 .. most * D]; + const vf = buf[most * D .. 2 * most * D]; + const w = buf[2 * most * D ..][0..most]; + const tau = buf[2 * most * D + most ..][0..most]; + const head = buf[2 * most * D + 2 * most ..]; + const qv = head[0..D]; + const pt = head[D .. 2 * D]; + const hf = head[2 * D .. 3 * D]; + const down = head[3 * D .. 4 * D]; + const up = head[4 * D .. 5 * D]; + const cut: f64 = if (p.case == .cutless) 0 else ar.operand_cut; + for (0..p.rows) |r| { + const ks = p.rowKeys(r, keys); + const n = ks.len; + const f = if (p.which == .sattn) sattnForm(n, cut) else attnForm(n, cut); + for (0..KVH) |kv| { + for (ks, 0..) |key, j| for (0..D) |e| { + kf[j * D + e] = c.fromBf16(p.k[p.at(kv, key) + e]); + vf[j * D + e] = c.fromBf16(p.v[p.at(kv, key) + e]); + }; + for (kv * GQA..(kv + 1) * GQA) |h| { + for (qv, 0..) |*x, e| x.* = c.fromBf16(p.q[(r * H + h) * D + e]); + ar.head(qv, kf, vf, n, D, scale, f, w, tau, pt, hf); + ar.cutSides(f, w, tau, vf, n, D, down, up); + for (0..D) |e| { + const i = (r * H + h) * D + e; + const gv: f64 = c.fromBf16(p.gp[r * PW + h * 2 * D + D + e]); + want[i] = bf(bf(pt[e]) * bf(1 / (1 + @exp(-gv)))); + mag[i] = if (std.math.isFinite(want[i])) @abs(want[i]) else 0; + const open = p.fill == .stress or !std.math.isFinite(pt[e]) or !std.math.isFinite(hf[e]); + const o = ar.ends(pt[e], hf[e], down[e], up[e]); + const ends = if (open) [2]f64{ -std.math.floatMax(f64), std.math.floatMax(f64) } else gated(o, sigRange(gv)); + p.lo[i] = ends[0]; + p.hi[i] = ends[1]; + } + } + } + } + } + + /// Each output's farther end from float64's value; gate 4 allows the rounding alone, a wide range any finite value. + pub fn bounds(p: *Probe, want: []const f64, _: []const f64, allow: []f64) !void { + const s = spec(p.which, p.case); + for (allow, want, p.lo, p.hi) |*l, w, lo, hi| { + if (p.fill == .stress) { + l.* = c.boundExact(s, w); + } else if (!std.math.isFinite(w)) { + l.* = 0; + } else l.* = @min(std.math.floatMax(f64), @max(@abs(w - lo), @abs(hi - w))); + } + } + + /// Each bf16 output within its range's ends: o's with the cut on its own side, through bf16 and the rounded gate. + pub fn intervals(p: *Probe, low: []f64, high: []f64) !void { + @memcpy(low, p.lo); + @memcpy(high, p.hi); + } + + /// Gate 4: the outputs if P x V's sum of the 1,024 values ran in `acc` (bsig is 1). + pub fn simulate(p: *Probe, acc: c.Accum, order: g.Order, out: []f64) !void { + var terms: [stress_on]f64 = undefined; + for (0..p.rows) |r| for (0..H) |h| for (0..D) |e| { + for (&terms, 0..) |*x, key| x.* = c.fromBf16(p.v[p.at(h / GQA, key) + e]); + out[(r * H + h) * D + e] = bf(contract.rows.sumIn(acc, order, &terms) / stress_on); + }; + } +}; + +const vt = g.hooks(Probe); + +pub const Suite = struct { + a: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + kernels: *std.ArrayList(g.Kernel), + probes: std.ArrayList(*Probe) = .empty, + + /// Only on the tensor units: the operand cut was measured there alone, so these cases register there. + pub fn collect(s: *Suite, only: ?[]const u8) !void { + if (!s.device.tensorUnits()) return; + for (std.enums.values(Which)) |w| for (std.enums.values(Case)) |k| { + if (only) |o| if (std.mem.indexOf(u8, name(w, k), o) == null) continue; + const p = try s.a.create(Probe); + errdefer s.a.destroy(p); + p.* = try Probe.init(s.a, s.device, s.queue, w, k); + try s.probes.append(s.a, p); + const control: ?c.Gate = if (k == .cutless or k == .edit) .shapes else null; + try s.kernels.append(s.a, .{ .spec = spec(w, k), .ctx = p, .vt = &vt, .control = control }); + }; + } + + pub fn deinit(s: *Suite) void { + for (s.probes.items) |p| { + p.deinit(); + s.a.destroy(p); + } + s.probes.deinit(s.a); + } +}; diff --git a/zig/tests/contract/metal_gemma4_attention.zig b/zig/tests/contract/metal_gemma4_attention.zig new file mode 100644 index 0000000000..7440584e51 --- /dev/null +++ b/zig/tests/contract/metal_gemma4_attention.zig @@ -0,0 +1,521 @@ +//! Gemma 4's prompt and decode attention on the M5's tensor units (prompt_attention.metal), as kernels.zig builds them. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); +const rig_ = @import("rig.zig"); +const c = contract.spec; +const g = contract.gates; +const ar = contract.attention_range; +const cfg = tf.gemma4.config; +const ring = tf.gemma4.state.ring; +const nan16 = 0x7fc0; +const log2e = 1.4426950408889634; +/// gemma_decode_attention's CD (keys a threadgroup) and BK (keys a step); the prompt kernel's RT (rows a tile). +const cd = 256; +const bk = 32; +const rt = 32; +/// Rows a prompt chunk takes (state.chunk_rows) and a decode call (forward.zig's attend_rows). +const prompt_rows = tf.gemma4.state.chunk_rows; +const decode_rows = tf.gemma4.state.max_rows; + +/// A layer kind's gemma_prompt_attention, or a global layer's gemma_decode_attention then gemma_attention_merge. +const Kind = enum { prompt, decode }; +/// Served; P spread over [0, 1], values in [0, 1); that range uncut; scores at 2 log2e; +Inf at a masked key. +const Case = enum { plain, spread, cutless, edit, masked }; + +const Variant = struct { geo: cfg.Geometry, global: bool, kind: Kind, case: Case = .plain }; + +/// Every pipeline the engine launches (the 12B's assistant compiles the 12B's global one), then the cases and controls. +const variants = [_]Variant{ + .{ .geo = cfg.dense31b, .global = false, .kind = .prompt }, + .{ .geo = cfg.dense31b, .global = true, .kind = .prompt }, + .{ .geo = cfg.unified12b, .global = false, .kind = .prompt }, + .{ .geo = cfg.unified12b, .global = true, .kind = .prompt }, + .{ .geo = cfg.dense31b, .global = true, .kind = .decode }, + .{ .geo = cfg.unified12b, .global = true, .kind = .decode }, + .{ .geo = cfg.dense31b, .global = true, .kind = .prompt, .case = .spread }, + .{ .geo = cfg.dense31b, .global = true, .kind = .prompt, .case = .cutless }, + .{ .geo = cfg.dense31b, .global = true, .kind = .decode, .case = .spread }, + .{ .geo = cfg.dense31b, .global = true, .kind = .decode, .case = .cutless }, + .{ .geo = cfg.dense31b, .global = false, .kind = .prompt, .case = .edit }, + .{ .geo = cfg.dense31b, .global = true, .kind = .decode, .case = .edit }, + .{ .geo = cfg.dense31b, .global = false, .kind = .prompt, .case = .masked }, + .{ .geo = cfg.dense31b, .global = true, .kind = .prompt, .case = .masked }, + .{ .geo = cfg.unified12b, .global = false, .kind = .prompt, .case = .masked }, + .{ .geo = cfg.unified12b, .global = true, .kind = .prompt, .case = .masked }, + .{ .geo = cfg.dense31b, .global = true, .kind = .decode, .case = .masked }, + .{ .geo = cfg.unified12b, .global = true, .kind = .decode, .case = .masked }, +}; + +fn nameOf(comptime v: Variant) []const u8 { + const what = if (v.kind == .decode) " decode attention global +merge" else if (v.global) " prompt attention global" else " prompt attention sliding"; + const base = "metal gemma4 " ++ v.geo.name ++ what; + return switch (v.case) { + .plain => base, + .spread => base ++ ", P spread over [0, 1]", + .cutless => "control: " ++ base ++ ", P spread, its range without the operand cut", + .edit => "control: " ++ base ++ ", scores at 2 log2e", + .masked => base ++ (if (v.kind == .decode) ", a non-finite value at a key masked in its step" else ", a non-finite value at a key masked in its tile"), + }; +} + +const names = blk: { + var out: [variants.len][]const u8 = undefined; + for (variants, &out) |v, *n| n.* = nameOf(v); + break :blk out; +}; + +/// Each kernel's planted control: one source edit gate 3 must catch. +fn edit(k: Kind) [2][]const u8 { + return if (k == .prompt) .{ "s[i][e] * kLog2e", "s[i][e] * (2.0f * kLog2e)" } else .{ "s[f][e] * 1.44269504089f", "s[f][e] * 2.88539008178f" }; +} + +/// kernels.zig's header for prompt_attention.metal at a layer kind. +fn header(a: std.mem.Allocator, geo: cfg.Geometry, global: bool) ![]u8 { + const dh: usize = if (global) cfg.global_dim else cfg.sliding_dim; + const nk: usize = if (global) geo.global_kv_heads else geo.sliding_kv_heads; + return std.fmt.allocPrint(a, "#define DH {d}\n#define NQ {d}\n#define G {d}\n#define SL {d}\n#define WINDOW {d}\n#define RING {d}\n", .{ dh, geo.query_heads, geo.query_heads / nk, dh / 128, @as(usize, if (global) 0 else cfg.window), @as(usize, if (global) 0 else ring) }); +} + +/// The library as kernels.zig builds it on this GPU: the header, then frags.source's text. +fn library(a: std.mem.Allocator, device: mtl.Device, geo: cfg.Geometry, global: bool) ![]u8 { + const h = try header(a, geo, global); + defer a.free(h); + const body = try tf.frags.source(device, a, sources.gemma4_prompt_attention); + defer a.free(body); + return std.mem.concat(a, u8, &.{ h, body }); +} + +/// Gate 5's texts for a record named "gemma4 attention ": here, and before the M5. +pub fn limitsText(a: std.mem.Allocator, device: mtl.Device, name: []const u8) !struct { here: []const u8, pre_m5: []const u8 } { + const geo = if (std.mem.indexOf(u8, name, " 31B ") != null) cfg.dense31b else if (std.mem.indexOf(u8, name, " 12B ") != null) cfg.unified12b else return error.UnknownLaunch; + const global = std.mem.indexOf(u8, name, " global ") != null; + const h = try header(a, geo, global); + defer a.free(h); + const body = try std.mem.replaceOwned(u8, a, sources.gemma4_prompt_attention, "#include \"../nax.h\"", sources.nax); + defer a.free(body); + const pre = try std.mem.concat(a, u8, &.{ h, "#define TF_SIMD_FRAGS 1\n", body }); + return .{ .here = if (device.tensorUnits()) try std.mem.concat(a, u8, &.{ h, body }) else pre, .pre_m5 = pre }; +} + +fn bf(v: f64) f64 { + return c.fromBf16(c.toBf16(@floatCast(v))); +} + +/// Gate 4's value for key `j` of `n`: big halves that cancel, four ones left, so only exact fp32 sums give 4. +fn stressValue(j: usize, n: usize) f32 { + const half = n / 2; + if (j >= 2 * half) return 0; + const at = if (j < half) j else j - half; + const big: f32 = if (at % 2 == 0) 127 else 16256; + const one: f32 = if (at < 8 and at % 2 == 0) 1 else 0; + return if (j < half) big + one else -big; +} + +/// gemma_prompt_attention over `blocks` blocks of `dh / 2` keys, its fp32 P cut by `cut` as it meets V. +fn promptForm(dh: usize, blocks: usize, cut: f64) ar.Form { + const b: f64 = @floatFromInt(blocks); + const ns: f64 = @floatFromInt(dh / 64); + return .{ + .score = @as(f64, @floatFromInt(dh / 16)) * ar.mma_step, // one accumulator over DH / 16 mma steps + .fast_exp = false, + .chunk_exp = false, + .p_rel = 0, + .p_cut = cut, + .num = (@as(f64, @floatFromInt(dh / 32)) * ar.mma_step + 1) * b, // DH / 32 P x V steps and a rescale a block + .den = 3 * b + 3 + (ns - 1), // a 16-key fold (4 levels, 2 adds) a block, its rescale, then NS simdgroups' sums + .score_round = 2, // s * kLog2e: the constant's rounding and the product's + .div = 4, // fast 1 / t (2.5 ulps), then acc * inv + .exp_ulps = .{ 4 * b, 2 * log2e + 0.5 }, // fast::exp2 (4 + 2 |x| ulps) of a weight and its block's rescales + }; +} + +/// gemma_decode_attention's 256-key chunks (8 steps of 32 keys over 4 128-dim slices), then the merge over `chunks`. +fn decodeForm(chunks: usize, cut: f64) ar.Form { + const n: f64 = @floatFromInt(chunks); + const steps: f64 = cd / bk; + return .{ + .score = 8 * ar.mma_step + 3, // 8 mma steps a slice, its 4 slices added in order + .fast_exp = false, + .chunk_exp = false, + .p_rel = 0, + .p_cut = cut, + .num = (2 * ar.mma_step + 1) * steps + n, // 2 P x V steps and a rescale a step, then the merge's fma chain + .den = 3 * steps + 3 + n, // two 16-key folds and a rescale a step, then the merge's fma chain + .score_round = 2 + 4, // s * log2e (2), and PM = top * ln2 then the merge's max, each 2 of |w| + x + .div = 3, // the merge's fast acc / lsum + // 8 fast::exp2s (4 + 2 |x| ulps), the merge's fast::exp (4 + 2 |x|), PM's and the merge's x roundings + .exp_ulps = .{ 4 * steps + 4, 2 * log2e + 0.5 + 2 + 0.5 + 2 }, + }; +} + +const prompt_sliding = [_]c.Shape{ .{ .rows = 77, .n = cfg.sliding_dim, .k = 0 }, .{ .rows = 64, .n = cfg.sliding_dim, .k = 1000 }, .{ .rows = 100, .n = cfg.sliding_dim, .k = 2000 }, .{ .rows = prompt_rows, .n = cfg.sliding_dim, .k = 0 } }; +const prompt_global = [_]c.Shape{ .{ .rows = 77, .n = cfg.global_dim, .k = 0 }, .{ .rows = 64, .n = cfg.global_dim, .k = 700 }, .{ .rows = 33, .n = cfg.global_dim, .k = 1500 }, .{ .rows = prompt_rows, .n = cfg.global_dim, .k = 0 } }; +const decode_global = [_]c.Shape{ .{ .rows = decode_rows, .n = cfg.global_dim, .k = 700 }, .{ .rows = 1, .n = cfg.global_dim, .k = 0 }, .{ .rows = 5, .n = cfg.global_dim, .k = 255 }, .{ .rows = 9, .n = cfg.global_dim, .k = 1100 }, .{ .rows = 1, .n = cfg.global_dim, .k = 3000 } }; + +/// Gate 1's known findings: a masked key's value still enters P x V, where P = 0 times Inf is NaN. +const prompt_finding = "a non-finite value at a key masked for a row of its tile reaches that row as NaN (P = 0 times Inf inside the P x V mma)"; +const decode_finding = "a non-finite value at a key masked for a decode row but in its 32-key step reaches that row as NaN (P = 0 times Inf inside the P x V mma)"; + +fn keysDepth(s: c.Shape) usize { + return s.n + s.k + s.rows; +} + +const Probe = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + v: Variant, + name: []const u8, + pipe: mtl.Pipeline, + merge: ?mtl.Pipeline = null, + rows: usize = 0, + p0: usize = 0, + cap: usize = 0, + stress: bool = false, + kp: []u16 = &.{}, // keys and values by position [kv heads, positions, dh] + vp: []u16 = &.{}, + q: []u16 = &.{}, // [rows (prompt: whole tiles), query heads, dh] + hi: []i32 = &.{}, // each prompt row's last key (an image's rows see the whole image) + point: []f64 = &.{}, + half: []f64 = &.{}, + down: []f64 = &.{}, + up: []f64 = &.{}, + refs: [5]rig_.Ref = undefined, + parts: [4]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, v: Variant, name: []const u8) !Probe { + const text = try library(a, device, v.geo, v.global); + defer a.free(text); + const src = if (v.case != .edit) try a.dupe(u8, text) else blk: { + const e = edit(v.kind); + if (std.mem.count(u8, text, e[0]) != 1) return error.ControlEditMissing; + break :blk try std.mem.replaceOwned(u8, a, text, e[0], e[1]); + }; + defer a.free(src); + const fun = if (v.kind == .prompt) "gemma_prompt_attention" else "gemma_decode_attention"; + var p: Probe = .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .v = v, .name = name, .pipe = try rig_.pipeline(device, src, fun) }; + errdefer p.pipe.deinit(); + if (v.kind == .decode) { + const glue = try tf.gemma4.kernels.glueSource(a, v.geo, true, false, 4); + defer a.free(glue); + p.merge = try rig_.pipeline(device, glue, "gemma_attention_merge"); + } + return p; + } + + fn free(p: *Probe) void { + for ([_][]u16{ p.kp, p.vp, p.q }) |x| p.a.free(x); + p.a.free(p.hi); + for ([_][]f64{ p.point, p.half, p.down, p.up }) |x| p.a.free(x); + p.kp = &.{}; + p.vp = &.{}; + p.q = &.{}; + p.hi = &.{}; + p.point = &.{}; + p.half = &.{}; + p.down = &.{}; + p.up = &.{}; + } + + fn deinit(p: *Probe) void { + p.free(); + p.rig.deinit(); + p.pipe.deinit(); + if (p.merge) |m| m.deinit(); + } + + fn dh(p: *const Probe) usize { + return if (p.v.global) cfg.global_dim else cfg.sliding_dim; + } + fn nq(p: *const Probe) usize { + return p.v.geo.query_heads; + } + fn nk(p: *const Probe) usize { + return if (p.v.global) p.v.geo.global_kv_heads else p.v.geo.sliding_kv_heads; + } + fn npos(p: *const Probe) usize { + return p.p0 + p.rows; + } + fn chunks(p: *const Probe) usize { + return (p.npos() + cd - 1) / cd; + } + + fn spec(p: *const Probe) c.Spec { + const served: []const c.Shape = if (p.v.kind == .decode) &decode_global else if (p.v.global) &prompt_global else &prompt_sliding; + const rows: c.Tile = if (p.v.kind == .decode) .{ .size = 16 / (p.nq() / p.nk()) } else .{ .size = rt }; + const keys: c.Tile = .{ .size = if (p.v.kind == .decode) cd else p.dh() / 2, .min = 0 }; + // ops: gate 4's exact sums meet only the final division's roundings + var s: c.Spec = .{ .name = p.name, .accum = .f32, .depth = keysDepth, .ops = 4, .out = .bf16, .tiles = .{ rows, .{}, keys }, .served = served }; + if (p.v.case == .masked) { + const at: usize = if (p.v.kind == .prompt) 1 else 0; // one served shape past key 0, for gate 4's cancelling sums + s.served = served[at..][0..1]; + s.known = .{ .gate = .nonfinite, .finding = if (p.v.kind == .prompt) prompt_finding else decode_finding }; + } + return s; + } + + /// The keys row `r` attends: a sliding row's window up to its last key, else every key up to its own. + fn keysOf(p: *const Probe, r: usize) [2]usize { + const pos = p.p0 + r; + if (p.v.global) return .{ 0, pos }; + return .{ (pos + 1) -| cfg.window, @intCast(p.hi[r]) }; + } + + /// Row `r`'s form: its prompt tile's blocks (the tile's first window start to its last row's last key), or chunks. + fn form(p: *const Probe, r: usize) ar.Form { + const cut: f64 = if (p.v.case == .cutless) 0 else ar.operand_cut; + if (p.v.kind == .decode) return decodeForm(p.chunks(), cut); + const r0 = r / rt * rt; + const last: usize = @intCast(p.hi[@min(r0 + rt, p.rows) - 1]); + const lo = if (p.v.global) 0 else (p.p0 + r0 + 1) -| cfg.window; + return promptForm(p.dh(), last / (p.dh() / 2) - lo / (p.dh() / 2) + 1, cut); + } + + pub fn prepare(p: *Probe, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + const most: usize = if (p.v.kind == .decode) decode_rows else prompt_rows; + if (shape.rows == 0 or shape.rows > most or shape.n != p.dh()) return false; + p.rig.reset(); + p.free(); + p.rows = shape.rows; + p.p0 = shape.k; + p.stress = fill == .stress; + const dh_ = p.dh(); + const n = p.npos(); + const pad = if (p.v.kind == .prompt) std.mem.alignForward(usize, p.rows, rt) else p.rows; + p.kp = try p.a.alloc(u16, p.nk() * n * dh_); + p.vp = try p.a.alloc(u16, p.nk() * n * dh_); + p.q = try p.a.alloc(u16, pad * p.nq() * dh_); + p.hi = try p.a.alloc(i32, p.rows); + for ([_]*[]f64{ &p.point, &p.half, &p.down, &p.up }) |x| x.* = try p.a.alloc(f64, p.count()); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for ([_][]u16{ p.kp, p.vp, p.q }) |x| for (x) |*y| { + y.* = c.toBf16(rng.float(f32) * 2 - 1); + }; + for (p.hi, 0..) |*h, r| h.* = @intCast(p.p0 + r); + if (!p.v.global and p.rows > 16) { // an image's rows: positions [at, at + len) all see the whole image + const at = p.rows / 4; + @memset(p.hi[at..][0 .. p.rows / 3], @intCast(p.p0 + at + p.rows / 3 - 1)); + } + switch (fill) { + .random => if (p.v.case == .spread or p.v.case == .cutless) p.spread(rng), + .nonfinite => { + // keys at p0 every row sees; a NaN key at p0 + 2 that rows 0 and 1 mask + p.q[(0 * p.nq() + 1) * dh_ + 3] = c.toBf16(std.math.nan(f32)); + p.vp[p.p0 * dh_ + 5] = c.toBf16(std.math.inf(f32)); + p.vp[p.p0 * dh_ + 6] = c.toBf16(-std.math.inf(f32)); + if (p.p0 + 2 < n) p.kp[(p.p0 + 2) * dh_] = c.toBf16(std.math.nan(f32)); + // the masked case: +Inf at p0 + 1, masked for row 0, read in its tile (prompt) or 32-key step (decode) + if (p.v.case == .masked) p.vp[(p.p0 + 1) * dh_ + 7] = c.toBf16(std.math.inf(f32)); + }, + .stress => { + // every score 0, so every weight 1; the values every row sees cancel but for four ones + @memset(p.q, 0); + @memset(p.kp, 0); + @memset(p.vp, 0); + const first = if (p.v.global) 0 else (p.p0 + p.rows) -| cfg.window; + for (0..p.nk()) |h| for (first..p.p0 + 1) |j| for (0..dh_) |e| { + p.vp[(h * n + j) * dh_ + e] = c.toBf16(stressValue(j - first, p.p0 + 1 - first)); + }; + }, + } + @memset(p.q[p.rows * p.nq() * dh_ ..], nan16); // a tile's rows past the chunk: read, never written + try p.bind(); + return true; + } + + /// Scores q0 k0 alone, so P = e^(s - max) spreads over about [2^-6, 1]; values in [0, 1) add the cut's bias up. + fn spread(p: *Probe, rng: std.Random) void { + const dh_ = p.dh(); + for (0..p.rows * p.nq()) |r| { + @memset(p.q[r * dh_ ..][0..dh_], 0); + p.q[r * dh_] = c.toBf16(0.5 + rng.float(f32)); + } + for (0..p.nk() * p.npos()) |j| { + @memset(p.kp[j * dh_ ..][0..dh_], 0); + p.kp[j * dh_] = c.toBf16(-3 * rng.float(f32)); + for (p.vp[j * dh_ ..][0..dh_]) |*x| x.* = c.toBf16(rng.float(f32)); + } + } + + /// The caches as the engine holds them: a sliding layer's ring, a global layer's linear slots; then the call. + fn bind(p: *Probe) !void { + const dh_ = p.dh(); + const n = p.npos(); + const sliding = !p.v.global; + p.cap = if (sliding) ring else std.mem.alignForward(usize, n, 256); + const kc = try p.a.alloc(u16, p.nk() * p.cap * dh_); + defer p.a.free(kc); + const vc = try p.a.alloc(u16, p.nk() * p.cap * dh_); + defer p.a.free(vc); + // past the last key: NaN (prompt never reads it in); zeros as caches start (decode reads 32 value rows whole) + if (p.v.kind == .prompt) { + @memset(kc, nan16); + @memset(vc, nan16); + } else { + @memset(kc, 0); + @memset(vc, 0); + } + for (0..p.nk()) |h| for (0..n) |j| { + const slot = if (sliding) j % ring else j; + @memcpy(kc[(h * p.cap + slot) * dh_ ..][0..dh_], p.kp[(h * n + j) * dh_ ..][0..dh_]); + @memcpy(vc[(h * p.cap + slot) * dh_ ..][0..dh_], p.vp[(h * n + j) * dh_ ..][0..dh_]); + }; + p.refs[0] = try p.rig.input(std.mem.sliceAsBytes(p.q)); + p.refs[1] = try p.rig.input(std.mem.sliceAsBytes(kc)); + p.refs[2] = try p.rig.input(std.mem.sliceAsBytes(vc)); + p.refs[3] = try p.rig.input(std.mem.sliceAsBytes(p.hi)); + if (p.v.kind == .decode) { + // forward.zig's meta at the sub-batch's row: the chunks and the rows, which the merge reads + var meta: [tf.gemma4.state.meta_ints]i32 = @splat(0); + meta[1] = @intCast(p.chunks()); + meta[3] = @intCast(p.rows); + p.refs[4] = try p.rig.input(std.mem.asBytes(&meta)); + const slots = p.nq() * p.rows * p.chunks(); + p.parts[0] = try p.rig.scratch(slots * 4); + p.parts[1] = try p.rig.scratch(slots * 4); + p.parts[2] = try p.rig.scratch(slots * dh_ * 4); + p.parts[3] = try p.rig.scratch(4); + } + p.out = try p.rig.output(p.count() * 2); + } + + pub fn count(p: *Probe) usize { + return p.rows * p.nq() * p.dh(); + } + + pub fn run(p: *Probe) !void { + const pass = p.rig.begin(); + const e = pass.e; + const dh_ = p.dh(); + e.setPipeline(p.pipe); + for (p.refs[0..3], 0..) |r, i| e.setBuffer(r.buf, r.off, i); + if (p.v.kind == .prompt) { + // as prompt.zig: (tiles, query heads) groups, 32 threads a 64-column simdgroup + e.setBuffer(p.out.buf, p.out.off, 3); + e.setValue([3]i32{ @intCast(p.rows), @intCast(p.p0), @intCast(p.cap) }, 4); + e.setBuffer(p.refs[3].buf, p.refs[3].off, 5); + e.dispatchGroups(.{ .width = (p.rows + rt - 1) / rt, .height = p.nq() }, .{ .width = 32 * dh_ / 64 }); + } else { + // as forward.zig: (chunks, key heads, tiles of 16 / G rows) groups, 32 threads a 128-dim slice; the merge + e.setValue([4]i32{ @intCast(p.rows), @intCast(p.p0), @intCast(p.cap), @intCast(p.chunks()) }, 3); + for (p.parts[0..3], 4..) |r, i| e.setBuffer(r.buf, r.off, i); + const grp = p.nq() / p.nk(); + e.dispatchGroups(.{ .width = p.chunks(), .height = p.nk(), .depth = (p.rows * grp + 15) / 16 }, .{ .width = 32 * dh_ / 128 }); + e.setPipeline(p.merge.?); + for (p.parts[0..3], 0..) |r, i| e.setBuffer(r.buf, r.off, i); + e.setBuffer(p.refs[4].buf, p.refs[4].off, 3); + e.setBuffer(p.out.buf, p.out.off, 4); + e.setBuffer(p.parts[3].buf, p.parts[3].off, 5); + e.setValue(@as(i32, 0), 6); // no o-projection sums: they follow the bf16 output, which is all checked here + e.dispatchGroups(.{ .width = 1, .height = p.nq(), .depth = p.rows }, .{ .width = 32 }); + } + try p.rig.finish(pass); + } + + /// Every head's float64 point, half width and the cut's sides over the keys its row attends. + pub fn reference(p: *Probe, want: []f64, mag: []f64) !void { + const dh_ = p.dh(); + const n = p.npos(); + const grp = p.nq() / p.nk(); + const buf = try p.a.alloc(f64, 2 * n * dh_ + 2 * n + dh_); + defer p.a.free(buf); + const kf = buf[0 .. n * dh_]; + const vf = buf[n * dh_ .. 2 * n * dh_]; + const w = buf[2 * n * dh_ ..][0..n]; + const tau = buf[2 * n * dh_ + n ..][0..n]; + const qv = buf[buf.len - dh_ ..]; + for (0..p.nk()) |hk| { + for (kf, vf, 0..) |*k, *v, i| { + k.* = c.fromBf16(p.kp[hk * n * dh_ + i]); + v.* = c.fromBf16(p.vp[hk * n * dh_ + i]); + } + for (0..p.rows) |r| { + const keys = p.keysOf(r); + const len = keys[1] + 1 - keys[0]; + const f = p.form(r); + const ks = kf[keys[0] * dh_ ..]; + const vs = vf[keys[0] * dh_ ..]; + for (hk * grp..(hk + 1) * grp) |h| { + const o = (r * p.nq() + h) * dh_; + for (qv, 0..) |*x, e| x.* = c.fromBf16(p.q[o + e]); + ar.head(qv, ks, vs, len, dh_, 1, f, w, tau, want[o..][0..dh_], p.half[o..][0..dh_]); + ar.cutSides(f, w, tau, vs, len, dh_, p.down[o..][0..dh_], p.up[o..][0..dh_]); + } + } + } + @memcpy(mag, p.half); + @memcpy(p.point, want); + } + + /// The range's allowance (attention_range.allowBf16); gate 4's exact sums keep the exact bound. + pub fn bounds(p: *Probe, want: []const f64, _: []const f64, allow: []f64) !void { + const s = p.spec(); + for (allow, want, p.half, p.down, p.up) |*l, w, h, d, u| l.* = if (p.stress) c.boundExact(s, w) else ar.allowBf16(w, h, d, u); + } + + /// Each bf16 output within its range's rounded ends (attention_range.intervalBf16); gate 4's are open. + pub fn intervals(p: *Probe, low: []f64, high: []f64) !void { + for (low, high, p.point, p.half, p.down, p.up) |*l, *hi, pt, h, d, u| { + const e = if (p.stress) [2]f64{ -std.math.floatMax(f64), std.math.floatMax(f64) } else ar.intervalBf16(pt, h, d, u); + l.* = e[0]; + hi.* = e[1]; + } + } + + /// Gate 4: every weight 1, so each output is its row's values summed in `acc` over their count. + pub fn simulate(p: *Probe, acc: c.Accum, order: g.Order, out: []f64) !void { + if (!p.stress) return error.StressRequiresEqualScores; + const dh_ = p.dh(); + const n = p.npos(); + const grp = p.nq() / p.nk(); + var terms: std.ArrayList(f64) = .empty; + defer terms.deinit(p.a); + for (0..p.rows) |r| { + const keys = p.keysOf(r); + for (0..p.nk()) |hk| for (0..dh_) |e| { + terms.clearRetainingCapacity(); + for (keys[0]..keys[1] + 1) |j| try terms.append(p.a, c.fromBf16(p.vp[(hk * n + j) * dh_ + e])); + const v = bf(contract.rows.sumIn(acc, order, terms.items) / @as(f64, @floatFromInt(terms.items.len))); + for (hk * grp..(hk + 1) * grp) |h| out[(r * p.nq() + h) * dh_ + e] = v; + }; + } + } +}; + +const vt = g.hooks(Probe); + +/// Every variant and control, as tf-kernel-contract collects them. +pub const Suite = struct { + a: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + kernels: *std.ArrayList(g.Kernel), + probes: std.ArrayList(*Probe) = .empty, + + /// Only on the tensor units: the operand cut these ranges state was measured there alone. + pub fn collect(s: *Suite, only: ?[]const u8) !void { + if (!s.device.tensorUnits()) return; + for (variants, names) |v, name| { + if (only) |o| if (std.mem.indexOf(u8, name, o) == null) continue; + const p = try s.a.create(Probe); + errdefer s.a.destroy(p); + p.* = try Probe.init(s.a, s.device, s.queue, v, name); + errdefer p.deinit(); + try s.probes.append(s.a, p); + const control = v.case == .cutless or v.case == .edit; + try s.kernels.append(s.a, .{ .spec = p.spec(), .ctx = p, .vt = &vt, .control = if (control) .shapes else null }); + } + } + + pub fn deinit(s: *Suite) void { + for (s.probes.items) |p| { + p.deinit(); + s.a.destroy(p); + } + s.probes.deinit(s.a); + } +}; diff --git a/zig/tests/contract/metal_limits.zig b/zig/tests/contract/metal_limits.zig index d35831128e..3c6d87d556 100644 --- a/zig/tests/contract/metal_limits.zig +++ b/zig/tests/contract/metal_limits.zig @@ -66,6 +66,14 @@ fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch) !La const t = try @import("metal_minimax.zig").limitsText(a, device, d.name); return .{ .decl = d, .function = fun, .here = t.here, .pre_m5 = t.pre_m5 }; } + if (std.mem.startsWith(u8, d.name, "spark25 attention ")) { + const t = try @import("metal_spark25_attention.zig").limitsText(a, device, d.name); + return .{ .decl = d, .function = fun, .here = t.here, .pre_m5 = t.pre_m5 }; + } + if (std.mem.startsWith(u8, d.name, "gemma4 ") and has(d.name, " attention ")) { + const t = try @import("metal_gemma4_attention.zig").limitsText(a, device, d.name); + return .{ .decl = d, .function = fun, .here = t.here, .pre_m5 = t.pre_m5 }; + } if (std.mem.startsWith(u8, d.name, "drama3 ")) { const text_d3 = try @import("metal_drama3.zig").source(a, d.name); return .{ .decl = d, .function = fun, .here = text_d3, .pre_m5 = text_d3 }; @@ -73,11 +81,14 @@ fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch) !La if (std.mem.startsWith(u8, fun, "tf_ctl_")) { text = sources.core_contract_controls; if (has(fun, "big_tg")) head = "#define TF_CTL_BIG_TG 1\n"; - } else if (std.mem.startsWith(u8, d.name, "h3 gemm ")) { + } else if (std.mem.startsWith(u8, d.name, "h3 gemm ") or std.mem.startsWith(u8, d.name, "nemotron ssd gemm ")) { text = for (sources.prefill) |f| { if (std.mem.eql(u8, f.name, "gemm_nax")) break f.text; } else return error.UnknownLaunch; frags = true; + } else if (std.mem.startsWith(u8, d.name, "flashnext prompt attention ")) { + text = if (std.mem.eql(u8, fun, "tf_sattn_nax")) sources.flashnext_attn else sources.flashnext_qmm6; + frags = true; } else if (std.mem.startsWith(u8, d.name, "h3 ")) { text = @import("h3_kernels").source; } else if (std.mem.startsWith(u8, d.name, "glm gather ") or std.mem.startsWith(u8, d.name, "glm rows ")) { diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 32d697898e..ad0158f04e 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -382,12 +382,30 @@ pub fn main(init: std.process.Init) !u8 { var nemotron_attention: @import("metal_nemotron_attention.zig").Suite = .{ .a = a, .device = device, .queue = queue, .kernels = &kernels }; defer nemotron_attention.deinit(); try nemotron_attention.collect(o.only); + var nemotron_prompt: @import("metal_nemotron_prompt_range.zig").Suite = .{ .a = a, .device = device, .queue = queue, .kernels = &kernels }; + defer nemotron_prompt.deinit(); + try nemotron_prompt.collect(o.only); + var flashnext_attn: @import("metal_flashnext_attn.zig").Suite = .{ .a = a, .device = device, .queue = queue, .kernels = &kernels }; + defer flashnext_attn.deinit(); + try flashnext_attn.collect(o.only); + var spark25_attention: @import("metal_spark25_attention.zig").Suite = .{ .a = a, .device = device, .queue = queue, .kernels = &kernels }; + defer spark25_attention.deinit(); + try spark25_attention.collect(o.only); + var gemma_attention: @import("metal_gemma4_attention.zig").Suite = .{ .a = a, .device = device, .queue = queue, .kernels = &kernels }; + defer gemma_attention.deinit(); + try gemma_attention.collect(o.only); var nemotron_unroll: @import("metal_nemotron_unroll.zig").Suite = .{ .a = a, .device = device, .queue = queue, .kernels = &kernels }; defer nemotron_unroll.deinit(); try nemotron_unroll.collect(o.only); var minimax: minimax_k.Suite = .{ .a = a, .device = device, .queue = queue, .kernels = &kernels }; defer minimax.deinit(); try minimax.collect(o.only); + var mma_operand: @import("metal_mma_operand.zig").Suite = .{ .a = a, .device = device, .queue = queue, .kernels = &kernels }; + defer mma_operand.deinit(); + try mma_operand.collect(o.only); + var nemotron_ssd: @import("metal_nemotron_ssd_gemm.zig").Suite = .{ .a = a, .device = device, .queue = queue, .kernels = &kernels }; + defer nemotron_ssd.deinit(); + try nemotron_ssd.collect(o.only); const weight_only = @import("metal_weight_only.zig"); const weight_source = try weight_only.text(a); defer a.free(weight_source); diff --git a/zig/tests/contract/metal_minimax_nax.zig b/zig/tests/contract/metal_minimax_nax.zig index 9467529dc3..c0bc4649b5 100644 --- a/zig/tests/contract/metal_minimax_nax.zig +++ b/zig/tests/contract/metal_minimax_nax.zig @@ -25,6 +25,7 @@ pub const Suite = struct { grouped: std.ArrayList(*Grouped) = .empty, index: ?*attn.Index = null, prompt: ?*attn.Prompt = null, + masked: ?*attn.Prompt = null, pub fn collect(s: *Suite, only: ?[]const u8) !void { for (grouped_forms) |f| { @@ -44,6 +45,13 @@ pub const Suite = struct { s.prompt.?.* = try attn.Prompt.init(s.a, s.device, s.queue); try s.kernels.append(s.a, s.prompt.?.kernel()); } + // the cut was measured on the tensor units only, so the masked case registers there + if (s.device.tensorUnits() and wanted(only, attn.Prompt.maskedSpec().name)) { + s.masked = try s.a.create(attn.Prompt); + s.masked.?.* = try attn.Prompt.init(s.a, s.device, s.queue); + s.masked.?.masked = true; + try s.kernels.append(s.a, s.masked.?.kernel()); + } } pub fn deinit(s: *Suite) void { for (s.grouped.items) |x| { @@ -55,10 +63,10 @@ pub const Suite = struct { x.deinit(); s.a.destroy(x); } - if (s.prompt) |x| { + for ([_]?*attn.Prompt{ s.prompt, s.masked }) |p| if (p) |x| { x.deinit(); s.a.destroy(x); - } + }; } }; diff --git a/zig/tests/contract/metal_minimax_nax_attn.zig b/zig/tests/contract/metal_minimax_nax_attn.zig index 8047c507ee..3272d8161f 100644 --- a/zig/tests/contract/metal_minimax_nax_attn.zig +++ b/zig/tests/contract/metal_minimax_nax_attn.zig @@ -240,10 +240,10 @@ pub const Index = struct { // --- mm_prompt_attention: softmax over the kept blocks' keys up to each row, bf16 out ---------------------------- -/// The range's fp32 arithmetic: mma scores, fp32 SCALE * log2e, fast exp2, 64-key blocks; P multiplies V as bf16. +/// The range's fp32 arithmetic: mma scores, fp32 SCALE * log2e, fast exp2, 64-key blocks; fp32 P cut as it meets V. fn form(n: usize) ar.Form { const blocks: f64 = @floatFromInt(n / 64 + 2); - return .{ .score = @as(f64, @floatFromInt(dh / 16)) * mma_step + 2, .fast_exp = true, .chunk_exp = false, .p_rel = 0x1p-8, .num = mma_step * 4 + blocks + 2, .den = blocks + 6, .score_round = 4, .nan_max = true }; + return .{ .score = @as(f64, @floatFromInt(dh / 16)) * mma_step + 2, .fast_exp = true, .chunk_exp = false, .p_rel = 0, .p_cut = ar.operand_cut, .num = mma_step * 4 + blocks + 2, .den = blocks + 6, .score_round = 4, .nan_max = true }; } /// Gate 4's keys: the first 1,024 score high (their probabilities exactly 1), every other key scores far below. @@ -260,11 +260,16 @@ pub const Prompt = struct { rows: usize = 0, p0: usize = 0, fill: g.Fill = .random, + /// The masked case's entry (maskedSpec), with its plant. + masked: bool = false, q: []u16 = &.{}, k: []u16 = &.{}, v: []u16 = &.{}, sel: []u32 = &.{}, + point: []f64 = &.{}, half: []f64 = &.{}, + down: []f64 = &.{}, + up: []f64 = &.{}, refs: [5]rig_.Ref = undefined, out: rig_.Ref = undefined, @@ -282,6 +287,16 @@ pub const Prompt = struct { }; } + /// The masked case: gate 1's Inf at key p0 + 1, unseen by row 0 but read in its tile, makes row 0's outputs NaN. + pub fn maskedSpec() c.Spec { + var s = spec(); + s.name = "metal minimax_m3 mm_prompt_attention, a non-finite value at a key masked in its tile"; + s.served = s.served[0..1]; + s.tiles = .{ .{}, .{}, .{} }; + s.known = .{ .gate = .nonfinite, .finding = "a non-finite value at a key masked for a row of its tile reaches that row as NaN (P = 0 times Inf inside the P x V mma)" }; + return s; + } + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue) !Prompt { const h = try promptHead(a); defer a.free(h); @@ -290,12 +305,15 @@ pub const Prompt = struct { fn free(t: *Prompt) void { for ([_][]u16{ t.q, t.k, t.v }) |x| t.a.free(x); t.a.free(t.sel); - t.a.free(t.half); + for ([_][]f64{ t.point, t.half, t.down, t.up }) |x| t.a.free(x); t.q = &.{}; t.k = &.{}; t.v = &.{}; t.sel = &.{}; + t.point = &.{}; t.half = &.{}; + t.down = &.{}; + t.up = &.{}; } pub fn deinit(t: *Prompt) void { t.free(); @@ -303,7 +321,7 @@ pub const Prompt = struct { t.pipe.deinit(); } pub fn kernel(t: *Prompt) g.Kernel { - return .{ .spec = spec(), .ctx = t, .vt = &vt }; + return .{ .spec = if (t.masked) maskedSpec() else spec(), .ctx = t, .vt = &vt }; } const vt = g.hooks(Prompt); @@ -356,7 +374,7 @@ pub const Prompt = struct { t.k = try t.a.alloc(u16, slots * kw); t.v = try t.a.alloc(u16, slots * kw); t.sel = try t.a.alloc(u32, pad * kvh * words); - t.half = try t.a.alloc(f64, t.count()); + for ([_]*[]f64{ &t.point, &t.half, &t.down, &t.up }) |x| x.* = try t.a.alloc(f64, t.count()); var prng = std.Random.DefaultPrng.init(seed); const rng = prng.random(); fillRandom(rng, t.q, 3); @@ -373,6 +391,8 @@ pub const Prompt = struct { t.v[t.p0 * kw + 0 * dh + 5] = c.toBf16(std.math.inf(f32)); t.v[t.p0 * kw + 1 * dh + 6] = c.toBf16(std.math.nan(f32)); t.k[(t.p0 + 2) * kw + 2 * dh] = c.toBf16(std.math.nan(f32)); + // the masked case: row 0 does not see key p0 + 1, which its 2-row tile reads + if (t.masked and t.rows > 1) t.v[(t.p0 + 1) * kw + 3 * dh + 7] = c.toBf16(std.math.inf(f32)); }, .stress => { // keys of 3s or -3s: the first 1,024 score high, the rest probability 0; values cancel but for 8 nudged @@ -439,22 +459,26 @@ pub const Prompt = struct { const hq = kv * grp + gh; for (qv, 0..) |*x, e| x.* = c.fromBf16(t.q[(r * nq + hq) * dh + e]); const at = (r * nq + hq) * dh; - ar.head(qv, ks, vs, n, dh, scale, form(n), w, tau, want[at..][0..dh], hf); + const f = form(n); + ar.head(qv, ks, vs, n, dh, scale, f, w, tau, want[at..][0..dh], hf); @memcpy(t.half[at..][0..dh], hf); + ar.cutSides(f, w, tau, vs, n, dh, t.down[at..][0..dh], t.up[at..][0..dh]); } }; + @memcpy(t.point, want); } - /// The range's half width plus the output's bf16 rounding; gate 4's exact sums allow the rounding alone. + /// The range's allowance (attention_range.allowBf16); gate 4's exact sums allow the rounding alone. pub fn bounds(t: *Prompt, want: []const f64, _: []const f64, allow: []f64) !void { - for (allow, want, t.half) |*l, w, h| { - if (t.fill == .stress) { - l.* = c.boundExact(spec(), w); - } else if (!std.math.isFinite(w)) { - l.* = 0; - } else if (!std.math.isFinite(h)) { - l.* = std.math.floatMax(f64); // scores too far apart for the range's first order: any finite value - } else l.* = h + 0x1p-8 * (@abs(w) + h) + 0x1p-133; + for (allow, want, t.half, t.down, t.up) |*l, w, h, d, u| l.* = if (t.fill == .stress) c.boundExact(spec(), w) else ar.allowBf16(w, h, d, u); + } + + /// Each bf16 output within its range's rounded ends, the cut on its own side; gate 4's stress fill is open. + pub fn intervals(t: *Prompt, low: []f64, high: []f64) !void { + for (low, high, t.point, t.half, t.down, t.up) |*l, *hi, p, h, d, u| { + const e = if (t.fill == .stress) [2]f64{ -std.math.floatMax(f64), std.math.floatMax(f64) } else ar.intervalBf16(p, h, d, u); + l.* = e[0]; + hi.* = e[1]; } } diff --git a/zig/tests/contract/metal_mma_operand.zig b/zig/tests/contract/metal_mma_operand.zig new file mode 100644 index 0000000000..e0684554f2 --- /dev/null +++ b/zig/tests/contract/metal_mma_operand.zig @@ -0,0 +1,249 @@ +//! M5 tensor units' fp32 mma_16x32 operands: low 13 mantissa bits cleared, bit-equal to float64; controls go red. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); + +const c = contract.spec; +const g = contract.gates; + +/// Which operands are fp32: the left (attention's P), the right, or both (gemm_nax's f32 kernels). +const Side = enum { left, right, both }; + +/// How the reference takes an fp32 operand: the measured cut, or a planted control's wrong statement of it. +const Assume = enum { cut, fp32, nearest }; + +const Case = struct { side: Side, assume: Assume, name: []const u8 }; + +const cases = [_]Case{ + .{ .side = .left, .assume = .cut, .name = "metal mma_16x32 fp32 left operand cut" }, + .{ .side = .right, .assume = .cut, .name = "metal mma_16x32 fp32 right operand cut" }, + .{ .side = .both, .assume = .cut, .name = "metal mma_16x32 fp32 operands cut" }, + .{ .side = .left, .assume = .fp32, .name = "control: metal mma_16x32 left operand stated as fp32" }, + .{ .side = .right, .assume = .fp32, .name = "control: metal mma_16x32 right operand stated as fp32" }, + .{ .side = .left, .assume = .nearest, .name = "control: metal mma_16x32 left operand stated as rounded to nearest" }, + .{ .side = .right, .assume = .nearest, .name = "control: metal mma_16x32 right operand stated as rounded to nearest" }, +}; + +const source = + \\#include "../nax.h" + \\using namespace tfp; + \\[[kernel]] void tf_operand_cut(const device LEFT_T* A [[buffer(0)]], const device RIGHT_T* B [[buffer(1)]], + \\ device float* C [[buffer(2)]], uint lane [[thread_index_in_simdgroup]]) { + \\ const short2 home = frag_home(ushort(lane)); + \\ frag a; + \\ frag b0, b1; + \\ frag_get(a, A, 16, 0, 0, home); + \\ frag_get(b0, B, 32, 0, 0, home); + \\ frag_get(b1, B, 32, 0, 16, home); + \\ frag c0 = frag(0), c1 = frag(0); + \\ mma_16x32(c0, c1, a, b0, b1); + \\ frag_put(c0, C, 32, 0, 0, home); + \\ frag_put(c1, C, 32, 0, 16, home); + \\} +; + +/// One op: C (16 x 32) = A (16 x 16) times B (16 x 32), fp32 accumulators. +const m = 16; +const n = 32; +const kd = 16; + +/// The measured cut: fp32's low 13 mantissa bits cleared; NaN stays NaN. +fn cut(x: f32) f64 { + if (std.math.isNan(x)) return x; + return @as(f32, @bitCast(@as(u32, @bitCast(x)) & 0xffffe000)); +} + +/// 10 mantissa bits rounded to nearest, ties to even: the nearest control's statement. +fn nearest(x: f32) f64 { + if (!std.math.isFinite(x)) return x; + const b: u32 = @bitCast(x); + const r = (b + 0xfff + ((b >> 13) & 1)) & 0xffffe000; + return @as(f32, @bitCast(r)); +} + +fn noSums(_: c.Shape) usize { + return 0; +} + +/// No sums are tested: each output is one nonzero product, exact in fp32, so every output must equal float64's. +fn spec(v: Case) c.Spec { + return .{ + .name = v.name, + .accum = .none, + .depth = noSums, + .out = .f32, + .served = &.{ .{ .rows = m, .n = n, .k = kd }, .{ .rows = 7, .n = n, .k = kd }, .{ .rows = 1, .n = n, .k = kd } }, + }; +} + +const Probe = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + v: Case, + rows: usize = 0, + /// Both operands as values, row-major: what the device was given. + x: [m * kd]f32 = undefined, + y: [kd * n]f32 = undefined, + refs: [2]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, v: Case) !Probe { + const head = switch (v.side) { + .left => "#define LEFT_T float\n#define RIGHT_T bfloat\n", + .right => "#define LEFT_T bfloat\n#define RIGHT_T float\n", + .both => "#define LEFT_T float\n#define RIGHT_T float\n", + }; + const body = try tf.frags.source(device, a, source); + defer a.free(body); + const src = try std.mem.concat(a, u8, &.{ head, body }); + defer a.free(src); + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .pipe = try rig_.pipeline(device, src, "tf_operand_cut"), .v = v }; + } + + pub fn deinit(p: *Probe) void { + p.pipe.deinit(); + p.rig.deinit(); + } + + pub fn kernel(p: *Probe) g.Kernel { + return .{ .spec = spec(p.v), .ctx = p, .vt = &vt, .control = if (p.v.assume == .cut) null else .shapes }; + } + + const vt = g.hooks(Probe); + + fn leftF32(p: *const Probe) bool { + return p.v.side != .right; + } + + fn rightF32(p: *const Probe) bool { + return p.v.side != .left; + } + + /// A probed fp32 value: random sign and all 23 mantissa bits, exponent in 2^-30 to 2^30, or a subnormal. + fn probed(rng: std.Random, subnormals: bool) f32 { + const sign = @as(u32, rng.int(u1)) << 31; + const mant = rng.int(u23); + if (subnormals and rng.uintLessThan(u32, 8) == 0) return @bitCast(sign | @max(mant, 1)); + const exp: u32 = 127 - 30 + rng.uintLessThan(u32, 61); + return @bitCast(sign | exp << 23 | mant); + } + + /// The other side's value: plus or minus 2^0 to 2^8, so a product is exact in fp32, subnormals included. + fn scale(rng: std.Random) f32 { + const s: f32 = if (rng.boolean()) -1 else 1; + return s * std.math.ldexp(@as(f32, 1), @intCast(rng.uintLessThan(u32, 9))); + } + + /// One product an output: left and both C[r][j] = A[r][j % 16] B[j % 16][j], right C[r][j] = A[r][r] B[r][j]. + pub fn prepare(p: *Probe, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows < 1 or shape.rows > m or shape.n != n or shape.k != kd) return false; + p.rig.reset(); + p.rows = shape.rows; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + @memset(&p.x, 0); + @memset(&p.y, 0); + const sub = p.v.side != .both; // a scale side keeps a cut subnormal's product exact + if (p.v.side == .right) { + for (0..shape.rows) |r| p.x[r * kd + r] = scale(rng); + for (&p.y) |*y| y.* = probed(rng, sub); + } else { + for (p.x[0 .. shape.rows * kd]) |*x| x.* = probed(rng, sub); + for (0..n) |j| p.y[(j % kd) * n + j] = if (p.v.side == .both) probed(rng, false) else scale(rng); + } + if (fill == .nonfinite) { + // NaN, +Inf, -Inf a row each and both infinities in a fourth, on the fp32 side the case probes + const at = [_][2]usize{ .{ 0, 3 }, .{ 1, 5 }, .{ 2, 7 }, .{ 3, 1 }, .{ 3, 9 } }; + const put = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32), std.math.inf(f32), -std.math.inf(f32) }; + for (at, put) |q, v| if (q[0] < shape.rows) { + if (p.v.side == .right) p.y[q[0] * n + q[1] * 2] = v else p.x[q[0] * kd + q[1]] = v; + }; + } + p.refs[0] = try p.upload(&p.x, p.leftF32()); + p.refs[1] = try p.upload(&p.y, p.rightF32()); + p.out = try p.rig.output(m * n * 4); + return true; + } + + fn upload(p: *Probe, vals: []const f32, f32_side: bool) !rig_.Ref { + if (f32_side) return p.rig.input(std.mem.sliceAsBytes(vals)); + var h: [kd * n]u16 = undefined; + for (vals, 0..) |v, i| h[i] = c.toBf16(v); + return p.rig.input(std.mem.sliceAsBytes(h[0..vals.len])); + } + + pub fn run(p: *Probe) !void { + const pass = p.rig.begin(); + pass.e.setPipeline(p.pipe); + pass.e.setBuffer(p.refs[0].buf, p.refs[0].off, 0); + pass.e.setBuffer(p.refs[1].buf, p.refs[1].off, 1); + pass.e.setBuffer(p.out.buf, p.out.off, 2); + pass.e.dispatchThreads(mtl.Size.of(32, 1, 1), mtl.Size.of(32, 1, 1)); + try p.rig.finish(pass); + } + + pub fn count(p: *Probe) usize { + return p.rows * n; + } + + /// An fp32 operand as the case states the tensor units take it (bf16 values are exact). + fn take(p: *const Probe, x: f32, f32_side: bool) f64 { + if (!f32_side) return x; + return switch (p.v.assume) { + .cut => cut(x), + .fp32 => x, + .nearest => nearest(x), + }; + } + + pub fn reference(p: *Probe, want: []f64, mag: []f64) !void { + for (0..p.rows) |r| for (0..n) |j| { + var s: f64 = 0; + var t: f64 = 0; + for (0..kd) |k| { + const prod = p.take(p.x[r * kd + k], p.leftF32()) * p.take(p.y[k * n + j], p.rightF32()); + s += prod; + t += @abs(prod); + } + want[r * n + j] = s; + mag[r * n + j] = t; + }; + } + + /// Bit-equal: one exact product an output leaves no rounding to allow for. + pub fn bounds(_: *Probe, _: []const f64, _: []const f64, allow: []f64) !void { + @memset(allow, 0); + } +}; + +pub const Suite = struct { + a: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + kernels: *std.ArrayList(g.Kernel), + probes: std.ArrayList(*Probe) = .empty, + + /// Only on the tensor units, where the cut was measured; before the M5 the operands' precision is unmeasured. + pub fn collect(s: *Suite, only: ?[]const u8) !void { + if (!s.device.tensorUnits()) return; + for (cases) |v| { + if (only) |o| if (std.mem.indexOf(u8, v.name, o) == null) continue; + const p = try s.a.create(Probe); + errdefer s.a.destroy(p); + p.* = try Probe.init(s.a, s.device, s.queue, v); + try s.probes.append(s.a, p); + try s.kernels.append(s.a, p.kernel()); + } + } + + pub fn deinit(s: *Suite) void { + for (s.probes.items) |p| { + p.deinit(); + s.a.destroy(p); + } + s.probes.deinit(s.a); + } +}; diff --git a/zig/tests/contract/metal_nemotron_attention.zig b/zig/tests/contract/metal_nemotron_attention.zig index d3d319e666..b0d366c216 100644 --- a/zig/tests/contract/metal_nemotron_attention.zig +++ b/zig/tests/contract/metal_nemotron_attention.zig @@ -9,13 +9,13 @@ const c = contract.spec; const g = contract.gates; const ref = tf.attention_ref; -const D = 128; -const kvh = 2; -const heads = 32; -const G = heads / kvh; +pub const D = 128; +pub const kvh = 2; +pub const heads = 32; +pub const G = heads / kvh; const maxd = 64; // tf_tree_tail's MAXD const chunk = 512; // the partials' and the tail's keys a chunk -const scale: f32 = 1.0 / @sqrt(@as(f32, D)); +pub const scale: f32 = 1.0 / @sqrt(@as(f32, D)); /// decode: attn_partial_N then attn_merge; tree: a partial, tf_tree_tail and tf_tree_merge; prompt: tf_attention_nax. pub const Which = enum { decode, tree, prompt, write, put, compact }; @@ -210,7 +210,7 @@ pub const Probe = struct { return .{ .ctx = p, .vt = &vt, .spec = spec(p.which, p.tiles, p.control), .control = if (p.control) .shapes else null }; } - fn at(p: *const Probe, values: bool, h: usize, pos: usize) usize { + pub fn at(p: *const Probe, values: bool, h: usize, pos: usize) usize { return ((@as(usize, @intFromBool(values)) * kvh + h) * p.cap + pos) * D; } @@ -481,12 +481,12 @@ pub const Probe = struct { } /// Output index of query `qi`, head `h`: [heads, rows, D] after a merge, [rows, heads, D] from the prompt kernel. - fn outAt(p: *const Probe, qi: usize, h: usize) usize { + pub fn outAt(p: *const Probe, qi: usize, h: usize) usize { return if (p.which == .prompt) (qi * heads + h) * D else (h * p.rows + qi) * D; } /// Query `qi` of head `h` as float64: Qp's row qi * G + g of its kv head, or the prompt's row. - fn query(p: *const Probe, qi: usize, h: usize, out: []f64) void { + pub fn query(p: *const Probe, qi: usize, h: usize, out: []f64) void { const at_ = if (p.which == .prompt) (qi * heads + h) * D else ((h / G) * 16 * p.rows + qi * G + h % G) * D; for (out, 0..) |*x, d| x.* = c.fromBf16(p.q[at_ + d]); } @@ -581,6 +581,7 @@ pub const Suite = struct { pub fn collect(s: *Suite, only: ?[]const u8) !void { for (std.enums.values(Which)) |w| for (1..if (w == .decode) 17 else 2) |tiles| for ([_]bool{ false, true }) |control| { if (only) |o| if (std.mem.indexOf(u8, name(w, tiles, control), o) == null) continue; + if (w == .prompt and !control and s.device.tensorUnits()) continue; // M5: metal_nemotron_prompt_range.zig const p = try s.a.create(Probe); errdefer s.a.destroy(p); p.* = try Probe.init(s.a, s.device, s.queue, w, tiles, control); diff --git a/zig/tests/contract/metal_nemotron_prompt_range.zig b/zig/tests/contract/metal_nemotron_prompt_range.zig new file mode 100644 index 0000000000..1bcc6ca692 --- /dev/null +++ b/zig/tests/contract/metal_nemotron_prompt_range.zig @@ -0,0 +1,176 @@ +//! Nemotron's tf_attention_nax on the M5: attention_range at its arithmetic, fp32 P cut as the tensor units take it. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const old = @import("metal_nemotron_attention.zig"); +const c = contract.spec; +const g = contract.gates; +const ar = contract.attention_range; +const D = old.D; + +/// P with the tensor units' cut; P spread over [0, 1]; that case without the cut; a non-finite value at a masked key. +const Case = enum { cut, spread, cutless, masked }; + +fn name(k: Case) []const u8 { + return switch (k) { + .cut => "metal nemotron attention prompt nax", + .spread => "metal nemotron attention prompt nax, P spread over [0, 1]", + .cutless => "control: metal nemotron attention prompt nax, P spread, its range without the operand cut", + .masked => "metal nemotron attention prompt nax, a non-finite value at a key masked in its tile", + }; +} + +/// tf_attention_nax: 8 mma steps a score, fp32 SCALE * log2e, fast exp2; a 32-key block: 2 P x V steps, rescale, folds. +fn form(n: usize, cut: f64) ar.Form { + const blocks: f64 = @floatFromInt(n / 32 + 2); + return .{ .score = 9 * ar.mma_step, .fast_exp = true, .chunk_exp = false, .p_rel = 0, .p_cut = cut, .num = (2 * ar.mma_step + 1) * blocks + ar.mma_step + 2, .den = 3 * blocks + 6, .score_round = 4, .div = 2 }; +} + +const Probe = struct { + inner: old.Probe, + case: Case, + rig: *@import("rig.zig").Rig, + point: []f64 = &.{}, + half: []f64 = &.{}, + down: []f64 = &.{}, + up: []f64 = &.{}, + + fn deinit(p: *Probe) void { + for ([_][]f64{ p.point, p.half, p.down, p.up }) |x| p.inner.a.free(x); + p.inner.deinit(); + } + + pub fn prepare(p: *Probe, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (!try p.inner.prepare(shape, fill, seed)) return false; + if ((p.case == .spread or p.case == .cutless) and fill == .random) p.spread(seed); + if (p.case == .masked and fill == .nonfinite) p.maskedInf(); + return true; + } + + /// +Inf in V at key len + 5 (head 0, column 7): rows 0 to 4 have it masked, yet their tile reads it. + fn maskedInf(p: *Probe) void { + const in = &p.inner; + in.host[in.at(true, 0, in.len + 5) + 7] = 0x7f80; + const cache = old.kvh * in.cap * D; + @memcpy(p.rig.bytes(in.kv[1], cache * 2), std.mem.sliceAsBytes(in.host[cache..][0..cache])); + } + + /// Scores q . k = q0 k0 alone, so P = 2^(s - max) spreads over about [2^-6, 1]; values in [0, 1) add the bias up. + fn spread(p: *Probe, seed: u64) void { + const in = &p.inner; + var prng = std.Random.DefaultPrng.init(seed ^ 0x1546); + const rnd = prng.random(); + for (0..in.rows * old.heads) |r| { + @memset(in.q[r * D ..][0..D], 0); + in.q[r * D] = c.toBf16(0.5 + rnd.float(f32)); + } + for (0..old.kvh) |h| for (0..in.cap) |j| { + in.host[in.at(false, h, j)] = c.toBf16(-32 * rnd.float(f32)); + for (0..D) |d| in.host[in.at(true, h, j) + d] = c.toBf16(rnd.float(f32)); + }; + // the inputs as bound: the caches' two halves and the queries + const cache = old.kvh * in.cap * D; + for (in.kv, 0..) |r, i| @memcpy(p.rig.bytes(r, cache * 2), std.mem.sliceAsBytes(in.host[i * cache ..][0..cache])); + @memcpy(p.rig.bytes(in.qref, in.q.len * 2), std.mem.sliceAsBytes(in.q)); + } + + pub fn run(p: *Probe) !void { + try p.inner.run(); + } + + pub fn count(p: *Probe) usize { + return p.inner.count(); + } + + pub fn simulate(p: *Probe, acc: c.Accum, order: g.Order, out: []f64) !void { + try p.inner.simulate(acc, order, out); + } + + /// Every head's float64 point, half width and the cut's sides over the keys up to its row. + pub fn reference(p: *Probe, want: []f64, mag: []f64) !void { + const a = p.inner.a; + const in = &p.inner; + for ([_]*[]f64{ &p.point, &p.half, &p.down, &p.up }) |x| { + a.free(x.*); + x.* = try a.alloc(f64, want.len); + } + const most = in.len + in.rows; + const buf = try a.alloc(f64, 2 * most * D + 2 * most); + defer a.free(buf); + const kf = buf[0 .. most * D]; + const vf = buf[most * D .. 2 * most * D]; + const w = buf[2 * most * D ..][0..most]; + const tau = buf[2 * most * D + most ..][0..most]; + const cut: f64 = if (p.case == .cutless) 0 else ar.operand_cut; + var q: [D]f64 = undefined; + for (0..old.kvh) |hk| { + for (0..most * D) |i| { + kf[i] = c.fromBf16(in.host[in.at(false, hk, 0) + i]); + vf[i] = c.fromBf16(in.host[in.at(true, hk, 0) + i]); + } + for (0..in.rows) |qi| for (hk * old.G..(hk + 1) * old.G) |h| { + const n = in.len + qi + 1; + const f = form(n, cut); + in.query(qi, h, &q); + const o = in.outAt(qi, h); + ar.head(&q, kf, vf, n, D, old.scale, f, w, tau, want[o..][0..D], p.half[o..][0..D]); + ar.cutSides(f, w, tau, vf, n, D, p.down[o..][0..D], p.up[o..][0..D]); + }; + } + @memcpy(mag, p.half); + @memcpy(p.point, want); + } + + /// The range's allowance (attention_range.allowBf16); gate 4's exact sums keep the exact bound. + pub fn bounds(p: *Probe, want: []const f64, _: []const f64, allow: []f64) !void { + const s = old.spec(.prompt, 1, false); + for (want, allow, p.half, p.down, p.up) |w, *l, h, d, u| l.* = if (p.inner.stress) c.boundExact(s, w) else ar.allowBf16(w, h, d, u); + } + + /// Each bf16 output within its range's rounded ends, the cut on its own side; gate 4's stress fill is open. + pub fn intervals(p: *Probe, low: []f64, high: []f64) !void { + for (low, high, p.point, p.half, p.down, p.up) |*l, *hi, pt, h, d, u| { + const e = if (p.inner.stress) [2]f64{ -std.math.floatMax(f64), std.math.floatMax(f64) } else ar.intervalBf16(pt, h, d, u); + l.* = e[0]; + hi.* = e[1]; + } + } +}; + +const vt = g.hooks(Probe); + +pub const Suite = struct { + a: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + kernels: *std.ArrayList(g.Kernel), + probes: std.ArrayList(*Probe) = .empty, + + /// Only on the tensor units, where the cut was measured; before the M5, attention_ref's bound stays. + pub fn collect(s: *Suite, only: ?[]const u8) !void { + if (!s.device.tensorUnits()) return; + for (std.enums.values(Case)) |k| { + if (only) |o| if (std.mem.indexOf(u8, name(k), o) == null) continue; + const p = try s.a.create(Probe); + errdefer s.a.destroy(p); + p.* = .{ .inner = try old.Probe.init(s.a, s.device, s.queue, .prompt, 1, false), .case = k, .rig = undefined }; + p.rig = &p.inner.rig; + try s.probes.append(s.a, p); + var spec = old.spec(.prompt, 1, false); + spec.name = name(k); + if (k == .masked) { + spec.served = spec.served[0..1]; + spec.known = .{ .gate = .nonfinite, .finding = "a non-finite value at a key masked for a row of its tile reaches that row as NaN (P = 0 times Inf inside the P x V mma)" }; + } + try s.kernels.append(s.a, .{ .spec = spec, .ctx = p, .vt = &vt, .control = if (k == .cutless) .shapes else null }); + } + } + + pub fn deinit(s: *Suite) void { + for (s.probes.items) |p| { + p.deinit(); + s.a.destroy(p); + } + s.probes.deinit(s.a); + } +}; diff --git a/zig/tests/contract/metal_nemotron_ssd_gemm.zig b/zig/tests/contract/metal_nemotron_ssd_gemm.zig new file mode 100644 index 0000000000..c0ee664685 --- /dev/null +++ b/zig/tests/contract/metal_nemotron_ssd_gemm.zig @@ -0,0 +1,363 @@ +//! Nemotron's SSD steps on gemm_nax's five f32 kernels on the M5: both fp32 operands cut, each product's cut one-sided. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const sources = @import("kernel_sources"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const c = contract.spec; +const g = contract.gates; +const ar = contract.attention_range; + +/// Nemotron 3.5 Lightning's Mamba heads, head dim and state (config.checkShapes), and prefill's 256-row SSD step. +const heads = 64; +const head_dim = 64; +const state = 128; +const step = 256; + +/// prefill_mixers.mamba's f32 GEMMs a step of r rows: y = sur x dtx (r x r, r x 64), dd^T x bh^T (64 x r, r x 128). +const Product = enum { y, state }; + +/// An entry point by its tags (A and B transposed; whole M, N and K blocks) and the shapes the step calls it at. +const Entry = struct { product: Product, tags: []const u8, served: []const c.Shape }; + +fn at(n: usize, k: usize) c.Shape { + return .{ .rows = heads, .n = n, .k = k }; +} + +/// A shape's rows are the heads (the launch's batch), n its columns, k the step's rows r: y's M is r, the state's 64. +const entries = [_]Entry{ + .{ .product = .y, .tags = "n_n_n_n_n", .served = &.{ at(head_dim, 100), at(head_dim, 255), at(head_dim, 63), at(head_dim, 33), at(head_dim, 1) } }, + .{ .product = .y, .tags = "n_n_t_n_n", .served = &.{ at(head_dim, 64), at(head_dim, 128), at(head_dim, 192) } }, + .{ .product = .y, .tags = "n_n_t_n_t", .served = &.{at(head_dim, step)} }, + .{ .product = .state, .tags = "t_t_t_t_n", .served = &.{ at(state, 100), at(state, 255), at(state, 192), at(state, 64), at(state, 1) } }, + .{ .product = .state, .tags = "t_t_t_t_t", .served = &.{at(state, step)} }, +}; + +/// Random operands; operands in (0, 1], where the cut's bias adds up; those stated as fp32; a source edit. +const Case = enum { plain, positive, fp32, edited }; + +const Reg = struct { entry: usize, case: Case, name: []const u8 }; + +const regs = [_]Reg{ + .{ .entry = 0, .case = .plain, .name = "metal nemotron ssd y gemm_nax_f32_n_n_n_n_n" }, + .{ .entry = 1, .case = .plain, .name = "metal nemotron ssd y gemm_nax_f32_n_n_t_n_n" }, + .{ .entry = 2, .case = .plain, .name = "metal nemotron ssd y gemm_nax_f32_n_n_t_n_t" }, + .{ .entry = 3, .case = .plain, .name = "metal nemotron ssd state gemm_nax_f32_t_t_t_t_n" }, + .{ .entry = 4, .case = .plain, .name = "metal nemotron ssd state gemm_nax_f32_t_t_t_t_t" }, + .{ .entry = 2, .case = .positive, .name = "metal nemotron ssd y gemm_nax_f32_n_n_t_n_t, operands in (0, 1]" }, + .{ .entry = 4, .case = .positive, .name = "metal nemotron ssd state gemm_nax_f32_t_t_t_t_t, operands in (0, 1]" }, + .{ .entry = 2, .case = .fp32, .name = "control: metal nemotron ssd y gemm_nax_f32_n_n_t_n_t, operands in (0, 1] stated as fp32" }, + .{ .entry = 4, .case = .fp32, .name = "control: metal nemotron ssd state gemm_nax_f32_t_t_t_t_t, operands in (0, 1] stated as fp32" }, + .{ .entry = 0, .case = .edited, .name = "control: metal nemotron ssd y gemm_nax_f32_n_n_n_n_n, its ragged K tail one short" }, +}; + +/// The edited control's one line: the last 32-deep step past the whole K blocks reads one column fewer. +const edit_from = "max(0, rest - kk)"; +const edit_to = "max(0, rest - kk - 1)"; + +/// gemm_step's two mma_16x32 k16 steps a 32-deep step, zero-padded past K, each mma_step roundings of |acc| + sum |ab|. +fn depth(s: c.Shape) usize { + return ar.mma_step * 2 * ((s.k + 31) / 32); +} + +fn spec(r: Reg) c.Spec { + return .{ .name = r.name, .accum = .f32, .depth = depth, .out = .f32, .served = entries[r.entry].served }; +} + +fn tag(flag: bool) []const u8 { + return if (flag) "t" else "n"; +} + +/// The entry point prefill_launch.gemm names for these sizes is this entry's. +fn fits(tags: []const u8, m: usize, n: usize, k: usize) bool { + var buf: [5]u8 = undefined; + const want = std.fmt.bufPrint(&buf, "{s}_{s}_{s}", .{ tag(m % 64 == 0), tag(n % 128 == 0), tag(k % 256 == 0) }) catch unreachable; + return std.mem.eql(u8, tags[4..], want); +} + +/// Each product's float64 cut: the operands' magnitudes shrink by up to 2^-10 each, so a product's by c (2 - c). +fn productCut(r: Reg) f64 { + const cut: f64 = if (r.case == .fp32) 0 else ar.operand_cut; + return cut * (2 - cut); +} + +/// float64's own roundings in the reference, and fp32's least step. +const slack = 1 + 0x1p-20; +const floor = 0x1p-149; + +const Probe = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + reg: Reg, + batch: usize = 0, + m: usize = 0, + n: usize = 0, + k: usize = 0, + stress: bool = false, + /// A and B as bound: y's sur [heads, r, r], dtx [r, heads * 64]; the state's dd [r, heads * 64], bh [heads, 128, r] + x: []f32 = &.{}, + w: []f32 = &.{}, + /// Per output: float64's value, the fp32 sums' allowance, and the cut's reach down (positive terms) and up. + point: []f64 = &.{}, + err: []f64 = &.{}, + down: []f64 = &.{}, + up: []f64 = &.{}, + params: [16]i32 = @splat(0), + refs: [2]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, r: Reg) !Probe { + const text = for (sources.prefill) |f| { + if (std.mem.eql(u8, f.name, "gemm_nax")) break f.text; + } else return error.UnknownSource; + const built = try tf.frags.source(device, a, text); + defer a.free(built); + var src: []u8 = built; + if (r.case == .edited) { + if (std.mem.count(u8, built, edit_from) != 1) return error.ControlEditMissing; + src = try std.mem.replaceOwned(u8, a, built, edit_from, edit_to); + } + defer if (r.case == .edited) a.free(src); + const fun = try std.fmt.allocPrint(a, "custom_kernel_tf_gemm_nax_f32_{s}_64_128_256_2_4_float_float_int32_t_float", .{entries[r.entry].tags}); + defer a.free(fun); + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .pipe = try rig_.pipeline(device, src, fun), .reg = r }; + } + + fn drop(p: *Probe) void { + p.a.free(p.x); + p.a.free(p.w); + for ([_][]f64{ p.point, p.err, p.down, p.up }) |v| p.a.free(v); + p.x = &.{}; + p.w = &.{}; + p.point = &.{}; + p.err = &.{}; + p.down = &.{}; + p.up = &.{}; + p.rig.reset(); + } + + fn deinit(p: *Probe) void { + p.drop(); + p.pipe.deinit(); + p.rig.deinit(); + } + + const vt = g.hooks(Probe); + + fn kernel(p: *Probe) g.Kernel { + const control: ?c.Gate = if (p.reg.case == .fp32 or p.reg.case == .edited) .shapes else null; + return .{ .spec = spec(p.reg), .ctx = p, .vt = &vt, .control = control }; + } + + fn isY(p: *const Probe) bool { + return entries[p.reg.entry].product == .y; + } + + fn lda(p: *const Probe) usize { + return if (p.isY()) p.k else p.batch * p.m; + } + + fn ldb(p: *const Probe) usize { + return if (p.isY()) p.batch * p.n else p.k; + } + + /// Head b's A[i][kk] and B[kk][j] in the stored arrays. + fn aAt(p: *const Probe, b: usize, i: usize, kk: usize) usize { + return if (p.isY()) (b * p.k + i) * p.k + kk else kk * p.lda() + b * p.m + i; + } + + fn bAt(p: *const Probe, b: usize, kk: usize, j: usize) usize { + return if (p.isY()) kk * p.ldb() + b * p.n + j else (b * p.n + j) * p.k + kk; + } + + pub fn prepare(p: *Probe, s: c.Shape, fill: g.Fill, seed: u64) !bool { + const m: usize = if (p.isY()) s.k else head_dim; + if (s.rows == 0 or s.n == 0 or s.k == 0 or !fits(entries[p.reg.entry].tags, m, s.n, s.k)) return false; + p.drop(); + p.batch = s.rows; + p.m = m; + p.n = s.n; + p.k = s.k; + p.stress = fill == .stress; + const outs = p.count(); + p.x = try p.a.alloc(f32, p.batch * p.m * p.k); + p.w = try p.a.alloc(f32, p.batch * p.k * p.n); + for ([_]*[]f64{ &p.point, &p.err, &p.down, &p.up }) |v| v.* = try p.a.alloc(f64, outs); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (p.stress) p.stressFill(rng) else { + const positive = p.reg.case == .positive or p.reg.case == .fp32; + for (p.x) |*v| v.* = if (positive) 1 - rng.float(f32) else rng.float(f32) * 2 - 1; + for (p.w) |*v| v.* = if (positive) 1 - rng.float(f32) else rng.float(f32) * 2 - 1; + } + if (fill == .nonfinite) { + // NaN, +Inf and -Inf in head 0's A rows 0 to 2, both infinities in row 3, -Inf in head 1's B column 1 + const k = p.k; + const put = [_][3]usize{ .{ 0, k / 3, 0 }, .{ 1, k / 2, 1 }, .{ 2, k - 1, 2 }, .{ 3, 0, 1 }, .{ 3, k - 1, 2 } }; + const vals = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32) }; + for (put) |q| if (q[0] < p.m) { + p.x[p.aAt(0, q[0], q[1])] = vals[q[2]]; + }; + if (p.batch > 1) p.w[p.bAt(1, 0, 1)] = -std.math.inf(f32); + } + const tn = (p.n + 127) / 128; + const tm = (p.m + 63) / 64; + const strides = if (p.isY()) [2]usize{ p.k * p.k, p.n } else [2]usize{ p.m, p.n * p.k }; + const v = [_]usize{ p.m, p.n, p.k, p.lda(), p.ldb(), p.n, tn, tm, 2, p.k / 256, strides[0], strides[1], p.m * p.n, 0, 0 }; + for (v, 0..) |x, i| p.params[i] = @intCast(x); + p.refs[0] = try p.rig.input(std.mem.sliceAsBytes(p.x)); + p.refs[1] = try p.rig.input(std.mem.sliceAsBytes(p.w)); + p.out = try p.rig.output(outs * 4); + return true; + } + + /// Gate 4: 11-bit integers times 5-bit ones, which the cut keeps: halves past fp16 and bf16, partials below 2^24. + fn stressFill(p: *Probe, rng: std.Random) void { + const h = p.k / 2; + @memset(p.x, 0); + @memset(p.w, 0); + for (0..p.batch) |b| { + for (0..p.m) |i| for (0..h) |kk| { + const v: f32 = @floatFromInt(1024 + rng.uintLessThan(u32, 1024)); + p.x[p.aAt(b, i, kk)] = v; + p.x[p.aAt(b, i, h + kk)] = v + @as(f32, if (kk < 8) 1 else 0); + }; + for (0..h) |kk| for (0..p.n) |j| { + const w: f32 = @floatFromInt(1 + rng.uintLessThan(u32, 31)); + p.w[p.bAt(b, kk, j)] = w; + p.w[p.bAt(b, h + kk, j)] = -w; + }; + } + } + + /// prefill_launch.gemm's dispatch: 32 x 4 x 2 threads a group, swizzle 2, one z pair a head. + pub fn run(p: *Probe) !void { + const tn = (p.n + 127) / 128; + const tm = (p.m + 63) / 64; + const pass = p.rig.begin(); + pass.e.setPipeline(p.pipe); + pass.e.setBuffer(p.refs[0].buf, p.refs[0].off, 0); + pass.e.setBuffer(p.refs[1].buf, p.refs[1].off, 1); + pass.e.setBytes(std.mem.asBytes(&p.params), 2); + pass.e.setBuffer(p.out.buf, p.out.off, 3); + pass.e.dispatchThreads(mtl.Size.of((tn << 2) * 32, (tm + 3) / 4 * 4, p.batch * 2), mtl.Size.of(32, 4, 2)); + try p.rig.finish(pass); + } + + pub fn count(p: *Probe) usize { + return p.batch * p.m * p.n; + } + + /// Head b's A as [m][k] and B as [n][k] in float64, each output's terms contiguous. + fn unpack(p: *const Probe, b: usize, am: []f64, bt: []f64) void { + for (0..p.m) |i| for (0..p.k) |kk| { + am[i * p.k + kk] = p.x[p.aAt(b, i, kk)]; + }; + for (0..p.n) |j| for (0..p.k) |kk| { + bt[j * p.k + kk] = p.w[p.bAt(b, kk, j)]; + }; + } + + /// D = A B in float64 on the operands as given, the fp32 sums' allowance and the cut's one-sided reach. + pub fn reference(p: *Probe, want: []f64, mag: []f64) !void { + const K = p.k; + const buf = try p.a.alloc(f64, (p.m + p.n) * K); + defer p.a.free(buf); + const am = buf[0 .. p.m * K]; + const bt = buf[p.m * K ..]; + const cut = productCut(p.reg); + const sums = c.gamma(@floatFromInt(depth(.{ .rows = p.batch, .n = p.n, .k = K })), 0x1p-24); + for (0..p.batch) |b| { + p.unpack(b, am, bt); + for (0..p.m) |i| for (0..p.n) |j| { + var s: f64 = 0; + var t: f64 = 0; + var pos: f64 = 0; + var neg: f64 = 0; + for (am[i * K ..][0..K], bt[j * K ..][0..K]) |x, y| { + const q = x * y; + s += q; + t += @abs(q); + if (q > 0) pos += q else neg -= q; + } + const o = (b * p.m + i) * p.n + j; + want[o] = s; + mag[o] = t; + p.err[o] = sums * t; + p.down[o] = cut * pos; + p.up[o] = cut * neg; + }; + } + @memcpy(p.point, want); + } + + /// The sums' allowance plus the cut's wider side; gate 4's exact sums keep the exact bound. + pub fn bounds(p: *Probe, want: []const f64, _: []const f64, allow: []f64) !void { + const s = spec(p.reg); + for (want, allow, p.err, p.down, p.up) |w, *l, e, d, u| { + l.* = if (p.stress) c.boundExact(s, w) else (e + @max(d, u)) * slack + floor; + } + } + + /// Each output within float64's value less the sums' allowance and the positive terms' cut, plus the negative's. + pub fn intervals(p: *Probe, low: []f64, high: []f64) !void { + const s = spec(p.reg); + for (low, high, p.point, p.err, p.down, p.up) |*l, *h, pt, e, d, u| { + if (p.stress) { + l.* = pt - c.boundExact(s, pt); + h.* = pt + c.boundExact(s, pt); + } else { + l.* = pt - (e + d) * slack - floor; + h.* = pt + (e + u) * slack + floor; + } + } + } + + /// Each output's products summed in the accumulator and order asked (gate 4's operands survive the cut whole). + pub fn simulate(p: *Probe, acc: c.Accum, order: g.Order, out: []f64) !void { + const K = p.k; + const buf = try p.a.alloc(f64, (p.m + p.n + 1) * K); + defer p.a.free(buf); + const am = buf[0 .. p.m * K]; + const bt = buf[p.m * K ..][0 .. p.n * K]; + const terms = buf[(p.m + p.n) * K ..]; + for (0..p.batch) |b| { + p.unpack(b, am, bt); + for (0..p.m) |i| for (0..p.n) |j| { + for (terms, am[i * K ..][0..K], bt[j * K ..][0..K]) |*q, x, y| q.* = x * y; + out[(b * p.m + i) * p.n + j] = contract.rows.sumIn(acc, order, terms); + }; + } + } +}; + +pub const Suite = struct { + a: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + kernels: *std.ArrayList(g.Kernel), + probes: std.ArrayList(*Probe) = .empty, + + /// Only on the tensor units, where the cut was measured; before the M5 the operands' precision is unmeasured. + pub fn collect(s: *Suite, only: ?[]const u8) !void { + if (!s.device.tensorUnits()) return; + for (regs) |r| { + if (only) |o| if (std.mem.indexOf(u8, r.name, o) == null) continue; + const p = try s.a.create(Probe); + errdefer s.a.destroy(p); + p.* = try Probe.init(s.a, s.device, s.queue, r); + try s.probes.append(s.a, p); + try s.kernels.append(s.a, p.kernel()); + } + } + + pub fn deinit(s: *Suite) void { + for (s.probes.items) |p| { + p.deinit(); + s.a.destroy(p); + } + s.probes.deinit(s.a); + } +}; diff --git a/zig/tests/contract/metal_spark25_attention.zig b/zig/tests/contract/metal_spark25_attention.zig new file mode 100644 index 0000000000..42f36c1bd2 --- /dev/null +++ b/zig/tests/contract/metal_spark25_attention.zig @@ -0,0 +1,408 @@ +//! Spark-X2.6-4B's prompt attention (spark25/attention_nax.metal) as metal_model.zig builds it, both layer kinds. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); +const rig_ = @import("rig.zig"); + +const c = contract.spec; +const g = contract.gates; +const ar = contract.attention_range; +const cfg = tf.spark25_metal.config; +const st = tf.spark25_metal.state; +const geo = cfg.spark4b; +const dh: usize = cfg.head_dim; +const nq: usize = geo.query_heads; +const kvh: usize = geo.kv_heads; +const grp: usize = nq / kvh; +const kw: usize = kvh * dh; +const scale: f64 = cfg.scale; +const rt = 32; // query rows a tile +const bkt = dh / 64 * 32; // keys a block: 32 a simdgroup, a simdgroup each 64 output columns +const nan16 = 0x7fc0; + +/// The two pipelines metal_model.zig builds: sliding layers (a 512-key window over a ring) and full layers. +pub const Layer = enum { sliding, full }; + +/// The engine's P with the M5's cut; P spread over [0, 1]; that without the cut; a bad edit; a masked key's Inf. +const Case = enum { plain, spread, cutless, edit, masked }; + +fn window(l: Layer) usize { + return if (l == .sliding) cfg.window else 0; +} + +fn ring(l: Layer) usize { + return if (l == .sliding) st.sliding_ring else 0; +} + +/// metal_model.zig's header for a layer kind. +fn header(a: std.mem.Allocator, l: Layer) ![]u8 { + return std.fmt.allocPrint(a, "#define DH {d}\n#define NQ {d}\n#define G {d}\n#define KVH {d}\n#define WINDOW {d}\n#define RING {d}\n#define SCALE {e}f\n", .{ cfg.head_dim, geo.query_heads, geo.query_heads / geo.kv_heads, geo.kv_heads, window(l), ring(l), cfg.scale }); +} + +/// Each layer's planted control: sliding sees one key past its window, full one key past its row. +fn edit(l: Layer) [2][]const u8 { + return switch (l) { + .sliding => .{ "col < row - WINDOW + 1", "col < row - WINDOW" }, + .full => .{ "const bool out = col > row ||", "const bool out = col > row + 1 ||" }, + }; +} + +/// The source as metal_model.zig builds it on this GPU (header, then frags.source's text), with a control's edit. +fn built(a: std.mem.Allocator, device: mtl.Device, l: Layer, case: Case) ![]u8 { + const head = try header(a, l); + defer a.free(head); + const body = try tf.frags.source(device, a, sources.spark25_attention_nax); + defer a.free(body); + const src = try std.mem.concat(a, u8, &.{ head, body }); + if (case != .edit) return src; + defer a.free(src); + const e = edit(l); + if (std.mem.count(u8, src, e[0]) != 1) return error.ControlEditMissing; + return std.mem.replaceOwned(u8, a, src, e[0], e[1]); +} + +/// Gate 5's texts for `spark25 attention ` launches, nax.h inlined (fragments before the M5). +pub fn limitsText(a: std.mem.Allocator, device: mtl.Device, launch: []const u8) !struct { here: []const u8, pre_m5: []const u8 } { + const l: Layer = if (std.mem.indexOf(u8, launch, " sliding ") != null) .sliding else if (std.mem.indexOf(u8, launch, " full ") != null) .full else return error.UnknownLaunch; + const head = try header(a, l); + const body = try std.mem.replaceOwned(u8, a, sources.spark25_attention_nax, "#include \"../nax.h\"", sources.nax); + const pre = try std.mem.concat(a, u8, &.{ head, "#define TF_SIMD_FRAGS 1\n", body }); + return .{ .here = if (device.tensorUnits()) try std.mem.concat(a, u8, &.{ head, body }) else pre, .pre_m5 = pre }; +} + +pub fn name(l: Layer, case: Case) []const u8 { + return switch (l) { + inline else => |x| switch (case) { + inline else => |k| comptime label(x, k), + }, + }; +} + +fn label(comptime l: Layer, comptime k: Case) []const u8 { + const base = "metal spark25 prompt attention " ++ @tagName(l); + return switch (k) { + .plain => base, + .spread => base ++ ", P spread over [0, 1]", + .cutless => "control: " ++ base ++ ", P spread, its range without the operand cut", + .edit => "control: " ++ base ++ (if (l == .sliding) ", a key past the window" else ", a key past the row"), + .masked => base ++ ", a non-finite value at a key masked in its tile", + }; +} + +/// The score's mma chain; the bounds hook below replaces the generic bound it would set. +fn depth(_: c.Shape) usize { + return dh / 16 * ar.mma_step; +} + +/// Rows: a prompt chunk's rows (the engine takes this path past 16); n: the head dim; k: the chunk's first position. +fn spec(l: Layer, case: Case) c.Spec { + const plain_shapes: []const c.Shape = switch (l) { + // a chunk mid-block, the shortest from the start, one deep in a wrapped ring, a whole chunk after a wrap + .sliding => &.{ .{ .rows = 33, .n = dh, .k = 777 }, .{ .rows = 17, .n = dh, .k = 0 }, .{ .rows = 300, .n = dh, .k = 100003 }, .{ .rows = st.prompt_rows, .n = dh, .k = 2600 } }, + // a chunk mid-block, the shortest from the start, a whole chunk from the start, one deep in the context + .full => &.{ .{ .rows = 33, .n = dh, .k = 300 }, .{ .rows = 17, .n = dh, .k = 0 }, .{ .rows = st.prompt_rows, .n = dh, .k = 0 }, .{ .rows = 17, .n = dh, .k = 32700 } }, + }; + const spread_shapes: []const c.Shape = switch (l) { + .sliding => &.{ .{ .rows = 64, .n = dh, .k = 700 }, .{ .rows = 33, .n = dh, .k = 0 } }, + .full => &.{ .{ .rows = 64, .n = dh, .k = 1100 }, .{ .rows = 33, .n = dh, .k = 0 } }, + }; + // one served shape, deep enough for gate 4's high keys + const masked_shapes = if (l == .sliding) plain_shapes[0..1] else plain_shapes[3..4]; + const spread = case == .spread or case == .cutless; + return .{ + .name = name(l, case), + .accum = .f32, + .depth = depth, + .out = .bf16, + .tiles = if (spread or case == .masked) .{ .{}, .{}, .{} } else .{ .{ .size = rt, .min = st.window_rows + 1 }, .{}, .{ .size = bkt, .min = 0 } }, + .served = if (spread) spread_shapes else if (case == .masked) masked_shapes else plain_shapes, + .sums = .{ .exact = true }, + .known = if (case == .masked) masked_key else null, + }; +} + +/// The masked case's gate-1 plant at key p0 + 5, which rows 0 to 4 do not see but their tile reads, makes them NaN. +const masked_key: c.Known = .{ .gate = .nonfinite, .finding = "a non-finite value at a key masked for a row of its tile reaches that row as NaN (P = 0 times Inf inside the P x V mma)" }; + +/// sp_prompt_attention over n keys, its counts from the source; `cut` the fp32 P's operand cut into P x V. +fn form(n: usize, cut: f64) ar.Form { + const blocks: f64 = @floatFromInt(n / bkt + 2); // 128-key blocks a row's tile visits, a window starting mid-block + return .{ + .score = @as(f64, @floatFromInt(dh / 16)) * ar.mma_step, // DH / 16 mma steps a score, from zero each block + .score_round = 4, // the fp32 SCALE * log2e and its product + .fast_exp = true, + // fast::exp2's own 3 + 1.2 |x|; each block's keep, exp2(top - top_new), adds 3, and their arguments sum to |x| + .exp_ulps = .{ 3 + 3 * blocks, 1.2 + 1.2 + 0.5 }, + .chunk_exp = false, + .p_rel = 0, // P is fp32 in threadgroup memory + .p_cut = cut, + .num = (bkt / 16 * ar.mma_step + 1) * blocks, // a block's 8 P x V steps on the running acc, then its keep + .den = 3 * blocks + 7, // a block's keep and two fold adds, a fold's 4-level tree; 3 adds over the simdgroups + .div = 2, // 1 / total, then the product + .nan_max = true, + }; +} + +/// Gate 4's keys: values repeat every 512 positions, a half and its negation but for 8 nudged keys a column. +const period = 512; +/// Gate 4's values are multiples of 64 up to 255 * 64: a 32-lane order's lane sums pass fp16's largest, all stay exact. +const step: f32 = 64; +/// Full layers' high keys (every other key scores far below); a sliding row's window is all high. +const stress_on = 1024; + +const Probe = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + layer: Layer, + case: Case, + cut: f64, + rows: usize = 0, + p0: usize = 0, + slots: usize = 0, + fill: g.Fill = .random, + q: []u16 = &.{}, + k: []u16 = &.{}, + v: []u16 = &.{}, + point: []f64 = &.{}, + half: []f64 = &.{}, + down: []f64 = &.{}, + up: []f64 = &.{}, + refs: [4]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, l: Layer, case: Case) !Probe { + const src = try built(a, device, l, case); + defer a.free(src); + const cut: f64 = if (case == .cutless) 0 else ar.operand_cut; // the cutless control states P whole + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .pipe = try rig_.pipeline(device, src, "sp_prompt_attention"), .layer = l, .case = case, .cut = cut }; + } + + fn free(p: *Probe) void { + for ([_][]u16{ p.q, p.k, p.v }) |x| p.a.free(x); + for ([_][]f64{ p.point, p.half, p.down, p.up }) |x| p.a.free(x); + p.q = &.{}; + p.k = &.{}; + p.v = &.{}; + p.point = &.{}; + p.half = &.{}; + p.down = &.{}; + p.up = &.{}; + } + + fn deinit(p: *Probe) void { + p.free(); + p.rig.deinit(); + p.pipe.deinit(); + } + + fn kernel(p: *Probe) g.Kernel { + const control: ?g.Gate = if (p.case == .cutless or p.case == .edit) .shapes else null; + return .{ .spec = spec(p.layer, p.case), .ctx = p, .vt = &vt, .control = control }; + } + const vt = g.hooks(Probe); + + /// Position x's cache slot: the ring's, or x itself. + fn slot(p: *const Probe, x: usize) usize { + return if (p.layer == .sliding) x % st.sliding_ring else x; + } + + /// The first key position row `pos` sees. + fn first(p: *const Probe, pos: usize) usize { + return if (p.layer == .sliding) (pos + 1) -| cfg.window else 0; + } + + pub fn prepare(p: *Probe, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.rows > st.prompt_rows or shape.n != dh) return false; + if (fill == .stress and shape.k < @as(usize, if (p.layer == .sliding) cfg.window else stress_on)) return false; + p.rig.reset(); + p.free(); + p.rows = shape.rows; + p.p0 = shape.k; + p.fill = fill; + const end = p.p0 + p.rows; // positions written once the chunk's keys are in + // the ring, or whole blocks: a block scores all its 128 keys, past the layer it reads the engine's values + p.slots = if (p.layer == .sliding) st.sliding_ring else std.mem.alignForward(usize, end, bkt); + const pad = std.mem.alignForward(usize, p.rows, rt); // a tile reads all its 32 query rows + p.q = try p.a.alloc(u16, pad * nq * dh); + p.k = try p.a.alloc(u16, p.slots * kw); + p.v = try p.a.alloc(u16, p.slots * kw); + for ([_]*[]f64{ &p.point, &p.half, &p.down, &p.up }) |x| x.* = try p.a.alloc(f64, p.count()); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + fillRandom(rng, p.q, 3); + fillRandom(rng, p.k, 3); + fillRandom(rng, p.v, 2); + switch (fill) { + .random => if (p.case == .spread or p.case == .cutless) p.spread(rng), + .nonfinite => { + // every row sees key p0, rows 2 on see key p0 + 2 + for ([_]usize{ 3, 1 * dh, 2 * dh + 5 }, [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32) }, 0..) |at, x, r| if (r < p.rows) { + p.q[(r * nq + r * grp + 1) * dh + at % dh] = c.toBf16(x); + }; + p.v[p.slot(p.p0) * kw + 0 * dh + 5] = c.toBf16(std.math.inf(f32)); + p.v[p.slot(p.p0) * kw + 1 * dh + 6] = c.toBf16(std.math.nan(f32)); + // the masked case: rows 0 to 4 do not see key p0 + 5, yet their tile reads it + if (p.case == .masked and p.rows > 5) p.v[p.slot(p.p0 + 5) * kw + 3 * dh + 7] = c.toBf16(std.math.inf(f32)); + p.k[p.slot(p.p0 + 2) * kw + 2 * dh] = c.toBf16(std.math.nan(f32)); + }, + .stress => try p.stress(rng), + } + // rows past the chunk and slots no position has reached hold anything: NaN here, which no output may see + @memset(p.q[p.rows * nq * dh ..], nan16); + if (end < p.slots) { + @memset(p.k[end * kw ..], nan16); + @memset(p.v[end * kw ..], nan16); + } + p.refs[0] = try p.rig.input(std.mem.sliceAsBytes(p.q)); + p.refs[1] = try p.rig.input(std.mem.sliceAsBytes(p.k)); + p.refs[2] = try p.rig.input(std.mem.sliceAsBytes(p.v)); + p.refs[3] = try p.rig.input(std.mem.asBytes(&[2]i32{ @intCast(p.rows), @intCast(p.p0) })); + p.out = try p.rig.output(p.count() * 2); + return true; + } + + /// Scores q . k = q0 k0 alone, so P = 2^(s - max) spreads over about [2^-9, 1]; values in [0, 1) add the bias up. + fn spread(p: *Probe, rng: std.Random) void { + for (0..p.rows * nq) |r| { + @memset(p.q[r * dh ..][0..dh], 0); + p.q[r * dh] = c.toBf16(0.5 + rng.float(f32)); + } + for (0..p.slots * kvh) |r| { + p.k[r * dh] = c.toBf16(-64 * rng.float(f32)); + for (p.v[r * dh ..][0..dh]) |*x| x.* = c.toBf16(rng.float(f32)); + } + } + + /// Equal high scores (P exactly 1) over a sliding window or a full layer's first 1,024 keys, the rest P = 0. + fn stress(p: *Probe, rng: std.Random) !void { + @memset(p.q, c.toBf16(3)); + const base = try p.a.alloc(f32, period / 2 * kw); + defer p.a.free(base); + for (base) |*x| x.* = step * @as(f32, @floatFromInt(1 + rng.uintLessThan(u32, 255))); + // a slot holds the latest position at it, so its value follows its position modulo 512 (the ring is 5 periods) + for (0..p.slots) |s| { + const high = p.layer == .sliding or s < stress_on; + @memset(p.k[s * kw ..][0..kw], c.toBf16(if (high) 3 else -3)); + const j = s % period; + for (0..kw) |e| p.v[s * kw + e] = c.toBf16(if (j < period / 2) base[j * kw + e] else -base[(j - period / 2) * kw + e] + step * @as(f32, @floatFromInt(@intFromBool(j - period / 2 < 8)))); + } + } + + pub fn run(p: *Probe) !void { + const pass = p.rig.begin(); + pass.e.setPipeline(p.pipe); + for (p.refs[0..3], 0..) |r, i| pass.e.setBuffer(r.buf, r.off, i); + pass.e.setBuffer(p.out.buf, p.out.off, 3); + pass.e.setBuffer(p.refs[3].buf, p.refs[3].off, 4); + // as metal_forward.zig: (rows / 32, query heads) groups, 32 threads a 64-column simdgroup + pass.e.dispatchGroups(mtl.Size.of((p.rows + rt - 1) / rt, nq, 1), mtl.Size.of(dh / 64 * 32, 1, 1)); + try p.rig.finish(pass); + } + + pub fn count(p: *const Probe) usize { + return p.rows * nq * dh; + } + + /// Each row's heads over the keys it sees (a window's or every one up to it): float64's points and ranges. + pub fn reference(p: *Probe, want: []f64, mag: []f64) !void { + const lo = p.first(p.p0); + const most = p.p0 + p.rows - lo; + const buf = try p.a.alloc(f64, 2 * most * dh + 2 * most + 2 * dh); + defer p.a.free(buf); + const ks = buf[0 .. most * dh]; + const vs = buf[most * dh .. 2 * most * dh]; + const w = buf[2 * most * dh ..][0..most]; + const tau = buf[2 * most * dh + most ..][0..most]; + const qv = buf[buf.len - 2 * dh ..][0..dh]; + const hf = buf[buf.len - dh ..]; + for (0..kvh) |kv| { + for (0..most) |j| for (0..dh) |e| { + const at = p.slot(lo + j) * kw + kv * dh + e; + ks[j * dh + e] = c.fromBf16(p.k[at]); + vs[j * dh + e] = c.fromBf16(p.v[at]); + }; + for (0..p.rows) |r| { + const pos = p.p0 + r; + const from = p.first(pos) - lo; + const n = pos + 1 - p.first(pos); + const f = form(n, p.cut); + const kr = ks[from * dh ..][0 .. n * dh]; + const vr = vs[from * dh ..][0 .. n * dh]; + for (kv * grp..(kv + 1) * grp) |hq| { + for (qv, 0..) |*x, e| x.* = c.fromBf16(p.q[(r * nq + hq) * dh + e]); + const o = (r * nq + hq) * dh; + ar.head(qv, kr, vr, n, dh, scale, f, w, tau, want[o..][0..dh], hf); + @memcpy(p.half[o..][0..dh], hf); + ar.cutSides(f, w, tau, vr, n, dh, p.down[o..][0..dh], p.up[o..][0..dh]); + } + } + } + @memcpy(mag, p.half); + @memcpy(p.point, want); + } + + /// The range's allowance (attention_range.allowBf16); gate 4 allows the rounding alone. + pub fn bounds(p: *Probe, want: []const f64, _: []const f64, allow: []f64) !void { + for (allow, want, p.half, p.down, p.up) |*l, w, h, d, u| l.* = if (p.fill == .stress) c.boundExact(spec(p.layer, p.case), w) else ar.allowBf16(w, h, d, u); + } + + /// Each bf16 output within its range's rounded ends (attention_range.intervalBf16); gate 4's are open. + pub fn intervals(p: *Probe, low: []f64, high: []f64) !void { + for (low, high, p.point, p.half, p.down, p.up) |*l, *hi, pt, h, d, u| { + const e = if (p.fill == .stress) [2]f64{ -std.math.floatMax(f64), std.math.floatMax(f64) } else ar.intervalBf16(pt, h, d, u); + l.* = e[0]; + hi.* = e[1]; + } + } + + /// Gate 4: the outputs if P x V's sum of the high keys' values ran in `acc`, in position order. + pub fn simulate(p: *Probe, acc: c.Accum, order: g.Order, out: []f64) !void { + var terms: [stress_on]f64 = undefined; + for (0..p.rows) |r| for (0..nq) |hq| for (0..dh) |e| { + const col = (hq / grp) * dh + e; + const pos = p.p0 + r; + const keys = if (p.layer == .sliding) terms[0..cfg.window] else terms[0..stress_on]; + const from = if (p.layer == .sliding) p.first(pos) else 0; + for (keys, 0..) |*x, j| x.* = c.fromBf16(p.v[p.slot(from + j) * kw + col]); + const sum = contract.rows.sumIn(acc, order, keys) / @as(f64, @floatFromInt(keys.len)); + out[(r * nq + hq) * dh + e] = c.fromBf16(c.toBf16(@floatCast(sum))); + }; + } +}; + +/// Random bf16 values in [-r, r). +fn fillRandom(rng: std.Random, out: []u16, r: f32) void { + for (out) |*x| x.* = c.toBf16((rng.float(f32) * 2 - 1) * r); +} + +/// Both layers' pipelines with their controls: the cut was measured on the tensor units only, so they register there. +pub const Suite = struct { + a: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + kernels: *std.ArrayList(g.Kernel), + probes: std.ArrayList(*Probe) = .empty, + + pub fn collect(s: *Suite, only: ?[]const u8) !void { + if (!s.device.tensorUnits()) return; + for (std.enums.values(Layer)) |l| for (std.enums.values(Case)) |k| { + if (only) |o| if (std.mem.indexOf(u8, name(l, k), o) == null) continue; + const p = try s.a.create(Probe); + errdefer s.a.destroy(p); + p.* = try Probe.init(s.a, s.device, s.queue, l, k); + try s.probes.append(s.a, p); + try s.kernels.append(s.a, p.kernel()); + }; + } + + pub fn deinit(s: *Suite) void { + for (s.probes.items) |p| { + p.deinit(); + s.a.destroy(p); + } + s.probes.deinit(s.a); + } +};