Skip to content

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

Merged
ashhart merged 17 commits into
ashhart:mainfrom
BobClawblaw:mma-operand-cut
Nov 1, 2026
Merged

ashhart merged 17 commits into
ashhart:mainfrom
BobClawblaw:mma-operand-cut

Conversation

@BobClawblaw

@BobClawblaw BobClawblaw commented Oct 31, 2026 •

Copy link
Copy Markdown
Contributor

Steps 1-3 of the fix in #1546 (#1546 (comment)). It changes only tests, specs and docs, so no kernel's output changes. @ashhart holds the claim on #1546; this is offered under the pending ATC pair proposal, as a draft, for you to take, split or close.

Step 1: the operand cut as a contract case

zig/tests/contract/metal_mma_operand.zig runs nax.h's mma_16x32 with an fp32 left operand, an fp32 right operand (left bf16) and both fp32. Each output is a single product, so each must be bit-equal to float64 on operands with their low 13 mantissa bits cleared (bits & 0xffffe000). That means 10 bits kept and rounded toward zero, with fp32's exponents and subnormals kept and NaN/Inf passed. The right operand is cut the same way as the left. Four controls state the operand as fp32 or as rounded to nearest at 10 bits, on each side, and go red at shapes. The cases register only on the tensor units: M1-M4 were not measured.

Step 2: a named one-sided term

