RFC: shared CUDA device primitives (one header set for fragments, FP4 decode, copies and reductions)

Why

Counted on main at 03c28201a, over the 215 .cu/.cuh files in zig/kernels/cuda:

ldmatrix: 17 files carry their own inline asm for it (23 asm lines), among them the four attention kernels (block, tree, window, prefill), the EXL3 prompt and block kernels, the GLM glue and qwen27_train.cu.
Fragment headers: four separate ones (qmm_frag.cuh, glm_qmm_frag.cuh, mma4.cuh, drama3_common.cuh), and mma.sync inline asm in 12 files.
cp.async: 70 inline-asm lines across 25 files.
FP4 decode: fp4pair is defined three times (prompt16.cu, qmmf.cu, nvfp4_experts.cu), and the copy in nvfp4_experts.cu already differs from the other two. e2m1 encoding is in two places (glm_quant.cu, nvfp4q.cuh).
warp_sum: 20 copies, which read as 7 different texts after normalising whitespace and the argument name.

Each copy is small, but together they cost in three ways:

Reviewing: a reviewer has to check every copy's rounding and lane order on its own.
Fixing: a fix or an arch port (e.g. sm_121, where tcgen05 does not exist) has to find every copy.
New families: a new family (Ling, the Gemma 4 E-series, Qwen4) copies the nearest file again.

The kernel contract checks each kernel, but nothing checks that two copies of the same primitive agree.

Proposal

One small header set, zig/kernels/cuda/prims/, holding only exact, order-defined primitives:

Header	Contents
`frag.cuh`	`ldmatrix` x1/x2/x4 (+ trans), the bf16/fp16 `mma.sync` m16n8k16 atoms, the swizzle address helper
`copy.cuh`	`cp.async` 4/8/16 bytes, commit/wait groups, an mbarrier ring for persistent producer/consumer kernels
`fp4.cuh` / `fp8.cuh`	exact e2m1 → bf16 decode (`prmt` byte tables), e4m3 scale broadcast to bf16x2, encode helpers
`reduce.cuh`	warp and block sums/max in one canonical order per width, named by that order

Rules:

Exact only: no approximate math and no lossy activation formats, matching the existing project rules.
One canonical order per reduction width, documented in the function name or its one-line comment.
Contract records: every primitive gets its own records, with a planted control that goes red.
Lean limits: files stay inside them (one-line comments, under 600 lines).

Migration

No kernel changes bytes because of this RFC. A kernel moves onto the prims only as a pure refactor, and the PR shows cuobjdump -sass identical for every function in the touched fatbins (sm_80/86/89/120/121). A copy that differs from the canonical one, such as one of the warp_sum variants, stays local until its family migrates deliberately under its own gates.
One family or kernel group per PR, starting where the copies are byte-identical today: fp4pair in prompt16.cu and qmmf.cu, and the ldmatrix wrappers.
New kernels (Ling, the E-series, Qwen4) use the prims from the start.

Non-goals

Not a DSL, a compiler or an autotuner: TF keeps one hand-written kernel per shape class.
Not a dependency on CUTLASS/CuTe, ThunderKittens or Humming. External libraries stay references for design, not includes, because of build time and because their reduction orders are not ours.

Questions

Is a shared prims/ directory acceptable, or should the primitives live next to an existing header (qmm_frag.cuh)?
Should the contract registry list primitives as their own entries, or only through the kernels that use them?
Any family you want left out of the first refactor PRs?

Related to the acceptance question on #1353: migrating a family onto a shared primitive whose summation order differs from the family's own changes that family's bytes, even though precision is unchanged. Which rule applies decides how far this RFC can go:

