[jschmied]

What this changes

Includes core/bf16_round.zig with #1056's sign fix (an underflow to zero keeps its sign, checked by raw bits), so the steps share one identical file.

Step 5 of #1048: window_attention serves head dim 256 with any number of query heads per key head (up to 16, so GQA 2 included), a per-row last key, and Metal's rounding points as an option. Nothing calls it yet.

It is one operator, as agreed on #1048. WindowAttention keeps its API, and decode and prompt take head dim 256 through the same Window. The new arithmetic lives in a new image, window_rows.cu, so window_attention.cu, its fatbin and every current launch stay as they are. The only caller today is Kolibri-1 (head dim 128); its kernels, arguments and timings are unchanged (receipt below).

• Decode: 256-key chunks at absolute key positions, one block for each (row, key head, chunk), then a merge in chunk order. A chunk scores its keys in fp32: each lane runs an fma chain over 8 dims, then a xor tree over the lanes. It takes the chunk's max, computes expf(s - m) and their sum, and then o = sum_j p_j v_j: up to 8 key runs, each an fma chain in key order, summed in run order. The merge computes o = sum_c a_c o_c / sum_c a_c l_c with a_c = expf(m_c - M), as an fma chain over the chunks in order. Probabilities stay fp32 throughout, as in Metal's decode. The split depends only on key positions, so a row's bits don't depend on the batch, the slot or the cache length (chunks with none of its keys are skipped exactly).

• Prompt: tensor-core tiles of 16 rows of one query head, with 2 warps per row band. Each warp takes a 128-dim slice: its share of q . k, summed in slice order, then its 128 output columns. The keys come in 32-key tiles at absolute positions, each folded into o by one fma. A row's bits don't depend on how the prompt is chunked.

• Window.p32 (off by default): P enters P x V as two bf16 halves, bf16(P) and bf16(P - bf16(P)). That is fp32's probabilities to 2^-16, as gemma4/prompt_attention.metal keeps them ("bf16 moved an 8k prompt's logits off mlx-lm"). With it off, P is one bf16, as window_attention.cu's prompt kernel rounds it.

• promptRows(..., hi): row i sees keys [p0 + i - win + 1, hi[i]] (int32, p0 + i <= hi[i] < p0 + n), so an image's rows see each other, as prompt_attention.metal's HI gives them. Without it (null, and in prompt), every row's last key is its own position.

• Scale is a parameter (Gemma uses 1.0). Window 0 runs over a linear cache (ring its length) in both forms.

• A NaN or infinite output re-runs its head in float64 (attention_ref's arithmetic, off the common path), so non-finite outputs match float64's class.

• core/attention_range.zig computes a head in float64 at Metal's rounding points (only the output's), plus the range each form's fp32 steps allow around it:

  • the score spread (fma chains and trees; for the mma, a stated model: a k16 step within 36 · 2^-24 of |acc| + sum |a b|);
  • expf (2 ulps) and __expf (2 + 1.173|x| ulps);
  • the chains of o and l, and P's representation.

  The checker rounds both ends of the range to bf16 and compares in value order, -0 below +0. A one-value range must match bit for bit, NaN only matches NaN, and infinities are compared by their bits (cuda norms: bf16 residual tails with a branch norm, a layer scalar and MLX's rounded rows #1056's checker rules).

Gemma 4 uses p32 with hi on its sliding layers (12B 16/8 heads, 31B 32/16; window 1024 over a ring of 1536).

Receipt

• Environment: TensorFold d1018b31 (1.1.1) plus this commit; one DGX Spark (GB10, sm_121, CUDA 13.0, nvcc 13.0.88). No checkpoint needed. Shapes come from families/gemma4/config.zig. Inputs follow mlx-community/gemma-4-12B-it-4bit's norm weights, which are constant within a layer: q rows at RMS 0.95 to 1.05, k at 0.06 to 0.13, v at unit RMS. Every few rows one key equals a query head at k's RMS, a score near 30, so softmax is peaked as well as flat.

• Exactness, tf-cuda-test window-rows:

  • Cases: 12B and 31B sliding decode (1, 3, 17 and 64 rows, positions to 40,000, three streams, wrapped rings); a no-window decode over a linear cache; 12B sliding prompt chunks of 512 at p0 0, 1,300 (with an image span, rows 40 to 300) and 5,000 (the ring wraps) with P32; a 31B chunk of 333 rows with bf16 P; and a no-window prompt over a linear cache.
  • Results: 2,606,808 outputs bit-equal to float64's value at the rounding points and 8,689,960 within their range (bf16 P is never decided at its 2^-8). Smaller decode calls' bits equal the widest call's, the linear cache declared 16x longer gives the same bits, and a prompt chunk split in two gives every row's bits.

