[jschmied]

What moves

Step 1 of #548. Nemotron's four norm kernels in nemotron_norms.cu move onto shared ops, and the file is deleted.

• tf_nemo_embed16: now embed_rows' bf16 row copy.
• tf_nemo_add_moe_norm: now norms.cu's tf_moe_add_rms_rows_bf16h, one launch. It is add_rms_h with the expert slots as its residual: each thread sums its 8 columns over the slots in slot order from 0 with fp32 fmas, rounds the sum to bf16, adds it to x, and then runs the shared norm. That is the arithmetic of moe_glue's combine followed by tf_add_rms_rows_bf16h, so the bytes are those of the two-launch form, without the second launch or the delta row. add_rms_h and its 16-byte helpers live in norms_vec.cuh, which norms.cu includes.
• tf_nemo_concat_norms and tf_nemo_group_rmsnorm: now two new norms.cu entry points, tf_rms_rows_bf16h_parts (rounded once) and tf_rms_rows_bf16h_parts_round (MLX rounding). Each runs one block per row part. Parts up to 512 wide (Mamba's groups) run on tf_rms_rows_bf16h_parts4 / _parts4_round instead: four parts a block, 64 threads each. Each thread's 8-term chain and the warp tree are the same, and warp 0 + warp 1 is the 8-warp tree's total when warps 2-7 hold zeros, so the bits are the one-block form's (y and the sums byte-equal over 224 calls, widths 64-512, 1-2,048 rows, both roundings, NaN, ±inf, tiny values). The one-block forms are thin wrappers that offset the pointers so norms_tail.cuh's tail<ONCE> / tail<MLX> (unchanged) lands each strided part.
• Step 1 no longer loads moe_glue: none of its launches remain here.

Bytes: this step changes Nemotron's bytes. The shared tail sums the squares in its own order, not in the old per-family order. Other families are untouched: every other serving fatbin is byte-identical to main, norms.fatbin's existing functions keep their SASS on sm_80/86/89/120/121, and contract_controls.fatbin gains one control.

Kernel contract

• The norms_tail harness covers Nemotron's norm kernels at its shapes:
  • add_y and add_h for tf_add_rms_rows_bf16h, 16 and 300 rows × 2,688. Its h is one rounded add, so it is checked element-wise with no sums.
  • moe_bf16_y, moe_bf16_h, moe_f32_y, moe_f32_h for tf_moe_add_rms_rows_bf16h, at the same shapes, with Nemotron's 8 slots: 6 routed weights (sigmoid scores normalized to 2.5) and 2 shared slots of weight 1, as bf16 slots (the prompt's down projection writes them) and as fp32 slots (decode's). h declares its sum: the slots' fp32 fma chain, depth 8, rounded to bf16 before the add; y's depth is that chain plus the squares'. The float64 check takes the combine as a range over the chain's gamma(8), rounded to bf16 at both ends, and checks h and y against it. Gate 4 gives h's kinds slot terms that are exact in fp32 in any order, with one term past fp16's range.
  • parts_once_y for 32 and 2 rows × 2,688; parts_round_y for 128 and 2,400 rows × 512 (Mamba's 8 groups). Each parts kind runs the one-block and the four-a-block form.
• A planted control, tf_ctl_tail_moe_keep_f32, runs the same add_rms_h body with the combine kept in fp32 before the add. It goes red at shapes, as planted (row 0 h[1] = 1.8671875 outside float64's [1.875, 1.875]).
• Records (nvcc 13.0, sm_80/86/89/120/121): tf_moe_add_rms_rows_bf16h 128 regs on all five, 32 B static, bound 256. tf_add_rms_rows_bf16h 80/80/80/98/98 regs, 32 B static. Both parts kernels 64 regs on all five, 32 B static, bound 256. _parts4 36/34/34/36/36 and _parts4_round 36/34/34/35/35 regs, 32 B static, bound 256. nvcc 13.2 compiles them the same.
• kernel-contract on the norm kinds: the 4 moe kinds and the 12 other norm kinds pass gates 1-4; the control is red as planted. kernel-contract-limits at nvcc 13.0 and 13.2: every record of this step current, and the failing lines are the same as on the base.

Nemotron's gates (Lightning MLX 4-bit, GB10)

gate	main	this step
drafted == plain, together 1/2/4/8 == solo, repeats, windows == pages, sampled repeat	equal	equal
widths, shared widths 8	63/63, 504/504	63/63, 504/504
prompt cuts, resume, paged 8k	pass	pass
glue vs host references	1470/1470	1470/1470
segments 1-4 (state and reply)	same, rounds differ	equal everywhere
greedy chat	haiku	haiku (new bytes)
16k needle	found	found

Distance to a float64 CPU reference: MLX weights dequantized exactly, every op in float64. It covers 3 passages × 320 scored rows after a 64-token prefix. The fused kernel gives the same bytes as the two-launch form these rows were measured on (kl dump sha c94ddc4a3f2f on both).

passage	rel-L2 mean main / step	argmax agree main / step	KL(ref‖tf) mean main / step
CONTRIBUTING.md	0.0366 / 0.0360	96.9% / 95.6%	0.0058 / 0.0052
README.md	0.0442 / 0.0462	95.3% / 93.8%	0.0073 / 0.0074
docs/kernel-contract.md	0.0408 / 0.0398	95.0% / 94.7%	0.0093 / 0.0086
all 960 rows	0.0406 / 0.0407	95.7% / 94.7%	0.0075 / 0.0071

Every argmax miss on either tree sits at a reference top-1 margin of at most 0.63 logits; nearly all are below 0.22. 23 rows miss only here and 13 miss only on main. McNemar gives p ≈ 0.1, so the difference is not significant at 960 rows. Mean rel-L2 is equal and mean KL is 5% lower. The teacher-forced loss over seven public texts (3,584 tokens) is +0.0008 ± 0.0046 against main.

Speed

Served A/B on a GB10, Lightning MLX 4-bit, main 2e33e297d against 8a77e3ec0, 12 rounds. Each arm and round gets a fresh server, in alternating order. Cold TTFT is the median of 3 a round; plain decode is a fresh --no-drafts server, 256 tokens, 3 reps. The table gives the median of the paired per-round ratios:

metric	main	this step	paired median	rounds this step is faster
cold TTFT 2k	265.1 ms	263.8 ms	-0.9%	8 of 12
cold TTFT 8k	1,025.8 ms	1,024.4 ms	-0.1%	7 of 12, 2 level
cold TTFT 32k	4,809 ms	4,806 ms	-0.1%	6 of 12, 1 level
plain decode, 1 stream, chat	108.5 tok/s	109.0 tok/s	+0.5%	12 of 12
plain decode, 1 stream, code	108.7 tok/s	109.2 tok/s	+0.4%	12 of 12
plain decode, 4 streams, chat	350.7 tok/s	351.1 tok/s	+0.2%	10 of 12, 2 level
plain decode, 4 streams, code	347.3 tok/s	347.8 tok/s	+0.1%	9 of 12, 1 level

Before the group-norm commit, cold TTFT was 0.1-0.3% slower than main; an nsys trace put that on the one-block parts norm (16.8 / 67.6 ms at 8k / 32k against tf_nemo_group_rmsnorm's 14.3 / 57.3). Drafted decode moves with acceptance, since the bytes differ from main's, so it is not a speed comparison here.

Receipt

• Head 8a77e3ec0 on origin/main 2e33e297d; GB10 (sm_121), CUDA 13.0, Zig 0.17.0.
• zig build test -Dnvcc=/usr/local/cuda/bin/nvcc: exit 0. lean_check: 0 new.

[ashhart]

Review of #1353 and #1354 (#548 steps 1 and 2) against main e9b69ed. Moving Nemotron's norms and router onto the shared glue is a good direction, and nothing here is lossy. But both steps change Nemotron's output bytes, and the float64 receipts get slightly worse rather than better:

• nemotron CUDA: RMSNorm rows run on the shared norms and combine (#548, step 1) #1353: argmax agreement 95.7% → 94.7%.
• nemotron CUDA: the router top-k is moe_glue's sigmoid top-k (#548, step 2) #1354: argmax misses 41 → 48, rel-L2 0.0406 → 0.0412, worse on all three passages.
  Prefill is also 1.1–1.8% slower with the same sign in every run; you attribute it to the extra combine launch, which could fuse into add_rms_h.

A refactor that changes bytes needs its error to be no worse than before, and we can't take a prefill slowdown. Two ways forward: keep Nemotron's current summation order inside the shared kernels (a variant that reproduces today's bytes), or show per-projection-class worst-column float64 error, prompt and decode separately, that is no worse. Also, dropping the +1e-20 turns 0 into NaN when every top-k sigmoid underflows; please keep a guard and add that case to the tests.

[jschmied]

One question before we rework #1353/#1354, on what "no worse" should mean.

The two steps change Nemotron's summation order, not its precision: the shared kernels sum the same f32 values in another order. A different order moves the last bits toward or away from the float64 value at random, so against a fixed main about half of any set of measurements comes out slightly worse by chance. Our receipts look like that noise rather than a trend:

• nemotron CUDA: the router top-k is moe_glue's sigmoid top-k (#548, step 2) #1354 vs main, 960 rows: rel-L2 0.04056 → 0.04119, paired +0.00063 ± 0.00096 (2 se). Argmax misses 41 → 48, McNemar p = 0.36.
• nemotron CUDA: RMSNorm rows run on the shared norms and combine (#548, step 1) #1353, 320 rows: argmax 96.88 % → 95.62 %. All 14 misses are at reference margins ≤ 0.22.
• The next step (Zig CUDA ports: shared layers below lanes.Backend, and one place to coordinate them #548 step 4) vs nemotron CUDA: the router top-k is moe_glue's sigmoid top-k (#548, step 2) #1354, same rows: rel-L2 0.04119 → 0.04021 (−0.00098 ± 0.00100), misses 48 → 45 (p = 0.76). The same kind of change landed on the other side.

A per-projection-class, worst-column test asked as "no worse in every class" is many comparisons of an extreme value. Two equally accurate orders would rarely pass it, so the only way to pass for sure is to reproduce main's bytes, i.e. keep Nemotron's own order inside the shared kernels.

Would "not worse beyond statistical noise" work for you instead? Concretely: paired differences per class (prompt and decode separately) within 2 se of zero, with no class significantly worse after correcting for the number of classes. Then the shared kernels keep one order for every family. If you need byte equality to main, we'll keep Nemotron's order in the shared kernels instead.

Either way, the prefill cost goes (the combine fuses back into the norm, as on main), and the +1e-20 guard comes back with an all-underflow test.

[jschmied]

One addition: this is the criterion #796 already used when Nemotron's add_rmsnorm moved onto the shared add_rms_h in 1.0.12. That was the same family and the same kind of change, with your note on landing it: "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." #1353/#1354 are the next two steps of that migration.

To make the noise level concrete rather than argued, I'll add two receipts on these PRs:

• the same teacher-forced loss over the seven public texts, main vs nemotron CUDA: RMSNorm rows run on the shared norms and combine (#548, step 1) #1353 vs nemotron CUDA: the router top-k is moe_glue's sigmoid top-k (#548, step 2) #1354;
• a placebo: main against main with one arbitrary reorder (the 8-warp butterfly reversed). It shows how far the per-class float64 error and argmax misses move from reordering alone.

[jschmied]

The two receipts. Lightning MLX 4-bit, GB10. Placebo = main e9b69ed with inv_rms_v's 8-warp fold run 1, 2, 4 instead of 4, 2, 1 (the norm Nemotron's dense layers use; same depth, local only). Seven public texts here = README, norms.zig, CPython textwrap.py, moe_glue.cu, GPL-3, xz(1) German man page, the 1.0.12 CHANGELOG: 512 tokens each, after a 16-token prompt. Your set isn't in the repo, so I'll rerun on yours if you send it.

teacher-forced NLL (nats/token), vs main, 2 se	#1353	#1354	placebo
all 3,584 tokens (main 1.8684)	+0.0008 ± 0.0046	+0.0024 ± 0.0047	+0.0020 ± 0.0047
per text, range	−0.0046 … +0.0057	−0.0031 … +0.0065	−0.0052 … +0.0077

vs float64, 960 rows, vs main (rel-L2 0.04056, 41 misses)	#1353	#1354	placebo
paired rel-L2, 2 se	+0.00012 ± 0.00095	+0.00063 ± 0.00096	−0.00061 ± 0.00091
argmax misses (McNemar p)	51 (0.13)	48 (0.36)	44 (0.77)

Per projection class, placebo vs main, 1,152 rows: paired rel-L2 moves −0.0007 … −0.0028 in prompt and −0.0008 … −0.0016 in decode. That is past 2 se in most classes, all on the better side, because rows within one run are correlated and the per-row se is too narrow. The worst column flips from 4.42 to 0.45 (in units of its layer's RMS) on prompt passage a. On main, one row (162, layer 18, column 2206) reads 52.5 where float64 has −1.18. The placebo doesn't hit that outlier, and the other two passages show ≈0.45 on both arms. An order change of the same precision moves every one of these metrics as far as #1353/#1354 do.

So by #796's criterion both steps are within noise. We'd keep the shared kernels' single order, fuse the combine back into add_rms_h for prefill, and restore the +1e-20 guard with its test, unless you need byte equality.

[ashhart]

Thanks, this is convincing. The placebo settles it: reversing one existing fold on main moves the float64 error, the argmax misses and the teacher-forced loss as far as #1353 and #1354 do, so a per-class worst-column test would mostly measure luck. "Not worse beyond noise" is the right bar for a reorder of the same precision, as on #796:

• paired per-class differences, prompt and decode separately, within noise after correcting for the number of classes;
• teacher-forced loss unchanged within noise;
• a placebo like yours as the yardstick.

Please keep the shared kernels' single order.

Two things from the earlier review still stand, because they aren't statistical:

1. Prefill speed. The 1.1–1.8% slowdown had the same sign in every run, so it needs to go. Fusing the combine back into add_rms_h for prefill, as you suggest, should do it; we'll A/B prefill and decode on a GB10.
2. The zero guard. Please restore the +1e-20 guard and add the all-underflow top-k case to the tests.

Every token check stays as it is: drafted equals plain, concurrent equals solo, and resumed equals fresh on the new order.

Separately, the row on main that reads 52.5 where float64 has −1.18 (row 162, layer 18, column 2206) looks like more than reordering noise. Could you open an issue with how you found it (prompt, passage and position)? We'd like to look at it on its own.

[jschmied]

Thanks. All three points are done. Both branches are pushed as new commits on top, with no rebase: #1353 is now at 402c89d06, #1354 at 962f749f3, and #1354 carries #1353's two commits.

1. Prefill. Two launches caused it, one in each step:

• nemotron CUDA: RMSNorm rows run on the shared norms and combine (#548, step 1) #1353 (6fabe2fa6, 93712b40a on nemotron CUDA: the router top-k is moe_glue's sigmoid top-k (#548, step 2) #1354): add_rms_h takes the expert slots as its residual. The new kernel is tf_moe_add_rms_rows_bf16h. Each thread sums its 8 columns over the slots in slot order with fp32 fmas, rounds to bf16 the way the combine stores it, and adds the result. The launch and the delta row are gone. The bytes are nemotron CUDA: RMSNorm rows run on the shared norms and combine (#548, step 1) #1353's: same kl dump sha (c94ddc4a3f2f) and same segment state.
• nemotron CUDA: the router top-k is moe_glue's sigmoid top-k (#548, step 2) #1354 (ffde19314): step 2 also had an extra launch, the combine that summed the router's K slices before the top-k. tf_moe_topk_sigmoid_score_warp now adds the slices itself, in slice order. That is the same chain as the combine's, so the bytes stay nemotron CUDA: the router top-k is moe_glue's sigmoid top-k (#548, step 2) #1354's (kl sha 19475174fbfe). Buffers.ones is gone too.

Both keep the shared kernels' single order. In norms.fatbin only the new function was added; in moe_glue.fatbin only the score top-k changed, on sm_80/86/89/120/121.

Served A/B on a GB10, Lightning MLX 4-bit, fresh server each arm-round, alternating order. Cold TTFT is the median of 3 per round; the table gives the median of the paired per-round ratios against main:

cold TTFT	2k	8k	32k
main (2e33e297d)	267 ms	1,039 ms	4,872 ms
#1353 before, 7 rounds	+1.1%	+1.1%	+0.9%
#1353 now, 7 rounds	-0.1%	+0.1%	+0.1%
#1353 now, 5 more rounds	-0.5%	+0.2%	+0.2%
#1354 now, 5 rounds	-0.6%	-0.6%	-0.6%

Plain decode (no drafts), 3 rounds, 1 and 4 streams:

• nemotron CUDA: RMSNorm rows run on the shared norms and combine (#548, step 1) #1353: -0.4% to +0.3% against main.
• nemotron CUDA: the router top-k is moe_glue's sigmoid top-k (#548, step 2) #1354: -0.9% to -1.3%. It isn't the slice sum: a microbenchmark has it no slower than combine plus top-k at 1 to 2,048 rows. The remaining candidate is step 2's top-k kernel against tf_nemo_topk, which I haven't isolated yet.

Drafted decode moves with acceptance, because the bytes differ from main's, so it isn't a speed comparison here.

2. Zero guard (b4c60b39e): tf_moe_topk_sigmoid_score_warp divides by sum + 1e-20 again. The float64 references (glue_ref.topk and the contract's score check) do the same. A new glue case routes three rows whose logits are two slices of -50 each. The ids match float64's, the routed weights are 0 rather than NaN, and the shared slots are 1.

Token checks, on both heads: drafted == plain, together 1/2/4/8 == solo, windows == pages, widths 63/63, shared widths 504/504, prompt cuts, resume and paged all pass. Glue is 1470/1470 on #1353 and 1472/1472 on #1354. The 16k needle is found. lean_check finds nothing new on either branch, and zig build test exits 0 on both.

3. The row-162 outlier is now #1489.

[ashhart]

Thanks for fusing the combine back into add_rms_h, and for opening #1489. We checked the head against our last review: the combine now runs inside tf_moe_add_rms_rows_bf16h (zig/kernels/cuda/norms.cu:469) in one launch with the same fma chain and rounding as before, and the shared kernels keep their single order. On a merge with current main, zig fmt, the lean check (no new findings), test-cuda-host and the Linux cross-compile all pass.

Two asks:

1. The contract's moe_add_y and moe_add_h kinds run one bf16 slot at weight 1 (zig/tests/cuda/contract_norms_tail.zig:198 and :220). There the combine is r itself, so the contract cannot see the slot order or the bf16 rounding before the add, the fp32-slot instance never runs, and moe_add_h's "no sums" (line 45) holds only for one slot. Please run them at Nemotron's 8 slots with real weights, fp32 and bf16, declare the slot chain, and add a planted control (say the combine kept in fp32 before the add) that goes red.
2. The description still shows the combine-then-norm pair, the old records and the -1.5 % prefill table. Please update it for the fused kernel (128 registers) and your new A/B.

A small point: after the fusion this step no longer calls MoeGlue (cuda_glue.zig:112, loaded at cuda_kernels.zig:113), and contract_moe.zig:169's 8 x 2,688 combine shapes describe a launch Nemotron no longer makes.

Neither ask changes the bytes, so we'll start the GB10 A/B of prefill and decode now, with the token checks on Lightning, Super VL and Bolt, whose Mamba group norms move too.

[jschmied]

Thanks. Both asks and the small point are done, pushed as three commits on top (fast-forward): #1353 is now at 8a77e3ec0. The third makes the prompt pass level with main.

1. The moe kinds (76ab46f02) now run Nemotron's 8 slots: 6 routed weights (sigmoid scores normalized to 2.5) and 2 shared slots of weight 1, at 16 and 300 rows × 2,688. There are four kinds: y and h, each with bf16 slots (the prompt's down projection) and fp32 slots (decode's).

• The slot chain is declared. h's kinds declare fp32 sums with depth 8: one fma a slot in slot order from 0, rounded to bf16 before the add. y's depth is that chain plus the squares'. "No sums" is left on add_h only.
• The float64 check takes the combine as a range: the exact sum ± gamma(8) of the terms' magnitude, rounded to bf16 at both ends. h and y are checked against that range.
• Gate 4 gives h's kinds slot terms that are exact in fp32 in any order, with one term past fp16's range. fp16 sums miss by inf×.
• The planted control, tf_ctl_tail_moe_keep_f32, runs the same body with the combine kept in fp32 before the add. It goes red at shapes: row 0 h[1] = 1.8671875 outside float64's [1.875, 1.875].
• To share that body, add_rms_h and its 16-byte helpers moved from norms.cu to norms_vec.cuh, with a KEEP template flag that is false in norms.cu. norms.fatbin's SASS is unchanged on sm_80/86/89/120/121, and contract_controls.fatbin only gains the control.
• kernel-contract on the norm kinds: the 4 moe kinds and the 12 other norm kinds pass gates 1-4, and the control is red as planted.

2. The description is updated for the fused kernel (128 registers on all five SMs), the new records and the new speed table.

Prompt time (8a77e3ec0): cold TTFT was still 0.1-0.3% slower than main. An nsys trace of an 8k and a 32k prompt pass put it on the Mamba group norm: tf_rms_rows_bf16h_parts_round took 16.8 / 67.6 ms where tf_nemo_group_rmsnorm took 14.3 / 57.3. It ran one 256-thread block a 512-wide group, with 64 threads holding data.

• tf_rms_rows_bf16h_parts4 and _parts4_round run parts up to 512 wide four a block, 64 threads each; rmsParts picks them by width. Each thread's 8-term chain and the warp tree are the same, and warp 0 + warp 1 is the 8-warp tree's total when warps 2-7 hold zeros, so the bits are the old kernel's: y and the sums byte-equal over 224 calls (widths 64-512, 1-2,048 rows, both roundings, NaN, ±inf, tiny values, with and without sums). At 8 groups × 512: 143 µs at 2,048 rows, 578 at 8,192, 2,261 at 32,768 (was 181, 710, 2,876).
• The contract's parts kinds run both forms. Records 36/34/34/36/36 (_parts4) and 36/34/34/35/35 (_parts4_round), 32 B static, bound 256, the same under 13.2. norms.fatbin's other functions keep their SASS on all five SMs.
• Served A/B against main 2e33e297d, 12 rounds, a fresh server for each arm and round, alternating order, medians of the paired ratios: cold TTFT -0.9% / -0.1% / -0.1% at 2k / 8k / 32k (8, 7 and 6 of 12 faster, 0, 2 and 1 level); plain decode +0.5% / +0.4% at 1 stream (12 of 12), +0.2% / +0.1% at 4 streams (chat / code).

Small point (70095272e): step 1 no longer loads MoeGlue. The engine, its Glue and the glue rig drop it, and the combine contract drops the 8 × 2,688 shapes. #1354 loads MoeGlue again in its own commit, since its top-k and Bolt's latent sum use it.

Checks on 8a77e3ec0 (the bytes are unchanged):

• Token checks: drafted == plain, together 1/2/4/8 == solo, windows == pages, chat sha 1a5f995f724e, kl dump sha c94ddc4a3f2f, segment state 8bddfcd16173. Widths 63/63, shared widths 504/504, prompt cuts, resume and paged pass. Glue 1470/1470.
• Records for nvcc 13.0 and 13.2 are current; the limit failures are the same lines as before.
• kernel-contract: the 12 norm kinds and the tail kinds pass gates 1-4, and the controls are red as planted.
• zig build test exits 0, and lean_check finds nothing new.

[ashhart]

Thanks @jschmied for the three commits. Before they arrived, we ran the head up to 76ab46f02 on a GB10, on main a1e6d3424, alternating main and the PR (base, new, new, base):

• Tokens differ from main. On Nemotron 3.5 Lightning (MLX 4-bit), 17 of 20 prompt rows give different tokens: 2k to 64k, from position 0, 5,000 and 10,000. On Nemotron 3.5 Super-VL (NVFP4), 16 of 20 do. Each build repeats itself exactly, and both main runs agree, as do both PR runs. So the PR's path isn't byte-equal to main's tf_nemo_* kernels on these checkpoints, although its parts may match the two-launch form they were checked against.
• Decode is slower. On Lightning, one stream is −7.6% and two −5.6%, and decode after a 64k prompt −13.8%. On Super-VL, decode after a 16k and 32k prompt is −6.1% and −8.7%. Prefill is within 0.8%.

Could you find where the bytes part from main? For example, compare each new launch's output with main's tf_nemo_embed16, tf_nemo_add_moe_norm, tf_nemo_concat_norms and tf_nemo_group_rmsnorm on a real layer's inputs (decode's fp32 slots and the prompt's bf16 slots), and say which kernel differs first. The contract checks bytes against the shared two-launch form, but not against what main runs for Nemotron. If the new order is closer to a float64 reference, that's worth showing, but for this step the bar is the same tokens as main. Then we'll rerun on 8a77e3ec0 or later: tokens on both checkpoints, and the decode speed.

[jschmied]

Thanks for the run. Fixed in one commit on top (fast-forward): #1353 is now at 8e1115e05. Its tokens equal main's on every row I could run.

Where the bytes parted. I built main with a local dump in the glue (never committed) and saved the first calls of each Nemotron glue kernel: the inputs before the launch and the outputs after, for decode's fp32 slots and the prompt's bf16 slots (Lightning MLX4, a 2,048-token prompt, drafted and plain). Then I replayed each call through the PR's kernels.

• The first output to differ is layer 0's Mamba group norm. tf_rms_rows_bf16h_parts4_round summed its squares as 8 contiguous elements a thread and a tree, where tf_nemo_group_rmsnorm strides them by 128 threads and adds its warps in order. The 64-value group sums (the next matmul's xs) ran 8 in order and then a lane tree, where main's run in one chain. At the prompt: 40 of 8.4M outputs and 1,956 of 131k sums differ. At one decode row: the outputs are equal, but 2 of 64 sums differ, so the first out_proj already differs.
• The MoE residual added the 2 shared slots into the routed fma chain, so ((r + s0) + s1) where main has r + (s0 + s1). Its squares and sums ran in the vector order too. The MTP norms took the tail's order.
• Eps, the rounding points and IEEE div/sqrt never differed. Only the order of fp32 sums did.

Fix (8e1115e05). norms_strided.cuh holds main's order as templates: thread t holds elements t, t + T, … (fma chain, xor butterfly, warps added in order, each 64-value group summed in order by one thread).

• tf_rms_rows_bf16h_parts = 256 threads × 16 (the MTP halves, one launch again). _parts_round = 128 × 8 (the Mamba groups).
• tf_moe_add_rms_rows_bf16h is the fused residual with the routed chain and the shared sum kept apart (a routed argument).
• The parts4 forms and the 16-byte combine are gone, and add_rms_h is main's again. The control tf_ctl_tail_moe_keep_f32 is the same body with the slot sum unrounded.
• Replay: every output of all 28 dumped calls (out, h, y, xs) equals main's bytes.

Tokens (#1353 merged into main db3fd5aba, against main). The rows are 2k, 4k, 8k, 16k, 32k, 48k and 64k long, from positions 0, 5,000 and 10,000, with 32 tokens each:

• Lightning MLX4: 21 of 21 rows drafted, 21 of 21 plain, 8 of 8 chat replies.
• Bolt: 21 of 21 rows, 8 of 8 replies.
• The old head merged the same way gave 13/21, 3/8, 8/21 and 4/8, which matches what you saw.
• Super-VL NVFP4 isn't on this box, so I couldn't check it.

Decode. On the code prompt, nsys shows the old head's slower decode came from its different tokens, not from kernel time: 75 rounds and 180 accepted against main's 63 and 192 (5.54 vs 5.20 ms a token), while its norm kernels were faster per call. With main's tokens the PR runs 63 rounds at 5.15 ms a token. The norm kernels match main's µs per call: 5.21 vs 5.27 (residual), 1.90 vs 1.86 (group), 3.10 vs 2.98 (MTP).

Served A/B against main db3fd5aba. Four rounds alternating base, new, new, base, with a fresh server for each arm and round, on a quiet host. Medians, + = faster:

• Cold TTFT 2k–64k: −3.0 / −4.4 / −4.1 / −4.0 / −2.1 / −0.9 %. The per-round minima overlap (16k: 2.147–2.195 s main, 2.161–2.197 s PR), and the same kernels with step 2 on top measure −0.2…+0.2 %.
• Decode: plain 1 stream +0.0 %. Drafted 1 stream +0.6 % (code) and −0.0 % (chat). 2 streams +5.5 % and +5.9 %.
• Decode after a 16k / 32k / 64k prompt: +11.7 / +3.2 / +5.6 %. That's 2 requests each, and these cells spread ±10 % within one arm.

Checks on 8e1115e05:

• drafted == plain, together 1/2/4/8 == solo, windows == pages, repeats. Widths 63/63, shared widths 504/504. Prompt cuts, resume and paged pass. Glue 1470/1470, against main's host references again.
• The kl dump's sha is 805afe768e0b, the same as main's. The chat and needle replies are main's.
• SASS against the base on sm_80/86/89/120/121: 0 changed functions. norms.fatbin gains the three strided forms, contract_controls gains the control, and nemotron_norms goes.
• Records are current for nvcc 13.0. Under 13.2 they are stale by 1 register on sm_120/121 for the residual, which still passes. The only limit failures are the existing tf_q35_attn_partial lines, which cuda_lane_pending.zig keeps red until they are recorded.
• kernel-contract: the 12 norm kinds pass gates 1–4 (the group sums are checked as one chain for the strided kinds), the tail kinds pass, and the MoE control is red.
• lean_check finds nothing new, and zig build test exits 0.