Byte equality to main: only copies that are already identical can move (the two equal fp4pair copies, most ldmatrix wrappers). The 7 different warp_sum texts and the per-family norm orders would stay local, or come in as named legacy-order variants.
Not worse beyond statistical noise (paired float64 error per class within 2 se, as proposed on nemotron CUDA: RMSNorm rows run on the shared norms and combine (#548, step 1) #1353): each family migrates once onto the canonical order under its usual gates, and the prims keep one order per width.

We'd apply whichever answer you give on #1353 here too, and write it into the RFC's migration section.

@jschmied

Thanks @jschmied. Part 1 is #1440, the pure move with the SASS comparison tool. Your answers there are in, and next we'll run the tool, the kernel contract and its limits on a merge with main.

On which rule applies to later parts: moving a kernel onto the prims has to leave its served tokens equal to main's, as for any kernel change. So the migration section stands as you wrote it: only copies that are already identical move (the two equal fp4pair copies and most ldmatrix wrappers), and the differing warp_sum texts and per-family norm orders stay local, or come in as named variants that keep their current order.

On your questions: common/, as #1440 has it, suits us, and no family needs leaving out of moves that keep the SASS identical. Whether primitives get contract entries of their own can wait for the first PR that builds a new kernel on them.

Thanks. Agreed for part 1 and for every pure move: SASS identical, nothing else to show.

Two points on the rule for the later parts.

warp_sum isn't a variant case. I overstated it in my earlier comment. The 20 float copies are all the same xor butterfly, 16 down to 1. The "7 texts" differ only in whitespace and the argument name, so they can all move with SASS unchanged. The copies that really carry an order are block_sum (two orders: the warps' sums added one after another, or a second butterfly), inv_rms (6 copies, 5 texts) and the activations (expf against tf_exp).

"Served tokens equal to main" as the bar for an order change. Main's order is one valid choice among several, not a reference.

The placebo: on nemotron CUDA: RMSNorm rows run on the shared norms and combine (#548, step 1) #1353 (receipt), reversing one existing fold on main (inv_rms_v, 4,2,1 → 1,2,4) moved the float64 error, the argmax misses and the teacher-forced loss as far as the PR did. You called that convincing (here). Main's own tokens would fail "equal to main" after that one-line reversal, with nothing gained or lost.
The near-tie: Nemotron 3.5 Lightning CUDA prefill: one row departs from float64 after a near-tie routing pick in layer 6 #1489 shows why. One routing pick in layer 6 is decided by 8.3e-5, about 22 fp32 ulps at 43. Main picks expert 46, and float64 picks 38. "Equal to main" would hold every later change to main's pick there, not to the one float64 makes.
What it costs the RFC: under that bar every family keeps its own reduction order indefinitely, and the prims can only collect copies that are already identical.

What equal tokens protects is real: a release shouldn't change replies silently, and bit-equality is a complete test that noise can't hide a bug behind. I'd keep that for every refactor, and handle an intended order change as its own step:

Pure moves and refactors: SASS identical (the cuda: shared device primitives in zig/kernels/cuda/common (#1409, part 1: SASS unchanged) #1440 tool), as now.
A deliberate order change: its own PR and its own changelog line saying replies change. It is checked against float64, not main, and must be no worse than a placebo reorder of main beyond statistical noise.
It must also keep every behavioural check exact: repeats, drafted == plain, concurrent == solo, resumed == fresh.
Where the new order is closer to float64, the PR says so.

Would that work for you? If you'd rather keep tokens-equal-to-main for now, I'll write that into the migration section and keep the differing orders as named variants.

Yes, that works, and it's the bar we set on #1353 on 27 Oct. My note this morning was too strict: "served tokens equal to main" is right for pure moves and refactors, where the #1440 tool shows identical SASS, but not for a deliberate reorder. So, for the later parts:

Pure moves and refactors: identical SASS, shown with the cuda: shared device primitives in zig/kernels/cuda/common (#1409, part 1: SASS unchanged) #1440 tool.
A deliberate order change at the same precision: its own PR, and its own changelog line saying replies change. It's judged against float64, not main: paired per-class differences, prompt and decode separately, corrected for the number of classes, and teacher-forced loss unchanged within noise, with a placebo reorder of main as the yardstick.
Every behavioural check stays exact: repeats, drafted equals plain, concurrent equals solo, and resumed equals fresh.
Closer to float64: where the new order is, the PR says so.

Two things don't change. Nothing lowers a precision (accumulators, activations or stored formats), and no served cell gets slower.

Thanks too for the warp_sum correction. If all 20 copies are the same butterfly, they can move as pure moves, and block_sum, inv_rms and the activations are the real order variants.

cuda: fp4pair, fp8pair and e4m3f in common/fp_pairs.cuh, SASS unchanged

cuda: sigmoidf_, softplusf_, silu and silu_grad in common/act.cuh, SA…

cuda: fp4pair, fp8pair and e4m3f in common/fp_pairs.cuh, SASS unchanged

cuda: sigmoidf_, softplusf_, silu and silu_grad in common/act.cuh, SA…

cuda: glm_mm takes common/qmm_frag.cuh, and its 117-line copy glm_qmm…

cuda: attn_train and train_attn compile one gated attention backward,…

cuda: cp.async commit and wait from common/cp_async.cuh, SASS unchanged

cuda: ldmatrix and the f16 mma from their one home each, SASS unchanged

cuda: cp.async commit and wait from common/cp_async.cuh, SASS unchanged

cuda: ldmatrix and the f16 mma from their one home each, SASS unchanged

@jschmied

Status: the pure-move part is mostly done. Tonight #1564, #1565, #1588, #1591 and #1550, #1555, #1556 landed. In each,
every kernel function they reach compiles to the same code as before.

Left: a few local copies (ldmatrix in h3_attention.cuh, the GLM glue and index headers and qwen27_train.cu, and
some block_sum and inv_rms copies); the deliberate order changes, each as its own PR under the rule we agreed;
and that rule written into the RFC's migration section. Thanks @jschmied.

Navigation Menu

Repository navigation

RFC: shared CUDA device primitives (one header set for fragments, FP4 decode, copies and reductions) #1409

Why

Proposal

Migration

Non-goals

Questions

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

RFC: shared CUDA device primitives (one header set for fragments, FP4 decode, copies and reductions) #1409

Description

Why

Proposal

Migration

Non-goals

Questions

Activity

jschmied commented on Oct 26, 2026

ashhart commented on Oct 28, 2026

jschmied commented on Oct 28, 2026

ashhart commented on Oct 28, 2026

ashhart commented on Oct 30, 2026

Metadata

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

Issue actions