feat: add an example on using autotuning.

.github/workflows/build_kernel.yaml

@@ -64,6 +64,7 @@ jobs:
             relu-kernel-cpu
             relu-backprop-compile-kernel
             relu-triton-kernel
+            gemm-triton-autotune-kernel
             silu-and-mul-kernel
 
   test:

docs/source/_toctree.yml

@@ -68,6 +68,10 @@
     - local: cli-skills
       title: Install agent skills
   title: kernel-builder CLI
+- sections: 
+    - local: builder/triton-autotune
+      title: Ship Triton autotune configurations
+  title: Advanced practices
 - sections:
     - local: kernel-requirements
       title: Kernel requirements

docs/source/builder/triton-autotune.md

@@ -0,0 +1,177 @@
+# Ship Triton autotune configurations
+
+Triton kernels typically have parameters, such as tile sizes, number of warps and 
+number of pipeline stages, whose optimal values depend on the GPU and the
+problem shape. Autotuning finds good values for these parameters, but doing
+it at runtime (e.g. with the
+[`@triton.autotune`](https://triton-lang.org/main/python-api/generated/triton.autotune.html)
+decorator) re-benchmarks every candidate configuration in each new process.
+
+
+However, one can run the tuner for the GPU models they like and store the
+best found configurations as files for using them later. This effectively
+reduces the potentially costly tuning time.
+
+`kernel-builder` support this by packaging these configurations as JSON files
+with the kernel. At runtime, the kernel looks up the configuration for the
+current GPU and shape and falls back to sensible defaults when there is no
+matching configuration.
+
+This is how, for example, the vLLM fused MoE kernel ships on the Hub: the
+[`RedHatAI/moe`](https://huggingface.co/RedHatAI/moe/tree/main/torch-ext/moe/configs)
+repository contains a `configs` directory with tuned configurations for many
+GPUs. This page walks through a small, complete example of the same pattern:
+the [`gemm-triton-autotune`](https://github.com/huggingface/kernels/tree/main/examples/kernels/gemm-triton-autotune)
+example kernel, a Triton GEMM published as
+[`kernels-test/gemm-triton-autotune`](https://huggingface.co/kernels-test/gemm-triton-autotune).
+
+## Shipping data files with a kernel
+
+Since autotune files are plain JSON files, we can store them anywhere inside the
+kernels main Python sources in `torch-ext/<kernel_name>`. For this example,
+we will use  `torch-ext/<kernel_name>/configs/`. By default, only `py` and `pyi`
+files are picked up from the kenel's Python source directory, so add `json` to the
+[`pyext` option](writing-kernels.md#torch-noarch) in `build.toml`:
+
+```toml
+[general]
+name = "gemm-triton-autotune"
+version = 1
+edition = 5
+license = "Apache-2.0"
+backends = ["cuda", "rocm", "xpu"]
+
+[general.hub]
+repo-id = "kernels-test/gemm-triton-autotune"
+
+[torch-noarch]
+pyext = ["json", "py"]
+```
+
+## Configuration file layout
+
+A GEMM computes `(M, K) @ (K, N)`. For a model, the weight dimensions `N`
+and `K` are known ahead of time, while `M` (e.g. the number of tokens)
+varies at runtime. The example therefore stores one file per `(N, K)` shape
+and GPU, following the same naming convention as the MoE kernel:
+
+```
+configs/N=4096,K=4096,device_name=NVIDIA_L4.json
+configs/N=14336,K=4096,device_name=NVIDIA_L4.json
+```
+
+Each file maps an `M` value to the best configuration that the autotuner
+found for that `M`:
+
+```json
+{
+    "1": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 32,
+        "GROUP_SIZE_M": 8,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "1024": {
+        "BLOCK_SIZE_M": 128,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 32,
+        "GROUP_SIZE_M": 8,
+        "num_warps": 4,
+        "num_stages": 3
+    }
+}
+```
+
+Since the device name is part of the file name, configurations tuned for one
+GPU are never applied to another. A configuration that would exceed the
+resources of a smaller GPU (e.g. shared memory) is therefore harmless to
+ship.
+
+## Looking up configurations at runtime
+
+At kernel launch, the kernel checks whether a configuration file exists for
+the current device and shape. If it does, the configuration with the nearest
+tuned `M` is used; otherwise the kernel falls back to a conservative default
+and logs a warning. The lookup is cached, so the file is read at most once
+per process:
+
+```python
+@functools.cache
+def _load_tuned_configs(N: int, K: int) -> Optional[Dict[int, Dict[str, int]]]:
+    path = _CONFIGS_DIR / config_file_name(N, K)
+    if path.exists():
+        logger.info("Using tuned GEMM configurations from %s.", path)
+        with open(path) as f:
+            return {int(m): config for m, config in json.load(f).items()}
+    logger.warning(
+        "No tuned GEMM configuration found for this device and shape (%s). "
+        "Falling back to heuristic defaults, performance may be suboptimal. "
+        "Generate a configuration with `tune_gemm(N=%d, K=%d)`.",
+        path.name,
+        N,
+        K,
+    )
+    return None
+
+
+def get_config(M: int, N: int, K: int) -> Dict[str, int]:
+    tuned = _load_tuned_configs(N, K)
+    if tuned is not None:
+        # Tuned Ms are spaced logarithmically, so pick the nearest in log space.
+        nearest_m = min(tuned, key=lambda m: abs(math.log(M / m)))
+        return tuned[nearest_m]
+    return default_config(M, N, K)
+```
+
+The configuration is then passed to the Triton kernel as its `constexpr`
+and launch parameters (see
+[`gemm.py`](https://github.com/huggingface/kernels/blob/main/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/gemm.py)
+in the example):
+
+```python
+def launch_gemm_kernel(a, b, out, config):
+    M, K = a.shape
+    N = b.shape[1]
+    grid = (
+        triton.cdiv(M, config["BLOCK_SIZE_M"]) * triton.cdiv(N, config["BLOCK_SIZE_N"]),
+    )
+    _gemm_kernel[grid](a, b, out, M, N, K, ..., **config)
+```
+
+## Generating the configurations
+
+The autotuner itself is ordinary benchmarking code: for each `M`, benchmark
+every candidate configuration with `triton.testing.do_bench` and keep the
+fastest. The example ships the tuner as part of the kernel (the `tune_gemm`
+function in
+[`tuning.py`](https://github.com/huggingface/kernels/blob/main/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/tuning.py)),
+so users can also generate configurations for GPUs that the kernel author
+did not tune. Candidates that do not fit the device (Triton raises
+`OutOfResources`) are skipped.
+
+The example repository contains a small script,
+[`tune.py`](https://github.com/huggingface/kernels/blob/main/examples/kernels/gemm-triton-autotune/tune.py),
+that runs the tuner and writes the configuration files to the source tree:
+
+```bash
+$ python tune.py --n 4096 --k 4096
+$ python tune.py --n 14336 --k 4096
+```
+
+Commit the generated files, rebuild, and the configurations ship with the
+kernel. When tuning a kernel that is not yet published, build it locally
+(see [Develop locally](local-dev.md)) and point `LOCAL_KERNELS` at the
+build:
+
+```bash
+$ LOCAL_KERNELS=kernels-test/gemm-triton-autotune=build python tune.py --n 4096 --k 4096
+```
+
+## Impact
+
+Tuned configurations are cheap to ship and can make a large difference. On
+an NVIDIA L4, the tuned configuration for a `(1024, 4096) @ (4096, 4096)`
+float16 GEMM is ~1.4× faster than the example's heuristic default (0.50 ms
+vs. 0.71 ms) — on par with cuBLAS for this shape.

docs/source/builder/writing-kernels.md

@@ -247,7 +247,9 @@ following options:
 
 - `src` (required): a list of source files and headers.
 - `pyext` (optional): the list of extensions for Python files. Default:
-  `["py", "pyi"]`.
+  `["py", "pyi"]`. Additional extensions can be listed to ship data files
+  with the kernel, such as `json` for
+  [Triton autotune configurations](triton-autotune.md).
 - `include` (optional): include directories relative to the project root.
   Default: `[]`.
 - `maxver` (optional): only build for this Torch version and earlier. Use cautiously, since this option produces
@@ -297,7 +299,9 @@ the supported archs. These are then exported to `metadata.json` for consumption
 by e.g. the Hugging Face Hub.
 
 - `pyext` (optional): the list of extensions for Python files. Default:
-  `["py", "pyi"]`.
+  `["py", "pyi"]`. Additional extensions can be listed to ship data files
+  with the kernel, such as `json` for
+  [Triton autotune configurations](triton-autotune.md).
 - `cuda-capabilities` (optional): a list of CUDA compute capabilities the
   kernel supports (e.g. `["9.0", "10.0"]`).
 - `rocm-archs` (optional): a list of ROCm architectures the kernel supports

examples/kernels/flake.nix

@@ -179,6 +179,11 @@
           path = ./relu-triton;
           drv = sys: out: out.packages.${sys}.redistributable.torch-cuda;
         }
+        {
+          name = "gemm-triton-autotune-kernel";
+          path = ./gemm-triton-autotune;
+          drv = sys: out: out.packages.${sys}.redistributable.torch-cuda;
+        }
       ];
 
       # ROCm kernels to build in CI.
@@ -188,6 +193,11 @@
           path = ./relu-triton;
           drv = sys: out: out.packages.${sys}.redistributable.torch-rocm;
         }
+        {
+          name = "gemm-triton-autotune-kernel";
+          path = ./gemm-triton-autotune;
+          drv = sys: out: out.packages.${sys}.redistributable.torch-rocm;
+        }
         {
           name = "relu-invalid-capability";
           path = ./relu-invalid-capability;
@@ -269,6 +279,11 @@
           path = ./relu-triton;
           drv = sys: out: out.packages.${sys}.redistributable.torch-xpu;
         }
+        {
+          name = "gemm-triton-autotune-kernel";
+          path = ./gemm-triton-autotune;
+          drv = sys: out: out.packages.${sys}.redistributable.torch-xpu;
+        }
         {
           name = "relu-kernel";
           path = ./relu;

examples/kernels/gemm-triton-autotune/CARD.md

@@ -0,0 +1,60 @@
+---
+library_name: kernels
+{% if license %}license: {{ license }}
+{% endif %}---
+
+This is the repository card of {{ repo_id }} that has been pushed on the Hub. It was built to be used with the [`kernels` library](https://github.com/huggingface/kernels). This card was automatically generated.
+
+## How to use
+{% if functions %}
+
+```python
+# make sure `kernels` is installed: `pip install -U kernels`
+from kernels import get_kernel
+
+# If the org / user isn't a trusted publisher, pass `trust_remote_code=True` to the
+# `get_kernel` call. You can find whether this kernel is from a trusted publisher
+# by going to the kernel's Hub page and finding the "Trusted publisher" status at
+# the top of the page.
+kernel_module = get_kernel("{{ repo_id }}", version={{ version }})
+{{ functions[0] }} = kernel_module.{{ functions[0] }}
+
+{{ functions[0] }}(...)
+```
+{% else %}
+
+Usage example not available.
+{% endif %}
+
+## Available functions
+{% if functions %}
+{% for func in functions %}
+- `{{ func }}`
+{% endfor %}
+{% else %}
+
+Function list not available.
+{% endif %}
+{% if layers %}
+
+## Available layers
+{% for layer in layers %}
+- `{{ layer }}`
+{% endfor %}
+{% endif %}
+
+## Benchmarks
+{% if has_benchmark %}
+
+Benchmarking script is available for this kernel. Run `kernels benchmark {{ repo_id }} --version {{ version }}`.
+{% else %}
+
+No benchmark available yet.
+{% endif %}
+{% if upstream %}
+
+## Source code
+
+Source code of this kernel originally comes from {{ upstream }} and it was repurposed for compatibility with `kernels`.
+{% endif %}
+

examples/kernels/gemm-triton-autotune/build.toml

@@ -0,0 +1,19 @@
+[general]
+name = "gemm-triton-autotune"
+version = 1
+edition = 5
+license = "Apache-2.0"
+backends = [
+    "cuda",
+    "rocm",
+    "xpu",
+]
+
+[general.hub]
+repo-id = "kernels-test/gemm-triton-autotune"
+
+[torch-noarch]
+pyext = [
+    "json",
+    "py",
+]

examples/kernels/gemm-triton-autotune/flake.nix

@@ -0,0 +1,17 @@
+{
+  description = "Flake for kernels tests";
+
+  inputs = {
+    kernel-builder.url = "path:../../..";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genKernelFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}

examples/kernels/gemm-triton-autotune/tests/conftest.py

@@ -0,0 +1,12 @@
+import pytest
+import torch
+
+
+@pytest.fixture(scope="session")
+def device() -> torch.device:
+    if hasattr(torch, "xpu") and torch.xpu.is_available():
+        return torch.device("xpu")
+    elif torch.cuda.is_available():
+        return torch.device("cuda")
+    else:
+        pytest.skip("No GPU available for Triton tests")