[jschmied]

The section-1 group 7 of #2251, the NaN-aware argmax. It follows #2792 (the bf16 matvec, merged), whose sources.zig lines it extends. Head 9188f37, one commit on main a0e2665.

What changes

• minimax_m3/glue.metal, spark25/glue.metal and llama/glue.metal each had the same comparator (mm_better/sp_better/ll_better: NaN first, the lower id among NaNs, then the larger value, ties to the lower id). Each also had the same 1024-thread *_argmax kernel over one bf16 logits row into ids[0].
• The new core/bf16_argmax.metal holds them once: tf_bf16_argmax_better and TF_BF16_ARGMAX(NAME). sources.zig puts it ahead of the three glue files. Each file makes its kernel with TF_BF16_ARGMAX(mm_argmax) / (sp_argmax) / (ll_argmax) where the kernel was, so names and the module's order are kept.
• Left as they are:
  • Nemotron's tf_argmax_bf16 is a different kernel (many rows, an optional id map). Its glue source also feeds Nemotron's kernel identity hash.
  • tf_dsv41_better has other rules (out-of-range ids, a control switch).
• From the rest of the issue's group 7:
  • stream_rinv: ten of its eleven copies, and two of tf_rotary_scalar's six, are Flash Next's recorded kernel texts (flashnext/*_<hash>.metal), which tools/zig/flashnext_checked.py writes verbatim. They stay.
  • The other tf_rotary_scalar copies (prefill/fn_select, fn_pool_pos, fn_mrope) are the text rotary_cache.zig finds and patches in each source. It refuses one that differs from fn_rotary_scalar.metal (RotarySourceMismatch), so they are already held to one text.
• One contract comment names the new comparator. No host code changes. 6 files, +51 −109.

Same code before and after (M3 Pro, main a0e2665 against this branch)

• Compiled code: the LLVM IR of minimax_m3_glue, spark25_glue and llama_glue (with the LL_DH define its loader prepends) is identical, with the comparator's name mapped back. This holds for -fmetal-math-mode=safe with both -fmetal-math-fp32-functions=fast and =precise.

• Kernel contract (tf-kernel-contract --rotary-accuracy rotary.json; main a0e2665 against this branch):

  • --only minimax: both runs give SUMMARY 50 kernels: 50 pass gates 1-4, 0 fail, 0 known findings; 0 controls: 0 red as planted, 0 missed; 0 limit problems.
  • --only llama: both runs give SUMMARY 22 kernels: 22 pass gates 1-4, 0 fail, 0 known findings; 12 controls: 12 red as planted, 0 missed; 0 limit problems.
  • Both outputs are identical apart from the SUMMARY's timings. They cover mm_argmax and ll_argmax and their NaN rule; Llama's control "a NaN logit never wins" goes red as planted.

• Tests: native build; test-spark25-host, test-lfm2-dspark, test-minimax-info-host, test-minimax-cuda-host and test-kernel-symbols-host pass. Lean 386 vs 386, 0 new; zig fmt --check and git diff --check clean.

• x86 CI (Linux build + host tests on 9188f37): green, 700/700 steps; check-native-cuda-host compile-cuda-tests green; pr_check 0 errors; fmt, diff --check clean; lean 386 vs 386, 0 new.

Rebased onto a0e2665 after #2792 merged: git range-diff shows the patch unchanged, only sources.zig context lines moved; all checks above were rerun on the rebased head.

Not run: seeded, resume and concurrency reply checks against main, because the compiled IR is identical, so no served output can change.

Proposed CHANGELOG line (for you to place):

• Metal: MiniMax-M3's, Spark2.5's and Llama's NaN-aware bf16 argmax come from one source (part of metal: unify duplicated kernels (attention partials, norms, gemm_nax wrappers, glue, RoPE) #2251). Thanks to @jschmied.

• No personal data, machine names, internal hosts or local paths in the diff or its commits. No AI attribution lines.

[ashhart]

Thanks @jschmied. We checked the claim independently: each library built from sources.zig exactly as the loader concatenates it (with LL_DH prepended for Llama), compiled to LLVM IR with the same flags on main f0d8c8d50 and on 9188f37ec. minimax_m3_glue (22 functions), spark25_glue (22) and llama_glue (33) are the same code; the only differences are the source path in the debug metadata. It's queued to land once CI is green.

• Synthena (Ash's Agent)

[ashhart]

Thank you @jschmied. 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]

Thanks @jschmied! This is in TensorFold 1.2.8, out now: https://github.com/ashhart/TensorFold/releases/tag/v1.2.8

• Synthena (Ash's Agent)