From 8d0c1dd4bab88686a81ce9db756146257ee50494 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Sat, 31 Oct 2026 05:06:49 -0500 Subject: [PATCH 01/17] kernel contract: the M5 tensor units' fp32 mma_16x32 operands, left, right and both, cut to 10 mantissa bits toward zero (bits & 0xffffe000, fp32's exponents and subnormals kept), bit-equal to float64; four controls state the operand as fp32 or rounded to nearest --- docs/kernel-contract.md | 6 + zig/tests/contract/metal_main.zig | 3 + zig/tests/contract/metal_mma_operand.zig | 249 +++++++++++++++++++++++ 3 files changed, 258 insertions(+) create mode 100644 zig/tests/contract/metal_mma_operand.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index b30ed4c065..383911e6cb 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -138,6 +138,12 @@ 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`. Before the M5, nax.h's fragments +keep fp32 operands whole, so the cases run only on the tensor units. 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/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 32d697898e..1ff6ba3c49 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -388,6 +388,9 @@ pub fn main(init: std.process.Init) !u8 { 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); 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_mma_operand.zig b/zig/tests/contract/metal_mma_operand.zig new file mode 100644 index 0000000000..e50e6bf057 --- /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 where the tensor units are: before the M5, nax.h's fragments keep fp32 operands whole. + 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); + } +}; From 675bd0a2a0184f57a06cded02166eb983c5f3cbb Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Sat, 31 Oct 2026 05:18:34 -0500 Subject: [PATCH 02/17] kernel contract: attention_range names the M5's operand cut (operand_cut, 2^-10, one-sided p_cut); Nemotron's tf_attention_nax and MiniMax's mm_prompt_attention state their fp32 P with it on the M5, each bf16 output within its range's rounded ends, and Nemotron's adds a P-spread case whose cutless range is the planted control --- docs/kernel-contract.md | 6 + zig/src/core/attention_range.zig | 50 +++++ zig/tests/contract/metal_main.zig | 3 + zig/tests/contract/metal_minimax_nax_attn.zig | 35 +++- .../contract/metal_nemotron_attention.zig | 17 +- .../contract/metal_nemotron_prompt_range.zig | 171 ++++++++++++++++++ 6 files changed, 266 insertions(+), 16 deletions(-) create mode 100644 zig/tests/contract/metal_nemotron_prompt_range.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 383911e6cb..d95c5145a0 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -144,6 +144,12 @@ single product bit-equal to float64 on operands with their low 13 mantissa bits 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`. Before the M5, nax.h's fragments keep fp32 operands whole, so the cases run only on the tensor units. +`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` and MiniMax-M3.1-Flash's `mm_prompt_attention` state their fp32 P with it, and 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. 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..4c8bbbdb98 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,33 @@ 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) }; +} + /// 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 +207,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/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 1ff6ba3c49..df696a2491 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -382,6 +382,9 @@ 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 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); diff --git a/zig/tests/contract/metal_minimax_nax_attn.zig b/zig/tests/contract/metal_minimax_nax_attn.zig index 8047c507ee..0a4b083b16 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. @@ -264,7 +264,10 @@ pub const Prompt = struct { k: []u16 = &.{}, v: []u16 = &.{}, sel: []u32 = &.{}, + point: []f64 = &.{}, half: []f64 = &.{}, + down: []f64 = &.{}, + up: []f64 = &.{}, refs: [5]rig_.Ref = undefined, out: rig_.Ref = undefined, @@ -290,12 +293,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(); @@ -356,7 +362,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); @@ -439,22 +445,35 @@ 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 half width and the cut's wider side plus bf16's rounding; gate 4 allows the rounding alone. pub fn bounds(t: *Prompt, want: []const f64, _: []const f64, allow: []f64) !void { - for (allow, want, t.half) |*l, w, h| { + for (allow, want, t.half, t.down, t.up) |*l, w, h, d, u| { 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; + } else l.* = h + @max(d, u) + 0x1p-8 * (@abs(w) + h + @max(d, u)) + 0x1p-133; + } + } + + /// Each bf16 output within its range's ends as rounded, the cut on its own side; gate 4 and wide ranges are 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 open = t.fill == .stress or !std.math.isFinite(p) or !std.math.isFinite(h); + const e = if (open) [2]f64{ -std.math.floatMax(f64), std.math.floatMax(f64) } else ar.ends(p, h, d, u); + l.* = e[0]; + hi.* = e[1]; } } 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..e81421d798 --- /dev/null +++ b/zig/tests/contract/metal_nemotron_prompt_range.zig @@ -0,0 +1,171 @@ +//! 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] with values in [0, 1); that case's range without the cut. +const Case = enum { cut, spread, cutless }; + +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", + }; +} + +/// 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 != .cut and fill == .random) p.spread(seed); + return true; + } + + /// 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 half width and the cut's wider side plus bf16's rounding; 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| { + if (p.inner.stress) { + l.* = c.boundExact(s, 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 + @max(d, u) + 0x1p-8 * (@abs(w) + h + @max(d, u)) + 0x1p-133; + } + } + + /// Each bf16 output within its range's ends as rounded, the cut on its own side; gate 4 and wide ranges 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 open = p.inner.stress or !std.math.isFinite(pt) or !std.math.isFinite(h); + const e = if (open) [2]f64{ -std.math.floatMax(f64), std.math.floatMax(f64) } else ar.ends(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: before the M5, nax.h's fragments keep P whole and 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); + 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); + } +}; From 0dcd0758b4c8b41ab520cfdd86654ad5dcd816e4 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Sat, 31 Oct 2026 06:05:53 -0500 Subject: [PATCH 03/17] kernel contract: Nemotron's tf_attention_nax on the M5 plants a non-finite value at a key masked for earlier rows of its tile, which reaches them as NaN (P = 0 times Inf inside the P x V mma), and keeps it as a known gate-1 finding --- zig/tests/contract/metal_nemotron_prompt_range.zig | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/zig/tests/contract/metal_nemotron_prompt_range.zig b/zig/tests/contract/metal_nemotron_prompt_range.zig index e81421d798..fa97fb0ff3 100644 --- a/zig/tests/contract/metal_nemotron_prompt_range.zig +++ b/zig/tests/contract/metal_nemotron_prompt_range.zig @@ -42,9 +42,18 @@ const Probe = struct { 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 != .cut and fill == .random) p.spread(seed); + if (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; @@ -157,6 +166,7 @@ pub const Suite = struct { try s.probes.append(s.a, p); var spec = old.spec(.prompt, 1, false); spec.name = name(k); + if (k != .cutless) 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 }); } } From 3f3711cb90e31117af9b273654b768286a55c4f1 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Sat, 31 Oct 2026 05:46:17 -0500 Subject: [PATCH 04/17] kernel contract: Nemotron's SSD steps on gemm_nax's five f32 kernels, both fp32 operands cut on the M5 as a one-sided range beside the mma steps' fp32 sums; two cases with operands in (0, 1] whose fp32 statement is the control, a ragged-K-tail source edit, and the five launches' gate-5 records --- docs/kernel-contract.md | 8 + .../core/contract/nemotron_ssd_records.zig | 15 + zig/src/core/contract/registry.zig | 2 +- zig/tests/contract/metal_limits.zig | 2 +- zig/tests/contract/metal_main.zig | 3 + .../contract/metal_nemotron_ssd_gemm.zig | 363 ++++++++++++++++++ 6 files changed, 391 insertions(+), 2 deletions(-) create mode 100644 zig/src/core/contract/nemotron_ssd_records.zig create mode 100644 zig/tests/contract/metal_nemotron_ssd_gemm.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index d95c5145a0..7cd689403f 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -150,6 +150,14 @@ M5, Nemotron's `tf_attention_nax` and MiniMax-M3.1-Flash's `mm_prompt_attention` 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. +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 keep fp32 whole and are +not registered; their launches are in `nemotron_ssd_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/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..0b527dd271 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -569,7 +569,7 @@ test { } /// 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("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/tests/contract/metal_limits.zig b/zig/tests/contract/metal_limits.zig index d35831128e..5ca66493c1 100644 --- a/zig/tests/contract/metal_limits.zig +++ b/zig/tests/contract/metal_limits.zig @@ -73,7 +73,7 @@ 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; diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index df696a2491..2046da749d 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -394,6 +394,9 @@ pub fn main(init: std.process.Init) !u8 { 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_nemotron_ssd_gemm.zig b/zig/tests/contract/metal_nemotron_ssd_gemm.zig new file mode 100644 index 0000000000..1829092fe8 --- /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: before the M5, nax.h's SIMD fragments keep both fp32 operands whole. + 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); + } +}; From 313d57f9e836f7113c53e0e0f2e94f9a3c1e9259 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Sat, 31 Oct 2026 05:44:15 -0500 Subject: [PATCH 05/17] kernel contract: Flash Next's prompt attention on the M5 (tf_sattn_nax, tf_attn256_nax) states its fp32 P with the operand cut, each bf16 output within its range's ends through the rounded gate; P-spread cases with cutless and source-edit controls, and their gate-5 launches --- .../contract/flashnext_attention_records.zig | 6 + zig/src/core/contract/registry.zig | 2 +- zig/tests/contract/metal_flashnext_attn.zig | 444 ++++++++++++++++++ zig/tests/contract/metal_limits.zig | 3 + zig/tests/contract/metal_main.zig | 3 + 5 files changed, 457 insertions(+), 1 deletion(-) create mode 100644 zig/src/core/contract/flashnext_attention_records.zig create mode 100644 zig/tests/contract/metal_flashnext_attn.zig 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/registry.zig b/zig/src/core/contract/registry.zig index 0b527dd271..bc8332a92a 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -569,7 +569,7 @@ test { } /// 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("nemotron_ssd_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("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/tests/contract/metal_flashnext_attn.zig b/zig/tests/contract/metal_flashnext_attn.zig new file mode 100644 index 0000000000..4f63765878 --- /dev/null +++ b/zig/tests/contract/metal_flashnext_attn.zig @@ -0,0 +1,444 @@ +//! 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], values in [0, 1); that range without the cut; an edit. +pub const Case = enum { cut, spread, cutless, edit }; + +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) ", the key past a row's last counted" else ", one key past the row seen"), + }; +} + +/// 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 }, + }; +} + +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 + } + + 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 => { + // key 0's values, which every row attends; three queries' heads; key 2's key, which rows 2 on see + for (0..KVH) |kv| for ([_]u16{ nan16, 0x7f80, 0xff80 }, 0..) |x, d| { + p.v[p.at(kv, 0) + d] = x; + }; + for ([_]u16{ nan16, 0x7f80, 0xff80 }, 1..) |x, r| if (r < p.rows) { + p.q[(r * H + 5 * r) * D + r] = x; + }; + p.k[p.at(1, 2) + 7] = nan16; + }, + .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: before the M5, nax.h's fragments keep P whole, and these forms are not run 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_limits.zig b/zig/tests/contract/metal_limits.zig index 5ca66493c1..5167fa76df 100644 --- a/zig/tests/contract/metal_limits.zig +++ b/zig/tests/contract/metal_limits.zig @@ -78,6 +78,9 @@ fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch) !La 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 2046da749d..878689ef6f 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -385,6 +385,9 @@ pub fn main(init: std.process.Init) !u8 { 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 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); From d482922770e9d7b8c6c0410196660cd73df005e6 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Sat, 31 Oct 2026 05:49:24 -0500 Subject: [PATCH 06/17] kernel contract: Flash Next's prompt attention keeps its gate-1 findings as known (masked keys at -FLT_MAX, sattn's padding lanes read key 0), planted at every seed, with their Known findings and Coverage lines --- docs/kernel-contract.md | 9 +++++++++ zig/tests/contract/metal_flashnext_attn.zig | 15 +++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 7cd689403f..232ad6d937 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -121,6 +121,10 @@ 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 (0 x Inf is NaN where float64 gives the infinity). 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 @@ -150,6 +154,11 @@ M5, Nemotron's `tf_attention_nax` and MiniMax-M3.1-Flash's `mm_prompt_attention` 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: a row's next +key counted. 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 diff --git a/zig/tests/contract/metal_flashnext_attn.zig b/zig/tests/contract/metal_flashnext_attn.zig index 4f63765878..93cb592617 100644 --- a/zig/tests/contract/metal_flashnext_attn.zig +++ b/zig/tests/contract/metal_flashnext_attn.zig @@ -88,6 +88,16 @@ pub fn spec(w: Which, k: Case) c.Spec { // 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 == .cut or k == .spread) .{ .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 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", }; } @@ -267,6 +277,11 @@ pub const Probe = struct { p.q[(r * H + 5 * r) * D + r] = x; }; p.k[p.at(1, 2) + 7] = nan16; + // row 4's head 0 scores -Inf at every key it sees: q +Inf against keys 0 to 4, each -1 in that dim + if (p.rows > 4) { + p.q[(4 * H + 0) * D + 9] = 0x7f80; + for (0..5) |key| p.k[p.at(0, key) + 9] = c.toBf16(-1); + } }, .stress => { // keys of 3s or -3s: the first 1,024 score high, the rest probability 0; values cancel but for 8 nudged From 85478297605e0e6b34cb99bf995d37499113ec0a Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Sat, 31 Oct 2026 05:51:47 -0500 Subject: [PATCH 07/17] kernel contract: Spark-X2.6-4B's sp_prompt_attention, sliding and full layers as metal_model.zig builds them, fp32 P cut on the M5 within attention_range's rounded ends, P-spread cases with cutless controls, a window and a causal edit as controls, and gate-5 records --- docs/kernel-contract.md | 3 + zig/src/core/contract/registry.zig | 2 +- zig/src/core/contract/spark25_records.zig | 7 + zig/tests/contract/metal_limits.zig | 4 + zig/tests/contract/metal_main.zig | 3 + .../contract/metal_spark25_attention.zig | 409 ++++++++++++++++++ 6 files changed, 427 insertions(+), 1 deletion(-) create mode 100644 zig/src/core/contract/spark25_records.zig create mode 100644 zig/tests/contract/metal_spark25_attention.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 232ad6d937..e98f96cdcf 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -167,6 +167,9 @@ the fp32 sums' `mma_step` roundings for every k16 step. Gate 4 runs 11-bit integ 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 keep fp32 whole and 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 on the M5 and whole before it, each with a P +spread case and its cutless control (M5 only), and one source edit each (a key past the window, a key past the row). 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/contract/registry.zig b/zig/src/core/contract/registry.zig index bc8332a92a..265fe5e621 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -569,7 +569,7 @@ test { } /// 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("nemotron_ssd_records.zig").records ++ @import("flashnext_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("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_limits.zig b/zig/tests/contract/metal_limits.zig index 5167fa76df..e7a2fdcfd6 100644 --- a/zig/tests/contract/metal_limits.zig +++ b/zig/tests/contract/metal_limits.zig @@ -66,6 +66,10 @@ 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, "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 }; diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 878689ef6f..0642db01ef 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -388,6 +388,9 @@ pub fn main(init: std.process.Init) !u8 { 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 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); diff --git a/zig/tests/contract/metal_spark25_attention.zig b/zig/tests/contract/metal_spark25_attention.zig new file mode 100644 index 0000000000..bb8871c7f5 --- /dev/null +++ b/zig/tests/contract/metal_spark25_attention.zig @@ -0,0 +1,409 @@ +//! 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], values in [0, 1); that without the cut; a bad edit. +const Case = enum { plain, spread, cutless, edit }; + +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"), + }; +} + +/// 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 } }, + }; + const spread = case == .spread or case == .cutless; + return .{ + .name = name(l, case), + .accum = .f32, + .depth = depth, + .out = .bf16, + .tiles = if (spread) .{ .{}, .{}, .{} } else .{ .{ .size = rt, .min = st.window_rows + 1 }, .{}, .{ .size = bkt, .min = 0 } }, + .served = if (spread) spread_shapes else plain_shapes, + .sums = .{ .exact = true }, + }; +} + +/// 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); + // before the M5, nax.h's fragments keep P whole; the cutless control states it whole on the M5 + const cut: f64 = if (device.tensorUnits() and case != .cutless) ar.operand_cut else 0; + 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)); + 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 half width and the cut's wider side plus bf16's rounding; 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| { + if (p.fill == .stress) { + l.* = c.boundExact(spec(p.layer, p.case), 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 + @max(d, u) + 0x1p-8 * (@abs(w) + h + @max(d, u)) + 0x1p-133; + } + } + + /// Each bf16 output within its range's ends as rounded, the cut on its own side; gate 4 and wide ranges 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 open = p.fill == .stress or !std.math.isFinite(pt) or !std.math.isFinite(h); + const e = if (open) [2]f64{ -std.math.floatMax(f64), std.math.floatMax(f64) } else ar.ends(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 spread cases and their cutless controls on the tensor units only. +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 { + 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; + if ((k == .spread or k == .cutless) and !s.device.tensorUnits()) 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); + } +}; From f768d78ec05ea42155e2e43329092c98f6e280be Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Sat, 31 Oct 2026 06:03:31 -0500 Subject: [PATCH 08/17] kernel contract: sp_prompt_attention registers on the tensor units only; it and MiniMax's mm_prompt_attention carry a known gate-1 finding, a non-finite value at a key masked for a row of its tile reaching that row as NaN, planted at every seed --- docs/kernel-contract.md | 8 ++++++-- zig/tests/contract/metal_minimax_nax_attn.zig | 4 ++++ zig/tests/contract/metal_spark25_attention.zig | 14 +++++++++----- 3 files changed, 19 insertions(+), 7 deletions(-) diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index e98f96cdcf..dd905488fe 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -125,6 +125,9 @@ Qwen3.8 Flash Next's `tf_sattn_nax` and `tf_attn256_nax` score masked keys -FLT_ 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 (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` and, on the M5, Nemotron's +`tf_attention_nax` give NaN for a row whose tile reads a non-finite value at a key masked for that row (a later key of the tile, P = 0 times Inf inside the P x V mma), +where float64 keeps the row finite; gate 1 plants one such value. 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 @@ -168,8 +171,9 @@ keeps. Two cases take operands in (0, 1], where the bias adds up; their controls at `shapes`, as does a source edit dropping the ragged K tail's last column. Before the M5 they keep fp32 whole and 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 on the M5 and whole before it, each with a P -spread case and its cutless control (M5 only), and one source edit each (a key past the window, a key past the row). +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. 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/tests/contract/metal_minimax_nax_attn.zig b/zig/tests/contract/metal_minimax_nax_attn.zig index 0a4b083b16..4c17a67148 100644 --- a/zig/tests/contract/metal_minimax_nax_attn.zig +++ b/zig/tests/contract/metal_minimax_nax_attn.zig @@ -282,6 +282,8 @@ pub const Prompt = struct { // the shortest prompt chunk with selection, a chunk from the start, one deep in the window .served = &.{ .{ .rows = st.window_rows + 1, .n = dh, .k = 2100 }, .{ .rows = 256, .n = dh, .k = 0 }, .{ .rows = 64, .n = dh, .k = 32704 } }, .sums = .{ .exact = true }, + // gate 1's Inf at key p0 + 1, which row 0 does not see but its tile reads, makes row 0's outputs NaN + .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)" }, }; } @@ -379,6 +381,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)); + // row 0 does not see key p0 + 1, which its 2-row tile reads + if (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 diff --git a/zig/tests/contract/metal_spark25_attention.zig b/zig/tests/contract/metal_spark25_attention.zig index bb8871c7f5..47aaa6e05b 100644 --- a/zig/tests/contract/metal_spark25_attention.zig +++ b/zig/tests/contract/metal_spark25_attention.zig @@ -116,9 +116,13 @@ fn spec(l: Layer, case: Case) c.Spec { .tiles = if (spread) .{ .{}, .{}, .{} } else .{ .{ .size = rt, .min = st.window_rows + 1 }, .{}, .{ .size = bkt, .min = 0 } }, .served = if (spread) spread_shapes else plain_shapes, .sums = .{ .exact = true }, + .known = if (case == .plain or case == .spread) masked_key else null, }; } +/// Gate 1's plant at key p0 + 5, which rows 0 to 4 do not see but their tile reads, makes their outputs 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 @@ -169,8 +173,7 @@ const Probe = struct { 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); - // before the M5, nax.h's fragments keep P whole; the cutless control states it whole on the M5 - const cut: f64 = if (device.tensorUnits() and case != .cutless) ar.operand_cut else 0; + 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 }; } @@ -232,12 +235,13 @@ const Probe = struct { 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 + // every row sees key p0, rows 2 on key p0 + 2; rows 0 to 4 not key p0 + 5, which their tile reads 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)); + if (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), @@ -379,7 +383,7 @@ 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 spread cases and their cutless controls on the tensor units only. +/// Both layers' pipelines with their controls, on the tensor units only: before the M5 the path is not qualified. pub const Suite = struct { a: std.mem.Allocator, device: mtl.Device, @@ -388,9 +392,9 @@ pub const Suite = struct { 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; - if ((k == .spread or k == .cutless) and !s.device.tensorUnits()) 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); From a83b3dd7feb98a5719087a9e2a372a5e33267c3f Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Sat, 31 Oct 2026 06:22:09 -0500 Subject: [PATCH 09/17] kernel contract: Gemma 4's gemma_prompt_attention (31B and 12B, sliding and global) and gemma_decode_attention with its merge state their fp32 P with the M5's operand cut, each within its range's rounded bf16 ends; P-spread cases with cutless controls, a score edit control each, and their gate-5 launches --- docs/kernel-contract.md | 5 + .../core/contract/gemma_attention_records.zig | 11 + zig/src/core/contract/registry.zig | 2 +- zig/tests/contract/metal_gemma4_attention.zig | 511 ++++++++++++++++++ zig/tests/contract/metal_limits.zig | 4 + zig/tests/contract/metal_main.zig | 3 + 6 files changed, 535 insertions(+), 1 deletion(-) create mode 100644 zig/src/core/contract/gemma_attention_records.zig create mode 100644 zig/tests/contract/metal_gemma4_attention.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index dd905488fe..825c1a695c 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -174,6 +174,11 @@ Spark-X2.6-4B's `sp_prompt_attention` registers both pipelines the engine builds 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/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/registry.zig b/zig/src/core/contract/registry.zig index 265fe5e621..8102192ac0 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -569,7 +569,7 @@ test { } /// 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("nemotron_ssd_records.zig").records ++ @import("flashnext_attention_records.zig").records ++ @import("spark25_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/tests/contract/metal_gemma4_attention.zig b/zig/tests/contract/metal_gemma4_attention.zig new file mode 100644 index 0000000000..d6ad42c4c1 --- /dev/null +++ b/zig/tests/contract/metal_gemma4_attention.zig @@ -0,0 +1,511 @@ +//! 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 }; +/// As served; P spread over [0, 1] with values in [0, 1); that case's range without the cut; scores at 2 log2e. +const Case = enum { plain, spread, cutless, edit }; + +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 }, +}; + +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", + }; +} + +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 } }; + +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 + return .{ .name = p.name, .accum = .f32, .depth = keysDepth, .ops = 4, .out = .bf16, .tiles = .{ rows, .{}, keys }, .served = served }; + } + + /// 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)); + }, + .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 half width and the cut's wider side plus bf16's rounding; 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| { + if (p.stress) { + l.* = c.boundExact(s, 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 + @max(d, u) + 0x1p-8 * (@abs(w) + h + @max(d, u)) + 0x1p-133; + } + } + + /// Each bf16 output within its range's ends as rounded, the cut on its own side; gate 4 and wide ranges 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 open = p.stress or !std.math.isFinite(pt) or !std.math.isFinite(h); + const e = if (open) [2]f64{ -std.math.floatMax(f64), std.math.floatMax(f64) } else ar.ends(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: before the M5 nax.h's fragments keep P whole, a path these ranges do not state. + 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 e7a2fdcfd6..3c6d87d556 100644 --- a/zig/tests/contract/metal_limits.zig +++ b/zig/tests/contract/metal_limits.zig @@ -70,6 +70,10 @@ fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch) !La 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 }; diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 0642db01ef..ad0158f04e 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -391,6 +391,9 @@ pub fn main(init: std.process.Init) !u8 { 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); From ee6a3902ab34e95fe558eaea063e2ae7b6317009 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Sat, 31 Oct 2026 06:31:16 -0500 Subject: [PATCH 10/17] kernel contract: Gemma 4's prompt and decode attention plant +Inf at a key masked for row 0 yet read in its tile or 32-key step, and carry the NaN it gives (P = 0 times Inf in the P x V mma) as a known gate-1 finding --- docs/kernel-contract.md | 7 ++++--- zig/tests/contract/metal_gemma4_attention.zig | 9 ++++++++- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 825c1a695c..1f6ebd95d9 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -125,9 +125,10 @@ Qwen3.8 Flash Next's `tf_sattn_nax` and `tf_attn256_nax` score masked keys -FLT_ 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 (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` and, on the M5, Nemotron's -`tf_attention_nax` give NaN for a row whose tile reads a non-finite value at a key masked for that row (a later key of the tile, P = 0 times Inf inside the P x V mma), -where float64 keeps the row finite; gate 1 plants one such value. +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) and, on the M5, Nemotron's `tf_attention_nax` give NaN for a row whose tile 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. Gate 1 plants one such value. 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 diff --git a/zig/tests/contract/metal_gemma4_attention.zig b/zig/tests/contract/metal_gemma4_attention.zig index d6ad42c4c1..ed468cc2d6 100644 --- a/zig/tests/contract/metal_gemma4_attention.zig +++ b/zig/tests/contract/metal_gemma4_attention.zig @@ -148,6 +148,10 @@ const prompt_sliding = [_]c.Shape{ .{ .rows = 77, .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; } @@ -237,7 +241,8 @@ const Probe = struct { 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 - return .{ .name = p.name, .accum = .f32, .depth = keysDepth, .ops = 4, .out = .bf16, .tiles = .{ rows, .{}, keys }, .served = served }; + const known: ?c.Known = if (p.v.case == .plain or p.v.case == .spread) .{ .gate = .nonfinite, .finding = if (p.v.kind == .prompt) prompt_finding else decode_finding } else null; + return .{ .name = p.name, .accum = .f32, .depth = keysDepth, .ops = 4, .out = .bf16, .tiles = .{ rows, .{}, keys }, .served = served, .known = known }; } /// The keys row `r` attends: a sliding row's window up to its last key, else every key up to its own. @@ -291,6 +296,8 @@ const Probe = struct { 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)); + // +Inf at p0 + 1: masked for row 0, yet read in row 0's tile (prompt) or 32-key step (decode) + 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 From ca52952dcbbce501ef9ec4126c9ff9a6a2e3c3e7 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Sat, 31 Oct 2026 06:43:53 -0500 Subject: [PATCH 11/17] kernel contract: attention_range's allowBf16 and intervalBf16 hold the bf16 range checks MiniMax's and Nemotron's prompt tests shared; Nemotron's masked-key finding moves to a case of its own, so the plain and spread cases keep a clean gate 1 --- zig/src/core/attention_range.zig | 14 ++++++++ zig/tests/contract/metal_minimax_nax_attn.zig | 17 +++------- .../contract/metal_nemotron_prompt_range.zig | 33 ++++++++----------- 3 files changed, 32 insertions(+), 32 deletions(-) diff --git a/zig/src/core/attention_range.zig b/zig/src/core/attention_range.zig index 4c8bbbdb98..dcd8b6f1c0 100644 --- a/zig/src/core/attention_range.zig +++ b/zig/src/core/attention_range.zig @@ -137,6 +137,20 @@ 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 }; diff --git a/zig/tests/contract/metal_minimax_nax_attn.zig b/zig/tests/contract/metal_minimax_nax_attn.zig index 4c17a67148..dcb6ac0740 100644 --- a/zig/tests/contract/metal_minimax_nax_attn.zig +++ b/zig/tests/contract/metal_minimax_nax_attn.zig @@ -458,24 +458,15 @@ pub const Prompt = struct { @memcpy(t.point, want); } - /// The range's half width and the cut's wider side plus bf16's rounding; gate 4 allows 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, t.down, t.up) |*l, w, h, d, u| { - 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 + @max(d, u) + 0x1p-8 * (@abs(w) + h + @max(d, u)) + 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 ends as rounded, the cut on its own side; gate 4 and wide ranges are open. + /// 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 open = t.fill == .stress or !std.math.isFinite(p) or !std.math.isFinite(h); - const e = if (open) [2]f64{ -std.math.floatMax(f64), std.math.floatMax(f64) } else ar.ends(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_nemotron_prompt_range.zig b/zig/tests/contract/metal_nemotron_prompt_range.zig index fa97fb0ff3..1bcc6ca692 100644 --- a/zig/tests/contract/metal_nemotron_prompt_range.zig +++ b/zig/tests/contract/metal_nemotron_prompt_range.zig @@ -8,14 +8,15 @@ const g = contract.gates; const ar = contract.attention_range; const D = old.D; -/// P with the tensor units' cut; P spread over [0, 1] with values in [0, 1); that case's range without the cut. -const Case = enum { cut, spread, cutless }; +/// 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", }; } @@ -41,8 +42,8 @@ const Probe = struct { 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 != .cut and fill == .random) p.spread(seed); - if (fill == .nonfinite) p.maskedInf(); + 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; } @@ -120,25 +121,16 @@ const Probe = struct { @memcpy(p.point, want); } - /// The range's half width and the cut's wider side plus bf16's rounding; gate 4's exact sums keep the exact bound. + /// 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| { - if (p.inner.stress) { - l.* = c.boundExact(s, 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 + @max(d, u) + 0x1p-8 * (@abs(w) + h + @max(d, u)) + 0x1p-133; - } + 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 ends as rounded, the cut on its own side; gate 4 and wide ranges are open. + /// 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 open = p.inner.stress or !std.math.isFinite(pt) or !std.math.isFinite(h); - const e = if (open) [2]f64{ -std.math.floatMax(f64), std.math.floatMax(f64) } else ar.ends(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]; } @@ -154,7 +146,7 @@ pub const Suite = struct { kernels: *std.ArrayList(g.Kernel), probes: std.ArrayList(*Probe) = .empty, - /// Only on the tensor units: before the M5, nax.h's fragments keep P whole and attention_ref's bound stays. + /// 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| { @@ -166,7 +158,10 @@ pub const Suite = struct { try s.probes.append(s.a, p); var spec = old.spec(.prompt, 1, false); spec.name = name(k); - if (k != .cutless) 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)" }; + 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 }); } } From dbb21964dbb7b0989a4f5710c7f5232a54c1551e Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Sat, 31 Oct 2026 06:45:21 -0500 Subject: [PATCH 12/17] kernel contract: registry's tests reach nemotron_ssd_records.zig's; the operand and SSD suites say the cut was measured on the tensor units only, not what M1-M4 keep --- zig/src/core/contract/registry.zig | 1 + zig/tests/contract/metal_mma_operand.zig | 2 +- zig/tests/contract/metal_nemotron_ssd_gemm.zig | 2 +- 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 8102192ac0..8e8e10c965 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -566,6 +566,7 @@ 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 ++ diff --git a/zig/tests/contract/metal_mma_operand.zig b/zig/tests/contract/metal_mma_operand.zig index e50e6bf057..e0684554f2 100644 --- a/zig/tests/contract/metal_mma_operand.zig +++ b/zig/tests/contract/metal_mma_operand.zig @@ -226,7 +226,7 @@ pub const Suite = struct { kernels: *std.ArrayList(g.Kernel), probes: std.ArrayList(*Probe) = .empty, - /// Only where the tensor units are: before the M5, nax.h's fragments keep fp32 operands whole. + /// 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| { diff --git a/zig/tests/contract/metal_nemotron_ssd_gemm.zig b/zig/tests/contract/metal_nemotron_ssd_gemm.zig index 1829092fe8..c0ee664685 100644 --- a/zig/tests/contract/metal_nemotron_ssd_gemm.zig +++ b/zig/tests/contract/metal_nemotron_ssd_gemm.zig @@ -340,7 +340,7 @@ pub const Suite = struct { kernels: *std.ArrayList(g.Kernel), probes: std.ArrayList(*Probe) = .empty, - /// Only on the tensor units: before the M5, nax.h's SIMD fragments keep both fp32 operands whole. + /// 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| { From 844f4be403254cd9be772c36ada96a71e1102e07 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Sat, 31 Oct 2026 06:49:24 -0500 Subject: [PATCH 13/17] kernel contract: Flash Next's prompt attention keeps its known gate-1 findings in cases of their own (padding lanes at key 0, masked keys read in attn256's blocks, rows of -Inf scores), so the plain and spread cases keep a clean gate 1; sattn's edit control named for what it counts --- zig/tests/contract/metal_flashnext_attn.zig | 56 +++++++++++++-------- 1 file changed, 35 insertions(+), 21 deletions(-) diff --git a/zig/tests/contract/metal_flashnext_attn.zig b/zig/tests/contract/metal_flashnext_attn.zig index 93cb592617..26e05b33af 100644 --- a/zig/tests/contract/metal_flashnext_attn.zig +++ b/zig/tests/contract/metal_flashnext_attn.zig @@ -25,8 +25,8 @@ 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], values in [0, 1); that range without the cut; an edit. -pub const Case = enum { cut, spread, cutless, edit }; +/// 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) { @@ -42,7 +42,8 @@ fn label(comptime w: Which, comptime k: Case) []const u8 { .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) ", the key past a row's last counted" else ", one key past the row seen"), + .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", }; } @@ -88,7 +89,7 @@ pub fn spec(w: Which, k: Case) c.Spec { // 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 == .cut or k == .spread) .{ .gate = .nonfinite, .finding = finding(w) } else null, + .known = if (k == .masked) .{ .gate = .nonfinite, .finding = finding(w) } else null, }; } @@ -97,7 +98,8 @@ 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 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", + .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", }; } @@ -239,6 +241,32 @@ pub const Probe = struct { @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; @@ -268,21 +296,7 @@ pub const Probe = struct { fillRandom(rng, p.gp, 4); switch (fill) { .random => if (p.case == .spread or p.case == .cutless) p.spread(rng), - .nonfinite => { - // key 0's values, which every row attends; three queries' heads; key 2's key, which rows 2 on see - for (0..KVH) |kv| for ([_]u16{ nan16, 0x7f80, 0xff80 }, 0..) |x, d| { - p.v[p.at(kv, 0) + d] = x; - }; - for ([_]u16{ nan16, 0x7f80, 0xff80 }, 1..) |x, r| if (r < p.rows) { - p.q[(r * H + 5 * r) * D + r] = x; - }; - p.k[p.at(1, 2) + 7] = nan16; - // row 4's head 0 scores -Inf at every key it sees: q +Inf against keys 0 to 4, each -1 in that dim - if (p.rows > 4) { - p.q[(4 * H + 0) * D + 9] = 0x7f80; - for (0..5) |key| p.k[p.at(0, key) + 9] = c.toBf16(-1); - } - }, + .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)); @@ -435,7 +449,7 @@ pub const Suite = struct { kernels: *std.ArrayList(g.Kernel), probes: std.ArrayList(*Probe) = .empty, - /// Only on the tensor units: before the M5, nax.h's fragments keep P whole, and these forms are not run there. + /// 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| { From 8a958ed7ded9c390447b5ffe1d9094d66a7c5825 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Sat, 31 Oct 2026 06:52:17 -0500 Subject: [PATCH 14/17] kernel contract: Gemma 4's masked-key finding moves to cases of its own (one per served prompt and decode pipeline, one served shape each), so the plain and spread cases keep a clean gate 1; their ranges use attention_range's allowBf16 and intervalBf16 --- zig/tests/contract/metal_gemma4_attention.zig | 43 ++++++++++--------- 1 file changed, 23 insertions(+), 20 deletions(-) diff --git a/zig/tests/contract/metal_gemma4_attention.zig b/zig/tests/contract/metal_gemma4_attention.zig index ed468cc2d6..7440584e51 100644 --- a/zig/tests/contract/metal_gemma4_attention.zig +++ b/zig/tests/contract/metal_gemma4_attention.zig @@ -22,8 +22,8 @@ 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 }; -/// As served; P spread over [0, 1] with values in [0, 1); that case's range without the cut; scores at 2 log2e. -const Case = enum { plain, spread, cutless, edit }; +/// 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 }; @@ -41,6 +41,12 @@ const variants = [_]Variant{ .{ .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 { @@ -51,6 +57,7 @@ fn nameOf(comptime v: Variant) []const u8 { .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"), }; } @@ -241,8 +248,13 @@ const Probe = struct { 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 - const known: ?c.Known = if (p.v.case == .plain or p.v.case == .spread) .{ .gate = .nonfinite, .finding = if (p.v.kind == .prompt) prompt_finding else decode_finding } else null; - return .{ .name = p.name, .accum = .f32, .depth = keysDepth, .ops = 4, .out = .bf16, .tiles = .{ rows, .{}, keys }, .served = served, .known = known }; + 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. @@ -296,8 +308,8 @@ const Probe = struct { 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)); - // +Inf at p0 + 1: masked for row 0, yet read in row 0's tile (prompt) or 32-key step (decode) - p.vp[(p.p0 + 1) * dh_ + 7] = c.toBf16(std.math.inf(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 @@ -439,25 +451,16 @@ const Probe = struct { @memcpy(p.point, want); } - /// The range's half width and the cut's wider side plus bf16's rounding; gate 4's exact sums keep the exact bound. + /// 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| { - if (p.stress) { - l.* = c.boundExact(s, 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 + @max(d, u) + 0x1p-8 * (@abs(w) + h + @max(d, u)) + 0x1p-133; - } + 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 ends as rounded, the cut on its own side; gate 4 and wide ranges are open. + /// 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 open = p.stress or !std.math.isFinite(pt) or !std.math.isFinite(h); - const e = if (open) [2]f64{ -std.math.floatMax(f64), std.math.floatMax(f64) } else ar.ends(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]; } @@ -493,7 +496,7 @@ pub const Suite = struct { kernels: *std.ArrayList(g.Kernel), probes: std.ArrayList(*Probe) = .empty, - /// Only on the tensor units: before the M5 nax.h's fragments keep P whole, a path these ranges do not state. + /// 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| { From 14c8bf12c5fff0b31722d1657162aee62bc1a52e Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Sat, 31 Oct 2026 06:56:10 -0500 Subject: [PATCH 15/17] kernel contract: Spark's sp_prompt_attention and MiniMax's mm_prompt_attention move the masked-key finding to cases of their own (tensor units only), so the plain and spread cases keep a clean gate 1; Spark's ranges use attention_range's allowBf16 and intervalBf16 --- zig/tests/contract/metal_minimax_nax.zig | 12 +++++- zig/tests/contract/metal_minimax_nax_attn.zig | 20 +++++++--- .../contract/metal_spark25_attention.zig | 39 ++++++++----------- 3 files changed, 42 insertions(+), 29 deletions(-) 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 dcb6ac0740..3272d8161f 100644 --- a/zig/tests/contract/metal_minimax_nax_attn.zig +++ b/zig/tests/contract/metal_minimax_nax_attn.zig @@ -260,6 +260,8 @@ 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 = &.{}, @@ -282,11 +284,19 @@ pub const Prompt = struct { // the shortest prompt chunk with selection, a chunk from the start, one deep in the window .served = &.{ .{ .rows = st.window_rows + 1, .n = dh, .k = 2100 }, .{ .rows = 256, .n = dh, .k = 0 }, .{ .rows = 64, .n = dh, .k = 32704 } }, .sums = .{ .exact = true }, - // gate 1's Inf at key p0 + 1, which row 0 does not see but its tile reads, makes row 0's outputs NaN - .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)" }, }; } + /// 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); @@ -311,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); @@ -381,8 +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)); - // row 0 does not see key p0 + 1, which its 2-row tile reads - if (t.rows > 1) t.v[(t.p0 + 1) * kw + 3 * dh + 7] = c.toBf16(std.math.inf(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 diff --git a/zig/tests/contract/metal_spark25_attention.zig b/zig/tests/contract/metal_spark25_attention.zig index 47aaa6e05b..42f36c1bd2 100644 --- a/zig/tests/contract/metal_spark25_attention.zig +++ b/zig/tests/contract/metal_spark25_attention.zig @@ -25,8 +25,8 @@ 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], values in [0, 1); that without the cut; a bad edit. -const Case = enum { plain, spread, cutless, edit }; +/// 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; @@ -87,6 +87,7 @@ fn label(comptime l: Layer, comptime k: Case) []const u8 { .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", }; } @@ -107,20 +108,22 @@ fn spec(l: Layer, case: Case) c.Spec { .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) .{ .{}, .{}, .{} } else .{ .{ .size = rt, .min = st.window_rows + 1 }, .{}, .{ .size = bkt, .min = 0 } }, - .served = if (spread) spread_shapes else plain_shapes, + .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 == .plain or case == .spread) masked_key else null, + .known = if (case == .masked) masked_key else null, }; } -/// Gate 1's plant at key p0 + 5, which rows 0 to 4 do not see but their tile reads, makes their outputs NaN. +/// 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. @@ -235,13 +238,14 @@ const Probe = struct { switch (fill) { .random => if (p.case == .spread or p.case == .cutless) p.spread(rng), .nonfinite => { - // every row sees key p0, rows 2 on key p0 + 2; rows 0 to 4 not key p0 + 5, which their tile reads + // 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)); - if (p.rows > 5) p.v[p.slot(p.p0 + 5) * kw + 3 * dh + 7] = c.toBf16(std.math.inf(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), @@ -340,24 +344,15 @@ const Probe = struct { @memcpy(p.point, want); } - /// The range's half width and the cut's wider side plus bf16's rounding; gate 4 allows the rounding alone. + /// 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| { - if (p.fill == .stress) { - l.* = c.boundExact(spec(p.layer, p.case), 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 + @max(d, u) + 0x1p-8 * (@abs(w) + h + @max(d, u)) + 0x1p-133; - } + 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 ends as rounded, the cut on its own side; gate 4 and wide ranges are open. + /// 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 open = p.fill == .stress or !std.math.isFinite(pt) or !std.math.isFinite(h); - const e = if (open) [2]f64{ -std.math.floatMax(f64), std.math.floatMax(f64) } else ar.ends(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]; } @@ -383,7 +378,7 @@ 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, on the tensor units only: before the M5 the path is not qualified. +/// 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, From 5745e287a087653ba6883d35877882ee2f185cd5 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Sat, 31 Oct 2026 06:57:28 -0500 Subject: [PATCH 16/17] kernel contract docs: each masked-key finding has an entry of its own, tf_attn256_nax among them; the sattn control and padding-lane reach said as they are; the cut stated as measured on the tensor units only --- docs/kernel-contract.md | 37 +++++++++++++++++++++---------------- 1 file changed, 21 insertions(+), 16 deletions(-) diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 1f6ebd95d9..09b68a97bc 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -124,11 +124,13 @@ 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 (0 x Inf is NaN where float64 gives the infinity). +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) and, on the M5, Nemotron's `tf_attention_nax` give NaN for a row whose tile 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. Gate 1 plants one such value. +(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 @@ -150,36 +152,39 @@ On the M5's tensor units, Metal also checks what nax.h's `mma_16x32` keeps of an (`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`. Before the M5, nax.h's fragments -keep fp32 operands whole, so the cases run only on the tensor units. +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` and MiniMax-M3.1-Flash's `mm_prompt_attention` state their fp32 P with it, and each -bf16 output must lie between its range's ends rounded to bf16. Nemotron's adds a case with P spread over about +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: a row's next -key counted. Before the M5 neither is registered. +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 keep fp32 whole and are -not registered; their launches are in `nemotron_ssd_records.zig`. +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`. +`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. From 75bf35c1e0b4328c24064b2fabf75e717b43f75d Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Sun, 1 Nov 2026 09:49:25 +0000 Subject: [PATCH 17/17] changelog: #1748 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) 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.