[jschmied]

What

First step of #1409. Pure refactor: no kernel changes its code. cuobjdump -sass is identical for every function of every image that includes a touched file.

Directory. The two headers that were already shared move to zig/kernels/cuda/common/, the layout the HIP backend already uses (zig/kernels/hip/common/):

• qmm_frag.cuh (22 includers): cp.async, ldmatrix, bf16 mma, nibble decoding, swizzle. Its first line pointed at zig/tests/cuda/copies.py, which left the tree with the Python engine; it now says what the header holds.
• glue_math.cuh (16 includers): the exact scalar math with a host twin in glue_math.zig.
• The two build-time generators that read their text follow the move: qwen4_ple_cuda_host.zig (embed path, the include it strips) and hybrid_ple_source.zig (the include it emits).

New headers, namespace tf_common:

• common/reduce.cuh: warp_sum and warp_max (xor butterfly, 16 down to 1); block_sum_tree and block_max_tree (each warp's butterfly, then one more butterfly over the warp results); group_sum8 (a thread's 8 values in order, then 8 lanes in a fixed tree).
• common/vec.cuh: pack8 / unpack8, eight bf16 values in one 16-byte word.

Added to qmm_frag.cuh: ldmatrix4t (x4 .trans), fp4pair (e2m1 pair to bf16), ld_nc (16-byte non-allocating global load).

Copies removed. Each unit drops its own definition for a using declaration and keeps its call sites, except the train kernels, whose block_sum / block_max calls now name the order (block_sum_tree, block_max_tree).

primitive	units now using the shared one
warp_sum	17, including glue_math.cuh's own copy
warp_max	6
block_sum_tree / block_max_tree	7 / 3 (the train kernels and train_mixers.cuh)
pack8 / unpack8 / group_sum8	5 / 5 / 3 (the act and norm row headers)
ldmatrix4t	5 (block, tree and window attention, window_rows.cuh, flashnext_attn.cu)
fp4pair, ld_nc	qmmf.cu, drama3_fp8.cu

67 files, +542 / −483: 65 under zig/ (+274 / −483), plus the SASS tool and its test.

Tool. tools/zig/cuda_sass_identity.py BASE HEAD builds every catalog image whose unit reaches a file changed between two revisions, from both trees with each tree's flags, and compares cuobjdump -sass function by function; --receipt writes the digests, and two planted controls must go red. tools/zig/cuda_sass_identity_test.py is its test.

Left alone on purpose

