From 36bb704a0649ab672c3406a44bba5efa517a1cde Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Tue, 28 Jul 2026 03:48:44 +0800 Subject: [PATCH 01/25] fix(rmsnorm): use custom eps in quantized forward kernels Thread caller-provided epsilon through dynamic, smooth, and fused-add quant builders, with coverage that detects silently hard-coded defaults. --- kernels/norm/rmsnorm_kernel.py | 14 ++++++-- tests/kernels/test_rmsnorm.py | 64 ++++++++++++++++++++++++++-------- 2 files changed, 62 insertions(+), 16 deletions(-) diff --git a/kernels/norm/rmsnorm_kernel.py b/kernels/norm/rmsnorm_kernel.py index 3b54bd430e..6e8edd9b12 100644 --- a/kernels/norm/rmsnorm_kernel.py +++ b/kernels/norm/rmsnorm_kernel.py @@ -743,6 +743,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) @@ -766,7 +767,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) @@ -1079,12 +1080,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, ) @@ -1092,12 +1095,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, ) @@ -1107,6 +1112,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") @@ -1135,7 +1141,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) @@ -1473,12 +1479,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, ) @@ -1486,12 +1494,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_rmsnorm.py b/tests/kernels/test_rmsnorm.py index c7175c2ca2..c27dafd11f 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -161,6 +161,19 @@ def _get_rmsnorm_large_configs(): ] +def _quant_test_eps(N: int) -> float: + """Use one non-default value to verify every quant builder forwards eps.""" + return 1e-2 if N == 256 else EPS + + +def _assert_rmsnorm_forward_copy_width(compiled_fn, N: int, dtype: str): + if dtype not in ("f16", "bf16"): + return + expects_vec8 = N >= 8 and (N <= rmsnorm_kernel_impl.SMALL_N_THRESHOLD or N % 8 == 0) + has_128b_io = "buffer_copy<128>, 16>" in compiled_fn._keepalive.source_ir + assert has_128b_io == expects_vec8, f"unexpected RMSNorm copy width for N={N}, dtype={dtype}" + + def run_test(M: int, N: int, dtype: str = "f32", weight_dtype: str | None = None): weight_dtype = dtype if weight_dtype is None else weight_dtype print(f"\nTesting RMSNorm (M={M}, N={N}, dtype={dtype}, weight_dtype={weight_dtype})") @@ -243,21 +256,24 @@ def kernel_launch(): return ok, flydsl_gpu_us -def run_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): +def run_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool, eps: float = EPS): mode = "smoothquant" if is_smooth else "dynamicquant" print(f"\nTesting RMSNorm {mode} (M={M}, N={N}, dtype={dtype})") try: if is_smooth: - launch_fn = build_rmsnorm_smoothquant_module(N, dtype) + launch_fn = build_rmsnorm_smoothquant_module(N, dtype, eps=eps) else: - launch_fn = build_rmsnorm_dynamicquant_module(N, dtype) + launch_fn = build_rmsnorm_dynamicquant_module(N, dtype, eps=eps) except Exception as e: print(f"[FAIL] Compile failed for {mode} (M={M}, N={N}, dtype={dtype}): " f"{type(e).__name__}: {e}") return False, None torch.manual_seed(42) input_t = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) + if eps != EPS: + # Keep mean(x^2) below eps so a hardcoded default is visible in YScale. + input_t = input_t * 1e-3 gamma_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) torch_dtype = _torch_dtype(dtype) @@ -317,6 +333,7 @@ def kernel_launch(): input_dev, gamma_dev, xscale_dev=xscale_dev, + eps=eps, ) q_out = output_dev.to(torch.int16) q_expected = q_ref.to(torch.int16) @@ -442,15 +459,15 @@ def kernel_launch(): return ok, flydsl_gpu_us -def run_fused_add_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): +def run_fused_add_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool, eps: float = EPS): mode = "smoothquant" if is_smooth else "dynamicquant" print(f"\nTesting FusedAdd RMSNorm {mode} (M={M}, N={N}, dtype={dtype})") try: if is_smooth: - launch_fn = build_fused_add_rmsnorm_smoothquant_module(N, dtype) + launch_fn = build_fused_add_rmsnorm_smoothquant_module(N, dtype, eps=eps) else: - launch_fn = build_fused_add_rmsnorm_dynamicquant_module(N, dtype) + launch_fn = build_fused_add_rmsnorm_dynamicquant_module(N, dtype, eps=eps) except Exception as e: print( f"[FAIL] Compile failed for fused_add rmsnorm {mode} " @@ -461,6 +478,10 @@ def run_fused_add_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): torch.manual_seed(42) input_t = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) residual_t = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) + if eps != EPS: + # Keep mean((x + residual)^2) below eps so YScale distinguishes eps. + input_t = input_t * 1e-3 + residual_t = residual_t * 1e-3 gamma_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) torch_dtype = _torch_dtype(dtype) @@ -568,6 +589,7 @@ def kernel_launch(): residual_in_dev, gamma_dev, xscale_dev=xscale_dev, + eps=eps, ) residual_out_ref = residual_out_dev.to(DTYPE_FP32) q_out = output_dev.to(torch.int16) @@ -603,14 +625,14 @@ def kernel_launch(): return ok, flydsl_gpu_us -def _reference_rmsnorm(input_dev, gamma_dev): +def _reference_rmsnorm(input_dev, gamma_dev, *, eps: float = EPS): x = input_dev.to(DTYPE_FP32) gamma = gamma_dev.to(DTYPE_FP32) - return ((x / torch.sqrt((x * x).mean(dim=1, keepdim=True) + EPS)) * gamma).to(DTYPE_FP32) + return ((x / torch.sqrt((x * x).mean(dim=1, keepdim=True) + eps)) * gamma).to(DTYPE_FP32) -def _reference_rmsnorm_quant(input_dev, gamma_dev, *, xscale_dev=None): - normalized = _reference_rmsnorm(input_dev, gamma_dev) +def _reference_rmsnorm_quant(input_dev, gamma_dev, *, xscale_dev=None, eps: float = EPS): + normalized = _reference_rmsnorm(input_dev, gamma_dev, eps=eps) if xscale_dev is not None: normalized = normalized * xscale_dev.to(DTYPE_FP32) @@ -634,6 +656,7 @@ def _reference_fused_add_rmsnorm_quant( gamma_dev, *, xscale_dev=None, + eps: float = EPS, ): added = input_dev + residual_in_dev residual_expected = added.to(DTYPE_FP32) @@ -641,6 +664,7 @@ def _reference_fused_add_rmsnorm_quant( added, gamma_dev, xscale_dev=xscale_dev, + eps=eps, ) return residual_expected, q, yscale @@ -1865,7 +1889,7 @@ def test_rmsnorm_dynamicquant(): failures = 0 for M, N, dtype in configs: - ok, flydsl_gpu_us = run_quant_test(M, N, dtype, is_smooth=False) + ok, flydsl_gpu_us = run_quant_test(M, N, dtype, is_smooth=False, eps=_quant_test_eps(N)) if not ok: failures += 1 @@ -1913,7 +1937,7 @@ def test_rmsnorm_smoothquant(): failures = 0 for M, N, dtype in configs: - ok, flydsl_gpu_us = run_quant_test(M, N, dtype, is_smooth=True) + ok, flydsl_gpu_us = run_quant_test(M, N, dtype, is_smooth=True, eps=_quant_test_eps(N)) if not ok: failures += 1 @@ -2000,7 +2024,13 @@ def test_fused_add_rmsnorm_dynamicquant(): failures = 0 for M, N, dtype in configs: - ok, flydsl_gpu_us = run_fused_add_quant_test(M, N, dtype, is_smooth=False) + ok, flydsl_gpu_us = run_fused_add_quant_test( + M, + N, + dtype, + is_smooth=False, + eps=_quant_test_eps(N), + ) if not ok: failures += 1 @@ -2046,7 +2076,13 @@ def test_fused_add_rmsnorm_smoothquant(): failures = 0 for M, N, dtype in configs: - ok, flydsl_gpu_us = run_fused_add_quant_test(M, N, dtype, is_smooth=True) + ok, flydsl_gpu_us = run_fused_add_quant_test( + M, + N, + dtype, + is_smooth=True, + eps=_quant_test_eps(N), + ) if not ok: failures += 1 From bc5d37440d45af3b1c42d63c236e7dab105826fb Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Tue, 28 Jul 2026 03:48:44 +0800 Subject: [PATCH 02/25] perf(rmsnorm): extend vec8 coverage in forward kernels Vectorize partial large-N tiles and small-N prefixes with scalar tails while preserving mixed weights and rstd behavior. Add boundary and source-IR coverage for the new routing. Co-authored-by: Cursor --- kernels/norm/rmsnorm_kernel.py | 222 +++++++++++++++++++++------------ tests/kernels/test_rmsnorm.py | 56 ++++++++- 2 files changed, 193 insertions(+), 85 deletions(-) diff --git a/kernels/norm/rmsnorm_kernel.py b/kernels/norm/rmsnorm_kernel.py index 6e8edd9b12..344eaad3ff 100644 --- a/kernels/norm/rmsnorm_kernel.py +++ b/kernels/norm/rmsnorm_kernel.py @@ -39,6 +39,7 @@ from kernels.norm.rmsnorm_common import load_vec as _load_vec from kernels.norm.rmsnorm_common import load_weight_vec as _load_weight_vec from kernels.norm.rmsnorm_common import make_reduction_storage as _make_reduction_storage +from kernels.norm.rmsnorm_common import make_single_reduction_storage as _make_single_reduction_storage from kernels.norm.rmsnorm_common import resolve_rmsnorm_weight_dtype as _resolve_rmsnorm_weight_dtype from kernels.norm.rmsnorm_common import store_scalar as _store_scalar from kernels.norm.rmsnorm_common import store_vec as _store_vec @@ -85,13 +86,15 @@ 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 weight_elem_bits = 32 if weight_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) + SharedStorage = _make_single_reduction_storage(RED_SLOTS) @flyc.kernel(**_kernel_kwargs) def rmsnorm_kernel( @@ -111,7 +114,6 @@ def rmsnorm_kernel( lds = fx.SharedAllocator().allocate(SharedStorage).peek() s_red = lds.s_red.view(fx.make_layout(RED_SLOTS, 1)) - s_red2 = lds.s_red2.view(fx.make_layout(RED_SLOTS, 1)) if const_expr(store_rstd): Rstd_buf = fx.rocdl.make_buffer_tensor(Rstd) @@ -127,45 +129,30 @@ def wave_reduce_add(x): return w def block_reduce_add(val): - dummy = fx.Float32(0.0) - r0, _ = block_reduce_add2(val, dummy) - return r0 - - def block_reduce_add2(val0, val1): if const_expr(RED_SLOTS == 1): - return wave_reduce_add(val0), wave_reduce_add(val1) + return wave_reduce_add(val) lane = tid % WARP_SIZE wave = tid // WARP_SIZE - - w0 = wave_reduce_add(val0) - w1 = wave_reduce_add(val1) - + w = wave_reduce_add(val) if lane == 0: - fx.memref_store(w0, s_red, wave) - fx.memref_store(w1, s_red2, wave) + fx.memref_store(w, s_red, wave) gpu.barrier() if wave == 0: in_range = lane < RED_SLOTS lane_safe = in_range.select(lane, 0) - v0 = fx.memref_load(s_red, lane_safe) - v1 = fx.memref_load(s_red2, 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) - + v = fx.memref_load(s_red, lane_safe) + ww = in_range.select(v, 0.0) + ww = wave_reduce_add(ww) if lane == 0: - fx.memref_store(ww0, s_red, 0) - fx.memref_store(ww1, s_red2, 0) + fx.memref_store(ww, s_red, 0) gpu.barrier() - return fx.memref_load(s_red, 0), fx.memref_load(s_red2, 0) + return fx.memref_load(s_red, 0) - # Fast path for complete 128-bit tiles. - if const_expr(N >= tile_cols and N % tile_cols == 0 and elem_bits <= 16): - num_tiles = N // tile_cols + # Vector path for every complete f16/bf16 vec8 row. + if const_expr(USE_VEC_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) @@ -182,21 +169,22 @@ def block_reduce_add2(val0, val1): c_zero_f = fx.Float32(0.0) thread_sumsq = c_zero_f - thread_dummy = c_zero_f 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) + sum_sq = block_reduce_add(thread_sumsq) mean_sq = sum_sq / n_float ms_eps = mean_sq + eps_c rrms = fmath.rsqrt(ms_eps, fastmath=fm_fast) @@ -206,17 +194,15 @@ def block_reduce_add2(val0, val1): _store_scalar(rstd_copy_atom, 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 + if idx < VEC_TILES: + g = _load_weight_vec(gamma_copy_atom, weight_dtype_str, weight_elem_dtype, gamma_div, idx) + x = in_local[tile_i].to(fx.Float32) - g = _load_weight_vec(gamma_copy_atom, weight_dtype_str, weight_elem_dtype, gamma_div, idx) - 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) - - out_idx = tid + tile_i * BLOCK_THREADS - _store_vec(copy_atom, VEC_WIDTH, elem_dtype, out_e, out_div, out_idx) + y = (x * rrms) * g + out_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, y) + _store_vec(copy_atom, VEC_WIDTH, elem_dtype, out_e, out_div, idx) else: # Scalar fallback for arbitrary N. @@ -351,6 +337,13 @@ def _build_rmsnorm_large_m_small_n_module( BLOCK_THREADS_SPECIAL = BLOCK_M * THREADS_PER_ROW elem_bits = 32 if dtype_str == "f32" else 16 weight_elem_bits = 32 if weight_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( @@ -385,19 +378,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, - ) - gamma_copy_atom_s = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if weight_elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - weight_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))): @@ -409,34 +389,118 @@ 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) + 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(_weight_vec_width(weight_dtype_str), 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) + gamma_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), weight_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) + gamma_copy_atom_s = fx.make_copy_atom( + fx.rocdl.BufferCopy16b() if weight_elem_bits <= 16 else fx.rocdl.BufferCopy32b(), + weight_elem_bits, + ) - 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) + 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, 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_weight_vec( + gamma_copy_atom, + weight_dtype_str, + weight_elem_dtype, + gamma_div, + idx, + ) + 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(gamma_copy_atom_s, weight_elem_dtype, gamma_div_s, tail_idx) + x_tail = x_tail_e.to(fx.Float32) + g_tail = g_tail_e if weight_dtype_str == "f32" else 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, 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, + ) + gamma_copy_atom_s = fx.make_copy_atom( + fx.rocdl.BufferCopy16b() if weight_elem_bits <= 16 else fx.rocdl.BufferCopy32b(), + weight_elem_bits, + ) - if const_expr(store_rstd): - if lane == 0: - _store_scalar(rstd_copy_atom, 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(gamma_copy_atom_s, weight_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 weight_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, 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, 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(gamma_copy_atom_s, weight_elem_dtype, gamma_div, idx) + x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) + g = g_e if weight_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, out_div, idx, y_e) if store_rstd: diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index c27dafd11f..d6450336e1 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -135,11 +135,23 @@ def _get_rmsnorm_configs(): if override is not None: return override return [ - (64, 256, "f32"), # f32 aligned - (32, 128, "f16"), # f16 aligned - (64, 2000, "f32"), # unaligned tail handling - (16, 512, "bf16"), # bf16 small shape - (64, 8192, "bf16"), # bf16 fast-path N with small M + (64, 256, "f32"), # small-N scalar path + (32, 128, "f16"), # small-N aligned vec8 path + (64, 2001, "f32"), # small-N scalar path with an unaligned tail + (16, 512, "bf16"), # small-N multi-row vec8 path + (64, 8192, "bf16"), # large-N aligned vec8 path + ] + + +def _get_rmsnorm_forward_configs(): + override = _get_rmsnorm_shape_override() + if override is not None: + return override + return _get_rmsnorm_configs() + [ + (1, 511, "bf16"), # small M plus vec8 prefix/scalar tail + (9, 2048, "f16"), # small-N dispatch boundary + (33, 3072, "bf16"), # vec8 row not divisible by BLOCK_THREADS * VEC_WIDTH + (7, 3073, "f16"), # large-N scalar fallback for an unaligned row ] @@ -212,6 +224,7 @@ def run_test(M: int, N: int, dtype: str = "f32", weight_dtype: str | None = None print("Launching kernel...") stream = torch.cuda.current_stream() compiled_fn = flyc.compile(launch_fn, input_dev, gamma_dev, output_dev, M, stream) + _assert_rmsnorm_forward_copy_width(compiled_fn, N, dtype) def kernel_launch(): compiled_fn(input_dev, gamma_dev, output_dev, M, stream) @@ -1392,7 +1405,7 @@ def test_rmsnorm(): print("Running RMSNorm Tests") print("=" * 80) - configs = _get_rmsnorm_configs() + configs = _get_rmsnorm_forward_configs() do_compare = os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1" perf_rows = [] @@ -1547,6 +1560,36 @@ def test_rmsnorm_eps_honored(): print(" -> PASSED") +def test_rmsnorm_small_n_mixed_weight_tail_and_rstd(): + """The small-N vec8 tail path preserves FP32 weights and training rstd.""" + device = torch.device("cuda", torch.cuda.current_device()) + M, N = 3, 511 + torch.manual_seed(0) + x = torch.randn((M, N), device=device, dtype=DTYPE_BF16) + weight = torch.rand((N,), device=device, dtype=DTYPE_FP32) + output = torch.empty_like(x) + rstd = torch.empty((M,), device=device, dtype=DTYPE_FP32) + stream = torch.cuda.current_stream(device) + + launcher = build_rmsnorm_module( + N, + "bf16", + store_rstd=True, + weight_dtype_str="f32", + ) + compiled = flyc.compile(launcher, x, weight, output, rstd, M, stream) + _assert_rmsnorm_forward_copy_width(compiled, N, "bf16") + assert "buffer_copy<128>, 32>" in compiled._keepalive.source_ir + compiled(x, weight, output, rstd, M, stream) + torch.cuda.synchronize(device) + + x_f32 = x.to(DTYPE_FP32) + expected_rstd = torch.rsqrt(x_f32.square().mean(dim=1) + EPS) + expected = x_f32 * expected_rstd.unsqueeze(1) * weight + torch.testing.assert_close(output.to(DTYPE_FP32), expected, rtol=1e-2, atol=2e-2) + torch.testing.assert_close(rstd, expected_rstd, rtol=1e-4, atol=1e-4) + + def test_rmsnorm_bwd_dispatch_boundary(): """The hybrid selector is the single authority for atomic vs two-stage.""" device = torch.device("cuda", torch.cuda.current_device()) @@ -2117,6 +2160,7 @@ def test_fused_add_rmsnorm_smoothquant(): test_rmsnorm_backward() test_rmsnorm_autograd() test_rmsnorm_eps_honored() + test_rmsnorm_small_n_mixed_weight_tail_and_rstd() if torch.cuda.device_count() >= 2: test_rmsnorm_multi_gpu() test_rmsnorm_dynamicquant() From d6bb83f2f928c67522410ad4d17fda27b9fd9a4a Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Tue, 28 Jul 2026 05:07:47 +0800 Subject: [PATCH 03/25] fix(layernorm): use custom eps in forward kernels Thread caller-provided epsilon through fused and quantized forward builders, with tests that detect silently hard-coded defaults. Co-authored-by: Cursor --- kernels/norm/layernorm_kernel.py | 18 ++++++-- tests/kernels/test_layernorm.py | 79 ++++++++++++++++++++++++-------- 2 files changed, 73 insertions(+), 24 deletions(-) diff --git a/kernels/norm/layernorm_kernel.py b/kernels/norm/layernorm_kernel.py index f2f3f517ad..4957b2083b 100644 --- a/kernels/norm/layernorm_kernel.py +++ b/kernels/norm/layernorm_kernel.py @@ -560,7 +560,7 @@ def launch_layernorm_bwd( return launch_layernorm_bwd -def build_fused_add_layernorm_module(N: int, dtype_str: str): +def build_fused_add_layernorm_module(N: int, dtype_str: str, eps: float = EPS): arch = get_rocm_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") @@ -583,7 +583,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)) @@ -786,6 +786,7 @@ 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 @@ -809,7 +810,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) @@ -1130,6 +1131,7 @@ 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") @@ -1158,7 +1160,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) @@ -1506,12 +1508,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 +1523,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 +1538,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 +1553,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/tests/kernels/test_layernorm.py b/tests/kernels/test_layernorm.py index fe7022cd9b..d17b33374b 100644 --- a/tests/kernels/test_layernorm.py +++ b/tests/kernels/test_layernorm.py @@ -108,6 +108,11 @@ def _get_layernorm_large_configs(): ] +def _quant_test_eps(N: int) -> float: + """Use a non-default eps to verify every quant builder forwards it.""" + return 1e-2 if N == 256 else EPS + + def run_test(M: int, N: int, dtype: str = "f32"): print(f"\nTesting LayerNorm (M={M}, N={N}, dtype={dtype})") @@ -206,21 +211,23 @@ def kernel_launch(): return ok, flydsl_gpu_us -def run_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): +def run_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool, eps: float = EPS): mode = "smoothquant" if is_smooth else "dynamicquant" print(f"\nTesting LayerNorm {mode} (M={M}, N={N}, dtype={dtype})") try: if is_smooth: - launch_fn = build_layernorm_smoothquant_module(N, dtype) + launch_fn = build_layernorm_smoothquant_module(N, dtype, eps=eps) else: - launch_fn = build_layernorm_dynamicquant_module(N, dtype) + launch_fn = build_layernorm_dynamicquant_module(N, dtype, eps=eps) except Exception as e: print(f"[FAIL] Compile failed for {mode} layernorm (M={M}, N={N}, dtype={dtype}): {type(e).__name__}: {e}") return False, None torch.manual_seed(42) input_t = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) + if eps != EPS: + input_t = input_t * 1e-3 gamma_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) beta_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) xscale_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + 0.5 if is_smooth else None @@ -244,6 +251,7 @@ def run_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): gamma_ref, beta_ref, xscale_dev=xscale_dev if is_smooth else None, + eps=eps, ) print("Launching kernel...") @@ -313,11 +321,11 @@ def kernel_launch(): return ok, flydsl_gpu_us -def run_fused_add_test(M: int, N: int, dtype: str = "f32"): +def run_fused_add_test(M: int, N: int, dtype: str = "f32", eps: float = EPS): print(f"\nTesting FusedAdd LayerNorm (M={M}, N={N}, dtype={dtype})") try: - launch_fn = build_fused_add_layernorm_module(N, dtype) + launch_fn = build_fused_add_layernorm_module(N, dtype, eps=eps) except Exception as e: print(f"[FAIL] Compile failed for fused_add layernorm (M={M}, N={N}, dtype={dtype}): {type(e).__name__}: {e}") return False, None @@ -344,7 +352,13 @@ def run_fused_add_test(M: int, N: int, dtype: str = "f32"): else: raise ValueError(f"unsupported dtype: {dtype}") - residual_expected, expected = _reference_fused_add_layernorm(input_dev, residual_dev, gamma_dev, beta_dev) + residual_expected, expected = _reference_fused_add_layernorm( + input_dev, + residual_dev, + gamma_dev, + beta_dev, + eps=eps, + ) print("Launching kernel...") stream = torch.cuda.current_stream() @@ -410,15 +424,15 @@ def kernel_launch(): return ok, flydsl_gpu_us -def run_fused_add_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): +def run_fused_add_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool, eps: float = EPS): mode = "smoothquant" if is_smooth else "dynamicquant" print(f"\nTesting FusedAdd LayerNorm {mode} (M={M}, N={N}, dtype={dtype})") try: if is_smooth: - launch_fn = build_fused_add_layernorm_smoothquant_module(N, dtype) + launch_fn = build_fused_add_layernorm_smoothquant_module(N, dtype, eps=eps) else: - launch_fn = build_fused_add_layernorm_dynamicquant_module(N, dtype) + launch_fn = build_fused_add_layernorm_dynamicquant_module(N, dtype, eps=eps) except Exception as e: print( f"[FAIL] Compile failed for fused_add {mode} layernorm " @@ -429,6 +443,9 @@ def run_fused_add_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): torch.manual_seed(42) input_t = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) residual_t = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) + if eps != EPS: + input_t = input_t * 1e-3 + residual_t = residual_t * 1e-3 gamma_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) beta_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) xscale_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + 0.5 if is_smooth else None @@ -460,6 +477,7 @@ def run_fused_add_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): gamma_dev, beta_dev, xscale_dev=xscale_dev if is_smooth else None, + eps=eps, ) print("Launching kernel...") @@ -578,17 +596,17 @@ 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): - normalized = _reference_layernorm(input_dev, gamma_dev, beta_dev) +def _reference_layernorm_quant(input_dev, gamma_dev, beta_dev, *, xscale_dev=None, eps: float = EPS): + normalized = _reference_layernorm(input_dev, gamma_dev, beta_dev, eps=eps) if xscale_dev is not None: normalized = normalized * xscale_dev.to(DTYPE_FP32) @@ -598,14 +616,22 @@ def _reference_layernorm_quant(input_dev, gamma_dev, beta_dev, *, xscale_dev=Non return q, yscale -def _reference_fused_add_layernorm(input_dev, residual_dev, gamma_dev, beta_dev): +def _reference_fused_add_layernorm(input_dev, residual_dev, gamma_dev, beta_dev, *, eps: float = EPS): added = input_dev + residual_dev residual_expected = added.to(DTYPE_FP32) - expected = _reference_layernorm(added, gamma_dev, beta_dev) + expected = _reference_layernorm(added, gamma_dev, beta_dev, eps=eps) return residual_expected, expected -def _reference_fused_add_layernorm_quant(input_dev, residual_dev, gamma_dev, beta_dev, *, xscale_dev=None): +def _reference_fused_add_layernorm_quant( + input_dev, + residual_dev, + gamma_dev, + beta_dev, + *, + xscale_dev=None, + eps: float = EPS, +): added = input_dev + residual_dev residual_expected = added.to(DTYPE_FP32) q, yscale = _reference_layernorm_quant( @@ -613,6 +639,7 @@ def _reference_fused_add_layernorm_quant(input_dev, residual_dev, gamma_dev, bet gamma_dev, beta_dev, xscale_dev=xscale_dev, + eps=eps, ) return residual_expected, q, yscale @@ -844,7 +871,7 @@ def test_layernorm_dynamicquant(): failures = 0 for M, N, dtype in configs: - ok, flydsl_gpu_us = run_quant_test(M, N, dtype, is_smooth=False) + ok, flydsl_gpu_us = run_quant_test(M, N, dtype, is_smooth=False, eps=_quant_test_eps(N)) if not ok: failures += 1 @@ -887,7 +914,7 @@ def test_layernorm_smoothquant(): failures = 0 for M, N, dtype in configs: - ok, flydsl_gpu_us = run_quant_test(M, N, dtype, is_smooth=True) + ok, flydsl_gpu_us = run_quant_test(M, N, dtype, is_smooth=True, eps=_quant_test_eps(N)) if not ok: failures += 1 @@ -930,7 +957,13 @@ def test_fused_add_layernorm_dynamicquant(): failures = 0 for M, N, dtype in configs: - ok, flydsl_gpu_us = run_fused_add_quant_test(M, N, dtype, is_smooth=False) + ok, flydsl_gpu_us = run_fused_add_quant_test( + M, + N, + dtype, + is_smooth=False, + eps=_quant_test_eps(N), + ) if not ok: failures += 1 @@ -973,7 +1006,13 @@ def test_fused_add_layernorm_smoothquant(): failures = 0 for M, N, dtype in configs: - ok, flydsl_gpu_us = run_fused_add_quant_test(M, N, dtype, is_smooth=True) + ok, flydsl_gpu_us = run_fused_add_quant_test( + M, + N, + dtype, + is_smooth=True, + eps=_quant_test_eps(N), + ) if not ok: failures += 1 From 0f2d86c8e5a38a95f86891578adc03c418b217d8 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Tue, 28 Jul 2026 05:10:40 +0800 Subject: [PATCH 04/25] perf(layernorm): extend vec8 coverage in forward kernels Use masked vector iterations for aligned f16 and bf16 row widths, including a partially active final iteration in the test matrix. Co-authored-by: Cursor --- kernels/norm/layernorm_kernel.py | 201 +++++++++++++++++-------------- tests/kernels/test_layernorm.py | 1 + 2 files changed, 110 insertions(+), 92 deletions(-) diff --git a/kernels/norm/layernorm_kernel.py b/kernels/norm/layernorm_kernel.py index 4957b2083b..3411514595 100644 --- a/kernels/norm/layernorm_kernel.py +++ b/kernels/norm/layernorm_kernel.py @@ -123,6 +123,9 @@ 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 + 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 SharedStorage = _make_reduction_storage(RED_SLOTS) @@ -221,12 +224,9 @@ def compute_mean_rstd(sum_val, sumsq_val): return mean, rstd # ================================================================== - # Fast path: N == BLOCK_THREADS * VEC_WIDTH * 4 - # Uses buffer_load / buffer_store for high-bandwidth vectorised - # memory access (same approach as preshuffle_gemm). + # Vector path for every complete f16/bf16 vec8 row. # ================================================================== - if const_expr(N == (BLOCK_THREADS * VEC_WIDTH * 4) and elem_bits <= 16): - num_tiles_py = 4 + if const_expr(USE_VEC_N): c_zero_f = fx.Float32(0.0) thread_sum = c_zero_f thread_sumsq = c_zero_f @@ -241,17 +241,19 @@ def compute_mean_rstd(sum_val, sumsq_val): copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) # ── Pass 1: load input, accumulate sum / sumsq ─────────────── - for tile_i in range_constexpr(num_tiles_py): + 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 red = x.reduce(ReductionOp.ADD, fastmath=fm_fast) red2 = x2.reduce(ReductionOp.ADD, fastmath=fm_fast) - thread_sum = thread_sum + red - thread_sumsq = thread_sumsq + red2 + thread_sum = thread_sum + is_valid.select(red, c_zero_f) + thread_sumsq = thread_sumsq + is_valid.select(red2, c_zero_f) sum_val, sumsq_val = block_reduce_add2(thread_sum, thread_sumsq) mean, rstd = compute_mean_rstd(sum_val, sumsq_val) @@ -261,31 +263,17 @@ def compute_mean_rstd(sum_val, sumsq_val): _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) - # ── Pass 2: normalize + affine + store ─────────────────────── - for tile_i in range_constexpr(num_tiles_py): - g_next = g_cur - b_next = b_cur - if const_expr(tile_i + 1 < num_tiles_py): - next_idx = tid + (tile_i + 1) * BLOCK_THREADS - g_next = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, next_idx).to(fx.Float32) - b_next = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, beta_div, next_idx).to(fx.Float32) - else: - g_next = g_cur - b_next = b_cur - - x = in_local[tile_i].to(fx.Float32) - y = (x - mean) * rstd - y = y * g_cur + b_cur - - 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) - - g_cur = g_next - b_cur = b_next + for tile_i in range_constexpr(NUM_VEC_ITERS): + idx = tid + tile_i * BLOCK_THREADS + if idx < VEC_TILES: + g = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, idx).to(fx.Float32) + b = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, beta_div, idx).to(fx.Float32) + x = in_local[tile_i].to(fx.Float32) + y = (x - mean) * rstd + y = y * g + b + out_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, y) + _store_vec(copy_atom, VEC_WIDTH, elem_dtype, out_e, out_div, idx) else: # ============================================================== @@ -566,6 +554,9 @@ def build_fused_add_layernorm_module(N: int, dtype_str: str, eps: float = EPS): 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 SharedStorage = _make_reduction_storage(RED_SLOTS) @@ -637,10 +628,9 @@ 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 + # Vector path for every complete f16/bf16 vec8 row. # ================================================================== - if const_expr(N == (BLOCK_THREADS * VEC_WIDTH * 4) and elem_bits <= 16): - num_tiles_py = 4 + if const_expr(USE_VEC_N): c_zero_f = fx.Float32(0.0) thread_sum = c_zero_f thread_sumsq = c_zero_f @@ -668,31 +658,37 @@ def compute_mean_rstd(sum_val, sumsq_val): copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) # Pass 1: add residual, cache/store it, and accumulate sum/sumsq. - for tile_i in range_constexpr(num_tiles_py): + for tile_i in range_constexpr(NUM_VEC_ITERS): idx = tid + tile_i * BLOCK_THREADS - x = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx).to(fx.Float32) - residual = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, residual_in_div, idx).to(fx.Float32) + is_valid = idx < VEC_TILES + idx_safe = is_valid.select(idx, 0) + x = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx_safe).to(fx.Float32) + residual = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, residual_in_div, idx_safe).to(fx.Float32) added_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, x + residual) added_local.append(added_e) added = added_e.to(fx.Float32) added2 = added * added - thread_sum = thread_sum + added.reduce(ReductionOp.ADD, fastmath=fm_fast) - thread_sumsq = thread_sumsq + added2.reduce(ReductionOp.ADD, fastmath=fm_fast) - _store_vec(copy_atom, VEC_WIDTH, elem_dtype, added_e, residual_out_div, idx) + sum_val = added.reduce(ReductionOp.ADD, fastmath=fm_fast) + sumsq_val = added2.reduce(ReductionOp.ADD, fastmath=fm_fast) + thread_sum = thread_sum + is_valid.select(sum_val, c_zero_f) + thread_sumsq = thread_sumsq + is_valid.select(sumsq_val, c_zero_f) + if idx < VEC_TILES: + _store_vec(copy_atom, VEC_WIDTH, elem_dtype, added_e, residual_out_div, idx) sum_val, sumsq_val = block_reduce_add2(thread_sum, thread_sumsq) mean, rstd = compute_mean_rstd(sum_val, sumsq_val) # Pass 2: normalize + affine + store, reusing cached added values. - for tile_i in range_constexpr(num_tiles_py): + for tile_i in range_constexpr(NUM_VEC_ITERS): idx = tid + tile_i * BLOCK_THREADS - added = added_local[tile_i].to(fx.Float32) - g = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, idx).to(fx.Float32) - b = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, beta_div, idx).to(fx.Float32) - y = (added - mean) * rstd - y = y * g + b - 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 idx < VEC_TILES: + added = added_local[tile_i].to(fx.Float32) + g = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, idx).to(fx.Float32) + b = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, beta_div, idx).to(fx.Float32) + y = (added - mean) * rstd + y = y * g + b + 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) else: # ============================================================== @@ -790,6 +786,9 @@ def _build_layernorm_quant_module( ): 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 quant_dtype_max = _quant_dtype_max(quant_dtype_str) SharedStorage = _make_reduction_storage(RED_SLOTS) @@ -895,10 +894,9 @@ def block_reduce_max(val): return fx.memref_load(s_sum, 0) # ================================================================== - # Fast path: N == BLOCK_THREADS * VEC_WIDTH * 4 + # Vector path for every complete f16/bf16 vec8 row. # ================================================================== - if const_expr(N == (BLOCK_THREADS * VEC_WIDTH * 4) and elem_bits <= 16): - num_tiles_py = 4 + if const_expr(USE_VEC_N): quant_half_width = VEC_WIDTH // 2 abs_mask = full(VEC_WIDTH, fx.Uint32(0x7FFFFFFF), fx.Uint32) @@ -929,14 +927,18 @@ def block_reduce_max(val): norm_input_local = [] # Pass 1: prepare normalization input and accumulate sum/sumsq. - for tile_i in range_constexpr(num_tiles_py): + for tile_i in range_constexpr(NUM_VEC_ITERS): idx = tid + tile_i * BLOCK_THREADS - x_e = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx) + is_valid = idx < VEC_TILES + idx_safe = is_valid.select(idx, 0) + x_e = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx_safe) norm_input_local.append(x_e) x_norm = x_e.to(fx.Float32) x2 = x_norm * x_norm - thread_sum = thread_sum + x_norm.reduce(ReductionOp.ADD, fastmath=fm_fast) - thread_sumsq = thread_sumsq + x2.reduce(ReductionOp.ADD, fastmath=fm_fast) + red = x_norm.reduce(ReductionOp.ADD, fastmath=fm_fast) + red2 = x2.reduce(ReductionOp.ADD, fastmath=fm_fast) + thread_sum = thread_sum + is_valid.select(red, c_zero_f) + thread_sumsq = thread_sumsq + is_valid.select(red2, c_zero_f) sum_val, sumsq_val = block_reduce_add2(thread_sum, thread_sumsq) mean = sum_val / n_float @@ -948,20 +950,22 @@ def block_reduce_max(val): y_local = [] # Pass 2: affine (+ optional smooth scale), cache y, accumulate row max. - for tile_i in range_constexpr(num_tiles_py): + for tile_i in range_constexpr(NUM_VEC_ITERS): idx = tid + tile_i * BLOCK_THREADS + is_valid = idx < VEC_TILES + idx_safe = is_valid.select(idx, 0) x = norm_input_local[tile_i].to(fx.Float32) - g = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, idx).to(fx.Float32) - b = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, beta_div, idx).to(fx.Float32) + g = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, idx_safe).to(fx.Float32) + b = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, beta_div, idx_safe).to(fx.Float32) y = (x - mean) * rstd y = y * g + b if const_expr(is_smooth): - s = _load_vec(copy_atom_xs, VEC_WIDTH, elem_dtype, xscale_div, idx).to(fx.Float32) + s = _load_vec(copy_atom_xs, VEC_WIDTH, elem_dtype, xscale_div, idx_safe).to(fx.Float32) y = y * s y_local.append(y) y_abs = (y.bitcast(fx.Uint32) & abs_mask).bitcast(fx.Float32) tile_max = y_abs.reduce(ReductionOp.MAX) - thread_row_max = thread_row_max.maximumf(tile_max) + thread_row_max = thread_row_max.maximumf(is_valid.select(tile_max, c_zero_f)) row_max = block_reduce_max(thread_row_max) scale = row_max / c_dtype_max @@ -973,14 +977,16 @@ def block_reduce_max(val): inv_scale = c_one_f / final_scale # Pass 3: quantize + store using per-row scale. - for tile_i in range_constexpr(num_tiles_py): - q = y_local[tile_i] * inv_scale - q_i8 = q.to(quant_dtype) - q_lo = q_i8.shuffle(q_i8, [0, 1, 2, 3]) - q_hi = q_i8.shuffle(q_i8, [4, 5, 6, 7]) - out_idx = tid * 2 + tile_i * BLOCK_THREADS * 2 - _store_vec(copy_atom_q, quant_half_width, quant_dtype, q_lo, out_div_q, out_idx) - _store_vec(copy_atom_q, quant_half_width, quant_dtype, q_hi, out_div_q, out_idx + 1) + for tile_i in range_constexpr(NUM_VEC_ITERS): + idx = tid + tile_i * BLOCK_THREADS + if idx < VEC_TILES: + q = y_local[tile_i] * inv_scale + q_i8 = q.to(quant_dtype) + q_lo = q_i8.shuffle(q_i8, [0, 1, 2, 3]) + q_hi = q_i8.shuffle(q_i8, [4, 5, 6, 7]) + out_idx = idx * 2 + _store_vec(copy_atom_q, quant_half_width, quant_dtype, q_lo, out_div_q, out_idx) + _store_vec(copy_atom_q, quant_half_width, quant_dtype, q_hi, out_div_q, out_idx + 1) else: # ============================================================== @@ -1138,6 +1144,9 @@ def _build_fused_add_layernorm_quant_module( 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 quant_dtype_max = _quant_dtype_max(quant_dtype_str) SharedStorage = _make_reduction_storage(RED_SLOTS) @@ -1245,10 +1254,9 @@ def block_reduce_max(val): return fx.memref_load(s_sum, 0) # ================================================================== - # Fast path: N == BLOCK_THREADS * VEC_WIDTH * 4 + # Vector path for every complete f16/bf16 vec8 row. # ================================================================== - if const_expr(N == (BLOCK_THREADS * VEC_WIDTH * 4) and elem_bits <= 16): - num_tiles_py = 4 + if const_expr(USE_VEC_N): quant_half_width = VEC_WIDTH // 2 abs_mask = full(VEC_WIDTH, fx.Uint32(0x7FFFFFFF), fx.Uint32) @@ -1285,17 +1293,22 @@ def block_reduce_max(val): norm_input_local = [] # Pass 1: add residual, store residual_out, and accumulate sum/sumsq. - for tile_i in range_constexpr(num_tiles_py): + for tile_i in range_constexpr(NUM_VEC_ITERS): idx = tid + tile_i * BLOCK_THREADS - x = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx).to(fx.Float32) - residual = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, residual_in_div, idx).to(fx.Float32) + is_valid = idx < VEC_TILES + idx_safe = is_valid.select(idx, 0) + x = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx_safe).to(fx.Float32) + residual = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, residual_in_div, idx_safe).to(fx.Float32) added_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, x + residual) norm_input_local.append(added_e) x_norm = added_e.to(fx.Float32) - _store_vec(copy_atom, VEC_WIDTH, elem_dtype, added_e, residual_out_div, idx) + if idx < VEC_TILES: + _store_vec(copy_atom, VEC_WIDTH, elem_dtype, added_e, residual_out_div, idx) x2 = x_norm * x_norm - thread_sum = thread_sum + x_norm.reduce(ReductionOp.ADD, fastmath=fm_fast) - thread_sumsq = thread_sumsq + x2.reduce(ReductionOp.ADD, fastmath=fm_fast) + red = x_norm.reduce(ReductionOp.ADD, fastmath=fm_fast) + red2 = x2.reduce(ReductionOp.ADD, fastmath=fm_fast) + thread_sum = thread_sum + is_valid.select(red, c_zero_f) + thread_sumsq = thread_sumsq + is_valid.select(red2, c_zero_f) sum_val, sumsq_val = block_reduce_add2(thread_sum, thread_sumsq) mean = sum_val / n_float @@ -1307,20 +1320,22 @@ def block_reduce_max(val): y_local = [] # Pass 2: affine (+ optional smooth scale), cache y, accumulate row max. - for tile_i in range_constexpr(num_tiles_py): + for tile_i in range_constexpr(NUM_VEC_ITERS): idx = tid + tile_i * BLOCK_THREADS + is_valid = idx < VEC_TILES + idx_safe = is_valid.select(idx, 0) x = norm_input_local[tile_i].to(fx.Float32) - g = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, idx).to(fx.Float32) - b = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, beta_div, idx).to(fx.Float32) + g = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, idx_safe).to(fx.Float32) + b = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, beta_div, idx_safe).to(fx.Float32) y = (x - mean) * rstd y = y * g + b if const_expr(is_smooth): - s = _load_vec(copy_atom_xs, VEC_WIDTH, elem_dtype, xscale_div, idx).to(fx.Float32) + s = _load_vec(copy_atom_xs, VEC_WIDTH, elem_dtype, xscale_div, idx_safe).to(fx.Float32) y = y * s y_local.append(y) y_abs = (y.bitcast(fx.Uint32) & abs_mask).bitcast(fx.Float32) tile_max = y_abs.reduce(ReductionOp.MAX) - thread_row_max = thread_row_max.maximumf(tile_max) + thread_row_max = thread_row_max.maximumf(is_valid.select(tile_max, c_zero_f)) row_max = block_reduce_max(thread_row_max) scale = row_max / c_dtype_max @@ -1332,14 +1347,16 @@ def block_reduce_max(val): inv_scale = c_one_f / final_scale # Pass 3: quantize + store using per-row scale. - for tile_i in range_constexpr(num_tiles_py): - q = y_local[tile_i] * inv_scale - q_i8 = q.to(quant_dtype) - q_lo = q_i8.shuffle(q_i8, [0, 1, 2, 3]) - q_hi = q_i8.shuffle(q_i8, [4, 5, 6, 7]) - out_idx = tid * 2 + tile_i * BLOCK_THREADS * 2 - _store_vec(copy_atom_q, quant_half_width, quant_dtype, q_lo, out_div_q, out_idx) - _store_vec(copy_atom_q, quant_half_width, quant_dtype, q_hi, out_div_q, out_idx + 1) + for tile_i in range_constexpr(NUM_VEC_ITERS): + idx = tid + tile_i * BLOCK_THREADS + if idx < VEC_TILES: + q = y_local[tile_i] * inv_scale + q_i8 = q.to(quant_dtype) + q_lo = q_i8.shuffle(q_i8, [0, 1, 2, 3]) + q_hi = q_i8.shuffle(q_i8, [4, 5, 6, 7]) + out_idx = idx * 2 + _store_vec(copy_atom_q, quant_half_width, quant_dtype, q_lo, out_div_q, out_idx) + _store_vec(copy_atom_q, quant_half_width, quant_dtype, q_hi, out_div_q, out_idx + 1) else: # ============================================================== diff --git a/tests/kernels/test_layernorm.py b/tests/kernels/test_layernorm.py index d17b33374b..0d3b4e04d9 100644 --- a/tests/kernels/test_layernorm.py +++ b/tests/kernels/test_layernorm.py @@ -97,6 +97,7 @@ def _get_layernorm_configs(): (32, 128, "f16"), # f16 aligned (64, 2000, "f32"), # unaligned tail handling (16, 512, "bf16"), # bf16 small shape + (8, 2056, "f16"), # vec8 path with a partially active final iteration (64, 8192, "bf16"), # bf16 fast-path N with small M ] return configs From 38ba12f2ad2bfd93d545a3ffba55e19069f74320 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Wed, 29 Jul 2026 17:42:13 +0800 Subject: [PATCH 05/25] test(rmsnorm): diagnose gfx1201 SmoothQuant scale paths Isolate scalar, single-tile, and multi-tile XScale handling so Navi CI can identify the failing path without running the quarantined benchmark matrix. Co-authored-by: Cursor --- tests/kernels/test_rmsnorm.py | 92 +++++++++++++++++++++++++++++++++++ 1 file changed, 92 insertions(+) diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index d6450336e1..dcb17cdded 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -1964,6 +1964,98 @@ def test_rmsnorm_dynamicquant(): raise SystemExit(1) +@pytest.mark.skipif(GPU_ARCH != "gfx1201", reason="gfx1201-specific RMSNorm SmoothQuant diagnostics") +@pytest.mark.parametrize( + "N,dtype,route", + [ + pytest.param(256, "f32", "scalar", id="f32-scalar"), + pytest.param(128, "f16", "scalar", id="f16-scalar"), + pytest.param(512, "bf16", "scalar", id="bf16-scalar"), + pytest.param(2048, "bf16", "vec8-single-tile", id="bf16-vec1"), + pytest.param(8192, "bf16", "vec8-multi-tile", id="bf16-vec4"), + ], +) +def test_rmsnorm_smoothquant_gfx1201_scale_paths(N, dtype, route): + """Distinguish scalar, vec8, and non-uniform XScale failures on gfx1201.""" + M = 2 + scale_tol = 1e-3 + torch.manual_seed(1201) + + torch_dtype = _torch_dtype(dtype) + input_dev = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32).to(torch_dtype).contiguous() + gamma_dev = torch.rand((N,), device="cuda", dtype=DTYPE_FP32).to(torch_dtype).contiguous() + random_xscale = (torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + 0.5).to(torch_dtype).contiguous() + xscale_dev = torch.ones((N,), device="cuda", dtype=torch_dtype) + output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_INT8) + yscale_dev = torch.empty((M,), device="cuda", dtype=DTYPE_FP32) + + launch_fn = build_rmsnorm_smoothquant_module(N, dtype) + stream = torch.cuda.current_stream() + compiled_fn = flyc.compile( + launch_fn, + input_dev, + gamma_dev, + xscale_dev, + output_dev, + yscale_dev, + M, + stream, + ) + + failures = [] + for scale_mode in ("ones", "random"): + if scale_mode == "ones": + xscale_dev.fill_(1.0) + else: + xscale_dev.copy_(random_xscale) + + # flyc.compile may execute the kernel once, so poison outputs before + # every measured launch to expose missing writes. + output_dev.fill_(-128) + yscale_dev.fill_(float("nan")) + compiled_fn(input_dev, gamma_dev, xscale_dev, output_dev, yscale_dev, M, stream) + torch.cuda.synchronize() + + q_expected, yscale_expected = _reference_rmsnorm_quant( + input_dev, + gamma_dev, + xscale_dev=xscale_dev, + ) + q_out = output_dev.to(torch.int16) + q_ref = q_expected.to(torch.int16) + quant_diff = (q_out - q_ref).abs() + scale_diff = (yscale_dev - yscale_expected).abs() + quant_error = quant_diff.max().item() + scale_error = scale_diff.max().item() + finite_yscale = torch.isfinite(yscale_dev).all().item() + mismatch_count = (quant_diff > 1).sum().item() + + print( + f"gfx1201 SmoothQuant route={route}, dtype={dtype}, N={N}, " + f"xscale={scale_mode}: quant_error={quant_error}, " + f"scale_error={scale_error:.3e}, mismatches={mismatch_count}, " + f"finite_yscale={finite_yscale}" + ) + + bad_indices = (quant_diff > 1).nonzero() + if bad_indices.numel() != 0: + row, col = bad_indices[0].tolist() + print( + f"first mismatch at ({row}, {col}): " + f"expected={q_ref[row, col].item()}, actual={q_out[row, col].item()}, " + f"xscale={xscale_dev[col].item()}" + ) + + if not finite_yscale or quant_error > 1 or not scale_error < scale_tol: + failures.append( + f"xscale={scale_mode}: quant_error={quant_error}, " + f"scale_error={scale_error:.3e}, mismatches={mismatch_count}, " + f"finite_yscale={finite_yscale}" + ) + + assert not failures, f"gfx1201 SmoothQuant {route} failed: {'; '.join(failures)}" + + @pytest.mark.skipif( GPU_ARCH == "gfx1201", reason="RMSNorm SmoothQuant is temporarily quarantined on gfx1201 pending correctness investigation", From 2ab9c0016a55a37994bc2f2071b5e8731c19354e Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Thu, 30 Jul 2026 00:42:59 +0800 Subject: [PATCH 06/25] test(rmsnorm): re-enable SmoothQuant on gfx1201 Run the full plain SmoothQuant matrix on Navi now that targeted scalar and vec8 diagnostics pass. Co-authored-by: Cursor --- tests/kernels/test_rmsnorm.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index dcb17cdded..f67c1affb1 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -2056,10 +2056,6 @@ def test_rmsnorm_smoothquant_gfx1201_scale_paths(N, dtype, route): assert not failures, f"gfx1201 SmoothQuant {route} failed: {'; '.join(failures)}" -@pytest.mark.skipif( - GPU_ARCH == "gfx1201", - reason="RMSNorm SmoothQuant is temporarily quarantined on gfx1201 pending correctness investigation", -) def test_rmsnorm_smoothquant(): print("=" * 80) print("Running RMSNorm SmoothQuant Tests") From d07bf282e4041e6a91423f2ca9ee5aa7163ab916 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Fri, 31 Jul 2026 11:44:00 +0800 Subject: [PATCH 07/25] test(rmsnorm): cover exact gfx1201 SmoothQuant regression Exercise the historical f32 shape and adjacent scalar boundaries so gfx1201 failures distinguish output, scale, and fused-residual corruption. Co-authored-by: Cursor --- tests/kernels/test_rmsnorm.py | 126 +++++++++++++++++++++++++++------- 1 file changed, 103 insertions(+), 23 deletions(-) diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index f67c1affb1..4f057625bf 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -1977,8 +1977,21 @@ def test_rmsnorm_dynamicquant(): ) def test_rmsnorm_smoothquant_gfx1201_scale_paths(N, dtype, route): """Distinguish scalar, vec8, and non-uniform XScale failures on gfx1201.""" - M = 2 + _run_rmsnorm_smoothquant_gfx1201_case(2, N, dtype, route, fused_add=False) + + +@pytest.mark.skipif(GPU_ARCH != "gfx1201", reason="gfx1201-specific RMSNorm SmoothQuant regression") +@pytest.mark.parametrize("M", [2, 64]) +@pytest.mark.parametrize("N", [2000, 2001, 2048]) +@pytest.mark.parametrize("fused_add", [False, True], ids=["plain", "fused-add"]) +def test_rmsnorm_smoothquant_gfx1201_exact_regression(M, N, fused_add): + """Exercise the historical f32 failure and its adjacent scalar-loop boundaries.""" + _run_rmsnorm_smoothquant_gfx1201_case(M, N, "f32", "f32-scalar-boundary", fused_add=fused_add) + + +def _run_rmsnorm_smoothquant_gfx1201_case(M, N, dtype, route, *, fused_add): scale_tol = 1e-3 + residual_tol = {"f32": 1e-4, "f16": 1e-2, "bf16": 2e-2}[dtype] torch.manual_seed(1201) torch_dtype = _torch_dtype(dtype) @@ -1988,19 +2001,38 @@ def test_rmsnorm_smoothquant_gfx1201_scale_paths(N, dtype, route): xscale_dev = torch.ones((N,), device="cuda", dtype=torch_dtype) output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_INT8) yscale_dev = torch.empty((M,), device="cuda", dtype=DTYPE_FP32) + residual_in_dev = None + residual_out_dev = None - launch_fn = build_rmsnorm_smoothquant_module(N, dtype) stream = torch.cuda.current_stream() - compiled_fn = flyc.compile( - launch_fn, - input_dev, - gamma_dev, - xscale_dev, - output_dev, - yscale_dev, - M, - stream, - ) + if fused_add: + residual_in_dev = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32).to(torch_dtype).contiguous() + residual_out_dev = torch.empty_like(input_dev) + launch_fn = build_fused_add_rmsnorm_smoothquant_module(N, dtype) + compiled_fn = flyc.compile( + launch_fn, + input_dev, + residual_in_dev, + gamma_dev, + xscale_dev, + output_dev, + residual_out_dev, + yscale_dev, + M, + stream, + ) + else: + launch_fn = build_rmsnorm_smoothquant_module(N, dtype) + compiled_fn = flyc.compile( + launch_fn, + input_dev, + gamma_dev, + xscale_dev, + output_dev, + yscale_dev, + M, + stream, + ) failures = [] for scale_mode in ("ones", "random"): @@ -2013,14 +2045,42 @@ def test_rmsnorm_smoothquant_gfx1201_scale_paths(N, dtype, route): # every measured launch to expose missing writes. output_dev.fill_(-128) yscale_dev.fill_(float("nan")) - compiled_fn(input_dev, gamma_dev, xscale_dev, output_dev, yscale_dev, M, stream) + if fused_add: + residual_out_dev.fill_(float("nan")) + compiled_fn( + input_dev, + residual_in_dev, + gamma_dev, + xscale_dev, + output_dev, + residual_out_dev, + yscale_dev, + M, + stream, + ) + else: + compiled_fn(input_dev, gamma_dev, xscale_dev, output_dev, yscale_dev, M, stream) torch.cuda.synchronize() - q_expected, yscale_expected = _reference_rmsnorm_quant( - input_dev, - gamma_dev, - xscale_dev=xscale_dev, - ) + if fused_add: + residual_expected, q_expected, yscale_expected = _reference_fused_add_rmsnorm_quant( + input_dev, + residual_in_dev, + gamma_dev, + xscale_dev=xscale_dev, + ) + residual_diff = (residual_out_dev.to(DTYPE_FP32) - residual_expected).abs() + residual_error = residual_diff.max().item() + finite_residual = torch.isfinite(residual_out_dev).all().item() + else: + q_expected, yscale_expected = _reference_rmsnorm_quant( + input_dev, + gamma_dev, + xscale_dev=xscale_dev, + ) + residual_error = 0.0 + finite_residual = True + q_out = output_dev.to(torch.int16) q_ref = q_expected.to(torch.int16) quant_diff = (q_out - q_ref).abs() @@ -2029,12 +2089,15 @@ def test_rmsnorm_smoothquant_gfx1201_scale_paths(N, dtype, route): scale_error = scale_diff.max().item() finite_yscale = torch.isfinite(yscale_dev).all().item() mismatch_count = (quant_diff > 1).sum().item() + unwritten_count = (output_dev == -128).sum().item() + variant = "fused-add" if fused_add else "plain" print( - f"gfx1201 SmoothQuant route={route}, dtype={dtype}, N={N}, " + f"gfx1201 SmoothQuant variant={variant}, route={route}, M={M}, dtype={dtype}, N={N}, " f"xscale={scale_mode}: quant_error={quant_error}, " f"scale_error={scale_error:.3e}, mismatches={mismatch_count}, " - f"finite_yscale={finite_yscale}" + f"unwritten={unwritten_count}, finite_yscale={finite_yscale}, " + f"residual_error={residual_error:.3e}, finite_residual={finite_residual}" ) bad_indices = (quant_diff > 1).nonzero() @@ -2046,14 +2109,31 @@ def test_rmsnorm_smoothquant_gfx1201_scale_paths(N, dtype, route): f"xscale={xscale_dev[col].item()}" ) - if not finite_yscale or quant_error > 1 or not scale_error < scale_tol: + bad_scale_rows = ((scale_diff >= scale_tol) | ~torch.isfinite(yscale_dev)).nonzero() + if bad_scale_rows.numel() != 0: + row = bad_scale_rows[0].item() + print( + f"first scale mismatch at row {row}: " + f"expected={yscale_expected[row].item()}, actual={yscale_dev[row].item()}" + ) + + failed = ( + not finite_yscale + or quant_error > 1 + or not scale_error < scale_tol + or unwritten_count != 0 + or not finite_residual + or not residual_error < residual_tol + ) + if failed: failures.append( f"xscale={scale_mode}: quant_error={quant_error}, " f"scale_error={scale_error:.3e}, mismatches={mismatch_count}, " - f"finite_yscale={finite_yscale}" + f"unwritten={unwritten_count}, finite_yscale={finite_yscale}, " + f"residual_error={residual_error:.3e}, finite_residual={finite_residual}" ) - assert not failures, f"gfx1201 SmoothQuant {route} failed: {'; '.join(failures)}" + assert not failures, f"gfx1201 SmoothQuant {variant} {route}, M={M}, N={N} failed: {'; '.join(failures)}" def test_rmsnorm_smoothquant(): From 6b7a130e616cfb7c6d6de978298f2769fc6f08ea Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Fri, 31 Jul 2026 18:02:25 +0800 Subject: [PATCH 08/25] test(rmsnorm): quarantine SmoothQuant on gfx1201 Keep the exact regression matrix as expected-failure diagnostics while preventing the known incorrect path from blocking unrelated Navi coverage. --- tests/kernels/test_rmsnorm.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index 4f057625bf..fadb377240 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -1981,6 +1981,11 @@ def test_rmsnorm_smoothquant_gfx1201_scale_paths(N, dtype, route): @pytest.mark.skipif(GPU_ARCH != "gfx1201", reason="gfx1201-specific RMSNorm SmoothQuant regression") +@pytest.mark.xfail( + GPU_ARCH == "gfx1201", + reason="known gfx1201 RMSNorm SmoothQuant correctness regression; keep diagnostics while quarantined", + strict=False, +) @pytest.mark.parametrize("M", [2, 64]) @pytest.mark.parametrize("N", [2000, 2001, 2048]) @pytest.mark.parametrize("fused_add", [False, True], ids=["plain", "fused-add"]) @@ -2136,6 +2141,10 @@ def _run_rmsnorm_smoothquant_gfx1201_case(M, N, dtype, route, *, fused_add): assert not failures, f"gfx1201 SmoothQuant {variant} {route}, M={M}, N={N} failed: {'; '.join(failures)}" +@pytest.mark.skipif( + GPU_ARCH == "gfx1201", + reason="RMSNorm SmoothQuant has a known correctness regression on gfx1201", +) def test_rmsnorm_smoothquant(): print("=" * 80) print("Running RMSNorm SmoothQuant Tests") From 85471e615121474dd6269ccd220cb382becda98b Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Fri, 31 Jul 2026 18:07:12 +0800 Subject: [PATCH 09/25] docs(norm): clarify SmoothQuant scale contract Document the activation-dtype XScale ABI and gfx1201 quarantine so callers do not rely on an unsupported path. Co-authored-by: Cursor --- docs/prebuilt_kernels_guide.md | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/docs/prebuilt_kernels_guide.md b/docs/prebuilt_kernels_guide.md index d92890126f..24445c2c71 100644 --- a/docs/prebuilt_kernels_guide.md +++ b/docs/prebuilt_kernels_guide.md @@ -74,6 +74,13 @@ per-row reciprocal std (`rstd`) for use by the backward pass. `weight_dtype_str` defaults to `dtype_str`; FP16/BF16 activations additionally support FP32 weights. +**Quantized variants:** The DynamicQuant and SmoothQuant builders emit int8 +`Output` and fp32 per-row `YScale`. SmoothQuant `XScale` must use the same +element dtype as `Input`, matching the LayerNorm SmoothQuant contract; this is +an ABI change from the older RMSNorm fp32-only `XScale` convention. RMSNorm +SmoothQuant remains unsupported on gfx1201 due to a known correctness +regression. + **Backward:** `build_rmsnorm_bwd_module(N, dtype_str, weight_dtype_str=None)` builds the fused RMSNorm backward kernel (grid `(M,)`, one block per row). Kernel signature From b47a4d0648baaf721500b2b2e2aa2bca84134e49 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Tue, 4 Aug 2026 12:57:32 +0800 Subject: [PATCH 10/25] restore pre-diagnostic SmoothQuant baseline --- docs/prebuilt_kernels_guide.md | 7 -- tests/kernels/test_rmsnorm.py | 179 +-------------------------------- 2 files changed, 1 insertion(+), 185 deletions(-) diff --git a/docs/prebuilt_kernels_guide.md b/docs/prebuilt_kernels_guide.md index 24445c2c71..d92890126f 100644 --- a/docs/prebuilt_kernels_guide.md +++ b/docs/prebuilt_kernels_guide.md @@ -74,13 +74,6 @@ per-row reciprocal std (`rstd`) for use by the backward pass. `weight_dtype_str` defaults to `dtype_str`; FP16/BF16 activations additionally support FP32 weights. -**Quantized variants:** The DynamicQuant and SmoothQuant builders emit int8 -`Output` and fp32 per-row `YScale`. SmoothQuant `XScale` must use the same -element dtype as `Input`, matching the LayerNorm SmoothQuant contract; this is -an ABI change from the older RMSNorm fp32-only `XScale` convention. RMSNorm -SmoothQuant remains unsupported on gfx1201 due to a known correctness -regression. - **Backward:** `build_rmsnorm_bwd_module(N, dtype_str, weight_dtype_str=None)` builds the fused RMSNorm backward kernel (grid `(M,)`, one block per row). Kernel signature diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index fadb377240..d6450336e1 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -1964,186 +1964,9 @@ def test_rmsnorm_dynamicquant(): raise SystemExit(1) -@pytest.mark.skipif(GPU_ARCH != "gfx1201", reason="gfx1201-specific RMSNorm SmoothQuant diagnostics") -@pytest.mark.parametrize( - "N,dtype,route", - [ - pytest.param(256, "f32", "scalar", id="f32-scalar"), - pytest.param(128, "f16", "scalar", id="f16-scalar"), - pytest.param(512, "bf16", "scalar", id="bf16-scalar"), - pytest.param(2048, "bf16", "vec8-single-tile", id="bf16-vec1"), - pytest.param(8192, "bf16", "vec8-multi-tile", id="bf16-vec4"), - ], -) -def test_rmsnorm_smoothquant_gfx1201_scale_paths(N, dtype, route): - """Distinguish scalar, vec8, and non-uniform XScale failures on gfx1201.""" - _run_rmsnorm_smoothquant_gfx1201_case(2, N, dtype, route, fused_add=False) - - -@pytest.mark.skipif(GPU_ARCH != "gfx1201", reason="gfx1201-specific RMSNorm SmoothQuant regression") -@pytest.mark.xfail( - GPU_ARCH == "gfx1201", - reason="known gfx1201 RMSNorm SmoothQuant correctness regression; keep diagnostics while quarantined", - strict=False, -) -@pytest.mark.parametrize("M", [2, 64]) -@pytest.mark.parametrize("N", [2000, 2001, 2048]) -@pytest.mark.parametrize("fused_add", [False, True], ids=["plain", "fused-add"]) -def test_rmsnorm_smoothquant_gfx1201_exact_regression(M, N, fused_add): - """Exercise the historical f32 failure and its adjacent scalar-loop boundaries.""" - _run_rmsnorm_smoothquant_gfx1201_case(M, N, "f32", "f32-scalar-boundary", fused_add=fused_add) - - -def _run_rmsnorm_smoothquant_gfx1201_case(M, N, dtype, route, *, fused_add): - scale_tol = 1e-3 - residual_tol = {"f32": 1e-4, "f16": 1e-2, "bf16": 2e-2}[dtype] - torch.manual_seed(1201) - - torch_dtype = _torch_dtype(dtype) - input_dev = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32).to(torch_dtype).contiguous() - gamma_dev = torch.rand((N,), device="cuda", dtype=DTYPE_FP32).to(torch_dtype).contiguous() - random_xscale = (torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + 0.5).to(torch_dtype).contiguous() - xscale_dev = torch.ones((N,), device="cuda", dtype=torch_dtype) - output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_INT8) - yscale_dev = torch.empty((M,), device="cuda", dtype=DTYPE_FP32) - residual_in_dev = None - residual_out_dev = None - - stream = torch.cuda.current_stream() - if fused_add: - residual_in_dev = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32).to(torch_dtype).contiguous() - residual_out_dev = torch.empty_like(input_dev) - launch_fn = build_fused_add_rmsnorm_smoothquant_module(N, dtype) - compiled_fn = flyc.compile( - launch_fn, - input_dev, - residual_in_dev, - gamma_dev, - xscale_dev, - output_dev, - residual_out_dev, - yscale_dev, - M, - stream, - ) - else: - launch_fn = build_rmsnorm_smoothquant_module(N, dtype) - compiled_fn = flyc.compile( - launch_fn, - input_dev, - gamma_dev, - xscale_dev, - output_dev, - yscale_dev, - M, - stream, - ) - - failures = [] - for scale_mode in ("ones", "random"): - if scale_mode == "ones": - xscale_dev.fill_(1.0) - else: - xscale_dev.copy_(random_xscale) - - # flyc.compile may execute the kernel once, so poison outputs before - # every measured launch to expose missing writes. - output_dev.fill_(-128) - yscale_dev.fill_(float("nan")) - if fused_add: - residual_out_dev.fill_(float("nan")) - compiled_fn( - input_dev, - residual_in_dev, - gamma_dev, - xscale_dev, - output_dev, - residual_out_dev, - yscale_dev, - M, - stream, - ) - else: - compiled_fn(input_dev, gamma_dev, xscale_dev, output_dev, yscale_dev, M, stream) - torch.cuda.synchronize() - - if fused_add: - residual_expected, q_expected, yscale_expected = _reference_fused_add_rmsnorm_quant( - input_dev, - residual_in_dev, - gamma_dev, - xscale_dev=xscale_dev, - ) - residual_diff = (residual_out_dev.to(DTYPE_FP32) - residual_expected).abs() - residual_error = residual_diff.max().item() - finite_residual = torch.isfinite(residual_out_dev).all().item() - else: - q_expected, yscale_expected = _reference_rmsnorm_quant( - input_dev, - gamma_dev, - xscale_dev=xscale_dev, - ) - residual_error = 0.0 - finite_residual = True - - q_out = output_dev.to(torch.int16) - q_ref = q_expected.to(torch.int16) - quant_diff = (q_out - q_ref).abs() - scale_diff = (yscale_dev - yscale_expected).abs() - quant_error = quant_diff.max().item() - scale_error = scale_diff.max().item() - finite_yscale = torch.isfinite(yscale_dev).all().item() - mismatch_count = (quant_diff > 1).sum().item() - unwritten_count = (output_dev == -128).sum().item() - variant = "fused-add" if fused_add else "plain" - - print( - f"gfx1201 SmoothQuant variant={variant}, route={route}, M={M}, dtype={dtype}, N={N}, " - f"xscale={scale_mode}: quant_error={quant_error}, " - f"scale_error={scale_error:.3e}, mismatches={mismatch_count}, " - f"unwritten={unwritten_count}, finite_yscale={finite_yscale}, " - f"residual_error={residual_error:.3e}, finite_residual={finite_residual}" - ) - - bad_indices = (quant_diff > 1).nonzero() - if bad_indices.numel() != 0: - row, col = bad_indices[0].tolist() - print( - f"first mismatch at ({row}, {col}): " - f"expected={q_ref[row, col].item()}, actual={q_out[row, col].item()}, " - f"xscale={xscale_dev[col].item()}" - ) - - bad_scale_rows = ((scale_diff >= scale_tol) | ~torch.isfinite(yscale_dev)).nonzero() - if bad_scale_rows.numel() != 0: - row = bad_scale_rows[0].item() - print( - f"first scale mismatch at row {row}: " - f"expected={yscale_expected[row].item()}, actual={yscale_dev[row].item()}" - ) - - failed = ( - not finite_yscale - or quant_error > 1 - or not scale_error < scale_tol - or unwritten_count != 0 - or not finite_residual - or not residual_error < residual_tol - ) - if failed: - failures.append( - f"xscale={scale_mode}: quant_error={quant_error}, " - f"scale_error={scale_error:.3e}, mismatches={mismatch_count}, " - f"unwritten={unwritten_count}, finite_yscale={finite_yscale}, " - f"residual_error={residual_error:.3e}, finite_residual={finite_residual}" - ) - - assert not failures, f"gfx1201 SmoothQuant {variant} {route}, M={M}, N={N} failed: {'; '.join(failures)}" - - @pytest.mark.skipif( GPU_ARCH == "gfx1201", - reason="RMSNorm SmoothQuant has a known correctness regression on gfx1201", + reason="RMSNorm SmoothQuant is temporarily quarantined on gfx1201 pending correctness investigation", ) def test_rmsnorm_smoothquant(): print("=" * 80) From 39d389724dc6f4896cb1619ed50b5d995f49cf5e Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Tue, 18 Aug 2026 14:23:21 +0800 Subject: [PATCH 11/25] perf(rmsnorm): vectorize unaligned large-N rows Use vec8 for the aligned prefix and scalar copies only for the tail, avoiding whole-row scalar fallback. Co-authored-by: Cursor --- kernels/norm/rmsnorm_kernel.py | 37 ++++++++++++++++++++++++++++++---- tests/kernels/test_rmsnorm.py | 4 ++-- 2 files changed, 35 insertions(+), 6 deletions(-) diff --git a/kernels/norm/rmsnorm_kernel.py b/kernels/norm/rmsnorm_kernel.py index 344eaad3ff..99f7e81864 100644 --- a/kernels/norm/rmsnorm_kernel.py +++ b/kernels/norm/rmsnorm_kernel.py @@ -6,8 +6,8 @@ RMSNorm(x) = x / sqrt(mean(x^2) + eps) * gamma Two paths: - - Fast path (N % tile_cols == 0): 128-bit buffer copies. - - Generic path (arbitrary N): scalar guarded copies. + - Vector path (f16/bf16, N >= 8): 128-bit vec8 copies plus a scalar tail. + - Scalar path (f32 or N < 8): guarded elementwise copies. """ import math @@ -89,8 +89,10 @@ def build_rmsnorm_module( RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) elem_bits = 32 if dtype_str == "f32" else 16 weight_elem_bits = 32 if weight_dtype_str == "f32" else 16 - USE_VEC_N = elem_bits <= 16 and N % VEC_WIDTH == 0 + USE_VEC_N = elem_bits <= 16 and N >= VEC_WIDTH VEC_TILES = N // VEC_WIDTH if USE_VEC_N else 0 + VEC_ELEMS = VEC_TILES * VEC_WIDTH + TAIL_ELEMS = N - VEC_ELEMS 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]} @@ -151,7 +153,7 @@ def block_reduce_add(val): return fx.memref_load(s_red, 0) - # Vector path for every complete f16/bf16 vec8 row. + # Vector path for the f16/bf16 vec8 prefix plus an optional scalar tail. if const_expr(USE_VEC_N): Input_buf = fx.rocdl.make_buffer_tensor(Input) Output_buf = fx.rocdl.make_buffer_tensor(Output) @@ -166,6 +168,15 @@ def block_reduce_add(val): copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) gamma_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), weight_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) + gamma_copy_atom_s = fx.make_copy_atom( + fx.rocdl.BufferCopy16b() if weight_elem_bits <= 16 else fx.rocdl.BufferCopy32b(), + weight_elem_bits, + ) c_zero_f = fx.Float32(0.0) thread_sumsq = c_zero_f @@ -184,6 +195,13 @@ def block_reduce_add(val): 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 tid < TAIL_ELEMS: + tail_idx = tid + 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 = block_reduce_add(thread_sumsq) mean_sq = sum_sq / n_float ms_eps = mean_sq + eps_c @@ -204,6 +222,17 @@ def block_reduce_add(val): out_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, y) _store_vec(copy_atom, VEC_WIDTH, elem_dtype, out_e, out_div, idx) + if const_expr(TAIL_ELEMS > 0): + if tid < TAIL_ELEMS: + tail_idx = tid + VEC_ELEMS + x_tail_e = _load_scalar(copy_atom_s, elem_dtype, row_div_s, tail_idx) + g_tail_e = _load_scalar(gamma_copy_atom_s, weight_elem_dtype, gamma_div_s, tail_idx) + x_tail = x_tail_e.to(fx.Float32) + g_tail = g_tail_e if weight_dtype_str == "f32" else 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, out_div_s, tail_idx, y_tail_e) + else: # Scalar fallback for arbitrary N. Input_buf = fx.rocdl.make_buffer_tensor(Input) diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index d6450336e1..f8398dd383 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -151,7 +151,7 @@ def _get_rmsnorm_forward_configs(): (1, 511, "bf16"), # small M plus vec8 prefix/scalar tail (9, 2048, "f16"), # small-N dispatch boundary (33, 3072, "bf16"), # vec8 row not divisible by BLOCK_THREADS * VEC_WIDTH - (7, 3073, "f16"), # large-N scalar fallback for an unaligned row + (7, 3073, "f16"), # large-N vec8 prefix/scalar tail ] @@ -181,7 +181,7 @@ def _quant_test_eps(N: int) -> float: def _assert_rmsnorm_forward_copy_width(compiled_fn, N: int, dtype: str): if dtype not in ("f16", "bf16"): return - expects_vec8 = N >= 8 and (N <= rmsnorm_kernel_impl.SMALL_N_THRESHOLD or N % 8 == 0) + expects_vec8 = N >= 8 has_128b_io = "buffer_copy<128>, 16>" in compiled_fn._keepalive.source_ir assert has_128b_io == expects_vec8, f"unexpected RMSNorm copy width for N={N}, dtype={dtype}" From 64879943e09d4b3799f276114626d40d8b2c19e7 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Wed, 19 Aug 2026 11:43:04 -0400 Subject: [PATCH 12/25] [Fix] Barrier the LDS scratch reuse between successive block reductions block_reduce_max (rmsnorm/layernorm quant) and block_reduce (softmax) store their per-wave partials into a scratch buffer that the preceding reduction is still broadcast-reading on exit. Nothing ordered that read against the new store, so a lagging wave could observe the next reduction's value and quantize a row with the wrong statistic. That race is the source of the intermittent gfx1201 RMSNorm/LayerNorm quant mismatches in CI: the window is visible in the generated ISA on both gfx1201 and gfx942, and a host-side model of the corrupted scale matched 539/539 observed bad rows. Add the missing barrier at the head of those helpers, which closes the window at under 1% cost, and lift the now-unnecessary gfx1201 quarantine from the SmoothQuant tests. --- kernels/norm/layernorm_kernel.py | 8 ++++++++ kernels/norm/rmsnorm_kernel.py | 8 ++++++++ kernels/norm/softmax_kernel.py | 4 ++++ tests/kernels/test_rmsnorm.py | 11 ----------- 4 files changed, 20 insertions(+), 11 deletions(-) diff --git a/kernels/norm/layernorm_kernel.py b/kernels/norm/layernorm_kernel.py index 9206323458..ca4dc56ef6 100644 --- a/kernels/norm/layernorm_kernel.py +++ b/kernels/norm/layernorm_kernel.py @@ -876,6 +876,10 @@ def block_reduce_max(val): lane = tid % WARP_SIZE wave = tid // WARP_SIZE + # s_sum still holds the sum reduction's result, which every wave + # broadcast-reads on exit. Only this barrier orders that read + # against the store below; without it a late wave reads this max. + gpu.barrier() w = wave_reduce_max(val) if lane == 0: fx.memref_store(w, s_sum, wave) @@ -1236,6 +1240,10 @@ def block_reduce_max(val): lane = tid % WARP_SIZE wave = tid // WARP_SIZE + # s_sum still holds the sum reduction's result, which every wave + # broadcast-reads on exit. Only this barrier orders that read + # against the store below; without it a late wave reads this max. + gpu.barrier() w = wave_reduce_max(val) if lane == 0: fx.memref_store(w, s_sum, wave) diff --git a/kernels/norm/rmsnorm_kernel.py b/kernels/norm/rmsnorm_kernel.py index 99f7e81864..907c82487a 100644 --- a/kernels/norm/rmsnorm_kernel.py +++ b/kernels/norm/rmsnorm_kernel.py @@ -935,6 +935,10 @@ def block_reduce_max(val): lane = tid % WARP_SIZE wave = tid // WARP_SIZE + # s_red still holds the sum reduction's result, which every wave + # broadcast-reads on exit. Only this barrier orders that read + # against the store below; without it a late wave reads this max. + gpu.barrier() w = wave_reduce_max(val) if lane == 0: fx.memref_store(w, s_red, wave) @@ -1309,6 +1313,10 @@ def block_reduce_max(val): lane = tid % WARP_SIZE wave = tid // WARP_SIZE + # s_red still holds the sum reduction's result, which every wave + # broadcast-reads on exit. Only this barrier orders that read + # against the store below; without it a late wave reads this max. + gpu.barrier() w = wave_reduce_max(val) if lane == 0: fx.memref_store(w, s_red, wave) diff --git a/kernels/norm/softmax_kernel.py b/kernels/norm/softmax_kernel.py index b68a25a3c0..d339a38279 100644 --- a/kernels/norm/softmax_kernel.py +++ b/kernels/norm/softmax_kernel.py @@ -80,6 +80,10 @@ def block_reduce(val, mode, s_red_buffer): wave = tid // WARP_SIZE neutral = c_neg_inf if mode == "max" else c_zero_f + # The buffer is reused across successive reductions, and every wave + # broadcast-reads the previous result on exit. Only this barrier + # orders that read against the store below. + gpu.barrier() w = wave_reduce(val, mode) if lane == 0: diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index f8398dd383..e95a841635 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -30,7 +30,6 @@ # so importing it earlier makes a torch-less collection fail (ImportError) instead of skip. import flydsl.compiler as flyc # noqa: E402 import kernels.norm.rmsnorm_kernel as rmsnorm_kernel_impl # noqa: E402 -from flydsl.runtime.device import get_rocm_arch # noqa: E402 from kernels.common.tensor_shim import _run_compiled # noqa: E402 from kernels.norm.rmsnorm_kernel import ( # noqa: E402 build_fused_add_rmsnorm_bwd_module, @@ -59,8 +58,6 @@ DTYPE_BF16 = torch.bfloat16 DTYPE_INT8 = torch.int8 -GPU_ARCH = str(get_rocm_arch()) - EPS: float = 1e-5 WARMUP_ITERS = 10 @@ -1964,10 +1961,6 @@ def test_rmsnorm_dynamicquant(): raise SystemExit(1) -@pytest.mark.skipif( - GPU_ARCH == "gfx1201", - reason="RMSNorm SmoothQuant is temporarily quarantined on gfx1201 pending correctness investigation", -) def test_rmsnorm_smoothquant(): print("=" * 80) print("Running RMSNorm SmoothQuant Tests") @@ -2103,10 +2096,6 @@ def test_fused_add_rmsnorm_dynamicquant(): raise SystemExit(1) -@pytest.mark.skipif( - GPU_ARCH == "gfx1201", - reason="RMSNorm SmoothQuant is temporarily quarantined on gfx1201 pending correctness investigation", -) def test_fused_add_rmsnorm_smoothquant(): print("=" * 80) print("Running FusedAdd RMSNorm SmoothQuant Tests") From d42e571f45564b550b8d1e9afd0566c36b0c56d1 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Thu, 20 Aug 2026 14:18:11 +0800 Subject: [PATCH 13/25] refactor: separate base and forward shape configurations --- tests/kernels/test_rmsnorm.py | 37 +++++++++++++++++++++-------------- 1 file changed, 22 insertions(+), 15 deletions(-) diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index e95a841635..72dfae0449 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -127,29 +127,36 @@ def _get_rmsnorm_shape_override(env_name="ROCDSL_RMSNORM_SHAPES"): return configs +_RMSNORM_CONFIGS = ( + (64, 256, "f32"), # small-N scalar path + (32, 128, "f16"), # small-N aligned vec8 path + (64, 2001, "f32"), # small-N scalar path with an unaligned tail + (16, 512, "bf16"), # small-N multi-row vec8 path + (64, 8192, "bf16"), # large-N aligned vec8 path +) + +_RMSNORM_FORWARD_EXTRA_CONFIGS = ( + (1, 511, "bf16"), # small M plus vec8 prefix/scalar tail + (9, 2048, "f16"), # small-N dispatch boundary + (33, 3072, "bf16"), # vec8 row not divisible by BLOCK_THREADS * VEC_WIDTH + (7, 3073, "f16"), # large-N vec8 prefix/scalar tail +) + + +def get_default_rmsnorm_forward_configs(): + return list(_RMSNORM_CONFIGS + _RMSNORM_FORWARD_EXTRA_CONFIGS) + + def _get_rmsnorm_configs(): override = _get_rmsnorm_shape_override() - if override is not None: - return override - return [ - (64, 256, "f32"), # small-N scalar path - (32, 128, "f16"), # small-N aligned vec8 path - (64, 2001, "f32"), # small-N scalar path with an unaligned tail - (16, 512, "bf16"), # small-N multi-row vec8 path - (64, 8192, "bf16"), # large-N aligned vec8 path - ] + return override if override is not None else list(_RMSNORM_CONFIGS) def _get_rmsnorm_forward_configs(): override = _get_rmsnorm_shape_override() if override is not None: return override - return _get_rmsnorm_configs() + [ - (1, 511, "bf16"), # small M plus vec8 prefix/scalar tail - (9, 2048, "f16"), # small-N dispatch boundary - (33, 3072, "bf16"), # vec8 row not divisible by BLOCK_THREADS * VEC_WIDTH - (7, 3073, "f16"), # large-N vec8 prefix/scalar tail - ] + return get_default_rmsnorm_forward_configs() def _get_rmsnorm_backward_configs(): From 72aaa42276d37ba5870477231c14d191b0397cfb Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Thu, 20 Aug 2026 17:36:23 +0800 Subject: [PATCH 14/25] test: remove gfx1201 quant norm quarantine Co-authored-by: Cursor --- tests/kernels/test_layernorm.py | 11 ----------- tests/kernels/test_rmsnorm.py | 4 ---- 2 files changed, 15 deletions(-) diff --git a/tests/kernels/test_layernorm.py b/tests/kernels/test_layernorm.py index 4a2c938253..0d3b4e04d9 100644 --- a/tests/kernels/test_layernorm.py +++ b/tests/kernels/test_layernorm.py @@ -19,7 +19,6 @@ import flydsl.compiler as flyc import flydsl.expr as fx -from flydsl.runtime.device import get_rocm_arch from kernels.norm.layernorm_kernel import ( build_fused_add_layernorm_dynamicquant_module, build_fused_add_layernorm_module, @@ -52,8 +51,6 @@ DTYPE_BF16 = torch.bfloat16 DTYPE_INT8 = torch.int8 -GPU_ARCH = str(get_rocm_arch()) - EPS: float = 1e-5 WARMUP_ITERS = 10 @@ -906,10 +903,6 @@ def test_layernorm_dynamicquant(): raise SystemExit(1) -@pytest.mark.skipif( - GPU_ARCH == "gfx1201", - reason="LayerNorm SmoothQuant is temporarily quarantined on gfx1201 pending correctness investigation", -) def test_layernorm_smoothquant(): print("=" * 80) print("Running LayerNorm SmoothQuant Tests") @@ -1002,10 +995,6 @@ def test_fused_add_layernorm_dynamicquant(): raise SystemExit(1) -@pytest.mark.skipif( - GPU_ARCH == "gfx1201", - reason="LayerNorm SmoothQuant is temporarily quarantined on gfx1201 pending correctness investigation", -) def test_fused_add_layernorm_smoothquant(): print("=" * 80) print("Running FusedAdd LayerNorm SmoothQuant Tests") diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index e4aa4ddcbe..72dfae0449 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -1924,10 +1924,6 @@ def test_rmsnorm_large_shape(): assert ok -@pytest.mark.skipif( - GPU_ARCH == "gfx1201", - reason="RMSNorm DynamicQuant is temporarily quarantined on gfx1201 pending correctness investigation", -) def test_rmsnorm_dynamicquant(): print("=" * 80) print("Running RMSNorm DynamicQuant Tests") From c5c8afd8f47b2e6ef0180f45f70e2126a44fc0bc Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Tue, 25 Aug 2026 13:43:37 +0800 Subject: [PATCH 15/25] perf(rmsnorm): restore runtime M dispatch for small N Select one-row or multi-row forward kernels from runtime m_in, preserving the M > 8192 crossover without specializing the build cache on M. Co-authored-by: Cursor --- kernels/norm/rmsnorm_kernel.py | 82 ++++++++++++++++++++++++++-------- tests/kernels/test_rmsnorm.py | 2 +- 2 files changed, 65 insertions(+), 19 deletions(-) diff --git a/kernels/norm/rmsnorm_kernel.py b/kernels/norm/rmsnorm_kernel.py index 1456e75d89..13863f2773 100644 --- a/kernels/norm/rmsnorm_kernel.py +++ b/kernels/norm/rmsnorm_kernel.py @@ -54,8 +54,8 @@ KERNEL_NAME = "rmsnorm" -# The small-N path derives its own block geometry and is not tuned. SMALL_N_THRESHOLD = 2048 +FWD_MULTI_ROW_THRESHOLD = 8192 def _quant_dtype_to_elem_type(dtype_str: str): @@ -79,9 +79,6 @@ def build_rmsnorm_module( weight_dtype_str: str | None = None, ): weight_dtype_str = _resolve_rmsnorm_weight_dtype(dtype_str, weight_dtype_str) - if N <= SMALL_N_THRESHOLD: - return _build_rmsnorm_large_m_small_n_module(N, dtype_str, store_rstd, eps, weight_dtype_str) - arch = get_rocm_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") @@ -289,6 +286,17 @@ def block_reduce_add(val): y_e = _to_elem_scalar(dtype_str, elem_dtype, y) _store_scalar(copy_atom_s, elem_dtype, out_div, idx, y_e) + if N <= SMALL_N_THRESHOLD: + return _build_rmsnorm_large_m_small_n_module( + N, + dtype_str, + store_rstd, + eps, + weight_dtype_str, + one_row_kernel=rmsnorm_kernel, + one_row_block_threads=BLOCK_THREADS, + ) + if store_rstd: @flyc.jit @@ -358,6 +366,9 @@ def _build_rmsnorm_large_m_small_n_module( store_rstd: bool = False, eps: float = EPS, weight_dtype_str: str | None = None, + *, + one_row_kernel, + one_row_block_threads: int, ): weight_dtype_str = _resolve_rmsnorm_weight_dtype(dtype_str, weight_dtype_str) BLOCK_N = 1 << (N - 1).bit_length() @@ -542,12 +553,30 @@ def launch_rmsnorm_large_m_small_n( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): - launcher = rmsnorm_large_m_small_n_kernel(Input, Gamma, Rstd, Output, m_in) - launcher.launch( - grid=((m_in + BLOCK_M - 1) // BLOCK_M, 1, 1), - block=(BLOCK_THREADS_SPECIAL, 1, 1), - stream=stream, - ) + def launch_one_row(): + launcher = one_row_kernel(Input, Gamma, Rstd, Output) + launcher.launch( + grid=(m_in, 1, 1), + block=(one_row_block_threads, 1, 1), + stream=stream, + ) + + def launch_multi_row(): + launcher = rmsnorm_large_m_small_n_kernel(Input, Gamma, Rstd, Output, m_in) + launcher.launch( + grid=((m_in + BLOCK_M - 1) // BLOCK_M, 1, 1), + block=(BLOCK_THREADS_SPECIAL, 1, 1), + stream=stream, + ) + + @flyc.jit + def dispatch(): + if m_in > fx.Int32(FWD_MULTI_ROW_THRESHOLD): + launch_multi_row() + else: + launch_one_row() + + dispatch() return launch_rmsnorm_large_m_small_n @@ -559,14 +588,31 @@ def launch_rmsnorm_large_m_small_n( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): - # store_rstd=False path: the Rstd slot is an unused placeholder here, so - # we pass Gamma to fill the argument (it is never dereferenced in-kernel). - launcher = rmsnorm_large_m_small_n_kernel(Input, Gamma, Gamma, Output, m_in) - launcher.launch( - grid=((m_in + BLOCK_M - 1) // BLOCK_M, 1, 1), - block=(BLOCK_THREADS_SPECIAL, 1, 1), - stream=stream, - ) + def launch_one_row(): + # The Rstd slot is unused when store_rstd=False. + launcher = one_row_kernel(Input, Gamma, Gamma, Output) + launcher.launch( + grid=(m_in, 1, 1), + block=(one_row_block_threads, 1, 1), + stream=stream, + ) + + def launch_multi_row(): + launcher = rmsnorm_large_m_small_n_kernel(Input, Gamma, Gamma, Output, m_in) + launcher.launch( + grid=((m_in + BLOCK_M - 1) // BLOCK_M, 1, 1), + block=(BLOCK_THREADS_SPECIAL, 1, 1), + stream=stream, + ) + + @flyc.jit + def dispatch(): + if m_in > fx.Int32(FWD_MULTI_ROW_THRESHOLD): + launch_multi_row() + else: + launch_one_row() + + dispatch() return launch_rmsnorm_large_m_small_n diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index 72dfae0449..6f8f7a5b7e 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -131,7 +131,7 @@ def _get_rmsnorm_shape_override(env_name="ROCDSL_RMSNORM_SHAPES"): (64, 256, "f32"), # small-N scalar path (32, 128, "f16"), # small-N aligned vec8 path (64, 2001, "f32"), # small-N scalar path with an unaligned tail - (16, 512, "bf16"), # small-N multi-row vec8 path + (16, 512, "bf16"), # small-N one-row vec8 path (64, 8192, "bf16"), # large-N aligned vec8 path ) From bbdde41a41ae6945e4ebaaef0ef20b99db63ec75 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Tue, 25 Aug 2026 16:19:34 +0800 Subject: [PATCH 16/25] perf(rmsnorm): optimize wide BF16 rows on gfx942 Reduce wide-row latency with streaming accesses, gamma reuse, fast BF16 packing, and runtime geometry dispatch while keeping M dynamic. Add a three-backend benchmark to reproduce FlyDSL, HIP, and Triton comparisons. Co-authored-by: Cursor --- kernels/norm/rmsnorm_common.py | 12 + kernels/norm/rmsnorm_kernel.py | 131 ++++++++- tests/kernels/benchmark_rmsnorm_backends.py | 288 ++++++++++++++++++++ 3 files changed, 427 insertions(+), 4 deletions(-) create mode 100644 tests/kernels/benchmark_rmsnorm_backends.py diff --git a/kernels/norm/rmsnorm_common.py b/kernels/norm/rmsnorm_common.py index 72fd273439..8109dd2d9e 100644 --- a/kernels/norm/rmsnorm_common.py +++ b/kernels/norm/rmsnorm_common.py @@ -117,6 +117,18 @@ def to_elem_vec(dtype_str: str, elem_dtype, use_hw_cvt_bf16: bool, y): return y.to(elem_dtype) +def to_elem_vec_rna(dtype_str: str, elem_dtype, use_hw_cvt_bf16: bool, y): + """Pack BF16 with round-to-nearest-away, avoiding the RNE tie fixup.""" + if const_expr(dtype_str == "bf16" and not use_hw_cvt_bf16): + bf16_bits = (y.bitcast(fx.Uint32) + 0x8000) >> 16 + even = bf16_bits.shuffle(bf16_bits, [0, 2, 4, 6]) + odd = bf16_bits.shuffle(bf16_bits, [1, 3, 5, 7]) + odd_sh = odd << 16 + packed = even | odd_sh + return packed.bitcast(elem_dtype) + return to_elem_vec(dtype_str, elem_dtype, use_hw_cvt_bf16, y) + + def resolve_rmsnorm_weight_dtype(dtype_str: str, weight_dtype_str: str | None = None) -> str: """Return the supported weight dtype for an activation specialization.""" weight_dtype_str = dtype_str if weight_dtype_str is None else weight_dtype_str diff --git a/kernels/norm/rmsnorm_kernel.py b/kernels/norm/rmsnorm_kernel.py index 13863f2773..a981d8ec14 100644 --- a/kernels/norm/rmsnorm_kernel.py +++ b/kernels/norm/rmsnorm_kernel.py @@ -45,6 +45,7 @@ from kernels.norm.rmsnorm_common import store_vec as _store_vec from kernels.norm.rmsnorm_common import to_elem_scalar as _to_elem_scalar from kernels.norm.rmsnorm_common import to_elem_vec as _to_elem_vec +from kernels.norm.rmsnorm_common import to_elem_vec_rna as _to_elem_vec_rna from kernels.norm.rmsnorm_common import weight_vec_width as _weight_vec_width try: @@ -56,6 +57,9 @@ SMALL_N_THRESHOLD = 2048 FWD_MULTI_ROW_THRESHOLD = 8192 +FWD_THROUGHPUT_MIN_ROWS = 8 * WARP_SIZE +FWD_THROUGHPUT_MAX_VEC_ITERS = 8 +DEFAULT_BLOCK_THREADS = BLOCK_THREADS def _quant_dtype_to_elem_type(dtype_str: str): @@ -77,6 +81,7 @@ def build_rmsnorm_module( eps: float = EPS, BLOCK_THREADS: int = BLOCK_THREADS, weight_dtype_str: str | None = None, + _enable_throughput_dispatch: bool = True, ): weight_dtype_str = _resolve_rmsnorm_weight_dtype(dtype_str, weight_dtype_str) arch = get_rocm_arch() @@ -93,6 +98,102 @@ def build_rmsnorm_module( 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]} + # On gfx942, stream large BF16 rows non-temporally so shared gamma remains + # cache-resident. Wide-row output stores are streaming as well. + USE_GFX942_BF16_FAST_PATH = ( + str(arch).startswith("gfx942") and dtype_str == "bf16" and weight_dtype_str == "bf16" and N >= 4096 + ) + PRELOAD_GAMMA = USE_GFX942_BF16_FAST_PATH and VEC_TILES >= BLOCK_THREADS and VEC_TILES % BLOCK_THREADS == 0 + INPUT_CACHE_MODIFIER = 2 if USE_GFX942_BF16_FAST_PATH else 0 + OUTPUT_CACHE_MODIFIER = ( + 2 if USE_GFX942_BF16_FAST_PATH and VEC_TILES > WARP_SIZE * FWD_THROUGHPUT_MAX_VEC_ITERS else 0 + ) + USE_FAST_BF16_OUTPUT = USE_GFX942_BF16_FAST_PATH + + one_wave_vec_iters = (VEC_TILES + WARP_SIZE - 1) // WARP_SIZE if USE_VEC_N else 0 + THROUGHPUT_BLOCK_THREADS = WARP_SIZE if 0 < one_wave_vec_iters <= FWD_THROUGHPUT_MAX_VEC_ITERS else BLOCK_THREADS + use_throughput_dispatch = ( + _enable_throughput_dispatch + and USE_GFX942_BF16_FAST_PATH + and N > SMALL_N_THRESHOLD + and BLOCK_THREADS == DEFAULT_BLOCK_THREADS + and THROUGHPUT_BLOCK_THREADS < BLOCK_THREADS + ) + if use_throughput_dispatch: + # Keep M runtime-only: compile both geometries into one launcher rather + # than selecting a different cached build for each row count. + latency_launch = build_rmsnorm_module( + N, + dtype_str, + store_rstd=store_rstd, + eps=eps, + BLOCK_THREADS=BLOCK_THREADS, + weight_dtype_str=weight_dtype_str, + _enable_throughput_dispatch=False, + ) + throughput_launch = build_rmsnorm_module( + N, + dtype_str, + store_rstd=store_rstd, + eps=eps, + BLOCK_THREADS=THROUGHPUT_BLOCK_THREADS, + weight_dtype_str=weight_dtype_str, + _enable_throughput_dispatch=False, + ) + + if store_rstd: + + @flyc.jit + def launch_rmsnorm( + Input: fx.Tensor, + Gamma: fx.Tensor, + Output: fx.Tensor, + Rstd: fx.Tensor, + m_in: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + def launch_throughput(): + throughput_launch(Input, Gamma, Output, Rstd, m_in, stream) + + def launch_default(): + latency_launch(Input, Gamma, Output, Rstd, m_in, stream) + + @flyc.jit + def dispatch(): + if m_in >= fx.Int32(FWD_THROUGHPUT_MIN_ROWS): + launch_throughput() + else: + launch_default() + + dispatch() + + return launch_rmsnorm + + @flyc.jit + def launch_rmsnorm( + Input: fx.Tensor, + Gamma: fx.Tensor, + Output: fx.Tensor, + m_in: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + def launch_throughput(): + throughput_launch(Input, Gamma, Output, m_in, stream) + + def launch_default(): + latency_launch(Input, Gamma, Output, m_in, stream) + + @flyc.jit + def dispatch(): + if m_in >= fx.Int32(FWD_THROUGHPUT_MIN_ROWS): + launch_throughput() + else: + launch_default() + + dispatch() + + return launch_rmsnorm + SharedStorage = _make_single_reduction_storage(RED_SLOTS) @flyc.kernel(**_kernel_kwargs) @@ -163,7 +264,14 @@ def block_reduce_add(val): out_div = fx.logical_divide(row_out, fx.make_layout(VEC_WIDTH, 1)) gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(_weight_vec_width(weight_dtype_str), 1)) - copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) + copy_atom = fx.make_copy_atom( + fx.rocdl.BufferCopy128b(INPUT_CACHE_MODIFIER), + elem_bits, + ) + output_copy_atom = fx.make_copy_atom( + fx.rocdl.BufferCopy128b(OUTPUT_CACHE_MODIFIER), + elem_bits, + ) gamma_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), weight_elem_bits) if const_expr(TAIL_ELEMS > 0): row_div_s = fx.logical_divide(row_in, fx.make_layout(1, 1)) @@ -178,6 +286,7 @@ def block_reduce_add(val): c_zero_f = fx.Float32(0.0) thread_sumsq = c_zero_f in_local = [] + gamma_local = [] # Pass 1: load + cache + sumsq for tile_i in range_constexpr(NUM_VEC_ITERS): @@ -186,6 +295,8 @@ def block_reduce_add(val): idx_safe = is_valid.select(idx, 0) vec = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx_safe) in_local.append(vec) + if const_expr(PRELOAD_GAMMA): + gamma_local.append(_load_vec(gamma_copy_atom, VEC_WIDTH, weight_elem_dtype, gamma_div, idx_safe)) x = vec.to(fx.Float32) x2 = x * x @@ -212,12 +323,23 @@ def block_reduce_add(val): for tile_i in range_constexpr(NUM_VEC_ITERS): idx = tid + tile_i * BLOCK_THREADS if idx < VEC_TILES: - g = _load_weight_vec(gamma_copy_atom, weight_dtype_str, weight_elem_dtype, gamma_div, idx) + if const_expr(PRELOAD_GAMMA): + g = gamma_local[tile_i].to(fx.Float32) + else: + g = _load_weight_vec(gamma_copy_atom, weight_dtype_str, weight_elem_dtype, gamma_div, idx) 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) - _store_vec(copy_atom, VEC_WIDTH, elem_dtype, out_e, out_div, idx) + if const_expr(USE_FAST_BF16_OUTPUT): + out_e = _to_elem_vec_rna( + dtype_str, + elem_dtype, + USE_HW_CVT_PK_BF16_F32, + y, + ) + else: + out_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, y) + _store_vec(output_copy_atom, VEC_WIDTH, elem_dtype, out_e, out_div, idx) if const_expr(TAIL_ELEMS > 0): if tid < TAIL_ELEMS: @@ -356,6 +478,7 @@ def rmsnorm_direct( dtype_str, BLOCK_THREADS=BLOCK_THREADS, weight_dtype_str=resolved_weight_dtype_str, + _enable_throughput_dispatch=False, ) launch(Input, Gamma, Output, m_in, stream) diff --git a/tests/kernels/benchmark_rmsnorm_backends.py b/tests/kernels/benchmark_rmsnorm_backends.py new file mode 100644 index 0000000000..230e6a0408 --- /dev/null +++ b/tests/kernels/benchmark_rmsnorm_backends.py @@ -0,0 +1,288 @@ +#!/usr/bin/env python3 + +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""Compare plain RMSNorm forward performance against AITER HIP and Triton. + +The dedicated BF16 matrix covers decode-sized row counts, the runtime +one-row/multi-row dispatch boundary, and representative model hidden sizes. +Compilation and correctness checks happen before timing; reported values +contain GPU execution time only. + +Example: + AITER_REPO=/path/to/aiter python3 tests/kernels/benchmark_rmsnorm_backends.py +""" + +from __future__ import annotations + +import argparse +import importlib.util +import os +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Callable +from unittest.mock import patch + +import torch + +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +EPS = 1e-5 + +# Keep performance coverage independent from the correctness matrix. The two +# small-N widths exercise the runtime M crossover; 4096 and 8192 represent +# common model hidden sizes that stay on the one-row implementation. +PLAIN_RMSNORM_PERF_M = (1, 16, 128, 1024, 8192, 8193, 32768) +PLAIN_RMSNORM_PERF_N = (512, 1024, 2048, 4096, 8192) +PLAIN_RMSNORM_PERF_CONFIGS = tuple((m, n, "bf16") for n in PLAIN_RMSNORM_PERF_N for m in PLAIN_RMSNORM_PERF_M) + + +@dataclass(frozen=True) +class BenchResult: + shape: str + dtype: str + flydsl_us: float + hip_us: float | None + triton_us: float | None + hip_status: str + triton_status: str + + +def _torch_dtype(dtype: str) -> torch.dtype: + return { + "f32": torch.float32, + "f16": torch.float16, + "bf16": torch.bfloat16, + }[dtype] + + +def _reference(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + x_f32 = x.float() + weight_f32 = weight.float() + rrms = torch.rsqrt(x_f32.square().mean(dim=-1, keepdim=True) + EPS) + return (x_f32 * rrms * weight_f32).to(x.dtype) + + +def _check_output(backend: str, output: torch.Tensor, reference: torch.Tensor, dtype: str) -> None: + tolerance = {"f32": 1e-4, "f16": 1e-2, "bf16": 2e-2}[dtype] + try: + torch.testing.assert_close( + output.float(), + reference.float(), + rtol=tolerance, + atol=tolerance, + ) + except AssertionError as exc: + max_error = (output.float() - reference.float()).abs().max().item() + raise AssertionError( + f"{backend} correctness failed for shape={tuple(output.shape)}, " + f"dtype={dtype}, max_abs_error={max_error:.3e}" + ) from exc + + +def _bench_kernel_us(run: Callable[[], None], *, warmup: int, iters: int) -> float: + """Return mean GPU kernel time, excluding Python/backend wrapper time.""" + for _ in range(warmup): + run() + torch.cuda.synchronize() + + import torch.profiler as tpf + + with tpf.profile( + activities=[tpf.ProfilerActivity.CPU, tpf.ProfilerActivity.CUDA], + profile_memory=False, + with_stack=False, + with_modules=False, + ) as prof: + for _ in range(iters): + run() + torch.cuda.synchronize() + + device_us = sum( + event.self_device_time_total for event in prof.events() if str(event.device_type).split(".")[-1] == "CUDA" + ) + if device_us <= 0: + raise RuntimeError("torch profiler did not record any GPU kernel time") + return device_us / iters + + +def _load_aiter(aiter_repo: str | None): + if aiter_repo: + repo_path = Path(aiter_repo).expanduser().resolve() + if not (repo_path / "aiter").is_dir(): + raise ValueError(f"--aiter-repo does not contain an aiter package: {repo_path}") + sys.path.insert(0, str(repo_path)) + + original_find_spec = importlib.util.find_spec + + def find_baseline_dependency(name, *args, **kwargs): + # Do not let AITER recursively load its optional FlyDSL kernels. This + # process uses the current checkout as the candidate implementation. + if name == "flydsl" or name.startswith("flydsl."): + return None + return original_find_spec(name, *args, **kwargs) + + try: + with patch("importlib.util.find_spec", side_effect=find_baseline_dependency): + import aiter + from aiter.ops.rmsnorm import rmsnorm as hip_rmsnorm + from aiter.ops.triton.normalization.rmsnorm import rms_norm as triton_rmsnorm + except Exception as exc: + raise RuntimeError( + "AITER HIP/Triton RMSNorm is unavailable. Install AITER or set AITER_REPO=/path/to/aiter." + ) from exc + + return aiter, hip_rmsnorm, triton_rmsnorm + + +def _compile_flydsl( + x: torch.Tensor, + weight: torch.Tensor, + output: torch.Tensor, + dtype: str, +) -> Callable[[], None]: + import flydsl.compiler as flyc + from kernels.norm.rmsnorm_kernel import build_rmsnorm_module + + m, n = x.shape + launch_fn = build_rmsnorm_module(n, dtype) + stream = torch.cuda.current_stream() + compiled = flyc.compile(launch_fn, x, weight, output, m, stream) + + def run(): + compiled(x, weight, output, m, stream) + + return run + + +def _run_case( + m: int, + n: int, + dtype: str, + hip_rmsnorm, + triton_rmsnorm, + warmup: int, + iters: int, +) -> BenchResult: + torch.manual_seed(42) + torch_dtype = _torch_dtype(dtype) + x = torch.randn((m, n), device="cuda", dtype=torch_dtype).contiguous() + weight = torch.rand((n,), device="cuda", dtype=torch_dtype).contiguous() + reference = _reference(x, weight) + + flydsl_output = torch.empty_like(x) + run_flydsl = _compile_flydsl(x, weight, flydsl_output, dtype) + run_flydsl() + torch.cuda.synchronize() + _check_output("FlyDSL", flydsl_output, reference, dtype) + + flydsl_us = _bench_kernel_us(run_flydsl, warmup=warmup, iters=iters) + hip_us = None + hip_status = "SKIP: f32 unsupported" + if dtype in ("f16", "bf16"): + try: + hip_output = torch.empty_like(x) + + def run_hip(): + hip_rmsnorm(hip_output, x, weight, EPS) + + run_hip() + torch.cuda.synchronize() + _check_output("HIP", hip_output, reference, dtype) + hip_us = _bench_kernel_us(run_hip, warmup=warmup, iters=iters) + hip_status = "ok" + except Exception as exc: + hip_status = f"ERR: {type(exc).__name__}: {exc}" + + triton_us = None + triton_status = "ok" + try: + triton_output = triton_rmsnorm(x, weight, EPS) + torch.cuda.synchronize() + _check_output("Triton", triton_output, reference, dtype) + triton_us = _bench_kernel_us(lambda: triton_rmsnorm(x, weight, EPS), warmup=warmup, iters=iters) + except Exception as exc: + triton_status = f"ERR: {type(exc).__name__}: {exc}" + + return BenchResult( + shape=f"{m}x{n}", + dtype=dtype, + flydsl_us=flydsl_us, + hip_us=hip_us, + triton_us=triton_us, + hip_status=hip_status, + triton_status=triton_status, + ) + + +def _fmt_us(value: float | None, status: str) -> str: + if value is not None: + return f"{value:.2f}" + return "SKIP" if status.startswith("SKIP") else "ERR" + + +def _fmt_speedup(flydsl_us: float, baseline_us: float | None) -> str: + return "-" if baseline_us is None else f"{baseline_us / flydsl_us:.2f}x" + + +def _print_results(results: list[BenchResult]) -> None: + print() + print( + f"{'shape':>12} {'dtype':>6} {'FlyDSL(us)':>12} {'HIP(us)':>10} " + f"{'Triton(us)':>12} {'FlyDSL/HIP':>12} {'FlyDSL/Triton':>16}" + ) + for result in results: + print( + f"{result.shape:>12} {result.dtype:>6} {result.flydsl_us:>12.2f} " + f"{_fmt_us(result.hip_us, result.hip_status):>10} " + f"{_fmt_us(result.triton_us, result.triton_status):>12} " + f"{_fmt_speedup(result.flydsl_us, result.hip_us):>12} " + f"{_fmt_speedup(result.flydsl_us, result.triton_us):>16}" + ) + print("\nSpeedups greater than 1.0x mean FlyDSL is faster than that baseline.") + for result in results: + if result.hip_status != "ok": + print(f"{result.shape} {result.dtype} HIP: {result.hip_status}") + if result.triton_status != "ok": + print(f"{result.shape} {result.dtype} Triton: {result.triton_status}") + + +def _parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--aiter-repo", + default=os.environ.get("AITER_REPO"), + help="AITER checkout to import when it is not installed (default: AITER_REPO).", + ) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--iters", type=int, default=100) + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("ROCm GPU is required") + if args.warmup < 0 or args.iters <= 0: + raise ValueError("--warmup must be >= 0 and --iters must be > 0") + + aiter, hip_rmsnorm, triton_rmsnorm = _load_aiter(args.aiter_repo) + configs = PLAIN_RMSNORM_PERF_CONFIGS + device = torch.cuda.current_device() + properties = torch.cuda.get_device_properties(device) + print(f"Device: {torch.cuda.get_device_name(device)} ({getattr(properties, 'gcnArchName', 'unknown')})") + print(f"AITER: {aiter.__file__}") + print(f"Cases: {configs}") + + results = [_run_case(m, n, dtype, hip_rmsnorm, triton_rmsnorm, args.warmup, args.iters) for m, n, dtype in configs] + _print_results(results) + if any(result.hip_status.startswith("ERR") or result.triton_status.startswith("ERR") for result in results): + raise SystemExit(1) + + +if __name__ == "__main__": + main() From b98ff768d95ffdbee4ad309803e686da51bf96bf Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Tue, 1 Sep 2026 12:34:50 +0800 Subject: [PATCH 17/25] Remove RMSNorm benchmark script from tracking --- tests/kernels/benchmark_rmsnorm_backends.py | 288 -------------------- 1 file changed, 288 deletions(-) delete mode 100644 tests/kernels/benchmark_rmsnorm_backends.py diff --git a/tests/kernels/benchmark_rmsnorm_backends.py b/tests/kernels/benchmark_rmsnorm_backends.py deleted file mode 100644 index 230e6a0408..0000000000 --- a/tests/kernels/benchmark_rmsnorm_backends.py +++ /dev/null @@ -1,288 +0,0 @@ -#!/usr/bin/env python3 - -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2025 FlyDSL Project Contributors - -"""Compare plain RMSNorm forward performance against AITER HIP and Triton. - -The dedicated BF16 matrix covers decode-sized row counts, the runtime -one-row/multi-row dispatch boundary, and representative model hidden sizes. -Compilation and correctness checks happen before timing; reported values -contain GPU execution time only. - -Example: - AITER_REPO=/path/to/aiter python3 tests/kernels/benchmark_rmsnorm_backends.py -""" - -from __future__ import annotations - -import argparse -import importlib.util -import os -import sys -from dataclasses import dataclass -from pathlib import Path -from typing import Callable -from unittest.mock import patch - -import torch - -_REPO_ROOT = Path(__file__).resolve().parents[2] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) - -EPS = 1e-5 - -# Keep performance coverage independent from the correctness matrix. The two -# small-N widths exercise the runtime M crossover; 4096 and 8192 represent -# common model hidden sizes that stay on the one-row implementation. -PLAIN_RMSNORM_PERF_M = (1, 16, 128, 1024, 8192, 8193, 32768) -PLAIN_RMSNORM_PERF_N = (512, 1024, 2048, 4096, 8192) -PLAIN_RMSNORM_PERF_CONFIGS = tuple((m, n, "bf16") for n in PLAIN_RMSNORM_PERF_N for m in PLAIN_RMSNORM_PERF_M) - - -@dataclass(frozen=True) -class BenchResult: - shape: str - dtype: str - flydsl_us: float - hip_us: float | None - triton_us: float | None - hip_status: str - triton_status: str - - -def _torch_dtype(dtype: str) -> torch.dtype: - return { - "f32": torch.float32, - "f16": torch.float16, - "bf16": torch.bfloat16, - }[dtype] - - -def _reference(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: - x_f32 = x.float() - weight_f32 = weight.float() - rrms = torch.rsqrt(x_f32.square().mean(dim=-1, keepdim=True) + EPS) - return (x_f32 * rrms * weight_f32).to(x.dtype) - - -def _check_output(backend: str, output: torch.Tensor, reference: torch.Tensor, dtype: str) -> None: - tolerance = {"f32": 1e-4, "f16": 1e-2, "bf16": 2e-2}[dtype] - try: - torch.testing.assert_close( - output.float(), - reference.float(), - rtol=tolerance, - atol=tolerance, - ) - except AssertionError as exc: - max_error = (output.float() - reference.float()).abs().max().item() - raise AssertionError( - f"{backend} correctness failed for shape={tuple(output.shape)}, " - f"dtype={dtype}, max_abs_error={max_error:.3e}" - ) from exc - - -def _bench_kernel_us(run: Callable[[], None], *, warmup: int, iters: int) -> float: - """Return mean GPU kernel time, excluding Python/backend wrapper time.""" - for _ in range(warmup): - run() - torch.cuda.synchronize() - - import torch.profiler as tpf - - with tpf.profile( - activities=[tpf.ProfilerActivity.CPU, tpf.ProfilerActivity.CUDA], - profile_memory=False, - with_stack=False, - with_modules=False, - ) as prof: - for _ in range(iters): - run() - torch.cuda.synchronize() - - device_us = sum( - event.self_device_time_total for event in prof.events() if str(event.device_type).split(".")[-1] == "CUDA" - ) - if device_us <= 0: - raise RuntimeError("torch profiler did not record any GPU kernel time") - return device_us / iters - - -def _load_aiter(aiter_repo: str | None): - if aiter_repo: - repo_path = Path(aiter_repo).expanduser().resolve() - if not (repo_path / "aiter").is_dir(): - raise ValueError(f"--aiter-repo does not contain an aiter package: {repo_path}") - sys.path.insert(0, str(repo_path)) - - original_find_spec = importlib.util.find_spec - - def find_baseline_dependency(name, *args, **kwargs): - # Do not let AITER recursively load its optional FlyDSL kernels. This - # process uses the current checkout as the candidate implementation. - if name == "flydsl" or name.startswith("flydsl."): - return None - return original_find_spec(name, *args, **kwargs) - - try: - with patch("importlib.util.find_spec", side_effect=find_baseline_dependency): - import aiter - from aiter.ops.rmsnorm import rmsnorm as hip_rmsnorm - from aiter.ops.triton.normalization.rmsnorm import rms_norm as triton_rmsnorm - except Exception as exc: - raise RuntimeError( - "AITER HIP/Triton RMSNorm is unavailable. Install AITER or set AITER_REPO=/path/to/aiter." - ) from exc - - return aiter, hip_rmsnorm, triton_rmsnorm - - -def _compile_flydsl( - x: torch.Tensor, - weight: torch.Tensor, - output: torch.Tensor, - dtype: str, -) -> Callable[[], None]: - import flydsl.compiler as flyc - from kernels.norm.rmsnorm_kernel import build_rmsnorm_module - - m, n = x.shape - launch_fn = build_rmsnorm_module(n, dtype) - stream = torch.cuda.current_stream() - compiled = flyc.compile(launch_fn, x, weight, output, m, stream) - - def run(): - compiled(x, weight, output, m, stream) - - return run - - -def _run_case( - m: int, - n: int, - dtype: str, - hip_rmsnorm, - triton_rmsnorm, - warmup: int, - iters: int, -) -> BenchResult: - torch.manual_seed(42) - torch_dtype = _torch_dtype(dtype) - x = torch.randn((m, n), device="cuda", dtype=torch_dtype).contiguous() - weight = torch.rand((n,), device="cuda", dtype=torch_dtype).contiguous() - reference = _reference(x, weight) - - flydsl_output = torch.empty_like(x) - run_flydsl = _compile_flydsl(x, weight, flydsl_output, dtype) - run_flydsl() - torch.cuda.synchronize() - _check_output("FlyDSL", flydsl_output, reference, dtype) - - flydsl_us = _bench_kernel_us(run_flydsl, warmup=warmup, iters=iters) - hip_us = None - hip_status = "SKIP: f32 unsupported" - if dtype in ("f16", "bf16"): - try: - hip_output = torch.empty_like(x) - - def run_hip(): - hip_rmsnorm(hip_output, x, weight, EPS) - - run_hip() - torch.cuda.synchronize() - _check_output("HIP", hip_output, reference, dtype) - hip_us = _bench_kernel_us(run_hip, warmup=warmup, iters=iters) - hip_status = "ok" - except Exception as exc: - hip_status = f"ERR: {type(exc).__name__}: {exc}" - - triton_us = None - triton_status = "ok" - try: - triton_output = triton_rmsnorm(x, weight, EPS) - torch.cuda.synchronize() - _check_output("Triton", triton_output, reference, dtype) - triton_us = _bench_kernel_us(lambda: triton_rmsnorm(x, weight, EPS), warmup=warmup, iters=iters) - except Exception as exc: - triton_status = f"ERR: {type(exc).__name__}: {exc}" - - return BenchResult( - shape=f"{m}x{n}", - dtype=dtype, - flydsl_us=flydsl_us, - hip_us=hip_us, - triton_us=triton_us, - hip_status=hip_status, - triton_status=triton_status, - ) - - -def _fmt_us(value: float | None, status: str) -> str: - if value is not None: - return f"{value:.2f}" - return "SKIP" if status.startswith("SKIP") else "ERR" - - -def _fmt_speedup(flydsl_us: float, baseline_us: float | None) -> str: - return "-" if baseline_us is None else f"{baseline_us / flydsl_us:.2f}x" - - -def _print_results(results: list[BenchResult]) -> None: - print() - print( - f"{'shape':>12} {'dtype':>6} {'FlyDSL(us)':>12} {'HIP(us)':>10} " - f"{'Triton(us)':>12} {'FlyDSL/HIP':>12} {'FlyDSL/Triton':>16}" - ) - for result in results: - print( - f"{result.shape:>12} {result.dtype:>6} {result.flydsl_us:>12.2f} " - f"{_fmt_us(result.hip_us, result.hip_status):>10} " - f"{_fmt_us(result.triton_us, result.triton_status):>12} " - f"{_fmt_speedup(result.flydsl_us, result.hip_us):>12} " - f"{_fmt_speedup(result.flydsl_us, result.triton_us):>16}" - ) - print("\nSpeedups greater than 1.0x mean FlyDSL is faster than that baseline.") - for result in results: - if result.hip_status != "ok": - print(f"{result.shape} {result.dtype} HIP: {result.hip_status}") - if result.triton_status != "ok": - print(f"{result.shape} {result.dtype} Triton: {result.triton_status}") - - -def _parse_args(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--aiter-repo", - default=os.environ.get("AITER_REPO"), - help="AITER checkout to import when it is not installed (default: AITER_REPO).", - ) - parser.add_argument("--warmup", type=int, default=10) - parser.add_argument("--iters", type=int, default=100) - return parser.parse_args() - - -def main() -> None: - args = _parse_args() - if not torch.cuda.is_available(): - raise RuntimeError("ROCm GPU is required") - if args.warmup < 0 or args.iters <= 0: - raise ValueError("--warmup must be >= 0 and --iters must be > 0") - - aiter, hip_rmsnorm, triton_rmsnorm = _load_aiter(args.aiter_repo) - configs = PLAIN_RMSNORM_PERF_CONFIGS - device = torch.cuda.current_device() - properties = torch.cuda.get_device_properties(device) - print(f"Device: {torch.cuda.get_device_name(device)} ({getattr(properties, 'gcnArchName', 'unknown')})") - print(f"AITER: {aiter.__file__}") - print(f"Cases: {configs}") - - results = [_run_case(m, n, dtype, hip_rmsnorm, triton_rmsnorm, args.warmup, args.iters) for m, n, dtype in configs] - _print_results(results) - if any(result.hip_status.startswith("ERR") or result.triton_status.startswith("ERR") for result in results): - raise SystemExit(1) - - -if __name__ == "__main__": - main() From cae767d63cfd1c13d949d13277001115deea9032 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Thu, 3 Sep 2026 16:54:57 +0800 Subject: [PATCH 18/25] restore the RNA semantic to RNE --- kernels/norm/rmsnorm_common.py | 12 ------------ kernels/norm/rmsnorm_kernel.py | 12 +----------- 2 files changed, 1 insertion(+), 23 deletions(-) diff --git a/kernels/norm/rmsnorm_common.py b/kernels/norm/rmsnorm_common.py index 8109dd2d9e..72fd273439 100644 --- a/kernels/norm/rmsnorm_common.py +++ b/kernels/norm/rmsnorm_common.py @@ -117,18 +117,6 @@ def to_elem_vec(dtype_str: str, elem_dtype, use_hw_cvt_bf16: bool, y): return y.to(elem_dtype) -def to_elem_vec_rna(dtype_str: str, elem_dtype, use_hw_cvt_bf16: bool, y): - """Pack BF16 with round-to-nearest-away, avoiding the RNE tie fixup.""" - if const_expr(dtype_str == "bf16" and not use_hw_cvt_bf16): - bf16_bits = (y.bitcast(fx.Uint32) + 0x8000) >> 16 - even = bf16_bits.shuffle(bf16_bits, [0, 2, 4, 6]) - odd = bf16_bits.shuffle(bf16_bits, [1, 3, 5, 7]) - odd_sh = odd << 16 - packed = even | odd_sh - return packed.bitcast(elem_dtype) - return to_elem_vec(dtype_str, elem_dtype, use_hw_cvt_bf16, y) - - def resolve_rmsnorm_weight_dtype(dtype_str: str, weight_dtype_str: str | None = None) -> str: """Return the supported weight dtype for an activation specialization.""" weight_dtype_str = dtype_str if weight_dtype_str is None else weight_dtype_str diff --git a/kernels/norm/rmsnorm_kernel.py b/kernels/norm/rmsnorm_kernel.py index a981d8ec14..a7ca360359 100644 --- a/kernels/norm/rmsnorm_kernel.py +++ b/kernels/norm/rmsnorm_kernel.py @@ -45,7 +45,6 @@ from kernels.norm.rmsnorm_common import store_vec as _store_vec from kernels.norm.rmsnorm_common import to_elem_scalar as _to_elem_scalar from kernels.norm.rmsnorm_common import to_elem_vec as _to_elem_vec -from kernels.norm.rmsnorm_common import to_elem_vec_rna as _to_elem_vec_rna from kernels.norm.rmsnorm_common import weight_vec_width as _weight_vec_width try: @@ -108,7 +107,6 @@ def build_rmsnorm_module( OUTPUT_CACHE_MODIFIER = ( 2 if USE_GFX942_BF16_FAST_PATH and VEC_TILES > WARP_SIZE * FWD_THROUGHPUT_MAX_VEC_ITERS else 0 ) - USE_FAST_BF16_OUTPUT = USE_GFX942_BF16_FAST_PATH one_wave_vec_iters = (VEC_TILES + WARP_SIZE - 1) // WARP_SIZE if USE_VEC_N else 0 THROUGHPUT_BLOCK_THREADS = WARP_SIZE if 0 < one_wave_vec_iters <= FWD_THROUGHPUT_MAX_VEC_ITERS else BLOCK_THREADS @@ -330,15 +328,7 @@ def block_reduce_add(val): x = in_local[tile_i].to(fx.Float32) y = (x * rrms) * g - if const_expr(USE_FAST_BF16_OUTPUT): - out_e = _to_elem_vec_rna( - dtype_str, - elem_dtype, - USE_HW_CVT_PK_BF16_F32, - y, - ) - else: - out_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, y) + out_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, y) _store_vec(output_copy_atom, VEC_WIDTH, elem_dtype, out_e, out_div, idx) if const_expr(TAIL_ELEMS > 0): From 21c0eae44a93f88851b44395bcfbac68d633ce5f Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Fri, 11 Sep 2026 15:07:14 +0800 Subject: [PATCH 19/25] perf(rmsnorm): dispatch gfx95x BF16 NT cache policy on row count Compile both cache policies into one gfx95x launcher and pick between them on the runtime row count, since streaming only helps once the working set outgrows the last-level cache. Widen the default block for wide gfx95x rows and mirror it in the autotune default config. Rename _build_rmsnorm_large_m_small_n_module to _build_rmsnorm_small_n_module. Co-authored-by: Cursor --- kernels/norm/rmsnorm_autotune.py | 13 ++-- kernels/norm/rmsnorm_kernel.py | 122 +++++++++++++++++++++++++++---- 2 files changed, 115 insertions(+), 20 deletions(-) diff --git a/kernels/norm/rmsnorm_autotune.py b/kernels/norm/rmsnorm_autotune.py index deed875791..930d3cb16d 100644 --- a/kernels/norm/rmsnorm_autotune.py +++ b/kernels/norm/rmsnorm_autotune.py @@ -4,11 +4,8 @@ """Two-track RMSNorm autotuning through the normal direct JIT path.""" from flydsl.autotune import Config, autotune -from kernels.norm.rmsnorm_common import ( - BLOCK_THREADS, - resolve_rmsnorm_weight_dtype, -) -from kernels.norm.rmsnorm_kernel import SMALL_N_THRESHOLD, rmsnorm_direct +from kernels.norm.rmsnorm_common import resolve_rmsnorm_weight_dtype +from kernels.norm.rmsnorm_kernel import SMALL_N_THRESHOLD, default_block_threads, rmsnorm_direct _SEARCH_CONFIGS = ( Config(BLOCK_THREADS=128), @@ -21,8 +18,10 @@ ) -def _default_config(*_args, **_kwargs): - return Config(BLOCK_THREADS=BLOCK_THREADS) +def _default_config(*_args, **kwargs): + # Mirror the untuned launcher so opting into autotuning without a search + # does not fall back to a narrower block than the plain path would pick. + return Config(BLOCK_THREADS=default_block_threads(kwargs.get("N", 0))) def _search_configs( diff --git a/kernels/norm/rmsnorm_kernel.py b/kernels/norm/rmsnorm_kernel.py index a7ca360359..d1f3cc7710 100644 --- a/kernels/norm/rmsnorm_kernel.py +++ b/kernels/norm/rmsnorm_kernel.py @@ -59,6 +59,24 @@ FWD_THROUGHPUT_MIN_ROWS = 8 * WARP_SIZE FWD_THROUGHPUT_MAX_VEC_ITERS = 8 DEFAULT_BLOCK_THREADS = BLOCK_THREADS +WIDE_ROW_MIN_N = 8192 +WIDE_ROW_BLOCK_THREADS = 512 + +# Both CDNA3 and CDNA4 share a 256MB last-level cache. Non-temporal streaming +# only pays off once input plus output stop fitting in it; below that it throws +# away reuse the cache would have served. The two policies measure as a tie at +# 1.125x LLC and streaming wins outright past it, uniformly across row widths. +LLC_BYTES = 256 * 1024 * 1024 +NT_MIN_WORKING_SET = LLC_BYTES * 9 // 8 + + +def default_block_threads(N: int, arch=None) -> int: + """Wide rows hold fewer resident vectors per thread at 512 than at 256.""" + if arch is None: + arch = get_rocm_arch() + if str(arch).startswith("gfx95") and N >= WIDE_ROW_MIN_N: + return WIDE_ROW_BLOCK_THREADS + return DEFAULT_BLOCK_THREADS def _quant_dtype_to_elem_type(dtype_str: str): @@ -78,13 +96,16 @@ def build_rmsnorm_module( dtype_str: str, store_rstd: bool = False, eps: float = EPS, - BLOCK_THREADS: int = BLOCK_THREADS, + BLOCK_THREADS: int | None = None, weight_dtype_str: str | None = None, _enable_throughput_dispatch: bool = True, + _force_nt: bool | None = None, ): weight_dtype_str = _resolve_rmsnorm_weight_dtype(dtype_str, weight_dtype_str) arch = get_rocm_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") + if BLOCK_THREADS is None: + BLOCK_THREADS = default_block_threads(N, arch) # BLOCK_THREADS controls storage, tiling, and launch geometry. RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) @@ -97,13 +118,21 @@ def build_rmsnorm_module( 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]} - # On gfx942, stream large BF16 rows non-temporally so shared gamma remains - # cache-resident. Wide-row output stores are streaming as well. - USE_GFX942_BF16_FAST_PATH = ( - str(arch).startswith("gfx942") and dtype_str == "bf16" and weight_dtype_str == "bf16" and N >= 4096 - ) + # Stream large BF16 rows non-temporally. x stays in registers across both + # passes, so x and y are each touched once and neither gains from being + # cached; letting x fills allocate in the memory-side cache only evicts y's + # pending stores. gfx942 streams unconditionally; gfx95x defers the choice + # to a runtime row-count dispatch below. + BF16_ROWS = dtype_str == "bf16" and weight_dtype_str == "bf16" + USE_GFX942_BF16_FAST_PATH = str(arch).startswith("gfx942") and BF16_ROWS and N >= 4096 + # Whether streaming wins is set by the working set, not the row width, so on + # gfx95x the runtime dispatch below covers every row the plain kernel handles. + use_nt_dispatch = str(arch).startswith("gfx95") and BF16_ROWS and N > SMALL_N_THRESHOLD and _force_nt is None + USE_NT = USE_GFX942_BF16_FAST_PATH or bool(_force_nt) PRELOAD_GAMMA = USE_GFX942_BF16_FAST_PATH and VEC_TILES >= BLOCK_THREADS and VEC_TILES % BLOCK_THREADS == 0 - INPUT_CACHE_MODIFIER = 2 if USE_GFX942_BF16_FAST_PATH else 0 + INPUT_CACHE_MODIFIER = 2 if USE_NT else 0 + # Streaming the stores too costs 6-8% on gfx95x; letting them settle in LLC + # drains them to HBM in better-scheduled bursts. gfx942 keeps its own tuning. OUTPUT_CACHE_MODIFIER = ( 2 if USE_GFX942_BF16_FAST_PATH and VEC_TILES > WARP_SIZE * FWD_THROUGHPUT_MAX_VEC_ITERS else 0 ) @@ -192,6 +221,73 @@ def dispatch(): return launch_rmsnorm + if use_nt_dispatch: + # Keep M runtime-only: compile both cache policies into one launcher + # rather than selecting a different cached build for each row count. + _nt_kwargs = dict( + store_rstd=store_rstd, + eps=eps, + BLOCK_THREADS=BLOCK_THREADS, + weight_dtype_str=weight_dtype_str, + _enable_throughput_dispatch=False, + ) + cached_launch = build_rmsnorm_module(N, dtype_str, _force_nt=False, **_nt_kwargs) + streaming_launch = build_rmsnorm_module(N, dtype_str, _force_nt=True, **_nt_kwargs) + NT_MIN_ROWS = NT_MIN_WORKING_SET // (N * 4) + + if store_rstd: + + @flyc.jit + def launch_rmsnorm( + Input: fx.Tensor, + Gamma: fx.Tensor, + Output: fx.Tensor, + Rstd: fx.Tensor, + m_in: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + def launch_streaming(): + streaming_launch(Input, Gamma, Output, Rstd, m_in, stream) + + def launch_cached(): + cached_launch(Input, Gamma, Output, Rstd, m_in, stream) + + @flyc.jit + def dispatch(): + if m_in > fx.Int32(NT_MIN_ROWS): + launch_streaming() + else: + launch_cached() + + dispatch() + + return launch_rmsnorm + + @flyc.jit + def launch_rmsnorm( + Input: fx.Tensor, + Gamma: fx.Tensor, + Output: fx.Tensor, + m_in: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + def launch_streaming(): + streaming_launch(Input, Gamma, Output, m_in, stream) + + def launch_cached(): + cached_launch(Input, Gamma, Output, m_in, stream) + + @flyc.jit + def dispatch(): + if m_in > fx.Int32(NT_MIN_ROWS): + launch_streaming() + else: + launch_cached() + + dispatch() + + return launch_rmsnorm + SharedStorage = _make_single_reduction_storage(RED_SLOTS) @flyc.kernel(**_kernel_kwargs) @@ -399,7 +495,7 @@ def block_reduce_add(val): _store_scalar(copy_atom_s, elem_dtype, out_div, idx, y_e) if N <= SMALL_N_THRESHOLD: - return _build_rmsnorm_large_m_small_n_module( + return _build_rmsnorm_small_n_module( N, dtype_str, store_rstd, @@ -473,7 +569,7 @@ def rmsnorm_direct( launch(Input, Gamma, Output, m_in, stream) -def _build_rmsnorm_large_m_small_n_module( +def _build_rmsnorm_small_n_module( N: int, dtype_str: str, store_rstd: bool = False, @@ -658,7 +754,7 @@ def group_reduce_add(x): if store_rstd: @flyc.jit - def launch_rmsnorm_large_m_small_n( + def launch_rmsnorm_small_n( Input: fx.Tensor, Gamma: fx.Tensor, Output: fx.Tensor, @@ -691,10 +787,10 @@ def dispatch(): dispatch() - return launch_rmsnorm_large_m_small_n + return launch_rmsnorm_small_n @flyc.jit - def launch_rmsnorm_large_m_small_n( + def launch_rmsnorm_small_n( Input: fx.Tensor, Gamma: fx.Tensor, Output: fx.Tensor, @@ -727,7 +823,7 @@ def dispatch(): dispatch() - return launch_rmsnorm_large_m_small_n + return launch_rmsnorm_small_n def build_fused_add_rmsnorm_module( From cf1151b9e39775a5ad672003622a0e2af2a3bef1 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Fri, 18 Sep 2026 05:26:08 +0800 Subject: [PATCH 20/25] perf(rmsnorm): extend vec8 coverage to fused-add and quant variants Use masked vector iterations for vec8-aligned row widths instead of requiring an exact tile multiple, including a partially active final iteration in the test matrix. Co-authored-by: Cursor --- kernels/norm/rmsnorm_kernel.py | 130 +++++++++++++++++++-------------- tests/kernels/test_rmsnorm.py | 2 + 2 files changed, 77 insertions(+), 55 deletions(-) diff --git a/kernels/norm/rmsnorm_kernel.py b/kernels/norm/rmsnorm_kernel.py index d1f3cc7710..2e59ba81d4 100644 --- a/kernels/norm/rmsnorm_kernel.py +++ b/kernels/norm/rmsnorm_kernel.py @@ -837,10 +837,12 @@ def build_fused_add_rmsnorm_module( arch = get_rocm_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") - 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 weight_elem_bits = 32 if weight_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 SharedStorage = _make_reduction_storage(RED_SLOTS) @@ -916,9 +918,8 @@ def block_reduce_add2(val0, val1): return fx.memref_load(s_red, 0), fx.memref_load(s_red2, 0) - # Fast path for complete 128-bit tiles. - if const_expr(N >= tile_cols and N % tile_cols == 0 and elem_bits <= 16): - num_tiles = N // tile_cols + # Vector path for every complete f16/bf16 vec8 row. + if const_expr(USE_VEC_N): Input_buf = fx.rocdl.make_buffer_tensor(Input) ResidualIn_buf = fx.rocdl.make_buffer_tensor(ResidualIn) Gamma_buf = fx.rocdl.make_buffer_tensor(Gamma) @@ -945,19 +946,22 @@ def block_reduce_add2(val0, val1): add_local = [] # Pass 1: add + cache + sumsq (also write residual_out) - for tile_i in range_constexpr(num_tiles): + for tile_i in range_constexpr(NUM_VEC_ITERS): idx = tid + tile_i * BLOCK_THREADS - x = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx).to(fx.Float32) - residual = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, residual_in_div, idx).to(fx.Float32) + is_valid = idx < VEC_TILES + idx_safe = is_valid.select(idx, 0) + x = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx_safe).to(fx.Float32) + residual = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, residual_in_div, idx_safe).to(fx.Float32) added_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, x + residual) add_local.append(added_e) added = added_e if dtype_str == "f32" else added_e.to(fx.Float32) added2 = added * added red2 = added2.reduce(ReductionOp.ADD, fastmath=fm_fast) - thread_sumsq = thread_sumsq + red2 + thread_sumsq = thread_sumsq + is_valid.select(red2, c_zero_f) - _store_vec(copy_atom, VEC_WIDTH, elem_dtype, added_e, residual_out_div, idx) + if idx < VEC_TILES: + _store_vec(copy_atom, VEC_WIDTH, elem_dtype, added_e, residual_out_div, idx) _, sum_sq = block_reduce_add2(thread_dummy, thread_sumsq) mean_sq = sum_sq / n_float @@ -969,13 +973,14 @@ def block_reduce_add2(val0, val1): _store_scalar(rstd_copy_atom, fx.Float32, rstd_div, bid, rrms) # Pass 2: normalize + gamma + store (reuse cached added values) - 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_weight_vec(gamma_copy_atom, weight_dtype_str, weight_elem_dtype, gamma_div, idx) - added = add_local[tile_i] if dtype_str == "f32" else add_local[tile_i].to(fx.Float32) - y = (added * 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 idx < VEC_TILES: + g = _load_weight_vec(gamma_copy_atom, weight_dtype_str, weight_elem_dtype, gamma_div, idx) + added = add_local[tile_i] if dtype_str == "f32" else add_local[tile_i].to(fx.Float32) + y = (added * 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) else: # Scalar fallback for arbitrary N. @@ -1093,10 +1098,12 @@ def _build_rmsnorm_quant_module( 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) elem_bits = 32 if dtype_str == "f32" else 16 quant_dtype_max = _quant_dtype_max(quant_dtype_str) + 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 SharedStorage = _make_reduction_storage(RED_SLOTS) @@ -1211,9 +1218,8 @@ def block_reduce_max(val): return fx.memref_load(s_red, 0) - # Fast path for complete 128-bit tiles. - if const_expr(N >= tile_cols and N % tile_cols == 0 and elem_bits <= 16): - num_tiles = N // tile_cols + # Vector path for every complete f16/bf16 vec8 row. + if const_expr(USE_VEC_N): quant_half_width = VEC_WIDTH // 2 abs_mask = fx.Vector.filled(VEC_WIDTH, 0x7FFFFFFF, fx.Uint32) Input_buf = fx.rocdl.make_buffer_tensor(Input) @@ -1241,14 +1247,16 @@ def block_reduce_max(val): 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 @@ -1259,20 +1267,22 @@ def block_reduce_max(val): y_local = [] # Pass 2: normalize + gamma (+ optional smooth scale), cache output, and accumulate row max - for tile_i in range_constexpr(num_tiles): + for tile_i in range_constexpr(NUM_VEC_ITERS): idx = tid + tile_i * BLOCK_THREADS + is_valid = idx < VEC_TILES + idx_safe = is_valid.select(idx, 0) - g = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, idx).to(fx.Float32) + g = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, idx_safe).to(fx.Float32) x = in_local[tile_i].to(fx.Float32) y = (x * rrms) * g if const_expr(is_smooth): - s = _load_vec(copy_atom_xs, VEC_WIDTH, elem_dtype, xscale_div, idx).to(fx.Float32) + s = _load_vec(copy_atom_xs, VEC_WIDTH, elem_dtype, xscale_div, idx_safe).to(fx.Float32) y = y * s y_local.append(y) y_abs = (y.bitcast(fx.Uint32) & abs_mask).bitcast(fx.Float32) tile_max = y_abs.reduce(ReductionOp.MAX) - thread_row_max = fx.max(thread_row_max, tile_max) + thread_row_max = fx.max(thread_row_max, is_valid.select(tile_max, c_zero_f)) row_max = block_reduce_max(thread_row_max) scale = row_max / c_dtype_max @@ -1284,14 +1294,16 @@ def block_reduce_max(val): inv_scale = c_one_f / final_scale # Pass 3: quantize + store using per-row scale - for tile_i in range_constexpr(num_tiles): - q = y_local[tile_i] * inv_scale - q_i8 = q.to(quant_dtype) - q_lo = q_i8.shuffle(q_i8, [0, 1, 2, 3]) - q_hi = q_i8.shuffle(q_i8, [4, 5, 6, 7]) - out_idx = tid * 2 + tile_i * BLOCK_THREADS * 2 - _store_vec(copy_atom_q, quant_half_width, quant_dtype, q_lo, out_div_q, out_idx) - _store_vec(copy_atom_q, quant_half_width, quant_dtype, q_hi, out_div_q, out_idx + 1) + for tile_i in range_constexpr(NUM_VEC_ITERS): + idx = tid + tile_i * BLOCK_THREADS + if idx < VEC_TILES: + q = y_local[tile_i] * inv_scale + q_i8 = q.to(quant_dtype) + q_lo = q_i8.shuffle(q_i8, [0, 1, 2, 3]) + q_hi = q_i8.shuffle(q_i8, [4, 5, 6, 7]) + out_idx = idx * 2 + _store_vec(copy_atom_q, quant_half_width, quant_dtype, q_lo, out_div_q, out_idx) + _store_vec(copy_atom_q, quant_half_width, quant_dtype, q_hi, out_div_q, out_idx + 1) else: # Scalar fallback for arbitrary N. @@ -1469,10 +1481,12 @@ def _build_fused_add_rmsnorm_quant_module( arch = get_rocm_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") - 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 quant_dtype_max = _quant_dtype_max(quant_dtype_str) + 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 SharedStorage = _make_reduction_storage(RED_SLOTS) @@ -1589,9 +1603,8 @@ def block_reduce_max(val): return fx.memref_load(s_red, 0) - # Fast path for complete 128-bit tiles. - if const_expr(N >= tile_cols and N % tile_cols == 0 and elem_bits <= 16): - num_tiles = N // tile_cols + # Vector path for every complete f16/bf16 vec8 row. + if const_expr(USE_VEC_N): quant_half_width = VEC_WIDTH // 2 abs_mask = fx.Vector.filled(VEC_WIDTH, 0x7FFFFFFF, fx.Uint32) Input_buf = fx.rocdl.make_buffer_tensor(Input) @@ -1625,17 +1638,20 @@ def block_reduce_max(val): add_local = [] # Pass 1: add + cache + sumsq (also write residual_out) - for tile_i in range_constexpr(num_tiles): + for tile_i in range_constexpr(NUM_VEC_ITERS): idx = tid + tile_i * BLOCK_THREADS - x = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx).to(fx.Float32) - residual = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, residual_in_div, idx).to(fx.Float32) + is_valid = idx < VEC_TILES + idx_safe = is_valid.select(idx, 0) + x = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx_safe).to(fx.Float32) + residual = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, residual_in_div, idx_safe).to(fx.Float32) added_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, x + residual) add_local.append(added_e) added = added_e if dtype_str == "f32" else added_e.to(fx.Float32) added2 = added * added red2 = added2.reduce(ReductionOp.ADD, fastmath=fm_fast) - thread_sumsq = thread_sumsq + red2 - _store_vec(copy_atom, VEC_WIDTH, elem_dtype, added_e, residual_out_div, idx) + thread_sumsq = thread_sumsq + is_valid.select(red2, c_zero_f) + if idx < VEC_TILES: + _store_vec(copy_atom, VEC_WIDTH, elem_dtype, added_e, residual_out_div, idx) _, sum_sq = block_reduce_add2(thread_dummy, thread_sumsq) mean_sq = sum_sq / n_float @@ -1646,19 +1662,21 @@ def block_reduce_max(val): y_local = [] # Pass 2: normalize + gamma (+ optional smooth scale), cache output, and accumulate row max - 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) + is_valid = idx < VEC_TILES + idx_safe = is_valid.select(idx, 0) + g = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, idx_safe).to(fx.Float32) added = add_local[tile_i] if dtype_str == "f32" else add_local[tile_i].to(fx.Float32) y = (added * rrms) * g if const_expr(is_smooth): - s = _load_vec(copy_atom_xs, VEC_WIDTH, elem_dtype, xscale_div, idx).to(fx.Float32) + s = _load_vec(copy_atom_xs, VEC_WIDTH, elem_dtype, xscale_div, idx_safe).to(fx.Float32) y = y * s y_local.append(y) y_abs = (y.bitcast(fx.Uint32) & abs_mask).bitcast(fx.Float32) tile_max = y_abs.reduce(ReductionOp.MAX) - thread_row_max = fx.max(thread_row_max, tile_max) + thread_row_max = fx.max(thread_row_max, is_valid.select(tile_max, c_zero_f)) row_max = block_reduce_max(thread_row_max) scale = row_max / c_dtype_max @@ -1670,14 +1688,16 @@ def block_reduce_max(val): inv_scale = c_one_f / final_scale # Pass 3: quantize + store using per-row scale - for tile_i in range_constexpr(num_tiles): - q = y_local[tile_i] * inv_scale - q_i8 = q.to(quant_dtype) - q_lo = q_i8.shuffle(q_i8, [0, 1, 2, 3]) - q_hi = q_i8.shuffle(q_i8, [4, 5, 6, 7]) - out_idx = tid * 2 + tile_i * BLOCK_THREADS * 2 - _store_vec(copy_atom_q, quant_half_width, quant_dtype, q_lo, out_div_q, out_idx) - _store_vec(copy_atom_q, quant_half_width, quant_dtype, q_hi, out_div_q, out_idx + 1) + for tile_i in range_constexpr(NUM_VEC_ITERS): + idx = tid + tile_i * BLOCK_THREADS + if idx < VEC_TILES: + q = y_local[tile_i] * inv_scale + q_i8 = q.to(quant_dtype) + q_lo = q_i8.shuffle(q_i8, [0, 1, 2, 3]) + q_hi = q_i8.shuffle(q_i8, [4, 5, 6, 7]) + out_idx = idx * 2 + _store_vec(copy_atom_q, quant_half_width, quant_dtype, q_lo, out_div_q, out_idx) + _store_vec(copy_atom_q, quant_half_width, quant_dtype, q_hi, out_div_q, out_idx + 1) else: # Scalar fallback for arbitrary N. diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index 6f8f7a5b7e..c9b19ca271 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -133,6 +133,8 @@ def _get_rmsnorm_shape_override(env_name="ROCDSL_RMSNORM_SHAPES"): (64, 2001, "f32"), # small-N scalar path with an unaligned tail (16, 512, "bf16"), # small-N one-row vec8 path (64, 8192, "bf16"), # large-N aligned vec8 path + (32, 5120, "bf16"), # vec8 path with a partially active final iteration + (32, 7168, "bf16"), # same, at the DeepSeek-V3 hidden size ) _RMSNORM_FORWARD_EXTRA_CONFIGS = ( From 0ad9aff31f4bb67b08f63c309c1d7f9faad70dc2 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Fri, 18 Sep 2026 07:47:46 +0800 Subject: [PATCH 21/25] fix(norm): assert the implicit quantized-norm dtype ABI The quantized RMSNorm and LayerNorm launchers specialize every non-quantized operand (Gamma, Beta, the fused-add residuals, and SmoothQuant XScale) on the activation dtype, but that ABI was implicit and unchecked: an operand passed with any other dtype was reinterpreted bit-for-bit and silently produced corrupted scales instead of failing, and callers written against the older fp32-only SmoothQuant XScale ABI broke the same silent way. Add a validate_norm_operand_dtypes helper, assert the ABI in all eight quantized launchers so a mismatch raises ValueError at compile time, and document the operand dtype requirement in the RMSNorm section of docs/prebuilt_kernels_guide.md. Co-authored-by: Cursor --- docs/prebuilt_kernels_guide.md | 8 ++++++++ kernels/norm/layernorm_kernel.py | 7 +++++++ kernels/norm/rmsnorm_common.py | 15 +++++++++++++++ kernels/norm/rmsnorm_kernel.py | 7 +++++++ 4 files changed, 37 insertions(+) diff --git a/docs/prebuilt_kernels_guide.md b/docs/prebuilt_kernels_guide.md index a515de2e48..6bd1ec4d6f 100644 --- a/docs/prebuilt_kernels_guide.md +++ b/docs/prebuilt_kernels_guide.md @@ -75,6 +75,14 @@ per-row reciprocal std (`rstd`) for use by the backward pass. `weight_dtype_str` defaults to `dtype_str`; FP16/BF16 activations additionally support FP32 weights. +**Quantized variants:** The DynamicQuant and SmoothQuant builders emit int8 +`Output` and fp32 per-row `YScale`. Every other operand — `Gamma`, the +fused-add `ResidualIn`/`ResidualOut`, and SmoothQuant `XScale` — must use the +same element dtype as `Input`. The launchers raise `ValueError` on a mismatch, +so a wrong dtype fails at compile time rather than silently producing corrupted +scales. Unlike the plain forward, the quantized builders do not accept FP32 +weights with FP16/BF16 activations. + **Backward:** `build_rmsnorm_bwd_module(N, dtype_str, weight_dtype_str=None)` builds the fused RMSNorm backward kernel (grid `(M,)`, one block per row). Kernel signature diff --git a/kernels/norm/layernorm_kernel.py b/kernels/norm/layernorm_kernel.py index 37239047aa..701973df41 100644 --- a/kernels/norm/layernorm_kernel.py +++ b/kernels/norm/layernorm_kernel.py @@ -20,6 +20,7 @@ from flydsl.expr.typing import ReductionOp, full from flydsl.runtime.device import get_rocm_arch from kernels.common.kernels_common import atomic_add, dtype_to_elem_type, get_warp_size +from kernels.norm.rmsnorm_common import validate_norm_operand_dtypes try: import torch @@ -1104,6 +1105,7 @@ def launch_layernorm_smoothquant( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): + validate_norm_operand_dtypes(Input, Gamma=Gamma, Beta=Beta, XScale=XScale) launcher = layernorm_quant_kernel(Input, Gamma, Beta, XScale, YScale, Output) launcher.launch( grid=(m_in, 1, 1), @@ -1125,6 +1127,7 @@ def launch_layernorm_dynamicquant( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): + validate_norm_operand_dtypes(Input, Gamma=Gamma, Beta=Beta) launcher = layernorm_quant_kernel(Input, Gamma, Beta, Gamma, YScale, Output) launcher.launch( grid=(m_in, 1, 1), @@ -1492,6 +1495,9 @@ def launch_fused_add_layernorm_smoothquant( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): + validate_norm_operand_dtypes( + Input, ResidualIn=ResidualIn, Gamma=Gamma, Beta=Beta, XScale=XScale, ResidualOut=ResidualOut + ) launcher = fused_add_layernorm_quant_kernel( Input, ResidualIn, Gamma, Beta, XScale, YScale, Output, ResidualOut ) @@ -1517,6 +1523,7 @@ def launch_fused_add_layernorm_dynamicquant( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): + validate_norm_operand_dtypes(Input, ResidualIn=ResidualIn, Gamma=Gamma, Beta=Beta, ResidualOut=ResidualOut) launcher = fused_add_layernorm_quant_kernel( Input, ResidualIn, Gamma, Beta, Gamma, YScale, Output, ResidualOut ) diff --git a/kernels/norm/rmsnorm_common.py b/kernels/norm/rmsnorm_common.py index 72fd273439..ca3fd3917a 100644 --- a/kernels/norm/rmsnorm_common.py +++ b/kernels/norm/rmsnorm_common.py @@ -52,6 +52,21 @@ class SharedStorage: return SharedStorage +def validate_norm_operand_dtypes(input_tensor, **operands) -> None: + """Validate quantized-norm operands against the input dtype. + + Quantized norm kernels specialize every non-quantized operand on the activation + dtype, so a mismatched tensor is reinterpreted bit-for-bit instead of failing. + """ + mismatched = {name: t.dtype for name, t in operands.items() if t.dtype != input_tensor.dtype} + if mismatched: + got = ", ".join(f"{name}={dtype}" for name, dtype in mismatched.items()) + raise ValueError( + f"quantized norm kernels require {'/'.join(operands)} to match the input dtype " + f"{input_tensor.dtype}, got {got}" + ) + + def load_scalar(copy_atom, elem_dtype, divided_tensor, index): view = fx.slice(divided_tensor, (None, index)) r = fx.make_rmem_tensor(1, elem_dtype) diff --git a/kernels/norm/rmsnorm_kernel.py b/kernels/norm/rmsnorm_kernel.py index 2e59ba81d4..9d67cabaeb 100644 --- a/kernels/norm/rmsnorm_kernel.py +++ b/kernels/norm/rmsnorm_kernel.py @@ -45,6 +45,7 @@ from kernels.norm.rmsnorm_common import store_vec as _store_vec from kernels.norm.rmsnorm_common import to_elem_scalar as _to_elem_scalar from kernels.norm.rmsnorm_common import to_elem_vec as _to_elem_vec +from kernels.norm.rmsnorm_common import validate_norm_operand_dtypes as _validate_norm_operand_dtypes from kernels.norm.rmsnorm_common import weight_vec_width as _weight_vec_width try: @@ -1410,6 +1411,7 @@ def launch_rmsnorm_smoothquant( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): + _validate_norm_operand_dtypes(Input, Gamma=Gamma, XScale=XScale) launcher = rmsnorm_quant_kernel(Input, Gamma, XScale, YScale, Output) launcher.launch( grid=(m_in, 1, 1), @@ -1430,6 +1432,7 @@ def launch_rmsnorm_dynamicquant( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): + _validate_norm_operand_dtypes(Input, Gamma=Gamma) launcher = rmsnorm_quant_kernel(Input, Gamma, Gamma, YScale, Output) launcher.launch( grid=(m_in, 1, 1), @@ -1819,6 +1822,9 @@ def launch_fused_add_rmsnorm_smoothquant( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): + _validate_norm_operand_dtypes( + Input, ResidualIn=ResidualIn, Gamma=Gamma, XScale=XScale, ResidualOut=ResidualOut + ) launcher = fused_add_rmsnorm_quant_kernel(Input, ResidualIn, Gamma, XScale, YScale, Output, ResidualOut) launcher.launch( grid=(m_in, 1, 1), @@ -1841,6 +1847,7 @@ def launch_fused_add_rmsnorm_dynamicquant( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): + _validate_norm_operand_dtypes(Input, ResidualIn=ResidualIn, Gamma=Gamma, ResidualOut=ResidualOut) launcher = fused_add_rmsnorm_quant_kernel(Input, ResidualIn, Gamma, Gamma, YScale, Output, ResidualOut) launcher.launch( grid=(m_in, 1, 1), From 9f2357fc26f04d450ca8b3710c6f7cc8a3acded2 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Fri, 18 Sep 2026 16:08:43 +0800 Subject: [PATCH 22/25] refactor(norm): load and store single elements by direct indexing The scalar paths moved one element through a copy atom, a (1, 1) logical divide and a one-element register-memory tensor, which is boilerplate for what buffer-tensor indexing already expresses. Replace the load_scalar/store_scalar helpers in rmsnorm_common.py and the _load_scalar/_store_scalar/_store_yscale helpers in layernorm_kernel.py with direct t[index] reads and writes at every call site, and drop the scalar copy atoms, (1, 1) divides and element-width locals that only fed them; this also retires the elem_dtype/store_dtype split _store_scalar needed. rmsnorm_bwd_kernel.py is touched because it imported the deleted helpers. The generated ISA is byte-identical to before the change (502 instructions, 36 VGPRs, no spills on the N=3073 and N=3001 bf16 configurations), so this is a source-level cleanup with no codegen impact. Co-authored-by: Cursor --- kernels/norm/layernorm_kernel.py | 197 +++++++-------------------- kernels/norm/rmsnorm_bwd_kernel.py | 122 ++++------------- kernels/norm/rmsnorm_common.py | 15 --- kernels/norm/rmsnorm_kernel.py | 205 +++++++---------------------- 4 files changed, 130 insertions(+), 409 deletions(-) diff --git a/kernels/norm/layernorm_kernel.py b/kernels/norm/layernorm_kernel.py index 701973df41..7efc3d6df1 100644 --- a/kernels/norm/layernorm_kernel.py +++ b/kernels/norm/layernorm_kernel.py @@ -46,21 +46,6 @@ class SharedStorage: return SharedStorage -def _load_scalar(copy_atom, elem_dtype, divided_tensor, index): - view = fx.slice(divided_tensor, (None, index)) - r = fx.make_rmem_tensor(1, elem_dtype) - fx.copy(copy_atom, view, r) - return fx.memref_load_vec(r)[0] - - -def _store_scalar(copy_atom, elem_dtype, store_dtype, divided_tensor, index, val): - r = fx.make_rmem_tensor(1, elem_dtype) - ts = full(1, store_dtype(val), store_dtype) - fx.memref_store_vec(ts, r) - view = fx.slice(divided_tensor, (None, index)) - fx.copy(copy_atom, r, view) - - def _load_vec(copy_atom, vec_width, elem_dtype, div_tensor, idx): r = fx.make_rmem_tensor(vec_width, elem_dtype) fx.copy(copy_atom, fx.slice(div_tensor, (None, idx)), r) @@ -99,13 +84,6 @@ def _to_elem_vec(dtype_str: str, elem_dtype, use_hw_cvt_bf16: bool, y): return y.to(elem_dtype) -def _store_yscale(scale_copy_atom, yscale_div, index, val): - r = fx.make_rmem_tensor(1, fx.Float32) - ts = full(1, fx.Float32(val), fx.Float32) - fx.memref_store_vec(ts, r) - fx.copy(scale_copy_atom, r, fx.slice(yscale_div, (None, index))) - - def _quant_dtype_to_elem_type(dtype_str: str): if dtype_str in ("i8", "int8"): return fx.Int8 @@ -167,9 +145,6 @@ def layernorm_kernel( 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): @@ -261,8 +236,8 @@ def 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) + mean_row[0] = mean + rstd_row[0] = rstd # ── Pass 2: normalize + affine + store ─────────────────────── for tile_i in range_constexpr(NUM_VEC_ITERS): @@ -284,22 +259,12 @@ def compute_mean_rstd(sum_val, sumsq_val): thread_sum = c_zero_f thread_sumsq = c_zero_f - 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(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(Output_buf, fx.make_layout(1, 1)) - # ── Pass 1: sum + sumsq ────────────────────────────────────── for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): idx = tid + 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_e = Input_buf[idx_safe] x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) x2 = x * x x_safe = is_valid.select(x, c_zero_f) @@ -312,16 +277,16 @@ def 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) + mean_row[0] = mean + rstd_row[0] = rstd # ── Pass 2: normalize + affine + store ─────────────────────── for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): idx = tid + 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) - b_e = _load_scalar(copy_atom_s, elem_dtype, beta_div, idx) + x_e = Input_buf[idx] + g_e = Gamma_buf[idx] + b_e = Beta_buf[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) b = b_e if dtype_str == "f32" else b_e.to(fx.Float32) @@ -330,7 +295,7 @@ def compute_mean_rstd(sum_val, sumsq_val): scaled = norm * g y = scaled + b y_e = _to_elem_scalar(dtype_str, elem_dtype, y) - _store_scalar(copy_atom_s, elem_dtype, elem_dtype, out_div, idx, y_e) + Output_buf[idx] = y_e # ── JIT host launcher ───────────────────────────────────────────────── if store_stats: @@ -390,7 +355,6 @@ def build_layernorm_bwd_module(N: int, dtype_str: str): 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 @@ -470,21 +434,8 @@ def block_reduce_add2(val0, val1): 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) + mean = mean_row[0] + rstd = rstd_row[0] dgamma_view = DGamma.view(fx.make_layout(N, 1)) dbias_view = DBias.view(fx.make_layout(N, 1)) @@ -495,9 +446,9 @@ def block_reduce_add2(val0, val1): 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_e = row_in[idx_safe] + dy_e = row_dy[idx_safe] + g_e = gamma[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) @@ -514,9 +465,9 @@ def block_reduce_add2(val0, val1): 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_e = row_in[idx] + dy_e = row_dy[idx] + g_e = gamma[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) @@ -524,7 +475,7 @@ def block_reduce_add2(val0, val1): 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) + row_dx[idx] = dx_e dgamma = dy * x_hat atomic_add(dgamma_view, idx, dgamma, dtype_bytes=4) @@ -707,18 +658,6 @@ def compute_mean_rstd(sum_val, sumsq_val): row_out = fx.slice(Output_buf, (bid, None)) row_residual_out = fx.slice(ResidualOut_buf, (bid, None)) - copy_atom_s = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - elem_bits, - ) - - in_div = fx.logical_divide(row_in, fx.make_layout(1, 1)) - residual_in_div = fx.logical_divide(row_residual_in, 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)) - residual_out_div = fx.logical_divide(row_residual_out, fx.make_layout(1, 1)) - c_zero_f = fx.Float32(0.0) thread_sum = c_zero_f thread_sumsq = c_zero_f @@ -727,8 +666,8 @@ def compute_mean_rstd(sum_val, sumsq_val): idx = tid + base_idx_int is_valid = idx < N idx_safe = is_valid.select(idx, 0) - x_e = _load_scalar(copy_atom_s, elem_dtype, in_div, idx_safe) - r_e = _load_scalar(copy_atom_s, elem_dtype, residual_in_div, idx_safe) + x_e = row_in[idx_safe] + r_e = row_residual_in[idx_safe] x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) residual = r_e if dtype_str == "f32" else r_e.to(fx.Float32) added_e = _to_elem_scalar(dtype_str, elem_dtype, x + residual) @@ -737,7 +676,7 @@ def compute_mean_rstd(sum_val, sumsq_val): thread_sum = thread_sum + added_safe thread_sumsq = thread_sumsq + is_valid.select(added * added, c_zero_f) if idx < N: - _store_scalar(copy_atom_s, elem_dtype, elem_dtype, residual_out_div, idx, added_e) + row_residual_out[idx] = added_e sum_val, sumsq_val = block_reduce_add2(thread_sum, thread_sumsq) mean, rstd = compute_mean_rstd(sum_val, sumsq_val) @@ -745,16 +684,16 @@ def compute_mean_rstd(sum_val, sumsq_val): for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): idx = tid + base_idx_int if idx < N: - added_e = _load_scalar(copy_atom_s, elem_dtype, residual_out_div, idx) - g_e = _load_scalar(copy_atom_s, elem_dtype, gamma_div, idx) - b_e = _load_scalar(copy_atom_s, elem_dtype, beta_div, idx) + added_e = row_residual_out[idx] + g_e = Gamma_buf[idx] + b_e = Beta_buf[idx] added = added_e if dtype_str == "f32" else added_e.to(fx.Float32) g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) b = b_e if dtype_str == "f32" else b_e.to(fx.Float32) y = (added - mean) * rstd y = y * g + b y_e = _to_elem_scalar(dtype_str, elem_dtype, y) - _store_scalar(copy_atom_s, elem_dtype, elem_dtype, out_div, idx, y_e) + row_out[idx] = y_e @flyc.jit def launch_fused_add_layernorm( @@ -820,10 +759,7 @@ def layernorm_quant_kernel( 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)) - YScale_buf = fx.rocdl.make_buffer_tensor(YScale) - yscale_div = fx.logical_divide(YScale_buf, fx.make_layout(1, 1)) - scale_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 32) def wave_reduce_add(x): w = x @@ -977,7 +913,7 @@ def block_reduce_max(val): final_scale = (scale == c_zero_f).select(c_one_f, scale) if tid == 0: - _store_yscale(scale_copy_atom, yscale_div, bid, final_scale) + YScale_buf[bid] = final_scale inv_scale = c_one_f / final_scale @@ -1007,19 +943,6 @@ def block_reduce_max(val): row_in = fx.slice(Input_buf, (bid, None)) row_out = fx.slice(Output_buf, (bid, None)) - copy_atom_s = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - elem_bits, - ) - copy_atom_qs = fx.make_copy_atom(fx.rocdl.BufferCopy(8), 8) - - in_div = fx.logical_divide(row_in, 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)) - if const_expr(is_smooth): - xscale_div = fx.logical_divide(XScale_buf, fx.make_layout(1, 1)) - def _abs_scalar(val): is_neg = val < c_zero_f neg_val = c_zero_f - val @@ -1032,7 +955,7 @@ def _abs_scalar(val): idx = tid + base_idx_int is_valid = idx < N idx_safe = is_valid.select(idx, 0) - x_e = _load_scalar(copy_atom_s, elem_dtype, in_div, idx_safe) + x_e = row_in[idx_safe] x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) x2 = x * x thread_sum = thread_sum + is_valid.select(x, c_zero_f) @@ -1049,16 +972,16 @@ def _abs_scalar(val): idx = tid + base_idx_int is_valid = idx < N idx_safe = is_valid.select(idx, 0) - x_e = _load_scalar(copy_atom_s, elem_dtype, in_div, idx_safe) - g_e = _load_scalar(copy_atom_s, elem_dtype, gamma_div, idx_safe) - b_e = _load_scalar(copy_atom_s, elem_dtype, beta_div, idx_safe) + x_e = row_in[idx_safe] + g_e = Gamma_buf[idx_safe] + b_e = Beta_buf[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) b = b_e if dtype_str == "f32" else b_e.to(fx.Float32) y = (x - mean) * rstd y = y * g + b if const_expr(is_smooth): - s_e = _load_scalar(copy_atom_s, elem_dtype, xscale_div, idx_safe) + s_e = XScale_buf[idx_safe] s = s_e if dtype_str == "f32" else s_e.to(fx.Float32) y = y * s y_abs = _abs_scalar(y) @@ -1069,28 +992,28 @@ def _abs_scalar(val): final_scale = (scale == c_zero_f).select(c_one_f, scale) if tid == 0: - _store_yscale(scale_copy_atom, yscale_div, bid, final_scale) + YScale_buf[bid] = final_scale inv_scale = c_one_f / final_scale for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): idx = tid + base_idx_int if idx < N: - x_e = _load_scalar(copy_atom_s, elem_dtype, in_div, idx) - g_e = _load_scalar(copy_atom_s, elem_dtype, gamma_div, idx) - b_e = _load_scalar(copy_atom_s, elem_dtype, beta_div, idx) + x_e = row_in[idx] + g_e = Gamma_buf[idx] + b_e = Beta_buf[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) b = b_e if dtype_str == "f32" else b_e.to(fx.Float32) y = (x - mean) * rstd y = y * g + b if const_expr(is_smooth): - s_e = _load_scalar(copy_atom_s, elem_dtype, xscale_div, idx) + s_e = XScale_buf[idx] s = s_e if dtype_str == "f32" else s_e.to(fx.Float32) y = y * s q = y * inv_scale q_i8 = q.to(quant_dtype) - _store_scalar(copy_atom_qs, quant_dtype, quant_dtype, out_div, idx, q_i8) + row_out[idx] = q_i8 if is_smooth: @@ -1188,8 +1111,6 @@ def fused_add_layernorm_quant_kernel( s_sumsq = lds.s_sumsq.view(fx.make_layout(RED_SLOTS, 1)) YScale_buf = fx.rocdl.make_buffer_tensor(YScale) - yscale_div = fx.logical_divide(YScale_buf, fx.make_layout(1, 1)) - scale_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 32) def wave_reduce_add(x): w = x @@ -1353,7 +1274,7 @@ def block_reduce_max(val): final_scale = (scale == c_zero_f).select(c_one_f, scale) if tid == 0: - _store_yscale(scale_copy_atom, yscale_div, bid, final_scale) + YScale_buf[bid] = final_scale inv_scale = c_one_f / final_scale @@ -1387,21 +1308,6 @@ def block_reduce_max(val): row_out = fx.slice(Output_buf, (bid, None)) row_residual_out = fx.slice(ResidualOut_buf, (bid, None)) - copy_atom_s = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - elem_bits, - ) - copy_atom_qs = fx.make_copy_atom(fx.rocdl.BufferCopy(8), 8) - - in_div = fx.logical_divide(row_in, fx.make_layout(1, 1)) - residual_in_div = fx.logical_divide(row_residual_in, 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)) - residual_out_div = fx.logical_divide(row_residual_out, fx.make_layout(1, 1)) - if const_expr(is_smooth): - xscale_div = fx.logical_divide(XScale_buf, fx.make_layout(1, 1)) - def _abs_scalar(val): is_neg = val < c_zero_f neg_val = c_zero_f - val @@ -1414,13 +1320,13 @@ def _abs_scalar(val): idx = tid + base_idx_int is_valid = idx < N idx_safe = is_valid.select(idx, 0) - x_e = _load_scalar(copy_atom_s, elem_dtype, in_div, idx_safe) - r_e = _load_scalar(copy_atom_s, elem_dtype, residual_in_div, idx_safe) + x_e = row_in[idx_safe] + r_e = row_residual_in[idx_safe] x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) residual = r_e if dtype_str == "f32" else r_e.to(fx.Float32) added_e = _to_elem_scalar(dtype_str, elem_dtype, x + residual) if idx < N: - _store_scalar(copy_atom_s, elem_dtype, elem_dtype, residual_out_div, idx, added_e) + row_residual_out[idx] = added_e x = added_e if dtype_str == "f32" else added_e.to(fx.Float32) x2 = x * x thread_sum = thread_sum + is_valid.select(x, c_zero_f) @@ -1431,22 +1337,21 @@ def _abs_scalar(val): var = sumsq_val / n_float - mean * mean var = (var < c_zero_f).select(c_zero_f, var) rstd = fmath.rsqrt(var + eps_c, fastmath=fm_fast) - thread_row_max = c_zero_f for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): idx = tid + base_idx_int is_valid = idx < N idx_safe = is_valid.select(idx, 0) - x_e = _load_scalar(copy_atom_s, elem_dtype, residual_out_div, idx_safe) - g_e = _load_scalar(copy_atom_s, elem_dtype, gamma_div, idx_safe) - b_e = _load_scalar(copy_atom_s, elem_dtype, beta_div, idx_safe) + x_e = row_residual_out[idx_safe] + g_e = Gamma_buf[idx_safe] + b_e = Beta_buf[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) b = b_e if dtype_str == "f32" else b_e.to(fx.Float32) y = (x - mean) * rstd y = y * g + b if const_expr(is_smooth): - s_e = _load_scalar(copy_atom_s, elem_dtype, xscale_div, idx_safe) + s_e = XScale_buf[idx_safe] s = s_e if dtype_str == "f32" else s_e.to(fx.Float32) y = y * s y_abs = _abs_scalar(y) @@ -1457,28 +1362,28 @@ def _abs_scalar(val): final_scale = (scale == c_zero_f).select(c_one_f, scale) if tid == 0: - _store_yscale(scale_copy_atom, yscale_div, bid, final_scale) + YScale_buf[bid] = final_scale inv_scale = c_one_f / final_scale for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): idx = tid + base_idx_int if idx < N: - x_e = _load_scalar(copy_atom_s, elem_dtype, residual_out_div, idx) - g_e = _load_scalar(copy_atom_s, elem_dtype, gamma_div, idx) - b_e = _load_scalar(copy_atom_s, elem_dtype, beta_div, idx) + x_e = row_residual_out[idx] + g_e = Gamma_buf[idx] + b_e = Beta_buf[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) b = b_e if dtype_str == "f32" else b_e.to(fx.Float32) y = (x - mean) * rstd y = y * g + b if const_expr(is_smooth): - s_e = _load_scalar(copy_atom_s, elem_dtype, xscale_div, idx) + s_e = XScale_buf[idx] s = s_e if dtype_str == "f32" else s_e.to(fx.Float32) y = y * s q = y * inv_scale q_i8 = q.to(quant_dtype) - _store_scalar(copy_atom_qs, quant_dtype, quant_dtype, out_div, idx, q_i8) + row_out[idx] = q_i8 if is_smooth: diff --git a/kernels/norm/rmsnorm_bwd_kernel.py b/kernels/norm/rmsnorm_bwd_kernel.py index 8874059063..5a5e8966f1 100644 --- a/kernels/norm/rmsnorm_bwd_kernel.py +++ b/kernels/norm/rmsnorm_bwd_kernel.py @@ -15,12 +15,10 @@ BLOCK_THREADS, VEC_WIDTH, WARP_SIZE, - load_scalar, load_vec, load_weight_vec, make_single_reduction_storage, resolve_rmsnorm_weight_dtype, - store_scalar, store_vec, to_elem_vec, weight_vec_width, @@ -51,8 +49,6 @@ def build_rmsnorm_bwd_module(N: int, dtype_str: str, weight_dtype_str: str | Non """ weight_dtype_str = resolve_rmsnorm_weight_dtype(dtype_str, weight_dtype_str) RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) - elem_bits = 32 if dtype_str == "f32" else 16 - weight_elem_bits = 32 if weight_dtype_str == "f32" else 16 SharedStorage = make_single_reduction_storage(RED_SLOTS) @flyc.kernel @@ -68,7 +64,6 @@ def rmsnorm_bwd_kernel( tid = fx.thread_idx.x elem_dtype = dtype_to_elem_type(dtype_str) - weight_elem_dtype = dtype_to_elem_type(weight_dtype_str) fm_fast = arith.FastMathFlags.fast n_float = float(N) c_zero_f = fx.Float32(0.0) @@ -114,23 +109,7 @@ def block_reduce_add(val): row_dy = fx.slice(DY_buf, (bid, None)) row_dx = fx.slice(DX_buf, (bid, None)) - copy_atom_s = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - elem_bits, - ) - gamma_copy_atom_s = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if weight_elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - weight_elem_bits, - ) - copy_atom_f32 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 32) - - row_div = fx.logical_divide(row_in, fx.make_layout(1, 1)) - dy_div = fx.logical_divide(row_dy, fx.make_layout(1, 1)) - gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(1, 1)) - dx_div = fx.logical_divide(row_dx, fx.make_layout(1, 1)) - rstd_div = fx.logical_divide(Rstd_buf, fx.make_layout(1, 1)) - - rstd = load_scalar(copy_atom_f32, fx.Float32, rstd_div, bid) + rstd = Rstd_buf[bid] # Pass 1: c1 = mean( x_hat * wdy ) = mean( (x*rstd) * (dy*gamma) ) thread_acc = c_zero_f @@ -138,9 +117,9 @@ def block_reduce_add(val): 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(gamma_copy_atom_s, weight_elem_dtype, gamma_div, idx_safe) + x_e = row_in[idx_safe] + dy_e = row_dy[idx_safe] + g_e = Gamma_buf[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 weight_dtype_str == "f32" else g_e.to(fx.Float32) @@ -156,9 +135,9 @@ def block_reduce_add(val): 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(gamma_copy_atom_s, weight_elem_dtype, gamma_div, idx) + x_e = row_in[idx] + dy_e = row_dy[idx] + g_e = Gamma_buf[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 weight_dtype_str == "f32" else g_e.to(fx.Float32) @@ -166,7 +145,7 @@ def block_reduce_add(val): wdy = dy * g dx = (wdy - x_hat * c1) * rstd dx_e = dx if dtype_str == "f32" else dx.to(elem_dtype) - store_scalar(copy_atom_s, elem_dtype, dx_div, idx, dx_e) + row_dx[idx] = dx_e dw = dy * x_hat atomic_add(DWeight, idx, dw, dtype_bytes=4) @@ -209,8 +188,6 @@ def build_fused_add_rmsnorm_bwd_module(N: int, dtype_str: str, weight_dtype_str: """ weight_dtype_str = resolve_rmsnorm_weight_dtype(dtype_str, weight_dtype_str) RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) - elem_bits = 32 if dtype_str == "f32" else 16 - weight_elem_bits = 32 if weight_dtype_str == "f32" else 16 SharedStorage = make_single_reduction_storage(RED_SLOTS) @flyc.kernel @@ -227,7 +204,6 @@ def fused_add_rmsnorm_bwd_kernel( tid = fx.thread_idx.x elem_dtype = dtype_to_elem_type(dtype_str) - weight_elem_dtype = dtype_to_elem_type(weight_dtype_str) fm_fast = arith.FastMathFlags.fast n_float = float(N) c_zero_f = fx.Float32(0.0) @@ -275,24 +251,7 @@ def block_reduce_add(val): row_dres_out = fx.slice(DResidualOut_buf, (bid, None)) row_dx = fx.slice(DX_buf, (bid, None)) - copy_atom_s = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - elem_bits, - ) - gamma_copy_atom_s = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if weight_elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - weight_elem_bits, - ) - copy_atom_f32 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 32) - - added_div = fx.logical_divide(row_added, fx.make_layout(1, 1)) - dy_div = fx.logical_divide(row_dy, fx.make_layout(1, 1)) - dres_out_div = fx.logical_divide(row_dres_out, fx.make_layout(1, 1)) - gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(1, 1)) - dx_div = fx.logical_divide(row_dx, fx.make_layout(1, 1)) - rstd_div = fx.logical_divide(Rstd_buf, fx.make_layout(1, 1)) - - rstd = load_scalar(copy_atom_f32, fx.Float32, rstd_div, bid) + rstd = Rstd_buf[bid] # Pass 1: c1 = mean( a_hat * wdy ) = mean( (added*rstd) * (dy*gamma) ) thread_acc = c_zero_f @@ -300,9 +259,9 @@ def block_reduce_add(val): idx = tid + base is_valid = idx < N idx_safe = is_valid.select(idx, 0) - a_e = load_scalar(copy_atom_s, elem_dtype, added_div, idx_safe) - dy_e = load_scalar(copy_atom_s, elem_dtype, dy_div, idx_safe) - g_e = load_scalar(gamma_copy_atom_s, weight_elem_dtype, gamma_div, idx_safe) + a_e = row_added[idx_safe] + dy_e = row_dy[idx_safe] + g_e = Gamma_buf[idx_safe] a = a_e if dtype_str == "f32" else a_e.to(fx.Float32) dy = dy_e if dtype_str == "f32" else dy_e.to(fx.Float32) g = g_e if weight_dtype_str == "f32" else g_e.to(fx.Float32) @@ -320,10 +279,10 @@ def block_reduce_add(val): for base in range_constexpr(0, N, BLOCK_THREADS): idx = tid + base if idx < N: - a_e = load_scalar(copy_atom_s, elem_dtype, added_div, idx) - dy_e = load_scalar(copy_atom_s, elem_dtype, dy_div, idx) - g_e = load_scalar(gamma_copy_atom_s, weight_elem_dtype, gamma_div, idx) - dres_out_e = load_scalar(copy_atom_s, elem_dtype, dres_out_div, idx) + a_e = row_added[idx] + dy_e = row_dy[idx] + g_e = Gamma_buf[idx] + dres_out_e = row_dres_out[idx] a = a_e if dtype_str == "f32" else a_e.to(fx.Float32) dy = dy_e if dtype_str == "f32" else dy_e.to(fx.Float32) g = g_e if weight_dtype_str == "f32" else g_e.to(fx.Float32) @@ -333,7 +292,7 @@ def block_reduce_add(val): d_added = (wdy - a_hat * c1) * rstd total = d_added + dres_out total_e = total if dtype_str == "f32" else total.to(elem_dtype) - store_scalar(copy_atom_s, elem_dtype, dx_div, idx, total_e) + row_dx[idx] = total_e dw = dy * a_hat atomic_add(DWeight, idx, dw, dtype_bytes=4) @@ -455,9 +414,6 @@ def block_reduce_add(val): DX_buf = fx.rocdl.make_buffer_tensor(DX) DWeightPartial_buf = fx.rocdl.make_buffer_tensor(DWeightPartial) - rstd_div = fx.logical_divide(Rstd_buf, fx.make_layout(1, 1)) - partial_div = fx.logical_divide(DWeightPartial_buf, fx.make_layout(1, 1)) - copy_atom_io = fx.make_copy_atom( ( fx.rocdl.BufferCopy128b() @@ -474,7 +430,6 @@ def block_reduce_add(val): ), weight_elem_bits, ) - copy_atom_f32 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 32) gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(WEIGHT_IO_WIDTH, 1)) gamma_local = [] @@ -507,7 +462,7 @@ def block_reduce_add(val): if const_expr(fused_add): dres_out_div = fx.logical_divide(row_dres_out, fx.make_layout(IO_WIDTH, 1)) - rstd = load_scalar(copy_atom_f32, fx.Float32, rstd_div, row) + rstd = Rstd_buf[row] thread_acc = c_zero_f source_local = [] dy_local = [] @@ -522,9 +477,9 @@ def block_reduce_add(val): dy_local.append(dy_e) gamma_e = gamma_local[tile_i] else: - source_e = load_scalar(copy_atom_io, elem_dtype, source_div, io_idx_safe) - dy_e = load_scalar(copy_atom_io, elem_dtype, dy_div, io_idx_safe) - gamma_e = load_scalar(gamma_copy_atom, weight_elem_dtype, gamma_div, io_idx_safe) + source_e = row_source[io_idx_safe] + dy_e = row_dy[io_idx_safe] + gamma_e = Gamma_buf[io_idx_safe] source = source_e if dtype_str == "f32" else source_e.to(fx.Float32) dy = dy_e if dtype_str == "f32" else dy_e.to(fx.Float32) @@ -549,9 +504,9 @@ def block_reduce_add(val): dy_e = dy_local[tile_i] gamma_e = gamma_local[tile_i] else: - source_e = load_scalar(copy_atom_io, elem_dtype, source_div, io_idx_safe) - dy_e = load_scalar(copy_atom_io, elem_dtype, dy_div, io_idx_safe) - gamma_e = load_scalar(gamma_copy_atom, weight_elem_dtype, gamma_div, io_idx_safe) + source_e = row_source[io_idx_safe] + dy_e = row_dy[io_idx_safe] + gamma_e = Gamma_buf[io_idx_safe] source = source_e if dtype_str == "f32" else source_e.to(fx.Float32) dy = dy_e if dtype_str == "f32" else dy_e.to(fx.Float32) @@ -564,7 +519,7 @@ def block_reduce_add(val): if const_expr(USE_VEC): dres = load_vec(copy_atom_io, IO_WIDTH, elem_dtype, dres_out_div, io_idx_safe).to(fx.Float32) else: - dres_e = load_scalar(copy_atom_io, elem_dtype, dres_out_div, io_idx_safe) + dres_e = row_dres_out[io_idx_safe] dres = dres_e if dtype_str == "f32" else dres_e.to(fx.Float32) dx = dx + dres @@ -574,7 +529,7 @@ def block_reduce_add(val): store_vec(copy_atom_io, IO_WIDTH, elem_dtype, dx_e, dx_div, io_idx) else: dx_e = dx if dtype_str == "f32" else dx.to(elem_dtype) - store_scalar(copy_atom_io, elem_dtype, dx_div, io_idx, dx_e) + row_dx[io_idx] = dx_e dw = dy * source_hat if const_expr(USE_VEC): @@ -593,13 +548,7 @@ def block_reduce_add(val): idx = io_idx * IO_WIDTH + lane partial_idx = bid * N + idx partial_value = dweight_partial[tile_i * IO_WIDTH + lane] - store_scalar( - copy_atom_f32, - fx.Float32, - partial_div, - partial_idx, - partial_value, - ) + DWeightPartial_buf[partial_idx] = partial_value @flyc.kernel def rmsnorm_bwd_dweight_reduce_kernel( @@ -617,13 +566,6 @@ def rmsnorm_bwd_dweight_reduce_kernel( weight_elem_dtype = dtype_to_elem_type(weight_dtype_str) DWeightPartial_buf = fx.rocdl.make_buffer_tensor(DWeightPartial) DWeight_buf = fx.rocdl.make_buffer_tensor(DWeight) - partial_div = fx.logical_divide(DWeightPartial_buf, fx.make_layout(1, 1)) - dweight_div = fx.logical_divide(DWeight_buf, fx.make_layout(1, 1)) - copy_atom_f32 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 32) - weight_copy_atom_s = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if weight_elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - weight_elem_bits, - ) lds = fx.SharedAllocator().allocate(DWeightReduceStorage).peek() s_partial = lds.s_red.view(fx.make_layout(DWEIGHT_REDUCE_THREADS, 1)) @@ -635,7 +577,7 @@ def rmsnorm_bwd_dweight_reduce_kernel( partial_valid = partial_row < num_programs partial_row_safe = partial_valid.select(partial_row, 0) partial_idx = partial_row_safe * N + col_safe - value = load_scalar(copy_atom_f32, fx.Float32, partial_div, partial_idx) + value = DWeightPartial_buf[partial_idx] acc = acc + partial_valid.select(value, c_zero_f) fx.memref_store(acc, s_partial, tid) gpu.barrier() @@ -646,13 +588,7 @@ def rmsnorm_bwd_dweight_reduce_kernel( for lane in range_constexpr(DWEIGHT_REDUCE_ROW_LANES): total = total + fx.memref_load(s_partial, lane * DWEIGHT_REDUCE_COLS + col_lane) out = total if weight_dtype_str == "f32" else total.to(weight_elem_dtype) - store_scalar( - weight_copy_atom_s, - weight_elem_dtype, - dweight_div, - col, - out, - ) + DWeight_buf[col] = out reduce_grid = (N + DWEIGHT_REDUCE_COLS - 1) // DWEIGHT_REDUCE_COLS diff --git a/kernels/norm/rmsnorm_common.py b/kernels/norm/rmsnorm_common.py index ca3fd3917a..834f5257fc 100644 --- a/kernels/norm/rmsnorm_common.py +++ b/kernels/norm/rmsnorm_common.py @@ -67,21 +67,6 @@ def validate_norm_operand_dtypes(input_tensor, **operands) -> None: ) -def load_scalar(copy_atom, elem_dtype, divided_tensor, index): - view = fx.slice(divided_tensor, (None, index)) - r = fx.make_rmem_tensor(1, elem_dtype) - fx.copy(copy_atom, view, r) - return fx.memref_load_vec(r)[0] - - -def store_scalar(copy_atom, elem_dtype, divided_tensor, index, val): - r = fx.make_rmem_tensor(1, elem_dtype) - ts = fx.Vector.filled(1, val, elem_dtype) - fx.memref_store_vec(ts, r) - view = fx.slice(divided_tensor, (None, index)) - fx.copy(copy_atom, r, view) - - def load_vec(copy_atom, vec_width, elem_dtype, div_tensor, idx): r = fx.make_rmem_tensor(vec_width, elem_dtype) fx.copy(copy_atom, fx.slice(div_tensor, (None, idx)), r) diff --git a/kernels/norm/rmsnorm_kernel.py b/kernels/norm/rmsnorm_kernel.py index 9d67cabaeb..d57908ae37 100644 --- a/kernels/norm/rmsnorm_kernel.py +++ b/kernels/norm/rmsnorm_kernel.py @@ -35,13 +35,11 @@ VEC_WIDTH, WARP_SIZE, ) -from kernels.norm.rmsnorm_common import load_scalar as _load_scalar from kernels.norm.rmsnorm_common import load_vec as _load_vec from kernels.norm.rmsnorm_common import load_weight_vec as _load_weight_vec from kernels.norm.rmsnorm_common import make_reduction_storage as _make_reduction_storage from kernels.norm.rmsnorm_common import make_single_reduction_storage as _make_single_reduction_storage from kernels.norm.rmsnorm_common import resolve_rmsnorm_weight_dtype as _resolve_rmsnorm_weight_dtype -from kernels.norm.rmsnorm_common import store_scalar as _store_scalar from kernels.norm.rmsnorm_common import store_vec as _store_vec from kernels.norm.rmsnorm_common import to_elem_scalar as _to_elem_scalar from kernels.norm.rmsnorm_common import to_elem_vec as _to_elem_vec @@ -312,8 +310,6 @@ def rmsnorm_kernel( if const_expr(store_rstd): Rstd_buf = fx.rocdl.make_buffer_tensor(Rstd) - rstd_div = fx.logical_divide(Rstd_buf, fx.make_layout(1, 1)) - rstd_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 32) def wave_reduce_add(x): w = x @@ -368,15 +364,6 @@ def block_reduce_add(val): elem_bits, ) gamma_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), weight_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) - gamma_copy_atom_s = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if weight_elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - weight_elem_bits, - ) c_zero_f = fx.Float32(0.0) thread_sumsq = c_zero_f @@ -401,7 +388,7 @@ def block_reduce_add(val): if const_expr(TAIL_ELEMS > 0): if tid < TAIL_ELEMS: tail_idx = tid + VEC_ELEMS - x_tail_e = _load_scalar(copy_atom_s, elem_dtype, row_div_s, tail_idx) + x_tail_e = row_in[tail_idx] x_tail = x_tail_e.to(fx.Float32) thread_sumsq = thread_sumsq + x_tail * x_tail @@ -412,7 +399,7 @@ def block_reduce_add(val): if const_expr(store_rstd): if tid == 0: - _store_scalar(rstd_copy_atom, fx.Float32, rstd_div, bid, rrms) + Rstd_buf[bid] = rrms # Pass 2: normalize + gamma + store (reuse cached input) for tile_i in range_constexpr(NUM_VEC_ITERS): @@ -431,13 +418,13 @@ def block_reduce_add(val): if const_expr(TAIL_ELEMS > 0): if tid < TAIL_ELEMS: tail_idx = tid + VEC_ELEMS - x_tail_e = _load_scalar(copy_atom_s, elem_dtype, row_div_s, tail_idx) - g_tail_e = _load_scalar(gamma_copy_atom_s, weight_elem_dtype, gamma_div_s, tail_idx) + x_tail_e = row_in[tail_idx] + g_tail_e = Gamma_buf[tail_idx] x_tail = x_tail_e.to(fx.Float32) g_tail = g_tail_e if weight_dtype_str == "f32" else 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, out_div_s, tail_idx, y_tail_e) + row_out[tail_idx] = y_tail_e else: # Scalar fallback for arbitrary N. @@ -448,19 +435,6 @@ def block_reduce_add(val): row_in = fx.slice(Input_buf, (bid, None)) row_out = fx.slice(Output_buf, (bid, None)) - copy_atom_s = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - elem_bits, - ) - gamma_copy_atom_s = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if weight_elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - weight_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)) - c_zero_f = fx.Float32(0.0) thread_sumsq = c_zero_f @@ -468,7 +442,7 @@ def block_reduce_add(val): idx = tid + 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_e = row_in[idx_safe] x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) x2 = x * x x2_safe = is_valid.select(x2, c_zero_f) @@ -481,19 +455,19 @@ def block_reduce_add(val): if const_expr(store_rstd): if tid == 0: - _store_scalar(rstd_copy_atom, fx.Float32, rstd_div, bid, rrms) + Rstd_buf[bid] = rrms for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): idx = tid + base_idx_int if idx < N: - x_e = _load_scalar(copy_atom_s, elem_dtype, row_div, idx) - g_e = _load_scalar(gamma_copy_atom_s, weight_elem_dtype, gamma_div, idx) + x_e = row_in[idx] + g_e = Gamma_buf[idx] x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) g = g_e if weight_dtype_str == "f32" else g_e.to(fx.Float32) norm = x * rrms y = norm * g y_e = _to_elem_scalar(dtype_str, elem_dtype, y) - _store_scalar(copy_atom_s, elem_dtype, out_div, idx, y_e) + row_out[idx] = y_e if N <= SMALL_N_THRESHOLD: return _build_rmsnorm_small_n_module( @@ -622,8 +596,6 @@ def rmsnorm_large_m_small_n_kernel( Output_buf = fx.rocdl.make_buffer_tensor(Output) if const_expr(store_rstd): Rstd_buf = fx.rocdl.make_buffer_tensor(Rstd) - rstd_div = fx.logical_divide(Rstd_buf, fx.make_layout(1, 1)) - rstd_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 32) row_in = fx.slice(Input_buf, (row, None)) row_out = fx.slice(Output_buf, (row, None)) @@ -646,15 +618,6 @@ def group_reduce_add(x): copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) gamma_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), weight_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) - gamma_copy_atom_s = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if weight_elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - weight_elem_bits, - ) in_local = [] for tile_i in range_constexpr(NUM_VEC_ROW_ITERS): @@ -671,7 +634,7 @@ def group_reduce_add(x): 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_e = row_in[tail_idx] x_tail = x_tail_e.to(fx.Float32) thread_sumsq = thread_sumsq + x_tail * x_tail @@ -682,7 +645,7 @@ def group_reduce_add(x): if const_expr(store_rstd): if lane == 0: - _store_scalar(rstd_copy_atom, fx.Float32, rstd_div, row, rrms) + Rstd_buf[row] = rrms for tile_i in range_constexpr(NUM_VEC_ROW_ITERS): idx = lane + tile_i * THREADS_PER_ROW @@ -702,32 +665,20 @@ def group_reduce_add(x): 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(gamma_copy_atom_s, weight_elem_dtype, gamma_div_s, tail_idx) + x_tail_e = row_in[tail_idx] + g_tail_e = Gamma_buf[tail_idx] x_tail = x_tail_e.to(fx.Float32) g_tail = g_tail_e if weight_dtype_str == "f32" else 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, out_div_s, tail_idx, y_tail_e) + row_out[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, - ) - gamma_copy_atom_s = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if weight_elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - weight_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)) 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_e = row_in[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) @@ -739,18 +690,18 @@ def group_reduce_add(x): if const_expr(store_rstd): if lane == 0: - _store_scalar(rstd_copy_atom, fx.Float32, rstd_div, row, rrms) + Rstd_buf[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(gamma_copy_atom_s, weight_elem_dtype, gamma_div, idx) + x_e = row_in[idx] + g_e = Gamma_buf[idx] x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) g = g_e if weight_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, out_div, idx, y_e) + row_out[idx] = y_e if store_rstd: @@ -871,8 +822,6 @@ def fused_add_rmsnorm_kernel( if const_expr(store_rstd): Rstd_buf = fx.rocdl.make_buffer_tensor(Rstd) - rstd_div = fx.logical_divide(Rstd_buf, fx.make_layout(1, 1)) - rstd_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 32) def wave_reduce_add(x): w = x @@ -971,7 +920,7 @@ def block_reduce_add2(val0, val1): if const_expr(store_rstd): if tid == 0: - _store_scalar(rstd_copy_atom, fx.Float32, rstd_div, bid, rrms) + Rstd_buf[bid] = rrms # Pass 2: normalize + gamma + store (reuse cached added values) for tile_i in range_constexpr(NUM_VEC_ITERS): @@ -996,21 +945,6 @@ def block_reduce_add2(val0, val1): row_out = fx.slice(Output_buf, (bid, None)) row_residual_out = fx.slice(ResidualOut_buf, (bid, None)) - copy_atom_s = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - elem_bits, - ) - gamma_copy_atom_s = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if weight_elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - weight_elem_bits, - ) - - row_div = fx.logical_divide(row_in, fx.make_layout(1, 1)) - residual_in_div = fx.logical_divide(row_residual_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)) - residual_out_div = fx.logical_divide(row_residual_out, fx.make_layout(1, 1)) - c_zero_f = fx.Float32(0.0) thread_sumsq = c_zero_f @@ -1018,13 +952,13 @@ def block_reduce_add2(val0, val1): idx = tid + 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) - residual_e = _load_scalar(copy_atom_s, elem_dtype, residual_in_div, idx_safe) + x_e = row_in[idx_safe] + residual_e = row_residual_in[idx_safe] x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) residual = residual_e if dtype_str == "f32" else residual_e.to(fx.Float32) added_e = _to_elem_scalar(dtype_str, elem_dtype, x + residual) if idx < N: - _store_scalar(copy_atom_s, elem_dtype, residual_out_div, idx, added_e) + row_residual_out[idx] = added_e added = added_e if dtype_str == "f32" else added_e.to(fx.Float32) added2 = added * added thread_sumsq = thread_sumsq + is_valid.select(added2, c_zero_f) @@ -1036,18 +970,18 @@ def block_reduce_add2(val0, val1): if const_expr(store_rstd): if tid == 0: - _store_scalar(rstd_copy_atom, fx.Float32, rstd_div, bid, rrms) + Rstd_buf[bid] = rrms for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): idx = tid + base_idx_int if idx < N: - g_e = _load_scalar(gamma_copy_atom_s, weight_elem_dtype, gamma_div, idx) - added_e = _load_scalar(copy_atom_s, elem_dtype, residual_out_div, idx) + g_e = Gamma_buf[idx] + added_e = row_residual_out[idx] g = g_e if weight_dtype_str == "f32" else g_e.to(fx.Float32) added = added_e if dtype_str == "f32" else added_e.to(fx.Float32) y = (added * rrms) * g y_e = _to_elem_scalar(dtype_str, elem_dtype, y) - _store_scalar(copy_atom_s, elem_dtype, out_div, idx, y_e) + row_out[idx] = y_e if store_rstd: @@ -1135,8 +1069,6 @@ def rmsnorm_quant_kernel( s_red2 = lds.s_red2.view(fx.make_layout(RED_SLOTS, 1)) YScale_buf = fx.rocdl.make_buffer_tensor(YScale) - yscale_div = fx.logical_divide(YScale_buf, fx.make_layout(1, 1)) - scale_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 32) def wave_reduce_add(x): w = x @@ -1290,7 +1222,7 @@ def block_reduce_max(val): final_scale = (scale == c_zero_f).select(c_one_f, scale) if tid == 0: - _store_scalar(scale_copy_atom, fx.Float32, yscale_div, bid, final_scale) + YScale_buf[bid] = final_scale inv_scale = c_one_f / final_scale @@ -1314,24 +1246,8 @@ def block_reduce_max(val): if const_expr(is_smooth): XScale_buf = fx.rocdl.make_buffer_tensor(XScale) - copy_atom_s = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - elem_bits, - ) - copy_atom_qs = fx.make_copy_atom(fx.rocdl.BufferCopy(8), 8) - if const_expr(is_smooth): - copy_atom_xs = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - elem_bits, - ) - row_in = fx.slice(Input_buf, (bid, None)) row_out = fx.slice(Output_buf, (bid, None)) - 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)) - if const_expr(is_smooth): - xscale_div = fx.logical_divide(XScale_buf, fx.make_layout(1, 1)) def _abs_scalar(val): is_neg = val < c_zero_f @@ -1345,7 +1261,7 @@ def _abs_scalar(val): idx = tid + 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_e = row_in[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) @@ -1361,13 +1277,13 @@ def _abs_scalar(val): idx = tid + 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) - g_e = _load_scalar(copy_atom_s, elem_dtype, gamma_div, idx_safe) + x_e = row_in[idx_safe] + g_e = Gamma_buf[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 if const_expr(is_smooth): - s_e = _load_scalar(copy_atom_xs, elem_dtype, xscale_div, idx_safe) + s_e = XScale_buf[idx_safe] s = s_e if dtype_str == "f32" else s_e.to(fx.Float32) y = y * s y_abs = _abs_scalar(y) @@ -1378,7 +1294,7 @@ def _abs_scalar(val): final_scale = (scale == c_zero_f).select(c_one_f, scale) if tid == 0: - _store_scalar(scale_copy_atom, fx.Float32, yscale_div, bid, final_scale) + YScale_buf[bid] = final_scale inv_scale = c_one_f / final_scale @@ -1386,18 +1302,18 @@ def _abs_scalar(val): for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): idx = tid + 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_e = row_in[idx] + g_e = Gamma_buf[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 if const_expr(is_smooth): - s_e = _load_scalar(copy_atom_xs, elem_dtype, xscale_div, idx) + s_e = XScale_buf[idx] s = s_e if dtype_str == "f32" else s_e.to(fx.Float32) y = y * s q = y * inv_scale q_i8 = q.to(quant_dtype) - _store_scalar(copy_atom_qs, quant_dtype, out_div, idx, q_i8) + row_out[idx] = q_i8 if is_smooth: @@ -1522,8 +1438,6 @@ def fused_add_rmsnorm_quant_kernel( s_red2 = lds.s_red2.view(fx.make_layout(RED_SLOTS, 1)) YScale_buf = fx.rocdl.make_buffer_tensor(YScale) - yscale_div = fx.logical_divide(YScale_buf, fx.make_layout(1, 1)) - scale_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 32) def wave_reduce_add(x): w = x @@ -1686,7 +1600,7 @@ def block_reduce_max(val): final_scale = (scale == c_zero_f).select(c_one_f, scale) if tid == 0: - _store_scalar(scale_copy_atom, fx.Float32, yscale_div, bid, final_scale) + YScale_buf[bid] = final_scale inv_scale = c_one_f / final_scale @@ -1712,30 +1626,11 @@ def block_reduce_max(val): if const_expr(is_smooth): XScale_buf = fx.rocdl.make_buffer_tensor(XScale) - copy_atom_s = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - elem_bits, - ) - copy_atom_qs = fx.make_copy_atom(fx.rocdl.BufferCopy(8), 8) - if const_expr(is_smooth): - copy_atom_xs = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - elem_bits, - ) - row_in = fx.slice(Input_buf, (bid, None)) row_residual_in = fx.slice(ResidualIn_buf, (bid, None)) row_out = fx.slice(Output_buf, (bid, None)) row_residual_out = fx.slice(ResidualOut_buf, (bid, None)) - row_div = fx.logical_divide(row_in, fx.make_layout(1, 1)) - residual_in_div = fx.logical_divide(row_residual_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)) - residual_out_div = fx.logical_divide(row_residual_out, fx.make_layout(1, 1)) - if const_expr(is_smooth): - xscale_div = fx.logical_divide(XScale_buf, fx.make_layout(1, 1)) - def _abs_scalar(val): is_neg = val < c_zero_f neg_val = c_zero_f - val @@ -1748,13 +1643,13 @@ def _abs_scalar(val): idx = tid + 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) - residual_e = _load_scalar(copy_atom_s, elem_dtype, residual_in_div, idx_safe) + x_e = row_in[idx_safe] + residual_e = row_residual_in[idx_safe] x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) residual = residual_e if dtype_str == "f32" else residual_e.to(fx.Float32) added_e = _to_elem_scalar(dtype_str, elem_dtype, x + residual) if idx < N: - _store_scalar(copy_atom_s, elem_dtype, residual_out_div, idx, added_e) + row_residual_out[idx] = added_e added = added_e if dtype_str == "f32" else added_e.to(fx.Float32) added2 = added * added thread_sumsq = thread_sumsq + is_valid.select(added2, c_zero_f) @@ -1770,13 +1665,13 @@ def _abs_scalar(val): idx = tid + base_idx_int is_valid = idx < N idx_safe = is_valid.select(idx, 0) - g_e = _load_scalar(copy_atom_s, elem_dtype, gamma_div, idx_safe) - added_e = _load_scalar(copy_atom_s, elem_dtype, residual_out_div, idx_safe) + g_e = Gamma_buf[idx_safe] + added_e = row_residual_out[idx_safe] g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) added = added_e if dtype_str == "f32" else added_e.to(fx.Float32) y = (added * rrms) * g if const_expr(is_smooth): - s_e = _load_scalar(copy_atom_xs, elem_dtype, xscale_div, idx_safe) + s_e = XScale_buf[idx_safe] s = s_e if dtype_str == "f32" else s_e.to(fx.Float32) y = y * s y_abs = _abs_scalar(y) @@ -1787,7 +1682,7 @@ def _abs_scalar(val): final_scale = (scale == c_zero_f).select(c_one_f, scale) if tid == 0: - _store_scalar(scale_copy_atom, fx.Float32, yscale_div, bid, final_scale) + YScale_buf[bid] = final_scale inv_scale = c_one_f / final_scale @@ -1795,18 +1690,18 @@ def _abs_scalar(val): for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): idx = tid + base_idx_int if idx < N: - g_e = _load_scalar(copy_atom_s, elem_dtype, gamma_div, idx) - added_e = _load_scalar(copy_atom_s, elem_dtype, residual_out_div, idx) + g_e = Gamma_buf[idx] + added_e = row_residual_out[idx] g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) added = added_e if dtype_str == "f32" else added_e.to(fx.Float32) y = (added * rrms) * g if const_expr(is_smooth): - s_e = _load_scalar(copy_atom_xs, elem_dtype, xscale_div, idx) + s_e = XScale_buf[idx] s = s_e if dtype_str == "f32" else s_e.to(fx.Float32) y = y * s q = y * inv_scale q_i8 = q.to(quant_dtype) - _store_scalar(copy_atom_qs, quant_dtype, out_div, idx, q_i8) + row_out[idx] = q_i8 if is_smooth: From fc799d9af2f8c52913c671020e2d7583fef377d1 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Mon, 21 Sep 2026 18:24:12 +0800 Subject: [PATCH 23/25] fix(norm): anchor the quantized-norm dtype guard to the builder dtype_str validate_norm_operand_dtypes compared every operand against Input.dtype but never against dtype_str, so a self-consistent operand set that did not match the builder specialization passed the guard and was reinterpreted bit-for-bit by the kernel: building a bf16 module and passing fp16 tensors produced corrupted residuals and int8 output with no error raised. Take dtype_str in the validator and reject Input.dtype != dtype_to_elem_type(dtype_str) ahead of the existing operand comparison, then pass dtype_str at all eight quantized launchers in rmsnorm_kernel.py and layernorm_kernel.py. Add negative tests covering the four RMSNorm and the four LayerNorm quantized launchers; both fail with DID NOT RAISE ValueError when the validator change is reverted. Update the RMSNorm quantized-variant paragraph in docs/prebuilt_kernels_guide.md, which stated the same Input-relative contract. Co-authored-by: Cursor --- docs/prebuilt_kernels_guide.md | 12 +++++------ kernels/norm/layernorm_kernel.py | 10 +++++---- kernels/norm/rmsnorm_common.py | 18 +++++++++++----- kernels/norm/rmsnorm_kernel.py | 8 +++---- tests/kernels/test_layernorm.py | 36 ++++++++++++++++++++++++++++++++ tests/kernels/test_rmsnorm.py | 35 +++++++++++++++++++++++++++++++ 6 files changed, 100 insertions(+), 19 deletions(-) diff --git a/docs/prebuilt_kernels_guide.md b/docs/prebuilt_kernels_guide.md index 6bd1ec4d6f..0de877b400 100644 --- a/docs/prebuilt_kernels_guide.md +++ b/docs/prebuilt_kernels_guide.md @@ -76,12 +76,12 @@ per-row reciprocal std (`rstd`) for use by the backward pass. support FP32 weights. **Quantized variants:** The DynamicQuant and SmoothQuant builders emit int8 -`Output` and fp32 per-row `YScale`. Every other operand — `Gamma`, the -fused-add `ResidualIn`/`ResidualOut`, and SmoothQuant `XScale` — must use the -same element dtype as `Input`. The launchers raise `ValueError` on a mismatch, -so a wrong dtype fails at compile time rather than silently producing corrupted -scales. Unlike the plain forward, the quantized builders do not accept FP32 -weights with FP16/BF16 activations. +`Output` and fp32 per-row `YScale`. `Input` must use the element dtype named by +the builder's `dtype_str`, and every other operand — `Gamma`, the fused-add +`ResidualIn`/`ResidualOut`, and SmoothQuant `XScale` — must match it. The +launchers raise `ValueError` on a mismatch, so a wrong dtype fails at compile +time rather than silently producing corrupted scales. Unlike the plain forward, +the quantized builders do not accept FP32 weights with FP16/BF16 activations. **Backward:** `build_rmsnorm_bwd_module(N, dtype_str, weight_dtype_str=None)` builds the fused RMSNorm backward kernel (grid `(M,)`, diff --git a/kernels/norm/layernorm_kernel.py b/kernels/norm/layernorm_kernel.py index 7efc3d6df1..c8486dd591 100644 --- a/kernels/norm/layernorm_kernel.py +++ b/kernels/norm/layernorm_kernel.py @@ -1028,7 +1028,7 @@ def launch_layernorm_smoothquant( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): - validate_norm_operand_dtypes(Input, Gamma=Gamma, Beta=Beta, XScale=XScale) + validate_norm_operand_dtypes(Input, dtype_str, Gamma=Gamma, Beta=Beta, XScale=XScale) launcher = layernorm_quant_kernel(Input, Gamma, Beta, XScale, YScale, Output) launcher.launch( grid=(m_in, 1, 1), @@ -1050,7 +1050,7 @@ def launch_layernorm_dynamicquant( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): - validate_norm_operand_dtypes(Input, Gamma=Gamma, Beta=Beta) + validate_norm_operand_dtypes(Input, dtype_str, Gamma=Gamma, Beta=Beta) launcher = layernorm_quant_kernel(Input, Gamma, Beta, Gamma, YScale, Output) launcher.launch( grid=(m_in, 1, 1), @@ -1401,7 +1401,7 @@ def launch_fused_add_layernorm_smoothquant( stream: fx.Stream = fx.Stream(None), ): validate_norm_operand_dtypes( - Input, ResidualIn=ResidualIn, Gamma=Gamma, Beta=Beta, XScale=XScale, ResidualOut=ResidualOut + Input, dtype_str, ResidualIn=ResidualIn, Gamma=Gamma, Beta=Beta, XScale=XScale, ResidualOut=ResidualOut ) launcher = fused_add_layernorm_quant_kernel( Input, ResidualIn, Gamma, Beta, XScale, YScale, Output, ResidualOut @@ -1428,7 +1428,9 @@ def launch_fused_add_layernorm_dynamicquant( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): - validate_norm_operand_dtypes(Input, ResidualIn=ResidualIn, Gamma=Gamma, Beta=Beta, ResidualOut=ResidualOut) + validate_norm_operand_dtypes( + Input, dtype_str, ResidualIn=ResidualIn, Gamma=Gamma, Beta=Beta, ResidualOut=ResidualOut + ) launcher = fused_add_layernorm_quant_kernel( Input, ResidualIn, Gamma, Beta, Gamma, YScale, Output, ResidualOut ) diff --git a/kernels/norm/rmsnorm_common.py b/kernels/norm/rmsnorm_common.py index 834f5257fc..e5e2ee78f9 100644 --- a/kernels/norm/rmsnorm_common.py +++ b/kernels/norm/rmsnorm_common.py @@ -12,7 +12,7 @@ import flydsl.expr as fx from flydsl.expr import const_expr -from kernels.common.kernels_common import get_warp_size +from kernels.common.kernels_common import dtype_to_elem_type, get_warp_size KERNEL_NAME = "rmsnorm" @@ -52,12 +52,20 @@ class SharedStorage: return SharedStorage -def validate_norm_operand_dtypes(input_tensor, **operands) -> None: - """Validate quantized-norm operands against the input dtype. +def validate_norm_operand_dtypes(input_tensor, dtype_str: str, **operands) -> None: + """Validate quantized-norm operands against the builder specialization. - Quantized norm kernels specialize every non-quantized operand on the activation - dtype, so a mismatched tensor is reinterpreted bit-for-bit instead of failing. + Quantized norm kernels specialize every non-quantized operand on ``dtype_str``, + so a mismatched tensor is reinterpreted bit-for-bit instead of failing. + ``Input`` and each operand are checked against ``dtype_str`` itself rather than + against one another. """ + expected = dtype_to_elem_type(dtype_str) + if input_tensor.dtype != expected: + raise ValueError( + f"quantized norm kernels specialize on dtype_str={dtype_str!r}, so Input must be " + f"{expected}, got {input_tensor.dtype}" + ) mismatched = {name: t.dtype for name, t in operands.items() if t.dtype != input_tensor.dtype} if mismatched: got = ", ".join(f"{name}={dtype}" for name, dtype in mismatched.items()) diff --git a/kernels/norm/rmsnorm_kernel.py b/kernels/norm/rmsnorm_kernel.py index d57908ae37..0c36b01751 100644 --- a/kernels/norm/rmsnorm_kernel.py +++ b/kernels/norm/rmsnorm_kernel.py @@ -1327,7 +1327,7 @@ def launch_rmsnorm_smoothquant( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): - _validate_norm_operand_dtypes(Input, Gamma=Gamma, XScale=XScale) + _validate_norm_operand_dtypes(Input, dtype_str, Gamma=Gamma, XScale=XScale) launcher = rmsnorm_quant_kernel(Input, Gamma, XScale, YScale, Output) launcher.launch( grid=(m_in, 1, 1), @@ -1348,7 +1348,7 @@ def launch_rmsnorm_dynamicquant( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): - _validate_norm_operand_dtypes(Input, Gamma=Gamma) + _validate_norm_operand_dtypes(Input, dtype_str, Gamma=Gamma) launcher = rmsnorm_quant_kernel(Input, Gamma, Gamma, YScale, Output) launcher.launch( grid=(m_in, 1, 1), @@ -1718,7 +1718,7 @@ def launch_fused_add_rmsnorm_smoothquant( stream: fx.Stream = fx.Stream(None), ): _validate_norm_operand_dtypes( - Input, ResidualIn=ResidualIn, Gamma=Gamma, XScale=XScale, ResidualOut=ResidualOut + Input, dtype_str, ResidualIn=ResidualIn, Gamma=Gamma, XScale=XScale, ResidualOut=ResidualOut ) launcher = fused_add_rmsnorm_quant_kernel(Input, ResidualIn, Gamma, XScale, YScale, Output, ResidualOut) launcher.launch( @@ -1742,7 +1742,7 @@ def launch_fused_add_rmsnorm_dynamicquant( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): - _validate_norm_operand_dtypes(Input, ResidualIn=ResidualIn, Gamma=Gamma, ResidualOut=ResidualOut) + _validate_norm_operand_dtypes(Input, dtype_str, ResidualIn=ResidualIn, Gamma=Gamma, ResidualOut=ResidualOut) launcher = fused_add_rmsnorm_quant_kernel(Input, ResidualIn, Gamma, Gamma, YScale, Output, ResidualOut) launcher.launch( grid=(m_in, 1, 1), diff --git a/tests/kernels/test_layernorm.py b/tests/kernels/test_layernorm.py index 0d3b4e04d9..ca0d7b20fc 100644 --- a/tests/kernels/test_layernorm.py +++ b/tests/kernels/test_layernorm.py @@ -1461,6 +1461,41 @@ def test_layernorm_dtype_mismatch(): print(" -> PASSED") +def test_quantized_layernorm_builder_dtype_mismatch(): + """Quantized launchers must reject an operand set that disagrees with ``dtype_str``.""" + print("=" * 80) + print("Running quantized LayerNorm builder-dtype guard Test") + print("=" * 80) + M, N = 4, 256 + stream = torch.cuda.current_stream() + # Self-consistent FP16 operands: only the builder specialization disagrees. + x = torch.randn((M, N), device="cuda", dtype=DTYPE_FP16) + residual = torch.randn((M, N), device="cuda", dtype=DTYPE_FP16) + gamma = torch.rand((N,), device="cuda", dtype=DTYPE_FP16) + beta = torch.rand((N,), device="cuda", dtype=DTYPE_FP16) + xscale = torch.rand((N,), device="cuda", dtype=DTYPE_FP16) + residual_out = torch.empty((M, N), device="cuda", dtype=DTYPE_FP16) + out = torch.empty((M, N), device="cuda", dtype=DTYPE_INT8) + yscale = torch.empty((M,), device="cuda", dtype=DTYPE_FP32) + + cases = ( + (build_layernorm_dynamicquant_module(N, "bf16"), (x, gamma, beta, out, yscale, M, stream)), + (build_layernorm_smoothquant_module(N, "bf16"), (x, gamma, beta, xscale, out, yscale, M, stream)), + ( + build_fused_add_layernorm_dynamicquant_module(N, "bf16"), + (x, residual, gamma, beta, out, residual_out, yscale, M, stream), + ), + ( + build_fused_add_layernorm_smoothquant_module(N, "bf16"), + (x, residual, gamma, beta, xscale, out, residual_out, yscale, M, stream), + ), + ) + for launch_fn, args in cases: + with pytest.raises(ValueError, match="dtype_str"): + flyc.compile(launch_fn, *args) + print(" -> PASSED") + + def test_layernorm_affine_shape_mismatch(): """Pointer views require exactly N affine elements, not merely a trailing N.""" N = 256 @@ -1485,6 +1520,7 @@ def test_layernorm_affine_shape_mismatch(): test_layernorm_eps_honored() test_layernorm_multi_gpu() test_layernorm_dtype_mismatch() + test_quantized_layernorm_builder_dtype_mismatch() test_fused_add_layernorm() test_layernorm_dynamicquant() test_layernorm_smoothquant() diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index c9b19ca271..1927289842 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -1262,6 +1262,40 @@ def test_fused_add_rmsnorm_dtype_mismatch(): print(" -> PASSED") +def test_quantized_rmsnorm_builder_dtype_mismatch(): + """Quantized launchers must reject an operand set that disagrees with ``dtype_str``.""" + print("=" * 80) + print("Running quantized RMSNorm builder-dtype guard Test") + print("=" * 80) + M, N = 4, 256 + stream = torch.cuda.current_stream() + # Self-consistent FP16 operands: only the builder specialization disagrees. + x = torch.randn((M, N), device="cuda", dtype=DTYPE_FP16) + residual = torch.randn((M, N), device="cuda", dtype=DTYPE_FP16) + gamma = torch.rand((N,), device="cuda", dtype=DTYPE_FP16) + xscale = torch.rand((N,), device="cuda", dtype=DTYPE_FP16) + residual_out = torch.empty((M, N), device="cuda", dtype=DTYPE_FP16) + out = torch.empty((M, N), device="cuda", dtype=DTYPE_INT8) + yscale = torch.empty((M,), device="cuda", dtype=DTYPE_FP32) + + cases = ( + (build_rmsnorm_dynamicquant_module(N, "bf16"), (x, gamma, out, yscale, M, stream)), + (build_rmsnorm_smoothquant_module(N, "bf16"), (x, gamma, xscale, out, yscale, M, stream)), + ( + build_fused_add_rmsnorm_dynamicquant_module(N, "bf16"), + (x, residual, gamma, out, residual_out, yscale, M, stream), + ), + ( + build_fused_add_rmsnorm_smoothquant_module(N, "bf16"), + (x, residual, gamma, xscale, out, residual_out, yscale, M, stream), + ), + ) + for launch_fn, args in cases: + with pytest.raises(ValueError, match="dtype_str"): + flyc.compile(launch_fn, *args) + print(" -> PASSED") + + @pytest.mark.multi_gpu def test_fused_add_rmsnorm_device_mismatch(): """Operands on different devices must be rejected (kernel binds to x.device).""" @@ -2167,6 +2201,7 @@ def test_fused_add_rmsnorm_smoothquant(): test_fused_add_rmsnorm_backward() test_fused_add_rmsnorm_autograd() test_fused_add_rmsnorm_dtype_mismatch() + test_quantized_rmsnorm_builder_dtype_mismatch() test_fused_add_rmsnorm_dynamicquant() test_fused_add_rmsnorm_smoothquant() if os.environ.get("ROCDSL_COMPARE_TORCH", "0") == "1": From 8c63e3be02a327897b935fd23503e1bda6af6e47 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Thu, 24 Sep 2026 17:44:06 +0800 Subject: [PATCH 24/25] docs(norm): resync the LayerNorm and RMSNorm guide sections with the kernels Several claims in the norm sections predate the current kernels. The LayerNorm vector path is gated on a 16-bit element type and `N % VEC_WIDTH == 0` rather than the old `N == BLOCK_THREADS * VEC_WIDTH * 4` fast path, and a row that is not divisible falls back to a fully scalar two-pass body, so both the guide bullets and the module docstring are rewritten around `USE_VEC_N`. `VEC_ALIGN` and `USE_NONTEMPORAL` were dropped from layernorm_kernel.py and no longer exist anywhere under kernels/, so their table rows go away, and `WARP_SIZE` now reads as `get_warp_size()` resolves it. RMSNorm's `BLOCK_THREADS` defaults to `None` and resolves through `default_block_threads` to 512 on gfx95x rows of `N >= 8192`, which the section now states in a table matching LayerNorm and Softmax. Its forward is two passes with the row held in registers, not three with an LDS row cache; LDS carries only the cross-wave reduction slots, sized by wave count rather than by N. Both "Kernel signature" blocks described a GPU_MODULE_NAME / self / __call__ class API that no longer exists and are replaced by the actual @flyc.kernel and @flyc.jit signatures. Co-authored-by: Cursor --- docs/prebuilt_kernels_guide.md | 59 +++++++++++++++++++------------- kernels/norm/layernorm_kernel.py | 7 ++-- 2 files changed, 40 insertions(+), 26 deletions(-) diff --git a/docs/prebuilt_kernels_guide.md b/docs/prebuilt_kernels_guide.md index 0de877b400..7769642c43 100644 --- a/docs/prebuilt_kernels_guide.md +++ b/docs/prebuilt_kernels_guide.md @@ -34,29 +34,28 @@ executor = build_layernorm_module(N=8192, dtype_str="bf16") | Constant | Value | Description | |---|---|---| | `BLOCK_THREADS` | 256 | Threads per block | -| `WARP_SIZE` | 64 | AMD wavefront size | +| `WARP_SIZE` | 64 on CDNA, 32 on RDNA | Wavefront size, resolved from the target arch | | `VEC_WIDTH` | 8 | Vector load/store width | -| `VEC_ALIGN` | 16 | Alignment for vector ops (bytes) | | `EPS` | 1e-5 | Numerical stability epsilon | -| `USE_NONTEMPORAL` | True | Non-temporal stores for output | **Algorithm:** - **Two-pass normalization**: Pass 1 computes mean and variance, Pass 2 applies affine transform -- **Fast path**: When `N == BLOCK_THREADS * VEC_WIDTH * 4` (for example, N=8192), uses fully register-resident computation with no scalar tail -- **Generic path**: Handles arbitrary N with vector body + scalar tail +- **Vectorized path**: When the element type is 16-bit and `N % VEC_WIDTH == 0`, the row is covered by `N / VEC_WIDTH` vector tiles with no scalar tail +- **Scalar path**: FP32, or any `N` not divisible by `VEC_WIDTH`, falls back to a fully scalar two-pass implementation - **bf16 handling**: Software round-to-nearest-even (RNE) pack on gfx942; hardware `cvt_pk_bf16_f32` on gfx950+ - **Warp reduction**: XOR-shuffle-based intra-wave reduction (shifts: 32, 16, 8, 4, 2, 1), then LDS-based cross-wave synchronization -**Kernel signature** (using `@flyc.kernel` API): -``` -GPU_MODULE_NAME = "layernorm_module" - -@kernel -layernorm_kernel(self, Input, Gamma, Beta, Output, m_in) +**Kernel signature:** +```python +@flyc.kernel +layernorm_kernel(Input, Gamma, Beta, Output, Mean, Rstd) -@jit -__call__(self, Input, Gamma, Beta, Output, m_in) +@flyc.jit +launch_layernorm(Input, Gamma, Beta, Output, m_in, stream=...) +# store_stats=True inserts Mean and Rstd before m_in ``` +The builder returns the `launch_layernorm` closure. The row count `m_in` is a +runtime launch argument, not a kernel parameter. ### 1.2 RMSNorm (`kernels/norm/rmsnorm_kernel.py`) @@ -70,7 +69,7 @@ executor = build_rmsnorm_module(N=8192, dtype_str="bf16", store_rstd=False) ``` `build_rmsnorm_module(N, dtype_str, store_rstd=False, eps=EPS, -BLOCK_THREADS=BLOCK_THREADS, weight_dtype_str=None)` optionally writes the +BLOCK_THREADS=None, weight_dtype_str=None)` optionally writes the per-row reciprocal std (`rstd`) for use by the backward pass. `weight_dtype_str` defaults to `dtype_str`; FP16/BF16 activations additionally support FP32 weights. @@ -92,19 +91,33 @@ grad). The forward bakes `eps` into `Rstd`, so the backward does not need it. The public plain and fused-add training wrappers return `dweight` in the original weight dtype. -**Configuration constants:** Same as LayerNorm (BLOCK_THREADS=256, VEC_WIDTH=8, etc.) +**Configuration constants:** +| Constant | Value | Description | +|---|---|---| +| `BLOCK_THREADS` | 256; 512 on gfx95x when `N >= 8192` | Resolved by `default_block_threads(N, arch)` when the builder argument is left at `None` | +| `WARP_SIZE` | 64 on CDNA, 32 on RDNA | Wavefront size, resolved from the target arch | +| `VEC_WIDTH` | 8 | Vector load/store width | +| `EPS` | 1e-5 | Numerical stability epsilon | + +**Algorithm (2-pass, row cached in registers):** +1. **Pass 1**: One vectorized global read per row; the input stays in registers + and the sum of squares is accumulated in the same pass. A scalar tail covers + the `N % VEC_WIDTH` leftover elements. +2. **Pass 2**: Normalize, multiply by gamma, and store, reusing the registers + from pass 1. `Gamma` is preloaded during pass 1 on the gfx942 BF16 fast path. -**Algorithm (3-pass with LDS caching):** -1. **Pass 0**: Global → LDS row cache (one-pass global read, vectorized) -2. **Pass 1**: Sum-of-squares computation from LDS row cache -3. **Pass 2**: Normalize + gamma multiply + store with software pipeline for Gamma prefetch +LDS holds only the cross-wave reduction slots, sized by the wave count rather +than by `N`; the row itself never passes through shared memory. The quantized +builders add a third pass that applies the per-row scale. **Kernel signature:** -``` -GPU_MODULE_NAME = "rmsnorm_module" +```python +@flyc.kernel +rmsnorm_kernel(Input, Gamma, Rstd, Output) -@kernel -rmsnorm_kernel(self, Input, Gamma, Output, m_in) +@flyc.jit +launch_rmsnorm(Input, Gamma, Output, m_in, stream=...) +# store_rstd=True inserts Rstd between Output and m_in ``` --- diff --git a/kernels/norm/layernorm_kernel.py b/kernels/norm/layernorm_kernel.py index c8486dd591..b6b63976e3 100644 --- a/kernels/norm/layernorm_kernel.py +++ b/kernels/norm/layernorm_kernel.py @@ -6,9 +6,10 @@ 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. - - Generic path (arbitrary N): scalar 2-pass implementation. + - Vector path (16-bit dtype, N % VEC_WIDTH == 0): buffer-backed tiled copy + over N / VEC_WIDTH tiles, input cached in registers across both passes. + - Scalar path (fp32, or N not divisible by VEC_WIDTH): 2-pass scalar + implementation. """ import math From da396b887700dc6361a4c8cc4de223b699943b3a Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Mon, 28 Sep 2026 18:55:44 +0800 Subject: [PATCH 25/25] fix(norm): guard IR assertions against cache-loaded artifacts #1159 (336d6ff1) dropped source_ir from CompiledArtifact.__getstate__, so it reads back as None on any artifact restored from disk and the five norm assertions on pre-lowering IR dereference it unconditionally. Only test_rmsnorm_large_shape trips on CI's cold cache: build_rmsnorm_module is a factory, so each call gets a fresh JitFunction with an empty in-memory cache, and M is not a builder argument, so it reloads the (8192, bf16) entry test_rmsnorm wrote earlier in the session. The other four survive only because their keys are unique to one test; a second run against the same build fails all five. Guarding loses nothing: manager_key hashes the kernel source through the launcher's closure, so any edit that could drop the wide copy invalidates the entry and the assertion runs against a real compile. Co-authored-by: Cursor --- tests/kernels/test_rmsnorm.py | 15 ++++++++++++--- tests/kernels/test_rmsnorm_autotune.py | 9 ++++++--- 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index 1927289842..ee98795f25 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -187,8 +187,13 @@ def _quant_test_eps(N: int) -> float: def _assert_rmsnorm_forward_copy_width(compiled_fn, N: int, dtype: str): if dtype not in ("f16", "bf16"): return + source_ir = compiled_fn._keepalive.source_ir + if source_ir is None: + # Pre-lowering IR is process-local, so a cache-loaded artifact carries none. + # Kernel edits change the cache key, so a real regression still gets a cold compile. + return expects_vec8 = N >= 8 - has_128b_io = "buffer_copy<128>, 16>" in compiled_fn._keepalive.source_ir + has_128b_io = "buffer_copy<128>, 16>" in source_ir assert has_128b_io == expects_vec8, f"unexpected RMSNorm copy width for N={N}, dtype={dtype}" @@ -1619,7 +1624,9 @@ def test_rmsnorm_small_n_mixed_weight_tail_and_rstd(): ) compiled = flyc.compile(launcher, x, weight, output, rstd, M, stream) _assert_rmsnorm_forward_copy_width(compiled, N, "bf16") - assert "buffer_copy<128>, 32>" in compiled._keepalive.source_ir + source_ir = compiled._keepalive.source_ir + if source_ir is not None: + assert "buffer_copy<128>, 32>" in source_ir compiled(x, weight, output, rstd, M, stream) torch.cuda.synchronize(device) @@ -1708,7 +1715,9 @@ def test_rmsnorm_bwd_mixed_weight_keeps_vector_io(fused_add): ) weight_copy_type = "!fly.copy_atom, 32>" - assert compiled._keepalive.source_ir.count(weight_copy_type) >= 3 + source_ir = compiled._keepalive.source_ir + if source_ir is not None: + assert source_ir.count(weight_copy_type) >= 3 @pytest.mark.parametrize("fused_add", (False, True), ids=("plain", "fused_add")) diff --git a/tests/kernels/test_rmsnorm_autotune.py b/tests/kernels/test_rmsnorm_autotune.py index 810378a598..af1bc672a9 100644 --- a/tests/kernels/test_rmsnorm_autotune.py +++ b/tests/kernels/test_rmsnorm_autotune.py @@ -69,13 +69,16 @@ def test_rmsnorm_direct_specializes_known_block_size(weight_dtype, weight_dtype_ compiled = flyc.compile(rmsnorm_direct, *compile_args) stream.synchronize() artifact = compiled._keepalive + # Pre-lowering IR is process-local, so a cache-loaded artifact carries none. + source_ir = artifact.source_ir - assert "known_block_size = array" in artifact.source_ir + if source_ir is not None: + assert "known_block_size = array" in source_ir match = re.search(r"max_flat_workgroup_size\\CD\\([0-9A-Fa-f]{2})\\([0-9A-Fa-f]{2})", artifact.ir) assert match is not None and int("".join(match.groups()), 16) == 512 - if weight_dtype == torch.float32: + if weight_dtype == torch.float32 and source_ir is not None: weight_copy_type = "!fly.copy_atom, 32>" - assert artifact.source_ir.count(weight_copy_type) >= 3 + assert source_ir.count(weight_copy_type) >= 3 _assert_close(out, ref)