cuda/norms: one tree a width, IEEE 1/sqrt; Nemotron's add_rmsnorm on the shared kernel

[jschmied]

What this changes

Four commits, following the notes on #736:

1. cuda/norms:

   • a bf16 row takes the 16-byte tree by its width alone; a misaligned buffer is refused (error.Misaligned)
     instead of running on the 1,024-thread tree, so a row's bits no longer depend on where its buffer sits;
   • 1 / sqrt by IEEE division and square root (__fdiv_rn, __fsqrt_rn) instead of rsqrtf, in the generic,
     16-byte and gated forms: correctly rounded on every architecture, as Nemotron's own norms already were;
   • if constexpr for the wide rows, so nvcc no longer warns "loop is not reachable".

2. Nemotron CUDA: add_rmsnorm runs on the shared add_rms_h; tf_nemo_add_rmsnorm is removed. The glue test's
   host reference now follows add_rms_h's order (8 values a chunk, chunk c on thread c % 256, the warp then the
   8-warp butterfly, the group sums' 8-lane tree), so it still compares every bit. norm-rows gains Nemotron's
   width (2688).

3. add_rms_h issues every load before any is used, and each warp folds the warps' sums itself (the same tree,
   so the same total, with no broadcast and one block sync where one norm runs): the same bits, one 5120-wide row in
   a CUDA graph 2.25 -> 1.84 µs a call; in the 27B's decode on a GB10 3.63 -> 2.65 µs (Triton's _add_rmsnorm 2.92).

4. Decode-sized gated norm calls (16 rows or fewer) take 32-thread blocks (launch.rowThreads, Triton's size for
   the 27B's _gated_norm), the same bits: on a GB10, one row, 1.95 -> 1.78 µs a call (Triton 1.32 µs). The rest
   of that gap is the norm's exactly rounded 1/sqrt, kept for the same bits on every GPU (cuda: SwiGLU, output gates, Gated DeltaNet inputs, QK norm and RoPE for every family, checked against float64 #797 has the numbers).

The shared kernel sums in another order than Nemotron's, at the same precision: Nemotron's replies change at
near-ties, as #736's note expected.

Receipt

• Environment: TensorFold c36b536d (1.0.9) plus these commits; one DGX Spark (GB10, sm_121, CUDA 13); Nemotron 3.5
  Lightning MLX 4-bit.
• Kernels: tf-cuda-test glue 1,242 of 1,242 comparisons byte-equal to the host references; norms and norm-rows
  within float64 tolerance (worst 0.996), every row count's bits equal, a misaligned bf16 row refused. Mutation:
  the cross-warp butterfly run 1, 2, 4 instead of 4, 2, 1 fails glue (byte-exact) and passes norm-rows
  (within tolerance).
• Kernel by itself, 4,096 random rows at 2688: 105 of 11,010,048 outputs differ from tf_nemo_add_rmsnorm (in 68
  rows), h identical; of the differing values Nemotron's is nearer float64 in 46, the shared one in 59, total error
  equal. Time per call: 1 and 4 rows 4.1 us both, 64 rows 6.08 -> 4.10 us, 4,096 rows 404 -> 391 us.
• Served, --parallel 4, 84 cells (2 prompts, greedy and 5 sampling settings x 4 seeds, drafted and plain, 96
  tokens, reasoning and content hashed): drafted == plain 42/42 on both builds; 30 of 84 replies equal to main's,
  the rest diverge after a near-tie.
• Teacher-forced, 4 documents x 512 tokens (docs, code, German, chat), against main: dNLL +0.0032 (per document
  -0.004 to +0.008), argmax agreement 98.04 %, KL 0.0038. The noise floor, main with only its own cross-warp sum
  reversed: dNLL +0.0015 (-0.010 to +0.017), argmax 97.80 %, KL 0.0046. The change sits inside it.
• Speed, A B B A: decode 512 tokens 142.8/143.6 and 143.4/144.3 tok/s (main) vs 146.8/147.9 and 147.4/148.3
  (this), +2.7 %; prompt 9.6k tokens 1.322/1.311 s vs 1.303/1.298 s, 35k tokens 5.469/5.367 s vs 5.397/5.365 s.
• Tests run: zig build test and zig build test-cuda-host (27/27 steps, 345/345 tests). zig fmt --check clean.
  lean_check.py: 538 findings, as on main.
• RTX 5090 (sm_120), run by @BobClawblaw (comment below), b3ad1b9d against 1e76a70e: all 16 tf-cuda-test suites pass, glue 1,242/1,242 byte-equal; the fold-order mutation fails glue as on the GB10. Served --parallel 4, 54 cells: drafted equals plain 27/27, concurrent equals solo 54/54, this PR equals its base 54/54. Decode +0.4 to 0.6 %; tf_add_rms_rows_bf16h 2.03 → 1.48 µs a call.
• Not run, and why: Metal (unchanged); sm_86.

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.)
• No precision traded for speed. If the bits change, the description says which sums change and why (the norm's sum order; same precision).
• 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]

