diff --git a/docs/prebuilt_kernels_guide.md b/docs/prebuilt_kernels_guide.md index bf70b86056..0c0ef2f5c3 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,11 +69,19 @@ 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. +**Quantized variants:** The DynamicQuant and SmoothQuant builders emit int8 +`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,)`, one block per row). Kernel signature @@ -84,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 150341eb06..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 @@ -20,6 +21,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 @@ -45,21 +47,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) @@ -98,13 +85,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 @@ -123,6 +103,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) @@ -163,9 +146,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): @@ -221,12 +201,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,51 +218,39 @@ 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) if const_expr(store_stats): if tid == 0: - _store_scalar(stats_copy_atom, fx.Float32, fx.Float32, mean_div, 0, mean) - _store_scalar(stats_copy_atom, fx.Float32, fx.Float32, rstd_div, 0, rstd) - - g_cur = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, tid).to(fx.Float32) - b_cur = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, beta_div, tid).to(fx.Float32) + mean_row[0] = mean + rstd_row[0] = rstd # ── 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: # ============================================================== @@ -295,22 +260,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) @@ -323,16 +278,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) @@ -341,7 +296,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: @@ -401,7 +356,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 @@ -481,21 +435,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)) @@ -506,9 +447,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) @@ -525,9 +466,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) @@ -535,7 +476,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) @@ -560,12 +501,15 @@ 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") 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) @@ -583,7 +527,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 +581,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 +611,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: # ============================================================== @@ -710,18 +659,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 @@ -730,8 +667,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) @@ -740,7 +677,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) @@ -748,16 +685,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( @@ -786,9 +723,13 @@ 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 + 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) @@ -809,7 +750,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) @@ -819,10 +760,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 @@ -876,6 +814,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) @@ -894,10 +836,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) @@ -928,14 +869,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 @@ -947,39 +892,43 @@ 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 = 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 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 # 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: # ============================================================== @@ -995,19 +944,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 @@ -1020,7 +956,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) @@ -1037,16 +973,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) @@ -1057,28 +993,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: @@ -1093,6 +1029,7 @@ def launch_layernorm_smoothquant( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): + 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), @@ -1114,6 +1051,7 @@ def launch_layernorm_dynamicquant( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): + 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), @@ -1130,12 +1068,16 @@ 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 + 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) @@ -1158,7 +1100,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) @@ -1170,8 +1112,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 @@ -1225,6 +1165,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) @@ -1243,10 +1187,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) @@ -1283,17 +1226,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 @@ -1305,39 +1253,43 @@ 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 = 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 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 # 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: # ============================================================== @@ -1357,21 +1309,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 @@ -1384,13 +1321,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) @@ -1401,22 +1338,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) @@ -1427,28 +1363,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: @@ -1465,6 +1401,9 @@ def launch_fused_add_layernorm_smoothquant( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): + validate_norm_operand_dtypes( + 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 ) @@ -1490,6 +1429,9 @@ def launch_fused_add_layernorm_dynamicquant( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): + 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 ) @@ -1506,12 +1448,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 +1463,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 +1478,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 +1493,14 @@ def build_fused_add_layernorm_smoothquant_module( N: int, dtype_str: str, quant_dtype_str: str = "i8", + eps: float = EPS, ): return _build_fused_add_layernorm_quant_module( N, dtype_str, is_smooth=True, quant_dtype_str=quant_dtype_str, + eps=eps, ) diff --git a/kernels/norm/rmsnorm_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_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 72fd273439..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,19 +52,27 @@ 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 validate_norm_operand_dtypes(input_tensor, dtype_str: str, **operands) -> None: + """Validate quantized-norm operands against the builder specialization. - -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) + 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()) + raise ValueError( + f"quantized norm kernels require {'/'.join(operands)} to match the input dtype " + f"{input_tensor.dtype}, got {got}" + ) def load_vec(copy_atom, vec_width, elem_dtype, div_tensor, idx): diff --git a/kernels/norm/rmsnorm_kernel.py b/kernels/norm/rmsnorm_kernel.py index 7c18d43312..0c36b01751 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 @@ -35,15 +35,15 @@ 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 +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: @@ -53,8 +53,29 @@ 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 +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): @@ -74,24 +95,199 @@ 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) - 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") + if BLOCK_THREADS is None: + BLOCK_THREADS = default_block_threads(N, arch) # 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 + 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]} - SharedStorage = _make_reduction_storage(RED_SLOTS) + # 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_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 + ) + + 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 + + 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) def rmsnorm_kernel( @@ -111,12 +307,9 @@ 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) - 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 @@ -127,45 +320,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 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) Gamma_buf = fx.rocdl.make_buffer_tensor(Gamma) @@ -177,46 +355,76 @@ def block_reduce_add2(val0, val1): 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) c_zero_f = fx.Float32(0.0) thread_sumsq = c_zero_f - thread_dummy = c_zero_f in_local = [] + gamma_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) + 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 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) + if const_expr(TAIL_ELEMS > 0): + if tid < TAIL_ELEMS: + tail_idx = tid + VEC_ELEMS + x_tail_e = row_in[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 rrms = fmath.rsqrt(ms_eps, fastmath=fm_fast) 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_tiles): + for tile_i in range_constexpr(NUM_VEC_ITERS): idx = tid + tile_i * BLOCK_THREADS + if idx < VEC_TILES: + 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) - 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(output_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 = 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) + row_out[tail_idx] = y_tail_e else: # Scalar fallback for arbitrary N. @@ -227,19 +435,6 @@ def block_reduce_add2(val0, val1): 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 @@ -247,7 +442,7 @@ 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) + 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) @@ -260,19 +455,30 @@ 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: - 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( + N, + dtype_str, + store_rstd, + eps, + weight_dtype_str, + one_row_kernel=rmsnorm_kernel, + one_row_block_threads=BLOCK_THREADS, + ) if store_rstd: @@ -333,16 +539,20 @@ 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) -def _build_rmsnorm_large_m_small_n_module( +def _build_rmsnorm_small_n_module( N: int, dtype_str: str, 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() @@ -351,6 +561,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( @@ -379,25 +596,10 @@ 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)) - 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,39 +611,102 @@ def group_reduce_add(x): c_zero_f = fx.Float32(0.0) thread_sumsq = c_zero_f - for base_idx_int in range_constexpr(0, BLOCK_N, THREADS_PER_ROW): - idx = lane + base_idx_int - is_valid = idx < N - idx_safe = is_valid.select(idx, 0) - x_e = _load_scalar(copy_atom_s, elem_dtype, row_div, idx_safe) - x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) - x2 = x * x - thread_sumsq = thread_sumsq + is_valid.select(x2, c_zero_f) - - sum_sq = group_reduce_add(thread_sumsq) - mean_sq = sum_sq / n_float - ms_eps = mean_sq + eps_c - rrms = fmath.rsqrt(ms_eps, fastmath=fm_fast) - - if const_expr(store_rstd): - if lane == 0: - _store_scalar(rstd_copy_atom, fx.Float32, rstd_div, row, rrms) + 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) + + 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 = row_in[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: + Rstd_buf[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 = 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) + row_out[tail_idx] = y_tail_e + else: - 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 = row_in[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: + 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 = 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) + row_out[idx] = y_e 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, @@ -449,33 +714,68 @@ 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() - return launch_rmsnorm_large_m_small_n + dispatch() + + 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, 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, + ) - return launch_rmsnorm_large_m_small_n + 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_small_n def build_fused_add_rmsnorm_module( @@ -489,10 +789,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) @@ -520,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 @@ -568,9 +868,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) @@ -597,19 +896,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 @@ -618,16 +920,17 @@ 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_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. @@ -642,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 @@ -664,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) @@ -682,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: @@ -743,11 +1031,14 @@ 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) 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) @@ -766,7 +1057,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) @@ -778,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 @@ -841,6 +1130,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) @@ -858,9 +1151,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) @@ -888,14 +1180,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 @@ -906,39 +1200,43 @@ 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 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 # 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. @@ -948,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 @@ -979,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) @@ -995,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) @@ -1012,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 @@ -1020,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: @@ -1045,6 +1327,7 @@ def launch_rmsnorm_smoothquant( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): + _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), @@ -1065,6 +1348,7 @@ def launch_rmsnorm_dynamicquant( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): + _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), @@ -1079,12 +1363,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 +1378,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,14 +1395,17 @@ 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") - 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) @@ -1135,7 +1426,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) @@ -1147,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 @@ -1210,6 +1499,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) @@ -1227,9 +1520,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) @@ -1263,17 +1555,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 @@ -1284,38 +1579,42 @@ 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 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 # 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. @@ -1327,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 @@ -1363,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) @@ -1385,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) @@ -1402,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 @@ -1410,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: @@ -1437,6 +1717,9 @@ def launch_fused_add_rmsnorm_smoothquant( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): + _validate_norm_operand_dtypes( + 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( grid=(m_in, 1, 1), @@ -1459,6 +1742,7 @@ def launch_fused_add_rmsnorm_dynamicquant( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): + _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), @@ -1473,12 +1757,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 +1772,14 @@ def build_fused_add_rmsnorm_smoothquant_module( N: int, dtype_str: str, quant_dtype_str: str = "i8", + eps: float = EPS, ): return _build_fused_add_rmsnorm_quant_module( N, dtype_str, is_smooth=True, quant_dtype_str=quant_dtype_str, + eps=eps, ) diff --git a/tests/kernels/test_layernorm.py b/tests/kernels/test_layernorm.py index a2d1f73041..ca0d7b20fc 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 @@ -100,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 @@ -111,6 +109,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})") @@ -209,21 +212,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 @@ -247,6 +252,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...") @@ -316,11 +322,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 @@ -347,7 +353,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() @@ -413,15 +425,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 " @@ -432,6 +444,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 @@ -463,6 +478,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...") @@ -581,17 +597,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) @@ -601,14 +617,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( @@ -616,6 +640,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 @@ -847,7 +872,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 @@ -878,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") @@ -894,7 +915,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 @@ -937,7 +958,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 @@ -968,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") @@ -984,7 +1007,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 @@ -1432,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 @@ -1456,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 b062be0cc0..ee98795f25 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 @@ -130,17 +127,38 @@ 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 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 = ( + (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() + 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 [ - (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 - ] + return get_default_rmsnorm_forward_configs() def _get_rmsnorm_backward_configs(): @@ -161,6 +179,24 @@ 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 + 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 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})") @@ -199,6 +235,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) @@ -243,21 +280,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 +357,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 +483,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 +502,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 +613,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 +649,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 +680,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 +688,7 @@ def _reference_fused_add_rmsnorm_quant( added, gamma_dev, xscale_dev=xscale_dev, + eps=eps, ) return residual_expected, q, yscale @@ -1219,6 +1267,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).""" @@ -1368,7 +1450,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 = [] @@ -1523,6 +1605,38 @@ 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") + 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) + + 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()) @@ -1601,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")) @@ -1853,10 +1969,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") @@ -1869,7 +1981,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 @@ -1901,10 +2013,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") @@ -1917,7 +2025,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 @@ -2004,7 +2112,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 @@ -2034,10 +2148,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") @@ -2050,7 +2160,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 @@ -2085,6 +2201,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() @@ -2093,6 +2210,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": 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)