Skip to content

cuda norms: nGPT rows, L2 norms with eps on the sum and the normalized lerp residual (Dory, #1233) - #1256

Closed
jschmied wants to merge 2 commits into
ashhart:mainfrom
jschmied:dory-l2-norms
Closed

jschmied wants to merge 2 commits into
ashhart:mainfrom
jschmied:dory-l2-norms

Conversation

@jschmied

@jschmied jschmied commented Oct 24, 2026 •

Copy link
Copy Markdown
Contributor

What this changes

Step 3, first kernel PR of the Dory family (#1233). Dory's residual stream lives on the unit sphere, so its rows need two kernels norms.cu does not have. Both are a block a row (256 threads, 16-byte loads, rows to 8,192 in 64s), so a row's bits never depend on the row count. No existing kernel or launch changes.

This PR closes nothing. The issue is #1233, and the family stays open there.

Stacking. The base commit of this branch is the contract include fix from #1245 (cuda contract: the softcap control includes the act_round helpers it calls). Without it contract_controls.cu does not build on main. It is the same patch as #1245. Drop that one commit when either lands. After it comes one commit, the norms.

Stack: PR1 (this) <- #1257 qk_rows <- #1258 act_rows <- #1259 dense_rows <- #1260 family <- #1261 serving.

What the kernels compute

  • tf_l2_rows: y = bf16(s / sqrt(sum(s^2) + eps)), eps = 1e-6 on the sum (not on the mean), fp32 sum, IEEE sqrt and division. s = x, or s = bf16(x + r) (a loop's entry, justnorm(state + input_h); with r null it is a loop's exit).
  • tf_slerp_rows: dory.py's slerp_norm. a = bf16(l2(res)), b = bf16(l2(out)), z = a + lr * (b - a) in fp32 steps with lr = |alpha * (alpha_init / std)| a host-formed fp32 vector, h = bf16(l2(z)).

Norms.l2 and Norms.slerp launch them.

Rounding points (the choice in #1233 question 3, open): a tensor that dory.py holds as bf16 is rounded to bf16, and all arithmetic between those points is fp32. Here: bf16(x + r), the output of each l2, a, b and h. The lerp is not rounded between its steps.

Kernel contract

Both are registered with the kernel contract. The check compares every output with float64's value at the bf16 rounding points, as a range over the fp32 run's spread (core/l2_ref.zig).

  • Gates 1-4 pass.
  • Three planted controls go red at gate 3, each dropping one rounding point: the sum before the norm, the slerp's a, the slerp's b.
  • Gate 4's integer rows make an fp16 sum overflow to a zero output (85x and 51x the bound). A bf16 sum's scale cancels in the norm, so that margin is stated, not proven, as for the RMS tails.
  • Limit records, nvcc 13.0 images for sm_80/86/89/120/121 (regs per thread, static shared 32 B, launch bound 256):
    • tf_l2_rows: 48/48/48/64/64 regs
    • tf_slerp_rows: 113 regs on all five
    • sm_120 and sm_121 are the last two entries of each.
  • Limit problems are unchanged from main.

Receipt

  • TensorFold d2a0e1b22 on top of origin/main 6244c6206; 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, includes this PR's kernels): 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

The open questions of #1233 are not answered there yet. This stack follows the defaults below, and each is easy to change:

  • Rounding convention (question 3): bf16 where dory.py casts, fp32 between. The planted controls name the points.
  • KV policy (question 2): per_loop only; not relevant to this PR.
  • Loop count (question 1): --loops as a load option; not relevant to this PR.

…d lerp residual

Dory's residual stream lives on the unit sphere, so its rows need two kernels norms.cu did
not have, both a block a row (256 threads, 16-byte loads, rows to 8,192 in 64s), so a row's
bits never depend on the row count:

- tf_l2_rows: y = bf16(s / sqrt(sum(s^2) + eps)), s = x, or bf16(x + r) (a loop's entry,
  justnorm(state + input_h); with r null, a loop's exit).
- tf_slerp_rows: dory.py's slerp_norm, h = bf16(l2(a + |alpha scale| (b - a))) with
  a = bf16(l2(res)) and b = bf16(l2(out)), the lerp in fp32 steps.

Norms.l2 and Norms.slerp launch them; no existing kernel or launch changes.

Both are registered with the kernel contract. Their check compares every output with
float64's value at the bf16 rounding points as a range over the fp32 run's spread
(core/l2_ref.zig); three controls each drop one rounding point (the sum before the norm,
the slerp's a, its b) and go red at gate 3. Gate 4's integer rows make an fp16 sum overflow
to a zero output (85x and 51x the bound); a bf16 sum's scale cancels in the norm, so that
margin is stated, not proven, as for the RMS tails. Launch records from nvcc 13.0's
sm_80/86/89/120/121 images.
@ashhart

ashhart commented Oct 25, 2026

Copy link
Copy Markdown
Owner

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants