[jschmied]

The first section-2 group of #2251. On main f8dacd7 (the bf16 argmax PR #2949 is merged; this no longer stacks on it). Two commits: the change, then a host test that pins Llama's fold (head 35bd11b).

What changes

• llama/attention_nax.metal was spark25/attention_nax.metal with two folds:
  • WINDOW and RING folded to 0: the window's first key block, the ring slot of a key and value row, and the window's mask term;
  • a different kernel name.
• The two become one core/prompt_attention_nax.metal, Spark2.5's text with the kernel named by TF_PROMPT_ATTENTION. sources.zig gives each family its name: spark25_attention_nax = #define TF_PROMPT_ATTENTION sp_prompt_attention ++ it, and llama_attention_nax = ... ll_prompt_attention ++ it.
• Llama's loader (llama/contracts.zig attentionSourceHeads, which the model and the contract share) adds #define WINDOW 0 and #define RING 0 to its geometry defines. With those, the shared text folds back to Llama's old kernel.
• One contract edit follows the text. Llama's control "the prompt sees the next key" edited col > row ? kMasked, which only the old Llama text had. It now edits the shared mask line, const bool out = col > row || to col > row + 1 ||, the same edit as Spark2.5's full-layer control. One contract comment names the new path.
• Kernel names, loaders' other defines, and host dispatch are unchanged. 6 files, +12 −214.
• Not in this PR: MiniMax-M3's attention_nax adds a block bitmap (SEL) and Gemma 4's is the parent with its own cache layout; they are the next steps of this template.

Same code before and after (M3 Pro)

• Compiled code: the LLVM IR is byte-identical, old against new, for three configurations, source path aside:

  • Llama (DH 128, NQ 56, G 7, KVH 8);
  • Spark2.5's full layer (WINDOW 0, RING 0);
  • Spark2.5's sliding layer (WINDOW 1024, RING 1088).

  Each source was its loader's defines ++ the source with nax.h inlined, compiled with xcrun metal -std=metal4.1 -S -emit-llvm under -fmetal-math-mode=safe, with -fmetal-math-fp32-functions=fast and =precise.

• Kernel contract (main f8dacd7 against this change): --only llama (22 kernels: 22 pass gates 1-4 ... 12 controls: 12 red as planted), --only spark25 and --only "Llama Nanbeige" (the masked-attention cases) all give output identical apart from the SUMMARY's timings.

  • On this Mac, the gates for these attention kernels and Llama's prompt control run on tensor units only, and they skip on main and on the branch alike.
  • The spark25 attention limit gate builds the shared text here, for both layer kinds, and passes.
  • An M5 run should show Llama's "prompt sees the next key" control still red with the new edit.

• Tests: native build; test-spark25-host (12/12), test-lfm2-dspark (57/57) and test-kernel-symbols-host (3/3) pass. Lean 341 vs 341, 0 new; zig fmt --check and git diff --check clean.

• Fold test (35bd11b): zig/tests/prompt_attention_fold.zig, run by test-kernel-symbols-host, pins Llama's fold. Every WINDOW/RING use in the shared text is guarded, and Llama's loader defines both as 0. Planted red first: dropping the loader's defines, or unguarding (WINDOW && col in the shared text, fails it.

• x86 CI (Linux build + host tests on 35bd11b, test-kernel-symbols-host included): green, 730/730 steps, 9470/9542 tests passed (72 skipped), 322/335 equal to server goldens with 0 unexpected differences, zig fmt --check clean.

• Served evidence: not run. The compiled IR is byte-identical, so seeded and greedy replies, resume and concurrency against main cannot change. No refusal is added.

Proposed CHANGELOG line (yours to place): Metal: Spark2.5 and Llama prompt attention now build from one shared source (#2251).

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

[ashhart]

Thanks @jschmied, this is a tidy, well-scoped consolidation. I re-derived the equivalence from the source rather than relying only on the IR claim: the shared text folds back to Llama's old kernel at WINDOW 0, RING 0. The only code differences are the lo/slot temporaries, the window/ring mask term and the value-row index, and the loader's defines collapse all three to the old expressions.

• The change is minimal: one source moved to core/, the 204-line Llama copy deleted, a define-based kernel name, one loader define pair and one control retarget. No dispatch, sampler or other define changes.
• The control edit is right: .prompt_sees_next_key now uses const bool out = col > row ||, the token Spark2.5's full-layer control already used, and the masked-attention fixtures build Llama at WINDOW 0, RING 0, so the control still means something.
• Host checks are green: test-spark25-host and test-kernel-symbols-host pass, lean is 341 before and after with nothing new, zig fmt --check and git diff --check are clean, and nothing still points at either old path.
• Your head merges cleanly with today's main, so no rebase is needed.

Must-fixes: none. The byte-identical IR comparison for the three configs and the M5 "prompt sees the next key" control are Mac-only, so for now they rest on your device run; we'll check them on our side before landing.

Optional: no host test pins the Llama fold today. A cheap assertion that the loader's WINDOW 0 / RING 0 defines plus the shared text give exactly the old mask token would keep this honest across later edits. The CHANGELOG line you offered is fine; we place it at merge.

• Synthena (Ash's Agent)

[jschmied]

Thanks for the review. I took the optional ask. 35bd11bc8 adds zig/tests/prompt_attention_fold.zig, which runs in test-kernel-symbols-host. It fails if:

• any line of the shared text that names WINDOW or RING is not one of its four guarded forms (lo, slot, the mask const bool out = col > row || (WINDOW && col < row - WINDOW + 1);, the value row), or the mask line is missing;
• Llama's attentionSourceHeads stops emitting #define WINDOW 0 and #define RING 0;
• llama_attention_nax is no longer the ll_prompt_attention define ++ the shared source.

I planted two faults and saw it go red before green: dropping the loader's two defines, and unguarding the window term in the shared text. A second test checks that the checker itself rejects both faults.

It's a second commit on top of c595104d2, with no rebase. x86 CI on 35bd11bc8: 730/730 steps, lean 341 vs 341 with 0 new, zig fmt --check clean. The description now names this head.