kernel contract: the M5's fp32 mma_16x32 operand cut, stated for every fp32-operand nax kernel (#1546)

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.

docs/kernel-contract.md

@@ -121,6 +121,16 @@ HIP's prompt attention tile (`tf_fa_paged_*`, `tf_fa_wide_*`) writes 0 for a row
 float64 gives NaN, as its running max is `fmaxf` and a row at max -inf takes p = 0; and an infinite value at a key
 whose fp32 weight is 0 (past the row's causal edge in its 64-key tile, or under fp32's least) gives NaN where float64
 gives inf or a finite value, as its P x V takes every key of the tile.
+Qwen3.8 Flash Next's `tf_sattn_nax` and `tf_attn256_nax` score masked keys -FLT_MAX, the value each row's running
+max starts from, so a row whose every score is -Inf weighs its masked keys 1 and gives a finite average where float64
+gives NaN. `tf_sattn_nax`'s padding lanes also read key 0 at p = 0, so a non-finite value in key 0's row reaches every
+row whose key count is not a multiple of 32 (0 x Inf is NaN where float64 gives the infinity).
+MiniMax-M3.1-Flash's `mm_prompt_attention`, Spark-X2.6-4B's `sp_prompt_attention`, Gemma 4's `gemma_prompt_attention`
+(31B and 12B, sliding and global), Nemotron's `tf_attention_nax` and Qwen3.8 Flash Next's `tf_attn256_nax` give NaN
+for a row whose tile or block reads a non-finite value at a key masked for that row, and Gemma 4's
+`gemma_decode_attention` for a decode row whose 32-key step reads one: P = 0 times Inf inside the P x V mma, where
+float64 leaves the masked key out. These were seen on the M5, where their cases run. Each kernel has an entry of its own
+that plants such a value and carries the finding, so its other entries keep a clean gate 1.
 
 HIP’s MLX stream tile keeps its finite dot2 sums unchanged. A non-finite result with an actual non-finite
 operand is recomputed by full quantization groups using the factored float64 formula. The plain tile is
@@ -138,6 +148,43 @@ bits) and simd_qmm_bits (3, 6 and 8 bits), and the lane projection (6-bit groups
 formats on the M5's tensor units, where its reg kernel must equal the cooperative one bit for bit), and tile_qmm's
 prompt tiles and their decode twin (4, 6 and 8 bits; the twin bit-equal to the tiles). Gate 5 also holds the rest of
 tile_qmm's pipelines and GLM-5.3-Flash's hc_expand_10.
+On the M5's tensor units, Metal also checks what nax.h's `mma_16x32` keeps of an fp32 operand
+(`metal_mma_operand.zig`): the left, the right and both (attention's fp32 P, `gemm_nax`'s f32 kernels), each output a
+single product bit-equal to float64 on operands with their low 13 mantissa bits cleared. That is 10 bits kept, rounded
+toward zero, fp32's exponents and subnormals kept, NaN and the infinities passed. Four controls state the operand as
+fp32 or as rounded to nearest at 10 bits, on each side, and must go red at `shapes`. The cut was measured on the tensor
+units only, so the cases run only there; what M1-M4 keep of an fp32 operand is not measured.
+`core/attention_range.zig` names that cut (`operand_cut`, 2^-10) as a form's `p_cut`: one-sided, since each p only
+shrinks, so it moves an output down by up to 2^-10 of its positive values' share and up by its negative values'. On the
+M5, Nemotron's `tf_attention_nax` states its fp32 P with it, and MiniMax-M3.1-Flash's `mm_prompt_attention` does on
+every chip (before the M5 that only widens its range); each bf16 output must lie between its range's ends rounded to
+bf16. Nemotron's adds a case with P spread over about
+[2^-6, 1] and values in [0, 1), where the cut's bias adds up; its control, the same range without the cut, goes red at
+`shapes`. Before the M5, `tf_attention_nax` keeps attention_ref's bound.
+Qwen3.8 Flash Next's prompt attention on the M5 does too: `tf_sattn_nax` (the indexer's dense and selected rows) and
+`tf_attn256_nax` (the causal chunk without the indexer), head dim 256, 24 heads over 2. Each bf16 output,
+bf16(bf16(o) * bsig(gate)), must lie between its range's ends rounded to bf16 and through the gate's rounded sigmoid.
+Each adds a P-spread case whose cutless range is a planted control, and a source edit gate 3 must catch: `tf_sattn_nax`
+counting its padding lane (key 0 again), `tf_attn256_nax` seeing one key past the row. Before the M5 neither is
+registered.
+On the M5, Nemotron's SSD steps (`prefill_mixers.zig`'s Mamba mixer) run `gemm_nax`'s five f32 kernels with both
+operands fp32, so each product enters cut, smaller by up to 1 - (1 - 2^-10)^2. Their entries
+(`metal_nemotron_ssd_gemm.zig`, y = sur x dtx and the state dd^T x bh^T a head, at the step's served row counts) state
+it as a one-sided range: float64's value less that share of the positive terms and plus that of the negative, beside
+the fp32 sums' `mma_step` roundings for every k16 step. Gate 4 runs 11-bit integers times 5-bit ones, which the cut
+keeps. Two cases take operands in (0, 1], where the bias adds up; their controls state the operands as fp32 and go red
+at `shapes`, as does a source edit dropping the ragged K tail's last column. Before the M5 they are not registered;
+their launches are in `nemotron_ssd_records.zig`.
+Spark-X2.6-4B's `sp_prompt_attention` registers both pipelines the engine builds (sliding layers over a 512-key window
+in a 2,560-slot ring, and full layers), its fp32 P stated with the cut, each with a P spread case
+and its cutless control, and one source edit each (a key past the window, a key past the row); before the M5 they are
+not registered.
+On the M5, Gemma 4's `gemma_prompt_attention` (the 31B's and 12B's sliding and global layers) and
+`gemma_decode_attention` with `gemma_attention_merge` (their global layers; the 12B's assistant launches the 12B's)
+state their fp32 P with the cut too, each form's counts taken from the kernel's own steps, at the sliding ring, image
+rows and decode chunks the engine serves. A P-spread case on the 31B's global prompt and decode has a cutless control
+each, red at `shapes`, and a score edit (2 log2e) in each kernel is a control. Their launches are in
+`gemma_attention_records.zig`.
 LFM2.5-24B-A2B's shared rounds: the greedy pick (`lfm2_argmax`) and the conv and attention kernels that run every
 stream's rows in one dispatch, each bit-equal to the per-stream kernels it replaces. Four controls edit the source: a tie
 to the highest id, a nonfinite logit not flagged, and every stream reading or writing the first stream's state or cache.

