diff --git a/kernels/norm/layernorm_kernel.py b/kernels/norm/layernorm_kernel.py index b7be89f8a..c27cdd3cf 100644 --- a/kernels/norm/layernorm_kernel.py +++ b/kernels/norm/layernorm_kernel.py @@ -19,7 +19,12 @@ from flydsl.expr import math as fmath from flydsl.expr.vector import ReductionOp, full from flydsl.runtime.device import get_rocm_arch -from kernels.common.kernels_common import dtype_to_elem_type, get_warp_size +from kernels.common.kernels_common import atomic_add, dtype_to_elem_type, get_warp_size + +try: + import torch +except ImportError: + torch = None KERNEL_NAME = "layernorm" @@ -112,7 +117,7 @@ def _quant_dtype_max(dtype_str: str) -> float: raise ValueError(f"unsupported quant dtype: {dtype_str!r} (expected 'i8' or 'int8')") -def build_layernorm_module(N: int, dtype_str: str): +def build_layernorm_module(N: int, dtype_str: str, store_stats: bool = False, eps: float = EPS): arch = get_rocm_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") @@ -124,22 +129,44 @@ def build_layernorm_module(N: int, dtype_str: str): # ── GPU kernel ──────────────────────────────────────────────────────── @flyc.kernel def layernorm_kernel( - Input: fx.Tensor, - Gamma: fx.Tensor, - Beta: fx.Tensor, - Output: fx.Tensor, + Input: fx.Pointer, + Gamma: fx.Pointer, + Beta: fx.Pointer, + Output: fx.Pointer, + Mean: fx.Pointer, + Rstd: fx.Pointer, ): bid = fx.block_idx.x tid = fx.thread_idx.x elem_dtype = dtype_to_elem_type(dtype_str) fm_fast = arith.FastMathFlags.fast - eps_c = EPS + eps_c = eps lds = fx.SharedAllocator().allocate(SharedStorage).peek() s_sum = lds.s_sum.view(fx.make_layout(RED_SLOTS, 1)) s_sumsq = lds.s_sumsq.view(fx.make_layout(RED_SLOTS, 1)) + # The wrapper guarantees contiguous [M, N] rows and contiguous [N] + # affine parameters. Reconstruct only this block's static views so the + # kernel ABI carries one address per argument, without shape/stride + # metadata for layouts already fixed by the contract. + row_layout = fx.make_layout(N, 1) + scalar_layout = fx.make_layout(1, 1) + bid_i64 = fx.Int64(bid) + row_offset = bid_i64 * N + Input_buf = fx.rocdl.make_buffer_tensor((Input + row_offset).view(row_layout)) + Gamma_buf = fx.rocdl.make_buffer_tensor(Gamma.view(row_layout)) + Beta_buf = fx.rocdl.make_buffer_tensor(Beta.view(row_layout)) + Output_buf = fx.rocdl.make_buffer_tensor((Output + row_offset).view(row_layout)) + + if const_expr(store_stats): + mean_row = fx.rocdl.make_buffer_tensor((Mean + bid_i64).view(scalar_layout)) + rstd_row = fx.rocdl.make_buffer_tensor((Rstd + bid_i64).view(scalar_layout)) + mean_div = fx.logical_divide(mean_row, scalar_layout) + rstd_div = fx.logical_divide(rstd_row, scalar_layout) + stats_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 32) + # ── helpers: wave / block reduction ─────────────────────────────── def wave_reduce_add(x): w = x @@ -206,16 +233,8 @@ def compute_mean_rstd(sum_val, sumsq_val): in_local = [] # ── Layout API: buffer-backed tensors + tiled access ───── - Input_buf = fx.rocdl.make_buffer_tensor(Input) - Output_buf = fx.rocdl.make_buffer_tensor(Output) - Gamma_buf = fx.rocdl.make_buffer_tensor(Gamma) - Beta_buf = fx.rocdl.make_buffer_tensor(Beta) - - row_in = fx.slice(Input_buf, (bid, None)) - row_out = fx.slice(Output_buf, (bid, None)) - - in_div = fx.logical_divide(row_in, fx.make_layout(VEC_WIDTH, 1)) - out_div = fx.logical_divide(row_out, fx.make_layout(VEC_WIDTH, 1)) + in_div = fx.logical_divide(Input_buf, fx.make_layout(VEC_WIDTH, 1)) + out_div = fx.logical_divide(Output_buf, fx.make_layout(VEC_WIDTH, 1)) gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(VEC_WIDTH, 1)) beta_div = fx.logical_divide(Beta_buf, fx.make_layout(VEC_WIDTH, 1)) @@ -237,6 +256,11 @@ def compute_mean_rstd(sum_val, sumsq_val): sum_val, sumsq_val = block_reduce_add2(thread_sum, thread_sumsq) mean, rstd = compute_mean_rstd(sum_val, sumsq_val) + if const_expr(store_stats): + if tid == 0: + _store_scalar(stats_copy_atom, fx.Float32, fx.Float32, mean_div, 0, mean) + _store_scalar(stats_copy_atom, fx.Float32, fx.Float32, rstd_div, 0, rstd) + g_cur = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, tid).to(fx.Float32) b_cur = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, beta_div, tid).to(fx.Float32) @@ -267,14 +291,6 @@ def compute_mean_rstd(sum_val, sumsq_val): # ============================================================== # Generic path: 2-pass scalar implementation for arbitrary N # ============================================================== - Input_buf = fx.rocdl.make_buffer_tensor(Input) - Output_buf = fx.rocdl.make_buffer_tensor(Output) - Gamma_buf = fx.rocdl.make_buffer_tensor(Gamma) - Beta_buf = fx.rocdl.make_buffer_tensor(Beta) - - row_in = fx.slice(Input_buf, (bid, None)) - row_out = fx.slice(Output_buf, (bid, None)) - c_zero_f = fx.Float32(0.0) thread_sum = c_zero_f thread_sumsq = c_zero_f @@ -284,10 +300,10 @@ def compute_mean_rstd(sum_val, sumsq_val): elem_bits, ) - row_div = fx.logical_divide(row_in, fx.make_layout(1, 1)) + row_div = fx.logical_divide(Input_buf, fx.make_layout(1, 1)) gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(1, 1)) beta_div = fx.logical_divide(Beta_buf, fx.make_layout(1, 1)) - out_div = fx.logical_divide(row_out, fx.make_layout(1, 1)) + out_div = fx.logical_divide(Output_buf, fx.make_layout(1, 1)) # ── Pass 1: sum + sumsq ────────────────────────────────────── for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): @@ -305,6 +321,11 @@ def compute_mean_rstd(sum_val, sumsq_val): sum_val, sumsq_val = block_reduce_add2(thread_sum, thread_sumsq) mean, rstd = compute_mean_rstd(sum_val, sumsq_val) + if const_expr(store_stats): + if tid == 0: + _store_scalar(stats_copy_atom, fx.Float32, fx.Float32, mean_div, 0, mean) + _store_scalar(stats_copy_atom, fx.Float32, fx.Float32, rstd_div, 0, rstd) + # ── Pass 2: normalize + affine + store ─────────────────────── for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): idx = tid + base_idx_int @@ -323,16 +344,38 @@ def compute_mean_rstd(sum_val, sumsq_val): _store_scalar(copy_atom_s, elem_dtype, elem_dtype, out_div, idx, y_e) # ── JIT host launcher ───────────────────────────────────────────────── + if store_stats: + + @flyc.jit + def launch_layernorm( + Input: fx.Pointer, + Gamma: fx.Pointer, + Beta: fx.Pointer, + Output: fx.Pointer, + Mean: fx.Pointer, + Rstd: fx.Pointer, + m_in: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + launcher = layernorm_kernel(Input, Gamma, Beta, Output, Mean, Rstd) + launcher.launch( + grid=(m_in, 1, 1), + block=(BLOCK_THREADS, 1, 1), + stream=stream, + ) + + return launch_layernorm + @flyc.jit def launch_layernorm( - Input: fx.Tensor, - Gamma: fx.Tensor, - Beta: fx.Tensor, - Output: fx.Tensor, + Input: fx.Pointer, + Gamma: fx.Pointer, + Beta: fx.Pointer, + Output: fx.Pointer, m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): - launcher = layernorm_kernel(Input, Gamma, Beta, Output) + launcher = layernorm_kernel(Input, Gamma, Beta, Output, Gamma, Gamma) launcher.launch( grid=(m_in, 1, 1), block=(BLOCK_THREADS, 1, 1), @@ -342,6 +385,181 @@ def launch_layernorm( return launch_layernorm +def build_layernorm_bwd_module(N: int, dtype_str: str): + """Fused LayerNorm backward: grid=(M,), one block per row. + + With x_hat = (x - mean)*rstd, wdy = dy*gamma: + c1 = mean_N(wdy) ; c2 = mean_N(wdy * x_hat) + Pass 2: + dx = (wdy - c1 - x_hat*c2) * rstd -> DX (elem dtype); + dgamma_elem = dy * x_hat (fp32) -> atomicAdd into DGamma[idx]; + dbias_elem = dy (fp32) -> atomicAdd into DBias[idx]. + eps is baked into Rstd by the forward, so it is not needed here. + + Perf follow-ups (deferred; correctness-complete as-is): generic scalar path + only — a vectorized fast path (mirroring the forward) and caching x/dy/gamma + between pass 1 and pass 2 would cut global traffic. + """ + 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 layernorm_bwd_kernel( + Input: fx.Pointer, + Gamma: fx.Pointer, + DY: fx.Pointer, + Mean: fx.Pointer, + Rstd: fx.Pointer, + DX: fx.Pointer, + DGamma: fx.Pointer, + DBias: fx.Pointer, + ): + 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_sum = lds.s_sum.view(fx.make_layout(RED_SLOTS, 1)) + s_sumsq = lds.s_sumsq.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_add2(val0, val1): + if const_expr(RED_SLOTS == 1): + return wave_reduce_add(val0), wave_reduce_add(val1) + + lane = tid % WARP_SIZE + wave = tid // WARP_SIZE + + w0 = wave_reduce_add(val0) + w1 = wave_reduce_add(val1) + + if lane == 0: + fx.memref_store(w0, s_sum, wave) + fx.memref_store(w1, s_sumsq, wave) + gpu.barrier() + + if wave == 0: + in_range = lane < RED_SLOTS + lane_safe = in_range.select(lane, 0) + v0 = fx.memref_load(s_sum, lane_safe) + v1 = fx.memref_load(s_sumsq, lane_safe) + ww0 = in_range.select(v0, 0.0) + ww1 = in_range.select(v1, 0.0) + ww0 = wave_reduce_add(ww0) + ww1 = wave_reduce_add(ww1) + + if lane == 0: + fx.memref_store(ww0, s_sum, 0) + fx.memref_store(ww1, s_sumsq, 0) + gpu.barrier() + + return fx.memref_load(s_sum, 0), fx.memref_load(s_sumsq, 0) + + # All layouts are fixed by the wrapper's contiguous contract. Rebuild + # only the row/scalar views used by this block so the kernel ABI carries + # one address per tensor and no dynamic shape/stride operands. + row_layout = fx.make_layout(N, 1) + scalar_layout = fx.make_layout(1, 1) + bid_i64 = fx.Int64(bid) + row_offset = bid_i64 * N + row_in = fx.rocdl.make_buffer_tensor((Input + row_offset).view(row_layout)) + gamma = fx.rocdl.make_buffer_tensor(Gamma.view(row_layout)) + row_dy = fx.rocdl.make_buffer_tensor((DY + row_offset).view(row_layout)) + mean_row = fx.rocdl.make_buffer_tensor((Mean + bid_i64).view(scalar_layout)) + rstd_row = fx.rocdl.make_buffer_tensor((Rstd + bid_i64).view(scalar_layout)) + row_dx = fx.rocdl.make_buffer_tensor((DX + row_offset).view(row_layout)) + + 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, fx.make_layout(1, 1)) + dx_div = fx.logical_divide(row_dx, fx.make_layout(1, 1)) + mean_div = fx.logical_divide(mean_row, scalar_layout) + rstd_div = fx.logical_divide(rstd_row, scalar_layout) + + mean = _load_scalar(copy_atom_f32, fx.Float32, mean_div, 0) + rstd = _load_scalar(copy_atom_f32, fx.Float32, rstd_div, 0) + dgamma_view = DGamma.view(fx.make_layout(N, 1)) + dbias_view = DBias.view(fx.make_layout(N, 1)) + + # Pass 1: c1 = mean(wdy) ; c2 = mean(wdy * x_hat) + thread_c1 = c_zero_f + thread_c2 = 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 - mean) * rstd + wdy = dy * g + thread_c1 = thread_c1 + is_valid.select(wdy, c_zero_f) + thread_c2 = thread_c2 + is_valid.select(wdy * x_hat, c_zero_f) + + sum_c1, sum_c2 = block_reduce_add2(thread_c1, thread_c2) + c1 = sum_c1 / n_float + c2 = sum_c2 / n_float + + # Pass 2: dx = (wdy - c1 - x_hat*c2) * rstd ; dgamma += dy*x_hat ; dbias += dy + 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 - mean) * rstd + wdy = dy * g + dx = (wdy - c1 - x_hat * c2) * 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) + + dgamma = dy * x_hat + atomic_add(dgamma_view, idx, dgamma, dtype_bytes=4) + atomic_add(dbias_view, idx, dy, dtype_bytes=4) + + @flyc.jit + def launch_layernorm_bwd( + Input: fx.Pointer, + Gamma: fx.Pointer, + DY: fx.Pointer, + Mean: fx.Pointer, + Rstd: fx.Pointer, + DX: fx.Pointer, + DGamma: fx.Pointer, + DBias: fx.Pointer, + m_in: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + launcher = layernorm_bwd_kernel(Input, Gamma, DY, Mean, Rstd, DX, DGamma, DBias) + launcher.launch(grid=(m_in, 1, 1), block=(BLOCK_THREADS, 1, 1), stream=stream) + + return launch_layernorm_bwd + + def build_fused_add_layernorm_module(N: int, dtype_str: str): arch = get_rocm_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") @@ -1334,3 +1552,197 @@ def build_fused_add_layernorm_smoothquant_module( is_smooth=True, quant_dtype_str=quant_dtype_str, ) + + +# ===================================================================== +# Python wrappers + autograd (quack-aligned). PR 3: plain layernorm. +# ===================================================================== +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. 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, bias, out, mean, rstd, M, N, dtype_str, store_stats, eps, stream): + key = (N, dtype_str, store_stats, float(eps), x.device) + entry = _FWD_CACHE.get(key) + if entry is None: + launch_fn = build_layernorm_module(N, dtype_str, store_stats=store_stats, eps=eps) + elem_dtype = dtype_to_elem_type(dtype_str) + x_ptr = flyc.from_c_void_p(elem_dtype, x.data_ptr()) + weight_ptr = flyc.from_c_void_p(elem_dtype, weight.data_ptr()) + bias_ptr = flyc.from_c_void_p(elem_dtype, bias.data_ptr()) + out_ptr = flyc.from_c_void_p(elem_dtype, out.data_ptr()) + if store_stats: + mean_ptr = flyc.from_c_void_p(fx.Float32, mean.data_ptr()) + rstd_ptr = flyc.from_c_void_p(fx.Float32, rstd.data_ptr()) + compiled = flyc.compile( + launch_fn, + x_ptr, + weight_ptr, + bias_ptr, + out_ptr, + mean_ptr, + rstd_ptr, + M, + stream, + ) + else: + compiled = flyc.compile(launch_fn, x_ptr, weight_ptr, bias_ptr, out_ptr, M, stream) + _FWD_CACHE[key] = compiled + entry = compiled + return entry + + def layernorm_fwd(x, weight, bias, eps=EPS, store_stats=False): + """Forward LayerNorm. Returns (out, mean, rstd). eps is baked into the kernel.""" + assert x.dim() == 2, "layernorm_fwd expects a 2D (M, N) input" + assert ( + x.is_contiguous() and weight.is_contiguous() and bias.is_contiguous() + ), "layernorm_fwd expects contiguous inputs" + # The kernel reads x/weight/bias with a single elem dtype derived from x; + # a mismatch would silently bit-reinterpret weight/bias bytes. + assert ( + x.dtype == weight.dtype == bias.dtype + ), f"x/weight/bias dtypes must match, got {x.dtype}/{weight.dtype}/{bias.dtype}" + assert weight.device == x.device and bias.device == x.device, "x/weight/bias must be on the same device" + M, N = x.shape + assert weight.dim() == 1 and bias.dim() == 1, "weight/bias must be 1D affine vectors" + assert weight.shape[0] == N and bias.shape[0] == N, "weight/bias length must equal x last dim (N)" + out = torch.empty_like(x) + mean = torch.empty((M,), device=x.device, dtype=torch.float32) if store_stats else None + rstd = torch.empty((M,), device=x.device, dtype=torch.float32) if store_stats else None + dtype_str = _torch_dtype_to_str(x.dtype) + # 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, bias, out, mean, rstd, M, N, dtype_str, store_stats, eps, stream) + if store_stats: + compiled( + x.data_ptr(), + weight.data_ptr(), + bias.data_ptr(), + out.data_ptr(), + mean.data_ptr(), + rstd.data_ptr(), + M, + stream, + ) + else: + compiled(x.data_ptr(), weight.data_ptr(), bias.data_ptr(), out.data_ptr(), M, stream) + return out, mean, rstd + + def layernorm_bwd(x, weight, dout, mean, rstd, eps=EPS): + """Backward LayerNorm. Returns (dx, dweight, dbias) cast to weight dtype. + + eps is not used directly here — it is already baked into `rstd`/`mean` by + the forward — but is accepted so callers can pass it symmetrically. + """ + assert x.dim() == 2, "layernorm_bwd expects a 2D (M, N) input" + assert all(t.is_contiguous() for t in (x, weight, dout, mean, rstd)), "layernorm_bwd expects contiguous inputs" + assert ( + x.dtype == weight.dtype == dout.dtype + ), f"x/weight/dout dtypes must match, got {x.dtype}/{weight.dtype}/{dout.dtype}" + M, N = x.shape + assert dout.shape == x.shape, "dout shape must equal x shape" + assert weight.dim() == 1 and weight.shape[0] == N, "weight must be a length-N 1D affine vector" + assert mean.numel() == M and rstd.numel() == M, "mean/rstd length must equal x rows (M)" + assert all( + t.device == x.device for t in (weight, dout, mean, rstd) + ), "x/weight/dout/mean/rstd must be on the same device" + # mean/rstd are per-row fp32 stats saved by the forward; the kernel reads + # them with an fp32 copy atom, so a non-fp32 dtype would be misread. + assert mean.dtype == torch.float32 and rstd.dtype == torch.float32, "mean/rstd must be fp32" + dtype_str = _torch_dtype_to_str(x.dtype) + elem_dtype = dtype_to_elem_type(dtype_str) + dx = torch.empty_like(x) + dweight = torch.zeros((N,), device=x.device, dtype=torch.float32) + dbias = torch.zeros((N,), device=x.device, dtype=torch.float32) + 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_layernorm_bwd_module(N, dtype_str) + # flyc.compile executes the kernel once during tracing, which would + # accumulate into DGamma/DBias; zero them AFTER compiling. + x_ptr = flyc.from_c_void_p(elem_dtype, x.data_ptr()) + weight_ptr = flyc.from_c_void_p(elem_dtype, weight.data_ptr()) + dout_ptr = flyc.from_c_void_p(elem_dtype, dout.data_ptr()) + mean_ptr = flyc.from_c_void_p(fx.Float32, mean.data_ptr()) + rstd_ptr = flyc.from_c_void_p(fx.Float32, rstd.data_ptr()) + dx_ptr = flyc.from_c_void_p(elem_dtype, dx.data_ptr()) + dweight_ptr = flyc.from_c_void_p(fx.Float32, dweight.data_ptr()) + dbias_ptr = flyc.from_c_void_p(fx.Float32, dbias.data_ptr()) + compiled = flyc.compile( + launch_fn, + x_ptr, + weight_ptr, + dout_ptr, + mean_ptr, + rstd_ptr, + dx_ptr, + dweight_ptr, + dbias_ptr, + M, + stream, + ) + _BWD_CACHE[key] = compiled + dweight.zero_() + dbias.zero_() + compiled( + x.data_ptr(), + weight.data_ptr(), + dout.data_ptr(), + mean.data_ptr(), + rstd.data_ptr(), + dx.data_ptr(), + dweight.data_ptr(), + dbias.data_ptr(), + M, + stream, + ) + return dx, dweight.to(weight.dtype), dbias.to(weight.dtype) + + class LayerNormFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, x, weight, bias, eps): + need_grad = x.requires_grad or weight.requires_grad or bias.requires_grad + out, mean, rstd = layernorm_fwd(x, weight, bias, eps=eps, store_stats=need_grad) + ctx.save_for_backward(x, weight, mean, rstd) + ctx.eps = eps + return out + + @staticmethod + def backward(ctx, dout): + x, weight, mean, rstd = ctx.saved_tensors + dx, dw, db = layernorm_bwd(x, weight, dout.contiguous(), mean, rstd, eps=ctx.eps) + return dx, dw, db, None + + def layernorm(x, weight, bias, eps=EPS): + """Public entry: plain LayerNorm with autograd.""" + assert weight is not None and bias is not None, "layernorm requires explicit weight and bias" + N = weight.shape[-1] + assert x.shape[-1] == N, f"x last dim {x.shape[-1]} != weight length {N}" + assert ( + x.dtype == weight.dtype == bias.dtype + ), f"x/weight/bias dtypes must match, got {x.dtype}/{weight.dtype}/{bias.dtype}" + # Raw-pointer kernels reconstruct dense rows and stride-1 affine + # vectors, so materialize those layouts at the public boundary. + x_flat = x.reshape(-1, N).contiguous() + weight_flat = weight.contiguous() + bias_flat = bias.contiguous() + out_flat = LayerNormFunction.apply(x_flat, weight_flat, bias_flat, eps) + return out_flat.reshape(x.shape) diff --git a/tests/kernels/benchmark_common.py b/tests/kernels/benchmark_common.py index 591861170..8770f3915 100644 --- a/tests/kernels/benchmark_common.py +++ b/tests/kernels/benchmark_common.py @@ -196,15 +196,31 @@ def _bench_flydsl_torch(*, op: str, M: int, N: int, dtype: str, warmup: int, ite return bench_gpu_us_torch(lambda: exe(x, y, M), warmup=warmup, iters=iters) if op == "layernorm": + import flydsl.compiler as flyc + import flydsl.expr as fx from kernels.norm.layernorm_kernel import build_layernorm_module m = build_layernorm_module(N, dtype) - exe = flydsl.compile(m) x = torch.randn((M, N), device="cuda", dtype=torch_dtype) gamma = torch.randn((N,), device="cuda", dtype=torch_dtype) beta = torch.randn((N,), device="cuda", dtype=torch_dtype) y = torch.empty((M, N), device="cuda", dtype=torch_dtype) - return bench_gpu_us_torch(lambda: exe(x, gamma, beta, y, M), warmup=warmup, iters=iters) + elem_dtype = {"f32": fx.Float32, "f16": fx.Float16, "bf16": fx.BFloat16}[dtype] + stream = torch.cuda.current_stream() + exe = flyc.compile( + m, + flyc.from_c_void_p(elem_dtype, x.data_ptr()), + flyc.from_c_void_p(elem_dtype, gamma.data_ptr()), + flyc.from_c_void_p(elem_dtype, beta.data_ptr()), + flyc.from_c_void_p(elem_dtype, y.data_ptr()), + M, + stream, + ) + return bench_gpu_us_torch( + lambda: exe(x.data_ptr(), gamma.data_ptr(), beta.data_ptr(), y.data_ptr(), M, stream), + warmup=warmup, + iters=iters, + ) if op == "rmsnorm": from kernels.norm.rmsnorm_kernel import build_rmsnorm_module diff --git a/tests/kernels/test_layernorm.py b/tests/kernels/test_layernorm.py index 2e00795f9..fe7022cd9 100644 --- a/tests/kernels/test_layernorm.py +++ b/tests/kernels/test_layernorm.py @@ -18,13 +18,16 @@ import pytest import flydsl.compiler as flyc +import flydsl.expr as fx from kernels.norm.layernorm_kernel import ( build_fused_add_layernorm_dynamicquant_module, build_fused_add_layernorm_module, build_fused_add_layernorm_smoothquant_module, + build_layernorm_bwd_module, build_layernorm_dynamicquant_module, build_layernorm_module, build_layernorm_smoothquant_module, + layernorm, ) from tests.kernels.benchmark_common import ( PerfRow, @@ -64,6 +67,20 @@ def _torch_dtype(dtype: str): raise ValueError(f"unsupported dtype: {dtype}") +def _flydsl_elem_dtype(dtype: str): + if dtype == "f32": + return fx.Float32 + if dtype == "f16": + return fx.Float16 + if dtype == "bf16": + return fx.BFloat16 + raise ValueError(f"unsupported dtype: {dtype}") + + +def _as_pointer(tensor, elem_dtype): + return flyc.from_c_void_p(elem_dtype, tensor.data_ptr()) + + def _get_layernorm_configs(): shapes_env = os.environ.get("ROCDSL_LAYERNORM_SHAPES", "").strip() if shapes_env: @@ -126,10 +143,26 @@ def run_test(M: int, N: int, dtype: str = "f32"): print("Launching kernel...") stream = torch.cuda.current_stream() - compiled_fn = flyc.compile(launch_fn, input_dev, gamma_dev, beta_dev, output_dev, M, stream) + elem_dtype = _flydsl_elem_dtype(dtype) + compiled_fn = flyc.compile( + launch_fn, + _as_pointer(input_dev, elem_dtype), + _as_pointer(gamma_dev, elem_dtype), + _as_pointer(beta_dev, elem_dtype), + _as_pointer(output_dev, elem_dtype), + M, + stream, + ) def kernel_launch(): - compiled_fn(input_dev, gamma_dev, beta_dev, output_dev, M, stream) + compiled_fn( + input_dev.data_ptr(), + gamma_dev.data_ptr(), + beta_dev.data_ptr(), + output_dev.data_ptr(), + M, + stream, + ) # One run for correctness visibility, then benchmark via shared harness. kernel_launch() @@ -971,8 +1004,447 @@ def test_fused_add_layernorm_smoothquant(): raise SystemExit(1) +def _reference_layernorm_bwd(x_dev, weight_dev, bias_dev, dy_dev): + """Eager layernorm backward via autograd. Returns dx, dgamma, dbias, mean, rstd (all fp32).""" + x = x_dev.detach().to(DTYPE_FP32).requires_grad_(True) + w = weight_dev.detach().to(DTYPE_FP32).requires_grad_(True) + b = bias_dev.detach().to(DTYPE_FP32).requires_grad_(True) + mean = x.mean(dim=1, keepdim=True) + var = x.var(dim=1, keepdim=True, unbiased=False) + rstd = torch.rsqrt(var + EPS) + y = (x - mean) * rstd * w + b + dx, dgamma, dbias = torch.autograd.grad(y, [x, w, b], grad_outputs=dy_dev.to(DTYPE_FP32)) + return ( + dx.detach(), + dgamma.detach(), + dbias.detach(), + mean.detach().squeeze(1).contiguous(), + rstd.detach().squeeze(1).contiguous(), + ) + + +def run_layernorm_bwd_test(M: int, N: int, dtype: str = "f32"): + print(f"\nTesting LayerNorm backward (M={M}, N={N}, dtype={dtype})") + + torch_dtype = _torch_dtype(dtype) + try: + fwd_fn = build_layernorm_module(N, dtype, store_stats=True) + bwd_fn = build_layernorm_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() + bias = 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, dgamma_ref, dbias_ref, mean_ref, rstd_ref = _reference_layernorm_bwd(x, weight, bias, dy) + + stream = torch.cuda.current_stream() + + # --- forward with store_stats: validates mean + rstd from the kernel --- + out = torch.empty((M, N), device="cuda", dtype=torch_dtype) + mean = torch.empty((M,), device="cuda", dtype=DTYPE_FP32) + rstd = torch.empty((M,), device="cuda", dtype=DTYPE_FP32) + elem_dtype = _flydsl_elem_dtype(dtype) + fwd_c = flyc.compile( + fwd_fn, + _as_pointer(x, elem_dtype), + _as_pointer(weight, elem_dtype), + _as_pointer(bias, elem_dtype), + _as_pointer(out, elem_dtype), + _as_pointer(mean, fx.Float32), + _as_pointer(rstd, fx.Float32), + M, + stream, + ) + fwd_c( + x.data_ptr(), + weight.data_ptr(), + bias.data_ptr(), + out.data_ptr(), + mean.data_ptr(), + rstd.data_ptr(), + M, + stream, + ) + torch.cuda.synchronize() + mean_err = (mean - mean_ref).abs().max().item() + rstd_err = (rstd - rstd_ref).abs().max().item() + + # --- backward: dx + dgamma + dbias --- + dx = torch.empty((M, N), device="cuda", dtype=torch_dtype) + dgamma = torch.zeros((N,), device="cuda", dtype=DTYPE_FP32) + dbias = torch.zeros((N,), device="cuda", dtype=DTYPE_FP32) + x_ptr = _as_pointer(x, elem_dtype) + weight_ptr = _as_pointer(weight, elem_dtype) + dy_ptr = _as_pointer(dy, elem_dtype) + mean_ptr = _as_pointer(mean, fx.Float32) + rstd_ptr = _as_pointer(rstd, fx.Float32) + dx_ptr = _as_pointer(dx, elem_dtype) + dgamma_ptr = _as_pointer(dgamma, fx.Float32) + dbias_ptr = _as_pointer(dbias, fx.Float32) + bwd_c = flyc.compile( + bwd_fn, + x_ptr, + weight_ptr, + dy_ptr, + mean_ptr, + rstd_ptr, + dx_ptr, + dgamma_ptr, + dbias_ptr, + M, + stream, + ) + dgamma.zero_() + dbias.zero_() + bwd_c( + x.data_ptr(), + weight.data_ptr(), + dy.data_ptr(), + mean.data_ptr(), + rstd.data_ptr(), + dx.data_ptr(), + dgamma.data_ptr(), + dbias.data_ptr(), + M, + stream, + ) + torch.cuda.synchronize() + + dx_err = (dx.to(DTYPE_FP32) - dx_ref).abs().max().item() + + # Tolerances (calibrated). dgamma/dbias summed over M -> larger magnitude -> relative. + stat_atol = 1e-3 + dx_atol = {"f32": 1e-3, "f16": 3e-2, "bf16": 2e-1}[dtype] + dg_rtol = {"f32": 1e-4, "f16": 3e-2, "bf16": 1e-1}[dtype] + dg_atol = {"f32": 1e-2, "f16": 1e-1, "bf16": 5e-1}[dtype] + + print(f" mean max abs err = {mean_err:.3e} (atol={stat_atol})") + print(f" rstd max abs err = {rstd_err:.3e} (atol={stat_atol})") + print(f" dx max abs err = {dx_err:.3e} (atol={dx_atol})") + + dg_ok = True + try: + torch.testing.assert_close(dgamma, dgamma_ref, rtol=dg_rtol, atol=dg_atol) + except AssertionError as e: + dg_ok = False + dg_err = (dgamma - dgamma_ref).abs().max().item() + print(f" dgamma max abs err = {dg_err:.3e} (rtol={dg_rtol}, atol={dg_atol})") + print(f" [dgamma mismatch] {e}") + else: + dg_err = (dgamma - dgamma_ref).abs().max().item() + print(f" dgamma max abs err = {dg_err:.3e} (rtol={dg_rtol}, atol={dg_atol})") + + db_ok = True + try: + torch.testing.assert_close(dbias, dbias_ref, rtol=dg_rtol, atol=dg_atol) + except AssertionError as e: + db_ok = False + db_err = (dbias - dbias_ref).abs().max().item() + print(f" dbias max abs err = {db_err:.3e} (rtol={dg_rtol}, atol={dg_atol})") + print(f" [dbias mismatch] {e}") + else: + db_err = (dbias - dbias_ref).abs().max().item() + print(f" dbias max abs err = {db_err:.3e} (rtol={dg_rtol}, atol={dg_atol})") + + ok = mean_err < stat_atol and rstd_err < stat_atol and dx_err < dx_atol and dg_ok and db_ok + print(f" -> {'PASSED' if ok else 'FAILED'}") + return ok + + +def run_layernorm_autograd_test(M: int, N: int, dtype: str = "f32"): + """End-to-end: public layernorm() grads on x + weight + bias, incl. 3D reshape.""" + print(f"\nTesting layernorm() 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) + bias = 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 = layernorm(x, weight, bias) + out.backward(dy) + dx_out, dw_out, db_out = x.grad.detach(), weight.grad.detach(), bias.grad.detach() + + # fp32 autograd reference + xf = x.detach().to(DTYPE_FP32).requires_grad_(True) + wf = weight.detach().to(DTYPE_FP32).requires_grad_(True) + bf = bias.detach().to(DTYPE_FP32).requires_grad_(True) + mean = xf.mean(dim=1, keepdim=True) + var = xf.var(dim=1, keepdim=True, unbiased=False) + rstd = torch.rsqrt(var + EPS) + yr = (xf - mean) * rstd * wf + bf + dxr, dwr, dbr = torch.autograd.grad(yr, [xf, wf, bf], 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() + db_err = (db_out.to(DTYPE_FP32) - dbr).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})") + print(f" db max abs err = {db_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 = layernorm(x3, weight, bias) + 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 db_err < dw_atol and shape_ok and grad_ok + print(f" -> {'PASSED' if ok else 'FAILED'}") + return ok + + +def test_layernorm_backward(): + print("=" * 80) + print("Running LayerNorm Backward Tests") + print("=" * 80) + + configs = [ + (64, 256, "f32"), # generic path, f32 aligned + (16, 512, "bf16"), # generic 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: + if not run_layernorm_bwd_test(M, N, dtype): + 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 test_layernorm_autograd(): + print("=" * 80) + print("Running layernorm() Autograd (end-to-end) Tests") + print("=" * 80) + + configs = [ + (64, 256, "f32"), # generic path + (128, 4096, "bf16"), # generic path + (128, 4096, "f16"), # generic path, f16 + ] + + failures = 0 + for M, N, dtype in configs: + if not run_layernorm_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) + + +def test_layernorm_eps_honored(): + """eps must be baked into the kernel, not silently replaced by the module EPS.""" + print("=" * 80) + print("Running LayerNorm 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) + b = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + + for eps in (1e-5, 1e-6, 1e-2): + y = layernorm(x, w, b, eps=eps) + mean = x.mean(dim=1, keepdim=True) + var = x.var(dim=1, keepdim=True, unbiased=False) + ref = (x - mean) * torch.rsqrt(var + eps) * w + b + 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 = (layernorm(x, w, b, eps=1e-2) - layernorm(x, w, b, 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.parametrize("M,N,dtype", [(3, 257, "f32"), (2, 8192, "bf16")]) +def test_layernorm_pointer_storage_offset(M, N, dtype): + """Pointer row math must be relative to each contiguous view's data_ptr.""" + torch_dtype = _torch_dtype(dtype) + + def offset_rand(shape): + numel = 1 + for dim in shape: + numel *= dim + tensor = torch.randn((numel + 1,), device="cuda", dtype=torch_dtype)[1:].view(shape) + assert tensor.is_contiguous() and tensor.storage_offset() == 1 + return tensor + + x = offset_rand((M, N)).detach().requires_grad_(True) + weight = offset_rand((N,)).detach().requires_grad_(True) + bias = offset_rand((N,)).detach().requires_grad_(True) + dout = offset_rand((M, N)) + + out = layernorm(x, weight, bias) + out.backward(dout) + + out_ref = _reference_layernorm(x.detach(), weight.detach(), bias.detach()) + dx_ref, dweight_ref, dbias_ref, _, _ = _reference_layernorm_bwd(x.detach(), weight.detach(), bias.detach(), dout) + out_atol = {"f32": 1e-4, "bf16": 2e-2}[dtype] + dx_atol = {"f32": 1e-3, "bf16": 2e-1}[dtype] + grad_rtol = {"f32": 1e-4, "bf16": 1e-1}[dtype] + grad_atol = {"f32": 1e-2, "bf16": 1.0}[dtype] + torch.testing.assert_close(out.to(DTYPE_FP32), out_ref, rtol=grad_rtol, atol=out_atol) + torch.testing.assert_close(x.grad.to(DTYPE_FP32), dx_ref, rtol=grad_rtol, atol=dx_atol) + torch.testing.assert_close(weight.grad.to(DTYPE_FP32), dweight_ref, rtol=grad_rtol, atol=grad_atol) + torch.testing.assert_close(bias.grad.to(DTYPE_FP32), dbias_ref, rtol=grad_rtol, atol=grad_atol) + + +def test_layernorm_strided_input_and_affine_params(): + """Public layernorm materializes layouts required by its raw-pointer kernels.""" + torch.manual_seed(0) + M, N = 3, 257 + + x_storage = torch.randn((2 * M, N), device="cuda", dtype=DTYPE_FP32) + weight_storage = torch.randn((2 * N,), device="cuda", dtype=DTYPE_FP32) + # Keep N elements allocated so a layout regression reads in-bounds data + # instead of risking an out-of-bounds access in optimized Python mode. + bias_storage = torch.randn((N,), device="cuda", dtype=DTYPE_FP32) + + x = x_storage[::2].detach().requires_grad_(True) + weight = weight_storage[::2].detach().requires_grad_(True) + bias = bias_storage[:1].expand(N).detach().requires_grad_(True) + dout = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) + + assert x.stride() == (2 * N, 1) + assert weight.stride() == (2,) + assert bias.stride() == (0,) + assert not x.is_contiguous() + assert not weight.is_contiguous() + assert not bias.is_contiguous() + + out = layernorm(x, weight, bias) + out.backward(dout) + + x_ref = x.detach().clone().requires_grad_(True) + weight_ref = weight.detach().clone().requires_grad_(True) + bias_ref = bias.detach().clone().requires_grad_(True) + out_ref = torch.nn.functional.layer_norm(x_ref, (N,), weight_ref, bias_ref, EPS) + out_ref.backward(dout) + + torch.testing.assert_close(out, out_ref, rtol=1e-4, atol=1e-4) + torch.testing.assert_close(x.grad, x_ref.grad, rtol=1e-4, atol=1e-3) + torch.testing.assert_close(weight.grad, weight_ref.grad, rtol=1e-4, atol=1e-2) + torch.testing.assert_close(bias.grad, bias_ref.grad, rtol=1e-4, atol=1e-2) + + +@pytest.mark.multi_gpu +def test_layernorm_multi_gpu(): + """Compiled-fn cache must not reuse a device-0 kernel on device 1 (would fault).""" + print("=" * 80) + print("Running LayerNorm 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) + b = torch.rand((N,), device=dev, dtype=DTYPE_FP32, requires_grad=True) + dy = torch.randn((16, N), device=dev, dtype=DTYPE_FP32) + y = layernorm(x, w, b) + y.backward(dy) + torch.cuda.synchronize(dev) + mean = x.detach().mean(1, keepdim=True) + var = x.detach().var(1, keepdim=True, unbiased=False) + ref = (x.detach() - mean) * torch.rsqrt(var + EPS) * w.detach() + b.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.multi_gpu +def test_layernorm_device_mismatch(): + """Raw pointer launches must reject cross-device affine parameters.""" + if torch.cuda.device_count() < 2: + pytest.skip("needs >=2 GPUs") + + x = torch.randn((4, 256), device="cuda:0", dtype=DTYPE_FP32) + w0 = torch.rand((256,), device="cuda:0", dtype=DTYPE_FP32) + b0 = torch.rand((256,), device="cuda:0", dtype=DTYPE_FP32) + w1 = w0.to("cuda:1") + b1 = b0.to("cuda:1") + + with pytest.raises(AssertionError, match="same device"): + layernorm(x, w1, b0) + with pytest.raises(AssertionError, match="same device"): + layernorm(x, w0, b1) + + +def test_layernorm_dtype_mismatch(): + """Mixed x/weight/bias dtypes must be rejected (kernel uses one elem dtype).""" + print("=" * 80) + print("Running layernorm dtype-mismatch guard Test") + print("=" * 80) + N = 256 + x = torch.randn((4, N), device="cuda", dtype=DTYPE_BF16) + w = torch.rand((N,), device="cuda", dtype=DTYPE_BF16) + b = torch.rand((N,), device="cuda", dtype=DTYPE_BF16) + for bad_w, bad_b in ((w.to(DTYPE_FP16), b), (w, b.to(DTYPE_FP32))): + try: + layernorm(x, bad_w, bad_b) + raise SystemExit("dtype mismatch was NOT rejected") + except AssertionError: + pass + print(" -> PASSED") + + +def test_layernorm_affine_shape_mismatch(): + """Pointer views require exactly N affine elements, not merely a trailing N.""" + N = 256 + x = torch.randn((4, N), device="cuda", dtype=DTYPE_FP32) + weight = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + bias = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + + with pytest.raises(AssertionError, match="1D"): + layernorm(x, weight.view(1, N), bias) + with pytest.raises(AssertionError, match="1D"): + layernorm(x, weight, bias.view(1, N)) + with pytest.raises(AssertionError, match="1D"): + layernorm(x, weight.view(1, N).expand(2, N).contiguous(), bias) + with pytest.raises(AssertionError, match="1D"): + layernorm(x, weight, bias.view(1, N).expand(2, N).contiguous()) + + if __name__ == "__main__": test_layernorm() + test_layernorm_backward() + test_layernorm_autograd() + test_layernorm_eps_honored() + test_layernorm_multi_gpu() + test_layernorm_dtype_mismatch() test_fused_add_layernorm() test_layernorm_dynamicquant() test_layernorm_smoothquant()