[2/5] autotune: add opt-in search and RMSNorm waves-per-EU tuning (#770)

kernels/norm/rmsnorm_autotune.py

@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2025 FlyDSL Project Contributors
+
+"""Autotuned RMSNorm — the first real adopter of ``flydsl.autotune``.
+
+RMSNorm bakes its structural knob (BLOCK_THREADS) at module-build time, so the
+tuner rebuilds the module per config via ``autotune_builder`` (builder mode).
+Normal runs use the ``get_default`` heuristic; ``FLYDSL_AUTOTUNE=1`` sweeps
+``get_all_configs``.
+
+    rmsnorm_autotuned(input, gamma, output, M, dtype_str="bf16", stream=stream)
+"""
+
+from flydsl.autotune import autotune_builder
+from kernels.norm.rmsnorm_config import get_all_configs, get_default
+from kernels.norm.rmsnorm_kernel import build_rmsnorm_module
+
+
+def _specialize(input_t, gamma, output, m_in, dtype_str="bf16", stream=None):
+    # Build/lookup axes; dtype_str must be here so bf16 vs f16 keys differ.
+    return {"N": int(input_t.shape[-1]), "dtype_str": dtype_str}
+
+
+rmsnorm_autotuned = autotune_builder(
+    name="rmsnorm",
+    build=build_rmsnorm_module,
+    specialize=_specialize,
+    configs=get_all_configs,
+    default=get_default,
+    structural=("BLOCK_THREADS",),
+)

kernels/norm/rmsnorm_config.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2025 FlyDSL Project Contributors
+
+"""Two-track tuning configs for RMSNorm (quack-style get_default + exhaustive):
+
+  - ``get_default`` — analytic BLOCK_THREADS, zero search (normal runs).
+  - ``get_all_configs`` — BLOCK_THREADS x waves_per_eu, swept under FLYDSL_AUTOTUNE=1.
+
+VEC_WIDTH is not tuned (pinned to 128//elem_bits by the 128-bit buffer copy).
+"""
+
+from flydsl.autotune import Config
+from kernels.norm.rmsnorm_kernel import SMALL_N_THRESHOLD
+
+# Candidate block sizes. All are multiples of the warp size (64 on CDNA) and
+# span both the <=256 (no known_block_size) and >256 (needs known_block_size)
+# regimes so the tuner can trade occupancy against per-thread work.
+_BLOCK_THREADS_CHOICES = (128, 256, 512, 1024)
+_WAVES_PER_EU_CHOICES = (0, 1, 2)  # 0 == leave to the compiler; nonzero lowers to the rocdl.waves_per_eu func attr
+
+
+def _elem_bits(dtype_str: str) -> int:
+    return 32 if dtype_str == "f32" else 16
+
+
+def get_default(N: int, dtype_str: str, arch: str = None) -> Config:
+    """Analytic default — a solid BLOCK_THREADS without searching.
+
+    Heuristic: pick the smallest block whose vectorized tiles cover the row in a
+    handful of iterations, clamped to [128, 1024]. Wider rows want more threads;
+    narrow rows keep the block small to preserve occupancy.
+
+    For bf16/f16 the pick is then shrunk to the largest candidate whose tile
+    (``BLOCK_THREADS * VEC_WIDTH``) evenly divides N, so the zero-search default
+    actually hits the vectorized fast path (the kernel gates it on
+    ``N % tile_cols == 0``) -- matching the divisibility filter get_all_configs
+    applies. Without this, a "tile-aware" block can still miss the fast path and
+    silently fall to the slow scalar loop for common N (e.g. bf16 N=5120 picks
+    256 whose tile 2048 does not divide 5120 -> scalar, while 128 would
+    vectorize). If nothing divides N, every block runs scalar anyway, so keep
+    the heuristic pick. f32 always uses the scalar loop, so no filter applies.
+    """
+    vec_width = 128 // _elem_bits(dtype_str)
+    # Aim for ~2 vectorized tiles per row: block ≈ N / (2 * vec_width).
+    target = N // max(1, (2 * vec_width))
+    block = 128
+    for choice in _BLOCK_THREADS_CHOICES:
+        if choice <= max(128, target):
+            block = choice
+
+    if _elem_bits(dtype_str) <= 16:
+        dividing = [b for b in _BLOCK_THREADS_CHOICES if b <= block and N >= b * vec_width and N % (b * vec_width) == 0]
+        if dividing:
+            block = max(dividing)
+
+    return Config(BLOCK_THREADS=block)
+
+
+def get_all_configs(N: int, dtype_str: str, arch: str = None):
+    """Exhaustive search space: BLOCK_THREADS x waves_per_eu.
+
+    bf16/f16 take the vectorized fast path (gated on ``elem_bits <= 16`` in the
+    kernel), so only BLOCK_THREADS whose tile ``BLOCK_THREADS * VEC_WIDTH``
+    evenly divides the row are kept. f32 never takes that path -- it uses the
+    scalar loop, which strides by BLOCK_THREADS and handles any N -- so every
+    BLOCK_THREADS is a valid, distinct f32 candidate (no tile filter). Previously
+    f32 was dropped entirely and silently collapsed to the single default."""
+    # Small-N kernel ignores BLOCK_THREADS, so there's nothing to sweep.
+    if N <= SMALL_N_THRESHOLD:
+        return [get_default(N, dtype_str, arch)]
+
+    vectorized = _elem_bits(dtype_str) <= 16
+    vec_width = 128 // _elem_bits(dtype_str)
+    configs = []
+    for block in _BLOCK_THREADS_CHOICES:
+        if vectorized:
+            tile_cols = block * vec_width
+            # bf16/f16: keep only blocks that hit the vectorized fast path for N.
+            if N < tile_cols or N % tile_cols != 0:
+                continue
+        for wpe in _WAVES_PER_EU_CHOICES:
+            waves = None if wpe == 0 else wpe
+            configs.append(Config(waves_per_eu=waves, BLOCK_THREADS=block))
+    # Fall back to the heuristic default if nothing fit (e.g. an odd bf16 N).
+    if not configs:
+        configs.append(get_default(N, dtype_str, arch))
+    return configs

kernels/norm/rmsnorm_kernel.py

@@ -47,6 +47,11 @@
 
 KERNEL_NAME = "rmsnorm"
 
+# N at or below this routes to the small-N kernel, whose block geometry is
+# derived analytically and ignores BLOCK_THREADS. Single source of truth so the
+# autotune config space (rmsnorm_config) stays in sync.
+SMALL_N_THRESHOLD = 2048
+
 
 def _store_yscale(scale_copy_atom, yscale_div, index, val):
     r = fx.make_rmem_tensor(1, fx.Float32)
@@ -67,20 +72,29 @@ def _quant_dtype_max(dtype_str: str) -> float:
     raise ValueError(f"unsupported quant dtype: {dtype_str!r} (expected 'i8' or 'int8')")
 
 
-def build_rmsnorm_module(N: int, dtype_str: str, store_rstd: bool = False, eps: float = EPS):
-    if N <= 2048:
+def build_rmsnorm_module(
+    N: int, dtype_str: str, store_rstd: bool = False, eps: float = EPS, BLOCK_THREADS: int = BLOCK_THREADS
+):
+    if N <= SMALL_N_THRESHOLD:
         return _build_rmsnorm_large_m_small_n_module(N, dtype_str, store_rstd, eps)
 
     arch = get_rocm_arch()
     USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95")
 
+    # BLOCK_THREADS is the block size (threads per row-block). It is a build-time
+    # structural knob: it sizes the shared reduction storage, the vectorized
+    # tile stride, and the launch block dim, so it is baked into the module
+    # rather than passed as a jit Constexpr. Autotune (builder mode) rebuilds
+    # this module per candidate BLOCK_THREADS. `known_block_size` is required
+    # on AMDGPU once the block exceeds 256.
     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
+    _kernel_kwargs = {} if BLOCK_THREADS <= 256 else {"known_block_size": [BLOCK_THREADS, 1, 1]}
 
     SharedStorage = _make_reduction_storage(RED_SLOTS)
 
-    @flyc.kernel
+    @flyc.kernel(**_kernel_kwargs)
     def rmsnorm_kernel(
         Input: fx.Tensor,
         Gamma: fx.Tensor,

python/flydsl/__init__.py

@@ -4,4 +4,8 @@
 
 __version__ = "0.2.4"
 
-from .autotune import Config as Config, autotune as autotune  # noqa: E402
+from .autotune import (  # noqa: E402
+    Config as Config,
+    autotune as autotune,
+    autotune_builder as autotune_builder,
+)