[jschmied]

What this changes

Stacked on #1123 (window_attention at head dim 256); this is its second commit. It includes core/bf16_round.zig with #1056's sign fix through that PR.

Step 6 of #1048, the global attention. As agreed there, it is not a new operator: window_attention runs with no window over a linear cache (win = 0, ring the cache length) at head dim 512, with key-tile splits and a merge for long caches. Nothing calls it yet.

• Head dim 512 in window_rows.cu, with any number of query heads per key head up to 16 (Gemma 4: 8 for the 31B, 16 for the 12B). Linear caches are served to 262,144 keys and beyond: positions are int32, and a cache of more than 2^31 rows is refused.
• Decode keeps the 256-key chunks at absolute positions and the merge in chunk order from cuda window_attention: head dim 256, any GQA, a row's last key and fp32 probabilities #1123. At head dim 512 the partials stage K and V in shared memory, two tiles of 32 keys loaded with cp.async. A block takes 4 query heads, and each of its 8 warps takes a key run: keys w, w + 8, ... of every tile. The scores have the same arithmetic as at 256: each lane's fma chain over its 16 dims, then a xor tree. The runs' o are added in run order. The arithmetic doesn't depend on the batch, the slot or the cache length, so neither do a row's bits; chunks past a row's keys are skipped exactly.
• The prompt tiles are cuda window_attention: head dim 256, any GQA, a row's last key and fp32 probabilities #1123's at 512: four 128-dim slices per 16-row band, with p32 and the per-row last key as at 256.
• The 256 kernels are unchanged: their SASS on sm_80, 86, 89, 120 and 121 equals cuda window_attention: head dim 256, any GQA, a row's last key and fp32 probabilities #1123's (20 of 20).

Why one operator holds: the current callers' kernels are not touched. window_attention.cu and its fatbin are byte-identical to main's, and Kolibri-1's timings don't move (receipt below). The new head dims are new instantiations behind the same Window and WindowAttention.decode and prompt. The current head-128 kernels were not folded into the new ones: that would change Kolibri-1's bytes (bf16 P and its own sum orders), which #1048 rules out.

Receipt

• Environment: TensorFold d1018b31 (1.1.1) plus cuda window_attention: head dim 256, any GQA, a row's last key and fp32 probabilities #1123 and 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, and value ranges follow mlx-community/gemma-4-12B-it-4bit's global layers: q at RMS 0.95 to 1.05, k at 0.06 to 0.07 (inputs drawn from 0.06 to 0.13), v at unit RMS, with peaked keys mixed in.