zig/src/core/attention_range.zig

@@ -18,6 +18,8 @@ pub const Form = struct {
     chunk_exp: bool,
     /// P's error into P x V, relative: bf16 2^-8, two bf16 halves 2^-16, fp32 0.
     p_rel: f64,
+    /// An fp32 P cut toward zero as it enters P x V, relative and one-sided (the M5's tensor units: operand_cut).
+    p_cut: f64 = 0,
     /// The chains of o's terms p_j v_j and of l's terms p_j, as roundings of their terms' magnitude.
     num: f64,
     den: f64,
@@ -33,6 +35,9 @@ pub const Form = struct {
 /// An mma k16 step's error in 2^-24 of |acc| + sum |a b|: 17 terms cut at the largest's ulp, then rounded down.
 pub const mma_step = 36;
 
+/// What the M5's tensor units keep of an fp32 operand: 10 mantissa bits, cut toward zero (metal_mma_operand.zig).
+pub const operand_cut = 0x1p-10;
+
 /// window_rows.cu's decode at head dim d over n keys: lanes' chains of d / 32 then a 5-level tree, 256-key chunks.
 pub fn decode(d: usize, n: usize) Form {
     const chunks: f64 = @floatFromInt(n / 256 + 2);
@@ -105,6 +110,47 @@ pub fn head(q: []const f64, k: []const f64, v: []const f64, n: usize, d: usize,
     }
 }
 
+/// P's cut after `head` (w, tau as it left them): each p only shrinks, so positive values move o down, negative up.
+pub fn cutSides(f: Form, w: []const f64, tau: []const f64, v: []const f64, n: usize, d: usize, down: []f64, up: []f64) void {
+    var l: f64 = 0;
+    var tmax: f64 = 0;
+    for (w[0..n], tau[0..n]) |p, t| {
+        l += p;
+        if (std.math.isFinite(t)) tmax = @max(tmax, t);
+    }
+    const kappa = (1 + tmax) / (1 - tmax);
+    for (0..d) |e| {
+        var pos: f64 = 0;
+        var neg: f64 = 0;
+        for (0..n) |j| {
+            const x = v[j * d + e];
+            if (x > 0) pos += w[j] * x else neg -= w[j] * x;
+        }
+        const finite = std.math.isFinite(pos) and std.math.isFinite(neg) and tmax < 0.5;
+        down[e] = if (finite) f.p_cut * pos / l * kappa * (1 + 0x1p-20) else 0;
+        up[e] = if (finite) f.p_cut * neg / l * kappa * (1 + 0x1p-20) else 0;
+    }
+}
+
+/// A bf16 output's ends, point - half - down to point + half + up, each rounded to nearest as the output is.
+pub fn ends(point: f64, half: f64, down: f64, up: f64) [2]f64 {
+    return .{ bfr.bf16(point - half - down), bfr.bf16(point + half + up) };
+}
+
+/// A bf16 output's allowance: the range's half width and the cut's wider side, plus the output's rounding.
+pub fn allowBf16(point: f64, half: f64, down: f64, up: f64) f64 {
+    if (!std.math.isFinite(point)) return 0;
+    if (!std.math.isFinite(half)) return std.math.floatMax(f64); // scores too far apart for the first order: any value
+    const h = half + @max(down, up);
+    return h + 0x1p-8 * (@abs(point) + h) + 0x1p-133;
+}
+
+/// A bf16 output's signed interval: its range's rounded ends, open where float64 or the range is not finite.
+pub fn intervalBf16(point: f64, half: f64, down: f64, up: f64) [2]f64 {
+    if (!std.math.isFinite(point) or !std.math.isFinite(half)) return .{ -std.math.floatMax(f64), std.math.floatMax(f64) };
+    return ends(point, half, down, up);
+}
+
 /// What a check saw: outputs whose range holds one bf16 value, several, and outputs non-finite in float64.
 pub const Stats = struct { decided: usize = 0, ranged: usize = 0, nonfinite: usize = 0 };
 
@@ -175,3 +221,21 @@ test "a weight's spread moves an output by about its share of |v - o|" {
     head(&.{1}, &.{ 30, 30 }, &.{ 0, 1 }, 2, 1, 1, prompt(256, 2, true), &w, &t, &pt, &hf);
     try std.testing.expect(hf[0] > small and hf[0] < 0x1p-9);
 }
+
+test "the cut widens a range on one side: down by positive values' share, up by negative values'" {
+    var w: [2]f64 = undefined;
+    var t: [2]f64 = undefined;
+    var pt: [2]f64 = undefined;
+    var hf: [2]f64 = undefined;
+    var down: [2]f64 = undefined;
+    var up: [2]f64 = undefined;
+    var f = prompt(256, 2, false);
+    f.p_rel = 0;
+    f.p_cut = operand_cut;
+    // equal scores, values 2 and 6 in one column, -2 and -6 in the other: o = 4 and -4
+    head(&.{ 1, 0 }, &.{ 1, 0, 1, 0 }, &.{ 2, -2, 6, -6 }, 2, 2, 1, f, &w, &t, &pt, &hf);
+    cutSides(f, &w, &t, &.{ 2, -2, 6, -6 }, 2, 2, &down, &up);
+    try std.testing.expectEqualSlices(f64, &.{ 4, -4 }, &pt);
+    try std.testing.expect(down[0] >= 4 * operand_cut and down[0] < 4.01 * operand_cut and up[0] == 0);
+    try std.testing.expect(up[1] >= 4 * operand_cut and up[1] < 4.01 * operand_cut and down[1] == 0);
+}

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 },
+};

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 },
+};

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));
+}

