From 36bb704a0649ab672c3406a44bba5efa517a1cde Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Tue, 28 Jul 2026 03:48:44 +0800 Subject: [PATCH 1/6] 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 3b54bd430..6e8edd9b1 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 c7175c2ca..c27dafd11 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 2/6] 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 6e8edd9b1..344eaad3f 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 c27dafd11..d6450336e 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 3/6] 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 f2f3f517a..4957b2083 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 fe7022cd9..d17b33374 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 4/6] 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 4957b2083..341151459 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 d17b33374..0d3b4e04d 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 5/6] 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 d6450336e..dcb17cdde 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 6/6] 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 dcb17cdde..f67c1affb 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")