From 2d8cd30d2e67869ca2914ad5eca6d1fe9dd23945 Mon Sep 17 00:00:00 2001 From: root Date: Fri, 3 Jul 2026 04:49:35 +0000 Subject: [PATCH 1/2] [rmsnorm] Add backward pass + forward store_rstd for training (#769) Makes RMSNorm training-capable (PR 1 of 3 for #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 --- kernels/rmsnorm_kernel.py | 333 +++++++++++++++++++++++++++++++++- tests/kernels/test_rmsnorm.py | 179 ++++++++++++++++++ 2 files changed, 507 insertions(+), 5 deletions(-) diff --git a/kernels/rmsnorm_kernel.py b/kernels/rmsnorm_kernel.py index 17af22fe6..4fa81afc7 100644 --- a/kernels/rmsnorm_kernel.py +++ b/kernels/rmsnorm_kernel.py @@ -14,12 +14,20 @@ import flydsl.compiler as flyc import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import fly, llvm from flydsl.expr import arith, const_expr, gpu, range_constexpr from flydsl.expr import math as fmath +from flydsl.expr.typing import T from flydsl.expr.vector import ReductionOp, full from flydsl.runtime.device import get_rocm_arch as get_hip_arch from kernels.kernels_common import dtype_to_elem_type, get_warp_size +try: + import torch +except ImportError: + torch = None + KERNEL_NAME = "rmsnorm" EPS = 1e-5 @@ -53,6 +61,18 @@ def _store_scalar(copy_atom, elem_dtype, store_dtype, divided_tensor, index, val fx.copy_atom_call(copy_atom, r, view) +def _get_llvm_ptr(ptr, offset, dtype_bytes, ptr_type=None): + if ptr_type is None: + ptr_type = ir.Type.parse("!llvm.ptr<1>") + base_ptr = fly.extract_aligned_pointer_as_index(ptr_type, ptr) + base_ptr = llvm.PtrToIntOp(T.i64, base_ptr).result + byte_offset = arith.index_cast(T.i64, fx.Index(offset) * fx.Index(dtype_bytes)) + llvm_ptr = llvm.AddOp(base_ptr, byte_offset, llvm.IntegerOverflowFlags(0)).result + llvm_ptr = llvm.IntToPtrOp(ptr_type, llvm_ptr).result + ptr_v = llvm_ptr._value if const_expr(hasattr(llvm_ptr, "_value")) else llvm_ptr + return ptr_v + + def _load_vec(copy_atom, vec_width, elem_dtype, div_tensor, idx): r = fx.make_rmem_tensor(vec_width, elem_dtype) fx.copy_atom_call(copy_atom, fx.slice(div_tensor, (None, idx)), r) @@ -110,9 +130,9 @@ def _quant_dtype_max(dtype_str: str) -> float: raise ValueError(f"unsupported quant dtype: {dtype_str!r} (expected 'i8' or 'int8')") -def build_rmsnorm_module(N: int, dtype_str: str): +def build_rmsnorm_module(N: int, dtype_str: str, store_rstd: bool = False): if N <= 2048: - return _build_rmsnorm_large_m_small_n_module(N, dtype_str) + return _build_rmsnorm_large_m_small_n_module(N, dtype_str, store_rstd) arch = get_hip_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") @@ -127,7 +147,7 @@ def build_rmsnorm_module(N: int, dtype_str: str): def rmsnorm_kernel( Input: fx.Tensor, Gamma: fx.Tensor, - _Unused: fx.Tensor, + Rstd: fx.Tensor, Output: fx.Tensor, ): bid = fx.block_idx.x @@ -142,6 +162,11 @@ def rmsnorm_kernel( s_red = lds.s_red.view(fx.make_layout(RED_SLOTS, 1)) s_red2 = lds.s_red2.view(fx.make_layout(RED_SLOTS, 1)) + if const_expr(store_rstd): + Rstd_buf = fx.rocdl.make_buffer_tensor(Rstd) + rstd_div = fx.logical_divide(Rstd_buf, fx.make_layout(1, 1)) + rstd_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 32) + def wave_reduce_add(x): w = x for _sh_exp in range_constexpr(int(math.log2(WARP_SIZE))): @@ -227,6 +252,10 @@ def block_reduce_add2(val0, val1): ms_eps = mean_sq + eps_c rrms = fmath.rsqrt(ms_eps, fastmath=fm_fast) + if const_expr(store_rstd): + if tid == 0: + _store_scalar(rstd_copy_atom, fx.Float32, fx.Float32, rstd_div, bid, rrms) + # Pass 2: normalize + gamma + store (reuse cached input) for tile_i in range_constexpr(num_tiles): idx = tid + tile_i * BLOCK_THREADS @@ -278,6 +307,10 @@ def block_reduce_add2(val0, val1): ms_eps = mean_sq + eps_c rrms = fmath.rsqrt(ms_eps, fastmath=fm_fast) + if const_expr(store_rstd): + if tid == 0: + _store_scalar(rstd_copy_atom, fx.Float32, fx.Float32, rstd_div, bid, rrms) + for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): idx = tid + base_idx_int if idx < N: @@ -290,6 +323,26 @@ def block_reduce_add2(val0, val1): y_e = _to_elem_scalar(dtype_str, elem_dtype, y) _store_scalar(copy_atom_s, elem_dtype, elem_dtype, out_div, idx, y_e) + if store_rstd: + + @flyc.jit + def launch_rmsnorm( + Input: fx.Tensor, + Gamma: fx.Tensor, + Output: fx.Tensor, + Rstd: fx.Tensor, + m_in: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + launcher = rmsnorm_kernel(Input, Gamma, Rstd, Output) + launcher.launch( + grid=(m_in, 1, 1), + block=(BLOCK_THREADS, 1, 1), + stream=stream, + ) + + return launch_rmsnorm + @flyc.jit def launch_rmsnorm( Input: fx.Tensor, @@ -308,7 +361,7 @@ def launch_rmsnorm( return launch_rmsnorm -def _build_rmsnorm_large_m_small_n_module(N: int, dtype_str: str): +def _build_rmsnorm_large_m_small_n_module(N: int, dtype_str: str, store_rstd: bool = False): BLOCK_N = 1 << (N - 1).bit_length() BLOCK_M = max(min(16384 // BLOCK_N, 32), 8) THREADS_PER_ROW = min(WARP_SIZE, 1024 // BLOCK_M) @@ -319,7 +372,7 @@ def _build_rmsnorm_large_m_small_n_module(N: int, dtype_str: str): def rmsnorm_large_m_small_n_kernel( Input: fx.Tensor, Gamma: fx.Tensor, - _Unused: fx.Tensor, + Rstd: fx.Tensor, Output: fx.Tensor, MIn: fx.Int32, ): @@ -339,6 +392,10 @@ def rmsnorm_large_m_small_n_kernel( Input_buf = fx.rocdl.make_buffer_tensor(Input) Gamma_buf = fx.rocdl.make_buffer_tensor(Gamma) Output_buf = fx.rocdl.make_buffer_tensor(Output) + if const_expr(store_rstd): + Rstd_buf = fx.rocdl.make_buffer_tensor(Rstd) + rstd_div = fx.logical_divide(Rstd_buf, fx.make_layout(1, 1)) + rstd_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 32) row_in = fx.slice(Input_buf, (row, None)) row_out = fx.slice(Output_buf, (row, None)) @@ -377,6 +434,10 @@ def group_reduce_add(x): ms_eps = mean_sq + eps_c rrms = fmath.rsqrt(ms_eps, fastmath=fm_fast) + if const_expr(store_rstd): + if lane == 0: + _store_scalar(rstd_copy_atom, fx.Float32, fx.Float32, rstd_div, row, rrms) + for base_idx_int in range_constexpr(0, BLOCK_N, THREADS_PER_ROW): idx = lane + base_idx_int if idx < N: @@ -388,6 +449,26 @@ def group_reduce_add(x): y_e = _to_elem_scalar(dtype_str, elem_dtype, y) _store_scalar(copy_atom_s, elem_dtype, elem_dtype, out_div, idx, y_e) + if store_rstd: + + @flyc.jit + def launch_rmsnorm_large_m_small_n( + Input: fx.Tensor, + Gamma: fx.Tensor, + Output: fx.Tensor, + Rstd: fx.Tensor, + m_in: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + launcher = rmsnorm_large_m_small_n_kernel(Input, Gamma, Rstd, Output, m_in) + launcher.launch( + grid=((m_in + fx.Int32(BLOCK_M - 1)) // fx.Int32(BLOCK_M), 1, 1), + block=(BLOCK_THREADS_SPECIAL, 1, 1), + stream=stream, + ) + + return launch_rmsnorm_large_m_small_n + @flyc.jit def launch_rmsnorm_large_m_small_n( Input: fx.Tensor, @@ -406,6 +487,154 @@ def launch_rmsnorm_large_m_small_n( return launch_rmsnorm_large_m_small_n +def build_rmsnorm_bwd_module(N: int, dtype_str: str): + """Fused RMSNorm backward: grid=(M,), one block per row. + + Pass 1: c1 = mean_N(x_hat * wdy), x_hat = x*rstd, wdy = dy*gamma. + Pass 2: dx = (wdy - x_hat*c1) * rstd -> DX (elem dtype); + dw_elem = dy * x_hat (fp32) -> atomicAdd into DWeight[idx] (fp32). + Uses the module EPS constant implicitly via Rstd (supplied by the forward). + """ + RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) + elem_bits = 32 if dtype_str == "f32" else 16 + SharedStorage = _make_reduction_storage(RED_SLOTS) + + @flyc.kernel + def rmsnorm_bwd_kernel( + Input: fx.Tensor, + Gamma: fx.Tensor, + DY: fx.Tensor, + Rstd: fx.Tensor, + DX: fx.Tensor, + DWeight: fx.Tensor, + ): + bid = fx.block_idx.x + tid = fx.thread_idx.x + + elem_dtype = dtype_to_elem_type(dtype_str) + fm_fast = arith.FastMathFlags.fast + n_float = float(N) + c_zero_f = fx.Float32(0.0) + + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + s_red = lds.s_red.view(fx.make_layout(RED_SLOTS, 1)) + + def wave_reduce_add(x): + w = x + for _sh_exp in range_constexpr(int(math.log2(WARP_SIZE))): + off = WARP_SIZE // (2 << _sh_exp) + peer = w.shuffle_xor(off, WARP_SIZE) + w = w.addf(peer, fastmath=fm_fast) + return w + + def block_reduce_add(val): + if const_expr(RED_SLOTS == 1): + return wave_reduce_add(val) + lane = tid % WARP_SIZE + wave = tid // WARP_SIZE + w = wave_reduce_add(val) + if lane == 0: + fx.memref_store(w, s_red, wave) + gpu.barrier() + if wave == 0: + in_range = lane < RED_SLOTS + lane_safe = in_range.select(lane, 0) + v = fx.memref_load(s_red, lane_safe) + ww = in_range.select(v, c_zero_f) + ww = wave_reduce_add(ww) + if lane == 0: + fx.memref_store(ww, s_red, 0) + gpu.barrier() + return fx.memref_load(s_red, 0) + + Input_buf = fx.rocdl.make_buffer_tensor(Input) + Gamma_buf = fx.rocdl.make_buffer_tensor(Gamma) + DY_buf = fx.rocdl.make_buffer_tensor(DY) + Rstd_buf = fx.rocdl.make_buffer_tensor(Rstd) + DX_buf = fx.rocdl.make_buffer_tensor(DX) + + row_in = fx.slice(Input_buf, (bid, None)) + row_dy = fx.slice(DY_buf, (bid, None)) + row_dx = fx.slice(DX_buf, (bid, None)) + + copy_atom_s = fx.make_copy_atom( + fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), + elem_bits, + ) + copy_atom_f32 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 32) + + row_div = fx.logical_divide(row_in, fx.make_layout(1, 1)) + dy_div = fx.logical_divide(row_dy, fx.make_layout(1, 1)) + gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(1, 1)) + dx_div = fx.logical_divide(row_dx, fx.make_layout(1, 1)) + rstd_div = fx.logical_divide(Rstd_buf, fx.make_layout(1, 1)) + + rstd = _load_scalar(copy_atom_f32, fx.Float32, rstd_div, bid) + + # Pass 1: c1 = mean( x_hat * wdy ) = mean( (x*rstd) * (dy*gamma) ) + thread_acc = c_zero_f + for base in range_constexpr(0, N, BLOCK_THREADS): + idx = tid + base + is_valid = idx < N + idx_safe = is_valid.select(idx, 0) + x_e = _load_scalar(copy_atom_s, elem_dtype, row_div, idx_safe) + dy_e = _load_scalar(copy_atom_s, elem_dtype, dy_div, idx_safe) + g_e = _load_scalar(copy_atom_s, elem_dtype, gamma_div, idx_safe) + x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) + dy = dy_e if dtype_str == "f32" else dy_e.to(fx.Float32) + g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) + x_hat = x * rstd + wdy = dy * g + prod = x_hat * wdy + thread_acc = thread_acc + is_valid.select(prod, c_zero_f) + + sum_prod = block_reduce_add(thread_acc) + c1 = sum_prod / n_float + + # Pass 2: dx = (wdy - x_hat*c1) * rstd ; dw = dy * x_hat (atomicAdd fp32) + for base in range_constexpr(0, N, BLOCK_THREADS): + idx = tid + base + if idx < N: + x_e = _load_scalar(copy_atom_s, elem_dtype, row_div, idx) + dy_e = _load_scalar(copy_atom_s, elem_dtype, dy_div, idx) + g_e = _load_scalar(copy_atom_s, elem_dtype, gamma_div, idx) + x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) + dy = dy_e if dtype_str == "f32" else dy_e.to(fx.Float32) + g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) + x_hat = x * rstd + wdy = dy * g + dx = (wdy - x_hat * c1) * rstd + dx_e = dx if dtype_str == "f32" else dx.to(elem_dtype) + _store_scalar(copy_atom_s, elem_dtype, elem_dtype, dx_div, idx, dx_e) + + dw = dy * x_hat + ptr = _get_llvm_ptr(DWeight, idx, 4) + llvm.AtomicRMWOp( + llvm.AtomicBinOp.fadd, + ptr, + dw.ir_value(), + llvm.AtomicOrdering.monotonic, + syncscope="agent", + alignment=4, + ) + + @flyc.jit + def launch_rmsnorm_bwd( + Input: fx.Tensor, + Gamma: fx.Tensor, + DY: fx.Tensor, + Rstd: fx.Tensor, + DX: fx.Tensor, + DWeight: fx.Tensor, + m_in: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + launcher = rmsnorm_bwd_kernel(Input, Gamma, DY, Rstd, DX, DWeight) + launcher.launch(grid=(m_in, 1, 1), block=(BLOCK_THREADS, 1, 1), stream=stream) + + return launch_rmsnorm_bwd + + def build_fused_add_rmsnorm_module(N: int, dtype_str: str): arch = get_hip_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") @@ -1386,3 +1615,97 @@ def build_fused_add_rmsnorm_smoothquant_module( is_smooth=True, quant_dtype_str=quant_dtype_str, ) + + +# ===================================================================== +# Python wrappers + autograd (quack-aligned). PR 1: plain rmsnorm. +# ===================================================================== +if torch is not None: + + def _torch_dtype_to_str(dt) -> str: + if dt == torch.float32: + return "f32" + if dt == torch.float16: + return "f16" + if dt == torch.bfloat16: + return "bf16" + raise ValueError(f"unsupported torch dtype: {dt}") + + # Compiled-fn caches keyed by (N, dtype_str, store_rstd) / (N, dtype_str). + _FWD_CACHE: dict = {} + _BWD_CACHE: dict = {} + + def _get_fwd_compiled(x, weight, out, rstd, M, N, dtype_str, store_rstd, stream): + key = (N, dtype_str, store_rstd) + entry = _FWD_CACHE.get(key) + if entry is None: + launch_fn = build_rmsnorm_module(N, dtype_str, store_rstd=store_rstd) + if store_rstd: + compiled = flyc.compile(launch_fn, x, weight, out, rstd, M, stream) + else: + compiled = flyc.compile(launch_fn, x, weight, out, M, stream) + _FWD_CACHE[key] = compiled + entry = compiled + return entry + + def rmsnorm_fwd(x, weight, eps=EPS, store_rstd=False): + """Forward RMSNorm. Returns (out, rstd). + + NOTE: the compiled kernel uses the module EPS constant (not `eps`); + `eps` is accepted for interface parity only (see module docstring). + """ + assert x.dim() == 2, "rmsnorm_fwd expects a 2D (M, N) input" + M, N = x.shape + out = torch.empty_like(x) + rstd = torch.empty((M,), device=x.device, dtype=torch.float32) if store_rstd else None + dtype_str = _torch_dtype_to_str(x.dtype) + stream = torch.cuda.current_stream() + compiled = _get_fwd_compiled(x, weight, out, rstd, M, N, dtype_str, store_rstd, stream) + if store_rstd: + compiled(x, weight, out, rstd, M, stream) + else: + compiled(x, weight, out, M, stream) + return out, rstd + + def rmsnorm_bwd(x, weight, dout, rstd, eps=EPS): + """Backward RMSNorm. Returns (dx, dw) with dw cast to weight dtype.""" + assert x.dim() == 2, "rmsnorm_bwd expects a 2D (M, N) input" + M, N = x.shape + dtype_str = _torch_dtype_to_str(x.dtype) + dx = torch.empty_like(x) + dweight = torch.zeros((N,), device=x.device, dtype=torch.float32) + stream = torch.cuda.current_stream() + key = (N, dtype_str) + compiled = _BWD_CACHE.get(key) + if compiled is None: + launch_fn = build_rmsnorm_bwd_module(N, dtype_str) + # flyc.compile executes the kernel once during tracing, which would + # accumulate into DWeight; zero it AFTER compiling. + compiled = flyc.compile(launch_fn, x, weight, dout, rstd, dx, dweight, M, stream) + _BWD_CACHE[key] = compiled + dweight.zero_() + compiled(x, weight, dout, rstd, dx, dweight, M, stream) + return dx, dweight.to(weight.dtype) + + class RMSNormFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, x, weight, eps): + need_grad = x.requires_grad or weight.requires_grad + out, rstd = rmsnorm_fwd(x, weight, eps=eps, store_rstd=need_grad) + ctx.save_for_backward(x, weight, rstd) + ctx.eps = eps + return out + + @staticmethod + def backward(ctx, dout): + x, weight, rstd = ctx.saved_tensors + dx, dw = rmsnorm_bwd(x, weight, dout.contiguous(), rstd, eps=ctx.eps) + return dx, dw, None + + def rmsnorm(x, weight=None, eps=EPS): + """Public entry: plain RMSNorm with autograd (weight required in PR 1).""" + assert weight is not None, "PR 1 rmsnorm requires an explicit weight" + N = weight.shape[-1] + x_flat = x.reshape(-1, N) + out_flat = RMSNormFunction.apply(x_flat, weight, eps) + return out_flat.reshape(x.shape) diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index 023280d57..77c03926b 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -22,9 +22,11 @@ build_fused_add_rmsnorm_dynamicquant_module, build_fused_add_rmsnorm_module, build_fused_add_rmsnorm_smoothquant_module, + build_rmsnorm_bwd_module, build_rmsnorm_dynamicquant_module, build_rmsnorm_module, build_rmsnorm_smoothquant_module, + rmsnorm, ) from tests.kernels.benchmark_common import ( PerfRow, @@ -569,6 +571,82 @@ def _reference_fused_add_rmsnorm_quant( return residual_expected, q, yscale +def _reference_rmsnorm_bwd(x_dev, weight_dev, dy_dev): + """Eager rmsnorm backward via autograd. Returns dx, dw, rstd (all fp32).""" + x = x_dev.detach().to(DTYPE_FP32).requires_grad_(True) + w = weight_dev.detach().to(DTYPE_FP32).requires_grad_(True) + rstd = torch.rsqrt((x * x).mean(dim=1, keepdim=True) + EPS) + y = x * rstd * w + dx, dw = torch.autograd.grad(y, [x, w], grad_outputs=dy_dev.to(DTYPE_FP32)) + return dx.detach(), dw.detach(), rstd.detach().squeeze(1).contiguous() + + +def run_bwd_test(M: int, N: int, dtype: str = "f32"): + print(f"\nTesting RMSNorm backward (M={M}, N={N}, dtype={dtype})") + + torch_dtype = _torch_dtype(dtype) + try: + fwd_fn = build_rmsnorm_module(N, dtype, store_rstd=True) + bwd_fn = build_rmsnorm_bwd_module(N, dtype) + except Exception as e: + print(f"[FAIL] Compile failed for bwd (M={M}, N={N}, dtype={dtype}): {type(e).__name__}: {e}") + return False + + torch.manual_seed(42) + x = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32).to(torch_dtype).contiguous() + weight = torch.rand((N,), device="cuda", dtype=DTYPE_FP32).to(torch_dtype).contiguous() + dy = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32).to(torch_dtype).contiguous() + + dx_ref, dw_ref, rstd_ref = _reference_rmsnorm_bwd(x, weight, dy) + + stream = torch.cuda.current_stream() + + # --- forward with store_rstd: validates rstd from the kernel --- + out = torch.empty((M, N), device="cuda", dtype=torch_dtype) + rstd = torch.empty((M,), device="cuda", dtype=DTYPE_FP32) + fwd_c = flyc.compile(fwd_fn, x, weight, out, rstd, M, stream) + fwd_c(x, weight, out, rstd, M, stream) + torch.cuda.synchronize() + rstd_err = (rstd - rstd_ref).abs().max().item() + + # --- backward: dx + dweight --- + dx = torch.empty((M, N), device="cuda", dtype=torch_dtype) + dweight = torch.zeros((N,), device="cuda", dtype=DTYPE_FP32) + bwd_c = flyc.compile(bwd_fn, x, weight, dy, rstd, dx, dweight, M, stream) + dweight.zero_() + bwd_c(x, weight, dy, rstd, dx, dweight, M, stream) + torch.cuda.synchronize() + + dx_err = (dx.to(DTYPE_FP32) - dx_ref).abs().max().item() + dw_mag = dw_ref.abs().max().item() + + # Tolerances (calibrated). dweight is summed over M -> larger magnitude -> relative. + rstd_atol = 1e-3 + dx_atol = {"f32": 1e-3, "f16": 3e-2, "bf16": 2e-1}[dtype] + dw_rtol = {"f32": 1e-4, "f16": 3e-2, "bf16": 1e-1}[dtype] + dw_atol = {"f32": 1e-2, "f16": 1e-1, "bf16": 5e-1}[dtype] + + print(f" rstd max abs err = {rstd_err:.3e} (atol={rstd_atol})") + print(f" dx max abs err = {dx_err:.3e} (atol={dx_atol})") + print(f" dweight |max| = {dw_mag:.3e}") + + dw_ok = True + try: + torch.testing.assert_close(dweight, dw_ref, rtol=dw_rtol, atol=dw_atol) + except AssertionError as e: + dw_ok = False + dw_err = (dweight - dw_ref).abs().max().item() + print(f" dweight max abs err = {dw_err:.3e} (rtol={dw_rtol}, atol={dw_atol})") + print(f" [dweight mismatch] {e}") + else: + dw_err = (dweight - dw_ref).abs().max().item() + print(f" dweight max abs err = {dw_err:.3e} (rtol={dw_rtol}, atol={dw_atol})") + + ok = rstd_err < rstd_atol and dx_err < dx_atol and dw_ok + print(f" -> {'PASSED' if ok else 'FAILED'}") + return ok + + def _bench_aiter_rmsnorm(M: int, N: int, dtype: str): torch_dtype = _torch_dtype(dtype) @@ -730,6 +808,105 @@ def test_rmsnorm(): raise SystemExit(1) +def test_rmsnorm_backward(): + print("=" * 80) + print("Running RMSNorm Backward Tests") + print("=" * 80) + + configs = [ + (64, 256, "f32"), # small-N path, f32 + (16, 512, "bf16"), # small-N path, bf16 + (4096, 4096, "bf16"), # generic path, large + (64, 2000, "f32"), # generic path, unaligned + (128, 4096, "f16"), # generic path, f16 + ] + + failures = 0 + for M, N, dtype in configs: + ok = run_bwd_test(M, N, dtype) + if not ok: + failures += 1 + + print("\n" + "=" * 80) + if failures == 0: + print("ALL TESTS PASSED") + else: + print(f"{failures} TESTS FAILED") + print("=" * 80) + if failures != 0: + raise SystemExit(1) + + +def run_autograd_test(M: int, N: int, dtype: str = "f32"): + """End-to-end: the public rmsnorm() autograd path (what quack calls), + including batched (>2D) input reshape and grads on x + weight.""" + print(f"\nTesting rmsnorm() autograd (M={M}, N={N}, dtype={dtype})") + torch_dtype = _torch_dtype(dtype) + torch.manual_seed(42) + + x = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32).to(torch_dtype).requires_grad_(True) + weight = torch.rand((N,), device="cuda", dtype=DTYPE_FP32).to(torch_dtype).requires_grad_(True) + dy = torch.randn((M, N), device="cuda", dtype=torch_dtype) + + out = rmsnorm(x, weight) + out.backward(dy) + dx_out, dw_out = x.grad.detach(), weight.grad.detach() + + # fp32 autograd reference + xf = x.detach().to(DTYPE_FP32).requires_grad_(True) + wf = weight.detach().to(DTYPE_FP32).requires_grad_(True) + rstd = torch.rsqrt((xf * xf).mean(dim=1, keepdim=True) + EPS) + yr = xf * rstd * wf + dxr, dwr = torch.autograd.grad(yr, [xf, wf], dy.to(DTYPE_FP32)) + + out_err = (out.detach().to(DTYPE_FP32) - yr.detach()).abs().max().item() + dx_err = (dx_out.to(DTYPE_FP32) - dxr).abs().max().item() + dw_err = (dw_out.to(DTYPE_FP32) - dwr).abs().max().item() + + out_atol = {"f32": 1e-3, "f16": 3e-2, "bf16": 2e-1}[dtype] + dx_atol = {"f32": 1e-3, "f16": 3e-2, "bf16": 2e-1}[dtype] + dw_atol = {"f32": 1e-2, "f16": 2e-1, "bf16": 1.0}[dtype] + + print(f" out max abs err = {out_err:.3e} (atol={out_atol})") + print(f" dx max abs err = {dx_err:.3e} (atol={dx_atol})") + print(f" dw max abs err = {dw_err:.3e} (atol={dw_atol})") + + # Batched (3D) input must reshape correctly through the public entry. + x3 = torch.randn((4, M // 4 if M >= 4 else 1, N), device="cuda", dtype=torch_dtype, requires_grad=True) + y3 = rmsnorm(x3, weight) + shape_ok = tuple(y3.shape) == tuple(x3.shape) + y3.sum().backward() + grad_ok = x3.grad is not None and tuple(x3.grad.shape) == tuple(x3.shape) + print(f" 3D reshape: out_shape_ok={shape_ok} grad_shape_ok={grad_ok}") + + ok = out_err < out_atol and dx_err < dx_atol and dw_err < dw_atol and shape_ok and grad_ok + print(f" -> {'PASSED' if ok else 'FAILED'}") + return ok + + +def test_rmsnorm_autograd(): + print("=" * 80) + print("Running rmsnorm() Autograd (end-to-end) Tests") + print("=" * 80) + + configs = [ + (64, 256, "f32"), # small-N path + (128, 4096, "bf16"), # generic path + (128, 4096, "f16"), # generic path, f16 + ] + + failures = 0 + for M, N, dtype in configs: + if not run_autograd_test(M, N, dtype): + failures += 1 + + print("\n" + "=" * 80) + print("ALL TESTS PASSED" if failures == 0 else f"{failures} TESTS FAILED") + print("=" * 80) + if failures != 0: + raise SystemExit(1) + + @pytest.mark.large_shape def test_rmsnorm_large_shape(): print("=" * 80) @@ -958,6 +1135,8 @@ def test_fused_add_rmsnorm_smoothquant(): if __name__ == "__main__": test_rmsnorm() + test_rmsnorm_backward() + test_rmsnorm_autograd() test_rmsnorm_dynamicquant() test_rmsnorm_smoothquant() test_fused_add_rmsnorm() From 0c5f94cf87ef03b694db6338a32723b10eb64516 Mon Sep 17 00:00:00 2001 From: root Date: Fri, 3 Jul 2026 05:45:57 +0000 Subject: [PATCH 2/2] [rmsnorm] Address review: honor eps, fix multi-GPU cache, dedup helper Fixes from the PR #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 --- kernels/kernels_common.py | 20 ++++++- kernels/rmsnorm_kernel.py | 101 +++++++++++++++++----------------- tests/kernels/test_rmsnorm.py | 51 +++++++++++++++++ 3 files changed, 121 insertions(+), 51 deletions(-) diff --git a/kernels/kernels_common.py b/kernels/kernels_common.py index 13a9fb44a..9d25cb3e8 100644 --- a/kernels/kernels_common.py +++ b/kernels/kernels_common.py @@ -13,14 +13,32 @@ from flydsl._mlir import ir from flydsl._mlir.dialects import arith as _std_arith from flydsl._mlir.dialects import builtin +from flydsl._mlir.dialects import fly as _fly from flydsl._mlir.dialects import gpu as _gpu from flydsl._mlir.dialects import llvm as _llvm from flydsl._mlir.dialects import scf as _scf -from flydsl.expr import buffer_ops +from flydsl.expr import arith as _expr_arith +from flydsl.expr import buffer_ops, const_expr from flydsl.expr.typing import T from flydsl.runtime.device import get_rocm_arch, is_rdna_arch +def get_llvm_ptr(ptr, offset, dtype_bytes, ptr_type=None): + """Build a global (address-space 1) ``!llvm.ptr`` at ``ptr + offset*dtype_bytes``. + + Shared home for the LLVM-ptr arithmetic used by atomic/global accesses + (previously duplicated in hgemm_splitk.py and rmsnorm_kernel.py). + """ + if ptr_type is None: + ptr_type = ir.Type.parse("!llvm.ptr<1>") + base_ptr = _fly.extract_aligned_pointer_as_index(ptr_type, ptr) + base_ptr = _llvm.PtrToIntOp(T.i64, base_ptr).result + byte_offset = _expr_arith.index_cast(T.i64, fx.Index(offset) * fx.Index(dtype_bytes)) + llvm_ptr = _llvm.AddOp(base_ptr, byte_offset, _llvm.IntegerOverflowFlags(0)).result + llvm_ptr = _llvm.IntToPtrOp(ptr_type, llvm_ptr).result + return llvm_ptr._value if const_expr(hasattr(llvm_ptr, "_value")) else llvm_ptr + + @contextmanager def _if_then(if_op, scf=None): """Context manager for SCF IfOp then-region across old/new Python APIs. diff --git a/kernels/rmsnorm_kernel.py b/kernels/rmsnorm_kernel.py index 4fa81afc7..5701a1892 100644 --- a/kernels/rmsnorm_kernel.py +++ b/kernels/rmsnorm_kernel.py @@ -14,14 +14,12 @@ import flydsl.compiler as flyc import flydsl.expr as fx -from flydsl._mlir import ir -from flydsl._mlir.dialects import fly, llvm +from flydsl._mlir.dialects import llvm from flydsl.expr import arith, const_expr, gpu, range_constexpr from flydsl.expr import math as fmath -from flydsl.expr.typing import T from flydsl.expr.vector import ReductionOp, full from flydsl.runtime.device import get_rocm_arch as get_hip_arch -from kernels.kernels_common import dtype_to_elem_type, get_warp_size +from kernels.kernels_common import dtype_to_elem_type, get_llvm_ptr, get_warp_size try: import torch @@ -61,18 +59,6 @@ def _store_scalar(copy_atom, elem_dtype, store_dtype, divided_tensor, index, val fx.copy_atom_call(copy_atom, r, view) -def _get_llvm_ptr(ptr, offset, dtype_bytes, ptr_type=None): - if ptr_type is None: - ptr_type = ir.Type.parse("!llvm.ptr<1>") - base_ptr = fly.extract_aligned_pointer_as_index(ptr_type, ptr) - base_ptr = llvm.PtrToIntOp(T.i64, base_ptr).result - byte_offset = arith.index_cast(T.i64, fx.Index(offset) * fx.Index(dtype_bytes)) - llvm_ptr = llvm.AddOp(base_ptr, byte_offset, llvm.IntegerOverflowFlags(0)).result - llvm_ptr = llvm.IntToPtrOp(ptr_type, llvm_ptr).result - ptr_v = llvm_ptr._value if const_expr(hasattr(llvm_ptr, "_value")) else llvm_ptr - return ptr_v - - def _load_vec(copy_atom, vec_width, elem_dtype, div_tensor, idx): r = fx.make_rmem_tensor(vec_width, elem_dtype) fx.copy_atom_call(copy_atom, fx.slice(div_tensor, (None, idx)), r) @@ -130,9 +116,9 @@ def _quant_dtype_max(dtype_str: str) -> float: raise ValueError(f"unsupported quant dtype: {dtype_str!r} (expected 'i8' or 'int8')") -def build_rmsnorm_module(N: int, dtype_str: str, store_rstd: bool = False): +def build_rmsnorm_module(N: int, dtype_str: str, store_rstd: bool = False, eps: float = EPS): if N <= 2048: - return _build_rmsnorm_large_m_small_n_module(N, dtype_str, store_rstd) + return _build_rmsnorm_large_m_small_n_module(N, dtype_str, store_rstd, eps) arch = get_hip_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") @@ -155,7 +141,7 @@ def rmsnorm_kernel( elem_dtype = dtype_to_elem_type(dtype_str) fm_fast = arith.FastMathFlags.fast - eps_c = EPS + eps_c = eps n_float = float(N) lds = fx.SharedAllocator().allocate(SharedStorage).peek() @@ -361,7 +347,7 @@ def launch_rmsnorm( return launch_rmsnorm -def _build_rmsnorm_large_m_small_n_module(N: int, dtype_str: str, store_rstd: bool = False): +def _build_rmsnorm_large_m_small_n_module(N: int, dtype_str: str, store_rstd: bool = False, eps: float = EPS): BLOCK_N = 1 << (N - 1).bit_length() BLOCK_M = max(min(16384 // BLOCK_N, 32), 8) THREADS_PER_ROW = min(WARP_SIZE, 1024 // BLOCK_M) @@ -386,7 +372,7 @@ def rmsnorm_large_m_small_n_kernel( if row < MIn: elem_dtype = dtype_to_elem_type(dtype_str) fm_fast = arith.FastMathFlags.fast - eps_c = EPS + eps_c = eps n_float = float(N) Input_buf = fx.rocdl.make_buffer_tensor(Input) @@ -493,7 +479,12 @@ def build_rmsnorm_bwd_module(N: int, dtype_str: str): Pass 1: c1 = mean_N(x_hat * wdy), x_hat = x*rstd, wdy = dy*gamma. Pass 2: dx = (wdy - x_hat*c1) * rstd -> DX (elem dtype); dw_elem = dy * x_hat (fp32) -> atomicAdd into DWeight[idx] (fp32). - Uses the module EPS constant implicitly via Rstd (supplied by the forward). + eps is baked into Rstd by the forward, so it is not needed here. + + Perf follow-ups (deferred; correctness-complete as-is): this is the generic + scalar path only — a vectorized fast path (mirroring the forward) and caching + x/dy/gamma between pass 1 and pass 2 (the forward caches `in_local`) would cut + global traffic. Left out of PR 1 to keep the first backward reviewable. """ RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) elem_bits = 32 if dtype_str == "f32" else 16 @@ -608,7 +599,7 @@ def block_reduce_add(val): _store_scalar(copy_atom_s, elem_dtype, elem_dtype, dx_div, idx, dx_e) dw = dy * x_hat - ptr = _get_llvm_ptr(DWeight, idx, 4) + ptr = get_llvm_ptr(DWeight, idx, 4) llvm.AtomicRMWOp( llvm.AtomicBinOp.fadd, ptr, @@ -1631,15 +1622,17 @@ def _torch_dtype_to_str(dt) -> str: return "bf16" raise ValueError(f"unsupported torch dtype: {dt}") - # Compiled-fn caches keyed by (N, dtype_str, store_rstd) / (N, dtype_str). + # Compiled-fn caches. Keys include device: a compiled function is bound to + # the device/context it was built on, so reusing it on another GPU faults. + # eps is a compile-time kernel constant, so it is part of the fwd key too. _FWD_CACHE: dict = {} _BWD_CACHE: dict = {} - def _get_fwd_compiled(x, weight, out, rstd, M, N, dtype_str, store_rstd, stream): - key = (N, dtype_str, store_rstd) + def _get_fwd_compiled(x, weight, out, rstd, M, N, dtype_str, store_rstd, eps, stream): + key = (N, dtype_str, store_rstd, float(eps), x.device) entry = _FWD_CACHE.get(key) if entry is None: - launch_fn = build_rmsnorm_module(N, dtype_str, store_rstd=store_rstd) + launch_fn = build_rmsnorm_module(N, dtype_str, store_rstd=store_rstd, eps=eps) if store_rstd: compiled = flyc.compile(launch_fn, x, weight, out, rstd, M, stream) else: @@ -1649,42 +1642,49 @@ def _get_fwd_compiled(x, weight, out, rstd, M, N, dtype_str, store_rstd, stream) return entry def rmsnorm_fwd(x, weight, eps=EPS, store_rstd=False): - """Forward RMSNorm. Returns (out, rstd). - - NOTE: the compiled kernel uses the module EPS constant (not `eps`); - `eps` is accepted for interface parity only (see module docstring). - """ + """Forward RMSNorm. Returns (out, rstd). eps is baked into the kernel.""" assert x.dim() == 2, "rmsnorm_fwd expects a 2D (M, N) input" + assert x.is_contiguous() and weight.is_contiguous(), "rmsnorm_fwd expects contiguous inputs" M, N = x.shape out = torch.empty_like(x) rstd = torch.empty((M,), device=x.device, dtype=torch.float32) if store_rstd else None dtype_str = _torch_dtype_to_str(x.dtype) - stream = torch.cuda.current_stream() - compiled = _get_fwd_compiled(x, weight, out, rstd, M, N, dtype_str, store_rstd, stream) - if store_rstd: - compiled(x, weight, out, rstd, M, stream) - else: - compiled(x, weight, out, M, stream) + # Bind compile + launch to the tensors' device so the compiled kernel and + # the stream belong to the right GPU/context (multi-GPU correctness). + with torch.cuda.device(x.device): + stream = torch.cuda.current_stream() + compiled = _get_fwd_compiled(x, weight, out, rstd, M, N, dtype_str, store_rstd, eps, stream) + if store_rstd: + compiled(x, weight, out, rstd, M, stream) + else: + compiled(x, weight, out, M, stream) return out, rstd def rmsnorm_bwd(x, weight, dout, rstd, eps=EPS): - """Backward RMSNorm. Returns (dx, dw) with dw cast to weight dtype.""" + """Backward RMSNorm. Returns (dx, dw) with dw cast to weight dtype. + + eps is not used directly here — it is already baked into `rstd` by the + forward — but is accepted so callers can pass it symmetrically. + """ assert x.dim() == 2, "rmsnorm_bwd expects a 2D (M, N) input" + assert x.is_contiguous() and dout.is_contiguous(), "rmsnorm_bwd expects contiguous inputs" M, N = x.shape dtype_str = _torch_dtype_to_str(x.dtype) dx = torch.empty_like(x) dweight = torch.zeros((N,), device=x.device, dtype=torch.float32) - stream = torch.cuda.current_stream() - key = (N, dtype_str) - compiled = _BWD_CACHE.get(key) - if compiled is None: - launch_fn = build_rmsnorm_bwd_module(N, dtype_str) - # flyc.compile executes the kernel once during tracing, which would - # accumulate into DWeight; zero it AFTER compiling. - compiled = flyc.compile(launch_fn, x, weight, dout, rstd, dx, dweight, M, stream) - _BWD_CACHE[key] = compiled - dweight.zero_() - compiled(x, weight, dout, rstd, dx, dweight, M, stream) + key = (N, dtype_str, x.device) + # Bind compile + launch to the tensors' device (multi-GPU correctness). + with torch.cuda.device(x.device): + stream = torch.cuda.current_stream() + compiled = _BWD_CACHE.get(key) + if compiled is None: + launch_fn = build_rmsnorm_bwd_module(N, dtype_str) + # flyc.compile executes the kernel once during tracing, which would + # accumulate into DWeight; zero it AFTER compiling. + compiled = flyc.compile(launch_fn, x, weight, dout, rstd, dx, dweight, M, stream) + _BWD_CACHE[key] = compiled + dweight.zero_() + compiled(x, weight, dout, rstd, dx, dweight, M, stream) return dx, dweight.to(weight.dtype) class RMSNormFunction(torch.autograd.Function): @@ -1706,6 +1706,7 @@ def rmsnorm(x, weight=None, eps=EPS): """Public entry: plain RMSNorm with autograd (weight required in PR 1).""" assert weight is not None, "PR 1 rmsnorm requires an explicit weight" N = weight.shape[-1] + assert x.shape[-1] == N, f"x last dim {x.shape[-1]} != weight length {N}" x_flat = x.reshape(-1, N) out_flat = RMSNormFunction.apply(x_flat, weight, eps) return out_flat.reshape(x.shape) diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index 77c03926b..37abaf1dc 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -907,6 +907,55 @@ def test_rmsnorm_autograd(): raise SystemExit(1) +def test_rmsnorm_eps_honored(): + """eps must be baked into the kernel, not silently replaced by the module EPS.""" + print("=" * 80) + print("Running RMSNorm eps-honored Test") + print("=" * 80) + torch.manual_seed(0) + M, N = 32, 256 + x = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) + w = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + + for eps in (1e-5, 1e-6, 1e-2): + y = rmsnorm(x, w, eps=eps) + ref = x / torch.sqrt((x * x).mean(dim=1, keepdim=True) + eps) * w + err = (y - ref).abs().max().item() + print(f" eps={eps:g}: max err vs torch ref = {err:.3e}") + assert err < 1e-4, f"eps={eps} not honored (err={err})" + + # A non-default eps must actually change the output (guards silent-ignore regressions). + diff = (rmsnorm(x, w, eps=1e-2) - rmsnorm(x, w, eps=1e-6)).abs().max().item() + print(f" eps 1e-2 vs 1e-6 output diff = {diff:.3e} (must be > 0)") + assert diff > 0, "eps appears to be ignored" + print(" -> PASSED") + + +@pytest.mark.multi_gpu +def test_rmsnorm_multi_gpu(): + """Compiled-fn cache must not reuse a device-0 kernel on device 1 (would fault).""" + print("=" * 80) + print("Running RMSNorm multi-GPU Test") + print("=" * 80) + if torch.cuda.device_count() < 2: + pytest.skip("needs >=2 GPUs") + + torch.manual_seed(0) + N = 256 + for dev in ("cuda:0", "cuda:1"): + x = torch.randn((16, N), device=dev, dtype=DTYPE_FP32, requires_grad=True) + w = torch.rand((N,), device=dev, dtype=DTYPE_FP32, requires_grad=True) + dy = torch.randn((16, N), device=dev, dtype=DTYPE_FP32) + y = rmsnorm(x, w) + y.backward(dy) + torch.cuda.synchronize(dev) + ref = x.detach() / torch.sqrt((x.detach() ** 2).mean(1, keepdim=True) + EPS) * w.detach() + err = (y.detach() - ref).abs().max().item() + print(f" {dev}: out err={err:.3e}, dx finite={torch.isfinite(x.grad).all().item()}") + assert err < 1e-4 and torch.isfinite(x.grad).all() + print(" -> PASSED") + + @pytest.mark.large_shape def test_rmsnorm_large_shape(): print("=" * 80) @@ -1137,6 +1186,8 @@ def test_fused_add_rmsnorm_smoothquant(): test_rmsnorm() test_rmsnorm_backward() test_rmsnorm_autograd() + test_rmsnorm_eps_honored() + test_rmsnorm_multi_gpu() test_rmsnorm_dynamicquant() test_rmsnorm_smoothquant() test_fused_add_rmsnorm()