[jschmied]

What this changes

Step 3, second kernel PR of the Dory family (#1233). tf_qk_rows_l2 takes dory.py's order for the attention inputs. Shape.l2 and Shape.q_scale select it; every other call keeps its kernel and its bits.

This PR closes nothing. The issue is #1233.

Stacking. Stacked on #1256 (norms), which carries the #1245 contract include fix as its base commit. This branch is one commit on top of #1256. Retarget to main once #1256 lands.

What the kernel computes

For each head of a row (head dim 256, 16 query heads over 4 key heads):

• q, k = bf16(rope(proj)): rotate-half RoPE on the projection's bf16 rows, fp32, rounded once. The angle is fl32(pos * inv) with the fp32 inverse 1 / theta^(2j/d) and cosf/sinf (2 ulp). vLLM forms the same fp32 angle for its cos/sin cache and then rounds cos/sin to bf16; this kernel keeps them fp32.
• qn = bf16(l2_head(q)), then q = bf16(qs * qn), with qs = (sqk * (sqrt(d)/std))^2 an fp32 vector formed on the host as torch forms it.
• k = bf16(l2_head(k)).
• v is stored as is. k and v go into the caches as tf_qk_rows writes them.

L2 is x * (1 / sqrt(sum(x^2) + 1e-6)), eps on the sum.

Rounding points: the rotated head, the normalized head, q after the scale.

Kernel contract

Registered with the kernel contract. Every output is compared with float64's range at its rounding points (cosf's ulps on the fp32 angle pos * inv, the inverse's fp32 spread), at head dim 256, 16 query heads over 4, with a third of the rows past position 199,800.

• Gates 1-4 pass.
• Two planted controls go red at gate 3, each dropping one rounding point: the rotated head (tf_ctl_qk_l2_drop_rope) and q's normalized head (tf_ctl_qk_l2_drop_norm).
• Gate 4's integer heads (no rotation) overflow an fp16 sum to a zero output.
• core/l2_ref.zig's intervals now keep infinite ends and widen by fp32's least subnormal step, so a planted infinity leaves only the outputs float64 makes non-finite.
• Limit record (tf_qk_rows_l2 at dim 256, 256 threads): 30 regs, 64 B static shared, launch bound 512 on sm_80/86/89/120/121 alike.
• Limit problems are unchanged from main.

Receipt

• TensorFold 8a2d188f9; NVIDIA GB10 (sm_121), CUDA 13.0 nvcc, Zig 0.17.0.
• zig build test -Dnvcc=...: exit 0 at this branch head (zig build test -Dnvcc=/usr/local/cuda/bin/nvcc, the full step).
• lean_check: 0 new against main.
• Full kernel contract on the top of the stack (PR6): on main plus the include fix, SUMMARY 280 kernels: 278 pass gates 1-4, 2 fail, 2 known findings; 53 controls: 52 red as planted, 1 missed; 1566 limit problems. On this stack's top: SUMMARY 290 kernels: 288 pass gates 1-4, 2 fail, 2 known; 62 controls: 61 red, 1 missed; 1566 limit problems. That is 10 Dory kernels and 9 Dory controls added, all passing or red; the 2 failures, the 1 missed control and the 1,566 limit problems are the same on main (qwen35 decode attention, tf_q35_gdn_slots records of 0). Two things on main keep the unmodified run from completing: w4a8_tests.zig does not compile against w4a8.zig's prompt signature, and the w4a8 kernels lose the CUDA context on their planted non-finite inputs. Both runs above skipped the w4a8 kernels and patched that test call locally, identically on both trees; neither patch is in this PR. All Dory launches pass the live limits check.

Choices ashhart can change

• Rounding convention (Nemotron-Labs Dory 1.0 4B on native CUDA: a looped nGPT family (dory) #1233 question 3): bf16 where dory.py casts, fp32 between. The exception is RoPE, which is computed in fp32 and rounded once, because vLLM's cos/sin cache is bf16 and its rotary internals differ between eager, the custom op and inductor.
• Loop count and KV policy (questions 1, 2) do not touch this PR.

[ashhart]

Landed on main in 6528fbe..81b476d: the whole Dory stack (#1256-#1261), with your records under your authorship plus five maintainer children. The children refuse host-only formats and refund failed upload ownership, pin release launch arguments, keep the unmeasured Dory launches pending in the contract, and own queued host inputs through stream completion. Every existing family file, contract entry and recorded launch is byte-identical. Thank you, Jürgen.

On a GB10 with nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-NVFP4 @ 41f02335, served at 2 loops with CUDA graphs:

• Prefill: 5,143, 5,354, 4,763 and 4,220 tok/s at 512, 2k, 8k and 16k.
• Decode: 44.2, 86.2, 164.1 and 298.6 tok/s aggregate at 1, 2, 4 and 8 streams. That matches your 43-44 tok/s a stream at L=2.
• Exactness and quality: repeats exact; chat with thinking off and on qualified (resume equals fresh); the 16k needle found.

- Synthema, Ash's AI Assistant.