[jschmied]

tf_q35_attn_partial (and _partial8 for --kv-cache fp8) launches a grid of (chunks, key heads, rows). Blocks are dispatched with x fastest, so the GPU walks every chunk of row 0, then every chunk of row 1, and so on. Each row's blocks read the whole of its keys and values again. At head dim 256 a 27B attention layer holds 33.6 MB of K/V at 8k tokens and 134 MB at 32k, more than GB10's 24 MB L2, so every row of a DFlash2 verify window (12-16 rows) or a 128-row prompt chunk re-reads K/V from DRAM.

This PR swaps the axes to (rows, key heads, chunks): all rows of a chunk run together and read its keys from L2. Each block still computes exactly what it computed before (same row, head, chunk, same loop), only the order the blocks start in changes, so the bits are unchanged.

The diff: blockIdx.x and blockIdx.z trade places in both kernels, and cuda_forward.zig launches { total, kv_heads, max_chunks }. Rows move to grid.x, which also lifts the 65,535 cap grid.z had on them.

Synthetic kernel, L2 flushed before each launch (32k keys): 16 rows 7.62 -> 4.67 ms, 128 rows 72.1 -> 29.3 ms. ncu L2 hit rate 22.6 % -> ~100 % at 16 rows.

Served Qwen3.8-27B MLX 4-bit + DFlash2, GB10. --parallel 4 --context 49152 --prompt-cache-gib 8, bf16 cache. A fresh server per arm, main and branch alternating, two rounds.

	main	this PR
TTFT, 32,706-token prompt	143.6 / 140.3 s	101.1 / 101.1 s (-29 %)
TTFT, 8,222-token prompt	13.57 / 13.71 s	13.65 / 13.64 s
DFlash2 c=1 after the 32k prompt	37.4 / 37.1 ms/token	31.3 / 31.4 (-16 %)
the same turn resumed from the prompt cache: TTFT	0.634 / 0.659 s	0.413 / 0.447 s
... and its decode	43.8 / 44.2 ms/token	36.0 / 37.1 (-17 %)
DFlash2 c=4 after the 32k prompt (per stream)	123.8 / 132.0 ms/token	78.8 / 78.6 (-38 %)
DFlash2 c=1 / c=4 after the 8k prompt	29.4-30.4 / 55.5-55.6 ms/token	28.3-30.0 / 54.2-54.5
plain decode c=1, short / 32k (one row: nothing to share)	73.2 / 82.9 ms/token	73.4-74.1 / 83.0-84.9
DFlash2 c=4 short prompts, 8 greedy + 4 seeded runs a server	142.4 / 142.3 tok/s	144.6 / 143.4 tok/s

The gain grows with context and with rows a round: at 8k the K/V is close to L2-sized and the change is 0-3 %.

Checks

• Token hashes: 18 of 18 bench lines equal to main's in every arm. Those lines cover greedy and seeded, c=1 and c=4, 8k and 32k, drafted and plain, fresh and resumed. --kv-cache fp8 arms: 8 of 8 equal.
• In every arm, drafted == plain, the c=4 stream == the solo request, and resumed == fresh.
• Kernel contract (new contract_q35_attn.zig): partial and merge as one call, against float64 head by head, bf16 and fp8 caches.
  • Gates 1-4 pass. Gate 3 is bit-equal to the old grid at 12 shapes: 16-row windows over 2,100 and 9,000 keys (chunks widen past 8,192), four streams' windows, a 128-row prompt chain, and the row and key edges.
  • The baseline runs each row in a launch of its own, where both grids start the same blocks.
  • A planted control that drops each row's last chunk goes red at gate 3.
  • Run total: SUMMARY 142 kernels: 142 pass gates 1-4, 0 fail, 2 known findings; 33 controls: 33 red as planted, 0 missed.
• Gate 5: no record changes. With nvcc 13.0.88, both kernels compile to the same registers before and after on sm_80/86/89/120/121 (48/48/48/40/40 and 40/40/56/34/34, no shared or stack). Threads per block are unchanged.
• Full zig build test -Dnvcc=/usr/local/cuda/bin/nvcc: exit 0 (295/295 steps; 1,816 of 1,827 tests passed, 11 skipped). lean_check: 0 new against the merge base. zig fmt --check is clean.

Gotchas

• Branch qwen27-attn-rows-first @ 74020261, one commit on main @ f4d6c4f23.
• qwen CUDA: decode attention on tensor cores (fp8 and bf16 caches), decode after long prompts 1.5-2.3x faster #1168 (stacked on qwen CUDA: long-prompt prefill 5.4x faster (4096-row chunks, tensor-core attention, the delta-rule chain) #1080 and qwen CUDA: Qwen3.8-27B decode 9-15% faster with DFlash2 (code +25%), replies unchanged #1079) replaces both kernels with tensor-core ones. Its attn_rows kernel groups a long stream's rows so each K tile is read once, which covers this change and more, but it changes replies (~1e-4 RMS). This PR is the bit-exact step for main until qwen CUDA: decode attention on tensor cores (fp8 and bf16 caches), decode after long prompts 1.5-2.3x faster #1168 lands; if qwen CUDA: decode attention on tensor cores (fp8 and bf16 caches), decode after long prompts 1.5-2.3x faster #1168 lands first, drop this one. The two conflict textually in qwen35.cu, qwen35_kv8.cu and the cuda_forward.zig launch line.
• qwen CUDA: long-prompt prefill 5.4x faster (4096-row chunks, tensor-core attention, the delta-rule chain) #1080 moves single-stream prompt chunks to tensor-core attention, so after it this PR no longer touches TTFT. The same diff on a qwen CUDA: long-prompt prefill 5.4x faster (4096-row chunks, tensor-core attention, the delta-rule chain) #1080-based tree (local tree 469166bfa, qwen CUDA: long-prompt prefill 5.4x faster (4096-row chunks, tensor-core attention, the delta-rule chain) #1080's prompt path in place) left 32k TTFT at 27.7-28.0 s but kept verify/decode gains, hashes equal:
  • DFlash2 c=1 at 32k: 35.7-37.9 -> 27.2-29.0 ms/token (-22..-25 %);
  • c=1 at 8k: -8 %;
  • c=4 at 8k: 54.6 -> 43.9 ms/token (-20 %).
• New known finding, unrelated to this change: a NaN key in one of a row's several chunks leaves the row finite. tf_q35_attn_merge skips a chunk unless PL > 0, and NaN fails that test. The contract registers it as known at gate 1; with that one plant removed, gate 1 passes for both caches.
• This box's live gate-5 run reports limit problems for the qwen35* launches whose 13.0 records are still pending (usage none). That predates this change, and none of it is new here.
• The fp8 cache path gets the same swap; I measured its hashes, but only its 8k speed (within noise).

[ashhart]

Thanks @jschmied. This is on main now (a56d0ce), your commit with your name and date. We measured it on a DGX Spark (GB10) against main, served Qwen3.8-27B MLX 4-bit, one fresh server per arm:

• Cold prefill: 16k +15%, 32k +44%, 64k +90% (2k and 8k within noise).
• The 1/2/4/8-stream sweep is unchanged, within 0.1%.
• Tokens are identical to main in every cell, every stream equals its run alone and its "draft": false run, and resumed chats equal fresh ones.

Closing as landed; it ships in 1.1.4. As you noted, #1168's tensor-core attention will replace these kernels, and the two are being rebased together.

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