diff --git a/docs/prebuilt_kernels_guide.md b/docs/prebuilt_kernels_guide.md index 210d4d454..184d06033 100644 --- a/docs/prebuilt_kernels_guide.md +++ b/docs/prebuilt_kernels_guide.md @@ -7,7 +7,7 @@ | Kernel | Builder Function | API Style | Dtypes | Key Feature | |---|---|---|---|---| | **LayerNorm** | `build_layernorm_module(N, dtype)` | Layout API (`@flyc.kernel`) | f32, f16, bf16 | Two-pass vectorized normalization | -| **RMSNorm** | `build_rmsnorm_module(N, dtype)` | Layout API (`@flyc.kernel`) | f32, f16, bf16 | LDS-cached 3-pass pipeline | +| **RMSNorm** | `build_rmsnorm_module(N, dtype)` | Layout API (`@flyc.kernel`) | f32, f16, bf16; optional fp32 weight | LDS-cached 3-pass pipeline | | **Softmax** | `build_softmax_module(M, N, dtype)` | Layout API (`@flyc.kernel`) | f32, f16, bf16 | Online softmax, adaptive block size | | **GEMM** | `compile_preshuffle_gemm(...)` | `@flyc.kernel` | fp8, int8, fp16, bf16 | Preshuffle B, ping-pong LDS, MFMA 16x16 | | **FlashAttention** | `build_flash_attn_func_module(...)` | `@flyc.kernel` | bf16, f16 (any arch); fp8 e4m3fn (gfx950, D=128, dense) | Dual-wave SWP fwd, GQA/MQA, causal, descale ABI | @@ -68,14 +68,20 @@ from kernels.norm.rmsnorm_kernel import build_rmsnorm_module executor = build_rmsnorm_module(N=8192, dtype_str="bf16", store_rstd=False) ``` -`build_rmsnorm_module(N, dtype_str, store_rstd=False, eps=EPS)` optionally -writes the per-row reciprocal std (`rstd`) for use by the backward pass. +`build_rmsnorm_module(N, dtype_str, store_rstd=False, eps=EPS, +BLOCK_THREADS=BLOCK_THREADS, 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. -**Backward:** `build_rmsnorm_bwd_module(N, dtype_str)` builds the fused RMSNorm -backward kernel (grid `(M,)`, one block per row). Kernel signature +**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 `rmsnorm_bwd_kernel(Input, Gamma, DY, Rstd, DX, DWeight)`: it reads the forward `Rstd`, writes `DX` (input grad) and atomic-adds into `DWeight` (fp32 weight grad). `eps` is baked into `Rstd` by the forward, so it is not needed here. +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.) diff --git a/kernels/norm/rmsnorm_autotune.py b/kernels/norm/rmsnorm_autotune.py index 7408b9dbc..f243ea160 100644 --- a/kernels/norm/rmsnorm_autotune.py +++ b/kernels/norm/rmsnorm_autotune.py @@ -4,7 +4,10 @@ """Two-track RMSNorm autotuning through the normal direct JIT path.""" from flydsl.autotune import Config, autotune -from kernels.norm.rmsnorm_common import BLOCK_THREADS +from kernels.norm.rmsnorm_common import ( + BLOCK_THREADS, + resolve_rmsnorm_weight_dtype, +) from kernels.norm.rmsnorm_kernel import SMALL_N_THRESHOLD, rmsnorm_direct _SEARCH_CONFIGS = ( @@ -22,7 +25,16 @@ def _default_config(*_args, **_kwargs): return Config(BLOCK_THREADS=BLOCK_THREADS) -def _search_configs(input_t, gamma, output, m_in, N, dtype_str="bf16", stream=None): +def _search_configs( + input_t, + gamma, + output, + m_in, + N, + dtype_str="bf16", + weight_dtype_str="bf16", + stream=None, +): if N <= SMALL_N_THRESHOLD: return [_default_config()] return list(_SEARCH_CONFIGS) @@ -30,14 +42,34 @@ def _search_configs(input_t, gamma, output, m_in, N, dtype_str="bf16", stream=No _rmsnorm_tuner = autotune( configs=_search_configs, - key=["m_in", "N", "dtype_str"], + key=["m_in", "N", "dtype_str", "weight_dtype_str"], default=_default_config, )(rmsnorm_direct) -def rmsnorm_autotuned(input_t, gamma, output, m_in, dtype_str="bf16", stream=None): +def rmsnorm_autotuned( + input_t, + gamma, + output, + m_in, + dtype_str=None, + stream=None, + weight_dtype_str=None, +): import torch + from kernels.norm.rmsnorm_common import torch_dtype_to_str + + input_dtype_str = torch_dtype_to_str(input_t.dtype) + if dtype_str is not None and dtype_str != input_dtype_str: + raise ValueError(f"dtype_str={dtype_str!r} does not match input dtype {input_dtype_str!r}") + dtype_str = input_dtype_str + + gamma_dtype_str = torch_dtype_to_str(gamma.dtype) + if weight_dtype_str is not None and weight_dtype_str != gamma_dtype_str: + raise ValueError(f"weight_dtype_str={weight_dtype_str!r} does not match gamma dtype {gamma_dtype_str!r}") + weight_dtype_str = resolve_rmsnorm_weight_dtype(dtype_str, gamma_dtype_str) + with torch.cuda.device(input_t.device): launch_stream = torch.cuda.current_stream() if stream is None else stream return _rmsnorm_tuner( @@ -47,5 +79,6 @@ def rmsnorm_autotuned(input_t, gamma, output, m_in, dtype_str="bf16", stream=Non m_in, N=int(input_t.shape[-1]), dtype_str=dtype_str, + weight_dtype_str=weight_dtype_str, stream=launch_stream, ) diff --git a/kernels/norm/rmsnorm_bwd_kernel.py b/kernels/norm/rmsnorm_bwd_kernel.py index f42a6090b..c0f0469b2 100644 --- a/kernels/norm/rmsnorm_bwd_kernel.py +++ b/kernels/norm/rmsnorm_bwd_kernel.py @@ -23,10 +23,13 @@ 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, ) # The staged path is selected by the Python wrapper only when M is large @@ -44,7 +47,7 @@ def is_rmsnorm_bwd_two_stage_vec_config(N: int, dtype_str: str) -> bool: return dtype_str in ("f16", "bf16") and N >= TWO_STAGE_PARTIAL_THREADS * VEC_WIDTH and N % VEC_WIDTH == 0 -def build_rmsnorm_bwd_module(N: int, dtype_str: str): +def build_rmsnorm_bwd_module(N: int, dtype_str: str, weight_dtype_str: str | None = None): """Fused RMSNorm backward: grid=(M,), one block per row. Pass 1: c1 = mean_N(x_hat * wdy), x_hat = x*rstd, wdy = dy*gamma. @@ -57,8 +60,10 @@ def build_rmsnorm_bwd_module(N: int, dtype_str: str): x/dy/gamma between pass 1 and pass 2 (the forward caches `in_local`) would cut global traffic. Left out of PR 1 to keep the first backward reviewable. """ + 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 @@ -74,6 +79,7 @@ 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) @@ -123,6 +129,10 @@ def block_reduce_add(val): 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)) @@ -141,10 +151,10 @@ def block_reduce_add(val): 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) + g_e = load_scalar(gamma_copy_atom_s, weight_elem_dtype, gamma_div, 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) + g = g_e if weight_dtype_str == "f32" else g_e.to(fx.Float32) x_hat = x * rstd wdy = dy * g prod = x_hat * wdy @@ -159,10 +169,10 @@ def block_reduce_add(val): 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) + g_e = load_scalar(gamma_copy_atom_s, weight_elem_dtype, gamma_div, idx) x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) 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) + g = g_e if weight_dtype_str == "f32" else g_e.to(fx.Float32) x_hat = x * rstd wdy = dy * g dx = (wdy - x_hat * c1) * rstd @@ -189,7 +199,7 @@ def launch_rmsnorm_bwd( return launch_rmsnorm_bwd -def build_fused_add_rmsnorm_bwd_module(N: int, dtype_str: str): +def build_fused_add_rmsnorm_bwd_module(N: int, dtype_str: str, weight_dtype_str: str | None = None): """Fused-add / prenorm RMSNorm backward: grid=(M,), one block per row. Inputs: Added (= residual_out from fwd), Gamma, DY (= dout), DResidualOut @@ -211,8 +221,10 @@ def build_fused_add_rmsnorm_bwd_module(N: int, dtype_str: str): Perf follow-ups (deferred; correctness-complete): generic scalar path only — a vectorized fast path + caching between passes would cut global traffic. """ + 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 @@ -229,6 +241,7 @@ 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) @@ -280,6 +293,10 @@ def block_reduce_add(val): 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)) @@ -299,10 +316,10 @@ def block_reduce_add(val): 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(copy_atom_s, elem_dtype, gamma_div, idx_safe) + g_e = load_scalar(gamma_copy_atom_s, weight_elem_dtype, gamma_div, 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 dtype_str == "f32" else g_e.to(fx.Float32) + g = g_e if weight_dtype_str == "f32" else g_e.to(fx.Float32) a_hat = a * rstd wdy = dy * g prod = a_hat * wdy @@ -319,11 +336,11 @@ def block_reduce_add(val): 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(copy_atom_s, elem_dtype, gamma_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 = 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 dtype_str == "f32" else g_e.to(fx.Float32) + g = g_e if weight_dtype_str == "f32" else g_e.to(fx.Float32) dres_out = dres_out_e if dtype_str == "f32" else dres_out_e.to(fx.Float32) a_hat = a * rstd wdy = dy * g @@ -359,6 +376,7 @@ def _build_rmsnorm_bwd_two_stage_module( num_programs: int, *, fused_add: bool, + weight_dtype_str: str | None = None, ): """Build the large-M persistent backward + dweight finalizer. @@ -377,10 +395,13 @@ def _build_rmsnorm_bwd_two_stage_module( if num_programs <= 0: raise ValueError(f"num_programs must be positive, got {num_programs}") + weight_dtype_str = resolve_rmsnorm_weight_dtype(dtype_str, weight_dtype_str) RED_SLOTS = max(1, (TWO_STAGE_PARTIAL_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 = is_rmsnorm_bwd_two_stage_vec_config(N, dtype_str) IO_WIDTH = VEC_WIDTH if USE_VEC else 1 + WEIGHT_IO_WIDTH = weight_vec_width(weight_dtype_str) if USE_VEC else 1 NUM_IO_TILES = (N + IO_WIDTH - 1) // IO_WIDTH NUM_IO_ITERS = (NUM_IO_TILES + TWO_STAGE_PARTIAL_THREADS - 1) // TWO_STAGE_PARTIAL_THREADS PARTIAL_ACC_SIZE = NUM_IO_ITERS * IO_WIDTH @@ -404,6 +425,7 @@ def rmsnorm_bwd_partial_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) @@ -458,16 +480,32 @@ def block_reduce_add(val): ), elem_bits, ) + gamma_copy_atom = fx.make_copy_atom( + ( + fx.rocdl.BufferCopy128b() + if USE_VEC + else (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) - gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(IO_WIDTH, 1)) + gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(WEIGHT_IO_WIDTH, 1)) gamma_local = [] if const_expr(USE_VEC): for tile_i in range_constexpr(NUM_IO_ITERS): io_idx = tid + tile_i * TWO_STAGE_PARTIAL_THREADS is_valid = io_idx < NUM_IO_TILES io_idx_safe = is_valid.select(io_idx, 0) - gamma_local.append(load_vec(copy_atom_io, IO_WIDTH, elem_dtype, gamma_div, io_idx_safe)) + gamma_local.append( + load_weight_vec( + gamma_copy_atom, + weight_dtype_str, + weight_elem_dtype, + gamma_div, + io_idx_safe, + ) + ) dweight_partial = fx.Vector.filled(PARTIAL_ACC_SIZE, 0.0, fx.Float32) for row in range(fx.Int32(bid), MIn, num_programs): @@ -500,11 +538,11 @@ def block_reduce_add(val): 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(copy_atom_io, elem_dtype, gamma_div, io_idx_safe) + gamma_e = load_scalar(gamma_copy_atom, weight_elem_dtype, gamma_div, 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) - gamma = gamma_e if dtype_str == "f32" else gamma_e.to(fx.Float32) + gamma = gamma_e if USE_VEC or weight_dtype_str == "f32" else gamma_e.to(fx.Float32) source_hat = source * rstd wdy = dy * gamma prod = source_hat * wdy @@ -527,11 +565,11 @@ def block_reduce_add(val): 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(copy_atom_io, elem_dtype, gamma_div, io_idx_safe) + gamma_e = load_scalar(gamma_copy_atom, weight_elem_dtype, gamma_div, 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) - gamma = gamma_e if dtype_str == "f32" else gamma_e.to(fx.Float32) + gamma = gamma_e if USE_VEC or weight_dtype_str == "f32" else gamma_e.to(fx.Float32) source_hat = source * rstd wdy = dy * gamma @@ -591,15 +629,15 @@ def rmsnorm_bwd_dweight_reduce_kernel( is_valid = col < N col_safe = is_valid.select(col, 0) - elem_dtype = dtype_to_elem_type(dtype_str) + 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) - copy_atom_s = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - elem_bits, + 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() @@ -621,8 +659,15 @@ def rmsnorm_bwd_dweight_reduce_kernel( total = fx.Float32(0.0) 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 dtype_str == "f32" else total.to(elem_dtype) - store_scalar(copy_atom_s, elem_dtype, elem_dtype, dweight_div, col, out) + out = total if weight_dtype_str == "f32" else total.to(weight_elem_dtype) + store_scalar( + weight_copy_atom_s, + weight_elem_dtype, + weight_elem_dtype, + dweight_div, + col, + out, + ) reduce_grid = (N + DWEIGHT_REDUCE_COLS - 1) // DWEIGHT_REDUCE_COLS @@ -704,11 +749,33 @@ def launch_rmsnorm_bwd_two_stage( return launch_rmsnorm_bwd_two_stage -def build_rmsnorm_bwd_two_stage_module(N: int, dtype_str: str, num_programs: int): +def build_rmsnorm_bwd_two_stage_module( + N: int, + dtype_str: str, + num_programs: int, + weight_dtype_str: str | None = None, +): """Large-M plain RMSNorm backward without global dweight atomics.""" - return _build_rmsnorm_bwd_two_stage_module(N, dtype_str, num_programs, fused_add=False) + return _build_rmsnorm_bwd_two_stage_module( + N, + dtype_str, + num_programs, + fused_add=False, + weight_dtype_str=weight_dtype_str, + ) -def build_fused_add_rmsnorm_bwd_two_stage_module(N: int, dtype_str: str, num_programs: int): +def build_fused_add_rmsnorm_bwd_two_stage_module( + N: int, + dtype_str: str, + num_programs: int, + weight_dtype_str: str | None = None, +): """Large-M fused-add RMSNorm backward without global dweight atomics.""" - return _build_rmsnorm_bwd_two_stage_module(N, dtype_str, num_programs, fused_add=True) + return _build_rmsnorm_bwd_two_stage_module( + N, + dtype_str, + num_programs, + fused_add=True, + weight_dtype_str=weight_dtype_str, + ) diff --git a/kernels/norm/rmsnorm_common.py b/kernels/norm/rmsnorm_common.py index 42fc63326..2c21fa797 100644 --- a/kernels/norm/rmsnorm_common.py +++ b/kernels/norm/rmsnorm_common.py @@ -74,6 +74,18 @@ def load_vec(copy_atom, vec_width, elem_dtype, div_tensor, idx): return fx.memref_load_vec(r) +def load_weight_vec(copy_atom, weight_dtype_str, weight_elem_dtype, div_tensor, idx): + """Load eight weights as fp32 using one 128-bit load per source vector.""" + if const_expr(weight_dtype_str == "f32"): + lo = load_vec(copy_atom, VEC_WIDTH // 2, weight_elem_dtype, div_tensor, idx * 2) + hi = load_vec(copy_atom, VEC_WIDTH // 2, weight_elem_dtype, div_tensor, idx * 2 + 1) + return fx.Vector.from_elements( + [lo[0], lo[1], lo[2], lo[3], hi[0], hi[1], hi[2], hi[3]], + fx.Float32, + ) + return load_vec(copy_atom, VEC_WIDTH, weight_elem_dtype, div_tensor, idx).to(fx.Float32) + + def store_vec(copy_atom, vec_width, elem_dtype, val, div_tensor, idx): r = fx.make_rmem_tensor(vec_width, elem_dtype) fx.memref_store_vec(val, r) @@ -106,6 +118,25 @@ def to_elem_vec(dtype_str: str, elem_dtype, use_hw_cvt_bf16: bool, y): return y.to(elem_dtype) +def resolve_rmsnorm_weight_dtype(dtype_str: str, weight_dtype_str: str | None = None) -> str: + """Return the supported weight dtype for an activation specialization.""" + weight_dtype_str = dtype_str if weight_dtype_str is None else weight_dtype_str + supported = dtype_str in ("f16", "bf16", "f32") and ( + weight_dtype_str == dtype_str or (dtype_str in ("f16", "bf16") and weight_dtype_str == "f32") + ) + if not supported: + raise ValueError( + "RMSNorm supports matching activation/weight dtypes or " + f"FP16/BF16 activations with FP32 weights, got {dtype_str}/{weight_dtype_str}" + ) + return weight_dtype_str + + +def weight_vec_width(weight_dtype_str: str) -> int: + """Number of weight elements transferred by one 128-bit copy.""" + return VEC_WIDTH // 2 if weight_dtype_str == "f32" else VEC_WIDTH + + if torch is not None: def torch_dtype_to_str(dt) -> str: diff --git a/kernels/norm/rmsnorm_kernel.py b/kernels/norm/rmsnorm_kernel.py index 6843615d9..07fbd40c6 100644 --- a/kernels/norm/rmsnorm_kernel.py +++ b/kernels/norm/rmsnorm_kernel.py @@ -37,11 +37,14 @@ ) 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 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 weight_vec_width as _weight_vec_width try: import torch @@ -74,10 +77,16 @@ def _quant_dtype_max(dtype_str: str) -> float: def build_rmsnorm_module( - N: int, dtype_str: str, store_rstd: bool = False, eps: float = EPS, BLOCK_THREADS: int = BLOCK_THREADS + N: int, + dtype_str: str, + store_rstd: bool = False, + eps: float = EPS, + BLOCK_THREADS: int = BLOCK_THREADS, + weight_dtype_str: str | None = None, ): + weight_dtype_str = _resolve_rmsnorm_weight_dtype(dtype_str, weight_dtype_str) if N <= SMALL_N_THRESHOLD: - return _build_rmsnorm_large_m_small_n_module(N, dtype_str, store_rstd, eps) + 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") @@ -86,6 +95,7 @@ def build_rmsnorm_module( 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 _kernel_kwargs = {} if BLOCK_THREADS <= 256 else {"known_block_size": [BLOCK_THREADS, 1, 1]} SharedStorage = _make_reduction_storage(RED_SLOTS) @@ -101,6 +111,7 @@ def rmsnorm_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 eps_c = eps n_float = float(N) @@ -174,9 +185,10 @@ def block_reduce_add2(val0, val1): in_div = fx.logical_divide(row_in, fx.make_layout(VEC_WIDTH, 1)) out_div = fx.logical_divide(row_out, fx.make_layout(VEC_WIDTH, 1)) - gamma_div = fx.logical_divide(Gamma_buf, 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) + 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 @@ -207,7 +219,7 @@ def block_reduce_add2(val0, val1): for tile_i in range_constexpr(num_tiles): idx = tid + tile_i * BLOCK_THREADS - g = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, idx).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 @@ -231,6 +243,10 @@ def block_reduce_add2(val0, val1): 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)) @@ -262,9 +278,9 @@ def block_reduce_add2(val0, val1): 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) + g_e = _load_scalar(gamma_copy_atom_s, weight_elem_dtype, gamma_div, idx) x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) - g = g_e if dtype_str == "f32" else g_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) @@ -320,18 +336,33 @@ def rmsnorm_direct( dtype_str: fx.Constexpr[str], BLOCK_THREADS: fx.Constexpr[int], stream: fx.Stream = fx.Stream(None), + weight_dtype_str: fx.Constexpr[str] = "", ): """Specialize the existing RMSNorm factory through JIT Constexpr inputs.""" - launch = build_rmsnorm_module(N, dtype_str, BLOCK_THREADS=BLOCK_THREADS) + resolved_weight_dtype_str = dtype_str if weight_dtype_str == "" else weight_dtype_str + launch = build_rmsnorm_module( + N, + dtype_str, + BLOCK_THREADS=BLOCK_THREADS, + weight_dtype_str=resolved_weight_dtype_str, + ) launch(Input, Gamma, Output, m_in, stream) -def _build_rmsnorm_large_m_small_n_module(N: int, dtype_str: str, store_rstd: bool = False, eps: float = EPS): +def _build_rmsnorm_large_m_small_n_module( + N: int, + dtype_str: str, + store_rstd: bool = False, + eps: float = EPS, + weight_dtype_str: str | None = None, +): + weight_dtype_str = _resolve_rmsnorm_weight_dtype(dtype_str, weight_dtype_str) BLOCK_N = 1 << (N - 1).bit_length() BLOCK_M = max(min(16384 // BLOCK_N, 32), 8) THREADS_PER_ROW = min(WARP_SIZE, 1024 // BLOCK_M) BLOCK_THREADS_SPECIAL = BLOCK_M * THREADS_PER_ROW elem_bits = 32 if dtype_str == "f32" else 16 + weight_elem_bits = 32 if weight_dtype_str == "f32" else 16 @flyc.kernel(known_block_size=[BLOCK_THREADS_SPECIAL, 1, 1]) def rmsnorm_large_m_small_n_kernel( @@ -350,6 +381,7 @@ def rmsnorm_large_m_small_n_kernel( if row < MIn: elem_dtype = dtype_to_elem_type(dtype_str) + weight_elem_dtype = dtype_to_elem_type(weight_dtype_str) fm_fast = arith.FastMathFlags.fast eps_c = eps n_float = float(N) @@ -369,6 +401,10 @@ def rmsnorm_large_m_small_n_kernel( 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)) @@ -407,9 +443,9 @@ def group_reduce_add(x): idx = lane + base_idx_int if idx < N: x_e = _load_scalar(copy_atom_s, elem_dtype, row_div, idx) - g_e = _load_scalar(copy_atom_s, elem_dtype, gamma_div, idx) + g_e = _load_scalar(gamma_copy_atom_s, weight_elem_dtype, gamma_div, idx) x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) - g = g_e if dtype_str == "f32" else g_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, elem_dtype, out_div, idx, y_e) @@ -454,13 +490,21 @@ def launch_rmsnorm_large_m_small_n( return launch_rmsnorm_large_m_small_n -def build_fused_add_rmsnorm_module(N: int, dtype_str: str, store_rstd: bool = False, eps: float = EPS): +def build_fused_add_rmsnorm_module( + N: int, + dtype_str: str, + store_rstd: bool = False, + eps: float = EPS, + weight_dtype_str: str | None = None, +): + weight_dtype_str = _resolve_rmsnorm_weight_dtype(dtype_str, weight_dtype_str) arch = get_rocm_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") 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 SharedStorage = _make_reduction_storage(RED_SLOTS) @@ -477,6 +521,7 @@ def fused_add_rmsnorm_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 eps_c = eps n_float = float(N) @@ -556,9 +601,10 @@ def block_reduce_add2(val0, val1): residual_in_div = fx.logical_divide(row_residual_in, fx.make_layout(VEC_WIDTH, 1)) out_div = fx.logical_divide(row_out, fx.make_layout(VEC_WIDTH, 1)) residual_out_div = fx.logical_divide(row_residual_out, fx.make_layout(VEC_WIDTH, 1)) - gamma_div = fx.logical_divide(Gamma_buf, 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) + 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 @@ -592,7 +638,7 @@ def block_reduce_add2(val0, val1): # Pass 2: normalize + gamma + store (reuse cached added values) for tile_i in range_constexpr(num_tiles): idx = tid + tile_i * BLOCK_THREADS - g = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, idx).to(fx.Float32) + 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) @@ -617,6 +663,10 @@ def block_reduce_add2(val0, val1): 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)) @@ -654,9 +704,9 @@ def block_reduce_add2(val0, val1): 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) + 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 = g_e if dtype_str == "f32" else g_e.to(fx.Float32) + 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) @@ -1529,11 +1579,29 @@ def _select_rmsnorm_bwd_config(M, N, dtype_str, device): return ("two_stage", min(M, num_programs)) return ("atomic", None) - def _get_fwd_compiled(x, weight, out, rstd, M, N, dtype_str, store_rstd, eps, stream): - key = (N, dtype_str, store_rstd, float(eps), x.device) + def _get_fwd_compiled( + x, + weight, + out, + rstd, + M, + N, + dtype_str, + weight_dtype_str, + store_rstd, + eps, + stream, + ): + key = (N, dtype_str, weight_dtype_str, store_rstd, float(eps), x.device) entry = _FWD_CACHE.get(key) if entry is None: - launch_fn = build_rmsnorm_module(N, dtype_str, store_rstd=store_rstd, eps=eps) + launch_fn = build_rmsnorm_module( + N, + dtype_str, + store_rstd=store_rstd, + eps=eps, + weight_dtype_str=weight_dtype_str, + ) if store_rstd: compiled = flyc.compile(launch_fn, x, weight, out, rstd, M, stream) else: @@ -1546,19 +1614,29 @@ def rmsnorm_fwd(x, weight, eps=EPS, store_rstd=False): """Forward RMSNorm. Returns (out, rstd). eps is baked into the kernel.""" assert x.dim() == 2, "rmsnorm_fwd expects a 2D (M, N) input" assert x.is_contiguous() and weight.is_contiguous(), "rmsnorm_fwd expects contiguous inputs" - # gamma is read at x's element width and the kernel launches on x's device, - # so a mismatched weight would silently corrupt output (or read out of bounds). assert weight.device == x.device, "rmsnorm_fwd: weight and x must be on the same device" - assert weight.dtype == x.dtype, "rmsnorm_fwd: weight dtype must match x dtype" M, N = x.shape out = torch.empty_like(x) rstd = torch.empty((M,), device=x.device, dtype=torch.float32) if store_rstd else None dtype_str = _torch_dtype_to_str(x.dtype) + weight_dtype_str = _resolve_rmsnorm_weight_dtype(dtype_str, _torch_dtype_to_str(weight.dtype)) # Bind compile + launch to the tensors' device so the compiled kernel and # the stream belong to the right GPU/context (multi-GPU correctness). with torch.cuda.device(x.device): stream = torch.cuda.current_stream() - compiled = _get_fwd_compiled(x, weight, out, rstd, M, N, dtype_str, store_rstd, eps, stream) + compiled = _get_fwd_compiled( + x, + weight, + out, + rstd, + M, + N, + dtype_str, + weight_dtype_str, + store_rstd, + eps, + stream, + ) if store_rstd: compiled(x, weight, out, rstd, M, stream) else: @@ -1566,7 +1644,7 @@ def rmsnorm_fwd(x, weight, eps=EPS, store_rstd=False): return out, rstd def rmsnorm_bwd(x, weight, dout, rstd, eps=EPS): - """Backward RMSNorm. Returns (dx, dw) with dw cast to weight dtype. + """Backward RMSNorm. Returns (dx, dw) with dw in weight dtype. eps is not used directly here — it is already baked into `rstd` by the forward — but is accepted so callers can pass it symmetrically. @@ -1576,25 +1654,29 @@ def rmsnorm_bwd(x, weight, dout, rstd, eps=EPS): assert x.dim() == 2, "rmsnorm_bwd expects a 2D (M, N) input" assert x.is_contiguous() and dout.is_contiguous(), "rmsnorm_bwd expects contiguous inputs" assert weight.is_contiguous(), "rmsnorm_bwd: weight must be contiguous" - # Same-device/same-dtype contract as the forward: gamma and dy are read at - # x's element width and the kernel launches on x's device (multi-GPU correctness). assert weight.device == device == dout.device == rstd.device, "rmsnorm_bwd: inputs must share a device" - assert weight.dtype == dtype and dout.dtype == dtype, "rmsnorm_bwd: weight/dout dtype must match x" + assert dout.dtype == dtype, "rmsnorm_bwd: dout dtype must match x" M, N = x.shape dtype_str = _torch_dtype_to_str(dtype) + weight_dtype_str = _resolve_rmsnorm_weight_dtype(dtype_str, _torch_dtype_to_str(weight.dtype)) dx = torch.empty_like(x) path, num_programs = _select_rmsnorm_bwd_config(M, N, dtype_str, device) if path == "two_stage": dweight = torch.empty_like(weight) partial = torch.empty((num_programs * N,), device=device, dtype=torch.float32) - key = (path, N, dtype_str, num_programs, device) + key = (path, N, dtype_str, weight_dtype_str, num_programs, device) launcher = _BWD_CACHE.get(key) if launcher is None: # The builder has an arch-dependent vec conversion choice, and # compilation/code-object loading are device-context bound. with torch.cuda.device(device): device_index = device.index - launcher = build_rmsnorm_bwd_two_stage_module(N, dtype_str, num_programs) + launcher = build_rmsnorm_bwd_two_stage_module( + N, + dtype_str, + num_programs, + weight_dtype_str=weight_dtype_str, + ) _run_compiled( launcher, x, @@ -1613,13 +1695,13 @@ def rmsnorm_bwd(x, weight, dout, rstd, eps=EPS): return dx, dweight dweight = torch.empty((N,), device=device, dtype=torch.float32) - key = (path, N, dtype_str, num_programs, device) + key = (path, N, dtype_str, weight_dtype_str, num_programs, device) launcher = _BWD_CACHE.get(key) dweight.zero_() if launcher is None: with torch.cuda.device(device): device_index = device.index - launcher = build_rmsnorm_bwd_module(N, dtype_str) + launcher = build_rmsnorm_bwd_module(N, dtype_str, weight_dtype_str=weight_dtype_str) _run_compiled( launcher, x, @@ -1670,12 +1752,30 @@ def rmsnorm(x, weight=None, eps=EPS): _FUSED_ADD_BWD_CACHE: dict = {} def _get_fused_add_fwd_compiled( - x, residual, weight, out, residual_out, rstd, M, N, dtype_str, store_rstd, eps, stream + x, + residual, + weight, + out, + residual_out, + rstd, + M, + N, + dtype_str, + weight_dtype_str, + store_rstd, + eps, + stream, ): - key = (N, dtype_str, store_rstd, float(eps), x.device) + key = (N, dtype_str, weight_dtype_str, store_rstd, float(eps), x.device) entry = _FUSED_ADD_FWD_CACHE.get(key) if entry is None: - launch_fn = build_fused_add_rmsnorm_module(N, dtype_str, store_rstd=store_rstd, eps=eps) + launch_fn = build_fused_add_rmsnorm_module( + N, + dtype_str, + store_rstd=store_rstd, + eps=eps, + weight_dtype_str=weight_dtype_str, + ) if store_rstd: compiled = flyc.compile(launch_fn, x, residual, weight, out, residual_out, rstd, M, stream) else: @@ -1695,11 +1795,7 @@ def fused_add_rmsnorm_fwd(x, residual, weight, eps=EPS, store_rstd=False): x.is_contiguous() and residual.is_contiguous() and weight.is_contiguous() ), "fused_add_rmsnorm_fwd expects contiguous inputs" assert x.shape == residual.shape, "x and residual must have the same shape" - # The kernel reads x/residual/weight with a single elem dtype derived from x; - # a mismatch would silently bit-reinterpret residual/weight bytes. - assert ( - x.dtype == residual.dtype == weight.dtype - ), f"x/residual/weight dtypes must match, got {x.dtype}/{residual.dtype}/{weight.dtype}" + assert x.dtype == residual.dtype, f"x/residual dtypes must match, got {x.dtype}/{residual.dtype}" # Only x.device gates compile/stream/cache; all operands must co-reside. assert ( x.device == residual.device == weight.device @@ -1709,10 +1805,23 @@ def fused_add_rmsnorm_fwd(x, residual, weight, eps=EPS, store_rstd=False): residual_out = torch.empty_like(x) rstd = torch.empty((M,), device=x.device, dtype=torch.float32) if store_rstd else None dtype_str = _torch_dtype_to_str(x.dtype) + weight_dtype_str = _resolve_rmsnorm_weight_dtype(dtype_str, _torch_dtype_to_str(weight.dtype)) with torch.cuda.device(x.device): stream = torch.cuda.current_stream() compiled = _get_fused_add_fwd_compiled( - x, residual, weight, out, residual_out, rstd, M, N, dtype_str, store_rstd, eps, stream + x, + residual, + weight, + out, + residual_out, + rstd, + M, + N, + dtype_str, + weight_dtype_str, + store_rstd, + eps, + stream, ) if store_rstd: compiled(x, residual, weight, out, residual_out, rstd, M, stream) @@ -1734,9 +1843,7 @@ def fused_add_rmsnorm_bwd(added, weight, dout, rstd, dresidual_out=None, eps=EPS dtype = added.dtype assert added.dim() == 2, "fused_add_rmsnorm_bwd expects a 2D (M, N) input" assert added.is_contiguous() and dout.is_contiguous(), "fused_add_rmsnorm_bwd expects contiguous inputs" - assert ( - dtype == weight.dtype == dout.dtype - ), f"added/weight/dout dtypes must match, got {dtype}/{weight.dtype}/{dout.dtype}" + assert dtype == dout.dtype, f"added/dout dtypes must match, got {dtype}/{dout.dtype}" assert ( device == weight.device == dout.device == rstd.device ), "fused_add_rmsnorm_bwd expects all tensors on the same device" @@ -1746,6 +1853,7 @@ def fused_add_rmsnorm_bwd(added, weight, dout, rstd, dresidual_out=None, eps=EPS assert dresidual_out.device == device, "dresidual_out must be on the same device as added" M, N = added.shape dtype_str = _torch_dtype_to_str(dtype) + weight_dtype_str = _resolve_rmsnorm_weight_dtype(dtype_str, _torch_dtype_to_str(weight.dtype)) if dresidual_out is None: dresidual_out = torch.zeros_like(added) dx = torch.empty_like(added) @@ -1753,12 +1861,17 @@ def fused_add_rmsnorm_bwd(added, weight, dout, rstd, dresidual_out=None, eps=EPS if path == "two_stage": dweight = torch.empty_like(weight) partial = torch.empty((num_programs * N,), device=device, dtype=torch.float32) - key = (path, N, dtype_str, num_programs, device) + key = (path, N, dtype_str, weight_dtype_str, num_programs, device) launcher = _FUSED_ADD_BWD_CACHE.get(key) if launcher is None: with torch.cuda.device(device): device_index = device.index - launcher = build_fused_add_rmsnorm_bwd_two_stage_module(N, dtype_str, num_programs) + launcher = build_fused_add_rmsnorm_bwd_two_stage_module( + N, + dtype_str, + num_programs, + weight_dtype_str=weight_dtype_str, + ) _run_compiled( launcher, added, @@ -1782,13 +1895,17 @@ def fused_add_rmsnorm_bwd(added, weight, dout, rstd, dresidual_out=None, eps=EPS return dx, dx, dweight dweight = torch.empty((N,), device=device, dtype=torch.float32) - key = (path, N, dtype_str, num_programs, device) + key = (path, N, dtype_str, weight_dtype_str, num_programs, device) launcher = _FUSED_ADD_BWD_CACHE.get(key) dweight.zero_() if launcher is None: with torch.cuda.device(device): device_index = device.index - launcher = build_fused_add_rmsnorm_bwd_module(N, dtype_str) + launcher = build_fused_add_rmsnorm_bwd_module( + N, + dtype_str, + weight_dtype_str=weight_dtype_str, + ) _run_compiled( launcher, added, diff --git a/scripts/run_benchmark.sh b/scripts/run_benchmark.sh index 9c4f6a0d3..5b967a9e1 100755 --- a/scripts/run_benchmark.sh +++ b/scripts/run_benchmark.sh @@ -70,6 +70,10 @@ LAYERNORM_SHAPES=' RMSNORM_SHAPES=' 32768,8192,bf16 ' +RMSNORM_MIXED_WEIGHT_SHAPES=' +4096,4096,bf16 +512,4096,f16 +' # FlashAttention shapes: # preferred: "batch,seq_len,num_heads,num_kv_heads,head_dim,dtype,causal" # legacy: "batch,seq_len,num_heads,head_dim,dtype,causal" (num_kv_heads=num_heads) @@ -620,6 +624,51 @@ if [ "${RUN_RMSNORM}" -eq 1 ]; then set -- $row _emit_row "$1" "$2" "$3" "$4" "$5" done + + # Training contract: FP16/BF16 activations with FP32 weights. These focused + # pytest entries avoid rerunning the full RMSNorm correctness matrix for each + # benchmark row while keeping gfx942/gfx950 dashboard coverage symmetric. + for shape in $RMSNORM_MIXED_WEIGHT_SHAPES; do + oldIFS=$IFS + IFS=, + # shellcheck disable=SC2086 # intentional word-splitting on IFS=, + set -- $shape + IFS=$oldIFS + M=$1; N=$2; activation_dtype=$3 + export ROCDSL_RMSNORM_MIXED_WEIGHT_BENCH_SHAPE="$shape" + + log="${BENCH_LOG_DIR}/rmsnorm_mixed_weight_${M}x${N}_${activation_dtype}.log" + if python3 -m pytest \ + tests/kernels/test_rmsnorm.py::test_rmsnorm_mixed_weight_forward_benchmark \ + -m benchmark -q -s >"${log}" 2>&1; then + SUCCESS_COUNT=$((SUCCESS_COUNT + 1)) + else + _fail_or_skip "${log}" "rmsnorm_mixed_weight" + fi + row="$(_py_parse_and_emit rmsnorm_mixed_weight "${M}x${N}" "${activation_dtype}+f32w" "${log}")" + set -- $row + _emit_row "$1" "$2" "$3" "$4" "$5" + + for mode in plain fused_add; do + export ROCDSL_RMSNORM_MIXED_WEIGHT_BWD_MODE="$mode" + op="rmsnorm_mixed_w_bwd" + if [ "$mode" = "fused_add" ]; then + op="rmsnorm_add_mixed_bwd" + fi + log="${BENCH_LOG_DIR}/${op}_${M}x${N}_${activation_dtype}.log" + if python3 -m pytest \ + tests/kernels/test_rmsnorm.py::test_rmsnorm_mixed_weight_backward_benchmark \ + -m benchmark -q -s >"${log}" 2>&1; then + SUCCESS_COUNT=$((SUCCESS_COUNT + 1)) + else + _fail_or_skip "${log}" "$op" + fi + row="$(_py_parse_and_emit "$op" "${M}x${N}" "${activation_dtype}+f32w" "${log}")" + set -- $row + _emit_row "$1" "$2" "$3" "$4" "$5" + done + done + unset ROCDSL_RMSNORM_MIXED_WEIGHT_BENCH_SHAPE ROCDSL_RMSNORM_MIXED_WEIGHT_BWD_MODE fi # FlashAttention / FMHA (CDNA only) diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index 5fbb1e8e9..c520c42fe 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -158,13 +158,17 @@ def _get_rmsnorm_large_configs(): ] -def run_test(M: int, N: int, dtype: str = "f32"): - print(f"\nTesting RMSNorm (M={M}, 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})") try: - launch_fn = build_rmsnorm_module(N, dtype) + launch_fn = build_rmsnorm_module(N, dtype, weight_dtype_str=weight_dtype) except Exception as e: - print(f"[FAIL] Compile failed for (M={M}, N={N}, dtype={dtype}): {type(e).__name__}: {e}") + print( + f"[FAIL] Compile failed for (M={M}, N={N}, dtype={dtype}, " + f"weight_dtype={weight_dtype}): {type(e).__name__}: {e}" + ) return False, None torch.manual_seed(42) @@ -172,8 +176,9 @@ def run_test(M: int, N: int, dtype: str = "f32"): gamma_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) torch_dtype = _torch_dtype(dtype) + weight_torch_dtype = _torch_dtype(weight_dtype) input_dev = input_t.to(torch_dtype).contiguous() - gamma_dev = gamma_t.to(torch_dtype).contiguous() + gamma_dev = gamma_t.to(weight_torch_dtype).contiguous() output_dev = torch.empty((M, N), device="cuda", dtype=torch_dtype) input_ref = input_dev.to(DTYPE_FP32) gamma_ref = gamma_dev.to(DTYPE_FP32) @@ -207,7 +212,8 @@ def kernel_launch(): # Bandwidth estimate: read input + read gamma + write output elem_bytes = 4 if dtype == "f32" else 2 - total_bytes = 2 * M * N * elem_bytes + weight_elem_bytes = 4 if weight_dtype == "f32" else 2 + total_bytes = 2 * M * N * elem_bytes + N * weight_elem_bytes bandwidth_gbs = total_bytes / (avg_us / 1e6) / 1e9 print(f"Kernel avg time: {avg_ms:.4f} ms via run_perftest (warmup={WARMUP_ITERS}, iters={BENCH_ITERS})") @@ -929,13 +935,14 @@ def run_fused_add_autograd_test(M: int, N: int, dtype: str = "f32"): return ok -def _bench_rmsnorm_backward_against_torch(M, N, dtype, *, fused_add): +def _bench_rmsnorm_backward_against_torch(M, N, dtype, *, fused_add, weight_dtype=None): """Benchmark public autograd backward; graph construction is excluded.""" torch_dtype = _torch_dtype(dtype) + weight_torch_dtype = _torch_dtype(dtype if weight_dtype is None else weight_dtype) torch.manual_seed(42) x = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() - weight = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() + weight = torch.rand((N,), device="cuda", dtype=weight_torch_dtype).contiguous() dy = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() if fused_add: @@ -987,8 +994,10 @@ def run_torch(): dx_atol = {"f32": 1e-3, "f16": 3e-2, "bf16": 2e-1}[dtype] dw_atol = {"f32": 1e-2, "f16": 2e-1, "bf16": 1.0}[dtype] for index, (fly_grad, torch_grad) in enumerate(zip(fly_grads, torch_grads)): - atol = dw_atol if index == len(fly_grads) - 1 else dx_atol - torch.testing.assert_close(fly_grad, torch_grad, rtol=grad_rtol, atol=atol) + is_dweight = index == len(fly_grads) - 1 + atol = (5e-2 if weight_torch_dtype == DTYPE_FP32 else dw_atol) if is_dweight else dx_atol + rtol = 5e-3 if is_dweight and weight_torch_dtype == DTYPE_FP32 else grad_rtol + torch.testing.assert_close(fly_grad, torch_grad, rtol=rtol, atol=atol) del fly_grads, torch_grads flydsl_us = bench_gpu_us_torch(run_flydsl, warmup=WARMUP_ITERS, iters=BENCH_ITERS) @@ -1006,15 +1015,58 @@ def _print_rmsnorm_torch_perf_table(rows): ) print("=" * 100) print( - f"{'op':18s} {'shape':18s} {'dtype':6s} {'FlyDSL(gpu us)':>16s} " + f"{'op':18s} {'shape':18s} {'act':6s} {'weight':6s} {'FlyDSL(gpu us)':>16s} " f"{'Torch(gpu us)':>14s} {'Torch/FlyDSL':>13s}" ) - for op, M, N, dtype, flydsl_us, torch_us in rows: + for op, M, N, dtype, weight_dtype, flydsl_us, torch_us in rows: speedup = torch_us / flydsl_us - print(f"{op:18s} {f'{M}x{N}':18s} {dtype:6s} " f"{flydsl_us:16.1f} {torch_us:14.1f} {speedup:9.2f}x") + print( + f"{op:18s} {f'{M}x{N}':18s} {dtype:6s} {weight_dtype:6s} " + f"{flydsl_us:16.1f} {torch_us:14.1f} {speedup:9.2f}x" + ) print("=" * 100 + "\n") +def _get_mixed_weight_benchmark_shape(): + shape = os.environ.get("ROCDSL_RMSNORM_MIXED_WEIGHT_BENCH_SHAPE", "").strip() + if not shape: + pytest.skip("set ROCDSL_RMSNORM_MIXED_WEIGHT_BENCH_SHAPE=M,N,dtype") + M, N, dtype = [part.strip() for part in shape.split(",")] + return int(M), int(N), dtype + + +@pytest.mark.benchmark +def test_rmsnorm_mixed_weight_forward_benchmark(): + """Emit the standard norm benchmark row for FP32-weight forward.""" + M, N, dtype = _get_mixed_weight_benchmark_shape() + ok, _ = run_test(M, N, dtype, weight_dtype="f32") + assert ok + + +@pytest.mark.benchmark +def test_rmsnorm_mixed_weight_backward_benchmark(): + """Emit effective bandwidth for plain or fused FP32-weight backward.""" + M, N, dtype = _get_mixed_weight_benchmark_shape() + mode = os.environ.get("ROCDSL_RMSNORM_MIXED_WEIGHT_BWD_MODE", "plain").strip() + if mode not in ("plain", "fused_add"): + raise ValueError(f"unsupported mixed-weight backward benchmark mode: {mode!r}") + + fused_add = mode == "fused_add" + flydsl_us, torch_us = _bench_rmsnorm_backward_against_torch( + M, + N, + dtype, + fused_add=fused_add, + weight_dtype="f32", + ) + elem_bytes = 4 if dtype == "f32" else 2 + logical_bytes = (3 + int(fused_add)) * M * N * elem_bytes + M * 4 + N * 8 + bandwidth_gbs = logical_bytes / (flydsl_us / 1e6) / 1e9 + print(f"Kernel avg time: {flydsl_us / 1000.0:.4f} ms") + print(f"Bandwidth: {bandwidth_gbs:.2f} GB/s") + print(f"PyTorch eager: {torch_us:.1f} us") + + @pytest.mark.benchmark def test_rmsnorm_backward_torch_benchmark(): """Opt-in representative sweep for the #800 backward follow-up.""" @@ -1030,7 +1082,7 @@ def test_rmsnorm_backward_torch_benchmark(): dtype, fused_add=False, ) - rows.append(("rmsnorm_bwd", M, N, dtype, flydsl_us, torch_us)) + rows.append(("rmsnorm_bwd", M, N, dtype, dtype, flydsl_us, torch_us)) flydsl_us, torch_us = _bench_rmsnorm_backward_against_torch( M, @@ -1038,7 +1090,7 @@ def test_rmsnorm_backward_torch_benchmark(): dtype, fused_add=True, ) - rows.append(("fused_add_bwd", M, N, dtype, flydsl_us, torch_us)) + rows.append(("fused_add_bwd", M, N, dtype, dtype, flydsl_us, torch_us)) _print_rmsnorm_torch_perf_table(rows) @@ -1084,8 +1136,69 @@ def test_fused_add_rmsnorm_autograd(): raise SystemExit(1) +@pytest.mark.parametrize("fused_add", (False, True), ids=("plain", "fused_add")) +@pytest.mark.parametrize( + "M,N,dtype", + ( + (16, 512, "f16"), # small-N forward + atomic backward + (64, 3001, "bf16"), # generic forward + atomic backward + (512, 4096, "f16"), # vec8 forward + staged backward + (513, 4097, "bf16"), # generic forward + staged scalar backward + ), +) +def test_rmsnorm_mixed_fp32_weight_autograd(M, N, dtype, fused_add): + """FP16/BF16 activations support FP32 weights through both training paths.""" + torch_dtype = _torch_dtype(dtype) + torch.manual_seed(42) + + x = torch.randn((M, N), device="cuda", dtype=torch_dtype, requires_grad=True) + weight = torch.rand((N,), device="cuda", dtype=DTYPE_FP32, requires_grad=True) + dout = torch.randn_like(x) + + if fused_add: + residual = torch.randn_like(x, requires_grad=True) + dresidual_out = torch.randn_like(x) + out, residual_out = fused_add_rmsnorm(x, residual, weight, prenorm=True) + torch.autograd.backward((out, residual_out), (dout, dresidual_out)) + source_ref = (x.detach() + residual.detach()).to(DTYPE_FP32).requires_grad_(True) + else: + residual = None + dresidual_out = None + out = rmsnorm(x, weight) + out.backward(dout) + source_ref = x.detach().to(DTYPE_FP32).requires_grad_(True) + + weight_ref = weight.detach().clone().requires_grad_(True) + rstd_ref = torch.rsqrt(source_ref.square().mean(dim=1, keepdim=True) + EPS) + out_ref = source_ref * rstd_ref * weight_ref + outputs = (out_ref, source_ref) if fused_add else (out_ref,) + grad_outputs = (dout.to(DTYPE_FP32), dresidual_out.to(DTYPE_FP32)) if fused_add else (dout.to(DTYPE_FP32),) + dsource_ref, dweight_ref = torch.autograd.grad(outputs, (source_ref, weight_ref), grad_outputs) + + assert out.dtype == torch_dtype + assert x.grad.dtype == torch_dtype + assert weight.grad.dtype == DTYPE_FP32 + if fused_add: + assert residual_out.dtype == torch_dtype + assert residual.grad.dtype == torch_dtype + + path, num_programs = rmsnorm_kernel_impl._select_rmsnorm_bwd_config(M, N, dtype, x.device) + fwd_cache = rmsnorm_kernel_impl._FUSED_ADD_FWD_CACHE if fused_add else rmsnorm_kernel_impl._FWD_CACHE + bwd_cache = rmsnorm_kernel_impl._FUSED_ADD_BWD_CACHE if fused_add else rmsnorm_kernel_impl._BWD_CACHE + assert (N, dtype, "f32", True, float(EPS), x.device) in fwd_cache + assert (path, N, dtype, "f32", num_programs, x.device) in bwd_cache + + grad_rtol = {"f16": 3e-2, "bf16": 1e-1}[dtype] + value_atol = {"f16": 3e-2, "bf16": 2e-1}[dtype] + torch.testing.assert_close(out.to(DTYPE_FP32), out_ref, rtol=grad_rtol, atol=value_atol) + torch.testing.assert_close(x.grad.to(DTYPE_FP32), dsource_ref, rtol=grad_rtol, atol=value_atol) + if fused_add: + torch.testing.assert_close(residual.grad.to(DTYPE_FP32), dsource_ref, rtol=grad_rtol, atol=value_atol) + torch.testing.assert_close(weight.grad, dweight_ref, rtol=5e-3, atol=5e-2) + + def test_fused_add_rmsnorm_dtype_mismatch(): - """Mixed x/residual/weight dtypes must be rejected (kernel uses one elem dtype).""" + """Residual must match x; weight may only differ by using FP32.""" print("=" * 80) print("Running fused_add_rmsnorm dtype-mismatch guard Test") print("=" * 80) @@ -1093,11 +1206,12 @@ def test_fused_add_rmsnorm_dtype_mismatch(): x = torch.randn((4, N), device="cuda", dtype=DTYPE_BF16) res_bad = torch.randn((4, N), device="cuda", dtype=DTYPE_FP16) w = torch.rand((N,), device="cuda", dtype=DTYPE_BF16) - for bad in (res_bad, w.to(DTYPE_FP32)): + weight_bad = w.to(DTYPE_FP16) + for bad in (res_bad, weight_bad): try: fused_add_rmsnorm(x, res_bad if bad is res_bad else torch.randn_like(x), bad if bad is not res_bad else w) raise SystemExit("dtype mismatch was NOT rejected") - except AssertionError: + except (AssertionError, ValueError): pass print(" -> PASSED") @@ -1122,8 +1236,9 @@ def test_fused_add_rmsnorm_device_mismatch(): print(" -> PASSED") -def _bench_aiter_rmsnorm(M: int, N: int, dtype: str): +def _bench_aiter_rmsnorm(M: int, N: int, dtype: str, weight_dtype: str | None = None): torch_dtype = _torch_dtype(dtype) + weight_torch_dtype = _torch_dtype(dtype if weight_dtype is None else weight_dtype) try: from aiter.ops.triton.rmsnorm import rms_norm as aiter_rms_norm @@ -1132,7 +1247,7 @@ def _bench_aiter_rmsnorm(M: int, N: int, dtype: str): return None x = torch.randn((M, N), device="cuda", dtype=torch_dtype) - w = torch.rand((N,), device="cuda", dtype=torch_dtype) + w = torch.rand((N,), device="cuda", dtype=weight_torch_dtype) def run_aiter(): aiter_rms_norm(x, w, EPS) @@ -1257,17 +1372,24 @@ def test_rmsnorm(): failures = 0 for M, N, dtype in configs: - ok, flydsl_gpu_us = run_test(M, N, dtype) + weight_dtype = os.environ.get("ROCDSL_RMSNORM_WEIGHT_DTYPE", "").strip() or dtype + ok, flydsl_gpu_us = run_test(M, N, dtype, weight_dtype) if not ok: failures += 1 if do_compare: aiter_us = None if maybe_enable_aiter(): - aiter_us = _bench_aiter_rmsnorm(M, N, dtype) + aiter_us = _bench_aiter_rmsnorm(M, N, dtype, weight_dtype) perf_rows.append( - PerfRow(op="rmsnorm", shape=f"{M}x{N}", dtype=dtype, flydsl_gpu_us=flydsl_gpu_us, aiter_gpu_us=aiter_us) + PerfRow( + op="rmsnorm", + shape=f"{M}x{N}", + dtype=f"{dtype}/{weight_dtype}", + flydsl_gpu_us=flydsl_gpu_us, + aiter_gpu_us=aiter_us, + ) ) print("\n" + "=" * 80) @@ -1417,6 +1539,68 @@ def test_rmsnorm_bwd_dispatch_boundary(): assert num_programs is None +@pytest.mark.parametrize("fused_add", (False, True), ids=("plain", "fused_add")) +def test_rmsnorm_bwd_mixed_weight_keeps_vector_io(fused_add): + """The staged path loads each FP32 weight tile as two 128-bit vectors.""" + device = torch.device("cuda", torch.cuda.current_device()) + M, N = 512, 4096 + path, num_programs = rmsnorm_kernel_impl._select_rmsnorm_bwd_config(M, N, "bf16", device) + assert path == "two_stage" + + source = torch.randn((M, N), device=device, dtype=DTYPE_BF16) + weight = torch.rand((N,), device=device, dtype=DTYPE_FP32) + dout = torch.randn_like(source) + rstd = torch.rsqrt(source.to(DTYPE_FP32).square().mean(1) + EPS) + dx = torch.empty_like(source) + dweight = torch.empty_like(weight) + partial = torch.empty((num_programs * N,), device=device, dtype=DTYPE_FP32) + stream = torch.cuda.current_stream(device) + + if fused_add: + launcher = build_fused_add_rmsnorm_bwd_two_stage_module( + N, + "bf16", + num_programs, + weight_dtype_str="f32", + ) + dresidual_out = torch.randn_like(source) + compiled = flyc.compile( + launcher, + source, + weight, + dout, + dresidual_out, + rstd, + dx, + dweight, + partial, + M, + stream, + ) + else: + launcher = build_rmsnorm_bwd_two_stage_module( + N, + "bf16", + num_programs, + weight_dtype_str="f32", + ) + compiled = flyc.compile( + launcher, + source, + weight, + dout, + rstd, + dx, + dweight, + partial, + M, + stream, + ) + + weight_copy_type = "!fly.copy_atom, 32>" + assert compiled._keepalive.source_ir.count(weight_copy_type) >= 3 + + @pytest.mark.parametrize("fused_add", (False, True), ids=("plain", "fused_add")) def test_rmsnorm_bwd_two_stage_cache_reuse_across_m(monkeypatch, fused_add): """One staged compiled callable must serve multiple runtime row counts.""" @@ -1429,7 +1613,7 @@ def test_rmsnorm_bwd_two_stage_cache_reuse_across_m(monkeypatch, fused_add): builder_name = "build_fused_add_rmsnorm_bwd_two_stage_module" if fused_add else "build_rmsnorm_bwd_two_stage_module" cache = rmsnorm_kernel_impl._FUSED_ADD_BWD_CACHE if fused_add else rmsnorm_kernel_impl._BWD_CACHE - key = (path, N, "bf16", num_programs, device) + key = (path, N, "bf16", "bf16", num_programs, device) original_builder = getattr(rmsnorm_kernel_impl, builder_name) build_count = 0 run_compiled_count = 0 @@ -1514,26 +1698,31 @@ def test_rmsnorm_bwd_two_stage_multistream_workspace(): torch.testing.assert_close(dweight.to(DTYPE_FP32), dw_ref, rtol=1e-1, atol=1.0) -def test_rmsnorm_bwd_vec8_contiguous_storage_offset(): +@pytest.mark.parametrize("weight_dtype", (DTYPE_BF16, DTYPE_FP32), ids=("bf16_weight", "fp32_weight")) +def test_rmsnorm_vec8_contiguous_storage_offset(weight_dtype): """Contiguous tensors need not have a 16-byte-aligned storage offset.""" device = torch.device("cuda", torch.cuda.current_device()) M, N = 512, 4096 - def offset_rand(shape): + def offset_rand(shape, dtype=DTYPE_BF16): numel = 1 for dim in shape: numel *= dim - tensor = torch.randn((numel + 1,), device=device, dtype=DTYPE_BF16)[1:].view(shape) + tensor = torch.randn((numel + 1,), device=device, dtype=dtype)[1:].view(shape) assert tensor.is_contiguous() and tensor.data_ptr() % 16 != 0 return tensor added = offset_rand((M, N)) - weight = offset_rand((N,)) + weight = offset_rand((N,), weight_dtype) dout = offset_rand((M, N)) dresidual_out = offset_rand((M, N)) rstd = torch.rsqrt(added.to(DTYPE_FP32).square().mean(1) + EPS) dx_ref, dweight_ref, _ = _reference_rmsnorm_bwd(added, weight, dout) + out = rmsnorm(added, weight) + out_ref = _reference_rmsnorm(added, weight) + torch.testing.assert_close(out.to(DTYPE_FP32), out_ref, rtol=1e-1, atol=2e-1) + dx, dweight = rmsnorm_kernel_impl.rmsnorm_bwd(added, weight, dout, rstd) torch.testing.assert_close(dx.to(DTYPE_FP32), dx_ref, rtol=1e-1, atol=2e-1) torch.testing.assert_close(dweight.to(DTYPE_FP32), dweight_ref, rtol=1e-1, atol=1.0) diff --git a/tests/kernels/test_rmsnorm_autotune.py b/tests/kernels/test_rmsnorm_autotune.py index 6ca0a3414..917370e12 100644 --- a/tests/kernels/test_rmsnorm_autotune.py +++ b/tests/kernels/test_rmsnorm_autotune.py @@ -37,10 +37,10 @@ def _reference(x, g): return xf * torch.rsqrt((xf * xf).mean(-1, keepdim=True) + EPS) * g.float() -def _inputs(M=32, N=8192): +def _inputs(M=32, N=8192, weight_dtype=torch.bfloat16): generator = torch.Generator(device="cuda").manual_seed(0) x = torch.randn(M, N, device="cuda", dtype=torch.bfloat16, generator=generator) - g = torch.rand(N, device="cuda", dtype=torch.bfloat16, generator=generator) + g = torch.rand(N, device="cuda", dtype=weight_dtype, generator=generator) return x, g, _reference(x, g) @@ -48,18 +48,28 @@ def _assert_close(out, ref): torch.testing.assert_close(out.float(), ref, rtol=0, atol=2e-2) -def test_rmsnorm_direct_specializes_known_block_size(): - x, g, ref = _inputs(M=1) +@pytest.mark.parametrize( + "weight_dtype,weight_dtype_str", + ((torch.bfloat16, "bf16"), (torch.float32, "f32")), +) +def test_rmsnorm_direct_specializes_known_block_size(weight_dtype, weight_dtype_str): + x, g, ref = _inputs(M=1, weight_dtype=weight_dtype) out = torch.empty_like(x) stream = torch.cuda.current_stream() - compiled = flyc.compile(rmsnorm_direct, x, g, out, x.shape[0], x.shape[1], "bf16", 512, stream) + compile_args = (x, g, out, x.shape[0], x.shape[1], "bf16", 512, stream) + if weight_dtype == torch.float32: + compile_args += (weight_dtype_str,) + compiled = flyc.compile(rmsnorm_direct, *compile_args) stream.synchronize() artifact = compiled._keepalive assert "known_block_size = array" in artifact.source_ir match = re.search(r"max_flat_workgroup_size\s*=\s*(\d+)", artifact.ir) assert match is not None and int(match.group(1)) == 512 + if weight_dtype == torch.float32: + weight_copy_type = "!fly.copy_atom, 32>" + assert artifact.source_ir.count(weight_copy_type) >= 3 _assert_close(out, ref) @@ -110,7 +120,12 @@ def bench_once(call, warmup, rep): assert completed == len(_SEARCH_CONFIGS) _assert_close(out, ref) - call_kwargs = {"N": x.shape[1], "dtype_str": "bf16", "stream": raw_stream} + call_kwargs = { + "N": x.shape[1], + "dtype_str": "bf16", + "weight_dtype_str": "bf16", + "stream": raw_stream, + } winner_key = _rmsnorm_tuner._make_key((x, g, out, x.shape[0]), call_kwargs) assert winner_key in _rmsnorm_tuner.cache other_m_key = _rmsnorm_tuner._make_key((x, g, out, x.shape[0] + 1), call_kwargs) @@ -128,3 +143,26 @@ def bench_once(call, warmup, rep): assert completed == len(_SEARCH_CONFIGS) _assert_close(cached, ref) + + +def test_rmsnorm_weight_dtype_has_distinct_tuning_identity(): + x, g, _ = _inputs(M=1) + _, g_fp32, mixed_ref = _inputs(M=1, weight_dtype=torch.float32) + out = torch.empty_like(x) + common = {"N": x.shape[1], "dtype_str": "bf16", "stream": torch.cuda.current_stream().cuda_stream} + + same_dtype_key = _rmsnorm_tuner._make_key( + (x, g, out, x.shape[0]), + {**common, "weight_dtype_str": "bf16"}, + ) + mixed_dtype_key = _rmsnorm_tuner._make_key( + (x, g_fp32, out, x.shape[0]), + {**common, "weight_dtype_str": "f32"}, + ) + + assert "weight_dtype_str" in _rmsnorm_tuner.key + assert mixed_dtype_key != same_dtype_key + + rmsnorm_autotuned(x, g_fp32, out, x.shape[0]) + torch.cuda.synchronize() + _assert_close(out, mixed_ref)