• Exactness, tf-cuda-test window-rows (cuda window_attention: head dim 256, any GQA, a row's last key and fp32 probabilities #1123's cases, plus):

  • 12B global decode over a linear cache of 262,144 keys (4 rows, one at the cache's end) and of 20,000 (17 rows); 31B global decode, two streams of 40,000 (17 rows);
  • 12B global prompt, 128 rows at 8,000 and 16 rows at 130,000 (P32); 31B global prompt, 300 rows at 0 (P32) and 96 rows at 3,000 (bf16 P).

  In all: 4,848,501 outputs bit-equal to float64's value at the rounding points and 14,566,539 within their range. Smaller decode calls' bits equal the widest call's, the 20,000-key cache declared 16x longer gives the same bits, and a prompt chunk split in two gives every row's bits.

  • The range widens with the context: the prompt folds one 32-key tile at a time, each by one fp32 fma. At 130,000 keys, 1.6 % of the outputs are decided.

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

  shape	1k	8k	32k	128k
  12B global decode, 1 row (16/1 x 512)	31.7	82.4	359.4	1,451.1
  12B global decode, 4 rows	54.5	339.1	1,309.7	5,291.9
  31B global decode, 1 row (32/4 x 512)	30.0	317.6	1,235.4	4,904.5
  12B sliding decode, 1 row (cuda window_attention: head dim 256, any GQA, a row's last key and fp32 probabilities #1123)	35.2	36.4	35.4	37.2
  31B sliding decode, 1 row (cuda window_attention: head dim 256, any GQA, a row's last key and fp32 probabilities #1123)	49.5	36.1	37.0	44.2

  At 128k, a 12B row reads 256 MiB (185 GB/s), 4 rows read 1 GiB (203 GB/s), and a 31B row 1 GiB (219 GB/s).

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

  shape	1k	8k	32k	128k
  12B global prompt, P32	586.8	5,586.2	23,384.4	95,513.5
  12B global prompt, bf16 P	545.1	5,135.7	21,341.6	87,265.6
  31B global prompt, P32	1,102.1	10,474.3	42,759.4	174,827.9

  The prompt tiles run at about 23 TFLOPS of useful work (cuda window_attention: head dim 256, any GQA, a row's last key and fp32 probabilities #1123's kernel at 512); a deeper key pipeline is the next lever.

• Current callers' speed: unchanged. Kolibri-1, ABBA x2 against main (the bench added), GPU µs, medians of 4:

  Kolibri-1 call	main	this branch	change
  decode tf_window_decode_tc, 1 row	12.6	12.6	+0.4 %
  decode, 4 rows	12.6	12.6	0.0 %
  decode, 16 rows	22.9	22.9	0.0 %
  decode, 64 rows	308.7	308.0	-0.2 %
  sliding prompt wattn_kernel<128,8,4,8>, 2,048 rows	567.1	565.6	-0.3 %
  sliding prompt, 8,192 rows at 8,192	2,381.6	2,384.6	+0.1 %
  full prompt, 2,048 rows at 8,192	9,047.2	9,051.1	+0.0 %
  full prompt, 2,048 rows at 30,720	30,281.5	30,284.9	+0.0 %

• Kernel contract, tf-cuda-test kernel-contract (built -Dnvcc, ReleaseFast) on the GB10:
  SUMMARY 91 kernels: 91 pass gates 1-4, 0 fail, 0 known findings; 14 controls: 14 red as planted, 0 missed; 1533 limit problems
  (cuda window_attention: head dim 256, any GQA, a row's last key and fp32 probabilities #1123: 88 and 13; 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: 12B global decode (16/1 heads of 512 over a linear cache of 32,768: 4 rows to its end, 16 rows to 3,000, one row at 300, and the 256-key chunks' edges); 12B global prompt with P32 (256 rows at 0, 128 at 4,096, 32 at 1,000); 31B global prompt with bf16 P (128 rows at 0, 64 at 3,000).
  • Gates 1 to 4 as in cuda window_attention: head dim 256, any GQA, a row's last key and fp32 probabilities #1123. Gate 4 without a window caps its keys at 1,000, so the mirrored integers' partial sums stay exact in fp32. bf16 sums miss by 43.7x (decode), 4.0x (P32) and 0.5x (bf16 P), and the spec says why.
  • 1 new control: the head dim 512 merge over up to 128 chunks without exp(m_c - M), red at gate 3. cuda window_attention: head dim 256, any GQA, a row's last key and fp32 probabilities #1123's six still go red.
  • Gate 5: 4 more launches recorded (tf_wrows_part_512 at 64 KiB dynamic shared memory, opt-in; tf_wrows_merge_512; tf_wrows_prompt_512 and _p32 at 80 KiB). All 8 window_rows records are current on sm_80, 86, 89, 120 and 121. The run over every image reports 877 problems, as main's does.

• Default path unchanged: every fatbin but window_rows and contract_controls is byte-identical to main's on all five SMs. glm_quant differs in bytes, not in SASS, exactly as in cuda window_attention: head dim 256, any GQA, a row's last key and fp32 probabilities #1123's build. window_attention's SASS: 25 of 25. The existing controls' SASS: 35 of 35. tf-cuda-test window passes.

• 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). 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. 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

The same as #1123's: no change to window_attention.cu or window_split.cuh. A test merge onto the #746/#798 scratch conflicts only in the registration lists (registry.zig's record line; build/cuda.zig and contract_tests.zig already conflict between that scratch and 1.1.1).

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.)
• 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 keeping the 512 work behind the same Window API, for leaving the 256 kernels' SASS unchanged, and for the extra 512 merge-rescale control. We are carrying both commits onto current main ourselves, kept as yours, after #1123 (whose fixes reach this branch too), and a separate commit of ours will add an odd query-to-key-head ratio to the 512 contract entries (zig/tests/cuda/contract_window_rows.zig:55 to :57) so the partial last head group in zig/kernels/cuda/window_rows512.cuh:30 runs, put the 17-row count back in the 256 decode cases (zig/tests/cuda/window_rows_tests.zig:70), and run #1123's neighbour-row test at both widths. One choice is yours: wide() accepts a windowed ring and image rows at 512 (zig/src/cuda/window_attention.zig:17, :86 and :109) but no case covers them, so would you rather we cover both at 512 or refuse them there for now? Your answer on #1123's float64 re-run covers this branch as well, since the 512 prompt kernels go through the same fallback (zig/kernels/cuda/window_rows.cu:62 and :71). Then it goes to the GB10 kernel contract and the Kolibri-1 served-equality check, and lands after #1123; you need not rebase, and if you are in the middle of a change, please tell us so we do not collide.

- Synthema, Ash's AI Assistant.

[jschmied]

Please refuse both at 512 for now. Gemma 4 is the only user of the 512 path, and its global layers are causal over the whole context with no window: the ring is the context length, and image rows see each other only on the sliding layers (head dim 256). Our 31B image branch (#1186) uses the same split, so neither case is reachable today. A refusal by name is better than an untested path; if a model needs either later, it can come with its own cases. Not in the middle of a change here.

[ashhart]

Thanks @jschmied. Head dim 512 over a linear cache for global attention is on main (a4c16a5) with the maintainer fix-up (odd query-to-key-head ratios, the 17-row decode count). The kernel contract passes, and Qwen3.8-27B tokens and prefill on a GB10 match main. 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.