diff --git a/kernels/norm/layernorm_kernel.py b/kernels/norm/layernorm_kernel.py index c27cdd3cf..685c4fc53 100644 --- a/kernels/norm/layernorm_kernel.py +++ b/kernels/norm/layernorm_kernel.py @@ -6,8 +6,8 @@ LayerNorm(x) = (x - mean) / sqrt(var + eps) * gamma + beta Two paths: - - Fast path (N == BLOCK_THREADS * VEC_WIDTH * 4): vectorised tiled copy, - register caching, pipelined gamma/beta loads. + - Fast path (N is a multiple of BLOCK_THREADS * VEC_WIDTH): vectorised + tiled copy, register caching, pipelined gamma/beta loads. - Generic path (arbitrary N): scalar 2-pass implementation. """ @@ -123,6 +123,7 @@ def build_layernorm_module(N: int, dtype_str: str, store_stats: bool = False, ep RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) elem_bits = 32 if dtype_str == "f32" else 16 + tile_cols = BLOCK_THREADS * VEC_WIDTH SharedStorage = _make_reduction_storage(RED_SLOTS) @@ -221,12 +222,12 @@ def compute_mean_rstd(sum_val, sumsq_val): return mean, rstd # ================================================================== - # Fast path: N == BLOCK_THREADS * VEC_WIDTH * 4 + # Fast path: N is a multiple of BLOCK_THREADS * VEC_WIDTH # Uses buffer_load / buffer_store for high-bandwidth vectorised # memory access (same approach as preshuffle_gemm). # ================================================================== - if const_expr(N == (BLOCK_THREADS * VEC_WIDTH * 4) and elem_bits <= 16): - num_tiles_py = 4 + if const_expr(N >= tile_cols and N % tile_cols == 0 and elem_bits <= 16): + num_tiles_py = N // tile_cols c_zero_f = fx.Float32(0.0) thread_sum = c_zero_f thread_sumsq = c_zero_f @@ -566,6 +567,7 @@ def build_fused_add_layernorm_module(N: int, dtype_str: str): RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) elem_bits = 32 if dtype_str == "f32" else 16 + tile_cols = BLOCK_THREADS * VEC_WIDTH SharedStorage = _make_reduction_storage(RED_SLOTS) @@ -583,7 +585,7 @@ def fused_add_layernorm_kernel( 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)) @@ -637,10 +639,10 @@ def compute_mean_rstd(sum_val, sumsq_val): return mean, fmath.rsqrt(var + eps_c, fastmath=fm_fast) # ================================================================== - # Fast path: N == BLOCK_THREADS * VEC_WIDTH * 4 + # Fast path: N is a multiple of BLOCK_THREADS * VEC_WIDTH # ================================================================== - if const_expr(N == (BLOCK_THREADS * VEC_WIDTH * 4) and elem_bits <= 16): - num_tiles_py = 4 + if const_expr(N >= tile_cols and N % tile_cols == 0 and elem_bits <= 16): + num_tiles_py = N // tile_cols c_zero_f = fx.Float32(0.0) thread_sum = c_zero_f thread_sumsq = c_zero_f @@ -786,9 +788,11 @@ def _build_layernorm_quant_module( *, is_smooth: bool, quant_dtype_str: str = "i8", + eps: float = EPS, ): RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) elem_bits = 32 if dtype_str == "f32" else 16 + tile_cols = BLOCK_THREADS * VEC_WIDTH quant_dtype_max = _quant_dtype_max(quant_dtype_str) SharedStorage = _make_reduction_storage(RED_SLOTS) @@ -809,7 +813,7 @@ def layernorm_quant_kernel( quant_dtype = _quant_dtype_to_elem_type(quant_dtype_str) fm_fast = arith.FastMathFlags.fast - eps_c = EPS + eps_c = eps n_float = float(N) c_zero_f = fx.Float32(0.0) c_one_f = fx.Float32(1.0) @@ -894,10 +898,10 @@ def block_reduce_max(val): return fx.memref_load(s_sum, 0) # ================================================================== - # Fast path: N == BLOCK_THREADS * VEC_WIDTH * 4 + # Fast path: N is a multiple of BLOCK_THREADS * VEC_WIDTH # ================================================================== - if const_expr(N == (BLOCK_THREADS * VEC_WIDTH * 4) and elem_bits <= 16): - num_tiles_py = 4 + if const_expr(N >= tile_cols and N % tile_cols == 0 and elem_bits <= 16): + num_tiles_py = N // tile_cols quant_half_width = VEC_WIDTH // 2 abs_mask = full(VEC_WIDTH, fx.Uint32(0x7FFFFFFF), fx.Uint32) @@ -1130,12 +1134,14 @@ def _build_fused_add_layernorm_quant_module( *, is_smooth: bool, quant_dtype_str: str = "i8", + eps: float = EPS, ): arch = get_rocm_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) elem_bits = 32 if dtype_str == "f32" else 16 + tile_cols = BLOCK_THREADS * VEC_WIDTH quant_dtype_max = _quant_dtype_max(quant_dtype_str) SharedStorage = _make_reduction_storage(RED_SLOTS) @@ -1158,7 +1164,7 @@ def fused_add_layernorm_quant_kernel( quant_dtype = _quant_dtype_to_elem_type(quant_dtype_str) fm_fast = arith.FastMathFlags.fast - eps_c = EPS + eps_c = eps n_float = float(N) c_zero_f = fx.Float32(0.0) c_one_f = fx.Float32(1.0) @@ -1243,10 +1249,10 @@ def block_reduce_max(val): return fx.memref_load(s_sum, 0) # ================================================================== - # Fast path: N == BLOCK_THREADS * VEC_WIDTH * 4 + # Fast path: N is a multiple of BLOCK_THREADS * VEC_WIDTH # ================================================================== - if const_expr(N == (BLOCK_THREADS * VEC_WIDTH * 4) and elem_bits <= 16): - num_tiles_py = 4 + if const_expr(N >= tile_cols and N % tile_cols == 0 and elem_bits <= 16): + num_tiles_py = N // tile_cols quant_half_width = VEC_WIDTH // 2 abs_mask = full(VEC_WIDTH, fx.Uint32(0x7FFFFFFF), fx.Uint32) @@ -1506,12 +1512,14 @@ def build_layernorm_dynamicquant_module( N: int, dtype_str: str, quant_dtype_str: str = "i8", + eps: float = EPS, ): return _build_layernorm_quant_module( N, dtype_str, is_smooth=False, quant_dtype_str=quant_dtype_str, + eps=eps, ) @@ -1519,12 +1527,14 @@ def build_layernorm_smoothquant_module( N: int, dtype_str: str, quant_dtype_str: str = "i8", + eps: float = EPS, ): return _build_layernorm_quant_module( N, dtype_str, is_smooth=True, quant_dtype_str=quant_dtype_str, + eps=eps, ) @@ -1532,12 +1542,14 @@ def build_fused_add_layernorm_dynamicquant_module( N: int, dtype_str: str, quant_dtype_str: str = "i8", + eps: float = EPS, ): return _build_fused_add_layernorm_quant_module( N, dtype_str, is_smooth=False, quant_dtype_str=quant_dtype_str, + eps=eps, ) @@ -1545,12 +1557,14 @@ def build_fused_add_layernorm_smoothquant_module( N: int, dtype_str: str, quant_dtype_str: str = "i8", + eps: float = EPS, ): return _build_fused_add_layernorm_quant_module( N, dtype_str, is_smooth=True, quant_dtype_str=quant_dtype_str, + eps=eps, ) diff --git a/kernels/norm/rmsnorm_kernel.py b/kernels/norm/rmsnorm_kernel.py index 6843615d9..4787bbed5 100644 --- a/kernels/norm/rmsnorm_kernel.py +++ b/kernels/norm/rmsnorm_kernel.py @@ -83,9 +83,13 @@ def build_rmsnorm_module( USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") # BLOCK_THREADS controls storage, tiling, and launch geometry. - tile_cols = BLOCK_THREADS * VEC_WIDTH RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) elem_bits = 32 if dtype_str == "f32" else 16 + USE_VEC_N = elem_bits <= 16 and N % VEC_WIDTH == 0 + VEC_TILES = N // VEC_WIDTH if USE_VEC_N else 0 + NUM_VEC_ITERS = ( + (VEC_TILES + BLOCK_THREADS - 1) // BLOCK_THREADS if USE_VEC_N else 0 + ) _kernel_kwargs = {} if BLOCK_THREADS <= 256 else {"known_block_size": [BLOCK_THREADS, 1, 1]} SharedStorage = _make_reduction_storage(RED_SLOTS) @@ -160,10 +164,9 @@ def block_reduce_add2(val0, val1): return fx.memref_load(s_red, 0), fx.memref_load(s_red2, 0) # ================================================================== - # Fast path: N is a multiple of tile_cols + # Fast path: vectorized full-row body for f16/bf16 N aligned to vec8 # ================================================================== - if const_expr(N >= tile_cols and N % tile_cols == 0 and elem_bits <= 16): - num_tiles = N // tile_cols + if const_expr(USE_VEC_N): # ── Layout API: buffer-backed tensors + tiled access ───── Input_buf = fx.rocdl.make_buffer_tensor(Input) Output_buf = fx.rocdl.make_buffer_tensor(Output) @@ -184,15 +187,17 @@ def block_reduce_add2(val0, val1): in_local = [] # Pass 1: load + cache + sumsq - for tile_i in range_constexpr(num_tiles): + for tile_i in range_constexpr(NUM_VEC_ITERS): idx = tid + tile_i * BLOCK_THREADS - vec = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx) + is_valid = idx < VEC_TILES + idx_safe = is_valid.select(idx, 0) + vec = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx_safe) in_local.append(vec) x = vec.to(fx.Float32) x2 = x * x red2 = x2.reduce(ReductionOp.ADD, fastmath=fm_fast) - thread_sumsq = thread_sumsq + red2 + thread_sumsq = thread_sumsq + is_valid.select(red2, c_zero_f) _, sum_sq = block_reduce_add2(thread_dummy, thread_sumsq) mean_sq = sum_sq / n_float @@ -204,17 +209,17 @@ def block_reduce_add2(val0, val1): _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): + for tile_i in range_constexpr(NUM_VEC_ITERS): idx = tid + tile_i * BLOCK_THREADS - g = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, idx).to(fx.Float32) - x = in_local[tile_i].to(fx.Float32) + if idx < VEC_TILES: + g = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, idx).to(fx.Float32) + x = in_local[tile_i].to(fx.Float32) - y = (x * rrms) * g - out_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, y) + y = (x * rrms) * g + out_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, y) - out_idx = tid + tile_i * BLOCK_THREADS - _store_vec(copy_atom, VEC_WIDTH, elem_dtype, out_e, out_div, out_idx) + _store_vec(copy_atom, VEC_WIDTH, elem_dtype, out_e, out_div, idx) else: # ============================================================== @@ -332,6 +337,15 @@ def _build_rmsnorm_large_m_small_n_module(N: int, dtype_str: str, store_rstd: bo THREADS_PER_ROW = min(WARP_SIZE, 1024 // BLOCK_M) BLOCK_THREADS_SPECIAL = BLOCK_M * THREADS_PER_ROW elem_bits = 32 if dtype_str == "f32" else 16 + USE_VEC_SMALL_N = elem_bits <= 16 and N >= VEC_WIDTH + VEC_TILES = N // VEC_WIDTH if USE_VEC_SMALL_N else 0 + VEC_ELEMS = VEC_TILES * VEC_WIDTH + TAIL_ELEMS = N - VEC_ELEMS + NUM_VEC_ROW_ITERS = ( + (VEC_TILES + THREADS_PER_ROW - 1) // THREADS_PER_ROW if USE_VEC_SMALL_N else 0 + ) + arch = get_rocm_arch() + USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") @flyc.kernel(known_block_size=[BLOCK_THREADS_SPECIAL, 1, 1]) def rmsnorm_large_m_small_n_kernel( @@ -365,15 +379,6 @@ def rmsnorm_large_m_small_n_kernel( row_in = fx.slice(Input_buf, (row, None)) row_out = fx.slice(Output_buf, (row, None)) - copy_atom_s = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - elem_bits, - ) - - row_div = fx.logical_divide(row_in, fx.make_layout(1, 1)) - gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(1, 1)) - out_div = fx.logical_divide(row_out, fx.make_layout(1, 1)) - def group_reduce_add(x): w = x for _sh_exp in range_constexpr(int(math.log2(THREADS_PER_ROW))): @@ -385,34 +390,104 @@ def group_reduce_add(x): c_zero_f = fx.Float32(0.0) thread_sumsq = c_zero_f - for base_idx_int in range_constexpr(0, BLOCK_N, THREADS_PER_ROW): - idx = lane + base_idx_int - is_valid = idx < N - idx_safe = is_valid.select(idx, 0) - x_e = _load_scalar(copy_atom_s, elem_dtype, row_div, idx_safe) - x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) - x2 = x * x - thread_sumsq = thread_sumsq + is_valid.select(x2, c_zero_f) - - sum_sq = group_reduce_add(thread_sumsq) - mean_sq = sum_sq / n_float - ms_eps = mean_sq + eps_c - rrms = fmath.rsqrt(ms_eps, fastmath=fm_fast) + if const_expr(USE_VEC_SMALL_N): + in_div = fx.logical_divide(row_in, fx.make_layout(VEC_WIDTH, 1)) + gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(VEC_WIDTH, 1)) + out_div = fx.logical_divide(row_out, fx.make_layout(VEC_WIDTH, 1)) + + copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) + if const_expr(TAIL_ELEMS > 0): + row_div_s = fx.logical_divide(row_in, fx.make_layout(1, 1)) + gamma_div_s = fx.logical_divide(Gamma_buf, fx.make_layout(1, 1)) + out_div_s = fx.logical_divide(row_out, fx.make_layout(1, 1)) + copy_atom_s = fx.make_copy_atom(fx.rocdl.BufferCopy16b(), elem_bits) + + in_local = [] + for tile_i in range_constexpr(NUM_VEC_ROW_ITERS): + idx = lane + tile_i * THREADS_PER_ROW + is_valid = idx < VEC_TILES + idx_safe = is_valid.select(idx, 0) + vec = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx_safe) + in_local.append(vec) + x = vec.to(fx.Float32) + + x2 = x * x + red2 = x2.reduce(ReductionOp.ADD, fastmath=fm_fast) + thread_sumsq = thread_sumsq + is_valid.select(red2, c_zero_f) + + if const_expr(TAIL_ELEMS > 0): + if lane < TAIL_ELEMS: + tail_idx = lane + VEC_ELEMS + x_tail_e = _load_scalar(copy_atom_s, elem_dtype, row_div_s, tail_idx) + x_tail = x_tail_e.to(fx.Float32) + thread_sumsq = thread_sumsq + (x_tail * x_tail) + + sum_sq = group_reduce_add(thread_sumsq) + mean_sq = sum_sq / n_float + 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 tile_i in range_constexpr(NUM_VEC_ROW_ITERS): + idx = lane + tile_i * THREADS_PER_ROW + if idx < VEC_TILES: + g = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, idx).to(fx.Float32) + x = in_local[tile_i].to(fx.Float32) + y = (x * rrms) * g + y_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, y) + _store_vec(copy_atom, VEC_WIDTH, elem_dtype, y_e, out_div, idx) + + if const_expr(TAIL_ELEMS > 0): + if lane < TAIL_ELEMS: + tail_idx = lane + VEC_ELEMS + x_tail_e = _load_scalar(copy_atom_s, elem_dtype, row_div_s, tail_idx) + g_tail_e = _load_scalar(copy_atom_s, elem_dtype, gamma_div_s, tail_idx) + x_tail = x_tail_e.to(fx.Float32) + g_tail = g_tail_e.to(fx.Float32) + y_tail = (x_tail * rrms) * g_tail + y_tail_e = _to_elem_scalar(dtype_str, elem_dtype, y_tail) + _store_scalar(copy_atom_s, elem_dtype, elem_dtype, out_div_s, tail_idx, y_tail_e) + else: + copy_atom_s = fx.make_copy_atom( + fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), + elem_bits, + ) - if const_expr(store_rstd): - if lane == 0: - _store_scalar(rstd_copy_atom, fx.Float32, fx.Float32, rstd_div, row, rrms) + row_div = fx.logical_divide(row_in, fx.make_layout(1, 1)) + gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(1, 1)) + out_div = fx.logical_divide(row_out, fx.make_layout(1, 1)) - for base_idx_int in range_constexpr(0, BLOCK_N, THREADS_PER_ROW): - idx = lane + base_idx_int - if idx < N: - x_e = _load_scalar(copy_atom_s, elem_dtype, row_div, idx) - g_e = _load_scalar(copy_atom_s, elem_dtype, gamma_div, idx) + for base_idx_int in range_constexpr(0, BLOCK_N, THREADS_PER_ROW): + idx = lane + base_idx_int + is_valid = idx < N + idx_safe = is_valid.select(idx, 0) + x_e = _load_scalar(copy_atom_s, elem_dtype, row_div, idx_safe) x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) - g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) - y = (x * rrms) * g - y_e = _to_elem_scalar(dtype_str, elem_dtype, y) - _store_scalar(copy_atom_s, elem_dtype, elem_dtype, out_div, idx, y_e) + x2 = x * x + thread_sumsq = thread_sumsq + is_valid.select(x2, c_zero_f) + + sum_sq = group_reduce_add(thread_sumsq) + mean_sq = sum_sq / n_float + 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: + x_e = _load_scalar(copy_atom_s, elem_dtype, row_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) + g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) + y = (x * rrms) * g + 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: @@ -710,6 +785,7 @@ def _build_rmsnorm_quant_module( *, is_smooth: bool, quant_dtype_str: str = "i8", + eps: float = EPS, ): tile_cols = BLOCK_THREADS * VEC_WIDTH RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) @@ -733,7 +809,7 @@ def rmsnorm_quant_kernel( quant_dtype = _quant_dtype_to_elem_type(quant_dtype_str) fm_fast = arith.FastMathFlags.fast - eps_c = EPS + eps_c = eps n_float = float(N) c_zero_f = fx.Float32(0.0) c_one_f = fx.Float32(1.0) @@ -1051,12 +1127,14 @@ def build_rmsnorm_dynamicquant_module( N: int, dtype_str: str, quant_dtype_str: str = "i8", + eps: float = EPS, ): return _build_rmsnorm_quant_module( N, dtype_str, is_smooth=False, quant_dtype_str=quant_dtype_str, + eps=eps, ) @@ -1064,12 +1142,14 @@ def build_rmsnorm_smoothquant_module( N: int, dtype_str: str, quant_dtype_str: str = "i8", + eps: float = EPS, ): return _build_rmsnorm_quant_module( N, dtype_str, is_smooth=True, quant_dtype_str=quant_dtype_str, + eps=eps, ) @@ -1079,6 +1159,7 @@ def _build_fused_add_rmsnorm_quant_module( *, is_smooth: bool, quant_dtype_str: str = "i8", + eps: float = EPS, ): arch = get_rocm_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") @@ -1107,7 +1188,7 @@ def fused_add_rmsnorm_quant_kernel( quant_dtype = _quant_dtype_to_elem_type(quant_dtype_str) fm_fast = arith.FastMathFlags.fast - eps_c = EPS + eps_c = eps n_float = float(N) c_zero_f = fx.Float32(0.0) c_one_f = fx.Float32(1.0) @@ -1450,12 +1531,14 @@ def build_fused_add_rmsnorm_dynamicquant_module( N: int, dtype_str: str, quant_dtype_str: str = "i8", + eps: float = EPS, ): return _build_fused_add_rmsnorm_quant_module( N, dtype_str, is_smooth=False, quant_dtype_str=quant_dtype_str, + eps=eps, ) @@ -1463,12 +1546,14 @@ def build_fused_add_rmsnorm_smoothquant_module( N: int, dtype_str: str, quant_dtype_str: str = "i8", + eps: float = EPS, ): return _build_fused_add_rmsnorm_quant_module( N, dtype_str, is_smooth=True, quant_dtype_str=quant_dtype_str, + eps=eps, ) diff --git a/tests/kernels/test_layernorm.py b/tests/kernels/test_layernorm.py index fe7022cd9..bd8f4dd06 100644 --- a/tests/kernels/test_layernorm.py +++ b/tests/kernels/test_layernorm.py @@ -97,6 +97,8 @@ def _get_layernorm_configs(): (32, 128, "f16"), # f16 aligned (64, 2000, "f32"), # unaligned tail handling (16, 512, "bf16"), # bf16 small shape + (64, 4096, "f16"), # f16 vectorized fast path below 8192 + (64, 8192, "f16"), # f16 fast-path N (64, 8192, "bf16"), # bf16 fast-path N with small M ] return configs @@ -578,13 +580,13 @@ def kernel_launch(): return ok, flydsl_gpu_us -def _reference_layernorm(input_dev, gamma_dev, beta_dev): +def _reference_layernorm(input_dev, gamma_dev, beta_dev, eps: float = EPS): x = input_dev.to(DTYPE_FP32) gamma = gamma_dev.to(DTYPE_FP32) beta = beta_dev.to(DTYPE_FP32) mean = x.mean(dim=1, keepdim=True) var = x.var(dim=1, keepdim=True, unbiased=False) - return ((x - mean) / torch.sqrt(var + EPS) * gamma + beta).to(DTYPE_FP32) + return ((x - mean) / torch.sqrt(var + eps) * gamma + beta).to(DTYPE_FP32) def _reference_layernorm_quant(input_dev, gamma_dev, beta_dev, *, xscale_dev=None): diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index 5fbb1e8e9..ba9d2a05b 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -136,6 +136,8 @@ def _get_rmsnorm_configs(): (32, 128, "f16"), # f16 aligned (64, 2000, "f32"), # unaligned tail handling (16, 512, "bf16"), # bf16 small shape + (64, 4096, "f16"), # f16 vectorized fast path below 8192 + (64, 8192, "f16"), # f16 fast-path N (64, 8192, "bf16"), # bf16 fast-path N with small M ]