quack-kernels 0.1.10__py3-none-any.whl → 0.2.0__py3-none-any.whl

quack/__init__.py

@@ -1,9 +1,16 @@
-__version__ = "0.1.10"
+__version__ = "0.2.0"
+
+import cutlass.cute as cute
 
 from quack.rmsnorm import rmsnorm
 from quack.softmax import softmax
 from quack.cross_entropy import cross_entropy
 
+import quack.cute_dsl_utils
+
+# Patch cute.compile to optionally dump SASS
+cute.compile = quack.cute_dsl_utils.cute_compile_patched
+
 __all__ = [
     "rmsnorm",
     "softmax",

quack/activation.py

@@ -0,0 +1,288 @@
+# Copyright (c) 2025, Tri Dao.
+
+import math
+from typing import Tuple
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Float32
+from cutlass.cutlass_dsl import T, dsl_user_op
+from cutlass._mlir.dialects import llvm
+
+
+@dsl_user_op
+def tanh(a: float | Float32, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        llvm.inline_asm(
+            T.f32(),
+            [Float32(a).ir_value(loc=loc, ip=ip)],
+            "tanh.approx.f32 $0, $1;",
+            "=f,f",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+
+
+@dsl_user_op
+def relu(x: Float32, *, loc=None, ip=None) -> Float32:
+    return cute.arch.fmax(x, Float32(0.0))
+
+
+@cute.jit
+@dsl_user_op
+def drelu(x: Float32, dout: Float32, *, loc=None, ip=None) -> Tuple[Float32, Float32]:
+    x_pos = cutlass.Boolean(x > 0)
+    return dout if x_pos else Float32(0.0), cute.arch.fmax(x, Float32(0.0))
+
+
+@dsl_user_op
+def relu_sq(x: Float32, *, loc=None, ip=None) -> Float32:
+    return cute.arch.fmax(x, Float32(0.0)) * x
+
+
+@cute.jit
+@dsl_user_op
+def drelu_sq(x: Float32, dout: Float32, *, loc=None, ip=None) -> Tuple[Float32, Float32]:
+    """
+    ReLU squared backward pass: computes gradient w.r.t. x and recomputes forward
+    Given: relu_sq_out = max(x, 0) * x, and dout = grad w.r.t. relu_sq_out
+    Returns: (dx, relu_sq_out) where:
+    - dx = dout * 2 * x if x > 0, else 0
+    - relu_sq_out = max(x, 0) * x
+    """
+    x_pos = cutlass.Boolean(x > 0)
+    relu_sq_out = cute.arch.fmax(x, Float32(0.0)) * x
+    # Derivative: d/dx[max(x,0) * x] = 2*x if x > 0, else 0
+    dx = (2.0 * dout * x) if x_pos else Float32(0.0)
+    return dx, relu_sq_out
+
+
+@dsl_user_op
+def gelu_tanh_approx(x: Float32, *, loc=None, ip=None) -> Float32:
+    """
+    gelu(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
+            = 0.5 * x * (1 + tanh(x * (0.797885 + 0.0356774 * x * x)))
+    """
+    sqrt_2_over_pi = math.sqrt(2 / math.pi)  # ~0.797885
+    sqrt_2_over_pi_coeff = 0.044715 * sqrt_2_over_pi  # ~0.0356774
+    return 0.5 * (x * (1 + tanh(x * (sqrt_2_over_pi + sqrt_2_over_pi_coeff * (x * x)))))
+
+
+@dsl_user_op
+def dgelu_tanh_approx(x: Float32, dout: Float32, *, loc=None, ip=None) -> Tuple[Float32, Float32]:
+    """
+    GELU tanh approximation backward pass: computes gradient w.r.t. x and recomputes forward
+    Given: gelu_out = 0.5 * x * (1 + tanh(x * (c1 + c2 * x^2))), and dout = grad w.r.t. gelu_out
+    Returns: (dx, gelu_out)
+
+    Derivative uses the chain rule:
+    d/dx[gelu(x)] = 0.5 * (1 + tanh(z)) + 0.5 * x * sech^2(z) * dz/dx
+    where z = x * (c1 + c2 * x^2), dz/dx = c1 + 3 * c2 * x^2
+    and sech^2(z) = 1 - tanh^2(z)
+    """
+    sqrt_2_over_pi = math.sqrt(2 / math.pi)  # c1 ~0.797885
+    sqrt_2_over_pi_coeff = 0.044715 * sqrt_2_over_pi  # c2 ~0.0356774
+    sqrt_2_over_pi_coeff_3 = 3.0 * sqrt_2_over_pi_coeff  # c3 ~0.01070322
+
+    # Compute z = x * (c1 + c2 * x^2)
+    x_sq = x * x
+    tanh_z = tanh(x * (sqrt_2_over_pi + sqrt_2_over_pi_coeff * x_sq))
+    half_tanh_z_plus_one = 0.5 + 0.5 * tanh_z
+    gelu_out = x * half_tanh_z_plus_one
+
+    # Compute gradient
+    # sech^2(z) = 1 - tanh^2(z)
+    sech2_z = 1 - tanh_z * tanh_z
+    # dz/dx = c1 + 3 * c2 * x^2
+    dz_dx = sqrt_2_over_pi + sqrt_2_over_pi_coeff_3 * x_sq
+    # d/dx[gelu(x)] = 0.5 * (1 + tanh(z)) + 0.5 * x * sech^2(z) * dz/dx
+    dgelu = half_tanh_z_plus_one + x * (0.5 * (sech2_z * dz_dx))
+
+    dx = dout * dgelu
+    return dx, gelu_out
+
+
+@dsl_user_op
+def silu(x: Float32, *, loc=None, ip=None) -> Float32:
+    """
+    silu(x) = x * sigmoid(x) = x * (1 + tanh(x / 2)) / 2 = (0.5 * x) * tanh(0.5 * x) + (0.5 * x)
+    This compiles down to 3 SASS instructions: FMUL to get 0.5 * x, MUFU.TANH, and FFMA.
+    """
+    x_half = 0.5 * x
+    return x_half * tanh(x_half) + x_half
+
+
+@dsl_user_op
+def swiglu(x: Float32, y: Float32, *, loc=None, ip=None) -> Float32:
+    return silu(x) * y
+
+
+@dsl_user_op
+def dswiglu(
+    x: Float32, y: Float32, dout: Float32, *, loc=None, ip=None
+) -> Tuple[Float32, Float32, Float32]:
+    """
+    SwiGLU backward pass: computes gradients w.r.t. x (gate) and y (up projection)
+    Given: swiglu_out = silu(x) * y, and dout = grad w.r.t. swiglu_out
+    Returns: (dx, dy, swiglu_out) where dx = dout * y * d_silu(x), dy = dout * silu(x)
+
+    d_silu(x) = sigmoid(x) * (1 + x * (1 - sigmoid(x)))
+
+    This has been optimized to use fewer instructions (i.e. we expand things out
+    to use FFMA instead of FADD and FMUL).
+    """
+    # Compute sigmoid(x) using tanh: sigmoid(x) = 0.5 * (1 + tanh(0.5 * x))
+    x_half = 0.5 * x  # FMUL
+    sigmoid_x = 0.5 + 0.5 * tanh(x_half)  # MUFU.TANH, then FFMA
+    silu_x = x * sigmoid_x  # FMUL
+    silu_x_dout = silu_x * dout  # FMUL
+    #   d_silu(x) * dout
+    # = sigmoid_x * (1 + x * (1 - sigmoid_x)) * dout
+    # = (sigmoid_x + sigmoid_x * x * (1 - sigmoid_x)) * dout
+    # = (sigmoid_x + silu_x * (1 - sigmoid_x)) * dout
+    # = (sigmoid_x + silu_x - silu_x * sigmoid_x) * dout
+    # = (sigmoid_x - silu_x * sigmoid_x) * dout + silu_x * dout
+    d_silu_x_dout = (sigmoid_x - silu_x * sigmoid_x) * dout + silu_x_dout  # FFMA, FFMA
+    dx = d_silu_x_dout * y  # FMUL
+    dy = silu_x_dout
+    swiglu_out = silu_x * y  # FMUL
+    # Overall it's 1 MUFU.TANH, 5 FMUL, 3 FFMA
+    return dx, dy, swiglu_out
+
+
+@dsl_user_op
+def swiglu_oai(x: Float32, y: Float32, alpha: float = 1.702, *, loc=None, ip=None) -> Float32:
+    """The swiglu variant used in gpt-oss, which has a scaling factor on x and bias of 1 to y.
+    https://github.com/openai/gpt-oss/blob/7be9334950053a888e24887a57dac797a17d6e00/gpt_oss/torch/model.py#L249
+    x * sigmoid(alpha * x) * (y + 1)
+    Compile down to FMUL, FMUL, TANH, FFMA, FFMA
+    """
+    # Compute sigmoid(alpha * x) using tanh: sigmoid(z) = 0.5 * (1 + tanh(z/2))
+    x_half = 0.5 * x
+    silu_x = x_half * tanh(alpha * x_half) + x_half
+    return silu_x * y + silu_x
+
+
+@dsl_user_op
+def dswiglu_oai(
+    x: Float32, y: Float32, dout: Float32, alpha: float = 1.702, *, loc=None, ip=None
+) -> Tuple[Float32, Float32, Float32]:
+    """
+    Swiglu OAI backward pass: computes gradients w.r.t. x and y
+    Given: swiglu_oai_out = x * sigmoid(alpha * x) * (y + 1), and dout = grad w.r.t. swiglu_oai_out
+    Returns: (dx, dy, swiglu_oai_out)
+
+    Derivative of x * sigmoid(alpha * x) w.r.t. x:
+    d/dx[x * sigmoid(alpha * x)] = sigmoid(alpha * x) + alpha * x * sigmoid(alpha * x) * (1 - sigmoid(alpha * x))
+    """
+    # Compute sigmoid(alpha * x) using tanh: sigmoid(z) = 0.5 * (1 + tanh(z/2))
+    alpha_x_half = (0.5 * alpha) * x  # FMUL
+    sigmoid_alpha_x = 0.5 + 0.5 * tanh(alpha_x_half)  # MUFU.TANH, then FFMA
+    silu_x = x * sigmoid_alpha_x  # FMUL
+    silu_x_dout = silu_x * dout  # FMUL
+    # FFMA, FFMA, FMUL
+    d_silu_x_dout = (sigmoid_alpha_x + alpha * (silu_x - silu_x * sigmoid_alpha_x)) * dout
+    dx = d_silu_x_dout * y + d_silu_x_dout  # FFMA, instead of multiply by y + 1
+    dy = silu_x_dout
+    swiglu_out = silu_x * y + silu_x  # FFMA, instead of multiply by y + 1
+    # Overall it's 1 MUFU.TANH, 4 FMUL, 5 FFMA
+    return dx, dy, swiglu_out
+
+
+@dsl_user_op
+def glu(x: Float32, y: Float32, *, loc=None, ip=None) -> Float32:
+    """GLU: Gated Linear Unit
+    glu(x, y) = sigmoid(x) * y
+    Using tanh to compute sigmoid: sigmoid(x) = 0.5 * (1 + tanh(x/2))
+    """
+    x_half = 0.5 * x  # FMUL
+    sigmoid_x = 0.5 + 0.5 * tanh(x_half)  # MUFU.TANH, then FFMA
+    return sigmoid_x * y  # FMUL
+
+
+@dsl_user_op
+def dglu(
+    x: Float32, y: Float32, dout: Float32, *, loc=None, ip=None
+) -> Tuple[Float32, Float32, Float32]:
+    """
+    GLU backward pass: computes gradients w.r.t. x (gate) and y (up projection)
+    Given: glu_out = sigmoid(x) * y, and dout = grad w.r.t. glu_out
+    Returns: (dx, dy, glu_out) where:
+    - dx = dout * y * sigmoid(x) * (1 - sigmoid(x))
+    - dy = dout * sigmoid(x)
+    - glu_out = sigmoid(x) * y
+    """
+    # Compute sigmoid(x) using tanh: sigmoid(x) = 0.5 * (1 + tanh(x/2))
+    x_half = 0.5 * x  # FMUL
+    sigmoid_x = 0.5 + 0.5 * tanh(x_half)  # MUFU.TANH, then FFMA
+    sigmoid_x_dout = sigmoid_x * dout  # FMUL
+    glu_out = sigmoid_x * y  # FMUL
+    # dx = y * sigmoid(x) * (1 - sigmoid(x)) * dout
+    #    = y * (1 - sigmoid(x)) * sigmoid_x_dout
+    #    = (y - y * sigmoid(x)) * sigmoid_x_dout
+    #    = (y - glu_out) * sigmoid_x_dout
+    dx = (y - glu_out) * sigmoid_x_dout  # FADD, FMUL
+    dy = sigmoid_x_dout
+    # Total: 1 MUFU.TANH, 4 FMUL, 1 FADD, 1 FFMA
+    return dx, dy, glu_out
+
+
+@dsl_user_op
+def reglu(x: Float32, y: Float32, *, loc=None, ip=None) -> Float32:
+    """ReGLU: ReLU Gated Linear Unit
+    reglu(x, y) = relu(x) * y = max(x, 0) * y
+    """
+    return cute.arch.fmax(x, Float32(0.0)) * y
+
+
+@cute.jit
+@dsl_user_op
+def dreglu(
+    x: Float32, y: Float32, dout: Float32, *, loc=None, ip=None
+) -> Tuple[Float32, Float32, Float32]:
+    """
+    ReGLU backward pass: computes gradients w.r.t. x (gate) and y (up projection)
+    Given: reglu_out = relu(x) * y, and dout = grad w.r.t. reglu_out
+    Returns: (dx, dy, reglu_out) where:
+    - dx = dout * y if x > 0, else 0
+    - dy = dout * relu(x)
+    - reglu_out = relu(x) * y
+    """
+    x_pos = cutlass.Boolean(x > 0)
+    relu_x = cute.arch.fmax(x, Float32(0.0))
+    dx = (dout * y) if x_pos else Float32(0.0)
+    dy = dout * relu_x
+    reglu_out = relu_x * y
+    return dx, dy, reglu_out
+
+
+@dsl_user_op
+def geglu(x: Float32, y: Float32, *, loc=None, ip=None) -> Float32:
+    """GeGLU: GELU Gated Linear Unit
+    geglu(x, y) = gelu(x) * y
+    Uses the tanh approximation of GELU
+    """
+    return gelu_tanh_approx(x) * y
+
+
+@dsl_user_op
+def dgeglu(
+    x: Float32, y: Float32, dout: Float32, *, loc=None, ip=None
+) -> Tuple[Float32, Float32, Float32]:
+    """
+    GeGLU backward pass: computes gradients w.r.t. x (gate) and y (up projection)
+    Given: geglu_out = gelu(x) * y, and dout = grad w.r.t. geglu_out
+    Returns: (dx, dy, geglu_out) where:
+    - dx = dout * y * d_gelu(x)
+    - dy = dout * gelu(x)
+    - geglu_out = gelu(x) * y
+    """
+    # Reuse dgelu_tanh_approx to compute d_gelu(x) * dout and gelu(x)
+    dgelu_x_dout, gelu_x = dgelu_tanh_approx(x, dout)
+    # Compute gradients for geglu
+    dx = dgelu_x_dout * y
+    dy = gelu_x * dout
+    geglu_out = gelu_x * y
+    return dx, dy, geglu_out

quack/autotuner.py

@@ -0,0 +1,310 @@
+# Adapted from https://github.com/triton-lang/triton/blob/main/python/triton/runtime/autotuner.py
+# Copyright (C) 2025, Tri Dao.
+from __future__ import annotations
+
+import builtins
+import os
+import time
+import inspect
+import base64
+import hashlib
+import json
+from pathlib import Path
+from functools import cached_property, partial
+from typing import Dict, Tuple
+
+import torch
+from torch import Tensor
+
+import triton
+
+from . import __version__
+
+
+PACKAGE_NAME = "quack"
+VERSION = __version__
+
+
+def get_home_dir():
+    return os.getenv(f"{PACKAGE_NAME.upper()}_HOME", Path.home())
+
+
+def default_cache_dir():
+    return os.path.join(get_home_dir(), f".{PACKAGE_NAME}", "cache")
+
+
+class FileCacheManager(triton.runtime.cache.FileCacheManager):
+    def __init__(self, key):
+        super().__init__(key)
+        self.cache_dir = (
+            os.getenv(f"{PACKAGE_NAME.upper()}_CACHE_DIR", "").strip() or default_cache_dir()
+        )
+        if self.cache_dir:
+            self.cache_dir = os.path.join(self.cache_dir, self.key)
+            self.lock_path = os.path.join(self.cache_dir, "lock")
+            os.makedirs(self.cache_dir, exist_ok=True)
+        else:
+            raise RuntimeError("Could not create or locate cache dir")
+
+
+def _base32(key):
+    # Assume key is a hex string.
+    return base64.b32encode(bytes.fromhex(key)).decode("utf-8").rstrip("=")
+
+
+class Autotuner:
+    def __init__(self, fn, key, configs, restore_value=None, do_bench=None, cache_results=False):
+        if not configs:
+            self.configs = [AutotuneConfig()]
+        else:
+            self.configs = configs
+        signature = inspect.signature(fn)
+        self.keys = key
+        self.cache: Dict[Tuple, AutotuneConfig] = {}
+        self.arg_names = list(signature.parameters.keys())
+        self.cache_results = (
+            cache_results or os.getenv(f"{PACKAGE_NAME.upper()}_CACHE_AUTOTUNING", None) == "1"
+        )
+
+        self.restore_value = []
+        if restore_value is not None:
+            self.restore_value = list(restore_value)
+
+        if len(self.restore_value) > 0:
+
+            def _pre_hook(kwargs):
+                self.restore_copies = {name: kwargs[name].clone() for name in self.restore_value}
+
+            self.pre_hook = _pre_hook
+        else:
+            self.pre_hook = None
+
+        if len(self.restore_value) > 0:
+
+            def _post_hook(kwargs, exception):
+                for name in self.restore_value:
+                    kwargs[name].copy_(self.restore_copies[name])
+                self.restore_copies = {}
+
+            self.post_hook = _post_hook
+        else:
+            self.post_hook = None
+
+        self.fn = fn
+        self._do_bench = do_bench
+
+    @cached_property
+    def do_bench(self):
+        if self._do_bench is None:
+            return partial(triton.testing.do_bench, warmup=5, rep=25)
+        return self._do_bench
+
+    def _bench(self, *args, config, **meta):
+        verbose = os.environ.get(f"{PACKAGE_NAME.upper()}_PRINT_AUTOTUNING", None) == "1"
+        if verbose:
+            print(f"Autotuning kernel {self.fn.__name__} with config {config}")
+
+        # check for conflicts, i.e. meta-parameters both provided
+        # as kwargs and by the autotuner
+        conflicts = meta.keys() & config.kwargs.keys()
+        if conflicts:
+            raise ValueError(
+                f"Conflicting meta-parameters: {', '.join(conflicts)}."
+                " Make sure that you don't re-define auto-tuned symbols."
+            )
+        # augment meta-parameters with tunable ones
+        current = dict(meta, **config.all_kwargs())
+        full_nargs = {**self.nargs, **current}
+
+        def kernel_call():
+            if self.pre_hook is not None:
+                self.pre_hook(full_nargs)
+            try:
+                self.fn.__call__(
+                    *args,
+                    **current,
+                )
+            except Exception as e:
+                try:
+                    if self.post_hook is not None:
+                        self.post_hook(full_nargs, exception=e)
+                finally:
+                    # Throw exception raised by `self.fn.run`
+                    raise
+
+            if self.post_hook is not None:
+                self.post_hook(full_nargs, exception=None)
+
+        try:
+            return self.do_bench(kernel_call, quantiles=(0.5, 0.2, 0.8))
+        except Exception as e:
+            if verbose:
+                print(f"Autotuning failed with {e}")
+            return [float("inf"), float("inf"), float("inf")]
+
+    @torch.compiler.disable
+    def check_disk_cache(self, tuning_key, configs, bench_fn):
+        if not tuning_key:
+            bench_fn()
+            return
+
+        fn = self.fn
+        config_str_list = [str(c) for c in configs]
+        assert len(config_str_list) == len(set(config_str_list)), "Config strings must be unique"
+        cache_key = [VERSION, str(tuning_key)] + config_str_list
+        cache_key = hashlib.sha256("-".join(cache_key).encode("utf-8")).hexdigest()
+        cache = FileCacheManager(_base32(cache_key))
+        file_name = f"{fn.__name__[:150]}.autotune.json"
+        path = cache.get_file(file_name)
+        # There's an environment variable to force cache update
+        if path and not os.environ.get(f"{PACKAGE_NAME.upper()}_FORCE_CACHE_UPDATE", False):
+            str2config = {s: c for s, c in zip(config_str_list, configs)}
+            with open(path, "r") as cached_configs:
+                timings = json.load(cached_configs)["configs_timings"]
+                timings = {str2config[config]: timing for config, timing in timings}
+                self.cache[tuning_key] = builtins.min(timings, key=timings.get)
+                self.configs_timings = timings
+                self.bench_time = 0
+            return
+
+        bench_fn()
+        cache.put(
+            json.dumps(
+                {
+                    "key": tuning_key,
+                    "configs_timings": [
+                        (str(config), timings) for config, timings in self.configs_timings.items()
+                    ],
+                }
+            ),
+            file_name,
+            binary=False,
+        )
+
+    def __call__(self, *args, **kwargs):
+        self.nargs = dict(zip(self.arg_names, args))
+        used_cached_result = True
+        if len(self.configs) > 1:
+            all_args = {**self.nargs, **kwargs}
+            _args = {k: v for (k, v) in all_args.items() if k in self.arg_names}
+            # Need "str" to make it json-serializable
+            key = [str(_args[key]) for key in self.keys if key in _args]
+            for _, arg in _args.items():
+                if isinstance(arg, Tensor):
+                    key.append(str(arg.shape))
+                    # If stride != 0, 1, we just cache it as 2
+                    key.append(str([s if s in {0, 1} else 2 for s in arg.stride()]))
+                    key.append(str(arg.dtype))
+            key = tuple(key)
+            if key not in self.cache:
+                used_cached_result = False
+
+                @torch.compiler.disable  # Don't want any tracing here
+                def benchmark():
+                    bench_start = time.time()
+                    timings = {
+                        config: self._bench(*args, config=config, **kwargs)
+                        for config in self.configs
+                    }
+                    bench_end = time.time()
+                    if os.getenv(f"{PACKAGE_NAME.upper()}_PRINT_AUTOTUNING", None) == "1":
+                        for config, time_ in timings.items():
+                            print(f"[{config}] -> {time_[0]:.3f}ms")
+                    self.bench_time = bench_end - bench_start
+                    self.cache[key] = builtins.min(timings, key=timings.get)
+                    self.configs_timings = timings
+
+                if self.cache_results:
+                    self.check_disk_cache(key, self.configs, benchmark)
+                else:
+                    benchmark()
+
+            config = self.cache[key]
+        else:
+            config = self.configs[0]
+        self.best_config = config
+        if (
+            os.getenv(f"{PACKAGE_NAME.upper()}_PRINT_AUTOTUNING", None) == "1"
+            and not used_cached_result
+        ):
+            print(
+                f"{PACKAGE_NAME} autotuning for function {self.fn.__name__} finished after "
+                f"{self.bench_time:.2f}s; best config selected: {self.best_config};"
+            )
+        ret = self.fn.__call__(
+            *args,
+            **kwargs,
+            **config.all_kwargs(),
+        )
+        self.nargs = None
+        return ret
+
+
+class AutotuneConfig:
+    """
+    An object that represents a possible kernel configuration for the auto-tuner to try.
+
+    :ivar kwargs: a dictionary of meta-parameters to pass to the kernel as keyword arguments.
+    :type kwargs: dict[Str, Any]
+    """
+
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+
+    def __setstate__(self, state):
+        self.kwargs = state.get("kwargs", {})
+
+    def all_kwargs(self):
+        return self.kwargs
+
+    def __str__(self):
+        res = []
+        for k, v in self.kwargs.items():
+            res.append(f"{k}: {v}")
+        return ", ".join(res)
+
+    def __hash__(self):
+        return hash(tuple(*self.all_kwargs().items()))
+
+    def __eq__(self, other):
+        self_tuple = tuple(*self.all_kwargs().items())
+        other_tuple = tuple(*other.all_kwargs().items())
+        return self_tuple == other_tuple
+
+
+def autotune(configs, key=None, restore_value=None, do_bench=None, cache_results=True):
+    f"""
+    Decorator for auto-tuning a function function.
+
+    .. highlight:: python
+
+    If the environment variable :code:`{PACKAGE_NAME.upper()}_PRINT_AUTOTUNING` is set to
+    :code:`"1"`, we will print a message to stdout after autotuning each
+    kernel, including the time spent autotuning and the best configuration.
+
+    :param configs: a list of :code:`AutotuneConfig` objects
+    :type configs: list[AutotuneConfig]
+    :param key: a list of argument names whose change in value will trigger the evaluation of all provided configs.
+    :type key: list[str]
+    :param restore_value: a list of argument names whose value will be restored after evaluating any configs.
+    :type restore_value: list[str]
+    :param do_bench: a benchmark function to measure the time of each run.
+    :type do_bench: lambda fn, quantiles
+    :param cache_results: whether to cache autotune timings to disk.  Defaults to False.
+    "type cache_results: bool
+    """
+
+    if key is None:
+        key = []
+
+    def decorator(fn):
+        return Autotuner(
+            fn,
+            key,
+            configs,
+            restore_value=restore_value,
+            do_bench=do_bench,
+            cache_results=cache_results,
+        )
+
+    return decorator