zig/src/core/contract/registry.zig

@@ -566,10 +566,11 @@ const paging = @import("paging_pending.zig");
 test {
     _ = paging;
     _ = @import("minimax_records.zig");
+    _ = @import("nemotron_ssd_records.zig");
 }
 /// The Metal collector keeps measured slots and the restored vision declarations.
 pub const all_metal = metal ++ @import("delta_slots_records.zig").records ++ @import("flashnext_pending.zig").records ++ @import("simd_pairs_records.zig").records ++ @import("metal_pair_records.zig").records ++ @import("glm_rows_records.zig").records ++ @import("weight_only_records.zig").metal ++ @import("hc_activation_records.zig").metal ++ @import("qwen4_records.zig").metal ++ @import("lfm2_records.zig").records ++ @import("nemotron_experts_records.zig").records ++ @import("nemotron_pages_records.zig").records ++ @import("nemotron_attention_records.zig").records ++
-    @import("moss_image_records.zig").records ++ @import("dsv41_metal_records.zig").records ++ @import("minimax_records.zig").records;
+    @import("nemotron_ssd_records.zig").records ++ @import("flashnext_attention_records.zig").records ++ @import("spark25_records.zig").records ++ @import("gemma_attention_records.zig").records ++ @import("moss_image_records.zig").records ++ @import("dsv41_metal_records.zig").records ++ @import("minimax_records.zig").records;
 const learning = @import("registry_learning.zig");
 pub const gradient_pending = learning.gradient_pending;
 pub const qwen_learning_pending = learning.qwen_learning_pending;

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 },
+};