[Perf] Optimize rmsnorm/layernorm

kernels/norm/layernorm_kernel.py

@@ -6,8 +6,8 @@
 LayerNorm(x) = (x - mean) / sqrt(var + eps) * gamma + beta
 
 Two paths:
-  - Fast path (N == BLOCK_THREADS * VEC_WIDTH * 4): vectorised tiled copy,
-    register caching, pipelined gamma/beta loads.
+  - Fast path (N is a multiple of BLOCK_THREADS * VEC_WIDTH): vectorised
+    tiled copy, register caching, pipelined gamma/beta loads.
   - Generic path (arbitrary N): scalar 2-pass implementation.
 """
 
@@ -123,6 +123,7 @@ def build_layernorm_module(N: int, dtype_str: str, store_stats: bool = False, ep
 
     RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE)
     elem_bits = 32 if dtype_str == "f32" else 16
+    tile_cols = BLOCK_THREADS * VEC_WIDTH
 
     SharedStorage = _make_reduction_storage(RED_SLOTS)
 
@@ -221,12 +222,12 @@ def compute_mean_rstd(sum_val, sumsq_val):
             return mean, rstd
 
         # ==================================================================
-        # Fast path: N == BLOCK_THREADS * VEC_WIDTH * 4
+        # Fast path: N is a multiple of BLOCK_THREADS * VEC_WIDTH
         # Uses buffer_load / buffer_store for high-bandwidth vectorised
         # memory access (same approach as preshuffle_gemm).
         # ==================================================================
-        if const_expr(N == (BLOCK_THREADS * VEC_WIDTH * 4) and elem_bits <= 16):
-            num_tiles_py = 4
+        if const_expr(N >= tile_cols and N % tile_cols == 0 and elem_bits <= 16):
+            num_tiles_py = N // tile_cols
             c_zero_f = fx.Float32(0.0)
             thread_sum = c_zero_f
             thread_sumsq = c_zero_f
@@ -566,6 +567,7 @@ def build_fused_add_layernorm_module(N: int, dtype_str: str):
 
     RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE)
     elem_bits = 32 if dtype_str == "f32" else 16
+    tile_cols = BLOCK_THREADS * VEC_WIDTH
 
     SharedStorage = _make_reduction_storage(RED_SLOTS)
 
@@ -583,7 +585,7 @@ def fused_add_layernorm_kernel(
 
         elem_dtype = dtype_to_elem_type(dtype_str)
         fm_fast = arith.FastMathFlags.fast
-        eps_c = EPS
+        eps_c = eps
 
         lds = fx.SharedAllocator().allocate(SharedStorage).peek()
         s_sum = lds.s_sum.view(fx.make_layout(RED_SLOTS, 1))
@@ -637,10 +639,10 @@ def compute_mean_rstd(sum_val, sumsq_val):
             return mean, fmath.rsqrt(var + eps_c, fastmath=fm_fast)
 
         # ==================================================================
-        # Fast path: N == BLOCK_THREADS * VEC_WIDTH * 4
+        # Fast path: N is a multiple of BLOCK_THREADS * VEC_WIDTH
         # ==================================================================
-        if const_expr(N == (BLOCK_THREADS * VEC_WIDTH * 4) and elem_bits <= 16):
-            num_tiles_py = 4
+        if const_expr(N >= tile_cols and N % tile_cols == 0 and elem_bits <= 16):
+            num_tiles_py = N // tile_cols
             c_zero_f = fx.Float32(0.0)
             thread_sum = c_zero_f
             thread_sumsq = c_zero_f
@@ -786,9 +788,11 @@ def _build_layernorm_quant_module(
     *,
     is_smooth: bool,
     quant_dtype_str: str = "i8",
+    eps: float = EPS,
 ):
     RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE)
     elem_bits = 32 if dtype_str == "f32" else 16
+    tile_cols = BLOCK_THREADS * VEC_WIDTH
     quant_dtype_max = _quant_dtype_max(quant_dtype_str)
 
     SharedStorage = _make_reduction_storage(RED_SLOTS)
@@ -809,7 +813,7 @@ def layernorm_quant_kernel(
         quant_dtype = _quant_dtype_to_elem_type(quant_dtype_str)
 
         fm_fast = arith.FastMathFlags.fast
-        eps_c = EPS
+        eps_c = eps
         n_float = float(N)
         c_zero_f = fx.Float32(0.0)
         c_one_f = fx.Float32(1.0)
@@ -894,10 +898,10 @@ def block_reduce_max(val):
             return fx.memref_load(s_sum, 0)
 
         # ==================================================================
-        # Fast path: N == BLOCK_THREADS * VEC_WIDTH * 4
+        # Fast path: N is a multiple of BLOCK_THREADS * VEC_WIDTH
         # ==================================================================
-        if const_expr(N == (BLOCK_THREADS * VEC_WIDTH * 4) and elem_bits <= 16):
-            num_tiles_py = 4
+        if const_expr(N >= tile_cols and N % tile_cols == 0 and elem_bits <= 16):
+            num_tiles_py = N // tile_cols
             quant_half_width = VEC_WIDTH // 2
             abs_mask = full(VEC_WIDTH, fx.Uint32(0x7FFFFFFF), fx.Uint32)
 
@@ -1130,12 +1134,14 @@ def _build_fused_add_layernorm_quant_module(
     *,
     is_smooth: bool,
     quant_dtype_str: str = "i8",
+    eps: float = EPS,
 ):
     arch = get_rocm_arch()
     USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95")
 
     RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE)
     elem_bits = 32 if dtype_str == "f32" else 16
+    tile_cols = BLOCK_THREADS * VEC_WIDTH
     quant_dtype_max = _quant_dtype_max(quant_dtype_str)
 
     SharedStorage = _make_reduction_storage(RED_SLOTS)
@@ -1158,7 +1164,7 @@ def fused_add_layernorm_quant_kernel(
         quant_dtype = _quant_dtype_to_elem_type(quant_dtype_str)
 
         fm_fast = arith.FastMathFlags.fast
-        eps_c = EPS
+        eps_c = eps
         n_float = float(N)
         c_zero_f = fx.Float32(0.0)
         c_one_f = fx.Float32(1.0)
@@ -1243,10 +1249,10 @@ def block_reduce_max(val):
             return fx.memref_load(s_sum, 0)
 
         # ==================================================================
-        # Fast path: N == BLOCK_THREADS * VEC_WIDTH * 4
+        # Fast path: N is a multiple of BLOCK_THREADS * VEC_WIDTH
         # ==================================================================
-        if const_expr(N == (BLOCK_THREADS * VEC_WIDTH * 4) and elem_bits <= 16):
-            num_tiles_py = 4
+        if const_expr(N >= tile_cols and N % tile_cols == 0 and elem_bits <= 16):
+            num_tiles_py = N // tile_cols
             quant_half_width = VEC_WIDTH // 2
             abs_mask = full(VEC_WIDTH, fx.Uint32(0x7FFFFFFF), fx.Uint32)
 
@@ -1506,51 +1512,59 @@ 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,
     )
 
 
 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,
     )
 
 
 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,
     )
 
 
 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,
     )