[Feature] Add rmsnorm backward (+ forward saves rstd) to enable training / downstream (quack) integration

[jhinpan]

Context

Spun out of RFC #749 (a contribution path for FlyDSL kernels into downstream OSS libraries). While analyzing the quack rmsnorm interface gap, we confirmed FlyDSL's norm family is forward-only:

• kernels/rmsnorm_kernel.py and kernels/layernorm_kernel.py expose only build_* / launch_* forward paths.
• There are no *bwd* / *backward* / *grad* files under kernels/, and a repo-wide autograd search returns 0.
• FlyDSL's own tests/kernels/test_rmsnorm.py exercises forward only (no gradient checks).

This is expected for an inference kernel library, but it is the blocker for training use cases and for backing a training-capable library (quack).

Goal

Add a correct, performant rmsnorm backward so FlyDSL can serve training, and so the downstream-backend PoC in #749 can target fwd+bwd rather than forward-only.

Scope / tasks

1. Backward kernel — compute dx, dweight (and dbias if bias is added). Accumulate weight/bias partials in fp32 for numerical accuracy (cf. quack's rmsnorm_bwd: per-SM fp32 partials, then reduce).
2. Forward must save backward state — the current forward does not emit rstd (the per-row 1/RMS). Backward needs it, so add an optional store_rstd / rstd output to the forward path instead of recomputing.
3. Autograd wrapper — a torch.autograd.Function (forward saves x, weight, rstd; backward calls the new kernel) so quack-style call sites get gradients transparently.
4. Fused-add (residual / prenorm) backward — kernels/rmsnorm_kernel.py already has a separate fused-add forward module; backward should handle the residual path too (grad w.r.t. residual_out).
5. Tests — extend tests/kernels/test_rmsnorm.py with gradient checks against a torch reference (bf16 / fp16 / fp32; several M/N including the N<=2048 small-N path). Reuse quack's tolerance style.
6. (stretch) layernorm backward — same pattern; can be a follow-up.

Why now

Direct enabler for #749 (training-capable downstream integration). Without this, any "fwd+bwd parity" with quack is net-new kernel work rather than interface adaptation.

Refs

• quack: rmsnorm_bwd (quack/rmsnorm.py:1252), RMSNormFunction (:1368).
• FlyDSL: build_rmsnorm_module (kernels/rmsnorm_kernel.py:113), forward has no rstd output; launch_fused_add_rmsnorm (:604).
• Related: [RFC]: a contribution path for FlyDSL kernels into downstream OSS libraries (PoC: Dao-AILab/quack) #749.

[jhinpan]

Implementation plan & PR breakdown

Spent some time reading the current forward-only implementation (kernels/rmsnorm_kernel.py, tests/kernels/test_rmsnorm.py) and mapping this out. Posting a concrete, staged plan so we can land it as a series of small, dependent PRs rather than one large change. Testing on MI355X (gfx950), 8× GPU.

Current state (for context)

• build_rmsnorm_module / build_fused_add_rmsnorm_module are forward-only. Forward already computes per-row rrms = rsqrt(mean(x²)+eps) but discards it.
• There is no backward/grad kernel anywhere in the repo, and no torch autograd / op-registration layer — the build_*_module factories are the public API.
• Reusable building blocks: in-block reductions (wave_reduce_add / block_reduce_add[2] over LDS), and an fp32 cross-block atomicAdd pattern already used in hgemm_splitk.py (llvm.AtomicRMWOp(fadd, …, syncscope="agent")). Note bf16 atomics are arch-gated (moe_blockscale_2stage.py), so weight/bias grads accumulate in fp32.

Math (matches quack rmsnorm_bwd_ref)

Per row, with x_hat = x·rstd, wdy = dy·γ:

• c1 = mean_N(x_hat · wdy)
• dx = (wdy − x_hat·c1) · rstd
• dweight += dy · x_hat (summed over rows M → length-N vector)

---

PR sequence (5 PRs, mostly dependent/inheriting)

PR 1 — Forward saves rstd (foundation, no dependency)

• Add optional store_rstd=False to build_rmsnorm_module. When on, an extra (M,) fp32 output is written (thread0 / lane0 per row) across all three paths (fast / generic / small-N N≤2048).
• Default off ⇒ existing signatures/tests unchanged (byte-for-byte).
• Small & self-contained; unblocks everything else.

PR 2 — RMSNorm backward kernel + tests (depends on PR 1)

• build_rmsnorm_bwd_module(N, dtype) mirroring forward's fast/generic/small-N paths, one block per row. Inputs x, γ, dy, rstd → outputs dx, dweight.
• Two-pass per row: reduce x_hat·wdy → c1, then write dx. dweight via fp32 atomicAdd into a zero-init (N,) buffer (single-kernel, matches hgemm_splitk pattern). All reductions/accumulation in fp32; only dx cast back to I/O dtype.
• Tests: _reference_rmsnorm_bwd + gradient checks vs torch across f32/f16/bf16, both aligned fast path and small-N N≤2048, existing tolerances.
• (Decision to confirm below: atomic vs two-pass scratch for dweight.)

PR 3 — Fused-add / residual backward + tests (depends on PR 2)

• build_fused_add_rmsnorm_bwd_module: same dx/dweight, plus dresidual and the grad-w.r.t.-residual_out (prenorm) path.
• Also needs the fwd fused-add to optionally save rstd (extend PR 1's approach to build_fused_add_rmsnorm_module).

PR 4 — Autograd wrapper (RMSNormFunction) (depends on PR 2, ideally PR 3)

• Thin torch.autograd.Function so gradients work transparently for quack-style calls (issue task 3).
• ⚠️ Design question: the repo has no torch-facing op layer today — we need to decide where this lives (e.g. a new python/flydsl/ops/ or a kernels/-level wrapper). Would like input here.

PR 5 — LayerNorm backward (stretch) (depends on PR 2's pattern; otherwise independent)

• Mirror the approach in layernorm_kernel.py, adding mean handling + dbias.

---

Open questions for reviewers

1. dweight reduction: fp32 atomicAdd (simple, 1 kernel, nondeterministic fp32 order) vs two-pass scratch-partials + finalizer kernel (deterministic, matches moe_reduction). I lean atomic for PR 2 and switch only if large-M contention shows up — objections?
2. Autograd layer location (PR 4) — where should the torch op wrapper live?
3. Anything in the quack interface we should match exactly for the [RFC]: a contribution path for FlyDSL kernels into downstream OSS libraries (PoC: Dao-AILab/quack) #749 integration to be drop-in?

@coderfeli — would appreciate a review of this staging and the two design calls above before I start on PR 1. Does the ordering / scope-per-PR look right to you?

[coderfeli]

@jhinpan 1. LGTM. 2. inside kernel is ok now? we can follow the style of quack? 3. we should try to align with them. For integration I also recommend 'Option 1 — Backend shim inside the downstream lib (recommended PoC).'

[jhinpan]

Update: experiment results + restructured to 3 PRs

Thanks @coderfeli! Addressing your points and reporting back on the one open design question we ran an experiment for.

On the 3 open questions

2. autograd layer → will keep it inside the kernel layer, quack-style (per your suggestion). No new op layer.
3. quack alignment → will align backward signatures / rstd conventions with quack; integration via Option 1 (backend shim in the downstream lib) as you recommended.
4. dweight cross-row reduction (atomic vs two-pass scratch) → we didn't want to guess, so we implemented both and benchmarked on MI355X. Details below.

dweight reduction experiment (MI355X, gfx950)

Both variants accumulate in fp32 and pass correctness vs torch.autograd.grad across f32/f16/bf16 (scratch is ~5× more accurate due to deterministic reduction, but both are well within tolerance).

Perf (µs, bf16), one block/row for atomic vs fixed 256-block scratch + finalizer:

M	N	atomic	scratch_total	winner
512	2048	23.6	29.5	atomic
4096	2048	142.8	61.5	scratch
32768	2048	866	383	scratch
32768	8192	923	1657	atomic
131072	8192	3304	6893	atomic

Findings: the tradeoff flips with N/M — scratch wins mid-M + small-N; atomic wins small-M and large-N. Notably MI355X fp32 global atomics are cheap enough that atomic never blows up (caps ~3.3ms even at M=131072), and it wins the large-N (4096–8192) regime — i.e. real LLM hidden sizes. Scratch's collapse at large M·N is mostly its fixed 256-block grid underutilizing the problem.

Decision: ship atomic as the default for PR 1 — simplest (no scratch buffer, no 2nd kernel), never catastrophic, and best in the regime that matters. We can add scratch later as an opt-in for determinism / mid-M-small-N (with a properly M-scaled grid) if a use case needs it.

Restructured PR plan (5 → 3 PRs)

Merging the mechanical steps into self-contained, usable units:

• PR 1 — RMSNorm training support: forward store_rstd + backward kernel (dx + dweight via fp32 atomic) + RMSNormFunction autograd wrapper (kernel-layer, quack-aligned) + gradient tests (f32/f16/bf16, both N paths incl. N≤2048). This alone makes rmsnorm drop-in trainable.
• PR 2 — Fused-add / residual backward: fused-add fwd saves rstd + residual backward (dx, dweight, dresidual) + wrapper + tests. Depends on PR 1's pattern.
• PR 3 — LayerNorm backward (stretch): mirror PR 1's approach + mean + dbias.

Starting on PR 1 now. Will self-review and run the full test/benchmark on device before opening it, and won't start PR 2/3 until PR 1 is validated.

[jhinpan]

Status: all 3 PRs open

The 3-PR series for this issue is up, each self-reviewed (multi-angle review + on-device verification on MI355X / gfx950) with fixes applied:

• PR 1 — [rmsnorm] Add backward pass + forward store_rstd for training (PR 1/3, #769) #795: rmsnorm backward (dx + fp32-atomic dweight) + forward store_rstd + RMSNormFunction autograd. Review fixes: eps honored (was silently ignored), multi-GPU cache fault fixed, get_llvm_ptr hoisted to kernels_common, shape/contiguity guards.
• PR 2 — [rmsnorm] Add fused-add / residual backward (PR 2/3, #769) #800: fused-add / residual (prenorm) backward (dx, dresidual, dgamma). Review fixes: dtype-mismatch guard, and dx==dresidual aliasing (dropped a redundant output buffer + double store).
• PR 3 — [layernorm] Add backward pass for training (PR 3/3, #769) #801: layernorm backward (dx, dgamma, dbias) + forward saves mean+rstd + LayerNormFunction. Review hardening: weight/bias length + fp32-stats asserts.

Stacking: #800 is stacked on #795, #801 on #800. Since base branches live in a fork, GitHub shows all three against main, so #800/#801 diffs currently include the earlier commits — merging #795 → #800 → #801 in order resolves that. Each PR's own change is a single commit.

Design decisions (from the earlier discussion + experiment): dweight/dgamma/dbias reductions use fp32 atomicAdd (benchmarked vs a scratch+finalizer variant — atomic wins the large-N / real-hidden-size regime and never blows up on MI355X). Autograd wrappers live in the kernel layer, quack-aligned, per @coderfeli's guidance.

Deferred (perf follow-ups, noted in code, not correctness): vectorized backward fast paths, pass-1/pass-2 load caching, reduction-helper dedup.

All backward passes verified against torch.autograd across bf16/fp16/fp32 including the small-N path. This unblocks the training / downstream (quack) integration in #749.

[jhinpan]

Status update — rmsnorm / layernorm backward series (as of Jul 8, 2026)

Implementing the plan agreed above, restructured to 3 dependent PRs (the autograd wrapper was folded into PR 1, and it lives inside the kernel layer, quack-style, per @coderfeli). All target main, meant to land in order. Developed and tested on MI355X (gfx950).

PR	Scope	State	CI	Merge
#795 [1/3]	rmsnorm backward (dx, fp32 dweight) + forward store_rstd + RMSNormFunction autograd wrapper + grad tests (f32/f16/bf16, fast + small-N ≤2048)	🟢 Open — heavily iterated (~10 review rounds)	✅ green (multi-GPU skipped)	MERGEABLE, awaiting final approving review
#800 [2/3]	Fused-add / residual (prenorm) backward — dx, dweight, dresidual + grad w.r.t. residual_out; backward kernels split into rmsnorm_bwd_kernel.py	🟢 Open — @coderfeli reviewing	pending (recomputing after latest push)	mergeability recomputing
#801 [3/3]	layernorm backward (adds mean handling + dbias)	🟠 Open — earliest stage	not yet run	needs rebase — conflicts with main

Design calls (settled)

• dweight reduction → fp32 atomicAdd. Confirmed by the MI355X benchmark above: atomic never blows up and wins the large-N (4k–8k) regime — i.e. real LLM hidden sizes. (Wrapped in the shared kernels_common.atomic_add helper per review.)
• Autograd wrapper stays in the kernel layer, quack-style (no new op layer).
• Signatures aligned to quack so the [RFC]: a contribution path for FlyDSL kernels into downstream OSS libraries (PoC: Dao-AILab/quack) #749 downstream integration is a drop-in backend shim.

Next actions

• [rmsnorm] Add backward pass + forward store_rstd for training (PR 1/3, #769) #795 is ready pending one approving review (CI green, review comments addressed).
• [rmsnorm] Add fused-add / residual backward (PR 2/3, #769) #800 — waiting on CI to finish + @coderfeli's re-review after the rmsnorm_bwd_kernel.py split.
• Rebase [layernorm] Add backward pass for training (PR 3/3, #769) #801 onto current main to clear conflicts; least-mature of the three.