From f754782f3342b2a33cc3dc68d6b7b4b0ae1e9ed1 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Mon, 10 Aug 2026 09:02:05 +0000 Subject: [PATCH 1/9] feat: add an example on using autotuning. --- .github/workflows/build_kernel.yaml | 1 + PLAN-issue-733.md | 168 ++++++++++++++++ docs/source/_toctree.yml | 2 + docs/source/builder/triton-autotune.md | 174 ++++++++++++++++ docs/source/builder/writing-kernels.md | 12 +- examples/kernels/flake.nix | 15 ++ examples/kernels/gemm-triton-autotune/CARD.md | 60 ++++++ .../kernels/gemm-triton-autotune/build.toml | 19 ++ .../kernels/gemm-triton-autotune/flake.nix | 17 ++ .../gemm-triton-autotune/tests/__init__.py | 0 .../gemm-triton-autotune/tests/conftest.py | 12 ++ .../gemm-triton-autotune/tests/test_gemm.py | 96 +++++++++ .../gemm_triton_autotune/__init__.py | 26 +++ .../N=14336,K=4096,device_name=NVIDIA_L4.json | 50 +++++ .../N=4096,K=4096,device_name=NVIDIA_L4.json | 50 +++++ .../torch-ext/gemm_triton_autotune/gemm.py | 103 ++++++++++ .../torch-ext/gemm_triton_autotune/tuning.py | 189 ++++++++++++++++++ examples/kernels/gemm-triton-autotune/tune.py | 69 +++++++ nix-builder/tests/Dockerfile.test-kernel | 1 + nix-builder/tests/run-tests.sh | 4 + 20 files changed, 1065 insertions(+), 3 deletions(-) create mode 100644 PLAN-issue-733.md create mode 100644 docs/source/builder/triton-autotune.md create mode 100644 examples/kernels/gemm-triton-autotune/CARD.md create mode 100644 examples/kernels/gemm-triton-autotune/build.toml create mode 100644 examples/kernels/gemm-triton-autotune/flake.nix create mode 100644 examples/kernels/gemm-triton-autotune/tests/__init__.py create mode 100644 examples/kernels/gemm-triton-autotune/tests/conftest.py create mode 100644 examples/kernels/gemm-triton-autotune/tests/test_gemm.py create mode 100644 examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/__init__.py create mode 100644 examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/configs/N=14336,K=4096,device_name=NVIDIA_L4.json create mode 100644 examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/configs/N=4096,K=4096,device_name=NVIDIA_L4.json create mode 100644 examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/gemm.py create mode 100644 examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/tuning.py create mode 100644 examples/kernels/gemm-triton-autotune/tune.py diff --git a/.github/workflows/build_kernel.yaml b/.github/workflows/build_kernel.yaml index 1ed9ac6a..953cef32 100644 --- a/.github/workflows/build_kernel.yaml +++ b/.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: diff --git a/PLAN-issue-733.md b/PLAN-issue-733.md new file mode 100644 index 00000000..12b1dd9a --- /dev/null +++ b/PLAN-issue-733.md @@ -0,0 +1,168 @@ +# Plan: Triton autotune configs example (issue #733) + +**Issue:** [huggingface/kernels#733](https://github.com/huggingface/kernels/issues/733) — +"Document/example: shipping Triton autotune configs with a kernel" + +**Goal:** Add an example Triton kernel to `examples/kernels/` that ships pre-computed +autotune configurations as JSON files, plus a script that generates those files and +documentation explaining the pattern. Reference implementation: +[`RedHatAI/moe`](https://huggingface.co/RedHatAI/moe/tree/main/torch-ext/moe) (vLLM fused MoE). + +## Background / findings from the repo + +- `examples/kernels/relu-triton/` is the existing Triton example: a `torch-noarch` + kernel with `torch.library.custom_op` wrappers. It is built in CI for CUDA, ROCm, + and XPU via the `ciKernels` / `ciRocmKernels` / `ciXpuKernels` lists in + `examples/kernels/flake.nix`, and its tests run on GPU through + `nix-builder/tests/Dockerfile.test-kernel` + `run-tests.sh` using the + `LOCAL_KERNELS` env var (`kernels-test/=`). +- `examples/kernels/extra-data/` already demonstrates shipping non-Python data by + adding `"json"` to `pyext` in `build.toml` — the exact mechanism we need for + config files. +- The vLLM MoE pattern (`fused_moe.py`): config files live in + `torch-ext/moe/configs/` named + `E=...,N=...,device_name=NVIDIA_H100_80GB_HBM3[,dtype=...].json`. Each file maps a + batch-size bucket (`M`) to a Triton config dict (`BLOCK_SIZE_M/N/K`, `GROUP_SIZE_M`, + `num_warps`, `num_stages`). At runtime an `lru_cache`d loader looks up the file for + the current device; if absent, it falls back to a heuristic default and logs a + warning. +- A local NVIDIA L4 GPU is available — the same GPU as the CI test runners + (`aws-g6-12xlarge`) — so we can generate and commit a real config that CI will + actually exercise. + +## Design decisions + +1. **Kernel: a Triton GEMM (`matmul`)**, not another relu. Autotuning is only + meaningful when tile sizes / warps / stages matter; GEMM is the canonical case and + mirrors the MoE reference. Kernel body can be the standard Triton matmul tutorial + kernel (fp16/bf16/fp32 inputs, fp16-accumulate-in-fp32). +2. **Ship JSON lookup tables instead of relying on `@triton.autotune`.** The + decorator re-benchmarks in every process (slow startup, nondeterministic) and its + cache is not portable. Instead: an offline tuning script does the + `@triton.autotune`-style grid search once, writes JSON, and the kernel loads the + JSON at call time. This is exactly what the issue asks to demonstrate. +3. **File naming convention** (adapted from vLLM MoE): since GEMM weight dims are + known ahead of time but the batch dim `M` varies at runtime, one file per + `(N, K, device)`: `configs/N=4096,K=4096,device_name=NVIDIA_L4.json`, whose keys + are `M` values and values are config dicts. Runtime picks the nearest `M` key + (same `min(keys, key=|log(M/key)|)` trick as vLLM). +4. **Tuning entry point ships with the kernel** (a `tune.py` module inside the + package) so users can regenerate configs for *their* GPU after + `get_kernel(...)`, plus a thin CLI script in the example root for kernel authors. + +## New files + +``` +examples/kernels/gemm-triton-autotune/ +├── build.toml # torch-noarch, pyext = ["py", "json"] +├── flake.nix # copied from relu-triton +├── CARD.md # copied from relu-triton +├── tune.py # CLI: python tune.py --n 4096 --k 4096 [--out ...] +│ # writes into torch-ext/gemm_triton_autotune/configs/ +├── torch-ext/gemm_triton_autotune/ +│ ├── __init__.py # exports gemm(), tune_gemm() +│ ├── gemm.py # @triton.jit matmul kernel + custom_op wrapper; +│ │ # per-call config lookup via tuning.get_config() +│ ├── tuning.py # device-name helper, config file naming, +│ │ # lru_cache'd JSON loader, default-config fallback +│ │ # (with one-time warning), tune_gemm() grid search +│ │ # using triton.testing.do_bench +│ └── configs/ +│ └── N=4096,K=4096,device_name=NVIDIA_L4.json # generated on local L4 +└── tests/ + ├── __init__.py + ├── conftest.py # device fixture, same as relu-triton + └── test_gemm.py # correctness vs torch.matmul across dtypes/shapes; + # shipped-config load test (monkeypatched device name); + # fallback-to-default test for unknown (N, K) +``` + +### `build.toml` sketch + +```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 = ["py", "json"] +``` + +### Runtime config lookup (in `tuning.py`) + +```python +@functools.lru_cache +def get_config(M: int, N: int, K: int) -> dict: + path = Path(__file__).parent / "configs" / _config_file_name(N, K) + if path.exists(): + configs = {int(m): cfg for m, cfg in json.loads(path.read_text()).items()} + return configs[min(configs, key=lambda m: abs(math.log(M / m)))] + warnings.warn(f"No tuned GEMM config for {path.name}, using defaults...", once) + return _default_config(M, N, K) +``` + +### Tuning script behavior + +- Candidate grid: the usual matmul space (`BLOCK_M/N/K ∈ {32..256}`, `GROUP_M`, + `num_warps ∈ {4, 8}`, `num_stages ∈ {2..5}`), pruned to valid combos. +- Benchmarks each candidate with `triton.testing.do_bench` for each `M` in + `{1, 16, 64, 256, 1024, 4096}` at fixed `(N, K)`. +- Writes `{M: best_config}` JSON to the package `configs/` dir, named with + `torch.cuda.get_device_name().replace(" ", "_")` (XPU equivalent when applicable). + +## Existing files to modify + +1. **`examples/kernels/flake.nix`** — register the new kernel in `ciKernels` + (`torch-cuda` noarch build, like `relu-triton-kernel`), `ciRocmKernels`, and + `ciXpuKernels`. +2. **`.github/workflows/build_kernel.yaml`** — add `gemm-triton-autotune-kernel` to + the uploaded-artifacts list. +3. **`nix-builder/tests/Dockerfile.test-kernel`** — `COPY + examples/kernels/gemm-triton-autotune/tests ./gemm_triton_autotune_tests`. +4. **`nix-builder/tests/run-tests.sh`** — run the new tests with + `LOCAL_KERNELS="kernels-test/gemm-triton-autotune=..."`. +5. **Docs:** + - New page `docs/source/builder/triton-autotune.md` — "Shipping Triton autotune + configurations": why ship configs, the JSON-per-device pattern, `pyext = ["json"]`, + the loader/fallback pattern, how to run the tune script, links to the example + and to `RedHatAI/moe`. + - Add the page to `docs/source/_toctree.yml` (kernel-builder section, after + `builder/writing-kernels`). + +## Implementation order + +1. Scaffold the example kernel (build.toml, flake.nix, CARD.md, package code). +2. Set up a local venv (`uv venv` + torch/cu126 + triton + kernels + pytest) and get the + kernel running directly from `torch-ext/` on the L4. +3. Run `tune.py` on the L4 for `(N=4096, K=4096)`; commit the generated JSON. +4. Write tests; run them locally against the local build + (`LOCAL_KERNELS=kernels-test/gemm-triton-autotune=` after a + `kernels build`/nix build, matching the CI invocation). +5. Wire up CI (flake.nix lists, workflow artifact list, Dockerfile, run-tests.sh). +6. Write the docs page + toctree entry. +7. `nix flake check` / build the example via + `nix build ./examples/kernels#ci-build-cuda` if feasible locally, else rely on CI. + +## Out of scope / maintainer follow-ups + +- Pushing the built kernel to the `kernels-test/gemm-triton-autotune` Hub repo + (needed for `get_kernel` without `LOCAL_KERNELS`) requires org access — CI tests + use `LOCAL_KERNELS`, so nothing blocks on this, but the Hub repo should be created + when merging (same as other `kernels-test/*` examples). +- Configs for ROCm/XPU devices can be contributed later by whoever has the hardware; + the fallback path covers them meanwhile (and the fallback is itself part of what + the example demonstrates). + +## Open questions + +1. Kernel/package name: `gemm-triton-autotune` (proposed) vs `matmul-triton-tune`. +2. Should the docs page live under kernel-builder docs (proposed) or as a section + appended to `writing-kernels.md`? +3. Single `(N, K)` shape for the committed config (proposed: 4096×4096) or a couple + of shapes to show multiple config files? diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 1b57b8da..008d0d21 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -47,6 +47,8 @@ - sections: - local: builder/writing-kernels title: Write kernels + - local: builder/triton-autotune + title: Ship Triton autotune configurations - local: builder/build title: Build with Nix - local: builder/local-dev diff --git a/docs/source/builder/triton-autotune.md b/docs/source/builder/triton-autotune.md new file mode 100644 index 00000000..27caa6e8 --- /dev/null +++ b/docs/source/builder/triton-autotune.md @@ -0,0 +1,174 @@ +# Ship Triton autotune configurations + +Triton kernels typically have parameters — tile sizes, number of warps, +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. +For kernels on the Hub, there is a better option: run the autotuner once per +GPU model, store the best configurations as JSON files, and ship those 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 + +Configuration files are plain JSON files inside the kernel's Python package, +in `torch-ext//configs/`. By default, only `py` and `pyi` +files are picked up from the package 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"] +``` + +This works the same for AOT-compiled kernels — the `torch` section also +supports `pyext`. + +## 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.lru_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: + # 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 +``` + +## Does it matter? + +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. diff --git a/docs/source/builder/writing-kernels.md b/docs/source/builder/writing-kernels.md index 905ba7ed..88d19849 100644 --- a/docs/source/builder/writing-kernels.md +++ b/docs/source/builder/writing-kernels.md @@ -37,7 +37,9 @@ format of the `build.toml` file, and some additional Python glue that `kernel-builder` provides. We will use a [simple ReLU kernel](https://github.com/huggingface/kernels/tree/main/examples/kernels/relu) as the running example. After reading this page, you may also want to have a look at the more realistic [ReLU kernel with backprop and `torch.compile`](https://github.com/huggingface/kernels/tree/main/examples/kernels/relu-backprop-compile) -support. +support. For Triton kernels, see +[Ship Triton autotune configurations](triton-autotune.md) for how to tune a +kernel and ship the tuned configurations with it. > [!TIP] > We maintain a set of conforming kernels in the @@ -247,7 +249,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 +301,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 diff --git a/examples/kernels/flake.nix b/examples/kernels/flake.nix index 2af64fa6..11ff0056 100644 --- a/examples/kernels/flake.nix +++ b/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; diff --git a/examples/kernels/gemm-triton-autotune/CARD.md b/examples/kernels/gemm-triton-autotune/CARD.md new file mode 100644 index 00000000..b970ce41 --- /dev/null +++ b/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 %} + diff --git a/examples/kernels/gemm-triton-autotune/build.toml b/examples/kernels/gemm-triton-autotune/build.toml new file mode 100644 index 00000000..39f79bc4 --- /dev/null +++ b/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", +] diff --git a/examples/kernels/gemm-triton-autotune/flake.nix b/examples/kernels/gemm-triton-autotune/flake.nix new file mode 100644 index 00000000..29c0cc8f --- /dev/null +++ b/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 = ./.; + }; +} diff --git a/examples/kernels/gemm-triton-autotune/tests/__init__.py b/examples/kernels/gemm-triton-autotune/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/kernels/gemm-triton-autotune/tests/conftest.py b/examples/kernels/gemm-triton-autotune/tests/conftest.py new file mode 100644 index 00000000..103c610f --- /dev/null +++ b/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") diff --git a/examples/kernels/gemm-triton-autotune/tests/test_gemm.py b/examples/kernels/gemm-triton-autotune/tests/test_gemm.py new file mode 100644 index 00000000..d6c7ed2c --- /dev/null +++ b/examples/kernels/gemm-triton-autotune/tests/test_gemm.py @@ -0,0 +1,96 @@ +import json + +import kernels +import pytest +import torch + +gemm_kernel = kernels.get_kernel("kernels-test/gemm-triton-autotune", version=1) + +# GEMM output values grow with sqrt(K) for randn inputs, so tolerances are +# scaled by the magnitude of the output values. +DTYPE_TOLERANCES = { + torch.float16: {"rtol": 2e-2, "atol": 5e-1}, + torch.bfloat16: {"rtol": 4e-2, "atol": 4.0}, +} + + +@pytest.mark.kernels_ci +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("M", [1, 17, 256]) +@pytest.mark.parametrize("N,K", [(4096, 4096), (512, 511)]) +def test_gemm(device, dtype, M, N, K): + a = torch.randn(M, K, device=device, dtype=dtype) + b = torch.randn(K, N, device=device, dtype=dtype) + ref = (a.float() @ b.float()).to(dtype) + torch.testing.assert_close(gemm_kernel.gemm(a, b), ref, **DTYPE_TOLERANCES[dtype]) + + out = torch.empty(M, N, device=device, dtype=dtype) + gemm_kernel.gemm(a, b, out=out) + torch.testing.assert_close(out, ref, **DTYPE_TOLERANCES[dtype]) + + +@pytest.mark.kernels_ci +def test_gemm_validation(device): + a = torch.randn(4, 8, device=device, dtype=torch.float16) + with pytest.raises(ValueError, match="Incompatible"): + gemm_kernel.gemm(a, torch.randn(4, 8, device=device, dtype=torch.float16)) + with pytest.raises(ValueError, match="dtype"): + gemm_kernel.gemm(a, torch.randn(8, 4, device=device, dtype=torch.bfloat16)) + with pytest.raises(ValueError, match="dtype"): + gemm_kernel.gemm(a.float(), torch.randn(8, 4, device=device)) + with pytest.raises(ValueError, match="2D"): + gemm_kernel.gemm(a.flatten(), a.flatten()) + + +@pytest.mark.kernels_ci +def test_nearest_tuned_m_is_used(device, monkeypatch, tmp_path): + tuning = gemm_kernel.tuning + + small = tuning.default_config(16, 128, 128) + large = tuning.default_config(1024, 128, 128) + assert small != large + config_path = tmp_path / tuning.config_file_name(128, 128) + config_path.write_text(json.dumps({"16": small, "1024": large})) + + monkeypatch.setattr(tuning, "_CONFIGS_DIR", tmp_path) + tuning._load_tuned_configs.cache_clear() + try: + assert tuning.get_config(1, 128, 128) == small + assert tuning.get_config(32, 128, 128) == small + assert tuning.get_config(500, 128, 128) == large + assert tuning.get_config(100000, 128, 128) == large + finally: + tuning._load_tuned_configs.cache_clear() + + +@pytest.mark.kernels_ci +def test_fallback_to_default_config(device): + tuning = gemm_kernel.tuning + # No configuration is shipped for this shape, so the heuristic default + # should be used. + N, K = 123, 321 + assert tuning.get_config(64, N, K) == tuning.default_config(64, N, K) + + +@pytest.mark.kernels_ci +def test_shipped_config_is_used(device): + tuning = gemm_kernel.tuning + # Configurations tuned on an L4 GPU ship with the kernel (which is also + # the GPU that CI runs on). + if tuning.device_name() != "NVIDIA_L4": + pytest.skip("Shipped configurations were tuned for NVIDIA L4") + for N, K in [(4096, 4096), (14336, 4096)]: + assert tuning._load_tuned_configs(N, K) is not None + + +@pytest.mark.kernels_ci +def test_tune_gemm(device, tmp_path): + tuning = gemm_kernel.tuning + candidates = [tuning.default_config(16, 256, 256), tuning.default_config(64, 256, 256)] + path = gemm_kernel.tune_gemm(N=256, K=256, Ms=(1, 64), save_dir=tmp_path, candidates=candidates) + + assert path.parent == tmp_path + configs = json.loads(path.read_text()) + assert set(configs.keys()) == {"1", "64"} + for config in configs.values(): + assert config in candidates diff --git a/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/__init__.py b/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/__init__.py new file mode 100644 index 00000000..4e239a0c --- /dev/null +++ b/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/__init__.py @@ -0,0 +1,26 @@ +from typing import Optional + +import torch + +from ._ops import ops +from .gemm import _gemm # noqa: F401 — imported to register the op. +from .tuning import tune_gemm + +_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16) + + +def gemm(a: torch.Tensor, b: torch.Tensor, out: Optional[torch.Tensor] = None) -> torch.Tensor: + """Compute ``a @ b`` using a Triton GEMM kernel with tuned configurations.""" + if a.ndim != 2 or b.ndim != 2: + raise ValueError("gemm expects 2D tensors") + if a.shape[1] != b.shape[0]: + raise ValueError(f"Incompatible GEMM shapes: {tuple(a.shape)} @ {tuple(b.shape)}") + if a.dtype not in _SUPPORTED_DTYPES or a.dtype != b.dtype: + raise ValueError(f"gemm expects two {_SUPPORTED_DTYPES} tensors of the same dtype") + if out is None: + out = torch.empty(a.shape[0], b.shape[1], device=a.device, dtype=a.dtype) + ops.gemm(out, a, b) + return out + + +__all__ = ["gemm", "tune_gemm"] diff --git a/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/configs/N=14336,K=4096,device_name=NVIDIA_L4.json b/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/configs/N=14336,K=4096,device_name=NVIDIA_L4.json new file mode 100644 index 00000000..45a90a97 --- /dev/null +++ b/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/configs/N=14336,K=4096,device_name=NVIDIA_L4.json @@ -0,0 +1,50 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 8, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "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 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/configs/N=4096,K=4096,device_name=NVIDIA_L4.json b/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/configs/N=4096,K=4096,device_name=NVIDIA_L4.json new file mode 100644 index 00000000..6fcaafa3 --- /dev/null +++ b/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/configs/N=4096,K=4096,device_name=NVIDIA_L4.json @@ -0,0 +1,50 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 8, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "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 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/gemm.py b/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/gemm.py new file mode 100644 index 00000000..1b7f728a --- /dev/null +++ b/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/gemm.py @@ -0,0 +1,103 @@ +from typing import Dict + +import torch +import triton +import triton.language as tl + +from ._ops import add_op_namespace_prefix +from .tuning import get_config + + +@triton.jit +def _gemm_kernel( + a_ptr, + b_ptr, + c_ptr, + M, + N, + K, + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) + acc = tl.dot(a, b, acc) + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + c = acc.to(c_ptr.dtype.element_ty) + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + + +def launch_gemm_kernel( + a: torch.Tensor, b: torch.Tensor, out: torch.Tensor, config: Dict[str, int] +) -> None: + """Launch the GEMM kernel with an explicit configuration. + + The tuner uses this to benchmark candidate configurations; regular calls + go through the ``gemm`` op, which looks up the configuration itself. + """ + 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, + a.stride(0), + a.stride(1), + b.stride(0), + b.stride(1), + out.stride(0), + out.stride(1), + **config, + ) + + +@torch.library.custom_op(add_op_namespace_prefix("gemm"), mutates_args={"out"}) +def _gemm(out: torch.Tensor, a: torch.Tensor, b: torch.Tensor) -> None: + M, K = a.shape + N = b.shape[1] + if M == 0 or N == 0: + return + launch_gemm_kernel(a, b, out, get_config(M, N, K)) + + +@_gemm.register_fake +def _(out: torch.Tensor, a: torch.Tensor, b: torch.Tensor) -> None: + pass diff --git a/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/tuning.py b/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/tuning.py new file mode 100644 index 00000000..941a6f03 --- /dev/null +++ b/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/tuning.py @@ -0,0 +1,189 @@ +"""Loading and generating tuned GEMM configurations. + +Tuned configurations are stored as JSON files in the ``configs`` directory +of this package, one file per ``(N, K, device)`` combination. Each file maps +an ``M`` value (the dimension that is typically only known at runtime, e.g. +the number of tokens) to the Triton configuration that performed best for +that ``M``: + + { + "16": {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, ..., "num_warps": 4}, + "1024": {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, ...}, + ... + } + +At kernel launch, the configuration with the nearest tuned ``M`` is used. +When no configuration file exists for the current device and shape, a +heuristic default is used instead and a warning is logged once. +""" + +import functools +import json +import logging +import math +from pathlib import Path +from typing import Dict, Iterable, Optional, Sequence + +import torch +import triton + +logger = logging.getLogger(__name__) + +_CONFIGS_DIR = Path(__file__).parent / "configs" + +# The M values that are tuned by default: GEMMs are often skinny (a few +# tokens during decoding) or wide (large batches during prefill), so we +# cover a logarithmic range. +DEFAULT_TUNE_MS = (1, 16, 64, 256, 1024, 4096) + + +def device_name() -> str: + """Name of the current accelerator, as used in configuration file names.""" + if hasattr(torch, "xpu") and torch.xpu.is_available(): + name = torch.xpu.get_device_name() + else: + name = torch.cuda.get_device_name() + return name.replace(" ", "_") + + +def config_file_name(N: int, K: int) -> str: + return f"N={N},K={K},device_name={device_name()}.json" + + +@functools.lru_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]: + """Return the best known configuration for a ``(M, N, K)`` GEMM.""" + tuned = _load_tuned_configs(N, K) + if tuned: + # Pick the configuration tuned for the M closest to the runtime M + # (in log space, since tuned Ms are spaced logarithmically). + nearest_m = min(tuned, key=lambda m: abs(math.log(M / m))) + return tuned[nearest_m] + return default_config(M, N, K) + + +def default_config(M: int, N: int, K: int) -> Dict[str, int]: + """Heuristic configuration for GEMMs without tuned configurations. + + Deliberately conservative so that it runs on every supported device. + """ + return { + "BLOCK_SIZE_M": 16 if M <= 16 else 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + } + + +def candidate_configs() -> Iterable[Dict[str, int]]: + """The configuration search space used by ``tune_gemm``. + + Candidates that do not fit the device (e.g. exceed shared memory) are + skipped during tuning. + """ + for block_m in (16, 32, 64, 128): + for block_n in (32, 64, 128): + for block_k in (32, 64): + for num_warps in (4, 8): + # Wide warps only help for large tiles. + if num_warps == 8 and block_m * block_n < 64 * 128: + continue + for num_stages in (2, 3, 4): + yield { + "BLOCK_SIZE_M": block_m, + "BLOCK_SIZE_N": block_n, + "BLOCK_SIZE_K": block_k, + "GROUP_SIZE_M": 8, + "num_warps": num_warps, + "num_stages": num_stages, + } + + +def tune_gemm( + N: int, + K: int, + Ms: Sequence[int] = DEFAULT_TUNE_MS, + dtype: torch.dtype = torch.float16, + save_dir: Optional[Path] = None, + candidates: Optional[Sequence[Dict[str, int]]] = None, +) -> Path: + """Benchmark candidate configurations and save the best ones. + + For each ``M`` in ``Ms``, every candidate configuration is benchmarked + for an ``(M, K) @ (K, N)`` GEMM and the fastest is recorded. The result + is written as a JSON file named after ``(N, K)`` and the current device. + + ``save_dir`` defaults to this package's ``configs`` directory, so that + subsequent ``gemm`` calls in the same environment pick up the tuned + configurations. Kernel authors should pass the ``configs`` directory of + the kernel *source tree* instead (see ``tune.py`` in the repository + root) and commit the result, so that the configurations ship with the + kernel. + """ + # Import here to avoid a circular import at module load time. + from .gemm import launch_gemm_kernel + + device = "xpu" if hasattr(torch, "xpu") and torch.xpu.is_available() else "cuda" + if candidates is None: + candidates = list(candidate_configs()) + + best_configs: Dict[int, Dict[str, int]] = {} + for M in Ms: + a = torch.randn(M, K, device=device, dtype=dtype) + b = torch.randn(K, N, device=device, dtype=dtype) + out = torch.empty(M, N, device=device, dtype=dtype) + + best_time = math.inf + for config in candidates: + try: + time = triton.testing.do_bench( + lambda: launch_gemm_kernel(a, b, out, config) + ) + except triton.runtime.errors.OutOfResources: + # Configuration does not fit this device, skip it. + continue + if time < best_time: + best_time, best_configs[M] = time, config + + if M not in best_configs: + raise RuntimeError(f"No candidate configuration fits M={M}, N={N}, K={K}") + logger.info( + "Best configuration for M=%d, N=%d, K=%d: %s (%.4f ms)", + M, + N, + K, + best_configs[M], + best_time, + ) + + if save_dir is None: + save_dir = _CONFIGS_DIR + save_dir.mkdir(parents=True, exist_ok=True) + path = save_dir / config_file_name(N, K) + with open(path, "w") as f: + json.dump({str(m): config for m, config in sorted(best_configs.items())}, f, indent=4) + f.write("\n") + + # Make sure that new configurations are picked up by subsequent calls. + _load_tuned_configs.cache_clear() + + return path diff --git a/examples/kernels/gemm-triton-autotune/tune.py b/examples/kernels/gemm-triton-autotune/tune.py new file mode 100644 index 00000000..141f8212 --- /dev/null +++ b/examples/kernels/gemm-triton-autotune/tune.py @@ -0,0 +1,69 @@ +#!/usr/bin/env python3 +"""Autotune the GEMM kernel and save the configurations to the source tree. + +The tuned configurations are written to +``torch-ext/gemm_triton_autotune/configs/``, named after the GEMM shape and +the current accelerator. Commit these files so that they ship with the +kernel. + +Tune the kernel as published on the Hub: + + python tune.py --n 4096 --k 4096 + +Or tune a local build (see the "Develop locally" chapter of the +kernel-builder documentation for building into ``build/``): + + LOCAL_KERNELS=kernels-test/gemm-triton-autotune=build python tune.py --n 4096 --k 4096 +""" + +import argparse +import logging +from pathlib import Path + +import torch + +import kernels + +SAVE_DIR = Path(__file__).parent / "torch-ext" / "gemm_triton_autotune" / "configs" + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--n", type=int, required=True, help="N dimension of the GEMM") + parser.add_argument("--k", type=int, required=True, help="K dimension of the GEMM") + parser.add_argument( + "--m", + type=int, + nargs="+", + default=None, + help="M values to tune (default: a logarithmic range)", + ) + parser.add_argument("--dtype", choices=["float16", "bfloat16"], default="float16") + parser.add_argument( + "--save-dir", + type=Path, + default=SAVE_DIR, + help="Directory to write the configuration file to", + ) + args = parser.parse_args() + + logging.basicConfig(level=logging.INFO) + + gemm_kernel = kernels.get_kernel("kernels-test/gemm-triton-autotune", version=1) + + tune_kwargs = {} + if args.m is not None: + tune_kwargs["Ms"] = tuple(args.m) + + path = gemm_kernel.tune_gemm( + N=args.n, + K=args.k, + dtype=getattr(torch, args.dtype), + save_dir=args.save_dir, + **tune_kwargs, + ) + print(f"Configurations written to {path}") + + +if __name__ == "__main__": + main() diff --git a/nix-builder/tests/Dockerfile.test-kernel b/nix-builder/tests/Dockerfile.test-kernel index 2a605290..6f7dae74 100644 --- a/nix-builder/tests/Dockerfile.test-kernel +++ b/nix-builder/tests/Dockerfile.test-kernel @@ -87,6 +87,7 @@ COPY examples/kernels/cutlass-gemm/tests ./cutlass_gemm_tests COPY examples/kernels/cutlass-gemm-tvm-ffi/tests ./cutlass_gemm_tvm_ffi_tests COPY examples/kernels/cpp20-symbols/tests ./cpp20_symbols_tests COPY examples/kernels/relu-triton/tests ./relu_triton_tests +COPY examples/kernels/gemm-triton-autotune/tests ./gemm_triton_autotune_tests # Run tests ADD nix-builder/tests/run-tests.sh ./run-tests.sh diff --git a/nix-builder/tests/run-tests.sh b/nix-builder/tests/run-tests.sh index 33426b81..7374aba1 100644 --- a/nix-builder/tests/run-tests.sh +++ b/nix-builder/tests/run-tests.sh @@ -9,6 +9,7 @@ RELU_TVM_FFI_PATH=$(echo relu-tvm-ffi-kernel/tvm-ffi*) CUTLASS_PATH=$(echo cutlass-gemm-kernel/torch*) CUTLASS_TVM_FFI_PATH=$(echo cutlass-gemm-tvm-ffi-kernel/tvm-ffi*) RELU_TRITON_PATH=$(echo relu-triton-kernel/torch*) +GEMM_TRITON_AUTOTUNE_PATH=$(echo gemm-triton-autotune-kernel/torch*) SILU_MUL_PATH=$(echo silu-and-mul-kernel/torch*) RELU_CPU_PATH=$(echo relu-kernel-cpu/torch*) CPP20_SYMBOLS_PATH=$(echo cpp20-symbols-kernel/torch*) @@ -19,6 +20,9 @@ LOCAL_KERNELS="kernels-test/extra-data=${EXTRA_DATA_PATH}:kernels-test/relu=${RE LOCAL_KERNELS="kernels-test/relu-triton=${RELU_TRITON_PATH}" \ .venv/bin/pytest relu_triton_tests +LOCAL_KERNELS="kernels-test/gemm-triton-autotune=${GEMM_TRITON_AUTOTUNE_PATH}" \ + .venv/bin/pytest gemm_triton_autotune_tests + # We only care about importing, the kernel is trivial. LOCAL_KERNELS="kernels-test/silu-and-mul=${SILU_MUL_PATH}" \ .venv/bin/python -c "import kernels; kernels.get_kernel('kernels-test/silu-and-mul', version=1)" From 8b46fb7febfe5023642888ba083e9f87069c5717 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Mon, 10 Aug 2026 09:11:03 +0000 Subject: [PATCH 2/9] update --- docs/source/builder/triton-autotune.md | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/docs/source/builder/triton-autotune.md b/docs/source/builder/triton-autotune.md index 27caa6e8..ae577451 100644 --- a/docs/source/builder/triton-autotune.md +++ b/docs/source/builder/triton-autotune.md @@ -6,8 +6,13 @@ 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. -For kernels on the Hub, there is a better option: run the autotuner once per -GPU model, store the best configurations as JSON files, and ship those files + + +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. + +`kernels` 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. @@ -42,9 +47,6 @@ repo-id = "kernels-test/gemm-triton-autotune" pyext = ["json", "py"] ``` -This works the same for AOT-compiled kernels — the `torch` section also -supports `pyext`. - ## Configuration file layout A GEMM computes `(M, K) @ (K, N)`. For a model, the weight dimensions `N` @@ -166,7 +168,7 @@ build: $ LOCAL_KERNELS=kernels-test/gemm-triton-autotune=build python tune.py --n 4096 --k 4096 ``` -## Does it matter? +## 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)` From 132dc1aadd780f30b503b8fbca327cf7a164f986 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Mon, 10 Aug 2026 09:11:21 +0000 Subject: [PATCH 3/9] remove plan --- PLAN-issue-733.md | 168 ---------------------------------------------- 1 file changed, 168 deletions(-) delete mode 100644 PLAN-issue-733.md diff --git a/PLAN-issue-733.md b/PLAN-issue-733.md deleted file mode 100644 index 12b1dd9a..00000000 --- a/PLAN-issue-733.md +++ /dev/null @@ -1,168 +0,0 @@ -# Plan: Triton autotune configs example (issue #733) - -**Issue:** [huggingface/kernels#733](https://github.com/huggingface/kernels/issues/733) — -"Document/example: shipping Triton autotune configs with a kernel" - -**Goal:** Add an example Triton kernel to `examples/kernels/` that ships pre-computed -autotune configurations as JSON files, plus a script that generates those files and -documentation explaining the pattern. Reference implementation: -[`RedHatAI/moe`](https://huggingface.co/RedHatAI/moe/tree/main/torch-ext/moe) (vLLM fused MoE). - -## Background / findings from the repo - -- `examples/kernels/relu-triton/` is the existing Triton example: a `torch-noarch` - kernel with `torch.library.custom_op` wrappers. It is built in CI for CUDA, ROCm, - and XPU via the `ciKernels` / `ciRocmKernels` / `ciXpuKernels` lists in - `examples/kernels/flake.nix`, and its tests run on GPU through - `nix-builder/tests/Dockerfile.test-kernel` + `run-tests.sh` using the - `LOCAL_KERNELS` env var (`kernels-test/=`). -- `examples/kernels/extra-data/` already demonstrates shipping non-Python data by - adding `"json"` to `pyext` in `build.toml` — the exact mechanism we need for - config files. -- The vLLM MoE pattern (`fused_moe.py`): config files live in - `torch-ext/moe/configs/` named - `E=...,N=...,device_name=NVIDIA_H100_80GB_HBM3[,dtype=...].json`. Each file maps a - batch-size bucket (`M`) to a Triton config dict (`BLOCK_SIZE_M/N/K`, `GROUP_SIZE_M`, - `num_warps`, `num_stages`). At runtime an `lru_cache`d loader looks up the file for - the current device; if absent, it falls back to a heuristic default and logs a - warning. -- A local NVIDIA L4 GPU is available — the same GPU as the CI test runners - (`aws-g6-12xlarge`) — so we can generate and commit a real config that CI will - actually exercise. - -## Design decisions - -1. **Kernel: a Triton GEMM (`matmul`)**, not another relu. Autotuning is only - meaningful when tile sizes / warps / stages matter; GEMM is the canonical case and - mirrors the MoE reference. Kernel body can be the standard Triton matmul tutorial - kernel (fp16/bf16/fp32 inputs, fp16-accumulate-in-fp32). -2. **Ship JSON lookup tables instead of relying on `@triton.autotune`.** The - decorator re-benchmarks in every process (slow startup, nondeterministic) and its - cache is not portable. Instead: an offline tuning script does the - `@triton.autotune`-style grid search once, writes JSON, and the kernel loads the - JSON at call time. This is exactly what the issue asks to demonstrate. -3. **File naming convention** (adapted from vLLM MoE): since GEMM weight dims are - known ahead of time but the batch dim `M` varies at runtime, one file per - `(N, K, device)`: `configs/N=4096,K=4096,device_name=NVIDIA_L4.json`, whose keys - are `M` values and values are config dicts. Runtime picks the nearest `M` key - (same `min(keys, key=|log(M/key)|)` trick as vLLM). -4. **Tuning entry point ships with the kernel** (a `tune.py` module inside the - package) so users can regenerate configs for *their* GPU after - `get_kernel(...)`, plus a thin CLI script in the example root for kernel authors. - -## New files - -``` -examples/kernels/gemm-triton-autotune/ -├── build.toml # torch-noarch, pyext = ["py", "json"] -├── flake.nix # copied from relu-triton -├── CARD.md # copied from relu-triton -├── tune.py # CLI: python tune.py --n 4096 --k 4096 [--out ...] -│ # writes into torch-ext/gemm_triton_autotune/configs/ -├── torch-ext/gemm_triton_autotune/ -│ ├── __init__.py # exports gemm(), tune_gemm() -│ ├── gemm.py # @triton.jit matmul kernel + custom_op wrapper; -│ │ # per-call config lookup via tuning.get_config() -│ ├── tuning.py # device-name helper, config file naming, -│ │ # lru_cache'd JSON loader, default-config fallback -│ │ # (with one-time warning), tune_gemm() grid search -│ │ # using triton.testing.do_bench -│ └── configs/ -│ └── N=4096,K=4096,device_name=NVIDIA_L4.json # generated on local L4 -└── tests/ - ├── __init__.py - ├── conftest.py # device fixture, same as relu-triton - └── test_gemm.py # correctness vs torch.matmul across dtypes/shapes; - # shipped-config load test (monkeypatched device name); - # fallback-to-default test for unknown (N, K) -``` - -### `build.toml` sketch - -```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 = ["py", "json"] -``` - -### Runtime config lookup (in `tuning.py`) - -```python -@functools.lru_cache -def get_config(M: int, N: int, K: int) -> dict: - path = Path(__file__).parent / "configs" / _config_file_name(N, K) - if path.exists(): - configs = {int(m): cfg for m, cfg in json.loads(path.read_text()).items()} - return configs[min(configs, key=lambda m: abs(math.log(M / m)))] - warnings.warn(f"No tuned GEMM config for {path.name}, using defaults...", once) - return _default_config(M, N, K) -``` - -### Tuning script behavior - -- Candidate grid: the usual matmul space (`BLOCK_M/N/K ∈ {32..256}`, `GROUP_M`, - `num_warps ∈ {4, 8}`, `num_stages ∈ {2..5}`), pruned to valid combos. -- Benchmarks each candidate with `triton.testing.do_bench` for each `M` in - `{1, 16, 64, 256, 1024, 4096}` at fixed `(N, K)`. -- Writes `{M: best_config}` JSON to the package `configs/` dir, named with - `torch.cuda.get_device_name().replace(" ", "_")` (XPU equivalent when applicable). - -## Existing files to modify - -1. **`examples/kernels/flake.nix`** — register the new kernel in `ciKernels` - (`torch-cuda` noarch build, like `relu-triton-kernel`), `ciRocmKernels`, and - `ciXpuKernels`. -2. **`.github/workflows/build_kernel.yaml`** — add `gemm-triton-autotune-kernel` to - the uploaded-artifacts list. -3. **`nix-builder/tests/Dockerfile.test-kernel`** — `COPY - examples/kernels/gemm-triton-autotune/tests ./gemm_triton_autotune_tests`. -4. **`nix-builder/tests/run-tests.sh`** — run the new tests with - `LOCAL_KERNELS="kernels-test/gemm-triton-autotune=..."`. -5. **Docs:** - - New page `docs/source/builder/triton-autotune.md` — "Shipping Triton autotune - configurations": why ship configs, the JSON-per-device pattern, `pyext = ["json"]`, - the loader/fallback pattern, how to run the tune script, links to the example - and to `RedHatAI/moe`. - - Add the page to `docs/source/_toctree.yml` (kernel-builder section, after - `builder/writing-kernels`). - -## Implementation order - -1. Scaffold the example kernel (build.toml, flake.nix, CARD.md, package code). -2. Set up a local venv (`uv venv` + torch/cu126 + triton + kernels + pytest) and get the - kernel running directly from `torch-ext/` on the L4. -3. Run `tune.py` on the L4 for `(N=4096, K=4096)`; commit the generated JSON. -4. Write tests; run them locally against the local build - (`LOCAL_KERNELS=kernels-test/gemm-triton-autotune=` after a - `kernels build`/nix build, matching the CI invocation). -5. Wire up CI (flake.nix lists, workflow artifact list, Dockerfile, run-tests.sh). -6. Write the docs page + toctree entry. -7. `nix flake check` / build the example via - `nix build ./examples/kernels#ci-build-cuda` if feasible locally, else rely on CI. - -## Out of scope / maintainer follow-ups - -- Pushing the built kernel to the `kernels-test/gemm-triton-autotune` Hub repo - (needed for `get_kernel` without `LOCAL_KERNELS`) requires org access — CI tests - use `LOCAL_KERNELS`, so nothing blocks on this, but the Hub repo should be created - when merging (same as other `kernels-test/*` examples). -- Configs for ROCm/XPU devices can be contributed later by whoever has the hardware; - the fallback path covers them meanwhile (and the fallback is itself part of what - the example demonstrates). - -## Open questions - -1. Kernel/package name: `gemm-triton-autotune` (proposed) vs `matmul-triton-tune`. -2. Should the docs page live under kernel-builder docs (proposed) or as a section - appended to `writing-kernels.md`? -3. Single `(N, K)` shape for the committed config (proposed: 4096×4096) or a couple - of shapes to show multiple config files? From 9805ed2c032ce827642b110cb8d82ba935d28aa6 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Mon, 10 Aug 2026 10:06:25 +0000 Subject: [PATCH 4/9] fix: copy gemm-triton-autotune built kernel into the test image The tests directory was copied into the Docker test image, but the built kernel artifact was not, so LOCAL_KERNELS pointed at an unexpanded glob. Co-Authored-By: Claude Fable 5 --- nix-builder/tests/Dockerfile.test-kernel | 1 + 1 file changed, 1 insertion(+) diff --git a/nix-builder/tests/Dockerfile.test-kernel b/nix-builder/tests/Dockerfile.test-kernel index 6f7dae74..7329b5cb 100644 --- a/nix-builder/tests/Dockerfile.test-kernel +++ b/nix-builder/tests/Dockerfile.test-kernel @@ -77,6 +77,7 @@ COPY relu-kernel-cpu ./relu-kernel-cpu COPY cutlass-gemm-kernel ./cutlass-gemm-kernel COPY cutlass-gemm-tvm-ffi-kernel ./cutlass-gemm-tvm-ffi-kernel COPY relu-triton-kernel ./relu-triton-kernel +COPY gemm-triton-autotune-kernel ./gemm-triton-autotune-kernel COPY silu-and-mul-kernel ./silu-and-mul-kernel COPY extra-data ./extra-data COPY cpp20-symbols-kernel ./cpp20-symbols-kernel From 2cd8593115c4bd43f1d73eadfb8c35520351f61d Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Mon, 17 Aug 2026 11:27:17 +0200 Subject: [PATCH 5/9] Apply suggestions from code review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Daniël de Kok Co-authored-by: Sayak Paul --- docs/source/builder/triton-autotune.md | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/docs/source/builder/triton-autotune.md b/docs/source/builder/triton-autotune.md index ae577451..04f61998 100644 --- a/docs/source/builder/triton-autotune.md +++ b/docs/source/builder/triton-autotune.md @@ -1,7 +1,7 @@ # Ship Triton autotune configurations -Triton kernels typically have parameters — tile sizes, number of warps, -number of pipeline stages — whose optimal values depend on the GPU and the +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) @@ -12,7 +12,7 @@ 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. -`kernels` support this by packaging these configurations as JSON files +`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. @@ -27,9 +27,10 @@ example kernel, a Triton GEMM published as ## Shipping data files with a kernel -Configuration files are plain JSON files inside the kernel's Python package, -in `torch-ext//configs/`. By default, only `py` and `pyi` -files are picked up from the package directory, so add `json` to the +Since autotune files are plain JSON files, we can store them anywhere inside the +kernels main Python sources in `torch-ext/`. For this example, +we will use `torch-ext//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 @@ -97,7 +98,7 @@ and logs a warning. The lookup is cached, so the file is read at most once per process: ```python -@functools.lru_cache +@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(): @@ -117,7 +118,7 @@ def _load_tuned_configs(N: int, K: int) -> Optional[Dict[int, Dict[str, int]]]: def get_config(M: int, N: int, K: int) -> Dict[str, int]: tuned = _load_tuned_configs(N, K) - if tuned: + 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] From f91727cc6fcf097b85588a3a7fca98708fb0ecfd Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Mon, 17 Aug 2026 11:30:08 +0200 Subject: [PATCH 6/9] remove tip from writing kernels. --- docs/source/builder/writing-kernels.md | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/docs/source/builder/writing-kernels.md b/docs/source/builder/writing-kernels.md index 88d19849..30679def 100644 --- a/docs/source/builder/writing-kernels.md +++ b/docs/source/builder/writing-kernels.md @@ -37,9 +37,7 @@ format of the `build.toml` file, and some additional Python glue that `kernel-builder` provides. We will use a [simple ReLU kernel](https://github.com/huggingface/kernels/tree/main/examples/kernels/relu) as the running example. After reading this page, you may also want to have a look at the more realistic [ReLU kernel with backprop and `torch.compile`](https://github.com/huggingface/kernels/tree/main/examples/kernels/relu-backprop-compile) -support. For Triton kernels, see -[Ship Triton autotune configurations](triton-autotune.md) for how to tune a -kernel and ship the tuned configurations with it. +support. > [!TIP] > We maintain a set of conforming kernels in the From 9829990f9b2a8b183c41c55d358a39810feb57bc Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Mon, 17 Aug 2026 11:32:48 +0200 Subject: [PATCH 7/9] move autotuning guide to advanced practices' section. --- docs/source/_toctree.yml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 008d0d21..4dd64ea9 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -70,6 +70,10 @@ - local: cli-skills title: Install agent skills title: kernel-builder CLI +- section: + - local: builder/triton-autotune + title: Ship Triton autotune configurations + title: Advanced practices - sections: - local: kernel-requirements title: Kernel requirements From 4b024ff40e47515ff64195c367926047de98798a Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 18 Aug 2026 13:42:28 +0200 Subject: [PATCH 8/9] remove duplicate section. --- docs/source/_toctree.yml | 2 -- 1 file changed, 2 deletions(-) diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 4dd64ea9..e75fddff 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -47,8 +47,6 @@ - sections: - local: builder/writing-kernels title: Write kernels - - local: builder/triton-autotune - title: Ship Triton autotune configurations - local: builder/build title: Build with Nix - local: builder/local-dev From 28b6d51150523862d943f9fe20cd4c0d7b6bba14 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 18 Aug 2026 13:48:51 +0200 Subject: [PATCH 9/9] fix sections --- docs/source/_toctree.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index e75fddff..c55f0c2e 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -68,7 +68,7 @@ - local: cli-skills title: Install agent skills title: kernel-builder CLI -- section: +- sections: - local: builder/triton-autotune title: Ship Triton autotune configurations title: Advanced practices