# Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao. import math from typing import Optional, Tuple, Type from functools import partial import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute from cutlass import Float32, Int32, const_expr import torch from torch import Tensor import quack.utils as utils import quack.copy_utils as copy_utils import quack.layout_utils as layout_utils from quack.compile_utils import make_fake_tensor as fake_tensor from quack.reduce import row_reduce from quack.reduction_base import ReductionBase from quack.cute_dsl_utils import torch2cute_dtype_map class RMSNorm(ReductionBase): def __init__(self, dtype: Type[cutlass.Numeric], N: int, is_layernorm: bool = False): super().__init__(dtype, N, stage=2 if is_layernorm else 1) self.is_layernorm = is_layernorm self.reload_from = None if N <= (16384 if is_layernorm else 8192) else "smem" self.delay_w_load = False def _threads_per_row(self): N = self.N for limit, threads in [(64, 8), (128, 16), (3072, 32), (6144, 64), (16384, 128)]: if N <= limit: return threads return 256 def _set_cluster_n(self): N = self.N # cluster_n = 4 is faster and cluster_n = 2 for N=64k for some reason # Similarly cluster_n = 8 is faster for N=128k if const_expr(self.dtype.width == 16): thresholds = [(16 * 1024, 1), (32 * 1024, 2), (64 * 1024, 4), (128 * 1024, 8)] else: thresholds = [(32 * 1024, 1), (64 * 1024, 2), (128 * 1024, 4), (256 * 1024, 8)] for limit, cluster in thresholds: if N <= limit: self.cluster_n = cluster return self.cluster_n = 16 @cute.jit def __call__( self, mX: cute.Tensor, mW: Optional[cute.Tensor], mB: Optional[cute.Tensor], mRes: Optional[cute.Tensor], mO: cute.Tensor, mResO: Optional[cute.Tensor], mRstd: Optional[cute.Tensor], mMean: Optional[cute.Tensor], eps: Float32, stream: cuda.CUstream, ): assert mX.element_type == self.dtype self._set_cluster_n() largest_dtype_width = const_expr( max(*(t.element_type.width for t in [mX, mRes, mW, mB, mO, mResO] if t is not None)) ) vecsize = math.gcd(self.N, 128 // largest_dtype_width) tiled_copy, tiler_mn, threads_per_row = self._get_tiled_copy(vecsize=vecsize) num_threads = tiled_copy.size mW, mB = [ layout_utils.expand(mT, dim=0, size=tiler_mn[0]) if const_expr(mT is not None) else None for mT in (mW, mB) ] mRstd, mMean = [ layout_utils.expand(mT, dim=1, size=self.N) if const_expr(mT is not None) else None for mT in (mRstd, mMean) ] self.kernel( mX, mW, mB, mRes, mO, mResO, mRstd, mMean, eps, tiler_mn, tiled_copy, threads_per_row ).launch( grid=[cute.ceil_div(mX.shape[0], tiler_mn[0]), self.cluster_n, 1], block=[num_threads, 1, 1], cluster=[1, self.cluster_n, 1] if const_expr(self.cluster_n > 1) else None, stream=stream, ) @cute.kernel def kernel( self, mX: cute.Tensor, mW: Optional[cute.Tensor], mB: Optional[cute.Tensor], mRes: Optional[cute.Tensor], mO: cute.Tensor, mResO: Optional[cute.Tensor], mRstd: Optional[cute.Tensor], mMean: Optional[cute.Tensor], eps: Float32, tiler_mn: cute.Shape, tiled_copy: cute.TiledCopy, threads_per_row: cutlass.Constexpr[int], ): tidx, _, _ = cute.arch.thread_idx() bidx, _, _ = cute.arch.block_idx() cluster_y = const_expr(0) if const_expr(self.cluster_n == 1) else cute.arch.block_idx()[1] tv_layout = tiled_copy.layout_tv_tiled smem = cutlass.utils.SmemAllocator() sX = smem.allocate_tensor( mX.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16 ) if const_expr(mRes is not None): sRes = smem.allocate_tensor( mRes.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16, ) reduction_buffer, mbar_ptr = self._allocate_reduction_buffer_and_mbar(smem, tv_layout) shape = mX.shape idX = cute.make_identity_tensor(shape) # slice for CTAs gX, gRes, gO, gResO, gRstd, gMean, cX = [ cute.local_tile(mT, tiler_mn, (bidx, cluster_y)) if mT is not None else None for mT in (mX, mRes, mO, mResO, mRstd, mMean, idX) ] gW, gB = [ cute.local_tile(mT, tiler_mn, (0, cluster_y)) if const_expr(mT is not None) else None for mT in (mW, mB) ] thr_copy_X = tiled_copy.get_slice(tidx) tXgW = thr_copy_X.partition_S(gW) if const_expr(mW is not None) else None tXgB = thr_copy_X.partition_S(gB) if const_expr(mB is not None) else None tXgX = thr_copy_X.partition_S(gX) tXsX = thr_copy_X.partition_D(sX) if const_expr(mRes is not None): tXgRes = thr_copy_X.partition_S(gRes) tXsRes = thr_copy_X.partition_D(sRes) tXgO = thr_copy_X.partition_D(gO) if const_expr(mResO is not None): tXgResO = thr_copy_X.partition_D(gResO) tXrRstd = thr_copy_X.partition_D(gRstd) if const_expr(mRstd is not None) else None tXrMean = thr_copy_X.partition_D(gMean) if const_expr(mMean is not None) else None tXcX = thr_copy_X.partition_S(cX)[(0, None), None, None] # allocate fragments for gmem->rmem tXrW = cute.make_fragment_like(tXgW) if const_expr(mW is not None) else None tXrB = cute.make_fragment_like(tXgB) if const_expr(mB is not None) else None tXrX, tXrO = [cute.make_fragment_like(t) for t in (tXgX, tXgO)] if const_expr(mRes is not None): tXrRes = cute.make_fragment_like(tXgRes) num_warps = cute.size(tiled_copy) // cute.arch.WARP_SIZE self._initialize_cluster(tidx, mbar_ptr, num_warps) is_even_N = const_expr(shape[1] == tiler_mn[1] * self.cluster_n) tXpX = ( copy_utils.predicate_k(thr_copy_X.partition_S(cX), limit=shape[1]) if not is_even_N else None ) # Each copy will use the same predicate copy = partial(copy_utils.copy, pred=tXpX) row = tXcX[0][0] if row < shape[0]: copy(tXgX, tXsX, is_async=True) if const_expr(mRes is not None): copy(tXgRes, tXsRes, is_async=True) cute.arch.cp_async_commit_group() if const_expr(not self.delay_w_load): if const_expr(mW is not None): copy(tXgW, tXrW) if const_expr(mB is not None): copy(tXgB, tXrB) cute.arch.cp_async_wait_group(0) cute.autovec_copy(tXsX, tXrX) x = tXrX.load().to(cute.Float32) if const_expr(mRes is not None): cute.autovec_copy(tXsRes, tXrRes) x += tXrRes.load().to(cute.Float32) if const_expr(mResO is not None): tXrResO = cute.make_fragment_like(tXgResO) tXrResO.store(x.to(tXrResO.element_type)) if row < shape[0]: copy(tXrResO, tXgResO) mean, rstd = None, None if const_expr(self.is_layernorm): # LayerNorm: compute mean first, then variance sum_x = row_reduce( x, cute.ReductionOp.ADD, threads_per_row, reduction_buffer[None, None, 0], mbar_ptr + 0 if const_expr(self.cluster_n > 1) else None, init_val=0.0, hook_fn=cute.arch.cluster_wait if const_expr(self.cluster_n > 1) else None, ) mean = sum_x / shape[1] if const_expr(mMean is not None): # Only the thread corresponding to column 0 writes out the mean to gmem if ( tXcX[0][1] == 0 and row < shape[0] and (self.cluster_n == 1 or cute.arch.block_idx_in_cluster() == 0) ): tXrMean[0] = mean if const_expr(self.reload_from == "smem"): cute.autovec_copy(tXsX, tXrX) x = tXrX.load().to(cute.Float32) if const_expr(mRes is not None): cute.autovec_copy(tXsRes, tXrRes) x += tXrRes.load().to(cute.Float32) elif const_expr(self.reload_from == "gmem"): copy(tXgX, tXrX) x = tXrX.load().to(cute.Float32) if const_expr(mRes is not None): copy(tXgRes, tXrRes) x += tXrRes.load().to(cute.Float32) sum_sq_x_sub_mean = row_reduce( (x - mean) * (x - mean), cute.ReductionOp.ADD, threads_per_row, reduction_buffer[None, None, 1], mbar_ptr + 1 if const_expr(self.cluster_n > 1) else None, init_val=0.0, ) rstd = cute.math.rsqrt(sum_sq_x_sub_mean / shape[1] + eps, fastmath=True) else: # RMSNorm: compute sum of squares directly mean = const_expr(0.0) sum_sq_x = row_reduce( x * x, cute.ReductionOp.ADD, threads_per_row, reduction_buffer[None, None, 0], mbar_ptr, init_val=0.0, hook_fn=cute.arch.cluster_wait if const_expr(self.cluster_n > 1) else None, ) rstd = cute.math.rsqrt(sum_sq_x / shape[1] + eps, fastmath=True) if const_expr(mRstd is not None): # Only the thread corresponding to column 0 writes out the rstd to gmem if ( tXcX[0][1] == 0 and row < shape[0] and (self.cluster_n == 1 or cute.arch.block_idx_in_cluster() == 0) ): tXrRstd[0] = rstd if const_expr(self.delay_w_load): if const_expr(mW is not None): copy(tXgW, tXrW) if const_expr(mB is not None): copy(tXgB, tXrB) if const_expr(self.reload_from == "smem" or self.reload_from == "gmem"): if const_expr(self.reload_from == "smem"): cute.autovec_copy(tXsX, tXrX) if const_expr(mRes is not None): cute.autovec_copy(tXsRes, tXrRes) else: copy(tXgX, tXrX) if const_expr(mRes is not None): copy(tXgRes, tXrRes) x = tXrX.load().to(cute.Float32) if const_expr(mRes is not None): x += tXrRes.load().to(cute.Float32) x_hat = (x - mean) * rstd if const_expr(self.is_layernorm) else x * rstd y = x_hat if const_expr(mW is not None): y *= tXrW.load().to(cute.Float32) if const_expr(mB is not None): y += tXrB.load().to(cute.Float32) tXrO.store(y.to(tXrO.element_type)) if row < shape[0]: copy(tXrO, tXgO) @torch.library.custom_op( "quack::_rmsnorm_fwd", mutates_args=("out", "rstd", "mean", "residual_out"), device_types="cuda", # We need to specify the schema manually since we're mutating an optional tensor schema="(Tensor x, Tensor? weight, Tensor(a2!) out, Tensor? bias, Tensor(a4!)? rstd, Tensor(a5!)? mean, Tensor? residual, Tensor(a7!)? residual_out, float eps=1e-6, bool is_layernorm=False) -> ()", ) def _rmsnorm_fwd( x: Tensor, weight: Optional[Tensor], out: Tensor, bias: Optional[Tensor] = None, rstd: Optional[Tensor] = None, mean: Optional[Tensor] = None, residual: Optional[Tensor] = None, residual_out: Optional[Tensor] = None, eps: float = 1e-6, is_layernorm: bool = False, ) -> None: """RMSNorm/LayerNorm forward pass. Args: x: Input tensor of shape (M, N) weight: Optional weight tensor of shape (N,) eps: Small value for numerical stability is_layernorm: If True, compute LayerNorm instead of RMSNorm Returns: Normalized output tensor of same shape as x """ # Don't need to check is_cuda since torch.library ensures that supported_types = {torch.float16, torch.bfloat16, torch.float32} assert x.dtype in supported_types, "Unsupported dtype" if weight is not None: assert weight.dtype in supported_types, "Weight must be float32, float16 or bfloat16" if residual is not None: assert residual.dtype in supported_types, "Residual must be float16, bfloat16, or float32" _, N = x.shape dtype, out_dtype, weight_dtype, bias_dtype, res_dtype, res_out_dtype = [ torch2cute_dtype_map[t.dtype] if t is not None else None for t in [x, out, weight, bias, residual, residual_out] ] compile_key = ( dtype, out_dtype, res_dtype, weight_dtype, bias_dtype, res_out_dtype, N, rstd is not None, mean is not None, is_layernorm, ) if compile_key not in _rmsnorm_fwd.compile_cache: batch_sym = cute.sym_int() all_dtypes = [dtype, out_dtype, res_dtype, weight_dtype, bias_dtype, res_out_dtype] div = math.gcd(N, *(128 // dt.width for dt in all_dtypes if dt is not None)) x_cute, out_cute, res_cute, res_out_cute = [ fake_tensor(dt, (batch_sym, N), div) for dt in [dtype, out_dtype, res_dtype, res_out_dtype] ] weight_cute, bias_cute = [fake_tensor(dt, (N,), div) for dt in [weight_dtype, bias_dtype]] rstd_cute = fake_tensor(Float32, (batch_sym,)) if rstd is not None else None mean_cute = fake_tensor(Float32, (batch_sym,)) if mean is not None else None _rmsnorm_fwd.compile_cache[compile_key] = cute.compile( RMSNorm(dtype, N, is_layernorm=is_layernorm), x_cute, weight_cute, bias_cute, res_cute, out_cute, res_out_cute, rstd_cute, mean_cute, Float32(0), # eps, just for compilation cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), options="--enable-tvm-ffi", ) _rmsnorm_fwd.compile_cache[compile_key]( x, weight, bias, residual, out, residual_out, rstd, mean, eps ) _rmsnorm_fwd.compile_cache = {} def rmsnorm_fwd( x: Tensor, weight: Optional[Tensor] = None, bias: Optional[Tensor] = None, residual: Optional[Tensor] = None, out_dtype: Optional[torch.dtype] = None, residual_dtype: Optional[torch.dtype] = None, eps: float = 1e-6, store_rstd: bool = False, ) -> Tuple[Tensor, Tensor, Optional[Tensor]]: # Need to wrap to handle the case where residual_out is a alias of x, which makes torch.library # and torch.compile unhappy. Also allocate memory for out and residual_out if they are None # so that _layer_norm_fwd_impl doesn't have to return them. out_dtype = x.dtype if out_dtype is None else out_dtype out = torch.empty_like(x, dtype=out_dtype) rstd = torch.empty(x.shape[0], device=x.device, dtype=torch.float32) if store_rstd else None if residual is not None: residual_dtype = residual.dtype if residual is not None or (residual_dtype is not None and residual_dtype != x.dtype): residual_out = torch.empty_like( x, dtype=residual_dtype if residual_dtype is not None else x.dtype ) else: residual_out = None _rmsnorm_fwd(x, weight, out, bias, rstd, None, residual, residual_out, eps, False) # residual_out is None if residual is None and residual_dtype == input_dtype and dropout_p == 0.0 if residual_out is None: residual_out = x return out, residual_out, rstd def rmsnorm_ref(x, w=None, bias=None, residual=None, eps=1e-6): x_f32 = x.float() if residual is not None: residual_f32 = residual.float() x_f32 += residual_f32 x_norm = x_f32 / (torch.sqrt(torch.mean(x_f32.square(), dim=-1, keepdim=True) + eps)) out = x_norm * w if w is not None else x_norm if bias is not None: out = out + bias.float() if residual is None: return out.to(x.dtype) else: return out.to(x.dtype), x_f32.to(residual.dtype) def rmsnorm_bwd_ref(x, w, dout, rstd, eps=1e-6): """Reference implementation for RMSNorm backward pass.""" x_f32 = x.float() x_hat = x_f32 * rstd.unsqueeze(1) if w is not None: wdy = dout * w else: wdy = dout c1 = (x_hat * wdy).mean(dim=-1, keepdim=True) dx = (wdy - x_hat * c1) * rstd.unsqueeze(1) # dL/dW if w is not None: dw = (dout * x_hat).sum(dim=0) return dx.to(x.dtype), dw.to(w.dtype) else: return dx.to(x.dtype), None class RMSNormBackward(ReductionBase): def __init__(self, dtype: cutlass.Numeric, N: int): # 2 stages for double buffering when computing mean of x_hat * wdy super().__init__(dtype, N, stage=2, reduction_dtype=Float32) self.reload_wdy = None if N <= 16 * 1024 else "smem" if self.N > 128 * 1024 and self.dtype.width >= 32: # Not enough smem raise ValueError("RMSNormBackward does not support N > 128k with dtype >= 32 bits") def _num_threads(self): return 128 if self.N <= 4096 else 256 def _threads_per_row(self): N = self.N for limit, threads in [(64, 8), (128, 16), (256, 32), (512, 64), (4096, 128)]: if N <= limit: return threads return 256 def _set_cluster_n(self): N = self.N for limit, cluster in [(8 * 1024, 1), (16 * 1024, 2), (32 * 1024, 4), (64 * 1024, 8)]: if N <= limit: self.cluster_n = cluster return self.cluster_n = 16 @cute.jit def __call__( self, mX: cute.Tensor, mW: Optional[cute.Tensor], mdO: cute.Tensor, mdResO: Optional[cute.Tensor], mRstd: cute.Tensor, mdX: cute.Tensor, mdW: Optional[cute.Tensor], mdRes: Optional[cute.Tensor], mdB: Optional[cute.Tensor], sm_count: Int32, stream: cuda.CUstream, ): assert mX.element_type == self.dtype self._set_cluster_n() largest_dtype_width = const_expr( max(*(t.element_type.width for t in [mX, mW, mdO, mdResO, mdX, mdRes] if t is not None)) ) vecsize = math.gcd(self.N, 128 // largest_dtype_width) tiled_copy, tiler_mn, threads_per_row = self._get_tiled_copy(vecsize=vecsize) num_threads = tiled_copy.size mW = ( layout_utils.expand(mW, dim=0, size=tiler_mn[0]) if const_expr(mW is not None) else None ) num_blocks = sm_count self.kernel( mX, mW, mdO, mdResO, mRstd, mdX, mdW, mdB, mdRes, tiler_mn, tiled_copy, threads_per_row ).launch( grid=[num_blocks, self.cluster_n, 1], block=[num_threads, 1, 1], cluster=[1, self.cluster_n, 1] if self.cluster_n > 1 else None, stream=stream, ) @cute.kernel def kernel( self, mX: cute.Tensor, mW: Optional[cute.Tensor], mdO: cute.Tensor, mdResO: Optional[cute.Tensor], mRstd: cute.Tensor, mdX: cute.Tensor, mdW: Optional[cute.Tensor], mdB: Optional[cute.Tensor], mdRes: Optional[cute.Tensor], tiler_mn: cute.Shape, tiled_copy: cute.TiledCopy, threads_per_row: cutlass.Constexpr[int], ): tidx, _, _ = cute.arch.thread_idx() bidx_start, _, _ = cute.arch.block_idx() gdim, _, _ = cute.arch.grid_dim() cluster_y = const_expr(0) if const_expr(self.cluster_n == 1) else cute.arch.block_idx()[1] tv_layout = tiled_copy.layout_tv_tiled shape = mX.shape M, N = shape[0], shape[1] is_even_N = const_expr(shape[1] == tiler_mn[1] * self.cluster_n) idX = cute.make_identity_tensor(shape) smem = cutlass.utils.SmemAllocator() smem_layout = cute.make_ordered_layout((tiler_mn[0], tiler_mn[1], 2), order=(1, 0, 2)) sX = smem.allocate_tensor(mX.element_type, smem_layout, byte_alignment=16) sdO = smem.allocate_tensor(mdO.element_type, smem_layout, byte_alignment=16) reduction_buffer, mbar_ptr = self._allocate_reduction_buffer_and_mbar( smem, tv_layout, is_persistent=True ) if const_expr(mbar_ptr is not None): mbar_full_ptr, mbar_empty_ptr = mbar_ptr, mbar_ptr + 2 else: mbar_full_ptr, mbar_empty_ptr = None, None thr_copy_X = tiled_copy.get_slice(tidx) gX, gdO, gdResO, gdX, gdRes, cX = [ cute.local_tile(mT, tiler_mn, (None, cluster_y)) if mT is not None else None for mT in (mX, mdO, mdResO, mdX, mdRes, idX) ] gW = cute.local_tile(mW, tiler_mn, (0, cluster_y)) if mW is not None else None gdW, gdB = [ cute.local_tile(mT, (1, tiler_mn[1]), (bidx_start, cluster_y)) if const_expr(mT is not None) else None for mT in (mdW, mdB) ] tXgX = thr_copy_X.partition_S(gX) tXsX = thr_copy_X.partition_D(sX) tXgdO = thr_copy_X.partition_S(gdO) tXsdO = thr_copy_X.partition_D(sdO) tXgdX = thr_copy_X.partition_D(gdX) if const_expr(mdResO is not None): tXgdResO = thr_copy_X.partition_S(gdResO) if const_expr(mdRes is not None): tXgdRes = thr_copy_X.partition_D(gdRes) tXcX = thr_copy_X.partition_S(cX)[(0, None), None, None, None] tXrX, tXrdO, tXrdX = [ cute.make_fragment_like(thr[None, None, None, 0]) for thr in (tXgX, tXgdO, tXgdX) ] tXrdResO = None if const_expr(mdResO is not None): tXrdResO = cute.make_fragment_like(tXgdResO[None, None, None, 0]) tXrdRes = None if const_expr(mdRes is not None): tXrdRes = cute.make_fragment_like(tXgdRes[None, None, None, 0]) # This doesn't change across iterations tXpX = ( None if is_even_N else copy_utils.predicate_k(thr_copy_X.partition_S(cX[None, None, 0]), limit=shape[1]) ) # Each copy will use the same number of elements as X copy = partial(copy_utils.copy, pred=tXpX) tXgdW, tXrdW = None, None tXgdB, tXrdB = None, None if const_expr(mdW is not None): tXgdW = thr_copy_X.partition_S(gdW) # Always compute partial weight gradients in fp32 tXrdW = cute.make_fragment_like(tXgdW, Float32) if const_expr(mdB is not None): tXgdB = thr_copy_X.partition_S(gdB) # Always compute partial bias gradients in fp32 tXrdB = cute.make_fragment_like(tXgdB, Float32) num_warps = cute.size(tiled_copy) // cute.arch.WARP_SIZE self._initialize_cluster(tidx, mbar_ptr, num_warps, is_persistent=True) tXrW = None if const_expr(mW is not None): tXgW = thr_copy_X.partition_S(gW) tXrW = cute.make_fragment_like(tXgW) # Need this, otherwise rW can have arbitrary values that changes the reduction if const_expr(not is_even_N): tXrW.fill(0.0) copy(tXgW, tXrW) # Prefetch the first batch row = tXcX[None, None, None, bidx_start][0][0] if row < M: copy(tXgX[None, None, None, bidx_start], tXsX[None, None, None, 0], is_async=True) copy(tXgdO[None, None, None, bidx_start], tXsdO[None, None, None, 0], is_async=True) else: if const_expr(tiler_mn[0] > 1): # Fill with zero, otherwise smem will be uninitialized, and we could read this back # later into registers, causing wrong dW. utils.fill_oob(tXsX[None, None, None, 0], None, fill_value=mX.element_type.zero) utils.fill_oob(tXsdO[None, None, None, 0], None, fill_value=mdO.element_type.zero) cute.arch.cp_async_commit_group() if const_expr(self.cluster_n > 1): cute.arch.cluster_wait() if const_expr(mdW is not None): tXrdW.fill(0.0) if const_expr(mdB is not None): tXrdB.fill(0.0) stage = Int32(0) producer_phase = Int32(1) consumer_phase = Int32(0) for bidx in cutlass.range(bidx_start, cute.ceil_div(M, tiler_mn[0]), gdim): row = tXcX[None, None, None, bidx][0][0] if row + gdim * tiler_mn[0] < M: # Prefetch the next batch copy( tXgX[None, None, None, bidx + gdim], tXsX[None, None, None, stage ^ 1], is_async=True, ) copy( tXgdO[None, None, None, bidx + gdim], tXsdO[None, None, None, stage ^ 1], is_async=True, ) else: if const_expr(tiler_mn[0] > 1): utils.fill_oob( tXsX[None, None, None, stage ^ 1], None, fill_value=mX.element_type.zero ) utils.fill_oob( tXsdO[None, None, None, stage ^ 1], None, fill_value=mdO.element_type.zero ) cute.arch.cp_async_commit_group() rstd = cutlass.Float.zero if row < M or tiler_mn[0] == 1: rstd = mRstd[row] if const_expr(mdResO is not None): if row < M or tiler_mn[0] == 1: copy(tXgdResO[None, None, None, bidx], tXrdResO) elif tiler_mn[0] > 1: tXrdResO.fill(0.0) cute.arch.cp_async_wait_group(1) cute.autovec_copy(tXsX[None, None, None, stage], tXrX) x = tXrX.load().to(cute.Float32) cute.autovec_copy(tXsdO[None, None, None, stage], tXrdO) dout = tXrdO.load().to(cute.Float32) x_hat = x * rstd wdy = dout if const_expr(mW is not None): wdy *= tXrW.load().to(Float32) if const_expr(self.cluster_n > 1): cute.arch.mbarrier_wait(mbar_empty_ptr + stage, producer_phase) mean_xhat_wdy = ( row_reduce( x_hat * wdy, cute.ReductionOp.ADD, threads_per_row, reduction_buffer[None, None, stage], (mbar_full_ptr + stage if const_expr(self.cluster_n > 1) else None), phase=consumer_phase, init_val=0.0, ) / shape[1] ) if const_expr(self.cluster_n > 1): # Need this fence since the STAS from the producer is using the async proxy. cute.arch.fence_proxy( cute.arch.ProxyKind.async_shared, space=cute.arch.SharedSpace.shared_cta ) # It's faster to have 1 lane per warp to signal the mbar, rather than all lanes # Requires adjusting the thread_count when initializing the mbar cute.arch.sync_warp() lane_idx = cute.arch.lane_idx() if lane_idx < self.cluster_n: cute.arch.mbarrier_arrive( mbar_empty_ptr + stage, peer_cta_rank_in_cluster=lane_idx ) if const_expr(self.reload_wdy == "smem"): cute.autovec_copy(tXsdO[None, None, None, stage], tXrdO) dout = tXrdO.load().to(cute.Float32) wdy = dout if const_expr(mW is not None): wdy *= tXrW.load().to(Float32) dx = (wdy - x_hat * mean_xhat_wdy) * rstd if const_expr(mdResO is not None): dx += tXrdResO.load().to(cute.Float32) tXrdX.store(dx.to(tXrdX.element_type)) if row < M or tiler_mn[0] == 1: copy(tXrdX, tXgdX[None, None, None, bidx]) if const_expr(mdRes is not None): tXrdRes.store(dx.to(tXrdRes.element_type)) if row < M or tiler_mn[0] == 1: copy(tXrdRes, tXgdRes[None, None, None, bidx]) if const_expr(mdW is not None): tXrdW.store(tXrdW.load() + dout * x_hat) if const_expr(mdB is not None): tXrdB.store(tXrdB.load() + dout) stage ^= 1 if stage == 0: consumer_phase ^= 1 producer_phase ^= 1 if const_expr(tiler_mn[0] > 1): if const_expr(mdW is not None): # reduction of dw_partial within the same threadblock sdW = cute.make_tensor( cute.recast_ptr(sX.iterator, dtype=cute.Float32), cute.make_ordered_layout(tiler_mn, order=(1, 0)), ) tXsdW = thr_copy_X.partition_D(sdW) cute.arch.barrier() row = tXcX[None, None, None, 0][0][0] if row > 0: cute.autovec_copy(tXrdW, tXsdW) cute.arch.barrier() if row == 0: for i in cutlass.range_constexpr(1, const_expr(tiler_mn[0])): tXrdW_other = cute.make_fragment_like(tXrdW) tXsdW_other = cute.make_tensor( tXsdW.iterator + i * sdW.stride[0], tXsdW.layout ) cute.autovec_copy(tXsdW_other, tXrdW_other) tXrdW.store(tXrdW.load() + tXrdW_other.load()) copy(tXrdW, tXgdW) cute.arch.barrier() if const_expr(mdB is not None): sdB = cute.make_tensor( cute.recast_ptr(sX.iterator, dtype=cute.Float32), cute.make_ordered_layout(tiler_mn, order=(1, 0)), ) tXsdB = thr_copy_X.partition_D(sdB) cute.arch.barrier() row = tXcX[None, None, None, 0][0][0] if row > 0: cute.autovec_copy(tXrdB, tXsdB) cute.arch.barrier() if row == 0: for i in cutlass.range_constexpr(1, const_expr(tiler_mn[0])): tXrdB_other = cute.make_fragment_like(tXrdB) tXsdB_other = cute.make_tensor( tXsdB.iterator + i * sdB.stride[0], tXsdB.layout ) cute.autovec_copy(tXsdB_other, tXrdB_other) tXrdB.store(tXrdB.load() + tXrdB_other.load()) copy(tXrdB, tXgdB) else: # dw is already in fp32, so we can directly copy to global memory if const_expr(mdW is not None): copy(tXrdW, tXgdW) if const_expr(mdB is not None): copy(tXrdB, tXgdB) if const_expr(self.cluster_n > 1): # Prevent cluster from exiting early # Assume state contains that next useful buffer # So we only need to advance to num_stages - 1 times to last used buffer stage ^= 1 if stage == 0: producer_phase ^= 1 cute.arch.mbarrier_wait(mbar_empty_ptr + stage, producer_phase) def _get_sm_count(N: int, device: torch.device) -> int: # This should be tuned on how many CTAs can be launched on each SM sm_count_multiple = ( 16 if N <= 256 else (8 if N <= 1024 else (4 if N <= 2048 else (2 if N <= 4096 else 1))) ) sm_count = torch.cuda.get_device_properties(device).multi_processor_count # By right, if we're using cluster, this should be cluster_count not sm_count. # But for cluster >= 4, due to quantization we would need to query active max cluster. # Instead we just do sm_count * 2, which is reasonably larger than active_cluster_count to # avoid wave quantization. sm_count = ( sm_count * sm_count_multiple if N <= 8192 else sm_count // 2 if N <= 16384 else sm_count * 2 ) return sm_count @torch.library.custom_op( "quack::_rmsnorm_bwd", mutates_args={"dx", "dw_partial", "db_partial", "dresidual"}, device_types="cuda", # We need to specify the schema manually since we're mutating an optional tensor schema="(Tensor x, Tensor? weight, Tensor dout, Tensor rstd, Tensor(a4!) dx, Tensor(a5!)? dw_partial, Tensor(a6!)? db_partial, Tensor? dresidual_out, Tensor(a8!)? dresidual, int? sm_count) -> ()", ) def _rmsnorm_bwd( x: Tensor, weight: Optional[Tensor], dout: Tensor, rstd: Tensor, dx: Tensor, dw_partial: Optional[Tensor], db_partial: Optional[Tensor] = None, dresidual_out: Optional[Tensor] = None, dresidual: Optional[Tensor] = None, sm_count: Optional[int] = None, ) -> None: """RMSNorm backward pass. Args: x: Input tensor of shape (M, N) weight: Optional weight tensor of shape (N,) dout: Upstream gradients tensor of shape (M, N) rstd: Reciprocal standard deviation tensor of shape (M,) Returns: Tuple of (dx, dw) where: - dx: Input gradients tensor of same shape as x - dw: Weight gradients tensor of same shape as weight (or None if weight is None) """ assert x.dim() == 2, "Input must be 2D" assert x.is_cuda, "Input tensor must be on CUDA device" supported_types = {torch.float16, torch.bfloat16, torch.float32} assert x.dtype in supported_types, "Unsupported dtype" if weight is not None: assert weight.dim() == 1, "Weight must be 1D" assert x.shape[-1] == weight.shape[0], "Last dimension of input must match weight dimension" assert weight.is_cuda, "Weight tensor must be on CUDA device" assert weight.dtype in supported_types, "Weight must be float32, float16 or bfloat16" if dresidual_out is not None: assert dresidual_out.shape == x.shape assert dresidual_out.is_cuda assert dresidual_out.dtype in supported_types, ( "Residual must be float16, bfloat16, or float32" ) if dresidual is not None: assert dresidual.shape == x.shape assert dresidual.is_cuda assert dresidual.dtype in supported_types, "Residual must be float16, bfloat16, or float32" N = x.size(1) if dw_partial is None and db_partial is None: assert sm_count is not None else: sm_count = dw_partial.shape[0] if dw_partial is not None else db_partial.shape[0] dtype, dout_dtype, dx_dtype, weight_dtype, dres_dtype, dres_out_dtype = [ torch2cute_dtype_map[t.dtype] if t is not None else None for t in [x, dout, dx, weight, dresidual, dresidual_out] ] compile_key = ( N, dtype, dout_dtype, dx_dtype, weight_dtype, db_partial is not None, dres_dtype, dres_out_dtype, ) if compile_key not in _rmsnorm_bwd.compile_cache: batch_sym, batch_partial_sym = cute.sym_int(), cute.sym_int() all_dtypes = [dtype, dout_dtype, dx_dtype, dres_dtype, dres_out_dtype] div = math.gcd(N, *(128 // dt.width for dt in all_dtypes if dt is not None)) x_cute, dout_cute, dx_cute, dres_out_cute, dres_cute = [ fake_tensor(dt, (batch_sym, N), div) for dt in [dtype, dout_dtype, dx_dtype, dres_out_dtype, dres_dtype] ] weight_cute = fake_tensor(weight_dtype, (N,), div) rstd_cute = fake_tensor(Float32, (batch_sym,)) dw_partial_cute = ( fake_tensor(Float32, (batch_partial_sym, N), div) if dw_partial is not None else None ) db_partial_cute = ( fake_tensor(Float32, (batch_partial_sym, N), div) if db_partial is not None else None ) _rmsnorm_bwd.compile_cache[compile_key] = cute.compile( RMSNormBackward(dtype, N), x_cute, weight_cute, dout_cute, dres_out_cute, rstd_cute, dx_cute, dw_partial_cute, dres_cute, db_partial_cute, sm_count, cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), options="--enable-tvm-ffi", ) _rmsnorm_bwd.compile_cache[compile_key]( x, weight, dout, dresidual_out, rstd, dx, dw_partial, dresidual, db_partial, sm_count ) _rmsnorm_bwd.compile_cache = {} def rmsnorm_bwd( x: Tensor, weight: Optional[Tensor], dout: Tensor, rstd: Tensor, dresidual_out: Optional[Tensor] = None, # grad wrt residual_out has_bias: bool = False, has_residual: bool = False, ) -> Tuple[Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor]]: device = x.device N = x.size(1) dx = torch.empty_like(x) if dresidual_out is not None and dresidual_out.dtype != dx.dtype: dresidual = torch.empty_like(x, dtype=dresidual_out.dtype) else: dresidual = None sm_count = _get_sm_count(N, device) if weight is not None: # Always store partial gradients in fp32 for numerical accuracy dw_partial = torch.empty(sm_count, N, device=device, dtype=torch.float32) else: dw_partial = None db_partial = torch.empty(sm_count, N, device=device, dtype=torch.float32) if has_bias else None _rmsnorm_bwd( x, weight, dout, rstd, dx, dw_partial, db_partial, dresidual_out, dresidual, sm_count ) # we have summed the partial gradients in fp32, now we convert back to the weight dtype dw = dw_partial.sum(dim=0).to(weight.dtype) if weight is not None else None db = db_partial.sum(dim=0).to(weight.dtype) if has_bias else None # dresidual is the same as dx in this case if has_residual and dresidual is None: dresidual = dx return dx, dw, db, dresidual class RMSNormFunction(torch.autograd.Function): @staticmethod def forward( ctx, x, weight, bias=None, residual=None, out_dtype=None, residual_dtype=None, eps=1e-6, prenorm=False, ): x_shape_og = x.shape # Flatten input x = x.reshape(-1, x.shape[-1]) if residual is not None: residual = residual.reshape(-1, residual.shape[-1]) need_grad = any(ctx.needs_input_grad[:3]) out, residual_out, rstd = rmsnorm_fwd( x, weight, bias=bias, residual=residual, out_dtype=out_dtype, residual_dtype=residual_dtype, eps=eps, store_rstd=need_grad, ) ctx.save_for_backward(x if residual is None else residual_out, weight, rstd) ctx.has_bias = bias is not None ctx.eps = eps ctx.x_shape_og = x_shape_og ctx.residual_dtype = residual.dtype if residual is not None else None ctx.prenorm = prenorm if residual_out is None or not prenorm: return out.reshape(x_shape_og) else: return out.reshape(x_shape_og), residual_out.reshape(x_shape_og) @staticmethod def backward(ctx, dout, *args): x, weight, rstd = ctx.saved_tensors has_bias = ctx.has_bias if ctx.prenorm and ctx.residual_dtype is not None: dresidual_out = args[0] dresidual_out = dresidual_out.reshape(-1, dresidual_out.shape[-1]) else: dresidual_out = None x_shape_og = ctx.x_shape_og # Reshape dout to match the flattened shape used in forward dout = dout.view(-1, dout.shape[-1]) dx, dw, db, dresidual = rmsnorm_bwd( x, weight, dout, rstd, dresidual_out, has_bias, has_residual=ctx.residual_dtype is not None, ) dx = dx.view(x_shape_og) if dresidual is not None: dresidual = dresidual.reshape(x_shape_og) return dx, dw, db, dresidual, *([None] * 4) def rmsnorm( x: Tensor, weight: Optional[Tensor] = None, bias: Optional[Tensor] = None, residual: Optional[Tensor] = None, out_dtype: Optional[torch.dtype] = None, residual_dtype: Optional[torch.dtype] = None, eps: float = 1e-6, prenorm: bool = False, ) -> Tensor: """RMSNorm with automatic differentiation support. Args: x: Input tensor of shape (M, N) weight: Optional weight tensor of shape (N,) eps: Small value for numerical stability Returns: Normalized output tensor of same shape as x """ return RMSNormFunction.apply(x, weight, bias, residual, out_dtype, residual_dtype, eps, prenorm) class QuackRMSNorm(torch.nn.RMSNorm): """RMSNorm module that behaves like torch.nn.RMSNorm. This class provides a drop-in replacement for torch.nn.RMSNorm that uses the quack.rmsnorm implementation under the hood. Args: dim (int): The dimension to normalize over eps (float, optional): A small constant for numerical stability. Default: 1e-6 Attributes: weight (torch.nn.Parameter): The learnable weight parameter eps (float): A small constant for numerical stability """ def __init__( self, dim: int, eps: float = 1e-6, elementwise_affine: bool = True, device=None, dtype=None ): super().__init__(dim, eps, elementwise_affine, device=device, dtype=dtype) def forward(self, x: Tensor) -> Tensor: """Apply RMSNorm to the input tensor. Args: x (Tensor): Input tensor Returns: Tensor: Normalized tensor """ return rmsnorm(x, self.weight, eps=self.eps) def layernorm_fwd( x: Tensor, weight: Tensor, bias: Optional[Tensor] = None, eps: float = 1e-6, return_rstd: bool = False, return_mean: bool = False, ): """LayerNorm forward pass using the unified RMSNorm/LayerNorm kernel. Args: x: Input tensor of shape (M, N) weight: Weight tensor of shape (N,). Must be float32. bias: Optional bias tensor of shape (N,). Must be float32. eps: Small value for numerical stability return_rstd: Whether to return the reciprocal standard deviation return_mean: Whether to return the mean Returns: Normalized output tensor of same shape as x If return_rstd is True, also returns rstd tensor of shape (M,) If return_mean is True, also returns mean tensor of shape (M,) """ assert x.dim() == 2, "Input must be 2D" assert weight.dim() == 1, "Weight must be 1D" assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported dtype" assert weight.dtype == torch.float32, "Weight must be float32" if bias is not None: assert bias.dim() == 1, "Bias must be 1D" assert bias.dtype == torch.float32, "Bias must be float32" M, N = x.shape device = x.device out = torch.empty_like(x) rstd = torch.empty(M, device=device, dtype=torch.float32) if return_rstd else None mean = torch.empty(M, device=device, dtype=torch.float32) if return_mean else None _rmsnorm_fwd(x, weight, out, bias, rstd, mean, None, None, eps, True) if return_rstd and return_mean: return out, rstd, mean elif return_rstd: return out, rstd elif return_mean: return out, mean return out def layernorm_ref(x: Tensor, w: Tensor, eps: float = 1e-6) -> Tensor: """Reference implementation for LayerNorm.""" x_f32 = x.float() return torch.nn.functional.layer_norm(x_f32, w.shape, w, None, eps).to(x.dtype) def layernorm_rstd_ref(x: torch.Tensor, eps: float = 1e-6): x_f32 = x.float() mean = x_f32.mean(dim=-1, keepdim=True) var = ((x_f32 - mean) ** 2).mean(dim=-1) return 1.0 / torch.sqrt(var + eps) def layernorm_mean_ref(x: torch.Tensor) -> torch.Tensor: return x.float().mean(dim=-1)