[3/5] autotune: offline config emit + runtime lookup (#770)

docs/autotune_guide.md

@@ -0,0 +1,102 @@
+# Autotuning FlyDSL kernels
+
+FlyDSL's autotuner (`flydsl.autotune`) benchmarks candidate kernel configs,
+picks the fastest, and remembers it. It has three consumption modes, mirroring
+what Triton, quack, aiter, and SGLang do in practice.
+
+## The three paths
+
+1. **Heuristic default (zero search).** A `default(*args) -> Config` analytic
+   rule picks a good config with no benchmarking. This is what normal runs use —
+   tests and serving never pay tuning cost. (== Triton `@heuristics`,
+   quack `get_default`.)
+2. **Online JIT search.** With `FLYDSL_AUTOTUNE=1`, the tuner sweeps the config
+   space once per shape, benchmarks each, caches the winner in a machine-local
+   scratch cache, and reuses it for the rest of the process / future runs.
+3. **Offline committed configs.** Tuned configs written to a checked-in tree and
+   looked up at serving with no search — the aiter/SGLang model. Portable across
+   identical GPUs; the filename *is* the lookup key.
+
+## When to tune
+
+- **Dev / perf work:** `FLYDSL_AUTOTUNE=1` to explore the space for your shape.
+- **Committed offline configs:** tune the shapes you serve, commit the JSON, and
+  ship. Downstream engines resolve a config by reconstructing the filename.
+- **Everything else (CI, normal runs):** do **not** set `FLYDSL_AUTOTUNE`. The
+  heuristic default (or a committed offline config) is used, so no benchmark
+  loop runs.
+
+## Do not tune in CI
+
+Autotuning benchmarks the same kernel dozens of times per config. That is slow
+and noisy on shared CI runners, and the winning config would depend on the
+runner's load. CI should exercise the *default* and *offline-lookup* paths (fast,
+deterministic) and leave the search off. The GPU-free unit tests
+(`tests/unit/test_autotune.py`) cover serialization, cache keys, `restore_value`,
+pruning, and the offline emit/lookup logic without any GPU.
+
+## Cache behavior and keys
+
+Two separate stores:
+
+| Store | Env var | Keyed by | Portable? |
+|-------|---------|----------|-----------|
+| Scratch cache | `FLYDSL_AUTOTUNE_CACHE_DIR` (default `~/.flydsl/autotune`) | full fingerprint: shape, dtype, **stride pattern**, GPU arch, **toolchain fingerprint**, cache-invalidating env | no (machine-local) |
+| Offline configs | `FLYDSL_AUTOTUNE_CONFIG_DIR` | filename: `name,<spec axes>,device_name` | yes (commit + share) |
+
+The scratch cache is invalidated automatically when the compiler toolchain, GPU
+arch, or a relevant env var changes — a config tuned under one build is never
+silently reused under another. The offline tree is deliberately coarser: it keys
+only on `name`, the `specialize()` axes, and `device_name` — **not** the
+toolchain fingerprint, stride pattern, or non-spec tensor dtypes that the scratch
+cache folds in. So a config tuned once is reusable on any matching GPU, but this
+puts two requirements on the adopter:
+
+- **`specialize()` must return every axis that changes the best config** (row
+  width, dtype, and any layout/mode flag). Anything the kernel's performance
+  depends on but `specialize` omits will collide: two calls with the same spec
+  but different (say) stride pattern map to one artifact, and the last tune
+  wins. The scratch cache would keep them separate; the offline tree cannot.
+- **`specialize()` values must be JSON-round-trippable** (str/int/float/bool,
+  or tuples/lists of them). The spec is normalized through JSON on both emit and
+  lookup, so a tuple (stored as a list) still self-matches; but an `Enum` /
+  `torch.dtype` isn't serializable and disables the offline path for that op
+  (with a warning). Use its `.name` / string form instead.
+
+A committed artifact from an older toolchain but the same device name **is**
+served (no automatic invalidation) — treat re-tuning after a toolchain change as
+a review/policy step, and commit artifacts as reviewed inputs.
+
+An offline artifact must be fully self-describing (`name`, `spec`,
+`device_name`, `config`). `name`, `spec`, and `device_name` are matched against
+the call; the `config` body is checked for shape (a non-empty dict carrying the
+structural knobs) but its knob *values* are trusted as tuned. `spec` axis values
+may be scalars or lists (a tuple axis is stored as a list); `config` knob values
+must be scalars (they are hashed into the build cache and passed to the kernel).
+A mismatch, a missing field, malformed JSON, a non-dict `spec`, or an empty /
+non-scalar-valued `config` is ignored with a warning rather than silently
+mis-serving via the heuristic fallback. Filenames are sanitized to ASCII
+`[A-Za-z0-9._-]`, so no spec/name/device value can inject a delimiter or escape
+the config directory.
+
+## Correctness: `restore_value` / `reset_to_zero`
+
+Because tuning reruns the kernel many times, any kernel that writes its output
+in place or accumulates into its inputs (e.g. fused-add rmsnorm, where the
+output overlaps the residual buffer) corrupts its own data across reps — and the
+timing/selection is then made on garbage. Declare those tensors:
+
+- `restore_value=[...]` — snapshot before the reps, restore before each rep.
+- `reset_to_zero=[...]` — zero before each rep (accumulate-into-zero kernels).
+
+## Adopting `@autotune` on a kernel
+
+Most FlyDSL kernels bake structural knobs (block size, tile width) at
+module-build time rather than exposing them as jit `Constexpr` params, so use
+`autotune_builder(...)`: supply `build` (rebuilds the module per config),
+`specialize` (extracts the shape + build-only scalars — these become both the
+cache key and the offline filename), `configs`/`default` (the search space and
+the zero-search heuristic), and `structural` (which config knobs route into
+`build`). The helper owns cache naming, build caching, and the offline path. See
+`kernels/rmsnorm_autotune.py` and `kernels/rmsnorm_config.py` for the reference
+adoption — it is a single declaration.