Thank you, @jschmied. Nemotron on CUDA adds each layer's residual and normalizes it in the shared kernels of cuda/norms.cu, and the family's own copy is gone. A bf16 row takes the 16-byte path by its width alone, a misaligned buffer is refused, and the reciprocal square root is an IEEE division and square root.

The sums run in another order of the same precision, so a CUDA reply can part from 1.0.9's at a near tie. Teacher-forced loss over the seven public texts is unchanged within noise, on MLX 4-bit and on NVFP4. Drafted replies equal "draft": false, 17 of 17 on every arm. Decode and prefill stay level.

Your commits are in 1.0.10 under your name (367ec59ff, 9b00d3a8e). Each patch matches your head (140ae71e0, 3e47fe5ca). The CHANGELOG line is 709f6aef5. Your head 8fe038e37 is merged at 7a21b889a with no change to the tree. The 32-thread launch is in 1.0.10 (dc9d8bb6c, the same patch as 8fe038e37).

Your branch has moved on since: its head is now b3ad1b9dd, where add_rms_h loads its inputs before the store and every warp folds the eight warp sums. Those commits are not in 1.0.10. They change the order of the residual norm's sums again, so they need the same bit, quality and speed round as the first part, and we'll write here when it has run. This pull request stays open for them.

[BobClawblaw]

@jschmied here's the sm_120 run your receipt lists as not run. RTX 5090 (170 SMs, 32 GB), driver CUDA 13.4, nvcc 13.3, Zig 0.17.0, built with -Dsm=120. I compared b3ad1b9d against its base on main, 1e76a70e. Nemotron 3.5 Lightning MLX 4-bit.

Kernels. All 16 tf-cuda-test suites pass on both builds: info, smoke, graph, launch-ex, ptx, symbols, cublaslt, norms, norm-rows, sample, glue, pocket, quant-ref, modelopt, exlr8 and exlr8-dense. glue is 1,242 of 1,242 byte-equal. norms and norm-rows are within tolerance (worst 0.996), every row count's bits are equal, and misaligned rows are refused. I also ran a mutation: inv_rms_v's new per-warp fold run 1, 2, 4 instead of 4, 2, 1. It fails glue (2 of 1,242 differ) and passes norms and norm-rows, as on your GB10.

Served, --parallel 4. I used 54 cells: 3 prompts (short prose, code, a 6k-character code summary) × greedy and 8 seeded sampling settings (t 0.7 / top_p 0.95 and t 1.0 / top_k 40, seeds 1–4), each drafted and with "draft": false, 128 tokens, reasoning and content hashed.

• Drafted equals plain in 27/27 on both builds.
• Four at a time equals solo in 54/54 on both builds.
• This PR equals its base in 54/54, solo and concurrent, which fits "the same bits".

Speed. Decode medians of 3 reps, 512 tokens, in the order base, PR, PR, base:

prompt	base	this PR
gpu-chat, t 1.0	587.9 / 588.5	591.4 / 591.6
fibonacci, t 0	1201.4 / 1201.6	1206.5 / 1206.0
gpu-chat, t 0	617.5 / 618.0	620.6 / 620.6

That's about +0.4 to 0.6 % end to end. Time to first token is the same (11–13 ms).

The kernel in decode, from nsys with --cuda-graph-trace=node, 256-token replies: tf_add_rms_rows_bf16h averages 2.03 µs a call on the base and 1.48 µs on this PR (median 2.08 → 1.50 µs, about 36.8k calls each). The other norms are unchanged: tf_nemo_group_rmsnorm 1.67 vs 1.66 µs and tf_nemo_add_moe_norm 6.60 vs 6.57 µs. Node-level tracing adds its own overhead, so these per-call times aren't comparable with your in-graph ones; only the direction and the rough size are.

Not checked:

• the nvcc "loop is not reachable" warning, because neither build printed it here, probably because the nvcc outputs were cached;
• prompt-length timings;
• the 27B, which isn't on main yet.

[jschmied]

@BobClawblaw, thank you, that's the run the receipt was missing: all 16 suites, your fold-order mutation caught by glue as on the GB10, the 54 served cells equal across drafted, concurrent and base, and the kernel's direction (2.03 → 1.48 µs a call) on sm_120. I've added it to the receipt.