Skip to content

cuda dense_rows: a per-column scaled epilogue for Dory's logits (Dory, #1233) - #1259

Closed
jschmied wants to merge 5 commits into
ashhart:mainfrom
jschmied:dory-dense-scaled
Closed

jschmied wants to merge 5 commits into
ashhart:mainfrom
jschmied:dory-dense-scaled

Conversation

@jschmied

Copy link
Copy Markdown
Contributor

What this changes

Step 3, last kernel PR of the Dory family (#1233). Dory's logits are the head's bf16 output times logit_scale / std per vocabulary column (dory.py: logits.float() * (logit_scale.float() / initializer_range)). tf_dense_rows_scaled_{1,4,16,32} are the same tiles as tf_dense_rows_* with that epilogue. DenseRows.rowsScaled launches them. tf_dense_rows_* keep their code, and their records from nvcc 13.0 are unchanged.

This PR closes nothing. The issue is #1233.

Stacking. Stacked on #1258 (act_rows), on #1257 (qk_rows), on #1256 (norms, carrying the #1245 contract include fix as its base commit). This branch is one commit on top of #1258.

What the kernels compute

out[n] = fp32(bf16(x . w[n])) * scale[n], in one fp32 product. The scale vector is formed on the host in fp32: ls = logit_scale * (1 / fl32(std)), as torch's CUDA division by a CPU scalar multiplies by the fp32 reciprocal (llama.cpp computes the same 1.0f / init_std).

Rounding points: the bf16 dot product, then one fp32 product.

Kernel contract

Registered as replacements: on the same inputs every output is bit-equal to its unscaled tile's fp32 outputs, rounded to bf16 and scaled on the host, at every served shape and edge (4,103 columns, 2,560 inputs, the row counts the wrapper sends each tile).

  • Gates 1, 2 and 4 pass against float64 of the same steps, with a NaN and an infinite scale planted.
  • No planted control: the exactness check is the bit-equality above, so a wrong scale or a missed rounding changes bits.
  • Limit records, sm_80/86/89/120/121 (regs; 0 static shared; launch bound 256):
    • _1: 30/40/40/38/38
    • _4: 38/38/38/45/45
    • _16: 79/79/79/98/98
    • _32: 127/127/127/180/180
  • Limit problems are unchanged from main.

Receipt

  • TensorFold 0346bdc40; 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

…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.
…2 norm and q's scale

tf_qk_rows_l2 takes dory.py's order: rotate-half RoPE on the projection's bf16 rows
(rounded to bf16), then each head's L2 norm with eps on the sum, then q times an fp32 scale
vector ((sqk f)^2, formed on the host as torch forms it) after the normalized row's
rounding; k is its normalized row, v as is, and k and v go into the caches as tf_qk_rows
writes them. Shape.l2 and Shape.q_scale select it; every other call keeps its kernel.

It is registered with the kernel contract: every output against 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, a third of the rows past position 199,800. Two controls
each drop one rounding point (the rotated head, q's normalized head) and go red at gate 3.
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.
…then L2-normed

Two activation rows for Dory, a block a row, 16-byte loads, fp32 scale vectors formed on
the host as torch forms them:

- tf_gate_l2_rows: out = bf16(l2(bf16(o * sigmoid(g * sg)))) over all heads of a row
  (dory.py's output gate, then justnorm over num_heads * head_dim);
- tf_swiglu_l2_rows: out = bf16(l2(bf16(silu(bf16(g * sug)) * bf16(u * suu)))) (its MLP).

sigmoid and silu are torch's fp32 forms, 1 / (1 + expf(-x)) and x / (1 + expf(-x)) with
IEEE division. ActRows.gateL2 and ActRows.swigluL2 launch them; no existing launch changes.

Both are registered with the kernel contract: every output against float64's range at its
rounding points (expf's ulps, fp32's exp range, the inverse's spread) at Dory's widths,
4,096 and 7,168; four controls each drop one rounding point (the gated row, the scaled gate,
the scaled up, the product) and go red at gate 3. Gate 4: the gate's rows are halves of
integers (sigmoid(0) = 1/2), exact in fp32; SwiGLU's squares are no integers, so its fp32
sums are checked against the shapes gate's bound, and an fp16 sum overflows either way.
Dory's logits are the head's bf16 output times logit_scale / std per vocabulary column
(dory.py: logits.float() * (logit_scale.float() / initializer_range)). tf_dense_rows_scaled_*
are the same tiles with that epilogue, out = fp32(bf16(x . w[n])) * scale[n] in one fp32
product; DenseRows.rowsScaled launches them. tf_dense_rows_* keep their code: their records
from nvcc 13.0 are unchanged.

Registered with the kernel contract as replacements: on the same inputs every output is
bit-equal to its unscaled tile's fp32 outputs, rounded to bf16 and scaled on the host, at
every served shape and edge (4,103 columns, 2,560 inputs, the row counts the wrapper sends
each tile); gates 1, 2 and 4 against float64 of the same steps, with a NaN and an infinite
scale planted.
@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.

@ashhart ashhart closed this Oct 25, 2026
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