• Decode speed, new shapes (tf-cuda-test window-bench gemma, GPU µs, median of 25, at positions 1k / 8k / 32k / 128k):

  shape	1k	8k	32k	128k
  12B sliding decode, 1 row (16/8 x 256)	34.3	35.2	35.0	35.1
  12B sliding decode, 4 rows	159.9	158.5	159.8	158.0
  31B sliding decode, 1 row (32/16 x 256)	43.4	40.1	44.4	40.0

  One 12B row reads 8 MiB of keys and values: 34 µs is about 245 GB/s.

• Prompt speed, new shapes (512-row chunk ending at the context):

  shape	1k	8k	32k	128k
  12B sliding prompt, P32	262.5	326.0	323.9	327.0
  12B sliding prompt, bf16 P	229.7	285.8	287.0	287.0

  P32 costs 13 to 14 %.

• Current callers' speed: unchanged. window_attention's only caller is Kolibri-1. Its served shapes were timed with the same bench on main (with the bench file added) and on this branch, ABBA x2 (A B B A A B B A), GPU µs, median of each side's 4 runs:

  Kolibri-1 call	main	this branch	change
  decode tf_window_decode_tc, 1 row	12.7	12.6	-0.4 %
  decode, 4 rows	12.7	12.6	-0.8 %
  decode, 16 rows	22.9	22.9	+0.2 %
  decode, 64 rows	308.6	308.2	-0.1 %
  sliding prompt wattn_kernel<128,8,4,8>, 2,048 rows at 0	566.1	572.5	+1.1 %
  sliding prompt, 8,192 rows at 8,192	2,373.4	2,379.1	+0.2 %
  full prompt, 2,048 rows at 8,192	9,050.4	9,047.5	-0.0 %
  full prompt, 2,048 rows at 30,720	30,281.5	30,282.8	+0.0 %

  The same kernels from the same bytes run in both, so the differences are run-to-run spread. The +1.1 % row's runs overlap: main 561.5 to 571.6, branch 562.3 to 576.1.

• Kernel contract (docs/kernel-contract.md), tf-cuda-test kernel-contract (built -Dnvcc, ReleaseFast) on the GB10:
  SUMMARY 88 kernels: 88 pass gates 1-4, 0 fail, 0 known findings; 13 controls: 13 red as planted, 0 missed; 1533 limit problems
  (main: SUMMARY 85 kernels: 85 pass gates 1-4, 0 fail, 0 known findings; 7 controls: 7 red as planted, 0 missed; 1533 limit problems).

  • 3 new entries (zig/tests/cuda/contract_window_rows.zig), at Gemma 4 12B's sliding shapes: decode (16 and 64 rows to position 40,000, one row at 1,100; key tiles' edges), P32 prompt with image rows, and bf16-P prompt (512 rows at p0 0 and 5,000, 96 rows at 1,300; row and key tiles' edges).
  • Gates 1 and 3 are the entry's own check: every output in float64's range at the rounding points. Gate 1 plants a NaN query and a NaN key, +inf in a value that only later rows see, +inf and -inf in one column, and +inf at a key whose float64 weight is e^-200 where fp32's is 0. That last one is the float64 path's case.
  • Gate 4: every key alike, so every weight is exactly 1 and the values are mirrored integers. fp16 sums overflow (sums.f16 = 100); bf16 sums miss by 77x (decode), 9.4x (P32) and 0.6x (bf16 P, whose own bound is P's 2^-8), so the spec says why (sums.bf16 = 0).
  • 6 controls, each red at gate 3: decode partials whose window starts one key early, the merge without exp(m_c - M), the merge rounding toward zero, P32 prompt rows whose window starts one key early, P32 prompt rows that ignore an image's last key, and P rounded to bf16 under the P32 range.
  • Gate 5: 4 launches recorded (tf_wrows_part_256, tf_wrows_merge_256, tf_wrows_prompt_256 and _p32; 48 to 255 registers, the prompt rows at 80 KiB dynamic shared memory, opt-in). tf-cuda-test kernel-contract-limits over fatbins for sm_80, 86, 89, 120 and 121 finds all records current. The run over every image reports 877 problems, as main's does.
  • The limit problems in the SUMMARY line are main's own (records 1.1.1 lists without a compiled usage yet); this commit adds none.

• Default path unchanged: window_attention.fatbin for sm_80, 86, 89, 120 and 121 is byte-identical to main's. cuobjdump -sass of its five kernels matches on all five SMs (25 of 25), as do tree_attention, block_attention, prefill_attention and nemotron_attention (470 of 470) and the existing controls (35 of 35). tf-cuda-test window passes as on main.

• Tests run: zig build test from an empty cache: 213/213 steps, 2,964 of 2,978 tests passed, 14 skipped (main: 211/211 steps, 2,959 of 2,973 tests passed, 14 skipped). tf-cuda-test window-rows and window pass. zig fmt --check is clean. lean_check.py: 444 findings, as on main.

• Not run, and why: Metal (no Metal change). sm_80, 86, 89 and 120 only through their recorded compiles (gate 5); the tensor-core sum model behind the prompt range was checked on sm_121 only. No model calls it yet (step 8 of Proposal: Gemma 4 on CUDA, built on the shared CUDA operators (shared extensions first) #1048).

Overlap with #798

#798 rewrites window_attention.cu: it moves the split decode to window_split.cuh and adds the NaN and infinity re-runs (attn_odd.cuh). This PR doesn't touch that file. Its new kernels are in window_rows.cu, with their own float64 re-run. A test merge onto a scratch merge of #746 and #798 (on d46d17a7) leaves window_attention.cu, window_attention.zig and window_split.cuh clean. One conflict is new: registry.zig's record line, where #798's records end and this PR appends its file; it resolves by keeping both. build/cuda.zig and contract_tests.zig already conflict between that scratch and 1.1.1. The remaining shared lines are the registration lists: contract_tests.zig, registry.zig's record line, build/cuda.zig's host tests and the end of contract_controls.cu.

Checklist

• Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
• Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. (No serving path changes; a row's bits are its own at any batch, slot, cache length or prompt chunking.)
• No precision traded for speed. If the bits change, the description says which sums change and why.
• A new or changed GPU kernel is registered with the kernel contract (docs/kernel-contract.md); each backend's SUMMARY line is in the receipt.
• Prompt processing is no slower than the last release.
• The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not.
• New tests fail before the change, pass after it, and skip cleanly without their dependency.
• Comments and docstrings are one line. No measurements or history in the source.
• No personal data, machine names, internal hosts or local paths. No AI attribution lines.

[ashhart]

Thanks @jschmied for putting the new arithmetic in its own window_rows.cu image so window_attention.cu and Kolibri-1's kernels stay byte-identical, and for the six planted controls. We are carrying this onto current main ourselves with your commit kept as yours, and a separate commit of ours will wire up the RMS test at zig/tests/cuda/window_rows_check.zig:160 that nothing runs today, keep zig/build/cuda.zig under 600 lines, and add contract entries for odd query-to-key-head ratios, 8 and 16, since the API takes any ratio up to 16 (zig/src/cuda/window_attention.zig:145) while the contract covers only 2. One choice is yours: the description says a non-finite output re-runs its head in float64, but the prompt kernel ORs the flag across the whole block (zig/kernels/cuda/window_rows.cuh:526) and odd() then recomputes every row and head in it (lines 541 to 547), so a finite row that shares a block with a bad one gets float64 bits that change when the chunk boundary moves; would you like the re-run limited to the affected row and head, as the description says, or is there a reason to keep the whole block? Whichever you choose, we will add a device test with a NaN query in a neighbouring row that compares the finite row alone, beside it and split across chunks, with both bf16 P and P32. Then it goes to the GB10 kernel contract and a Kolibri-1 served-equality check, since every WindowAttention load now loads the new image, and lands; you need not rebase, and if you are in the middle of a change here or on #1124, please tell us so we do not collide.

- Synthema, Ash's AI Assistant.

[jschmied]

Thanks for carrying it. Please limit the float64 re-run to the affected row and head, as the description says. There is no reason to keep the whole block: the block-wide OR was a shortcut, and it makes a finite row's bits depend on where the chunk boundary falls, which breaks resumed == fresh across different cuts. The neighbouring-NaN device test you describe, alone, beside and split across chunks with bf16 P and P32, is the right check for it. Not in the middle of a change here or on #1124.

[ashhart]

Thanks @jschmied. window_attention at head dim 256 (any GQA, a row's last key, fp32 probabilities) is on main (a4c16a5) with the maintainer fix-up, and the kernel contract passes (79 of 79, every planted control red). On a GB10, Qwen3.8-27B tokens are identical to main and cold prefill is unchanged. Closing as landed; it ships in 1.1.4.

- Synthema, Ash's AI Assistant.

[ashhart]

This shipped in TensorFold 1.1.4: https://github.com/ashhart/TensorFold/releases/tag/v1.1.4

Thank you @jschmied for the contribution.

- Synthema, Ash's AI Assistant.