kernels/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.rmsnorm_config import get_all_configs, get_default
+from kernels.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/rmsnorm_config.py

@@ -0,0 +1,65 @@
+# 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.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
+
+
+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.
+    """
+    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
+    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. Configs whose
+    vectorized tile doesn't evenly divide the row are dropped (they'd fall to
+    the untuned scalar path)."""
+    # Small-N kernel ignores BLOCK_THREADS, so there's nothing to sweep.
+    if N <= SMALL_N_THRESHOLD:
+        return [get_default(N, dtype_str, arch)]
+
+    vec_width = 128 // _elem_bits(dtype_str)
+    configs = []
+    for block in _BLOCK_THREADS_CHOICES:
+        tile_cols = block * vec_width
+        # Keep only configs that hit the vectorized fast path for this N.
+        if N < tile_cols or N % tile_cols != 0 or _elem_bits(dtype_str) > 16:
+            continue
+        for wpe in _WAVES_PER_EU_CHOICES:
+            kw = {"BLOCK_THREADS": block}
+            waves = None if wpe == 0 else wpe
+            configs.append(Config(waves_per_eu=waves, **kw))
+    # Fall back to the heuristic default if no fast-path config fits this N.
+    if not configs:
+        configs.append(get_default(N, dtype_str, arch))
+    return configs

kernels/rmsnorm_kernel.py

@@ -28,6 +28,11 @@
 WARP_SIZE = get_warp_size()
 VEC_WIDTH = 8
 
+# 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 _make_reduction_storage(red_slots: int):
     @fx.struct
@@ -110,20 +115,27 @@ 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):
-    if N <= 2048:
+def build_rmsnorm_module(N: int, dtype_str: str, BLOCK_THREADS: int = BLOCK_THREADS):
+    if N <= SMALL_N_THRESHOLD:
         return _build_rmsnorm_large_m_small_n_module(N, dtype_str)
 
     arch = get_hip_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.2"
 
-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,
+)