core/attention_range.zig adds operand_cut (2^-10), the Form field p_cut, cutSides (each p only shrinks, so an output moves down by up to 2^-10 of its positive values' share and up by its negative values'), ends (a range's bf16-rounded ends), and allowBf16/intervalBf16, which the attention tests share. On the M5, Nemotron's tf_attention_nax (metal_nemotron_prompt_range.zig) states its fp32 P with the term. It also has a case with P spread over about [2^-6, 1] and values in [0, 1), where the bias adds up; its control, the same range without the cut, goes red.

Step 3: gates 1-4 for kernels 2-7, the term for kernel 8, gate-5 records

Kernel Test Controls
gemma_prompt_attention (31B, 12B; sliding, global) and gemma_decode_attention + gemma_attention_merge metal_gemma4_attention.zig P-spread cutless (prompt, decode), score at 2 log2e (each)
sp_prompt_attention (sliding, full) metal_spark25_attention.zig P-spread cutless, a key past the window / past the row
tf_sattn_nax, tf_attn256_nax metal_flashnext_attn.zig P-spread cutless, padding lane counted / one key past the row
the five tf_gemm_nax_f32_* (Nemotron SSD y and state) metal_nemotron_ssd_gemm.zig operands stated as fp32, ragged K tail's last column dropped
mm_prompt_attention metal_minimax_nax_attn.zig its form said P multiplies V as bf16; P is fp32, so p_rel 0 and p_cut

Each form's counts come from the kernel's own steps, with a one-line comment each. Gate-5 records are added in nemotron_ssd_records.zig, gemma_attention_records.zig, spark25_records.zig and flashnext_attention_records.zig. No bound was widened to pass.

Known findings (shipped kernels, unchanged here)

Each has an entry of its own that plants the value, so every other entry keeps a clean gate 1:

  • A non-finite value at a key masked for a row but read in its tile, block or 32-key step gives that row NaN (P = 0 times Inf inside the P x V mma). This affects mm_prompt_attention, sp_prompt_attention, gemma_prompt_attention, gemma_decode_attention, tf_attention_nax and tf_attn256_nax.
  • tf_sattn_nax's padding lanes read key 0, so a non-finite value there reaches every row whose key count is not a multiple of 32.
  • tf_sattn_nax and tf_attn256_nax start the running max at -FLT_MAX, so a row whose every score is -Inf weighs its masked keys 1 and gives a finite average where float64 gives NaN.

Measured (M5 Ultra, macOS 27)

  • tf-kernel-contract --only metal: SUMMARY 342 kernels: 342 pass gates 1-4, 0 fail, 15 known findings; 368 controls: 368 red as planted, 0 missed; 0 limit problems. Gate 5: 439 registered Metal launches against M1-M5, 0 problems. (Head 3263e4b, rebased onto main 3ab49e3.)
  • zig build test, test-kernel-contract, test-golden, test-server-security test-pull-security and fuzz-cortheon exit 0. zig fmt --check is clean. The lean check reports 0 new problems (424 on main, 424 here).
  • Not run: M1-M4. Every new case registers on the tensor units only, so no records were built for M1/M2 with --offline.

Open points

  • MiniMax's form num = 4*mma_step + blocks + 2 may undercount the P x V chain and the keeps' exp error. It passes, and it is unchanged here.
  • Gemma's sliding decode pipeline is compiled but never launched, so it is not covered.
  • Gemma's prompt kernel reads K up to the block end, and its decode kernel reads K and V up to a 32-key step end. With a global capacity that is not a multiple of 256 (prompt) or 32 (decode), the last kv head reads past the cache buffer. This comes from reading the source and is untested: the tests use aligned capacities.

@ashhart

ashhart commented Nov 1, 2026

Copy link
Copy Markdown
Owner

Thanks @BobClawblaw, and sorry for the wait. We'll take it as one PR, no split: please mark it ready. It changes only tests, specs, records and docs. The operand case pins what the M5 does (10 mantissa bits, cut toward zero, bit-equal to float64 on cut operands), and the nearest-rounding control going red shows it's a cut. Every fp32-P kernel now states the one-sided term, with a cutless control that goes red. Thanks also for planting the three known findings separately, which keeps every other gate 1 clean. Whether to change any outputs is #1772's question, not this PR's. On your open points: we'll check whether Gemma's served cache capacity can ever be unaligned. If it can, the read past the buffer gets its own issue and fix. The MiniMax num form and Gemma's unlaunched sliding decode pipeline can be follow-ups.

@BobClawblaw

Copy link
Copy Markdown
Contributor Author

Thanks @ashhart. Marked ready. I rebased onto current main (3ab49e3) so CI picks up main's contract_images fix; the red zig jobs were that test failing on the old base. On the M5, tf-kernel-contract --only metal now gives SUMMARY 342 kernels: 342 pass gates 1-4, 0 fail, 15 known findings; 368 controls: 368 red as planted, 0 missed; 0 limit problems, and 439 registered Metal launches against M1-M5 with 0 problems. The counts grew with main's new registrations. zig build test, test-kernel-contract, test-golden, the security tests and fuzz all exit 0, fmt is clean, and lean shows 424 problems, the same as main. I'll take the MiniMax num form and the Gemma sliding decode pipeline as follow-ups, and I'll watch for your answer on Gemma's cache alignment.

@ashhart

ashhart commented Nov 1, 2026

Copy link
Copy Markdown
Owner

Thanks @BobClawblaw, for marking it ready and for the rebase. The 16 commits are the ones we read, so our call
stands, and the SUMMARY growth is main's own registrations. Main has moved again since (#1706's records and a HIP
known finding), which gives two small appended conflicts in registry.zig's all_metal line and the known-findings
paragraph. We'll keep both sides when we land it, so no need to rebase again. We'll answer on Gemma's cache
alignment once we've checked the served capacities.

BobClawblaw and others added 17 commits November 1, 2026 11:04
…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
…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
…inite 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
… 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
…x, 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
…ngs 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
…l 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
…ly; 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
…ng 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
…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
…e 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
…he operand and SSD suites say the cut was measured on the tensor units only, not what M1-M4 keep
… 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
…wn (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
…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
…, 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
@ashhart
ashhart merged commit 7651ce6 into ashhart:main Nov 1, 2026
8 of 9 checks passed
@ashhart

ashhart commented Nov 1, 2026

Copy link
Copy Markdown
Owner

Thank you @BobClawblaw. This is merged into main with your commits under your own name, so you're listed as a TensorFold contributor. It ships in the next release, credited to you in the changelog and the release notes.

@ashhart

ashhart commented Nov 1, 2026

Copy link
Copy Markdown
Owner

Thanks @BobClawblaw! The M5 fp32 operand cut, now stated in the Metal kernel contract for every kernel with one, is in 1.2.1: https://github.com/ashhart/TensorFold/releases/tag/v1.2.1

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants