diff --git a/docs/autotune_guide.md b/docs/autotune_guide.md new file mode 100644 index 000000000..e0beb4753 --- /dev/null +++ b/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,,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. diff --git a/kernels/norm/rmsnorm_kernel.py b/kernels/norm/rmsnorm_kernel.py index 424f06468..69425a762 100644 --- a/kernels/norm/rmsnorm_kernel.py +++ b/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,33 @@ 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, + BLOCK_THREADS: int = BLOCK_THREADS, + store_rstd: bool = False, + eps: float = EPS, +): + 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, diff --git a/kernels/rmsnorm_autotune.py b/kernels/rmsnorm_autotune.py new file mode 100644 index 000000000..e2336c5cd --- /dev/null +++ b/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.norm.rmsnorm_kernel import build_rmsnorm_module +from kernels.rmsnorm_config import get_all_configs, get_default + + +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",), +) diff --git a/kernels/rmsnorm_config.py b/kernels/rmsnorm_config.py new file mode 100644 index 000000000..1f0c5fac3 --- /dev/null +++ b/kernels/rmsnorm_config.py @@ -0,0 +1,73 @@ +# 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, VEC_WIDTH + +# 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)] + + # The BLOCK_THREADS sweep only tunes the vectorized fast path, which packs + # VEC_WIDTH 16-bit elements into a 128-bit buffer copy. f32 (and any wider + # dtype) runs the generic scalar path this sweep can't tune — return the + # heuristic default explicitly, rather than a single-config "search" that + # looks like tuning but isn't. + if _elem_bits(dtype_str) > 16: + return [get_default(N, dtype_str, arch)] + + # tile_cols reuses the kernel's VEC_WIDTH (single source of truth) so the + # divisibility filter matches the kernel's actual vectorized tile. + configs = [] + for block in _BLOCK_THREADS_CHOICES: + tile_cols = block * VEC_WIDTH + # Keep only configs whose vectorized tile evenly divides the row. + 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(BLOCK_THREADS=block, waves_per_eu=waves)) + # 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 diff --git a/python/flydsl/__init__.py b/python/flydsl/__init__.py index 3909c79d8..b728b0452 100644 --- a/python/flydsl/__init__.py +++ b/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, +) diff --git a/python/flydsl/autotune.py b/python/flydsl/autotune.py index b243d528b..35b18b1f2 100644 --- a/python/flydsl/autotune.py +++ b/python/flydsl/autotune.py @@ -3,9 +3,14 @@ """FlyDSL autotuner - benchmark multiple kernel configs, pick the fastest.""" +import functools import inspect import json +import logging +import math import os +import tempfile +import threading from pathlib import Path from typing import Callable, Dict, List @@ -14,6 +19,41 @@ except ImportError: torch = None +# Module logger. Warnings on the (zero-search) serving path go through this so a +# serving stack can level/silence/redirect them instead of unconditional stdout. +logger = logging.getLogger("flydsl.autotune") + + +def _atomic_write_text(path: Path, text: str) -> None: + """Write *text* to *path* atomically (same-dir tempfile + ``os.replace``). + + A concurrent reader (e.g. serving resolving an offline artifact) never sees a + truncated file, and a crashed/interrupted writer can't leave a half-written + one — both real under multi-process / multi-GPU tuning into a shared tree. + """ + path.parent.mkdir(parents=True, exist_ok=True) + fd, tmp = tempfile.mkstemp(dir=str(path.parent), suffix=".tmp") + try: + with os.fdopen(fd, "w") as f: + f.write(text) + os.replace(tmp, path) + except BaseException: + try: + os.unlink(tmp) + except OSError: + pass + raise + + +def _tuning_enabled() -> bool: + """Whether to run the full search when a heuristic default exists. + + Off by default so normal runs (tests, serving) pay zero search cost and use + the analytic default. Opt in with ``FLYDSL_AUTOTUNE=1`` to actually tune. + Autotuners without a ``default`` always search (there is no fallback). + """ + return os.environ.get("FLYDSL_AUTOTUNE", "0") not in ("0", "", "false", "False") + def _env_fingerprint() -> tuple: """Sorted cache-invalidating env vars (reuses the JIT's canonical list).""" @@ -51,6 +91,61 @@ def _device_fingerprint() -> str: return "" +def _device_name() -> str: + """Human-readable device name for offline filenames (e.g. + 'AMD_Instinct_MI300X'); falls back to the arch string.""" + name = "" + if torch is not None: + try: + if torch.cuda.is_available(): + # current device, not 0 — correct under mixed/multi-GPU processes. + name = torch.cuda.get_device_name(torch.cuda.current_device()) + except Exception: + name = "" + return (name or _device_fingerprint() or "unknown").replace(" ", "_") + + +def _offline_config_dir(): + """Committed offline-config tree (FLYDSL_AUTOTUNE_CONFIG_DIR), or None. + + The aiter/SGLang "offline" model: tune once, commit, look up at serving with + no search. Portable, human-named artifacts — unlike the machine-local + fingerprint scratch cache (FLYDSL_AUTOTUNE_CACHE_DIR).""" + d = os.environ.get("FLYDSL_AUTOTUNE_CONFIG_DIR") + return Path(d) if d else None + + +def _safe_component(value) -> str: + """Map a filename component to strict ASCII [A-Za-z0-9._-] so no shape/name/ + device value can inject a ',' / '=' delimiter or a path separator.""" + safe = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789._-" + return "".join(c if c in safe else "_" for c in str(value)) + + +def _is_json_scalar(v) -> bool: + """A JSON scalar. Config knob values must be scalars: they are hashable and + flow into the build-cache key (a list would be unhashable) and reach the + kernel as a launch constant. (Spec axis values, by contrast, may be lists — + they are only compared by equality — but that check lives with the offline + config guard that needs it.)""" + return isinstance(v, (str, int, float, bool)) + + +def offline_config_filename(name, spec, device_name=None): + """SGLang-style filename-as-key: ``name,k1=v1,...,device_name=X.json``. + + ``spec`` is the sequence of (axis, value) pairs from the tuner's key (shape + + build-only scalars like dtype). The filename fully identifies the config, so + a downstream engine resolves it by reconstructing the name.""" + if device_name is None: + device_name = _device_name() + parts = [_safe_component(name)] + for k, v in spec: + parts.append(f"{_safe_component(k)}={_safe_component(v)}") + parts.append(f"device_name={_safe_component(device_name)}") + return ",".join(parts) + ".json" + + def _normalize_strides(t) -> tuple: """Bucket strides to {0, 1, other}: the layout *pattern* (broadcast / contiguous / strided) affects the best config, the exact numbers don't.""" @@ -111,6 +206,9 @@ def __repr__(self): return f"Config({', '.join(parts)})" def to_dict(self): + # Note: pre_hook is intentionally not serialized (it's a callable, not + # JSON), so a pre_hook that affects correctness won't survive the disk + # cache — keep pre_hook for timing side-effects only. d = dict(self.kwargs) for k in ("num_warps", "waves_per_eu", "maxnreg"): v = getattr(self, k) @@ -165,9 +263,15 @@ def __init__( pre_hook=None, post_hook=None, do_bench_fn=None, + build_fn=None, + default=None, + name=None, + key_fn=None, + arg_names=None, + structural=None, ): - self.fn = fn # JitFunction instance - self.configs = configs + self.fn = fn # JitFunction instance (None in builder mode) + self.configs = configs # list, or callable(*args, **kwargs) -> [Config] self.key = key or [] self.warmup = warmup self.rep = rep @@ -177,24 +281,73 @@ def __init__( self.pre_hook = pre_hook self.post_hook = post_hook self._do_bench = do_bench_fn or do_bench + # Resolved serving decision per key (a tuned best, an offline artifact, + # or a default). Memoizing all three keeps the serving path zero-I/O + # after the first call — otherwise every call would re-stat/parse the + # offline artifact and re-warn about a rejected one. self.cache: Dict[tuple, Config] = {} - - # Infer arg names from the underlying function - if hasattr(fn, "func"): - self.arg_names = list(inspect.signature(fn.func).parameters.keys()) + # Subset of `cache` keys that came from an actual search. Only these are + # persisted to the disk cache, so a memoized default/offline config can't + # be written to disk and shadow a committed offline artifact next process. + self._tuned_keys: set = set() + + # Serializes the compiler-hint mutate/run/restore on the shared launch fn + # (see _run_with_hints) so concurrent serving threads can't compile/run + # under each other's hints. Only the hinted path takes it; the common + # hint-free default/serving path is untouched. GPU launches are async + # enqueues, so the lock is held only briefly. + self._hint_lock = threading.Lock() + + # Tokens already warned about (bad offline artifacts, dropped builder + # kwargs), so a persistent problem warns once instead of on every call. + self._warned_once: set = set() + + # Builder mode: build_fn rebuilds the module per config. `structural` + # names the config kwargs build_fn consumes; the build cache keys only on + # those, so hint-only variants (waves_per_eu) reuse one built module. + self.build_fn = build_fn + self.default = default + self.structural = tuple(structural) if structural is not None else None + self._build_cache: Dict[tuple, object] = {} + + # key_fn(*args, **kwargs) -> ((name, value), ...): the specialization + # axes. When set it replaces the self.key name lookup in _make_key, so + # build-only scalars (dtype_str, causal, ...) enter the key too. + self.key_fn = key_fn + + # Arg names for reset/restore/filter lookup: explicit > jit fn sig > + # build_fn sig minus leading 'config'. + if arg_names is not None: + self.arg_names = list(arg_names) + elif fn is not None: + src = fn.func if hasattr(fn, "func") else fn + self.arg_names = list(inspect.signature(src).parameters.keys()) + elif build_fn is not None: + src = build_fn.func if hasattr(build_fn, "func") else build_fn + self.arg_names = list(inspect.signature(src).parameters.keys())[1:] else: - self.arg_names = list(inspect.signature(fn).parameters.keys()) + self.arg_names = [] - # Disk cache - fn_name = getattr(fn, "__name__", None) or getattr(fn, "func", None) - if fn_name is not None and not isinstance(fn_name, str): - fn_name = getattr(fn_name, "__name__", "unknown") - fn_name = fn_name or "unknown" - cache_dir = Path(os.environ.get("FLYDSL_AUTOTUNE_CACHE_DIR", os.path.expanduser("~/.flydsl/autotune"))) - self._cache_file = cache_dir / f"{fn_name}.json" + # Disk cache. Prefer an explicit name (required for builder mode, where + # fn is None — otherwise every builder tuner would share unknown.json). + cache_name = name + if cache_name is None: + cache_name = getattr(fn, "__name__", None) or getattr(fn, "func", None) + if cache_name is not None and not isinstance(cache_name, str): + cache_name = getattr(cache_name, "__name__", None) + self.name = cache_name or "unknown" self._load_disk_cache() + @property + def _cache_file(self) -> Path: + # Resolved per access so FLYDSL_AUTOTUNE_CACHE_DIR can change between + # calls (a module-level tuner isn't pinned to the import-time dir). + cache_dir = Path(os.environ.get("FLYDSL_AUTOTUNE_CACHE_DIR", os.path.expanduser("~/.flydsl/autotune"))) + # Sanitize like offline filenames: a name with '/' (or '..') must not + # escape the cache dir or collide with a path separator. + return cache_dir / f"{_safe_component(self.name)}.json" + def _make_key(self, args, kwargs): """Cache key over shape/dtype/stride + arch + toolchain + env. A config tuned under any of these axes must not be reused under another.""" @@ -202,14 +355,18 @@ def _make_key(self, args, kwargs): sig_args.update(kwargs) key_vals = [] - for k in self.key: - v = sig_args.get(k) - if hasattr(v, "shape"): - key_vals.append(tuple(v.shape)) - elif hasattr(v, "dtype"): - key_vals.append(str(v.dtype)) - else: - key_vals.append(v) + if self.key_fn is not None: + # Explicit specialization axes (includes build-only scalars). + key_vals.append(tuple(self.key_fn(*args, **kwargs))) + else: + for k in self.key: + v = sig_args.get(k) + if hasattr(v, "shape"): + key_vals.append(tuple(v.shape)) + elif hasattr(v, "dtype"): + key_vals.append(str(v.dtype)) + else: + key_vals.append(v) # Tensor dtypes + stride patterns, sorted so kwarg order doesn't change # the key (else identical calls would tune twice). @@ -258,7 +415,10 @@ def _snapshot_tensors(self, args, kwargs): snapshot = {} for name in self.restore_value: t = sig_args.get(name) - if t is not None and hasattr(t, "clone"): + # Need both clone (to snapshot) and copy_ (to restore); gating on + # clone alone would AttributeError in _restore_tensors for a + # tensor-like that can clone but not copy_ in place. + if t is not None and hasattr(t, "clone") and hasattr(t, "copy_"): snapshot[name] = (t, t.clone()) return snapshot @@ -268,6 +428,14 @@ def _restore_tensors(snapshot): for _name, (dst, src) in snapshot.items(): dst.copy_(src) + def _warn_once(self, token, msg): + """Emit a warning at most once per *token* (dedup key), so a persistent + problem on the serving path doesn't spam every call.""" + if token in self._warned_once: + return + self._warned_once.add(token) + logger.warning("%s", msg) + def _prune(self, configs, args, kwargs): if self.prune_configs_by is not None: sig_args = dict(zip(self.arg_names, args)) @@ -275,11 +443,50 @@ def _prune(self, configs, args, kwargs): return self.prune_configs_by(configs, sig_args) return configs - def _bench_one(self, config, args, kwargs): + def _resolve_fn(self, config, key, args, kwargs): + """Return the launch callable for a config. + + Direct mode: the wrapped jit fn. Builder mode: build the module (cached), + keyed on the spec + the structural knobs build_fn consumes — NOT the + whole config, so configs differing only in compiler hints (waves_per_eu) + share one build instead of recompiling per hint. + """ + if self.build_fn is None: + return self.fn + # Builder mode consumes only `structural` knobs (via build_fn) and + # compiler hints (via compiler_opts); any other config kwarg or num_warps + # is silently dropped. Warn once so a mis-specified config isn't a no-op. + if self.structural is not None: + dropped = [k for k in config.kwargs if k not in self.structural] + if config.num_warps is not None: + dropped.append("num_warps") + if dropped: + self._warn_once( + ("builder_drop", tuple(sorted(dropped))), + f"[autotune] {self.name}: builder mode ignores non-structural " + f"config value(s) {sorted(dropped)} (only {list(self.structural)} " + f"+ compiler hints are applied)", + ) + if self.structural is not None: + knob_key = tuple((k, config.kwargs.get(k)) for k in self.structural) + else: + knob_key = repr(config) # unknown knobs: fall back to full identity + cache_key = (key, knob_key) + built = self._build_cache.get(cache_key) + if built is None: + built = self.build_fn(config, *args, **kwargs) + self._build_cache[cache_key] = built + return built + + def _bench_one(self, config, key, args, kwargs): """Compile and benchmark one config. Returns time in ms.""" - merged_kwargs = dict(kwargs) - merged_kwargs.update(config.all_kwargs()) compiler_opts = config.compiler_opts() + fn = self._resolve_fn(config, key, args, kwargs) + merged_kwargs = dict(self._filter_call_kwargs(fn, kwargs)) + # In builder mode the config's structural kwargs (e.g. BLOCK_THREADS) + # are consumed by build_fn, not passed to the launch call. + if self.build_fn is None: + merged_kwargs.update(config.all_kwargs()) # Snapshot once before any rep runs, so restores are from pristine input. snapshot = self._snapshot_tensors(args, merged_kwargs) @@ -287,14 +494,14 @@ def _bench_one(self, config, args, kwargs): def kernel_call(): # Order: restore/reset the inputs first, THEN run the pre_hooks, so a # hook that sets up state (incl. mutating a tensor) isn't clobbered - # by the restore. Each benchmark rep starts from clean inputs. + # by the restore. (Matches Triton: pre_hook runs on clean inputs.) self._restore_tensors(snapshot) self._reset_tensors(args, merged_kwargs) if config.pre_hook: config.pre_hook(merged_kwargs) if self.pre_hook: self.pre_hook(merged_kwargs) - self._run_with_hints(compiler_opts, args, merged_kwargs) + self._run_with_hints(fn, compiler_opts, args, merged_kwargs) if self.post_hook: self.post_hook(merged_kwargs) @@ -305,38 +512,119 @@ def kernel_call(): if snapshot: self._restore_tensors(snapshot) - def _run_with_hints(self, compiler_opts, args, kwargs): - """Run the kernel with optional compiler hints. Import is deferred so - the core stays importable without the compiled bindings when unused.""" - if compiler_opts: - from .compiler.kernel_function import CompilationContext - - with CompilationContext.compile_hints(compiler_opts): - self.fn(*args, **kwargs) - else: - self.fn(*args, **kwargs) - - def _run_config(self, config, args, kwargs): - """Run the chosen config as a real (non-benchmark) call. Re-applies - reset_to_zero so cache hits and the post-tune run behave like a single - clean run (restore_value tensors are already restored by _bench_one).""" - merged = dict(kwargs) - merged.update(config.all_kwargs()) + def _run_with_hints(self, fn, compiler_opts, args, kwargs): + """Run fn with optional compiler hints. Hints are folded into + fn.compile_hints (a shared attribute that enters the JIT cache key) and + restored after, so each distinct waves_per_eu / maxnreg compiles a + distinct binary instead of reusing a cached one. Import is deferred so + the core stays importable unbuilt.""" + if not compiler_opts: + return fn(*args, **kwargs) + + from .compiler.kernel_function import CompilationContext + + prev_hints = getattr(fn, "compile_hints", None) + # Serialize the mutate/run/restore: fn.compile_hints is shared (builder + # mode reuses one built fn across hint-only config variants), so two + # threads serving different hinted configs must not interleave here and + # compile/run under each other's hints. The lock spans only the (async) + # kernel enqueue, not the GPU execution. + with self._hint_lock: + if prev_hints is not None: + # JitFunction: fold hints into its cache key so each distinct + # (waves_per_eu, maxnreg) compiles and caches a distinct binary. + fn.compile_hints = {**prev_hints, **compiler_opts} + try: + with CompilationContext.compile_hints(compiler_opts): + return fn(*args, **kwargs) + finally: + if prev_hints is not None: + fn.compile_hints = prev_hints + + def _filter_call_kwargs(self, fn, kwargs): + """Drop kwargs the launch fn doesn't accept. In builder mode the caller + may pass build-only kwargs (e.g. dtype_str) that route to build_fn but + aren't launch params; the built jit fn binds strictly.""" + if self.build_fn is None: + return kwargs + src = fn.func if hasattr(fn, "func") else fn + try: + params = inspect.signature(src).parameters + except (TypeError, ValueError): + return kwargs + if any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()): + return kwargs + return {k: v for k, v in kwargs.items() if k in params} + + def _run_config(self, config, key, args, kwargs): + """Run the chosen config as a real (non-benchmark) call. Resolves the + launch fn (builder mode rebuilds per config), applies config kwargs + + hints, and re-applies reset_to_zero so cache hits and the post-tune run + behave like a single clean run (restore_value already handled).""" + fn = self._resolve_fn(config, key, args, kwargs) + merged = dict(self._filter_call_kwargs(fn, kwargs)) + # Builder mode: structural kwargs (e.g. BLOCK_THREADS) go to build_fn, + # not the launch call. + if self.build_fn is None: + merged.update(config.all_kwargs()) self._reset_tensors(args, merged) - return self._run_with_hints(config.compiler_opts(), args, merged) + return self._run_with_hints(fn, config.compiler_opts(), args, merged) def __call__(self, *args, **kwargs): + self._load_disk_cache() # pick up the current cache dir (may be set post-init) key = self._make_key(args, kwargs) - if key in self.cache: - return self._run_config(self.cache[key], args, kwargs) - # Benchmark all configs - configs = self._prune(self.configs, args, kwargs) + # FLYDSL_AUTOTUNE=1 forces a fresh search: bypass the in-memory/disk + # cache and the heuristic default so an explicit tune re-benchmarks and + # re-emits, instead of short-circuiting on a stale cached best. + force = _tuning_enabled() + + # 1. Memoized serving decision for this key (tuned best from a prior tune + # or the disk cache, an offline artifact, or a default resolved on a + # prior call) — keeps the serving path zero-I/O after the first call. + if not force and key in self.cache: + return self._run_config(self.cache[key], key, args, kwargs) + + # 2. Committed offline artifact (aiter/SGLang model): no search, portable + # across identical GPUs. Validated before use (see _load_offline_config). + if not force: + cfg = self._load_offline_config(args, kwargs) + if cfg is not None: + try: + result = self._run_config(cfg, key, args, kwargs) + except Exception as e: + # Validation can't catch every bad artifact (e.g. a knob + # value the builder rejects at build time). Never crash the + # serving call: warn once and fall through to default/search. + self._warn_once( + ("offline_runtime", self.name, key), + f"[autotune] {self.name}: offline config failed at runtime " + f"({e}); falling back to default/search", + ) + else: + self.cache[key] = cfg # memoized (not tuned → not persisted) + logger.info("[autotune] %s: serving offline config", self.name) + return result + + # 3. Two-track heuristic: unless tuning is explicitly requested, take + # the analytic default and skip the search entirely (zero-search + # normal run). Mirrors Triton @heuristics / quack get_default. Memoize + # it so subsequent calls short-circuit at step 1 instead of re-resolving + # the offline path (stat + JSON) and re-warning on every serving call. + if not force and self.default is not None: + cfg = self.default(*args, **kwargs) + self.cache[key] = cfg # memoized (not tuned → not persisted) + return self._run_config(cfg, key, args, kwargs) + + # 4. Full search: benchmark every config, pick fastest, cache. configs + # may be a callable(*args) -> [Config] to build the space per shape. + configs = self.configs(*args, **kwargs) if callable(self.configs) else self.configs + configs = self._prune(configs, args, kwargs) print(f"[autotune] tuning {len(configs)} configs...") results = [] for i, config in enumerate(configs): try: - t = self._bench_one(config, args, kwargs) + t = self._bench_one(config, key, args, kwargs) results.append((config, t)) print(f" [{i+1}/{len(configs)}] {config} -> {t:.3f} ms") except Exception as e: @@ -349,27 +637,213 @@ def __call__(self, *args, **kwargs): print(f"[autotune] best: {best_config} ({best_time:.3f} ms)") self.cache[key] = best_config + self._tuned_keys.add(key) # only searched bests are persisted to disk self._save_disk_cache() - - return self._run_config(best_config, args, kwargs) + self._emit_offline_config(best_config, args, kwargs) + + return self._run_config(best_config, key, args, kwargs) + + # --- Offline config artifacts (aiter/SGLang model) --- + def _offline_spec(self, args, kwargs): + """The (axis, value) spec dict for this call, or None if offline mode is + off. Reuses key_fn so the offline key == the tuning key axes, and passes + values through JSON so the in-memory spec compares equal to the one + reloaded from disk (e.g. a tuple becomes a list on both sides).""" + if _offline_config_dir() is None or self.key_fn is None: + return None + try: + raw = dict(self.key_fn(*args, **kwargs)) + return json.loads(json.dumps(raw)) # normalize to JSON types + except Exception as e: + # Offline mode is configured but the spec isn't JSON-serializable + # (e.g. an Enum/torch.dtype axis) — warn rather than silently + # disabling the offline path for this op. + self._warn_once( + ("offline_spec", self.name), + f"[autotune] offline disabled for {self.name}: spec not JSON-serializable ({e})", + ) + return None + + def _offline_path(self, args, kwargs): + spec = self._offline_spec(args, kwargs) + if spec is None: + return None + cfg_dir = _offline_config_dir() + path = cfg_dir / offline_config_filename(self.name, spec.items()) + # Sanitized components can't escape, but confirm anyway. + if path.parent.resolve() != cfg_dir.resolve(): + return None + return cfg_dir, path, spec + + def _load_offline_config(self, args, kwargs): + """Load + validate a committed offline Config, or None on a genuine miss. + A file that exists but is corrupt or describes a different call is warned + about and ignored, so a broken artifact can't silently mis-serve.""" + resolved = self._offline_path(args, kwargs) + if resolved is None: + return None + _, path, spec = resolved + tok = str(path) + # exists()/read can raise OSError (e.g. ENAMETOOLONG for an over-long + # filename, which pathlib does NOT swallow) or JSON errors — treat any of + # them as a miss to fall back on, never an exception that escapes serving. + try: + if not path.exists(): + return None + data = json.loads(path.read_text()) + except Exception as e: + self._warn_once(tok, f"[autotune] offline config {path.name} unreadable ({e}); ignoring") + return None + if not isinstance(data, dict): + self._warn_once(tok, f"[autotune] offline config {path.name} is not an object; ignoring") + return None + required = ("name", "spec", "device_name", "config") + missing = [f for f in required if f not in data] + if missing: + self._warn_once(tok, f"[autotune] offline config {path.name} missing {missing}; ignoring") + return None + # Validate every axis against this call. Any malformed field (e.g. a + # non-dict spec) is a mismatch to ignore, never an exception to escape. + expected = {"name": self.name, "spec": spec, "device_name": _device_name()} + for field, want in expected.items(): + got = data[field] + if got != want: + self._warn_once(tok, f"[autotune] offline config {path.name} {field} {got!r}!={want!r}; ignoring") + return None + # The config body must be a non-empty dict and, in builder mode, carry + # the structural knobs — otherwise an empty/partial "config": {} would + # load a bare Config and silently run build_fn's *default* knob value. + cfg_body = data["config"] + if not isinstance(cfg_body, dict) or not cfg_body: + self._warn_once(tok, f"[autotune] offline config {path.name} has an empty/invalid config; ignoring") + return None + # Every config value must be a JSON leaf; a nested/non-leaf value (e.g. + # {"BLOCK": {}}) would otherwise reach the kernel and blow up at launch. + bad = [k for k, v in cfg_body.items() if not _is_json_scalar(v)] + if bad: + self._warn_once( + tok, f"[autotune] offline config {path.name} has non-scalar config value(s) {bad}; ignoring" + ) + return None + # Reject non-finite floats (json.loads parses NaN/Infinity by default); + # they pass the scalar check but are meaningless as a kernel constant. + nonfinite = [k for k, v in cfg_body.items() if isinstance(v, float) and not math.isfinite(v)] + if nonfinite: + self._warn_once( + tok, f"[autotune] offline config {path.name} has non-finite config value(s) {nonfinite}; ignoring" + ) + return None + if self.structural: + missing_knobs = [k for k in self.structural if k not in cfg_body] + if missing_knobs: + self._warn_once( + tok, f"[autotune] offline config {path.name} missing structural {missing_knobs}; ignoring" + ) + return None + # Type-check knob values against the heuristic default's config: the same + # knob must carry the same scalar type. A hand-edited value of the wrong + # type (e.g. "256" instead of 256) is a valid JSON scalar and passes + # every check above, but would reach build_fn and crash the serving call. + if self.default is not None: + try: + ref = self.default(*args, **kwargs).to_dict() + except Exception: + ref = {} + # bool is a subclass of int, so exact type() identity is intended. + mistyped = [k for k, v in cfg_body.items() if k in ref and type(v) is not type(ref[k])] + if mistyped: + self._warn_once( + tok, f"[autotune] offline config {path.name} has mistyped value(s) {mistyped}; ignoring" + ) + return None + try: + return Config.from_dict(cfg_body) + except Exception as e: + self._warn_once(tok, f"[autotune] offline config {path.name} malformed ({e}); ignoring") + return None + + def _emit_offline_config(self, config, args, kwargs): + """Write the tuned Config as a self-describing committed artifact.""" + resolved = self._offline_path(args, kwargs) + if resolved is None: + return + cfg_dir, path, spec = resolved + payload = { + "name": self.name, + "spec": spec, # already JSON-normalized by _offline_spec + "device_name": _device_name(), + "config": config.to_dict(), + } + # Detect a filename-sanitization collision: an existing artifact at this + # path that describes a DIFFERENT (name, spec, device) is about to be + # clobbered. Content re-validation keeps this from mis-serving, but the + # loser's config would be silently lost, so surface it. + try: + if path.exists(): + old = json.loads(path.read_text()) + if isinstance(old, dict) and ( + old.get("name"), + old.get("spec"), + old.get("device_name"), + ) != (payload["name"], payload["spec"], payload["device_name"]): + logger.warning( + "[autotune] offline config %s already exists for a different " + "spec (filename collision); overwriting", + path.name, + ) + except OSError: + pass + try: + _atomic_write_text(path, json.dumps(payload, indent=2, sort_keys=True)) + print(f"[autotune] wrote offline config {path.name}") + except Exception as e: + logger.warning("[autotune] failed to write offline config: %s", e) # --- Disk cache --- def _load_disk_cache(self): - if self._cache_file.exists(): + # Invalidate memoized offline/default decisions if the committed-config + # dir changed (a long-lived process may repoint FLYDSL_AUTOTUNE_CONFIG_DIR): + # offline configs are memoized with no other tie to their source dir. + # Tuned bests are independent of it, so they're kept. + _UNSET = "\0unset" + offline_dir = os.environ.get("FLYDSL_AUTOTUNE_CONFIG_DIR") + if getattr(self, "_loaded_config_dir", _UNSET) != offline_dir: + if getattr(self, "_loaded_config_dir", _UNSET) != _UNSET: + for k in [k for k in self.cache if k not in self._tuned_keys]: + del self.cache[k] + self._warned_once.clear() + self._loaded_config_dir = offline_dir + + # Re-load when the resolved path changes (FLYDSL_AUTOTUNE_CACHE_DIR may be + # set after a module-level tuner is constructed), so loads track the same + # dir that saves write to — not just the import-time default. Clear the + # in-memory cache too, or entries tuned under the old dir would be served + # after switching to a new (possibly empty) dir. + path = self._cache_file + if getattr(self, "_loaded_cache_path", None) == path: + return + if getattr(self, "_loaded_cache_path", None) is not None: + self.cache.clear() + self._tuned_keys.clear() + self._loaded_cache_path = path + if path.exists(): try: - data = json.loads(self._cache_file.read_text()) + data = json.loads(path.read_text()) for key_str, cfg_dict in data.items(): key = tuple(json.loads(key_str)) self.cache[key] = Config.from_dict(cfg_dict) + self._tuned_keys.add(key) # disk entries are tuned bests except Exception: pass def _save_disk_cache(self): - self._cache_file.parent.mkdir(parents=True, exist_ok=True) + # Persist only searched bests (not memoized defaults/offline configs), + # and write atomically so a concurrent reader never sees a torn file. data = {} for key, config in self.cache.items(): - data[json.dumps(list(key))] = config.to_dict() - self._cache_file.write_text(json.dumps(data, indent=2)) + if key in self._tuned_keys: + data[json.dumps(list(key))] = config.to_dict() + _atomic_write_text(self._cache_file, json.dumps(data, indent=2)) def autotune( @@ -383,16 +857,30 @@ def autotune( pre_hook: Callable = None, post_hook: Callable = None, do_bench: Callable = None, + build_fn: Callable = None, + default: Callable = None, ): """Autotune decorator for @jit functions. - Usage: + Direct mode (structural knobs are jit Constexpr params):: + @autotune(configs=[Config(BLOCK=128), Config(BLOCK=256)], key=['n']) @flyc.jit def myKernel(..., BLOCK: fx.Constexpr[int], ...): ... + For kernels whose structural knobs are baked at module-build time (as every + current FlyDSL kernel is), prefer ``autotune_builder`` — it wires build_fn / + default / configs / structural for you. The low-level ``build_fn`` / ``default`` + args below are the primitives it builds on. + Args: + build_fn: build_fn(config, *args, **kwargs) -> launch_callable. Enables + builder mode; built modules are cached per (key, config). + default: two-track heuristic default(*args, **kwargs) -> Config. When + set, normal runs use it and skip the search (zero-search); set + FLYDSL_AUTOTUNE=1 to force the full search. Without a default, every + uncached run searches. restore_value: tensor args the kernel mutates in place (output overlaps input, or accumulation). Snapshotted and restored before each bench rep so every config is measured on identical inputs. Required when @@ -414,6 +902,113 @@ def decorator(fn): pre_hook=pre_hook, post_hook=post_hook, do_bench_fn=do_bench, + build_fn=build_fn, + default=default, ) return decorator + + +def autotune_builder( + *, + name, + build, + specialize, + configs, + default=None, + structural=(), + warmup=10, + rep=50, + restore_value=None, + reset_to_zero=None, + do_bench_fn=None, +): + """One-call adoption for kernels whose knobs are baked at module-build time. + + You supply the kernel-specific pieces; this owns the Autotuner mechanics. + + rmsnorm_autotuned = autotune_builder( + name="rmsnorm", + build=build_rmsnorm_module, # build(**spec, **structural) -> launch fn + specialize=lambda inp, g, out, m, dtype_str="bf16", **kw: { + "N": inp.shape[-1], "dtype_str": dtype_str}, + configs=get_all_configs, # (**spec) -> [Config] + default=get_default, # (**spec) -> Config + structural=("BLOCK_THREADS",), # config kwargs routed into build() + ) + + Args: + name: artifact / disk-cache identity (required — builder tuners share a + cache dir, so each needs a distinct name). + build: build(**spec, **structural_knobs) -> launch callable. + specialize: specialize(*args, **kwargs) -> dict of the build/lookup axes + (shape + build-only scalars). Its items become the cache key AND the + offline filename, so it must (a) include every axis the best config + depends on — an omitted axis silently collides in the offline tree — + and (b) use JSON-scalar values (str/int/float/bool; tuples are ok, + stored as lists) so an emitted artifact matches on later lookup. + configs / default: called as configs(**spec) / default(**spec). + structural: config kwarg names passed to build() (vs compiler hints, + which flow through Config.compiler_opts()). + """ + if not name or not isinstance(name, str): + raise ValueError("autotune_builder requires a non-empty string name (the cache identity)") + structural = tuple(structural) + + def _spec(args, kwargs): + return specialize(*args, **kwargs) + + def _key_fn(*args, **kwargs): + return tuple(sorted(_spec(args, kwargs).items())) + + def _build_fn(config, *args, **kwargs): + spec = _spec(args, kwargs) + knobs = {k: config.kwargs[k] for k in structural if k in config.kwargs} + return build(**spec, **knobs) + + def _configs(*args, **kwargs): + return configs(**_spec(args, kwargs)) + + _default = None + if default is not None: + + def _default(*args, **kwargs): + return default(**_spec(args, kwargs)) + + # specialize's signature names the full call, so restore_value/reset_to_zero + # can look tensors up by name. + src = specialize.func if hasattr(specialize, "func") else specialize + sig = inspect.signature(src) + launch_arg_names = list(sig.parameters.keys()) + + tuner = Autotuner( + fn=None, + configs=_configs, + key=None, + warmup=warmup, + rep=rep, + restore_value=restore_value, + reset_to_zero=reset_to_zero, + build_fn=_build_fn, + default=_default, + name=name, + key_fn=_key_fn, + arg_names=launch_arg_names, + structural=structural, + do_bench_fn=do_bench_fn, + ) + + # A build-only scalar (a param that is also a spec key, e.g. dtype_str) must + # be keyword — positionally it would shift the launch args to the wrong slot. + @functools.wraps(src) + def call(*args, **kwargs): + leaked = [n for n in list(sig.parameters)[: len(args)] if n in _spec(args, kwargs)] + if leaked: + raise TypeError( + f"{name}: build-only arg(s) {leaked} must be passed by keyword, not positionally " + f"(they route to build/specialize, not the launch call)" + ) + return tuner(*args, **kwargs) + + call.tuner = tuner # expose for tests / introspection + return call diff --git a/tests/kernels/test_rmsnorm_autotune.py b/tests/kernels/test_rmsnorm_autotune.py new file mode 100644 index 000000000..32e6c8198 --- /dev/null +++ b/tests/kernels/test_rmsnorm_autotune.py @@ -0,0 +1,156 @@ +#!/usr/bin/env python3 + +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""GPU integration test for autotuned RMSNorm (first autotune_builder adopter, #770). + +Verifies the two-track builder-mode autotuner end to end: + - zero-search default run produces correct output + - forced-search (FLYDSL_AUTOTUNE=1) sweeps configs, picks one, stays correct + - the tuned result is cached (a second call does not re-tune) +""" + +import json +import os + +import pytest + +pytestmark = [pytest.mark.l2_device, pytest.mark.rocm_lower] + +try: + import torch +except ImportError: + torch = None +if torch is None or not torch.cuda.is_available(): + pytest.skip("CUDA/ROCm not available. Skipping GPU tests.", allow_module_level=True) + +from kernels.rmsnorm_autotune import rmsnorm_autotuned # noqa: E402 + +EPS = 1e-5 + + +@pytest.fixture(autouse=True) +def _fresh_cache(): + """Clear the tuner's in-memory cache so tuned results don't leak across + tests (the disk cache is isolated per test via FLYDSL_AUTOTUNE_CACHE_DIR).""" + rmsnorm_autotuned.tuner.cache.clear() + yield + rmsnorm_autotuned.tuner.cache.clear() + + +def _reference(x, g): + xf = x.float() + return (xf * torch.rsqrt((xf * xf).mean(-1, keepdim=True) + EPS)) * g.float() + + +def _run(M, N, autotune_env, tmp_cache): + os.environ["FLYDSL_AUTOTUNE_CACHE_DIR"] = str(tmp_cache) + if autotune_env: + os.environ["FLYDSL_AUTOTUNE"] = "1" + else: + os.environ.pop("FLYDSL_AUTOTUNE", None) + + torch.manual_seed(0) + x = torch.randn(M, N, device="cuda").to(torch.bfloat16) + g = torch.rand(N, device="cuda").to(torch.bfloat16) + ref = _reference(x, g) + s = torch.cuda.current_stream() + + out = torch.empty(M, N, device="cuda", dtype=torch.bfloat16) + rmsnorm_autotuned(x, g, out, M, dtype_str="bf16", stream=s) + torch.cuda.synchronize() + err = (out.float() - ref).abs().max().item() + return err, (x, g, ref, s) + + +def test_rmsnorm_autotuned_default(tmp_path): + """Zero-search default run is correct.""" + err, _ = _run(4096, 8192, autotune_env=False, tmp_cache=tmp_path) + assert err < 2e-2, f"default run max_err={err}" + + +def test_rmsnorm_autotuned_search_and_cache(tmp_path, monkeypatch): + """Forced search is correct, and a subsequent normal call does NOT re-tune + (it reuses the cached best) — proven by counting benchmark invocations.""" + err, (x, g, ref, s) = _run(4096, 8192, autotune_env=True, tmp_cache=tmp_path) + assert err < 2e-2, f"tuned run max_err={err}" + + # A tuned-config JSON must have been persisted. + files = list(tmp_path.glob("*.json")) + assert files, "no tuned-config cache file written" + + # Second call with tuning OFF: must serve from cache, no benchmarking. + monkeypatch.delenv("FLYDSL_AUTOTUNE", raising=False) + n_bench = {"n": 0} + orig = rmsnorm_autotuned.tuner._bench_one + + def counting_bench(*a, **k): + n_bench["n"] += 1 + return orig(*a, **k) + + monkeypatch.setattr(rmsnorm_autotuned.tuner, "_bench_one", counting_bench) + + out2 = torch.empty_like(ref, dtype=torch.bfloat16) + rmsnorm_autotuned(x, g, out2, x.shape[0], dtype_str="bf16", stream=s) + torch.cuda.synchronize() + err2 = (out2.float() - ref).abs().max().item() + assert err2 < 2e-2, f"cached run max_err={err2}" + assert n_bench["n"] == 0, "second call re-tuned instead of using the cache" + + +def test_rmsnorm_offline_emit_and_lookup(tmp_path, monkeypatch): + """Forced tune emits a committed offline artifact; a normal run then serves + from it (no search), proven by hooking the offline loader + benchmark.""" + M, N = 4096, 8192 + cfg_dir = tmp_path / "configs" + monkeypatch.setenv("FLYDSL_AUTOTUNE_CONFIG_DIR", str(cfg_dir)) + monkeypatch.setenv("FLYDSL_AUTOTUNE_CACHE_DIR", str(tmp_path / "cache")) + + torch.manual_seed(0) + x = torch.randn(M, N, device="cuda").to(torch.bfloat16) + g = torch.rand(N, device="cuda").to(torch.bfloat16) + ref = _reference(x, g) + s = torch.cuda.current_stream() + + # 1. Force a tune -> emits offline config. + monkeypatch.setenv("FLYDSL_AUTOTUNE", "1") + out = torch.empty(M, N, device="cuda", dtype=torch.bfloat16) + rmsnorm_autotuned(x, g, out, M, dtype_str="bf16", stream=s) + torch.cuda.synchronize() + emitted = list(cfg_dir.glob("rmsnorm,N=8192,dtype_str=bf16,device_name=*.json")) + assert emitted, f"no offline config emitted in {cfg_dir}" + + # 2. Tuning OFF, empty in-memory + fresh scratch cache -> must use offline. + monkeypatch.delenv("FLYDSL_AUTOTUNE", raising=False) + monkeypatch.setenv("FLYDSL_AUTOTUNE_CACHE_DIR", str(tmp_path / "cache_fresh")) + rmsnorm_autotuned.tuner.cache.clear() + got = {"cfg": None} + benched = {"n": 0} + orig_load, orig_bench = rmsnorm_autotuned.tuner._load_offline_config, rmsnorm_autotuned.tuner._bench_one + + def hooked_load(a, k): + cfg = orig_load(a, k) + got["cfg"] = cfg + return cfg + + monkeypatch.setattr(rmsnorm_autotuned.tuner, "_load_offline_config", hooked_load) + monkeypatch.setattr( + rmsnorm_autotuned.tuner, + "_bench_one", + lambda *a, **k: (benched.__setitem__("n", benched["n"] + 1), orig_bench(*a, **k))[1], + ) + out2 = torch.empty(M, N, device="cuda", dtype=torch.bfloat16) + rmsnorm_autotuned(x, g, out2, M, dtype_str="bf16", stream=s) + torch.cuda.synchronize() + assert (out2.float() - ref).abs().max().item() < 2e-2 + assert benched["n"] == 0 # no search + # Prove offline was the actual SOURCE (not a None->default fallback): the + # loader returned a real Config carrying the tuned BLOCK_THREADS. + assert got["cfg"] is not None and "BLOCK_THREADS" in got["cfg"].kwargs + emitted_block = json.loads(emitted[0].read_text())["config"]["BLOCK_THREADS"] + assert got["cfg"].kwargs["BLOCK_THREADS"] == emitted_block + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v", "-s"])) diff --git a/tests/unit/test_autotune.py b/tests/unit/test_autotune.py index a4ae8459d..13ad28704 100644 --- a/tests/unit/test_autotune.py +++ b/tests/unit/test_autotune.py @@ -16,10 +16,11 @@ """ import json +import logging import pytest -from flydsl.autotune import Autotuner, Config, _normalize_strides, autotune +from flydsl.autotune import Autotuner, Config, _normalize_strides, autotune, autotune_builder @pytest.fixture(autouse=True) @@ -29,6 +30,28 @@ def _isolate_disk_cache(tmp_path, monkeypatch): monkeypatch.setenv("FLYDSL_AUTOTUNE_CACHE_DIR", str(tmp_path / "autotune_cache")) +@pytest.fixture(autouse=True) +def _keep_flydsl_log_propagation(): + """Keep the ``flydsl`` logger propagating so ``caplog`` sees autotune warnings. + + ``flydsl.utils.logger._init_logger()`` sets + ``logging.getLogger("flydsl").propagate = False`` the first time anything calls + ``log()``, and it latches for the whole process. In the full test suite that + happens (via the jit/compiler modules) before these tests run, so + ``flydsl.autotune`` records no longer reach the root logger where ``caplog`` + listens and the offline-config warning asserts fail with an empty ``caplog``. + Run in isolation nothing flips it, which is why this only breaks in-suite and + is order-dependent. Force propagation on for each test and restore it after. + """ + flydsl_logger = logging.getLogger("flydsl") + saved = flydsl_logger.propagate + flydsl_logger.propagate = True + try: + yield + finally: + flydsl_logger.propagate = saved + + # ── Fakes ──────────────────────────────────────────────────────────────── class FakeTensor: """Minimal tensor stand-in with the attributes _make_key / restore_value use.""" @@ -361,5 +384,644 @@ def fake_jit(a, out, **kw): assert [c.kwargs["BLOCK"] for c in tuned.configs] == [128] +# ── builder mode + two-track default ───────────────────────────────────── +def _bench_run_all(call, warmup, rep): + # deterministic fake do_bench: run once, return a constant time + call() + return 1.0 + + +def test_builder_mode_rebuilds_per_config(): + """build_fn is called once per config; the returned fn is what runs.""" + built = [] + + def build_fn(config, a, out, **kw): + block = config.kwargs["BLOCK"] + built.append(block) + + def launch(a, out, **kw): + out._data[0] = float(block) # record which build ran + + return launch + + t = Autotuner( + fn=None, + configs=[Config(BLOCK=64), Config(BLOCK=128)], + key=["a"], + warmup=1, + rep=1, + build_fn=build_fn, + do_bench_fn=_bench_run_all, + ) + a = FakeTensor((8,)) + out = FakeTensor((1,)) + t(a, out) + # both configs built + benchmarked exactly once + assert sorted(built) == [64, 128] + # arg_names inferred from build_fn with leading 'config' stripped + assert t.arg_names[:2] == ["a", "out"] + + +def test_builder_mode_caches_built_modules(): + """Re-running the same (key, config) does not rebuild.""" + n_builds = {"n": 0} + + def build_fn(config, a, out, **kw): + n_builds["n"] += 1 + return lambda a, out, **kw: None + + t = Autotuner( + fn=None, + configs=[Config(BLOCK=64)], + key=["a"], + warmup=1, + rep=1, + build_fn=build_fn, + do_bench_fn=_bench_run_all, + ) + a, out = FakeTensor((8,)), FakeTensor((1,)) + t(a, out) # tune: builds once + t(a, out) # cached best: reuses build + assert n_builds["n"] == 1 + + +def test_default_skips_search(monkeypatch): + """With a default heuristic and FLYDSL_AUTOTUNE off, no benchmarking runs.""" + monkeypatch.delenv("FLYDSL_AUTOTUNE", raising=False) + benched = {"n": 0} + + def bench(call, warmup, rep): + benched["n"] += 1 + call() + return 1.0 + + ran = {"block": None} + + def build_fn(config, a, out, **kw): + return lambda a, out, **kw: ran.__setitem__("block", config.kwargs["BLOCK"]) + + def default(a, out, **kw): + return Config(BLOCK=999) + + t = Autotuner( + fn=None, + configs=[Config(BLOCK=64), Config(BLOCK=128)], + key=["a"], + warmup=1, + rep=1, + build_fn=build_fn, + default=default, + do_bench_fn=bench, + ) + t(FakeTensor((8,)), FakeTensor((1,))) + assert benched["n"] == 0 # no search + assert ran["block"] == 999 # heuristic default was used + + +def test_default_forced_search_with_env(monkeypatch): + """FLYDSL_AUTOTUNE=1 forces the full search even when a default exists.""" + monkeypatch.setenv("FLYDSL_AUTOTUNE", "1") + benched = {"n": 0} + + def bench(call, warmup, rep): + benched["n"] += 1 + call() + return 1.0 + + t = Autotuner( + fn=None, + configs=[Config(BLOCK=64), Config(BLOCK=128)], + key=["a"], + warmup=1, + rep=1, + build_fn=build_fn_noop, + default=lambda a, out, **kw: Config(BLOCK=64), + do_bench_fn=bench, + ) + t(FakeTensor((8,)), FakeTensor((1,))) + assert benched["n"] == 2 # both configs searched + + +def build_fn_noop(config, a, out, **kw): + return lambda a, out, **kw: None + + +def test_filter_call_kwargs_drops_build_only_kwargs(): + """Builder-only kwargs (e.g. dtype_str) must not reach the launch fn.""" + + def build_fn(config, a, out, dtype_str="bf16", **kw): + # launch fn only accepts a, out — not dtype_str + def launch(a, out): + pass + + return launch + + t = Autotuner( + fn=None, + configs=[Config(BLOCK=64)], + key=["a"], + warmup=1, + rep=1, + build_fn=build_fn, + default=lambda a, out, **kw: Config(BLOCK=64), + do_bench_fn=_bench_run_all, + ) + # dtype_str would raise TypeError if not filtered out before the launch call + t(FakeTensor((8,)), FakeTensor((1,)), dtype_str="f16") + + +def test_tuning_enabled_env(monkeypatch): + from flydsl.autotune import _tuning_enabled + + monkeypatch.delenv("FLYDSL_AUTOTUNE", raising=False) + assert _tuning_enabled() is False + monkeypatch.setenv("FLYDSL_AUTOTUNE", "0") + assert _tuning_enabled() is False + monkeypatch.setenv("FLYDSL_AUTOTUNE", "1") + assert _tuning_enabled() is True + + +# ── autotune_builder (one-call adoption) ───────────────────────────────── +def _make_builder(**over): + """A fake-kernel autotune_builder: build() records which config ran into the + output tensor's [0] slot; specialize keys on N + dtype_str.""" + built = over.pop("_built_log", []) + + def build(N, dtype_str, BLOCK=0): + built.append((N, dtype_str, BLOCK)) + + def launch(a, out, dtype_str="bf16", stream=None): + out._data[0] = float(BLOCK) + + return launch + + def specialize(a, out, dtype_str="bf16", stream=None): + return {"N": a.shape[-1], "dtype_str": dtype_str} + + kw = dict( + name="fakeop", + build=build, + specialize=specialize, + configs=lambda N, dtype_str: [Config(BLOCK=64), Config(BLOCK=128)], + default=lambda N, dtype_str: Config(BLOCK=7), + structural=("BLOCK",), + warmup=1, + rep=1, + ) + kw.update(over) + t = autotune_builder(**kw) + return t, built + + +def test_builder_default_runs_without_search(monkeypatch): + monkeypatch.delenv("FLYDSL_AUTOTUNE", raising=False) + t, built = _make_builder() + a, out = FakeTensor((16, 512)), FakeTensor((1,)) + t(a, out, dtype_str="bf16") + assert out._data[0] == 7.0 # heuristic default's BLOCK + assert built == [(512, "bf16", 7)] # built once, from the default + + +def test_builder_search_sweeps_space(monkeypatch, tmp_path): + monkeypatch.setenv("FLYDSL_AUTOTUNE_CACHE_DIR", str(tmp_path)) + monkeypatch.setenv("FLYDSL_AUTOTUNE", "1") + t, built = _make_builder(do_bench_fn=lambda call, warmup, rep: (call(), 1.0)[1]) + a, out = FakeTensor((16, 512)), FakeTensor((1,)) + t(a, out, dtype_str="bf16") + swept = sorted(b for _, _, b in built) + assert swept == [64, 128] # both configs from get_all_configs were built + + +def test_builder_scalar_enters_key(): + """dtype_str is build-only but must change the cache key (bug: scalars that + change codegen without changing shape were dropped from the key).""" + t, _ = _make_builder() + a = FakeTensor((16, 512)) + k_bf16 = t.tuner._make_key((), {"a": a, "out": FakeTensor((1,)), "dtype_str": "bf16"}) + k_f16 = t.tuner._make_key((), {"a": a, "out": FakeTensor((1,)), "dtype_str": "f16"}) + assert k_bf16 != k_f16 + + +def test_builder_requires_name(): + """Builder mode must carry a distinct cache name (else tuners collide on + unknown.json).""" + t, _ = _make_builder(name="softmax") + assert t.tuner.name == "softmax" + assert t.tuner._cache_file.name == "softmax.json" + # empty / missing name must be rejected (else tuners collide on unknown.json) + with pytest.raises(ValueError): + _make_builder(name="") + + +def test_builder_positional_scalar_rejected(monkeypatch): + """A build-only scalar (dtype_str) passed positionally is rejected up front + (codex#1: silently binding it to the wrong launch slot is worse).""" + monkeypatch.delenv("FLYDSL_AUTOTUNE", raising=False) + + def build(N, dtype_str, BLOCK=0): + return lambda a, out, stream=None: None + + def specialize(a, out, dtype_str="bf16", stream=None): + return {"N": a.shape[-1], "dtype_str": dtype_str} + + t = autotune_builder( + name="op", + build=build, + specialize=specialize, + configs=lambda N, dtype_str: [Config(BLOCK=1)], + default=lambda N, dtype_str: Config(BLOCK=1), + structural=("BLOCK",), + warmup=1, + rep=1, + ) + # dtype_str keyword: fine. + t(FakeTensor((16, 512)), FakeTensor((1,)), dtype_str="f16") + # dtype_str positional (5th arg): rejected with a clear error. + with pytest.raises(TypeError, match="by keyword"): + t(FakeTensor((16, 512)), FakeTensor((1,)), "f16") + + +def test_builder_build_cache_ignores_compiler_hints(monkeypatch, tmp_path): + """configs differing only in waves_per_eu must build the module once, not + once per hint (C1: build cache was over-keyed on repr(config)).""" + monkeypatch.setenv("FLYDSL_AUTOTUNE_CACHE_DIR", str(tmp_path)) + monkeypatch.setenv("FLYDSL_AUTOTUNE", "1") + n_build = {"n": 0} + + def build(N, dtype_str, BLOCK=0): + n_build["n"] += 1 + return lambda a, out, stream=None: None + + def specialize(a, out, dtype_str="bf16", stream=None): + return {"N": a.shape[-1], "dtype_str": dtype_str} + + # 3 configs, same BLOCK, differing only in waves_per_eu. + space = [Config(BLOCK=64, waves_per_eu=w) for w in (None, 1, 2)] + t = autotune_builder( + name="op", + build=build, + specialize=specialize, + configs=lambda N, dtype_str: space, + structural=("BLOCK",), + warmup=1, + rep=1, + do_bench_fn=lambda call, warmup, rep: (call(), 1.0)[1], + ) + t(FakeTensor((16, 512)), FakeTensor((1,)), dtype_str="bf16") + assert n_build["n"] == 1, f"rebuilt {n_build['n']}x for hint-only variants (should be 1)" + + +def test_run_with_hints_sets_and_restores_compile_hints(): + """_run_with_hints must set fn.compile_hints during the call (so the hint + enters the JIT cache key) and restore it after (no cross-config leak).""" + + class FakeJit: + def __init__(self): + self.compile_hints = {} + self.seen = None + + def __call__(self, *a, **k): + self.seen = dict(self.compile_hints) + + # No compiler bindings needed: with hints present, _run_with_hints imports + # CompilationContext, so skip when the compiled bindings are absent. + pytest.importorskip("flydsl._mlir._mlir_libs._mlirDialectsFly") + fn = FakeJit() + t = _make_tuner(fn=lambda a, out, **kw: None, configs=[Config(BLOCK=1)]) + t._run_with_hints(fn, Config(BLOCK=1, waves_per_eu=2).compiler_opts(), (), {}) + assert fn.seen == {"waves_per_eu": 2} # in the cache-key dict during compile + assert fn.compile_hints == {} # restored afterward + + +def test_cache_dir_change_does_not_serve_stale_config(tmp_path, monkeypatch): + """Switching FLYDSL_AUTOTUNE_CACHE_DIR must drop the in-memory config tuned + under the old dir. The fake tune picks BLOCK=64; the default is BLOCK=7. + After switching to an empty dir with tuning OFF, the call must fall to the + default (7) — proving the stale dir-A best (64) was cleared, not served.""" + monkeypatch.setenv("FLYDSL_AUTOTUNE_CACHE_DIR", str(tmp_path / "A")) + + # 1. Force a tune into dir A -> in-memory best BLOCK=64. + monkeypatch.setenv("FLYDSL_AUTOTUNE", "1") + t, _ = _make_builder(do_bench_fn=lambda call, warmup, rep: (call(), 1.0)[1]) + out = FakeTensor((1,)) + t(FakeTensor((16, 512)), out, dtype_str="bf16") + assert out._data[0] == 64.0 # tuned config in memory + + # 2. Switch to empty dir B, tuning OFF: stale in-memory best must be dropped, + # so this serves the heuristic default (7), not dir-A's 64. + monkeypatch.delenv("FLYDSL_AUTOTUNE", raising=False) + monkeypatch.setenv("FLYDSL_AUTOTUNE_CACHE_DIR", str(tmp_path / "B")) + out2 = FakeTensor((1,)) + t(FakeTensor((16, 512)), out2, dtype_str="bf16") + assert out2._data[0] == 7.0, "served a stale in-memory config from the old cache dir" + + +# ── offline artifacts (aiter/SGLang model) ─────────────────────────────── +def test_offline_filename_is_key_and_sanitized(): + from flydsl.autotune import offline_config_filename + + name = offline_config_filename("op", [("N", 8192), ("dtype_str", "bf16")], device_name="AMD MI300X") + assert name == "op,N=8192,dtype_str=bf16,device_name=AMD_MI300X.json" + # injection is neutralized: no raw delimiters / path separators survive + evil = offline_config_filename("../x", [("N", "1,dtype_str=y")], device_name="a/b") + assert "/" not in evil and evil.count(".json") == 1 + + +def _write_offline(cfg_dir, spec_items, config, name="fakeop", device=None, raw=None): + """Write an offline artifact; returns its path. `raw` overrides the payload.""" + from flydsl.autotune import _device_name, offline_config_filename + + cfg_dir.mkdir(parents=True, exist_ok=True) + device = device or _device_name() + path = cfg_dir / offline_config_filename(name, spec_items, device_name=device) + payload = ( + raw + if raw is not None + else { + "name": name, + "spec": dict(spec_items), + "device_name": device, + "config": config, + } + ) + path.write_text(json.dumps(payload)) + return path + + +def test_offline_emit_then_served_is_the_offline_config(tmp_path, monkeypatch): + """Prove the OFFLINE config is what serves — not the heuristic default. + + The fake build writes the chosen BLOCK into out[0]. Forced tune picks BLOCK=64 + (first config, equal timings); the default is BLOCK=7. A fresh tuner with + tuning OFF and a *separate* scratch cache dir must produce 64, so we know the + value came from the committed artifact, not the default and not disk cache.""" + monkeypatch.setenv("FLYDSL_AUTOTUNE_CONFIG_DIR", str(tmp_path / "cfg")) + + # 1. Force a tune -> emits the artifact (BLOCK=64). + monkeypatch.setenv("FLYDSL_AUTOTUNE", "1") + monkeypatch.setenv("FLYDSL_AUTOTUNE_CACHE_DIR", str(tmp_path / "scratch1")) + t, _ = _make_builder(do_bench_fn=lambda call, warmup, rep: (call(), 1.0)[1]) + t(FakeTensor((16, 512)), FakeTensor((1,)), dtype_str="bf16") + payload = json.loads(next((tmp_path / "cfg").glob("*.json")).read_text()) + assert payload["spec"] == {"N": 512, "dtype_str": "bf16"} + assert payload["config"]["BLOCK"] == 64 + + # 2. Fresh tuner, tuning OFF, SEPARATE scratch dir (so a disk-cache hit can't + # masquerade as an offline hit), and no benchmarking allowed. + monkeypatch.delenv("FLYDSL_AUTOTUNE", raising=False) + monkeypatch.setenv("FLYDSL_AUTOTUNE_CACHE_DIR", str(tmp_path / "scratch2")) + t2, _ = _make_builder(do_bench_fn=lambda call, warmup, rep: (call(), 1.0)[1]) + n = {"bench": 0} + orig = t2.tuner._bench_one + t2.tuner._bench_one = lambda *a, **k: (n.__setitem__("bench", n["bench"] + 1), orig(*a, **k))[1] + out = FakeTensor((1,)) + t2(FakeTensor((16, 512)), out, dtype_str="bf16") + assert n["bench"] == 0 # no search + assert out._data[0] == 64.0 # served the OFFLINE config (64), not default (7) + + +def test_offline_stale_artifact_rejected(tmp_path, monkeypatch, caplog): + caplog.set_level(logging.WARNING, logger="flydsl.autotune") + monkeypatch.setenv("FLYDSL_AUTOTUNE_CONFIG_DIR", str(tmp_path / "cfg")) + monkeypatch.delenv("FLYDSL_AUTOTUNE", raising=False) + # filename matches N=512, but content claims N=4096 (stale / hand-edited) + _write_offline( + tmp_path / "cfg", + [("N", 512), ("dtype_str", "bf16")], + {"BLOCK": 1}, + raw={ + "name": "fakeop", + "spec": {"N": 4096, "dtype_str": "bf16"}, + "device_name": __import__("flydsl.autotune", fromlist=["_device_name"])._device_name(), + "config": {"BLOCK": 1}, + }, + ) + t, _ = _make_builder(do_bench_fn=lambda call, warmup, rep: (call(), 1.0)[1]) + out = FakeTensor((1,)) + t(FakeTensor((16, 512)), out, dtype_str="bf16") + assert "ignoring" in caplog.text # rejected + assert out._data[0] == 7.0 # fell back to default (BLOCK=7), not the stale config + + +def test_offline_malformed_spec_does_not_crash(tmp_path, monkeypatch, caplog): + """spec=null (present but not a dict) must warn + fall back, never raise.""" + caplog.set_level(logging.WARNING, logger="flydsl.autotune") + monkeypatch.setenv("FLYDSL_AUTOTUNE_CONFIG_DIR", str(tmp_path / "cfg")) + monkeypatch.delenv("FLYDSL_AUTOTUNE", raising=False) + _write_offline( + tmp_path / "cfg", + [("N", 512), ("dtype_str", "bf16")], + None, + raw={ + "name": "fakeop", + "spec": None, + "device_name": __import__("flydsl.autotune", fromlist=["_device_name"])._device_name(), + "config": {"BLOCK": 1}, + }, + ) + t, _ = _make_builder(do_bench_fn=lambda call, warmup, rep: (call(), 1.0)[1]) + out = FakeTensor((1,)) + t(FakeTensor((16, 512)), out, dtype_str="bf16") # must not raise + assert "ignoring" in caplog.text + assert out._data[0] == 7.0 # default fallback + + +def test_offline_empty_config_rejected(tmp_path, monkeypatch, caplog): + """A matching artifact with an empty/partial config must be rejected, not + served (else builder mode silently runs build_fn's default structural knob). + The fake build's structural knob is BLOCK; an empty config lacks it.""" + caplog.set_level(logging.WARNING, logger="flydsl.autotune") + monkeypatch.setenv("FLYDSL_AUTOTUNE_CONFIG_DIR", str(tmp_path / "cfg")) + monkeypatch.delenv("FLYDSL_AUTOTUNE", raising=False) + _write_offline( + tmp_path / "cfg", + [("N", 512), ("dtype_str", "bf16")], + {}, # empty config body + raw={ + "name": "fakeop", + "spec": {"N": 512, "dtype_str": "bf16"}, + "device_name": __import__("flydsl.autotune", fromlist=["_device_name"])._device_name(), + "config": {}, + }, + ) + t, _ = _make_builder(do_bench_fn=lambda call, warmup, rep: (call(), 1.0)[1]) + out = FakeTensor((1,)) + t(FakeTensor((16, 512)), out, dtype_str="bf16") + assert "ignoring" in caplog.text + assert out._data[0] == 7.0 # fell back to default (BLOCK=7), not the empty artifact + + +@pytest.mark.parametrize("bad_value", [{}, [1, 2], None]) +def test_offline_non_scalar_config_value_rejected(tmp_path, monkeypatch, caplog, bad_value): + """A config knob whose value is not a scalar (dict, list, or null) must be + rejected + fall back. A list would be unhashable in the build-cache key and + crash before fallback if it slipped through — config values must be scalar + even though spec axes may be lists.""" + caplog.set_level(logging.WARNING, logger="flydsl.autotune") + monkeypatch.setenv("FLYDSL_AUTOTUNE_CONFIG_DIR", str(tmp_path / "cfg")) + monkeypatch.delenv("FLYDSL_AUTOTUNE", raising=False) + _write_offline( + tmp_path / "cfg", + [("N", 512), ("dtype_str", "bf16")], + {"BLOCK": bad_value}, + raw={ + "name": "fakeop", + "spec": {"N": 512, "dtype_str": "bf16"}, + "device_name": __import__("flydsl.autotune", fromlist=["_device_name"])._device_name(), + "config": {"BLOCK": bad_value}, + }, + ) + t, _ = _make_builder(do_bench_fn=lambda call, warmup, rep: (call(), 1.0)[1]) + out = FakeTensor((1,)) + t(FakeTensor((16, 512)), out, dtype_str="bf16") # must not raise + assert "non-scalar" in caplog.text + assert out._data[0] == 7.0 # default fallback + + +def test_offline_corrupt_and_missing_fields(tmp_path, monkeypatch, caplog): + from flydsl.autotune import _device_name, offline_config_filename + + caplog.set_level(logging.WARNING, logger="flydsl.autotune") + cfg = tmp_path / "cfg" + cfg.mkdir() + monkeypatch.setenv("FLYDSL_AUTOTUNE_CONFIG_DIR", str(cfg)) + monkeypatch.delenv("FLYDSL_AUTOTUNE", raising=False) + fname = offline_config_filename("fakeop", [("N", 512), ("dtype_str", "bf16")], device_name=_device_name()) + # unparseable JSON + (cfg / fname).write_text("{ not json") + t, _ = _make_builder(do_bench_fn=lambda call, warmup, rep: (call(), 1.0)[1]) + t(FakeTensor((16, 512)), FakeTensor((1,)), dtype_str="bf16") + assert "unreadable" in caplog.text + # present file, missing required fields + caplog.clear() + (cfg / fname).write_text(json.dumps({"config": {"BLOCK": 1}})) + t2, _ = _make_builder(do_bench_fn=lambda call, warmup, rep: (call(), 1.0)[1]) + t2(FakeTensor((16, 512)), FakeTensor((1,)), dtype_str="bf16") + assert "missing" in caplog.text + + +def test_offline_tuple_spec_round_trips(tmp_path, monkeypatch): + """A spec value that JSON turns into a list (tuple) must still self-match on + lookup — the artifact must not reject its own emitted config.""" + + def build(N, tile, BLOCK=0): + return lambda a, out, stream=None: out._data.__setitem__(0, float(BLOCK)) + + def specialize(a, out, stream=None): + return {"N": a.shape[-1], "tile": (16, 32)} + + monkeypatch.setenv("FLYDSL_AUTOTUNE_CONFIG_DIR", str(tmp_path / "cfg")) + monkeypatch.setenv("FLYDSL_AUTOTUNE", "1") + monkeypatch.setenv("FLYDSL_AUTOTUNE_CACHE_DIR", str(tmp_path / "s1")) + mk = dict( + name="op", + build=build, + specialize=specialize, + configs=lambda N, tile: [Config(BLOCK=5)], + default=lambda N, tile: Config(BLOCK=9), + structural=("BLOCK",), + warmup=1, + rep=1, + do_bench_fn=lambda call, warmup, rep: (call(), 1.0)[1], + ) + autotune_builder(**mk)(FakeTensor((8, 512)), FakeTensor((1,))) + monkeypatch.delenv("FLYDSL_AUTOTUNE", raising=False) + monkeypatch.setenv("FLYDSL_AUTOTUNE_CACHE_DIR", str(tmp_path / "s2")) + out = FakeTensor((1,)) + autotune_builder(**mk)(FakeTensor((8, 512)), out) + assert out._data[0] == 5.0 # offline (5) served, not default (9) + + +def test_offline_mistyped_scalar_rejected(tmp_path, monkeypatch, caplog): + """A knob value that is a valid JSON scalar but the WRONG type (e.g. "64" + instead of 64) passes every structural check yet would reach build_fn and + crash serving. It must be rejected against the default's type and fall back.""" + caplog.set_level(logging.WARNING, logger="flydsl.autotune") + monkeypatch.setenv("FLYDSL_AUTOTUNE_CONFIG_DIR", str(tmp_path / "cfg")) + monkeypatch.delenv("FLYDSL_AUTOTUNE", raising=False) + _write_offline( + tmp_path / "cfg", + [("N", 512), ("dtype_str", "bf16")], + {"BLOCK": "64"}, + raw={ + "name": "fakeop", + "spec": {"N": 512, "dtype_str": "bf16"}, + "device_name": __import__("flydsl.autotune", fromlist=["_device_name"])._device_name(), + "config": {"BLOCK": "64"}, # str, but the default's BLOCK is int + }, + ) + t, _ = _make_builder(do_bench_fn=lambda call, warmup, rep: (call(), 1.0)[1]) + out = FakeTensor((1,)) + t(FakeTensor((16, 512)), out, dtype_str="bf16") # must not crash + assert "mistyped" in caplog.text + assert out._data[0] == 7.0 # default fallback (7), not the mistyped artifact + + +def test_offline_reject_does_not_reread_or_respam(tmp_path, monkeypatch, caplog): + """A present-but-rejected artifact must be read + warned about once, not on + every call: the default is memoized so later calls short-circuit (no per-call + stat/parse/warn on the serving hot path).""" + caplog.set_level(logging.WARNING, logger="flydsl.autotune") + cfg_dir = tmp_path / "cfg" + monkeypatch.setenv("FLYDSL_AUTOTUNE_CONFIG_DIR", str(cfg_dir)) + monkeypatch.delenv("FLYDSL_AUTOTUNE", raising=False) + _write_offline( + cfg_dir, + [("N", 512), ("dtype_str", "bf16")], + {"BLOCK": 1}, + raw={ # content spec != filename spec -> rejected as stale + "name": "fakeop", + "spec": {"N": 4096, "dtype_str": "bf16"}, + "device_name": __import__("flydsl.autotune", fromlist=["_device_name"])._device_name(), + "config": {"BLOCK": 1}, + }, + ) + t, _ = _make_builder(do_bench_fn=lambda call, warmup, rep: (call(), 1.0)[1]) + reads = {"n": 0} + orig = t.tuner._load_offline_config + t.tuner._load_offline_config = lambda *a, **k: (reads.__setitem__("n", reads["n"] + 1), orig(*a, **k))[1] + for _ in range(5): + out = FakeTensor((1,)) + t(FakeTensor((16, 512)), out, dtype_str="bf16") + assert out._data[0] == 7.0 # default fallback every call + assert reads["n"] == 1, "offline artifact re-read every call (no negative cache)" + assert caplog.text.count("ignoring") == 1, "rejected artifact re-warned every call" + + +def test_rmsnorm_adopter_imports_resolve(): + """Regression guard for the RMSNorm adopter's kernel import path. The GPU + integration test skips before importing the adopter on a non-GPU runner, so a + wrong path (kernels.rmsnorm_kernel, which does not exist) hides there. This + GPU-free check reads the source so it runs on every CI runner.""" + import pathlib + + import kernels + + root = pathlib.Path(kernels.__file__).parent + for fname in ("rmsnorm_autotune.py", "rmsnorm_config.py"): + src = (root / fname).read_text() + assert "from kernels.rmsnorm_kernel import" not in src, f"{fname} imports non-existent kernels.rmsnorm_kernel" + assert "kernels.norm.rmsnorm_kernel" in src, f"{fname} must import from kernels.norm.rmsnorm_kernel" + + +def test_offline_force_bypasses_and_reemits(tmp_path, monkeypatch): + """FLYDSL_AUTOTUNE=1 must ignore an existing offline artifact, re-benchmark, + and re-emit — not short-circuit on the committed config.""" + monkeypatch.setenv("FLYDSL_AUTOTUNE_CONFIG_DIR", str(tmp_path / "cfg")) + monkeypatch.setenv("FLYDSL_AUTOTUNE_CACHE_DIR", str(tmp_path / "s1")) + monkeypatch.setenv("FLYDSL_AUTOTUNE", "1") + t, _ = _make_builder(do_bench_fn=lambda call, warmup, rep: (call(), 1.0)[1]) + t(FakeTensor((16, 512)), FakeTensor((1,)), dtype_str="bf16") # emits artifact + + # Force again with a fresh scratch dir: must still benchmark (not serve offline). + monkeypatch.setenv("FLYDSL_AUTOTUNE_CACHE_DIR", str(tmp_path / "s2")) + t2, _ = _make_builder(do_bench_fn=lambda call, warmup, rep: (call(), 1.0)[1]) + n = {"bench": 0} + orig = t2.tuner._bench_one + t2.tuner._bench_one = lambda *a, **k: (n.__setitem__("bench", n["bench"] + 1), orig(*a, **k))[1] + t2(FakeTensor((16, 512)), FakeTensor((1,)), dtype_str="bf16") + assert n["bench"] > 0 # force bypassed the offline artifact + + if __name__ == "__main__": raise SystemExit(pytest.main([__file__, "-v"]))