• Files in open PRs keep their copies until those land: prefill_attention.cu, qwen35_attn_tc.cuh, prompt16.cu, nvfp4_experts.cu, fp8_experts.cu, norms.cu.
• Files whose text a generator parses: flashnext_gdn.cu, flashnext_batch.cu, flashnext_hc.cu, the qwen4_* units, and flashnext_attn.cu's block_sum.
• The serial block order (warp sums added one after another in warp order). Its copies sit in generator-parsed files. h3.cu has the same order, but its h3_norm did not compile to the same SASS through a shared function, so it keeps its own.
• Copies that differ in signature or reduction order (drama3's block_sum and block_max, the double-precision deepseek copies) and the precision-bearing helpers (sigmoid/silu/softplus on expf vs tf_exp, inv_rms). These wait on the nemotron CUDA: RMSNorm rows run on the shared norms and combine (#548, step 1) #1353 rule.
• The bf16 rounding helper. It is named bf in some files and rbf in others, and bf means widening in others still, so it needs call-site renames. That is a separate PR.

Verification

• SASS: tools/zig/cuda_sass_identity.py db3fd5aba b6ea38468 (current main against its merge with this branch), nvcc 13.0.88: every image whose unit includes a changed file, 53 images (the catalog's plus h3 and contract_controls) with 2,890 functions, built with the build's flags for sm_80/86/89/120/121. 0 changed, 0 added, 0 removed functions in every image; exit status 0. The digests cover every instruction line at any address width and each instruction's second encoding word. Receipt: one row per image, arch and function with both SHA-256 digests.
• Planted controls go red (each built against the head):
  • an altered fp4pair constant (common/qmm_frag.cuh) changes 52 functions in 1 of the 26 images that include it (qmmf);
  • a reversed warp_sum butterfly (common/reduce.cuh) changes 580 functions in 24 of the 34 images that include it.
• lean: 427 problems vs 428 on the base, 0 new.
• zig build test -Dnvcc=/usr/local/cuda/bin/nvcc: RC 0, 0 errors.
• The tool's own test: python3 -m unittest cuda_sass_identity_test (from tools/zig) compares a canned dump of two functions that differ only at address 0x10000, and only in a second encoding word. Both cases fail with the tool's first version.

Conflicts: #1353 and #1354 delete nemotron_norms.cu, whose #include line this PR rewrites. Whichever lands second drops the file; nothing else conflicts with an open PR.

Part of #1409.

[ashhart]

Thanks for starting #1409 with a pure move, and for planting controls that show the comparison can go red. We checked a merge with current main: it is clean, nothing still includes the old paths, zig fmt passes, the lean check drops one finding (the old qmm_frag.cuh header line), and test-cuda-host and the Linux cross-compile pass. The copies we compared match the shared bodies token for token, and no contract record changes, so the records hold as long as the SASS does.

Two asks:

1. The SASS comparison is the evidence for the whole PR, but only its counts are here, and it ran against 8c36359d5, 86 commits behind main. Please commit what you ran (a script under tools/ that builds the touched images for the five SMs from two trees and compares cuobjdump -sass per function, with the two planted controls), or its digests as a receipt, and rerun it on a merge with current main. We'll run it on our side too.
2. zig/tests/fixtures/gradient-serving.tsv:34 and :55 still pin glue_math.cuh and qmm_frag.cuh at their old paths, so zig build test-gradient-serving now stops at FileNotFound. Please point them at common/.

A small point: glue_math.cuh's own warp_sum used __fadd_rn, while the shared one uses += (common/reduce.cuh:14), which nvcc may fuse with a multiply under its default --fmad=true. Every glue includer passes --fmad=false, so nothing changes now; __fadd_rn would keep the glue's IEEE promise without the flag, if your comparison shows no change.

Then we'll run the kernel contract and kernel-contract-limits on the merge.

[jschmied]

Thanks. Pushed 4e5d837a8, rebased on main 5897e0bca. Main has gained five commits since then (41abcf402), none under zig/kernels, and the merge is clean. Why I rebased: the pin fix below changes prompt16.cu's pin line, which main also changed, so a merge would have conflicted.

1. SASS comparison, committed as tools/zig/cuda_sass_identity.py. Run it as python3 tools/zig/cuda_sass_identity.py BASE HEAD [--receipt FILE]; BASE and HEAD can be any revisions, a merge commit included. It works in five steps:

   • Exports both trees' zig/kernels and zig/build.
   • Selects every image whose unit reaches, through quoted includes, a file changed under zig/kernels: the catalog's, plus those other build files name (h3, contract_controls).
   • Builds each image from both trees with that tree's flags, torch's nvcc flags and -I zig/kernels/cuda for sm_80/86/89/120/121, as cuda_fatbin.zig does (min_sm, a targets).
   • Compares cuobjdump -sass function by function.
   • Plants two controls against the head, each of which must change functions: fp4pair's constant 0x7E807E80 → 0x7E817E80, and warp_sum's butterfly reversed to 1..16.

   It exits 0 only if every function matches and both controls go red.

   Result on main 41abcf402 against this head (its kernels are those of the merge with main): 53 images, 2,890 functions, 0 changed / added / removed. The controls changed 52 functions in 1 of 26 images (fp4pair) and 580 functions in 24 of 34 images (warp_sum). nvcc 13.0.88.

   Receipt. A per-function digest receipt (image, arch, function, base and head sha256; 2,890 rows, all equal) is in sass-1440-receipt.tsv, which I can attach or commit if you'd like it in the tree, sha256 cfbbbdeeaacfa6792531824b508464310398e706c6de74abf12a23eb93e626de.

   Caveat. qwen4_gdn and qwen4_ple are compiled from their checked-in units, not from the text the build generates. The tool marks them in its output.

2. gradient-serving.tsv. Both pins now name common/ with the moved files' bytes. The move also changes the include lines of 12 units whose pins matched main (ck_experts, ck_lane, gdn, glm_kda, lfm2, the five nemotron_*, prompt16, train_mixers), and those pins take the new bytes. route_rows.cu passes on main through pocket-serving-approved.tsv; its approval takes the new bytes, with a note naming cuda: shared device primitives in zig/kernels/cuda/common (#1409, part 1: SASS unchanged) #1440.

   Main fails this step on its own. zig build test-gradient-serving on main 5897e0bca stops at serving source changed: zig/kernels/cuda/contract_controls.cu. By my count, 63 of the 168 pins match neither their hash nor an approval there. On this head the step stops at the same file, and the failing set is the same 63, with none added. I left those pins alone: refreshing them is a separate decision, and I can open it if you want.

3. warp_sum and __fadd_rn. I tried it, and the SASS is not unchanged. The comparison on common/reduce.cuh alone changes tf_train_project in train on all five SMs (5 functions). train.cu is built with -O3 and no --fmad=false, so nvcc fuses its += there. Every other includer is identical. I dropped the change, so the shared warp_sum keeps += and the PR stays a pure move. __fadd_rn would need its own PR that re-qualifies train.

Checks on the head: zig fmt --check zig clean; lean 0 new (427 against 428 on main); zig build test -Dnvcc=... RC 0; ATC overlap shows no conflicts from this branch.

[ashhart]

Thanks @jschmied for committing the comparison as a tool, and for trying __fadd_rn before dropping it. Keeping += in the shared warp_sum is the right call for a pure move; __fadd_rn can come in its own PR that re-qualifies train.

What we checked:

• The pins now name common/ and match the moved files (zig/tests/fixtures/gradient-serving.tsv:34, :55).
• test-gradient-serving now stops at contract_controls.cu, exactly where it stops on main. Hashing every pin gives the same 63 failures on main and on the merge, so this PR adds none. Those 63 are ours to sort out separately.
• Your run against 41abcf402 covers current main, because no CUDA kernel has changed on main since then.
• On the merge, these all pass: zig fmt, the lean check (one finding fewer), test-cuda-host, test-runtime-host, test-qwen4-ple-host and the Linux cross-compiles.

Two asks before we rely on the tool on our side:

1. tools/zig/cuda_sass_identity.py:107 only takes instruction lines whose address has four hex digits. cuobjdump prints offsets from 0x10000 on with five, so anything past a function's first 4,096 instructions is left out of its digest. Please match [0-9a-f]{4,}, and also keep the /* 0x… */ continuation lines (each instruction's second encoding word), as zig/tests/cuda/oracle/sass_compare.py:24 does. A small test on a canned dump, two functions that differ only at a five-digit address, would fail before the fix. Then please rerun on a merge with current main and update the receipt.
2. The description still gives base 8c36359d5, 2,880 functions and the old control counts. Please replace them with the tool's run.

Then we'll run the tool ourselves, and the kernel contract and kernel-contract-limits on the merge.

[jschmied]

Thanks. Both done, pushed as 79a275b5c.

1. The tool. functions() now matches [0-9a-f]{4,} and keeps the /* 0x… */ continuation lines, as sass_compare.py:24 does. tools/zig/cuda_sass_identity_test.py (python3 -m unittest cuda_sass_identity_test from tools/zig) runs a canned dump of two functions: one case differs only at 0x10000, one only in an instruction's second encoding word. Both fail on the old matcher and pass now.

   Rerun on b6ea38468, the merge of this branch with current main db3fd5aba (made before a docstring-only amend; its zig/ tree is the same), base db3fd5aba, nvcc 13.0.88: 53 images, 2,890 functions, 0 changed, 0 added, 0 removed, exit 0. The controls are still red: fp4pair changes 52 functions in 1 of 26 images, and warp_sum 580 in 24 of 34. Receipt attached (https://gist.github.com/jschmied/d68c5b6514d42334f9ee6e898d12ab24).

   One thing to know before you rely on it. My first run on the same two revisions reported one difference: nemotron_route, sm_86, tf_nemo_router_tiles, with a base digest that no build has produced since. That file's only change here is its include path. Six more builds of that image, three from each tree, and the full rerun above all give ae587bc2…, which is also the head's digest in the first run. So in that run nvcc's sm_86 output was not the same from build to build, with four builds running side by side and other jobs on the box. If you see a lone red function, rebuild that image before reading it as a change. I can add an automatic second build of any differing image if you want it in the tool.

2. Description. Updated: base db3fd5aba and head b6ea38468, 53 images and 2,890 functions, the controls as the tool prints them, and the current file count.

Checks on the branch: zig fmt --check zig, lean 427 problems vs 428 on the base (0 new), zig build test -Dnvcc=/usr/local/cuda/bin/nvcc RC 0 with 0 errors.

[ashhart]

Thanks @jschmied, both are done:

• The matcher now takes five-digit addresses and each instruction's second encoding word.
• The canned-dump test fails on the old version.
• The rerun on a merge with db3fd5aba finds all 2,890 functions in the 53 images unchanged, with both controls red.
  No CUDA kernel has changed on main since, so that run still covers it.

On the lone red function: this PR doesn't need an automatic second build, and a follow-up would be welcome. If you add
it, here is what we'd want:

• Rebuild a differing image from both trees, one build at a time.
• Count a function as changed only if it differs again.
• If two builds of the same tree disagree, report it as unstable and exit non-zero, so a flaky build can never pass
  quietly.

Next on our side:

• the tool on a merge with current main;
• a comparison of the qwen4_gdn, qwen4_ple and Flash Next hybrid PLE images the build generates, main against the
  merge, since the tool compiles their checked-in units instead;
• the kernel contract and kernel-contract-limits on a GB10;
• one served model against main by token hash.

[jschmied]

The rebuild follow-up is up as #1532, a draft stacked on this PR: a differing image is rebuilt from both trees one build at a time, a function counts only if it differs again, and a tree whose two builds disagree fails as unstable.

[ashhart]

Thanks @jschmied, this is merged. We built every CUDA image for main and for this branch at sm_121 and sm_80 and compared cuobjdump -sass function by function: all 1,866 functions in all 104 images are identical, with none changed, added or removed. Flash Next FP4 on a GB10 also gives main's tokens in every prompt row. With the machine code identical, the kernel contract can't come out differently, so we didn't run it again. Thanks for the clean first step of #1409, and for the SASS identity tool.

[ashhart]

Thanks @jschmied: the shared CUDA headers shipped in TensorFold 1.1.9.