[rmsnorm] Add backward pass + forward store_rstd for training (PR 1/3, #769) by jhinpan · Pull Request #795 · ROCm/FlyDSL

jhinpan · 2026-07-03T04:49:59Z

Summary

PR 1 of 3 for #769 — makes RMSNorm training-capable. This PR adds the core rmsnorm backward pass, the forward rstd output it needs, and a quack-aligned autograd wrapper.

Follow-ups (tracked in #769): PR 2 = fused-add/residual backward; PR 3 = layernorm backward (stretch).

What's included

Forward saves rstd: build_rmsnorm_module(N, dtype, store_rstd=False). When enabled, writes per-row rstd (1/RMS, fp32, (M,)) needed by backward. Default off ⇒ existing launcher signature and all callers byte-for-byte unchanged. Covers fast, generic, and small-N (N≤2048) paths.
Backward kernel: build_rmsnorm_bwd_module — fused single kernel, one block per row. Pass 1 computes c1 = mean_N(x_hat·wdy); pass 2 stores dx = (wdy − x_hat·c1)·rstd and atomicAdds dw = dy·x_hat into an fp32 DWeight[N]. All reduction + weight-grad accumulation in fp32; only dx cast back to I/O dtype.
Autograd wrapper: RMSNormFunction + public rmsnorm(x, weight, eps), quack-aligned, in the kernel layer (per review discussion). Math matches quack rmsnorm_bwd_ref.

Design note: dweight reduction = fp32 atomicAdd

The cross-row dweight = Σ_rows(dy·x_hat) reduction was chosen by experiment, not assumption. We implemented both fp32 atomicAdd and a two-pass scratch+finalizer variant and benchmarked on MI355X (full data in #769):

atomic never blows up (caps ~3.3ms even at M=131072) and wins the large-N regime (4096–8192, i.e. real LLM hidden sizes);
scratch wins mid-M/small-N but collapses at large M·N with a fixed grid.

Shipping atomic as the default (simplest — no scratch buffer, no 2nd kernel). Scratch can be added later as an opt-in for determinism if needed.

Testing (MI355X, gfx950)

Forward regression: unchanged, all pass.
test_rmsnorm_backward: dx + dweight + kernel-rstd vs torch.autograd.grad, across f32/f16/bf16, both N paths (incl. unaligned 64×2000 and N≤2048).
test_rmsnorm_autograd: end-to-end public rmsnorm() path — grads on x + weight, and batched (>2D) reshape.
Full non-large suite: 8/8 passed. Verified dweight has no accumulation leak (stable across repeated calls; only atomic-order fp32 noise ~1e-5 f32).

Fixes from the PR ROCm#795 code review: - eps is now baked into the forward kernel (build_rmsnorm_module / small-N builder gain an `eps` param) instead of being silently ignored in favor of the module EPS=1e-5. rmsnorm(x, w, eps=1e-6) now produces correct numerics; eps is part of the fwd compile-cache key. - Multi-GPU correctness: compiled-fn cache keys now include x.device, and compile+launch run under `with torch.cuda.device(x.device)` so a kernel built on cuda:0 is never launched on cuda:1 (previously faulted with hipErrorInvalidDevice + memory access fault). - Deduplicate the LLVM-ptr helper: hoist get_llvm_ptr into kernels_common.py (was a byte-for-byte copy of hgemm_splitk's helper, which CLAUDE.md says belongs in kernels_common) and use it. - Public rmsnorm() now asserts x.shape[-1] == weight length; rmsnorm_fwd/ rmsnorm_bwd assert contiguous inputs (make_buffer_tensor assumes row-major contiguous). - Document the backward kernel's deferred perf follow-ups (vectorized fast path + pass-1/pass-2 caching) instead of leaving them implicit. - Tests: add test_rmsnorm_eps_honored (eps must change output / match a torch ref at 1e-5/1e-6/1e-2) and test_rmsnorm_multi_gpu (marked multi_gpu; runs rmsnorm fwd+bwd on cuda:0 and cuda:1). Deferred (perf/altitude, not correctness; noted in code): vectorized backward fast path, pass-2 load caching, and the store_rstd launcher duplication (kept to preserve the byte-for-byte store_rstd=False path). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

jhinpan · 2026-07-03T21:52:42Z

Ran a self-review pass and pushed fixes in 0c5f94c:

eps was silently ignored (kernel hardcoded 1e-5) → now baked into the forward kernel and part of the compile-cache key. rmsnorm(x, w, eps=1e-6) now matches a torch ref (err ~2e-7). (verified on MI355X)
Multi-GPU fault: compiled-fn cache keys omitted device, so a cuda:0 kernel reused on cuda:1 crashed with hipErrorInvalidDevice + memory access fault → cache key now includes x.device and compile/launch run under torch.cuda.device(x.device). (verified cuda:0 → cuda:1)
Dedup: hoisted get_llvm_ptr into kernels_common.py (was a byte-for-byte copy of hgemm_splitk's; CLAUDE.md says it belongs there).
Guards: public rmsnorm() now asserts x.shape[-1] == weight length; rmsnorm_fwd/bwd assert contiguous inputs.
Tests: added test_rmsnorm_eps_honored + test_rmsnorm_multi_gpu (marked multi_gpu). Full non-large suite: 10/10 pass.

Deferred as noted in code (perf/altitude, not correctness): vectorized backward fast path, pass-2 load caching, and the store_rstd launcher duplication (kept to preserve the byte-for-byte store_rstd=False path). These are follow-ups.

Proceeding to PR 2 (fused-add/residual backward) and PR 3 (layernorm backward).

coderfeli · 2026-07-07T11:55:55Z

+
+                dw = dy * x_hat
+                ptr = get_llvm_ptr(DWeight, idx, 4)
+                llvm.AtomicRMWOp(


need llvm.atomic? wrap it into mem ops .py?

Good call — dropped the raw llvm.AtomicRMWOp from the kernel and wrapped it in a shared atomic_add() helper in kernels_common.py, right next to its companion get_llvm_ptr (per CLAUDE.md, LLVM-ptr / atomic kernel helpers live there — there's no separate mem_ops.py, and expr/ is target-neutral). The backward now just calls:

atomic_add(DWeight, idx, dw, dtype_bytes=4)

It builds the global !llvm.ptr via get_llvm_ptr and emits llvm.atomicrmw, auto-selecting fadd for a float operand / integer add otherwise and returning the old value, so the other inline atomicrmw sites (hgemm split-K, etc.) can adopt it in a follow-up. Generated IR is unchanged — llvm.atomicrmw fadd %ptr, %v syncscope("agent") monotonic {alignment = 4} — verified with a full COMPILE_ONLY lowering on gfx942 for f32/f16/bf16.

Also merged latest main to clear the conflict: the kernels/ subpackage reorg moved this file to kernels/norm/rmsnorm_kernel.py and the helper module to kernels/common/kernels_common.py.

Done — the inline llvm.atomicrmw is now wrapped in a shared atomic_add() helper in kernels/common/kernels_common.py (commit 0680d0d). The kernel just calls atomic_add(DWeight, idx, dw, dtype_bytes=4); the helper picks fadd vs integer add from the operand type and owns the get_llvm_ptr + llvm.atomicrmw pair.

coderfeli · 2026-07-07T12:00:34Z

@jhinpan conflict

…OCm#795) Address coderfeli's review: the backward kernel emitted a raw llvm.AtomicRMWOp inline. Hoist that pattern into a shared atomic_add() helper in kernels_common.py, next to its companion get_llvm_ptr (per CLAUDE.md, LLVM-ptr / atomic kernel helpers live in kernels_common). - atomic_add(dst, offset, value, *, dtype_bytes=4, syncscope="agent", ...) builds the global !llvm.ptr via get_llvm_ptr and emits llvm.atomicrmw, picking fadd for a float operand and integer add otherwise, and returns the previous value. Emits the same op the kernel had inline: llvm.atomicrmw fadd %ptr, %v syncscope("agent") monotonic {alignment = 4} - rmsnorm backward now calls atomic_add(DWeight, idx, dw, dtype_bytes=4) and drops the direct llvm-dialect import and get_llvm_ptr usage. The helper is intentionally general so the duplicated raw atomicrmw sites (hgemm_splitk / split-K GEMMs) can adopt it in a follow-up. Verified: byte-identical generated IR and full COMPILE_ONLY lowering (gfx942) for the backward kernel across f32/f16/bf16. Co-authored-by: Cursor <cursoragent@cursor.com>

coderfeli · 2026-07-08T04:01:44Z

 from flydsl.runtime.device import get_rocm_arch, is_rdna_arch


+def get_llvm_ptr(ptr, offset, dtype_bytes, ptr_type=None):


this still here? can't reuse current？

Ah I see. Sorry about that. I will try to fix them all now.

Good catch. It's now reused: I removed the duplicated local get_llvm_ptr copies from gemm/hgemm_splitk.py, gemm/small_m_hgemm.py and gemm/splitk_hgemm.py and pointed all of them at the single shared kernels.common.kernels_common.get_llvm_ptr (commit b4dd6b5). All three local copies used the address-space-1 !llvm.ptr<1> / i64 form that the shared helper defaults to, so behavior is unchanged and the net diff is -23 lines.

coderfeli · 2026-07-08T04:08:36Z

 from flydsl.runtime.device import get_rocm_arch, is_rdna_arch


+def get_llvm_ptr(ptr, offset, dtype_bytes, ptr_type=None):


Dedup is incomplete: this adds a shared get_llvm_ptr/atomic_add but the byte-for-byte copies in gemm/hgemm_splitk.py:285, gemm/small_m_hgemm.py:623, and gemm/splitk_hgemm.py:271 are left in place — the count went up, not down. Either migrate those call sites here, or reword the commit message (it claims it dedups hgemm's helper).

Fixed — the dedup is now complete. I migrated the byte-for-byte get_llvm_ptr copies in gemm/hgemm_splitk.py:285, gemm/small_m_hgemm.py:623 and gemm/splitk_hgemm.py:271 to the shared kernels.common.kernels_common.get_llvm_ptr and deleted the local defs (commit b4dd6b5). The count now goes down: 3 files changed, +3/-26 lines. All three used the same !llvm.ptr<1> / i64 form the shared helper defaults to, so no behavior change.

coderfeli · 2026-07-08T04:08:36Z

        Input: fx.Tensor,
        Gamma: fx.Tensor,
-        _Unused: fx.Tensor,
+        Rstd: fx.Tensor,


After the _Unused→Rstd rename, the store_rstd=False launchers pass Gamma into this Rstd slot (rmsnorm_kernel(Input, Gamma, Gamma, Output) ~L339, same for the small-N launcher ~L465). Harmless (the slot is never dereferenced when store_rstd=False) but now reads as a bug — add a one-line comment noting it's an unused placeholder.

Added the placeholder comment at both call sites (commit f38ced6). The store_rstd=False launchers for rmsnorm_kernel (~L339) and rmsnorm_large_m_small_n_kernel (~L468) now carry a one-line note that the Rstd slot is an unused placeholder and Gamma is passed only to fill the argument (never dereferenced in-kernel).

Makes RMSNorm training-capable (PR 1 of 3 for ROCm#769): - build_rmsnorm_module gains store_rstd=False: when enabled, writes per-row rstd (1/RMS, fp32, shape (M,)) needed by backward. Default off keeps the existing launcher signature and all callers byte-for-byte unchanged; covers fast, generic, and small-N (N<=2048) paths. - build_rmsnorm_bwd_module: fused single kernel, one block per row. Pass 1 computes c1 = mean_N(x_hat*wdy); pass 2 stores dx = (wdy - x_hat*c1)*rstd and atomicAdds dw = dy*x_hat into an fp32 DWeight[N]. All reduction and weight-grad accumulation in fp32; only dx cast back to I/O dtype. - fp32 atomicAdd chosen for the cross-row dweight reduction after benchmarking it against a two-pass scratch+finalizer variant on MI355X: atomic never blows up and wins the large-N (real LLM hidden size) regime. - RMSNormFunction (torch.autograd.Function) + public rmsnorm(x, weight, eps), quack-aligned, in the kernel layer. Math matches quack rmsnorm_bwd_ref. - Tests: gradient checks vs torch.autograd across f32/f16/bf16, both N paths (incl. unaligned and N<=2048), plus an end-to-end rmsnorm() autograd test covering batched (>2D) reshape and grads on x + weight. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Fixes from the PR ROCm#795 code review: - eps is now baked into the forward kernel (build_rmsnorm_module / small-N builder gain an `eps` param) instead of being silently ignored in favor of the module EPS=1e-5. rmsnorm(x, w, eps=1e-6) now produces correct numerics; eps is part of the fwd compile-cache key. - Multi-GPU correctness: compiled-fn cache keys now include x.device, and compile+launch run under `with torch.cuda.device(x.device)` so a kernel built on cuda:0 is never launched on cuda:1 (previously faulted with hipErrorInvalidDevice + memory access fault). - Deduplicate the LLVM-ptr helper: hoist get_llvm_ptr into kernels_common.py (was a byte-for-byte copy of hgemm_splitk's helper, which CLAUDE.md says belongs in kernels_common) and use it. - Public rmsnorm() now asserts x.shape[-1] == weight length; rmsnorm_fwd/ rmsnorm_bwd assert contiguous inputs (make_buffer_tensor assumes row-major contiguous). - Document the backward kernel's deferred perf follow-ups (vectorized fast path + pass-1/pass-2 caching) instead of leaving them implicit. - Tests: add test_rmsnorm_eps_honored (eps must change output / match a torch ref at 1e-5/1e-6/1e-2) and test_rmsnorm_multi_gpu (marked multi_gpu; runs rmsnorm fwd+bwd on cuda:0 and cuda:1). Deferred (perf/altitude, not correctness; noted in code): vectorized backward fast path, pass-2 load caching, and the store_rstd launcher duplication (kept to preserve the byte-for-byte store_rstd=False path). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Resolve the conflict from the kernels/ directory reorg landed on main, which moved the flat kernel modules into topical subpackages: kernels/kernels_common.py -> kernels/common/kernels_common.py kernels/rmsnorm_kernel.py -> kernels/norm/rmsnorm_kernel.py The only content conflict was the shared import line in the rmsnorm kernel. Resolved to the new kernels.common.kernels_common path while keeping this PR's get_llvm_ptr import. All PR additions are preserved (rmsnorm backward, forward store_rstd, autograd wrapper, tests), and kernels/common/kernels_common.py keeps both this PR's get_llvm_ptr and upstream's new _if_else helper. Test import updated by upstream to kernels.norm.rmsnorm_kernel. Co-authored-by: Cursor <cursoragent@cursor.com>

…OCm#795) Address coderfeli's review: the backward kernel emitted a raw llvm.AtomicRMWOp inline. Hoist that pattern into a shared atomic_add() helper in kernels_common.py, next to its companion get_llvm_ptr (per CLAUDE.md, LLVM-ptr / atomic kernel helpers live in kernels_common). - atomic_add(dst, offset, value, *, dtype_bytes=4, syncscope="agent", ...) builds the global !llvm.ptr via get_llvm_ptr and emits llvm.atomicrmw, picking fadd for a float operand and integer add otherwise, and returns the previous value. Emits the same op the kernel had inline: llvm.atomicrmw fadd %ptr, %v syncscope("agent") monotonic {alignment = 4} - rmsnorm backward now calls atomic_add(DWeight, idx, dw, dtype_bytes=4) and drops the direct llvm-dialect import and get_llvm_ptr usage. The helper is intentionally general so the duplicated raw atomicrmw sites (hgemm_splitk / split-K GEMMs) can adopt it in a follow-up. Verified: byte-identical generated IR and full COMPILE_ONLY lowering (gfx942) for the backward kernel across f32/f16/bf16. Co-authored-by: Cursor <cursoragent@cursor.com>

…unchers The _Unused->Rstd rename left the store_rstd=False launchers passing Gamma into the Rstd slot (rmsnorm_kernel and rmsnorm_large_m_small_n_kernel). The slot is never dereferenced when store_rstd=False, but it now reads as a bug. Add a one-line comment at each call site marking it an unused placeholder. Addresses review comment ROCm#795 (r3541145967). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Complete the dedup: the byte-for-byte get_llvm_ptr copies in gemm/hgemm_splitk.py, gemm/small_m_hgemm.py and gemm/splitk_hgemm.py are removed and replaced with an import of the shared kernels.common.kernels_common.get_llvm_ptr. All three local copies used the address-space-1 !llvm.ptr<1> / i64 form that the shared helper defaults to, so behavior is unchanged. Net -23 lines. Addresses review comment ROCm#795 (r3541145963). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Copilot AI review requested due to automatic review settings July 3, 2026 04:50

Copilot AI reviewed Jul 3, 2026

This was referenced Jul 3, 2026

[rmsnorm] Add fused-add / residual backward (PR 2/3, #769) #800

Open

[layernorm] Add backward pass for training (PR 3/3, #769) #801

Open

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

Open

coderfeli reviewed Jul 7, 2026

View reviewed changes

coderfeli reviewed Jul 8, 2026

View reviewed changes

jhinpan and others added 4 commits July 8, 2026 05:10

jhinpan force-pushed the feat/rmsnorm-backward-769 branch from 5029619 to 0680d0d Compare July 8, 2026 05:11

jhinpan and others added 2 commits July 8, 2026 05:12

		from flydsl.runtime.device import get_rocm_arch, is_rdna_arch


		def get_llvm_ptr(ptr, offset, dtype_bytes, ptr_type=None):

Uh oh!

Conversation

jhinpan commented Jul 3, 2026

Summary

What's included

Design note: dweight reduction = fp32 atomicAdd

Testing (MI355X, gfx950)

Related

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

jhinpan commented Jul 3, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

coderfeli commented Jul 7, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants